from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='gpt_oss_kernels',
    ext_modules=[
        CUDAExtension('gpt_oss_kernels', [
            'kernels/flash_attention_hopper.cu',
            'kernels/flash_attention_hopper_bwd.cu',
            'kernels/mxfp4_quant.cu',
            'kernels/flash_attention.cpp',
        ], extra_compile_args={
            'cxx': ['-O3', '-std=c++17'],
            'nvcc': [
                '-O3', '-std=c++17',
                '-arch=sm_90a', # Hopper
                '-DCUDA_ENABLE_TMA=1',
                '-D__CUDA_BF16_H__',
                '-lineinfo',
                '-maxrregcount=255',
                '--ptxas-options=-v'
            ]
        })
    ],
    cmdclass={'build_ext': BuildExtension}
)