From bffe108a51ba26d190822fdc8f74051fdae61ac1 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 09:24:15 +0000 Subject: [PATCH 1/2] [wip] blackwell: sm100 FP8 blockwise GEMM dispatch stub Placeholder commit for draft PR. Implementation tracked in: batchgen-agent-metadata/batchgen_design/blackwell/blackwell-kernel-port-v1.md Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/BLACKWELL_KERNELS_WIP.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 docs/BLACKWELL_KERNELS_WIP.md diff --git a/docs/BLACKWELL_KERNELS_WIP.md b/docs/BLACKWELL_KERNELS_WIP.md new file mode 100644 index 00000000..2e85057d --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 FP8 blockwise GEMM dispatch stub From 10204361425a66e5c88df091e075cf422f5933b9 Mon Sep 17 00:00:00 2001 From: v-tairan Copilot user Date: Sat, 30 May 2026 14:44:40 +0000 Subject: [PATCH 2/2] feat(sm100): FP8 blockwise grouped GEMM fallback via row-wise _scaled_mm The compiled SM90a CuTe FP8 blockwise kernel is unavailable on Blackwell, and cuBLAS in torch 2.9+cu129 rejects 1x128/128x128 blockwise FP8 scaling (heuristic returns CUBLAS_STATUS_NOT_SUPPORTED; fast-accum unsupported for 1x128). Only row-wise/tensor-wise FP8 GEMM is available. Add an sm100 branch in grouped_fp8_blockwise_gemm() that emulates deepseek blockwise scaling exactly: split the K contraction into 128-wide blocks and issue one row-wise torch._scaled_mm per block, accumulating partials in fp32. Within a K-block the act scale is constant per row (1x128) and the weight scale is constant per 128-output-row block (128x128, expanded to per-column), so the row-wise GEMM is numerically identical to true blockwise for that block. - Processes the full uniform mtp reserved rows per expert -> static control flow, CUDA-graph compatible (no host sync on seqlens, no data-dependent shapes). Padding-row outputs are ignored by downstream gather. - Covers s1 (gate+up), s3 (down) and fused_s1 automatically since all route through grouped_fp8_blockwise_gemm (fused_s1 falls back to s1_silu on sm100 as its _C kernel is absent). - Dim asserts (K,N multiples of 128; scale shapes) guard silent mis-scaling. Verified on B200 (GLM-5 / MiniMax-M2.5 style dims, E=4, mtp=64, K=2048, N=1536, incl. a zero-token expert): GEMM max rel 0.0039, S1+SiLU 0.0056, fused_s1 matches s1_silu (< 1e-2 acceptance). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- batchgen/moe/grouped_fp8_blockwise_moe.py | 78 +++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/batchgen/moe/grouped_fp8_blockwise_moe.py b/batchgen/moe/grouped_fp8_blockwise_moe.py index c52c47b4..f21d6bdf 100644 --- a/batchgen/moe/grouped_fp8_blockwise_moe.py +++ b/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(