Repository navigation
test: add unit tests for return_prompt_tokens serialization - #5
Merged
sidsingh-nvidia merged 1 commit intoJul 21, 2026
Conversation
Cover the new DynamicInferenceRequest.serialize() drop/keep contract from PR NVIDIA#5918: prompt_length is always reported (needed for usage.prompt_tokens), the prompt_tokens tensor is dropped from the wire payload unless SamplingParams.return_prompt_tokens=True, self.prompt_tokens is restored on the local instance, the routing_indices shape check still uses the saved prompt_len when self.prompt_tokens is temporarily nulled, and the None-prompt edge case does not crash. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: shanmugamr1992 <shanmugamr@nvidia.com>
sidsingh-nvidia
approved these changes
Jul 21, 2026
sidsingh-nvidia
marked this pull request as ready for review
July 21, 2026 21:54
sidsingh-nvidia
merged commit Jul 21, 2026
9555af0
into
sidsingh-nvidia:siddharth/trim-prompt-tokens-from-engine
1 check failed
sidsingh-nvidia
added a commit
that referenced
this pull request
Sep 1, 2026
The coordinator's hash table is a guess about what each engine still holds. It sees blocks being routed and never blocks being evicted, so the guess only ever gets more optimistic: an entry written once stays forever, and the coordinator keeps sending requests to a rank for a prefix it dropped long ago, paying a cold prefill and passing up a rank that could have served it. The table also grows without bound. Entries untouched for `prefix_cache_ttl_seconds` (default 300) are now dropped. The queue is insertion-ordered, so expired entries form a prefix of it and the sweep stops at the first live one -- the cost is what it evicts, not the size of the table. Re-routing a block refreshes it, and the stale queue entry left behind is recognized by its older timestamp and skipped, so a block still in use survives. Sweeping happens on the routing path rather than on a timer, because the coordinator is a single event loop with nothing else to run one, and expiry only has to keep pace with the traffic that grows the table. Two things differ from the version this came from. That one was dead code -- `_expire_rank_hashes` was defined but never called anywhere, so the TTL never actually ran. And its queue was keyed on `(time, hash, rank_idx)`, which is unsafe here: removing an engine renumbers the rank indices, so a queued entry would afterwards name a different rank and could evict the wrong one. Keying on `(time, hash)` and expiring every rank entry older than the queued timestamp is index-free and survives a renumber. The assignment counter the table used for timestamps is replaced by a monotonic clock, since ordering alone cannot say whether an entry is older than the TTL. Nothing read the old values: the routing score stopped consulting recency when the tiebreak moved to least-loaded. Verified against nvidia/main across all of tests/unit_tests/inference including engines/: the set of failing tests is identical, 218 on both sides. Five tests added, covering expiry, refresh-on-reuse, that an expired entry stops winning routing, that expiry survives an engine removal, and that entries carry the time they were routed. Signed-off-by: Siddharth Singh <sidsingh@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds unit-test coverage for the new
DynamicInferenceRequest.serialize()drop/keep contract in NVIDIA#5918 (return_prompt_tokens).test_dynamic_inference_request_serialize_return_prompt_tokens(parametrized onreturn_prompt_tokens=False/True):prompt_lengthis always populated on the serialized payload (needed forusage.prompt_tokens).return_prompt_tokens=False):prompt_tokensis dropped from the wire (present butNone).return_prompt_tokens=True):prompt_tokensis preserved as the("tensor", list)wrapper.self.prompt_tokensis restored afterserialize()— the drop is wire-only.routing_indicesshape check still passes with the drop path (uses the savedprompt_len, not the temporarily-nulledself.prompt_tokens).test_dynamic_inference_request_serialize_prompt_length_absent: edge case whereprompt_tokens is Noneon the request — no crash,prompt_lengthreported asNone.Added to the existing sister file
tests/unit_tests/inference/test_inference_request.pyrather than a new file, matching the file's existing style (well-commented docstrings,_make_dynamic_requestfixture, one parametrized test per contract).Verification
Ran on cw-dfw via cog (8 GPUs × torch.distributed.run, job
14220125):(8 pre-existing tests + 3 new test cases = 11 passed, on every rank.)
Test plan
import pytest).🤖 Generated with Claude Code