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 01/19] [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 0000000000..7fe5eb7893 --- /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 0000000000..436b9bfebc --- /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 09a514322e..3f97cfa2e3 100644 --- a/tests/models/qwen4_exp/test_host_kv.py +++ b/tests/models/qwen4_exp/test_host_kv.py @@ -9,6 +9,61 @@ 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 check(gather=True): 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 run(): 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 6206462509..6f8bdd910e 100644 --- a/tests/v1/worker/test_step_profiler.py +++ b/tests/v1/worker/test_step_profiler.py @@ -106,6 +106,39 @@ def forbidden(**kwargs): 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 b274f9bdf5..99ff21faf2 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 @@ _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 @@ def resolve(self, layer: int, ids: torch.Tensor) -> torch.Tensor: 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 c01c47222d..03d2957eb0 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 f4abdb330d..bc92fd03f2 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 e150e9da8f4d8817a52d18ee34ec7883576f9570 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 03:43:21 +0800 Subject: [PATCH 02/19] [Kernel] Reduce Flash-Next prefill compute and cold expert traffic Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../OPTIMIZATION_20261014.md | 442 ++++++++++++++++++ benchmarks/flashnext_prefill/cold_dma.py | 92 ++++ benchmarks/flashnext_prefill/collectives.py | 79 ++++ benchmarks/flashnext_prefill/dense_blas.py | 116 +++++ benchmarks/flashnext_prefill/hc_pipeline.py | 138 ++++++ benchmarks/flashnext_prefill/hc_projection.py | 59 +++ .../flashnext_prefill/measure_prefill.py | 173 +++++++ benchmarks/flashnext_prefill/moe_dma.py | 122 +++++ benchmarks/flashnext_prefill/moe_streams.py | 189 ++++++++ benchmarks/flashnext_prefill/moe_tiles.py | 346 ++++++++++++++ .../flashnext_prefill/nsys_breakdown.py | 30 +- benchmarks/flashnext_prefill/quality.py | 125 +++++ csrc/gguf_expert_cache.cpp | 39 ++ .../kernels/gemm/kernel/sm70_884_4.cu | 54 +++ .../quantization/test_gguf_grouped_prefill.py | 152 ++++++ .../test_gguf_turbomind_dequant.py | 60 ++- .../qwen4_exp/test_gguf_expert_cache.py | 40 ++ .../linear/mixed_precision/sm70_gguf.py | 6 + .../layers/quantization/gguf_expert_cache.py | 11 +- 19 files changed, 2242 insertions(+), 31 deletions(-) create mode 100644 benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md create mode 100644 benchmarks/flashnext_prefill/cold_dma.py create mode 100644 benchmarks/flashnext_prefill/collectives.py create mode 100644 benchmarks/flashnext_prefill/dense_blas.py create mode 100644 benchmarks/flashnext_prefill/hc_pipeline.py create mode 100644 benchmarks/flashnext_prefill/hc_projection.py create mode 100644 benchmarks/flashnext_prefill/measure_prefill.py create mode 100644 benchmarks/flashnext_prefill/moe_dma.py create mode 100644 benchmarks/flashnext_prefill/moe_streams.py create mode 100644 benchmarks/flashnext_prefill/moe_tiles.py create mode 100644 benchmarks/flashnext_prefill/quality.py create mode 100644 tests/kernels/quantization/test_gguf_grouped_prefill.py diff --git a/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md b/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md new file mode 100644 index 0000000000..bc189ce3ae --- /dev/null +++ b/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md @@ -0,0 +1,442 @@ +# Flash-Next prefill: 6K investigation + +Integration baseline: `onecat/main`, `6f3557c973b4602905037e086fd4acba9b0d867f` +(PR #1203). Work is isolated on +`codex/v100-flashnext-prefill-6k-9683-20261013-160548`. + +## Workload and evidence boundaries + +- Four V100-SXM2 16 GB GPUs, NV2 links, CUDA 12.8, driver 570.211.01, + Torch 2.10.0+cu128, Python 3.12.3. One 18-core/36-thread E5-2686 v4 socket. +- Flash-Next IQ3_S, TP4, MTP4, maximum context **262144 tokens**, + maximum scheduled chunk **8192 tokens**, prefix caching and thinking enabled. + Experts retain 80% GPU / 20% host residency; QSA KV stays on host, PLE on disk. +- Production sampling remains temperature 0.6 / top-p 0.95 / top-k 20. + The historical `ttft.py` is an explicitly greedy, one-output-token timing test. +- Historical `ttft.py 32000` uses random words, approximately 34700 tokens, + **five** chunks (7344 x 4 plus approximately 5347). Its 7.50–7.55 s figure + measures a complete HTTP response, not streaming TTFT or pure prefill. +- Historical `profile_request.py` recreates its RNG inside the word generator, + producing one repeated word and **four** chunks. Its 6.48 s trace cannot close + the wall-clock accounting for the different 7.5 s request. The old table also + omitted 0.91 s of other kernels when only its first five rows were summed. +- For approximately 34700 tokens, 6000 tokens/s corresponds to approximately + **5.8 s**: approximately 1.7 s / 23% below the historical response latency. + No new end-to-end speed claim is established by the microbenchmarks below. + +## Resource findings from the existing DMA trace + +The critical GPU's main-stream kernel/copy union occupies 87.35% of the +6.482 s trace interval. Across all streams, kernel union occupies 88.98% and +kernel/copy union 97.56%. Main-stream gaps are not whole-device idle time. +Cold-expert copies overlap kernels for 85–87% of their duration; their summed +3.45–3.70 s per GPU must not be added to end-to-end latency. + +Matching collective ordinals across all four ranks gives 400 all-reduces +and 812 all-gathers per rank. GPU0 arrives last at 278/400 all-reduces. +The sum from the latest rank arrival to the latest completion is 400.8 ms +for all-reduce, compared with GPU0's 419.5 ms kernel sum. The earliest-to-latest +arrival spreads sum to 258.6 ms; those waits mostly occur on earlier ranks. +They cannot simply be subtracted from the critical rank's latency. These are +trace timing estimates, not link bandwidth counters or a new workload result. + +The original analysis joined CUDA correlation IDs across TP processes. +`nsys_breakdown.py` now keys them by process and correlation ID, and keys the +main stream by process and stream ID. Replaying the existing DMA trace moves +20 ms from the host-gap bucket to GPU waits: **133 ms host gaps, 687 ms GPU +waits**, with the same 6482 ms span and kernel totals. These are gap +attributions, not hardware utilization counters. + +The trace has launch geometry and registers/shared-memory metadata, but no SM, +Tensor Core, DRAM, or NVLink hardware utilization counters. Static occupancy +limits are not measured resource saturation. For example, QSA rows use +120 registers/thread and approximately 34 KiB shared memory per block, allowing +two 256-thread blocks/SM. Increasing this to three was already rejected in the +previous investigation. + +The saved real QSA input has 8160 queries, six heads, dimension 256, and 2051 +index positions per query. Only 12240 of 16736160 indices are negative. +Reading each valid K/V pair once would request approximately 8.56 GB; dividing +by kernel time is **requested payload bandwidth**, not measured HBM traffic +(cache hits are not known). + +## Experiments and decision gates + +### 1. Narrow grouped MoE tiles + +Hypothesis: expert batches contain roughly 100–160 rows at the current chunks. +Additional 64x64 tiles may reduce masked work, accumulator registers, and N=320 +padding. Reuse the existing TurboMind grouped provider, codec transforms, +autotuner, workspace, and group-sized K iteration. Existing kernels remain +available as fallback. No graph, cache-state, or sampling behavior is changed. + +The experimental tiles use the existing narrow-prefill admission +(`M >= 512` routed rows, `N <= 512`, `N % 64 == 0`). Existing provider validation +still applies. Measured cases are much narrower: TP4 rank-zero real weights, +512 experts, top-10, hidden 2560, intermediate 160, N=320 gate/up; down executes +the production 4096-token sub-batches and unroute operation. This does not +qualify every admitted shape or other models. + +`moe_tiles.py` tests all gate/up types present in this model: +18 (17 layers), 22 (20 layers), 21 (10 layers), 23 (one layer). Down uses +20 (39 layers) or 42 (nine layers). Timing excludes routing, transcode, +host DMA, and multi-GPU communication. It includes both GEMMs, SiLU, down +sub-batching, allocation calls, and output unroute. + +Predeclared microbenchmark numerical gates: finite hidden/output tensors; +maximum checked-row relative L2 <= 0.01 against independently dequantized +FP32 gate/up and complete MoE references; candidate/control checked-row +relative L2 <= 0.003. These tolerate FP16 accumulation/rounding, but do not +replace paired model task-quality, repetition/EOS, or long-context checks. +Require a stable chain improvement before spending a service restart on it. + +Baseline, median of five CUDA-event samples, five iterations/sample: + +| Layer / gate type | Tokens | Gate/up + SiLU ms | Down + unroute ms | Complete chain ms | +|---|---:|---:|---:|---:| +| 0 / 18 | 7344 | 4.327 | 3.156 | 7.367 | +| 0 / 18 | 5347 | 3.759 | 2.294 | 5.628 | +| 1 / 22 | 7344 | 4.042 | 2.908 | 6.952 | +| 1 / 22 | 5347 | 2.482 | 2.061 | 4.506 | +| 17 / 21 | 7344 | 4.458 | 3.150 | 7.465 | +| 17 / 21 | 5347 | 3.373 | 2.280 | 5.700 | +| 47 / 23 | 7344 | 3.875 | 3.200 | 6.833 | +| 47 / 23 | 5347 | 2.924 | 2.323 | 5.175 | + +All eight baseline cases pass; maximum checked output error is 0.000844. +Native dispatch confirms type-18 currently selects 32x128, so a benefit cannot +be inferred merely from reducing the M tile. Separate-stage times should not +be summed as a replacement for the measured chain (GPU clocks vary). + +Rejected before GPU testing: 64x64 with four warps cannot instantiate the +current metadata thread map (zero-width access for this layout). Retain the +two-warp candidate for measurement; do not change shared iterator contracts +just to admit the extra tuning candidate. + +A second artifact adds padded activation supply using the existing quantized +batch-family operand. Hypothesis: at K32/K64, the unpadded shared-memory row +stride repeats banks across the distinct rows of each eight-lane MMA load +group. Padding by eight half values separates those accesses. This is a +static layout hypothesis, not hardware-counter evidence. Candidates are +32x128x32 with 32x64 epilogue subtiles and 64x64x64 with four warps. The K64 +metadata tile has enough elements for four warps. Decoder, group size, +routing, bank storage, and fallback kernels remain shared with the existing +family. Separate artifacts retain the first candidate for comparison. + +Measured complete MoE chains (ms): + +| Layer / tokens | Control | 64x64 | Padded candidates | Padded reduction | +|---|---:|---:|---:|---:| +| 0 / 7344 | 7.367 | 7.047 | 6.995 | 5.1% | +| 0 / 5347 | 5.628 | 5.124 | 4.956 | 11.9% | +| 1 / 7344 | 6.952 | 6.868 | 6.318 | 9.1% | +| 1 / 5347 | 4.506 | 4.537 | 4.404 | 2.3% | +| 17 / 7344 | 7.465 | 7.224 | 6.711 | 10.1% | +| 17 / 5347 | 5.700 | 5.172 | 5.006 | 12.2% | +| 47 / 7344 | 6.833 | 6.736 | 6.329 | 7.4% | +| 47 / 5347 | 5.175 | 4.808 | 4.507 | 12.9% | + +All eight cases pass for both candidates. Checked hidden/output rows are +bit-exact to control. Weighting these isolated chain deltas by 48 layer types +and five chunks gives only about 0.13 s for the padded artifact; this is a +screening estimate, not an end-to-end result. No claim of 6K follows. +The added grouped-prefill regression covers M=511/512/513, interior M=2049, +N=64/320/512, N32 and wide fallbacks, K=256/512/768/2560, uneven/empty experts, +indexed and direct input, fused SiLU, output guards, and changed-input graph +replay. Its predeclared maximum-row relative-L2 limit is 0.005 versus an +independent GGUF-dequantized FP32 reference. + +**The second padded artifact is rejected despite those eight passes.** The +boundary regression finds non-finite/wrong results in five of 24 cases. +Its 32x128 CTA used a 32x64 epilogue subtile to distinguish the tuning key, +but that subtile is narrower than the output thread map's horizontal span. +Some threads therefore read outside the intended shared-memory tile. The +third candidate removes it and uses a full 32x64x64 CTA/epilogue with a 2x2 +warp map, alongside the full 64x64x64 candidate. No shared epilogue or +iterator contract is changed. A 1x4 warp map was rejected at compilation: +its N=16/warp is below this MMA family's N=32 atom. + +The verified original artifact passes **24/24**. Pytest initially prepended +the candidate source root through the tests package, so a nominal control +invocation actually loaded the padded artifact; that comparison is invalid +and retained only as a harness failure. The corrected runner imports vLLM +and `_C` before invoking pytest, asserts their paths against the requested +artifact, and prints the native SHA256. It confirms 24 original passes and +five failures for the second padded artifact. Microbenchmark JSON now also +records loaded package/native paths and native hash. + +The corrected third artifact passes **24/24 boundary cases**. Its full-chain +times for the eight rows above are 6.723, 4.753, 6.201, 4.437, 6.672, 4.824, +6.404, and 4.592 ms: reductions of 1.5–15.5%, with all checked hidden/output +rows bit-exact to control. The layer/chunk-weighted isolated delta is 0.163 s. +This artifact, rather than the rejected second one, proceeds to model testing. +Its `_C` SHA256 is +`5d389732b5fb828b947ffdb74a3ec6c80aa70cb60bdc6b1ca65e437a9514148e`. + +Two smaller hypotheses were rejected on the original artifact: + +- Increasing down sub-batches from 4096 to 8192 tokens: layers 0/1 at + 7344/5347 tokens give complete-chain times of 7.590/5.921 and 7.269/4.878 ms, + versus 7.367/5.628 and 6.952/4.506 ms. All checked outputs are bit-exact to + control; speed is 3–8% worse, with a larger temporary output allocation. +- Tuning gate/up at the actual 7344 rows rather than 8192: complete chains + are 7.378 and 6.908 ms for layers 0/1, versus 7.367 and 6.952 ms. There is + no stable material improvement; checked outputs are bit-exact to control. + +### 2. HC narrow projection + +Hypothesis: the M~8K, N=88, K=10240 HC projection exposes too few concurrent +blocks. `hc_projection.py` compares the existing cuBLAS call with two, four, +and eight K partitions, FP32 partials and FP32 reduction. It includes strided +batch handling and reduction costs. Synthetic data screens throughput only; +any admitted candidate needs real activations/weights and paired quality. +The numerical screen is finite output and maximum row relative L2 <= 0.001 +against FP32 and against the current projection. + +Rejected: for 7344 rows, control is 0.503 ms and K partitions 2/4/8 are +0.532/0.537/0.568 ms; for 5347 rows, control is 0.360 ms and partitions are +0.425/0.419/0.411 ms. All outputs are finite; maximum relative L2 is 0.000298 +against FP32 and 0.000602 against control. More parallel blocks do not offset +the partial-output/reduction overhead. + +### 3. Larger custom all-reduce admission + +`collectives.py` uses all four idle GPUs, rank-max CUDA-event times, FP16 +rows x 2560 messages, and the production PyNccl/CustomAllreduce providers. +It includes eager custom-AR staging into its IPC buffer. Raising its capacity +to 48 MiB adds 80 MiB of GPU buffers per rank relative to the default 8 MiB. + +| Rows | NCCL AR ms | Custom AR ms | NCCL AG, width 88 ms | NCCL AG, width 640 ms | +|---|---:|---:|---:|---:| +| 7344 | 0.547 | 0.585 | 0.084 | 0.318 | +| 5347 | 0.411 | 0.431 | 0.076 | 0.254 | + +Rejected: increasing custom-AR capacity alone is 5–7% slower. Checked rank-zero +outputs are finite, relative L2 <= 0.000370. This is a synthetic throughput +screen, not full collective correctness or model qualification. Isolated NCCL +is faster than historical in-service call averages: rank arrival skew and +concurrent resource use need to be separated from transfer cost before another +communication change. + +### 4. Existing QSA rows route + +The saved real prefill input confirms the current rows route at 12.54 ms, +versus 23.07 ms for the old tiled route (1.84x). Maximum absolute output error +is 0.000977 and maximum per-head relative L2 is 0.000774; no NaNs. Requested +payload divided by runtime is about 683 GB/s, not a measured HBM counter. +This gain is already in the baseline and must not be counted again. + +### 5. HC communication/compute pipeline + +`hc_pipeline.py` uses the existing HC packet, gate-mix, and assembly kernels +with cuBLAS and PyNccl. It splits rows into two/four chunks and overlaps the +two gathers with subsequent projection work on a side stream. The timing +includes submissions, events, assembly, and final concatenation. + +Rejected: 7344-row sequential time is 1.343 ms, versus 2.390/4.985 ms for +two/four chunks. At 5347 rows, times are 1.056 versus 2.642/4.913 ms. This +Python pipeline adds more overhead than it hides. Output versus the same +tiles executed serially is bit-exact, isolating stream ordering from cuBLAS +shape-dependent arithmetic. The original maximum-row relative-L2 screen +of 0.001 also fails for the four-element injection output in three cases +(maximum 0.004036); it is not relaxed. Model-output errors are <= 0.0000761 +and all outputs are finite. No runtime change or quality qualification follows. + +### 6. Per-projection prefill streams + +Hypothesis: the shared cold stream can hold down DMA behind a gate/up bank +wait, and the shared resident stream orders next-layer gate/up transcode +before current-layer down transcode. `moe_streams.py` screens separate cold +streams and separate cold/resident streams independently, using the actual +native expert cache with 402 resident experts plus eight spare slots and +136 host slots. Canonical banks and staging bytes do not increase. + +Eight logical layers repeat four real TP4-rank-zero layer formats, wrapping +the four-view ring. Output and restarted-pass checks must be bit-exact; +native decode misses then change residency before another comparison. The +split-resident experiment explicitly joins each view's down readers before +reusing it, including leftover prefetch after an interrupted pass. Optional +GPU sleep models only a between-layer delay, not model compute contention. +This remains a benchmark-only hypothesis pending measured benefit. + +Rejected: with no artificial gap, eight-layer pass medians are 130.18 ms +shared, 130.59 ms cold-only split, 127.23 ms fully split, and 130.86 ms on +the repeated shared control. With a 20-million-cycle layer gap, all modes +take 186.3–186.9 ms. The small isolated benefit disappears once between-layer +work hides transfers. All output/restart comparisons are bit-exact, including +after 256 expert loads, 208 evictions/write-backs, and zero cache errors. +There is no runtime stream change. + +### 7. Reuse the existing dequantize-plus-cuBLAS dense provider + +The old trace attributes 450 ms to Q6_K bit-plane dense kernels and 149 ms to +Q5_K, across four chunks. They use 128x256 CTAs, approximately 248–255 +registers/thread, and one resident CTA/SM. These are static resource limits; +no Tensor Core utilization counter was collected. Dispatch inspection confirms +the affine fused route uses the default policy for M > 32, so an 8192-row +warmup does not autotune the measured 7344/5347 shapes. + +`dense_blas.py` uses real rank-zero TP4 weights and includes weight +dequantization into temporary scratch on every call. It compares fused, +dequantize+BLAS, then fused again. Maximum row-relative L2 against the FP32 +reference and paired output must be <= 0.003, with finite outputs. All ten +large-M cases and 25 threshold-screen cases pass. + +| Projection / source | M | Fused ms | Dequantize+BLAS ms | Fused repeat ms | +|---|---:|---:|---:|---:| +| QKV / Q6_K | 7344 | 1.867 | 1.191 | 1.839 | +| QKV / Q6_K | 5347 | 1.370 | 0.917 | 1.362 | +| QKV / Q5_K | 7344 | 1.908 | 1.143 | 1.921 | +| QKV / Q5_K | 5347 | 1.411 | 0.900 | 1.408 | +| SSM output / Q6_K | 7344 | 1.024 | 0.711 | 1.049 | +| SSM output / Q6_K | 5347 | 0.797 | 0.563 | 0.795 | +| Attention gate / Q6_K | 7344 | 1.929 | 0.697 | 1.975 | +| Attention gate / Q6_K | 5347 | 0.979 | 0.587 | 0.975 | + +Only four physical layouts are added to the existing calibration table: +(bits, group, K, N) = (5,32,2560,2560), (6,16,2560,2560), +(6,16,1536,2560), and (6,16,2560,1536), all at **M >= 2048 tokens**. +At M=512/1024 several layouts regress, so those retain fused computation. +The N=160 shared-expert screen is inconsistent and is not admitted. +Uncalibrated layouts, missing native operators, and failed workspace allocation +retain the existing fused fallback. No codec, provider, startup flag, or +workspace owner is added. The existing shared maximum-size scratch is retained; +its model-level memory impact is part of the pending service check. + +The affected GPU suite passes **15 tests**, including actual native route +observations at M=2047/2048/2049/4097, both eager and compiled calls, the +serving opaque projection, shared scratch, and changed-input graph replay. +Correctness support remains the existing affine provider's valid FP16 matrix +contract. Acceleration qualification is those four layouts; measured speed +points are M=256/512/1024/2048/4096/5347/7344, not every possible M or TP rank. + +### 8. Skip cold-expert holes while preserving staging offsets + +After decode, 110 cold experts occupy an extent of 135–136 host slots. The +old staging copy transfers that whole extent, including slots that the cold +transcode never reads. `cold_dma.py` first screens fragmented 110/104-slot +maps at real shard sizes: full copies of 73.98 MB take 16.05–16.37 ms; +selected copies take 13.31/12.71 ms including Python/native submissions. + +`ExpertCacheLane.stage_cold` now derives contiguous occupied ranges from the +same saved residency map consumed by the transcode, keeping every destination +offset unchanged. It reuses the existing host arena, staging buffer, stream +dependencies, and cache owner. It does not change residency, arithmetic, +spare slots, or GPU memory. A host slot must be valid for every cold expert; +invalid mappings raise an explicit error. + +`moe_dma.py` compares both copy policies in the actual eight-layer prefill +pipeline, with all other operations identical. Times are medians in ms: + +| Cache state / between-layer gap | Full extent | Occupied ranges | Full repeat | +|---|---:|---:|---:| +| Initial contiguous / none | 132.857 | 133.717 | 134.331 | +| Initial contiguous / 20M GPU cycles | 181.523 | 182.574 | 182.467 | +| Fragmented after decode / none | 161.027 | 136.131 | 162.643 | +| Fragmented after decode / 20M GPU cycles | 187.707 | 182.438 | 187.692 | + +All outputs and interrupted-pass restarts are bit-exact, including after 888 +expert loads, evictions, and write-backs, with zero cache errors. The gap is +synthetic delay, not real model compute. The complete affected cache suite +passes **9 tests**, including graph decode, a single hardware queue, mixed +prefill/decode, fragmented/slack-free host arenas, untouched staging holes, +and using the saved snapshot after the live CPU residency table changes. +Full-model latency remains a separate acceptance gate. + +### Matched request measurement + +`measure_prefill.py` uses deterministic word sequences and per-request hashes, +clears prefix caching on an exclusively owned server, and records streaming +first choice, first text, complete response, actual token count, and service +metrics. Service metric sums are attributed only when their count increases +by exactly one. Server prefill includes scheduled wall time, not just GPU +kernel time. `--profile` uses identical request generation; profiled wall time +is not an unprofiled speed result. + +Fresh baseline on the owned server, after two excluded warmups (13.35/8.01 s): + +| Prompt tokens | Client first text s | Server prefill s | Server TTFT s | +|---:|---:|---:|---:| +| 34688 | 7.895 | 7.544 | 7.757 | +| 34606 | 7.513 | 7.186 | 7.448 | +| 34631 | 7.429 | 7.165 | 7.376 | + +Median client TTFT is **7.513 s / 4605.9 tokens/s**. In the middle request, +server TTFT exceeds server prefill by 0.261 s, and client first text follows +server TTFT by 0.066 s. These metrics share one request, unlike the historical +four-chunk profile. The first measured request still shows warming drift. +The existing 30K-token retrieval passes 2/2. `quality.py` adds paired archive +retrieval plus arithmetic, production T0.6/top-p0.95/top-k20 with thinking, +strict expected JSON, normal EOS, and a repetition check. The 32761-token +baseline passes with 166 generated tokens; near-256K qualification is pending. + +The first near-256K baseline was interrupted before completion to keep the +cheap optimization screen first. Above the configured 131072-token QSA +history-stage limit, the original runtime falls back to direct host-history +attention and slows sharply (around 32 layer snapshots per 30 s, all GPUs +busy but around 85–90 W). No long-context quality result or candidate +regression is inferred. Complete this baseline/candidate boundary check +before promotion; the configured limit and deployment remain unchanged. + +First combined candidate, same five prompt hashes and token counts as control: + +| Prompt tokens | Client first text s | Server prefill s | Server TTFT s | +|---:|---:|---:|---:| +| 34688 | 7.502 | 7.171 | 7.387 | +| 34606 | 7.091 | 6.823 | 7.034 | +| 34631 | 7.099 | 6.790 | 7.021 | + +Median TTFT is **7.099 s / 4878.1 tokens/s**, improving 0.414 s (5.5%) from +the paired 7.513 s baseline. Excluded warmups take 14.13 and 7.56 s. +This is unprofiled; microbenchmark deltas are not substituted for this result. +It does **not** meet 6K tokens/s. The 32K quality prompt has the same hash and +passes the same expected JSON, EOS and repetition checks (167 generated tokens +versus control's 166). The 131070-token candidate also passes (161 generated +tokens, 60.44 s complete response); a matching 128K control is still pending, +so this does not establish a long-context performance comparison. + +The isolated cold-cache candidate startup took about 22 minutes, including +SM70+SM75 FlashQLA fallback compilation, checkpoint loading, and fresh main/MTP +graph compilation. This is not the warmed deployment's startup time. No OOM +occurred. The owned candidate was stopped after these checks for a bounded +expert-batch FP16 MoE screen and a matched Nsight profile. Its descendants +were recorded and only its verified API process received SIGTERM. + +## Execution discipline + +Microbenchmarks acquire the existing GPU locks and start only on an empty GPU. +Other tasks' services are left to their owners. Build output, native tuning, +Triton, and Python compiler caches are task-owned. The shared venv is unchanged. +Candidate native code is built through the normal CMake `_C` target; no private +sidecar kernel or preload is an accepted artifact. + +The first native candidate installs `_C` SHA256 +`02dc9ff2fca7fb2c154448b5f21d537ee29a547b0ade6f5b5e4018759a56bb63`; +control is `5f0d7a83415532ec817c3e23a7b24b5015d7cbacafcd3d452fabbc28efeadd5b`. +The padded candidate installs `_C` SHA256 +`fc0c457f39c2bb7b788c614b1d67c8226922a34ec44aef6e20d005c3c7c0e7cf`. +The build reused three unchanged Marlin objects after verifying identical +normalized compiler commands, all 15 relocated source dependencies per +object, object hashes, and an up-to-date original Ninja dependency graph +(5864 dependencies/object). Those objects were copied into an owned cache; +an audited compiler launcher lets the normal CMake `_C` target consume them +and records the relocated dependencies. It invokes the real compiler on any +mismatch. All modified code was freshly compiled. This is an incremental +build, not a claim that every unchanged object was rebuilt from scratch. +The installed extension has no RPATH/RUNPATH or private DSO dependencies; +its dependencies are the declared CUDA/Torch libraries and system C++/C +runtime. Build-cache files are not runtime dependencies. + +The first isolated baseline startup compiled FlashQLA in an empty task cache. +When checkpoint loading overlapped four native compiler processes, the 5 GiB +host-memory guard stopped only the owned service tree. No benchmark request +completed in that attempt. Subsequent full-model startup must wait until the +large native build completes; low-memory microbenchmarks can run concurrently. +The effective unchanged residency log is 402/512 experts plus eight spare GPU +slots, 136 host slots, 9.46 GiB GPU expert slots and 3.14 GiB pinned host per +rank, matching the historical deployment with configured fraction 0.8. + +Pending: candidate model quality, matched unprofiled end-to-end measurement, +and paired near-256K qualification. The original 256K deployment has not been +promoted to a candidate. The isolated candidate server is starting with the +corrected MoE tiles, four dense calibrations, and occupied-slot staging. diff --git a/benchmarks/flashnext_prefill/cold_dma.py b/benchmarks/flashnext_prefill/cold_dma.py new file mode 100644 index 0000000000..5bc35a8955 --- /dev/null +++ b/benchmarks/flashnext_prefill/cold_dma.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen copying only occupied cold slots, retaining their staging offsets. + +Uses three real TP4 layer-zero shard sizes and a synthetic fragmented slot +map. Timing includes Python/native submission of every contiguous range. +This is a DMA screen; model transcode overlap must be measured separately. +""" + +import argparse +import json +from functools import partial +from pathlib import Path + +import torch +from moe_tiles import artifact_identity, measure + +from vllm.model_executor.layers.quantization.gguf_expert_cache import _native +from vllm.transformers_utils.gguf_tensor_reader import quant_size + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + gate_block, gate_size = quant_size(18) + down_block, down_size = quant_size(20) + slot_bytes = [160 * (2560 // gate_block) * gate_size] * 2 + [ + 2560 * (160 // down_block) * down_size + ] + slots = 136 + host = [ + torch.empty(slots, size, dtype=torch.uint8).pin_memory() for size in slot_bytes + ] + for tensor in host: + tensor.random_(0, 256) + stage = [torch.empty_like(tensor, device="cuda") for tensor in host] + copy = _native().copy_h2d + stream = torch.cuda.current_stream().cuda_stream + report = {"artifact": artifact_identity(), "slot_bytes": slot_bytes, "cases": []} + for cold in (110, 104): + selected = torch.randperm(slots, generator=torch.Generator().manual_seed(9683))[ + :cold + ] + occupied = set(selected.tolist()) + extent = max(occupied) + 1 + ranges = [] + for slot in sorted(occupied): + if ranges and ranges[-1][1] == slot: + ranges[-1][1] += 1 + else: + ranges.append([slot, slot + 1]) + + def submit(spans): + for src, dst, size in zip(host, stage, slot_bytes): + for begin, end in spans: + copy( + dst.data_ptr() + begin * size, + src.data_ptr() + begin * size, + (end - begin) * size, + stream, + ) + + bulk = partial(submit, [(0, extent)]) + sparse = partial(submit, ranges) + for dst in stage: + dst.zero_() + sparse() + torch.accelerator.synchronize() + for src, dst in zip(host, stage): + actual = dst.cpu() + torch.testing.assert_close(actual[selected], src[selected], rtol=0, atol=0) + holes = torch.tensor(sorted(set(range(slots)) - occupied)) + assert (actual[holes] == 0).all() + row = { + "cold_slots": cold, + "extent": extent, + "ranges": len(ranges), + "bulk_bytes": extent * sum(slot_bytes), + "selected_bytes": cold * sum(slot_bytes), + "bulk": measure(bulk, 10), + "selected": measure(sparse, 10), + "bulk_repeat": measure(bulk, 10), + "cold_exact_holes_untouched": True, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/collectives.py b/benchmarks/flashnext_prefill/collectives.py new file mode 100644 index 0000000000..b538f66158 --- /dev/null +++ b/benchmarks/flashnext_prefill/collectives.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP4 prefill-sized collective screen; launch with torchrun on four idle GPUs. + +Uses the same PyNccl and custom-AR implementations as serving. Custom AR's larger +scratch allocation is experimental here, not a change to runtime admission. +""" + +import json +import os +import statistics +from functools import partial + +import torch +import torch.distributed as dist + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + + +def measure(fn): + samples = [] + for _ in range(5): + fn() + torch.accelerator.synchronize() + dist.barrier() + for _ in range(5): + begin, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + begin.record() + for _ in range(20): + fn() + end.record() + end.synchronize() + times = [None] * dist.get_world_size() + dist.all_gather_object(times, begin.elapsed_time(end) / 20) + samples.append(max(times)) + return {"median_ms": statistics.median(samples), "samples_ms": samples} + + +def main(): + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + nccl = PyNcclCommunicator(dist.group.WORLD, rank) + ca = CustomAllreduce(dist.group.WORLD, rank, max_size=48 * 1024**2) + assert not ca.disabled and not nccl.disabled + try: + for rows in (7344, 5347): + x = torch.randn(rows, 2560, device="cuda", dtype=torch.half) * 0.03 + control, candidate = torch.empty_like(x), torch.empty_like(x) + run_nccl = partial(nccl.all_reduce, x, control) + run_custom = partial(ca.all_reduce, x, out=candidate) + result = { + "tokens": rows, + "reduce_bytes": x.numel() * x.element_size(), + "nccl_all_reduce": measure(run_nccl), + "custom_all_reduce_eager": measure(run_custom), + } + difference = (candidate.float() - control.float()).norm() + result["relative_l2"] = (difference / control.float().norm()).item() + result["finite"] = bool(torch.isfinite(candidate).all()) + # Two HC packets per module: N=88 and N=640 per rank. + for n in (88, 640): + packet = torch.randn(rows, n, device="cuda", dtype=torch.half) + gathered = torch.empty(4 * rows, n, device="cuda", dtype=torch.half) + result[f"nccl_all_gather_n{n}"] = measure( + partial(nccl.all_gather, gathered, packet) + ) + if rank == 0: + print(json.dumps(result), flush=True) + finally: + torch.accelerator.synchronize() + ca.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/dense_blas.py b/benchmarks/flashnext_prefill/dense_blas.py new file mode 100644 index 0000000000..68e4cfd015 --- /dev/null +++ b/benchmarks/flashnext_prefill/dense_blas.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare existing fused and dequantize+cuBLAS dense providers on real TP4 weights. + +BLAS timing includes dequantization into its scratch on every call. This is a +provider screen, not a model qualification or a persistent FP16 weight cache. +""" + +import argparse +import json +from functools import partial +from pathlib import Path + +import numpy as np +import torch +from moe_tiles import artifact_identity, measure, relative_l2 + +from vllm.model_executor.layers.quantization.gguf_transcode import transcode_affine +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, dequantize + +CASES = [ + ("blk.0.attn_qkv.weight", 0), + ("blk.1.attn_qkv.weight", 0), + ("blk.0.ssm_out.weight", 1), + ("blk.1.attn_gate.weight", 0), + ("blk.5.ffn_up_shexp.weight", 0), +] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347]) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + reader = GGUFReader(args.model) + report = {"artifact": artifact_identity(), "cases": []} + for name, axis in CASES: + tensor = next(t for t in reader.tensors if t.name == name) + raw = np.ascontiguousarray(tensor.data) + kind = int(tensor.tensor_type) + canonical = transcode_affine(raw, kind).tp_slice(0, 4, axis=axis) + n, k = canonical.codes.shape + dense = dequantize(raw, kind) + dense = dense[:n] if axis == 0 else dense[:, :k] + reference_weight = torch.from_numpy(np.ascontiguousarray(dense)).cuda() + w, s, meta = torch.ops._C.gguf_affine_sm70_prepare( + torch.from_numpy(canonical.codes).cuda(), + torch.from_numpy(canonical.scales).cuda(), + torch.from_numpy(canonical.mins).cuda(), + canonical.bits, + canonical.group_size, + ) + k_ld, s_ld = meta.tolist() + scratch = torch.empty(k, n, dtype=torch.half, device="cuda") + for m in args.tokens: + torch.manual_seed(9683 + m) + x = (torch.randn(m, k, device="cuda") * 0.125).half() + fused_out = torch.empty(m, n, dtype=torch.half, device="cuda") + blas_out = torch.empty_like(fused_out) + + fused = partial( + torch.ops._C.gguf_affine_gemm_sm70_out, + fused_out, + x, + w, + s, + canonical.bits, + k_ld, + s_ld, + canonical.group_size, + ) + blas = partial( + torch.ops._C.gguf_affine_blas_sm70_out, + blas_out, + x, + w, + s, + canonical.bits, + scratch, + canonical.group_size, + ) + + fused() + blas() + expected = x[:16].float() @ reference_weight.float().T + errors = { + "fused_fp32_max_rel_l2": relative_l2(fused_out[:16], expected), + "blas_fp32_max_rel_l2": relative_l2(blas_out[:16], expected), + "paired_max_rel_l2": relative_l2(blas_out, fused_out), + } + passed = ( + bool(torch.isfinite(fused_out).all() and torch.isfinite(blas_out).all()) + and max(errors.values()) <= 0.003 + ) + row = { + "tensor": name, + "source_type": kind, + "tp_axis": axis, + "shape_mnk": [m, n, k], + "scratch_bytes": scratch.numel() * scratch.element_size(), + "fused": measure(fused, args.iterations), + "dequant_blas": measure(blas, args.iterations), + "fused_repeat": measure(fused, args.iterations), + **errors, + "passed": passed, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise SystemExit(0 if all(row["passed"] for row in report["cases"]) else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/hc_pipeline.py b/benchmarks/flashnext_prefill/hc_pipeline.py new file mode 100644 index 0000000000..7b841988ee --- /dev/null +++ b/benchmarks/flashnext_prefill/hc_pipeline.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen overlapping TP4 HC projections and gathers, without a model restart. + +Launch with torchrun on four idle V100s. This uses the production HC arithmetic +and NCCL provider with synthetic checkpoint-shaped inputs. It is a throughput +screen, not a model quality or runtime-state qualification. +""" + +import json +import os + +import torch +import torch.distributed as dist +from collectives import measure +from moe_tiles import relative_l2 + +from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator +from vllm.models.qwen4_exp.nvidia.ops.hc import ( + hc_assemble_ranks, + hc_gate_mix_local, + hc_packet_silu, +) + +ORDER = (0, 1, 2, 3) + + +def sequential(nccl, x, down, up, rank): + local = torch.mm(x, down.T) + packet = local.new_empty((4 * x.shape[0], 88)) + nccl.all_gather(packet, local) + lora, injection = hc_packet_silu(packet, ORDER, 4) + local_gate = torch.mm(lora, up.T) + mixed = hc_gate_mix_local(x, local_gate, rank, 4) + gathered = mixed.new_empty((4 * x.shape[0], 640)) + nccl.all_gather(gathered, mixed) + return hc_assemble_ranks(gathered, ORDER), injection + + +class Pipeline: + def __init__(self, nccl, x, down, up, rank, chunks, overlap=True): + self.nccl, self.down, self.up, self.rank = nccl, down, up, rank + self.x = x.chunk(chunks) + self.comm = torch.cuda.Stream() if overlap else torch.cuda.current_stream() + self.ready = [torch.cuda.Event() for _ in self.x] + self.packet_done = [torch.cuda.Event() for _ in self.x] + self.mixed_ready = [torch.cuda.Event() for _ in self.x] + self.gather_done = [torch.cuda.Event() for _ in self.x] + self.packet = [v.new_empty((4 * v.shape[0], 88)) for v in self.x] + self.gathered = [v.new_empty((4 * v.shape[0], 640)) for v in self.x] + + def __call__(self): + main = torch.cuda.current_stream() + # Retain the inputs until all their asynchronous gathers are queued. + locals_, mixed, injections = [], [], [] + for i, x in enumerate(self.x): + locals_.append(torch.mm(x, self.down.T)) + self.ready[i].record(main) + with torch.cuda.stream(self.comm): + self.comm.wait_event(self.ready[i]) + self.nccl.all_gather(self.packet[i], locals_[-1]) + locals_[-1].record_stream(self.comm) + self.packet_done[i].record(self.comm) + for i, x in enumerate(self.x): + main.wait_event(self.packet_done[i]) + lora, injection = hc_packet_silu(self.packet[i], ORDER, 4) + injections.append(injection) + gate = torch.mm(lora, self.up.T) + mixed.append(hc_gate_mix_local(x, gate, self.rank, 4)) + self.mixed_ready[i].record(main) + with torch.cuda.stream(self.comm): + self.comm.wait_event(self.mixed_ready[i]) + self.nccl.all_gather(self.gathered[i], mixed[-1]) + mixed[-1].record_stream(self.comm) + self.gather_done[i].record(self.comm) + outputs = [] + for i, gathered in enumerate(self.gathered): + main.wait_event(self.gather_done[i]) + outputs.append(hc_assemble_ranks(gathered, ORDER)) + return torch.cat(outputs), torch.cat(injections) + + +def main(): + rank = int(os.environ["LOCAL_RANK"]) + torch.accelerator.set_device_index(rank) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + nccl = PyNcclCommunicator(dist.group.WORLD, rank) + assert not nccl.disabled + passed = True + try: + for rows in (7344, 5347): + torch.manual_seed(rows) + x = torch.randn(rows, 10240, device="cuda", dtype=torch.half) * 0.3 + torch.manual_seed(100 + rank) + down = torch.randn(88, 10240, device="cuda", dtype=torch.half) * 0.01 + up = torch.randn(2560, 320, device="cuda", dtype=torch.half) * 0.01 + + def control(x=x, down=down, up=up): + return sequential(nccl, x, down, up, rank) + + expected = control() + result = {"tokens": rows, "sequential": measure(control)} + for chunks in (2, 4): + candidate = Pipeline(nccl, x, down, up, rank, chunks) + actual = candidate() + same_tiles = Pipeline(nccl, x, down, up, rank, chunks, overlap=False)() + error = [relative_l2(a, b) for a, b in zip(actual, expected)] + serial_error = max( + relative_l2(a, b) for a, b in zip(actual, same_tiles) + ) + finite = all(bool(torch.isfinite(a).all()) for a in actual) + errors = [None] * 4 + dist.all_gather_object(errors, (error, serial_error, finite)) + numerical_pass = all( + max(e) <= 0.001 and s == 0 and f for e, s, f in errors + ) + passed &= numerical_pass + # Failed numerical screens remain failed. Timing them decides + # whether further diagnosis is worthwhile, not admission. + result[f"chunks_{chunks}"] = measure(candidate) | { + "max_rank_output_relative_l2": max(e[0] for e, _, _ in errors), + "max_rank_injection_relative_l2": max(e[1] for e, _, _ in errors), + "overlap_vs_serial_same_tiles": max(s for _, s, _ in errors), + "finite_all_ranks": all(f for _, _, f in errors), + "numerical_pass": numerical_pass, + } + if rank == 0: + print(json.dumps(result), flush=True) + finally: + torch.accelerator.synchronize() + dist.destroy_process_group() + if not passed: + raise SystemExit("Numerical screen failed; candidate is not qualified") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/hc_projection.py b/benchmarks/flashnext_prefill/hc_projection.py new file mode 100644 index 0000000000..df3c412aa5 --- /dev/null +++ b/benchmarks/flashnext_prefill/hc_projection.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test whether HC's narrow N=88 projection benefits from more K partitions. + +Synthetic FP16 operands, serving dimensions; all packing/reduction costs are +included. This is a rejection/admission experiment, not model qualification. +""" + +import argparse +import json +from functools import partial + +import torch +from moe_tiles import measure, relative_l2 + + +def partitioned(x, weight, splits): + rows, k = x.shape + n = weight.shape[0] + partials = torch.bmm( + x.view(rows, splits, k // splits).permute(1, 0, 2), + weight.view(n, splits, k // splits).permute(1, 2, 0), + out_dtype=torch.float32, + ) + return partials.sum(dim=0).half() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347]) + args = parser.parse_args() + torch.backends.cuda.matmul.allow_tf32 = False + for rows in args.tokens: + gen = torch.Generator(device="cuda").manual_seed(rows) + x = torch.randn(rows, 10240, device="cuda", dtype=torch.half, generator=gen) + weight = ( + torch.randn(88, 10240, device="cuda", generator=gen) / 10240**0.5 + ).half() + expected = x.float() @ weight.float().T + baseline = torch.mm(x, weight.T) + result = {"tokens": rows, "candidates": []} + result["baseline"] = measure(partial(torch.mm, x, weight.T), 20) + result["baseline_rel_l2"] = relative_l2(baseline, expected) + for splits in (2, 4, 8): + candidate = partitioned(x, weight, splits) + result["candidates"].append( + { + "splits": splits, + **measure(partial(partitioned, x, weight, splits), 20), + "fp32_rel_l2": relative_l2(candidate, expected), + "baseline_rel_l2": relative_l2(candidate, baseline), + "finite": bool(torch.isfinite(candidate).all()), + } + ) + print(json.dumps(result), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/measure_prefill.py b/benchmarks/flashnext_prefill/measure_prefill.py new file mode 100644 index 0000000000..fbebad9d23 --- /dev/null +++ b/benchmarks/flashnext_prefill/measure_prefill.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Matched, cold-prefix prefill requests against an exclusively owned server. + +Use the same --seed/--words/--warmup/--trials for control and candidate. Word +count is not token count: throughput uses the server's reported prompt_tokens. +--profile wraps only measured requests in the server's profiler window; do not +use profiled timing as an unprofiled end-to-end result. +""" + +import argparse +import hashlib +import json +import random +import statistics +import time +import urllib.request +from pathlib import Path + +WORDS = "alpha river stone cloud engine paper signal forest copper lantern orbit meadow" +METRICS = ( + "vllm:request_prefill_time_seconds", + "vllm:time_to_first_token_seconds", + "vllm:request_queue_time_seconds", +) + + +def request(base, path, body=None): + data = None if body is None else json.dumps(body).encode() + req = urllib.request.Request( + base + path, data, {"Content-Type": "application/json"} + ) + return urllib.request.urlopen(req, timeout=1800) + + +def metrics(base): + with request(base, "/metrics") as response: + lines = response.read().decode().splitlines() + result = {} + for line in lines: + if not line or line.startswith("#"): + continue + name = line.split("{", 1)[0].split(" ", 1)[0] + if name in {m + suffix for m in METRICS for suffix in ("_sum", "_count")}: + result[name] = result.get(name, 0) + float(line.rsplit(" ", 1)[1]) + return result + + +def metric_deltas(base, before): + deadline = time.monotonic() + 8 + while True: + after = metrics(base) + counts = { + m: after.get(m + "_count", 0) - before.get(m + "_count", 0) for m in METRICS + } + if all(n >= 1 for n in counts.values()) or time.monotonic() >= deadline: + break + time.sleep(0.2) + # More than one request means this measurement cannot be attributed. + return { + m.removeprefix("vllm:"): { + "count": counts[m], + "seconds": ( + after.get(m + "_sum", 0) - before.get(m + "_sum", 0) + if counts[m] == 1 + else None + ), + } + for m in METRICS + } + + +def run_one(args, index): + rng = random.Random(args.seed + index) + # An early request-specific nonce also avoids sharing warmup prefixes. + prompt = f"[{rng.randrange(10**9)}] " + " ".join( + rng.choice(WORDS.split()) for _ in range(args.words) + ) + body = { + "model": args.model, + "prompt": prompt, + "temperature": 0, + "max_tokens": 1, + "stream": True, + "stream_options": {"include_usage": True}, + } + # This endpoint is intended only for a server owned by this benchmark. + with request(args.base_url, "/reset_prefix_cache", {}) as response: + response.read() + before = metrics(args.base_url) + start = time.perf_counter() + first_choice = first_text = None + usage, choices = None, [] + with request(args.base_url, "/v1/completions", body) as response: + for line in response: + if not line.startswith(b"data: "): + continue + data = line[6:].strip() + if data == b"[DONE]": + break + event = json.loads(data) + if "error" in event: + raise RuntimeError(event["error"]) + if event.get("choices"): + choices.extend(event["choices"]) + elapsed = time.perf_counter() - start + if first_choice is None: + first_choice = elapsed + if first_text is None and any(c.get("text") for c in event["choices"]): + first_text = elapsed + if event.get("usage"): + usage = event["usage"] + total = time.perf_counter() - start + if first_choice is None or usage is None: + raise RuntimeError("Missing completion choice or final usage") + return { + "index": index, + "prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(), + "usage": usage, + "first_choice_seconds": first_choice, + "first_text_seconds": first_text, + "response_seconds": total, + "prompt_tokens_per_first_choice_second": usage["prompt_tokens"] / first_choice, + "choices": choices, + "server_metrics": metric_deltas(args.base_url, before), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8002") + parser.add_argument("--model", default="flash-next") + parser.add_argument("--words", type=int, default=32000) + parser.add_argument("--seed", type=int, default=9683) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--profile", action="store_true") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.words < 1 or args.warmup < 0 or args.trials < 1: + parser.error("words/trials must be positive and warmup nonnegative") + report = {"config": vars(args) | {"output": str(args.output)}, "requests": []} + for i in range(args.warmup + args.trials): + profiling = args.profile and i >= args.warmup + if profiling: + with request(args.base_url, "/start_profile", {}) as response: + response.read() + try: + result = run_one(args, i) + finally: + if profiling: + with request(args.base_url, "/stop_profile", {}) as response: + response.read() + result["warmup"] = i < args.warmup + report["requests"].append(result) + print(json.dumps(result), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + measured = [r for r in report["requests"] if not r["warmup"]] + report["summary"] = { + "profiled": args.profile, + "median_first_choice_seconds": statistics.median( + r["first_choice_seconds"] for r in measured + ), + "median_prompt_tokens_per_first_choice_second": statistics.median( + r["prompt_tokens_per_first_choice_second"] for r in measured + ), + } + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report["summary"]), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_dma.py b/benchmarks/flashnext_prefill/moe_dma.py new file mode 100644 index 0000000000..812cd20356 --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_dma.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare full-extent and occupied-slot DMA in the actual prefill pipeline. + +The eight-layer cache, residency, staging allocation, transcode, and GEMMs are +identical in both modes. Initial and fragmented post-decode maps are timed. +""" + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from moe_streams import build_cache +from moe_tiles import EXPERTS, HIDDEN, TOPK, artifact_identity, measure + +from vllm.model_executor.layers.quantization.gguf_expert_cache import SHARDS, _native +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader + + +class LaneProbe: + def __init__(self, cache): + self.cache, self.lane = cache, cache.lane + self.mode = "bulk" + + def __getattr__(self, name): + return getattr(self.lane, name) + + def stage_cold(self, i, shard, dst, stream): + if self.mode == "sparse": + return self.lane.stage_cold(i, shard, dst, stream) + cache = self.cache + layer = cache.layers[i] + rows = cache.tm.extent[i % cache.tm.VIEWS] + if rows: + _native().copy_h2d( + dst, + cache.host[layer, SHARDS[shard]].data_ptr(), + rows * cache.sizes[layer, SHARDS[shard]], + stream, + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.manual_seed(9683) + cache = build_cache(GGUFReader(args.model)) + cache.lane = LaneProbe(cache) + x = (torch.randn(7344, HIDDEN, device="cuda") * 0.5).half() + scores, ids = torch.randn(7344, EXPERTS, device="cuda").topk(TOPK, dim=1) + weights = scores.softmax(dim=1) + native = Path(_native().__file__) + report = { + "artifact": artifact_identity(), + "cache_native_sha256": hashlib.sha256(native.read_bytes()).hexdigest(), + "cases": [], + "maps": {}, + } + + def forward(gap=0, count=8): + outputs = [] + for layer in cache.layers[:count]: + outputs.append(cache.prefill_moe(layer, x, ids, weights)) + if gap: + torch.cuda._sleep(gap) + return outputs + + def check(): + for actual, expected in zip(forward(), reference): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + assert cache.tm.error.item() == 0 + + try: + reference = [out.cpu() for out in forward()] + for phase in ("initial", "after_decode"): + if phase == "after_decode": + for _ in range(10): + for layer in cache.layers: + cache.resolve( + layer, torch.randperm(EXPERTS, device="cuda")[:50] + ) + torch.accelerator.synchronize() + report["maps"][phase] = [] + for i in range(len(cache.layers)): + host = cache.host_slot_of[i][cache.slot_of[i] < 0] + report["maps"][phase].append( + {"cold": host.numel(), "extent": int(host.max()) + 1} + ) + for gap in (0, 20000000): + for mode in ("bulk", "sparse", "bulk"): + torch.accelerator.synchronize() + cache.lane.mode = mode + cache.tm.prefetched.clear() + cache.tm.start() + check() + forward(gap, count=4) + check() # Restart while the interrupted pass has pending copies. + row = { + "phase": phase, + "mode": mode, + "gap_cycles": gap, + **measure(lambda gap=gap: forward(gap), 2), + "exact_output_and_restart": True, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + report["cache_stats"] = cache.stats() + assert report["cache_stats"]["errors"] == 0 + args.output.write_text(json.dumps(report, indent=2) + "\n") + finally: + torch.accelerator.synchronize() + cache.close() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_streams.py b/benchmarks/flashnext_prefill/moe_streams.py new file mode 100644 index 0000000000..48e27a62ca --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_streams.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen prefill stream dependencies with real weights and the native cache. + +Uses eight logical layers (four real TP4 rank-zero layers repeated) to wrap +the residency-view ring. All modes share kernels, weights, expert residency, +and bank/staging capacity. Optional GPU sleep models a between-layer delay, +not concurrent model compute. This is not an end-to-end throughput result. +""" + +import argparse +import json +import statistics +from pathlib import Path + +import torch +from moe_tiles import EXPERTS, HIDDEN, INTERMEDIATE, TOPK, artifact_identity, read_layer + +from vllm.model_executor.layers.quantization.gguf_expert_cache import ( + SHARDS, + ExpertCache, + TurboMindPrefill, +) +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader + + +class StreamProbe(TurboMindPrefill): + """Benchmark-only policies; production _issue/_locate/compute are reused.""" + + mode = "shared" + + def start(self): + super().start() + device = self.cache.device + self.cold_by_key = { + "gu": self.cold_stream, + "dn": torch.cuda.Stream(device, priority=-1), + } + self.resident_by_key = { + "gu": self.resident_stream, + "dn": torch.cuda.Stream(device, priority=-1), + } + self.located = [torch.cuda.Event() for _ in range(self.VIEWS)] + self.view_done = [torch.cuda.Event() for _ in range(self.VIEWS)] + for event in self.view_done: + event.record() + + def _locate(self, i): + p = i % self.VIEWS + if self.mode == "split": + # Down may still use this view after an interrupted pass. + self.resident_stream.wait_event(self.view_done[p]) + super()._locate(i) + self.located[p].record(self.resident_stream) + + def _issue(self, key, i): + self.cold_stream = self.cold_by_key["gu" if self.mode == "shared" else key] + self.resident_stream = self.resident_by_key[ + key if self.mode == "split" else "gu" + ] + if key == "dn" and self.mode == "split": + self.resident_stream.wait_event(self.located[i % self.VIEWS]) + super()._issue(key, i) + if key == "dn": + self.view_done[i % self.VIEWS].record(self.resident_stream) + + +def build_cache(reader): + real_layers = [0, 1, 17, 47] + sources = [read_layer(reader, layer) for layer in real_layers] + banks = [] + for layer in range(8): + raw, gate_type, down_type, left = sources[layer % len(sources)] + for shard, part in zip(SHARDS, ("gate", "up", "down")): + source = ( + (down_type, INTERMEDIATE, left) + if part == "down" + else (gate_type, HIDDEN, 0) + ) + banks.append((layer, shard, raw[part].shape, 0, source)) + cache = ExpertCache(banks, EXPERTS, 0.8, 8, 0.05, None, "cuda:0", "turbomind", 4096) + for layer in cache.layers: + raw = sources[layer % len(sources)][0] + for shard, part in zip(SHARDS, ("gate", "up", "down")): + cache.store_bank(layer, shard, torch.from_numpy(raw[part])) + cache.tm.__class__ = StreamProbe + cache.start() + return cache + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--tokens", type=int, default=7344) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--gap-cycles", type=int, nargs="+", default=[0, 20000000]) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.tokens < 1 or args.iterations < 1 or min(args.gap_cycles) < 0: + parser.error("invalid tokens, iterations, or gap cycles") + torch.manual_seed(9683) + cache = build_cache(GGUFReader(args.model)) + gen = torch.Generator(device="cuda").manual_seed(9683) + x = (torch.randn(args.tokens, HIDDEN, device="cuda", generator=gen) * 0.5).half() + scores, ids = torch.randn(args.tokens, EXPERTS, device="cuda", generator=gen).topk( + TOPK, dim=1 + ) + weights = scores.softmax(dim=1) + report = { + "artifact": artifact_identity(), + "tokens": args.tokens, + "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(), + "resident_experts": cache.resident0, + "host_slots": cache.host_slots, + "canonical_and_staging_bytes": cache.tm.storage_bytes(), + "cases": [], + } + + def forward(gap, count=8): + outputs = [] + for layer in cache.layers[:count]: + outputs.append(cache.prefill_moe(layer, x, ids, weights)) + if gap: + torch.cuda._sleep(gap) + return outputs + + def select(mode): + torch.accelerator.synchronize() + cache.tm.prefetched.clear() + cache.tm.mode = mode + cache.tm.start() + + def compare(outputs, reference): + # Stream changes preserve arithmetic: exact output is the contract. + for actual, expected in zip(outputs, reference): + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + assert cache.tm.error.item() == 0 + + try: + # Warm native tuning, Triton, allocation pools and every source format. + reference = [out.cpu() for out in forward(0)] + for gap in args.gap_cycles: + for mode in ("shared", "cold_split", "split", "shared"): + select(mode) + compare(forward(gap), reference) + # Restart without draining side-stream work from this pass. + forward(gap, count=4) + compare(forward(gap), reference) + samples = [] + for _ in range(5): + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + for _ in range(args.iterations): + forward(gap) + end.record() + end.synchronize() + samples.append(begin.elapsed_time(end) / args.iterations) + case = { + "mode": mode, + "gap_cycles": gap, + "median_ms": statistics.median(samples), + "samples_ms": samples, + "exact_output_and_restart": True, + } + report["cases"].append(case) + print(json.dumps(case), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + # Exercise real decode misses, host write-backs, and a new snapshot. + for layer in cache.layers: + cache.resolve(layer, torch.arange(480, 512, device="cuda")) + torch.accelerator.synchronize() + for mode in ("shared", "cold_split", "split"): + select(mode) + compare(forward(0), reference) + report["after_decode_residency_change_exact"] = True + report["cache_stats"] = cache.stats() + assert report["cache_stats"]["errors"] == 0 + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report["cache_stats"]), flush=True) + finally: + torch.accelerator.synchronize() + cache.close() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_tiles.py b/benchmarks/flashnext_prefill/moe_tiles.py new file mode 100644 index 0000000000..b13fd544d2 --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_tiles.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Real-weight TP4 MoE microbenchmark, using the serving route/subroute code. + +Run control/candidate in fresh processes with separate native tuning caches. +This isolates resident canonical-bank compute: it does not measure host DMA, +transcoding overlap, TP communication, or end-to-end model quality. +""" + +import argparse +import hashlib +import json +import statistics +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +import vllm._C_gguf # noqa: F401 + +import vllm._custom_ops # noqa: F401 +from vllm.model_executor.layers.quantization.gguf_expert_cache import ( + TurboMindPrefill, + _canonical, +) +from vllm.model_executor.layers.quantization.gguf_native import packed_tp_span +from vllm.transformers_utils.gguf_tensor_reader import ( + GGUFReader, + dequantize, + quant_size, +) + +EXPERTS, TOPK, HIDDEN, INTERMEDIATE = 512, 10, 2560, 160 + + +def artifact_identity(): + import vllm._C + + import vllm + + native = Path(vllm._C.__file__).resolve() + return { + "package": str(Path(vllm.__file__).resolve()), + "native": str(native), + "native_sha256": hashlib.sha256(native.read_bytes()).hexdigest(), + } + + +def measure(fn, iterations): + for _ in range(3): + fn() + torch.accelerator.synchronize() + samples = [] + for _ in range(5): + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + for _ in range(iterations): + fn() + end.record() + end.synchronize() + samples.append(begin.elapsed_time(end) / iterations) + return {"median_ms": statistics.median(samples), "samples_ms": samples} + + +def pointer_table(tensor, stride, ld): + table = torch.empty(EXPERTS, 2, dtype=torch.int64, device="cuda") + table[:, 0] = tensor.data_ptr() + torch.arange(EXPERTS, device="cuda") * stride + table[:, 1] = ld + return table.view(torch.uint8).flatten() + + +def relative_l2(actual, expected): + actual, expected = actual.float(), expected.float() + return ( + ((actual - expected).norm(dim=-1) / expected.norm(dim=-1).clamp_min(1e-6)) + .max() + .item() + ) + + +def read_layer(reader, layer): + tensors = { + part: next( + t for t in reader.tensors if t.name == f"blk.{layer}.ffn_{part}_exps.weight" + ) + for part in ("gate", "up", "down") + } + gtype, dtype = (int(tensors[p].tensor_type) for p in ("gate", "down")) + assert int(tensors["up"].tensor_type) == gtype + raw = { + p: np.ascontiguousarray(np.asarray(tensors[p].data)[:, :INTERMEDIATE]) + for p in ("gate", "up") + } + full = np.asarray(tensors["down"].data) + block, size = quant_size(dtype) + first, last, left, _ = packed_tp_span(full.shape[2] // size * block, dtype, 4, 0) + raw["down"] = np.ascontiguousarray(full[:, :, first:last]) + return raw, gtype, dtype, left + + +def load_layer(reader, layer): + raw, gtype, dtype, left = read_layer(reader, layer) + source = {p: torch.from_numpy(a).cuda() for p, a in raw.items()} + info = { + "gate": (gtype,), + "down": (dtype,), + "rows": INTERMEDIATE, + "hidden": HIDDEN, + "gu": _canonical(gtype, 2 * INTERMEDIATE, HIDDEN), + "dn": _canonical(dtype, HIDDEN, INTERMEDIATE), + } + banks, tables = {}, {} + ids = torch.arange(EXPERTS, dtype=torch.int32, device="cuda") + count = torch.full((1,), EXPERTS, dtype=torch.int32, device="cuda") + error = torch.zeros(1, dtype=torch.int32, device="cuda") + pointers = { + p: t.data_ptr() + + torch.arange(EXPERTS, device="cuda", dtype=torch.int64) * t[0].numel() + for p, t in source.items() + } + for key, part, n, k, skip in ( + ("gu", "gate", 2 * INTERMEDIATE, HIDDEN, 0), + ("dn", "down", HIDDEN, INTERMEDIATE, left), + ): + _, _, _, weight_bytes, stats_bytes = info[key] + weight = torch.empty(EXPERTS * weight_bytes, dtype=torch.uint8, device="cuda") + stats = torch.empty(EXPERTS * stats_bytes, dtype=torch.uint8, device="cuda") + banks[key] = (weight, stats) + tables[key] = ( + pointer_table(weight, weight_bytes, 32 * k), + pointer_table(stats, stats_bytes, n), + ) + torch.ops._C.gguf_prefill_transcode_sm70( + weight, + stats, + weight_bytes, + stats_bytes, + pointers[part], + pointers["up"] if key == "gu" else None, + ids, + count, + 1024, + gtype if key == "gu" else dtype, + n, + k, + skip, + raw[part].shape[2], + error, + ) + assert error.item() == 0, "transcode reported an invalid expert" + return info, tables, banks, raw, left + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--layers", type=int, nargs="+", default=[0, 1, 17, 47]) + parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347]) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--tune-tokens", type=int, default=8192) + parser.add_argument("--down-tokens", type=int, default=4096) + parser.add_argument("--routing", choices=["uniform", "skewed"], default="uniform") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--save-reference", type=Path) + parser.add_argument("--compare-reference", type=Path) + args = parser.parse_args() + assert args.iterations > 0 and min(args.tokens) > 0 and args.down_tokens > 0 + torch.backends.cuda.matmul.allow_tf32 = False + reader = GGUFReader(args.model) + references = ( + torch.load(args.compare_reference, map_location="cpu", weights_only=True) + if args.compare_reference + else {} + ) + saved, results = {}, [] + report = { + "artifact": artifact_identity(), + "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(), + "routing": args.routing, + "tune_tokens": args.tune_tokens, + "down_tokens": args.down_tokens, + "cases": results, + } + for layer in args.layers: + info, tables, banks, raw, left = load_layer(reader, layer) + # Reuse the production sorting, sub-batching, fused GEMM, and unroute. + runner = TurboMindPrefill.__new__(TurboMindPrefill) + runner.cache = SimpleNamespace(experts=EXPERTS, prefill_tokens=args.down_tokens) + runner.expert_ids = torch.arange(EXPERTS + 1, dtype=torch.int32, device="cuda") + runner.tuned = set() + probe = torch.zeros(args.tune_tokens, HIDDEN, dtype=torch.half, device="cuda") + runner._tune(layer, info, tables, probe) + del probe + dequantized = {} + + def ref_weight( + expert, part, dequantized=dequantized, info=info, raw=raw, left=left + ): + key = (expert, part) + if key not in dequantized: + ty = info["down" if part == "down" else "gate"][0] + array = dequantize(raw[part][expert], ty) + if part == "down": + array = array[:, left : left + INTERMEDIATE] + dequantized[key] = ( + torch.from_numpy(np.ascontiguousarray(array)).cuda().float() + ) + return dequantized[key] + + for tokens in args.tokens: + gen = torch.Generator(device="cuda").manual_seed(tokens) + x = (torch.randn(tokens, HIDDEN, device="cuda", generator=gen) * 0.5).half() + scores = torch.randn(tokens, EXPERTS, device="cuda", generator=gen) + if args.routing == "skewed": + scores[:, :32] += 2 + scores, ids = scores.topk(TOPK, dim=1) + weights = scores.softmax(dim=1) + route = runner._route(ids, 0, EXPERTS) + subroutes = ( + [ + ( + start, + min(tokens, start + args.down_tokens), + runner._subroute( + route, start, min(tokens, start + args.down_tokens) + ), + ) + for start in range(0, tokens, args.down_tokens) + ] + if tokens > args.down_tokens + else [(0, tokens, route)] + ) + out = torch.empty_like(x) + + def gu(runner=runner, info=info, tables=tables, x=x, route=route): + return runner._gate_up(info, tables, x, route, False) + + hidden = gu() + + def down( + h=hidden, + runner=runner, + info=info, + tables=tables, + subroutes=subroutes, + weights=weights, + out=out, + ): + for start, end, subroute in subroutes: + runner._down( + info, + tables, + h, + subroute, + weights[start:end], + out[start:end], + False, + ) + + def chain(gu=gu, down=down): + down(gu()) + + result = { + "layer": layer, + "gate_type": info["gate"][0], + "down_type": info["down"][0], + "tokens": tokens, + "gate_up": measure(gu, args.iterations), + "down_unroute": measure(down, args.iterations), + "chain": measure(chain, args.iterations), + } + down() + # FP32 reference at expert boundaries catches ragged-row/layout errors. + offsets = route[0].cpu().tolist() + positions = sorted( + { + p + for e in (0, 1, 31, 32, 127, 255, 510, 511) + if offsets[e] < offsets[e + 1] + for p in (offsets[e], offsets[e + 1] - 1) + } + ) + hidden_errors = [] + for position in positions: + token = int(route[2][position]) + expert = int(ids.flatten()[route[1][position]]) + expected = torch.nn.functional.silu( + ref_weight(expert, "gate") @ x[token].float() + ) + expected *= ref_weight(expert, "up") @ x[token].float() + hidden_errors.append(relative_l2(hidden[position], expected)) + chosen_tokens = torch.tensor( + sorted({0, 1, tokens // 2, tokens - 1}), device="cuda" + ).clamp_max(tokens - 1) + sample = { + "hidden": hidden[positions].cpu(), + "out": out[chosen_tokens].cpu(), + } + expected_rows = [] + for token in chosen_tokens.tolist(): + expected_row = torch.zeros(HIDDEN, device="cuda") + for slot in range(TOPK): + expert = int(ids[token, slot]) + expected_hidden = torch.nn.functional.silu( + ref_weight(expert, "gate") @ x[token].float() + ) * (ref_weight(expert, "up") @ x[token].float()) + expected_down = ( + (ref_weight(expert, "down") @ expected_hidden.half().float()) + .half() + .float() + ) + expected_row += weights[token, slot] * expected_down + expected_rows.append(expected_row.half()) + output_error = relative_l2(out[chosen_tokens], torch.stack(expected_rows)) + key = f"{layer}/{tokens}/{args.routing}" + result["hidden_fp32_max_rel_l2"] = max(hidden_errors) + result["output_fp32_max_rel_l2"] = output_error + result["finite"] = bool( + torch.isfinite(hidden).all() & torch.isfinite(out).all() + ) + passed = ( + result["finite"] and max(hidden_errors) <= 0.01 and output_error <= 0.01 + ) + if args.compare_reference: + errors = { + name: relative_l2(value, references[key][name]) + for name, value in sample.items() + } + result["control_max_rel_l2"] = errors + passed &= max(errors.values()) <= 0.003 + result["passed"] = passed + saved[key] = sample + results.append(result) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(result), flush=True) + del banks, tables, dequantized, hidden, out, runner + if args.save_reference: + torch.save(saved, args.save_reference) + if not all(case["passed"] for case in results): + raise SystemExit("numerical gate failed") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/nsys_breakdown.py b/benchmarks/flashnext_prefill/nsys_breakdown.py index 7af351a9c2..7d0e3270e2 100644 --- a/benchmarks/flashnext_prefill/nsys_breakdown.py +++ b/benchmarks/flashnext_prefill/nsys_breakdown.py @@ -14,22 +14,27 @@ c = sqlite3.connect(sys.argv[1]) dev = int(sys.argv[2]) names = dict(c.execute("select id, value from StringIds").fetchall()) -api = dict( - c.execute("select correlationId, end from CUPTI_ACTIVITY_KIND_RUNTIME").fetchall() -) +# CUDA correlation ids are process-local. Nsight packs the process id above +# the low 24 thread-id bits; joining only correlationId mixes TP workers. +api = { + (tid >> 24, cid): end + for tid, cid, end in c.execute( + "select globalTid, correlationId, end from CUPTI_ACTIVITY_KIND_RUNTIME" + ) +} ks = c.execute( - "select start,end,streamId,correlationId,demangledName,gridX,gridY " + "select start,end,streamId,correlationId,demangledName,gridX,gridY,globalPid " "from CUPTI_ACTIVITY_KIND_KERNEL where deviceId=?", (dev,), ).fetchall() ms = c.execute( - "select start,end,streamId,correlationId,NULL,0,0 " + "select start,end,streamId,correlationId,NULL,0,0,globalPid " "from CUPTI_ACTIVITY_KIND_MEMCPY where deviceId=?", (dev,), ).fetchall() -st = collections.Counter(r[2] for r in ks) +st = collections.Counter((r[-1], r[2]) for r in ks) main = st.most_common(1)[0][0] -m = sorted(r for r in ks + ms if r[2] == main) +m = sorted(r for r in ks + ms if (r[-1], r[2]) == main) phase = "gu" @@ -128,11 +133,12 @@ def seg(n, gx, gy): wait = collections.defaultdict(float) host = 0 end = m[0][0] -for s, e, _, cid, nm, gx, gy in m: +for s, e, _, cid, nm, gx, gy, pid in m: name = names.get(nm) if nm is not None else None g = seg(name, gx, gy) if s > end: - if api.get(cid) is not None and api[cid] <= end: + submitted = api.get((pid >> 24, cid)) + if submitted is not None and submitted <= end: wait[g] += s - end else: host += s - end @@ -151,9 +157,9 @@ def seg(n, gx, gy): f"wait-before {wait[g] / 1e6:5.0f} n={n[g]}" ) side = collections.defaultdict(float) -for s, e, sid, cid, nm, gx, gy in ks: - if sid != main: - side[(sid, (names.get(nm) or "")[:40])] += e - s +for s, e, sid, cid, nm, gx, gy, pid in ks: + if (pid, sid) != main: + side[(pid >> 24, sid, (names.get(nm) or "")[:40])] += e - s print( "side streams:", [(k, round(v / 1e6)) for k, v in sorted(side.items(), key=lambda x: -x[1])[:5]], diff --git a/benchmarks/flashnext_prefill/quality.py b/benchmarks/flashnext_prefill/quality.py new file mode 100644 index 0000000000..52a237e907 --- /dev/null +++ b/benchmarks/flashnext_prefill/quality.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired retrieval/arithmetic and EOS checks at specified prompt-token budgets. + +Uses production sampling and thinking. Reset only an exclusively owned server. +Record the same prompt hash, seed and actual token length for each artifact. +""" + +import argparse +import hashlib +import json +import time +from collections import Counter +from pathlib import Path + +import regex as re +from measure_prefill import request + +FILLER = ( + "Computing records describe mechanical calculators, vacuum tubes, transistors, " + "integrated circuits, storage systems and communication networks. " + "These background records do not contain any vault or captain information.\n" +) +EXPECTED = {"vault": "48213", "captain": "Orlanda Vexley", "packages": 45} + + +def prompt(repeats): + return ( + "Read this archive and answer the question at the end.\n" + + FILLER * (repeats // 8) + + "The vault code is 48213.\n" + + FILLER * (repeats // 2) + + "The ship captain is Orlanda Vexley.\n" + + FILLER * (repeats - repeats // 8 - repeats // 2) + + "A warehouse has 17 boxes with 24 items each. After selling 93 items, " + "all remaining items are packed into groups of 7. How many full packages " + "are made? Return only one JSON object containing the vault code as a " + 'string, captain name, and number of packages: {"vault": ..., ' + '"captain": ..., "packages": ...}. Think briefly.\n' + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8002") + parser.add_argument("--budgets", type=int, nargs="+", default=[32768]) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if min(args.budgets) < 512: + parser.error("prompt budgets must be at least 512 tokens") + results = [] + for budget in args.budgets: + + def count(repeats): + messages = [{"role": "user", "content": prompt(repeats)}] + with request( + args.base_url, + "/tokenize", + { + "model": "flash-next", + "messages": messages, + "chat_template_kwargs": {"enable_thinking": True}, + }, + ) as response: + return json.load(response)["count"], messages + + fixed, _ = count(0) + ten, _ = count(10) + repeats = max(0, int((budget - fixed) * 10 / (ten - fixed))) + tokens, messages = count(repeats) + while tokens > budget: + repeats -= 1 + tokens, messages = count(repeats) + body = { + "model": "flash-next", + "messages": messages, + "max_tokens": 1024, + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "seed": 9683, + "chat_template_kwargs": {"enable_thinking": True}, + } + with request(args.base_url, "/reset_prefix_cache", {}) as response: + response.read() + start = time.perf_counter() + with request(args.base_url, "/v1/chat/completions", body) as response: + result = json.load(response) + choice = result["choices"][0] + message = choice["message"] + content = message.get("content") or "" + match = re.search(r"\{[^{}]*\}", content) + try: + answer = json.loads(match.group()) if match else None + except json.JSONDecodeError: + answer = None + text = " ".join( + str(message.get(k) or "") + for k in ("content", "reasoning", "reasoning_content") + ) + words = text.split() + ngrams = Counter(tuple(words[i : i + 10]) for i in range(len(words) - 9)) + max_repeat = max(ngrams.values(), default=0) + row = { + "budget_tokens": budget, + "tokenize_count": tokens, + "prompt_sha256": hashlib.sha256( + messages[0]["content"].encode() + ).hexdigest(), + "seconds": time.perf_counter() - start, + "usage": result["usage"], + "choice": choice, + "max_ten_word_repeat": max_repeat, + "passed": answer == EXPECTED + and choice["finish_reason"] == "stop" + and max_repeat <= 5, + } + results.append(row) + args.output.write_text(json.dumps(results, indent=2, ensure_ascii=False) + "\n") + print(json.dumps(row, ensure_ascii=False), flush=True) + raise SystemExit(0 if all(row["passed"] for row in results) else 1) + + +if __name__ == "__main__": + main() diff --git a/csrc/gguf_expert_cache.cpp b/csrc/gguf_expert_cache.cpp index 012e11249e..a3310326d6 100644 --- a/csrc/gguf_expert_cache.cpp +++ b/csrc/gguf_expert_cache.cpp @@ -257,6 +257,43 @@ class ExpertCacheLane { return extent; } + // Stage only host slots used by the most recent residency snapshot. Keep + // their offsets unchanged: the GPU source map still addresses those slots. + // The caller orders the copy after decode write-backs and the preceding + // cold transcode, exactly as for a full-arena copy. + void stage_cold(int64_t layer, int64_t shard, int64_t destination, + int64_t stream_ptr) { + TORCH_CHECK(layer >= 0 && layer < layers_ && shard >= 0 && shard < kShards, + "expert cache: bad cold stage bank"); + auto stream = reinterpret_cast(stream_ptr); + std::lock_guard guard(mutex_); + const int32_t* map = residency_maps_ + 2 * layer * experts_; + std::vector occupied(host_slots_, 0); + for (int64_t e = 0; e < experts_; ++e) { + if (map[e] >= 0) continue; + const int32_t h = map[experts_ + e]; + TORCH_CHECK(h >= 0 && h < host_slots_, + "expert cache: non-resident expert has no host slot"); + occupied[h] = 1; + } + check(cudaSetDevice(static_cast(device_)), "set device"); + const int64_t stride = bytes(layer, shard); + for (int64_t begin = 0; begin < host_slots_;) { + if (!occupied[begin]) { + ++begin; + continue; + } + int64_t end = begin + 1; + while (end < host_slots_ && occupied[end]) ++end; + check(cudaMemcpyAsync( + reinterpret_cast(destination) + begin * stride, + host(layer, shard, begin), (end - begin) * stride, + cudaMemcpyHostToDevice, stream), + "stage cold experts"); + begin = end; + } + } + private: int64_t bytes(int64_t layer, int shard) const { return slot_bytes_[layer * kShards + shard]; @@ -533,6 +570,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("prefill", &ExpertCacheLane::prefill, pybind11::call_guard()) .def("residency", &ExpertCacheLane::residency, + pybind11::call_guard()) + .def("stage_cold", &ExpertCacheLane::stage_cold, pybind11::call_guard()); m.def("host_register", &host_register); m.def("host_unregister", &host_unregister); diff --git a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu index f4a632a5ff..41be11393b 100644 --- a/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu +++ b/csrc/sm70_turbomind/lmdeploy/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu @@ -18,6 +18,19 @@ using D = cache_policy::Default; namespace { +// Reuse the existing decoder and indexed grouped layout, with the padded +// activation supply already used by the quantized batch family. +template +struct PaddedActivationConfig; + +template +struct PaddedActivationConfig> { + using type = Sm70_s884, TransformA, U, B, + TransformB, V, order_C, Tc, raster_order, group_axis>; +}; + // Keep shape-specific small-N tactics out of prefill and unrelated CUDA graph // shapes. template @@ -124,6 +137,31 @@ void Registry::sm70_884_4() { 1, Group, 64, 64>; Add(std::make_unique< GgufPrefillNarrowKernelImpl>()); + // Ragged expert batches often have fewer than 128 rows. Halving M + // reduces both masked MMA work and the accumulator register demand; + // retain the existing quantization-group-sized K iteration. + using Prefill64x64 = + typename G::template Type<64, 64, Group, 2, 1, 1, D, D, 2, true, + 1, Group, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + using P = typename PaddedActivationConfig::type; + // K32/K64 unpadded activation rows repeat shared-memory banks across + // the four distinct rows read by each eight-lane MMA load group. + // Eight half values of row padding spread those accesses apart. + // A full-width epilogue is required: N64 subtiles of an N128 CTA + // are smaller than the epilogue thread map's horizontal span. + // K64 also distinguishes these tactics in serialized tuning caches. + using Padded32x64K64 = + typename P::template Type<32, 64, 64, 2, 2, 1, D, D, 2, true, + 1, Group, 32, 64>; + using Padded64x64K64 = + typename P::template Type<64, 64, 64, 2, 2, 1, D, D, 2, true, + 1, Group, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); } Add>(); @@ -159,6 +197,22 @@ void Registry::sm70_884_4() { GroupSize, 64, 64>; Add(std::make_unique< GgufPrefillNarrowKernelImpl>()); + using Prefill64x64 = + typename G::template Type<64, 64, 32, 2, 1, 1, D, D, 2, true, 1, + GroupSize, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + using P = typename PaddedActivationConfig::type; + using Padded32x64K64 = + typename P::template Type<32, 64, 64, 2, 2, 1, D, D, 2, true, + 1, GroupSize, 32, 64>; + using Padded64x64K64 = + typename P::template Type<64, 64, 64, 2, 2, 1, D, D, 2, true, + 1, GroupSize, 64, 64>; + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); + Add(std::make_unique< + GgufPrefillNarrowKernelImpl>()); } Add>(); diff --git a/tests/kernels/quantization/test_gguf_grouped_prefill.py b/tests/kernels/quantization/test_gguf_grouped_prefill.py new file mode 100644 index 0000000000..778c5b38a5 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_grouped_prefill.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Narrow indexed GGUF prefill: padding, empty groups and graph replay.""" + +import gguf +import numpy as np +import pytest +import torch +from test_gguf_lattice_transcode import source as lattice_source +from test_gguf_lut_transcode import source as lut_source + +from vllm import _custom_ops # noqa: F401 +from vllm.model_executor.layers.quantization.gguf_expert_cache import _canonical + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("weight_type", [18, 21, 22, 23]) +@pytest.mark.parametrize( + "m,n,k,gated,indexed", + [ + (511, 320, 256, True, True), # below the new tile admission + (512, 64, 512, False, False), + (513, 320, 768, True, True), + (2049, 512, 2560, False, True), + (769, 352, 512, True, True), # N32 padding: existing fallback + (513, 576, 256, False, False), # wide projection: existing fallback + ], +) +def test_narrow_grouped_prefill(weight_type, m, n, k, gated, indexed): + experts = 5 + rows = n // 2 if gated else n + parts = [] + dense = [] + for part in range(2 if gated else 1): + source = lut_source if weight_type == 23 else lattice_source + data = np.stack( + [ + source( + weight_type, + scale=0.0009765625, + n=rows, + k=k, + seed=9683 + part * experts + expert, + ) + for expert in range(experts) + ] + ) + parts.append(torch.from_numpy(data).cuda()) + dense.append( + torch.from_numpy( + gguf.quants.dequantize(data, gguf.GGMLQuantizationType(weight_type)) + ).cuda() + ) + family, decoder, group, weight_bytes, stats_bytes = _canonical(weight_type, n, k) + weights = torch.empty(experts * weight_bytes, dtype=torch.uint8, device="cuda") + stats = torch.empty(experts * stats_bytes, dtype=torch.uint8, device="cuda") + sequence = torch.arange(experts, device="cuda", dtype=torch.int64) + pointers = [p.data_ptr() + sequence * p[0].numel() for p in parts] + error = torch.zeros(1, dtype=torch.int32, device="cuda") + torch.ops._C.gguf_prefill_transcode_sm70( + weights, + stats, + weight_bytes, + stats_bytes, + pointers[0], + pointers[1] if gated else None, + sequence.int(), + torch.tensor([experts], dtype=torch.int32, device="cuda"), + 128, + weight_type, + n, + k, + 0, + parts[0].shape[-1], + error, + ) + assert error.item() == 0 + tables = [] + for bank, stride, ld in ((weights, weight_bytes, 32 * k), (stats, stats_bytes, n)): + table = torch.empty(experts, 2, dtype=torch.int64, device="cuda") + table[:, 0] = bank.data_ptr() + sequence * stride + table[:, 1] = ld + tables.append(table.view(torch.uint8).flatten()) + bounds = [0, 0, 17, m - 1, m, m] + offsets = torch.tensor(bounds, dtype=torch.int32, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(9683 + m) + x = ( + torch.randn(m + 7 if indexed else m, k, device="cuda", generator=gen) * 0.125 + ).half() + indices = ( + torch.randint(x.shape[0], (m,), device="cuda", dtype=torch.int32, generator=gen) + if indexed + else None + ) + storage = torch.full((m * rows + 32,), 37, dtype=torch.half, device="cuda") + out = storage[16:-16].view(m, rows) + + def run(tune=False): + torch.ops._C.gguf_grouped_prefill_gemm_sm70_out( + out, + x, + offsets, + *tables, + family, + decoder, + n, + experts, + group, + gated, + tune, + indices, + ) + + def check(): + actual_x = x if indices is None else x[indices.long()] + expected = torch.empty_like(out, dtype=torch.float32) + for expert, (begin, end) in enumerate(zip(bounds, bounds[1:])): + a = actual_x[begin:end].float() + projected = a @ dense[0][expert].float().T + if gated: + projected = torch.nn.functional.silu(projected) * ( + a @ dense[1][expert].float().T + ) + expected[begin:end] = projected + assert torch.isfinite(expected.half()).all(), "reference exceeds FP16 range" + finite = torch.isfinite(out) + assert finite.all(), { + "bad_indices": (~finite).nonzero()[:8].tolist(), + "expected_abs_max": expected.abs().max().item(), + } + relative = (out.float() - expected).norm(dim=1) / expected.norm( + dim=1 + ).clamp_min(1e-6) + assert relative.max().item() < 0.005 + assert (storage[:16] == 37).all() and (storage[-16:] == 37).all() + + run() # no pre-existing native tuning entry + check() + run(tune=True) + run() + check() + if m == 513: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + # Address-stable replay with different activations and gather indices. + x.mul_(-0.75) + if indices is not None: + indices.copy_(torch.flip(indices, (0,))) + graph.replay() + check() diff --git a/tests/kernels/quantization/test_gguf_turbomind_dequant.py b/tests/kernels/quantization/test_gguf_turbomind_dequant.py index b2db6b1f52..1dff641d1c 100644 --- a/tests/kernels/quantization/test_gguf_turbomind_dequant.py +++ b/tests/kernels/quantization/test_gguf_turbomind_dequant.py @@ -108,18 +108,30 @@ def run(x=x, out=out): ) -def test_calibrated_prefill_framework_routes_and_shared_workspace(): +@pytest.mark.parametrize( + "bits,group,k,n", + [ + (5, 32, 5120, 1536), + (5, 32, 2560, 2560), + (6, 16, 2560, 2560), + (6, 16, 2560, 1536), + (6, 16, 1536, 2560), + ], +) +def test_calibrated_prefill_framework_routes_and_shared_workspace(bits, group, k, n): + from tools.sm70.native_trace import NativeDispatchTrace + from vllm.model_executor.layers.quantization import gguf_turbomind # noqa: F401 + torch._dynamo.reset() - k, n = 5120, 1536 config = Sm70GgufAffineConfig( (k, n), (k, n), - scalar_types.uint5, + scalar_types.uint5 if bits == 5 else scalar_types.uint6, torch.float16, - 32, + group, True, False, - source_type=13, + source_type=13 if bits == 5 else 14, ) layers, kernels = [], [] for _ in range(2): @@ -127,9 +139,9 @@ def test_calibrated_prefill_framework_routes_and_shared_workspace(): kernel = TurboMindGgufAffineKernel(config, "codes", "scales", "mins") for name, data in { "codes": torch.zeros((n, k), dtype=torch.uint8, device="cuda"), - "scales": torch.ones((n, k // 32), dtype=torch.float16, device="cuda"), + "scales": torch.ones((n, k // group), dtype=torch.float16, device="cuda"), "mins": torch.full( - (n, k // 32), 0.00390625, dtype=torch.float16, device="cuda" + (n, k // group), 0.00390625, dtype=torch.float16, device="cuda" ), }.items(): layer.register_parameter( @@ -144,11 +156,19 @@ def test_calibrated_prefill_framework_routes_and_shared_workspace(): layers[0].gguf_tm_blas_workspace.data_ptr() == layers[1].gguf_tm_blas_workspace.data_ptr() ) - for m in (512, 2048): + for m in (512, 2048) if k == 5120 else (2047, 2048, 2049, 4097): x = (torch.randn((m, k), device="cuda") * 0.125).half() expected = (x.float().sum(-1, keepdim=True) * 0.00390625).expand(m, n) run = lambda x: kernels[0].apply_weights(layers[0], x) - torch.testing.assert_close(run(x).float(), expected, rtol=0.003, atol=0.0002) + with NativeDispatchTrace() as trace: + eager = run(x) + operator = "blas" if m >= 2048 else "gemm" + assert any( + r["operator"] == f"_C::gguf_affine_{operator}_sm70_out" + and r["evidence"] == "cuda_dispatch" + for r in trace.records + ) + torch.testing.assert_close(eager.float(), expected, rtol=0.003, atol=0.0002) compiled = torch.compile(run, backend="eager", fullgraph=True) torch.testing.assert_close( compiled(x).float(), expected, rtol=0.003, atol=0.0002 @@ -163,3 +183,25 @@ def test_calibrated_prefill_framework_routes_and_shared_workspace(): out = run(x) graph.replay() torch.testing.assert_close(out.float(), expected, rtol=0.003, atol=0.0002) + # The serving graph passes the calibrated bands through its opaque op. + layer = layers[0] + opaque = torch.ops.vllm.prepared_gguf_projection( + x, + layer.codes, + layer.scales, + None, + 0, + bits, + group, + layer.gguf_tm_k_ld, + layer.gguf_tm_q_ld, + n, + n, + [], + [2048, -1], + ) + torch.testing.assert_close(opaque.float(), expected, rtol=0.003, atol=0.0002) + x.mul_(-0.5) + graph.replay() + expected = (x.float().sum(-1, keepdim=True) * 0.00390625).expand(m, n) + torch.testing.assert_close(out.float(), expected, rtol=0.003, atol=0.0002) diff --git a/tests/models/qwen4_exp/test_gguf_expert_cache.py b/tests/models/qwen4_exp/test_gguf_expert_cache.py index 4d36786136..80f3a76b81 100644 --- a/tests/models/qwen4_exp/test_gguf_expert_cache.py +++ b/tests/models/qwen4_exp/test_gguf_expert_cache.py @@ -139,6 +139,46 @@ def test_prefill_bank_matches_and_interleaves_with_decode(): assert cache.stats()["errors"] == 0 +@pytest.mark.parametrize("slack", [0.0, 0.25]) +def test_cold_stage_uses_snapshot_and_preserves_slot_offsets(slack): + cache, original = _cache(slack=slack) + generator = torch.Generator().manual_seed(9683) + try: + # Fill holes and exercise write-backs before taking the prefill snapshot. + for _ in range(10): + for layer in LAYERS: + cache.resolve(layer, _random_ids(generator, 5)) + torch.accelerator.synchronize() + stream = torch.cuda.current_stream().cuda_stream + for i, layer in enumerate(LAYERS): + snapshot = torch.empty(2, EXPERTS, dtype=torch.int32, device="cuda") + cache.lane.residency(i, snapshot.data_ptr(), stream) + torch.accelerator.synchronize() + gpu_slots, host_slots = snapshot.cpu() + cold = torch.nonzero(gpu_slots < 0).flatten() + selected = host_slots[cold].long() + # Source selection must use the saved snapshot, not a later CPU map. + saved = cache.slot_of[i].clone() + cache.slot_of[i].fill_(0) + try: + for s, shard in enumerate(gguf_expert_cache.SHARDS): + stage = torch.full_like( + cache.host[layer, shard], 0xA5, device="cuda" + ) + cache.lane.stage_cold(i, s, stage.data_ptr(), stream) + torch.testing.assert_close( + stage[selected], original[layer, shard][cold], rtol=0, atol=0 + ) + holes = torch.ones(cache.host_slots, dtype=torch.bool) + holes[selected] = False + assert (stage[holes] == 0xA5).all() + finally: + cache.slot_of[i].copy_(saved) + finally: + torch.accelerator.synchronize() + cache.close() + + HIDDEN, INTER, TOP_K = 256, 32, 10 diff --git a/vllm/model_executor/kernels/linear/mixed_precision/sm70_gguf.py b/vllm/model_executor/kernels/linear/mixed_precision/sm70_gguf.py index 9e45329642..55c7f045a4 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/sm70_gguf.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/sm70_gguf.py @@ -44,6 +44,12 @@ (5, 32, 5120, 2560): 2048, (6, 16, 6144, 5120): 2048, (6, 16, 1536, 5120): 2048, + # Flash-Next TP4 Q5_K/Q6_K projections; dequantization is included in + # the paired timings in benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md. + (5, 32, 2560, 2560): 2048, + (6, 16, 2560, 2560): 2048, + (6, 16, 2560, 1536): 2048, + (6, 16, 1536, 2560): 2048, } _AFFINE_BLAS_WORKSPACE_ELEMENTS = max(k * n for _, _, k, n in _AFFINE_BLAS_MIN_M) _affine_blas_workspaces: weakref.WeakValueDictionary = weakref.WeakValueDictionary() diff --git a/vllm/model_executor/layers/quantization/gguf_expert_cache.py b/vllm/model_executor/layers/quantization/gguf_expert_cache.py index b274f9bdf5..30e789b140 100644 --- a/vllm/model_executor/layers/quantization/gguf_expert_cache.py +++ b/vllm/model_executor/layers/quantization/gguf_expert_cache.py @@ -1014,17 +1014,8 @@ def _issue(self, key: str, i: int): # (decode write-backs precede it) and the previous layer's cold # transcode (stream order), not the bank: it runs ahead. cold.wait_event(self.fresh) - rows = self.extent[p] for s in (0, 1) if key == "gu" else (2,): - shard = SHARDS[s] - nbytes = rows * cache.sizes[layer, shard] - if nbytes: - _native().copy_h2d( - self.stage_at[s], - cache.host[layer, shard].data_ptr(), - nbytes, - cold.cuda_stream, - ) + cache.lane.stage_cold(i, s, self.stage_at[s], cold.cuda_stream) cold.wait_stream(stream) spec = info[key] weight, stats = self.bank(key, spec, slot[1]) From 3faaa58d084c1467dc530a2559636b116b64b41b Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 04:33:48 +0800 Subject: [PATCH 03/19] [Bugfix] Preserve GGUF prefill policy with segment storage Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../OPTIMIZATION_20261014.md | 135 +++++++++- .../flashnext_prefill/dense_segments.py | 103 ++++++++ benchmarks/flashnext_prefill/moe_blas.py | 237 ++++++++++++++++++ .../quantization/test_gguf_segment_prefill.py | 120 +++++++++ .../layers/quantization/gguf_dense_hmma.py | 48 +++- .../layers/quantization/gguf_turbomind.py | 43 +++- 6 files changed, 663 insertions(+), 23 deletions(-) create mode 100644 benchmarks/flashnext_prefill/dense_segments.py create mode 100644 benchmarks/flashnext_prefill/moe_blas.py create mode 100644 tests/kernels/quantization/test_gguf_segment_prefill.py diff --git a/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md b/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md index bc189ce3ae..975939075a 100644 --- a/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md +++ b/benchmarks/flashnext_prefill/OPTIMIZATION_20261014.md @@ -343,6 +343,66 @@ prefill/decode, fragmented/slack-free host arenas, untouched staging holes, and using the saved snapshot after the live CPU residency table changes. Full-model latency remains a separate acceptance gate. +### 9. Reject padded expert-batch FP16 BLAS + +`moe_blas.py` screens real layer-zero TP4 weights at 7344 tokens, top-10, +512 experts and 192 padded rows/expert. It includes per-batch dequantization, +padded activation gather, gate/up, SiLU, down, scatter and weighted unroute. +Contiguous compressed GPU banks are an optimistic input layout; production +would also need to gather resident/cold slots. Both arms omit routing and DMA. + +The current quantized chain takes 6.414 ms (repeat 6.478). Batches of 32, 64, +and 128 experts take 19.860, 17.301 and 16.972 ms respectively. All outputs +are finite with paired maximum row-relative L2 0.001142, within the preset +0.003 bound, but performance rejects this implementation. For 64 experts, +gate/up and down GEMMs alone take 0.688 and 0.255 ms, or about 7.54 ms for +all eight batches before dequantization/gather/scatter. FP16 gate/up output +only reduces its 0.688 ms to 0.661; it cannot close the gap. No runtime route +is added. Raw artifacts: `moe-blas.json` and `moe-blas-first.json`. + +### 10. Restore the calibrated policy after segment-bank restoration + +The matched serving trace exposed a missing integration path: no affine BLAS +or affine dequantization kernels execute, despite the four admitted physical +layouts. With `small_m_hmma=true`, initialization replaces canonical weights +with segment storage. Its prefill fallback restores canonical weights but +always invokes fused GEMM, and serialization discards the BLAS bands. +Thus the first 7.099 s result below does not demonstrate dense BLAS execution. +The earlier direct-kernel and ordinary opaque-path tests did not cover this +production storage mode. + +The follow-up carries the existing admitted bands through single and mixed +segment wrappers, retaining the small-M route. After canonical restoration, +it reuses the existing affine BLAS operator and engine-owned scratch only +inside those bands; all other inputs retain fused GEMM. Mixed output views +use a contiguous temporary because the existing native BLAS API requires it. +No native operator, configuration flag or extra persistent workspace is added. +`dense_segments.py` measures this actual restoration chain on real weights; +`test_gguf_segment_prefill.py` covers threshold neighbors, a range interior, +return to decode, compiled opaque calls, changed-input graph replay, mixed +output views and an uncalibrated fallback. **44 GPU/dispatch tests pass**, +including the existing segment codec, restoration, small-M graph and runtime +opacity regressions. Actual native BLAS dispatch is observed in the restored +chain; opaque and compiled surfaces produce the same qualified outputs. + +Real-weight restored-chain medians (ms), including canonical restoration and +weight dequantization on every call: + +| Projection / source | M | Restored fused | Restored BLAS | Fused repeat | +|---|---:|---:|---:|---:| +| blk.0.attn_qkv.weight / 14 | 7344 | 1.959 | 1.309 | 1.938 | +| blk.0.attn_qkv.weight / 14 | 5347 | 1.464 | 1.016 | 1.398 | +| blk.1.attn_qkv.weight / 13 | 7344 | 1.991 | 1.229 | 1.918 | +| blk.1.attn_qkv.weight / 13 | 5347 | 1.440 | 0.980 | 1.475 | +| blk.0.ssm_out.weight / 14 | 7344 | 1.148 | 0.824 | 1.144 | +| blk.0.ssm_out.weight / 14 | 5347 | 0.882 | 0.610 | 0.882 | +| blk.1.attn_gate.weight / 14 | 7344 | 2.077 | 0.814 | 2.038 | +| blk.1.attn_gate.weight / 14 | 5347 | 1.061 | 0.777 | 1.063 | + +All eight cases are finite and meet the preset 0.003 maximum row-relative +L2 bounds against both the FP32 oracle and the paired fused output. Full-model +route observation, quality and unprofiled speed with this fix remain pending. + ### Matched request measurement `measure_prefill.py` uses deterministic word sequences and per-request hashes, @@ -378,7 +438,8 @@ busy but around 85–90 W). No long-context quality result or candidate regression is inferred. Complete this baseline/candidate boundary check before promotion; the configured limit and deployment remain unchanged. -First combined candidate, same five prompt hashes and token counts as control: +First candidate (before the segment-policy fix), same five prompt hashes and +token counts as control: | Prompt tokens | Client first text s | Server prefill s | Server TTFT s | |---:|---:|---:|---:| @@ -402,6 +463,70 @@ occurred. The owned candidate was stopped after these checks for a bounded expert-batch FP16 MoE screen and a matched Nsight profile. Its descendants were recorded and only its verified API process received SIGTERM. +### Matched candidate trace and boundary quality + +The first-candidate profile uses the same index-2 request as the unprofiled +runs: 34688 tokens, hash +`126bc5319be9170da4f07bd2bca263e083d2d636da17d0697d7a19d746b04342`. +Client first text is 7.312 s, server prefill 6.952 s and server TTFT 7.209 s. +The four main-stream spans are 6.914 s. This capture also enables NCCL +INIT/TUNING logging, so its wall time is not an unprofiled performance result. +The profiled prefill chunk row counts are 7344, 6528 and 6128; do not assume +that every request uses the historical 7344/5347 microbenchmark shapes. + +GPU0 closes as 6045.9 ms main-stream kernel/copy union, approximately 793 ms submitted +GPU wait and 75 ms host gap. The 397.9 ms outside the stream span remains +an unattributed HTTP/server residual; it is not added to kernel categories. +Main-stream categories from `nsys_breakdown.py`, in ms: + +| Work | Kernel/copy time | +|---|---:| +| TurboMind dense projections | 987 | +| MoE gate/up + down | 662 + 501 | +| MoE unroute + routing | 122 + 68 | +| QSA sparse partial + other + merge | 624 + 354 + 12 | +| AllReduce + AllGather | 501 + 381 | +| HC combine/norm + other HC | 348 + 125 | +| Other cuBLAS, including HC down | 589 | +| GDN | 350 | +| Elementwise/other, MTP, memcpy, dense helpers | 258 + 66 + 55 + 42 | + +The HC down cuBLAS kernel alone is 272.2 ms / 500 launches. GPU0 union +coverage is 87.44% main-stream kernel/copy, 89.88% all kernels, and 97.88% +all kernels/copies. These are timeline coverage, not SM, Tensor Core or HBM +utilization counters. NCU counters remain unavailable. Cold-stream H2D is +13.620 GB/rank over 3.553–3.766 s, about 3.62–3.83 GB/s. It overlaps +computation and must not be added to the wall-clock table. + +Across ranks, 500 main FP16 all-reduces spend 434.5 ms after the last rank's +arrival; arrival spread sums to 245.1 ms, with GPU0 arriving last 366 times. +The 1030 all-gathers have 335.9 ms after the last arrival and 133.3 ms spread. +NCCL 2.27.5 TUNING logs explicitly select RING/SIMPLE for the prefill payloads +(37,601,280-byte reduce; 1,292,544/9,400,320-byte per-rank gather inputs). +The trace kernel name ends in RING_LL even for these calls: that name alone +is not evidence of actual protocol. This explains why forcing Simple in the +previous handoff gave no benefit; no protocol override is promoted. + +The dominant grouped down kernel uses 255 registers/thread, 128 threads +and 32784 bytes shared memory per CTA. Register capacity bounds it to two +resident CTAs/SM, or eight warps (12.5% of 64); this is a resource ceiling, +not measured achieved occupancy. It motivates a bounded short-K tile screen +inside the existing provider, with smaller row tiles and the existing padded +activation loader. Its accelerated subset is grouped M>=512 routed rows, +N=2560 and K=160, for IQ4/group32 and affine U2/group32; other shapes retain +the existing tuner candidates. New GPU correctness and timing are pending. + +The first candidate's 260089-token quality prompt passes retrieval, arithmetic, +normal EOS and repetition checks (181 output tokens; maximum repeated ten-word +span 1). Hash: +`2c78b9e46e325dfd708b4d8c8f8c7c8fa2b9696bcb3d54f2a62313c70fb76665`. +The full response takes 821.64 s, under the existing 131072-token staging limit +and direct-host fallback beyond it. This is a boundary quality pass, not a +paired speed comparison; matching control is still pending. The owned API, +workers and Nsight launcher are stopped, with process ids archived in +`profile-stop-tree.json`. The remote trace files were removed only after +matching local SHA256 checks; local raw artifacts are `candidate-matched.*`. + ## Execution discipline Microbenchmarks acquire the existing GPU locks and start only on an empty GPU. @@ -436,7 +561,7 @@ The effective unchanged residency log is 402/512 experts plus eight spare GPU slots, 136 host slots, 9.46 GiB GPU expert slots and 3.14 GiB pinned host per rank, matching the historical deployment with configured fraction 0.8. -Pending: candidate model quality, matched unprofiled end-to-end measurement, -and paired near-256K qualification. The original 256K deployment has not been -promoted to a candidate. The isolated candidate server is starting with the -corrected MoE tiles, four dense calibrations, and occupied-slot staging. +Draft PR: . The first candidate +passed matched 32K quality and unprofiled performance checks; 6K tokens/s is +not achieved. Paired 128K/near-256K qualification and the newly identified +segment-policy fix remain pending. The original deployment is unchanged. diff --git a/benchmarks/flashnext_prefill/dense_segments.py b/benchmarks/flashnext_prefill/dense_segments.py new file mode 100644 index 0000000000..471735bd3f --- /dev/null +++ b/benchmarks/flashnext_prefill/dense_segments.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Measure prefill after the serving segment-bank restoration, including DQ. + +Both arms keep small-M segment storage. The control omits only the admitted +BLAS bands, reproducing the missed serving route without another checkpoint. +""" + +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import torch +from dense_blas import CASES +from moe_tiles import artifact_identity, measure, relative_l2 + +from tools.sm70.native_trace import NativeDispatchTrace +from vllm.config import set_current_vllm_config +from vllm.config.kernel import KernelConfig +from vllm.model_executor.layers.quantization.gguf_turbomind import ( + _prepared_gguf_mixed_projection, + prepare_gguf_projections, + prepared_projection_arguments, +) +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, dequantize + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + config = SimpleNamespace( + kernel_config=KernelConfig(), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + report = {"artifact": artifact_identity(), "cases": []} + reader = GGUFReader(args.model) + with set_current_vllm_config(config): + for name, axis in CASES[:4]: + tensor = next(t for t in reader.tensors if t.name == name) + raw = np.asarray(tensor.data) + raw = raw[: raw.shape[0] // 4] if axis == 0 else raw[:, : raw.shape[1] // 4] + raw = np.ascontiguousarray(raw) + source = int(tensor.tensor_type) + reference_weight = torch.from_numpy(dequantize(raw, source).copy()).cuda() + projection = prepare_gguf_projections( + [(torch.from_numpy(raw).cuda(), source)], torch.float16, True, 8 + )[0] + assert hasattr(projection, "segment_format") + actual = prepared_projection_arguments([projection]) + assert actual[3][0] == 3 and actual[5] == [2048, -1] + control = list(actual) + control[3] = actual[3][:-1] + [0] + control[5] = [] + k, n = projection.kernel.config.partition_weight_shape + for m in (7344, 5347): + torch.manual_seed(9683 + m) + x = (torch.randn(m, k, device="cuda") * 0.125).half() + + def fused(x=x, control=control): + return torch.ops.vllm.prepared_gguf_mixed_projection(x, *control) + + def blas(x=x, actual=actual): + return torch.ops.vllm.prepared_gguf_mixed_projection(x, *actual) + + old = fused() + with NativeDispatchTrace() as trace: + observed = _prepared_gguf_mixed_projection(x, *actual) + new = blas() + torch.testing.assert_close(new, observed, rtol=0, atol=0) + assert any( + r["operator"] == "_C::gguf_affine_blas_sm70_out" + and r["evidence"] == "cuda_dispatch" + for r in trace.records + ) + reference = x[:16].float() @ reference_weight.float().T + errors = { + "control_fp32_max_rel_l2": relative_l2(old[:16], reference), + "candidate_fp32_max_rel_l2": relative_l2(new[:16], reference), + "paired_max_rel_l2": relative_l2(new, old), + } + row = { + "tensor": name, + "source_type": source, + "shape_mnk": [m, n, k], + "restore_fused": measure(fused, 5), + "restore_blas": measure(blas, 5), + "restore_fused_repeat": measure(fused, 5), + **errors, + "passed": bool(torch.isfinite(new).all()) + and max(errors.values()) <= 0.003, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise SystemExit(0 if all(r["passed"] for r in report["cases"]) else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/moe_blas.py b/benchmarks/flashnext_prefill/moe_blas.py new file mode 100644 index 0000000000..09b12d221e --- /dev/null +++ b/benchmarks/flashnext_prefill/moe_blas.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen bounded expert-batch FP16 GEMMs against the quantized MoE chain. + +Real TP4 weights and routes; both arms exclude routing and host DMA. The BLAS +arm includes dequantization, padded input gather, scatter and weighted unroute. +Its contiguous raw GPU banks are an optimistic source layout: production's +resident/cold slots would need additional gathering. This is not a new backend. +""" + +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +import torch +from moe_tiles import ( + EXPERTS, + HIDDEN, + INTERMEDIATE, + TOPK, + artifact_identity, + load_layer, + measure, + relative_l2, +) + +from vllm.model_executor.layers.quantization.gguf import dequantize_gguf_rows +from vllm.model_executor.layers.quantization.gguf_expert_cache import ( + TurboMindPrefill, + _unroute, +) +from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, quant_size +from vllm.triton_utils import tl, triton + + +@triton.jit +def _pad_gather( + x, + rows, + offsets, + padded, + FIRST: tl.constexpr, + PAD: tl.constexpr, + K: tl.constexpr, + BLOCK: tl.constexpr, +): + expert = tl.program_id(0) + pos = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + r, col = pos // K, pos % K + begin = tl.load(offsets + FIRST + expert) + end = tl.load(offsets + FIRST + expert + 1) + token = tl.load(rows + begin + r, (r < PAD) & (begin + r < end), 0) + value = tl.load(x + token * K + col, (r < PAD) & (begin + r < end), 0) + tl.store(padded + expert * PAD * K + pos, value, r < PAD) + + +@triton.jit +def _scatter( + padded, + offsets, + out, + FIRST: tl.constexpr, + PAD: tl.constexpr, + K: tl.constexpr, + BLOCK: tl.constexpr, +): + expert = tl.program_id(0) + pos = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + r, col = pos // K, pos % K + begin = tl.load(offsets + FIRST + expert) + end = tl.load(offsets + FIRST + expert + 1) + valid = (r < PAD) & (begin + r < end) + value = tl.load(padded + expert * PAD * K + pos, valid, 0) + tl.store(out + (begin + r) * K + col, value, valid) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model") + parser.add_argument("--layer", type=int, default=0) + parser.add_argument("--tokens", type=int, default=7344) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + info, tables, banks, raw, left = load_layer(GGUFReader(args.model), args.layer) + runner = TurboMindPrefill.__new__(TurboMindPrefill) + runner.cache = SimpleNamespace(experts=EXPERTS, prefill_tokens=4096) + runner.expert_ids = torch.arange(EXPERTS + 1, dtype=torch.int32, device="cuda") + runner.tuned = set() + runner._tune( + args.layer, + info, + tables, + torch.zeros(8192, HIDDEN, device="cuda", dtype=torch.half), + ) + torch.manual_seed(args.tokens) + x = (torch.randn(args.tokens, HIDDEN, device="cuda") * 0.5).half() + scores, ids = torch.randn(args.tokens, EXPERTS, device="cuda").topk(TOPK, dim=1) + weights = scores.softmax(dim=1) + route = runner._route(ids, 0, EXPERTS) + offsets, _, rows, _, inverse, _ = route + pad = triton.cdiv(int((offsets[1:] - offsets[:-1]).max()), 32) * 32 + subroutes = [ + ( + t, + min(t + 4096, args.tokens), + runner._subroute(route, t, min(t + 4096, args.tokens)), + ) + for t in range(0, args.tokens, 4096) + ] + control = torch.empty_like(x) + + def baseline(): + hidden = runner._gate_up(info, tables, x, route, False) + for first, last, sub in subroutes: + runner._down( + info, + tables, + hidden, + sub, + weights[first:last], + control[first:last], + False, + ) + + baseline() + raw_gu = torch.cat([torch.from_numpy(raw[p]).cuda() for p in ("gate", "up")], 1) + raw_down = torch.from_numpy(raw["down"]).cuda() + block, size = quant_size(info["down"][0]) + down_k = raw_down.shape[-1] // size * block + routed = torch.empty(args.tokens * TOPK, HIDDEN, device="cuda", dtype=torch.half) + output = torch.empty_like(x) + report = { + "artifact": artifact_identity(), + "layer": args.layer, + "tokens": args.tokens, + "padded_rows_per_expert": pad, + "control": measure(baseline, 3), + "cases": [], + } + for count in (32, 64, 128): + padded = torch.empty(count, pad, HIDDEN, device="cuda", dtype=torch.half) + + def candidate(count=count, padded=padded): + for first in range(0, EXPERTS, count): + gu = dequantize_gguf_rows( + raw_gu[first : first + count].flatten(0, 1), + info["gate"][0], + HIDDEN, + torch.half, + ).view(count, 2 * INTERMEDIATE, HIDDEN) + dn = dequantize_gguf_rows( + raw_down[first : first + count].flatten(0, 1), + info["down"][0], + down_k, + torch.half, + ).view(count, HIDDEN, down_k) + dn = dn[:, :, left : left + INTERMEDIATE] + grid = (count, triton.cdiv(pad * HIDDEN, 2048)) + _pad_gather[grid](x, rows, offsets, padded, first, pad, HIDDEN, 2048) + projected = torch.bmm( + padded, gu.transpose(1, 2), out_dtype=torch.float32 + ) + hidden = ( + torch.nn.functional.silu(projected[:, :, :INTERMEDIATE]) + * projected[:, :, INTERMEDIATE:] + ).half() + down = torch.bmm(hidden, dn.transpose(1, 2)) + _scatter[grid](down, offsets, routed, first, pad, HIDDEN, 2048) + _unroute[(args.tokens, triton.cdiv(HIDDEN, 1024))]( + routed, weights, inverse, output, HIDDEN, TOPK=TOPK, BLOCK=1024 + ) + + candidate() + error = relative_l2(output, control) + finite = bool(torch.isfinite(output).all()) + row = { + "experts_per_batch": count, + **measure(candidate, 3), + "control_max_rel_l2": error, + "finite": finite, + "passed": finite and error <= 0.003, + } + report["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + report["control_repeat"] = measure(baseline, 3) + # Localize a rejected whole-chain screen before considering native work. + count = 64 + packed_gu = raw_gu[:count].flatten(0, 1) + packed_dn = raw_down[:count].flatten(0, 1) + + def dequant_gu(): + return dequantize_gguf_rows( + packed_gu, info["gate"][0], HIDDEN, torch.half + ).view(count, 2 * INTERMEDIATE, HIDDEN) + + def dequant_dn(): + return dequantize_gguf_rows( + packed_dn, info["down"][0], down_k, torch.half + ).view(count, HIDDEN, down_k)[:, :, left : left + INTERMEDIATE] + + gu, dn = dequant_gu(), dequant_dn() + padded = torch.empty(count, pad, HIDDEN, device="cuda", dtype=torch.half) + grid = (count, triton.cdiv(pad * HIDDEN, 2048)) + + def gather(): + _pad_gather[grid](x, rows, offsets, padded, 0, pad, HIDDEN, 2048) + + gather() + projected = torch.bmm(padded, gu.transpose(1, 2), out_dtype=torch.float32) + hidden = ( + torch.nn.functional.silu(projected[:, :, :INTERMEDIATE]) + * projected[:, :, INTERMEDIATE:] + ).half() + down = torch.bmm(hidden, dn.transpose(1, 2)) + parts = { + "dequant_gate_up": dequant_gu, + "dequant_down": dequant_dn, + "gather": gather, + "gate_up_bmm_fp32_output": lambda: torch.bmm( + padded, gu.transpose(1, 2), out_dtype=torch.float32 + ), + "gate_up_bmm_fp16_output": lambda: torch.bmm(padded, gu.transpose(1, 2)), + "down_bmm": lambda: torch.bmm(hidden, dn.transpose(1, 2)), + "scatter": lambda: _scatter[grid](down, offsets, routed, 0, pad, HIDDEN, 2048), + } + report["first_64_expert_components"] = { + name: measure(fn, 20) for name, fn in parts.items() + } + print(json.dumps(report["first_64_expert_components"]), flush=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + raise SystemExit(0 if all(row["passed"] for row in report["cases"]) else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/kernels/quantization/test_gguf_segment_prefill.py b/tests/kernels/quantization/test_gguf_segment_prefill.py new file mode 100644 index 0000000000..77f44f6f3f --- /dev/null +++ b/tests/kernels/quantization/test_gguf_segment_prefill.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Segment storage retains calibrated prefill policy across runtime M changes.""" + +from types import SimpleNamespace + +import gguf +import pytest +import torch + +from tests.kernels.quantization.test_gguf_dense_segments import raw_weights +from tools.sm70.native_trace import NativeDispatchTrace +from vllm.config import set_current_vllm_config +from vllm.config.kernel import KernelConfig +from vllm.model_executor.layers.quantization.gguf_turbomind import ( + _prepared_gguf_mixed_projection, + apply_prepared_gguf_projections, + prepare_gguf_projections, + prepared_projection_arguments, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize( + "source,k,n", + [(13, 2560, 2560), (14, 2560, 2560), (14, 2560, 1536), (14, 1536, 2560)], +) +def test_segment_storage_prefill_and_decode(source, k, n): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + torch._dynamo.reset() + config = SimpleNamespace( + kernel_config=KernelConfig(), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + raw = raw_weights(source, n, k) + reference_weight = torch.from_numpy( + gguf.quants.dequantize(raw, gguf.GGMLQuantizationType(source)) + ).cuda() + with set_current_vllm_config(config): + projections = prepare_gguf_projections( + [(torch.from_numpy(raw).cuda(), source)], torch.float16, True, 8 + ) + projection = projections[0] + assert hasattr(projection, "segment_format") + serialized = prepared_projection_arguments(projections) + assert serialized[3][0] == 3 and serialized[5] == [2048, -1] + + def run(x): + return apply_prepared_gguf_projections(x, projections) + + compiled = torch.compile(run, backend="eager", fullgraph=True, dynamic=True) + for m in (2047, 2048, 2049, 4097, 8, 20, 33): + x = (torch.randn(m, k, device="cuda") * 0.125).half() + operator = ( + "gguf_affine_blas_sm70_out" + if m >= 2048 + else "gguf_dense_segments_sm70_out" + if m in (8, 20) + else "gguf_affine_gemm_sm70_out" + ) + with NativeDispatchTrace() as trace: + out = _prepared_gguf_mixed_projection(x, *serialized) + assert any( + r["operator"] == f"_C::{operator}" and r["evidence"] == "cuda_dispatch" + for r in trace.records + ) + + def check(value, x=x): + expected = x[:32].float() @ reference_weight.float().T + error = (value[:32].float() - expected).norm(dim=-1) + relative = error / expected.norm(dim=-1).clamp_min(1e-12) + assert bool(torch.isfinite(value).all()) + assert float(relative.max()) <= 0.003 + + check(out) + check(compiled(x)) + check(projection(x)) + if m == 2048: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run(x) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + graph_out = run(x) + x.mul_(-0.5) + graph.replay() + check(graph_out) + + +def test_mixed_segment_output_views_and_unqualified_fallback(): + config = SimpleNamespace( + kernel_config=KernelConfig(), + compilation_config=SimpleNamespace(static_forward_context={}), + ) + with set_current_vllm_config(config): + # A calibrated Q6 projection plus an uncalibrated Q4 projection: + # only one segment uses BLAS, whose native API requires contiguous out. + sources = [ + (torch.from_numpy(raw_weights(source, n, 2560)).cuda(), source) + for source, n in ((14, 1536), (12, 128)) + ] + projections = prepare_gguf_projections(sources, torch.float16, True, 8) + assert all(hasattr(p, "segment_format") for p in projections) + x = (torch.randn(2049, 2560, device="cuda") * 0.125).half() + with NativeDispatchTrace() as trace: + mixed = _prepared_gguf_mixed_projection( + x, *prepared_projection_arguments(projections) + ) + operators = [ + r["operator"] for r in trace.records if r["evidence"] == "cuda_dispatch" + ] + assert "_C::gguf_affine_blas_sm70_out" in operators + assert "_C::gguf_affine_gemm_sm70_out" in operators + independent = torch.cat([p(x) for p in projections], dim=-1) + torch.testing.assert_close(mixed, independent, rtol=0, atol=0) + assert mixed.is_contiguous() and bool(torch.isfinite(mixed).all()) diff --git a/vllm/model_executor/layers/quantization/gguf_dense_hmma.py b/vllm/model_executor/layers/quantization/gguf_dense_hmma.py index ec62617d60..9fe648ce37 100644 --- a/vllm/model_executor/layers/quantization/gguf_dense_hmma.py +++ b/vllm/model_executor/layers/quantization/gguf_dense_hmma.py @@ -7,6 +7,9 @@ from vllm.config import get_current_vllm_config_or_none from vllm.model_executor.kernels.gguf import GGUFOperatorCapability, decoder_family +from vllm.model_executor.kernels.linear.mixed_precision.sm70_gguf import ( + _get_affine_blas_workspace, +) from vllm.model_executor.layers.quantization.gguf_dense_hmma_formats import ( pack_device, ) @@ -110,7 +113,9 @@ def prepare_segment_bank(projection, canonical, device): return True -def apply_segments(rows, codes, scales, high, formats, ns, output, views): +def apply_segments( + rows, codes, scales, high, formats, ns, output, views, *, blas_bands=() +): k = rows.shape[1] shared = k == 160 or sum(ns) <= 320 scratch = workspace(rows.device, shared) @@ -135,12 +140,27 @@ def apply_segments(rows, codes, scales, high, formats, ns, output, views): None, ) return output - for c, s, h, fmt, n, destination in zip(codes, scales, high, formats, ns, views): - restore_and_apply(rows, c, s, h, fmt, n, destination) + for i, (c, s, h, fmt, n, destination) in enumerate( + zip(codes, scales, high, formats, ns, views) + ): + restore_and_apply( + rows, + c, + s, + h, + fmt, + n, + destination, + blas_bands=blas_bands[i] if blas_bands else (), + ) return output -def restore_and_apply(rows, codes, scale, high, fmt, n, output, k_ld=None, q_ld=None): +def restore_and_apply( + rows, codes, scale, high, fmt, n, output, k_ld=None, q_ld=None, *, blas_bands=() +): + from .gguf_turbomind import _supports_band + k = rows.shape[1] scratch = workspace(rows.device, k == 160 or n <= 320) bits = (4, 5, 6, 4, 8)[fmt] @@ -163,7 +183,25 @@ def restore_and_apply(rows, codes, scale, high, fmt, n, output, k_ld=None, q_ld= # is replaced. Defaults are supplied only by the multi-segment wrapper. if k_ld is None: k_ld, q_ld = descriptor(k, n, fmt) - if fmt == 3: + if fmt != 3 and _supports_band(rows.shape[0], blas_bands): + # Segment storage changes the resident layout, not the admitted + # large-M compute policy. Reuse its prepared, engine-owned scratch. + blas_workspace = _get_affine_blas_workspace(codes) + if blas_workspace is None: + raise RuntimeError("Admitted GGUF BLAS workspace is unavailable") + destination = output if output.is_contiguous() else rows.new_empty(output.shape) + torch.ops._C.gguf_affine_blas_sm70_out( + destination, + rows, + weight, + stats, + bits, + blas_workspace[: k * n].view(k, n), + group, + ) + if destination is not output: + output.copy_(destination) + elif fmt == 3: torch.ops._C.gguf_lut4_gemm_sm70_out( output, rows, weight, stats, 0, k_ld, q_ld, group ) diff --git a/vllm/model_executor/layers/quantization/gguf_turbomind.py b/vllm/model_executor/layers/quantization/gguf_turbomind.py index 631fa237f5..f50ec4b921 100644 --- a/vllm/model_executor/layers/quantization/gguf_turbomind.py +++ b/vllm/model_executor/layers/quantization/gguf_turbomind.py @@ -107,7 +107,16 @@ def _prepared_gguf_projection( ) else: restore_and_apply( - rows, codes, stats, cache, decoder, output_size, out, k_ld, q_ld + rows, + codes, + stats, + cache, + decoder, + output_size, + out, + k_ld, + q_ld, + blas_bands=blas_bands, ) return out.reshape(*x.shape[:-1], logical_size) if cache is not None and _supports_band(rows.shape[0], cache_bands): @@ -184,16 +193,6 @@ def _prepared_gguf_mixed_projection( ) -> torch.Tensor: rows = x.reshape(-1, x.shape[-1]).contiguous() specs = [descriptors[i : i + 9] for i in range(0, len(descriptors), 9)] - if specs and all(spec[0] == 3 for spec in specs): - from .gguf_dense_hmma import apply_segments - - ns = [spec[6] for spec in specs] - output = torch.empty((rows.shape[0], sum(ns)), dtype=x.dtype, device=x.device) - views = list(output.split(ns, dim=1)) - apply_segments( - rows, codes, stats, caches, [spec[1] for spec in specs], ns, output, views - ) - return output.reshape(*x.shape[:-1], sum(ns)) cache_offset = blas_offset = 0 policies = [] for spec in specs: @@ -206,6 +205,24 @@ def _prepared_gguf_mixed_projection( ) cache_offset += cache_count blas_offset += blas_count + if specs and all(spec[0] == 3 for spec in specs): + from .gguf_dense_hmma import apply_segments + + ns = [spec[6] for spec in specs] + output = torch.empty((rows.shape[0], sum(ns)), dtype=x.dtype, device=x.device) + views = list(output.split(ns, dim=1)) + apply_segments( + rows, + codes, + stats, + caches, + [spec[1] for spec in specs], + ns, + output, + views, + blas_bands=[bb for _, bb in policies], + ) + return output.reshape(*x.shape[:-1], sum(ns)) # Retain the calibrated per-projection policy outside measured target # verification sizes. This decision must use actual M inside the op. direct = ( @@ -353,7 +370,7 @@ def prepared_projection_arguments(projections): cb = [] if hasattr(projection, "segment_format"): family, decoder = 3, projection.segment_format - cb = bb = [] + cb = [] caches.append( projection.dmv_high if family == 4 @@ -703,7 +720,7 @@ def forward(self, x): config.partition_weight_shape[1], self.logical_output_size, _admitted_bands(self.cache_capabilities), - [] if family == 3 else _admitted_bands(blas), + _admitted_bands(blas), ) # Imported lazily because the GGUF method owns fallback dispatch. from vllm.model_executor.layers.quantization.gguf import fused_mul_mat_gguf 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 04/19] [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 | 106 +++++++++--------- 5 files changed, 171 insertions(+), 76 deletions(-) diff --git a/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py b/benchmarks/kernels/benchmark_sm70_gguf_cache_copy.py index 7fe5eb7893..b140fb8a9a 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.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 436b9bfebc..47d3a8aab9 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 3f97cfa2e3..e1d913b6e8 100644 --- a/tests/models/qwen4_exp/test_host_kv.py +++ b/tests/models/qwen4_exp/test_host_kv.py @@ -64,6 +64,64 @@ def 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") diff --git a/vllm/model_executor/layers/quantization/gguf_expert_cache.py b/vllm/model_executor/layers/quantization/gguf_expert_cache.py index 99ff21faf2..6f112ac9d8 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 @@ _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 @@ def resolve(self, layer: int, ids: torch.Tensor) -> torch.Tensor: 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 03d2957eb0..9059e19de9 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] + # 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, ) - * Dim - + dims[None, None, :], - need_history[:, None, None], - other=0, - ) - if FP8: - scales = tl.load( - Scales + tokens[:, None] * 2 + kv[None, :], - need_history[:, None], + 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, ) - decoded = (fp8_e4m3fn_bits_to_fp32_bitcast(codes) * scales[:, :, None]).to( - tl.float16 + 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], ) - 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 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 05/19] [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 47d3a8aab9..2b78b4d7e1 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 9ae2ae6db5..477b7f6f00 100644 --- a/tests/kernels/test_sm70_hc_sharded_fallback.py +++ b/tests/kernels/test_sm70_hc_sharded_fallback.py @@ -16,35 +16,65 @@ ) -@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 5decfd764d..45eb39de12 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 project(x): 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 __init__(self): def enable_partial_inputs(self, name, runtime=None): self.partial = True + self.runtime = runtime class MoE: def __init__(self): @@ -246,6 +248,7 @@ class Decoder: 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 __init__(self, index, moe): 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 99a0195bcb..97daf96c31 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 06/19] [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 2b78b4d7e1..19251eec2e 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 45eb39de12..5decfd764d 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 project(x): 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 __init__(self): def enable_partial_inputs(self, name, runtime=None): self.partial = True - self.runtime = runtime class MoE: def __init__(self): @@ -248,7 +246,6 @@ class Decoder: 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 __init__(self, index, moe): 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 97daf96c31..99a0195bcb 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 07/19] [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 d332eb2ce0..42ee3e3a3b 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 66c5bc04ff..6be2ea2ece 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 19251eec2e..684c98eaee 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 7b624dedbe..371e205ef5 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 captured(): 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 f83e58c965..bc6ce2b4d4 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -884,6 +884,11 @@ def resolve_sm70_rmsnorm_gated(self, *, qualified: bool) -> None: 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 f7d578d23a..e343555704 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 9059e19de9..6239affbaf 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( ) 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 @@ def write(self, key: torch.Tensor, value: torch.Tensor, slots: torch.Tensor): self.dim, self.sets, self.fp8, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) @@ -611,6 +659,8 @@ def write_staged(self, key, value, slots, plan, history, scales): self.sets, self.fp8, HOT=False, + HOT_FP8=self.hot_fp8, + HotScales=self.hot_scales, num_warps=4, ) @@ -859,6 +909,8 @@ def _resolve(self, indices, block_table, token_to_req, positions, lengths, stage 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 3195f40dac..e19d265b26 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 94423e914e..86ca086192 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 9d249c68af..d088008efb 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 ce206b5f15..c2b521d699 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -729,6 +729,15 @@ def __init__( # 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 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: 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, From 5545c87ee9cad07d7b0a365c69da0cc6a0c670c8 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 16:08:52 +0800 Subject: [PATCH 08/19] [Deploy] Add the 256K Flash-Next GGUF abliterated launcher Combine the qualified prefill and decode dependencies in an owned local deployment with MTP4, prefix caching, tools, encoded target hot KV and independent draft capacity. Record the model revision, direct mirror download contract and focused GPU qualification. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../flashnext_abliterated_local_20261014.md | 81 ++++++++++++++++ examples/serve_flashnext_gguf_256k.sh | 96 +++++++++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 docs/design/flashnext_abliterated_local_20261014.md create mode 100755 examples/serve_flashnext_gguf_256k.sh diff --git a/docs/design/flashnext_abliterated_local_20261014.md b/docs/design/flashnext_abliterated_local_20261014.md new file mode 100644 index 0000000000..10ebc66f8e --- /dev/null +++ b/docs/design/flashnext_abliterated_local_20261014.md @@ -0,0 +1,81 @@ +# Local Flash-Next IQ3_S abliterated deployment + +This deployment combines the latest integration base `6f3557c973`, the decode +implementation from PR #1204 at `46f5c8c197`, and the qualified prefill +implementation from PR #1205 at `3faaa58d08`. Later prefill experiments remain +outside this deployment. + +The model is +[SC117/Qwen3.8-Flash-Next-GSQ-RCO-abliterated-GGUF](https://huggingface.co/SC117/Qwen3.8-Flash-Next-GSQ-RCO-abliterated-GGUF), +IQ3_S, revision `eec4e10a2c29b440a2ae85591fa11569ce752963`. Both GGUF shards +are downloaded through `https://hf-mirror.com`, with resumable ranges and the +repository's SHA256 values checked before serving. The initial environment-only +proxy bypass did not bypass the host's transparent TUN. Subsequent requests use +real DNS and sockets bound to the physical NIC; the routing mistake and corrected +connection evidence are retained in the external report. +The matching stock BF16 MTP tensors are copied from the existing local +checkpoint unchanged and verified against the previously qualified compact +checkpoint hash. The runtime prepares its existing Q8_0 MTP expert route. +The repository's standalone Q8_0 GGUF draft file is also retained locally. + +## Runtime contract + +Use [the launcher](../../examples/serve_flashnext_gguf_256k.sh) from its owning +source checkout and an environment with the matching native extensions. + +- Four V100-SXM2-16GB GPUs, physical GPUs 4–7, TP4, one active sequence. +- Python 3.12.13, Torch 2.10.0+cu128, Triton 3.6.0, CUDA toolkit 12.8. +- Capacity is **262144 total input and output tokens**, with chunk size 8192. +- Prefix caching, Qwen3 reasoning and Qwen3-Coder tool calling enabled. +- Temperature 0.6, top-p 0.95, top-k 20; thinking enabled. +- MTP4, greedy draft, FULL target graphs at M5, existing draft single graph. +- Original experts retain 80% GPU / 20% host placement. PLE remains on disk. +- Authoritative target KV uses host E4M3. Encoded target hot capacity is 32768; + the independent draft cache retains FP16 and 16384 tokens. +- The accepted grouped-prefill and dense restoration policy, cold-expert DMA, + balanced chunks and host-KV zeroing skip are enabled through existing paths. +- Rejected HC/PLE, indexer, tile and larger-chunk screens are not added. +- Native components are built by normal CMake targets into the source package. + No private native sidecar, source overlay or preload is required. +- Compiler caches and the service's memory guard are owned by this deployment. + +These settings preserve existing runtime support and fallback boundaries. +They do not qualify a different TP size, concurrent requests, all quantization +tiers, or every length within the configured capacity. + +## Reproduce + +Set an API key and the local paths before launching: + +```bash +export MODEL_FILE=/path/to/IQ3_S/Qwen3.8-Flash-Next-GSQ-RCO-abliterated-IQ3_S-00001-of-00002.gguf +export MTP_DIR=/path/to/mtp-stock-bf16 +export CACHE_DIR=/path/to/dedicated-runtime-cache +export PYTHON_BIN=/path/to/venv/bin/python +export CUDA_HOME=/path/to/cuda-12.8 +export VLLM_API_KEY='your-key' +export GPU_IDS=4,5,6,7 CPU_SET=24-47,72-95 +bash examples/serve_flashnext_gguf_256k.sh +``` + +Acquire the selected GPUs' locks and start the memory guard as part of the +service manager. The local deployment retains the existing user service and +public tunnel, with a dedicated internal API port. Public forwarding is enabled +after authenticated generation, tools, prefix reuse and capacity checks pass. +Credentials stay outside the source repository. + +## Evidence + +The prefill dependency's original-model TP4 screen measured **5010 tokens/s** +at 34688 input tokens, median TTFT 6.907 s. The decode dependency's original +32K screen measured conservative mean verifier cost 15.494 ms. These are +dependency results on the 9683 host, not performance measurements of this +abliterated model on the local machine. + +The combined build passes 81 focused GPU cases covering encoded/protected KV, +expert-cache swaps, grouped-prefill boundaries and segment-prefill restoration. +All twelve normal native components are built and their library search paths +are audited. Route checks, unprofiled prefill/decode, tools, prefix reuse +and long-context results are recorded after validation. Raw download journals, +checksums, builds and model responses are retained in the task's external +`flashnext-abliterated-local-20261014` report directory. diff --git a/examples/serve_flashnext_gguf_256k.sh b/examples/serve_flashnext_gguf_256k.sh new file mode 100755 index 0000000000..47eaf8888b --- /dev/null +++ b/examples/serve_flashnext_gguf_256k.sh @@ -0,0 +1,96 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Qualified Flash-Next GGUF TP4 settings: 256K total tokens, host KV and disk PLE. +set -euo pipefail +SOURCE_DIR=${SOURCE_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)} +PYTHON_BIN=${PYTHON_BIN:-$SOURCE_DIR/.venv/bin/python} +: "${MODEL_FILE:?Set MODEL_FILE to the first IQ3_S GGUF shard}" +: "${MTP_DIR:?Set MTP_DIR to the matching stock MTP safetensors directory}" +: "${CACHE_DIR:?Set CACHE_DIR to a dedicated runtime cache}" +: "${VLLM_API_KEY:?Set VLLM_API_KEY for the public API}" +PORT=${PORT:-8000} +GPU_IDS=${GPU_IDS:-4,5,6,7} +CPU_SET=${CPU_SET:-24-47,72-95} +EXPERT_PROFILE=${EXPERT_PROFILE:-} +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.8} +export PATH="$CUDA_HOME/bin:$(dirname "$PYTHON_BIN"):$PATH" +export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES="$GPU_IDS" +export PYTHONPATH="$SOURCE_DIR" +export TORCH_CUDA_ARCH_LIST=7.0 +export PYTORCH_ALLOC_CONF=expandable_segments:True OMP_NUM_THREADS=1 MALLOC_ARENA_MAX=2 +export VLLM_NO_USAGE_STATS=1 DO_NOT_TRACK=1 +export VLLM_EXPERT_CACHE_STACKS=1 VLLM_BALANCED_PREFILL_CHUNKS=1 +export VLLM_SM70_MTP_PROFILE=0 VLLM_NVTX_SCOPES_FOR_PROFILING=0 +export MAX_JOBS=2 TORCHINDUCTOR_COMPILE_THREADS=1 +export TOKENIZERS_PARALLELISM=false +export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 +export VLLM_CACHE_ROOT="$CACHE_DIR/vllm" TRITON_CACHE_DIR="$CACHE_DIR/triton" +export TORCHINDUCTOR_CACHE_DIR="$CACHE_DIR/inductor" +export TORCH_EXTENSIONS_DIR="$CACHE_DIR/torch_extensions" +export VLLM_SM70_GEMM_LUT_PATH="$CACHE_DIR/gemm-lut-{device}.bin" +unset LD_PRELOAD LD_AUDIT ONECAT_VLLM_EXTENSION_DIR +mkdir -p "$CACHE_DIR" +KC=$("$PYTHON_BIN" - "$EXPERT_PROFILE" <<'PY' +import json +import sys +config = { + "sm70_gguf": { + "expert_storage": "original", "expert_arena": True, + "embedding_storage": "original", "dense_storage": "canonical", + "small_m_dp4a": True, "q8_expert_intermediate": True, + "small_m_hmma": True, "lut4_expert_dp4a": True, + "device_transcode": True, "dequant_workspace_bytes": 33554432, + "expert_cache_fraction": 0.8, "expert_cache_host_slack": 0.05, + "expert_cache_prefill": "turbomind", "expert_cache_prefill_tokens": 4096, + }, + "sm70_turbomind_workspace_bytes": 8388608, + "sm70_router_weight_storage": "row_major", "sm70_mtp_q8_0_experts": True, + "hc_weight_storage": "sharded", "hc_ll_shard": True, + "hc_ll_optimized_loads": False, "sm70_hcx": True, + "sm70_hcx_output_projection": False, "sm70_hcx_local_schedule": True, + "sm70_qsa_shared_key": True, "sm70_qsa_device_history": True, + "sm70_qsa_host_native": True, "qsa_host_kv": True, + "qsa_host_indexer_history": False, "qsa_host_kv_dtype": "fp8_e4m3", + "qsa_host_kv_draft_dtype": "float16", "qsa_host_kv_device_reference": False, + "qsa_host_kv_hot_tokens": 32768, "qsa_host_kv_hot_dtype": "fp8_e4m3", + "qsa_host_kv_draft_hot_tokens": 16384, "qsa_host_kv_staging_rows": 8, + "qsa_host_kv_prefill_stage_tokens": 131072, + "sm70_sparse": {"qsa_score_tile_mb": 32}, "qsa_host_kv_state_blocks": 29, + "qsa_auto_e4m3": False, "ple_row_cache_mib": 512, + "ple_input_prepare": True, "ple_pinned_decode": False, + "ple_fast_lane": True, "ple_draft_prefetch": True, + "sm70_greedy_verify": True, "sm70_sparse_target_sampling": True, + "sm70_sparse_target_reference_check": False, "sm70_draft_single_graph": True, + "sm70_fused_side_projections": True, "sm70_top1x": True, + "sm70_draft_hot_vocab": 0, +} +if sys.argv[1]: + config["sm70_gguf"]["expert_cache_profile"] = sys.argv[1] +print(json.dumps(config)) +PY +) +SPEC=$("$PYTHON_BIN" - "$MTP_DIR" <<'PY' +import json +import sys +print(json.dumps({ + "method": "mtp", "model": sys.argv[1], "num_speculative_tokens": 4, + "draft_load_config": {"load_format": "safetensors"}, + "draft_sample_method": "greedy", +})) +PY +) +cd "$SOURCE_DIR" +exec taskset -c "$CPU_SET" "$PYTHON_BIN" -m vllm.entrypoints.cli.main serve "$MODEL_FILE" \ + --served-model-name flash-next-abliterated flash-next qwen3.8-flash-next \ + --host 127.0.0.1 --port "$PORT" --quantization gguf --tensor-parallel-size 4 \ + --dtype half --kv-cache-dtype float16 --mamba-ssm-cache-dtype float32 \ + --max-model-len 262144 --max-num-seqs 1 --max-num-batched-tokens 8192 \ + --gpu-memory-utilization 0.95 --enable-prefix-caching \ + --mamba-cache-mode align --no-mamba-tail-checkpoint --language-model-only \ + --enable-auto-tool-choice --tool-call-parser qwen3_coder --reasoning-parser qwen3 \ + --default-chat-template-kwargs '{"enable_thinking":true}' \ + --override-generation-config '{"temperature":0.6,"top_p":0.95,"top_k":20}' \ + --compilation-config '{"mode":3,"cudagraph_mode":"FULL","cudagraph_capture_sizes":[5],"inductor_compile_config":{"benchmark_combo_kernel":false}}' \ + --kernel-config "$KC" --speculative-config "$SPEC" + From 5ddfe43dc601d5faebe13e8d9b8f88d1a289686e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 17:42:43 +0800 Subject: [PATCH 09/19] [Doc] Record the qualified local Flash-Next deployment Record measured prefill and pure decode, near-256K task completion, prefix hit counters and public API validation. Retain the measured prefix admission limitation and transparent proxy routing correction. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../flashnext_abliterated_local_20261014.md | 47 ++++++++++++++++--- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/docs/design/flashnext_abliterated_local_20261014.md b/docs/design/flashnext_abliterated_local_20261014.md index 10ebc66f8e..fa1603b716 100644 --- a/docs/design/flashnext_abliterated_local_20261014.md +++ b/docs/design/flashnext_abliterated_local_20261014.md @@ -60,9 +60,11 @@ bash examples/serve_flashnext_gguf_256k.sh Acquire the selected GPUs' locks and start the memory guard as part of the service manager. The local deployment retains the existing user service and -public tunnel, with a dedicated internal API port. Public forwarding is enabled -after authenticated generation, tools, prefix reuse and capacity checks pass. -Credentials stay outside the source repository. +public tunnel, with a dedicated internal API port. The existing bridge now +forwards to the verified model instance. Read the tunnel client's published +endpoint when connecting: its assigned public port can differ from the port +requested in the preset configuration. Credentials stay outside the source +repository. ## Evidence @@ -75,7 +77,40 @@ abliterated model on the local machine. The combined build passes 81 focused GPU cases covering encoded/protected KV, expert-cache swaps, grouped-prefill boundaries and segment-prefill restoration. All twelve normal native components are built and their library search paths -are audited. Route checks, unprofiled prefill/decode, tools, prefix reuse -and long-context results are recorded after validation. Raw download journals, -checksums, builds and model responses are retained in the task's external +are audited. All model files pass their complete SHA256 checks. Worker logs +confirm native encoded QSA, Q8_0 MTP expert storage and the existing GDN/PLE paths. + +Five unprofiled cold-prefix requests use two warmups and three measured trials, +with the same 32000-word/seed-9683 workload and an early per-request nonce. +Production development cache-reset routes remain disabled; the unique first +block prevents prefix reuse. Measured inputs are 34688, 34606 and 34631 tokens. +Median first-choice latency is **6.805 s**, or **5088.8 input tokens/s**. Server +prefill time is recorded separately. This is the local abliterated model's +measurement, rather than a projected gain from the dependency results. + +Eight prefix-warm, unprofiled 32774-input / 192-output requests at production +sampling measure **6.454 ms/output token**, or **154.9 output tokens/s**. +Mean decode time is 1.233 s, excluding 0.893-s TTFT. The whole-round HTTP +estimate is 18.128 ms with 2.838 emitted tokens per round. It includes draft and +scheduling and must not be labelled an isolated GPU verifier timer. The +recorded affinity changes affect only this service's API/worker descendants. + +Arithmetic and automatic tool-call/follow-up checks pass with natural EOS. +Retrieval/arithmetic tasks at 32761, 131070 and 260977 input tokens pass all +answer, EOS and repetition checks, with 181, 165 and 188 output tokens. The +largest case uses 261165 total tokens and finishes in 286.959 s without OOM +or a memory-guard intervention. These are capacity and task checks, not cold +throughput baselines at each length. + +Prefix-cache counters prove 31824 cached tokens on the third identical +32774-input request after the long-context checks. Total request time falls +from 8.921/9.304 s to 2.731 s. A separate 18282-token repeated prompt does not +produce a hit in two trials; that observation is retained. Prefix caching uses +the existing Mamba alignment/admission policy and finite cache capacity. This +evidence establishes reuse for the measured workload, not every prompt length +or guaranteed residency after other requests. + +Public authentication, model discovery and streaming arithmetic are checked +through the existing tunnel before delivery. Raw download journals, checksums, +builds, counters and model responses are retained in the task's external `flashnext-abliterated-local-20261014` report directory. From 699c864d3c3a841b8626b657a1c8f236457ab1b5 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:20:17 +0800 Subject: [PATCH 10/19] [Deploy] Qualify the abliterated model on 9683 Expose the existing host-indexer cache setting in the launcher while preserving its default. Record destination-native GPU checks, memory qualification, near-256K task completion, matched prefill, prefix reuse and public authentication/stream evidence. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../flashnext_abliterated_9683_20261014.md | 111 ++++++++++++++++++ .../flashnext_abliterated_local_20261014.md | 7 ++ examples/serve_flashnext_gguf_256k.sh | 8 +- 3 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 docs/design/flashnext_abliterated_9683_20261014.md diff --git a/docs/design/flashnext_abliterated_9683_20261014.md b/docs/design/flashnext_abliterated_9683_20261014.md new file mode 100644 index 0000000000..ca5181649e --- /dev/null +++ b/docs/design/flashnext_abliterated_9683_20261014.md @@ -0,0 +1,111 @@ +# Flash-Next abliterated deployment on 9683 + +This deployment reuses the source and model contract of the +[qualified local deployment](flashnext_abliterated_local_20261014.md) on the +9683 host. It combines integration `6f3557c973`, decode `46f5c8c197` and the +qualified prefill implementation `3faaa58d08`. The copied source package is +verified against 2862 source/native file hashes. Native components are the +normal CMake build described in the local report, loaded from the deployed +package through their normal loaders; no private kernel overlay or preload is +required. The shared remote Python environment is not modified. + +## Runtime and model contract + +- Ubuntu 24.04, Python 3.12.3, Torch 2.10.0+cu128, CUDA toolkit 12.8. +- Four V100-SXM2-16GB GPUs with NVLink, physical GPUs 0–3, TP4, 46 GiB RAM. +- SC117 Flash-Next GSQ-RCO abliterated IQ3_S, revision + `eec4e10a2c29b440a2ae85591fa11569ce752963`. +- 262144 total input/output-token capacity, 8192-token chunks, one active + sequence; thinking, prefix caching and Qwen3-Coder tools enabled. +- T0.6, top-p0.95, top-k20; matching stock BF16 MTP with the existing Q8_0 + preparation route, greedy MTP4 and FULL M5 target graphs. +- Experts retain 80% GPU / 20% host placement. PLE remains on disk, target + history remains host E4M3 with encoded HOT32768, draft FP16 HOT16384. +- `QSA_HOST_INDEXER_HISTORY=true` also places compressed indexer history on + host, freeing about 230 MiB device memory. The launcher default remains + `false` for the separately measured local deployment. +- Private runtime caches, GPU ownership locks and a process-group memory guard + are retained. The guard requires 5 GiB host memory available, matching the + existing host's memory-protection threshold. + +Correctness support and fallbacks remain those of the existing implementation. +This qualification does not cover concurrent execution, other TP sizes or +quantization tiers, or every context length within the configured capacity. + +## Download and public-service transition + +The new first shard is downloaded directly on the destination from +`https://hf-mirror.com`, with real DNS and sockets bound to the physical NIC. +The host also has a transparent proxy; clearing proxy environment variables +alone is insufficient. The physical socket endpoints are retained as evidence. +The attempted machine-to-machine weight transfer is stopped at the user's +request. The final first-shard checksum is +`40f369eca301ec6ea318dbb671b3132c26a2f0b89aa5458a7051b33215ee296c`. + +The second shard is byte-identical to the existing stock shard, verified as +`316b46f3a2dbd68c900f43136ab9449f9dcc3725dfd8c794847c204bc161e113`. +It is retained through an owned hard link without another download or disk +allocation. The new first shard uses the second disk because the root disk is +nearly full. Existing model weights and source remain available for rollback. +The matching compact BF16 MTP checksum remains +`d50d273a87f51d9b7de3142a3f93f2c2823cd36a6fbf0f53f08303ce33029c17`. + +The old Strata unit owns this host's public API and GPU locks. It is stopped +only after the replacement model has downloaded and passed its checksum. The +new model uses the existing public tunnel and an authenticated proxy, with the +same credential supplied to the backend. Credentials remain outside Git. +Two duplicate tunnel clients are reduced to one, eliminating their repeated +control-channel disconnects. + +## Verification + +All shipped native components load through their normal package loaders. +Twenty-one focused GPU cases pass on the destination, covering encoded hot +storage, changed graph inputs, missing-op fallback and misaligned rejection +(22 unaffected cases deselected). The existing local qualification retains its +81-case evidence separately. + +The first configuration passes arithmetic and tools but runs out of device +memory on its first long prefill: PLE convolution needs a 144 MiB output +allocation with about 92 MiB free. The existing host-indexer setting addresses +this without reducing context capacity, expert placement, MTP depth or the +8192-token chunk. Three focused cache-budget tests pass, including the existing +host-indexer pool case. The first cold startup also exceeds the controller's +20-minute readiness timeout; its model process continues compiling normally +and completes without being restarted by that timeout. These failed attempts +and their raw logs are retained. + +With host indexer history enabled, arithmetic, automatic tool calls and replies +pass with natural EOS. Cold-prefix prefill uses the same 32000-word/seed-9683 +contract, two warmups and three measured trials, early unique nonces and no +profiler. Actual measured inputs are 34688, 34606 and 34631 tokens. Median TTFT +is **6.944 s / 4987.2 input tokens/s**. After applying the existing host's CPU +isolation pattern only to owned API/worker descendants, a matching repeat gives +**6.945 s / 4982.5 tokens/s**: no measurable prefill gain is claimed. These +measurements are close to, but do not establish the strict 5000-tokens/s target. + +The third identical 32774-input prefix request proves 31824 cached tokens; +total request time falls from 11.476/10.886 s to 4.621 s. Retrieval/arithmetic +tasks at 32761, 131070 and 260977 input tokens pass answer, EOS and repetition +checks (196, 199 and 210 output tokens). The largest completes with 261187 +total tokens in 310.776 s, without OOM or a memory-guard intervention under +the final configuration. These are task/capacity checks, not cold throughput +baselines at each length. + +Public authentication rejects an unauthenticated model query; authenticated +discovery and streamed arithmetic succeed with the correct 1517 and EOS. +The initial external generation queues behind the long-context check and +takes 49.689 s to its first chunk; after the owned test queue is empty, the +same short check reaches its first streamed chunk in **1.050 s**, completing +in 1.341 s. First HTTP/SSE chunks, first generated tokens and first visible +answer text remain separate concepts. Single active-sequence capacity means +simultaneous Agent requests can still queue. + +The new model and proxy units are enabled for startup; the old Strata unit is +disabled with its definition and prior enabled state retained for rollback. +The CPU isolation helper is retained in the deployment wrapper for subsequent +starts. Production development cache-reset endpoints remain disabled. + +The external `flashnext-abliterated-9683-20261014` report retains checksums, +source and native identity, failed/abandoned attempts, GPU tests, runtime logs +and request responses. Use its handoff for the live service paths and endpoint. diff --git a/docs/design/flashnext_abliterated_local_20261014.md b/docs/design/flashnext_abliterated_local_20261014.md index fa1603b716..f77658466a 100644 --- a/docs/design/flashnext_abliterated_local_20261014.md +++ b/docs/design/flashnext_abliterated_local_20261014.md @@ -55,6 +55,7 @@ export PYTHON_BIN=/path/to/venv/bin/python export CUDA_HOME=/path/to/cuda-12.8 export VLLM_API_KEY='your-key' export GPU_IDS=4,5,6,7 CPU_SET=24-47,72-95 +export QSA_HOST_INDEXER_HISTORY=false bash examples/serve_flashnext_gguf_256k.sh ``` @@ -66,6 +67,12 @@ endpoint when connecting: its assigned public port can differ from the port requested in the preset configuration. Credentials stay outside the source repository. +`QSA_HOST_INDEXER_HISTORY` accepts `true` or `false` and defaults to `false`, +preserving this local measurement's configuration. `true` selects the existing +host-backed compressed-indexer cache, reserving more device memory for prefill +temporaries. The [9683 deployment](flashnext_abliterated_9683_20261014.md) +records the separate qualification of that setting. + ## Evidence The prefill dependency's original-model TP4 screen measured **5010 tokens/s** diff --git a/examples/serve_flashnext_gguf_256k.sh b/examples/serve_flashnext_gguf_256k.sh index 47eaf8888b..f1a8b8161d 100755 --- a/examples/serve_flashnext_gguf_256k.sh +++ b/examples/serve_flashnext_gguf_256k.sh @@ -13,6 +13,7 @@ PORT=${PORT:-8000} GPU_IDS=${GPU_IDS:-4,5,6,7} CPU_SET=${CPU_SET:-24-47,72-95} EXPERT_PROFILE=${EXPERT_PROFILE:-} +QSA_HOST_INDEXER_HISTORY=${QSA_HOST_INDEXER_HISTORY:-false} export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.8} export PATH="$CUDA_HOME/bin:$(dirname "$PYTHON_BIN"):$PATH" export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES="$GPU_IDS" @@ -31,9 +32,11 @@ export TORCH_EXTENSIONS_DIR="$CACHE_DIR/torch_extensions" export VLLM_SM70_GEMM_LUT_PATH="$CACHE_DIR/gemm-lut-{device}.bin" unset LD_PRELOAD LD_AUDIT ONECAT_VLLM_EXTENSION_DIR mkdir -p "$CACHE_DIR" -KC=$("$PYTHON_BIN" - "$EXPERT_PROFILE" <<'PY' +KC=$("$PYTHON_BIN" - "$EXPERT_PROFILE" "$QSA_HOST_INDEXER_HISTORY" <<'PY' import json import sys +if sys.argv[2].lower() not in ("true", "false"): + raise ValueError("QSA_HOST_INDEXER_HISTORY must be true or false") config = { "sm70_gguf": { "expert_storage": "original", "expert_arena": True, @@ -51,7 +54,8 @@ config = { "sm70_hcx_output_projection": False, "sm70_hcx_local_schedule": True, "sm70_qsa_shared_key": True, "sm70_qsa_device_history": True, "sm70_qsa_host_native": True, "qsa_host_kv": True, - "qsa_host_indexer_history": False, "qsa_host_kv_dtype": "fp8_e4m3", + "qsa_host_indexer_history": sys.argv[2].lower() == "true", + "qsa_host_kv_dtype": "fp8_e4m3", "qsa_host_kv_draft_dtype": "float16", "qsa_host_kv_device_reference": False, "qsa_host_kv_hot_tokens": 32768, "qsa_host_kv_hot_dtype": "fp8_e4m3", "qsa_host_kv_draft_hot_tokens": 16384, "qsa_host_kv_staging_rows": 8, From c7c0eb9c5f936957288f320db83a748c5a5ec74e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 15 Oct 2026 03:51:44 +0800 Subject: [PATCH 11/19] [Perf] Reduce Flash-Next MTP4 verifier batch inflation Reuse existing SM70 kernels in bounded verifier tiles, compact recurrent state banks, and score replicated host indexer history once per TP group. Keep acceleration options off by default and retain graph-safe fallbacks. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../flashnext_prefill/analyze_mtp_verifier.py | 101 +++++++ .../measure_mtp_concurrency.py | 211 +++++++++++++++ .../flashnext_prefill/mtp_trace_worker.py | 85 ++++++ docs/design/architecture/runtime_reference.md | 4 +- .../flashnext_mtp_batch_16gb_20261015.md | 247 ++++++++++++++++++ .../config/test_flashnext_mtp_batch_policy.py | 94 +++++++ .../quantization/test_gguf_compact_storage.py | 2 + .../quantization/test_gguf_mtp_batch.py | 106 ++++++++ .../kernels/test_sm70_hcx_compiled_payload.py | 7 +- .../qwen4_exp/test_device_kv_attention.py | 93 +++++++ tests/models/qwen4_exp/test_qsa_cache.py | 24 +- .../qwen4_exp/test_qsa_indexer_single_rank.py | 136 ++++++++++ .../qwen4_exp/test_qsa_shared_key_batch.py | 164 ++++++++++++ tests/v1/core/test_qwen4_exp_kv_cache.py | 84 +++++- .../worker/test_cpu_block_table_dispatch.py | 18 ++ .../test_gpu_model_runner_v2_kv_zero.py | 18 +- tests/v1/worker/test_gpu_warmup_blocks.py | 50 ++++ tests/v1/worker/test_kv_block_zeroer.py | 101 +++++++ vllm/config/kernel.py | 40 +++ .../layers/quantization/gguf_moe.py | 53 ++-- .../qwen4_exp/nvidia/hyperconnection.py | 13 +- vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 56 +++- vllm/models/qwen4_exp/nvidia/model.py | 19 +- vllm/models/qwen4_exp/nvidia/ops/qsa.py | 28 +- .../qwen4_exp/nvidia/ops/qsa_shared_key.py | 73 ++++++ vllm/models/qwen4_exp/nvidia/qsa.py | 15 +- vllm/models/qwen4_exp/nvidia/sm70_hcx.py | 84 +++--- vllm/v1/core/block_pool.py | 82 ++++-- vllm/v1/core/kv_cache_coordinator.py | 1 + vllm/v1/core/kv_cache_utils.py | 78 +++++- vllm/v1/kv_cache_interface.py | 33 +++ vllm/v1/worker/block_table.py | 15 ++ vllm/v1/worker/gpu/block_table.py | 9 + vllm/v1/worker/gpu/model_runner.py | 7 + vllm/v1/worker/gpu/warmup.py | 40 ++- vllm/v1/worker/gpu_model_runner.py | 8 + vllm/v1/worker/utils.py | 33 ++- 37 files changed, 2095 insertions(+), 137 deletions(-) create mode 100644 benchmarks/flashnext_prefill/analyze_mtp_verifier.py create mode 100644 benchmarks/flashnext_prefill/measure_mtp_concurrency.py create mode 100644 benchmarks/flashnext_prefill/mtp_trace_worker.py create mode 100644 docs/design/flashnext_mtp_batch_16gb_20261015.md create mode 100644 tests/config/test_flashnext_mtp_batch_policy.py create mode 100644 tests/kernels/quantization/test_gguf_mtp_batch.py create mode 100644 tests/models/qwen4_exp/test_qsa_indexer_single_rank.py create mode 100644 tests/models/qwen4_exp/test_qsa_shared_key_batch.py diff --git a/benchmarks/flashnext_prefill/analyze_mtp_verifier.py b/benchmarks/flashnext_prefill/analyze_mtp_verifier.py new file mode 100644 index 0000000000..d6f3c20e93 --- /dev/null +++ b/benchmarks/flashnext_prefill/analyze_mtp_verifier.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diagnostic MTP4 verifier phases, aligned by step ordinal across all TP4 ranks.""" + +import argparse +import json +import statistics +from pathlib import Path + +import regex as re + +p = argparse.ArgumentParser() +p.add_argument("log", type=Path) +p.add_argument("requests", type=Path) +p.add_argument("output", type=Path) +a = p.parse_args() + + +def stats(v): + x = sorted(v) + return dict( + n=len(x), + mean=statistics.mean(x), + p50=statistics.median(x), + p90=x[round(0.9 * (len(x) - 1))], + p99=x[round(0.99 * (len(x) - 1))], + max=x[-1], + ) + + +raw = a.log.read_bytes() +result = [] +for request in json.loads(a.requests.read_text()): + c = request["concurrency"] + groups = {} + for line in ( + raw[request["profile_byte_start"] : request["profile_byte_end"]] + .decode(errors="replace") + .splitlines() + ): + match = re.search( + r"Worker_TP(\d+).*SM70 V2 MTP profile interval_avg_ms " + r"calls=(\d+) interval_calls=1 tokens=(\d+) drafts=(\d+) (.*)", + line, + ) + if not match: + continue + rank, ordinal, tokens, drafts = map(int, match.groups()[:4]) + if (tokens, drafts) != (5 * c, 4 * c): + continue + values = {k: float(v) for k, v in re.findall(r"([a-z_]+)=([0-9.]+)", match[5])} + if rank in groups.setdefault(ordinal, {}): + raise ValueError("Duplicate rank ordinal") + groups[ordinal][rank] = values + complete = [(n, g) for n, g in sorted(groups.items()) if set(g) == {0, 1, 2, 3}] + if request["repetition"] < 0: + continue + if len(complete) < 16: + raise ValueError(("Too few full-concurrency verifier rounds", c, len(complete))) + records = [] + for ordinal, ranks in complete: + critical = max(ranks, key=lambda r: ranks[r]["target_verifier_gpu"]) + rec = dict( + ordinal=ordinal, critical_rank=critical, ranks=ranks, **ranks[critical] + ) + rec["rank_max_verifier_wall_ms"] = max( + r["target_verifier_wall_cpu"] for r in ranks.values() + ) + records.append(rec) + fields = ( + "target_verifier_gpu", + "rank_max_verifier_wall_ms", + "target_forward", + "target_sample", + "target_state_update", + "draft_total", + "total_gpu", + ) + row = dict( + concurrency=c, + repetition=request["repetition"], + complete_rounds=len(complete), + incomplete_rounds=len(groups) - len(complete), + phases={k: stats([r[k] for r in records]) for k in fields}, + per_round=records, + ) + result.append(row) + print(c, request["repetition"], json.dumps(row["phases"]), flush=True) +a.output.write_text( + json.dumps( + dict( + definition=( + "Per-ordinal TP4 maximum GPU event sum: target forward + sampling " + "+ state commit. Diagnostic events/fences are enabled. Draft is " + "separate; stage statistics are taken on the critical verifier rank." + ), + cases=result, + ), + indent=2, + ) +) diff --git a/benchmarks/flashnext_prefill/measure_mtp_concurrency.py b/benchmarks/flashnext_prefill/measure_mtp_concurrency.py new file mode 100644 index 0000000000..a1b5f249af --- /dev/null +++ b/benchmarks/flashnext_prefill/measure_mtp_concurrency.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Matched concurrent MTP requests, with an explicit common decode interval. + +Token timestamps are HTTP arrival times, not GPU verifier timers. Warm the same +prompts first; use the worker step profiler separately for phase attribution. +""" + +import argparse +import hashlib +import json +import os +import statistics +import threading +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import regex as re + + +def request(base, path, body=None): + data = json.dumps(body).encode() if body is not None else None + headers = {"Content-Type": "application/json"} + if key := os.environ.get("VLLM_API_KEY"): + headers["Authorization"] = "Bearer " + key + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + return opener.open(urllib.request.Request(base + path, data, headers), timeout=900) + + +def counters(base): + with request(base, "/metrics") as response: + text = response.read().decode() + return { + name: sum( + float(value) + for value in re.findall( + rf"^vllm:spec_decode_num_{name}_total\{{[^}}]*\}} ([0-9.e+]+)$", + text, + re.M, + ) + ) + for name in ("drafts", "draft_tokens", "accepted_tokens") + } + + +def run_one(base, case, output_tokens, barrier, seed): + body = { + "model": "flash-next-abliterated", + "messages": [{"role": "user", "content": case["prompt"]}], + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "seed": seed, + "max_tokens": output_tokens, + "chat_template_kwargs": {"enable_thinking": True}, + "stream": True, + "stream_options": {"include_usage": True}, + "return_token_ids": True, + } + barrier.wait() + started = time.perf_counter() + events = [] + content = [] + reasoning = [] + usage = None + finish = None + done = False + with request(base, "/v1/chat/completions", body) as response: + for raw in response: + line = raw.decode().strip() + if not line.startswith("data:"): + continue + if line == "data: [DONE]": + done = True + continue + obj = json.loads(line[5:]) + if obj.get("usage"): + usage = obj["usage"] + for choice in obj.get("choices", []): + finish = choice.get("finish_reason") or finish + delta = choice.get("delta", {}) + ids = choice.get("token_ids") or delta.get("token_ids") or [] + if ids: + events.append({"at": time.perf_counter(), "token_ids": ids}) + content.append(delta.get("content") or "") + reasoning.append(delta.get("reasoning") or "") + ended = time.perf_counter() + if not done or usage is None or not events: + raise RuntimeError("Incomplete stream or missing emitted token IDs") + emitted = sum(len(event["token_ids"]) for event in events) + if emitted != usage["completion_tokens"]: + raise RuntimeError(f"Token accounting differs: {emitted} versus {usage}") + first, last = events[0]["at"], events[-1]["at"] + decode_tokens = emitted - len(events[0]["token_ids"]) + return { + "seed": seed, + "prompt_sha256": hashlib.sha256(case["prompt"].encode()).hexdigest(), + "usage": usage, + "finish_reason": finish, + "ttft_s": first - started, + "total_s": ended - started, + "decode_s": last - first, + "decode_tokens": decode_tokens, + "output_token_ms": 1000 * (last - first) / max(decode_tokens, 1), + "events": events, + "content": "".join(content), + "reasoning": "".join(reasoning), + } + + +def run_batch(base, case, concurrency, output_tokens, repetition, independent_prefix): + before = counters(base) + barrier = threading.Barrier(concurrency) + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=concurrency) as pool: + futures = [] + for i in range(concurrency): + worker_case = dict(case) + if independent_prefix: + worker_case["prompt"] = f"[Independent request {i}]\n" + case["prompt"] + futures.append( + pool.submit( + run_one, base, worker_case, output_tokens, barrier, 20261014 + i + ) + ) + rows = [future.result() for future in futures] + time.sleep(0.2) # Let the final asynchronous scheduler counters settle. + after = counters(base) + # Exclude every request's first and last emission. All requests must still + # be decoding in the shared interval; this avoids calling tail drain C8. + common_start = max(row["events"][0]["at"] for row in rows) + common_end = min(row["events"][-1]["at"] for row in rows) + seconds = common_end - common_start + common_tokens = sum( + len(event["token_ids"]) + for row in rows + for event in row["events"] + if common_start < event["at"] < common_end + ) + valid_overlap = seconds > 0 and common_tokens >= concurrency * 16 + if repetition >= 0 and not valid_overlap: + raise RuntimeError( + "Insufficient common decode interval; increase output length" + ) + delta = {key: after[key] - before[key] for key in before} + return { + "concurrency": concurrency, + "prefix_mode": "independent" if independent_prefix else "shared", + "repetition": repetition, + "batch_elapsed_s": time.perf_counter() - start, + "mean_ttft_s": statistics.mean(row["ttft_s"] for row in rows), + "mean_output_token_ms": statistics.mean(row["output_token_ms"] for row in rows), + "common_decode_s": seconds if valid_overlap else None, + "common_decode_tokens": common_tokens, + "common_aggregate_tokens_per_s": common_tokens / seconds + if valid_overlap + else None, + "common_per_request_token_ms": ( + 1000 * seconds * concurrency / common_tokens if valid_overlap else None + ), + "accepted_length": 1 + delta["accepted_tokens"] / max(delta["drafts"], 1), + "mtp_metrics_delta": delta, + "requests": rows, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:18041") + parser.add_argument("--cases", type=Path, required=True) + parser.add_argument("--case", default="i32768") + parser.add_argument("--concurrency", nargs="+", type=int, default=[1, 2, 4, 6, 8]) + parser.add_argument("--output-tokens", type=int, default=384) + parser.add_argument("--repetitions", type=int, default=3) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--independent-prefix", action="store_true") + parser.add_argument("--service-log", type=Path) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + case = next( + row for row in json.loads(args.cases.read_text()) if row["name"] == args.case + ) + results = [] + for concurrency in args.concurrency: + for repetition in range(-args.warmups, args.repetitions): + log_start = args.service_log.stat().st_size if args.service_log else None + row = run_batch( + args.base_url, + case, + concurrency, + args.output_tokens, + repetition, + args.independent_prefix, + ) + results.append(row) + if args.service_log: + row.update( + profile_byte_start=log_start, + profile_byte_end=args.service_log.stat().st_size, + ) + args.out.write_text(json.dumps(results, indent=2)) + print( + json.dumps({k: v for k, v in row.items() if k != "requests"}), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/flashnext_prefill/mtp_trace_worker.py b/benchmarks/flashnext_prefill/mtp_trace_worker.py new file mode 100644 index 0000000000..d2748b5dcc --- /dev/null +++ b/benchmarks/flashnext_prefill/mtp_trace_worker.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diagnostic NVTX annotations installed after graph capture. + +Use only for the Nsight composition pass. Ordinary speed admission uses the +normal worker, without these wrappers or synchronized step profiling. +""" + +from functools import wraps + +import torch + +from vllm.v1.worker.gpu_worker import Worker + + +def annotate(owner, name, label): + original = getattr(owner, name) + + @wraps(original) + def observed(*args, **kwargs): + text = label(*args, **kwargs) if callable(label) else label + with torch.cuda.nvtx.range(text): + return original(*args, **kwargs) + + setattr(owner, name, observed) + + +class TraceWorker(Worker): + def profile(self, is_start=True, profile_prefix=None): + if is_start: + # The benchmark reaches this boundary with an idle request queue. + # Release unused eager-prefill allocations before CUPTI allocates + # its buffers; graph-owned tensors and pointers remain live. + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + return super().profile(is_start, profile_prefix) + + def compile_or_warm_up_model(self): + result = super().compile_or_warm_up_model() + runner = self.model_runner + prepared = {} + prepare = runner.prepare_inputs + + @wraps(prepare) + def prepare_inputs(*args, **kwargs): + with torch.cuda.nvtx.range("flashnext_batch.prepare"): + batch = prepare(*args, **kwargs) + prepared.update( + requests=batch.num_reqs, + tokens=batch.num_tokens, + drafts=batch.num_draft_tokens, + ) + return batch + + runner.prepare_inputs = prepare_inputs + annotate(runner, "sample", "flashnext_batch.target_sample") + annotate(runner, "postprocess_sampled", "flashnext_batch.target_state") + annotate( + runner.cudagraph_manager, + "run_fullgraph", + lambda desc: ( + f"flashnext_batch.target/M{desc.num_tokens}/R{desc.num_reqs}" + f"/actualR{prepared['requests']}/T{prepared['tokens']}" + f"/D{prepared['drafts']}" + ), + ) + draft = runner.speculator + annotate( + draft, + "propose", + lambda batch, *args, **kwargs: f"flashnext_batch.draft/R{batch.num_reqs}", + ) + for manager, name in ( + (draft.prefill_cudagraph_manager, "draft_first"), + (draft.decode_cudagraph_manager, "draft_rest"), + (getattr(draft, "multistep_cudagraph_manager", None), "draft_chain"), + ): + if manager is None: + continue + annotate( + manager, + "run_fullgraph", + lambda desc, name=name: f"flashnext_batch.{name}/M{desc.num_tokens}", + ) + return result diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index 6be2ea2ece..a17ac1896c 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -38,12 +38,12 @@ 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#L1049) | [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#L1079) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L625) | | `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) | | `vllm.config.kernel.Sm70NvFp4Config` | kernel_config.sm70_nvfp4 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L277); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L236) | -| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L608) | +| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L614) | | `vllm.config.kv_events.KVEventsConfig` | kv_events_config | __post_init__: [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L50) | [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L11) | | `vllm.config.kv_transfer.KVTransferConfig` | kv_transfer_config | compute_hash: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L75); __post_init__: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L93) | [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L23) | | `vllm.config.legacy_inputs.LegacyInputs` | kernel_config.sm70_awq.legacy, kernel_config.sm70_fp8.legacy, kernel_config.sm70_moe.awq.legacy, kernel_config.sm70_moe.fp8.legacy, kernel_config.sm70_moe.mxfp4.legacy, kernel_config.sm70_moe.nvfp4.legacy, kernel_config.sm70_nvfp4.legacy, scheduler_config.sm70_inputs | value: [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L39) | [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L15) | diff --git a/docs/design/flashnext_mtp_batch_16gb_20261015.md b/docs/design/flashnext_mtp_batch_16gb_20261015.md new file mode 100644 index 0000000000..599cd81606 --- /dev/null +++ b/docs/design/flashnext_mtp_batch_16gb_20261015.md @@ -0,0 +1,247 @@ +# Flash-Next MTP4 batch scaling on four 16GB V100 GPUs + +The optimization target is the MTP4 verifier, with the qualified C1 cost of +approximately 16 ms as the reference. One verification round processes five +target rows per request. Its cost includes target forward, target sampling and +state commit; draft generation is separate. HTTP time per emitted token and +prefill throughput cannot substitute for this measurement. + +## Workload and evidence boundaries + +Integration base: `onecat/main`, `6f3557c973b4602905037e086fd4acba9b0d867f`. +The owned branch starts from the qualified deployment composition +`699c864d3c3a841b8626b657a1c8f236457ab1b5` (decode #1204, prefill #1205 and +deployment #1206). Native libraries were rebuilt from that source in an owned +CMake directory. All 12 packaged libraries have recorded SHA256 values and +no dependencies under private build directories or another worktree. + +Hardware: local GPUs 4–7, Tesla V100-SXM2-16GB, TP4, driver 580.178.04, +CUDA 12.8, Torch 2.10.0+cu128, Python 3.12.13. The model is SC117's +Qwen3.8-Flash-Next-GSQ-RCO-abliterated IQ3_S with the matching stock MTP +weights converted to Q8_0. Context capacity is 262144 total tokens, prefix +caching, aligned recurrent state, disk PLE and host QSA E4M3 history. +Experts retain the required 80% device / 20% host placement. +Sampling remains temperature 0.6, top-p 0.95, top-k 20, seeds starting at +20261014; thinking and natural EOS remain enabled. Draft length is four with +greedy draft sampling. The measured long prompt has 32774 tokens. + +The earlier 16 ms evidence used the original GSQ-RCO model on 9683 and kept +compressed indexer history on device. The first local batch diagnostic uses +host compressed indexer history and a 512-token prefill chunk to fit the +larger state pool. Its absolute C1 cost must not be substituted for the +qualified 16 ms result. The original-layout C1 control is measured separately. + +The owned-source original-layout C1 node capture confirms this distinction: +36 interior complete rounds have a 16.009 ms median and 16.234 ms mean, with a +16.243 ms TP4 envelope mean. Trace service on the critical rank is 3.295 ms +dense/shared projections, 2.882 ms expert projections, 2.322 ms HC/fused +communication and 1.985 ms QSA history/attention. The inactive-capture HTTP +control emits 143.1 tokens/s in its common decode interval with acceptance +length 2.838. Nsight injection is still present in that control; it is not a +fresh uninstrumented production result. Full per-round data and exclusions +remain in `c1-original-layout-node-breakdown.json`. + +## Initial verifier diagnostic + +Synchronized worker events retain all four TP ranks. Align the same step +ordinal, require `tokens=5*C` and `drafts=4*C`, and select the rank with the +largest verifier event sum. Stage values below come from that same rank. +Event profiling adds synchronization and is diagnostic evidence, not an +unprofiled production speed result. + +| Actual concurrent requests | Complete rounds | Verifier mean ms | Target forward ms | Sampling ms | State commit ms | Draft ms | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| 1, repetition 0 | 112 | 27.251 | 26.829 | 0.405 | 0.017 | 4.707 | +| 1, repetition 1 | 112 | 27.247 | 26.824 | 0.405 | 0.017 | 4.706 | +| 2, repetition 0 | 93 | 69.715 | 69.262 | 0.433 | 0.020 | 7.230 | +| 2, repetition 1 | 97 | 69.580 | 69.127 | 0.433 | 0.019 | 7.232 | +| 4, first actual overlap | 32 | 102.197 | 101.650 | 0.524 | 0.024 | 10.583 | + +The C4 row is from the first warmup batch's actual four-request overlap. Its +later cached batch never reached M20, and the HTTP common-window check failed. +That failed attempt is retained and cannot be reported as sustained C4. +C2 verifier inflation is about 2.56x; C4 is about 3.75x. Nearly all of it is in +target forward. GPU kernel attribution and paired candidate timings are needed +before selecting a performance default. + +## Full-width MTP4 batch measurements + +The compact-state/tiled-kernel candidate with grouped shared-key scoring +produced sustained full-width verifier rounds at all five concurrency levels. +The existing scheduler option `--mixed-prefill-min-tokens 512` was fixed for +this pass: its adaptive floor of 16 previously starved new prefills, so C6/C8 +never overlapped. This is a lab configuration change, not a new scheduler +default. Only rounds with five target rows and four draft rows per request +are included below. All four ranks must be present. + +| Concurrency | Complete rounds | Verifier mean ms | p50 ms | p90 ms | p99 ms | Forward ms | Sampling ms | State ms | Draft ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 27.422 | 27.268 | 28.019 | 29.934 | 26.997 | 0.408 | 0.017 | 4.798 | +| 2 | 191 | 52.964 | 52.745 | 53.764 | 56.661 | 52.510 | 0.435 | 0.019 | 6.760 | +| 4 | 206 | 93.346 | 92.797 | 95.453 | 101.870 | 92.796 | 0.527 | 0.023 | 10.753 | +| 6 | 182 | 138.721 | 138.250 | 140.552 | 150.466 | 138.115 | 0.580 | 0.026 | 15.069 | +| 8 | 168 | 185.889 | 185.226 | 190.645 | 196.910 | 184.995 | 0.864 | 0.030 | 19.122 | + +With grouped scoring disabled, C2 was 60.521 ms and C4 was 102.246 ms under +the earlier tiling pass. The first control C2 was approximately 69.6 ms. +Those C4 controls have fewer overlap rounds and are not a paired sustained-C4 +production result. Per-request verifier share above is 27.4/26.5/23.3/23.1/23.2 +ms. It remains above the qualified 16 ms device-indexer C1 reference. + +Node attribution localizes the remaining cost to compressed-indexer scoring, +rather than main selected-history attention. In the short retained trace, +C1/C2 indexer GPU service is 13.644/28.931 ms on the critical verifier rank; +the native shared-key reader accounts for 13.008/28.292 ms. Main selected +history attention is only 0.663/1.033 ms. These are traced service sums from +64 C1 and seven C2 rounds, not the event-profile absolute timings above. +The trace retains a separate wall residual and incomplete-round exclusions. + +## Candidate changes and retained fallbacks + +* QSA prefill staging now checks per-request query length, rather than treating + every batch above 32 total rows as prefill. MTP4 C8 is M40 with query length + five. The existing native reader handles its 32+8 fragments. True prefill + retains staging; unsupported direct inputs retain the protected reader. +* `sm70_hcx_batch_max_m` defaults to 8. An explicit limit up to 40 uses existing + native eight-row HCX tiles. The limit belongs to each module, not mutable + shared runtime policy. Larger batches keep projection/reduction fallbacks. +* `sm70_gguf.small_m_dp4a_max_m` defaults to 20. An explicit limit up to 40 + uses existing native experts in fragments of at most 20 rows. Fragment size + also respects the cache's resident slots and miss ring. Each fragment finishes + before the next resolution can evict its weights. Batches above the configured + limit retain TurboMind. The opt-in also affects prefill tails of 21–40 rows + and therefore requires numerical and model-quality evidence. +* `qsa_host_kv_state_compact` defaults to false. The bounded host-KV allocator + gives each recurrent group independent device banks and distinct scheduler ID + ranges. The block tables and zeroer translate those IDs into local bank IDs; + null ID zero remains zero. History IDs, prefix hashes and reference lifetimes + remain global. PP or CP above one is explicitly rejected. This avoids charging + every layer for every recurrent group's IDs. For C8, the tested layout saves + 463.46 MiB per rank, rather than the initially estimated 1.6 GiB. +* `sm70_qsa_shared_key_batch` defaults to false. It reuses the existing + single-request native scorer for each request in a small batch, then keeps + shared batched top-k/expansion. Admission requires FP16 H4/D128 on SM70, + two to eight queries per request, valid CPU query starts and at most 64 rows. + All requests are admitted before any scorer executes; other layouts retain + the paged scorer. No new native provider or kernel is introduced. +* `sm70_qsa_indexer_single_rank` is an experimental opt-in, default false. + It binds only for SM70 FP16 H4/KV1/D128, host compressed-indexer storage and + TP4/PP1/DCP1/PCP1. With one to 64 rows, rank zero selects indices and the + existing TP communicator broadcasts them. Every rank still writes its own + raw/compressed keys. Larger batches retain independent selection. Device + indexer history, including the qualified C1 layout, remains on its existing + path. The 32K task gate passes; fresh ordinary-worker paired performance is pending. + +These are explicit opt-ins until the paired end-to-end gates pass. No kernel +sidecars, copied extensions or preloads are required. + +## Validation and negative results + +The five QSA GPU cases cover M32, M33, M40, M64 and M65, including changed KV +pages, selection indices and lengths under graph replay. They pass against +the protected reference reader. The compact pool tests cover C1/C2/C4/C6/C8, +the adjacent insufficient capacity, unique per-group quotas, null IDs and +prefix-cache lifetime. CPU cache/dispatch tests pass (46 cases), and typed +policy/CPU block-table tests pass (12 cases). + +Four-card GPU tests pass (23 cases): compact-state zeroing with guarded foreign +IDs and changed graph inputs; block-table request permutation and padding; +tiled GGUF experts against independently dequantized Q8 oracles at M20/M21/M30/M40; +and compiled HCX payload ownership at M5/M8/M9/M10/M20/M30/M40/M41, including +changed inputs and final mixing fallback. GGUF error bounds are rtol/atol 0.003; +HCX retains the existing 0.002 relative-norm allowance. These operator tests +do not establish full-model quality or performance. + +Grouped scoring passes 14 GPU cases with device and mapped-host histories, +C2/C4/C6/C8, request reorder, missing pages, zero lengths and padded rows. +Changed-input graph replay uses rtol 0.003/atol 0.0002 for scorer reductions; +expanded selection matches the independently invoked native request scorer +exactly in the tested cases. The single-rank prototype passes a four-rank +graph test at M1/M5/M10/M20/M30/M40/M64/M65: only rank zero executes selection +through M64, while M65 falls back on all ranks. Changed requests/pages/lengths +match independent selection exactly. Typed policy/hash tests pass too. + +An aligned main-history cache microbenchmark rejected increasing hot capacity +from 32768 to 49152: both are about 0.219 ms for the representative expanded +indices. An earlier unaligned synthetic selection produced about 10 ms, but +does not represent the indexer's four-token expansion and is not model evidence. + +Combined synchronized worker events and Nsight node tracing failed twice +(one CUDA launch/Xid43, one CUDA-event recording segmentation fault during +mixed prefill). All five event-only measurement batches completed before the +second trace failure. Subsequent event and node passes are separate; partial +traces are used only for attribution. Those failures are not discarded or +reported as successful full-width trace coverage. + +The first compact C8 start captured M5/M10/M20/M30/M40 but its hand-built +startup warmup allocated IDs outside the new state-group ranges. Warmup now +uses the existing BlockPool and respects each group's capacity; 17 warmup +tests pass. The next start exposed an active-M33/padded-M40 output gate shape +error in the existing host QSA wrapper. Active gates are now sliced with the +other active inputs; six CPU wrapper cases cover empty input, M3/M5, M7/M10 +and M33/M40 padding. These failed starts remain separate from performance data. + +Rejected unmodified C8 starts include 1024-token and 512-token chunks that ran +out of device memory. A 64-token chunk with M30 graphs instead entered +TurboMind prefill and attempted an event dependency on an uncaptured stream. +This is distinct from CUDA IPC handle failures; #1208 is not a prerequisite +for that reproducer. A later M5-only C8 start also ran out of memory. Raw failed +logs and contracts are preserved; none is a speed baseline. + +## Reproduction and retained artifacts + +Run the service with the qualified sampling/MTP/KV settings and an owned GPU +group. Use `VLLM_SM70_MTP_PROFILE=1` and +`VLLM_SM70_MTP_PROFILE_INTERVAL=1` only for the diagnostic phase pass. +The benchmark requires a JSON list of prompt cases with `name` and `prompt`. + +```bash +python benchmarks/flashnext_prefill/measure_mtp_concurrency.py \ + --base-url http://127.0.0.1:18041 --cases /path/to/prompt-cases.json \ + --concurrency 1 2 4 6 8 --service-log /path/to/worker.log \ + --out /path/to/requests.json +python benchmarks/flashnext_prefill/analyze_mtp_verifier.py \ + /path/to/worker.log /path/to/requests.json /path/to/verifier.json +``` + +For node attribution, launch the same service under Nsight with +`--cuda-graph-trace=node:host-only --capture-range=cudaProfilerApi` and +`--profiler-config '{"profiler":"cuda"}'`. The optional diagnostic worker +`mtp_trace_worker.TraceWorker` annotates target and draft graph submissions +after graph capture. Capture only short warmed generation intervals. Report +GPU service sums, same-rank verifier wall and TP4 envelope separately; node +tracing is composition evidence, not an accepted absolute latency. + +Local raw artifacts are retained at +`/home/ymzx/桌面/1cat-vllm/reports/flashnext-mtp-batch-20261014/`. +The local public API is stopped. The 9683 API is outside this experiment and +remains running. No shared environment packages or model files are modified. + +## Single-rank host-indexer candidate + +Only the single-rank selection option changes relative to the full-width +grouped-scoring event pass. All five measured request streams complete; +the collector performs no Nsight capture. Values remain synchronized worker +diagnostics, not accepted uninstrumented throughput. + +| Concurrency | Complete rounds | Before ms | Candidate mean ms | p50 ms | p90 ms | p99 ms | Forward ms | Sampling ms | State ms | Draft ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 27.422 | 20.979 | 20.823 | 21.678 | 23.439 | 20.557 | 0.404 | 0.017 | 4.124 | +| 2 | 197 | 52.964 | 38.099 | 37.782 | 38.939 | 43.372 | 37.648 | 0.432 | 0.019 | 5.330 | +| 4 | 169 | 93.346 | 58.831 | 58.283 | 61.249 | 65.323 | 58.283 | 0.525 | 0.023 | 7.673 | +| 6 | 168 | 138.721 | 84.812 | 84.353 | 86.517 | 91.666 | 84.198 | 0.588 | 0.026 | 10.236 | +| 8 | 178 | 185.889 | 111.006 | 110.241 | 113.423 | 123.378 | 110.115 | 0.861 | 0.030 | 12.568 | + +C2 improves a further 28.1%, C4 37.0%, C6 38.9% and C8 40.3%. +C8 is 6.94 times the approximately 16 ms device-indexer C1 reference; +its 13.876 ms per-request verifier share is not a 13.876 ms batch wall. +The candidate host-indexer C1 is 20.979 ms, so it does not meet the qualified +device-indexer C1 latency. Draft costs are separate. Eight submitted 32K retrieval/arithmetic requests, executable order-preserving +Python deduplication, automatic weather-tool calls and tool replies pass with +thinking, production sampling and natural EOS. The tool test initially required +`finish_reason=stop` for a successful tool call; that harness mistake was fixed +to admit the expected `tool_calls` completion, and tool/reply checks were rerun. +Fresh ordinary-worker paired task and emitted-token measurements are pending. +The affected final CPU suite passes 73 cases. A stale warmup fixture originally +patched the removed runner platform/environment hooks; it now exercises the +existing bound warmup policy. No runtime compatibility workaround was added. diff --git a/tests/config/test_flashnext_mtp_batch_policy.py b/tests/config/test_flashnext_mtp_batch_policy.py new file mode 100644 index 0000000000..a7bb878c81 --- /dev/null +++ b/tests/config/test_flashnext_mtp_batch_policy.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch limits are typed engine policy and survive worker serialization.""" + +import copy +import multiprocessing + +import pytest + +from vllm.config.kernel import KernelConfig + + +def _send_worker_policy(connection, policy): + connection.send(policy) + connection.close() + + +def test_mtp_batch_policy_is_owned_hashed_and_serializable(): + legacy = KernelConfig() + tiled = KernelConfig( + sm70_hcx=True, + sm70_hcx_batch_max_m=40, + qsa_host_kv_state_compact=True, + sm70_qsa_shared_key_batch=True, + sm70_gguf={"small_m_dp4a": True, "small_m_dp4a_max_m": 40}, + qsa_host_kv=True, + qsa_host_indexer_history=True, + sm70_qsa_indexer_single_rank=True, + ) + tiled.sm70_gguf.active = True + context = multiprocessing.get_context("spawn") + receiver, sender = context.Pipe(duplex=False) + worker = context.Process(target=_send_worker_policy, args=(sender, tiled)) + worker.start() + sender.close() + try: + assert receiver.poll(60), "Worker policy transfer timed out" + transferred = receiver.recv() + finally: + receiver.close() + worker.join(10) + if worker.is_alive(): + worker.terminate() + worker.join() + assert worker.exitcode == 0 + assert legacy.sm70_hcx_batch_max_m == 8 + assert legacy.sm70_gguf.small_m_dp4a_max_m == 20 + assert not legacy.qsa_host_kv_state_compact + assert transferred == tiled + assert transferred.compute_hash() == tiled.compute_hash() + for field in ( + "sm70_hcx_batch_max_m", + "qsa_host_kv_state_compact", + "sm70_qsa_shared_key_batch", + "sm70_qsa_indexer_single_rank", + ): + changed = copy.deepcopy(tiled) + setattr(changed, field, getattr(legacy, field)) + assert changed.compute_hash() != tiled.compute_hash() + transferred.sm70_gguf.small_m_dp4a_max_m = 20 + assert transferred.compute_hash() != tiled.compute_hash() + assert tiled.sm70_gguf.small_m_dp4a_max_m == 40 + + +def test_disabled_parents_keep_inactive_batch_options_out_of_hash(): + a = KernelConfig(sm70_hcx=False, sm70_qsa_shared_key=False) + b = KernelConfig( + sm70_hcx=False, + sm70_qsa_shared_key=False, + sm70_hcx_batch_max_m=40, + sm70_qsa_shared_key_batch=True, + ) + assert a.compute_hash() == b.compute_hash() + + +@pytest.mark.parametrize( + "host, history", [(False, False), (False, True), (True, False)] +) +def test_single_rank_indexer_is_inactive_without_host_history(host, history): + a = KernelConfig(qsa_host_kv=host, qsa_host_indexer_history=history) + b = KernelConfig( + qsa_host_kv=host, + qsa_host_indexer_history=history, + sm70_qsa_indexer_single_rank=True, + ) + assert a.compute_hash() == b.compute_hash() + + +@pytest.mark.parametrize("limit", [0, 41]) +def test_mtp_batch_policy_rejects_unsupported_limits(limit): + with pytest.raises(ValueError): + KernelConfig(sm70_hcx_batch_max_m=limit) + with pytest.raises(ValueError): + KernelConfig(sm70_gguf={"small_m_dp4a_max_m": limit}) diff --git a/tests/kernels/quantization/test_gguf_compact_storage.py b/tests/kernels/quantization/test_gguf_compact_storage.py index 650951d4c1..736b8a73ff 100644 --- a/tests/kernels/quantization/test_gguf_compact_storage.py +++ b/tests/kernels/quantization/test_gguf_compact_storage.py @@ -100,6 +100,8 @@ def test_original_experts_keep_fast_route_for_graph_padding(monkeypatch, m): method = object.__new__(GGUFNativeMoEMethod) method.dp4a_admitted = True + method.expert_cache = None + method.expert_trace_tokens = 0 method.weight_types = {"w1": 21, "w3": 21, "w2": 42} method.input_padding = {"w2": (32, 192)} method.q8_intermediate = True diff --git a/tests/kernels/quantization/test_gguf_mtp_batch.py b/tests/kernels/quantization/test_gguf_mtp_batch.py new file mode 100644 index 0000000000..15162be6b5 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_mtp_batch.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tiled raw expert execution against independent dequantized Q8 oracles.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from tests.kernels.quantization.test_gguf_dp4a import quantize_reference +from vllm.model_executor.layers.fused_moe import MoEActivation +from vllm.model_executor.layers.quantization.gguf_moe import GGUFNativeMoEMethod +from vllm.model_executor.layers.quantization.gguf_native import packed_tp_span +from vllm.model_executor.layers.quantization.gguf_raw import RawGGUFProjection +from vllm.transformers_utils.gguf_tensor_reader import dequantize, quant_size + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("rows", [20, 21, 30, 40]) +def test_tiled_mtp_experts_match_q8_oracle_and_changed_graph(rows): + experts, hidden, intermediate, top_k = 4, 256, 160, 2 + rng = np.random.default_rng(20261015) + + def weight(kind, n, k): + block, size = quant_size(kind) + raw = rng.integers(0, 256, (experts, n, k // block, size), dtype=np.uint8) + raw[..., :2] = ( + rng.uniform(0.001, 0.005, raw.shape[:-1]) + .astype("mtn", decoded_x, w[ids.long()]).half() + for w in (decoded_gate, decoded_up) + ] + activated = (torch.nn.functional.silu(a.float()) * b.float()).half() + q, scales, _ = quantize_reference(activated.reshape(rows * top_k, intermediate)) + decoded_intermediate = (q.float() * scales[..., None].float()).reshape_as( + activated + ) + projected = torch.einsum( + "mtk,mtnk->mtn", decoded_intermediate, decoded_down[ids.long()] + ).half() + expected = (projected.float() * probabilities[..., None]).sum(1).half() + torch.testing.assert_close(actual, expected, rtol=0.003, atol=0.003) + assert torch.isfinite(actual).all() + + check(run()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run() + for shift in (1, 3): + x.copy_(torch.randn_like(x) * 0.125) + ids.copy_((ids + shift) % experts) + probabilities.copy_(torch.softmax(torch.randn_like(probabilities), -1)) + graph.replay() + check(output) diff --git a/tests/kernels/test_sm70_hcx_compiled_payload.py b/tests/kernels/test_sm70_hcx_compiled_payload.py index 3809c5840d..8e538147de 100644 --- a/tests/kernels/test_sm70_hcx_compiled_payload.py +++ b/tests/kernels/test_sm70_hcx_compiled_payload.py @@ -46,7 +46,7 @@ def _worker_run(rank, port, results): runtime = get_hcx_runtime(torch.device("cuda", rank)) assert runtime.enabled, runtime.reason runtime.diagnostic = True - module.enable_partial_inputs("payload", runtime) + module.enable_partial_inputs("payload", runtime, batch_max_m=40) module._hcx_moe_payload = True final_module = ( GatedResidual(config, use_combine=False, prefix="final").cuda().half() @@ -54,11 +54,12 @@ def _worker_run(rank, port, results): with torch.no_grad(): for parameter in final_module.parameters(): parameter.copy_(torch.randn_like(parameter) * 0.02) - final_module.enable_partial_inputs("final") + final_module.enable_partial_inputs("final", batch_max_m=40) final_module._hcx_moe_payload = True register_moe_runner( "payload", SimpleNamespace( + _hcx_batch_max_m=40, _maybe_sm70_moe_sum2_allreduce=lambda a, b, n: ( tensor_model_parallel_all_reduce_sum2(a, b) ), @@ -110,7 +111,7 @@ def ple_snapshots(x): ple_layer._PLE_DIAGNOSTIC_BUFFERS["after"], ple_result ) torch.testing.assert_close(ple_scratch, ple_input * 17) - for rows in (20, 5): + for rows in (5, 8, 9, 10, 20, 30, 40, 41): torch.manual_seed(1) hidden = (torch.randn(rows, 10240, device="cuda") * 0.5).half() injection = torch.randn(rows, 4, device="cuda").half() diff --git a/tests/models/qwen4_exp/test_device_kv_attention.py b/tests/models/qwen4_exp/test_device_kv_attention.py index 371e205ef5..57f1d676d1 100644 --- a/tests/models/qwen4_exp/test_device_kv_attention.py +++ b/tests/models/qwen4_exp/test_device_kv_attention.py @@ -45,6 +45,99 @@ def decoded_history(state): return (values * scales[..., None]).half() +@pytest.mark.parametrize("rows", [32, 33, 40, 64, 65]) +def test_batched_verifier_does_not_stage_prefill_history(rows, monkeypatch): + """Batch size must not put decode into a CPU-planned prefill route.""" + from vllm.models.qwen4_exp.nvidia.qsa import Qwen4ExpQSAAttention + + torch.manual_seed(rows) + device = torch.device("cuda:0") + state = HostQSAKV( + 8, + 256, + 256, + device, + width=65, + hot_tokens=128, + rows=8, + direct_device=True, + hot_fp8=True, + ) + layer = SimpleNamespace( + host_kv=state, + host_kv_prefill_stage_tokens=131072, + host_kv_decode_query_len=5, + host_kv_is_draft=False, + ) + key = torch.randn(2048, 1, 256, device=device, dtype=torch.float16) + value = torch.randn_like(key) + slots = torch.arange(2048, device=device) + state.write(key, value, slots) + query = torch.randn(rows, 6, 256, device=device, dtype=torch.float16) + gate = torch.randn_like(query) + indices = torch.arange(65, device=device, dtype=torch.int32).repeat(rows, 1) + table = torch.arange(8, device=device, dtype=torch.int32).view(8, 1) + requests = torch.arange(rows, device=device, dtype=torch.int32) % 8 + positions = torch.full((rows,), 100, device=device, dtype=torch.int64) + lengths = torch.full((8,), 200, device=device, dtype=torch.int32) + actual, reference = torch.empty_like(query), torch.empty_like(query) + metadata = SimpleNamespace(max_query_len=5) + + def forbidden(*args, **kwargs): + raise AssertionError("decode attempted a host prefill staging plan") + + monkeypatch.setattr(state, "stage_plan", forbidden) + + def run(): + Qwen4ExpQSAAttention.host_kv_forward( + layer, + query, + indices, + table, + requests, + positions, + lengths, + actual, + gate, + metadata, + ) + + def check(): + for start in range(0, rows, state.resolve_rows): + stop = min(start + state.resolve_rows, rows) + host_qsa_attention( + query[start:stop], + state, + indices[start:stop], + table, + requests[start:stop], + positions[start:stop], + lengths, + reference[start:stop], + gate[start:stop], + ) + torch.accelerator.synchronize() + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, reference, rtol=0.003, atol=0.0002) + + run() + check() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + # Reorder request pages, change selections/causal lengths and rewrite KV. + # Captured replay must observe live metadata, rather than a captured plan. + for length in (63, 177): + table.copy_(table.flip(0)) + lengths.fill_(length) + positions.fill_(length - 1) + indices.add_(3) + query.mul_(0.75) + state.write(key * 0.5, value * 0.5, slots) + graph.replay() + check() + + @pytest.mark.parametrize("rows", [1, 5, 20]) @pytest.mark.parametrize("dtype", [torch.uint8, torch.float16]) @pytest.mark.parametrize("strided", [False, True]) diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py index 01b89e75f0..326c987b97 100644 --- a/tests/models/qwen4_exp/test_qsa_cache.py +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -266,20 +266,27 @@ def test_qsa_e4m3_forward_cannot_bypass_scale_finalization() -> None: ) -@pytest.mark.parametrize("active", [0, 3, 5]) -def test_host_qsa_only_initializes_graph_padding(active: int) -> None: +@pytest.mark.parametrize( + "active,padded", [(0, 5), (3, 5), (5, 5), (7, 10), (33, 40), (40, 40)] +) +def test_host_qsa_only_initializes_graph_padding(active: int, padded: int) -> None: impl = object.__new__(Qwen4ExpQSAFlashAttentionImpl) impl.alibi_slopes = None impl.sinks = None impl.sliding_window = (-1, -1) - query = torch.ones(5, 2, 8, dtype=torch.float16) + query = torch.ones(padded, 2, 8, dtype=torch.float16) + gate = torch.ones(padded, 16, dtype=torch.float16) output = torch.full_like(query, float("nan")) calls = [] - def host_forward(q, indices, table, requests, positions, lengths, out, gate): + def host_forward( + q, indices, table, requests, positions, lengths, out, gate, metadata + ): # The producer owns the active output. It must not depend on a fill. assert torch.isnan(out).all() assert out.shape[0] == active + assert gate.view_as(q).shape[0] == active + assert metadata.num_actual_tokens == active out.copy_(q * 2) calls.append(active) @@ -287,7 +294,7 @@ def host_forward(q, indices, table, requests, positions, lengths, out, gate): SimpleNamespace( host_kv_enabled=True, host_kv_forward=host_forward, - topk_indices_buffer=torch.zeros(5, 4, dtype=torch.int32), + topk_indices_buffer=torch.zeros(padded, 4, dtype=torch.int32), ), query, query[:, :1], @@ -295,9 +302,10 @@ def host_forward(q, indices, table, requests, positions, lengths, out, gate): torch.empty(0), SimpleNamespace(num_actual_tokens=active, block_table=torch.zeros(1, 1)), output, - torch.zeros(5, dtype=torch.int32), - query_positions=torch.arange(5), - sequence_lengths=torch.tensor([5]), + torch.zeros(padded, dtype=torch.int32), + query_positions=torch.arange(padded), + sequence_lengths=torch.tensor([padded]), + output_gate=gate, ) assert result is output assert torch.equal(output[:active], query[:active] * 2) diff --git a/tests/models/qwen4_exp/test_qsa_indexer_single_rank.py b/tests/models/qwen4_exp/test_qsa_indexer_single_rank.py new file mode 100644 index 0000000000..6b454e1da6 --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_indexer_single_rank.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Replicated selection uses one reader, with live paged state on graph replay.""" + +import socket +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp + + +def _worker(rank, port, results): + import traceback + + try: + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed import ( + get_tp_group, + graph_capture, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.models.qwen4_exp.nvidia.indexer_qsa import QSAIndexer + from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + + torch.accelerator.set_device_index(rank) + with set_current_vllm_config(VllmConfig()): + init_distributed_environment( + 4, rank, f"tcp://127.0.0.1:{port}", rank, "nccl" + ) + initialize_model_parallel(4, 1) + indexer = QSAIndexer.__new__(QSAIndexer) + torch.nn.Module.__init__(indexer) + indexer.token_topk, indexer.compress_ratio = 2048, 4 + indexer.shared_key_scoring = True + indexer.shared_key_batch_scoring = True + indexer.selection_tp_group = get_tp_group() + report = [] + for rows in (1, 5, 10, 20, 30, 40, 64, 65): + requests = (rows + 4) // 5 + torch.manual_seed(20261015 + rows) + host = ( + torch.randn(requests * 16, 64, 1, 128, dtype=torch.float16) * 0.125 + ).pin_memory() + cache = get_accelerator_view_from_cpu_tensor(host) + q = ( + torch.randn(rows, 4, 128, device="cuda", dtype=torch.float16) + * 0.125 + ) + table = torch.arange( + requests * 16, device="cuda", dtype=torch.int32 + ).view(requests, 16) + mapping = torch.arange(rows, device="cuda", dtype=torch.int32) // 5 + lengths = torch.full( + (requests,), 4096, device="cuda", dtype=torch.int32 + ) + positions = torch.full((rows,), 4095, device="cuda", dtype=torch.int64) + starts = torch.tensor([*range(0, rows, 5), rows], dtype=torch.int32) + metadata = SimpleNamespace( + block_table=table, + token_to_req=mapping, + logical_positions=positions, + seq_lens=lengths, + query_start_loc_cpu=starts, + max_seq_len=4096, + ) + out = torch.empty(rows, 2051, dtype=torch.int32, device="cuda") + expected = torch.empty_like(out) + indexer._select_local_tokens(q, cache, metadata, expected) + indexer._select_tokens(q, cache, metadata, out) + calls: list[int] = [] + local_select = indexer._select_local_tokens + + def counted(*args, calls=calls, local_select=local_select): + calls.append(1) + return local_select(*args) + + indexer._select_local_tokens = counted + graph = torch.cuda.CUDAGraph() + with ( + graph_capture(device=torch.device("cuda", rank)) as context, + torch.cuda.graph(graph, stream=context.stream), + ): + result = indexer._select_tokens(q, cache, metadata, out) + indexer._select_local_tokens = local_select + assert len(calls) == int(rows > 64 or rank == 0), (rank, rows, calls) + assert result.data_ptr() == out.data_ptr() + for shift in (1, 2): + q.mul_(0.75) + table.copy_(table.roll(shift, 0)) + lengths[0] = 0 + positions.sub_(47) + mapping[-1] = -1 + graph.replay() + indexer._select_local_tokens(q, cache, metadata, expected) + torch.accelerator.synchronize() + torch.testing.assert_close(result, expected, rtol=0, atol=0) + report.append(rows) + results.put((rank, report)) + except BaseException: + results.put((rank, {"error": traceback.format_exc()})) + raise + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_single_rank_indexer_graph_replay_and_capacity_fallback(): + if torch.accelerator.device_count() != 4 or torch.cuda.get_device_capability() != ( + 7, + 0, + ): + pytest.skip("Four visible SM70 GPUs required") + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + context = mp.get_context("spawn") + results = context.Queue() + workers = [ + context.Process(target=_worker, args=(r, port, results)) for r in range(4) + ] + for worker in workers: + worker.start() + try: + reports = [results.get(timeout=180) for _ in workers] + finally: + for worker in workers: + worker.join(timeout=10) + if worker.is_alive(): + worker.terminate() + worker.join() + assert all(isinstance(report, list) for _, report in reports), reports + assert all(worker.exitcode == 0 for worker in workers) + assert sorted(rank for rank, _ in reports) == list(range(4)) diff --git a/tests/models/qwen4_exp/test_qsa_shared_key_batch.py b/tests/models/qwen4_exp/test_qsa_shared_key_batch.py new file mode 100644 index 0000000000..3544d5e9b8 --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_shared_key_batch.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Grouped shared-key scoring preserves paged request identity and live state.""" + +import pytest +import torch + +from vllm.models.qwen4_exp.nvidia.ops.qsa import qsa_mqa_paged, qsa_select_paged_tokens +from vllm.models.qwen4_exp.nvidia.ops.qsa_shared_key import ( + maybe_shared_key_batch_scores, +) +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="SM70 CUDA required" +) + + +@pytest.mark.parametrize("requests", [2, 4, 6, 8]) +@pytest.mark.parametrize("host", [False, True]) +def test_shared_key_batch_scores_and_replay_match_retained_paged_reader(requests, host): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + rows, page = requests * 5, 204 + torch.manual_seed(20261015 + requests) + data = torch.randn(requests * 3, page, 1, 128, dtype=torch.float16) * 0.125 + if host: + data = data.pin_memory() + cache = get_accelerator_view_from_cpu_tensor(data) + else: + cache = data.cuda() + q = torch.randn(rows, 4, 128, dtype=torch.float16, device="cuda") * 0.125 + table = torch.full((requests, 4), -1, dtype=torch.int32, device="cuda") + table[:, :3] = torch.arange(requests * 3, dtype=torch.int32, device="cuda").view( + requests, 3 + ) + table[-1, 1] = -1 + req = torch.arange(requests, dtype=torch.int32, device="cuda").repeat_interleave(5) + positions = torch.full((rows,), 2399, dtype=torch.int64, device="cuda") + lengths = ( + torch.arange(1, requests + 1, dtype=torch.int32, device="cuda") * 128 + 1024 + ) + starts = torch.arange(0, rows + 1, 5, dtype=torch.int32) + + def run(enabled): + return qsa_mqa_paged( + q, + cache, + table, + req, + positions, + lengths, + 4, + shared_key_scoring=enabled, + shared_key_query_start_loc_cpu=starts if enabled else None, + ) + + def check(actual, visible): + expected, expected_visible = run(False) + torch.testing.assert_close(visible, expected_visible, rtol=0, atol=0) + for row, n in enumerate(visible.tolist()): + a, b = actual[row, :n], expected[row, :n] + assert torch.equal(torch.isfinite(a), torch.isfinite(b)) + finite = torch.isfinite(b) + torch.testing.assert_close(a[finite], b[finite], rtol=0.003, atol=0.0002) + + output = run(True) + check(*output) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = run(True) + for shift in (1, 2): + q.mul_(0.75) + table[:, :3].copy_(table[:, :3].roll(shift, 0)) + positions.sub_(47) + lengths[0] = 0 + req[-1] = -1 + graph.replay() + check(*output) + + +@pytest.mark.parametrize("requests", [2, 8]) +def test_grouped_selection_matches_independent_native_requests(requests): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + rows = requests * 5 + torch.manual_seed(20261016) + q = torch.randn(rows, 4, 128, device="cuda", dtype=torch.float16) * 0.125 + cache = ( + torch.randn(requests * 16, 64, 1, 128, device="cuda", dtype=torch.float16) + * 0.125 + ) + table = torch.arange(requests * 16, dtype=torch.int32, device="cuda").view( + requests, 16 + ) + mapping = torch.arange( + requests, dtype=torch.int32, device="cuda" + ).repeat_interleave(5) + positions = torch.full((rows,), 4095, device="cuda", dtype=torch.int64) + lengths = torch.full((requests,), 4096, device="cuda", dtype=torch.int32) + starts = torch.arange(0, rows + 1, 5, dtype=torch.int32) + + def run(): + return qsa_select_paged_tokens( + q, + cache, + table, + mapping, + positions, + lengths, + 2048, + 4, + query_start_loc_cpu=starts, + shared_key_scoring=True, + shared_key_batch_scoring=True, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = run() + for shift in (1, 2): + table.copy_(table.roll(shift, 0)) + q.mul_(0.75) + lengths[0] = 0 + graph.replay() + expected = torch.cat( + [ + qsa_select_paged_tokens( + q[5 * i : 5 * i + 5], + cache, + table[i : i + 1], + mapping[5 * i : 5 * i + 5] - i, + positions[5 * i : 5 * i + 5], + lengths[i : i + 1], + 2048, + 4, + shared_key_scoring=True, + ) + for i in range(requests) + ] + ) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("starts", [[0, 1, 10], [0, 9, 10], [0, 5, 9], [0, 7, 6, 10]]) +def test_shared_key_batch_declines_unqualified_host_layouts(starts): + q = torch.zeros(10, 4, 128, dtype=torch.float16, device="cuda") + table = torch.zeros(len(starts) - 1, 2, dtype=torch.int32, device="cuda") + assert ( + maybe_shared_key_batch_scores( + q, + q.new_zeros(2, 64, 1, 128), + table, + torch.zeros(10, dtype=torch.int32, device="cuda"), + torch.zeros(10, dtype=torch.int64, device="cuda"), + torch.zeros(len(starts) - 1, dtype=torch.int32, device="cuda"), + 4, + 128, + 128**0.5, + torch.tensor(starts, dtype=torch.int32), + ) + is None + ) diff --git a/tests/v1/core/test_qwen4_exp_kv_cache.py b/tests/v1/core/test_qwen4_exp_kv_cache.py index c887a364ab..91deb4416d 100644 --- a/tests/v1/core/test_qwen4_exp_kv_cache.py +++ b/tests/v1/core/test_qwen4_exp_kv_cache.py @@ -262,24 +262,20 @@ def test_host_state_pool_rejects_insufficient_active_capacity(): _bounded_host_cache_config(state_blocks=2) -@pytest.mark.parametrize( - "state_blocks, device_budget_mib, host_indexer", - [(29, 504, False), (32, 560, False), (29, 274, True)], -) -def test_tp4_mtp4_full_context_pool_fits_bounded_device_budget( - state_blocks, device_budget_mib, host_indexer -): +def _tp4_mtp4_cache_groups(state_blocks, host_indexer, requests=1, compact=False): config = _vllm_config() config.model_config = SimpleNamespace( max_model_len=262144, get_num_kv_heads=lambda _: 1, get_total_num_hidden_layers=lambda: 49, ) - config.scheduler_config.max_num_seqs = 1 + config.scheduler_config.max_num_seqs = requests config.scheduler_config.max_num_batched_tokens = 512 config.cache_config.mamba_cache_mode = "align" config.cache_config.enable_prefix_caching = True - config.kernel_config = SimpleNamespace(qsa_host_kv_state_blocks=state_blocks) + config.kernel_config = SimpleNamespace( + qsa_host_kv_state_blocks=state_blocks, qsa_host_kv_state_compact=compact + ) specs = {} for layer in range(49): prefix = f"model.layers.{layer}" if layer < 48 else "mtp.layers.48" @@ -325,6 +321,17 @@ def test_tp4_mtp4_full_context_pool_fits_bounded_device_budget( tp_replicated=True, ) groups = get_kv_cache_groups(config, specs) + return config, groups + + +@pytest.mark.parametrize( + "state_blocks, device_budget_mib, host_indexer", + [(29, 504, False), (32, 560, False), (29, 274, True)], +) +def test_tp4_mtp4_full_context_pool_fits_bounded_device_budget( + state_blocks, device_budget_mib, host_indexer +): + config, groups = _tp4_mtp4_cache_groups(state_blocks, host_indexer) # Includes target/draft history, compressor and speculative recurrent state. cache = get_kv_cache_config_from_groups(config, groups, device_budget_mib * 1024**2) device_bytes = sum(t.size for t in cache.kv_cache_tensors if not t.host_backed) @@ -341,6 +348,65 @@ def test_tp4_mtp4_full_context_pool_fits_bounded_device_budget( assert get_max_concurrency_for_kv_cache_config(config, cache) >= 1 +@pytest.mark.parametrize("requests", [1, 2, 4, 6, 8]) +def test_compact_state_banks_preserve_history_and_concurrency(requests): + state_blocks = 7 * requests + 1 + config, groups = _tp4_mtp4_cache_groups( + state_blocks, True, requests=requests, compact=True + ) + cache = get_kv_cache_config_from_groups(config, groups, 4 * 1024**3) + legacy_config, legacy_groups = _tp4_mtp4_cache_groups( + 28 * requests + 1, True, requests=requests + ) + legacy = get_kv_cache_config_from_groups(legacy_config, legacy_groups, 4 * 1024**3) + assert cache.num_blocks == legacy.num_blocks + assert cache.logical_state_blocks == legacy.device_state_blocks + assert cache.device_state_group_offsets == tuple(7 * requests * i for i in range(4)) + assert get_max_concurrency_for_kv_cache_config(config, cache) >= requests + device_bytes = sum(t.size for t in cache.kv_cache_tensors if not t.host_backed) + old_device_bytes = sum(t.size for t in legacy.kv_cache_tensors if not t.host_backed) + assert device_bytes < old_device_bytes + assert sum(t.size for t in cache.kv_cache_tensors if t.host_backed) == sum( + t.size for t in legacy.kv_cache_tensors if t.host_backed + ) + state_names = [ + name for i in cache.device_state_group_ids for name in groups[i].layer_names + ] + banks = [t for t in cache.kv_cache_tensors if t.physical_num_blocks] + assert sorted(name for t in banks for name in t.shared_by) == sorted(state_names) + assert all(len(t.shared_by) == 1 for t in banks) + # Exact capacity succeeds; the adjacent smaller pool is rejected. + config.kernel_config.qsa_host_kv_state_blocks -= 1 + with pytest.raises(ValueError, match=f"at least {state_blocks}"): + get_kv_cache_config_from_groups(config, groups, 4 * 1024**3) + + +def test_compact_state_queues_keep_unique_ids_and_independent_quotas(): + pool = BlockPool( + 16, + True, + 16, + device_state_blocks=4, + device_state_group_ids=(1, 2), + device_state_group_offsets=(0, 3), + ) + first = pool.get_new_blocks(3, 1) + second = pool.get_new_blocks(3, 2) + history = pool.get_new_blocks(3, 0) + assert [b.block_id for b in first] == [1, 2, 3] + assert [b.block_id for b in second] == [4, 5, 6] + assert [b.block_id for b in history] == [7, 8, 9] + assert not pool.can_allocate_group_blocks({1: 1}) + pool.free_blocks(first) + assert pool.can_allocate_group_blocks({1: 3}) + assert not pool.can_allocate_group_blocks({2: 1}) + pool.touch(first[:1]) + assert pool.get_num_free_blocks(1) == 2 + pool.free_blocks(first[:1] + second + history) + assert pool.get_num_free_blocks() == 15 + assert pool.reset_prefix_cache() + + def test_partitioned_block_pool_keeps_prefix_lifetime_and_quota(): init_none_hash(sha256) pool = BlockPool(12, True, 16, device_state_blocks=4, device_state_group_ids=(1, 2)) diff --git a/tests/v1/worker/test_cpu_block_table_dispatch.py b/tests/v1/worker/test_cpu_block_table_dispatch.py index 9518147036..9cd7de0fb7 100644 --- a/tests/v1/worker/test_cpu_block_table_dispatch.py +++ b/tests/v1/worker/test_cpu_block_table_dispatch.py @@ -10,6 +10,24 @@ from vllm.v1.worker import block_table +def test_compact_state_block_table_preserves_null_and_history_identity(): + table = block_table.BlockTable(16, 2, 4, 8, False, torch.device("cpu"), 16, 1) + table.state_block_id_offset = 3 + table.device_state_blocks = 4 + table.add_row([0, 4, 6], 0) + assert table.block_table.np[0, :3].tolist() == [0, 1, 3] + table.add_row([5], 1) + table.swap_row(0, 1) + assert table.block_table.np[0, 0] == 2 + assert table.block_table.np[1, :3].tolist() == [0, 1, 3] + for invalid in (3, 7, -1): + with pytest.raises(ValueError, match="State"): + table.add_row([invalid], 0) + history = block_table.BlockTable(16, 2, 4, 8, False, torch.device("cpu"), 16, 1) + history.add_row([0, 7, 9], 0) + assert history.block_table.np[0, :3].tolist() == [0, 7, 9] + + @pytest.mark.parametrize("num_reqs", [1, 2]) @pytest.mark.parametrize("num_tokens", [1, 3, 5]) def test_repeated_cpu_slot_mapping_keeps_registered_fallback( diff --git a/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py b/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py index 7612954c5d..08b0bc5023 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_kv_zero.py @@ -7,7 +7,7 @@ import torch -from vllm import envs +from vllm.platforms import current_platform from vllm.v1.worker.gpu import model_runner as mrv2 from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator @@ -34,6 +34,9 @@ def test_v2_initializes_kv_zeroer_from_all_attention_groups(monkeypatch) -> None runner.attn_groups = groups runner._kernel_block_sizes = [16, 32] runner.cache_config = SimpleNamespace(cache_dtype="auto") + runner.kv_cache_config = SimpleNamespace( + kv_cache_tensors=[], state_block_id_offsets={} + ) runner.compilation_config = SimpleNamespace( static_forward_context=static_forward_context ) @@ -50,6 +53,8 @@ def test_v2_initializes_kv_zeroer_from_all_attention_groups(monkeypatch) -> None "cache_dtype": "auto", "runner_only_attn_layers": set(), "static_forward_context": static_forward_context, + "host_backed_layers": set(), + "state_block_id_offsets": {}, } @@ -122,16 +127,11 @@ def test_v2_warms_qwen38_mtp_moe_prefill_and_decode_shapes(monkeypatch) -> None: runner._dummy_run = mock.Mock() with ( - mock.patch.object( - mrv2.current_platform, "is_device_capability", return_value=True - ), + mock.patch.object(current_platform, "is_device_capability", return_value=True), mock.patch( - "vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn." - "_warmup_sm70_qwen_gdn_causal_conv1d", - return_value=False, + "vllm.v1.worker.gpu.spec_decode.eagle.speculator.unquantized_moe_policy", + return_value=SimpleNamespace(value=lambda name: name == "mtp_tuned"), ), - mock.patch.object(envs, "VLLM_SM70_AUX_KERNEL_WARMUP", True), - mock.patch.object(envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", True), mock.patch.object(torch.accelerator, "synchronize"), ): runner._warmup_sm70_aux_kernels() diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index 3e84c9f2c9..8cec08e1db 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -93,6 +93,56 @@ def test_kernel_prefill_warmup_runs_default_batch_and_extra_profile( assert len(samples) == 4 +def test_compact_warmup_respects_group_ranges_and_physical_capacity(monkeypatch): + from vllm.v1.kv_cache_interface import map_device_state_block_ids + + state = MambaSpec( + block_size=4, + shapes=((1,),), + dtypes=(torch.float16,), + mamba_cache_mode="align", + num_speculative_blocks=0, + ) + runner = SimpleNamespace( + num_speculative_steps=4, + scheduler_config=SimpleNamespace(max_num_seqs=2, max_num_batched_tokens=128), + max_model_len=64, + compilation_config=SimpleNamespace(static_forward_context={}), + kv_cache_config=SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace(kv_cache_spec=s) + for s in (state, _full_attention_spec(), state) + ], + num_blocks=64, + device_state_blocks=4, + device_state_group_ids=(0, 2), + device_state_group_offsets=(0, 3), + ), + vllm_config=SimpleNamespace(num_lookahead_tokens=4), + model_state=SimpleNamespace(max_encoder_len=0), + is_pooling_model=False, + is_last_pp_rank=True, + model_config=SimpleNamespace(get_vocab_size=lambda: 64), + kv_connector=SimpleNamespace(set_disabled=lambda _: None), + ) + outputs: list[Any] = [] + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + warmup_kernels(runner, outputs.append, lambda _: None) + prefill, decode, cleanup = outputs + assert len(prefill.scheduled_new_reqs) == 1 + ids = prefill.scheduled_new_reqs[0].block_ids + assert ids == ([1, 2], [7], [4, 5]) + additions = decode.scheduled_cached_reqs.new_block_ids[0] + assert additions == ([3], [], [6]) + for group, offset in ((0, 0), (2, 3)): + assert map_device_state_block_ids(ids[group] + additions[group], offset, 4) == [ + 1, + 2, + 3, + ] + assert cleanup.finished_req_ids + + def _speculative_config(method: str) -> SpeculativeConfig: config = object.__new__(SpeculativeConfig) object.__setattr__(config, "method", method) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index adac54d208..c3d94e4d86 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -29,12 +29,113 @@ N_BLOCKS = 4 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compact_state_ids_zero_only_the_owning_bank_and_graph_replay(): + from vllm.v1.worker.utils import _zero_kv_blocks_kernel + + capacity = 4 + banks = [ + torch.full((capacity, 8), 7, dtype=torch.int32, device="cuda") for _ in range(3) + ] + groups = [ + SimpleNamespace( + kv_cache_spec=MambaSpec( + block_size=16, shapes=((8,),), dtypes=(torch.int32,) + ), + kv_cache_group_id=i, + layer_names=[f"state.{i}"], + ) + for i in range(3) + ] + zeroer = KVBlockZeroer(torch.device("cuda"), pin_memory=False) + zeroer.init_meta( + groups, + [16] * 3, + "auto", + set(), + { + f"state.{i}": SimpleNamespace(kv_cache=[bank]) + for i, bank in enumerate(banks) + }, + state_block_id_offsets={0: 0, 1: 3, 2: 6}, + ) + zeroer.zero_block_ids([1, 6, 8, 10, -1]) + torch.accelerator.synchronize() + for bank, local in zip(banks, (1, 3, 2)): + expected = torch.full_like(bank, 7) + expected[local] = 0 + torch.testing.assert_close(bank, expected, rtol=0, atol=0) + + addresses, strides, spans, counts, chunks, block_size, segments = zeroer._meta + ids = torch.tensor([2, 4, 9], device="cuda", dtype=torch.int64) + + def run(): + _zero_kv_blocks_kernel[(len(ids), segments, chunks)]( + addresses, + strides, + spans, + counts, + zeroer._segment_block_id_offsets, + ids, + BLOCK_SIZE=block_size, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for updated, expected_local in (([3, 5, 7], (3, 2, 1)), ([0, 10, -1], (0, 0, 0))): + for bank in banks: + bank.fill_(7) + ids.copy_(torch.tensor(updated, device="cuda")) + graph.replay() + torch.accelerator.synchronize() + for bank, local in zip(banks, expected_local): + expected = torch.full_like(bank, 7) + expected[local] = 0 + torch.testing.assert_close(bank, expected, rtol=0, atol=0) + + class _BlockFirstAttentionBackend: @staticmethod def get_kv_cache_block_dim(*args, **kwargs) -> int: return 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_compact_v2_tables_keep_group_identity_and_request_reorder_in_graph(): + from vllm.v1.worker.gpu.block_table import BlockTables + + tables = BlockTables( + [16, 16, 16], + 3, + 16, + [4, 4, 4], + torch.device("cuda"), + [16, 16, 16], + state_block_id_offsets={0: 3, 1: 0}, + device_state_blocks=4, + ) + tables.append_block_ids(0, ([0, 4, 6], [0, 1, 3], [0, 9, 12]), overwrite=True) + tables.append_block_ids(1, ([5], [2], [15]), overwrite=True) + tables.apply_staged_writes() + order = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + tables.gather_block_tables(order, 3) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = tables.gather_block_tables(order, 3) + for permutation in ([0, 1], [1, 0]): + order.copy_(torch.tensor(permutation, dtype=torch.int32, device="cuda")) + graph.replay() + torch.accelerator.synchronize() + for group, expected in zip(output, ([0, 1, 3], [0, 1, 3], [0, 9, 12])): + row = permutation.index(0) + assert group[row, :3].tolist() == expected + assert not group[2].count_nonzero() + with pytest.raises(ValueError, match="owning"): + tables.append_block_ids(2, ([7], [1], [19]), overwrite=True) + + @pytest.mark.parametrize( ("spec", "group_kernel_bs", "expected"), [ diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index bc6ce2b4d4..3af36adda5 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -572,6 +572,12 @@ class Sm70GgufConfig: small_m_dp4a: bool = True """Use Q8_1 activations and FP32 integer dots for calibrated small GGUF batches.""" + small_m_dp4a_max_m: int = Field(default=20, ge=1, le=40) + """Maximum logical batch; native fragments remain at most 20 rows. + + Values above 20 opt into tiled original-block experts, including small + prefill tails. Larger batches retain the existing TurboMind path. + """ lut4_expert_dp4a: bool = True """Admit canonical IQ4 gate/up integer dots at calibrated expert shapes.""" @@ -830,6 +836,12 @@ def resolve_sm70_rmsnorm_gated(self, *, qualified: bool) -> None: sm70_hcx: bool = False """Leave block outputs as TP partials and run all-reduce, HC combine/norm, HC down and HC up as one SM70 kernel for verification batches up to 8.""" + sm70_hcx_batch_max_m: int = Field(default=8, ge=1, le=40) + """Maximum HCX batch, executed with the qualified eight-row native tiles. + + Values above 8 opt into tiling; larger batches retain the existing + projection and communication fallback. Bound independently on each module. + """ sm70_hcx_local_schedule: bool = False """Experimentally delay HC up-weight prefetch and distribute gate-mix for @@ -870,6 +882,18 @@ def resolve_sm70_rmsnorm_gated(self, *, qualified: bool) -> None: M1..20, H6, D256. Speculative draft attention retains its existing path.""" sm70_qsa_shared_key: bool = True """Share FP16 indexer keys across M2..8 queries of one request on SM70.""" + sm70_qsa_shared_key_batch: bool = False + """Reuse that scorer per request for small verification batches. + + Requires ``sm70_qsa_shared_key`` and a valid host query-start layout. + Unqualified layouts retain the existing batched scorer. Off by default. + """ + sm70_qsa_indexer_single_rank: bool = False + """Select host-backed replicated QSA indices on TP rank zero and broadcast. + + Applies to qualified SM70 FP16 TP4/PP1/CP1 small batches. Each rank still + updates its own key caches. Other shapes retain independent selection. + """ qsa_host_kv: bool = False """Keep QSA attention history in pinned host storage on SM70.""" qsa_host_kv_dtype: Literal["fp8_e4m3", "float16"] = "fp8_e4m3" @@ -908,6 +932,12 @@ def resolve_sm70_rmsnorm_gated(self, *, qualified: bool) -> None: Zero retains the shared capacity. A nonzero pool retains active states and bounded prefix checkpoints while history IDs use a separate range. """ + qsa_host_kv_state_compact: bool = False + """Give bounded recurrent groups independent, compact device state banks. + + Requires the host-KV layout, PP1 and CP1. Scheduler block IDs remain unique; + block tables map each state group to its own physical bank. Off by default. + """ qsa_host_kv_active: bool = Field(default=False, init=False) """Whether the host QSA cache geometry has been admitted.""" qsa_host_kv_reason: str | None = Field(default=None, init=False) @@ -1081,6 +1111,16 @@ def compute_hash(self) -> str: } if not self.sm70_skinny_moe_applicable: ignored_factors.add("sm70_skinny_moe") + if not self.sm70_hcx: + ignored_factors.add("sm70_hcx_batch_max_m") + if not self.sm70_qsa_shared_key: + ignored_factors.add("sm70_qsa_shared_key_batch") + if ( + not self.qsa_host_kv + or not self.qsa_host_indexer_history + or self.qsa_host_kv_device_reference + ): + ignored_factors.add("sm70_qsa_indexer_single_rank") if not self.fused_fp16_aux_gemv_applicable: ignored_factors.add("fused_fp16_aux_gemv") if not self.qsa_auto_e4m3_active: diff --git a/vllm/model_executor/layers/quantization/gguf_moe.py b/vllm/model_executor/layers/quantization/gguf_moe.py index 3f59846d39..7b6b5d7ca8 100644 --- a/vllm/model_executor/layers/quantization/gguf_moe.py +++ b/vllm/model_executor/layers/quantization/gguf_moe.py @@ -134,6 +134,9 @@ def __init__(self, quant_config: "GGUFConfig", moe): config.kernel_config.sm70_gguf.enabled if config is not None else True ) self.small_m_dp4a = bool(config and config.kernel_config.sm70_gguf.small_m_dp4a) + self.small_m_dp4a_max_m = ( + config.kernel_config.sm70_gguf.small_m_dp4a_max_m if config else 20 + ) self.lut4_expert_dp4a = bool( config and config.kernel_config.sm70_gguf.lut4_expert_dp4a ) @@ -337,7 +340,8 @@ def process_weights_after_loading(self, layer): if self.dp4a_admitted else "original_dp4a_shape_or_operator_unavailable", "min_m": 1, - "max_m": 20, + "max_m": self.small_m_dp4a_max_m, + "native_max_m": 20, "storage": "original", "accumulation": "FP32", } @@ -376,25 +380,42 @@ def apply( w1, w3, w2 = layer.gguf_w1, layer.gguf_w3, layer.gguf_w2 if ( self.dp4a_admitted - and 1 <= tokens <= 20 + and 1 <= tokens <= getattr(self, "small_m_dp4a_max_m", 20) and layer.expert_map is None and topk_weights.dtype == torch.float32 ): + step = 20 if cache is not None: - ids = cache.resolve(self.cache_layer, ids) - return torch.ops.vllm.gguf_original_expert_dp4a( - x.contiguous(), - ids, - topk_weights.contiguous(), - w1, - w3, - w2, - self.weight_types["w1"], - self.weight_types["w2"], - self.input_padding.get("w2", (0, 0))[0], - self.q8_intermediate, - self.weight_types["w1"] == 21, - ) + capacity = min(cache.resident0, cache.ring.shape[1] - 8) + step = min(step, capacity // top_k) + if step < 1: + raise ValueError("Expert cache cannot hold one routed token") + output = torch.empty_like(x) if tokens > step else None + for start in range(0, tokens, step): + stop = min(start + step, tokens) + fragment_ids = ids[start:stop].contiguous() + if cache is not None: + # Finish each fragment before resolving the next one, which + # may evict its slots. The existing model stream owns both. + fragment_ids = cache.resolve(self.cache_layer, fragment_ids) + fragment = torch.ops.vllm.gguf_original_expert_dp4a( + x.contiguous() if output is None else x[start:stop].contiguous(), + fragment_ids, + topk_weights[start:stop].contiguous(), + w1, + w3, + w2, + self.weight_types["w1"], + self.weight_types["w2"], + self.input_padding.get("w2", (0, 0))[0], + self.q8_intermediate, + self.weight_types["w1"] == 21, + ) + if output is None: + return fragment + output[start:stop].copy_(fragment) + assert output is not None + return output if cache is not None and cache.tm is not None: return cache.prefill_moe( self.cache_layer, x.contiguous(), ids, topk_weights diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py index 13af9ad2de..42046f693f 100644 --- a/vllm/models/qwen4_exp/nvidia/hyperconnection.py +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -204,9 +204,10 @@ def mix( return hidden_states, block_input, injection - def enable_partial_inputs(self, name: str, runtime=None) -> None: + def enable_partial_inputs(self, name: str, runtime=None, *, batch_max_m=8) -> None: """Block outputs reaching this module are TP partials from now on.""" self._partial_inputs = True + self._hcx_batch_max_m = batch_max_m self._hcx_name = name _PARTIAL_MODULES[name] = self if runtime is not None and self.use_combine: @@ -230,7 +231,7 @@ def _reduce_partial(self, block_output: torch.Tensor) -> torch.Tensor: if not getattr(self, "_partial_inputs", False): return block_output first, second = self._partial_pair(block_output) - if first.shape[0] > 8: + if first.shape[0] > getattr(self, "_hcx_batch_max_m", 8): return first if second is not None: from vllm.distributed import tensor_model_parallel_all_reduce_sum2 @@ -316,7 +317,8 @@ def _hcx_combine_and_mix( runtime = getattr(module, "_hcx", None) oproj = getattr(module, "_hcx_oproj", None) block_output, secondary = module._partial_pair(block_output) - if runtime is not None and 0 < block_output.shape[0] <= 8: + maximum = getattr(module, "_hcx_batch_max_m", 8) + if runtime is not None and 0 < block_output.shape[0] <= maximum: diagnostic = ( {"snapshot_name": name} if getattr(runtime, "diagnostic", False) else {} ) @@ -330,16 +332,17 @@ def _hcx_combine_and_mix( module._hcx_up, None if oproj is None else oproj[1:], secondary=secondary, + batch_max_m=maximum, **diagnostic, ) from vllm.distributed import tensor_model_parallel_all_reduce - if oproj is not None and block_output.shape[0] <= 8: + if oproj is not None and block_output.shape[0] <= maximum: # Deferred projection: the block output carries the o-proj input. layer, k = oproj[0], oproj[1] block_output, _ = layer(block_output[:, :k].contiguous()) reduced = block_output - if block_output.shape[0] <= 8: + if block_output.shape[0] <= maximum: if secondary is not None: from vllm.distributed import tensor_model_parallel_all_reduce_sum2 diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index ababa83d6f..67f98847e0 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -15,6 +15,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope +from vllm.platforms import current_platform from vllm.transformers_utils.configs.qwen4_exp import ( Qwen4ExpTextConfig, ) @@ -130,6 +131,29 @@ def __init__( # while continuing to update the QSA side cache. self.skip_topk = False self.shared_key_scoring = vllm_config.kernel_config.sm70_qsa_shared_key + self.shared_key_batch_scoring = ( + vllm_config.kernel_config.sm70_qsa_shared_key_batch + ) + self.selection_tp_group = None + parallel = vllm_config.parallel_config + if ( + vllm_config.kernel_config.sm70_qsa_indexer_single_rank + and vllm_config.kernel_config.qsa_host_kv + and vllm_config.kernel_config.qsa_host_kv_active + and vllm_config.kernel_config.qsa_host_indexer_history + and not vllm_config.kernel_config.qsa_host_kv_device_reference + and vllm_config.model_config.dtype == torch.float16 + and current_platform.is_device_capability(70) + and (self.index_n_heads, self.index_kv_heads, self.index_head_dim) + == (4, 1, 128) + and parallel.tensor_parallel_size == 4 + and parallel.pipeline_parallel_size == 1 + and parallel.decode_context_parallel_size == 1 + and parallel.prefill_context_parallel_size == 1 + ): + from vllm.distributed import get_tp_group + + self.selection_tp_group = get_tp_group() self.index_qk_proj = ReplicatedLinear( int(config.hidden_size), @@ -222,7 +246,6 @@ def forward( from .ops.qsa import ( qsa_compress_groups_with_ratio, - qsa_select_paged_tokens, qsa_store_cache_rows, ) @@ -348,7 +371,35 @@ def forward( raise RuntimeError("QSA top-k reuse requires an output buffer") return out - # Score compressed keys, select blocks, then expand them to token indices. + return self._select_tokens(q, compressed_key_cache, compressed_metadata, out) + + def _select_tokens( + self, + q: torch.Tensor, + compressed_key_cache: torch.Tensor, + compressed_metadata: QSAForwardMetadata, + out: torch.Tensor | None, + ) -> torch.Tensor: + group = self.selection_tp_group if 0 < q.shape[0] <= 64 else None + if out is None: + out = torch.empty( + q.shape[0], self.output_width, dtype=torch.int32, device=q.device + ) + if group is None or group.rank_in_group == 0: + self._select_local_tokens(q, compressed_key_cache, compressed_metadata, out) + if group is not None: + group.broadcast(out, src=0) + return out + + def _select_local_tokens( + self, + q: torch.Tensor, + compressed_key_cache: torch.Tensor, + compressed_metadata: QSAForwardMetadata, + out: torch.Tensor, + ) -> torch.Tensor: + from .ops.qsa import qsa_select_paged_tokens + return qsa_select_paged_tokens( q, compressed_key_cache, @@ -361,6 +412,7 @@ def forward( out, query_start_loc_cpu=compressed_metadata.query_start_loc_cpu, shared_key_scoring=self.shared_key_scoring, + shared_key_batch_scoring=self.shared_key_batch_scoring, max_seq_len=compressed_metadata.max_seq_len, ) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 99a0195bcb..e727e45d42 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -514,6 +514,7 @@ def enable_sm70_hcx( fuse_output_projection: bool = True, diagnostic: bool = False, local_schedule: bool = False, + batch_max_m: int = 8, ) -> bool: """Leave supported small-M outputs partial; retain larger-batch reductions.""" from .sm70_hcx import ( @@ -531,7 +532,8 @@ def enable_sm70_hcx( reason=runtime.reason, scope="hc_operator_capability", min_m=1, - max_m=8, + max_m=batch_max_m, + native_max_m=8, prepared_modules=0, deferred_projections=0, output_projection_policy="fused" if fuse_output_projection else "separate", @@ -577,6 +579,7 @@ def enable_sm70_hcx( else layer.self_attn.o_proj ) attention.reduce_results = False + attention._hcx_batch_max_m = batch_max_m planes = pack_output_projection(attention) if fuse_output_projection else None width = getattr(attention, "input_size_per_partition", None) if planes is not None and planes[0] != width: @@ -609,19 +612,26 @@ def enable_sm70_hcx( status["deferred_projections"] += 1 if isinstance(layer.mlp, Qwen4ExpSparseMoeBlock): runner = layer.mlp.experts.runner + runner._hcx_batch_max_m = batch_max_m runner.output_transform = partial(partial_moe_output, prefix) layer.mlp.sm70_hcx_packed_outputs = True 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) + layer.attn_hyper_connection.enable_partial_inputs( + prefix + ".attn", runtime, batch_max_m=batch_max_m + ) status["prepared_modules"] += 1 - layer.mlp_hyper_connection.enable_partial_inputs(prefix + ".mlp", runtime) + layer.mlp_hyper_connection.enable_partial_inputs( + prefix + ".mlp", runtime, batch_max_m=batch_max_m + ) status["prepared_modules"] += 1 previous_moe = isinstance(layer.mlp, Qwen4ExpSparseMoeBlock) if previous_moe and model.hyper_connection_mixer is not None: model.hyper_connection_mixer._hcx_moe_payload = True - model.hyper_connection_mixer.enable_partial_inputs("final") + model.hyper_connection_mixer.enable_partial_inputs( + "final", batch_max_m=batch_max_m + ) status.update(enabled=True, reason=None) logger.info_once("SM70 HCX loaded operator capability: %s.", str(status)) return True @@ -684,6 +694,7 @@ def _maybe_enable_sm70_peer_paths(vllm_config: VllmConfig, model: nn.Module) -> fuse_output_projection=kernel_config.sm70_hcx_output_projection, diagnostic=kernel_config.sm70_hcx_diagnostics, local_schedule=kernel_config.sm70_hcx_local_schedule, + batch_max_m=kernel_config.sm70_hcx_batch_max_m, ) kernel_config.collective_kernel_selections["hcx:target"] = model.sm70_hcx_status diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index d088008efb..9e09a1912e 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -1178,6 +1178,7 @@ def qsa_mqa_paged( num_columns: int | None = None, score_scale: float | None = None, shared_key_scoring: bool = False, + shared_key_query_start_loc_cpu: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute QSA scores directly from a paged compressed-key cache.""" @@ -1209,7 +1210,26 @@ def qsa_mqa_paged( if columns < 0: raise ValueError("QSA score width must be non-negative") if shared_key_scoring and 0 < columns <= capacity and q.shape[0]: - from .qsa_shared_key import maybe_shared_key_scores + from .qsa_shared_key import ( + maybe_shared_key_batch_scores, + maybe_shared_key_scores, + ) + + if shared_key_query_start_loc_cpu is not None: + result = maybe_shared_key_batch_scores( + q, + k_cache, + page_table, + token_to_req, + query_positions, + sequence_lengths, + compress_ratio, + columns, + score_divisor, + shared_key_query_start_loc_cpu, + ) + if result is not None: + return result result = maybe_shared_key_scores( q, @@ -1598,6 +1618,7 @@ def qsa_select_paged_tokens( query_start_loc_cpu: torch.Tensor | None = None, shared_key_scoring: bool = False, max_seq_len: int | None = None, + shared_key_batch_scoring: bool = False, ) -> torch.Tensor: """Score, select, and expand QSA indices without host synchronization. @@ -1708,6 +1729,11 @@ def qsa_select_paged_tokens( sequence_lengths, compress_ratio, shared_key_scoring=shared_key_scoring, + shared_key_query_start_loc_cpu=( + query_start_loc_cpu + if shared_key_batch_scoring and row_start == 0 and row_end == rows + else None + ), ) blocks = blocks_buffer[: row_end - row_start] use_cooperative_topk = ( diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py index 8b5f9b884d..83efbfb1d1 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py @@ -81,3 +81,76 @@ def maybe_shared_key_scores( ) logger.info_once("Using SM70 QSA shared-key indexer for M=%d.", q.shape[0]) return scores, visible + + +def maybe_shared_key_batch_scores( + q, + cache, + table, + requests, + positions, + lengths, + compress, + columns, + divisor, + query_start_loc_cpu, +): + """Reuse the single-request provider without changing batched top-k. + + Query starts are the owner's host layout, including capture padding. + All device metadata remains live across graph replays. An unqualified + group sends the entire batch to the existing scorer. + """ + if ( + table.shape[0] < 2 + or not 1 <= q.shape[0] <= 64 + or query_start_loc_cpu.device.type != "cpu" + or query_start_loc_cpu.ndim != 1 + or query_start_loc_cpu.dtype not in (torch.int32, torch.int64) + ): + return None + starts = query_start_loc_cpu.tolist() + if len(starts) != table.shape[0] + 1 or starts[0] != 0 or starts[-1] != q.shape[0]: + return None + if any(end < start for start, end in zip(starts, starts[1:])): + return None + groups = [ + (i, start, end) + for i, (start, end) in enumerate(zip(starts, starts[1:])) + if end > start + ] + if not groups or any(not 2 <= end - start <= 8 for _, start, end in groups): + return None + for i, start, end in groups: + if ( + shared_key_reason( + q[start:end], + cache, + table[i : i + 1], + requests[start:end], + positions[start:end], + lengths[i : i + 1], + ) + is not None + ): + return None + scores = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) + visible = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) + for i, start, end in groups: + local_requests = (requests[start:end] - i).contiguous() + torch.ops.vllm_sm70_qsa_indexer.run( + scores[start:end], + visible[start:end], + q[start:end], + cache, + table[i : i + 1], + local_requests, + positions[start:end], + lengths[i : i + 1], + float(divisor), + compress, + ) + logger.info_once( + "Using SM70 shared-key indexer for %d request groups.", len(groups) + ) + return scores, visible diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index c2b521d699..096efd917c 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -390,7 +390,7 @@ def forward_qsa( query_positions[:num_tokens], sequence_lengths, output[:num_tokens], - output_gate, + output_gate[:num_tokens] if output_gate is not None else None, attn_metadata, ) return output @@ -729,6 +729,12 @@ def __init__( # 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 + speculative_config = vllm_config.speculative_config + self.host_kv_decode_query_len = 1 + ( + (speculative_config.num_speculative_tokens or 0) + if speculative_config is not None + else 0 + ) if self.host_kv_is_draft: self.host_kv_hot_tokens = ( vllm_config.kernel_config.qsa_host_kv_draft_hot_tokens @@ -955,6 +961,13 @@ def host_kv_forward( if ( self.host_kv_prefill_stage_tokens and query.shape[0] > state.resolve_rows + # A verifier batch can exceed resolve_rows while every request + # still has only its decode query. Staging plans read metadata on + # the CPU and cannot be captured or reused for those requests. + and ( + metadata is None + or metadata.max_query_len > self.host_kv_decode_query_len + ) and query.shape[1:] == (6, 256) and indices.shape[1] <= state.width ): diff --git a/vllm/models/qwen4_exp/nvidia/sm70_hcx.py b/vllm/models/qwen4_exp/nvidia/sm70_hcx.py index e850874664..8c36117654 100644 --- a/vllm/models/qwen4_exp/nvidia/sm70_hcx.py +++ b/vllm/models/qwen4_exp/nvidia/sm70_hcx.py @@ -47,10 +47,11 @@ def _output_projection(x: torch.Tensor, name: str) -> torch.Tensor: # Resolve actual M inside the opaque op, rather than freezing a decode # choice while compiling a range that also includes C4 and prefill. layer, defer = _OUTPUT_PROJECTIONS[name] - if 0 < x.shape[0] <= HCX_MAX_M and defer: + maximum = getattr(layer, "_hcx_batch_max_m", HCX_MAX_M) + if 0 < x.shape[0] <= maximum and defer: return torch.nn.functional.pad(x, (0, layer.output_size - x.shape[1])) projected, _ = layer(x) - if x.shape[0] > HCX_MAX_M: + if x.shape[0] > maximum: from vllm.distributed import tensor_model_parallel_all_reduce projected = tensor_model_parallel_all_reduce(projected) @@ -66,14 +67,15 @@ def _moe_output( shared: torch.Tensor | None, fused: torch.Tensor, name: str, trunc_size: int ) -> torch.Tensor: rows = fused.shape[0] - if 0 < rows <= HCX_MAX_M: + runner = _MOE_RUNNERS[name] + maximum = getattr(runner, "_hcx_batch_max_m", HCX_MAX_M) + if 0 < rows <= maximum: # An owned payload makes both contributions explicit to the compiler # and graph allocator. Hidden references to producer temporaries are # unsafe: the memory planner cannot see those additional consumers. first = fused[..., :trunc_size] second = torch.zeros_like(first) if shared is None else shared[..., :trunc_size] return torch.cat((first, second), dim=0) - runner = _MOE_RUNNERS[name] reduced = runner._maybe_sm70_moe_sum2_allreduce(shared, fused, trunc_size) if reduced is None: summed = fused if shared is None else shared + fused @@ -324,6 +326,7 @@ def run( oproj=None, secondary=None, snapshot_name=None, + batch_max_m=HCX_MAX_M, ): m = partial.shape[0] ox = ocodes = ohigh = oscale = None @@ -350,40 +353,45 @@ def run( if "epoch" not in snapshot: snapshot["epoch"] = torch.empty_like(self.seq) snapshot["epoch"].copy_(self.seq) - torch.ops._C.sm70_hcx_out( - partial if ox is not None else partial.contiguous(), - secondary, - hidden.contiguous(), - injection.contiguous(), - norm_weight, - eps, - packed_down, - packed_up, - hidden_out, - block, - injection_out, - self.xn, - self.sq, - self.dpart, - self.bar, - self.seq, - self.ar, - self.lora, - self.hb, - self.logical_rank, - None, - int(self.full), - ox, - ocodes, - ohigh, - oscale, - ofmt, - None, - None, - 1e-6, - None, - self.local_schedule, - ) + if not 0 < m <= batch_max_m: + raise ValueError("HCX batch exceeds its initialized policy") + for start in range(0, m, HCX_MAX_M): + stop = min(start + HCX_MAX_M, m) + tile = slice(start, stop) + torch.ops._C.sm70_hcx_out( + partial[tile] if ox is not None else partial[tile].contiguous(), + secondary[tile] if secondary is not None else None, + hidden[tile].contiguous(), + injection[tile].contiguous(), + norm_weight, + eps, + packed_down, + packed_up, + hidden_out[tile], + block[tile], + injection_out[tile], + self.xn, + self.sq, + self.dpart, + self.bar, + self.seq, + self.ar, + self.lora, + self.hb, + self.logical_rank, + None, + int(self.full), + ox[tile] if ox is not None else None, + ocodes, + ohigh, + oscale, + ofmt, + None, + None, + 1e-6, + None, + self.local_schedule, + ) if snapshot is not None: for key, value in ( ("hidden_out", hidden_out), diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index ff876ed0ea..77b2182712 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -135,38 +135,75 @@ class _PartitionedBlockQueue: without growing GPU state tensors to the host-history capacity. """ - def __init__(self, blocks, state_blocks, state_group_ids): + def __init__(self, blocks, state_blocks, state_group_ids, state_group_offsets=()): self.state_blocks = state_blocks self.state_group_ids = frozenset(state_group_ids) - self.state = FreeKVCacheBlockQueue(blocks[1:state_blocks]) - self.history = FreeKVCacheBlockQueue(blocks[state_blocks:]) + if state_group_offsets: + if tuple(state_group_offsets) != tuple( + i * (state_blocks - 1) for i in range(len(state_group_ids)) + ): + raise ValueError("Compact state groups require consecutive ID ranges") + self.states = { + group_id: FreeKVCacheBlockQueue( + blocks[offset + 1 : offset + state_blocks] + ) + for group_id, offset in zip( + state_group_ids, state_group_offsets, strict=True + ) + } + self.logical_state_blocks = 1 + (state_blocks - 1) * len(state_group_ids) + else: + state = FreeKVCacheBlockQueue(blocks[1:state_blocks]) + self.states = dict.fromkeys(state_group_ids, state) + self.logical_state_blocks = state_blocks + self._state_queues = list(dict.fromkeys(self.states.values())) + self._state_group_order = list(self.states) + self.history = FreeKVCacheBlockQueue(blocks[self.logical_state_blocks :]) @property def num_free_blocks(self): - return self.state.num_free_blocks + self.history.num_free_blocks + return ( + sum(q.num_free_blocks for q in self._state_queues) + + self.history.num_free_blocks + ) def for_group(self, group_id): if group_id is None: raise ValueError("Partitioned cache allocation requires a group ID") - return self.state if group_id in self.state_group_ids else self.history + return self.states.get(group_id, self.history) + + def for_block(self, block): + if block.block_id >= self.logical_state_blocks: + return self.history + if len(self._state_queues) == 1: + return self._state_queues[0] + index = (block.block_id - 1) // (self.state_blocks - 1) + return self.states[self._state_group_order[index]] def popleft_n(self, n, group_id=None): return self.for_group(group_id).popleft_n(n) def remove(self, block): - queue = self.state if block.block_id < self.state_blocks else self.history - queue.remove(block) + self.for_block(block).remove(block) def append_n(self, blocks): - self.state.append_n([b for b in blocks if b.block_id < self.state_blocks]) - self.history.append_n([b for b in blocks if b.block_id >= self.state_blocks]) + for queue in [*self._state_queues, self.history]: + queue.append_n( + [block for block in blocks if self.for_block(block) is queue] + ) def prepend_n(self, blocks): - self.state.prepend_n([b for b in blocks if b.block_id < self.state_blocks]) - self.history.prepend_n([b for b in blocks if b.block_id >= self.state_blocks]) + for queue in [*self._state_queues, self.history]: + queue.prepend_n( + [block for block in blocks if self.for_block(block) is queue] + ) def get_all_free_blocks(self): - return self.state.get_all_free_blocks() + self.history.get_all_free_blocks() + return [ + block + for q in [*self._state_queues, self.history] + for block in q.get_all_free_blocks() + ] class BlockPool: @@ -197,6 +234,7 @@ def __init__( metrics_collector: KVCacheMetricsCollector | None = None, device_state_blocks: int = 0, device_state_group_ids: tuple[int, ...] = (), + device_state_group_offsets: tuple[int, ...] = (), ): assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0 self.num_gpu_blocks = num_gpu_blocks @@ -227,8 +265,13 @@ def __init__( if not 1 < device_state_blocks < num_gpu_blocks: raise ValueError("State partition must leave history capacity") self.free_block_queue = _PartitionedBlockQueue( - self.blocks, device_state_blocks, device_state_group_ids + self.blocks, + device_state_blocks, + device_state_group_ids, + device_state_group_offsets, ) + if self.free_block_queue.logical_state_blocks >= num_gpu_blocks: + raise ValueError("State partition must leave history capacity") self.enable_kv_cache_events = enable_kv_cache_events self.kv_event_queue: list[KVCacheEvent] = [] @@ -591,13 +634,12 @@ def can_allocate_group_blocks(self, demand: dict[int, int]) -> bool: """Check partition quotas; callers still check aggregate free capacity.""" if not isinstance(self.free_block_queue, _PartitionedBlockQueue): return True - state = sum( - n for i, n in demand.items() if i in self.free_block_queue.state_group_ids - ) - history = sum(demand.values()) - state - return ( - state <= self.free_block_queue.state.num_free_blocks - and history <= self.free_block_queue.history.num_free_blocks + queue_demand: dict[FreeKVCacheBlockQueue, int] = {} + for group_id, count in demand.items(): + queue = self.free_block_queue.for_group(group_id) + queue_demand[queue] = queue_demand.get(queue, 0) + count + return all( + count <= queue.num_free_blocks for queue, count in queue_demand.items() ) def get_usage(self) -> float: diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index c28caf63f4..92fdc155b6 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -61,6 +61,7 @@ def __init__( metrics_collector=metrics_collector, device_state_blocks=kv_cache_config.device_state_blocks, device_state_group_ids=kv_cache_config.device_state_group_ids, + device_state_group_offsets=kv_cache_config.device_state_group_offsets, ) # KV cache group indices that get the EAGLE last-block drop. diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 516c20555f..a49bec12b9 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -932,12 +932,16 @@ def get_max_concurrency_for_kv_cache_config( + 1 ) if i in state_groups: - state_per_request += demand + state_per_request = ( + max(state_per_request, demand) + if kv_cache_config.device_state_group_offsets + else state_per_request + demand + ) else: history_per_request += demand return min( (kv_cache_config.device_state_blocks - 1) / state_per_request, - (kv_cache_config.num_blocks - kv_cache_config.device_state_blocks) + (kv_cache_config.num_blocks - kv_cache_config.logical_state_blocks) / history_per_request, ) num_blocks_per_request = sum( @@ -1786,6 +1790,11 @@ def _host_kv_state_blocks(vllm_config: VllmConfig) -> int: return int(getattr(policy, "qsa_host_kv_state_blocks", 0)) +def _host_kv_compact_state(vllm_config: VllmConfig) -> bool: + policy = getattr(vllm_config, "kernel_config", None) + return bool(getattr(policy, "qsa_host_kv_state_compact", False)) + + def _get_kv_cache_config_csa_linear( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1796,6 +1805,7 @@ def _get_kv_cache_config_csa_linear( return None state_blocks = _host_kv_state_blocks(vllm_config) + compact_state = _host_kv_compact_state(vllm_config) state_groups = tuple( i for i, group in enumerate(kv_cache_groups) if group in layout.mamba_groups ) @@ -1817,10 +1827,13 @@ def demand(group): + 1 ) - active_states = 1 + requests * sum( + state_demands = [ demand(group) for i, group in enumerate(kv_cache_groups) if i in state_groups + ] + active_states = 1 + requests * ( + max(state_demands) if compact_state else sum(state_demands) ) if state_blocks < active_states: raise ValueError(f"State pool needs at least {active_states} blocks") @@ -1829,12 +1842,24 @@ def demand(group): for i, group in enumerate(kv_cache_groups) if i not in state_groups ) - state_bytes = ( - sum(layout.owner_page_size(i) for i in range(len(layout.main_kv_owners))) - * state_blocks + state_page_bytes = ( + sum( + group.kv_cache_spec.page_size_bytes * len(group.layer_names) + for group in layout.mamba_groups + ) + if compact_state + else sum( + layout.owner_page_size(i) for i in range(len(layout.main_kv_owners)) + ) ) + state_bytes = state_page_bytes * state_blocks history_page = layout.device_compressed_page_bytes - num_blocks = state_blocks + history_blocks + logical_states = ( + 1 + (state_blocks - 1) * len(state_groups) + if compact_state + else state_blocks + ) + num_blocks = logical_states + history_blocks required = state_bytes + history_page * num_blocks if required > available_memory: raise ValueError( @@ -1884,7 +1909,7 @@ def demand(group): device_members = [name for name in members if name not in host_members] + [ group.layer_names[index] for group in layout.mamba_groups - if index < len(group.layer_names) + if not compact_state and index < len(group.layer_names) ] if device_members: kv_cache_tensors.append( @@ -1895,6 +1920,18 @@ def demand(group): physical_num_blocks=state_blocks or None, ) ) + if compact_state: + if not state_blocks: + raise ValueError("Compact host state requires a bounded device state pool") + kv_cache_tensors.extend( + KVCacheTensor( + size=group.kv_cache_spec.page_size_bytes * state_blocks, + shared_by=[layer_name], + physical_num_blocks=state_blocks, + ) + for group in layout.mamba_groups + for layer_name in group.layer_names + ) kv_cache_tensors.extend( KVCacheTensor( size=layout.compressed_page_sizes[index] * num_blocks, @@ -1922,6 +1959,17 @@ def get_kv_cache_config_from_groups( Returns: The generated KVCacheConfig """ + if _host_kv_compact_state(vllm_config): + compact_layout = _get_csa_linear_tensor_layout(kv_cache_groups) + if compact_layout is None or not _host_kv_state_blocks(vllm_config): + raise ValueError("Compact host state requires the bounded host-KV layout") + parallel = vllm_config.parallel_config + if ( + parallel.pipeline_parallel_size != 1 + or parallel.decode_context_parallel_size != 1 + or parallel.prefill_context_parallel_size != 1 + ): + raise ValueError("Compact host state currently requires PP1 and CP1") if len(kv_cache_groups) == 0: # Attention free models do not have KV cache. # Return num_blocks=1 as BlockPool always needs a null_block. @@ -2045,15 +2093,21 @@ def get_kv_cache_config_from_groups( if layout is not None and layout.host_main_kv_names else 0 ) + state_group_ids = tuple( + i + for i, group in enumerate(kv_cache_groups) + if layout is not None and state_blocks and group in layout.mamba_groups + ) return KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=kv_cache_tensors, kv_cache_groups=kv_cache_groups, device_state_blocks=state_blocks, - device_state_group_ids=tuple( - i - for i, group in enumerate(kv_cache_groups) - if layout is not None and state_blocks and group in layout.mamba_groups + device_state_group_ids=state_group_ids, + device_state_group_offsets=( + tuple(i * (state_blocks - 1) for i in range(len(state_group_ids))) + if _host_kv_compact_state(vllm_config) + else () ), ) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index a63cc9fc08..27e706ba20 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -1017,6 +1017,27 @@ class KVCacheConfig: """ device_state_blocks: int = 0 device_state_group_ids: tuple[int, ...] = () + device_state_group_offsets: tuple[int, ...] = () + + @property + def logical_state_blocks(self) -> int: + if self.device_state_group_offsets: + return 1 + (self.device_state_blocks - 1) * len(self.device_state_group_ids) + return self.device_state_blocks + + @property + def state_block_id_offsets(self) -> dict[int, int]: + return ( + dict( + zip( + self.device_state_group_ids, + self.device_state_group_offsets, + strict=True, + ) + ) + if self.device_state_group_offsets + else {} + ) @property def has_mamba_layers(self) -> bool: @@ -1025,3 +1046,15 @@ def has_mamba_layers(self) -> bool: @property def needs_kv_cache_zeroing(self) -> bool: return self.has_mamba_layers + + +def map_device_state_block_ids( + block_ids: list[int], offset: int, capacity: int +) -> list[int]: + """Map one group's unique scheduler IDs to its independent device bank.""" + mapped = [block - offset if block else 0 for block in block_ids] + if any( + block and not 0 < local < capacity for block, local in zip(block_ids, mapped) + ): + raise ValueError("State block ID is outside its owning cache group") + return mapped diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index ce479e5ee4..04e8e1c436 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -11,6 +11,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import map_device_state_block_ids from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.cp_utils import get_total_cp_world_size @@ -64,6 +65,8 @@ def __init__( Will be the same as `block_size` if `block_size` is supported by the attention kernel. """ + self.state_block_id_offset: int | None = None + self.device_state_blocks = 0 self.max_num_reqs = max_num_reqs self.max_num_batched_tokens = max_num_batched_tokens self.pin_memory = pin_memory @@ -141,6 +144,10 @@ def append_row( ) -> None: if not block_ids: return + if self.state_block_id_offset is not None: + block_ids = map_device_state_block_ids( + block_ids, self.state_block_id_offset, self.device_state_blocks + ) if self.use_hybrid_blocks: block_ids = self.map_to_kernel_blocks( @@ -448,6 +455,14 @@ def __init__( ) ] + def set_device_state_mapping(self, offsets: dict[int, int], capacity: int) -> None: + """Bind the cache layout before any request writes its block tables.""" + for i, table in enumerate(self.block_tables): + if table.num_blocks_per_row.any(): + raise ValueError("Cannot replace a live state block-table mapping") + table.state_block_id_offset = offsets.get(i) + table.device_state_blocks = capacity + def append_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None: for i, block_table in enumerate(self.block_tables): block_table.append_row(block_ids[i], row_idx) diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 748280cb75..337b02dbbf 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -6,6 +6,7 @@ from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import map_device_state_block_ids from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor @@ -23,12 +24,16 @@ def __init__( cp_interleave: int = 1, slot_mapping_enabled: list[bool] | None = None, dcp_sharded: list[bool] | None = None, + state_block_id_offsets: dict[int, int] | None = None, + device_state_blocks: int = 0, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes self.max_num_reqs = max_num_reqs self.max_num_batched_tokens = max_num_batched_tokens self.device = device + self.state_block_id_offsets = state_block_id_offsets or {} + self.device_state_blocks = device_state_blocks self.cp_size = cp_size self.cp_rank = cp_rank @@ -118,6 +123,10 @@ def append_block_ids( for i in range(self.num_kv_cache_groups): start = self.num_blocks.np[i, req_index] if not overwrite else 0 block_ids = new_block_ids[i] + if i in self.state_block_id_offsets: + block_ids = map_device_state_block_ids( + block_ids, self.state_block_id_offsets[i], self.device_state_blocks + ) bpk = self.blocks_per_kv_block[i] if bpk > 1: block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 2dd1256d73..ac73b59b1a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -596,6 +596,10 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cp_interleave=self.cp_interleave, slot_mapping_enabled=slot_mapping_enabled, dcp_sharded=dcp_sharded, + state_block_id_offsets=getattr( + self.kv_cache_config, "state_block_id_offsets", {} + ), + device_state_blocks=self.kv_cache_config.device_state_blocks, ) initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config @@ -671,6 +675,9 @@ def _init_kv_zero_meta(self) -> None: if tensor.host_backed for name in tensor.shared_by }, + state_block_id_offsets=getattr( + self.kv_cache_config, "state_block_id_offsets", {} + ), ) def _zero_block_ids(self, block_ids: list[int]) -> None: diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 9ad49e90a2..adf1deef88 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -9,6 +9,7 @@ from vllm import PoolingParams, SamplingParams from vllm.utils.math_utils import cdiv +from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.sched.output import ( CachedRequestData, GrammarOutput, @@ -156,14 +157,42 @@ def warmup_kernels( ) if profile_idx: num_reqs = min(num_reqs, 1) + config = model_runner.kv_cache_config + state_blocks = getattr(config, "device_state_blocks", 0) + block_pool = ( + BlockPool( + config.num_blocks, + False, + 1, + device_state_blocks=state_blocks, + device_state_group_ids=config.device_state_group_ids, + device_state_group_offsets=config.device_state_group_offsets, + ) + if state_blocks + else None + ) + if block_pool is not None: + while num_reqs and not block_pool.can_allocate_group_blocks( + {i: n * num_reqs for i, n in enumerate(decode_block_counts)} + ): + num_reqs -= 1 if num_reqs <= 0: continue req_ids = [f"_warmup_{profile_idx}_{i}_" for i in range(num_reqs)] next_block_id = 1 - def _alloc_blocks(num_blocks: int) -> list[int]: + def _alloc_blocks( + group_id: int, num_blocks: int, *, pool=block_pool + ) -> list[int]: nonlocal next_block_id + if pool is not None: + return [ + block.block_id + for block in pool.get_new_blocks( + num_blocks, kv_cache_group_id=group_id + ) + ] return list( range(next_block_id, next_block_id := next_block_id + num_blocks) ) @@ -176,7 +205,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: sampling_params, pooling_params, ), - block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), + block_ids=tuple( + _alloc_blocks(g, n) for g, n in enumerate(prefill_block_counts) + ), prefill_token_ids=prompt_token_ids, ) for i in range(num_reqs) @@ -214,7 +245,10 @@ def _alloc_blocks(num_blocks: int) -> list[int]: new_block = any(decode_block_deltas) cached_req_data.new_block_ids = [ ( - tuple(_alloc_blocks(n) for n in decode_block_deltas) + tuple( + _alloc_blocks(g, n) + for g, n in enumerate(decode_block_deltas) + ) if new_block else None ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 02b86aa3df..55b56b4ac2 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2008,6 +2008,11 @@ def _init_kv_zero_meta(self) -> None: Delegates to KVBlockZeroer.init_meta with the runner's state. Called from gpu_worker.py outside the CuMem pool context. """ + if self.kv_cache_config.state_block_id_offsets: + self.input_batch.block_table.set_device_state_mapping( + self.kv_cache_config.state_block_id_offsets, + self.kv_cache_config.device_state_blocks, + ) self._kv_block_zeroer = KVBlockZeroer(self.device, self.pin_memory) self._kv_block_zeroer.init_meta( attn_groups_iter=self._kv_cache_spec_attn_group_iterator(), @@ -2021,6 +2026,9 @@ def _init_kv_zero_meta(self) -> None: if tensor.host_backed for name in tensor.shared_by }, + state_block_id_offsets=getattr( + self.kv_cache_config, "state_block_id_offsets", {} + ), ) def _zero_block_ids(self, block_ids: list[int]) -> None: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index effdb8e204..ee99b18d37 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -158,6 +158,7 @@ def _zero_kv_blocks_kernel( seg_block_strides_ptr, seg_page_sizes_ptr, seg_num_blocks_ptr, + seg_block_id_offsets_ptr, block_ids_ptr, BLOCK_SIZE: tl.constexpr, ): @@ -189,7 +190,13 @@ def _zero_kv_blocks_kernel( chunk_offset = chunk_index.to(tl.int64) * BLOCK_SIZE if chunk_offset >= page_size_el: return - block_id = tl.load(block_ids_ptr + block_index) + logical_id = tl.load(block_ids_ptr + block_index) + offset = 0 + if seg_block_id_offsets_ptr is not None: + offset = tl.load(seg_block_id_offsets_ptr + seg_index) + block_id = tl.where(logical_id == 0, 0, logical_id - offset) + if logical_id != 0 and block_id <= 0: + return if block_id < 0 or block_id >= tl.load(seg_num_blocks_ptr + seg_index): return seg_addr = tl.load(seg_addrs_ptr + seg_index) @@ -215,6 +222,7 @@ class KVBlockZeroer: def __init__(self, device: torch.device, pin_memory: bool): self.device = device self.pin_memory = pin_memory + self._segment_block_id_offsets: torch.Tensor | None = None self._meta: ( tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None @@ -228,6 +236,7 @@ def init_meta( runner_only_attn_layers: set[str], static_forward_context: dict[str, Any], host_backed_layers: set[str] | None = None, + state_block_id_offsets: dict[int, int] | None = None, ) -> None: """One-time precomputation for zero_block_ids. @@ -251,14 +260,20 @@ def init_meta( (~13 ms per 8K-token prefill step on x8 Gen3). """ host_backed_layers = host_backed_layers or set() + state_block_id_offsets = state_block_id_offsets or {} seen_ptrs: dict[int, int] = {} seg_addrs: list[int] = [] seg_page_sizes: list[int] = [] # Logical blocks each segment's tensor actually holds. seg_num_blocks: list[int] = [] + seg_block_id_offsets: list[int] = [] - def add_segment(address: int, page_size_el: int, num_blocks: int) -> None: + def add_segment( + address: int, page_size_el: int, num_blocks: int, block_id_offset: int + ) -> None: if (index := seen_ptrs.get(address)) is not None: + if seg_block_id_offsets[index] != block_id_offset: + raise ValueError("Independent state groups must not alias storage") seg_page_sizes[index] = max(seg_page_sizes[index], page_size_el) seg_num_blocks[index] = min(seg_num_blocks[index], num_blocks) return @@ -266,8 +281,10 @@ def add_segment(address: int, page_size_el: int, num_blocks: int) -> None: seg_addrs.append(address) seg_page_sizes.append(page_size_el) seg_num_blocks.append(num_blocks) + seg_block_id_offsets.append(block_id_offset) for group in attn_groups_iter: + block_id_offset = state_block_id_offsets.get(group.kv_cache_group_id, 0) spec = group.kv_cache_spec if group.kv_cache_group_id >= len(kernel_block_sizes): continue @@ -310,7 +327,9 @@ def add_segment(address: int, page_size_el: int, num_blocks: int) -> None: num_blocks = kv.shape[block_dim] // ratio for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)): off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) - add_segment(dp + off_bytes, cur_page_el, num_blocks) + add_segment( + dp + off_bytes, cur_page_el, num_blocks, block_id_offset + ) elif isinstance(spec, MambaSpec): if not isinstance(kv, (list, tuple)) or not kv: continue @@ -324,7 +343,7 @@ def add_segment(address: int, page_size_el: int, num_blocks: int) -> None: cur_page_el = page_bytes // 4 # A bounded state pool exposes fewer pages than the # scheduler's block-ID range. - add_segment(dp, cur_page_el, first_state.shape[0]) + add_segment(dp, cur_page_el, first_state.shape[0], block_id_offset) logger.info( "KV block zeroer: %d segments, %.1f MiB per block (skipped %d " @@ -362,6 +381,11 @@ def add_segment(address: int, page_size_el: int, num_blocks: int) -> None: dtype=torch.int64, device=self.device, ) + self._segment_block_id_offsets = ( + torch.tensor(seg_block_id_offsets, dtype=torch.int64, device=self.device) + if any(seg_block_id_offsets) + else None + ) self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), @@ -404,6 +428,7 @@ def zero_block_ids(self, block_ids: list[int]) -> None: seg_block_strides, seg_page_sizes, seg_num_blocks, + self._segment_block_id_offsets, idx, BLOCK_SIZE=blk_size, ) From 8c9d0dc488511d260e8f63b504998fe7b70d6f53 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 15 Oct 2026 07:30:37 +0800 Subject: [PATCH 12/19] [Perf] Stage shared indexer pages and reuse MTP batch weights Bound host indexer staging under its engine owner and reuse original dense/expert decoders for qualified batches. Keep opt-ins off by default, preserve legacy native and prepared-operand interfaces, and document whole-verifier distributions and rejected experiments. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 13 ++ .../ops/gguf_dense_segments_sm70.cuh | 102 +++++++--- csrc/sm70_turbomind/ops/gguf_dp4a.cuh | 87 ++++++++ csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu | 11 +- .../ops/gguf_grouped_gate_sm70.cuh | 190 ++++++++++++++++++ csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu | 102 +++++++--- csrc/torch_bindings.cpp | 15 ++ docs/design/architecture/runtime_reference.md | 4 +- .../flashnext_mtp_batch_16gb_20261015.md | 156 +++++++++++++- .../config/test_flashnext_mtp_batch_policy.py | 18 +- .../quantization/test_gguf_dense_segments.py | 86 ++++++++ .../quantization/test_gguf_grouped_gate.py | 66 ++++++ .../quantization/test_gguf_mtp_batch.py | 1 + .../quantization/test_gguf_segment_prefill.py | 39 ++++ .../qwen4_exp/test_qsa_indexer_single_rank.py | 5 + .../qwen4_exp/test_qsa_shared_key_batch.py | 22 +- .../qwen4_exp/test_sm70_hcx_dispatch.py | 9 +- vllm/config/kernel.py | 12 ++ .../layers/quantization/gguf_dense_hmma.py | 23 ++- .../layers/quantization/gguf_dmv.py | 18 ++ .../layers/quantization/gguf_grouped_gate.py | 79 ++++++++ .../layers/quantization/gguf_moe.py | 98 +++++++-- .../layers/quantization/gguf_turbomind.py | 94 ++++++++- vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 25 +++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 5 + .../nvidia/ops/qsa_indexer_staging.py | 162 +++++++++++++++ .../qwen4_exp/nvidia/ops/qsa_shared_key.py | 67 +++++- 27 files changed, 1417 insertions(+), 92 deletions(-) create mode 100644 csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh create mode 100644 tests/kernels/quantization/test_gguf_grouped_gate.py create mode 100644 vllm/model_executor/layers/quantization/gguf_grouped_gate.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/qsa_indexer_staging.py diff --git a/csrc/ops.h b/csrc/ops.h index fd1dd131f2..69a1176adb 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1045,6 +1045,7 @@ void gguf_dp4a_gate_up_sm70_out(torch::Tensor out, torch::Tensor activation, torch::Tensor up, int64_t source_type, bool activated, int64_t lanes_per_row, bool bank_aware); + void gguf_dp4a_lut4_gate_up_sm70_out( torch::Tensor out, torch::Tensor activation, torch::Tensor ids, torch::Tensor gate, torch::Tensor gate_stats, torch::Tensor up, @@ -1056,6 +1057,12 @@ void gguf_dense_segments_sm70_out( std::vector out, std::vector fmt, std::vector n, int64_t k, int64_t split, int64_t warps, torch::Tensor ws, torch::Tensor cnt, std::optional sgate); +void gguf_dense_segments_batch_sm70_out( + torch::Tensor x, std::vector codes, + std::vector high, std::vector scale, + std::vector out, std::vector fmt, + std::vector n, int64_t k, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, std::optional sgate); void gguf_shared_gate_up_sm70_out(torch::Tensor x, std::vector gate, std::vector up, @@ -1186,3 +1193,9 @@ void gguf_dmvq_sm70_out(torch::Tensor output, torch::Tensor input, torch::Tensor counters, torch::Tensor table, int64_t type, int64_t kw, int64_t split); void gguf_dmvq_book_sm70_out(torch::Tensor table, int64_t type); + +void gguf_dp4a_gate_up_grouped_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode); diff --git a/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh index e8c9cf4059..d048bd489d 100644 --- a/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dense_segments_sm70.cuh @@ -368,18 +368,39 @@ __global__ void __launch_bounds__(32 * W) template void launch(const Segs& segs, const half* x, int ldx, int M, int K, int S, - int G, int split, float* ws, int* cnt, int tiles, cudaStream_t st) { - // Small-M verification needs three token tiles, but each uses the same - // weights. Keep the original path for other shapes and split-K schedules. + int G, int split, float* ws, int* cnt, int tiles, cudaStream_t st, + bool batch) { + // Reuse one weight decode across verification token tiles. Unsupported + // shapes and split-K schedules keep the existing row-tiled path. int total_n = 0; for (int i = 0; i < segs.nseg; ++i) total_n += segs.s[i].n; if constexpr (W == 4 || W == 8) { - if (M == 20 && split == 1 && - ((K == 2560 && (total_n == 4096 || total_n == 3584)) || - (K == 1536 && total_n == 2560))) { - dense_weight_major - <<>>( - segs, x, ldx, M, K, S, G, split, ws, cnt); + const bool legacy = + M == 20 && ((K == 2560 && (total_n == 4096 || total_n == 3584)) || + (K == 1536 && total_n == 2560)); + const bool batched = + batch && M > 8 && M <= 40 && + ((K == 2560 && total_n > 320) || (K == 1536 && total_n == 2560)); + if (split == 1 && (legacy || batched)) { +#define LAUNCH_WEIGHT_MAJOR(TM) \ + dense_weight_major \ + <<>>( \ + segs, x, ldx, M, K, S, G, split, ws, cnt) + switch ((M + 7) / 8) { + case 2: + LAUNCH_WEIGHT_MAJOR(2); + break; + case 3: + LAUNCH_WEIGHT_MAJOR(3); + break; + case 4: + LAUNCH_WEIGHT_MAJOR(4); + break; + case 5: + LAUNCH_WEIGHT_MAJOR(5); + break; + } +#undef LAUNCH_WEIGHT_MAJOR C10_CUDA_KERNEL_LAUNCH_CHECK(); return; } @@ -410,15 +431,20 @@ void launch(const Segs& segs, const half* x, int ldx, int M, int K, int S, } // namespace // codes/high/scale: per segment planes; out: per segment [M, n] views. -void gguf_dense_segments_sm70_out( - torch::Tensor x, std::vector codes, - std::vector high, std::vector scale, - std::vector out, std::vector fmt, - std::vector n, int64_t K, int64_t split, int64_t warps, - torch::Tensor ws, torch::Tensor cnt, std::optional sgate) { - TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kHalf && x.dim() == 2 && - x.stride(1) == 1 && x.size(0) > 0 && x.size(0) <= 32, - "input must be CUDA FP16 [1..32, K] with contiguous rows"); +template +void gguf_dense_segments_impl(torch::Tensor x, std::vector codes, + std::vector high, + std::vector scale, + std::vector out, + std::vector fmt, std::vector n, + int64_t K, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, + std::optional sgate) { + TORCH_CHECK( + x.is_cuda() && x.scalar_type() == at::kHalf && x.dim() == 2 && + x.stride(1) == 1 && x.size(0) > 0 && x.size(0) <= (Batch ? 40 : 32), + Batch ? "input must be CUDA FP16 [1..40, K] with contiguous rows" + : "input must be CUDA FP16 [1..32, K] with contiguous rows"); const c10::cuda::CUDAGuard guard(x.device()); const auto* properties = at::cuda::getCurrentDeviceProperties(); TORCH_CHECK(properties->major == 7 && properties->minor == 0, @@ -492,7 +518,7 @@ void gguf_dense_segments_sm70_out( cnt.numel() >= tiles * ((x.size(0) + 7) / 8), "split workspace or counter storage is too small"); const int M = static_cast(x.size(0)); - TORCH_CHECK(M <= 32); + TORCH_CHECK(M <= 40); const int S = static_cast((K + 31) / 32); const int G = (S + 3) / 4; auto st = at::cuda::getCurrentCUDAStream(); @@ -502,31 +528,51 @@ void gguf_dense_segments_sm70_out( int* cp = cnt.data_ptr(); switch (warps) { case 1: - launch<1>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<1>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 2: - launch<2>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<2>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 4: - launch<4>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<4>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 16: - launch<16>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<16>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 5: - launch<5>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<5>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 6: - launch<6>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<6>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 10: - launch<10>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<10>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; case 12: - launch<12>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<12>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; default: - launch<8>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st); + launch<8>(segs, xp, ldx, M, K, S, G, split, wsp, cp, tiles, st, Batch); break; } } + +void gguf_dense_segments_sm70_out( + torch::Tensor x, std::vector codes, + std::vector high, std::vector scale, + std::vector out, std::vector fmt, + std::vector n, int64_t K, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, std::optional sgate) { + gguf_dense_segments_impl(x, codes, high, scale, out, fmt, n, K, split, + warps, ws, cnt, sgate); +} + +void gguf_dense_segments_batch_sm70_out( + torch::Tensor x, std::vector codes, + std::vector high, std::vector scale, + std::vector out, std::vector fmt, + std::vector n, int64_t K, int64_t split, int64_t warps, + torch::Tensor ws, torch::Tensor cnt, std::optional sgate) { + gguf_dense_segments_impl(x, codes, high, scale, out, fmt, n, K, split, + warps, ws, cnt, sgate); +} diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh index ce28e96624..437b08182e 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh @@ -55,6 +55,93 @@ struct LatticeDot { return book[index]; } + struct Decoded { + int words[8]; + float scale; + int scale0, scale1; + }; + + __device__ static Decoded decode(const uint8_t* row, int group, + const uint32_t* book, + const uint32_t* masks) { + Decoded operand{}; + const uint8_t* b = row + (group / 8) * kBlockBytes; + const int sub = group % 8; + const uint32_t low = load_u32_2(b + 2 + sub * (Type == 22 ? 4 : 8)); + uint32_t high = 0, signs = 0; + if constexpr (Type != 22) high = load_u32_2(b + 6 + sub * 8); + if constexpr (Type == 18) + signs = load_u32_2(b + 66 + sub * 4); + else + signs = load_u32_2(b + (Type == 21 ? 74 : 34) + sub * 4); + const int qh = Type == 18 ? 0 : b[66 + sub]; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + int first, second, sign; + if constexpr (Type == 22) { + first = + (((low >> (octet * 8)) & 255) | (((qh >> (2 * octet)) & 3) << 8)) * + 2; + second = first + 1; + } else { + const uint32_t codes = octet < 2 ? low : high; + const int shift = (octet % 2) * 16; + first = ((codes >> shift) & 255); + second = ((codes >> (shift + 8)) & 255); + if constexpr (Type == 21) { + first |= ((qh >> (2 * octet)) & 1) << 8; + second |= ((qh >> (2 * octet + 1)) & 1) << 8; + } + } + if constexpr (Type == 18) { + sign = (signs >> (7 * octet)) & 127; + sign |= (__popc(sign) & 1) << 7; + } else { + sign = (signs >> (8 * octet)) & 255; + } + const uint32_t s0 = BankAware ? sign_mask(sign & 15) : masks[sign & 15]; + const uint32_t s1 = BankAware ? sign_mask(sign >> 4) : masks[sign >> 4]; + operand.words[2 * octet] = __vsub4(word(book, first) ^ s0, s0); + operand.words[2 * octet + 1] = __vsub4(word(book, second) ^ s1, s1); + } + operand.scale = __half2float(*reinterpret_cast(b)); + if constexpr (Type == 21) { + operand.scale0 = 1 + 2 * ((b[106 + sub / 2] >> (4 * (sub % 2))) & 15); + } else if constexpr (Type == 18) { + operand.scale0 = 1 + 2 * (signs >> 28); + } else { + operand.scale0 = 1 + 2 * (b[74 + sub] & 15); + operand.scale1 = 1 + 2 * (b[74 + sub] >> 4); + } + return operand; + } + + __device__ static float dot_decoded(const Decoded& operand, const Q8_1& x) { + const int* activation = reinterpret_cast(x.qs); + int sum0 = 0, sum1 = 0; +#pragma unroll + for (int octet = 0; octet < 4; ++octet) { + if constexpr (Type == 22) { + if (octet >= 2) { + sum1 = __dp4a(operand.words[2 * octet], activation[2 * octet], sum1); + sum1 = __dp4a(operand.words[2 * octet + 1], activation[2 * octet + 1], + sum1); + continue; + } + } + sum0 = __dp4a(operand.words[2 * octet], activation[2 * octet], sum0); + sum0 = + __dp4a(operand.words[2 * octet + 1], activation[2 * octet + 1], sum0); + } + const float d = operand.scale * __low2float(x.ds); + if constexpr (Type == 21) + return d * float(sum0 * operand.scale0); + else if constexpr (Type == 18) + return d * (float(sum0) * float(operand.scale0) * .25f); + else + return d * float(sum0 * operand.scale0 + sum1 * operand.scale1) * .125f; + } + __device__ static float dot(const uint8_t* row, int group, const Q8_1& x, const uint32_t* book, const uint32_t* masks) { const uint8_t* b = row + (group / 8) * kBlockBytes; diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu index faa16955c5..5563db43b0 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu +++ b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu @@ -23,12 +23,17 @@ __global__ void quantize_q8(Q8_1* out, const half* input, int k) { template + bool ScalarLut = false, bool BankAware = false, int FixedGroups = 0, + bool SparseFallback = false> __global__ void gate_up(void* output, const Q8_1* activation, const Index* ids, const uint8_t* gate, const uint8_t* up, int n, int k, int stride, int top_k, const StridedPtr* gate_stats = nullptr, - const StridedPtr* up_stats = nullptr) { + const StridedPtr* up_stats = nullptr, + const int* grouping_mode = nullptr) { + if constexpr (SparseFallback) { + if (grouping_mode[0] != 0) return; + } using Dot = std::conditional_t< Canonical, vllm::sm70_gguf::CanonicalIntegerDot, std::conditional_t, @@ -712,3 +717,5 @@ void gguf_dp4a_scalar_lut_gate_up_sm70_out( dispatch_scalar_lut_gate_up<22>(out, activation, ids, gate, up, activated); C10_CUDA_KERNEL_LAUNCH_CHECK(); } + +#include "gguf_grouped_gate_sm70.cuh" diff --git a/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh new file mode 100644 index 0000000000..65f76051a4 --- /dev/null +++ b/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Sorted route groups reuse original GGUF weight/codebook decoding. Integer +// dots, group reduction, FP16 activation boundaries and Q8_1 encoding retain +// the single-route order. The original gate/up operator stays unchanged. +namespace { +template +__global__ __launch_bounds__(512, 1) void gate_up_grouped( + Q8_1* output, const Q8_1* activation, const int* ids, + const int* route_tokens, const int* order, const int* original_ids, + const int* mode, const uint8_t* gate, const uint8_t* up, int m, int n, + int k, int stride, int top_k, int experts) { + using Dot = LatticeDot; + __shared__ uint32_t book[Dot::kBookWords], masks[16]; + __shared__ half intermediate[4][32]; + __shared__ int active, count, tokens[4], expert; + extern __shared__ Q8_1 shared_x[]; + const int route = blockIdx.y, routes = m * top_k; + if (mode[0] == 0) return; + if (threadIdx.x == 0) { + const int e = ids[route]; + int first = route; + while (first > 0 && ids[first - 1] == e) --first; + active = ((route - first) % 4 == 0) && e >= 0 && e < experts; + count = 1; + while (count < 4 && route + count < routes && ids[route + count] == e) + ++count; + for (int q = 0; q < count; ++q) + tokens[q] = min(max(route_tokens[route + q], 0), m - 1); + expert = e; + } + __syncthreads(); + if (!active) return; + const int groups = k / 32; + for (int i = threadIdx.x; i < count * groups * 9; i += blockDim.x) { + const int token = tokens[i / (groups * 9)]; + reinterpret_cast(shared_x)[i] = + reinterpret_cast( + activation)[int64_t(token) * groups * 9 + i % (groups * 9)]; + } + Dot::initialize(book, masks); + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int local_row = warp * 2 + lane / 16, row = blockIdx.x * 32 + local_row; + const uint8_t* g = gate + (int64_t(expert) * n + row) * stride; + const uint8_t* u = up + (int64_t(expert) * n + row) * stride; + float gs[4] = {}, us[4] = {}; + for (int group = lane % 16; group < groups; group += 16) { + if (count == 1) { + gs[0] += Dot::dot(g, group, shared_x[group], book, masks); + us[0] += Dot::dot(u, group, shared_x[group], book, masks); + } else { + const auto gd = Dot::decode(g, group, book, masks); + const auto ud = Dot::decode(u, group, book, masks); +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < count) { + gs[q] += Dot::dot_decoded(gd, shared_x[q * groups + group]); + us[q] += Dot::dot_decoded(ud, shared_x[q * groups + group]); + } + } + } + } +#pragma unroll + for (int offset = 8; offset; offset >>= 1) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + gs[q] += __shfl_down_sync(0xffffffff, gs[q], offset, 16); + us[q] += __shfl_down_sync(0xffffffff, us[q], offset, 16); + } + } + if (lane % 16 == 0) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < count) { + const float gv = __half2float(__float2half_rn(gs[q])); + intermediate[q][local_row] = __hmul( + __float2half_rn(gv / (1.f + expf(-gv))), __float2half_rn(us[q])); + } + } + } + __syncthreads(); + if (threadIdx.x < 32) { +#pragma unroll + for (int q = 0; q < 4; ++q) { + if (q < count) + vllm::sm70_gguf::quantize_q8_1_warp( + output + int64_t(order[route + q]) * (n / 32) + blockIdx.x, + __half2float(intermediate[q][threadIdx.x])); + } + } +} + +template +void launch_grouped_gate(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, torch::Tensor tokens, + torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + gate_up_grouped<<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(activation.data_ptr()), ids.data_ptr(), + tokens.data_ptr(), order.data_ptr(), + original_ids.data_ptr(), mode.data_ptr(), + gate.data_ptr(), up.data_ptr(), activation.size(0), + gate.size(1), activation.size(1) * 32, gate.size(2), ids.size(1), + gate.size(0)); + // Sparse routing uses the retained single-route kernel, with only a device + // admission guard. It must not inherit the grouped kernel's four-row register + // and shared-memory footprint. Both guards are graph replay inputs. + gate_up + <<>>( + out.data_ptr(), reinterpret_cast(activation.data_ptr()), + original_ids.data_ptr(), gate.data_ptr(), + up.data_ptr(), gate.size(1), activation.size(1) * 32, + gate.size(2), ids.size(1), nullptr, nullptr, mode.data_ptr()); +} +} // namespace + +void gguf_dp4a_gate_up_grouped_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + TORCH_CHECK(source_type == 18 || source_type == 21 || source_type == 22, + "Paired gate/up requires an original lattice format"); + const c10::cuda::CUDAGuard guard(activation.device()); + require_sm70(); + TORCH_CHECK(activation.is_cuda() && activation.scalar_type() == at::kByte && + activation.is_contiguous() && activation.dim() == 3 && + activation.size(0) > 0 && activation.size(0) <= 20 && + activation.size(2) == sizeof(Q8_1) && + activation.size(1) % 8 == 0, + "Invalid paired Q8_1 activation"); + const int k = activation.size(1) * 32; + const int bytes = source_type == 18 ? 98 : source_type == 21 ? 110 : 82; + TORCH_CHECK(gate.device() == activation.device() && + up.device() == activation.device() && + gate.scalar_type() == at::kByte && + up.scalar_type() == at::kByte && gate.is_contiguous() && + up.is_contiguous() && gate.dim() == 3 && + gate.sizes() == up.sizes() && gate.size(0) > 0 && + gate.size(1) > 0 && gate.size(1) % 32 == 0 && + (gate.size(2) == k / 256 * bytes || + gate.size(2) == (k / 256 * bytes + 7) / 8 * 8), + "Invalid paired original weight banks"); + TORCH_CHECK( + ids.device() == activation.device() && ids.scalar_type() == at::kInt && + ids.is_contiguous() && ids.dim() == 2 && + ids.size(0) == activation.size(0) && ids.size(1) > 0 && + ids.size(1) <= 16 && route_tokens.device() == activation.device() && + route_tokens.scalar_type() == at::kInt && + route_tokens.is_contiguous() && route_tokens.numel() == ids.numel(), + "Invalid sorted route or token mapping"); + TORCH_CHECK(order.device() == activation.device() && + order.scalar_type() == at::kInt && order.is_contiguous() && + order.numel() == ids.numel() && + original_ids.device() == activation.device() && + original_ids.scalar_type() == at::kInt && + original_ids.is_contiguous() && + original_ids.sizes() == ids.sizes() && + mode.device() == activation.device() && + mode.scalar_type() == at::kInt && mode.is_contiguous() && + mode.numel() == 1, + "Invalid grouped routing workspace"); + TORCH_CHECK( + out.device() == activation.device() && out.scalar_type() == at::kByte && + out.is_contiguous() && out.dim() == 4 && + out.size(0) == activation.size(0) && out.size(1) == ids.size(1) && + out.size(2) == gate.size(1) / 32 && out.size(3) == sizeof(Q8_1), + "Invalid paired Q8_1 output"); +#define PAIR(TYPE) \ + if (bank_aware) \ + launch_grouped_gate(out, activation, ids, gate, up, \ + route_tokens, order, original_ids, mode); \ + else \ + launch_grouped_gate(out, activation, ids, gate, up, \ + route_tokens, order, original_ids, mode) + if (source_type == 18) { + PAIR(18); + } else if (source_type == 21) { + PAIR(21); + } else { + PAIR(22); + } +#undef PAIR + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu b/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu index d0be9cf100..e7f17a61a5 100644 --- a/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu +++ b/csrc/sm70_turbomind/ops/qsa_indexer_sm70.cu @@ -20,14 +20,13 @@ namespace { "+f"(C[5]), "+f"(C[6]), "+f"(C[7]) \ : "r"(A0), "r"(A1), "r"(B0), "r"(B1)) -template -__global__ void mqa_multiquery(const half* q, const half* cache, - const int* table, const int* requests, - const Position* positions, const int* lengths, - float* scores, int* visible, int m, int columns, - int pages, int page_size, int table_width, - int sq0, int sq1, int sq2, int sc0, int sc1, - int sc3, int st1, float divisor, int compress) { +template +__global__ void mqa_multiquery( + const half* q, const half* cache, const int* table, const int* requests, + const Position* positions, const int* lengths, float* scores, int* visible, + int m, int columns, int pages, int page_size, int table_width, int sq0, + int sq1, int sq2, int sc0, int sc1, int sc3, int st1, float divisor, + int compress, const half* staged, const int* page_slots, int staged_pages) { const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; const int quad = (lane >> 2) & 3; const int r = (lane & 3) + ((lane & 16) ? 4 : 0); @@ -49,6 +48,20 @@ __global__ void mqa_multiquery(const half* q, const half* cache, ? table[(col / page_size) * st1] : -1; const bool key_valid = col < end && page >= 0 && page < pages; + const half* key = nullptr; + int key_stride = sc3; + if (key_valid) { + int resident = -1; + if constexpr (Staged) resident = page_slots[page]; + if (resident >= 0 && resident < staged_pages) { + key = + staged + + (static_cast(resident) * page_size + col % page_size) * 128; + key_stride = 1; + } else { + key = cache + static_cast(page) * sc0 + (col % page_size) * sc1; + } + } const int qhead = quad * 8 + r, qrow = qhead / 4; float acc[8] = {}; #pragma unroll @@ -59,14 +72,13 @@ __global__ void mqa_multiquery(const half* q, const half* cache, half* bv = reinterpret_cast(&b); half* lv = reinterpret_cast(&lo); half* hv = reinterpret_cast(&hi); - if (sc3 == 1 && sq2 == 1 && sc0 % 8 == 0 && sc1 % 8 == 0 && sq0 % 8 == 0 && - sq1 % 8 == 0 && + if (key_stride == 1 && sq2 == 1 && sc0 % 8 == 0 && sc1 % 8 == 0 && + sq0 % 8 == 0 && sq1 % 8 == 0 && (reinterpret_cast(q) | reinterpret_cast(cache)) % 16 == 0) { if (key_valid) { - const half* k = cache + static_cast(page) * sc0 + - (col % page_size) * sc1 + g * 16; + const half* k = key + g * 16; a = *reinterpret_cast(k); b = *reinterpret_cast(k + 8); } @@ -79,10 +91,8 @@ __global__ void mqa_multiquery(const half* q, const half* cache, #pragma unroll for (int j = 0; j < 8; ++j) { if (key_valid) { - const half* k = cache + static_cast(page) * sc0 + - (col % page_size) * sc1; - av[j] = k[(g * 16 + j) * sc3]; - bv[j] = k[(g * 16 + 8 + j) * sc3]; + av[j] = key[(g * 16 + j) * key_stride]; + bv[j] = key[(g * 16 + 8 + j) * key_stride]; } if (qrow < m) { const half* query = q + qrow * sq0 + (qhead % 4) * sq1; @@ -123,10 +133,13 @@ __global__ void mqa_multiquery(const half* q, const half* cache, #undef QSA_MMA -void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, - torch::Tensor cache, torch::Tensor table, torch::Tensor requests, - torch::Tensor positions, torch::Tensor lengths, double divisor, - int64_t compress) { +template +void mqa_out_impl(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, + torch::Tensor cache, torch::Tensor table, + torch::Tensor requests, torch::Tensor positions, + torch::Tensor lengths, double divisor, int64_t compress, + const torch::Tensor& staged, + const torch::Tensor& page_slots) { const c10::cuda::CUDAGuard guard(q.device()); TORCH_CHECK(q.is_cuda() && q.scalar_type() == at::kHalf && q.dim() == 3 && q.size(0) > 0 && q.size(0) <= 8 && q.size(1) == 4 && @@ -155,9 +168,25 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, "MQA row metadata must be contiguous"); const auto* props = at::cuda::getCurrentDeviceProperties(); TORCH_CHECK(props->major == 7 && props->minor == 0, "SM70 required"); + const half* staged_ptr = nullptr; + const int* slots_ptr = nullptr; + int staged_pages = 0; + if constexpr (Staged) { + TORCH_CHECK( + staged.device() == q.device() && staged.scalar_type() == at::kHalf && + staged.dim() == 4 && staged.is_contiguous() && + staged.size(1) == cache.size(1) && staged.size(2) == 1 && + staged.size(3) == 128 && page_slots.device() == q.device() && + page_slots.scalar_type() == at::kInt && + page_slots.is_contiguous() && page_slots.numel() >= cache.size(0), + "Invalid staged indexer page workspace"); + staged_ptr = reinterpret_cast(staged.data_ptr()); + slots_ptr = page_slots.data_ptr(); + staged_pages = staged.size(0); + } if (!columns) return; if (positions.scalar_type() == at::kLong) { - mqa_multiquery + mqa_multiquery <<<(columns + 31) / 32, 128, 0, at::cuda::getCurrentCUDAStream()>>>( reinterpret_cast(q.data_ptr()), reinterpret_cast(cache.data_ptr()), @@ -166,9 +195,10 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, scores.data_ptr(), visible.data_ptr(), m, columns, cache.size(0), cache.size(1), table.size(1), q.stride(0), q.stride(1), q.stride(2), cache.stride(0), cache.stride(1), - cache.stride(3), table.stride(1), divisor, compress); + cache.stride(3), table.stride(1), divisor, compress, staged_ptr, + slots_ptr, staged_pages); } else { - mqa_multiquery + mqa_multiquery <<<(columns + 31) / 32, 128, 0, at::cuda::getCurrentCUDAStream()>>>( reinterpret_cast(q.data_ptr()), reinterpret_cast(cache.data_ptr()), @@ -177,10 +207,28 @@ void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, scores.data_ptr(), visible.data_ptr(), m, columns, cache.size(0), cache.size(1), table.size(1), q.stride(0), q.stride(1), q.stride(2), cache.stride(0), cache.stride(1), - cache.stride(3), table.stride(1), divisor, compress); + cache.stride(3), table.stride(1), divisor, compress, staged_ptr, + slots_ptr, staged_pages); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } + +void mqa_out(torch::Tensor scores, torch::Tensor visible, torch::Tensor q, + torch::Tensor cache, torch::Tensor table, torch::Tensor requests, + torch::Tensor positions, torch::Tensor lengths, double divisor, + int64_t compress) { + mqa_out_impl(scores, visible, q, cache, table, requests, positions, + lengths, divisor, compress, {}, {}); +} + +void mqa_staged_out(torch::Tensor scores, torch::Tensor visible, + torch::Tensor q, torch::Tensor cache, torch::Tensor table, + torch::Tensor requests, torch::Tensor positions, + torch::Tensor lengths, double divisor, int64_t compress, + torch::Tensor staged, torch::Tensor page_slots) { + mqa_out_impl(scores, visible, q, cache, table, requests, positions, + lengths, divisor, compress, staged, page_slots); +} } // namespace TORCH_LIBRARY(vllm_sm70_qsa_indexer, m) { m.def( @@ -188,5 +236,11 @@ TORCH_LIBRARY(vllm_sm70_qsa_indexer, m) { "Tensor table, Tensor requests, Tensor positions, Tensor lengths, " "float divisor, int compress) -> ()"); m.impl("run", torch::kCUDA, &mqa_out); + m.def( + "run_staged(Tensor(a!) scores, Tensor(b!) visible, Tensor q, Tensor " + "cache, " + "Tensor table, Tensor requests, Tensor positions, Tensor lengths, " + "float divisor, int compress, Tensor staged, Tensor page_slots) -> ()"); + m.impl("run_staged", torch::kCUDA, &mqa_staged_out); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 32783ee3f0..4c34a26516 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -396,6 +396,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int warps, Tensor(b!) ws, Tensor(c!) cnt, Tensor? sgate=None) -> ()"); ops.impl("gguf_dense_segments_sm70_out", torch::kCUDA, &gguf_dense_segments_sm70_out); + ops.def( + "gguf_dense_segments_batch_sm70_out(Tensor x, Tensor[] codes, Tensor[] " + "high, " + "Tensor[] scale, Tensor(a!)[] out, int[] fmt, int[] n, int k, int split, " + "int warps, Tensor(b!) ws, Tensor(c!) cnt, Tensor? sgate=None) -> ()"); + ops.impl("gguf_dense_segments_batch_sm70_out", torch::kCUDA, + &gguf_dense_segments_batch_sm70_out); ops.def( "gguf_shared_gate_up_sm70_out(Tensor x, Tensor[] gate, Tensor[] up, " "int[] fmts, Tensor wg, Tensor(a!) h, Tensor(b!) sg, Tensor(c!) ws, " @@ -432,6 +439,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "int lanes_per_row=16, bool bank_aware=False) -> ()"); ops.impl("gguf_dp4a_gate_up_sm70_out", torch::kCUDA, &gguf_dp4a_gate_up_sm70_out); + + ops.def( + "gguf_dp4a_gate_up_grouped_sm70_out(Tensor(a!) out, Tensor activation, " + "Tensor ids, Tensor gate, Tensor up, int source_type, bool bank_aware, " + "Tensor route_tokens, Tensor order, Tensor original_ids, Tensor mode) -> " + "()"); + ops.impl("gguf_dp4a_gate_up_grouped_sm70_out", torch::kCUDA, + &gguf_dp4a_gate_up_grouped_sm70_out); ops.def( "gguf_dp4a_lut4_gate_up_sm70_out(Tensor(a!) out, Tensor activation, " "Tensor ids, Tensor gate, Tensor gate_stats, Tensor up, Tensor up_stats, " diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index a17ac1896c..deaec251b8 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -38,12 +38,12 @@ 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#L1079) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L625) | +| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1089) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L629) | | `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) | | `vllm.config.kernel.Sm70NvFp4Config` | kernel_config.sm70_nvfp4 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L277); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L236) | -| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L614) | +| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L618) | | `vllm.config.kv_events.KVEventsConfig` | kv_events_config | __post_init__: [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L50) | [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L11) | | `vllm.config.kv_transfer.KVTransferConfig` | kv_transfer_config | compute_hash: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L75); __post_init__: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L93) | [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L23) | | `vllm.config.legacy_inputs.LegacyInputs` | kernel_config.sm70_awq.legacy, kernel_config.sm70_fp8.legacy, kernel_config.sm70_moe.awq.legacy, kernel_config.sm70_moe.fp8.legacy, kernel_config.sm70_moe.mxfp4.legacy, kernel_config.sm70_moe.nvfp4.legacy, kernel_config.sm70_nvfp4.legacy, scheduler_config.sm70_inputs | value: [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L39) | [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L15) | diff --git a/docs/design/flashnext_mtp_batch_16gb_20261015.md b/docs/design/flashnext_mtp_batch_16gb_20261015.md index 599cd81606..78b5e3dc8d 100644 --- a/docs/design/flashnext_mtp_batch_16gb_20261015.md +++ b/docs/design/flashnext_mtp_batch_16gb_20261015.md @@ -130,7 +130,7 @@ The trace retains a separate wall residual and incomplete-round exclusions. existing TP communicator broadcasts them. Every rank still writes its own raw/compressed keys. Larger batches retain independent selection. Device indexer history, including the qualified C1 layout, remains on its existing - path. The 32K task gate passes; fresh ordinary-worker paired performance is pending. + path. The 32K task gate and fresh ordinary-worker paired task/performance gates pass (see the updated evidence below). These are explicit opt-ins until the paired end-to-end gates pass. No kernel sidecars, copied extensions or preloads are required. @@ -241,7 +241,159 @@ Python deduplication, automatic weather-tool calls and tool replies pass with thinking, production sampling and natural EOS. The tool test initially required `finish_reason=stop` for a successful tool call; that harness mistake was fixed to admit the expected `tool_calls` completion, and tool/reply checks were rerun. -Fresh ordinary-worker paired task and emitted-token measurements are pending. +Fresh ordinary-worker paired task and emitted-token measurements completed; see the updated evidence below. Further batch latency optimization remains necessary. The affected final CPU suite passes 73 cases. A stale warmup fixture originally patched the removed runner platform/environment hooks; it now exercises the existing bound warmup policy. No runtime compatibility workaround was added. + +## Updated acceptance target and completed clean control + +The requested costs refer to the **whole MTP4 verifier batch**, not its share +per request: C1 approximately 16 ms, C2 17–18 ms, C4 approximately 22 ms and +C8 approximately 28 ms. The previous 111 ms C8 result does not satisfy that +objective even when divided by eight. All further results retain five target +rows and four draft rows per request, and separate draft time. + +Fresh ordinary workers with profiling and NVTX disabled completed paired +root-selection controls, using the same capacity-eight host-indexer setup. +Common decode intervals exclude prefill and generation boundaries: + +| Concurrency | Retained independent selection tok/s | Root selection tok/s | +| ---: | ---: | ---: | +| 1 | 76.46 | 107.91 | +| 2 | 99.35 | 116.72 | +| 4 | 94.74 | 157.36 | +| 6 | 95.92 | 144.65 | +| 8 | 98.71 | 150.58 | + +Both runs pass the same task gates: eight seeded 32K exact-field retrieval +answers with natural EOS, executable deduplication code, automatic weather +tool invocation and its response. These are emitted-token speed evidence, +not replacement measurements for the verifier cost. The local GPU-indexer +C1 control is 14.655 ms and C2 25.472 ms at capacity two. A capacity-four +GPU-indexer experiment failed near full VRAM and is not accepted speed evidence. + +## Bounded shared-page indexer staging + +`sm70_qsa_indexer_stage_keys` defaults to false. A qualified host FP16 +H4/KV1/D128 indexer borrows one engine-owned device bank, bounded to 32768 +compressed keys rounded to physical pages (about 8 MiB at page size 204). +Every selection resets its map, deduplicates live physical pages across up to +eight requests, and refreshes their keys. The host history stays authoritative; +missing pages and overflow retain the existing scorer's paged-reader semantics. +The native `run_staged` entry reuses its original arithmetic; `run` keeps its ABI. +PP/DCP/PCP above one, other query layouts and absent native entries retain their +qualified paths. The bank never exposes borrowed output views and does not +cross engine configuration transfers. + +The short scorer test at 32K improves C8 shared-prefix scoring from 3.74 to +0.36 ms; independent histories improve approximately 1.9x. This is not a claim +that arbitrary independent contexts obtain the shared-prefix gain. Twenty-three +GPU tests pass, including changed key contents, page reorder, padding, missing +pages, graph replay, exact selection and a deliberately overflowing one-page bank. + +The full model diagnostic uses an 8 MiB QSA score tile and the same sampler, +256K capacity, prefix cache, 80% GPU / 20% host experts and MTP4. A previous +32 MiB score-tile run exhausted VRAM on a 64 MiB prefill score allocation; +that failure is retained. The bounded setting completes all 32K shared-prefix +streams with no incomplete TP4 verifier rounds: + +| Concurrency | Rounds | Whole verifier mean ms | p50 ms | p90 ms | p99 ms | Draft mean ms | +| ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 17.887 | 17.733 | 18.616 | 19.950 | 3.881 | +| 2 | 193 | 29.413 | 29.124 | 30.535 | 32.590 | 4.610 | +| 4 | 191 | 36.992 | 36.497 | 38.333 | 41.594 | 5.935 | +| 6 | 181 | 52.411 | 51.824 | 54.096 | 60.300 | 7.513 | +| 8 | 178 | 66.309 | 65.201 | 69.296 | 80.457 | 8.807 | + +These are diagnostic GPU events with critical-rank alignment, not clean +production throughput. They remain above the requested whole-batch costs. +The staging candidate's fresh clean paired throughput, long-output drift and +256K boundary gates remain pending; no default is promoted. + +## Follow-up candidates and rejected directions + +Large dense projections now have an opt-in `sm70_gguf.small_m_hmma_batch`. +It reuses the existing segment format and native decoder across token tiles; +qualified DMV Q4/IQ4 banks provide their equivalent full scale representation. +The six-array prepared-operand ABI is retained. The accelerated initialization +subset is K2560/N>=1536 or K1536/N2560 in the existing qualified formats and +layouts. Runtime acceleration is limited to M9..16; M1..8 and M17..40 +retain the previous projection routes, and larger batches retain canonical +restoration. The native batch operator has separately tested M1..40 coverage. The old native entry retains +its M1..32 contract; the batch entry is explicit. Eighty-six dense native GPU +cases and 111 affected storage/runtime cases pass, including capture/replay +and the M40/M41 fallback boundary. The initial wider M9..40 model experiment +produced C1/C2/C4/C6/C8 costs 17.931/26.655/38.265/51.922/63.368 ms: +C2 improves, but C4 regresses. That wider runtime admission is rejected; +M9..16 retains its C2 opportunity and preserves the prior M20 route. The +narrowed policy still needs matched full-model evidence. + +A single-grid HC experiment passed its small numerical check but did not +improve C8 (approximately 0.192 ms for both retained and candidate calls). +Its implementation is removed from the runtime and retained as a research +patch in the private artifacts. Existing eight-row HC tiles remain qualified. + +A separate default-off `sm70_gguf.grouped_gate_up` candidate reuses the +original lattice weight decoder for up to four routes to one expert. The +accelerated shape is M9..20, source types18/21/22, N divisible by32 and K by256, +with the existing Q8_1 intermediate. Routing is sorted and admitted on the GPU +each replay; outputs return directly to their original route order. Sparse +routes retain the original single-route algorithm. Changed-input graph tests +exercise both live admission outcomes and exact Q8_1 bytes; full-model benefit +remains unqualified. The first merged grouped/sparse kernel regressed sparse +routes because they inherited its register and shared-memory footprint; it is +rejected in favor of a separately guarded retained kernel. + +Rejected experiments are preserved in the task artifacts: + +* Parallel mapped-host indexer scoring in request-major and interleaved grids + slowed independent-history scoring approximately threefold. They are removed. +* The retained full-batch HC matrix path is slower at C1/C2 and improves C8 + only slightly; it is not promoted as a scaling fix. +* Expert lane-width and bank-layout changes improve the leading integer kernels + by only about 10%; they do not establish the requested scaling. +* Upstream MoE MMQ at M40 takes 1.6–1.9 ms for gate/up in representative real + banks, versus about 0.22 ms for the retained integer route; it is rejected. +* Sorting routed experts alone passes exact Q8 output checks but slows the + complete scorer by approximately 2–4%; its private native entry is removed. +* The expert-cache entry already holds 248 route IDs. At top-k ten the C8 + native path uses two twenty-row fragments. The 136 count in startup logs is + host weight slots, not request-entry capacity; it does not justify a ring fix. + +Hardware counter attribution is not available to the local unprivileged user +(`RmProfilingAdminOnly=1`). Cache/decode hypotheses are therefore based on +paired exact-shape experiments, not claimed Nsight Compute counters. + +## Combined narrowed dense and grouped-expert diagnostic + +The narrowed dense policy (M9..16) plus grouped experts completes all five +32K shared-prefix cases under the same capacity-eight, score-tile 8 MiB, +chunk 512, host-history contract. Whole verifier includes target forward, +sampling and state commit; draft remains separate. + +| Concurrency | Complete rounds | Whole verifier mean ms | p99 ms | Draft mean ms | +| ---: | ---: | ---: | ---: | ---: | +| 1 | 202 | 17.768 | 20.321 | 3.871 | +| 2 | 202 | 26.779 | 31.678 | 4.601 | +| 4 | 203 | 36.705 | 41.995 | 5.915 | +| 6 | 203 | 51.769 | 59.463 | 7.472 | +| 8 | 190 | 64.895 | 80.106 | 8.808 | + +There are no incomplete TP4 rounds. This combined experiment does not isolate +the contribution of grouping: the preceding wider dense policy also differs. +Relative to page staging alone, its C8 improvement is about 2.1%; this is not +the requested 28 ms C8 target. At C8 forward is 64.005 ms; sampling 0.860 ms and +state commit 0.030 ms cannot explain the inflation. A new matched node trace +will decide the next bottleneck rather than attributing it from micro gains. + +The eight submitted 32K retrieval/arithmetic tasks, executable Python function, +automatic weather-tool call and its reply all pass with natural EOS. The +quality run still has worker diagnostic fences, so clean emitted-token +qualification and the 256K boundary remain pending. The first model run uses +the owned native build immediately before a symbol rename/formatting pass; +the final formatted source artifacts are subsequently rebuilt and installed +atomically. The final artifacts pass 54 affected GPU cases, including changed-route +graph replay, independent expert output oracles and four-rank indexer state +checks. Clean model speed and longer-context gates remain separate before +any promotion. diff --git a/tests/config/test_flashnext_mtp_batch_policy.py b/tests/config/test_flashnext_mtp_batch_policy.py index a7bb878c81..d6a6613726 100644 --- a/tests/config/test_flashnext_mtp_batch_policy.py +++ b/tests/config/test_flashnext_mtp_batch_policy.py @@ -22,10 +22,16 @@ def test_mtp_batch_policy_is_owned_hashed_and_serializable(): sm70_hcx_batch_max_m=40, qsa_host_kv_state_compact=True, sm70_qsa_shared_key_batch=True, - sm70_gguf={"small_m_dp4a": True, "small_m_dp4a_max_m": 40}, + sm70_gguf={ + "small_m_dp4a": True, + "small_m_dp4a_max_m": 40, + "small_m_hmma_batch": True, + "grouped_gate_up": True, + }, qsa_host_kv=True, qsa_host_indexer_history=True, sm70_qsa_indexer_single_rank=True, + sm70_qsa_indexer_stage_keys=True, ) tiled.sm70_gguf.active = True context = multiprocessing.get_context("spawn") @@ -53,13 +59,21 @@ def test_mtp_batch_policy_is_owned_hashed_and_serializable(): "qsa_host_kv_state_compact", "sm70_qsa_shared_key_batch", "sm70_qsa_indexer_single_rank", + "sm70_qsa_indexer_stage_keys", ): changed = copy.deepcopy(tiled) setattr(changed, field, getattr(legacy, field)) assert changed.compute_hash() != tiled.compute_hash() transferred.sm70_gguf.small_m_dp4a_max_m = 20 assert transferred.compute_hash() != tiled.compute_hash() + + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.grouped_gate_up = False + assert transferred.compute_hash() != tiled.compute_hash() assert tiled.sm70_gguf.small_m_dp4a_max_m == 40 + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.small_m_hmma_batch = False + assert transferred.compute_hash() != tiled.compute_hash() def test_disabled_parents_keep_inactive_batch_options_out_of_hash(): @@ -69,6 +83,7 @@ def test_disabled_parents_keep_inactive_batch_options_out_of_hash(): sm70_qsa_shared_key=False, sm70_hcx_batch_max_m=40, sm70_qsa_shared_key_batch=True, + sm70_qsa_indexer_stage_keys=True, ) assert a.compute_hash() == b.compute_hash() @@ -82,6 +97,7 @@ def test_single_rank_indexer_is_inactive_without_host_history(host, history): qsa_host_kv=host, qsa_host_indexer_history=history, sm70_qsa_indexer_single_rank=True, + sm70_qsa_indexer_stage_keys=True, ) assert a.compute_hash() == b.compute_hash() diff --git a/tests/kernels/quantization/test_gguf_dense_segments.py b/tests/kernels/quantization/test_gguf_dense_segments.py index b7307fb747..5ecaf4a929 100644 --- a/tests/kernels/quantization/test_gguf_dense_segments.py +++ b/tests/kernels/quantization/test_gguf_dense_segments.py @@ -114,6 +114,92 @@ def test_reject_m33_without_launch(): ) +@pytest.mark.parametrize("source", [8, 12, 13, 14, 20, 23]) +@pytest.mark.parametrize( + "m,k,n,split,warps", + [ + (9, 2560, 512, 1, 4), + (10, 2560, 4096, 1, 4), + (17, 2560, 3584, 1, 8), + (20, 2560, 3328, 1, 4), + (25, 1536, 2560, 1, 8), + (30, 2560, 160, 5, 8), + (33, 2560, 4096, 1, 4), + (40, 2560, 3328, 1, 4), + ], +) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_batch_weight_reuse_official_and_changed_graph(source, m, k, n, split, warps): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + import vllm._C # noqa: F401 + + raw = raw_weights(source, n=n, k=k) + fmt, q, scale, minimum, group = decode(raw, source) + payload = [torch.from_numpy(p).cuda() for p in pack(fmt, q, scale, minimum, group)] + x = torch.randn(m, k, device="cuda", dtype=torch.float16) * 0.125 + output = torch.empty(m, n + 5, device="cuda", dtype=torch.float16)[:, 3 : n + 3] + tiles = (n + 31) // 32 + ws = torch.empty( + tiles * ((m + 7) // 8) * split * 256, device="cuda", dtype=torch.float32 + ) + counter = torch.zeros(tiles * ((m + 7) // 8), device="cuda", dtype=torch.int32) + official = torch.from_numpy(reconstruct(fmt, q, scale, minimum, group)).cuda() + + def run(): + torch.ops._C.gguf_dense_segments_batch_sm70_out( + x, + [payload[0]], + [payload[1]], + [payload[2]], + [output], + [fmt], + [n], + k, + split, + warps, + ws, + counter, + None, + ) + + run() + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for shift in (0.0, 0.125): + x.normal_(mean=shift, std=0.125) + graph.replay() + reference = x.float() @ official.float().T + torch.accelerator.synchronize() + assert bool(torch.isfinite(output).all()) + assert float((output.float() - reference).norm() / reference.norm()) < 0.003 + assert not bool(counter.any()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_batch_rejects_m41_before_launch(): + import vllm._C # noqa: F401 + + with pytest.raises(RuntimeError, match="1..40"): + torch.ops._C.gguf_dense_segments_batch_sm70_out( + torch.empty(41, 256, device="cuda", dtype=torch.float16), + [], + [], + [], + [], + [], + [], + 256, + 1, + 4, + torch.empty(1, device="cuda"), + torch.empty(1, device="cuda", dtype=torch.int32), + None, + ) + + @pytest.mark.parametrize("source", [8, 12, 13, 14, 20, 23]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_transient_restore_matches_turbomind_storage(source): diff --git a/tests/kernels/quantization/test_gguf_grouped_gate.py b/tests/kernels/quantization/test_gguf_grouped_gate.py new file mode 100644 index 0000000000..d6458e0114 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_grouped_gate.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Live grouping changes preserve the retained Q8_1 expert output bytes.""" + +import numpy as np +import pytest +import torch +import vllm._C # noqa: F401 + +from vllm.model_executor.layers.quantization.gguf_grouped_gate import grouped_gate_up +from vllm.model_executor.layers.quantization.gguf_raw import RawGGUFProjection +from vllm.transformers_utils.gguf_tensor_reader import quant_size + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("source", [18, 21, 22]) +@pytest.mark.parametrize("m", [9, 10, 20]) +@pytest.mark.parametrize("index_dtype", [torch.int32, torch.int64]) +def test_grouped_experts_changed_route_graph_exact(source, m, index_dtype): + if torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 required") + experts, n, k, top_k = 256, 32, 256, 10 + block, size = quant_size(source) + rng = np.random.default_rng(20261015 + source) + weights = [] + for _ in range(2): + raw = rng.integers(0, 256, (experts * n, k // block, size), dtype=np.uint8) + raw[..., :2] = np.array([0.0007], dtype="= 0).sum()) == 1 + assert int(count.item()) <= cache.shape[0] @pytest.mark.parametrize("requests", [2, 8]) @@ -99,6 +117,7 @@ def test_grouped_selection_matches_independent_native_requests(requests): positions = torch.full((rows,), 4095, device="cuda", dtype=torch.int64) lengths = torch.full((requests,), 4096, device="cuda", dtype=torch.int32) starts = torch.arange(0, rows + 1, 5, dtype=torch.int32) + owner = QSAIndexerStagingWorkspace() def run(): return qsa_select_paged_tokens( @@ -113,6 +132,7 @@ def run(): query_start_loc_cpu=starts, shared_key_scoring=True, shared_key_batch_scoring=True, + indexer_staging_workspace=owner, ) run() diff --git a/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py b/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py index 5decfd764d..fe722b3c2b 100644 --- a/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py +++ b/tests/models/qwen4_exp/test_sm70_hcx_dispatch.py @@ -69,10 +69,12 @@ def test_small_m_hc_consumes_unrounded_moe_pair(): # both originals and sum in FP32 instead of consuming that rounded sum. fused = torch.full((5, 4), 2048, dtype=torch.float16) shared = torch.ones_like(fused) + hcx.register_moe_runner("p", SimpleNamespace(_hcx_batch_max_m=8)) partial = hcx._moe_output(shared, fused, "p", 4) seen = [] - def run(first, *args, secondary): + def run(first, *args, secondary, batch_max_m): + assert batch_max_m == 8 seen.append((first, secondary)) return args[0], first, args[1] @@ -235,7 +237,7 @@ def __init__(self): ) self.partial = False - def enable_partial_inputs(self, name, runtime=None): + def enable_partial_inputs(self, name, runtime=None, *, batch_max_m=8): self.partial = True class MoE: @@ -283,7 +285,8 @@ def test_compiled_moe_payload_keeps_both_dependencies_live(): if not torch._C._dispatch_has_kernel_for_dispatch_key("vllm::" + name, "CPU"): library.impl(name, function) - def native(first, hidden, injection, *args, secondary): + def native(first, hidden, injection, *args, secondary, batch_max_m): + assert batch_max_m == 8 block = ((first.float() + secondary.float()) * 0.5).half() return hidden.clone(), block, injection.clone() diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 3af36adda5..f6a6cc117d 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -584,6 +584,10 @@ class Sm70GgufConfig: small_m_hmma: bool = True """Use one coalesced integer bank and fused FP16 MMA projections at M1..8.""" + small_m_hmma_batch: bool = False + """Reuse large resident projection weights at M9..16 on qualified shapes.""" + grouped_gate_up: bool = False + """Reuse original expert decoding across M9..20 routed rows on SM70.""" q8_expert_intermediate: bool = True """Encode routed intermediates once in qualified integer expert gate/up.""" @@ -894,6 +898,12 @@ def resolve_sm70_rmsnorm_gated(self, *, qualified: bool) -> None: Applies to qualified SM70 FP16 TP4/PP1/CP1 small batches. Each rank still updates its own key caches. Other shapes retain independent selection. """ + sm70_qsa_indexer_stage_keys: bool = False + """Copy live shared physical indexer pages once per small verification batch. + + Qualified SM70 FP16 host indexer history borrows a bounded engine-owned + device bank. Overflow pages retain the authoritative paged reader. + """ qsa_host_kv: bool = False """Keep QSA attention history in pinned host storage on SM70.""" qsa_host_kv_dtype: Literal["fp8_e4m3", "float16"] = "fp8_e4m3" @@ -1115,12 +1125,14 @@ def compute_hash(self) -> str: ignored_factors.add("sm70_hcx_batch_max_m") if not self.sm70_qsa_shared_key: ignored_factors.add("sm70_qsa_shared_key_batch") + ignored_factors.add("sm70_qsa_indexer_stage_keys") if ( not self.qsa_host_kv or not self.qsa_host_indexer_history or self.qsa_host_kv_device_reference ): ignored_factors.add("sm70_qsa_indexer_single_rank") + ignored_factors.add("sm70_qsa_indexer_stage_keys") if not self.fused_fp16_aux_gemv_applicable: ignored_factors.add("fused_fp16_aux_gemv") if not self.qsa_auto_e4m3_active: diff --git a/vllm/model_executor/layers/quantization/gguf_dense_hmma.py b/vllm/model_executor/layers/quantization/gguf_dense_hmma.py index 9fe648ce37..9ee5a5acfb 100644 --- a/vllm/model_executor/layers/quantization/gguf_dense_hmma.py +++ b/vllm/model_executor/layers/quantization/gguf_dense_hmma.py @@ -108,23 +108,40 @@ def prepare_segment_bank(projection, canonical, device): projection.stats = torch.nn.Parameter(payload[2], False) projection.register_parameter("segment_high", torch.nn.Parameter(payload[1], False)) projection.segment_format = fmt + projection.gguf_batch_stats = ( + projection.stats + if cfg is not None + and cfg.kernel_config.sm70_gguf.small_m_hmma_batch + and hasattr(torch.ops._C, "gguf_dense_segments_batch_sm70_out") + and ((k == 2560 and n >= 1536) or (k, n) == (1536, 2560)) + else None + ) projection.cache_capabilities = () workspace(device, k == 160 or n <= 320) return True def apply_segments( - rows, codes, scales, high, formats, ns, output, views, *, blas_bands=() + rows, codes, scales, high, formats, ns, output, views, *, blas_bands=(), batch=False ): k = rows.shape[1] shared = k == 160 or sum(ns) <= 320 scratch = workspace(rows.device, shared) - if 1 <= rows.shape[0] <= 8 or rows.shape[0] == 20: + if ( + 1 <= rows.shape[0] <= 8 + or rows.shape[0] == 20 + or (batch and rows.shape[0] <= 40) + ): warps = 4 if k == 2560 and sum(ns) > 320 else 8 if k == 1536: warps = 8 split = 5 if k == 2560 and sum(ns) <= 320 else 1 - torch.ops._C.gguf_dense_segments_sm70_out( + op = ( + torch.ops._C.gguf_dense_segments_batch_sm70_out + if batch and rows.shape[0] > 8 + else torch.ops._C.gguf_dense_segments_sm70_out + ) + op( rows, codes, high, diff --git a/vllm/model_executor/layers/quantization/gguf_dmv.py b/vllm/model_executor/layers/quantization/gguf_dmv.py index 51b07912d1..628a7bff78 100644 --- a/vllm/model_executor/layers/quantization/gguf_dmv.py +++ b/vllm/model_executor/layers/quantization/gguf_dmv.py @@ -177,6 +177,24 @@ def reorder(a): ), ) projection.dmv_format = fmt + projection.register_buffer("gguf_batch_stats", None, persistent=False) + cfg = get_current_vllm_config_or_none() + if ( + cfg is not None + and cfg.kernel_config.sm70_gguf.small_m_hmma + and cfg.kernel_config.sm70_gguf.small_m_hmma_batch + and fmt in (0, 3) + and ((k == 2560 and n >= 1536) or (k, n) == (1536, 2560)) + and hasattr(torch.ops._C, "gguf_dense_segments_batch_sm70_out") + ): + full_scale = scale + if fmt == 3: + # Invert compact_lut4_scale without reconstructing the weights. + half_bits = scale.view(np.uint16).astype(np.uint32) + full_scale = (half_bits | (half_bits << 16)).view(np.uint8) + projection.gguf_batch_stats = torch.from_numpy( + np.ascontiguousarray(full_scale) + ).to(raw_weight.device) projection.cache_capabilities = () projection.dmv_capability = GGUFOperatorCapability( decoder_family(kind), diff --git a/vllm/model_executor/layers/quantization/gguf_grouped_gate.py b/vllm/model_executor/layers/quantization/gguf_grouped_gate.py new file mode 100644 index 0000000000..74ee493ca8 --- /dev/null +++ b/vllm/model_executor/layers/quantization/gguf_grouped_gate.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Grouped original integer experts, with live sparsity admission on device.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _sort_routes( + Ids, + Sorted, + Tokens, + Order, + Mode, + Routes: tl.constexpr, + TopK: tl.constexpr, + Block: tl.constexpr, +): + i = tl.arange(0, Block) + ids = tl.load(Ids + i, i < Routes, other=2147483647) + packed = tl.where(i < Routes, ids * Block + i, 2147483647) + ranked = tl.sort(packed, descending=False) + experts, original = ranked // Block, ranked % Block + previous = tl.gather(experts, tl.maximum(i - 1, 0), 0) + distinct = tl.sum( + ((i < Routes) & ((i == 0) | (experts != previous))).to(tl.int32), 0 + ) + # At least one third of routes must reuse an expert before grouping pays. + grouped = distinct * 3 <= Routes * 2 + tl.store(Mode, grouped.to(tl.int32)) + tl.store(Sorted + i, experts, i < Routes) + tl.store(Tokens + i, original // TopK, i < Routes) + tl.store(Order + i, original, i < Routes) + + +def grouped_gate_up(out, activation, ids, gate, up, source_type, bank_aware): + """Return False before launch when the native grouping contract is absent.""" + if ( + not 9 <= ids.shape[0] <= 20 + or ids.dtype not in (torch.int32, torch.int64) + or source_type not in (18, 21, 22) + or out.dtype != torch.uint8 + or not hasattr(torch.ops._C, "gguf_dp4a_gate_up_grouped_sm70_out") + ): + return False + # Router IDs are bounded by the expert bank. Keep the native route mapping + # int32 without reading device sparsity or IDs on the host. + ids = ids.to(dtype=torch.int32).contiguous() + sorted_ids = torch.empty_like(ids) + tokens = torch.empty_like(ids) + order = torch.empty_like(ids) + mode = torch.empty(1, dtype=torch.int32, device=ids.device) + _sort_routes[(1,)]( + ids, + sorted_ids, + tokens, + order, + mode, + ids.numel(), + ids.shape[1], + triton.next_power_of_2(ids.numel()), + num_warps=4, + ) + torch.ops._C.gguf_dp4a_gate_up_grouped_sm70_out( + out, + activation, + sorted_ids, + gate, + up, + source_type, + bank_aware, + tokens, + order, + ids, + mode, + ) + return True diff --git a/vllm/model_executor/layers/quantization/gguf_moe.py b/vllm/model_executor/layers/quantization/gguf_moe.py index 7b6b5d7ca8..38e7561045 100644 --- a/vllm/model_executor/layers/quantization/gguf_moe.py +++ b/vllm/model_executor/layers/quantization/gguf_moe.py @@ -49,6 +49,7 @@ def _original_expert_dp4a( left, quantized_hidden, bank_aware, + grouped_gate: bool = False, ): m, top_k = ids.shape n = gate.shape[1] @@ -60,17 +61,23 @@ def _original_expert_dp4a( ) output = torch.empty_like(x) torch.ops._C.gguf_quantize_q8_1_sm70_out(q8, x) - torch.ops._C.gguf_dp4a_gate_up_sm70_out( - hidden, - q8, - ids, - gate, - up, - source_type, - True, - 4 if quantized_hidden and source_type in (20, 23) else 16, - bank_aware, - ) + grouped = False + if grouped_gate and quantized_hidden: + from .gguf_grouped_gate import grouped_gate_up + + grouped = grouped_gate_up(hidden, q8, ids, gate, up, source_type, bank_aware) + if not grouped: + torch.ops._C.gguf_dp4a_gate_up_sm70_out( + hidden, + q8, + ids, + gate, + up, + source_type, + True, + 4 if quantized_hidden and source_type in (20, 23) else 16, + bank_aware, + ) torch.ops._C.gguf_dp4a_raw_down_unroute_sm70_out( output, hidden, @@ -95,6 +102,7 @@ def _original_expert_dp4a_fake( left, quantized_hidden, bank_aware, + grouped_gate: bool = False, ): return torch.empty_like(x) @@ -109,6 +117,45 @@ def _original_expert_dp4a_fake( ) +def _grouped_expert_dp4a( + x, + ids, + probabilities, + gate, + up, + down, + source_type, + down_type, + left, + quantized_hidden, + bank_aware, +): + return _original_expert_dp4a( + x, + ids, + probabilities, + gate, + up, + down, + source_type, + down_type, + left, + quantized_hidden, + bank_aware, + grouped_gate=True, + ) + + +register_policy_op( + "gguf_grouped_expert_dp4a", + "(Tensor x, Tensor ids, Tensor probabilities, Tensor gate, Tensor up, " + "Tensor down, int source_type, int down_type, int left, " + "bool quantized_hidden, bool bank_aware) -> Tensor", + _grouped_expert_dp4a, + _original_expert_dp4a_fake, +) + + def _original_expert_reduce(down, probabilities): return moe_fused_mul_sum(down, probabilities) @@ -143,6 +190,9 @@ def __init__(self, quant_config: "GGUFConfig", moe): self.q8_intermediate = bool( config and config.kernel_config.sm70_gguf.q8_expert_intermediate ) + self.grouped_gate_up = bool( + config and config.kernel_config.sm70_gguf.grouped_gate_up + ) self.expert_trace_tokens = ( config.kernel_config.sm70_gguf.expert_trace_tokens if config else 0 ) @@ -345,6 +395,25 @@ def process_weights_after_loading(self, layer): "storage": "original", "accumulation": "FP32", } + grouped_shape = ( + self.dp4a_admitted + and self.q8_intermediate + and self.weight_types["w1"] in (18, 21, 22) + ) + if ( + self.grouped_gate_up + and grouped_shape + and not hasattr(torch.ops._C, "gguf_dp4a_gate_up_grouped_sm70_out") + ): + raise ValueError("Grouped GGUF experts require the matching native build") + self.grouped_gate_up = self.grouped_gate_up and grouped_shape + self.native_admission["grouped_gate_up"] = { + "enabled": self.grouped_gate_up, + "min_m": 9, + "max_m": 20, + "device_admission": "distinct_experts_at_most_two_thirds_of_routes", + "fallback": "original_single_route_integer_dot", + } def get_fused_moe_quant_config(self, layer): return None @@ -391,6 +460,11 @@ def apply( if step < 1: raise ValueError("Expert cache cannot hold one routed token") output = torch.empty_like(x) if tokens > step else None + operation = ( + torch.ops.vllm.gguf_grouped_expert_dp4a + if getattr(self, "grouped_gate_up", False) + else torch.ops.vllm.gguf_original_expert_dp4a + ) for start in range(0, tokens, step): stop = min(start + step, tokens) fragment_ids = ids[start:stop].contiguous() @@ -398,7 +472,7 @@ def apply( # Finish each fragment before resolving the next one, which # may evict its slots. The existing model stream owns both. fragment_ids = cache.resolve(self.cache_layer, fragment_ids) - fragment = torch.ops.vllm.gguf_original_expert_dp4a( + fragment = operation( x.contiguous() if output is None else x[start:stop].contiguous(), fragment_ids, topk_weights[start:stop].contiguous(), diff --git a/vllm/model_executor/layers/quantization/gguf_turbomind.py b/vllm/model_executor/layers/quantization/gguf_turbomind.py index f50ec4b921..245e50a13c 100644 --- a/vllm/model_executor/layers/quantization/gguf_turbomind.py +++ b/vllm/model_executor/layers/quantization/gguf_turbomind.py @@ -79,6 +79,31 @@ def _prepared_gguf_projection( # Keep M-dependent policy behind an opaque op: vLLM's range compilation # drops Dynamo guards and would otherwise freeze the prefill branch. rows = x.reshape(-1, x.shape[-1]).contiguous() + # Batch variants keep their original small-row and prefill layouts. + # Family 6 carries compact DMV stats plus its explicit full-scale bank. + batch_rows = 9 <= rows.shape[0] <= 16 + if family == 6: + if cache is None: + raise ValueError("Batch DMV requires its explicit full-scale bank") + if batch_rows: + from .gguf_dense_hmma import apply_segments + + out = rows.new_empty((rows.shape[0], output_size)) + apply_segments( + rows, + [codes], + [cache], + [cache[:0]], + [decoder], + [output_size], + out, + [out], + batch=True, + ) + return out.reshape(*x.shape[:-1], logical_size) + family, cache = 4, cache[:0] + elif family == 5 and not batch_rows: + family = 3 if family == 4: from .gguf_dmv import restore @@ -97,13 +122,25 @@ def _prepared_gguf_projection( blas_bands, ) return output.reshape(*x.shape[:-1], logical_size) - if family == 3: + if family in (3, 5): from .gguf_dense_hmma import apply_segments, restore_and_apply out = torch.empty((rows.shape[0], output_size), dtype=x.dtype, device=x.device) - if rows.shape[0] <= 8 or rows.shape[0] == 20: + if ( + rows.shape[0] <= 8 + or rows.shape[0] == 20 + or (family == 5 and rows.shape[0] <= 40) + ): apply_segments( - rows, [codes], [stats], [cache], [decoder], [output_size], out, [out] + rows, + [codes], + [stats], + [cache], + [decoder], + [output_size], + out, + [out], + batch=family == 5, ) else: restore_and_apply( @@ -193,6 +230,18 @@ def _prepared_gguf_mixed_projection( ) -> torch.Tensor: rows = x.reshape(-1, x.shape[-1]).contiguous() specs = [descriptors[i : i + 9] for i in range(0, len(descriptors), 9)] + batch_rows = 9 <= rows.shape[0] <= 16 + if not batch_rows: + caches = list(caches) + for i, spec in enumerate(specs): + if spec[0] == 5: + spec[0] = 3 + elif spec[0] == 6: + spec[0] = 4 + full_scale = caches[i] + if full_scale is None: + raise ValueError("Batch DMV requires its explicit full-scale bank") + caches[i] = full_scale[:0] cache_offset = blas_offset = 0 policies = [] for spec in specs: @@ -205,7 +254,8 @@ def _prepared_gguf_mixed_projection( ) cache_offset += cache_count blas_offset += blas_count - if specs and all(spec[0] == 3 for spec in specs): + batched = batch_rows and all(spec[0] in (5, 6) for spec in specs) + if specs and (all(spec[0] == 3 for spec in specs) or batched): from .gguf_dense_hmma import apply_segments ns = [spec[6] for spec in specs] @@ -214,19 +264,26 @@ def _prepared_gguf_mixed_projection( apply_segments( rows, codes, - stats, - caches, + [ + cache if spec[0] == 6 else stat + for cache, stat, spec in zip(caches, stats, specs) + ], + [ + cache[:0] if spec[0] == 6 and cache is not None else cache + for cache, spec in zip(caches, specs) + ], [spec[1] for spec in specs], ns, output, views, blas_bands=[bb for _, bb in policies], + batch=batched, ) return output.reshape(*x.shape[:-1], sum(ns)) # Retain the calibrated per-projection policy outside measured target # verification sizes. This decision must use actual M inside the op. direct = ( - not any(spec[0] in (3, 4) for spec in specs) + not any(spec[0] in (3, 4, 5, 6) for spec in specs) and rows.shape[0] in (5, 20) and all( not _supports_band(rows.shape[0], cb) @@ -371,11 +428,21 @@ def prepared_projection_arguments(projections): if hasattr(projection, "segment_format"): family, decoder = 3, projection.segment_format cb = [] + if getattr(projection, "gguf_batch_stats", None) is not None: + family, decoder = ( + 6 if hasattr(projection, "dmv_format") else 5, + getattr( + projection, "segment_format", getattr(projection, "dmv_format", 0) + ), + ) + cb = [] caches.append( - projection.dmv_high + projection.gguf_batch_stats + if family == 6 + else projection.dmv_high if family == 4 else projection.segment_high - if family == 3 + if family in (3, 5) else projection.fp16_cache ) descriptors.extend( @@ -701,16 +768,21 @@ def forward(self, x): family, decoder = 4, self.dmv_format if hasattr(self, "segment_format"): family, decoder = 3, self.segment_format + batch_stats = getattr(self, "gguf_batch_stats", None) + if batch_stats is not None: + family = 6 if hasattr(self, "dmv_format") else 5 capabilities = getattr(kernel, "operator_capabilities", ()) blas = tuple(c for c in capabilities if "blas" in c.operator) return torch.ops.vllm.prepared_gguf_projection( x, self.codes, self.stats, - self.dmv_high + batch_stats + if family == 6 + else self.dmv_high if family == 4 else self.segment_high - if family == 3 + if family in (3, 5) else self.fp16_cache, family, decoder, diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index 67f98847e0..a15ed056e2 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -135,7 +135,31 @@ def __init__( vllm_config.kernel_config.sm70_qsa_shared_key_batch ) self.selection_tp_group = None + self.indexer_staging_workspace = None parallel = vllm_config.parallel_config + if ( + vllm_config.kernel_config.sm70_qsa_indexer_stage_keys + and self.shared_key_scoring + and vllm_config.kernel_config.qsa_host_kv + and vllm_config.kernel_config.qsa_host_kv_active + and vllm_config.kernel_config.qsa_host_indexer_history + and not vllm_config.kernel_config.qsa_host_kv_device_reference + and vllm_config.model_config.dtype == torch.float16 + and current_platform.is_device_capability(70) + and (self.index_n_heads, self.index_kv_heads, self.index_head_dim) + == (4, 1, 128) + and parallel.pipeline_parallel_size == 1 + and parallel.decode_context_parallel_size == 1 + and parallel.prefill_context_parallel_size == 1 + ): + from vllm.runtime_resources import runtime_resources_for + + from .ops.qsa_indexer_staging import QSAIndexerStagingWorkspace + + resources = runtime_resources_for(vllm_config) + self.indexer_staging_workspace = resources.setdefault( + "qsa_indexer_staging", QSAIndexerStagingWorkspace() + ) if ( vllm_config.kernel_config.sm70_qsa_indexer_single_rank and vllm_config.kernel_config.qsa_host_kv @@ -414,6 +438,7 @@ def _select_local_tokens( shared_key_scoring=self.shared_key_scoring, shared_key_batch_scoring=self.shared_key_batch_scoring, max_seq_len=compressed_metadata.max_seq_len, + indexer_staging_workspace=self.indexer_staging_workspace, ) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index 9e09a1912e..d8f06db519 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -1179,6 +1179,7 @@ def qsa_mqa_paged( score_scale: float | None = None, shared_key_scoring: bool = False, shared_key_query_start_loc_cpu: torch.Tensor | None = None, + indexer_staging_workspace=None, ) -> tuple[torch.Tensor, torch.Tensor]: """Compute QSA scores directly from a paged compressed-key cache.""" @@ -1227,6 +1228,7 @@ def qsa_mqa_paged( columns, score_divisor, shared_key_query_start_loc_cpu, + indexer_staging_workspace, ) if result is not None: return result @@ -1241,6 +1243,7 @@ def qsa_mqa_paged( compress_ratio, columns, score_divisor, + indexer_staging_workspace, ) if result is not None: return result @@ -1619,6 +1622,7 @@ def qsa_select_paged_tokens( shared_key_scoring: bool = False, max_seq_len: int | None = None, shared_key_batch_scoring: bool = False, + indexer_staging_workspace=None, ) -> torch.Tensor: """Score, select, and expand QSA indices without host synchronization. @@ -1734,6 +1738,7 @@ def qsa_select_paged_tokens( if shared_key_batch_scoring and row_start == 0 and row_end == rows else None ), + indexer_staging_workspace=indexer_staging_workspace, ) blocks = blocks_buffer[: row_end - row_start] use_cooperative_topk = ( diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer_staging.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer_staging.py new file mode 100644 index 0000000000..95429f284c --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer_staging.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bounded, per-engine indexer staging; shared physical pages are copied once. + +The host history remains authoritative. Every selection rebuilds the page +plan and copies live keys, so tentative overwrites and page reuse cannot leave +a stale device copy. Pages outside the bounded bank retain the paged reader. +Model execution orders borrowers on its stream; no staged views escape. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _plan_pages( + Table, + Requests, + Positions, + Lengths, + Slots, + Pages, + Count, + TableStride: tl.constexpr, + TableColumnStride: tl.constexpr, + Width: tl.constexpr, + Rows: tl.constexpr, + CachePages: tl.constexpr, + Capacity: tl.constexpr, + Page: tl.constexpr, + Columns: tl.constexpr, + Compress: tl.constexpr, + BlockRows: tl.constexpr, +): + request = tl.program_id(1) + rows = tl.arange(0, BlockRows) + req = tl.load(Requests + rows, rows < Rows, other=-1) + position = tl.load(Positions + rows, rows < Rows, other=-1) + length = tl.load(Lengths + request) + end = tl.max(tl.where(req == request, tl.minimum(position + 1, length), 0), 0) + end = tl.minimum(tl.maximum(end // Compress, 0), Columns) + logical = tl.program_id(0) * 64 + tl.arange(0, 64) + physical = tl.load( + Table + request * TableStride + logical * TableColumnStride, + (logical < Width) & (logical * Page < end), + other=-1, + ) + valid = (physical >= 0) & (physical < CachePages) + old = tl.atomic_cas( + Slots + tl.where(valid, physical, CachePages), + tl.full((64,), -1, tl.int32), + tl.full((64,), -2, tl.int32), + sem="relaxed", + ) + claimed = valid & (old == -1) + slot = tl.atomic_add( + Count + tl.zeros((64,), tl.int32), claimed.to(tl.int32), sem="relaxed" + ) + tl.store(Pages + slot, physical, claimed & (slot < Capacity)) + tl.store( + Slots + tl.maximum(physical, 0), + tl.where(slot < Capacity, slot, -3), + claimed, + ) + + +@triton.jit +def _copy_pages( + Cache, + Pages, + Count, + Staged, + StridePage: tl.constexpr, + StrideToken: tl.constexpr, + StrideDim: tl.constexpr, + Page: tl.constexpr, + Capacity: tl.constexpr, +): + slot = tl.program_id(0) + count = tl.load(Count) + if slot < tl.minimum(count, Capacity): + physical = tl.load(Pages + slot) + token = tl.program_id(1) * 16 + tl.arange(0, 16) + dim = tl.arange(0, 128) + key = tl.load( + Cache + + physical.to(tl.int64) * StridePage + + token[:, None] * StrideToken + + dim[None, :] * StrideDim, + (token < Page)[:, None], + other=0.0, + ) + tl.store( + Staged + (slot * Page + token[:, None]) * 128 + dim[None, :], + key, + (token < Page)[:, None], + ) + + +class QSAIndexerStagingWorkspace: + """One bounded bank borrowed by successive indexer layers of an engine.""" + + def __init__(self, tokens: int = 32768): + if tokens <= 0: + raise ValueError("Indexer staging capacity must be positive") + self.tokens = tokens + self.banks: dict[tuple, tuple[torch.Tensor, ...]] = {} + + def prepare(self, cache, table, requests, positions, lengths, compress, columns): + page = cache.shape[1] + capacity = min(cache.shape[0], triton.cdiv(self.tokens, page)) + key = (cache.device, cache.shape[0], page) + if key not in self.banks: + self.banks[key] = ( + torch.empty( + (capacity, page, 1, 128), dtype=torch.float16, device=cache.device + ), + torch.empty(cache.shape[0] + 1, dtype=torch.int32, device=cache.device), + torch.empty(capacity, dtype=torch.int32, device=cache.device), + torch.empty(1, dtype=torch.int32, device=cache.device), + ) + staged, slots, pages, count = self.banks[key] + slots.fill_(-1) + slots[-1:].fill_(-3) + count.zero_() + _plan_pages[(triton.cdiv(table.shape[1], 64), table.shape[0])]( + table, + requests, + positions, + lengths, + slots, + pages, + count, + table.stride(0), + table.stride(1), + table.shape[1], + requests.numel(), + cache.shape[0], + capacity, + page, + columns, + compress, + triton.next_power_of_2(requests.numel()), + num_warps=4, + ) + _copy_pages[(capacity, triton.cdiv(page, 16))]( + cache, + pages, + count, + staged, + cache.stride(0), + cache.stride(1), + cache.stride(3), + page, + capacity, + num_warps=4, + ) + return staged, slots + + def close(self): + self.banks.clear() diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py index 83efbfb1d1..9df6a47f42 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_shared_key.py @@ -35,6 +35,25 @@ def fake( ): return None + if hasattr(torch.ops.vllm_sm70_qsa_indexer, "run_staged"): + + @torch.library.register_fake("vllm_sm70_qsa_indexer::run_staged") + def fake_staged( + scores, + visible, + q, + cache, + table, + requests, + positions, + lengths, + divisor, + compress, + staged, + page_slots, + ): + return None + return True @@ -59,7 +78,16 @@ def shared_key_reason(q, cache, table, requests, positions, lengths): def maybe_shared_key_scores( - q, cache, table, requests, positions, lengths, compress, columns, divisor + q, + cache, + table, + requests, + positions, + lengths, + compress, + columns, + divisor, + staging_workspace=None, ): reason = shared_key_reason(q, cache, table, requests, positions, lengths) if reason is not None: @@ -67,7 +95,21 @@ def maybe_shared_key_scores( return None scores = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) visible = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) - torch.ops.vllm_sm70_qsa_indexer.run( + staged = ( + staging_workspace.prepare( + cache, table, requests, positions, lengths, compress, columns + ) + if staging_workspace is not None + and table.shape[0] <= 8 + and hasattr(torch.ops.vllm_sm70_qsa_indexer, "run_staged") + else () + ) + op = ( + torch.ops.vllm_sm70_qsa_indexer.run_staged + if staged + else torch.ops.vllm_sm70_qsa_indexer.run + ) + op( scores, visible, q, @@ -78,6 +120,7 @@ def maybe_shared_key_scores( lengths, float(divisor), compress, + *staged, ) logger.info_once("Using SM70 QSA shared-key indexer for M=%d.", q.shape[0]) return scores, visible @@ -94,6 +137,7 @@ def maybe_shared_key_batch_scores( columns, divisor, query_start_loc_cpu, + staging_workspace=None, ): """Reuse the single-request provider without changing batched top-k. @@ -136,9 +180,23 @@ def maybe_shared_key_batch_scores( return None scores = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) visible = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) + staged = ( + staging_workspace.prepare( + cache, table, requests, positions, lengths, compress, columns + ) + if staging_workspace is not None + and table.shape[0] <= 8 + and hasattr(torch.ops.vllm_sm70_qsa_indexer, "run_staged") + else () + ) + op = ( + torch.ops.vllm_sm70_qsa_indexer.run_staged + if staged + else torch.ops.vllm_sm70_qsa_indexer.run + ) for i, start, end in groups: local_requests = (requests[start:end] - i).contiguous() - torch.ops.vllm_sm70_qsa_indexer.run( + op( scores[start:end], visible[start:end], q[start:end], @@ -149,8 +207,11 @@ def maybe_shared_key_batch_scores( lengths[i : i + 1], float(divisor), compress, + *staged, ) logger.info_once( "Using SM70 shared-key indexer for %d request groups.", len(groups) ) + if staged: + logger.info_once("Using bounded SM70 indexer page staging with paged overflow.") return scores, visible From b5d7d7a1623cbb8ef3e5f337a84a910bd7f23557 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:31:52 +0800 Subject: [PATCH 13/19] [Perf] Reuse routed expert weights and fuse qualified HC batches Preserve legacy operators and default policies while adding bounded forty-row cache leases, grouped original down projections, and the full-mesh HC batch path. Fix mixed dense batch fusion and qualify graph/state ownership and paired task behavior. Keep C2 on retained HC tiles after the ordinary-worker acceptance regression; document whole-verifier curves, emitted-token evidence, and the unpromoted shared-expert experiment separately. Assisted-by: Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- csrc/ops.h | 25 ++ csrc/sm70_turbomind/ops/gguf_dp4a.cuh | 55 ++++ csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu | 2 + .../ops/gguf_grouped_down_sm70.cuh | 140 ++++++++++ .../ops/gguf_grouped_gate_sm70.cuh | 35 ++- csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh | 259 ++++++++++++++++++ csrc/sm70_turbomind/ops/hcx_sm70.cu | 59 +++- csrc/torch_bindings.cpp | 25 ++ docs/design/architecture/runtime_reference.md | 4 +- .../flashnext_mtp_batch_16gb_20261015.md | 245 ++++++++++++++++- .../config/test_flashnext_mtp_batch_policy.py | 23 ++ .../quantization/test_gguf_grouped_down.py | 96 +++++++ .../quantization/test_gguf_grouped_gate.py | 39 ++- .../quantization/test_gguf_mtp_batch.py | 11 +- .../quantization/test_gguf_segment_prefill.py | 56 ++++ .../quantization/test_gguf_shared_batch.py | 144 ++++++++++ .../kernels/test_sm70_hcx_compiled_payload.py | 17 +- .../qwen4_exp/test_gguf_expert_cache.py | 49 +++- .../qwen4_exp/test_sm70_hcx_dispatch.py | 52 ++++ vllm/config/kernel.py | 37 ++- .../layers/quantization/gguf_dense_hmma.py | 78 +++--- .../layers/quantization/gguf_expert_cache.py | 25 +- .../layers/quantization/gguf_grouped_gate.py | 58 +++- .../layers/quantization/gguf_moe.py | 185 +++++++++++-- .../layers/quantization/gguf_turbomind.py | 18 +- vllm/models/qwen4_exp/nvidia/sm70_hcx.py | 50 +++- 26 files changed, 1675 insertions(+), 112 deletions(-) create mode 100644 csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh create mode 100644 csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh create mode 100644 tests/kernels/quantization/test_gguf_grouped_down.py create mode 100644 tests/kernels/quantization/test_gguf_shared_batch.py diff --git a/csrc/ops.h b/csrc/ops.h index 69a1176adb..d0cb52d697 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -1199,3 +1199,28 @@ void gguf_dp4a_gate_up_grouped_sm70_out( torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, torch::Tensor mode); + +void gguf_dp4a_gate_up_grouped_batch_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode); + +void gguf_dp4a_raw_down_grouped_sm70_out( + torch::Tensor out, torch::Tensor projected, torch::Tensor input, + torch::Tensor ids, torch::Tensor order, torch::Tensor probabilities, + torch::Tensor weights, int64_t source_type, int64_t left); + +void sm70_hcx_batch_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule); diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh index 437b08182e..b850830c1a 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a.cuh +++ b/csrc/sm70_turbomind/ops/gguf_dp4a.cuh @@ -220,6 +220,61 @@ using IQ3SDot = LatticeDot<21>; // canonical reader, without expanding a second expert bank at load time. template struct OriginalIntegerDot { + struct Decoded { + int values[8]; + float scale; + }; + __device__ static Decoded decode(const uint8_t* row, int group) { + const uint8_t* b; + const uint8_t* codes; + float scale; + if constexpr (Type == 23) { + b = row + (group / 8) * 136; + const int sub = group % 8; + const uint16_t high = *reinterpret_cast(b + 2); + const int subscale = ((b[4 + sub / 2] >> (4 * (sub % 2))) & 15) | + (((high >> (2 * sub)) & 3) << 4); + // Match the canonical LUT reader's expanded FP16 coefficients. + scale = __half2float(__float2half_rn( + __half2float(*reinterpret_cast(b)) * (subscale - 32))); + codes = b + 8 + sub * 16; + } else { + b = row + (Type == 20 ? group : group / 2) * 18; + scale = __half2float(*reinterpret_cast(b)); + codes = b + 2; + } + Decoded decoded; + decoded.scale = scale; +#pragma unroll + for (int packet = 0; packet < 8; ++packet) { + uint32_t values = 0; + if constexpr (Type == 20 || Type == 23) { + const uint32_t bytes = + load_u32_2(codes + (packet % 4) * 4) >> ((packet / 4) * 4); + const uint32_t nibbles = (bytes & 15) | ((bytes >> 4) & 0xf0) | + ((bytes >> 8) & 0xf00) | + ((bytes >> 12) & 0xf000); + using Lut = turbomind::gemm::Transform_HMMA_SM70_Lut4<0>; + values = Lut::iq_values(nibbles) ^ 0x80808080U; + } else { + const uint32_t byte = b[2 + (group % 2) * 8 + packet]; + const uint32_t codes = (byte & 3) | ((byte << 6) & 0x300) | + ((byte << 12) & 0x30000) | + ((byte << 18) & 0x3000000); + values = __vsub4(codes, 0x01010101U); + } + decoded.values[packet] = static_cast(values); + } + return decoded; + } + __device__ static float dot_decoded(const Decoded& decoded, const Q8_1& x) { + const int* activation = reinterpret_cast(x.qs); + int sum = 0; +#pragma unroll + for (int packet = 0; packet < 8; ++packet) + sum = __dp4a(decoded.values[packet], activation[packet], sum); + return float(sum) * (decoded.scale * __low2float(x.ds)); + } static_assert(Type == 20 || Type == 23 || Type == 42); __device__ static float dot(const uint8_t* row, int group, const Q8_1& x) { const uint8_t* b; diff --git a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu index 5563db43b0..e1e9df682a 100644 --- a/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu +++ b/csrc/sm70_turbomind/ops/gguf_dp4a_sm70.cu @@ -719,3 +719,5 @@ void gguf_dp4a_scalar_lut_gate_up_sm70_out( } #include "gguf_grouped_gate_sm70.cuh" + +#include "gguf_grouped_down_sm70.cuh" diff --git a/csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh new file mode 100644 index 0000000000..43c0425867 --- /dev/null +++ b/csrc/sm70_turbomind/ops/gguf_grouped_down_sm70.cuh @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +namespace { +template +__global__ void raw_down_grouped_routes(half* projected, const Q8_1* input, + const int* ids, const int* order, + const uint8_t* weights, int routes, + int n, int groups, int stride, + int left_groups, int experts) { + __shared__ int active, count, expert, original[4]; + extern __shared__ Q8_1 shared_x[]; + const int route = blockIdx.y; + if (threadIdx.x == 0) { + expert = ids[route]; + int first = route; + while (first > 0 && ids[first - 1] == expert) --first; + active = ((route - first) % 4 == 0) && expert >= 0 && expert < experts; + count = 1; + while (count < 4 && route + count < routes && ids[route + count] == expert) + ++count; + for (int q = 0; q < count; ++q) original[q] = order[route + q]; + } + __syncthreads(); + if (!active) return; + for (int i = threadIdx.x; i < count * groups * 9; i += blockDim.x) { + const int src = original[i / (groups * 9)]; + reinterpret_cast(shared_x)[i] = + reinterpret_cast( + input)[int64_t(src) * groups * 9 + i % (groups * 9)]; + } + __syncthreads(); + const int col = blockIdx.x * 128 + threadIdx.x; + if (col >= n) return; + const uint8_t* row = weights + (int64_t(expert) * n + col) * stride; + float acc[4] = {}; + for (int group = 0; group < groups; ++group) { + const auto decoded = vllm::sm70_gguf::OriginalIntegerDot::decode( + row, group + left_groups); +#pragma unroll + for (int q = 0; q < 4; ++q) + if (q < count) + acc[q] += vllm::sm70_gguf::OriginalIntegerDot::dot_decoded( + decoded, shared_x[q * groups + group]); + } +#pragma unroll + for (int q = 0; q < 4; ++q) + if (q < count) + projected[int64_t(original[q]) * n + col] = __float2half_rn(acc[q]); +} + +__global__ void raw_down_grouped_unroute(half* output, const half* projected, + const float* weights, int n, + int top_k) { + const int token = blockIdx.y; + const int col = blockIdx.x * 128 + threadIdx.x; + if (col >= n) return; + // Retain the four warp-owned subtotals and their original addition order. + float sum = 0.f; +#pragma unroll + for (int owner = 0; owner < 4; ++owner) { + float partial = 0.f; + for (int route = owner; route < top_k; route += 4) { + const int slot = token * top_k + route; + partial += + __half2float(projected[int64_t(slot) * n + col]) * weights[slot]; + } + sum += partial; + } + output[int64_t(token) * n + col] = __float2half_rn(sum); +} + +template +void launch_grouped_raw_down(torch::Tensor out, torch::Tensor projected, + torch::Tensor input, torch::Tensor ids, + torch::Tensor order, torch::Tensor probabilities, + torch::Tensor weights, int left) { + const int n = out.size(1), top_k = ids.size(1), groups = input.size(2); + const auto stream = at::cuda::getCurrentCUDAStream(); + raw_down_grouped_routes<<>>( + reinterpret_cast(projected.data_ptr()), + reinterpret_cast(input.data_ptr()), ids.data_ptr(), + order.data_ptr(), weights.data_ptr(), ids.numel(), n, + groups, weights.size(2), left / 32, weights.size(0)); + raw_down_grouped_unroute<<>>( + reinterpret_cast(out.data_ptr()), + reinterpret_cast(projected.data_ptr()), + probabilities.data_ptr(), n, top_k); +} +} // namespace + +void gguf_dp4a_raw_down_grouped_sm70_out( + torch::Tensor out, torch::Tensor projected, torch::Tensor input, + torch::Tensor ids, torch::Tensor order, torch::Tensor probabilities, + torch::Tensor weights, int64_t source_type, int64_t left) { + const c10::cuda::CUDAGuard guard(input.device()); + require_sm70(); + TORCH_CHECK( + (source_type == 20 || source_type == 42) && left >= 0 && left % 32 == 0, + "Grouped down requires original IQ4_NL or Q2_0 rows"); + TORCH_CHECK(input.is_cuda() && input.scalar_type() == at::kByte && + input.is_contiguous() && input.dim() == 4 && + input.size(0) > 0 && input.size(0) <= 40 && + input.size(1) > 0 && input.size(1) <= 16 && + input.size(2) > 0 && input.size(3) == sizeof(Q8_1), + "Invalid grouped down Q8_1 inputs"); + const int m = input.size(0), top_k = input.size(1), k = input.size(2) * 32; + TORCH_CHECK(ids.device() == input.device() && ids.scalar_type() == at::kInt && + ids.is_contiguous() && ids.dim() == 2 && ids.size(0) == m && + ids.size(1) == top_k && order.device() == input.device() && + order.scalar_type() == at::kInt && order.is_contiguous() && + order.sizes() == ids.sizes() && + probabilities.device() == input.device() && + probabilities.scalar_type() == at::kFloat && + probabilities.is_contiguous() && + probabilities.sizes() == ids.sizes(), + "Invalid grouped down route mapping"); + TORCH_CHECK( + out.device() == input.device() && out.scalar_type() == at::kHalf && + out.is_contiguous() && out.dim() == 2 && out.size(0) == m && + out.size(1) > 0 && weights.device() == input.device() && + weights.scalar_type() == at::kByte && weights.is_contiguous() && + weights.dim() == 3 && weights.size(0) > 0 && + weights.size(1) == out.size(1) && + weights.size(2) >= ((k + left + (source_type == 20 ? 31 : 63)) / + (source_type == 20 ? 32 : 64)) * + 18 && + projected.device() == input.device() && + projected.scalar_type() == at::kHalf && projected.is_contiguous() && + projected.numel() == int64_t(m) * top_k * out.size(1), + "Invalid grouped down banks or output"); + if (source_type == 20) + launch_grouped_raw_down<20>(out, projected, input, ids, order, + probabilities, weights, left); + else + launch_grouped_raw_down<42>(out, projected, input, ids, order, + probabilities, weights, left); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh b/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh index 65f76051a4..18211ad064 100644 --- a/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh +++ b/csrc/sm70_turbomind/ops/gguf_grouped_gate_sm70.cuh @@ -119,18 +119,21 @@ void launch_grouped_gate(torch::Tensor out, torch::Tensor activation, } } // namespace -void gguf_dp4a_gate_up_grouped_sm70_out( - torch::Tensor out, torch::Tensor activation, torch::Tensor ids, - torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, - torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, - torch::Tensor mode) { +template +static void grouped_gate_out_impl(torch::Tensor out, torch::Tensor activation, + torch::Tensor ids, torch::Tensor gate, + torch::Tensor up, int64_t source_type, + bool bank_aware, torch::Tensor route_tokens, + torch::Tensor order, + torch::Tensor original_ids, + torch::Tensor mode) { TORCH_CHECK(source_type == 18 || source_type == 21 || source_type == 22, "Paired gate/up requires an original lattice format"); const c10::cuda::CUDAGuard guard(activation.device()); require_sm70(); TORCH_CHECK(activation.is_cuda() && activation.scalar_type() == at::kByte && activation.is_contiguous() && activation.dim() == 3 && - activation.size(0) > 0 && activation.size(0) <= 20 && + activation.size(0) > 0 && activation.size(0) <= MaxRows && activation.size(2) == sizeof(Q8_1) && activation.size(1) % 8 == 0, "Invalid paired Q8_1 activation"); @@ -188,3 +191,23 @@ void gguf_dp4a_gate_up_grouped_sm70_out( #undef PAIR C10_CUDA_KERNEL_LAUNCH_CHECK(); } + +void gguf_dp4a_gate_up_grouped_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + grouped_gate_out_impl<20>(out, activation, ids, gate, up, source_type, + bank_aware, route_tokens, order, original_ids, + mode); +} + +void gguf_dp4a_gate_up_grouped_batch_sm70_out( + torch::Tensor out, torch::Tensor activation, torch::Tensor ids, + torch::Tensor gate, torch::Tensor up, int64_t source_type, bool bank_aware, + torch::Tensor route_tokens, torch::Tensor order, torch::Tensor original_ids, + torch::Tensor mode) { + grouped_gate_out_impl<40>(out, activation, ids, gate, up, source_type, + bank_aware, route_tokens, order, original_ids, + mode); +} diff --git a/csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh b/csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh new file mode 100644 index 0000000000..bc72bc6f38 --- /dev/null +++ b/csrc/sm70_turbomind/ops/hcx_batch_sm70.cuh @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Included in namespace hcx, after the retained HCX primitives. Full-mesh +// TP4, separate output projection, local schedule, M9..40 only. One resident +// grid performs the exchange and both grid rendezvous for the whole batch. +constexpr int BATCH_ROWS = 40; + +__global__ __launch_bounds__(256, 1) void hcx_batch_kernel(Args a) { + const int i = blockIdx.x, t = threadIdx.x, warp = t >> 5, lane = t & 31; + const int rk = a.rank, p1 = rk ^ 1, p2 = rk ^ 2, p3 = rk ^ 3; + const unsigned tag = ldv(a.seq) + 1, base = ldv(a.bar + 3); + const int r8 = (lane & 3) + ((lane & 16) ? 4 : 0), quad = (lane >> 2) & 3; + constexpr int AR_STRIDE = BATCH_ROWS * 32; + __shared__ __align__(16) half xs[BATCH_ROWS][136]; + __shared__ __align__(16) half ls[BATCH_ROWS][328]; + __shared__ float red[5][8][33]; + __shared__ float rr[BATCH_ROWS][4]; + + for (int row = warp; row < a.M; row += 8) { + const int col = 32 * i + lane; + float v = __half2float(a.p0[row * HD + col]); + if (a.p1) v += __half2float(a.p1[row * HD + col]); + const int slot = (rk * NC + i) * AR_STRIDE + row * 32 + lane; + st2(a.ar[p1] + slot, __float_as_uint(v), tag); + st2(a.ar[p2] + slot, __float_as_uint(v), tag); + st2(a.ar[p3] + slot, __float_as_uint(v), tag); + } + uint4 dlo[4], dhi[4], ulo[4], uhi[4]; + if (warp < 6) { +#pragma unroll + for (int s = 0; s < 4; ++s) { + const uint4* p = a.wd + ((i * 6 + warp) * 4 + s) * 64; + dlo[s] = ldgv(p + lane); + dhi[s] = ldgv(p + 32 + lane); + } + } + if (warp >= 3) { +#pragma unroll + for (int s = 0; s < 4; ++s) { + const uint4* p = a.wu + ((i * 5 + warp - 3) * 4 + s) * 64; + ulo[s] = ldgv(p + lane); + uhi[s] = ldgv(p + 32 + lane); + } + } + for (int row = warp; row < BATCH_ROWS; row += 8) { + const int col = 32 * i + lane; + float sums[4] = {}; +#pragma unroll + for (int b = 0; b < 4; ++b) xs[row][b * 32 + lane] = __float2half_rn(0.f); + if (row < a.M) { + float v = __half2float(a.p0[row * HD + col]); + if (a.p1) v += __half2float(a.p1[row * HD + col]); + const int slot = i * AR_STRIDE + row * 32 + lane; + const float x1 = + __uint_as_float(poll2(a.ar[rk] + p1 * NC * AR_STRIDE + slot, tag)); + const float x2 = + __uint_as_float(poll2(a.ar[rk] + p2 * NC * AR_STRIDE + slot, tag)); + const float x3 = + __uint_as_float(poll2(a.ar[rk] + p3 * NC * AR_STRIDE + slot, tag)); + const float near = (rk & 1) ? x1 + v : v + x1; + const float far = (rk & 1) ? x3 + x2 : x2 + x3; + const float reduced = + __half2float(__float2half_rn((rk & 2) ? far + near : near + far)); +#pragma unroll + for (int b = 0; b < 4; ++b) { + const float g = 2.f * sigm(__half2float(a.inj[row * 4 + b]) / 4.f); + const half h = __float2half_rn( + __half2float(a.res[row * KD + b * HD + col]) + reduced * g); + const float o = __half2float(h); + const half nw = a.nw[a.nw_full ? b * HD + col : col]; + a.res_out[row * KD + b * HD + col] = h; + xs[row][b * 32 + lane] = __float2half_rn(o + o * __half2float(nw)); + sums[b] = o * o; + } + } +#pragma unroll + for (int s = 16; s > 0; s >>= 1) +#pragma unroll + for (int b = 0; b < 4; ++b) + sums[b] += __shfl_xor_sync(0xffffffff, sums[b], s); + if (lane == 0 && row < a.M) + *reinterpret_cast(a.sq + (i * BATCH_ROWS + row) * 4) = + make_float4(sums[0], sums[1], sums[2], sums[3]); + } + __syncthreads(); + if (warp < 6) { + const int tile = warp >> 1, kh = warp & 1; +#pragma unroll + for (int token_tile = 0; token_tile < 5; ++token_tile) { + if (token_tile * 8 >= a.M) break; + float acc0[8] = {}, acc1[8] = {}; +#pragma unroll + for (int s = 0; s < 4; ++s) { + const int k = (kh * 4 + s) * 16; + const uint4 x0 = + *reinterpret_cast(&xs[token_tile * 8 + r8][k]); + const uint4 x1 = + *reinterpret_cast(&xs[token_tile * 8 + r8][k + 8]); + float (&acc)[8] = s < 2 ? acc0 : acc1; + mma(acc, x0.x, x0.y, dlo[s].x, dlo[s].y); + mma(acc, x0.z, x0.w, dlo[s].z, dlo[s].w); + mma(acc, x1.x, x1.y, dhi[s].x, dhi[s].y); + mma(acc, x1.z, x1.w, dhi[s].z, dhi[s].w); + } +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int rw = + token_tile * 8 + (e & 2) + ((lane & 16) ? 4 : 0) + (lane & 1); + const int n = tile * 32 + quad * 8 + + ((e & 1) | (((lane >> 1) & 1) << 1) | ((e >> 2) << 2)); + if (rw < a.M) + *reinterpret_cast( + a.dpart + ((i * BATCH_ROWS + rw) * 96 + n) * 4 + 2 * kh) = + make_float2(acc0[e], acc1[e]); + } + } + } + grid_bar_epoch(a.bar + 2, base + NC); + for (int row = warp; row < a.M; row += 8) { + float4 qs[3]; +#pragma unroll + for (int j = 0; j < 3; ++j) + qs[j] = lane + 32 * j < NC + ? __ldcg(reinterpret_cast(a.sq) + + (lane + 32 * j) * BATCH_ROWS + row) + : make_float4(0.f, 0.f, 0.f, 0.f); + float sums[4] = { + (qs[0].x + qs[1].x) + qs[2].x, (qs[0].y + qs[1].y) + qs[2].y, + (qs[0].z + qs[1].z) + qs[2].z, (qs[0].w + qs[1].w) + qs[2].w}; +#pragma unroll + for (int s = 16; s > 0; s >>= 1) +#pragma unroll + for (int b = 0; b < 4; ++b) + sums[b] += __shfl_xor_sync(0xffffffff, sums[b], s); +#pragma unroll + for (int b = 0; b < 4; ++b) { + const float inverse = rsqrtf(sums[b] / HD + a.eps); + if (lane == 0) rr[row][b] = inverse; + const float o = + __half2float(a.res_out[row * KD + b * HD + 32 * i + lane]); + float y = o * inverse; + y += y * __half2float( + a.nw[a.nw_full ? b * HD + 32 * i + lane : 32 * i + lane]); + a.xn[row * KD + b * HD + 32 * i + lane] = __float2half_rn(y); + } + } + __syncthreads(); + for (int oi = i + NC * warp; oi < a.M * 96; oi += NC * 8) { + const int rw = oi / 96, n = oi % 96; + if (!(n < 80 || (rk == 3 && n < 84))) continue; + float pb[4] = {}; +#pragma unroll + for (int j = 0; j < 3; ++j) { + if (lane + 32 * j < NC) { + const float4 pv = __ldcg(reinterpret_cast(a.dpart) + + ((lane + 32 * j) * BATCH_ROWS + rw) * 96 + n); + pb[0] += pv.x; + pb[1] += pv.y; + pb[2] += pv.z; + pb[3] += pv.w; + } + } +#pragma unroll + for (int s = 16; s > 0; s >>= 1) +#pragma unroll + for (int b = 0; b < 4; ++b) + pb[b] += __shfl_xor_sync(0xffffffff, pb[b], s); + float q = 0.f; +#pragma unroll + for (int b = 0; b < 4; ++b) q = fmaf(rr[rw][b], pb[b], q); + half hv = __float2half_rn(q); + if (n < 80) { + const float v = divf(__half2float(hv), 4.f); + hv = __float2half_rn(v * sigm(v)); + } + const int pos = rw * LW + (n < 80 ? rk * 80 + n : 320 + n - 80); + const unsigned value = __half_as_ushort(hv); + if (lane == 0) st2(a.lora[rk] + pos, value, tag); + if (lane == 1) st2(a.lora[p1] + pos, value, tag); + if (lane == 2) st2(a.lora[p2] + pos, value, tag); + if (lane == 3) st2(a.lora[p3] + pos, value, tag); + } + for (int idx = i * 256 + t; idx < a.M * 324; idx += NC * 256) + poll2(a.lora[rk] + (idx / 324) * LW + idx % 324, tag); + grid_bar_epoch(a.bar + 2, base + 2 * NC); + for (int idx = t; idx < BATCH_ROWS * 320; idx += 256) + ls[idx / 320][idx % 320] = + idx < a.M * 320 + ? __ushort_as_half(static_cast( + __ldcg(a.lora[rk] + (idx / 320) * LW + idx % 320).x)) + : __float2half_rn(0.f); + if (i == 0 && t < a.M * 4) + a.inj_out[t] = __ushort_as_half(static_cast( + __ldcg(a.lora[rk] + (t >> 2) * LW + 320 + (t & 3)).x)); + __syncthreads(); +#pragma unroll + for (int token_tile = 0; token_tile < 5; ++token_tile) { + if (token_tile * 8 >= a.M) break; + if (warp >= 3) { + const int uw = warp - 3; + float acc[8] = {}; +#pragma unroll + for (int s = 0; s < 4; ++s) { + const int k = (uw * 4 + s) * 16; + const uint4 x0 = + *reinterpret_cast(&ls[token_tile * 8 + r8][k]); + const uint4 x1 = + *reinterpret_cast(&ls[token_tile * 8 + r8][k + 8]); + mma(acc, x0.x, x0.y, ulo[s].x, ulo[s].y); + mma(acc, x0.z, x0.w, ulo[s].z, ulo[s].w); + mma(acc, x1.x, x1.y, uhi[s].x, uhi[s].y); + mma(acc, x1.z, x1.w, uhi[s].z, uhi[s].w); + } +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int rw = (e & 2) | ((lane & 16) ? 4 : 0) | (lane & 1); + red[uw][rw][quad * 8 + ((e & 1) | (lane & 2) | (e & 4))] = acc[e]; + } + } + __syncthreads(); + const int row = token_tile * 8 + warp; + if (row < a.M) { + const int h = 640 * rk + 8 * i + (lane & 7); + float acc = red[0][warp][lane]; +#pragma unroll + for (int q = 1; q < 5; ++q) acc += red[q][warp][lane]; + const float sg = sigm(__half2float(__float2half_rn(acc))); + const float xv = + __half2float(__ldcg(a.xn + row * KD + (lane >> 3) * HD + h)); + float mixed = 0.f; +#pragma unroll + for (int b = 0; b < 4; ++b) + mixed = fmaf(__shfl_sync(0xffffffff, sg, b * 8 + (lane & 7)), + __shfl_sync(0xffffffff, xv, b * 8 + (lane & 7)), mixed); + const unsigned bits = __half_as_ushort(__float2half_rn(divf(mixed, 4.f))); + const unsigned packed = + bits | (__shfl_down_sync(0xffffffff, bits, 1) << 16); + const unsigned upper = __shfl_down_sync(0xffffffff, packed, 2); + if (lane < 8 && !(lane & 3)) { + st4(a.hb[p1] + (row * HD + h) / 2, packed, upper, tag); + st4(a.hb[p2] + (row * HD + h) / 2, packed, upper, tag); + st4(a.hb[p3] + (row * HD + h) / 2, packed, upper, tag); + *reinterpret_cast(a.blk_out + row * HD + h) = + make_uint2(packed, upper); + } + } + __syncthreads(); + } + for (int u = t; u < 6 * a.M; u += 256) { + const int src = u % 3 == 0 ? p1 : u % 3 == 1 ? p2 : p3; + const int row = (u / 3) >> 1, h = 640 * src + 8 * i + ((u / 3) & 1) * 4; + *reinterpret_cast(a.blk_out + row * HD + h) = + poll4(a.hb[rk] + (row * HD + h) / 2, tag); + } + if (i == 0 && t == 0) { + a.bar[3] = base + 2 * NC; + a.seq[0] = tag; + } +} diff --git a/csrc/sm70_turbomind/ops/hcx_sm70.cu b/csrc/sm70_turbomind/ops/hcx_sm70.cu index 8bb2461974..22d1af4fd5 100644 --- a/csrc/sm70_turbomind/ops/hcx_sm70.cu +++ b/csrc/sm70_turbomind/ops/hcx_sm70.cu @@ -736,9 +736,11 @@ __global__ __launch_bounds__(256, 1) void hcx_kernel(Args a) { } } } +#include "hcx_batch_sm70.cuh" } // namespace hcx -void sm70_hcx_out( +template +void sm70_hcx_impl( torch::Tensor p0, std::optional p1, torch::Tensor res, torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, @@ -779,6 +781,23 @@ void sm70_hcx_out( a.dbg = dbg ? reinterpret_cast(dbg->data_ptr()) : nullptr; a.M = static_cast(p0.size(0)); + if constexpr (Batch) { + TORCH_CHECK(a.M >= 9 && a.M <= 40 && full && local_schedule && !ox && !gz, + "HCX batch requires M9..40, full mesh, local schedule and " + "separate projection"); + const auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0 && + props->multiProcessorCount >= hcx::NC, + "HCX batch requires SM70 with at least 80 SMs"); + TORCH_CHECK(xn.numel() >= 40 * hcx::KD && sq.numel() >= hcx::NC * 40 * 4 && + dpart.numel() >= hcx::NC * 40 * 96 * 4 && bar.numel() >= 4, + "HCX batch workspace too small"); + hcx:: + hcx_batch_kernel<<>>( + a); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return; + } TORCH_CHECK(a.M >= 1 && a.M <= 8, "HCX: M1..8"); auto st = at::cuda::getCurrentCUDAStream(); if (ox) { @@ -878,3 +897,41 @@ void sm70_dmv13_out(torch::Tensor x, std::vector codes, dmvns::run(x, codes, high, scale, out, fmt, n, K, split, warps, ws, cnt, tp, std::nullopt, std::nullopt, extra_weight, extra_out); } + +void sm70_hcx_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule) { + sm70_hcx_impl(p0, p1, res, inj, nw, eps, wd, wu, res_out, blk_out, + inj_out, xn, sq, dpart, bar, seq, ar, lora, hb, rank, + dbg, full, ox, ocodes, ohigh, oscale, ofmt, gz, gw, geps, + gscr, local_schedule); +} + +void sm70_hcx_batch_out( + torch::Tensor p0, std::optional p1, torch::Tensor res, + torch::Tensor inj, torch::Tensor nw, double eps, torch::Tensor wd, + torch::Tensor wu, torch::Tensor res_out, torch::Tensor blk_out, + torch::Tensor inj_out, torch::Tensor xn, torch::Tensor sq, + torch::Tensor dpart, torch::Tensor bar, torch::Tensor seq, + std::vector ar, std::vector lora, std::vector hb, + int64_t rank, std::optional dbg, int64_t full, + std::optional ox, std::optional ocodes, + std::optional ohigh, std::optional oscale, + int64_t ofmt, std::optional gz, + std::optional gw, double geps, + std::optional gscr, bool local_schedule) { + sm70_hcx_impl(p0, p1, res, inj, nw, eps, wd, wu, res_out, blk_out, + inj_out, xn, sq, dpart, bar, seq, ar, lora, hb, rank, dbg, + full, ox, ocodes, ohigh, oscale, ofmt, gz, gw, geps, gscr, + local_schedule); +} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 4c34a26516..e88f8aee49 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -102,6 +102,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ", Tensor? oscale, int ofmt, Tensor? gz, Tensor? gw, float geps, Tensor" "(i!)? gscr, bool local_schedule=False) -> ()"); ops.impl("sm70_hcx_out", torch::kCUDA, &sm70_hcx_out); + ops.def( + "sm70_hcx_batch_out(Tensor p0, Tensor? p1, Tensor res, Tensor inj, " + "Tensor nw," + " float eps, Tensor wd, Tensor wu, Tensor(a!) res_out, Tensor(b!) blk_o" + "ut, Tensor(c!) inj_out, Tensor(d!) xn, Tensor(e!) sq, Tensor(f!) dpart" + ", Tensor(g!) bar, Tensor(h!) seq, int[] ar, int[] lora, int[] hb, int " + "rank, Tensor? dbg, int full, Tensor? ox, Tensor? ocodes, Tensor? ohigh" + ", Tensor? oscale, int ofmt, Tensor? gz, Tensor? gw, float geps, Tensor" + "(i!)? gscr, bool local_schedule=False) -> ()"); + ops.impl("sm70_hcx_batch_out", torch::kCUDA, &sm70_hcx_batch_out); ops.def( "sm70_dmv13_out(Tensor x, Tensor[] codes, Tensor[] high, Tensor[] scale, " "Tensor(a!)[] out, int[] fmt, int[] n, int K, int split, int warps, " @@ -447,6 +457,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "()"); ops.impl("gguf_dp4a_gate_up_grouped_sm70_out", torch::kCUDA, &gguf_dp4a_gate_up_grouped_sm70_out); + ops.def( + "gguf_dp4a_gate_up_grouped_batch_sm70_out(Tensor(a!) out, Tensor " + "activation, " + "Tensor ids, Tensor gate, Tensor up, int source_type, bool bank_aware, " + "Tensor route_tokens, Tensor order, Tensor original_ids, Tensor mode) -> " + "()"); + ops.impl("gguf_dp4a_gate_up_grouped_batch_sm70_out", torch::kCUDA, + &gguf_dp4a_gate_up_grouped_batch_sm70_out); + ops.def( + "gguf_dp4a_raw_down_grouped_sm70_out(Tensor(a!) out, " + "Tensor(b!) projected, Tensor input, Tensor ids, Tensor order, " + "Tensor probabilities, Tensor weights, int source_type, int left) -> ()"); + ops.impl("gguf_dp4a_raw_down_grouped_sm70_out", torch::kCUDA, + &gguf_dp4a_raw_down_grouped_sm70_out); + ops.def( "gguf_dp4a_lut4_gate_up_sm70_out(Tensor(a!) out, Tensor activation, " "Tensor ids, Tensor gate, Tensor gate_stats, Tensor up, Tensor up_stats, " diff --git a/docs/design/architecture/runtime_reference.md b/docs/design/architecture/runtime_reference.md index deaec251b8..42e427886f 100644 --- a/docs/design/architecture/runtime_reference.md +++ b/docs/design/architecture/runtime_reference.md @@ -38,12 +38,12 @@ 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#L1089) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L629) | +| `vllm.config.kernel.KernelConfig` | kernel_config | compute_hash: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L1119) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L653) | | `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) | | `vllm.config.kernel.Sm70NvFp4Config` | kernel_config.sm70_nvfp4 | resolve: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L277); capture_inputs: [vllm/config/kernel.py](../../../vllm/config/kernel.py#L216) | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L236) | -| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L618) | +| `vllm.config.kernel.Sm70RingConfig` | kernel_config.sm70_ring | | [vllm/config/kernel.py](../../../vllm/config/kernel.py#L642) | | `vllm.config.kv_events.KVEventsConfig` | kv_events_config | __post_init__: [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L50) | [vllm/config/kv_events.py](../../../vllm/config/kv_events.py#L11) | | `vllm.config.kv_transfer.KVTransferConfig` | kv_transfer_config | compute_hash: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L75); __post_init__: [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L93) | [vllm/config/kv_transfer.py](../../../vllm/config/kv_transfer.py#L23) | | `vllm.config.legacy_inputs.LegacyInputs` | kernel_config.sm70_awq.legacy, kernel_config.sm70_fp8.legacy, kernel_config.sm70_moe.awq.legacy, kernel_config.sm70_moe.fp8.legacy, kernel_config.sm70_moe.mxfp4.legacy, kernel_config.sm70_moe.nvfp4.legacy, kernel_config.sm70_nvfp4.legacy, scheduler_config.sm70_inputs | value: [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L39) | [vllm/config/legacy_inputs.py](../../../vllm/config/legacy_inputs.py#L15) | diff --git a/docs/design/flashnext_mtp_batch_16gb_20261015.md b/docs/design/flashnext_mtp_batch_16gb_20261015.md index 78b5e3dc8d..4f45e6920a 100644 --- a/docs/design/flashnext_mtp_batch_16gb_20261015.md +++ b/docs/design/flashnext_mtp_batch_16gb_20261015.md @@ -318,9 +318,9 @@ It reuses the existing segment format and native decoder across token tiles; qualified DMV Q4/IQ4 banks provide their equivalent full scale representation. The six-array prepared-operand ABI is retained. The accelerated initialization subset is K2560/N>=1536 or K1536/N2560 in the existing qualified formats and -layouts. Runtime acceleration is limited to M9..16; M1..8 and M17..40 -retain the previous projection routes, and larger batches retain canonical -restoration. The native batch operator has separately tested M1..40 coverage. The old native entry retains +layouts. Runtime acceleration covers M9..16 and M25..40; M1..8 and +M17..24 retain the previous projection routes. Compatible small segment +peers share the same batch launch. Larger batches retain canonical restoration. The native batch operator has separately tested M1..40 coverage. The old native entry retains its M1..32 contract; the batch entry is explicit. Eighty-six dense native GPU cases and 111 affected storage/runtime cases pass, including capture/replay and the M40/M41 fallback boundary. The initial wider M9..40 model experiment @@ -397,3 +397,242 @@ atomically. The final artifacts pass 54 affected GPU cases, including changed-ro graph replay, independent expert output oracles and four-rank indexer state checks. Clean model speed and longer-context gates remain separate before any promotion. + +## Matched node attribution and mixed-segment fusion + +Five decode-only node captures use the final rebuilt source and identical +C1/C2/C4/C6/C8 contracts. Their interior counts are150/113/87/60/49. +Node tracing changes absolute latency: traced whole verifier means are +22.105/29.934/39.397/57.967/71.533ms, separately from event measurements. + +| TP0 GPU service ms | C1 | C2 | C4 | C6 | C8 | +| --- | ---: | ---: | ---: | ---: | ---: | +| Dense/shared projections | 3.489 | 6.879 | 6.224 | 9.291 | 14.847 | +| Routed expert projections | 2.880 | 4.639 | 7.573 | 12.015 | 15.439 | +| HC and its fused exchange | 2.504 | 4.648 | 7.564 | 10.575 | 13.877 | +| Indexer and top-k | 3.556 | 3.876 | 4.434 | 5.109 | 5.955 | +| Expert-cache resolution/wait | 1.659 | 1.398 | 1.498 | 3.285 | 3.761 | + +These are service sums, with overlapping streams; they are not an additive +wall decomposition. C8 peer broadcast includes waiting for TP0 indexer work. +The actual local GPU4–7 mesh has mixed NV1/NV2 links; the remote9683 NV2 mesh +is a separate hardware contract. CPU affinity matches NUMA node1. + +The captured operand descriptors expose a dispatch problem: QKV banks mix a +large batch-enabled segment (family5) with existing small K/V segments +(family3). Requiring every segment to have a batch tag splits compatible +projections and restores the small peers before separate calls. The candidate +admits one native batch launch when all segments have the compatible layout +and at least one calibrated large segment enables batching. The old M20 +fusion remains unchanged. M9..16 and M25..40 are qualified separately from +larger prefill fallback. + +Real-weight QKV screens across four layers improve M10 by2.07–3.63x, +M30 by1.48–2.04x and M40 by1.98–2.81x. M20 arithmetic is unchanged. +Relative output errors remain below0.00077 versus the FP32 dequantized-weight +reference, under the predeclared0.003 norm bound. The affected111 GPU runtime +cases pass; three additional mixed-format cases exercise11 boundary/interior +M values, changed-input replay and exactly one native batch dispatch at the +accelerated sizes. The matching complete model curve is being collected; +these micro gains are not reported as whole-verifier gains. + +Two original-down scheduling screens are rejected: cooperative two/four-lane +source reads usually regress despite exact outputs, and one-warp-per-route +regresses IQ4_NL while offering no representative M20 gain for Q2_0. Both +research patches and all15 real-weight cases are retained outside the runtime. + +The mixed-segment model experiment completes all five 32K shared-prefix +cases with181/214/182/178/165 complete TP4 verifier rounds and no incomplete +rounds. Whole means are17.513/26.114/36.692/50.150/61.131ms; C8 p99 is68.697ms +and draft mean8.768ms. C8 improves5.8% relative to the preceding combined +64.895ms diagnostic. Retrieval, executable code and automatic tool call/reply +all pass. This remains above the requested18/22/28ms targets and uses worker +fences; ordinary-worker qualification remains separate. Source/native sampler +acceptance varies between runs, so HTTP gains alone are not attributed to the +projection change. + +An isolated HC grid-barrier screen compares the existing cumulative arrival +counter with distributed arrival flags. Both take approximately1.62us per +barrier and pass ordering/wrap checks in that screen. No native protocol change +is justified; the private prototype is retained as a negative result. + +## Closed verifier interval audit + +The five captures use the pre-mixed-fusion grouped-twenty-row candidate +(the untraced C8 event mean was 64.895 ms). They were swept by GPU timestamp on each +selected critical rank. Each interval belongs to exactly one category, a +separate overlap bucket, or the idle residual. Categories below sum to verifier +wall; they are not summed kernel service. Node tracing perturbs absolute latency +and these numbers must not replace the untraced event measurements above. + +| Concurrency | Complete interior rounds | Verifier wall ms | HC ms | Experts ms | Dense/shared ms | TP collectives ms | Overlap ms | Idle ms | +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| 1 | 150 | 22.105 | 2.390 | 2.879 | 2.489 | 3.661 | 0.999 | 2.150 | +| 2 | 113 | 29.934 | 4.592 | 2.411 | 3.577 | 3.988 | 5.689 | 2.590 | +| 4 | 87 | 39.397 | 7.399 | 7.560 | 5.488 | 4.823 | 0.167 | 3.035 | +| 6 | 60 | 57.967 | 10.565 | 11.416 | 5.086 | 5.821 | 4.334 | 4.364 | +| 8 | 49 | 71.533 | 13.702 | 12.595 | 7.753 | 6.724 | 9.138 | 4.338 | + +The remaining wall belongs to selected-history attention, GDN, routing, cache +resolution, normalization, sampling/state commit and uncategorized kernels. +Raw `scaling-nodes-{1..5}-closed.json` files retain those categories and each +round; `scaling-nodes-closed-summary.json` retains their means. The closure +assertion tolerates at most 1e-7 ms. This audit includes target forward, +sampling and state commit, and excludes MTP draft calls. Expert and HC scaling +remain primary optimization targets. A collective can include a wait on the +indexer root, so its duration does not establish intrinsic link cost. + +## Forty-row expert/cache experiment + +The default fragment and mailbox capacity remain twenty rows. An explicit +`sm70_gguf.grouped_gate_up_max_m=40` permits one 400-route lease with a +512-integer mailbox. Reply capacity accounts for all request entries plus +spare restoration; it grows to the next power of two when required. Retained +quantization and ordinary gate/up/down kernels still run in twenty-row tiles. +Grouped original gate/up has a separate native forty-row entry point, keeping +the legacy operator and opaque schemas unchanged. + +The matched diagnostic event curve with forty-row gate/up, without grouped +down, is 17.752 / 26.270 / 37.187 / 47.641 / 61.115 ms for C1/2/4/6/8, +with 202 / 200 / 190 / 196 / 174 complete rounds and no incomplete rounds. +C6 improves versus the mixed-fusion twenty-row control; C8 does not. This is +not evidence that a faster isolated gate/up microbenchmark speeds the service. +Eight concurrent 32K exact retrieval/arithmetic responses, executable Python, +automatic weather tool selection/reply and natural EOS checks all pass. + +The next default-off `sm70_gguf.grouped_down` candidate groups original +IQ4_NL/Q2_0 down weights across at most four equal-expert routes. It keeps +FP16 projection boundaries, the retained four warp-owned route subtotals and +their addition order. Only M25..40 is accelerated, with the forty-row lease +policy and admitted Q8 intermediate/native expert shape. Smaller fragments +keep the retained down path. Unsupported shapes use that path or its existing +TurboMind continuation; a requested qualified route without its native build +rejects initialization. + +Twenty-seven real-weight down screens and twenty-one threshold/interior +screens preserve the prior bytes, including the saved prior-binary references. +At M40, sort-inclusive speedups are 1.12..1.92x for IQ4_NL and 1.17..1.57x +for Q2_0. Sparse M20 IQ4_NL regresses and remains outside acceleration. Sparse +M25 IQ4_NL only improves 1.013x; this boundary result is not a broad speed claim. +The native, cache, state/graph and complete expert-chain checks pass 141 GPU +cases after correcting a test parameter that shadowed its weight tensor. +The combined service screen below establishes measured 32K performance and +task evidence; 256K quality qualification remains outstanding. + +Additional rejected screens are retained in the report directory. A bounded +mailbox-plan scan changed isolated cache time by less than 1% and was removed. +A four-float HC partial-workspace pad appeared faster in a one-call graph, +but the 64-call graph screen removed launch skew and showed no C8 gain +(0.13680 ms retained versus 0.13803 ms padded). Its source was removed. + +## Rechecked whole-batch HC path + +The earlier whole-batch HC rejection did not establish a reliable native-path +comparison: its launcher did not explicitly enable the fusion policy, and a +single-call graph mixed GPU work with peer launch skew. Rechecking the native +operator directly, with observed CUDA dispatch and sixty-four calls per graph, +shows bit-identical outputs and critical-rank per-call times: + +| Concurrency | Retained tiled HC ms | Whole-batch HC ms | Speedup | +| --- | --- | --- | --- | +| 1 | 0.021842 | 0.021777 | 1.003x (retained fallback) | +| 2 | 0.045594 | 0.043849 | 1.040x | +| 4 | 0.073743 | 0.065897 | 1.119x | +| 6 | 0.103555 | 0.092633 | 1.118x | +| 8 | 0.136066 | 0.121190 | 1.123x | + +This supersedes the earlier HC rejection for the measured full-mesh TP4, +separate-projection, local-schedule path. The default-off typed policy +`sm70_hcx_batch_fusion` owns its larger peer buffers and local workspaces. +Producer-fused output projections, non-full-mesh execution, unsupported local +modes and M1..16 retain the original eight-row path. M9..16 was excluded after +the ordinary-worker acceptance regression described below. M>40 retains the existing +projection/collective continuation. Requested qualified fusion rejects a +missing native build. Original operator signatures remain unchanged. + +The four-rank Inductor/graph payload test passes all M5/8/9/10/17/20/24/25/30/ +32/33/39/40/41 cases, including input mutation, retained/fast alternation, +materialized/final HC consumers and output ownership checks. These results are +kernel/state evidence. The combined service performance and task gates are +being measured separately; 256K boundary qualification remains outstanding. + +## Combined expert-down and HC service screen + +With mixed dense fusion, forty-row expert leases, grouped down and whole-batch +HC enabled, the matched MTP4 diagnostic GPU-event verifier costs are: + +| Concurrency | Complete rounds | Incomplete | Mean ms | P99 ms | Draft mean ms | +| --- | --- | --- | --- | --- | --- | +| 1 | 202 | 0 | 17.650 | 20.388 | 3.875 | +| 2 | 220 | 0 | 25.681 | 28.084 | 4.580 | +| 4 | 214 | 0 | 35.286 | 42.640 | 5.936 | +| 6 | 186 | 0 | 45.974 | 58.584 | 7.455 | +| 8 | 193 | 0 | 54.660 | 64.421 | 8.749 | + +C8 improves from 61.131 to 54.660 ms (10.6%) versus the prior mixed-fusion +screen. C4 improves from 36.692 to 35.286 ms and C6 from 50.150 to 45.974 ms. +These are whole-batch costs, never costs divided by concurrency. The workload +keeps 256K capacity, 32K shared input, MTP4, C8 capacity, the same host QSA and +80% GPU expert storage, 512-token prefill chunks and natural-EOS sampling. +Profiling events/fences are enabled. The 16/18/22/28 ms C1/2/4/8 targets remain +unmet. Ordinary-worker throughput/acceptance and paired task checks are a +separate gate; no default promotion follows from this screen. + +## Ordinary-worker acceptance audit + +Fresh ordinary workers, without event fences, NVTX tracing or worker wrappers, +completed matched mixed-fusion control and combined expert-down/HC passes. +Both use the same source/native artifact, 384 output tokens, 32774 input +tokens, capacity eight, sampling and per-request seeds. The control retains +twenty-row gate/up and tiled HC; the combined pass enables forty-row gate/up, +grouped down and whole-batch HC. Each row is the measured common interval +with all submitted requests still decoding; rates are aggregate emitted tokens. + +| Concurrency | Control tok/s | Combined tok/s | Control acceptance length | Combined acceptance length | +| --- | ---: | ---: | ---: | ---: | +| 1 | 110.15 | 125.34 | 2.477 | 2.824 | +| 2 | 169.49 | 133.84 | 2.630 | 2.121 | +| 4 | 213.31 | 235.15 | 2.339 | 2.473 | +| 6 | 251.40 | 277.44 | 2.458 | 2.481 | +| 8 | 259.09 | 305.02 | 2.259 | 2.451 | + +The C2 emitted-token regression is approximately 21%, despite a slightly +faster diagnostic verifier. This candidate is not accepted as a general MTP +speed improvement. C1 acceptance also changes, so its rate improvement cannot +be attributed to verifier compute alone. Both passes complete the eight-way +32K retrieval/arithmetic, executable Python, automatic tool-call/reply and +natural-EOS gates. These bounded task checks do not establish unchanged MTP +acceptance or broad quality. The runtime HC acceleration range is consequently +narrowed to M17..40 while C2 acceptance is investigated. The native operator's +correctness range remains M9..40; admission and measured benefit are distinct. + +The subsequent shared-expert experiment reuses the existing fused gate/up and +down operations in bounded twenty-row fragments. It is explicitly enabled by +`sm70_gguf.shared_expert_batch`, defaults off, admits Q8_0 down projections at +M9..40, and preserves the canonical continuation for other projections and +larger batches. The existing two-integer internal down descriptor remains +valid; the explicit opt-in adds a third policy integer without changing the +registered operator schema. Real-weight screens across six layers show +1.34..2.18x isolated M40 speedups and paired relative-norm error at most +0.000622 against the predeclared 0.003 bound. M5/M8/M20 retain their prior +bytes. GPU graph tests cover M5/8/9/10/16/17/20/21/24/25/30/32/33/39/40/41, +mixed source formats, the independent FP32 oracle, changed inputs and native +execution/fallback. Eleven format combinations pass. + +The first ordinary-worker shared-expert pass emits 110.24 / 151.25 / 244.01 / +253.33 / 292.45 tok/s at C1/2/4/6/8. C2's additional same-seed repetition emits +159.54 tok/s, with acceptance length changing from 2.444 to 2.579. Eight-way +retrieval/arithmetic, coding and tool/reply gates pass. C6/C8 do not exceed the +previous combined pass, and C2 remains below control. This is an unpromoted +experiment; an isolated shared-expert speedup is insufficient to accept it. +Its separate synchronized-verifier curve is still being collected. + +The latest combined C8 node capture closes a 60.184 ms traced verifier wall +with 12.465 ms HC/fused communication, 9.069 ms expert projections, 6.599 ms +dense/shared projections, 6.453 ms TP collectives, 3.874 ms selected-history +attention, 2.423 ms GDN, 9.191 ms overlap, 3.399 ms idle and 6.711 ms in the +remaining categories. Categories are exclusive wall intervals; traced absolute +time is not the 54.660 ms diagnostic result. The observed native expert route +groups average approximately four rows per expert. Neither these counts nor +collective durations establish memory bandwidth or intrinsic NVLink cost; +Nsight Compute profiling counters remain unavailable to this user. diff --git a/tests/config/test_flashnext_mtp_batch_policy.py b/tests/config/test_flashnext_mtp_batch_policy.py index d6a6613726..aa90504ca6 100644 --- a/tests/config/test_flashnext_mtp_batch_policy.py +++ b/tests/config/test_flashnext_mtp_batch_policy.py @@ -20,6 +20,7 @@ def test_mtp_batch_policy_is_owned_hashed_and_serializable(): tiled = KernelConfig( sm70_hcx=True, sm70_hcx_batch_max_m=40, + sm70_hcx_batch_fusion=True, qsa_host_kv_state_compact=True, sm70_qsa_shared_key_batch=True, sm70_gguf={ @@ -27,6 +28,9 @@ def test_mtp_batch_policy_is_owned_hashed_and_serializable(): "small_m_dp4a_max_m": 40, "small_m_hmma_batch": True, "grouped_gate_up": True, + "grouped_gate_up_max_m": 40, + "grouped_down": True, + "shared_expert_batch": True, }, qsa_host_kv=True, qsa_host_indexer_history=True, @@ -51,11 +55,13 @@ def test_mtp_batch_policy_is_owned_hashed_and_serializable(): assert worker.exitcode == 0 assert legacy.sm70_hcx_batch_max_m == 8 assert legacy.sm70_gguf.small_m_dp4a_max_m == 20 + assert legacy.sm70_gguf.grouped_gate_up_max_m == 20 assert not legacy.qsa_host_kv_state_compact assert transferred == tiled assert transferred.compute_hash() == tiled.compute_hash() for field in ( "sm70_hcx_batch_max_m", + "sm70_hcx_batch_fusion", "qsa_host_kv_state_compact", "sm70_qsa_shared_key_batch", "sm70_qsa_indexer_single_rank", @@ -67,6 +73,16 @@ def test_mtp_batch_policy_is_owned_hashed_and_serializable(): transferred.sm70_gguf.small_m_dp4a_max_m = 20 assert transferred.compute_hash() != tiled.compute_hash() + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.grouped_gate_up_max_m = 20 + assert transferred.compute_hash() != tiled.compute_hash() + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.grouped_down = False + assert transferred.compute_hash() != tiled.compute_hash() + transferred = copy.deepcopy(tiled) + transferred.sm70_gguf.shared_expert_batch = False + assert transferred.compute_hash() != tiled.compute_hash() + transferred = copy.deepcopy(tiled) transferred.sm70_gguf.grouped_gate_up = False assert transferred.compute_hash() != tiled.compute_hash() @@ -82,6 +98,7 @@ def test_disabled_parents_keep_inactive_batch_options_out_of_hash(): sm70_hcx=False, sm70_qsa_shared_key=False, sm70_hcx_batch_max_m=40, + sm70_hcx_batch_fusion=True, sm70_qsa_shared_key_batch=True, sm70_qsa_indexer_stage_keys=True, ) @@ -108,3 +125,9 @@ def test_mtp_batch_policy_rejects_unsupported_limits(limit): KernelConfig(sm70_hcx_batch_max_m=limit) with pytest.raises(ValueError): KernelConfig(sm70_gguf={"small_m_dp4a_max_m": limit}) + + +@pytest.mark.parametrize("limit", [19, 21, 39, 41]) +def test_grouped_gate_policy_rejects_unqualified_fragments(limit): + with pytest.raises(ValueError): + KernelConfig(sm70_gguf={"grouped_gate_up_max_m": limit}) diff --git a/tests/kernels/quantization/test_gguf_grouped_down.py b/tests/kernels/quantization/test_gguf_grouped_down.py new file mode 100644 index 0000000000..2d7fed6116 --- /dev/null +++ b/tests/kernels/quantization/test_gguf_grouped_down.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Grouped original down retains TP block boundaries and route-sum order.""" + +import numpy as np +import pytest +import torch +import vllm._C # noqa: F401 + +from tests.kernels.quantization.test_gguf_dp4a import quantize_reference +from vllm.model_executor.layers.quantization.gguf_grouped_gate import grouped_down +from vllm.model_executor.layers.quantization.gguf_native import packed_tp_span +from vllm.transformers_utils.gguf_tensor_reader import dequantize, quant_size + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("source", [20, 42]) +@pytest.mark.parametrize("rank", [0, 1, 2, 3]) +@pytest.mark.parametrize("m", [25, 30, 40]) +def test_grouped_down_changed_graph_exact_and_official_oracle(source, rank, m): + experts, n, full_k, k, top_k = 512, 32, 640, 160, 10 + block, size = quant_size(source) + rng = np.random.default_rng(20261015 + source) + raw = rng.integers(0, 256, (experts, n, full_k // block, size), dtype=np.uint8) + raw[..., :2] = np.array([0.002], dtype="