Implement ring attention (arXiv:2310.01889) - #450

Draft

mawad-amd with Copilot wants to merge 30 commits into

mainfrom

copilot/implement-ring-attention

Commits on Mar 12, 2026

Initial plan
Copilot
committed
d21bb20
Implement Ring Attention (arxiv:2310.01889)

Copilot
and
mawad-amd
committed
445e4ec
Replace torch.distributed with iris.put for ring KV rotation

Copilot
and
mawad-amd
committed
32c79d7
Fuse K+V put into single kernel; validate on AMD GPUs (16/16 tests pass)

Copilot
and
mawad-amd
committed
2eefc41
Add benchmark script with roofline analysis and measured results on MI300X

Copilot
and
mawad-amd
committed
17c9856
Address code review: typed buffer cache, named constants, fix --causal arg semantics

Copilot
and
mawad-amd
committed
bbdac9b

Commits on Mar 13, 2026

Add strong/weak scaling benchmark (1-8 GPUs) and update results figure

Copilot
and
mawad-amd
committed
0f01c45

Commits on Mar 23, 2026

Optimize ring attention: comm-compute overlap + bug fixes

mawad-amd
and
claude
committed
5357874
Apply Ruff auto-fixes
github-actions[bot]
committed
7b76443
ring_attention: fuse KV rotation into attention kernel

mawad-amd
and
claude
committed
a118644
Apply Ruff auto-fixes
github-actions[bot]
committed
647b53d
ring_attention: add per-step profiling script

mawad-amd
and
claude
committed
5037b94
Apply Ruff auto-fixes
github-actions[bot]
committed
f68bd36
ring_attention: eliminate causal load imbalance

mawad-amd
and
claude
committed
6f9fa80
Apply Ruff auto-fixes
github-actions[bot]
committed
0430a23
perf: use fp16 MFMA and causal early exit in ring attention kernel

mawad-amd
and
claude
committed
061903c
ring_attention: persistent kernel with signal-flag synchronization

mawad-amd
and
claude
committed
651d68a
Apply Ruff auto-fixes
github-actions[bot]
committed
dc7b345
ring_attention: fix put visibility — use sys scope on completion counter

mawad-amd
and
claude
committed
8b9ba4c
ring_attention: simplify causal path to match original kernel structure

mawad-amd
and
claude
committed
929e90c
DEBUG: disable inner-loop skip to test causal mask

mawad-amd
and
claude
committed
cf49885
fix: pass pre-computed kv_rank_starts array to avoid in-kernel constexpr * runtime

mawad-amd
and
claude
committed
7eb6ab0
Apply Ruff auto-fixes
github-actions[bot]
committed
6709eb0
re-enable causal early-exit optimization for KV blocks

mawad-amd
and
claude
committed
a6f701e

Commits on Mar 24, 2026

ring_attention: fix put visibility race + relax FP16 tolerances

mawad-amd
and
claude
committed
6b60492
ring_attention: per-CTA remote signaling for put visibility

mawad-amd
and
claude
committed
0808688
Apply Ruff auto-fixes
github-actions[bot]
committed
6907194
Add backward signaling to prevent write-before-read race in persistent kernel

mawad-amd
and
claude
committed
fc69249
Apply Ruff auto-fixes
github-actions[bot]
committed
a11a350
Fix signal flag reset ordering to prevent deadlock across iterations

mawad-amd
and
claude
committed
0d07d5d