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/docs/design/sm70_qwen38_awq_active_grouped_decode.md b/docs/design/sm70_qwen38_awq_active_grouped_decode.md new file mode 100644 index 0000000000..f3be0ea1a2 --- /dev/null +++ b/docs/design/sm70_qwen38_awq_active_grouped_decode.md @@ -0,0 +1,225 @@ +# Qwen3.8 SM70 AWQ active grouped decode + +## Scope and implementation + +This is a narrow replacement for the larger implementation in +[1CatAI/1Cat-vLLM#491](https://github.com/1CatAI/1Cat-vLLM/pull/491), +which was closed over maintenance cost versus measured benefit, not a +correctness failure. It reuses `awq_moe_active_dense_stage_sm70_out`, its +existing active-segment builder, and the existing grouped GEMM implementation. +There is no new public operator, Python operator wrapper, or GEMM kernel. + +Admission requires SM70 AWQ, TP4, 512 experts, top-k 10, effective group size 32, +hidden size 2560, W13 `(K,N)=(2560,320)`, W2 `(K,N)=(160,2560)`, and 2–8 +input tokens. Runtime and pre-capture warmup share the admission policy. +GPU validation uses a native-g32 checkpoint; remapped checkpoints are not a +separate validated quality claim. Single-token and unmatched shapes retain +their existing routes. Explicit +dense/exact diagnostic routes and the existing decode token cap take priority. + +Repeated expert IDs form multi-row segments; unused trailing offsets describe +empty segments. The scheduler remains offsets-based, not a one-row-per-route +NVFP4 dispatch. Scratch offsets and IDs are marked mutable in the existing +Torch schema. Rebuild the extension with the Python changes: an older binary +can still expose the same operator name without the new implementation. + +The model-specific grouped route defaults on for admitted contracts. To roll +back, set this before starting a fresh engine: + +```bash +export VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE=0 +``` + +## Optional existing autotune + +For the validated Qwen3.8 TP4 deployment, the candidate configuration is: + +```bash +export VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE=1 +export VLLM_SM70_AWQ_TUNE_SMALL_SHAPES=1 +``` + +`VLLM_SM70_AWQ_TUNE_SMALL_SHAPES` already exists and remains **off by default**. +This PR does not change the global tuning, FP16 reduction, NCCL, or scheduler +defaults. The tuning flag is broader than the narrow admission gate, so its +results must not be generalized to all AWQ models. Retain the existing +preserve-default-splits settings. For a clean comparison, start independent +engines with private empty caches and no imported GEMM LUT; toggling a flag +inside a process does not invalidate previously selected tactics. Roll tuning +back separately with `VLLM_SM70_AWQ_TUNE_SMALL_SHAPES=0` before startup. + +## Measurement contract and separate attribution + +Evidence was collected on four V100 PCIe 32 GB GPUs, Qwen3.8 Flash-Next AWQ +g32, TP4/MTP0, FP16 activations and KV, maximum length 131328, eight maximum +sequences, 8192 batched tokens, chunked prefill on, prefix caching and async +scheduling off, GPU memory utilization 0.89, language-only MRv2 +`FULL_AND_PIECEWISE`. All arms use the same QSA page4 logical-order correction +tracked separately in +[1CatAI/1Cat-vLLM#494](https://github.com/1CatAI/1Cat-vLLM/pull/494). +That attention correction is not included here. Runtime results are from the +frozen `fbcef6e2f9`-based validation stack, not a retest of later upstream main. + +All measured arms also freeze the existing `VLLM_SM70_QWEN38_FP16_GEMV=1`, +`VLLM_SM70_QWEN38_FUSED_HC_FP16=1`, +`VLLM_SM70_QWEN38_FUSED_GDN_INPUT_FP16=1`, and +`VLLM_SM70_QSA_INDEXER_CUBLAS=1` opt-ins. Runtime-lossy online QPN8 is off. +Actual collectives on this PCIe validation stack use PyNCCL; the later +fully-connected custom-AR HC-sharding results from SXM2 systems do not imply +that the same optimized route was active here. + +Each cell has a 16-token warmup and one scored request batch with a 320-token +output limit, `ignore_eos=false`, `min_tokens=0`, greedy decoding, and frozen +prompt token IDs. The numbered-list prompts request enough output to reach +steady state without suppressing EOS. All scored requests reached 320 tokens. +Pure aggregate throughput counts token deliveries strictly inside the common +window after every request has started decoding and before any request ends. +Each request contributes 319/294/304 tokens for C1/C4/C8 respectively; there +were no multi-token deliveries. This is neither per-request speed nor E2E. + +### Grouped route alone: matched r3 comparison + +Autotune is off in both AWQ arms. These are the implementation-only gains. + +| Cell | Grouped OFF tok/s | Grouped ON tok/s | Gain | NVFP4 tok/s | +|---|---:|---:|---:|---:| +| C1×64K | 48.65 | 48.60 | −0.09% | 54.65 | +| C4×64K | 110.02 | 116.20 | +5.62% | 130.91 | +| C8×16K | 207.29 | 211.94 | +2.24% | 237.58 | + +### Existing autotune: subsequent fresh-process comparison + +Grouped decode is on in both arms. This isolates the configuration benefit, +not additional code added by this PR. + +| Cell | Tuning OFF tok/s | Tuning ON tok/s | Gain | Historical r3 NVFP4 tok/s | +|---|---:|---:|---:|---:| +| C1×64K | 48.7874 | 48.5898 | −0.40% | 54.6544 | +| C4×64K | 116.1290 | 129.2741 | +11.32% | 130.9051 | +| C8×16K | 212.0154 | 242.7240 | +14.48% | 237.5780 | + +The new OFF baseline differs from the previous grouped ON by only ++0.38%/−0.06%/+0.04%. NVFP4 was not rerun as a third arm in this second +experiment; one score per cell does not establish statistical superiority. +Do not add percentages from the two experiments or extrapolate to C8×64K. + +| Cell | Prefill+mixed seconds OFF→ON | Pure-window seconds OFF→ON | E2E seconds OFF→ON | +|---|---:|---:|---:| +| C1×64K | 15.4253→15.4243 | 6.5386→6.5652 | 21.9639→21.9895 | +| C4×64K | 83.8265→83.5213 | 10.1267→9.0970 | 94.6550→93.2887 | +| C8×16K | 32.8447→32.8341 | 11.4709→10.0196 | 44.8195→43.3021 | + +E2E also includes the final drain. C4/C8 E2E duration falls by only +1.44%/3.39%; mixed-phase mean ITL remains about 2.868/2.092 seconds. This is +not a fix for long-prefill interference or a scheduler-policy change. + +## Kernel evidence, quality and costs + +After scoring, each tuning arm ran a separate C4 diagnostic profile. All eight +rank traces contain 16 CPU execute annotations, 16 GPU execute annotations, +16 CUDA Graph replays and 48 W13/W2 pairs per step, with four generation +requests and zero prefill requests. W13 changes from M64×128×32 to M8×256×64 +on all ranks. Its mean per-rank GPU time falls from 5.531 to 2.109 ms/step; +W2 stays near 1.31 ms, FP16 projections/HC near 10.42 ms, and QSA near +3.49 ms. W2 launch grids are not identical across all ranks. These are +profiled kernel durations, not new unprofiled ITLs or a sum across ranks. + +- Grouped-only r3: 54 short quality requests across three arms passed basic + answer checks and stopped naturally at EOS. AWQ OFF/ON complete token IDs + matched in 17/18 cases; the remaining arithmetic wording was `and` versus + `+`, with the same answer 156. +- Tuning A/B: 36/36 short requests passed basic answer checks, stopped at EOS, + and had finite recorded logprobs. C4 self-repeat matched complete token IDs + and recorded top-5 logprobs in 4/4 cases in each arm. Across arms 16/18 + complete token sequences matched; two C8 responses differed in wording. + Both new arms also differ from the older grouped-ON reference in 2/18 + cases, so fresh-process baseline drift precludes assigning every change + exclusively to tuning. +- Real-weight operator replay showed small W13 rounding differences; a + separate tuning probe reached maximum absolute difference 0.0009765625. + This is not proof of model equivalence. No new functional failure was + observed in the bounded prompts, but bitwise equality and broad quality + acceptance are **not** claimed. +- The precision-sensitive cross-batch/process issue is recorded and deferred. + This PR deliberately does not change HC/NCCL precision or attempt a general + numerical-determinism repair. It is separate from QSA ordering correctness. +- Tuning OFF/ON both report 386,392 KV tokens, 509 blocks and 5,210,075,136 + KV tensor bytes per rank, and 0.37 GiB graph memory. Initialization-to-ready + took 431.090/427.330 seconds. There was no observed extra total startup or + KV-capacity cost in this pair; host cache variability and full memory peaks + were not controlled well enough to claim startup acceleration or identical + peak memory. + +## Validation and follow-up + +The measured source checkpoint is `5cceeaad89d6ead1474ca834afd9aaf3a7bd413c`. +The extension built successfully; prior validation retained 51 GPU-directed +test passes and 32 dynamic-route comparisons, including repeated experts and +graph replay. The focused CPU policy/warmup regression command is: + +```bash +.venv/bin/python -m pytest -q tests/quantization/test_sm70_awq_active_grouped_decode.py +``` + +The natural-EOS, per-token, startup/LUT and per-rank trace artifacts are +retained under experiment IDs `awq-narrow-qsa-fixed-ab-20260904` (r3), +`awq-nvfp4-decode-gap-profile-20260904` (r3), and +`awq-autotune-model-ab-20260905`. No production deployment is implied. + +For fork review, the patch was ported onto synchronized main +`755baae1d075ee04fa9096b23fc0225b23589a86`. Conflict resolution preserves the +new indexed-prefill admission and compact-metadata initialization alongside +the grouped-decode flag. Added boundary tests verify that indexed prefill and +grouped decode never both admit the same token count. The original validated +branch is retained. The new base also changes HC, NVFP4 dispatch and scratch +lifetimes: the historical GPU results above are **not** a GPU acceptance of +this new integrated tree. The separate integration run below closes the bounded +model acceptance gate; final human review remains required before promotion. + +### Aligned-tree integration acceptance (2026-09-05, r2) + +The native extensions were rebuilt from `b50bb1f5037d` on base `755baae1d075`. +Both formats additionally use the separate QSA ordering port `05c46b2cd0` and +its matching Flash-V100 Python package and rebuilt extension. No production +changes were made after these binaries were built. Each AWQ cell was scored +once with grouped decode and the optional existing autotune enabled; shared +expert overlap was disabled. This is integration acceptance, not a new matched +grouped-OFF/ON experiment or evidence that all changes come from this PR. + +| AWQ cell | Pure aggregate tok/s | Mean pure ITL ms | Prefill+mixed s | Pure window s | E2E s | +|---|---:|---:|---:|---:|---:| +| C1×64K | 51.2897 | 19.4971 | 14.5356 | 6.2196 | 20.7552 | +| C4×64K | 131.7116 | 30.3694 | 79.8826 | 8.9286 | 89.4635 | +| C8×16K | 247.5878 | 32.3118 | 31.3695 | 9.8228 | 41.6304 | + +The frozen inputs, natural-EOS 320-token cap and pure-window accounting are +unchanged: 319/294/304 tokens per request, no multi-token deliveries. All +18 short quality requests stopped at EOS, passed basic answer checks and had +finite recorded logprobs. Same-arm C4 repeats matched complete token IDs and +top-5 logprobs in 4/4 cases. The post-score C1 profile matched the scored output. +This is bounded quality evidence, not cross-version bitwise or broad quality +acceptance; the precision deferral above remains unchanged. + +All four workers reported 48 admitted grouped layers and tuning enabled. +Startup took 438.4613 seconds. The engine reported 427,385 KV tokens, with +563 configured cache blocks and 0.37 GiB graph memory. Correction: the earlier +563×784 extrapolation was not the engine-reported token capacity. +The model exited with status 0 and no OOM. Python +resource-tracker cleanup warnings at shutdown are retained in the raw log; +the completed scores and profile are not reclassified as a clean shutdown log. + +The first integration attempt failed before long-context scoring: QSA imported +an image-local 18-argument XQA binding, while the hash check inspected a +different top-level 19-argument module. Fixing the harness package path and +checking actual binding identity plus argument conversion resolved that +runtime mismatch. No attention route, quality gate or precision setting was +disabled. Failed and successful evidence remain separate under experiment ID +`qwen38-c1-main-20260905-r2` (successful AWQ/NVFP4) and its unsuffixed failed run. +This investigation did not require another change to the PR's five production +files. GPU integration acceptance does not imply production deployment. + +C1 is a separate follow-up for **both AWQ and NVFP4**, not merely an attempt +to reach NVFP4's current speed. Investigate shared projection/HC, attention +and launch/reduction overhead alongside format-specific MoE preparation. +Neither the common costs nor NVFP4's current performance establish how much +can actually be recovered. Keep that investigation out of this PR. diff --git a/tests/quantization/test_sm70_awq_active_grouped_decode.py b/tests/quantization/test_sm70_awq_active_grouped_decode.py new file mode 100644 index 0000000000..b0aa1960d9 --- /dev/null +++ b/tests/quantization/test_sm70_awq_active_grouped_decode.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm import envs +from vllm.model_executor.layers.quantization.awq_sm70_moe import ( + _qwen38_active_grouped_layer_contract, + _use_qwen38_active_grouped_decode, + _use_qwen38_indexed_prefill, +) +from vllm.model_executor.warmup import awq_sm70_warmup as warmup + +pytestmark = pytest.mark.skip_global_cleanup + + +@pytest.fixture(autouse=True) +def default_policy(monkeypatch): + for name in ( + "VLLM_SM70_AWQ_MOE_BATCHED_DECODE_MAX_TOKENS", + "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", + "VLLM_SM70_AWQ_MOE_BATCHED_EXACT_W2", + "VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2", + ): + monkeypatch.setenv(name, "0") + + +def _layer(): + return SimpleNamespace( + moe_config=SimpleNamespace(tp_size=4), + sm70_awq_qwen38_active_grouped_decode=True, + sm70_awq_moe_batched_gemm=True, + sm70_num_experts=512, + sm70_hidden_logical_size=2560, + sm70_w13_k_dim=2560, + sm70_w13_n_dim=320, + sm70_w2_k_dim=160, + sm70_w2_n_dim=2560, + ) + + +def test_default_and_rollback(monkeypatch): + name = "VLLM_SM70_AWQ_QWEN38_MOE_COMPACT_GROUPED_DECODE" + monkeypatch.delenv(name, raising=False) + assert getattr(envs, name) + monkeypatch.setenv(name, "0") + assert not getattr(envs, name) + layer = _layer() + layer.sm70_awq_qwen38_active_grouped_decode = False + assert not _use_qwen38_active_grouped_decode(layer, 4, 10) + + +@pytest.mark.parametrize("tokens", [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) +def test_token_gate(tokens): + assert _use_qwen38_active_grouped_decode(_layer(), tokens, 10) == (2 <= tokens <= 8) + + +@pytest.mark.parametrize("group_size", [32, 64, 128]) +def test_group_size_gate(group_size): + assert _qwen38_active_grouped_layer_contract(_layer(), group_size) == ( + group_size == 32 + ) + + +@pytest.mark.parametrize("tokens", [1, 2, 8, 9, 127, 128]) +def test_grouped_decode_and_indexed_prefill_are_disjoint(tokens): + layer = _layer() + layer.sm70_awq_qwen38_indexed_prefill = True + layer.sm70_awq_checkpoint_group_size = 32 + layer.sm70_awq_group_size = 32 + layer.sm70_intermediate_size = 160 + x = torch.empty(tokens, 2560, dtype=torch.float16) + topk_ids = torch.empty(tokens, 10, dtype=torch.int32) + grouped = _use_qwen38_active_grouped_decode(layer, tokens, 10) + indexed = _use_qwen38_indexed_prefill(layer, x, topk_ids) + assert grouped == (2 <= tokens <= 8) + assert indexed == (tokens >= 128) + assert not (grouped and indexed) + + +@pytest.mark.parametrize( + "attribute,value", + [ + ("sm70_num_experts", 256), + ("sm70_hidden_logical_size", 2592), + ("sm70_w13_k_dim", 2592), + ("sm70_w13_n_dim", 384), + ("sm70_w2_k_dim", 192), + ("sm70_w2_n_dim", 2592), + ], +) +def test_shape_gate(attribute, value): + layer = _layer() + setattr(layer, attribute, value) + assert not _qwen38_active_grouped_layer_contract(layer, 32) + + +def test_topology_and_router_gate(): + layer = _layer() + assert _qwen38_active_grouped_layer_contract(layer, 32) + layer.moe_config.tp_size = 2 + assert not _qwen38_active_grouped_layer_contract(layer, 32) + assert not _use_qwen38_active_grouped_decode(layer, 4, 8) + layer.sm70_awq_moe_batched_gemm = False + assert not _use_qwen38_active_grouped_decode(layer, 4, 10) + + +@pytest.mark.parametrize( + "name", + [ + "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", + "VLLM_SM70_AWQ_MOE_BATCHED_EXACT_W2", + "VLLM_SM70_AWQ_MOE_BATCHED_ACTIVE_EXACT_W2", + ], +) +def test_explicit_routes_take_precedence(monkeypatch, name): + monkeypatch.setenv(name, "1") + assert not _use_qwen38_active_grouped_decode(_layer(), 4, 10) + + +def test_decode_cap(monkeypatch): + monkeypatch.setenv("VLLM_SM70_AWQ_MOE_BATCHED_DECODE_MAX_TOKENS", "4") + assert _use_qwen38_active_grouped_decode(_layer(), 4, 10) + assert not _use_qwen38_active_grouped_decode(_layer(), 8, 10) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_warmup_reuses_active_op_and_runtime_policy(monkeypatch, strict): + monkeypatch.setenv( + "VLLM_SM70_AWQ_MOE_BATCHED_SINGLE_TOKEN_DENSE_W13", str(int(strict)) + ) + layer = _layer() + layer._awq_moe_buf_top_k = 10 + layer.w13_tm_scales = torch.empty((512, 80, 320), dtype=torch.float16) + for name in ( + "w13_strided_ptrs_w", + "w13_strided_ptrs_s", + "w2_strided_ptrs_w", + "w2_strided_ptrs_s", + ): + setattr(layer, name, torch.empty(1, dtype=torch.uint8)) + dense_calls, active_calls = [], [] + monkeypatch.setattr( + torch.ops._C, "awq_moe_dense_stage_sm70_out", object(), raising=False + ) + monkeypatch.setattr( + warmup.sm70_ops, + "awq_moe_dense_stage_sm70_out", + lambda *a: dense_calls.append(a), + ) + monkeypatch.setattr( + warmup.sm70_ops, + "awq_moe_active_dense_stage_sm70_out", + lambda *a: active_calls.append(a), + ) + monkeypatch.setattr(warmup, "_silu_and_mul_w13", lambda *a: None) + assert warmup._warmup_moe_dense_stage_layers([layer], [1, 2, 4, 8, 9]) == 10 + assert [a[7] for a in active_calls] == ([] if strict else [20, 20, 40, 40, 80, 80]) + assert len(dense_calls) == (10 if strict else 4) + for w13, w2 in zip(active_calls[::2], active_calls[1::2]): + assert w13[2].tolist() == list(range(w13[7])) + assert w13[3].numel() == w13[7] + 1 + assert w13[4].numel() == w13[7] + assert w13[3] is w2[3] and w13[4] is w2[4] diff --git a/vllm/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,