Skip to content

fix(rocm): clamp QSA KV tile to the device LDS budget on RDNA (gfx11/gfx12) - #587

Merged
jason-fxz merged 1 commit into
FlashML-org:mainfrom
zhangnju:rocm714_torch214
Oct 16, 2026
Merged

fix(rocm): clamp QSA KV tile to the device LDS budget on RDNA (gfx11/gfx12)#587
jason-fxz merged 1 commit into
FlashML-org:mainfrom
zhangnju:rocm714_torch214

Conversation

@zhangnju

@zhangnju zhangnju commented Oct 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What

qsa_sparse_paged_attention picks KV staging tiles tuned for the GB300's ~228KB shared memory. RDNA parts (gfx1100 / gfx1201) have ~64KB/CU, so k+v staging overflows LDS and the kernel fails to launch on AMD. This queries the device's actual shared_memory_per_block and shrinks block_n only when the tile doesn't fit — ROCm-gated (torch.version.hip is not None) so NVIDIA tiles stay byte-identical.

Scope change

This branch previously also carried two gguf ROCm-compile fixes (a 64-bit warp mask in the sgl-kernel donor dispatch.h, and stripping nvcc-only cflags in the gguf build loader). Both are now obsolete: #690 re-based the gguf backend onto vendored ggml-cuda, which ships its own vendors/hip.h HIP shuffles and a separate DEFAULT_HIP_CFLAGS (no --expt-relaxed-constexpr / -ccbin) in kernel/utils.py. So this PR now reduces to the single QSA LDS clamp, which is still missing upstream (attend.py unchanged on main).

Validation

Full pytest passes on W7900 (gfx1100) and R9700 (gfx1201). All changes ROCm-gated; the CUDA path is untouched.

@zhangnju

zhangnju commented Oct 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

@jason-fxz @andy-yang-1 could you take a look when you have a chance? This enables the ROCm / RDNA3+RDNA4 path on the latest stacks (torch 2.14 + ROCm 7.14 and torch 2.15 + ROCm 10), all ROCm-gated so the CUDA path is unchanged. Full pytest passes on W7900 (gfx1100) and R9700 (gfx1201) on both stacks. Happy to share more detail or split anything out. Thanks!

@github-actions github-actions Bot mentioned this pull request Oct 1, 2026
Draft
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 3, 2026
claude2 @claude
feat(rocm): run on torch 2.14/2.15 with ROCm 7.14/10 on RDNA3/RDNA4
911d405 
- moe: a pure-torch, graph-capturable moe_align_block_size on ROCm, where
  Triton 3.8's AMD backend drops the masked scatter-store of the Triton
  kernel and misaligns expert_ids with sorted_token_ids
- cpu_moe: call the HIP stream memops directly (hipStreamWaitValue64 takes
  a mask), and probe a captured WAIT->WRITE graph replay at startup; ROCm
  7.14 does not replay memop nodes, so it falls back to the host-callback sync
- qsa: shrink the KV tile to the device's LDS and run single-stage on ROCm
- gguf: unmasked shuffles and hipcc-compatible flags on ROCm
- build: C++20 only for ROCm torch >= 2.14; scripts/rocm_overrides.py pins
  the image's torch/triton for uv so the CUDA pins stay unchanged

Alternative to upstream PR FlashML-org#587: keeps the CUDA dependency pins and C++17
build untouched instead of widening them for everyone, detects memop graph
replay functionally instead of by runtime version, and drops the hiprtc
spin-kernel handshake (a busy CU) in favor of the existing fallback.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 3, 2026
claude2 @claude
chore(fork): record upstream PR decisions
fa9f634 
Adds the fork's PR triage tooling (fork/pr_queue.py) and logs FlashML-org#591, FlashML-org#590,
FlashML-org#589, FlashML-org#587 and FlashML-org#585.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 3, 2026
KarrAcaRn @claude @zhangnju
feat(rocm): run on torch 2.14/2.15 with ROCm 7.14/10 on RDNA3/RDNA4
c362171 
- moe: a pure-torch, graph-capturable moe_align_block_size on ROCm, where
  Triton 3.8's AMD backend drops the masked scatter-store of the Triton
  kernel and misaligns expert_ids with sorted_token_ids
- cpu_moe: call the HIP stream memops directly (hipStreamWaitValue64 takes
  a mask), and probe a captured WAIT->WRITE graph replay at startup; ROCm
  7.14 does not replay memop nodes, so it falls back to the host-callback sync
- qsa: shrink the KV tile to the device's LDS and run single-stage on ROCm
- gguf: unmasked shuffles and hipcc-compatible flags on ROCm
- build: C++20 only for ROCm torch >= 2.14; scripts/rocm_overrides.py pins
  the image's torch/triton for uv so the CUDA pins stay unchanged

Alternative to upstream PR FlashML-org#587: keeps the CUDA dependency pins and C++17
build untouched instead of widening them for everyone, detects memop graph
replay functionally instead of by runtime version, and drops the hiprtc
spin-kernel handshake (a busy CU) in favor of the existing fallback.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Zhang Jason <ning.zhang2@amd.com>
Co-authored-by: zhangnju <zhangnju@126.com>
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 3, 2026
KarrAcaRn @claude
chore(fork): record upstream PR decisions
48a10eb 
Adds the fork's PR triage tooling (fork/pr_queue.py) and logs FlashML-org#591, FlashML-org#590,
FlashML-org#589, FlashML-org#587 and FlashML-org#585.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
KarrAcaRn pushed a commit to KarrAcaRn/FreeToken-ByAI that referenced this pull request Oct 7, 2026
KarrAcaRn
Merge upstream PR FlashML-org#241: rocm: fix native extensions, kerne…
8ec7388 
…l JIT builds, and Triton PTX fallbacks

Conflicts: next already gives gguf unmasked ROCm shuffles and hipcc flags (c362171, our take on FlashML-org#587); kept those. Takes the hipify .gitignore entries and the RDNA decode BLOCK_H >= 16 padding (masked lanes, ROCm only).

Assisted-by: Claude Opus 5.5
This was referenced Oct 7, 2026
Closed
Merged
LZ-QWQ added a commit to LZ-QWQ/FreeToken that referenced this pull request Oct 8, 2026
@LZ-QWQ
merge(rocm): integrate PR 587 after roadmap support
ae6bd21 
Apply PR 587 after FlashML-org#133, FlashML-org#134, and the incremental FlashML-org#378 commit. Retain the FlashML-org#378 CPU-MoE handshake and its runtime-only linkage because the overlapping FlashML-org#587 path previously failed on ROCm 10. Keep the FlashML-org#134 installation document; record the actual workspace environment outside this branch.
LZ-QWQ added a commit to LZ-QWQ/FreeToken that referenced this pull request Oct 8, 2026
@LZ-QWQ
merge(rocm): integrate PR 587 after the full roadmap stack
9357825 
Integrate FlashML-org#133, FlashML-org#134, FlashML-org#135, FlashML-org#378 and FlashML-org#491 before FlashML-org#587. Preserve the tested FlashML-org#378 handshake and runtime-only linkage, and retain the FlashML-org#134 installation document. Exclude CI and closed GGUF PR FlashML-org#136.
LZ-QWQ added a commit to LZ-QWQ/FreeToken that referenced this pull request Oct 9, 2026
@LZ-QWQ
merge(rocm): integrate remaining support after upstream AMD merges
a8b9840 
Use upstream FlashML-org#638 and FlashML-org#134; integrate incremental FlashML-org#135 and FlashML-org#378 plus the updated standalone FlashML-org#491 before FlashML-org#587. Retain the probed FlashML-org#378 handshake and upstream installation document.
@LZ-QWQ LZ-QWQ mentioned this pull request Oct 15, 2026
Merged
@zihaomu zihaomu mentioned this pull request Oct 15, 2026
Merged
@zhangnju zhangnju changed the title feat(rocm): run on RDNA3/RDNA4 with torch 2.14/2.15 and ROCm 7.14/10 fix(rocm): RDNA fixes still missing upstream (QSA LDS clamp + gguf ROCm compile) Oct 15, 2026
@github-actions github-actions Bot added the amd label Oct 15, 2026
@zhangnju
fix(rocm): clamp QSA KV tile to the device LDS budget on RDNA
3ebfaac 
The qsa_sparse_paged_attention tiles are tuned for the GB300's 228KB shared
memory; RDNA gfx11/gfx12 have ~64KB/CU, so k+v staging overflows LDS. Query the
device's actual shared-memory limit and shrink the KV tile only when it doesn't
fit, ROCm-gated so NVIDIA tiles stay byte-identical.

The gguf ROCm-compile fixes this branch also carried are now covered upstream by
the ggml-cuda vendoring in FlashML-org#690 (separate DEFAULT_HIP_CFLAGS + vendors/hip.h), so
this branch reduces to the QSA clamp only.
@zhangnju zhangnju changed the title fix(rocm): RDNA fixes still missing upstream (QSA LDS clamp + gguf ROCm compile) fix(rocm): clamp QSA KV tile to the device LDS budget on RDNA (gfx11/gfx12) Oct 16, 2026
@zhangnju

zhangnju commented Oct 16, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto latest main and reduced this to the single QSA LDS clamp. The two gguf ROCm-compile fixes this branch carried are now covered upstream by the ggml-cuda vendoring in #690 (separate DEFAULT_HIP_CFLAGS in kernel/utils.py + ggml-cuda/vendors/hip.h), so I dropped them. The QSA clamp is still needed — attend.py is unchanged on main and overflows the ~64KB RDNA LDS with the GB300-sized tiles. @jason-fxz @andy-yang-1 PTAL, thanks!

@jason-fxz

jason-fxz commented Oct 16, 2026

Copy link
Copy Markdown
Collaborator

LGTM, merging.

@jason-fxz
jason-fxz merged commit 7d6490c into FlashML-org:main Oct 16, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

amd

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@zhangnju @jason-fxz