[blackwell][draft] sm100: Fused INT4 grouped + SiLU Triton kernel (#231)

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/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)
 

batchgen_kernels/triton/__init__.py

@@ -33,3 +33,13 @@
 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,
+)
+from batchgen_kernels.triton.fused_int4_grouped_silu import (
+    fused_int4_grouped_silu,
+    silu_mul,
+    int4_expert_mlp,
+    int4_grouped_moe_forward,
+)

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)

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

docs/BLACKWELL_KERNELS_WIP.md

@@ -0,0 +1 @@
+# WIP: [wip] blackwell: sm100 fused INT4+SiLU Triton stub