Summary
GroupedGemmPersistent3WGKernel.forward runs two memory-bound ops inside every call that the GEMM itself does not require:
C = torch.zeros(numel, N) — allocates and zero-fills the output each call, but the kernel overwrites every valid row (full tiles via TMA-store, partial via predicated STG), so the memset is pure overhead.F.pad(A, ...) — pads A by block_m guard rows every call (required_rows = numel + block_m > A.shape[0] always holds), i.e. a full numel×K copy. The guard rows only matter when an expert's row count is not a multiple of block_m; on aligned workloads no tile overreads numel, so the pad protects nothing.
On a GLM-up shape (numel=262144, N=4096, K=6144, bf16, H200) these add ~2.8 ms (~13%) on top of a ~20 ms GEMM. They also distort grouped-GEMM benchmarks: a kernel timed through this entry point pays the tax while baselines such as DeepGEMM (pre-allocated output D, alignment required) do not.
Proposed change
- Add an
out= parameter to forward() so callers pre-allocate/reuse the output (matches DeepGEMM's out-param and steady-state serving); use torch.empty instead of torch.zeros internally.
- Detect
block_m alignment (true_sizes % block_m == 0); when aligned, compile the unpadded A_shape=(numel, K) variant and skip the F.pad copy. A compile-time flag threads through the JIT builder and both templates; unaligned inputs keep the existing padded path. Mirrors the _b_copy_ok A_shape conditioning already in moe_grouped_gemm_nopad.py.
Impact
- No behavior change for unaligned inputs (defensive pad retained).
- Aligned workloads (the common MoE contiguous-layout case) drop the per-call pad + memset.
Summary
GroupedGemmPersistent3WGKernel.forwardruns two memory-bound ops inside every call that the GEMM itself does not require:C = torch.zeros(numel, N)— allocates and zero-fills the output each call, but the kernel overwrites every valid row (full tiles via TMA-store, partial via predicated STG), so the memset is pure overhead.F.pad(A, ...)— pads A byblock_mguard rows every call (required_rows = numel + block_m > A.shape[0]always holds), i.e. a fullnumel×Kcopy. The guard rows only matter when an expert's row count is not a multiple ofblock_m; on aligned workloads no tile overreadsnumel, so the pad protects nothing.On a GLM-up shape (numel=262144, N=4096, K=6144, bf16, H200) these add ~2.8 ms (~13%) on top of a ~20 ms GEMM. They also distort grouped-GEMM benchmarks: a kernel timed through this entry point pays the tax while baselines such as DeepGEMM (pre-allocated output
D, alignment required) do not.Proposed change
out=parameter toforward()so callers pre-allocate/reuse the output (matches DeepGEMM's out-param and steady-state serving); usetorch.emptyinstead oftorch.zerosinternally.block_malignment (true_sizes % block_m == 0); when aligned, compile the unpaddedA_shape=(numel, K)variant and skip theF.padcopy. A compile-time flag threads through the JIT builder and both templates; unaligned inputs keep the existing padded path. Mirrors the_b_copy_okA_shape conditioning already inmoe_grouped_gemm_nopad.py.Impact