[Example] Add sequence parallel all-to-all examples

[Rachmanino]

Summary

• add TileScale IPC pre-attention all-to-all with and without fused transpose
• add the inverse post-attention all-to-all with fused transpose
• retain the tuned 65,536-CTA target for the pre-attention transpose path
• add usage documentation and 4-GPU correctness tests

Testing

• pre-commit run --files <changed files>
• 4x H200: all three distributed pytest cases passed
• default shape (B=2, H=32, S=8192, D=128, 4 GPUs): 131-139 us; no regression versus the migrated legacy examples

Summary by CodeRabbit

• New Features

  • Added distributed sequence-parallel examples for pre- and post-attention all-to-all communication.
  • Added fused transpose examples with configurable tensor shapes, data types, validation, benchmarking, and optional generated-source output.
  • Added setup and usage documentation, including four-GPU launch requirements.

• Tests

  • Added four-process tests covering all sequence-parallel communication examples.

[github-actions]

👋 Hi! Thank you for contributing to the TileScale project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: a4bbc3cb-de50-466d-bf09-e0c66abd8769

📥 Commits

Reviewing files that changed from the base of the PR and between 4e42b83 and f2f0122.

📒 Files selected for processing (5)

• examples/distributed/experimental/sequence_parallel/README.md
• examples/distributed/experimental/sequence_parallel/example_post_attn_all2all_transpose.py
• examples/distributed/experimental/sequence_parallel/example_pre_attn_all2all.py
• examples/distributed/experimental/sequence_parallel/example_pre_attn_all2all_transpose.py
• examples/distributed/experimental/sequence_parallel/test_example_sequence_parallel_all2all.py

---

📝 Walkthrough

Walkthrough

Added three distributed sequence-parallel all-to-all examples. The examples include TileLang kernels, PyTorch references, peer-memory allocation, correctness checks, benchmarks, multiprocess tests, and runtime documentation.

Changes

Sequence-parallel all-to-all examples

Layer / File(s)	Summary
Pre-attention all-to-all implementation examples/distributed/experimental/sequence_parallel/example_pre_attn_all2all.py	Adds sequence-to-head redistribution with a PyTorch reference, a T.put_block TileLang kernel, distributed setup, validation, benchmarking, and CLI controls.
Fused transpose variants examples/distributed/experimental/sequence_parallel/example_pre_attn_all2all_transpose.py, examples/distributed/experimental/sequence_parallel/example_post_attn_all2all_transpose.py	Adds fused pre-attention and post-attention redistribution kernels with tensor-layout transposes, peer-backed buffers, correctness checks, benchmarks, and multiprocess launch handling.
Distributed tests and usage documentation examples/distributed/experimental/sequence_parallel/test_example_sequence_parallel_all2all.py, examples/distributed/experimental/sequence_parallel/README.md	Adds four-process tests for all three examples and documents tensor layouts, execution, and runtime requirements.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Launcher
  participant ProcessGroup
  participant DistributedAllocator
  participant TileLangKernel
  participant Validation
  Launcher->>ProcessGroup: Spawn and initialize ranks
  ProcessGroup->>DistributedAllocator: Create peer-backed tensors
  DistributedAllocator->>TileLangKernel: Provide source and destination buffers
  TileLangKernel->>DistributedAllocator: Exchange blocks with T.put_block
  TileLangKernel->>Validation: Return redistributed output
  Validation->>ProcessGroup: Synchronize and compare with reference

Loading

Possibly related PRs

• tile-ai/tilescale#16: Adds related distributed sequence-parallel all-to-all examples.
• tile-ai/tilescale#18: Introduces distributed launching, allocation, and peer-memory infrastructure used by these examples.
• tile-ai/tilescale#21: Expands related pre-attention and post-attention transpose examples.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the added sequence-parallel all-to-all examples.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch add-sp-alltoall-examples

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}