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/torch_bindings.cpp b/csrc/torch_bindings.cpp index 241f438180..c4d59e3ff3 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -696,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..d46c4fe84c --- /dev/null +++ b/docs/design/sm70_awq_qpn_m1.md @@ -0,0 +1,69 @@ +# 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 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. 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/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,