[blackwell][draft] sm100: MXFP4 dispatch -> Triton fused_mxfp4_gemm (#227)

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. "

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(

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
 # ──────────────────────────────────────────────────────────────────────────────

docs/BLACKWELL_KERNELS_WIP.md

@@ -0,0 +1 @@
+# WIP: [wip] blackwell: sm100 MXFP4 Triton dispatch stub