Repository navigation
Conversation
- dot2 attention (rocmkernels.h/.hip): K/V as scalar-cache SGPR operands of v_dot2_f32_f16, lazily rescaled online softmax, prep kernel that compacts on-board keys and fuses RoPE; key splits for small batches. Config: rocmUseDot2Attention (auto = archs with v_dot2_f32_f16). - Optional CK DPP GEMM (rocmckgemm.h/.hip, gfx103x builds with CK installed) for FP16 transformer projections and 1x1 convs, chosen per GEMM shape by timing against hipBLAS on first use. Config: rocmUseCkGemm. - Vectorized FP16 NHWC scale/bias/activation and masked residual add, wave-per-row RMSNorm, LDS-tiled batched transpose. gfx1030 (RX 6800 XT) benchmarknn, nnEval/s at batch 8 / 32: tf3-b10c512 443->741 / 535->946, tf3-b11c768 249->386 / 286->449, tf2-b10c384 787->1532 / 1018->2025. testgpuerror within limits.
- With the CK GEMM kernels enabled, Q/K/V (when dot2 attention consumes them) and the FFN linear1/gate projections write one buffer interleaved per token, so each can run as a single CK GEMM; hipBLAS writes the same layout through strided-batched GEMMs with an output leading dimension. SwiGLU then works in place at that row stride and linear2 reads it there. - CK instances use the gridwise DPP GEMM directly with C column-major in CK's terms (16-byte C stores) and an Add variant, so accumulating GEMMs (nbt post convs, outProj/linear2 with requireMaxBoardSize) can use CK too. - CK vs hipBLAS (and which CK instance) is chosen per GEMM shape by timing the first real calls with events instead of isolated back-to-back runs, which ranked candidates differently from the model's interleaved calls. - Vectorized (16-byte) SwiGLU kernel. gfx1030, tf3-b10c512 benchmarknn -require-exact-nnlen ms/batch: b4 7.50->6.2, b8 11.0->9.9, b16 18.6->17.9, b32 34.2->33.0.
- rocmUseDot2Attention=auto now only enables dot2 attention on gfx103x (rocmDot2AttentionAutoArch), the archs it was measured on; other archs with v_dot2_f32_f16 need rocmUseDot2Attention=true. - The vectorized scale/bias, masked residual add, wave-per-row RMSNorm and SwiGLU kernels are ROCm-only, so the CUDA backend runs exactly the kernels it did before. - CK GEMM tuning no longer blocks: it keeps rotating candidates until the last sample's event has completed (hipEventQuery) instead of hipEventSynchronize. Documented the transient/non-reproducible choice and the FP16 rounding of CK's accumulate. - CK GEMM configure probe: compiles all targeted gfx103x archs with the build-type flags, signature covers build type/flags/archs, and reports missing CK headers. - Log messages name the dot2 kernel when CK fused attention falls back to it. - Docs/config: dot2 auto scope, gfx1030-only measurements.
…elow the reference When every score of a row's first key pair (of the whole row, or of a key-split chunk) is more than 8 below the initial reference 0, the rescale moved the reference down with alpha = exp2(-shift) = inf, and 0 * inf turned the still empty l and acc into NaN. Seen as "Got nonfinite for policy sum" with b11c768h12nbt3tflrs-fson-silu. Clamp alpha to 1: it only ever exceeds 1 when nothing has been accumulated yet.
zhaob1n
marked this pull request as ready for review
October 21, 2026 15:12
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On RDNA2 (gfx103x: no matrix cores, no CK FMHA kernels) the ROCm backend ran transformer nets with the generic online-softmax attention kernel (~45-50% of transformer forward time on a gfx1030) and hipBLAS for every projection. This PR adds an RDNA2-specific attention kernel built on
v_dot2_f32_f16, an optional Composable Kernel DPP GEMM path for the FP16 token GEMMs, and a few ROCm-only bandwidth fixes. The attention and GEMM changes are on by default only on gfx103x; other AMD archs get only the ROCm vectorized elementwise kernels (see the table below), and CUDA runs the same kernels as before.Changes
cpp/neuralnet/rocmkernels.h/.hip, new). Each lane owns whole query rows (online softmax stays lane-local), and K/V reach the ALU as scalar (SGPR) operands ofv_dot2_f32_f16loaded through the scalar cache, since LDS/VMEM broadcast tops out well below the dot2 rate. A prep kernel compacts on-board keys (masked boards cost nothing), interleaves V in key pairs, pre-scales K bylog2(e)/sqrt(d)and applies RoPE (table or learnable recompute), replacing Step 3's RoPE launch. Small batches split the keys over up to 4 wave groups merged through LDS. Head dims 32 and 64, GQA by head grouping.rocmUseDot2Attention(auto/true/false). auto = gfx1030-gfx1036 only. Other archs withv_dot2_f32_f16only with= true. Where CK FMHA is attempted, it still runs first; dot2 replaces the plain kernel as its fallback only on the archs where dot2 is enabled.cpp/neuralnet/rocmckgemm.h/.hip, new; CMake section "Optional CK DPP GEMM (RDNA2)"). Built only for the gfx103x archs of the build when CK headers are present and a cached configure probe compiles it (same pattern as the CK FMHA probe); otherwiseKATAGO_ROCM_HAS_CK_GEMM=0and nothing changes. Instances are CK's gridwise DPP GEMM with C column-major in CK's terms (16-byte C stores) and Set/Add variants.[tokens, q|k|v]buffer (the existing combined-QKV layout, now also on ROCm), and FFN linear1/gate as[tokens, linear1|gate]with SwiGLU in place; hipBLAS produces the same layouts via strided-batched GEMMs with an output leading dimension.rocmUseCkGemm(auto/true/false), auto = on where built.cudaandrocmhelpers.inc, underKATAGO_GPU_HIP): 16-byte vectorized FP16 scale/bias/activation, masked residual add, SwiGLU (bit-identical arithmetic), wave-per-row RMSNorm (same per-element math, different sum-of-squares order), and an LDS-tiled NCHW/NHWC transpose (the existing ROCm-only transpose, which the MIOpen NCHW fallback made hot).Compiling.md,gtp_example.cfg,analysis_example.cfg, CK FMHA README.Scope by architecture
MatMulLayer::apply(..., inputLd),applySharedInputStridedMatMuls(..., outputLd),shouldCombineQKVrestructured with the CUDA condition unchanged). Not compiled on CUDA here (no CUDA toolchain); please build-check.CI configs (
gfx942;gfx1100, Linux ROCm and Windows TheRock) do not build CK GEMM; a local build forgfx942;gfx1100compiles cleanly.Measurements
RX 6800 XT (gfx1030), ROCm 7.2.4 (CK from the same release), 19x19, FP16,
katago benchmark -v 800 -t 8,16,24,32 -n 10, best visits/s across the thread sweep (including the extra server-thread and batch-size configurationsbenchmarktries on its own):Correctness
On gfx1030:
runnnlayertestspass;testbackendreferencePASSED on tf3-b10c512, tf2-b10c384, b28c512nbt;testgpuerror(19x19, 9x9, andrequireMaxBoardSize=true) within limits for the transformer nets above and two convnets; a standalone harness checks dot2 attention against an FP32 reference with logits scaled up to 1000x. FP16 error is somewhat higher than master's, still far from the limits.testgpuerror -quick -boardsize 19, closest margins (99% / max; the max column is a single worst position and noisy):Both parts add a little: dot2 attention multiplies V by FP16 softmax weights (p <= 256, so 11 significant bits) where the plain kernel keeps them in FP32, and CK's accumulate rounds to FP16 before adding (below).
Caveats / questions
outProj/linear2with no mask, and nbt post convs) rounds the product to FP16 before the add, like the existing masked-board residual path, unlike hipBLASbeta=1. On fulltestgpuerrorwithrequireMaxBoardSize=true(b10c512h8), winrate error mean was 0.0531% with CK accumulate excluded and 0.0552% with it. Without it, exact-board inference (benchmarknn -require-exact-nnlen, b10c512h8) is 13% / 6% / 5% / 4% slower at batch 4 / 8 / 16 / 32.