Skip to content

[Model] Support Decision 3.0 prefix_bidir attention (vllm-sr/d3-turbo) - #62584

Open
Xunzhuo wants to merge 1 commit into
vllm-project:mainfrom
Xunzhuo:d3-turbo-prefix-bidir
Open

Xunzhuo wants to merge 1 commit into
vllm-project:mainfrom
Xunzhuo:d3-turbo-prefix-bidir

Conversation

@Xunzhuo

@Xunzhuo Xunzhuo commented Oct 20, 2026 •

Copy link
Copy Markdown
Member

Overview

Serve vllm-sr/d3-turbo (27B) with vllm serve vllm-sr/d3-turbo and no other flags, and ask it typed questions on /v1/systemone. Today it stops at start-up with ValueError: unsupported decision_config.json attention_mode 'prefix_bidir': its checkpoint declares an attention mode that D3ForSequenceClassification does not take yet. Follow-up to #61565 and #61778, which added Decision 3.0 and its other sizes.

Claims

  • prefix_bidir gets 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/systemone gives 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.
  • No flags needed. The checkpoint declares the length it was validated at (validated_length, 32,768 tokens per question), and it becomes the default max-model-len when none is given; --max-model-len 131072 serves long contexts. The other Decision 3.0 checkpoints declare none and are unaffected.
  • D3Strategy sends 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.
  • The other Decision 3.0 models are unchanged (their GPU tests pass). The docs list d3-turbo and its long-context limit, the test registry lists it, and a new GPU test runs prefix_bidir in 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 (main 59a51259c plus this PR; rebased onto 1a74cc7bb before opening, with no changes to these files) copied over the installed one; no kernels were built. It ran vllm serve vllm-sr/d3-turbo (32,768 tokens per question, the new default; --max-model-len 131072 for 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 its requirements.txt pins (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. Checkpoint vllm-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.

Set Questions Decisions changed Median abs dp p90 Max
text and JSON 212 1 (reference margin 0.0056) 0.0007 0.0072 0.019
one image 25 0 0.0001 0.0057 0.025
several images 10 0 0.0001 0.013 0.013
video 9 0 0.0010 0.026 0.026
24K-token state 4 0 0.0023 0.0027 0.0027
64K, 96K, 117K-token states 6 0 0.0024 0.014 0.014

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.

Request Runtime vLLM
1 text question 60 to 61 30 to 35
4 text questions 77 to 107 58 to 85
16 text questions 176 to 242 225 to 271
16 questions, JSON state 218 376
64 text questions 756 810
255-option choice 398 370
1 photo, 1 question 303 191
1 photo, 8 questions 254 to 360 909 to 1,310
3 photos, 6 questions 682 2,489
11 s video, 6 questions 447 1,456
24K-token state, 4 questions 3,579 24,862
64K-token state, 2 questions 10,574 55,449
117K-token state, 2 questions 26,275 168,277

Start-up: vllm serve vllm-sr/d3-turbo with 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 new test_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 run on the changed files: all hooks pass (ruff, mypy, typos, markdownlint and the repository checks).
  • VLLM_GPU_SYNC_CHECK=error on the server: 62 requests (text, one and several images, video) with no synchronization error.

Details

  • The mask: the Triton, FlexAttention and FlashAttention 4 backends already OR per-request bidirectional ranges into the causal mask for multimodal prefix-LM models. Two ranges per request, [0, s) for the state block and [s, n) for the question block, give exactly the runtime's prefix_bidir mask, so no kernel changes are needed.
  • Where the ranges come from: D3Strategy counts each prompt's state block (images and videos included, since they come first) and sends it as PoolingParams.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 a DefaultModelState.get_mm_prefix_ranges hook that prepare_attn calls; its default is the previous inline code.
  • Detection and config: the arch config convertor already detects Decision 3.0 checkpoints by their decision_config.json; for prefix_bidir it also sets is_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 declared validated_length sets max_model_len only when the user gives none.
  • Tradeoffs and limits
    • State per question: the runtime reads a request's state once and answers every question from it. vLLM reads it again for each question, one engine request per question as for the other Decision 3.0 models, so many questions about a large state (several images, long text) cost more than in the runtime (latency table). The media themselves are processed once per request and the vision encoder output is cached.
    • Long prompts: at 64K tokens one question takes about 28 s in vLLM and the runtime's state pass about 10 s; most of the gap is the attention kernel with ranges, which is slower than the runtime's SDPA at that length on this GPU.
    • Model Runner V2 on CUDA or ROCm GPUs only: Model Runner V1 does not give the ranges to the attention, so loading a prefix_bidir checkpoint 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.
    • /classify and /pooling cannot carry the state block size, so a prefix_bidir prompt sent there reads as a single block; LLM.classify takes it through PoolingParams, as the new test does.
    • Not in this PR: Box and Track questions and images inside questions or options, which /v1/systemone does not offer; NVIDIA measurements beyond CI (no NVIDIA GPU was available to me); tensor parallelism; LoRA.
  • CI: the new test runs in the optional "Multimodal Models (Extended Pooling)" step on d3-edge's weights (already downloaded there) plus six small runtime files from d3-turbo at a pinned revision. d3-turbo's own 52 GB checkpoint is covered by the measurements above, not CI.
  • Not a duplicate: no open or closed PR or issue covers d3-turbo or 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/systemone tokenizes a shared state and also edits D3Strategy.read(); whichever lands second rebases, and the state block size here is computed on the same prompt either way.
  • AI assistance was used for the code, tests, measurements and this description. The PR stays a draft until the submitter has reviewed every line.

Follow-ups

  • State once: answer all of a request's questions from one pass over its state, either as one engine request per state or with prefix caching restricted to KV blocks that lie entirely inside the state block (the question block reads itself both ways, so blocks that cross into it cannot be shared).
  • A faster attention path for long prefix_bidir prompts on ROCm.
  • Box and Track questions and images inside questions and options on /v1/systemone.

Pull Request Checklist
  • I used vLLM's /pr-checklist skill. (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.

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>
@mergify

mergify Bot commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

Documentation preview: https://vllm--62584.org.readthedocs.build/en/62584/

@mergify mergify Bot added documentation Improvements or additions to documentation frontend pooling multi-modality Related to multi-modality (#4194) mrv2 Model Runner V2 specific labels Oct 20, 2026
@Xunzhuo Xunzhuo mentioned this pull request Oct 20, 2026
6 tasks done
@Xunzhuo
Xunzhuo marked this pull request as ready for review October 20, 2026 02:05
Copilot AI balanced review requested due to automatic review settings October 20, 2026 02:05

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify

mergify Bot commented Oct 22, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Xunzhuo.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@gaby

gaby commented Oct 26, 2026

Copy link
Copy Markdown
Contributor

@Xunzhuo Any news on this?

@QuantumCamper

QuantumCamper commented Oct 27, 2026 •

Copy link
Copy Markdown

@Xunzhuo Looking forward to it

@zamog

zamog commented Oct 28, 2026

Copy link
Copy Markdown
Contributor

Hi @Xunzhuo, thanks for prefix_bidir. We built prefix caching of the D3 state block on top of it (RFC:
#64267). Two things from this PR's code:

  1. A bool is not a state block size. The state length was checked with isinstance(state, int), which accepts
    a bool, so extra_kwargs={"d3_state_tokens": True} made a one-token state block. Accept only an int, as
    PrefixLMScheduler does for prefix_lm_len: zamog@99fd132
  2. Keep the hybrid model state. D3PrefixBidirModelState derives from DefaultModelState, so a prefix_bidir
    checkpoint bypasses the MambaHybridModelState the runner picks for Qwen3.5 otherwise.

Numbers with the state cache: vllm-sr/d3-turbo, a 26-question page of ~47K tokens, answers checked against
uncached runs.

A100-80GB H200
Decisions changed vs. uncached (30K/60K states ×1/4/16/64 questions; 117K on H200) 0 0
Computed tokens per page, state cached −25% —
Throughput at c16, pages/s (with CPU offload of evicted states) 0.845 2.19

AI assistance (Claude Code) was used to write these changes, and AI agents reviewed them. Not all of the code has
been reviewed line by line by a human yet. Reviews are very welcome.

— Gilad Zamoscinski

@Manfredss

Copy link
Copy Markdown
Contributor

Hi @Xunzhuo, a question about vllm-sr/d3-ultra, released after this PR. Its decision_config.json uses the same prefix_bidir attention as d3-turbo, so main rejects it at start-up the same way (unsupported decision_config.json attention_mode 'prefix_bidir', vllm/model_executor/models/config.py). It also has two fields that d3-turbo does not have and that neither main nor this PR reads:

  • "prompt_layout": "header-v1"
  • "temperature_by_questions": {"max_questions": [1, 4, 16], "temperature": [1.046, 1.0, 1.0, 1.0]}

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.

@gaby

gaby commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

@Xunzhuo @noooop Any updates on this? We now have two models that can't be used with vLLM:

  • vllm-sr/d3-turbo
  • vllm-sr/d3-ultra

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation frontend mrv2 Model Runner V2 specific multi-modality Related to multi-modality (#4194) needs-rebase pooling

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants