From 0cd6c8078df6009203e07bffec9de85ed577f98f Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:58:58 -0700 Subject: [PATCH 1/9] [Perf][SM70] Narrow Qwen3.8 AWQ active grouped decode Reuse the existing active-stage operator for TP4 g32 M2-8 decode, align warmup admission, and annotate scratch mutations. Keep C1 and unmatched contracts on existing routes; add focused coverage without new public operators or GEMM kernels. Validated with 51 GPU-directed tests, 32 dynamic-route operator comparisons, and QSA-fixed matched model runs. Single-run pure-decode gains: C4x64K +5.62%, C8x16K +2.24%. Natural-EOS AWQ outputs match 17/18; the retained and/+ near-tie and W13 rounding differences are not claimed bitwise exact. No precision-default changes or deployment. Co-authored-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- csrc/sm70_turbomind/ops/awq_sm70_gemm.cu | 17 ++ csrc/torch_bindings.cpp | 3 +- .../test_sm70_awq_active_grouped_decode.py | 150 ++++++++++++++++++ vllm/envs.py | 6 + .../layers/quantization/awq_sm70_moe.py | 77 ++++++--- vllm/model_executor/warmup/awq_sm70_warmup.py | 36 +++-- 6 files changed, 260 insertions(+), 29 deletions(-) create mode 100644 tests/quantization/test_sm70_awq_active_grouped_decode.py 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..241f438180 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); 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..7a72a7dc3f --- /dev/null +++ b/tests/quantization/test_sm70_awq_active_grouped_decode.py @@ -0,0 +1,150 @@ +# 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, +) +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( + "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/envs.py b/vllm/envs.py index f686e9b0e0..71570456f7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -119,6 +119,7 @@ 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_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 +1663,11 @@ 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")) + ), "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_sm70_moe.py b/vllm/model_executor/layers/quantization/awq_sm70_moe.py index e8d6fc9b16..8596567ac7 100644 --- a/vllm/model_executor/layers/quantization/awq_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/awq_sm70_moe.py @@ -78,6 +78,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 +834,10 @@ 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) + ) self._allocate_buffers(layer) del layer.w13_qweight, layer.w13_scales, layer.w13_qzeros @@ -1469,7 +1504,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 +1541,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 +1558,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 +1639,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, From 96b929346ca890f6b007e9064966423771130b92 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:27:09 -0700 Subject: [PATCH 2/9] docs: record narrow AWQ grouped decode and opt-in tuning evidence Separate implementation-only and existing-autotune gains; document EOS rounding limitations, capacity and startup observations, rollback, and the dual-format C1 follow-up. Co-authored-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- .../sm70_qwen38_awq_active_grouped_decode.md | 163 ++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 docs/design/sm70_qwen38_awq_active_grouped_decode.md 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..b1614be954 --- /dev/null +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -0,0 +1,163 @@ +# 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, native 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. +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. + +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 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. + +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. From 021ea59e9d11ba2e94478c463af1b8884b39a846 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:29:53 -0700 Subject: [PATCH 3/9] test: preserve indexed-prefill admission on aligned AWQ base Record the new main integration gate separately from frozen GPU results; keep prefill and grouped decode admission disjoint. Co-authored-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- .../sm70_qwen38_awq_active_grouped_decode.md | 18 +++++++++++++++--- .../test_sm70_awq_active_grouped_decode.py | 17 +++++++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/docs/design/sm70_qwen38_awq_active_grouped_decode.md b/docs/design/sm70_qwen38_awq_active_grouped_decode.md index b1614be954..83224e83c2 100644 --- a/docs/design/sm70_qwen38_awq_active_grouped_decode.md +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -9,10 +9,12 @@ 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, native group size 32, +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. -Single-token and unmatched shapes retain their existing routes. Explicit +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 @@ -142,7 +144,7 @@ profiled kernel durations, not new unprofiled ITLs or a sum across ranks. ## Validation and follow-up -The source checkpoint is `5cceeaad89d6ead1474ca834afd9aaf3a7bd413c`. +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: @@ -156,6 +158,16 @@ 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. Keep the fork PR in draft until that integration +has been built and GPU-validated; CPU checks cannot close this gate. + 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. diff --git a/tests/quantization/test_sm70_awq_active_grouped_decode.py b/tests/quantization/test_sm70_awq_active_grouped_decode.py index 7a72a7dc3f..b0aa1960d9 100644 --- a/tests/quantization/test_sm70_awq_active_grouped_decode.py +++ b/tests/quantization/test_sm70_awq_active_grouped_decode.py @@ -10,6 +10,7 @@ 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 @@ -64,6 +65,22 @@ def test_group_size_gate(group_size): ) +@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", [ From b50bb1f5037d7dcad3b1ed720a00c6dccda18acb Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:33:27 -0700 Subject: [PATCH 4/9] docs: pin shared FP16 opt-ins and PCIe collective provenance Co-authored-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- docs/design/sm70_qwen38_awq_active_grouped_decode.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/design/sm70_qwen38_awq_active_grouped_decode.md b/docs/design/sm70_qwen38_awq_active_grouped_decode.md index 83224e83c2..ffdfdf4e8d 100644 --- a/docs/design/sm70_qwen38_awq_active_grouped_decode.md +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -60,6 +60,14 @@ tracked separately in 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 From bc4686eadd9acf097366b32ddf4685507c200a41 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 09:06:28 -0700 Subject: [PATCH 5/9] docs: record aligned AWQ model integration acceptance Assisted-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- .../sm70_qwen38_awq_active_grouped_decode.md | 44 ++++++++++++++++++- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/docs/design/sm70_qwen38_awq_active_grouped_decode.md b/docs/design/sm70_qwen38_awq_active_grouped_decode.md index ffdfdf4e8d..57e243eec2 100644 --- a/docs/design/sm70_qwen38_awq_active_grouped_decode.md +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -173,8 +173,48 @@ 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. Keep the fork PR in draft until that integration -has been built and GPU-validated; CPU checks cannot close this gate. +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, KV capacity was 441,392 tokens (563×784), and +graph memory was 0.37 GiB. 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 From fbaa195f4790fddd11b9141b6c375e154bac2ce4 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 10:58:23 -0700 Subject: [PATCH 6/9] docs: correct engine-reported AWQ KV capacity Use the retained engine log instead of extrapolating cache block count. No production or benchmark changes. Co-authored-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> --- docs/design/sm70_qwen38_awq_active_grouped_decode.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/design/sm70_qwen38_awq_active_grouped_decode.md b/docs/design/sm70_qwen38_awq_active_grouped_decode.md index 57e243eec2..f3be0ea1a2 100644 --- a/docs/design/sm70_qwen38_awq_active_grouped_decode.md +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -201,8 +201,10 @@ 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, KV capacity was 441,392 tokens (563×784), and -graph memory was 0.37 GiB. The model exited with status 0 and no OOM. Python +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. From 9ae1a291c3691cc068471b78f96b292f500beeb1 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:19:50 -0700 Subject: [PATCH 7/9] [Kernel][SM70] Add native-g32 AWQ QPN M1 operator Add the typed inference operator, prepared-bank decoder, FP32 reduction kernels and standalone GPU tests without changing model dispatch. Keep 3B scalar loading here so cooperative metadata access can be reviewed separately. Assisted-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> (cherry picked from commit 27a27d3300571a9efcc8ab20d6112ce50911c297) (cherry picked from commit 5b4135fd8bb1f89d0f090c17e7144c09d4f5ca37) --- CMakeLists.txt | 1 + csrc/ops.h | 7 + csrc/sm70_turbomind/ops/awq_qpn_m1_sm70.cu | 341 +++++++++++++++++++++ csrc/torch_bindings.cpp | 6 + docs/design/sm70_awq_qpn_m1.md | 69 +++++ tests/kernels/test_sm70_awq_qpn_m1.py | 132 ++++++++ vllm/_sm70_ops.py | 39 ++- 7 files changed, 592 insertions(+), 3 deletions(-) create mode 100644 csrc/sm70_turbomind/ops/awq_qpn_m1_sm70.cu create mode 100644 docs/design/sm70_awq_qpn_m1.md create mode 100644 tests/kernels/test_sm70_awq_qpn_m1.py 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, From 4d95867e3eb55b36475fd4734cef367ae5ee1017 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:21:18 -0700 Subject: [PATCH 8/9] [Perf][SM70] Add opt-in native-g32 AWQ QPN M1 dispatch Gate the native operator to the existing Qwen3.8 TP4 E512 M1 contract, retaining default-off behavior and neighboring routes. Document numerical attribution, observed quality changes, and incomplete production acceptance. Assisted-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> (cherry picked from commit 37fb5d2a394bb2b77b666e17236318ca1bf2dd55) --- docs/design/sm70_awq_qpn_m1.md | 109 ++++++++++++ tests/quantization/test_awq_qpn_sm70.py | 162 ++++++++++++++++++ vllm/envs.py | 6 + .../layers/quantization/awq_qpn_sm70.py | 46 +++++ .../layers/quantization/awq_sm70_moe.py | 24 +++ 5 files changed, 347 insertions(+) create mode 100644 tests/quantization/test_awq_qpn_sm70.py create mode 100644 vllm/model_executor/layers/quantization/awq_qpn_sm70.py diff --git a/docs/design/sm70_awq_qpn_m1.md b/docs/design/sm70_awq_qpn_m1.md index d46c4fe84c..43cfd00070 100644 --- a/docs/design/sm70_awq_qpn_m1.md +++ b/docs/design/sm70_awq_qpn_m1.md @@ -67,3 +67,112 @@ 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/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/vllm/envs.py b/vllm/envs.py index 71570456f7..0570a93e16 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -120,6 +120,7 @@ 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 @@ -1668,6 +1669,11 @@ def _resolve_rust_frontend_path() -> str | None: "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 8596567ac7..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, @@ -838,6 +842,9 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: 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 @@ -1172,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).", From 356c1bafff769d8bb1a48f8402703333ff6c45e7 Mon Sep 17 00:00:00 2001 From: Leonccaa <166551845+Leonccaa@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:25:01 -0700 Subject: [PATCH 9/9] [Docs][SM70] Define controlled numerical perturbation acceptance Evaluate both AWQ arithmetic paths against an independent fixed-weight reference; distinguish existing quantization error, local arithmetic bounds, and task-level quality. Assisted-by: OpenAI Codex Signed-off-by: Leonccaa <166551845+Leonccaa@users.noreply.github.com> (cherry picked from commit 6f73eac8c019a55a27a103bbd8b6bdfb2bd04ece) --- docs/design/sm70_awq_qpn_m1.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/design/sm70_awq_qpn_m1.md b/docs/design/sm70_awq_qpn_m1.md index 43cfd00070..c474a562fe 100644 --- a/docs/design/sm70_awq_qpn_m1.md +++ b/docs/design/sm70_awq_qpn_m1.md @@ -49,6 +49,17 @@ retains `half_fma(q, scale, half(-zero * scale))`; replacing this with ## 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