[Perf][DSV4.1] Skinny GEMM dispatch for tiny-M decode projections

[gcanlin]

Purpose

At tiny M (1–4 rows) the dense projections on DeepSeek-V4.1's decode path are dispatch-mismatched: MXFP8 CuTeDSL and BF16 cuBLAS kernels sized for bulk traffic run weight-bound rounds. This PR adds a shape-keyed skinny dispatch table (SM100/B200 only, TP8 shapes) that swaps in the measured winners exactly where they win — a Triton SIMT GEMM for the three MXFP8 projections, CuTe skinny for the BF16 indexer/lm_head projections, and ll_bf16 dotprod for the FP32 compressor scores, 20 (shape, M) cells in total.

Semantics are preserved exactly: MXFP8 candidates keep the per-32-element E8M0 activation scale and FP8 rounding (round_mx in the SIMT kernel reproduces the baseline MXFP8 quantization), and FP32 compressor outputs stay FP32. Everything falls back to the original kernels off-table, off-device, non-BF16 input, with LoRA attached, with a bias, or under VLLM_BATCH_INVARIANT. The dispatch runs at model init for the three unquantized BF16 modules (K3/GLM-5.2 quant-method replacement precedent) and at the model-layer call sites for the fused/wo_a/compressor paths, so module forward and any wrappers are untouched off the table.

Test Plan

Deployment (DeepSeek-V4.1-Flash, 8xB200, TP8/EP, MegaMoE+SP, FlashInfer sparse MLA, FP8 KV, C1/C4/C16 client suites, vllm bench serve, seed 42, ignore-EOS, tokenizer trusted):

vllm serve /gpfs/mszn/models/deepseek-ai/DeepSeek-V4.1-Flash \
  --trust-remote-code --served-model-name dsv4 --tensor-parallel-size 8 \
  --distributed-executor-backend mp --enable-expert-parallel \
  --tokenizer-mode deepseek_v41 --attention-backend FLASHINFER_MLA_SPARSE_DSV41 \
  --kv-cache-dtype fp8 --max-model-len 1048576 --max-num-batched-tokens 65536 \
  --max-num-seqs 96 --gpu-memory-utilization 0.90 \
  --kernel-config '{"moe_backend":"deep_gemm_mega_moe"}' --no-enable-flashinfer-autotune

for C in 1 4; do vllm bench serve --backend vllm --endpoint /v1/completions \
  --model dsv4 --dataset-name random --random-input-len 1024 \
  --random-output-len 256 --random-range-ratio 0 --random-prefix-len 0 \
  --num-prompts 100 --max-concurrency $C --request-rate inf --ignore-eos \
  --seed 42; done
vllm bench serve ... --random-input-len 8192 --random-output-len 1024 \
  --num-prompts 100 --max-concurrency 16  # identical flags otherwise

Unit/kernel tests:

python -m pytest tests/kernels/test_bf16_skinny_gemm.py tests/kernels/test_dsv41_mxfp8_skinny_gemm.py -q

Test Result

• pytest on 8xB200 (SM100): 204 passed, 267 skipped (the skips are the existing SM90/SM103-gated K3/GLM cases; all DSV4.1 cases and every shared case for this arch pass). Coverage includes the 20-cell table vs selected.json configs, install gating (SM100/dtype/LoRA/shape/bias/M-miss fallbacks byte-equal to the original calls), SIMT-vs-reference rel-RMSE across 5 activation scales (1e-5..100), the wo_a chain vs the real DeepGEMM fp8_einsum baseline, and compressor ll_bf16 vs torch.mm(out_dtype=fp32).
• pre-commit: ruff check, ruff format, mypy-3.10, check-torch-cuda-call all pass on the touched files.
• Serving A/B (same main revision b7e0cdac5d, baseline = original kernels, PR = skinny dispatch; 0 failed requests in all runs):

random 1024/256, C1:

metric	baseline	PR	delta
out tok/s	117.5	120.3	+2.3%
TPOT p50 (ms)	8.001	7.801	-2.5%
E2EL p50 (ms)	2176.3	2127.1	-2.3%

random 1024/256, C4:

metric	baseline	PR	delta
out tok/s	439.9	442.0	+0.5%
TPOT p50 (ms)	8.523	8.473	-0.6%

random 8192/1024, C16:

metric	baseline	PR	delta
out tok/s	1375.6	1385.5	+0.7%
TPOT p50 (ms)	9.428	9.385	-0.5%

Prefix-cache suite (prefix_repetition 32768/2048/1024, C96) is unchanged within the run-to-run spread (out 4064.7 -> 4051.4 tok/s, -0.3%; TPOT 18.04 -> 18.48 ms); the dispatch only covers M<=4 cells, which C96 decode does not satisfy per step.

C1 was measured twice on fresh servers: round 1 out 120.3 tok/s (+2.3%) / TPOT 7.801 ms; round 2 out 119.9 tok/s (+2.0%) / TPOT 7.802 ms against the same baseline. The direction and size hold across server starts.

• Quality: 150-question GSM8K smoke (5-shot, flexible-extract, temperature 0, identical harness to [Perf][DSV4.1] Pad shared experts for native MegaMoE fusion #56568's eval): 98.0% / strict-match 83.3% on the PR configuration (model family trend on this main revision is 96-98.7%, so no quality regression is visible). Numerics identical-by-construction is asserted by the scale-preserving specializations and covered by the two pytest suites above. A V1-runner startup plus GSM8K smoke on the same tree also passes (piecewise graph capture succeeds on all ranks).

This PR was developed with AI assistance (Kimi Code CLI): the measured winning dispatch table comes from a prior experimental campaign by this team (documented in-tree in benchmarks/kernels/experimental/dsv41_skinny/), and I integrated, verified and measured it; the human submitter reviewed every changed line.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gcanlin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork