Repository navigation
Conversation
Serve vllm-sr/d3-turbo, whose decision_config.json declares prefix_bidir attention. Each request's state block and question block become the PrefixLM ranges of the full attention layers through a Model Runner V2 state, and the System One strategy sends each prompt's state block size with the checkpoint runtime's rule. A declared validated_length becomes the default max_model_len. Signed-off-by: Xunzhuo Liu <xunzhuo.liu@amd.com>
|
Documentation preview: https://vllm--62584.org.readthedocs.build/en/62584/ |
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
|
This pull request has merge conflicts that must be resolved before it can be |
|
@Xunzhuo Any news on this? |
|
@Xunzhuo Looking forward to it |
|
Hi @Xunzhuo, thanks for
Numbers with the state cache: vllm-sr/d3-turbo, a 26-question page of ~47K tokens, answers checked against
AI assistance (Claude Code) was used to write these changes, and AI agents reviewed them. Not all of the code has — Gilad Zamoscinski |
|
Hi @Xunzhuo, a question about
Are you planning to add d3-ultra here or in a follow-up? If not, I'd be glad to do the follow-up on top of this PR and validate it against the checkpoint's own runtime on an 80 GB GPU. |
Overview
Serve
vllm-sr/d3-turbo(27B) withvllm serve vllm-sr/d3-turboand no other flags, and ask it typed questions on/v1/systemone. Today it stops at start-up withValueError: unsupported decision_config.json attention_mode 'prefix_bidir': its checkpoint declares an attention mode thatD3ForSequenceClassificationdoes not take yet. Follow-up to #61565 and #61778, which added Decision 3.0 and its other sizes.Claims
prefix_bidirgets the checkpoint's own mask. A prompt is a state block (the system turn, the media,State:and the state) followed by a question block. Each request's two blocks become the PrefixLM ranges of the full attention layers, the mask vLLM already applies to multimodal prefix-LM models (causal, or both ways within a range): the state reads itself both ways, the question reads the state and itself both ways, and the state never sees a question. The linear-attention layers stay causal./v1/systemonegives the checkpoint runtime's answers on one AMD Instinct MI325X: 1 of 260 decisions changed (a near-tie: the reference's two best answers are 0.0056 apart), median |dp| 0.0007, max 0.026, over text, JSON, one image, several images, a video and a 24K-token state. At long context (64K to 117K-token states,--max-model-len 131072) 0 of 6 changed, max |dp| 0.014.validated_length, 32,768 tokens per question), and it becomes the defaultmax-model-lenwhen none is given;--max-model-len 131072serves long contexts. The other Decision 3.0 checkpoints declare none and are unaffected.D3Strategysends each prompt's state block size, counted with the runtime's own rule: the lead tokens that end by the end of the state text, so a token joining the state's last character to the question block belongs to the question.prefix_bidirin CI on a small checkpoint against d3-turbo's own runtime.Validation
Setup: one AMD Instinct MI325X per side. vLLM is the ROCm nightly image
vllm/vllm-openai-rocm:nightly-rocm72-e68efc685ab1e6f1eda004ced7157e3ca411da2f(PyTorch 2.13, ROCm 7.2) with this branch's Python tree (main59a51259cplus this PR; rebased onto1a74cc7bbbefore opening, with no changes to these files) copied over the installed one; no kernels were built. It ranvllm serve vllm-sr/d3-turbo(32,768 tokens per question, the new default;--max-model-len 131072for long context) with no other flags: Model Runner V2 and the Triton attention backend, which vLLM selects for the PrefixLM ranges. The reference is the checkpoint's own runtime,AutoModel.from_pretrained(..., trust_remote_code=True), with the versions itsrequirements.txtpins (transformers 5.17.0, flash-linear-attention 0.5.2) in the same image, on its default path: the state block once per request, then the questions on it, through its fused ROCm kernels. Checkpointvllm-sr/d3-turbo@689b3af1(tag v1.1).Answers vs the reference (
model.system_one(...)in process): a seeded set of 32 requests and 260 questions. Text: six states with 1, 4 and 16 questions each, a 16-question JSON state, a 64-question request, a 40-option and a 255-option choice, and an empty state. Media: three photos with 8 questions each, a one-question photo request, requests with 3 and with 2 images, an 11 s video with 6 questions, and a video with an image. Then a 24K-token state with 4 questions, and for long context 64K, 96K and 117K-token states with 2 questions each (reference:long_context=True). The decision is the top choice,noul >= 0.5, or a score's most likely level; |dp| is the largest absolute probability difference over a question's options.Latency, one request at a time, p50 over 3 runs after a warm-up (1 run for the 64K to 117K states), ms. The runtime runs in process; vLLM is measured over HTTP.
Start-up:
vllm serve vllm-sr/d3-turbowith no flags and an empty Hugging Face cache answered its first request 270 s after launch, the 52 GB download included.Throughput, closed loop over the same requests (the 23 text requests, or the 8 image and video requests), 30 to 60 s per level: vLLM serves the text mix at 5.95 req/s with 1 client, 7.56 with 8 and 7.89 with 64 (55.5, 69.6 and 73.2 questions/s), and the media mix at 0.87 and 0.95 req/s with 1 and 8 clients (5.0 and 5.3 questions/s). The runtime's mean latency on the same requests corresponds to 6.15 req/s (56.7 questions/s) on text and 2.45 req/s (13.5 questions/s) on media, one request at a time.
Tests
pytest tests/models/multimodal/pooling/test_d3.py(MI325X): 3 passed in 234 s: d3-edge and d3-nano as before, and the newtest_d3_prefix_bidir_matches_the_checkpoint_runtime. Without the state block sizes in the pooling params the new test fails (0.035 > 0.02).pytest tests/entrypoints/generate/structured_decisions/test_strategies.py: 62 passed, including the new state block test.pytest tests/config/test_model_arch_config.py -k decision_config: 8 passed;pytest tests/models/test_registry.py -k D3: 1 passed.pre-commit runon the changed files: all hooks pass (ruff, mypy, typos, markdownlint and the repository checks).VLLM_GPU_SYNC_CHECK=erroron the server: 62 requests (text, one and several images, video) with no synchronization error.Details
[0, s)for the state block and[s, n)for the question block, give exactly the runtime'sprefix_bidirmask, so no kernel changes are needed.D3Strategycounts each prompt's state block (images and videos included, since they come first) and sends it asPoolingParams.extra_kwargs["d3_state_tokens"].D3PrefixBidirModelState, the model's Model Runner V2 state, turns it into the request's two ranges. The only change outside the model is aDefaultModelState.get_mm_prefix_rangeshook thatprepare_attncalls; its default is the previous inline code.decision_config.json; forprefix_bidirit also setsis_mm_prefix_lm, so attention backends that take the ranges are selected. The config hook keeps the text layers causal (decoder attention that reads the ranges) and marks the model non-causal, so chunked prefill and prefix caching stay off as for the other Decision 3.0 models. A declaredvalidated_lengthsetsmax_model_lenonly when the user gives none.prefix_bidircheckpoint there raises a clear error. FlexAttention does not take this model's hybrid KV cache layout (784-token blocks), so on ROCm the Triton backend serves it./classifyand/poolingcannot carry the state block size, so aprefix_bidirprompt sent there reads as a single block;LLM.classifytakes it throughPoolingParams, as the new test does./v1/systemonedoes not offer; NVIDIA measurements beyond CI (no NVIDIA GPU was available to me); tensor parallelism; LoRA.prefix_bidir(searched for d3-turbo,prefix_bidir, Decision 3.0,D3ForSequenceClassification). [Frontend] Tokenize the /v1/systemone state once per decision and let big decisions take turns #61251 changes how/v1/systemonetokenizes a shared state and also editsD3Strategy.read(); whichever lands second rebases, and the state block size here is computed on the same prompt either way.Follow-ups
prefix_bidirprompts on ROCm./v1/systemone.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.