From a4eccac38ffd04f9fed5aa2dd5b245b793ad686d Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:13 +0000 Subject: [PATCH 1/2] [wip] blackwell: sm100 setup.py _sm100_extensions infra skeleton Placeholder commit for draft PR. Implementation tracked in: batchgen-agent-metadata/batchgen_design/blackwell/blackwell-kernel-port-v1.md Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/BLACKWELL_KERNELS_WIP.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 docs/BLACKWELL_KERNELS_WIP.md diff --git a/docs/BLACKWELL_KERNELS_WIP.md b/docs/BLACKWELL_KERNELS_WIP.md new file mode 100644 index 00000000..02c84896 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 setup.py _sm100_extensions infra skeleton From a811b75a3bdff8ebb50194c5f7800644e334a435 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:09:30 +0000 Subject: [PATCH 2/2] blackwell(sm100): add _sm100_extensions build infra + recompile generic kernels setup.py: - Add _sm100_flags (-arch=sm_100a, -O3, --use_fast_math) and _sm100_extensions list containing the 8 generic-CUDA / SM80-mma.sync extensions that are SM100-portable (no SM90a WGMMA/TMA): marlin_grouped_gemm (carries -DUSE_BF16_COMPUTE), fp8_blockwise_ops, marlin_transform, _C_routing (generic kernels only), dispatch_scatter_3d, fused_kv_norm_rope, fused_q_absorb, fused_q_split. - Wire assembly: if _build_sm90a / elif _build_sm100 / always _sm80. - WGMMA sources (fused_gate.cu, qkv_wgmma.cu, glm5 uses generic so kept) are excluded from the SM100 routing bundle; glm5_router_gemm.cu is generic CUDA (warp-shuffle only) and retained. routing_extension.cc: - Guard the WGMMA-backed fused_gate_* pybind bindings behind #ifdef BATCHGEN_HAS_FUSED_GATE so the SM100 build (which omits fused_gate.cu) does not emit an undefined-symbol import error. SM90a build defines the macro. Verified: BUILD_ARCH=sm100 builds all 8 extensions; all load on B200 (sm_100). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen_kernels/setup.py | 95 ++++++++++++++++++- .../src/moe/routing/routing_extension.cc | 2 + 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/batchgen_kernels/setup.py b/batchgen_kernels/setup.py index 7f6de247..308c10d2 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 528a2eba..9de6a618 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)",