[Bugfix][Spec Decode] Handle out-of-range IDs in heterogeneous vocab mapping

[shaohuaxi]

Purpose

Mixed prompt embeddings can leave placeholder token IDs outside the model vocabulary. These positions use the supplied embeddings in the target model, but heterogeneous speculative decoding still looks up their IDs in a mapping table sized from the model config. The lookup fails before the existing UNK fallback can run.

This change checks the ID range before lookup, uses bounded indices for the gather, and maps out-of-range IDs to UNK. It applies the same handling in both directions and preserves the existing fallback for tokens outside the shared vocabulary, without adding a tensor-dependent Python branch.

Duplicate-work check on September 28, 2026 found no existing fix for this case. #57396 removes CPU/GPU synchronization from these methods; #51564 addresses the separate local-argmax path that bypasses vocabulary mapping.

Test Plan

Using the CPU test environment with this checkout on PYTHONPATH (TEST_PYTHON is its Python executable):

PYTHONDONTWRITEBYTECODE=1 \
PYTHONPATH="$PWD" VLLM_TARGET_DEVICE=cpu \
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
"$TEST_PYTHON" -m pytest -q -p no:cacheprovider \
  tests/v1/spec_decode/test_vocab_mapping.py -o addopts=''

The existing CPU/CUDA fixture now includes tokenizer entries beyond the model vocabulary. The added cases cover both mapping directions, negative IDs, and empty inputs, while retaining the existing dtype, device, and GPU synchronization assertions.

Also ran the unchanged test file on an NVIDIA A10 using an existing CUDA 13 image, with the committed mapping and GPU sync-check modules mounted read-only:

/tmp/.venv/bin/python -m pytest \
  --rootdir=/bundle/suite --confcutdir=/bundle/suite \
  /bundle/suite/test_vocab_mapping.py -vv -p no:cacheprovider \
  --junitxml=/results/fixed/tests.xml

Test Result

• A10: 29 passed, 0 skipped, including all 12 CUDA cases and the existing GPU sync checks. The earlier CPU run was 17 passed, 12 CUDA cases skipped.
• The four new boundary cases fail against the unchanged main mapping on both CPU and CUDA. On CUDA, out-of-range lookups trigger device-side asserts; negative indices return incorrect tokens. Each failing CUDA case ran in a separate process.
• A separate GPT-2 tokenizer check on CUDA maps added token ID 50257 to UNK 50256 in both directions, for int32 and int64 inputs. CUDA Graph capture and replay with changed inputs also passed for all four combinations.
• Changed-file pre-commit and mypy-3.12 checks passed.

AI assistance: Codex assisted with the implementation and validation.

[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.

[shaohuaxi]

Hi @tomasruizt, since you reviewed the vocab mapping and CUDA sync checks in #57396, could you take a look at this PR when you have a chance? It maps out-of-range IDs to UNK in both directions while preserving the no-host-sync path. All 29 tests passed on A10, including the CUDA sync checks. I'd appreciate your feedback. Thanks!

[tomasruizt]

@shaohuaxi I can review it next week

[shaohuaxi]

Thanks!