Longcat platform#1214
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces tensor parallel support for Flux2 models, including weight/bias splitting and rank-0 broadcast mechanisms, and adds sequence parallel support for LongCat Image models along with platform-specific attention fallbacks for Hunyuan Video. Feedback on these changes highlights a critical shape mismatch bug in the NPU attention implementation, potential AttributeError exceptions when calling empty_cache on CPU devices, and several potential KeyErrors from direct dictionary access of seq_parallel in configuration objects.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| out = torch.zeros(batch * seqlen, heads * dim, dtype=out_unpad.dtype, device=out_unpad.device) | ||
| out[indices] = out_unpad | ||
| return out.reshape(batch, seqlen, heads, dim) |
There was a problem hiding this comment.
There is a shape mismatch bug when assigning out_unpad to out[indices]. out_unpad has a 3D shape of (total_unpad_tokens, heads, dim), but out is initialized as a 2D tensor of shape (batch * seqlen, heads * dim). This will raise a RuntimeError at runtime due to the dimension mismatch. To fix this, initialize out as a 3D tensor of shape (batch * seqlen, heads, dim) so that the slice assignment matches the shape of out_unpad.
| out = torch.zeros(batch * seqlen, heads * dim, dtype=out_unpad.dtype, device=out_unpad.device) | |
| out[indices] = out_unpad | |
| return out.reshape(batch, seqlen, heads, dim) | |
| out = torch.zeros(batch * seqlen, heads, dim, dtype=out_unpad.dtype, device=out_unpad.device) | |
| out[indices] = out_unpad | |
| return out.reshape(batch, seqlen, heads, dim) |
| if self.text_encoders is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | ||
| del self.text_encoders | ||
| self.text_encoders = None | ||
| torch_device_module.empty_cache() | ||
| gc.collect() |
There was a problem hiding this comment.
If AI_DEVICE is set to "cpu", torch_device_module will be torch.cpu, which does not have an empty_cache attribute. This will cause an AttributeError at runtime. It is safer to check if torch_device_module has the empty_cache attribute before calling it.
| if self.text_encoders is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | |
| del self.text_encoders | |
| self.text_encoders = None | |
| torch_device_module.empty_cache() | |
| gc.collect() | |
| if self.text_encoders is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | |
| del self.text_encoders | |
| self.text_encoders = None | |
| if hasattr(torch_device_module, "empty_cache"): | |
| torch_device_module.empty_cache() | |
| gc.collect() |
| if self.vae is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | ||
| del self.vae | ||
| self.vae = None | ||
| torch_device_module.empty_cache() | ||
| gc.collect() |
There was a problem hiding this comment.
If AI_DEVICE is set to "cpu", torch_device_module will be torch.cpu, which does not have an empty_cache attribute. This will cause an AttributeError at runtime. It is safer to check if torch_device_module has the empty_cache attribute before calling it.
| if self.vae is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | |
| del self.vae | |
| self.vae = None | |
| torch_device_module.empty_cache() | |
| gc.collect() | |
| if self.vae is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): | |
| del self.vae | |
| self.vae = None | |
| if hasattr(torch_device_module, "empty_cache"): | |
| torch_device_module.empty_cache() | |
| gc.collect() |
| max_seqlen_kv=total_len, | ||
| model_cls="longcat_image", | ||
| ) | ||
| if self.config["seq_parallel"]: |
There was a problem hiding this comment.
Accessing seq_parallel directly via self.config["seq_parallel"] will raise a KeyError if the key is not present in the configuration dictionary. It is safer to use self.config.get("seq_parallel", False) to handle cases where the key is missing.
| if self.config["seq_parallel"]: | |
| if self.config.get("seq_parallel", False): |
| max_seqlen_kv=total_len, | ||
| model_cls="longcat_image", | ||
| ) | ||
| if self.config["seq_parallel"]: |
There was a problem hiding this comment.
Accessing seq_parallel directly via self.config["seq_parallel"] will raise a KeyError if the key is not present in the configuration dictionary. It is safer to use self.config.get("seq_parallel", False) to handle cases where the key is missing.
| if self.config["seq_parallel"]: | |
| if self.config.get("seq_parallel", False): |
| encoder_hidden_states=prompt_embeds, | ||
| ) | ||
|
|
||
| if self.config["seq_parallel"]: |
There was a problem hiding this comment.
Accessing seq_parallel directly via self.config["seq_parallel"] will raise a KeyError if the key is not present in the configuration dictionary. It is safer to use self.config.get("seq_parallel", False) to handle cases where the key is missing.
| if self.config["seq_parallel"]: | |
| if self.config.get("seq_parallel", False): |
|
|
||
| noise_pred = self.post_infer.infer(self.post_weight, hidden_states, pre_infer_out.temb) | ||
|
|
||
| if self.config["seq_parallel"]: |
There was a problem hiding this comment.
Accessing seq_parallel directly via self.config["seq_parallel"] will raise a KeyError if the key is not present in the configuration dictionary. It is safer to use self.config.get("seq_parallel", False) to handle cases where the key is missing.
| if self.config["seq_parallel"]: | |
| if self.config.get("seq_parallel", False): |
|
|
||
| # Attention calculation module | ||
| self.add_module("calculate", ATTN_WEIGHT_REGISTER[self.attn_type]()) | ||
| if self.config["seq_parallel"]: |
There was a problem hiding this comment.
Accessing seq_parallel directly via self.config["seq_parallel"] will raise a KeyError if the key is not present in the configuration dictionary. It is safer to use self.config.get("seq_parallel", False) to handle cases where the key is missing.
| if self.config["seq_parallel"]: | |
| if self.config.get("seq_parallel", False): |
|
|
||
| # Attention calculation module | ||
| self.add_module("calculate", ATTN_WEIGHT_REGISTER[self.attn_type]()) | ||
| if self.config["seq_parallel"]: |
There was a problem hiding this comment.
Accessing seq_parallel directly via self.config["seq_parallel"] will raise a KeyError if the key is not present in the configuration dictionary. It is safer to use self.config.get("seq_parallel", False) to handle cases where the key is missing.
| if self.config["seq_parallel"]: | |
| if self.config.get("seq_parallel", False): |
This reverts commit e00a254.
1 participant
No description provided.