Add CPU controls for paged QSA scorer scheduling

tools/qwen4_exp/qsa_scorer_study/README.md

@@ -0,0 +1,121 @@
+# Paged QSA scorer CPU study
+
+This standalone NumPy study compares two scheduling changes for small
+multi-request FP16 QSA verification batches. It prepares input fixtures and
+comparison gates before GPU implementation. It does not change runtime dispatch
+or import vLLM, PyTorch, Triton, CUDA, or model weights.
+
+## Run the CPU checks
+
+Use a project virtual environment. From the repository root:
+
+```bash
+uv venv --python 3.12 .venv
+uv pip install --python .venv/bin/python numpy==2.2.6
+.venv/bin/python -m unittest tools.qwen4_exp.qsa_scorer_study.test_cpu -v
+.venv/bin/python -m tools.qwen4_exp.qsa_scorer_study.run check \
+  --suite quick --export qsa-results/inputs --output qsa-results/quick.json
+.venv/bin/python -m tools.qwen4_exp.qsa_scorer_study.run check \
+  --suite long --output qsa-results/long.json
+.venv/bin/python -m tools.qwen4_exp.qsa_scorer_study.run matrix \
+  --output qsa-results/gpu-matrix.json
+```
+
+`quick` covers ragged and repeated request runs, padding, negative positions,
+invalid pages, strided tensors, ties, truncated widths, empty visibility, page
+tails, and 16/20/35/40/64-row batches. `long` adds C4 at 16/20 rows and C8 at
+40 rows, each at 32768/131072/262144 token contexts. `full` combines them.
+`--case NAME` selects one case. CPU wall time is not a GPU performance metric.
+
+## Algorithms and comparison contract
+
+- `row_reference` gathers each query row independently and computes every head
+  with explicit FP32 K-dimension accumulation. Head reduction is left-to-right.
+- `tile_grouped` maps one, four, or eight existing 64-key tiles to one program.
+  Each row retains its own request, visible length, and output coordinates.
+  Grouping changes ownership and query reuse, not the CPU arithmetic contract.
+- `request_segmented` splits contiguous runs into at most eight rows, reads keys
+  once per segment, and writes scores back to their original query rows. Each
+  row keeps its own causal limit. Repeated runs remain separate.
+
+The reference has independent gather and score code. All variants must match
+FP32 score words, visible lengths, one-write coverage, untouched padding, and
+exact top-512 indices. Selection uses score descending/index ascending ties and
+returns selected indices in ascending order, with -1 padding for short rows.
+Invalid physical pages have -inf scores within the visible prefix.
+
+The initial arithmetic scope is finite FP16 Q/K, four heads, dimension 128,
+positive finite divisor, integer metadata, and score width within page-table
+capacity. Arbitrary positive input strides and negative padded positions are
+included. CUDA output beyond the visible prefix is unspecified; CPU poison
+checks are stricter ownership tests. Zero-column device visibility is also
+unspecified and is not a GPU comparison case.
+
+These are CPU scheduling prototypes, not bitwise emulators of SM70 MMA or
+Triton. In particular, the existing single-request shared-key test permits
+`rtol=2e-6, atol=2e-6` for score comparison. Passing that test does not satisfy
+the proposed multi-request bitwise gate. Keep K/head reduction and MMA choices
+as separate numerical questions during GPU implementation.
+
+## Launch and memory accounting
+
+The shipped paged path uses one tile/program for M<=32 and eight for M>32.
+Thus tile8 is already the paged baseline at M40; it is not a new optimization
+there. Single-request M2..8 can instead use the existing shared-key operator
+when enabled. Record the actual route before comparing a candidate.
+
+For a synthetic 16-row, 131456-column envelope with visibility <=8293, tile1
+launches 32864 programs, of which at least 30784 are empty. Tile4 launches 8224;
+tile8 launches 4112. These counts prove neither speedup nor reduced HBM traffic.
+131456 is a rounded grid envelope, not a universal deployment cache capacity.
+
+The segmented prototype allocates a fixed `[M,3]` int32 descriptor buffer and
+models a grid with M slots. Empty descriptor slots still launch programs. The
+geometry report also shows the smaller *compact ideal* count. Achieving that
+count requires a safe segment-capacity bound from existing CPU scheduler
+metadata; reading a segment count back from the GPU is not an acceptable plan.
+Include descriptor generation, transfer if any, empty slots, and extra launches
+in the full scorer cost. Descriptor storage is 12*M bytes, before count/scan
+workspace. Logical key-fetch counts are not measured memory transactions.
+
+## Deferred GPU validation
+
+The matrix contains 12 scorer-only shapes: C1/M8, C4/M16, C4/M20, C8/M40 at
+32K/128K/256K. Long-context synthetic scorer inputs do not establish that an
+inference engine can admit the corresponding number of requests. GPU candidate
+kernels and runtime integration are not supplied by this CPU study.
+
+1. Fix one source revision, FP16 inputs, compression, page layout, GPU power,
+   clocks and capture mode. Preserve shipped dispatch as one control, and call
+   the paged kernel with tile1/tile4/tile8 as explicit controls. Hold selection
+   unchanged while evaluating the scorer.
+2. Capture the paged baseline and each candidate on the same inputs. Require
+   bitwise visible-prefix scores, visible lengths, and top-512 IDs. Check page
+   isolation, invalid pages, ragged runs, ties and partial tiles first. Any
+   numerical relaxation needs a separately approved error and quality contract.
+3. Replay one fixed-capacity CUDA graph after changing queries, cache contents,
+   page tables, request IDs, positions, lengths and active descriptor count.
+   Check for new allocations, host readback, stale descriptors and out-of-bounds
+   writes. Include the metadata builder in graph replay and scorer timing.
+4. Measure scorer-only and scorer+selection, including workspace and all extra
+   launches. Alternate baseline/candidate measurements in the same window and
+   report distributions. Do not derive speed from CPU time or kernel counts.
+5. Only then compare model decode on valid C4/C8 windows, offload and long
+   contexts. Verify actual concurrency from per-step scheduler markers. A
+   blocked or invalid C8 long-context case remains unavailable, not a pass.
+
+Exported NPZ inputs preserve content and stride metadata; `load_inputs` restores
+the views and verifies their digest. `cpu_reference_*` arrays are debugging
+oracles, not expected GPU bit patterns. For two later GPU captures, write
+`scores` (FP32), `visible` (integer), `top512` (integer `[M,512]`), and identical
+nonempty `fixture_sha256` and `experiment_id` strings, then run:
+
+```bash
+.venv/bin/python -m tools.qwen4_exp.qsa_scorer_study.run compare-gpu \
+  qsa-results/baseline.npz qsa-results/candidate.npz \
+  --output qsa-results/comparison.json
+```
+
+The reader compares GPU with GPU, ignores undefined score padding, and exits
+nonzero on a difference. It cannot verify capture provenance or GPU timing.
+Its unit tests use explicitly synthetic captures; they do not run CUDA.

tools/qwen4_exp/qsa_scorer_study/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU-only experiments for paged QSA scorer scheduling."""

tools/qwen4_exp/qsa_scorer_study/fixtures.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Deterministic synthetic inputs and a deferred GPU comparison matrix."""
+
+from dataclasses import dataclass
+from math import ceil
+
+import numpy as np
+
+from .prototype import Inputs, trunc_div
+
+
+@dataclass(frozen=True)
+class Case:
+    name: str
+    widths: tuple[int, ...]
+    contexts: tuple[int, ...]
+    columns: int
+    page_size: int = 204
+    strided: bool = False
+    invalid_pages: bool = False
+    interleaved: bool = False
+    padding: bool = False
+    ties: bool = False
+
+
+def make_inputs(case: Case, seed: int = 1029) -> Inputs:
+    if len(case.widths) != len(case.contexts) or not case.widths:
+        raise ValueError("one context is required per request")
+    if min(case.widths) < 1 or min(case.contexts) < 0:
+        raise ValueError("invalid case shape")
+    rng = np.random.default_rng(seed)
+    rows, requests = sum(case.widths), len(case.widths)
+    table_width = max(1, ceil(case.columns / case.page_size))
+    needed = [
+        max(1, min(table_width, ceil(length / (4 * case.page_size))))
+        for length in case.contexts
+    ]
+    pages = sum(needed)
+    if case.strided:
+        q = rng.standard_normal((rows + 1, 8, 256)).astype(np.float16)[1:, ::2, ::2]
+        cache = rng.standard_normal((pages, case.page_size * 2, 1, 256)).astype(
+            np.float16
+        )[:, ::2, :, ::2]
+        table = np.full((requests * 2, table_width * 2), -1, np.int32)[::2, 1::2]
+    else:
+        q = rng.standard_normal((rows, 4, 128)).astype(np.float16)
+        cache = rng.standard_normal((pages, case.page_size, 1, 128)).astype(np.float16)
+        table = np.full((requests, table_width), -1, np.int32)
+    page_ids = rng.permutation(pages).astype(np.int32)
+    offset = 0
+    for request, n in enumerate(needed):
+        table[request, :n] = page_ids[offset : offset + n]
+        offset += n
+    if case.invalid_pages:
+        table[:, ::7] = -1
+        table[:, 3::11] = pages
+    row_requests = np.repeat(np.arange(requests, dtype=np.int32), case.widths)
+    positions = np.concatenate(
+        [
+            np.arange(length - width, length, dtype=np.int64)
+            for width, length in zip(case.widths, case.contexts)
+        ]
+    )
+    if case.interleaved:
+        order = np.argsort(
+            np.concatenate([np.arange(n) for n in case.widths]), kind="stable"
+        )
+        row_requests, positions, q = row_requests[order], positions[order], q[order]
+    if case.padding:
+        row_requests[::7] = -1
+        row_requests[3::11] = requests
+        positions[::5] = -6
+    if case.ties:
+        q.fill(0)
+    return Inputs(
+        q,
+        cache,
+        table,
+        row_requests,
+        positions,
+        np.asarray(case.contexts, np.int32),
+        case.columns,
+    )
+
+
+def quick_cases() -> list[Case]:
+    return [
+        Case("single8", (8,), (3073,), 1025),
+        Case("c4_rows16", (4,) * 4, (4097,) * 4, 1100),
+        Case("c4_rows20", (5,) * 4, (4097,) * 4, 1100),
+        Case("c8_rows40", (5,) * 8, (3073,) * 8, 1025),
+        Case("c7_rows35", (5,) * 7, (3073,) * 7, 1025),
+        Case("c8_rows64_stress", (8,) * 8, (2053,) * 8, 517),
+        Case("ragged", (8, 5, 2), (1600, 2400, 3200), 1025),
+        Case("split_long_run", (17, 3), (2600, 3100), 1025),
+        Case("interleaved", (5, 4, 3), (2500, 3000, 3500), 1025, interleaved=True),
+        Case("strided", (4, 5), (2500, 3200), 1031, strided=True),
+        Case("invalid_pages", (5, 3), (2800, 3100), 1025, invalid_pages=True),
+        Case("padding_negative", (8, 5, 2), (2500, 3000, 3500), 1025, padding=True),
+        Case("all_ties", (5, 5), (3100, 3100), 1025, ties=True),
+        Case("width_truncated", (5, 4), (2500, 3100), 63),
+        Case("zero_visible", (4, 4), (0, 0), 65),
+        Case("zero_columns", (3,), (10,), 0),
+        Case("page16_tail", (5, 5), (2053, 2057), 515, page_size=16),
+    ]
+
+
+def long_cases() -> list[Case]:
+    # Full-width allocations are synthetic. 131456 is a rounded grid envelope,
+    # not a claim that every deployment has this exact cache capacity.
+    return [
+        Case(f"{name}_{context}", (width,) * n, (context,) * n, 131456)
+        for name, n, width in [
+            ("c4_rows16", 4, 4),
+            ("c4_rows20", 4, 5),
+            ("c8_rows40", 8, 5),
+        ]
+        for context in (32768, 131072, 262144)
+    ]
+
+
+def geometry_cases() -> list[dict]:
+    """No cache allocation is needed to prepare the full GPU shape matrix."""
+    cases = [
+        Case(f"{name}_{length}", (width,) * n, (length,) * n, 131456)
+        for name, n, width in [
+            ("single8", 1, 8),
+            ("c4_rows16", 4, 4),
+            ("c4_rows20", 4, 5),
+            ("c8_rows40", 8, 5),
+        ]
+        for length in (32768, 131072, 262144)
+    ]
+    out = []
+    for case in cases:
+        requests = [r for r, n in enumerate(case.widths) for _ in range(n)]
+        visible = [
+            trunc_div(position + 1, 4)
+            for n, length in zip(case.widths, case.contexts)
+            for position in range(length - n, length)
+        ]
+        out.append(
+            {
+                "case": case.name,
+                "rows_per_request": list(case.widths),
+                "contexts": list(case.contexts),
+                "columns": case.columns,
+                "requests": requests,
+                "visible": visible,
+                "page_size": case.page_size,
+                "runtime_paged_tiles_per_program": 1 if sum(case.widths) <= 32 else 8,
+                "existing_single_request_shared_key_shape": len(case.widths) == 1,
+                "gpu_validation": "not_run",
+            }
+        )
+    return out