Reduce direct-token latency with candidate-only projection

[AngeLouCN]

Summary

Direct-token inference projects the full vocabulary even though the response uses only the candidate labels. Select the requested LM-head rows before projection and cache the candidate token indices on the head's device, reducing the readout to the number of available choices.

This applies to lm_token evaluation. Training continues to produce full-vocabulary logits; evaluation retains temperature calibration and the existing device/dtype handling. The derived index buffer is rebuilt from checkpoint metadata and is not persisted in checkpoint weights.

Accuracy and latency

Ours means candidate-only projection plus cached candidate indices. Parentheses show the change relative to the named baseline: accuracy uses percentage points (pp), latency uses percent, and negative latency deltas mean faster. Deltas are calculated from unrounded measurements.

Model	Configuration	Delta baseline	Transfer acc (Δ pp)	JevBench acc (Δ pp)	Median ms (Δ %)	Mean ms (Δ %)
4B	Eager	—	78.39%	80.52%	78.096	84.918
4B	Eager + Ours	Eager	78.39% (0.00)	80.52% (0.00)	77.768 (-0.42%)	84.604 (-0.37%)
4B	CUDA Graph	Eager	78.20% (-0.19)	81.39% (+0.87)	21.669 (-72.25%)	29.703 (-65.02%)
4B	CUDA Graph + Ours	CUDA Graph	78.20% (0.00)	81.39% (0.00)	21.318 (-1.62%)	29.319 (-1.29%)
27B	Eager	—	77.72%	84.42%	93.563	98.739
27B	Eager + Ours	Eager	77.72% (0.00)	84.42% (0.00)	92.816 (-0.80%)	97.985 (-0.76%)
27B	CUDA Graph	Eager	77.63% (-0.10)	84.42% (0.00)	33.577 (-64.11%)	42.346 (-57.11%)
27B	CUDA Graph + Ours	CUDA Graph	77.63% (0.00)	84.42% (0.00)	32.878 (-2.08%)	41.752 (-1.40%)

Adding Ours changed no answers in either execution mode for these two measured configurations. CUDA Graph itself has small accuracy differences from Eager, shown separately above.

Measurement scope

• Both experiments used H200, BF16, SDPA, PyTorch 2.8.0 + CUDA 12.8 and Transformers 5.17.0. Each configuration measured the same 1,495 requests. Accuracy covers 1,046 clean knowable Transfer-v9 questions and 231 JevBench questions; latency is the median/mean across all 1,495 requests.
• 4B: the trained SimpleJev/JevAny-Qwen3.5-4B-Direct-Token-LoRA checkpoint, with paired baseline/optimized readouts on the same model. Native graphs were capped at 2,048 tokens: 1,459 requests replayed and 36 used Eager fallback per configuration. The table uses evaluation-traversal timing, which includes first encounters with new shapes and audit-logit transfer.
• 27B: original Qwen/Qwen3.8-27B weights, revision 1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0, using zero-shot candidate scoring in a separate benchmark harness. This is not the trained 27B Pointer release. Both Graph configurations shared the captured text backbone, covered inputs through 4,096 tokens, and replayed all 1,495 requests without fallback. All 665 input-length/candidate-count shapes were warmed before timing; configuration order rotated by request, and audit-logit transfer was excluded. The separate fixed-input panel also measured 8 rounds × 10 calls per configuration and input.
• These timings include input processing, the full model forward, readout and probability return, but exclude model loading and graph capture. Because training and timing protocols differ, compare optimization deltas within each model, rather than ranking 4B against 27B by absolute values.
• Candidate projection is algebraically equivalent to selecting the same rows after a full projection, but floating-point logits, probabilities and API confidence can change. The observed zero accuracy delta is limited to these evaluations; it is not a bitwise-equivalence guarantee or a claim that every workload improves.

Validation

• Independent branch based directly on main; only jevany/model.py, tests/test_decision_modes.py and tests/test_backbones.py change.
• CPU regression: python -m pytest tests -m 'not server' -q -ra --strict-config --strict-markers — 431 passed, 37 skipped, 7 deselected.
• python -m build --no-isolation — source distribution and wheel both built successfully.
• The H200 figures above were collected before this PR using the same readout optimization applied at runtime. They are not new GPU measurements of this branch; the 27B benchmark harness does not change the product's CUDA Graph configuration.

Base: main at 625aed9.