[Perf][SM70] Optimize Qwen3.8 Flash Next prefill

[yangzhuxinyzx]

Public migration note (2026-08-27): this Draft is the public continuation of private PR #23, mirrored at exact head a271361 after public main resumed through #342. The original implementation notes and evidence are preserved below.

This is stacked on the migrated Qwen3.8 Flash Next adaptation and inherits its consolidation dependency on public #338.

Objective

Improve and validate Qwen3.8 Flash Next NVFP4 TP4 prefill performance on four V100-SXM2-32GB GPUs. Decode/MTP optimization is explicitly out of scope.

Base

• Integration branch: codex/v100-qwen38-flash-next-nvfp4-20260826-140311
• BASE_SHA: 383ff458dbc8af1dda79a1a29bf700019030f054

Planned gates

• Exact TP4/FP16/Flash-V100 route audit
• 8K correctness, TTFT, and pure prefill baseline
• Profile-driven exact-shape microbench before source changes
• Focused CPU/GPU numerical tests
• Same-contract 8K/32K/128K trend validation

Status

Draft campaign boundary created; no prefill source change yet.

[yangzhuxinyzx]

SM70 QSA indexer hotspot update

Commit 9d62e5db98 replaces the exact Qwen3.8 single-request FP16 QSA index
scoring shape with a bounded cuBLAS Tensor Core route while retaining the
paged Triton fallback for decode, short work, BF16, multi-request and generic
shapes.

The cached SM70 PTX showed that both QSA selected attention and index scoring
contained no HMMA/WMMA. A selected-attention WMMA prototype was rejected: it
was only 1.0748x at 8192 rows and failed the numerical gate. The index scorer
is the useful route:

• 1024 rows, 131K context, scoring plus paged gather: 35.767 -> 2.108 ms
  (16.964x), max absolute error 8.11e-6, relative L2 2.90e-7.
• Complete score/top-k/expand: 38.792 -> 3.118 ms (12.441x).
• Crossover sweep: 1K 0.733x (fallback), 2K 2.164x, 4K 3.658x, 8K 4.900x,
  16K 8.064x, 32K 11.102x.
• Production admission therefore requires at least 512 rows and 1,048,576
  score elements. The route is default-on only for exact SM70 FP16
  [rows,4,128] against a single-request FP16 paged MQA cache.
• Focused QSA suites: 7 passed. Ruff and diff checks pass.

The exact 8192-token/top-10 NVFP4 MoE tuning screen was also closed: default
W13/W2 is 6.702 ms and measured tuning is 6.725 ms (0.9966x), so the large
MoE tune limit remains rejected.

The matched TP4 token-hash/quality gate and one Nsight Systems capture are
still pending shared-machine availability; no endpoint speedup is claimed by
this comment.

[yangzhuxinyzx]

Source audit update against the repaired #345 stack:

• Scope is limited to the exact Qwen4Exp NVFP4 grouped-prefill/QSA routes, plus an explicit prebuilt FlashQLA extension path; decode, short prompts, unsupported shapes, and generic models retain their existing fallbacks.
• Changed-file pre-commit, Python compile, and tests/models/qwen4_exp/test_qsa_ops.py all pass (5 tests).
• Recorded same-contract evidence: 8K grouped+QSA prefill 21.339760s -> 2.315786s (9.2149x); exact QSA kernel 1320.204ms -> 85.261ms; BN32 improves 32K/64K/131K by 15.79%/14.00%/11.49% with matching token hashes.
• Numerical evidence is bounded: tile max diff 0.000122, no non-finite values; cuBLAS indexer max abs 8.11e-6, rel-L2 2.9e-7, 1023/1024 top-k selections identical with the sole change at a near-tie boundary, and final sparse cosine 0.9999992.
• The faster cuBLAS indexer is default-on only behind strict SM70/FP16/exact-shape/long-prefill gates and has VLLM_SM70_QSA_INDEXER_CUBLAS=0 rollback. The rejected D256 WMMA path (rel-L2 0.5627) is not included.

Decision under the project policy: the large prefill gain with bounded numerical drift and narrow fallback-safe gates is acceptable as default-on; greedy identity is not required. The prior remote failure came from #345 before its #359 fixes, not from this PR. I will resync/re-run CI after #345 lands, then merge if the final tree remains identical to the audited projection.

[yangzhuxinyzx]

SM70 prefill update (commit 8f151b5):

• Retained TP4/V2/no-MTP/NVFP4/FP16-KV/FlashAttention gate, one model startup.
• Pure prefill: 32K 4532.07 tok/s, 64K 4446.64 tok/s, 131K 4108.16 tok/s.
• Versus retained N32 gate: 1.2509x / 1.4600x / 1.7753x; all output token streams and hashes exactly match.
• Nsight Systems identifies selected QSA as the largest remaining kernel hotspot (36.42% aggregate kernel time, 55.151 ms/layer/rank at the 8192-token capture).
• The committed N16/two-warp SM70 launch profile reduces isolated selected-QSA latency 7.15% and raises the occupancy ceiling from two to three CTAs/SM. Short-prefill and non-SM70 profiles are unchanged.
• A one-shot locality diagnostic found 81.92% selected-block overlap between adjacent rows (Jaccard 0.6980). The next candidate will exploit two-row K/V reuse and must clear isolated correctness/latency gates before another full-model load.