diff --git a/CMakeLists.txt b/CMakeLists.txt index 48f53249b0..c07dce2779 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -398,6 +398,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA") "${VLLM_SM70_TURBOMIND_ROOT}/ops/tm_registry_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/fp8_qpn8_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/mxfp4_qpn_m1_sm70.cu" + "${VLLM_SM70_TURBOMIND_ROOT}/ops/awq_qpn_m1_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/nvfp4_grouped_decode_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/nvfp4_qpn4_sm70.cu" "${VLLM_SM70_TURBOMIND_ROOT}/ops/qwen38_prefill_cutlass.cu" diff --git a/csrc/ops.h b/csrc/ops.h index 83d33c1def..36edae2f0d 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -518,6 +518,13 @@ void awq_moe_single_token_sm70_out( torch::Tensor inv_permuted_idx, int64_t w13_k, int64_t w13_n, int64_t w2_k, int64_t w2_n, int64_t group_size, int64_t hidden_logical_size); +void awq_moe_qpn_m1_sm70_out(torch::Tensor out, torch::Tensor intermediate, + const torch::Tensor& input, + const torch::Tensor& w13, const torch::Tensor& s13, + const torch::Tensor& w2, const torch::Tensor& s2, + const torch::Tensor& ids, + const torch::Tensor& topk); + void fp8_moe_gemm_sm70_out(torch::Tensor out, torch::Tensor sorted_input, torch::Tensor expert_offsets, torch::Tensor strided_ptrs_w, diff --git a/csrc/sm70_turbomind/ops/awq_qpn_m1_sm70.cu b/csrc/sm70_turbomind/ops/awq_qpn_m1_sm70.cu new file mode 100644 index 0000000000..d28a3e0aa4 --- /dev/null +++ b/csrc/sm70_turbomind/ops/awq_qpn_m1_sm70.cu @@ -0,0 +1,341 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright contributors to the vLLM project +// +// Qwen3.8 TP4 native-g32 AWQ M=1. The quadpair-N m8n8k4 dataflow is +// derived from mxfp4_qpn_m1_sm70.cu / dnv2003/v100-skinny (MIT). +// See adjacent LICENSE.v100-skinny for the retained notice. +// Preserve FP16 scale/bias dequantization and per-route rounding; the +// CTA-local FP32 reduction is numerically, not bitwise, equivalent to +// the legacy TurboMind serial split-K route. + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kK = 2560; +constexpr int kN = 320; +constexpr int kExperts = 512; +constexpr int kRoutes = 10; +constexpr int kSplit = 16; + +// Formats: existing 3B scalar (0) / 4B scale+bias metadata (2). +template +__device__ __forceinline__ uint32_t read_awq_stats(const uint8_t* stats, + int group, int tile, int col, + int n) { + uint32_t bits = 0; + if constexpr (Format == 2) { + return __ldg(reinterpret_cast(stats) + group * n + + tile * 32 + col); + } else { + const auto* record = + stats + (static_cast(group) * n + tile * 32 + col) * 3; + bits = static_cast(__ldg(record)) | + (static_cast(__ldg(record + 1)) << 8) | + (static_cast(__ldg(record + 2)) << 16); + } + const half scale = __ushort_as_half(static_cast(bits)); + const half zero = __int2half_rn(static_cast((bits >> 16) & 0xff)); + const half bias = __hmul(__hneg(zero), scale); + return (bits & 0xffffu) | + (static_cast(__half_as_ushort(bias)) << 16); +} + +__device__ __forceinline__ void dequant_awq_u4x8(uint32_t packed, + uint32_t stats, + half2* decoded) { + // Exact current TurboMind U4 -> half conversion, with the native FP16 bias + // boundary. In particular, do not substitute (q - zero) * scale. + uint32_t h[4]; + const uint32_t upper = __byte_perm(packed, 0, 0x4321); + constexpr uint32_t lut = (0xf0 & 0xcc) | 0xaa; + constexpr uint32_t bottom_mask = 0x000f000f; + constexpr uint32_t top_mask = 0x00f000f0; + constexpr uint32_t magic0 = 0x64006400; + constexpr uint32_t magic1 = 0x54005400; + asm("lop3.b32 %0, %1, %2, %3, %4;" + : "=r"(h[0]) + : "r"(packed), "n"(bottom_mask), "n"(magic0), "n"(lut)); + asm("lop3.b32 %0, %1, %2, %3, %4;" + : "=r"(h[1]) + : "r"(packed), "n"(top_mask), "n"(magic1), "n"(lut)); + asm("lop3.b32 %0, %1, %2, %3, %4;" + : "=r"(h[2]) + : "r"(upper), "n"(bottom_mask), "n"(magic0), "n"(lut)); + asm("lop3.b32 %0, %1, %2, %3, %4;" + : "=r"(h[3]) + : "r"(upper), "n"(top_mask), "n"(magic1), "n"(lut)); + asm("sub.f16x2 %0, %1, %2;" : "=r"(h[0]) : "r"(h[0]), "r"(magic0)); + asm("sub.f16x2 %0, %1, %2;" : "=r"(h[1]) : "r"(h[1]), "r"(magic1)); + asm("sub.f16x2 %0, %1, %2;" : "=r"(h[2]) : "r"(h[2]), "r"(magic0)); + asm("sub.f16x2 %0, %1, %2;" : "=r"(h[3]) : "r"(h[3]), "r"(magic1)); + const half scale = __ushort_as_half(static_cast(stats)); + const half bias = __ushort_as_half(static_cast(stats >> 16)); +#pragma unroll + for (int i = 0; i < 4; ++i) { + decoded[i] = + __hfma2(*reinterpret_cast(&h[i]), + __halves2half2(scale, scale), __halves2half2(bias, bias)); + } +} + +#define QPN_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};" \ + : "+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)) + +template +__global__ void shared_qpn_w13_kernel(const half* input, + const uint32_t* weights, + const uint8_t* metadata, + const int32_t* expert_ids, half* output) { + __shared__ float partials[kSplit][32]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + const int tile = blockIdx.x; + const int route = blockIdx.y; + const int expert = __ldg(expert_ids + route); + if (expert < 0 || expert >= kExperts) { + if (threadIdx.x < 16) + output[route * 160 + tile * 16 + threadIdx.x] = __float2half(0.f); + return; + } + const int quadpair = (lane >> 2) & 3; + const int a_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int packed_col = quadpair * 8 + a_row; + const uint32_t* expert_weights = + weights + static_cast(expert) * kK * kN / 8; + constexpr int stats_bytes = (kK / 32) * kN * (Format == 2 ? 4 : 3); + const uint8_t* expert_stats = + metadata + static_cast(expert) * stats_bytes; + float accum[8] = {}; + uint32_t stats = 0; +#pragma unroll 4 + for (int group = warp * 10; group < warp * 10 + 10; ++group) { + const size_t base = + (static_cast(tile) * (kK / 8) + group * 2) * 32 + packed_col; + const uint32_t packed0 = __ldcs(expert_weights + base); + const uint32_t packed1 = __ldcs(expert_weights + base + 32); + half2 decoded[8]; + + // Each warp begins on an even K16 group. Reuse g32 metadata for both + // halves without changing the common K16/MMA/FP32 accumulation order. + if ((group & 1) == 0) + stats = + read_awq_stats(expert_stats, group / 2, tile, packed_col, kN); + dequant_awq_u4x8(packed0, stats, decoded); + dequant_awq_u4x8(packed1, stats, decoded + 4); + + const auto* b = reinterpret_cast(decoded); + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (a_row == 0) { + input01 = *reinterpret_cast(input + group * 16); + input23 = *reinterpret_cast(input + group * 16 + 8); + } + const auto* a0 = reinterpret_cast(&input01); + const auto* a1 = reinterpret_cast(&input23); + QPN_MMA(accum, a0[0], a0[1], b[0], b[1]); + QPN_MMA(accum, a0[2], a0[3], b[2], b[3]); + QPN_MMA(accum, a1[0], a1[1], b[4], b[5]); + QPN_MMA(accum, a1[2], a1[3], b[6], b[7]); + } + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + partials[warp][quadpair * 8 + col] = accum[index]; + } + } + } + __syncthreads(); + if (warp == 0) { + float value = 0.f; +#pragma unroll + for (int part = 0; part < kSplit; ++part) value += partials[part][lane]; + const half rounded = __float2half(value); + const unsigned rounded_bits = __half_as_ushort(rounded); + const int source_lane = (lane & 15) * 2; + const half gate = __ushort_as_half(static_cast( + __shfl_sync(0xffffffffu, rounded_bits, source_lane))); + const half up = __ushort_as_half(static_cast( + __shfl_sync(0xffffffffu, rounded_bits, source_lane + 1))); + if (lane < 16) { + const float gate_f = __half2float(gate); + const half silu = __float2half(gate_f / (1.f + expf(-gate_f))); + output[route * 160 + tile * 16 + lane] = __hmul(silu, up); + } + } +} + +template +__global__ void shared_qpn_w2_reduce_kernel( + const half* __restrict__ input, const uint32_t* __restrict__ weights, + const uint8_t* __restrict__ metadata, + const int32_t* __restrict__ expert_ids, + const float* __restrict__ topk_weights, half* __restrict__ output) { + constexpr int k = 160; + constexpr int n = 2560; + __shared__ half route_outputs[kRoutes][32]; + const int lane = threadIdx.x & 31; + const int route = threadIdx.x >> 5; + const int tile = blockIdx.x; + const int expert = __ldg(expert_ids + route); + float accum[8] = {}; + if (expert >= 0 && expert < kExperts) { + const int quadpair = (lane >> 2) & 3; + const int a_row = (lane & 3) + ((lane & 16) ? 4 : 0); + const int packed_col = quadpair * 8 + a_row; + const uint32_t* expert_weights = + weights + static_cast(expert) * k * n / 8; + constexpr int bytes = (k / 32) * n * (Format == 2 ? 4 : 3); + const uint8_t* expert_stats = + metadata + static_cast(expert) * bytes; + const half* input_row = input + route * k; + uint32_t stats = 0; +#pragma unroll + for (int group = 0; group < k / 16; ++group) { + const size_t base = + (static_cast(tile) * (k / 8) + group * 2) * 32 + packed_col; + const uint32_t packed0 = __ldcs(expert_weights + base); + const uint32_t packed1 = __ldcs(expert_weights + base + 32); + half2 decoded[8]; + + if ((group & 1) == 0) + stats = read_awq_stats(expert_stats, group / 2, tile, + packed_col, n); + dequant_awq_u4x8(packed0, stats, decoded); + dequant_awq_u4x8(packed1, stats, decoded + 4); + + const auto* b = reinterpret_cast(decoded); + uint4 input01 = make_uint4(0, 0, 0, 0); + uint4 input23 = make_uint4(0, 0, 0, 0); + if (a_row == 0) { + input01 = *reinterpret_cast(input_row + group * 16); + input23 = *reinterpret_cast(input_row + group * 16 + 8); + } + const auto* a0 = reinterpret_cast(&input01); + const auto* a1 = reinterpret_cast(&input23); + QPN_MMA(accum, a0[0], a0[1], b[0], b[1]); + QPN_MMA(accum, a0[2], a0[3], b[2], b[3]); + QPN_MMA(accum, a1[0], a1[1], b[4], b[5]); + QPN_MMA(accum, a1[2], a1[3], b[6], b[7]); + } + if ((lane & 17) == 0) { +#pragma unroll + for (int pair = 0; pair < 2; ++pair) { +#pragma unroll + for (int offset = 0; offset < 2; ++offset) { + const int index = pair * 4 + offset; + const int col = offset | (((lane >> 1) & 1) << 1) | (pair << 2); + route_outputs[route][quadpair * 8 + col] = __float2half(accum[index]); + } + } + } + } else if (lane < 4) { +#pragma unroll + for (int offset = 0; offset < 8; ++offset) { + route_outputs[route][lane * 8 + offset] = __float2half(0.f); + } + } + __syncthreads(); + if (route == 0) { + float weighted = 0.f; +#pragma unroll + for (int selected = 0; selected < kRoutes; ++selected) { + // Preserve original router order and per-route FP16 materialization. + weighted = fmaf(__ldg(topk_weights + selected), + __half2float(route_outputs[selected][lane]), weighted); + } + output[tile * 32 + lane] = __float2half(weighted); + } +} + +template +void launch(const at::Tensor& input, const at::Tensor& w13, + const at::Tensor& s13, const at::Tensor& w2, const at::Tensor& s2, + const at::Tensor& ids, const at::Tensor& topk, + at::Tensor& intermediate, at::Tensor& out, cudaStream_t stream) { + shared_qpn_w13_kernel<<>>( + reinterpret_cast(input.const_data_ptr()), + reinterpret_cast(w13.const_data_ptr()), + reinterpret_cast(s13.const_data_ptr()), + ids.const_data_ptr(), + reinterpret_cast(intermediate.mutable_data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + shared_qpn_w2_reduce_kernel<<<80, 320, 0, stream>>>( + reinterpret_cast(intermediate.const_data_ptr()), + reinterpret_cast(w2.const_data_ptr()), + reinterpret_cast(s2.const_data_ptr()), + ids.const_data_ptr(), topk.const_data_ptr(), + reinterpret_cast(out.mutable_data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void check_tensor(const at::Tensor& tensor, at::ScalarType dtype, + at::IntArrayRef shape, const at::Device& device, + const char* name, uintptr_t alignment = 16) { + TORCH_CHECK(tensor.is_cuda() && tensor.device() == device, + "awq_qpn_m1: ", name, " must be CUDA on the input device"); + TORCH_CHECK(tensor.scalar_type() == dtype && tensor.sizes() == shape && + tensor.is_contiguous(), + "awq_qpn_m1: ", name, " dtype/shape/contiguity mismatch"); + TORCH_CHECK( + reinterpret_cast(tensor.const_data_ptr()) % alignment == 0, + "awq_qpn_m1: ", name, " alignment mismatch"); +} + +} // namespace + +void awq_moe_qpn_m1_sm70_out(at::Tensor out, at::Tensor intermediate, + const at::Tensor& input, const at::Tensor& w13, + const at::Tensor& s13, const at::Tensor& w2, + const at::Tensor& s2, const at::Tensor& ids, + const at::Tensor& topk) { + const auto device = input.device(); + check_tensor(input, at::kHalf, {1, 2560}, device, "input"); + check_tensor(out, at::kHalf, {1, 2560}, device, "output"); + check_tensor(intermediate, at::kHalf, {10, 160}, device, "intermediate"); + check_tensor(w13, at::kInt, {512, 2560, 40}, device, "W13"); + check_tensor(w2, at::kInt, {512, 160, 320}, device, "W2"); + const bool compact = s13.scalar_type() == at::kByte; + if (compact) { + check_tensor(s13, at::kByte, {512, 80, 320, 3}, device, "W13 metadata"); + check_tensor(s2, at::kByte, {512, 5, 2560, 3}, device, "W2 metadata"); + } else { + check_tensor(s13, at::kInt, {512, 80, 320}, device, "W13 metadata"); + check_tensor(s2, at::kInt, {512, 5, 2560}, device, "W2 metadata"); + } + check_tensor(ids, at::kInt, {1, 10}, device, "expert IDs", 4); + check_tensor(topk, at::kFloat, {1, 10}, device, "router weights", 4); + for (const auto* tensor : {&input, &w13, &s13, &w2, &s2, &ids, &topk}) { + at::assert_no_overlap(out, *tensor); + at::assert_no_overlap(intermediate, *tensor); + } + at::assert_no_overlap(out, intermediate); + const c10::cuda::CUDAGuard guard(device); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "awq_qpn_m1 requires SM70"); + const auto stream = at::cuda::getCurrentCUDAStream(); + // Reuse the prepared bank; no load-time repack or additional weight copy. + if (compact) { + launch<0, 0>(input, w13, s13, w2, s2, ids, topk, intermediate, out, stream); + } else { + launch<2, 2>(input, w13, s13, w2, s2, ids, topk, intermediate, out, stream); + } +} diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index 0ebaaf3c0c..3d68006737 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -8405,6 +8405,23 @@ void awq_moe_active_dense_stage_sm70_out( "SM70 AWQ MoE active dense-stage path enabled C++ op reached", input, total_slots, total_slots); + // Reuse the active-segment ABI for the narrow Qwen3.8 TP4 decode shapes. + // Equal trailing offsets describe empty groups; nonempty groups may contain + // multiple rows. Keep offsets-based scheduling (not one-row slot dispatch). + const char* grouped = + std::getenv("VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE"); + const char* exact_w2 = + std::getenv("VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2"); + if ((grouped == nullptr || std::atoi(grouped) != 0) && + (exact_w2 == nullptr || std::atoi(exact_w2) == 0) && group_size == 32 && + total_slots >= 20 && total_slots <= 80 && total_slots % 10 == 0 && + ((k == 2560 && n == 320) || (k == 160 && n == 2560))) { + awq_moe_gemm_sm70_out_impl(out, input, active_expert_offsets, ptrs_w, + ptrs_s, total_slots, k, n, group_size, false, + active_expert_ids, true); + return; + } + for (int segment = 0; segment < static_cast(total_slots); ++segment) { torch::Tensor offsets = active_expert_offsets.narrow(0, segment, 2); torch::Tensor expert_idx = active_expert_ids.narrow(0, segment, 1); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index f487996750..c4d59e3ff3 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -613,7 +613,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def( "awq_moe_active_dense_stage_sm70_out(" "Tensor(a!) out, Tensor input, Tensor permuted_experts_id, " - "Tensor active_expert_offsets, Tensor active_expert_ids, Tensor ptrs_w, " + "Tensor(b!) active_expert_offsets, Tensor(c!) active_expert_ids, Tensor " + "ptrs_w, " "Tensor ptrs_s, int total_slots, int k, int n, int group_size) -> ()"); ops.impl("awq_moe_active_dense_stage_sm70_out", torch::kCUDA, &awq_moe_active_dense_stage_sm70_out); @@ -695,6 +696,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("awq_moe_single_token_sm70_out", torch::kCUDA, &awq_moe_single_token_sm70_out); + ops.def( + "awq_moe_qpn_m1_sm70_out(Tensor(a!) out, Tensor(b!) intermediate, " + "Tensor input, Tensor w13, Tensor s13, Tensor w2, Tensor s2, " + "Tensor ids, Tensor topk) -> ()"); + ops.impl("awq_moe_qpn_m1_sm70_out", torch::kCUDA, &awq_moe_qpn_m1_sm70_out); + ops.def( "fp8_moe_gemm_sm70_out(Tensor(a!) out, Tensor sorted_input, " "Tensor expert_offsets, Tensor strided_ptrs_w, Tensor strided_ptrs_s, " diff --git a/docs/design/sm70_awq_qpn_m1.md b/docs/design/sm70_awq_qpn_m1.md new file mode 100644 index 0000000000..c474a562fe --- /dev/null +++ b/docs/design/sm70_awq_qpn_m1.md @@ -0,0 +1,189 @@ +# SM70 AWQ QPN single-token operator + +## Kernel layer + +`_C::awq_moe_qpn_m1_sm70_out` implements the native-group-32 Qwen3.8 +TP4 routed-expert geometry. This layer registers an inference-only operator; +it does not select it in the model runtime or change any default route. + +The quadpair-N Tensor Core dataflow is derived from the existing NVFP4 QPN +implementation and its retained `LICENSE.v100-skinny` notice. This is a +model-specific AWQ adaptation, not an enablement of a generic Skinny backend. + +The two launches are: + +1. W13: selected experts in original router order, CTA-local FP32 reduction, + FP16 gate/up materialization, then FP16 SwiGLU intermediate output. +2. W2: FP32 dot products, per-route FP16 output materialization, then ordered + FP32 router-weight accumulation into the FP16 output. + +No selected-weight bank, input replication, checkpoint rewrite or persistent +weight copy is added. The operator consumes the existing prepared banks and +caller-owned intermediate/output buffers. + +### Admission contract + +| Argument | Shape | Type | +| --- | --- | --- | +| Input / output | `(1, 2560)` | FP16 | +| Intermediate | `(10, 160)` | FP16 | +| W13 prepared weight | `(512, 2560, 40)` | INT32 | +| W2 prepared weight | `(512, 160, 320)` | INT32 | +| W13 4-byte metadata | `(512, 80, 320)` | INT32 | +| W2 4-byte metadata | `(512, 5, 2560)` | INT32 | +| W13 3-byte metadata | `(512, 80, 320, 3)` | UINT8 | +| W2 3-byte metadata | `(512, 5, 2560, 3)` | UINT8 | +| Expert IDs / router weights | `(1, 10)` | INT32 / FP32 | + +All arguments must be contiguous on the same SM70 CUDA device. Weight, +metadata and activation pointers require 16-byte alignment; IDs and router +weights require 4-byte alignment. Output/intermediate must not overlap each +other or any input. Negative or out-of-range expert IDs contribute zero; +duplicate valid expert IDs retain their separate router weights. + +The 4-byte layout stores the existing FP16 scale and rounded FP16 bias. +The 3-byte layout stores a FP16 scale and UINT8 zero point, read scalarly in +this layer. Bias is reconstructed at the same FP16 boundary. Dequantization +retains `half_fma(q, scale, half(-zero * scale))`; replacing this with +`half((q - zero) * scale)` is not an equivalent rounding contract. + +## Numerical boundary and tests + +The acceptance objective is **controlled additional numerical perturbation**. +Both implementations approximate the same computation; legacy output is not +ground truth. Holding the AWQ checkpoint fixed separates the arithmetic +change from the quantization error already present in both paths. Compare +each path against the same independently decoded-weight/FP64 reference, +then assess their paired differences and task-level behavior. This does not +measure either path's error against the original unquantized model. +Passing the retained local bounds means both paths have bounded error on +the tested inputs, not that their errors are identical or cancel. Whole-model +quality and amplification across states/routing remain separate checks. + +The CTA-local reduction changes FP32 summation order relative to the legacy +TurboMind split-K route. Bitwise equality to legacy AWQ is not promised, and +neither path is declared the mathematical reference merely because it existed +first. Full-model acceptance must examine fixed-prefix raw logits and paired +quality, separately from speed and free-running token-stream equality. + +Run the portable prepared-layout test on a V100 native build: + +```bash +.venv/bin/python -m pytest -q tests/kernels/test_sm70_awq_qpn_m1.py +``` + +It independently constructs prepared metadata/weight tiles for both layouts, +checks one-hot reads across K/group boundaries, an FP64 W2 dot reference with +explicit FP16 rounding allowance, changing CUDA Graph inputs, duplicate and +invalid expert IDs, aliased/misaligned arguments, and the registered fake op. +It requires SM70; a CPU skip is not a GPU test pass. Shape-specific kernel +tests do not by themselves establish full-model quality or throughput. + +## Runtime layer + +`VLLM_SM70_AWQ_QWEN38_QPN_M1=1` opts in at model initialization. The default +is `0`; other values are rejected. Use a native build containing the operator +and restart the engine after changing the setting, including for rollback. +Changing an environment variable does not replace an already captured graph. +There is no research sidecar, runtime compilation or external DSO loader. + +Admission requires the existing TP4/E512/native-group-32 geometry, batched +TurboMind weights, interleaved W13 and the legacy single-token compact path. +Both the checkpoint and prepared group sizes must be 32. An explicit opt-in +with an unsupported layer contract or missing native operator fails closed. +At execution, only contiguous FP16 `(1, 2560)` inputs with ten INT32 expert +IDs and FP32 router weights select QPN. Other physical batch sizes, including +padded CUDA Graph batches, retain their existing route. This does not change +grouped decode, prefill, attention, shared experts, router selection or MTP. +The existing prepared banks and per-call buffers are passed by reference. + +Run CPU admission and neighboring dispatch tests: + +```bash +.venv/bin/python -m pytest -q \ + tests/quantization/test_awq_qpn_sm70.py \ + tests/quantization/test_sm70_awq_active_grouped_decode.py \ + tests/quantization/test_sm70_awq_indexed_prefill.py \ + tests/quantization/test_sm70_awq_compact_metadata.py +``` + +### Validation snapshot and limitations + +The September 5-6, 2026 investigation used four V100s, TP4/MTP0, the same +native-group-32 AWQ checkpoint and frozen prompt token IDs, FP16 activations +and KV, 4-byte metadata, prefix caching off, and `ignore_eos=false`. The +runtime inherited the separate grouped-decode change and used a separate +QSA page4 logical-order fix in both arms. Those are baseline dependencies, +not changes made by this proposal. Raw artifacts distinguish the built core, +Python sources, QSA extension, checkpoint, tokenizer and evaluation tools. + +Uninstrumented full-model cells were each measured once with an output limit +of 320 tokens; these are observations, not confidence intervals. The metric +below is aggregate **pure-decode** throughput, excluding prefill overlap: + +| Workload | QPN off (tok/s) | QPN on (tok/s) | +| --- | --- | --- | +| C1 x 64K | 50.0326 | 59.0313 | +| C4 x 64K | 131.3485 | 129.5357 | +| C8 x 16K | 246.3619 | 245.7809 | + +The separate NVFP4 reference was 60.2609 tok/s for C1 x 64K: about a 2.04% +gap, not exact parity. All these full-model numbers use 4-byte metadata; +they must not be attributed to cooperative 3-byte metadata loading. The +older 54.6544 tok/s baseline and its 68.3180 tok/s (+25%) research target +remain distinct from this paired experiment; that target is not achieved. +Instrumented dual-path diagnostic timings are not performance measurements. + +The initial cross-process quality pair had an IFEval pass-to-fail change +(4/5 to 3/5); its unchanged-route controls also differed, so it could not +isolate QPN as the cause. A later same-runtime 65-case pair recorded: + +| Evaluation subset | QPN off | QPN on | +| --- | --- | --- | +| HumanEval | 5/5 | 5/5 | +| MBPP | 4/5 | 4/5 | +| IFEval strict | 3/5 | 4/5 | +| GSM8K | 29/32 | 28/32 | +| Tool selection | 10/12 | 10/12 | +| Needle retrieval | 6/6 | 6/6 | + +There was one GSM regression and one IFEval improvement; 47/65 output token +streams were exact. This is **not** a zero-regression or statistically +non-inferior quality result. The original budgets and failed cases remain +part of the evidence, rather than being replaced by the focused diagnostic. + +For attribution, the focused same-process/same-graph experiment executed +both local MoE implementations on identical inputs/routes and selected the +returned output with a device flag. Shapes were prewarmed and GEMM LUTs +remained fixed. Same-arm repeats and unchanged-route C4 controls were exact. +Independent checkpoint/FP64 checks covered 1,152 captured layer/rank/arm +samples at three fixed prefixes; both paths' W13 and W2 passed the retained +rounding bounds. Neither implementation was used as the other's oracle. + +The GSM first divergence occurred at index 120, where the legacy logits for +token IDs 4003 and 16526 were exactly tied at 21.59375. The focused legacy +run also reproduced the earlier QPN failed answer at the unchanged 256-token +budget. This supports numerical trajectory variation, not a demonstrated +kernel defect. It does not make a truncated or incorrect answer acceptable +by definition. + +Teacher forcing prevents different sampled prefixes from confounding the +comparison; it does not force hidden states or expert membership to match. +At the captured 64K step, 13/48 layers changed expert membership while each +selection respected its own scores. The first change reversed a 0.01171875 +score margin. Global logit differences can therefore be larger than local +rounding errors. The focused 64K maximum was 1.751953, and an earlier full +trace reached 5.052734; these are not described as a few ULPs. Not every +historical worst-logit step was captured for independent local replay. + +The analysis tool was also corrected to use `argmax`, not the first index +from `topk(2)`, for the greedy tie rule. Reanalysis found an IFEval flip at +index 80 where QPN's top two logits tie; an earlier short-prompt flip was +a tie-order reporting artifact. Raw tensors and task scores did not change. + +Current scope is an opt-in review proposal. The tested local numerical +evidence does not justify rewriting the kernel solely to reproduce legacy +tokens, but broad quality non-inferiority and production readiness remain +unproven. Future acceptance must retain task-level quality checks alongside +numerical bounds; neither single-question changes nor their aggregate +cancellation alone settle that decision. diff --git a/docs/design/sm70_qwen38_awq_active_grouped_decode.md b/docs/design/sm70_qwen38_awq_active_grouped_decode.md new file mode 100644 index 0000000000..f3be0ea1a2 --- /dev/null +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -0,0 +1,225 @@ +# Qwen3.8 SM70 AWQ active grouped decode + +## Scope and implementation + +This is a narrow replacement for the larger implementation in +[1CatAI/1Cat-vLLM#491](https://github.com/1CatAI/1Cat-vLLM/pull/491), +which was closed over maintenance cost versus measured benefit, not a +correctness failure. It reuses `awq_moe_active_dense_stage_sm70_out`, its +existing active-segment builder, and the existing grouped GEMM implementation. +There is no new public operator, Python operator wrapper, or GEMM kernel. + +Admission requires SM70 AWQ, TP4, 512 experts, top-k 10, effective group size 32, +hidden size 2560, W13 `(K,N)=(2560,320)`, W2 `(K,N)=(160,2560)`, and 2–8 +input tokens. Runtime and pre-capture warmup share the admission policy. +GPU validation uses a native-g32 checkpoint; remapped checkpoints are not a +separate validated quality claim. Single-token and unmatched shapes retain +their existing routes. Explicit +dense/exact diagnostic routes and the existing decode token cap take priority. + +Repeated expert IDs form multi-row segments; unused trailing offsets describe +empty segments. The scheduler remains offsets-based, not a one-row-per-route +NVFP4 dispatch. Scratch offsets and IDs are marked mutable in the existing +Torch schema. Rebuild the extension with the Python changes: an older binary +can still expose the same operator name without the new implementation. + +The model-specific grouped route defaults on for admitted contracts. To roll +back, set this before starting a fresh engine: + +```bash +export VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE=0 +``` + +## Optional existing autotune + +For the validated Qwen3.8 TP4 deployment, the candidate configuration is: + +```bash +export VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE=1 +export VLLM_SM70_AWQ_TUNE_SMALL_SHAPES=1 +``` + +`VLLM_SM70_AWQ_TUNE_SMALL_SHAPES` already exists and remains **off by default**. +This PR does not change the global tuning, FP16 reduction, NCCL, or scheduler +defaults. The tuning flag is broader than the narrow admission gate, so its +results must not be generalized to all AWQ models. Retain the existing +preserve-default-splits settings. For a clean comparison, start independent +engines with private empty caches and no imported GEMM LUT; toggling a flag +inside a process does not invalidate previously selected tactics. Roll tuning +back separately with `VLLM_SM70_AWQ_TUNE_SMALL_SHAPES=0` before startup. + +## Measurement contract and separate attribution + +Evidence was collected on four V100 PCIe 32 GB GPUs, Qwen3.8 Flash-Next AWQ +g32, TP4/MTP0, FP16 activations and KV, maximum length 131328, eight maximum +sequences, 8192 batched tokens, chunked prefill on, prefix caching and async +scheduling off, GPU memory utilization 0.89, language-only MRv2 +`FULL_AND_PIECEWISE`. All arms use the same QSA page4 logical-order correction +tracked separately in +[1CatAI/1Cat-vLLM#494](https://github.com/1CatAI/1Cat-vLLM/pull/494). +That attention correction is not included here. Runtime results are from the +frozen `fbcef6e2f9`-based validation stack, not a retest of later upstream main. + +All measured arms also freeze the existing `VLLM_SM70_QWEN38_FP16_GEMV=1`, +`VLLM_SM70_QWEN38_FUSED_HC_FP16=1`, +`VLLM_SM70_QWEN38_FUSED_GDN_INPUT_FP16=1`, and +`VLLM_SM70_QSA_INDEXER_CUBLAS=1` opt-ins. Runtime-lossy online QPN8 is off. +Actual collectives on this PCIe validation stack use PyNCCL; the later +fully-connected custom-AR HC-sharding results from SXM2 systems do not imply +that the same optimized route was active here. + +Each cell has a 16-token warmup and one scored request batch with a 320-token +output limit, `ignore_eos=false`, `min_tokens=0`, greedy decoding, and frozen +prompt token IDs. The numbered-list prompts request enough output to reach +steady state without suppressing EOS. All scored requests reached 320 tokens. +Pure aggregate throughput counts token deliveries strictly inside the common +window after every request has started decoding and before any request ends. +Each request contributes 319/294/304 tokens for C1/C4/C8 respectively; there +were no multi-token deliveries. This is neither per-request speed nor E2E. + +### Grouped route alone: matched r3 comparison + +Autotune is off in both AWQ arms. These are the implementation-only gains. + +| Cell | Grouped OFF tok/s | Grouped ON tok/s | Gain | NVFP4 tok/s | +|---|---:|---:|---:|---:| +| C1×64K | 48.65 | 48.60 | −0.09% | 54.65 | +| C4×64K | 110.02 | 116.20 | +5.62% | 130.91 | +| C8×16K | 207.29 | 211.94 | +2.24% | 237.58 | + +### Existing autotune: subsequent fresh-process comparison + +Grouped decode is on in both arms. This isolates the configuration benefit, +not additional code added by this PR. + +| Cell | Tuning OFF tok/s | Tuning ON tok/s | Gain | Historical r3 NVFP4 tok/s | +|---|---:|---:|---:|---:| +| C1×64K | 48.7874 | 48.5898 | −0.40% | 54.6544 | +| C4×64K | 116.1290 | 129.2741 | +11.32% | 130.9051 | +| C8×16K | 212.0154 | 242.7240 | +14.48% | 237.5780 | + +The new OFF baseline differs from the previous grouped ON by only ++0.38%/−0.06%/+0.04%. NVFP4 was not rerun as a third arm in this second +experiment; one score per cell does not establish statistical superiority. +Do not add percentages from the two experiments or extrapolate to C8×64K. + +| Cell | Prefill+mixed seconds OFF→ON | Pure-window seconds OFF→ON | E2E seconds OFF→ON | +|---|---:|---:|---:| +| C1×64K | 15.4253→15.4243 | 6.5386→6.5652 | 21.9639→21.9895 | +| C4×64K | 83.8265→83.5213 | 10.1267→9.0970 | 94.6550→93.2887 | +| C8×16K | 32.8447→32.8341 | 11.4709→10.0196 | 44.8195→43.3021 | + +E2E also includes the final drain. C4/C8 E2E duration falls by only +1.44%/3.39%; mixed-phase mean ITL remains about 2.868/2.092 seconds. This is +not a fix for long-prefill interference or a scheduler-policy change. + +## Kernel evidence, quality and costs + +After scoring, each tuning arm ran a separate C4 diagnostic profile. All eight +rank traces contain 16 CPU execute annotations, 16 GPU execute annotations, +16 CUDA Graph replays and 48 W13/W2 pairs per step, with four generation +requests and zero prefill requests. W13 changes from M64×128×32 to M8×256×64 +on all ranks. Its mean per-rank GPU time falls from 5.531 to 2.109 ms/step; +W2 stays near 1.31 ms, FP16 projections/HC near 10.42 ms, and QSA near +3.49 ms. W2 launch grids are not identical across all ranks. These are +profiled kernel durations, not new unprofiled ITLs or a sum across ranks. + +- Grouped-only r3: 54 short quality requests across three arms passed basic + answer checks and stopped naturally at EOS. AWQ OFF/ON complete token IDs + matched in 17/18 cases; the remaining arithmetic wording was `and` versus + `+`, with the same answer 156. +- Tuning A/B: 36/36 short requests passed basic answer checks, stopped at EOS, + and had finite recorded logprobs. C4 self-repeat matched complete token IDs + and recorded top-5 logprobs in 4/4 cases in each arm. Across arms 16/18 + complete token sequences matched; two C8 responses differed in wording. + Both new arms also differ from the older grouped-ON reference in 2/18 + cases, so fresh-process baseline drift precludes assigning every change + exclusively to tuning. +- Real-weight operator replay showed small W13 rounding differences; a + separate tuning probe reached maximum absolute difference 0.0009765625. + This is not proof of model equivalence. No new functional failure was + observed in the bounded prompts, but bitwise equality and broad quality + acceptance are **not** claimed. +- The precision-sensitive cross-batch/process issue is recorded and deferred. + This PR deliberately does not change HC/NCCL precision or attempt a general + numerical-determinism repair. It is separate from QSA ordering correctness. +- Tuning OFF/ON both report 386,392 KV tokens, 509 blocks and 5,210,075,136 + KV tensor bytes per rank, and 0.37 GiB graph memory. Initialization-to-ready + took 431.090/427.330 seconds. There was no observed extra total startup or + KV-capacity cost in this pair; host cache variability and full memory peaks + were not controlled well enough to claim startup acceleration or identical + peak memory. + +## Validation and follow-up + +The measured source checkpoint is `5cceeaad89d6ead1474ca834afd9aaf3a7bd413c`. +The extension built successfully; prior validation retained 51 GPU-directed +test passes and 32 dynamic-route comparisons, including repeated experts and +graph replay. The focused CPU policy/warmup regression command is: + +```bash +.venv/bin/python -m pytest -q tests/quantization/test_sm70_awq_active_grouped_decode.py +``` + +The natural-EOS, per-token, startup/LUT and per-rank trace artifacts are +retained under experiment IDs `awq-narrow-qsa-fixed-ab-20260904` (r3), +`awq-nvfp4-decode-gap-profile-20260904` (r3), and +`awq-autotune-model-ab-20260905`. No production deployment is implied. + +For fork review, the patch was ported onto synchronized main +`755baae1d075ee04fa9096b23fc0225b23589a86`. Conflict resolution preserves the +new indexed-prefill admission and compact-metadata initialization alongside +the grouped-decode flag. Added boundary tests verify that indexed prefill and +grouped decode never both admit the same token count. The original validated +branch is retained. The new base also changes HC, NVFP4 dispatch and scratch +lifetimes: the historical GPU results above are **not** a GPU acceptance of +this new integrated tree. The separate integration run below closes the bounded +model acceptance gate; final human review remains required before promotion. + +### Aligned-tree integration acceptance (2026-09-05, r2) + +The native extensions were rebuilt from `b50bb1f5037d` on base `755baae1d075`. +Both formats additionally use the separate QSA ordering port `05c46b2cd0` and +its matching Flash-V100 Python package and rebuilt extension. No production +changes were made after these binaries were built. Each AWQ cell was scored +once with grouped decode and the optional existing autotune enabled; shared +expert overlap was disabled. This is integration acceptance, not a new matched +grouped-OFF/ON experiment or evidence that all changes come from this PR. + +| AWQ cell | Pure aggregate tok/s | Mean pure ITL ms | Prefill+mixed s | Pure window s | E2E s | +|---|---:|---:|---:|---:|---:| +| C1×64K | 51.2897 | 19.4971 | 14.5356 | 6.2196 | 20.7552 | +| C4×64K | 131.7116 | 30.3694 | 79.8826 | 8.9286 | 89.4635 | +| C8×16K | 247.5878 | 32.3118 | 31.3695 | 9.8228 | 41.6304 | + +The frozen inputs, natural-EOS 320-token cap and pure-window accounting are +unchanged: 319/294/304 tokens per request, no multi-token deliveries. All +18 short quality requests stopped at EOS, passed basic answer checks and had +finite recorded logprobs. Same-arm C4 repeats matched complete token IDs and +top-5 logprobs in 4/4 cases. The post-score C1 profile matched the scored output. +This is bounded quality evidence, not cross-version bitwise or broad quality +acceptance; the precision deferral above remains unchanged. + +All four workers reported 48 admitted grouped layers and tuning enabled. +Startup took 438.4613 seconds. The engine reported 427,385 KV tokens, with +563 configured cache blocks and 0.37 GiB graph memory. Correction: the earlier +563×784 extrapolation was not the engine-reported token capacity. +The model exited with status 0 and no OOM. Python +resource-tracker cleanup warnings at shutdown are retained in the raw log; +the completed scores and profile are not reclassified as a clean shutdown log. + +The first integration attempt failed before long-context scoring: QSA imported +an image-local 18-argument XQA binding, while the hash check inspected a +different top-level 19-argument module. Fixing the harness package path and +checking actual binding identity plus argument conversion resolved that +runtime mismatch. No attention route, quality gate or precision setting was +disabled. Failed and successful evidence remain separate under experiment ID +`qwen38-c1-main-20260905-r2` (successful AWQ/NVFP4) and its unsuffixed failed run. +This investigation did not require another change to the PR's five production +files. GPU integration acceptance does not imply production deployment. + +C1 is a separate follow-up for **both AWQ and NVFP4**, not merely an attempt +to reach NVFP4's current speed. Investigate shared projection/HC, attention +and launch/reduction overhead alongside format-specific MoE preparation. +Neither the common costs nor NVFP4's current performance establish how much +can actually be recovered. Keep that investigation out of this PR. diff --git a/tests/kernels/test_sm70_awq_qpn_m1.py b/tests/kernels/test_sm70_awq_qpn_m1.py new file mode 100644 index 0000000000..1e5b5e3be8 --- /dev/null +++ b/tests/kernels/test_sm70_awq_qpn_m1.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Independent prepared-layout oracle for the opt-in Qwen3.8 AWQ M1 op.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires SM70 and its native extension", +) + + +def _bank(k, n, compact, experts): + weight = torch.zeros((512, k, n // 8), dtype=torch.int32, device="cuda") + metadata = torch.zeros( + (512, k // 32, n, 3) if compact else (512, k // 32, n), + dtype=torch.uint8 if compact else torch.int32, + device="cuda", + ) + decoded = {} + for expert in experts: + codes = torch.randint(0, 16, (k, n), device="cuda") + zeros = torch.randint(0, 16, (k // 32, n), device="cuda") + scales = (torch.rand(k // 32, n, device="cuda") * 0.01 + 0.001).half() + bias = (-zeros.half() * scales).half() + # Independently build the N32/K8 prepared tile layout and nibble order. + values = codes.reshape(k // 8, 8, n // 32, 32).permute(2, 0, 3, 1) + packed = torch.zeros((n // 32, k // 8, 32), dtype=torch.int64, device="cuda") + for logical, physical in enumerate((0, 4, 1, 5, 2, 6, 3, 7)): + packed |= values[..., logical] << (4 * physical) + weight[expert].copy_(packed.int().reshape(k, n // 8)) + scale_bits = scales.view(torch.int16).int() & 0xFFFF + if compact: + metadata[expert].copy_( + torch.stack((scale_bits & 255, scale_bits >> 8, zeros), -1).byte() + ) + else: + bias_bits = bias.view(torch.int16).int() & 0xFFFF + metadata[expert].copy_(scale_bits | (bias_bits << 16)) + group = torch.arange(k, device="cuda") // 32 + decoded[expert] = ( + codes.double() * scales[group].double() + bias[group].double() + ).half() + return weight, metadata, decoded + + +@pytest.mark.parametrize("compact", [False, True]) +def test_awq_qpn_m1_reference_graph_and_admission(compact): + from vllm import _sm70_ops as ops + + assert hasattr(torch.ops._C, "awq_moe_qpn_m1_sm70_out") + torch.manual_seed(731) + experts = (0, 1, 257, 511) + w13, s13, ref13 = _bank(2560, 320, compact, experts) + w2, s2, ref2 = _bank(160, 2560, compact, experts) + x = torch.zeros((1, 2560), dtype=torch.float16, device="cuda") + ids = torch.tensor( + [[511, 0, 257, 1, 511, 1, 0, 257, 0, 511]], + dtype=torch.int32, + device="cuda", + ) + topk = torch.softmax(torch.randn(1, 10, device="cuda"), dim=-1) + out = torch.empty_like(x) + intermediate = torch.empty((10, 160), dtype=x.dtype, device="cuda") + + def run(): + ops.awq_moe_qpn_m1_sm70_out(out, intermediate, x, w13, s13, w2, s2, ids, topk) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run() + # One-hot inputs isolate weight reads at group/partition boundaries. + for index in (0, 31, 32, 2559): + x.zero_() + x[0, index] = 1 + run() + eager_out, eager_mid = out.clone(), intermediate.clone() + expected_mid, expected_routes = [], [] + for expert in ids[0].tolist(): + value = ref13[expert][index].float() + silu = (value[::2] / (1 + torch.exp(-value[::2]))).half() + activation = (silu * value[1::2].half()).half() + expected_mid.append(activation) + expected_routes.append((activation.double() @ ref2[expert].double()).half()) + torch.testing.assert_close( + intermediate, torch.stack(expected_mid), rtol=0, atol=0 + ) + reference = ( + (torch.stack(expected_routes).double() * topk[0].double().unsqueeze(1)) + .sum(0) + .half() + .unsqueeze(0) + ) + # FP64 dot oracle is not the legacy reduction. Allow the materialized + # FP16 route/output roundings, including the minimum subnormal floor. + difference = (out.double() - reference.double()).abs() + assert difference.max() <= reference.double().abs().max() * 2e-3 + 2**-24 + assert ( + difference.norm() <= reference.double().norm() * 2e-3 + 2560**0.5 * 2**-24 + ) + graph.replay() + assert torch.equal(out, eager_out) + assert torch.equal(intermediate, eager_mid) + ids.fill_(-1) + graph.replay() + assert torch.count_nonzero(out) == 0 + assert torch.count_nonzero(intermediate) == 0 + ids.fill_(512) + graph.replay() + assert torch.count_nonzero(out) == 0 + assert torch.count_nonzero(intermediate) == 0 + # Tracing must resolve the native fake implementation without an external + # research DSO. Actual model Inductor/graph validation is a separate gate. + torch.compile(run, backend="eager", fullgraph=True)() + valid = [out, intermediate, x, w13, s13, w2, s2, ids, topk] + unaligned = torch.empty(2561, dtype=x.dtype, device="cuda")[1:].view_as(x) + for index, replacement in ( + (0, x), + (2, unaligned), + (7, ids.long()), + (8, topk.half()), + ): + args = list(valid) + args[index] = replacement + with pytest.raises(RuntimeError): + ops.awq_moe_qpn_m1_sm70_out(*args) diff --git a/tests/quantization/test_awq_qpn_sm70.py b/tests/quantization/test_awq_qpn_sm70.py new file mode 100644 index 0000000000..7d4341c966 --- /dev/null +++ b/tests/quantization/test_awq_qpn_sm70.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from vllm.model_executor.layers.quantization import awq_qpn_sm70 as qpn +from vllm.model_executor.layers.quantization import awq_sm70_moe as moe + +pytestmark = pytest.mark.skip_global_cleanup + + +def layer(): + return SimpleNamespace( + sm70_awq_moe_batched_gemm=True, + sm70_awq_moe_w13_interleaved=True, + sm70_awq_moe_legacy_single_token_compact=True, + sm70_awq_moe_compact_metadata=True, + sm70_awq_checkpoint_group_size=32, + sm70_awq_group_size=32, + sm70_awq_qwen38_qpn_m1=True, + ) + + +def test_default_off_never_loads(monkeypatch): + monkeypatch.delenv("VLLM_SM70_AWQ_QWEN38_QPN_M1", raising=False) + monkeypatch.setattr( + qpn, "_has_native_op", lambda: pytest.fail("unexpected native lookup") + ) + assert not qpn.initialize_qpn_m1(object(), False) + + +@pytest.mark.parametrize("value", ["2", "true", "", "-1"]) +def test_bad_flag_rejected(monkeypatch, value): + monkeypatch.setenv("VLLM_SM70_AWQ_QWEN38_QPN_M1", value) + with pytest.raises(ValueError): + qpn.initialize_qpn_m1(layer(), True) + + +@pytest.mark.parametrize( + "attribute,value", + [ + ("sm70_awq_moe_batched_gemm", False), + ("sm70_awq_moe_w13_interleaved", False), + ("sm70_awq_moe_legacy_single_token_compact", False), + ("sm70_awq_checkpoint_group_size", 128), + ("sm70_awq_group_size", 64), + ], +) +def test_explicit_unsupported_layer_fails(monkeypatch, attribute, value): + monkeypatch.setenv("VLLM_SM70_AWQ_QWEN38_QPN_M1", "1") + current = layer() + setattr(current, attribute, value) + with pytest.raises(RuntimeError): + qpn.initialize_qpn_m1(current, True) + + +def test_explicit_shape_and_native_build(monkeypatch): + monkeypatch.setenv("VLLM_SM70_AWQ_QWEN38_QPN_M1", "1") + with pytest.raises(RuntimeError): + qpn.initialize_qpn_m1(layer(), False) + monkeypatch.setattr(qpn, "_has_native_op", lambda: False) + with pytest.raises(RuntimeError, match="native build"): + qpn.initialize_qpn_m1(layer(), True) + monkeypatch.setattr(qpn, "_has_native_op", lambda: True) + assert qpn.initialize_qpn_m1(layer(), True) + + +@pytest.mark.parametrize("compact", [False, True]) +def test_both_existing_metadata_layouts(monkeypatch, compact): + monkeypatch.setenv("VLLM_SM70_AWQ_QWEN38_QPN_M1", "1") + monkeypatch.setattr(qpn, "_has_native_op", lambda: True) + current = layer() + current.sm70_awq_moe_compact_metadata = compact + assert qpn.initialize_qpn_m1(current, True) + + +def test_no_sidecar_or_implicit_build(monkeypatch): + monkeypatch.setenv("VLLM_SM70_AWQ_QWEN38_QPN_M1", "1") + monkeypatch.setenv("VLLM_SM70_AWQ_QPN_EXTENSION_PATH", "/old/research.so") + monkeypatch.setattr(qpn, "_has_native_op", lambda: False) + monkeypatch.setattr( + torch.ops, "load_library", lambda _: pytest.fail("unexpected DSO load") + ) + with pytest.raises(RuntimeError, match="native build"): + qpn.initialize_qpn_m1(layer(), True) + + +@pytest.mark.parametrize("tokens", [0, 1, 2, 4, 5, 8, 128, 8192]) +def test_only_one_physical_token(tokens): + x = torch.empty(tokens, 2560, dtype=torch.float16) + ids = torch.empty(tokens, 10, dtype=torch.int32) + weights = torch.empty(tokens, 10, dtype=torch.float32) + assert qpn.use_qpn_m1(layer(), x, weights, ids) == (tokens == 1) + + +def test_runtime_layout_and_rollback(): + current = layer() + x = torch.empty(1, 2560, dtype=torch.float16) + ids = torch.empty(1, 10, dtype=torch.int32) + weights = torch.empty(1, 10, dtype=torch.float32) + assert qpn.use_qpn_m1(current, x, weights, ids) + assert not qpn.use_qpn_m1(current, x.bfloat16(), weights, ids) + assert not qpn.use_qpn_m1(current, x, weights.half(), ids) + assert not qpn.use_qpn_m1(current, x, weights, ids.long()) + strided = torch.empty(1, 5120, dtype=torch.float16)[:, ::2] + assert not qpn.use_qpn_m1(current, strided, weights, ids) + current.sm70_awq_qwen38_qpn_m1 = False + assert not qpn.use_qpn_m1(current, x, weights, ids) + + +@pytest.mark.parametrize("compact", [False, True]) +def test_real_moe_branch_calls_native_with_existing_banks(monkeypatch, compact): + current = layer() + current.sm70_awq_moe_compact_metadata = compact + for name in ("w13_tm_weight", "w13_tm_scales", "w2_tm_weight", "w2_tm_scales"): + setattr(current, name, object()) + x = torch.empty(1, 2560, dtype=torch.float16) + ids = torch.empty(1, 10, dtype=torch.int32) + weights = torch.empty(1, 10, dtype=torch.float32) + out = torch.empty_like(x) + intermediate = torch.empty(10, 160, dtype=x.dtype) + seen: list[tuple[Any, ...]] = [] + + def native(*args): + seen.append(args) + args[0].fill_(7) + + monkeypatch.setattr(moe.sm70_ops, "awq_moe_qpn_m1_sm70_out", native) + monkeypatch.setattr( + moe.sm70_ops, + "awq_moe_single_token_sm70_out", + lambda *args: pytest.fail("unexpected legacy route"), + ) + result = moe.AWQSM70MoEMethod._apply_legacy_single_token_compact( + cast(Any, object()), + cast(Any, current), + x, + weights, + ids, + {"intermediate": intermediate}, + 10, + out, + ) + expected = ( + out, + intermediate, + x, + current.w13_tm_weight, + current.w13_tm_scales, + current.w2_tm_weight, + current.w2_tm_scales, + ids, + weights, + ) + assert len(seen) == 1 + assert all(actual is wanted for actual, wanted in zip(seen[0], expected)) + assert result is out + assert torch.all(out == 7) diff --git a/tests/quantization/test_sm70_awq_active_grouped_decode.py b/tests/quantization/test_sm70_awq_active_grouped_decode.py new file mode 100644 index 0000000000..b0aa1960d9 --- /dev/null +++ b/tests/quantization/test_sm70_awq_active_grouped_decode.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm import envs +from vllm.model_executor.layers.quantization.awq_sm70_moe import ( + _qwen38_active_grouped_layer_contract, + _use_qwen38_active_grouped_decode, + _use_qwen38_indexed_prefill, +) +from vllm.model_executor.warmup import awq_sm70_warmup as warmup + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True) +def default_policy(monkeypatch): + for name in ( + "VLLM_SM70_AWQ_MOE_BATCHED_DECODE_MAX_TOKENS", + "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", + "VLLM_SM70_AWQ_MOE_BATCHED_EXACT_W2", + "VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2", + ): + monkeypatch.setenv(name, "0") + + +def _layer(): + return SimpleNamespace( + moe_config=SimpleNamespace(tp_size=4), + sm70_awq_qwen38_active_grouped_decode=True, + sm70_awq_moe_batched_gemm=True, + sm70_num_experts=512, + sm70_hidden_logical_size=2560, + sm70_w13_k_dim=2560, + sm70_w13_n_dim=320, + sm70_w2_k_dim=160, + sm70_w2_n_dim=2560, + ) + + +def test_default_and_rollback(monkeypatch): + name = "VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE" + monkeypatch.delenv(name, raising=False) + assert getattr(envs, name) + monkeypatch.setenv(name, "0") + assert not getattr(envs, name) + layer = _layer() + layer.sm70_awq_qwen38_active_grouped_decode = False + assert not _use_qwen38_active_grouped_decode(layer, 4, 10) + + +@pytest.mark.parametrize("tokens", [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) +def test_token_gate(tokens): + assert _use_qwen38_active_grouped_decode(_layer(), tokens, 10) == (2 <= tokens <= 8) + + +@pytest.mark.parametrize("group_size", [32, 64, 128]) +def test_group_size_gate(group_size): + assert _qwen38_active_grouped_layer_contract(_layer(), group_size) == ( + group_size == 32 + ) + + +@pytest.mark.parametrize("tokens", [1, 2, 8, 9, 127, 128]) +def test_grouped_decode_and_indexed_prefill_are_disjoint(tokens): + layer = _layer() + layer.sm70_awq_qwen38_indexed_prefill = True + layer.sm70_awq_checkpoint_group_size = 32 + layer.sm70_awq_group_size = 32 + layer.sm70_intermediate_size = 160 + x = torch.empty(tokens, 2560, dtype=torch.float16) + topk_ids = torch.empty(tokens, 10, dtype=torch.int32) + grouped = _use_qwen38_active_grouped_decode(layer, tokens, 10) + indexed = _use_qwen38_indexed_prefill(layer, x, topk_ids) + assert grouped == (2 <= tokens <= 8) + assert indexed == (tokens >= 128) + assert not (grouped and indexed) + + +@pytest.mark.parametrize( + "attribute,value", + [ + ("sm70_num_experts", 256), + ("sm70_hidden_logical_size", 2592), + ("sm70_w13_k_dim", 2592), + ("sm70_w13_n_dim", 384), + ("sm70_w2_k_dim", 192), + ("sm70_w2_n_dim", 2592), + ], +) +def test_shape_gate(attribute, value): + layer = _layer() + setattr(layer, attribute, value) + assert not _qwen38_active_grouped_layer_contract(layer, 32) + + +def test_topology_and_router_gate(): + layer = _layer() + assert _qwen38_active_grouped_layer_contract(layer, 32) + layer.moe_config.tp_size = 2 + assert not _qwen38_active_grouped_layer_contract(layer, 32) + assert not _use_qwen38_active_grouped_decode(layer, 4, 8) + layer.sm70_awq_moe_batched_gemm = False + assert not _use_qwen38_active_grouped_decode(layer, 4, 10) + + +@pytest.mark.parametrize( + "name", + [ + "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", + "VLLM_SM70_AWQ_MOE_BATCHED_EXACT_W2", + "VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2", + ], +) +def test_explicit_routes_take_precedence(monkeypatch, name): + monkeypatch.setenv(name, "1") + assert not _use_qwen38_active_grouped_decode(_layer(), 4, 10) + + +def test_decode_cap(monkeypatch): + monkeypatch.setenv("VLLM_SM70_AWQ_MOE_BATCHED_DECODE_MAX_TOKENS", "4") + assert _use_qwen38_active_grouped_decode(_layer(), 4, 10) + assert not _use_qwen38_active_grouped_decode(_layer(), 8, 10) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_warmup_reuses_active_op_and_runtime_policy(monkeypatch, strict): + monkeypatch.setenv( + "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", str(int(strict)) + ) + layer = _layer() + layer._awq_moe_buf_top_k = 10 + layer.w13_tm_scales = torch.empty((512, 80, 320), dtype=torch.float16) + for name in ( + "w13_strided_ptrs_w", + "w13_strided_ptrs_s", + "w2_strided_ptrs_w", + "w2_strided_ptrs_s", + ): + setattr(layer, name, torch.empty(1, dtype=torch.uint8)) + dense_calls, active_calls = [], [] + monkeypatch.setattr( + torch.ops._C, "awq_moe_dense_stage_sm70_out", object(), raising=False + ) + monkeypatch.setattr( + warmup.sm70_ops, + "awq_moe_dense_stage_sm70_out", + lambda *a: dense_calls.append(a), + ) + monkeypatch.setattr( + warmup.sm70_ops, + "awq_moe_active_dense_stage_sm70_out", + lambda *a: active_calls.append(a), + ) + monkeypatch.setattr(warmup, "_silu_and_mul_w13", lambda *a: None) + assert warmup._warmup_moe_dense_stage_layers([layer], [1, 2, 4, 8, 9]) == 10 + assert [a[7] for a in active_calls] == ([] if strict else [20, 20, 40, 40, 80, 80]) + assert len(dense_calls) == (10 if strict else 4) + for w13, w2 in zip(active_calls[::2], active_calls[1::2]): + assert w13[2].tolist() == list(range(w13[7])) + assert w13[3].numel() == w13[7] + 1 + assert w13[4].numel() == w13[7] + assert w13[3] is w2[3] and w13[4] is w2[4] diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 0c8bc76064..d23f02abe8 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -1954,7 +1954,7 @@ def _nvfp4_qwen38_w2_direct_reduce_out_sidecar_fake( scale_codes, global_scales, interleaved_w13, - fast_decode_rounding: None + fast_decode_rounding: (None) ) if hasattr(_raw_namespace, "nvfp4_moe_qpn_raw_w13_swiglu_batch_sm70_out"): register_fake(f"{_raw_prefix}::nvfp4_moe_qpn_raw_w13_swiglu_batch_sm70_out")( @@ -1964,7 +1964,7 @@ def _nvfp4_qwen38_w2_direct_reduce_out_sidecar_fake( scale_codes, global_scales, expert_ids, - interleaved: None + interleaved: (None) ) if hasattr(_raw_namespace, "nvfp4_moe_qpn_raw_w2_reduce_sm70_out"): register_fake(f"{_raw_prefix}::nvfp4_moe_qpn_raw_w2_reduce_sm70_out")( @@ -1974,7 +1974,7 @@ def _nvfp4_qwen38_w2_direct_reduce_out_sidecar_fake( scale_codes, global_scales, expert_ids, - topk_weights: None + topk_weights: (None) ) @@ -3761,6 +3761,39 @@ def _awq_moe_single_token_weighted_reduce_out_fake( return None +def awq_moe_qpn_m1_sm70_out( + out: torch.Tensor, + intermediate: torch.Tensor, + input: torch.Tensor, + w13: torch.Tensor, + s13: torch.Tensor, + w2: torch.Tensor, + s2: torch.Tensor, + ids: torch.Tensor, + topk: torch.Tensor, +) -> None: + _op("awq_moe_qpn_m1_sm70_out")( + out, intermediate, input, w13, s13, w2, s2, ids, topk + ) + + +if hasattr(torch.ops._C, "awq_moe_qpn_m1_sm70_out"): + + @register_fake("_C::awq_moe_qpn_m1_sm70_out") + def _awq_moe_qpn_m1_sm70_out_fake( + out: torch.Tensor, + intermediate: torch.Tensor, + input: torch.Tensor, + w13: torch.Tensor, + s13: torch.Tensor, + w2: torch.Tensor, + s2: torch.Tensor, + ids: torch.Tensor, + topk: torch.Tensor, + ) -> None: + return None + + def awq_moe_single_token_sm70_out( out: torch.Tensor, x: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index f686e9b0e0..0570a93e16 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -119,6 +119,8 @@ VLLM_SM70_AWQ_MOE_DISABLE: bool = False VLLM_SM70_AWQ_MOE_BATCHED_GEMM: bool = True VLLM_SM70_AWQ_QWEN38_MOE_INDEXED_PREFILL: bool = True + VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE: bool = True + VLLM_SM70_AWQ_QWEN38_QPN_M1: bool = False VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13: bool = False VLLM_SM70_AWQ_MOE_BATCHED_EXACT_W2: bool = False VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2: bool = False @@ -1662,6 +1664,16 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_AWQ_QWEN38_MOE_INDEXED_PREFILL": lambda: bool( int(os.getenv("VLLM_SM70_AWQ_QWEN38_MOE_INDEXED_PREFILL", "1")) ), + # Qwen3.8 TP4 g32 small-batch MoE: group active expert segments through + # the existing active-stage op. Set to 0 before startup for the old route. + "VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE": lambda: bool( + int(os.getenv("VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE", "1")) + ), + # Opt-in Qwen3.8 TP4/native-g32 single-token W13/W2 QPN route. + # Keeps prepared metadata and FP16 boundaries, but changes reduction order. + "VLLM_SM70_AWQ_QWEN38_QPN_M1": lambda: ( + env_with_choices("VLLM_SM70_AWQ_QWEN38_QPN_M1", "0", ["0", "1"])() == "1" + ), "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13": lambda: bool( int(os.getenv("VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", "0")) ), diff --git a/vllm/model_executor/layers/quantization/awq_qpn_sm70.py b/vllm/model_executor/layers/quantization/awq_qpn_sm70.py new file mode 100644 index 0000000000..56e7274b1c --- /dev/null +++ b/vllm/model_executor/layers/quantization/awq_qpn_sm70.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Default-off native-g32 Qwen3.8 TP4 AWQ QPN M1 admission. + +The changed FP32 reduction order is not a bitwise-equivalence guarantee. +""" + +import torch + +from vllm import envs + + +def _has_native_op() -> bool: + return hasattr(torch.ops._C, "awq_moe_qpn_m1_sm70_out") + + +def initialize_qpn_m1(layer, shape_contract: bool) -> bool: + if not envs.VLLM_SM70_AWQ_QWEN38_QPN_M1: + return False + if not ( + shape_contract + and layer.sm70_awq_moe_batched_gemm + and layer.sm70_awq_moe_w13_interleaved + and layer.sm70_awq_moe_legacy_single_token_compact + and layer.sm70_awq_checkpoint_group_size == 32 + and layer.sm70_awq_group_size == 32 + ): + raise RuntimeError("AWQ QPN M1 requires native-g32 TP4 E512 C1") + if not _has_native_op(): + raise RuntimeError("AWQ QPN M1 requires a native build with CUDA arch 7.0") + return True + + +def use_qpn_m1(layer, x, topk_weights, topk_ids) -> bool: + return bool( + getattr(layer, "sm70_awq_qwen38_qpn_m1", False) + and x.shape == (1, 2560) + and x.dtype == torch.float16 + and x.is_contiguous() + and topk_ids.shape == (1, 10) + and topk_ids.dtype == torch.int32 + and topk_ids.is_contiguous() + and topk_weights.shape == (1, 10) + and topk_weights.dtype == torch.float32 + and topk_weights.is_contiguous() + ) diff --git a/vllm/model_executor/layers/quantization/awq_sm70_moe.py b/vllm/model_executor/layers/quantization/awq_sm70_moe.py index e8d6fc9b16..88c66d0a87 100644 --- a/vllm/model_executor/layers/quantization/awq_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/awq_sm70_moe.py @@ -21,6 +21,10 @@ SharedExperts, ) from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.quantization.awq_qpn_sm70 import ( + initialize_qpn_m1, + use_qpn_m1, +) from vllm.model_executor.layers.quantization.sm70_moe_router import ( Sm70MoeStageRoute, select_sm70_quantized_moe_route, @@ -78,6 +82,37 @@ def _log_runtime_route_once(message: str, *args) -> None: logger.info_once(message, *args) +def _qwen38_active_grouped_layer_contract( + layer: RoutedExperts, group_size: int +) -> bool: + return bool( + int(layer.moe_config.tp_size) == 4 + and layer.sm70_num_experts == 512 + and group_size == 32 + and layer.sm70_hidden_logical_size == layer.sm70_w13_k_dim == 2560 + and layer.sm70_w13_n_dim == 320 + and layer.sm70_w2_k_dim == 160 + and layer.sm70_w2_n_dim == 2560 + ) + + +def _use_qwen38_active_grouped_decode( + layer: RoutedExperts, num_tokens: int, top_k: int +) -> bool: + """Share the runtime admission policy with pre-capture warmup.""" + max_tokens = envs.VLLM_SM70_AWQ_MOE_BATCHED_DECODE_MAX_TOKENS + return bool( + getattr(layer, "sm70_awq_qwen38_active_grouped_decode", False) + and layer.sm70_awq_moe_batched_gemm + and 2 <= num_tokens <= 8 + and top_k == 10 + and (max_tokens <= 0 or num_tokens <= max_tokens) + and not envs.VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13 + and not envs.VLLM_SM70_AWQ_MOE_BATCHED_EXACT_W2 + and not envs.VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2 + ) + + def _use_temporary_buffers_for_dummy_or_capture() -> bool: # Dummy/profile and CUDA graph capture allocate temporary tensors. Captured # addresses subsequently remain fixed in the graph pool; normal eager @@ -803,6 +838,13 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: and indexed_prefill_requested and indexed_prefill_available ) + layer.sm70_awq_qwen38_active_grouped_decode = bool( + envs.VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE + and _qwen38_active_grouped_layer_contract(layer, self.group_size) + ) + layer.sm70_awq_qwen38_qpn_m1 = initialize_qpn_m1( + layer, _qwen38_active_grouped_layer_contract(layer, self.group_size) + ) self._allocate_buffers(layer) del layer.w13_qweight, layer.w13_scales, layer.w13_qzeros @@ -1137,6 +1179,23 @@ def _apply_legacy_single_token_compact( top_k: int, output: torch.Tensor, ) -> torch.Tensor: + if use_qpn_m1(layer, x, topk_weights, topk_ids_i32): + _log_runtime_route_once( + "SM70 AWQ Qwen3.8 QPN M1 W13/W2 enabled " + "(existing prepared banks, direct route order)." + ) + sm70_ops.awq_moe_qpn_m1_sm70_out( + output, + buffers["intermediate"], + x, + layer.w13_tm_weight, + layer.w13_tm_scales, + layer.w2_tm_weight, + layer.w2_tm_scales, + topk_ids_i32, + topk_weights, + ) + return output _log_runtime_route_once( "SM70 AWQ MoE legacy single-token monolithic compact path enabled " "(top_k=%d, experts=%d).", @@ -1469,7 +1528,9 @@ def apply( use_batched_moe_gemm = route_plan.use_batched_moe_gemm use_batched_active_exact_w2 = route_plan.use_batched_active_exact_w2 use_batched_exact_w2 = route_plan.use_batched_exact_w2 - use_active_exact_small_batched_moe = False + use_active_exact_small_batched_moe = _use_qwen38_active_grouped_decode( + layer, num_tokens, top_k + ) compare_dense_step = None compare_dense_w13_stats = None compare_dense_w2_stats = None @@ -1504,16 +1565,16 @@ def apply( ) elif use_active_exact_small_batched_moe: _log_runtime_route_once( - "SM70 AWQ MoE batched path using active-route exact " - "dense-stage route (tokens=%d, routes=%d).", + "SM70 Qwen3.8 AWQ active grouped decode (tokens=%d, routed_slots=%d).", num_tokens, total_slots, ) - sm70_ops.awq_moe_single_token_dense_stage_sm70_out( + sm70_ops.awq_moe_active_dense_stage_sm70_out( buffers["gate_up"], buffers["permuted_input"], - buffers["active_expert_offsets"], buffers["permuted_experts_id"], + buffers["active_expert_offsets"], + buffers["sorted_expert_ids"], layer.w13_strided_ptrs_w, layer.w13_strided_ptrs_s, total_slots, @@ -1521,6 +1582,21 @@ def apply( layer.sm70_w13_n_dim, self.group_size, ) + if compare_dense_step is not None: + dense_gate_up = torch.empty_like(buffers["gate_up"]) + sm70_ops.awq_moe_dense_stage_sm70_out( + dense_gate_up, + buffers["permuted_input"], + buffers["expert_offsets"], + layer._awq_moe_buf_dense_expert_ids, + layer.w13_strided_ptrs_w, + layer.w13_strided_ptrs_s, + layer.sm70_num_experts, + layer.sm70_w13_k_dim, + layer.sm70_w13_n_dim, + self.group_size, + ) + compare_dense_w13_stats = _diff_stats(buffers["gate_up"], dense_gate_up) elif route_plan.w13 == Sm70MoeStageRoute.PER_EXPERT_DISPATCH: _log_runtime_route_once( "SM70 AWQ MoE batched W13 using per-expert dispatch " @@ -1587,20 +1663,7 @@ def apply( buffers["intermediate"] = _dump_awq_moe_buffer( layer, buffers["intermediate"], "silu_out" ) - if use_active_exact_small_batched_moe: - sm70_ops.awq_moe_single_token_dense_stage_sm70_out( - buffers["sorted_output"], - buffers["intermediate"], - buffers["active_expert_offsets"], - buffers["permuted_experts_id"], - layer.w2_strided_ptrs_w, - layer.w2_strided_ptrs_s, - total_slots, - layer.sm70_w2_k_dim, - layer.sm70_w2_n_dim, - self.group_size, - ) - elif use_batched_active_exact_w2: + if use_active_exact_small_batched_moe or use_batched_active_exact_w2: _log_runtime_route_once( "SM70 AWQ MoE batched path using grouped-active exact W2 (routes=%d).", total_slots, diff --git a/vllm/model_executor/warmup/awq_sm70_warmup.py b/vllm/model_executor/warmup/awq_sm70_warmup.py index 6e04c86f88..527b23cee0 100644 --- a/vllm/model_executor/warmup/awq_sm70_warmup.py +++ b/vllm/model_executor/warmup/awq_sm70_warmup.py @@ -5,7 +5,7 @@ from __future__ import annotations import tempfile -from collections.abc import Iterable +from collections.abc import Callable, Iterable from pathlib import Path from typing import TYPE_CHECKING, Any @@ -15,6 +15,9 @@ from vllm import _sm70_ops as sm70_ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm +from vllm.model_executor.layers.quantization.awq_sm70_moe import ( + _use_qwen38_active_grouped_decode, +) from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( _prepare_compact_slot_groups, _use_compact_grouped, @@ -613,6 +616,23 @@ def _warmup_moe_dense_stage_layers( for num_tokens in token_counts: total_slots = num_tokens * top_k expert_offsets = _build_balanced_offsets(total_slots, num_experts, device) + active_grouped = _use_qwen38_active_grouped_decode(layer, num_tokens, top_k) + stage_op: Callable[..., None] = sm70_ops.awq_moe_dense_stage_sm70_out + stage_metadata: tuple[torch.Tensor, ...] = ( + expert_offsets, + dense_expert_ids, + ) + stage_experts = num_experts + if active_grouped: + # One row per expert warms the same dynamic-offset GEMM used + # when repeated experts form multi-row segments at replay. + stage_op = sm70_ops.awq_moe_active_dense_stage_sm70_out + stage_metadata = ( + dense_expert_ids[:total_slots], + torch.empty(total_slots + 1, dtype=torch.int32, device=device), + torch.empty(total_slots, dtype=torch.int32, device=device), + ) + stage_experts = total_slots permuted_input = torch.empty( (total_slots, int(layer.sm70_w13_k_dim)), dtype=torch.float16, @@ -634,27 +654,25 @@ def _warmup_moe_dense_stage_layers( device=device, ) - sm70_ops.awq_moe_dense_stage_sm70_out( + stage_op( gate_up, permuted_input, - expert_offsets, - dense_expert_ids, + *stage_metadata, layer.w13_strided_ptrs_w, layer.w13_strided_ptrs_s, - num_experts, + stage_experts, int(layer.sm70_w13_k_dim), int(layer.sm70_w13_n_dim), group_size, ) _silu_and_mul_w13(layer, intermediate, gate_up) - sm70_ops.awq_moe_dense_stage_sm70_out( + stage_op( sorted_output, intermediate, - expert_offsets, - dense_expert_ids, + *stage_metadata, layer.w2_strided_ptrs_w, layer.w2_strided_ptrs_s, - num_experts, + stage_experts, int(layer.sm70_w2_k_dim), int(layer.sm70_w2_n_dim), group_size,