E5DF [a2av] Align length of major dimension in output of 2D a2av by kwen2501 · Pull Request #155172 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[a2av] Align length of major dimension in output of 2D a2av #155172

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kwen2501 wants to merge 10 commits into from
Closed

[a2av] Align length of major dimension in output of 2D a2av #155172

Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
79094e2
Update
kwen2501 Jun 4, 2025
ecfeaf4
Update
kwen2501 Jun 4, 2025
754d6bd
Update
kwen2501 Jun 5, 2025
93041d6
Update
kwen2501 Jun 5, 2025
b8db459
Update
kwen2501 Jun 5, 2025
8a889c0
Update
kwen2501 Jun 5, 2025
ba6b78d
Update
kwen2501 Jun 5, 2025
75cc801
Update
kwen2501 Jun 5, 2025
10BC0 741b5a8
Update
kwen2501 Jun 5, 2025
fd596f4
Update
kwen2501 Jun 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Update 
[ghstack-poisoned]
  • Loading branch information
@kwen2501
kwen2501 committed Jun 5, 2025
commit ba6b78db3dc8531bc9495c695aeed0280f4b6644
3 changes: 3 additions & 0 deletions torch/csrc/distributed/c10d/nvshmem_extension.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -518,8 +518,11 @@ at::Tensor nvshmem_all_to_all_vdev_2d(
auto split_shape = in_out_splits.sizes();
TORCH_CHECK(split_shape.size() == 2 && split_shape[0] == 3, "in_out_splits must be 2D with 3 rows");
TORCH_CHECK(split_shape[1] % world_size == 0, "Each row of in_out_splits must be a multiple of world_size");

// Number of experts per rank
int ne = split_shape[1] / world_size;
constexpr int NUM_TILES = THREADS_PER_BLOCK / A2AV_TILE_SIZE;
TORCH_CHECK(ne <= NUM_TILES, "Number of experts must be smaller than NUM_TILES", NUM_TILES);

// Set device context for getting the stream and launching kernels below
c10::cuda::CUDAGuard guard(input.device());
Expand Down
Loading
0