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 1/2] [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 a96321de6bda9e822a84c0692fedffaa0b43fb94 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:21:26 +0000 Subject: [PATCH 2/2] 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 # ──────────────────────────────────────────────────────────────────────────────