From aca594c24aef000d29c7b218d6acebf75d850474 Mon Sep 17 00:00:00 2001 From: ghazni Date: Tue, 25 Aug 2026 14:46:49 +0000 Subject: [PATCH] perf(GFX1100-TG200): T6a adds the warp-per-row cooperative GDN scan arm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The donor scan mapping prices decode at ~61us/call on gfx1100: one block per (slot, head), only min(dv, 256) of 256 threads active, each walking its state row's dk SERIALLY through dk-strided private cache lines — ~17 GB/s effective on the model's 1 MiB/call bf16 state stream (24 layers x 60.7us = 1.46 ms/token). GdnScanCoopK keeps the elementwise recurrence identical (decayed still Ld*decay computed per pass, donor vp/o formulas) but maps one WARP per state row: lanes walk ki coalesced, the two dots reduce through a fixed shfl_down tree, rows iterate warp-strided. Opt-in behind VT_GDN_SCAN_COOP=1 (default OFF keeps the donor walk); the reduction-order change makes outputs NMSE-equal, not bit-exact: cross_device passes 24/25 with the one failure pre-existing and identical under both paths, and the engine A/B records the near-tie adjudication. Acceptance A/B interleaved x5: median 73.061 vs 69.820 tok/s (+4.6%), COOP wins all five pairs. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:ox-alpha [omp] Assisted-by: AGENT:OMEN-ALPHA [OMP] --- docs/ENVIRONMENT.md | 1 + src/vt/rocm/rocm_gdn_scan.hip | 126 ++++++++++++++++++++++++++++++---- 2 files changed, 113 insertions(+), 14 deletions(-) diff --git a/docs/ENVIRONMENT.md b/docs/ENVIRONMENT.md index 96d2f7231..83118f0d0 100644 --- a/docs/ENVIRONMENT.md +++ b/docs/ENVIRONMENT.md @@ -226,6 +226,7 @@ portable/reference path. In normal operation leave them unset. | `VT_ROCM_QUANT_WMMA` | on | ROCm keep-quant Q6_K and Q4_K prefill GEMM: the RDNA4 rocWMMA int8 tile arm (`KQuantGemmKWmmaQ6K`, `KQuantGemmKWmmaQ4K`), queue-device-resolved `gfx1200`/`gfx1201` only and only when `m >= 16 && n >= 16` (at least one full 16-wide tile in each dimension, not exact alignment — the scalar `KQuantGemmK` remaps its launched index to enumerate only the remainder the WMMA corner leaves untouched (the bottom strip past the aligned row boundary, then the right strip past the aligned column boundary), so a non-16-multiple M/N still takes the WMMA arm over its floor(M/16)xfloor(N/16) corner without paying a full-`m*n`-grid launch for the fill). `0` forces the scalar `Dp4a` arm this row's spec (`KERNEL-QUANT-CIQ-GEMM-ROCM-RDNA4`, issue #2109) is chasing a performance gap against, on every architecture and every shape. Bit-identical to the scalar arm by construction (the WMMA tile's raw int8 dot is scaled and reduced in the same integer arithmetic the scalar path uses; only the one f32 scale product per superblock is shared with it) — this is a same-binary A/B, not a correctness fallback | | `VT_GEMV_MMVQ` | off | Exact `1` routes dense `m == 1` Q4_K, Q5_K, and Q6_K ROCm keep-quant calls through the MMVQ GEMV arm. The arm folds activation quantization into the GEMV prologue when the output row count and scratch size fit the fold limits. Unset and every other value keep the baseline route. `VT_GEMV_MMVQ_FOLD_MAX` sets the inclusive output-row limit. Its default is `512`; empty, malformed, nonpositive, and overflowing values use the default. Both values are read for each call, so graph capture records the selected route. | | `VT_NORM_QUANT_FUSED` | off | Exact `1` lets a contiguous ROCm `RmsNorm` f32 or bf16 output produce Q8_K activation scratch for one immediately matching dense Q4_K, Q5_K, or Q6_K `MatmulBTQuant` call on the same queue and host thread. Unset and every other value keep standalone activation quantization. A mismatch consumes the one-shot token and keeps the standalone route. | +| `VT_GDN_SCAN_COOP` | off | `=1` selects the warp-per-row cooperative GDN scan arm (`GdnScanCoopK`) on ROCm, mapping one warp per output row instead of the donor walk (`GdnScanK`); read once per process like the sibling arms | | `VT_GDN_PACKED_DECODE` | on (CUDA GDN) | Unpacked GDN decode path | | `VT_GDN_DECODE_BV` | `32` (CUDA GDN decode experiment) | Exact `16` selects the byte-identical 16-value fused-recurrence tile; unset and every other spelling keep the 32-value schedule. Experimental opt-in; no release or cross-hardware default change | | `VT_GDN_DECODE_SWIZZLE` | `0` (CUDA GDN decode experiment) | Exact `1` enables the shared-memory bank swizzle only for the `BV=16`, `Dv=Dk=128`, eight-lane production geometry; all other values and shapes keep the incumbent layout | diff --git a/src/vt/rocm/rocm_gdn_scan.hip b/src/vt/rocm/rocm_gdn_scan.hip index 4c1bd90a0..6212f8026 100644 --- a/src/vt/rocm/rocm_gdn_scan.hip +++ b/src/vt/rocm/rocm_gdn_scan.hip @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -103,6 +104,78 @@ __global__ void GdnScanK(Tout* out, const Tin* q, const Tin* k, const Tin* v, __syncthreads(); } } +// ── Cooperative scan (TG200 T6a, opt-in VT_GDN_SCAN_COOP=1) ───────────────── +// Same recurrence, different WORK MAPPING: one WARP per state row instead of +// one thread. The donor mapping prices decode at ~61us/call on gfx1100: only +// min(dv, blockDim) of 256 threads do anything, each walking its row's dk +// SERIALLY through dk*sizeof(TState)-strided addresses, so every access is a +// private cache line and the walk is latency-bound (~17 GB/s effective on a +// 1 MiB/call state stream). Here lanes walk ki COALESCED across the row, +// the two dots reduce through a fixed shfl_down tree, and rows iterate +// warp-strided so every wavefront of the state read/write is dense. +// NUMERIC CONTRACT: the elementwise math is unchanged — decayed is still +// Ld*decay computed twice (pass 1 dot, pass 2 update), vp/o formulas are the +// donor's — but the two dot reductions change ORDER (tree vs serial), so +// outputs are NMSE-equal, not bit-exact: this op is cross_device NMSE-gated, +// and engine-level A/B records near-tie adjudication per campaign doctrine. +template +__global__ void GdnScanCoopK(Tout* out, const Tin* q, const Tin* k, + const Tin* v, const float* g, const float* beta, + TState* state, const int32_t* qsl, + const int32_t* state_idx, int64_t state_slots, + int64_t hk_n, int64_t dk, int64_t hv_n, int64_t dv, + float scale) { + constexpr int NWARPS = kBlock / 32; + const int64_t s = blockIdx.y; + const int64_t hv = blockIdx.x; + const int64_t hk = hv / (hv_n / hk_n); + const int64_t state_slot = state_idx != nullptr ? state_idx[s] : s; + if (state_slot < 0 || state_slot >= state_slots) { + const int64_t begin = qsl != nullptr ? qsl[s] : s; + const int64_t end = qsl != nullptr ? qsl[s + 1] : s + 1; + for (int64_t t = begin; t < end; ++t) + for (int64_t vi = threadIdx.x; vi < dv; vi += blockDim.x) + St(out, (t * hv_n + hv) * dv + vi, 0.0f); + return; + } + extern __shared__ float smem[]; // [dk] q' then [dk] k + float* q_sh = smem; + float* k_sh = smem + dk; + TState* s_head = state + (state_slot * hv_n + hv) * dv * dk; + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + const int64_t begin = qsl != nullptr ? qsl[s] : s; + const int64_t end = qsl != nullptr ? qsl[s + 1] : s + 1; + for (int64_t t = begin; t < end; ++t) { + for (int64_t i = threadIdx.x; i < dk; i += blockDim.x) { + q_sh[i] = Ld(q, (t * hk_n + hk) * dk + i) * scale; + k_sh[i] = Ld(k, (t * hk_n + hk) * dk + i); + } + __syncthreads(); + const float decay = expf(g[t * hv_n + hv]); + const float beta_t = beta[t * hv_n + hv]; + for (int64_t vi = warp; vi < dv; vi += NWARPS) { + TState* s_row = s_head + vi * dk; + float dot = 0.0f; + for (int64_t ki = lane; ki < dk; ki += 32) + dot += Ld(s_row, ki) * decay * k_sh[ki]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) dot += __shfl_down(dot, off); + dot = __shfl(dot, 0); + const float vp = (Ld(v, (t * hv_n + hv) * dv + vi) - dot) * beta_t; + float o = 0.0f; + for (int64_t ki = lane; ki < dk; ki += 32) { + const float updated = Ld(s_row, ki) * decay + vp * k_sh[ki]; + St(s_row, ki, updated); + o += updated * q_sh[ki]; + } + #pragma unroll + for (int off = 16; off > 0; off >>= 1) o += __shfl_down(o, off); + if (lane == 0) St(out, (t * hv_n + hv) * dv + vi, o); + } + __syncthreads(); + } +} template void LaunchGdnScanState(hipStream_t s, Tensor& out, const Tensor& q_in, @@ -114,25 +187,51 @@ void LaunchGdnScanState(hipStream_t s, Tensor& out, const Tensor& q_in, const int64_t hv_n = v.shape[1], dv = v.shape[2]; const dim3 grid(static_cast(hv_n), static_cast(n)); const size_t shmem = 2 * static_cast(dk) * sizeof(float); + // T6a opt-in (read once per process like the sibling arms): the + // warp-per-row cooperative mapping. Default OFF keeps the donor walk. + static const bool scan_coop = [] { + const char* e = std::getenv("VT_GDN_SCAN_COOP"); + return e != nullptr && e[0] == '1' && e[1] == '\0'; + }(); if (state.dtype == DType::kF16) { - GdnScanK<<>>( - out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), - g.Ptr(), beta.Ptr(), state.Ptr<__half>(), qsl, state_idx, - state.shape[0], hk_n, dk, hv_n, dv, args.scale); + if (scan_coop) { + GdnScanCoopK<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), + g.Ptr(), beta.Ptr(), state.Ptr<__half>(), qsl, state_idx, + state.shape[0], hk_n, dk, hv_n, dv, args.scale); + } else { + GdnScanK<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), + g.Ptr(), beta.Ptr(), state.Ptr<__half>(), qsl, state_idx, + state.shape[0], hk_n, dk, hv_n, dv, args.scale); + } } else if (state.dtype == DType::kBF16) { - GdnScanK<<>>( - out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), - g.Ptr(), beta.Ptr(), state.Ptr<__hip_bfloat16>(), qsl, - state_idx, state.shape[0], hk_n, dk, hv_n, dv, args.scale); + if (scan_coop) { + GdnScanCoopK<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), + g.Ptr(), beta.Ptr(), state.Ptr<__hip_bfloat16>(), qsl, + state_idx, state.shape[0], hk_n, dk, hv_n, dv, args.scale); + } else { + GdnScanK<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), + g.Ptr(), beta.Ptr(), state.Ptr<__hip_bfloat16>(), qsl, + state_idx, state.shape[0], hk_n, dk, hv_n, dv, args.scale); + } } else { - GdnScanK<<>>( - out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), - g.Ptr(), beta.Ptr(), state.Ptr(), qsl, state_idx, - state.shape[0], hk_n, dk, hv_n, dv, args.scale); + if (scan_coop) { + GdnScanCoopK<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), + g.Ptr(), beta.Ptr(), state.Ptr(), qsl, state_idx, + state.shape[0], hk_n, dk, hv_n, dv, args.scale); + } else { + GdnScanK<<>>( + out.Ptr(), q_in.Ptr(), k.Ptr(), v.Ptr(), + g.Ptr(), beta.Ptr(), state.Ptr(), qsl, state_idx, + state.shape[0], hk_n, dk, hv_n, dv, args.scale); + } } Check(hipGetLastError(), "gdn scan launch"); } - template void LaunchGdnScanIn(hipStream_t s, Tensor& out, const Tensor& q_in, const Tensor& k, const Tensor& v, const Tensor& g, @@ -145,7 +244,6 @@ void LaunchGdnScanIn(hipStream_t s, Tensor& out, const Tensor& q_in, LaunchGdnScanState(s, out, q_in, k, v, g, beta, state, qsl, state_idx, n, args); } - } // namespace void GdnPrefillKernelRocm(Queue& q, Tensor& out, const Tensor& q_in, const Tensor& k,