[Perf][DSV4.1] Optimize MegaMoE staging and NVFP4 cache gathers

[WoosukKwon]

Purpose

Reduce DeepSeek V4.1 MegaMoE/MegaAttention prefill overhead: tile multiple tokens
per MegaMoE staging program for sufficiently large inputs, and size the NVFP4
gather grid by request count and output capacity. Keep conservative launch
choices for small inputs. Both transformations preserve the existing arithmetic
and add no persistent buffers.

Dispatch uses eight-token staging tiles from 64 tokens, otherwise one token.
On GB200 at H5120, forced eight-token tiles regress 2–16-token kernels, tie at
32, and improve from 64 onward (both scale layouts):

Tokens	Native staging us	Eight-token tile us
16	2.400	2.752
32	2.880	2.816
64	3.328	3.136
128	4.672	3.296
8192	175.616	32.416
16384	348.993	60.736

NVFP4 gathers retain 128 workers for short rows and chunks of >=16 requests.
Otherwise use min(2048 // num_reqs, max(128, compressed_capacity)) workers.
At 100K gathered rows, 1/2/4/8-request kernel latency falls by 88%/76%/50%/18%;
16/32 requests are unchanged. The caller passes a view of the compressed region
so SWA workspace does not inflate its capacity. No allocation, device-to-host
synchronization or actual-KV decode dispatch is added.

Kernel sweeps use CUPTI, CUDA graphs and cold L2, checking bitwise equality
before timing. Final branch boundary/stride checks are recorded below.

Performance and validation

GB200×4, DEP4/MegaMoE/MegaAttention, model revision
dba1be0a40aa45a94ad051997016db3960a90277, FlashInfer 0.6.18.post1.
Earlier paired ten-sample fixed-forward confirmations with matching tuned GEMMs:

Workload per DP rank	Baseline ms	Candidate ms	Reduction
Q16K, KV100K	415.559	400.955	3.51%
B31Q1 + B1Q8K, KV100K	224.438	207.711	7.45%

Mixed baseline SD 9.124 ms versus candidate 0.697 ms; all samples retained.
The earlier three-sample mixed screen improved 6.14%. B8Q2K improves 1.73%;
decode/draft are essentially unchanged. No new model or workspace buffers.
These are synthetic fixed-forward GPU timings, not serving throughput.

Matched full real-weight GSM8K: 1272→1271 /1319 (96.44→96.36% flexible extraction;
7 improve, 8 regress). DSpark acceptance length 3.8281→3.8289. Real-prefill
traces confirm both changed paths execute. Both accuracy arms carry the separate
sparse-remap correctness guard, which is not included in this PR.

GPU validation at e2a2356a6229: 29 staging tests and 24 gather tests pass.
The gather cases cover 1/8/32 requests, 70/128/129/2051 rows, partial gathers,
non-contiguous request strides, untouched bounds, and CUDA-graph replay after
changing lengths in persistent buffers.

.venv/bin/python -m pytest tests/models/test_deepseek_v4_mega_moe.py \
  -k 'fused_input_staging or stages_shared_scale_tma_layout' -q -rs
.venv/bin/python -m pytest tests/kernels/test_compressor_kv_cache.py \
  -k v41_nvfp4_gather_matches_insert -q -rs

All pre-commit hooks pass. The follow-up removes only benchmark scripts;
production code and correctness tests retain the validated contents. Final dispatched-kernel checks preserve short gathers
and small staging inputs. At H5120/N64, the final staging comparison is
3.360→3.264 us; the gain near the boundary is modest. At H7168/N64,
3.712→3.168 us supports the same cutoff. Large-input gains remain intact.
The earlier real-weight quality evaluation is reused: this dispatch refinement
preserves the evaluated arithmetic; the narrower view changes launch bounds only.

Validation commit e2a2356a6229 full-model checks pass with finite outputs and zero
eager/graph, changed-input and restored-input replay errors:

• Q16K: 400.490 ms (three samples), consistent with the earlier confirmation.
• DEP B32 adaptive verification, draft length 5 and seeded lengths 1–6:
  112 total tokens; same-instance original/automatic staging measures
  18.310→18.285 ms (0.14%), with matched inputs, KV and GEMM tactics.
  A post-timing trace confirms all 40 staging calls: summed kernel time
  144.00→112.16 us; MegaMoE time remains 11.703 ms in both arms.

An earlier unpaired B32 check measured 30.287 ms with high spread. It is retained
as an unexplained slow run, not used to estimate the staging effect. The paired
check above isolates staging and does not reproduce that regression. Compilation
is excluded from timing and guarded against during timed work.

Duplicate-work check

No open PR found for tiled DeepSeek MegaMoE staging or CUDA NVFP4 gather sizing.
#56720 sizes a ROCm FP8 gather grid and explicitly preserves CUDA dispatch;
this change targets the separate CUDA NVFP4 branch. #43592 targets MiniMax-M2
prepared-input fusion rather than this DeepSeek staging kernel. #55109 fixes
block-table strides; #57489 adds gfx950 FP4 attention. Neither implements these
CUDA launch heuristics.

AI assistance was used to implement, benchmark and test this change.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[WoosukKwon]

/ci run

[github-actions]

❌ This PR is 3 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89887 for commit 08d438b242c1.