add-choose-vit-backend #1191
add-choose-vit-backend #1191
Conversation
Summary of ChangesHello @SangChengC, 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 significantly enhances the LightLLM framework by introducing a flexible and extensible architecture for Vision Transformer (ViT) attention computations. It abstracts the underlying attention implementation, enabling the dynamic selection of optimized backends like FA3 or Triton. This change not only improves the modularity of the system but also paves the way for better performance and easier integration of future ViT attention optimizations within visual models. 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 a mechanism to select different backends for Vision Transformer (ViT) attention, similar to the existing implementation for language model attention. The changes include adding base classes and selection utilities for ViT attention backends, and implementing fa3 and triton backends.
The overall approach is good, but I've found several issues that need to be addressed:
- There are critical implementation errors in the Triton backend (
triton/fp.py) that would cause runtimeTypeErrors due to incorrect method signatures and not implementing the required abstract methods. - There are type hint inconsistencies in the FA3 backend (
fa3/fp.py). - Some debugging
printstatements have been left in the code. - A comment in a non-English language suggests a temporary solution that should be properly documented with a TODO.
I've left specific comments with suggestions to fix these issues. After these are addressed, the PR should be in good shape.
| @dataclasses.dataclass | ||
| class TritonVitAttState(BaseVitAttState): | ||
| def init_state(self): | ||
| pass | ||
|
|
||
| def _vit_att( | ||
| self, | ||
| q: torch.Tensor, | ||
| k: torch.Tensor, | ||
| v: torch.Tensor, | ||
| o: torch.Tensor, | ||
| cu_seqlens: torch.Tensor, | ||
| max_seqlen: int, | ||
| alloc_func=torch.empty, | ||
| ): | ||
| from lightllm.models.vit.triton_kernel.flashattention_nopad import _flash_attention_triton_fwd | ||
|
|
||
| _flash_attention_triton_fwd( | ||
| q, | ||
| k, | ||
| v, | ||
| o, | ||
| cu_seqlens, # q k v cu_seqlens, | ||
| max_seqlen, | ||
| ) | ||
| return |
There was a problem hiding this comment.
The class TritonVitAttState does not correctly implement the BaseVitAttState interface, which will lead to runtime errors.
- The abstract method
vit_attis not implemented. A private method_vit_attis implemented instead. This will raise aTypeErroron instantiation. - The method should return the output tensor
oto be consistent with the base class. - The
alloc_funcparameter is unused and can be removed. - The method is missing a return type hint.
@dataclasses.dataclass
class TritonVitAttState(BaseVitAttState):
def init_state(self):
pass
def vit_att(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
o: torch.Tensor,
cu_seqlens: torch.Tensor,
max_seqlen: int,
) -> torch.Tensor:
from lightllm.models.vit.triton_kernel.flashattention_nopad import _flash_attention_triton_fwd
_flash_attention_triton_fwd(
q,
k,
v,
o,
cu_seqlens, # q k v cu_seqlens,
max_seqlen,
)
return o| o: torch.Tensor, | ||
| cu_seqlens: torch.Tensor, | ||
| max_seqlen: int, | ||
| ) -> None: |
There was a problem hiding this comment.
The return type hint for vit_att is None, but the function returns the tensor o. This should be corrected to torch.Tensor to match the base class BaseVitAttState.
| ) -> None: | |
| ) -> torch.Tensor: |
| for backend_name in priority_list: | ||
| if validate(backend_name): | ||
| logger.info(f"Auto-selected {backend_name} backend (validated)") | ||
| print(f"llm_dtype is {llm_dtype}, backend_name is {backend_name} ") |
There was a problem hiding this comment.
This print statement appears to be for debugging and should be removed from the final code.
| for backend_name in priority_list: | ||
| if validate(backend_name): | ||
| logger.info(f"Auto-selected {backend_name} backend (validated) for ViT") | ||
| print(f"llm_dtype is {llm_dtype}, backend_name is {backend_name} ") |
There was a problem hiding this comment.
This print statement appears to be for debugging and should be removed from the final code.
| cu_seqlens: torch.Tensor, | ||
| max_seqlen: int, | ||
| ) -> None: | ||
| self.backend: Fa3ViTAttBackend = self.backend # for typing |
There was a problem hiding this comment.
This line is redundant for type hinting, as the backend attribute is already typed in the class definition on line 257. It can be removed for cleaner code.
| self.cache_client._channel.stream.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) | ||
| self.data_type = kvargs["data_type"] | ||
|
|
||
| set_global_rank(kvargs["tp_rank_id"]) # 这里看后面怎么改 |
There was a problem hiding this comment.
The comment # 这里看后面怎么改 ("see how to change this later") suggests this line is a temporary solution. It's better to use a logger.warning or add a TODO comment with more context for future developers. Also, comments should be in English for better maintainability.
| set_global_rank(kvargs["tp_rank_id"]) # 这里看后面怎么改 | |
| # TODO: Revisit the global rank setting for the visual model process. | |
| set_global_rank(kvargs["tp_rank_id"]) |
No description provided.