From 7faa7f682ea4b23921241ae01c0179dcc3fde058 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:46:56 +0800 Subject: [PATCH 1/9] [Bugfix] Guard DFlash2 sampling cutoffs and repair operator audit oracles Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../benchmark_sm70_quasar_nvfp4_oracle.py | 24 +- .../sm70_quasar_dflash2_operator_audit.md | 244 ++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 12 + .../test_sm70_quasar_nvfp4_oracle.py | 42 +++ .../kernels/core/test_sm70_qwen_layer_dump.py | 43 +++ .../v1/sample/test_topk_topp_tied_cutoffs.py | 62 +++++ tests/v1/spec_decode/test_dflash2.py | 4 +- .../layers/vocab_parallel_embedding.py | 12 +- vllm/model_executor/models/qwen3_next.py | 5 +- vllm/v1/sample/ops/topk_topp_triton.py | 35 +++ .../spec_decode/dflash2/sparse_rejection.py | 44 +++- 11 files changed, 516 insertions(+), 11 deletions(-) create mode 100644 docs/design/sm70_quasar_dflash2_operator_audit.md create mode 100644 tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py create mode 100644 tests/kernels/core/test_sm70_qwen_layer_dump.py create mode 100644 tests/v1/sample/test_topk_topp_tied_cutoffs.py diff --git a/benchmarks/kernels/benchmark_sm70_quasar_nvfp4_oracle.py b/benchmarks/kernels/benchmark_sm70_quasar_nvfp4_oracle.py index 9f39faf38f..b4513f805d 100644 --- a/benchmarks/kernels/benchmark_sm70_quasar_nvfp4_oracle.py +++ b/benchmarks/kernels/benchmark_sm70_quasar_nvfp4_oracle.py @@ -74,10 +74,26 @@ def _load_column_parallel( scales = checkpoint.tensor(scale_key) if packed.shape[0] % tp_size: raise ValueError(f"{prefix}: output rows do not divide TP{tp_size}") - rows = packed.shape[0] // tp_size - start = tp_rank * rows - packed_parts.append(packed[start : start + rows].contiguous()) - scale_parts.append(scales[start : start + rows].contiguous()) + output_sizes = (packed.shape[0],) + if prefix.endswith(".in_proj_qkv"): + config = json.loads((checkpoint.model / "config.json").read_text()) + config = config.get("text_config", config) + key_dim = config["linear_num_key_heads"] * config["linear_key_head_dim"] + value_dim = ( + config["linear_num_value_heads"] * config["linear_value_head_dim"] + ) + output_sizes = (key_dim, key_dim, value_dim) + if sum(output_sizes) != packed.shape[0]: + raise ValueError(f"{prefix}: Q/K/V sizes do not match checkpoint") + offset = 0 + for output_size in output_sizes: + if output_size % tp_size: + raise ValueError(f"{prefix}: logical projection does not divide TP") + rows = output_size // tp_size + start = offset + tp_rank * rows + packed_parts.append(packed[start : start + rows].contiguous()) + scale_parts.append(scales[start : start + rows].contiguous()) + offset += output_size weight_divisors.append(float(checkpoint.tensor(weight_key).flatten()[0])) input_divisors.append(float(checkpoint.tensor(input_key).flatten()[0])) if len(set(weight_divisors)) != 1 or len(set(input_divisors)) != 1: diff --git a/docs/design/sm70_quasar_dflash2_operator_audit.md b/docs/design/sm70_quasar_dflash2_operator_audit.md new file mode 100644 index 0000000000..ea1e94d33b --- /dev/null +++ b/docs/design/sm70_quasar_dflash2_operator_audit.md @@ -0,0 +1,244 @@ +# QUASAR + DFlash2 operator quality audit, 2026-09-06 + +This audit fixes two sampling-boundary defects and two defects in the +measurement path. It does **not** establish that the checkpoint, FP8 KV +cache, or approximate LM-head candidate search preserves unquantized model +quality. Keep the change in Draft: reference fallback has a measured latency +cost, and one of three selected MBPP Plus cases regressed at the fixed seed. + +## Frozen contract + +- Integration: `onecat/main`, base + `755baae1d075ee04fa9096b23fc0225b23589a86`. +- Target: `QUASAR-QAT/Qwen3.8-27B-QUASAR-NVFP4`, revision + `d8e6fbfa3e3a78899b440222b827430045a05b44`. +- Draft revision: `dedf8df68adfb1afeaf7b7480c0a0243108177b4`; + checkpoint SHA256: + `67fc76d68dc5a9415511a4f394ef744d67510cd20e93b37cc2cc7d28e4bab65c`. +- Four V100-SXM2-32GB GPUs, TP4, FP16 activations, NVFP4 checkpoint weights + executed through TurboMind W4A16/QPN2; this is not native NVFP4 arithmetic. + Python 3.12.13, Torch 2.10.0+cu128, CUDA 12.8, Triton 3.6.0. +- V2 runner, `FLASH_ATTN_V100` for both models, target E5M2 KV, draft FP16 + KV, seven probabilistic speculative tokens, QPN8 top-64 candidate search + followed by FP16 rerank with dense-vocabulary ordering. +- Maximum length 262144, batch-token budget 4096, maximum sequences 4, + GPU memory fraction 0.8, prefix caching enabled, Mamba cache mode + `align`, FULL_AND_PIECEWISE graphs. The audited verifier has eight rows. +- Sampling: temperature 1, top-k 20, top-p 0.95, natural EOS. Speed requests + use seed 20260925 and maximum output 1024. The three diagnostic MBPP cases + use seed 0 and maximum output 16384. Thinking is enabled, effort xhigh. + +Worker logs confirm Flash-V100 grouped q8 verification, XQA paged decode, +FlashQLA, TurboMind QPN2 and QPN8 rerank dispatch. Launch scripts and queue +manifests retain the complete environment. Speed runs contain no capture +hooks. Native binaries were reused; this change needs no native rebuild. +An attempted final live-worker library snapshot occurred after shutdown and +captured zero workers; it is not evidence of loaded-library identity. +Resolved native-file hashes and the original baseline loaded-library +manifest are retained separately. + +## Reference and evidence rules + +The C2 capture contains 64 target layers, all four TP ranks, and verification +steps 0, 7 and 31 of one real request: 1131 saved tensors per rank and step. +References use the **same saved input, logical TP weight shard, position, +physical cache pages, gate, and state selection**. They do not compare +unrelated generations. Linear references dequantize the same checkpoint +into FP32; they do not measure the loss from the original unquantized model. + +Relative L2 is `||actual - reference||2 / ||reference||2`. Tables show the +maximum across cases, not an average and not a single shared worst case. +Residual and transport checks use their staged dtype contracts; matrix, +attention and recurrence references otherwise use FP32 arithmetic. +Distribution error uses total variation `sum(abs(p-q))/2`. + +Small nonzero errors are measurements, not automatic acceptance. In +particular, a representation or reduction-order difference requires its own +propagation evidence before promotion. + +## Confirmed defects and changes + +1. **Compact top-20 loses the full-vocabulary tie contract.** PyTorch top-k + masking retains all entries equal to the kth threshold, potentially more + than 20. A fixed list of 20 cannot represent that distribution. A top-p + boundary can also split equal logits in a different vocabulary order. + Request a 21st reranked candidate, detect kth ties, split nucleus ties and + near-boundary CDF rounding, then use existing full-vocabulary verification + for an ambiguous block. Keep separate contiguous buffers for top-16, + top-20 and top-21. The check is outside model graph capture. +2. **Dense Triton masking mishandles ambiguous pivots.** On a compiled V100 + fixture with two high logits and 18 tied logits, top-p 0.95 retained 20 + candidates where PyTorch retained 19; distribution TV was 0.04266838. + The kernel now flags ambiguous rows and the public wrapper remasks those + rows using the untouched full-vocabulary PyTorch reference. Direct graph + capture uses that reference without a host decision inside capture. + The internal pivot algorithm is not claimed to have been replaced. +3. **Layer-dump custom-op output aliased its input against its schema.** + `torch.library.opcheck(test_schema)` rejected the old implementation. + Its AOT buffer reuse produced false residual errors at layers 6, 34 and + 62 in the first capture. Return owned storage, including the fake + implementation. This only repairs enabled diagnostics; it is not evidence + that ordinary serving previously had those residual errors. +4. **GDN QKV oracle sharded the concatenation incorrectly.** The checkpoint + concatenates logical Q/K/V. Shard each logical segment before + concatenation, matching the runtime loader, and test all four ranks. + +The compact guard adds one host decision and copies only eight by 21 values. +An exact-shape microbenchmark reduced guard median time from 255.85 us for +multiple GPU operations plus a host flag to 62.33 us for one copy and CPU +calculation, with zero decision disagreements on 300 random checks. This +does not remove the cost of full-vocabulary fallback. + +## Target operators on real inputs + +| Operator | Cases | Maximum absolute error | Maximum relative L2 | +|---|---:|---:|---:| +| Attention QKV projection | 192 | 0.02303 | 5.482e-4 | +| Attention output projection | 192 | 0.02396 | 5.854e-4 | +| GDN QKV/Z/B/A projection | 576 | 0.02952 | 5.802e-4 | +| GDN output projection | 576 | 0.01605 | 7.005e-4 | +| MLP gate/up with fused SiLU | 768 | 0.09627 | 1.283e-3 | +| MLP down projection | 768 | 0.03246 | 5.977e-4 | +| Input RMSNorm, staged reference | 768 | 0.001953 | 1.184e-5 | +| Post-attention RMSNorm, staged reference | 768 | 0.001953 | 1.333e-5 | +| Residual addition, staged reference | 768 | 0 | 0 | +| Q normalization and partial RoPE | 192 | 0.008986 | 2.458e-4 | +| K normalization and partial RoPE | 192 | 0.007598 | 3.032e-4 | +| Attention on identical cached KV | 192 | 0.03008 | 3.042e-4 | +| Attention sigmoid gate | 192 | 0.0002441 | 2.130e-4 | +| Attention gated multiply | 192 | 0.007690 | 2.461e-4 | +| KV key encoder vs E5M2 roundtrip | 192 | 0 | 0 | +| KV value encoder vs E5M2 roundtrip | 192 | 0 | 0 | +| E5M2 key representation vs original FP16 K | 192 | 0.9922 | 5.840e-2 | +| E5M2 value representation vs original FP16 V | 192 | 3.969 | 5.932e-2 | +| FP16 LM head vs FP32 same-weight arithmetic | 12 shards | 0.007971 | 2.085e-4 | +| TP reduction vs FP32 sum | 384 | 0.03125 | 2.736e-4 | + +All 7116 target reference cases are finite. All 384 TP reductions exactly +equal FP32 summation followed by FP16 rounding, and agree across ranks. + +The KV encoder is exact under its E5M2 contract. The roughly 5.9% conversion +loss belongs to that precision choice. A synthetic attention comparison +against original FP16 KV gives 4.34–5.34% relative output error, substantially +larger than attention arithmetic on the same cached KV. This does not prove +that KV precision caused the user's observed text-quality regression. + +The LM-head example illustrates why scalar tolerances and dataset scores +are insufficient: all 24 real rows preserve top-1, but one row changes its +top-p support against FP32 arithmetic. Maximum full-softmax TV is 0.002119; +after top-20/top-p it is **0.021532**. Separately, indexed FP16 rerank differs +from dense FP16 logits by 0.015625 on four of those rows, while all captured +top-20 ID sets agree. Neither FP32 references nor the tie guard establish +global coverage of the approximate QPN8 top-64 candidate search. + +## State, boundary and draft operators + +| Operator or boundary | Evidence | Result | +|---|---|---| +| GDN QKV packing | strided M=1/8/137 | Bitwise equal | +| GDN causal convolution + SiLU | all 48 layer weights, TP0 real q8 input, synthetic history; accepted selectors 1/8 | 96 cases; max relative L2 3.384e-4; whole updated state exact | +| GDN gated RMS | 48 real layer weights and real Z; synthetic core input | max relative L2 2.224e-4 | +| FlashQLA prefill recurrence | T=8/64/256/1024, FP32 state, independent sequential reference | output max abs 7.624e-6; state 1.192e-7 | +| FlashQLA fused decode | two seeded inputs, untouched-slot check | output max abs 1.763e-6; state 2.235e-8 | +| Packed q8 GDN verification | 64 graph replays, permuted state slots, accepted selectors 1–8; same precomputed gates and production matching flags | Output and whole state bitwise equal to split recurrent reference on every replay | +| Target grouped attention | q8, Hq6/Hkv1/D256, randomized physical pages; contexts 1032/32768/131072/262144 | same-cache FP32 relative L2 3.18–3.34e-4; XQA comparison is not bitwise equal | +| Draft noncausal SWA | q8, Hq8/Hkv2/D128, window 2047 each side, contexts 1024/2048/2055/4097/32768 | max relative L2 about 2.88e-4 | +| Draft FP16 projections | 120 real-input layer/rank cases | max relative L2 3.713e-4; max abs 1.849 on large intermediate values | +| Draft BF16 SwiGLU transport | 20 real-input cases | Captured input, rerun kernel and staged reference exact; row scales exact | +| Draft grouped convolution | 20 real-weight/input cases, both sides | max relative L2 6.11e-4 | +| Draft BF16 RMS/residual | 10 real-weight cases, two magnitudes | Exact, including residual magnitudes beyond FP16 finite range | +| Draft per-layer context K RMS | all five layer weights | max relative L2 2.065e-4 | +| Draft selector edge scores | real codebooks, candidate fixture, synthetic hidden rows | max relative L2 1.693e-6 | + +The 1024 initial real-weight M8 projection probes had exact repeated eager +and graph results. A separate 32-case dynamic-input replay check is exact. +These are not 1024 real-input or production-aligned GDN measurements: +the first GDN oracle had the sharding defect described above. + +Real-logit sampling replay covers 24 rows at top-p 1/0.95/0.6, isolating +sampling from LM-head arithmetic. Protected masks match the reference in +all 72 row cases and measured distribution TV is zero. The kth tie in this +capture carried negligible mass, so the large distribution defect is +established by the compiled adversarial fixtures, not overstated as a +large observed error on these 24 real rows. + +This is coverage of the principal q8 verifier and draft operator families, +not exhaustive acceptance of every prefill shape, embedding/context-FC +path, scheduler transition, sampling penalty, arbitrary RNG trajectory, +or checkpoint quantization loss. Those unmeasured contracts remain open. + +## Validation and cost + +Focused GPU validation: + +- Oracle and layer-dump tests: 6 passed. +- New tied-cutoff tests plus existing top-k/top-p suite: + 112 passed, 35 platform/configuration skips. +- Existing DFlash2 rerank/compact-top-k regression selection: 22 passed, + 144 deselected. +- After reducing guard overhead: 14 tied-cutoff tests passed again. +- Focused Ruff checks/format and `git diff --check` pass. + +Example commands from the owned checkout: + +```bash +python -m pytest -q tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py tests/kernels/core/test_sm70_qwen_layer_dump.py +python -m pytest -q tests/v1/sample/test_topk_topp_tied_cutoffs.py tests/v1/sample/test_topk_topp_sampler.py +python -m pytest -q tests/v1/spec_decode/test_dflash2.py -k 'rerank or compact_topk' +``` + +Unprofiled, one warmup then three measured repetitions; the following uses +the median consistently. Complete-round cost is engine pure-decode time +divided by speculative rounds, **not target-verifier GPU time**. + +| Request | Baseline round ms | Final round ms | Baseline / final decode tok/s | Baseline / final TTFT s | Emitted tokens / rounds | +|---|---:|---:|---:|---:|---:| +| Release 1K, actual prompt 1019 | 18.037 | 20.781 | 152.826 / 132.646 | 0.3522 / 0.3464 | 318 / 115 | +| MBPP 28, actual prompt 135 | 17.588 | 19.941 | 268.541 / 236.857 | 0.1051 / 0.1004 | 308 / 65 | + +Median round cost increases about 15.2% and 13.4%. The initial GPU-operation +guard was slower still: 21.037/20.068 ms. Request token counts and acceptance +counts are unchanged on these two speed fixtures. No end-to-end speedup or +restoration of the 17.6 ms release baseline is claimed. + +Three selected diagnostic MBPP cases all stop naturally: Base 3/3, Plus +0/3. The old DFlash baseline scored Plus 1/3 on these same cases; target-only +scored Plus 0/3. A real JSON-object request returns valid JSON with result +42. This small smoke cannot justify a model-quality improvement. + +## Retained artifacts and rejected evidence + +Artifact bundle: `v100-quasar-dflash2-operator-audit-20260906`. +The local Chinese audit report records its absolute location and the owned +worktree. Raw tensors, model paths and compiler caches are not checked in. + +- `captures-v2/`: C2 tensors, norms, route inventory, cache pages, full and + compact logits, per-shard FP32 head references. +- `results/real-v2-rank{0,1,2,3}.json`: per-layer/operator metrics; + `operator-summary.json` provides the aggregate. +- `math-scan.json`, `draft-math.json`, `gdn-q8-state.json`, + `gdn-conv-norm.json`, `tp-reduction.json`, + `lm-head-propagation.json`: independent operator and state evidence. +- `sampler-before.json`, `sampler-after.json`, + `sampler-real-final.json`: pre-fix and protected distribution checks. + The after file intentionally still reports the unguarded compact helper, + alongside the repaired dense public operator. +- `final-speed-*.json`, `fixed-quality-subset.json`, + `fixed-evalplus-subset.json`, `fixed-json.json`: service evidence. +- `scripts/`, `queue/`, `logs/`, `contract.json`, + `final-native-files.json`, `cleanup-final.json`: exact commands, + environments, negative attempts, binary hashes, and cleanup. + +Do not reuse C1 numerical conclusions: it was affected by the diagnostic +alias bug. Do not reuse the original FP16 multiply/divide restore probe as +graph corruption evidence: that transform was not lossless. The initial +uniform-logit masking test failed and was fixed; the first convolution/norm +probe failed because it instantiated a CustomOp without a vLLM config and +was corrected before producing the final 144 cases. An occupied test port +was avoided without interrupting its owner. + +Task-owned services and GPU leases have been released; the unrelated +service on port 8000 remained healthy. Before promotion, resolve the +fallback performance cost and investigate precision propagation at the +LM head and KV boundary with matched inputs. Do not replace those checks +with a larger aggregate-score run. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 8f8a52fed2..c0d6743f7f 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -2,6 +2,18 @@ Date: 2026-05-30 +## QUASAR + DFlash2 operator audit, 2026-09-06 + +See [the operator audit](sm70_quasar_dflash2_operator_audit.md) for the frozen +`755baae1d075ee04fa9096b23fc0225b23589a86` baseline, per-operator error +tables, sampling-boundary fixes, and rejected diagnostic evidence. C1 +captures and the original concatenated GDN QKV oracle are invalid for +production numerical conclusions. C2 residuals and q8 state updates are +exact under their staged contracts. E5M2 conversion and LM-head rounding +remain distinct precision concerns. Keep this change in Draft: final +complete-round costs are 20.781/19.941 ms versus 18.037/17.588 ms, and the +three-case Plus smoke is 0/3 versus the old DFlash 1/3. + ## Objective Target tree: diff --git a/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py b/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py new file mode 100644 index 0000000000..09faba70f3 --- /dev/null +++ b/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from types import SimpleNamespace + +import pytest +import torch + +from benchmarks.kernels.benchmark_sm70_quasar_nvfp4_oracle import ( + _load_column_parallel, +) + + +@pytest.mark.parametrize("rank", range(4)) +def test_gdn_qkv_shards_each_logical_projection(tmp_path, rank): + (tmp_path / "config.json").write_text( + json.dumps( + { + "text_config": { + "linear_num_key_heads": 4, + "linear_key_head_dim": 1, + "linear_num_value_heads": 12, + "linear_value_head_dim": 1, + } + } + ) + ) + prefix = "model.language_model.layers.0.linear_attn.in_proj_qkv" + packed = torch.arange(20, dtype=torch.uint8)[:, None].expand(-1, 2) + scales = torch.arange(20, dtype=torch.float32)[:, None] + tensors = { + prefix + ".weight_packed": packed, + prefix + ".weight_scale": scales, + prefix + ".weight_global_scale": torch.tensor(2.0), + prefix + ".input_global_scale": torch.tensor(3.0), + } + checkpoint = SimpleNamespace(model=tmp_path, tensor=tensors.__getitem__) + projection = _load_column_parallel(checkpoint, "gdn_qkv", (prefix,), rank, 4) + expected_rows = torch.tensor([rank, 4 + rank, *range(8 + 3 * rank, 11 + 3 * rank)]) + assert torch.equal(projection.packed, packed[expected_rows]) + assert torch.equal(projection.scales, scales[expected_rows]) diff --git a/tests/kernels/core/test_sm70_qwen_layer_dump.py b/tests/kernels/core/test_sm70_qwen_layer_dump.py new file mode 100644 index 0000000000..61542f7dfe --- /dev/null +++ b/tests/kernels/core/test_sm70_qwen_layer_dump.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_layer_dump_obeys_nonaliasing_schema(): + from vllm.model_executor.models import qwen3_next # noqa: F401 + + op = torch.ops.vllm.sm70_qwen_layer_dump.default + x = torch.randn(8, 32, device="cuda", dtype=torch.float16) + torch.library.opcheck( + op, + (x, "schema_test", 0, "linear_attention"), + test_utils=("test_schema", "test_faketensor"), + ) + out = op(x, "schema_test", 0, "linear_attention") + assert torch.equal(out, x) + out.zero_() + assert torch.count_nonzero(x) > 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qwen_layer_dump_preserves_compiled_residual(): + from vllm.model_executor.models import qwen3_next # noqa: F401 + + def forward(x, residual): + saved = x.float() + residual + value = torch.ops.vllm.sm70_qwen_layer_dump( + x, "compiled_test", 0, "linear_attention" + ) + updated = value.float() + residual + normalized = updated * torch.rsqrt(updated.square().mean(-1, True) + 1e-6) + return saved, updated + normalized + + x = torch.randn(8, 32, device="cuda", dtype=torch.float16) + residual = torch.randn_like(x, dtype=torch.float32) + expected = forward(x, residual) + actual = torch.compile(forward, backend="aot_eager", fullgraph=True)(x, residual) + for a, e in zip(actual, expected): + torch.testing.assert_close(a, e, atol=0, rtol=0) diff --git a/tests/v1/sample/test_topk_topp_tied_cutoffs.py b/tests/v1/sample/test_topk_topp_tied_cutoffs.py new file mode 100644 index 0000000000..0cbe22e05e --- /dev/null +++ b/tests/v1/sample/test_topk_topp_tied_cutoffs.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch +from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton +from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( + _compact_target_requires_reference, +) + + +@pytest.mark.parametrize("top_p", [1.0, 0.95, 0.6]) +@pytest.mark.parametrize("case", ["k_tie", "p_tie", "uniform", "unique"]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_tied_cutoffs_match_full_vocabulary_reference(case, top_p): + x = torch.full((8, 32768), -20.0, device="cuda") + if case == "k_tie": + x[:, :24] = 1.0 + x[:, :18] = 2.0 + elif case == "p_tie": + x[:, :20] = 1.0 + x[:, :2] = 2.0 + elif case == "uniform": + x.fill_(1.0) + else: + x[:, :32] = torch.arange(32, 0, -1, device="cuda") / 8 + k = torch.full((8,), 20, dtype=torch.int32, device="cuda") + p = torch.full((8,), top_p, device="cuda") + expected = apply_top_k_top_p_pytorch(x.clone(), k, p) + actual = apply_top_k_top_p_triton(x.clone(), k, p) + assert torch.equal(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_standalone_topp_ties_and_graph_capture(): + x = torch.full((2, 32768), -20.0, device="cuda") + x[:, :20] = 1.0 + x[:, :2] = 2.0 + p = torch.full((2,), 0.95, device="cuda") + expected = apply_top_k_top_p_pytorch(x.clone(), None, p) + actual = apply_top_k_top_p_triton(x.clone(), None, p) + assert torch.equal(actual, expected) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = apply_top_k_top_p_triton(x.clone(), None, p) + graph.replay() + assert torch.equal(captured, expected) + + +def test_compact_guard_only_rejects_ambiguous_cutoffs(): + unique = torch.arange(21, 0, -1, dtype=torch.float32)[None] / 8 + assert not _compact_target_requires_reference(unique, 1.0, 0.95) + k_tie = unique.clone() + k_tie[:, -1] = k_tie[:, -2] + assert _compact_target_requires_reference(k_tie, 1.0, 1.0) + p_tie = torch.ones(1, 21) + p_tie[:, :2] = 2.0 + p_tie[:, -1] = -20.0 + assert _compact_target_requires_reference(p_tie, 1.0, 0.95) + assert not _compact_target_requires_reference(p_tie, 1.0, 1.0) diff --git a/tests/v1/spec_decode/test_dflash2.py b/tests/v1/spec_decode/test_dflash2.py index 1881552e36..3d7b3ffc2b 100644 --- a/tests/v1/spec_decode/test_dflash2.py +++ b/tests/v1/spec_decode/test_dflash2.py @@ -1617,11 +1617,11 @@ def test_lm_head_candidate_interface_falls_back_when_rerank_is_disabled(monkeypa ) -@pytest.mark.parametrize("selector_k", [16, 20]) +@pytest.mark.parametrize("selector_k", [16, 20, 21]) @pytest.mark.parametrize("num_rows", [1, 7, 8]) def test_qpn8_rerank_output_buffers_are_contiguous(selector_k, num_rows): layer = SimpleNamespace() - for top_k in (16, 20): + for top_k in (16, 20, 21): setattr( layer, f"_sm70_dflash2_rerank_values_{top_k}", diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index b8a0819795..83b99abfa5 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -266,7 +266,7 @@ def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: # [max_rows, 20] allocation for top-16 leaves a row stride of 20 and makes # the result non-contiguous. The TP all-gather requires contiguous inputs, # and inserting a runtime contiguous() copy would add work to both graphs. - for selector_k in (16, 20): + for selector_k in (16, 20, 21): layer.register_buffer( f"_sm70_dflash2_rerank_values_{selector_k}", torch.empty((max_rows, selector_k), dtype=torch.float16, device=device), @@ -299,7 +299,7 @@ def _sm70_dflash2_rerank_output_buffers( num_rows: int, selector_k: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Select graph-stable, contiguous rerank outputs for top-16 or top-20.""" + """Select contiguous rerank outputs, including the target tie sentinel.""" if selector_k == 16: values = layer._sm70_dflash2_rerank_values_16[:num_rows] positions = layer._sm70_dflash2_rerank_positions_16[:num_rows] @@ -308,6 +308,10 @@ def _sm70_dflash2_rerank_output_buffers( values = layer._sm70_dflash2_rerank_values_20[:num_rows] positions = layer._sm70_dflash2_rerank_positions_20[:num_rows] ids = layer._sm70_dflash2_rerank_ids_20[:num_rows] + elif selector_k == 21: + values = layer._sm70_dflash2_rerank_values_21[:num_rows] + positions = layer._sm70_dflash2_rerank_positions_21[:num_rows] + ids = layer._sm70_dflash2_rerank_ids_21[:num_rows] else: raise ValueError(f"Unsupported DFlash2 rerank top-k: {selector_k}") return values, positions, ids @@ -522,7 +526,9 @@ def _maybe_sm70_dflash2_qpn8_rerank( return None if not getattr(layer, "_sm70_dflash2_qpn8_rerank_prepared", False): return None - if selector_k not in (16, 20) or bias is not None: + if selector_k not in (16, 20, 21) or bias is not None: + return None + if selector_k == 21 and not _sm70_dflash2_use_dense_order(): return None if x.dtype != torch.float16 or not x.is_cuda: return None diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b54454adad..b32171fcad 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -151,6 +151,9 @@ def _sm70_qwen_layer_dump_impl( layer_idx: int, layer_type: str, ) -> torch.Tensor: + # The custom-op schema has no output alias annotation. Returning the input + # can corrupt AOT buffer reuse and change the tensors being diagnosed. + tensor = tensor.clone() dump_dir = os.getenv("VLLM_SM70_DUMP_QWEN_LAYER_DIR") graph_buffers = os.getenv("VLLM_SM70_DUMP_QWEN_LAYER_GRAPH_BUFFERS") == "1" target_labels = { @@ -285,7 +288,7 @@ def _sm70_qwen_layer_dump_fake( layer_idx: int, layer_type: str, ) -> torch.Tensor: - return tensor + return torch.empty_like(tensor) direct_register_custom_op( diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 0766b70ce2..3196658c54 100644 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -148,6 +148,7 @@ def _topk_topp_kernel( BLOCK_SIZE_TRUNC: tl.constexpr, TOPK_ENABLED: tl.constexpr, TOPP_ENABLED: tl.constexpr, + REFERENCE_ROWS=None, ): NUM_TILES: tl.constexpr = (VOCAB_SIZE + BLOCK_SIZE - 1) // BLOCK_SIZE pid = tl.program_id(0) @@ -161,6 +162,7 @@ def _topk_topp_kernel( num_duplicate_logit = tl.zeros((), dtype=tl.uint32) num_keep = tl.zeros((), dtype=tl.uint32) num_kept = tl.zeros((), dtype=tl.uint32) + needs_reference = tl.full((), False, tl.int1) max_logit = -float("inf") min_logit = float("inf") @@ -406,6 +408,7 @@ def _topk_topp_kernel( duplicate_logit = min_larger num_duplicate_logit = num_min_larger num_keep = num_duplicate_logit - (k_pivots_num - k) + needs_reference = num_keep < num_duplicate_logit num_kept = tl.zeros((), dtype=tl.uint32) # Top-k only path. If there are fewer finite values @@ -973,6 +976,14 @@ def _topk_topp_kernel( # Top-p only path final_pivot = tl.log(p_pivot * sum_exp_logits) + max_sample + # Pivot-space duplicate selection does not preserve the reference's + # token tie order, and log/exp reconstruction can miss the cutoff. + # Mark only these ambiguous rows for exact reference masking. + needs_reference = needs_reference | (num_keep < num_duplicate_logit) + needs_reference = needs_reference | ~(final_pivot < max_logit) + if REFERENCE_ROWS is not None: + tl.store(REFERENCE_ROWS + row_id, needs_reference) + # Sixth pass: Apply mask and store final output. # If the pivot >= max logit (or is NaN), no token would # survive the strict `>` keep_mask. Skip masking. @@ -1038,6 +1049,16 @@ def apply_top_k_top_p_triton( if batch_size == 0 or not (topk_enabled or topp_enabled): return logits + # Sampling normally runs outside the model graphs. Keep direct graph + # callers correct without introducing a device-to-host fence in capture. + if logits.is_cuda and torch.cuda.is_current_stream_capturing(): + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch + + reference_logits = apply_top_k_top_p_pytorch(logits, k, p) + if mask_value != float("-inf"): + reference_logits.masked_fill_(torch.isneginf(reference_logits), mask_value) + return reference_logits + # The Triton kernel supports arbitrary row strides, but it still assumes # the vocab dimension is laid out contiguously within each row. if logits.stride(1) != 1: @@ -1099,6 +1120,8 @@ def apply_top_k_top_p_triton( # rows are default-on; the measured MTP verifier rows remain opt-in. launch_kwargs["num_warps"] = 8 + original_logits = logits.clone() + reference_rows = torch.empty(batch_size, dtype=torch.bool, device=logits.device) _topk_topp_kernel[(NUM_PROGRAMS,)]( logits, logits.stride(0), @@ -1114,9 +1137,21 @@ def apply_top_k_top_p_triton( BLOCK_SIZE_TRUNC=block_size_trunc, TOPK_ENABLED=topk_enabled, TOPP_ENABLED=topp_enabled, + REFERENCE_ROWS=reference_rows, **launch_kwargs, ) + if reference_rows.any(): + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch + + reference_logits = apply_top_k_top_p_pytorch( + original_logits[reference_rows], + k[reference_rows] if k is not None else None, + p[reference_rows] if p is not None else None, + ) + if mask_value != float("-inf"): + reference_logits.masked_fill_(torch.isneginf(reference_logits), mask_value) + logits[reference_rows] = reference_logits return logits diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py index 01b42e3968..6fa1f1d51d 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py @@ -31,6 +31,35 @@ _SELECTOR_ALIGNMENT_STEP = 0 +def _compact_target_requires_reference( + probe_logits: torch.Tensor, + temperature: float, + top_p: float, +) -> bool: + """Keep ambiguous cutoffs on the full-vocabulary sampling contract. + + The 21st candidate detects a tie crossing top-20. Ties wholly inside the + retained nucleus are harmless; ties split by top-p need the reference's + vocabulary tie order. The small CDF guard also covers FP32 scan rounding. + This is called outside the model CUDA graphs, once per B1 verification. + """ + # The branch needs one host decision anyway. Copy the tiny B1 probe once + # instead of launching a chain of GPU reductions followed by the same fence. + probe = probe_logits.detach().cpu().float().numpy() + logits = probe[:, :_TARGET_TOP_K] / temperature + exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True)) + probs = exp_logits / exp_logits.sum(axis=-1, keepdims=True) + before = probs.cumsum(axis=-1) - probs + keep = before < top_p + cutoff_tie = probe[:, -2] == probe[:, -1] + nucleus_tie = ( + (logits[:, :-1] == logits[:, 1:]) & (keep[:, :-1] != keep[:, 1:]) + ).any(axis=-1) + near_cutoff = np.abs(before - top_p).min(axis=-1) <= (16 * np.finfo(np.float32).eps) + ambiguous = cutoff_tie | ((nucleus_tie | near_cutoff) & (top_p < 1.0)) + return bool(ambiguous.any()) + + def _parse_alignment_steps(raw_steps: str | None) -> set[int] | None: if not raw_steps: return None @@ -219,8 +248,21 @@ def try_dflash2_sparse_target_rejection( draft_topk_ids, draft_topk_logits = sparse_draft_logits target_topk_ids, target_topk_logits = model.get_topk_tokens_and_logits( sample_hidden_states, - _TARGET_TOP_K, + _TARGET_TOP_K + 1, ) + idx = input_batch.idx_mapping_np[0] + states = rejection_sampler.sampler.sampling_states + if _compact_target_requires_reference( + target_topk_logits, + float(states.temperature.np[idx]), + float(states.top_p.np[idx]), + ): + logger.info_once( + "DFlash2 target cutoff requires full-vocabulary reference sampling." + ) + return None + target_topk_ids = target_topk_ids[:, :_TARGET_TOP_K] + target_topk_logits = target_topk_logits[:, :_TARGET_TOP_K] draft_sampled = input_batch.input_ids[input_batch.logits_indices] pos = input_batch.positions[input_batch.logits_indices] sampled, num_sampled = dflash2_sparse_topk_rejection_sample( From 6ec27bec9c2e0aec24037597710045b3b8b25e5d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:20:17 +0800 Subject: [PATCH 2/9] [Kernel] Add E4M3 KV and FP32 logits for the DFlash2 quality path Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../sm70_quasar_dflash2_operator_audit.md | 4 + docs/design/sm70_quasar_e4m3_fp32_logits.md | 198 ++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 11 + .../flash_attn_v100/flash_attn_interface.py | 6 + .../kernel/flash_decode_paged.cu | 59 ++++-- flash-attention-v100/kernel/fused_mha_api.cpp | 1 + .../attention/test_sm70_grouped_e4m3.py | 82 ++++++++ tests/kernels/core/test_sm70_fp32_lm_head.py | 46 ++++ .../attention/test_sm70_flash_v100_policy.py | 23 +- vllm/envs.py | 6 + .../layers/sm70_fp32_lm_head.py | 52 +++++ .../layers/vocab_parallel_embedding.py | 68 ++++-- vllm/v1/attention/backends/flash_attn_v100.py | 24 ++- 13 files changed, 535 insertions(+), 45 deletions(-) create mode 100644 docs/design/sm70_quasar_e4m3_fp32_logits.md create mode 100644 tests/kernels/attention/test_sm70_grouped_e4m3.py create mode 100644 tests/kernels/core/test_sm70_fp32_lm_head.py create mode 100644 vllm/model_executor/layers/sm70_fp32_lm_head.py diff --git a/docs/design/sm70_quasar_dflash2_operator_audit.md b/docs/design/sm70_quasar_dflash2_operator_audit.md index ea1e94d33b..1c7fc6a9b7 100644 --- a/docs/design/sm70_quasar_dflash2_operator_audit.md +++ b/docs/design/sm70_quasar_dflash2_operator_audit.md @@ -1,5 +1,9 @@ # QUASAR + DFlash2 operator quality audit, 2026-09-06 +This is the retained first-stage audit. The +[E4M3 KV and FP32-logit follow-up](sm70_quasar_e4m3_fp32_logits.md) addresses +the precision findings and records subsequent performance and quality checks. + This audit fixes two sampling-boundary defects and two defects in the measurement path. It does **not** establish that the checkpoint, FP8 KV cache, or approximate LM-head candidate search preserves unquantized model diff --git a/docs/design/sm70_quasar_e4m3_fp32_logits.md b/docs/design/sm70_quasar_e4m3_fp32_logits.md new file mode 100644 index 0000000000..b6dd458d58 --- /dev/null +++ b/docs/design/sm70_quasar_e4m3_fp32_logits.md @@ -0,0 +1,198 @@ +# QUASAR + DFlash2: E4M3 KV and FP32 logits + +This is the precision follow-up to the +[operator audit](sm70_quasar_dflash2_operator_audit.md). It addresses the +observed 0.021532 top-p distribution error and adds E4M3 storage to the +grouped q8 verifier. The opt-in mode improves measured operator precision; +it is not a claim that the original unquantized model or aggregate answer +quality has been recovered. + +## Implementation + +`VLLM_SM70_DFLASH2_FP32_LOGITS=1` retains FP32 logits for the existing +SM70 TP4 DFlash2 LM-head contract: FP16 weights, local shape 62080 by 5120, +QPN8 top-64 candidate search, and up to eight verifier rows. A small Triton +kernel evaluates each selected original weight row with FP32 multiplication +and reduction, storing directly into FP32 candidate buffers. It avoids both +FP16 output rounding and expanded cross-row products in the old rerank. +Dense-vocabulary ordering and the audited sampling cutoff protection remain +enabled. Full-vocabulary fallback also returns FP32 logits through +`torch.mm(..., out_dtype=torch.float32)`; this is tested with Torch 2.10. +The flag defaults off, making the changed precision contract explicit. + +Explicit `--kv-cache-dtype fp8_e4m3` now uses the grouped one-pass q8 kernel +for supported Hq6/Hkv1/D256 layouts, instead of forcing independent XQA +rows. Reuse the existing exact E4M3 paired conversion used by XQA. Other +grouped sparse-page paths retain their conversion. The new entry requires +16-byte-aligned KV block/token strides; the backend gates that contract. +The E4M3 route is limited to q8/one-pass in model dispatch and the native +API. Other query lengths retain their existing fallback. + +The native module advertises E4M3 grouped support. Python paired with an +older native extension keeps the older routing behavior rather than +calling an unsupported entry. Rebuild Flash-V100 to use the new route; +no TurboMind native rebuild is needed. + +## Same-input numerical evidence + +References reuse the previous valid C2 captures: all four TP ranks, +verification steps 0/7/31, original checkpoint weights and hidden states. +They measure runtime arithmetic against that same quantized checkpoint, +not checkpoint quantization loss against an unavailable original model. + +| Measurement | Previous FP16/E5M2 | New FP32/E4M3 | +|---|---:|---:| +| Top-p TV on the previously problematic row | 0.02153214 | 3.6694e-7 | +| Maximum top-p TV across the 24 real rows | 0.02153214 | 1.1325e-6 | +| Top-p support matches FP32 reference | 23/24 rows | 24/24 rows | +| Candidate logit maximum absolute error, actual QPN8 support | FP16 rounding up to one logit ULP | 3.8147e-6 | +| Missing local top-21 in actual QPN8 top-64 support | separately audited | 0 across 12 rank/step cases | +| Maximum real KV conversion relative L2 | 0.0593241 | 0.0290236 | +| Native KV encoder vs matching Torch FP8 conversion | exact | exact | + +The problematic distribution difference is reduced by about five orders +of magnitude on that row, not mathematically zero. Dense FP32 fallback has +a different reduction order from the independent FP32 reference: maximum +absolute logit error 0.00051308 and maximum top-p TV 2.9933e-5, with matching +support on all 24 rows. Candidate-path results must not be substituted for +dense-fallback results in a precision claim. + +KV checks use 384 K/V cases per dtype across the 16 full-attention layers, +four ranks and three captured steps. No checked source value exceeds E4M3's +finite magnitude range, and no NaN/Inf was produced. Native cache write +and decode roundtrips match the corresponding dtype reference. This is +evidence for these inputs, not a proof that all future KV values fit an +unscaled E4M3 cache. The checkpoint supplies no KV scale tensors. + +Independent randomized-page attention checks cover 1032, 32768, 131072 and +262144 context lengths. At 262144, representation-induced relative output +error falls from about 0.04856 to 0.02556; arithmetic on the same decoded +cache stays around 3.3e-4. Graph and eager outputs agree. + +The faster E4M3 conversion is bitwise equal to the scalar conversion on all +four context lengths and a fixture exercising every finite E4M3 encoding, +including signed zeros and subnormals. No attention arithmetic or reduction +schedule is changed by that optimization. + +## Performance evidence and limits + +A real-hidden candidate microbenchmark on V100, M=8/K=5120/64 candidates, +measures about 100.76 us for packed FP16 rerank and 5.68 us for FP32 candidate +dots. This is a graph microbenchmark with repeated inputs and resident +candidate weights; it is not an end-to-end speedup claim. + +Same-input native E4M3 conversion A/B: + +| Context | Scalar grouped us | Paired grouped us | Output | +|---|---:|---:|---| +| 1032 | 45.57 | 36.25 | Bitwise equal | +| 32768 | 225.18 | 159.33 | Bitwise equal | +| 131072 | 824.52 | 562.28 | Bitwise equal | +| 262144 | 1617.41 | 1089.13 | Bitwise equal | + +These are complete grouped-attention kernel-path times, not isolated +conversion instructions. Absolute timings from different microbenchmark +executions are kept separate; use the same-input A/B for the paired +conversion delta. + +The retained final service result is recorded below after validation. +Complete-round cost is engine pure-decode time divided by speculative +rounds. It is not the target verifier alone. Changed precision can change +generated length and acceptance, so include them with throughput and TTFT. + +| Configuration | Release 1K round ms / decode tok/s | MBPP28 round ms / decode tok/s | Output tokens (1K / MBPP) | +|---|---:|---:|---:| +| Original E5M2/FP16 baseline | 18.037 / 152.826 | 17.588 / 268.541 | 318 / 308 | +| E5M2/FP32 control | 20.457 / 166.521 | 19.979 / 252.811 | 529 / 299 | +| E4M3/FP32 scalar conversion | 20.516 / 146.230 | 20.254 / 224.833 | 283 / 297 | +| E4M3/FP32 paired, final confirmation | 20.149 / 135.033 | 19.745 / 267.695 | 303 / 260 | + +Final confirmation TTFT medians are 0.3512/0.1033 s. Accepted/emitted +lengths are 2.7297/5.3061 tokens per round, with 111/49 rounds. An earlier +paired-service startup measured 20.099/19.623 ms; the confirmation measured +20.149/19.745 ms and reproduced the complete token hashes of both speed +fixtures. The original 17.6–18.0 ms round cost is **not fully recovered**. + +Scalar and paired service transcripts differ, even though all same-input +conversion probes agree bitwise. Added replay on 192 real-query cases with +C2 cached KV re-encoded to E4M3 also agrees bitwise. These checks do not prove +end-to-end trajectory equivalence, so the final paired service was separately +checked for natural-stop text and JSON quality instead of inheriting the +scalar service score. + +E5M2/FP32, scalar E4M3/FP32 and final paired E4M3/FP32 services completed the same three +diagnostic MBPP requests naturally. Base tests pass 3/3, Plus tests pass +0/3, and a JSON-object request returns valid JSON with result 42. These +selected hard cases do not establish an improvement in answer quality. +The original DFlash baseline had Plus 1/3 on those cases, and its target-only +control had 0/3. Keep this limitation visible. + +## Validation and reproduction + +Environment: four V100-SXM2-32GB GPUs, TP4, Python 3.12.13, Torch 2.10.0+cu128, +CUDA 12.8, Triton 3.6.0. Target checkpoint revision +`d8e6fbfa3e3a78899b440222b827430045a05b44`; draft revision +`dedf8df68adfb1afeaf7b7480c0a0243108177b4`. +Starting precision-audit commit: +`7faa7f682ea4b23921241ae01c0179dcc3fde058`, based on +`onecat/main` at `755baae1d075ee04fa9096b23fc0225b23589a86`. + +Focused checks: + +- FP32 candidate graph replay/rounding and E4M3 scaled random-page tests: + 7 passed. +- Existing grouped E5M2 regression tests: 25 passed. +- Routing and DFlash2 rerank/compact-top-k selection: + 31 passed, 257 deselected. +- After paired conversion: new E4M3 and existing grouped tests: + 29 passed. +- Final native shape/stride contract and E4M3 reference/graph checks: 6 passed. +- Mypy, Ruff and required commit hooks are checked on the final source. + +Example serving contract, using the pinned local checkpoints and this +checkout's rebuilt Flash-V100: + +```bash +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +VLLM_USE_V2_MODEL_RUNNER=1 \ +VLLM_SM70_DFLASH2_FP32_LOGITS=1 \ +vllm serve "$QUASAR_MODEL" \ + --dtype half --tensor-parallel-size 4 \ + --attention-backend FLASH_ATTN_V100 --kv-cache-dtype fp8_e4m3 \ + --max-model-len 262144 --max-num-batched-tokens 4096 --max-num-seqs 4 \ + --gpu-memory-utilization 0.8 --enable-prefix-caching --mamba-cache-mode align \ + --limit-mm-per-prompt '{"image":0,"video":0}' \ + --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_coder \ + --default-chat-template-kwargs '{"enable_thinking":true,"reasoning_effort":"xhigh"}' \ + --speculative-config "{\"method\":\"dflash\",\"model\":\"$DFLASH2_MODEL\",\"revision\":\"dedf8df68adfb1afeaf7b7480c0a0243108177b4\",\"num_speculative_tokens\":7,\"kv_cache_dtype\":\"auto\",\"attention_backend\":\"FLASH_ATTN_V100\",\"draft_sample_method\":\"probabilistic\",\"enforce_eager\":false}" \ + --seed 0 +``` + +Speed requests use temperature 1, top-k 20, top-p 0.95, seed 20260925, +maximum output 1024 and natural EOS, with one warmup and three measured +repetitions. Quality requests use seed 0 and maximum output 16384. + +Raw bundle: `v100-quasar-e4m3-fp32logits-20260906`. The local Chinese +report retains absolute locations. Important files include: + +- `results/head-pipeline-real.json`: actual QPN8 search, FP32 candidate and + dense-fallback comparisons. +- `results/kv-real.json`: native encoder and representation errors. +- `results/e4m3-boundary.json`, `results/e4m3-pair-compare.json`: long + boundaries and bitwise paired-conversion A/B. +- `results/*-speed-*.json`, `*-quality-subset.json`, `*-evalplus.json`: + unprofiled service results, text, tokens and scores. +- `scripts/`, `queue/`, `logs/`: full commands and negative attempts. +- Native build manifests and live worker manifests identify the loaded + extension, source hashes and task-owned cache locations. + +Rejected attempts are retained: the first head probe imported the wrong +Python operator facade; the first prepared-buffer replay lacked inference +mode; and a paired-conversion probe was queued before its library finished +building. These failed probes produced no accepted numerical evidence. +The initial scalar E4M3 route is retained as the slower correctness control. +No new model weights or generated binaries are committed. + +The final native host validation admits exactly q8. Its extracted CUDA +device fatbinary is SHA256-identical to the library used for final service +validation; only the host shape guard was tightened. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index c0d6743f7f..53c1911175 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -2,6 +2,17 @@ Date: 2026-05-30 +## QUASAR E4M3 KV and FP32 logits, 2026-09-06 + +The [precision follow-up](sm70_quasar_e4m3_fp32_logits.md) adds explicit +FP32 candidate/dense logits and E4M3 grouped q8 verification. The measured +problem-row distribution TV drops from 0.021532 to 3.6694e-7, all 24 checked +nucleus supports match the FP32 reference, and maximum real KV conversion +relative L2 drops from 0.059324 to 0.029024. Paired E4M3 conversion preserves +the scalar route bitwise across 1K–256K kernel probes. Keep the opt-in +precision contract and the answer-quality limitations visible; do not +equate these operator improvements with recovery of unquantized-model quality. + ## QUASAR + DFlash2 operator audit, 2026-09-06 See [the operator audit](sm70_quasar_dflash2_operator_audit.md) for the frozen diff --git a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py index 433cd1fee7..c617bbfd39 100644 --- a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py +++ b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py @@ -1090,6 +1090,12 @@ def flash_attn_grouped_verify_paged( ) +# Older native extensions retain the E5M2 route until rebuilt. +flash_attn_grouped_verify_paged.supports_e4m3 = bool( # type: ignore[attr-defined] + getattr(flash_attn_v100_cuda, "grouped_verify_e4m3", False) +) + + def flash_attn_decode_paged_xqa( q: torch.Tensor, k_cache: torch.Tensor, diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 7384538f75..a9a71de529 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -2125,7 +2125,10 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ : tile_start; load_xqa_tc_kv_panel( + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 || + (KV_DTYPE == + flash_v100::KV_CACHE_DTYPE_FP8_E4M3 && + !SPARSE_PAGE4)>( shared_kv, k_cache, page_ids, valid_k_rows, kPanelStrideVec, kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, k_block_stride, k_token_stride, k_head_stride, 0); @@ -2213,7 +2216,9 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ : tile_start; load_xqa_tc_kv_panel( + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 || + (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 && + !SPARSE_PAGE4)>( shared_kv, k_cache, page_ids, valid_k_rows, kPanelStrideVec, kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, k_block_stride, k_token_stride, k_head_stride, 0); @@ -2317,7 +2322,9 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ load_xqa_tc_kv_panel( + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E5M2 || + (KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 && + !SPARSE_PAGE4)>( shared_kv, v_cache, page_ids, valid_k_rows, kPanelStrideVec, kSharedStrideVec, tile_page_offset, 0, page_block_size, 0, v_block_stride, v_token_stride, v_head_stride, 0); @@ -4054,11 +4061,17 @@ at::Tensor flash_attention_grouped_verify_paged( TORCH_CHECK(partial_out.is_cuda() && partial_lse.is_cuda(), "grouped verify workspaces must be CUDA tensors"); TORCH_CHECK(q.dtype() == torch::kFloat16, "grouped verify q must be fp16"); - TORCH_CHECK(kv_cache_dtype == "fp8_e5m2", - "grouped verify prototype supports fp8_e5m2 KV only"); + TORCH_CHECK(kv_cache_dtype == "fp8_e5m2" || kv_cache_dtype == "fp8_e4m3", + "grouped verify requires E5M2 or E4M3 KV"); + TORCH_CHECK(kv_cache_dtype != "fp8_e4m3" || (q.size(0) == 8 && one_pass), + "E4M3 grouped verify requires q=8 and one_pass=true"); + TORCH_CHECK(kv_cache_dtype != "fp8_e4m3" || + (k_cache.stride(0) % 16 == 0 && k_cache.stride(1) % 16 == 0 && + v_cache.stride(0) % 16 == 0 && v_cache.stride(1) % 16 == 0), + "E4M3 grouped verify requires 16-byte-aligned KV strides"); TORCH_CHECK( k_cache.dtype() == torch::kUInt8 && v_cache.dtype() == torch::kUInt8, - "grouped verify E5M2 cache must use uint8 storage"); + "grouped verify FP8 cache must use uint8 storage"); TORCH_CHECK( block_table.dtype() == torch::kInt32 && seq_lens.dtype() == torch::kInt32, "grouped verify block_table/seq_lens must be int32"); @@ -4128,12 +4141,12 @@ at::Tensor flash_attention_grouped_verify_paged( const size_t partial_shared_mem = sizeof(GroupedVerifySmem); #define LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, \ SINGLE_QUERY, CONTIGUOUS_LAYOUT, \ - STAGE_PAGE_IDS) \ + STAGE_PAGE_IDS, KV_TYPE) \ do { \ auto partial_kernel = \ (void*)flash_attention_grouped_verify_e5m2_partial_kernel< \ MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, SINGLE_QUERY, \ - CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS>; \ + CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS, KV_TYPE>; \ const cudaError_t smem_status = cudaFuncSetAttribute( \ partial_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, \ partial_shared_mem); \ @@ -4147,7 +4160,7 @@ at::Tensor flash_attention_grouped_verify_paged( cudaGetErrorString(carveout_status)); \ flash_attention_grouped_verify_e5m2_partial_kernel< \ MAX_QUERY_TOKENS, TWO_PASS, PAGE_SIZE, SINGLE_QUERY, \ - CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS> \ + CONTIGUOUS_LAYOUT, STAGE_PAGE_IDS, KV_TYPE> \ <<>>( \ reinterpret_cast(q.data_ptr()), k_cache.data_ptr(), \ v_cache.data_ptr(), block_table.data_ptr(), \ @@ -4179,30 +4192,42 @@ at::Tensor flash_attention_grouped_verify_paged( fixed_interleaved_layout && dflash2_grouped_stage_page_ids_enabled(); \ if (stage_page_ids && page_size == 1648) { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ - SINGLE_QUERY, true, true); \ + SINGLE_QUERY, true, true, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } else if (stage_page_ids) { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ - SINGLE_QUERY, true, true); \ + SINGLE_QUERY, true, true, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } else if (fixed_interleaved_layout && page_size == 1648) { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ - SINGLE_QUERY, true, false); \ + SINGLE_QUERY, true, false, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } else if (fixed_interleaved_layout) { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ - SINGLE_QUERY, true, false); \ + SINGLE_QUERY, true, false, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } else if (page_size == 1648) { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 1648, \ - SINGLE_QUERY, false, false); \ + SINGLE_QUERY, false, false, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } else if (page_size == 3296) { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 3296, \ - SINGLE_QUERY, false, false); \ + SINGLE_QUERY, false, false, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } else { \ LAUNCH_GROUPED_VERIFY_PARTIAL(MAX_QUERY_TOKENS, TWO_PASS, 0, \ - SINGLE_QUERY, false, false); \ + SINGLE_QUERY, false, false, \ + flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ } \ } while (0) const bool single_query = q.size(0) == 1; - if (wide_query && one_pass) { + if (kv_cache_dtype == "fp8_e4m3") { + // Reuse the grouped q8 schedule and existing E4M3 vector conversion. + // Runtime strides cover both separate and interleaved hybrid KV pages. + LAUNCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ8MaxQ, false, 0, false, false, + false, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); + } else if (wide_query && one_pass) { DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ16MaxQ, false, false); } else if (wide_query) { DISPATCH_GROUPED_VERIFY_PARTIAL(kGroupedVerifyQ16MaxQ, true, false); diff --git a/flash-attention-v100/kernel/fused_mha_api.cpp b/flash-attention-v100/kernel/fused_mha_api.cpp index 76859871ac..598178bbd8 100644 --- a/flash-attention-v100/kernel/fused_mha_api.cpp +++ b/flash-attention-v100/kernel/fused_mha_api.cpp @@ -24,6 +24,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Staged FlashAttention XQA decode over paged KV cache (Volta)"); m.def("grouped_verify_paged_fwd", &flash_attention_grouped_verify_paged, "Exact grouped DFlash2 verification over paged KV cache (Volta)"); + m.attr("grouped_verify_e4m3") = true; m.def("grouped_verify_max_query_tokens", &flash_attention_grouped_verify_max_query_tokens, "Maximum query length supported by grouped DFlash2 verification"); diff --git a/tests/kernels/attention/test_sm70_grouped_e4m3.py b/tests/kernels/attention/test_sm70_grouped_e4m3.py new file mode 100644 index 0000000000..4c78ddc8ce --- /dev/null +++ b/tests/kernels/attention/test_sm70_grouped_e4m3.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("page_size", [16, 3296]) +@pytest.mark.parametrize("length", [1032, 32768]) +@torch.inference_mode() +def test_e4m3_grouped_scaled_random_pages_and_graph(page_size, length): + fa = pytest.importorskip("flash_attn_v100") + if not getattr(fa.flash_attn_grouped_verify_paged, "supports_e4m3", False): + pytest.skip("Rebuild Flash-V100 for E4M3 grouped verification") + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + torch.manual_seed(length + page_size) + pages = (length + page_size - 1) // page_size + 2 + raw = torch.randn(pages, 2, page_size, 1, 256, device="cuda").half() * 0.25 + cache = raw.to(torch.float8_e4m3fn).view(torch.uint8) + k, v = cache.unbind(1) + table = torch.randperm(pages, device="cuda", dtype=torch.int32)[None] + q = torch.randn(8, 6, 256, device="cuda").half() * 0.25 + seq = torch.tensor([length], device="cuda", dtype=torch.int32) + out = torch.empty_like(q) + + def run(): + return fa.flash_attn_grouped_verify_paged( + q, + k, + v, + table, + seq, + out=out, + kv_cache_dtype="fp8_e4m3", + k_scale=0.5, + v_scale=2.0, + one_pass=True, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for current in [length, length - 5]: + seq.fill_(current) + graph.replay() + actual = out.clone() + run() + assert torch.equal(actual, out) + keys = k[table[0].long()].view(torch.float8_e4m3fn) + values = v[table[0].long()].view(torch.float8_e4m3fn) + keys = keys.reshape(-1, 256)[:current].float() * 0.5 + values = values.reshape(-1, 256)[:current].float() * 2 + scores = q.transpose(0, 1).float() @ keys.t() / 16 + pos = torch.arange(current - 8, current, device="cuda") + mask = torch.arange(current, device="cuda")[None] > pos[:, None] + scores.masked_fill_(mask[None], -torch.inf) + expected = (scores.softmax(-1) @ values).transpose(0, 1) + relative = (actual.float() - expected).norm() / expected.norm() + assert relative < 5e-4 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("invalid", ["query_length", "stride"]) +def test_e4m3_grouped_rejects_unsupported_contract(invalid): + fa = pytest.importorskip("flash_attn_v100") + if not getattr(fa.flash_attn_grouped_verify_paged, "supports_e4m3", False): + pytest.skip("Rebuild Flash-V100 for E4M3 grouped verification") + query_len = 1 if invalid == "query_length" else 8 + width = 257 if invalid == "stride" else 256 + cache = torch.zeros(2, 2, 16, 1, width, device="cuda", dtype=torch.uint8) + k, v = cache[..., :256].unbind(1) + q = torch.zeros(query_len, 6, 256, device="cuda", dtype=torch.float16) + table = torch.zeros(1, 1, device="cuda", dtype=torch.int32) + seq = torch.tensor([8], device="cuda", dtype=torch.int32) + message = "requires q=8" if invalid == "query_length" else "aligned KV strides" + with pytest.raises(RuntimeError, match=message): + fa.flash_attn_grouped_verify_paged( + q, k, v, table, seq, kv_cache_dtype="fp8_e4m3", one_pass=True + ) diff --git a/tests/kernels/core/test_sm70_fp32_lm_head.py b/tests/kernels/core/test_sm70_fp32_lm_head.py new file mode 100644 index 0000000000..12b505957e --- /dev/null +++ b/tests/kernels/core/test_sm70_fp32_lm_head.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.model_executor.layers.sm70_fp32_lm_head import indexed_fp32_logits + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("rows", [1, 8]) +@torch.inference_mode() +def test_indexed_logits_preserve_fp32_and_graph_replay(rows): + torch.manual_seed(817) + x = torch.randn(rows, 5120, device="cuda", dtype=torch.float16) * 0.1 + weight = torch.randn(1024, 5120, device="cuda", dtype=torch.float16) * 0.1 + ids = torch.randint(1024, (rows, 64), device="cuda") + out = torch.empty(rows, 64, device="cuda", dtype=torch.float32) + indexed_fp32_logits(x, weight, ids, out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + indexed_fp32_logits(x, weight, ids, out) + for _ in range(3): + x.normal_(std=0.1) + ids.random_(1024) + graph.replay() + expected = (x.float() @ weight.float().t()).gather(1, ids) + torch.testing.assert_close(out, expected, atol=3e-6, rtol=3e-6) + assert out.dtype == torch.float32 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_indexed_logits_do_not_collapse_half_rounding_boundary(): + x = torch.zeros(1, 5120, device="cuda", dtype=torch.float16) + x[0, 0] = 1 + x[0, 1] = 2**-12 + weight = torch.zeros(64, 5120, device="cuda", dtype=torch.float16) + weight[:, 0] = 1 + weight[:, 1] = torch.arange(64, device="cuda") / 64 + ids = torch.arange(64, device="cuda")[None] + out = torch.empty(1, 64, device="cuda", dtype=torch.float32) + indexed_fp32_logits(x, weight, ids, out) + expected = x.float() @ weight.float().t() + assert torch.equal(out, expected) + assert torch.unique(out).numel() == 64 + assert torch.unique(out.half()).numel() == 1 diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 46e1f49fd2..a3bfe228ba 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -1723,8 +1723,9 @@ def fake_decode( @pytest.mark.parametrize("query_len", [8, 16]) @pytest.mark.parametrize("page_size", [3296, 3456]) +@pytest.mark.parametrize("kv_dtype", ["fp8_e5m2", "fp8_e4m3"]) def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata( - query_len: int, page_size: int + query_len: int, page_size: int, kv_dtype: str ): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl @@ -1735,7 +1736,7 @@ def test_flash_v100_dflash2_grouped_verify_uses_original_request_metadata( num_kv_heads=1, alibi_slopes=None, sliding_window=None, - kv_cache_dtype="fp8_e5m2", + kv_cache_dtype=kv_dtype, ) impl.use_dflash2_grouped_verify = True impl.dflash2_grouped_verify_max_query_tokens = 16 @@ -1787,6 +1788,24 @@ def grouped_verify( num_query_tokens=query_len, ) attn_metadata.max_model_len = 32768 + if kv_dtype == "fp8_e4m3": + assert not impl._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=query_len, + ) + grouped_verify.supports_e4m3 = True # type: ignore[attr-defined] + if query_len == 16: + assert not impl._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=query_len, + ) + return result = impl._flash_v100_small_query_prefill_as_decode( layer, query, diff --git a/vllm/envs.py b/vllm/envs.py index f686e9b0e0..b2c10b7d3f 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -224,6 +224,7 @@ VLLM_SM70_LM_HEAD_TOP1: bool = True VLLM_SM70_LM_HEAD_TOP1_TC: bool = False VLLM_SM70_DFLASH2_QPN8_RERANK: bool = False + VLLM_SM70_DFLASH2_FP32_LOGITS: bool = False VLLM_SM70_DFLASH2_QPN8_RERANK_SHADOW: bool = False VLLM_SM70_DFLASH2_QPN8_DENSE_ORDER: bool = True VLLM_SM70_DFLASH2_QPN8_ALLOW_CANDIDATE_ORDER: bool = False @@ -2108,6 +2109,11 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DFLASH2_QPN8_RERANK": lambda: bool( int(os.getenv("VLLM_SM70_DFLASH2_QPN8_RERANK", "0")) ), + # Explicit precision contract: retain FP32 candidate and dense logits + # for the SM70 TP4 DFlash2 LM head, including reference fallback. + "VLLM_SM70_DFLASH2_FP32_LOGITS": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FP32_LOGITS", "0")) + ), # Audit-only eager mode: execute QPN8+rerank, compare it with the dense # local top-k, and return the dense result so the baseline trajectory is # unchanged. This intentionally synchronizes for diagnostics. diff --git a/vllm/model_executor/layers/sm70_fp32_lm_head.py b/vllm/model_executor/layers/sm70_fp32_lm_head.py new file mode 100644 index 0000000000..98bae8ac63 --- /dev/null +++ b/vllm/model_executor/layers/sm70_fp32_lm_head.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Candidate LM-head dots without an intermediate FP16 logit rounding.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _indexed_fp32_logits( + X, + W, + IDS, + OUT, + K: tl.constexpr, + C: tl.constexpr, + BLOCK_K: tl.constexpr, +): + row = tl.program_id(0) + candidate = tl.program_id(1) + token = tl.load(IDS + row * C + candidate) + k = tl.arange(0, BLOCK_K) + x = tl.load(X + row * K + k, k < K, 0).to(tl.float32) + weight = tl.load(W + token * K + k, k < K, 0).to(tl.float32) + value = tl.sum(x * weight, 0) + tl.store(OUT + row * C + candidate, value) + + +def indexed_fp32_logits( + x: torch.Tensor, + weight: torch.Tensor, + candidate_ids: torch.Tensor, + out: torch.Tensor, +) -> None: + """Evaluate contiguous FP16 rows into an owned FP32 candidate buffer. + + The caller gates the SM70 TP4 DFlash2 shape and candidate bounds. Each + program computes one dot, avoiding both expanded cross-row products and + FP16 storage of the logits used for top-k/top-p. + """ + _indexed_fp32_logits[(x.shape[0], candidate_ids.shape[1])]( + x, + weight, + candidate_ids, + out, + x.shape[1], + candidate_ids.shape[1], + triton.next_power_of_2(x.shape[1]), + num_warps=4, + ) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 83b99abfa5..45bb294b32 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -202,6 +202,9 @@ def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: torch.accelerator.empty_cache() device = weight.device + fp32_logits = envs.VLLM_SM70_DFLASH2_FP32_LOGITS + layer._sm70_dflash2_fp32_logits = fp32_logits + rerank_dtype = torch.float32 if fp32_logits else torch.float16 max_rows = _SM70_DFLASH2_QPN8_MAX_ROWS candidates = _SM70_DFLASH2_QPN8_CANDIDATES layer.register_buffer("_sm70_dflash2_qpn8_codes", codes, persistent=False) @@ -223,7 +226,7 @@ def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: ) layer.register_buffer( "_sm70_dflash2_rerank_logits", - torch.empty((max_rows, candidates), dtype=torch.float16, device=device), + torch.empty((max_rows, candidates), dtype=rerank_dtype, device=device), persistent=False, ) selected_rows = max_rows * candidates @@ -254,7 +257,7 @@ def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: ) layer.register_buffer( "_sm70_dflash2_rerank_dense_logits", - torch.empty((max_rows, rows), dtype=torch.float16, device=device), + torch.empty((max_rows, rows), dtype=rerank_dtype, device=device), persistent=False, ) layer.register_buffer( @@ -262,14 +265,14 @@ def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: torch.empty((max_rows, candidates), dtype=torch.int64, device=device), persistent=False, ) - # Keep distinct top-16 and top-20 outputs. Slicing the columns of one + # Keep distinct top-16, top-20 and top-21 outputs. Slicing the columns of one # [max_rows, 20] allocation for top-16 leaves a row stride of 20 and makes # the result non-contiguous. The TP all-gather requires contiguous inputs, # and inserting a runtime contiguous() copy would add work to both graphs. for selector_k in (16, 20, 21): layer.register_buffer( f"_sm70_dflash2_rerank_values_{selector_k}", - torch.empty((max_rows, selector_k), dtype=torch.float16, device=device), + torch.empty((max_rows, selector_k), dtype=rerank_dtype, device=device), persistent=False, ) layer.register_buffer( @@ -289,7 +292,8 @@ def _prepare_sm70_dflash2_qpn8_rerank(layer: torch.nn.Module) -> bool: ) layer._sm70_dflash2_qpn8_rerank_prepared = True logger.info_once( - "SM70 DFlash2 QPN8 top-64 plus exact TurboMind FP16 rerank layout prepared." + "SM70 DFlash2 QPN8 top-64 rerank layout prepared (%s logits).", + "FP32" if fp32_logits else "FP16", ) return True @@ -399,6 +403,12 @@ def _maybe_sm70_lm_head_forward( x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor | None: + if getattr(layer, "_sm70_dflash2_fp32_logits", False): + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + out = torch.mm(x_2d, layer.weight.t(), out_dtype=torch.float32) + if bias is not None: + out = out + bias.float() + return out.reshape(*x.shape[:-1], out.shape[-1]) if not _sm70_env_bool("VLLM_SM70_ENABLE_LM_HEAD_FASTPATH", False): return None if not getattr(layer, "_sm70_f16_prepared", False): @@ -528,7 +538,11 @@ def _maybe_sm70_dflash2_qpn8_rerank( return None if selector_k not in (16, 20, 21) or bias is not None: return None - if selector_k == 21 and not _sm70_dflash2_use_dense_order(): + if ( + selector_k == 21 + and not _sm70_dflash2_use_dense_order() + and not getattr(layer, "_sm70_dflash2_fp32_logits", False) + ): return None if x.dtype != torch.float16 or not x.is_cuda: return None @@ -567,23 +581,35 @@ def _maybe_sm70_dflash2_qpn8_rerank( out=(qpn8_values, qpn8_ids), ) - sm70_ops.sm70_f16_indexed_rerank_packed_out( - layer._sm70_dflash2_rerank_logits[:num_rows], - x_2d, - layer._sm70_f16_tm_weight, - qpn8_ids, - layer._sm70_dflash2_rerank_selected_packed, - layer._sm70_dflash2_rerank_expanded, - layer._sm70_dflash2_rerank_partials, - layer._sm70_dflash2_rerank_barriers, - _SM70_DFLASH2_RERANK_CTA_N, - _SM70_DFLASH2_RERANK_SPLIT_K, - ) + fp32_logits = getattr(layer, "_sm70_dflash2_fp32_logits", False) + if fp32_logits: + from vllm.model_executor.layers.sm70_fp32_lm_head import indexed_fp32_logits + + indexed_fp32_logits( + x_2d, + layer.weight, + qpn8_ids, + layer._sm70_dflash2_rerank_logits[:num_rows], + ) + logger.info_once("SM70 DFlash2 FP32 candidate logits enabled.") + else: + sm70_ops.sm70_f16_indexed_rerank_packed_out( + layer._sm70_dflash2_rerank_logits[:num_rows], + x_2d, + layer._sm70_f16_tm_weight, + qpn8_ids, + layer._sm70_dflash2_rerank_selected_packed, + layer._sm70_dflash2_rerank_expanded, + layer._sm70_dflash2_rerank_partials, + layer._sm70_dflash2_rerank_barriers, + _SM70_DFLASH2_RERANK_CTA_N, + _SM70_DFLASH2_RERANK_SPLIT_K, + ) rerank_logits = layer._sm70_dflash2_rerank_logits[:num_rows] values, _positions, ids = _sm70_dflash2_rerank_output_buffers( layer, num_rows, selector_k ) - use_dense_order = _sm70_dflash2_use_dense_order() + use_dense_order = fp32_logits or _sm70_dflash2_use_dense_order() if use_dense_order: _sm70_dflash2_dense_order_topk( layer._sm70_dflash2_rerank_dense_logits[:num_rows], @@ -640,8 +666,8 @@ def _maybe_sm70_dflash2_qpn8_rerank( ) logger.info_once( - "SM70 DFlash2 QPN8 top-64 plus exact packed TurboMind FP16 rerank " - "path enabled (dense_order=%s).", + "SM70 DFlash2 QPN8 top-64 plus %s rerank path enabled (dense_order=%s).", + "FP32 candidate" if fp32_logits else "packed TurboMind FP16", use_dense_order, ) output_shape = (*x.shape[:-1], selector_k) diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index 0f921ebd05..d7ce5e1f79 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -785,10 +785,10 @@ def _log_kv_dtype_contract(kv_cache_dtype: str) -> None: "dtype is independent of model weight quantization." ) elif kv_cache_dtype == "fp8_e4m3": - logger.warning( + logger.info( "SM70 Flash-V100 is using explicitly requested E4M3 KV cache. " - "The optimized V100 quantized-KV route uses E5M2. KV-cache dtype " - "is independent of model weight quantization." + "The decode route depends on the native extension and tensor " + "layout. KV-cache dtype is independent of model weight quantization." ) elif kv_cache_dtype == "fp8_e5m2": logger.info( @@ -5279,7 +5279,20 @@ def _dflash2_grouped_verify_allowed( and value_cache.dtype == torch.uint8 and key_cache.stride(-1) == 1 and value_cache.stride(-1) == 1 - and self.kv_cache_dtype == "fp8_e5m2" + and ( + self.kv_cache_dtype == "fp8_e5m2" + or ( + self.kv_cache_dtype == "fp8_e4m3" + and num_query_tokens == 8 + and key_cache.stride(0) % 16 == 0 + and key_cache.stride(1) % 16 == 0 + and value_cache.stride(0) % 16 == 0 + and value_cache.stride(1) % 16 == 0 + and getattr( + self.flash_attn_grouped_verify_paged, "supports_e4m3", False + ) + ) + ) and block_table is not None and block_table.ndim == 2 and block_table.shape[0] == 1 @@ -5341,8 +5354,9 @@ def _call_dflash2_grouped_verify( if not _logged_prefill_smallq_grouped_verify: logger.info( "FLASH_ATTN_V100 DFlash2 exact grouped verifier active " - "(q%d/H6/Hkv1/D256, FP8 E5M2 KV, one-pass).", + "(q%d/H6/Hkv1/D256, %s KV, one-pass).", query.shape[0], + self.kv_cache_dtype, ) _logged_prefill_smallq_grouped_verify = True self.flash_attn_grouped_verify_paged( From f22ac115d0ac0cc8a13bd042cf1472c33add00c4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 11:56:25 +0800 Subject: [PATCH 3/9] [Kernel] Recover DFlash2 latency with context and metadata graphs Generated-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- ...0_quasar_dflash2_quality_speed_recovery.md | 180 ++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 26 +++ .../test_dflash2_context_pipeline.py | 90 +++++++++ .../test_rejection_sampler_utils.py | 11 ++ vllm/envs.py | 8 + vllm/model_executor/models/qwen3_dflash.py | 26 ++- vllm/v1/attention/backends/flash_attn_v100.py | 29 ++- vllm/v1/worker/gpu/model_runner.py | 5 + .../gpu/spec_decode/dflash/speculator.py | 57 ++++-- .../spec_decode/dflash2/sparse_rejection.py | 11 +- .../gpu/spec_decode/dflash2/speculator.py | 159 ++++++++++++++++ .../spec_decode/rejection_sampler_utils.py | 11 +- 12 files changed, 582 insertions(+), 31 deletions(-) create mode 100644 docs/design/sm70_quasar_dflash2_quality_speed_recovery.md create mode 100644 tests/kernels/attention/test_dflash2_context_pipeline.py diff --git a/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md b/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md new file mode 100644 index 0000000000..b949808f1f --- /dev/null +++ b/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md @@ -0,0 +1,180 @@ +# QUASAR DFlash2 quality-preserving latency recovery + +## Result and acceptance boundary + +The E4M3 KV / FP32-logit configuration recovers complete verification-round +latency from about 20 ms to **18.925 ms on release1k and 18.467 ms on MBPP28**. +These are unprofiled medians, with one warmup and three measured requests. +The previous E5M2 / FP16-logit peaks were 18.037 and 17.588 ms respectively. +**The 17.6–18 ms objective is not yet met.** The remaining gap is about 0.9 ms. + +All control/candidate comparisons within each service retained the complete +token sequence and acceptance counts. The long-output comparison also retained +all 3,559 / 1,400 / 1,856 tokens for the three selected MBPP inputs. This is +evidence about these optimizations, not proof that every source of model +quality loss has been eliminated. Independent service startups still produced +different fixed-seed trajectories; cross-start reproducibility remains an +unresolved gate and must not be confused with the within-process comparison. + +The preceding [precision repair](sm70_quasar_e4m3_fp32_logits.md) remains active: +FP32 candidate logits, E4M3 target KV, and exact-reference fallback for ambiguous +top-k/top-p boundaries. This change does not relax their numerical contract. + +## Implementation + +1. Capture the fixed eight-row draft context computation in private CUDA graph + pools. Variable prefill shapes retain the ordinary path. +2. Separate context projection from cache insertion. For B1/q8 decode, project + the target hidden states and compute context K/V before the target sampling + fence. Only insert K/V after acceptance has populated the original slot + mappings. Rejected and evicted rows retain `PAD_SLOT_ID` and do not write. +3. Replay the three persistent metadata copies used by the captured non-causal + Flash-V100 paged draft graph. Other shapes, causal models, CP, alternate + backends, and anchored attention keep the normal metadata builder. +4. Preserve the top-20 view of the 21-column cutoff probe. The rejection kernel + already accepts a row stride; this removes two copies. B1 queries whose + sampling rows cover all real tokens also avoid two identity index gathers. + +The projection, normalization, RoPE, KV write operators, and sampling arithmetic +remain the same. The eager context method and the split methods share the same +implementation. Persistent graph outputs are retained explicitly; compute, +cache insertion, and draft query graphs have separate allocation pools. + +The additional opt-in is: + +```bash +VLLM_SM70_DFLASH2_FP32_LOGITS=1 \ +VLLM_SM70_DFLASH2_CONTEXT_PIPELINE=1 \ +vllm serve \ + --tensor-parallel-size 4 --dtype half \ + --attention-backend FLASH_ATTN_V100 --kv-cache-dtype fp8_e4m3 \ + --max-model-len 262144 --max-num-batched-tokens 4096 \ + --max-num-seqs 4 --gpu-memory-utilization 0.8 \ + --enable-prefix-caching --mamba-cache-mode align \ + --speculative-config '' +``` + +`VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH=1` enables only the context-graph stage for +isolated A/B comparison. The pipeline flag includes that stage. Both new flags +default off; these measurements do not promote other checkpoints or hardware. + +## Frozen workload + +- Integration base: `755baae1d075ee04fa9096b23fc0225b23589a86`. +- Precision control: `6ec27bec9c2e0aec24037597710045b3b8b25e5d`, Draft PR #517. +- Target: `QUASAR-QAT/Qwen3.8-27B-QUASAR-NVFP4`, revision + `d8e6fbfa3e3a78899b440222b827430045a05b44`. +- DFlash2 revision: `dedf8df68adfb1afeaf7b7480c0a0243108177b4`; draft weight + SHA256 `67fc76d68dc5a9415511a4f394ef744d67510cd20e93b37cc2cc7d28e4bab65c`. +- Four V100-SXM2-32GB GPUs, TP4, CUDA 12.8, Torch 2.10.0+cu128, + Python 3.12.13, Triton 3.6.0. Target NVFP4 uses TurboMind W4A16/QPN2. +- V2 runner; target and draft Flash-V100; target E4M3 KV, draft FP16 KV; + seven probabilistic draft tokens, eight-row verification; full/piecewise + target graphs and full draft graphs. +- Context limit 262,144; token budget 4,096; capacity four requests; one live + request; memory utilization 0.8; prefix cache and Mamba align enabled. +- Temperature 1, top-k 20, top-p 0.95; thinking enabled, effort `xhigh`; + no image/video inputs. Natural EOS, 1,024 output-token cap for speed. +- release1k uses 1,019 prompt tokens and seed 20260925. MBPP28 uses 135 prompt + tokens and seed 0. Preserve these different seeds when reproducing. + +## Unprofiled comparison + +The control below is the repeated control at the end of the same-process +control → pipeline → control experiment. Both arms include the removal of +identity gathers and stride copies; only the new graphs/pipeline are toggled. + +| Workload | Control round ms | Pipeline round ms | Control decode tok/s | Pipeline decode tok/s | Tokens / rounds | +| --- | ---: | ---: | ---: | ---: | ---: | +| release1k | 20.035 | 18.925 | 142.457 | 150.817 | 275 / 96 | +| MBPP28 | 19.642 | 18.467 | 243.926 | 259.445 | 437 / 91 | + +Median TTFT was 0.349 / 0.352 s for release1k and 0.101 / 0.106 s for MBPP28 +(control / pipeline). These are small samples, not a TTFT improvement claim. +Decode rates exclude TTFT/prefill. Round latency is not latency per emitted +token; acceptance is shown separately and remained identical in the A/B arms. + +The earlier context-graph-only experiment recovered 20.168 / 19.834 ms to +19.421 / 19.166 ms. A direct same-service comparison then measured context-only +19.424 / 19.124 ms versus staged computation 19.304 / 18.991 ms. Do not combine +different startup trajectories into an acceptance-rate comparison. + +## Operator and text validation + +- Actual loaded draft weights: 16 changing hidden-state/position/slot cases on + each TP rank, covering accepted lengths 1–8 and positions near 1K, 32K, + 128K, and the 256K boundary. Context graph and split compute/write graph + produced bitwise-identical cache entries to eager execution on all ranks. +- Real metadata builders: all three persistent inputs matched ordinary rebuilds + bitwise at 1K, 32K, 128K, and 256K on all four ranks, including changed + physical block-table contents. +- Context microbenchmark: approximately 0.63 ms eager submission/completion + versus 0.066–0.069 ms graph replay, using 200 repeated resident-input calls. + This is a component measurement, not the whole-round speedup. +- Selected MBPP indices 3, 7, and 24: all full token hashes identical between + control and pipeline. Base tests 3/3, Plus tests 1/3; identical transcripts + imply identical scores for the control. The score is not evidence of an + improvement over the previous startup, which had different transcripts. +- Structured JSON returned the requested integer 42. All quality cases ended + naturally, without truncation or replacement characters. +- Focused GPU regressions: 32 passed. Flash-V100 metadata/policy tests: + 15 passed. Ruff and local mypy passed. + +```bash +.venv/bin/python -m pytest -q \ + tests/kernels/attention/test_dflash2_context_pipeline.py \ + tests/v1/spec_decode/test_rejection_sampler_utils.py \ + tests/v1/spec_decode/test_dflash2.py \ + -k 'context or dflash2_sparse_topk or compact' +.venv/bin/python -m pytest -q \ + tests/v1/attention/test_sm70_flash_v100_policy.py -k 'metadata or draft' +``` + +The new kernel regression checks that early computation does not mutate the +cache, then changes acceptance and physical slots before every write replay. +It compares the entire cache, including invalid entries, with masked eager +execution. Padded-candidate rejection tests exercise the retained sentinel +stride while placing a large excluded value in the sentinel column. + +256K validation here is at operator/metadata level. No new 256K endpoint speed +or full long-context model-quality result is claimed. + +## Profiling and rejected paths + +Nsight Systems 2025.3.1 captured CUDA graph nodes on all four TP ranks. Aligning +the explicit verification-round ordinals gives 12 complete groups; excluding +the two edges leaves 10 steady groups. Diagnostic round intervals averaged +20.815 ms, and host `DFlash2Speculator.propose` ranges averaged 0.801 ms. +These traced values include profiling overhead and are not accepted speed. +The compact cutoff range includes the preceding GPU-completion wait; its +12.443 ms average must not be called CPU boundary-check computation. + +- Nsight 2026.4 produced no CUDA report even for a minimal probe. Its bundled + release notes explicitly remove Pascal/Volta support starting at 2025.4. + A verified NVIDIA 2025.3.1 package fixed this environment issue. +- The first 2025.3 trace had all worker NVTX ranges but GPU events only for + rank zero. It is retained as partial evidence, not a TP4 critical-path table. + Starting/stopping the CUDA profiler on every rank produced the complete trace. +- Pinning the four worker main threads to separate NUMA-local cores changed + medians by only about 0.03–0.04 ms in a direct test. Original affinities were + restored; pinning is not part of the accepted configuration. +- The boundary guard still detects real ambiguous inputs. Disabling it, lowering + logit precision, or reverting KV dtype was not admitted as a speed fix. + +## Retained evidence + +Task bundle: `v100-quasar-quality-speed-recovery-20260906`. + +- `results/metadata-ab-{control,control-repeat,pipeline}-speed-*.json`. +- `results/final-{control,pipeline}-quality-subset.json`, + `results/final-pipeline-json.json`, `results/final-evalplus.json`. +- `results/metadata-context-real-oracle.json`, `results/metadata-real-oracle.json`. +- `profile/final-tp4.nsys-rep`, original SQLite, explicitly aligned SQLite, + `final-tp4.rounds.json`, and `final-gpu-breakdown.{json,csv,md}`. + Each trace row is a verification round, not an emitted token. +- `scripts/serve-pipeline.sh`, A/B clients, real-weight worker oracles, task + lease/ownership record, worker/native provenance, and the tested source diff. + +The next performance gate is the remaining sampling/state/launch dependency +after target completion. It needs its own exact-distribution and mutable-state +validation; a projected saving is not evidence that the old peak is restored. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 53c1911175..fef31a4fae 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45789,3 +45789,29 @@ Interpretation: fused variants compile with 31 registers/16 bytes shared/zero stack or spills. Compare complete HC with the newly registered up fusion held fixed, including actual auxiliary sum2 and post-wrap checks. GPU gate pending. + +## 2026-09-06 QUASAR E4M3 / FP32 DFlash2 latency recovery + +- Continue owned Draft PR #517 from precision-control `6ec27bec9c`; integration + base remains `755baae1d0`. Keep E4M3 target KV, FP32 candidate logits, and + ambiguous-cutoff reference sampling. +- Capture q8 context computation and accepted-slot writes; stage independent + context work before sampling; refresh the non-causal paged graph's persistent + metadata directly. Remove four identity gather/stride-copy launches. +- Same-process control/pipeline/control medians: release1k 20.035 → 18.925 ms; + MBPP28 19.642 → 18.467 ms. Full token sequences and acceptance counts match. + The historical 18.037 / 17.588 ms peaks are still about 0.9 ms faster. +- Real weights: all 64 changing-context cases and all 16 metadata cases match + eager reference bitwise, including accepted lengths 1–8 and the 256K boundary. + Three long code responses (3559 / 1400 / 1856 tokens) match completely; MBPP + Base 3/3, Plus 1/3 in both arms; structured JSON 42 passes. Independent startup + trajectories differ, so cross-start reproducibility is still unclosed. +- Focused GPU tests 32 passed; metadata/policy tests 15 passed; Ruff/mypy pass. +- Profiling environment: use Nsight Systems 2025.3.1 for V100. 2026.4 does not + support Volta. Start the profiler in every TP worker; a rank-zero-only CUDA + trace must not be reported as TP4 GPU evidence. Final trace has all four ranks. +- NUMA pinning gave only ~0.03–0.04 ms; original affinity restored. Do not repeat + that experiment or disable sampling guards to manufacture a peak result. +- Full contract, exact tests, acceptance limits, and artifact bundle are in + [the recovery report](sm70_quasar_dflash2_quality_speed_recovery.md). Raw bundle: + `v100-quasar-quality-speed-recovery-20260906`. diff --git a/tests/kernels/attention/test_dflash2_context_pipeline.py b/tests/kernels/attention/test_dflash2_context_pipeline.py new file mode 100644 index 0000000000..fa2de17009 --- /dev/null +++ b/tests/kernels/attention/test_dflash2_context_pipeline.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.models.qwen3_dflash import DFlashQwen3Model + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@torch.inference_mode() +def test_context_pipeline_defers_writes_and_refreshes_accepted_slots(): + """Computing rejected rows early must never write their scratch K/V.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(717) + model = DFlashQwen3Model.__new__(DFlashQwen3Model) + torch.nn.Module.__init__(model) + layers, hidden, heads, dim, block = 3, 256, 2, 128, 16 + model._num_attn_layers = layers + model._kv_size = heads * dim + model._head_dim = dim + model._num_kv_heads = heads + model._rms_norm_eps = 1e-6 + model._hidden_norm_weight = torch.ones(hidden, device=device, dtype=torch.float16) + model._fused_kv_weight = ( + torch.randn(layers * 2 * heads * dim, hidden, generator=gen, device=device) + * 0.03 + ).half() + model._fused_kv_bias = None + model._k_norm_weights = ( + 1 + torch.randn(layers, dim, generator=gen, device=device) * 0.1 + ).half() + model._rope_head_size = dim + model._rope_is_neox = True + freq = 10000.0 ** (-torch.arange(0, dim, 2, device=device).float() / dim) + angles = torch.arange(4096, device=device).float()[:, None] * freq + model._rope_cos_sin_cache = torch.cat((angles.cos(), angles.sin()), dim=1).half() + model._sm70_context_k_debugged = False + scale = torch.tensor(1.0, device=device) + + def update(attn, key, value, cache, slots): + ops.reshape_and_cache_flash( + key, value, cache[:, 0], cache[:, 1], slots, "auto", scale, scale + ) + + caches = [ + torch.full((4, 2, block, heads, dim), 0.125, device=device, dtype=torch.float16) + for _ in range(layers) + ] + model._attn_layers = [ + SimpleNamespace(kv_cache=c, impl=SimpleNamespace(do_kv_cache_update=update)) + for c in caches + ] + states = torch.randn(8, hidden, generator=gen, device=device).half() + positions = torch.arange(8, device=device) + slots = [torch.full((8,), -1, device=device, dtype=torch.int64) for _ in caches] + model.precompute_and_store_context_kv(states, positions, slots) + compute, write = torch.cuda.CUDAGraph(), torch.cuda.CUDAGraph() + with torch.cuda.graph(compute): + key, value = model.compute_context_kv(states, positions) + with torch.cuda.graph(write): + model.store_context_kv(key, value, slots) + + for step in range(16): + accepted = step % 8 + 1 + states.copy_(torch.randn(states.shape, generator=gen, device=device).half()) + positions.copy_(torch.arange(4080 - step, 4088 - step, device=device)) + before = [c.clone() for c in caches] + compute.replay() + # Acceptance and cache placement are intentionally unknown until after + # projection. Change both between every replay, including invalid slots. + for layer, mapping in enumerate(slots): + mapping.copy_( + block * (1 + layer % 3) + torch.arange(8, device=device).roll(step) + ) + mapping[accepted:] = -1 + for actual, expected in zip(caches, before): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + write.replay() + actual = [c.clone() for c in caches] + for c, snapshot in zip(caches, before): + c.copy_(snapshot) + reference_positions = positions.clone() + reference_positions[accepted:] = 0 + model.precompute_and_store_context_kv(states, reference_positions, slots) + for got, expected in zip(actual, caches): + torch.testing.assert_close(got, expected, rtol=0, atol=0) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 8edc3624b7..770c7c71a9 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -134,10 +134,12 @@ def _assert_distribution_match( @pytest.mark.parametrize("num_speculative_steps", [3, 7]) @pytest.mark.parametrize("top_p", [1.0, 0.95]) @pytest.mark.parametrize("temperature", [0.6, 1.0]) +@pytest.mark.parametrize("padded_candidates", [False, True]) def test_dflash2_sparse_topk_matches_dense_rejection( num_speculative_steps: int, top_p: float, temperature: float, + padded_candidates: bool, ): """Compact p/q support must preserve the dense DFlash2 decision path.""" torch.manual_seed(20260823) @@ -245,6 +247,15 @@ def test_dflash2_sparse_topk_matches_dense_rejection( seeds, num_speculative_steps, ) + if padded_candidates: + padded_ids = target_topk_ids.new_zeros((num_logits, target_top_k + 1)) + padded_logits = target_topk_logits.new_full( + (num_logits, target_top_k + 1), 1000.0 + ) + padded_ids[:, :target_top_k].copy_(target_topk_ids) + padded_logits[:, :target_top_k].copy_(target_topk_logits) + target_topk_ids = padded_ids[:, :target_top_k] + target_topk_logits = padded_logits[:, :target_top_k] sparse_sampled, sparse_num_sampled = dflash2_sparse_topk_rejection_sample( target_topk_ids, target_topk_logits, diff --git a/vllm/envs.py b/vllm/envs.py index b2c10b7d3f..9905ca7be8 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -241,6 +241,8 @@ VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS: bool = False VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION: bool = False VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC: bool = False + VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH: bool = False + VLLM_SM70_DFLASH2_CONTEXT_PIPELINE: bool = False VLLM_SM70_GLM53_MHC_NATIVE_VERIFY: bool = False VLLM_SM70_GLM53_MHC_FUSED_POST_DOT_Q8: bool = False VLLM_SM70_GLM53_MOE_FUSED_PERMUTE_Q8: bool = True @@ -2209,6 +2211,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC": lambda: bool( int(os.getenv("VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC", "0")) ), + "VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH", "0")) + ), + "VLLM_SM70_DFLASH2_CONTEXT_PIPELINE": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_CONTEXT_PIPELINE", "0")) + ), # Native SM70 final stage for the GLM-5.3 q8 mHC verifier. Audited model and # topology contracts enable it while the global default remains off. "VLLM_SM70_GLM53_MHC_NATIVE_VERIFY": lambda: bool( diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index c3eaf63cc6..c1b594460c 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -722,6 +722,15 @@ def precompute_and_store_context_kv( When context_slot_mapping is None (e.g. during dummy_run) only the computation runs, and no K/V is written to cache. """ + all_k, all_v = self.compute_context_kv(context_states, context_positions) + self.store_context_kv(all_k, all_v, context_slot_mapping) + + def compute_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Project context without modifying the KV cache.""" if not hasattr(self, "_num_attn_layers"): logger.warning_once( "DFlash buffer initialization was skipped. If dummy weights are not " @@ -755,13 +764,20 @@ def precompute_and_store_context_kv( self._rope_is_neox, ) + all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) + return all_k_final, all_v + + def store_context_kv( + self, + all_k: torch.Tensor, + all_v: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + """Write only the accepted, resident context slots.""" if context_slot_mapping is None: return - - # --- Per-layer cache insert --- - all_k_final = all_k_flat.view(L, num_ctx, nkv, hd) per_layer = isinstance(context_slot_mapping, (list, tuple)) - for i in range(L): + for i in range(self._num_attn_layers): slot_mapping = ( context_slot_mapping[i] if per_layer else context_slot_mapping ) @@ -771,7 +787,7 @@ def precompute_and_store_context_kv( kv_cache = attn.kv_cache attn.impl.do_kv_cache_update( attn, - all_k_final[i], + all_k[i], all_v[i], kv_cache, slot_mapping, diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index d7ce5e1f79..9736dd0eca 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -3515,20 +3515,37 @@ def _stabilize_draft_graph_metadata( assert self._draft_seq_lens is not None assert self._draft_query_start_loc is not None - self._draft_block_table[:num_reqs].copy_(block_table, non_blocking=True) - self._draft_seq_lens[:num_reqs].copy_( + self.copy_dflash_graph_metadata( + block_table, attn_metadata.seq_lens[:num_reqs], - non_blocking=True, - ) - self._draft_query_start_loc[: num_reqs + 1].copy_( attn_metadata.query_start_loc[: num_reqs + 1], - non_blocking=True, ) attn_metadata.block_table = self._draft_block_table[:num_reqs] attn_metadata.seq_lens = self._draft_seq_lens[:num_reqs] attn_metadata.query_start_loc = self._draft_query_start_loc[: num_reqs + 1] + def copy_dflash_graph_metadata( + self, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + ) -> None: + """Refresh the three persistent inputs of a non-causal DFlash graph.""" + num_reqs = seq_lens.numel() + assert self._draft_block_table is not None + assert self._draft_seq_lens is not None + assert self._draft_query_start_loc is not None + self._draft_block_table[:num_reqs].copy_(block_table, non_blocking=True) + self._draft_seq_lens[:num_reqs].copy_( + seq_lens, + non_blocking=True, + ) + self._draft_query_start_loc[: num_reqs + 1].copy_( + query_start_loc, + non_blocking=True, + ) + def _configured_smallq_max_query_len(self) -> int: return int(os.getenv("VLLM_FLASH_V100_SMALLQ_DECODE_MAX_Q", "16")) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d77b753c44..45ece776a2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -117,6 +117,7 @@ from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import ( try_dflash2_sparse_target_rejection, ) +from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, ) @@ -1726,6 +1727,10 @@ def sample_tokens( return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) # Last rank: sample tokens + if isinstance(self.speculator, DFlash2Speculator): + self.speculator.prepare_target_context( + input_batch, hidden_states, aux_hidden_states + ) mtp_sample_start = self._sm70_v2_mtp_profile_start(mtp_profile_ctx) sampler_output, num_sampled, num_rejected = self.sample( hidden_states, input_batch, grammar_output diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index e0fb68aee3..21eb0be449 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -565,7 +565,10 @@ def propose( # number of rejected tokens, we maintain the size of input_ids and # hidden_states the same as the target model's. This means, we pad each # request's query length to include any rejected positions. - if aux_hidden_states: + prepared_hidden = self._get_prepared_context_hidden(input_batch) + if prepared_hidden is not None: + hidden_states = prepared_hidden + elif aux_hidden_states: if ( getattr(self, "_debug_proposal_stages", False) and getattr(self, "_debug_real_proposal", False) @@ -692,7 +695,7 @@ def propose( context_slots = self._context_slot_mappings[0][:num_target_tokens] self._debug_proposal_stage("context kv begin") with record_function_or_nullcontext("dflash: materialize context kv"): - self.model.precompute_and_store_context_kv( + self._precompute_context_kv( self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], context_slots, @@ -741,20 +744,26 @@ def propose( num_reqs_padded = batch_desc.num_reqs or num_reqs num_tokens_padded = batch_desc.num_tokens - # Rebuild the draft attention metadata even when replaying the FULL - # graph so that any attention metadata builder state is updated. - draft_attn_metadata = self._build_draft_attn_metadata( - num_reqs=num_reqs, - num_reqs_padded=num_reqs_padded, - num_tokens_padded=num_tokens_padded, - seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, - step=self.num_query_per_req, - causal=self._group_causal, - ) - draft_slot_mappings_by_layer = build_slot_mappings_by_layer( - query_slot_mappings[:, :num_tokens_padded], - self.kv_cache_config, - ) + # Refresh persistent metadata before FULL replay. The supported + # DFlash2 path copies its captured inputs; other routes rebuild. + if batch_desc.cg_mode == CUDAGraphMode.FULL and ( + self._refresh_draft_graph_metadata(num_reqs_padded, num_tokens_padded) + ): + draft_attn_metadata = None + draft_slot_mappings_by_layer = None + else: + draft_attn_metadata = self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_tokens_padded=num_tokens_padded, + seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound, + step=self.num_query_per_req, + causal=self._group_causal, + ) + draft_slot_mappings_by_layer = build_slot_mappings_by_layer( + query_slot_mappings[:, :num_tokens_padded], + self.kv_cache_config, + ) self._debug_proposal_stage("query preparation end") if batch_desc.cg_mode == CUDAGraphMode.FULL: @@ -787,6 +796,22 @@ def propose( return self.draft_tokens[:num_reqs] + def _precompute_context_kv( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + slots: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + self.model.precompute_and_store_context_kv(hidden_states, positions, slots) + + def _get_prepared_context_hidden( + self, input_batch: InputBatch + ) -> torch.Tensor | None: + return None + + def _refresh_draft_graph_metadata(self, num_reqs: int, num_tokens: int) -> bool: + return False + @triton.jit def _prepare_dflash_inputs_kernel( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py index 6fa1f1d51d..fca6ac1afb 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py @@ -263,8 +263,15 @@ def try_dflash2_sparse_target_rejection( return None target_topk_ids = target_topk_ids[:, :_TARGET_TOP_K] target_topk_logits = target_topk_logits[:, :_TARGET_TOP_K] - draft_sampled = input_batch.input_ids[input_batch.logits_indices] - pos = input_batch.positions[input_batch.logits_indices] + num_rows = target_topk_ids.shape[0] + if input_batch.num_tokens == num_rows: + # For the gated B1 decode, logits_indices spans the entire real query. + # Keep views instead of launching two identity gather kernels. + draft_sampled = input_batch.input_ids[:num_rows] + pos = input_batch.positions[:num_rows] + else: + draft_sampled = input_batch.input_ids[input_batch.logits_indices] + pos = input_batch.positions[input_batch.logits_indices] sampled, num_sampled = dflash2_sparse_topk_rejection_sample( target_topk_ids, target_topk_logits, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 935febea60..3dbde1f8b0 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -11,6 +11,7 @@ from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dflash2.lookup import ( @@ -353,6 +354,11 @@ class DFlash2Speculator(DFlashSpeculator): def __init__(self, vllm_config: VllmConfig, device: torch.device): super().__init__(vllm_config, device) + self._context_kv_graph: torch.cuda.CUDAGraph | None = None + self._context_compute_graph: torch.cuda.CUDAGraph | None = None + self._context_store_graph: torch.cuda.CUDAGraph | None = None + self._prepared_context_batch: InputBatch | None = None + self._draft_metadata_graph: torch.cuda.CUDAGraph | None = None self._debug_token_dump_count = 0 draft_config = self.draft_model_config.hf_config.dflash_config self.selector_top_k = int(draft_config["selector_top_k"]) @@ -742,6 +748,159 @@ def next_num_draft_tokens(self) -> int: reason = "strong-copy" if self._lookup_long_active else "adaptive-default" return self._record_lookup_width(width, reason) + def capture(self) -> None: + super().capture() + if ( + not ( + envs.VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH + or envs.VLLM_SM70_DFLASH2_CONTEXT_PIPELINE + ) + or self.device.type != "cuda" + or torch.cuda.get_device_capability(self.device) != (7, 0) + or self.num_query_per_req != 8 + or self.query_cudagraph_manager is None + or not self.query_cudagraph_manager.graphs + ): + return + slots = ( + [self._context_slot_mappings[i][:8] for i in self._layer_group_idx] + if self._layer_group_idx is not None + else self._context_slot_mappings[0][:8] + ) + graph = torch.cuda.CUDAGraph() + # All inputs are persistent draft buffers refreshed by propose(). + # A separate pool keeps these intermediates independent of query graphs. + with torch.cuda.graph(graph): + super()._precompute_context_kv( + self.hidden_states[:8], self.context_positions[:8], slots + ) + self._context_kv_graph = graph + logger.info("SM70 DFlash2 q8 context KV CUDA graph captured.") + if not envs.VLLM_SM70_DFLASH2_CONTEXT_PIPELINE: + return + self._context_target_positions = torch.zeros_like(self.context_positions[:8]) + compute = torch.cuda.CUDAGraph() + with torch.cuda.graph(compute): + all_k, all_v = self.model.model.compute_context_kv( + self.hidden_states[:8], self._context_target_positions + ) + write = torch.cuda.CUDAGraph() + with torch.cuda.graph(write): + self.model.model.store_context_kv(all_k, all_v, slots) + self._context_compute_graph = compute + self._context_store_graph = write + self._context_projected_kv = (all_k, all_v) + logger.info("SM70 DFlash2 context computation is staged before sampling.") + self._capture_draft_metadata_graph() + + def _capture_draft_metadata_graph(self) -> None: + from vllm.v1.attention.backends.flash_attn_v100 import ( + FlashAttnV100Impl, + FlashAttnV100MetadataBuilder, + ) + + has_causal = ( + any(self._group_causal.values()) + if isinstance(self._group_causal, dict) + else self._group_causal + ) + if self.block_tables.cp_size != 1 or has_causal: + return + if any( + not isinstance(a.impl, FlashAttnV100Impl) + or a.impl.use_triton_prefill + or not a.impl.use_flash_v100_prefill_paged + or a.impl.prefix_anchored_decode_window is not None + for a in self.model.model._attn_layers + ): + return + builders = [] + for gid, groups in enumerate(self.attn_groups): + for group in groups: + builder = group.get_metadata_builder(0) + if ( + not isinstance(builder, FlashAttnV100MetadataBuilder) + or not builder._is_dflash_draft_model + or builder._flash_draft_buffer_shape is None + ): + return + builders.append((gid, builder)) + if not builders: + return + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for gid, builder in builders: + builder.copy_dflash_graph_metadata( + self.block_tables.input_block_tables[gid][:1], + self.input_buffers.seq_lens[:1], + self.input_buffers.query_start_loc[:2], + ) + self._draft_metadata_graph = graph + logger.info("SM70 DFlash2 B1 paged graph metadata refresh captured.") + + def _refresh_draft_graph_metadata(self, num_reqs: int, num_tokens: int) -> bool: + if self._draft_metadata_graph is None or (num_reqs, num_tokens) != (1, 8): + return False + # The non-causal paged query graph reads only these persistent buffers; + # prepare_dflash_inputs has already refreshed its query slots and rows. + self._draft_metadata_graph.replay() + return True + + def prepare_target_context( + self, + input_batch: InputBatch, + hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + ) -> None: + self._prepared_context_batch = None + if ( + self._context_compute_graph is None + or input_batch.num_reqs != 1 + or input_batch.num_tokens != 8 + or input_batch.num_draft_tokens != 7 + or input_batch.is_prefilling_np[0] + ): + return + if aux_hidden_states: + hidden_states = self.model.combine_hidden_states( + torch.cat(aux_hidden_states, dim=-1) + ) + self.hidden_states[:8].copy_(hidden_states[:8]) + self._context_target_positions.copy_(input_batch.positions[:8]) + # Context projection does not depend on the acceptance decision. Raw + # positions equal the later masked positions for every accepted row. + # Rejected rows remain scratch data and never reach the KV cache. + self._context_compute_graph.replay() + self._prepared_context_batch = input_batch + logger.info_once("Using SM70 DFlash2 context pipeline before target sampling.") + + def _get_prepared_context_hidden( + self, input_batch: InputBatch + ) -> torch.Tensor | None: + if self._prepared_context_batch is input_batch: + return self.hidden_states[:8] + return None + + def _precompute_context_kv( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + slots: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + if self._prepared_context_batch is not None and slots is not None: + assert self._context_store_graph is not None + self._context_store_graph.replay() + self._prepared_context_batch = None + return + if ( + self._context_kv_graph is not None + and hidden_states.shape[0] == 8 + and slots is not None + ): + self._context_kv_graph.replay() + return + super()._precompute_context_kv(hidden_states, positions, slots) + def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: # The selector walk and rejection sampler must consume identical scores. # BF16 rounding measurably changes candidate order, so keep this FP32. diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 253ff39553..5624bd9428 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -778,8 +778,15 @@ def dflash2_sparse_topk_rejection_sample( if not 0 < draft_topk_ids.shape[2] <= 64: raise ValueError("draft top-k width must be in [1, 64]") - target_topk_ids = target_topk_ids.contiguous() - target_topk_logits = target_topk_logits.contiguous() + # The cutoff probe retains 21 columns before slicing to top-20. The kernel + # already accepts a row stride, so preserve that view and avoid two copies. + if ( + target_topk_ids.stride(-1) != 1 + or target_topk_logits.stride(-1) != 1 + or target_topk_ids.stride(0) != target_topk_logits.stride(0) + ): + target_topk_ids = target_topk_ids.contiguous() + target_topk_logits = target_topk_logits.contiguous() draft_topk_ids = draft_topk_ids.contiguous() draft_topk_logits = draft_topk_logits.contiguous() num_reqs = cu_num_logits.shape[0] - 1 From 6f07be1cce338b87e7e8d1a714bf9e30e5424fd8 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 12:01:50 +0800 Subject: [PATCH 4/9] [Doc] Record production confirmation of DFlash2 latency recovery Generated-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- ...0_quasar_dflash2_quality_speed_recovery.md | 30 ++++++++++++++++--- docs/design/sm70_v100_migration_control.md | 10 +++++-- 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md b/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md index b949808f1f..b82c621f49 100644 --- a/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md +++ b/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md @@ -3,7 +3,8 @@ ## Result and acceptance boundary The E4M3 KV / FP32-logit configuration recovers complete verification-round -latency from about 20 ms to **18.925 ms on release1k and 18.467 ms on MBPP28**. +latency from about 20 ms to **18.892 ms on release1k and 18.435 ms on MBPP28** +in an independent production startup without diagnostic worker extensions. These are unprofiled medians, with one warmup and three measured requests. The previous E5M2 / FP16-logit peaks were 18.037 and 17.588 ms respectively. **The 17.6–18 ms objective is not yet met.** The remaining gap is about 0.9 ms. @@ -12,9 +13,10 @@ All control/candidate comparisons within each service retained the complete token sequence and acceptance counts. The long-output comparison also retained all 3,559 / 1,400 / 1,856 tokens for the three selected MBPP inputs. This is evidence about these optimizations, not proof that every source of model -quality loss has been eliminated. Independent service startups still produced -different fixed-seed trajectories; cross-start reproducibility remains an -unresolved gate and must not be confused with the within-process comparison. +quality loss has been eliminated. The production confirmation reproduced both +speed-fixture hashes and all 4,939 / 5,904 / 633 quality-fixture tokens from the +previous precision configuration. Task-only diagnostic services produced other +fixed-seed trajectories; the cause of that difference remains unresolved. The preceding [precision repair](sm70_quasar_e4m3_fp32_logits.md) remains active: FP32 candidate logits, E4M3 target KV, and exact-reference fallback for ambiguous @@ -80,6 +82,20 @@ default off; these measurements do not promote other checkpoints or hardware. ## Unprofiled comparison +The independent production confirmation uses implementation commit +`f22ac115d0ac0cc8a13bd042cf1472c33add00c4`, no profiler or development mode, and +fresh task-owned compile caches. Against the preceding precision repair: + +| Workload | Precision control round ms | Optimized round ms | Control / optimized decode tok/s | Tokens / rounds | +| --- | ---: | ---: | ---: | ---: | +| release1k | 20.149 | 18.892 | 135.033 / 144.015 | 303 / 111 | +| MBPP28 | 19.745 | 18.435 | 267.695 / 286.721 | 260 / 49 | + +The full token hashes and acceptance counts match the preceding control on all +four requests (warmup plus three measured runs) for both fixtures. Final TTFT +medians are 0.350 / 0.109 s. The round cost drops by 6.24% / 6.64%; the remaining +gap against the old E5M2/FP16 peaks is 0.855 / 0.847 ms. + The control below is the repeated control at the end of the same-process control → pipeline → control experiment. Both arms include the removal of identity gathers and stride copies; only the new graphs/pipeline are toggled. @@ -117,6 +133,10 @@ different startup trajectories into an acceptance-rate comparison. improvement over the previous startup, which had different transcripts. - Structured JSON returned the requested integer 42. All quality cases ended naturally, without truncation or replacement characters. +- Independent production confirmation: MBPP outputs 4,939 / 5,904 / 633 tokens, + bitwise identical token IDs to the preceding precision configuration; Base + 3/3, Plus 0/3, and JSON result 42. The diagnostic service's Plus 1/3 result + above must not be reported as a production quality improvement. - Focused GPU regressions: 32 passed. Flash-V100 metadata/policy tests: 15 passed. Ruff and local mypy passed. @@ -168,6 +188,8 @@ Task bundle: `v100-quasar-quality-speed-recovery-20260906`. - `results/metadata-ab-{control,control-repeat,pipeline}-speed-*.json`. - `results/final-{control,pipeline}-quality-subset.json`, `results/final-pipeline-json.json`, `results/final-evalplus.json`. +- `results/production-confirm-{speed-release1k,speed-mbpp28,quality-subset,json,evalplus}.json`, + `production-worker-provenance.json` (four workers, native hashes and 184 JIT cubins). - `results/metadata-context-real-oracle.json`, `results/metadata-real-oracle.json`. - `profile/final-tp4.nsys-rep`, original SQLite, explicitly aligned SQLite, `final-tp4.rounds.json`, and `final-gpu-breakdown.{json,csv,md}`. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index fef31a4fae..0421777367 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45801,11 +45801,17 @@ Interpretation: - Same-process control/pipeline/control medians: release1k 20.035 → 18.925 ms; MBPP28 19.642 → 18.467 ms. Full token sequences and acceptance counts match. The historical 18.037 / 17.588 ms peaks are still about 0.9 ms faster. +- Independent production confirmation at `f22ac115d0`: 18.892 / 18.435 ms, + 144.015 / 286.721 pure decode tok/s; 303 / 260 tokens and 111 / 49 rounds. + Both speed hashes match the preceding precision control; round cost drops + 6.24% / 6.64%. No profiler, development mode, or diagnostic worker extension. - Real weights: all 64 changing-context cases and all 16 metadata cases match eager reference bitwise, including accepted lengths 1–8 and the 256K boundary. Three long code responses (3559 / 1400 / 1856 tokens) match completely; MBPP - Base 3/3, Plus 1/3 in both arms; structured JSON 42 passes. Independent startup - trajectories differ, so cross-start reproducibility is still unclosed. + Base 3/3, Plus 1/3 in both diagnostic arms; structured JSON 42 passes. + Production outputs 4939 / 5904 / 633 tokens match the preceding precision + configuration completely, with Base 3/3, Plus 0/3 and JSON 42. The diagnostic + service's different trajectories remain unexplained; no quality-score gain. - Focused GPU tests 32 passed; metadata/policy tests 15 passed; Ruff/mypy pass. - Profiling environment: use Nsight Systems 2025.3.1 for V100. 2026.4 does not support Volta. Start the profiler in every TP worker; a rank-zero-only CUDA From fcb6dada5882d8b40a8ca255eda26ab4ba68b9aa Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 13:31:33 +0800 Subject: [PATCH 5/9] [Bugfix] Preserve NVFP4 alignment and FP32 logits on TP2 Generated-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_sm70_nvfp4_output_padding.py | 49 +++++++++++++++++++ .../test_sm70_fp32_lm_head_admission.py | 24 +++++++++ .../test_sm70_turbomind_adapter.py | 19 ++++--- .../layers/quantization/sm70_turbomind.py | 4 +- .../layers/vocab_parallel_embedding.py | 10 ++++ 5 files changed, 97 insertions(+), 9 deletions(-) create mode 100644 tests/kernels/quantization/test_sm70_nvfp4_output_padding.py create mode 100644 tests/quantization/test_sm70_fp32_lm_head_admission.py diff --git a/tests/kernels/quantization/test_sm70_nvfp4_output_padding.py b/tests/kernels/quantization/test_sm70_nvfp4_output_padding.py new file mode 100644 index 0000000000..723a752693 --- /dev/null +++ b/tests/kernels/quantization/test_sm70_nvfp4_output_padding.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 converter/consumer alignment, including Qwen GDN TP2 width.""" + +import pytest +import torch + +from vllm.model_executor.layers.quantization import sm70_turbomind as tm + + +@pytest.mark.parametrize("n", [48, 8240]) +@pytest.mark.parametrize("m", [1, 8, 17]) +@torch.inference_mode() +def test_nvfp4_unaligned_output_matches_dense_and_graph(n, m): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("requires SM70") + if not hasattr(torch.ops._C, "nvfp4_sm70_prepare"): + pytest.skip("requires native TurboMind NVFP4") + gen = torch.Generator(device="cuda").manual_seed(713) + k = 512 + packed = torch.randint( + 0, 256, (n, k // 2), device="cuda", dtype=torch.uint8, generator=gen + ) + scales = torch.full((n, k // 16), 0.125, device="cuda", dtype=torch.float16) + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter(packed, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False) + layer.weight_global_scale = torch.nn.Parameter( + torch.tensor(0.5, device="cuda"), requires_grad=False + ) + codes = torch.stack((packed & 15, packed >> 4), -1).flatten(-2) + lut = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device="cuda") + weight = lut[(codes & 7).long()] * torch.where((codes & 8) != 0, -1.0, 1.0) * 0.0625 + x = torch.randn((m, k), device="cuda", generator=gen).half() * 0.1 + tm.prepare_nvfp4_linear(layer) + state = getattr(layer, tm.STATE_ATTR) + assert state.padded_output_size == (n + 31) // 32 * 32 + eager = tm.apply_prepared_linear(layer, x, None) + expected = x.float() @ weight.t() + torch.testing.assert_close(eager.float(), expected, atol=2e-3, rtol=1e-3) + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph): + out = tm.apply_prepared_linear(layer, x, None) + x.mul_(0.5) + graph.replay() + torch.testing.assert_close( + out, tm.apply_prepared_linear(layer, x, None), atol=0, rtol=0 + ) diff --git a/tests/quantization/test_sm70_fp32_lm_head_admission.py b/tests/quantization/test_sm70_fp32_lm_head_admission.py new file mode 100644 index 0000000000..71e2143a0e --- /dev/null +++ b/tests/quantization/test_sm70_fp32_lm_head_admission.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers import vocab_parallel_embedding as vocab + + +@pytest.mark.parametrize("tp", [2, 4]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_fp32_head_admission_without_qpn8_layout(monkeypatch, tp, enabled): + layer = SimpleNamespace( + tp_size=tp, + weight=torch.empty((248320 // tp, 5120), device="meta"), + ) + monkeypatch.setattr(vocab, "_is_sm70_lm_head_fastpath_eligible", lambda _: True) + monkeypatch.setattr(vocab, "_sm70_lm_head_packed_layout_requested", lambda: False) + monkeypatch.setattr(vocab.envs, "VLLM_SM70_DFLASH2_FP32_LOGITS", enabled) + assert vocab.maybe_prepare_sm70_lm_head_top1(layer) + assert getattr(layer, "_sm70_dflash2_fp32_logits", False) == enabled + assert not getattr(layer, "_sm70_dflash2_qpn8_rerank_prepared", False) diff --git a/tests/quantization/test_sm70_turbomind_adapter.py b/tests/quantization/test_sm70_turbomind_adapter.py index 5d56668fb6..296ef759e2 100644 --- a/tests/quantization/test_sm70_turbomind_adapter.py +++ b/tests/quantization/test_sm70_turbomind_adapter.py @@ -5,6 +5,7 @@ import sys from pathlib import Path +import pytest import torch @@ -83,14 +84,15 @@ def test_mxfp4_unpack_flattens_last_two_block_dims_like_lmdeploy(): assert weight.tolist() == [[0, 4], [1, 5], [2, 6], [3, 7]] -def test_nvfp4_prepare_pads_output_to_converter_alignment(monkeypatch): +@pytest.mark.parametrize("logical_n", [24, 48]) +def test_nvfp4_prepare_pads_output_to_converter_alignment(monkeypatch, logical_n): tm = _load_adapter() layer = torch.nn.Module() layer.weight = torch.nn.Parameter( - torch.zeros((24, 16), dtype=torch.uint8), requires_grad=False + torch.zeros((logical_n, 16), dtype=torch.uint8), requires_grad=False ) layer.weight_scale = torch.nn.Parameter( - torch.ones((24, 2), dtype=torch.float16), requires_grad=False + torch.ones((logical_n, 2), dtype=torch.float16), requires_grad=False ) layer.weight_global_scale = torch.nn.Parameter( torch.tensor(0.25, dtype=torch.float32), requires_grad=False @@ -99,6 +101,7 @@ def test_nvfp4_prepare_pads_output_to_converter_alignment(monkeypatch): from vllm import _sm70_ops as sm70_ops prepared = [] + physical_n = (logical_n + 31) // 32 * 32 def fake_prepare(qweight, scales, group_size, interleave_gated_silu): prepared.append( @@ -110,9 +113,9 @@ def fake_prepare(qweight, scales, group_size, interleave_gated_silu): ) ) return ( - torch.empty((32, 4), dtype=torch.int32), - torch.empty((2, 32), dtype=torch.float16), - torch.tensor([32, 32], dtype=torch.int64), + torch.empty((32, physical_n // 8), dtype=torch.int32), + torch.empty((2, physical_n), dtype=torch.float16), + torch.tensor([32, physical_n], dtype=torch.int64), ) monkeypatch.setattr(sm70_ops, "nvfp4_sm70_prepare", fake_prepare) @@ -120,8 +123,8 @@ def fake_prepare(qweight, scales, group_size, interleave_gated_silu): tm.prepare_nvfp4_linear(layer) state = getattr(layer, tm.STATE_ATTR) - assert prepared == [((32, 32), (2, 32), 16, False)] - assert state.output_size == 24 + assert prepared == [((32, physical_n), (2, physical_n), 16, False)] + assert state.output_size == logical_n def test_nvfp4_apply_crops_converter_padding(monkeypatch): diff --git a/vllm/model_executor/layers/quantization/sm70_turbomind.py b/vllm/model_executor/layers/quantization/sm70_turbomind.py index 0faeb07f6e..889a99873c 100644 --- a/vllm/model_executor/layers/quantization/sm70_turbomind.py +++ b/vllm/model_executor/layers/quantization/sm70_turbomind.py @@ -14,7 +14,9 @@ COMPRESSED_UINT4_GROUP_SIZES = (32, 128) MXFP4_GROUP_SIZE = 32 NVFP4_GROUP_SIZE = 16 -NVFP4_OUTPUT_ALIGNMENT = 16 +# SM70 packed NVFP4 GEMM needs complete 32-column tiles. N=8240 (Qwen +# GDN on TP2) is 16-aligned but corrupts the result without this padding. +NVFP4_OUTPUT_ALIGNMENT = 32 NVFP4_QPN4_DENSE_WORKSPACE_ELEMENTS = 5120 * 8704 STATE_ATTR = "_sm70_turbomind_linear" SM70QuantBackend = Literal["auto", "marlin", "turbomind"] diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 45bb294b32..a729ed4043 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -369,6 +369,16 @@ def maybe_prepare_sm70_lm_head_top1(layer: torch.nn.Module) -> bool: if not _is_sm70_lm_head_fastpath_eligible(layer): return False + tp_size = getattr(layer, "tp_size", 1) + if ( + envs.VLLM_SM70_DFLASH2_FP32_LOGITS + and tp_size in (2, 4) + and tuple(layer.weight.shape) == (248320 // tp_size, 5120) + ): + # Dense FP32 output must not depend on the TP4-only QPN8 candidate + # layout being available. TP2 otherwise silently keeps FP16 logits. + layer._sm70_dflash2_fp32_logits = True + raw_top1_requested = _sm70_env_bool( "VLLM_SM70_LM_HEAD_TOP1", _sm70_lm_head_top1_default() ) From 41a9018e9f715989b25b8f3c3ff436ffe57a13fe Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 14:01:26 +0800 Subject: [PATCH 6/9] [Audit] Trace QUASAR TP partition drift through target sampling Generated-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../benchmark_sm70_quasar_tp_quality.py | 343 ++++++++++++++++++ docs/design/sm70_quasar_tp2_tp4_quality.md | 230 ++++++++++++ docs/design/sm70_v100_migration_control.md | 33 ++ .../test_sm70_quasar_nvfp4_oracle.py | 14 +- .../test_sm70_fp32_lm_head_admission.py | 20 + .../layers/vocab_parallel_embedding.py | 1 + 6 files changed, 637 insertions(+), 4 deletions(-) create mode 100644 benchmarks/kernels/benchmark_sm70_quasar_tp_quality.py create mode 100644 docs/design/sm70_quasar_tp2_tp4_quality.md diff --git a/benchmarks/kernels/benchmark_sm70_quasar_tp_quality.py b/benchmarks/kernels/benchmark_sm70_quasar_tp_quality.py new file mode 100644 index 0000000000..247103c233 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_quasar_tp_quality.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Same real activation/weight inputs; isolate TP partition and kernel route. + +No distributed model rollout here. Row-parallel reductions use FP32 summation +followed by FP16, with partial outputs retained for a real-collective check. +""" + +import argparse +import importlib.util +import json +import sys +from pathlib import Path + +import torch +import torch.nn.functional as F + +from vllm import _sm70_ops +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + ref_nvfp4_quant_dequant, +) + +_spec = importlib.util.spec_from_file_location( + "quasar_tp_reference", + Path(__file__).with_name("benchmark_sm70_quasar_nvfp4_oracle.py"), +) +assert _spec is not None and _spec.loader is not None +ref = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = ref +_spec.loader.exec_module(ref) + + +def metrics(actual, expected): + a, e = actual.float(), expected.float() + d = a - e + return { + "max_abs": d.abs().max().item(), + "relative_l2": (d.norm() / e.norm().clamp_min(1e-30)).item(), + "max_row_relative_l2": (d.norm(dim=-1) / e.norm(dim=-1).clamp_min(1e-30)) + .max() + .item(), + "rmse": d.square().mean().sqrt().item(), + "unequal": int((a != e).sum()), + "elements": a.numel(), + "finite": bool(torch.isfinite(a).all()), + "reference_absmax": e.abs().max().item(), + } + + +def projections(cp, layer): + p = f"model.language_model.layers.{layer}" + if layer % 4 != 3: + yield ( + "gdn_qkvzba", + "linear_attn.in_proj_qkvz", + "column", + tuple( + p + ".linear_attn." + s + for s in ["in_proj_qkv", "in_proj_z", "in_proj_b", "in_proj_a"] + ), + [2048, 2048, 6144, 6144, 48, 48], + ) + yield ( + "gdn_out", + "linear_attn.out_proj", + "row", + (p + ".linear_attn.out_proj",), + [5120], + ) + else: + yield ( + "attention_qkv", + "self_attn.qkv_proj", + "column", + tuple(p + ".self_attn." + s for s in ["q_proj", "k_proj", "v_proj"]), + [12288, 1024, 1024], + ) + yield ( + "attention_out", + "self_attn.o_proj", + "row", + (p + ".self_attn.o_proj",), + [5120], + ) + yield ( + "mlp_gate_up", + "mlp.gate_up_proj", + "column", + (p + ".mlp.gate_proj", p + ".mlp.up_proj"), + [17408, 17408], + ) + yield "mlp_down", "mlp.down_proj", "row", (p + ".mlp.down_proj",), [5120] + + +def column_rows(widths, tp, rank): + parts, offset = [], 0 + for width in widths: + assert width % tp == 0 + parts.append( + torch.arange(offset + rank * width // tp, offset + (rank + 1) * width // tp) + ) + offset += width + return torch.cat(parts) + + +def native(pr, x, route, tm_alignment=16): + n, k = pr.packed.shape[0], pr.packed.shape[1] * 2 + if route == "tm": + pn = (n + tm_alignment - 1) // tm_alignment * tm_alignment + qw = ref._unpack_codes(pr.packed).cuda() + scales = ( + (pr.scales.t().float() / pr.weight_global_divisor) + .half() + .contiguous() + .cuda() + ) + if pn != n: + qw = F.pad(qw, (0, pn - n)) + scales = F.pad(scales, (0, pn - n)) + tw, ts, meta = _sm70_ops.nvfp4_sm70_prepare(qw, scales, 16, False) + out = torch.empty((x.shape[0], pn), device=x.device, dtype=x.dtype) + _sm70_ops.nvfp4_gemm_sm70_out( + out, x, tw, ts, 16, int(meta[0]), int(meta[1]), False + ) + y = out[:, :n].clone() + gated = None + if pr.name == "mlp_gate_up": + gated = torch.empty((x.shape[0], n // 2), device=x.device, dtype=x.dtype) + torch.ops._C.silu_and_mul(gated, y) + return y, gated + pn = (n + 31) // 32 * 32 + packed, scales = pr.packed.cuda(), pr.scales.cuda() + if pn != n: + p = torch.zeros((pn, packed.shape[1]), device=x.device, dtype=packed.dtype) + s = torch.zeros((pn, scales.shape[1]), device=x.device, dtype=scales.dtype) + p[:n].copy_(packed) + s[:n].copy_(scales) + packed, scales = p, s + codes, qs = _sm70_ops.nvfp4_qpn2_prepare_sm70(packed, scales) + split = 8 if pr.name in ["gdn_out", "attention_out", "mlp_gate_up"] else 16 + assert (k // 16) % split == 0 + out = torch.empty((x.shape[0], pn), device=x.device, dtype=x.dtype) + _sm70_ops.nvfp4_qpn2_gemm_sm70_out( + out, x, codes, qs, 1 / pr.weight_global_divisor, split, 2 + ) + y = out[:, :n].clone() + gated = None + if pr.name == "mlp_gate_up": + assert pn == n + gated = torch.empty((x.shape[0], n // 2), device=x.device, dtype=x.dtype) + _sm70_ops.nvfp4_qpn2_gated_sm70_out( + gated, x, codes, qs, 1 / pr.weight_global_divisor, split, 2 + ) + return y, gated + + +@torch.inference_mode() +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", type=Path, required=True) + ap.add_argument("--capture-dir", type=Path, required=True) + ap.add_argument("--step", type=int, default=7) + ap.add_argument("--layers", nargs="+", type=int, default=list(range(64))) + ap.add_argument("--out", type=Path, required=True) + ap.add_argument("--tm-alignment", type=int, choices=[16, 32], default=32) + args = ap.parse_args() + torch.set_num_threads(1) + torch.backends.cuda.matmul.allow_tf32 = False + cp = ref.Checkpoint(args.model) + tables = [] + positions = None + input_ids = None + for rank in range(4): + d = torch.load( + args.capture_dir / f"rank{rank}-step{args.step}.pt", weights_only=True + ) + if positions is None: + positions, input_ids = d["positions"], d["input_ids"] + assert torch.equal(positions, d["positions"]) and torch.equal( + input_ids, d["input_ids"] + ) + tables.append( + {(t["layer_idx"], t["label"]): t["tensor"] for t in d["tensors"].values()} + ) + rows, partials = [], {} + for layer in args.layers: + for name, suffix, parallel, prefixes, widths in projections(cp, layer): + tag = f"quant_input:language_model.model.layers.{layer}.{suffix}" + ( + ":gated" if name == "mlp_gate_up" else ":linear" + ) + xs = [t[layer, tag] for t in tables] + if parallel == "column": + # Captured replicated tensors can already differ by rounding. + # Fix rank-zero input for every shard in this operator oracle. + fullx = xs[0].cuda() + full = ref._load_column_parallel(cp, name, prefixes, 0, 1) + else: + fullx = torch.cat(xs, dim=-1).cuda() + full = ref._load_row_parallel(cp, name, prefixes[0], 0, 1) + weight = ref._dequantize_weight(full, torch.device("cuda")) + assert tuple(weight.shape) == (sum(widths), fullx.shape[-1]) + oracle = fullx.float() @ weight.t() + full_half = fullx.float() @ weight.half().float().t() + qat_x = ref_nvfp4_quant_dequant( + fullx, + torch.tensor( + full.input_global_divisor, device=fullx.device, dtype=torch.float32 + ), + 16, + ) + qat_out = qat_x.float() @ weight.t() + assembled, gated_all = {}, {} + part_sets = {} + row = { + "layer": layer, + "operator": name, + "parallel": parallel, + "input_shape": list(fullx.shape), + "weight_shape": list(weight.shape), + "fp16_materialized_weight_vs_checkpoint": metrics( + weight.half(), weight + ), + "single_dense_fp16_weight_vs_checkpoint_output": metrics( + full_half, oracle + ), + "local_shapes": {}, + } + if parallel == "column": + row["captured_rank_inputs_vs_rank0"] = [metrics(x, xs[0]) for x in xs] + row["qat_w4a4_activation_vs_w4a16"] = metrics(qat_x, fullx) + row["qat_w4a4_output_vs_w4a16"] = metrics(qat_out, oracle) + for tp in [2, 4]: + products = {"tm": [], "qpn2": [], "fp32": [], "fp16": []} + gates = {"tm": [], "qpn2": []} + for rank in range(tp): + if parallel == "column": + pr = ref._load_column_parallel(cp, name, prefixes, rank, tp) + ids = column_rows(widths, tp, rank) + assert torch.equal(pr.packed, full.packed[ids]) + assert torch.equal( + pr.scales.view(torch.uint8), + full.scales.view(torch.uint8)[ids], + ) + x = fullx + local_ref = oracle[:, ids.cuda()] + else: + pr = ref._load_row_parallel(cp, name, prefixes[0], rank, tp) + sl = slice( + rank * fullx.shape[-1] // tp, + (rank + 1) * fullx.shape[-1] // tp, + ) + assert torch.equal( + pr.packed, + full.packed[:, slice(sl.start // 2, sl.stop // 2)], + ) + x = fullx[:, sl].contiguous() + local_ref = x.float() @ weight[:, sl].t() + row["local_shapes"][f"tp{tp}"] = { + "M": x.shape[0], + "N": pr.packed.shape[0], + "K": x.shape[-1], + } + products["fp32"].append(local_ref) + products["fp16"].append(local_ref.half()) + for route in ["tm", "qpn2"]: + y, g = native(pr, x, route, tm_alignment=args.tm_alignment) + products[route].append(y) + if g is not None: + gates[route].append(g) + for route, ys in products.items(): + if parallel == "row": + total = torch.stack([y.float() for y in ys]).sum(0) + assembled[f"tp{tp}_{route}"] = ( + total if route == "fp32" else total.half() + ) + if route in ["tm", "qpn2"]: + part_sets[f"tp{tp}_{route}"] = torch.stack(ys).cpu() + else: + total = torch.empty_like(oracle) + for rank, y in enumerate(ys): + total[:, column_rows(widths, tp, rank).cuda()] = y.float() + assembled[f"tp{tp}_{route}"] = total + for route, gs in gates.items(): + if gs: + gated_all[f"tp{tp}_{route}"] = torch.cat(gs, dim=-1) + row["vs_checkpoint_fp32"] = { + k: metrics(v, oracle) for k, v in assembled.items() + } + row["tp2_vs_tp4"] = { + route: metrics(assembled["tp4_" + route], assembled["tp2_" + route]) + for route in ["tm", "qpn2", "fp32", "fp16"] + } + row["production_tp4_qpn2_vs_tp2_tm"] = metrics( + assembled["tp4_qpn2"], assembled["tp2_tm"] + ) + row["tp4_qpn2_vs_tm"] = metrics(assembled["tp4_qpn2"], assembled["tp4_tm"]) + if gated_all: + gate, up = oracle.chunk(2, -1) + expected = F.silu(gate) * up + row["gated_vs_fp32"] = { + k: metrics(v, expected) for k, v in gated_all.items() + } + row["gated_tp2_vs_tp4"] = { + route: metrics(gated_all["tp4_" + route], gated_all["tp2_" + route]) + for route in ["tm", "qpn2"] + } + row["gated_production_tp4_qpn2_vs_tp2_tm"] = metrics( + gated_all["tp4_qpn2"], gated_all["tp2_tm"] + ) + partials[f"{layer}:{name}"] = part_sets + rows.append(row) + print( + layer, + name, + "TP delta", + row["production_tp4_qpn2_vs_tp2_tm"]["relative_l2"], + "same TM", + row["tp2_vs_tp4"]["tm"]["relative_l2"], + flush=True, + ) + args.out.write_text( + json.dumps( + { + "model": str(args.model), + "activation_capture": str(args.capture_dir), + "step": args.step, + "tm_alignment": args.tm_alignment, + "note": ( + "TP2 QPN2 is a kernel counterfactual, not production " + "dispatch. FP32 summation emulates row collective; check " + "collectives separately. Reference uses checkpoint " + "dequantization, not the BF16 teacher." + ), + "rows": rows, + }, + indent=2, + ) + ) + torch.save(partials, args.out.with_suffix(".partials.pt")) + + +if __name__ == "__main__": + main() diff --git a/docs/design/sm70_quasar_tp2_tp4_quality.md b/docs/design/sm70_quasar_tp2_tp4_quality.md new file mode 100644 index 0000000000..a88689080b --- /dev/null +++ b/docs/design/sm70_quasar_tp2_tp4_quality.md @@ -0,0 +1,230 @@ +# QUASAR QAT: TP2/TP4 arithmetic and distribution audit + +## Findings + +TP size changes both matrix partitions and the current SM70 execution route. +The audit found and repaired two TP2 defects, then demonstrated remaining +sampling-token changes on identical prefixes after those repairs. Greedy top-1 +agreement does not establish distribution agreement. + +1. **NVFP4 output alignment:** the TP2 GDN input projection has logical N=8240. + It satisfies the old 16-column adapter alignment but corrupts the native + packed TurboMind result. Padding to N=8256 (32 columns) restores accuracy. + Eight real shards across layers 0, 1, 32 and 62 had relative L2 + **0.326777–0.512499 before, 0.000249–0.000323 after**. TP4's existing + 4120→4128 physical layout is unchanged. +2. **FP32 LM-head admission:** the precision flag previously took effect only + while preparing the TP4 QPN8 candidate layout. TP2 fell back to FP16 dense + logits even when the flag was set. The explicit FP32 flag now independently + admits the pinned vocabulary/hidden shape on TP2 and TP4. QPN8 admission + remains TP4-only; no global precision default changes. +3. **Remaining TP drift:** after fixing both defects, two same-prefix probes + retained all 258 greedy top-1 IDs, but target sampling distributions reached + **1.029% / 2.208% maximum TV**. Production token-keyed Gumbel noise produced + **21 differences in 16,512 paired draws**, including one with seed 0. +4. **QAT semantics:** the checkpoint describes W4A4, while the SM70 route uses + W4A16. Applying its activation quantizer on the same captured inputs changes + individual projection outputs by up to **12.24% relative L2**. Changing only + the last MLP down projection to that activation contract produces up to + **5.042% sampling TV** downstream. These are execution differences, not + percentages of text-quality loss or proof that W4A4 is better on this host. + +The audit does not establish overall QAT quality relative to its BF16 teacher: +that exact teacher checkpoint is unavailable locally. The operator reference +decodes the *QAT checkpoint's* weights; it is not an unquantized-model oracle. + +## Frozen inputs and source + +- Integration: `onecat/main`, `755baae1d075ee04fa9096b23fc0225b23589a86`. +- Starting precision/performance branch: `6f07be1cce338b87e7e8d1a714bf9e30e5424fd8`. +- Production-code repairs: `fcb6dada58` in owned Draft PR #517. +- QUASAR checkpoint revision: `d8e6fbfa3e3a78899b440222b827430045a05b44`. +- DFlash2 revision: `dedf8df68adfb1afeaf7b7480c0a0243108177b4`. +- V100-SXM2-32GB; TP4 uses GPUs 0–3, TP2 uses GPUs 0–1. CUDA 12.8, + Torch 2.10.0+cu128, Python 3.12.13, Triton 3.6.0, FP16 activations. +- Target/draft FLASH_ATTN_V100, target E4M3 KV, draft FP16 KV; V2 runner; + draft7/q8; full/piecewise target and full draft graphs; context pipeline on. +- Max length 262144, batch tokens 4096, capacity four requests, one live + request, memory utilization 0.8, prefix cache and Mamba align enabled. + +The matrix audit reuses alias-corrected C2 captures at step 7 from the preceding +operator audit. Their trajectories were generated with E5M2, but the exact same +input vector is supplied to each compared TP partition. Column projections +use rank-zero input on every shard; row inputs concatenate all four logical +shards. Logical Q/K/V and gate/up segments are sharded independently, and +packed checkpoint bytes are checked against the full projection. + +The end-to-end probes use fresh E4M3 services. Token tapes come from the prior +natural-stop production MBPP28 and MBPP3 outputs. Each comparison covers the +prefill's last position plus 16 q8 verification groups: 129 rows per input. +Every compared position, token ID and group boundary is checked. The runner +forces the tape and accepts its query rows, preventing an early sample change +from feeding different prefixes to the two models. These requests are capped +diagnostics, not free generation, quality scores or performance benchmarks. + +## Matrix dimensions and execution + +Dimensions below are local logical `(N, K)` for `X[M,K] @ W[N,K].T`, with M=8 +for the decode operator comparison. + +| Projection | TP2 `(N,K)` | TP4 `(N,K)` | Partition | +| --- | --- | --- | --- | +| GDN qkvzba | 8240, 5120 | 4120, 5120 | Output rows, logical segments | +| Attention qkv | 7168, 5120 | 3584, 5120 | Output rows, logical segments | +| MLP gate/up | 17408, 5120 | 8704, 5120 | Output rows, gate/up separately | +| GDN/attention output | 5120, 3072 | 5120, 1536 | Input columns and sum | +| MLP down | 5120, 8704 | 5120, 4352 | Input columns and sum | +| LM-head | 124160, 5120 | 62080, 5120 | Vocabulary rows | + +TP4's admitted q8 target projections use QPN2, with fixed split-K/accumulator +settings. TP2 uses TurboMind. QPN2 internally uses FP32 HMMA accumulators and +FP32 partial reduction, then writes FP16. Consequently, row-parallel layers +round each local partial before the cross-rank sum. An accurate collective +cannot recover those discarded bits. + +The comparison therefore includes both the production route difference +(TP2 TurboMind versus TP4 QPN2) and an isolated partition test using the same +kernel family on both sides. Direct QPN2 calls on TP2 are counterfactual +operator probes; this change does not enable QPN2 for TP2 production. + +## Operator results + +All 64 layers were inspected: 256 fused projection invocations, covering the +checkpoint's 496 logical linear layers. The table reports maximum relative L2 +between the two TP assemblies on identical input vectors, after alignment repair. + +| Projection | Cases | Same TurboMind family | Same QPN2 configuration rule | TP2 TM vs TP4 QPN2 | +| --- | ---: | ---: | ---: | ---: | +| GDN qkvzba | 48 | 0.00323% | 0 | 0.06678% | +| Attention qkv | 16 | 0.00242% | 0 | 0.06513% | +| MLP gate/up, before activation | 64 | 0.00182% | 0 | 0.06416% | +| GDN output | 48 | 0.04404% | 0.04775% | 0.08108% | +| Attention output | 16 | 0.04654% | 0.04564% | 0.07763% | +| MLP down | 64 | 0.04423% | 0.04420% | 0.07423% | + +The fused gate/SILU/up output reaches 0.11928% between production routes. +Changing only the rounding of exact FP32 local partials already causes up to +0.05107% TP2/TP4 difference in MLP down. Thus there is an arithmetic mechanism +independent of a broken collective or changed checkpoint shards. + +Real partial tensors from three row projections, two kernel routes and both +TP sizes were also passed through native captured collectives. All six cases +per rank matched FP32 summation followed by FP16 bitwise. This validates the +assembly reference on those inputs; it does not prove every collective shape. + +Older C2 replicated norm inputs contained a few rank-dependent rounded values. +The matrix oracle explicitly fixes those inputs instead of comparing different +vectors. In the fresh end-to-end probes, replicated final hidden states are +bitwise identical among ranks within each TP configuration on all compared rows. + +## Distribution and sampled-token changes + +Both end-to-end services use the repaired FP32 dense LM-head for observation. +An independent FP32 dot-product oracle is also computed per vocabulary shard. +The head-only comparison evaluates the old FP16 path on the *same* hidden states. + +| Check | MBPP28, 129 positions | MBPP3, 129 positions | +| --- | ---: | ---: | +| Final hidden relative L2, TP2 vs TP4 | 0.9576% | 0.9778% | +| Greedy top-1 flips, TP2 vs TP4 | 0 | 0 | +| Changed top-p support, TP2 vs TP4 | 0 | 1 | +| Maximum sampling TV, FP32 oracle TP2 vs TP4 | 1.0291% | 2.2077% | +| Old TP2 FP16 head vs same-hidden FP32: maximum TV | 2.5391% | 1.4526% | +| Repaired TP2 head vs same-hidden FP32: maximum TV | 0.001708% | 0.002081% | +| Changed head-only top-p supports, old → repaired | 1 → 0 | 1 → 0 | + +The first decode group's layer-zero input matches exactly between TP sizes. +Its MLP output already differs by about 0.16%; differences then propagate +through attention, recurrent state, norms and subsequent projections. The +final-state difference is not attributable solely to one matrix or solely to +TP count: production kernel families and shapes both differ. + +Sampling uses T=1, k=20 and p=0.95 with seeds fixed in advance to 0–63. The +production `gumbel_noised_argmax` primitive keys noise by seed, absolute +position and token ID. Sparse support enumeration is checked against the +full-vocabulary `gumbel_sample` at seed 0 before counting flips. + +| Paired draws | MBPP28 / 8256 | MBPP3 / 8256 | +| --- | ---: | ---: | +| TP2 vs TP4 FP32 distributions: different sampled IDs | 10 | 11 | +| Old TP2 FP16 head vs FP32 oracle | 5 | 5 | +| Repaired TP2 FP32 head vs FP32 oracle | 0 | 0 | + +For example, at MBPP28 position 167 and seed 0, TP2 and TP4 choose token IDs +13 and 1132 respectively under the same target-sampling noise. This is a +reproduced sampling flip with unchanged greedy top-1. It is not a full +DFlash rejection/acceptance trace or a claim about free-running answer quality. + +## QAT execution contract + +The pinned model card and `quantization_config` declare W4A4 NVFP4 with +16-element activation groups, E4M3 group scales and a checkpoint global scale. +The accepted SM70 TurboMind/QPN2 route retains quantized weights but consumes +FP16 activations; it does not apply that activation quantizer. + +The audit applies the repository's NVFP4 activation quantize/dequantize +reference using the stored input global divisor, then multiplies by the same +FP32-dequantized QAT weights. Maximum output differences are 6.88% for GDN +input, 8.90% for GDN output, 5.96% for attention qkv, 12.24% for attention +output, 9.11% for MLP gate/up and 9.01% for MLP down. + +To trace this through logits without replacing every upstream state, a second +probe replaces only layer 63's MLP down result on the same captured input and +residual, then runs the final norm and FP32 LM-head. Across 24 rows at three +captured steps, top-1 stays unchanged, one top-p support changes, and maximum +TV reaches 5.042%. By contrast, changing only that projection between repaired +TP2 TurboMind and TP4 QPN2 reaches 0.04112% maximum TV in these probes. + +This evidence supports auditing the QAT activation contract separately from +ordinary GEMM rounding. It does **not** justify automatically enabling costly +W4A4 emulation, assuming higher activation precision always improves this +trained model, or quoting the model card's native-FP4 scores for the V100 route. +No new BF16-teacher comparison or end-to-end W4A4 quality score was obtained. + +## Validation, reproducibility and limits + +- Native padding and adapter regressions: **19 passed**, including N=8240, + M=1/8/17 and changed-input CUDA Graph replay. +- FP32 head admission: **6 passed**. Logical GDN TP2/TP4 sharding: **6 passed**. +- The tracked matrix-audit CLI ran successfully against real layer-zero data; + full 64-layer results were produced by its retained task-script precursor. +- Ruff, mypy and applicable commit hooks pass for the fixes. +- No native source or native binary changed. Existing TP4 physical widths, + QPN2 accumulation and FP32 head execution stay the same. The preceding + 18.892/18.435 ms result is historical evidence; no new latency claim is made. + +```bash +PYTHONPATH=. .venv/bin/python benchmarks/kernels/benchmark_sm70_quasar_tp_quality.py \ + --model --capture-dir \ + --step 7 --tm-alignment 32 --out +``` + +The old alignment can be selected in this diagnostic CLI with +`--tm-alignment 16`; it is not a serving rollback recommendation. + +One TP4 diagnostic service stalled during a request transition. The initial +hook reread a global active-case file during every sampler call; the async +runner can still have a round in flight after the API reaches its output cap. +Removing/switching that file can send ranks through different diagnostic +collectives. The retained hook now binds the tape to request ID. The accepted +comparisons validate the complete common prefill/q8 groups and exclude extra +in-flight groups; MBPP3 TP4 was collected in a separate startup. The revised +request-binding hook has not received a fresh multi-request endpoint gate. +This diagnostic failure is not classified as a production model defect. + +The first standalone TP2 collective probe lacked a current `VllmConfig` and +failed before validation; a configured rerun passed. A debugger attach was +unavailable on this host. Failed paths and task-script revisions are retained. + +Artifacts: `v100-quasar-tp2-tp4-quality-audit-20260906`, including +`operator-all-layers-fixed.json`, `padding-probe.json`, +`teacher-comparison.json`, `sampling-flips.json`, +`last-layer-propagation-summary.json`, actual module manifests, native/JIT +inventory, exact token tapes, launch scripts and logs. Results are under its +`results/` directory; large captures remain external to Git. Task services and +GPU leases have been released. + +The remaining quality gate is to compare the pinned QAT W4A4 contract and the +V100 W4A16 path against the exact BF16 teacher on shared prefixes, then test +whether delaying row-partial FP16 rounding reduces the measured drift without +losing the established round-time budget. Neither is declared solved here. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 0421777367..f2c7c7dfdd 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45821,3 +45821,36 @@ Interpretation: - Full contract, exact tests, acceptance limits, and artifact bundle are in [the recovery report](sm70_quasar_dflash2_quality_speed_recovery.md). Raw bundle: `v100-quasar-quality-speed-recovery-20260906`. + +## 2026-09-06 QUASAR TP2/TP4 quality and QAT execution audit + +- Continue owned Draft PR #517 from `6f07be1cce`, integration `755baae1d0`. + Repairs at `fcb6dada58`: align TurboMind NVFP4 physical output to 32 columns + and independently honor explicit FP32 dense LM-head output on TP2. +- TP2 GDN N=8240 was 16-aligned but corrupt. Eight real shards show relative + L2 32.68%–51.25%; padding to 8256 restores 0.0249%–0.0323%. TP4's 4128 + physical width is unchanged. Do not describe this as ordinary rounding. +- All 64 target layers / 256 fused projections tested on identical C2 inputs. + Same-family TP changes affect row projections through local FP16 rounding; + native captured TP2/TP4 collectives match FP32 sum then FP16 on the six + real-partial cases per rank. QPN2 is still a TP4 production route. +- Fresh E4M3 services, two fixed-prefix tapes, 129 positions each: TP2/TP4 + greedy IDs all match, maximum sampling TV 1.029% / 2.208%; 21 differing + target Gumbel draws across 16,512 paired seed/position tests. Final hidden + relative L2 is 0.958% / 0.978%. This is not full speculative sampling or + a free-generation score, and includes the production kernel-family change. +- The TP2 head-only repair removes both changed nucleus supports and reduces + differing Gumbel draws from 10 to 0 against the same-hidden FP32 oracle. +- QAT declares W4A4; SM70 currently executes W4A16. Same-input activation + quantization changes projection outputs up to 12.24% relative L2; replacing + only the final down projection yields up to 5.042% sampling TV. No BF16 + teacher or end-to-end W4A4 quality comparison; do not claim either is better. +- Tests: 19 native/adapter, 6 head admission, 6 logical shard tests passed. + Matrix CLI and actual target-prefix/production Gumbel probes completed. +- Diagnostic request switching used a global file and could race in-flight + rounds after the API cap. Bind diagnostics to request IDs; exclude extra + rounds and validate exact input positions. TP4 MBPP3 used a separate startup. + The revised hook has no fresh multi-request endpoint validation yet. +- No new speed claim. Full evidence, reproducible matrix command and remaining + gates: [TP quality audit](sm70_quasar_tp2_tp4_quality.md). Bundle: + `v100-quasar-tp2-tp4-quality-audit-20260906`. Owned GPU services/leases stopped. diff --git a/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py b/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py index 09faba70f3..c89714244e 100644 --- a/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py +++ b/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py @@ -12,8 +12,8 @@ ) -@pytest.mark.parametrize("rank", range(4)) -def test_gdn_qkv_shards_each_logical_projection(tmp_path, rank): +@pytest.mark.parametrize("tp,rank", [(tp, rank) for tp in (2, 4) for rank in range(tp)]) +def test_gdn_qkv_shards_each_logical_projection(tmp_path, tp, rank): (tmp_path / "config.json").write_text( json.dumps( { @@ -36,7 +36,13 @@ def test_gdn_qkv_shards_each_logical_projection(tmp_path, rank): prefix + ".input_global_scale": torch.tensor(3.0), } checkpoint = SimpleNamespace(model=tmp_path, tensor=tensors.__getitem__) - projection = _load_column_parallel(checkpoint, "gdn_qkv", (prefix,), rank, 4) - expected_rows = torch.tensor([rank, 4 + rank, *range(8 + 3 * rank, 11 + 3 * rank)]) + projection = _load_column_parallel(checkpoint, "gdn_qkv", (prefix,), rank, tp) + expected_rows = torch.tensor( + [ + *range(rank * 4 // tp, (rank + 1) * 4 // tp), + *range(4 + rank * 4 // tp, 4 + (rank + 1) * 4 // tp), + *range(8 + rank * 12 // tp, 8 + (rank + 1) * 12 // tp), + ] + ) assert torch.equal(projection.packed, packed[expected_rows]) assert torch.equal(projection.scales, scales[expected_rows]) diff --git a/tests/quantization/test_sm70_fp32_lm_head_admission.py b/tests/quantization/test_sm70_fp32_lm_head_admission.py index 71e2143a0e..c74b68ed68 100644 --- a/tests/quantization/test_sm70_fp32_lm_head_admission.py +++ b/tests/quantization/test_sm70_fp32_lm_head_admission.py @@ -9,6 +9,26 @@ from vllm.model_executor.layers import vocab_parallel_embedding as vocab +@pytest.mark.parametrize("enabled", [False, True]) +def test_explicit_fp32_flag_independent_of_other_head_fastpaths(monkeypatch, enabled): + layer = SimpleNamespace( + prefix="language_model.lm_head", + weight=SimpleNamespace( + dtype=torch.float16, + is_cuda=True, + device=torch.device("cuda", 0), + ndim=2, + shape=(124160, 5120), + ), + ) + monkeypatch.setattr(vocab, "_sm70_env_bool", lambda *args: False) + monkeypatch.setattr(vocab, "_sm70_dflash2_qpn8_rerank_requested", lambda: False) + monkeypatch.setattr(vocab.current_platform, "is_cuda_alike", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _: (7, 0)) + monkeypatch.setattr(vocab.envs, "VLLM_SM70_DFLASH2_FP32_LOGITS", enabled) + assert vocab._is_sm70_lm_head_fastpath_eligible(layer) == enabled + + @pytest.mark.parametrize("tp", [2, 4]) @pytest.mark.parametrize("enabled", [False, True]) def test_fp32_head_admission_without_qpn8_layout(monkeypatch, tp, enabled): diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index a729ed4043..09c0ca797e 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -107,6 +107,7 @@ def _is_sm70_lm_head_fastpath_eligible(layer: torch.nn.Module) -> bool: _sm70_env_bool("VLLM_SM70_ENABLE_LM_HEAD_FASTPATH", False) or _sm70_env_bool("VLLM_SM70_LM_HEAD_TOP1", _sm70_lm_head_top1_default()) or _sm70_env_bool("VLLM_SM70_LM_HEAD_TOP1_TC", False) + or envs.VLLM_SM70_DFLASH2_FP32_LOGITS or _sm70_dflash2_qpn8_rerank_requested() ): _trace_sm70_lm_head_skip("disabled") From 53be620005fb0f7664dcd27a512979191b342c73 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 14:36:52 +0800 Subject: [PATCH 7/9] [Bugfix] Preserve proposal probabilities in DFlash2 lookup fusion Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_dflash2_fastpath_numerics.md | 198 ++++++++++++++++++ docs/design/sm70_dflash2_lookup_augmented.md | 8 +- docs/design/sm70_v100_migration_control.md | 33 +++ tests/v1/spec_decode/test_dflash2_lookup.py | 144 +++++++++++++ .../worker/gpu/spec_decode/dflash2/lookup.py | 10 + .../gpu/spec_decode/dflash2/speculator.py | 1 + 6 files changed, 393 insertions(+), 1 deletion(-) create mode 100644 docs/design/sm70_dflash2_fastpath_numerics.md diff --git a/docs/design/sm70_dflash2_fastpath_numerics.md b/docs/design/sm70_dflash2_fastpath_numerics.md new file mode 100644 index 0000000000..60d7e57263 --- /dev/null +++ b/docs/design/sm70_dflash2_fastpath_numerics.md @@ -0,0 +1,198 @@ +# DFlash2 fast-path numerical audit, 2026-09-06 + +## Scope and frozen execution + +Continue Draft PR #517 from `41a9018e9f715989b25b8f3c3ff436ffe57a13fe`, +integration `onecat/main` at `755baae1d075ee04fa9096b23fc0225b23589a86`. +The production change in this audit only repairs probability bookkeeping in +the optional lookup path. Target GEMMs, attention, KV formats, selector +arithmetic and the context-pipeline implementation are unchanged. + +Target: QUASAR Qwen3.8-27B all-NVFP4 checkpoint revision +`d8e6fbfa3e3a78899b440222b827430045a05b44`, executing W4A16 on SM70. +Draft: DFlash2 revision `dedf8df68adfb1afeaf7b7480c0a0243108177b4`, five +layers, hidden 5120, block8/draft7, selector rank 256/top16. The environment +is V100-SXM2-32GB GPUs 0–3, TP4, CUDA 12.8, Torch 2.10.0+cu128, +Python 3.12.13, Triton 3.6.0. Target KV is E4M3; draft KV is FP16. +`VLLM_SM70_DFLASH2_FP32_LOGITS=1` preserves the previous head precision fix. + +The real request uses MBPP28, input 135, temperature 1, top-k 20, top-p 0.95, +seed 0, natural EOS, cap 512, thinking/xhigh. Model length 262144, batch +budget 4096, maximum requests 4, one active request, memory 0.8, prefix cache +and Mamba align. The V2 runner uses full/piecewise target graphs and full +draft graphs. Worker logs confirm Flash-V100 E4M3 q8 grouped verification, +FP32 candidate rerank, sparse rejection, BF16 draft emulation and TP4 +output-sharded context FC. The paired context-graph flags are recorded with +each launch. Numerical capture synchronizes CPU copies and is excluded from +latency evidence. + +## Confirmed lookup probability defect and repair + +Affected configuration: probabilistic lookup-augmented drafting, +`ngram_assist=true`, verifier width greater than the trained draft width, +`VLLM_DFLASH2_LOOKUP_AGREE > 0`, and a match shorter than +`VLLM_DFLASH2_LOOKUP_NSTRONG` that qualifies through neural-prefix agreement. +The standard q8 request does not enable lookup. The default agreement +threshold 0 also does not exhibit this defect. + +Previously, the lookup fusion decided whether to replace proposals using +the just-sampled neural prefix, then rewrote every replaced row's proposal +distribution as a point mass. The agreeing prefix already contained those +tokens, but changing its probability from the neural q to 1 invalidated the +acceptance ratio. This is an algorithmic distribution error, not FP16 noise. + +A two-token counterexample uses q(A)=q(B)=0.5 and p(A)=0.8. A weak historical +continuation begins with A and requires one agreeing proposal. When the +draft draws A, the old code records q(A)=1; when it draws B, q(B) remains 0.5. +The resulting output probability for A is +`0.5 * 0.8 + 0.5 * (1 - 0.2 / 0.5) = 0.7`, rather than 0.8. + +The repair preserves the original q for the agreeing random prefix. Only +positions after that prefix become point masses conditional on the already +sampled prefix. Strong matches can still replace the whole block because +their decision depends on request history. Proposed token sequences remain +identical; the correction probabilities change where necessary. + +Actual production lookup, point-mass and dense-rejection kernels, 100000 +independent seeds, show: + +| Proposal policy | Observed P(A) | Target P(A) | +| --- | ---: | ---: | +| History-only lookup control | 0.799110 | 0.8 | +| Agreement-conditioned lookup, old bookkeeping | 0.700270 | 0.8 | +| Agreement-conditioned lookup, repaired bookkeeping | 0.800720 | 0.8 | + +The committed regression repeats the repaired experiment for both dense and +sparse rejection. It also verifies strong/weak matches, greedy/probabilistic +fusion, unchanged proposed tokens, and the prefix's point-mass mask. All 12 +tests in `tests/v1/spec_decode/test_dflash2_lookup.py` pass on V100. + +## Real selector arithmetic and sampling + +Ten actual request snapshots contain 70 conditional proposal rows. Startup +`_warmup_*` requests are excluded. The FP32 arithmetic reference uses the +same captured backbone states, FP16-materialized checkpoint codebooks and +hidden-projection weights, and the same captured FP32 unary logits. It +changes accumulation/intermediate rounding within the selector, not the +candidate support, LM-head or draft backbone. It is not a BF16 teacher. + +| Comparison | Maximum or count | +| --- | ---: | +| Compiled selector vs eager, same arithmetic | Bitwise equal | +| Selector lattice vs FP32 arithmetic, max absolute | 0.0117207 | +| Bilinear edges vs FP32 arithmetic, relative L2 | 0.000362991 | +| Conditional proposal TV at temperature 1 | 0.000953654 (0.0953654%) | +| Conditional greedy top-1 changes | 0 / 70 | +| Temperature1, seeds 0–255, FP32 counterfactual draft-token changes | 3 / 17920 | +| Temperature0.6 diagnostic, seeds 0–255, counterfactual changes | 10 / 17920 | + +The 13 differing draft positions occur in five proposal chains; some are +downstream consequences of an earlier selector change. Seed0 does not flip +in these samples. These are changes in the draft proposal, not 13 final +target-token errors. They can affect acceptance and latency. Correct +rejection sampling must correct the actual q drawn by the selector; a draft +FP32 conversion alone is not evidence of improved target output quality. + +Separately, the existing fused prefix/tail selector and its persistent CUDA +Graph were compared against sequential full-vocabulary Gumbel sampling, +using actual candidate IDs and positions, seeds 0–63, temperatures1 and0.6. +All 8960 proposal positions match. Realized scores, request-slot sparse cache +and their dense-cache entries match exactly, including temperature scaling. +The seed 0/temperature 1 replay also matches the captured model proposal. + +An additional 96-round cache-overwrite probe alternates eager/graph launches, +intersecting and reordered supports, permuted request slots, an invalid +request row, and padded dense strides. All 288 checked request-slot outputs +match an independent dense scatter; padding remains intact. No cache-scatter +ordering defect was reproduced. This finite probe is not a proof for every +supported shape or concurrent request transition. + +## Context FC partition arithmetic + +The FC projects concatenated target layers `[5,19,33,47,61]` from 25600 to 5120. +At M8, TP4 changes local output width to 1280 and uses TurboMind plus an +all-gather. It does not split the K reduction. Nevertheless, the changed N +can change the GEMM schedule and FP32 accumulation order before FP16 storage. + +Three fresh decode captures (snapshot indices 3,7,11) reproduce the observed +TP4 projection bitwise from independently packed local checkpoint shards. +The following are local arithmetic comparisons on identical full inputs; +the TP2 variant is a counterfactual output partition, not a TP2 endpoint. + +| Comparison | Maximum relative L2 | +| --- | ---: | +| TP4 TurboMind vs FP32 matrix multiplication | 0.000220945 | +| TP4 TurboMind vs replicated cuBLAS FP16 | 0.0000177354 | +| Same comparison after BF16 context normalization | 0.000105950 | +| TurboMind TP2 vs TP4 output partitions | 0.00000921025 | + +TP2/TP4 maximum absolute difference is 0.0625. BF16 rounding can amplify a +small FC difference around a rounding boundary. Changed-input graph replay +matches eager bitwise over 24 replays. Prior C2 snapshots independently show +the same shape-dependent effect; the current table uses fresh E4M3 data. +No end-to-end target-quality attribution to this FC difference is established. + +## Cost and quality boundaries + +Two fresh services differ only in enabling +`VLLM_SM70_DFLASH2_CONTEXT_PIPELINE=1` and +`VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH=1`. Final worker logs confirm the staged +context computation, deferred writes and B1 metadata refresh graph. Across +all ten paired real-request snapshots, target/auxiliary states, FC outputs, +draft inputs and backbone states, candidates, unary/lattice scores, proposal +tokens, realized/cached q scores and acceptance/rejection counts are bitwise +equal. The snapshots include partial acceptance (2, 3, 4, 6 and 7 emitted +tokens) and full q8 emission. Both services produce the same complete +260-token output, 49 rounds and 211 accepted draft tokens. This extends the +earlier isolated context-graph check through the real selector and sampler +boundary. It remains a short B1 request, not a long-context or request-reuse +proof. + +The lookup repair adds no kernel launch and affects no arithmetic in the +ordinary q8 baseline. A captured fusion-kernel benchmark uses 256 nodes per +graph, eight replays per observation and seven observations. B1 medians are +1.735–1.736 us before and 1.741–1.742 us after the repair; B4 medians are +1.818–1.829 us before and 1.818–1.834 us after. The largest paired increase +is 0.006 us. This is only fusion-kernel cost, not end-to-end round latency. + +The fresh baseline request naturally emits 260 tokens over 49 draft rounds, +with 211 accepted draft tokens. Its complete token hash matches the previous +precision/recovery control: +`78648509da0a573cb79264412ffd477af13cf1be615019ab573c4975ad4ec908`. +This single code request is an execution check, not a broad quality score. + +Historical unprofiled E4M3/FP32 medians remain 18.435 ms for MBPP28 and +18.892 ms for release1k, at `f22ac115d0`. This audit does not claim recovery +to 17.6–18 ms. No new long-context endpoint, QAT teacher comparison, or +acceptance-rate improvement from FP32 selector arithmetic is claimed. + +## Reproduction and retained artifacts + +Owned worktree: +`/home/ymzx/桌面/1cat-vllm/worktrees/v100-quasar-dflash2-operator-audit-20260905-172402`. +Raw bundle: +`/data/minimax-h3/task-cache/v100-dflash2-fastpath-numerics-20260906`. +The bundle contains exact job commands/environment in `queue/*.done.json`, +launchers and analysis scripts, logs, real snapshots, JSON results, and +native/JIT hashes in `provenance.json`. Run GPU scripts under an owned lease +with `job-env.json` and the worktree's `.venv/bin/python`. + +- `scripts/cache_lookup_probe.py`: old-policy counterexample and cache probe; + old lookup behavior is preserved by `probabilistic=False` in the wrapper. +- `scripts/cache_lookup_fixed_probe.py`: repaired probability experiment. +- `scripts/real_selector_probe.py`: current graph/cache vs sequential oracle. +- `scripts/selector_fp32_flip_probe.py`: changed-arithmetic draft sampling. +- `scripts/context_fc_fresh_probe.py`: fresh same-input FC comparisons. +- `scripts/compare_pipeline.py`: paired real boundary and full token equality. +- `scripts/lookup_cost.py`: captured fusion-kernel timing. +- `scripts/serve-capture.sh` and `scripts/serve-optimized-capture.sh`: baseline + and context-pipeline capture services on task-owned port 18145. + +The first real-selector analysis correctly rejected four-request startup +warmup data. The retained version filters request IDs, not just tensor +shapes. `context-fc-fresh-v2.json` similarly excludes warmup snapshots; do +not use the earlier warmup-containing table as real-request evidence. One +generated offline probe had a syntax error before GPU execution; the +corrected script and successful job are retained. Neither diagnostic failure +is classified as a model defect. Keep the PR Draft pending the broader +quality/performance gates. diff --git a/docs/design/sm70_dflash2_lookup_augmented.md b/docs/design/sm70_dflash2_lookup_augmented.md index c8409e3ea3..fe5ce100d4 100644 --- a/docs/design/sm70_dflash2_lookup_augmented.md +++ b/docs/design/sm70_dflash2_lookup_augmented.md @@ -29,7 +29,13 @@ the patch mechanically. Development is tracked in public Draft PR #355. overlapping matches are legal. - The lookup proposal is fused with the neural proposal. Filled probabilistic rows become point masses in the existing sparse draft-logit cache, including - complete erase metadata for the following step. + complete erase metadata for the following step. When a weak match relies on + `VLLM_DFLASH2_LOOKUP_AGREE > 0`, the agreeing neural prefix retains its + original proposal scores. Only subsequent positions become point masses: + conditioning the prefix's own correction on its sampled tokens biases the + target distribution. Strong history-only matches and the default agreement + threshold of zero retain the existing behavior. See the + [numerical audit](sm70_dflash2_fastpath_numerics.md) for the counterexample. - Structured-output and prefill batches retain q8 and do not use lookup. - The host controller enters q16 only after two consecutive strong copy signals. B1 may coast for three steps; batches larger than one never keep diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index f2c7c7dfdd..4ab19d8840 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45854,3 +45854,36 @@ Interpretation: - No new speed claim. Full evidence, reproducible matrix command and remaining gates: [TP quality audit](sm70_quasar_tp2_tp4_quality.md). Bundle: `v100-quasar-tp2-tp4-quality-audit-20260906`. Owned GPU services/leases stopped. + +## 2026-09-06 DFlash2 proposal and context fast-path numerical audit + +- Continue Draft PR #517 from `41a9018e9f`, integration `755baae1d0`. + Optional probabilistic lookup with a positive agreement threshold rewrote + the deciding random prefix's q as point masses. A production-kernel + counterexample changes target P(A)=0.8 to 0.70027 over 100000 seeds. + Preserve that prefix's original q and only replace subsequent positions: + repaired P(A)=0.80072. Default q8 and agreement threshold0 are unaffected. +- Twelve lookup GPU regressions pass, including sparse and dense statistical + correction. The fusion kernel adds no launch; paired graph median cost + increases at most 0.006 us on B1/B4. This is not whole-round latency. +- Ten fresh E4M3/FP32 request snapshots: 70 conditional selector rows have + unchanged greedy top-1, but FP32 selector arithmetic gives maximum proposal + TV 0.0954%. Token-keyed draws change 3/17920 at official temperature1 and + 10/17920 at diagnostic temperature0.6. These are draft-token changes, not + final target-token errors. Do not claim a selector precision upgrade + improves final quality without measuring acceptance and target correction. +- Fused selector graph vs sequential dense Gumbel: all 8960 positions match; + realized/sparse/dense q caches are exact. Cache overwrite under reordered, + intersecting supports and permuted request slots passes 288 checks. +- Fresh context FC outputs reproduce the live TP4 path bitwise. Same-kernel + counterfactual TP2/TP4 output partition error reaches 9.21e-6 relative L2; + TP4-vs-cuBLAS FC drift is 1.77e-5 and becomes 1.06e-4 after BF16 norm. +- Context pipeline/KV/metadata graphs on vs off: ten complete real boundaries + match bitwise through candidate scores, q caches and accept/reject counts. + Both full outputs match the prior 260-token MBPP28 control (49 rounds). + Synchronized capture timing is excluded; retain historical 18.435/18.892 ms + medians and the unmet 17.6–18 ms objective. +- Exclude `_warmup_*` captures by request ID; first two dumps are startup + warmups. Use `context-fc-fresh-v2.json`. Raw bundle: + `v100-dflash2-fastpath-numerics-20260906`. Full reproduction and scope limits: + [fast-path numerical audit](sm70_dflash2_fastpath_numerics.md). diff --git a/tests/v1/spec_decode/test_dflash2_lookup.py b/tests/v1/spec_decode/test_dflash2_lookup.py index 13001125d8..af0faae67a 100644 --- a/tests/v1/spec_decode/test_dflash2_lookup.py +++ b/tests/v1/spec_decode/test_dflash2_lookup.py @@ -164,6 +164,150 @@ def test_lookup_controller_flag_requires_lookup_and_full_q8_emission() -> None: assert output.tolist() == [1, 0, 0] +@pytest.mark.parametrize("probabilistic", [False, True]) +@pytest.mark.parametrize("match_length", [4, 6]) +def test_agreement_prefix_retains_its_proposal_distribution( + probabilistic: bool, match_length: int +) -> None: + """Only positions after a random lookup decision may become point masses.""" + k, draft_block, agree_min = 15, 7, 3 + lookup = torch.arange(k, device="cuda", dtype=torch.int64).view(1, k) + draft = lookup.clone() + draft[:, agree_min:] += 100 + use = torch.zeros_like(draft, dtype=torch.int32) + hits = torch.zeros((), device="cuda", dtype=torch.int64) + fuse_draft( + draft, + lookup, + torch.tensor([match_length], device="cuda", dtype=torch.int32), + torch.tensor([k], device="cuda", dtype=torch.int32), + use, + torch.zeros(1, device="cuda", dtype=torch.int32), + hits, + 1, + k, + draft_block=draft_block, + nmin=4, + nstrong=6, + agree_min=agree_min, + probabilistic=probabilistic, + ) + # The repair changes q bookkeeping, not the proposed token sequence. + assert torch.equal(draft, lookup) + assert torch.all(use[:, agree_min:] == 1) + expected_prefix_use = int(not probabilistic or match_length >= 6) + assert torch.all(use[:, :agree_min] == expected_prefix_use) + + +@pytest.mark.parametrize("sparse", [False, True]) +def test_agreement_lookup_rejection_preserves_target_distribution(sparse: bool) -> None: + """The old point-mass rewrite turned p(A)=0.8 into p(A)=0.7.""" + from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample + from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( + dflash2_sparse_topk_rejection_sample, + rejection_sample, + ) + + num_reqs, steps, vocab_size = 100_000, 2, 32 + idx = torch.arange(num_reqs, device="cuda", dtype=torch.int32) + seeds = idx.to(torch.int64) + temperature = torch.ones(num_reqs, device="cuda") + positions = torch.arange(steps + 1, device="cuda").repeat(num_reqs) + 100 + dense_q = torch.full((num_reqs, steps, vocab_size), -float("inf"), device="cuda") + dense_q[:, :, :2] = 0.0 + draft = torch.zeros((num_reqs, steps), device="cuda", dtype=torch.int64) + draft[:, 0] = gumbel_sample( + dense_q[:, 0], + idx, + temperature, + seeds, + positions.view(num_reqs, -1)[:, 0].contiguous(), + apply_temperature=False, + is_drafting=True, + ) + cached_ids = torch.arange(2, device="cuda").repeat(num_reqs, steps, 1) + cached_scores = torch.zeros_like(cached_ids, dtype=torch.float32) + use = torch.zeros_like(draft, dtype=torch.int32) + fuse_draft( + draft, + torch.zeros_like(draft), + torch.full_like(idx, 4), + torch.full_like(idx, steps), + use, + idx, + torch.zeros((), device="cuda", dtype=torch.int64), + num_reqs, + steps, + draft_block=1, + nmin=4, + nstrong=6, + agree_min=1, + probabilistic=True, + ) + _point_mass_draft_logits_kernel[(num_reqs * steps,)]( + dense_q, + cached_ids, + cached_scores, + draft, + draft.stride(0), + use, + idx, + 1, + cached_ids.stride(0), + cached_ids.stride(1), + dense_q.stride(0), + dense_q.stride(1), + num_steps=steps, + top_k=2, + BLOCK_K=2, + CACHE_SCORES=True, + num_warps=1, + ) + sampled_input = torch.zeros((num_reqs, steps + 1), device="cuda", dtype=torch.int64) + sampled_input[:, 1:] = draft + cu = torch.arange(num_reqs + 1, device="cuda", dtype=torch.int32) * (steps + 1) + target_values = torch.tensor([0.8, 0.2], device="cuda").log() + target_values = target_values.repeat(num_reqs * (steps + 1), 1) + target_ids = torch.arange(2, device="cuda").expand_as(target_values).contiguous() + if sparse: + sampled, _ = dflash2_sparse_topk_rejection_sample( + target_ids, + target_values, + cached_ids, + cached_scores, + sampled_input.flatten(), + cu, + positions, + idx, + temperature, + torch.ones_like(temperature), + seeds, + steps, + ) + else: + target_dense = torch.full( + (num_reqs * (steps + 1), vocab_size), -float("inf"), device="cuda" + ) + target_dense[:, :2] = target_values + sampled, _ = rejection_sample( + target_dense, + dense_q, + sampled_input.flatten(), + cu, + positions, + idx, + idx.repeat_interleave(steps + 1), + torch.arange(steps + 1, device="cuda", dtype=torch.int32).repeat(num_reqs), + temperature, + seeds, + steps, + ) + observed = (sampled[:, 0] == 0).float().mean().item() + # Eight standard deviations; the old approximately 0.10 bias is far larger. + tolerance = 8 * (0.8 * 0.2 / num_reqs) ** 0.5 + assert abs(observed - 0.8) < tolerance, observed + + def test_point_mass_rewrite_preserves_sparse_cache_invariant() -> None: device = torch.device("cuda") num_steps, top_k, vocab_size = 15, 16, 64 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py index 0bc065907f..930d9bbabc 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/lookup.py @@ -132,6 +132,7 @@ def _fuse_draft_kernel( draft_block, k: tl.constexpr, BLOCK_K: tl.constexpr, + PROBABILISTIC: tl.constexpr, ): req = tl.program_id(0) tl.store(take_flags_ptr + req, 0) @@ -163,6 +164,13 @@ def _fuse_draft_kernel( tail = offsets >= draft_block take_tail = (match_len >= nmin_tail) & (take_head | (agreement >= draft_block)) from_lookup = tl.where(tail, take_tail, take_head) & (offsets < valid) & mask + if PROBABILISTIC: + # A weak lookup is selected by the first agree_min random proposals. + # They already equal the lookup tokens, but must retain their original + # q scores. Replacing those scores with point masses conditions on the + # draws being corrected and biases rejection sampling. Only subsequent + # positions may use a point mass conditional on that sampled prefix. + from_lookup &= tail | (match_len >= nstrong) | (offsets >= agree_min) tl.store( draft_tokens_ptr + req * draft_stride + offsets, looked_up, @@ -296,6 +304,7 @@ def fuse_draft( nmin_tail: int = 4, long_min: int = 6, take_flags: torch.Tensor | None = None, + probabilistic: bool = False, ) -> None: """Fuse lookup continuations into a model-drafted prefix in place.""" if take_flags is None: @@ -322,5 +331,6 @@ def fuse_draft( draft_block, k=num_draft_tokens, BLOCK_K=triton.next_power_of_2(num_draft_tokens), + PROBABILISTIC=probabilistic, num_warps=1, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 3dbde1f8b0..02f76d361a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -1142,6 +1142,7 @@ def _apply_lookup(self, num_reqs: int) -> None: nmin_tail=self._lookup_nmin_tail, long_min=self._lookup_long_min, take_flags=self._lookup_take_flags, + probabilistic=self.draft_logits is not None, ) draft_logits = self.draft_logits From 11d6b6b9dce15d8bf89d6f4509b0f8136274a653 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 15:46:17 +0800 Subject: [PATCH 8/9] [Audit] Validate DFlash2 verifier routes and repair diagnostics Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../sm70_dflash2_verifier_route_audit.md | 205 ++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 38 ++++ .../core/test_sm70_gdn_projection_dump.py | 55 +++++ .../test_dflash2_alignment_rank.py | 32 +++ .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 8 +- .../spec_decode/dflash2/sparse_rejection.py | 4 + 6 files changed, 339 insertions(+), 3 deletions(-) create mode 100644 docs/design/sm70_dflash2_verifier_route_audit.md create mode 100644 tests/kernels/core/test_sm70_gdn_projection_dump.py create mode 100644 tests/v1/spec_decode/test_dflash2_alignment_rank.py diff --git a/docs/design/sm70_dflash2_verifier_route_audit.md b/docs/design/sm70_dflash2_verifier_route_audit.md new file mode 100644 index 0000000000..8ab17fe485 --- /dev/null +++ b/docs/design/sm70_dflash2_verifier_route_audit.md @@ -0,0 +1,205 @@ +# DFlash2 verifier acceleration audit, 2026-09-06 + +## Contract and interpretation + +Continue Draft PR #517 from `53be620005fb0f7664dcd27a512979191b342c73`, +integration `755baae1d075ee04fa9096b23fc0225b23589a86`. This audit measures +cost and numerical behavior separately. Production changes are limited to +two diagnostic defects; no inference arithmetic or precision default changes. + +The model pair remains QUASAR Qwen3.8-27B all-NVFP4 revision +`d8e6fbfa3e3a78899b440222b827430045a05b44` and DFlash2 revision +`dedf8df68adfb1afeaf7b7480c0a0243108177b4`. Execution is SM70 W4A16, +FP32 logits, E4M3 target KV, FP16 draft KV, Flash-V100, probabilistic +draft7/q8, TP4, maximum length 262144, batch budget 4096, maximum requests 4, +memory 0.8, prefix cache and Mamba align. V2 runner, target full/piecewise +graphs, full draft graphs, context pipeline and context KV graphs are enabled. +The environment is CUDA 12.8, Torch 2.10.0+cu128, Python 3.12.13, Triton 3.6.0. + +GPUs 4–7 are exclusively leased because GPUs 0–3 belong to another task. +Local component comparisons use logical CUDA 0, physical GPU 4, while loading +each of the four logical weight shards in turn. They exclude TP communication. +Do not compare their absolute latency to a different GPU set as a speedup. + +References use identical inputs and checkpoint weights. FP32/FP64 arithmetic +references are not an unquantized QAT teacher. No dataset quality-loss +percentage is inferred from relative L2, TV, token flips or shorter output. + +## Route-by-route ledger + +| Acceleration path | Numerical evidence | Cost evidence and disposition | +| --- | --- | --- | +| NVFP4 QPN2 and fused gate/up/SiLU | Previous 64-layer audit covers 256 fused / 496 logical projections against the same checkpoint. Nonzero rounding and TP partition effects remain. | Retain previous source and matrix evidence; no new isolated timing. A BF16 teacher comparison remains open. | +| GDN QKV packing, split and integer metadata | Previous packing and q8 accepted-state probes match their references, including accepted selectors 1–8 and untouched states. | Retained evidence; no new isolated timing. Do not infer floating-point quality from an integer route hit. | +| GDN one-pass gated norm | New 24 real-input rank/step cases reproduce the live fused output exactly; fused/native FP16 outputs can differ. | New q8 local graph timing: 2.231 us fused vs 25.428 us eager native. Full-model attribution is qualified below. | +| Gemma fused residual/RMS | New 144 admitted real-input cases reproduce live fused outputs exactly; residual sums equal the staged reference. Normalized FP16 outputs can differ. | New local graph timing: 3.620 us fused vs 26.846 us eager reference. The reference timing is not the compiled production fallback. | +| Grouped Flash-V100 E4M3 verifier | Previous paired/scalar E4M3 conversion path is bitwise equal through 256K. Comparison to FP32 attention on the same KV has ordinary reduction error. | Retained grouped-path timings:36.25 vs45.57 us at 1032 context;1089.13 vs1617.41 us at 262144. These are not end-to-end round times. | +| TP all-reduce | Previous real partials match FP32 summation followed by FP16 output, across all ranks. Earlier local partial rounding is not recovered. | No new isolated collective timing. Do not count synchronization waits as arithmetic cost. | +| QPN8 top64 plus FP32 head rerank | New 535 target/draft rows: no local top21 or required global top-k misses; no target support changes. | New local q8 graph timing: 573.4–575.1 us vs 992.6–1010.6 us native dense FP32 head, excluding TP collectives. Retain the search with finite-coverage limits. | +| Top-k/top-p boundary protection | Previous ambiguous-cutoff regressions cover ties and scan rounding. All target groups in the new535-row head probe are unambiguous. | Retain the guard. Its earlier small-probe CPU measurement must not be confused with the preceding GPU wait. | +| Sparse target rejection | New 60 independent real q8 rounds cover emitted counts 1–8; token decisions and counts equal dense rejection and the actual captured sampler. |15.242 us sparse including top-p vs 43.530 us dense rejection alone. The dense timing excludes its separate top-k/top-p step. | +| Identity-index and padded-stride copy removal | Previous padded-sentinel tests and the new real sampler replay preserve the same inputs and decisions. | Retain; do not add component deltas to whole-round latency. | +| Output-sharded context FC | Previous real TP4 output reproduced bitwise; same-family TP2/TP4 partitions show small rounding differences. | Retained previous evidence; this is not a new TP2 endpoint. | +| Context compute/store and metadata graphs | Previous paired complete request boundaries are bitwise equal through candidates, q caches and acceptance, with matching full output. | Retain previous real-boundary and production timing evidence. | +| Lookup augmentation | Previous positive-agreement probability-bookkeeping repair preserves the deciding prefix's q. Default q8 does not use lookup. | No new lookup change or timing in this audit. | + +Previous evidence and its exact limitations are in the +[operator audit](sm70_quasar_dflash2_operator_audit.md), +[E4M3/FP32 report](sm70_quasar_e4m3_fp32_logits.md), +[TP/QAT audit](sm70_quasar_tp2_tp4_quality.md), and +[proposal/context audit](sm70_dflash2_fastpath_numerics.md). + +## Head coverage and real rejection + +The head probe uses 465 target rows and 70 actual draft rows from the retained +fixed-prefix and real-request captures. Target support is global top20; +draft support is top16. All four local head shards use the actual QPN8 +packing, top64 search, indexed FP32 rerank and dense-vocabulary ordering. +The independent oracle multiplies identical FP16-materialized weights in +FP32. All local top21 candidates are covered, required global top-k sets +match, and maximum protected target sampling TV is 1.2456439e-6. +There are zero changed target top-p supports. This extends observed coverage; +it does not prove that approximate top64 search can never miss a candidate. + +The rejection probe uses actual aligned p, q, proposed IDs, request slots and +positions from the natural MBPP28 request. It scatters those supports into +the full 248320-token vocabulary, applies the public top-k/top-p reference, +and executes the production dense rejection kernels. All 60 independent +rounds reproduce sparse tokens and emitted counts, including rejection at +every depth and full acceptance. The diagnostic directory originally holds +240 files because all four workers incorrectly identified as rank0; these +are 60 rounds with four replicas, not 240 independent samples. + +Head timing uses 16 graph nodes and eight replays per observation, five +observations. Rejection timing uses 32 nodes and 16 replays per observation, +seven observations. Both use resident inputs and CUDA events, with warmup. +They are component measurements, not wall-clock service improvements. + +## Norm arithmetic and the attribution limit + +Real norm inputs come from three q8 decode snapshots, all four ranks, seven +target layers for Gemma and two GDN layers. The first layer's post-attention +residual is FP16 and does not enter the fused Gemma gate; it is excluded. +The admitted 168 cases all reproduce their live fused outputs exactly. + +| Operator | Fused vs staged FP32 reference | Correctly rounded FP64 comparison | +| --- | --- | --- | +| Gemma residual/RMS | Residual exact; max normalized relative L2 1.05e-5, max absolute 0.001953125 | 468 fused vs 444 fallback output elements differ from FP64 over 144 cases; neither path is exact | +| GDN gated RMS | Maximum relative L2 approximately 6.94e-6, max absolute 7.63e-6 over 24 cases | 32 fused vs 23 fallback elements differ from FP64; maximum absolute error is 3.05e-5 for both | + +The FP64 Gemma oracle retains the FP32 residual-addition boundary and then +evaluates normalization in FP64. The GDN oracle evaluates RMS, weights and +SiLU in FP64. These finite local comparisons do not establish which full +model has better text quality. Do not simply switch to the slower reference. + +Full-model experiments use identical forced prefixes for MBPP28 and MBPP3: +one prefill logit followed by 16 q8 rounds, 129 positions per case. Inputs, +positions and query boundaries are checked explicitly. Selected snapshots +also verify identical replicas within each TP configuration. Temperature 1, +top-k20 and top-p0.95 are applied to the saved native FP32 target logits. + +| Diagnostic comparison | MBPP28 maximum TV | MBPP3 maximum TV | Greedy changes / top-p support changes | +| --- | ---: | ---: | --- | +| Gemma fusion off vs optimized |4.5527%|1.9551%|0/258 greedy;3 support rows | +| GDN norm fusion off vs optimized |2.6499%|1.5832%|0/258 greedy;2 support rows | +| Gemma off with tensor-copy capture disabled |4.3382%|1.7690%|0/258 greedy;4 support rows | +| Same optimized configuration, repeated startup |4.3292%|0.6145%|0/258 greedy;1 support row | + +**The same-configuration control also drifts. These measurements cannot +isolate either norm switch as the cause of the full-model differences.** +For example, in the first q8 Gemma comparison, layer 0/rank 2 GDN core output +already differs at 52 elements, before any affected Gemma fusion. Its local +input norm and Z are equal; the difference spreads through projection and +all-reduce. The repeated startup has the same complete natural 297-token +MBPP28 output, while its forced-prefix logits still differ. This is a +repeatability gate for the state/prefix diagnostic, not proof that ordinary +production has a 4.33% sampling defect. + +The tensor-copy-disabled variant still retains Python forward wrappers and +the forced-prefix intervention. Both corresponding natural requests reach +the diagnostic 512-token cap; they are not accepted natural-stop quality +evidence. A separate service without a worker extension, layer/selector +capture or forced prefixes is used for production closure. + +The next attribution check should isolate the first GDN state divergence +with a fully reset single request and matching prefill state, then repeat +the same q8 update before changing a norm implementation. An aggregate +benchmark score or a stable top-1 cannot replace that check. + +Coupled production target-Gumbel draws at seeds 0–63 also change: the Gemma +diagnostic has 20 differing draws over 16512 positions (one at seed 0), and +the GDN diagnostic has 20 (none at seed 0). These are conditional target draws +on the saved prefixes, not full DFlash rejection or a free-generation score. +The repeated-startup confound also applies to their attribution. +The same-configuration repeated startup itself changes20/16512 paired +target draws, including one at seed0, so those flip counts alone cannot +identify either norm as the cause. + +## Uninstrumented production closure + +A separate service uses no worker extension, layer/selector capture or +forced-prefix hook. The pinned sampling contract is temperature 1, top-k20, +top-p0.95, natural EOS, maximum output 1024; release1k uses seed 20260925 and +MBPP28 seed 0. After one warmup, three measured repetitions give: + +| Request | Median round ms | Pure decode tokens/s | TTFT seconds | Output tokens / rounds | Emitted tokens per round | +| --- | ---: | ---: | ---: | ---: | ---: | +| release1k, input 1019 |19.505|154.434|0.3563|248 / 82|3.024 | +| MBPP28, input 135 |19.092|234.829|0.1187|270 / 60|4.500 | + +All three repetitions of each request have identical complete token hashes, +natural stop and a nonempty final answer. MBPP28 returns the correct +`n * (n + 1) * (n + 2) // 6` implementation. The final hashes are +`0759c9a5199126539653edae14393addf3bfd8ae6e0a42a56457ea49bd8ee584` +and `e965b6253b702563d53be0a3084030304eff43a7e34c43df9c17fc15e8f46fbc`. + +These physical GPUs 4–7 and output trajectories differ from the historical +GPUs 0–3 record (303/260 tokens, 18.892/18.435 ms). This is a new hardware-set +baseline, not a matched performance regression or a 17.6–18 ms restoration. +Four resolved native-library hashes match the preceding audit. Cross-startup +and cross-GPU-set output parity remains open; neither the diagnostic results +nor this short stable production repetition proves broad output quality. + +## Diagnostic repairs and validation + +`sm70_gdn_projection_dump` returned its input despite a non-aliasing custom-op +schema. `torch.library.opcheck(test_schema)` reports the alias violation. +Its eager/captured and fake implementations now return owned storage. +Tests cover schema/FakeTensor, AOT preservation of a live input, and changed +inputs under CUDA Graph replay. This affects enabled diagnostics only. + +The alignment dumper used only environment rank variables. Multiprocess +workers without those variables all fell back to 0. It now uses the initialized +distributed rank first and preserves the environment fallback before process +group initialization. Tests cover ranks 0–3, stale environment values and the +uninitialized fallback. Three GDN dump GPU tests and seven rank tests pass. + +The retained old alignment files are explicitly deduplicated; use +`sparse-real-cost-v2.json`. The first norm table included the non-admitted +first-layer residual case; use `norm-real-cost-v2.json`. The initial head +probe failed before measurements because inference buffers were updated +outside inference mode; the corrected probe uses the same inference-mode +contract as serving. These are diagnostic failures, not target-model defects. + +## Reproduction + +Raw bundle: `v100-dflash2-verifier-route-audit-20260906`, under the task cache +recorded in the local handoff. `queue/*.done.json` records exact commands, +environment, GPU ownership and logs; `provenance.json` retains native/JIT +hashes. Use the owned worktree's `.venv/bin/python` under an exclusive lease. + +- `head_coverage_cost.py`: native candidate coverage, FP32 reference and costs. +- `sparse_real_cost.py`: independent aligned rounds, actual decisions and cost. +- `norm_real_cost.py` and `norm_fp64_oracle.py`: admitted real norm shapes. +- `run_matrix.py`, `run_minimal_matrix.py`, `run_repeat.py`: separately started + diagnostic configurations, with request-ID-bound forced-prefix tapes. +- `analyze_routes.py`, `analyze_minimal_routes.py`, `analyze_repeat.py`: strict + prefix validation and native target-distribution comparisons. +- `route_sample_flips.py`: paired production Gumbel draws on saved distributions, + checked against dense sampling at seed 0; not full speculative decoding. +- `serve-production.sh` and `run_production.py`: independent uninstrumented + production closure, natural EOS, one warmup and three measured requests. + +Keep Draft. Native arithmetic is unchanged, the broader QAT teacher question +remains open, and single-path full-model attribution requires the state and +prefix repeatability gate above. diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 4ab19d8840..fdb0e4b269 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45887,3 +45887,41 @@ Interpretation: warmups. Use `context-fc-fresh-v2.json`. Raw bundle: `v100-dflash2-fastpath-numerics-20260906`. Full reproduction and scope limits: [fast-path numerical audit](sm70_dflash2_fastpath_numerics.md). + +## 2026-09-06 DFlash2 verifier route costs and quality attribution + +- Continue owned Draft PR #517 from `53be620005`, integration `755baae1d0`. + GPU 0–3 are occupied by another task; this audit leases GPU 4–7. Do not label + changed hardware-set measurements as recovery of the earlier speed peak. +- QPN8 support/FP32 rerank: 535 real rows (465 target, 70 draft), no local top21 + or required global top-k misses, no target top-p support changes; maximum + target TV 1.2456e-6. Local q8 head cost is 573–575 us vs 993–1011 us dense + FP32, excluding TP communication. +- Sparse rejection: 60 independent real q8 rounds, emitted counts 1–8, exactly + equal to dense rejection and the captured output. Local graph cost 15.242 us + vs 43.530 us dense rejection alone; the latter excludes separate top-p. +- Real admitted norm cases: 144 Gemma and 24 GDN reproduce live fused outputs + exactly, but can differ from staged FP32/FP64 references. Standalone fused + vs eager costs: Gemma 3.620/26.846 us; GDN 2.231/25.428 us. Do not extrapolate + the eager reference timing to the compiled full-model fallback. +- Full-model fixed-prefix norm-switch comparisons show up to 4.55% TV with + unchanged greedy top-1. However, the same optimized configuration restarted + also differs by up to 4.33%. Layer 0/rank 2 GDN core already differs before + the first affected Gemma fusion. Attribution to either norm is not closed; + repeatability of prefill state and the forced-prefix diagnostic comes first. +- Repair two diagnostics only: GDN projection dump violates its non-aliasing + schema; alignment rank detection can duplicate every TP replica. Three GPU + schema/AOT/graph tests and seven distributed-rank/fallback tests pass. + Deduplicate the original 240 alignment files into 60 independent rounds. + Use `norm-real-cost-v2.json` to exclude first-layer FP16 residuals that do + not enter the fused Gemma gate. +- Uninstrumented GPU 4–7 production closure: medians 19.505/19.092 ms per round, + 154.434/234.829 decode tokens/s, 248/270 output tokens, 82/60 rounds. Each + request's three measured repetitions match and stop naturally. The token + trajectories differ from the historical GPU 0–3 run; 17.6–18 ms and broad + cross-startup output parity remain open. Native hashes are unchanged. +- Route-by-route ledger, controls and failed-diagnostic exclusions: + [verifier route audit](sm70_dflash2_verifier_route_audit.md). Raw bundle: + `v100-dflash2-verifier-route-audit-20260906`. Runtime arithmetic and precision + defaults are unchanged; keep the existing Draft PR pending the remaining + state, QAT teacher and performance gates. diff --git a/tests/kernels/core/test_sm70_gdn_projection_dump.py b/tests/kernels/core/test_sm70_gdn_projection_dump.py new file mode 100644 index 0000000000..aec1411610 --- /dev/null +++ b/tests/kernels/core/test_sm70_gdn_projection_dump.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +LAYER_NAME = "language_model.model.layers.0.linear_attn" + + +def _dump(x): + from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: F401 + + return torch.ops.vllm.sm70_gdn_projection_dump(x, "proj_core_in", LAYER_NAME) + + +def test_gdn_projection_dump_obeys_nonaliasing_schema(): + x = torch.randn(8, 128, device="cuda", dtype=torch.float16) + out = _dump(x) + torch.library.opcheck( + torch.ops.vllm.sm70_gdn_projection_dump.default, + (x, "proj_core_in", LAYER_NAME), + test_utils=("test_schema", "test_faketensor"), + ) + assert torch.equal(out, x) + out.zero_() + assert torch.count_nonzero(x) > 0 + + +def test_gdn_projection_dump_preserves_compiled_live_input(): + def forward(x, z): + saved = x.float() * z.float() + projected = _dump(x) + return saved, projected.float() + z.float() + + x = torch.randn(8, 128, device="cuda", dtype=torch.float16) + z = torch.randn_like(x) + expected = forward(x, z) + actual = torch.compile(forward, backend="aot_eager", fullgraph=True)(x, z) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, atol=0, rtol=0) + + +def test_gdn_projection_dump_graph_keeps_owned_output(): + x = torch.randn(8, 128, device="cuda", dtype=torch.float16) + _dump(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _dump(x) + for value in (1.0, -2.0, 4.0): + x.fill_(value) + graph.replay() + assert torch.equal(out, x) + out.zero_() + assert torch.all(x == value) diff --git a/tests/v1/spec_decode/test_dflash2_alignment_rank.py b/tests/v1/spec_decode/test_dflash2_alignment_rank.py new file mode 100644 index 0000000000..61eac3883c --- /dev/null +++ b/tests/v1/spec_decode/test_dflash2_alignment_rank.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode.dflash2.sparse_rejection import _diagnostic_rank + + +@pytest.mark.parametrize("rank", [0, 1, 2, 3]) +def test_alignment_dump_uses_initialized_process_rank(monkeypatch, rank): + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("LOCAL_RANK", "0") + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: rank) + assert _diagnostic_rank() == rank + + +@pytest.mark.parametrize( + "rank,local,expected", [(None, None, 0), (None, "2", 2), ("3", "1", 3)] +) +def test_alignment_dump_rank_fallback_before_initialization( + monkeypatch, rank, local, expected +): + monkeypatch.delenv("RANK", raising=False) + monkeypatch.delenv("LOCAL_RANK", raising=False) + if rank is not None: + monkeypatch.setenv("RANK", rank) + if local is not None: + monkeypatch.setenv("LOCAL_RANK", local) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + assert _diagnostic_rank() == expected diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 685cadbc94..fc99c3a880 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1265,7 +1265,7 @@ def _sm70_gdn_projection_dump_impl( ) -> torch.Tensor: _sm70_gdn_graph_buffer_copy(label, layer_name, tensor, "proj") if torch.cuda.is_current_stream_capturing(): - return tensor + return tensor.clone() dump_dir = os.getenv("VLLM_SM70_DUMP_GDN_PROJ_DIR") enable_file = os.getenv("VLLM_SM70_DUMP_GDN_PROJ_ENABLE_FILE") if ( @@ -1302,7 +1302,9 @@ def _sm70_gdn_projection_dump_impl( }, path, ) - return tensor + # This custom op has a non-aliasing output schema. Returning the input + # lets AOT reuse live projection storage across diagnostic boundaries. + return tensor.clone() def _sm70_gdn_projection_dump_fake( @@ -1310,7 +1312,7 @@ def _sm70_gdn_projection_dump_fake( label: str, layer_name: LayerNameType, ) -> torch.Tensor: - return tensor + return tensor.clone() direct_register_custom_op( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py index fca6ac1afb..4474ba672e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py @@ -91,6 +91,10 @@ def _safe_dump_tag(raw_tag: str) -> str: def _diagnostic_rank() -> int: + # Multiprocess workers need not export RANK/LOCAL_RANK. Falling back to + # zero there dumps every TP replica and overcounts independent samples. + if torch.distributed.is_initialized(): + return torch.distributed.get_rank() return int(os.getenv("RANK", os.getenv("LOCAL_RANK", "0"))) From a9ff3354c5e8f8c2b7c723df686095d27970d9ee Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Sun, 6 Sep 2026 17:13:19 +0800 Subject: [PATCH 9/9] [Audit] Verify prefill repair ancestry for DFlash2 quality Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../sm70_dflash2_verifier_route_audit.md | 56 +++++++++++++++++++ docs/design/sm70_v100_migration_control.md | 14 +++++ 2 files changed, 70 insertions(+) diff --git a/docs/design/sm70_dflash2_verifier_route_audit.md b/docs/design/sm70_dflash2_verifier_route_audit.md index 8ab17fe485..265c77ebe2 100644 --- a/docs/design/sm70_dflash2_verifier_route_audit.md +++ b/docs/design/sm70_dflash2_verifier_route_audit.md @@ -160,6 +160,62 @@ Four resolved native-library hashes match the preceding audit. Cross-startup and cross-GPU-set output parity remains open; neither the diagnostic results nor this short stable production repetition proves broad output quality. +## Prefill PR history and repeatability recheck + +The follow-up reviewed the public PR records against main +`95205a2d9952813aa7469f63ff65b8f2813c027a` and audit runtime source +`11d6b6b9dce15d8bf89d6f4509b0f8136274a653`. The Flash-V100 companion +repository has no PR records. No applicable, validated, unmerged prefill +repair was identified for this QUASAR dense GDN/full-attention route. + +| PR | Repair or evidence | Applicability and integration | +| --- | --- | --- | +| [#202](https://github.com/1CatAI/1Cat-vLLM/pull/202) | Two-phase P commit removes a paged-prefill shared-memory race for D64/D128. | Already in main and the tested audit source. D256 uses separate P storage. | +| [#226](https://github.com/1CatAI/1Cat-vLLM/pull/226) | Aligns WMMA accumulators and shared-memory base; six-replay D128 regression. | Already in both trees; 32-byte accumulator alignment and the assertion remain present. | +| [#219](https://github.com/1CatAI/1Cat-vLLM/pull/219), [#350](https://github.com/1CatAI/1Cat-vLLM/pull/350) | FP32 XQA probabilities for FP16 KV; restoration of D256 prefill operators. | Already in both trees; neither is a new E4M3/GDN-state fix. | +| [#403](https://github.com/1CatAI/1Cat-vLLM/pull/403) | Records Flash Next QSA prefill quality evidence. | Already integrated; its matched quality claim concerns a different model route. | +| [#434](https://github.com/1CatAI/1Cat-vLLM/pull/434), [#408](https://github.com/1CatAI/1Cat-vLLM/pull/408) | Legacy-runner hybrid prefill dispatch; Flash Next/cache correctness repairs. | Already integrated. The legacy-runner dispatch fix does not execute in this MRV2 run. | +| [#494](https://github.com/1CatAI/1Cat-vLLM/pull/494), [#525](https://github.com/1CatAI/1Cat-vLLM/pull/525) | QSA logical-page ordering and its NVFP4 model validation. | Both are in main; this Qwen3.5-family 27B model does not execute QSA. | +| [#524](https://github.com/1CatAI/1Cat-vLLM/pull/524) | Experimental E4M3 grouped attention with FP32 partials. | Remains Draft with a failed model token gate; not admitted as a quality repair. | + +Main advanced from `4366d9d5fe80eeaf79575b51ec36a6a032673df0` when another +task merged #525 during this review. Its shared CUDA-file change is limited +to the Qwen4Exp HC down scatter kernel; it does not change this model's +ordinary TP all-reduce or the paged-prefill kernel. + +The runtime library is the retained `lib-final` artifact with SHA256 +`c3f3bef28a21f681d3d3d84e65d5f208b9d2c282b2c4bfe7cb5f7e221d55802e`. +Its retained compile and link logs build the paged-prefill object from this +owned source tree. The paged-prefill source is unchanged against current +main and contains both old fixes. Thus the observed 4.33% drift was already +measured with those repairs; merging them again supplies no new intervention. + +CPU reanalysis of the retained captures also narrows the causal gap: + +- Both runs execute the full identical 135-token MBPP28 prompt. Their final + prefill hidden states already differ, before any forced q8 acceptance + update. The first q8 layer-0 QKV projection is exact across starts, while + rank 2's GDN core has 52 different elements, maximum absolute 3.8147e-6, + before the audited output normalization. +- At position 212, the top-20 distribution's first two candidates have + cumulative mass 0.94912833 versus 0.95171565. Crossing top-p 0.95 removes + token 23, whose sampled mass was 0.04329225; recomputed TV is 0.04329227. + Full-vocabulary softmax TV at that position is already 0.02030133. +- The original probe omitted prefill layer tensors and incoming conv/SSM + state. It cannot distinguish a prefill arithmetic difference from state + propagation or core execution. It does not prove stale state or a race. + The fixed FP8 tuning experiment in #524 concerns another checkpoint: + this NVFP4 prompt M135/M153 exceeds the default dense tuning maximum M16. + Do not transfer that experiment's causal conclusion to this model. + +No runtime source, production flag or integration branch changes were made +by this history recheck. No fresh GPU replay was run while all GPUs were +owned by other tasks. Preserve the pending state-replay gate; the historical +two passing D128 tests in #226 are not a new model-level validation. +PR snapshots and ancestry/library checks are retained in bundle +`v100-dflash2-prefill-pr-audit-20260906`; the CPU reanalysis is in the previous +verifier bundle's `repeatability-cause` directory and exited successfully. + ## Diagnostic repairs and validation `sm70_gdn_projection_dump` returned its input despite a non-aliasing custom-op diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index fdb0e4b269..ed399c64c8 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -45925,3 +45925,17 @@ Interpretation: `v100-dflash2-verifier-route-audit-20260906`. Runtime arithmetic and precision defaults are unchanged; keep the existing Draft PR pending the remaining state, QAT teacher and performance gates. + +### Prefill repair history recheck + +At main `95205a2d9952813aa7469f63ff65b8f2813c027a`, the Flash-V100 +paged-prefill race/alignment fixes #202/#226 are already integrated and +present in the source used for the retained 4.33% TV run and rebuilt native +library. QSA allocation repair #494 and its validation #525 are in main but +do not execute in +the QUASAR 27B GDN/full-attention model. Draft #524 has a failed model token +gate. No applicable validated pending prefill repair was found or merged. +The [PR history recheck](sm70_dflash2_verifier_route_audit.md) records the +scope and ancestry checks, prefill divergence, top-p boundary amplification, +and the remaining conv/SSM state-replay requirement. CPU capture analysis +passed; no fresh GPU run or performance claim accompanies this recheck.