Skip to content

fix(dsa): fall back for unsupported odd cuDNN top-k - #19

Merged
pstefa1707 merged 1 commit into
trainers-mainfrom
jackrao/fix-cudnn-odd-topk
Jul 17, 2026
Merged

fix(dsa): fall back for unsupported odd cuDNN top-k#19
pstefa1707 merged 1 commit into
trainers-mainfrom
jackrao/fix-cudnn-odd-topk

Conversation

@JackRao123

@JackRao123 JackRao123 commented Jul 17, 2026

Copy link
Copy Markdown

Summary

  • route odd indexer top-K selection through masked torch.topk to avoid cuDNN Frontend 1.26.0's vector-width assertion
  • retain cuDNN score generation and FlashMLA sparse attention, and keep even K on the fused cuDNN selector
  • cover the CP32 regression shape, varlen masking, score returns, and even-K dispatch

Test plan

  • Run Ruff checks on the changed files
  • Compile both changed Python files
  • Focused pytest passes on the B200 devbox (4 passed)
  • Exact [235, 705], top_k=705 patched primitive passes on B200
  • Repeat the exact CP32 15,000-token trainer forward after building the patched image

Keep fused score generation and sparse attention while selecting odd top-k values with PyTorch, avoiding cuDNN Frontend's vector-width assertion for packed CP segments.

Co-authored-by: Cursor <cursoragent@cursor.com>
topk_scores.masked_fill_(~valid, float("-inf"))
else:
topk_scores = None
return {"indices": topk_indices, "values": topk_scores}

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

This will b v slow if its raw pytorch right? When do we hit this case?
Should we have a triton kernel to cover it?

@pstefa1707
pstefa1707 merged commit d3932e7 into trainers-main Jul 17, 2026
2 checks passed
@pstefa1707
pstefa1707 deleted the jackrao/fix-cudnn-odd-topk branch July 17, 2026 02:57
jerryhong21 pushed a commit that referenced this pull request Aug 20, 2026
Keep fused score generation and sparse attention while selecting odd top-k values with PyTorch, avoiding cuDNN Frontend's vector-width assertion for packed CP segments.

Co-authored-by: Cursor <cursoragent@cursor.com>
jerryhong21 pushed a commit that referenced this pull request Aug 20, 2026
Keep fused score generation and sparse attention while selecting odd top-k values with PyTorch, avoiding cuDNN Frontend's vector-width assertion for packed CP segments.

Co-authored-by: Cursor <cursoragent@cursor.com>
jerryhong21 pushed a commit that referenced this pull request Aug 21, 2026
Keep fused score generation and sparse attention while selecting odd top-k values with PyTorch, avoiding cuDNN Frontend's vector-width assertion for packed CP segments.

Co-authored-by: Cursor <cursoragent@cursor.com>
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