[Kernel][WIP] Optimize exact single-request Qwen3.8 QSA and router

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(),
+            }
+        )
+    )

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 <torch/all.h>
+#include <torch/library.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+
+#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.size(0), vllm::qsa::kLexicographicTopKThreads, 0, stream>>>(
+            logits.data_ptr<float>(), lengths.data_ptr<int32_t>(),
+            output.data_ptr<int32_t>(), logits.size(0), logits.size(1),
+            logits.stride(0));
+  } else {
+    vllm::qsa::launch_qsa_lexicographic_topk<512>(
+        logits.data_ptr<float>(), lengths.data_ptr<int32_t>(),
+        output.data_ptr<int32_t>(), 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);
+}

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()