From e905b4b6cc0bece33a2ba9e17d0afb074fa9434f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:45:05 +0800 Subject: [PATCH] [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 Co-authored-by: Claude Opus 5.5 --- csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu | 63 +++++++- csrc/sm70_turbomind/ops/qpn_pair_sm70.cuh | 148 ++++++++++++++++-- tests/kernels/test_sm70_fused_gdn_ba_batch.py | 47 ++++++ .../layers/quantization/sm70_gdn_ba_verify.py | 19 ++- 4 files changed, 254 insertions(+), 23 deletions(-) create mode 100644 tests/kernels/test_sm70_fused_gdn_ba_batch.py diff --git a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu b/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu index bce357af6..cc033ba7f 100644 --- a/csrc/sm70_turbomind/ops/fp8_qpn8_sm70.cu +++ b/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 + 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( + 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(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(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(z_out.data_ptr()), + reinterpret_cast(ba_weight.data_ptr()), + reinterpret_cast(b_out.data_ptr()), + reinterpret_cast(a_out.data_ptr()), + static_cast(qkv_n), static_cast(ba_n)}; + if (m <= 16) { + vllm::sm70::launch_qpn_pair_m16_fused_ba( + qkv_out, input, codes, group_scales, static_cast(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<<(n / 32 + (m + 7) / 8), 256, 0, stream>>>( + codes.data_ptr(), + reinterpret_cast(group_scales.data_ptr()), + packed_ptr, reinterpret_cast(qkv_out.data_ptr()), + nullptr, static_cast(n), static_cast(k), static_cast(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()); diff --git a/csrc/sm70_turbomind/ops/qpn_pair_sm70.cuh b/csrc/sm70_turbomind/ops/qpn_pair_sm70.cuh index aaf7f9971..ed95a2c2d 100644 --- a/csrc/sm70_turbomind/ops/qpn_pair_sm70.cuh +++ b/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 +__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(group) * m + row) * 16; + a01 = *reinterpret_cast(a); + a23 = *reinterpret_cast(a + 8); + } + if (column < ba.ba_n) { + const half* b = ba.weight + static_cast(column) * k + group * 16; + b01 = *reinterpret_cast(b); + b23 = *reinterpret_cast(b + 8); + } + const unsigned* a0 = reinterpret_cast(&a01); + const unsigned* a1 = reinterpret_cast(&a23); + const unsigned* b0 = reinterpret_cast(&b01); + const unsigned* b1 = reinterpret_cast(&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(out_row) * half_n + out_col] = __float2half(sum); + } else { + ba.a[static_cast(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(row) * ba.qkv_n + column] = value; + } else { + ba.z[static_cast(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 +template __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(storage); + if constexpr (FusedBA) { + const int pairs = width / 64; + if (blockIdx.x >= pairs) { + fused_ba_tile(input, fused_ba, reinterpret_cast(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(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(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(row) * width + tile1 * 32 + e % 32] = + __float2half(up); } - } - if constexpr (!Gated) { - output[static_cast(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 +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; + 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(input.size(0)); + kernel<<>>( + reinterpret_cast(codes.data_ptr()), scales.data_ptr(), a, + reinterpret_cast(out.data_ptr()), nullptr, width, + static_cast(input.size(1)), m, 1.0f, fused_ba); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace vllm::sm70 diff --git a/tests/kernels/test_sm70_fused_gdn_ba_batch.py b/tests/kernels/test_sm70_fused_gdn_ba_batch.py new file mode 100644 index 000000000..e8c702e61 --- /dev/null +++ b/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() diff --git a/vllm/model_executor/layers/quantization/sm70_gdn_ba_verify.py b/vllm/model_executor/layers/quantization/sm70_gdn_ba_verify.py index 35cb408f3..3696fe6d3 100644 --- a/vllm/model_executor/layers/quantization/sm70_gdn_ba_verify.py +++ b/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,