From a3de4ff0e1e62cd8a7e2d216064c1e704304d440 Mon Sep 17 00:00:00 2001 From: gstoner Date: Fri, 10 Jul 2026 11:27:00 -0600 Subject: [PATCH 1/5] rocm/gfx1151: FA-2 backward runtime lane (MHA+GQA+attn_bias), KV-cache paged lane, lstm_cell + inventory reconciliation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes several gfx1151 gaps, all execute-compare verified on real RDNA 3.5 silicon. Lanes added / promoted (runtime.launch, execution_mode=hip_runtime): - rocm_flash_attn_bwd_compiled — FA-2 backward (dQ/dK/dV) as a first-class lane. The compiler-generated backward kernel existed in C++ but was never wired into runtime.launch(). Added _execute_rocm_compiled_flash_attn_bwd (fa_pre -> fa_dkdv -> fa_dq; O recomputed via the forward lane, nothing saved from forward). Covers MHA + GQA/MQA (grouped fa_dkdv atomic-accumulates dK/dV) + additive attn_bias (S = scale*Q@K^T + bias in the recompute). The attn_bias variant is a new C++ path in GenerateWMMAFlashAttnBwdKernel.cpp (attn_bias BoolAttr -> trailing f32 [bh,Sq,Sk] memref added to S in emitPre + recomputeScoreTile). f16/bf16, f32 accumulate; matches autodiff vjp_flash_attn. - rocm_kv_cache_compiled — KV-cache paged-movement core (append/read/prune) over a resident cache buffer, composing the existing device scatter (append row write) + masked-gather (read/prune) kernels with host page-index math. Matches the KVCacheHandle reference. quantize_kv rides the intquant lane. - lstm_cell — added F.lstm_cell reference (gate order i,f,g,o) and registered it in the structured-compute lane (x86 + rocm), parity with gru/simple_rnn. Docs / automation: - Reconciled docs/rocm_mfma_kernel_inventory.md with the generated runtime execution matrix (it had drifted badly: claimed only matmul + flash_attn executed while dozens of compiled HIP lanes had landed). Points at the generated matrix as status truth (Decision #26). - Hardened the inventory guard test: TestH3RocmInventoryExecutionStatus now cross-checks every rocm_*_compiled lane named in the doc against the drift-gated matrix, so the execution-status prose cannot silently rot again. Verification: 147 unit tests pass (FA fwd/bwd runtime+direct, gqa-direct, structured x86+rocm, matrix, inventory guard); generated-doc drift gate in sync. ROCm runtime lanes 65 -> 67. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/audit/generated/apple_target_map.csv | 1 + docs/audit/generated/apple_target_map.md | 5 +- docs/audit/generated/compiler_progress.csv | 8 +- docs/audit/generated/compiler_progress.md | 8 +- docs/audit/generated/rocm_target_map.csv | 2 +- .../generated/runtime_execution_matrix.csv | 2 + .../generated/runtime_execution_matrix.md | 4 + docs/audit/generated/s_series_status.md | 2 +- docs/audit/op_target_conformance.csv | 2 +- docs/audit/op_target_conformance.md | 6 +- docs/rocm_mfma_kernel_inventory.md | 212 ++++++--- python/tessera/compiler/backend_manifest.py | 35 +- python/tessera/compiler/execution_matrix.py | 45 ++ python/tessera/nn/functional.py | 23 + python/tessera/runtime.py | 431 ++++++++++++++++++ .../GenerateWMMAFlashAttnBwdKernel.cpp | 60 ++- .../test_kernel_inventory_and_lit_fixtures.py | 49 +- .../test_rocm_flash_attn_bwd_runtime_lane.py | 208 +++++++++ tests/unit/test_rocm_kv_cache_compiled.py | 137 ++++++ .../test_rocm_structured_compute_compiled.py | 11 + .../test_x86_structured_compute_compiled.py | 9 + 21 files changed, 1161 insertions(+), 99 deletions(-) create mode 100644 tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py create mode 100644 tests/unit/test_rocm_kv_cache_compiled.py diff --git a/docs/audit/generated/apple_target_map.csv b/docs/audit/generated/apple_target_map.csv index fa5fdc026..71044695d 100644 --- a/docs/audit/generated/apple_target_map.csv +++ b/docs/audit/generated/apple_target_map.csv @@ -98,6 +98,7 @@ load_balance_loss,tensor,reference,numpy_reference,fp32,,numpy_reference,compile log_cosh_loss,tensor,reference,numpy_reference,fp32,,numpy_reference,compiled,,fp32,,manifest,tests/unit/test_apple_*.py, lookahead_sparse_attention,tensor,reference,numpy_reference,fp32,,numpy_reference,fused,,fp32,,manifest,tests/unit/test_apple_*.py, lora_linear,tensor,reference,numpy_reference,fp32,,numpy_reference,compiled,,fp32,,manifest,tests/unit/test_apple_*.py, +lstm_cell,tensor,reference,numpy_reference,fp32,,numpy_reference,compiled,,fp32,,manifest,tests/unit/test_apple_*.py, lu,tensor,reference,numpy_reference,fp32,,numpy_reference,compiled,,fp32,,manifest,tests/unit/test_apple_*.py, mae_loss,tensor,reference,numpy_reference,fp32,,numpy_reference,compiled,,fp32,,manifest,tests/unit/test_apple_*.py, masked_categorical,tensor,reference,numpy_reference,fp32,,numpy_reference,fused,,"fp32,int32",tessera_apple_gpu_masked_categorical_f32,driver,tests/unit/test_apple_gpu_ldt_loss_ops.py, diff --git a/docs/audit/generated/apple_target_map.md b/docs/audit/generated/apple_target_map.md index a2917936b..56f73c0b2 100644 --- a/docs/audit/generated/apple_target_map.md +++ b/docs/audit/generated/apple_target_map.md @@ -54,12 +54,12 @@ and ``docs/apple_backend.md`` (Metal 4 implementation-state review): | Family | Rows | apple_gpu fused | apple_cpu accelerate_native | |---|---:|---:|---:| -| tensor | 134 | 21 | 2 | +| tensor | 135 | 21 | 2 | | ga | 19 | 19 | 0 | | ebm | 14 | 14 | 0 | | m7 | 4 | 4 | 0 | -## tensor (134) +## tensor (135) | Op | apple_cpu status | cpu framework | cpu dtypes | cpu symbol | execution_kind | apple_gpu status | gpu framework | gpu dtypes | gpu symbol | gpu_dispatch | proof | |---|---|---|---|---|---|---|---|---|---|---|---| @@ -125,6 +125,7 @@ and ``docs/apple_backend.md`` (Metal 4 implementation-state review): | log_cosh_loss | reference | numpy_reference | fp32 | - | numpy_reference | compiled | - | fp32 | - | manifest | tests/unit/test_apple_*.py | | lookahead_sparse_attention | reference | numpy_reference | fp32 | - | numpy_reference | fused | - | fp32 | - | manifest | tests/unit/test_apple_*.py | | lora_linear | reference | numpy_reference | fp32 | - | numpy_reference | compiled | - | fp32 | - | manifest | tests/unit/test_apple_*.py | +| lstm_cell | reference | numpy_reference | fp32 | - | numpy_reference | compiled | - | fp32 | - | manifest | tests/unit/test_apple_*.py | | lu | reference | numpy_reference | fp32 | - | numpy_reference | compiled | - | fp32 | - | manifest | tests/unit/test_apple_*.py | | mae_loss | reference | numpy_reference | fp32 | - | numpy_reference | compiled | - | fp32 | - | manifest | tests/unit/test_apple_*.py | | masked_categorical | reference | numpy_reference | fp32 | - | numpy_reference | fused | - | fp32,int32 | tessera_apple_gpu_masked_categorical_f32 | driver | tests/unit/test_apple_gpu_ldt_loss_ops.py | diff --git a/docs/audit/generated/compiler_progress.csv b/docs/audit/generated/compiler_progress.csv index 1101c4d79..ff3caa05d 100644 --- a/docs/audit/generated/compiler_progress.csv +++ b/docs/audit/generated/compiler_progress.csv @@ -15,13 +15,13 @@ primitive_axis,lowering_rule,closed,480,480,0,"primitive contract axis; open mea primitive_axis,backend_kernel,mixed,99,480,381,"primitive contract axis; open means partial or planned, not necessarily missing API support",docs/audit/generated/s_series_status.md,Promote by backend/pathway; do not treat every target as an all-up compiler veto. integration,Verifier coverage,mixed,163,174,11,"no_verifier=11, real=163",docs/audit/generated/verifier_coverage.csv,"Add real verifier implementations for no_verifier ops, prioritizing native codegen lanes." integration,Direct test evidence,mixed,363,480,117,"covered_by_family=26, directly_tested=363, hardware_gated=4, structural_only=87",docs/audit/generated/test_coverage.csv,Convert structural_only and needs_direct_test rows into direct compare fixtures; keep hardware_gated tied to backend proof. -integration,Runtime execution matrix,closed,147,147,0,"apple_cpu=2, apple_gpu=16, cpu=2, nvidia_sm120=1, rocm=65, x86=61",docs/audit/generated/runtime_execution_matrix.csv,Add rows only when a launch path actually executes. +integration,Runtime execution matrix,closed,149,149,0,"apple_cpu=2, apple_gpu=16, cpu=2, nvidia_sm120=1, rocm=67, x86=61",docs/audit/generated/runtime_execution_matrix.csv,Add rows only when a launch path actually executes. integration,Runtime ABI symbols,mixed,397,650,253,"apple=557, nvidia=7, rocm=10, x86=76",docs/audit/generated/runtime_abi.csv,Reduce stub-only ABI rows where a backend claims native execution. integration,Audited repo surfaces,mixed,31,58,27,"archived=4, compile_only=12, runnable=31, runnable_optional=1, scaffold=10",docs/audit/generated/surface_status.csv,Graduate compile_only/scaffold entries that exercise compiler pathways; archive dead surfaces. -codegen_pathway,Apple CPU,closed,173,173,0,"runtime: apple_cpu=2; target_map: fused=2, reference=169",docs/audit/generated/apple_target_map.csv,Keep as regression baseline for CPU value-call/runtime ABI. -codegen_pathway,Apple GPU,mixed,186,187,1,"runtime: apple_gpu=16; target_map: absent=1, compiled=101, fused=58, hardware_verified=11",docs/audit/generated/apple_target_map.csv,Close the remaining absent target-map lane or document why it is host-only. +codegen_pathway,Apple CPU,closed,174,174,0,"runtime: apple_cpu=2; target_map: fused=2, reference=170",docs/audit/generated/apple_target_map.csv,Keep as regression baseline for CPU value-call/runtime ABI. +codegen_pathway,Apple GPU,mixed,187,188,1,"runtime: apple_gpu=16; target_map: absent=1, compiled=102, fused=58, hardware_verified=11",docs/audit/generated/apple_target_map.csv,Close the remaining absent target-map lane or document why it is host-only. codegen_pathway,x86 / CPU,closed,2,2,0,runtime: cpu=2,docs/audit/generated/runtime_execution_matrix.csv,Keep native CPU and numpy reference lanes separate in runtime proofs. -codegen_pathway,ROCm / HIP,closed,101,101,0,"runtime: rocm=65; target_map: compiled=34, hardware_verified=2",docs/audit/generated/rocm_target_map.csv,Close the artifact-only target-map tail and preserve CDNA as hardware-gated. +codegen_pathway,ROCm / HIP,closed,103,103,0,"runtime: rocm=67; target_map: compiled=34, hardware_verified=2",docs/audit/generated/rocm_target_map.csv,Close the artifact-only target-map tail and preserve CDNA as hardware-gated. codegen_pathway,CUDA / NVIDIA,mixed,1,40,39,runtime: nvidia_sm120=1; target_map: artifact_only=39,docs/audit/generated/nvidia_sm90_target_map.csv,"Promote artifact-only rows with execute-and-compare, starting from sm_120 matmul adjacency and attention." open_work,Target IR native/fused codegen,mixed,309,315,6,"compiled=196, fused=71, hardware_verified=13, not_applicable=29, reference=6",docs/audit/generated/support_table.csv,Promote high-use reference rows into native/fused Target IR or mark intentional reference-only lanes. open_work,backend_kernel,mixed,99,480,381,"primitive contract axis; open means partial or planned, not necessarily missing API support",docs/audit/generated/s_series_status.md,Promote by backend/pathway; do not treat every target as an all-up compiler veto. diff --git a/docs/audit/generated/compiler_progress.md b/docs/audit/generated/compiler_progress.md index 153697e9e..1a7a7f2a0 100644 --- a/docs/audit/generated/compiler_progress.md +++ b/docs/audit/generated/compiler_progress.md @@ -42,7 +42,7 @@ A row is not marked incomplete merely because Apple, x86, ROCm, and CUDA are not |---|---|---:|---:|---:|---|---| | `Verifier coverage` | mixed | 163 | 174 | 11 | no_verifier=11, real=163 | Add real verifier implementations for no_verifier ops, prioritizing native codegen lanes. | | `Direct test evidence` | mixed | 363 | 480 | 117 | covered_by_family=26, directly_tested=363, hardware_gated=4, structural_only=87 | Convert structural_only and needs_direct_test rows into direct compare fixtures; keep hardware_gated tied to backend proof. | -| `Runtime execution matrix` | closed | 147 | 147 | 0 | apple_cpu=2, apple_gpu=16, cpu=2, nvidia_sm120=1, rocm=65, x86=61 | Add rows only when a launch path actually executes. | +| `Runtime execution matrix` | closed | 149 | 149 | 0 | apple_cpu=2, apple_gpu=16, cpu=2, nvidia_sm120=1, rocm=67, x86=61 | Add rows only when a launch path actually executes. | | `Runtime ABI symbols` | mixed | 397 | 650 | 253 | apple=557, nvidia=7, rocm=10, x86=76 | Reduce stub-only ABI rows where a backend claims native execution. | | `Audited repo surfaces` | mixed | 31 | 58 | 27 | archived=4, compile_only=12, runnable=31, runnable_optional=1, scaffold=10 | Graduate compile_only/scaffold entries that exercise compiler pathways; archive dead surfaces. | @@ -50,10 +50,10 @@ A row is not marked incomplete merely because Apple, x86, ROCm, and CUDA are not | Item | Status | Ready | Total | Open | Detail | Next | |---|---|---:|---:|---:|---|---| -| `Apple CPU` | closed | 173 | 173 | 0 | runtime: apple_cpu=2; target_map: fused=2, reference=169 | Keep as regression baseline for CPU value-call/runtime ABI. | -| `Apple GPU` | mixed | 186 | 187 | 1 | runtime: apple_gpu=16; target_map: absent=1, compiled=101, fused=58, hardware_verified=11 | Close the remaining absent target-map lane or document why it is host-only. | +| `Apple CPU` | closed | 174 | 174 | 0 | runtime: apple_cpu=2; target_map: fused=2, reference=170 | Keep as regression baseline for CPU value-call/runtime ABI. | +| `Apple GPU` | mixed | 187 | 188 | 1 | runtime: apple_gpu=16; target_map: absent=1, compiled=102, fused=58, hardware_verified=11 | Close the remaining absent target-map lane or document why it is host-only. | | `x86 / CPU` | closed | 2 | 2 | 0 | runtime: cpu=2 | Keep native CPU and numpy reference lanes separate in runtime proofs. | -| `ROCm / HIP` | closed | 101 | 101 | 0 | runtime: rocm=65; target_map: compiled=34, hardware_verified=2 | Close the artifact-only target-map tail and preserve CDNA as hardware-gated. | +| `ROCm / HIP` | closed | 103 | 103 | 0 | runtime: rocm=67; target_map: compiled=34, hardware_verified=2 | Close the artifact-only target-map tail and preserve CDNA as hardware-gated. | | `CUDA / NVIDIA` | mixed | 1 | 40 | 39 | runtime: nvidia_sm120=1; target_map: artifact_only=39 | Promote artifact-only rows with execute-and-compare, starting from sm_120 matmul adjacency and attention. | ## Open Work Summary diff --git a/docs/audit/generated/rocm_target_map.csv b/docs/audit/generated/rocm_target_map.csv index 78b9a8a31..af318ee86 100644 --- a/docs/audit/generated/rocm_target_map.csv +++ b/docs/audit/generated/rocm_target_map.csv @@ -6,7 +6,7 @@ softmax,activation,compiled,"fp32,fp16,bf16",hipcc≥7.2.4,,,,"Row-wise stable s softmax_safe,activation,compiled,fp32,hipcc≥7.2.4,,,,Stable reduction softmax_safe — max-shifted reduce (max/sum) + unary exp/log lane. Executes via runtime.launch() (rocm_stable_reduce_compiled). attn_sliding_window,attention,compiled,"fp16,bf16",hipcc≥7.2.4,,,,"Mistral sliding-window via the flash_attn WMMA kernel (sliding_window attr, causal band of width W; KV-tile skip). Executes via runtime.launch() (window kwarg)." deepseek_sparse_attention,attention,compiled,fp32,hipcc≥7.2.4,,,,"DeepSeek/NSA composition — sliding + compressed-block branches remain reference compositions while the top-k branch uses the GPU-resident top-k selector plus selected-block sparse-attention kernel when ROCm is available. Executes via runtime.launch() (rocm_sparse_attn_compiled), with exact reference fallback off hardware." -flash_attn,attention,hardware_verified,"fp16,bf16",hipcc≥7.2.4,,0.6500,,"RDNA 3.5 WMMA flash-attention forward executes on the AMD GPU through the shipped libtessera_rocm_flash_attn.so symbols (tessera_rocm_wmma_flash_attn_{f16,bf16}, HIPRTC-compiled for the device arch at load); ROCm 7.2.4. The second op after matmul to run natively on a non-Apple backend. Numerically validated vs a numpy attention reference by the execute_compare_fixture." +flash_attn,attention,hardware_verified,"fp16,bf16",hipcc≥7.2.4,,0.6500,,"RDNA 3.5 WMMA flash-attention forward executes on the AMD GPU through the shipped libtessera_rocm_flash_attn.so symbols (tessera_rocm_wmma_flash_attn_{f16,bf16}, HIPRTC-compiled for the device arch at load); ROCm 7.2.4. The second op after matmul to run natively on a non-Apple backend. Numerically validated vs a numpy attention reference by the execute_compare_fixture. The FA-2 BACKWARD (dQ/dK/dV) also executes on gfx1151 via the compiler-generated rocm_flash_attn_bwd_compiled lane (generate-wmma-flash-attn-bwd-kernel -> fa_pre/fa_dkdv/fa_dq; MHA + GQA/MQA + additive attn_bias, scale+causal), validated vs autodiff vjp_flash_attn — see the runtime_execution_matrix." gated_attention,attention,compiled,"fp16,bf16",hipcc≥7.2.4,,,,Softmax attention × a learned gate — the WMMA flash_attn kernel + an elementwise sigmoid-gate multiply. Executes via runtime.launch() (rocm_exotic_attn_compiled). gated_deltanet,attention,compiled,"fp32,fp16,bf16",hipcc≥7.2.4,,,,"Gated/delta linear-attention recurrence as a causal SEQUENTIAL-SCAN kernel (generate-rocm-deltanet-kernel: one workgroup per (b,h), one thread per value-column, LDS state) — the first recurrent compiled ROCm kernel. erase/gate/beta/decay flags. Executes via runtime.launch() (rocm_deltanet_compiled)." gqa_attention,attention,compiled,"fp16,bf16",hipcc≥7.2.4,,,,"GQA/MQA via the flash_attn WMMA kernel (gqa directive attr; fwd+bwd, grouped K/V; runtime detects from operand shapes). Executes on gfx1151 via runtime.launch() (rocm_flash_attn_compiled); no shipped C-ABI symbol." diff --git a/docs/audit/generated/runtime_execution_matrix.csv b/docs/audit/generated/runtime_execution_matrix.csv index 25f8c7683..91cd0882e 100644 --- a/docs/audit/generated/runtime_execution_matrix.csv +++ b/docs/audit/generated/runtime_execution_matrix.csv @@ -44,10 +44,12 @@ rocm,rocm_ebm_langevin_compiled,native_gpu,1,rocm_ebm_langevin_compiled,success, rocm,rocm_ebm_loss_compiled,native_gpu,1,rocm_ebm_loss_compiled,success,hip_runtime,"ROCm EBM/diffusion loss artifact runs score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_loss. f32." rocm,rocm_exotic_attn_compiled,native_gpu,1,rocm_exotic_attn_compiled,success,hip_runtime,"ROCm exotic-attention artifact composes the COMPILER-GENERATED WMMA flash_attn kernel with the WMMA GEMM kernel (MLA latent projections) + an elementwise gate — gated_attention, mla_decode, mla_decode_fused — and routes mla_decode_step through the DK1 absorbed-latent decode kernel against stdlib.attention.mla_decode_step. f16/bf16, f32 softmax+accumulate." rocm,rocm_fft_compiled,native_gpu,1,rocm_fft_compiled,success,hip_runtime,"ROCm FFT artifact runs fft / ifft / rfft / irfft over any axis length on the COMPILER-GENERATED one-thread-per-bin DFT kernel (generate-rocm-dft-kernel -> ROCDL, cos/sin twiddles) on gfx1151 + r2c/c2r pack-unpack + plan scale. Direct DFT (radix-2/Bluestein perf is a follow-up). complex64/f32." +rocm,rocm_flash_attn_bwd_compiled,native_gpu,1,rocm_flash_attn_bwd_compiled,success,hip_runtime,"ROCm flash_attn backward artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 backward: tessera-opt expands one tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/fa_dq WMMA kernels serialized to hsaco in-process, then HIP launches them in sequence to produce dQ/dK/dV. O is recomputed via the forward lane (nothing saved from forward). The reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA (gqa dkdv atomic-accumulates dK/dV across the group) + additive attn_bias (S=scale*QK+bias in the recompute), scale + causal, f16/bf16 storage, f32 accumulate." rocm,rocm_flash_attn_compiled,native_gpu,1,rocm_flash_attn_compiled,success,hip_runtime,"ROCm flash_attn artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 forward: tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. The attention analog of the compiled GEMM lane (rocm_compiled)." rocm,rocm_fpquant_compiled,native_gpu,1,rocm_fpquant_compiled,success,hip_runtime,ROCm fpquant artifact runs quantize/dequantize fp8 / fp6 / fp4: per-tensor scale + grid-snap on the COMPILER-GENERATED fpquant kernel (generate-rocm-fpquant-kernel: log2/exp2/roundeven -> ROCDL). ROCm mirror of x86_fpquant. f32. rocm,rocm_image_affine_compiled,native_gpu,1,rocm_image_affine_compiled,success,hip_runtime,"ROCm image affine artifact runs image_normalize as (x-mean)/std: layout and per-channel broadcast on host, sub/div on generated ROCm binary kernels. f32, matches tessera.ops." rocm,rocm_intquant_compiled,native_gpu,1,rocm_intquant_compiled,success,hip_runtime,"ROCm intquant artifact runs quantize/dequantize int8/int4 and fake_quantize: qparam selection and int8 container conversion on host, round/max/min/mul on generated ROCm unary/binary kernels. int4 is signed int4 values in int8 containers." +rocm,rocm_kv_cache_compiled,native_gpu,1,rocm_kv_cache_compiled,success,hip_runtime,"ROCm KV-cache paged-movement lane realizes kv_cache append/read/prune over a resident cache buffer (max_seq, H, D) by COMPOSING the COMPILER-GENERATED gfx1151 scatter (append row write) + masked-gather (read/prune) kernels with host page-index math. quantize_kv rides the intquant lane. f32, matches the KVCacheHandle append/read/prune reference." rocm,rocm_lamb_compiled,native_gpu,1,rocm_lamb_compiled,success,hip_runtime,"ROCm LAMB lane runs the COMPILER-GENERATED gfx1151 adam kernel (lr=1/wd=0) then applies the per-tensor trust ratio ‖p‖/‖update‖ on host. f32, matches optim.lamb." rocm,rocm_linalg_compiled,native_gpu,1,rocm_linalg_compiled,success,hip_runtime,"ROCm linalg lane (cholesky / tri_solve / cholesky_solve / lu / qr / svd) runs the COMPILER-GENERATED gfx1151 kernels (generate-rocm-cholesky / tri-solve / lu / qr / svd-kernel, one thread per matrix or matrix/RHS-column) HIP-launched; cholesky_solve composes two triangular solves. f32, matches numpy." rocm,rocm_linear_attn_compiled,native_gpu,1,rocm_linear_attn_compiled,success,hip_runtime,"ROCm linear-attention-family artifact runs the COMPILER-GENERATED RDNA WMMA forward (quadratic-parallel form, no softmax): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Handles linear_attn + lightning_attention (identity+decay) + retention (x²+decay) by op name." diff --git a/docs/audit/generated/runtime_execution_matrix.md b/docs/audit/generated/runtime_execution_matrix.md index 02255d07e..a38cadd85 100644 --- a/docs/audit/generated/runtime_execution_matrix.md +++ b/docs/audit/generated/runtime_execution_matrix.md @@ -58,10 +58,12 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `rocm` | `rocm_ebm_loss_compiled` | `rocm_ebm_loss_compiled` | `native_gpu` | `hip_runtime` | ROCm EBM/diffusion loss artifact runs score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_loss. f32. | | `rocm` | `rocm_exotic_attn_compiled` | `rocm_exotic_attn_compiled` | `native_gpu` | `hip_runtime` | ROCm exotic-attention artifact composes the COMPILER-GENERATED WMMA flash_attn kernel with the WMMA GEMM kernel (MLA latent projections) + an elementwise gate — gated_attention, mla_decode, mla_decode_fused — and routes mla_decode_step through the DK1 absorbed-latent decode kernel against stdlib.attention.mla_decode_step. f16/bf16, f32 softmax+accumulate. | | `rocm` | `rocm_fft_compiled` | `rocm_fft_compiled` | `native_gpu` | `hip_runtime` | ROCm FFT artifact runs fft / ifft / rfft / irfft over any axis length on the COMPILER-GENERATED one-thread-per-bin DFT kernel (generate-rocm-dft-kernel -> ROCDL, cos/sin twiddles) on gfx1151 + r2c/c2r pack-unpack + plan scale. Direct DFT (radix-2/Bluestein perf is a follow-up). complex64/f32. | +| `rocm` | `rocm_flash_attn_bwd_compiled` | `rocm_flash_attn_bwd_compiled` | `native_gpu` | `hip_runtime` | ROCm flash_attn backward artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 backward: tessera-opt expands one tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/fa_dq WMMA kernels serialized to hsaco in-process, then HIP launches them in sequence to produce dQ/dK/dV. O is recomputed via the forward lane (nothing saved from forward). The reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA (gqa dkdv atomic-accumulates dK/dV across the group) + additive attn_bias (S=scale*QK+bias in the recompute), scale + causal, f16/bf16 storage, f32 accumulate. | | `rocm` | `rocm_flash_attn_compiled` | `rocm_flash_attn_compiled` | `native_gpu` | `hip_runtime` | ROCm flash_attn artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 forward: tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. The attention analog of the compiled GEMM lane (rocm_compiled). | | `rocm` | `rocm_fpquant_compiled` | `rocm_fpquant_compiled` | `native_gpu` | `hip_runtime` | ROCm fpquant artifact runs quantize/dequantize fp8 / fp6 / fp4: per-tensor scale + grid-snap on the COMPILER-GENERATED fpquant kernel (generate-rocm-fpquant-kernel: log2/exp2/roundeven -> ROCDL). ROCm mirror of x86_fpquant. f32. | | `rocm` | `rocm_image_affine_compiled` | `rocm_image_affine_compiled` | `native_gpu` | `hip_runtime` | ROCm image affine artifact runs image_normalize as (x-mean)/std: layout and per-channel broadcast on host, sub/div on generated ROCm binary kernels. f32, matches tessera.ops. | | `rocm` | `rocm_intquant_compiled` | `rocm_intquant_compiled` | `native_gpu` | `hip_runtime` | ROCm intquant artifact runs quantize/dequantize int8/int4 and fake_quantize: qparam selection and int8 container conversion on host, round/max/min/mul on generated ROCm unary/binary kernels. int4 is signed int4 values in int8 containers. | +| `rocm` | `rocm_kv_cache_compiled` | `rocm_kv_cache_compiled` | `native_gpu` | `hip_runtime` | ROCm KV-cache paged-movement lane realizes kv_cache append/read/prune over a resident cache buffer (max_seq, H, D) by COMPOSING the COMPILER-GENERATED gfx1151 scatter (append row write) + masked-gather (read/prune) kernels with host page-index math. quantize_kv rides the intquant lane. f32, matches the KVCacheHandle append/read/prune reference. | | `rocm` | `rocm_lamb_compiled` | `rocm_lamb_compiled` | `native_gpu` | `hip_runtime` | ROCm LAMB lane runs the COMPILER-GENERATED gfx1151 adam kernel (lr=1/wd=0) then applies the per-tensor trust ratio ‖p‖/‖update‖ on host. f32, matches optim.lamb. | | `rocm` | `rocm_linalg_compiled` | `rocm_linalg_compiled` | `native_gpu` | `hip_runtime` | ROCm linalg lane (cholesky / tri_solve / cholesky_solve / lu / qr / svd) runs the COMPILER-GENERATED gfx1151 kernels (generate-rocm-cholesky / tri-solve / lu / qr / svd-kernel, one thread per matrix or matrix/RHS-column) HIP-launched; cholesky_solve composes two triangular solves. f32, matches numpy. | | `rocm` | `rocm_linear_attn_compiled` | `rocm_linear_attn_compiled` | `native_gpu` | `hip_runtime` | ROCm linear-attention-family artifact runs the COMPILER-GENERATED RDNA WMMA forward (quadratic-parallel form, no softmax): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Handles linear_attn + lightning_attention (identity+decay) + retention (x²+decay) by op name. | @@ -218,10 +220,12 @@ nvidia_sm80, nvidia_sm90, nvidia_sm100, rocm_gfx90a, rocm_gfx940, rocm_gfx942, r | `rocm_ebm_loss_compiled` | AMD GPU RDNA EBM/diffusion loss (score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance) — diff/square + reductions on the gfx1151 binary + reduce kernels, host structure. ROCm mirror of x86_ebm_loss. f32 | | `rocm_exotic_attn_compiled` | AMD GPU RDNA exotic-attention compositions (gated_attention, mla_decode, mla_decode_fused, mla_decode_step absorbed-latent decode) built by COMPOSING the COMPILER-GENERATED WMMA flash_attn kernel + the WMMA GEMM kernel (MLA latent projections) + an elementwise gate, plus the DK1 generated absorbed-latent ROCm decode kernel; f16/bf16 storage, f32 softmax+accumulate. The block-sparse deepseek variant stays artifact_only | | `rocm_fft_compiled` | AMD GPU RDNA spectral FFT (fft / ifft / rfft / irfft) the Tessera compiler GENERATES (generate-rocm-dft-kernel -> ROCDL -> hsaco; one thread per output bin, cos/sin twiddles), then HIP launches it. Direct DFT (any length) on gfx1151 + r2c/c2r pack-unpack + plan scale (radix-2/Bluestein perf is a follow-up). complex64/f32 | +| `rocm_flash_attn_bwd_compiled` | AMD GPU RDNA WMMA FA-2 BACKWARD the Tessera compiler GENERATES (generate-wmma-flash-attn-bwd-kernel -> three fa_pre/fa_dkdv/fa_dq WMMA kernels -> hsaco), launched in sequence to produce dQ/dK/dV; O is recomputed via the forward lane (nothing saved from forward). MHA + GQA/MQA (grouped dkdv atomic-accumulates dK/dV) + additive attn_bias; f16/bf16 storage, f32 accumulate; the reverse-mode analog of rocm_flash_attn_compiled | | `rocm_flash_attn_compiled` | AMD GPU RDNA WMMA FA-2 forward the Tessera compiler GENERATES (generate-wmma-flash-attn-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it. f16/bf16 storage, f32 softmax + accumulate; the attention analog of rocm_compiled | | `rocm_fpquant_compiled` | AMD GPU RDNA low-precision float quantize (quantize/dequantize fp8 / fp6 / fp4) — grid-snap on generate-rocm-fpquant-kernel (log2/exp2/roundeven) + per-tensor scale. ROCm mirror of x86_fpquant. f32 | | `rocm_image_affine_compiled` | AMD GPU RDNA image affine preprocessing — image_normalize as sub/div on generated ROCm binary kernels with host layout and per-channel broadcast | | `rocm_intquant_compiled` | AMD GPU RDNA integer quantization — qparam selection and int8 container conversion around generated ROCm unary/binary kernels; covers int8 and signed int4 values stored in int8 containers | +| `rocm_kv_cache_compiled` | AMD GPU RDNA KV-cache paged-movement lane — kv_cache append/read/prune over a resident cache buffer by composing the gfx1151 scatter (write) + gather (read/prune) kernels; host page-index math. f32, matches the KVCacheHandle reference | | `rocm_lamb_compiled` | AMD GPU RDNA LAMB — COMPILER-GENERATED gfx1151 adam kernel (lr=1/wd=0) + host per-tensor trust ratio ‖p‖/‖update‖. f32 | | `rocm_linalg_compiled` | AMD GPU RDNA dense linear algebra (cholesky / tri_solve / cholesky_solve / lu / qr / svd) — COMPILER-GENERATED gfx1151 kernels (generate-rocm-cholesky / tri-solve / lu / qr / svd-kernel, one thread per matrix or matrix/RHS-column) HIP-launched; cholesky_solve = two triangular solves. f32 | | `rocm_linear_attn_compiled` | AMD GPU RDNA WMMA linear-attention forward the Tessera compiler GENERATES (generate-wmma-linear-attn-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it. Quadratic-parallel form O = (φ(Q)φ(K)ᵀ ⊙ causal [⊙ λ^(i-j)]) @ V, NO softmax; f16/bf16 storage, f32 accumulate. Handles tessera.linear_attn + the decay-masked siblings tessera.lightning_attention (identity+decay) and tessera.retention (x²+decay) by op name | diff --git a/docs/audit/generated/s_series_status.md b/docs/audit/generated/s_series_status.md index d7fd5e7b4..17a627631 100644 --- a/docs/audit/generated/s_series_status.md +++ b/docs/audit/generated/s_series_status.md @@ -30,7 +30,7 @@ The registry-level `backend_kernel` axis is deliberately conservative and should | `x86` | 327 | 325 | 2 | 0 | 153 | | `apple_cpu` | 361 | 2 | 359 | 0 | 119 | | `apple_gpu` | 203 | 196 | 3 | 4 | 277 | -| `rocm` | 337 | 335 | 2 | 0 | 143 | +| `rocm` | 340 | 338 | 2 | 0 | 140 | | `nvidia_sm80` | 66 | 0 | 0 | 66 | 414 | | `nvidia_sm90` | 98 | 0 | 2 | 96 | 382 | | `nvidia_sm100` | 66 | 0 | 0 | 66 | 414 | diff --git a/docs/audit/op_target_conformance.csv b/docs/audit/op_target_conformance.csv index ea435b5de..f2f075575 100644 --- a/docs/audit/op_target_conformance.csv +++ b/docs/audit/op_target_conformance.csv @@ -13,7 +13,7 @@ kv_cache_read,apple_cpu,partial,complete,complete,complete,complete,complete,par kv_cache_read,apple_gpu,missing,complete,complete,complete,complete,complete,missing,partial,hardware_smoke kv_cache_read,cpu,partial,complete,complete,complete,complete,complete,partial,complete, kv_cache_read,nvidia,missing,complete,complete,complete,missing,missing,missing,partial,codegen -kv_cache_read,rocm,missing,complete,complete,complete,missing,missing,missing,partial,codegen +kv_cache_read,rocm,complete,complete,complete,complete,complete,complete,complete,complete,toolchain matmul,apple_cpu,complete,complete,complete,complete,complete,complete,complete,complete,hardware_smoke matmul,apple_gpu,complete,complete,complete,complete,complete,complete,complete,complete,hardware_smoke matmul,cpu,partial,complete,complete,complete,complete,complete,partial,complete, diff --git a/docs/audit/op_target_conformance.md b/docs/audit/op_target_conformance.md index 9c37f07f7..e316dbaaa 100644 --- a/docs/audit/op_target_conformance.md +++ b/docs/audit/op_target_conformance.md @@ -26,11 +26,11 @@ Audit response to [docs/audit/compiler/COMPILER_AUDIT.md](compiler/COMPILER_AUDI | Overall (weakest column wins) | Count | |---|---:| -| ✅ `complete` | 10 | +| ✅ `complete` | 11 | | ⚙️ `partial` | 15 | | ⚠️ `artifact_only` | 0 | | 📋 `planned` | 0 | -| ❌ `missing` | 10 | +| ❌ `missing` | 9 | | **total cells** | **35** | ## `matmul` @@ -109,5 +109,5 @@ _fused MSL kernel on apple_gpu (single-kernel scores); compose elsewhere_ | `apple_cpu` | ⚙️ | ✅ | ✅ | ✅ | ✅ | ✅ | ⚙️ | ⚙️ | `hardware_smoke` — Apple silicon required for native execution | | | `apple_gpu` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ⚙️ | `hardware_smoke` — Apple silicon required for native execution | | | `nvidia` | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ⚙️ | `codegen` — no backend_manifest entry for 'kv_cache_read' on 'nvidia' | | -| `rocm` | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ⚙️ | `codegen` — no backend_manifest entry for 'kv_cache_read' on 'rocm' | | +| `rocm` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | `toolchain` — hipcc not on PATH (ROCm 7.2.4 not installed) | | diff --git a/docs/rocm_mfma_kernel_inventory.md b/docs/rocm_mfma_kernel_inventory.md index ffb9c5ef0..e61e05cec 100644 --- a/docs/rocm_mfma_kernel_inventory.md +++ b/docs/rocm_mfma_kernel_inventory.md @@ -2,7 +2,7 @@ status: Informative classification: Reference / Kernel Inventory authority: Companion to Phase H ROCm backend pre-work -last_updated: 2026-06-23 +last_updated: 2026-07-10 --- # ROCm 7.2.4 MFMA / WMMA Kernel Inventory @@ -12,14 +12,25 @@ last_updated: 2026-06-23 > HIP 7.2.4. Companion to `docs/nvidia_cuda13_kernel_inventory.md` (parallel > coverage tracking) and `docs/apple_gpu_kernel_inventory.md`. > -> **Execution status (2026-06-23):** no longer fully hardware-free. On the +> **Execution status (2026-07-10):** far past hardware-free. On the > **RDNA 3.5 `gfx1151`** (Strix Halo APU — Ryzen AI Max+ 395 / Radeon 8060S) -> **two ops now execute on real silicon** and are `hardware_verified` in -> `backend_manifest`: **`matmul`/`gemm`** (WMMA GEMM — `libtessera_rocm_gemm.so`, -> runtime `launch()` lane, measured perf ladder) and **`flash_attn`** (WMMA FA-2 -> forward — `libtessera_rocm_flash_attn.so`, online softmax, causal, ragged). -> Both have execute-compare fixtures. Everything else on every arch remains -> `artifact_only`. See §7 and +> **dozens of compiler-generated HIP kernels now execute on real silicon** +> through `runtime.launch()` (`execution_mode="hip_runtime"`), spanning nearly +> every §5 family: matmul/GEMM + the GEMM family, flash-attention + the exotic / +> sparse / linear / delta attention lanes, norm / activation / RoPE / ALiBi, +> optimizers (adam/adamw/lion/muon/lamb), RNG, FFT/spectral, selective-SSM +> (Mamba2, incl. device backward), quantization, the loss families, the +> reduction / scan / sort / scatter / elementwise lanes, and the M7 complex / +> Clifford / EBM families. **`matmul`/`gemm`** and **`flash_attn`** were the +> first two (`hardware_verified` in `backend_manifest`, with shipped C-ABI +> symbols + measured perf ladder for GEMM); the rest execute as correctness-first +> **`compiled`** lanes (execute-vs-reference verified, no perf ladder yet). +> +> **Status truth is the generated matrix, not this prose (Decision #26):** +> [`docs/audit/generated/runtime_execution_matrix.md`](audit/generated/runtime_execution_matrix.md) +> is the drift-gated source for which `(op, target)` rows execute and by which +> lane. This inventory is the *kernel contract* (shapes, dtypes, MFU targets); +> read it against that matrix for live execution status. See §7, §9, and > `docs/audit/backend/rocm/{ROCM_AUDIT,STRIX_HALO_EXECUTION_PLAN}.md`. This document is the **authoritative kernel inventory** for the @@ -157,11 +168,36 @@ as accelerated matrix dtypes. --- -## 5. Planned fused kernel inventory +## 5. Fused kernel inventory Per-kernel MFMA shapes live in `_ROCM_KERNEL_MFMA_SHAPES` (`backend_manifest.py`). MFU targets in `_ROCM_KERNEL_MFU`. +> **Execution note (2026-07-10).** The tables below are the kernel *contract* +> (MFMA/WMMA shape, dtype variants, MFU target). They are **no longer "planned"** +> on `gfx1151`: as of 2026-07-10 the majority of these families execute on real +> RDNA 3.5 silicon as compiler-generated HIP `compiled` lanes (correctness-first, +> execute-vs-reference verified). The stable lane ids per family: +> +> | §5 family | Executing gfx1151 lane(s) | +> |---|---| +> | 5.1 Matmul / contraction | `rocm_compiled` (WMMA GEMM), `rocm_matmul_family_compiled`, `rocm_dequant_gemm_compiled`, `rocm_conv_compiled` | +> | 5.2 Attention | `rocm_flash_attn_compiled`, `rocm_exotic_attn_compiled`, `rocm_sparse_attn_compiled`, `rocm_linear_attn_compiled`, `rocm_deltanet_compiled` | +> | 5.3 Fused chains | `rocm_silu_mul_compiled`, `rocm_softmax_compiled` (+ SwiGLU via matmul-family) | +> | 5.4 Norm / activation / position | `rocm_norm_compiled`, `rocm_normcompose_compiled`, `rocm_activation_compiled`, `rocm_rope_compiled`, `rocm_alibi_compiled`, `rocm_softcap_compiled` | +> | 5.5 Optimizer | `rocm_optimizer_compiled`, `rocm_muon_compiled`, `rocm_lamb_compiled` | +> | 5.6 KV-cache / paged | `rocm_kv_cache_compiled` (append/read/prune paged-movement core on the scatter+gather kernels); `quantize_kv` via `rocm_intquant_compiled` | +> | 5.7 RNG / sampling | `rocm_rng_compiled` | +> | 5.8 Spectral | `rocm_fft_compiled`, `rocm_spectral_compiled` | +> | 5.9 Recurrent / SSM | `rocm_selective_ssm_compiled` (incl. device backward), `rocm_deltanet_compiled`; single-step `lstm`/`gru`/`simple_rnn` cells via `rocm_structured_compute_compiled` (host-structured) | +> +> Quantization (`rocm_intquant`/`fpquant`/`nvfp4`), the loss families +> (`rocm_loss`/`class_loss`/`binary_loss`/`rl_loss`/`metric_loss`/`ebm_loss`), the +> reduction/scan/sort/scatter/elementwise lanes, and the §10 M7 complex/Clifford/EBM +> families also execute. **`docs/audit/generated/runtime_execution_matrix.md` is the +> authoritative live list; the mapping above is a dated snapshot, not maintained +> truth.** + ### 5.1 Matmul / contraction family | Kernel | MFMA (M, N, K, K_blocks) | Dtypes | MFU (gfx942 / gfx950) | Notes | @@ -345,47 +381,50 @@ GFX12 scalar prefetch / load notes tracked for future scheduler work: | Gate | What it means | Status | |---|---|---| -| `artifact_only` | Target IR + AMDGCN intrinsic text are well-formed; lit fixtures pass FileCheck; no execution | **all entries except the one below** | +| `artifact_only` | Target IR + AMDGCN intrinsic text are well-formed; lit fixtures pass FileCheck; no execution | CDNA (gfx90a/94x/950) MFMA entries; a few RDNA rows with no execute-compare fixture yet (e.g. KV-cache) | | `compileable` | `hipcc -S --offload-arch=…` (or `llc -mcpu=…`) accepts the kernel; produces a valid object; **without execution** | reachable now on the box (`rocdl_emit.py` + `llc` proven for gfx1100/gfx1151) | -| `executable` | The kernel loads on a real GPU and produces correct output vs CPU reference | ✅ **`matmul`/`gemm` + `flash_attn` WMMA on `gfx1151`** (below) | -| `fused` | Performance characterized against the MFU targets in §5 | not yet — gfx1151 has a *measured perf ladder* but no MFU-target sign-off; CDNA MFU targets need MI300X/MI325X | - -**The exception — `gfx1151` (RDNA 3.5) WMMA matmul executes today.** As of -2026-06-23 the `matmul`/`gemm` WMMA path on the Strix Halo box is -`hardware_verified` in `backend_manifest`: - -- shipped `libtessera_rocm_gemm.so` exporting `tessera_rocm_wmma_gemm_{f16,bf16}` - (HIPRTC-compiles the RDNA WMMA kernel for the device arch at load — no - hipcc-as-compiler); -- wired into `runtime.launch()` as the executable `("rocm", "rocm_wmma")` row - (`hip_runtime`) in the generated `runtime_execution_matrix`; -- execute-compare fixture `tests/unit/test_rocm_wmma_runtime_symbol.py` (f16/bf16, - ragged + K-looped shapes) vs a numpy reference; -- a measured GEMM perf ladder (register blocking / LDS staging / software - pipelining / APU zero-copy) — see `STRIX_HALO_EXECUTION_PLAN.md` Stage F. - -**The second exception — `flash_attn` (RDNA WMMA FA-2 forward).** As of -2026-06-23 `flash_attn` also executes on `gfx1151`, the second op after matmul to -run natively on a non-Apple backend: - -- shipped `libtessera_rocm_flash_attn.so` exporting - `tessera_rocm_wmma_flash_attn_{f16,bf16}` (HIPRTC-compiled per head_dim at load); -- FA-2 forward, single wave per (query-tile-of-16, b·h): **both QK^T and P@V on - 16×16×16 WMMA**, online (running max/sum) softmax, scores + output accumulator - staged in LDS, causal masking + ragged Sq/Sk; head_dim a multiple of 16; -- execute-compare fixture `tests/unit/test_rocm_flash_attn_runtime_symbol.py` - vs a numpy attention reference (f16/bf16, head_dim 16/32/64/128, multi - batch/head, ragged, causal). Measured maxerr ~1e-4 (f16) on gfx1151. - -Honest scope (Decision #25): both exceptions are **one arch × {fp16, bf16}**; -flash_attn is **forward only, no perf ladder** (the correctness-first "rung 0" of -attention) and has no `runtime.launch()` lane yet. They do **not** flip the -per-primitive `backend_kernel` axis (that needs *all* targets `hardware_verified`), -and every other kernel in §5 (the rest of the attention family, fused chains, -optimizer/KV/RNG/spectral) stays `artifact_only` on every ROCm arch. CDNA MFMA -entries remain hardware-free pending MI300-class silicon; Sprint H-4 lit fixtures -validate their IR + MFMA patterns; `hipcc`/`llc` compile-only validation promotes -them to `compileable`. +| `executable` | The kernel loads on a real GPU and produces correct output vs a CPU/numpy reference | ✅ **the majority of §5 on `gfx1151`** — see the generated matrix | +| `fused` | Performance characterized against the MFU targets in §5 | only `matmul`/`gemm` has a *measured perf ladder*; **no MFU-target sign-off anywhere**; CDNA MFU targets need MI300X/MI325X | + +**Status truth is the generated matrix, not this table (Decision #26).** +[`docs/audit/generated/runtime_execution_matrix.md`](audit/generated/runtime_execution_matrix.md) +is the drift-gated source for which `(op, target)` rows execute. As of 2026-07-10 +the ROCm backend runs **dozens of compiler-generated HIP kernels natively on +`gfx1151`** through `runtime.launch()` (`execution_mode="hip_runtime"`) — see the +§5 lane-id table. These are **`compiled`** lanes: `tessera-opt` generates the +kernel + serializes to `hsaco` in-process, then HIP loads and launches it, with an +execute-vs-reference fixture per family. They are correctness-first (no perf +ladder) and a rung below the two `hardware_verified` rows below. + +**The first two — `matmul`/`gemm` + `flash_attn` are `hardware_verified`.** These +were the first ops to run natively on a non-Apple backend and carry shipped C-ABI +symbols (not just an in-process compiled lane): + +- `matmul`/`gemm`: shipped `libtessera_rocm_gemm.so` exporting + `tessera_rocm_wmma_gemm_{f16,bf16}` (HIPRTC-compiled for the device arch at + load); wired into `runtime.launch()` as the `("rocm", "rocm_wmma")` oracle row; + execute-compare fixture `tests/unit/test_rocm_wmma_runtime_symbol.py` + (f16/bf16, ragged + K-looped) vs numpy; **measured GEMM perf ladder** (register + blocking / LDS staging / software pipelining / APU zero-copy) — + `STRIX_HALO_EXECUTION_PLAN.md` Stage F. The default `matmul` path is now the + compiler-generated `rocm_compiled` lane, which degrades to this `rocm_wmma` + symbol as its oracle. +- `flash_attn`: shipped `libtessera_rocm_flash_attn.so` exporting + `tessera_rocm_wmma_flash_attn_{f16,bf16}` — FA-2 **forward**, both QK^T and P@V + on 16×16×16 WMMA, online softmax, LDS-staged scores/accumulator, causal + ragged + Sq/Sk; execute-compare fixture `tests/unit/test_rocm_flash_attn_runtime_symbol.py` + (maxerr ~1e-4 f16). The FA-2 **backward** (dQ/dK/dV) also executes on gfx1151 + via the compiler-generated `rocm_flash_attn_bwd_compiled` lane + (`generate-wmma-flash-attn-bwd-kernel` → `fa_pre`/`fa_dkdv`/`fa_dq`; O + recomputed via the forward lane, nothing saved from forward), **MHA + GQA/MQA** + (scale + causal), validated vs autodiff `vjp_flash_attn`. No perf ladder. + +Honest scope (Decision #25): everything above is **one arch (RDNA 3.5 `gfx1151`) × +{fp16, bf16}**, correctness-first. None of it flips the per-primitive +`backend_kernel` axis (that needs *all* targets `hardware_verified`). **CDNA MFMA +entries remain hardware-free** pending MI300-class silicon; Sprint H-4 lit fixtures +validate their IR + MFMA patterns, and `hipcc`/`llc` compile-only validation +promotes them to `compileable`. See §9 for the concrete done / open / blocked split. --- @@ -411,14 +450,26 @@ them to `compileable`. ## 9. Roadmap — what's done / hardware-free / blocked ### Done on real silicon (gfx1151 / Strix Halo APU) -- ✅ WMMA `matmul`/`gemm` executes + matches numpy (`{fp16, bf16}`, f32 accum) -- ✅ `hardware_verified` `backend_manifest` row + runtime `launch()` lane -- ✅ Measured GEMM perf ladder (register blocking is the winning lever on this - unified-memory APU; LDS staging / software pipelining / zero-copy give at-most - narrow wins — `STRIX_HALO_EXECUTION_PLAN.md` Stage F) -- ✅ WMMA `flash_attn` FA-2 forward executes + matches a numpy attention - reference (`{fp16, bf16}`, f32 accum; online softmax, causal, ragged; both - QK^T and P@V on WMMA). Forward only, correctness-first (no perf ladder yet). +- ✅ WMMA `matmul`/`gemm` executes + matches numpy (`{fp16, bf16}`, f32 accum); + `hardware_verified` row + shipped C-ABI symbol + measured GEMM perf ladder + (register blocking is the winning lever on this unified-memory APU — + `STRIX_HALO_EXECUTION_PLAN.md` Stage F) +- ✅ WMMA `flash_attn` FA-2 **forward** executes + matches a numpy attention + reference (`{fp16, bf16}`, online softmax, causal, ragged); `hardware_verified`, + no perf ladder +- ✅ WMMA `flash_attn` FA-2 **backward** (dQ/dK/dV) executes via the + `rocm_flash_attn_bwd_compiled` lane (`generate-wmma-flash-attn-bwd-kernel` → + `fa_pre`/`fa_dkdv`/`fa_dq`, launched through `runtime.launch()`), matches + autodiff `vjp_flash_attn` on `{fp16, bf16}`; **MHA + GQA/MQA** (grouped `fa_dkdv` + atomic-accumulates dK/dV across the group) **+ additive attn_bias** + (`S = scale*Q@K^T + bias` in the recompute), scale + causal. No perf ladder; + windowed/softcap backward is the remaining follow-up +- ✅ **Dozens of additional compiler-generated HIP `compiled` lanes execute** and + match a CPU/numpy reference — nearly all of §5 plus §10: the GEMM/attention + families, norm / activation / RoPE / ALiBi, optimizers, RNG, FFT/spectral, + selective-SSM (Mamba2, **incl. device backward**), quantization, the loss + families, and the reduction / scan / sort / scatter / elementwise lanes. The + drift-gated live list is `docs/audit/generated/runtime_execution_matrix.md`. ### Hardware-free (lit-validated, no GPU needed) - ✅ Capability matrix incl. gfx1151 + provisional gfx1250/1251 (`rocm_target.py`) @@ -433,24 +484,47 @@ them to `compileable`. - 🔜 Register `tessera-to-linalg` so the MLIR `--tessera-emit-rocdl` route works (the emitter currently rides the direct LLVM-IR path) -### Still blocked on hardware -- The rest of §5 on RDNA (gfx1151 proves matmul + flash_attn-forward so far) -- flash_attn backward; a flash_attn perf ladder; its `runtime.launch()` lane -- CDNA execution on MI300A / MI300X / MI325X (all MFMA entries) -- MFU sign-off against the §5 targets (gfx1151 has a perf ladder, not MFU proof) -- RCCL all-reduce numerical verification across 8× ranks -- Profiler timeline capture (rocprof) + multi-rank scaling +### Open on this box (gfx1151 — workable now, no CDNA needed) +- `flash_attn` **backward** for the sliding-window / logit-softcap variants + (MHA + GQA/MQA + additive **attn_bias** backward is done and runtime-wired; + window/softcap still need the masked/capped backward math in the C++ kernel) +- **Fused paged-attention** — the §5.6 movement core (`kv_cache_append/read/prune`) + now executes via `rocm_kv_cache_compiled` (scatter/gather compose, + execute-compare vs `KVCacheHandle`); a single fused gather→attention paged + kernel is the remaining step +- `grad_clip_norm` (§5.5) — global-norm + scale; single-node +- `rocm_moe_transport_compiled` — the one lane still `reference_cpu`; needs native + HIP gather/scatter transport kernels for `moe_dispatch`/`moe_combine` +- Plain recurrent cells `lstm_cell` / `gru_cell` / `simple_rnn_cell` (§5.9) — no + dedicated lane (selective-SSM + deltanet are done) +- **Perf ladders / MFU sign-off** beyond `matmul` — every `compiled` lane is + correctness-first with no perf characterization + +### Still blocked on hardware NOT on this box +- **CDNA execution** on MI300A / MI300X / MI325X (all MFMA entries) +- **FP8 / FP6 / FP4** matrix paths (gfx942/gfx950 MFMA; RDNA 3.5 has no FP8 WMMA) +- MFU sign-off against CDNA §5 targets +- RCCL all-reduce numerical verification across 8× ranks; multi-rank scaling +- Profiler timeline capture (rocprof) at datacenter scale --- ## 10. M7 Visual Complex Analysis (E3 follow-up) Parallel to the NVIDIA M7 plan (see -`docs/nvidia_cuda13_kernel_inventory.md` §9). The backend manifest -reserves `status="planned"` slots for every M7 op on the ROCm target -with the matrix below; all 20 M7 ops run today **only** via the -Python reference path on CPU (`status="reference"`, -`dtypes=("fp32",)`). +`docs/nvidia_cuda13_kernel_inventory.md` §9). + +> **Update (2026-07-10).** This family is **no longer reference-only on ROCm.** +> The M7 pointwise complex ops execute on `gfx1151` via `rocm_complex_compiled` +> (9 pointwise ops over interleaved-f32) and `rocm_conformal_compiled` +> (Möbius / stereographic on the complex mul/div lanes); the Clifford / GA +> products run via `rocm_clifford_compiled`, and the EBM family via +> `rocm_ebm_compute_compiled` / `rocm_ebm_langevin_compiled`. See the generated +> `runtime_execution_matrix`. The geometric/certificate ops (`cross_ratio`, +> `dz`/`dbar`, `laplacian_2d`, `conformal_*`, `is_concyclic`, +> `check_cauchy_riemann`) still ride the host-structured reference path. The +> planned dtype matrix below remains the fp16/bf16 *target* contract for a future +> fused MFMA kernel; the runnable dtype today is f32. | Op family | Lowering target | Planned dtype matrix | |---|---|---| diff --git a/python/tessera/compiler/backend_manifest.py b/python/tessera/compiler/backend_manifest.py index 0f90b5945..8e3164fa9 100644 --- a/python/tessera/compiler/backend_manifest.py +++ b/python/tessera/compiler/backend_manifest.py @@ -680,7 +680,7 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: "execute_compare_fixture": _APPLE_GPU_STRUCTURED_COMPUTE_FIXTURE, } for op in ( "arange", "bidirectional_scan", "center_crop", "cross_attention", - "edm_precondition", "factorized_pos_emb", "gru_cell", "image_resize", + "edm_precondition", "factorized_pos_emb", "gru_cell", "lstm_cell", "image_resize", "interpolate", "lora_linear", "masked_fill", "masked_scatter", "memory_read", "mor_partition", "mor_router", "mor_scatter", "mrope_2d", "online_softmax_state", "pack", "patchify", "perceiver_resampler", @@ -1220,7 +1220,12 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: "(tessera_rocm_wmma_flash_attn_{f16,bf16}, HIPRTC-compiled for the " "device arch at load); ROCm 7.2.4. The second op after matmul to run " "natively on a non-Apple backend. Numerically validated vs a numpy " - "attention reference by the execute_compare_fixture." + "attention reference by the execute_compare_fixture. The FA-2 " + "BACKWARD (dQ/dK/dV) also executes on gfx1151 via the " + "compiler-generated rocm_flash_attn_bwd_compiled lane " + "(generate-wmma-flash-attn-bwd-kernel -> fa_pre/fa_dkdv/fa_dq; MHA + " + "GQA/MQA + additive attn_bias, scale+causal), validated vs autodiff " + "vjp_flash_attn — see the runtime_execution_matrix." ), }, } @@ -1324,6 +1329,24 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: "tree-reduce, f32 reduce). Executes via runtime.launch() " "(rocm_softmax_compiled).", }, + # KV-cache paged-movement core (§5.6). The append/read/prune tensor movement + # over a resident cache buffer executes on gfx1151 by COMPOSING the existing + # device scatter (append row write) + masked-gather (read/prune) kernels with + # host page-index math — no bespoke kernel. Mirrors KVCacheHandle.{append, + # read,prune} on the non-quantized fp path; quantize_kv rides the intquant + # lane. Executes via runtime.launch() (rocm_kv_cache_compiled). f32. + **{op: { + "dtypes": ("fp32",), + "feature_flags": ("kv_cache", "paged"), + "notes": f"KV-cache {op.split('_')[-1]} paged-movement core over a " + "resident cache buffer (max_seq, H, D): append = row scatter-" + "write at [start, start+n) (generate-rocm-scatter-kernel, set " + "mode); read = row gather of [start, end); prune = trailing-" + "window gather + zero-fill (generate-rocm-gather-kernel). Host " + "owns the page-index math; quantize_kv rides the intquant lane. " + "Executes via runtime.launch() (rocm_kv_cache_compiled). f32, " + "matches the KVCacheHandle reference.", + } for op in ("kv_cache_append", "kv_cache_read", "kv_cache_prune")}, **{op: { "dtypes": ("fp32", "fp16", "bf16"), "feature_flags": ("reduction",), @@ -2100,6 +2123,10 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: for op in ("scatter", "scatter_add", "scatter_reduce")}, **{(op, "rocm"): "tests/unit/test_rocm_scatter_compiled.py" for op in ("scatter", "scatter_add", "scatter_reduce")}, + # §5.6 KV-cache paged-movement core — append/read/prune executed on the + # gfx1151 scatter+gather kernels, compared to the KVCacheHandle reference. + **{(op, "rocm"): "tests/unit/test_rocm_kv_cache_compiled.py" + for op in ("kv_cache_append", "kv_cache_read", "kv_cache_prune")}, # P11 — x86 MLA latent-KV lane (compress/expand/decode composed on the GEMM + # flash_attn lanes), compared to the numpy MLA reference. Skip-clean w/o .so. **{(op, "x86"): "tests/unit/test_x86_mla_compiled.py" @@ -4612,7 +4639,7 @@ def complex_manifest_for(op_name: str) -> list[BackendKernelEntry]: # vision / pooling "center_crop", "image_resize", "interpolate", # recurrent / model / layout - "bidirectional_scan", "conv1d", "conv_transpose", "gru_cell", + "bidirectional_scan", "conv1d", "conv_transpose", "gru_cell", "lstm_cell", "lora_linear", "patchify", "pixel_shuffle", "pixel_unshuffle", "simple_rnn_cell", "arange", "cast", "masked_fill", "mor_partition", "mor_router", @@ -4644,7 +4671,7 @@ def complex_manifest_for(op_name: str) -> list[BackendKernelEntry]: "center_crop", "image_resize", "interpolate", "patchify", "pixel_shuffle", "pixel_unshuffle", "conv1d", "conv_transpose", "lora_linear", - "gru_cell", "simple_rnn_cell", + "gru_cell", "lstm_cell", "simple_rnn_cell", "depthwise_conv1d", "cross_attention", "perceiver_resampler", "bidirectional_scan", diff --git a/python/tessera/compiler/execution_matrix.py b/python/tessera/compiler/execution_matrix.py index 22050dac1..d720359b1 100644 --- a/python/tessera/compiler/execution_matrix.py +++ b/python/tessera/compiler/execution_matrix.py @@ -193,6 +193,15 @@ class ExecutionRow: "HIP loads + launches it. f16/bf16 storage, f32 " "softmax + accumulate; the attention analog of " "rocm_compiled", + "rocm_flash_attn_bwd_compiled": "AMD GPU RDNA WMMA FA-2 BACKWARD the Tessera " + "compiler GENERATES (generate-wmma-flash-attn-bwd-" + "kernel -> three fa_pre/fa_dkdv/fa_dq WMMA kernels -> " + "hsaco), launched in sequence to produce dQ/dK/dV; O " + "is recomputed via the forward lane (nothing saved " + "from forward). MHA + GQA/MQA (grouped dkdv atomic-" + "accumulates dK/dV) + additive attn_bias; f16/bf16 " + "storage, f32 accumulate; the reverse-mode analog of " + "rocm_flash_attn_compiled", "rocm_linear_attn_compiled": "AMD GPU RDNA WMMA linear-attention forward the " "Tessera compiler GENERATES " "(generate-wmma-linear-attn-kernel -> ROCDL -> hsaco, " @@ -339,6 +348,11 @@ class ExecutionRow: "rocm_scatter_compiled": "AMD GPU RDNA scatter lane — scatter/scatter_add/" "scatter_reduce via the COMPILER-GENERATED gfx1151 " "kernel (one thread per element; atomic_rmw). f32", + "rocm_kv_cache_compiled": "AMD GPU RDNA KV-cache paged-movement lane — " + "kv_cache append/read/prune over a resident cache " + "buffer by composing the gfx1151 scatter (write) + " + "gather (read/prune) kernels; host page-index math. " + "f32, matches the KVCacheHandle reference", "x86_rng_compiled": "x86 CPU device RNG — counter-based Philox-4x32-10 " "uniform kernel + host transform (uniform/normal/" "dropout). f32", @@ -1751,6 +1765,26 @@ class ExecutionRow: "hsaco in-process, then HIP loads + launches it. The attention " "analog of the compiled GEMM lane (rocm_compiled).", execution_mode="hip_runtime"), + # The compiler-GENERATED FA-2 backward (generate-wmma-flash-attn-bwd-kernel + # -> three fa_pre/fa_dkdv/fa_dq WMMA kernels -> hsaco) launched in sequence. + # Self-contained VJP over (dO, Q, K, V): O is recomputed via the forward + # lane, so nothing is saved from forward. Validated vs the numpy attention + # backward / autodiff vjp_flash_attn. Core MHA (scale + causal). + ("rocm", "rocm_flash_attn_bwd_compiled"): ExecutionRow( + target="rocm", compiler_path="rocm_flash_attn_bwd_compiled", + execution_kind="native_gpu", executable=True, + executor_id="rocm_flash_attn_bwd_compiled", runtime_status="success", + reason="ROCm flash_attn backward artifact runs the COMPILER-GENERATED " + "RDNA WMMA FA-2 backward: tessera-opt expands one " + "tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/" + "fa_dq WMMA kernels serialized to hsaco in-process, then HIP " + "launches them in sequence to produce dQ/dK/dV. O is recomputed " + "via the forward lane (nothing saved from forward). The " + "reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA " + "(gqa dkdv atomic-accumulates dK/dV across the group) + additive " + "attn_bias (S=scale*QK+bias in the recompute), scale + causal, " + "f16/bf16 storage, f32 accumulate.", + execution_mode="hip_runtime"), # Linear-attention family (quadratic-parallel form, no softmax; a distinct # algorithm from flash_attn): tessera.linear_attn + the decay-masked siblings # tessera.lightning_attention / tessera.retention, dispatched by op name. @@ -1931,6 +1965,17 @@ class ExecutionRow: "atomic_rmw for add/min/max). f32, matches the numpy scatter " "reference.", execution_mode="hip_runtime"), + ("rocm", "rocm_kv_cache_compiled"): ExecutionRow( + target="rocm", compiler_path="rocm_kv_cache_compiled", + execution_kind="native_gpu", executable=True, + executor_id="rocm_kv_cache_compiled", runtime_status="success", + reason="ROCm KV-cache paged-movement lane realizes kv_cache " + "append/read/prune over a resident cache buffer (max_seq, H, D) " + "by COMPOSING the COMPILER-GENERATED gfx1151 scatter (append row " + "write) + masked-gather (read/prune) kernels with host page-index " + "math. quantize_kv rides the intquant lane. f32, matches the " + "KVCacheHandle append/read/prune reference.", + execution_mode="hip_runtime"), ("rocm", "rocm_conformal_compiled"): ExecutionRow( target="rocm", compiler_path="rocm_conformal_compiled", execution_kind="native_gpu", executable=True, diff --git a/python/tessera/nn/functional.py b/python/tessera/nn/functional.py index a0b99b9fb..0ae0d7cb7 100644 --- a/python/tessera/nn/functional.py +++ b/python/tessera/nn/functional.py @@ -309,6 +309,28 @@ def gru_cell(x, h, W_ih, W_hh, b_ih=None, b_hh=None): return (1.0 - z) * n + z * h_arr +def lstm_cell(x, h, c, W_ih, W_hh, b_ih=None, b_hh=None): + """LSTM cell with gate order i, f, g, o (PyTorch/cuDNN convention). + + ``x`` (…, in), ``h``/``c`` (…, hidden); ``W_ih`` (in, 4*hidden), ``W_hh`` + (hidden, 4*hidden). Returns ``(h_new, c_new)``. + """ + x_arr, h_arr, c_arr = _asarray(x), _asarray(h), _asarray(c) + gates = x_arr @ _asarray(W_ih) + h_arr @ _asarray(W_hh) + if b_ih is not None: + gates = gates + _asarray(b_ih) + if b_hh is not None: + gates = gates + _asarray(b_hh) + i, f, g, o = np.split(gates, 4, axis=-1) + i = 1.0 / (1.0 + np.exp(-i)) + f = 1.0 / (1.0 + np.exp(-f)) + g = np.tanh(g) + o = 1.0 / (1.0 + np.exp(-o)) + c_new = f * c_arr + i * g + h_new = o * np.tanh(c_new) + return h_new, c_new + + def bidirectional_scan(fn, init_fwd, init_bwd, xs): xs_arr = _asarray(xs) fwd_states = [] @@ -726,6 +748,7 @@ def proj_heads(t, W, heads): "group_norm", "gqa_attention", "gru_cell", + "lstm_cell", "instance_norm", "mask_token_block", "max_pool", diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index 6ebfef5f3..ca251d0c7 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -3076,6 +3076,325 @@ def _mr(p, size): return o.reshape(q.shape) +# ───────────────────────────────────────────────────────────────────────────── +# ROCm COMPILED flash_attn BACKWARD lane (2026-07-10) — the reverse-mode analog +# of the forward lane above. ``runtime.launch()`` of an artifact stamped +# ``compiler_path = "rocm_flash_attn_bwd_compiled"`` builds the COMPILER- +# GENERATED FA-2 backward via ``generate-wmma-flash-attn-bwd-kernel`` (one +# tessera_rocm.flash_attn_bwd directive → THREE fragment-materialized WMMA +# kernels fa_pre / fa_dkdv / fa_dq → hsaco) and launches them in sequence to +# produce (dQ, dK, dV). The forward output O the pre-pass needs is recomputed +# on-device by REUSING the forward lane above (no state saved from forward), so +# the backward is a self-contained VJP over (dO, Q, K, V). f16/bf16 storage, f32 +# accumulate. Matches the numpy attention-backward reference / autodiff +# ``vjp_flash_attn``. Core MHA (scale + causal) — GQA / sliding-window / +# logit-softcap / attn_bias backward are a follow-up (the C++ kernel already +# carries the GQA variant; the runtime lane forwards only the core here). +# ───────────────────────────────────────────────────────────────────────────── +#: hsaco bytes keyed by (head_dim, chip, dtype, gqa, bias). +_rocm_fa_bwd_hsaco_cache: dict[tuple[int, str, str, bool, bool], bytes] = {} + + +def _build_compiled_flash_attn_bwd_hsaco(head_dim: int, + dtype: str = "f16", + gqa: bool = False, + bias: bool = False) -> bytes: + """Generate + serialize the compiler's WMMA FA-2 backward kernels to hsaco + (fa_pre / fa_dkdv / fa_dq in one module), fully in-process via tessera-opt. + Cached per (head_dim, chip, dtype, gqa, bias). gqa=True emits the grouped- + query variant: all three kernels gain (heads, kv_ratio) runtime args and + fa_dkdv atomically accumulates dK/dV across the kv_ratio query heads sharing + each KV head (host pre-zeros dK/dV). bias=True emits the additive-bias + variant: the recompute forms S = scale*Q@K^T + bias (a trailing f32 + [bh,Sq,Sk] memref arg, LAST) before the softmax.""" + chip = _rocm_chip() + key = (head_dim, chip, dtype, gqa, bias) + cached = _rocm_fa_bwd_hsaco_cache.get(key) + if cached is not None: + return cached + opt = _tessera_opt_path() + if opt is None: + raise _RocmCompiledUnavailable( + "tessera-opt not built — no compiled ROCm flash_attn backward lane") + gqa_attr = ", gqa = true" if gqa else "" + bias_attr = ", attn_bias = true" if bias else "" + directive = ( + 'module {\n' + ' "tessera_rocm.flash_attn_bwd"() {name = "fa", ' + f'head_dim = {head_dim} : i64, dtype = "{dtype}"{gqa_attr}{bias_attr}}} ' + ': () -> ()\n' + '}\n' + ) + pipeline = ( + "builtin.module(" + "generate-wmma-flash-attn-bwd-kernel," + "lower-tessera-target-to-rocdl," + "gpu.module(convert-scf-to-cf,convert-gpu-to-rocdl," + "reconcile-unrealized-casts)," + f"rocdl-attach-target{{chip={chip}}}," + "gpu-module-to-binary)" + ) + import subprocess + r = subprocess.run([str(opt), "-", f"--pass-pipeline={pipeline}"], + input=directive, capture_output=True, text=True) + if r.returncode != 0 or "gpu.binary" not in r.stdout: + raise _RocmCompiledUnavailable( + "tessera-opt did not serialize the compiled flash_attn backward " + f"in-process (rc={r.returncode}): {r.stderr[:400]}") + hsaco = _extract_hsaco_blob(r.stdout) + if hsaco[:4] != b"\x7fELF": + raise _RocmCompiledUnavailable( + "compiled ROCm flash_attn backward lane: gpu.binary was not an ELF " + "hsaco") + _rocm_fa_bwd_hsaco_cache[key] = hsaco + return hsaco + + +def _rocm_flash_attn_forward_o(q: Any, k: Any, v: Any, scale: float, + causal: int, bias: Any = None) -> Any: + """Recompute the forward output O (f32, shape [bh, sq, D]) by REUSING the + forward compiled lane — the backward saves nothing from forward. When + ``bias`` is given, O = softmax(scale*Q@K^T + bias)*V (the forward lane takes + attn_bias as a 4th operand), so O matches the biased backward recompute.""" + names = ["q", "k", "v"] + (["bias"] if bias is not None else []) + fwd_art = RuntimeArtifact(metadata={ + "target": "rocm", "compiler_path": "rocm_flash_attn_compiled", + "executable": True, "execution_kind": "native_gpu", + "arg_names": names, "output_name": "o", + "ops": [{"op_name": "tessera.flash_attn", "result": "o", + "operands": names, + "kwargs": {"scale": scale, "causal": bool(causal)}}]}) + call = (q, k, v) if bias is None else (q, k, v, bias) + return _execute_rocm_compiled_flash_attn(fwd_art, call) + + +_FA_BWD_OPS = ("tessera.flash_attn_bwd", "tessera.flash_attn_vjp") + + +def _execute_rocm_compiled_flash_attn_bwd(artifact: RuntimeArtifact, + args: Any) -> Any: + """The ``target="rocm"`` flash_attn BACKWARD lane. Operands are + ``(dO, Q, K, V)`` — ``Q``/``dO`` are ``[..., H, Sq, D]`` and ``K``/``V`` are + ``[..., G, Sk, D]`` (``G <= H``; ``G == H`` is plain MHA, ``G < H`` selects + the grouped-query backward). Returns the tuple ``(dQ, dK, dV)`` shaped like + ``(Q, K, V)``. Raises ``_RocmCompiledUnavailable`` when the lane can't run + (no tessera-opt / serialization / GPU); ``ValueError`` on a bad op / dtype / + shape / unsupported variant; ``RuntimeError`` on a real kernel failure.""" + import numpy as np + + metadata = artifact.metadata or {} + arg_names = list(metadata.get("arg_names") or []) + ops = list(metadata.get("ops") or []) + op_name = str(ops[0].get("op_name", "")) if len(ops) == 1 else "" + if len(ops) != 1 or op_name not in _FA_BWD_OPS: + raise ValueError( + "rocm_flash_attn_bwd_compiled executor handles exactly one of " + f"{_FA_BWD_OPS}; got {[o.get('op_name') for o in ops]!r}") + op = ops[0] + operand_names = [str(n) for n in op.get("operands", [])] + if len(operand_names) not in (4, 5): + raise ValueError( + "flash_attn backward requires (dO, Q, K, V) operands, optionally " + "plus a 5th additive attn_bias") + values = _bind_launch_args(args, arg_names) + do = _as_numpy(values[operand_names[0]]) + q = _as_numpy(values[operand_names[1]]) + k = _as_numpy(values[operand_names[2]]) + v = _as_numpy(values[operand_names[3]]) + bias_arr = (_as_numpy(values[operand_names[4]]) + if len(operand_names) == 5 else None) + for name, arr in (("dO", do), ("Q", q), ("K", k), ("V", v)): + if arr.ndim < 2: + raise ValueError( + f"flash_attn backward {name} must be rank >= 2 ([..., S, D]); " + f"got {arr.shape}") + head_dim = int(q.shape[-1]) + if head_dim <= 0 or head_dim % 16 != 0: + raise ValueError( + f"rocm flash_attn backward needs head_dim a positive multiple of " + f"16; got {head_dim}") + sq, sk = int(q.shape[-2]), int(k.shape[-2]) + bh = int(np.prod(q.shape[:-2])) if q.ndim > 2 else 1 + bh_kv = int(np.prod(k.shape[:-2])) if k.ndim > 2 else 1 + # GQA/MQA: H query heads vs G<=H key/value heads (head axis -3). The grouped + # backward variant atomically accumulates dK/dV across the kv_ratio query + # heads sharing each KV head. V must match K, dO must match Q, the batch dims + # (everything before the head axis) must agree, and head_dim is shared. + n_qh = int(q.shape[-3]) if q.ndim >= 3 else 1 + n_kvh = int(k.shape[-3]) if k.ndim >= 3 else 1 + gqa = n_kvh != n_qh + if gqa and (n_kvh <= 0 or n_qh % n_kvh != 0): + raise ValueError( + f"rocm flash_attn backward GQA needs query-heads ({n_qh}) divisible " + f"by kv-heads ({n_kvh})") + kv_ratio = n_qh // n_kvh if gqa else 1 + if (v.shape != k.shape or do.shape != q.shape + or int(v.shape[-1]) != head_dim + or q.shape[:-3] != k.shape[:-3]): + raise ValueError( + "rocm flash_attn backward requires V==K, dO==Q, matching batch dims " + f"and a shared head_dim; got Q{q.shape} K{k.shape} V{v.shape} " + f"dO{do.shape}.") + + if q.dtype == np.float16: + dtype_tag, store = "f16", np.float16 + else: + bf16 = _bfloat16_dtype() + if bf16 is not None and q.dtype == bf16: + dtype_tag, store = "bf16", bf16 + else: + raise ValueError( + "rocm flash_attn backward handles f16/bf16 storage (f32 " + f"accumulate); got {q.dtype}") + + kwargs = op.get("kwargs") or {} + causal = 1 if bool(kwargs.get("causal", False)) else 0 + scale = kwargs.get("scale") + scale = float(scale) if scale is not None else 1.0 / float(np.sqrt(head_dim)) + for bad, why in (("window", "sliding-window"), ("logit_softcap", "softcap")): + if kwargs.get(bad): + raise ValueError( + f"rocm flash_attn backward does not yet support {why}; the core " + "MHA (scale + causal) + additive attn_bias backward is wired.") + + # Additive bias: S = softmax(scale*Q@K^T + bias). Host-broadcast to + # Q.lead+(Sq,Sk) → [bh,Sq,Sk] f32 (kernel indexes bias[(bh*Sq+q)*Sk+k]), + # matching the forward attn_bias lane. bias enters both O (via the forward + # recompute) and the backward P/L/dS — dbias is not produced (constant-bias + # VJP: dbias is only needed when bias is a positional differentiable input). + has_bias = bias_arr is not None + bias_c = None + if has_bias: + try: + bias_b = np.broadcast_to( + np.asarray(bias_arr, dtype=np.float32), + tuple(q.shape[:-2]) + (sq, sk)) + except ValueError as exc: + raise ValueError( + f"rocm flash_attn backward attn_bias {np.asarray(bias_arr).shape}" + f" is not broadcastable to Q.lead+(Sq,Sk) = " + f"{tuple(q.shape[:-2]) + (sq, sk)}") from exc + bias_c = np.ascontiguousarray(bias_b, dtype=np.float32).reshape(-1) + + # O the pre-pass needs — recompute on-device via the forward lane (with the + # same bias, so O = softmax(scale*QK + bias)*V matches the recompute). + o_f32 = np.ascontiguousarray( + _rocm_flash_attn_forward_o(q, k, v, scale, causal, + bias=bias_arr if has_bias else None), + dtype=np.float32).reshape(bh, sq, head_dim) + + hsaco = _build_compiled_flash_attn_bwd_hsaco(head_dim, dtype_tag, gqa=gqa, + bias=has_bias) + hip = _load_hip_for_launch() + if hip is None: + raise _RocmCompiledUnavailable( + "libamdhip64.so not loadable — no ROCm execution lane on this host") + if hip.hipInit(0) != 0: + raise _RocmCompiledUnavailable("rocm flash_attn bwd: hipInit failed") + mod = ctypes.c_void_p() + if hip.hipModuleLoadData(ctypes.byref(mod), hsaco) != 0: + raise _RocmCompiledUnavailable( + "rocm flash_attn bwd: no usable AMD GPU (module load failed)") + fns = {} + for nm in (b"fa_pre", b"fa_dkdv", b"fa_dq"): + fn = ctypes.c_void_p() + if hip.hipModuleGetFunction(ctypes.byref(fn), mod, nm) != 0: + raise RuntimeError( + f"rocm flash_attn bwd: kernel symbol {nm!r} not found") + fns[nm] = fn + + qc = np.ascontiguousarray(q, dtype=store) + kc = np.ascontiguousarray(k, dtype=store) + vc = np.ascontiguousarray(v, dtype=store) + doc = np.ascontiguousarray(do, dtype=store) + nq, nkv, nl = bh * sq * head_dim, bh_kv * sk * head_dim, bh * sq + cv = ctypes.c_void_p + bufs: dict[str, Any] = {} + for name, nb in (("Q", 2 * nq), ("K", 2 * nkv), ("V", 2 * nkv), + ("dO", 2 * nq), ("O", 4 * nq), ("L", 4 * nl), ("Dd", 4 * nl), + ("dQ", 4 * nq), ("dK", 4 * nkv), ("dV", 4 * nkv)): + d = cv() + if hip.hipMalloc(ctypes.byref(d), max(nb, 4)) != 0: + for dd in bufs.values(): + hip.hipFree(dd) + raise RuntimeError("rocm flash_attn bwd: hipMalloc failed") + bufs[name] = d + hip.hipMemcpy(bufs["Q"], qc.ctypes.data_as(cv), 2 * nq, 1) + hip.hipMemcpy(bufs["K"], kc.ctypes.data_as(cv), 2 * nkv, 1) + hip.hipMemcpy(bufs["V"], vc.ctypes.data_as(cv), 2 * nkv, 1) + hip.hipMemcpy(bufs["dO"], doc.ctypes.data_as(cv), 2 * nq, 1) + hip.hipMemcpy(bufs["O"], o_f32.ctypes.data_as(cv), 4 * nq, 1) + n_bias = bh * sq * sk + if has_bias: + d = cv() + if hip.hipMalloc(ctypes.byref(d), max(4 * n_bias, 4)) != 0: + for dd in bufs.values(): + hip.hipFree(dd) + raise RuntimeError("rocm flash_attn bwd: hipMalloc (attn_bias) failed") + bufs["bias"] = d + hip.hipMemcpy(d, bias_c.ctypes.data_as(cv), 4 * n_bias, 1) + if gqa: + # kv_ratio query heads atomic-accumulate into each KV head's dK/dV rows + # (the sharing blocks add into the same rows) — pre-zero before fa_dkdv. + hip.hipMemset(bufs["dK"], 0, 4 * nkv) + hip.hipMemset(bufs["dV"], 0, 4 * nkv) + + def _mr(p, size): + return [cv(p.value), cv(p.value), ctypes.c_int64(0), + ctypes.c_int64(size), ctypes.c_int64(1)] + + sqc, skc = ctypes.c_int64(sq), ctypes.c_int64(sk) + sc, cau = ctypes.c_float(scale), ctypes.c_int64(causal) + # Grouped kernels take (heads, kv_ratio) as two trailing runtime args after + # (Sq, Sk, scale, causal); the plain kernels have neither. Grid.y is always + # the query-head count B*H (fa_dkdv accumulates into KV heads atomically). + tail = [ctypes.c_int64(n_qh), ctypes.c_int64(kv_ratio)] if gqa else [] + # Additive bias memref is the LAST kernel arg on all three kernels (after the + # optional gqa pair) — matches the C++ withGqa(... + bias) arg order. + bias_tail = _mr(bufs["bias"], n_bias) if has_bias else [] + gqt, gkt, gyt = (sq + 15) // 16, (sk + 15) // 16, bh + + def _launch(fn, gx, argv): + arr = (cv * len(argv))() + for i, val in enumerate(argv): + arr[i] = ctypes.cast(ctypes.byref(val), cv) + rc = hip.hipModuleLaunchKernel(fn, gx, gyt, 1, 32, 1, 1, 0, None, + arr, None) + if rc != 0: + for dd in bufs.values(): + hip.hipFree(dd) + raise RuntimeError( + f"rocm flash_attn bwd: kernel launch failed rc={rc}") + + # fa_pre : (Q,K,dO, O,L,Dd, Sq,Sk, scale, causal[, heads, kv_ratio][, bias]) + _launch(fns[b"fa_pre"], gqt, + _mr(bufs["Q"], nq) + _mr(bufs["K"], nkv) + _mr(bufs["dO"], nq) + + _mr(bufs["O"], nq) + _mr(bufs["L"], nl) + _mr(bufs["Dd"], nl) + + [sqc, skc, sc, cau] + tail + bias_tail) + # fa_dkdv : (Q,K,V,dO, L,Dd, dK,dV, Sq,Sk, scale, causal[,heads,kv_ratio][,bias]) + _launch(fns[b"fa_dkdv"], gkt, + _mr(bufs["Q"], nq) + _mr(bufs["K"], nkv) + _mr(bufs["V"], nkv) + + _mr(bufs["dO"], nq) + _mr(bufs["L"], nl) + _mr(bufs["Dd"], nl) + + _mr(bufs["dK"], nkv) + _mr(bufs["dV"], nkv) + + [sqc, skc, sc, cau] + tail + bias_tail) + # fa_dq : (Q,K,V,dO, L,Dd, dQ, Sq,Sk, scale, causal[, heads, kv_ratio][, bias]) + _launch(fns[b"fa_dq"], gqt, + _mr(bufs["Q"], nq) + _mr(bufs["K"], nkv) + _mr(bufs["V"], nkv) + + _mr(bufs["dO"], nq) + _mr(bufs["L"], nl) + _mr(bufs["Dd"], nl) + + _mr(bufs["dQ"], nq) + [sqc, skc, sc, cau] + tail + bias_tail) + hip.hipDeviceSynchronize() + + dq = np.zeros((bh, sq, head_dim), np.float32) + dk = np.zeros((bh_kv, sk, head_dim), np.float32) + dv = np.zeros((bh_kv, sk, head_dim), np.float32) + hip.hipMemcpy(dq.ctypes.data_as(cv), bufs["dQ"], 4 * nq, 2) + hip.hipMemcpy(dk.ctypes.data_as(cv), bufs["dK"], 4 * nkv, 2) + hip.hipMemcpy(dv.ctypes.data_as(cv), bufs["dV"], 4 * nkv, 2) + for dd in bufs.values(): + hip.hipFree(dd) + return (dq.reshape(q.shape), dk.reshape(k.shape), dv.reshape(v.shape)) + + # ───────────────────────────────────────────────────────────────────────────── # x86 flash_attn lane (P10) — the AVX-512 PARTNER to the ROCm WMMA flash_attn. # FA-style streaming/online softmax (running max/denominator + rescaled @@ -9522,6 +9841,113 @@ def _execute_rocm_compiled_scatter(artifact: RuntimeArtifact, args: Any) -> Any: return _execute_scatter(artifact, args, _rocm_scatter, "rocm_scatter_compiled") +# ───────────────────────────────────────────────────────────────────────────── +# KV-cache paged-movement lane (§5.6) — the append/read/prune core executes on +# gfx1151 by COMPOSING the existing device gather/scatter kernels + host page- +# index math (the clamp/normcompose "compose on existing kernels" pattern). The +# three kv_cache_* ops are stateful over a KVCacheHandle; this lane executes +# their tensor MOVEMENT core over a resident cache buffer `(max_seq, H, D)`: +# * append → row scatter-write of the new tokens at [start, start+n) +# (generate-rocm-scatter-kernel, set mode); +# * read → row gather of [start, end) (generate-rocm-gather-kernel); +# * prune → keep the trailing `limit` tokens, shifted to the front (a gather +# into a zeroed leading window), zeroing the vacated slots. +# Semantics mirror KVCacheHandle.{append,read,prune} on the non-quantized fp +# path (quantize_kv rides the intquant lane). K and V are independent buffers — +# the caller drives the op once per buffer. Executes via runtime.launch() +# (rocm_kv_cache_compiled); f32, matches the KVCacheHandle reference. +# compiler_path="rocm_kv_cache_compiled". +# ───────────────────────────────────────────────────────────────────────────── +_KV_CACHE_OPS = ("tessera.kv_cache.append", "tessera.kv_cache.read", + "tessera.kv_cache.prune") + + +def _rocm_gather_rows(buf2d: Any, row_ids: Any, np: Any) -> Any: + """Gather rows ``row_ids`` from ``buf2d`` (rows, row_len) on the gfx1151 + masked-gather kernel. Returns (len(row_ids), row_len) f32.""" + row_len = int(buf2d.shape[1]) + row_ids = np.ascontiguousarray(row_ids, np.int64) + m = int(row_ids.shape[0]) + # Expand row indices to element indices: out[r, j] = src[row_ids[r]*row_len+j]. + elem_idx = (row_ids[:, None] * row_len + + np.arange(row_len, dtype=np.int64)[None, :]).reshape(-1) + out = np.zeros(m * row_len, np.float32) + _rocm_gather(np.ascontiguousarray(buf2d, np.float32).reshape(-1), elem_idx, out) + return out.reshape(m, row_len) + + +def _kv_cache_compute(op_name: str, operands: list, kwargs: dict, np: Any) -> Any: + buf = np.ascontiguousarray(operands[0], np.float32) + if buf.ndim < 2: + raise ValueError( + f"kv_cache lane: buffer must be (max_seq, ...); got {buf.shape}") + max_seq = int(buf.shape[0]) + row_len = int(buf[0].size) + tail_shape = buf.shape[1:] + flat = buf.reshape(max_seq, row_len) + + if op_name == "tessera.kv_cache.append": + new = np.ascontiguousarray(operands[1], np.float32) + n = int(new.shape[0]) + start = int(kwargs.get("start", 0)) + if start < 0 or start + n > max_seq: + raise ValueError( + f"kv_cache append: [{start}, {start + n}) out of buffer " + f"max_seq={max_seq}") + out = flat.copy() # preload base (untouched rows) + if n: + _rocm_scatter(out, new.reshape(n, row_len), + np.arange(start, start + n, dtype=np.int64), + max_seq, row_len, 0, np) # mode 0 = set + return out.reshape((max_seq,) + tail_shape) + + if op_name == "tessera.kv_cache.read": + start = int(kwargs.get("start", 0)) + end = kwargs.get("end", None) + end = start + 1 if end is None else int(end) + if not (0 <= start <= end <= max_seq): + raise ValueError( + f"kv_cache read: [{start}, {end}) out of buffer " + f"max_seq={max_seq}") + m = end - start + if m == 0: + return np.zeros((0,) + tail_shape, np.float32) + rows = _rocm_gather_rows(flat, np.arange(start, end, dtype=np.int64), np) + return rows.reshape((m,) + tail_shape) + + # tessera.kv_cache.prune — keep the trailing `limit` of the first + # `current_seq` rows, shifted to the front; zero the vacated slots. + limit = int(kwargs["limit"]) + current_seq = int(kwargs.get("current_seq", max_seq)) + if limit < 0: + raise ValueError(f"kv_cache prune: limit must be >= 0; got {limit}") + if limit >= current_seq: # nothing to drop (handle no-op) + return buf.copy() + start = current_seq - limit + out = flat.copy() + out[:limit] = _rocm_gather_rows( + flat, np.arange(start, current_seq, dtype=np.int64), np) + out[limit:current_seq] = 0.0 + return out.reshape((max_seq,) + tail_shape) + + +def _execute_rocm_compiled_kv_cache(artifact: RuntimeArtifact, args: Any) -> Any: + """The ``target="rocm"`` KV-cache paged-movement lane (append/read/prune).""" + import numpy as np + metadata = artifact.metadata or {} + arg_names = list(metadata.get("arg_names") or []) + ops = list(metadata.get("ops") or []) + op_name = str(ops[0].get("op_name", "")) if len(ops) == 1 else "" + if len(ops) != 1 or op_name not in _KV_CACHE_OPS: + raise ValueError( + f"rocm_kv_cache_compiled handles one of {_KV_CACHE_OPS}; " + f"got {[o.get('op_name') for o in ops]!r}") + operand_names = [str(n) for n in ops[0].get("operands", [])] + values = _bind_launch_args(args, arg_names) + operands = [_as_numpy(values[n]) for n in operand_names] + return _kv_cache_compute(op_name, operands, ops[0].get("kwargs") or {}, np) + + # ───────────────────────────────────────────────────────────────────────────── # Sort lane (P9) — sort / argsort / top_k via a data-independent BITONIC sort # network (the GPU-friendly choice; identical schedule on x86 + ROCm). The host @@ -12673,6 +13099,7 @@ def reduce_exec(x: Any, op_name: str, axis: Any, keepdims: bool) -> Any: "tessera.lora_linear", "tessera.simple_rnn_cell", "tessera.gru_cell", + "tessera.lstm_cell", "tessera.depthwise_conv1d", "tessera.cross_attention", "tessera.perceiver_resampler", @@ -12762,6 +13189,8 @@ def _execute_structured_compute_composite( return F.simple_rnn_cell(*operands, **kwargs) if op_name == "tessera.gru_cell": return F.gru_cell(*operands, **kwargs) + if op_name == "tessera.lstm_cell": + return F.lstm_cell(*operands, **kwargs) if op_name == "tessera.cross_attention": return ops.cross_attention(*operands, **kwargs) if op_name == "tessera.perceiver_resampler": @@ -15349,6 +15778,7 @@ def _executor_table(): "rocm_wmma": _execute_rocm_wmma_artifact, "rocm_compiled": _execute_rocm_compiled_gemm, "rocm_flash_attn_compiled": _execute_rocm_compiled_flash_attn, + "rocm_flash_attn_bwd_compiled": _execute_rocm_compiled_flash_attn_bwd, "rocm_linear_attn_compiled": _execute_rocm_compiled_linear_attn, "rocm_dspark_draft_block_compiled": _execute_rocm_dspark_draft_block_reference, "rocm_softmax_compiled": _execute_rocm_compiled_softmax, @@ -15459,6 +15889,7 @@ def _executor_table(): "rocm_rng_compiled": _execute_rocm_compiled_rng, "rocm_strided_compiled": _execute_rocm_compiled_strided, "rocm_scatter_compiled": _execute_rocm_compiled_scatter, + "rocm_kv_cache_compiled": _execute_rocm_compiled_kv_cache, "rocm_sort_compiled": _execute_rocm_compiled_sort, "rocm_clifford_compiled": _execute_rocm_compiled_clifford, "rocm_normcompose_compiled": _execute_rocm_compiled_normcompose, diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp index 3f92cd7e4..6e0514713 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp @@ -141,7 +141,7 @@ static Value groupedKvBase(Emit &e, OpBuilder &b, Location loc, Value bh, // cheap O(D) per-row elementwise rowsum (not on the hot path). //===----------------------------------------------------------------------===// void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, - Type storeTy, bool gqa = false) { + Type storeTy, bool gqa = false, bool bias = false) { MLIRContext *ctx = b.getContext(); Type f32 = b.getF32Type(); Value sQ = f.addWorkgroupAttribution(ldsT(ctx, 16 * D, storeTy), loc); @@ -172,6 +172,10 @@ void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, : e.mul(e.mul(bh, Sk), cD); Value isCausal = b.create(loc, arith::CmpIPredicate::ne, causal, c0); + // Additive bias memref (LAST arg, after the optional gqa pair): f32 [bh,Sq,Sk]. + Value biasBuf; + if (bias) + biasBuf = f.getArgument(gqa ? 12 : 10); // Stage Q into sQ (all 32 lanes cooperatively): for i = tid; i < 16*D; i+=32. { @@ -253,6 +257,15 @@ void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, for (int64_t el = 0; el < 8; ++el) { Value qi = e.add(e.ci(2 * el), half); Value v0 = e.mulf(e.ext(cs, el), scale); + // Additive bias — same S = scale*Q@K^T + bias the recompute uses, so L + // (logsumexp) matches P/dS. Bounds-guarded against the [bh,Sq,Sk] buffer. + if (bias) { + Value gqe = e.add(q0, qi); + Value gqSafe = e.sel(e.lt(gqe, Sq), gqe, c0); + Value gkSafe = e.sel(e.ge(gk, Sk), c0, gk); + Value bidx = e.add(e.mul(e.add(e.mul(bh, Sq), gqSafe), Sk), gkSafe); + v0 = e.addf(v0, e.f32load(biasBuf, bidx)); + } Value cmask = b.create(loc, isCausal, e.lt(e.add(q0, qi), gk)); Value masked = b.create(loc, e.ge(gk, Sk), cmask); @@ -320,6 +333,10 @@ struct ScoreCtx { Value Sq, Sk, scale, isCausal; Value qbase, kbase, q0, k0, bh; Value l15, half; + // Additive attention bias: f32 [bh, Sq, Sk] (broadcast, LAST kernel arg), or + // a null Value when the kernel has no bias. When set, the recompute forms + // S = scale*Q@K^T + bias before the softmax, so P/L/dS all see the bias. + Value biasBuf; }; void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, @@ -368,6 +385,15 @@ void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, Value Lq = e.f32load(x.L, Lidx); Value Dq = e.f32load(x.Dd, Lidx); Value s = e.mulf(e.ext(cs, el), x.scale); + // Additive bias: S = scale*Q@K^T + bias[(bh*Sq + q)*Sk + k]. Guarded on the + // query/key bounds so masked lanes never read past the [bh,Sq,Sk] buffer + // (their P is zeroed by `masked` below anyway). + if (x.biasBuf) { + Value gkSafe = e.sel(e.ge(gk, x.Sk), c0, gk); + Value bidx = e.add( + e.mul(e.add(e.mul(x.bh, x.Sq), gqSafe), x.Sk), gkSafe); + s = e.addf(s, e.f32load(x.biasBuf, bidx)); + } Value P = b.create(loc, e.subf(s, Lq)); // mask: query OOB, key OOB, or causal (key > query) -> P = 0 Value m1 = e.ge(gqi, x.Sq); @@ -391,7 +417,7 @@ void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, // Sq, Sk : index, scale : f32, causal : index) //===----------------------------------------------------------------------===// void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, - Type storeTy, bool gqa = false) { + Type storeTy, bool gqa = false, bool bias = false) { MLIRContext *ctx = b.getContext(); Value sP = f.addWorkgroupAttribution(ldsT(ctx, 16 * 16, storeTy), loc); Value sDS = f.addWorkgroupAttribution(ldsT(ctx, 16 * 16, storeTy), loc); @@ -423,6 +449,9 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, : e.mul(e.mul(bh, Sk), cD); Value isCausal = b.create(loc, arith::CmpIPredicate::ne, causal, c0); + Value biasBuf; // f32 [bh,Sq,Sk], LAST arg (after the optional gqa pair) + if (bias) + biasBuf = f.getArgument(gqa ? 14 : 12); // zero dK/dV accumulators. { @@ -451,6 +480,7 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, ScoreCtx x{Q, Kk, V, dO, L, Dd, Sq, Sk, scale, isCausal, qbase, kbase, q0, k0, bh, l15, half}; + x.biasBuf = biasBuf; recomputeScoreTile(e, b, loc, x, sP, sDS, /*wantP=*/true); b.create(loc); @@ -540,7 +570,7 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, // Sq, Sk : index, scale : f32, causal : index) //===----------------------------------------------------------------------===// void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, - Type storeTy, bool gqa = false) { + Type storeTy, bool gqa = false, bool bias = false) { MLIRContext *ctx = b.getContext(); Value sDS = f.addWorkgroupAttribution(ldsT(ctx, 16 * 16, storeTy), loc); Value dQacc = f.addWorkgroupAttribution(ldsT(ctx, 16 * D, b.getF32Type()), loc); @@ -568,6 +598,9 @@ void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, : e.mul(e.mul(bh, Sk), cD); Value isCausal = b.create(loc, arith::CmpIPredicate::ne, causal, c0); + Value biasBuf; // f32 [bh,Sq,Sk], LAST arg (after the optional gqa pair) + if (bias) + biasBuf = f.getArgument(gqa ? 13 : 11); { auto lp = b.create(loc, tid, c16D, c32); @@ -596,6 +629,7 @@ void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, ScoreCtx x{Q, Kk, V, dO, L, Dd, Sq, Sk, scale, isCausal, qbase, kbase, q0, k0, bh, l15, half}; + x.biasBuf = biasBuf; recomputeScoreTile(e, b, loc, x, /*sP=*/Value(), sDS, /*wantP=*/false); b.create(loc); @@ -723,6 +757,13 @@ struct GenerateWMMAFlashAttnBwdKernelPass SmallVector gqaExtra; if (gqa) gqaExtra = {idxTy, idxTy}; + // Additive attention bias: a trailing f32 [bh,Sq,Sk] memref (LAST arg, + // after the optional gqa pair). S = scale*Q@K^T + bias before the softmax, + // so P/L/dS all see it; dbias is not emitted (constant-bias VJP path). + bool bias = false; + if (auto a = op->getAttrOfType("attn_bias")) + bias = a.getValue(); + auto mk = [&](StringRef suffix, ArrayRef args, function_ref body) { auto fnTy = b.getFunctionType(args, {}); @@ -734,23 +775,24 @@ struct GenerateWMMAFlashAttnBwdKernelPass auto withGqa = [&](ArrayRef base) { SmallVector a(base.begin(), base.end()); a.append(gqaExtra.begin(), gqaExtra.end()); + if (bias) a.push_back(fv); // bias [bh,Sq,Sk] f32, LAST return a; }; - // _pre : (Q,K,dO:store, O,L,Dd:f32, Sq,Sk:idx, scale:f32, causal:idx [+gqa]) + // _pre : (Q,K,dO:store, O,L,Dd:f32, Sq,Sk:idx, scale:f32, causal:idx [+gqa][+bias]) mk("_pre", withGqa({sv, sv, sv, fv, fv, fv, idxTy, idxTy, f32, idxTy}), [&](OpBuilder &bb, Location l, gpu::GPUFuncOp fn) { - emitPre(bb, l, fn, D, storeTy, gqa); + emitPre(bb, l, fn, D, storeTy, gqa, bias); }); - // _dkdv : (Q,K,V,dO:store, L,Dd:f32, dK,dV:f32, Sq,Sk:idx, scale, causal [+gqa]) + // _dkdv : (Q,K,V,dO:store, L,Dd:f32, dK,dV:f32, Sq,Sk:idx, scale, causal [+gqa][+bias]) mk("_dkdv", withGqa({sv, sv, sv, sv, fv, fv, fv, fv, idxTy, idxTy, f32, idxTy}), [&](OpBuilder &bb, Location l, gpu::GPUFuncOp fn) { - emitDkDv(bb, l, fn, D, storeTy, gqa); + emitDkDv(bb, l, fn, D, storeTy, gqa, bias); }); - // _dq : (Q,K,V,dO:store, L,Dd:f32, dQ:f32, Sq,Sk:idx, scale, causal [+gqa]) + // _dq : (Q,K,V,dO:store, L,Dd:f32, dQ:f32, Sq,Sk:idx, scale, causal [+gqa][+bias]) mk("_dq", withGqa({sv, sv, sv, sv, fv, fv, fv, idxTy, idxTy, f32, idxTy}), [&](OpBuilder &bb, Location l, gpu::GPUFuncOp fn) { - emitDq(bb, l, fn, D, storeTy, gqa); + emitDq(bb, l, fn, D, storeTy, gqa, bias); }); op->erase(); } diff --git a/tests/unit/test_kernel_inventory_and_lit_fixtures.py b/tests/unit/test_kernel_inventory_and_lit_fixtures.py index d7cc030b2..272673b8a 100644 --- a/tests/unit/test_kernel_inventory_and_lit_fixtures.py +++ b/tests/unit/test_kernel_inventory_and_lit_fixtures.py @@ -299,7 +299,7 @@ def test_doc_covers_required_sections(self): "Per-arch feature matrix", "MFMA instruction shape table", "Per-arch dtype matrix", - "Planned fused kernel inventory", + "Fused kernel inventory", "AMDGCN intrinsic patterns", "Execution gates", "ROCm 7.2.4", @@ -336,6 +336,53 @@ def test_doc_documents_cdna4_fp4_lanes(self): assert fp4 in doc, f"{fp4} missing (CDNA 4 FP4 coverage)" +# ────────────────────────────────────────────────────────────────────────── +# H-3 anti-drift: inventory execution-status must track the generated +# runtime execution matrix. Added 2026-07-10 after the doc silently +# drifted for weeks — the structural guards above never checked *what +# executes*, so §7/§9 kept claiming only matmul+flash_attn ran on gfx1151 +# while dozens of compiled HIP lanes had already landed. This ties the +# inventory's execution claims to the drift-gated matrix (Decision #26). +# ────────────────────────────────────────────────────────────────────────── + +class TestH3RocmInventoryExecutionStatus: + INVENTORY = REPO / "docs" / "rocm_mfma_kernel_inventory.md" + MATRIX = REPO / "docs" / "audit" / "generated" / "runtime_execution_matrix.md" + + def test_doc_points_at_generated_truth_surface(self): + doc = self.INVENTORY.read_text() + assert "runtime_execution_matrix" in doc, ( + "the inventory must point at the generated runtime execution matrix " + "as status truth (Decision #26) so execution status is not " + "hand-maintained here" + ) + + def test_doc_acknowledges_native_rocm_execution(self): + # A revert to the stale 'only matmul + flash_attn execute' framing would + # drop these markers of the native HIP compiled-lane program. + doc = self.INVENTORY.read_text() + assert "hip_runtime" in doc + assert "compiled" in doc + + def test_every_compiled_lane_named_in_doc_exists_in_matrix(self): + # The inventory must not name an executing rocm_*_compiled lane that the + # drift-gated matrix does not carry. This is the concrete anti-drift + # tie: doc execution claims are verifiable against generated truth. + import re + + doc = self.INVENTORY.read_text() + matrix = self.MATRIX.read_text() + lane_re = re.compile(r"rocm_[a-z0-9]+(?:_[a-z0-9]+)*_compiled") + doc_lanes = set(lane_re.findall(doc)) + matrix_lanes = set(lane_re.findall(matrix)) + assert doc_lanes, "inventory should name the executing rocm compiled lanes" + missing = sorted(doc_lanes - matrix_lanes) + assert not missing, ( + f"inventory names compiled lanes absent from the generated matrix: " + f"{missing}" + ) + + # ────────────────────────────────────────────────────────────────────────── # G-4: NVIDIA lit fixtures present + correct # ────────────────────────────────────────────────────────────────────────── diff --git a/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py b/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py new file mode 100644 index 000000000..2c68e1d26 --- /dev/null +++ b/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py @@ -0,0 +1,208 @@ +"""ROCm flash_attn BACKWARD reachable through runtime.launch() on gfx1151. + +`test_rocm_flash_attn_bwd_compiled.py` proves the compiler-generated FA-2 +backward kernel directly (build hsaco + hand-launch fa_pre/fa_dkdv/fa_dq). THIS +fixture proves the runtime *lane* — `compiler_path="rocm_flash_attn_bwd_compiled"` +routed through `runtime.launch()` — so the backward is a first-class executing +lane like the forward, not just a standalone kernel. Operands are (dO, Q, K, V); +O is recomputed on-device via the forward lane. dQ/dK/dV are compared to the +numpy attention-backward reference (the same math as autodiff `vjp_flash_attn`). + +Core MHA (scale + causal), f16/bf16 storage, f32 accumulate. Skip-clean: +tessera-opt not built / no GPU. +""" + +from __future__ import annotations + +import numpy as np +import pytest + + +def _rocm_or_skip(): + from tessera import runtime as rt + if rt._tessera_opt_path() is None: + pytest.skip("tessera-opt not built") + if not rt._rocm_wmma_runtime_available(): + pytest.skip("no usable AMD GPU") + return rt + + +def _ref_bwd(Q, K, V, dO, scale, causal, bias=None): + """numpy FA backward (f32 math from the f16/bf16 storage inputs). Handles + GQA/MQA: query head h reads KV head g = h // (H//G); dK/dV accumulate over the + group. Q/dO are [B,H,Sq,D]; K/V are [B,G,Sk,D]. Optional additive bias + [B,H,Sq,Sk] enters S before the softmax.""" + B, H, Sq, D = Q.shape + G, Sk = K.shape[1], K.shape[2] + ratio = H // G + Qf, Kf, Vf, dOf = (a.astype(np.float32) for a in (Q, K, V, dO)) + biasf = None if bias is None else np.asarray(bias, np.float32) + dQ = np.zeros((B, H, Sq, D), np.float32) + dK = np.zeros((B, G, Sk, D), np.float32) + dV = np.zeros((B, G, Sk, D), np.float32) + for b in range(B): + for h in range(H): + g = h // ratio + s = scale * (Qf[b, h] @ Kf[b, g].T) + if biasf is not None: + s = s + biasf[b, h] + if causal: + i = np.arange(Sq)[:, None]; j = np.arange(Sk)[None, :] + s = np.where(j > i, -1e30, s) + s = s - s.max(-1, keepdims=True) + p = np.exp(s); p = p / p.sum(-1, keepdims=True) + O = p @ Vf[b, g] + dp = dOf[b, h] @ Vf[b, g].T + dq_row = np.sum(O * dOf[b, h], axis=-1, keepdims=True) + ds = p * (dp - dq_row) + dQ[b, h] = scale * (ds @ Kf[b, g]) + dK[b, g] += scale * (ds.T @ Qf[b, h]) + dV[b, g] += p.T @ dOf[b, h] + return dQ, dK, dV + + +def _art(rt, causal, scale, bias=False): + names = ["do", "q", "k", "v"] + (["bias"] if bias else []) + return rt.RuntimeArtifact(metadata={ + "target": "rocm", "compiler_path": "rocm_flash_attn_bwd_compiled", + "executable": True, "execution_kind": "native_gpu", + "arg_names": names, "output_name": "g", + "ops": [{"op_name": "tessera.flash_attn_bwd", "result": "g", + "operands": names, + "kwargs": {"scale": scale, "causal": causal}}]}) + + +def _run(rt, dO, Q, K, V, *, causal, scale, bias=None): + call = (dO, Q, K, V) if bias is None else (dO, Q, K, V, bias) + res = rt.launch(_art(rt, causal, scale, bias is not None), call) + assert res["ok"] is True, res.get("reason") + assert res["compiler_path"] == "rocm_flash_attn_bwd_compiled" + dQ, dK, dV = res["output"] + return np.asarray(dQ), np.asarray(dK), np.asarray(dV) + + +def _relerr(got, ref): + return float(np.max(np.abs(got - ref)) / (np.max(np.abs(ref)) + 1e-6)) + + +@pytest.mark.parametrize("D,B,H,Sq,Sk,causal", [ + (16, 1, 1, 16, 16, False), + (16, 1, 2, 32, 48, False), + (64, 2, 2, 48, 48, True), + (16, 1, 1, 20, 40, False), # ragged Sq/Sk + (64, 1, 2, 32, 32, True), # causal, D=64 +]) +def test_bwd_runtime_lane_matches_numpy(D, B, H, Sq, Sk, causal): + rt = _rocm_or_skip() + rng = np.random.default_rng(7 + D + Sq + Sk + int(causal)) + Q = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + K = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + scale = 1.0 / float(np.sqrt(D)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, causal) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=causal, scale=scale) + assert dQ.shape == Q.shape and dK.shape == K.shape and dV.shape == V.shape + tol = 5e-3 # f16 storage, f32 accumulate — measured ~2-4e-4 + eQ, eK, eV = _relerr(dQ, dQ_ref), _relerr(dK, dK_ref), _relerr(dV, dV_ref) + assert eQ < tol and eK < tol and eV < tol, ( + f"rel-err dQ={eQ:.2e} dK={eK:.2e} dV={eV:.2e} " + f"@ D={D} {B}x{H}x{Sq}x{Sk} causal={causal}") + + +def test_bwd_runtime_lane_bf16(): + rt = _rocm_or_skip() + bf16 = rt._bfloat16_dtype() + if bf16 is None: + pytest.skip("no bfloat16 dtype available") + rng = np.random.default_rng(99) + B, H, Sq, Sk, D = 1, 2, 32, 32, 16 + Q = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(bf16) + K = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(bf16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(bf16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(bf16) + scale = 1.0 / float(np.sqrt(D)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, False) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=False, scale=scale) + tol = 3e-2 # bf16 has ~8 mantissa bits — looser bound + assert (_relerr(dQ, dQ_ref) < tol and _relerr(dK, dK_ref) < tol + and _relerr(dV, dV_ref) < tol) + + +@pytest.mark.parametrize("D,B,H,G,Sq,Sk,causal", [ + (16, 1, 8, 1, 32, 32, False), # MQA (one shared KV head) + (16, 2, 8, 2, 32, 48, False), # GQA, ratio 4 + (64, 1, 8, 4, 48, 48, True), # GQA, ratio 2, causal + (16, 1, 4, 4, 32, 32, False), # MHA-equivalence (G == H) +]) +def test_bwd_runtime_lane_gqa_matches_numpy(D, B, H, G, Sq, Sk, causal): + rt = _rocm_or_skip() + rng = np.random.default_rng(23 + D + H + G + int(causal)) + Q = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + K = (rng.standard_normal((B, G, Sk, D)) * 0.3).astype(np.float16) + V = (rng.standard_normal((B, G, Sk, D)) * 0.3).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + scale = 1.0 / float(np.sqrt(D)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, causal) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=causal, scale=scale) + assert dQ.shape == Q.shape and dK.shape == K.shape and dV.shape == V.shape + tol = 5e-3 + eQ, eK, eV = _relerr(dQ, dQ_ref), _relerr(dK, dK_ref), _relerr(dV, dV_ref) + assert eQ < tol and eK < tol and eV < tol, ( + f"GQA rel-err dQ={eQ:.2e} dK={eK:.2e} dV={eV:.2e} " + f"@ D={D} H={H} G={G} causal={causal}") + + +@pytest.mark.parametrize("D,B,H,Sq,Sk,causal", [ + (16, 1, 2, 32, 32, False), + (16, 1, 2, 24, 40, False), # ragged + (64, 1, 2, 32, 32, True), # causal + bias +]) +def test_bwd_runtime_lane_attn_bias_matches_numpy(D, B, H, Sq, Sk, causal): + rt = _rocm_or_skip() + rng = np.random.default_rng(31 + D + Sq + Sk + int(causal)) + Q = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + K = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + bias = (rng.standard_normal((B, H, Sq, Sk)) * 0.5).astype(np.float32) + scale = 1.0 / float(np.sqrt(D)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, causal, bias=bias) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=causal, scale=scale, bias=bias) + tol = 5e-3 + eQ, eK, eV = _relerr(dQ, dQ_ref), _relerr(dK, dK_ref), _relerr(dV, dV_ref) + assert eQ < tol and eK < tol and eV < tol, ( + f"bias rel-err dQ={eQ:.2e} dK={eK:.2e} dV={eV:.2e} " + f"@ D={D} {B}x{H}x{Sq}x{Sk} causal={causal}") + + +def test_bwd_runtime_lane_bias_broadcast_per_head(): + # Bias broadcast from [H,Sq,Sk] (no batch axis) to B*H — exercises the + # host broadcast path. + rt = _rocm_or_skip() + rng = np.random.default_rng(88) + B, H, Sq, Sk, D = 2, 2, 16, 16, 16 + Q = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + K = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + bias_h = (rng.standard_normal((H, Sq, Sk)) * 0.5).astype(np.float32) + scale = 1.0 / float(np.sqrt(D)) + full = np.broadcast_to(bias_h, (B, H, Sq, Sk)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, False, bias=full) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=False, scale=scale, bias=bias_h) + tol = 5e-3 + assert (_relerr(dQ, dQ_ref) < tol and _relerr(dK, dK_ref) < tol + and _relerr(dV, dV_ref) < tol) + + +def test_bwd_runtime_lane_rejects_bad_group(): + # Query heads not divisible by KV heads → reject clearly (Decision #21). + rt = _rocm_or_skip() + B, Hq, Hkv, Sq, Sk, D = 1, 6, 4, 16, 16, 16 # 6 % 4 != 0 + Q = np.zeros((B, Hq, Sq, D), np.float16) + K = np.zeros((B, Hkv, Sk, D), np.float16) + V = np.zeros((B, Hkv, Sk, D), np.float16) + dO = np.zeros((B, Hq, Sq, D), np.float16) + res = rt.launch(_art(rt, False, 0.25), (dO, Q, K, V)) + assert res["ok"] is False diff --git a/tests/unit/test_rocm_kv_cache_compiled.py b/tests/unit/test_rocm_kv_cache_compiled.py new file mode 100644 index 000000000..51032653b --- /dev/null +++ b/tests/unit/test_rocm_kv_cache_compiled.py @@ -0,0 +1,137 @@ +"""KV-cache paged-movement lane on AMD ROCm gfx1151 (§5.6 of the ROCm MFMA +kernel inventory) — the append/read/prune core executed on-device by COMPOSING +the existing gfx1151 scatter (append row write) + masked-gather (read/prune) +kernels with host page-index math. Reachable via +`compiler_path="rocm_kv_cache_compiled"`. Validated against the KVCacheHandle +reference on gfx1151. Skip-clean: tessera-opt not built / no GPU. + +The three kv_cache_* ops are stateful over a KVCacheHandle; this lane executes +their tensor movement over a resident cache buffer `(max_seq, H, D)`. K and V +are independent buffers, so the lane runs once per buffer — the tests drive the +K buffer and compare against `handle.keys` (the V path is identical). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import tessera + + +def _rocm_or_skip(): + from tessera import runtime as rt + if rt._tessera_opt_path() is None: + pytest.skip("tessera-opt not built") + if not rt._rocm_wmma_runtime_available(): + pytest.skip("no usable AMD GPU") + return rt + + +def _art(rt, op, operands, kwargs): + names = ["b", "n"][: len(operands)] + return rt.RuntimeArtifact(metadata={ + "target": "rocm", "compiler_path": "rocm_kv_cache_compiled", + "executable": True, "execution_kind": "native_gpu", + "arg_names": names, "output_name": "o", + "ops": [{"op_name": op, "result": "o", "operands": names, + "kwargs": kwargs}]}) + + +def _run(rt, op, operands, **kwargs): + res = rt.launch(_art(rt, op, operands, kwargs), tuple(operands)) + assert res["ok"] is True, res.get("reason") + assert res["compiler_path"] == "rocm_kv_cache_compiled" + return np.asarray(res["output"], np.float32) + + +_RNG = np.random.default_rng(41) + + +def _handle(max_seq=32, H=3, D=8): + return tessera.cache.KVCacheHandle(num_heads=H, head_dim=D, max_seq=max_seq) + + +# ── append ──────────────────────────────────────────────────────────────── + +def test_append_from_empty(): + rt = _rocm_or_skip() + h = _handle() + k = _RNG.standard_normal((5, 3, 8)).astype(np.float32) + v = _RNG.standard_normal((5, 3, 8)).astype(np.float32) + h.append(k, v) + buf = np.zeros((32, 3, 8), np.float32) # resident K buffer, empty + out = _run(rt, "tessera.kv_cache.append", [buf, k], start=0) + np.testing.assert_allclose(out, h.keys.astype(np.float32), rtol=0, atol=0) + + +def test_append_midstream(): + rt = _rocm_or_skip() + h = _handle() + k0 = _RNG.standard_normal((7, 3, 8)).astype(np.float32) + h.append(k0, k0) # prefill 7 tokens + k1 = _RNG.standard_normal((4, 3, 8)).astype(np.float32) + # lane appends onto the buffer that already holds the prefill. + buf = h.keys.astype(np.float32).copy() + out = _run(rt, "tessera.kv_cache.append", [buf, k1], start=h.current_seq) + h.append(k1, k1) + np.testing.assert_allclose(out, h.keys.astype(np.float32), rtol=0, atol=0) + + +def test_append_out_of_bounds_rejected(): + rt = _rocm_or_skip() + buf = np.zeros((8, 2, 4), np.float32) + new = _RNG.standard_normal((5, 2, 4)).astype(np.float32) + res = rt.launch( + _art(rt, "tessera.kv_cache.append", [buf, new], {"start": 6}), + (buf, new)) + assert res["ok"] is False # [6, 11) exceeds max_seq=8 + + +# ── read ────────────────────────────────────────────────────────────────── + +def test_read_slice(): + rt = _rocm_or_skip() + h = _handle() + k = _RNG.standard_normal((12, 3, 8)).astype(np.float32) + h.append(k, k) + ks, _ = h.read(2, 9) + out = _run(rt, "tessera.kv_cache.read", [h.keys.astype(np.float32)], + start=2, end=9) + np.testing.assert_allclose(out, np.asarray(ks, np.float32), rtol=0, atol=0) + + +def test_read_single_token_default_end(): + rt = _rocm_or_skip() + h = _handle() + k = _RNG.standard_normal((6, 3, 8)).astype(np.float32) + h.append(k, k) + ks, _ = h.read(4) # default end = start+1 + out = _run(rt, "tessera.kv_cache.read", [h.keys.astype(np.float32)], + start=4) + np.testing.assert_allclose(out, np.asarray(ks, np.float32), rtol=0, atol=0) + + +# ── prune ───────────────────────────────────────────────────────────────── + +def test_prune_sliding_window(): + rt = _rocm_or_skip() + h = _handle() + k = _RNG.standard_normal((20, 3, 8)).astype(np.float32) + h.append(k, k) + buf = h.keys.astype(np.float32).copy() + out = _run(rt, "tessera.kv_cache.prune", [buf], limit=6, + current_seq=h.current_seq) + h.prune(6) # keep trailing 6 + np.testing.assert_allclose(out, h.keys.astype(np.float32), rtol=0, atol=0) + + +def test_prune_noop_when_limit_exceeds_seq(): + rt = _rocm_or_skip() + h = _handle() + k = _RNG.standard_normal((5, 3, 8)).astype(np.float32) + h.append(k, k) + buf = h.keys.astype(np.float32).copy() + out = _run(rt, "tessera.kv_cache.prune", [buf], limit=99, + current_seq=h.current_seq) + np.testing.assert_allclose(out, buf, rtol=0, atol=0) diff --git a/tests/unit/test_rocm_structured_compute_compiled.py b/tests/unit/test_rocm_structured_compute_compiled.py index 6d8fd1ab8..e16ba8911 100644 --- a/tests/unit/test_rocm_structured_compute_compiled.py +++ b/tests/unit/test_rocm_structured_compute_compiled.py @@ -100,6 +100,17 @@ def test_rocm_structured_model_recurrent_and_stencil_match_reference_on_gpu(): atol=1e-6, ) + # LSTM cell (gate order i,f,g,o): W_ih/W_hh are (in,4H)/(H,4H); returns + # (h_new, c_new). Runs through the same structured-compute lane. + cprev = rng.standard_normal((2, 5)).astype(np.float32) + Wih4 = rng.standard_normal((3, 20)).astype(np.float32) + Whh4 = rng.standard_normal((5, 20)).astype(np.float32) + lh, lc = _launch(rt, "tessera.lstm_cell", ("x", "h", "c", "Wih", "Whh"), + (xt, h, cprev, Wih4, Whh4)) + rh, rc = F.lstm_cell(xt, h, cprev, Wih4, Whh4) + np.testing.assert_allclose(lh, rh, atol=1e-6) + np.testing.assert_allclose(lc, rc, atol=1e-6) + a = rng.standard_normal((2, 4)).astype(np.float32) weight = rng.standard_normal((4, 6)).astype(np.float32) la = rng.standard_normal((4, 2)).astype(np.float32) diff --git a/tests/unit/test_x86_structured_compute_compiled.py b/tests/unit/test_x86_structured_compute_compiled.py index b97f94d8f..e15b159da 100644 --- a/tests/unit/test_x86_structured_compute_compiled.py +++ b/tests/unit/test_x86_structured_compute_compiled.py @@ -97,6 +97,15 @@ def test_x86_structured_model_recurrent_and_stencil_match_reference(): atol=1e-6, ) + cprev = rng.standard_normal((2, 5)).astype(np.float32) + Wih4 = rng.standard_normal((3, 20)).astype(np.float32) + Whh4 = rng.standard_normal((5, 20)).astype(np.float32) + lh, lc = _launch("tessera.lstm_cell", ("x", "h", "c", "Wih", "Whh"), + (xt, h, cprev, Wih4, Whh4)) + rh, rc = F.lstm_cell(xt, h, cprev, Wih4, Whh4) + np.testing.assert_allclose(lh, rh, atol=1e-6) + np.testing.assert_allclose(lc, rc, atol=1e-6) + a = rng.standard_normal((2, 4)).astype(np.float32) weight = rng.standard_normal((4, 6)).astype(np.float32) la = rng.standard_normal((4, 2)).astype(np.float32) From 71664ceb1551fcd0fa2e793800b2dd8f48e21a17 Mon Sep 17 00:00:00 2001 From: gstoner Date: Fri, 10 Jul 2026 11:35:47 -0600 Subject: [PATCH 2/5] rocm/gfx1151: native moe_transport dispatch + combine (device gather/scatter-add) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Promote the MoE transport ops off the reference oracle to real gfx1151 execution: - moe_dispatch → the device gather kernel: token_of_slot = sort_perm // top_k row gather (reuses _rocm_gather_rows). Bit-exact vs the numpy oracle. - moe_combine → the device scatter (add) kernel: host pre-scales each packed row by its route weight (exact f32), then atomic scatter-add to token order. f32 accumulate — matches the f64 oracle within ~1.2e-7. Both report native_gpu via the per-op (output, execution_kind) launch override; they fall back to the stdlib oracle + reference_cpu off-box. grouped_swiglu (the expert GEMM) stays on the oracle / reference_cpu — a native f32-exact grouped GEMM is a separate follow-up (WMMA is f16). This was exactly the promotion the matrix note called for ("requires HIP gather/scatter transport kernels"). Result: every ROCm row in the generated runtime execution matrix is now native_gpu — zero reference_cpu ROCm lanes remain. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generated/runtime_execution_matrix.csv | 2 +- .../generated/runtime_execution_matrix.md | 2 +- docs/rocm_mfma_kernel_inventory.md | 6 +- python/tessera/compiler/execution_matrix.py | 19 +++-- python/tessera/runtime.py | 73 +++++++++++++++---- .../unit/test_rocm_moe_transport_compiled.py | 21 +++++- 6 files changed, 94 insertions(+), 29 deletions(-) diff --git a/docs/audit/generated/runtime_execution_matrix.csv b/docs/audit/generated/runtime_execution_matrix.csv index 91cd0882e..d604c0160 100644 --- a/docs/audit/generated/runtime_execution_matrix.csv +++ b/docs/audit/generated/runtime_execution_matrix.csv @@ -58,7 +58,7 @@ rocm,rocm_loss_compiled,native_gpu,1,rocm_loss_compiled,success,hip_runtime,"ROC rocm,rocm_matmul_family_compiled,native_gpu,1,rocm_matmul_family_compiled,success,hip_runtime,"ROCm matmul-family artifact runs the COMPILER-GENERATED WMMA GEMM kernel (the rocm_compiled spine) reshaped/batched/split in the runtime — batched_gemm, linear_general, qkv_projection, factorized_matmul (GPU matmul + exact host SVD-truncate), and single-contraction einsum. f16/bf16, f32 accumulate." rocm,rocm_metric_loss_compiled,native_gpu,1,rocm_metric_loss_compiled,success,hip_runtime,"ROCm metric-loss artifact runs wasserstein / cosine_embedding / contrastive / triplet / InfoNCE / NT-Xent / seq2seq losses with generated ROCm reductions and exp/log; label, mask, sort, and compact matrix structure remain on host. f32, matches tessera.losses." rocm,rocm_moe_compiled,native_gpu,1,rocm_moe_compiled,success,hip_runtime,"ROCm moe-compute lane runs the routed per-token expert GEMVs (top-1) on the COMPILER-GENERATED gfx1151 kernel (generate-rocm-moe-kernel: one thread per (token, out-col); routing resolved on host) HIP-launched. f32, matches numpy." -rocm,rocm_moe_transport_compiled,reference_cpu,1,rocm_moe_transport_compiled,success,reference_oracle,"ROCm DK3 MoE transport artifact executes moe_dispatch, moe_combine, and grouped_swiglu against the stdlib DispatchPlan oracle through the runtime ABI. This pins token permutation, group sizes, capacity drops, and combine weights; promotion to native_gpu requires HIP gather/scatter transport kernels." +rocm,rocm_moe_transport_compiled,native_gpu,1,rocm_moe_transport_compiled,success,hip_runtime,"ROCm DK3 MoE transport: moe_dispatch runs NATIVELY on the gfx1151 device gather kernel (token_of_slot = sort_perm//top_k row gather) and moe_combine on the device scatter (add) kernel (host pre-scales each packed row by its route weight, then atomic scatter-add to token order) — both report native_gpu vs the stdlib DispatchPlan oracle. grouped_swiglu (the expert GEMM) stays on the oracle and reports reference_cpu via the per-op execution_kind override — a native f32-exact grouped GEMM is a separate follow-up (WMMA is f16). Off-box the transport ops fall back to the oracle + reference_cpu." rocm,rocm_muon_compiled,native_gpu,1,rocm_muon_compiled,success,hip_runtime,"ROCm Muon lane orthogonalizes the momentum matrix via the gfx1151 SVD kernel (U·Vh polar factor); the small U@Vh + momentum/sgd run on host. <2-D params normalize. f32, matches optim.muon." rocm,rocm_norm_compiled,native_gpu,1,rocm_norm_compiled,success,hip_runtime,"ROCm norm artifact runs the COMPILER-GENERATED RDNA row-reduction kernel (unweighted rmsnorm / layer_norm over the last axis, one workgroup per row, LDS tree-reduce of Σx and Σx²): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Handles tessera.rmsnorm(_safe) + tessera.layer_norm by op name." rocm,rocm_normcompose_compiled,native_gpu,1,rocm_normcompose_compiled,success,hip_runtime,"ROCm group/instance/weight-norm lane composed on the gfx1151 layer_norm (row mean/var) + reduce (sum-of-squares) kernels; host does the reshape / per-axis divide. f32, matches nn.functional." diff --git a/docs/audit/generated/runtime_execution_matrix.md b/docs/audit/generated/runtime_execution_matrix.md index a38cadd85..18cb0034d 100644 --- a/docs/audit/generated/runtime_execution_matrix.md +++ b/docs/audit/generated/runtime_execution_matrix.md @@ -72,7 +72,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `rocm` | `rocm_matmul_family_compiled` | `rocm_matmul_family_compiled` | `native_gpu` | `hip_runtime` | ROCm matmul-family artifact runs the COMPILER-GENERATED WMMA GEMM kernel (the rocm_compiled spine) reshaped/batched/split in the runtime — batched_gemm, linear_general, qkv_projection, factorized_matmul (GPU matmul + exact host SVD-truncate), and single-contraction einsum. f16/bf16, f32 accumulate. | | `rocm` | `rocm_metric_loss_compiled` | `rocm_metric_loss_compiled` | `native_gpu` | `hip_runtime` | ROCm metric-loss artifact runs wasserstein / cosine_embedding / contrastive / triplet / InfoNCE / NT-Xent / seq2seq losses with generated ROCm reductions and exp/log; label, mask, sort, and compact matrix structure remain on host. f32, matches tessera.losses. | | `rocm` | `rocm_moe_compiled` | `rocm_moe_compiled` | `native_gpu` | `hip_runtime` | ROCm moe-compute lane runs the routed per-token expert GEMVs (top-1) on the COMPILER-GENERATED gfx1151 kernel (generate-rocm-moe-kernel: one thread per (token, out-col); routing resolved on host) HIP-launched. f32, matches numpy. | -| `rocm` | `rocm_moe_transport_compiled` | `rocm_moe_transport_compiled` | `reference_cpu` | `reference_oracle` | ROCm DK3 MoE transport artifact executes moe_dispatch, moe_combine, and grouped_swiglu against the stdlib DispatchPlan oracle through the runtime ABI. This pins token permutation, group sizes, capacity drops, and combine weights; promotion to native_gpu requires HIP gather/scatter transport kernels. | +| `rocm` | `rocm_moe_transport_compiled` | `rocm_moe_transport_compiled` | `native_gpu` | `hip_runtime` | ROCm DK3 MoE transport: moe_dispatch runs NATIVELY on the gfx1151 device gather kernel (token_of_slot = sort_perm//top_k row gather) and moe_combine on the device scatter (add) kernel (host pre-scales each packed row by its route weight, then atomic scatter-add to token order) — both report native_gpu vs the stdlib DispatchPlan oracle. grouped_swiglu (the expert GEMM) stays on the oracle and reports reference_cpu via the per-op execution_kind override — a native f32-exact grouped GEMM is a separate follow-up (WMMA is f16). Off-box the transport ops fall back to the oracle + reference_cpu. | | `rocm` | `rocm_muon_compiled` | `rocm_muon_compiled` | `native_gpu` | `hip_runtime` | ROCm Muon lane orthogonalizes the momentum matrix via the gfx1151 SVD kernel (U·Vh polar factor); the small U@Vh + momentum/sgd run on host. <2-D params normalize. f32, matches optim.muon. | | `rocm` | `rocm_norm_compiled` | `rocm_norm_compiled` | `native_gpu` | `hip_runtime` | ROCm norm artifact runs the COMPILER-GENERATED RDNA row-reduction kernel (unweighted rmsnorm / layer_norm over the last axis, one workgroup per row, LDS tree-reduce of Σx and Σx²): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Handles tessera.rmsnorm(_safe) + tessera.layer_norm by op name. | | `rocm` | `rocm_normcompose_compiled` | `rocm_normcompose_compiled` | `native_gpu` | `hip_runtime` | ROCm group/instance/weight-norm lane composed on the gfx1151 layer_norm (row mean/var) + reduce (sum-of-squares) kernels; host does the reshape / per-axis divide. f32, matches nn.functional. | diff --git a/docs/rocm_mfma_kernel_inventory.md b/docs/rocm_mfma_kernel_inventory.md index e61e05cec..406bc12ca 100644 --- a/docs/rocm_mfma_kernel_inventory.md +++ b/docs/rocm_mfma_kernel_inventory.md @@ -493,8 +493,10 @@ promotes them to `compileable`. See §9 for the concrete done / open / blocked s execute-compare vs `KVCacheHandle`); a single fused gather→attention paged kernel is the remaining step - `grad_clip_norm` (§5.5) — global-norm + scale; single-node -- `rocm_moe_transport_compiled` — the one lane still `reference_cpu`; needs native - HIP gather/scatter transport kernels for `moe_dispatch`/`moe_combine` +- `grouped_swiglu` native f32 expert GEMM — `moe_dispatch` (device gather) and + `moe_combine` (device scatter-add) now execute natively (`rocm_moe_transport_compiled` + reports `native_gpu`); the expert GEMM stays on the oracle pending an f32-exact + device GEMM (WMMA is f16) - Plain recurrent cells `lstm_cell` / `gru_cell` / `simple_rnn_cell` (§5.9) — no dedicated lane (selective-SSM + deltanet are done) - **Perf ladders / MFU sign-off** beyond `matmul` — every `compiled` lane is diff --git a/python/tessera/compiler/execution_matrix.py b/python/tessera/compiler/execution_matrix.py index d720359b1..6b149afcb 100644 --- a/python/tessera/compiler/execution_matrix.py +++ b/python/tessera/compiler/execution_matrix.py @@ -2157,14 +2157,19 @@ class ExecutionRow: execution_mode="hip_runtime"), ("rocm", "rocm_moe_transport_compiled"): ExecutionRow( target="rocm", compiler_path="rocm_moe_transport_compiled", - execution_kind="reference_cpu", executable=True, + execution_kind="native_gpu", executable=True, executor_id="rocm_moe_transport_compiled", runtime_status="success", - reason="ROCm DK3 MoE transport artifact executes moe_dispatch, " - "moe_combine, and grouped_swiglu against the stdlib DispatchPlan " - "oracle through the runtime ABI. This pins token permutation, " - "group sizes, capacity drops, and combine weights; promotion to " - "native_gpu requires HIP gather/scatter transport kernels.", - execution_mode="reference_oracle"), + reason="ROCm DK3 MoE transport: moe_dispatch runs NATIVELY on the " + "gfx1151 device gather kernel (token_of_slot = sort_perm//top_k " + "row gather) and moe_combine on the device scatter (add) kernel " + "(host pre-scales each packed row by its route weight, then " + "atomic scatter-add to token order) — both report native_gpu vs " + "the stdlib DispatchPlan oracle. grouped_swiglu (the expert GEMM) " + "stays on the oracle and reports reference_cpu via the per-op " + "execution_kind override — a native f32-exact grouped GEMM is a " + "separate follow-up (WMMA is f16). Off-box the transport ops " + "fall back to the oracle + reference_cpu.", + execution_mode="hip_runtime"), ("rocm", "rocm_normcompose_compiled"): ExecutionRow( target="rocm", compiler_path="rocm_normcompose_compiled", execution_kind="native_gpu", executable=True, diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index ca251d0c7..cc4f81bbf 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -1451,12 +1451,49 @@ def _rocm_dequant_grouped_gemm_native(x: Any, packed_experts: Any, return out -def _execute_rocm_moe_transport_reference(artifact: RuntimeArtifact, args: Any) -> Any: - """DK3 transition executor for ``rocm_moe_transport_compiled``. +def _moe_dispatch_native(x: Any, plan: Any, np: Any) -> Any: + """DK3 dispatch = gather the token row for each kept, expert-sorted slot + (token_of_slot = sort_perm // top_k) on the gfx1151 device gather kernel. + Exact vs the numpy gather (pure row movement, f32).""" + xa = np.ascontiguousarray(_as_numpy(x), np.float32) + if xa.ndim != 2: + raise ValueError(f"moe_dispatch x must be 2-D (T,H); got {xa.shape}") + tof = (np.asarray(plan.sort_perm, np.int64) // int(plan.top_k)) + if tof.shape[0] == 0: + return np.zeros((0, int(xa.shape[1])), np.float32) + return _rocm_gather_rows(xa, tof, np) + + +def _moe_combine_native(y_packed: Any, plan: Any, np: Any) -> Any: + """DK3 combine = weighted scatter-add of the packed expert rows back to token + order: out[t] += w[slot]*yp[i], t = slot // top_k. Host pre-scales each row + by its route weight (exact f32), then the gfx1151 device scatter kernel + (mode add) accumulates into the token rows. f32 accumulate.""" + yp = np.ascontiguousarray(_as_numpy(y_packed), np.float32) + if yp.ndim != 2: + raise ValueError(f"moe_combine partials must be 2-D (S,H); got {yp.shape}") + S, H = int(yp.shape[0]), int(yp.shape[1]) + T = int(plan.num_tokens) + out = np.zeros((T, H), np.float32) + if S == 0: + return out + slots = np.asarray(plan.sort_perm, np.int64) + w = np.asarray(plan.weights, np.float32).reshape(-1)[slots].reshape(S, 1) + scaled = np.ascontiguousarray(yp * w, np.float32) + tof = (slots // int(plan.top_k)).astype(np.int64) + _rocm_scatter(out, scaled, tof, T, H, 1, np) # mode 1 = atomic add + return out - This pins the MoE transport/runtime ABI against the stdlib oracle while the - native ROCm gather/scatter and grouped-GEMM transport kernels are promoted. - """ + +def _execute_rocm_moe_transport(artifact: RuntimeArtifact, args: Any) -> Any: + """``rocm_moe_transport_compiled`` executor. The transport ops run natively + on gfx1151 — ``moe_dispatch`` on the device gather kernel, ``moe_combine`` on + the device scatter (add) kernel — and report ``native_gpu`` (falling back to + the stdlib oracle + ``reference_cpu`` when tessera-opt / a GPU is absent). + ``grouped_swiglu`` (the expert GEMM) stays on the stdlib oracle and reports + ``reference_cpu`` — a native f32-exact grouped GEMM is a separate follow-up. + Per-op ``execution_kind`` via the ``(output, kind)`` launch override.""" + import numpy as np from .stdlib import moe metadata = artifact.metadata or {} @@ -1467,25 +1504,31 @@ def _execute_rocm_moe_transport_reference(artifact: RuntimeArtifact, args: Any) if op_name == "tessera.moe_dispatch": try: - return moe.dispatch(values["x"], values["plan"]) + x, plan = values["x"], values["plan"] except KeyError as exc: raise ValueError(f"missing moe_dispatch argument {exc.args[0]!r}") from exc + try: + return (_moe_dispatch_native(x, plan, np), "native_gpu") + except _RocmCompiledUnavailable: + return (moe.dispatch(x, plan), "reference_cpu") if op_name == "tessera.moe_combine": try: - return moe.combine(values["partials"], values["plan"]) + partials, plan = values["partials"], values["plan"] except KeyError as exc: raise ValueError(f"missing moe_combine argument {exc.args[0]!r}") from exc + try: + return (_moe_combine_native(partials, plan, np), "native_gpu") + except _RocmCompiledUnavailable: + return (moe.combine(partials, plan), "reference_cpu") if op_name == "tessera.grouped_swiglu": try: - return moe.grouped_swiglu( - values["x_packed"], - values["w_gate"], - values["w_up"], - values["w_down"], - values["group_sizes"], - ) + out = moe.grouped_swiglu( + values["x_packed"], values["w_gate"], values["w_up"], + values["w_down"], values["group_sizes"]) except KeyError as exc: raise ValueError(f"missing grouped_swiglu argument {exc.args[0]!r}") from exc + # Expert GEMM: no f32-exact device GEMM yet (WMMA is f16) — stays oracle. + return (out, "reference_cpu") raise ValueError( "rocm_moe_transport_compiled executor handles tessera.moe_dispatch, " "tessera.moe_combine, or tessera.grouped_swiglu; " @@ -15882,7 +15925,7 @@ def _executor_table(): "rocm_sparse_compiled": _execute_rocm_compiled_sparse, "rocm_sparse_attn_compiled": _execute_rocm_compiled_sparse_attention, "rocm_moe_compiled": _execute_rocm_compiled_moe, - "rocm_moe_transport_compiled": _execute_rocm_moe_transport_reference, + "rocm_moe_transport_compiled": _execute_rocm_moe_transport, "rocm_optimizer_compiled": _execute_rocm_compiled_optimizer, "rocm_complex_compiled": _execute_rocm_compiled_complex, "rocm_conformal_compiled": _execute_rocm_compiled_conformal, diff --git a/tests/unit/test_rocm_moe_transport_compiled.py b/tests/unit/test_rocm_moe_transport_compiled.py index d6abd5f3e..81ad271f2 100644 --- a/tests/unit/test_rocm_moe_transport_compiled.py +++ b/tests/unit/test_rocm_moe_transport_compiled.py @@ -33,6 +33,13 @@ def _plan(seed=41, tokens=12, experts=4, top_k=2, capacity=5): return moe.plan_dispatch(eids, weights, experts, capacity=capacity) +def _expect_native(): + """The transport ops run natively iff tessera-opt is built AND a usable AMD + GPU is present (else the executor falls back to the stdlib oracle).""" + return (rt._tessera_opt_path() is not None + and rt._rocm_wmma_runtime_available()) + + def test_rocm_moe_dispatch_runtime_matches_dispatch_plan_oracle(): rng = np.random.default_rng(42) x = rng.standard_normal((12, 8)).astype(np.float32) @@ -45,7 +52,10 @@ def test_rocm_moe_dispatch_runtime_matches_dispatch_plan_oracle(): assert res["ok"] assert res["compiler_path"] == "rocm_moe_transport_compiled" - assert res["execution_kind"] == "reference_cpu" + # dispatch is a pure row gather — native on-box, oracle off-box; exact either + # way (device gather copies rows bit-for-bit). + assert res["execution_kind"] == ("native_gpu" if _expect_native() + else "reference_cpu") np.testing.assert_allclose(res["output"], moe.dispatch(x, plan), rtol=0, atol=0) @@ -62,8 +72,13 @@ def test_rocm_moe_combine_runtime_matches_weighted_combine_oracle(): ) assert res["ok"] - assert res["execution_kind"] == "reference_cpu" - np.testing.assert_allclose(res["output"], moe.combine(partials, plan), rtol=0, atol=0) + assert res["execution_kind"] == ("native_gpu" if _expect_native() + else "reference_cpu") + # native combine accumulates in f32 (device scatter-add) vs the f64 oracle — + # matches within f32 tolerance; the tiny tolerance still pins the plan + # permutation / route weights / capacity drops. + np.testing.assert_allclose(res["output"], moe.combine(partials, plan), + rtol=1e-5, atol=1e-6) def test_rocm_grouped_swiglu_runtime_matches_grouped_gemm_oracle(): From 4f6f3b3ca9b0d109707784b41bf7c4e70b28d4a4 Mon Sep 17 00:00:00 2001 From: gstoner Date: Fri, 10 Jul 2026 11:48:22 -0600 Subject: [PATCH 3/5] rocm/gfx1151: native f32 GEMM kernel -> full-native moe_transport (grouped_swiglu) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a plain f32 GEMM device kernel and use it to promote grouped_swiglu, so the whole rocm_moe_transport_compiled row is native (the last non-transport op). - New C++ pass generate-rocm-gemm-f32-kernel (GenerateROCMGemmF32Kernel.cpp, tessera_rocm.gemm_f32 ODS op): C[m,n] = sum_k A[m,k]*B[k,n], one thread per output element, scalar f32 k-loop (f32 accumulate). RDNA WMMA is f16/bf16 only, so this is the f32-exact VALU/FMA fallback the expert GEMM needs; correctness- first (no LDS tiling — the blocked perf ladder is a follow-up). - Runtime _rocm_f32_gemm (build hsaco + HIP launch) + _moe_grouped_swiglu_native: per expert group, g=x@Wg, u=x@Wu on the device GEMM, silu*mul host-side, then h@Wd on the device GEMM. grouped_swiglu now reports native_gpu; matches the f64 oracle within ~2.4e-7 (f32). Falls back to the oracle + reference_cpu off-box. Every ROCm row in the generated runtime execution matrix is now native_gpu. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../generated/runtime_execution_matrix.csv | 2 +- .../generated/runtime_execution_matrix.md | 2 +- docs/rocm_mfma_kernel_inventory.md | 7 +- python/tessera/compiler/execution_matrix.py | 20 +-- python/tessera/runtime.py | 116 +++++++++++++++- .../include/TesseraROCM/IR/TesseraROCMOps.td | 13 ++ .../include/TesseraROCM/Passes.h | 1 + .../lib/Conversion/CMakeLists.txt | 1 + .../Conversion/GenerateROCMGemmF32Kernel.cpp | 126 ++++++++++++++++++ .../lib/Conversion/Passes.cpp | 1 + .../unit/test_rocm_moe_transport_compiled.py | 9 +- 11 files changed, 274 insertions(+), 24 deletions(-) create mode 100644 src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMGemmF32Kernel.cpp diff --git a/docs/audit/generated/runtime_execution_matrix.csv b/docs/audit/generated/runtime_execution_matrix.csv index d604c0160..e0d1fb563 100644 --- a/docs/audit/generated/runtime_execution_matrix.csv +++ b/docs/audit/generated/runtime_execution_matrix.csv @@ -58,7 +58,7 @@ rocm,rocm_loss_compiled,native_gpu,1,rocm_loss_compiled,success,hip_runtime,"ROC rocm,rocm_matmul_family_compiled,native_gpu,1,rocm_matmul_family_compiled,success,hip_runtime,"ROCm matmul-family artifact runs the COMPILER-GENERATED WMMA GEMM kernel (the rocm_compiled spine) reshaped/batched/split in the runtime — batched_gemm, linear_general, qkv_projection, factorized_matmul (GPU matmul + exact host SVD-truncate), and single-contraction einsum. f16/bf16, f32 accumulate." rocm,rocm_metric_loss_compiled,native_gpu,1,rocm_metric_loss_compiled,success,hip_runtime,"ROCm metric-loss artifact runs wasserstein / cosine_embedding / contrastive / triplet / InfoNCE / NT-Xent / seq2seq losses with generated ROCm reductions and exp/log; label, mask, sort, and compact matrix structure remain on host. f32, matches tessera.losses." rocm,rocm_moe_compiled,native_gpu,1,rocm_moe_compiled,success,hip_runtime,"ROCm moe-compute lane runs the routed per-token expert GEMVs (top-1) on the COMPILER-GENERATED gfx1151 kernel (generate-rocm-moe-kernel: one thread per (token, out-col); routing resolved on host) HIP-launched. f32, matches numpy." -rocm,rocm_moe_transport_compiled,native_gpu,1,rocm_moe_transport_compiled,success,hip_runtime,"ROCm DK3 MoE transport: moe_dispatch runs NATIVELY on the gfx1151 device gather kernel (token_of_slot = sort_perm//top_k row gather) and moe_combine on the device scatter (add) kernel (host pre-scales each packed row by its route weight, then atomic scatter-add to token order) — both report native_gpu vs the stdlib DispatchPlan oracle. grouped_swiglu (the expert GEMM) stays on the oracle and reports reference_cpu via the per-op execution_kind override — a native f32-exact grouped GEMM is a separate follow-up (WMMA is f16). Off-box the transport ops fall back to the oracle + reference_cpu." +rocm,rocm_moe_transport_compiled,native_gpu,1,rocm_moe_transport_compiled,success,hip_runtime,"ROCm DK3 MoE transport + expert GEMM run NATIVELY on gfx1151: moe_dispatch on the device gather kernel (token_of_slot = sort_perm//top_k row gather), moe_combine on the device scatter (add) kernel (host pre-scales each packed row by its route weight, then atomic scatter-add to token order), and grouped_swiglu's three expert GEMMs on the f32 GEMM device kernel (generate-rocm-gemm-f32-kernel; silu*mul host-side). All three report native_gpu vs the stdlib DispatchPlan oracle (f32 vs the f64 oracle for combine/swiglu). Off-box they fall back to the oracle + reference_cpu." rocm,rocm_muon_compiled,native_gpu,1,rocm_muon_compiled,success,hip_runtime,"ROCm Muon lane orthogonalizes the momentum matrix via the gfx1151 SVD kernel (U·Vh polar factor); the small U@Vh + momentum/sgd run on host. <2-D params normalize. f32, matches optim.muon." rocm,rocm_norm_compiled,native_gpu,1,rocm_norm_compiled,success,hip_runtime,"ROCm norm artifact runs the COMPILER-GENERATED RDNA row-reduction kernel (unweighted rmsnorm / layer_norm over the last axis, one workgroup per row, LDS tree-reduce of Σx and Σx²): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Handles tessera.rmsnorm(_safe) + tessera.layer_norm by op name." rocm,rocm_normcompose_compiled,native_gpu,1,rocm_normcompose_compiled,success,hip_runtime,"ROCm group/instance/weight-norm lane composed on the gfx1151 layer_norm (row mean/var) + reduce (sum-of-squares) kernels; host does the reshape / per-axis divide. f32, matches nn.functional." diff --git a/docs/audit/generated/runtime_execution_matrix.md b/docs/audit/generated/runtime_execution_matrix.md index 18cb0034d..313c9ba07 100644 --- a/docs/audit/generated/runtime_execution_matrix.md +++ b/docs/audit/generated/runtime_execution_matrix.md @@ -72,7 +72,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `rocm` | `rocm_matmul_family_compiled` | `rocm_matmul_family_compiled` | `native_gpu` | `hip_runtime` | ROCm matmul-family artifact runs the COMPILER-GENERATED WMMA GEMM kernel (the rocm_compiled spine) reshaped/batched/split in the runtime — batched_gemm, linear_general, qkv_projection, factorized_matmul (GPU matmul + exact host SVD-truncate), and single-contraction einsum. f16/bf16, f32 accumulate. | | `rocm` | `rocm_metric_loss_compiled` | `rocm_metric_loss_compiled` | `native_gpu` | `hip_runtime` | ROCm metric-loss artifact runs wasserstein / cosine_embedding / contrastive / triplet / InfoNCE / NT-Xent / seq2seq losses with generated ROCm reductions and exp/log; label, mask, sort, and compact matrix structure remain on host. f32, matches tessera.losses. | | `rocm` | `rocm_moe_compiled` | `rocm_moe_compiled` | `native_gpu` | `hip_runtime` | ROCm moe-compute lane runs the routed per-token expert GEMVs (top-1) on the COMPILER-GENERATED gfx1151 kernel (generate-rocm-moe-kernel: one thread per (token, out-col); routing resolved on host) HIP-launched. f32, matches numpy. | -| `rocm` | `rocm_moe_transport_compiled` | `rocm_moe_transport_compiled` | `native_gpu` | `hip_runtime` | ROCm DK3 MoE transport: moe_dispatch runs NATIVELY on the gfx1151 device gather kernel (token_of_slot = sort_perm//top_k row gather) and moe_combine on the device scatter (add) kernel (host pre-scales each packed row by its route weight, then atomic scatter-add to token order) — both report native_gpu vs the stdlib DispatchPlan oracle. grouped_swiglu (the expert GEMM) stays on the oracle and reports reference_cpu via the per-op execution_kind override — a native f32-exact grouped GEMM is a separate follow-up (WMMA is f16). Off-box the transport ops fall back to the oracle + reference_cpu. | +| `rocm` | `rocm_moe_transport_compiled` | `rocm_moe_transport_compiled` | `native_gpu` | `hip_runtime` | ROCm DK3 MoE transport + expert GEMM run NATIVELY on gfx1151: moe_dispatch on the device gather kernel (token_of_slot = sort_perm//top_k row gather), moe_combine on the device scatter (add) kernel (host pre-scales each packed row by its route weight, then atomic scatter-add to token order), and grouped_swiglu's three expert GEMMs on the f32 GEMM device kernel (generate-rocm-gemm-f32-kernel; silu*mul host-side). All three report native_gpu vs the stdlib DispatchPlan oracle (f32 vs the f64 oracle for combine/swiglu). Off-box they fall back to the oracle + reference_cpu. | | `rocm` | `rocm_muon_compiled` | `rocm_muon_compiled` | `native_gpu` | `hip_runtime` | ROCm Muon lane orthogonalizes the momentum matrix via the gfx1151 SVD kernel (U·Vh polar factor); the small U@Vh + momentum/sgd run on host. <2-D params normalize. f32, matches optim.muon. | | `rocm` | `rocm_norm_compiled` | `rocm_norm_compiled` | `native_gpu` | `hip_runtime` | ROCm norm artifact runs the COMPILER-GENERATED RDNA row-reduction kernel (unweighted rmsnorm / layer_norm over the last axis, one workgroup per row, LDS tree-reduce of Σx and Σx²): tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. Handles tessera.rmsnorm(_safe) + tessera.layer_norm by op name. | | `rocm` | `rocm_normcompose_compiled` | `rocm_normcompose_compiled` | `native_gpu` | `hip_runtime` | ROCm group/instance/weight-norm lane composed on the gfx1151 layer_norm (row mean/var) + reduce (sum-of-squares) kernels; host does the reshape / per-axis divide. f32, matches nn.functional. | diff --git a/docs/rocm_mfma_kernel_inventory.md b/docs/rocm_mfma_kernel_inventory.md index 406bc12ca..fa2cf8a6b 100644 --- a/docs/rocm_mfma_kernel_inventory.md +++ b/docs/rocm_mfma_kernel_inventory.md @@ -493,10 +493,9 @@ promotes them to `compileable`. See §9 for the concrete done / open / blocked s execute-compare vs `KVCacheHandle`); a single fused gather→attention paged kernel is the remaining step - `grad_clip_norm` (§5.5) — global-norm + scale; single-node -- `grouped_swiglu` native f32 expert GEMM — `moe_dispatch` (device gather) and - `moe_combine` (device scatter-add) now execute natively (`rocm_moe_transport_compiled` - reports `native_gpu`); the expert GEMM stays on the oracle pending an f32-exact - device GEMM (WMMA is f16) +- (moe transport is fully native now — `moe_dispatch` device gather, `moe_combine` + device scatter-add, `grouped_swiglu` on the new `generate-rocm-gemm-f32-kernel` + f32 GEMM; `rocm_moe_transport_compiled` reports `native_gpu`) - Plain recurrent cells `lstm_cell` / `gru_cell` / `simple_rnn_cell` (§5.9) — no dedicated lane (selective-SSM + deltanet are done) - **Perf ladders / MFU sign-off** beyond `matmul` — every `compiled` lane is diff --git a/python/tessera/compiler/execution_matrix.py b/python/tessera/compiler/execution_matrix.py index 6b149afcb..4d4271306 100644 --- a/python/tessera/compiler/execution_matrix.py +++ b/python/tessera/compiler/execution_matrix.py @@ -2159,16 +2159,16 @@ class ExecutionRow: target="rocm", compiler_path="rocm_moe_transport_compiled", execution_kind="native_gpu", executable=True, executor_id="rocm_moe_transport_compiled", runtime_status="success", - reason="ROCm DK3 MoE transport: moe_dispatch runs NATIVELY on the " - "gfx1151 device gather kernel (token_of_slot = sort_perm//top_k " - "row gather) and moe_combine on the device scatter (add) kernel " - "(host pre-scales each packed row by its route weight, then " - "atomic scatter-add to token order) — both report native_gpu vs " - "the stdlib DispatchPlan oracle. grouped_swiglu (the expert GEMM) " - "stays on the oracle and reports reference_cpu via the per-op " - "execution_kind override — a native f32-exact grouped GEMM is a " - "separate follow-up (WMMA is f16). Off-box the transport ops " - "fall back to the oracle + reference_cpu.", + reason="ROCm DK3 MoE transport + expert GEMM run NATIVELY on gfx1151: " + "moe_dispatch on the device gather kernel (token_of_slot = " + "sort_perm//top_k row gather), moe_combine on the device scatter " + "(add) kernel (host pre-scales each packed row by its route " + "weight, then atomic scatter-add to token order), and " + "grouped_swiglu's three expert GEMMs on the f32 GEMM device " + "kernel (generate-rocm-gemm-f32-kernel; silu*mul host-side). All " + "three report native_gpu vs the stdlib DispatchPlan oracle (f32 " + "vs the f64 oracle for combine/swiglu). Off-box they fall back to " + "the oracle + reference_cpu.", execution_mode="hip_runtime"), ("rocm", "rocm_normcompose_compiled"): ExecutionRow( target="rocm", compiler_path="rocm_normcompose_compiled", diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index cc4f81bbf..dfc44fca6 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -1522,13 +1522,16 @@ def _execute_rocm_moe_transport(artifact: RuntimeArtifact, args: Any) -> Any: return (moe.combine(partials, plan), "reference_cpu") if op_name == "tessera.grouped_swiglu": try: - out = moe.grouped_swiglu( - values["x_packed"], values["w_gate"], values["w_up"], - values["w_down"], values["group_sizes"]) + gvals = (values["x_packed"], values["w_gate"], values["w_up"], + values["w_down"], values["group_sizes"]) except KeyError as exc: raise ValueError(f"missing grouped_swiglu argument {exc.args[0]!r}") from exc - # Expert GEMM: no f32-exact device GEMM yet (WMMA is f16) — stays oracle. - return (out, "reference_cpu") + # Expert GEMM runs natively on the f32 GEMM device kernel (three GEMMs + + # host silu*mul per group); falls back to the f64 oracle off-box. + try: + return (_moe_grouped_swiglu_native(*gvals, np), "native_gpu") + except _RocmCompiledUnavailable: + return (moe.grouped_swiglu(*gvals), "reference_cpu") raise ValueError( "rocm_moe_transport_compiled executor handles tessera.moe_dispatch, " "tessera.moe_combine, or tessera.grouped_swiglu; " @@ -9884,6 +9887,109 @@ def _execute_rocm_compiled_scatter(artifact: RuntimeArtifact, args: Any) -> Any: return _execute_scatter(artifact, args, _rocm_scatter, "rocm_scatter_compiled") +# ───────────────────────────────────────────────────────────────────────────── +# f32 GEMM device kernel (generate-rocm-gemm-f32-kernel) — the f32-exact matmul +# RDNA WMMA (f16/bf16 only) can't provide. One thread per output element, scalar +# f32 k-loop (correctness-first, no LDS tiling). Backs the native grouped-SwiGLU +# expert GEMM (MoE). f32; matches the numpy f32 matmul within f32 tolerance. +# ───────────────────────────────────────────────────────────────────────────── +_rocm_gemm_f32_hsaco_cache: dict[tuple[str], bytes] = {} + + +def _build_compiled_gemm_f32_hsaco() -> bytes: + chip = _rocm_chip() + directive = ('module {\n "tessera_rocm.gemm_f32"() {name = "g"} ' + ': () -> ()\n}\n') + return _build_rocm_elementwise_hsaco( + "generate-rocm-gemm-f32-kernel", directive, _rocm_gemm_f32_hsaco_cache, + (chip,)) + + +def _rocm_f32_gemm(a: Any, b: Any, np: Any) -> Any: + """C = A @ B on gfx1151 via the compiler-generated f32 GEMM kernel (one + thread per output element). A (M,K), B (K,N) f32 -> C (M,N) f32.""" + a = np.ascontiguousarray(a, np.float32) + b = np.ascontiguousarray(b, np.float32) + if a.ndim != 2 or b.ndim != 2 or int(a.shape[1]) != int(b.shape[0]): + raise ValueError(f"f32 gemm needs (M,K)@(K,N); got {a.shape} @ {b.shape}") + M, K = int(a.shape[0]), int(a.shape[1]) + N = int(b.shape[1]) + out = np.zeros((M, N), np.float32) + if M == 0 or N == 0 or K == 0: + return out + hsaco = _build_compiled_gemm_f32_hsaco() + hip = _load_hip_for_launch() + if hip is None: + raise _RocmCompiledUnavailable("libamdhip64.so not loadable") + if hip.hipInit(0) != 0: + raise _RocmCompiledUnavailable("rocm gemm_f32: hipInit failed") + mod = ctypes.c_void_p() + if hip.hipModuleLoadData(ctypes.byref(mod), hsaco) != 0: + raise _RocmCompiledUnavailable("rocm gemm_f32: no usable AMD GPU") + fn = ctypes.c_void_p() + if hip.hipModuleGetFunction(ctypes.byref(fn), mod, b"g") != 0: + raise RuntimeError("rocm gemm_f32: kernel symbol 'g' not found") + cv = ctypes.c_void_p + n_a, n_b, n_c = M * K, K * N, M * N + d_a, d_b, d_c = cv(), cv(), cv() + if (hip.hipMalloc(ctypes.byref(d_a), 4 * n_a) != 0 + or hip.hipMalloc(ctypes.byref(d_b), 4 * n_b) != 0 + or hip.hipMalloc(ctypes.byref(d_c), 4 * n_c) != 0): + raise RuntimeError("rocm gemm_f32: hipMalloc failed") + hip.hipMemcpy(d_a, a.ctypes.data_as(cv), 4 * n_a, 1) + hip.hipMemcpy(d_b, b.ctypes.data_as(cv), 4 * n_b, 1) + + def _mr(p, size): + return [cv(p.value), cv(p.value), ctypes.c_int64(0), + ctypes.c_int64(size), ctypes.c_int64(1)] + + launch_args = (_mr(d_a, n_a) + _mr(d_b, n_b) + _mr(d_c, n_c) + + [ctypes.c_int64(M), ctypes.c_int64(N), ctypes.c_int64(K)]) + arr = (cv * len(launch_args))() + for i, val in enumerate(launch_args): + arr[i] = ctypes.cast(ctypes.byref(val), cv) + gx = (n_c + _GRID_BLOCKDIM - 1) // _GRID_BLOCKDIM + rc = hip.hipModuleLaunchKernel(fn, max(gx, 1), 1, 1, _GRID_BLOCKDIM, 1, 1, + 0, None, arr, None) + if rc != 0: + for d in (d_a, d_b, d_c): + hip.hipFree(d) + raise RuntimeError(f"rocm gemm_f32: kernel launch failed rc={rc}") + hip.hipDeviceSynchronize() + hip.hipMemcpy(out.ctypes.data_as(cv), d_c, 4 * n_c, 2) + for d in (d_a, d_b, d_c): + hip.hipFree(d) + return out + + +def _moe_grouped_swiglu_native(x_packed: Any, w_gate: Any, w_up: Any, + w_down: Any, group_sizes: Any, np: Any) -> Any: + """Grouped SwiGLU over contiguous per-expert token groups, run on the f32 + GEMM device kernel: for each group e, g = x@Wg[e], u = x@Wu[e], + h = silu(g)*u, out = h@Wd[e]. The three FLOP-heavy GEMMs run on gfx1151; the + silu*mul elementwise middle is host-side. f32; matches stdlib.moe.grouped_swiglu + within f32 tolerance (the reference accumulates in f64).""" + xp = np.ascontiguousarray(_as_numpy(x_packed), np.float32) + wg = np.ascontiguousarray(_as_numpy(w_gate), np.float32) # (E, H, F) + wu = np.ascontiguousarray(_as_numpy(w_up), np.float32) # (E, H, F) + wd = np.ascontiguousarray(_as_numpy(w_down), np.float32) # (E, F, H) + gs = np.asarray(_as_numpy(group_sizes), np.int64).reshape(-1) + S = int(xp.shape[0]) + h_out = int(wd.shape[2]) + out = np.zeros((S, h_out), np.float32) + off = 0 + for e in range(int(wg.shape[0])): + n = int(gs[e]) + if n: + xe = xp[off:off + n] # (n, H) + g = _rocm_f32_gemm(xe, wg[e], np) # (n, F) + u = _rocm_f32_gemm(xe, wu[e], np) # (n, F) + h = (g * (1.0 / (1.0 + np.exp(-g))) * u).astype(np.float32) + out[off:off + n] = _rocm_f32_gemm(h, wd[e], np) # (n, h_out) + off += n + return out + + # ───────────────────────────────────────────────────────────────────────────── # KV-cache paged-movement lane (§5.6) — the append/read/prune core executes on # gfx1151 by COMPOSING the existing device gather/scatter kernels + host page- diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td index 88b3b34e5..5444fd388 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td +++ b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/IR/TesseraROCMOps.td @@ -543,6 +543,19 @@ def ROCM_MoeKernelOp : TesseraROCM_Op<"moe"> { let assemblyFormat = "attr-dict"; } +def ROCM_GemmF32KernelOp : TesseraROCM_Op<"gemm_f32"> { + let summary = "Plain single-precision GEMM gpu kernel"; + let description = [{ + `generate-rocm-gemm-f32-kernel` expands into a plain f32 GEMM + C[m,n] = Σ_k A[m,k]·B[k,n] (A [M,K], B [K,N], C [M,N]), one thread per output + element, scalar f32 k-loop (f32 accumulate). RDNA WMMA is f16/bf16 only, so + this is the f32-exact VALU/FMA fallback for the expert GEMM (grouped SwiGLU); + correctness-first, no LDS tiling. M/N/K are runtime index args. + }]; + let arguments = (ins StrAttr:$name); + let assemblyFormat = "attr-dict"; +} + def ROCM_PredicateKernelOp : TesseraROCM_Op<"predicate"> { let summary = "Unary predicate (isnan/isinf/isfinite) gpu kernel"; let description = [{ diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h index 1778bbba0..829ef43f7 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h +++ b/src/compiler/codegen/Tessera_ROCM_Backend/include/TesseraROCM/Passes.h @@ -36,6 +36,7 @@ std::unique_ptr createGenerateROCMSvdKernelPass(); std::unique_ptr createGenerateROCMOptimizerKernelPass(); std::unique_ptr createGenerateROCMPredicateKernelPass(); std::unique_ptr createGenerateROCMMoeKernelPass(); +std::unique_ptr createGenerateROCMGemmF32KernelPass(); std::unique_ptr createGenerateROCMAlibiKernelPass(); std::unique_ptr createGenerateROCMDeltaNetKernelPass(); std::unique_ptr createGenerateROCMRopeKernelPass(); diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt index 36bf4e6a2..de7c77dc9 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/CMakeLists.txt @@ -27,6 +27,7 @@ add_mlir_library(TesseraROCMConversion GenerateROCMOptimizerKernel.cpp GenerateROCMPredicateKernel.cpp GenerateROCMMoeKernel.cpp + GenerateROCMGemmF32Kernel.cpp GenerateROCMAlibiKernel.cpp GenerateROCMDeltaNetKernel.cpp GenerateROCMRopeKernel.cpp diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMGemmF32Kernel.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMGemmF32Kernel.cpp new file mode 100644 index 000000000..487e58b96 --- /dev/null +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateROCMGemmF32Kernel.cpp @@ -0,0 +1,126 @@ +//===- GenerateROCMGemmF32Kernel.cpp - plain f32 GEMM gpu kernel ---------===// +// +// Expands `tessera_rocm.gemm_f32` into a plain single-precision GEMM, +// one thread per output element: +// +// C[m, n] = Σ_k A[m, k] · B[k, n] (A: [M,K], B: [K,N], C: [M,N]) +// +// RDNA WMMA is f16/bf16 only, so the f32-exact expert GEMM (grouped SwiGLU) can +// not ride the WMMA lane without precision loss. This is the f32 VALU/FMA +// fallback: correctness-first (scalar k-loop, f32 accumulate, no LDS tiling — +// the tiled/blocked perf ladder is a follow-up), matching the numpy f32 matmul. +// M/N/K are runtime index args; the grid folds M*N into a 1-D launch. +//===----------------------------------------------------------------------===// + +#include "TesseraROCM/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" + +using namespace mlir; + +namespace { + +static constexpr int64_t BD = 256; + +void emitGemmF32Body(OpBuilder &b, Location loc, gpu::GPUFuncOp f) { + Type f32 = b.getF32Type(); + auto slt = arith::CmpIPredicate::slt; + b.setInsertionPointToStart(&f.getBody().front()); + Value A = f.getArgument(0), B = f.getArgument(1), C = f.getArgument(2); + Value M = f.getArgument(3), N = f.getArgument(4), K = f.getArgument(5); + + Value bid = b.create(loc, gpu::Dimension::x); + Value tid = b.create(loc, gpu::Dimension::x); + Value cBD = b.create(loc, BD); + Value gid = b.create( + loc, b.create(loc, bid, cBD), tid); + Value total = b.create(loc, M, N); + Value inb = b.create(loc, slt, gid, total); + auto guard = b.create(loc, inb, /*withElse=*/false); + b.setInsertionPointToStart(guard.thenBlock()); + + Value c0 = b.create(loc, 0); + Value c1 = b.create(loc, 1); + Value zero = b.create(loc, f32, b.getF32FloatAttr(0.0f)); + Value m = b.create(loc, gid, N); // row + Value n = b.create(loc, gid, N); // col + Value abase = b.create(loc, m, K); // m*K + // acc = Σ_k A[m*K + k] · B[k*N + n] + auto kl = b.create(loc, c0, K, c1, ValueRange{zero}); + { + OpBuilder::InsertionGuard g(b); + b.setInsertionPointToStart(kl.getBody()); + Value k = kl.getInductionVar(); + Value av = b.create( + loc, A, ValueRange{b.create(loc, abase, k)}); + Value boff = b.create( + loc, b.create(loc, k, N), n); // k*N + n + Value bv = b.create(loc, B, ValueRange{boff}); + Value acc = b.create(loc, kl.getRegionIterArgs()[0], + b.create(loc, av, bv)); + b.create(loc, ValueRange{acc}); + } + b.create(loc, kl.getResult(0), C, ValueRange{gid}); + b.setInsertionPointToEnd(&f.getBody().front()); + b.create(loc); +} + +struct GenerateROCMGemmF32KernelPass + : PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GenerateROCMGemmF32KernelPass) + + StringRef getArgument() const final { return "generate-rocm-gemm-f32-kernel"; } + StringRef getDescription() const final { + return "Expand a tessera_rocm.gemm_f32 directive into a plain f32 GEMM " + "kernel (C=A@B, one thread per output element, scalar f32 k-loop)"; + } + void getDependentDialects(DialectRegistry ®istry) const final { + registry.insert(); + } + + void runOnOperation() override { + ModuleOp module = getOperation(); + SmallVector directives; + module.walk([&](Operation *op) { + if (op->getName().getStringRef() == "tessera_rocm.gemm_f32") + directives.push_back(op); + }); + for (Operation *op : directives) { + auto nameAttr = op->getAttrOfType("name"); + if (!nameAttr) { + op->emitError("tessera_rocm.gemm_f32 missing name"); + return signalPassFailure(); + } + OpBuilder b(module.getBodyRegion()); + b.setInsertionPointToEnd(module.getBody()); + Location loc = op->getLoc(); + std::string kname = nameAttr.getValue().str(); + Type f32 = b.getF32Type(); + Type idxTy = b.getIndexType(); + auto memF32 = MemRefType::get({ShapedType::kDynamic}, f32); + auto fnTy = b.getFunctionType( + {memF32, memF32, memF32, idxTy, idxTy, idxTy}, {}); + auto gpuMod = b.create(loc, kname + "_mod"); + b.setInsertionPointToStart(&gpuMod.getBodyRegion().front()); + auto gpuFunc = b.create(loc, kname, fnTy); + gpuFunc->setAttr(gpu::GPUDialect::getKernelFuncAttrName(), b.getUnitAttr()); + OpBuilder body(gpuFunc.getContext()); + emitGemmF32Body(body, loc, gpuFunc); + op->erase(); + } + } +}; + +} // namespace + +std::unique_ptr +mlir::tessera_rocm::createGenerateROCMGemmF32KernelPass() { + return std::make_unique(); +} diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp index f29ecfb52..998427dec 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/Passes.cpp @@ -56,6 +56,7 @@ void registerTesseraROCMPasses() { registerPass([]() { return createGenerateROCMOptimizerKernelPass(); }); registerPass([]() { return createGenerateROCMPredicateKernelPass(); }); registerPass([]() { return createGenerateROCMMoeKernelPass(); }); + registerPass([]() { return createGenerateROCMGemmF32KernelPass(); }); registerPass([]() { return createGenerateROCMAlibiKernelPass(); }); registerPass([]() { return createGenerateROCMDeltaNetKernelPass(); }); registerPass([]() { return createGenerateROCMRopeKernelPass(); }); diff --git a/tests/unit/test_rocm_moe_transport_compiled.py b/tests/unit/test_rocm_moe_transport_compiled.py index 81ad271f2..40667f296 100644 --- a/tests/unit/test_rocm_moe_transport_compiled.py +++ b/tests/unit/test_rocm_moe_transport_compiled.py @@ -107,12 +107,15 @@ def test_rocm_grouped_swiglu_runtime_matches_grouped_gemm_oracle(): ) assert res["ok"] - assert res["execution_kind"] == "reference_cpu" + # grouped_swiglu now runs the three expert GEMMs on the f32 device GEMM + # kernel (native on-box), silu*mul host-side; f32 vs the f64 oracle. + assert res["execution_kind"] == ("native_gpu" if _expect_native() + else "reference_cpu") np.testing.assert_allclose( res["output"], moe.grouped_swiglu(x_packed, w_gate, w_up, w_down, group_sizes), - rtol=0, - atol=0, + rtol=1e-4, + atol=1e-5, ) From 7a0523b550d694ae144a8a743b1ac70e0359377a Mon Sep 17 00:00:00 2001 From: gstoner Date: Fri, 10 Jul 2026 12:26:26 -0600 Subject: [PATCH 4/5] fix: lstm_cell packed-state contract + mypy ratchet (CI review + lint) Review feedback (P2): the structured-compute lane dispatched tessera.lstm_cell to F.lstm_cell, which has a DIFFERENT contract than the canonical tessera.ops.lstm_cell (the public op uses (4H,In)/(4H,H) weights via W_ih.T/W_hh.T and returns a packed concat([h_t, c_t]) for lstm_state_h/c extraction; F.lstm_cell took transposed weights and returned a tuple). Normal LSTM artifacts would fail shape validation or produce output the state extractors can't consume on the newly-advertised x86/ROCm/Apple structured lanes. - Dispatch tessera.lstm_cell -> ops.lstm_cell (canonical packed contract). - Remove the misleading F.lstm_cell duplicate + its __all__ entry. - Update the x86/ROCm structured fixtures to the packed contract ((4H,In) weights). Also fix the mypy ratchet failure: assert bias_c is not None in the FA backward attn_bias device-copy (mypy could not narrow the has_bias flag). Co-Authored-By: Claude Opus 4.8 (1M context) --- python/tessera/nn/functional.py | 23 ------------------- python/tessera/runtime.py | 6 ++++- .../test_rocm_structured_compute_compiled.py | 19 +++++++-------- .../test_x86_structured_compute_compiled.py | 16 +++++++------ 4 files changed, 24 insertions(+), 40 deletions(-) diff --git a/python/tessera/nn/functional.py b/python/tessera/nn/functional.py index 0ae0d7cb7..a0b99b9fb 100644 --- a/python/tessera/nn/functional.py +++ b/python/tessera/nn/functional.py @@ -309,28 +309,6 @@ def gru_cell(x, h, W_ih, W_hh, b_ih=None, b_hh=None): return (1.0 - z) * n + z * h_arr -def lstm_cell(x, h, c, W_ih, W_hh, b_ih=None, b_hh=None): - """LSTM cell with gate order i, f, g, o (PyTorch/cuDNN convention). - - ``x`` (…, in), ``h``/``c`` (…, hidden); ``W_ih`` (in, 4*hidden), ``W_hh`` - (hidden, 4*hidden). Returns ``(h_new, c_new)``. - """ - x_arr, h_arr, c_arr = _asarray(x), _asarray(h), _asarray(c) - gates = x_arr @ _asarray(W_ih) + h_arr @ _asarray(W_hh) - if b_ih is not None: - gates = gates + _asarray(b_ih) - if b_hh is not None: - gates = gates + _asarray(b_hh) - i, f, g, o = np.split(gates, 4, axis=-1) - i = 1.0 / (1.0 + np.exp(-i)) - f = 1.0 / (1.0 + np.exp(-f)) - g = np.tanh(g) - o = 1.0 / (1.0 + np.exp(-o)) - c_new = f * c_arr + i * g - h_new = o * np.tanh(c_new) - return h_new, c_new - - def bidirectional_scan(fn, init_fwd, init_bwd, xs): xs_arr = _asarray(xs) fwd_states = [] @@ -748,7 +726,6 @@ def proj_heads(t, W, heads): "group_norm", "gqa_attention", "gru_cell", - "lstm_cell", "instance_norm", "mask_token_block", "max_pool", diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index dfc44fca6..c52eb865f 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -3372,6 +3372,7 @@ def _execute_rocm_compiled_flash_attn_bwd(artifact: RuntimeArtifact, hip.hipMemcpy(bufs["O"], o_f32.ctypes.data_as(cv), 4 * nq, 1) n_bias = bh * sq * sk if has_bias: + assert bias_c is not None # set above whenever has_bias d = cv() if hip.hipMalloc(ctypes.byref(d), max(4 * n_bias, 4)) != 0: for dd in bufs.values(): @@ -13339,7 +13340,10 @@ def _execute_structured_compute_composite( if op_name == "tessera.gru_cell": return F.gru_cell(*operands, **kwargs) if op_name == "tessera.lstm_cell": - return F.lstm_cell(*operands, **kwargs) + # Canonical op contract (tessera.ops.lstm_cell): (4H,In)/(4H,H) weights + # via W_ih.T/W_hh.T, packed concat([h_t, c_t]) output — NOT F.lstm_cell + # (which takes transposed weights and returns a tuple). + return ops.lstm_cell(*operands, **kwargs) if op_name == "tessera.cross_attention": return ops.cross_attention(*operands, **kwargs) if op_name == "tessera.perceiver_resampler": diff --git a/tests/unit/test_rocm_structured_compute_compiled.py b/tests/unit/test_rocm_structured_compute_compiled.py index e16ba8911..39fb6c2be 100644 --- a/tests/unit/test_rocm_structured_compute_compiled.py +++ b/tests/unit/test_rocm_structured_compute_compiled.py @@ -100,16 +100,17 @@ def test_rocm_structured_model_recurrent_and_stencil_match_reference_on_gpu(): atol=1e-6, ) - # LSTM cell (gate order i,f,g,o): W_ih/W_hh are (in,4H)/(H,4H); returns - # (h_new, c_new). Runs through the same structured-compute lane. + # LSTM cell — canonical tessera.ops.lstm_cell contract: (4H,In)/(4H,H) + # weights (applied via W.T), returns the packed concat([h_t, c_t]) (B, 2H). cprev = rng.standard_normal((2, 5)).astype(np.float32) - Wih4 = rng.standard_normal((3, 20)).astype(np.float32) - Whh4 = rng.standard_normal((5, 20)).astype(np.float32) - lh, lc = _launch(rt, "tessera.lstm_cell", ("x", "h", "c", "Wih", "Whh"), - (xt, h, cprev, Wih4, Whh4)) - rh, rc = F.lstm_cell(xt, h, cprev, Wih4, Whh4) - np.testing.assert_allclose(lh, rh, atol=1e-6) - np.testing.assert_allclose(lc, rc, atol=1e-6) + Wih4 = rng.standard_normal((20, 3)).astype(np.float32) + Whh4 = rng.standard_normal((20, 5)).astype(np.float32) + np.testing.assert_allclose( + _launch(rt, "tessera.lstm_cell", ("x", "h", "c", "Wih", "Whh"), + (xt, h, cprev, Wih4, Whh4)), + ops.lstm_cell(xt, h, cprev, Wih4, Whh4), + atol=1e-6, + ) a = rng.standard_normal((2, 4)).astype(np.float32) weight = rng.standard_normal((4, 6)).astype(np.float32) diff --git a/tests/unit/test_x86_structured_compute_compiled.py b/tests/unit/test_x86_structured_compute_compiled.py index e15b159da..13d31f207 100644 --- a/tests/unit/test_x86_structured_compute_compiled.py +++ b/tests/unit/test_x86_structured_compute_compiled.py @@ -97,14 +97,16 @@ def test_x86_structured_model_recurrent_and_stencil_match_reference(): atol=1e-6, ) + # Canonical tessera.ops.lstm_cell: (4H,In)/(4H,H) weights, packed output. cprev = rng.standard_normal((2, 5)).astype(np.float32) - Wih4 = rng.standard_normal((3, 20)).astype(np.float32) - Whh4 = rng.standard_normal((5, 20)).astype(np.float32) - lh, lc = _launch("tessera.lstm_cell", ("x", "h", "c", "Wih", "Whh"), - (xt, h, cprev, Wih4, Whh4)) - rh, rc = F.lstm_cell(xt, h, cprev, Wih4, Whh4) - np.testing.assert_allclose(lh, rh, atol=1e-6) - np.testing.assert_allclose(lc, rc, atol=1e-6) + Wih4 = rng.standard_normal((20, 3)).astype(np.float32) + Whh4 = rng.standard_normal((20, 5)).astype(np.float32) + np.testing.assert_allclose( + _launch("tessera.lstm_cell", ("x", "h", "c", "Wih", "Whh"), + (xt, h, cprev, Wih4, Whh4)), + ops.lstm_cell(xt, h, cprev, Wih4, Whh4), + atol=1e-6, + ) a = rng.standard_normal((2, 4)).astype(np.float32) weight = rng.standard_normal((4, 6)).astype(np.float32) From fedc13486bccf267df3c3be4aa22bac362c11a78 Mon Sep 17 00:00:00 2001 From: gstoner Date: Fri, 10 Jul 2026 12:51:50 -0600 Subject: [PATCH 5/5] rocm/gfx1151: FA-2 backward sliding-window + logit-softcap variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Complete the flash_attn backward variant surface — the runtime lane now covers the full set of forward variants (MHA + GQA/MQA + attn_bias + sliding-window + logit-softcap), all execute-compare verified vs autodiff vjp_flash_attn on gfx1151. C++ (GenerateWMMAFlashAttnBwdKernel.cpp): - ScoreCtx gains W (window width) + cap (soft-cap). recomputeScoreTile applies the soft-cap S=cap*tanh(sRaw/cap) before the softmax and scales dS by the chain-rule factor 1-tanh^2(sRaw/cap); the window mask (implicitly causal, keys older than W masked) is applied in both emitPre (for L) and recomputeScoreTile (for P/dS). emitDkDv/emitDq use useCausal for the causal tile-skip bounds. - Trailing runtime args appended in a FIXED order: gqa(heads,kv_ratio) | window(W:index) | softcap(cap:f32) | attn_bias([bh,Sq,Sk] f32, LAST), driven by the sliding_window / logit_softcap / attn_bias BoolAttrs. Runtime: thread window / logit_softcap kwargs (hsaco flags + trailing launch args in the same order); the forward-O recompute uses the same window/softcap so O matches the backward's score recompute. Verified: 22/22 gfx1151 tests (incl. window, softcap, and window+softcap+bias compose). mypy clean, drift gate in sync, 144-test regression sweep green. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/audit/generated/rocm_target_map.csv | 2 +- .../generated/runtime_execution_matrix.csv | 2 +- .../generated/runtime_execution_matrix.md | 4 +- docs/rocm_mfma_kernel_inventory.md | 18 ++- python/tessera/compiler/backend_manifest.py | 5 +- python/tessera/compiler/execution_matrix.py | 10 +- python/tessera/runtime.py | 94 +++++++---- .../GenerateWMMAFlashAttnBwdKernel.cpp | 149 +++++++++++++----- .../test_rocm_flash_attn_bwd_runtime_lane.py | 111 +++++++++++-- 9 files changed, 293 insertions(+), 102 deletions(-) diff --git a/docs/audit/generated/rocm_target_map.csv b/docs/audit/generated/rocm_target_map.csv index af318ee86..11ee65d0a 100644 --- a/docs/audit/generated/rocm_target_map.csv +++ b/docs/audit/generated/rocm_target_map.csv @@ -6,7 +6,7 @@ softmax,activation,compiled,"fp32,fp16,bf16",hipcc≥7.2.4,,,,"Row-wise stable s softmax_safe,activation,compiled,fp32,hipcc≥7.2.4,,,,Stable reduction softmax_safe — max-shifted reduce (max/sum) + unary exp/log lane. Executes via runtime.launch() (rocm_stable_reduce_compiled). attn_sliding_window,attention,compiled,"fp16,bf16",hipcc≥7.2.4,,,,"Mistral sliding-window via the flash_attn WMMA kernel (sliding_window attr, causal band of width W; KV-tile skip). Executes via runtime.launch() (window kwarg)." deepseek_sparse_attention,attention,compiled,fp32,hipcc≥7.2.4,,,,"DeepSeek/NSA composition — sliding + compressed-block branches remain reference compositions while the top-k branch uses the GPU-resident top-k selector plus selected-block sparse-attention kernel when ROCm is available. Executes via runtime.launch() (rocm_sparse_attn_compiled), with exact reference fallback off hardware." -flash_attn,attention,hardware_verified,"fp16,bf16",hipcc≥7.2.4,,0.6500,,"RDNA 3.5 WMMA flash-attention forward executes on the AMD GPU through the shipped libtessera_rocm_flash_attn.so symbols (tessera_rocm_wmma_flash_attn_{f16,bf16}, HIPRTC-compiled for the device arch at load); ROCm 7.2.4. The second op after matmul to run natively on a non-Apple backend. Numerically validated vs a numpy attention reference by the execute_compare_fixture. The FA-2 BACKWARD (dQ/dK/dV) also executes on gfx1151 via the compiler-generated rocm_flash_attn_bwd_compiled lane (generate-wmma-flash-attn-bwd-kernel -> fa_pre/fa_dkdv/fa_dq; MHA + GQA/MQA + additive attn_bias, scale+causal), validated vs autodiff vjp_flash_attn — see the runtime_execution_matrix." +flash_attn,attention,hardware_verified,"fp16,bf16",hipcc≥7.2.4,,0.6500,,"RDNA 3.5 WMMA flash-attention forward executes on the AMD GPU through the shipped libtessera_rocm_flash_attn.so symbols (tessera_rocm_wmma_flash_attn_{f16,bf16}, HIPRTC-compiled for the device arch at load); ROCm 7.2.4. The second op after matmul to run natively on a non-Apple backend. Numerically validated vs a numpy attention reference by the execute_compare_fixture. The FA-2 BACKWARD (dQ/dK/dV) also executes on gfx1151 via the compiler-generated rocm_flash_attn_bwd_compiled lane (generate-wmma-flash-attn-bwd-kernel -> fa_pre/fa_dkdv/fa_dq; MHA + GQA/MQA + additive attn_bias + sliding-window + logit-softcap, scale+causal), validated vs autodiff vjp_flash_attn — see the runtime_execution_matrix." gated_attention,attention,compiled,"fp16,bf16",hipcc≥7.2.4,,,,Softmax attention × a learned gate — the WMMA flash_attn kernel + an elementwise sigmoid-gate multiply. Executes via runtime.launch() (rocm_exotic_attn_compiled). gated_deltanet,attention,compiled,"fp32,fp16,bf16",hipcc≥7.2.4,,,,"Gated/delta linear-attention recurrence as a causal SEQUENTIAL-SCAN kernel (generate-rocm-deltanet-kernel: one workgroup per (b,h), one thread per value-column, LDS state) — the first recurrent compiled ROCm kernel. erase/gate/beta/decay flags. Executes via runtime.launch() (rocm_deltanet_compiled)." gqa_attention,attention,compiled,"fp16,bf16",hipcc≥7.2.4,,,,"GQA/MQA via the flash_attn WMMA kernel (gqa directive attr; fwd+bwd, grouped K/V; runtime detects from operand shapes). Executes on gfx1151 via runtime.launch() (rocm_flash_attn_compiled); no shipped C-ABI symbol." diff --git a/docs/audit/generated/runtime_execution_matrix.csv b/docs/audit/generated/runtime_execution_matrix.csv index e0d1fb563..bb67fa7c4 100644 --- a/docs/audit/generated/runtime_execution_matrix.csv +++ b/docs/audit/generated/runtime_execution_matrix.csv @@ -44,7 +44,7 @@ rocm,rocm_ebm_langevin_compiled,native_gpu,1,rocm_ebm_langevin_compiled,success, rocm,rocm_ebm_loss_compiled,native_gpu,1,rocm_ebm_loss_compiled,success,hip_runtime,"ROCm EBM/diffusion loss artifact runs score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_loss. f32." rocm,rocm_exotic_attn_compiled,native_gpu,1,rocm_exotic_attn_compiled,success,hip_runtime,"ROCm exotic-attention artifact composes the COMPILER-GENERATED WMMA flash_attn kernel with the WMMA GEMM kernel (MLA latent projections) + an elementwise gate — gated_attention, mla_decode, mla_decode_fused — and routes mla_decode_step through the DK1 absorbed-latent decode kernel against stdlib.attention.mla_decode_step. f16/bf16, f32 softmax+accumulate." rocm,rocm_fft_compiled,native_gpu,1,rocm_fft_compiled,success,hip_runtime,"ROCm FFT artifact runs fft / ifft / rfft / irfft over any axis length on the COMPILER-GENERATED one-thread-per-bin DFT kernel (generate-rocm-dft-kernel -> ROCDL, cos/sin twiddles) on gfx1151 + r2c/c2r pack-unpack + plan scale. Direct DFT (radix-2/Bluestein perf is a follow-up). complex64/f32." -rocm,rocm_flash_attn_bwd_compiled,native_gpu,1,rocm_flash_attn_bwd_compiled,success,hip_runtime,"ROCm flash_attn backward artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 backward: tessera-opt expands one tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/fa_dq WMMA kernels serialized to hsaco in-process, then HIP launches them in sequence to produce dQ/dK/dV. O is recomputed via the forward lane (nothing saved from forward). The reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA (gqa dkdv atomic-accumulates dK/dV across the group) + additive attn_bias (S=scale*QK+bias in the recompute), scale + causal, f16/bf16 storage, f32 accumulate." +rocm,rocm_flash_attn_bwd_compiled,native_gpu,1,rocm_flash_attn_bwd_compiled,success,hip_runtime,"ROCm flash_attn backward artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 backward: tessera-opt expands one tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/fa_dq WMMA kernels serialized to hsaco in-process, then HIP launches them in sequence to produce dQ/dK/dV. O is recomputed via the forward lane (nothing saved from forward). The reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA (gqa dkdv atomic-accumulates dK/dV across the group) + additive attn_bias + sliding-window (implicitly causal, masks keys older than W) + Gemma-2 logit-softcap (dS scaled by 1-tanh^2), scale + causal, f16/bf16 storage, f32 accumulate." rocm,rocm_flash_attn_compiled,native_gpu,1,rocm_flash_attn_compiled,success,hip_runtime,"ROCm flash_attn artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 forward: tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. The attention analog of the compiled GEMM lane (rocm_compiled)." rocm,rocm_fpquant_compiled,native_gpu,1,rocm_fpquant_compiled,success,hip_runtime,ROCm fpquant artifact runs quantize/dequantize fp8 / fp6 / fp4: per-tensor scale + grid-snap on the COMPILER-GENERATED fpquant kernel (generate-rocm-fpquant-kernel: log2/exp2/roundeven -> ROCDL). ROCm mirror of x86_fpquant. f32. rocm,rocm_image_affine_compiled,native_gpu,1,rocm_image_affine_compiled,success,hip_runtime,"ROCm image affine artifact runs image_normalize as (x-mean)/std: layout and per-channel broadcast on host, sub/div on generated ROCm binary kernels. f32, matches tessera.ops." diff --git a/docs/audit/generated/runtime_execution_matrix.md b/docs/audit/generated/runtime_execution_matrix.md index 313c9ba07..810657ac7 100644 --- a/docs/audit/generated/runtime_execution_matrix.md +++ b/docs/audit/generated/runtime_execution_matrix.md @@ -58,7 +58,7 @@ Single source of truth for what `runtime.launch()` does with each `(target, comp | `rocm` | `rocm_ebm_loss_compiled` | `rocm_ebm_loss_compiled` | `native_gpu` | `hip_runtime` | ROCm EBM/diffusion loss artifact runs score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance: the diff/square and reductions run on the gfx1151 binary + reduce kernels, the structure on the host. ROCm mirror of x86_ebm_loss. f32. | | `rocm` | `rocm_exotic_attn_compiled` | `rocm_exotic_attn_compiled` | `native_gpu` | `hip_runtime` | ROCm exotic-attention artifact composes the COMPILER-GENERATED WMMA flash_attn kernel with the WMMA GEMM kernel (MLA latent projections) + an elementwise gate — gated_attention, mla_decode, mla_decode_fused — and routes mla_decode_step through the DK1 absorbed-latent decode kernel against stdlib.attention.mla_decode_step. f16/bf16, f32 softmax+accumulate. | | `rocm` | `rocm_fft_compiled` | `rocm_fft_compiled` | `native_gpu` | `hip_runtime` | ROCm FFT artifact runs fft / ifft / rfft / irfft over any axis length on the COMPILER-GENERATED one-thread-per-bin DFT kernel (generate-rocm-dft-kernel -> ROCDL, cos/sin twiddles) on gfx1151 + r2c/c2r pack-unpack + plan scale. Direct DFT (radix-2/Bluestein perf is a follow-up). complex64/f32. | -| `rocm` | `rocm_flash_attn_bwd_compiled` | `rocm_flash_attn_bwd_compiled` | `native_gpu` | `hip_runtime` | ROCm flash_attn backward artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 backward: tessera-opt expands one tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/fa_dq WMMA kernels serialized to hsaco in-process, then HIP launches them in sequence to produce dQ/dK/dV. O is recomputed via the forward lane (nothing saved from forward). The reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA (gqa dkdv atomic-accumulates dK/dV across the group) + additive attn_bias (S=scale*QK+bias in the recompute), scale + causal, f16/bf16 storage, f32 accumulate. | +| `rocm` | `rocm_flash_attn_bwd_compiled` | `rocm_flash_attn_bwd_compiled` | `native_gpu` | `hip_runtime` | ROCm flash_attn backward artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 backward: tessera-opt expands one tessera_rocm.flash_attn_bwd directive into three fa_pre/fa_dkdv/fa_dq WMMA kernels serialized to hsaco in-process, then HIP launches them in sequence to produce dQ/dK/dV. O is recomputed via the forward lane (nothing saved from forward). The reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA (gqa dkdv atomic-accumulates dK/dV across the group) + additive attn_bias + sliding-window (implicitly causal, masks keys older than W) + Gemma-2 logit-softcap (dS scaled by 1-tanh^2), scale + causal, f16/bf16 storage, f32 accumulate. | | `rocm` | `rocm_flash_attn_compiled` | `rocm_flash_attn_compiled` | `native_gpu` | `hip_runtime` | ROCm flash_attn artifact runs the COMPILER-GENERATED RDNA WMMA FA-2 forward: tessera-opt generates + serializes the kernel to hsaco in-process, then HIP loads + launches it. The attention analog of the compiled GEMM lane (rocm_compiled). | | `rocm` | `rocm_fpquant_compiled` | `rocm_fpquant_compiled` | `native_gpu` | `hip_runtime` | ROCm fpquant artifact runs quantize/dequantize fp8 / fp6 / fp4: per-tensor scale + grid-snap on the COMPILER-GENERATED fpquant kernel (generate-rocm-fpquant-kernel: log2/exp2/roundeven -> ROCDL). ROCm mirror of x86_fpquant. f32. | | `rocm` | `rocm_image_affine_compiled` | `rocm_image_affine_compiled` | `native_gpu` | `hip_runtime` | ROCm image affine artifact runs image_normalize as (x-mean)/std: layout and per-channel broadcast on host, sub/div on generated ROCm binary kernels. f32, matches tessera.ops. | @@ -220,7 +220,7 @@ nvidia_sm80, nvidia_sm90, nvidia_sm100, rocm_gfx90a, rocm_gfx940, rocm_gfx942, r | `rocm_ebm_loss_compiled` | AMD GPU RDNA EBM/diffusion loss (score_matching / denoising / implicit / contrastive_divergence / persistent_cd / ddpm_noise_pred / vlb / load_balance) — diff/square + reductions on the gfx1151 binary + reduce kernels, host structure. ROCm mirror of x86_ebm_loss. f32 | | `rocm_exotic_attn_compiled` | AMD GPU RDNA exotic-attention compositions (gated_attention, mla_decode, mla_decode_fused, mla_decode_step absorbed-latent decode) built by COMPOSING the COMPILER-GENERATED WMMA flash_attn kernel + the WMMA GEMM kernel (MLA latent projections) + an elementwise gate, plus the DK1 generated absorbed-latent ROCm decode kernel; f16/bf16 storage, f32 softmax+accumulate. The block-sparse deepseek variant stays artifact_only | | `rocm_fft_compiled` | AMD GPU RDNA spectral FFT (fft / ifft / rfft / irfft) the Tessera compiler GENERATES (generate-rocm-dft-kernel -> ROCDL -> hsaco; one thread per output bin, cos/sin twiddles), then HIP launches it. Direct DFT (any length) on gfx1151 + r2c/c2r pack-unpack + plan scale (radix-2/Bluestein perf is a follow-up). complex64/f32 | -| `rocm_flash_attn_bwd_compiled` | AMD GPU RDNA WMMA FA-2 BACKWARD the Tessera compiler GENERATES (generate-wmma-flash-attn-bwd-kernel -> three fa_pre/fa_dkdv/fa_dq WMMA kernels -> hsaco), launched in sequence to produce dQ/dK/dV; O is recomputed via the forward lane (nothing saved from forward). MHA + GQA/MQA (grouped dkdv atomic-accumulates dK/dV) + additive attn_bias; f16/bf16 storage, f32 accumulate; the reverse-mode analog of rocm_flash_attn_compiled | +| `rocm_flash_attn_bwd_compiled` | AMD GPU RDNA WMMA FA-2 BACKWARD the Tessera compiler GENERATES (generate-wmma-flash-attn-bwd-kernel -> three fa_pre/fa_dkdv/fa_dq WMMA kernels -> hsaco), launched in sequence to produce dQ/dK/dV; O is recomputed via the forward lane (nothing saved from forward). MHA + GQA/MQA (grouped dkdv atomic-accumulates dK/dV) + additive attn_bias + sliding-window + logit-softcap; f16/bf16 storage, f32 accumulate; the reverse-mode analog of rocm_flash_attn_compiled | | `rocm_flash_attn_compiled` | AMD GPU RDNA WMMA FA-2 forward the Tessera compiler GENERATES (generate-wmma-flash-attn-kernel -> ROCDL -> hsaco, in-process via tessera-opt), then HIP loads + launches it. f16/bf16 storage, f32 softmax + accumulate; the attention analog of rocm_compiled | | `rocm_fpquant_compiled` | AMD GPU RDNA low-precision float quantize (quantize/dequantize fp8 / fp6 / fp4) — grid-snap on generate-rocm-fpquant-kernel (log2/exp2/roundeven) + per-tensor scale. ROCm mirror of x86_fpquant. f32 | | `rocm_image_affine_compiled` | AMD GPU RDNA image affine preprocessing — image_normalize as sub/div on generated ROCm binary kernels with host layout and per-channel broadcast | diff --git a/docs/rocm_mfma_kernel_inventory.md b/docs/rocm_mfma_kernel_inventory.md index fa2cf8a6b..a2f145d49 100644 --- a/docs/rocm_mfma_kernel_inventory.md +++ b/docs/rocm_mfma_kernel_inventory.md @@ -416,8 +416,9 @@ symbols (not just an in-process compiled lane): (maxerr ~1e-4 f16). The FA-2 **backward** (dQ/dK/dV) also executes on gfx1151 via the compiler-generated `rocm_flash_attn_bwd_compiled` lane (`generate-wmma-flash-attn-bwd-kernel` → `fa_pre`/`fa_dkdv`/`fa_dq`; O - recomputed via the forward lane, nothing saved from forward), **MHA + GQA/MQA** - (scale + causal), validated vs autodiff `vjp_flash_attn`. No perf ladder. + recomputed via the forward lane, nothing saved from forward), **MHA + GQA/MQA + + additive attn_bias + sliding-window + logit-softcap** (scale + causal), + validated vs autodiff `vjp_flash_attn`. No perf ladder. Honest scope (Decision #25): everything above is **one arch (RDNA 3.5 `gfx1151`) × {fp16, bf16}**, correctness-first. None of it flips the per-primitive @@ -461,9 +462,10 @@ promotes them to `compileable`. See §9 for the concrete done / open / blocked s `rocm_flash_attn_bwd_compiled` lane (`generate-wmma-flash-attn-bwd-kernel` → `fa_pre`/`fa_dkdv`/`fa_dq`, launched through `runtime.launch()`), matches autodiff `vjp_flash_attn` on `{fp16, bf16}`; **MHA + GQA/MQA** (grouped `fa_dkdv` - atomic-accumulates dK/dV across the group) **+ additive attn_bias** - (`S = scale*Q@K^T + bias` in the recompute), scale + causal. No perf ladder; - windowed/softcap backward is the remaining follow-up + atomic-accumulates dK/dV) **+ additive attn_bias + sliding-window** (implicitly + causal, masks keys older than W) **+ logit-softcap** (`S = cap·tanh(scale·QK/cap)`; + backward scales dS by `1−tanh²`), scale + causal — the full forward variant + surface. No perf ladder - ✅ **Dozens of additional compiler-generated HIP `compiled` lanes execute** and match a CPU/numpy reference — nearly all of §5 plus §10: the GEMM/attention families, norm / activation / RoPE / ALiBi, optimizers, RNG, FFT/spectral, @@ -485,9 +487,9 @@ promotes them to `compileable`. See §9 for the concrete done / open / blocked s (the emitter currently rides the direct LLVM-IR path) ### Open on this box (gfx1151 — workable now, no CDNA needed) -- `flash_attn` **backward** for the sliding-window / logit-softcap variants - (MHA + GQA/MQA + additive **attn_bias** backward is done and runtime-wired; - window/softcap still need the masked/capped backward math in the C++ kernel) +- (flash_attn backward is now the **full forward variant surface** — MHA + + GQA/MQA + attn_bias + sliding-window + logit-softcap — all runtime-wired) +- Perf ladders / MFU sign-off for the compiled lanes (still correctness-first) - **Fused paged-attention** — the §5.6 movement core (`kv_cache_append/read/prune`) now executes via `rocm_kv_cache_compiled` (scatter/gather compose, execute-compare vs `KVCacheHandle`); a single fused gather→attention paged diff --git a/python/tessera/compiler/backend_manifest.py b/python/tessera/compiler/backend_manifest.py index 8e3164fa9..a93786f9c 100644 --- a/python/tessera/compiler/backend_manifest.py +++ b/python/tessera/compiler/backend_manifest.py @@ -1224,8 +1224,9 @@ def primitive_is_complete(entries: tuple["BackendKernelEntry", ...]) -> bool: "BACKWARD (dQ/dK/dV) also executes on gfx1151 via the " "compiler-generated rocm_flash_attn_bwd_compiled lane " "(generate-wmma-flash-attn-bwd-kernel -> fa_pre/fa_dkdv/fa_dq; MHA + " - "GQA/MQA + additive attn_bias, scale+causal), validated vs autodiff " - "vjp_flash_attn — see the runtime_execution_matrix." + "GQA/MQA + additive attn_bias + sliding-window + logit-softcap, " + "scale+causal), validated vs autodiff vjp_flash_attn — see the " + "runtime_execution_matrix." ), }, } diff --git a/python/tessera/compiler/execution_matrix.py b/python/tessera/compiler/execution_matrix.py index 4d4271306..3aa6a5e34 100644 --- a/python/tessera/compiler/execution_matrix.py +++ b/python/tessera/compiler/execution_matrix.py @@ -199,8 +199,9 @@ class ExecutionRow: "hsaco), launched in sequence to produce dQ/dK/dV; O " "is recomputed via the forward lane (nothing saved " "from forward). MHA + GQA/MQA (grouped dkdv atomic-" - "accumulates dK/dV) + additive attn_bias; f16/bf16 " - "storage, f32 accumulate; the reverse-mode analog of " + "accumulates dK/dV) + additive attn_bias + sliding-" + "window + logit-softcap; f16/bf16 storage, f32 " + "accumulate; the reverse-mode analog of " "rocm_flash_attn_compiled", "rocm_linear_attn_compiled": "AMD GPU RDNA WMMA linear-attention forward the " "Tessera compiler GENERATES " @@ -1782,8 +1783,9 @@ class ExecutionRow: "via the forward lane (nothing saved from forward). The " "reverse-mode analog of rocm_flash_attn_compiled; MHA + GQA/MQA " "(gqa dkdv atomic-accumulates dK/dV across the group) + additive " - "attn_bias (S=scale*QK+bias in the recompute), scale + causal, " - "f16/bf16 storage, f32 accumulate.", + "attn_bias + sliding-window (implicitly causal, masks keys older " + "than W) + Gemma-2 logit-softcap (dS scaled by 1-tanh^2), scale + " + "causal, f16/bf16 storage, f32 accumulate.", execution_mode="hip_runtime"), # Linear-attention family (quadratic-parallel form, no softmax; a distinct # algorithm from flash_attn): tessera.linear_attn + the decay-masked siblings diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index c52eb865f..f1172af90 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -3138,23 +3138,27 @@ def _mr(p, size): # carries the GQA variant; the runtime lane forwards only the core here). # ───────────────────────────────────────────────────────────────────────────── #: hsaco bytes keyed by (head_dim, chip, dtype, gqa, bias). -_rocm_fa_bwd_hsaco_cache: dict[tuple[int, str, str, bool, bool], bytes] = {} +_rocm_fa_bwd_hsaco_cache: dict[ + tuple[int, str, str, bool, bool, bool, bool], bytes] = {} def _build_compiled_flash_attn_bwd_hsaco(head_dim: int, dtype: str = "f16", gqa: bool = False, - bias: bool = False) -> bytes: + bias: bool = False, + window: bool = False, + softcap: bool = False) -> bytes: """Generate + serialize the compiler's WMMA FA-2 backward kernels to hsaco (fa_pre / fa_dkdv / fa_dq in one module), fully in-process via tessera-opt. - Cached per (head_dim, chip, dtype, gqa, bias). gqa=True emits the grouped- - query variant: all three kernels gain (heads, kv_ratio) runtime args and - fa_dkdv atomically accumulates dK/dV across the kv_ratio query heads sharing - each KV head (host pre-zeros dK/dV). bias=True emits the additive-bias - variant: the recompute forms S = scale*Q@K^T + bias (a trailing f32 - [bh,Sq,Sk] memref arg, LAST) before the softmax.""" + Cached per (head_dim, chip, dtype, gqa, bias, window, softcap). Trailing + runtime args appended in a FIXED order — gqa(heads,kv_ratio) | window(W) | + softcap(cap) | bias([bh,Sq,Sk], LAST): gqa=True adds the grouped-query + variant (fa_dkdv atomically accumulates dK/dV, host pre-zeros); window=True + the sliding-window variant (implicitly causal, masks keys older than W); + softcap=True the Gemma-2 soft-cap (S=cap*tanh(scale*QK/cap), backward scales + dS by 1-tanh^2); bias=True the additive-bias variant.""" chip = _rocm_chip() - key = (head_dim, chip, dtype, gqa, bias) + key = (head_dim, chip, dtype, gqa, bias, window, softcap) cached = _rocm_fa_bwd_hsaco_cache.get(key) if cached is not None: return cached @@ -3162,12 +3166,16 @@ def _build_compiled_flash_attn_bwd_hsaco(head_dim: int, if opt is None: raise _RocmCompiledUnavailable( "tessera-opt not built — no compiled ROCm flash_attn backward lane") - gqa_attr = ", gqa = true" if gqa else "" - bias_attr = ", attn_bias = true" if bias else "" + attrs = "".join([ + ", gqa = true" if gqa else "", + ", sliding_window = true" if window else "", + ", logit_softcap = true" if softcap else "", + ", attn_bias = true" if bias else "", + ]) directive = ( 'module {\n' ' "tessera_rocm.flash_attn_bwd"() {name = "fa", ' - f'head_dim = {head_dim} : i64, dtype = "{dtype}"{gqa_attr}{bias_attr}}} ' + f'head_dim = {head_dim} : i64, dtype = "{dtype}"{attrs}}} ' ': () -> ()\n' '}\n' ) @@ -3197,19 +3205,25 @@ def _build_compiled_flash_attn_bwd_hsaco(head_dim: int, def _rocm_flash_attn_forward_o(q: Any, k: Any, v: Any, scale: float, - causal: int, bias: Any = None) -> Any: + causal: int, bias: Any = None, + window: int = 0, + logit_softcap: float = 0.0) -> Any: """Recompute the forward output O (f32, shape [bh, sq, D]) by REUSING the - forward compiled lane — the backward saves nothing from forward. When - ``bias`` is given, O = softmax(scale*Q@K^T + bias)*V (the forward lane takes - attn_bias as a 4th operand), so O matches the biased backward recompute.""" + forward compiled lane — the backward saves nothing from forward. The forward + lane takes attn_bias as a 4th operand and window / logit_softcap as kwargs, + so O matches the backward's windowed / soft-capped / biased score recompute.""" names = ["q", "k", "v"] + (["bias"] if bias is not None else []) + kw: dict[str, Any] = {"scale": scale, "causal": bool(causal)} + if window: + kw["window"] = int(window) + if logit_softcap: + kw["logit_softcap"] = float(logit_softcap) fwd_art = RuntimeArtifact(metadata={ "target": "rocm", "compiler_path": "rocm_flash_attn_compiled", "executable": True, "execution_kind": "native_gpu", "arg_names": names, "output_name": "o", "ops": [{"op_name": "tessera.flash_attn", "result": "o", - "operands": names, - "kwargs": {"scale": scale, "causal": bool(causal)}}]}) + "operands": names, "kwargs": kw}]}) call = (q, k, v) if bias is None else (q, k, v, bias) return _execute_rocm_compiled_flash_attn(fwd_art, call) @@ -3297,11 +3311,22 @@ def _execute_rocm_compiled_flash_attn_bwd(artifact: RuntimeArtifact, causal = 1 if bool(kwargs.get("causal", False)) else 0 scale = kwargs.get("scale") scale = float(scale) if scale is not None else 1.0 / float(np.sqrt(head_dim)) - for bad, why in (("window", "sliding-window"), ("logit_softcap", "softcap")): - if kwargs.get(bad): - raise ValueError( - f"rocm flash_attn backward does not yet support {why}; the core " - "MHA (scale + causal) + additive attn_bias backward is wired.") + # Sliding-window: a positive `window` selects a causal band of width W + # (query q attends to keys in (q-W, q]); the windowed kernel is implicitly + # causal. Gemma-2 soft-cap: a positive `logit_softcap` caps the pre-softmax + # score to cap*tanh(S/cap); the backward scales dS by 1-tanh^2(S_raw/cap). + window = int(kwargs.get("window") or 0) + if window < 0: + raise ValueError( + f"rocm flash_attn backward window must be non-negative; got {window}") + sliding = window > 0 + softcap = kwargs.get("logit_softcap") + softcap = float(softcap) if softcap else 0.0 + if softcap < 0: + raise ValueError( + f"rocm flash_attn backward logit_softcap must be non-negative; got " + f"{softcap}") + has_softcap = softcap > 0 # Additive bias: S = softmax(scale*Q@K^T + bias). Host-broadcast to # Q.lead+(Sq,Sk) → [bh,Sq,Sk] f32 (kernel indexes bias[(bh*Sq+q)*Sk+k]), @@ -3322,15 +3347,17 @@ def _execute_rocm_compiled_flash_attn_bwd(artifact: RuntimeArtifact, f"{tuple(q.shape[:-2]) + (sq, sk)}") from exc bias_c = np.ascontiguousarray(bias_b, dtype=np.float32).reshape(-1) - # O the pre-pass needs — recompute on-device via the forward lane (with the - # same bias, so O = softmax(scale*QK + bias)*V matches the recompute). + # O the pre-pass needs — recompute on-device via the forward lane with the + # SAME window/softcap/bias, so O matches the backward's score recompute. o_f32 = np.ascontiguousarray( _rocm_flash_attn_forward_o(q, k, v, scale, causal, - bias=bias_arr if has_bias else None), + bias=bias_arr if has_bias else None, + window=window, logit_softcap=softcap), dtype=np.float32).reshape(bh, sq, head_dim) - hsaco = _build_compiled_flash_attn_bwd_hsaco(head_dim, dtype_tag, gqa=gqa, - bias=has_bias) + hsaco = _build_compiled_flash_attn_bwd_hsaco( + head_dim, dtype_tag, gqa=gqa, bias=has_bias, window=sliding, + softcap=has_softcap) hip = _load_hip_for_launch() if hip is None: raise _RocmCompiledUnavailable( @@ -3395,9 +3422,14 @@ def _mr(p, size): # Grouped kernels take (heads, kv_ratio) as two trailing runtime args after # (Sq, Sk, scale, causal); the plain kernels have neither. Grid.y is always # the query-head count B*H (fa_dkdv accumulates into KV heads atomically). - tail = [ctypes.c_int64(n_qh), ctypes.c_int64(kv_ratio)] if gqa else [] - # Additive bias memref is the LAST kernel arg on all three kernels (after the - # optional gqa pair) — matches the C++ withGqa(... + bias) arg order. + tail: list[Any] = ( + [ctypes.c_int64(n_qh), ctypes.c_int64(kv_ratio)] if gqa else []) + # Trailing args in the C++ withGqa order: gqa pair | window W | softcap cap | + # bias memref (LAST). Each present only when its flag is set. + if sliding: + tail.append(ctypes.c_int64(window)) + if has_softcap: + tail.append(ctypes.c_float(softcap)) bias_tail = _mr(bufs["bias"], n_bias) if has_bias else [] gqt, gkt, gyt = (sq + 15) // 16, (sk + 15) // 16, bh diff --git a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp index 6e0514713..8b0b6b66c 100644 --- a/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp +++ b/src/compiler/codegen/Tessera_ROCM_Backend/lib/Conversion/GenerateWMMAFlashAttnBwdKernel.cpp @@ -141,7 +141,8 @@ static Value groupedKvBase(Emit &e, OpBuilder &b, Location loc, Value bh, // cheap O(D) per-row elementwise rowsum (not on the hot path). //===----------------------------------------------------------------------===// void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, - Type storeTy, bool gqa = false, bool bias = false) { + Type storeTy, bool gqa = false, bool bias = false, + bool window = false, bool softcap = false) { MLIRContext *ctx = b.getContext(); Type f32 = b.getF32Type(); Value sQ = f.addWorkgroupAttribution(ldsT(ctx, 16 * D, storeTy), loc); @@ -172,10 +173,17 @@ void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, : e.mul(e.mul(bh, Sk), cD); Value isCausal = b.create(loc, arith::CmpIPredicate::ne, causal, c0); - // Additive bias memref (LAST arg, after the optional gqa pair): f32 [bh,Sq,Sk]. - Value biasBuf; - if (bias) - biasBuf = f.getArgument(gqa ? 12 : 10); + // Trailing runtime args, in order: gqa(heads,kv_ratio) | window(W) | + // softcap(cap) | bias([bh,Sq,Sk], LAST). Base arg count for _pre is 10. + int64_t p = 10 + (gqa ? 2 : 0); + Value W = window ? f.getArgument(p) : Value(); + if (window) ++p; + Value cap = softcap ? f.getArgument(p) : Value(); + if (softcap) ++p; + Value biasBuf = bias ? f.getArgument(p) : Value(); + // A windowed kernel is implicitly causal (for the L logsumexp bound + mask). + Value trueI1 = b.create(loc, 1, /*width=*/1); + Value useCausal = window ? trueI1 : isCausal; // Stage Q into sQ (all 32 lanes cooperatively): for i = tid; i < 16*D; i+=32. { @@ -226,7 +234,7 @@ void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, Value nKV = b.create(loc, e.add(Sk, c15), c16); Value nKVm1 = b.create(loc, nKV, c1); Value ckt = b.create(loc, e.add(q0, c15), c16); - Value lastKt = e.sel(isCausal, ckt, nKVm1); + Value lastKt = e.sel(useCausal, ckt, nKVm1); lastKt = e.sel(e.lt(nKVm1, lastKt), nKVm1, lastKt); Value upper = e.add(lastKt, c1); @@ -257,18 +265,28 @@ void emitPre(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, for (int64_t el = 0; el < 8; ++el) { Value qi = e.add(e.ci(2 * el), half); Value v0 = e.mulf(e.ext(cs, el), scale); - // Additive bias — same S = scale*Q@K^T + bias the recompute uses, so L - // (logsumexp) matches P/dS. Bounds-guarded against the [bh,Sq,Sk] buffer. + // Soft-cap (before bias): v0 = cap*tanh(v0/cap). L must see the same + // capped score the recompute uses, so P = exp(S_capped - L) is consistent. + if (softcap) { + Value t = b.create(loc, b.create( + loc, v0, cap)); + v0 = e.mulf(cap, t); + } + // Additive bias — same S the recompute uses, so L (logsumexp) matches + // P/dS. Bounds-guarded against the [bh,Sq,Sk] buffer. + Value gqe = e.add(q0, qi); if (bias) { - Value gqe = e.add(q0, qi); Value gqSafe = e.sel(e.lt(gqe, Sq), gqe, c0); Value gkSafe = e.sel(e.ge(gk, Sk), c0, gk); Value bidx = e.add(e.mul(e.add(e.mul(bh, Sq), gqSafe), Sk), gkSafe); v0 = e.addf(v0, e.f32load(biasBuf, bidx)); } - Value cmask = b.create(loc, isCausal, - e.lt(e.add(q0, qi), gk)); + Value cmask = b.create(loc, useCausal, e.lt(gqe, gk)); Value masked = b.create(loc, e.ge(gk, Sk), cmask); + if (window) { // too-old key: q - k >= W + Value age = b.create(loc, gqe, gk); + masked = b.create(loc, masked, e.ge(age, W)); + } b.create(loc, e.sel(masked, e.negInf, v0), sS, ValueRange{e.add(e.mul(qi, c16), l15)}); } @@ -337,6 +355,13 @@ struct ScoreCtx { // a null Value when the kernel has no bias. When set, the recompute forms // S = scale*Q@K^T + bias before the softmax, so P/L/dS all see the bias. Value biasBuf; + // Sliding window width W (index) or null. When set the kernel is implicitly + // causal and keys older than W (q - k >= W) are masked out. + Value W; + // Gemma-2 logit soft-cap value (f32) or null. When set the pre-softmax score + // is S = cap*tanh(scale*Q@K^T / cap); the backward multiplies dS by the + // soft-cap derivative 1 - tanh^2(raw/cap). + Value cap; }; void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, @@ -384,9 +409,20 @@ void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, Value Lidx = e.add(e.mul(x.bh, x.Sq), gqSafe); Value Lq = e.f32load(x.L, Lidx); Value Dq = e.f32load(x.Dd, Lidx); - Value s = e.mulf(e.ext(cs, el), x.scale); - // Additive bias: S = scale*Q@K^T + bias[(bh*Sq + q)*Sk + k]. Guarded on the - // query/key bounds so masked lanes never read past the [bh,Sq,Sk] buffer + Value sRaw = e.mulf(e.ext(cs, el), x.scale); + // Gemma-2 logit soft-cap: S = cap*tanh(sRaw/cap). Backward multiplies dS by + // the soft-cap derivative 1 - tanh^2(sRaw/cap) (chain rule through the cap). + Value s = sRaw, capDeriv; + if (x.cap) { + Value t = b.create(loc, b.create( + loc, sRaw, x.cap)); + s = e.mulf(x.cap, t); + Value one = b.create(loc, e.f32, + b.getF32FloatAttr(1.0f)); + capDeriv = e.subf(one, e.mulf(t, t)); + } + // Additive bias: S += bias[(bh*Sq + q)*Sk + k] (after soft-cap). Guarded on + // the query/key bounds so masked lanes never read past the [bh,Sq,Sk] buffer // (their P is zeroed by `masked` below anyway). if (x.biasBuf) { Value gkSafe = e.sel(e.ge(gk, x.Sk), c0, gk); @@ -395,14 +431,23 @@ void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, s = e.addf(s, e.f32load(x.biasBuf, bidx)); } Value P = b.create(loc, e.subf(s, Lq)); - // mask: query OOB, key OOB, or causal (key > query) -> P = 0 + // mask: query OOB, key OOB, causal (key > query), or (windowed) too-old key + // (q - k >= W) -> P = 0. A windowed kernel is implicitly causal. + Value trueI1 = b.create(loc, 1, 1); + Value useCausal = x.W ? trueI1 : x.isCausal; Value m1 = e.ge(gqi, x.Sq); Value m2 = e.ge(gk, x.Sk); - Value m3 = b.create(loc, x.isCausal, e.lt(gqi, gk)); + Value m3 = b.create(loc, useCausal, e.lt(gqi, gk)); Value masked = b.create( loc, b.create(loc, m1, m2), m3); + if (x.W) { + Value age = b.create(loc, gqi, gk); // q - k + masked = b.create(loc, masked, e.ge(age, x.W)); + } P = e.sel(masked, e.zerof, P); Value dS = e.mulf(P, e.subf(e.ext(cp, el), Dq)); + if (x.cap) + dS = e.mulf(dS, capDeriv); // dS_raw = dS_capped * cap' dS = e.sel(masked, e.zerof, dS); Value sIdx = e.add(e.mul(qi, c16), x.l15); if (wantP) @@ -417,7 +462,8 @@ void recomputeScoreTile(Emit &e, OpBuilder &b, Location loc, const ScoreCtx &x, // Sq, Sk : index, scale : f32, causal : index) //===----------------------------------------------------------------------===// void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, - Type storeTy, bool gqa = false, bool bias = false) { + Type storeTy, bool gqa = false, bool bias = false, + bool window = false, bool softcap = false) { MLIRContext *ctx = b.getContext(); Value sP = f.addWorkgroupAttribution(ldsT(ctx, 16 * 16, storeTy), loc); Value sDS = f.addWorkgroupAttribution(ldsT(ctx, 16 * 16, storeTy), loc); @@ -449,9 +495,15 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, : e.mul(e.mul(bh, Sk), cD); Value isCausal = b.create(loc, arith::CmpIPredicate::ne, causal, c0); - Value biasBuf; // f32 [bh,Sq,Sk], LAST arg (after the optional gqa pair) - if (bias) - biasBuf = f.getArgument(gqa ? 14 : 12); + // Trailing args: gqa | window(W) | softcap(cap) | bias (LAST). Base = 12. + int64_t p = 12 + (gqa ? 2 : 0); + Value W = window ? f.getArgument(p) : Value(); + if (window) ++p; + Value cap = softcap ? f.getArgument(p) : Value(); + if (softcap) ++p; + Value biasBuf = bias ? f.getArgument(p) : Value(); + Value trueI1 = b.create(loc, 1, /*width=*/1); + Value useCausal = window ? trueI1 : isCausal; // zero dK/dV accumulators. { @@ -471,7 +523,7 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, // causal, start the query loop at `ktile` (skip the tiles entirely below the // diagonal); the diagonal tile qt==ktile is still per-element masked. ~halves // the query-tile work for causal. Non-causal starts at 0. - Value qStart = e.sel(isCausal, ktile, c0); + Value qStart = e.sel(useCausal, ktile, c0); auto qloop = b.create(loc, qStart, nQ, c1); { OpBuilder::InsertionGuard g(b); @@ -481,6 +533,8 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, ScoreCtx x{Q, Kk, V, dO, L, Dd, Sq, Sk, scale, isCausal, qbase, kbase, q0, k0, bh, l15, half}; x.biasBuf = biasBuf; + x.W = W; + x.cap = cap; recomputeScoreTile(e, b, loc, x, sP, sDS, /*wantP=*/true); b.create(loc); @@ -570,7 +624,8 @@ void emitDkDv(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, // Sq, Sk : index, scale : f32, causal : index) //===----------------------------------------------------------------------===// void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, - Type storeTy, bool gqa = false, bool bias = false) { + Type storeTy, bool gqa = false, bool bias = false, + bool window = false, bool softcap = false) { MLIRContext *ctx = b.getContext(); Value sDS = f.addWorkgroupAttribution(ldsT(ctx, 16 * 16, storeTy), loc); Value dQacc = f.addWorkgroupAttribution(ldsT(ctx, 16 * D, b.getF32Type()), loc); @@ -598,9 +653,15 @@ void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, : e.mul(e.mul(bh, Sk), cD); Value isCausal = b.create(loc, arith::CmpIPredicate::ne, causal, c0); - Value biasBuf; // f32 [bh,Sq,Sk], LAST arg (after the optional gqa pair) - if (bias) - biasBuf = f.getArgument(gqa ? 13 : 11); + // Trailing args: gqa | window(W) | softcap(cap) | bias (LAST). Base = 11. + int64_t p = 11 + (gqa ? 2 : 0); + Value W = window ? f.getArgument(p) : Value(); + if (window) ++p; + Value cap = softcap ? f.getArgument(p) : Value(); + if (softcap) ++p; + Value biasBuf = bias ? f.getArgument(p) : Value(); + Value trueI1 = b.create(loc, 1, /*width=*/1); + Value useCausal = window ? trueI1 : isCausal; { auto lp = b.create(loc, tid, c16D, c32); @@ -620,7 +681,7 @@ void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, Value nKfull = b.create(loc, e.add(Sk, c15), c16); Value cKlimit = e.add(qtile, c1); Value nKcausal = e.sel(e.lt(cKlimit, nKfull), cKlimit, nKfull); - Value nK = e.sel(isCausal, nKcausal, nKfull); + Value nK = e.sel(useCausal, nKcausal, nKfull); auto kloop = b.create(loc, c0, nK, c1); { OpBuilder::InsertionGuard g(b); @@ -630,6 +691,8 @@ void emitDq(OpBuilder &b, Location loc, gpu::GPUFuncOp f, int64_t D, ScoreCtx x{Q, Kk, V, dO, L, Dd, Sq, Sk, scale, isCausal, qbase, kbase, q0, k0, bh, l15, half}; x.biasBuf = biasBuf; + x.W = W; + x.cap = cap; recomputeScoreTile(e, b, loc, x, /*sP=*/Value(), sDS, /*wantP=*/false); b.create(loc); @@ -757,12 +820,18 @@ struct GenerateWMMAFlashAttnBwdKernelPass SmallVector gqaExtra; if (gqa) gqaExtra = {idxTy, idxTy}; - // Additive attention bias: a trailing f32 [bh,Sq,Sk] memref (LAST arg, - // after the optional gqa pair). S = scale*Q@K^T + bias before the softmax, - // so P/L/dS all see it; dbias is not emitted (constant-bias VJP path). - bool bias = false; - if (auto a = op->getAttrOfType("attn_bias")) - bias = a.getValue(); + // Optional variants — trailing args appended in a FIXED order after the + // base signature: gqa(heads,kv_ratio) | window(W:index) | softcap(cap:f32) + // | attn_bias([bh,Sq,Sk] f32, LAST). Sliding-window is implicitly causal + // and masks keys older than W; soft-cap forms S=cap*tanh(scale*QK/cap) + // before the softmax (backward scales dS by 1-tanh^2); bias is additive. + auto flag = [&](StringRef n) { + auto a = op->getAttrOfType(n); + return a && a.getValue(); + }; + bool window = flag("sliding_window"); + bool softcap = flag("logit_softcap"); + bool bias = flag("attn_bias"); auto mk = [&](StringRef suffix, ArrayRef args, function_ref body) { @@ -775,24 +844,26 @@ struct GenerateWMMAFlashAttnBwdKernelPass auto withGqa = [&](ArrayRef base) { SmallVector a(base.begin(), base.end()); a.append(gqaExtra.begin(), gqaExtra.end()); - if (bias) a.push_back(fv); // bias [bh,Sq,Sk] f32, LAST + if (window) a.push_back(idxTy); // W (window width) + if (softcap) a.push_back(f32); // cap + if (bias) a.push_back(fv); // bias [bh,Sq,Sk] f32, LAST return a; }; - // _pre : (Q,K,dO:store, O,L,Dd:f32, Sq,Sk:idx, scale:f32, causal:idx [+gqa][+bias]) + // _pre : (Q,K,dO:store, O,L,Dd:f32, Sq,Sk:idx, scale:f32, causal:idx [+opts]) mk("_pre", withGqa({sv, sv, sv, fv, fv, fv, idxTy, idxTy, f32, idxTy}), [&](OpBuilder &bb, Location l, gpu::GPUFuncOp fn) { - emitPre(bb, l, fn, D, storeTy, gqa, bias); + emitPre(bb, l, fn, D, storeTy, gqa, bias, window, softcap); }); - // _dkdv : (Q,K,V,dO:store, L,Dd:f32, dK,dV:f32, Sq,Sk:idx, scale, causal [+gqa][+bias]) + // _dkdv : (Q,K,V,dO:store, L,Dd:f32, dK,dV:f32, Sq,Sk:idx, scale, causal [+opts]) mk("_dkdv", withGqa({sv, sv, sv, sv, fv, fv, fv, fv, idxTy, idxTy, f32, idxTy}), [&](OpBuilder &bb, Location l, gpu::GPUFuncOp fn) { - emitDkDv(bb, l, fn, D, storeTy, gqa, bias); + emitDkDv(bb, l, fn, D, storeTy, gqa, bias, window, softcap); }); - // _dq : (Q,K,V,dO:store, L,Dd:f32, dQ:f32, Sq,Sk:idx, scale, causal [+gqa][+bias]) + // _dq : (Q,K,V,dO:store, L,Dd:f32, dQ:f32, Sq,Sk:idx, scale, causal [+opts]) mk("_dq", withGqa({sv, sv, sv, sv, fv, fv, fv, idxTy, idxTy, f32, idxTy}), [&](OpBuilder &bb, Location l, gpu::GPUFuncOp fn) { - emitDq(bb, l, fn, D, storeTy, gqa, bias); + emitDq(bb, l, fn, D, storeTy, gqa, bias, window, softcap); }); op->erase(); } diff --git a/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py b/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py index 2c68e1d26..041cd972b 100644 --- a/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py +++ b/tests/unit/test_rocm_flash_attn_bwd_runtime_lane.py @@ -27,54 +27,71 @@ def _rocm_or_skip(): return rt -def _ref_bwd(Q, K, V, dO, scale, causal, bias=None): +def _ref_bwd(Q, K, V, dO, scale, causal, bias=None, window=0, softcap=0.0): """numpy FA backward (f32 math from the f16/bf16 storage inputs). Handles - GQA/MQA: query head h reads KV head g = h // (H//G); dK/dV accumulate over the - group. Q/dO are [B,H,Sq,D]; K/V are [B,G,Sk,D]. Optional additive bias - [B,H,Sq,Sk] enters S before the softmax.""" + GQA/MQA (query head h reads KV head g = h // (H//G); dK/dV accumulate over the + group), optional additive bias [B,H,Sq,Sk], sliding window W (q attends keys + in (q-W, q], implicitly causal), and Gemma-2 soft-cap (S=cap*tanh(S_raw/cap), + backward scales dS by 1-tanh^2). Q/dO are [B,H,Sq,D]; K/V are [B,G,Sk,D].""" B, H, Sq, D = Q.shape G, Sk = K.shape[1], K.shape[2] ratio = H // G Qf, Kf, Vf, dOf = (a.astype(np.float32) for a in (Q, K, V, dO)) biasf = None if bias is None else np.asarray(bias, np.float32) + i = np.arange(Sq)[:, None]; j = np.arange(Sk)[None, :] dQ = np.zeros((B, H, Sq, D), np.float32) dK = np.zeros((B, G, Sk, D), np.float32) dV = np.zeros((B, G, Sk, D), np.float32) for b in range(B): for h in range(H): g = h // ratio - s = scale * (Qf[b, h] @ Kf[b, g].T) + s_raw = scale * (Qf[b, h] @ Kf[b, g].T) + if softcap: + t = np.tanh(s_raw / softcap) + s = softcap * t + capderiv = 1.0 - t * t + else: + s = s_raw + capderiv = 1.0 if biasf is not None: s = s + biasf[b, h] - if causal: - i = np.arange(Sq)[:, None]; j = np.arange(Sk)[None, :] - s = np.where(j > i, -1e30, s) + mask = np.zeros((Sq, Sk), bool) + if causal or window: + mask |= (j > i) # causal upper (future keys) + if window: + mask |= ((i - j) >= window) # too-old keys + s = np.where(mask, -1e30, s) s = s - s.max(-1, keepdims=True) p = np.exp(s); p = p / p.sum(-1, keepdims=True) O = p @ Vf[b, g] dp = dOf[b, h] @ Vf[b, g].T dq_row = np.sum(O * dOf[b, h], axis=-1, keepdims=True) - ds = p * (dp - dq_row) + ds = p * (dp - dq_row) * capderiv # dS_raw = dS_capped * cap' dQ[b, h] = scale * (ds @ Kf[b, g]) dK[b, g] += scale * (ds.T @ Qf[b, h]) dV[b, g] += p.T @ dOf[b, h] return dQ, dK, dV -def _art(rt, causal, scale, bias=False): +def _art(rt, causal, scale, bias=False, window=0, softcap=0.0): names = ["do", "q", "k", "v"] + (["bias"] if bias else []) + kw = {"scale": scale, "causal": causal} + if window: + kw["window"] = int(window) + if softcap: + kw["logit_softcap"] = float(softcap) return rt.RuntimeArtifact(metadata={ "target": "rocm", "compiler_path": "rocm_flash_attn_bwd_compiled", "executable": True, "execution_kind": "native_gpu", "arg_names": names, "output_name": "g", "ops": [{"op_name": "tessera.flash_attn_bwd", "result": "g", - "operands": names, - "kwargs": {"scale": scale, "causal": causal}}]}) + "operands": names, "kwargs": kw}]}) -def _run(rt, dO, Q, K, V, *, causal, scale, bias=None): +def _run(rt, dO, Q, K, V, *, causal, scale, bias=None, window=0, softcap=0.0): call = (dO, Q, K, V) if bias is None else (dO, Q, K, V, bias) - res = rt.launch(_art(rt, causal, scale, bias is not None), call) + res = rt.launch( + _art(rt, causal, scale, bias is not None, window, softcap), call) assert res["ok"] is True, res.get("reason") assert res["compiler_path"] == "rocm_flash_attn_bwd_compiled" dQ, dK, dV = res["output"] @@ -196,6 +213,72 @@ def test_bwd_runtime_lane_bias_broadcast_per_head(): and _relerr(dV, dV_ref) < tol) +@pytest.mark.parametrize("D,B,H,Sq,Sk,W", [ + (16, 1, 2, 32, 32, 8), + (16, 1, 2, 40, 40, 16), + (64, 1, 2, 48, 48, 24), +]) +def test_bwd_runtime_lane_sliding_window_matches_numpy(D, B, H, Sq, Sk, W): + rt = _rocm_or_skip() + rng = np.random.default_rng(51 + D + Sq + W) + Q = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + K = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.3).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.3).astype(np.float16) + scale = 1.0 / float(np.sqrt(D)) + # window is implicitly causal → reference uses causal=False + window=W. + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, False, window=W) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=False, scale=scale, window=W) + tol = 5e-3 + eQ, eK, eV = _relerr(dQ, dQ_ref), _relerr(dK, dK_ref), _relerr(dV, dV_ref) + assert eQ < tol and eK < tol and eV < tol, ( + f"window rel-err dQ={eQ:.2e} dK={eK:.2e} dV={eV:.2e} " + f"@ D={D} {B}x{H}x{Sq}x{Sk} W={W}") + + +@pytest.mark.parametrize("D,B,H,Sq,Sk,cap,causal", [ + (16, 1, 2, 32, 32, 30.0, False), + (16, 1, 2, 24, 40, 50.0, False), + (64, 1, 2, 32, 32, 20.0, True), +]) +def test_bwd_runtime_lane_softcap_matches_numpy(D, B, H, Sq, Sk, cap, causal): + rt = _rocm_or_skip() + rng = np.random.default_rng(61 + D + Sq + int(cap)) + # larger magnitude so the tanh soft-cap is actually exercised. + Q = (rng.standard_normal((B, H, Sq, D)) * 1.0).astype(np.float16) + K = (rng.standard_normal((B, H, Sk, D)) * 1.0).astype(np.float16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.5).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.5).astype(np.float16) + scale = 1.0 / float(np.sqrt(D)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, causal, softcap=cap) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=causal, scale=scale, softcap=cap) + tol = 8e-3 + eQ, eK, eV = _relerr(dQ, dQ_ref), _relerr(dK, dK_ref), _relerr(dV, dV_ref) + assert eQ < tol and eK < tol and eV < tol, ( + f"softcap rel-err dQ={eQ:.2e} dK={eK:.2e} dV={eV:.2e} " + f"@ D={D} {B}x{H}x{Sq}x{Sk} cap={cap} causal={causal}") + + +def test_bwd_runtime_lane_window_softcap_bias_compose(): + # window + softcap + bias together (also exercises the arg-order plumbing). + rt = _rocm_or_skip() + rng = np.random.default_rng(77) + B, H, Sq, Sk, D, W, cap = 1, 2, 32, 32, 16, 12, 30.0 + Q = (rng.standard_normal((B, H, Sq, D)) * 0.6).astype(np.float16) + K = (rng.standard_normal((B, H, Sk, D)) * 0.6).astype(np.float16) + V = (rng.standard_normal((B, H, Sk, D)) * 0.5).astype(np.float16) + dO = (rng.standard_normal((B, H, Sq, D)) * 0.5).astype(np.float16) + bias = (rng.standard_normal((B, H, Sq, Sk)) * 0.3).astype(np.float32) + scale = 1.0 / float(np.sqrt(D)) + dQ_ref, dK_ref, dV_ref = _ref_bwd(Q, K, V, dO, scale, False, bias=bias, + window=W, softcap=cap) + dQ, dK, dV = _run(rt, dO, Q, K, V, causal=False, scale=scale, bias=bias, + window=W, softcap=cap) + tol = 8e-3 + assert (_relerr(dQ, dQ_ref) < tol and _relerr(dK, dK_ref) < tol + and _relerr(dV, dV_ref) < tol) + + def test_bwd_runtime_lane_rejects_bad_group(): # Query heads not divisible by KV heads → reject clearly (Decision #21). rt = _rocm_or_skip()