diff --git a/batchgen/moe/routing/cuda_routing.py b/batchgen/moe/routing/cuda_routing.py index 062cbb81c..0894d0105 100644 --- a/batchgen/moe/routing/cuda_routing.py +++ b/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( diff --git a/batchgen_kernels/triton/__init__.py b/batchgen_kernels/triton/__init__.py index 27874ab1f..ab8d30621 100644 --- a/batchgen_kernels/triton/__init__.py +++ b/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 diff --git a/batchgen_kernels/triton/fused_router_gemm.py b/batchgen_kernels/triton/fused_router_gemm.py new file mode 100644 index 000000000..f742dff46 --- /dev/null +++ b/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 diff --git a/docs/BLACKWELL_KERNELS_WIP.md b/docs/BLACKWELL_KERNELS_WIP.md new file mode 100644 index 000000000..2e5ce54b7 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 router GEMM Triton kernel stub