fix coderabbit comments

belgarten-nv · belgarten-nv · commit 5475eccdd17c · 2026-06-02T22:04:17.000Z
Signed-off-by: Brenden Elgarten &lt;belgarten@nvidia.com&gt;
diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py b/tensorrt_llm/_torch/visual_gen/models/flux/joint_proj.py
@@ -257,7 +257,7 @@ def range_size(r):
                     "up": (local_mlp_hidden_start, local_mlp_hidden_end),
                 },
             )
-            self.local_qkv_dim = (q_dim + 2 * kv_dim) // self.tp_size
+            self.local_qkv_dim = local_q_dim + 2 * local_kv_dim
             self.local_mlp_dim = local_mlp_hidden_size
 
     def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

Original file line number	Diff line number	Diff line change
`@@ -257,7 +257,7 @@ def range_size(r):`
`257`	`257`	`"up": (local_mlp_hidden_start, local_mlp_hidden_end),`
`258`	`258`	`},`
`259`	`259`	`)`
`260`		`- self.local_qkv_dim = (q_dim + 2 * kv_dim) // self.tp_size`
	`260`	`+ self.local_qkv_dim = local_q_dim + 2 * local_kv_dim`
`261`	`261`	`self.local_mlp_dim = local_mlp_hidden_size`
`262`	`262`
`263`	`263`	`def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:`