[Perf][SM70] Group Qwen3.8 AWQ concurrent decode

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",

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,

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 <typename index_t>
 __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<int>(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<unsigned> 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<int>(), compact_offsets.data_ptr<int>(),
+      compact_expert_ids.data_ptr<int>(), static_cast<int>(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,

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, "

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]