## Summary
The paged text path projects the output head over **every** packed row,
but sampling reads only one row per
prefill segment (`_build_paged_row_targets` takes `cu_seqlens[num_decode
+ j + 1] - 1`). For a 2048-token prompt
that is 2047 wasted rows of `[hidden_dim × vocab_size]` matmul and a
`[2048 × 151936]` tensor materialized to
sample one row.
`ModelAdapter.target_forward` gains an optional `logits_indices`:
prefill steps gather the sampled rows and run
the head on those alone; pure-decode steps pass `None` and are
byte-for-byte unchanged. `_PagedForwardState`
carries a new `logits_cu_seqlens`, since logits are no longer
row-aligned with the input — hidden-state consumers
keep `cu_seqlens` for the spec drafter. Implements #525.
**A model's head is not always reproducible outside its `__call__`, and
getting that wrong yields silently wrong
logits, not a crash** (Cohere's `logit_scale`, Granite's
`logits_scaling`). Rather than an `mlx_lm` allowlist,
`supports_selective_logits` probes once at load and demands bit-exact
equality with the model's own output;
mismatches keep full logits and log why.
## Benchmarks
`vllm bench serve --model Qwen/Qwen3-0.6B --dataset-name random
--random-input-len 2048 --random-output-len 16
--num-prompts 40`, M3 Pro, vllm 0.26.0. 5 iterations per arm,
**interleaved** with a restart per arm, medians;
40/40 requests OK across all 10 runs.
| metric | before | after | delta |
|---|---|---|---|
| Output throughput | 20.26 tok/s | 23.46 tok/s | **+15.8%** |
| Mean TTFT | 17644.89 ms | 15375.80 ms | **−12.9%** |
| Mean TPOT | 571.96 ms | 485.36 ms | **−15.1%** |
| Duration | 31.58 s | 27.28 s | **−13.6%** |
CV ≤ 3.2% and the arms' ranges are **disjoint** on all four. The machine
was not idle — hence the interleaving,
so load drift could not land in one arm — so absolute numbers are low
for this hardware and the ratio is the
result. Isolated head projection at 2048 rows: 107.00 → 2.51 ms, 593.5
MiB avoided.
## Testing
```
pytest tests/test_model_adapter.py tests/test_v1_model_runner_generate.py tests/test_grammar_bitmask.py tests/test_decode_pipeline.py
224 passed
```
Covers index selection, `logits_cu_seqlens` compaction, selective and
full projections agreeing on sampled rows,
the probe rejecting a model that scales its head output, and bitmask
application under the compacted layout.
Prompt logprobs are unimplemented here (`prompt_logprobs_dict={}` is
hardcoded), so dropping non-last prefill
rows regresses nothing. Also updates #588's new `_paged_state` helper
and `fake_target_forward` stub (6 lines) to
match the file's eight others.
---
AI assistance was used in developing this change. The problem analysis,
design, benchmarks and tests were
reviewed and verified by me, and I am responsible for all submitted
changes.
---------
Signed-off-by: Roberto Meroni <inbox1rob@gmail.com>
Signed-off-by: Ranran <ranranhaoranzhang@gmail.com>
Co-authored-by: Ranran <ranranhaoranzhang@gmail.com>
