[cameronr]

Problem and result

The MTP drafter's 2-bit draft-head copy (its fast top-1 argmax) capped at 49,152 rows and threw draft_proxy_top1: too many rows above that, but the drafter enabled the copy on head type and width alone - it never checked the row count. Any model whose draft_vocab.bin is larger than 49,152 tokens - upstream's default 106,299-token CJK subset - threw on every draft step, so the server retried each token.

After: the copy scores any draft vocab. The parity selftest passes at 40,525 / 49,184 / 106,299 rows for Q6_K and Q8_0, and the draft-head argmax at 106,299 rows is ~2.4-3.1x faster than the exact head.

Related discussion

none. This is fork-only: the 2-bit draft proxy is not in upstream, so the bug does not exist there (upstream always uses the exact head, which has no row cap).

Approach

The rescore kernel swept its candidates over a fixed RS_SLOTS(12) x 512 grid = 6,144 candidates = 49,152 rows. The sweep is now grid-stride: start = blockIdx.x*RS_W + warp, step = gridDim.x*RS_W. That is the exact slot sequence the fixed sweep used, so the rows that once fit rescore in the same order and the sum is bitwise identical; the loop only keeps going past the old cap. cand[] is small and L2-resident, so dropping the register preload costs nothing (the row reads dominate, and the window test skips all but the rows near the best regardless of size).

draft_proxy_fits() now only rejects a count the kernel's int candidate index cannot represent; bind_head keeps it as the fallback gate, and STRATA_DRAFT_PROXY=0 stays as the manual off switch.

Supporting changes so the selftest can actually run: STRATA_BUILD_TESTS=ON failed to configure in this fork because tests/ is Python-only and two bench/micro sources are absent - both are now guarded with EXISTS, matching the pattern the file already uses for other optional sources. Added a one-command selftest (selftest-draft-proxy.sh + docker/Dockerfile.selftest), since docker build has no GPU and the runtime image is built STRATA_BUILD_TESTS=OFF.

Validation

Passed:

• sudo ./selftest-draft-proxy.sh @ 12089e3, CUDA 13.0, RTX 4090. All sizes ok; probability within 0.02-0.03 of exact; 0 near-ties picked otherwise.
  • 40,525 rows (control, the shipped subset): copy 48 us / exact 110-137 us, 2.3-2.9x
  • 49,184 rows (one tile past the old cap, the shape that used to throw): copy 56 us / exact 129-162 us, 2.3-2.9x
  • 106,299 rows (CJK subset): copy 108-110 us / exact 261-332 us, 2.4-3.1x

Failed:

• none observed.

Not run:

• End-to-end throughput effect not isolated. tools/strata_bench.py A/B (same image, STRATA_DRAFT_PROXY 0 vs 1, 106k CJK subset) showed no significant difference: the suite's noise floor (~2%, measured via prefill, which the proxy cannot affect) exceeds the expected ~1-2% decode effect, and decode had ±20% run-to-run variance at n=3. The parity timing above is the authoritative measurement of the draft-head cost.

Limits

• The end-to-end gain is ~0-2% and below this suite's noise floor; the primary value is correctness (the CJK draft vocab works instead of throwing / falling back), with the ~2.5x draft-head speedup as a bonus.
• The copy adds ~85 MiB VRAM at 106,299 rows, on top of the exact head it still keeps for the near-best rescore.’

[Copilot]

🟡 Changes recommended

The eligibility gate accepts row counts that exceed the scoring kernel’s signed integer range.

1 open issue

• Reject row counts above INT32_MAX before tile arithmetic · New

What changed in this PR

This PR removes the draft proxy’s fixed 49,152-row rescoring cap so the MTP drafter can handle larger draft vocabularies, including the 106,299-token CJK subset.

Changes:

• Replaces the fixed rescoring sweep with grid-stride iteration and adds an eligibility gate.
• Expands parity checks and timing measurements across three vocabulary sizes for Q6_K and Q8_0.
• Adds a Docker-based selftest runner and guards optional CMake test sources.

File	Description
src/​kernels/​draft_proxy_parity.cpp	Tests and times three vocabulary sizes.
src/​kernels/​cuda/​draft_proxy.cu	Adds grid-stride rescoring and row-count eligibility checking.
src/​core/​mtp.cpp	Gates proxy use and logs exact-head fallback.
selftest-draft-proxy.sh	Builds and runs the GPU selftest in Docker.
include/​strata/​kernels/​draft_proxy.hpp	Declares and documents the eligibility API.
docker/​Dockerfile.selftest	Provides the selftest development environment.
CMakeLists.txt	Guards optional test sources and the C++ test directory.

_{Review effort: Balanced}

[eddoursul]

Thanks, merged. This branch's own draft subset is 40,525 tokens (its data/draft_vocab.bin), so it never reached the
cap, but setup's Cyrillic subset (58,963 tokens) did: here a generate run stopped with "too many rows" before its
first token. With your change both models draft over it, and over upstream's CJK subset too.

What I checked on this machine (RTX 3090 + RTX 5070 Ti, Windows):

• The selftest passes at all three sizes on both cards.
• Window hashes over 1,200 tokens of fixed text are bitwise the same as before on both models.
• Your loop is faster than the preload it replaces: at 40,525 rows on the 3090, 50.9 against 52.4 us (Q6_K) and 50.2
  against 52.0 us (Q8_0), medians of 6 alternating runs of your test. Loading 12 candidates ahead per pass, in the same
  slot order, was 1-2 us slower at all three sizes.
• End to end with the CJK subset, on fixed text: the copy drafts in 1.32 ms a round against the exact head's 1.71,
  which makes generation 3.0-4.1% faster on thinking text and 2.0% faster on code.

One follow-up (a85301c):

• draft_proxy_fits is gone. As the review found, it let counts past INT32_MAX through to the kernels' int row ids,
  and no draft head comes near that limit (24 GB holds 9-12 M rows of Q8_0 or Q6_K). draft_proxy_top1 now refuses
  n_rows outside 1..INT32_MAX itself. The start message for a head type the copy doesn't take stays.
• The comments describe the current sweep and test sizes rather than the old cap. CMakeLists.txt's comments say what
  the published source leaves out: the Catch2 suite's tests/CMakeLists.txt (tests/ itself has C++ tests too).

STRATA_BUILD_TESTS=ON now configures on Windows as well. I didn't run the Docker selftest here, but I read it
against the Dockerfile's build stage.