Support draft attention #673
Conversation
Summary of ChangesHello @Musisoul, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request introduces a "draft attention" mechanism aimed at enhancing the computational efficiency of transformer models, particularly for processing spatial data like video frames. It achieves this by employing a sparse attention strategy that samples query-key pairs and applies a percentile-based mask to concentrate attention on the most significant interactions. This new attention type is seamlessly integrated into the transformer's self-attention layers, offering flexible control over its sparsity through model configuration. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![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 support for a new sparse attention mechanism called "draft attention". The core logic is implemented in the new file lightx2v/common/ops/attn/draft_attn.py, with integrations into the model's inference and weight loading parts. My review focuses on correctness, portability, and maintainability. I've identified some critical issues, primarily hardcoded cuda device strings that will prevent the code from running on other hardware like CPU or MPS. I've also included several medium-severity suggestions to improve code clarity, reduce duplication, and adhere to best practices.
| try: | ||
| from flash_attn.flash_attn_interface import flash_attn_varlen_func as _func | ||
|
|
||
| flash_attn_varlen_func = _func | ||
| except ImportError: | ||
| logger.info("flash_attn_varlen_func not found, please install flash_attn2 first") | ||
|
|
||
|
|
||
| try: | ||
| from flash_attn_interface import flash_attn_varlen_func as _func | ||
|
|
||
| flash_attn_varlen_func = _func | ||
| except ImportError: | ||
| logger.info("flash_attn_varlen_func_v3 not found, please install flash_attn3 first") |
There was a problem hiding this comment.
The current logic for importing flash_attn_varlen_func can be confusing. It attempts to import from flash_attn (v2) and then from flash_attn_interface (v3), overwriting the v2 function if v3 is also present. A clearer pattern is to try importing the preferred version (v3) first, and fall back to the older version (v2) if the import fails. This makes the preference explicit and improves readability. For example:
flash_attn_varlen_func = None
try:
# Prefer flash_attn v3
from flash_attn_interface import flash_attn_varlen_func as _func
flash_attn_varlen_func = _func
except ImportError:
logger.info("flash_attn_varlen_func_v3 not found, trying to import from flash_attn2.")
try:
from flash_attn.flash_attn_interface import flash_attn_varlen_func as _func
flash_attn_varlen_func = _func
except ImportError:
logger.info("flash_attn_varlen_func not found in flash_attn2 either. Please install flash-attn v2 or v3.")| def __init__(self): | ||
| self.config = {} |
There was a problem hiding this comment.
This class inherits from AttnWeightTemplate, but its __init__ method does not call super().__init__(). This breaks the inheritance chain and can lead to subtle bugs if the base class initialization is important. Please add a call to super().__init__(). Note that you may need to adjust the call to match the base class constructor's signature.
| if self.config["self_attn_1_type"] == "draft_attn": | ||
| attn_out = phase.self_attn_1.apply( | ||
| q=q, | ||
| k=k, | ||
| v=v, | ||
| cu_seqlens_q=self.self_attn_cu_seqlens_qkv, | ||
| cu_seqlens_kv=self.self_attn_cu_seqlens_qkv, | ||
| max_seqlen_q=img_qkv_len, | ||
| max_seqlen_kv=img_qkv_len, | ||
| frame_h=self.scheduler.latents.shape[2] // self.scheduler.patch_size[1], | ||
| frame_w=self.scheduler.latents.shape[3] // self.scheduler.patch_size[2], | ||
| block_idx=self.block_idx, | ||
| ) | ||
| else: | ||
| attn_out = phase.self_attn_1.apply( | ||
| q=q, | ||
| k=k, | ||
| v=v, | ||
| cu_seqlens_q=self.self_attn_cu_seqlens_qkv, | ||
| cu_seqlens_kv=self.self_attn_cu_seqlens_qkv, | ||
| max_seqlen_q=img_qkv_len, | ||
| max_seqlen_kv=img_qkv_len, | ||
| model_cls=self.config["model_cls"], | ||
| ) |
There was a problem hiding this comment.
There's significant code duplication in the if/else block for handling different attention types. This can be refactored to improve readability and maintainability. Consider creating a dictionary with the common arguments and then conditionally adding the specific ones for each attention type before making the apply call.
attn_kwargs = {
"q": q,
"k": k,
"v": v,
"cu_seqlens_q": self.self_attn_cu_seqlens_qkv,
"cu_seqlens_kv": self.self_attn_cu_seqlens_qkv,
"max_seqlen_q": img_qkv_len,
"max_seqlen_kv": img_qkv_len,
}
if self.config["self_attn_1_type"] == "draft_attn":
attn_kwargs["frame_h"] = self.scheduler.latents.shape[2] // self.scheduler.patch_size[1]
attn_kwargs["frame_w"] = self.scheduler.latents.shape[3] // self.scheduler.patch_size[2]
attn_kwargs["block_idx"] = self.block_idx
else:
attn_kwargs["model_cls"] = self.config["model_cls"]
attn_out = phase.self_attn_1.apply(**attn_kwargs)
No description provided.