[Perf][Spec Decode] Filter rejection sampler output rows in Python instead of per-row NumPy

[ZhengGong-amd]

Purpose

RejectionSampler.parse_output builds its result with a NumPy fancy index and a
tolist() per row:

outputs = [
    row[valid_mask[i]].tolist() for i, row in enumerate(output_token_ids_np)
]

A row holds max_spec_len + 1 tokens — typically 2 to 6 — so each request pays two NumPy
dispatches to filter a handful of values, and the cost grows linearly with the running
batch. Convert the batch once and filter in Python instead, which is faster at every batch
size measured, and drop the valid_mask allocation entirely when logprobs are off, since
nothing else reads it.

The function is the max_gen_len > 1 branch of get_output() and execute_model(), so it
runs once per speculative-decoding step for ngram, EAGLE/EAGLE3, MTP and draft-model alike.

Test Plan

Both versions are lifted out of their trees by text and run in one interpreter so a single
loop alternates between them. Grid: batch 1–512 × max_spec_len 1/3/5 × acceptance
0.3/0.6/0.9, vocab_size=151936, int32, 5 interleaved pairs per shape with the order
alternated, min-of-pairs, within-arm spread ≤ 1.5 %.

Equivalence: 2560 cases (batch 0–64, 1–6 columns, two vocab sizes, all-placeholder rows,
out-of-vocabulary and negative ids, discard_req_indices as list / ndarray / tensor / every
row, with and without logprobs_tensors) through the real function in both trees.

pytest tests/v1/sample/test_rejection_sampler.py -k parse_output    # 33 passed
vllm bench serve --dataset-name random --random-input-len 128 --random-output-len 512 \
  --max-concurrency {16,64} --ignore-eos --seed 1234                # Qwen3-8B, ngram, TP=1

Test Result

Equivalence: 0 mismatches over 2560 cases, byte-identical outputs including the
logprobs path.

parse_output with a host-side input, max_spec_len=3, acceptance 0.6:

batch	before	after	Δ
1	2.55 µs	0.96 µs	−62.5 %
16	6.62 µs	2.21 µs	−66.7 %
64	20.52 µs	6.82 µs	−66.8 %
256	77.64 µs	27.27 µs	−64.9 %
512	158.70 µs	66.75 µs	−57.9 %

Over the full 90-shape grid: median −64.6 %, worst case −48.3 %, no shape regresses.

Under async scheduling get_output() passes a host tensor, so the table above is the whole
function. The synchronous path passes a device tensor and also pays .cpu(); the same
change reads −12.6 % at batch 4, −25.2 % at 16, −39.8 % at 64, −56.5 % at 256.

End-to-end (Qwen3-8B, ngram, TP=1, 4 interleaved pairs, pair 0 dropped) is flat: +4.2 %
at concurrency 16 and +0.6 % at 64, against within-arm spreads of 6.4 % and 4.4 %. No
end-to-end speedup is claimed — the saving is microseconds against a multi-millisecond
decode step. The claim is the host-side cost of the function, and that it no longer grows
with the running batch.

The obvious alternative — one masked gather plus a Python split — is better at batch 512
(33 µs against 67 µs) but costs +15 % at batch 1 and regresses 18 of 90 grid shapes.

Scope

Platform-independent — no custom op, no backend dispatch, no current_platform branch — so
CUDA gets the same change by the same mechanism. The filter predicate, the empty rows for
discarded requests and the pre-discard logprobs boundaries are unchanged, and the
max_gen_len == 1 path does not reach this function.

Reproducing the numbers

Patch bf74f8e39c on 9598487bc1. Measured on vLLM v0.29.0 (98dff2a81d), torch
2.12.0+git6bbd260; parse_output's body is identical between that tag and main apart
from docstring blank lines added by #52136.

---

Found by Hyperloom; reviewed and measured by hand.

[ZhengGong-amd]

Impact surface and end-to-end

Who reaches it. parse_output is the max_gen_len > 1 branch, i.e. every
speculative-decoding step, for ngram, EAGLE/EAGLE3, MTP and draft-model alike. Scanning
the checkpoints here, 32 carry MTP draft layers, covering every current flagship family —
DeepSeek-V4/V4.1, GLM-5.2/5.3, MiniMax-M3, Qwen3.5/3.6/3.8. It is architecture-agnostic:
nothing in the function looks at the attention class, quantization or platform.

Which denominator. The two call sites pass different tensors. Under async scheduling
get_output() passes sampled_token_ids_cpu, already on the host, and the micro table in
the description is the whole function. The synchronous path passes a device tensor, so
parse_output also pays .cpu():

batch	before	after	Δ
4	17.00 µs	14.86 µs	−12.6 %
16	21.70 µs	16.23 µs	−25.2 %
64	34.18 µs	20.58 µs	−39.8 %
256	97.17 µs	42.24 µs	−56.5 %

End-to-end: no signal, as expected. Qwen3-8B + ngram (num_speculative_tokens=3),
MI355X, TP=1, vllm bench serve, 128 in / 512 out, 4 interleaved pairs with the arm order
alternated, pair 0 dropped:

shape	output tok/s Δ (median of 3 pairs)	per-pair	base within-arm spread
conc 16	+4.2 %	−5.0 / +6.4 / +4.2	6.4 %
conc 64	+0.6 %	+0.9 / +0.6 / −3.4	4.4 %

The spread is larger than the effect in both shapes — ngram acceptance varies run to run —
so this is flat, not a win, and I am not claiming the micro number carries over. It
should not: a few microseconds against a multi-millisecond decode step. The claim is the
host-side cost of the function itself, and that it no longer grows with the running batch.

No GSM8K run, deliberately. The two implementations are not statistically equivalent,
they are identical: 2560 differential cases (batch 0–64, 1–6 columns, two vocab sizes,
all-placeholder rows, out-of-vocabulary and negative ids, discard_req_indices as
list / ndarray / tensor / every row, with and without logprobs_tensors) produce
byte-identical outputs on both trees. A sampled accuracy benchmark would be a weaker
statement than the exhaustive one.

[ZhengGong-amd]

CUDA side

This change is platform-independent and does affect CUDA, in the same direction and by
the same mechanism.

parse_output is pure NumPy and Python — no custom op, no backend dispatch, no
current_platform branch — so both arms execute identical code on every platform. The
micro numbers in the description were taken with host tensors and are therefore a CPU
measurement, not a ROCm one; the only platform-dependent part is the .cpu() copy at the
synchronous call site, and that is unchanged.

Async scheduling is the case that matters most on CUDA: get_output() hands the function
a tensor that is already on the host, so the whole call is host time on the critical path
between steps, and it grows linearly with the running batch on main.

I do not have an NVIDIA GPU here, so there is no CUDA timing in this PR. I would rather
say that than imply otherwise — but the claim does not need one: the code path is the same
one, and the equivalence evidence (2560 differential cases, byte-identical) is
platform-independent too.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.