[BCG] Restore Qwen3.5 MRoPE fusion under breakable CUDA graph (#27918)

This commit is contained in:
Kaixi
2026-07-09 14:30:05 -07:00
committed by GitHub
parent bda1dc0d95
commit 10bb2eff3d
@@ -458,9 +458,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
),
):
if self.layer_model is not None:
positions = self._get_layer_model_positions(forward_batch)
return self.layer_model.forward(
forward_batch.input_ids,
forward_batch.positions,
positions,
forward_batch,
forward_batch.input_embeds,
)
@@ -470,6 +471,21 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
forward_batch,
)
def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tensor:
"""Mirror outer multimodal wrappers when BCG captures layer_model directly."""
if forward_batch.mrope_positions is None:
return forward_batch.positions
model = self.model_runner.model
if getattr(model, "is_mrope_enabled", False):
return forward_batch.mrope_positions
language_model = getattr(model, "language_model", None)
if getattr(language_model, "is_mrope_enabled", False):
return forward_batch.mrope_positions
return forward_batch.positions
def _run_dummy_forward(self, num_tokens: int) -> None:
"""Build a dummy ForwardBatch at this shape, init attn metadata,
run forward once. Used by TcPiecewiseCudaGraphBackend.prepare