[Perf] Reduce DFlash2 attention cost and screen exact QK reuse

[yangzhuxinyzx]

Purpose

Reduce DFlash2 target q8 attention at 256K while preserving full context, K16 compensated FP32 sums, N32 online updates, 80 logical splits and the complete FP32 workspace. The user now requires all sixteen target attention layers <=7 ms; complete verification rounds must also reach <=22 ms. Neither target is met.

The selected private E4M3 lookup plus compact P/residual layout improves sixteen-layer attention 14.893 → 12.983 ms. Actual service graph-node attribution confirms 13.077 ms per rank/round. One same-startup A/B improves complete rounds 36.134 → 33.263 ms, with all 12 request pairs identical in tokens, finish reasons and draft acceptance.

New private QK producers reuse each K fragment across three independent M16 products. The best QK-only result improves 6.521 → 5.545 ms, excluding softmax, PV, merge and service execution. It is not a 7-ms attention result. Full attention integration of this producer is incomplete.

Prerequisites #586 and #587 are merged into main. This PR now targets main 213bd9c31ed1ebf2d06ac865a76c5f302ae82b9b; review head is be8f8ca4b931e940fffcd054ee183951273b5c0d. Latest main changes have been merged into the owned branch without conflicts. The measured private native sources remain frozen; the numerical/performance results below refer to those exact libraries. All candidates are hashed private extensions; no serving default changes. Assisted by Codex.

Test Plan

• Build isolated CUDA 12.8 / SM70 extensions with owned caches and rear-GPU leases. Use V100-SXM2-32GB, Torch 2.10.0+cu128, TP4, FP16 activations, E4M3 KV and FP32 state/logits.
• Compare output bytes, full FP32 numerator/max/sum workspaces and canaries through 262144 tokens. For independent QK, compare every FP32 tiled score and guard tile against the frozen producer before timing.
• Time sixteen distinct layer KV allocations with q8 / Hq6 / Hkv1 / D256, page3296 and alternating graph replay trials. Keep operator, QK-stage, traced kernel and complete-round timing boundaries separate.
• Require q2–q7 checks, sanitizers, actual route hits and same-startup output/acceptance comparisons before a new full-attention candidate enters serving. Retain native FA2/original GDN prefill and compact scalar q1 in both arms.

Test Result

The service-checked P-layout candidate passes 65 native checks, 60 tail checks and 25 cases each under memcheck/racecheck/synccheck, with zero errors or hazards. All 12 service request pairs match. Its trace has sixteen full-q8 and sixteen combine kernels in each of eight interior rounds on four ranks. Unprofiled pure decode improves 128.310 → 139.384 tokens/s; the startup's generation trajectory differs from earlier trials, so comparisons remain within the paired startup.

New seven-millisecond screens:

• Eleven QK variants, 264 complete FP32 score comparisons: all exact. N128 panels with K-fragment reuse and separate independent products give the best stage result, 6.521 → 5.545 ms at261888 and 3.539 → 3.040 ms at128K.
• The packaged QK runner passes another24 comparisons with explicit intact-canary assertions. The selected producer also passes24 memcheck comparisons with zero errors. Racecheck, synccheck, full-attention integration and service admission remain incomplete for this producer.
• P-only XOR layout passes65 checks and improves native attention13.004 → 12.827 ms. Combined K/P XOR also passes65 but regresses to13.219 ms. P-only XOR is native-only and has no new service admission.
• Unpadded two-head CTAs pass65 complete-attention checks but regress13.002 → 23.632 ms; rejected before sanitizers or service trials.
• Direct global operands, a context-major grid, six-warp grouping, an exactly decoded FP16 K mirror and wider panels without fragment reuse all lose to their paired QK reference. The FP16 mirror excludes conversion cost and does not justify a serving cache change.

Ten later QK sources and the compact-CTA source regenerate byte for byte. The first historical direct-QK source predates explicit carveout requests and formatting; its measured frozen source/manifest are retained separately. All applicable pre-commit hooks pass. Compiler resource limits are not achieved occupancy; Nsight Compute counters remain unavailable under the driver policy.

Source/DSO hashes, configurations, raw report names, rejected variants and reproduction commands are in docs/design/sm70_dflash2_attention_resources_20260910.md. Absolute artifact paths and GPU ownership are retained in the local handoff.

The user explicitly requested merging the current optimization work on 2026-09-10. The selected service result remains13.077 ms attention and33.263 ms per complete round. The7-ms/22-ms targets, repeated-startup and expanded natural-output admission are still open. No numerical-contract change is admitted.

Main integration validation

This PR is the final eight-file benchmark/builder/documentation delta after #586 and #587. It does not duplicate their runtime loader or graph-selection changes. The combined main tree 0046aea49749fcaa83f0fa2a2af73f5dfed6b3af passes CUDA_VISIBLE_DEVICES='' .venv/bin/python -m pytest -q tests/v1/worker/test_sm70_long_attention_graphs.py tests/kernels/core/test_sm70_dflash2_state_audit.py (26 passed, 1 skipped) and all applicable pre-commit hooks across the full35-file DFlash2 stack. Both prerequisite merge trees exactly match the locally reviewed three-way merges. The final branch retains the same tested source tree after incorporating their GitHub merge commits.

The current main merge is explicitly user-authorized. Keep experimental options opt-in and retain outstanding admission/7-ms/22-ms limits. No additional service deployment or default activation is included.