diff --git a/benchmarks/benchmark_sm70_turbomind_exactness.py b/benchmarks/benchmark_sm70_turbomind_exactness.py index af1c1cdb86..27f59da6f0 100644 --- a/benchmarks/benchmark_sm70_turbomind_exactness.py +++ b/benchmarks/benchmark_sm70_turbomind_exactness.py @@ -36,6 +36,7 @@ "batched_w2_per_expert_dispatch", "dense_out", "dense_graphsafe", + "compact_grouped", "single_token_dense", "single_token_indexed", "single_token_compact_w13", @@ -948,6 +949,13 @@ def _check_awq_moe( _require_torch_op("awq_moe_gemm_sm70_per_expert_dispatch_out") if actual_impl in ("dense_graphsafe", "batched_w2_per_expert_dispatch"): _require_torch_op("awq_moe_dense_stage_sm70_out") + if actual_impl == "compact_grouped": + _require_torch_op("awq_moe_compact_grouped_dense_stage_sm70_out") + _require_torch_op("awq_moe_prepare_compact_expert_groups_sm70_out") + if not (2 <= m <= 8 and num_experts == 512 and top_k == 10): + raise ValueError( + "compact_grouped requires Qwen3.8 E512/K10 with --m in [2, 8]." + ) if actual_impl == "active_dense_stage": _require_torch_op("awq_moe_dense_stage_sm70_out") _require_torch_op("awq_moe_active_dense_stage_sm70_out") @@ -1011,6 +1019,15 @@ def _check_awq_moe( sorted_expert_ids, order = torch.sort(logical_expert_ids) expert_offsets, expert_offsets64 = _expert_offsets(sorted_expert_ids, num_experts) dense_expert_ids = torch.arange(num_experts, dtype=torch.int32, device=device) + compact_offsets = torch.empty(total_slots + 1, dtype=torch.int32, device=device) + compact_expert_ids = torch.empty(total_slots, dtype=torch.int32, device=device) + if actual_impl == "compact_grouped": + sm70_ops.awq_moe_prepare_compact_expert_groups_sm70_out( + sorted_expert_ids.to(torch.int32).contiguous(), + compact_offsets, + compact_expert_ids, + total_slots, + ) hidden_size = int(w13_qweight.shape[1]) if actual_impl in AWQ_SINGLE_TOKEN_ACTUALS: @@ -1051,6 +1068,19 @@ def _check_awq_moe( w13_n, group_size, ) + elif actual_impl == "compact_grouped": + sm70_ops.awq_moe_compact_grouped_dense_stage_sm70_out( + gate_up_actual, + sorted_input, + compact_offsets, + compact_expert_ids, + w13_ptrs_w, + w13_ptrs_s, + total_slots, + int(w13_tm_weight.shape[1]), + w13_n, + group_size, + ) elif actual_impl == "batched_w2_per_expert_dispatch" or actual_impl in ( "dense_graphsafe", "active_dense_stage", @@ -1250,6 +1280,19 @@ def _check_awq_moe( group_size, False, ) + elif actual_impl == "compact_grouped": + sm70_ops.awq_moe_compact_grouped_dense_stage_sm70_out( + sorted_output_actual, + intermediate_actual, + compact_offsets, + compact_expert_ids, + w2_ptrs_w, + w2_ptrs_s, + total_slots, + int(w2_tm_weight.shape[1]), + hidden_out, + group_size, + ) elif actual_impl == "dense_graphsafe": sm70_ops.awq_moe_dense_stage_sm70_out( sorted_output_actual, @@ -2545,6 +2588,7 @@ def _parse_args() -> argparse.Namespace: "batched_w2_per_expert_dispatch", "dense_out", "dense_graphsafe", + "compact_grouped", "active_dense_stage", "single_token_dense", "single_token_indexed", diff --git a/csrc/ops.h b/csrc/ops.h index 57d4bfb1c6..0a22200fb1 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -423,6 +423,15 @@ void awq_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, int64_t num_experts, int64_t k, int64_t n, int64_t group_size); +void awq_moe_compact_grouped_dense_stage_sm70_out( + torch::Tensor out, torch::Tensor input, torch::Tensor compact_offsets, + torch::Tensor routed_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_groups, int64_t k, int64_t n, int64_t group_size); + +void awq_moe_prepare_compact_expert_groups_sm70_out( + torch::Tensor sorted_expert_ids, torch::Tensor compact_offsets, + torch::Tensor compact_expert_ids, int64_t total_slots); + void awq_moe_active_dense_stage_sm70_out( torch::Tensor out, torch::Tensor input, torch::Tensor permuted_experts_id, torch::Tensor active_expert_offsets, torch::Tensor active_expert_ids, diff --git a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu index c4d0f45451..df6a54fa0d 100644 --- a/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu +++ b/csrc/sm70_turbomind/ops/awq_sm70_gemm.cu @@ -6512,7 +6512,7 @@ void awq_moe_gemm_sm70_out_impl( int64_t num_experts, int64_t k, int64_t n, int64_t group_size, bool gated_silu, torch::Tensor b_group_indices, bool per_expert_dispatch, torch::Tensor reduce_out, torch::Tensor sorted_weights, - bool weighted_reduce); + bool weighted_reduce, bool compact_grouped_rows); template __global__ void awq_moe_single_token_prepare_kernel( @@ -7331,7 +7331,7 @@ void awq_moe_gemm_sm70_out_impl( bool per_expert_dispatch = false, torch::Tensor reduce_out = torch::Tensor(), torch::Tensor sorted_weights = torch::Tensor(), - bool weighted_reduce = false) { + bool weighted_reduce = false, bool compact_grouped_rows = false) { TORCH_CHECK( sorted_input.is_cuda() && sorted_input.scalar_type() == torch::kFloat16, "awq_moe_gemm_sm70: input must be CUDA float16."); @@ -7510,7 +7510,12 @@ void awq_moe_gemm_sm70_out_impl( op.quant_a = {turbomind::gemm::QuantType::kNone, 0}; op.quant_b = {turbomind::gemm::QuantType::kK, static_cast(group_size)}; op.batch_dim = 0; - op.dispatch_num_override = per_expert_dispatch ? 1 : 0; + op.dispatch_num_override = + (per_expert_dispatch || compact_grouped_rows) ? 1 : 0; + // Compact expert segments can own multiple rows and leave graph-dynamic + // empty groups in the tail. Keep the single-group dispatch choice while + // letting the offsets scheduler discover those bounds on device. + op.active_group_count = 0; auto& workspace_holder = vllm::awq_sm70::get_workspace(device, stream); auto& gemm = vllm::awq_sm70::get_gemm(device); @@ -7613,6 +7618,61 @@ void awq_moe_dense_stage_sm70_out(torch::Tensor out, torch::Tensor input, } } +void awq_moe_compact_grouped_dense_stage_sm70_out( + torch::Tensor out, torch::Tensor input, torch::Tensor compact_offsets, + torch::Tensor routed_expert_ids, torch::Tensor ptrs_w, torch::Tensor ptrs_s, + int64_t num_groups, int64_t k, int64_t n, int64_t group_size) { + TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kFloat16 && + input.is_contiguous(), + "awq_moe_compact_grouped_dense_stage_sm70_out: input must be " + "contiguous CUDA float16."); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16 && + out.is_contiguous(), + "awq_moe_compact_grouped_dense_stage_sm70_out: out must be " + "contiguous CUDA float16."); + TORCH_CHECK( + num_groups >= 20 && num_groups <= 80 && num_groups % 10 == 0, + "awq_moe_compact_grouped_dense_stage_sm70_out: exact Qwen3.8 C2-C8 " + "decode requires 20-80 routed groups in multiples of 10."); + TORCH_CHECK( + group_size == 32 && ((k == 2560 && n == 320) || (k == 160 && n == 2560)), + "awq_moe_compact_grouped_dense_stage_sm70_out: exact Qwen3.8 " + "TP4 AWQ g32 W13/W2 shape is required."); + TORCH_CHECK( + input.dim() == 2 && input.size(0) == num_groups && input.size(1) == k, + "awq_moe_compact_grouped_dense_stage_sm70_out: input shape " + "mismatch."); + TORCH_CHECK(out.dim() == 2 && out.size(0) == num_groups && out.size(1) == n, + "awq_moe_compact_grouped_dense_stage_sm70_out: out shape " + "mismatch."); + TORCH_CHECK(compact_offsets.is_cuda() && + compact_offsets.scalar_type() == torch::kInt32 && + compact_offsets.is_contiguous() && + compact_offsets.numel() >= num_groups + 1, + "awq_moe_compact_grouped_dense_stage_sm70_out: compact offsets " + "must be contiguous CUDA int32 with num_groups+1 entries."); + TORCH_CHECK(routed_expert_ids.is_cuda() && + routed_expert_ids.scalar_type() == torch::kInt32 && + routed_expert_ids.is_contiguous() && + routed_expert_ids.numel() >= num_groups, + "awq_moe_compact_grouped_dense_stage_sm70_out: routed expert " + "IDs must be contiguous CUDA int32."); + TORCH_CHECK(ptrs_w.is_cuda() && ptrs_s.is_cuda(), + "awq_moe_compact_grouped_dense_stage_sm70_out: ptr rows must be " + "CUDA."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + static std::atomic logged_awq_compact_grouped{0u}; + maybe_log_sm70_moe_route_once( + logged_awq_compact_grouped, + "SM70 Qwen3.8 AWQ compact grouped decode path enabled C++ op reached", + input, input.size(0), num_groups); + awq_moe_gemm_sm70_out_impl(out, input, compact_offsets, ptrs_w, ptrs_s, + num_groups, k, n, group_size, false, + routed_expert_ids, false, torch::Tensor(), + torch::Tensor(), false, true); +} + namespace { __global__ void awq_moe_build_active_expert_segments_kernel( @@ -7649,6 +7709,41 @@ __global__ void awq_moe_build_active_expert_segments_kernel( } // namespace +void awq_moe_prepare_compact_expert_groups_sm70_out( + torch::Tensor sorted_expert_ids, torch::Tensor compact_offsets, + torch::Tensor compact_expert_ids, int64_t total_slots) { + TORCH_CHECK(sorted_expert_ids.is_cuda() && + sorted_expert_ids.scalar_type() == torch::kInt32 && + sorted_expert_ids.is_contiguous(), + "awq_moe_prepare_compact_expert_groups_sm70_out: sorted expert " + "IDs must be contiguous CUDA int32."); + TORCH_CHECK(compact_offsets.is_cuda() && + compact_offsets.scalar_type() == torch::kInt32 && + compact_offsets.is_contiguous(), + "awq_moe_prepare_compact_expert_groups_sm70_out: offsets must " + "be contiguous CUDA int32."); + TORCH_CHECK(compact_expert_ids.is_cuda() && + compact_expert_ids.scalar_type() == torch::kInt32 && + compact_expert_ids.is_contiguous(), + "awq_moe_prepare_compact_expert_groups_sm70_out: compact expert " + "IDs must be contiguous CUDA int32."); + TORCH_CHECK(total_slots >= 20 && total_slots <= 80 && total_slots % 10 == 0, + "awq_moe_prepare_compact_expert_groups_sm70_out: exact Qwen3.8 " + "C2-C8 routed-slot count is required."); + TORCH_CHECK(sorted_expert_ids.numel() >= total_slots && + compact_offsets.numel() >= total_slots + 1 && + compact_expert_ids.numel() >= total_slots, + "awq_moe_prepare_compact_expert_groups_sm70_out: index buffer " + "too small."); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(sorted_expert_ids)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + awq_moe_build_active_expert_segments_kernel<<<1, 1, 0, stream>>>( + sorted_expert_ids.data_ptr(), compact_offsets.data_ptr(), + compact_expert_ids.data_ptr(), static_cast(total_slots)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + void awq_moe_active_dense_stage_sm70_out( torch::Tensor out, torch::Tensor input, torch::Tensor permuted_experts_id, torch::Tensor active_expert_offsets, torch::Tensor active_expert_ids, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 8721632102..809dc7a032 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -568,6 +568,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("awq_moe_dense_stage_sm70_out", torch::kCUDA, &awq_moe_dense_stage_sm70_out); + ops.def( + "awq_moe_compact_grouped_dense_stage_sm70_out(" + "Tensor(a!) out, Tensor input, Tensor compact_offsets, " + "Tensor routed_expert_ids, Tensor ptrs_w, Tensor ptrs_s, " + "int num_groups, int k, int n, int group_size) -> ()"); + ops.impl("awq_moe_compact_grouped_dense_stage_sm70_out", torch::kCUDA, + &awq_moe_compact_grouped_dense_stage_sm70_out); + + ops.def( + "awq_moe_prepare_compact_expert_groups_sm70_out(" + "Tensor sorted_expert_ids, Tensor(a!) compact_offsets, " + "Tensor(b!) compact_expert_ids, int total_slots) -> ()"); + ops.impl("awq_moe_prepare_compact_expert_groups_sm70_out", torch::kCUDA, + &awq_moe_prepare_compact_expert_groups_sm70_out); + ops.def( "awq_moe_active_dense_stage_sm70_out(" "Tensor(a!) out, Tensor input, Tensor permuted_experts_id, " diff --git a/tests/quantization/test_sm70_awq_compact_grouped_decode.py b/tests/quantization/test_sm70_awq_compact_grouped_decode.py new file mode 100644 index 0000000000..3e3dedb805 --- /dev/null +++ b/tests/quantization/test_sm70_awq_compact_grouped_decode.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm import envs +from vllm.model_executor.layers.quantization.awq_sm70_moe import ( + _QWEN38_COMPACT_GROUPED_MAX_SLOTS, + _use_qwen38_compact_grouped_decode, +) +from vllm.model_executor.warmup import awq_sm70_warmup as warmup + +pytestmark = pytest.mark.skip_global_cleanup + + +def _qwen38_layer() -> SimpleNamespace: + return SimpleNamespace( + moe_config=SimpleNamespace(tp_size=4), + sm70_awq_qwen38_compact_grouped_decode=True, + sm70_awq_moe_batched_gemm=True, + sm70_awq_group_size=32, + sm70_num_experts=512, + sm70_hidden_logical_size=2560, + sm70_intermediate_size=160, + sm70_w13_k_dim=2560, + sm70_w13_n_dim=320, + sm70_w2_k_dim=160, + sm70_w2_n_dim=2560, + ) + + +def test_qwen38_awq_compact_grouped_decode_defaults_on_with_rollback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + name = "VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE" + monkeypatch.delenv(name, raising=False) + assert envs.VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE + + monkeypatch.setenv(name, "0") + assert not envs.VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE + + +def test_qwen38_awq_compact_grouped_decode_gate_is_exact() -> None: + layer = _qwen38_layer() + + assert _QWEN38_COMPACT_GROUPED_MAX_SLOTS == 80 + assert not _use_qwen38_compact_grouped_decode(layer, 1, 10) + assert _use_qwen38_compact_grouped_decode(layer, 2, 10) + assert _use_qwen38_compact_grouped_decode(layer, 4, 10) + assert _use_qwen38_compact_grouped_decode(layer, 8, 10) + assert not _use_qwen38_compact_grouped_decode(layer, 9, 10) + assert not _use_qwen38_compact_grouped_decode(layer, 8, 8) + + layer.moe_config.tp_size = 2 + assert not _use_qwen38_compact_grouped_decode(layer, 4, 10) + layer.moe_config.tp_size = 4 + + layer.sm70_awq_group_size = 128 + assert not _use_qwen38_compact_grouped_decode(layer, 4, 10) + layer.sm70_awq_group_size = 32 + + layer.sm70_w2_n_dim = 2592 + assert not _use_qwen38_compact_grouped_decode(layer, 4, 10) + layer.sm70_w2_n_dim = 2560 + + layer.sm70_awq_qwen38_compact_grouped_decode = False + assert not _use_qwen38_compact_grouped_decode(layer, 4, 10) + + +def _warmup_layer() -> nn.Module: + layer = nn.Module() + layer._awq_moe_buf_top_k = 10 + layer.sm70_num_experts = 512 + layer.sm70_w13_k_dim = 2560 + layer.sm70_w13_n_dim = 320 + layer.sm70_w2_k_dim = 160 + layer.sm70_w2_n_dim = 2560 + layer.sm70_awq_moe_w13_interleaved = False + layer.sm70_awq_qwen38_compact_grouped_decode = True + layer.sm70_awq_compact_grouped_max_slots = 80 + layer.w13_tm_scales = torch.empty((80, 320), dtype=torch.float16) + layer.w13_strided_ptrs_w = torch.empty(1, dtype=torch.uint8) + layer.w13_strided_ptrs_s = torch.empty(1, dtype=torch.uint8) + layer.w2_strided_ptrs_w = torch.empty(1, dtype=torch.uint8) + layer.w2_strided_ptrs_s = torch.empty(1, dtype=torch.uint8) + return layer + + +def test_awq_warmup_uses_compact_groups_only_through_c8( + monkeypatch: pytest.MonkeyPatch, +) -> None: + layer = _warmup_layer() + dense_calls: list[tuple] = [] + compact_calls: list[tuple] = [] + monkeypatch.setattr( + torch.ops._C, + "awq_moe_compact_grouped_dense_stage_sm70_out", + object(), + raising=False, + ) + 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 *args: dense_calls.append(args), + ) + monkeypatch.setattr( + warmup.sm70_ops, + "awq_moe_compact_grouped_dense_stage_sm70_out", + lambda *args: compact_calls.append(args), + ) + monkeypatch.setattr( + warmup, + "_silu_and_mul_w13", + lambda layer, out, gate_up: out.zero_(), + ) + + assert warmup._warmup_moe_dense_stage_layers([layer], [1, 4, 8, 9]) == 8 + + assert [call[6] for call in compact_calls] == [40, 40, 80, 80] + assert all(call[2].tolist() == list(range(call[6] + 1)) for call in compact_calls) + assert all(call[3].tolist() == list(range(call[6])) for call in compact_calls) + assert [call[6] for call in dense_calls] == [512, 512, 512, 512] diff --git a/vllm/_sm70_ops.py b/vllm/_sm70_ops.py index 6fc4440ca1..2da796bfbd 100644 --- a/vllm/_sm70_ops.py +++ b/vllm/_sm70_ops.py @@ -2729,6 +2729,76 @@ def _awq_moe_dense_stage_sm70_out_fake( return None +def awq_moe_compact_grouped_dense_stage_sm70_out( + out: torch.Tensor, + input: torch.Tensor, + compact_offsets: torch.Tensor, + routed_expert_ids: torch.Tensor, + ptrs_w: torch.Tensor, + ptrs_s: torch.Tensor, + num_groups: int, + k: int, + n: int, + group_size: int, +) -> None: + _op("awq_moe_compact_grouped_dense_stage_sm70_out")( + out, + input, + compact_offsets, + routed_expert_ids, + ptrs_w, + ptrs_s, + num_groups, + k, + n, + group_size, + ) + + +if hasattr(torch.ops._C, "awq_moe_compact_grouped_dense_stage_sm70_out"): + + @register_fake("_C::awq_moe_compact_grouped_dense_stage_sm70_out") + def _awq_moe_compact_grouped_dense_stage_sm70_out_fake( + out: torch.Tensor, + input: torch.Tensor, + compact_offsets: torch.Tensor, + routed_expert_ids: torch.Tensor, + ptrs_w: torch.Tensor, + ptrs_s: torch.Tensor, + num_groups: int, + k: int, + n: int, + group_size: int, + ) -> None: + return None + + +def awq_moe_prepare_compact_expert_groups_sm70_out( + sorted_expert_ids: torch.Tensor, + compact_offsets: torch.Tensor, + compact_expert_ids: torch.Tensor, + total_slots: int, +) -> None: + _op("awq_moe_prepare_compact_expert_groups_sm70_out")( + sorted_expert_ids, + compact_offsets, + compact_expert_ids, + total_slots, + ) + + +if hasattr(torch.ops._C, "awq_moe_prepare_compact_expert_groups_sm70_out"): + + @register_fake("_C::awq_moe_prepare_compact_expert_groups_sm70_out") + def _awq_moe_prepare_compact_expert_groups_sm70_out_fake( + sorted_expert_ids: torch.Tensor, + compact_offsets: torch.Tensor, + compact_expert_ids: torch.Tensor, + total_slots: int, + ) -> None: + return None + + def awq_moe_active_dense_stage_sm70_out( out: torch.Tensor, input: torch.Tensor, diff --git a/vllm/envs.py b/vllm/envs.py index d50e8dd634..33da702842 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -118,6 +118,7 @@ VLLM_SM70_COMPRESSED_TENSORS_TURBOMIND: bool = False VLLM_SM70_AWQ_MOE_DISABLE: bool = False VLLM_SM70_AWQ_MOE_BATCHED_GEMM: 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 @@ -1627,6 +1628,13 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_AWQ_MOE_BATCHED_GEMM": lambda: bool( int(os.getenv("VLLM_SM70_AWQ_MOE_BATCHED_GEMM", "1")) ), + # Match the compact grouped execution contract used by the SM70 NVFP4 + # backend for Qwen3.8 TP4 decode. Contiguous rows routed to the same expert + # share one active group; unsupported shapes retain the existing + # 512-expert batched 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 bab590483e..7da8cd8807 100644 --- a/vllm/model_executor/layers/quantization/awq_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/awq_sm70_moe.py @@ -28,6 +28,7 @@ logger = init_logger(__name__) _DEFAULT_PERSISTENT_MAX_TOKENS = 32 +_QWEN38_COMPACT_GROUPED_MAX_SLOTS = 80 def _log_runtime_route_once(message: str, *args) -> None: @@ -36,6 +37,37 @@ def _log_runtime_route_once(message: str, *args) -> None: logger.info_once(message, *args) +def _qwen38_compact_grouped_layer_contract(layer: RoutedExperts) -> bool: + return bool( + int(layer.moe_config.tp_size) == 4 + and int(layer.sm70_num_experts) == 512 + and int(layer.sm70_hidden_logical_size) == 2560 + and int(layer.sm70_intermediate_size) == 160 + and int(layer.sm70_w13_k_dim) == 2560 + and int(layer.sm70_w13_n_dim) == 320 + and int(layer.sm70_w2_k_dim) == 160 + and int(layer.sm70_w2_n_dim) == 2560 + and int(layer.sm70_awq_group_size) == 32 + ) + + +def _use_qwen38_compact_grouped_decode( + layer: RoutedExperts, + num_tokens: int, + top_k: int, +) -> bool: + """Admit only exact Qwen3.8 TP4 C2-C8 routed-slot groups.""" + total_slots = num_tokens * top_k + return bool( + getattr(layer, "sm70_awq_qwen38_compact_grouped_decode", False) + and getattr(layer, "sm70_awq_moe_batched_gemm", False) + and num_tokens > 1 + and top_k == 10 + and total_slots <= _QWEN38_COMPACT_GROUPED_MAX_SLOTS + and _qwen38_compact_grouped_layer_contract(layer) + ) + + def _use_temporary_buffers_for_dummy_or_capture() -> bool: # CUDA graph replay is address-fixed. Use the per-layer persistent buffers # during capture too, so the captured indexed MoE scratch/output lifetimes @@ -645,10 +677,56 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.sm70_w2_k_ld = w2_k_ld layer.sm70_w2_q_ld = w2_q_ld layer.sm70_intermediate_size = layer.sm70_w2_k_dim + layer.sm70_awq_group_size = self.group_size layer.sm70_awq_moe_batched_gemm = batched_gemm layer.sm70_awq_moe_layer_id = _get_layer_id(layer) layer.sm70_awq_moe_w13_interleaved = w13_interleaved layer.sm70_awq_moe_legacy_single_token_compact = build_legacy_w13 + compact_grouped_requested = bool( + envs.VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE + ) + compact_grouped_ops = { + "awq_moe_compact_grouped_dense_stage_sm70_out": hasattr( + torch.ops._C, "awq_moe_compact_grouped_dense_stage_sm70_out" + ), + "awq_moe_prepare_compact_expert_groups_sm70_out": hasattr( + torch.ops._C, + "awq_moe_prepare_compact_expert_groups_sm70_out", + ), + } + compact_grouped_available = all(compact_grouped_ops.values()) + compact_grouped_explicit = ( + "VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE" in os.environ + ) + compact_grouped_contract = _qwen38_compact_grouped_layer_contract(layer) + if ( + compact_grouped_contract + and compact_grouped_requested + and not compact_grouped_available + ): + if compact_grouped_explicit: + raise RuntimeError( + "The explicit SM70 Qwen3.8 AWQ compact grouped decode route " + "requires " + + ", ".join( + name + for name, available in compact_grouped_ops.items() + if not available + ) + + "." + ) + logger.warning_once( + "The default SM70 Qwen3.8 AWQ compact grouped decode route is " + "not present in the loaded extension; falling back to the " + "512-expert batched route. Explicitly setting " + "VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE=1 fails closed." + ) + layer.sm70_awq_qwen38_compact_grouped_decode = bool( + compact_grouped_contract + and compact_grouped_requested + and compact_grouped_available + ) + layer.sm70_awq_compact_grouped_max_slots = _QWEN38_COMPACT_GROUPED_MAX_SLOTS self._allocate_buffers(layer) del layer.w13_qweight, layer.w13_scales, layer.w13_qzeros @@ -1011,6 +1089,9 @@ def apply( num_tokens = x.shape[0] top_k = topk_ids.shape[1] total_slots = num_tokens * top_k + compact_grouped_decode = _use_qwen38_compact_grouped_decode( + layer, num_tokens, top_k + ) buffers = self._get_buffers(layer, total_slots, num_tokens) output = buffers["output"] output.zero_() @@ -1267,7 +1348,6 @@ 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 compare_dense_step = None compare_dense_w13_stats = None compare_dense_w2_stats = None @@ -1280,18 +1360,24 @@ def apply( if num_tokens <= 8 and use_batched_moe_gemm: compare_dense_step = _compare_dense_decode_step(layer) - if use_active_exact_small_batched_moe: + if compact_grouped_decode: + sm70_ops.awq_moe_prepare_compact_expert_groups_sm70_out( + buffers["permuted_experts_id"], + buffers["active_expert_offsets"], + buffers["sorted_expert_ids"], + total_slots, + ) _log_runtime_route_once( - "SM70 AWQ MoE batched path using active-route exact " - "dense-stage route (tokens=%d, routes=%d).", + "SM70 Qwen3.8 AWQ compact active-expert decode enabled " + "(tokens=%d, routed_slots=%d).", num_tokens, total_slots, ) - sm70_ops.awq_moe_single_token_dense_stage_sm70_out( + sm70_ops.awq_moe_compact_grouped_dense_stage_sm70_out( buffers["gate_up"], buffers["permuted_input"], buffers["active_expert_offsets"], - buffers["permuted_experts_id"], + buffers["sorted_expert_ids"], layer.w13_strided_ptrs_w, layer.w13_strided_ptrs_s, total_slots, @@ -1299,6 +1385,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 " @@ -1365,12 +1466,12 @@ 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( + if compact_grouped_decode: + sm70_ops.awq_moe_compact_grouped_dense_stage_sm70_out( buffers["sorted_output"], buffers["intermediate"], buffers["active_expert_offsets"], - buffers["permuted_experts_id"], + buffers["sorted_expert_ids"], layer.w2_strided_ptrs_w, layer.w2_strided_ptrs_s, total_slots, diff --git a/vllm/model_executor/warmup/awq_sm70_warmup.py b/vllm/model_executor/warmup/awq_sm70_warmup.py index 6e04c86f88..3758852b1a 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 @@ -612,7 +612,31 @@ def _warmup_moe_dense_stage_layers( dense_expert_ids = torch.arange(num_experts, dtype=torch.int32, device=device) for num_tokens in token_counts: total_slots = num_tokens * top_k - expert_offsets = _build_balanced_offsets(total_slots, num_experts, device) + stage_op: Callable[..., None] + compact_grouped = bool( + getattr(layer, "sm70_awq_qwen38_compact_grouped_decode", False) + and num_tokens > 1 + and total_slots + <= int(getattr(layer, "sm70_awq_compact_grouped_max_slots", 0)) + and hasattr( + torch.ops._C, + "awq_moe_compact_grouped_dense_stage_sm70_out", + ) + ) + if compact_grouped: + expert_offsets = torch.arange( + total_slots + 1, dtype=torch.int32, device=device + ) + stage_expert_ids = dense_expert_ids[:total_slots] + stage_experts = total_slots + stage_op = sm70_ops.awq_moe_compact_grouped_dense_stage_sm70_out + else: + expert_offsets = _build_balanced_offsets( + total_slots, num_experts, device + ) + stage_expert_ids = dense_expert_ids + stage_experts = num_experts + stage_op = sm70_ops.awq_moe_dense_stage_sm70_out permuted_input = torch.empty( (total_slots, int(layer.sm70_w13_k_dim)), dtype=torch.float16, @@ -634,27 +658,27 @@ 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_expert_ids, 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_expert_ids, 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,