upd

sgl-project · Jan 30, 2025 · d04f525 · d04f525
1 parent fc24280
commit d04f525
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/sgl-kernel/setup.py b/sgl-kernel/setup.py
@@ -128,6 +128,7 @@ def _get_version():
         nvcc_flags.append("-gencode=arch=compute_90a,code=sm_90a")
         nvcc_flags.append("-DCOMPILE_HOPPER_TMA_GEMMS")
         nvcc_flags.append("-DCUTLASS_ARCH_MMA_SM90_SUPPORTED")
+        nvcc_flags.append("-DCUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED")
         sources.extend(sources_cutlass_moe_gemm)
     if sm_version >= 90:
         nvcc_flags.extend(nvcc_flags_fp8)
@@ -142,6 +143,7 @@ def _get_version():
         nvcc_flags.append("-gencode=arch=compute_90a,code=sm_90a")
         nvcc_flags.append("-DCOMPILE_HOPPER_TMA_GEMMS")
         nvcc_flags.append("-DCUTLASS_ARCH_MMA_SM90_SUPPORTED")
+        nvcc_flags.append("-DCUTLASS_ARCH_MMA_MODIFIABLE_TMA_SM90_SUPPORTED")
         sources.extend(sources_cutlass_moe_gemm)
     if enable_fp8:
         nvcc_flags.extend(nvcc_flags_fp8)