perf(GFX1100-TG200): T9 cooperative gated norm arm

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

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 (`<<<rows, 1>>>`): 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 <typename Tx, typename Tout>
+__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<int>(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<float>(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 <typename Tattn>
@@ -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<unsigned>(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<float, float><<<grid, 256, 0, s>>>(
+            out.Ptr<float>(), x.Ptr<float>(), gate.Ptr<float>(), w.Ptr<float>(), t, d,
+            gate_group, gate_outer, args.eps, args.sigmoid_gate);
+      } else {
+        RmsNormGatedCoopK<float, __hip_bfloat16><<<grid, 256, 0, s>>>(
+            out.Ptr<__hip_bfloat16>(), x.Ptr<float>(), gate.Ptr<float>(), w.Ptr<float>(),
+            t, d, gate_group, gate_outer, args.eps, args.sigmoid_gate);
+      }
+    } else {
+      if (out.dtype == DType::kF32) {
+        RmsNormGatedCoopK<__hip_bfloat16, float><<<grid, 256, 0, s>>>(
+            out.Ptr<float>(), 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><<<grid, 256, 0, s>>>(
+            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<float, float><<<grid, 1, 0, s>>>(

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<unsigned>(d));
+    std::vector<uint16_t> abf(rows * d), gb(rows * d), gw(d);
+    for (auto& v : abf) v = vt::F32ToBF16(static_cast<float>(static_cast<int>(rng() % 2001) - 1000) / 500.0F);
+    for (auto& v : gb) v = vt::F32ToBF16(static_cast<float>(static_cast<int>(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<unsigned char> plain(abf.size() * 2), coop(abf.size() * 2);
+    void* d_o = gpu.Alloc(abf.size() * 2);
+    {
+      ::unsetenv("VT_GDN_NORMGATED_COOP");
+      run(static_cast<char*>(d_o));
+      gpu.Copy(gq, plain.data(), d_o, plain.size());
+      ::setenv("VT_GDN_NORMGATED_COOP", "1", 1);
+      run(static_cast<char*>(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<uint16_t>(pb));
+      const float c = vt::BF16ToF32(static_cast<uint16_t>(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<unsigned char> again(abf.size() * 2);
+    ::unsetenv("VT_GDN_NORMGATED_COOP");
+    run(static_cast<char*>(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);
+}