diff --git a/scripts/env-doc-allowlist.txt b/scripts/env-doc-allowlist.txt index c3995c167..f5a641ad6 100644 --- a/scripts/env-doc-allowlist.txt +++ b/scripts/env-doc-allowlist.txt @@ -82,6 +82,7 @@ VT_GDN_MERGED_BA VT_GDN_MERGED_PROJ VT_GDN_MERGED_QKVZ VT_GDN_MERGED_QKVZ_FP8 +VT_GDN_NORMGATED_COOP VT_GDN_OCC_BLOCK VT_GDN_OUT_FP8_FUSE VT_GDN_PACKED_DECODE_FP8_TOWER diff --git a/src/vt/rocm/rocm_gdn_fused.hip b/src/vt/rocm/rocm_gdn_fused.hip index 76e64f8fe..65d73c3ef 100644 --- a/src/vt/rocm/rocm_gdn_fused.hip +++ b/src/vt/rocm/rocm_gdn_fused.hip @@ -80,6 +80,52 @@ __global__ void RmsNormGatedK(Tout* out, const Tx* x, const Tx* gate, } } +// The donor kernel runs ONE THREAD PER ROW (`<<>>`): each row walks +// d twice serially — 24 launches/tok x 17.25us = 0.414 ms/tok of pure +// single-thread latency. This arm gives each row a whole 256-thread block: +// strided-per-thread sumsq with a wavefront-shfl reduction (width from +// warpSize), then a strided gated store. The float association CHANGES +// (opt-in flag, adjudication owed like its T6/T8 siblings); no quant +// contract exists on this op. +template +__global__ void RmsNormGatedCoopK(Tout* out, const Tx* x, const Tx* gate, + const Tx* w, int64_t rows, int64_t d, + int64_t gate_group, int64_t gate_outer, float eps, + bool sigmoid_gate) { + const int64_t i = blockIdx.x; // one row per block + if (i >= rows) return; + const int tid = static_cast(threadIdx.x); + const int waveSz = warpSize; + __shared__ float partial[256 / 32]; // sized for the narrowest wavefront + const Tx* xrow = x + i * d; + // Pass 1: strided-per-thread sumsq — lanes touch consecutive addresses, + // which IS the coalesced pattern for a streaming pass. + float acc = 0.0f; + for (int64_t j = tid; j < d; j += 256) { + const float v = Ld(xrow, j); + acc += v * v; + } + for (int off = waveSz / 2; off > 0; off >>= 1) + acc += __shfl_down_sync(0xffffffffULL, acc, off); + if (tid % waveSz == 0) partial[tid / waveSz] = acc; + __syncthreads(); + if (tid == 0) { + float t = 0.0f; + const int nw = 256 / waveSz; + for (int wv = 0; wv < nw; ++wv) t += partial[wv]; + partial[0] = t; + } + __syncthreads(); + const float inv = 1.0f / sqrtf(partial[0] / static_cast(d) + eps); + // Pass 2: gated scaled store, same strided pattern. + const int64_t gbase = (i / gate_group) * gate_outer + (i % gate_group) * d; + for (int64_t j = tid; j < d; j += 256) { + const float z = Ld(gate, gbase + j); + const float act = sigmoid_gate ? Sigmoid(z) : Silu(z); + St(out, i * d + j, Ld(xrow, j) * inv * Ld(w, j) * act); + } +} + // ops.cpp contract: out bf16, attn f32-or-bf16, gate f32 (unrounded sigmoid // input). Tattn only. template @@ -169,7 +215,6 @@ __global__ void AttnQkNormRopeGateK(Tqk* q_out, Tqk* k_out, Tgate* gate_out, } } // namespace - void RmsNormGatedKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor& gate, const Tensor& w, const RmsNormGatedArgs& args) { VT_CHECK(x.dtype == DType::kF32 || x.dtype == DType::kBF16, @@ -183,6 +228,37 @@ void RmsNormGatedKernelRocm(Queue& q, Tensor& out, const Tensor& x, const Tensor const int64_t gate_outer = gate.stride[0]; hipStream_t s = AsStream(q); const unsigned grid = static_cast(t); + // T9 opt-in arm (read PER CALL like the sibling flags): one block per + // row, cooperative reduction. Default OFF keeps the donor kernel. + const char* coop_e = std::getenv("VT_GDN_NORMGATED_COOP"); + const bool coop = coop_e != nullptr && coop_e[0] == '1' && coop_e[1] == '\0'; + if (coop) { + if (x.dtype == DType::kF32) { + if (out.dtype == DType::kF32) { + RmsNormGatedCoopK<<>>( + out.Ptr(), x.Ptr(), gate.Ptr(), w.Ptr(), t, d, + gate_group, gate_outer, args.eps, args.sigmoid_gate); + } else { + RmsNormGatedCoopK<<>>( + out.Ptr<__hip_bfloat16>(), x.Ptr(), gate.Ptr(), w.Ptr(), + t, d, gate_group, gate_outer, args.eps, args.sigmoid_gate); + } + } else { + if (out.dtype == DType::kF32) { + RmsNormGatedCoopK<__hip_bfloat16, float><<>>( + out.Ptr(), x.Ptr<__hip_bfloat16>(), gate.Ptr<__hip_bfloat16>(), + w.Ptr<__hip_bfloat16>(), t, d, gate_group, gate_outer, args.eps, + args.sigmoid_gate); + } else { + RmsNormGatedCoopK<__hip_bfloat16, __hip_bfloat16><<>>( + out.Ptr<__hip_bfloat16>(), x.Ptr<__hip_bfloat16>(), gate.Ptr<__hip_bfloat16>(), + w.Ptr<__hip_bfloat16>(), t, d, gate_group, gate_outer, args.eps, + args.sigmoid_gate); + } + } + Check(hipGetLastError(), "rmsnorm_gated coop launch"); + return; + } if (x.dtype == DType::kF32) { if (out.dtype == DType::kF32) { RmsNormGatedK<<>>( diff --git a/tests/vt/test_rocm_quant_dot.cpp b/tests/vt/test_rocm_quant_dot.cpp index dece62c70..2e71ad369 100644 --- a/tests/vt/test_rocm_quant_dot.cpp +++ b/tests/vt/test_rocm_quant_dot.cpp @@ -32,6 +32,7 @@ #include "vt/quant.h" #include "vt/rocm/rocm_mmvq_policy.h" #include "vt/rocm/rocm_norm_quant_bridge.h" +#include "vt/rocm/rocm_runtime.h" #include "vt/tensor.h" // Compile the exact scratch policy used by the HIP provider with host fakes. @@ -1140,3 +1141,87 @@ TEST_CASE( #endif gpu.DestroyQueue(gq); } + + +// T9 (GFX1100-TG200): cooperative gated-norm remap (VT_GDN_NORMGATED_COOP=1). +// The donor kernel runs ONE THREAD PER ROW; the arm gives each row a +// 256-thread block with a wavefront-shfl reduction. The reduction +// association changes, so outputs may move within float ULPs -- held to an +// NMSE band vs the plain kernel here, with flag-inertness asserted +// byte-level. RED-first: before the arm existed COOP=1 was inert and the +// byte-equality could not witness it; the ULP-band leg is nonzero only +// when the arm ENGAGES, so the pair (inert bytes equal when unset, band +// non-tight failure risk when broken) is the witness. +TEST_CASE("T9 COOP gated-norm: output within ULP band of donor kernel; flag inert when unset") { + if (!vt::rocm::DeviceAvailable()) { + MESSAGE("no AMD GPU on this host; ROCm keep-quant gate skipped"); + return; + } + Backend& gpu = vt::GetBackend(DeviceType::kROCM); + Queue gq = gpu.CreateQueue(); + for (int64_t d : {int64_t{256}, int64_t{2560}}) { + const int64_t rows = 4; + CAPTURE(d); + std::mt19937 rng(0x7C00U + static_cast(d)); + std::vector abf(rows * d), gb(rows * d), gw(d); + for (auto& v : abf) v = vt::F32ToBF16(static_cast(static_cast(rng() % 2001) - 1000) / 500.0F); + for (auto& v : gb) v = vt::F32ToBF16(static_cast(static_cast(rng() % 2001) - 1000) / 500.0F); + for (auto& v : gw) v = vt::F32ToBF16(0.5F); + void* d_a = gpu.Alloc(abf.size() * 2); + void* d_g = gpu.Alloc(gb.size() * 2); + void* d_w = gpu.Alloc(gw.size() * 2); + gpu.Copy(gq, d_a, abf.data(), abf.size() * 2); + gpu.Copy(gq, d_g, gb.data(), gb.size() * 2); + gpu.Copy(gq, d_w, gw.data(), gw.size() * 2); + + auto run = [&](char* dst) { + Tensor xt = DevTensor(d_a, DType::kBF16, {rows, d}); + Tensor gt = DevTensor(d_g, DType::kBF16, {rows, d}); + Tensor wt = DevTensor(d_w, DType::kBF16, {d}); + Tensor ot = DevTensor(dst, DType::kBF16, {rows, d}); + vt::RmsNormGated(gq, ot, xt, gt, wt, vt::RmsNormGatedArgs{1e-6f, false}); + gpu.Synchronize(gq); + }; + std::vector plain(abf.size() * 2), coop(abf.size() * 2); + void* d_o = gpu.Alloc(abf.size() * 2); + { + ::unsetenv("VT_GDN_NORMGATED_COOP"); + run(static_cast(d_o)); + gpu.Copy(gq, plain.data(), d_o, plain.size()); + ::setenv("VT_GDN_NORMGATED_COOP", "1", 1); + run(static_cast(d_o)); + gpu.Copy(gq, coop.data(), d_o, coop.size()); + gpu.Synchronize(gq); + } + double num = 0.0, den = 0.0; + bool identical = true; + for (size_t i = 0; i < abf.size(); ++i) { + const unsigned pb = plain[i * 2] | (plain[i * 2 + 1] << 8); + const unsigned cb = coop[i * 2] | (coop[i * 2 + 1] << 8); + if (pb != cb) identical = false; + const float p = vt::BF16ToF32(static_cast(pb)); + const float c = vt::BF16ToF32(static_cast(cb)); + num += (p - c) * (p - c); + den += p * p; + } + // Informational only: whether the reassociation flips a rounded bit is + // data-dependent. ENGAGEMENT is witnessed by the rocpd kernel symbol in + // the acceptance window, not here. + CAPTURE(identical); + const double nmse = den > 0 ? num / den : 0.0; + CAPTURE(nmse); + CHECK(nmse <= 1e-6); + // Inert leg: flag truly unset reproduces the first run bit-for-bit. + std::vector again(abf.size() * 2); + ::unsetenv("VT_GDN_NORMGATED_COOP"); + run(static_cast(d_o)); + gpu.Copy(gq, again.data(), d_o, again.size()); + gpu.Synchronize(gq); + CHECK(again == plain); + gpu.Free(d_o); + gpu.Free(d_a); + gpu.Free(d_g); + gpu.Free(d_w); + } + gpu.DestroyQueue(gq); +}