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/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_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_dflash2_verifier_route_audit.md b/docs/design/sm70_dflash2_verifier_route_audit.md new file mode 100644 index 0000000000..a4aec516d3 --- /dev/null +++ b/docs/design/sm70_dflash2_verifier_route_audit.md @@ -0,0 +1,264 @@ +# 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. + +## 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 +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. + +The broader QAT teacher question remains open, and single-path full-model +attribution requires the state and prefix repeatability gate above. The user +subsequently requested mainline integration of PR #517's validated repairs +with these limits retained. That integration does not close the remaining +quality investigation or promote the optional precision/performance flags. +Merge-time checks and the exact integrated revision are recorded on #517. 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..1c7fc6a9b7 --- /dev/null +++ b/docs/design/sm70_quasar_dflash2_operator_audit.md @@ -0,0 +1,248 @@ +# 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 +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_quasar_dflash2_quality_speed_recovery.md b/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md new file mode 100644 index 0000000000..b82c621f49 --- /dev/null +++ b/docs/design/sm70_quasar_dflash2_quality_speed_recovery.md @@ -0,0 +1,202 @@ +# 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.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. + +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. 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 +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 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. + +| 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. +- 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. + +```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/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}`. + 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_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_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 0a0f5ef0b1..aed89a39f5 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -2,6 +2,29 @@ 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 +`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: @@ -45766,6 +45789,157 @@ 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. +- 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 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 + 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`. + +## 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. + +## 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). + +## 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. + +### 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. + - The exact down packet screen completed and is rejected. Same-source complete-HC control `1.989379 ms`, CUDA down plus separate gather `2.072549 ms`, fused down/gather `2.218926 ms`. All intermediate/final @@ -45959,3 +46133,24 @@ Interpretation: code, including a warp-ordering attempt, so the intended lookahead was not established. Do not run these as purported prefetch variants or repeat the previously rejected ordinary CUDA128 down. No GPU startup for this screen. + +## 2026-09-06 DFlash2 quality repair mainline integration + +The user explicitly requested that the existing output-quality repair +PR #517 be integrated into main after the remaining limits were reported. Its +validated scope includes sampling cutoff boundaries, TP2 NVFP4 alignment, +lookup proposal probabilities, independent FP32 logits, E4M3 q8 support, +context/metadata graph options, and diagnostic ownership/rank fixes. + +Synchronizing main at `95205a2d9952813aa7469f63ff65b8f2813c027a` preserves +the independent QSA ordering, HC/router, AWQ and PP work. The only merge +conflict is this append-only ledger; both histories are retained. Native +Flash-V100 integration is rebuilt for SM70 and the scoped regressions are +recorded with final status, exact head and artifact hashes on PR #517. + +FP32 logits and the new context graph/pipeline switches remain opt-in; E4M3 +requires an explicit KV setting and a rebuilt native library. This admission +does not certify QAT-versus-BF16 quality, solve the 4.33% fixed-prefix +repeatability issue, or establish recovery to 17.6–18 ms. Historical Draft +notes describe the investigation at their recorded revisions. The remaining +state/prefix and performance goals continue after implementation integration. 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 aa43664a5a..91624ea96f 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -2126,7 +2126,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); @@ -2214,7 +2217,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); @@ -2318,7 +2323,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); @@ -4106,11 +4113,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"); @@ -4180,12 +4193,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); \ @@ -4199,7 +4212,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(), \ @@ -4231,30 +4244,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/benchmarks/test_sm70_quasar_nvfp4_oracle.py b/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py new file mode 100644 index 0000000000..c89714244e --- /dev/null +++ b/tests/benchmarks/test_sm70_quasar_nvfp4_oracle.py @@ -0,0 +1,48 @@ +# 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("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( + { + "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, 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/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/kernels/attention/test_sm70_flash_v100_paged_prefill_determinism.py b/tests/kernels/attention/test_sm70_flash_v100_paged_prefill_determinism.py index e4904ded58..44b2b9db57 100644 --- a/tests/kernels/attention/test_sm70_flash_v100_paged_prefill_determinism.py +++ b/tests/kernels/attention/test_sm70_flash_v100_paged_prefill_determinism.py @@ -109,6 +109,8 @@ def test_sm70_flash_v100_paged_prefill_d128_replay_is_bit_exact( True, -1, -1, + None, + 0, ) torch.accelerator.synchronize() direct_outputs.append(direct_output) 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/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/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/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..c74b68ed68 --- /dev/null +++ b/tests/quantization/test_sm70_fp32_lm_head_admission.py @@ -0,0 +1,44 @@ +# 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("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): + 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/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/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/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/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/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 d83592a035..d13bf8e8fe 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -226,6 +226,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 @@ -242,6 +243,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 @@ -2124,6 +2127,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. @@ -2219,6 +2227,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/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/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/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 b8a0819795..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") @@ -202,6 +203,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 +227,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 +258,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 +266,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): + 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 +293,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 @@ -299,7 +304,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 +313,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 @@ -361,6 +370,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() ) @@ -395,6 +414,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): @@ -522,7 +547,13 @@ 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() + and not getattr(layer, "_sm70_dflash2_fp32_logits", False) + ): return None if x.dtype != torch.float16 or not x.is_cuda: return None @@ -561,23 +592,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], @@ -634,8 +677,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/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/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/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index 0f921ebd05..9736dd0eca 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( @@ -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")) @@ -5279,7 +5296,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 +5371,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( 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/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9ab591206e..1456f69be8 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, ) @@ -1728,6 +1729,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/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/sparse_rejection.py b/vllm/v1/worker/gpu/spec_decode/dflash2/sparse_rejection.py index 01b42e3968..4474ba672e 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 @@ -62,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"))) @@ -219,10 +252,30 @@ 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, ) - draft_sampled = input_batch.input_ids[input_batch.logits_indices] - pos = input_batch.positions[input_batch.logits_indices] + 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] + 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..02f76d361a 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. @@ -983,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 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