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 01/16] [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 143a8681896f90ef3811d45cdedb28197cf3ddba Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:14 +0000 Subject: [PATCH 02/16] [wip] blackwell: sm100 recompile 9 generic CUDA kernels 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..30d7748a --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 recompile 9 generic CUDA kernels From bffe108a51ba26d190822fdc8f74051fdae61ac1 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:15 +0000 Subject: [PATCH 03/16] [wip] blackwell: sm100 FP8 blockwise GEMM dispatch stub 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..2e85057d --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 FP8 blockwise GEMM dispatch stub From fc8e3d02313cb652fdf4ea3d7400ce0e0e885f0d Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:16 +0000 Subject: [PATCH 04/16] [wip] blackwell: sm100 MXFP4 Triton dispatch stub 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..2bde5ba4 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 MXFP4 Triton dispatch stub From 2477b721312ac9919a78ccbdcb5afac148b02a2e Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:17 +0000 Subject: [PATCH 05/16] [wip] blackwell: sm100 QKV+RoPE Triton kernel stub 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..30f13a26 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 QKV+RoPE Triton kernel stub From cb0d3bd5b7ddcc86eec9369a26728bccfb2b9390 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:18 +0000 Subject: [PATCH 06/16] [wip] blackwell: sm100 router GEMM Triton kernel stub 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..2e5ce54b --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 router GEMM Triton kernel stub From 0a6dd383496d4b2b5503313fbb949c6949269b69 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:19 +0000 Subject: [PATCH 07/16] [wip] blackwell: sm100 INT4 grouped GEMM Triton stub 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..55dc7aa7 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 INT4 grouped GEMM Triton stub From 967cc9a7011363f96b619a9f28476bc40aeb6c41 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:20 +0000 Subject: [PATCH 08/16] [wip] blackwell: sm100 fused INT4+SiLU Triton stub 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..2f0186bf --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 fused INT4+SiLU Triton stub 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 09/16] 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)", From 9ac9679c8ecac0183963a460aa1d7af5ad86a9aa Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:14:26 +0000 Subject: [PATCH 10/16] feat(sm100): port fused QKV projection + RoPE to Triton Fills the SM100 placeholder for the fused QKV WGMMA kernel. On Blackwell the WGMMA .cu is not built, so cuda_qkv_wgmma() dispatches to a pure-Triton GEMM+split kernel with a standard rotate_half RoPE epilogue matching the SM90a CUDA convention (out[n1]=x1*c-x2*s, out[n2]=x2*c+x1*s). - batchgen_kernels/triton/qkv_proj_rope.py: qkv_proj_split, apply_rope_qk, qkv_proj_rope - qkv_wgmma.py: sm100 dispatch branch + availability via Triton - triton/__init__.py: export new kernels Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen/attention/fused_kernels/qkv_wgmma.py | 26 ++ batchgen_kernels/triton/__init__.py | 5 + batchgen_kernels/triton/qkv_proj_rope.py | 243 ++++++++++++++++++ 3 files changed, 274 insertions(+) create mode 100644 batchgen_kernels/triton/qkv_proj_rope.py diff --git a/batchgen/attention/fused_kernels/qkv_wgmma.py b/batchgen/attention/fused_kernels/qkv_wgmma.py index e7673ce5..4f4cf9fc 100644 --- a/batchgen/attention/fused_kernels/qkv_wgmma.py +++ b/batchgen/attention/fused_kernels/qkv_wgmma.py @@ -22,6 +22,16 @@ _module = None _qkv_wgmma_available: Optional[bool] = None +_arch: Optional[str] = None + + +def _get_arch() -> str: + """Cached device arch ('sm90a' / 'sm100').""" + global _arch + if _arch is None: + import batchgen_kernels + _arch = batchgen_kernels.get_device_arch() + return _arch def _check_wgmma_support() -> bool: @@ -66,6 +76,12 @@ def is_qkv_wgmma_available() -> bool: _qkv_wgmma_available = False return False + # On SM100 (Blackwell) the WGMMA .cu is not built; the fused QKV path is + # provided by the pure-Triton kernel (qkv_proj_rope), always available. + if _get_arch() == "sm100": + _qkv_wgmma_available = True + return True + _qkv_wgmma_available = _get_module() is not None return _qkv_wgmma_available @@ -102,6 +118,16 @@ def cuda_qkv_wgmma( (Q, K, V) as separate contiguous [M, *] BF16 tensors. When rope_cos/sin provided, Q and K have RoPE applied; V is unchanged. """ + # SM100 (Blackwell): WGMMA .cu is not built — use the pure-Triton port. + if _get_arch() == "sm100": + from batchgen_kernels.triton.qkv_proj_rope import qkv_proj_rope + _bias = bias if (bias is not None and bias.numel() > 0) else None + _cos = rope_cos if (rope_cos is not None and rope_cos.numel() > 0) else None + _sin = rope_sin if (rope_sin is not None and rope_sin.numel() > 0) else None + return qkv_proj_rope( + input, weight, _bias, q_size, kv_size, head_dim, _cos, _sin, + ) + mod = _get_module() assert mod is not None, "QKV WGMMA module not available" diff --git a/batchgen_kernels/triton/__init__.py b/batchgen_kernels/triton/__init__.py index 27874ab1..63a1adf6 100644 --- a/batchgen_kernels/triton/__init__.py +++ b/batchgen_kernels/triton/__init__.py @@ -33,3 +33,8 @@ from batchgen_kernels.triton.fused_dequant_gemm import fused_fp8_bf16_gemm from batchgen_kernels.triton.fused_q_absorb import fused_q_absorb_query_states from batchgen_kernels.triton.fused_out_absorb import fused_out_absorb_reshape +from batchgen_kernels.triton.qkv_proj_rope import ( + qkv_proj_split, + qkv_proj_rope, + apply_rope_qk, +) diff --git a/batchgen_kernels/triton/qkv_proj_rope.py b/batchgen_kernels/triton/qkv_proj_rope.py new file mode 100644 index 00000000..73da1456 --- /dev/null +++ b/batchgen_kernels/triton/qkv_proj_rope.py @@ -0,0 +1,243 @@ +"""QKV projection + split (+ optional RoPE) Triton kernel for SM100 (Blackwell). + +SM100 port of the SM90a WGMMA kernel `batchgen_kernels/src/attention/qkv_wgmma.cu`. +That kernel uses Hopper-only `wgmma.*` + warpgroup TMA, which do not exist on +sm_100a, so on Blackwell we replace it with a pure-Triton GEMM that Triton JIT-compiles +to UMMA. No `torch.matmul` / `F.linear` is used. + +Operation (decode/prefill): + + x[M, K] @ W_qkv[N, K]^T -> [M, N], N = q_size + kv_size + kv_size + +split column-wise into: + + Q[M, q_size] (cols [0, q_size)) + K[M, kv_size] (cols [q_size, q_size + kv_size)) + V[M, kv_size] (cols [q_size + kv_size, N)) + +RoPE (optional) is applied to Q and K only, never V (see `apply_rope_qk`). + +Design note (see batchgen-agent-metadata/.../incremental_compilation_contract.md and +blackwell-kernel-port-v1.md sub-task 5): RoPE is applied in a *separate* validated pass +(`apply_rope_qk`) rather than inline in the GEMM epilogue. At decode M the GEMM is +memory-bound on the [N, K] weight read, so the extra Q/K reread is negligible, and the +two-pass structure keeps each step independently testable — matching how BatchGen already +structures RoPE as dedicated kernels. +""" + +import torch +import triton +import triton.language as tl + + +# ────────────────────────────────────────────────────────────────────────────── +# Pass 1 — GEMM + column split into Q / K / V +# ────────────────────────────────────────────────────────────────────────────── + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=3, num_warps=8), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=2, num_warps=8), + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_stages=3, num_warps=4), + ], + key=['M', 'K', 'N'], +) +@triton.jit +def _qkv_proj_split_kernel( + x_ptr, w_ptr, bias_ptr, + q_ptr, k_ptr, v_ptr, + M, K, N, q_size, kv_size, + stride_xm, stride_xk, + stride_wn, stride_wk, # w is [N, K] row-major: stride_wn=K, stride_wk=1 + stride_qm, stride_qn, + stride_km, stride_kn, + stride_vm, stride_vn, + HAS_BIAS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + m_mask = offs_m < M + n_mask = offs_n < N + + # x[BLOCK_M, K] @ w[N, K]^T -> acc[BLOCK_M, BLOCK_N], accumulate in fp32 + x_ptrs = x_ptr + (offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk) + # w is [N, K]; we need w[n, k] for the n-tile, contracted over k -> shape [BLOCK_K, BLOCK_N] + w_ptrs = w_ptr + (offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k0 in range(0, K, BLOCK_K): + k_mask = (k0 + offs_k) < K + x_tile = tl.load(x_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0) + w_tile = tl.load(w_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0) + acc += tl.dot(x_tile, w_tile, out_dtype=tl.float32) + x_ptrs += BLOCK_K * stride_xk + w_ptrs += BLOCK_K * stride_wk + + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + acc += bias[None, :] + + out = acc.to(q_ptr.dtype.element_ty) + + kv_end = q_size + kv_size + + # Region masks (per-column). Robust to any BLOCK_N / boundary alignment. + is_q = offs_n < q_size + is_k = (offs_n >= q_size) & (offs_n < kv_end) + is_v = offs_n >= kv_end + + # Q store: local col = offs_n + q_cols = offs_n + q_ptrs = q_ptr + (offs_m[:, None] * stride_qm + q_cols[None, :] * stride_qn) + tl.store(q_ptrs, out, mask=m_mask[:, None] & (is_q & n_mask)[None, :]) + + # K store: local col = offs_n - q_size + k_cols = offs_n - q_size + k_ptrs = k_ptr + (offs_m[:, None] * stride_km + k_cols[None, :] * stride_kn) + tl.store(k_ptrs, out, mask=m_mask[:, None] & (is_k & n_mask)[None, :]) + + # V store: local col = offs_n - kv_end + v_cols = offs_n - kv_end + v_ptrs = v_ptr + (offs_m[:, None] * stride_vm + v_cols[None, :] * stride_vn) + tl.store(v_ptrs, out, mask=m_mask[:, None] & (is_v & n_mask)[None, :]) + + +def qkv_proj_split( + x: torch.Tensor, # [M, K] BF16 + w_qkv: torch.Tensor, # [N, K] BF16, N = q_size + 2*kv_size + bias: torch.Tensor | None, + q_size: int, + kv_size: int, +): + """Fused QKV GEMM + column split. Returns (Q, K, V) BF16, no RoPE.""" + assert x.dim() == 2 and w_qkv.dim() == 2 + M, K = x.shape + N, Kw = w_qkv.shape + assert Kw == K, f"K mismatch: x K={K}, w K={Kw}" + assert N == q_size + 2 * kv_size, f"N={N} != q_size({q_size}) + 2*kv_size({kv_size})" + + Q = torch.empty((M, q_size), dtype=x.dtype, device=x.device) + Kt = torch.empty((M, kv_size), dtype=x.dtype, device=x.device) + V = torch.empty((M, kv_size), dtype=x.dtype, device=x.device) + + HAS_BIAS = bias is not None + bias_arg = bias if HAS_BIAS else x # dummy ptr when unused + + grid = lambda meta: (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N'])) + _qkv_proj_split_kernel[grid]( + x, w_qkv, bias_arg, + Q, Kt, V, + M, K, N, q_size, kv_size, + x.stride(0), x.stride(1), + w_qkv.stride(0), w_qkv.stride(1), + Q.stride(0), Q.stride(1), + Kt.stride(0), Kt.stride(1), + V.stride(0), V.stride(1), + HAS_BIAS=HAS_BIAS, + ) + return Q, Kt, V + + +# ────────────────────────────────────────────────────────────────────────────── +# Pass 2 — RoPE for Q and K (standard rotate_half / contiguous-half convention) +# ────────────────────────────────────────────────────────────────────────────── +# cos/sin are [M, head_dim] (already gathered per-row by position_id upstream). +# rotate_half: for head split into [x1 | x2] (each head_dim/2 wide), +# out = [x1*cos1 - x2*sin1 , x2*cos2 + x1*sin2] +# with cos = [cos1 | cos2], sin = [sin1 | sin2] and cos1==cos2, sin1==sin2 in the +# standard HF layout. This matches torch: +# (t * cos) + (rotate_half(t) * sin), rotate_half(t) = cat(-t2, t1) + +@triton.jit +def _rope_inplace_kernel( + t_ptr, # [M, n_heads * head_dim] + cos_ptr, sin_ptr, # [M, head_dim] + M, n_heads, + head_dim: tl.constexpr, + half: tl.constexpr, + stride_tm, stride_td, + stride_cm, stride_cd, + BLOCK_H: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_h = tl.program_id(1) # head index + if pid_m >= M: + return + + h = pid_h + if h >= n_heads: + return + + half_idx = tl.arange(0, half) + base = pid_m * stride_tm + h * head_dim * stride_td + + # x1 = first half, x2 = second half + x1 = tl.load(t_ptr + base + half_idx * stride_td) + x2 = tl.load(t_ptr + base + (half + half_idx) * stride_td) + + cos1 = tl.load(cos_ptr + pid_m * stride_cm + half_idx * stride_cd) + sin1 = tl.load(sin_ptr + pid_m * stride_cm + half_idx * stride_cd) + cos2 = tl.load(cos_ptr + pid_m * stride_cm + (half + half_idx) * stride_cd) + sin2 = tl.load(sin_ptr + pid_m * stride_cm + (half + half_idx) * stride_cd) + + x1f = x1.to(tl.float32) + x2f = x2.to(tl.float32) + out1 = x1f * cos1.to(tl.float32) - x2f * sin1.to(tl.float32) + out2 = x2f * cos2.to(tl.float32) + x1f * sin2.to(tl.float32) + + tl.store(t_ptr + base + half_idx * stride_td, out1.to(t_ptr.dtype.element_ty)) + tl.store(t_ptr + base + (half + half_idx) * stride_td, out2.to(t_ptr.dtype.element_ty)) + + +def apply_rope_qk( + t: torch.Tensor, # [M, n_heads * head_dim] BF16, modified in place + cos: torch.Tensor, # [M, head_dim] + sin: torch.Tensor, # [M, head_dim] + head_dim: int, +): + """Apply standard rotate_half RoPE in place to a [M, n_heads*head_dim] tensor.""" + M, total = t.shape + assert total % head_dim == 0 + n_heads = total // head_dim + half = head_dim // 2 + grid = (M, n_heads) + _rope_inplace_kernel[grid]( + t, cos, sin, + M, n_heads, head_dim, half, + t.stride(0), t.stride(1), + cos.stride(0), cos.stride(1), + BLOCK_H=1, + ) + return t + + +def qkv_proj_rope( + x: torch.Tensor, # [M, K] BF16 + w_qkv: torch.Tensor, # [N, K] BF16 + bias: torch.Tensor | None, + q_size: int, + kv_size: int, + head_dim: int, + rope_cos: torch.Tensor | None, # [M, head_dim] + rope_sin: torch.Tensor | None, # [M, head_dim] +): + """QKV projection + split, with optional standard rotate_half RoPE on Q and K. + + Returns (Q, K, V), each BF16. RoPE is applied to Q and K only when + rope_cos/rope_sin are provided. V is never rotated. + """ + Q, Kt, V = qkv_proj_split(x, w_qkv, bias, q_size, kv_size) + if rope_cos is not None and rope_sin is not None: + apply_rope_qk(Q, rope_cos, rope_sin, head_dim) + apply_rope_qk(Kt, rope_cos, rope_sin, head_dim) + return Q, Kt, V From a96321de6bda9e822a84c0692fedffaa0b43fb94 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:21:26 +0000 Subject: [PATCH 11/16] feat(sm100): MXFP4 expert + grouped MoE dispatch to Triton On Blackwell the Hopper WGMMA/TMA MXFP4 .cu kernels are not built. Route both MXFP4 MoE paths to the existing pure-Triton fused_mxfp4_gemm kernels: - fused_wgmma_expert.py: is_wgmma_available() returns True on sm100 without loading _C; fused_mxfp4_expert_forward() dispatches to fused_mxfp4_mlp_forward - fused_wgmma_grouped.py: is_grouped_wgmma_available() True on sm100; new fused_mxfp4_grouped_moe_forward_triton() does a per-expert masked loop with fp32 accumulation and slot-specific routing-weight reduction - model.py _grouped_forward: sm100 branch uses the stacked per-expert weight tensors (device-pointer arrays can't be reversed) via the Triton helper Verified on B200: grouped output bit-exact vs per-token reference; expert path matches direct Triton call; no _C_*_mxfp4_wgmma import on sm100. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen/models/openai/gpt_oss_120b/model.py | 15 ++++ batchgen/moe/fused_wgmma_expert.py | 37 +++++++++ batchgen/moe/fused_wgmma_grouped.py | 82 ++++++++++++++++++++ 3 files changed, 134 insertions(+) diff --git a/batchgen/models/openai/gpt_oss_120b/model.py b/batchgen/models/openai/gpt_oss_120b/model.py index 10e4f9da..a9cc3b32 100644 --- a/batchgen/models/openai/gpt_oss_120b/model.py +++ b/batchgen/models/openai/gpt_oss_120b/model.py @@ -941,6 +941,21 @@ def _grouped_forward( ) -> torch.Tensor: """Grouped kernel for persistent experts.""" if self.weight_format == "mxfp4": + # SM100 (Blackwell): no Hopper WGMMA/TMA grouped .cu — use the + # pure-Triton per-expert path over the stacked weight tensors. + import batchgen_kernels as _bk + if _bk.get_device_arch() == "sm100": + from batchgen.moe.fused_wgmma_grouped import ( + fused_mxfp4_grouped_moe_forward_triton, + ) + return fused_mxfp4_grouped_moe_forward_triton( + hidden_flat, topk_indices, topk_weights, + self.persistent_expert_indices, + self.gate_weights, self.gate_scales, + self.up_weights, self.up_scales, + self.down_weights, self.down_scales, + self.gate_biases, self.up_biases, self.down_biases, + ) if not (_HAS_WGMMA_GROUPED and _HAS_CUDA_ROUTING): raise RuntimeError( "Grouped WGMMA MXFP4 kernel not available. " diff --git a/batchgen/moe/fused_wgmma_expert.py b/batchgen/moe/fused_wgmma_expert.py index 0dbaa48f..28e1b392 100644 --- a/batchgen/moe/fused_wgmma_expert.py +++ b/batchgen/moe/fused_wgmma_expert.py @@ -39,6 +39,16 @@ _wgmma_available = None _module_bf16_moe = None _module_mxfp4_moe = None +_arch = None + + +def _get_arch() -> str: + """Cached device arch ('sm90a' / 'sm100').""" + global _arch + if _arch is None: + import batchgen_kernels + _arch = batchgen_kernels.get_device_arch() + return _arch @@ -96,6 +106,13 @@ def is_wgmma_available() -> bool: _wgmma_available = False return False + # SM100 (Blackwell): the WGMMA .cu is not built. The fused MXFP4 expert + # path is provided by the pure-Triton kernel (fused_mxfp4_mlp_forward), + # which JIT-compiles for sm100. Do NOT attempt to load the _C extension. + if _get_arch() == "sm100": + _wgmma_available = True + return True + # Try to load the module mod = _load_mxfp4_module() _wgmma_available = mod is not None @@ -144,6 +161,26 @@ def fused_mxfp4_expert_forward( Raises: RuntimeError: If WGMMA kernels are not available """ + # SM100 (Blackwell): use the pure-Triton fused MXFP4 MLP (no WGMMA .cu). + if _get_arch() == "sm100": + from batchgen.triton_kernels.fused_mxfp4_gemm import fused_mxfp4_mlp_forward + x = hidden_states + original_shape = None + if x.dim() == 3: + original_shape = x.shape + x = x.view(-1, x.shape[-1]) + if x.dtype != torch.bfloat16: + x = x.to(torch.bfloat16) + out = fused_mxfp4_mlp_forward( + x.contiguous(), + gate_packed, gate_scales, gate_bias, + up_packed, up_scales, up_bias, + down_packed, down_scales, down_bias, + ) + if original_shape is not None: + out = out.view(original_shape[0], original_shape[1], -1) + return out + mod = _load_mxfp4_module() if mod is None: raise RuntimeError( diff --git a/batchgen/moe/fused_wgmma_grouped.py b/batchgen/moe/fused_wgmma_grouped.py index ece624ba..17fa7d63 100644 --- a/batchgen/moe/fused_wgmma_grouped.py +++ b/batchgen/moe/fused_wgmma_grouped.py @@ -55,6 +55,18 @@ def _check_wgmma_support() -> bool: return True +_arch = None + + +def _get_arch() -> str: + """Cached device arch ('sm90a' / 'sm100').""" + global _arch + if _arch is None: + import batchgen_kernels + _arch = batchgen_kernels.get_device_arch() + return _arch + + def _load_grouped_module(): """Load the grouped WGMMA CUDA module (pre-compiled via pip install).""" global _grouped_module @@ -88,11 +100,81 @@ def is_grouped_wgmma_available() -> bool: _grouped_wgmma_available = False return False + # SM100 (Blackwell): grouped MXFP4 is served by the model-level Triton path + # (fused_mxfp4_grouped_moe_forward_triton). Report available without loading + # the Hopper-only _C extension. + if _get_arch() == "sm100": + _grouped_wgmma_available = True + return True + mod = _load_grouped_module() _grouped_wgmma_available = mod is not None return _grouped_wgmma_available +def fused_mxfp4_grouped_moe_forward_triton( + hidden_states: torch.Tensor, # [num_tokens, hidden] BF16 + topk_indices: torch.Tensor, # [num_tokens, topk] int + topk_weights: torch.Tensor, # [num_tokens, topk] float + expert_indices, # iterable of global expert idx to process + gate_weights, gate_scales, # List[Tensor] indexed by global expert idx + up_weights, up_scales, + down_weights, down_scales, + gate_biases=None, up_biases=None, down_biases=None, +) -> torch.Tensor: + """SM100 grouped MXFP4 MoE forward using the pure-Triton expert MLP. + + Correctness-first port of the Hopper grouped WGMMA path for Blackwell. The + CUDA path sorts tokens by expert and applies a weighted scatter-add reduce; + here we process each expert over its routed tokens via a boolean mask and + accumulate the slot-specific routing-weighted contribution. Mathematically + equivalent to the grouped CUDA reduce (per-token contributions summed across + the experts it routes to). + + Returns: + Output [num_tokens, hidden] BF16 (routing-weighted sum of expert outputs). + """ + from batchgen.triton_kernels.fused_mxfp4_gemm import fused_mxfp4_mlp_forward + + num_tokens, hidden_size = hidden_states.shape + # Accumulate in fp32 to match the CUDA grouped reduce precision, then cast. + output = torch.zeros( + num_tokens, hidden_size, + dtype=torch.float32, device=hidden_states.device, + ) + + # Single CPU-GPU sync: which experts have any routed token. + active_experts = set(topk_indices.flatten().tolist()) + + for e in expert_indices: + if e not in active_experts: + continue + + mask = (topk_indices == e).any(dim=-1) + x_e = hidden_states[mask].contiguous() + + out_e = fused_mxfp4_mlp_forward( + x_e, + gate_weights[e], gate_scales[e], + gate_biases[e] if gate_biases is not None else None, + up_weights[e], up_scales[e], + up_biases[e] if up_biases is not None else None, + down_weights[e], down_scales[e], + down_biases[e] if down_biases is not None else None, + ) + + # Slot-specific routing weight for expert e on each selected token. + sel_idx = topk_indices[mask] + sel_w = topk_weights[mask] + w_e = torch.where( + sel_idx == e, sel_w, torch.zeros_like(sel_w) + ).sum(dim=-1).float() + + output[mask] += out_e.float() * w_e.unsqueeze(-1) + + return output.to(hidden_states.dtype) + + # ────────────────────────────────────────────────────────────────────────────── # Low-Level Python Wrappers # ────────────────────────────────────────────────────────────────────────────── From 2e15ed5f12528df33bf2edb13ff19abde8196a05 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:24:31 +0000 Subject: [PATCH 12/16] feat(sm100): router GEMM Triton kernel + fused-gate dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On Blackwell the WGMMA fused-gate .cu is not built. Provide the router logits GEMM as a pure-Triton kernel and route FusedGateContext through it: - batchgen_kernels/triton/fused_router_gemm.py: router_gemm_bias(hidden, weight, bias) -> logits [N,E] FP32 (FP32 tensor-core accumulation) - triton/__init__.py: export router_gemm_bias - cuda_routing.py: FusedGateContext is sm100-aware — stores weight/bias in Python, warmup() no-op, forward() runs router_gemm_bias + the generic-CUDA gate_topk_softmax (compiled for sm100 in sub-task 2) Verified on B200: logits FP32, max_abs_err ~5e-5 vs float32 reference; FusedGateContext top-k indices match reference. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen/moe/routing/cuda_routing.py | 37 +++++- batchgen_kernels/triton/__init__.py | 1 + batchgen_kernels/triton/fused_router_gemm.py | 117 +++++++++++++++++++ 3 files changed, 153 insertions(+), 2 deletions(-) create mode 100644 batchgen_kernels/triton/fused_router_gemm.py diff --git a/batchgen/moe/routing/cuda_routing.py b/batchgen/moe/routing/cuda_routing.py index 062cbb81..0894d010 100644 --- a/batchgen/moe/routing/cuda_routing.py +++ b/batchgen/moe/routing/cuda_routing.py @@ -271,6 +271,22 @@ def __init__(self, router_weight, router_bias, topk=4): router_bias: [E] BF16 (or empty tensor if no bias) topk: number of top experts (2, 4, or 8) """ + import batchgen_kernels as _bk + self._is_sm100 = _bk.get_device_arch() == "sm100" + self._topk = topk + + if self._is_sm100: + # SM100 (Blackwell): the WGMMA fused-gate .cu is not built. Keep the + # weight/bias in Python and run a Triton GEMM + generic-CUDA TopK. + self._weight = router_weight + if router_bias is not None and router_bias.numel() > 0: + self._bias = router_bias + else: + self._bias = None + self._ctx = None + self._ext = None + return + ext = _get_ext() bias = router_bias if router_bias is not None else torch.empty( 0, dtype=torch.bfloat16, device=router_weight.device) @@ -288,6 +304,8 @@ def warmup(self, base_buffer): Args: base_buffer: [WB_max, H] BF16 — the full allocated buffer """ + if self._is_sm100: + return # no TMA descriptor on the Triton path self._ext.fused_gate_warmup(self._ctx, base_buffer) def forward(self, hidden_states, logits=None, topk_indices=None, @@ -305,6 +323,20 @@ def forward(self, hidden_states, logits=None, topk_indices=None, topk_indices: [N, K] int32 topk_weights: [N, K] FP32 """ + if self._is_sm100: + from batchgen_kernels.triton.fused_router_gemm import router_gemm_bias + computed_logits = router_gemm_bias(hidden_states, self._weight, self._bias) + if logits is not None and logits.numel() > 0: + logits.copy_(computed_logits) + computed_logits = logits + return gate_topk_softmax_cuda( + computed_logits, + topk_indices=topk_indices, + topk_weights=topk_weights, + k=self._topk, + num_valid_tokens=num_valid_tokens, + ) + _empty = torch.empty(0) result = self._ext.fused_gate_forward( self._ctx, @@ -317,8 +349,9 @@ def forward(self, hidden_states, logits=None, topk_indices=None, return result[0], result[1] def __del__(self): - if hasattr(self, '_ctx') and hasattr(self, '_ext'): - self._ext.destroy_fused_gate_context(self._ctx) + if not getattr(self, '_is_sm100', False) and hasattr(self, '_ctx') and hasattr(self, '_ext'): + if self._ctx is not None and self._ext is not None: + self._ext.destroy_fused_gate_context(self._ctx) def reduce_weighted_scatter_cuda( diff --git a/batchgen_kernels/triton/__init__.py b/batchgen_kernels/triton/__init__.py index 27874ab1..ab8d3062 100644 --- a/batchgen_kernels/triton/__init__.py +++ b/batchgen_kernels/triton/__init__.py @@ -33,3 +33,4 @@ from batchgen_kernels.triton.fused_dequant_gemm import fused_fp8_bf16_gemm from batchgen_kernels.triton.fused_q_absorb import fused_q_absorb_query_states from batchgen_kernels.triton.fused_out_absorb import fused_out_absorb_reshape +from batchgen_kernels.triton.fused_router_gemm import router_gemm_bias diff --git a/batchgen_kernels/triton/fused_router_gemm.py b/batchgen_kernels/triton/fused_router_gemm.py new file mode 100644 index 00000000..f742dff4 --- /dev/null +++ b/batchgen_kernels/triton/fused_router_gemm.py @@ -0,0 +1,117 @@ +"""Router GEMM + bias Triton kernel for SM100 (Blackwell). + +Pure-Triton replacement for the WGMMA half of the fused-gate kernel +(`fused_gate_forward`) whose Hopper-only `.cu` is not built on sm_100a. +Computes the router logits + + logits[N, E] = hidden[N, K_dim] @ weight[E, K_dim].T (+ bias[E]) + +in FP32 (FP32 tensor-core accumulation) to match the SM90a kernel's output +dtype and precision. The subsequent TopK+Softmax stage is a generic CUDA +kernel (`gate_topk_softmax`) already compiled for sm100. +""" + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=2, num_warps=8), + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_stages=3, num_warps=8), + ], + key=['N', 'K_dim', 'E'], +) +@triton.jit +def _router_gemm_bias_kernel( + x_ptr, w_ptr, bias_ptr, out_ptr, + N, K_dim, E, + stride_xn, stride_xk, + stride_wk, stride_we, + stride_on, stride_oe, + HAS_BIAS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + x_ptrs = x_ptr + (offs_m[:, None] * stride_xn + offs_k[None, :] * stride_xk) + w_ptrs = w_ptr + (offs_k[:, None] * stride_wk + offs_n[None, :] * stride_we) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k_start in range(0, K_dim, BLOCK_K): + k_mask = offs_k[None, :] < (K_dim - k_start) + x_tile = tl.load(x_ptrs, mask=(offs_m[:, None] < N) & k_mask, other=0.0) + w_tile = tl.load( + w_ptrs, + mask=(offs_k[:, None] < (K_dim - k_start)) & (offs_n[None, :] < E), + other=0.0, + ) + acc += tl.dot(x_tile, w_tile, out_dtype=tl.float32) + x_ptrs += BLOCK_K * stride_xk + w_ptrs += BLOCK_K * stride_wk + + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=offs_n < E, other=0.0) + acc += bias[None, :].to(tl.float32) + + out_ptrs = out_ptr + (offs_m[:, None] * stride_on + offs_n[None, :] * stride_oe) + tl.store(out_ptrs, acc, mask=(offs_m[:, None] < N) & (offs_n[None, :] < E)) + + +def router_gemm_bias( + hidden: torch.Tensor, # [N, K_dim] BF16 + weight: torch.Tensor, # [E, K_dim] BF16 (nn.Linear weight, row-major) + bias: torch.Tensor = None, # [E] BF16/FP32 or None +) -> torch.Tensor: + """Router GEMM + optional bias. Returns logits [N, E] FP32. + + Args: + hidden: Input activations [N, K_dim] BF16. + weight: Router weight [E, K_dim] BF16 (nn.Linear layout). + bias: Optional router bias [E]. + + Returns: + logits [N, E] FP32 (FP32 accumulation, matches SM90a output dtype). + """ + assert hidden.dim() == 2, f"hidden must be 2D, got {hidden.shape}" + N, K_dim = hidden.shape + E = weight.shape[0] + assert weight.shape[1] == K_dim, ( + f"weight K mismatch: hidden K={K_dim}, weight K={weight.shape[1]}" + ) + + hidden = hidden.contiguous() + # [E, K_dim] -> [K_dim, E] so the inner GEMM dim is contiguous-friendly. + w_t = weight.t().contiguous() + + has_bias = bias is not None + bias_arg = bias if has_bias else hidden # dummy ptr when no bias + if has_bias and bias.dtype != torch.float32 and bias.dtype != torch.bfloat16: + bias_arg = bias.to(torch.float32) + + logits = torch.empty((N, E), dtype=torch.float32, device=hidden.device) + + grid = lambda meta: ( + triton.cdiv(N, meta['BLOCK_M']), + triton.cdiv(E, meta['BLOCK_N']), + ) + _router_gemm_bias_kernel[grid]( + hidden, w_t, bias_arg, logits, + N, K_dim, E, + hidden.stride(0), hidden.stride(1), + w_t.stride(0), w_t.stride(1), + logits.stride(0), logits.stride(1), + HAS_BIAS=has_bias, + ) + return logits From 3d6941a737319c002bfd7b806ea2f70ed9ff7d36 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:28:51 +0000 Subject: [PATCH 13/16] feat(sm100): INT4 (W4A16) grouped GEMM Triton kernel (K2.5 decode) Pure-Triton port of the Hopper WGMMA INT4 grouped MoE GEMM building block. Matches the SM90a .cu weight layout: w_packed [N, K//2] uint8 (low nibble = even-K, high nibble = odd-K), scale [N, K//32] bf16, group_size=32, dequant (nibble-8)*scale. GEMM via even/odd column split + two bf16 tl.dot with FP32 accumulation (avoids in-kernel nibble interleave). - batchgen_kernels/triton/int4_grouped_gemm.py: int4_grouped_gemm (single expert) + int4_moe_grouped_gemm (per-expert sorted-token dispatch) - triton/__init__.py: export both Verified on B200: max_abs_err 0.07-0.18 (rel ~0.3%) vs dequant reference, well under the 0.5 INT4-noise threshold; grouped path matches. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen_kernels/triton/__init__.py | 4 + batchgen_kernels/triton/int4_grouped_gemm.py | 163 +++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 batchgen_kernels/triton/int4_grouped_gemm.py diff --git a/batchgen_kernels/triton/__init__.py b/batchgen_kernels/triton/__init__.py index 27874ab1..0a3838be 100644 --- a/batchgen_kernels/triton/__init__.py +++ b/batchgen_kernels/triton/__init__.py @@ -33,3 +33,7 @@ from batchgen_kernels.triton.fused_dequant_gemm import fused_fp8_bf16_gemm from batchgen_kernels.triton.fused_q_absorb import fused_q_absorb_query_states from batchgen_kernels.triton.fused_out_absorb import fused_out_absorb_reshape +from batchgen_kernels.triton.int4_grouped_gemm import ( + int4_grouped_gemm, + int4_moe_grouped_gemm, +) diff --git a/batchgen_kernels/triton/int4_grouped_gemm.py b/batchgen_kernels/triton/int4_grouped_gemm.py new file mode 100644 index 00000000..adf92538 --- /dev/null +++ b/batchgen_kernels/triton/int4_grouped_gemm.py @@ -0,0 +1,163 @@ +"""INT4 (W4A16) GEMM Triton kernels for SM100 (Blackwell). + +Pure-Triton port of the Hopper WGMMA INT4 grouped MoE kernels (K2.5 decode) +whose `.cu` uses `wgmma`/TMA intrinsics not built for sm_100a. + +Weight / scale layout (matches the SM90a `.cu`, e.g. +`single_expert_int4_wgmma.cu::load_decode_rhs_int4_swizzled`): + * `w_packed` : [N, K // 2] uint8 — output-channel major, two INT4 values per + byte along K. The **low** nibble `(byte & 0x0F)` is the even K index, the + **high** nibble `((byte >> 4) & 0x0F)` is the odd K index. + * `scale` : [N, K // group_size] bf16 — one scale per group of + `group_size` (=32) contiguous K values, shared across the group. + * dequant : `w[n, k] = (nibble - 8) * scale[n, k // group_size]` + * GEMM : `out[m, n] = sum_k x[m, k] * w[n, k]` (i.e. `x @ dequant(W).T`). +""" + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 64, 'BLOCK_KP': 64}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128, 'BLOCK_KP': 64}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 128, 'BLOCK_KP': 64}, num_stages=4, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_KP': 32}, num_stages=3, num_warps=8), + ], + key=['M', 'K', 'N', 'GROUP_SIZE'], +) +@triton.jit +def _int4_gemm_kernel( + x_ptr, w_ptr, scale_ptr, out_ptr, + M, K, N, + stride_xm, stride_xk, + stride_wn, stride_wk, # w_packed: [N, K//2] uint8 + stride_sn, stride_sk, # scale: [N, K//group_size] bf16 + stride_om, stride_on, + GROUP_SIZE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_KP: tl.constexpr, # packed-K (bytes) per iteration; real-K = 2*BLOCK_KP +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_kp = tl.arange(0, BLOCK_KP) + + Kp = K // 2 + # packed-K bytes per scale group (group_size K-values -> group_size//2 bytes) + GROUP_BYTES = GROUP_SIZE // 2 + + m_mask = offs_m[:, None] < M + n_mask_out = offs_n[None, :] < N # [1, BLOCK_N] for output/x-side + n_mask_w = offs_n[:, None] < N # [BLOCK_N, 1] for weight/scale rows + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for kp_start in range(0, Kp, BLOCK_KP): + kp = kp_start + offs_kp # [BLOCK_KP] packed indices + kp_mask = kp < Kp + + # x even / odd columns: x[:, 2*kp] and x[:, 2*kp+1] + x_even = tl.load( + x_ptr + offs_m[:, None] * stride_xm + (2 * kp[None, :]) * stride_xk, + mask=m_mask & kp_mask[None, :], other=0.0, + ) + x_odd = tl.load( + x_ptr + offs_m[:, None] * stride_xm + (2 * kp[None, :] + 1) * stride_xk, + mask=m_mask & kp_mask[None, :], other=0.0, + ) + + # packed weight byte [BLOCK_N, BLOCK_KP] + wb = tl.load( + w_ptr + offs_n[:, None] * stride_wn + kp[None, :] * stride_wk, + mask=n_mask_w & kp_mask[None, :], other=0, + ) + lo = ((wb & 0x0F).to(tl.int32) - 8).to(tl.bfloat16) + hi = (((wb >> 4) & 0x0F).to(tl.int32) - 8).to(tl.bfloat16) + + # scale [BLOCK_N, BLOCK_KP]: group index = (2*kp) // group_size = kp // GROUP_BYTES + sg = kp // GROUP_BYTES + sc = tl.load( + scale_ptr + offs_n[:, None] * stride_sn + sg[None, :] * stride_sk, + mask=n_mask_w & kp_mask[None, :], other=0.0, + ) + w_lo = (lo * sc) # [BLOCK_N, BLOCK_KP] bf16 + w_hi = (hi * sc) + + acc += tl.dot(x_even, tl.trans(w_lo), out_dtype=tl.float32) + acc += tl.dot(x_odd, tl.trans(w_hi), out_dtype=tl.float32) + + out_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on + tl.store(out_ptrs, acc.to(tl.bfloat16), mask=m_mask & n_mask_out) + + +def int4_grouped_gemm( + x: torch.Tensor, # [M, K] BF16 + w_packed: torch.Tensor, # [N, K // 2] uint8 (INT4 packed, output-major) + scales: torch.Tensor, # [N, K // group_size] BF16 + group_size: int = 32, +) -> torch.Tensor: + """Single-expert INT4×BF16 GEMM with per-group dequant. Returns [M, N] BF16. + + Computes `out = x @ dequant(w_packed, scales).T` where + `dequant[n, k] = ((nibble(n,k) - 8) * scales[n, k // group_size])`. + """ + assert x.dim() == 2, f"x must be 2D, got {x.shape}" + M, K = x.shape + if w_packed.dtype == torch.int32: + w_packed = w_packed.view(torch.uint8) + assert w_packed.dtype == torch.uint8, f"w_packed must be uint8, got {w_packed.dtype}" + N, Kp = w_packed.shape + assert Kp == K // 2, f"w_packed K mismatch: x K={K} -> expected {K // 2}, got {Kp}" + assert K % group_size == 0, f"K={K} not divisible by group_size={group_size}" + + x = x.contiguous() + w_packed = w_packed.contiguous() + scales = scales.contiguous() + + out = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) + grid = lambda meta: ( + triton.cdiv(M, meta['BLOCK_M']), + triton.cdiv(N, meta['BLOCK_N']), + ) + _int4_gemm_kernel[grid]( + x, w_packed, scales, out, + M, K, N, + x.stride(0), x.stride(1), + w_packed.stride(0), w_packed.stride(1), + scales.stride(0), scales.stride(1), + out.stride(0), out.stride(1), + GROUP_SIZE=group_size, + ) + return out + + +def int4_moe_grouped_gemm( + x: torch.Tensor, # [total_M, K] BF16 (sorted by expert) + w_packed_list, # list[Tensor [N, K//2] uint8] per expert + scales_list, # list[Tensor [N, K//group_size] bf16] per expert + m_offsets, # [num_experts + 1] int (row offsets into x) + group_size: int = 32, +) -> torch.Tensor: + """Batched INT4 MoE GEMM across experts (decode: small M per expert). + + Tokens are assumed pre-sorted by expert; `m_offsets[e]:m_offsets[e+1]` is the + contiguous row range for expert `e`. Returns [total_M, N] BF16. + """ + total_M, K = x.shape + N = w_packed_list[0].shape[0] + out = torch.empty((total_M, N), dtype=torch.bfloat16, device=x.device) + num_experts = len(w_packed_list) + offs = m_offsets.tolist() if torch.is_tensor(m_offsets) else list(m_offsets) + for e in range(num_experts): + lo, hi = offs[e], offs[e + 1] + if hi <= lo: + continue + out[lo:hi] = int4_grouped_gemm( + x[lo:hi], w_packed_list[e], scales_list[e], group_size=group_size, + ) + return out From ee7765253d85f337feeae8425231c2ed1d211c14 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:33:24 +0000 Subject: [PATCH 14/16] feat(sm100): fused INT4 grouped + SiLU MoE Triton path (K2.5) Completes the SM100 INT4 MoE expert path on top of sub-task 7's GEMM: - batchgen_kernels/triton/fused_int4_grouped_silu.py: * silu_mul epilogue (silu(gate)*up, FP32 math) matching the SM90a stage-1 * fused_int4_grouped_silu (stage-1 SwiGLU activation) * int4_expert_mlp (full gate+up+SiLU+down expert) * int4_grouped_moe_forward (per-expert masked loop, fp32 accumulation, slot-specific routing-weight reduction) - int4_single_expert_wgmma.py: single_expert_int4_forward dispatches to the Triton int4_expert_mlp on sm100 (no _C load) - fused_int4_wgmma_grouped.py: is_int4_grouped_wgmma_available() True on sm100 via Triton without loading the Hopper _C extension - triton/__init__.py: export the new kernels Verified on B200: silu_mul/expert_mlp rel err ~0.4-0.6%; single-expert sm100 dispatch matches int4_expert_mlp; grouped MoE matches per-token reference. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen/moe/fused_int4_wgmma_grouped.py | 19 +++ batchgen/moe/int4_single_expert_wgmma.py | 20 +++ batchgen_kernels/triton/__init__.py | 6 + .../triton/fused_int4_grouped_silu.py | 119 ++++++++++++++++++ 4 files changed, 164 insertions(+) create mode 100644 batchgen_kernels/triton/fused_int4_grouped_silu.py diff --git a/batchgen/moe/fused_int4_wgmma_grouped.py b/batchgen/moe/fused_int4_wgmma_grouped.py index 51a015c9..aaced5df 100644 --- a/batchgen/moe/fused_int4_wgmma_grouped.py +++ b/batchgen/moe/fused_int4_wgmma_grouped.py @@ -62,6 +62,18 @@ def _check_wgmma_support() -> bool: return True +_arch = None + + +def _get_arch() -> str: + """Cached device arch ('sm90a' / 'sm100').""" + global _arch + if _arch is None: + import batchgen_kernels + _arch = batchgen_kernels.get_device_arch() + return _arch + + def _load_int4_grouped_module(): """Load the pre-compiled grouped INT4 WGMMA CUDA module (Stage 1 + Stage 2).""" global _int4_grouped_module @@ -97,6 +109,13 @@ def is_int4_grouped_wgmma_available() -> bool: _int4_grouped_wgmma_available = False return False + # SM100 (Blackwell): grouped INT4 is served by the model-level Triton path + # (int4_grouped_moe_forward). Report available without loading the + # Hopper-only _C extension. + if _get_arch() == "sm100": + _int4_grouped_wgmma_available = True + return True + mod = _load_int4_grouped_module() _int4_grouped_wgmma_available = mod is not None return _int4_grouped_wgmma_available diff --git a/batchgen/moe/int4_single_expert_wgmma.py b/batchgen/moe/int4_single_expert_wgmma.py index 983f384b..6bdbec9e 100644 --- a/batchgen/moe/int4_single_expert_wgmma.py +++ b/batchgen/moe/int4_single_expert_wgmma.py @@ -19,6 +19,16 @@ import torch _single_expert_module = None +_arch = None + + +def _get_arch() -> str: + """Cached device arch ('sm90a' / 'sm100').""" + global _arch + if _arch is None: + import batchgen_kernels + _arch = batchgen_kernels.get_device_arch() + return _arch def _get_single_expert_module(): @@ -56,6 +66,16 @@ def single_expert_int4_forward( Returns: output: [M, K=7168] bf16 """ + # SM100 (Blackwell): the WGMMA INT4 .cu is not built — use the Triton MLP. + if _get_arch() == "sm100": + from batchgen_kernels.triton.fused_int4_grouped_silu import int4_expert_mlp + gp = gate_packed.view(torch.uint8) if gate_packed.dtype == torch.int32 else gate_packed + upp = up_packed.view(torch.uint8) if up_packed.dtype == torch.int32 else up_packed + dp = down_packed.view(torch.uint8) if down_packed.dtype == torch.int32 else down_packed + return int4_expert_mlp( + hidden, gp, gate_scale, upp, up_scale, dp, down_scale, group_size=32, + ) + mod = _get_single_expert_module() empty_bias = torch.empty(0, dtype=torch.bfloat16, device=hidden.device) diff --git a/batchgen_kernels/triton/__init__.py b/batchgen_kernels/triton/__init__.py index 0a3838be..251242ca 100644 --- a/batchgen_kernels/triton/__init__.py +++ b/batchgen_kernels/triton/__init__.py @@ -37,3 +37,9 @@ int4_grouped_gemm, int4_moe_grouped_gemm, ) +from batchgen_kernels.triton.fused_int4_grouped_silu import ( + fused_int4_grouped_silu, + silu_mul, + int4_expert_mlp, + int4_grouped_moe_forward, +) diff --git a/batchgen_kernels/triton/fused_int4_grouped_silu.py b/batchgen_kernels/triton/fused_int4_grouped_silu.py new file mode 100644 index 00000000..97c4c7dd --- /dev/null +++ b/batchgen_kernels/triton/fused_int4_grouped_silu.py @@ -0,0 +1,119 @@ +"""Fused INT4 grouped + SiLU (SwiGLU) Triton kernels for SM100 (Blackwell). + +Pure-Triton port of the Hopper WGMMA INT4 MoE expert path (K2.5 decode). The +SM90a stage-1 kernel computes `silu(gate) * up` and stage-2 the down +projection. On sm_100a those `.cu` kernels are not built, so we compose the +sub-task-7 INT4 GEMM building block with a small `silu_mul` epilogue and a +second INT4 GEMM for the down projection. + +SwiGLU convention (matches `single_expert_int4_wgmma.cu` epilogue): + out = (gate * sigmoid(gate)) * up # silu(gate) * up, computed in FP32 +""" + +import torch +import triton +import triton.language as tl + +from batchgen_kernels.triton.int4_grouped_gemm import int4_grouped_gemm + + +@triton.jit +def _silu_mul_kernel(gate_ptr, up_ptr, out_ptr, n_elems, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elems + gate = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) + up = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) + out = (gate * (1.0 / (1.0 + tl.exp(-gate)))) * up # silu(gate) * up + tl.store(out_ptr + offs, out.to(tl.bfloat16), mask=mask) + + +def silu_mul(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """Elementwise SwiGLU activation: silu(gate) * up (FP32 math, BF16 out).""" + assert gate.shape == up.shape, f"shape mismatch {gate.shape} vs {up.shape}" + out = torch.empty_like(gate, dtype=torch.bfloat16) + n = gate.numel() + if n == 0: + return out + BLOCK = 1024 + grid = (triton.cdiv(n, BLOCK),) + _silu_mul_kernel[grid](gate.contiguous(), up.contiguous(), out, n, BLOCK=BLOCK) + return out + + +def fused_int4_grouped_silu( + x: torch.Tensor, + wg_packed: torch.Tensor, wg_scales: torch.Tensor, + wu_packed: torch.Tensor, wu_scales: torch.Tensor, + group_size: int = 32, +) -> torch.Tensor: + """silu(x @ dequant(Wg).T) * (x @ dequant(Wu).T) with INT4 weights. + + Returns the stage-1 SwiGLU activation [M, N_intermediate] BF16. + """ + gate = int4_grouped_gemm(x, wg_packed, wg_scales, group_size) + up = int4_grouped_gemm(x, wu_packed, wu_scales, group_size) + return silu_mul(gate, up) + + +def int4_expert_mlp( + x: torch.Tensor, # [M, K] BF16 + gate_packed: torch.Tensor, gate_scale: torch.Tensor, + up_packed: torch.Tensor, up_scale: torch.Tensor, + down_packed: torch.Tensor, down_scale: torch.Tensor, + group_size: int = 32, +) -> torch.Tensor: + """Full INT4 expert MLP: stage1 (gate+up+SiLU) + stage2 (down). + + Pure-Triton equivalent of `single_expert_int4_forward`. Returns [M, K] BF16. + """ + intermediate = fused_int4_grouped_silu( + x, gate_packed, gate_scale, up_packed, up_scale, group_size, + ) + return int4_grouped_gemm(intermediate, down_packed, down_scale, group_size) + + +def int4_grouped_moe_forward( + hidden_states: torch.Tensor, # [num_tokens, K] BF16 + topk_indices: torch.Tensor, # [num_tokens, topk] int + topk_weights: torch.Tensor, # [num_tokens, topk] float + expert_indices, # iterable of global expert idx to process + gate_packed, gate_scale, # List[Tensor] indexed by global expert idx + up_packed, up_scale, + down_packed, down_scale, + group_size: int = 32, +) -> torch.Tensor: + """SM100 INT4 grouped MoE forward (correctness-first per-expert masked loop). + + Mirrors the MXFP4 sm100 path: each routed expert runs the full INT4 MLP over + its masked tokens, accumulated in FP32 with slot-specific routing weights. + + Returns: + Output [num_tokens, K] BF16 (routing-weighted sum of expert outputs). + """ + num_tokens, K = hidden_states.shape + output = torch.zeros(num_tokens, K, dtype=torch.float32, device=hidden_states.device) + + active_experts = set(topk_indices.flatten().tolist()) + for e in expert_indices: + if e not in active_experts: + continue + mask = (topk_indices == e).any(dim=-1) + x_e = hidden_states[mask].contiguous() + + out_e = int4_expert_mlp( + x_e, + gate_packed[e], gate_scale[e], + up_packed[e], up_scale[e], + down_packed[e], down_scale[e], + group_size, + ) + + sel_idx = topk_indices[mask] + sel_w = topk_weights[mask] + w_e = torch.where( + sel_idx == e, sel_w, torch.zeros_like(sel_w) + ).sum(dim=-1).float() + output[mask] += out_e.float() * w_e.unsqueeze(-1) + + return output.to(hidden_states.dtype) From 10204361425a66e5c88df091e075cf422f5933b9 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:44:40 +0000 Subject: [PATCH 15/16] feat(sm100): FP8 blockwise grouped GEMM fallback via row-wise _scaled_mm The compiled SM90a CuTe FP8 blockwise kernel is unavailable on Blackwell, and cuBLAS in torch 2.9+cu129 rejects 1x128/128x128 blockwise FP8 scaling (heuristic returns CUBLAS_STATUS_NOT_SUPPORTED; fast-accum unsupported for 1x128). Only row-wise/tensor-wise FP8 GEMM is available. Add an sm100 branch in grouped_fp8_blockwise_gemm() that emulates deepseek blockwise scaling exactly: split the K contraction into 128-wide blocks and issue one row-wise torch._scaled_mm per block, accumulating partials in fp32. Within a K-block the act scale is constant per row (1x128) and the weight scale is constant per 128-output-row block (128x128, expanded to per-column), so the row-wise GEMM is numerically identical to true blockwise for that block. - Processes the full uniform mtp reserved rows per expert -> static control flow, CUDA-graph compatible (no host sync on seqlens, no data-dependent shapes). Padding-row outputs are ignored by downstream gather. - Covers s1 (gate+up), s3 (down) and fused_s1 automatically since all route through grouped_fp8_blockwise_gemm (fused_s1 falls back to s1_silu on sm100 as its _C kernel is absent). - Dim asserts (K,N multiples of 128; scale shapes) guard silent mis-scaling. Verified on B200 (GLM-5 / MiniMax-M2.5 style dims, E=4, mtp=64, K=2048, N=1536, incl. a zero-token expert): GEMM max rel 0.0039, S1+SiLU 0.0056, fused_s1 matches s1_silu (< 1e-2 acceptance). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen/moe/grouped_fp8_blockwise_moe.py | 78 +++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/batchgen/moe/grouped_fp8_blockwise_moe.py b/batchgen/moe/grouped_fp8_blockwise_moe.py index c52c47b4..f21d6bdf 100644 --- a/batchgen/moe/grouped_fp8_blockwise_moe.py +++ b/batchgen/moe/grouped_fp8_blockwise_moe.py @@ -24,6 +24,80 @@ _warned_import = False _warned_fused_s1 = False +_arch = None + + +def _get_arch() -> str: + """Cached device arch ("sm100" / "sm90a" / ...) via batchgen_kernels.""" + global _arch + if _arch is None: + import batchgen_kernels as _bk + _arch = _bk.get_device_arch() + return _arch + + +def _grouped_fp8_blockwise_gemm_sm100( + x_fp8: Tensor, + weight_3d: Tensor, + x_scale: Tensor, + w_scale_3d: Tensor, + output: Optional[Tensor] = None, +) -> Tensor: + """SM100 (Blackwell) fallback for the FP8 blockwise grouped GEMM. + + The compiled SM90a CuTe kernel is unavailable on sm_100, and cuBLAS in + torch 2.9+cu129 does not yet support 1x128/128x128 blockwise FP8 scaling + (the heuristic returns CUBLAS_STATUS_NOT_SUPPORTED). Only row-wise FP8 + GEMM is supported, so we emulate deepseek-style blockwise scaling exactly + by splitting the contraction dim K into 128-wide blocks and issuing one + row-wise ``torch._scaled_mm`` per block, accumulating partials in fp32. + + Within a single K-block the activation scale is constant per token row + (1x128) and the weight scale is constant per 128-output-row block + (128x128, expanded here to per-output-column), so the row-wise GEMM is + numerically identical to true blockwise scaling for that block. + + Processes the full uniform ``mtp`` reserved rows for every expert so the + control flow is static (CUDA-graph compatible — no data-dependent shapes + or host syncs on ``seqlens``). Padding rows produce values in output rows + that downstream gather ignores. + """ + E, N, K = weight_3d.shape + g = 128 + assert K % g == 0, f"FP8 sm100 GEMM requires K (={K}) multiple of 128" + assert N % g == 0, f"FP8 sm100 GEMM requires N (={N}) multiple of 128" + EM = x_fp8.shape[0] + assert EM % E == 0, f"x_fp8 rows (={EM}) not divisible by E (={E})" + mtp = EM // E + nblk = K // g + assert x_scale.shape[0] >= nblk, ( + f"x_scale dim0 (={x_scale.shape[0]}) < K/128 (={nblk})") + assert w_scale_3d.shape[1] == N // g, ( + f"w_scale dim1 (={w_scale_3d.shape[1]}) != N/128 (={N // g})") + assert w_scale_3d.shape[2] >= nblk, ( + f"w_scale dim2 (={w_scale_3d.shape[2]}) < K/128 (={nblk})") + + if output is None: + output = torch.empty((EM, N), dtype=torch.bfloat16, device=x_fp8.device) + + for e in range(E): + start = e * mtp + x_e = x_fp8[start:start + mtp] # [mtp, K] fp8 + w_e = weight_3d[e] # [N, K] fp8 + xs_e = x_scale[:, start:start + mtp] # [>=nblk, mtp] f32 (transposed) + ws_e = w_scale_3d[e] # [N/128, >=nblk] f32 + acc = torch.zeros((mtp, N), dtype=torch.float32, device=x_fp8.device) + for j in range(nblk): + a_blk = x_e[:, j * g:(j + 1) * g] # [mtp, 128] row-major view + b_blk = w_e[:, j * g:(j + 1) * g].t() # [128, N] col-major view + sa = xs_e[j].contiguous().view(mtp, 1) # [mtp, 1] act scale + sb = ws_e[:, j].repeat_interleave(g)[:N].contiguous().view(1, N) + o = torch._scaled_mm( + a_blk, b_blk, scale_a=sa, scale_b=sb, out_dtype=torch.bfloat16) + acc += o.float() + output[start:start + mtp] = acc.to(torch.bfloat16) + return output + def _get_kernel(): """Load the compiled FP8 blockwise GEMM kernel.""" @@ -88,6 +162,10 @@ def grouped_fp8_blockwise_gemm( Returns: [E*mtp, N] bf16 output """ + if _get_arch() == "sm100": + return _grouped_fp8_blockwise_gemm_sm100( + x_fp8, weight_3d, x_scale, w_scale_3d, output) + kernel = _get_kernel() if kernel is None: raise RuntimeError( From 9b21172a37436c38dc68692d0db72b7443d3abe3 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 15:13:11 +0000 Subject: [PATCH 16/16] chore: union-merge WIP marker doc to ease integration merges --- .gitattributes | 1 + 1 file changed, 1 insertion(+) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 00000000..65803ddb --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +docs/BLACKWELL_KERNELS_WIP.md merge=ours