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..b140fb8a9 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py @@ -0,0 +1,285 @@ +# 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, 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, parallel_shards = POLICIES[policy] + plan, arena, mapped, device_offsets, host_offsets, words = buffers + grid = (programs, split, 3) if parallel_shards else (programs, split) + _apply_plan[grid]( + plan, + arena, + mapped, + device_offsets, + host_offsets, + words, + 0, + OPS_PROGRAMS=programs, + SPLIT=split, + BLOCK=block, + PARALLEL_SHARDS=parallel_shards, + 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") + 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") + 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/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 new file mode 100644 index 000000000..684c98eae --- /dev/null +++ b/docs/design/flashnext_decode_verifier_16gb_20261014.md @@ -0,0 +1,407 @@ +# 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. 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 + +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 **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; 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. + +## 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.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. + +## 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 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. 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, 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 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 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.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 +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. + +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 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 +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. + +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 +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. + +## 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/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_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/tests/models/qwen4_exp/test_host_kv.py b/tests/models/qwen4_exp/test_host_kv.py index 09a514322..e1d913b6e 100644 --- a/tests/models/qwen4_exp/test_host_kv.py +++ b/tests/models/qwen4_exp/test_host_kv.py @@ -9,6 +9,119 @@ 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]) +@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") @@ -82,8 +195,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 +353,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 +368,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 +391,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 +442,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/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/model_executor/layers/quantization/gguf_expert_cache.py b/vllm/model_executor/layers/quantization/gguf_expert_cache.py index b274f9bdf..6f112ac9d 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) +# 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 @@ -109,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). @@ -121,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) @@ -582,8 +586,9 @@ class ExpertCache: i, OPS_PROGRAMS=_COPY_PROGRAMS[0], SPLIT=_COPY_PROGRAMS[1], - BLOCK=1024, - num_warps=8, + BLOCK=256, + PARALLEL_SHARDS=True, + num_warps=4, ) out = torch.empty_like(ids) _remap[(1,)]( 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 c01c47222..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]) @@ -343,46 +363,85 @@ 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,)) - cached = tl.load( - Hot - + (hot_tokens[:, None, None] * 2 + kv[None, :, None]) * Dim - + dims[None, None, :], - token_hits[:, 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, :], - (valid & ~token_hits)[:, None, None], - other=0, - ) - if FP8: - scales = tl.load( - Scales + tokens[:, None] * 2 + kv[None, :], - (valid & ~token_hits)[:, None], + 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. + # 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 + 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 + + ( + ( + 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) - token_installs = tl.reshape(tl.broadcast_to(install[:, None], (4, 4)), (16,)) - 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, :], + 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: + 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( @@ -395,17 +454,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) @@ -435,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") @@ -451,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 @@ -485,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) @@ -556,6 +615,8 @@ class HostQSAKV: self.dim, self.sets, self.fp8, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) @@ -598,6 +659,8 @@ class HostQSAKV: self.sets, self.fp8, HOT=False, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) @@ -846,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, 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, )