[Kernel][SM70] Share FP16 attention fast paths with E4M3 KV

[yangzhuxinyzx]

Purpose

E4M3 FP8 KV cache could not consistently reuse the qualified SM70 FP16 attention fast paths. Long prefix prefill could miss the Q8192 75T kernel, decode had an artificial batch-16 cap, mixed chunked-prefill batches sent resident q=1..16 rows through serial paged prefill, and the page-800 wave loader addressed only two logical pages even when p896/p1664 spans three or four. The last issue produced a maximum absolute error of 1.2546e-3 at 256K.

This change gives the Qwen3.8 GQA6/D256 E4M3 route the corresponding FP16 schedules:

• gather E4M3 prefix KV once into the shared FP16 workspace, then dispatch Q8000/Q8192 Tensor Core attention with FP32 accumulation;
• keep E4M3 softmax probabilities, PV state, partition outputs, and partition reduction in FP32 until the public FP16 output boundary;
• enable device-side p64/p256/long-wave XQA selection by default for aligned page sizes, including page 800 and 1568;
• remove the E4M3 B<=16 admission cap and route small resident rows in mixed prefill/decode batches through one paged XQA call;
• extend the default no-MTP full CUDA Graph capture set through B32 while keeping larger scheduler batches on the same accelerated attention route through piecewise graphs;
• preserve the dedicated grouped FP32 route for DFlash selector rows and the existing E5M2 numerical contract;
• make the cold Q8192 benchmark CUDA-graph-only, allow checkpoint weight-quantization metadata including NVFP4, and withhold a decode claim until at least 63 token intervals are measured;
• correct CPython ABI suffix declarations for the SM70 pybind extensions so a complete source build copies the generated modules.

All new fast paths are default on and have explicit environment rollback controls. The E4M3 FP32 partial workspace adds about 25 MiB for B1/256K and about 200 MiB for B32/256K on this shape.

Test Plan

• Build all CUDA/C++ extensions from this worktree with CUDA 12.8, TORCH_CUDA_ARCH_LIST=7.0, VLLM_FLASH_ATTN_SM70=ON, and VLLM_SM70_79T_PREFILL=ON.
• Run focused planner, prefix/decode concurrency, E4M3 numerical, policy, and CUDA Graph dispatcher tests on V100.
• Compare default E4M3 XQA against the scalar FP32-state route at page 800/1568, B1/B32, and 256K under torch.cuda.CUDAGraph.
• Benchmark the source-built Q8192 operator at KV128K and KV256K against a sampled FP32 reference.
• Run Qwen3.8-27B NVFP4/compressed-tensors on TP4 V100 with FP16 execution, E4M3 KV, Q8192 chunks, max length 262144, prefix caching off, no speculative decoding, and FULL_AND_PIECEWISE CUDA graphs. Do not use --enforce-eager.
• Run standard vllm bench serve cohorts at concurrency 2/4/8/16/32 with exact 2048-token input and 256-token output, then check a 32-request natural-language answer burst.

Test Result

• Source build_ext --inplace: exit 0 for the SM70 core, stable libtorch, FA2/Q8192, Flash-V100, H3, sparse-attention, MoE, and bundled FlashQLA extensions. The optional Rust frontend was skipped because cargo is absent; it is not required by the SM70 CUDA runtime.
• Final pre-commit: all hooks passed. Python compileall and git diff --check passed.
• Focused tests before packaging: 148 passed. Final source-built DSO rerun: 41 + 40 + 37 = 118 passed.
• Q8192 source-final operator on one V100-SXM2-32GB:
  • KV128K: 80.9762 ms, 77.0142 TFLOP/s, relative L2 0.002387, max abs 6.06e-5;
  • KV256K: 166.9019 ms, 75.9654 TFLOP/s, relative L2 0.002607, max abs 7.94e-5;
  • both outputs finite.
• Source-final CUDA Graph decode:
  • B1/page800/256K: scalar 3.2975 ms, XQA 0.4834 ms, 6.82x, bitwise equal;
  • B1/page1568/256K: scalar 3.4051 ms, XQA 0.4922 ms, 6.92x, max abs 2.38e-7;
  • B32/page800/256K: scalar 96.4159 ms, XQA 14.6154 ms, 6.60x, max abs 4.77e-7.
• TP4 full-model CUDA Graph gate:
  • 16K: 3.8562-second TTFT, 4149.15 prompt tok/s, answer 海蓝石榴；木星;
  • 256K natural-EOS quality run: 102.7519-second TTFT, 2491.44 prompt tok/s, complete answer 校验词是「海蓝石榴」，太阳系最大的行星是木星。; its 15 decode intervals are not reported as a speed baseline;
  • 256K fixed 256-token speed run: 102.9135-second TTFT, 2487.53 prompt tok/s, 255 decode intervals in 5.3902 seconds, 47.308 tok/s / 21.138 ms TPOT;
  • both retrieval and knowledge checks passed, cached_tokens=0, and the log contains enforce_eager=False with all four ranks completing FULL graph capture;
  • every rank recorded 480 Q8192 FP32-accumulated calls and 496 E4M3 bridge calls for the 256K request; shutdown summaries recorded 48 dynamic page-800 E4M3 XQA decode calls per rank;
  • no NaN, Inf, CUDA error, OOM, traceback, or worker failure.
• Default CUDA Graph concurrency and serving matrix:

  • startup captured both FULL and PIECEWISE graphs for [1, 2, 4, 8, 16, 32]; B32 recorded the page-800 E4M3 XQA route;

  • standard vllm bench serve, exact 2048 input / 256 output, prefix cache off, graph mode, zero failures:

    C	Median TTFT	Derived pure decode	ITL median/P90	Output TPS	Median request wall
    2	0.8554 s	115.725 tok/s	17.282 / 17.406 ms	92.288 tok/s	5.539 s
    4	1.6962 s	223.730 tok/s	17.879 / 18.038 ms	151.222 tok/s	6.761 s
    8	5.2214 s	418.723 tok/s	19.106 / 19.248 ms	203.168 tok/s	10.076 s
    16	10.7109 s	708.450 tok/s	22.585 / 22.837 ms	248.859 tok/s	16.451 s
    32	21.7727 s	983.986 tok/s	32.521 / 32.824 ms	272.868 tok/s	30.008 s

  • all 62 requests generated exactly 256 tokens; all 32 natural-language responses returned the same complete 海蓝石榴；木星 answer and natural EOS;

  • B2/B4/B8/B16/B32 page-800/256K CUDA Graph operator checks were finite, differed from scalar E4M3 by at most 4.77e-7, and were 3.64x/6.12x/6.35x/6.46x/6.59x faster;

  • the full policy suite passes 153/153 with the source Flash-V100 package, and pre-commit passes all hooks.

• Final SHA256: core 2557f6b7...fa4c4, stable libtorch 622af596...d162, FA2/Q8192 aa657e16...5add, Flash-V100 66df783d...70b3. Local imports and process maps resolve the attention DSOs from this owned worktree.