switch GDN to use attn-gym kernels for titanRL - #4363
Open
liangel-02 wants to merge 1 commit into
Open
Conversation
meta-cla
Bot
added
the
CLA Signed
liangel-02
commented
Aug 28, 2026
liangel-02
commented
Aug 28, 2026
liangel-02
force-pushed
the
gdn_attn_gym
branch
from
August 28, 2026 03:42
94fbc7d to
653302e
Compare
Contributor
|
question: does this also work on gb gpus only? |
Contributor
Author
|
@shuhuayu only blackwell |
liangel-02
force-pushed
the
gdn_attn_gym
branch
from
August 28, 2026 16:23
653302e to
71feff5
Compare
liangel-02
changed the title
liangel-02
marked this pull request as ready for review
August 28, 2026 16:46
liangel-02
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 28, 2026 16:46
Contributor
|
B200 and B300 to be grace or non grace |
wwwjn
reviewed
Aug 28, 2026
liangel-02
force-pushed
the
gdn_attn_gym
branch
2 times, most recently
from
August 28, 2026 21:20
ed38e0d to
15912af
Compare
drisspg
reviewed
Aug 28, 2026
drisspg
reviewed
Aug 28, 2026
liangel-02
force-pushed
the
gdn_attn_gym
branch
3 times, most recently
from
August 28, 2026 22:27
9c2b1d8 to
ed345f4
Compare
liangel-02
force-pushed
the
gdn_attn_gym
branch
from
August 29, 2026 00:16
ed345f4 to
c04cd02
Compare
liangel-02
force-pushed
the
gdn_attn_gym
branch
from
August 29, 2026 03:28
c04cd02 to
fb0ad75
Compare
tianyu-l
reviewed
Aug 29, 2026
| x=x_TD.unsqueeze(0), | ||
| weight=weight.squeeze(1), | ||
| bias=None, | ||
| del cu_seqlens_cpu |
Contributor
There was a problem hiding this comment.
with attn-gym, do we no longer need this any more? If so let's delete the code that generates it in get_attention_masks
Comment on lines
+18
to
+22
| from attn_gym.linear import ( | ||
| causal_conv1d as _attn_gym_causal_conv1d, | ||
| l2norm as _attn_gym_l2norm, | ||
| recurrent_gdn as _attn_gym_recurrent_gdn, | ||
| ) |
Contributor
There was a problem hiding this comment.
I'd be happy with removing _attn_gym_ prefix
| "the paged convolution history has contiguous channels." | ||
| ) | ||
|
|
||
| self.cache_config.mamba_ssm_cache_dtype = "float32" |
Contributor
There was a problem hiding this comment.
should this be fp32 regardless of batch invariant mode?
Contributor
There was a problem hiding this comment.
without looking into details, I had expected that with attn-gym we can make the code shorter (e.g. we don't need to copy the state cache out -> computation -> put the result back in cache), but it doesn't look like so?
Contributor
There was a problem hiding this comment.
this is because we are stilll supporting prefill(chunked) through fla for gdn now; im working on fixing this
This was referenced Sep 1, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
performance improvements:
