Skip to content

fix(rollout): reject shared prefill with continuous batching - #8583

Open
modelpath-dev wants to merge 2 commits into
deepspeedai:masterfrom
modelpath-dev:fix/8458-reject-shared-prefill-continuous
Open

modelpath-dev wants to merge 2 commits into
deepspeedai:masterfrom
modelpath-dev:fix/8458-reject-shared-prefill-continuous

Conversation

@modelpath-dev

@modelpath-dev modelpath-dev commented Sep 18, 2026

Copy link
Copy Markdown

Why

HybridEngineRollout.generate selects the continuous path whenever SamplingConfig.continuous_batch_size is set, before any shared-prefill logic runs. With HybridEngineRolloutConfig.use_shared_prefill=True, that flag was therefore ignored with no error. This PR rejects the combination explicitly, matching the approach already taken for CUDA graph capture on the continuous path.

Fixes #8458.

Scope

  • HybridEngineRollout._validate_continuous_inputs: raise ValueError when use_shared_prefill is enabled.
  • Unit test for the incompatible pair.
  • docs/code-docs/source/inference-engine.rst: document that shared prefill and continuous batching cannot be combined.

Out of scope: implementing shared-prefill semantics inside continuous batching.

Tradeoffs

Rejection over a compatible shared-prefill continuous mode. Continuous batching already requires n_samples_per_prompt=1, so shared prefill has no clear benefit there today.

Blast Radius

Only the experimental continuous path changes, and only when shared prefill is also requested. Default generate() and shared-prefill-only configs are unchanged.

Verification

Before the fix, _validate_continuous_inputs accepted use_shared_prefill=True with continuous_batch_size=1 (silent ignore). After the fix it raises ValueError: continuous batching does not support shared prompt prefill.

TRANSFORMERS_NO_TF=1 PYTHONPATH=. python3 -m pytest \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py::test_continuous_generation_rejects_shared_prefill \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py::test_continuous_generation_rejects_unsupported_inputs \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py::test_shared_prefill_rejects_graph_capture \
  tests/unit/runtime/rollout/test_hybrid_engine_rollout.py::test_generate_uses_shared_prefill_for_multiple_samples \
  -q
# 4 passed

modelpath-dev and others added 2 commits September 18, 2026 13:20
Add a failing unit test for DeepSpeed deepspeedai#8458 before the validator change.

Co-authored-by: Cursor <cursoragent@cursor.com>
Raise ValueError when HybridEngineRolloutConfig.use_shared_prefill is set
together with SamplingConfig.continuous_batch_size, instead of selecting
the continuous path and silently ignoring shared prefill. Document the
incompatibility in the inference-engine notes.

Fixes deepspeedai#8458.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Do not silently ignore use_shared_prefill with continuous batching

1 participant