[BCG] Restore Qwen3.5 MRoPE fusion under breakable CUDA graph (#27918)
This commit is contained in:
@@ -458,9 +458,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
),
|
),
|
||||||
):
|
):
|
||||||
if self.layer_model is not None:
|
if self.layer_model is not None:
|
||||||
|
positions = self._get_layer_model_positions(forward_batch)
|
||||||
return self.layer_model.forward(
|
return self.layer_model.forward(
|
||||||
forward_batch.input_ids,
|
forward_batch.input_ids,
|
||||||
forward_batch.positions,
|
positions,
|
||||||
forward_batch,
|
forward_batch,
|
||||||
forward_batch.input_embeds,
|
forward_batch.input_embeds,
|
||||||
)
|
)
|
||||||
@@ -470,6 +471,21 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
forward_batch,
|
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:
|
def _run_dummy_forward(self, num_tokens: int) -> None:
|
||||||
"""Build a dummy ForwardBatch at this shape, init attn metadata,
|
"""Build a dummy ForwardBatch at this shape, init attn metadata,
|
||||||
run forward once. Used by TcPiecewiseCudaGraphBackend.prepare
|
run forward once. Used by TcPiecewiseCudaGraphBackend.prepare
|
||||||
|
|||||||
Reference in New Issue
Block a user