From 85c71e9bfea0202e19f2178fdcd76a00dbc4ac8e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 03:18:04 +0800 Subject: [PATCH 1/5] [Perf] Reduce host transfers in Flash-Next decode verification Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_gguf_cache_copy.py | 276 ++++++++++++++++++ ...flashnext_decode_verifier_16gb_20261014.md | 122 ++++++++ tests/models/qwen4_exp/test_host_kv.py | 87 +++++- tests/v1/worker/test_step_profiler.py | 33 +++ .../layers/quantization/gguf_expert_cache.py | 8 +- vllm/models/qwen4_exp/nvidia/ops/host_kv.py | 15 +- vllm/v1/spec_decode/profiling.py | 4 +- 7 files changed, 522 insertions(+), 23 deletions(-) create mode 100644 benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py create mode 100644 docs/design/flashnext_decode_verifier_16gb_20261014.md diff --git a/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py new file mode 100644 index 000000000..7fe5eb789 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""TP4 paired screens of the existing expert-cache byte-copy kernel.""" + +import argparse +import json +import os +import statistics +from functools import partial +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm import _custom_ops # noqa: F401 -- register mapped-view operator +from vllm.model_executor.layers.quantization.gguf_expert_cache import _apply_plan + +POLICIES = { + "reference": (16, 8, 1024, 8), + "o2s64b256w4": (2, 64, 256, 4), + "o4s32b256w4": (4, 32, 256, 4), + "o4s64b256w4": (4, 64, 256, 4), + "o8s16b512w4": (8, 16, 512, 4), + "o2s32b512w4": (2, 32, 512, 4), +} + + +def launch(buffers, policy): + programs, split, block, warps = POLICIES[policy] + plan, arena, mapped, device_offsets, host_offsets, words = buffers + _apply_plan[(programs, split)]( + plan, + arena, + mapped, + device_offsets, + host_offsets, + words, + 0, + OPS_PROGRAMS=programs, + SPLIT=split, + BLOCK=block, + num_warps=warps, + ) + + +def reset_data(arena, host, original_arena, original_host): + arena.copy_(original_arena) + host.copy_(original_host) + + +def summarize(values): + ordered = sorted(values) + return { + "mean": statistics.mean(values), + "p50": statistics.median(values), + "p90": ordered[round(0.9 * (len(ordered) - 1))], + "values": values, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--shapes", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--samples", type=int, default=12) + parser.add_argument("--repeats", type=int, default=16) + parser.add_argument("--boundary-only", action="store_true") + args = parser.parse_args() + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + device = torch.device("cuda", rank) + # Flash-Next GSQ-RCO IQ3_S TP4 words per expert. Q2_0's down projection + # retains its 192-column physical padding for the 160-column local input. + layouts = { + 0: dict(zip(("w1", "w3", "w2"), (19600, 19600, 28800))), + 1: dict(zip(("w1", "w3", "w2"), (16400, 16400, 17280))), + 17: dict(zip(("w1", "w3", "w2"), (22000, 22000, 28800))), + 47: dict(zip(("w1", "w3", "w2"), (27200, 27200, 28800))), + } + if args.shapes: + layouts = {} + for entry in json.loads(args.shapes.read_text()): + layer = int(entry["name"].split(".")[1]) + shard = ( + "w2" + if "down" in entry["name"] + else "w1" + if "gate" in entry["name"] + else "w3" + ) + assert entry["rank_bytes"] % 8 == 0 + layouts.setdefault(layer, {})[shard] = entry["rank_bytes"] // 8 + if args.boundary_only: + layouts = {-1: {"w1": 255, "w3": 256, "w2": 257}} + output = { + "world_size": 4, + "policies": POLICIES, + "samples": args.samples, + "repeats": args.repeats, + "cases": [], + "correctness": [], + } + + for layer, shards in sorted(layouts.items()): + sizes = [shards[k] for k in ("w1", "w3", "w2")] + # Production slot counts; tail guards separate projection banks. + slots, host_slots = 410, 136 + offsets, host_offsets, total, host_total = [], [], 0, 0 + for n in sizes: + total = ((total + 31) // 32) * 32 + offsets.append(total) + host_offsets.append(host_total) + total += slots * n + 32 + host_total += host_slots * n + arena = torch.arange(total, dtype=torch.int64, device=device) + host = torch.empty(host_total, dtype=torch.int64, pin_memory=True) + host.copy_(torch.arange(host_total, dtype=torch.int64).mul_(-3).sub_(7)) + original_arena, original_host = arena.clone(), host.clone() + mapped = torch.ops._C.get_cuda_view_from_cpu_tensor(host) + assert mapped.device == device + device_offsets = torch.tensor(offsets, dtype=torch.int64, device=device) + cpu_offsets = torch.tensor(host_offsets, dtype=torch.int64, device=device) + words = torch.tensor(sizes, dtype=torch.int64, device=device) + plan = torch.zeros(8 + 5 * 256, dtype=torch.int32, device=device) + + buffers = (plan, arena, mapped, device_offsets, cpu_offsets, words) + invoke = partial(launch, buffers) + reset = partial(reset_data, arena, host, original_arena, original_host) + + # Separate loads/writebacks model the spare-slot policy; a swap is + # also valid when an incoming expert gives its host slot to a victim. + cases = ( + ("empty", 0), + ("paired", 1), + ("paired", 2), + ("paired", 4), + ("paired", 8), + ("paired", 20), + ("swap", 1), + ("swap", 4), + ) + if args.boundary_only: + cases = tuple(("swap", n) for n in (0, 1, 2, 3, 4, 5, 15, 16, 17, 128)) + for mode, count in cases: + cpu_plan = torch.zeros_like(plan, device="cpu") + ops = count * 2 if mode == "paired" else count + cpu_plan[0] = ops + for i in range(count): + entries = ( + [(i, i, i)] + if mode == "swap" + else [(i, i, -1), (slots - count + i, -1, host_slots - count + i)] + ) + for j, (slot, load, back) in enumerate(entries): + op = i * len(entries) + j + cpu_plan[8 + op * 5 : 11 + op * 5] = torch.tensor( + [slot, load, back] + ) + plan.copy_(cpu_plan) + reset() + invoke("reference") + torch.accelerator.synchronize() + oracle_arena, oracle_host = arena.clone(), host.clone() + graphs = {} + for name in POLICIES: + reset() + invoke(name) + torch.accelerator.synchronize() + assert torch.equal(arena, oracle_arena), ( + rank, + layer, + mode, + count, + name, + "GPU", + ) + assert torch.equal(host, oracle_host), ( + rank, + layer, + mode, + count, + name, + "host", + ) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(args.repeats): + invoke(name) + graphs[name] = graph + reset() + graph.replay() + torch.accelerator.synchronize() + # Even swap counts restore the original data after a graph. + expected_arena = ( + original_arena + if mode == "swap" and args.repeats % 2 == 0 + else oracle_arena + ) + expected_host = ( + original_host + if mode == "swap" and args.repeats % 2 == 0 + else oracle_host + ) + assert torch.equal(arena, expected_arena), (rank, name, "graph GPU") + assert torch.equal(host, expected_host), (rank, name, "graph host") + reset() + arena.add_(19) + host.add_(23) + for _ in range(args.repeats): + invoke("reference") + torch.accelerator.synchronize() + changed_arena, changed_host = arena.clone(), host.clone() + reset() + arena.add_(19) + host.add_(23) + graph.replay() + torch.accelerator.synchronize() + assert torch.equal(arena, changed_arena), (rank, name, "changed GPU") + assert torch.equal(host, changed_host), (rank, name, "changed host") + del changed_arena, changed_host + + timings = {name: [] for name in POLICIES} + for sample in range(args.samples): + order = list(POLICIES) + if sample % 2: + order.reverse() + for name in order: + reset() + dist.barrier() + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + graphs[name].replay() + end.record() + end.synchronize() + elapsed = torch.tensor( + start.elapsed_time(end) * 1000 / args.repeats, + dtype=torch.float64, + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX) + timings[name].append(elapsed.item()) + case = { + "layer": layer, + "words": sizes, + "mode": mode, + "count": count, + "plan_ops": ops, + "critical_rank_us": { + name: summarize(values) for name, values in timings.items() + }, + } + output["cases"].append(case) + output["correctness"].append( + { + "layer": layer, + "mode": mode, + "count": count, + "eager_and_graph_bit_exact": True, + } + ) + if rank == 0: + print(json.dumps(case), flush=True) + args.output.write_text(json.dumps(output, indent=2)) + del graphs, oracle_arena, oracle_host + del invoke, reset, buffers, arena, host, mapped, original_arena, original_host + torch.accelerator.empty_cache() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/docs/design/flashnext_decode_verifier_16gb_20261014.md b/docs/design/flashnext_decode_verifier_16gb_20261014.md new file mode 100644 index 000000000..436b9bfeb --- /dev/null +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -0,0 +1,122 @@ +# Flash-Next verifier diagnosis on four 16GB V100 GPUs + +The objective is target verification below 16 ms at approximately 32K input +tokens, while retaining emitted-token speed and output quality. Verification +includes target forward, target sampling and state update. Draft proposal is +reported separately. Model qualification of the candidate is in progress; +the operator results below do not establish the 16-ms objective. + +## Fixed workload + +The reference source is `6f3557c973b4602905037e086fd4acba9b0d867f`. +Hardware is four V100-SXM2-16GB GPUs with a full NV2 mesh, driver 570.211.01, +CUDA 12.8, Torch 2.10.0+cu128 and Python 3.12.3. The model is Flash-Next +GSQ-RCO IQ3_S with Q8_0 MTP4, greedy draft sampling, TP4 and one request. +Target sampling is temperature 0.6, top-p 0.95, top-k 20, seed 20261014; +thinking and EOS remain enabled. + +Context capacity remains 262144 **input tokens**, chunk size 8192, prefix +caching enabled, and FULL target decode graphs at M5. Expert storage retains +the configured 80% GPU / 20% host placement. Target QSA history uses host +E4M3 with an 8192-token device hot cache; draft history uses host FP16. +PLE remains on disk with its native lane and 512-MiB CPU row cache. +Measured prompts contain 574 and 32774 tokens, with 192 output tokens. +These measurements do not qualify 256K speed or concurrent requests. + +## Reference measurements + +Ordinary streaming HTTP excludes TTFT from decode. Metrics count actual +draft-verification steps, so the round estimator divides by draft count, +without subtracting a prefill step. Retain raw observations and first-run +variation; the HTTP estimate is not an independent verifier timer. + +| Input tokens | Mean ms/output token | Estimated ms/round | Tokens/round | +| --- | ---: | ---: | ---: | +| 574 | 9.110 | 18.710 | 2.054 | +| 32774 | 8.721 | 21.256 | 2.450 | + +CUDA-only Nsight graph-node capture observes 1120 target kernels on every +rank. Short/long target counts are 94/76, versus independent committed +verification counts of 93/75; retain the extra pipeline-end replay. Draft +graphs contain 70/71 kernels and 207/210 kernels. Requests are separated by +their prefill boundary, and TP replays are aligned by ordinal. + +The following means close the same GPU0 target-start to next-target-start +window at 32K. Node tracing perturbs latency; this is composition evidence. + +| Phase | Mean ms | +| --- | ---: | +| Target graph activity envelope | 22.124 | +| Target end to first draft graph | 0.862 | +| Dependent draft graph chain | 3.895 | +| Draft end to next target graph | 1.173 | +| Complete round | 28.053 | + +Expert copies contribute 2.437 ms of target kernel service, expert publish / +CPU-plan wait 1.206 ms, and QSA host gather 1.990 ms. Dense MMA contributes +2.272 ms and 94 complete HC boundaries contribute 2.151 ms. Service can +overlap; these values must not be added into an endpoint wall table. + +Explicit synchronized step timing retains every TP worker. At 32K, the two +reference requests have mean per-round maximum verifier CPU walls of +20.289 / 17.701 ms, p90 25.496 / 20.319 ms, and p99 40.704 / 25.687 ms. +Target sampling is approximately 0.42 ms and state update 0.02 ms. +Event collection, fences and interval-one logging perturb endpoint latency; +ordinary HTTP remains the accepted emitted-token speed measurement. + +## Candidate and operator evidence + +The existing expert copy kernel changes from 16 operation programs with +8 word splits, BLOCK1024/W8, to 4 operation programs with 64 splits, +BLOCK256/W4. This increases parallel work for small miss plans. Whole-word +storage, swap ordering and masked tails retain their exact byte contract. + +The paired TP4 screen uses real per-expert byte geometry from layers 0, 1, +17 and 47, with production 410 device / 136 host slot counts. All six +policies pass eager and changed-input graph equality over complete host and +device banks, including guards. Common one/two-miss cases improve roughly +6–20%; one layer's eight-miss case regresses 6.4%. End-to-end admission is +therefore required. Word lengths 255/256/257 and operation counts around +4 and 16, through 128, also pass four-rank byte checks. + +For QSA direct resolution (STAGE0), values that cannot be installed are +read later by the authoritative native attention reader. The candidate +omits the unused host read/decode and mixed-tile cached-value read. It +preserves tags, ownership, counters, resolved indices and consumed values. +Staging modes 1/2 retain their original value-read masks. + +Correctness support retains the existing valid page4/D256 geometry and +FP16/E4M3 storage. The changed subset is direct resolution with failed +installations or mixed hit/miss tiles. Shape admission and existing native +and staged fallbacks are unchanged. No numerical tolerance is relaxed. +The host KV suite passes 104 cases on both reference and candidate, +including protected-capacity fallback and changed-input graph replay. +The step-profiler suite passes 13 cases on GPU0. + +Existing tests had compared the full 32-row workspace against a smaller +active request and passed capacity-sized metadata to attention. The +original kernel reproduces the failure. Tests now slice the active rows; +the original and candidate pass the same corrected oracle. + +## Reproduce and retained negative paths + +Run the installed source benchmark on an idle, locked TP4 GPU set: + +```bash +.venv/bin/python -m torch.distributed.run --standalone --nproc_per_node 4 \ + benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py --output copy.json +.venv/bin/python -m torch.distributed.run --standalone --nproc_per_node 4 \ + benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py \ + --boundary-only --samples 1 --repeats 2 --output boundary.json +``` + +Use the prepared environment's Python and its normal native artifacts. +Raw Nsight report, SQLite, HTTP responses, per-rank logs and copy JSON are +retained in the task's external `flashnext-decode-verifier-9683-20261014` +report directory, outside Git. + +Optional NVTX scopes failed because the Python package is absent; CUDA-only +node tracing succeeds. A concurrent prefill compiler exhausted host headroom +and memguard terminated an early capture; the user chose compiler priority. +The serving CLI does not expose `--observability-config`; the existing +captured MTP profiling compatibility flags enable the diagnostic instead. diff --git a/tests/models/qwen4_exp/test_host_kv.py b/tests/models/qwen4_exp/test_host_kv.py index 09a514322..3f97cfa2e 100644 --- a/tests/models/qwen4_exp/test_host_kv.py +++ b/tests/models/qwen4_exp/test_host_kv.py @@ -9,6 +9,61 @@ from vllm.models.qwen4_exp.nvidia.ops.host_kv import HostQSAKV pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.uint8, torch.float16]) +@pytest.mark.parametrize("rows", [1, 5]) +def test_direct_resolve_protected_capacity_fallback_and_changed_graph(dtype, rows): + device = torch.device("cuda:0") + state = HostQSAKV( + 9, 256, 256, device, hot_tokens=16, rows=rows, width=2051, dtype=dtype + ) + key = torch.randn(2304, 1, 256, device=device, dtype=torch.float16) + value = torch.randn_like(key) + state.write(key, value, torch.arange(2304, device=device)) + torch.accelerator.synchronize() + cached = state.host[0, :, :16, 0].transpose(0, 1) + if dtype == torch.uint8: + cached = cached.view(torch.float8_e4m3fn).float() + cached = (cached * state.host_scales[:16, :, None]).half() + state.hot_values.copy_(cached) + state.tags.view(-1).copy_(torch.arange(4, device=device)) + state.page_slots[:4] = torch.arange(4, device=device) + before = state.hot_values.clone() + state.staging.fill_(float("nan")) + indices = torch.arange(2051, dtype=torch.int32, device=device).repeat(rows, 1) + table = torch.arange(9, dtype=torch.int32, device=device).view(1, -1) + requests = torch.zeros(rows, dtype=torch.int32, device=device) + positions = torch.full((rows,), 2050, dtype=torch.int64, device=device) + lengths = torch.full((1,), 2051, dtype=torch.int32, device=device) + + def run(): + state.resolve_hot(indices, table, requests, positions, lengths) + + def check(): + torch.accelerator.synchronize() + valid = (indices >= 0) & (indices <= positions[:, None]) + expected = torch.where(valid, -1, -2) + expected = torch.where(valid & (indices < 16), indices, expected) + assert torch.equal(state.resolved[:rows], expected) + assert torch.equal(state.hot_values, before) + assert torch.isnan(state.staging).all() + assert not (state.page_slots == -2).any() + assert not (state.stamps == -1).any() + hits, misses, contended = state.stats.tolist() + assert hits > 0 and misses > 0 and contended > 0 + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + indices[:, -3:] = torch.tensor([-1, 2047, 2050], device=device) + positions.fill_(2047) + for _ in range(3): + graph.replay() + check() + state.staging.zero_() + + @pytest.mark.parametrize("dtype", [torch.uint8, torch.float16]) def test_shared_physical_page_rewrite_updates_every_logical_alias(dtype): device = torch.device("cuda:0") @@ -82,8 +137,8 @@ def test_host_gather_collisions_rejection_and_graph(rows, page, dtype): indices, table, requests, positions, lengths ) else: - actual_k, actual_v = state.staging.unbind(1) - remap = state.remapped + actual_k, actual_v = state.staging[:rows].unbind(1) + remap = state.remapped[:rows] torch.accelerator.synchronize() cpu_indices = indices.cpu().long() valid = ( @@ -240,11 +295,11 @@ def test_host_attention_causal_boundary_and_padding(visible): k, v, remapped, - state.table, - state.requests, + state.table[:5], + state.requests[:5], output, - query_positions=state.positions, - sequence_lengths=state.lengths, + query_positions=state.positions[:5], + sequence_lengths=state.lengths[:5], ) scores = ( torch.einsum("mhd,mkd->mhk", query.float(), k[:, : len(selected), 0].float()) @@ -255,7 +310,9 @@ def test_host_attention_causal_boundary_and_padding(visible): ) assert torch.isfinite(output).all() assert ((output.float() - reference).norm() / reference.norm()).item() < 0.002 - assert torch.equal(state.lengths, torch.full_like(state.lengths, len(selected))) + assert torch.equal( + state.lengths[:5], torch.full_like(state.lengths[:5], len(selected)) + ) def test_host_attention_empty_padded_rows(): @@ -276,11 +333,11 @@ def test_host_attention_empty_padded_rows(): k, v, remapped, - state.table, - state.requests, + state.table[:5], + state.requests[:5], output, - query_positions=state.positions, - sequence_lengths=state.lengths, + query_positions=state.positions[:5], + sequence_lengths=state.lengths[:5], ) assert torch.isfinite(output).all() assert torch.count_nonzero(output).item() == 0 @@ -327,11 +384,11 @@ def test_direct_host_attention_matches_staging_on_misses_and_replay( k, v, remap, - state.table, - state.requests, + state.table[:rows], + state.requests[:rows], reference, - query_positions=state.positions, - sequence_lengths=state.lengths, + query_positions=state.positions[:rows], + sequence_lengths=state.lengths[:rows], output_gate=gate, ) host_qsa_attention( diff --git a/tests/v1/worker/test_step_profiler.py b/tests/v1/worker/test_step_profiler.py index 620646250..6f8bdd910 100644 --- a/tests/v1/worker/test_step_profiler.py +++ b/tests/v1/worker/test_step_profiler.py @@ -106,6 +106,39 @@ def test_disabled_context_never_allocates_events(monkeypatch): assert profiler.calls == 0 and not profiler.totals +@pytest.mark.parametrize("role", ["runner", "runner_v2", "proposer"]) +def test_v2_reports_nonzero_tp_rank_timings(monkeypatch, role): + monkeypatch.setattr( + "vllm.v1.spec_decode.profiling.is_last_pp_first_tp_rank", lambda: False + ) + config = SimpleNamespace( + observability_config=ObservabilityConfig( + step_profiler=StepProfilerConfig(enabled=True, interval=1) + ), + speculative_config=SimpleNamespace(method="mtp"), + ) + messages = [] + profiler = create_step_profiler( + config, + torch.device("cuda"), + role=role, + logger=SimpleNamespace(info=lambda msg, *args: messages.append(msg % args)), + ) + if role == "runner_v2": + profiler.report( + dict( + events=[("target_forward", Event(7), Event())], + num_tokens=5, + num_draft_tokens=4, + target_verifier_wall_cpu=8, + total_wall_start=0, + ) + ) + assert any("target_forward=7.000" in message for message in messages) + else: + assert not profiler.report_rank() + + @pytest.mark.parametrize("reverse", [False, True]) def test_config_snapshot_priority_serialization_and_hash(monkeypatch, reverse): monkeypatch.setenv("VLLM_SM70_MTP_PROFILE", "1") diff --git a/vllm/model_executor/layers/quantization/gguf_expert_cache.py b/vllm/model_executor/layers/quantization/gguf_expert_cache.py index b274f9bdf..99ff21faf 100644 --- a/vllm/model_executor/layers/quantization/gguf_expert_cache.py +++ b/vllm/model_executor/layers/quantization/gguf_expert_cache.py @@ -38,7 +38,9 @@ SHARDS = ("w1", "w3", "w2") _ENTRY_HEADER = 8 _RING_ENTRIES = 1024 _MAX_OPS = 256 -_COPY_PROGRAMS = (16, 8) +# Small miss plans need more word splits than operation programs to expose +# enough independent mapped-host transfers. See the paired TP4 copy screen. +_COPY_PROGRAMS = (4, 64) _LAYER = re.compile(r"layers\.(\d+)\.") _MEMLOG = envs.VLLM_EXPERT_CACHE_MEMLOG @@ -582,8 +584,8 @@ class ExpertCache: i, OPS_PROGRAMS=_COPY_PROGRAMS[0], SPLIT=_COPY_PROGRAMS[1], - BLOCK=1024, - num_warps=8, + BLOCK=256, + num_warps=4, ) out = torch.empty_like(ids) _remap[(1,)]( diff --git a/vllm/models/qwen4_exp/nvidia/ops/host_kv.py b/vllm/models/qwen4_exp/nvidia/ops/host_kv.py index c01c47222..03d2957eb 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/host_kv.py +++ b/vllm/models/qwen4_exp/nvidia/ops/host_kv.py @@ -343,11 +343,19 @@ def _gather( ) hot_tokens = tl.reshape(slots[:, None] * 4 + page_lanes[None, :], (16,)) token_hits = tl.reshape(tl.broadcast_to(hit[:, None], (4, 4)), (16,)) + token_installs = tl.reshape(tl.broadcast_to(install[:, None], (4, 4)), (16,)) + # Direct readers consume existing hot values themselves, and read a + # failed installation from authoritative history. Resolve only needs + # values that it actually writes to a new hot slot or staging. + need_cached = token_hits & (STAGE > 0) + need_history = valid & ~token_hits + if STAGE == 0: + need_history &= token_installs cached = tl.load( Hot + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim + dims[None, None, :], - token_hits[:, None, None], + need_cached[:, None, None], other=0, ) codes = tl.load( @@ -359,13 +367,13 @@ def _gather( ) * Dim + dims[None, None, :], - (valid & ~token_hits)[:, None, None], + need_history[:, None, None], other=0, ) if FP8: scales = tl.load( Scales + tokens[:, None] * 2 + kv[None, :], - (valid & ~token_hits)[:, None], + need_history[:, None], other=0, ) decoded = (fp8_e4m3fn_bits_to_fp32_bitcast(codes) * scales[:, :, None]).to( @@ -374,7 +382,6 @@ def _gather( else: decoded = codes.to(tl.float16) values = tl.where(token_hits[:, None, None], cached, decoded) - token_installs = tl.reshape(tl.broadcast_to(install[:, None], (4, 4)), (16,)) tl.store( Hot + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim diff --git a/vllm/v1/spec_decode/profiling.py b/vllm/v1/spec_decode/profiling.py index f4abdb330..bc92fd03f 100644 --- a/vllm/v1/spec_decode/profiling.py +++ b/vllm/v1/spec_decode/profiling.py @@ -108,5 +108,7 @@ def create_step_profiler( interval=policy.interval, report=_REPORTS[role], logger=logger, - report_rank=is_last_pp_first_tp_rank, + # V2 verification timings must retain TP stragglers. Worker log + # prefixes identify each rank; ordinary inference has no context. + report_rank=(lambda: True) if role == "runner_v2" else is_last_pp_first_tp_rank, ) From c80c2f4a74a2df1a640b1aa11f160c1f8799b01f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 04:42:46 +0800 Subject: [PATCH 2/5] [Perf] Parallelize Flash-Next expert copies and skip empty KV installs Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../kernels/benchmark_sm70_gguf_cache_copy.py | 33 ++++-- ...flashnext_decode_verifier_16gb_20261014.md | 39 +++++-- tests/models/qwen4_exp/test_host_kv.py | 58 ++++++++++ .../layers/quantization/gguf_expert_cache.py | 11 +- vllm/models/qwen4_exp/nvidia/ops/host_kv.py | 108 +++++++++--------- 5 files changed, 172 insertions(+), 77 deletions(-) diff --git a/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py index 7fe5eb789..b140fb8a9 100644 --- a/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py +++ b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py @@ -17,19 +17,23 @@ from vllm import _custom_ops # noqa: F401 -- register mapped-view operator from vllm.model_executor.layers.quantization.gguf_expert_cache import _apply_plan POLICIES = { - "reference": (16, 8, 1024, 8), - "o2s64b256w4": (2, 64, 256, 4), - "o4s32b256w4": (4, 32, 256, 4), - "o4s64b256w4": (4, 64, 256, 4), - "o8s16b512w4": (8, 16, 512, 4), - "o2s32b512w4": (2, 32, 512, 4), + "reference": (16, 8, 1024, 8, False), + "o2s64b256w4": (2, 64, 256, 4, False), + "o4s32b256w4": (4, 32, 256, 4, False), + "o4s64b256w4": (4, 64, 256, 4, False), + "o8s16b512w4": (8, 16, 512, 4, False), + "o2s32b512w4": (2, 32, 512, 4, False), + "o2s32b256w4p3": (2, 32, 256, 4, True), + "o4s32b256w4p3": (4, 32, 256, 4, True), + "o4s64b256w4p3": (4, 64, 256, 4, True), } def launch(buffers, policy): - programs, split, block, warps = POLICIES[policy] + programs, split, block, warps, parallel_shards = POLICIES[policy] plan, arena, mapped, device_offsets, host_offsets, words = buffers - _apply_plan[(programs, split)]( + grid = (programs, split, 3) if parallel_shards else (programs, split) + _apply_plan[grid]( plan, arena, mapped, @@ -40,6 +44,7 @@ def launch(buffers, policy): OPS_PROGRAMS=programs, SPLIT=split, BLOCK=block, + PARALLEL_SHARDS=parallel_shards, num_warps=warps, ) @@ -66,7 +71,11 @@ def main(): parser.add_argument("--samples", type=int, default=12) parser.add_argument("--repeats", type=int, default=16) parser.add_argument("--boundary-only", action="store_true") + parser.add_argument("--policies", nargs="+", choices=POLICIES) args = parser.parse_args() + policies = ( + {name: POLICIES[name] for name in args.policies} if args.policies else POLICIES + ) rank = int(os.environ["LOCAL_RANK"]) torch.accelerator.set_device_index(rank) dist.init_process_group("gloo") @@ -97,7 +106,7 @@ def main(): layouts = {-1: {"w1": 255, "w3": 256, "w2": 257}} output = { "world_size": 4, - "policies": POLICIES, + "policies": policies, "samples": args.samples, "repeats": args.repeats, "cases": [], @@ -165,7 +174,7 @@ def main(): torch.accelerator.synchronize() oracle_arena, oracle_host = arena.clone(), host.clone() graphs = {} - for name in POLICIES: + for name in policies: reset() invoke(name) torch.accelerator.synchronize() @@ -222,9 +231,9 @@ def main(): assert torch.equal(host, changed_host), (rank, name, "changed host") del changed_arena, changed_host - timings = {name: [] for name in POLICIES} + timings = {name: [] for name in policies} for sample in range(args.samples): - order = list(POLICIES) + order = list(policies) if sample % 2: order.reverse() for name in order: diff --git a/docs/design/flashnext_decode_verifier_16gb_20261014.md b/docs/design/flashnext_decode_verifier_16gb_20261014.md index 436b9bfeb..47d3a8aab 100644 --- a/docs/design/flashnext_decode_verifier_16gb_20261014.md +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -67,32 +67,45 @@ ordinary HTTP remains the accepted emitted-token speed measurement. ## Candidate and operator evidence The existing expert copy kernel changes from 16 operation programs with -8 word splits, BLOCK1024/W8, to 4 operation programs with 64 splits, -BLOCK256/W4. This increases parallel work for small miss plans. Whole-word +8 word splits, BLOCK1024/W8, to 4 operation programs with 32 word splits +and 3 parallel projections, BLOCK256/W4. This increases parallel work for +small miss plans without serializing all three projections in each CTA. Whole-word storage, swap ordering and masked tails retain their exact byte contract. The paired TP4 screen uses real per-expert byte geometry from layers 0, 1, 17 and 47, with production 410 device / 136 host slot counts. All six policies pass eager and changed-input graph equality over complete host and -device banks, including guards. Common one/two-miss cases improve roughly -6–20%; one layer's eight-miss case regresses 6.4%. End-to-end admission is -therefore required. Word lengths 255/256/257 and operation counts around -4 and 16, through 128, also pass four-rank byte checks. +device banks, including guards. The initial 4-by-64 word split improves +common one/two-miss cases roughly 6–20%, with a 6.4% regression in one +eight-miss case. The projection-parallel screen additionally improves paired +miss plans by 0.5–17.5% versus that first candidate, while some swap cases +regress. End-to-end admission is required. Word lengths 255/256/257 and +operation counts around 4 and 16, through 128, also pass four-rank byte checks. For QSA direct resolution (STAGE0), values that cannot be installed are read later by the authoritative native attention reader. The candidate -omits the unused host read/decode and mixed-tile cached-value read. It +omits the unused host read/decode and mixed-tile cached-value read, and +skips the entire value-materialization block when no page is installed. It preserves tags, ownership, counters, resolved indices and consumed values. -Staging modes 1/2 retain their original value-read masks. +Staging modes 1/2 retain their original value-read masks; staging writes +complete before the same release of cache metadata. Correctness support retains the existing valid page4/D256 geometry and FP16/E4M3 storage. The changed subset is direct resolution with failed installations or mixed hit/miss tiles. Shape admission and existing native and staged fallbacks are unchanged. No numerical tolerance is relaxed. -The host KV suite passes 104 cases on both reference and candidate, -including protected-capacity fallback and changed-input graph replay. +The first host KV suite passes 104 cases on both reference and candidate, +including protected-capacity fallback and changed-input graph replay. The +second candidate passes 108 host KV cases and 7 expert cache cases, including +probe-pointer signed/unsigned wrap, graph replay and decode/prefill interleaving. The step-profiler suite passes 13 cases on GPU0. +The first candidate's 32K second request improves maximum verifier CPU-wall +mean from 17.701 to 16.704 ms, p90 from 20.319 to 18.388 ms, and p99 from +25.687 to 19.113 ms. Its first request is 20.598 ms with an 87.139-ms maximum, +versus the reference 20.289-ms mean; retain this first-request variation. +The second candidate model measurement is pending. + Existing tests had compared the full 32-row workspace against a smaller active request and passed capacity-sized metadata to attention. The original kernel reproduces the failure. Tests now slice the active rows; @@ -120,3 +133,9 @@ node tracing succeeds. A concurrent prefill compiler exhausted host headroom and memguard terminated an early capture; the user chose compiler priority. The serving CLI does not expose `--observability-config`; the existing captured MTP profiling compatibility flags enable the diagnostic instead. + +An eight-warp direct-resolve screen regresses FP8 M5 churn by approximately +30% and other M5 cases by approximately 12%; retain W4. A scalar reservation +of four CLOCK probe ranges improves FP8 M5 only 0–2%, with noisy FP16 +regressions; retain the existing per-page atomics. These rejected choices do +not change the final runtime or introduce startup switches. diff --git a/tests/models/qwen4_exp/test_host_kv.py b/tests/models/qwen4_exp/test_host_kv.py index 3f97cfa2e..e1d913b6e 100644 --- a/tests/models/qwen4_exp/test_host_kv.py +++ b/tests/models/qwen4_exp/test_host_kv.py @@ -64,6 +64,64 @@ def test_direct_resolve_protected_capacity_fallback_and_changed_graph(dtype, row state.staging.zero_() +@pytest.mark.parametrize("dtype", [torch.uint8, torch.float16]) +@pytest.mark.parametrize("initial_hand", [2**31 - 16, -16]) +def test_direct_resolve_probe_range_wrap_and_changed_graph(dtype, initial_hand): + device = torch.device("cuda:0") + state = HostQSAKV(2, 64, 256, device, hot_tokens=64, rows=5, width=16, dtype=dtype) + key = torch.randn(128, 1, 256, device=device, dtype=torch.float16) + state.write(key, -key, torch.arange(128, device=device)) + indices = torch.arange(80, dtype=torch.int32, device=device).view(5, 16) + table = torch.arange(2, dtype=torch.int32, device=device).view(1, -1) + requests = torch.zeros(5, dtype=torch.int32, device=device) + positions = torch.full((5,), 127, dtype=torch.int64, device=device) + lengths = torch.full((1,), 128, dtype=torch.int32, device=device) + + def run(): + state.tags.fill_(-1) + state.page_slots.fill_(-1) + state.stamps.zero_() + state.hands.fill_(initial_hand) + state.resolve_hot(indices, table, requests, positions, lengths) + + def check(): + torch.accelerator.synchronize() + resolved = state.resolved[:5] + hits = resolved >= 0 + assert hits.any() and (resolved == -1).any() + assert not (state.page_slots == -2).any() + assert not (state.stamps == -1).any() + probe_count = (state.hands.item() - initial_hand) % 2**32 + assert probe_count > 0 and probe_count % 16 == 0 + tokens = indices.cpu().long() + codes = state.host.view(128 // 64, 2, 64, 256)[tokens // 64, :, tokens % 64] + if dtype == torch.uint8: + expected = ( + codes.view(torch.float8_e4m3fn).float() + * state.host_scales[tokens, :, None] + ).half() + else: + expected = codes + actual = state.hot_values[resolved.clamp_min(0).long()].cpu() + assert torch.equal(actual[hits.cpu()], expected[hits.cpu()]) + pages = state.tags.view(-1) + live = pages >= 0 + assert torch.equal( + state.page_slots[pages[live].long()], + torch.arange(pages.numel(), device=device, dtype=torch.int32)[live], + ) + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + indices.add_(48) + for _ in range(3): + graph.replay() + check() + + @pytest.mark.parametrize("dtype", [torch.uint8, torch.float16]) def test_shared_physical_page_rewrite_updates_every_logical_alias(dtype): device = torch.device("cuda:0") diff --git a/vllm/model_executor/layers/quantization/gguf_expert_cache.py b/vllm/model_executor/layers/quantization/gguf_expert_cache.py index 99ff21faf..6f112ac9d 100644 --- a/vllm/model_executor/layers/quantization/gguf_expert_cache.py +++ b/vllm/model_executor/layers/quantization/gguf_expert_cache.py @@ -38,9 +38,9 @@ SHARDS = ("w1", "w3", "w2") _ENTRY_HEADER = 8 _RING_ENTRIES = 1024 _MAX_OPS = 256 -# Small miss plans need more word splits than operation programs to expose -# enough independent mapped-host transfers. See the paired TP4 copy screen. -_COPY_PROGRAMS = (4, 64) +# Split both words and projections to expose independent mapped-host +# transfers for small miss plans. See the paired TP4 copy screen. +_COPY_PROGRAMS = (4, 32, 3) _LAYER = re.compile(r"layers\.(\d+)\.") _MEMLOG = envs.VLLM_EXPERT_CACHE_MEMLOG @@ -111,6 +111,7 @@ def _apply_plan( OPS_PROGRAMS: tl.constexpr, SPLIT: tl.constexpr, BLOCK: tl.constexpr, + PARALLEL_SHARDS: tl.constexpr = False, ): """Write back, then load, the planned experts (int64 words). @@ -123,7 +124,8 @@ def _apply_plan( slot = tl.load(plan + 8 + op * 5).to(tl.int64) load = tl.load(plan + 9 + op * 5).to(tl.int64) back = tl.load(plan + 10 + op * 5).to(tl.int64) - for shard in tl.static_range(3): + for shard_index in tl.static_range(1 if PARALLEL_SHARDS else 3): + shard = tl.program_id(2) if PARALLEL_SHARDS else shard_index n = tl.load(words + layer * 3 + shard) device = tl.load(device_offsets + layer * 3 + shard) + slot * n host_base = tl.load(host_offsets + layer * 3 + shard) @@ -585,6 +587,7 @@ class ExpertCache: OPS_PROGRAMS=_COPY_PROGRAMS[0], SPLIT=_COPY_PROGRAMS[1], BLOCK=256, + PARALLEL_SHARDS=True, num_warps=4, ) out = torch.empty_like(ids) diff --git a/vllm/models/qwen4_exp/nvidia/ops/host_kv.py b/vllm/models/qwen4_exp/nvidia/ops/host_kv.py index 03d2957eb..9059e19de 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/host_kv.py +++ b/vllm/models/qwen4_exp/nvidia/ops/host_kv.py @@ -347,49 +347,66 @@ def _gather( # Direct readers consume existing hot values themselves, and read a # failed installation from authoritative history. Resolve only needs # values that it actually writes to a new hot slot or staging. - need_cached = token_hits & (STAGE > 0) - need_history = valid & ~token_hits - if STAGE == 0: - need_history &= token_installs - cached = tl.load( - Hot - + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim - + dims[None, None, :], - need_cached[:, None, None], - other=0, - ) - codes = tl.load( - History - + ( - (tl.maximum(blocks, 0).to(tl.int64)[:, None, None] * 2 + kv[None, :, None]) - * Page - + (tl.maximum(logical, 0) % Page)[:, None, None] - ) - * Dim - + dims[None, None, :], - need_history[:, None, None], - other=0, - ) - if FP8: - scales = tl.load( - Scales + tokens[:, None] * 2 + kv[None, :], - need_history[:, None], + # A failed direct installation has no values to materialize. + if STAGE > 0 or tl.sum(install.to(tl.int32), 0) > 0: + need_cached = token_hits & (STAGE > 0) + need_history = valid & ~token_hits + if STAGE == 0: + need_history &= token_installs + cached = tl.load( + Hot + + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim + + dims[None, None, :], + need_cached[:, None, None], other=0, ) - decoded = (fp8_e4m3fn_bits_to_fp32_bitcast(codes) * scales[:, :, None]).to( - tl.float16 + codes = tl.load( + History + + ( + ( + tl.maximum(blocks, 0).to(tl.int64)[:, None, None] * 2 + + kv[None, :, None] + ) + * Page + + (tl.maximum(logical, 0) % Page)[:, None, None] + ) + * Dim + + dims[None, None, :], + need_history[:, None, None], + other=0, ) - else: - decoded = codes.to(tl.float16) - values = tl.where(token_hits[:, None, None], cached, decoded) - tl.store( - Hot - + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim - + dims[None, None, :], - values, - token_installs[:, None, None], - ) - tl.debug_barrier() + if FP8: + scales = tl.load( + Scales + tokens[:, None] * 2 + kv[None, :], + need_history[:, None], + other=0, + ) + decoded = (fp8_e4m3fn_bits_to_fp32_bitcast(codes) * scales[:, :, None]).to( + tl.float16 + ) + else: + decoded = codes.to(tl.float16) + values = tl.where(token_hits[:, None, None], cached, decoded) + tl.store( + Hot + + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim + + dims[None, None, :], + values, + token_installs[:, None, None], + ) + if STAGE > 0: + store_values = columns < Padded + if STAGE == 2: + store_values &= valid & ~(token_hits | token_installs) + tl.store( + Out + + ((row * 2 + kv[None, :, None]) * Padded + columns[:, None, None]) + * Dim + + dims[None, None, :], + values, + store_values[:, None, None], + ) + tl.debug_barrier() tl.atomic_cas(Tags + slots, tl.where(install, old_tag, -2), first, sem="release") tl.atomic_cas(PageSlots + first, tl.where(install, -2, -3), slots, sem="release") tl.atomic_cas( @@ -402,17 +419,6 @@ def _gather( tl.where(valid, tl.where(token_hits | token_installs, hot_tokens, -1), -2), columns < Width, ) - if STAGE > 0: - store_values = columns < Padded - if STAGE == 2: - store_values &= valid & ~(token_hits | token_installs) - tl.store( - Out - + ((row * 2 + kv[None, :, None]) * Padded + columns[:, None, None]) * Dim - + dims[None, None, :], - values, - store_values[:, None, None], - ) counter = (row * tl.cdiv(Width, 16) + tile) * 3 counters = tl.arange(0, 4) hits = tl.sum((valid & token_hits).to(tl.int64), 0) From 708692b6a0727411db54574f4345b5a94d952da9 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:11:11 +0800 Subject: [PATCH 3/5] [Perf] Reuse LL projections at PLE HyperConnection boundaries Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- ...flashnext_decode_verifier_16gb_20261014.md | 81 +++++++++++++++++-- .../kernels/test_sm70_hc_sharded_fallback.py | 66 ++++++++++----- .../qwen4_exp/test_sm70_hcx_dispatch.py | 8 +- vllm/models/qwen4_exp/nvidia/model.py | 10 ++- 4 files changed, 137 insertions(+), 28 deletions(-) diff --git a/docs/design/flashnext_decode_verifier_16gb_20261014.md b/docs/design/flashnext_decode_verifier_16gb_20261014.md index 47d3a8aab..2b78b4d7e 100644 --- a/docs/design/flashnext_decode_verifier_16gb_20261014.md +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -59,7 +59,7 @@ overlap; these values must not be added into an endpoint wall table. Explicit synchronized step timing retains every TP worker. At 32K, the two reference requests have mean per-round maximum verifier CPU walls of -20.289 / 17.701 ms, p90 25.496 / 20.319 ms, and p99 40.704 / 25.687 ms. +20.354 / 17.701 ms, p90 25.496 / 20.319 ms, and p99 40.704 / 25.687 ms. Target sampling is approximately 0.42 ms and state update 0.02 ms. Event collection, fences and interval-one logging perturb endpoint latency; ordinary HTTP remains the accepted emitted-token speed measurement. @@ -103,8 +103,11 @@ The step-profiler suite passes 13 cases on GPU0. The first candidate's 32K second request improves maximum verifier CPU-wall mean from 17.701 to 16.704 ms, p90 from 20.319 to 18.388 ms, and p99 from 25.687 to 19.113 ms. Its first request is 20.598 ms with an 87.139-ms maximum, -versus the reference 20.289-ms mean; retain this first-request variation. -The second candidate model measurement is pending. +versus the reference 20.354-ms mean; retain this first-request variation. +The second candidate's first/second 32K means are 20.054 / 16.864 ms; its +second-request p90/p99 are 18.373 / 25.294 ms. This does not establish a +further gain or satisfy 16 ms. Ordinary candidate HTTP qualification remains +pending while optimization continues. Existing tests had compared the full 32-row workspace against a smaller active request and passed capacity-sized metadata to attention. The @@ -134,8 +137,70 @@ and memguard terminated an early capture; the user chose compiler priority. The serving CLI does not expose `--observability-config`; the existing captured MTP profiling compatibility flags enable the diagnostic instead. -An eight-warp direct-resolve screen regresses FP8 M5 churn by approximately -30% and other M5 cases by approximately 12%; retain W4. A scalar reservation -of four CLOCK probe ranges improves FP8 M5 only 0–2%, with noisy FP16 -regressions; retain the existing per-page atomics. These rejected choices do -not change the final runtime or introduce startup switches. +Initial direct-resolve screens retained one KV epoch across replays. They +measure a protected-cache regime: selected pages remain pinned and later +misses fall back instead of installing. They do not represent normal decode +churn. The corrected fixture advances the epoch before every resolve, +matching the protection transition of a decode KV write. + +With advancing epochs, the adversarial FP8 M5 churn screen measures W4/W8 +at 693 / 660 us, while shared/overflow selections regress about 10% at W8. +Eight-token tiles improve that churn case only about 2%; four-token tiles +regress both shared selections and churn. Retain the existing W4, 16-token +tile and per-page CLOCK atomics. Register counts are 127 for W4/T16, 76 for +W8/T16 and 72 for W4/T8, with no spills. This is a synthetic operator screen; +register reduction alone does not qualify model speed. No new tiling or warp +startup switches are added. + +An exact-word HC gather experiment encoded Half storage words as integers in +FP32 and used the existing peer reduction. Every Half bit pattern and changed +graph passed exact equality, but the complete M5 boundary regressed from +69.824 to 79.296 us; the model's second 32K verifier mean was 17.472 ms. The +helper was reverted. Collective substitution alone is not a speedup. + +The existing 32768-token hot-cache setting also failed admission: the first +real request needed a 144-MiB allocation with only 75.94 MiB free. It retained +the expert ratio and context capacity but cannot run on this 16GB deployment. +The 16384-token setting runs successfully and measures first/second 32K +verifier means of 18.435 / 16.059 ms, with second-request p90/p99 of +17.209 / 18.848 ms. It is a separate configuration screen and still misses +16 ms. Arithmetic, Python evaluation, two 32K retrieval tasks and a 1809-token +CPU explanation pass their task/EOS gates; the explanation has no six-word +phrase repeated more than twice. Ordinary throughput qualification is pending. + +The two target-graph NCCL gathers occur after the early PLE boundary. PLE +materializes the pending MoE residual before adding its result, then calls +standalone HC `mix()`. Retaining the six-dimensional HCX fusion pack at that +boundary forces unpacking and two gathers, because standalone LL projection +consumes the existing five-dimensional pack. A candidate leaves partial +residual ownership intact and prepares the LL bank at PLE boundaries instead +of the unused HCX bank. This changes one target boundary on the measured +model and allocates no duplicate bank. HCX schedules at other boundaries, +partial residual ownership, normalization and request state are unchanged. + +The actual layer-1 checkpoint screen on four ranks measures complete M5 +boundaries at 76.672 / 19.140 us for reference / LL storage. M1/M8/M20 are +also faster; M21 and M256 retain fallback with exact outputs. Native arithmetic +uses the existing FP16 HC error gate, rtol 0.003 and atol 0.003, fixed before +the screen. Both variants replay changed inputs exactly against their own +eager results, with finite outputs. CPU dispatch/lifetime tests pass 21 cases; +the corrected fallback suite passes 17 GPU cases across two rank orders, +both storage layouts, M5/M20/M21 and M512. The model screen retains HOT8192 +to separate this source change from cache tuning. Model quality and endpoint +measurements remain pending. + +The old fallback oracle omitted the physical rank and simulated peers by +scaling an already mixed local result. That does not model each peer's +different input columns and nonlinear gates; it fails the unchanged HCX +fallback too. The corrected oracle computes remote projections and columns, +includes the actual local packet, and retains exact assertions and changed +input graph checks. + +An asynchronous final verifier replay can finish after the HTTP stream. +The original before-counter read occasionally assigned that replay to the +next request. Retained API POST boundaries align all four workers and keep +every final replay with its preceding request; aligned counts differ from +independent committed draft metrics by at most one. Raw and aligned ledgers +are both preserved. New event probes wait for all-rank counters to remain +stable for 200 ms between requests. This does not change measured decode +timestamps or ordinary production requests. diff --git a/tests/kernels/test_sm70_hc_sharded_fallback.py b/tests/kernels/test_sm70_hc_sharded_fallback.py index 9ae2ae6db..477b7f6f0 100644 --- a/tests/kernels/test_sm70_hc_sharded_fallback.py +++ b/tests/kernels/test_sm70_hc_sharded_fallback.py @@ -16,35 +16,65 @@ pytestmark = pytest.mark.skipif( ) -@pytest.mark.parametrize("tokens", [5, 20, 512]) +@pytest.mark.parametrize("tokens", [5, 20, 21, 512]) @pytest.mark.parametrize("order", [(0, 1, 2, 3), (0, 2, 1, 3)]) -def test_hcx_sharded_fallback_graph_matches_rank_reorder(tokens, order, monkeypatch): - communicator = SimpleNamespace(status={"enabled": True}, order=order) - - def gather(value, dim): - return torch.cat([value * (rank + 1) for rank in range(4)], dim=dim) +@pytest.mark.parametrize("storage", ["hcx", "ll"]) +def test_hcx_sharded_fallback_graph_matches_rank_reorder( + tokens, order, storage, monkeypatch +): + communicator = SimpleNamespace( + status={"enabled": True}, order=order, apply=lambda *_: None + ) group = SimpleNamespace( + rank_in_group=0, device_communicator=SimpleNamespace(hc_ll_comm=communicator), - all_gather=gather, ) monkeypatch.setattr("vllm.distributed.parallel_state.get_tp_group", lambda: group) down_weight = torch.randn(336, 10240, device="cuda", dtype=torch.float16) * 0.002 up_weight = torch.randn(10240, 320, device="cuda", dtype=torch.float16) * 0.002 - down, up = pack_down(down_weight, 0), pack_up(up_weight, 0) - local_down, local_up = hc._unpack_hc_storage(down, up) - index = torch.tensor(order, device="cuda", dtype=torch.int64) + packs = [] + for rank in range(4): + if storage == "hcx": + packs.append((pack_down(down_weight, rank), pack_up(up_weight, rank))) + else: + packs.append( + ( + hc._pack_hc_batch_weight(down_weight, "down", rank), + hc._pack_hc_batch_weight(up_weight, "up", rank), + ) + ) + down, up = packs[0] + local_weights = [hc._unpack_hc_storage(*pack) for pack in packs] x = torch.randn(tokens, 10240, device="cuda", dtype=torch.float16) def reference(): - local = x @ local_down.T - packet = gather(local, -1).reshape(tokens, 4, 88).index_select(1, index) - lora = hc_silu(packet[:, :, :80].reshape(tokens, 320), 4) - injection = packet[:, 3, 80:84].contiguous() - local_gate = lora @ local_up.T - gate = gather(local_gate, -1).reshape(tokens, 4, 4, 640) - gate = gate.index_select(1, index).permute(0, 2, 1, 3).reshape(tokens, 10240) - return hc_gate_mix(x, gate, 4), injection + packets = [x @ weights[0].T for weights in local_weights] + lora = hc_silu(torch.cat([packet[:, :80] for packet in packets], dim=1), 4) + gate = torch.cat( + [ + (lora @ weights[1].T).reshape(tokens, 4, 640) + for weights in local_weights + ], + dim=2, + ).reshape(tokens, 10240) + return hc_gate_mix(x, gate, 4), packets[3][:, 80:84].contiguous() + + def gather(value, dim): + assert dim == 0 + # Model each peer's projection and input columns before its nonlinear + # mix, and keep the actual local packet as part of the collective. + if value.shape[1] == 88: + logical = [x @ weights[0].T for weights in local_weights] + else: + assert value.shape[1] == 640 + block, _ = reference() + logical = list(block.split(640, dim=1)) + physical = [logical[order.index(rank)] for rank in range(4)] + physical[group.rank_in_group] = value + return torch.cat(physical, dim=0) + + group.all_gather = gather # Compile the unchanged arithmetic before capture, as model warmup does. expected = reference() diff --git a/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py b/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py index 5decfd764..45eb39de1 100644 --- a/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py +++ b/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py @@ -216,7 +216,8 @@ def test_compiled_projection_resolves_actual_rows(monkeypatch): del library -def test_hcx_only_consumes_partial_producers(monkeypatch): +@pytest.mark.parametrize("ple", [False, True]) +def test_hcx_only_consumes_partial_producers(monkeypatch, ple): from vllm.v1.attention.backends import fa_utils monkeypatch.setattr(fa_utils, "get_flash_attn_version", lambda *a, **k: 2) @@ -237,6 +238,7 @@ def test_hcx_only_consumes_partial_producers(monkeypatch): def enable_partial_inputs(self, name, runtime=None): self.partial = True + self.runtime = runtime class MoE: def __init__(self): @@ -246,6 +248,7 @@ def test_hcx_only_consumes_partial_producers(monkeypatch): def __init__(self, index, moe): self.layer_idx = index self.layer_type = "linear_attention" + self.ple = object() if ple and index == 1 else None self.linear_attn = SimpleNamespace( out_proj=SimpleNamespace(reduce_results=True) ) @@ -268,9 +271,12 @@ def test_hcx_only_consumes_partial_producers(monkeypatch): assert model_module.enable_sm70_hcx(model, torch.device("cpu")) assert not model.layers[0].attn_hyper_connection.partial assert model.layers[1].attn_hyper_connection._hcx_moe_payload + assert (model.layers[1].attn_hyper_connection.runtime is None) == ple + assert all(layer.mlp_hyper_connection.runtime is not None for layer in model.layers) assert not model.layers[2].attn_hyper_connection.partial assert model.hyper_connection_mixer._hcx_moe_payload assert model.sm70_hcx_status["prepared_modules"] == 4 + assert model.sm70_hcx_status["ple_mix_modules"] == int(ple) def test_compiled_moe_payload_keeps_both_dependencies_live(): diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 99a0195bc..97daf96c3 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -533,6 +533,7 @@ def enable_sm70_hcx( min_m=1, max_m=8, prepared_modules=0, + ple_mix_modules=0, deferred_projections=0, output_projection_policy="fused" if fuse_output_projection else "separate", local_schedule=bool(local_schedule and runtime.enabled and runtime.full), @@ -614,7 +615,14 @@ def enable_sm70_hcx( register_moe_runner(prefix, runner) if previous_moe: layer.attn_hyper_connection._hcx_moe_payload = True - layer.attn_hyper_connection.enable_partial_inputs(prefix + ".attn", runtime) + # PLE materializes the pending residual before its addition and + # then calls mix(), so this boundary cannot use fused HCX. Keep + # its existing LL projection pack rather than unpacking HCX banks + # and gathering checkpoint projections on every decode step. + layer.attn_hyper_connection.enable_partial_inputs( + prefix + ".attn", runtime if layer.ple is None else None + ) + status["ple_mix_modules"] += int(layer.ple is not None) status["prepared_modules"] += 1 layer.mlp_hyper_connection.enable_partial_inputs(prefix + ".mlp", runtime) status["prepared_modules"] += 1 From 953835998302be2927d51843e30bc2f278dea0ad Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:31:04 +0800 Subject: [PATCH 4/5] [Perf] Withdraw PLE projection change after model screen Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../flashnext_decode_verifier_16gb_20261014.md | 16 +++++++++++++--- tests/models/qwen4_exp/test_sm70_hcx_dispatch.py | 8 +------- vllm/models/qwen4_exp/nvidia/model.py | 10 +--------- 3 files changed, 15 insertions(+), 19 deletions(-) diff --git a/docs/design/flashnext_decode_verifier_16gb_20261014.md b/docs/design/flashnext_decode_verifier_16gb_20261014.md index 2b78b4d7e..19251eec2 100644 --- a/docs/design/flashnext_decode_verifier_16gb_20261014.md +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -176,7 +176,7 @@ consumes the existing five-dimensional pack. A candidate leaves partial residual ownership intact and prepares the LL bank at PLE boundaries instead of the unused HCX bank. This changes one target boundary on the measured model and allocates no duplicate bank. HCX schedules at other boundaries, -partial residual ownership, normalization and request state are unchanged. +partial residual ownership, normalization and request state were unchanged. The actual layer-1 checkpoint screen on four ranks measures complete M5 boundaries at 76.672 / 19.140 us for reference / LL storage. M1/M8/M20 are @@ -186,8 +186,18 @@ the screen. Both variants replay changed inputs exactly against their own eager results, with finite outputs. CPU dispatch/lifetime tests pass 21 cases; the corrected fallback suite passes 17 GPU cases across two rank orders, both storage layouts, M5/M20/M21 and M512. The model screen retains HOT8192 -to separate this source change from cache tuning. Model quality and endpoint -measurements remain pending. +to separate this source change from cache tuning. + +That model screen rejects the PLE storage change. First/second 32K verifier +means are 18.077 / 17.048 ms; the second-request p90/p99 are 19.262 / 23.022 ms. +Committed draft steps rise to 91 for 192 output tokens, versus 75 in the +preceding HOT8192 candidate. Four integer tasks pass, but the CPU explanation +remains in reasoning at the fixed 2048-token limit. This does not establish a +quality regression by itself, but gives no qualification for promotion. +The runtime and dispatch-test changes are reverted to the preceding candidate. +Retain the corrected fallback oracle and its passed GPU coverage. The large +operator speedup and node-trace NCCL waits did not predict lower live verifier +cost; do not promote this route based on the microbenchmark alone. The old fallback oracle omitted the physical rank and simulated peers by scaling an already mixed local result. That does not model each peer's diff --git a/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py b/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py index 45eb39de1..5decfd764 100644 --- a/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py +++ b/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py @@ -216,8 +216,7 @@ def test_compiled_projection_resolves_actual_rows(monkeypatch): del library -@pytest.mark.parametrize("ple", [False, True]) -def test_hcx_only_consumes_partial_producers(monkeypatch, ple): +def test_hcx_only_consumes_partial_producers(monkeypatch): from vllm.v1.attention.backends import fa_utils monkeypatch.setattr(fa_utils, "get_flash_attn_version", lambda *a, **k: 2) @@ -238,7 +237,6 @@ def test_hcx_only_consumes_partial_producers(monkeypatch, ple): def enable_partial_inputs(self, name, runtime=None): self.partial = True - self.runtime = runtime class MoE: def __init__(self): @@ -248,7 +246,6 @@ def test_hcx_only_consumes_partial_producers(monkeypatch, ple): def __init__(self, index, moe): self.layer_idx = index self.layer_type = "linear_attention" - self.ple = object() if ple and index == 1 else None self.linear_attn = SimpleNamespace( out_proj=SimpleNamespace(reduce_results=True) ) @@ -271,12 +268,9 @@ def test_hcx_only_consumes_partial_producers(monkeypatch, ple): assert model_module.enable_sm70_hcx(model, torch.device("cpu")) assert not model.layers[0].attn_hyper_connection.partial assert model.layers[1].attn_hyper_connection._hcx_moe_payload - assert (model.layers[1].attn_hyper_connection.runtime is None) == ple - assert all(layer.mlp_hyper_connection.runtime is not None for layer in model.layers) assert not model.layers[2].attn_hyper_connection.partial assert model.hyper_connection_mixer._hcx_moe_payload assert model.sm70_hcx_status["prepared_modules"] == 4 - assert model.sm70_hcx_status["ple_mix_modules"] == int(ple) def test_compiled_moe_payload_keeps_both_dependencies_live(): diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 97daf96c3..99a0195bc 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -533,7 +533,6 @@ def enable_sm70_hcx( min_m=1, max_m=8, prepared_modules=0, - ple_mix_modules=0, deferred_projections=0, output_projection_policy="fused" if fuse_output_projection else "separate", local_schedule=bool(local_schedule and runtime.enabled and runtime.full), @@ -615,14 +614,7 @@ def enable_sm70_hcx( register_moe_runner(prefix, runner) if previous_moe: layer.attn_hyper_connection._hcx_moe_payload = True - # PLE materializes the pending residual before its addition and - # then calls mix(), so this boundary cannot use fused HCX. Keep - # its existing LL projection pack rather than unpacking HCX banks - # and gathering checkpoint projections on every decode step. - layer.attn_hyper_connection.enable_partial_inputs( - prefix + ".attn", runtime if layer.ple is None else None - ) - status["ple_mix_modules"] += int(layer.ple is not None) + layer.attn_hyper_connection.enable_partial_inputs(prefix + ".attn", runtime) status["prepared_modules"] += 1 layer.mlp_hyper_connection.enable_partial_inputs(prefix + ".mlp", runtime) status["prepared_modules"] += 1 From 46f5c8c197ed805a4700fdfecfdb29ea91ba39ad Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:26:25 +0800 Subject: [PATCH 5/5] [Perf] Keep encoded target KV in the Flash-Next hot cache Add opt-in E4M3 target hot storage with native and protected-reader support, preserving FP16 attention operands and the default cache policy. Qualify 32K mean verifier bounds, ordinary decode, quality and near-256K capacity on TP4 V100 16GB. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../ops/qsa_device_history_sm70.cu | 70 ++++-- docs/design/architecture/runtime_reference.md | 2 +- ...flashnext_decode_verifier_16gb_20261014.md | 201 +++++++++++++++++- .../qwen4_exp/test_device_kv_attention.py | 184 ++++++++++++++++ vllm/config/kernel.py | 5 + .../nvidia/ops/device_kv_attention.py | 18 +- vllm/models/qwen4_exp/nvidia/ops/host_kv.py | 56 ++++- .../qwen4_exp/nvidia/ops/host_kv_attention.py | 2 + .../qwen4_exp/nvidia/ops/host_kv_reader.py | 16 +- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 4 + vllm/models/qwen4_exp/nvidia/qsa.py | 12 +- 11 files changed, 545 insertions(+), 25 deletions(-) diff --git a/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu b/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu index d332eb2ce..42ee3e3a3 100644 --- a/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu +++ b/csrc/sm70_turbomind/ops/qsa_device_history_sm70.cu @@ -23,8 +23,9 @@ struct Args { const int64_t* positions; const int* lengths; const half* gate; - const half* hot; // [hot_tokens, 2, D] FP16, or null - const int* resolved; // [M, width] hot token, -1 host, -2 invalid + const half* hot; // [hot_tokens, 2, D] FP16 or encoded E4M3 bytes + const float* hot_scales; // [hot_tokens, 2], encoded hot storage only + const int* resolved; // [M, width] hot token, -1 host, -2 invalid half* out; float* partial; float* state; @@ -140,14 +141,24 @@ __device__ __forceinline__ uint4 decode8(const unsigned char* source, // Read one 8-value chunk of token K (kv=0) or V (kv=1): from the protected // hot cache when resolved there, else from the authoritative history, which // may be pinned host memory mapped into the device address space (UVA). -template +template __device__ __forceinline__ uint4 load_kv(const Args& a, int64_t slot, int hot, int kv, int dim, const float* lut) { if (slot < 0) return make_uint4(0, 0, 0, 0); if constexpr (HOT) { - if (hot >= 0) - return __ldg(reinterpret_cast( - a.hot + (int64_t(hot) * 2 + kv) * D + dim)); + if (hot >= 0) { + if constexpr (HOT_FP8) { + const auto* source = reinterpret_cast(a.hot) + + (int64_t(hot) * 2 + kv) * D + dim; + // Match the software bit decoder used to materialize FP16 hot values, + // including the FP16 rounding before QK/PV arithmetic. + return decode8_bits(__ldg(reinterpret_cast(source)), + __ldg(a.hot_scales + int64_t(hot) * 2 + kv)); + } else { + return __ldg(reinterpret_cast( + a.hot + (int64_t(hot) * 2 + kv) * D + dim)); + } + } } const int block = slot / a.page, off = slot % a.page; const int64_t element = ((int64_t(block) * 2 + kv) * a.page + off) * D + dim; @@ -158,7 +169,7 @@ __device__ __forceinline__ uint4 load_kv(const Args& a, int64_t slot, int hot, reinterpret_cast(a.history) + element)); } -template +template __global__ __launch_bounds__(256) void partial(Args a) { const int row = blockIdx.y, split = blockIdx.x; const int t = threadIdx.x, warp = t >> 5, lane = t & 31; @@ -211,7 +222,7 @@ __global__ __launch_bounds__(256) void partial(Args a) { for (int v = t; v < C * 32; v += 256) { const int token = v / 32, dim = (v % 32) * 8; *reinterpret_cast(&kv[token][dim]) = - load_kv(a, slots[token], hots[token], 0, dim, lut); + load_kv(a, slots[token], hots[token], 0, dim, lut); } __syncthreads(); // Without the hot cache, issue this thread's V loads now so their @@ -296,8 +307,8 @@ __global__ __launch_bounds__(256) void partial(Args a) { if constexpr (HOT) { for (int v = t; v < C * 32; v += 256) { const int token = v / 32, dim = (v % 32) * 8; - *reinterpret_cast(&kv[token][dim]) = - load_kv(a, slots[token], hots[token], 1, dim, lut); + *reinterpret_cast(&kv[token][dim]) = load_kv( + a, slots[token], hots[token], 1, dim, lut); } } else { #pragma unroll @@ -779,7 +790,8 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, std::optional hot = std::nullopt, std::optional resolved = std::nullopt, int64_t tiles = 1, bool prefill = false, - int64_t whole_rows = 0) { + int64_t whole_rows = 0, + std::optional hot_scales = std::nullopt) { const c10::cuda::CUDAGuard guard(q.device()); TORCH_CHECK(q.is_cuda() && q.scalar_type() == at::kHalf && q.dim() == 3 && q.size(0) > 0 && (whole_rows || q.size(0) <= 64) && @@ -813,16 +825,29 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, TORCH_CHECK(tensor.device() == q.device()); if (gate) TORCH_CHECK(gate->device() == q.device()); TORCH_CHECK(hot.has_value() == resolved.has_value()); + TORCH_CHECK(!hot_scales || (hot && history.scalar_type() == at::kByte), + "Encoded hot storage requires E4M3 authoritative history"); if (hot) { TORCH_CHECK(hot->is_cuda() && hot->device() == q.device() && - hot->scalar_type() == at::kHalf && hot->dim() == 3 && - hot->size(1) == 2 && hot->size(2) == 256 && + hot->scalar_type() == (hot_scales ? at::kByte : at::kHalf) && + hot->dim() == 3 && hot->size(1) == 2 && hot->size(2) == 256 && hot->is_contiguous()); + TORCH_CHECK( + reinterpret_cast(hot->data_ptr()) % (hot_scales ? 8 : 16) == + 0, + "QSA hot storage requires aligned vector loads"); TORCH_CHECK(resolved->device() == q.device() && resolved->scalar_type() == at::kInt && resolved->dim() == 2 && resolved->size(0) >= q.size(0) && resolved->size(1) >= indices.size(1) && resolved->stride(1) == 1); + if (hot_scales) { + TORCH_CHECK(hot_scales->device() == q.device() && + hot_scales->scalar_type() == at::kFloat && + hot_scales->dim() == 2 && + hot_scales->size(0) == hot->size(0) && + hot_scales->size(1) == 2 && hot_scales->is_contiguous()); + } } TORCH_CHECK(reinterpret_cast(q.data_ptr()) % 16 == 0 && q.stride(0) % 8 == 0 && q.stride(1) % 8 == 0); @@ -839,6 +864,7 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, a.lengths = lengths.data_ptr(); a.gate = gate ? reinterpret_cast(gate->data_ptr()) : nullptr; a.hot = hot ? reinterpret_cast(hot->data_ptr()) : nullptr; + a.hot_scales = hot_scales ? hot_scales->data_ptr() : nullptr; a.resolved = resolved ? resolved->data_ptr() : nullptr; a.rs0 = resolved ? resolved->stride(0) : 0; a.out = reinterpret_cast(out.data_ptr()); @@ -896,6 +922,8 @@ void device_qsa_out(torch::Tensor q, torch::Tensor history, } if (prefill && fp8 && !hot && a.tiles == 1) device_qsa::partial_fp8_lite<<>>(a); + else if (fp8 && hot_scales) + device_qsa::partial<<>>(a); else if (fp8 && hot) device_qsa::partial<<>>(a); else if (fp8) @@ -985,6 +1013,22 @@ TORCH_LIBRARY(vllm_sm70_qsa_device, m) { positions, lengths, out, gate, partial, state, hot, resolved); }); + m.def( + "run_hot_fp8(Tensor q, Tensor history, Tensor scales, Tensor indices, " + "Tensor table, Tensor requests, Tensor positions, Tensor lengths, " + "Tensor(a!) out, Tensor? gate, Tensor(b!) partial, Tensor(c!) state, " + "Tensor hot, Tensor resolved, Tensor hot_scales) -> ()"); + m.impl("run_hot_fp8", torch::kCUDA, + [](torch::Tensor q, torch::Tensor history, torch::Tensor scales, + torch::Tensor indices, torch::Tensor table, torch::Tensor requests, + torch::Tensor positions, torch::Tensor lengths, torch::Tensor out, + std::optional gate, torch::Tensor partial, + torch::Tensor state, torch::Tensor hot, torch::Tensor resolved, + torch::Tensor hot_scales) { + device_qsa_out(q, history, scales, indices, table, requests, + positions, lengths, out, gate, partial, state, hot, + resolved, 1, false, 0, hot_scales); + }); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index 66c5bc04f..6be2ea2ec 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -38,7 +38,7 @@ identify initialization/hash code; they do not evaluate its semantics. | `vllm.config.gdn_state.GdnStateConfig` | kernel_config.gdn.state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L39); graph_options: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L52) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L23) | | `vllm.config.gdn_state.GdnStateTraceConfig` | observability_config.gdn_state | resolve: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L99); compile_ignored_aliases: [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L142) | [vllm/config/gdn_state.py](../../../vllm/config/gdn_state.py#L59) | | `vllm.config.kernel.IrOpPriorityConfig` | kernel_config.ir_op_priority | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L84) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L67) | -| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1044) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) | +| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1049) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L619) | | `vllm.config.kernel.Sm70AwqConfig` | kernel_config.sm70_awq | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L338); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L317) | | `vllm.config.kernel.Sm70Fp8Config` | kernel_config.sm70_fp8 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L433); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L369) | | `vllm.config.kernel.Sm70GgufConfig` | kernel_config.sm70_gguf | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L498) | diff --git a/docs/design/flashnext_decode_verifier_16gb_20261014.md b/docs/design/flashnext_decode_verifier_16gb_20261014.md index 19251eec2..684c98eae 100644 --- a/docs/design/flashnext_decode_verifier_16gb_20261014.md +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -3,8 +3,11 @@ The objective is target verification below 16 ms at approximately 32K input tokens, while retaining emitted-token speed and output quality. Verification includes target forward, target sampling and state update. Draft proposal is -reported separately. Model qualification of the candidate is in progress; -the operator results below do not establish the 16-ms objective. +reported separately. The final opt-in encoded target hot cache measures an +all-round mean of 15.482–15.494 ms at 32K, including conservative bounds for +two missing graph starts. This meets the mean objective; p90 and p99 remain +17.340 and 31.689 ms. Ordinary decode and near-256K completion also pass the +measured cases described below. ## Fixed workload @@ -15,10 +18,13 @@ GSQ-RCO IQ3_S with Q8_0 MTP4, greedy draft sampling, TP4 and one request. Target sampling is temperature 0.6, top-p 0.95, top-k 20, seed 20261014; thinking and EOS remain enabled. -Context capacity remains 262144 **input tokens**, chunk size 8192, prefix -caching enabled, and FULL target decode graphs at M5. Expert storage retains +Context capacity remains 262144 **total input and output tokens**, chunk size +8192, prefix caching enabled, and FULL target decode graphs at M5. Expert storage retains the configured 80% GPU / 20% host placement. Target QSA history uses host -E4M3 with an 8192-token device hot cache; draft history uses host FP16. +E4M3; the reference has an 8192-token device hot cache and the configuration +candidate has 16384 tokens. Draft history uses host FP16. Device hot values +remain FP16 for both history types; the larger cache adds about 104 MiB per +rank including draft storage. PLE remains on disk with its native lane and 512-MiB CPU row cache. Measured prompts contain 574 and 32774 tokens, with 192 output tokens. These measurements do not qualify 256K speed or concurrent requests. @@ -214,3 +220,188 @@ independent committed draft metrics by at most one. Raw and aligned ledgers are both preserved. New event probes wait for all-rank counters to remain stable for 200 ms between requests. This does not change measured decode timestamps or ordinary production requests. + +## Ordinary whole-graph measurements and completeness + +CUDA graph-level tracing records graph envelopes without per-round CUDA +Events or fences. The verifier window runs from the earliest target graph +GPU start to the latest `_postprocess_mamba_align_kernel` end across TP4. +Source inspection establishes `_post_update_kernel`, then +`_scatter_num_accepted_kernel`, snapshot copy and Mamba alignment as target +sampling/state commit. Draft preparation follows this boundary. CPU graph +launch-to-draft time includes previously queued work and is not an isolated +verifier measurement. + +The matched ordinary HTTP pair retains identical prompts, seeds, sampling, +MTP, native artifacts and warmup sequence. Reference HOT8192 versus source +candidate HOT16384 measures 574-token mean TPOT at 8.957 / 7.714 ms, and +32774-token mean TPOT at 8.331 / 8.519 ms. The latter is about 2.3% slower; +there is no unconditional emitted-token speed claim. Tokens per committed +round differ, so output TPOT and verifier latency must remain separate. + +Commit-aligned four-request traces measure first/second 32K verifier means +of 16.967 / 16.169 ms for the reference and 15.816 / 17.788 ms for the +candidate. The candidate's 86.557-ms tail remains included. Every committed +round is retained; one additional asynchronous pipeline-end graph is +recorded separately per request. + +An eight-request steady HOT16384 capture has 584 complete state commits on +every rank: each request contains one prefill commit, 71 committed verifier +rounds and one pipeline-end replay. Five graph activity records are missing +across the ranks, affecting four of the 568 committed rounds. Direct target +graph ordinal alignment is invalid for this report and its old global +statistics are withdrawn. The three commit kernels, target graphs and +draft graphs all use the same CUDA context and stream on each rank. +Associating each graph with its following commit aligns complete request +and round ledgers. For a missing start, the preceding same-stream commit +end bounds the target start from below; this conservatively includes draft +and preparation time instead of imputing or dropping the missing window. +All remaining graph records map uniquely to a commit and all independent +committed-step counts reconcile. + +The resulting eight-request all-round verifier mean lies between 16.014 +and 16.043 ms; upper-bound p50/p90/p99/max are 15.265 / 18.984 / 24.134 / +31.897 ms. This remains above the 16-ms objective. Nsight's incomplete-event +warning and four conservatively bounded rounds are retained. A fresh single +capture range is required for a fully exact distribution. Interactive rearm +returned `StopRequest` errors and did not reliably honor a new output/range; +do not mix its report with the later HTTP requests. + +HOT20480 passes the five screened task/EOS cases but measures synchronized +32K verifier means of 17.864 / 16.590 ms, slower than HOT16384 and with less +memory headroom. W2/T16 at HOT16384 uses 202 registers without spills but +regresses shared, overflow and advancing-epoch churn selections; retain W4. +No corresponding runtime or default tuning change is promoted. + +The ordinary pair passes arithmetic, Python evaluation and both 32K needle +retrieval tasks. Its open-ended CPU explanation reaches the 2048-token cap +in both arms; a candidate 4096-token follow-up also reaches its cap. These +runs do not qualify EOS or long-output drift. The earlier HOT16384 screen's +1809-token natural EOS is retained separately. A 261959-input-token capacity +request reaches the 300-second client timeout without an OOM; it does not +establish successful near-256K completion. Keep both gates explicit before +promoting the deployment configuration. + +An empty-plan expert-ID remap screen tests counts 0/1/5/10/50/128 and plan +sizes 0/1/2/4/20/256 with W4 and W1, with and without an empty-metadata branch. +All 36 cases per policy retain exact table/output bytes in eager execution +and changed-input graphs. The measured M5 empty-plan kernel changes from +1.832 to 1.992 us with the W4 branch, so no runtime change is retained. + +The 564 exactly recorded steady rounds show mean maximum individual-rank +verifier duration of 15.656 ms, versus 16.017 ms for the complete TP4 window. +GPU target-start skew has mean/p90/p99 of 1.081 / 3.543 / 6.824 ms. CPU target +enqueue skew averages 8.579 ms; GPU queue-ahead averages 11.246 ms and includes +prior work. Enqueue skew correlates with GPU-start skew (0.607), motivating +an owned-process CPU affinity screen rather than treating queued time as +additional verifier computation. + +## Encoded hot-cache model screen + +The opt-in target setting `qsa_host_kv_hot_dtype="fp8_e4m3"` retains the +encoded bytes and original per-vector FP32 scales in device hot storage. +Both the existing native QSA family and its protected Triton reader decode +to the same FP16 operands used by the decoded cache. The authoritative host +history, page ownership, request masks, tentative writes, release ordering +and staging format remain unchanged. Default hot storage remains FP16. +Encoded hot storage requires E4M3 history; incompatible history is rejected. +An unavailable `run_hot_fp8` native artifact uses the qualified protected +reader, and unaligned native hot storage is explicitly rejected. + +The optional `qsa_host_kv_draft_hot_tokens` independently bounds the draft +cache; its default None preserves the existing shared capacity policy. +The measured prototype doubles target hot capacity from 16384 to 32768 +while retaining the 16384-token FP16 draft cache. Target storage changes +from 1024 to 520 bytes per cached token. Across twelve target QSA layers, +this configuration adds about 3 MiB per rank including scales, while leaving +the expert ratio, host histories, disk PLE and total context capacity intact. + +GPU qualification passes 151 cases: the original 130 host/native tests, +18 new native/protected-reader cases over M1/M5/M20/M21/M32, selection +widths 1/15/16/17/65/2051/4096, hot capacities 48/128/32768, request masks, +causal bounds and changed-input graphs, plus unavailable-native fallback +and two alignment-rejection cases. Live hot bytes and scales match the +host source exactly. Native attention uses its existing FP16 comparison +gate, rtol 0.003 / atol 0.0002, unchanged before the candidate runs. +The first prototype's eight fallback failures are retained: its generic +reader still interpreted encoded bytes as FP16, and three fixtures used +invalid staging capacities. The corrected reader and fixtures pass. + +The native module is built with the configured target's exact NVCC/link +arguments, writing only task-owned outputs and replacing the deployed +source module atomically. Shared venv packages are untouched. Original +native SHA256 is `ed1d3a854e673cec3c646076bbc7d3b3010d5840fa4c17867d1fc445a2e41eac`; +prototype SHA256 is `76fa88cc1fe323f019c0ad6942756dde84e5396bf46ae07b00cd27dc99115fd2`. +Rebuilding the final formatted source produces SHA256 +`bf0dd515b90a2132f95ae408ab4e23bec4e0e19b81b1199f2d74b9d644b7a54b`. +All 24 affected encoded-cache, unavailable-native, alignment-rejection and +legacy E4M3-code tests pass again against that final artifact. +Rebuilding the final formatted source gives SHA256 +`bf0dd515b90a2132f95ae408ab4e23bec4e0e19b81b1199f2d74b9d644b7a54b`. +Its 24 targeted GPU regressions pass, covering encoded/protected-reader +graphs, missing-native fallback, alignment rejection and all E4M3 codes; +the preceding 151-case suite remains the broader qualification evidence. + +The TP4 M5 screen includes resolve plus native attention at 32774 input +tokens, width 2051, advancing KV epochs and alternating paired measurements. +FP16 HOT16384 / encoded HOT32768 shared-selection medians are +51.456/52.992, 43.264/45.808, 48.896/53.760 and 47.104/49.920 us across +ranks. Sixteen changing selections measure 719.776/595.888, +731.136/592.207, 713.152/598.064 and 716.224/590.608 us. Churn improves +17–20%, with shared selections 3–10% slower. This passes the predeclared +operator screen of at least 5% churn improvement and at most 10% shared +regression. Model admission uses the separate measurements below; the +repository's default hot storage remains FP16. + +The encoded configuration's eight-request model capture retains all 640 +committed verification rounds. State commits are complete on TP4; 638 +windows are exact and two missing graph starts use conservative bounds. +The all-round mean lies between 15.482 and 15.494 ms, below the 16-ms mean +objective even at the conservative endpoint. Upper-bound p50/p90/p99/max +are 14.687 / 17.340 / 31.689 / 35.538 ms. Request means are 15.209, 14.780, +15.894, 16.804, 15.303, 15.072, 16.119 and 14.772 ms: this is an average +objective, not an every-round or p99 guarantee. The preceding FP16 HOT16384 +capture under the same affinity policy has mean bounds of 16.079–16.091 ms; +the runs have different warmup histories, so this is consecutive-screen +comparison rather than an unconditional throughput gain. + +Arithmetic, Python evaluation and both 32K retrieval tasks pass with natural +EOS. The paired concise 32K CPU task passes both seeds 20261014/20261015: +the encoded run ends naturally at 1088/682 output tokens versus 1333/1348 in +the preceding FP16-hot run. Both final answers correctly explain renaming, +RAW readiness, the ROB's ordered retirement and precise exceptions, with +concrete dependency examples. Reasoning and final output have no six-word +phrase repeated more than twice. The open-ended 2048-token CPU task still +hits its cap and remains an inconclusive long-answer stress observation. + +Eight ordinary streaming requests, with capture inactive and the same 32774 +input / 192 output tokens, measure mean TPOT 7.109 ms, or 140.7 output +tokens/s. Mean decode time is 1.358 s, excluding 1.309-s prefix-warm TTFT. +The HTTP round estimator is 18.859 ms with 2.681 emitted tokens per round; +it includes draft and scheduling and is distinct from GPU verification. +All requests emit 192 tokens. This is a candidate observation, not a paired +throughput gain over the earlier runs, whose acceptance counts differ. + +The same capacity probe, with its timeout raised from 300 to 900 seconds +and unchanged 128-output limit, completes with 261959 input and 71 output +tokens (262030 total), natural EOS and the correct `SAFFRON45793` marker. +TTFT is 320.034 s and total time is 320.435 s. This establishes completion +near the 262144-total-token boundary for this request, not 256K speed or +concurrency. No OOM or memguard termination occurs. GPU memory reaches +16143 MiB per rank afterward, leaving roughly 17 MiB free: capacity passes +with very little device headroom. The earlier 300-second timeout remains +retained as an incomplete observation rather than a successful capacity run. + +The final ordinary server runs without Nsight injection or profiler CLI +configuration, with target encoded HOT32768 and draft FP16 HOT16384. After +three warmup requests, task-owned worker main and helper threads are placed +on separate physical CPU cores using the recorded affinity policy. One +32774-input / 192-output smoke measures TPOT 7.836 ms, or 127.6 tokens/s, +and estimated whole-round cost 18.252 ms. Its acceptance is 0.335, versus +the earlier eight-request screen's different acceptance history. This is +a single deployment smoke, not a new mean-verifier distribution. +The key-gated public proxy returns 401 without authentication and 200 for +the authenticated model list and generation. Arithmetic generation returns +1517 with natural EOS. Memory protection remains active. The standard +deployment launcher retains its GPU locks and now records the selected +target/draft cache settings; its original script is backed up outside Git. diff --git a/tests/models/qwen4_exp/test_device_kv_attention.py b/tests/models/qwen4_exp/test_device_kv_attention.py index 7b624dedb..371e205ef 100644 --- a/tests/models/qwen4_exp/test_device_kv_attention.py +++ b/tests/models/qwen4_exp/test_device_kv_attention.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib +from types import SimpleNamespace + import pytest import torch @@ -259,3 +262,184 @@ def test_fp32_attention_oracle_aliases_masks_writes_and_graph(rows, dtype): graph.replay() torch.accelerator.synchronize() check() + + +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize( + "rows,width,capacity", + [ + (1, 1, 48), + (5, 15, 48), + (5, 16, 48), + (5, 17, 48), + (5, 65, 128), + (20, 2051, 128), + (21, 2051, 128), + (32, 4096, 128), + (5, 2051, 32768), + ], +) +def test_encoded_hot_storage_operands_and_changed_graph(rows, width, capacity, native): + """Raw storage is exact; native attention keeps its existing FP16 gate.""" + torch.manual_seed(73) + device = torch.device("cuda:0") + states = [ + HostQSAKV( + 2, + 816, + 256, + device, + width=width, + rows=rows, + hot_tokens=capacity, + direct_device=native, + hot_fp8=encoded, + ) + for encoded in (False, True) + ] + key = torch.randn(1632, 1, 256, device=device, dtype=torch.float16) + value = torch.randn_like(key) + slots = torch.arange(1632, device=device) + for state in states: + state.write(key, value, slots) + indices = torch.arange(width, device=device, dtype=torch.int32).repeat(rows, 1) + if width > 1: + indices[:, -1] = -1 + table = torch.tensor([[0, 1], [0, -1]], device=device, dtype=torch.int32) + requests = torch.arange(rows, device=device, dtype=torch.int32) % 2 + if rows > 1: + requests[-1] = -1 + positions = torch.full((rows,), 1540, device=device, dtype=torch.int64) + lengths = torch.tensor([1600, 1500], device=device, dtype=torch.int32) + query = torch.randn(rows, 6, 256, device=device, dtype=torch.float16) + gate = torch.randn_like(query) + outputs = [torch.empty_like(query) for _ in states] + + def run(): + for state, output in zip(states, outputs): + state.write(key[:2], value[:2], slots[:2]) + host_qsa_attention( + query, state, indices, table, requests, positions, lengths, output, gate + ) + + def check(): + assert all(torch.isfinite(output).all() for output in outputs) + torch.testing.assert_close(outputs[1], outputs[0], rtol=0.003, atol=0.0002) + assert torch.equal(states[0].host, states[1].host) + assert torch.equal(states[0].host_scales, states[1].host_scales) + for state in states: + pages = state.tags.flatten().cpu() + live = pages >= 0 + if not live.any(): + continue + slots_live = torch.arange(pages.numel())[live] + offsets = torch.arange(4) + tokens = (pages[live, None] * 4 + offsets).flatten().long() + cached_tokens = (slots_live[:, None] * 4 + offsets).flatten().long() + codes = state.host[tokens // 816, :, tokens % 816, 0] + scales = state.host_scales[tokens] + cache = state.hot_values[cached_tokens.to(device)].cpu() + if state.hot_fp8: + assert torch.equal(cache, codes) + cached_scales = state.hot_scales[cached_tokens.to(device)].cpu() + assert torch.equal(cached_scales, scales) + else: + decoded = ( + codes.view(torch.float8_e4m3fn).float() * scales[..., None] + ).half() + assert torch.equal(cache, decoded) + assert not (state.page_slots == -2).any() + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + value[:2].mul_(-0.125) + query.mul_(0.25) + positions.fill_(63) + for _ in range(3): + graph.replay() + torch.accelerator.synchronize() + check() + + +def test_encoded_hot_fallback_without_new_native_op(monkeypatch): + importlib.import_module("vllm._sm70_qsa_device_C") + original = torch.ops.vllm_sm70_qsa_device + monkeypatch.setattr( + torch.ops, + "vllm_sm70_qsa_device", + SimpleNamespace(run_hot=original.run_hot), + ) + state = HostQSAKV( + 1, + 64, + 256, + torch.device("cuda:0"), + hot_tokens=64, + rows=1, + width=64, + hot_fp8=True, + direct_device=True, + ) + assert state.device_history_workspace is None + assert state.device_history_reason == "native_encoded_hot_cache_unavailable" + key = torch.zeros(64, 1, 256, device="cuda", dtype=torch.float16) + state.write(key, torch.full_like(key, 3.25), torch.arange(64, device="cuda")) + query = torch.zeros(1, 6, 256, device="cuda", dtype=torch.float16) + output = torch.empty_like(query) + host_qsa_attention( + query, + state, + torch.arange(64, device="cuda", dtype=torch.int32)[None], + torch.zeros(1, 1, device="cuda", dtype=torch.int32), + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([63], device="cuda", dtype=torch.int64), + torch.tensor([64], device="cuda", dtype=torch.int32), + output, + ) + torch.testing.assert_close(output, torch.full_like(output, 3.25), rtol=0, atol=0) + + +@pytest.mark.parametrize("encoded", [False, True]) +def test_native_hot_rejects_misaligned_storage(encoded): + importlib.import_module("vllm._sm70_qsa_device_C") + state = HostQSAKV( + 1, + 64, + 256, + torch.device("cuda:0"), + width=16, + rows=1, + hot_tokens=16, + hot_fp8=encoded, + direct_device=True, + ) + query = torch.zeros(1, 6, 256, device="cuda", dtype=torch.float16) + bad = torch.empty(16 * 512 + 1, device="cuda", dtype=state.hot_values.dtype)[1:] + bad = bad.view(16, 2, 256) + args = ( + query, + state.history, + state.scales, + torch.arange(16, device="cuda", dtype=torch.int32)[None], + torch.zeros(1, 1, device="cuda", dtype=torch.int32), + torch.zeros(1, device="cuda", dtype=torch.int32), + torch.tensor([15], device="cuda", dtype=torch.int64), + torch.tensor([16], device="cuda", dtype=torch.int32), + torch.empty_like(query), + None, + *state.device_history_workspace, + bad, + torch.zeros(1, 16, device="cuda", dtype=torch.int32), + ) + op = ( + torch.ops.vllm_sm70_qsa_device.run_hot_fp8 + if encoded + else torch.ops.vllm_sm70_qsa_device.run_hot + ) + if encoded: + args += (state.hot_scales,) + with pytest.raises(RuntimeError, match="aligned vector loads"): + op(*args) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index f83e58c96..bc6ce2b4d 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -884,6 +884,11 @@ class KernelConfig: device and the remainder directly from mapped host history.""" qsa_host_kv_hot_tokens: int = Field(default=8192, gt=0, multiple_of=16) """Per-layer device hot-page capacity; collisions use exact host gathers.""" + qsa_host_kv_hot_dtype: Literal["float16", "fp8_e4m3"] = "float16" + """Target hot storage; encoded E4M3 requires E4M3 authoritative history. + Encoded hot reads round to the same FP16 operands as the decoded cache.""" + qsa_host_kv_draft_hot_tokens: int | None = Field(default=None, gt=0, multiple_of=16) + """Optional independent draft hot capacity; None preserves target capacity.""" qsa_host_kv_staging_rows: int = Field(default=32, gt=0, le=32) """Shared host-miss staging capacity; larger batches use bounded row chunks.""" qsa_host_kv_prefill_stage_tokens: int = Field(default=0, ge=0) diff --git a/vllm/models/qwen4_exp/nvidia/ops/device_kv_attention.py b/vllm/models/qwen4_exp/nvidia/ops/device_kv_attention.py index f7d578d23..e34355570 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/device_kv_attention.py +++ b/vllm/models/qwen4_exp/nvidia/ops/device_kv_attention.py @@ -44,6 +44,12 @@ def initialize_device_history_attention(state, enabled: bool | None) -> None: torch.ops.vllm_sm70_qsa_device, "run_hot" ): reason = "native_hot_cache_unavailable" + elif ( + not state.device_reference + and getattr(state, "hot_fp8", False) + and not hasattr(torch.ops.vllm_sm70_qsa_device, "run_hot_fp8") + ): + reason = "native_encoded_hot_cache_unavailable" state.device_history_reason = reason # Host-authoritative history: resolve selected pages into the protected # hot cache, then read hot tokens on device and the rest from mapped host. @@ -153,7 +159,13 @@ def device_history_attention( gate = gate.view_as(query) if gate is not None else None if hot: resolved = state.resolve_hot(indices, table, requests, positions, lengths) - torch.ops.vllm_sm70_qsa_device.run_hot( + op = ( + torch.ops.vllm_sm70_qsa_device.run_hot_fp8 + if getattr(state, "hot_fp8", False) + else torch.ops.vllm_sm70_qsa_device.run_hot + ) + extra = (state.hot_scales,) if getattr(state, "hot_fp8", False) else () + op( query, state.history, state.scales, @@ -167,11 +179,13 @@ def device_history_attention( *workspace, state.hot_values, resolved, + *extra, ) state.device_history_reason = None logger.info_once( "Using SM70 native QSA attention over host history with a device " - "hot cache (FP32 PV)" + "hot cache (%s storage, FP32 PV)", + "E4M3" if getattr(state, "hot_fp8", False) else "FP16", ) return True torch.ops.vllm_sm70_qsa_device.run( diff --git a/vllm/models/qwen4_exp/nvidia/ops/host_kv.py b/vllm/models/qwen4_exp/nvidia/ops/host_kv.py index 9059e19de..6239affba 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/host_kv.py +++ b/vllm/models/qwen4_exp/nvidia/ops/host_kv.py @@ -72,6 +72,8 @@ def _write( Sets: tl.constexpr, FP8: tl.constexpr, HOT: tl.constexpr = True, + HOT_FP8: tl.constexpr = False, + HotScales=None, ): # HOT=False: a plain history copy (a prefill stage buffer): no epoch, # no hot-cache update. @@ -109,6 +111,13 @@ def _write( decoded = (fp8_e4m3fn_bits_to_fp32_bitcast(codes) * scales[:, None]).to( tl.float16 ) + if HOT_FP8: + decoded = codes + tl.store( + HotScales + tl.maximum(resident, 0) * 8 + slot % 4 * 2 + kv, + scales, + (resident >= 0) & (tag == page), + ) tl.store( Hot + ((tl.maximum(resident, 0) * 4 + slot % 4) * 2 + kv[:, None]) * Dim @@ -215,6 +224,8 @@ def _gather( Sets: tl.constexpr, FP8: tl.constexpr, STAGE: tl.constexpr, + HOT_FP8: tl.constexpr = False, + HotScales=None, ): row, tile = tl.program_id(0), tl.program_id(1) lanes = tl.arange(0, 16) @@ -257,6 +268,15 @@ def _gather( initial_hit[:, None, None], other=0, ) + if HOT_FP8: + scales = tl.load( + HotScales + tl.maximum(initial, 0)[:, None] * 2 + kv[None, :], + initial_hit[:, None], + other=0, + ) + values = ( + fp8_e4m3fn_bits_to_fp32_bitcast(values) * scales[:, :, None] + ).to(tl.float16) tl.store( Out + ((row * 2 + kv[None, :, None]) * Padded + columns[:, None, None]) @@ -360,6 +380,15 @@ def _gather( need_cached[:, None, None], other=0, ) + if HOT_FP8: + cached_scales = tl.load( + HotScales + hot_tokens[:, None] * 2 + kv[None, :], + need_cached[:, None], + other=0, + ) + cached = ( + fp8_e4m3fn_bits_to_fp32_bitcast(cached) * cached_scales[:, :, None] + ).to(tl.float16) codes = tl.load( History + ( @@ -391,9 +420,15 @@ def _gather( Hot + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim + dims[None, None, :], - values, + codes if HOT_FP8 else values, token_installs[:, None, None], ) + if HOT_FP8: + tl.store( + HotScales + hot_tokens[:, None] * 2 + kv[None, :], + scales, + token_installs[:, None], + ) if STAGE > 0: store_values = columns < Padded if STAGE == 2: @@ -448,6 +483,7 @@ class HostQSAKV: direct_device: bool | None = None, is_speculative_draft: bool = False, resolve_rows: int = 32, + hot_fp8: bool = False, ): if blocks <= 0 or page_size <= 0 or page_size % 4 or dim != 256: raise ValueError("Host QSA KV requires positive page4 geometry and D256") @@ -464,6 +500,9 @@ class HostQSAKV: if dtype not in (torch.uint8, torch.float16): raise ValueError("Host history requires E4M3 bytes or FP16 values") self.fp8 = dtype == torch.uint8 + if hot_fp8 and not self.fp8: + raise ValueError("Encoded hot cache requires E4M3 authoritative history") + self.hot_fp8 = hot_fp8 self.device_reference = device_reference self.is_speculative_draft = is_speculative_draft self.host = None @@ -498,7 +537,14 @@ class HostQSAKV: ) self.scales = get_accelerator_view_from_cpu_tensor(self.host_scales) self.hot_values = torch.empty( - (hot_tokens, 2, dim), dtype=torch.float16, device=device + (hot_tokens, 2, dim), + dtype=torch.uint8 if hot_fp8 else torch.float16, + device=device, + ) + self.hot_scales = torch.empty( + (hot_tokens, 2) if hot_fp8 else (0,), + dtype=torch.float32, + device=device, ) self.tags = torch.full((self.sets, 4), -1, dtype=torch.int32, device=device) self.stamps = torch.zeros_like(self.tags) @@ -569,6 +615,8 @@ class HostQSAKV: self.dim, self.sets, self.fp8, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) @@ -611,6 +659,8 @@ class HostQSAKV: self.sets, self.fp8, HOT=False, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) @@ -859,6 +909,8 @@ class HostQSAKV: self.sets, self.fp8, stage, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) diff --git a/vllm/models/qwen4_exp/nvidia/ops/host_kv_attention.py b/vllm/models/qwen4_exp/nvidia/ops/host_kv_attention.py index 3195f40da..e19d265b2 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/host_kv_attention.py +++ b/vllm/models/qwen4_exp/nvidia/ops/host_kv_attention.py @@ -102,6 +102,8 @@ def host_qsa_attention( HOST_INDICES=resolved, HOST_VALID_COUNTS=state.lengths, HOST_CACHE=True, + HOT_SCALES=state.hot_scales, + HOST_CACHE_E4M3=state.hot_fp8, num_warps=warps, num_stages=2, ) diff --git a/vllm/models/qwen4_exp/nvidia/ops/host_kv_reader.py b/vllm/models/qwen4_exp/nvidia/ops/host_kv_reader.py index 94423e914..86ca08619 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/host_kv_reader.py +++ b/vllm/models/qwen4_exp/nvidia/ops/host_kv_reader.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Read protected FP16 hot values and separately decoded misses.""" +"""Read protected hot values as FP16 and separately decoded misses.""" +from vllm.models.deepseek_v4.common.ops.fp8_software import ( + fp8_e4m3fn_bits_to_fp32_bitcast, +) from vllm.triton_utils import tl, triton @@ -16,6 +19,8 @@ def load_host_kv( dims, PADDED: tl.constexpr, DIM: tl.constexpr, + HotScales=None, + HOT_FP8: tl.constexpr = False, ): hits = valid & (hot_tokens >= 0) safe_hot = tl.maximum(hot_tokens, 0).to(tl.int64) @@ -29,6 +34,15 @@ def load_host_kv( hits[:, None], other=0, ) + if HOT_FP8: + k_scales = tl.load(HotScales + safe_hot * 2, hits, other=0) + v_scales = tl.load(HotScales + safe_hot * 2 + 1, hits, other=0) + cached_k = (fp8_e4m3fn_bits_to_fp32_bitcast(cached_k) * k_scales[None, :]).to( + tl.float16 + ) + cached_v = (fp8_e4m3fn_bits_to_fp32_bitcast(cached_v) * v_scales[:, None]).to( + tl.float16 + ) misses = valid & ~hits source_k = tl.load( Fallback + (row * 2 * PADDED + columns[None, :]) * DIM + dims[:, None], diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 9d249c68a..d088008ef 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -641,6 +641,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( HISTORY_LENGTHS=None, DEVICE_HISTORY: tl.constexpr = False, HISTORY_E4M3: tl.constexpr = False, + HOT_SCALES=None, + HOST_CACHE_E4M3: tl.constexpr = False, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -739,6 +741,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( dim_offsets, tl.cdiv(TOPK, 4) * 4, HEAD_DIM, + HotScales=HOT_SCALES, + HOT_FP8=HOST_CACHE_E4M3, ) else: keys = tl.load( diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index ce206b5f1..c2b521d69 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -729,6 +729,15 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): # Cache binding runs under the target worker's config, also for draft # layers. Retain each owner's policy from model construction. self.host_kv_is_draft = vllm_config.is_speculative_draft + if self.host_kv_is_draft: + self.host_kv_hot_tokens = ( + vllm_config.kernel_config.qsa_host_kv_draft_hot_tokens + or self.host_kv_hot_tokens + ) + self.host_kv_hot_fp8 = ( + not self.host_kv_is_draft + and vllm_config.kernel_config.qsa_host_kv_hot_dtype == "fp8_e4m3" + ) self.host_kv_direct_device = vllm_config.kernel_config.sm70_qsa_device_history self.host_kv_dtype = ( vllm_config.kernel_config.qsa_host_kv_draft_dtype @@ -912,13 +921,14 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): device_reference=self.host_kv_device_reference, direct_device=self.host_kv_direct_device, is_speculative_draft=self.host_kv_is_draft, + hot_fp8=self.host_kv_hot_fp8, ) if not self.host_kv_is_draft: _TARGET_HOST_KV.append(self.host_kv) logger.info_once( "QSA encoded history initialized: storage=%s, dtype=%s, hot_tokens=%d; " "direct device M1..20/H6/D256=%s, fallback_reason=%s; " - "other shapes use protected FP16 hot pages and miss staging.", + "other shapes use protected hot pages and FP16 miss staging.", "device_reference" if self.host_kv_device_reference else "host", self.host_kv_dtype, self.host_kv_hot_tokens,