Skip to content

LeanAttention code modularization #765

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

Jump to bottom
Merged
valechen merged 7 commits into from
Aug 6, 2025
Merged

LeanAttention code modularization #765

valechen merged 7 commits into from
Aug 6, 2025

Conversation

@valechen
Copy link
Contributor

@valechen valechen commented Aug 4, 2025 •
edited
Loading

  1. Move "calculate_max_output_tiles_analytically" to lean_atten.py
  2. Combine "get_num_splits_and_buffer_sizes" and "get_lean_attention_params"
  3. Create _attention_inner() @jit for inner loop attention calculation
  4. Add handling for total_tiles < num_SMs case

@valechen valechen requested a review from rahulbatra85 August 4, 2025 22:28
@valechen valechen requested a review from vgokhale August 5, 2025 16:26
@valechen valechen merged commit 86a6fd2 into main Aug 6, 2025
14 checks passed
@valechen valechen deleted the la_opt2 branch August 6, 2025 16:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@vgokhale vgokhale vgokhale approved these changes

@rahulbatra85 rahulbatra85 Awaiting requested review from rahulbatra85

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@valechen @vgokhale