feat(rocm): matmul-family compiled lane (batched_gemm/linear_general/qkv_projection/factorized_matmul/einsum) — gfx1151 - #129
Conversation
…/qkv_projection/factorized_matmul/einsum on gfx1151
All five matmul-family ops execute on the SAME compiler-generated WMMA GEMM
kernel (the rocm_compiled spine), reshaped/batched/split in the runtime — no new
MLIR pass:
- batched_gemm — loop the gemm over leading batch dims (np.matmul semantics)
- linear_general — axis=-1 reshape [...,K]→[M,K] + gemm + optional bias
- qkv_projection — packed x@W_qkv projection (the 3-way split is a host view)
- factorized_matmul — GPU matmul + exact host rank-r SVD-truncate epilogue
- einsum — single-contraction two-operand specs → (batched) gemm;
other specs emit a stable "unsupported" diagnostic (#21)
Shared lane rocm_matmul_family_compiled; GEMM-family compiled entries keep the
unified MMA descriptor. f16/bf16, f32 accumulate. Validated on gfx1151 vs numpy.
(Rebuilt clean onto post-#125 main; supersedes the original #126 branch which
GitHub closed when its stacked base branch was deleted during the merge.)
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e/mla_decode_fused on gfx1151 (#130) 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: gstoner <angstroms01@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a8ed73aef1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| y = _rocm_wmma_gemm_2d(x2, w).reshape(*x.shape[:-1], n3) | ||
| # The op returns (Q, K, V); the runtime returns the packed projection and | ||
| # the 3-way last-axis split is a trivial host view. | ||
| return y |
There was a problem hiding this comment.
When a tessera.qkv_projection artifact is launched through this new ROCm executor, the public API/CPU and Apple lanes return (Q, K, V), but this path returns the packed [..., 3N] array. Any caller that unpacks q, k, v = rt.launch(...)["output"] will either fail or read slices along the wrong axis when the leading dimension happens to be 3, so the lane is not semantically equivalent to the op it advertises; split y before returning and reject non-3-way widths.
Useful? React with 👍 / 👎.
| if b_batch.shape[0] != nb_a: | ||
| raise ValueError( | ||
| f"batched_gemm batch mismatch; got {a.shape} @ {b.shape}") |
There was a problem hiding this comment.
For broadcastable batch operands such as A.shape == (5, M, K) and B.shape == (1, K, N), tessera.ops.batched_gemm is defined as the np.matmul reference and should broadcast B across the batch. This branch flattens B and requires the flattened batch count to equal A's unless B is rank-2, so valid inputs raise batched_gemm batch mismatch; the same issue affects einsum specs that rely on NumPy batch broadcasting.
Useful? React with 👍 / 👎.
| a2 = lhs_t.reshape(int(np.prod(bsh or [1])), int(np.prod(msh or [1])), kdim) | ||
| b2 = rhs_t.reshape(int(np.prod(bsh or [1])), kdim, int(np.prod(nsh or [1]))) | ||
| prod = _rocm_batched_gemm(a2, b2) # [B, M, N] | ||
| canon = prod.reshape(*(bsh + msh + nsh)) # batch + free_l + free_r |
There was a problem hiding this comment.
A vector dot such as i,i-> is a two-operand single-contraction einsum that maps cleanly to a 1xK @ Kx1 GEMM, but here bsh + msh + nsh is empty, so this expands to prod.reshape() and raises TypeError instead of returning the scalar result. Use an explicit empty shape, e.g. reshape(()), for the no-free-index case.
Useful? React with 👍 / 👎.
1 participant
Replacement for #126 (GitHub closed the original when its stacked base branch
rocm/silu-mul-alibiwas deleted after #125 merged). Rebuilt clean onto post-#125main— identical content, basemain.All five matmul-family ops execute on the same compiler-generated WMMA GEMM kernel (the
rocm_compiledspine), reshaped/batched/split in the runtime — no new MLIR pass:batched_gemmlinear_generalqkv_projectionx@W_qkv(3-way split is a host view)factorized_matmuleinsumShared lane
rocm_matmul_family_compiled; GEMM-family compiled entries keep the unified MMA descriptor (fixes thetest_backend_manifest_rocm_gemm_carries_mma_descriptorregression). f16/bf16, f32 accumulate. 32 tests pass on gfx1151 + AMD-wiring; mypy clean; 17 dashboards in sync.🤖 Generated with Claude Code