From 11829ccd75e67167254249d149d3ab01c86ff25a Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Mon, 15 Jun 2026 19:40:06 +0800 Subject: [PATCH] [Bugfix] Fix LTX-2.3 tensor-parallel gated attention (#4439) Signed-off-by: mglyn <1203789601@qq.com> --- vllm_omni/diffusion/models/ltx2/ltx2_transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 4e5659c54..dcf848390 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -691,7 +691,7 @@ class LTX2Attention(torch.nn.Module): # LTX-2.3: per-head gated attention # leave unquantized for this linear if apply_gated_attention: - self.to_gate_logits = nn.Linear(query_dim, heads, bias=True) + self.to_gate_logits = nn.Linear(query_dim, self.query_num_heads, bias=True) else: self.to_gate_logits = None