Skip to content

[WIP] Add Log-Linear Attention #524

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 23, 2025
+1,319 −0
Merged

[WIP] Add Log-Linear Attention #524

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
01827c3
initial kernels
2022tgoel Jul 12, 2025
390aaad
Merge branch 'main' into main
sustcsonglin Jul 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions fla/ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from .hgrn import fused_recurrent_hgrn
from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn
from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn
from .log_linear_attn import chunk_log_linear_attn, LogLinearAttentionState
from .mesa_net import chunk_mesa_net
from .nsa import parallel_nsa
from .path_attn import parallel_path_attention
Expand All @@ -41,6 +42,7 @@
'fused_recurrent_hgrn',
'chunk_lightning_attn', 'fused_recurrent_lightning_attn',
'chunk_linear_attn', 'fused_chunk_linear_attn', 'fused_recurrent_linear_attn',
'chunk_log_linear_attn', 'LogLinearAttentionState',
'chunk_mesa_net',
'parallel_nsa',
'parallel_path_attention',
Expand Down
6 changes: 6 additions & 0 deletions fla/ops/log_linear_attn/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
from .chunk import chunk_log_linear_attn, LogLinearAttentionState

__all__ = [
'chunk_log_linear_attn',
'LogLinearAttentionState'
]
Comment on lines +1 to +6
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Fix formatting issues: import order and end-of-file.

The pipeline indicates two issues:

  1. isort would modify the import ordering
  2. Missing newline at end of file

Please run the pre-commit hooks to fix these issues.

 from .chunk import chunk_log_linear_attn, LogLinearAttentionState
 
 __all__ = [
     'chunk_log_linear_attn',
     'LogLinearAttentionState'
-]
+]
+
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
from .chunk import chunk_log_linear_attn, LogLinearAttentionState
__all__ = [
'chunk_log_linear_attn',
'LogLinearAttentionState'
]
from .chunk import chunk_log_linear_attn, LogLinearAttentionState
__all__ = [
'chunk_log_linear_attn',
'LogLinearAttentionState'
]
🧰 Tools
🪛 GitHub Actions: lint

[error] 1-1: Pre-commit hook 'end-of-file-fixer' modified this file to fix end of file issues.

[error] 1-1: Pre-commit hook 'isort' modified this file to fix import sorting.

🤖 Prompt for AI Agents 
In fla/ops/log_linear_attn/__init__.py at lines 1 to 6, fix the import order
according to isort standards by running the isort pre-commit hook, and add a
newline at the end of the file to comply with formatting requirements. Run the
full set of pre-commit hooks to automatically correct these formatting issues.

Loading
Loading