Skip to content

Fix ragged all-to-all with ragged buffer factor in DeepSeek-V3.#4397

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_944752214
Open

Fix ragged all-to-all with ragged buffer factor in DeepSeek-V3.#4397
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_944752214

Conversation

@copybara-service

Copy link
Copy Markdown
Contributor

Fix ragged all-to-all with ragged buffer factor in DeepSeek-V3.

  • Implement traffic matrix truncation in forward/backward all-to-all passes when ragged_buffer_factor > 0 to prevent out-of-bounds writes in ragged_all_to_all.
  • Partition GMM v2 VMEM scratchpad (partial_out_ref) by N program ID to prevent conflicts when num_n > 1.
  • Add unit tests for the truncation logic.

- Implement traffic matrix truncation in forward/backward all-to-all passes when `ragged_buffer_factor > 0` to prevent out-of-bounds writes in `ragged_all_to_all`.
- Partition GMM v2 VMEM scratchpad (`partial_out_ref`) by N program ID to prevent conflicts when `num_n > 1`.
- Add unit tests for the truncation logic.

PiperOrigin-RevId: 944752214
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.

1 participant