Repository navigation
Conversation
Qwen3-VL routes an ``input_deepstack_embeds`` kwarg to the language
model's forward on image requests. Under breakable-CUDA-graph replay
the tensor was silently dropped: the captured graph never traced the
DeepStack ``add_`` branch, and ``replay_layer_forward`` did not copy
the live tensor anywhere, so image outputs on any DeepStack-carrying
model on the BCG allowlist matched a zero-DeepStack run instead of
the eager reference. No shipped configuration triggers this today
(``Qwen/Qwen3.5-*`` releases ship with ``deepstack_visual_indexes=[]``
and Qwen3-VL is not on ``multimodal_breakable_cuda_graph_supported_model_archs``),
so this is a latent regression that would activate on either config
change; closing it defensively removes the trap.
Fix mirrors the existing ``input_embeds`` slot-and-copy pattern:
* ``Qwen3VLForConditionalGeneration`` declares an explicit class
attribute ``supports_bcg_deepstack_replay = True``. Its MoE
subclass ``Qwen3VLMoeForConditionalGeneration`` inherits the flag.
Non-Qwen3-VL archs (Qwen2.5-VL, Qwen3, Cohere2Vision, KimiK25,
MiniMaxM3) do not set the attribute and default to ``False`` via
``getattr``.
* ``PrefillCudaGraphRunner.__init__`` computes
``deepstack_replay_width = hidden_size × num_deepstack_embeddings``
only when the capability flag is ``True`` AND the instance reports
``num_deepstack_embeddings > 0`` — otherwise ``0``. The width is
passed to both ``PrefillInputBuffers.create`` (which allocates a
single ``[max_num_tokens, width]`` bf16 buffer) and
``build_prefill_registry`` (which registers the matching optional
``GraphSlot``).
* ``_run_forward``'s BCG capture branch passes the slot buffer as
the ``input_deepstack_embeds`` kwarg when ``has_slot`` reports the
slot is registered, so Dynamo traces the ``add_`` branch into the
captured graph. Otherwise it falls through to today's four-
positional call unchanged.
* ``replay_layer_forward`` mirrors the ``input_embeds`` copy: when
the slot is registered, it validates the live kwarg (present,
``numel() > 0``, matching per-row shape, dtype, device) and copies
into the slot with a zeroed padded tail; missing / empty /
mismatched inputs land in an ``else`` branch that zeros the whole
slice so a stale image contribution cannot leak across bucket-
reuse into a text-only request.
Zero cost for models that do not declare the capability: no
``PrefillInputBuffers`` field allocated, no registry slot, no
captured-graph kwarg, no replay-side copy — the entire block is
guarded by ``has_slot()``.
Refs sgl-project#30868 (PCG DeepStack Dynamo warmup) and sgl-project#30872 (input_embeds
BCG slot). Preserves ``run_dummy_multimodal_deepstack_forward`` for
the TC piecewise path unchanged.
CPU-only unit tests covering the three-site contract of the fix in
``fix(bcg): copy Qwen3-VL input_deepstack_embeds into a stable
replay slot``:
* ``build_prefill_registry`` registers the ``input_deepstack_embeds``
slot iff ``is_multimodal AND register_input_embeds AND
deepstack_replay_width > 0``; the slot is skipped when the eager-
extend path passes ``register_input_embeds=False``; the slot's
shape and dtype match ``embed_dtype`` and the width kwarg; the
existing ``input_embeds`` slot registration is undisturbed.
* ``PrefillInputBuffers.create`` allocates the matching
``input_deepstack_embeds`` field iff both is_multimodal AND
deepstack_replay_width > 0. Adoption by the registry via
``source=self.buffers`` succeeds — the buffer field IS the same
tensor the slot exposes.
* ``Qwen3VLForConditionalGeneration`` declares the explicit opt-in
class attribute. ``Qwen3VLMoeForConditionalGeneration`` inherits
it. ``Qwen2_5_VLForConditionalGeneration`` and text-only
``Qwen3ForCausalLM`` do not — getattr defaults to False and the
runner's whole DeepStack code path is skipped for them.
Runs under ``pytest test/registered/unit/model_executor/test_deepstack_replay_slot.py``;
registered on the ``base-a-test-cpu`` CI suite with ``est_time=5``.
No GPU / no CUDA context required.
…eview
Behavior-preserving simplifications identified during the pre-PR
adversarial review; verified by:
* 46 existing model_executor unit tests still pass on the patched
branch (test_cuda_graph_buffer_registry / test_prefill_cuda_graph_runner
/ test_prefill_cuda_graph_runner_helpers);
* 11 new unit tests in test_deepstack_replay_slot.py still pass
(consolidated from 14 — see below);
* GPU 4-arm 2x2 rerun on upstream/main SHA eac1f78 produced
verdict PASS_BCG_CORRECT with signature bit-identical to the pre-
simplification M4b run:
bcg_normal vs eager_normal: (15, True, 0.071) — unchanged
bcg_normal vs bcg_zero: (7, False, 1.148) — unchanged
bcg_zero vs eager_zero: (20, True, 0.066) — unchanged
eager_zero vs eager_normal: (7, False, 1.138) — unchanged
Changes (source, net -37 lines from +119 to +82):
* qwen3_vl.py: trim capability-flag comment from 10 lines to 3;
drop the `: bool` type annotation to match `supports_lora = True`
style used elsewhere (gemma3_causal.py:736, torch_native_llama.py:400,
dflash.py:338).
* cuda_graph_buffer_registry.py: replace the lambda closure-capture
default-arg trick with a plain lambda over the function argument;
trim the slot comment from 9 lines to 3.
* prefill_cuda_graph_runner.py: collapse the capture-pass if/else
(which duplicated the same 5-arg call) into a single call with
kwargs unpacking (~9 lines saved); trim the width-computation
comment; trim the replay-side comment while keeping the load-
bearing note about why zero-on-empty is required.
* buffers.py: trim allocation comment from 6 lines to 1.
Changes (tests, -21 lines from 214 to 193):
* Drop test_input_embeds_slot_registered_alongside — coexistence
is implicit in every other test that instantiates the registry.
* Consolidate the three negative slot-registration tests (not-
multimodal, width-zero, register_input_embeds=False) into a
single test_slot_absent_when_any_gate_is_off using subTests.
11 tests remain, still covering: slot-absent (all 3 gate-off cases
via subTests), slot-present, slot shape+dtype contract, buffer
field allocation (3 cases), registry adoption, Qwen3-VL declares
flag, Qwen3-VL-Moe inherits flag, Qwen2.5-VL does not declare
flag, text-only Qwen3 does not declare flag.
No semantic change to the fix itself; no behavior visible to callers
changes. Preservation invariants intact.
Pre-commit's black-jupyter hook noted two whitespace nits on the files touched by fd4c4cb / c9d6d89 / 9410775: * prefill_cuda_graph_runner.py: collapse a 3-line ``getattr`` into one line (fits in 88 chars). * test_deepstack_replay_slot.py: add missing blank line between two class definitions. Both are pure formatting; no semantic change. Verified idempotent by rerunning ``pre-commit run black-jupyter --files ...`` — the second run passes with no further changes. Committing as a focused new commit rather than amending prior commits per the reviewer's do-not-amend / do-not-force-push policy.
JustinTong0323
left a comment
JustinTong0323
left a comment
There was a problem hiding this comment.
Cleaned up the DeepStack replay patch: trimmed redundant inline comments, added a logger.warning in the replay slot rejection branch so a non-None de that fails dtype/shape/device validation now surfaces instead of silently replaying zero contribution.
Verified locally: the 11 unit tests pass in test_deepstack_replay_slot.py; pre-commit is green on the touched Python files (rust hooks skipped locally, no rust code changed here).
Approving — the fix is correct and well-scoped. CI will need a maintainer's run-ci label / authorization to actually run on this fork PR.
There was a problem hiding this comment.
Thanks for pushing this. I reviewed the replay path closely because we are actively working on SGLang-Omni Qwen3-Omni prefill BCG performance, and correct image/video DeepStack replay is an upstream dependency for that work.
The overall abstraction here looks right to me: explicit model opt in, a separate token axis stable slot, capture time storage, replay time refresh, and tail clearing so shorter requests cannot inherit stale DeepStack rows. The four arm H200 experiment is also exactly the semantic control I would want for this kind of silent replay bug.
I found one correctness behavior I think is worth tightening before merge:
Malformed not None DeepStack should fail closed.
In replay_layer_forward, if input_deepstack_embeds is not None but fails the shape/dtype/device checks, the current path warns, zeros the stable slot, and still replays the graph. That knowingly changes the request to zero-DeepStack semantics, which is very close to the silent behavior this PR is fixing.
A genuinely absent DeepStack contribution can safely clear the slot. If an empty tensor is an intentional absence sentinel, that can be handled explicitly too. For a non-None tensor that is incompatible with the captured slot, I would prefer either:
- eager fallback if validation can happen before graph commitment; or
- otherwise a
RuntimeErrorfrom replay rather than returning a semantically altered result.
I would also strongly recommend one focused regression around the runtime slot lifecycle itself:
- nonzero DeepStack is copied into the stable slot;
- a shorter same bucket request zeros the unused tail;
- a request with no DeepStack contribution clears the slot;
- malformed not None shape/dtype/device fails closed.
The current allocation/registry tests are useful, and the H200 four arm result is convincing. The additional lifecycle regression would directly protect the stateful same bucket behavior that can otherwise turn into cross request stale data.
For downstream context, Qwen3-Omni in SGLang-Omni produces DeepStack features for both image and video. We are working on the SGLang-Omni performance path to late-bind both final input_embeds and full token aligned input_deepstack_embeds through BCG, so this PR is directly relevant to that work. Once this contract is settled, I am happy to run Qwen3-Omni image/video validation against the exact PR head and report replay/parity evidence here if that would help.
|
Hi Charlie, thanks for the review. You're right that the warn-and-zero path for a malformed non-None DeepStack tensor is basically the same silent-semantic-drift bug this PR is trying to kill. I'll fix it so only None (and an explicit empty-tensor sentinel, if we end up wanting one) clears the slot silently. A non-None tensor that fails the shape/dtype/device checks will either fall back to eager or raise a RuntimeError at replay instead of quietly running with zeroed DeepStack. I'll also add the lifecycle regression test you outlined, nonzero copy-in, tail clearing on a shorter same-bucket request, clean clearing when there's no contribution, and fail-closed on malformed input, next to the existing allocation/registry tests. Once this is settled I'd really appreciate the Qwen3-Omni image/video validation against the PR head, I'll ping you here when it's ready. |
The DeepStack replay slot previously warned and zeroed itself whenever
input_deepstack_embeds failed its shape/dtype/device checks, then
replayed anyway. For a non-empty tensor that is a silent downgrade to
zero-DeepStack semantics — the same class of quiet corruption the slot
was added to prevent, with a log line in front of it.
Split the cases apart:
* None or an empty tensor state absence; the slot is cleared so a
previous replay's rows cannot bleed into a text-only or shorter
request through the LM's add_.
* A fitting tensor is copied in and the unused tail cleared.
* A non-empty tensor that does not fit the captured slot now raises
RuntimeError before any write, so the request fails loudly instead
of being answered without its visual contribution.
Validation cannot move ahead of graph commitment at this site:
input_deepstack_embeds is not materialized until get_input_embeddings
runs inside the outer model.forward, which is already past
_prepare_forward_metadata_for_replay and inside the replay context.
Raising is safe for runner state — the layer_model.forward patch is
restored in a finally.
The refresh logic moves out of the replay closure into
_refresh_deepstack_replay_slot so it is unit-testable; the closure drops
from 33 lines to 8.
Tests cover the slot lifecycle across replays: stale-tail clearing on a
shorter same-bucket request, clearing on absence, the empty-tensor
sentinel, and fail-closed on width/dtype/row-count mismatch. The four
fail-closed cases were verified to fail against the pre-fix branch; the
lifecycle cases guard behavior that already held but had no coverage.
|
@charliechenye Fixed in f2596d5. Malformed non-empty input_deepstack_embeds now raises at replay instead of zeroing and continuing; None and empty tensors still clear the slot as genuine absence. I chose the raise over eager fallback because the tensor isn't materialized until separate_deepstack_embeds runs inside the model's forward (via general_mm_embed_routine), which is already past _prepare_forward_metadata_for_replay and inside the replay context, so there's no pre-commit point to validate at. It's state-safe: validation precedes any write to the slot, and the layer_model.forward patch is restored in a finally. The refresh also moved into a pure helper so it's unit-testable, and your lifecycle regression is in. The fail-closed cases were verified to fail against the pre-fix branch. Would still love the Qwen3-Omni validation once CI settles. |
Drop cases that cannot turn red under any plausible change: - The empty-tensor sentinel: without the numel branch an empty tensor still passes the fit check and clears the slot identically. - Row/width fail-closed subcases: copy_ itself raises RuntimeError on shape mismatch, so they pin PyTorch, not the guard. - The off-device subcase has no CPU-only expression: a meta-device tensor fails copy_ with NotImplementedError, itself a RuntimeError subclass, so the assert passes with the guard deleted. Keep and strengthen: - Merge gate on/off and capability opt-in cases into table tests. - _leaves_slot_unmodified feeds a dtype-mismatched input, the only rejected input copy_ would have modified the slot with, so it is what proves validate-before-write. - The remaining malformed case is dtype drift, the one mismatch copy_ absorbs silently; verified red when the validation guard is removed.
The docstring keeps the two non-obvious invariants — stale rows bleed through the LM's add_, and a non-fitting contribution raises instead of being zeroed and served silently — and drops restatements of the signature; the closure comment adds nothing over the function name.
With the DeepStack replay slot in place, the known arch-specific fault for Qwen3-VL (image prefill replayed without its DeepStack contribution) is fixed, so the arch meets the bar for the BCG multimodal opt-in allowlist and gets prefill graphs by default. Qwen3VLMoe inherits both the allowlist entry and supports_bcg_deepstack_replay; hardware e2e so far covered the dense 8B only. The allowlist unit test stops using Qwen3VLForConditionalGeneration as its negative example (now Qwen2_5_VL) and covers both new archs.
JustinTong0323
added
the
run-ci-extra
|
@JustinTong0323 Description updated with all four items, including the explicit subclass list for the inheritance point. MoE smoke done on Qwen3-VL-30B-A3B-Instruct (H200, greedy, deepstack_visual_indexes=[8,16,24] so non-vacuous): BCG and eager give byte-identical completions on image and text prompts, with the bcg arm logging cuda graph: True on both substantive prefills where eager is False. |
|
Merged Re-validated on an up-to-date stack, as asked earlier: 1× H200, torch 2.14.1+cu130, sglang-kernel 0.4.9, flashinfer 0.7.0.post1, source install of this branch, no Unit tests pass ( |
…validation - Compute deepstack_replay_width only for BREAKABLE/FULL; tc_piecewise captures the outer model.forward and never binds or refreshes the slot, so it no longer allocates the (max_num_tokens, width) buffer (192 MiB for Qwen3-VL-8B at 8192 tokens). - Reject rank-mismatched input_deepstack_embeds in _refresh_deepstack_replay_slot so a 0-dim tensor raises the descriptive RuntimeError instead of an IndexError. - Tests: parameterize malformed-input cases (dtype, broadcastable width, excess rows, scalar), cover empty-tensor clearing, and add a wiring regression that drives _execute_body_capture to prove the replay closure refreshes the captured slot and that malformed input prevents backend replay.
There was a problem hiding this comment.
[Codex]: Code review of 4d7d9cc64e41ed54a56f11de6cc724462949d4d0.
The fix looks sound on source inspection. Capture binds a stable DeepStack buffer, replay copies the current request's features into it, and absence/tail clearing prevents stale contributions from earlier requests. Malformed non-empty inputs are rejected before replay, and the temporary forward replacement is restored in finally. I found no confirmed correctness defect in these changes.
The new tests cover allocation, registration/adoption, refresh and clearing, rejection of malformed inputs, and replay wiring. Two non-blocking follow-ups:
- Consider an explicit capability interface/defaults instead of the newly added
getattrfallbacks, and move the new test's function-local imports to module scope. This is a maintainability suggestion, not a correctness blocker.
Oasis-Git
left a comment
Oasis-Git
left a comment
There was a problem hiding this comment.
[Codex]: Two non-blocking code-style comments inline.
| * getattr(self.model_runner.model, "num_deepstack_embeddings", 0) | ||
| if self._uses_eager_prefill_tail() | ||
| and getattr(self.model_runner.model, "supports_bcg_deepstack_replay", False) |
There was a problem hiding this comment.
[Codex]: Please avoid getattr fallbacks for these capability fields. Define an explicit capability interface with defaults for unsupported models, then access the declared fields directly. This keeps a missing or misspelled field from silently disabling DeepStack replay. Please use the same explicit interface in the capability test.
There was a problem hiding this comment.
Addressed in d73b050.
Get rid of the two getattr fallbacks, and now use a small protocol (_SupportsDeepStackReplay) plus a resolver to get the slot width. Qwen3-VL implements get_num_deepstack_embeddings() instead of the class flag, and the capability test checks the same protocol.
However, this PR still did not fully solve the silent fail. If the method name is wrong on the model, the runner still gets 0 and no error; only the unit test catches it. We can make the runner raise in that case, but I may leave it for a follow-up to keep this PR small. Happy to add it here if you want.
Also, the implementation relies heavily on the model defining this method. Other models with DeepStack but without it will get no slot. That is fine today since they are not on the breakable allowlist, but if someone adds one later without the method, DeepStack will be dropped silently during replay.
The runner sized the replay slot from two getattr lookups with defaults (num_deepstack_embeddings, supports_bcg_deepstack_replay), so a renamed or misspelled field would have dropped the slot without an error. Replace both with a runtime-checkable protocol and a keyword-only resolver, the shape misc_utils already uses for resolve_aux_hidden_states_width: Qwen3VLForConditionalGeneration implements get_num_deepstack_embeddings() in place of the class flag, and _resolve_deepstack_replay_width returns 0 for a backend without an eager prefill tail or a model that does not implement the protocol. Behavior is unchanged; subclasses inherit the opt-in as before. The capability test asserts through the same protocol, and a new case pins the resolver's two zero branches.
TestQwen3VLCapabilityOptIn imported its four model classes inside the test method. Move them to the module's import block; the test body is unchanged.
|
[Codex]: As requested, run only TestQwen3VLServer.test_deepstack_grounding_hits_target_box on the current PR head (841be30), using the selective /rerun-test command. This tests Qwen3-VL-30B-A3B image grounding with the PR's default BCG prefill configuration. |
|
/rerun-test test/registered/vlm/test_vision_openai_server_a.py::TestQwen3VLServer.test_deepstack_grounding_hits_target_box |
|
🚀 |
|
[Codex]: Targeted CI passed on |
|
[Codex]: As requested, use /run-full-ci to trigger full CI (Base plus Extra) for the current PR head |
|
/run-full-ci |
bowenwan6
changed the title
S1 (experiments/slo_sched): plan, JITServe and SGLang background, the client and runners, session reports P1-P5 and the texts of PR-A (sgl-project/sglang#43967, merged) and PR-C (sgl-project/sglang#45643, open); PR-B was measured and dropped. M12: revalidation of sgl-project/sglang#33726 at ccfda8dfd9 on 2x H200. plan.md gains a status table (section 0) and the outcomes of both tracks; the README tracks are brought up to date.
Motivation
Prefill breakable CUDA graph (BCG) replay should preserve the same model semantics as eager execution.
Qwen3-VL passes intermediate visual features through
input_deepstack_embedsand adds them at selected language-model layers. However, BCG did not copy this request-specific tensor into a stable graph input slot.As a result, normal BCG execution behaved like an explicit zero-DeepStack run rather than eager execution. Requests still completed successfully, making this a silent correctness issue.
Reproduction
Environment:
Qwen/Qwen3-VL-8B-Instruct0c351dd01ed87e9c1b53cbc748cba10e6187ff3bmax_new_tokens=20The GPU test harness temporarily enabled Qwen3-VL for BCG testing.
The experiment compares four arms:
eager_normaleager_zero_deepstackbcg_normalbcg_zero_deepstackThe zero-DeepStack arms replace
input_deepstack_embedswith zeros immediately before the language-model forward call.Each result is formatted as:
bcg_normalvseager_normal(7, False, 1.154)(15, True, 0.071)bcg_normalvsbcg_zero_deepstack(20, True, 0.000)(7, False, 1.148)bcg_zero_deepstackvseager_zero_deepstack(20, True, 0.066)(20, True, 0.066)eager_zero_deepstackvseager_normal(7, False, 1.138)unchangedBefore the fix,
bcg_normalwas identical tobcg_zero_deepstackand diverged fromeager_normal.After the fix,
bcg_normalmatcheseager_normalat the token level and correctly diverges frombcg_zero_deepstack. The zero-DeepStack control remains unchanged.Instrumentation also confirmed that
input_deepstack_embedswas present and non-zero at the language-model entry before replay.The RGB image is used only as a deterministic semantic-equivalence test, not as a visual-quality benchmark.
Note on text-only numerics: the slot buffer changes the captured-graph layout, so text-only completions under BCG shift slightly relative to eager (divergence around char 78 in hardware testing, with the slot verified zero throughout). This is the same class as ordinary BCG-vs-eager drift, but anyone diffing text outputs across this PR should expect it.
Modifications
supports_bcg_deepstack_replaycapability to Qwen3-VL.input_deepstack_embedsbuffer and graph slot.Absence (
Noneor an empty tensor) clears the slot silently; a non-empty tensor with a mismatched shape, dtype, or device raises rather than silently downgrading the request to zero-DeepStack semantics.The
RuntimeErrorfires inside the replayed forward, so if the guard ever trips it takes down the scheduler rather than failing a single request. That is the intended trade for an unreachable contract violation, where the alternative is serving a silently wrong answer. Eager fallback is not available at this site:input_deepstack_embedsis only materialized insidemodel.forward(viageneral_mm_embed_routine→separate_deepstack_embeds), already past_prepare_forward_metadata_for_replayand inside the replay context, so there is no pre-commit validation point.input_embedscannot be reused because DeepStack has a different shape and is consumed at selected intermediate language-model layers.supports_bcg_deepstack_replayis inherited by every subclass ofQwen3VLForConditionalGeneration: Qwen3.5, Qwen3.5-MoE, InternS1-Pro, InternS2-Mobius, InternS2-Preview, and the Qwen3-Omni MoE Thinker. This is deliberate. Those archs consume DeepStack through the same code path, so any checkpoint of theirs that ships non-emptydeepstack_visual_indexesis covered without further change. Released checkpoints for those archs currently ship empty indexes, so the replay path is a no-op for them today.MuseGlimmerForConditionalGenerationis on the BCG allowlist but derives from a different base and does not consume DeepStack at all, so it is unaffected.Qwen3-VL and Qwen3-VL MoE now get prefill graphs by default. This changes replay-vs-eager numerics to the new baseline, and anyone who previously forced breakable prefill was silently running zero-DeepStack and should see grounding improve.
Accuracy Tests
Result:
This includes 8 new tests and 47 existing adjacent tests. Allowlist membership is covered separately in
test/registered/unit/server_args/test_server_args.py.The GPU four-arm validation results are shown in the Reproduction section.
Speed Tests and Profiling
This is a correctness fix, not a speed optimization.
Models that do not activate DeepStack replay do not allocate the new buffer or perform the replay copy.
For Qwen3-VL-8B with
max_num_tokens=8192,hidden_size=4096, three DeepStack embeddings, and BF16, the maximum additional buffer is:No replay-latency benchmark was performed.
Checklist
CI States
Latest PR Test (Base): ❌ Run #39644407713
Latest PR Test (Extra): ❌ Run #39644407677
Latest PR Test (AMD ROCm 10): ❌ Run #39644407708