Repository navigation
[Feat]support sequence parallel all2all with transpose example - #21
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughAdds two new distributed example scripts demonstrating all-to-all transpose (pre-attn) and reverse all-to-all transpose (post-attn) with PyTorch golden references, TileLang NVSHMEM kernels, CLI, and verification. Also inserts a memory fence after an NVSHMEM signal in the existing pre-attn all-to-all example. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User
participant Script as Pre-attn Transpose Script
participant Torch as PyTorch Ref
participant Kernel as TileLang Kernel
participant NV as NVSHMEM Runtime
participant PEs as PEs (Ranks)
User->>Script: run test_transpose_all_to_all_with_golden_reference()
Script->>Torch: all_to_all + reshape (golden reference)
Script->>Kernel: compile(generate params)
Script->>NV: allocate/copy inputs
Kernel->>NV: signal + (fence) + wait loops
NV->>PEs: per-(batch, head) transfers (all-to-all transpose)
PEs-->>NV: data delivered
NV-->>Script: output buffer
Script->>Script: verify_results(tilelang vs torch)
sequenceDiagram
autonumber
participant User
participant Script as Post-attn Reverse Transpose Script
participant Torch as PyTorch Ref
participant Kernel as TileLang Reverse Kernel
participant NV as NVSHMEM Runtime
participant PEs as PEs (Ranks)
User->>Script: run test_reverse_transpose_all_to_all_with_golden_reference()
Script->>Torch: reverse all_to_all + reorder (golden ref)
Script->>Kernel: compile(generate params)
Script->>NV: allocate/copy inputs
Kernel->>NV: putmem_nbi_block + signals + waits
NV->>PEs: per-(batch, seq) cross-PE transfers
PEs-->>NV: data delivered
NV-->>Script: output buffer
Script->>Script: verify_results(tilelang vs torch)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
There was a problem hiding this comment.
Pull Request Overview
This PR introduces support for sequence parallel all-to-all communication with transpose functionality in distributed attention mechanisms. The implementation provides both pre-attention and post-attention all-to-all operations that convert between sequence-parallel and head-parallel tensor layouts.
Key changes:
- Adds pre-attention all-to-all with transpose converting from sequence-parallel to head-parallel format
- Adds post-attention all-to-all with transpose converting from head-parallel back to sequence-parallel format
- Fixes a missing fence operation in the existing pre-attention all-to-all example
Reviewed Changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| examples/distributed/example_pre_attn_all2all_transpose.py | New example implementing pre-attention all-to-all with transpose from sequence-parallel to head-parallel format |
| examples/distributed/example_pre_attn_all2all.py | Adds missing fence operation after signal for proper synchronization |
| examples/distributed/example_post_attn_all2all_transpose.py | New example implementing post-attention all-to-all with transpose from head-parallel back to sequence-parallel format |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
examples/distributed/example_pre_attn_all2all_transpose.py (1)
115-123: Make transfer size dtype-agnostic (avoid hardcoded 2 bytes)The code assumes fp16. Compute element size from
dtypeso fp32/bf16 work too. This was flagged previously—addressing it here closes the loop.- transfer_size = HEAD_DIM * 2 # float16 = 2 bytes + transfer_size = HEAD_DIM * ELEM_NBYTESAdd this near the top of
sequence_parallel_all_to_all_transpose(after HEADS_PER_PE):ELEM_NBYTES = torch.tensor([], dtype=dtype_map[dtype]).element_size()
🧹 Nitpick comments (14)
examples/distributed/example_pre_attn_all2all.py (1)
129-137: Fence after signal: safe, but likely redundant; prefer explicit compare enum for wait
- Adding
T.fence()afterT.signal_opis safe, yet the fence before signaling already orders the preceding puts; the post-signal fence adds little. Keep if you’ve observed ordering issues, otherwise consider removing to reduce overhead.- Small consistency fix: use
T.CmpType.EQinstead of magic0to match the other examples.- for k in T.serial(PE_num): - T.signal_wait_until(T.address_of(signal[k]), 0, NUM_BLOCKS_X) + for k in T.serial(PE_num): + T.signal_wait_until(T.address_of(signal[k]), T.CmpType.EQ, NUM_BLOCKS_X)examples/distributed/example_pre_attn_all2all_transpose.py (6)
28-30: Remove unused variablerank
rankis assigned but never used. Drop it to silence lint and reduce noise.- world_size = dist.get_world_size(group) - rank = dist.get_rank(group) + world_size = dist.get_world_size(group)
125-137: Ensure completion before kernel exitAfter the waits, add
T.quiet()to guarantee all nonblocking puts are complete before the kernel returns.T.fence() # Wait for all blocks to complete all sequence positions T.signal_wait_until(T.address_of(signal[target_pe]), T.CmpType.EQ, NUM_BLOCKS_X) + T.quiet()
1-10: Trim unused imports
os,TensorSupplyType, andperf_fnare unused here.-import os -import tilelang -import tilelang.language as T -from tilelang.profiler import TensorSupplyType -from tilelang.distributed.utils import init_distributed, dtype_map, perf_fn +import os # REMOVE +import tilelang +import tilelang.language as T +from tilelang.profiler import TensorSupplyType # REMOVE +from tilelang.distributed.utils import init_distributed, dtype_map, perf_fn # drop perf_fn
187-190: Prefer explicit validation overassertin examplesReplace runtime
assertwith a clear exception so the script behaves the same under optimized Python.- assert args.seq_len % PE_num == 0, f"seq_len {args.seq_len} must be divisible by PE_num {PE_num}" - assert args.num_heads % PE_num == 0, f"num_heads {args.num_heads} must be divisible by PE_num {PE_num}" + if args.seq_len % PE_num != 0: + raise ValueError(f"seq_len {args.seq_len} must be divisible by PE_num {PE_num}") + if args.num_heads % PE_num != 0: + raise ValueError(f"num_heads {args.num_heads} must be divisible by PE_num {PE_num}")
248-251: Optional: add NVSHMEM barrier for symmetry with post-attn exampleNot strictly required given the signal/waits, but adding a barrier aligns both examples’ host-side sync.
# Execute kernel kernel(data_src, data_dst, signal) - - return data_dst + pynvshmem.nvshmem_barrier_all() + return data_dst
264-264: Drop unused variable
verification_passedisn’t used.- verification_passed = verify_results(tilelang_output, torch_output, RANK) + verify_results(tilelang_output, torch_output, RANK)examples/distributed/example_post_attn_all2all_transpose.py (7)
28-30: Remove unused variablerank
rankis never used.- world_size = dist.get_world_size(group) - rank = dist.get_rank(group) + world_size = dist.get_world_size(group)
130-139: Comment and accounting mismatchThe comment says “signal the number of head chunks processed,” but the code signals
1per block (not per head). Either increment byHEADS_PER_PEor adjust the comment to “signal block completion.” Current logic is fine since wait usesNUM_BLOCKS_X.
233-257: Barrier usage: consistent host-side syncYou call
nvshmem_barrier_all()here but not in the pre-attn counterpart. Consider keeping both consistent (either both have it or both rely solely on signal/waits).
269-269: Drop unused variable
verification_passedisn’t used.- verification_passed = verify_results(tilelang_output, torch_output, RANK) + verify_results(tilelang_output, torch_output, RANK)
191-194: Prefer explicit validation overassertin examplesUse exceptions to ensure checks aren’t skipped with optimized Python.
- assert args.seq_len % PE_num == 0, f"seq_len {args.seq_len} must be divisible by PE_num {PE_num}" - assert args.num_heads % PE_num == 0, f"num_heads {args.num_heads} must be divisible by PE_num {PE_num}" + if args.seq_len % PE_num != 0: + raise ValueError(f"seq_len {args.seq_len} must be divisible by PE_num {PE_num}") + if args.num_heads % PE_num != 0: + raise ValueError(f"num_heads {args.num_heads} must be divisible by PE_num {PE_num}")
1-10: Trim unused imports
os,TensorSupplyType, andperf_fnaren’t used in this file.-import os -import tilelang -import tilelang.language as T -from tilelang.profiler import TensorSupplyType -from tilelang.distributed.utils import init_distributed, dtype_map, perf_fn +import os # REMOVE +import tilelang +import tilelang.language as T +from tilelang.profiler import TensorSupplyType # REMOVE +from tilelang.distributed.utils import init_distributed, dtype_map, perf_fn # drop perf_fn
35-45: Stale comment about transposeThe comment mentions a preparatory transpose, but the code sends contiguous
[B, H, S_chunk, D]. Update the comment for accuracy.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
examples/distributed/example_post_attn_all2all_transpose.py(1 hunks)examples/distributed/example_pre_attn_all2all.py(1 hunks)examples/distributed/example_pre_attn_all2all_transpose.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
examples/distributed/example_pre_attn_all2all.py (1)
tilelang/language/distributed/multi_device/nvshmem.py (1)
fence(61-63)
examples/distributed/example_post_attn_all2all_transpose.py (5)
tilelang/utils/tensor.py (1)
TensorSupplyType(15-22)tilelang/distributed/utils.py (1)
perf_fn(217-238)tilelang/env.py (1)
disable_cache(247-248)examples/distributed/example_pre_attn_all2all_transpose.py (2)
main(91-136)torch_sequence_all_to_all_transpose_reference(14-71)tilelang/language/distributed/multi_device/nvshmem.py (6)
get_pe(6-8)putmem_nbi_block(101-109)fence(61-63)signal_op(164-172)signal_wait_until(175-183)quiet(56-58)
examples/distributed/example_pre_attn_all2all_transpose.py (6)
tilelang/utils/tensor.py (1)
TensorSupplyType(15-22)tilelang/distributed/utils.py (1)
perf_fn(217-238)tilelang/env.py (1)
disable_cache(247-248)examples/distributed/example_post_attn_all2all_transpose.py (2)
main(94-140)verify_results(145-168)examples/distributed/example_pre_attn_all2all.py (3)
main(92-137)verify_results(142-165)tilelang_all_to_all(238-255)tilelang/language/distributed/multi_device/nvshmem.py (5)
get_pe(6-8)putmem_nbi_block(101-109)fence(61-63)signal_op(164-172)signal_wait_until(175-183)
🪛 Ruff (0.12.2)
examples/distributed/example_post_attn_all2all_transpose.py
29-29: Local variable rank is assigned to but never used
Remove assignment to unused variable rank
(F841)
49-49: Loop control variable pe_idx not used within loop body
(B007)
192-192: Use of assert detected
(S101)
193-193: Use of assert detected
(S101)
269-269: Local variable verification_passed is assigned to but never used
Remove assignment to unused variable verification_passed
(F841)
examples/distributed/example_pre_attn_all2all_transpose.py
29-29: Local variable rank is assigned to but never used
Remove assignment to unused variable rank
(F841)
47-47: Loop control variable pe_idx not used within loop body
(B007)
188-188: Use of assert detected
(S101)
189-189: Use of assert detected
(S101)
264-264: Local variable verification_passed is assigned to but never used
Remove assignment to unused variable verification_passed
(F841)
pre attn is all to all Before attention and post attn is all to all after attetnion for sequence parallel
Summary by CodeRabbit
New Features
Bug Fixes