Skip to content

Comments

lse support#7

Merged
divchenko merged 1 commit intofwaifrom
lse_support
Nov 15, 2025
Merged

lse support#7
divchenko merged 1 commit intofwaifrom
lse_support

Conversation

@divchenko
Copy link

@divchenko divchenko commented Nov 14, 2025

Add LSE suport in all attention kernels.
Allow strides for retuned LSEs.

@divchenko divchenko changed the base branch from main to fwai November 14, 2025 20:32
enable_pdl: Optional[bool] = None,
sinks: Optional[List[torch.Tensor]] = None,
kv_cache_scales: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
return_lse: bool = False,
Copy link

Choose a reason for hiding this comment

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

Gate return lse just based on if set to None or not?

Copy link
Author

Choose a reason for hiding this comment

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

they already have this API for ragged_attention, I'm just making it consistent

@divchenko divchenko merged commit 2999167 into fwai Nov 15, 2025
1 of 2 checks passed
divchenko added a commit that referenced this pull request Nov 19, 2025
Add LSE suport in all attention kernels
divchenko added a commit that referenced this pull request Dec 6, 2025
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Dec 15, 2025
Add LSE suport in all attention kernels
divchenko added a commit that referenced this pull request Dec 17, 2025
Add LSE suport in all attention kernels
divchenko added a commit that referenced this pull request Dec 18, 2025
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 14, 2026
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 16, 2026
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 26, 2026
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 26, 2026
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 26, 2026
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 26, 2026
Add LSE suport in all attention kernels
murphymatt pushed a commit that referenced this pull request Jan 27, 2026
Add LSE suport in all attention kernels
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants