Skip to content

feat(rocm): exotic-attention compositions (gated_attention/mla_decode/mla_decode_fused) — gfx1151 - #130

Merged
gstoner merged 1 commit into
mainfrom
rocm/exotic-attn-r
Jun 26, 2026
Merged

gstoner merged 1 commit into
mainfrom
rocm/exotic-attn-r

Conversation

@gstoner

@gstoner gstoner commented Jun 26, 2026

Copy link
Copy Markdown
Owner

Replacement for #127 (GitHub closed the original when its stacked base chain was deleted during the merge). Rebuilt clean onto post-#129 main — identical content, base main.

The attention analog of the matmul-family lane — each op composes already-compiled kernels (the WMMA flash_attn kernel + the WMMA GEMM kernel for MLA's latent projections + an elementwise gate), no new MLIR pass:

op composition
gated_attention flash_attn × sigmoid(gate)
mla_decode latent K/V projections (WMMA GEMM) → flash_attn
mla_decode_fused c=x@w_dkv; K=c@w_uk; V=c@w_uv (WMMA GEMM) → flash_attn

Shared lane rocm_exotic_attn_compiled. f16/bf16, f32 softmax+accumulate. 20 tests pass on gfx1151 + AMD-wiring; mypy clean; 17 dashboards in sync.

The recurrent DeltaNet variants (gated_deltanet/kimi_delta_attention/modified_delta_attention) + block-sparse deepseek_sparse_attention stay artifact_only with their technical blockers documented in ROCM_AUDIT (Decision #25 — sequential-scan / sparse-gather kernels, not compositions).

🤖 Generated with Claude Code

…e/mla_decode_fused on gfx1151

The attention analog of the matmul-family lane — each op composes
already-compiled kernels (the WMMA flash_attn kernel + the WMMA GEMM kernel for
MLA's latent projections + an elementwise gate), no new MLIR pass:

- gated_attention   — flash_attn × elementwise sigmoid-gate
- mla_decode        — latent K/V projections (WMMA GEMM) + flash_attn
- mla_decode_fused  — down/up projections (c=x@w_dkv; K=c@w_uk; V=c@w_uv) + flash_attn

Shared lane rocm_exotic_attn_compiled. f16/bf16, f32 softmax+accumulate.
Validated on gfx1151 vs the numpy attention reference. The recurrent DeltaNet
variants + block-sparse deepseek stay artifact_only with their technical
blockers documented in ROCM_AUDIT (Decision #25 — not marked compiled).

(Rebuilt clean onto post-#129 main; supersedes the original #127 branch which
GitHub closed when its stacked base chain was deleted during the merge.)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@gstoner
gstoner merged commit a42f586 into main Jun 26, 2026
12 checks passed
@gstoner
gstoner deleted the rocm/exotic-attn-r branch June 26, 2026 00:05
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.

1 participant