Skip to content

[Copy to main] fix(ssm): handle alignment padding in GDN packed seq + CP - #4289

Closed
svcnvidia-nemo-ci wants to merge 1 commit into
mainfrom
copy-pr-4230-to-main
Closed

[Copy to main] fix(ssm): handle alignment padding in GDN packed seq + CP#4289
svcnvidia-nemo-ci wants to merge 1 commit into
mainfrom
copy-pr-4230-to-main

Conversation

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🤖 This PR was auto-generated by Claude via the /claude copy command.

Cherry-picked from #4230.


What does this PR do ?

Fix shape mismatch in GDN packed sequence + CP > 1 when alignment padding is present.

Summary

When cu_seqlens_q_padded is not provided, GDN falls back to cu_seqlens_q (actual token boundaries) for the per-sequence all-to-all unpack. Alignment padding tokens are lost, making the output shorter than hidden_states, crashing at the residual connection.

Fix: Extract cu_seqlens resolution into _resolve_cu_seqlens():

  • cu_seqlens_padded available → use it (unchanged behavior)
  • CP = 1 → use actual (no all-to-all, no issue)
  • Single sequence + padding → auto-extend to cover trailing padding
  • Multi sequence + padding → raise ValueError requesting cu_seqlens_padded

Test results

Unit tests: 5/5 passed

(2× H100, CP=2):

  • Single seq + padding + no cu_seqlens_padded → auto-extend, output shape correct
  • Multi seq + padding + no cu_seqlens_paddedValueError raised correctly
  • With cu_seqlens_padded set → standard path works
  • No padding → passes through unchanged

Fixes #4194

@copy-pr-bot

copy-pr-bot Bot commented Apr 14, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions

Copy link
Copy Markdown
Contributor

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

  1. Add the oncall reviewer (optional reviewer)
  2. Add required review teams based on your changes

See the contribution guide for more details.

@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft April 14, 2026 01:04
@Phlip79 Phlip79 changed the title [Copy to main] [dev] fix(ssm): handle alignment padding in GDN packed seq + CP [Copy to main] fix(ssm): handle alignment padding in GDN packed seq + CP Apr 14, 2026
@yuzhongw-nvidia

Copy link
Copy Markdown
Contributor

This fix could be involved in #2645. Close this one.

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.

GDN packed sequence (THD format) + CP + SP causes shape mismatch at residual connection.

2 participants