[BugFix][HunyuanImage3] Set MRoPE dynamic_arg_dims so graph mode can compile (#3630)
Signed-off-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
1 file changed
+10
-1
@@ -96,7 +96,16 @@ from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjecto
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@support_torch_compile
|
||||
@support_torch_compile(
|
||||
dynamic_arg_dims={
|
||||
"input_ids": 0,
|
||||
# positions is (3, num_tokens) for mRoPE, so the token dimension is
|
||||
# the last dimension rather than dim 0.
|
||||
"positions": -1,
|
||||
"intermediate_tensors": 0,
|
||||
"inputs_embeds": 0,
|
||||
}
|
||||
)
|
||||
class HunyuanModel(HunYuanModel):
|
||||
def _split_qkv_weight(self, qkv: torch.Tensor):
|
||||
num_attention_heads = self.config.num_attention_heads
|
||||
|
||||
Reference in new issue
Block a user