[Perf][QSA] SM70: score multi-request decode batches by request segments

tests/config/test_revolta_kernel_profile.py

@@ -27,7 +27,7 @@ def test_cascade_profile_preserves_explicit_choices(
     profile.apply_revolta_kernel_profile(kernel, enabled=enabled)
     assert kernel.ple_disk_cascade is expected
     audit = json.loads(profile.profile_audit(enabled=enabled, kernel_config=kernel))
-    assert len(audit["values"]) == 45
+    assert len(audit["values"]) == 46
     assert "ple_disk_cascade" not in audit["values"]
     assert audit["kernel_config"]["ple_disk_cascade"] == {
         "value": expected,

tests/config/test_revolta_profile.py

@@ -23,6 +23,7 @@
     "VLLM_SM70_GREEDY_VERIFY": "1",
     "VLLM_SM70_SKIP_REJECTED_BATCH_PACKS": "1",
     "VLLM_SM70_QSA_SHARED_KEY": "1",
+    "VLLM_SM70_QSA_REQUEST_SEGMENTS": "1",
     "VLLM_SM70_QSA_COMPACT_TOPK": "1",
     "VLLM_SM70_PLE_NGRAM_BATCH": "1",
     "VLLM_SM70_PLE_INPUT_PREPARE": "1",
@@ -196,7 +197,7 @@ def test_sources_after_upstream_auto_defaults(profile, monkeypatch):
     monkeypatch.setenv("NCCL_GRAPH_REGISTER", "0")  # env_override
     audit = json.loads(module.profile_audit(enabled=True))
     assert audit["enabled"]
-    assert len(audit["values"]) == 45
+    assert len(audit["values"]) == 46
     assert audit["values"][FP16[0]]["source"] == "upstream_default"
     assert audit["values"]["NCCL_GRAPH_REGISTER"]["source"] == "upstream_default"
     assert audit["values"]["VLLM_SM70_AWQ_QWEN38_GROUPED_QPN_SPLIT"] == {
@@ -494,7 +495,7 @@ def test_frontend_interconnect_mode_and_audit(
     # policy must not read "pex_p2p" while the workers log nvlink.
     assert audit["interconnect_mode"] == logged
     assert audit["interconnect_policy"] == ("auto" if requested is None else requested)
-    assert len(audit["values"]) == 45
+    assert len(audit["values"]) == 46
 
 
 def _applied_profile(profile, monkeypatch, requested, disabled):

tests/kernels/qsa_request_segments.md

@@ -0,0 +1,78 @@
+# Request-segmented QSA scorer
+
+`vllm/models/qwen4_exp/nvidia/ops/qsa_request_segments.py` scores multi-request
+SM70 QSA batches. The paged scorer (`qsa.qsa_mqa_paged`) loads every 64-key
+tile once per query row. Rows of one request share a page table, so this scorer
+splits each contiguous request run into segments of up to four rows and loads
+each tile once per segment. Four rows by four heads form one 128x16 query block.
+The kernel keeps the paged kernel's 64x128 by 128x16 dot, 2 warps, pipelined
+tile loop and per-row 16-lane reduction (the other rows' lanes are zeroed). On
+SM70 the scores and visible lengths are bitwise equal to the paged scorer.
+
+## Runtime use
+
+`qsa.qsa_select_paged_tokens` uses this scorer in place of `qsa_mqa_paged` when
+`VLLM_SM70_QSA_REQUEST_SEGMENTS=1` (the ReVolta profile sets it; an explicit
+value wins) and `request_segments_unsupported` returns `None`:
+
+- two or more requests (`single_request` keeps the shared-key and cuBLAS paths);
+- at most 256 rows in the logits chunk;
+- at least two rows per request on average (`one_row_per_request`: MTP drafting
+  steps, including CUDA-graph padding, gain nothing from one-row segments);
+- FP16 Q and indexer K with H4/D128, int32 page table, request map and lengths,
+  int32 or int64 positions, positive strides.
+
+The variable is not a torch.compile factor: the route is chosen inside the QSA
+custom op at run time. A device builder writes a fixed `[M,3]` descriptor buffer
+on every call (unused slots hold -1 and exit), so there is no host readback and
+the call can be captured in a CUDA graph. Buffers are allocated per call, as in
+the paged scorer.
+
+The segment limit is four rows. Eight-row segments (two query blocks), query
+blocks loaded inside the tile loop, a data-dependent branch in the loop body, or
+more warps either spill the key tile at 2 warps or change the reduction order.
+
+## CPU checks
+
+From the checkout root, in an environment with PyTorch and Triton:
+
+```bash
+CUDA_VISIBLE_DEVICES='' TRITON_INTERPRET=1 \
+  .venv/bin/python tests/kernels/test_qsa_request_segments_cpu.py
+.venv/bin/python -m pytest tests/kernels/test_qsa_request_segments_dispatch.py
+```
+
+The first runs the builder and scorer in Triton's CPU interpreter against the
+paged kernel of this checkout and a NumPy oracle (scores, visible lengths,
+stable top-512 IDs, untouched padding, descriptors; remapped requests, inactive
+rows, invalid pages, negative positions, ragged and repeated runs, strides, page
+sizes 16/204/416). The interpreter does not establish SM70 rounding, launch
+safety or CUDA-graph behavior. The second checks when the runtime uses the
+scorer.
+
+## SM70 checks
+
+On one SM70 GPU, with the native exact top-k extension installed:
+
+```bash
+.venv/bin/python tests/kernels/qsa_request_segments_gpu.py --out kernel.json
+.venv/bin/python tests/kernels/qsa_request_segments_select_gpu.py --out select.json --time
+```
+
+`qsa_request_segments_gpu.py` runs 36 kernel cases (C1/M8, C4/M16, C4/M20,
+C8/M40 at 32768/131072/262144 tokens, segment limits 1/2/4) against the paged
+scorer, eagerly and through CUDA graphs replayed after the queries, keys, page
+tables, positions, lengths and request runs change. FP32 score words, visible
+lengths and exact top-512 IDs must match; guard storage surrounds the buffers.
+
+`qsa_request_segments_select_gpu.py` runs `qsa_select_paged_tokens` with the
+variable off and on and compares the selected token indices: decode batches,
+one row per request, ragged rows, a prefill chunk next to decode rows (with and
+without the host request split), two logits chunks and CUDA-graph padding rows.
+It also checks which batches use the scorer. `--time` adds the graph replay
+time of the whole selection (scoring, top-512, expansion).
+
+Measured on a V100-PCIE-32GB at 250 W, whole selection per layer, new / paged:
+4 requests x 4 rows 0.39-0.40, 4 x 5 rows 0.47-0.50, 8 x 4 rows 0.32-0.36,
+8 x 5 rows 0.49-0.55 (32K-256K tokens); one row per request 0.97-1.15, which is
+why those batches keep the paged scorer.

tests/kernels/qsa_request_segments_gpu.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Deferred SM70 scorer qualification. --matrix writes the plan without a GPU.
+
+Run only in an approved GPU window. This does not run an inference engine or
+measure performance. Failure is fatal: no numerical tolerance is applied.
+"""
+
+import argparse
+import hashlib
+import importlib
+import json
+from pathlib import Path
+from typing import Any
+
+import torch
+from test_qsa_request_segments_cpu import ROOT, inputs, load_candidate
+
+
+def matrix():
+    return [
+        {
+            "requests": c,
+            "rows_per_request": r,
+            "context": ctx,
+            "page_size": 416,
+            "table_width": 316,
+            "segment_rows": segment,
+        }
+        for c, r in ((1, 8), (4, 4), (4, 5), (8, 5))
+        for ctx in (32768, 131072, 262144)
+        for segment in (1, 2, 4)
+    ]
+
+
+def guarded_workspace(mod, q, columns):
+    # Guard all caller-owned buffers; keep the slices themselves contiguous.
+    stores = []
+    views = []
+    for shape, dtype in (
+        ((len(q), 3), torch.int32),
+        ((len(q),), torch.int32),
+        ((len(q), columns), torch.float32),
+    ):
+        n = 1
+        for dim in shape:
+            n *= dim
+        store = torch.full((n + 128,), 991, dtype=dtype, device=q.device)
+        stores.append(store)
+        views.append(store[64:-64].view(shape))
+    return mod.SegmentWorkspace(*views), stores
+
+
+def run_case(mod, qsa, shape):
+    c, r, ctx = (shape[k] for k in ("requests", "rows_per_request", "context"))
+    columns = shape["page_size"] * shape["table_width"]
+    x = [
+        t.cuda()
+        for t in inputs((r,) * c, page=shape["page_size"], columns=columns, normal=True)
+    ]
+    q, cache, table, requests, positions, lengths = x
+    positions.copy_(torch.arange(ctx - r, ctx, device=q.device).repeat(c))
+    lengths.fill_(ctx)
+    ws, guards = guarded_workspace(mod, q, columns)
+
+    def candidate():
+        return mod.qsa_mqa_request_segments(*x, 4, ws, max_rows=shape["segment_rows"])
+
+    def baseline():
+        return qsa.qsa_mqa_paged(*x, 4, shared_key_scoring=False)
+
+    for _ in range(3):
+        candidate()
+        baseline()
+    torch.accelerator.synchronize()
+    graphs = []
+    outputs = []
+    for fn in (baseline, candidate):
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            outputs.append(fn())
+        graphs.append(graph)
+    checks = []
+    for change in range(4):
+        if change:
+            q.neg_()
+            cache.neg_()
+            table.copy_(table.flip(1))
+            positions.sub_(37)
+            lengths.sub_(37)
+        if change == 1:
+            requests.zero_()  # Change active request and descriptor counts.
+        elif change == 2:
+            requests.copy_(torch.arange(len(q), device=q.device).remainder(c))
+            requests[-1] = -1
+            positions[-1] = -5
+            table[0, 1] = -1  # Invalid visible page, not merely an unused tail.
+        elif change == 3:
+            requests.fill_(c - 1)
+            lengths[-1] = 79
+            positions.fill_(78)
+        ws.scores.fill_(float("nan"))
+        ws.descriptors.fill_(991)
+        for graph in graphs:
+            graph.replay()
+        torch.accelerator.synchronize()
+        fresh = baseline()
+        for name, (score, visible) in zip(("paged_graph", "segment_graph"), outputs):
+            assert torch.equal(visible, fresh[1]), (shape, change, name, "visible")
+            mask = torch.arange(columns, device=q.device)[None, :] < fresh[1][:, None]
+            bits_equal = torch.equal(
+                score.view(torch.int32)[mask], fresh[0].view(torch.int32)[mask]
+            )
+            assert bits_equal, (shape, change, name, "FP32 words")
+            selected = []
+            for scores, vis in ((score, visible), fresh):
+                ids = torch.empty((len(q), 512), dtype=torch.int32, device=q.device)
+                qsa._sm70_qsa_lexicographic_topk_op()(scores, vis, ids, 512)
+                selected.append(ids)
+            assert torch.equal(*selected), (shape, change, name, "top512")
+        for store in guards:
+            assert bool((store[:64] == 991).all())
+            assert bool((store[-64:] == 991).all())
+        checks.append(
+            {
+                "change": change,
+                "scores_bitwise": True,
+                "visible_equal": True,
+                "top512_equal": True,
+                "buffer_guards": True,
+            }
+        )
+    return {"shape": shape, "checks": checks}
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--matrix", type=Path)
+    ap.add_argument("--out", type=Path)
+    args = ap.parse_args()
+    if args.matrix:
+        args.matrix.write_text(json.dumps(matrix(), indent=2) + "\n")
+        return
+    if args.out is None:
+        ap.error("--out is required for GPU qualification")
+    assert torch.cuda.is_available() and torch.cuda.get_device_capability() == (7, 0)
+    import vllm._custom_ops  # noqa: F401
+
+    qsa = importlib.import_module("vllm.models.qwen4_exp.nvidia.ops.qsa")
+    assert qsa.__file__ is not None
+    installed = hashlib.sha256(Path(qsa.__file__).read_bytes()).hexdigest()
+    source = hashlib.sha256(
+        (ROOT / "vllm/models/qwen4_exp/nvidia/ops/qsa.py").read_bytes()
+    ).hexdigest()
+    assert installed == source, "installed paged baseline must match the checkout"
+    mod = load_candidate()
+    result: dict[str, Any] = {
+        "baseline_sha256": source,
+        "candidate_sha256": hashlib.sha256(
+            (
+                ROOT / "vllm/models/qwen4_exp/nvidia/ops/qsa_request_segments.py"
+            ).read_bytes()
+        ).hexdigest(),
+        "engine_validation": "not_run",
+        "performance": "not_measured",
+        "cases": [],
+    }
+    for shape in matrix():
+        result["cases"].append(run_case(mod, qsa, shape))
+        args.out.write_text(json.dumps(result, indent=2) + "\n")
+        print(json.dumps(result["cases"][-1]), flush=True)
+
+
+if __name__ == "__main__":
+    main()