[Perf][SM70] Run GDN b/a in the QKVZ launch for 9-32 verify rows (#1368)

Above eight rows the GDN input fell back to a cuBLAS GEMM with its reduction,
plus a zero fill and a tail copy: three to four extra launches per layer. Appended
CTAs of the QKVZ launch now compute b/a over the K16 activations and write qkv, z,
b and a directly, as the eight-row route already does.

Co-authored-by: yangzhuxinyzx <Susannah_Carterect@gmx.ca>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu

@@ -493,7 +493,7 @@ struct Fp8PairReader {
 // reduction retain the exact p0 + ... + p(SplitK-1) order. Each projection uses
 // (SplitK / 2 + 1) * 4 KiB of reduction storage and reuses B over all 32 rows.
 template <int SplitK, bool PackedInput = false, bool Gated = false,
-          bool FullRows = false>
+          bool FullRows = false, bool FusedBA = false>
 __global__ __launch_bounds__(
     32 * (SplitK / 2) * (Gated ? 2 : 1),
     FullRows
@@ -503,7 +503,8 @@ __global__ __launch_bounds__(
                                                     const half* __restrict__ input,
                                                     half* __restrict__ output,
                                                     half* __restrict__ k16_output,
-                                                    int n, int k, int m) {
+                                                    int n, int k, int m,
+                                                    vllm::sm70::Sm70FusedBa fused_ba) {
   static_assert(SplitK == 8 || SplitK == 12 || SplitK == 16,
                 "M32 two-phase QPN8 supports split-8, split-12 or split-16");
   if constexpr (FullRows) m = 32;
@@ -518,6 +519,15 @@ __global__ __launch_bounds__(
   constexpr int kProjections = Gated ? 2 : 1;
   __shared__ float reduction_storage[kProjections][kPhysicalWarps + 1]
                                     [kOutputElements];
+  if constexpr (FusedBA) {
+    const int tiles = n / 32;
+    if (blockIdx.x >= tiles) {
+      vllm::sm70::fused_ba_tile<kPhysicalWarps>(
+          input, fused_ba, &reduction_storage[0][0][0],
+          (blockIdx.x - tiles) * 8, k, m);
+      return;
+    }
+  }
 
   const int lane = threadIdx.x & 31;
   const int warp =
@@ -635,8 +645,13 @@ __global__ __launch_bounds__(
         if (output_row < m) {
           if constexpr (Gated) value = (value / (1.0f + __expf(-value))) * up;
           const half result = __float2half(value);
-          output[static_cast<size_t>(output_row) * n + blockIdx.x * 32 +
-                 output_col] = result;
+          if constexpr (FusedBA) {
+            vllm::sm70::fused_ba_store(output, fused_ba, n, output_row,
+                                       blockIdx.x * 32 + output_col, result);
+          } else {
+            output[static_cast<size_t>(output_row) * n + blockIdx.x * 32 +
+                   output_col] = result;
+          }
           if constexpr (Gated) {
             if (k16_output != nullptr) {
               k16_output[vllm::sm70::k16_half_index(
@@ -663,7 +678,8 @@ void launch_fp8_qpn8_m32_twophase_sm70(const uint8_t* codes,
     }
   }
   kernel<<<(n / 32), (32 * kPhysicalWarps * (Gated ? 2 : 1)), 0, stream>>>(
-      codes, channel_scales, input, output, k16_output, n, k, m);
+      codes, channel_scales, input, output, k16_output, n, k, m,
+      vllm::sm70::Sm70FusedBa{});
 }
 
 void launch_fp8_qpn8_ba_split_sm70(const uint8_t* codes,
@@ -1777,7 +1793,8 @@ void fp8_qpn8_dispatch_ba_split_sm70_out(
   TORCH_CHECK(ba_weight.dim() == 2 && ba_weight.size(0) == ba_n &&
                   ba_weight.size(1) == k,
               "fp8_qpn8_dispatch_ba_split_sm70_out: b/a weight mismatch");
-  TORCH_CHECK(m <= 8 || dense_weight_ptr != 0,
+  const bool fused_batch = m > 8 && m <= 32 && k == 5120;
+  TORCH_CHECK(m <= 8 || fused_batch || dense_weight_ptr != 0,
               "fp8_qpn8_dispatch_ba_split_sm70_out: large-M requires the "
               "dense QPN8 workspace");
 
@@ -1815,6 +1832,40 @@ void fp8_qpn8_dispatch_ba_split_sm70_out(
     return;
   }
 
+  if (fused_batch) {
+    // Verify batches: b/a runs in appended CTAs of the QKVZ launch, which also
+    // writes z directly. This drops the GEMM, its reduction and the tail copy.
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    const vllm::sm70::Sm70FusedBa fused{
+        reinterpret_cast<half*>(z_out.data_ptr<at::Half>()),
+        reinterpret_cast<const half*>(ba_weight.data_ptr<at::Half>()),
+        reinterpret_cast<half*>(b_out.data_ptr<at::Half>()),
+        reinterpret_cast<half*>(a_out.data_ptr<at::Half>()),
+        static_cast<int>(qkv_n), static_cast<int>(ba_n)};
+    if (m <= 16) {
+      vllm::sm70::launch_qpn_pair_m16_fused_ba<Fp8PairReader, 16>(
+          qkv_out, input, codes, group_scales, static_cast<int>(n), fused,
+          stream);
+      return;
+    }
+    torch::Tensor packed_input;
+    auto* packed_ptr =
+        vllm::sm70::k16_packed_input(input, packed_input, stream);
+    auto kernel =
+        m == 32
+            ? fp8_qpn8_m32_twophase_sm70_kernel<16, true, false, true, true>
+            : fp8_qpn8_m32_twophase_sm70_kernel<16, true, false, false, true>;
+    kernel<<<static_cast<int>(n / 32 + (m + 7) / 8), 256, 0, stream>>>(
+        codes.data_ptr<uint8_t>(),
+        reinterpret_cast<const half*>(group_scales.data_ptr<at::Half>()),
+        packed_ptr, reinterpret_cast<half*>(qkv_out.data_ptr<at::Half>()),
+        nullptr, static_cast<int>(n), static_cast<int>(k), static_cast<int>(m),
+        fused);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+    return;
+  }
+
   fp8_qpn8_dispatch_sm70_out(qkvz_staging, dense_weight_ptr, input, codes,
                              group_scales, 16, 2, false, false);
   at::mm_out(ba_staging, input, ba_weight.t());

csrc/sm70_turbomind/ops/qpn_pair_sm70.cuh

@@ -6,19 +6,118 @@
 
 namespace vllm::sm70 {
 
+// Where a QKVZ projection writes its z tail and b/a when both run in one launch.
+struct Sm70FusedBa {
+  half* z = nullptr;
+  const half* weight = nullptr;
+  half* b = nullptr;
+  half* a = nullptr;
+  int qkv_n = 0;
+  int ba_n = 0;
+};
+
+#define VLLM_QPN_BA_MMA(C, A0, A1, B0, B1)                          \
+  asm volatile(                                                     \
+      "mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 "            \
+      "{%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, "             \
+      "{%0,%1,%2,%3,%4,%5,%6,%7};\n"                                \
+      : "+f"(C[0]), "+f"(C[1]), "+f"(C[2]), "+f"(C[3]), "+f"(C[4]), \
+        "+f"(C[5]), "+f"(C[6]), "+f"(C[7])                          \
+      : "r"(A0), "r"(A1), "r"(B0), "r"(B1))
+
+// One appended CTA of a QKVZ launch: eight rows of the short dense FP16 b/a
+// projection over the K16 activations, using the launch's shared scratch.
+template <int Warps>
+__device__ __forceinline__ void fused_ba_tile(const half* input,
+                                              const Sm70FusedBa& ba,
+                                              float* partial, int first_row,
+                                              int k, int m) {
+  const int lane = threadIdx.x & 31;
+  const int warp = threadIdx.x >> 5;
+  const int quadpair = (lane >> 2) & 3;
+  const int local_row = (lane & 3) + ((lane & 16) ? 4 : 0);
+  const int row = first_row + local_row;
+  const int column = quadpair * 8 + local_row;
+  const int per_warp = (k >> 4) / Warps;
+  float accum[2][8] = {};
+  for (int group = warp * per_warp; group < (warp + 1) * per_warp; ++group) {
+    uint4 a01 = make_uint4(0, 0, 0, 0), a23 = a01, b01 = a01, b23 = a01;
+    if (row < m) {
+      const half* a = input + (static_cast<size_t>(group) * m + row) * 16;
+      a01 = *reinterpret_cast<const uint4*>(a);
+      a23 = *reinterpret_cast<const uint4*>(a + 8);
+    }
+    if (column < ba.ba_n) {
+      const half* b = ba.weight + static_cast<size_t>(column) * k + group * 16;
+      b01 = *reinterpret_cast<const uint4*>(b);
+      b23 = *reinterpret_cast<const uint4*>(b + 8);
+    }
+    const unsigned* a0 = reinterpret_cast<const unsigned*>(&a01);
+    const unsigned* a1 = reinterpret_cast<const unsigned*>(&a23);
+    const unsigned* b0 = reinterpret_cast<const unsigned*>(&b01);
+    const unsigned* b1 = reinterpret_cast<const unsigned*>(&b23);
+    VLLM_QPN_BA_MMA(accum[0], a0[0], a0[1], b0[0], b0[1]);
+    VLLM_QPN_BA_MMA(accum[1], a0[2], a0[3], b0[2], b0[3]);
+    VLLM_QPN_BA_MMA(accum[0], a1[0], a1[1], b1[0], b1[1]);
+    VLLM_QPN_BA_MMA(accum[1], a1[2], a1[3], b1[2], b1[3]);
+  }
+#pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    const int r = (i & 2) + ((lane & 16) ? 4 : 0) + (lane & 1);
+    const int c = (i & 1) | (((lane >> 1) & 1) << 1) | ((i >> 2) << 2);
+    partial[warp * 256 + r * 32 + quadpair * 8 + c] = accum[0][i] + accum[1][i];
+  }
+  __syncthreads();
+  for (int e = threadIdx.x; e < 256; e += blockDim.x) {
+    const int out_row = first_row + e / 32;
+    const int out_col = e % 32;
+    if (out_row >= m || out_col >= ba.ba_n) continue;
+    float sum = 0.0f;
+#pragma unroll
+    for (int w = 0; w < Warps; ++w) sum += partial[w * 256 + e];
+    const int half_n = ba.ba_n / 2;
+    if (out_col < half_n) {
+      ba.b[static_cast<size_t>(out_row) * half_n + out_col] = __float2half(sum);
+    } else {
+      ba.a[static_cast<size_t>(out_row) * half_n + out_col - half_n] =
+          __float2half(sum);
+    }
+  }
+}
+
+// Stores one QKVZ element into the contiguous qkv or z output.
+__device__ __forceinline__ void fused_ba_store(half* qkv, const Sm70FusedBa& ba,
+                                               int n, int row, int column,
+                                               half value) {
+  if (column < ba.qkv_n) {
+    qkv[static_cast<size_t>(row) * ba.qkv_n + column] = value;
+  } else {
+    ba.z[static_cast<size_t>(row) * (n - ba.qkv_n) + column - ba.qkv_n] = value;
+  }
+}
+
 // Two adjacent projections share each activation fragment. The reader consumes
 // the existing compressed layout; no persistent weight copy is required.
-template <class Reader, int Split, bool Gated>
+template <class Reader, int Split, bool Gated, bool FusedBA = false>
 __global__ void qpn_pair_m16_kernel(const uint8_t* __restrict__ codes,
                                     const void* __restrict__ scales,
                                     const half* __restrict__ input,
                                     half* __restrict__ output,
                                     half* __restrict__ k16_output, int width,
-                                    int k, int m, float global_scale) {
+                                    int k, int m, float global_scale,
+                                    Sm70FusedBa fused_ba = {}) {
   constexpr int kRows = 2;
   constexpr int kElements = kRows * 256;
   extern __shared__ __align__(16) unsigned char storage[];
   auto& partial = *reinterpret_cast<float (*)[2][Split][kElements]>(storage);
+  if constexpr (FusedBA) {
+    const int pairs = width / 64;
+    if (blockIdx.x >= pairs) {
+      fused_ba_tile<Split>(input, fused_ba, reinterpret_cast<float*>(storage),
+                           (blockIdx.x - pairs) * 8, k, m);
+      return;
+    }
+  }
   const int lane = threadIdx.x & 31;
   const int warp = threadIdx.x >> 5;
   const int qp = (lane >> 2) & 3;
@@ -101,15 +200,22 @@ __global__ void qpn_pair_m16_kernel(const uint8_t* __restrict__ codes,
       } else if constexpr (Gated) {
         result = __float2half((sum / (1.0f + __expf(-sum))) * up);
       }
-      output[static_cast<size_t>(row) * width + tile0 * 32 + e % 32] = result;
-      if constexpr (Gated) {
-        if (k16_output != nullptr) {
-          k16_output[k16_half_index(row, tile0 * 32 + e % 32, m)] = result;
+      if constexpr (FusedBA) {
+        fused_ba_store(output, fused_ba, width, row, tile0 * 32 + e % 32,
+                       result);
+        fused_ba_store(output, fused_ba, width, row, tile1 * 32 + e % 32,
+                       __float2half(up));
+      } else {
+        output[static_cast<size_t>(row) * width + tile0 * 32 + e % 32] = result;
+        if constexpr (Gated) {
+          if (k16_output != nullptr) {
+            k16_output[k16_half_index(row, tile0 * 32 + e % 32, m)] = result;
+          }
+        }
+        if constexpr (!Gated) {
+          output[static_cast<size_t>(row) * width + tile1 * 32 + e % 32] =
+              __float2half(up);
         }
-      }
-      if constexpr (!Gated) {
-        output[static_cast<size_t>(row) * width + tile1 * 32 + e % 32] =
-            __float2half(up);
       }
     }
   }
@@ -137,4 +243,26 @@ void launch_qpn_pair_m16(torch::Tensor out, torch::Tensor input,
   if (k16_output != nullptr) k16_publish(out, twin);
 }
 
+// QKVZ pairs plus b/a in one launch. `out` is qkv; the z tail goes to fused_ba.z.
+template <class Reader, int Split>
+void launch_qpn_pair_m16_fused_ba(torch::Tensor out, torch::Tensor input,
+                                  torch::Tensor codes, torch::Tensor scales,
+                                  int width, const Sm70FusedBa& fused_ba,
+                                  cudaStream_t stream) {
+  constexpr int kShared = 2 * Split * 512 * sizeof(float);
+  auto kernel = qpn_pair_m16_kernel<Reader, Split, false, true>;
+  if constexpr (kShared > 48 * 1024) {
+    C10_CUDA_CHECK(cudaFuncSetAttribute(
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kShared));
+  }
+  torch::Tensor packed;
+  auto* a = k16_packed_input(input, packed, stream);
+  const int m = static_cast<int>(input.size(0));
+  kernel<<<width / 64 + (m + 7) / 8, 32 * Split, kShared, stream>>>(
+      reinterpret_cast<const uint8_t*>(codes.data_ptr()), scales.data_ptr(), a,
+      reinterpret_cast<half*>(out.data_ptr<at::Half>()), nullptr, width,
+      static_cast<int>(input.size(1)), m, 1.0f, fused_ba);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
 }  // namespace vllm::sm70

tests/kernels/test_sm70_fused_gdn_ba_batch.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The fused M9-M32 b/a route against the QPN8 projection plus a dense GEMM."""
+
+import pytest
+import torch
+
+K, N, QKV_N, BA_N = 5120, 4096, 2560, 24
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA GPU")
+@pytest.mark.parametrize("rows", [9, 12, 16, 17, 24, 31, 32])
+def test_fused_batch_matches_separate_projections(rows):
+    op = getattr(torch.ops._C, "fp8_qpn8_dispatch_ba_split_sm70_out", None)
+    if op is None or torch.cuda.get_device_capability() != (7, 0):
+        pytest.skip("requires the SM70 native ops")
+    torch.manual_seed(5)
+    device = "cuda"
+    codes = torch.randint(0, 0x60, (K, N), dtype=torch.uint8, device=device)
+    scales = (torch.rand(1, N, device=device) * 0.01 + 0.001).to(torch.float16)
+    ba_weight = (torch.randn(BA_N, K, device=device) * 0.02).to(torch.float16)
+    x = (torch.randn(rows, K, device=device) * 0.5).to(torch.float16)
+    qkv, z, b, a = (
+        x.new_full((rows, n), float("nan")) for n in (QKV_N, N - QKV_N, 12, 12)
+    )
+    op(
+        qkv,
+        z,
+        b,
+        a,
+        x.new_empty((rows, N)),
+        x.new_empty((rows, BA_N)),
+        0,
+        x,
+        codes,
+        scales,
+        ba_weight,
+    )
+    expected = x.new_empty((rows, N))
+    torch.ops._C.fp8_qpn8_gemm_sm70_out(expected, x, codes, scales, 16, 2, True, False)
+    torch.cuda.synchronize()
+    assert torch.equal(qkv, expected[:, :QKV_N])
+    assert torch.equal(z, expected[:, QKV_N:])
+    expected_ba = x.float() @ ba_weight.float().t()
+    error = (torch.cat([b, a], dim=1).float() - expected_ba).abs().max()
+    # One FP16 rounding of the largest value, plus the accumulation order.
+    assert error <= 2 ** -10 * expected_ba.abs().max()

vllm/model_executor/layers/quantization/sm70_gdn_ba_verify.py

@@ -18,7 +18,9 @@ def apply_gdn_ba_verify(layer, x, outputs=None):
         or not isinstance(ba, LinearBase)
         or not current_platform.is_device_capability(70)
         or not getattr(qkvz, "sm70_fp8_qpn8", False)
-        or x.shape != (8, 5120)
+        or x.ndim != 2
+        or x.shape[1] != 5120
+        or not 8 <= x.shape[0] <= 32
         or x.dtype != torch.float16
         or not x.is_contiguous()
         or qkvz.weight.shape != (5120, 4096)
@@ -31,10 +33,11 @@ def apply_gdn_ba_verify(layer, x, outputs=None):
         or not hasattr(torch.ops._C, "fp8_qpn8_dispatch_ba_split_sm70_out")
     ):
         return None
+    rows = x.shape[0]
     if outputs is None:
-        qkv = x.new_empty((8, 2560))
-        z = x.new_empty((8, 1536))
-        b = x.new_empty((8, 12))
+        qkv = x.new_empty((rows, 2560))
+        z = x.new_empty((rows, 1536))
+        b = x.new_empty((rows, 12))
         a = torch.empty_like(b)
     else:
         qkv, z, b, a = outputs
@@ -43,7 +46,9 @@ def apply_gdn_ba_verify(layer, x, outputs=None):
             or tensor.dtype != x.dtype
             or tensor.device != x.device
             or not tensor.is_contiguous()
-            for tensor, shape in zip(outputs, ((8, 2560), (8, 1536), (8, 12), (8, 12)))
+            for tensor, shape in zip(
+                outputs, ((rows, 2560), (rows, 1536), (rows, 12), (rows, 12))
+            )
         ):
             return None
     torch.ops._C.fp8_qpn8_dispatch_ba_split_sm70_out(
@@ -51,8 +56,8 @@ def apply_gdn_ba_verify(layer, x, outputs=None):
         z,
         b,
         a,
-        x.new_empty((8, 4096)),
-        x.new_empty((8, 24)),
+        x.new_empty((rows, 4096)),
+        x.new_empty((rows, 24)),
         0,
         x,
         qkvz.weight,