[Kernel] Reduce DFlash2 weight and scale memory on SM70

[yangzhuxinyzx]

Purpose

DFlash2 QPN2 stores target NVFP4 codes separately from TurboMind, and FP32 LM-head serving still prepares unused FP16 packed copies. Add shared-code and compact-scale paths, and omit unused LM-head packing while retaining accelerated QPN2/QPN8 execution.

Under the production TP4 configuration, the latest follow-up reduces model loading from 8.203934 to 6.286138 GiB/rank after code sharing, saving 7.671183 GiB across TP4. Automatic KV budget increases from 15.76 to 17.70 GiB/rank. The original dual-code loading measurement was 11.08 GiB/rank.

• VLLM_SM70_NVFP4_QPN2_SHARED_WEIGHT=1 reads TurboMind codes directly, removing 2.835693 GiB/rank of duplicate target codes. Shared scheduling/cache recovery retains the existing arithmetic and dynamic-M fallback.
• VLLM_SM70_NVFP4_QPN2_SHARED_SCALES=1 keeps QPN2 E4M3 scales and restores exact FP16 scales temporarily for TurboMind fallback. Small-M QPN2 decode and bounded dense prefill are unchanged. Admission requires rebuilt operators, shared codes, the DFlash2 TP4 q7/no-DBO contract, and graph capture sizes <=32. Larger graphs retain persistent scales to avoid capturing a separate temporary for every projection. Both switches default off.
• FP32 dense/candidate LM-head paths retain the original FP16 parameter and QPN8 screening layout. Skip unused FP16 packing and rerank scratch; explicit packed Tensor Core top1 still receives its layout. Allocation stacks identify the second 606.25 MiB/rank copy as draft placeholder-head packing retained by the native cache after target-head sharing. Avoiding unused packing prevents both copies without changing candidate selection or logits arithmetic.

Base: e5d63c51f0fcc1ddf75d229e3df06bf52df206f5 on onecat/main. Head: d26804a9ccce2f2ac3fa4f3a69874575c1aa98e1. This extends the existing memory recovery PR rather than creating a duplicate. AI assistance was used; human review remains pending. Keep this PR Draft.

Test Plan

Verify loader aliases, capability and graph-size gates, explicit packed top1, warmup dispatch, real checkpoint scale/output bits, and changed-input CUDA Graph replay. Measure operator costs independently from whole-model throughput. Preserve production serving: 4×V100-SXM2-32GB, Torch2.10.0+cu128/CUDA12.8, TP4, FP16 model/draft KV, maxlen262144, utilization0.8, automatic E4M3 KV, chunk4096, maxseq4, DFlash2 q7/probabilistic, FP32 logits, context pipeline/KV graph and CUDA Graph. Sampling remains temperature1/top_p0.95/top_k20/seed0, xhigh thinking and natural EOS.

Test Result

• .venv/bin/python -m pytest tests/quantization/test_sm70_warmup.py tests/quantization/test_sm70_nvfp4_qpn2.py tests/model_executor/layers/test_sm70_lm_head.py -q: 43 passed.
• CUDA sidecar builds and targeted pre-commit checks pass, including Ruff, changed-line clang-format, mypy and Markdown.
• Previous shared-code recovery passes 336 real-shard cases, with weighted repeated-operator M8/M16/M32 ratios1.0013/0.9869/0.9591. A six-tensor working set gives1.0259/0.9915/0.9019; these are projection measurements, not model speed.
• Compact scales pass 224 ordinary/gated cases on24 real TP4 shards at M8/16/32/33/135/512/1023/1024. Restored FP16 scales and outputs are bitwise identical, including changed-input graph replay. benchmark_sm70_nvfp4_shared_weight.py --compact-scales reproduces the comparison against shared codes with persistent scales.
• Compact/persistent projection ratios are M8 1.0008, M16 1.0004, M32 1.0002. Fallback costs are explicit: M33 1.1394, M135 1.0587, M512 1.0307, M1023 1.0163. M1024 bounded prefill is 0.9996. Scale restoration adds about2.87–3.43ms across the isolated fallback projection sequence; this is not TTFT or model latency.
• Six real LM-head cases preserve QPN8 code/scale bits, candidate IDs and FP32 logits in eager and changed-input graph replay. Candidate/control timing ratios range 0.9969–1.0010.

Latest completed production allocation:

Metric	Prior shared-code cohort	Compact scales + head cleanup
Model loading, GiB/rank	8.203934	6.286138
Model loading, GiB/TP4	32.815735	25.144552
Automatic KV budget, GiB/rank	15.76	17.70
Logical KV tokens	1,479,578	1,661,426
Graph capture increment, GiB/rank	0.26	0.26
Idle worker NVML, MiB/rank	26,030	26,036

All four ranks agree. Census operations do not change allocated bytes. Both unused head matrices and persistent TurboMind scales are absent; the QPN8 acceleration copy remains. Idle snapshots show no active/waiting requests and zero KV usage. Automatic KV capacity grows12.29%, so total residency stays near the existing memory budget.

The latest MBPP28 warmup/three measured requests match all 754 tokens of the archived shared-code cohort; MBPP0 matches all 1093 tokens. Both finish naturally with nonempty answers. Median pure decode is 225.66tok/s, verification round 19.068ms, TTFT 107.95ms. Archived same-output values were221.40tok/s,19.435ms and112.37ms. No focused slowdown appears, but these are not contemporaneous controls; do not claim the approximately2% difference as a speedup.

Remaining acceptance and evidence boundaries

The current paired control was terminated during compilation and has no endpoint result. Earlier source-aligned cohorts produced different tokens; the latest archived-cohort match does not explain those differences or admit broader concurrency/long context. The switches remain opt-in and this PR remains Draft.

A benchmark-only sidecar option aligns legacy/shared QPN2 native implementations. The older installed core dispatched M9–32 to TurboMind despite newer Python route logs; that validation gap is separately documented. Failed lane mapping, old Flash-V100 extension, harness import/inference-mode errors, interrupted services and excluded fixed2GiB/8K/E5M2 diagnostics remain recorded.

docs/design/sm70_dflash2_shared_nvfp4.md and the migration ledger retain contracts, reproduction, storage accounting, measured deltas and rejected variants. Task artifacts contain commands, manifests, source/diff and binary hashes, allocation stacks, operator samples and endpoint responses. Test services have exited and released their locks.

Latest speed and quality follow-up

Inference source remains 2d683cc; the latest commit adds evidence only. The two completed candidate final answers pass MBPP base and EvalPlus hidden tests (2/2 each, dataset hash ee43ecabebf20deef4bb776a405ac5b1).

Two current same-source A/B attempts were interrupted before candidate execution, after six speed requests and only two/three quality requests. The latest parent received SIGINT and its cleanup sent SIGTERM to the child; there was no logged OOM or CUDA computation failure, and the sender remains unidentified. The planned 32-task and four-concurrent checks have no passing result.

Controls retaining old scales and head layouts show restart variation: MBPP28 returns 270 versus 634 tokens, first differing at zero-based token 8, while each process repeats its own output identically six times. Different-output decode medians 233.57/253.62 tokens/s are not a memory-optimization speed comparison. No new overall speed or broad quality parity is claimed. Preserve the interrupted cohorts, resolve uninterrupted GPU ownership and control reproducibility, and keep this PR Draft. No owned GPU/API or scorer processes remain.