onnxruntime
395c8b23 - [CUDA] Skip padded rows in the DeepGEMM MoE pack and SwiGLU kernels

Commit
71 days ago
[CUDA] Skip padded rows in the DeepGEMM MoE pack and SwiGLU kernels The decode MoE chain lays tokens out as [32 experts, 64 padded rows, K] but only masked_m[expert] rows are real, typically 1-3 of 64. PackInputKernel and InterleavedSwiGLUKernel both walked the full padded extent, so their cost was constant regardless of how much work there actually was: ncu measured them at a flat 11.0 us and 14.4 us per layer with DRAM SOL of 0.18% and 20.9%. Bound both loops by the per-expert row count instead. The padded rows are now left stale rather than zero-filled, which is safe because the GEMMs are row-independent and UnpackOutputKernel copies back only the first count rows. This was verified bitwise against a poisoned workspace (NaN, +/-Inf, 1e30 across 128 patterns, ragged and zero counts): the compact output is identical, even though 40-60% of the padded FC2 rows do become NaN. Also give SwiGLU a 2D grid over experts and load the adjacent gate/linear pair as one __nv_bfloat162; that part is an exact refactor, bit-identical on valid rows. The two kernels now scale with real work and sit within 0.35 us of the empty kernel launch floor: 1.55 us and 1.51 us at the realistic decode point. Bounding by count removes an implicit clamp the old padded loop bound gave for free, so pin the invariant with a static_assert alongside the alignment one the vectorized load now needs. DeepSeek-V4-Flash, 8xH200, world 8, 1024 prompt / 128 gen: decode 92.12 -> 96.23 tps, 10.86 -> 10.39 ms/token. Acceptance rate 0.329 and 2.65 tokens/step are unchanged, and rank argmax disagreements stay 0.
Author
Parents
Loading