Skip to content

Add ragged_conv1d for ragged chunked kernel integration for gdn attention#2099

Merged
helloworld1 merged 3 commits intomainfrom
sharding_issue_6
Apr 1, 2026
Merged

Add ragged_conv1d for ragged chunked kernel integration for gdn attention#2099
helloworld1 merged 3 commits intomainfrom
sharding_issue_6

Conversation

@helloworld1
Copy link
Copy Markdown
Collaborator

@helloworld1 helloworld1 commented Mar 31, 2026

Description

Add ragged_conv1d for ragged chunked kernel integration for gdn attention

  • Add ragged_conv1d
  • Add interface for ragged gated delta rule
  • Use reorder_concatenated_tensor_for_sharding for sharding conversion to simplify code
  • Updated argument comments to make them clear

Tests

Tested quality and performance of qwen 3.5 model. At gives right answers in examples/offline_inference.py
Tested numerical on par with previous implementation

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@github-actions
Copy link
Copy Markdown

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

  • why is this change being made,
  • the problem being solved and any relevant context,
  • why this is a good solution,
  • some information about the specific implementation,
  • shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@helloworld1 helloworld1 force-pushed the sharding_issue_6 branch 2 times, most recently from 3cbd2ea to ac431bd Compare April 1, 2026 02:32
@helloworld1 helloworld1 marked this pull request as ready for review April 1, 2026 02:33
@helloworld1 helloworld1 added the ready ONLY add when PR is ready to merge/full CI is needed label Apr 1, 2026
Signed-off-by: Jacob Platin <jacobplatin@google.com>
Signed-off-by: Haowen Ning <hning@google.com>
@helloworld1 helloworld1 merged commit 2e7f853 into main Apr 1, 2026
50 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants