diff --git a/benchmarks/kernels/build_sm70_qsa_topk_sidecar.py b/benchmarks/kernels/build_sm70_qsa_topk_sidecar.py new file mode 100644 index 0000000000..402e967319 --- /dev/null +++ b/benchmarks/kernels/build_sm70_qsa_topk_sidecar.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build current QSA decode specialization for source-overlay validation.""" + +import argparse +import hashlib +import json +from pathlib import Path + +from torch.utils.cpp_extension import load + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--build-dir", type=Path, required=True) + args = parser.parse_args() + args.build_dir.mkdir(parents=True, exist_ok=True) + source = Path(__file__).with_name("sm70_qsa_topk_sidecar.cu") + header = source.parents[2] / "csrc/qsa_lexicographic_topk.cuh" + library = load( + name="vllm_qsa_decode_topk_sm70", + sources=[str(source)], + extra_cuda_cflags=["-O3", "-lineinfo"], + build_directory=str(args.build_dir.resolve()), + is_python_module=False, + verbose=True, + ) + print( + json.dumps( + { + "library": library, + "library_sha256": hashlib.sha256( + Path(library).read_bytes() + ).hexdigest(), + "header_sha256": hashlib.sha256(header.read_bytes()).hexdigest(), + } + ) + ) diff --git a/benchmarks/kernels/sm70_qsa_topk_sidecar.cu b/benchmarks/kernels/sm70_qsa_topk_sidecar.cu new file mode 100644 index 0000000000..7ee0e2b0ad --- /dev/null +++ b/benchmarks/kernels/sm70_qsa_topk_sidecar.cu @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include +#include + +#include "../../csrc/qsa_lexicographic_topk.cuh" + +namespace { +void topk(torch::Tensor logits, torch::Tensor lengths, torch::Tensor output, + int64_t k, bool control) { + TORCH_CHECK(logits.is_cuda() && lengths.is_cuda() && output.is_cuda(), + "QSA tensors must be CUDA"); + TORCH_CHECK(logits.device() == lengths.device() && + logits.device() == output.device(), "QSA device mismatch"); + TORCH_CHECK(logits.scalar_type() == torch::kFloat32 && + lengths.scalar_type() == torch::kInt32 && + output.scalar_type() == torch::kInt32, "QSA dtype mismatch"); + TORCH_CHECK(k == 512 && logits.dim() == 2 && lengths.dim() == 1 && + output.dim() == 2 && lengths.numel() == logits.size(0) && + output.size(0) == logits.size(0) && output.size(1) == k && + logits.stride(1) == 1 && lengths.is_contiguous() && + output.is_contiguous(), "QSA shape mismatch"); + if (!logits.size(0)) return; + const c10::cuda::CUDAGuard guard(logits.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + if (control) { + vllm::qsa::qsa_lexicographic_topk_kernel<512> + <<>>( + logits.data_ptr(), lengths.data_ptr(), + output.data_ptr(), logits.size(0), logits.size(1), + logits.stride(0)); + } else { + vllm::qsa::launch_qsa_lexicographic_topk<512>( + logits.data_ptr(), lengths.data_ptr(), + output.data_ptr(), logits.size(0), logits.size(1), + logits.stride(0), stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +void candidate(torch::Tensor x, torch::Tensor n, torch::Tensor y, int64_t k) { + topk(x, n, y, k, false); +} +void baseline(torch::Tensor x, torch::Tensor n, torch::Tensor y, int64_t k) { + topk(x, n, y, k, true); +} +int64_t version() { return 1; } +} // namespace + +TORCH_LIBRARY_FRAGMENT(_C_qsa_sm70, ops) { + ops.def("qsa_lexicographic_topk(Tensor logits, Tensor lengths, " + "Tensor(a!) output, int top_k) -> ()"); + ops.impl("qsa_lexicographic_topk", torch::kCUDA, &candidate); + ops.def("decode_specialization_version() -> int", &version); +} +TORCH_LIBRARY_FRAGMENT(_C_qsa_verify, ops) { + ops.def("baseline(Tensor logits, Tensor lengths, Tensor(a!) output, int k) -> ()"); + ops.impl("baseline", torch::kCUDA, &baseline); +} diff --git a/benchmarks/kernels/verify_sm70_qsa_resolved.py b/benchmarks/kernels/verify_sm70_qsa_resolved.py new file mode 100644 index 0000000000..b7eea92b32 --- /dev/null +++ b/benchmarks/kernels/verify_sm70_qsa_resolved.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen address resolution plus unchanged sparse attention/merge, M1 TP4.""" + +import argparse +import json +from functools import partial +from pathlib import Path +from statistics import median + +import torch +from verify_sm70_qsa_router_exact import capture + +from vllm.models.qwen4_exp.nvidia.ops import qsa + +ORIGINAL_GATE = qsa._use_sm70_qsa_resolved_indices + + +def paired(ga, gb): + for _ in range(2500): + ga.replay() + gb.replay() + torch.cuda.synchronize() + values = {"control": [], "resolved": []} + for turn in range(8): + pairs = [("control", ga), ("resolved", gb)] + if turn % 2: + pairs.reverse() + for label, graph in pairs: + for _ in range(20): + graph.replay() + a, b = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + a.record() + for _ in range(100): + graph.replay() + b.record() + b.synchronize() + values[label].append(a.elapsed_time(b) / 100) + return { + "samples_ms": values, + "median_ms": {k: median(v) for k, v in values.items()}, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("out", type=Path) + parser.add_argument("--interleaved-kv", action="store_true") + parser.add_argument("--skip-timing", action="store_true") + args = parser.parse_args() + torch.cuda.set_device(0) + torch.manual_seed(20260905) + qsa._SM70_QSA_XQA_PAGE4 = False + results = [] + for context in (8192, 32768, 262144): + page, layers, width = 400, 12, 2051 + blocks = (context + page - 1) // page + if args.interleaved_kv: + # Match the worker ABI: [blocks, 2, page, KV heads, head dim]. + kv = torch.randn( + layers, blocks, 2, page, 1, 256, device="cuda", dtype=torch.float16 + ) + k, v = kv.unbind(2) + else: + kv = None + k = torch.randn( + layers, blocks, page, 1, 256, device="cuda", dtype=torch.float16 + ) + v = torch.randn_like(k) + queries = torch.randn(layers, 1, 6, 256, device="cuda", dtype=torch.float16) + gates = torch.randn_like(queries) + indices = torch.randint( + context, (layers, 1, width), device="cuda", dtype=torch.int32 + ) + tables = torch.stack( + [ + torch.randperm(blocks, device="cuda", dtype=torch.int32) + for _ in range(layers) + ] + ).view(layers, 1, blocks) + requests = torch.zeros(layers, 1, device="cuda", dtype=torch.int32) + a, b = torch.empty_like(queries), torch.empty_like(queries) + + def run( + output, + resolved, + state=(queries, k, v, indices, tables, requests, gates), + layers=layers, + ): + queries, k, v, indices, tables, requests, gates = state + qsa._use_sm70_qsa_resolved_indices = ( + ORIGINAL_GATE if resolved else lambda *args: False + ) + try: + for i in range(layers): + qsa.qsa_sparse_paged_attention( + queries[i], + k[i], + v[i], + indices[i], + tables[i], + requests[i], + out=output[i], + output_gate=gates[i], + ) + finally: + qsa._use_sm70_qsa_resolved_indices = ORIGINAL_GATE + + ga, gb = capture(partial(run, a, False)), capture(partial(run, b, True)) + for scenario in range(8): + indices.random_(context) + queries.normal_() + tables.copy_(tables.roll(1, dims=-1)) + requests.zero_() + if scenario == 1: + indices[:, :, ::7] = -1 + if scenario == 2: + indices[:, :, ::7] = blocks * page + if scenario == 3: + tables[:, :, 0] = -1 + if scenario == 4: + tables[:, :, 1] = blocks + if scenario == 5: + indices.zero_() + if scenario == 6: + requests.fill_(-1) + if scenario == 7: + requests.fill_(1) + b.fill_(float("nan")) + ga.replay() + gb.replay() + torch.cuda.synchronize() + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), ( + context, + scenario, + (a - b).abs().max().item(), + ) + # Restore valid, varied metadata for timing, not the all-invalid case. + requests.zero_() + indices.random_(context) + tables.copy_( + torch.stack( + [ + torch.randperm(blocks, device="cuda", dtype=torch.int32) + for _ in range(layers) + ] + ).view_as(tables) + ) + result = { + "context": context, + "layers": layers, + "page_size": page, + "cache_layout": "interleaved_kv" if args.interleaved_kv else "separate_kv", + "key_strides": list(k[0].stride()), + "bitwise_graph_scenarios": 8, + **({} if args.skip_timing else paired(ga, gb)), + } + results.append(result) + args.out.write_text(json.dumps(results, indent=2) + "\n") + print(json.dumps(result), flush=True) + del ga, gb, run, kv, k, v, queries, gates, indices, tables, requests, a, b + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/verify_sm70_qsa_router_exact.py b/benchmarks/kernels/verify_sm70_qsa_router_exact.py new file mode 100644 index 0000000000..c3bd8f27dc --- /dev/null +++ b/benchmarks/kernels/verify_sm70_qsa_router_exact.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paired exactness/graph screen; no model weights or serving process.""" + +import argparse +import json +from pathlib import Path +from statistics import median + +import torch + +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import ( + _sm70_qwen38_router_topk_kernel as router, +) + + +def capture(fn): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + return graph + + +def timing(graph, loops=200): + samples = [] + for _ in range(5): + for _ in range(20): + graph.replay() + start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + start.record() + for _ in range(loops): + graph.replay() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) / loops) + return {"median_ms": median(samples), "samples_ms": samples} + + +def paired_timing(control, candidate): + # Interleave A/B and reverse order to expose clock/cache drift. + for _ in range(3000): + control.replay() + candidate.replay() + torch.cuda.synchronize() + samples = {"control": [], "candidate": []} + for repeat in range(8): + pairs = [("control", control), ("candidate", candidate)] + if repeat % 2: + pairs.reverse() + for label, graph in pairs: + for _ in range(20): + graph.replay() + a, b = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + a.record() + for _ in range(200): + graph.replay() + b.record() + b.synchronize() + samples[label].append(a.elapsed_time(b) / 200) + return { + label: {"median_ms": median(v), "samples_ms": v} for label, v in samples.items() + } + + +def router_screen(): + def out(rows): + return ( + torch.empty(rows, 10, dtype=torch.float32, device="cuda"), + torch.empty(rows, 10, dtype=torch.int32, device="cuda"), + torch.empty(rows, 10, dtype=torch.int32, device="cuda"), + ) + + def launch(x, dst, packed): + router[(x.shape[0],)]( + x, + *dst, + E=512, + K=10, + M=x.shape[0], + BLOCK_E=512, + PACKED_HALF_KEY=packed, + num_warps=8, + ) + + # Exhaust all raw half encodings, then finite-only permutations and ties. + raw = torch.arange(65536, dtype=torch.int32, device="cuda").to(torch.int16) + exhaustive = raw.view(torch.float16).reshape(128, 512) + finite = torch.nan_to_num(exhaustive, nan=0, posinf=65504, neginf=-65504) + cases = [ + exhaustive, + finite, + finite.flatten()[torch.randperm(65536, device="cuda")].reshape(128, 512), + ] + for scale in (0.001, 0.1, 1.0, 10.0): + cases.append(torch.randn(128, 512, dtype=torch.float16, device="cuda") * scale) + cases += [torch.zeros_like(finite), torch.full_like(finite, -float("inf"))] + cases.append((torch.arange(512, device="cuda") % 7).half().repeat(128, 1)) + samples = 0 + for x in cases: + a, b = out(x.shape[0]), out(x.shape[0]) + launch(x, a, False) + launch(x, b, True) + for ref, got in zip(a, b): + assert torch.equal(ref.view(torch.int32), got.view(torch.int32)) + samples += x.shape[0] + + x = torch.randn(48, 512, dtype=torch.float16, device="cuda") + a, b = out(48), out(48) + + def run(dst, packed): + for i in range(48): + launch(x[i : i + 1], tuple(t[i : i + 1] for t in dst), packed) + + ga, gb = capture(lambda: run(a, False)), capture(lambda: run(b, True)) + for _ in range(16): + x.normal_() + for value in b: + value.fill_(-777) + ga.replay() + gb.replay() + assert all( + torch.equal(p.view(torch.int32), q.view(torch.int32)) for p, q in zip(a, b) + ) + times = paired_timing(ga, gb) + return { + "bitwise_rows": samples, + "changing_graph_replays": 16, + "control": times["control"], + "packed32": times["candidate"], + "calls": 48, + } + + +def qsa_screen(): + base = torch.ops._C_qsa_verify.baseline + fast = torch.ops._C_qsa_sm70.qsa_lexicographic_topk + assert torch.ops._C_qsa_sm70.decode_specialization_version() == 1 + count = 0 + for rows in (1, 2): + for length in (0, 1, 511, 512, 2048, 2304, 2305, 4096, 65536): + width = max(512, length) + x = torch.randn(rows, width, device="cuda") + n = torch.full((rows,), length, dtype=torch.int32, device="cuda") + a = torch.empty(rows, 512, dtype=torch.int32, device="cuda") + b = torch.empty_like(a) + for pattern in ("random", "ties", "zeros", "special"): + if pattern == "ties": + x.copy_(torch.arange(width, device="cuda") % 7) + elif pattern == "zeros": + x.zero_() + x[:, ::2] = -0.0 + elif pattern == "special": + x[:, 0], x[:, 1], x[:, 2] = ( + float("nan"), + float("inf"), + -float("inf"), + ) + a.fill_(-111) + b.fill_(-111) + base(x, n, a, 512) + fast(x, n, b, 512) + assert torch.equal(a, b), (rows, length, pattern) + count += 1 + x = torch.randn(12, 1, 4096, device="cuda") + n = torch.tensor( + [[2048 + i * 11] for i in range(12)], dtype=torch.int32, device="cuda" + ) + a = torch.empty(12, 1, 512, dtype=torch.int32, device="cuda") + b = torch.empty_like(a) + + def run(op, out): + for i in range(12): + op(x[i], n[i], out[i], 512) + + ga, gb = capture(lambda: run(base, a)), capture(lambda: run(fast, b)) + for _ in range(16): + x.normal_() + b.fill_(-777) + ga.replay() + gb.replay() + assert torch.equal(a, b) + times = paired_timing(ga, gb) + return { + "exact_cases": count, + "changing_graph_replays": 16, + "control": times["control"], + "decode": times["candidate"], + "calls": 12, + } + + +def qsa_dynamic_screen(base=None): + """Change device lengths across fast/fallback boundaries in one graph.""" + if base is None: + base = torch.ops._C_qsa_verify.baseline + fast = torch.ops._C_qsa_sm70.qsa_lexicographic_topk + rows, width, replays = 12, 65536, 128 + # Padded row storage also tests that an M1 view need not have stride=width. + storage = torch.empty(rows, width + 32, device="cuda") + x = storage[:, :width] + n = torch.zeros(rows, 1, dtype=torch.int32, device="cuda") + a = torch.empty(rows, 1, 512, dtype=torch.int32, device="cuda") + b = torch.empty_like(a) + lengths = torch.tensor( + [-1, 0, 511, 512, 513, 2048, 2169, 2304, 2305, 4096, 65536, 65537], + dtype=torch.int32, + device="cuda", + ) + + def run(op, out): + for i in range(rows): + op(x[i : i + 1], n[i], out[i], 512) + + ga, gb = capture(lambda: run(base, a)), capture(lambda: run(fast, b)) + for replay in range(replays): + n[:, 0].copy_(lengths.roll(replay % rows)) + x.normal_() + pattern = replay % 4 + if pattern == 0: + x.relu_() # Dense zero ties, as in ReLU-based indexer scores. + elif pattern == 1: + x.mul_(1e-6).add_(1) # Shared exponent and many close scores. + elif pattern == 2: + x.round_() + else: + x[:, ::19] = float("nan") + x[:, ::23] = -float("inf") + x[:, ::29] = float("inf") + a.fill_(-111) + b.fill_(-777) + ga.replay() + gb.replay() + assert torch.equal(a, b), ("dynamic lengths", replay, pattern) + return { + "changing_length_graph_replays": replays, + "row_comparisons": replays * rows, + "lengths": lengths.cpu().tolist(), + "patterns": ["relu", "near_ties", "integer_ties", "nonfinite"], + "padded_row_stride": storage.stride(0), + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--qsa-library", required=True) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + torch.cuda.set_device(0) + assert torch.cuda.get_device_capability() == (7, 0) + torch.manual_seed(20260905) + torch.ops.load_library(args.qsa_library) + result = {"scope": "model-free GPU operator/graph screen, not endpoint speed"} + result["router"] = router_screen() + args.out.write_text(json.dumps(result, indent=2) + "\n") + result["qsa"] = qsa_screen() + result["qsa_dynamic"] = qsa_dynamic_screen() + result.update( + torch=torch.__version__, + cuda=torch.version.cuda, + gpu=torch.cuda.get_device_name(), + ) + args.out.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2), flush=True) diff --git a/docs/design/sm70_qwen38_exact_decode_followup.md b/docs/design/sm70_qwen38_exact_decode_followup.md new file mode 100644 index 0000000000..0b3052574f --- /dev/null +++ b/docs/design/sm70_qwen38_exact_decode_followup.md @@ -0,0 +1,230 @@ +# Qwen3.8 exact single-request decode follow-up + +Integration line: `public/main`, base `755baae1d075ee04fa9096b23fc0225b23589a86`. +This task is stacked on HC PR #506 (`4ae6a0005a`) to preserve the measured +single-request contract; it does not duplicate the M4/M8/M16 work in #504 or +the Page4 relocation-order fix in #494. Human review is required. AI assistance +was used (OpenAI Codex). + +## Baseline and scope + +TP4 V100-SXM2-32GB GPU0–3, Torch2.10.0+cu128, Qwen3.8-Flash-Next-NVFP4, +FP16 activations/KV, native NVFP4 experts, no MTP/prefix cache, hybrid PLE, +V2 dual CUDA graph, max length262144, prefill chunk8192, single request. +8K/513 before-capture engine baseline: 96.394713/96.417843/96.420508 tok/s, +aggregate96.411020, TPOT10.372258ms. Configured256K is not an input-quality gate. +Three greedy baseline token sequences equal the previous baseline; two short +official-sampling/thinking/natural-EOS smokes passed. Full256K quality pending. + +The matching 8K/32 graph-node trace has 29 middle replay windows/rank. Complete +HC is2.033180ms; QSA1.219981ms including top-k0.307393ms; router top-k0.486903ms. +These are GPU service times, not additive end-to-end latency. + +Raw baseline directory (outside Git): +`/home/ymzx/桌面/1cat-vllm/worktrees/v100-qwen38-nomtp-token-trace-20260903-173451/.artifacts/hc_trace_vector_20260905_retry1/`. + +## Three bounded directions + +1. Build the existing QSA single-row decode top-k from the frozen source and + verify that the actual runtime uses it. The baseline trace contains the + generic kernel even at grid1; source selects the specialized kernel for M1. +2. Screen a lossless 32-bit router sort key for FP16 logits/E512/K10/M1. + Keep max/exp/normalization reduction, tie rules and invalid-row semantics. + Other dtypes and widths retain the original 64-bit-key path. +3. Screen precomputed sparse KV addresses, keeping logical token order, + split/merge arithmetic, dtype and FP16 boundaries. Do not assume a gather + or persistent-kernel rewrite is profitable. + +No lower precision, changed accumulation order, top-k truncation, MTP or +top1-only LM-head shortcuts. First prove numerical/graph/shape gates on small +operators, then use a combined whole-model run when justified. Never preempt +unrelated GPU owners. Keep failed paths and raw evidence. + +## Current status + +Draft PR #507 is stacked on #506. The three operator paths are implemented +and hit the combined model trace. A whole-model run measured98.965175 tok/s. +The frozen baseline's greedy output was not reproducible on an isolated +old-source control either. No new operator mismatch was found, but do not +promote this as complete deterministic-output or broad quality acceptance. +Detailed evidence and the existing repeatability issue follow. + +The router keeps the original 8-warp FP32 max/exp/normalization tree. Its +32-bit key is selected only for existing FP16 E512/K10 M1 routing. Signed-zero, +NaN/Inf, ties, all raw half encodings and finite shuffled inputs are covered: +1,280 rows are bitwise equal, plus16 changed-input/poisoned graph replays. +After warmup, eight alternating-order A/B pairs measure48 calls at +**0.252652 ->0.203530ms** (save0.049121ms). An earlier non-interleaved screen +drifted in clock state and is retained but not used for the accepted delta. + +The QSA sidecar builder compiles the production header, exposes a version +marker and a literal old generic-kernel control.72 cases cover M1/M2, lengths +0/1/511/512/2048/2304/2305/4096/65536, ties, signed-zero and non-finite scores; +all selected IDs are exact.16 changed-input/poisoned graph replays pass. +Twelve calls at live lengths2048–2169 measure **0.258196 ->0.111665ms**. +Source-overlay services must pin the freshly built library through existing +`VLLM_SM70_QSA_TOPK_LIBRARY`; do not infer native-binary freshness from Python +source or an old unversioned log message. The endpoint trace must prove the +decode-specialized symbol is present. Header SHA256: +`e09d4af611894d2c3613ea1d5ac50e1fd2606f729e6fa7c45eb8079fcc6b9508`. +Sidecar SHA256: +`e56d7874877dddb88589a7e08ecbe9074f7f7532af1be764ab5f47a9b8aa8165`. + +QSA address resolution retains logical ordering/duplicates and validity, +precomputes physical token slots, and removes the dependent page-table load +from the unchanged partial attention arithmetic. It is limited to SM70, +FP16 M1/Q6/KV1/D256/page400/selection2051 and signed-int32-safe physical slots. +Other shapes, cache formats and prefill keep the original path. Eight changing +graph scenarios per length test page relocation, invalid pages/indices and +requests, duplicates and poisoned outputs; all outputs are bitwise identical. +Public production-dispatch screen (12 attention+merge calls, resolver included): + +| Cache context | Original ms | Resolved ms | +|---|---:|---:| +|8192|0.345016|0.318909| +|32768|0.384205|0.367094| +|262144|0.365860|0.323968| + +These are operator service times, not additive endpoint savings. The256K row +is an operator cache-size check, not a256K model-input quality acceptance. +CPU dispatcher/QSA launch suites:49 passed,1 skipped (GPU-only). Targeted Ruff +checks pass. No new lower-precision weights, KV or arithmetic introduced. +An additional correctness-only run uses the real worker's interleaved K/V +layout, strides `[204800,256,256,1]`; eight graph scenarios at each of8K/32K/ +256K are bitwise equal. The timing table above uses separate contiguous K/V +storage; do not relabel it as an interleaved-layout latency measurement. +Reproduce the additional gate with `--interleaved-kv --skip-timing`; raw output +is `.artifacts/three_paths/address_interleaved_exact.json`. + +Reproduction (project Python environment, SM70 GPU ownership required): + +```bash +CUDA_HOME=/usr TORCH_CUDA_ARCH_LIST=7.0 .venv/bin/python \ + benchmarks/kernels/build_sm70_qsa_topk_sidecar.py --build-dir .artifacts/qsa-build +.venv/bin/python benchmarks/kernels/verify_sm70_qsa_router_exact.py \ + --qsa-library .artifacts/qsa-build/vllm_qsa_decode_topk_sm70.so \ + --out .artifacts/operators.json +.venv/bin/python benchmarks/kernels/verify_sm70_qsa_resolved.py \ + .artifacts/address.json +``` + +Raw task artifacts: `.artifacts/three_paths/operators_interleaved.json`, +`address_production.json`, build/queue logs. All standalone GPU screens exited. +## Combined endpoint, source d2c8401c22 + +One model initialization produced `.artifacts/endpoint/result.json`, +`quality.json`, `contract.txt`, `run.log`, and the8K/32 graph-node report. +Workers explicitly log QSA source-overlay decode specialization version1; +the trace contains `qsa_lexicographic_decode_topk_kernel` and the new +`_qsa_resolve_physical_indices_kernel`. The old HC/W2 binaries and M2-disable +compatibility pin are unchanged. + +| Unprofiled8K/513, three repeats | Frozen baseline | Candidate | +|---|---:|---:| +| Steady decode tok/s |96.411020|98.965175| +| TPOT ms |10.372258|10.104565| + +Candidate repeats:98.939675/98.977150/98.978711 tok/s. The observed difference +is0.267694ms/token (2.65% throughput), not a three-repeat same-output delta. +This does not meet100 tok/s. Mean TTFT1.169634s and prefill1.166232s; the +fixed8192-token prompt corresponds to approximately7024 prefill tok/s. + +The graph trace has31 replays/rank,29 middle windows. Rank-average diagnostic +GPU service (ms/token; do not sum these to close endpoint wall time): + +| Work | Frozen trace | Candidate trace | +|---|---:|---:| +| QSA selected-block top-k |0.307393|0.175693| +| Router top-k |0.486903|0.356512| +| QSA partial attention |0.436250|0.382302| +| QSA split merge |0.117537|0.116910| +| New physical-address resolver |0|0.028287| +| Whole QSA category |1.219981|1.060166| + +The candidate trace includes a21.585578ms interval and7.718280ms following +rank-start skew. Keep these samples: mean rank-max interval11.328738ms, +median10.840832ms. Complete HC mean2.100456ms/p502.022327ms versus baseline +2.033180ms/p502.031939ms; its code is unchanged and waiting distorts the mean. +Do not call this an HC regression/improvement or replace unprofiled TPOT with +the trace interval. The previous HC1.5ms target remains unmet. + +Health checks use official temperature1/top-p0.95/top-k20, seed0, thinking, +natural EOS. Short arithmetic, exact record copy, and261632-token padded +arithmetic all pass (125/110/140 output tokens). The task is at the end of the +long prompt: this is not long-range retrieval or comprehensive quality. +The262143-input +1-output request completes the262144 boundary; its50.758678s +prefill is approximately5164 tok/s, and it has no steady decode iteration. + +### Whole-model repeatability investigation + +Do not hide this difference behind the three passing health checks. The old +greedy/ignore-EOS513-token performance repeats were identical. All three new +repeats have different hashes. Repeats1/3 match the old natural prefix through +first EOS at index28 and diverge at index37, after forced continuation. +Repeat2 emits EOS at index8, so the mismatch is **not only after EOS**. +Precision/arithmetic source contracts and operator checks pass; that alone +does not prove unchanged full-model output. + +Frozen token hash: +`7385dacbed6a3d06576993bda99375b51c6a2e6132ee4f2fc0079646b461fca1`. +Candidate hashes: + +- `f1ad6ecc74c99ff5980d993884b44e7e147d477637353c9e750310c9ccaa08a8` +- `8419c18737e8ef9135e1a9f5e3cd788ea6561d63b3168a8f307f8deb48874899` +- `1888a2872baf55ec768df79cb7b8877b4e5911970ee949ffa213f815cdc55bb2` + +Model-free follow-up against the **installed old native QSA binary**, rather +than only the rebuilt control:72 cases and16 changing-logit graph replays +pass. An additional128 changing-length graph replays/1536 row comparisons +pass, including fast/fallback transitions, ReLU zero ties, near ties, +non-finite scores and padded row storage. Public reproduction function: +`qsa_dynamic_screen(torch.ops._C.qsa_lexicographic_topk)` after loading the +candidate sidecar. Raw files are `installed_topk_oracle.json` and +`dynamic_topk_oracle.json` under `.artifacts/three_paths/`. + +Known test-condition differences: the new run uses task-owned fresh compiler +caches and performs the261632-token health request **before** the speed +repeats; the old baseline had only short health requests. A second and final +model initialization tested these conditions on clean old-source4ae6a0005a, +without any of the three candidates and without Nsight capture. + +Control `.artifacts/control/`: two short health requests,513/32-token warmups, +three8K/513 repeats, the same261632-token health request, then three8K/513 +repeats. Both sets use the same loaded engine. All health checks pass naturally +(109/167/90 tokens). Before the long request the repeats have EOS indices +8/28/28 and two distinct hashes; after it they have8/8/8 and three hashes. +Therefore long-request history is **not necessary** for the observed drift. +The old-source control reproduces the candidate's early-EOS hash +`8419c18737e8ef9135e1a9f5e3cd788ea6561d63b3168a8f307f8deb48874899` +in all513 tokens, both before and after the long request. This specific +repeatability symptom demonstrably predates the three optimizations; it is +not evidence by itself of a new precision reduction or new wrong top-k. +It does not prove that every other generation difference is harmless. + +Control mean decode is96.107235 tok/s before and96.083709 after the long +request. One identical-output old/new pair measures96.106397 versus98.977150 +tok/s (0.301793ms/token less); this is one pair, not a replacement three-repeat +controlled acceptance benchmark. The primary observed candidate result remains +98.965175 tok/s. The CPU audit is `.artifacts/repeatability_audit.json`. + +Related existing [PR #494](https://github.com/1CatAI/1Cat-vLLM/pull/494) fixes +page4 physical-allocation-dependent attention ordering and is still open, +not merged, at this audit. Its established issue is a relevant next suspect, +not a proven explanation for every NVFP4 output above. Reuse/review that fix +before inventing another planner rewrite; do not silently adopt a new FP32 +arithmetic or NCCL policy to force matching text. Actual-tensor capture and a +same-engine intervention would be needed for causal attribution here. + +All two model jobs, operator jobs, samplers and queues are terminal. GPU0–3 +were released; GPU4–7 belong to an unrelated service. Final integration is +`public/main`; #507 remains Draft, stacked on #506, pending human review and +resolution of the relevant existing quality/integration gate. No main push +or resident API was performed. Both100 tok/s and HC1.5ms remain unmet. + +Retained SHA256 checksums: + +- Candidate result: `e4b27b6e54298cf26cbd5acde28046a5dfed6210265b752ce7df1ac75ad2c8cb` +- Candidate health: `2a79fdfc172d709c9245d911a195d88c285dffd947011a2e306b4914b3ec204f` +- Candidate trace: `d1afc22e0d91d5c7ce94bfaa3132675302af4ef74e2be18456dec4983b897cd2` +- Old control result: `95c577b208090aa605db3402641e563a8d3892ef55823d0dc4bcb80478786ccb` +- Old control health: `fe3b26835418ca7508fcbef275560c293073f2f611872656785f5fa5631c9944` diff --git a/tests/kernels/moe/test_sm70_router_key_dispatch.py b/tests/kernels/moe/test_sm70_router_key_dispatch.py new file mode 100644 index 0000000000..965569a17b --- /dev/null +++ b/tests/kernels/moe/test_sm70_router_key_dispatch.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU route guards for the lossless FP16 E512/K10 sort-key specialization.""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.router import fused_topk_router as mod + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("rows", [1, 2, 5, 16]) +def test_packed_half_key_is_m1_half_only(monkeypatch, dtype, rows): + calls = [] + + class Kernel: + def __getitem__(self, grid): + def launch(*args, **kwargs): + calls.append((grid, kwargs)) + + return launch + + monkeypatch.setattr(mod, "_sm70_qwen38_router_topk_kernel", Kernel()) + x = torch.empty(rows, 512, dtype=dtype) + weights = torch.empty(rows, 10, dtype=torch.float32) + ids = torch.empty(rows, 10, dtype=torch.int32) + mod._sm70_qwen38_router_topk(weights, ids, torch.empty_like(ids), x) + assert len(calls) == 1 + grid, kwargs = calls[0] + assert grid == (rows,) + assert kwargs["PACKED_HALF_KEY"] == (rows == 1 and dtype == torch.float16) + assert kwargs["num_warps"] == 8 # Keep the FP32 normalization reduction. diff --git a/tests/models/qwen4_exp/test_qsa_resolved_dispatch.py b/tests/models/qwen4_exp/test_qsa_resolved_dispatch.py new file mode 100644 index 0000000000..05a28b1920 --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_resolved_dispatch.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace as NS + +import pytest +import torch + +from vllm.models.qwen4_exp.nvidia.ops import qsa + + +@pytest.mark.parametrize("blocks", [21, 82, 656]) +def test_exact_m1_cache_geometry(monkeypatch, blocks): + monkeypatch.setattr(qsa.current_platform, "is_device_capability", lambda c: c == 70) + q = NS(shape=(1, 6, 256), dtype=torch.float16) + k = NS(shape=(blocks, 400, 1, 256), dtype=torch.float16) + ids = NS(shape=(1, 2051), dtype=torch.int32) + assert qsa._use_sm70_qsa_resolved_indices(q, k, ids, "float16") + + +@pytest.mark.parametrize( + "bad", ["arch", "rows", "dtype", "page", "width", "overflow", "kv"] +) +def test_unsupported_routes_unchanged(monkeypatch, bad): + monkeypatch.setattr( + qsa.current_platform, "is_device_capability", lambda c: bad != "arch" + ) + q = NS(shape=(1, 6, 256), dtype=torch.float16) + k = NS(shape=(656, 400, 1, 256), dtype=torch.float16) + ids = NS(shape=(1, 2051), dtype=torch.int32) + if bad == "rows": + q.shape = (2, 6, 256) + if bad == "dtype": + q.dtype = torch.bfloat16 + if bad == "page": + k.shape = (400, 784, 1, 256) + if bad == "width": + ids.shape = (1, 512) + if bad == "overflow": + k.shape = (2**31, 400, 1, 256) + dtype = "fp8_e4m3" if bad == "kv" else "float16" + assert not qsa._use_sm70_qsa_resolved_indices(q, k, ids, dtype) diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py index 6ee6ad5e25..f648d0b00a 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py @@ -30,6 +30,7 @@ def _sm70_qwen38_router_topk_kernel( K: tl.constexpr, M: tl.constexpr, BLOCK_E: tl.constexpr, + PACKED_HALF_KEY: tl.constexpr = False, ) -> None: """Sort one exact Qwen3.8 decode or MTP verifier row per program.""" @@ -54,19 +55,39 @@ def _sm70_qwen38_router_topk_kernel( # Transform float32 into an ascending-sortable key. Packing the expert ID # into the low bits preserves the generic kernel's lower-ID tie break. - min_i32: tl.constexpr = -2147483648 - logit_bits = sort_logits.to(tl.int32, bitcast=True) - sign = logit_bits >> 31 - key = tl.where(sign == 0, logit_bits ^ -1, logit_bits ^ min_i32) - key = tl.where(valid, key, 0x7FFFFFFF) - packed = ((key.to(tl.int64) & 0xFFFFFFFF) << 32) | offsets.to(tl.int64) - sorted_packed = tl.sort(packed, descending=False) - - sorted_keys = ((sorted_packed >> 32) & 0xFFFFFFFF).to(tl.int32) - sorted_ids = (sorted_packed & 0xFFFFFFFF).to(tl.int32) - sorted_sign = sorted_keys >> 31 - sorted_bits = tl.where(sorted_sign < 0, sorted_keys ^ -1, sorted_keys ^ min_i32) - sorted_logits = sorted_bits.to(tl.float32, bitcast=True) + if PACKED_HALF_KEY: + # FP16 -> FP32 above is exact. Sort the original 16-bit values plus + # nine expert-ID bits in one int32, without quantizing any logits. + # Degenerate rows use -offsets (0..511), also exactly representable. + tl.static_assert(E == 512 and BLOCK_E == 512) + bits = sort_logits.to(tl.float16).to(tl.int16, bitcast=True).to(tl.int32) + key = tl.where(bits < 0, bits ^ 0x8000, bits ^ 0xFFFF) & 0xFFFF + # Original int64 sort is signed: flip the key sign bit when moving + # to a positive 25-bit key so positive logits still precede negatives. + packed = ((key ^ 0x8000) << 9) | offsets + sorted_packed = tl.sort(packed, descending=False) + sorted_keys = (sorted_packed >> 9) ^ 0x8000 + sorted_ids = sorted_packed & 0x1FF + sorted_bits = tl.where( + (sorted_keys & 0x8000) != 0, + sorted_keys ^ 0xFFFF, + sorted_keys ^ 0x8000, + ).to(tl.uint16) + sorted_logits = sorted_bits.to(tl.float16, bitcast=True).to(tl.float32) + else: + min_i32: tl.constexpr = -2147483648 + logit_bits = sort_logits.to(tl.int32, bitcast=True) + sign = logit_bits >> 31 + key = tl.where(sign == 0, logit_bits ^ -1, logit_bits ^ min_i32) + key = tl.where(valid, key, 0x7FFFFFFF) + packed = ((key.to(tl.int64) & 0xFFFFFFFF) << 32) | offsets.to(tl.int64) + sorted_packed = tl.sort(packed, descending=False) + + sorted_keys = ((sorted_packed >> 32) & 0xFFFFFFFF).to(tl.int32) + sorted_ids = (sorted_packed & 0xFFFFFFFF).to(tl.int32) + sorted_sign = sorted_keys >> 31 + sorted_bits = tl.where(sorted_sign < 0, sorted_keys ^ -1, sorted_keys ^ min_i32) + sorted_logits = sorted_bits.to(tl.float32, bitcast=True) raw_weights = tl.math.exp2((sorted_logits - max_logit) * 1.4426950408889634) raw_weights = tl.where(invalid_row, 0.0, raw_weights) @@ -100,6 +121,7 @@ def _sm70_qwen38_router_topk( K=10, M=num_tokens, BLOCK_E=512, + PACKED_HALF_KEY=(gating_output.dtype == torch.float16 and num_tokens == 1), num_warps=8, ) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index c02b3505c3..3d5940e003 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -24,6 +24,13 @@ _SM70_QSA_TOPK_LIBRARY = os.getenv("VLLM_SM70_QSA_TOPK_LIBRARY") if _SM70_QSA_TOPK_LIBRARY is not None: torch.ops.load_library(_SM70_QSA_TOPK_LIBRARY) + _topk_version = getattr( + torch.ops._C_qsa_sm70, "decode_specialization_version", None + ) + if _topk_version is not None: + logger.info( + "SM70 QSA source-overlay decode specialization version %d.", _topk_version() + ) if hasattr(torch.ops._C_qsa_sm70, "qsa_lexicographic_topk"): @@ -529,6 +536,50 @@ def _qsa_xqa_page4_table_kernel( ) +@triton.jit +def _qsa_resolve_physical_indices_kernel( + indices_ptr, + table_ptr, + token_to_req_ptr, + output_ptr, + stride_indices_row, + stride_table_req, + num_cache_blocks, + num_requests, + TOPK: tl.constexpr, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + columns = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + token = tl.load( + indices_ptr + row * stride_indices_row + columns, mask=columns < TOPK, other=-1 + ) + safe_token = tl.maximum(token, 0) + page = safe_token // PAGE_SIZE + valid = ( + (columns < TOPK) + & (request >= 0) + & (request < num_requests) + & (token >= 0) + & (page < PAGE_TABLE_WIDTH) + ) + physical = tl.load( + table_ptr + safe_request * stride_table_req + page, mask=valid, other=-1 + ) + valid &= (physical >= 0) & (physical < num_cache_blocks) + # Keep logical order, duplicates and invalid slots. Never sort by page. + slot = physical.to(tl.int64) * PAGE_SIZE + safe_token % PAGE_SIZE + tl.store( + output_ptr + row * TOPK + columns, + tl.where(valid, slot, -1), + mask=columns < TOPK, + ) + + @triton.jit(do_not_specialize=["num_requests"]) def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, @@ -571,6 +622,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, KV_E4M3: tl.constexpr, + RESOLVED_INDICES: tl.constexpr = False, ) -> None: row = tl.program_id(0) kv_head = tl.program_id(1) @@ -614,19 +666,18 @@ def _qsa_sparse_paged_gqa_splitk_kernel( safe_token = tl.maximum(logical_token, 0) logical_page = safe_token // PAGE_SIZE page_offset = safe_token % PAGE_SIZE - valid = ( - (request >= 0) - & (request < num_requests) - & (logical_token >= 0) - & (logical_page < PAGE_TABLE_WIDTH) - ) - physical_page = tl.load( - block_table_ptr - + safe_request * stride_table_req - + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1), - mask=valid, - other=-1, - ) + valid = (request >= 0) & (request < num_requests) & (logical_token >= 0) + if RESOLVED_INDICES: + physical_page = logical_page + else: + valid &= logical_page < PAGE_TABLE_WIDTH + physical_page = tl.load( + block_table_ptr + + safe_request * stride_table_req + + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1), + mask=valid, + other=-1, + ) valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) # physical_page * block stride can overflow int32 for large caches. safe_page = tl.maximum(physical_page, 0).to(tl.int64) @@ -2006,6 +2057,20 @@ def _qsa_sparse_paged_attention_sm70_xqa_page4( ) +def _use_sm70_qsa_resolved_indices(q, k_cache, indices, kv_cache_dtype): + """Admit only the measured checkpoint-FP16 M1 TP4 cache geometry.""" + return bool( + current_platform.is_device_capability(70) + and q.shape == (1, 6, 256) + and q.dtype == k_cache.dtype == torch.float16 + and k_cache.shape[1:] == (400, 1, 256) + and k_cache.shape[0] * 400 < 2**31 + and indices.shape == (1, 2051) + and indices.dtype == torch.int32 + and kv_cache_dtype in ("auto", "float16") + ) + + def qsa_sparse_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -2107,6 +2172,28 @@ def qsa_sparse_paged_attention( _qsa_output_gate(xqa_output, output_gate_view) return xqa_output + resolved_indices = _use_sm70_qsa_resolved_indices( + q, k_cache, logical_indices, kv_cache_dtype + ) + if resolved_indices: + physical_indices = torch.empty_like(logical_indices) + _qsa_resolve_physical_indices_kernel[(1, triton.cdiv(2051, 256))]( + logical_indices, + block_table, + token_to_req, + physical_indices, + logical_indices.stride(0), + block_table.stride(0), + k_cache.shape[0], + block_table.shape[0], + TOPK=2051, + PAGE_SIZE=400, + PAGE_TABLE_WIDTH=block_table.shape[1], + BLOCK=256, + num_warps=4, + ) + logical_indices = physical_indices + group_size = q.shape[1] // k_cache.shape[2] block_m = triton.next_power_of_2(group_size) base_programs = q.shape[0] * k_cache.shape[2] @@ -2184,6 +2271,7 @@ def qsa_sparse_paged_attention( BLOCK_M=block_m, BLOCK_N=block_n, KV_E4M3=kv_e4m3, + RESOLVED_INDICES=resolved_indices, num_warps=partial_warps, num_stages=2, )