[enhancement] Integrate C++ compilation in setup.py to streamline installation

[tscholak]

🧐 Problem Description

The current installation process for Fast-LLM requires a separate make step to compile a single C++ file located in fast_llm/csrc/. This can be frustrating for users installing Fast-LLM directly via pip, as it introduces an extra step. This step is usually hidden in the Docker image workflow, but becomes problematic outside of it, especially during direct installs.

Example command:

make -C ./fast_llm/csrc/

💡 Proposed Solution

Integrate the compilation step within setup.py using setuptools, making installation seamless and consistent across environments. This would align with typical practices for compiling extension modules and eliminate the need for the manual make step.

🔄 Alternatives Considered

One alternative is to separate this component into a standalone module or dependency, but this might overcomplicate the installation process and the codebase. Integrating it into setup.py is likely the most straightforward and user-friendly solution.

📈 Potential Benefits

• Streamlined installation: Simplifies the installation for users who prefer direct installs without using Docker.
• Improved usability: Removes the need for users to manually compile C++ code during setup.
• Consistency: Matches the setup experience of other Python packages with compiled extensions.

📝 Additional Context

Currently writing documentation for Fast-LLM, and this step complicates installation instructions. Simplifying the setup would make documentation more straightforward and reduce potential user issues.

[tscholak]

Here's a minimal setup.py that illustrates the change:

from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import sys
import os
import pybind11

# Define the C++ extension with required flags
cpp_extension = Extension(
    "fast_llm.csrc.data",  # Specify the path for the compiled module
    sources=["fast_llm/csrc/data.cpp"],  # Path to the C++ source file
    include_dirs=[pybind11.get_include()],  # Include path for pybind11
    extra_compile_args=["-O3", "-Wall", "-shared", "-std=c++11", "-fPIC"],
    extra_link_args=["-fdiagnostics-color"],  # Additional linking flags, if needed
    language="c++",
)

# Customize build_ext to apply build settings
class BuildExt(build_ext):
    def build_extensions(self):
        # Optional: Configure more custom compiler settings here
        super().build_extensions()

# Setup function
setup(
    name="fast-llm",
    version="0.1",
    packages=["fast_llm"],  # Adjust this according to your package structure
    ext_modules=[cpp_extension],
    cmdclass={"build_ext": BuildExt},  # Use the customized build command
    install_requires=[
        "pybind11>=2.5.0",  # Ensure pybind11 is installed
    ],
)