[blackwell][draft] sm100: Router GEMM Triton kernel fused_router_gemm.py (#229)

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(

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

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

docs/BLACKWELL_KERNELS_WIP.md

@@ -0,0 +1 @@
+# WIP: [wip] blackwell: sm100 router GEMM Triton kernel stub