chore(setup): properly package the repository as a Python package

[XuehaiPan]

The current package layout is:

• flash_mla Python code is installed in the site-packages directory. The in-tree source flash_mla has higher import precedence if the working directory is the project root.
• flash_mla_cuda is a separate .so extension placed in the site-packages directory which is neither under site-packages/flash_mla nor under <Project Root>/flash_mla.

This PR makes the package pip installable and can be redistributed as wheels.

Changes:

• Add pyproject.toml file to specify the Python build system to replace the deprecated setup.py method.
• Change the install command from deprecated python setup.py install to pip install [--no-build-isolation] --editable ..
• Move the .so shared library under the flash_mla namespace.
• Add typing stub flash_mla_cuda.pyi.
• Add auxiliary packaging files MANIFEST.in, py.typed.

Install from source:

$ python3 -m pip install --upgrade pip setuptools
$ python3 -m pip install torch pybind11 --index-url https://download.pytorch.org/whl/cu126
$ python3 -m pip install --no-build-isolation --editable .
Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple
Obtaining file:///...Projects/flash-mla
  Preparing metadata (setup.py) ... done
Installing collected packages: flash-mla
  Attempting uninstall: flash-mla
    Found existing installation: flash-mla 1.0.0+2d049c8
    Uninstalling flash-mla-1.0.0+2d049c8:
      Successfully uninstalled flash-mla-1.0.0+2d049c8
  Running setup.py develop for flash-mla
Successfully installed flash-mla-1.0.0+dc080a3

Build wheels:

$ python3 -m build --wheel --no-isolation
* Getting build dependencies for wheel...
running egg_info
writing flash_mla.egg-info/PKG-INFO
writing dependency_links to flash_mla.egg-info/dependency_links.txt
writing top-level names to flash_mla.egg-info/top_level.txt
reading manifest file 'flash_mla.egg-info/SOURCES.txt'
reading manifest template 'MANIFEST.in'
adding license file 'LICENSE'
writing manifest file 'flash_mla.egg-info/SOURCES.txt'
* Building wheel...
running bdist_wheel
running build
running build_py
copying flash_mla/__init__.py -> build/lib.linux-x86_64-cpython-312/flash_mla
copying flash_mla/flash_mla_interface.py -> build/lib.linux-x86_64-cpython-312/flash_mla
copying flash_mla/flash_mla_cuda.pyi -> build/lib.linux-x86_64-cpython-312/flash_mla
copying flash_mla/py.typed -> build/lib.linux-x86_64-cpython-312/flash_mla
running build_ext
building 'flash_mla.flash_mla_cuda' extension
Emitting ninja build file /.../Projects/flash-mla/build/temp.linux-x86_64-cpython-312/build.ninja...
Compiling objects...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
/usr/bin/g++ -fno-strict-overflow -Wsign-compare -DNDEBUG -g -O3 -Wall -shared /.../Projects/flash-mla/build/temp.linux-x86_64-cpython-312/csrc/flash_api.o /.../Projects/flash-mla/build/temp.linux-x86_64-cpython-312/csrc/flash_fwd_mla_bf16_sm90.o -L/.../Projects/flash-mla/venv/lib/python3.12/site-packages/torch/lib -L/usr/local/cuda/lib64 -L/home/linuxbrew/.linuxbrew/opt/[email protected]/lib -lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda -o build/lib.linux-x86_64-cpython-312/flash_mla/flash_mla_cuda.cpython-312-x86_64-linux-gnu.so
installing to build/bdist.linux-x86_64/wheel
running install
running install_lib
creating build/bdist.linux-x86_64/wheel
creating build/bdist.linux-x86_64/wheel/flash_mla
copying build/lib.linux-x86_64-cpython-312/flash_mla/py.typed -> build/bdist.linux-x86_64/wheel/./flash_mla
copying build/lib.linux-x86_64-cpython-312/flash_mla/__init__.py -> build/bdist.linux-x86_64/wheel/./flash_mla
copying build/lib.linux-x86_64-cpython-312/flash_mla/flash_mla_interface.py -> build/bdist.linux-x86_64/wheel/./flash_mla
copying build/lib.linux-x86_64-cpython-312/flash_mla/flash_mla_cuda.cpython-312-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/wheel/./flash_mla
copying build/lib.linux-x86_64-cpython-312/flash_mla/flash_mla_cuda.pyi -> build/bdist.linux-x86_64/wheel/./flash_mla
running install_egg_info
running egg_info
writing flash_mla.egg-info/PKG-INFO
writing dependency_links to flash_mla.egg-info/dependency_links.txt
writing top-level names to flash_mla.egg-info/top_level.txt
reading manifest file 'flash_mla.egg-info/SOURCES.txt'
reading manifest template 'MANIFEST.in'
adding license file 'LICENSE'
writing manifest file 'flash_mla.egg-info/SOURCES.txt'
Copying flash_mla.egg-info to build/bdist.linux-x86_64/wheel/./flash_mla-1.0.0+dc080a3-py3.12.egg-info
running install_scripts
creating build/bdist.linux-x86_64/wheel/flash_mla-1.0.0+dc080a3.dist-info/WHEEL
creating '/.../Projects/flash-mla/dist/.tmp-k52wtra0/flash_mla-1.0.0+dc080a3-cp312-cp312-linux_x86_64.whl' and adding 'build/bdist.linux-x86_64/wheel' to it
adding 'flash_mla/__init__.py'
adding 'flash_mla/flash_mla_cuda.cpython-312-x86_64-linux-gnu.so'
adding 'flash_mla/flash_mla_cuda.pyi'
adding 'flash_mla/flash_mla_interface.py'
adding 'flash_mla/py.typed'
adding 'flash_mla-1.0.0+dc080a3.dist-info/LICENSE'
adding 'flash_mla-1.0.0+dc080a3.dist-info/METADATA'
adding 'flash_mla-1.0.0+dc080a3.dist-info/WHEEL'
adding 'flash_mla-1.0.0+dc080a3.dist-info/top_level.txt'
adding 'flash_mla-1.0.0+dc080a3.dist-info/RECORD'
removing build/bdist.linux-x86_64/wheel
Successfully built flash_mla-1.0.0+dc080a3-cp312-cp312-linux_x86_64.whl

[XuehaiPan]

Any thoughts or comments?