From 8aecc18ab4df89650745437211dc031254a41927 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:21:25 -0700 Subject: [PATCH] Add CPU controls for paged QSA scorer scheduling Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- tools/qwen4_exp/qsa_scorer_study/README.md | 121 +++++++ tools/qwen4_exp/qsa_scorer_study/__init__.py | 3 + tools/qwen4_exp/qsa_scorer_study/fixtures.py | 157 +++++++++ tools/qwen4_exp/qsa_scorer_study/prototype.py | 324 ++++++++++++++++++ tools/qwen4_exp/qsa_scorer_study/run.py | 242 +++++++++++++ tools/qwen4_exp/qsa_scorer_study/test_cpu.py | 302 ++++++++++++++++ 6 files changed, 1149 insertions(+) create mode 100644 tools/qwen4_exp/qsa_scorer_study/README.md create mode 100644 tools/qwen4_exp/qsa_scorer_study/__init__.py create mode 100644 tools/qwen4_exp/qsa_scorer_study/fixtures.py create mode 100644 tools/qwen4_exp/qsa_scorer_study/prototype.py create mode 100644 tools/qwen4_exp/qsa_scorer_study/run.py create mode 100644 tools/qwen4_exp/qsa_scorer_study/test_cpu.py diff --git a/tools/qwen4_exp/qsa_scorer_study/README.md b/tools/qwen4_exp/qsa_scorer_study/README.md new file mode 100644 index 0000000000..f2179ad490 --- /dev/null +++ b/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. diff --git a/tools/qwen4_exp/qsa_scorer_study/__init__.py b/tools/qwen4_exp/qsa_scorer_study/__init__.py new file mode 100644 index 0000000000..acf527ace6 --- /dev/null +++ b/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.""" diff --git a/tools/qwen4_exp/qsa_scorer_study/fixtures.py b/tools/qwen4_exp/qsa_scorer_study/fixtures.py new file mode 100644 index 0000000000..a1ea9740cb --- /dev/null +++ b/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 diff --git a/tools/qwen4_exp/qsa_scorer_study/prototype.py b/tools/qwen4_exp/qsa_scorer_study/prototype.py new file mode 100644 index 0000000000..879452535e --- /dev/null +++ b/tools/qwen4_exp/qsa_scorer_study/prototype.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Independent row reference, grouped tiles, and request-segment prototypes. + +The explicit FP32 CPU reduction is a comparison contract, not an emulator of +Triton or SM70 MMA. No function imports the runtime, CUDA, or a model. +""" + +from dataclasses import dataclass +from math import ceil, isfinite, sqrt +from typing import Any + +import numpy as np +from numpy.typing import NDArray + +Array = NDArray[Any] +POISON = np.uint32(0x7FC00042) + + +@dataclass(frozen=True) +class Inputs: + q: Array + cache: Array + table: Array + requests: Array + positions: Array + lengths: Array + columns: int + compress: int = 4 + divisor: float = sqrt(128) + + def validate(self) -> None: + if self.q.ndim != 3 or self.q.shape[1:] != (4, 128): + raise ValueError("queries must be [rows, 4, 128]") + if self.cache.ndim != 4 or self.cache.shape[2:] != (1, 128): + raise ValueError("cache must be [pages, page_size, 1, 128]") + if self.q.dtype != np.float16 or self.cache.dtype != np.float16: + raise ValueError("queries and keys must use FP16") + if not np.isfinite(self.q).all() or not np.isfinite(self.cache).all(): + raise ValueError("this CPU arithmetic contract requires finite Q/K") + if self.table.ndim != 2 or self.lengths.shape != (len(self.table),): + raise ValueError("one length is required per page-table request") + if any(x.shape != (len(self.q),) for x in (self.requests, self.positions)): + raise ValueError("row metadata must match query rows") + if any( + x.dtype.kind not in "iu" + for x in (self.table, self.requests, self.positions, self.lengths) + ): + raise ValueError("metadata must use integer arrays") + if self.compress <= 0 or not isfinite(self.divisor) or self.divisor <= 0: + raise ValueError("compression and score divisor must be positive") + if self.columns < 0 or self.columns > self.table.shape[1] * self.cache.shape[1]: + raise ValueError("score columns must fit the page-table capacity") + if len(self.q) and ( + not len(self.table) + or not all(self.cache.shape[:2]) + or not self.table.shape[1] + ): + raise ValueError("nonempty queries need nonempty paged storage") + + +@dataclass(frozen=True) +class Segment: + request: int + begin: int + end: int + + +@dataclass +class Result: + scores: Array + visible: Array + writes: Array + logical_key_fetches: int = 0 + launched_programs: int = 0 + active_programs: int = 0 + + +def trunc_div(value: int, divisor: int) -> int: + return value // divisor if value >= 0 else -((-value) // divisor) + + +def visible_lengths(x: Inputs) -> Array: + """Match signed division, including negative padded positions.""" + values = [] + for request, position in zip(x.requests, x.positions): + request = int(request) + length = int(x.lengths[request]) if 0 <= request < len(x.lengths) else 0 + values.append( + min(trunc_div(int(position) + 1, x.compress), trunc_div(length, x.compress)) + ) + return np.asarray(values, dtype=np.int64) + + +def make_segments( + requests: Array, num_requests: int, max_rows: int = 8 +) -> list[Segment]: + """CPU oracle for a future device builder; never a GPU-to-host readback. + + Preserve row order, split long runs, and keep separate runs of one request + separate. At most one descriptor per row is needed, including padding. + """ + if max_rows < 1 or max_rows > 8: + raise ValueError("shared-key segments support one to eight rows") + if num_requests < 0 or requests.ndim != 1 or requests.dtype.kind not in "iu": + raise ValueError("invalid request metadata") + segments = [] + begin = 0 + while begin < len(requests): + request = int(requests[begin]) + if not 0 <= request < num_requests: + begin += 1 + continue + end = begin + 1 + while ( + end < len(requests) + and end - begin < max_rows + and int(requests[end]) == request + ): + end += 1 + segments.append(Segment(request, begin, end)) + begin = end + return segments + + +def segment_descriptors( + requests: Array, num_requests: int, capacity: int | None = None +) -> Array: + """Fixed-capacity [request, begin, end] buffer; -1 marks an unused slot. + + The conservative capacity is the row count, so a future device builder can + change the segment count without a host readback or replay allocation. + A smaller capacity requires a bound from existing CPU scheduler metadata. + """ + capacity = len(requests) if capacity is None else capacity + segments = make_segments(requests, num_requests) + if capacity < len(segments) or capacity < 0: + raise ValueError("descriptor capacity is too small") + out = np.full((capacity, 3), -1, dtype=np.int32) + for slot, segment in enumerate(segments): + out[slot] = (segment.request, segment.begin, segment.end) + return out + + +def _empty(x: Inputs) -> Result: + x.validate() + scores = np.full((len(x.q), x.columns), POISON, dtype=np.uint32).view(np.float32) + return Result(scores, visible_lengths(x), np.zeros(scores.shape, dtype=np.uint8)) + + +def _keys(x: Inputs, request: int, columns: Array) -> tuple[Array, Array]: + """Candidate gather; ordinary indexing preserves all input strides.""" + pages = x.table[request, columns // x.cache.shape[1]].astype(np.int64) + valid = (pages >= 0) & (pages < len(x.cache)) + keys = np.zeros((len(columns), 128), dtype=np.float16) + keys[valid] = x.cache[pages[valid], columns[valid] % x.cache.shape[1], 0] + return keys, valid + + +def _score_queries(queries: Array, keys: Array, divisor: float) -> Array: + # Vectorize independent rows/heads/columns, never the ordered K reduction. + accum = np.zeros((len(queries), 4, len(keys)), dtype=np.float32) + for dim in range(128): + product = ( + queries[:, :, dim].astype(np.float32)[:, :, None] + * keys[:, dim].astype(np.float32)[None, None, :] + ) + accum += product + np.maximum(accum, np.float32(0), out=accum) + return (((accum[:, 0] + accum[:, 1]) + accum[:, 2]) + accum[:, 3]) / np.float32( + divisor + ) + + +def row_reference(x: Inputs) -> Result: + """Independent gather and scalar head/K loops; vectorize columns only.""" + result = _empty(x) + # Intentionally do not call the candidate gather or score helpers. + for row in range(len(x.q)): + request = int(x.requests[row]) + if not 0 <= request < len(x.lengths): + continue + length = min(int(x.positions[row]) + 1, int(x.lengths[request])) + end = max(0, min(x.columns, trunc_div(length, x.compress))) + columns = np.arange(end) + page_ids = x.table[request][columns // x.cache.shape[1]] + valid = (page_ids >= 0) & (page_ids < x.cache.shape[0]) + vectors = np.zeros((end, 128), dtype=np.float32) + vectors[valid] = x.cache[page_ids[valid], columns[valid] % x.cache.shape[1], 0] + heads = [] + for head in range(4): + value = np.zeros(end, dtype=np.float32) + for dim in range(128): + value = value + np.float32(x.q[row, head, dim]) * vectors[:, dim] + heads.append(np.maximum(value, np.float32(0))) + scores = ((heads[0] + heads[1]) + heads[2]) + heads[3] + scores = scores / np.float32(x.divisor) + result.scores[row, :end] = np.where(valid, scores, -np.inf) + result.writes[row, :end] = 1 + result.logical_key_fetches += int(valid.sum()) + return result + + +def tile_grouped(x: Inputs, tiles_per_program: int) -> Result: + """Keep 64-column score tiles; group 1, 4, or 8 tiles per row/program. + + CPU arithmetic across independent tiles is vectorized. This checks the + resulting column ownership, not the GPU compiler's instruction schedule. + """ + if tiles_per_program not in (1, 4, 8): + raise ValueError("expected one, four, or eight tiles per program") + result = _empty(x) + width = 64 * tiles_per_program + grid = ceil(x.columns / width) + result.launched_programs = len(x.q) * grid + for row in range(len(x.q)): + end = max(0, min(x.columns, int(result.visible[row]))) + for program in range(grid): + begin = program * width + if begin >= end: + continue + columns = np.arange(begin, min(begin + width, end)) + keys, valid = _keys(x, int(x.requests[row]), columns) + values = _score_queries(x.q[row : row + 1], keys, x.divisor)[0] + result.scores[row, columns] = np.where(valid, values, -np.inf) + result.writes[row, columns] += 1 + result.logical_key_fetches += int(valid.sum()) + result.active_programs += 1 + return result + + +def request_segmented(x: Inputs) -> Result: + """Share paged keys within each request run, using 32-column native tiles.""" + result = _empty(x) + descriptors = segment_descriptors(x.requests, len(x.lengths)) + grid = ceil(x.columns / 32) + result.launched_programs = len(descriptors) * grid + for request, begin_row, end_row in descriptors: + if request < 0: + continue + segment = Segment(int(request), int(begin_row), int(end_row)) + visible = result.visible[segment.begin : segment.end] + end = max(0, min(x.columns, int(visible.max()))) + for program in range(grid): + begin = program * 32 + if begin >= end: + continue + columns = np.arange(begin, min(begin + 32, end)) + keys, valid = _keys(x, segment.request, columns) + values = _score_queries(x.q[segment.begin : segment.end], keys, x.divisor) + for local, row in enumerate(range(segment.begin, segment.end)): + live = columns < result.visible[row] + result.scores[row, columns[live]] = np.where( + valid, values[local], -np.inf + )[live] + result.writes[row, columns[live]] += 1 + result.logical_key_fetches += int(valid.sum()) + result.active_programs += 1 + return result + + +def select_topk(result: Result, k: int = 512) -> Array: + """Score descending/index ascending selection, final indices ascending. + + Invalid physical pages still have a defined -inf score in the visible + prefix. They are not silently removed from the reference selection. + """ + if k < 1: + raise ValueError("top-k must be positive") + selected = np.full((len(result.visible), k), -1, dtype=np.int32) + for row, visible in enumerate(result.visible): + end = max(0, min(result.scores.shape[1], int(visible))) + indices = np.arange(end, dtype=np.int32) + order = np.lexsort((indices, -result.scores[row, :end])) + ids = np.sort(indices[order[:k]]) + selected[row, : len(ids)] = ids + return selected + + +def assert_equivalent(reference: Result, actual: Result) -> dict[str, int]: + """Check values, writes, untouched padding, and exact top-512 IDs.""" + np.testing.assert_array_equal(actual.visible, reference.visible) + np.testing.assert_array_equal(actual.writes, reference.writes) + np.testing.assert_array_equal( + actual.scores.view(np.uint32), reference.scores.view(np.uint32) + ) + np.testing.assert_array_equal(select_topk(actual), select_topk(reference)) + assert np.all(actual.writes <= 1), "duplicate score write" + return { + "score_words": int(actual.scores.size), + "visible_score_words": int(actual.writes.sum()), + "top512_rows": len(actual.visible), + } + + +def launch_geometry( + requests: Array, visible: Array, columns: int, num_requests: int +) -> dict[str, dict[str, int]]: + """Schedule arithmetic only; counts are not GPU timings or HBM traffic.""" + if columns < 0 or visible.shape != requests.shape: + raise ValueError("invalid geometry") + ends = np.clip(visible, 0, columns).astype(np.int64) + plans = {} + for group in (1, 4, 8): + launched = len(requests) * ceil(columns / (64 * group)) + active = sum(ceil(int(end) / (64 * group)) for end in ends) + plans[f"tile{group}"] = { + "launched": launched, + "active": active, + "empty": launched - active, + } + segments = make_segments(requests, num_requests) + launched = len(requests) * ceil(columns / 32) + active = sum(ceil(int(ends[s.begin : s.end].max()) / 32) for s in segments) + plans["segmented"] = { + "launched": launched, + "active": active, + "empty": launched - active, + "segments": len(segments), + "compact_ideal_launched": len(segments) * ceil(columns / 32), + "descriptor_capacity": len(requests), + "descriptor_bytes": len(requests) * 3 * 4, + } + return plans diff --git a/tools/qwen4_exp/qsa_scorer_study/run.py b/tools/qwen4_exp/qsa_scorer_study/run.py new file mode 100644 index 0000000000..5e4f8c8dff --- /dev/null +++ b/tools/qwen4_exp/qsa_scorer_study/run.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run CPU comparisons, export inputs, and compare deferred GPU captures.""" + +import argparse +import hashlib +import json +from dataclasses import asdict +from functools import partial +from pathlib import Path + +import numpy as np + +from .fixtures import geometry_cases, long_cases, make_inputs, quick_cases +from .prototype import ( + Array, + Inputs, + Result, + assert_equivalent, + launch_geometry, + request_segmented, + row_reference, + select_topk, + tile_grouped, +) + +ARRAY_NAMES = ("q", "cache", "table", "requests", "positions", "lengths") + + +def input_digest(x: Inputs) -> str: + h = hashlib.sha256() + for name in ARRAY_NAMES: + a = getattr(x, name) + h.update(json.dumps([name, a.dtype.str, a.shape, a.strides]).encode()) + h.update(a.tobytes()) + h.update(json.dumps([x.columns, x.compress, x.divisor]).encode()) + return h.hexdigest() + + +def export_inputs(path: Path, x: Inputs, reference: Result) -> None: + layout = {name: list(getattr(x, name).strides) for name in ARRAY_NAMES} + path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + path, + **{name: getattr(x, name) for name in ARRAY_NAMES}, + columns=x.columns, + compress=x.compress, + divisor=x.divisor, + layout=json.dumps(layout), + fixture_sha256=input_digest(x), + cpu_reference_scores=reference.scores, + cpu_reference_visible=reference.visible, + cpu_reference_top512=select_topk(reference), + ) + + +def load_inputs(path: Path) -> Inputs: + """Restore positive-stride layouts; np.savez alone loses those strides.""" + with np.load(path, allow_pickle=False) as data: + layout = json.loads(str(data["layout"])) + arrays = [] + for name in ARRAY_NAMES: + values = data[name] + strides = tuple(layout[name]) + if len(strides) != values.ndim or any(s < 0 for s in strides): + raise ValueError("unsupported fixture strides") + span = values.dtype.itemsize + sum( + max(0, n - 1) * s for n, s in zip(values.shape, strides) + ) + array: Array = np.ndarray( + values.shape, values.dtype, buffer=bytearray(span), strides=strides + ) + array[...] = values + arrays.append(array) + x = Inputs( + q=arrays[0], + cache=arrays[1], + table=arrays[2], + requests=arrays[3], + positions=arrays[4], + lengths=arrays[5], + columns=int(data["columns"]), + compress=int(data["compress"]), + divisor=float(data["divisor"]), + ) + x.validate() + if input_digest(x) != str(data["fixture_sha256"]): + raise ValueError("fixture digest mismatch") + return x + + +def compare_gpu_captures(baseline: Path, candidate: Path) -> dict: + """Compare GPU with GPU; the CPU score oracle is never a GPU bitwise gate. + + Capture producers must provide the same fixture digest and experiment id. + This offline reader validates arrays, not the truth of capture provenance. + """ + with ( + np.load(baseline, allow_pickle=False) as base, + np.load(candidate, allow_pickle=False) as got, + ): + for key in ("fixture_sha256", "experiment_id"): + if not str(base[key]) or str(base[key]) != str(got[key]): + raise ValueError(f"capture {key} mismatch") + scores, actual = base["scores"], got["scores"] + visible, av = base["visible"], got["visible"] + selected, ids = base["top512"], got["top512"] + if scores.dtype != np.float32 or actual.dtype != np.float32: + raise ValueError("captures must contain FP32 scores") + if scores.shape != actual.shape or scores.ndim != 2: + raise ValueError("capture score shape mismatch") + if visible.shape != (len(scores),) or av.shape != visible.shape: + raise ValueError("capture visible shape mismatch") + if visible.dtype.kind not in "iu" or av.dtype.kind not in "iu": + raise ValueError("capture visible lengths must be integers") + if selected.shape != (len(scores), 512) or ids.shape != selected.shape: + raise ValueError("capture selection shape mismatch") + if selected.dtype.kind not in "iu" or ids.dtype.kind not in "iu": + raise ValueError("capture selections must be integers") + mask = np.arange(scores.shape[1])[None, :] < visible[:, None] + compared = int(mask.sum()) + bit_differences = int( + np.count_nonzero( + scores.view(np.uint32)[mask] != actual.view(np.uint32)[mask] + ) + ) + visible_equal = bool(np.array_equal(visible, av)) + selected_equal = bool(np.array_equal(selected, ids)) + return { + "fixture_sha256": str(base["fixture_sha256"]), + "experiment_id": str(base["experiment_id"]), + "compared_score_words": compared, + "different_score_words": bit_differences, + "visible_equal": visible_equal, + "top512_equal": selected_equal, + "qualified": bit_differences == 0 and visible_equal and selected_equal, + "timing_evaluated": False, + } + + +def prepare_matrix() -> list[dict]: + out = [] + for case in geometry_cases(): + plans = launch_geometry( + np.asarray(case["requests"]), + np.asarray(case["visible"]), + case["columns"], + len(case["contexts"]), + ) + out.append({**case, "plans": plans}) + return out + + +def check(args: argparse.Namespace) -> dict: + cases = [] + if args.suite in ("quick", "full"): + cases += quick_cases() + if args.suite in ("long", "full"): + cases += long_cases() + if args.case: + cases = [c for c in cases if c.name == args.case] + if not cases: + raise ValueError("unknown case in selected suite") + results = [] + totals = {"score_words": 0, "visible_score_words": 0, "top512_rows": 0} + for index, case in enumerate(cases): + x = make_inputs(case, seed=1029 + index) + expected = row_reference(x) + variants = {} + for name, implementation in [ + ("tile1", partial(tile_grouped, x, 1)), + ("tile4", partial(tile_grouped, x, 4)), + ("tile8", partial(tile_grouped, x, 8)), + ("segmented", partial(request_segmented, x)), + ]: + actual = implementation() + counts = assert_equivalent(expected, actual) + for key in totals: + totals[key] += counts[key] + variants[name] = { + **counts, + "logical_key_fetches": actual.logical_key_fetches, + "launched_programs": actual.launched_programs, + "active_programs": actual.active_programs, + } + print(f"PASS {case.name} {name}", flush=True) + if args.export: + export_inputs(Path(args.export) / f"{case.name}.npz", x, expected) + roundtrip = load_inputs(Path(args.export) / f"{case.name}.npz") + assert input_digest(roundtrip) == input_digest(x) + results.append( + { + "case": asdict(case), + "input_sha256": input_digest(x), + "reference_logical_key_fetches": expected.logical_key_fetches, + "variants": variants, + } + ) + return { + "backend": "numpy_cpu", + "numeric_cases": len(results), + "candidate_comparisons": len(results) * 4, + "all_equal": True, + "totals": totals, + "cases": results, + "gpu_matrix": prepare_matrix(), + "gpu_numerics_validated": False, + "gpu_performance_validated": False, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + run = sub.add_parser("check", help="compare CPU algorithms") + run.add_argument("--suite", choices=("quick", "long", "full"), default="quick") + run.add_argument("--case") + run.add_argument("--export", help="export NPZ input fixtures with layout metadata") + matrix = sub.add_parser("matrix", help="prepare deferred GPU shapes without Q/K") + compare = sub.add_parser( + "compare-gpu", help="compare baseline/candidate NPZ captures" + ) + compare.add_argument("baseline", type=Path) + compare.add_argument("candidate", type=Path) + for cmd in (run, matrix, compare): + cmd.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + if args.command == "check": + result = check(args) + elif args.command == "matrix": + result = {"gpu_matrix": prepare_matrix(), "gpu_validation": "not_run"} + else: + result = compare_gpu_captures(args.baseline, args.candidate) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(f"Wrote {args.output}", flush=True) + if result.get("qualified") is False: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/qwen4_exp/qsa_scorer_study/test_cpu.py b/tools/qwen4_exp/qsa_scorer_study/test_cpu.py new file mode 100644 index 0000000000..3e67089eb6 --- /dev/null +++ b/tools/qwen4_exp/qsa_scorer_study/test_cpu.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU contract and adversarial checks; synthetic captures do not test CUDA.""" + +import sys +import tempfile +import unittest +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np + +from .fixtures import Case, make_inputs, quick_cases +from .prototype import ( + Inputs, + Result, + assert_equivalent, + launch_geometry, + make_segments, + request_segmented, + row_reference, + segment_descriptors, + select_topk, + tile_grouped, + visible_lengths, +) +from .run import ( + compare_gpu_captures, + export_inputs, + input_digest, + load_inputs, + prepare_matrix, +) + + +class ScorerTest(unittest.TestCase): + def compare_all(self, x): + expected = row_reference(x) + for group in (1, 4, 8): + assert_equivalent(expected, tile_grouped(x, group)) + assert_equivalent(expected, request_segmented(x)) + return expected + + def test_hand_computed_scores(self): + q = np.zeros((2, 4, 128), np.float16) + q[:, 0] = 1 + q[:, 1] = -1 + cache = np.ones((2, 16, 1, 128), np.float16) + cache[1] = 2 + x = Inputs( + q, + cache, + np.array([[0, 1]], np.int32), + np.zeros(2, np.int32), + np.array([127, 79], np.int64), + np.array([128], np.int32), + 32, + divisor=1, + ) + result = self.compare_all(x) + np.testing.assert_array_equal(result.scores[0, :16], 128) + np.testing.assert_array_equal(result.scores[0, 16:], 256) + np.testing.assert_array_equal(result.scores[1, :16], 128) + np.testing.assert_array_equal(result.scores[1, 16:20], 256) + self.assertEqual(result.writes[1, 20:].sum(), 0) + + def test_awkward_layouts_and_request_runs(self): + names = {"strided", "interleaved", "padding_negative", "split_long_run"} + for case in quick_cases(): + if case.name in names: + with self.subTest(case=case.name): + x = make_inputs(case) + if case.strided: + self.assertFalse(x.q.flags.c_contiguous) + self.assertFalse(x.cache.flags.c_contiguous) + self.assertFalse(x.table.flags.c_contiguous) + self.compare_all(x) + + def test_signed_division_and_invalid_requests(self): + x = make_inputs(Case("negative", (5,), (10,), 16)) + x = replace(x, positions=np.array([-2, -3, -4, -5, -6], np.int64)) + np.testing.assert_array_equal(visible_lengths(x), [0, 0, 0, -1, -1]) + self.compare_all(x) + x = replace(x, requests=np.array([-1, 1, 0, 0, 0], np.int32)) + np.testing.assert_array_equal(visible_lengths(x), [0, 0, 0, -1, -1]) + self.compare_all(x) + + def test_head_reduction_tree_is_part_of_the_contract(self): + x = make_inputs(Case("rounding", (2,), (4,), 1)) + x.q.fill(0) + x.cache.fill(0) + x.q[:, 0, 0] = 4096 + x.q[:, 1:, 0] = np.float16(2**-12) + x.cache[:, :, 0, 0] = 4096 + result = self.compare_all( + replace(x, divisor=1, positions=np.full(2, 3, np.int64)) + ) + self.assertEqual(result.scores[0, 0], np.float32(2**24)) + balanced = np.float32(np.float32(2**24 + 1) + np.float32(1 + 1)) + self.assertNotEqual(result.scores[0, 0], balanced) + + def test_empty_query_batch(self): + x = Inputs( + np.empty((0, 4, 128), np.float16), + np.empty((0, 16, 1, 128), np.float16), + np.empty((0, 0), np.int32), + np.empty(0, np.int32), + np.empty(0, np.int64), + np.empty(0, np.int32), + 0, + ) + self.compare_all(x) + + def test_invalid_pages_and_top512_ties(self): + x = make_inputs( + Case("ties", (5, 5), (3100, 3100), 1025, invalid_pages=True, ties=True) + ) + result = self.compare_all(x) + # Independent scalar selection includes -inf invalid-page scores, then + # uses the original logical column as the tie breaker. + for row, visible in enumerate(result.visible): + ids = sorted( + range(int(visible)), key=lambda c: (-float(result.scores[row, c]), c) + )[:512] + np.testing.assert_array_equal(select_topk(result)[row], sorted(ids)) + + def test_each_visible_column_written_once(self): + # Two rows straddle every 64-key and grouped-program tail. A page size + # of 16 also puts page boundaries inside a scoring tile. + for columns in (1, 63, 64, 65, 255, 256, 257, 511, 512, 513): + x = make_inputs( + Case("tail", (2,), (4 * (columns + 3),), columns, page_size=16) + ) + expected = self.compare_all(x) + np.testing.assert_array_equal(expected.writes, 1) + + def test_segments_partition_valid_rows(self): + rng = np.random.default_rng(41) + for rows in (0, 1, 5, 8, 9, 17, 20, 35, 40, 64): + for _ in range(30): + requests = rng.integers(-1, 5, rows, dtype=np.int32) + counts = np.zeros(rows, np.int32) + segments = make_segments(requests, 4) + self.assertLessEqual(len(segments), rows) + for segment in segments: + self.assertTrue(1 <= segment.end - segment.begin <= 8) + np.testing.assert_array_equal( + requests[segment.begin : segment.end], segment.request + ) + counts[segment.begin : segment.end] += 1 + np.testing.assert_array_equal( + counts, ((requests >= 0) & (requests < 4)).astype(np.int32) + ) + # This long run cannot be represented as one native M>8 call. + self.assertEqual( + [(s.begin, s.end) for s in make_segments(np.zeros(20, np.int32), 1)], + [(0, 8), (8, 16), (16, 20)], + ) + + def test_fixed_capacity_descriptors_and_metadata_changes(self): + requests = np.repeat(np.arange(4, dtype=np.int32), 5) + descriptors = segment_descriptors(requests, 4) + self.assertEqual(descriptors.shape, (20, 3)) + self.assertEqual(np.count_nonzero(descriptors[:, 0] >= 0), 4) + requests[::2] = -1 + changed = segment_descriptors(requests, 4) + self.assertEqual(changed.shape, descriptors.shape) + self.assertGreater(np.count_nonzero(changed[:, 0] >= 0), 4) + with self.assertRaises(ValueError): + segment_descriptors(requests, 4, capacity=4) + + def test_request_isolation_and_wrong_table_negative_control(self): + x = make_inputs(Case("isolation", (5, 5), (3100, 3100), 1025)) + expected = row_reference(x) + wrong = request_segmented(replace(x, table=x.table[::-1])) + with self.assertRaises(AssertionError): + assert_equivalent(expected, wrong) + cache = x.cache.copy() + pages = x.table[1][x.table[1] >= 0] + cache[pages] *= np.float16(2) + changed = request_segmented(replace(x, cache=cache)) + np.testing.assert_array_equal( + changed.scores[:5].view(np.uint32), expected.scores[:5].view(np.uint32) + ) + self.assertTrue( + np.any( + changed.scores[5:].view(np.uint32) + != expected.scores[5:].view(np.uint32) + ) + ) + + def test_comparison_detects_corruption_and_padding_writes(self): + x = make_inputs(Case("corrupt", (2,), (100,), 65)) + expected = row_reference(x) + for column in (0, 64): + wrong = Result( + expected.scores.copy(), expected.visible.copy(), expected.writes.copy() + ) + wrong.scores.view(np.uint32)[0, column] ^= np.uint32(1) + with self.assertRaises(AssertionError): + assert_equivalent(expected, wrong) + wrong = Result( + expected.scores.copy(), expected.visible.copy(), expected.writes.copy() + ) + wrong.writes[0, 0] = 2 + with self.assertRaises(AssertionError): + assert_equivalent(expected, wrong) + + def test_bad_input_contracts(self): + x = make_inputs(Case("bad", (2,), (100,), 65)) + for changes in ( + {"compress": 0}, + {"divisor": 0}, + {"divisor": float("nan")}, + {"columns": -1}, + {"columns": 999999}, + {"q": x.q.astype(np.float32)}, + {"requests": x.requests.astype(float)}, + ): + with self.subTest(changes=list(changes)), self.assertRaises(ValueError): + row_reference(replace(x, **changes)) + for rows in (0, 9): + with self.assertRaises(ValueError): + make_segments(x.requests, 1, rows) + + def test_geometry_and_observed_grid_envelope(self): + # Use a closed-form independent count, including zero visible rows. + for columns in (0, 1, 63, 64, 65, 8192, 131456): + for visible in (-6, 0, 1, 63, 64, 65, 8293, columns, columns + 17): + plans = launch_geometry( + np.zeros(16, np.int32), np.full(16, visible), columns, 1 + ) + end = min(columns, max(0, visible)) + for group in (1, 4, 8): + width = 64 * group + p = plans[f"tile{group}"] + self.assertEqual( + p["launched"], 16 * ((columns + width - 1) // width) + ) + self.assertEqual(p["active"], 16 * ((end + width - 1) // width)) + self.assertEqual(p["empty"], p["launched"] - p["active"]) + example = launch_geometry( + np.repeat(np.arange(4), 4), np.full(16, 8293), 131456, 4 + ) + self.assertEqual( + example["tile1"], {"launched": 32864, "active": 2080, "empty": 30784} + ) + self.assertEqual(example["tile4"]["launched"], 8224) + self.assertEqual(example["tile8"]["launched"], 4112) + + def test_export_restores_strides_and_detects_changed_input(self): + x = make_inputs(Case("export", (4, 5), (2500, 3200), 1031, strided=True)) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "input.npz" + export_inputs(path, x, row_reference(x)) + got = load_inputs(path) + self.assertEqual(input_digest(x), input_digest(got)) + for name in ("q", "cache", "table"): + self.assertEqual(getattr(x, name).strides, getattr(got, name).strides) + got.q[0, 0, 0] += np.float16(1) + self.assertNotEqual(input_digest(x), input_digest(got)) + + def test_capture_comparator_with_synthetic_data(self): + x = make_inputs(Case("capture", (2,), (100,), 65)) + expected = row_reference(x) + fields: dict[str, Any] = { + "fixture_sha256": "a" * 64, + "experiment_id": "synthetic-reader-test", + "scores": expected.scores.copy(), + "visible": expected.visible, + "top512": select_topk(expected), + } + with tempfile.TemporaryDirectory() as directory: + baseline, candidate = ( + Path(directory) / "base.npz", + Path(directory) / "got.npz", + ) + np.savez(baseline, **fields) + fields["scores"][0, 64] = ( + 123 # Undefined padding is ignored for GPU captures. + ) + np.savez(candidate, **fields) + self.assertTrue(compare_gpu_captures(baseline, candidate)["qualified"]) + fields["scores"].view(np.uint32)[0, 0] ^= np.uint32(1) + np.savez(candidate, **fields) + self.assertFalse(compare_gpu_captures(baseline, candidate)["qualified"]) + fields["fixture_sha256"] = "b" * 64 + np.savez(candidate, **fields) + with self.assertRaises(ValueError): + compare_gpu_captures(baseline, candidate) + + def test_matrix_is_deferred_and_requires_no_runtime(self): + matrix = prepare_matrix() + self.assertEqual(len(matrix), 12) + self.assertTrue(all(c["gpu_validation"] == "not_run" for c in matrix)) + for name in ("torch", "vllm", "triton"): + self.assertNotIn(name, sys.modules) + + +if __name__ == "__main__": + unittest.main()