[blackwell] sm100 kernel port — integration (sub-tasks 1–8, #204)

.gitattributes

@@ -0,0 +1 @@
+docs/BLACKWELL_KERNELS_WIP.md merge=ours

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"
 

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_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

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

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(

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)