diff --git a/batchgen_kernels/setup.py b/batchgen_kernels/setup.py index 7f6de2477..308c10d24 100644 --- a/batchgen_kernels/setup.py +++ b/batchgen_kernels/setup.py @@ -111,6 +111,94 @@ def _setup_ccache(): _sm80_flags = ["-std=c++17", "-O3", "--threads", _nvcc_threads] + _sm80_gencode +# ── SM100 (Blackwell / B200) build flags ───────────────────────────────────── + +_sm100_flags = [ + "-std=c++17", + "-arch=sm_100a", + "-O3", + "--use_fast_math", + "-lineinfo", + "--ptxas-options=-v", + "--threads", _nvcc_threads, +] + +# ── SM100 extensions: generic CUDA + SM80 mma.sync kernels (no WGMMA) ──────── +# Kernels that use SM90a WGMMA/TMA are EXCLUDED from this list and replaced by +# Triton kernels (qkv_proj_rope, fused_router_gemm) or torch._scaled_mm (fp8). + +_sm100_extensions = [ + # Marlin W4A16 grouped GEMM — SM80 mma.sync, backward compatible with SM100 + CUDAExtension( + name="batchgen_kernels.moe._C_marlin_grouped_gemm", + sources=["src/moe/marlin_grouped_gemm.cu"], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": _sm100_flags + [ + "-DUSE_BF16_COMPUTE", + "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", + ], + }, + ), + # FP8 blockwise pipeline ops (act_quant_3d, silu_mul_3d, fused_silu_quant_3d) + CUDAExtension( + name="batchgen_kernels.moe._C_fp8_blockwise_ops", + sources=["src/moe/fp8_blockwise/fp8_blockwise_ops.cu"], + extra_compile_args={"cxx": ["-O3"], "nvcc": _sm100_flags}, + ), + # Marlin weight-format transform — generic CUDA, no ISA deps + CUDAExtension( + name="batchgen_kernels.moe._C_marlin_transform", + sources=["src/moe/marlin_transform_kernel.cu"], + extra_compile_args={"cxx": ["-O3"], "nvcc": _sm100_flags}, + ), + # Routing bundle (generic kernels only; fused_gate.cu EXCLUDED — WGMMA source; + # replaced by batchgen_kernels/triton/fused_router_gemm.py on SM100) + CUDAExtension( + name="batchgen_kernels.moe._C_routing", + sources=[ + "src/moe/routing/routing_extension.cc", + "src/moe/routing/gate_topk_softmax.cu", + "src/moe/routing/dispatch_count_gather.cu", + "src/moe/routing/reduce_weighted_scatter.cu", + "src/moe/routing/router_epilogue.cu", + "src/moe/routing/gate_sigmoid_topk.cu", + "src/moe/routing/glm5_router_gemm.cu", # generic CUDA, no WGMMA + ], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": _sm100_flags, + }, + ), + # 3D dispatch scatter + reduce (strided MoE buffer) — generic CUDA + CUDAExtension( + name="batchgen_kernels.moe._C_dispatch_scatter_3d", + sources=["src/moe/dispatch_scatter_3d.cu"], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": _sm100_flags + ["-U__CUDA_NO_BFLOAT16_CONVERSIONS__"], + }, + ), + # Fused RMSNorm + RoPE + KV cache write — generic CUDA + CUDAExtension( + name="batchgen_kernels.attention._C_fused_kv_norm_rope", + sources=["src/attention/fused_kv_norm_rope_cache.cu"], + extra_compile_args={"cxx": ["-O3"], "nvcc": _sm100_flags}, + ), + # Fused q_absorb GEMV + q_pe copy — generic CUDA + CUDAExtension( + name="batchgen_kernels.attention._C_fused_q_absorb", + sources=["src/attention/fused_q_absorb.cu"], + extra_compile_args={"cxx": ["-O3"], "nvcc": _sm100_flags}, + ), + # Fused q_b split into q_nope + q_pe — generic CUDA + CUDAExtension( + name="batchgen_kernels.attention._C_fused_q_split", + sources=["src/attention/fused_q_split.cu"], + extra_compile_args={"cxx": ["-O3"], "nvcc": _sm100_flags}, + ), +] + # ── Build extension list ── _sm90a_extensions = [ @@ -216,8 +304,9 @@ def _setup_ccache(): "src/moe/routing/fused_gate.cu", ], extra_compile_args={ - "cxx": ["-O3"], + "cxx": ["-O3", "-DBATCHGEN_HAS_FUSED_GATE"], "nvcc": ["-O3", "--use_fast_math", "-std=c++17", + "-DBATCHGEN_HAS_FUSED_GATE", "-gencode", "arch=compute_90a,code=sm_90a", "--threads", _nvcc_threads], }, @@ -376,8 +465,10 @@ def _setup_ccache(): _ext_modules = [] if _build_sm90a: _ext_modules.extend(_sm90a_extensions) +elif _build_sm100: + _ext_modules.extend(_sm100_extensions) else: - print(f"[batchgen_kernels] BUILD_ARCH={_build_arch}: skipping SM90a-only kernels") + print(f"[batchgen_kernels] BUILD_ARCH={_build_arch}: skipping arch-specific kernels") _ext_modules.extend(_sm80_extensions) setup( diff --git a/batchgen_kernels/src/moe/routing/routing_extension.cc b/batchgen_kernels/src/moe/routing/routing_extension.cc index 528a2eba6..9de6a6189 100644 --- a/batchgen_kernels/src/moe/routing/routing_extension.cc +++ b/batchgen_kernels/src/moe/routing/routing_extension.cc @@ -54,6 +54,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("bucket_size"), py::arg("world_size")); +#ifdef BATCHGEN_HAS_FUSED_GATE m.def("create_fused_gate_context", &create_fused_gate_context, "Create cached fused gate context (SM90a WGMMA)", py::arg("router_weight"), @@ -77,6 +78,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("topk_indices"), py::arg("topk_weights"), py::arg("num_valid_tokens") = -1); +#endif // BATCHGEN_HAS_FUSED_GATE m.def("reduce_weighted_scatter", &reduce_weighted_scatter_cuda, "Reduce: weighted scatter-add (CUDA)", diff --git a/docs/BLACKWELL_KERNELS_WIP.md b/docs/BLACKWELL_KERNELS_WIP.md new file mode 100644 index 000000000..02c848966 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 setup.py _sm100_extensions infra skeleton