[sm100] Packed FP4×FP4 mega-MoE kernel (W4A4) with per-band BLOCK_K

[Romaosir]

Adds a dedicated W4A4 (FP4 weights × FP4 activations) fused mega-MoE kernel for SM100 — sm100_fp4_mega_moe — a packed-E2M1 specialization that issues tcgen05.mma.kind::mxf4, plus a per-band BLOCK_K selection that exploits packed-FP4's lower SMEM footprint for a measured kernel-level speedup over the existing FP8-acts and 0.1.0 W4A4 paths.

Paired with sgl-project/sglang#28210 (sgl-project/sglang#28210) — version bump + one-line env propagation; the existing W4A4 env path transparently routes to this kernel, no other sglang code changes.

What's new

• sm100_fp4_mega_moe_impl: packed float_e2m1_t A and B (no multi-mode branches), hand-built packed-byte SMEM descriptors, dedicated packed_fp4_2sm_tma_load_2d (byte-typed tensor maps, per-CTA cta_group::2 issues with leader-centralized 2× tx accounting), SwiGLU epilogue cast straight to packed FP4. Half the activation bytes + SMEM of the FP8-acts path.
• Per-band BLOCK_K=256 + 128B swizzle for the block_m ∈ {32, 64, 192} bands ({96, 128} stay 128). Packed FP4's 256-deep K stage costs ~31 KB SMEM (vs ~57 KB for FP8 acts), so it sustains 6–10 pipeline stages where the FP8-acts retrofit cannot afford the knob. Larger, more-contiguous TMA transactions raise DRAM efficiency. Overrides: DG_BLOCK_K (128/256), DG_NUM_STAGES (clamp).
• mega_moe_pre_dispatch: packs activations to E2M1 + UE8M0 SF into the symmetric buffer.
• csrc/apis/mega.hpp: env-selected dispatch (DG_MEGA_MOE_FP4 / DG_USE_FP4_ACTS); symm-buffer halves the activation regions under FP4 acts.
• Tests: FP4-acts coverage + --shape-ntoks/--bench-reps sweep with in-process interleaved A/B (DG_BLOCKK_AB/DG_STAGES_AB/DG_WAVE_AB).

How it differs from the existing W4A4 path (sm100_fp8_fp4_mega_moe<kUseFp4Acts, kUseMxf4Kind>)

Same architecture (warp roles, scheduler, comm/barriers, packed-E2M1 symm buffer) and the same final tcgen05.mma.kind::mxf4 instruction — the difference is a dedicated specialization (enabling the per-band BLOCK_K win) instead of a template-flag retrofit of the FP8 kernel:

aspect	existing W4A4 (retrofit)	this kernel (specialized)
code shape	1 kernel + 4 template bools; conditional dtypes, swizzle/SMEM ternaries, per-phase MMA-desc dtype flips	unconditional packed float_e2m1_t A and B; no multi-mode branches
BLOCK_K	fixed 128	per-band 256/128 (the measured speedup)
SMEM descriptors	generic make_umma_desc via _unpacksmem/_ALIGN16B	hand-built packed-byte strides (generic helper miscounts K for packed operands)
TMA loads	generic tma::copy	dedicated packed_fp4_2sm_tma_load_2d (byte tensor maps, 64B/128B swizzle atoms, leader-centralized 2× tx)
MMA wrapper	generic instr-desc + dtype-field flip	dedicated ptx::SM100_MMA_MXFP4_2x1SM_SS
standalone validation	—	sm100_fp4_fp4_gemm_1d1d single-GPU GEMM clone

Performance (4× B200, DSV4-Flash: h=4096 / i=2048 / E=256 / topk=6; bit-exact numerics vs the BLOCK_K=128 path, torch.equal)

Per-band BLOCK_K vs all-128 — in-process 4-rank interleaved A/B, 20 reps (the precision claim):
32-band ~+10%, 64-band +2–6%, 192-band ~+4%. Mechanism (NCU): DRAM efficiency from larger contiguous transactions (64% → 85% busy at ntok=512).

vs W4A4-old — the existing 0.1.0-release W4A4 kernel (sgl-project/DeepGEMM#27, sm100_fp8_fp4_mega_moe<kUseFp4Acts, kUseMxf4Kind>) — both kernels measured through the same binding (DG_FORCE_FP8FP4 toggle) to isolate kernel from packaging:

• 1-rank NCU 537.6 µs vs 712.5 µs (+24.6%), DRAM 85.2% vs 64.3% — the block_k=256 data path (the 0.1.0 retrofit is fixed at block_k=128).
• Fairness: the W4A4-old arm is true W4A4 — L2 contracts FP4-width (L2_SHAPE_K identical, compile-time); its DRAM 64.3% = the expected block_k=128 efficiency, no traffic confound.
• 1-rank over-weights the memory regime; end-to-end the kernel win dilutes through Amdahl (mega-MoE is a fraction of the workload — the decode path is ~82% sm100_tf32_hc_prenorm_gemm, shared with both baselines).

End-to-end 5-concurrency sweep (TP4+DP4, sequential, same GPU group, 1 rep/cell after per-concurrency warmup, num_prompts raised for steady state). Ours and the 0.1.0 W4A4 ("old") run via the same fork pybind; W4A8 (FP8 acts) is the 0.1.0-release reference on its own tvm-ffi binding (quant-scheme baseline, not a same-binding comparison):

Prefill (in=1024, out=1) — input throughput (tok/s) · median TTFT (ms):

conc	Ours tput	old tput	W4A8 tput	Δ tput Ours/old	Ours TTFT	old TTFT	Δ TTFT Ours/old
1	5847	5886	5683	−0.7%	156.8	154.8	+1.4%
4	14802	14478	10095	+2.2%	289.5	292.2	−0.9%
32	54114	54755	52762	−1.2%	417.6	402.2	+3.8%
128	126557	117838	123033	+7.4%	1110.6	1145.8	−3.1%
512	133116	130299	124708	+2.2%	4575.7	4672.6	−2.1%

Mixed (in=1024, out=256) — total throughput (tok/s) · median TPOT (ms):

conc	Ours tput	old tput	W4A8 tput	Δ tput Ours/old	Ours TPOT	old TPOT	Δ TPOT Ours/old
1	342	338	340	+1.1%	14.02	14.17	−1.1%
4	1268	1291	1274	−1.8%	14.55	14.82	−1.8%
32	8212	8055	8149	+2.0%	17.82	18.11	−1.6%
128	24982	24544	24204	+1.8%	22.40	22.63	−1.0%
512	51295	50201	50363	+2.2%	33.50	37.81	−11.4%

Reading: the kernel win surfaces at production concurrency (c ≥ 128) where mega-MoE's share is largest — prefill input-throughput +2–7%, mixed throughput +2%, TPOT down to −11% at c=512. Low concurrency (c ≤ 4) is within the e2e noise band, and the c=1/4 mixed cells (num_prompts = 12) are under-sampled.

NCU Speed-of-Light (1-rank isolation, binding-controlled, ntok=512)

Profiling the fused 4-rank kernel with NCU is invalid (persistent kernel dominated by cross-rank NVLink-barrier spinning → SOL collapses). Single-rank removes the barrier and gives true SOL. All three arms measured at the identical shape (hidden=4096 / inter=2048 / E=256 / topk=6, ntok=512) through the same fork binding:

metric	W4A4-Ours (packed, bk256)	W4A4-old (fp8_fp4, bk128)	W4A8 (fp8 acts)
Duration (µs)	537.6	712.5	753.2
Compute (SM) %	33.6	28.6	42.5
Memory %	85.2	64.3	61.4
DRAM throughput (TB/s)	6.54	4.93	4.71
L2 %	56.1	42.4	40.3
Achieved occupancy %	22.5	22.4	22.5
Registers/thread	128	128	128

Ours +24.6% vs W4A4-old, +28.6% vs W4A8 (1-rank). Mechanism: same occupancy/registers (identical launch config) → the win is the data path. Ours' block_k=256 sustains 85% memory SOL / 6.54 TB/s DRAM vs W4A4-old's block_k=128 at 64% / 4.93 TB/s — the larger-contiguous-transaction DRAM-efficiency win. W4A8 carries the heaviest compute (42.5% SM, FP8 acts = 1 byte/act) and the least DRAM headroom. 1-rank over-weights the memory regime (255 experts local, ~2 tokens/expert); use it for the mechanism, the e2e sweep above for the realistic magnitude.

Accuracy (eval_v2, DSV4-Flash)

eval	W4A4-Ours	W4A4-old	W4A8
GSM8K	0.966	0.962	0.962
HumanEval (pass@1)	0.898	0.888	0.887

W4A4-Ours is at or above both the W4A4-old and W4A8 baselines on every eval run — the packed FP4×FP4 path + per-band BLOCK_K preserve accuracy.

Notes

• The existing W4A4 env path transparently routes to this kernel — see the paired sglang PR [deepseek-v4] Bump sgl-deep-gemm for packed FP4×FP4 mega-MoE (W4A4) sglang#28210 ([deepseek-v4] Bump sgl-deep-gemm for packed FP4×FP4 mega-MoE (W4A4) sglang#28210). No other sglang code changes.
• The bit-exact legacy-baseline comparison (torch.equal vs deep_ep + grouped GEMM) needs an internal deep_ep build with ElasticBuffer; rerun in CI before merge.