Skip to content

[Improvement] CopyType::Vector in concatenate if axis=0 and contiguous #3663

Merged
nastya236 merged 1 commit into
mainfrom
concatenate-vector
Jun 13, 2026
Merged

[Improvement] CopyType::Vector in concatenate if axis=0 and contiguous #3663
nastya236 merged 1 commit into
mainfrom
concatenate-vector

Conversation

@nastya236

@nastya236 nastya236 commented Jun 12, 2026

Copy link
Copy Markdown
Collaborator

Tiny improvement for concatenating along axis 0 in contiguous case. This could be useful in fsdp / average_gradients (we concatenate weights and launch reduction on concatenation).

Shape n Old time (ms) Old GB/s New time (ms) New GB/s
[65536] 2 0.0578 9.07 0.0518 10.13
[65536] 16 0.1598 26.24 0.1264 33.20
[1048576] 16 0.1894 354.41 0.1540 435.91
[1024,1024] 4 0.0834 201.11 0.0671 250.03
[4096,4096] 4 0.1954 1373.50 0.1789 1500.79
[8192,8192] 10 1.0496 2557.58 1.0154 2643.70

@nastya236 nastya236 requested a review from zcbenz June 12, 2026 13:13

@angeloskath angeloskath left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Nice and general!

@nastya236 nastya236 merged commit 8906447 into main Jun 13, 2026
16 checks passed
@nastya236 nastya236 deleted the concatenate-vector branch June 13, 2026 19:32
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