@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e126687a9a

[Model] Support Qwen3.8-Flash-Next (#53896)

Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com> Signed-off-by: andyluo7 <andy.luo@amd.com> Signed-off-by: lishunyang12 <lishunyang03@gmail.com> Signed-off-by: Harry Huang <huanghaoyan.hhy@alibaba-inc.com> Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: KungYork <firelikest@163.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: andyluo7 <43718156+andyluo7@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: andyluo7 <andy.luo@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: lishunyang12 <lishunyang03@gmail.com> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Harry Huang committed Aug 31, 2026 at 13:57 UTC e126687a9a828d513c01a07cd69f025f27d63280
124 files changed +19425 -451
.buildkite/test_areas/lm_eval.yaml
+36
@@ -331,6 +331,42 @@ steps:
331 - uv pip install --system 'gpt-oss[eval]==0.0.5'
332 - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
333
334 +- label: ":nvidia: (B200) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
335 + key: accuracy-eval-qwen3-8-flash-next-fp8-b200
336 + timeout_in_minutes: 60
337 + device: b200-k8s
338 + optional: true
339 + num_devices: 4
340 + source_file_dependencies:
341 + - csrc/libtorch_stable/gdn/
342 + - tests/evals/qwen4_exp/
343 + - vllm/model_executor/layers/mamba/
344 + - vllm/models/qwen4_exp/
345 + - vllm/transformers_utils/configs/qwen4_exp.py
346 + - vllm/v1/attention/backends/short_conv_attn.py
347 + - vllm/v1/spec_decode/qwen4_exp.py
348 + commands:
349 + - uv pip install --system 'evalscope==1.10.0'
350 + - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-b200.txt
351 +
352 +- label: ":nvidia: (H200) Qwen3.8-Flash-Next-FP8 Accuracy Eval"
353 + key: accuracy-eval-qwen3-8-flash-next-fp8-h200
354 + timeout_in_minutes: 60
355 + device: h200
356 + optional: true
357 + num_devices: 4
358 + source_file_dependencies:
359 + - csrc/libtorch_stable/gdn/
360 + - tests/evals/qwen4_exp/
361 + - vllm/model_executor/layers/mamba/
362 + - vllm/models/qwen4_exp/
363 + - vllm/transformers_utils/configs/qwen4_exp.py
364 + - vllm/v1/attention/backends/short_conv_attn.py
365 + - vllm/v1/spec_decode/qwen4_exp.py
366 + commands:
367 + - uv pip install --system 'evalscope==1.10.0'
368 + - pytest -s -v evals/qwen4_exp/test_accuracy.py --config-list-file=configs/models-h200.txt
369 +
370 - label: ":nvidia: (DGX) Spark GPQA Eval (GPT-OSS)"
371 key: gpqa-eval-gpt-oss-spark
372 timeout_in_minutes: 35
.buildkite/test_areas/models_basic.yaml
+37
@@ -86,6 +86,43 @@ steps:
86 # The native NVIDIA Kimi K3 kernels require the SM100 family.
87 - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
88
89 +- label: ":nvidia: (H200) Qwen4 Exp"
90 + key: qwen4-exp-unit-tests
91 + timeout_in_minutes: 40
92 + device: h200_35gb
93 + source_file_dependencies:
94 + - vllm/models/qwen4_exp/
95 + - vllm/transformers_utils/configs/qwen4_exp.py
96 + - vllm/config/compilation.py
97 + - vllm/v1/attention/backends/short_conv_attn.py
98 + - tests/models/qwen4_exp/
99 + commands:
100 + - pytest -v -s models/qwen4_exp/test_hc_ops.py models/qwen4_exp/test_qsa_pre_indexer.py models/qwen4_exp/test_qsa_reference.py
101 + mirror:
102 + amd:
103 + label: ":amd: (MI300) Qwen4 Exp"
104 + dind: false
105 + device: mi300_1
106 + depends_on:
107 + - image-build-amd
108 + commands:
109 + - pytest -v -s models/qwen4_exp/test_qsa_amd.py
110 +
111 +- label: ":computer: (CPU) Qwen4 Exp"
112 + key: qwen4-exp-unit-tests-cpu
113 + depends_on:
114 + - image-build-cpu
115 + timeout_in_minutes: 20
116 + device: cpu-small
117 + source_file_dependencies:
118 + - vllm/models/qwen4_exp/
119 + - vllm/transformers_utils/configs/qwen4_exp.py
120 + - vllm/config/compilation.py
121 + - vllm/v1/attention/backends/short_conv_attn.py
122 + - tests/models/qwen4_exp/
123 + commands:
124 + - pytest -v -s models/qwen4_exp/test_config.py models/qwen4_exp/test_ple.py
125 +
126 - label: ":computer: (CPU) Basic Models Other"
127 key: basic-models-test-other-cpu
128 depends_on:
.buildkite/test_areas/spec_decode.yaml
+1
@@ -11,6 +11,7 @@ steps:
11 - vllm/distributed/
12 - vllm/inputs/
13 - vllm/model_executor/
14 + - vllm/models/qwen4_exp/
15 - vllm/platforms/
16 - vllm/sampling_params.py
17 - vllm/transformers_utils/
csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu
+26 -10
@@ -4,6 +4,7 @@
4 */
5
6 #include <cstdint>
7 +#include <string>
8 #include <cuda_bf16.h>
9 #include <cuda_fp16.h>
10 #include <cuda_runtime.h>
@@ -146,7 +147,7 @@ __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) {
147 return {x, y};
148 }
149
149 -template <typename StateT, int ValueHeadsPerKeyHead>
150 +template <typename StateT, int ValueHeadsPerKeyHead, bool SigmoidGate>
151 __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
152 const __nv_bfloat16* __restrict__ mixed_qkv,
153 const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b,
@@ -354,9 +355,11 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
355 #pragma unroll
356 for (int i = 0; i < 4; ++i) {
357 const int value = lane + i * 32;
357 - const float gate = silu_fast(__bfloat162float(
358 + const float gate_input = __bfloat162float(
359 output_gate[static_cast<int64_t>(token) * strides.gate_row +
359 - value_head * kDimV + value]));
360 + value_head * kDimV + value]);
361 + const float gate =
362 + SigmoidGate ? sigmoid_fast(gate_input) : silu_fast(gate_input);
363 const float weight =
364 norm_weight_is_bf16
365 ? __bfloat162float(
@@ -370,7 +373,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel(
373 }
374 }
375
373 -template <typename StateT, int ValueHeadsPerKeyHead>
376 +template <typename StateT, int ValueHeadsPerKeyHead, bool SigmoidGate>
377 void launch_gdn_decode_post_conv_mtp(
378 torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log,
379 torch::stable::Tensor const& dt_bias,
@@ -395,7 +398,7 @@ void launch_gdn_decode_post_conv_mtp(
398 get_current_cuda_stream(mixed_qkv.get_device_index());
399 const int num_requests = static_cast<int>(state_indices.size(0));
400 const dim3 grid(num_requests, num_value_heads);
398 - gdn_decode_post_conv_mtp_kernel<StateT, ValueHeadsPerKeyHead>
401 + gdn_decode_post_conv_mtp_kernel<StateT, ValueHeadsPerKeyHead, SigmoidGate>
402 <<<grid, kThreads, 0, stream>>>(
403 static_cast<const __nv_bfloat16*>(mixed_qkv.data_ptr()), a, b,
404 static_cast<const float*>(a_log.data_ptr()), dt_bias.data_ptr(),
@@ -425,7 +428,7 @@ void fused_gdn_decode_post_conv_mtp(
428 torch::stable::Tensor const& num_accepted_tokens,
429 torch::stable::Tensor& state, torch::stable::Tensor const& output_gate,
430 torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out,
428 - double scale, double norm_eps) {
431 + double scale, double norm_eps, const std::string& output_gate_activation) {
432 using torch::headeronly::ScalarType;
433
434 STD_TORCH_CHECK(
@@ -466,6 +469,9 @@ void fused_gdn_decode_post_conv_mtp(
469 "norm_weight must be a CUDA float32 or bfloat16 tensor");
470 STD_TORCH_CHECK(out.is_cuda() && out.scalar_type() == ScalarType::BFloat16,
471 "out must be a CUDA bfloat16 tensor");
472 + STD_TORCH_CHECK(
473 + output_gate_activation == "silu" || output_gate_activation == "sigmoid",
474 + "output_gate_activation must be 'silu' or 'sigmoid'");
475
476 STD_TORCH_CHECK(mixed_qkv.dim() == 2,
477 "mixed_qkv must have shape [L, 2 * H * 128 + HV * 128]");
@@ -547,8 +553,9 @@ void fused_gdn_decode_post_conv_mtp(
553 const auto* b_ptr = static_cast<const __nv_bfloat16*>(b.data_ptr());
554 const auto* output_gate_ptr =
555 static_cast<const __nv_bfloat16*>(output_gate.data_ptr());
550 - const auto launch = [&]<typename StateT, int ValueHeadsPerKeyHead>() {
551 - launch_gdn_decode_post_conv_mtp<StateT, ValueHeadsPerKeyHead>(
556 + const auto launch = [&]<typename StateT, int ValueHeadsPerKeyHead,
557 + bool SigmoidGate>() {
558 + launch_gdn_decode_post_conv_mtp<StateT, ValueHeadsPerKeyHead, SigmoidGate>(
559 mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens,
560 num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr,
561 output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps,
@@ -556,9 +563,18 @@ void fused_gdn_decode_post_conv_mtp(
563 };
564 const auto dispatch_state_type = [&]<int ValueHeadsPerKeyHead>() {
565 if (state_scalar_type == ScalarType::Float) {
559 - launch.template operator()<float, ValueHeadsPerKeyHead>();
566 + if (output_gate_activation == "sigmoid") {
567 + launch.template operator()<float, ValueHeadsPerKeyHead, true>();
568 + } else {
569 + launch.template operator()<float, ValueHeadsPerKeyHead, false>();
570 + }
571 } else {
561 - launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead>();
572 + if (output_gate_activation == "sigmoid") {
573 + launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead, true>();
574 + } else {
575 + launch
576 + .template operator()<__nv_bfloat16, ValueHeadsPerKeyHead, false>();
577 + }
578 }
579 };
580 switch (value_heads_per_key_head) {
csrc/libtorch_stable/ops.h
+1 -1
@@ -405,7 +405,7 @@ void fused_gdn_decode_post_conv_mtp(
405 torch::stable::Tensor const& num_accepted_tokens,
406 torch::stable::Tensor& state, torch::stable::Tensor const& output_gate,
407 torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out,
408 - double scale, double norm_eps);
408 + double scale, double norm_eps, const std::string& output_gate_activation);
409
410 #endif
411
csrc/libtorch_stable/torch_bindings.cpp
+2 -1
@@ -535,7 +535,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
535 "Tensor mixed_qkv, Tensor a, Tensor b, Tensor A_log, Tensor dt_bias, "
536 "Tensor state_indices, Tensor cu_seqlens, Tensor num_accepted_tokens, "
537 "Tensor! state, Tensor output_gate, Tensor norm_weight, Tensor! out, "
538 - "float scale, float norm_eps=1e-5) -> ()");
538 + "float scale, float norm_eps=1e-5, "
539 + "str output_gate_activation='silu') -> ()");
540 #endif
541
542 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
tests/config/test_config_utils.py
+19
@@ -7,6 +7,7 @@ from enum import Enum
7 import pytest
8
9 from vllm.config.cache import CacheConfig
10 +from vllm.config.scheduler import SchedulerConfig
11 from vllm.config.utils import get_hash_factors, hash_factors, normalize_value
12
13 # Helpers
@@ -216,6 +217,24 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
217 assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
218
219
220 +def test_scheduler_config_hash_includes_max_num_seqs():
221 + """Per-request workspace sizes must invalidate compiled graphs."""
222 + base_hash = SchedulerConfig(
223 + max_model_len=8192,
224 + is_encoder_decoder=False,
225 + max_num_batched_tokens=8192,
226 + max_num_seqs=128,
227 + ).compute_hash()
228 + larger_batch_hash = SchedulerConfig(
229 + max_model_len=8192,
230 + is_encoder_decoder=False,
231 + max_num_batched_tokens=8192,
232 + max_num_seqs=1024,
233 + ).compute_hash()
234 +
235 + assert larger_batch_hash != base_hash
236 +
237 +
238 def test_cache_config_hash_ignores_prefix_cache_retention_interval():
239 base_hash = CacheConfig().compute_hash()
240 assert CacheConfig(prefix_cache_retention_interval=64).compute_hash() == base_hash
tests/config/test_speculative_draft_hf_overrides.py
+52 -1
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for SpeculativeConfig.compose_draft_hf_overrides.
3 +"""Tests for draft config overrides used by SpeculativeConfig.
4
5 Callable ``hf_overrides`` on the target model config (e.g. the
6 ``dummy_hf_overrides`` shrink used by ``tests/models/test_initialization.py``)
@@ -12,10 +12,12 @@ when the target itself is shrunk — which is what kept spec-decode archs like
12 """
13
14 import functools
15 +from unittest.mock import MagicMock, patch
16
17 import pytest
18 from transformers import PretrainedConfig
19
20 +from vllm.config.parallel import ParallelConfig
21 from vllm.config.speculative import SpeculativeConfig
22
23
@@ -134,3 +136,52 @@ def test_composed_override_is_picklable():
136
137 out = composed(_make_hf_config())
138 assert out.num_hidden_layers == 1
139 +
140 +
141 +def _make_mtp_speculative_config(
142 + override: bool | None,
143 + checkpoint_value: bool,
144 +) -> SpeculativeConfig:
145 + draft_hf_config = _make_hf_config(
146 + architectures=["Qwen4ExpMTP"],
147 + model_type="qwen4_exp_mtp",
148 + n_predict=1,
149 + index_share_for_mtp_iteration=checkpoint_value,
150 + )
151 + draft_model_config = MagicMock(
152 + model="draft",
153 + hf_config=draft_hf_config,
154 + architectures=draft_hf_config.architectures,
155 + max_model_len=128,
156 + )
157 + target_model_config = MagicMock(
158 + model="target",
159 + max_model_len=128,
160 + quantization=None,
161 + hf_overrides={},
162 + )
163 +
164 + with patch("vllm.config.speculative.ModelConfig", return_value=draft_model_config):
165 + return SpeculativeConfig(
166 + model="draft",
167 + method="mtp",
168 + num_speculative_tokens=1,
169 + index_share_for_mtp_iteration=override,
170 + target_model_config=target_model_config,
171 + target_parallel_config=ParallelConfig(),
172 + )
173 +
174 +
175 +@pytest.mark.cpu_test
176 +@pytest.mark.parametrize(
177 + ("override", "checkpoint_value", "expected"),
178 + [(None, True, True), (False, True, False), (True, False, True)],
179 +)
180 +def test_mtp_index_share_override(
181 + override: bool | None, checkpoint_value: bool, expected: bool
182 +):
183 + speculative_config = _make_mtp_speculative_config(override, checkpoint_value)
184 + assert (
185 + speculative_config.draft_model_config.hf_config.index_share_for_mtp_iteration
186 + is expected
187 + )
tests/distributed/test_custom_all_reduce.py
+23
@@ -24,6 +24,29 @@ for i, v in enumerate(test_sizes):
24 test_sizes[i] -= v % 8
25
26
27 +@pytest.mark.parametrize(
28 + ("dtype", "expected"),
29 + [
30 + (torch.float32, True),
31 + (torch.float16, True),
32 + (torch.bfloat16, True),
33 + (torch.int8, False),
34 + (torch.float8_e4m3fn, False),
35 + ],
36 +)
37 +def test_custom_allreduce_filters_dtype(
38 + dtype: torch.dtype,
39 + expected: bool,
40 +) -> None:
41 + communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
42 + communicator.disabled = False
43 + communicator.world_size = 2
44 + communicator.max_size = 1024
45 + communicator._ptr = 0
46 +
47 + assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) is expected
48 +
49 +
50 @pytest.mark.parametrize(
51 ("major", "local_multicast", "expected"),
52 [
tests/evals/qwen4_exp/README.md new
+14
@@ -0,0 +1,14 @@
1 +# Qwen4Exp accuracy evaluation
2 +
3 +This suite starts a Qwen3.8-Flash-Next-FP8 OpenAI-compatible server once and
4 +uses EvalScope to evaluate GSM8K and AIME25.
5 +
6 +```bash
7 +# B200
8 +pytest -s -v tests/evals/qwen4_exp/test_accuracy.py \
9 + --config-list-file=configs/models-b200.txt
10 +
11 +# H200
12 +pytest -s -v tests/evals/qwen4_exp/test_accuracy.py \
13 + --config-list-file=configs/models-h200.txt
14 +```
tests/evals/qwen4_exp/__init__.py new
+3
@@ -0,0 +1,3 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8.yaml new
+26
@@ -0,0 +1,26 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +model_name: "Qwen/Qwen3.8-Flash-Next-FP8"
4 +datasets:
5 + gsm8k:
6 + metric_threshold: 0.98
7 + tolerance: 0.02
8 + aime25:
9 + metric_threshold: 0.90
10 + tolerance: 0.05
11 +eval_batch_size: 32
12 +startup_max_wait_seconds: 1800
13 +generation_config:
14 + seed: 1236
15 + do_sample: true
16 + temperature: 0.7
17 + top_p: 0.8
18 + top_k: 20
19 + repetition_penalty: 1.0
20 + presence_penalty: 1.5
21 + max_tokens: 32768
22 + stream: true
23 + extra_body:
24 + chat_template_kwargs:
25 + enable_thinking: true
26 +server_args: "--tensor-parallel-size 4"
tests/evals/qwen4_exp/configs/models-b200.txt new
+2
@@ -0,0 +1,2 @@
1 +# B200 configuration for Qwen3.8-Flash-Next-FP8 accuracy evaluation
2 +Qwen3.8-Flash-Next-FP8.yaml
tests/evals/qwen4_exp/configs/models-h200.txt new
+2
@@ -0,0 +1,2 @@
1 +# H200 configuration for Qwen3.8-Flash-Next-FP8 accuracy evaluation
2 +Qwen3.8-Flash-Next-FP8.yaml
tests/evals/qwen4_exp/conftest.py new
+42
@@ -0,0 +1,42 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from pathlib import Path
5 +
6 +
7 +def pytest_addoption(parser):
8 + parser.addoption(
9 + "--config-list-file",
10 + required=True,
11 + help="File containing Qwen4Exp evaluation config files",
12 + )
13 +
14 +
15 +def pytest_generate_tests(metafunc):
16 + if "config_filename" not in metafunc.fixturenames:
17 + return
18 +
19 + config_list_path = Path(metafunc.config.getoption("--config-list-file"))
20 + if not config_list_path.is_absolute():
21 + test_relative_path = Path(__file__).parent / config_list_path
22 + config_list_path = (
23 + test_relative_path
24 + if test_relative_path.exists()
25 + else Path.cwd() / config_list_path
26 + )
27 +
28 + config_dir = config_list_path.parent
29 + config_files = [
30 + config_dir / line.strip()
31 + for line in config_list_path.read_text(encoding="utf-8").splitlines()
32 + if line.strip() and not line.lstrip().startswith("#")
33 + ]
34 + missing = [path for path in config_files if not path.exists()]
35 + if missing:
36 + raise FileNotFoundError(f"Missing evaluation configs: {missing}")
37 +
38 + metafunc.parametrize(
39 + "config_filename",
40 + config_files,
41 + ids=[config_file.stem for config_file in config_files],
42 + )
tests/evals/qwen4_exp/test_accuracy.py new
+68
@@ -0,0 +1,68 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""GSM8K and AIME25 accuracy evaluation for Qwen3.8-Flash-Next-FP8."""
4 +
5 +import shlex
6 +from pathlib import Path
7 +from typing import Any
8 +
9 +import yaml
10 +
11 +from tests.utils import RemoteOpenAIServer
12 +
13 +
14 +def run_evalscope(
15 + eval_config: dict[str, Any], base_url: str, work_dir: Path
16 +) -> dict[str, float]:
17 + from evalscope.run import run_task
18 +
19 + reports = run_task(
20 + {
21 + "model": eval_config["model_name"],
22 + "api_url": base_url,
23 + "api_key": "EMPTY_TOKEN",
24 + "datasets": list(eval_config["datasets"]),
25 + "eval_batch_size": eval_config.get("eval_batch_size", 32),
26 + "generation_config": eval_config["generation_config"],
27 + "work_dir": str(work_dir),
28 + "no_timestamp": True,
29 + }
30 + )
31 + if not isinstance(reports, dict):
32 + raise TypeError(f"Expected EvalScope reports to be a dict, got {type(reports)}")
33 +
34 + return {dataset: float(report.score) for dataset, report in reports.items()}
35 +
36 +
37 +def test_qwen4_exp_accuracy(config_filename: Path, tmp_path: Path):
38 + eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
39 + server_args = shlex.split(eval_config.get("server_args", ""))
40 + server_args.extend(["--trust-remote-code", "--disable-uvicorn-access-log"])
41 +
42 + model_name = eval_config["model_name"]
43 + print(f"Starting Qwen4Exp evaluation for model: {model_name}")
44 + print(f"Datasets: {', '.join(eval_config['datasets'])}")
45 + print(f"Server args: {' '.join(server_args)}")
46 +
47 + with RemoteOpenAIServer(
48 + model_name,
49 + server_args,
50 + env_dict=eval_config.get("env"),
51 + max_wait_seconds=eval_config.get("startup_max_wait_seconds", 1800),
52 + ) as remote_server:
53 + scores = run_evalscope(eval_config, remote_server.url_for("v1"), tmp_path)
54 +
55 + for dataset, metric_config in eval_config["datasets"].items():
56 + score = scores[dataset]
57 + threshold = metric_config["metric_threshold"]
58 + tolerance = metric_config.get("tolerance", 0.0)
59 + minimum_score = threshold - tolerance
60 +
61 + print(
62 + f"{dataset}: measured={score:.4f}, expected={threshold:.4f}, "
63 + f"tolerance={tolerance:.4f}"
64 + )
65 + assert score >= minimum_score, (
66 + f"{dataset} score too low: {score:.4f} < {threshold:.4f} - "
67 + f"{tolerance:.4f} = {minimum_score:.4f}"
68 + )
tests/kernels/mamba/test_gdn_fused_mtp.py
+8 -3
@@ -57,13 +57,13 @@ EPS = 1e-6
57
58
59 class _TestGatedNorm:
60 - def __init__(self, weight: torch.Tensor) -> None:
60 + def __init__(self, weight: torch.Tensor, activation: str) -> None:
61 self.weight = weight
62 self.bias = None
63 self.eps = EPS
64 self.group_size = None
65 self.norm_before_gate = True
66 - self.activation = "silu"
66 + self.activation = activation
67
68 def __call__(self, x: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
69 return rmsnorm_fn(
@@ -99,6 +99,7 @@ def _build_layer(
99 dt_bias: torch.Tensor,
100 conv_weight: torch.Tensor,
101 norm_weight: torch.Tensor,
102 + output_gate_activation: str,
103 ):
104 layer = types.SimpleNamespace(
105 prefix=PREFIX,
@@ -116,7 +117,7 @@ def _build_layer(
117 dt_bias=dt_bias,
118 conv1d=types.SimpleNamespace(weight=conv_weight, bias=None),
119 kv_cache=(conv_state, ssm_state),
119 - norm=_TestGatedNorm(norm_weight),
120 + norm=_TestGatedNorm(norm_weight, output_gate_activation),
121 layer_norm_epsilon=EPS,
122 gdn_decode_kernel="cuda",
123 )
@@ -250,12 +251,14 @@ def test_fused_forward_uses_packed_entrypoint() -> None:
251 pytest.param([128], [1], [-1], 0, id="pure-decode"),
252 ],
253 )
254 +@pytest.mark.parametrize("output_gate_activation", ["silu", "sigmoid"])
255 @torch.inference_mode()
256 def test_fused_model_path_matches_reference(
257 seq_lens: list[int],
258 query_lens: list[int],
259 draft_tokens: list[int],
260 expected_fused_calls: int,
261 + output_gate_activation: str,
262 ) -> None:
263 """Fused MTP and its mixed/prefill/decode fallbacks match the reference."""
264 torch.manual_seed(1)
@@ -334,6 +337,7 @@ def test_fused_model_path_matches_reference(
337 dt_bias,
338 conv_weight,
339 norm_weight,
340 + output_gate_activation,
341 )
342 reference_out = torch.zeros_like(output_gate)
343 with patch.object(
@@ -355,6 +359,7 @@ def test_fused_model_path_matches_reference(
359 dt_bias,
360 conv_weight,
361 norm_weight,
362 + output_gate_activation,
363 )
364 context.no_compile_layers = {PREFIX: fused_layer}
365 fused_out = torch.zeros_like(output_gate)
tests/kernels/mamba/test_precopy_mamba_align.py
+2 -2
@@ -366,7 +366,7 @@ def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bi
366 input_batch=scalar_case[1],
367 requests=scalar_case[2],
368 forward_context={},
369 - mamba_state_copy_funcs=(),
369 + mamba_state_copy_funcs={},
370 copy_bufs=scalar_copy_bufs,
371 )
372
@@ -384,7 +384,7 @@ def test_preprocess_fused_align_matches_scalar_bookkeeping(monkeypatch, token_bi
384 input_batch=fused_case[1],
385 requests=fused_case[2],
386 forward_context={},
387 - mamba_state_copy_funcs=(),
387 + mamba_state_copy_funcs={},
388 copy_bufs=fused_copy_bufs,
389 align_ctx=ctx,
390 )
tests/kernels/moe/test_block_fp8.py
+104
@@ -351,3 +351,107 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)
351 torch.accelerator.synchronize()
352
353 torch.testing.assert_close(out, ref_out, atol=0.035, rtol=0.035)
354 +
355 +
356 +def _dequant_block_fp8(
357 + w_fp8: torch.Tensor, w_s: torch.Tensor, block_shape: list[int]
358 +) -> torch.Tensor:
359 + """Dequantize a [N, K] fp8 weight with [Nb, Kb] block scales to fp32."""
360 + n, k = w_fp8.shape
361 + bn, bk = block_shape
362 + s = w_s.repeat_interleave(bn, dim=0).repeat_interleave(bk, dim=1)
363 + return w_fp8.to(torch.float32) * s[:n, :k]
364 +
365 +
366 +@pytest.mark.parametrize("tp_rank", [0, 1, 2, 3])
367 +@torch.inference_mode()
368 +def test_w8a8_block_fp8_fused_moe_refined_block_scales(tp_rank, workspace_init):
369 + """TP-misaligned blockwise FP8 MoE (e.g. Qwen4Exp, intermediate
370 + 640 at TP=4 -> 160 per rank, not divisible by the checkpoint's 128 block).
371 +
372 + The per-rank scale grid is refined from 128 to 32 (lossless: each 32-block
373 + lies inside one global 128-block) so the Triton kernel can consume exact
374 + per-shard scales. Simulates one TP rank's shard and checks:
375 + 1. the refined+sharded scales dequantize exactly like the global scales;
376 + 2. the Triton fused MoE kernel with block_shape=[32, 32] matches the
377 + native-torch blockwise reference on the same shard.
378 + """
379 + torch.manual_seed(0)
380 + dtype = torch.bfloat16
381 + M, E, topk = 8, 4, 2
382 + hidden = 256
383 + inter_full = 640
384 + tp_size = 4
385 + n_shard = inter_full // tp_size # 160
386 + ckpt_block = [128, 128]
387 + refined_block = [32, 32]
388 + factor = ckpt_block[0] // refined_block[0] # 4
389 +
390 + a = torch.randn((M, hidden), dtype=dtype) / 10
391 + score = torch.randn((M, E), dtype=dtype)
392 +
393 + (_, w1_full, w1_s_full, _), (_, w2_full, w2_s_full, _) = make_test_weights(
394 + E,
395 + inter_full,
396 + hidden,
397 + dtype,
398 + torch.float8_e4m3fn,
399 + per_out_ch_quant=False,
400 + block_shape=ckpt_block,
401 + )
402 + assert w1_s_full is not None
403 + assert w2_s_full is not None
404 +
405 + # TP shard the weights (exact slicing, no scale involvement).
406 + lo, hi = tp_rank * n_shard, (tp_rank + 1) * n_shard
407 + w1 = torch.cat(
408 + [w1_full[:, lo:hi], w1_full[:, inter_full + lo : inter_full + hi]], dim=1
409 + )
410 + w2 = w2_full[:, :, lo:hi]
411 +
412 + # Refine the global 128-block scales to 32 blocks, then take the shard's
413 + # slice -- mirrors the Fp8MoEMethod/RoutedExperts loading path.
414 + w1_s32 = w1_s_full.repeat_interleave(factor, dim=-2).repeat_interleave(
415 + factor, dim=-1
416 + )
417 + nb = n_shard // refined_block[0] # 5 local 32-blocks per projection
418 + gate_hi = w1_s32.shape[1] // 2
419 + w1_s = torch.cat(
420 + [
421 + w1_s32[:, tp_rank * nb : (tp_rank + 1) * nb],
422 + w1_s32[:, gate_hi + tp_rank * nb : gate_hi + (tp_rank + 1) * nb],
423 + ],
424 + dim=1,
425 + )
426 + w2_s32 = w2_s_full.repeat_interleave(factor, dim=-2).repeat_interleave(
427 + factor, dim=-1
428 + )
429 + w2_s = w2_s32[:, :, tp_rank * nb : (tp_rank + 1) * nb]
430 +
431 + # The refined per-shard scales must reproduce the global-scale dequant
432 + # exactly (dequant full-width, then slice to the shard).
433 + for e in range(E):
434 + d32 = _dequant_block_fp8(w1[e, :n_shard], w1_s[e, :nb], refined_block)
435 + d128 = _dequant_block_fp8(w1_full[e], w1_s_full[e], ckpt_block)[lo:hi]
436 + assert torch.equal(d32, d128)
437 + d32 = _dequant_block_fp8(w2[e], w2_s[e], refined_block)
438 + d128 = _dequant_block_fp8(w2_full[e], w2_s_full[e], ckpt_block)[:, lo:hi]
439 + assert torch.equal(d32, d128)
440 +
441 + quant_config = fp8_w8a8_moe_quant_config(
442 + w1_scale=w1_s,
443 + w2_scale=w2_s,
444 + block_shape=refined_block,
445 + )
446 +
447 + topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
448 +
449 + with set_current_vllm_config(vllm_config):
450 + ref_out = torch_w8a8_block_fp8_moe(
451 + a, w1, w2, w1_s, w2_s, topk_weights, topk_ids, refined_block
452 + )
453 + out = fused_experts(
454 + a, w1, w2, topk_weights, topk_ids, quant_config=quant_config
455 + )
456 +
457 + torch.testing.assert_close(out, ref_out, atol=0.035, rtol=0.035)
tests/kernels/moe/test_trtllm_bf16_moe.py
+10 -6
@@ -28,8 +28,9 @@ from vllm.model_executor.layers.fused_moe.config import (
28 from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
29 TrtLlmBf16ExpertsModular,
30 )
31 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
32 - convert_moe_weights_to_flashinfer_trtllm_block_layout,
31 +from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
32 + UnquantizedMoeBackend,
33 + convert_to_unquantized_kernel_format,
34 )
35 from vllm.platforms import current_platform
36 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
@@ -46,9 +47,8 @@ if pytest and (
47 )
48
49 # (m, n, k) = (tokens, intermediate_size_per_partition, hidden_dim).
49 -# Covers larger NvFP4-like shapes while keeping BF16's FlashInfer TRTLLM
50 -# intermediate-size multiple-of-128 requirement.
50 MNK_FACTORS = [
51 + (2, 160, 2560),
52 (2, 1024, 1024),
53 (64, 2048, 1536),
54 (64, 1024, 4096),
@@ -97,11 +97,15 @@ def test_trtllm_bf16_moe_modular_no_graph(
97 max_num_tokens=next_power_of_2(m),
98 )
99
100 - trtllm_w1, trtllm_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout(
101 - {},
100 + trtllm_w1, trtllm_w2 = convert_to_unquantized_kernel_format(
101 + UnquantizedMoeBackend.FLASHINFER_TRTLLM,
102 + moe_config,
103 w1,
104 w2,
105 )
106 + expected_n = (n + 127) // 128 * 128
107 + assert moe_config.intermediate_size_per_partition == expected_n
108 + assert trtllm_w2.numel() == e * k * expected_n
109
110 trtllm_experts = mk.FusedMoEKernel(
111 maybe_make_prepare_finalize(
tests/kernels/test_fused_gdn_post_conv.py
+4 -1
@@ -267,6 +267,7 @@ def test_fused_post_conv_l0():
267 ),
268 ],
269 )
270 +@pytest.mark.parametrize("output_gate_activation", ["silu", "sigmoid"])
271 @torch.inference_mode()
272 def test_fused_gdn_decode_post_conv_mtp_head_ratios(
273 head_ratio: int,
@@ -274,6 +275,7 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios(
275 query_lengths: tuple[int, ...],
276 state_dtype: torch.dtype,
277 norm_dtype: torch.dtype,
278 + output_gate_activation: str,
279 ) -> None:
280 if torch.cuda.get_device_capability() < (8, 0):
281 pytest.skip("fused GDN decode MTP requires compute capability 8.0+")
@@ -357,7 +359,7 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios(
359 z=output_gate,
360 eps=eps,
361 norm_before_gate=True,
360 - activation="silu",
362 + activation=output_gate_activation,
363 )
364 actual = ops.fused_gdn_decode_post_conv_mtp(
365 mixed_qkv=mixed_qkv,
@@ -374,6 +376,7 @@ def test_fused_gdn_decode_post_conv_mtp_head_ratios(
376 out=torch.empty_like(output_gate),
377 scale=scale,
378 norm_eps=eps,
379 + output_gate_activation=output_gate_activation,
380 )
381
382 output_error = (actual.float() - expected.float()).norm()
tests/model_executor/layers/test_fused_shared_expert.py
+80
@@ -18,8 +18,12 @@ from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts
18 from vllm.model_executor.layers.quantization.fp8 import Fp8Config
19 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
20 from vllm.model_executor.layers.quantization.utils.config_utils import (
21 + get_quark_ocp_mx_group_size,
22 is_shared_expert_quant_fse_compatible,
23 )
24 +from vllm.model_executor.models.qwen3_next import (
25 + _should_replicate_misaligned_shared_expert,
26 +)
27 from vllm.model_executor.models.utils import PPMissingLayer
28 from vllm.models.deepseek_v4 import quant_config as deepseek_v4_quant_config
29 from vllm.models.minimax_m3.amd import model as minimax_m3_model
@@ -605,6 +609,82 @@ def test_quark_shared_expert_fse_compatibility(
609 )
610
611
612 +def test_quark_ocp_mx_group_size_for_layer() -> None:
613 + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []})
614 + assert (
615 + get_quark_ocp_mx_group_size(
616 + quant_config,
617 + "model.layers.0.mlp.shared_expert.down_proj",
618 + )
619 + == 32
620 + )
621 +
622 + quant_config.quant_config["exclude"] = ["model.layers.0.mlp.shared_expert_gate"]
623 + assert (
624 + get_quark_ocp_mx_group_size(
625 + quant_config,
626 + "model.layers.0.mlp.shared_expert.down_proj",
627 + )
628 + == 32
629 + )
630 +
631 + quant_config.quant_config["exclude"] = [
632 + "model.layers.0.mlp.shared_expert.down_proj"
633 + ]
634 + assert (
635 + get_quark_ocp_mx_group_size(
636 + quant_config,
637 + "model.layers.0.mlp.shared_expert.down_proj",
638 + )
639 + is None
640 + )
641 +
642 +
643 +@pytest.mark.parametrize("tp_size", [1, 2, 4])
644 +def test_quark_ocp_mx_shared_expert_uses_supported_tp_shards(tp_size: int) -> None:
645 + assert not _should_replicate_misaligned_shared_expert(
646 + intermediate_size=640,
647 + tp_size=tp_size,
648 + group_size=32,
649 + enable_expert_parallel=False,
650 + is_sequence_parallel=False,
651 + )
652 +
653 +
654 +def test_quark_ocp_mx_shared_expert_rejects_plain_tp8() -> None:
655 + with pytest.raises(
656 + ValueError,
657 + match="TP size 8 produces a partition of 80",
658 + ):
659 + _should_replicate_misaligned_shared_expert(
660 + intermediate_size=640,
661 + tp_size=8,
662 + group_size=32,
663 + enable_expert_parallel=False,
664 + is_sequence_parallel=False,
665 + )
666 +
667 +
668 +def test_quark_ocp_mx_shared_expert_is_replicated_for_tep8() -> None:
669 + assert _should_replicate_misaligned_shared_expert(
670 + intermediate_size=640,
671 + tp_size=8,
672 + group_size=32,
673 + enable_expert_parallel=True,
674 + is_sequence_parallel=False,
675 + )
676 +
677 +
678 +def test_quark_ocp_mx_shared_expert_is_replicated_when_not_tp_divisible() -> None:
679 + assert _should_replicate_misaligned_shared_expert(
680 + intermediate_size=672,
681 + tp_size=8,
682 + group_size=32,
683 + enable_expert_parallel=True,
684 + is_sequence_parallel=False,
685 + )
686 +
687 +
688 def test_quark_shared_expert_fse_requires_matching_layer_quant_configs() -> None:
689 global_quant_config = {"weight": {"dtype": "fp4"}}
690 quant_config = QuarkConfig(
tests/models/qwen4_exp/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/models/qwen4_exp/test_config.py new
+200
@@ -0,0 +1,200 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from unittest.mock import patch
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.config.speculative import SpeculativeConfig
11 +from vllm.model_executor.models.config import (
12 + Qwen3_5ForConditionalGenerationConfig,
13 + Qwen4ExpForConditionalGenerationConfig,
14 +)
15 +from vllm.models.qwen4_exp.config import (
16 + Qwen4ExpConfig,
17 + Qwen4ExpTextConfig,
18 +)
19 +from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState
20 +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
21 +
22 +
23 +def _text_config(**kwargs) -> Qwen4ExpTextConfig:
24 + values = {
25 + "vocab_size": 64,
26 + "hidden_size": 16,
27 + "intermediate_size": 32,
28 + "num_hidden_layers": 2,
29 + "num_attention_heads": 2,
30 + "num_key_value_heads": 1,
31 + "head_dim": 8,
32 + "layer_types": ["linear_attention", "full_attention"],
33 + "linear_num_key_heads": 2,
34 + "linear_num_value_heads": 2,
35 + "linear_key_head_dim": 8,
36 + "linear_value_head_dim": 8,
37 + "num_experts": 0,
38 + "hc_count": 2,
39 + "hc_lowrank": 4,
40 + "ple_layer_ids": [1],
41 + "mtp_num_hidden_layers": 1,
42 + "mtp": {"hybrid": True},
43 + }
44 + values.update(kwargs)
45 + return Qwen4ExpTextConfig(**values)
46 +
47 +
48 +def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None:
49 + from vllm.models.qwen4_exp.nvidia.mtp import (
50 + Qwen4ExpMultiTokenPredictor,
51 + )
52 +
53 + model = object.__new__(Qwen4ExpMultiTokenPredictor)
54 + torch.nn.Module.__init__(model)
55 + model.hc_count = 2
56 + model.hidden_size = 4
57 + model.num_mtp_layers = 1
58 + model.layers = [
59 + lambda **kwargs: (
60 + kwargs["hidden_states"],
61 + kwargs["hidden_states"],
62 + torch.zeros(kwargs["hidden_states"].shape[0], 2),
63 + ),
64 + ]
65 + model.hyper_connection_mixer = SimpleNamespace(
66 + combine_and_mix=lambda hidden_states, block_output, injection: (
67 + hidden_states,
68 + hidden_states.unflatten(-1, (2, 4)).mean(-2),
69 + None,
70 + ),
71 + )
72 + multi_hidden = torch.arange(16, dtype=torch.float32).reshape(2, 8)
73 + pp_group = SimpleNamespace(is_first_rank=False, is_last_rank=True)
74 +
75 + with patch(
76 + "vllm.models.qwen4_exp.nvidia.mtp.get_pp_group",
77 + return_value=pp_group,
78 + ):
79 + sample_hidden, returned_multi_hidden = model.forward(
80 + input_ids=None,
81 + positions=torch.arange(2),
82 + intermediate_tensors={"hidden_states": multi_hidden},
83 + )
84 +
85 + torch.testing.assert_close(
86 + sample_hidden,
87 + multi_hidden.unflatten(-1, (2, 4)).mean(dim=-2),
88 + )
89 + assert returned_multi_hidden is multi_hidden
90 +
91 +
92 +@pytest.mark.parametrize("wrapped_config", [False, True])
93 +def test_qwen4_exp_mtp_override_sets_draft_config(
94 + wrapped_config: bool,
95 +) -> None:
96 + text_config = _text_config(
97 + architectures=["Qwen4ExpForCausalLM"],
98 + index_share_for_mtp_iteration=True,
99 + )
100 + config = (
101 + Qwen4ExpConfig(
102 + architectures=["Qwen4ExpForConditionalGeneration"],
103 + text_config=text_config,
104 + )
105 + if wrapped_config
106 + else text_config
107 + )
108 +
109 + draft_config = SpeculativeConfig.hf_config_override(config)
110 +
111 + assert draft_config.index_share_for_mtp_iteration is True
112 + assert draft_config.model_type == "qwen4_exp_mtp"
113 + assert draft_config.architectures == ["Qwen4ExpMTP"]
114 + assert draft_config.hc_mult == 2
115 + assert draft_config.n_predict == 1
116 +
117 +
118 +@pytest.mark.parametrize("ple_layer_ids", [[1], []])
119 +def test_qwen4_exp_rejects_pipeline_parallel_only_with_ple(ple_layer_ids) -> None:
120 + """PLE needs raw input_ids, which non-first pipeline ranks never see. The
121 + rest of the architecture is PP-capable, so the refusal must be conditional
122 + -- and must land before the engine spends time loading weights."""
123 + vllm_config = SimpleNamespace(
124 + model_config=SimpleNamespace(
125 + hf_text_config=_text_config(ple_layer_ids=ple_layer_ids),
126 + multimodal_config=None,
127 + ),
128 + parallel_config=SimpleNamespace(
129 + pipeline_parallel_size=2, enable_dbo=False, ubatch_size=1
130 + ),
131 + speculative_config=None,
132 + )
133 + with patch.object(
134 + Qwen3_5ForConditionalGenerationConfig, "verify_and_update_config"
135 + ):
136 + if ple_layer_ids:
137 + with pytest.raises(NotImplementedError, match="pipeline_parallel_size=1"):
138 + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(
139 + vllm_config
140 + )
141 + else:
142 + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config)
143 +
144 +
145 +def test_qwen4_exp_model_state_prepares_ngram_context() -> None:
146 + model_state = object.__new__(Qwen4ExpModelState)
147 + model_state.uses_ngram_embedding = True
148 + model_state.ngram_context_len = 3
149 + model_state.ngram_eos_token_id = 99
150 + model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32)
151 + model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64)
152 + model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)
153 +
154 + input_batch = SimpleNamespace(
155 + num_reqs=2,
156 + num_reqs_after_padding=3,
157 + idx_mapping=torch.tensor([1, 0]),
158 + query_start_loc=torch.tensor([0, 2, 3, 3], dtype=torch.int32),
159 + )
160 + req_states = SimpleNamespace(
161 + num_computed_tokens=SimpleNamespace(gpu=torch.tensor([3, 1])),
162 + all_token_ids=SimpleNamespace(
163 + gpu=torch.tensor([[1, 2, 3, 4], [20, 21, 22, 23]], dtype=torch.int32)
164 + ),
165 + )
166 +
167 + with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}):
168 + model_inputs = model_state.prepare_inputs(input_batch, req_states)
169 +
170 + torch.testing.assert_close(
171 + model_inputs["query_start_loc"],
172 + torch.tensor([0, 2, 3, 3], dtype=torch.int32),
173 + )
174 + torch.testing.assert_close(
175 + model_inputs["ngram_context"],
176 + torch.tensor([[99, 99, 20], [1, 2, 3], [99, 99, 99]], dtype=torch.int32),
177 + )
178 +
179 +
180 +def test_qwen4_exp_model_state_prepares_stable_dummy_ngram_inputs() -> None:
181 + model_state = object.__new__(Qwen4ExpModelState)
182 + model_state.uses_ngram_embedding = True
183 + model_state.ngram_eos_token_id = 99
184 + model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32)
185 + model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32)
186 +
187 + with patch.object(MambaHybridModelState, "prepare_dummy_inputs", return_value={}):
188 + first = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=8)
189 + query_start_loc_ptr = first["query_start_loc"].data_ptr()
190 + ngram_context_ptr = first["ngram_context"].data_ptr()
191 + second = model_state.prepare_dummy_inputs(num_reqs=3, num_tokens=8)
192 +
193 + torch.testing.assert_close(
194 + second["query_start_loc"], torch.tensor([0, 2, 5, 8], dtype=torch.int32)
195 + )
196 + torch.testing.assert_close(
197 + second["ngram_context"], torch.full((3, 3), 99, dtype=torch.int32)
198 + )
199 + assert second["query_start_loc"].data_ptr() == query_start_loc_ptr
200 + assert second["ngram_context"].data_ptr() == ngram_context_ptr
tests/models/qwen4_exp/test_hc_ops.py new
+94
@@ -0,0 +1,94 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.models.qwen4_exp.nvidia.ops.hc import (
8 + grouped_gemma_rmsnorm,
9 + hc_combine,
10 + hc_combine_norm,
11 + hc_gate_mix,
12 +)
13 +from vllm.platforms import current_platform
14 +from vllm.triton_utils import HAS_TRITON
15 +
16 +pytestmark = pytest.mark.skipif(
17 + not current_platform.is_cuda() or not HAS_TRITON,
18 + reason="HC kernels require CUDA and Triton",
19 +)
20 +
21 +HC = 4
22 +HIDDEN_SIZE = 2560
23 +HYPER_HIDDEN_SIZE = HC * HIDDEN_SIZE
24 +EPS = 1e-6
25 +
26 +
27 +def test_grouped_gemma_rmsnorm() -> None:
28 + torch.manual_seed(0)
29 + x = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
30 + weight = torch.randn(HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
31 +
32 + actual = grouped_gemma_rmsnorm(x, weight, EPS, HC)
33 +
34 + grouped = x.float().unflatten(-1, (HC, HIDDEN_SIZE))
35 + variance = grouped.square().mean(-1, keepdim=True)
36 + expected = grouped * torch.rsqrt(variance + EPS)
37 + expected = expected.flatten(-2) * (1.0 + weight.float())
38 + torch.testing.assert_close(actual, expected.to(torch.bfloat16))
39 +
40 +
41 +def test_hc_gate_mix() -> None:
42 + torch.manual_seed(0)
43 + x = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
44 + gate = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
45 +
46 + actual = hc_gate_mix(x, gate, HC)
47 + expected = (
48 + torch.sigmoid(gate.float().unflatten(-1, (HC, HIDDEN_SIZE)))
49 + * x.float().unflatten(-1, (HC, HIDDEN_SIZE))
50 + ).mean(-2)
51 +
52 + torch.testing.assert_close(actual, expected.to(torch.bfloat16))
53 +
54 +
55 +def test_hc_combine() -> None:
56 + torch.manual_seed(0)
57 + block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
58 + residual = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
59 + injection = torch.randn(2, HC, dtype=torch.bfloat16, device="cuda")
60 +
61 + actual = hc_combine(residual, block_output, injection, HC)
62 + injection_weight = 2.0 * torch.sigmoid(injection.float() / HC)
63 + expected = residual.float().unflatten(-1, (HC, HIDDEN_SIZE))
64 + expected = expected + block_output.float().unsqueeze(
65 + -2
66 + ) * injection_weight.unsqueeze(-1)
67 +
68 + torch.testing.assert_close(actual, expected.flatten(-2).to(torch.bfloat16))
69 +
70 +
71 +def test_hc_combine_norm() -> None:
72 + torch.manual_seed(0)
73 + block_output = torch.randn(2, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
74 + residual = torch.randn(2, HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
75 + injection = torch.randn(2, HC, dtype=torch.bfloat16, device="cuda")
76 + weight = torch.randn(HYPER_HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda")
77 +
78 + actual, actual_norm = hc_combine_norm(
79 + residual, block_output, injection, weight, EPS, HC
80 + )
81 +
82 + injection_weight = 2.0 * torch.sigmoid(injection.float() / HC)
83 + expected = residual.float().unflatten(-1, (HC, HIDDEN_SIZE))
84 + expected = expected + block_output.float().unsqueeze(
85 + -2
86 + ) * injection_weight.unsqueeze(-1)
87 + expected = expected.flatten(-2).to(residual.dtype)
88 + grouped = expected.float().unflatten(-1, (HC, HIDDEN_SIZE))
89 + variance = grouped.square().mean(-1, keepdim=True)
90 + expected_norm = grouped * torch.rsqrt(variance + EPS)
91 + expected_norm = expected_norm.flatten(-2) * (1.0 + weight.float())
92 +
93 + torch.testing.assert_close(actual, expected)
94 + torch.testing.assert_close(actual_norm, expected_norm.to(torch.bfloat16))
tests/models/qwen4_exp/test_ple.py new
+401
@@ -0,0 +1,401 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from functools import partial
5 +from types import SimpleNamespace
6 +
7 +import pytest
8 +import torch
9 +from torch import nn
10 +from torch.nn import functional as F
11 +
12 +import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module
13 +import vllm.model_executor.parameter as parameter_module
14 +from vllm.model_executor.layers.quantization.fp8 import Fp8Config
15 +from vllm.models.qwen4_exp.common.ple import (
16 + PLEShardOverlap,
17 + compute_ple_shard_overlap,
18 + copy_ple_embedding_shard_,
19 +)
20 +from vllm.models.qwen4_exp.nvidia import ple_layer as ple_layer_module
21 +from vllm.models.qwen4_exp.nvidia.ple_layer import (
22 + Qwen4ExpNGramEmbedding,
23 + Qwen4ExpPLEFp8EmbeddingMethod,
24 + Qwen4ExpPLELayer,
25 + _get_ple_embedding_quant_method,
26 +)
27 +
28 +
29 +def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
30 + module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
31 + nn.Module.__init__(module)
32 + module.split_ngram_parts = 2
33 + module.register_buffer("layer_multipliers", torch.zeros(1, dtype=torch.long))
34 + module.register_buffer("ngram_heads_offsets", torch.zeros(1, dtype=torch.long))
35 + module.register_buffer("ngram_heads_vocab_sizes", torch.zeros(1, dtype=torch.long))
36 + module.ngram_embedding = SimpleNamespace(
37 + org_vocab_size=8,
38 + embedding_dim=2,
39 + weight=nn.Parameter(torch.full((4, 2), -1.0)),
40 + shard_indices=SimpleNamespace(
41 + org_vocab_start_index=2,
42 + org_vocab_end_index=6,
43 + ),
44 + )
45 + _set_test_embedding_weight_loader(module.ngram_embedding)
46 + return module
47 +
48 +
49 +def _set_test_embedding_weight_loader(embedding) -> None:
50 + embedding.weight.weight_loader = partial(
51 + copy_ple_embedding_shard_,
52 + tp_start=embedding.shard_indices.org_vocab_start_index,
53 + tp_end=embedding.shard_indices.org_vocab_end_index,
54 + )
55 +
56 +
57 +def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding:
58 + module = _make_ngram_embedding_for_load_test()
59 + embedding = nn.Module()
60 + embedding.org_vocab_size = 8
61 + embedding.embedding_dim = 2
62 + embedding.shard_indices = SimpleNamespace(
63 + org_vocab_start_index=2,
64 + org_vocab_end_index=6,
65 + )
66 + embedding.register_parameter(
67 + "weight",
68 + nn.Parameter(
69 + torch.full((4, 2), -1.0).to(torch.float8_e4m3fn),
70 + requires_grad=False,
71 + ),
72 + )
73 + embedding.register_parameter(
74 + "weight_scale",
75 + nn.Parameter(torch.zeros(1, dtype=torch.bfloat16), requires_grad=False),
76 + )
77 + _set_test_embedding_weight_loader(embedding)
78 + module.ngram_embedding = embedding
79 + return module
80 +
81 +
82 +def test_ple_shard_overlap_and_copy() -> None:
83 + overlap = compute_ple_shard_overlap(
84 + checkpoint_start=2, checkpoint_rows=5, tp_start=4, tp_end=8
85 + )
86 + assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3)
87 +
88 + destination = torch.full((4, 2), -1.0)
89 + loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2)
90 + copied = copy_ple_embedding_shard_(
91 + destination,
92 + loaded,
93 + checkpoint_start=2,
94 + tp_start=4,
95 + tp_end=8,
96 + )
97 +
98 + assert copied == 3
99 + torch.testing.assert_close(destination[:3], loaded[2:5].float())
100 + torch.testing.assert_close(destination[3], torch.tensor([-1.0, -1.0]))
101 +
102 +
103 +def test_ple_shard_copy_is_a_noop_without_overlap() -> None:
104 + destination = torch.ones(4, 2)
105 + copied = copy_ple_embedding_shard_(
106 + destination,
107 + torch.zeros(2, 2),
108 + checkpoint_start=10,
109 + tp_start=4,
110 + tp_end=8,
111 + )
112 +
113 + assert copied == 0
114 + assert torch.equal(destination, torch.ones_like(destination))
115 +
116 +
117 +def test_ngram_embedding_loads_shards_and_ignores_legacy_token_lookup() -> None:
118 + module = _make_ngram_embedding_for_load_test()
119 + shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2)
120 + shard_1 = torch.arange(8, 16, dtype=torch.float32).reshape(4, 2)
121 +
122 + loaded = module.load_weights(
123 + [
124 + ("ngram_embedding.shard_0.weight", shard_0),
125 + ("ngram_embedding.shard_1.weight", shard_1),
126 + ("token_lookup", torch.tensor([2, 1, 0])),
127 + ]
128 + )
129 +
130 + assert loaded == {"ngram_embedding.weight"}
131 + torch.testing.assert_close(
132 + module.ngram_embedding.weight,
133 + torch.cat((shard_0[2:4], shard_1[0:2])),
134 + )
135 +
136 +
137 +def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None:
138 + module = _make_fp8_ngram_embedding_for_load_test()
139 + shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn)
140 + shard_1 = (
141 + torch.arange(8, 16, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn)
142 + )
143 + weight_scale = torch.tensor([0.25], dtype=torch.bfloat16)
144 +
145 + loaded = module.load_weights(
146 + [
147 + ("ngram_embedding.shard_0.weight", shard_0),
148 + ("ngram_embedding.shard_1.weight", shard_1),
149 + ("ngram_embedding.weight_scale", weight_scale),
150 + ]
151 + )
152 +
153 + assert loaded == {"ngram_embedding.weight", "ngram_embedding.weight_scale"}
154 + assert module.ngram_embedding.weight.dtype == torch.float8_e4m3fn
155 + assert torch.equal(
156 + module.ngram_embedding.weight.float(),
157 + torch.cat((shard_0[2:4], shard_1[0:2])).float(),
158 + )
159 + assert torch.equal(module.ngram_embedding.weight_scale, weight_scale)
160 +
161 +
162 +def _make_fp8_embedding_layer(
163 + monkeypatch: pytest.MonkeyPatch,
164 +) -> embedding_module.VocabParallelEmbedding:
165 + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
166 + monkeypatch.setattr(
167 + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1
168 + )
169 + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)
170 + monkeypatch.setattr(
171 + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1
172 + )
173 + method = Qwen4ExpPLEFp8EmbeddingMethod()
174 + layer = embedding_module.VocabParallelEmbedding(
175 + 3,
176 + 2,
177 + params_dtype=torch.bfloat16,
178 + padding_size=1,
179 + quant_method=method,
180 + )
181 + weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
182 + layer.weight.data.copy_(weight.to(torch.float8_e4m3fn))
183 + layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=torch.bfloat16))
184 + return layer
185 +
186 +
187 +def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None:
188 + layer = _make_fp8_embedding_layer(monkeypatch)
189 + quantized_output = layer(torch.tensor([2, 0]))
190 + ple_layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
191 + nn.Module.__init__(ple_layer)
192 + ple_layer.ple_embedding = nn.Module()
193 + ple_layer.ple_embedding.ngram_embedding = layer
194 +
195 + output = ple_layer._dequantize_embeddings(
196 + quantized_output,
197 + torch.bfloat16,
198 + )
199 +
200 + assert layer.weight.dtype == torch.float8_e4m3fn
201 + assert layer.weight_scale.dtype == torch.bfloat16
202 + assert quantized_output.dtype == torch.float8_e4m3fn
203 + assert output.dtype == torch.bfloat16
204 + weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]])
205 + torch.testing.assert_close(output, (weight[[2, 0]] * 0.25).bfloat16())
206 +
207 +
208 +def test_ple_fp8_embedding_uses_int8_for_tp_reduce(monkeypatch) -> None:
209 + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)
210 + monkeypatch.setattr(
211 + embedding_module, "get_tensor_model_parallel_world_size", lambda: 2
212 + )
213 + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)
214 + monkeypatch.setattr(
215 + parameter_module, "get_tensor_model_parallel_world_size", lambda: 2
216 + )
217 + monkeypatch.setattr(
218 + embedding_module,
219 + "get_masked_input_and_mask",
220 + lambda *args: (
221 + torch.tensor([0, 0]),
222 + torch.tensor([False, True]),
223 + ),
224 + )
225 + reduced_dtypes = []
226 +
227 + def all_reduce(tensor: torch.Tensor) -> torch.Tensor:
228 + reduced_dtypes.append(tensor.dtype)
229 + return tensor.clone()
230 +
231 + monkeypatch.setattr(
232 + embedding_module,
233 + "tensor_model_parallel_all_reduce",
234 + all_reduce,
235 + )
236 + layer = embedding_module.VocabParallelEmbedding(
237 + 4,
238 + 2,
239 + params_dtype=torch.bfloat16,
240 + padding_size=1,
241 + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(),
242 + )
243 + layer.weight.data.copy_(
244 + torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e4m3fn)
245 + )
246 +
247 + output = layer(torch.tensor([0, 2]))
248 +
249 + assert reduced_dtypes == [torch.int8]
250 + assert output.dtype == torch.float8_e4m3fn
251 + torch.testing.assert_close(output[0].float(), layer.weight[0].float())
252 + assert torch.count_nonzero(output[1].float()) == 0
253 +
254 +
255 +def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None:
256 + prefix = "model.layers.1.ple.ple_embedding.ngram_embedding"
257 + quant_config = Fp8Config(
258 + is_checkpoint_fp8_serialized=True,
259 + ignored_layers=[],
260 + weight_block_size=[128, 128],
261 + )
262 + assert isinstance(
263 + _get_ple_embedding_quant_method(quant_config, prefix),
264 + Qwen4ExpPLEFp8EmbeddingMethod,
265 + )
266 +
267 + quant_config.ignored_layers = [f"{prefix}.shard_0"]
268 + assert _get_ple_embedding_quant_method(quant_config, prefix) is None
269 +
270 +
271 +def test_ple_ngram_ids_custom_op_uses_current_request_layout(monkeypatch) -> None:
272 + class RuntimeNGramEmbedding(nn.Module):
273 + def compute_ngram_ids(
274 + self,
275 + input_ids: torch.Tensor,
276 + query_start_loc: torch.Tensor,
277 + ngram_context: torch.Tensor,
278 + ) -> torch.Tensor:
279 + del input_ids, ngram_context
280 + num_reqs = query_start_loc.numel() - 1
281 + return torch.full((4, 2), num_reqs, dtype=torch.long)
282 +
283 + layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
284 + nn.Module.__init__(layer)
285 + layer.ple_embedding = RuntimeNGramEmbedding()
286 + monkeypatch.setattr(
287 + ple_layer_module,
288 + "get_forward_context",
289 + lambda: SimpleNamespace(no_compile_layers={"ple": layer}),
290 + )
291 + input_ids = torch.arange(4)
292 + ngram_context = torch.zeros(2, 2, dtype=torch.long)
293 + output = torch.empty(4, 2, dtype=torch.long)
294 +
295 + ple_layer_module.qwen4_exp_compute_ple_ngram_ids(
296 + input_ids,
297 + torch.tensor([0, 4]),
298 + ngram_context,
299 + output,
300 + "ple",
301 + )
302 + assert torch.equal(output, torch.ones_like(output))
303 +
304 + ple_layer_module.qwen4_exp_compute_ple_ngram_ids(
305 + input_ids,
306 + torch.tensor([0, 2, 4]),
307 + ngram_context,
308 + output,
309 + "ple",
310 + )
311 + assert torch.equal(output, torch.full_like(output, 2))
312 +
313 +
314 +def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None:
315 + module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
316 + nn.Module.__init__(module)
317 + module.conv_state_len = 6
318 + module.short_conv_dilation = 2
319 +
320 + conv_weights = torch.tensor([[0.25, -0.5, 0.75, 1.0]])
321 + conv_state = torch.zeros(2, 1, 9)
322 + conv_state[1] = torch.arange(1, 10, dtype=torch.float32).reshape(1, 9)
323 + first_inputs = torch.tensor([[10.0], [20.0], [30.0], [40.0]])
324 + initial_state = conv_state[1:].clone()
325 + first_history = torch.cat(
326 + (initial_state[..., : module.conv_state_len], first_inputs.T.unsqueeze(0)),
327 + dim=-1,
328 + )
329 +
330 + graph_padded_inputs = F.pad(first_inputs, (0, 0, 0, 4))
331 + first_output = module._short_conv_dilated_spec_batched(
332 + graph_padded_inputs,
333 + conv_state,
334 + conv_weights,
335 + torch.tensor([1, 0]),
336 + torch.tensor([0, 4, 4]),
337 + torch.tensor([1, 0]),
338 + spec_query_len=4,
339 + )
340 +
341 + expected_first_output = F.silu(
342 + F.conv1d(
343 + first_history,
344 + conv_weights.unsqueeze(1),
345 + groups=1,
346 + dilation=module.short_conv_dilation,
347 + )
348 + ).transpose(1, 2)[0]
349 + expected_first_state = first_history[..., 1:10]
350 + torch.testing.assert_close(first_output[:4], expected_first_output)
351 + assert torch.count_nonzero(first_output[4:]) == 0
352 + assert torch.count_nonzero(conv_state[0]) == 0
353 + torch.testing.assert_close(conv_state[1:], expected_first_state)
354 +
355 + second_inputs = torch.tensor([[50.0], [60.0]])
356 + rollback_state = expected_first_state[..., 1:7]
357 + padded_second_inputs = F.pad(second_inputs.T.unsqueeze(0), (0, 2))
358 + second_history = torch.cat((rollback_state, padded_second_inputs), dim=-1)
359 + expected_second_state = expected_first_state.clone()
360 + expected_second_state[..., :7] = second_history[..., 1:8]
361 +
362 + second_output = module._short_conv_dilated_spec_batched(
363 + second_inputs,
364 + conv_state,
365 + conv_weights,
366 + torch.tensor([1]),
367 + torch.tensor([0, 2]),
368 + torch.tensor([2]),
369 + spec_query_len=4,
370 + )
371 +
372 + expected_second_output = F.silu(
373 + F.conv1d(
374 + second_history,
375 + conv_weights.unsqueeze(1),
376 + groups=1,
377 + dilation=module.short_conv_dilation,
378 + )
379 + ).transpose(1, 2)[0, :2]
380 + torch.testing.assert_close(second_output, expected_second_output)
381 + torch.testing.assert_close(conv_state[1:], expected_second_state)
382 +
383 +
384 +def test_ple_short_conv_uses_fallback_when_profile_metadata_is_omitted(
385 + monkeypatch: pytest.MonkeyPatch,
386 +) -> None:
387 + module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
388 + nn.Module.__init__(module)
389 + module.prefix = "model.layers.1.ple"
390 + inputs = torch.arange(6, dtype=torch.float32).reshape(3, 2)
391 + expected = inputs + 1
392 + monkeypatch.setattr(module, "_short_conv_fallback", lambda _: expected)
393 + monkeypatch.setattr(
394 + ple_layer_module,
395 + "get_forward_context",
396 + lambda: SimpleNamespace(attn_metadata={"model.layers.0.self_attn": object()}),
397 + )
398 +
399 + output = module._short_conv(inputs)
400 +
401 + assert output is expected
tests/models/qwen4_exp/test_qsa_amd.py new
+297
@@ -0,0 +1,297 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from types import SimpleNamespace
6 +from typing import Any
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
12 +from vllm.models.qwen4_exp.amd import (
13 + model as _qwen4_exp_model, # noqa: F401
14 +)
15 +from vllm.models.qwen4_exp.amd import ple_layer as ple_layer_module
16 +from vllm.models.qwen4_exp.amd.indexer_qsa import (
17 + apply_qsa_rmsnorm,
18 + apply_qsa_rope,
19 +)
20 +from vllm.models.qwen4_exp.amd.ops import qsa as qsa_ops
21 +from vllm.platforms import current_platform
22 +from vllm.triton_utils import HAS_TRITON
23 +
24 +pytestmark = pytest.mark.skipif(
25 + not current_platform.is_rocm(),
26 + reason="AMD QSA requires ROCm",
27 +)
28 +
29 +requires_qsa_kernels = pytest.mark.skipif(
30 + not HAS_TRITON,
31 + reason="AMD QSA kernels require Triton",
32 +)
33 +
34 +
35 +def test_ple_ngram_embedding_custom_op_uses_resident_weight(
36 + monkeypatch: pytest.MonkeyPatch,
37 +) -> None:
38 + layer_name = "model.layers.0.ple"
39 + layer = ple_layer_module.Qwen4ExpPLELayer.__new__(ple_layer_module.Qwen4ExpPLELayer)
40 + torch.nn.Module.__init__(layer)
41 + layer.ple_embedding = torch.nn.Module()
42 + layer.ple_embedding.ngram_embedding = torch.nn.Embedding(8, 3)
43 + context = SimpleNamespace(no_compile_layers={layer_name: layer})
44 + monkeypatch.setattr(ple_layer_module, "get_forward_context", lambda: context)
45 +
46 + ngram_ids = torch.tensor([[0, 1], [2, 3]])
47 + output = torch.empty(2, 6)
48 + ple_layer_module.qwen4_exp_amd_ple_ngram_embedding(
49 + ngram_ids,
50 + output,
51 + layer_name,
52 + )
53 +
54 + expected = layer.ple_embedding.ngram_embedding(ngram_ids).flatten(-2)
55 + torch.testing.assert_close(output, expected)
56 +
57 +
58 +def _qsa_sparse_paged_attention_reference(
59 + q: torch.Tensor,
60 + k_cache: torch.Tensor,
61 + v_cache: torch.Tensor,
62 + logical_indices: torch.Tensor,
63 + block_table: torch.Tensor,
64 + token_to_req: torch.Tensor,
65 + softmax_scale: float,
66 +) -> torch.Tensor:
67 + output = torch.zeros_like(q)
68 + repeats = q.shape[1] // k_cache.shape[2]
69 + page_size = k_cache.shape[1]
70 + for row in range(q.shape[0]):
71 + logical = logical_indices[row]
72 + logical = logical[logical >= 0].long()
73 + if not logical.numel():
74 + continue
75 + request = token_to_req[row].long()
76 + pages = block_table[request, logical // page_size].long()
77 + offsets = logical % page_size
78 + keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1)
79 + values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1)
80 + scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float())
81 + probabilities = torch.softmax(scores * softmax_scale, dim=-1)
82 + output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to(
83 + q.dtype
84 + )
85 + return output
86 +
87 +
88 +def test_qsa_rope_uses_platform_dispatch() -> None:
89 + tensor = torch.arange(16, dtype=torch.float32).reshape(2, 2, 4)
90 + positions = torch.tensor([0, 1])
91 + calls = []
92 +
93 + def apply_rotary_emb(
94 + rotary_input: torch.Tensor,
95 + cos: torch.Tensor,
96 + sin: torch.Tensor,
97 + ) -> torch.Tensor:
98 + calls.append((rotary_input, cos, sin))
99 + return rotary_input + 1
100 +
101 + rotary_emb = SimpleNamespace(
102 + rotary_dim=2,
103 + apply_rotary_emb=apply_rotary_emb,
104 + _match_cos_sin_cache_dtype=lambda _: torch.zeros(2, 4),
105 + )
106 +
107 + output = apply_qsa_rope(rotary_emb, positions, tensor)
108 +
109 + assert len(calls) == 1
110 + torch.testing.assert_close(output[..., :2], tensor[..., :2] + 1)
111 + torch.testing.assert_close(output[..., 2:], tensor[..., 2:])
112 +
113 +
114 +def test_qsa_rmsnorm_uses_portable_implementation(default_vllm_config) -> None:
115 + norm = GemmaRMSNorm(4, eps=1e-6)
116 + norm.weight.data.copy_(torch.tensor([0.1, -0.2, 0.3, -0.4]))
117 + tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4)
118 + output = apply_qsa_rmsnorm(norm, tensor)
119 +
120 + torch.testing.assert_close(output, norm.forward_native(tensor))
121 +
122 +
123 +def test_qsa_selection_uses_portable_topk_on_rocm(
124 + monkeypatch: pytest.MonkeyPatch,
125 +) -> None:
126 + rows = 2
127 + token_topk = 8
128 + compress_ratio = 4
129 + block_topk = token_topk // compress_ratio
130 + blocks = torch.empty((rows, block_topk), dtype=torch.int32)
131 + visible_blocks = torch.tensor([4, 6], dtype=torch.int32)
132 + logits = torch.empty((rows, 8), dtype=torch.float32)
133 + selection_output = torch.empty(
134 + (rows, token_topk + compress_ratio - 1), dtype=torch.int32
135 + )
136 + topk_call: dict[str, Any] = {}
137 +
138 + monkeypatch.setattr(
139 + qsa_ops,
140 + "qsa_mqa_paged",
141 + lambda *args, **kwargs: (logits, visible_blocks),
142 + )
143 + monkeypatch.setattr(qsa_ops.current_platform, "is_cuda", lambda: False)
144 +
145 + def top_k_per_row_decode(
146 + topk_logits,
147 + next_n,
148 + seq_lens,
149 + raw_topk_indices,
150 + num_rows,
151 + stride0,
152 + stride1,
153 + topk_tokens,
154 + ) -> None:
155 + topk_call.update(
156 + logits=topk_logits,
157 + next_n=next_n,
158 + seq_lens=seq_lens,
159 + raw_topk_indices=raw_topk_indices,
160 + num_rows=num_rows,
161 + strides=(stride0, stride1),
162 + topk_tokens=topk_tokens,
163 + )
164 + raw_topk_indices.zero_()
165 +
166 + def expand_qsa_block_indices(*args) -> None:
167 + args[-1].fill_(-1)
168 +
169 + monkeypatch.setattr(qsa_ops.ops, "top_k_per_row_decode", top_k_per_row_decode)
170 + monkeypatch.setattr(
171 + qsa_ops,
172 + "expand_qsa_block_indices_cuda",
173 + expand_qsa_block_indices,
174 + )
175 +
176 + output = qsa_ops.qsa_select_paged_tokens(
177 + torch.empty((rows, 1, 1)),
178 + torch.empty((1, 4, 1, 1)),
179 + torch.empty((1, 2), dtype=torch.int32),
180 + torch.zeros(rows, dtype=torch.int32),
181 + torch.arange(rows, dtype=torch.int32),
182 + torch.tensor([8], dtype=torch.int32),
183 + token_topk,
184 + compress_ratio,
185 + selection_output,
186 + )
187 +
188 + assert output is selection_output
189 + assert torch.all(output == -1)
190 + assert topk_call["logits"] is logits
191 + assert topk_call["next_n"] == 1
192 + assert topk_call["seq_lens"] is visible_blocks
193 + assert topk_call["raw_topk_indices"].shape == blocks.shape
194 + assert topk_call["num_rows"] == rows
195 + assert topk_call["strides"] == (logits.stride(0), logits.stride(1))
196 + assert topk_call["topk_tokens"] == block_topk
197 +
198 +
199 +@requires_qsa_kernels
200 +@pytest.mark.parametrize(
201 + ("num_rows", "num_query_heads", "num_kv_heads", "page_size"),
202 + [
203 + pytest.param(1, 24, 2, 1792, id="tp1_split64"),
204 + pytest.param(16, 12, 1, 1792, id="tp2_split32"),
205 + pytest.param(32, 6, 1, 1024, id="tp4_split8"),
206 + pytest.param(257, 6, 1, 1024, id="tp4_split4"),
207 + pytest.param(513, 6, 1, 1024, id="tp4_split1"),
208 + ],
209 +)
210 +def test_qsa_sparse_paged_attention_matches_reference(
211 + num_rows: int,
212 + num_query_heads: int,
213 + num_kv_heads: int,
214 + page_size: int,
215 +) -> None:
216 + torch.manual_seed(2)
217 + head_dim = 256
218 + num_requests = 2
219 + num_selected_pages = 64
220 + num_pages_per_request = num_selected_pages + 1
221 + num_cache_blocks = num_requests * num_pages_per_request
222 + indexer_budget = 2048
223 + indexer_compress_ratio = 4
224 + selection_width = indexer_budget + indexer_compress_ratio - 1
225 + q = torch.randn(
226 + num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16
227 + )
228 + kv_cache = torch.randn(
229 + num_cache_blocks,
230 + page_size,
231 + num_kv_heads,
232 + 2 * head_dim,
233 + device="cuda",
234 + dtype=torch.bfloat16,
235 + )
236 + k_cache, v_cache = kv_cache.split(head_dim, dim=-1)
237 + block_table = (
238 + torch.randperm(num_cache_blocks, device="cuda")
239 + .reshape(num_requests, num_pages_per_request)
240 + .to(torch.int32)
241 + )
242 + rows_per_request = math.ceil(num_rows / num_requests)
243 + row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32)
244 + token_to_req = row_indices // rows_per_request
245 + request_row_counts = torch.tensor(
246 + [rows_per_request, num_rows - rows_per_request],
247 + device="cuda",
248 + dtype=torch.int32,
249 + )
250 +
251 + context_length = num_pages_per_request * page_size - 1
252 + block_topk = indexer_budget // indexer_compress_ratio
253 + compressed_blocks_per_page = page_size // indexer_compress_ratio
254 + selection = torch.arange(block_topk, device="cuda")
255 + selected_pages = selection % num_selected_pages
256 + selected_offsets = selection // num_selected_pages
257 + row_shifts = 2 * row_indices.unsqueeze(1)
258 + selected_offsets = (selected_offsets + row_shifts) % compressed_blocks_per_page
259 + block_indices = (selected_pages * compressed_blocks_per_page + selected_offsets).to(
260 + torch.int32
261 + )
262 + rows_within_request = row_indices % rows_per_request
263 + query_positions = (
264 + context_length - request_row_counts[token_to_req.long()] + rows_within_request
265 + ).to(torch.int64)
266 + sequence_lengths = torch.full(
267 + (num_requests,), context_length, device="cuda", dtype=torch.int32
268 + )
269 + logical_indices = qsa_ops.expand_qsa_block_indices_cuda(
270 + block_indices,
271 + query_positions,
272 + sequence_lengths,
273 + token_to_req,
274 + indexer_compress_ratio,
275 + indexer_budget,
276 + )
277 + assert logical_indices.shape == (num_rows, selection_width)
278 +
279 + actual = qsa_ops.qsa_sparse_paged_attention(
280 + q,
281 + k_cache,
282 + v_cache,
283 + logical_indices,
284 + block_table,
285 + token_to_req,
286 + )
287 + expected = _qsa_sparse_paged_attention_reference(
288 + q,
289 + k_cache,
290 + v_cache,
291 + logical_indices,
292 + block_table,
293 + token_to_req,
294 + q.shape[-1] ** -0.5,
295 + )
296 +
297 + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
tests/models/qwen4_exp/test_qsa_pre_indexer.py new
+297
@@ -0,0 +1,297 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for the fused QSA pre-indexer."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.models.qwen4_exp.common.qsa_cache import (
9 + canonical_qsa_rope_positions,
10 + circular_qsa_slot_mapping,
11 + compressed_qsa_slot_mapping,
12 +)
13 +from vllm.models.qwen4_exp.nvidia.indexer_qsa import apply_qsa_rope
14 +from vllm.models.qwen4_exp.nvidia.ops.qsa import (
15 + qsa_compress_groups_with_ratio,
16 + qsa_store_cache_rows,
17 +)
18 +from vllm.models.qwen4_exp.nvidia.ops.qsa_pre_indexer import (
19 + qsa_pre_indexer,
20 +)
21 +from vllm.platforms import current_platform
22 +from vllm.triton_utils import HAS_TRITON
23 +
24 +requires_qsa_kernels = pytest.mark.skipif(
25 + not current_platform.is_cuda() or not HAS_TRITON,
26 + reason="QSA kernels require CUDA and Triton",
27 +)
28 +
29 +HQ, D = 4, 128
30 +CR = 4
31 +MROPE_SECTION = (11, 11, 10)
32 +EPS = 1e-6
33 +BLOCK_SIZE = 16
34 +COMP_PAGE = BLOCK_SIZE // CR
35 +ROPE_POS_OFFSET = D
36 +RTOL = 1.6e-2
37 +ATOL = 1e-2
38 +MIXED_BATCH = ([260, 259, 138], [1, 1, 37], [8, 8, 8])
39 +
40 +
41 +def _make_block_table(block_counts):
42 + num_blocks = sum(block_counts)
43 + table = torch.full((len(block_counts), max(block_counts)), -1, dtype=torch.int32)
44 + physical_blocks = torch.randperm(num_blocks)
45 + offset = 0
46 + for request, count in enumerate(block_counts):
47 + table[request, :count] = physical_blocks[offset : offset + count]
48 + offset += count
49 + return table, num_blocks
50 +
51 +
52 +@requires_qsa_kernels
53 +@pytest.mark.usefixtures("default_vllm_config")
54 +@pytest.mark.parametrize(
55 + "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens",
56 + [
57 + pytest.param(True, True, True, 4, *MIXED_BATCH, id="mrope"),
58 + pytest.param(False, False, False, 4, *MIXED_BATCH, id="text"),
59 + pytest.param(True, False, True, 8, *MIXED_BATCH, id="mrope-model-1d"),
60 + pytest.param(True, True, False, 4, *MIXED_BATCH, id="mrope-no-position-cache"),
61 + pytest.param(
62 + True,
63 + False,
64 + False,
65 + 8,
66 + *MIXED_BATCH,
67 + id="mrope-model-1d-no-position-cache",
68 + ),
69 + pytest.param(
70 + False, False, True, 4, *MIXED_BATCH, id="text-with-position-cache"
71 + ),
72 + pytest.param(True, True, True, 4, [37], [37], [0], id="fresh"),
73 + pytest.param(True, True, True, 4, [4097], [4097], [0], id="tiled"),
74 + ],
75 +)
76 +def test_qsa_fused_pre_indexer_matches_unfused(
77 + mrope,
78 + is_2d_positions,
79 + cache_rope_positions,
80 + state_size,
81 + seq_lens,
82 + query_lens,
83 + history_lens,
84 +) -> None:
85 + from flashinfer.norm import gemma_rmsnorm
86 +
87 + from vllm.model_executor.layers.rotary_embedding import get_rope
88 +
89 + device = "cuda"
90 + rope_params = {
91 + "partial_rotary_factor": 0.25,
92 + "rope_theta": 10000000,
93 + "rope_type": "default",
94 + }
95 + if mrope:
96 + rope_params["mrope_interleaved"] = True
97 + rope_params["mrope_section"] = list(MROPE_SECTION)
98 + with torch.device(device):
99 + rope = get_rope(
100 + head_size=256,
101 + max_position=32768,
102 + rope_parameters=rope_params,
103 + dtype=torch.bfloat16,
104 + )
105 +
106 + token_to_req = torch.cat(
107 + [
108 + torch.full((length,), request, dtype=torch.int32)
109 + for request, length in enumerate(query_lens)
110 + ]
111 + ).to(device)
112 + logical_positions = torch.cat(
113 + [
114 + torch.arange(seq_len - query_len, seq_len, dtype=torch.int64)
115 + for seq_len, query_len in zip(seq_lens, query_lens)
116 + ]
117 + ).to(device)
118 + num_tokens = logical_positions.numel()
119 + query_start_loc = torch.tensor(
120 + [0, *torch.tensor(query_lens).cumsum(0).tolist()], dtype=torch.int32
121 + ).to(device)
122 + positions = (
123 + torch.stack(
124 + [
125 + logical_positions,
126 + logical_positions // 7 + 3,
127 + logical_positions // 13 + 11,
128 + ]
129 + )
130 + if is_2d_positions
131 + else logical_positions
132 + )
133 + position_rows = (
134 + canonical_qsa_rope_positions(positions)
135 + if cache_rope_positions
136 + else logical_positions.view(-1, 1, 1).expand(-1, 1, 3)
137 + )
138 +
139 + compressed_block_counts = [
140 + (seq_len // CR + COMP_PAGE - 1) // COMP_PAGE for seq_len in seq_lens
141 + ]
142 + raw_block_table, num_raw_blocks = _make_block_table([1] * len(seq_lens))
143 + compressed_block_table, num_compressed_blocks = _make_block_table(
144 + compressed_block_counts
145 + )
146 + raw_block_table = raw_block_table.to(device)
147 + raw_slots = circular_qsa_slot_mapping(
148 + raw_block_table,
149 + token_to_req,
150 + logical_positions,
151 + state_size,
152 + query_start_loc,
153 + )
154 + compressed_slots = compressed_qsa_slot_mapping(
155 + compressed_block_table.to(device),
156 + token_to_req,
157 + logical_positions,
158 + COMP_PAGE,
159 + CR,
160 + )
161 + group_counts = torch.tensor(
162 + [
163 + seq_len // CR - (seq_len - query_len) // CR
164 + for seq_len, query_len in zip(seq_lens, query_lens)
165 + ],
166 + dtype=torch.int32,
167 + device=device,
168 + )
169 + k_work_counts = torch.maximum(group_counts, torch.ones_like(group_counts))
170 + k_start_loc = torch.cat([k_work_counts.new_zeros(1), k_work_counts.cumsum(0)])
171 + work_requests = torch.repeat_interleave(
172 + torch.arange(len(query_lens), dtype=torch.int32, device=device),
173 + k_work_counts,
174 + )
175 + local_work = torch.arange(
176 + int(k_start_loc[-1]), dtype=torch.int32, device=device
177 + ) - torch.repeat_interleave(k_start_loc[:-1], k_work_counts)
178 + max_k_work = (num_tokens + (CR - 1) * len(query_lens)) // CR
179 + k_work_metadata = torch.full((max_k_work, 2), -1, dtype=torch.int32, device=device)
180 + k_work_metadata[: work_requests.numel()] = torch.stack(
181 + (work_requests, local_work), dim=1
182 + )
183 +
184 + raw_width = D + 12 if cache_rope_positions else D
185 + # Match vLLM's padded-page cache layout: rows are contiguous, while physical
186 + # blocks have a larger stride than their logical contents.
187 + raw_page_elements = state_size * raw_width
188 + fused_raw_storage = torch.zeros(
189 + num_raw_blocks,
190 + raw_page_elements + 16,
191 + dtype=torch.bfloat16,
192 + device=device,
193 + )
194 + fused_raw = torch.as_strided(
195 + fused_raw_storage,
196 + (num_raw_blocks, state_size, 1, raw_width),
197 + (raw_page_elements + 16, raw_width, raw_width, 1),
198 + )
199 + for request, history_len in enumerate(history_lens):
200 + history_end = seq_lens[request] - query_lens[request]
201 + for position in range(history_end - history_len, history_end):
202 + block = int(raw_block_table[request, 0])
203 + row = fused_raw[block, position % state_size, 0]
204 + row[:D] = torch.randn(D, dtype=torch.bfloat16, device=device)
205 + if cache_rope_positions:
206 + row[ROPE_POS_OFFSET:].view(torch.int64).copy_(
207 + torch.tensor(
208 + [position, position // 7 + 3, position // 13 + 11],
209 + dtype=torch.int64,
210 + device=device,
211 + )
212 + )
213 + unfused_raw = fused_raw.clone()
214 + compressed_page_elements = COMP_PAGE * D
215 + fused_compressed_storage = torch.zeros(
216 + num_compressed_blocks,
217 + compressed_page_elements + 16,
218 + dtype=torch.bfloat16,
219 + device=device,
220 + )
221 + fused_compressed = torch.as_strided(
222 + fused_compressed_storage,
223 + (num_compressed_blocks, COMP_PAGE, 1, D),
224 + (compressed_page_elements + 16, D, D, 1),
225 + )
226 + unfused_compressed = fused_compressed.clone()
227 +
228 + projected_qk = torch.randn(
229 + num_tokens, (HQ + 1) * D, dtype=torch.bfloat16, device=device
230 + )
231 + q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
232 + k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
233 +
234 + fused_query = torch.empty(num_tokens, HQ, D, dtype=torch.bfloat16, device=device)
235 + qsa_pre_indexer(
236 + projected_qk[:, : HQ * D],
237 + projected_qk[:, HQ * D :],
238 + positions,
239 + rope.cos_sin_cache,
240 + q_weight,
241 + k_weight,
242 + EPS,
243 + fused_query,
244 + fused_raw,
245 + raw_slots,
246 + raw_block_table,
247 + query_start_loc,
248 + logical_positions,
249 + fused_compressed,
250 + compressed_slots,
251 + k_work_metadata,
252 + compress_ratio=CR,
253 + mrope_section=MROPE_SECTION if mrope else None,
254 + rope_pos_offset=ROPE_POS_OFFSET if cache_rope_positions else None,
255 + )
256 +
257 + unfused_query = projected_qk[:, : HQ * D].reshape(num_tokens, HQ, D)
258 + unfused_query = gemma_rmsnorm(
259 + unfused_query.reshape(-1, D), q_weight, EPS
260 + ).reshape_as(unfused_query)
261 + unfused_query = apply_qsa_rope(rope, positions, unfused_query)
262 +
263 + raw_keys = unfused_raw[..., :D]
264 + rope_positions = (
265 + unfused_raw[..., ROPE_POS_OFFSET:].view(torch.int64)
266 + if cache_rope_positions
267 + else None
268 + )
269 + pooled, first_positions = qsa_compress_groups_with_ratio(
270 + projected_qk[:, HQ * D :].reshape(-1, 1, D),
271 + position_rows,
272 + raw_keys,
273 + raw_block_table,
274 + token_to_req,
275 + query_start_loc,
276 + logical_positions,
277 + compressed_slots,
278 + CR,
279 + rope_positions,
280 + )
281 + compressed_rows = gemma_rmsnorm(pooled.reshape(-1, D), k_weight, EPS).reshape(
282 + -1, 1, D
283 + )
284 + group_positions = (
285 + first_positions.transpose(0, 1) if mrope else first_positions[:, 0]
286 + )
287 + compressed_rows = apply_qsa_rope(rope, group_positions, compressed_rows)
288 + qsa_store_cache_rows(unfused_compressed, compressed_slots, compressed_rows)
289 + qsa_store_cache_rows(raw_keys, raw_slots, projected_qk[:, HQ * D :])
290 + if rope_positions is not None:
291 + qsa_store_cache_rows(rope_positions, raw_slots, position_rows)
292 +
293 + torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)
294 + assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16))
295 + torch.testing.assert_close(
296 + fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL
297 + )
tests/models/qwen4_exp/test_qsa_reference.py new
+973
@@ -0,0 +1,973 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import math
5 +from types import SimpleNamespace
6 +
7 +import pytest
8 +import torch
9 +
10 +from vllm.models.qwen4_exp.common import qsa_cache
11 +from vllm.models.qwen4_exp.common.qsa_cache import QSAMetadataBuilder
12 +from vllm.models.qwen4_exp.nvidia import indexer_qsa
13 +from vllm.models.qwen4_exp.nvidia import (
14 + model as _qwen4_exp_model, # noqa: F401
15 +)
16 +from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops
17 +from vllm.platforms import current_platform
18 +from vllm.triton_utils import HAS_TRITON
19 +
20 +requires_qsa_kernels = pytest.mark.skipif(
21 + not current_platform.is_cuda() or not HAS_TRITON,
22 + reason="QSA kernels require CUDA and Triton",
23 +)
24 +
25 +
26 +def test_qsa_mtp_index_share_updates_cache_but_skips_selection(
27 + monkeypatch: pytest.MonkeyPatch,
28 +) -> None:
29 + rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32)
30 + raw_metadata = SimpleNamespace(
31 + num_actual_tokens=2,
32 + slot_mapping=torch.arange(2),
33 + block_table=torch.empty(0),
34 + query_start_loc=torch.arange(3),
35 + logical_positions=torch.arange(2),
36 + )
37 + compressed_metadata = SimpleNamespace(
38 + num_actual_tokens=2,
39 + slot_mapping=torch.arange(2),
40 + k_work_metadata=torch.empty(0),
41 + )
42 + updates = []
43 + selections = []
44 + indexer = SimpleNamespace(
45 + skip_topk=True,
46 + _metadata=lambda: (raw_metadata, compressed_metadata),
47 + index_qk_proj=lambda hidden: (torch.zeros(2, 2), None),
48 + index_n_heads=1,
49 + index_kv_heads=1,
50 + index_head_dim=1,
51 + raw_key_cache=SimpleNamespace(
52 + kv_cache=torch.empty(0),
53 + rope_position_cache=None,
54 + rope_position_offset=0,
55 + ),
56 + compressed_key_cache=SimpleNamespace(kv_cache=torch.empty(0)),
57 + use_fused_pre_indexer=True,
58 + rotary_emb=SimpleNamespace(cos_sin_cache=torch.empty(0)),
59 + q_layernorm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6),
60 + k_layernorm=SimpleNamespace(weight=torch.ones(1)),
61 + compress_ratio=2,
62 + )
63 +
64 + monkeypatch.setattr(
65 + indexer_qsa,
66 + "qsa_pre_indexer",
67 + lambda *args, **kwargs: updates.append((args, kwargs)),
68 + )
69 + monkeypatch.setattr(
70 + qsa_ops,
71 + "qsa_select_paged_tokens",
72 + lambda *args, **kwargs: selections.append((args, kwargs)),
73 + )
74 +
75 + actual = indexer_qsa.QSAIndexer.forward(
76 + indexer,
77 + torch.zeros(2, 4),
78 + torch.tensor([7, 8]),
79 + rows,
80 + )
81 +
82 + assert actual is rows
83 + assert len(updates) == 1
84 + assert not selections
85 +
86 +
87 +def _qsa_mqa_paged_reference(
88 + q: torch.Tensor,
89 + k_cache: torch.Tensor,
90 + page_table: torch.Tensor,
91 + token_to_req: torch.Tensor,
92 + visible_lengths: torch.Tensor,
93 +) -> torch.Tensor:
94 + pages = page_table.index_select(0, token_to_req.long()).long()
95 + keys = k_cache[pages, :, 0, :].flatten(1, 2)
96 + scores = torch.einsum("rhd,rnd->rnh", q.float(), keys.float())
97 + logits = torch.relu(scores).sum(dim=-1) / math.sqrt(q.shape[-1])
98 + positions = torch.arange(keys.shape[1], device=q.device).unsqueeze(0)
99 + return logits.masked_fill(positions >= visible_lengths.unsqueeze(1), -torch.inf)
100 +
101 +
102 +def _qsa_relative_topk_reference(
103 + logits: torch.Tensor,
104 + row_starts: torch.Tensor,
105 + row_ends: torch.Tensor,
106 + topk: int,
107 +) -> torch.Tensor:
108 + output = torch.full(
109 + (logits.shape[0], topk), -1, dtype=torch.int32, device=logits.device
110 + )
111 + for row in range(logits.shape[0]):
112 + start = int(row_starts[row].item())
113 + length = int((row_ends[row] - row_starts[row]).item())
114 + width = min(length, topk)
115 + if width:
116 + output[row, :width] = torch.topk(
117 + logits[row, start : start + length], width
118 + ).indices.to(torch.int32)
119 + return output
120 +
121 +
122 +def _expand_qsa_indices_reference(
123 + block_indices: torch.Tensor,
124 + query_positions: torch.Tensor,
125 + sequence_lengths: torch.Tensor,
126 + compress_ratio: int,
127 + token_topk: int,
128 +) -> torch.Tensor:
129 + rows = block_indices.shape[0]
130 + block_topk = token_topk // compress_ratio
131 + output_width = token_topk + compress_ratio - 1
132 + offsets = torch.arange(compress_ratio, device=block_indices.device)
133 + blocks = block_indices.long()
134 + expanded = blocks.unsqueeze(-1) * compress_ratio + offsets
135 + expanded = torch.where(
136 + blocks.unsqueeze(-1) >= 0, expanded, torch.full_like(expanded, -1)
137 + ).reshape(rows, block_topk * compress_ratio)
138 + expanded = expanded[:, :token_topk]
139 + expanded = torch.where(
140 + (expanded >= 0) & (expanded < sequence_lengths.unsqueeze(1)),
141 + expanded,
142 + torch.full_like(expanded, -1),
143 + )
144 +
145 + tail_offsets = torch.arange(compress_ratio - 1, device=block_indices.device)
146 + visible_tokens = query_positions + 1
147 + tail_start = visible_tokens // compress_ratio * compress_ratio
148 + tail = tail_start.unsqueeze(1) + tail_offsets.unsqueeze(0)
149 + tail_count = (visible_tokens - tail_start).unsqueeze(1)
150 + tail_valid = (tail_offsets.unsqueeze(0) < tail_count) & (
151 + tail < sequence_lengths.unsqueeze(1)
152 + )
153 + tail = torch.where(tail_valid, tail, torch.full_like(tail, -1))
154 +
155 + result = torch.cat((expanded, tail), dim=1)
156 + order = torch.arange(output_width, device=result.device).expand(rows, -1)
157 + sort_key = torch.where(result >= 0, order, order + output_width)
158 + return result.gather(1, torch.argsort(sort_key, dim=1, stable=True)).to(torch.int32)
159 +
160 +
161 +def _qsa_select_paged_tokens_reference(
162 + q: torch.Tensor,
163 + k_cache: torch.Tensor,
164 + page_table: torch.Tensor,
165 + token_to_req: torch.Tensor,
166 + query_positions: torch.Tensor,
167 + sequence_lengths: torch.Tensor,
168 + token_topk: int,
169 + compress_ratio: int,
170 +) -> torch.Tensor:
171 + row_sequence_lengths = sequence_lengths.index_select(0, token_to_req.long())
172 + visible_blocks = torch.minimum(
173 + (query_positions + 1) // compress_ratio,
174 + row_sequence_lengths // compress_ratio,
175 + ).to(torch.int32)
176 + logits = _qsa_mqa_paged_reference(
177 + q,
178 + k_cache,
179 + page_table,
180 + token_to_req,
181 + visible_blocks,
182 + )
183 + starts = torch.zeros_like(visible_blocks)
184 + blocks = _qsa_relative_topk_reference(
185 + logits,
186 + starts,
187 + visible_blocks,
188 + token_topk // compress_ratio,
189 + )
190 + return _expand_qsa_indices_reference(
191 + blocks,
192 + query_positions,
193 + row_sequence_lengths,
194 + compress_ratio,
195 + token_topk,
196 + )
197 +
198 +
199 +def _qsa_sparse_paged_attention_reference(
200 + q: torch.Tensor,
201 + k_cache: torch.Tensor,
202 + v_cache: torch.Tensor,
203 + logical_indices: torch.Tensor,
204 + block_table: torch.Tensor,
205 + token_to_req: torch.Tensor,
206 + softmax_scale: float,
207 +) -> torch.Tensor:
208 + output = torch.zeros_like(q)
209 + repeats = q.shape[1] // k_cache.shape[2]
210 + page_size = k_cache.shape[1]
211 + for row in range(q.shape[0]):
212 + logical = logical_indices[row]
213 + logical = logical[logical >= 0].long()
214 + if not logical.numel():
215 + continue
216 + request = token_to_req[row].long()
217 + pages = block_table[request, logical // page_size].long()
218 + offsets = logical % page_size
219 + keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1)
220 + values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1)
221 + scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float())
222 + probabilities = torch.softmax(scores * softmax_scale, dim=-1)
223 + output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to(
224 + q.dtype
225 + )
226 + return output
227 +
228 +
229 +@requires_qsa_kernels
230 +def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None:
231 + device = torch.device("cuda")
232 + builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
233 + builder.compress_ratio = 1
234 + builder.is_circular_buffer = False
235 + builder.storage_block_size = 64
236 + builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device)
237 + builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device)
238 + builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device)
239 + builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device)
240 + query_start_loc = torch.tensor([0, 4, 8, 12, 12], dtype=torch.int32, device=device)
241 + token_to_req = torch.tensor([0] * 4 + [1] * 4 + [2] * 4 + [0] * 4, device=device)
242 + common = SimpleNamespace(
243 + num_actual_tokens=16,
244 + query_start_loc=query_start_loc,
245 + query_start_loc_cpu=query_start_loc.cpu(),
246 + seq_lens=torch.tensor([68, 68, 68, 0], dtype=torch.int32, device=device),
247 + slot_mapping=torch.tensor(list(range(12)) + [-1] * 4, device=device),
248 + block_table_tensor=torch.empty((4, 0), dtype=torch.int32, device=device),
249 + token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),
250 + )
251 +
252 + metadata = builder.build(0, common)
253 +
254 + assert metadata.logical_positions.tolist() == [
255 + 64,
256 + 65,
257 + 66,
258 + 67,
259 + 64,
260 + 65,
261 + 66,
262 + 67,
263 + 64,
264 + 65,
265 + 66,
266 + 67,
267 + -1,
268 + -1,
269 + -1,
270 + -1,
271 + ]
272 + assert metadata.slot_mapping.tolist() == list(range(12)) + [-1] * 4
273 +
274 +
275 +@requires_qsa_kernels
276 +def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None:
277 + device = torch.device("cuda")
278 + builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
279 + builder.compress_ratio = 4
280 + builder.is_circular_buffer = True
281 + builder.kv_cache_spec = SimpleNamespace(block_size=4)
282 + builder.storage_block_size = 4
283 + builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device)
284 + builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device)
285 + builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device)
286 + builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device)
287 + query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device)
288 + token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device)
289 + block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device)
290 + common = SimpleNamespace(
291 + num_actual_tokens=16,
292 + query_start_loc=query_start_loc,
293 + query_start_loc_cpu=query_start_loc.cpu(),
294 + seq_lens=torch.tensor([9, 11, 0], dtype=torch.int32, device=device),
295 + slot_mapping=torch.full((16,), -1, dtype=torch.int64, device=device),
296 + block_table_tensor=block_table,
297 + token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),
298 + )
299 +
300 + metadata = builder.build(0, common)
301 + expected = [
302 + -1,
303 + -1,
304 + -1,
305 + 5,
306 + 6,
307 + 7,
308 + 4,
309 + -1,
310 + -1,
311 + 3,
312 + 0,
313 + 1,
314 + 2,
315 + -1,
316 + -1,
317 + -1,
318 + ]
319 +
320 + assert metadata.slot_mapping.tolist() == expected
321 +
322 +
323 +@pytest.mark.parametrize("chunk_start", list(range(8)))
324 +def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> None:
325 + """A speculative step must not overwrite the open group's committed keys.
326 +
327 + The step stores every row it computes, drafts included, before acceptance
328 + is known, while the next step still reads the earlier members of the group
329 + being compressed from the ring. A ring sized at the compression ratio makes
330 + those rows alias, so a rejected draft silently replaces a committed key.
331 + """
332 + compress_ratio = 4
333 + num_spec = 3
334 + capacity = compress_ratio * -(-(compress_ratio + num_spec) // compress_ratio)
335 + query_len = num_spec + 1
336 +
337 + slots = qsa_cache.circular_qsa_slot_mapping(
338 + torch.tensor([[0]], dtype=torch.int32),
339 + torch.zeros(query_len, dtype=torch.int32),
340 + torch.arange(chunk_start, chunk_start + query_len),
341 + capacity,
342 + query_start_loc=torch.tensor([0, query_len], dtype=torch.int32),
343 + )
344 +
345 + committed = torch.arange(chunk_start - chunk_start % compress_ratio, chunk_start)
346 + assert set(slots.tolist()).isdisjoint((committed % capacity).tolist())
347 +
348 +
349 +def _qsa_key_cache(block_size: int, compress_ratio: int) -> qsa_cache.QSAKeyStateCache:
350 + return qsa_cache.QSAKeyStateCache(
351 + head_size=64,
352 + dtype=torch.bfloat16,
353 + cache_config=SimpleNamespace(block_size=block_size),
354 + prefix=f"raw.{block_size}.{compress_ratio}",
355 + vllm_config=SimpleNamespace(
356 + compilation_config=SimpleNamespace(static_forward_context={})
357 + ),
358 + compress_ratio=compress_ratio,
359 + )
360 +
361 +
362 +def test_qsa_state_caches_adapt_the_unified_logical_layout() -> None:
363 + raw_cache = _qsa_key_cache(block_size=32, compress_ratio=4)
364 + compressed_cache = qsa_cache.QSACompressedKeyCache(
365 + head_size=64,
366 + dtype=torch.bfloat16,
367 + cache_config=SimpleNamespace(block_size=32),
368 + prefix="compressed.bind",
369 + vllm_config=SimpleNamespace(
370 + compilation_config=SimpleNamespace(static_forward_context={})
371 + ),
372 + compress_ratio=4,
373 + )
374 + raw_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16)
375 + compressed_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16)
376 +
377 + raw_cache.bind_kv_cache(raw_view)
378 + compressed_cache.bind_kv_cache(compressed_view)
379 +
380 + assert raw_cache.kv_cache.shape == (2, 8, 1, 64)
381 + assert compressed_cache.kv_cache.shape == (2, 8, 1, 64)
382 + assert raw_cache.kv_cache.data_ptr() == raw_view.data_ptr()
383 + assert compressed_cache.kv_cache.data_ptr() == compressed_view.data_ptr()
384 +
385 +
386 +@pytest.mark.parametrize(
387 + ("compress_ratio", "num_spec", "expected"),
388 + [(4, 0, 4), (4, 1, 8), (4, 3, 8), (4, 4, 8), (4, 5, 12), (2, 3, 6)],
389 +)
390 +def test_qsa_ring_capacity_covers_one_speculative_step(
391 + compress_ratio: int, num_spec: int, expected: int
392 +) -> None:
393 + """Capacity spans the open group plus one speculative step, in whole groups."""
394 + spec = _qsa_key_cache(
395 + block_size=48, compress_ratio=compress_ratio
396 + ).get_kv_cache_spec(SimpleNamespace(num_speculative_tokens=num_spec))
397 + assert spec.block_size == expected
398 +
399 +
400 +@requires_qsa_kernels
401 +def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:
402 + device = torch.device("cuda")
403 + builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)
404 + builder.compress_ratio = 4
405 + builder.is_circular_buffer = False
406 + builder.storage_block_size = 16
407 + builder.token_to_req_buffer = torch.empty(8, dtype=torch.int32, device=device)
408 + builder.slot_mapping_buffer = torch.empty(8, dtype=torch.int64, device=device)
409 + builder.logical_positions_buffer = torch.empty(8, dtype=torch.int64, device=device)
410 + # Simulate max_num_seqs exceeding the three live requests below.
411 + builder.request_capacity = 8
412 + builder.k_work_metadata_buffer = torch.empty(4, 2, dtype=torch.int32, device=device)
413 + query_start_loc = torch.tensor([0, 3, 3, 8], dtype=torch.int32, device=device)
414 + token_to_req = torch.tensor(
415 + [0, 0, 0, 2, 2, 2, 2, 2], dtype=torch.int32, device=device
416 + )
417 + common = SimpleNamespace(
418 + num_actual_tokens=8,
419 + query_start_loc=query_start_loc,
420 + query_start_loc_cpu=query_start_loc.cpu(),
421 + seq_lens=torch.tensor([7, 0, 12], dtype=torch.int32, device=device),
422 + slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device),
423 + block_table_tensor=torch.zeros((3, 1), dtype=torch.int32, device=device),
424 + token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),
425 + )
426 +
427 + metadata = builder.build(0, common)
428 +
429 + assert metadata.slot_mapping.tolist() == [-1] * 8
430 + assert metadata.k_work_metadata.tolist() == [[0, 0], [2, 0], [2, 1], [-1, -1]]
431 +
432 +
433 +@requires_qsa_kernels
434 +@pytest.mark.parametrize("compress_ratio", [1, 4])
435 +@pytest.mark.parametrize("num_reqs", [2, 3, 4, 7, 8, 9])
436 +def test_qsa_triton_metadata_matches_pytorch(
437 + compress_ratio: int, num_reqs: int
438 +) -> None:
439 + device = torch.device("cuda")
440 + num_tokens = 8
441 + query_start_loc = torch.tensor(
442 + [0, 3, *([3] * (num_reqs - 2)), 8], dtype=torch.int32, device=device
443 + )
444 + token_to_req = torch.tensor(
445 + [0, 0, 0, *([num_reqs - 1] * 5)],
446 + dtype=torch.int32,
447 + device=device,
448 + )
449 + block_table_rows = torch.tensor(
450 + [
451 + [4, -1, 8, -1, 12, -1],
452 + [1, -1, 2, -1, 3, -1],
453 + [7, -1, 9, -1, 11, -1],
454 + ],
455 + dtype=torch.int32,
456 + device=device,
457 + )
458 + block_table_storage = block_table_rows[
459 + torch.arange(num_reqs, device=device) % block_table_rows.shape[0]
460 + ]
461 + seq_lens = torch.zeros(num_reqs, dtype=torch.int32, device=device)
462 + seq_lens[0] = 10
463 + seq_lens[-1] = 20
464 + common = SimpleNamespace(
465 + num_actual_tokens=num_tokens,
466 + query_start_loc=query_start_loc,
467 + query_start_loc_cpu=query_start_loc.cpu(),
468 + seq_lens=seq_lens,
469 + slot_mapping=torch.tensor(
470 + [0, 1, -1, 3, 4, -1, -1, -1], dtype=torch.int64, device=device
471 + ),
472 + block_table_tensor=block_table_storage[:, ::2],
473 + token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),
474 + )
475 +
476 + def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
477 + return (
478 + torch.empty(num_tokens, dtype=torch.int32, device=device),
479 + torch.empty(num_tokens, dtype=torch.int64, device=device),
480 + torch.empty(num_tokens, dtype=torch.int64, device=device),
481 + )
482 +
483 + max_num_work = (
484 + (num_tokens + (compress_ratio - 1) * num_reqs) // compress_ratio
485 + if compress_ratio != 1
486 + else 0
487 + )
488 + actual_k_work = (
489 + torch.empty(max_num_work, 2, dtype=torch.int32, device=device)
490 + if max_num_work
491 + else None
492 + )
493 + actual_buffers = make_buffers()
494 + actual = qsa_cache.build_qsa_metadata_triton(
495 + common,
496 + *actual_buffers,
497 + storage_block_size=2,
498 + compress_ratio=compress_ratio,
499 + k_work_metadata_buffer=actual_k_work,
500 + request_capacity=num_reqs,
501 + )
502 +
503 + expected_k_work = (
504 + torch.empty_like(actual_k_work) if actual_k_work is not None else None
505 + )
506 + expected_buffers = make_buffers()
507 + expected = qsa_cache._build_qsa_metadata_torch(
508 + common,
509 + *expected_buffers,
510 + storage_block_size=2,
511 + compress_ratio=compress_ratio,
512 + k_work_metadata_buffer=expected_k_work,
513 + request_capacity=num_reqs,
514 + )
515 +
516 + for actual_tensor, expected_tensor in zip(actual, expected):
517 + torch.testing.assert_close(actual_tensor, expected_tensor)
518 + if actual_k_work is not None:
519 + torch.testing.assert_close(actual_k_work, expected_k_work)
520 +
521 +
522 +@requires_qsa_kernels
523 +def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None:
524 + device = torch.device("cuda")
525 + num_mapped_tokens = 4096
526 + num_tokens = 4224
527 + query_start_loc = torch.tensor(
528 + [0, num_mapped_tokens], dtype=torch.int32, device=device
529 + )
530 + common = SimpleNamespace(
531 + num_actual_tokens=num_tokens,
532 + query_start_loc=query_start_loc,
533 + query_start_loc_cpu=query_start_loc.cpu(),
534 + seq_lens=torch.tensor(
535 + [num_mapped_tokens + 32], dtype=torch.int32, device=device
536 + ),
537 + block_table_tensor=torch.arange(256, dtype=torch.int32, device=device)[None],
538 + slot_mapping=torch.tensor(
539 + [0] * num_mapped_tokens + [-1] * (num_tokens - num_mapped_tokens),
540 + dtype=torch.int64,
541 + device=device,
542 + ),
543 + token_to_req_indices=lambda buffer: buffer.zero_(),
544 + )
545 +
546 + def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
547 + return (
548 + torch.empty(num_tokens, dtype=torch.int32, device=device),
549 + torch.empty(num_tokens, dtype=torch.int64, device=device),
550 + torch.empty(num_tokens, dtype=torch.int64, device=device),
551 + )
552 +
553 + max_num_work = (num_tokens + 3) // 4
554 + actual_k_work = torch.empty(max_num_work, 2, dtype=torch.int32, device=device)
555 + expected_k_work = torch.empty_like(actual_k_work)
556 + actual = qsa_cache.build_qsa_metadata_triton(
557 + common,
558 + *make_buffers(),
559 + storage_block_size=8,
560 + compress_ratio=4,
561 + k_work_metadata_buffer=actual_k_work,
562 + )
563 + expected = qsa_cache._build_qsa_metadata_torch(
564 + common,
565 + *make_buffers(),
566 + storage_block_size=8,
567 + compress_ratio=4,
568 + k_work_metadata_buffer=expected_k_work,
569 + )
570 +
571 + for actual_tensor, expected_tensor in zip(actual, expected):
572 + torch.testing.assert_close(actual_tensor, expected_tensor)
573 + torch.testing.assert_close(actual_k_work, expected_k_work)
574 +
575 +
576 +@requires_qsa_kernels
577 +@pytest.mark.parametrize(
578 + "num_rows",
579 + [
580 + pytest.param(3, id="one_tile_per_program"),
581 + pytest.param(33, id="looped_tiles"),
582 + ],
583 +)
584 +def test_qsa_mqa_paged_matches_test_reference(num_rows: int) -> None:
585 + torch.manual_seed(1)
586 + head_dim = 128
587 + q = torch.randn(num_rows, 4, head_dim, device="cuda", dtype=torch.bfloat16)
588 + cache = torch.randn(40, 4, 1, head_dim, device="cuda", dtype=torch.bfloat16)
589 + page_table = torch.randperm(40, device="cuda", dtype=torch.int32).reshape(2, 20)
590 + request_rows = [(num_rows + 1) // 2, num_rows // 2]
591 + token_to_req = torch.repeat_interleave(
592 + torch.arange(2, device="cuda", dtype=torch.int32),
593 + torch.tensor(request_rows, device="cuda"),
594 + )
595 + sequence_length_values = [320, 264]
596 + sequence_lengths = torch.tensor(
597 + sequence_length_values, device="cuda", dtype=torch.int32
598 + )
599 + query_positions = torch.cat(
600 + [
601 + torch.arange(length - rows, length, device="cuda", dtype=torch.int32)
602 + for rows, length in zip(request_rows, sequence_length_values, strict=True)
603 + ]
604 + )
605 + compress_ratio = 4
606 + visible_lengths = (query_positions + 1) // compress_ratio
607 +
608 + actual, actual_visible_blocks = qsa_ops.qsa_mqa_paged(
609 + q,
610 + cache,
611 + page_table,
612 + token_to_req,
613 + query_positions,
614 + sequence_lengths,
615 + compress_ratio=compress_ratio,
616 + )
617 + expected = _qsa_mqa_paged_reference(
618 + q, cache, page_table, token_to_req, visible_lengths
619 + )
620 +
621 + torch.testing.assert_close(actual_visible_blocks, visible_lengths)
622 + # Top-k is bounded by visible_blocks; later columns are intentionally unwritten.
623 + columns = torch.arange(actual.shape[1], device=actual.device)
624 + visible = columns[None, :] < actual_visible_blocks[:, None]
625 + torch.testing.assert_close(actual[visible], expected[visible], rtol=1e-3, atol=1e-3)
626 +
627 +
628 +@requires_qsa_kernels
629 +def test_qsa_block_expansion_matches_test_reference() -> None:
630 + blocks = torch.tensor([[0, -1], [1, 0]], device="cuda", dtype=torch.int32)
631 + query_positions = torch.tensor([5, 10], device="cuda")
632 + sequence_lengths = torch.tensor([6, 11], device="cuda")
633 + token_to_req = torch.tensor([0, 1], device="cuda", dtype=torch.int32)
634 +
635 + actual = qsa_ops.expand_qsa_block_indices_cuda(
636 + blocks,
637 + query_positions,
638 + sequence_lengths,
639 + token_to_req,
640 + compress_ratio=4,
641 + token_topk=8,
642 + )
643 + expected = _expand_qsa_indices_reference(
644 + blocks,
645 + query_positions,
646 + sequence_lengths,
647 + compress_ratio=4,
648 + token_topk=8,
649 + )
650 +
651 + torch.testing.assert_close(actual, expected)
652 +
653 +
654 +@requires_qsa_kernels
655 +@pytest.mark.parametrize(
656 + ("num_rows", "num_query_heads", "num_kv_heads", "page_size"),
657 + [
658 + # Kernel-visible pages with --block-size 256 and hybrid-cache alignment.
659 + pytest.param(1, 24, 2, 1792, id="tp1_split64"),
660 + pytest.param(16, 12, 1, 1792, id="tp2_split32"),
661 + pytest.param(32, 6, 1, 1024, id="tp4_split8"),
662 + pytest.param(257, 6, 1, 1024, id="tp4_split4"),
663 + pytest.param(513, 6, 1, 1024, id="tp4_split1"),
664 + ],
665 +)
666 +def test_qsa_sparse_paged_attention_matches_test_reference(
667 + num_rows: int,
668 + num_query_heads: int,
669 + num_kv_heads: int,
670 + page_size: int,
671 +) -> None:
672 + torch.manual_seed(2)
673 + head_dim = 256
674 + num_requests = 2
675 + num_selected_pages = 64
676 + # Keep the newest page outside the synthetic top-k as causal headroom.
677 + num_pages_per_request = num_selected_pages + 1
678 + num_cache_blocks = num_requests * num_pages_per_request
679 + indexer_budget = 2048
680 + indexer_compress_ratio = 4
681 + selection_width = indexer_budget + indexer_compress_ratio - 1
682 + q = torch.randn(
683 + num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16
684 + )
685 + kv_cache = torch.randn(
686 + num_cache_blocks,
687 + page_size,
688 + num_kv_heads,
689 + 2 * head_dim,
690 + device="cuda",
691 + dtype=torch.bfloat16,
692 + )
693 + k_cache, v_cache = kv_cache.split(head_dim, dim=-1)
694 + block_table = (
695 + torch.randperm(num_cache_blocks, device="cuda")
696 + .reshape(num_requests, num_pages_per_request)
697 + .to(torch.int32)
698 + )
699 + rows_per_request = math.ceil(num_rows / num_requests)
700 + row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32)
701 + token_to_req = row_indices // rows_per_request
702 + request_row_counts = torch.tensor(
703 + [rows_per_request, num_rows - rows_per_request],
704 + device="cuda",
705 + dtype=torch.int32,
706 + )
707 +
708 + context_length = num_pages_per_request * page_size - 1
709 + block_topk = indexer_budget // indexer_compress_ratio
710 + compressed_blocks_per_page = page_size // indexer_compress_ratio
711 + selection = torch.arange(block_topk, device="cuda")
712 + selected_pages = selection % num_selected_pages
713 + selected_offsets = selection // num_selected_pages
714 + row_shifts = 2 * row_indices.unsqueeze(1)
715 + # Eight blocks per page; adjacent rows overlap by six of those eight.
716 + selected_offsets = (selected_offsets + row_shifts) % compressed_blocks_per_page
717 + block_indices = (selected_pages * compressed_blocks_per_page + selected_offsets).to(
718 + torch.int32
719 + )
720 + rows_within_request = row_indices % rows_per_request
721 + query_positions = (
722 + context_length - request_row_counts[token_to_req.long()] + rows_within_request
723 + ).to(torch.int64)
724 + sequence_lengths = torch.full(
725 + (num_requests,), context_length, device="cuda", dtype=torch.int32
726 + )
727 + logical_indices = qsa_ops.expand_qsa_block_indices_cuda(
728 + block_indices,
729 + query_positions,
730 + sequence_lengths,
731 + token_to_req,
732 + indexer_compress_ratio,
733 + indexer_budget,
734 + )
735 + assert logical_indices.shape == (num_rows, selection_width)
736 + scale = q.shape[-1] ** -0.5
737 +
738 + actual = qsa_ops.qsa_sparse_paged_attention(
739 + q,
740 + k_cache,
741 + v_cache,
742 + logical_indices,
743 + block_table,
744 + token_to_req,
745 + )
746 + expected = _qsa_sparse_paged_attention_reference(
747 + q,
748 + k_cache,
749 + v_cache,
750 + logical_indices,
751 + block_table,
752 + token_to_req,
753 + scale,
754 + )
755 +
756 + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
757 +
758 +
759 +@requires_qsa_kernels
760 +def test_qsa_selection_chunks_workspace_and_matches_test_reference(
761 + monkeypatch: pytest.MonkeyPatch,
762 + workspace_init,
763 +) -> None:
764 + rows, keys, heads, head_dim = 65, 640, 4, 16
765 + token_topk, compress_ratio = 2048, 4
766 + torch.manual_seed(3)
767 + q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)
768 + cache = torch.randn(40, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16)
769 + page_table = torch.randperm(40, device="cuda", dtype=torch.int32).unsqueeze(0)
770 + token_to_req = torch.zeros(rows, device="cuda", dtype=torch.int32)
771 + query_positions = torch.full((rows,), 2559, device="cuda", dtype=torch.int32)
772 + sequence_lengths = torch.tensor([2560], device="cuda", dtype=torch.int32)
773 + monkeypatch.setattr(qsa_ops, "_LOGITS_WORKSPACE_BYTES", 32 * keys * 4)
774 + original_score = qsa_ops.qsa_mqa_paged
775 + scored_row_counts = []
776 +
777 + def record_score(query: torch.Tensor, *args, **kwargs):
778 + scored_row_counts.append(query.shape[0])
779 + return original_score(query, *args, **kwargs)
780 +
781 + monkeypatch.setattr(qsa_ops, "qsa_mqa_paged", record_score)
782 +
783 + actual = qsa_ops.qsa_select_paged_tokens(
784 + q,
785 + cache,
786 + page_table,
787 + token_to_req,
788 + query_positions,
789 + sequence_lengths,
790 + token_topk,
791 + compress_ratio,
792 + )
793 + expected = _qsa_select_paged_tokens_reference(
794 + q,
795 + cache,
796 + page_table,
797 + token_to_req,
798 + query_positions,
799 + sequence_lengths,
800 + token_topk,
801 + compress_ratio,
802 + )
803 +
804 + torch.testing.assert_close(actual.sort().values, expected.sort().values)
805 + assert scored_row_counts == [32, 32, 1]
806 +
807 +
808 +@requires_qsa_kernels
809 +def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) -> None:
810 + q = torch.zeros(2, 4, 8, device="cuda", dtype=torch.bfloat16)
811 + cache = torch.zeros(1, 16, 1, 8, device="cuda", dtype=torch.bfloat16)
812 + page_table = torch.zeros(1, 1, device="cuda", dtype=torch.int32)
813 + token_to_req = torch.zeros(2, device="cuda", dtype=torch.int32)
814 + query_positions = torch.tensor([1, 2], device="cuda", dtype=torch.int32)
815 + sequence_lengths = torch.tensor([3], device="cuda", dtype=torch.int32)
816 +
817 + selected = qsa_ops.qsa_select_paged_tokens(
818 + q,
819 + cache,
820 + page_table,
821 + token_to_req,
822 + query_positions,
823 + sequence_lengths,
824 + token_topk=2048,
825 + compress_ratio=4,
826 + )
827 +
828 + assert selected[0, :2].tolist() == [0, 1]
829 + assert selected[1, :3].tolist() == [0, 1, 2]
830 + assert torch.all(selected[0, 2:] == -1)
831 + assert torch.all(selected[1, 3:] == -1)
832 +
833 +
834 +@requires_qsa_kernels
835 +def test_qsa_streaming_compression_and_compressor_state_store_match_reference() -> None:
836 + head_dim = 8
837 + current_pairs = [
838 + *((0, position) for position in range(2, 9)),
839 + *((1, position) for position in range(5, 11)),
840 + ]
841 +
842 + def key_row(request: int, position: int) -> torch.Tensor:
843 + return (
844 + torch.arange(head_dim, dtype=torch.float32) + request * 1000 + position * 10
845 + )
846 +
847 + def position_row(request: int, position: int) -> torch.Tensor:
848 + return torch.tensor(
849 + [
850 + request * 1000 + position,
851 + request * 1000 + position + 100,
852 + request * 1000 + position + 200,
853 + ],
854 + dtype=torch.int64,
855 + )
856 +
857 + raw_keys = (
858 + torch.stack([key_row(request, position) for request, position in current_pairs])
859 + .unsqueeze(1)
860 + .to(device="cuda", dtype=torch.bfloat16)
861 + )
862 + raw_positions = (
863 + torch.stack(
864 + [position_row(request, position) for request, position in current_pairs]
865 + )
866 + .unsqueeze(1)
867 + .to(device="cuda")
868 + )
869 + token_to_req = torch.tensor(
870 + [request for request, _ in current_pairs],
871 + dtype=torch.int32,
872 + device="cuda",
873 + )
874 + logical_positions = torch.tensor(
875 + [position for _, position in current_pairs],
876 + dtype=torch.int64,
877 + device="cuda",
878 + )
879 + query_start_loc = torch.tensor([0, 7, 13], dtype=torch.int32, device="cuda")
880 + compressor_state_block_table = torch.tensor(
881 + [[1], [0]], dtype=torch.int32, device="cuda"
882 + )
883 + compressor_state_cache = torch.zeros(
884 + 2, 4, 1, head_dim, dtype=torch.bfloat16, device="cuda"
885 + )
886 + rope_cache = torch.zeros(2, 4, 1, 3, dtype=torch.int64, device="cuda")
887 + for request, position, block in ((0, 0, 1), (0, 1, 1), (1, 4, 0)):
888 + compressor_state_cache[block, position % 4, 0] = key_row(request, position).to(
889 + device="cuda", dtype=torch.bfloat16
890 + )
891 + rope_cache[block, position % 4, 0] = position_row(request, position).to("cuda")
892 +
893 + compressed_slots = torch.full(
894 + (len(current_pairs),), -1, dtype=torch.int64, device="cuda"
895 + )
896 + valid_rows = torch.tensor([1, 5, 9], dtype=torch.int64, device="cuda")
897 + compressed_slots[valid_rows] = torch.arange(3, device="cuda")
898 + pooled, first_positions = qsa_ops.qsa_compress_groups_with_ratio(
899 + raw_keys,
900 + raw_positions,
901 + compressor_state_cache,
902 + compressor_state_block_table,
903 + token_to_req,
904 + query_start_loc,
905 + logical_positions,
906 + compressed_slots,
907 + compress_ratio=4,
908 + rope_cache=rope_cache,
909 + )
910 + pooled_without_rope, scalar_first_positions = (
911 + qsa_ops.qsa_compress_groups_with_ratio(
912 + raw_keys,
913 + raw_positions,
914 + compressor_state_cache,
915 + compressor_state_block_table,
916 + token_to_req,
917 + query_start_loc,
918 + logical_positions,
919 + compressed_slots,
920 + compress_ratio=4,
921 + )
922 + )
923 +
924 + groups = [
925 + [(0, position) for position in range(0, 4)],
926 + [(0, position) for position in range(4, 8)],
927 + [(1, position) for position in range(4, 8)],
928 + ]
929 + expected_pooled = (
930 + torch.stack(
931 + [
932 + torch.stack([key_row(*pair) for pair in group]).mean(dim=0)
933 + for group in groups
934 + ]
935 + )
936 + .unsqueeze(1)
937 + .to(device="cuda", dtype=torch.bfloat16)
938 + )
939 + expected_positions = torch.stack(
940 + [position_row(0, 0), position_row(0, 4), position_row(1, 4)]
941 + ).to("cuda")
942 + expected_scalar_positions = torch.tensor(
943 + [[0, 0, 0], [4, 4, 4], [4, 4, 4]],
944 + dtype=torch.int64,
945 + device="cuda",
946 + )
947 +
948 + torch.testing.assert_close(pooled[valid_rows], expected_pooled)
949 + torch.testing.assert_close(pooled_without_rope[valid_rows], expected_pooled)
950 + torch.testing.assert_close(first_positions[valid_rows], expected_positions)
951 + torch.testing.assert_close(
952 + scalar_first_positions[valid_rows], expected_scalar_positions
953 + )
954 +
955 + compressor_state_slots = torch.tensor(
956 + [-1, -1, -1, 5, 6, 7, 4, -1, -1, 3, 0, 1, 2],
957 + dtype=torch.int64,
958 + device="cuda",
959 + )
960 + qsa_ops.qsa_store_cache_rows(
961 + compressor_state_cache, compressor_state_slots, raw_keys
962 + )
963 + qsa_ops.qsa_store_cache_rows(rope_cache, compressor_state_slots, raw_positions)
964 + for request, positions, block in ((0, range(5, 9), 1), (1, range(7, 11), 0)):
965 + for position in positions:
966 + torch.testing.assert_close(
967 + compressor_state_cache[block, position % 4, 0],
968 + key_row(request, position).to(device="cuda", dtype=torch.bfloat16),
969 + )
970 + torch.testing.assert_close(
971 + rope_cache[block, position % 4, 0],
972 + position_row(request, position).to("cuda"),
973 + )
tests/models/registry.py
+6
@@ -508,6 +508,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
508 extras={"native-prefix": "imdatta0/small_qwen3_5_20b"},
509 max_model_len=4096,
510 ),
511 + "Qwen4ExpForCausalLM": _HfExamplesInfo("", is_available_online=False),
512 "MellumForCausalLM": _HfExamplesInfo("JetBrains/Mellum2-12B-A2.5B-Base"),
513 "Qwen3NextForCausalLM": _HfExamplesInfo(
514 "Qwen/Qwen3-Next-80B-A3B-Instruct",
@@ -1348,6 +1349,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1349 "Qwen/Qwen3.5-35B-A3B",
1350 max_model_len=4096,
1351 ),
1352 + "Qwen4ExpForConditionalGeneration": _HfExamplesInfo(
1353 + "",
1354 + is_available_online=False,
1355 + ),
1356 "Qwen3OmniMoeForConditionalGeneration": _HfExamplesInfo(
1357 "Qwen/Qwen3-Omni-30B-A3B-Instruct",
1358 max_model_len=4096,
@@ -1797,6 +1802,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1802 "Qwen/Qwen3.5-35B-A3B",
1803 speculative_model="Qwen/Qwen3.5-35B-A3B",
1804 ),
1805 + "Qwen4ExpMTP": _HfExamplesInfo("", is_available_online=False),
1806 "Step3p5MTP": _HfExamplesInfo(
1807 "stepfun-ai/Step-3.5-Flash",
1808 speculative_model="stepfun-ai/Step-3.5-Flash",
tests/quantization/test_fp8.py
+44
@@ -6,6 +6,7 @@ Run `pytest tests/quantization/test_fp8.py --forked`.
6 """
7
8 import logging
9 +from types import SimpleNamespace
10
11 import pytest
12 import regex as re
@@ -31,6 +32,10 @@ from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
32 from vllm.model_executor.layers.quantization.online.fp8 import (
33 Fp8PerTensorOnlineLinearMethod,
34 )
35 +from vllm.model_executor.layers.quantization.utils import flashinfer_utils
36 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
37 + prepare_fp8_moe_layer_for_fi,
38 +)
39 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
40 process_fp8_input_tensor_strategy_moe,
41 )
@@ -50,6 +55,45 @@ MODELS = [
55 ]
56
57
58 +def test_prepare_gated_trtllm_fp8_moe_weights_pads_each_projection(monkeypatch):
59 + monkeypatch.setattr(
60 + flashinfer_utils,
61 + "rotate_weights_for_fi_trtllm_fp8_per_tensor_moe",
62 + lambda *args: None,
63 + )
64 + intermediate = 17
65 + padded_intermediate = 32
66 + hidden_size = 4
67 + gate = torch.ones((1, intermediate, hidden_size), dtype=torch.float8_e4m3fn)
68 + up = torch.full_like(gate, 2)
69 + w13 = torch.cat((gate, up), dim=1)
70 + w2 = torch.ones((1, hidden_size, intermediate), dtype=torch.float8_e4m3fn)
71 + layer = SimpleNamespace(
72 + activation=SimpleNamespace(is_gated=True),
73 + moe_config=SimpleNamespace(
74 + is_act_and_mul=True,
75 + intermediate_size_per_partition=intermediate,
76 + ),
77 + )
78 +
79 + padded_w31, _, _, _ = prepare_fp8_moe_layer_for_fi(
80 + layer,
81 + w13,
82 + w2,
83 + w13_scale=torch.ones(1),
84 + w13_input_scale=torch.ones(1),
85 + w2_scale=torch.ones(1),
86 + w2_input_scale=torch.ones(1),
87 + is_trtllm=True,
88 + )
89 +
90 + expected = w13.new_zeros((1, 2 * padded_intermediate, hidden_size))
91 + expected[:, :intermediate] = up
92 + expected[:, padded_intermediate : padded_intermediate + intermediate] = gate
93 + assert layer.moe_config.intermediate_size_per_partition == padded_intermediate
94 + assert torch.equal(padded_w31, expected)
95 +
96 +
97 def test_static_fp8_moe_input_scales_remain_scalar() -> None:
98 a1_scale, a2_scale = process_fp8_input_tensor_strategy_moe(
99 torch.tensor([0.25, 0.5]),
tests/quantization/test_gfx950_moe.py
+86
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.fused_moe.config import (
16 )
17 from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
18 Mxfp4MoeBackend,
19 + _requires_qwen38_tep8_emulation,
20 select_mxfp4_moe_backend,
21 )
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -72,6 +73,91 @@ def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
73 )
74
75
76 +def _make_qwen38_tep8_moe_config() -> FusedMoEConfig:
77 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
78 +
79 + return FusedMoEConfig(
80 + num_experts=512,
81 + experts_per_token=9,
82 + hidden_dim=2560,
83 + intermediate_size=640,
84 + num_local_experts=64,
85 + num_logical_experts=512,
86 + moe_parallel_config=FusedMoEParallelConfig(
87 + tp_size=1,
88 + tp_rank=0,
89 + pcp_size=1,
90 + pcp_rank=0,
91 + dp_size=1,
92 + dp_rank=0,
93 + ep_size=8,
94 + ep_rank=0,
95 + sp_size=1,
96 + use_ep=True,
97 + all2all_backend="allgather_reducescatter",
98 + enable_eplb=False,
99 + ),
100 + activation=MoEActivation.SILU,
101 + in_dtype=torch.bfloat16,
102 + device="cuda",
103 + routing_method=RoutingMethodType.Renormalize,
104 + )
105 +
106 +
107 +def test_qwen38_tep8_requires_emulation_only_on_gfx950(monkeypatch):
108 + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle
109 +
110 + config = _make_qwen38_tep8_moe_config()
111 + monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
112 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
113 +
114 + assert _requires_qwen38_tep8_emulation(config, kMxfp4Dynamic)
115 +
116 + config.moe_parallel_config.ep_size = 4
117 + config.num_local_experts = 128
118 + assert not _requires_qwen38_tep8_emulation(config, kMxfp4Dynamic)
119 +
120 + config = _make_qwen38_tep8_moe_config()
121 + monkeypatch.setattr(mxfp4_oracle.current_platform, "is_rocm", lambda: False)
122 + assert not _requires_qwen38_tep8_emulation(config, kMxfp4Dynamic)
123 +
124 +
125 +@pytest.mark.parametrize(
126 + "requested_backend,expected_backend",
127 + [
128 + ("auto", Mxfp4MoeBackend.EMULATION),
129 + ("aiter", Mxfp4MoeBackend.AITER_MXFP4_MXFP4),
130 + ],
131 +)
132 +def test_qwen38_tep8_auto_fallback_respects_explicit_backend(
133 + requested_backend,
134 + expected_backend,
135 + monkeypatch,
136 +):
137 + import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle
138 +
139 + class SupportedExperts:
140 + @staticmethod
141 + def is_supported_config(*args, **kwargs):
142 + return True, None
143 +
144 + config = _make_qwen38_tep8_moe_config()
145 + config.moe_backend = requested_backend
146 + monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
147 + monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
148 + monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None)
149 + monkeypatch.setattr(
150 + mxfp4_oracle, "backend_to_kernel_cls", lambda backend: [SupportedExperts]
151 + )
152 +
153 + backend, experts_cls = select_mxfp4_moe_backend(
154 + config, activation_key=kMxfp4Dynamic
155 + )
156 +
157 + assert backend == expected_backend
158 + assert experts_cls is SupportedExperts
159 +
160 +
161 @pytest.fixture
162 def mxfp4_oracle_config():
163 """Stub the config the oracle reads (``model_config.quantization_config``)
tests/quantization/test_online_mxfp4.py
+21
@@ -604,3 +604,24 @@ def test_online_mxfp4_dense_matches_quark(
604 keys=("weight", "weight_scale"),
605 packed_weight_keys=("weight",),
606 )
607 +
608 +
609 +def test_quark_ocp_mx_dense_rejects_unaligned_input_partition():
610 + scheme = object.__new__(QuarkOCP_MX)
611 + layer = torch.nn.Module()
612 + layer.prefix = "model.layers.0.mlp.shared_expert.down_proj"
613 +
614 + with pytest.raises(
615 + ValueError,
616 + match=(
617 + "input size per partition of 80, which must be divisible by "
618 + "the OCP MX group size 32"
619 + ),
620 + ):
621 + scheme.create_weights(
622 + layer=layer,
623 + output_partition_sizes=[2560],
624 + input_size_per_partition=80,
625 + params_dtype=torch.bfloat16,
626 + weight_loader=lambda *_: None,
627 + )
tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py
+51
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
11 prepare_nvfp4_moe_layer_for_fi_or_cutlass,
12 )
13 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
14 + align_fp4_moe_weights_for_fi,
15 align_trtllm_fp4_moe_hidden_dim_for_fi,
16 )
17
@@ -107,3 +108,53 @@ def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple():
108 assert torch.count_nonzero(out_w13_scale[:, :, hidden_dim // 16 :]) == 0
109 assert torch.count_nonzero(out_w2[:, hidden_dim:, :]) == 0
110 assert torch.count_nonzero(out_w2_scale[:, hidden_dim:, :]) == 0
111 +
112 +
113 +def test_align_trtllm_fp4_moe_intermediate_pads_gate_and_up_separately():
114 + # Qwen3.8-Flash-Next has intermediate_size=640. At TP4 each rank owns
115 + # 160 values, which TRTLLM pads to the next 64-value boundary (192).
116 + intermediate = 160
117 + padded_intermediate = 192
118 + hidden_dim = 32
119 +
120 + gate = torch.ones((1, intermediate, hidden_dim // 2), dtype=torch.uint8)
121 + up = torch.full_like(gate, 2)
122 + w13 = torch.cat((gate, up), dim=1)
123 + gate_scale = torch.full((1, intermediate, hidden_dim // 16), 3, dtype=torch.uint8)
124 + up_scale = torch.full_like(gate_scale, 4)
125 + w13_scale = torch.cat((gate_scale, up_scale), dim=1)
126 + w2 = torch.full((1, hidden_dim, intermediate // 2), 5, dtype=torch.uint8)
127 + w2_scale = torch.full((1, hidden_dim, intermediate // 16), 6, dtype=torch.uint8)
128 +
129 + out_w13, out_w13_scale, out_w2, out_w2_scale, out_intermediate = (
130 + align_fp4_moe_weights_for_fi(
131 + w13,
132 + w13_scale,
133 + w2,
134 + w2_scale,
135 + is_act_and_mul=True,
136 + min_alignment=64,
137 + )
138 + )
139 +
140 + assert out_intermediate == padded_intermediate
141 + torch.testing.assert_close(out_w13[:, :intermediate], gate)
142 + torch.testing.assert_close(
143 + out_w13[:, padded_intermediate : padded_intermediate + intermediate], up
144 + )
145 + torch.testing.assert_close(out_w13_scale[:, :intermediate], gate_scale)
146 + torch.testing.assert_close(
147 + out_w13_scale[:, padded_intermediate : padded_intermediate + intermediate],
148 + up_scale,
149 + )
150 +
151 + assert torch.count_nonzero(out_w13[:, intermediate:padded_intermediate]) == 0
152 + assert torch.count_nonzero(out_w13[:, padded_intermediate + intermediate :]) == 0
153 + assert torch.count_nonzero(out_w13_scale[:, intermediate:padded_intermediate]) == 0
154 + assert (
155 + torch.count_nonzero(out_w13_scale[:, padded_intermediate + intermediate :]) == 0
156 + )
157 + torch.testing.assert_close(out_w2[:, :, : intermediate // 2], w2)
158 + torch.testing.assert_close(out_w2_scale[:, :, : intermediate // 16], w2_scale)
159 + assert torch.count_nonzero(out_w2[:, :, intermediate // 2 :]) == 0
160 + assert torch.count_nonzero(out_w2_scale[:, :, intermediate // 16 :]) == 0
tests/v1/core/test_contiguous_kv_packing.py
+312
@@ -9,23 +9,34 @@ block dim (a contiguous region per layer) or inside it (all layers' pages within
9 block); the allocation is the same either way.
10 """
11
12 +from dataclasses import replace
13 from unittest.mock import MagicMock
14
15 import pytest
16 import torch
17
18 from vllm.config import CacheConfig
19 +from vllm.v1.core.kv_cache_manager import KVCacheManager
20 from vllm.v1.core.kv_cache_utils import (
21 _get_kv_cache_bytes_per_block,
22 + _get_packed_kv_cache_groups,
23 _pool_bytes_per_block,
24 + generate_scheduler_kv_cache_config,
25 get_kv_cache_config_from_groups,
26 + get_kv_cache_groups,
27 + resolve_kv_cache_block_sizes,
28 )
29 from vllm.v1.kv_cache_interface import (
30 + CircularBufferSpec,
31 FullAttentionSpec,
32 KVCacheGroupSpec,
33 KVCacheLayout,
34 + KVCacheSpec,
35 + MambaSpec,
36 MLAAttentionSpec,
37 + SlidingWindowMLASpec,
38 UniformTypeKVCacheSpecs,
39 + iter_layer_specs,
40 )
41 from vllm.v1.worker.utils import allocate_kv_cache
42
@@ -86,6 +97,307 @@ def _bind(config, layout: str):
97 return allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout[layout], None)
98
99
100 +MAIN_KV_PAGE_BYTES = 2_048
101 +COMPRESSED_PAGE_BYTES = 128
102 +NUM_CACHE_TUPLES = 3
103 +BYTES_PER_BLOCK = NUM_CACHE_TUPLES * (MAIN_KV_PAGE_BYTES + COMPRESSED_PAGE_BYTES)
104 +
105 +
106 +def _main_kv_name(layer_index: int) -> str:
107 + return f"model.layers.{layer_index}.self_attn"
108 +
109 +
110 +def _compressed_name(layer_index: int) -> str:
111 + return f"model.layers.{layer_index}.self_attn.indexer.compressed_key_cache"
112 +
113 +
114 +def _compressor_state_name(layer_index: int) -> str:
115 + return f"model.layers.{layer_index}.self_attn.indexer.raw_key_cache"
116 +
117 +
118 +def _make_csa_linear_specs(
119 + num_mamba: int = 7,
120 + num_tuples: int = NUM_CACHE_TUPLES,
121 +) -> dict[str, KVCacheSpec]:
122 + specs: dict[str, KVCacheSpec] = {}
123 + for layer_index in range(num_mamba):
124 + specs[f"model.layers.{layer_index}.linear_attn"] = MambaSpec(
125 + block_size=16,
126 + shapes=((32,),),
127 + dtypes=(torch.bfloat16,),
128 + )
129 + specs["model.layers.0.ple"] = MambaSpec(
130 + block_size=16,
131 + shapes=((24,),),
132 + dtypes=(torch.bfloat16,),
133 + tp_replicated=True,
134 + )
135 + for layer_index in range(num_tuples):
136 + specs[_main_kv_name(layer_index)] = FullAttentionSpec(
137 + block_size=16,
138 + num_kv_heads=2,
139 + head_size=16,
140 + head_size_v=16,
141 + dtype=torch.bfloat16,
142 + )
143 + specs[_compressed_name(layer_index)] = MLAAttentionSpec(
144 + block_size=16,
145 + num_kv_heads=1,
146 + head_size=16,
147 + dtype=torch.bfloat16,
148 + tokens_per_state=4,
149 + )
150 + specs[_compressor_state_name(layer_index)] = CircularBufferSpec(
151 + block_size=4,
152 + num_kv_heads=1,
153 + head_size=8,
154 + head_size_v=0,
155 + dtype=torch.bfloat16,
156 + )
157 + return specs
158 +
159 +
160 +def _shared_layout_config():
161 + config = _mock_vllm_config("BLNHC")
162 + config.scheduler_config.disable_hybrid_kv_cache_manager = False
163 + config.speculative_config = None
164 + config.model_config.max_model_len = 16
165 + config.model_config.get_total_num_hidden_layers.return_value = 64
166 + config.model_config.get_total_num_kv_heads.return_value = 2
167 + config.model_config.get_num_kv_heads.return_value = 2
168 + config.parallel_config.pipeline_parallel_size = 1
169 + config.parallel_config.decode_context_parallel_size = 1
170 + config.cache_config.block_size = 16
171 + config.cache_config.enable_prefix_caching = False
172 + config.cache_config.prefix_match_unit = None
173 + config.cache_config.mamba_cache_mode = "none"
174 + return config
175 +
176 +
177 +class TestCSALinearGrouping:
178 + """A CSA + linear-attention model (sparse attention with a compressor ring,
179 + plus sharded GDN and one TP-replicated PLE state) goes through the generic
180 + packed-group path; no model-specific branch is involved."""
181 +
182 + @staticmethod
183 + def _mamba_groups(groups):
184 + """(group, per-layer mamba spec) for every group holding mamba state."""
185 + out = []
186 + for group in groups:
187 + spec = next(iter(iter_layer_specs(group.kv_cache_spec)))
188 + if isinstance(spec, MambaSpec):
189 + out.append((group, spec))
190 + return out
191 +
192 + def test_replicated_state_gets_its_own_group(self):
193 + """The PLE state is TP-replicated and a different size, so it must not
194 + share a manager group with the sharded GDN states (the NIXL worker
195 + requires exactly one single-layer replicated group)."""
196 + groups = get_kv_cache_groups(_shared_layout_config(), _make_csa_linear_specs())
197 +
198 + mamba_groups = self._mamba_groups(groups)
199 + replicated = [g for g, spec in mamba_groups if spec.tp_replicated]
200 + assert len(replicated) == 1
201 + assert replicated[0].layer_names == ["model.layers.0.ple"]
202 + sharded = [g for g, spec in mamba_groups if not spec.tp_replicated]
203 + assert sharded, "GDN states must keep their own groups"
204 + assert sorted(n for g in sharded for n in g.layer_names) == sorted(
205 + f"model.layers.{i}.linear_attn" for i in range(7)
206 + )
207 +
208 + def test_roles_land_in_separate_groups(self):
209 + groups = get_kv_cache_groups(_shared_layout_config(), _make_csa_linear_specs())
210 +
211 + owner = next(g for g in groups if _main_kv_name(0) in g.layer_names)
212 + assert sorted(owner.layer_names) == sorted(
213 + [
214 + *(_main_kv_name(i) for i in range(NUM_CACHE_TUPLES)),
215 + *(_compressed_name(i) for i in range(NUM_CACHE_TUPLES)),
216 + ]
217 + )
218 + assert owner.kv_cache_spec.prefix_cacheable
219 +
220 + scratch = next(g for g in groups if _compressor_state_name(0) in g.layer_names)
221 + assert scratch.layer_names == [
222 + _compressor_state_name(i) for i in range(NUM_CACHE_TUPLES)
223 + ]
224 + assert not scratch.kv_cache_spec.prefix_cacheable
225 +
226 + def test_every_group_fits_one_packed_block(self):
227 + config = _shared_layout_config()
228 + groups = get_kv_cache_groups(config, _make_csa_linear_specs())
229 + bytes_per_block = _get_kv_cache_bytes_per_block(groups)
230 +
231 + pages = _pages(groups)
232 + for group in groups:
233 + assert sum(pages[n] for n in group.layer_names) <= bytes_per_block
234 +
235 + kv_cache_config = get_kv_cache_config_from_groups(
236 + config, groups, available_memory=bytes_per_block * 32
237 + )
238 + assert kv_cache_config.num_blocks == 32
239 + # Groups overlay from byte 0 of each block, so a layer never addresses
240 + # past the block it belongs to.
241 + for tensor in kv_cache_config.kv_cache_tensors:
242 + assert tensor.block_stride == bytes_per_block
243 + assert tensor.offset < bytes_per_block
244 +
245 + def test_scratch_group_survives_computed_block_truncation(self):
246 + """The scratch group contributes no computed blocks, so truncating a
247 + lookup result must skip it: its block size is the ring capacity, which
248 + neither divides the hit length nor bounds the (empty) block list."""
249 + config = _shared_layout_config()
250 + config.cache_config.enable_prefix_caching = True
251 + groups = get_kv_cache_groups(config, _make_csa_linear_specs())
252 + kv_cache_config = get_kv_cache_config_from_groups(
253 + config, groups, available_memory=BYTES_PER_BLOCK * 64
254 + )
255 + scheduler_config = generate_scheduler_kv_cache_config([kv_cache_config])
256 + manager = KVCacheManager(
257 + scheduler_config,
258 + max_model_len=8192,
259 + enable_caching=True,
260 + hash_block_size=16,
261 + scheduler_block_size=16,
262 + )
263 + blocks = manager.create_kv_cache_blocks(
264 + tuple(
265 + manager.block_pool.get_new_blocks(3)
266 + if group.kv_cache_spec.prefix_cacheable
267 + else []
268 + for group in scheduler_config.kv_cache_groups
269 + )
270 + )
271 + truncated = manager.truncate_computed_blocks(blocks, 48)
272 + scratch_index = next(
273 + i
274 + for i, group in enumerate(scheduler_config.kv_cache_groups)
275 + if not group.kv_cache_spec.prefix_cacheable
276 + )
277 + assert truncated.blocks[scratch_index] == []
278 +
279 + def test_prefix_hits_respect_compression_alignment(self):
280 + config = _shared_layout_config()
281 + config.cache_config.enable_prefix_caching = True
282 + config.cache_config.prefix_match_unit = 16
283 + config.cache_config.mamba_cache_mode = "align"
284 + specs = {
285 + name: replace(spec, mamba_cache_mode="align")
286 + if isinstance(spec, MambaSpec)
287 + else spec
288 + for name, spec in _make_csa_linear_specs(num_tuples=1).items()
289 + }
290 + groups = get_kv_cache_groups(config, specs)
291 + kv_cache_config = get_kv_cache_config_from_groups(
292 + config, groups, available_memory=8 * MAIN_KV_PAGE_BYTES
293 + )
294 + assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (16, 16)
295 +
296 + # 2 tokens is not a multiple of the compression ratio 4: a prefix hit
297 + # could land inside a partially filled compressed state.
298 + config.cache_config.prefix_match_unit = 2
299 + with pytest.raises(ValueError, match="per-state compression"):
300 + resolve_kv_cache_block_sizes(kv_cache_config, config)
301 +
302 + def test_scratch_ring_does_not_drag_hash_granularity(self):
303 + config = _shared_layout_config()
304 + config.cache_config.enable_prefix_caching = True
305 + config.cache_config.mamba_cache_mode = "align"
306 + specs = {
307 + name: replace(spec, mamba_cache_mode="align")
308 + if isinstance(spec, MambaSpec)
309 + else spec
310 + for name, spec in _make_csa_linear_specs(num_tuples=1).items()
311 + }
312 + groups = get_kv_cache_groups(config, specs)
313 + kv_cache_config = get_kv_cache_config_from_groups(
314 + config, groups, available_memory=8 * MAIN_KV_PAGE_BYTES
315 + )
316 + # The hash granularity is the GCD over prefix-cacheable groups only;
317 + # the 4-token scratch ring is excluded (it would drag it to 4).
318 + assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (16, 16)
319 +
320 + def test_compressed_attention_hashes_can_be_finer_than_cache_hits(self):
321 + config = _shared_layout_config()
322 + config.cache_config.enable_prefix_caching = True
323 + specs = {
324 + "compressed.4": MLAAttentionSpec(
325 + block_size=256,
326 + num_kv_heads=1,
327 + head_size=16,
328 + dtype=torch.bfloat16,
329 + tokens_per_state=4,
330 + ),
331 + "compressed.128": MLAAttentionSpec(
332 + block_size=256,
333 + num_kv_heads=1,
334 + head_size=16,
335 + dtype=torch.bfloat16,
336 + tokens_per_state=128,
337 + ),
338 + "compressor_state.4": SlidingWindowMLASpec(
339 + block_size=4,
340 + num_kv_heads=1,
341 + head_size=8,
342 + head_size_v=0,
343 + dtype=torch.bfloat16,
344 + sliding_window=8,
345 + ),
346 + }
347 + groups = get_kv_cache_groups(config, specs)
348 + kv_cache_config = get_kv_cache_config_from_groups(
349 + config, groups, available_memory=8 * MAIN_KV_PAGE_BYTES
350 + )
351 +
352 + # Hashes are computed every 4 tokens, but without an align-mode Mamba
353 + # group cache hits remain on the 256-token scheduler boundary.
354 + assert resolve_kv_cache_block_sizes(kv_cache_config, config) == (256, 4)
355 +
356 + @pytest.mark.parametrize("wide", ["unbalanced_attention", "unsplittable_state"])
357 + def test_mamba_split_measures_the_block_the_other_groups_already_force(self, wide):
358 + """Whatever fixes the block stride -- a bucket with unequal layer counts
359 + per page size, which is emitted whole, or a single-layer state that
360 + cannot be split at all -- the mamba layers must be sized against it.
361 + Sizing them against a narrower bucket splits them past what the block
362 + already fits, spending a pool block per extra group for no saving."""
363 + config = _mock_vllm_config("BLNHC")
364 + config.speculative_config = None
365 + specs = {}
366 + if wide == "unbalanced_attention":
367 + for i in range(3):
368 + specs[f"wide.{i}"] = _mla(1024)
369 + specs["wide.odd"] = _mla(800)
370 + else:
371 + specs["wide.state"] = MambaSpec(
372 + block_size=16,
373 + shapes=((51_200,),),
374 + dtypes=(torch.bfloat16,),
375 + tp_replicated=True,
376 + )
377 + # Mixed and balanced, but far narrower than the bucket above.
378 + for i in range(10):
379 + specs[f"narrow.a.{i}"] = FullAttentionSpec(
380 + block_size=16, num_kv_heads=1, head_size=8, dtype=torch.uint8
381 + )
382 + specs[f"narrow.b.{i}"] = FullAttentionSpec(
383 + block_size=16, num_kv_heads=1, head_size=16, dtype=torch.uint8
384 + )
385 + for i in range(100):
386 + specs[f"gdn.{i}"] = MambaSpec(
387 + block_size=16, shapes=((512,),), dtypes=(torch.bfloat16,)
388 + )
389 +
390 + groups = _get_packed_kv_cache_groups(config, specs)
391 + gdn = [g for g in groups if g.layer_names[0].startswith("gdn.")]
392 +
393 + assert _get_kv_cache_bytes_per_block(groups) == sum(
394 + specs[name].page_size_bytes for name in specs if name.startswith("wide.")
395 + )
396 + # That block holds every GDN state at once, so the repeat pattern alone
397 + # decides the split; the cap must not add groups on top of it.
398 + assert len(gdn) == 10
399 +
400 +
401 class TestDensePacking:
402 def test_bytes_per_block_is_largest_group(self):
403 groups, g1, g2 = _mixed_page_groups()
tests/v1/core/test_kv_cache_utils.py
+39 -48
@@ -12,7 +12,7 @@ import pytest
12 import torch
13
14 import vllm.v1.core.kv_cache_utils as kv_cache_utils
15 -from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
15 +from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
16 from vllm.config.kv_events import KVEventsConfig
17 from vllm.lora.request import LoRARequest
18 from vllm.multimodal.inputs import (
@@ -37,7 +37,6 @@ from vllm.v1.core.kv_cache_utils import (
37 get_kv_cache_groups,
38 get_max_concurrency_for_kv_cache_config,
39 get_request_block_hasher,
40 - group_and_unify_kv_cache_specs,
40 hash_block_tokens,
41 init_none_hash,
42 is_kv_cache_spec_uniform,
@@ -2226,38 +2225,6 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None):
2225 )
2226
2227
2229 -def test_group_and_unify_kv_cache_specs_no_swa_mla_returns_none():
2230 - # Without any SlidingWindowMLASpec the function does not apply.
2231 - specs = {"mla.0": new_mla_spec(), "mla.1": new_mla_spec()}
2232 - assert group_and_unify_kv_cache_specs(specs) is None
2233 -
2234 -
2235 -def test_group_and_unify_kv_cache_specs_uniform_page_size_returns_none():
2236 - # A non-DeepseekV4 model that mixes full MLA and sliding-window MLA layers
2237 - # with a uniform page size must not fall into the DeepseekV4 tuple-packing
2238 - # path; it should defer to the generic uniform-page-size grouping instead.
2239 - mla_spec = new_mla_spec()
2240 - swa_spec = new_swa_mla_spec()
2241 - assert mla_spec.page_size_bytes == swa_spec.page_size_bytes
2242 - specs = {"mla.0": mla_spec, "mla.1": new_mla_spec(), "swa.0": swa_spec}
2243 - assert group_and_unify_kv_cache_specs(specs) is None
2244 -
2245 -
2246 -def test_group_and_unify_kv_cache_specs_mixed_page_size_groups():
2247 - # DeepseekV4-style: differing page sizes across MLA and sliding-window MLA
2248 - # layers do require tuple packing, so grouping must still be produced.
2249 - mla_spec = new_mla_spec()
2250 - swa_spec = new_swa_mla_spec(head_size=1024)
2251 - assert mla_spec.page_size_bytes != swa_spec.page_size_bytes
2252 - specs = {"mla.0": mla_spec, "mla.1": new_mla_spec(), "swa.0": swa_spec}
2253 - grouped = group_and_unify_kv_cache_specs(specs)
2254 - assert grouped is not None
2255 - # One MLA group plus one sliding-window MLA group.
2256 - assert len(grouped) == 2
2257 - layer_names = {name for g in grouped for name in g.kv_cache_specs}
2258 - assert layer_names == {"mla.0", "mla.1", "swa.0"}
2259 -
2260 -
2228 def new_indexer_mla_spec(block_size=16):
2229 # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller
2230 # page size than the main MLA attention (uint8, small head), so their pages
@@ -2270,10 +2237,30 @@ def new_indexer_mla_spec(block_size=16):
2237 )
2238
2239
2240 +def test_mixed_page_size_groups_use_spec_compatibility():
2241 + specs = {}
2242 + for i in range(3):
2243 + specs[f"mla.{i}"] = new_mla_spec()
2244 + specs[f"indexer.{i}"] = new_indexer_mla_spec()
2245 + specs.update({f"swa.{i}": new_swa_mla_spec(head_size=1024) for i in range(5)})
2246 +
2247 + config = _grouping_config()
2248 + config.cache_config = CacheConfig()
2249 + config.cache_config.kv_cache_layout = "BLNHC"
2250 + groups = get_kv_cache_groups(config, specs)
2251 +
2252 + assert len(groups) == 3
2253 + assert {name for group in groups for name in group.layer_names} == set(specs)
2254 + assert sorted(len(group.layer_names) for group in groups) == [2, 3, 6]
2255 +
2256 +
2257 def _grouping_config():
2258 + cache_config = CacheConfig()
2259 + cache_config.kv_cache_layout = "LBNHC"
2260 return SimpleNamespace(
2261 scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
2262 speculative_config=None,
2263 + cache_config=cache_config,
2264 )
2265
2266
@@ -3281,10 +3268,16 @@ def test_iter_layer_specs_returns_group_members():
3268 assert list(iter_layer_specs(wrapped)) == [full, mla]
3269
3270
3284 -def _spec_decode_grouping_config(method="dspark"):
3271 +def _spec_decode_grouping_config(method="dspark", model_type=None):
3272 """Grouping config with an EAGLE-family speculative method enabled."""
3273 return SimpleNamespace(
3274 scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
3275 + cache_config=SimpleNamespace(
3276 + get_resolved_kv_cache_layout=lambda: SimpleNamespace(
3277 + is_block_outermost=True
3278 + )
3279 + ),
3280 + model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)),
3281 speculative_config=SimpleNamespace(
3282 method=method,
3283 use_eagle=lambda: True,
@@ -3322,13 +3315,12 @@ def test_draft_group_annotated_on_hybrid_general_path():
3315
3316 flagged = [g for g in groups if g.is_eagle_group]
3317 assert len(flagged) == 1
3325 - assert flagged[0].layer_names == ["draft.attn.0"]
3318 + assert "draft.attn.0" in flagged[0].layer_names
3319
3320
3321 def test_mamba_groups_never_flagged_even_when_draft_shares_a_group():
3329 - # MLAAttentionSpec.merge ORs non_causal_multi_token_decode rather than
3330 - # requiring equality, so a draft layer can share a group with target
3331 - # layers; that combined group still holds volatile draft KV and must be
3322 + # Packed uniform-type groups can contain distinct target and draft layer
3323 + # specs; the combined group still holds volatile draft KV and must be
3324 # flagged. What must never happen is a Mamba group being flagged: that
3325 # widens its lookup window to two consecutive chunks, which align-mode
3326 # checkpointing never produces, zeroing every lookup.
@@ -3347,10 +3339,8 @@ def test_mamba_groups_never_flagged_even_when_draft_shares_a_group():
3339 def test_draft_group_not_annotated_without_spec_decode():
3340 # The marker alone must not flag anything; the eagle semantics only apply
3341 # when a speculative method is actually enabled.
3350 - config = SimpleNamespace(
3351 - scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
3352 - speculative_config=None,
3353 - )
3342 + config = _spec_decode_grouping_config()
3343 + config.speculative_config = None
3344 groups = get_kv_cache_groups(config, _hybrid_specs_with_draft(draft=True))
3345
3346 assert not any(g.is_eagle_group for g in groups)
@@ -3396,12 +3386,13 @@ def _deepseek_v4_specs(model_version="deepseek_v4"):
3386 }
3387
3388
3399 -def test_deepseek_v4_draft_group_annotated_on_group_and_unify_path():
3389 +def test_deepseek_v4_draft_group_annotated_on_packed_path():
3390 # DeepseekV4's MTP block reuses the target's decoder layer, so its spec
3391 # carries no draft marker and only the positional rule can find it. This
3392 # pins the pre-existing behaviour that the unified annotator must preserve.
3393 groups = get_kv_cache_groups(
3404 - _spec_decode_grouping_config(method="mtp"), _deepseek_v4_specs()
3394 + _spec_decode_grouping_config(method="mtp", model_type="deepseek_v4"),
3395 + _deepseek_v4_specs(model_version=None),
3396 )
3397
3398 flagged = [g for g in groups if g.is_eagle_group]
@@ -3409,13 +3400,13 @@ def test_deepseek_v4_draft_group_annotated_on_group_and_unify_path():
3400 assert "model.layers.3.self_attn.attn" in flagged[0].layer_names
3401
3402
3412 -def test_deepseek_v4_annotation_requires_model_version():
3403 +def test_deepseek_v4_annotation_requires_model_type():
3404 # The positional rule is only sound for DeepseekV4, where the draft layer
3405 # is known to be registered last. Without that model gate nothing may be
3406 # flagged, however the grouping happens to fall out.
3407 groups = get_kv_cache_groups(
3417 - _spec_decode_grouping_config(method="mtp"),
3418 - _deepseek_v4_specs(model_version=None),
3408 + _spec_decode_grouping_config(method="mtp", model_type="other"),
3409 + _deepseek_v4_specs(),
3410 )
3411
3412 assert not any(g.is_eagle_group for g in groups)
tests/v1/core/test_scheduler.py
+46
@@ -1761,6 +1761,52 @@ def test_spec_decode_padding_skipped_for_diffusion():
1761 assert r2.request_id not in out.scheduled_spec_decode_tokens
1762
1763
1764 +@pytest.mark.parametrize("aligned_tokens", [1, 2, 3])
1765 +def test_spec_decode_padding_dropped_when_recurrent_alignment_clips(
1766 + aligned_tokens: int,
1767 +):
1768 + """A clipped padded request must drop its padding, not shorten it.
1769 +
1770 + The sampler derives a request's row window from its draft count, so a
1771 + padded request that keeps fewer than 1 + num_spec query rows would select
1772 + the preceding request's hidden states.
1773 + """
1774 + num_spec = 3
1775 + scheduler = create_scheduler(
1776 + num_speculative_tokens=num_spec,
1777 + enable_prefix_caching=True,
1778 + block_size=16,
1779 + )
1780 + r1, r2 = create_requests(
1781 + num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16
1782 + )
1783 +
1784 + scheduler.add_request(r1)
1785 + out = scheduler.schedule()
1786 + assert out.num_scheduled_tokens[r1.request_id] == 33
1787 + _model_output(scheduler, out, [[100]])
1788 + scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]]))
1789 +
1790 + # Model the recurrent align split that clips a padded 1+3 prompt tail at
1791 + # the next cache boundary. The real QSA block size produces all three cases
1792 + # depending on prompt length modulo four.
1793 + scheduler.need_mamba_block_aligned_split = True
1794 + original_split = scheduler._mamba_block_aligned_split
1795 +
1796 + def align_prompt_tail(request, num_new_tokens, *args):
1797 + if request is r2:
1798 + return aligned_tokens
1799 + return original_split(request, num_new_tokens, *args)
1800 +
1801 + scheduler._mamba_block_aligned_split = align_prompt_tail
1802 + scheduler.add_request(r2)
1803 + out = scheduler.schedule()
1804 +
1805 + assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3]
1806 + assert out.num_scheduled_tokens[r2.request_id] == 1
1807 + assert r2.request_id not in out.scheduled_spec_decode_tokens
1808 +
1809 +
1810 def test_spec_decode_padding_skipped_with_prefill_in_batch():
1811 """Padding is skipped when the batch contains a prefill chunk: the batch is
1812 already mixed/non-uniform, so padding a new decode request buys nothing.
tests/v1/core/test_single_type_kv_cache_manager.py
+63
@@ -14,6 +14,7 @@ from vllm.v1.core.kv_cache_utils import (
14 )
15 from vllm.v1.core.single_type_kv_cache_manager import (
16 ChunkedLocalAttentionManager,
17 + CircularBufferManager,
18 FullAttentionManager,
19 MambaManager,
20 RSWAManager,
@@ -21,6 +22,7 @@ from vllm.v1.core.single_type_kv_cache_manager import (
22 )
23 from vllm.v1.kv_cache_interface import (
24 ChunkedLocalAttentionSpec,
25 + CircularBufferSpec,
26 FullAttentionSpec,
27 MambaSpec,
28 RSWASpec,
@@ -136,6 +138,67 @@ def get_chunked_local_attention_manager(
138 )
139
140
141 +def test_circular_buffer_allocates_one_block_for_the_request_lifetime():
142 + block_size = 4
143 + spec = CircularBufferSpec(
144 + block_size=block_size,
145 + num_kv_heads=1,
146 + head_size=8,
147 + dtype=torch.bfloat16,
148 + )
149 + block_pool = BlockPool(
150 + num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size
151 + )
152 + manager = CircularBufferManager(
153 + spec,
154 + block_pool=block_pool,
155 + enable_caching=True,
156 + kv_cache_group_id=0,
157 + scheduler_block_size=block_size,
158 + needs_kv_cache_zeroing=True,
159 + max_admission_blocks_per_request=1,
160 + )
161 + request_id = "request"
162 +
163 + assert (
164 + manager.get_num_blocks_to_allocate(
165 + request_id,
166 + num_tokens=128,
167 + new_computed_blocks=[block_pool.blocks[1]],
168 + total_computed_tokens=0,
169 + num_local_computed_tokens=0,
170 + num_tokens_main_model=128,
171 + )
172 + == 1
173 + )
174 + blocks = manager.allocate_new_blocks(
175 + request_id, num_tokens=128, num_tokens_main_model=128
176 + )
177 + assert len(blocks) == 1
178 + assert not manager.records_new_block_ids
179 + assert manager.take_new_block_ids() == []
180 +
181 + for num_tokens in (256, 1024):
182 + assert (
183 + manager.get_num_blocks_to_allocate(
184 + request_id,
185 + num_tokens,
186 + [],
187 + total_computed_tokens=num_tokens - 1,
188 + num_local_computed_tokens=0,
189 + num_tokens_main_model=num_tokens,
190 + )
191 + == 0
192 + )
193 + assert manager.allocate_new_blocks(request_id, num_tokens, num_tokens) == []
194 +
195 + manager.cache_blocks(request_id, 1024)
196 + manager.remove_skipped_blocks(request_id, 1024)
197 + assert manager.get_num_common_prefix_blocks(request_id) == 0
198 + assert manager.get_num_skipped_tokens(1024) == 0
199 + assert manager.req_to_blocks[request_id] == blocks
200 +
201 +
202 def test_sliding_window_records_new_blocks_for_zeroing():
203 block_size = 2
204 spec = SlidingWindowSpec(
tests/v1/e2e/general/test_mamba_prefix_cache.py
+12 -9
@@ -234,7 +234,8 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):
234 iter(scheduler_output.num_scheduled_tokens.values())
235 )
236 assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
237 - mamba_group_ids, mamba_spec = get_mamba_groups(self.kv_cache_config)
237 + mamba_groups = get_mamba_groups(self.kv_cache_config)
238 + mamba_spec, mamba_group_ids = next(iter(mamba_groups.items()))
239 mamba_group_id = mamba_group_ids[0]
240 mamba_layer_name = self.kv_cache_config.kv_cache_groups[
241 mamba_group_id
@@ -321,7 +322,8 @@ def get_fake_process_mamba_fn(
322 assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 0
323 else:
324 assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 2
324 - mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
325 + mamba_groups = get_mamba_groups(kv_cache_config)
326 + mamba_spec, mamba_group_ids = next(iter(mamba_groups.items()))
327 mamba_group_id = mamba_group_ids[0]
328 mamba_layer_name = kv_cache_config.kv_cache_groups[
329 mamba_group_id
@@ -936,12 +938,12 @@ def _run_mamba_prefix_cache_mrv1(
938 cleanup_dist_env_and_memory()
939
940
939 -@create_new_process_for_each_test()
941 +@create_new_process_for_each_test("spawn")
942 def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
943 _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False)
944
945
944 -@create_new_process_for_each_test()
946 +@create_new_process_for_each_test("spawn")
947 def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch):
948 _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=True)
949
@@ -975,11 +977,12 @@ def _run_mamba_prefix_cache_mrv2(
977 forward_context = (
978 model_state.vllm_config.compilation_config.static_forward_context
979 )
978 - group_ids, _ = get_mamba_groups(kv_cache_config)
979 - for group_id in group_ids:
980 - block_table = block_tables[group_id]
981 - for layer_name in kv_cache_config.kv_cache_groups[group_id].layer_names:
982 - yield forward_context[layer_name].kv_cache[-1], block_table
980 + mamba_groups = get_mamba_groups(kv_cache_config)
981 + for group_ids in mamba_groups.values():
982 + for group_id in group_ids:
983 + block_table = block_tables[group_id]
984 + for layer_name in kv_cache_config.kv_cache_groups[group_id].layer_names:
985 + yield forward_context[layer_name].kv_cache[-1], block_table
986
987 def temporal_block(temporal_state, block_table, col):
988 # Resolving the block id for assertions is a deliberate D2H.
tests/v1/kv_connector/unit/test_hidden_states_connector.py
+52 -7
@@ -15,6 +15,8 @@ from vllm.v1.kv_cache_interface import (
15 FullAttentionSpec,
16 HiddenStateCacheSpec,
17 KVCacheGroupSpec,
18 + KVCacheLayout,
19 + MambaSpec,
20 MLAAttentionSpec,
21 SlidingWindowMLASpec,
22 )
@@ -99,12 +101,12 @@ def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config():
101 assert block_size == 16
102
103
102 -# ---- MLA-verifier absorption ------------------------------------------------
104 +# ---- Packed MLA grouping ----------------------------------------------------
105
106
105 -def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier():
106 - # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding-
107 - # window MLA verifier absorbs it into the MLA group instead of isolating it.
107 +def test_hidden_state_group_isolated_from_packed_mla_groups():
108 + # HiddenStateCacheSpec subclasses MLAAttentionSpec, but grouping must pull
109 + # it out before packing compatible MLA cache specs.
110 dt = torch.bfloat16
111 spec = {
112 "layers.0.mla": MLAAttentionSpec(
@@ -116,11 +118,54 @@ def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier():
118 "cache_only_layers.61": _hidden(64),
119 }
120 vllm_config = SimpleNamespace(
121 + cache_config=SimpleNamespace(
122 + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC
123 + ),
124 scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
125 speculative_config=None,
126 )
127 groups = get_kv_cache_groups(vllm_config, spec)
123 - assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups)
128 + hidden_group_ids = [
129 + i
130 + for i, group in enumerate(groups)
131 + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
132 + ]
133 + assert len(hidden_group_ids) == 1
134 cfg = SimpleNamespace(kv_cache_groups=groups)
125 - with pytest.raises(ValueError, match="MLA verifiers are unsupported"):
126 - ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
135 + assert (
136 + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == hidden_group_ids[0]
137 + )
138 +
139 +
140 +def test_hidden_state_group_isolated_from_packed_mixed_page_groups():
141 + # Packed grouping keeps groups with unequal page sizes (blocks are strided
142 + # by the widest group), so the hidden group is appended without padding.
143 + hidden = _hidden(16)
144 + spec = {
145 + "layers.0.attn": _full(16),
146 + "layers.1.mamba": MambaSpec(
147 + block_size=16, shapes=((1024,),), dtypes=(torch.float32,)
148 + ),
149 + "cache_only_layers.61": hidden,
150 + }
151 + vllm_config = SimpleNamespace(
152 + cache_config=SimpleNamespace(
153 + get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC
154 + ),
155 + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
156 + speculative_config=None,
157 + )
158 + groups = get_kv_cache_groups(vllm_config, spec)
159 + page_sizes = {g.kv_cache_spec.page_size_bytes for g in groups}
160 + assert len(page_sizes) > 1, "spec must exercise the packed grouping path"
161 + hidden_group_ids = [
162 + i
163 + for i, group in enumerate(groups)
164 + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
165 + ]
166 + assert len(hidden_group_ids) == 1
167 + assert groups[hidden_group_ids[0]].kv_cache_spec == hidden
168 + cfg = SimpleNamespace(kv_cache_groups=groups)
169 + assert (
170 + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == hidden_group_ids[0]
171 + )
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+5
@@ -865,6 +865,10 @@ def _make_mock_worker_for_desc_ids(
865 worker._group_spec_types = group_spec_types
866 worker.block_len_per_layer = block_len_per_layer or [100]
867 worker._conv_decomp = None
868 + worker._is_csa_linear = False
869 + worker._ssm_region_indices = [0] if has_mamba else []
870 + worker._ple_group_index = None
871 + worker._ple_region_index = None
872 if has_mamba:
873 from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( # noqa: E501
874 MambaConvSplitInfo,
@@ -1823,6 +1827,7 @@ def test_push_write_hybrid_mla_replicates_attention():
1827 worker.shutdown = lambda: None # skeleton worker: silence __del__
1828 worker.use_mla = True
1829 worker._has_mamba = True
1830 + worker._is_csa_linear = False
1831 worker._group_spec_types = (MLAAttentionSpec, MambaSpec)
1832 worker.transfer_topo = MagicMock()
1833 worker.transfer_topo.tp_ratio.return_value = -2
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+299
@@ -645,3 +645,302 @@ def test_mismatched_mla_kernel_page_rejected_for_mla_hybrid():
645 meta_r.block_lens = [x // 2 for x in worker.block_len_per_layer]
646 with pytest.raises((AssertionError, RuntimeError)):
647 worker.add_remote_agent(meta_r, remote_tp_rank=0, remote_tp_size=2)
648 +
649 +
650 +def _make_csa_linear_ple_worker(scratch_aliases: str = "compressed"):
651 + """``scratch_aliases`` selects which pages the compressor ring overlays."""
652 + from unittest.mock import MagicMock
653 +
654 + from vllm.config import set_current_vllm_config
655 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
656 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
657 + NixlConnectorWorker,
658 + )
659 + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
660 + from vllm.v1.kv_cache_interface import (
661 + CircularBufferSpec,
662 + FullAttentionSpec,
663 + KVCacheConfig,
664 + KVCacheGroupSpec,
665 + KVCacheTensor,
666 + MambaSpec,
667 + MLAAttentionSpec,
668 + UniformTypeKVCacheSpecs,
669 + )
670 +
671 + main_kv_specs = {
672 + f"main_kv.{index}": FullAttentionSpec(
673 + block_size=4,
674 + num_kv_heads=1,
675 + head_size=16,
676 + dtype=torch.float16,
677 + )
678 + for index in range(2)
679 + }
680 + compressed_specs = {
681 + f"compressed.{index}": MLAAttentionSpec(
682 + block_size=4,
683 + num_kv_heads=1,
684 + head_size=64,
685 + dtype=torch.float16,
686 + tokens_per_state=2,
687 + )
688 + for index in range(2)
689 + }
690 + compressor_state_specs = {
691 + f"compressor_state.{index}": CircularBufferSpec(
692 + block_size=2,
693 + num_kv_heads=1,
694 + head_size=32,
695 + dtype=torch.float16,
696 + )
697 + for index in range(2)
698 + }
699 + sharded_mamba_spec = MambaSpec(
700 + block_size=1,
701 + shapes=((8, 3), (1, 4, 4)),
702 + dtypes=(torch.float16, torch.float32),
703 + page_size_padded=256,
704 + mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
705 + )
706 + ple_spec = MambaSpec(
707 + block_size=1,
708 + shapes=((6, 3),),
709 + dtypes=(torch.float16,),
710 + page_size_padded=256,
711 + mamba_type=MambaAttentionBackendEnum.SHORT_CONV,
712 + tp_replicated=True,
713 + )
714 + compressed_sparse_specs = {
715 + name: spec
716 + for index in range(2)
717 + for name, spec in (
718 + (f"main_kv.{index}", main_kv_specs[f"main_kv.{index}"]),
719 + (f"compressed.{index}", compressed_specs[f"compressed.{index}"]),
720 + )
721 + }
722 +
723 + def group(specs):
724 + uniform = UniformTypeKVCacheSpecs.from_specs(specs)
725 + assert uniform is not None
726 + return KVCacheGroupSpec(list(specs), uniform)
727 +
728 + tensor_regions: tuple[tuple[str, ...], ...]
729 + if scratch_aliases == "compressed":
730 + tensor_regions = (
731 + ("main_kv.0", "mamba.sharded", "mamba.ple"),
732 + ("main_kv.1",),
733 + ("compressed.0", "compressor_state.0"),
734 + ("compressed.1", "compressor_state.1"),
735 + )
736 + else:
737 + assert scratch_aliases == "main_kv"
738 + tensor_regions = (
739 + ("main_kv.0", "mamba.sharded", "mamba.ple", "compressor_state.0"),
740 + ("main_kv.1", "compressor_state.1"),
741 + ("compressed.0",),
742 + ("compressed.1",),
743 + )
744 + region_size = 512
745 + page_size = 256
746 + kv_cache_config = KVCacheConfig(
747 + num_blocks=2,
748 + kv_cache_tensors=[
749 + KVCacheTensor(
750 + size=len(tensor_regions) * region_size,
751 + layers=[layer_name],
752 + layer_stride=region_size,
753 + block_stride=page_size,
754 + offset=region_index * region_size,
755 + )
756 + for region_index, layer_names in enumerate(tensor_regions)
757 + for layer_name in layer_names
758 + ],
759 + kv_cache_groups=[
760 + group(compressed_sparse_specs),
761 + group(compressor_state_specs),
762 + KVCacheGroupSpec(["mamba.sharded"], sharded_mamba_spec),
763 + KVCacheGroupSpec(["mamba.ple"], ple_spec),
764 + ],
765 + )
766 +
767 + vllm_config = create_vllm_config(block_size=4)
768 + vllm_config.cache_config.enable_prefix_caching = False
769 + vllm_config.kv_transfer_config.kv_buffer_device = "cuda"
770 + fake_backend = MagicMock()
771 + fake_backend.get_supported_kernel_block_sizes.return_value = [4]
772 + fake_backend.get_name.return_value = "TEST_ATTN"
773 + fake_backend.full_cls_name.return_value = "test.AttentionBackend"
774 + fake_platform = MagicMock()
775 + fake_platform.device_type = "cuda"
776 + fake_platform.get_nixl_memory_type.return_value = "VRAM"
777 +
778 + with (
779 + patch.object(bw, "NixlWrapper", _RecordingNixl),
780 + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),
781 + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),
782 + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]),
783 + patch.object(bw, "current_platform", fake_platform),
784 + patch(
785 + "vllm.model_executor.layers.mamba.mamba_utils.get_conv_state_layout",
786 + return_value="DS",
787 + ),
788 + set_current_vllm_config(vllm_config),
789 + ):
790 + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config)
791 + tensors = [torch.zeros((2, 256), dtype=torch.uint8) for _ in range(4)]
792 + worker.register_kv_caches(
793 + {
794 + layer_name: tensors[region_index]
795 + for region_index, layer_names in enumerate(tensor_regions)
796 + for layer_name in layer_names
797 + }
798 + )
799 + return worker
800 +
801 +
802 +@pytest.mark.cpu_test
803 +def test_csa_linear_registration_discovers_shared_regions():
804 + worker = _make_csa_linear_ple_worker()
805 +
806 + assert worker._ssm_region_indices == [0]
807 + assert worker._ple_group_index == 3
808 + assert worker._ple_region_index == 0
809 + assert worker._region_is_mla == [False, False, True, True]
810 + assert worker._scratch_region_indices == [2, 3]
811 +
812 +
813 +@pytest.mark.cpu_test
814 +def test_csa_linear_scratch_descs_follow_the_pages_they_overlay():
815 + """The scratch ring is addressed by the regions it actually registered in,
816 + not by inferring that it must sit on the MLA pages: overlaying the main KV
817 + pages instead moves its descriptors with it."""
818 + worker = _make_csa_linear_ple_worker(scratch_aliases="main_kv")
819 +
820 + # Same MLA regions as before, but the ring now lives in the non-MLA ones.
821 + assert worker._region_is_mla == [False, False, True, True]
822 + assert worker._scratch_region_indices == [0, 1]
823 +
824 + desc_ids = worker._compute_desc_ids(
825 + block_ids=([1], [0], [1], [1]),
826 + dst_num_blocks=2,
827 + block_size_ratio=None,
828 + physical_blocks_per_logical=1,
829 + )
830 + # Scratch group (block 0) now maps to regions 0 and 1 -> descs 0 and 2,
831 + # where the compressed-page placement gave 4 and 6.
832 + assert desc_ids.tolist()[4:6] == [0, 2]
833 +
834 +
835 +@pytest.mark.cpu_test
836 +def test_csa_linear_selects_compressed_sparse_compressor_state_and_mamba_regions():
837 + worker = _make_csa_linear_ple_worker()
838 +
839 + desc_ids = worker._compute_desc_ids(
840 + block_ids=([1], [0], [1], [1]),
841 + dst_num_blocks=2,
842 + block_size_ratio=None,
843 + physical_blocks_per_logical=1,
844 + )
845 + assert desc_ids.tolist() == [
846 + 1,
847 + 3,
848 + 5,
849 + 7,
850 + 4,
851 + 6,
852 + 9,
853 + 11,
854 + 13,
855 + 15,
856 + 17,
857 + ]
858 +
859 + mamba_only = worker._compute_desc_ids(
860 + block_ids=([], [], [1], [1]),
861 + dst_num_blocks=2,
862 + block_size_ratio=None,
863 + physical_blocks_per_logical=1,
864 + )
865 + assert mamba_only.tolist() == [9, 11, 13, 15, 17]
866 +
867 +
868 +@pytest.mark.cpu_test
869 +def test_csa_linear_ple_descriptor_is_not_split():
870 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
871 + TPMapping,
872 + )
873 +
874 + worker = _make_csa_linear_ple_worker()
875 + bases = [0x10000, 0x20000, 0x30000, 0x40000]
876 + attention = worker._build_fa_local(bases, block_size_ratio=1)
877 + mamba = worker._build_mamba_local(bases)
878 + assert attention.shape == (8, 3)
879 + assert mamba.shape == (10, 3)
880 + assert mamba[-2:, 0].tolist() == [0x10000, 0x10100]
881 + assert mamba[-2:, 1].tolist() == [256, 256]
882 +
883 + mapping = TPMapping(
884 + source_ranks_per_group=((0,), (0,), (0, 1), (0, 1)),
885 + all_source_ranks=(0, 1),
886 + rank_to_attention_slot={0: 0, 1: 0},
887 + rank_offset_factor=0,
888 + )
889 + splits = list(
890 + worker._build_local_splits_from_plan(
891 + mapping,
892 + np.concatenate([attention, mamba]),
893 + num_fa_descs=len(attention),
894 + )
895 + )
896 + assert len(splits) == 2
897 + assert [row[1] for row in splits[0][-2:]] == [256, 256]
898 + assert [row[1] for row in splits[1][-2:]] == [256, 256]
899 + assert splits[0][-3][1] * 2 == mamba[-3, 1]
900 + assert splits[1][-3][1] * 2 == mamba[-3, 1]
901 +
902 +
903 +@pytest.mark.cpu_test
904 +def test_csa_linear_remote_ple_is_copied_whole():
905 + from types import SimpleNamespace
906 +
907 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
908 + NixlAgentMetadata,
909 + )
910 +
911 + worker = _make_csa_linear_ple_worker()
912 + metadata = NixlAgentMetadata(
913 + engine_id="remote",
914 + agent_metadata=b"agent",
915 + kv_caches_base_addr=[
916 + 0x10000,
917 + 0x20000,
918 + 0x30000,
919 + 0x40000,
920 + ],
921 + device_id=0,
922 + num_blocks=2,
923 + block_lens=[256] * 4,
924 + block_strides=[256] * 4,
925 + kv_cache_layout="HND",
926 + block_size=4,
927 + ssm_sizes=(24, 32),
928 + attn_backend_name="test",
929 + physical_blocks_per_logical_kv_block=1,
930 + )
931 + descriptors = worker._build_mamba_remote(
932 + metadata,
933 + tp_ratio=-2,
934 + transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
935 + )
936 + assert descriptors.shape == (10, 3)
937 + assert descriptors[-2:, 0].tolist() == [0x10000, 0x10100]
938 + assert descriptors[-2:, 1].tolist() == [256, 256]
939 +
940 + metadata.block_lens[0] = 128
941 + with pytest.raises(ValueError, match="PLE pages require identical"):
942 + worker._build_mamba_remote(
943 + metadata,
944 + tp_ratio=-2,
945 + transfer_info=SimpleNamespace(remote_physical_blocks_per_logical=1),
946 + )
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+70
@@ -359,6 +359,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
359 w._physical_blocks_per_logical_kv_block = 1
360 # Single non-hybrid attention group, matching the stub block id lists.
361 w._has_mamba = False
362 + w._is_csa_linear = False
363 w._group_spec_types = (FullAttentionSpec,)
364 w._engine_ttl = 0.0
365 w._engine_last_active = {}
@@ -1198,6 +1199,75 @@ class TestPushWriterMlaReplication:
1199 # engine thread never sees a partial set and double-frees the request.
1200 assert sorted(w._sending_transfers["p-req"]) == [1000, 1001]
1201
1202 + def test_csa_linear_writes_every_group_to_each_target(self):
1203 + from types import SimpleNamespace
1204 +
1205 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
1206 + RemoteMeta,
1207 + ReqMeta,
1208 + )
1209 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
1210 + TPMapping,
1211 + )
1212 +
1213 + engine_id = "decode-engine"
1214 + w = _StubWriterWorker.fresh()
1215 + w.world_size = 2
1216 + w.use_mla = False
1217 + w._has_mamba = True
1218 + w._is_csa_linear = True
1219 + w.nixl_wrapper = MagicMock()
1220 + w.transfer_topo = MagicMock()
1221 + w.transfer_topo.get_engine_info.return_value = SimpleNamespace(
1222 + remote_tp_size=4,
1223 + remote_block_size=16,
1224 + remote_physical_blocks_per_logical=1,
1225 + )
1226 + w.transfer_topo.tp_ratio.return_value = -2
1227 + w.transfer_topo.handshake_target_ranks.return_value = (0, 1)
1228 + w.tp_mappings = {
1229 + engine_id: TPMapping(
1230 + source_ranks_per_group=((0,), (0,), (0, 1), (0, 1)),
1231 + all_source_ranks=(0, 1),
1232 + rank_to_attention_slot={0: 0, 1: 0},
1233 + rank_offset_factor=0,
1234 + )
1235 + }
1236 + w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids
1237 + w.dst_xfer_side_handles = {engine_id: {0: 1000, 1: 1001}}
1238 + w.src_xfer_handles_by_tp_ratio = {(-2, 16): [2000, 2001]}
1239 +
1240 + writes = []
1241 +
1242 + def _fake_xfer(**kwargs):
1243 + writes.append(kwargs)
1244 + return 3000 + kwargs["read_spec"].remote_rank
1245 +
1246 + w._xfer_blocks = _fake_xfer
1247 + block_ids = ([10], [20], [30], [40])
1248 + meta = ReqMeta(
1249 + local_block_ids=block_ids,
1250 + local_physical_block_ids=block_ids,
1251 + tp_size=2,
1252 + remote=RemoteMeta(
1253 + block_ids=([50], [60], [70], [80]),
1254 + host="",
1255 + port=0,
1256 + engine_id=engine_id,
1257 + request_id="d-req",
1258 + ),
1259 + )
1260 +
1261 + w._xfer_blocks_for_req(req_id="p-req", meta=meta)
1262 +
1263 + assert [write["read_spec"].remote_rank for write in writes] == [0, 1]
1264 + assert [write["read_spec"].local_block_ids for write in writes] == [
1265 + list(block_ids),
1266 + list(block_ids),
1267 + ]
1268 + assert [write["local_xfer_side_handle"] for write in writes] == [2000, 2001]
1269 + assert sorted(w._sending_transfers["p-req"]) == [3000, 3001]
1270 +
1271
1272 class TestPushPrefixCaching:
1273 """Partial prefix-cache hit on D: D preallocates only its *uncomputed*
tests/v1/kv_connector/unit/test_tp_mapping.py
+31
@@ -154,6 +154,9 @@ def _make_mock_worker_for_splits(group_spec_types):
154 worker.block_len_per_layer = []
155 worker.num_regions = 0
156 worker._region_is_mla = []
157 + worker._conv_decomp = SimpleNamespace(local_conv_offsets=())
158 + worker._ssm_region_indices = [0] if MambaSpec in group_spec_types else []
159 + worker._ple_region_index = None
160 return worker
161
162
@@ -282,3 +285,31 @@ class TestMambaPlanSplitHandles:
285
286 with pytest.raises(AssertionError, match="Head-sharded"):
287 list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2, 2))
288 +
289 +
290 +@pytest.mark.parametrize(
291 + ("total_kv_heads", "local_tp", "remote_tp", "compatible"),
292 + [
293 + # 2 heads: TP1 packs both heads into a page, TP >= 2 packs one.
294 + (2, 1, 2, False),
295 + (2, 2, 1, False),
296 + (2, 1, 1, True),
297 + (2, 2, 4, True),
298 + (2, 4, 2, True),
299 + # 4 heads: the boundary moves with the head count.
300 + (4, 2, 4, False),
301 + (4, 2, 2, True),
302 + (4, 4, 8, True),
303 + ],
304 +)
305 +def test_csa_linear_tp_layout_boundary(total_kv_heads, local_tp, remote_tp, compatible):
306 + worker = object.__new__(NixlConnectorWorker)
307 + worker.world_size = local_tp
308 + worker._is_csa_linear = True
309 + worker.transfer_topo = SimpleNamespace(total_num_kv_heads=total_kv_heads)
310 +
311 + if compatible:
312 + worker._validate_csa_linear_tp_layout(remote_tp)
313 + else:
314 + with pytest.raises(ValueError, match="KV-head sharding boundary"):
315 + worker._validate_csa_linear_tp_layout(remote_tp)
tests/v1/worker/test_gpu_autoregressive_speculator.py
+46 -3
@@ -104,6 +104,47 @@ def _make_speculator(
104 return speculator
105
106
107 +@pytest.mark.parametrize(("hc_mult", "expected"), [(None, 64), (4, 256)])
108 +def test_speculator_uses_draft_model_hidden_size(monkeypatch, hc_mult, expected):
109 + # Qwen4Exp targets expose multi-stream HC residuals to the drafter.
110 + monkeypatch.setattr(base_spec_module, "_target_feeds_hc_residual", lambda _: True)
111 + hf_config = SimpleNamespace()
112 + if hc_mult is not None:
113 + hf_config.hc_mult = hc_mult
114 + draft_model_config = SimpleNamespace(
115 + hf_config=hf_config,
116 + get_hidden_size=lambda: 64,
117 + get_vocab_size=lambda: 32,
118 + )
119 + speculative_config = SimpleNamespace(
120 + method="mtp",
121 + num_speculative_tokens=3,
122 + draft_model_config=draft_model_config,
123 + use_local_argmax_reduction=False,
124 + draft_sample_method="greedy",
125 + )
126 + vllm_config = SimpleNamespace(
127 + speculative_config=speculative_config,
128 + scheduler_config=SimpleNamespace(
129 + max_num_seqs=2,
130 + max_num_batched_tokens=8,
131 + ),
132 + model_config=SimpleNamespace(
133 + max_model_len=32,
134 + dtype=torch.float32,
135 + use_fp64_gumbel=False,
136 + ),
137 + parallel_config=SimpleNamespace(
138 + data_parallel_size=1,
139 + data_parallel_rank=0,
140 + ),
141 + )
142 +
143 + speculator = _TestSpeculator(vllm_config, torch.device("cpu"))
144 +
145 + assert speculator.hidden_size == expected
146 +
147 +
148 def test_mm_support_configured_after_model_load(monkeypatch):
149 target_model_config = object()
150 draft_model_config = object()
@@ -197,9 +238,11 @@ def test_load_model_disables_mm_support_for_text_only_drafter(monkeypatch):
238
239 assert not speculator.supports_mm_inputs
240 assert warning_messages == [
200 - "Draft model _TextOnlyDraftModel does not support external multimodal "
201 - "embeddings. Embeddings from the target model will not be passed to the "
202 - "drafter; using text-only draft inputs instead."
241 + (
242 + "Draft model _TextOnlyDraftModel does not support external multimodal "
243 + "embeddings. Embeddings from the target model will not be passed to the "
244 + "drafter; using text-only draft inputs instead."
245 + )
246 ]
247
248
tests/v1/worker/test_gpu_block_table.py
+36
@@ -132,6 +132,42 @@ def test_block_tables_apply_staged_writes_single_group():
132 )
133
134
135 +def test_block_tables_skip_custom_slot_mapping_groups():
136 + device = torch.device("cuda")
137 + block_tables = BlockTables(
138 + block_sizes=[8, 262144],
139 + max_num_reqs=1,
140 + max_num_batched_tokens=4,
141 + max_num_blocks_per_group=[1, 1],
142 + device=device,
143 + kernel_block_sizes=[8, 262144],
144 + slot_mapping_enabled=[False, True],
145 + )
146 + block_tables.append_block_ids(
147 + req_index=0,
148 + new_block_ids=([7], [12]),
149 + overwrite=True,
150 + )
151 + block_tables.apply_staged_writes()
152 +
153 + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device)
154 + query_start_loc = torch.tensor([0, 2], dtype=torch.int32, device=device)
155 + positions = torch.tensor([153797, 165757], dtype=torch.int64, device=device)
156 + slot_mappings = block_tables.compute_slot_mappings(
157 + idx_mapping,
158 + query_start_loc,
159 + positions,
160 + num_tokens_padded=2,
161 + )
162 + torch.accelerator.synchronize()
163 +
164 + assert slot_mappings[0].tolist() == [-1, -1]
165 + assert slot_mappings[1].tolist() == [
166 + 12 * 262144 + 153797,
167 + 12 * 262144 + 165757,
168 + ]
169 +
170 +
171 @pytest.mark.parametrize("cp_rank", range(4))
172 def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int):
173 """DCP interleave is expressed in logical-block token coordinates."""
tests/v1/worker/test_gpu_model_runner_v2.py
+92
@@ -8,15 +8,107 @@ import torch
8
9 import vllm.v1.worker.gpu.model_runner as model_runner_module
10 from vllm.v1.kv_cache_interface import (
11 + CircularBufferSpec,
12 FullAttentionSpec,
13 KVCacheConfig,
14 KVCacheGroupSpec,
15 MambaSpec,
16 + UniformTypeKVCacheSpecs,
17 )
18 from vllm.v1.worker.gpu.block_table import BlockTables
19 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
20
21
22 +def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
23 + runner = GPUModelRunner.__new__(GPUModelRunner)
24 + runner.max_model_len = 262144
25 + runner.is_encoder_decoder = False
26 + runner.dcp_size = 1
27 + runner.dcp_rank = 0
28 + runner.cp_interleave = 1
29 + runner.cache_config = SimpleNamespace(enable_prefix_caching=True)
30 + parallel_config = SimpleNamespace(
31 + decode_context_parallel_size=1,
32 + cp_kv_cache_interleave_size=1,
33 + )
34 + runner.parallel_config = parallel_config
35 + runner.vllm_config = SimpleNamespace(
36 + parallel_config=parallel_config,
37 + cache_config=SimpleNamespace(mamba_cache_mode="none"),
38 + )
39 + runner.model_state = SimpleNamespace(
40 + get_additional_cg_support=lambda: (),
41 + num_new_sampled_tokens_per_step=1,
42 + )
43 + runner.speculator = None
44 + runner.req_states = []
45 + runner.input_buffers = SimpleNamespace(query_start_loc=None)
46 + runner.vocab_size = 1
47 + runner.max_num_reqs = 1
48 + runner.max_num_tokens = 2
49 + runner.device = torch.device("cuda")
50 +
51 + raw_spec = CircularBufferSpec(
52 + block_size=8,
53 + num_kv_heads=1,
54 + head_size=128,
55 + dtype=torch.bfloat16,
56 + )
57 + compressed_spec = FullAttentionSpec(
58 + block_size=262144,
59 + num_kv_heads=1,
60 + head_size=128,
61 + dtype=torch.bfloat16,
62 + )
63 + kv_cache_config = KVCacheConfig(
64 + num_blocks=1,
65 + kv_cache_tensors=[],
66 + kv_cache_groups=[
67 + KVCacheGroupSpec(
68 + layer_names=["raw"],
69 + kv_cache_spec=UniformTypeKVCacheSpecs(
70 + block_size=8,
71 + kv_cache_specs={"raw": raw_spec},
72 + ),
73 + ),
74 + KVCacheGroupSpec(layer_names=["compressed"], kv_cache_spec=compressed_spec),
75 + ],
76 + )
77 +
78 + class FakeAttnCGSupport:
79 + def narrow(self, *args):
80 + return self
81 +
82 + attn_cg_support = FakeAttnCGSupport()
83 + monkeypatch.setattr(
84 + model_runner_module,
85 + "init_attn_backend",
86 + lambda *args: ([], attn_cg_support, [8, 262144]),
87 + )
88 + monkeypatch.setattr(
89 + model_runner_module,
90 + "maybe_create_adaptive_verification_manager",
91 + lambda **kwargs: None,
92 + )
93 +
94 + captured = {}
95 +
96 + class BlockTablesCaptured(Exception):
97 + pass
98 +
99 + def capture_block_tables(**kwargs):
100 + captured.update(kwargs)
101 + raise BlockTablesCaptured
102 +
103 + monkeypatch.setattr(model_runner_module, "BlockTables", capture_block_tables)
104 +
105 + with pytest.raises(BlockTablesCaptured):
106 + runner.initialize_kv_cache(kv_cache_config)
107 +
108 + assert captured["max_num_blocks_per_group"] == [1, 1]
109 + assert captured["slot_mapping_enabled"] == [False, True]
110 +
111 +
112 @pytest.mark.parametrize(
113 ("mamba_cache_mode", "num_speculative_blocks", "expected"),
114 [
tests/v1/worker/test_gpu_warmup_blocks.py
+15 -2
@@ -20,6 +20,7 @@ from vllm.config.vllm import VllmConfig
20 from vllm.platforms import current_platform
21 from vllm.utils.math_utils import cdiv
22 from vllm.v1.kv_cache_interface import (
23 + CircularBufferSpec,
24 FullAttentionSpec,
25 KVCacheConfig,
26 KVCacheGroupSpec,
@@ -53,6 +54,18 @@ def _attention_group() -> KVCacheGroupSpec:
54 )
55
56
57 +def _circular_group() -> KVCacheGroupSpec:
58 + return KVCacheGroupSpec(
59 + ["circular"],
60 + CircularBufferSpec(
61 + block_size=BLOCK_SIZE,
62 + num_kv_heads=1,
63 + head_size=1,
64 + dtype=torch.float32,
65 + ),
66 + )
67 +
68 +
69 def _mamba_group(mamba_cache_mode: str) -> KVCacheGroupSpec:
70 # Name carries the mode so several groups can coexist in one config.
71 return KVCacheGroupSpec(
@@ -212,8 +225,7 @@ def test_warmup_reserves_mamba_speculative_blocks(mamba_cache_mode):
225
226
227 def _hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
215 - """Full attention plus one Mamba group per cache mode, so a single manager
216 - exercises every branch `_reserved_block_count` has.
228 + """Attention, circular, and Mamba groups exercise every reservation branch.
229
230 "none" and "all" reach the same branch of both `_reserved_block_count` and
231 `MambaManager`, which tests only for "align". They are still both listed:
@@ -226,6 +238,7 @@ def _hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
238 kv_cache_tensors=[],
239 kv_cache_groups=[
240 _attention_group(),
241 + _circular_group(),
242 _mamba_group("none"),
243 _mamba_group("all"),
244 _mamba_group("align"),
tests/v1/worker/test_mamba_utils.py
+188 -12
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Callable
3 from dataclasses import dataclass
4 from typing import Any
5 from unittest.mock import MagicMock, patch
@@ -10,11 +9,20 @@ import pytest
9 import torch
10
11 from vllm.model_executor.layers.mamba.mamba_utils import (
12 + MambaStateCopyFunc,
13 + MambaStateCopyFuncsByType,
14 get_conv_copy_spec,
15 get_temporal_copy_spec,
16 )
17 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
18 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
17 -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec
19 +from vllm.v1.kv_cache_interface import (
20 + KVCacheConfig,
21 + KVCacheGroupSpec,
22 + MambaSpec,
23 + UniformTypeKVCacheSpecs,
24 +)
25 +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
26 from vllm.v1.worker.mamba_utils import (
27 MambaCopyBuffers,
28 MambaSpecDecodeGPUContext,
@@ -22,17 +30,21 @@ from vllm.v1.worker.mamba_utils import (
30 batch_memcpy,
31 collect_mamba_copy_meta,
32 do_mamba_copy_block,
33 + get_mamba_groups,
34 preprocess_mamba,
35 stage_postprocess_inputs_to_gpu,
36 )
37
29 -MambaStateCopyFunc = Callable[..., Any]
30 -
38 # Conv + temporal copy specs, in the order the tests' MambaSpec shapes expect.
32 -_COPY_FUNCS: tuple[MambaStateCopyFunc, ...] = (
39 +_DEFAULT_COPY_FUNCS: tuple[MambaStateCopyFunc, ...] = (
40 get_conv_copy_spec,
41 get_temporal_copy_spec,
42 )
43 +_COPY_FUNCS: MambaStateCopyFuncsByType = {
44 + MambaAttentionBackendEnum.MAMBA2: _DEFAULT_COPY_FUNCS,
45 + MambaAttentionBackendEnum.GDN_ATTN: _DEFAULT_COPY_FUNCS,
46 + MambaAttentionBackendEnum.SHORT_CONV: (get_conv_copy_spec,),
47 +}
48
49
50 def postprocess_mamba(
@@ -41,7 +53,7 @@ def postprocess_mamba(
53 input_batch: Any,
54 requests: dict[str, Any],
55 forward_context: dict[str, Any],
44 - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
56 + mamba_state_copy_funcs: MambaStateCopyFuncsByType,
57 copy_bufs: "MambaCopyBuffers",
58 ):
59 """CPU reference for the align-mode postprocess.
@@ -147,7 +159,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
159 input_batch,
160 {}, # requests
161 {}, # forward_context
150 - (), # mamba_state_copy_funcs
162 + {}, # mamba_state_copy_funcs
163 copy_bufs,
164 )
165
@@ -430,7 +442,10 @@ def _make_requests(
442
443
444 def _make_copy_bufs(
433 - cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device
445 + cfg: _TestConfig,
446 + kv_cache_config: KVCacheConfig,
447 + device: torch.device,
448 + copy_funcs: MambaStateCopyFuncsByType = _COPY_FUNCS,
449 ) -> MambaCopyBuffers:
450 """Create MambaCopyBuffers for the Python path."""
451
@@ -440,13 +455,16 @@ def _make_copy_bufs(
455 return MambaCopyBuffers.create(
456 max_num_reqs=cfg.max_num_reqs,
457 kv_cache_config=kv_cache_config,
443 - copy_funcs=(get_conv_copy_spec, get_temporal_copy_spec),
458 + copy_funcs=copy_funcs,
459 make_buffer=make_buffer,
460 )
461
462
463 def _make_gpu_ctx(
449 - cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device
464 + cfg: _TestConfig,
465 + kv_cache_config: KVCacheConfig,
466 + device: torch.device,
467 + copy_funcs: MambaStateCopyFuncsByType = _COPY_FUNCS,
468 ) -> MambaSpecDecodeGPUContext:
469 """Create MambaSpecDecodeGPUContext for the GPU path."""
470
@@ -456,12 +474,170 @@ def _make_gpu_ctx(
474 return MambaSpecDecodeGPUContext.create(
475 max_num_reqs=cfg.max_num_reqs,
476 kv_cache_config=kv_cache_config,
459 - num_state_types=2,
477 + copy_funcs=copy_funcs,
478 device=device,
479 make_buffer=make_buffer,
480 )
481
482
483 +def test_mamba_groups_support_different_state_specs():
484 + gdn_spec = MambaSpec(
485 + block_size=16,
486 + shapes=((4, 4), (2, 4, 4)),
487 + dtypes=(torch.float16, torch.float32),
488 + mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
489 + mamba_cache_mode="align",
490 + num_speculative_blocks=3,
491 + )
492 + short_conv_spec = MambaSpec(
493 + block_size=16,
494 + shapes=((12, 8),),
495 + dtypes=(torch.float16,),
496 + mamba_type=MambaAttentionBackendEnum.SHORT_CONV,
497 + mamba_cache_mode="align",
498 + num_speculative_blocks=3,
499 + )
500 + kv_cache_config = KVCacheConfig(
501 + num_blocks=8,
502 + kv_cache_tensors=[],
503 + kv_cache_groups=[
504 + KVCacheGroupSpec(["gdn.0", "gdn.1"], gdn_spec),
505 + KVCacheGroupSpec(["ple.0"], short_conv_spec),
506 + ],
507 + )
508 +
509 + mamba_groups = get_mamba_groups(kv_cache_config)
510 + assert mamba_groups == {gdn_spec: [0], short_conv_spec: [1]}
511 +
512 + model_state = object.__new__(MambaHybridModelState)
513 + model_state._mamba_group_ids = []
514 + model_state._mamba_spec = None
515 + group_ids, representative_spec = model_state._get_mamba_group_info(kv_cache_config)
516 + assert group_ids == [0, 1]
517 + assert representative_spec is gdn_spec
518 +
519 + cfg = _TestConfig()
520 + copy_bufs = _make_copy_bufs(cfg, kv_cache_config, torch.device("cpu"))
521 + assert copy_bufs.src_ptrs.cpu.numel() == cfg.max_num_reqs * 5
522 + assert copy_bufs.dst_ptrs.cpu.numel() == cfg.max_num_reqs * 5
523 + assert copy_bufs.sizes.cpu.numel() == cfg.max_num_reqs * 5
524 +
525 + ctx = _make_gpu_ctx(cfg, kv_cache_config, torch.device("cpu"))
526 + assert ctx.num_states == 5
527 +
528 + forward_context = {}
529 + gdn_spec_layer_names = ["gdn.0", "gdn.1"]
530 + for layer_name in gdn_spec_layer_names:
531 + attention = MagicMock()
532 + attention.kv_cache = (
533 + torch.empty(8, 4, 4, dtype=torch.float16),
534 + torch.empty(8, 2, 4, 4, dtype=torch.float32),
535 + )
536 + forward_context[layer_name] = attention
537 + ple_attention = MagicMock()
538 + ple_attention.kv_cache = (torch.empty(8, 12, 8, dtype=torch.float16),)
539 + forward_context["ple.0"] = ple_attention
540 + assert gdn_spec_layer_names == kv_cache_config.kv_cache_groups[0].layer_names
541 +
542 + block_tables = [
543 + torch.zeros(4, 8, dtype=torch.int32),
544 + torch.zeros(4, 8, dtype=torch.int32),
545 + ]
546 + ctx.initialize_from_forward_context(
547 + kv_cache_config,
548 + forward_context,
549 + _COPY_FUNCS,
550 + block_tables,
551 + )
552 + assert ctx.is_initialized
553 + assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 1]
554 + assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12]
555 +
556 +
557 +def test_mamba_groups_support_mixed_specs_in_uniform_group():
558 + gdn_spec = MambaSpec(
559 + block_size=16,
560 + shapes=((4, 4), (2, 4, 4)),
561 + dtypes=(torch.float16, torch.float32),
562 + mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
563 + mamba_cache_mode="align",
564 + num_speculative_blocks=3,
565 + )
566 + short_conv_spec = MambaSpec(
567 + block_size=16,
568 + shapes=((12, 8),),
569 + dtypes=(torch.float16,),
570 + mamba_type=MambaAttentionBackendEnum.SHORT_CONV,
571 + mamba_cache_mode="align",
572 + num_speculative_blocks=3,
573 + )
574 + layer_specs = {
575 + "gdn.0": gdn_spec,
576 + "gdn.1": gdn_spec,
577 + "ple.0": short_conv_spec,
578 + }
579 + kv_cache_config = KVCacheConfig(
580 + num_blocks=8,
581 + kv_cache_tensors=[],
582 + kv_cache_groups=[
583 + KVCacheGroupSpec(
584 + list(layer_specs),
585 + UniformTypeKVCacheSpecs(
586 + block_size=16,
587 + kv_cache_specs=layer_specs,
588 + ),
589 + )
590 + ],
591 + )
592 +
593 + assert kv_cache_config.has_mamba_layers
594 + assert get_mamba_groups(kv_cache_config) == {
595 + gdn_spec: [0],
596 + short_conv_spec: [0],
597 + }
598 +
599 + # The group is listed once per spec, so its id must not be repeated.
600 + model_state = object.__new__(MambaHybridModelState)
601 + model_state._mamba_group_ids = []
602 + model_state._mamba_spec = None
603 + group_ids, representative_spec = model_state._get_mamba_group_info(kv_cache_config)
604 + assert group_ids == [0]
605 + assert representative_spec is gdn_spec
606 +
607 + cfg = _TestConfig()
608 + copy_bufs = _make_copy_bufs(cfg, kv_cache_config, torch.device("cpu"))
609 + assert copy_bufs.mamba_group_ids == [0]
610 + assert copy_bufs.src_ptrs.cpu.numel() == cfg.max_num_reqs * 5
611 +
612 + ctx = _make_gpu_ctx(cfg, kv_cache_config, torch.device("cpu"))
613 + assert ctx.mamba_group_ids == [0]
614 + assert ctx.num_states == 5
615 +
616 + forward_context = {
617 + "gdn.0": MagicMock(
618 + kv_cache=(
619 + torch.empty(8, 4, 4, dtype=torch.float16),
620 + torch.empty(8, 2, 4, 4, dtype=torch.float32),
621 + )
622 + ),
623 + "gdn.1": MagicMock(
624 + kv_cache=(
625 + torch.empty(8, 4, 4, dtype=torch.float16),
626 + torch.empty(8, 2, 4, 4, dtype=torch.float32),
627 + )
628 + ),
629 + "ple.0": MagicMock(kv_cache=(torch.empty(8, 12, 8, dtype=torch.float16),)),
630 + }
631 + ctx.initialize_from_forward_context(
632 + kv_cache_config,
633 + forward_context,
634 + _COPY_FUNCS,
635 + [torch.zeros(4, 8, dtype=torch.int32)],
636 + )
637 + assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 0]
638 + assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12]
639 +
640 +
641 # -----------------------------------------------------------------------------
642 # stage_postprocess_inputs_to_gpu: single-pass staging into pinned views
643 # -----------------------------------------------------------------------------
@@ -620,7 +796,7 @@ def _run_gpu_postprocess(
796 *,
797 kv_cache_config: KVCacheConfig,
798 forward_context: dict[str, Any],
623 - copy_funcs: tuple,
799 + copy_funcs: MambaStateCopyFuncsByType,
800 block_table: torch.Tensor,
801 req_ids: list[str],
802 num_accepted_tokens: list[int],
vllm/_custom_ops.py
+2
@@ -2817,6 +2817,7 @@ def fused_gdn_decode_post_conv_mtp(
2817 out: torch.Tensor | None = None,
2818 scale: float = 128**-0.5,
2819 norm_eps: float = 1e-5,
2820 + output_gate_activation: str = "silu",
2821 ) -> torch.Tensor:
2822 if out is None:
2823 out = torch.empty_like(output_gate)
@@ -2835,6 +2836,7 @@ def fused_gdn_decode_post_conv_mtp(
2836 out,
2837 scale,
2838 norm_eps,
2839 + output_gate_activation,
2840 )
2841 return out
2842
vllm/config/compilation.py
+3
@@ -768,6 +768,9 @@ class CompilationConfig:
768 "vllm::mamba_mixer2",
769 "vllm::mamba_mixer",
770 "vllm::short_conv",
771 + "vllm::qwen4_exp_compute_ple_ngram_ids",
772 + "vllm::qwen4_exp_ple_short_conv",
773 + "vllm::qwen4_exp_qsa_with_output",
774 "vllm::linear_attention",
775 "vllm::qwen_gdn_attention_core",
776 "vllm::qwen_gdn_attention_core_fused_norm_packed",
vllm/config/scheduler.py
+4
@@ -215,6 +215,10 @@ class SchedulerConfig:
215 # https://github.com/vllm-project/vllm/issues/29585
216 factors.append(self.max_num_batched_tokens)
217
218 + # PLE and other model components allocate static per-request buffers.
219 + # Their shapes are captured in compiled graphs.
220 + factors.append(self.max_num_seqs)
221 +
222 hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
223 return hash_str
224
vllm/config/speculative.py
+43
@@ -47,6 +47,7 @@ MTPModelTypes = Literal[
47 "exaone_moe_mtp",
48 "exaone4_5_mtp",
49 "qwen3_next_mtp",
50 + "qwen4_exp_mtp",
51 "qwen3_5_mtp",
52 "longcat_flash_mtp",
53 "bailing_hybrid_v3_mtp",
@@ -425,6 +426,9 @@ class SpeculativeConfig:
426 """The specific revision to use for the draft model code on Hugging Face
427 Hub. It can be a branch name, a tag name, or a commit id. If unspecified,
428 will use the default version."""
429 + index_share_for_mtp_iteration: bool | None = None
430 + """Override whether MTP iterations reuse the first step's sparse indices.
431 + If `None`, use the value from the draft model's Hugging Face config."""
432
433 # Advanced control
434 disable_padded_drafter_batch: bool = False
@@ -623,6 +627,15 @@ class SpeculativeConfig:
627 # Convert to tuple to make it hashable
628 factors.append(tuple(layer_ids))
629
630 + if self.method == "mtp" and self.draft_model_config is not None:
631 + factors.append(
632 + getattr(
633 + self.draft_model_config.hf_config,
634 + "index_share_for_mtp_iteration",
635 + False,
636 + )
637 + )
638 +
639 hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
640 return hash_str
641
@@ -805,6 +818,27 @@ class SpeculativeConfig:
818 hf_config.update(
819 {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]}
820 )
821 + if hf_config.model_type in {"qwen4_exp", "qwen4_exp_text"}:
822 + hf_config.model_type = "qwen4_exp_mtp"
823 + if hf_config.model_type == "qwen4_exp_mtp":
824 + text_config = get_hf_text_config(hf_config)
825 + n_predict = getattr(
826 + text_config,
827 + "mtp_num_hidden_layers",
828 + getattr(text_config, "num_nextn_predict_layers", None),
829 + )
830 + share_mtp_indices = getattr(
831 + text_config, "index_share_for_mtp_iteration", False
832 + )
833 + hf_config.update(
834 + {
835 + # hc_count is the HC stream multiplier for Qwen MTP feedback.
836 + "hc_mult": int(text_config.hc_count),
837 + "n_predict": n_predict,
838 + "architectures": ["Qwen4ExpMTP"],
839 + "index_share_for_mtp_iteration": share_mtp_indices,
840 + }
841 + )
842
843 architectures = getattr(hf_config, "architectures", []) or []
844 if initial_architecture == "BailingMoeV3ForCausalLM":
@@ -1461,6 +1495,15 @@ class SpeculativeConfig:
1495 )
1496 )
1497
1498 + if self.index_share_for_mtp_iteration is not None:
1499 + if self.method != "mtp" or self.draft_model_config is None:
1500 + raise ValueError(
1501 + "index_share_for_mtp_iteration is only supported with method='mtp'"
1502 + )
1503 + self.draft_model_config.hf_config.index_share_for_mtp_iteration = (
1504 + self.index_share_for_mtp_iteration
1505 + )
1506 +
1507 if self.method != "dspark" and self.enable_adaptive_verification:
1508 raise ValueError("Adaptive verification only supported with DSpark")
1509
vllm/distributed/device_communicators/custom_all_reduce.py
+2
@@ -405,6 +405,8 @@ class CustomAllreduce:
405 def should_custom_ar(self, inp: torch.Tensor):
406 if self.disabled or self.world_size > 8:
407 return False
408 + if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16):
409 + return False
410 inp_size = inp.numel() * inp.element_size()
411 # custom allreduce requires input byte size to be multiples of 16
412 if inp_size % 16 != 0:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+188 -51
@@ -72,12 +72,15 @@ from vllm.utils.gpu_sync_debug import gpu_sync_allowed
72 from vllm.utils.network_utils import make_zmq_path
73 from vllm.utils.torch_utils import async_tensor_h2d
74 from vllm.v1.kv_cache_interface import (
75 + CircularBufferSpec,
76 FullAttentionSpec,
77 KVCacheLayout,
78 + KVCacheSpec,
79 MambaSpec,
80 MLAAttentionSpec,
81 SlidingWindowMLASpec,
82 UniformTypeKVCacheSpecs,
83 + iter_layer_specs,
84 )
85 from vllm.v1.worker.block_table import BlockTable
86 from vllm.v1.worker.utils import select_common_block_size
@@ -118,7 +121,12 @@ class NixlBaseConnectorWorker:
121 # NIXL regions per SSM layer = conv sub-projections + 1 SSM temporal
122 # (Mamba2/GDN: 3+1=4; Mamba1: 1+1=2).
123 ssm_regions_per_layer = len(self._conv_decomp.local_conv_offsets) + 1
121 - num_ssm_regions = len(self.block_len_per_layer) * ssm_regions_per_layer
124 + # Count only the regions that actually hold SSM state; must match
125 + # the descriptors emitted by _build_mamba_local.
126 + num_ssm_regions = (
127 + len(self._ssm_region_indices or self.block_len_per_layer)
128 + * ssm_regions_per_layer
129 + )
130
131 num_blocks = dst_num_blocks
132 if block_size_ratio is not None:
@@ -146,12 +154,19 @@ class NixlBaseConnectorWorker:
154 all_descs: list[np.ndarray] = []
155 for i, group in enumerate(block_ids):
156 group_arr = np.asarray(group)
149 - if _is_attention_spec(self._group_spec_types[i]):
150 - fa_region_ids = np.arange(self.num_regions)[:, None]
157 + spec_type = self._group_spec_types[i]
158 + if _is_attention_spec(spec_type):
159 + # A scratch cache lives only in its own regions; every other
160 + # attention group spans all of them.
161 + fa_region_ids = (
162 + np.asarray(self._scratch_region_indices, dtype=np.int64)
163 + if spec_type is CircularBufferSpec
164 + else np.arange(self.num_regions)
165 + )[:, None]
166 all_descs.append(
167 (fa_region_ids * num_blocks + group_arr[None, :]).flatten()
168 )
154 - elif _is_ssm_spec(self._group_spec_types[i]):
169 + elif _is_ssm_spec(spec_type):
170 # NOTE (NickLucche) SSM and Attention block regions can
171 # be exchanged arbitrarily by manager. Therefore, descs
172 # are laid out as:
@@ -159,7 +174,11 @@ class NixlBaseConnectorWorker:
174 # num_fa_descs offset must be computed per-engine since
175 # P and D can have different num_blocks (and thus
176 # different FA desc counts).
162 - ssm_region_ids = np.arange(num_ssm_regions)[:, None]
177 + ssm_region_ids = (
178 + np.arange(num_ssm_regions, num_ssm_regions + 1)
179 + if i == self._ple_group_index
180 + else np.arange(num_ssm_regions)
181 + )[:, None]
182 all_descs.append(
183 (
184 ssm_region_ids * logical_blocks
@@ -206,6 +225,10 @@ class NixlBaseConnectorWorker:
225
226 # Per-FA-descriptor replicate flag, in _build_fa_local emission order.
227 fa_desc_replicated = self._fa_desc_replicated(num_fa_descs)
228 + sharded_desc_end = len(src_blocks_data) - (
229 + self._logical_num_blocks if self._ple_region_index is not None else 0
230 + )
231 + assert num_fa_descs <= sharded_desc_end
232
233 assert block_size_ratio == 1 or fa_num_splits == 1 or all(fa_desc_replicated), (
234 "Head-sharded attention reads with P_TP > D_TP and heterogeneous "
@@ -226,9 +249,11 @@ class NixlBaseConnectorWorker:
249 # SPLIT (full-attn): this rank's head slice.
250 chunk = local_len // fa_num_splits
251 handle.append((addr + fa_slot * chunk, chunk, dev))
229 - else:
252 + elif j < sharded_desc_end:
253 chunk = local_len // ssm_num_splits
254 handle.append((addr + p_idx * chunk, chunk, dev))
255 + else:
256 + handle.append((addr, local_len, dev))
257 yield handle
258
259 def _needs_split_local_xfer_handles(self, tp_ratio: int, plan: TPMapping) -> bool:
@@ -307,19 +332,22 @@ class NixlBaseConnectorWorker:
332 )
333 )
334
335 + self.kv_cache_config = kv_cache_config
336 + # Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers.
337 + self._layer_specs: dict[str, KVCacheSpec] = {}
338 + for group in kv_cache_config.transfer_groups:
339 + group_spec = group.kv_cache_spec
340 + if isinstance(group_spec, UniformTypeKVCacheSpecs):
341 + self._layer_specs.update(group_spec.kv_cache_specs)
342 + else:
343 + self._layer_specs.update(dict.fromkeys(group.layer_names, group_spec))
344 self._is_hma_required = (
345 not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
346 and any(
313 - not isinstance(g.kv_cache_spec, FullAttentionSpec)
314 - for g in kv_cache_config.transfer_groups
347 + not isinstance(spec, FullAttentionSpec)
348 + for spec in self._layer_specs.values()
349 )
350 )
317 - self.kv_cache_config = kv_cache_config
318 - self._layer_specs = {
319 - layer: group.kv_cache_spec
320 - for group in kv_cache_config.transfer_groups
321 - for layer in group.layer_names
322 - }
351
352 # ---- Model state (derived from model config) ----
353 mamba_ssm_size = (0, 0)
@@ -328,9 +356,29 @@ class NixlBaseConnectorWorker:
356 # conv sub-projections are contiguous in memory.
357 self._conv_decomp: MambaConvSplitInfo | None = None
358 self._has_mamba = any(
331 - isinstance(g.kv_cache_spec, MambaSpec)
332 - for g in kv_cache_config.transfer_groups
359 + isinstance(spec, MambaSpec) for spec in self._layer_specs.values()
360 )
361 + self._is_csa_linear = any(
362 + get_representative_spec_type(group.kv_cache_spec) is CircularBufferSpec
363 + for group in kv_cache_config.kv_cache_groups
364 + )
365 + self._ple_group_index: int | None = None
366 + if self._is_csa_linear:
367 + ple_groups = [
368 + (index, group)
369 + for index, group in enumerate(kv_cache_config.kv_cache_groups)
370 + if group.layer_names
371 + and all(
372 + isinstance(spec, MambaSpec) and spec.tp_replicated
373 + for spec in iter_layer_specs(group.kv_cache_spec)
374 + )
375 + ]
376 + if len(ple_groups) != 1 or len(ple_groups[0][1].layer_names) != 1:
377 + raise ValueError(
378 + "CSA-linear NIXL requires exactly one PLE cache owner."
379 + )
380 + self._ple_group_index = ple_groups[0][0]
381 +
382 if self._has_mamba:
383 assert self._is_hma_required
384 from vllm.model_executor.layers.mamba.mamba_utils import (
@@ -345,6 +393,7 @@ class NixlBaseConnectorWorker:
393 spec
394 for spec in self._layer_specs.values()
395 if isinstance(spec, MambaSpec)
396 + and (not self._is_csa_linear or not spec.tp_replicated)
397 )
398 self._conv_decomp = derive_mamba_conv_split(
399 mamba_spec,
@@ -579,6 +628,12 @@ class NixlBaseConnectorWorker:
628 # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models
629 # combining both (e.g. GQA main + MLA Eagle-3 draft).
630 self._region_is_mla = list[bool]()
631 + self._ssm_region_indices = list[int]()
632 + # Regions holding a scratch cache (the CSA compressor circular buffer).
633 + # Tracked explicitly because a scratch page shares its address with the
634 + # page it overlays, so its regions cannot be recovered from spec type.
635 + self._scratch_region_indices = list[int]()
636 + self._ple_region_index: int | None = None
637
638 # Enable different block lengths for different layers *only* when MLA is used.
639 # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
@@ -632,6 +687,24 @@ class NixlBaseConnectorWorker:
687 self.block_size = kernel_block_size
688 self.num_blocks *= self._physical_blocks_per_logical_kv_block
689
690 + def _validate_csa_linear_tp_layout(self, remote_tp_size: int) -> None:
691 + """Reject P/D pairs whose main-KV pages have different head layouts.
692 +
693 + A rank at TP < total KV heads holds several heads per page, while a
694 + rank at TP >= total KV heads holds exactly one (replicated beyond
695 + that). CSA-linear pages are shared tensors that transfer as whole
696 + block units, so endpoints on opposite sides of that boundary lay
697 + them out differently and cannot exchange them.
698 + """
699 + assert self.transfer_topo is not None
700 + total_kv_heads = self.transfer_topo.total_num_kv_heads
701 + if (self.world_size < total_kv_heads) != (remote_tp_size < total_kv_heads):
702 + raise ValueError(
703 + "CSA-linear NIXL requires both endpoints on the same side of "
704 + f"the KV-head sharding boundary ({total_kv_heads} KV heads): "
705 + f"got local TP {self.world_size}, remote TP {remote_tp_size}."
706 + )
707 +
708 def _nixl_handshake(
709 self,
710 host: str,
@@ -643,6 +716,8 @@ class NixlBaseConnectorWorker:
716 notif_agents_only: bool = False,
717 ) -> tuple[dict[tuple[int, int], str], float]:
718 """Do a NIXL handshake with a remote instance."""
719 + if self._is_csa_linear:
720 + self._validate_csa_linear_tp_layout(remote_tp_size)
721
722 # the first time we connect to a remote agent.
723 # be careful, the handshake happens in a background thread.
@@ -1034,6 +1109,11 @@ class NixlBaseConnectorWorker:
1109 transfer_mode=self._TRANSFER_MODE,
1110 )
1111
1112 + if self._is_csa_linear and self.use_host_buffer:
1113 + raise NotImplementedError(
1114 + "NIXL host staging does not preserve CSA-linear shared tensors."
1115 + )
1116 +
1117 if self.use_host_buffer:
1118 self.initialize_host_xfer_buffer(kv_caches=kv_caches)
1119 assert len(self.host_xfer_buffers) == len(kv_caches), (
@@ -1059,8 +1139,14 @@ class NixlBaseConnectorWorker:
1139 caches_data = []
1140 seen_storage_addresses: set[int] = set()
1141 seen_base_addresses: list[int] = []
1142 + self._ssm_region_indices = []
1143 + self._scratch_region_indices = []
1144 + self._ple_region_index = None
1145
1146 packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values()))
1147 + # CSA-linear needs separate logical regions for attention and state
1148 + # aliases even though every view shares one block-major allocation.
1149 + packed_storage = packed_storage and not self._is_csa_linear
1150
1151 # K and V are packed into the content dim, so each attention layer is a
1152 # single NIXL region whose block transfers as one unit. Mamba layers instead
@@ -1112,14 +1198,16 @@ class NixlBaseConnectorWorker:
1198 )
1199
1200 if isinstance(layer_spec, MambaSpec):
1201 + physical_ratio = self._physical_blocks_per_logical_kv_block
1202 + block_len = physical_page_size // physical_ratio
1203 + block_stride = physical_page_size
1204 + if self._is_csa_linear and cache.ndim != 1:
1205 + # CSA-linear packs multiple cache owners in each block, so
1206 + # state views advance by the packed block stride.
1207 + block_stride = cache.stride(0) * cache.element_size()
1208 + assert block_stride % physical_ratio == 0
1209 region_specs = [
1116 - (
1117 - cache.data_ptr(),
1118 - physical_page_size
1119 - // self._physical_blocks_per_logical_kv_block,
1120 - physical_page_size
1121 - // self._physical_blocks_per_logical_kv_block,
1122 - )
1210 + (cache.data_ptr(), block_len, block_stride // physical_ratio)
1211 ]
1212 else:
1213 if cache.ndim == 1:
@@ -1131,19 +1219,18 @@ class NixlBaseConnectorWorker:
1219 block_stride = cache.stride(0) * cache.element_size()
1220 storage_is_block_major = num_blocks * block_stride == storage.nbytes()
1221 # A layer whose [H, N, C] interior is dense addresses its own page
1134 - # as one chunk, so it can own a region. When it does not, or when
1135 - # every layer is packed into this one storage, the block's whole
1136 - # row is the only contiguous unit.
1222 + # as one chunk. Otherwise the block's whole row is the only
1223 + # contiguous transfer unit.
1224 hnc_contiguous = (
1225 cache.ndim == 4
1226 and cache.stride(2) == cache.shape[3]
1227 and cache.stride(1) == cache.shape[2] * cache.shape[3]
1228 )
1229 if storage_is_block_major and (
1143 - (packed_storage and is_mla_region) or not hnc_contiguous
1230 + (packed_storage and is_mla_region)
1231 + or (not hnc_contiguous and not self._is_csa_linear)
1232 ):
1145 - # TODO(Lucas): handle TP slicing for packed_storage; for now
1146 - # restrict to MLA (DSv4) where kv is replicated.
1233 + # Packed MLA layouts transfer the complete storage row.
1234 storage_block_len = storage.nbytes() // num_blocks
1235 region_specs = [
1236 (storage_addr, storage_block_len, storage_block_len)
@@ -1170,17 +1257,27 @@ class NixlBaseConnectorWorker:
1257
1258 for base_addr, block_len, block_stride in region_specs:
1259 if base_addr in seen_base_addresses:
1173 - if is_mla_region:
1174 - # Dual-purpose HMA tensor: an MLA layer shares a region
1175 - # that a non-MLA layer registered first. MLA is not
1176 - # head-sharded, so the region must be flagged MLA.
1177 - idx = seen_base_addresses.index(base_addr)
1178 - self._region_is_mla[idx] = True
1179 - continue
1180 - seen_base_addresses.append(base_addr)
1181 - self.block_len_per_layer.append(block_len)
1182 - self.block_stride_per_layer.append(block_stride)
1183 - self._region_is_mla.append(is_mla_region)
1260 + region_index = seen_base_addresses.index(base_addr)
1261 + self._region_is_mla[region_index] |= is_mla_region
1262 + else:
1263 + region_index = len(seen_base_addresses)
1264 + seen_base_addresses.append(base_addr)
1265 + self.block_len_per_layer.append(block_len)
1266 + self.block_stride_per_layer.append(block_stride)
1267 + self._region_is_mla.append(is_mla_region)
1268 +
1269 + if isinstance(layer_spec, MambaSpec):
1270 + if layer_spec.tp_replicated:
1271 + assert self._is_csa_linear
1272 + assert self._ple_region_index in (None, region_index)
1273 + self._ple_region_index = region_index
1274 + elif region_index not in self._ssm_region_indices:
1275 + self._ssm_region_indices.append(region_index)
1276 + elif (
1277 + isinstance(layer_spec, CircularBufferSpec)
1278 + and region_index not in self._scratch_region_indices
1279 + ):
1280 + self._scratch_region_indices.append(region_index)
1281
1282 # When there's a mismatch between kbs<>bs, we rely on HMA to ensure
1283 # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
@@ -1210,6 +1307,8 @@ class NixlBaseConnectorWorker:
1307 == len(self._region_is_mla)
1308 == len(self.block_stride_per_layer)
1309 )
1310 + # Descriptor ids must be region-ordered, matching the remote side.
1311 + self._scratch_region_indices.sort()
1312
1313 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
1314 self.num_regions = len(seen_base_addresses)
@@ -1324,17 +1423,26 @@ class NixlBaseConnectorWorker:
1423 physical_per_logical = self._physical_blocks_per_logical_kv_block
1424 device_id = self.device_id
1425 block_arange = np.arange(num_blocks, dtype=np.uint64)
1327 -
1426 parts: list[np.ndarray] = []
1329 - for i, base_addr in enumerate(base_addresses):
1330 - # Jump one page_size, but ssm page_size may be bigger when kernel
1331 - # locks block size to a specific value (physical_per_logical scale).
1332 - page_stride = self.block_len_per_layer[i] * physical_per_logical
1333 - blk_addrs = base_addr + block_arange * page_stride
1427 +
1428 + region_indices = self._ssm_region_indices or range(len(base_addresses))
1429 + for i in region_indices:
1430 + base_addr = base_addresses[i]
1431 + block_stride = self.block_stride_per_layer[i] * physical_per_logical
1432 + blk_addrs = base_addr + block_arange * block_stride
1433 for off, sz in conv_offsets:
1434 parts.append(self._stack_descs(blk_addrs + off, sz, device_id))
1435 # SSM temporal state follows the conv state.
1436 parts.append(self._stack_descs(blk_addrs + conv_size, ssm_size, device_id))
1437 +
1438 + if (region_index := self._ple_region_index) is not None:
1439 + block_len = self.block_len_per_layer[region_index] * physical_per_logical
1440 + block_stride = (
1441 + self.block_stride_per_layer[region_index] * physical_per_logical
1442 + )
1443 + block_addrs = base_addresses[region_index] + block_arange * block_stride
1444 + parts.append(self._stack_descs(block_addrs, block_len, self.device_id))
1445 +
1446 return np.concatenate(parts)
1447
1448 def _build_mamba_remote(
@@ -1370,14 +1478,44 @@ class NixlBaseConnectorWorker:
1478 parts: list[np.ndarray] = []
1479 # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
1480 # block lengths vary across layers (e.g. MLA).
1373 - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1374 - page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical
1375 - blk_addrs = base_addr + block_arange * page_stride
1481 + region_indices = self._ssm_region_indices or range(
1482 + len(nixl_agent_meta.kv_caches_base_addr)
1483 + )
1484 + for i in region_indices:
1485 + base_addr = nixl_agent_meta.kv_caches_base_addr[i]
1486 + block_stride = (
1487 + nixl_agent_meta.block_strides[i] * remote_physical_per_logical
1488 + )
1489 + blk_addrs = base_addr + block_arange * block_stride
1490 for off, sz in conv_offsets:
1491 parts.append(self._stack_descs(blk_addrs + off, sz, device_id))
1492 # SSM temporal state is also TP-sharded on the heads dimension.
1493 ssm_addrs = blk_addrs + conv_size_remote + local_offset * ssm_read_size
1494 parts.append(self._stack_descs(ssm_addrs, ssm_read_size, device_id))
1495 +
1496 + if (region_index := self._ple_region_index) is not None:
1497 + local_block_len = (
1498 + self.block_len_per_layer[region_index]
1499 + * self._physical_blocks_per_logical_kv_block
1500 + )
1501 + remote_block_len = (
1502 + nixl_agent_meta.block_lens[region_index] * remote_physical_per_logical
1503 + )
1504 + if local_block_len != remote_block_len:
1505 + raise ValueError(
1506 + "PLE pages require identical P/D geometry: "
1507 + f"local={local_block_len}, remote={remote_block_len}."
1508 + )
1509 + remote_block_stride = (
1510 + nixl_agent_meta.block_strides[region_index]
1511 + * remote_physical_per_logical
1512 + )
1513 + block_addrs = (
1514 + nixl_agent_meta.kv_caches_base_addr[region_index]
1515 + + block_arange * remote_block_stride
1516 + )
1517 + parts.append(self._stack_descs(block_addrs, remote_block_len, device_id))
1518 +
1519 return np.concatenate(parts)
1520
1521 @staticmethod
@@ -2410,8 +2548,7 @@ class NixlBaseConnectorWorker:
2548 group_specs = self.kv_cache_config.transfer_groups
2549 physical_block_ids = []
2550 for i, group in enumerate(block_ids):
2413 - spec = group_specs[i].kv_cache_spec
2414 - if isinstance(spec, MambaSpec):
2551 + if _is_ssm_spec(get_representative_spec_type(group_specs[i].kv_cache_spec)):
2552 physical_block_ids.append(group)
2553 else:
2554 physical_block_ids.append(
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+2 -1
@@ -536,7 +536,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
536 def group_ids(block_ids: BlockIds, rank: int) -> BlockIds:
537 return [
538 list(block_ids[g])
539 - if (replicate_attn and _is_attention_spec(self._group_spec_types[g]))
539 + if (self._is_csa_linear and tp_ratio < 0)
540 + or (replicate_attn and _is_attention_spec(self._group_spec_types[g]))
541 or rank in plan.source_ranks_per_group[g]
542 else []
543 for g in range(num_groups)
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+14
@@ -39,6 +39,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
39 is_deep_gemm_e8m0_used,
40 )
41 from vllm.model_executor.layers.quantization.utils.quant_utils import (
42 + FP8_DTYPE,
43 QuantKey,
44 kFp8Dynamic128Sym,
45 kFp8DynamicTensorSym,
@@ -153,6 +154,19 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
154 (kFp8StaticTensorSym, kFp8StaticTensorSym),
155 (kFp8StaticTensorSym, kFp8DynamicTensorSym),
156 ]
157 + # Block-quantized FP8 with arbitrary block shape: the Triton
158 + # kernels take the block shape as a runtime argument.
159 + if (
160 + weight_key is not None
161 + and activation_key == kFp8Dynamic128Sym
162 + and weight_key.dtype == FP8_DTYPE
163 + and weight_key.symmetric
164 + and weight_key.scale.static
165 + and weight_key.scale.dtype == torch.float32
166 + and weight_key.scale.group_shape.row > 1
167 + and weight_key.scale.group_shape.col > 1
168 + ):
169 + return True
170 return (weight_key, activation_key) in supported
171
172 @staticmethod
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+34
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import math
4 from enum import Enum
5 from typing import Any
6
@@ -268,6 +269,39 @@ def map_fp8_backend(runner_backend: MoEBackend) -> Fp8MoeBackend:
269 )
270
271
272 +def refine_fp8_moe_block_shape(
273 + config: FusedMoEConfig,
274 + weight_block_size: list[int],
275 +) -> list[int] | None:
276 + """
277 + Compute a refined block shape for block-quantized FP8 MoE weights whose
278 + checkpoint blocks cannot be sharded exactly across TP ranks.
279 +
280 + TP shards the intermediate dim of the expert weights, so a per-shard size
281 + that is not a multiple of the checkpoint's block size makes the
282 + checkpoint's block scales impossible to shard exactly. When a finer block
283 + size (>= 32) divides both the checkpoint blocks and all involved dims,
284 + the weight scales can be refined to that granularity at load time (a
285 + lossless upsampling, since the refined block divides the checkpoint
286 + block). Only Triton-based kernels can consume the refined block shape:
287 + they take it as a runtime argument, while the other backends require the
288 + native 128x128 blocks. The refined shape is encoded in the QuantKey used
289 + for backend selection, so backends that only support 128x128 blocks are
290 + rejected by the oracle automatically.
291 +
292 + Returns the refined [block_n, block_k] shape, or None if no refinement
293 + is needed or possible.
294 + """
295 + block_n, block_k = weight_block_size
296 + ispp = config.intermediate_size_per_partition
297 + if ispp % block_n == 0 and (config.tp_size == 1 or ispp % block_k == 0):
298 + return None
299 + refine = math.gcd(block_n, block_k, ispp, config.hidden_dim)
300 + if refine < 32:
301 + return None
302 + return [refine, refine]
303 +
304 +
305 def select_fp8_moe_backend(
306 config: FusedMoEConfig,
307 weight_key: QuantKey | None,
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+48
@@ -474,6 +474,39 @@ def _get_requested_backends(
474 return _filter_by_activation(backends, requested_activation_key)
475
476
477 +def _requires_qwen38_tep8_emulation(
478 + config: FusedMoEConfig,
479 + activation_key: QuantKey | None,
480 +) -> bool:
481 + """Avoid the inaccurate native gfx950 kernel for Qwen3.8 TEP8.
482 +
483 + Qwen3.8 Flash Next's routed experts use the distinctive
484 + ``E=512, H=2560, N=640`` W4A4 shape. With eight-way expert parallelism,
485 + AITER operates on 64 local experts and pads ``N`` to 768. That native path
486 + is not numerically reliable on gfx950, while OCP MX emulation preserves
487 + model accuracy. Keep this guard exact so other MXFP4 shapes and smaller EP
488 + configurations retain the native backend.
489 + """
490 + parallel = config.moe_parallel_config
491 + if not (
492 + activation_key == kMxfp4Dynamic
493 + and parallel.use_ep
494 + and parallel.ep_size == 8
495 + and config.num_experts == 512
496 + and config.num_local_experts == 64
497 + and config.hidden_dim == 2560
498 + and config.intermediate_size == 640
499 + ):
500 + return False
501 +
502 + if not current_platform.is_rocm():
503 + return False
504 +
505 + from vllm.platforms.rocm import on_gfx950
506 +
507 + return on_gfx950()
508 +
509 +
510 def select_mxfp4_moe_backend(
511 config: FusedMoEConfig,
512 activation_key: QuantKey | None = None,
@@ -532,6 +565,21 @@ def select_mxfp4_moe_backend(
565 assert last_error is not None
566 raise last_error
567
568 + if _requires_qwen38_tep8_emulation(config, requested_activation_key):
569 + backend = Mxfp4MoeBackend.EMULATION
570 + logger.warning_once(
571 + "Using OCP MX emulation for the Qwen3.8 Flash Next TEP8 routed "
572 + "experts on gfx950 because the native AITER W4A4 kernel is not "
573 + "numerically reliable for this shape. Performance will be lower."
574 + )
575 + return _return_or_raise(
576 + backend,
577 + config,
578 + kMxfp4Static,
579 + requested_activation_key,
580 + activation_format,
581 + )
582 +
583 # Select kernels in order of backend.
584 AVAILABLE_BACKENDS = _filter_by_activation(
585 _get_priority_backends_for_gpt_oss(), requested_activation_key
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+7 -8
@@ -345,14 +345,13 @@ def convert_to_unquantized_kernel_format(
345
346 elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
347 is_act_and_mul = moe_config.is_act_and_mul
348 - if not is_act_and_mul:
349 - # Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK
350 - # weight layout uses block_k=128). Pad along the intermediate dim when
351 - # the model + TP split don't satisfy the constraint.
352 - w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi(
353 - w13_weight, w2_weight, is_act_and_mul, min_alignment=128
354 - )
355 - moe_config.intermediate_size_per_partition = padded_intermediate
348 + # Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK
349 + # weight layout uses block_k=128). Pad along the intermediate dim when
350 + # the model + TP split don't satisfy the constraint.
351 + w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi(
352 + w13_weight, w2_weight, is_act_and_mul, min_alignment=128
353 + )
354 + moe_config.intermediate_size_per_partition = padded_intermediate
355
356 _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
357 w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(
vllm/model_executor/layers/fused_moe/routed_experts.py
+14
@@ -832,6 +832,20 @@ class RoutedExperts(PluggableLayer):
832 FusedMoeWeightScaleSupported.GROUP.value,
833 FusedMoeWeightScaleSupported.BLOCK.value,
834 ]:
835 + scale_refine = getattr(self.quant_method, "weight_scale_refine", None)
836 + if (
837 + quant_method == FusedMoeWeightScaleSupported.BLOCK.value
838 + and scale_refine is not None
839 + ):
840 + # FP8 block scales are stored per (block_n, block_k) tile
841 + # of the unsharded weight, while the TP-sharded parameters
842 + # use a refined block grid (see Fp8MoEMethod). Upsample the
843 + # scales to the refined grid (lossless: the refined block
844 + # divides the checkpoint block) so per-rank slicing stays
845 + # exact. Dim -2 is the weight's N dim, dim -1 is K.
846 + loaded_weight = loaded_weight.repeat_interleave(
847 + scale_refine[0], dim=-2
848 + ).repeat_interleave(scale_refine[1], dim=-1)
849 self._load_model_weight_or_group_weight_scale(
850 shard_id=shard_id,
851 shard_dim=shard_dim,
vllm/model_executor/layers/mamba/abstract.py
+5
@@ -56,6 +56,10 @@ class MambaBase(AttentionLayerBase):
56 def mamba_type(self) -> MambaAttentionBackendEnum:
57 pass
58
59 + @property
60 + def is_kv_cache_tp_replicated(self) -> bool:
61 + return False
62 +
63 @abstractmethod
64 def get_state_dtype(self) -> tuple[torch.dtype, ...]:
65 pass
@@ -70,6 +74,7 @@ class MambaBase(AttentionLayerBase):
74 block_size=mamba_block_size,
75 page_size_padded=page_size_padded,
76 mamba_type=self.mamba_type,
77 + tp_replicated=self.is_kv_cache_tp_replicated,
78 mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
79 # RecoverSSM verifies the whole window off one checkpoint, so it
80 # never writes the baseline's per-draft-token state slots.
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+5 -4
@@ -519,7 +519,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
519 self.gqa_interleaved_layout
520 or self.head_k_dim != 128
521 or self.head_v_dim != 128
522 - or self.norm.activation != "silu"
522 + or self.norm.activation not in ("silu", "sigmoid")
523 or vllm_config.model_config.dtype != torch.bfloat16
524 or conv_state_dtype != torch.bfloat16
525 or recurrent_state_dtype not in FUSED_GDN_STATE_DTYPES
@@ -527,9 +527,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
527 ):
528 return (
529 "the fused CUDA kernel requires a BF16 GDN model with "
530 - "K=V=128, SiLU gating, non-interleaved GQA layout, BF16 "
531 - "convolution cache, BF16 or FP32 recurrent state, and a "
532 - "GPU with compute capability 8.0+"
530 + "K=V=128, SiLU or sigmoid gating, non-interleaved GQA "
531 + "layout, BF16 convolution cache, BF16 or FP32 recurrent "
532 + "state, and a GPU with compute capability 8.0+"
533 )
534 if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"):
535 return "torch.ops._C.fused_gdn_decode_post_conv_mtp is not built"
@@ -1766,6 +1766,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1766 out=core_attn_out,
1767 scale=self.head_k_dim**-0.5,
1768 norm_eps=self.layer_norm_epsilon,
1769 + output_gate_activation=self.norm.activation,
1770 )
1771
1772 def _forward_core_fused_norm_packed(
vllm/model_executor/layers/mamba/mamba_utils.py
+5
@@ -17,6 +17,7 @@ from vllm.utils.torch_utils import (
17 STR_DTYPE_TO_TORCH_DTYPE,
18 get_kv_cache_torch_dtype,
19 )
20 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
21
22 logger = init_logger(__name__)
23
@@ -352,6 +353,10 @@ Parameters:
353 num_accepted_tokens: int - number of accepted tokens used to compute the copy offset.
354 Range: 1 .. 1 + num_speculative_tokens (inclusive).
355 """
356 +MambaStateCopyFuncs: TypeAlias = tuple[MambaStateCopyFunc, ...]
357 +MambaStateCopyFuncsByType: TypeAlias = dict[
358 + MambaAttentionBackendEnum, MambaStateCopyFuncs
359 +]
360
361
362 def get_conv_copy_spec(
vllm/model_executor/layers/quantization/fp8.py
+53 -13
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
31 convert_to_fp8_moe_kernel_format,
32 make_fp8_moe_kernel,
33 make_fp8_moe_quant_config,
34 + refine_fp8_moe_block_shape,
35 select_fp8_moe_backend,
36 )
37 from vllm.model_executor.layers.linear import (
@@ -63,7 +64,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
64 kFp8Dynamic128Sym,
65 kFp8DynamicTensorSym,
66 kFp8DynamicTokenSym,
66 - kFp8Static128BlockSym,
67 kFp8StaticTensorSym,
68 )
69 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -481,9 +481,41 @@ class Fp8MoEMethod(FusedMoEMethodBase):
481 "weight_scale_inv" if self.block_quant else "weight_scale"
482 )
483
484 + self.weight_scale_refine: tuple[int, int] | None = None
485 + self.moe_block_shape = self.weight_block_size
486 +
487 # Set weight key and activation key for kernel compatibility
488 if self.block_quant:
486 - weight_key = kFp8Static128BlockSym
489 + assert self.weight_block_size is not None
490 + # TP shards the intermediate dim of the expert weights, so a
491 + # per-shard size that is not a multiple of the checkpoint's block
492 + # size makes the checkpoint's block scales impossible to shard
493 + # exactly. When a finer block size (>= 32) divides both the
494 + # checkpoint blocks and all involved dims, the weight scales are
495 + # refined to that granularity at load time (a lossless upsampling,
496 + # since the refined block divides the checkpoint block). The
497 + # refined block shape is encoded in the weight key, so the oracle
498 + # only selects kernels that support it (e.g. Triton, which takes
499 + # the block shape as a runtime argument).
500 + refined_shape = refine_fp8_moe_block_shape(self.moe, self.weight_block_size)
501 + if refined_shape is not None:
502 + block_n, block_k = self.weight_block_size
503 + self.weight_scale_refine = (
504 + block_n // refined_shape[0],
505 + block_k // refined_shape[1],
506 + )
507 + self.moe_block_shape = refined_shape
508 + logger.info_once(
509 + "FP8 MoE block scales refined from %s to %s to fit "
510 + "the TP-sharded intermediate size %d.",
511 + str(self.weight_block_size),
512 + str(refined_shape),
513 + self.moe.intermediate_size_per_partition,
514 + )
515 + assert self.moe_block_shape is not None
516 + weight_key = create_fp8_quant_key(
517 + static=True, group_shape=GroupShape(*self.moe_block_shape)
518 + )
519 activation_key = kFp8Dynamic128Sym
520 else:
521 weight_key = kFp8StaticTensorSym
@@ -519,6 +551,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
551
552 if self.block_quant:
553 assert self.weight_block_size is not None
554 + assert self.moe_block_shape is not None
555 + moe_block_shape = self.moe_block_shape
556 layer.weight_block_size = self.weight_block_size
557 tp_size = get_tensor_model_parallel_world_size()
558 block_n, block_k = (
@@ -530,18 +564,24 @@ class Fp8MoEMethod(FusedMoEMethodBase):
564 # layers must be divisible by block_n.
565 # Required by column parallel or enabling merged weights
566 if intermediate_size_per_partition % block_n != 0:
533 - raise ValueError(
534 - f"The output_size of gate's and up's weight = "
535 - f"{intermediate_size_per_partition} is not divisible by "
536 - f"weight quantization block_n = {block_n}."
537 - )
567 + if self.weight_scale_refine is None:
568 + raise ValueError(
569 + f"The output_size of gate's and up's weight = "
570 + f"{intermediate_size_per_partition} is not divisible by "
571 + f"weight quantization block_n = {block_n}."
572 + )
573 + # Use the refined block grid for the scale parameters; the
574 + # loader upsamples the checkpoint scales accordingly.
575 + block_n, block_k = moe_block_shape
576 if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
577 # Required by row parallel
540 - raise ValueError(
541 - f"The input_size of down's weight = "
542 - f"{intermediate_size_per_partition} is not divisible by "
543 - f"weight quantization block_k = {block_k}."
544 - )
578 + if self.weight_scale_refine is None:
579 + raise ValueError(
580 + f"The input_size of down's weight = "
581 + f"{intermediate_size_per_partition} is not divisible by "
582 + f"weight quantization block_k = {block_k}."
583 + )
584 + block_n, block_k = moe_block_shape
585
586 # WEIGHTS
587 w13_weight = torch.nn.Parameter(
@@ -747,7 +787,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
787 w2_scale=w2_scale,
788 a1_scale=a1_scale,
789 a2_scale=a2_scale,
750 - block_shape=self.weight_block_size,
790 + block_shape=self.moe_block_shape,
791 swiglu_limit=getattr(layer, "swiglu_limit", None),
792 gemm1_alpha=getattr(layer, "swiglu_alpha", None),
793 gemm1_beta=getattr(layer, "swiglu_beta", None),
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+10
@@ -161,6 +161,16 @@ class QuarkOCP_MX(QuarkScheme):
161 weight_loader: Callable,
162 **kwargs,
163 ):
164 + if input_size_per_partition % OCP_MX_BLOCK_SIZE != 0:
165 + layer_name = getattr(layer, "prefix", "") or type(layer).__name__
166 + raise ValueError(
167 + f"OCP MX linear layer {layer_name!r} has an input size per "
168 + f"partition of {input_size_per_partition}, which must be "
169 + f"divisible by the OCP MX group size {OCP_MX_BLOCK_SIZE}. "
170 + "Choose a compatible tensor-parallel size or avoid "
171 + "tensor-parallel sharding for this layer."
172 + )
173 +
174 if self.dynamic_mxfp4_quant:
175 weight = ModelWeightParameter(
176 data=torch.empty(
vllm/model_executor/layers/quantization/utils/config_utils.py
+35
@@ -9,6 +9,41 @@ if TYPE_CHECKING:
9 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
10
11
12 +def get_quark_ocp_mx_group_size(
13 + quant_config: "QuantizationConfig | None",
14 + layer_name: str,
15 +) -> int | None:
16 + """Return the OCP MX group size for a quantized Quark linear layer."""
17 + if quant_config is None:
18 + return None
19 +
20 + from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
21 + from vllm.model_executor.layers.quantization.quark.utils import should_ignore_layer
22 +
23 + if not isinstance(quant_config, QuarkConfig):
24 + return None
25 +
26 + if should_ignore_layer(
27 + layer_name,
28 + ignore=quant_config.quant_config.get("exclude") or [],
29 + fused_mapping=quant_config.packed_modules_mapping,
30 + ):
31 + return None
32 +
33 + layer_quant_config = (
34 + quant_config.get_layer_quant_config_from_name(layer_name)
35 + or quant_config.quant_config.get("global_quant_config")
36 + or {}
37 + )
38 + weight_quant = layer_quant_config.get("weight")
39 + input_quant = layer_quant_config.get("input_tensors")
40 + if not quant_config._is_w_ocp_mx_a_x(weight_quant, input_quant):
41 + return None
42 +
43 + assert weight_quant is not None
44 + return int(weight_quant["group_size"])
45 +
46 +
47 def is_shared_expert_quant_fse_compatible(
48 quant_config: "QuantizationConfig | None",
49 expert_prefix: str,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+24 -3
@@ -232,7 +232,16 @@ def align_fp4_moe_weights_for_fi(
232
233 # Pad w13 and w2 along its intermediate dimension.
234 padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size // 2))
235 - padded_w13[:, : w13.shape[1], :] = w13
235 + if is_act_and_mul:
236 + # Keep the two logical projections independently aligned. Copying the
237 + # fused [gate, up] tensor contiguously would move the up projection into
238 + # the padded tail of gate when intermediate is rounded up.
239 + padded_w13[:, :intermediate, :] = w13[:, :intermediate, :]
240 + padded_w13[:, padded_intermediate : padded_intermediate + intermediate, :] = (
241 + w13[:, intermediate:, :]
242 + )
243 + else:
244 + padded_w13[:, : w13.shape[1], :] = w13
245
246 padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate // 2))
247 padded_w2[:, :, : w2.shape[2]] = w2
@@ -240,7 +249,13 @@ def align_fp4_moe_weights_for_fi(
249 padded_w13_scale = w13_scale.new_zeros(
250 (num_experts, padded_gate_up_dim, hidden_size // 16)
251 )
243 - padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale
252 + if is_act_and_mul:
253 + padded_w13_scale[:, :intermediate, :] = w13_scale[:, :intermediate, :]
254 + padded_w13_scale[
255 + :, padded_intermediate : padded_intermediate + intermediate, :
256 + ] = w13_scale[:, intermediate:, :]
257 + else:
258 + padded_w13_scale[:, : w13_scale.shape[1], :] = w13_scale
259
260 padded_w2_scale = w2_scale.new_zeros(
261 (num_experts, hidden_size, padded_intermediate // 16)
@@ -322,7 +337,13 @@ def align_moe_weights_for_fi(
337
338 # Pad w13 and w2 along its intermediate dimension.
339 padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size))
325 - padded_w13[:, : w13.shape[1], :] = w13
340 + if is_act_and_mul:
341 + padded_w13[:, :intermediate, :] = w13[:, :intermediate, :]
342 + padded_w13[:, padded_intermediate : padded_intermediate + intermediate, :] = (
343 + w13[:, intermediate:, :]
344 + )
345 + else:
346 + padded_w13[:, :intermediate, :] = w13
347
348 padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate))
349 padded_w2[:, :, :intermediate] = w2
vllm/model_executor/layers/vocab_parallel_embedding.py
+14 -3
@@ -233,6 +233,7 @@ class VocabParallelEmbedding(PluggableLayer):
233 quant_config: quant config for the layer
234 prefix: full name of the layer in the state dict
235 disable_tp: If true, tensor parallelism will be disabled for this layer.
236 + quant_method: Preselected quantization method for model-specific layers.
237 """ # noqa: E501
238
239 # --8<-- [end:vocab_parallel_embedding]
@@ -248,6 +249,7 @@ class VocabParallelEmbedding(PluggableLayer):
249 prefix: str = "",
250 *,
251 disable_tp: bool = False,
252 + quant_method: QuantizeMethodBase | None = None,
253 ):
254 super().__init__()
255
@@ -281,8 +283,9 @@ class VocabParallelEmbedding(PluggableLayer):
283 )
284 self.embedding_dim = embedding_dim
285
284 - quant_method = None
285 - if quant_config is not None:
286 + # Avoid overriding a preselected model-specific method with generic
287 + # config-based dispatch.
288 + if quant_method is None and quant_config is not None:
289 quant_method = quant_config.get_quant_method(self, prefix=prefix)
290 if quant_method is None:
291 quant_method = UnquantizedEmbeddingMethod()
@@ -501,8 +504,16 @@ class VocabParallelEmbedding(PluggableLayer):
504 output_parallel = self.quant_method.embedding(self, masked_input.long())
505 # Mask the output embedding.
506 if self.tp_size > 1:
507 + if output_parallel.dtype in (
508 + torch.float8_e4m3fn,
509 + torch.float8_e5m2,
510 + ):
511 + # Each vocab token has one owner, so FP8 bytes can use int8 SUM.
512 + comm_output = output_parallel.view(torch.int8)
513 + comm_output.masked_fill_(input_mask.unsqueeze(-1), 0)
514 + output = tensor_model_parallel_all_reduce(comm_output)
515 + return output.view(output_parallel.dtype)
516 output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
505 - # Reduce across all the model parallel GPUs.
517 return tensor_model_parallel_all_reduce(output_parallel)
518 return output_parallel
519
vllm/model_executor/model_loader/weight_utils.py
+7 -2
@@ -289,12 +289,17 @@ def get_quant_config(
289 # if hf_quant_config is None, we will try to get config from
290 # hf_overrides
291 hf_overrides = model_config.hf_overrides
292 + if callable(hf_overrides):
293 + # A callable hf_overrides is a config-to-config transform (e.g. the
294 + # one SpeculativeConfig installs on draft model configs); it cannot
295 + # carry quantization config entries, so treat it as no overrides.
296 + hf_overrides = {}
297 if not isinstance(hf_overrides, dict):
298 raise ValueError(
299 "hf_overrides must be a dict for get_quant_config "
300 "to get the quantization config from it."
301 )
297 - quantization_config_file = hf_overrides.get("quantization_config_file", None)
302 + quantization_config_file = hf_overrides.get("quantization_config_file")
303 if quantization_config_file is not None:
304 if hasattr(quant_cls, "from_config_file"):
305 return quant_cls.from_config_file(quantization_config_file)
@@ -304,7 +309,7 @@ def get_quant_config(
309 "but quant_cls.from_config_file is not implemented in "
310 f"{quant_cls}"
311 )
307 - quantization_config_json = hf_overrides.get("quantization_config_dict_json", None)
312 + quantization_config_json = hf_overrides.get("quantization_config_dict_json")
313 if quantization_config_json is not None:
314 if hasattr(quant_cls, "from_config_dict_json"):
315 return quant_cls.from_config_dict_json(quantization_config_json)
vllm/model_executor/models/config.py
+80
@@ -839,6 +839,83 @@ class Qwen3_5ForCausalLMConfig(Qwen3_5ForConditionalGenerationConfig):
839 rope_parameters.pop("mrope_interleaved", None)
840
841
842 +def _strip_qwen4_exp_mrope(model_config: "ModelConfig") -> None:
843 + configs = {
844 + id(config): config
845 + for config in (
846 + getattr(model_config, "hf_config", None),
847 + model_config.hf_text_config,
848 + )
849 + if config is not None
850 + }
851 + for config in configs.values():
852 + rope_parameters = getattr(config, "rope_parameters", None)
853 + if rope_parameters is not None:
854 + rope_parameters.pop("mrope_section", None)
855 + rope_parameters.pop("mrope_interleaved", None)
856 +
857 +
858 +class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConfig):
859 + """Apply the Qwen3.5 hybrid-cache contract to Qwen4Exp."""
860 +
861 + @staticmethod
862 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
863 + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config)
864 + text_config = vllm_config.model_config.hf_text_config
865 + if text_config.hc_count <= 1:
866 + raise ValueError("Qwen4Exp requires hc_count > 1")
867 + parallel_config = vllm_config.parallel_config
868 + uses_ple_or_qsa = bool(text_config.ple_layer_ids) or (
869 + getattr(text_config, "indexer_n_heads", None) is not None
870 + )
871 + if uses_ple_or_qsa and (
872 + parallel_config.enable_dbo or parallel_config.ubatch_size > 1
873 + ):
874 + raise NotImplementedError(
875 + "Qwen4Exp PLE/QSA does not support dual-batch overlap or microbatching"
876 + )
877 + # Checked again in Qwen4ExpModelState; rejecting it here keeps the
878 + # engine from loading weights first.
879 + if text_config.ple_layer_ids and parallel_config.pipeline_parallel_size > 1:
880 + raise NotImplementedError(
881 + "Qwen4Exp N-gram PLE embedding requires pipeline_parallel_size=1 "
882 + "because non-first pipeline ranks do not receive the raw input_ids "
883 + "it needs. Please run with PP=1."
884 + )
885 + multimodal_config = vllm_config.model_config.multimodal_config
886 + if multimodal_config is not None and multimodal_config.language_model_only:
887 + _strip_qwen4_exp_mrope(vllm_config.model_config)
888 + spec_config = vllm_config.speculative_config
889 + if spec_config is not None and spec_config.method not in {
890 + "mtp",
891 + "ngram",
892 + "ngram_gpu",
893 + }:
894 + raise NotImplementedError(
895 + "Qwen4Exp speculative decoding supports only its native MTP "
896 + "checkpoint and linear n-gram proposers"
897 + )
898 +
899 +
900 +class Qwen4ExpForCausalLMConfig(Qwen4ExpForConditionalGenerationConfig):
901 + @staticmethod
902 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
903 + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config)
904 +
905 + _strip_qwen4_exp_mrope(vllm_config.model_config)
906 +
907 +
908 +class Qwen4ExpMTPConfig(Qwen4ExpForConditionalGenerationConfig):
909 + """Preserve MRoPE for a VL target and use 1D RoPE for a text target."""
910 +
911 + @staticmethod
912 + def verify_and_update_config(vllm_config: "VllmConfig") -> None:
913 + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config)
914 + if hasattr(vllm_config.model_config.hf_config, "vision_config"):
915 + return
916 + _strip_qwen4_exp_mrope(vllm_config.model_config)
917 +
918 +
919 class ColQwen3_5Config(Qwen3_5ForConditionalGenerationConfig):
920 """Apply the attention contract declared by a ColQwen3.5 checkpoint."""
921
@@ -963,6 +1040,9 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
1040 "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
1041 "Qwen3_5MoeForCausalLM": Qwen3_5ForCausalLMConfig,
1042 "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig,
1043 + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig,
1044 + "Qwen4ExpForConditionalGeneration": (Qwen4ExpForConditionalGenerationConfig),
1045 + "Qwen4ExpMTP": Qwen4ExpMTPConfig,
1046 "UnlimitedOCRForCausalLM": UnlimitedOCRForCausalLMConfig,
1047 "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig,
1048 "XLMRobertaModel": JinaRobertaModelConfig,
vllm/model_executor/models/interfaces.py
+27 -1
@@ -49,13 +49,17 @@ if TYPE_CHECKING:
49 from vllm.lora.model_manager import LoRAModelManager
50 from vllm.model_executor.layers.fused_moe import MoERunner
51 from vllm.model_executor.layers.logits_processor import LogitsProcessor
52 - from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
52 + from vllm.model_executor.layers.mamba.mamba_utils import (
53 + MambaStateCopyFunc,
54 + MambaStateCopyFuncsByType,
55 + )
56 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
57 from vllm.model_executor.models.utils import WeightsMapper
58 from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem
59 from vllm.multimodal.registry import _ProcessorFactories
60 from vllm.sequence import IntermediateTensors
61 from vllm.tasks import ScoreType
62 + from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
63 from vllm.v1.worker.encoder_cudagraph_defs import (
64 EncoderCudaGraphCaptureInputs,
65 EncoderCudaGraphConfig,
@@ -970,6 +974,14 @@ class IsHybrid(Protocol):
974 """
975 ...
976
977 + @classmethod
978 + def get_mamba_state_copy_funcs(
979 + cls,
980 + mamba_types: set["MambaAttentionBackendEnum"],
981 + ) -> "MambaStateCopyFuncsByType":
982 + copy_funcs = cls.get_mamba_state_copy_func()
983 + return {mamba_type: copy_funcs for mamba_type in mamba_types}
984 +
985
986 @overload
987 def is_hybrid(model: type[object]) -> TypeIs[type[IsHybrid]]: ...
@@ -1131,6 +1143,20 @@ class SupportsMambaPrefixCaching(Protocol):
1143
1144 supports_mamba_prefix_caching: ClassVar[Literal[True]] = True
1145
1146 + @classmethod
1147 + def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]:
1148 + """Return copy functions for the model's Mamba states."""
1149 + ...
1150 +
1151 + @classmethod
1152 + def get_mamba_state_copy_funcs(
1153 + cls,
1154 + mamba_types: set["MambaAttentionBackendEnum"],
1155 + ) -> "MambaStateCopyFuncsByType":
1156 + """Map legacy copy functions to each requested Mamba backend."""
1157 + copy_funcs = cls.get_mamba_state_copy_func()
1158 + return {mamba_type: copy_funcs for mamba_type in mamba_types}
1159 +
1160
1161 @overload
1162 def supports_mamba_prefix_caching(
vllm/model_executor/models/qwen2_moe.py
+4 -2
@@ -80,15 +80,17 @@ class Qwen2MoeMLP(nn.Module):
80 reduce_results: bool = True,
81 expert_gate: torch.nn.Linear | None = None,
82 is_sequence_parallel: bool = False,
83 + disable_tp: bool = False,
84 prefix: str = "",
85 ) -> None:
86 super().__init__()
87 + disable_tp = disable_tp or is_sequence_parallel
88 self.gate_up_proj = MergedColumnParallelLinear(
89 hidden_size,
90 [intermediate_size] * 2,
91 bias=False,
92 quant_config=quant_config,
91 - disable_tp=is_sequence_parallel,
93 + disable_tp=disable_tp,
94 prefix=f"{prefix}.gate_up_proj",
95 )
96 self.down_proj = RowParallelLinear(
@@ -97,7 +99,7 @@ class Qwen2MoeMLP(nn.Module):
99 bias=False,
100 quant_config=quant_config,
101 reduce_results=reduce_results,
100 - disable_tp=is_sequence_parallel,
102 + disable_tp=disable_tp,
103 prefix=f"{prefix}.down_proj",
104 )
105 if hidden_act != "silu":
vllm/model_executor/models/qwen3_next.py
+62 -3
@@ -43,6 +43,9 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
43 MambaStateShapeCalculator,
44 )
45 from vllm.model_executor.layers.quantization import QuantizationConfig
46 +from vllm.model_executor.layers.quantization.utils.config_utils import (
47 + get_quark_ocp_mx_group_size,
48 +)
49 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
50 from vllm.model_executor.layers.vocab_parallel_embedding import (
51 ParallelLMHead,
@@ -90,6 +93,38 @@ def _should_use_sequence_parallel(vllm_config: VllmConfig) -> bool:
93 )
94
95
96 +def _should_replicate_misaligned_shared_expert(
97 + intermediate_size: int,
98 + tp_size: int,
99 + group_size: int | None,
100 + enable_expert_parallel: bool,
101 + is_sequence_parallel: bool,
102 +) -> bool:
103 + if intermediate_size <= 0 or group_size is None:
104 + return False
105 +
106 + partition_size, remainder = divmod(intermediate_size, tp_size)
107 + if remainder == 0 and partition_size % group_size == 0:
108 + return False
109 +
110 + if enable_expert_parallel or is_sequence_parallel:
111 + return True
112 +
113 + if remainder != 0:
114 + raise ValueError(
115 + f"Shared-expert intermediate size {intermediate_size} must be "
116 + f"divisible by tensor-parallel size {tp_size}."
117 + )
118 +
119 + raise ValueError(
120 + "The Quark OCP MX shared expert cannot be tensor-parallelized: "
121 + f"intermediate size {intermediate_size} with TP size {tp_size} "
122 + f"produces a partition of {partition_size}, which is not divisible by "
123 + f"the OCP MX group size {group_size}. Choose a compatible "
124 + "tensor-parallel size or enable expert parallelism."
125 + )
126 +
127 +
128 class Qwen3NextSparseMoeBlock(nn.Module):
129 def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
130 super().__init__()
@@ -106,6 +141,17 @@ class Qwen3NextSparseMoeBlock(nn.Module):
141
142 self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
143
144 + shared_expert_group_size = get_quark_ocp_mx_group_size(
145 + quant_config,
146 + f"{prefix}.shared_expert.down_proj",
147 + )
148 + self.replicate_shared_expert = _should_replicate_misaligned_shared_expert(
149 + config.shared_expert_intermediate_size,
150 + self.tp_size,
151 + shared_expert_group_size,
152 + parallel_config.enable_expert_parallel,
153 + self.is_sequence_parallel,
154 + )
155 if self.tp_size > config.num_experts:
156 raise ValueError(
157 f"Tensor parallel size {self.tp_size} is greater than "
@@ -138,7 +184,10 @@ class Qwen3NextSparseMoeBlock(nn.Module):
184 )
185
186 self.is_fused_shared_expert_enabled = False
141 - if config.shared_expert_intermediate_size > 0:
187 + if (
188 + config.shared_expert_intermediate_size > 0
189 + and not self.replicate_shared_expert
190 + ):
191 self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert(
192 quant_config,
193 prefix,
@@ -159,11 +208,14 @@ class Qwen3NextSparseMoeBlock(nn.Module):
208 reduce_results=False,
209 expert_gate=self.shared_expert_gate,
210 is_sequence_parallel=self.is_sequence_parallel,
211 + disable_tp=self.replicate_shared_expert,
212 prefix=f"{prefix}.shared_expert",
213 )
214
215 self.experts = FusedMoEFactory(
166 - shared_experts=self.shared_expert,
216 + shared_experts=(
217 + None if self.replicate_shared_expert else self.shared_expert
218 + ),
219 gate=self.gate,
220 num_experts=self.n_routed_experts,
221 top_k=config.num_experts_per_tok,
@@ -195,9 +247,16 @@ class Qwen3NextSparseMoeBlock(nn.Module):
247 if self.is_sequence_parallel and not already_sequence_parallel:
248 hidden_states = sequence_parallel_chunk(hidden_states)
249
250 + replicated_shared_output = (
251 + self.shared_expert(hidden_states)
252 + if self.replicate_shared_expert and self.shared_expert is not None
253 + else None
254 + )
255 final_hidden_states = self.experts(
256 hidden_states=hidden_states, router_logits=hidden_states
257 )
258 + if replicated_shared_output is not None:
259 + final_hidden_states += replicated_shared_output
260
261 if self.is_sequence_parallel and not already_sequence_parallel:
262 final_hidden_states = tensor_model_parallel_all_gather(
@@ -488,7 +547,7 @@ class Qwen3NextDecoderLayer(nn.Module):
547 self,
548 hidden_states: torch.Tensor,
549 residual: torch.Tensor | None,
491 - positions: torch.Tensor = None,
550 + positions: torch.Tensor,
551 **kwargs: object,
552 ):
553 full_num_tokens = positions.shape[-1]
vllm/model_executor/models/registry.py
+9
@@ -111,6 +111,10 @@ _TEXT_GENERATION_MODELS = {
111 "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),
112 "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"),
113 "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"),
114 + "Qwen4ExpForCausalLM": (
115 + "vllm.models.qwen4_exp",
116 + "Qwen4ExpForCausalLM",
117 + ),
118 "GlmForCausalLM": ("glm", "GlmForCausalLM"),
119 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
120 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
@@ -573,6 +577,10 @@ _MULTIMODAL_MODELS = {
577 "qwen3_5",
578 "Qwen3_5MoeForConditionalGeneration",
579 ),
580 + "Qwen4ExpForConditionalGeneration": (
581 + "vllm.models.qwen4_exp",
582 + "Qwen4ExpForConditionalGeneration",
583 + ),
584 "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),
585 "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"),
586 "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"),
@@ -659,6 +667,7 @@ _SPECULATIVE_DECODING_MODELS = {
667 "MedusaModel": ("medusa", "Medusa"),
668 "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"),
669 "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"),
670 + "Qwen4ExpMTP": ("vllm.models.qwen4_exp", "Qwen4ExpMTP"),
671 "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"),
672 "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"),
673 "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"),
vllm/models/qwen4_exp/__init__.py new
+61
@@ -0,0 +1,61 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen4Exp model package."""
4 +
5 +from typing import TYPE_CHECKING, Any
6 +
7 +from .common.hyperconnection import (
8 + GatedResidual,
9 + GroupedGemmaRMSNorm,
10 + HyperConnectionBase,
11 + HyperConnectionConfig,
12 +)
13 +
14 +if TYPE_CHECKING:
15 + from .nvidia.model import (
16 + Qwen4ExpForCausalLM,
17 + Qwen4ExpForConditionalGeneration,
18 + )
19 + from .nvidia.mtp import Qwen4ExpMTP
20 +
21 +
22 +def __getattr__(name: str) -> Any:
23 + if name in {
24 + "Qwen4ExpForCausalLM",
25 + "Qwen4ExpForConditionalGeneration",
26 + "Qwen4ExpMTP",
27 + }:
28 + from vllm.platforms import current_platform
29 +
30 + if current_platform.is_xpu() or current_platform.is_tpu():
31 + raise NotImplementedError("Qwen4Exp currently supports CUDA and ROCm only")
32 + if current_platform.is_rocm():
33 + from .amd.model import (
34 + Qwen4ExpForCausalLM,
35 + Qwen4ExpForConditionalGeneration,
36 + )
37 + from .amd.mtp import Qwen4ExpMTP
38 + else:
39 + from .nvidia.model import (
40 + Qwen4ExpForCausalLM,
41 + Qwen4ExpForConditionalGeneration,
42 + )
43 + from .nvidia.mtp import Qwen4ExpMTP
44 +
45 + return {
46 + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLM,
47 + "Qwen4ExpForConditionalGeneration": (Qwen4ExpForConditionalGeneration),
48 + "Qwen4ExpMTP": Qwen4ExpMTP,
49 + }[name]
50 + raise AttributeError(name)
51 +
52 +
53 +__all__ = [
54 + "GatedResidual",
55 + "GroupedGemmaRMSNorm",
56 + "HyperConnectionBase",
57 + "HyperConnectionConfig",
58 + "Qwen4ExpForCausalLM",
59 + "Qwen4ExpForConditionalGeneration",
60 + "Qwen4ExpMTP",
61 +]
vllm/models/qwen4_exp/amd/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/qwen4_exp/amd/hyperconnection.py new
+205
@@ -0,0 +1,205 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HyperConnection (Gated Residual) utilities for the AMD model variant.
4 +
5 +Implements the HyperConnection residual scheme proposed in
6 +"HyperConnections" (https://arxiv.org/abs/2409.19606). This AMD variant
7 +delays each HC combine to the following HC mix boundary. HC glue kernels,
8 +including fused combine+RMSNorm, live in ``ops/hc.py``; projections remain
9 +standard vLLM Linear modules.
10 +
11 +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner
12 +(HC outer, HS inner — checkpoint-native layout).
13 +
14 +Typical usage inside a transformer decoder layer::
15 +
16 + self.attn_hc = GatedResidual(hc_config)
17 +
18 + hidden_states, block_input, injection = self.attn_hc.mix(hidden_states)
19 + attention_output = attention(block_input)
20 + hidden_states, block_input, injection = self.mlp_hc.combine_and_mix(
21 + hidden_states, attention_output, injection
22 + )
23 +"""
24 +
25 +import torch
26 +from torch import nn
27 +
28 +from vllm.model_executor.layers.linear import (
29 + MergedColumnParallelLinear,
30 + ReplicatedLinear,
31 +)
32 +from vllm.model_executor.models.utils import maybe_prefix
33 +
34 +from ..common.hyperconnection import (
35 + GroupedGemmaRMSNorm,
36 + HyperConnectionConfig,
37 +)
38 +from .ops.hc import (
39 + grouped_gemma_rmsnorm,
40 + hc_combine,
41 + hc_combine_norm,
42 + hc_gate_mix,
43 + hc_silu,
44 +)
45 +
46 +
47 +# ---------------------------------------------------------------------------
48 +# Gated-residual variant
49 +# ---------------------------------------------------------------------------
50 +class GatedResidual(nn.Module):
51 + """Gated HyperConnection with learnable low-rank mixing and injection.
52 +
53 + ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged
54 + low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean
55 + over the HC streams). When passed a pending block output and an injection,
56 + it fuses their residual combine with the RMSNorm. Final mixers use
57 + ``use_combine=False`` and do not produce a new injection.
58 +
59 + Weights: the norm owns the grouped GemmaRMSNorm affine; the projections
60 + are vLLM Linear modules (merged replicated linear for down+inject), so
61 + GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the
62 + standard quant_method mechanism.
63 + """
64 +
65 + def __init__(
66 + self,
67 + config: HyperConnectionConfig,
68 + use_combine: bool = True,
69 + prefix: str = "",
70 + ) -> None:
71 + super().__init__()
72 + self.config = config
73 + self.lora_rank = config.hc_lowrank
74 + self.hc_count = config.hc_count
75 + self.hidden_size = config.hidden_size
76 + self.use_combine = use_combine
77 +
78 + norm_size = (
79 + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size
80 + )
81 + group_size = config.hidden_size if config.hc_per_branch_norm else None
82 + # Normalize each H-sized HC stream independently while retaining a
83 + # separate affine weight for every element of the HC*H layout.
84 + self.hc_norm = GroupedGemmaRMSNorm(
85 + norm_size,
86 + eps=config.rms_norm_eps,
87 + group_size=group_size,
88 + dtype=config.params_dtype,
89 + )
90 +
91 + # -- vLLM Linear weights --------------------------------------------
92 + # The merged skinny-GEMM shape is physically padded to 16 rows for
93 + # alignment and efficient backend dispatch.
94 + self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0
95 + if use_combine:
96 + self.input_mix_weight_down_block_inject = MergedColumnParallelLinear(
97 + self.hyper_hidden_size,
98 + [self.lora_rank, self.hc_count]
99 + + ([self.pad_size] if self.pad_size else []),
100 + bias=False,
101 + params_dtype=config.params_dtype,
102 + quant_config=None,
103 + prefix=maybe_prefix(prefix, "input_mix_weight_down_block_inject"),
104 + return_bias=False,
105 + disable_tp=True,
106 + )
107 + else:
108 + self.input_mix_weight_down = ReplicatedLinear(
109 + self.hyper_hidden_size,
110 + self.lora_rank,
111 + bias=False,
112 + params_dtype=config.params_dtype,
113 + quant_config=None,
114 + prefix=maybe_prefix(prefix, "input_mix_weight_down"),
115 + return_bias=False,
116 + )
117 + self.input_mix_weight_up = ReplicatedLinear(
118 + self.lora_rank,
119 + self.hyper_hidden_size,
120 + bias=False,
121 + params_dtype=config.params_dtype,
122 + quant_config=None,
123 + prefix=maybe_prefix(prefix, "input_mix_weight_up"),
124 + return_bias=False,
125 + )
126 +
127 + def mix(
128 + self, hidden_states: torch.Tensor
129 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
130 + xn = grouped_gemma_rmsnorm(
131 + hidden_states,
132 + self.hc_norm.weight,
133 + self.config.rms_norm_eps,
134 + self.hc_count,
135 + )
136 +
137 + if self.use_combine:
138 + # produce injection logits for combine
139 + split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
140 + down_and_injection = self.input_mix_weight_down_block_inject(xn)
141 + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
142 + else:
143 + lora = self.input_mix_weight_down(xn)
144 + injection = None
145 +
146 + lora = hc_silu(lora, self.hc_count)
147 + gate = self.input_mix_weight_up(lora) # [M, D]
148 + block_input = hc_gate_mix(xn, gate, self.hc_count)
149 +
150 + return hidden_states, block_input, injection
151 +
152 + def combine_and_mix(
153 + self,
154 + hidden_states: torch.Tensor,
155 + prev_block_output: torch.Tensor,
156 + prev_injection: torch.Tensor,
157 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
158 + """Consume a pending combine, then prepare the next block input.
159 +
160 + ``hidden_states`` is the multi-stream state from before the pending
161 + block's mix. Its combine with ``block_output`` is fused with this
162 + module's input RMSNorm.
163 + """
164 + hidden_states, xn = hc_combine_norm(
165 + hidden_states,
166 + prev_block_output,
167 + prev_injection,
168 + self.hc_norm.weight,
169 + self.config.rms_norm_eps,
170 + self.hc_count,
171 + )
172 +
173 + if self.use_combine:
174 + # produce injection logits for combine
175 + split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
176 + down_and_injection = self.input_mix_weight_down_block_inject(xn)
177 + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
178 + else:
179 + lora = self.input_mix_weight_down(xn)
180 + injection = None
181 +
182 + lora = hc_silu(lora, self.hc_count)
183 + gate = self.input_mix_weight_up(lora) # [M, D]
184 + block_input = hc_gate_mix(xn, gate, self.hc_count)
185 +
186 + return hidden_states, block_input, injection
187 +
188 + def combine(
189 + self,
190 + hidden_states: torch.Tensor,
191 + block_output: torch.Tensor,
192 + injection: torch.Tensor,
193 + ) -> torch.Tensor:
194 + return hc_combine(hidden_states, block_output, injection, self.hc_count)
195 +
196 + @property
197 + def hyper_hidden_size(self) -> int:
198 + return self.hc_count * self.hidden_size
199 +
200 +
201 +__all__ = [
202 + "GatedResidual",
203 + "GroupedGemmaRMSNorm",
204 + "HyperConnectionConfig",
205 +]
vllm/models/qwen4_exp/amd/indexer_qsa.py new
+321
@@ -0,0 +1,321 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen4Exp weight-free QSA indexer."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import cast
8 +
9 +import torch
10 +from torch import nn
11 +
12 +from vllm.config import VllmConfig
13 +from vllm.forward_context import get_forward_context
14 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
15 +from vllm.model_executor.layers.linear import ReplicatedLinear
16 +from vllm.model_executor.layers.quantization import QuantizationConfig
17 +from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope
18 +from vllm.transformers_utils.configs.qwen4_exp import (
19 + Qwen4ExpTextConfig,
20 +)
21 +
22 +from ..common.qsa_cache import (
23 + QSACompressedKeyCache,
24 + QSAForwardMetadata,
25 + QSAKeyStateCache,
26 + canonical_qsa_rope_positions,
27 +)
28 +
29 +
30 +def apply_qsa_rope(
31 + rotary_emb: nn.Module,
32 + positions: torch.Tensor,
33 + tensor: torch.Tensor,
34 +) -> torch.Tensor:
35 + """Apply the main attention's exact 1D/MRoPE composition to QSA heads."""
36 +
37 + num_tokens, _, head_dim = tensor.shape
38 + rotary_dim = rotary_emb.rotary_dim
39 + cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001
40 + cos_sin = cache[positions]
41 + cos, sin = cos_sin.chunk(2, dim=-1)
42 + if positions.ndim == 2:
43 + shape = tensor.shape
44 + tensor, _ = triton_mrope(
45 + tensor.reshape(num_tokens, -1),
46 + tensor.new_empty((num_tokens, head_dim)),
47 + cos,
48 + sin,
49 + rotary_emb.mrope_section,
50 + head_dim,
51 + rotary_dim,
52 + rotary_emb.mrope_interleaved,
53 + rotary_emb.is_neox_style,
54 + )
55 + return tensor.reshape(shape)
56 +
57 + rotated = rotary_emb.apply_rotary_emb(
58 + tensor[..., :rotary_dim],
59 + cos,
60 + sin,
61 + )
62 + return torch.cat((rotated, tensor[..., rotary_dim:]), dim=-1)
63 +
64 +
65 +def apply_qsa_rmsnorm(
66 + norm: GemmaRMSNorm,
67 + tensor: torch.Tensor,
68 +) -> torch.Tensor:
69 + """Use vLLM's portable RMSNorm implementation on ROCm."""
70 +
71 + return cast(torch.Tensor, norm(tensor))
72 +
73 +
74 +class QSAIndexer(nn.Module):
75 + """Replicated Q/K projection plus paged, weight-free QSA selection.
76 +
77 + ``prefix`` must be the checkpoint's indexer prefix, normally
78 + ``model.layers.N.self_attn.indexer``. Consequently the trainable names are
79 + ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix.
80 + """
81 +
82 + def __init__(
83 + self,
84 + *,
85 + vllm_config: VllmConfig,
86 + config: Qwen4ExpTextConfig,
87 + layer_id: int,
88 + rotary_emb: nn.Module,
89 + quant_config: QuantizationConfig | None = None,
90 + prefix: str = "",
91 + ) -> None:
92 + super().__init__()
93 + if vllm_config.cache_config is None:
94 + raise ValueError("QSA requires a paged KV cache")
95 + if vllm_config.model_config.dtype != torch.bfloat16:
96 + raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
97 +
98 + self.layer_id = int(layer_id)
99 + self.index_n_heads = int(config.indexer_n_heads)
100 + self.index_kv_heads = int(config.indexer_kv_heads)
101 + self.index_head_dim = int(config.indexer_head_dim)
102 + self.token_topk = int(config.indexer_budget)
103 + self.compress_ratio = int(config.indexer_compress_ratio)
104 + self.rotary_emb = rotary_emb
105 + self.prefix = prefix
106 + # MTP step 0 selects the target-aligned rows; later steps reuse them
107 + # while continuing to update the QSA side cache.
108 + self.skip_topk = False
109 +
110 + self.index_qk_proj = ReplicatedLinear(
111 + int(config.hidden_size),
112 + (self.index_n_heads + self.index_kv_heads) * self.index_head_dim,
113 + bias=False,
114 + quant_config=quant_config,
115 + prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj",
116 + )
117 + self.q_layernorm = GemmaRMSNorm(
118 + self.index_head_dim,
119 + eps=float(getattr(config, "rms_norm_eps", 1e-6)),
120 + )
121 + self.k_layernorm = GemmaRMSNorm(
122 + self.index_head_dim,
123 + eps=float(getattr(config, "rms_norm_eps", 1e-6)),
124 + )
125 +
126 + cache_config = vllm_config.cache_config
127 + cache_prefix = f"{prefix}." if prefix else ""
128 + self.raw_key_cache = QSAKeyStateCache(
129 + head_size=self.index_head_dim,
130 + dtype=torch.bfloat16,
131 + cache_rope_positions=vllm_config.model_config.uses_mrope,
132 + prefix=f"{cache_prefix}raw_key_cache",
133 + cache_config=cache_config,
134 + compress_ratio=self.compress_ratio,
135 + vllm_config=vllm_config,
136 + )
137 + self.compressed_key_cache = QSACompressedKeyCache(
138 + head_size=self.index_head_dim,
139 + dtype=torch.bfloat16,
140 + compress_ratio=self.compress_ratio,
141 + prefix=f"{cache_prefix}compressed_key_cache",
142 + cache_config=cache_config,
143 + vllm_config=vllm_config,
144 + )
145 +
146 + @property
147 + def output_width(self) -> int:
148 + return self.token_topk + self.compress_ratio - 1
149 +
150 + def project_qk(
151 + self,
152 + hidden_states: torch.Tensor,
153 + positions: torch.Tensor,
154 + ) -> tuple[torch.Tensor, torch.Tensor]:
155 + """Project replicated Q/K, normalize+rotate Q, and preserve raw K."""
156 +
157 + qk, _ = self.index_qk_proj(hidden_states)
158 + q_raw, token_k = qk.split(
159 + (
160 + self.index_n_heads * self.index_head_dim,
161 + self.index_kv_heads * self.index_head_dim,
162 + ),
163 + dim=-1,
164 + )
165 + q = q_raw.reshape(-1, self.index_n_heads, self.index_head_dim)
166 + q = apply_qsa_rmsnorm(
167 + self.q_layernorm,
168 + q.reshape(-1, self.index_head_dim),
169 + ).reshape_as(q)
170 + q = apply_qsa_rope(self.rotary_emb, positions, q)
171 + return q, token_k.reshape(-1, 1, self.index_head_dim)
172 +
173 + def normalize_compressed_keys(
174 + self,
175 + compressed_keys: torch.Tensor,
176 + first_rope_positions: torch.Tensor,
177 + ) -> torch.Tensor:
178 + """Normalize pooled K and apply the first token's exact group position."""
179 +
180 + keys = compressed_keys.reshape(-1, self.index_head_dim)
181 + keys = apply_qsa_rmsnorm(self.k_layernorm, keys).reshape(
182 + -1, 1, self.index_head_dim
183 + )
184 + if getattr(self.rotary_emb, "mrope_section", None):
185 + positions = first_rope_positions.transpose(0, 1)
186 + else:
187 + positions = first_rope_positions[:, 0]
188 + return apply_qsa_rope(self.rotary_emb, positions, keys)
189 +
190 + def _metadata(
191 + self,
192 + ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None:
193 + metadata = get_forward_context().attn_metadata
194 + if isinstance(metadata, list):
195 + metadata = metadata[0]
196 + if not isinstance(metadata, dict):
197 + return None
198 + raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix])
199 + compressed = cast(
200 + QSAForwardMetadata, metadata[self.compressed_key_cache.prefix]
201 + )
202 + if raw.num_actual_tokens != compressed.num_actual_tokens:
203 + raise RuntimeError("QSA side-cache metadata token counts disagree")
204 + if not raw.logical_positions.is_cuda and (
205 + not torch.equal(raw.logical_positions, compressed.logical_positions)
206 + ):
207 + raise RuntimeError("QSA side-cache metadata positions disagree")
208 + return raw, compressed
209 +
210 + def _update_and_compress(
211 + self,
212 + token_k: torch.Tensor,
213 + positions: torch.Tensor,
214 + raw_metadata: QSAForwardMetadata,
215 + compressed_metadata: QSAForwardMetadata,
216 + ) -> None:
217 + num_tokens = raw_metadata.num_actual_tokens
218 + raw_key_cache = self.raw_key_cache.key_cache
219 + rope_position_cache = self.raw_key_cache.rope_position_cache
220 + from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows
221 +
222 + if rope_position_cache is None:
223 + position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand(
224 + -1, 1, 3
225 + )
226 + else:
227 + position_rows = canonical_qsa_rope_positions(positions)[:num_tokens].to(
228 + device=raw_key_cache.device
229 + )
230 + pooled, first_positions = qsa_compress_groups_with_ratio(
231 + token_k[:num_tokens],
232 + position_rows,
233 + raw_key_cache,
234 + raw_metadata.block_table,
235 + raw_metadata.token_to_req,
236 + raw_metadata.query_start_loc,
237 + raw_metadata.logical_positions,
238 + compressed_metadata.slot_mapping,
239 + self.compress_ratio,
240 + rope_position_cache,
241 + )
242 + normalized = self.normalize_compressed_keys(pooled, first_positions)
243 + qsa_store_cache_rows(
244 + self.compressed_key_cache.kv_cache,
245 + compressed_metadata.slot_mapping,
246 + normalized,
247 + )
248 + qsa_store_cache_rows(
249 + raw_key_cache,
250 + raw_metadata.slot_mapping,
251 + token_k[:num_tokens],
252 + )
253 + if rope_position_cache is not None:
254 + qsa_store_cache_rows(
255 + rope_position_cache,
256 + raw_metadata.slot_mapping,
257 + position_rows,
258 + )
259 +
260 + def _select(
261 + self,
262 + q: torch.Tensor,
263 + metadata: QSAForwardMetadata,
264 + out: torch.Tensor | None,
265 + ) -> torch.Tensor:
266 + from .ops.qsa import qsa_select_paged_tokens
267 +
268 + return qsa_select_paged_tokens(
269 + q,
270 + self.compressed_key_cache.kv_cache,
271 + metadata.block_table,
272 + metadata.token_to_req,
273 + metadata.logical_positions,
274 + metadata.seq_lens,
275 + self.token_topk,
276 + self.compress_ratio,
277 + out,
278 + )
279 +
280 + def forward(
281 + self,
282 + hidden_states: torch.Tensor,
283 + positions: torch.Tensor,
284 + out: torch.Tensor | None = None,
285 + ) -> torch.Tensor:
286 + """Return fixed-width request-relative token indices padded with ``-1``."""
287 +
288 + metadata = self._metadata()
289 + if metadata is None:
290 + # Preserve step-0 indices when later MTP steps reuse the buffer.
291 + if self.skip_topk and out is not None:
292 + return out
293 + result = torch.full(
294 + (hidden_states.shape[0], self.output_width),
295 + -1,
296 + dtype=torch.int32,
297 + device=hidden_states.device,
298 + )
299 + if out is not None:
300 + out.copy_(result)
301 + return out
302 + return result
303 + raw_metadata, compressed_metadata = metadata
304 + num_tokens = raw_metadata.num_actual_tokens
305 + q, token_k = self.project_qk(
306 + hidden_states[:num_tokens], positions[..., :num_tokens]
307 + )
308 + self._update_and_compress(
309 + token_k,
310 + positions[..., :num_tokens],
311 + raw_metadata,
312 + compressed_metadata,
313 + )
314 + if self.skip_topk:
315 + if out is None:
316 + raise RuntimeError("QSA top-k reuse requires an output buffer")
317 + return out
318 + return self._select(q, compressed_metadata, out)
319 +
320 +
321 +__all__ = ["QSAIndexer", "apply_qsa_rope"]
vllm/models/qwen4_exp/amd/low_latency_gemm.py new
+15
@@ -0,0 +1,15 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen4Exp low-latency GEMM hook for AMD ROCm."""
4 +
5 +import torch
6 +from torch import nn
7 +
8 +
9 +def enable_qwen4_exp_low_latency_gemm(
10 + module: nn.Module,
11 + dtype: torch.dtype,
12 +) -> None:
13 + """Keep the standard vLLM linear methods on AMD ROCm."""
14 +
15 + del module, dtype
vllm/models/qwen4_exp/amd/model.py new
+1065
@@ -0,0 +1,1065 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Qwen4Exp model."""
4 +
5 +from collections.abc import Iterable
6 +from itertools import islice
7 +
8 +import torch
9 +from torch import nn
10 +
11 +from vllm.compilation.decorators import support_torch_compile
12 +from vllm.config import VllmConfig
13 +from vllm.distributed import get_pp_group
14 +from vllm.model_executor.layers.fused_moe.utils import (
15 + is_model_fused_shared_expert_compatible,
16 +)
17 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
18 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
19 + QwenGatedDeltaNetAttention,
20 +)
21 +from vllm.model_executor.layers.mamba.mamba_utils import (
22 + MambaStateCopyFunc,
23 + MambaStateCopyFuncCalculator,
24 + MambaStateCopyFuncsByType,
25 + MambaStateDtypeCalculator,
26 + MambaStateShapeCalculator,
27 +)
28 +from vllm.model_executor.layers.quantization import QuantizationConfig
29 +from vllm.model_executor.layers.vocab_parallel_embedding import (
30 + ParallelLMHead,
31 + VocabParallelEmbedding,
32 +)
33 +from vllm.model_executor.models.interfaces import (
34 + HasInnerState,
35 + IsHybrid,
36 + MixtureOfExperts,
37 + MultiModalEmbeddings,
38 + SupportsLoRA,
39 + SupportsMRoPE,
40 + SupportsPP,
41 + _require_is_multimodal,
42 +)
43 +from vllm.model_executor.models.qwen3_5 import (
44 + Qwen3_5ForConditionalGeneration,
45 + Qwen3_5Model,
46 +)
47 +from vllm.model_executor.models.qwen3_next import (
48 + Qwen3NextAttention,
49 + Qwen3NextMLP,
50 + Qwen3NextSparseMoeBlock,
51 +)
52 +from vllm.model_executor.models.qwen3_vl import (
53 + Qwen3_VisionTransformer,
54 + Qwen3VLDummyInputsBuilder,
55 + Qwen3VLMultiModalProcessor,
56 + Qwen3VLProcessingInfo,
57 +)
58 +from vllm.model_executor.models.utils import (
59 + AutoWeightsLoader,
60 + StageMissingLayer,
61 + WeightsMapper,
62 + _merge_multimodal_embeddings,
63 + extract_layer_index,
64 + make_empty_intermediate_tensors_factory,
65 + make_layers,
66 + maybe_fuse_shared_experts,
67 + maybe_prefix,
68 +)
69 +from vllm.multimodal import MULTIMODAL_REGISTRY
70 +from vllm.multimodal.inputs import MultiModalFeatureSpec
71 +from vllm.sequence import IntermediateTensors
72 +from vllm.tokenizers.registry import cached_tokenizer_from_config
73 +from vllm.transformers_utils.configs.qwen4_exp import (
74 + Qwen4ExpTextConfig,
75 +)
76 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
77 +from vllm.v1.kv_cache_interface import MambaSpec
78 +
79 +from ..config import Qwen4ExpConfig
80 +from .hyperconnection import GatedResidual, HyperConnectionConfig
81 +from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
82 +from .ple_layer import Qwen4ExpPLELayer
83 +from .qsa import Qwen4ExpQSAAttention
84 +
85 +
86 +def without_modelopt_fp4(
87 + quant_config: QuantizationConfig | None,
88 +) -> QuantizationConfig | None:
89 + """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4."""
90 +
91 + if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
92 + return None
93 + return quant_config
94 +
95 +
96 +def _remap_qsa_cache_scale_name(
97 + name: str,
98 + qsa_layer_ids: frozenset[int],
99 +) -> str:
100 + """Map serialized main-cache scales onto the merged QSA owner.
101 +
102 + Regular attention keeps cache scales below its ``attn`` child. QSA owns
103 + that cache directly, so only QSA layers need the final path component
104 + moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers.
105 + """
106 +
107 + scale_suffixes = {
108 + "k_proj.k_scale": "_k_scale",
109 + "k_proj.output_scale": "_k_scale",
110 + "attn.k_scale": "_k_scale",
111 + "attn._k_scale": "_k_scale",
112 + "k_scale": "_k_scale",
113 + "_k_scale": "_k_scale",
114 + "v_proj.v_scale": "_v_scale",
115 + "v_proj.output_scale": "_v_scale",
116 + "attn.v_scale": "_v_scale",
117 + "attn._v_scale": "_v_scale",
118 + "v_scale": "_v_scale",
119 + "_v_scale": "_v_scale",
120 + }
121 + for layer_id in qsa_layer_ids:
122 + marker = f"layers.{layer_id}.self_attn."
123 + marker_start = name.find(marker)
124 + if marker_start < 0 or (marker_start > 0 and name[marker_start - 1] != "."):
125 + continue
126 + suffix = name[marker_start + len(marker) :]
127 + mapped_suffix = scale_suffixes.get(suffix)
128 + if mapped_suffix is not None:
129 + return f"{name[: marker_start + len(marker)]}{mapped_suffix}"
130 + return name
131 +
132 +
133 +_QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [
134 + ".bias",
135 + "_bias",
136 + ".k_scale",
137 + "_k_scale",
138 + ".v_scale",
139 + "_v_scale",
140 + "_weight_scale",
141 + "_input_scale",
142 +]
143 +
144 +# The checkpoint keeps down and injection projections separate; runtime packs
145 +# them into adjacent logical shards of one MergedColumnParallelLinear.
146 +_HC_WEIGHTS_MAPPER = WeightsMapper(
147 + orig_to_new_stacked={
148 + "hyper_connection.input_mix_weight_down.weight": (
149 + "hyper_connection.input_mix_weight_down_block_inject.weight",
150 + 0,
151 + ),
152 + "hyper_connection.block_inject_weight.weight": (
153 + "hyper_connection.input_mix_weight_down_block_inject.weight",
154 + 1,
155 + ),
156 + }
157 +)
158 +
159 +
160 +class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock):
161 + """Qwen3Next MoE with Qwen4Exp HC validation."""
162 +
163 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
164 + parallel_config = vllm_config.parallel_config
165 + if parallel_config.use_sequence_parallel_moe:
166 + raise NotImplementedError(
167 + "Qwen4Exp HC does not support sequence-parallel MoE"
168 + )
169 + super().__init__(vllm_config=vllm_config, prefix=prefix)
170 + config = vllm_config.model_config.hf_text_config
171 + self.n_shared_experts = int(config.shared_expert_intermediate_size > 0)
172 +
173 +
174 +class Qwen4ExpDecoderLayer(nn.Module):
175 + def __init__(
176 + self,
177 + vllm_config: VllmConfig,
178 + layer_type: str,
179 + prefix: str = "",
180 + ) -> None:
181 + super().__init__()
182 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
183 + model_config = vllm_config.model_config
184 + cache_config = vllm_config.cache_config
185 + quant_config = vllm_config.quant_config
186 +
187 + self.config = config
188 + self.layer_type = layer_type
189 + self.layer_idx = extract_layer_index(prefix)
190 + if vllm_config.parallel_config.use_sequence_parallel_moe:
191 + raise NotImplementedError(
192 + "Qwen4Exp HC does not support sequence-parallel MoE"
193 + )
194 + self.ple: Qwen4ExpPLELayer | None = None
195 + ple_layer_ids = config.ple_layer_ids
196 + if (self.layer_idx + 1) in ple_layer_ids:
197 + ple_layer_ids_sorted = sorted(set(ple_layer_ids))
198 + ple_dense_layer_id_map = {
199 + abs_id: idx for idx, abs_id in enumerate(ple_layer_ids_sorted)
200 + }
201 + ple_dense_layer_id = ple_dense_layer_id_map[self.layer_idx + 1]
202 + self.ple = Qwen4ExpPLELayer(
203 + config,
204 + vllm_config=vllm_config,
205 + layer_idx=self.layer_idx,
206 + ple_dense_layer_id=ple_dense_layer_id,
207 + prefix=f"{prefix}.ple",
208 + )
209 +
210 + if layer_type == "linear_attention":
211 + self.linear_attn = QwenGatedDeltaNetAttention(
212 + config,
213 + vllm_config=vllm_config,
214 + prefix=f"{prefix}.linear_attn",
215 + gqa_interleaved_layout=False,
216 + )
217 + elif layer_type == "full_attention":
218 + use_qsa = getattr(config, "indexer_n_heads", None) is not None
219 + if not use_qsa:
220 + self.self_attn = Qwen3NextAttention(
221 + config,
222 + model_config=model_config,
223 + cache_config=cache_config,
224 + quant_config=quant_config,
225 + prefix=f"{prefix}.self_attn",
226 + )
227 + else:
228 + self.self_attn = Qwen4ExpQSAAttention(
229 + vllm_config=vllm_config,
230 + config=config,
231 + layer_id=self.layer_idx,
232 + quant_config=quant_config,
233 + prefix=f"{prefix}.self_attn",
234 + )
235 + else:
236 + raise ValueError(f"Invalid layer_type {layer_type}")
237 +
238 + mlp_only_layers = getattr(config, "mlp_only_layers", [])
239 + num_experts = getattr(config, "num_experts", 0) or 0
240 + absolute_layer_id = self.layer_idx + 1
241 + is_moe_layer = self.layer_idx not in mlp_only_layers and (
242 + num_experts > 0 and absolute_layer_id % config.decoder_sparse_step == 0
243 + )
244 + if is_moe_layer:
245 + self.mlp = Qwen4ExpSparseMoeBlock(
246 + vllm_config=vllm_config, prefix=f"{prefix}.mlp"
247 + )
248 + else:
249 + self.mlp = Qwen3NextMLP(
250 + hidden_size=config.hidden_size,
251 + intermediate_size=config.intermediate_size,
252 + hidden_act=config.hidden_act,
253 + quant_config=quant_config,
254 + prefix=f"{prefix}.mlp",
255 + )
256 +
257 + hc_config = HyperConnectionConfig(
258 + hc_count=config.hc_count,
259 + hidden_size=config.hidden_size,
260 + params_dtype=torch.bfloat16,
261 + hc_lowrank=config.hc_lowrank,
262 + rms_norm_eps=config.rms_norm_eps,
263 + hc_per_branch_norm=True,
264 + )
265 + self.attn_hyper_connection = GatedResidual(
266 + hc_config,
267 + prefix=maybe_prefix(prefix, "attn_hyper_connection"),
268 + )
269 + self.mlp_hyper_connection = GatedResidual(
270 + hc_config,
271 + prefix=maybe_prefix(prefix, "mlp_hyper_connection"),
272 + )
273 +
274 + def forward(
275 + self,
276 + hidden_states: torch.Tensor,
277 + prev_block_output: torch.Tensor | None,
278 + prev_injection: torch.Tensor | None,
279 + positions: torch.Tensor,
280 + *,
281 + input_ids: torch.Tensor | None,
282 + query_start_loc: torch.Tensor | None,
283 + ngram_context: torch.Tensor | None,
284 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
285 + attn_hc = self.attn_hyper_connection
286 + if self.ple is not None:
287 + # PLE adds directly to the multi-stream state, so pending HC state
288 + # must be materialized before the addition.
289 + if prev_block_output is not None and prev_injection is not None:
290 + hidden_states = attn_hc.combine(
291 + hidden_states, prev_block_output, prev_injection
292 + )
293 + prev_block_output = prev_injection = None
294 +
295 + if input_ids is None or query_start_loc is None or ngram_context is None:
296 + raise RuntimeError("PLE inputs were not prepared")
297 + hidden_states = hidden_states + self.ple(
298 + hidden_states,
299 + input_ids,
300 + query_start_loc,
301 + ngram_context,
302 + )
303 +
304 + # Fuse a pending combine with this HC module's mix when possible.
305 + if prev_block_output is not None and prev_injection is not None:
306 + hidden_states, block_input, injection = attn_hc.combine_and_mix(
307 + hidden_states, prev_block_output, prev_injection
308 + )
309 + else:
310 + hidden_states, block_input, injection = attn_hc.mix(hidden_states)
311 +
312 + if self.layer_type == "linear_attention":
313 + attn_out = self.linear_attn(hidden_states=block_input)
314 + elif self.layer_type == "full_attention":
315 + attn_out = self.self_attn(
316 + hidden_states=block_input,
317 + positions=positions,
318 + )
319 + else:
320 + raise ValueError("Invalid layer_type")
321 +
322 + mlp_hc = self.mlp_hyper_connection
323 + hidden_states, block_input, injection = mlp_hc.combine_and_mix(
324 + hidden_states, attn_out, injection
325 + )
326 + mlp_out = self.mlp(block_input)
327 + return hidden_states, mlp_out, injection
328 +
329 +
330 +class Qwen4ExpMixtureOfExperts(MixtureOfExperts):
331 + """Expose Qwen4Exp routed experts through vLLM's EPLB protocol."""
332 +
333 + def set_moe_parameters(self, layers: Iterable[nn.Module]) -> None:
334 + self.moe_layers = []
335 + self.moe_mlp_layers = []
336 + example_moe = None
337 + for layer in layers:
338 + if isinstance(layer, Qwen4ExpDecoderLayer) and isinstance(
339 + layer.mlp, Qwen4ExpSparseMoeBlock
340 + ):
341 + example_moe = layer.mlp
342 + self.moe_mlp_layers.append(layer.mlp)
343 + self.moe_layers.append(layer.mlp.experts)
344 +
345 + self.num_moe_layers = len(self.moe_layers)
346 + if example_moe is None:
347 + self.num_expert_groups = 0
348 + self.num_shared_experts = 0
349 + self.num_logical_experts = 0
350 + self.num_physical_experts = 0
351 + self.num_local_physical_experts = 0
352 + self.num_routed_experts = 0
353 + self.num_redundant_experts = 0
354 + return
355 +
356 + self.num_expert_groups = 1
357 + self.num_shared_experts = example_moe.n_shared_experts
358 + self.num_logical_experts = example_moe.n_logical_experts
359 + self.num_physical_experts = example_moe.n_physical_experts
360 + self.num_local_physical_experts = example_moe.n_local_physical_experts
361 + self.num_routed_experts = example_moe.n_routed_experts
362 + self.num_redundant_experts = example_moe.n_redundant_experts
363 +
364 + def update_physical_experts_metadata(
365 + self,
366 + num_physical_experts: int,
367 + num_local_physical_experts: int,
368 + ) -> None:
369 + assert self.num_local_physical_experts == num_local_physical_experts
370 + self.num_physical_experts = num_physical_experts
371 + self.num_local_physical_experts = num_local_physical_experts
372 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
373 + for moe in self.moe_mlp_layers:
374 + moe.n_physical_experts = num_physical_experts
375 + moe.n_local_physical_experts = num_local_physical_experts
376 + moe.n_redundant_experts = self.num_redundant_experts
377 + moe.experts.update_expert_map()
378 +
379 +
380 +@support_torch_compile(
381 + dynamic_arg_dims={
382 + "input_ids": 0,
383 + "positions": -1,
384 + "intermediate_tensors": 0,
385 + "inputs_embeds": 0,
386 + "query_start_loc": 0,
387 + "ngram_context": 0,
388 + "deepstack_input_embeds": 0,
389 + }
390 +)
391 +class Qwen4ExpModel(nn.Module):
392 + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
393 +
394 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
395 + super().__init__()
396 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
397 + self.config = config
398 + self.num_redundant_experts = (
399 + vllm_config.parallel_config.eplb_config.num_redundant_experts
400 + )
401 + self.vocab_size = config.vocab_size
402 + self._qsa_layer_ids = frozenset(
403 + layer_idx
404 + for layer_idx, layer_type in enumerate(config.layer_types)
405 + if layer_type == "full_attention"
406 + and getattr(config, "indexer_n_heads", None) is not None
407 + )
408 + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size)
409 +
410 + def get_layer(prefix: str) -> Qwen4ExpDecoderLayer:
411 + layer_idx = extract_layer_index(prefix)
412 + return Qwen4ExpDecoderLayer(
413 + vllm_config,
414 + layer_type=config.layer_types[layer_idx],
415 + prefix=prefix,
416 + )
417 +
418 + self.start_layer, self.end_layer, self.layers = make_layers(
419 + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
420 + )
421 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
422 + self.layers,
423 + Qwen4ExpSparseMoeBlock,
424 + "mlp",
425 + )
426 + intermediate_size = config.hidden_size * config.hc_count
427 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
428 + ["hidden_states"], intermediate_size
429 + )
430 +
431 + self.hyper_connection_mixer: GatedResidual | None
432 + if get_pp_group().is_last_rank:
433 + hc_config = HyperConnectionConfig(
434 + hc_count=config.hc_count,
435 + hidden_size=config.hidden_size,
436 + params_dtype=torch.bfloat16,
437 + hc_lowrank=config.hc_lowrank,
438 + rms_norm_eps=config.rms_norm_eps,
439 + hc_per_branch_norm=True,
440 + )
441 + self.hyper_connection_mixer = GatedResidual(
442 + hc_config,
443 + use_combine=False,
444 + prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
445 + )
446 + else:
447 + self.hyper_connection_mixer = None
448 +
449 + spec_config = vllm_config.speculative_config
450 + # MTP HC multi-stream outputs: when speculative method=="mtp" and the
451 + # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
452 + # hidden state [T, hc_count*H] so the MTP drafter can feed a real
453 + # multi-stream backbone hidden on its first step (scheme A). Derived
454 + # purely from config (NOT node identity) so P/D nodes stay consistent.
455 + needs_mtp_hidden = (
456 + spec_config is not None
457 + and getattr(spec_config, "method", None) == "mtp"
458 + and get_pp_group().is_last_rank
459 + )
460 + if needs_mtp_hidden:
461 + self._mtp_hidden_buffer = torch.empty(
462 + vllm_config.scheduler_config.max_num_batched_tokens,
463 + config.hc_count * config.hidden_size,
464 + dtype=vllm_config.model_config.dtype,
465 + )
466 + else:
467 + self._mtp_hidden_buffer = None
468 +
469 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
470 + return self.embed_tokens(input_ids)
471 +
472 + def forward(
473 + self,
474 + input_ids: torch.Tensor | None,
475 + positions: torch.Tensor,
476 + intermediate_tensors: IntermediateTensors | None = None,
477 + inputs_embeds: torch.Tensor | None = None,
478 + query_start_loc: torch.Tensor | None = None,
479 + ngram_context: torch.Tensor | None = None,
480 + deepstack_input_embeds: IntermediateTensors | None = None,
481 + ) -> torch.Tensor | IntermediateTensors:
482 + if get_pp_group().is_first_rank:
483 + if inputs_embeds is not None:
484 + hidden_states = inputs_embeds
485 + else:
486 + if input_ids is None:
487 + raise ValueError("input_ids or inputs_embeds is required")
488 + hidden_states = self.embed_input_ids(input_ids)
489 + hidden_states = hidden_states.repeat(1, self.config.hc_count)
490 + else:
491 + if intermediate_tensors is None:
492 + raise ValueError("pipeline stage requires intermediate tensors")
493 + hidden_states = intermediate_tensors["hidden_states"]
494 +
495 + block_output = None
496 + injection = None
497 + last_layer = None
498 + for layer_idx, layer in islice(
499 + enumerate(self.layers), self.start_layer, self.end_layer
500 + ):
501 + last_layer = layer
502 + hidden_states, block_output, injection = layer(
503 + hidden_states=hidden_states,
504 + prev_block_output=block_output,
505 + prev_injection=injection,
506 + positions=positions,
507 + input_ids=input_ids,
508 + query_start_loc=query_start_loc,
509 + ngram_context=ngram_context,
510 + )
511 + if deepstack_input_embeds is not None and layer_idx < len(
512 + deepstack_input_embeds
513 + ):
514 + deepstack_embed = deepstack_input_embeds[
515 + f"deepstack_input_embeds_{layer_idx}"
516 + ]
517 + deepstack_embed = (
518 + deepstack_embed.unsqueeze(-2)
519 + .expand(
520 + *deepstack_embed.shape[:-1],
521 + self.config.hc_count,
522 + self.config.hidden_size,
523 + )
524 + .flatten(-2)
525 + )
526 + # Deepstack is an external addition to the materialized
527 + # multi-stream state and therefore terminates delayed combine.
528 + hidden_states = layer.mlp_hyper_connection.combine(
529 + hidden_states, block_output, injection
530 + )
531 + block_output = None
532 + injection = None
533 + hidden_states = hidden_states + deepstack_embed
534 +
535 + if not get_pp_group().is_last_rank:
536 + # PP transports one tensor, not the delayed HC tuple. Materialize
537 + # with the HC module that produced the pending injection.
538 + if last_layer is not None and block_output is not None:
539 + hidden_states = last_layer.mlp_hyper_connection.combine(
540 + hidden_states, block_output, injection
541 + )
542 + return IntermediateTensors({"hidden_states": hidden_states})
543 +
544 + # The final mixer consumes the last pending combine and returns both
545 + # the sampled single stream and the materialized multi-stream state.
546 + final_mixer = self.hyper_connection_mixer
547 + assert final_mixer is not None
548 + multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix(
549 + hidden_states, block_output, injection
550 + )
551 + if self._mtp_hidden_buffer is not None:
552 + # Capture the pre-final-mixer multi-stream hidden state
553 + # [T, hc_count*H] for the MTP drafter (zero extra compute:
554 + # this tensor is needed by the final mixer regardless).
555 + num_tokens = multi_hidden.shape[0]
556 + self._mtp_hidden_buffer[:num_tokens].copy_(multi_hidden)
557 + return sample_hidden_states
558 +
559 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
560 + weights = (
561 + (
562 + _remap_qsa_cache_scale_name(name, self._qsa_layer_ids),
563 + weight,
564 + )
565 + for name, weight in weights
566 + )
567 + weights = maybe_fuse_shared_experts(
568 + weights,
569 + enabled=self.is_fused_shared_expert_enabled,
570 + n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
571 + n_shared_experts=1,
572 + ckpt_prefix="mlp.shared_expert",
573 + )
574 + # Non-persistent PLE state rebuilt in __init__; skip any ckpt
575 + # column for them.
576 + skip_substrs = (
577 + "hashstats_",
578 + "token_lookup",
579 + "hyper_connection_mixer.block_inject_weight",
580 + )
581 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
582 + orig_to_new_substr={substr: None for substr in skip_substrs}
583 + )
584 + # The final HC mixer only exists on the last PP rank; earlier ranks
585 + # must drop its checkpoint weights instead of failing to place them.
586 + ignore_prefixes = (
587 + None
588 + if self.hyper_connection_mixer is not None
589 + else ["hyper_connection_mixer."]
590 + )
591 + loader = AutoWeightsLoader(
592 + self,
593 + ignore_unexpected_prefixes=ignore_prefixes,
594 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
595 + )
596 + loaded = loader.load_weights(
597 + weights,
598 + mapper=mapper,
599 + )
600 + return loaded
601 +
602 +
603 +class Qwen4ExpForCausalLM(
604 + nn.Module,
605 + HasInnerState,
606 + SupportsLoRA,
607 + SupportsMRoPE,
608 + SupportsPP,
609 + Qwen4ExpMixtureOfExperts,
610 + IsHybrid,
611 +):
612 + packed_modules_mapping = {
613 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
614 + "gate_up_proj": ["gate_proj", "up_proj"],
615 + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
616 + "in_proj_ba": ["in_proj_b", "in_proj_a"],
617 + "input_mix_weight_down_block_inject": [
618 + "input_mix_weight_down",
619 + "block_inject_weight",
620 + "_input_mix_padding",
621 + ],
622 + }
623 + hf_to_vllm_mapper = WeightsMapper(
624 + orig_to_new_prefix={"model.language_model.": "model."}
625 + )
626 + requires_raw_input_tokens = True
627 +
628 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
629 + super().__init__()
630 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
631 + self.vllm_config = vllm_config
632 + self.model_config = vllm_config.model_config
633 + self.quant_config = vllm_config.quant_config
634 + self.config = config
635 + self.scheduler_config = vllm_config.scheduler_config
636 + if vllm_config.cache_config.mamba_cache_mode == "all":
637 + raise NotImplementedError(
638 + "Qwen4Exp currently does not support 'all' prefix caching, "
639 + "please use '--mamba-cache-mode=align' instead"
640 + )
641 + self.model = Qwen4ExpModel(
642 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
643 + )
644 + self.lm_head = ParallelLMHead(
645 + config.vocab_size,
646 + config.hidden_size,
647 + prefix=maybe_prefix(prefix, "lm_head"),
648 + )
649 + self.logits_processor = LogitsProcessor(config.vocab_size)
650 + self.make_empty_intermediate_tensors = (
651 + self.model.make_empty_intermediate_tensors
652 + )
653 + self.set_moe_parameters(self.model.layers)
654 + enable_qwen4_exp_low_latency_gemm(self, self.model_config.dtype)
655 +
656 + @staticmethod
657 + def get_model_state_cls():
658 + from .model_state import Qwen4ExpModelState
659 +
660 + return Qwen4ExpModelState
661 +
662 + def forward(
663 + self,
664 + input_ids: torch.Tensor | None,
665 + positions: torch.Tensor,
666 + intermediate_tensors: IntermediateTensors | None = None,
667 + inputs_embeds: torch.Tensor | None = None,
668 + **kwargs: object,
669 + ) -> torch.Tensor | IntermediateTensors:
670 + # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs
671 + # path (query_start_loc / ngram_context) reaches Qwen4ExpModel.
672 + return self.model(
673 + input_ids,
674 + positions,
675 + intermediate_tensors,
676 + inputs_embeds,
677 + **kwargs,
678 + )
679 +
680 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
681 + return self.model.embed_input_ids(input_ids)
682 +
683 + @classmethod
684 + def get_ple_mamba_state_dtype_from_config(
685 + cls,
686 + vllm_config: VllmConfig,
687 + ) -> tuple[torch.dtype, ...]:
688 + return MambaStateDtypeCalculator.short_conv_state_dtype(
689 + vllm_config.model_config.dtype,
690 + vllm_config.cache_config.mamba_cache_dtype,
691 + )
692 +
693 + @classmethod
694 + def get_ple_mamba_state_shape_from_config(
695 + cls, vllm_config: VllmConfig
696 + ) -> tuple[tuple[int, int]]:
697 + hf_config = vllm_config.model_config.hf_text_config
698 + conv_kernel_size = hf_config.ple_conv_kernel_size
699 + short_conv_dilation = hf_config.ngram_size
700 + conv_state_len = (conv_kernel_size - 1) * short_conv_dilation
701 + num_spec = (
702 + vllm_config.speculative_config.num_speculative_tokens
703 + if vllm_config.speculative_config
704 + else 0
705 + )
706 + hc_count = hf_config.hc_count
707 + hc_hidden_size = hf_config.hidden_size * hc_count
708 + return MambaStateShapeCalculator.short_conv_state_shape(
709 + tp_world_size=1,
710 + intermediate_size=hc_hidden_size,
711 + conv_kernel=conv_state_len + 1,
712 + num_spec=num_spec,
713 + )
714 +
715 + @classmethod
716 + def get_gdn_mamba_state_dtype_from_config(
717 + cls, vllm_config: VllmConfig
718 + ) -> tuple[torch.dtype, torch.dtype]:
719 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
720 + vllm_config.model_config.dtype,
721 + vllm_config.cache_config.mamba_cache_dtype,
722 + vllm_config.cache_config.mamba_ssm_cache_dtype,
723 + )
724 +
725 + @classmethod
726 + def get_gdn_mamba_state_shape_from_config(
727 + cls, vllm_config: VllmConfig
728 + ) -> tuple[tuple[int, int], tuple[int, int]]:
729 + parallel_config = vllm_config.parallel_config
730 + hf_config = vllm_config.model_config.hf_text_config
731 + tp_size = parallel_config.tensor_parallel_size
732 + num_spec = (
733 + vllm_config.speculative_config.num_speculative_tokens
734 + if vllm_config.speculative_config
735 + else 0
736 + )
737 + return MambaStateShapeCalculator.gated_delta_net_state_shape(
738 + tp_size,
739 + hf_config.linear_num_key_heads,
740 + hf_config.linear_num_value_heads,
741 + hf_config.linear_key_head_dim,
742 + hf_config.linear_value_head_dim,
743 + hf_config.linear_conv_kernel_dim,
744 + num_spec,
745 + )
746 +
747 + @classmethod
748 + def get_mamba_state_dtype_from_config(
749 + cls,
750 + vllm_config: VllmConfig,
751 + ) -> tuple[torch.dtype, torch.dtype]:
752 + return cls.get_gdn_mamba_state_dtype_from_config(vllm_config)
753 +
754 + @classmethod
755 + def get_mamba_state_shape_from_config(
756 + cls, vllm_config: VllmConfig
757 + ) -> tuple[tuple[int, int], tuple[int, int]]:
758 + return cls.get_gdn_mamba_state_shape_from_config(vllm_config)
759 +
760 + @classmethod
761 + def get_mamba_state_copy_func(
762 + cls,
763 + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
764 + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
765 +
766 + @classmethod
767 + def get_mamba_state_copy_funcs(
768 + cls,
769 + mamba_types: set[MambaAttentionBackendEnum],
770 + ) -> MambaStateCopyFuncsByType:
771 + copy_funcs_by_type = {
772 + MambaAttentionBackendEnum.GDN_ATTN: cls.get_mamba_state_copy_func(),
773 + MambaAttentionBackendEnum.SHORT_CONV: (
774 + MambaStateCopyFuncCalculator.short_conv_state_copy_func()
775 + ),
776 + }
777 + missing_types = mamba_types - copy_funcs_by_type.keys()
778 + assert not missing_types, f"missing state copy funcs for {missing_types}"
779 + return {
780 + mamba_type: copy_funcs_by_type[mamba_type] for mamba_type in mamba_types
781 + }
782 +
783 + @classmethod
784 + def get_mamba_specs_from_config(
785 + cls, vllm_config: VllmConfig
786 + ) -> tuple[MambaSpec, ...]:
787 + """Return all MambaSpecs for this model (GDN layers + PLE layer).
788 +
789 + The PLE layer uses a separate short_conv MambaSpec whose page_size_bytes
790 + may exceed the GDN spec; callers should take the maximum.
791 + """
792 + return (
793 + MambaSpec(
794 + shapes=cls.get_gdn_mamba_state_shape_from_config(vllm_config),
795 + dtypes=cls.get_gdn_mamba_state_dtype_from_config(vllm_config),
796 + block_size=-1,
797 + ),
798 + MambaSpec(
799 + shapes=cls.get_ple_mamba_state_shape_from_config(vllm_config),
800 + dtypes=cls.get_ple_mamba_state_dtype_from_config(vllm_config),
801 + block_size=-1,
802 + tp_replicated=True,
803 + ),
804 + )
805 +
806 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
807 + return self.logits_processor(self.lm_head, hidden_states)
808 +
809 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
810 + return self.model._mtp_hidden_buffer
811 +
812 + def get_mrope_input_positions(
813 + self,
814 + input_tokens: list[int],
815 + mm_features: list[MultiModalFeatureSpec],
816 + ) -> tuple[torch.Tensor, int]:
817 + positions = torch.arange(len(input_tokens), dtype=torch.long)
818 + return positions.unsqueeze(0).expand(3, -1), 0
819 +
820 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
821 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
822 + orig_to_new_substr={"mtp.": None}
823 + )
824 + loader = AutoWeightsLoader(
825 + self,
826 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
827 + )
828 + return loader.load_weights(weights, mapper=mapper)
829 +
830 +
831 +class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo):
832 + def get_hf_config(self) -> Qwen4ExpConfig:
833 + return self.ctx.get_hf_config(Qwen4ExpConfig)
834 +
835 +
836 +@MULTIMODAL_REGISTRY.register_processor(
837 + Qwen3VLMultiModalProcessor,
838 + info=Qwen4ExpProcessingInfo,
839 + dummy_inputs=Qwen3VLDummyInputsBuilder,
840 +)
841 +class Qwen4ExpForConditionalGeneration(
842 + Qwen3_5ForConditionalGeneration,
843 + HasInnerState,
844 + Qwen4ExpMixtureOfExperts,
845 +):
846 + """Qwen3-VL vision tower backed by the Qwen4Exp language model."""
847 +
848 + requires_raw_input_tokens = True
849 +
850 + packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
851 + "input_mix_weight_down_block_inject": [
852 + "input_mix_weight_down",
853 + "block_inject_weight",
854 + "_input_mix_padding",
855 + ]
856 + }
857 +
858 + @staticmethod
859 + def get_model_state_cls():
860 + from .model_state import Qwen4ExpModelState
861 +
862 + return Qwen4ExpModelState
863 +
864 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
865 + nn.Module.__init__(self)
866 + config: Qwen4ExpConfig = vllm_config.model_config.hf_config
867 + quant_config = vllm_config.quant_config
868 + multimodal_config = vllm_config.model_config.multimodal_config
869 + if multimodal_config is None:
870 + raise ValueError(
871 + "Qwen4ExpForConditionalGeneration requires multimodal_config"
872 + )
873 +
874 + self.config = config
875 + self.model_config = vllm_config.model_config
876 + self.multimodal_config = multimodal_config
877 + self.language_model_only = multimodal_config.language_model_only
878 + if self.language_model_only:
879 + self.use_data_parallel = False
880 + self.is_multimodal_pruning_enabled = False
881 + self.video_pruning_method = None
882 + self.video_pruning_rate = 0.0
883 + self._tokenizer = None
884 + self.visual = StageMissingLayer("vision_tower")
885 + self._tower_model_names = []
886 + else:
887 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
888 + self._init_video_pruning(multimodal_config)
889 + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
890 +
891 + with self._mark_tower_model(vllm_config, {"image", "video"}):
892 + self.visual = Qwen3_VisionTransformer(
893 + config.vision_config,
894 + norm_eps=config.text_config.rms_norm_eps,
895 + quant_config=quant_config,
896 + prefix=maybe_prefix(prefix, "visual"),
897 + )
898 +
899 + self.use_deepstack = (
900 + not self.language_model_only
901 + and bool(config.vision_config.deepstack_visual_indexes)
902 + and not isinstance(self.visual, StageMissingLayer)
903 + )
904 + self.deepstack_num_level = (
905 + len(config.vision_config.deepstack_visual_indexes)
906 + if self.use_deepstack
907 + else 0
908 + )
909 + self.visual_dim = config.vision_config.out_hidden_size
910 + self.multiscale_dim = self.visual_dim * self.deepstack_num_level
911 +
912 + if self.use_deepstack:
913 + self.deepstack_input_embeds = [
914 + torch.zeros(
915 + vllm_config.scheduler_config.max_num_batched_tokens,
916 + config.text_config.hidden_size,
917 + )
918 + for _ in range(self.deepstack_num_level)
919 + ]
920 + self.deepstack_input_embeds_num_tokens = 0
921 +
922 + with self._mark_language_model(vllm_config):
923 + self.language_model = Qwen4ExpForCausalLM(
924 + vllm_config=vllm_config,
925 + prefix=maybe_prefix(prefix, "language_model"),
926 + )
927 +
928 + self.make_empty_intermediate_tensors = (
929 + self.language_model.make_empty_intermediate_tensors
930 + )
931 + if not get_pp_group().is_first_rank and self.use_deepstack:
932 + assert self.language_model.model.start_layer >= len(
933 + config.vision_config.deepstack_visual_indexes
934 + ), (
935 + "start_layer should be greater than or equal to "
936 + "len(deepstack_visual_indexes)"
937 + )
938 + self.set_moe_parameters(self.language_model.model.layers)
939 +
940 + def embed_input_ids(
941 + self,
942 + input_ids: torch.Tensor,
943 + multimodal_embeddings: MultiModalEmbeddings | None = None,
944 + *,
945 + is_multimodal: torch.Tensor | None = None,
946 + ) -> torch.Tensor:
947 + inputs_embeds = self._embed_text_input_ids(
948 + input_ids,
949 + self.language_model.embed_input_ids,
950 + is_multimodal=is_multimodal,
951 + )
952 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
953 + return inputs_embeds
954 + if self.language_model_only:
955 + raise ValueError(
956 + "Qwen4Exp language_model_only does not accept multimodal embeddings"
957 + )
958 +
959 + is_multimodal = _require_is_multimodal(is_multimodal)
960 + if self.use_deepstack:
961 + deepstack_input_embeds, multimodal_embeddings = (
962 + self._compute_deepstack_embeds(
963 + inputs_embeds=inputs_embeds,
964 + multimodal_embeddings=multimodal_embeddings,
965 + is_multimodal=is_multimodal,
966 + )
967 + )
968 + else:
969 + deepstack_input_embeds = None
970 +
971 + inputs_embeds = _merge_multimodal_embeddings(
972 + inputs_embeds=inputs_embeds,
973 + multimodal_embeddings=multimodal_embeddings,
974 + is_multimodal=is_multimodal,
975 + )
976 + if deepstack_input_embeds is not None:
977 + self._set_deepstack_input_embeds(deepstack_input_embeds)
978 + return inputs_embeds
979 +
980 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
981 + return self.language_model.get_mtp_target_hidden_states()
982 +
983 + def forward(
984 + self,
985 + input_ids: torch.Tensor | None,
986 + positions: torch.Tensor,
987 + intermediate_tensors: IntermediateTensors | None = None,
988 + inputs_embeds: torch.Tensor | None = None,
989 + **kwargs: object,
990 + ) -> torch.Tensor | IntermediateTensors:
991 + if intermediate_tensors is not None:
992 + inputs_embeds = None
993 + if inputs_embeds is not None and get_pp_group().is_first_rank:
994 + deepstack_input_embeds = self._get_deepstack_input_embeds(
995 + inputs_embeds.size(0)
996 + )
997 + else:
998 + deepstack_input_embeds = None
999 +
1000 + hidden_states = self.language_model.model(
1001 + input_ids=input_ids,
1002 + positions=positions,
1003 + intermediate_tensors=intermediate_tensors,
1004 + inputs_embeds=inputs_embeds,
1005 + query_start_loc=kwargs.get("query_start_loc"),
1006 + ngram_context=kwargs.get("ngram_context"),
1007 + deepstack_input_embeds=deepstack_input_embeds,
1008 + )
1009 + if inputs_embeds is not None and get_pp_group().is_first_rank:
1010 + self._clear_deepstack_input_embeds(inputs_embeds.size(0))
1011 + return hidden_states
1012 +
1013 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1014 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
1015 + orig_to_new_substr={"mtp.": None},
1016 + orig_to_new_prefix={"visual.": None} if self.language_model_only else {},
1017 + )
1018 + loader = AutoWeightsLoader(
1019 + self,
1020 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
1021 + )
1022 + return loader.load_weights(weights, mapper=mapper)
1023 +
1024 + @classmethod
1025 + def get_mamba_state_dtype_from_config(
1026 + cls,
1027 + vllm_config: VllmConfig,
1028 + ) -> tuple[torch.dtype, torch.dtype]:
1029 + return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
1030 +
1031 + @classmethod
1032 + def get_mamba_state_shape_from_config(
1033 + cls,
1034 + vllm_config: VllmConfig,
1035 + ) -> tuple[tuple[int, int], tuple[int, int]]:
1036 + return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config)
1037 +
1038 + @classmethod
1039 + def get_mamba_state_copy_func(
1040 + cls,
1041 + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
1042 + return Qwen4ExpForCausalLM.get_mamba_state_copy_func()
1043 +
1044 + @classmethod
1045 + def get_mamba_state_copy_funcs(
1046 + cls,
1047 + mamba_types: set[MambaAttentionBackendEnum],
1048 + ) -> MambaStateCopyFuncsByType:
1049 + return Qwen4ExpForCausalLM.get_mamba_state_copy_funcs(mamba_types)
1050 +
1051 + @classmethod
1052 + def get_mamba_specs_from_config(
1053 + cls, vllm_config: VllmConfig
1054 + ) -> tuple[MambaSpec, ...]:
1055 + return Qwen4ExpForCausalLM.get_mamba_specs_from_config(vllm_config)
1056 +
1057 +
1058 +__all__ = [
1059 + "Qwen4ExpDecoderLayer",
1060 + "Qwen4ExpForCausalLM",
1061 + "Qwen4ExpForConditionalGeneration",
1062 + "Qwen4ExpMixtureOfExperts",
1063 + "Qwen4ExpModel",
1064 + "Qwen4ExpSparseMoeBlock",
1065 +]
vllm/models/qwen4_exp/amd/model_state.py new
+143
@@ -0,0 +1,143 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Model-runner state for Qwen4Exp PLE inputs."""
4 +
5 +from typing import Any
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.v1.worker.gpu.input_batch import InputBatch
12 +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
13 +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
14 +from vllm.v1.worker.gpu.states import RequestState
15 +
16 +
17 +class Qwen4ExpModelState(MambaHybridModelState):
18 + """Add rollback-safe PLE n-gram context to the model inputs."""
19 +
20 + def __init__(
21 + self,
22 + vllm_config: VllmConfig,
23 + model: nn.Module,
24 + encoder_cache: EncoderCache | None,
25 + device: torch.device,
26 + ) -> None:
27 + super().__init__(vllm_config, model, encoder_cache, device)
28 + config = self.model_config.hf_text_config
29 + self.uses_ngram_embedding = bool(config.ple_layer_ids)
30 + if not self.uses_ngram_embedding:
31 + self.ngram_context_len = 0
32 + self.ngram_eos_token_id = 0
33 + return
34 +
35 + if vllm_config.parallel_config.pipeline_parallel_size > 1:
36 + raise RuntimeError(
37 + "N-gram PLE embedding currently requires "
38 + "pipeline_parallel_size=1 because non-first pipeline ranks do "
39 + "not receive the raw input_ids required by PLE. Please run "
40 + "with PP=1."
41 + )
42 +
43 + self.ngram_context_len = int(config.ngram_size) - 1
44 + if self.ngram_context_len <= 0:
45 + raise ValueError("N-gram embedding requires context length >= 1.")
46 + self.ngram_eos_token_id = int(config.eos_token_id)
47 + self.ngram_context = torch.full(
48 + (self.max_num_reqs, self.ngram_context_len),
49 + self.ngram_eos_token_id,
50 + dtype=torch.int32,
51 + device=self.device,
52 + )
53 + self.ngram_context_offsets = torch.arange(
54 + -self.ngram_context_len,
55 + 0,
56 + dtype=torch.int64,
57 + device=self.device,
58 + )
59 + self.ple_query_start_loc = torch.zeros(
60 + self.max_num_reqs + 1,
61 + dtype=torch.int32,
62 + device=self.device,
63 + )
64 +
65 + def _prepare_ngram_context(
66 + self,
67 + input_batch: InputBatch,
68 + req_states: RequestState,
69 + ) -> torch.Tensor:
70 + num_reqs = input_batch.num_reqs
71 + num_reqs_padded = input_batch.num_reqs_after_padding
72 + context = self.ngram_context[:num_reqs_padded]
73 + context.fill_(self.ngram_eos_token_id)
74 + if num_reqs == 0:
75 + return context
76 +
77 + request_indices = input_batch.idx_mapping[:num_reqs].long()
78 + context_end = req_states.num_computed_tokens.gpu[request_indices].long()
79 + token_indices = context_end.unsqueeze(1) + self.ngram_context_offsets
80 + valid_tokens = token_indices >= 0
81 + token_indices.clamp_min_(0)
82 + context_tokens = req_states.all_token_ids.gpu[
83 + request_indices.unsqueeze(1), token_indices
84 + ]
85 + context[:num_reqs].copy_(
86 + torch.where(
87 + valid_tokens,
88 + context_tokens,
89 + context_tokens.new_full((), self.ngram_eos_token_id),
90 + )
91 + )
92 + return context
93 +
94 + def prepare_inputs(
95 + self,
96 + input_batch: InputBatch,
97 + req_states: RequestState,
98 + ) -> dict[str, Any]:
99 + model_inputs = super().prepare_inputs(input_batch, req_states)
100 + if not self.uses_ngram_embedding:
101 + return model_inputs
102 +
103 + num_reqs_padded = input_batch.num_reqs_after_padding
104 + query_start_loc = self.ple_query_start_loc[: num_reqs_padded + 1]
105 + query_start_loc.copy_(input_batch.query_start_loc[: num_reqs_padded + 1])
106 + model_inputs.update(
107 + query_start_loc=query_start_loc,
108 + ngram_context=self._prepare_ngram_context(input_batch, req_states),
109 + )
110 + return model_inputs
111 +
112 + def prepare_dummy_inputs(
113 + self,
114 + num_reqs: int,
115 + num_tokens: int,
116 + ) -> dict[str, Any]:
117 + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
118 + if not self.uses_ngram_embedding:
119 + return model_inputs
120 +
121 + query_start_loc = self.ple_query_start_loc[: num_reqs + 1]
122 + query_start_loc[0] = 0
123 + tokens_per_req, num_extra_tokens = divmod(num_tokens, num_reqs)
124 + query_lens = torch.full(
125 + (num_reqs,),
126 + tokens_per_req,
127 + dtype=query_start_loc.dtype,
128 + device=query_start_loc.device,
129 + )
130 + if num_extra_tokens > 0:
131 + query_lens[-num_extra_tokens:] += 1
132 + torch.cumsum(query_lens, dim=0, out=query_start_loc[1:])
133 +
134 + ngram_context = self.ngram_context[:num_reqs]
135 + ngram_context.fill_(self.ngram_eos_token_id)
136 + model_inputs.update(
137 + query_start_loc=query_start_loc,
138 + ngram_context=ngram_context,
139 + )
140 + return model_inputs
141 +
142 +
143 +__all__ = ["Qwen4ExpModelState"]
vllm/models/qwen4_exp/amd/mtp.py new
+462
@@ -0,0 +1,462 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Qwen4Exp MTP (Multi-Token Predictor) model.
4 +
5 +The MTP draft model reuses the Qwen4Exp backbone (PLE/HC/MoE) but:
6 + - drops all multi-modal handling (text-only),
7 + - forces PLE off while keeping the main model's HC stream count,
8 + - fuses the backbone hidden and the new-token embedding via
9 + ``residual_linear_shared`` (fc_embedding + shared fc_hidden) instead of
10 + the ``Linear(2H, H)`` + repeat used by other MTP variants,
11 + - emits TWO hidden streams per step (scheme A): a single stream [T, H]
12 + (final-mixer collapsed, fed to the LM head) and a pre-final-mixer
13 + multi stream [T, hc_count*H] (fed to the next draft step).
14 +"""
15 +
16 +from collections.abc import Iterable
17 +
18 +import regex as re
19 +import torch
20 +from torch import nn
21 +
22 +from vllm.compilation.decorators import support_torch_compile
23 +from vllm.config import VllmConfig, replace, set_current_vllm_config
24 +from vllm.distributed import get_pp_group
25 +from vllm.model_executor.layers.fused_moe.utils import (
26 + is_model_fused_shared_expert_compatible,
27 +)
28 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
29 +from vllm.model_executor.layers.linear import ColumnParallelLinear
30 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
31 +from vllm.model_executor.layers.vocab_parallel_embedding import (
32 + ParallelLMHead,
33 + VocabParallelEmbedding,
34 +)
35 +from vllm.model_executor.model_loader.utils import configure_quant_config
36 +from vllm.model_executor.models.interfaces import SupportsPP
37 +from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
38 +from vllm.model_executor.models.utils import (
39 + AutoWeightsLoader,
40 + PPMissingLayer,
41 + WeightsMapper,
42 + get_draft_quant_config,
43 + make_empty_intermediate_tensors_factory,
44 + maybe_fuse_shared_experts,
45 + maybe_prefix,
46 +)
47 +from vllm.sequence import IntermediateTensors
48 +from vllm.transformers_utils.configs.qwen4_exp import (
49 + Qwen4ExpTextConfig,
50 +)
51 +
52 +from .hyperconnection import GatedResidual, HyperConnectionConfig
53 +from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
54 +from .model import (
55 + _HC_WEIGHTS_MAPPER,
56 + _QWEN4_EXP_IGNORED_MISSING_SUFFIXES,
57 + Qwen4ExpDecoderLayer,
58 + Qwen4ExpMixtureOfExperts,
59 + Qwen4ExpSparseMoeBlock,
60 +)
61 +
62 +
63 +def _remap_ignored_layers(
64 + ignored_layers: list[str],
65 + mtp_start_layer_idx: int,
66 +) -> list[str]:
67 + remapped: list[str] = []
68 + for name in ignored_layers:
69 + if name.startswith("mtp."):
70 + new_name = re.sub(
71 + r"(?<=\.layers\.)\d+",
72 + lambda m: str(mtp_start_layer_idx + int(m.group(0))),
73 + name,
74 + )
75 + remapped.append(new_name)
76 + else:
77 + remapped.append(name)
78 + return remapped
79 +
80 +
81 +def _remap_mtp_weight_name(name: str) -> str | None:
82 + """Map Qwen4Exp checkpoint paths into the standalone draft model."""
83 +
84 + for checkpoint_prefix in (
85 + "model.language_model.",
86 + "language_model.",
87 + ):
88 + if name.startswith(checkpoint_prefix):
89 + name = name.removeprefix(checkpoint_prefix)
90 + break
91 +
92 + if name.startswith("embed_tokens."):
93 + name = f"model.{name}"
94 + if name.startswith("model.mtp."):
95 + name = name.removeprefix("model.")
96 + if name.startswith("mtp.shared_head.head."):
97 + return name.replace("mtp.shared_head.head.", "lm_head.", 1)
98 + if name.startswith("model.shared_head.head."):
99 + return name.replace("model.shared_head.head.", "lm_head.", 1)
100 + if name.startswith("shared_head.head."):
101 + return name.replace("shared_head.head.", "lm_head.", 1)
102 + if name.startswith("model.lm_head."):
103 + return name.removeprefix("model.")
104 + if name.startswith("mtp."):
105 + return name.replace("mtp.", "model.", 1)
106 + if name.startswith("model.embed_tokens.") or name.startswith("lm_head."):
107 + return name
108 + return None
109 +
110 +
111 +def _make_draft_vllm_config(
112 + vllm_config: VllmConfig,
113 + mtp_start_layer_idx: int,
114 +) -> VllmConfig:
115 + """Ensure that the draft model config is set in the vLLM config."""
116 + speculative_config = vllm_config.speculative_config
117 + if speculative_config is None or speculative_config.draft_model_config is None:
118 + raise ValueError("speculative_config.draft_model_config must be set")
119 +
120 + draft_quant_config = get_draft_quant_config(vllm_config)
121 +
122 + # inject packed and ignored modules to the quantization config of draft model
123 + if draft_quant_config is not None:
124 + configure_quant_config(draft_quant_config, Qwen4ExpMTP)
125 + ignored_layers = getattr(draft_quant_config, "ignored_layers", None)
126 + if ignored_layers:
127 + setattr( # noqa: B010
128 + draft_quant_config,
129 + "ignored_layers",
130 + _remap_ignored_layers(ignored_layers, mtp_start_layer_idx),
131 + )
132 + exclude_modules = getattr(draft_quant_config, "exclude_modules", None)
133 + if exclude_modules:
134 + setattr( # noqa: B010
135 + draft_quant_config,
136 + "exclude_modules",
137 + _remap_ignored_layers(exclude_modules, mtp_start_layer_idx),
138 + )
139 +
140 + draft_vllm_config = replace(
141 + vllm_config,
142 + model_config=speculative_config.draft_model_config,
143 + )
144 + # VllmConfig post-init derives the target quant config, so restore the
145 + # independently resolved draft quant config after replacement.
146 + draft_vllm_config.quant_config = draft_quant_config
147 + return draft_vllm_config
148 +
149 +
150 +@support_torch_compile(
151 + dynamic_arg_dims={
152 + "input_ids": 0,
153 + "positions": -1,
154 + "intermediate_tensors": 0,
155 + "inputs_embeds": 0,
156 + "hidden_states": 0,
157 + }
158 +)
159 +class Qwen4ExpMultiTokenPredictor(nn.Module):
160 + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
161 +
162 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
163 + super().__init__()
164 +
165 + model_config = vllm_config.model_config
166 + config: Qwen4ExpTextConfig = model_config.hf_text_config
167 +
168 + self.config = config
169 + self.vocab_size = config.vocab_size
170 +
171 + self.mtp_start_layer_idx = config.num_hidden_layers
172 + self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
173 +
174 + self.hidden_size = config.hidden_size
175 + self.hc_count = config.hc_count
176 +
177 + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size)
178 + draft_vllm_config = _make_draft_vllm_config(
179 + vllm_config,
180 + self.mtp_start_layer_idx,
181 + )
182 + with set_current_vllm_config(draft_vllm_config, prefix=prefix):
183 + # residual_linear_shared fusion: fc_embedding projects the token
184 + # embedding, fc_hidden (shared across HC branches) projects the
185 + # backbone hidden; the embedding is added as a residual to every
186 + # branch (see mtp_residual_linear_shared.md).
187 + self.fc_embedding = ColumnParallelLinear(
188 + self.hidden_size,
189 + self.hidden_size,
190 + gather_output=True,
191 + bias=False,
192 + return_bias=False,
193 + quant_config=draft_vllm_config.quant_config,
194 + prefix=f"{prefix}.fc_embedding",
195 + )
196 + self.fc_hidden = ColumnParallelLinear(
197 + self.hidden_size,
198 + self.hidden_size,
199 + gather_output=True,
200 + bias=False,
201 + return_bias=False,
202 + quant_config=draft_vllm_config.quant_config,
203 + prefix=f"{prefix}.fc_hidden",
204 + )
205 + self.layers = nn.ModuleList(
206 + Qwen4ExpDecoderLayer(
207 + draft_vllm_config,
208 + layer_type="full_attention",
209 + prefix=f"{prefix}.layers.{self.mtp_start_layer_idx + idx}",
210 + )
211 + for idx in range(self.num_mtp_layers)
212 + )
213 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
214 + self.layers,
215 + Qwen4ExpSparseMoeBlock,
216 + "mlp",
217 + )
218 +
219 + self.pre_fc_norm_embedding = GemmaRMSNorm(
220 + self.hidden_size, eps=config.rms_norm_eps
221 + )
222 + self.pre_fc_norm_hidden = GemmaRMSNorm(
223 + self.hidden_size * self.hc_count, eps=config.rms_norm_eps
224 + )
225 + # HC final mixer collapses the multi stream into [T, H] for the LM head.
226 + hc_config = HyperConnectionConfig(
227 + hc_count=config.hc_count,
228 + hidden_size=config.hidden_size,
229 + params_dtype=torch.bfloat16,
230 + hc_lowrank=config.hc_lowrank,
231 + rms_norm_eps=config.rms_norm_eps,
232 + hc_per_branch_norm=True,
233 + )
234 + self.hyper_connection_mixer = GatedResidual(
235 + hc_config,
236 + use_combine=False,
237 + prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
238 + )
239 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
240 + ["hidden_states"], self.hidden_size * self.hc_count
241 + )
242 +
243 + def _iter_qsa_attentions(self):
244 + """Yield MTP attention modules that own a QSA indexer."""
245 +
246 + for layer in self.layers:
247 + attention = getattr(layer, "self_attn", None)
248 + if (
249 + attention is not None
250 + and getattr(attention, "indexer", None) is not None
251 + ):
252 + yield attention
253 +
254 + def set_skip_topk(self, skip: bool) -> None:
255 + """Select on MTP step 0 and reuse its QSA indices on later steps."""
256 +
257 + for attention in self._iter_qsa_attentions():
258 + attention.indexer.skip_topk = skip
259 +
260 + def compact_topk_indices(self, row_indices: torch.Tensor) -> None:
261 + """Keep each request's target-aligned step-0 sparse-index row."""
262 +
263 + num_rows = row_indices.numel()
264 + for attention in self._iter_qsa_attentions():
265 + buffer = attention.topk_indices_buffer
266 + selected = buffer.index_select(0, row_indices)
267 + buffer[:num_rows].copy_(selected)
268 +
269 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
270 + return self.embed_tokens(input_ids)
271 +
272 + def forward(
273 + self,
274 + input_ids: torch.Tensor | None,
275 + positions: torch.Tensor,
276 + hidden_states: torch.Tensor | None = None,
277 + intermediate_tensors: IntermediateTensors | None = None,
278 + inputs_embeds: torch.Tensor | None = None,
279 + spec_step_idx: int = 0,
280 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
281 + hc_count = self.hc_count
282 + hidden_size = self.hidden_size
283 +
284 + if get_pp_group().is_first_rank:
285 + assert hidden_states is not None
286 + if inputs_embeds is None:
287 + assert input_ids is not None
288 + inputs_embeds = self.embed_input_ids(input_ids)
289 + # Embedding branch: pre-norm -> fc_embedding -> [T, H].
290 + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
291 + inputs_embeds = self.fc_embedding(inputs_embeds)
292 +
293 + # Backbone hidden is multi-stream [T, hc_count*H] (scheme A:
294 + # the main model truly emits the pre-final-mixer multi stream
295 + # on the first step; subsequent steps reuse the prior draft
296 + # step's multi stream).
297 + num_tokens = hidden_states.shape[0]
298 + hidden_states = hidden_states.view(num_tokens, hc_count, hidden_size)
299 + hidden_states = self.pre_fc_norm_hidden(hidden_states.flatten(-2)).view(
300 + num_tokens, hc_count, hidden_size
301 + )
302 + hidden_states = self.fc_hidden(hidden_states)
303 + # Add the embedding residual to every branch, then fold back
304 + # to [T, hc_count*H] (HC outer, HS inner) for the HC decoder.
305 + hidden_states = inputs_embeds.unsqueeze(-2) + hidden_states
306 + hidden_states = hidden_states.flatten(-2)
307 + else:
308 + assert intermediate_tensors is not None
309 + hidden_states = intermediate_tensors["hidden_states"]
310 +
311 + current_step_idx = spec_step_idx % self.num_mtp_layers
312 + layer = self.layers[current_step_idx]
313 + hidden_states, block_output, injection = layer(
314 + hidden_states=hidden_states,
315 + prev_block_output=None,
316 + prev_injection=None,
317 + positions=positions,
318 + input_ids=None,
319 + query_start_loc=None,
320 + ngram_context=None,
321 + )
322 + if not get_pp_group().is_last_rank:
323 + # As in the target model, PP carries a materialized tensor rather
324 + # than the delayed hidden/output/injection tuple.
325 + hidden_states = layer.mlp_hyper_connection.combine(
326 + hidden_states, block_output, injection
327 + )
328 + return IntermediateTensors({"hidden_states": hidden_states})
329 +
330 + # Last PP rank finalize. Keep both:
331 + # (A) sample_hidden_states [T, H] -> single stream for the LM head
332 + # (B) multi_hidden [T, hc_count*H] -> pre-final-mixer multi stream
333 + # for the next draft step (zero extra compute, just kept).
334 + multi_hidden, sample_hidden_states, _ = (
335 + self.hyper_connection_mixer.combine_and_mix(
336 + hidden_states, block_output, injection
337 + )
338 + )
339 + return sample_hidden_states, multi_hidden
340 +
341 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
342 + weights = maybe_fuse_shared_experts(
343 + weights,
344 + enabled=self.is_fused_shared_expert_enabled,
345 + n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
346 + n_shared_experts=1,
347 + ckpt_prefix="mlp.shared_expert",
348 + )
349 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
350 + orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
351 + )
352 + loader = AutoWeightsLoader(
353 + self,
354 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
355 + )
356 + return loader.load_weights(weights, mapper=mapper)
357 +
358 +
359 +@support_torch_compile(
360 + dynamic_arg_dims={
361 + "input_ids": 0,
362 + "positions": -1,
363 + "intermediate_tensors": 0,
364 + "inputs_embeds": 0,
365 + "hidden_states": 0,
366 + }
367 +)
368 +class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts):
369 + packed_modules_mapping = {
370 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
371 + "gate_up_proj": ["gate_proj", "up_proj"],
372 + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
373 + "in_proj_ba": ["in_proj_b", "in_proj_a"],
374 + "input_mix_weight_down_block_inject": [
375 + "input_mix_weight_down",
376 + "block_inject_weight",
377 + "_input_mix_padding",
378 + ],
379 + }
380 +
381 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
382 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
383 + self.vllm_config = vllm_config
384 + cache_config = vllm_config.cache_config
385 + if cache_config.mamba_cache_mode == "all":
386 + raise NotImplementedError(
387 + "Qwen4ExpMTP currently does not support 'all' prefix caching, "
388 + "please use '--mamba-cache-mode=align' instead"
389 + )
390 +
391 + self.quant_config = vllm_config.quant_config
392 +
393 + super().__init__()
394 + self.config = config
395 + self.model = Qwen4ExpMultiTokenPredictor(
396 + vllm_config=vllm_config,
397 + prefix=maybe_prefix(prefix, "mtp"),
398 + )
399 +
400 + if get_pp_group().is_last_rank:
401 + if config.tie_word_embeddings:
402 + self.lm_head = self.model.embed_tokens
403 + else:
404 + self.lm_head = ParallelLMHead(
405 + config.vocab_size,
406 + config.hidden_size,
407 + prefix=maybe_prefix(prefix, "lm_head"),
408 + )
409 + else:
410 + self.lm_head = PPMissingLayer()
411 +
412 + self.logits_processor = LogitsProcessor(config.vocab_size)
413 + self.make_empty_intermediate_tensors = (
414 + self.model.make_empty_intermediate_tensors
415 + )
416 + self.set_moe_parameters(self.model.layers)
417 + enable_qwen4_exp_low_latency_gemm(self, vllm_config.model_config.dtype)
418 +
419 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
420 + return self.model.embed_input_ids(input_ids)
421 +
422 + def forward(
423 + self,
424 + input_ids: torch.Tensor | None,
425 + positions: torch.Tensor,
426 + hidden_states: torch.Tensor | None = None,
427 + intermediate_tensors: IntermediateTensors | None = None,
428 + inputs_embeds: torch.Tensor | None = None,
429 + spec_step_idx: int = 0,
430 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
431 + return self.model(
432 + input_ids,
433 + positions,
434 + hidden_states,
435 + intermediate_tensors,
436 + inputs_embeds,
437 + spec_step_idx=spec_step_idx,
438 + )
439 +
440 + def compute_logits(
441 + self, hidden_states: torch.Tensor, spec_step_idx: int = 0
442 + ) -> torch.Tensor | None:
443 + return self.logits_processor(self.lm_head, hidden_states)
444 +
445 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
446 + def remap_weight_names():
447 + for name, weight in weights:
448 + remapped_name = _remap_mtp_weight_name(name)
449 + if remapped_name is not None:
450 + yield remapped_name, weight
451 +
452 + mapper = WeightsMapper(
453 + orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
454 + )
455 + loader = AutoWeightsLoader(
456 + self,
457 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
458 + )
459 + return loader.load_weights(remap_weight_names(), mapper=mapper)
460 +
461 +
462 +__all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"]
vllm/models/qwen4_exp/amd/ops/__init__.py new
+3
@@ -0,0 +1,3 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AMD ROCm Qwen4Exp kernels; import leaf modules directly."""
vllm/models/qwen4_exp/amd/ops/hc.py new
+487
@@ -0,0 +1,487 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AMD ROCm HyperConnection kernels for Qwen4Exp."""
4 +
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +from vllm.triton_utils import tl, triton
9 +from vllm.utils.torch_utils import direct_register_custom_op
10 +
11 +
12 +@triton.jit
13 +def _grouped_gemma_rmsnorm_kernel(
14 + x_ptr,
15 + w_ptr,
16 + y_ptr,
17 + stride_x,
18 + stride_y,
19 + DIM: tl.constexpr,
20 + NUM_GROUPS: tl.constexpr,
21 + W_SHARED: tl.constexpr,
22 + EPS: tl.constexpr,
23 + launch_pdl: tl.constexpr,
24 +) -> None:
25 + GROUP_DIM: tl.constexpr = DIM // NUM_GROUPS
26 + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(GROUP_DIM)
27 +
28 + pid = tl.program_id(0)
29 + group_id = pid % NUM_GROUPS
30 + row = pid // NUM_GROUPS
31 +
32 + offs_g = tl.arange(0, BLOCK_SIZE)
33 + offsets = group_id * GROUP_DIM + offs_g
34 + mask = offs_g < GROUP_DIM
35 + # A [GROUP_DIM] affine is shared; a [DIM] affine follows the grouped
36 + # checkpoint layout.
37 + w_offs = offs_g if W_SHARED else offsets
38 +
39 + if launch_pdl:
40 + tl.extra.cuda.gdc_wait()
41 +
42 + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0).to(tl.float32)
43 + w = tl.load(w_ptr + w_offs, mask, other=0.0)
44 +
45 + rrms = tl.rsqrt(tl.sum(x * x) / GROUP_DIM + EPS)
46 + # Gemma's (1 + w) affine is written this way to lower to an FMA.
47 + y = x * rrms
48 + y += y * w.to(tl.float32)
49 +
50 + if launch_pdl:
51 + tl.extra.cuda.gdc_launch_dependents()
52 + tl.store(y_ptr + row * stride_y + offsets, y, mask)
53 +
54 +
55 +def _grouped_gemma_rmsnorm(
56 + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
57 +) -> torch.Tensor:
58 + N, DIM = x.shape
59 + assert x.stride(1) == 1, "grouped Gemma RMSNorm requires unit inner stride"
60 + assert weight.is_contiguous(), "grouped Gemma RMSNorm weight must be contiguous"
61 + assert DIM % num_groups == 0
62 + group_dim = DIM // num_groups
63 + assert weight.numel() in (group_dim, DIM)
64 +
65 + y = x.new_empty(x.shape)
66 + _grouped_gemma_rmsnorm_kernel[(N * num_groups,)](
67 + x,
68 + weight,
69 + y,
70 + x.stride(0),
71 + y.stride(0),
72 + DIM,
73 + num_groups,
74 + W_SHARED=weight.numel() == group_dim,
75 + EPS=eps,
76 + launch_pdl=current_platform.is_arch_support_pdl(),
77 + )
78 + return y
79 +
80 +
81 +@triton.jit
82 +def _hc_silu_kernel(
83 + x_ptr,
84 + y_ptr,
85 + stride_x,
86 + stride_y,
87 + DIM: tl.constexpr,
88 + HC: tl.constexpr,
89 + launch_pdl: tl.constexpr,
90 +) -> None:
91 + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM)
92 +
93 + row = tl.program_id(0)
94 + offs = tl.arange(0, BLOCK_SIZE)
95 + mask = offs < DIM
96 +
97 + if launch_pdl:
98 + tl.extra.cuda.gdc_wait()
99 +
100 + x = tl.load(x_ptr + row * stride_x + offs, mask).to(tl.float32) / HC
101 + y = x * tl.sigmoid(x)
102 +
103 + if launch_pdl:
104 + tl.extra.cuda.gdc_launch_dependents()
105 + tl.store(y_ptr + row * stride_y + offs, y, mask)
106 +
107 +
108 +def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
109 + num_tokens, DIM = x.shape
110 + assert x.stride(1) == 1
111 +
112 + output = x.new_empty(x.shape)
113 + _hc_silu_kernel[(num_tokens,)](
114 + x,
115 + output,
116 + x.stride(0),
117 + output.stride(0),
118 + DIM=DIM,
119 + HC=hc_count,
120 + launch_pdl=current_platform.is_arch_support_pdl(),
121 + )
122 + return output
123 +
124 +
125 +@triton.jit
126 +def _hc_gate_mix_kernel(
127 + x_ptr,
128 + g_ptr,
129 + y_ptr,
130 + stride_x,
131 + stride_g,
132 + stride_y,
133 + DIM: tl.constexpr,
134 + HC: tl.constexpr,
135 + BLOCK_SIZE: tl.constexpr,
136 + launch_pdl: tl.constexpr,
137 +) -> None:
138 + HC_DIM: tl.constexpr = DIM // HC
139 +
140 + row = tl.program_id(0)
141 + tile_id = tl.program_id(1)
142 + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
143 + mask = offs_inner < HC_DIM
144 +
145 + if launch_pdl:
146 + tl.extra.cuda.gdc_wait()
147 +
148 + # The constexpr loop is unrolled and keeps one stream live at a time.
149 + # Materializing [HC, BLOCK_SIZE] more than doubles latency at large M.
150 + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
151 + for stream in tl.static_range(HC):
152 + offsets = stream * HC_DIM + offs_inner
153 + g = tl.load(g_ptr + row * stride_g + offsets, mask, other=0.0)
154 + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0)
155 + acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32)
156 + acc /= HC
157 +
158 + if launch_pdl:
159 + tl.extra.cuda.gdc_launch_dependents()
160 + tl.store(y_ptr + row * stride_y + offs_inner, acc, mask)
161 +
162 +
163 +def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
164 + N, DIM = gate.shape
165 + assert x.shape == gate.shape
166 + assert DIM % hc_count == 0
167 + assert x.stride(1) == 1
168 + assert gate.stride(1) == 1
169 +
170 + HC_DIM = DIM // hc_count
171 + out = x.new_empty(N, HC_DIM)
172 + BLOCK_SIZE = 512
173 + _hc_gate_mix_kernel[(N, triton.cdiv(HC_DIM, BLOCK_SIZE))](
174 + x,
175 + gate,
176 + out,
177 + x.stride(0),
178 + gate.stride(0),
179 + out.stride(0),
180 + DIM,
181 + hc_count,
182 + BLOCK_SIZE,
183 + launch_pdl=current_platform.is_arch_support_pdl(),
184 + )
185 + return out
186 +
187 +
188 +@triton.jit
189 +def _hc_combine_kernel(
190 + block_ptr,
191 + res_ptr,
192 + inj_ptr,
193 + out_ptr,
194 + stride_block,
195 + stride_res,
196 + stride_inj,
197 + stride_out,
198 + HC_DIM: tl.constexpr,
199 + HC: tl.constexpr,
200 + BLOCK_SIZE: tl.constexpr,
201 + launch_pdl: tl.constexpr,
202 +) -> None:
203 + HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
204 +
205 + row = tl.program_id(0)
206 + tile_id = tl.program_id(1)
207 +
208 + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
209 + mask_inner = offs_inner < HC_DIM
210 + offs_hc = tl.arange(0, HC_PAD)
211 + mask_hc = offs_hc < HC
212 + offs = offs_hc[:, None] * HC_DIM + offs_inner[None, :]
213 + mask = mask_hc[:, None] & mask_inner[None, :]
214 +
215 + if launch_pdl:
216 + tl.extra.cuda.gdc_wait()
217 +
218 + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)
219 + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)
220 + res = tl.load(res_ptr + row * stride_res + offs, mask, other=0.0)
221 +
222 + # Keeping HC as a broadcast dimension is faster here than four separate
223 + # residual load/store sequences.
224 + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)
225 + out = res.to(tl.float32) + block.to(tl.float32)[None, :] * inj[:, None]
226 +
227 + if launch_pdl:
228 + tl.extra.cuda.gdc_launch_dependents()
229 + tl.store(out_ptr + row * stride_out + offs, out, mask=mask)
230 +
231 +
232 +def _hc_combine(
233 + residual: torch.Tensor,
234 + block_output: torch.Tensor,
235 + injection_logits: torch.Tensor,
236 + hc_count: int,
237 +) -> torch.Tensor:
238 + N, DIM = residual.shape
239 + assert DIM % hc_count == 0
240 + hc_dim = DIM // hc_count
241 + assert block_output.shape == (N, hc_dim)
242 + assert injection_logits.shape == (N, hc_count)
243 + assert residual.stride(1) == 1
244 + assert block_output.stride(1) == 1
245 + assert injection_logits.stride(1) == 1
246 +
247 + out = residual.new_empty(residual.shape)
248 + BLOCK_SIZE = 512
249 + _hc_combine_kernel[(N, triton.cdiv(hc_dim, BLOCK_SIZE))](
250 + block_output,
251 + residual,
252 + injection_logits,
253 + out,
254 + block_output.stride(0),
255 + residual.stride(0),
256 + injection_logits.stride(0),
257 + out.stride(0),
258 + hc_dim,
259 + hc_count,
260 + BLOCK_SIZE,
261 + launch_pdl=current_platform.is_arch_support_pdl(),
262 + )
263 + return out
264 +
265 +
266 +@triton.jit
267 +def _hc_combine_norm_kernel(
268 + block_ptr,
269 + res_ptr,
270 + inj_ptr,
271 + w_ptr,
272 + out_ptr,
273 + y_ptr,
274 + stride_block,
275 + stride_res,
276 + stride_inj,
277 + stride_out,
278 + stride_y,
279 + HC_DIM: tl.constexpr,
280 + HC: tl.constexpr,
281 + W_SHARED: tl.constexpr,
282 + EPS: tl.constexpr,
283 + BLOCK_SIZE: tl.constexpr,
284 + launch_pdl: tl.constexpr,
285 +) -> None:
286 + HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
287 + NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE)
288 + NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES)
289 +
290 + row = tl.program_id(0)
291 + stream = tl.program_id(1)
292 + offs_hc = tl.arange(0, HC_PAD)
293 + mask_hc = offs_hc < HC
294 + tile_ids = tl.arange(0, NUM_TILES_PAD)
295 + offs_inner = tile_ids[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :]
296 + mask_inner = offs_inner < HC_DIM
297 + offs = stream * HC_DIM + offs_inner
298 + # Shared norm weights repeat across streams; per-branch weights use the
299 + # same flattened HC layout as the residual.
300 + w_offs = offs_inner if W_SHARED else offs
301 +
302 + if launch_pdl:
303 + tl.extra.cuda.gdc_wait()
304 +
305 + # Start the uncached residual load first, then issue the other combine
306 + # loads before consuming any of them.
307 + res = tl.load(res_ptr + row * stride_res + offs, mask_inner, other=0.0)
308 + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)
309 + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)
310 + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)
311 + inj = tl.sum(tl.where(offs_hc == stream, inj, 0.0))
312 + # Round the materialized combine result before normalization. This matches
313 + # the unfused combine -> RMSNorm boundary.
314 + out = (res.to(tl.float32) + block.to(tl.float32) * inj).to(out_ptr.dtype.element_ty)
315 + tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner)
316 +
317 + out = out.to(tl.float32)
318 + # Keep the two-axis reduction: flattening the padded tile is ~40% slower
319 + # at decode sizes.
320 + sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0)
321 + rrms = tl.rsqrt(sum_sq / HC_DIM + EPS)
322 +
323 + if launch_pdl:
324 + tl.extra.cuda.gdc_launch_dependents()
325 +
326 + # Loading the weight earlier helps decode but keeps the tile live across
327 + # the reduction and regresses larger batches, so defer it to the norm.
328 + w = tl.load(w_ptr + w_offs, mask_inner, other=0.0)
329 + y = out * rrms
330 + y += y * w.to(tl.float32)
331 + tl.store(y_ptr + row * stride_y + offs, y, mask_inner)
332 +
333 +
334 +def _hc_combine_norm(
335 + residual: torch.Tensor,
336 + block_output: torch.Tensor,
337 + injection_logits: torch.Tensor,
338 + norm_weight: torch.Tensor,
339 + eps: float,
340 + hc_count: int,
341 +) -> tuple[torch.Tensor, torch.Tensor]:
342 + N, DIM = residual.shape
343 + assert DIM % hc_count == 0
344 + hc_dim = DIM // hc_count
345 + assert block_output.shape == (N, hc_dim)
346 + assert injection_logits.shape == (N, hc_count)
347 + assert residual.stride(1) == 1
348 + assert block_output.stride(1) == 1
349 + assert injection_logits.stride(1) == 1
350 + assert norm_weight.is_contiguous()
351 + assert norm_weight.numel() in (hc_dim, DIM)
352 +
353 + out = residual.new_empty(residual.shape)
354 + y = residual.new_empty(residual.shape)
355 + BLOCK_SIZE = 512
356 + _hc_combine_norm_kernel[(N, hc_count)](
357 + block_output,
358 + residual,
359 + injection_logits,
360 + norm_weight,
361 + out,
362 + y,
363 + block_output.stride(0),
364 + residual.stride(0),
365 + injection_logits.stride(0),
366 + out.stride(0),
367 + y.stride(0),
368 + hc_dim,
369 + hc_count,
370 + W_SHARED=norm_weight.numel() == hc_dim,
371 + EPS=eps,
372 + BLOCK_SIZE=BLOCK_SIZE,
373 + launch_pdl=current_platform.is_arch_support_pdl(),
374 + )
375 + return out, y
376 +
377 +
378 +def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor:
379 + return x.new_empty(x.shape)
380 +
381 +
382 +def _hc_gate_mix_fake(
383 + x: torch.Tensor, gate: torch.Tensor, hc_count: int
384 +) -> torch.Tensor:
385 + del gate
386 + return x.new_empty((x.shape[0], x.shape[1] // hc_count))
387 +
388 +
389 +def _hc_combine_fake(
390 + residual: torch.Tensor,
391 + block_output: torch.Tensor,
392 + injection_logits: torch.Tensor,
393 + hc_count: int,
394 +) -> torch.Tensor:
395 + del block_output, injection_logits, hc_count
396 + return residual.new_empty(residual.shape)
397 +
398 +
399 +def _hc_combine_norm_fake(
400 + residual: torch.Tensor,
401 + block_output: torch.Tensor,
402 + injection_logits: torch.Tensor,
403 + norm_weight: torch.Tensor,
404 + eps: float,
405 + hc_count: int,
406 +) -> tuple[torch.Tensor, torch.Tensor]:
407 + del block_output, injection_logits, norm_weight, eps, hc_count
408 + return residual.new_empty(residual.shape), residual.new_empty(residual.shape)
409 +
410 +
411 +direct_register_custom_op(
412 + op_name="qwen4_exp_grouped_gemma_rmsnorm",
413 + op_func=_grouped_gemma_rmsnorm,
414 + fake_impl=_same_shape_fake,
415 +)
416 +direct_register_custom_op(
417 + op_name="qwen4_exp_hc_silu",
418 + op_func=_hc_silu,
419 + fake_impl=_same_shape_fake,
420 +)
421 +direct_register_custom_op(
422 + op_name="qwen4_exp_hc_gate_mix",
423 + op_func=_hc_gate_mix,
424 + fake_impl=_hc_gate_mix_fake,
425 +)
426 +direct_register_custom_op(
427 + op_name="qwen4_exp_hc_combine",
428 + op_func=_hc_combine,
429 + fake_impl=_hc_combine_fake,
430 +)
431 +direct_register_custom_op(
432 + op_name="qwen4_exp_hc_combine_norm",
433 + op_func=_hc_combine_norm,
434 + fake_impl=_hc_combine_norm_fake,
435 +)
436 +
437 +
438 +def grouped_gemma_rmsnorm(
439 + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
440 +) -> torch.Tensor:
441 + return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups)
442 +
443 +
444 +def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
445 + return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count)
446 +
447 +
448 +def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
449 + return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count)
450 +
451 +
452 +def hc_combine(
453 + residual: torch.Tensor,
454 + block_output: torch.Tensor,
455 + injection_logits: torch.Tensor,
456 + hc_count: int,
457 +) -> torch.Tensor:
458 + return torch.ops.vllm.qwen4_exp_hc_combine(
459 + residual, block_output, injection_logits, hc_count
460 + )
461 +
462 +
463 +def hc_combine_norm(
464 + residual: torch.Tensor,
465 + block_output: torch.Tensor,
466 + injection_logits: torch.Tensor,
467 + norm_weight: torch.Tensor,
468 + eps: float,
469 + hc_count: int,
470 +) -> tuple[torch.Tensor, torch.Tensor]:
471 + return torch.ops.vllm.qwen4_exp_hc_combine_norm(
472 + residual,
473 + block_output,
474 + injection_logits,
475 + norm_weight,
476 + eps,
477 + hc_count,
478 + )
479 +
480 +
481 +__all__ = [
482 + "grouped_gemma_rmsnorm",
483 + "hc_combine",
484 + "hc_combine_norm",
485 + "hc_gate_mix",
486 + "hc_silu",
487 +]
vllm/models/qwen4_exp/amd/ops/qsa.py new
+1128
@@ -0,0 +1,1128 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton kernels for the Qwen4Exp weight-free QSA path."""
4 +
5 +from __future__ import annotations
6 +
7 +import math
8 +
9 +import torch
10 +
11 +from vllm import _custom_ops as ops
12 +from vllm.platforms import current_platform
13 +from vllm.triton_utils import HAS_TRITON, tl, triton
14 +
15 +_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024
16 +_TOPK_WORKSPACE_BYTES = 1024 * 1024
17 +
18 +
19 +@triton.jit
20 +def _qsa_mqa_paged_kernel(
21 + q_ptr,
22 + k_cache_ptr,
23 + page_table_ptr,
24 + token_to_req_ptr,
25 + query_positions_ptr,
26 + sequence_lengths_ptr,
27 + visible_blocks_ptr,
28 + logits_ptr,
29 + stride_q_row,
30 + stride_q_head,
31 + stride_q_dim,
32 + stride_cache_block,
33 + stride_cache_token,
34 + stride_cache_dim,
35 + stride_table_req,
36 + stride_table_page,
37 + stride_logits_row,
38 + num_rows,
39 + num_columns,
40 + num_pages,
41 + num_requests,
42 + score_divisor,
43 + PAGE_SIZE: tl.constexpr,
44 + PAGE_TABLE_WIDTH: tl.constexpr,
45 + NUM_HEADS: tl.constexpr,
46 + HEAD_DIM: tl.constexpr,
47 + BLOCK_N: tl.constexpr,
48 + BLOCK_D: tl.constexpr,
49 + COMPRESS_RATIO: tl.constexpr,
50 +) -> None:
51 + row = tl.program_id(0)
52 + columns = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
53 + dims = tl.arange(0, BLOCK_D)
54 + request = tl.load(token_to_req_ptr + row)
55 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
56 + query_position = tl.load(query_positions_ptr + row)
57 + sequence_length = tl.load(
58 + sequence_lengths_ptr + safe_request,
59 + mask=(request >= 0) & (request < num_requests),
60 + other=0,
61 + )
62 + visible = tl.minimum(
63 + (query_position + 1) // COMPRESS_RATIO,
64 + sequence_length // COMPRESS_RATIO,
65 + )
66 + if tl.program_id(1) == 0:
67 + tl.store(visible_blocks_ptr + row, visible)
68 + logical_page = columns // PAGE_SIZE
69 + page_offset = columns % PAGE_SIZE
70 + valid = (
71 + (row < num_rows)
72 + & (columns < num_columns)
73 + & (columns < visible)
74 + & (request >= 0)
75 + & (request < num_requests)
76 + & (logical_page < PAGE_TABLE_WIDTH)
77 + )
78 + safe_logical_page = tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1)
79 + physical_page = tl.load(
80 + page_table_ptr
81 + + safe_request * stride_table_req
82 + + safe_logical_page * stride_table_page,
83 + mask=valid,
84 + other=-1,
85 + )
86 + valid &= (physical_page >= 0) & (physical_page < num_pages)
87 + # physical_page * block stride can overflow int32 for large caches.
88 + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64)
89 + score = tl.zeros((BLOCK_N,), dtype=tl.float32)
90 +
91 + for head in tl.static_range(0, NUM_HEADS):
92 + query = tl.load(
93 + q_ptr + row * stride_q_row + head * stride_q_head + dims * stride_q_dim,
94 + mask=dims < HEAD_DIM,
95 + other=0.0,
96 + ).to(tl.float32)
97 + keys = tl.load(
98 + k_cache_ptr
99 + + safe_physical_page[:, None] * stride_cache_block
100 + + page_offset[:, None] * stride_cache_token
101 + + dims[None, :] * stride_cache_dim,
102 + mask=valid[:, None] & (dims[None, :] < HEAD_DIM),
103 + other=0.0,
104 + ).to(tl.float32)
105 + dot = tl.sum(keys * query[None, :], axis=1)
106 + score += tl.maximum(dot, 0.0)
107 +
108 + score /= score_divisor
109 + tl.store(
110 + logits_ptr + row * stride_logits_row + columns,
111 + tl.where(valid, score, -float("inf")),
112 + mask=(row < num_rows) & (columns < num_columns),
113 + )
114 +
115 +
116 +@triton.jit
117 +def _expand_qsa_indices_kernel(
118 + block_indices_ptr,
119 + query_positions_ptr,
120 + sequence_lengths_ptr,
121 + token_to_req_ptr,
122 + output_ptr,
123 + stride_blocks_row,
124 + stride_blocks_column,
125 + stride_output_row,
126 + stride_output_column,
127 + rows,
128 + num_requests,
129 + BLOCK_TOPK: tl.constexpr,
130 + COMPRESS_RATIO: tl.constexpr,
131 + TOKEN_TOPK: tl.constexpr,
132 + OUTPUT_WIDTH: tl.constexpr,
133 + COLUMN_BLOCK: tl.constexpr,
134 +) -> None:
135 + row = tl.program_id(0)
136 + columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK)
137 + query_position = tl.load(query_positions_ptr + row)
138 + request = tl.load(token_to_req_ptr + row)
139 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
140 + sequence_length = tl.load(
141 + sequence_lengths_ptr + safe_request,
142 + mask=(request >= 0) & (request < num_requests),
143 + other=0,
144 + )
145 + complete_blocks = tl.minimum(
146 + tl.minimum(
147 + (query_position + 1) // COMPRESS_RATIO,
148 + sequence_length // COMPRESS_RATIO,
149 + ),
150 + BLOCK_TOPK,
151 + )
152 + expanded_count = complete_blocks * COMPRESS_RATIO
153 + tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO
154 + tail_count = (query_position + 1) - tail_start
155 +
156 + is_expanded = columns < expanded_count
157 + block_rank = columns // COMPRESS_RATIO
158 + offset = columns % COMPRESS_RATIO
159 + safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1)
160 + block = tl.load(
161 + block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column,
162 + mask=(row < rows) & is_expanded,
163 + other=-1,
164 + )
165 + expanded = block * COMPRESS_RATIO + offset
166 + tail_offset = columns - expanded_count
167 + is_tail = (
168 + (columns >= expanded_count)
169 + & (tail_offset < tail_count)
170 + & (tail_offset < COMPRESS_RATIO - 1)
171 + )
172 + token = tl.where(is_expanded, expanded, tail_start + tail_offset)
173 + valid = (
174 + (row < rows)
175 + & (columns < OUTPUT_WIDTH)
176 + & (is_expanded | is_tail)
177 + & (token >= 0)
178 + & (token < sequence_length)
179 + )
180 + tl.store(
181 + output_ptr + row * stride_output_row + columns * stride_output_column,
182 + tl.where(valid, token, -1),
183 + mask=(row < rows) & (columns < OUTPUT_WIDTH),
184 + )
185 +
186 +
187 +@triton.jit
188 +def _qsa_sparse_paged_gqa_splitk_kernel(
189 + q_ptr,
190 + k_cache_ptr,
191 + v_cache_ptr,
192 + indices_ptr,
193 + block_table_ptr,
194 + token_to_req_ptr,
195 + partial_output_ptr,
196 + partial_lse_ptr,
197 + output_ptr,
198 + stride_q_row,
199 + stride_q_head,
200 + stride_k_block,
201 + stride_k_token,
202 + stride_k_head,
203 + stride_v_block,
204 + stride_v_token,
205 + stride_v_head,
206 + stride_indices_row,
207 + stride_table_req,
208 + stride_output_row,
209 + stride_output_head,
210 + num_rows,
211 + num_cache_blocks,
212 + num_requests,
213 + TOPK: tl.constexpr,
214 + PAGE_SIZE: tl.constexpr,
215 + PAGE_TABLE_WIDTH: tl.constexpr,
216 + GROUP_SIZE: tl.constexpr,
217 + HEAD_DIM: tl.constexpr,
218 + NUM_QUERY_HEADS: tl.constexpr,
219 + NUM_SPLITS: tl.constexpr,
220 + NUM_TILES: tl.constexpr,
221 + BLOCK_M: tl.constexpr,
222 + BLOCK_N: tl.constexpr,
223 +) -> None:
224 + row = tl.program_id(0)
225 + kv_head = tl.program_id(1)
226 + split_id = tl.program_id(2)
227 + request = tl.load(token_to_req_ptr + row)
228 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
229 +
230 + head_offsets = tl.arange(0, BLOCK_M)
231 + dim_offsets = tl.arange(0, HEAD_DIM)
232 + column_offsets = tl.arange(0, BLOCK_N)
233 + first_head = kv_head * GROUP_SIZE
234 + query = tl.load(
235 + q_ptr
236 + + row * stride_q_row
237 + + (first_head + head_offsets[:, None]) * stride_q_head
238 + + dim_offsets[None, :],
239 + mask=head_offsets[:, None] < GROUP_SIZE,
240 + other=0.0,
241 + )
242 +
243 + max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32)
244 + normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32)
245 + accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)
246 + softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634
247 +
248 + # Dynamic bounds avoid padded main-loop iterations for uneven splits.
249 + split_tile_start = split_id * NUM_TILES // NUM_SPLITS
250 + split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS
251 + for tile in range(split_tile_start, split_tile_end):
252 + columns = tile * BLOCK_N + column_offsets
253 + logical_token = tl.load(
254 + indices_ptr + row * stride_indices_row + columns,
255 + mask=columns < TOPK,
256 + other=-1,
257 + )
258 + safe_token = tl.maximum(logical_token, 0)
259 + logical_page = safe_token // PAGE_SIZE
260 + page_offset = safe_token % PAGE_SIZE
261 + valid = (
262 + (request >= 0)
263 + & (request < num_requests)
264 + & (logical_token >= 0)
265 + & (logical_page < PAGE_TABLE_WIDTH)
266 + )
267 + physical_page = tl.load(
268 + block_table_ptr
269 + + safe_request * stride_table_req
270 + + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1),
271 + mask=valid,
272 + other=-1,
273 + )
274 + valid &= (physical_page >= 0) & (physical_page < num_cache_blocks)
275 + # physical_page * block stride can overflow int32 for large caches.
276 + safe_page = tl.maximum(physical_page, 0).to(tl.int64)
277 + keys = tl.load(
278 + k_cache_ptr
279 + + safe_page[None, :] * stride_k_block
280 + + page_offset[None, :] * stride_k_token
281 + + kv_head * stride_k_head
282 + + dim_offsets[:, None],
283 + mask=valid[None, :],
284 + other=0.0,
285 + )
286 + values = tl.load(
287 + v_cache_ptr
288 + + safe_page[:, None] * stride_v_block
289 + + page_offset[:, None] * stride_v_token
290 + + kv_head * stride_v_head
291 + + dim_offsets[None, :],
292 + mask=valid[:, None],
293 + other=0.0,
294 + )
295 + scores = tl.dot(query, keys)
296 + # Scaling scores avoids re-quantizing a scaled query to BF16.
297 + scores *= softmax_scale_log2
298 + scores = tl.where(valid[None, :], scores, -1.0e20)
299 + next_max = tl.maximum(max_value, tl.max(scores, axis=1))
300 + alpha = tl.math.exp2(max_value - next_max)
301 + probabilities = tl.where(
302 + valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0
303 + )
304 + accumulator = tl.dot(
305 + probabilities.to(values.dtype),
306 + values,
307 + acc=accumulator * alpha[:, None],
308 + )
309 + normalizer = normalizer * alpha + tl.sum(probabilities, axis=1)
310 + max_value = next_max
311 +
312 + has_values = normalizer > 0
313 + normalized_output = tl.where(
314 + has_values[:, None],
315 + accumulator / tl.maximum(normalizer[:, None], 1.0e-20),
316 + 0.0,
317 + )
318 + output_mask = head_offsets[:, None] < GROUP_SIZE
319 + if NUM_SPLITS == 1:
320 + tl.store(
321 + output_ptr
322 + + row * stride_output_row
323 + + (first_head + head_offsets[:, None]) * stride_output_head
324 + + dim_offsets[None, :],
325 + normalized_output,
326 + mask=output_mask,
327 + )
328 + else:
329 + partial_lse = tl.where(
330 + has_values,
331 + max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)),
332 + -float("inf"),
333 + )
334 + tl.store(
335 + partial_output_ptr
336 + + (
337 + (split_id * num_rows + row) * NUM_QUERY_HEADS
338 + + first_head
339 + + head_offsets[:, None]
340 + )
341 + * HEAD_DIM
342 + + dim_offsets[None, :],
343 + normalized_output,
344 + mask=output_mask,
345 + )
346 + tl.store(
347 + partial_lse_ptr
348 + + (split_id * num_rows + row) * NUM_QUERY_HEADS
349 + + first_head
350 + + head_offsets,
351 + partial_lse,
352 + mask=head_offsets < GROUP_SIZE,
353 + )
354 +
355 +
356 +@triton.jit
357 +def _qsa_merge_splitk_kernel(
358 + partial_output_ptr,
359 + partial_lse_ptr,
360 + output_ptr,
361 + stride_output_row,
362 + stride_output_head,
363 + num_rows,
364 + HEAD_DIM: tl.constexpr,
365 + NUM_QUERY_HEADS: tl.constexpr,
366 + NUM_SPLITS: tl.constexpr,
367 + BLOCK_SPLITS: tl.constexpr,
368 +) -> None:
369 + row = tl.program_id(0)
370 + head = tl.program_id(1)
371 + split_offsets = tl.arange(0, BLOCK_SPLITS)
372 + dim_offsets = tl.arange(0, HEAD_DIM)
373 + split_mask = split_offsets < NUM_SPLITS
374 + lse = tl.load(
375 + partial_lse_ptr + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head,
376 + mask=split_mask,
377 + other=-float("inf"),
378 + )
379 + lse_max = tl.max(lse, axis=0)
380 + has_values = lse_max > -float("inf")
381 + shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf"))
382 + weights = tl.math.exp2(shifted)
383 + denominator = tl.sum(weights, axis=0)
384 + partial_output = tl.load(
385 + partial_output_ptr
386 + + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head)
387 + * HEAD_DIM
388 + + dim_offsets[None, :],
389 + mask=split_mask[:, None],
390 + other=0.0,
391 + )
392 + merged = tl.sum(partial_output * weights[:, None], axis=0)
393 + merged = tl.where(denominator > 0, merged / denominator, 0.0)
394 + tl.store(
395 + output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets,
396 + merged,
397 + )
398 +
399 +
400 +@triton.jit
401 +def _store_qsa_rows_kernel(
402 + cache_ptr,
403 + slots_ptr,
404 + rows_ptr,
405 + stride_cache_block,
406 + stride_cache_token,
407 + stride_cache_dim,
408 + stride_rows_row,
409 + stride_rows_dim,
410 + num_rows,
411 + num_blocks,
412 + PAGE_SIZE: tl.constexpr,
413 + WIDTH: tl.constexpr,
414 + BLOCK_D: tl.constexpr,
415 +) -> None:
416 + row = tl.program_id(0)
417 + dims = tl.arange(0, BLOCK_D)
418 + slot = tl.load(slots_ptr + row)
419 + valid = (row < num_rows) & (slot >= 0) & (slot < num_blocks * PAGE_SIZE)
420 + block = tl.maximum(slot, 0) // PAGE_SIZE
421 + token = tl.maximum(slot, 0) % PAGE_SIZE
422 + values = tl.load(
423 + rows_ptr + row * stride_rows_row + dims * stride_rows_dim,
424 + mask=valid & (dims < WIDTH),
425 + other=0,
426 + )
427 + tl.store(
428 + cache_ptr
429 + + block * stride_cache_block
430 + + token * stride_cache_token
431 + + dims * stride_cache_dim,
432 + values,
433 + mask=valid & (dims < WIDTH),
434 + )
435 +
436 +
437 +@triton.jit
438 +def _compress_qsa_groups_kernel(
439 + raw_keys_ptr, # this step's raw key rows, straight from activations
440 + raw_positions_ptr, # this step's per-token positions
441 + compressor_state_cache_ptr, # per-request ring of previous raw keys
442 + rope_cache_ptr, # packed RoPE position tail of the ring
443 + compressor_state_table_ptr,
444 + token_to_req_ptr,
445 + query_start_loc_ptr,
446 + logical_positions_ptr,
447 + compressed_slots_ptr,
448 + pooled_ptr,
449 + first_positions_ptr,
450 + stride_raw_row,
451 + stride_raw_dim,
452 + stride_raw_positions_row,
453 + stride_raw_positions_dim,
454 + stride_compressor_state_block,
455 + stride_compressor_state_token,
456 + stride_compressor_state_dim,
457 + stride_rope_block,
458 + stride_rope_token,
459 + stride_rope_dim,
460 + stride_compressor_state_table_req,
461 + stride_pooled_row,
462 + stride_pooled_dim,
463 + stride_positions_row,
464 + stride_positions_dim,
465 + num_rows,
466 + num_compressor_state_blocks,
467 + num_requests,
468 + COMPRESSOR_STATE_SIZE: tl.constexpr,
469 + COMPRESS_RATIO: tl.constexpr,
470 + HEAD_DIM: tl.constexpr,
471 + BLOCK_D: tl.constexpr,
472 + LOAD_ROPE_POSITIONS: tl.constexpr,
473 +) -> None:
474 + row = tl.program_id(0)
475 + dims = tl.arange(0, BLOCK_D)
476 + request = tl.load(token_to_req_ptr + row)
477 + end_position = tl.load(logical_positions_ptr + row)
478 + compressed_slot = tl.load(compressed_slots_ptr + row)
479 + valid_request = (request >= 0) & (request < num_requests)
480 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
481 + query_row_start = tl.load(
482 + query_start_loc_ptr + safe_request, mask=valid_request, other=0
483 + )
484 + query_row_end = tl.load(
485 + query_start_loc_ptr + safe_request + 1, mask=valid_request, other=0
486 + )
487 + chunk_start_position = end_position - (row - query_row_start)
488 + compressor_state_block = tl.load(
489 + compressor_state_table_ptr + safe_request * stride_compressor_state_table_req,
490 + mask=valid_request,
491 + other=-1,
492 + )
493 + valid_compressor_state_block = (compressor_state_block >= 0) & (
494 + compressor_state_block < num_compressor_state_blocks
495 + )
496 + valid_row = (
497 + (row < num_rows)
498 + & valid_request
499 + & (row >= query_row_start)
500 + & (row < query_row_end)
501 + & (end_position >= COMPRESS_RATIO - 1)
502 + & (compressed_slot >= 0)
503 + )
504 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32)
505 +
506 + # A group can span the compressor-state ring (older members) and this
507 + # step's raw rows (members at positions >= chunk_start_position).
508 + for group_offset in tl.range(0, COMPRESS_RATIO):
509 + position = end_position - (COMPRESS_RATIO - 1 - group_offset)
510 + use_raw = position >= chunk_start_position
511 + raw_row = query_row_start + position - chunk_start_position
512 + raw_values = tl.load(
513 + raw_keys_ptr + raw_row * stride_raw_row + dims * stride_raw_dim,
514 + mask=valid_row
515 + & use_raw
516 + & (raw_row >= query_row_start)
517 + & (raw_row < query_row_end)
518 + & (raw_row < num_rows)
519 + & (dims < HEAD_DIM),
520 + other=0.0,
521 + ).to(tl.float32)
522 + compressor_state_values = tl.load(
523 + compressor_state_cache_ptr
524 + + tl.maximum(compressor_state_block, 0).to(tl.int64)
525 + * stride_compressor_state_block
526 + + (position % COMPRESSOR_STATE_SIZE) * stride_compressor_state_token
527 + + dims * stride_compressor_state_dim,
528 + mask=valid_row
529 + & ~use_raw
530 + & valid_compressor_state_block
531 + & (dims < HEAD_DIM),
532 + other=0.0,
533 + ).to(tl.float32)
534 + accumulator += tl.where(use_raw, raw_values, compressor_state_values)
535 +
536 + tl.store(
537 + pooled_ptr + row * stride_pooled_row + dims * stride_pooled_dim,
538 + accumulator / COMPRESS_RATIO,
539 + mask=(row < num_rows) & (dims < HEAD_DIM),
540 + )
541 +
542 + position_dims = tl.arange(0, 4)
543 + first_position = end_position - COMPRESS_RATIO + 1
544 + if LOAD_ROPE_POSITIONS:
545 + first_from_raw = first_position >= chunk_start_position
546 + raw_first_row = query_row_start + first_position - chunk_start_position
547 + raw_position_values = tl.load(
548 + raw_positions_ptr
549 + + raw_first_row * stride_raw_positions_row
550 + + position_dims * stride_raw_positions_dim,
551 + mask=valid_row
552 + & first_from_raw
553 + & (raw_first_row >= query_row_start)
554 + & (raw_first_row < query_row_end)
555 + & (raw_first_row < num_rows)
556 + & (position_dims < 3),
557 + other=0,
558 + )
559 + compressor_state_position_values = tl.load(
560 + rope_cache_ptr
561 + + tl.maximum(compressor_state_block, 0).to(tl.int64) * stride_rope_block
562 + + (first_position % COMPRESSOR_STATE_SIZE) * stride_rope_token
563 + + position_dims * stride_rope_dim,
564 + mask=valid_row
565 + & ~first_from_raw
566 + & valid_compressor_state_block
567 + & (position_dims < 3),
568 + other=0,
569 + )
570 + position_values = tl.where(
571 + first_from_raw,
572 + raw_position_values,
573 + compressor_state_position_values,
574 + )
575 + else:
576 + position_values = tl.where(valid_row, first_position, 0)
577 + tl.store(
578 + first_positions_ptr
579 + + row * stride_positions_row
580 + + position_dims * stride_positions_dim,
581 + position_values,
582 + mask=(row < num_rows) & (position_dims < 3),
583 + )
584 +
585 +
586 +def _validate_mqa(q: torch.Tensor) -> None:
587 + if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0:
588 + raise ValueError("QSA query must be [rows, heads, head_dim]")
589 +
590 +
591 +def qsa_mqa_paged(
592 + q: torch.Tensor,
593 + k_cache: torch.Tensor,
594 + page_table: torch.Tensor,
595 + token_to_req: torch.Tensor,
596 + query_positions: torch.Tensor,
597 + sequence_lengths: torch.Tensor,
598 + compress_ratio: int,
599 + num_columns: int | None = None,
600 + score_scale: float | None = None,
601 +) -> tuple[torch.Tensor, torch.Tensor]:
602 + """Compute QSA scores directly from a paged compressed-key cache."""
603 +
604 + _validate_mqa(q)
605 + if not q.is_cuda or not HAS_TRITON:
606 + raise RuntimeError("paged QSA scoring requires a GPU and Triton")
607 + if k_cache.ndim != 4 or k_cache.shape[2] != 1:
608 + raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]")
609 + if k_cache.shape[3] != q.shape[2]:
610 + raise ValueError("QSA query and cache dimensions must match")
611 + if page_table.ndim != 2:
612 + raise ValueError("QSA page table must be two-dimensional")
613 + if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)):
614 + raise ValueError("QSA paged scoring cache and page table must be nonempty")
615 + if token_to_req.shape != (q.shape[0],):
616 + raise ValueError("QSA request mapping must match query rows")
617 + if query_positions.shape != (q.shape[0],):
618 + raise ValueError("QSA query positions must match query rows")
619 + if sequence_lengths.shape != (page_table.shape[0],):
620 + raise ValueError("QSA sequence lengths must match page-table requests")
621 + if compress_ratio <= 0:
622 + raise ValueError("QSA compression ratio must be positive")
623 + score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale
624 + if score_divisor <= 0:
625 + raise ValueError("QSA score scale must be positive")
626 +
627 + capacity = page_table.shape[1] * k_cache.shape[1]
628 + columns = capacity if num_columns is None else num_columns
629 + if columns < 0:
630 + raise ValueError("QSA score width must be non-negative")
631 + logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device)
632 + visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device)
633 + if not q.shape[0] or not columns:
634 + return logits, visible_blocks
635 + block_n = 32
636 + _qsa_mqa_paged_kernel[(q.shape[0], triton.cdiv(columns, block_n))](
637 + q,
638 + k_cache,
639 + page_table,
640 + token_to_req,
641 + query_positions,
642 + sequence_lengths,
643 + visible_blocks,
644 + logits,
645 + q.stride(0),
646 + q.stride(1),
647 + q.stride(2),
648 + k_cache.stride(0),
649 + k_cache.stride(1),
650 + k_cache.stride(3),
651 + page_table.stride(0),
652 + page_table.stride(1),
653 + logits.stride(0),
654 + q.shape[0],
655 + columns,
656 + k_cache.shape[0],
657 + page_table.shape[0],
658 + float(score_divisor),
659 + PAGE_SIZE=k_cache.shape[1],
660 + PAGE_TABLE_WIDTH=page_table.shape[1],
661 + NUM_HEADS=q.shape[1],
662 + HEAD_DIM=q.shape[2],
663 + BLOCK_N=block_n,
664 + BLOCK_D=triton.next_power_of_2(q.shape[2]),
665 + COMPRESS_RATIO=compress_ratio,
666 + num_warps=4,
667 + )
668 + return logits, visible_blocks
669 +
670 +
671 +def expand_qsa_block_indices_cuda(
672 + block_indices: torch.Tensor,
673 + query_positions: torch.Tensor,
674 + sequence_lengths: torch.Tensor,
675 + token_to_req: torch.Tensor,
676 + compress_ratio: int,
677 + token_topk: int,
678 + out: torch.Tensor | None = None,
679 +) -> torch.Tensor:
680 + """Expand compressed blocks and compact the causal tail of the open group."""
681 +
682 + if not block_indices.is_cuda or not HAS_TRITON:
683 + raise RuntimeError("QSA index expansion requires a GPU and Triton")
684 + if token_topk % compress_ratio:
685 + raise ValueError("QSA token top-k must be divisible by compression ratio")
686 + block_topk = token_topk // compress_ratio
687 + output_width = token_topk + compress_ratio - 1
688 + if block_indices.shape != (query_positions.numel(), block_topk):
689 + raise ValueError("QSA compressed top-k has an invalid shape")
690 + if token_to_req.shape != query_positions.shape:
691 + raise ValueError("QSA request mapping must match query positions")
692 + if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]:
693 + raise ValueError("QSA request sequence lengths must be nonempty")
694 + if out is None:
695 + out = torch.empty(
696 + (block_indices.shape[0], output_width),
697 + dtype=torch.int32,
698 + device=block_indices.device,
699 + )
700 + elif out.shape != (block_indices.shape[0], output_width):
701 + raise ValueError("QSA expansion output has an invalid shape")
702 + if not block_indices.shape[0]:
703 + return out
704 + column_block = 256
705 + _expand_qsa_indices_kernel[
706 + (block_indices.shape[0], triton.cdiv(output_width, column_block))
707 + ](
708 + block_indices,
709 + query_positions,
710 + sequence_lengths,
711 + token_to_req,
712 + out,
713 + block_indices.stride(0),
714 + block_indices.stride(1),
715 + out.stride(0),
716 + out.stride(1),
717 + block_indices.shape[0],
718 + sequence_lengths.shape[0],
719 + BLOCK_TOPK=block_topk,
720 + COMPRESS_RATIO=compress_ratio,
721 + TOKEN_TOPK=token_topk,
722 + OUTPUT_WIDTH=output_width,
723 + COLUMN_BLOCK=column_block,
724 + num_warps=4,
725 + )
726 + return out
727 +
728 +
729 +def qsa_select_paged_tokens(
730 + q: torch.Tensor,
731 + k_cache: torch.Tensor,
732 + page_table: torch.Tensor,
733 + token_to_req: torch.Tensor,
734 + query_positions: torch.Tensor,
735 + sequence_lengths: torch.Tensor,
736 + token_topk: int,
737 + compress_ratio: int,
738 + out: torch.Tensor | None = None,
739 +) -> torch.Tensor:
740 + """Score, select, and expand QSA indices without host synchronization."""
741 +
742 + rows = q.shape[0]
743 + output_width = token_topk + compress_ratio - 1
744 + if out is None:
745 + out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device)
746 + if out.shape != (rows, output_width):
747 + raise ValueError("QSA selection output has an invalid shape")
748 + if not rows:
749 + return out
750 +
751 + columns = page_table.shape[1] * k_cache.shape[1]
752 + block_topk = token_topk // compress_ratio
753 + rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1))
754 + chunk_rows = min(rows, rows_per_chunk)
755 + blocks_buffer = torch.empty(
756 + (chunk_rows, block_topk), dtype=torch.int32, device=q.device
757 + )
758 + topk_workspace = torch.empty(
759 + (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device
760 + )
761 + for row_start in range(0, rows, rows_per_chunk):
762 + row_end = min(row_start + rows_per_chunk, rows)
763 + row_slice = slice(row_start, row_end)
764 + logits, visible_blocks = qsa_mqa_paged(
765 + q[row_slice],
766 + k_cache,
767 + page_table,
768 + token_to_req[row_slice],
769 + query_positions[row_slice],
770 + sequence_lengths,
771 + compress_ratio,
772 + )
773 + blocks = blocks_buffer[: row_end - row_start]
774 + use_cooperative_topk = (
775 + current_platform.is_cuda()
776 + and blocks.shape[0] <= 32
777 + and logits.stride(0) % 4 == 0
778 + and current_platform.has_device_capability(90)
779 + and not current_platform.is_device_capability_family(120)
780 + )
781 + if use_cooperative_topk:
782 + torch.ops._C.cooperative_topk(
783 + logits,
784 + visible_blocks,
785 + blocks,
786 + topk_workspace,
787 + block_topk,
788 + columns,
789 + )
790 + elif current_platform.is_cuda():
791 + torch.ops._C.persistent_topk(
792 + logits,
793 + visible_blocks,
794 + blocks,
795 + topk_workspace,
796 + block_topk,
797 + columns,
798 + )
799 + else:
800 + ops.top_k_per_row_decode(
801 + logits,
802 + 1,
803 + visible_blocks,
804 + blocks,
805 + blocks.shape[0],
806 + logits.stride(0),
807 + logits.stride(1),
808 + block_topk,
809 + )
810 + expand_qsa_block_indices_cuda(
811 + blocks,
812 + query_positions[row_slice],
813 + sequence_lengths,
814 + token_to_req[row_slice],
815 + compress_ratio,
816 + token_topk,
817 + out[row_slice],
818 + )
819 + return out
820 +
821 +
822 +def qsa_sparse_paged_attention(
823 + q: torch.Tensor,
824 + k_cache: torch.Tensor,
825 + v_cache: torch.Tensor,
826 + logical_indices: torch.Tensor,
827 + block_table: torch.Tensor,
828 + token_to_req: torch.Tensor,
829 + out: torch.Tensor | None = None,
830 +) -> torch.Tensor:
831 + """Run sparse GQA directly over paged BF16 K/V caches."""
832 +
833 + if not q.is_cuda or not HAS_TRITON:
834 + raise RuntimeError("paged QSA sparse attention requires a GPU and Triton")
835 + if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape:
836 + raise ValueError("QSA sparse attention received invalid Q/K/V shapes")
837 + if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]:
838 + raise ValueError("QSA indices must have one row per query")
839 + if token_to_req.shape != (q.shape[0],) or block_table.ndim != 2:
840 + raise ValueError("QSA sparse attention metadata has invalid shapes")
841 + if not all(k_cache.shape[:3]) or not all(block_table.shape):
842 + raise ValueError("QSA sparse attention cache and block table must be nonempty")
843 + if logical_indices.shape[1] <= 0:
844 + raise ValueError("QSA sparse attention requires a positive selection width")
845 + if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]:
846 + raise ValueError("QSA sparse attention requires valid grouped-query heads")
847 + head_dim = q.shape[2]
848 + assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0
849 + assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16
850 + assert logical_indices.dtype == block_table.dtype == torch.int32
851 + assert token_to_req.dtype == torch.int32
852 + assert q.device == k_cache.device == v_cache.device
853 + assert q.device == logical_indices.device == block_table.device
854 + assert q.device == token_to_req.device
855 + assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1
856 + assert logical_indices.stride(1) == block_table.stride(1) == 1
857 + assert token_to_req.stride(0) == 1
858 + if out is None:
859 + out = torch.empty_like(q)
860 + if out.shape != q.shape:
861 + raise ValueError("QSA sparse output must match its query")
862 + assert out.dtype == q.dtype and out.device == q.device
863 + assert out.stride(2) == 1
864 + if not q.shape[0]:
865 + return out
866 +
867 + group_size = q.shape[1] // k_cache.shape[2]
868 + block_m = triton.next_power_of_2(group_size)
869 + base_programs = q.shape[0] * k_cache.shape[2]
870 + small_profile_limit = 8 if block_m <= 8 else 4
871 +
872 + # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes.
873 + # Narrow tiles favor decode; wide tiles improve throughput for prefill.
874 + if base_programs <= small_profile_limit:
875 + block_n, target_splits, partial_warps = 16, 64, 4
876 + elif base_programs < 32:
877 + block_n, target_splits, partial_warps = 16, 32, 4
878 + elif base_programs <= 256:
879 + block_n, target_splits, partial_warps = 64, 8, 2
880 + elif base_programs <= 512:
881 + block_n, target_splits, partial_warps = 64, 4, 2
882 + else:
883 + block_n, target_splits, partial_warps = 64, 1, 2
884 + # gfx942 and gfx950 have a 64 KiB LDS limit. One software-pipelining
885 + # stage keeps the wide TP4 tile within that shared-memory budget.
886 + partial_stages = 1 if current_platform.is_rocm() else 2
887 +
888 + num_tiles = triton.cdiv(logical_indices.shape[1], block_n)
889 + # Avoid empty splits when the selection width is smaller than the profile.
890 + max_useful_splits = 1 << (num_tiles.bit_length() - 1)
891 + num_splits = min(max_useful_splits, target_splits)
892 +
893 + # Split=1 writes output directly and compiles out all workspace accesses.
894 + if num_splits == 1:
895 + partial_output = out
896 + partial_lse = out
897 + else:
898 + # FP32 partials preserve accuracy when merging independently normalized
899 + # splits.
900 + partial_output = torch.empty(
901 + (num_splits, *q.shape), dtype=torch.float32, device=q.device
902 + )
903 + partial_lse = torch.empty(
904 + (num_splits, q.shape[0], q.shape[1]),
905 + dtype=torch.float32,
906 + device=q.device,
907 + )
908 +
909 + partial_grid = (q.shape[0], k_cache.shape[2], num_splits)
910 + _qsa_sparse_paged_gqa_splitk_kernel[partial_grid](
911 + q,
912 + k_cache,
913 + v_cache,
914 + logical_indices,
915 + block_table,
916 + token_to_req,
917 + partial_output,
918 + partial_lse,
919 + out,
920 + q.stride(0),
921 + q.stride(1),
922 + k_cache.stride(0),
923 + k_cache.stride(1),
924 + k_cache.stride(2),
925 + v_cache.stride(0),
926 + v_cache.stride(1),
927 + v_cache.stride(2),
928 + logical_indices.stride(0),
929 + block_table.stride(0),
930 + out.stride(0),
931 + out.stride(1),
932 + q.shape[0],
933 + k_cache.shape[0],
934 + block_table.shape[0],
935 + TOPK=logical_indices.shape[1],
936 + PAGE_SIZE=k_cache.shape[1],
937 + PAGE_TABLE_WIDTH=block_table.shape[1],
938 + GROUP_SIZE=group_size,
939 + HEAD_DIM=q.shape[2],
940 + NUM_QUERY_HEADS=q.shape[1],
941 + NUM_SPLITS=num_splits,
942 + NUM_TILES=num_tiles,
943 + BLOCK_M=block_m,
944 + BLOCK_N=block_n,
945 + num_warps=partial_warps,
946 + num_stages=partial_stages,
947 + )
948 + if num_splits == 1:
949 + return out
950 +
951 + _qsa_merge_splitk_kernel[(q.shape[0], q.shape[1])](
952 + partial_output,
953 + partial_lse,
954 + out,
955 + out.stride(0),
956 + out.stride(1),
957 + q.shape[0],
958 + HEAD_DIM=q.shape[2],
959 + NUM_QUERY_HEADS=q.shape[1],
960 + NUM_SPLITS=num_splits,
961 + BLOCK_SPLITS=triton.next_power_of_2(num_splits),
962 + num_warps=2,
963 + num_stages=1,
964 + )
965 + return out
966 +
967 +
968 +def qsa_store_cache_rows(
969 + cache: torch.Tensor,
970 + slot_mapping: torch.Tensor,
971 + rows: torch.Tensor,
972 +) -> None:
973 + """Store fixed-width rows in a QSA cache without boolean indexing."""
974 +
975 + if not cache.is_cuda or not HAS_TRITON:
976 + raise RuntimeError("QSA cache stores require a GPU and Triton")
977 + if cache.ndim != 4 or cache.shape[2] != 1:
978 + raise ValueError("QSA cache must be [pages, page_size, 1, width]")
979 + if not all(cache.shape):
980 + raise ValueError("QSA cache dimensions must be nonzero")
981 + if rows.ndim == 3:
982 + if rows.shape[1] != 1:
983 + raise ValueError("QSA cache rows must have one head")
984 + rows = rows[:, 0]
985 + if rows.shape != (slot_mapping.numel(), cache.shape[3]):
986 + raise ValueError("QSA cache rows and slots have incompatible shapes")
987 + if not rows.shape[0]:
988 + return
989 + _store_qsa_rows_kernel[(rows.shape[0],)](
990 + cache,
991 + slot_mapping,
992 + rows,
993 + cache.stride(0),
994 + cache.stride(1),
995 + cache.stride(3),
996 + rows.stride(0),
997 + rows.stride(1),
998 + rows.shape[0],
999 + cache.shape[0],
1000 + PAGE_SIZE=cache.shape[1],
1001 + WIDTH=cache.shape[3],
1002 + BLOCK_D=triton.next_power_of_2(cache.shape[3]),
1003 + num_warps=4,
1004 + )
1005 +
1006 +
1007 +def qsa_compress_groups_with_ratio(
1008 + raw_keys: torch.Tensor, # this step's raw key rows [rows, 1, head_size]
1009 + raw_positions: torch.Tensor, # this step's positions [rows, 1, 3] int64
1010 + compressor_state_cache: torch.Tensor,
1011 + compressor_state_block_table: torch.Tensor,
1012 + token_to_req: torch.Tensor,
1013 + query_start_loc: torch.Tensor,
1014 + logical_positions: torch.Tensor,
1015 + compressed_slots: torch.Tensor,
1016 + compress_ratio: int,
1017 + rope_cache: torch.Tensor | None = None,
1018 +) -> tuple[torch.Tensor, torch.Tensor]:
1019 + """Pool completed groups from the compressor-state ring and raw token rows."""
1020 +
1021 + if not raw_keys.is_cuda or not HAS_TRITON:
1022 + raise RuntimeError("QSA compression requires a GPU and Triton")
1023 + rows = token_to_req.numel()
1024 + if compress_ratio <= 0:
1025 + raise ValueError("QSA compression ratio must be positive")
1026 + if raw_keys.ndim != 3 or raw_keys.shape[:2] != (rows, 1):
1027 + raise ValueError("QSA raw keys must be [rows, 1, head_size]")
1028 + if raw_positions.shape != (rows, 1, 3) or raw_positions.dtype != torch.int64:
1029 + raise ValueError("QSA raw positions must be [rows, 1, 3] int64")
1030 + if logical_positions.shape != (rows,) or compressed_slots.shape != (rows,):
1031 + raise ValueError("QSA compression metadata must match token rows")
1032 + if compressor_state_cache.ndim != 4 or compressor_state_cache.shape[2] != 1:
1033 + raise ValueError("QSA compressor-state cache has an invalid shape")
1034 + if (
1035 + # The ring is wider than one group so speculative rows cannot alias
1036 + # onto the committed keys of the group still being collected.
1037 + compressor_state_cache.shape[1] < compress_ratio
1038 + or compressor_state_cache.shape[3] != raw_keys.shape[2]
1039 + or compressor_state_cache.dtype != raw_keys.dtype
1040 + ):
1041 + raise ValueError(
1042 + "QSA compressor-state cache does not match the compression layout"
1043 + )
1044 + if (
1045 + compressor_state_block_table.ndim != 2
1046 + or compressor_state_block_table.shape[1] < 1
1047 + ):
1048 + raise ValueError(
1049 + "QSA compressor-state block table must contain one block per request"
1050 + )
1051 + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2:
1052 + raise ValueError("QSA query starts must contain a terminal offset")
1053 + num_requests = query_start_loc.shape[0] - 1
1054 + if compressor_state_block_table.shape[0] < num_requests:
1055 + raise ValueError("QSA compressor-state block table has too few request rows")
1056 + if rope_cache is not None and (
1057 + rope_cache.ndim != 4
1058 + or rope_cache.shape[:3] != compressor_state_cache.shape[:3]
1059 + or rope_cache.shape[3] != 3
1060 + or rope_cache.dtype != torch.int64
1061 + ):
1062 + raise ValueError("QSA packed position view has an invalid shape or dtype")
1063 + if rows and (
1064 + not all(compressor_state_cache.shape)
1065 + or not all(compressor_state_block_table.shape)
1066 + ):
1067 + raise ValueError("QSA compressor-state cache and block table must be nonempty")
1068 + pooled = torch.empty(
1069 + (rows, 1, raw_keys.shape[2]),
1070 + dtype=raw_keys.dtype,
1071 + device=raw_keys.device,
1072 + )
1073 + first_positions = torch.empty((rows, 3), dtype=torch.int64, device=raw_keys.device)
1074 + if not rows:
1075 + return pooled, first_positions
1076 + if rope_cache is None:
1077 + rope_cache = compressor_state_cache
1078 + load_rope_positions = False
1079 + else:
1080 + load_rope_positions = True
1081 + _compress_qsa_groups_kernel[(rows,)](
1082 + raw_keys,
1083 + raw_positions,
1084 + compressor_state_cache,
1085 + rope_cache,
1086 + compressor_state_block_table,
1087 + token_to_req,
1088 + query_start_loc,
1089 + logical_positions,
1090 + compressed_slots,
1091 + pooled,
1092 + first_positions,
1093 + raw_keys.stride(0),
1094 + raw_keys.stride(2),
1095 + raw_positions.stride(0),
1096 + raw_positions.stride(2),
1097 + compressor_state_cache.stride(0),
1098 + compressor_state_cache.stride(1),
1099 + compressor_state_cache.stride(3),
1100 + rope_cache.stride(0),
1101 + rope_cache.stride(1),
1102 + rope_cache.stride(3),
1103 + compressor_state_block_table.stride(0),
1104 + pooled.stride(0),
1105 + pooled.stride(2),
1106 + first_positions.stride(0),
1107 + first_positions.stride(1),
1108 + rows,
1109 + compressor_state_cache.shape[0],
1110 + num_requests,
1111 + COMPRESSOR_STATE_SIZE=compressor_state_cache.shape[1],
1112 + COMPRESS_RATIO=compress_ratio,
1113 + HEAD_DIM=raw_keys.shape[2],
1114 + LOAD_ROPE_POSITIONS=load_rope_positions,
1115 + BLOCK_D=triton.next_power_of_2(raw_keys.shape[2]),
1116 + num_warps=4,
1117 + )
1118 + return pooled, first_positions
1119 +
1120 +
1121 +__all__ = [
1122 + "expand_qsa_block_indices_cuda",
1123 + "qsa_compress_groups_with_ratio",
1124 + "qsa_mqa_paged",
1125 + "qsa_select_paged_tokens",
1126 + "qsa_sparse_paged_attention",
1127 + "qsa_store_cache_rows",
1128 +]
vllm/models/qwen4_exp/amd/ple_layer.py new
+1131
@@ -0,0 +1,1131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""GPU-resident Qwen4Exp position-learning enhancement layers."""
4 +
5 +import math
6 +from collections.abc import Iterable, Sequence
7 +
8 +import torch
9 +import torch.nn.functional as F
10 +from torch import nn
11 +
12 +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
13 +from vllm.forward_context import get_forward_context
14 +from vllm.model_executor.layers.linear import ReplicatedLinear
15 +from vllm.model_executor.layers.mamba.abstract import MambaBase
16 +from vllm.model_executor.layers.mamba.mamba_utils import (
17 + MambaStateDtypeCalculator,
18 + MambaStateShapeCalculator,
19 + is_conv_state_dim_first,
20 +)
21 +from vllm.model_executor.models.utils import AutoWeightsLoader
22 +from vllm.transformers_utils.configs.qwen4_exp import (
23 + Qwen4ExpTextConfig,
24 +)
25 +from vllm.utils.torch_utils import direct_register_custom_op
26 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
27 +from vllm.v1.attention.backends.short_conv_attn import (
28 + PleShortConvAttentionBackend,
29 + PleShortConvAttentionMetadata,
30 +)
31 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
32 +
33 +from ..common.ple import PLEVocabParallelEmbedding
34 +
35 +
36 +class Qwen4ExpPLEGroupedNorm(nn.Module):
37 + def __init__(
38 + self,
39 + hidden_size: int,
40 + eps: float,
41 + group_size: int | None,
42 + dtype: torch.dtype | None,
43 + ) -> None:
44 + super().__init__()
45 + if group_size is not None and hidden_size % group_size:
46 + raise ValueError(
47 + f"hidden_size ({hidden_size}) must be divisible by "
48 + f"group_size ({group_size})"
49 + )
50 + self.eps = eps
51 + self.group_size = group_size
52 + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype))
53 +
54 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
55 + input_dtype = hidden_states.dtype
56 + hidden_states = hidden_states.float()
57 + if self.group_size is None:
58 + variance = hidden_states.square().mean(dim=-1, keepdim=True)
59 + normalized = hidden_states * torch.rsqrt(variance + self.eps)
60 + else:
61 + grouped = hidden_states.unflatten(
62 + -1, (hidden_states.shape[-1] // self.group_size, self.group_size)
63 + )
64 + variance = grouped.square().mean(dim=-1, keepdim=True)
65 + normalized = (grouped * torch.rsqrt(variance + self.eps)).flatten(-2)
66 + return (normalized * (1.0 + self.weight.float())).to(input_dtype)
67 +
68 +
69 +class Qwen4ExpNGramEmbedding(nn.Module):
70 + _MASK64 = (1 << 64) - 1
71 + _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
72 + _SPLITMIX_M1 = 0xBF58476D1CE4E5B9
73 + _SPLITMIX_M2 = 0x94D049BB133111EB
74 + _PLE_LAYER_PRIME = 10007
75 +
76 + @classmethod
77 + def _splitmix64(cls, value: int) -> int:
78 + """Mix an integer into a deterministic unsigned 64-bit value."""
79 + value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64
80 + value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64
81 + value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64
82 + return (value ^ (value >> 31)) & cls._MASK64
83 +
84 + @staticmethod
85 + def _is_prime_64(value: int) -> bool:
86 + """Return whether a 64-bit integer is prime."""
87 + if value < 2:
88 + return False
89 + for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
90 + if value % prime == 0:
91 + return value == prime
92 + exponent = value - 1
93 + shifts = 0
94 + while exponent % 2 == 0:
95 + exponent //= 2
96 + shifts += 1
97 + for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022):
98 + if base % value == 0:
99 + continue
100 + witness = pow(base, exponent, value)
101 + if witness in (1, value - 1):
102 + continue
103 + for _ in range(shifts - 1):
104 + witness = pow(witness, 2, value)
105 + if witness == value - 1:
106 + break
107 + else:
108 + return False
109 + return True
110 +
111 + @classmethod
112 + def _nth_prime_after(cls, start: int, count: int) -> int:
113 + """Return the ``count``-th prime strictly greater than ``start``."""
114 + prime = int(start)
115 + for _ in range(count):
116 + candidate = prime + 1
117 + if candidate <= 2:
118 + prime = 2
119 + continue
120 + if candidate % 2 == 0:
121 + candidate += 1
122 + while not cls._is_prime_64(candidate):
123 + candidate += 2
124 + prime = candidate
125 + return prime
126 +
127 + @classmethod
128 + def _make_layer_multipliers(
129 + cls,
130 + *,
131 + ngram_size: int,
132 + unigram_vocab_size: int,
133 + seed: int,
134 + ple_dense_layer_id: int,
135 + ) -> list[int]:
136 + """Build deterministic hash multipliers for one PLE layer."""
137 + max_multiplier = ((1 << 63) - 1) // unigram_vocab_size
138 + half_bound = max(1, max_multiplier // 2)
139 + base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id
140 + multipliers = []
141 + for index in range(ngram_size):
142 + value = base_seed + cls._SPLITMIX_GAMMA * (index + 1)
143 + multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1)
144 + return multipliers
145 +
146 + @classmethod
147 + def _make_vocab_layout(
148 + cls,
149 + *,
150 + ngram_vocab_size_base: int,
151 + ngram_heads: int,
152 + ple_dense_layer_id: int,
153 + ) -> tuple[list[int], list[int], int]:
154 + """Build per-head vocabulary sizes, offsets, and total row count."""
155 + sizes: list[int] = []
156 + offsets: list[int] = []
157 + offset = 0
158 + for local_head in range(ngram_heads):
159 + global_head = ple_dense_layer_id * ngram_heads + local_head
160 + size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1)
161 + sizes.append(size)
162 + offsets.append(offset)
163 + offset += size
164 + return sizes, offsets, offset
165 +
166 + def __init__(
167 + self,
168 + config: Qwen4ExpTextConfig,
169 + embedding_dim: int,
170 + ple_dense_layer_id: int,
171 + max_total_tokens: int,
172 + max_num_reqs: int,
173 + prefix: str,
174 + layer_name: str,
175 + ) -> None:
176 + super().__init__()
177 + self.embedding_dim = embedding_dim
178 + self.layer_name = layer_name
179 + self.ngram_size = int(config.ngram_size)
180 + self.heads_per_ngram = int(config.heads_per_ngram)
181 + self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
182 + if self.ngram_size < 2:
183 + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}")
184 + if self.heads_per_ngram <= 0:
185 + raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}")
186 + if embedding_dim % self.ngram_heads:
187 + raise ValueError(
188 + "ple_embed_dim must be divisible by total ngram heads: "
189 + f"{embedding_dim} % {self.ngram_heads} != 0"
190 + )
191 + self.head_dim = embedding_dim // self.ngram_heads
192 + self.eos_token_id = int(config.eos_token_id)
193 + self.unigram_vocab_size = int(config.vocab_size)
194 + self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512))
195 + if self.split_ngram_parts <= 0:
196 + raise ValueError("split_ngram_parts must be positive")
197 +
198 + multipliers = self._make_layer_multipliers(
199 + ngram_size=self.ngram_size,
200 + unigram_vocab_size=self.unigram_vocab_size,
201 + seed=int(getattr(config, "seed", 1234)),
202 + ple_dense_layer_id=ple_dense_layer_id,
203 + )
204 + self.register_buffer(
205 + "layer_multipliers",
206 + torch.tensor(multipliers, dtype=torch.long),
207 + persistent=True,
208 + )
209 +
210 + sizes, offsets, total_vocab_size = self._make_vocab_layout(
211 + ngram_vocab_size_base=int(config.ngram_vocab_size_base),
212 + ngram_heads=self.ngram_heads,
213 + ple_dense_layer_id=ple_dense_layer_id,
214 + )
215 + self.register_buffer(
216 + "ngram_heads_vocab_sizes",
217 + torch.tensor(sizes, dtype=torch.long),
218 + persistent=True,
219 + )
220 + self.register_buffer(
221 + "ngram_heads_offsets",
222 + torch.tensor(offsets, dtype=torch.long),
223 + persistent=True,
224 + )
225 + divisor = int(config.make_ngram_vocab_size_divisible_by)
226 + padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
227 + self.ngram_embedding = PLEVocabParallelEmbedding(
228 + padded_vocab_size,
229 + self.head_dim,
230 + padding_size=divisor,
231 + prefix=f"{prefix}.ngram_embedding",
232 + )
233 + self.register_buffer(
234 + "positions_buffer",
235 + torch.arange(max_total_tokens, dtype=torch.int64),
236 + persistent=False,
237 + )
238 + self.register_buffer(
239 + "padded_buffer",
240 + torch.full(
241 + (max_num_reqs, max_total_tokens),
242 + self.eos_token_id,
243 + dtype=torch.int64,
244 + ),
245 + persistent=False,
246 + )
247 +
248 + @staticmethod
249 + def _shift_precompute(
250 + tokens: torch.Tensor, eos_token_id: int
251 + ) -> tuple[torch.Tensor, torch.Tensor]:
252 + if tokens.dim() != 2:
253 + raise ValueError("tokens must be a 2D tensor")
254 + batch_size, seq_len = tokens.shape
255 + positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64)
256 + eos_positions = torch.where(tokens == eos_token_id, positions, -1)
257 + previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
258 + previous_eos = torch.cat(
259 + [
260 + eos_positions.new_full((batch_size, 1), -1),
261 + previous_eos_inclusive[:, :-1],
262 + ],
263 + dim=1,
264 + )
265 + return positions, positions.unsqueeze(0) - previous_eos - 1
266 +
267 + @staticmethod
268 + def _shift_apply(
269 + tokens: torch.Tensor,
270 + positions: torch.Tensor,
271 + position_in_segment: torch.Tensor,
272 + shift: int,
273 + eos_token_id: int,
274 + ) -> torch.Tensor:
275 + if shift == 0:
276 + return tokens
277 + source = positions - shift
278 + gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1)
279 + shifted = tokens.gather(1, gather_indices)
280 + valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift)
281 + return torch.where(valid, shifted, tokens.new_full((), eos_token_id))
282 +
283 + def forward(
284 + self,
285 + input_ids: torch.Tensor,
286 + query_start_loc: torch.Tensor,
287 + ngram_context: torch.Tensor,
288 + ) -> torch.Tensor:
289 + input_ids = input_ids.reshape(-1).long()
290 + query_start_loc = query_start_loc.long()
291 + num_reqs = query_start_loc.numel() - 1
292 + num_tokens = input_ids.shape[0]
293 + if num_tokens > self.positions_buffer.numel():
294 + raise ValueError(
295 + f"PLE received {num_tokens} tokens, but its workspace supports "
296 + f"at most {self.positions_buffer.numel()}"
297 + )
298 + if num_reqs > self.padded_buffer.shape[0]:
299 + raise ValueError(
300 + f"PLE received {num_reqs} requests, but its workspace supports "
301 + f"at most {self.padded_buffer.shape[0]}"
302 + )
303 +
304 + positions = self.positions_buffer[:num_tokens]
305 + packed = self.padded_buffer[:num_reqs]
306 + packed.fill_(self.eos_token_id)
307 + request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
308 + request_indices.clamp_(max=num_reqs - 1)
309 + columns = (positions - query_start_loc[request_indices]).clamp(
310 + 0, packed.shape[1] - 1
311 + )
312 + packed[request_indices, columns] = input_ids
313 + ngram_context = ngram_context[:num_reqs].to(
314 + device=input_ids.device, dtype=torch.long
315 + )
316 +
317 + context = torch.cat([ngram_context, packed], dim=-1)
318 + positions_2d, position_in_segment = self._shift_precompute(
319 + context, self.eos_token_id
320 + )
321 + shifted = [context]
322 + for shift in range(1, self.ngram_size):
323 + shifted.append(
324 + self._shift_apply(
325 + context,
326 + positions_2d,
327 + position_in_segment,
328 + shift,
329 + self.eos_token_id,
330 + )
331 + )
332 + adjusted_columns = columns + self.ngram_size - 1
333 + id_blocks = []
334 + for ngram in range(2, self.ngram_size + 1):
335 + start = (ngram - 2) * self.heads_per_ngram
336 + end = start + self.heads_per_ngram
337 + mixed = shifted[0] * self.layer_multipliers[0]
338 + for index in range(1, ngram):
339 + mixed = torch.bitwise_xor(
340 + mixed, shifted[index] * self.layer_multipliers[index]
341 + )
342 + sizes = self.ngram_heads_vocab_sizes[start:end]
343 + offsets = self.ngram_heads_offsets[start:end]
344 + ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
345 + id_blocks.append(ids[request_indices, adjusted_columns])
346 + ngram_ids = torch.cat(id_blocks, dim=-1)
347 + output = ngram_ids.new_empty(
348 + (ngram_ids.shape[0], self.embedding_dim),
349 + dtype=self.ngram_embedding.params_dtype,
350 + )
351 + torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding(
352 + ngram_ids,
353 + output,
354 + self.layer_name,
355 + )
356 + return output
357 +
358 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
359 + """Load hash buffers and checkpoint-split embedding rows."""
360 +
361 + persistent_buffers = {
362 + "layer_multipliers": self.layer_multipliers,
363 + "ngram_heads_offsets": self.ngram_heads_offsets,
364 + "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
365 + }
366 + loaded: set[str] = set()
367 + regular_weights: list[tuple[str, torch.Tensor]] = []
368 + shard_prefix = "ngram_embedding.shard_"
369 +
370 + for name, loaded_weight in weights:
371 + leaf_name = name.rsplit(".", 1)[-1]
372 + if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
373 + continue
374 + if name in persistent_buffers:
375 + buffer = persistent_buffers[name]
376 + if buffer.shape != loaded_weight.shape:
377 + raise ValueError(
378 + f"Shape mismatch for {name}: expected "
379 + f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
380 + )
381 + buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
382 + loaded.add(name)
383 + continue
384 + if name.startswith(shard_prefix) and name.endswith(".weight"):
385 + shard_text = name[len(shard_prefix) : -len(".weight")]
386 + if not shard_text.isdigit():
387 + regular_weights.append((name, loaded_weight))
388 + continue
389 + shard_index = int(shard_text)
390 + if shard_index >= self.split_ngram_parts:
391 + raise ValueError(
392 + f"PLE embedding shard index {shard_index} exceeds "
393 + f"split_ngram_parts={self.split_ngram_parts}"
394 + )
395 + embedding = self.ngram_embedding
396 + shard_size = (
397 + embedding.org_vocab_size + self.split_ngram_parts - 1
398 + ) // self.split_ngram_parts
399 + checkpoint_start = shard_index * shard_size
400 + expected_rows = max(
401 + 0,
402 + min(shard_size, embedding.org_vocab_size - checkpoint_start),
403 + )
404 + expected_shape = (expected_rows, embedding.embedding_dim)
405 + if tuple(loaded_weight.shape) != expected_shape:
406 + raise ValueError(
407 + f"Shape mismatch for PLE embedding shard {shard_index}: "
408 + f"expected {expected_shape}, got "
409 + f"{tuple(loaded_weight.shape)}"
410 + )
411 + embedding.weight.weight_loader(
412 + embedding.weight,
413 + loaded_weight,
414 + checkpoint_start=checkpoint_start,
415 + )
416 + loaded.add("ngram_embedding.weight")
417 + continue
418 + regular_weights.append((name, loaded_weight))
419 +
420 + if regular_weights:
421 + loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
422 + return loaded
423 +
424 +
425 +class Qwen4ExpPLELayer(nn.Module, MambaBase):
426 + def __init__(
427 + self,
428 + config: Qwen4ExpTextConfig,
429 + vllm_config: VllmConfig,
430 + layer_idx: int = 0,
431 + ple_dense_layer_id: int | None = None,
432 + prefix: str = "",
433 + ) -> None:
434 + super().__init__()
435 + model_config = vllm_config.model_config
436 + cache_config = vllm_config.cache_config
437 + quant_config = vllm_config.quant_config
438 + self.model_config: ModelConfig = model_config
439 + self.cache_config: CacheConfig = cache_config
440 + self.layer_idx = layer_idx
441 + self.ple_dense_layer_id = (
442 + int(ple_dense_layer_id)
443 + if ple_dense_layer_id is not None
444 + else int(layer_idx)
445 + )
446 + self.prefix = prefix
447 + self.hidden_size = int(config.hidden_size)
448 + self.hc_count = config.hc_count
449 + self.hc_hidden_size = self.hidden_size * self.hc_count
450 + self.conv_kernel_size = int(config.ple_conv_kernel_size)
451 + self.short_conv_dilation = int(config.ngram_size)
452 + self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation
453 + self.num_spec_tokens = vllm_config.num_speculative_tokens
454 + self.activation = "silu"
455 + self.ple_embedding: nn.Module = Qwen4ExpNGramEmbedding(
456 + config,
457 + int(config.ple_embed_dim),
458 + self.ple_dense_layer_id,
459 + vllm_config.scheduler_config.max_num_batched_tokens,
460 + vllm_config.scheduler_config.max_num_seqs,
461 + f"{prefix}.ple_embedding",
462 + prefix,
463 + )
464 + self.key_proj = ReplicatedLinear(
465 + int(config.ple_embed_dim),
466 + self.hc_hidden_size,
467 + bias=False,
468 + quant_config=quant_config,
469 + prefix=f"{prefix}.key_proj",
470 + )
471 + self.value_proj = ReplicatedLinear(
472 + int(config.ple_embed_dim),
473 + self.hidden_size,
474 + bias=False,
475 + quant_config=quant_config,
476 + prefix=f"{prefix}.value_proj",
477 + )
478 + norm_args = (
479 + self.hc_hidden_size,
480 + config.rms_norm_eps,
481 + self.hidden_size,
482 + model_config.dtype,
483 + )
484 + self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args)
485 + self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args)
486 + self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args)
487 + self.conv1d = nn.Conv1d(
488 + self.hc_hidden_size,
489 + self.hc_hidden_size,
490 + self.conv_kernel_size,
491 + groups=self.hc_hidden_size,
492 + padding=self.conv_state_len,
493 + dilation=self.short_conv_dilation,
494 + bias=False,
495 + dtype=model_config.dtype,
496 + )
497 + nn.init.zeros_(self.conv1d.weight)
498 + self.conv1d.weight._no_reinit = True
499 + self.kv_cache = (torch.tensor([]),)
500 + compilation_config = get_current_vllm_config().compilation_config
501 + if prefix in compilation_config.static_forward_context:
502 + raise ValueError(f"Duplicate layer name: {prefix}")
503 + compilation_config.static_forward_context[prefix] = self
504 +
505 + @property
506 + def mamba_type(self) -> MambaAttentionBackendEnum:
507 + return MambaAttentionBackendEnum.SHORT_CONV
508 +
509 + @property
510 + def is_kv_cache_tp_replicated(self) -> bool:
511 + return True
512 +
513 + def get_attn_backend(self) -> type[PleShortConvAttentionBackend]:
514 + return PleShortConvAttentionBackend
515 +
516 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:
517 + return MambaStateDtypeCalculator.short_conv_state_dtype(
518 + self.model_config.dtype, self.cache_config.mamba_cache_dtype
519 + )
520 +
521 + def get_state_shape(self) -> Sequence[tuple[int, ...]]:
522 + return MambaStateShapeCalculator.short_conv_state_shape(
523 + tp_world_size=1,
524 + intermediate_size=self.hc_hidden_size,
525 + conv_kernel=self.conv_state_len + 1,
526 + num_spec=self.num_spec_tokens,
527 + )
528 +
529 + def _apply_norm(
530 + self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor
531 + ) -> torch.Tensor:
532 + shape = hidden_states.shape
533 + return norm(hidden_states.flatten(-2)).reshape(shape)
534 +
535 + def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor:
536 + # Profiling / CUDA graph capture only; conv state is not updated.
537 + inputs_t = inputs.transpose(0, 1).unsqueeze(0)
538 + output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)]
539 + return F.silu(output).squeeze(0).transpose(0, 1)
540 +
541 + def _short_conv_dilated_decode_batched(
542 + self,
543 + x_d: torch.Tensor,
544 + conv_state: torch.Tensor,
545 + conv_weights: torch.Tensor,
546 + state_indices_tensor_d: torch.Tensor,
547 + has_initial_states_d: torch.Tensor | None,
548 + ) -> torch.Tensor:
549 + state_indices = state_indices_tensor_d.to(
550 + device=conv_state.device, dtype=torch.int64
551 + )
552 + # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to
553 + # slot 0 for a safe gather, then zero output and skip write-back.
554 + valid_state = state_indices != NULL_BLOCK_ID
555 + state_indices = torch.where(
556 + valid_state, state_indices, torch.zeros_like(state_indices)
557 + )
558 + if has_initial_states_d is None:
559 + has_initial_state = valid_state
560 + else:
561 + if has_initial_states_d.numel() < state_indices_tensor_d.numel():
562 + raise ValueError(
563 + "has_initial_states_d size mismatch: "
564 + f"got {has_initial_states_d.numel()}, "
565 + f"need >= {state_indices_tensor_d.numel()}."
566 + )
567 + has_initial_state = has_initial_states_d[
568 + : state_indices_tensor_d.numel()
569 + ].to(device=conv_state.device, dtype=torch.bool)
570 + has_initial_state = has_initial_state & valid_state
571 +
572 + cached_state = conv_state.index_select(0, state_indices)
573 + state = cached_state[..., : self.conv_state_len].to(x_d.dtype)
574 + if self.conv_state_len > 0:
575 + initial_state = torch.where(
576 + has_initial_state.view(-1, 1, 1),
577 + state,
578 + torch.zeros_like(state),
579 + )
580 + history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1)
581 + else:
582 + history = x_d.unsqueeze(-1)
583 +
584 + conv_output = F.conv1d(
585 + history,
586 + conv_weights.unsqueeze(1).contiguous(),
587 + groups=history.size(1),
588 + dilation=self.short_conv_dilation,
589 + ).squeeze(-1)
590 + output = F.silu(conv_output)
591 + output = output * valid_state.view(-1, 1).to(output.dtype)
592 +
593 + if self.conv_state_len > 0:
594 + next_state = history[..., -self.conv_state_len :]
595 + # Padded rows are remapped to the reserved null slot. Preserve its
596 + # existing value while writing the new states for valid rows.
597 + existing_base_state = cached_state[..., : self.conv_state_len]
598 + safe_next_state = torch.where(
599 + valid_state.view(-1, 1, 1),
600 + next_state.to(conv_state.dtype),
601 + existing_base_state,
602 + )
603 + cached_state[..., : self.conv_state_len] = safe_next_state
604 + conv_state.index_copy_(0, state_indices, cached_state)
605 +
606 + return output
607 +
608 + def _short_conv_dilated_prefill_batched(
609 + self,
610 + x_p: torch.Tensor,
611 + metadata: PleShortConvAttentionMetadata,
612 + conv_state: torch.Tensor,
613 + conv_weights: torch.Tensor,
614 + state_indices_tensor_p: torch.Tensor,
615 + num_prefills: int,
616 + num_decode_tokens: int,
617 + num_prefill_tokens: int,
618 + ) -> torch.Tensor:
619 + # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill)
620 + # requests and equals ``query_start_loc`` when spec-decode is inactive.
621 + non_spec_query_start_loc = metadata.non_spec_query_start_loc
622 + if non_spec_query_start_loc is None:
623 + raise ValueError("query_start_loc is required for prefill short-conv")
624 + query_start_loc_p = (
625 + non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens
626 + )
627 + # The metadata builder guarantees that the prefill query offsets start
628 + # at 0 and end at num_prefill_tokens. Avoid reading those values here,
629 + # since doing so would force a device-to-host synchronization.
630 + has_initial_states_p = metadata.has_initial_states_p
631 + if has_initial_states_p is None:
632 + raise ValueError("has_initial_states_p is required for prefill short-conv")
633 +
634 + output = torch.empty_like(x_p)
635 + q_starts = query_start_loc_p.to(torch.int64)
636 + if state_indices_tensor_p.numel() < num_prefills:
637 + raise ValueError(
638 + "state_indices_tensor_p size mismatch: "
639 + f"got {state_indices_tensor_p.numel()}, "
640 + f"need >= {num_prefills}."
641 + )
642 + if has_initial_states_p.numel() < num_prefills:
643 + raise ValueError(
644 + "has_initial_states_p size mismatch: "
645 + f"got {has_initial_states_p.numel()}, "
646 + f"need >= {num_prefills}."
647 + )
648 + if num_prefills == 0 or x_p.numel() == 0:
649 + return output
650 + lengths = q_starts[1:] - q_starts[:-1]
651 + # Use the CPU-computed packing width from the metadata builder instead
652 + # of synchronizing on lengths.max().
653 + max_len = metadata.max_prefill_query_len
654 + if max_len <= 0:
655 + return output
656 +
657 + hidden_size = x_p.shape[1]
658 + positions = torch.arange(
659 + num_prefill_tokens, device=x_p.device, dtype=torch.int64
660 + )
661 + req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
662 + col_indices = positions - q_starts[req_indices]
663 +
664 + packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size))
665 + packed_tokens[req_indices, col_indices] = x_p
666 + packed_tokens = packed_tokens.transpose(1, 2).contiguous()
667 +
668 + state_indices = state_indices_tensor_p[:num_prefills].to(
669 + device=conv_state.device, dtype=torch.int64
670 + )
671 + valid_state = state_indices != NULL_BLOCK_ID
672 + state_indices = torch.where(
673 + valid_state, state_indices, torch.zeros_like(state_indices)
674 + )
675 + has_initial = has_initial_states_p[:num_prefills].to(
676 + device=conv_state.device, dtype=torch.bool
677 + )
678 + if self.conv_state_len > 0:
679 + if conv_state.shape[0] == 0:
680 + state = conv_state.new_zeros(
681 + (num_prefills, hidden_size, self.conv_state_len),
682 + dtype=x_p.dtype,
683 + )
684 + else:
685 + state = conv_state.index_select(0, state_indices)[
686 + ..., : self.conv_state_len
687 + ].to(x_p.dtype)
688 + use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1)
689 + initial_state = torch.where(
690 + use_initial_mask,
691 + state,
692 + torch.zeros_like(state),
693 + )
694 + history = torch.cat((initial_state, packed_tokens), dim=-1)
695 + else:
696 + history = packed_tokens
697 +
698 + conv_output = F.conv1d(
699 + history,
700 + conv_weights.unsqueeze(1).contiguous(),
701 + groups=history.size(1),
702 + dilation=self.short_conv_dilation,
703 + )
704 + conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
705 +
706 + token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64)
707 + valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1)
708 + valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view(
709 + num_prefills, 1
710 + )
711 + conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0)
712 + output.copy_(conv_output[req_indices, col_indices])
713 +
714 + if self.conv_state_len > 0 and conv_state.shape[0] > 0:
715 + state_starts = lengths.to(device=history.device, dtype=torch.int64).view(
716 + num_prefills, 1, 1
717 + )
718 + state_offsets = torch.arange(
719 + self.conv_state_len, device=history.device, dtype=torch.int64
720 + ).view(1, 1, self.conv_state_len)
721 + next_state = history.gather(
722 + dim=2,
723 + index=(state_starts + state_offsets).expand(-1, history.size(1), -1),
724 + )
725 + # Write back without a host synchronization. Valid, non-empty rows
726 + # receive their new state; padding and zero-length rows keep the
727 + # current cache value.
728 + existing_state = conv_state.index_select(0, state_indices)
729 + existing_base_state = existing_state[..., : self.conv_state_len]
730 + update_mask = valid_state & (lengths.to(device=conv_state.device) > 0)
731 + safe_next_state = torch.where(
732 + update_mask.view(num_prefills, 1, 1),
733 + next_state.to(conv_state.dtype),
734 + existing_base_state,
735 + )
736 + existing_state[..., : self.conv_state_len] = safe_next_state
737 + conv_state.index_copy_(0, state_indices, existing_state)
738 + return output
739 +
740 + def _short_conv_dilated_spec_batched(
741 + self,
742 + x_spec: torch.Tensor,
743 + conv_state: torch.Tensor,
744 + conv_weights: torch.Tensor,
745 + spec_state_indices_tensor: torch.Tensor,
746 + spec_query_start_loc: torch.Tensor,
747 + num_accepted_tokens: torch.Tensor,
748 + spec_query_len: int,
749 + ) -> torch.Tensor:
750 + """Dilated short-conv for speculative-decode (MTP) requests.
751 +
752 + Each spec request feeds multiple (draft + 1) query tokens. The conv
753 + outputs are computed causally after rolling back the previous draft
754 + state by ``num_accepted_tokens - 1``. The current candidate inputs stay
755 + in the extended cache for the next forward, matching
756 + ``causal_conv1d_update``.
757 +
758 + ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query
759 + length and is a Python int, so no host synchronization is needed; this
760 + keeps the path safe for full CUDA-graph capture/replay where the buffers
761 + are padded at the request level.
762 + """
763 + num_reqs = spec_state_indices_tensor.numel()
764 + hidden_size = x_spec.size(-1)
765 + # Use a fixed packing width instead of synchronizing on lengths.max().
766 + max_len = spec_query_len
767 + # Full CUDA graphs can pad these buffers. Only the first num_reqs
768 + # accepted-token counts belong to actual speculative requests.
769 + num_accepted_tokens = num_accepted_tokens[:num_reqs]
770 + q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64)
771 + # Keep the number of real speculative tokens on the device.
772 + total_real_tokens = q_starts[num_reqs]
773 +
774 + state_indices = spec_state_indices_tensor.to(
775 + device=conv_state.device, dtype=torch.int64
776 + )
777 + valid_state = state_indices != NULL_BLOCK_ID
778 + state_indices = torch.where(
779 + valid_state, state_indices, torch.zeros_like(state_indices)
780 + )
781 + positions = torch.arange(
782 + x_spec.size(0), device=x_spec.device, dtype=torch.int64
783 + )
784 + # Route graph-padded token rows to the discarded dummy request so that
785 + # they cannot overwrite real packed data.
786 + req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
787 + valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs)
788 + clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0))
789 + col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1)
790 + pack_req_indices = torch.where(
791 + valid_tokens,
792 + clamped_req_indices,
793 + torch.full_like(req_indices, num_reqs),
794 + )
795 + pack_col_indices = torch.where(
796 + valid_tokens, col_indices, torch.zeros_like(col_indices)
797 + )
798 +
799 + # The last request row is the dummy sink for graph padding.
800 + packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size))
801 + packed[pack_req_indices, pack_col_indices] = x_spec
802 + packed = packed.transpose(1, 2).contiguous()
803 +
804 + if self.conv_state_len > 0:
805 + cached_state = conv_state.index_select(0, state_indices)
806 + rollback_offsets = num_accepted_tokens.to(
807 + device=conv_state.device, dtype=torch.int64
808 + ).sub(1)
809 + rollback_offsets = torch.where(
810 + valid_state,
811 + rollback_offsets.clamp_(0, max_len - 1),
812 + torch.zeros_like(rollback_offsets),
813 + )
814 + state_offsets = torch.arange(
815 + self.conv_state_len, device=conv_state.device, dtype=torch.int64
816 + ).view(1, 1, self.conv_state_len)
817 + rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets
818 + state = cached_state.gather(
819 + 2, rollback_indices.expand(-1, hidden_size, -1)
820 + ).to(x_spec.dtype)
821 + state = torch.where(
822 + valid_state.view(num_reqs, 1, 1),
823 + state,
824 + torch.zeros_like(state),
825 + )
826 + # Append a zeroed dummy-row state to match the [num_reqs + 1] pack.
827 + dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len))
828 + state_full = torch.cat((state, dummy_state), dim=0)
829 + history = torch.cat((state_full, packed), dim=-1)
830 + else:
831 + history = packed
832 +
833 + conv_output = F.conv1d(
834 + history,
835 + conv_weights.unsqueeze(1).contiguous(),
836 + groups=history.size(1),
837 + dilation=self.short_conv_dilation,
838 + )
839 + conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
840 +
841 + output = conv_output[pack_req_indices, pack_col_indices]
842 + output = output * valid_tokens.view(-1, 1).to(output.dtype)
843 +
844 + # Keep all current candidate inputs in the extended state. On the next
845 + # target forward, ``num_accepted_tokens - 1`` selects the rollback
846 + # window before processing the newly scheduled tokens.
847 + if self.conv_state_len > 0:
848 + state_capacity = self.conv_state_len + max_len - 1
849 + if conv_state.size(-1) < state_capacity:
850 + raise RuntimeError(
851 + "PLE short-conv cache cannot retain speculative tokens: "
852 + f"got {conv_state.size(-1)}, need {state_capacity}."
853 + )
854 + candidate_state = history[:num_reqs, :, 1 : state_capacity + 1]
855 + query_lengths = q_starts[1:] - q_starts[:-1]
856 + state_positions = torch.arange(
857 + state_capacity, device=history.device, dtype=torch.int64
858 + ).view(1, 1, state_capacity)
859 + update_lengths = (self.conv_state_len + query_lengths - 1).view(
860 + num_reqs, 1, 1
861 + )
862 + update_mask = valid_state.view(num_reqs, 1, 1) & (
863 + state_positions < update_lengths
864 + )
865 + existing_state = cached_state[..., :state_capacity]
866 + next_state = torch.where(
867 + update_mask,
868 + candidate_state.to(conv_state.dtype),
869 + existing_state,
870 + )
871 + cached_state[..., :state_capacity] = next_state
872 + conv_state.index_copy_(0, state_indices, cached_state)
873 +
874 + return output
875 +
876 + def _short_conv_dilated_dispatch(
877 + self,
878 + inputs: torch.Tensor,
879 + metadata: PleShortConvAttentionMetadata,
880 + conv_state: torch.Tensor,
881 + conv_weights: torch.Tensor,
882 + ) -> torch.Tensor:
883 + num_prefills = metadata.num_prefills
884 + num_decodes = metadata.num_decodes
885 + num_decode_tokens = metadata.num_decode_tokens
886 + num_prefill_tokens = metadata.num_prefill_tokens
887 + has_prefill = num_prefills > 0
888 + has_decode = num_decodes > 0
889 + has_spec = metadata.spec_sequence_masks is not None
890 + x = inputs[: metadata.num_actual_tokens]
891 +
892 + # Split spec / non-spec tokens.
893 + if has_spec:
894 + if has_prefill or has_decode:
895 + assert metadata.spec_token_indx is not None
896 + assert metadata.non_spec_token_indx is not None
897 + x_spec = x.index_select(0, metadata.spec_token_indx.long())
898 + x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long())
899 + else:
900 + x_spec = x
901 + x_non_spec = None
902 + else:
903 + x_spec = None
904 + x_non_spec = x
905 +
906 + spec_output = None
907 + # 1. Run the multi-query speculative-decode part.
908 + if has_spec:
909 + assert metadata.spec_state_indices_tensor is not None
910 + assert metadata.spec_query_start_loc is not None
911 + assert metadata.num_accepted_tokens is not None
912 + spec_output = self._short_conv_dilated_spec_batched(
913 + x_spec=x_spec,
914 + conv_state=conv_state,
915 + conv_weights=conv_weights,
916 + spec_state_indices_tensor=metadata.spec_state_indices_tensor[
917 + : metadata.num_spec_decodes
918 + ],
919 + spec_query_start_loc=metadata.spec_query_start_loc,
920 + num_accepted_tokens=metadata.num_accepted_tokens,
921 + spec_query_len=metadata.spec_query_len,
922 + )
923 +
924 + # 2. Run regular decode and prefill requests.
925 + conv_out_non_spec = None
926 + state_indices_tensor = metadata.state_indices_tensor
927 + if x_non_spec is not None:
928 + assert state_indices_tensor is not None
929 + if has_prefill:
930 + state_indices_tensor_d, state_indices_tensor_p = torch.split(
931 + state_indices_tensor,
932 + [num_decodes, num_prefills],
933 + dim=0,
934 + )
935 + x_d, x_p = torch.split(
936 + x_non_spec,
937 + [num_decode_tokens, num_prefill_tokens],
938 + dim=0,
939 + )
940 + non_spec_parts: list[torch.Tensor] = []
941 + if has_decode:
942 + non_spec_parts.append(
943 + self._short_conv_dilated_decode_batched(
944 + x_d=x_d,
945 + conv_state=conv_state,
946 + conv_weights=conv_weights,
947 + state_indices_tensor_d=state_indices_tensor_d,
948 + has_initial_states_d=metadata.has_initial_states_d,
949 + )
950 + )
951 + non_spec_parts.append(
952 + self._short_conv_dilated_prefill_batched(
953 + x_p=x_p,
954 + metadata=metadata,
955 + conv_state=conv_state,
956 + conv_weights=conv_weights,
957 + state_indices_tensor_p=state_indices_tensor_p,
958 + num_prefills=num_prefills,
959 + num_decode_tokens=num_decode_tokens,
960 + num_prefill_tokens=num_prefill_tokens,
961 + )
962 + )
963 + conv_out_non_spec = torch.vstack(non_spec_parts)
964 + else:
965 + conv_out_non_spec = self._short_conv_dilated_decode_batched(
966 + x_d=x_non_spec,
967 + conv_state=conv_state,
968 + conv_weights=conv_weights,
969 + state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)],
970 + has_initial_states_d=metadata.has_initial_states_d,
971 + )
972 +
973 + # 3. Merge both parts back into the original token order.
974 + if has_spec and conv_out_non_spec is not None:
975 + assert metadata.spec_token_indx is not None
976 + assert metadata.non_spec_token_indx is not None
977 + assert spec_output is not None
978 + output = x.new_empty((metadata.num_actual_tokens, x.size(-1)))
979 + output.index_copy_(0, metadata.spec_token_indx, spec_output)
980 + output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec)
981 + return output
982 + elif has_spec:
983 + assert spec_output is not None
984 + return spec_output
985 + if conv_out_non_spec is None:
986 + return x
987 + return conv_out_non_spec
988 +
989 + def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor:
990 + forward_context = get_forward_context()
991 + attn_metadata = forward_context.attn_metadata
992 + if attn_metadata is None:
993 + return self._short_conv_fallback(inputs)
994 +
995 + if not isinstance(attn_metadata, dict):
996 + raise RuntimeError(
997 + "PLE short-conv expects per-layer attention metadata dict "
998 + f"during inference, got {type(attn_metadata).__name__}."
999 + )
1000 +
1001 + layer_attn_metadata = attn_metadata.get(self.prefix)
1002 + if layer_attn_metadata is None:
1003 + # MRV2 omits Mamba-family metadata during profile warmup.
1004 + return self._short_conv_fallback(inputs)
1005 + if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
1006 + raise TypeError(
1007 + "Expected PleShortConvAttentionMetadata for layer "
1008 + f"'{self.prefix}', got "
1009 + f"{type(layer_attn_metadata).__name__}."
1010 + )
1011 +
1012 + conv_state = self.kv_cache[0]
1013 + if not is_conv_state_dim_first():
1014 + conv_state = conv_state.transpose(-1, -2)
1015 + conv_weights = self.conv1d.weight.squeeze(1)
1016 +
1017 + state_capacity = self.conv_state_len + self.num_spec_tokens
1018 + if state_capacity > 0:
1019 + if conv_state.size(-1) < state_capacity:
1020 + raise RuntimeError(
1021 + "PLE short-conv cache is smaller than expected for "
1022 + f"dilated convolution: got {conv_state.size(-1)}, "
1023 + f"expect at least {state_capacity}."
1024 + )
1025 + conv_state = conv_state[..., -state_capacity:]
1026 + return self._short_conv_dilated_dispatch(
1027 + inputs,
1028 + layer_attn_metadata,
1029 + conv_state,
1030 + conv_weights.to(dtype=inputs.dtype),
1031 + )
1032 +
1033 + def forward(
1034 + self,
1035 + hidden_states: torch.Tensor,
1036 + input_ids: torch.Tensor,
1037 + query_start_loc: torch.Tensor,
1038 + ngram_context: torch.Tensor,
1039 + ) -> torch.Tensor:
1040 + input_ids = input_ids.reshape(-1)
1041 + if input_ids.shape[0] != hidden_states.shape[0]:
1042 + raise ValueError(
1043 + "PLE expects input_ids and hidden_states to have the same "
1044 + f"token length, got {input_ids.shape[0]} and "
1045 + f"{hidden_states.shape[0]}"
1046 + )
1047 + embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
1048 + key, _ = self.key_proj(embeddings)
1049 + value, _ = self.value_proj(embeddings)
1050 + token_count = hidden_states.shape[0]
1051 + key = key.reshape(token_count, self.hc_count, self.hidden_size)
1052 + query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size)
1053 + key = self._apply_norm(self.norm_key, key)
1054 + query = self._apply_norm(self.norm_query, query)
1055 + gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size)
1056 + gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt())
1057 + gated_value = gate * value.unsqueeze(-2)
1058 + normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2)
1059 + conv_output = torch.zeros_like(normalized)
1060 + torch.ops.vllm.qwen4_exp_ple_short_conv(
1061 + normalized,
1062 + conv_output,
1063 + self.prefix,
1064 + )
1065 + return gated_value.flatten(-2) + conv_output
1066 +
1067 +
1068 +def qwen4_exp_amd_ple_ngram_embedding(
1069 + ngram_ids: torch.Tensor,
1070 + output: torch.Tensor,
1071 + layer_name: str,
1072 +) -> None:
1073 + """Run the large PLE embedding lookup outside Inductor's FX graph.
1074 +
1075 + Keeping the embedding weight in ``static_forward_context`` prevents AOT
1076 + compile-time autotuning from materializing a synthetic copy of the weight.
1077 + """
1078 + layer = get_forward_context().no_compile_layers[layer_name]
1079 + if not isinstance(layer, Qwen4ExpPLELayer):
1080 + raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner")
1081 + result = layer.ple_embedding.ngram_embedding(ngram_ids).flatten(-2)
1082 + output.copy_(result)
1083 +
1084 +
1085 +def qwen4_exp_amd_ple_ngram_embedding_fake(
1086 + ngram_ids: torch.Tensor,
1087 + output: torch.Tensor,
1088 + layer_name: str,
1089 +) -> None:
1090 + return
1091 +
1092 +
1093 +def qwen4_exp_ple_short_conv(
1094 + inputs: torch.Tensor,
1095 + output: torch.Tensor,
1096 + layer_name: str,
1097 +) -> None:
1098 + layer = get_forward_context().no_compile_layers[layer_name]
1099 + result = layer._short_conv(inputs)
1100 + output[: result.shape[0]].copy_(result)
1101 +
1102 +
1103 +def qwen4_exp_ple_short_conv_fake(
1104 + inputs: torch.Tensor,
1105 + output: torch.Tensor,
1106 + layer_name: str,
1107 +) -> None:
1108 + return
1109 +
1110 +
1111 +direct_register_custom_op(
1112 + op_name="qwen4_exp_amd_ple_ngram_embedding",
1113 + op_func=qwen4_exp_amd_ple_ngram_embedding,
1114 + mutates_args=["output"],
1115 + fake_impl=qwen4_exp_amd_ple_ngram_embedding_fake,
1116 +)
1117 +
1118 +
1119 +direct_register_custom_op(
1120 + op_name="qwen4_exp_ple_short_conv",
1121 + op_func=qwen4_exp_ple_short_conv,
1122 + mutates_args=["output"],
1123 + fake_impl=qwen4_exp_ple_short_conv_fake,
1124 +)
1125 +
1126 +
1127 +__all__ = [
1128 + "Qwen4ExpNGramEmbedding",
1129 + "Qwen4ExpPLEGroupedNorm",
1130 + "Qwen4ExpPLELayer",
1131 +]
vllm/models/qwen4_exp/amd/qsa.py new
+474
@@ -0,0 +1,474 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""AMD ROCm QSA owner with Triton kernels."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import ClassVar, cast
8 +
9 +import torch
10 +from torch import nn
11 +
12 +from vllm.config import VllmConfig
13 +from vllm.config.cache import CacheDType
14 +from vllm.distributed import get_tensor_model_parallel_world_size
15 +from vllm.forward_context import get_forward_context
16 +from vllm.model_executor.layers.attention.attention import (
17 + set_default_quant_scales,
18 +)
19 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
20 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
21 +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
22 +from vllm.model_executor.layers.quantization import QuantizationConfig
23 +from vllm.model_executor.layers.rotary_embedding import get_rope
24 +from vllm.model_executor.models.qwen3_next import Qwen3NextAttention
25 +from vllm.platforms import current_platform
26 +from vllm.transformers_utils.configs.qwen4_exp import (
27 + Qwen4ExpTextConfig,
28 +)
29 +from vllm.utils.torch_utils import (
30 + LayerNameType,
31 + _encode_layer_name,
32 + _resolve_layer_name,
33 + canonicalize_singleton_dim_strides,
34 + direct_register_custom_op,
35 + kv_cache_dtype_str_to_dtype,
36 +)
37 +from vllm.v1.attention.backend import (
38 + AttentionBackend,
39 + AttentionCGSupport,
40 + AttentionType,
41 + MultipleOf,
42 +)
43 +from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
44 +from vllm.v1.attention.backends.flash_attn import (
45 + FlashAttentionBackend,
46 + FlashAttentionImpl,
47 + FlashAttentionMetadata,
48 + FlashAttentionMetadataBuilder,
49 +)
50 +from vllm.v1.kv_cache_interface import (
51 + FullAttentionSpec,
52 + KVCacheSpec,
53 + get_kv_quant_mode,
54 +)
55 +
56 +from ..common.qsa_cache import QSAForwardMetadata
57 +from . import model
58 +from .indexer_qsa import QSAIndexer
59 +
60 +
61 +class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
62 + """Flash metadata supporting uniform decode and target-verify graphs."""
63 +
64 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
65 +
66 +
67 +class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
68 + """FullAttentionSpec backend used by the merged QSA owner."""
69 +
70 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
71 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
72 +
73 + @staticmethod
74 + def get_name() -> str:
75 + return "QWEN4_EXP_QSA_TRITON"
76 +
77 + @staticmethod
78 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
79 + # QSA consumes manager pages directly and does not use FA4 paged attention.
80 + return [MultipleOf(16)]
81 +
82 + @staticmethod
83 + def get_impl_cls() -> type[Qwen4ExpQSAFlashAttentionImpl]:
84 + return Qwen4ExpQSAFlashAttentionImpl
85 +
86 + @staticmethod
87 + def get_builder_cls() -> type[Qwen4ExpQSAMetadataBuilder]:
88 + return Qwen4ExpQSAMetadataBuilder
89 +
90 + @classmethod
91 + def is_sparse(cls) -> bool:
92 + return True
93 +
94 + @classmethod
95 + def supports_kv_connector(cls) -> bool:
96 + return False
97 +
98 +
99 +class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
100 + """Run paged sparse GQA with the QSA Triton kernel."""
101 +
102 + supports_dcp: bool = False
103 + supports_pcp: bool = False
104 +
105 + def __init__(self, *args, **kwargs) -> None:
106 + super().__init__(*args, **kwargs)
107 + if not is_flash_attn_varlen_func_available():
108 + raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")
109 + if self.dcp_world_size != 1:
110 + raise NotImplementedError(
111 + "Qwen4Exp QSA does not support decode context parallelism"
112 + )
113 + if self.kv_cache_dtype not in ("auto", "bfloat16"):
114 + raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
115 + self.supports_quant_query_input = False
116 +
117 + def forward_qsa(
118 + self,
119 + layer: torch.nn.Module,
120 + query: torch.Tensor,
121 + key: torch.Tensor,
122 + value: torch.Tensor,
123 + kv_cache: torch.Tensor,
124 + attn_metadata: FlashAttentionMetadata,
125 + output: torch.Tensor,
126 + token_to_req: torch.Tensor,
127 + output_scale: torch.Tensor | None = None,
128 + output_block_scale: torch.Tensor | None = None,
129 + ) -> torch.Tensor:
130 + del key, value
131 + if output_scale is not None or output_block_scale is not None:
132 + raise NotImplementedError("QSA does not support fused output quantization")
133 + if self.alibi_slopes is not None or self.sinks is not None:
134 + raise NotImplementedError("QSA does not support ALiBi or attention sinks")
135 + if self.sliding_window != (-1, -1):
136 + raise NotImplementedError("QSA does not support sliding-window attention")
137 +
138 + num_tokens = attn_metadata.num_actual_tokens
139 + output.zero_()
140 + if num_tokens == 0:
141 + return output
142 +
143 + topk_buffer = getattr(layer, "topk_indices_buffer", None)
144 + if topk_buffer is None:
145 + raise RuntimeError("QSA owner did not provide its top-k buffer")
146 + logical_indices = topk_buffer[:num_tokens]
147 + token_to_req = token_to_req[:num_tokens]
148 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
149 + key_cache = canonicalize_singleton_dim_strides(key_cache)
150 + value_cache = canonicalize_singleton_dim_strides(value_cache)
151 + if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:
152 + raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")
153 +
154 + from .ops.qsa import qsa_sparse_paged_attention
155 +
156 + qsa_sparse_paged_attention(
157 + query[:num_tokens],
158 + key_cache,
159 + value_cache,
160 + logical_indices,
161 + attn_metadata.block_table,
162 + token_to_req,
163 + output[:num_tokens],
164 + )
165 + return output
166 +
167 +
168 +class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
169 + """Merged Qwen full-attention owner with a QSA index side branch."""
170 +
171 + supports_dcp = False
172 +
173 + def __init__(
174 + self,
175 + *,
176 + vllm_config: VllmConfig,
177 + config: Qwen4ExpTextConfig,
178 + layer_id: int,
179 + quant_config: QuantizationConfig | None = None,
180 + reduce_results: bool = True,
181 + prefix: str = "",
182 + ) -> None:
183 + nn.Module.__init__(self)
184 + cache_config = vllm_config.cache_config
185 + model_config = vllm_config.model_config
186 + if cache_config is None:
187 + raise ValueError("Qwen4Exp QSA requires a paged KV cache")
188 + if model_config.dtype != torch.bfloat16:
189 + raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
190 + if cache_config.cache_dtype not in ("auto", "bfloat16"):
191 + raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
192 + if getattr(quant_config, "kv_cache_scheme", None) is not None:
193 + raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
194 + parallel_config = vllm_config.parallel_config
195 + if (
196 + parallel_config.prefill_context_parallel_size > 1
197 + or parallel_config.decode_context_parallel_size > 1
198 + ):
199 + raise NotImplementedError(
200 + "Qwen4Exp QSA does not support context parallelism"
201 + )
202 + if not getattr(config, "is_causal", True):
203 + raise NotImplementedError("Qwen4Exp QSA requires causal decoder attention")
204 +
205 + self.config = config
206 + self.hidden_size = int(config.hidden_size)
207 + tp_size = get_tensor_model_parallel_world_size()
208 + self.total_num_heads = int(config.num_attention_heads)
209 + if self.total_num_heads % tp_size:
210 + raise ValueError("QSA attention heads must be divisible by TP size")
211 + self.num_heads = self.total_num_heads // tp_size
212 + self.total_num_kv_heads = int(config.num_key_value_heads)
213 + if self.total_num_kv_heads >= tp_size:
214 + if self.total_num_kv_heads % tp_size:
215 + raise ValueError("QSA KV heads must be divisible by TP size")
216 + elif tp_size % self.total_num_kv_heads:
217 + raise ValueError("TP size must be divisible by replicated QSA KV heads")
218 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
219 + self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads)
220 + self.q_size = self.num_heads * self.head_dim
221 + self.kv_size = self.num_kv_heads * self.head_dim
222 + self.scaling = self.head_dim**-0.5
223 + self.dual_chunk_attention_config = getattr(
224 + config, "dual_chunk_attention_config", None
225 + )
226 + if self.dual_chunk_attention_config is not None:
227 + raise NotImplementedError("Qwen4Exp QSA does not support dual-chunk RoPE")
228 + # Qwen4Exp full-attention checkpoints always pack a sigmoid output
229 + # gate next to Q, even when an inherited config default says otherwise.
230 + self.attn_output_gate = True
231 +
232 + self.qkv_proj = QKVParallelLinear(
233 + self.hidden_size,
234 + self.head_dim,
235 + self.total_num_heads * (1 + self.attn_output_gate),
236 + self.total_num_kv_heads,
237 + bias=False,
238 + quant_config=model.without_modelopt_fp4(quant_config),
239 + prefix=f"{prefix}.qkv_proj",
240 + )
241 + self.o_proj = RowParallelLinear(
242 + self.total_num_heads * self.head_dim,
243 + self.hidden_size,
244 + bias=False,
245 + reduce_results=reduce_results,
246 + quant_config=quant_config,
247 + prefix=f"{prefix}.o_proj",
248 + )
249 + self.rotary_emb = get_rope(
250 + head_size=self.head_dim,
251 + max_position=config.max_position_embeddings,
252 + rope_parameters=config.rope_parameters,
253 + )
254 + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
255 + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
256 +
257 + mm_config = model_config.multimodal_config
258 + text_only = mm_config is None or mm_config.language_model_only
259 + self.use_fused_qk_norm_rope_gate = (
260 + self.attn_output_gate
261 + and getattr(self.rotary_emb, "is_neox_style", False)
262 + and current_platform.is_cuda()
263 + and text_only
264 + )
265 +
266 + self.layer_name = f"{prefix}.attn"
267 + self.attn_type = AttentionType.DECODER
268 + self.kv_cache_dtype = cache_config.cache_dtype
269 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
270 + self.kv_cache_dtype, model_config
271 + )
272 + if self.kv_cache_torch_dtype != torch.bfloat16:
273 + raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")
274 + self.kv_sharing_target_layer_name = None
275 + self.kv_cache = torch.tensor([])
276 + set_default_quant_scales(self, register_buffer=True)
277 +
278 + self.attn_backend = Qwen4ExpQSAFlashAttentionBackend
279 + self.impl = Qwen4ExpQSAFlashAttentionImpl(
280 + self.num_heads,
281 + self.head_dim,
282 + self.scaling,
283 + self.num_kv_heads,
284 + None,
285 + None,
286 + self.kv_cache_dtype,
287 + None,
288 + AttentionType.DECODER,
289 + None,
290 + )
291 + self.indexer = QSAIndexer(
292 + vllm_config=vllm_config,
293 + config=config,
294 + layer_id=layer_id,
295 + rotary_emb=self.rotary_emb,
296 + quant_config=quant_config,
297 + prefix=f"{prefix}.indexer",
298 + )
299 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
300 + self.register_buffer(
301 + "topk_indices_buffer",
302 + torch.empty(
303 + max_tokens,
304 + self.indexer.output_width,
305 + dtype=torch.int32,
306 + ),
307 + persistent=False,
308 + )
309 +
310 + static_context = vllm_config.compilation_config.static_forward_context
311 + if self.layer_name in static_context:
312 + raise ValueError(f"Duplicate layer name: {self.layer_name}")
313 + static_context[self.layer_name] = self
314 +
315 + def get_attn_backend(self) -> type[AttentionBackend]:
316 + return self.attn_backend
317 +
318 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
319 + return FullAttentionSpec(
320 + block_size=vllm_config.cache_config.block_size,
321 + num_kv_heads=self.num_kv_heads,
322 + head_size=self.head_dim,
323 + head_size_v=self.head_dim,
324 + dtype=self.kv_cache_torch_dtype,
325 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
326 + )
327 +
328 + def _run_qsa(
329 + self,
330 + hidden_states: torch.Tensor,
331 + positions: torch.Tensor,
332 + query: torch.Tensor,
333 + key: torch.Tensor,
334 + value: torch.Tensor,
335 + output: torch.Tensor,
336 + ) -> None:
337 + metadata = get_forward_context().attn_metadata
338 + if isinstance(metadata, list):
339 + metadata = metadata[0]
340 + if not isinstance(metadata, dict):
341 + output.zero_()
342 + return
343 + main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name])
344 + if self.kv_cache.numel() == 0:
345 + raise RuntimeError("QSA main K/V cache is not bound")
346 +
347 + num_tokens = main_metadata.num_actual_tokens
348 + side_metadata = cast(
349 + QSAForwardMetadata,
350 + metadata[self.indexer.raw_key_cache.prefix],
351 + )
352 + if side_metadata.num_actual_tokens != num_tokens:
353 + raise RuntimeError("QSA main and side metadata token counts disagree")
354 + selected = self.indexer(
355 + hidden_states,
356 + positions,
357 + self.topk_indices_buffer[:num_tokens],
358 + )
359 + if selected.shape != (
360 + num_tokens,
361 + self.indexer.output_width,
362 + ):
363 + raise RuntimeError("QSA indexer returned an invalid selection shape")
364 + impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
365 + impl.do_kv_cache_update(
366 + self,
367 + key,
368 + value,
369 + self.kv_cache,
370 + main_metadata.slot_mapping,
371 + )
372 + impl.forward_qsa(
373 + self,
374 + query,
375 + key,
376 + value,
377 + self.kv_cache,
378 + main_metadata,
379 + output,
380 + token_to_req=side_metadata.token_to_req,
381 + )
382 +
383 + def forward(
384 + self,
385 + positions: torch.Tensor,
386 + hidden_states: torch.Tensor,
387 + ) -> torch.Tensor:
388 + qkv, _ = self.qkv_proj(hidden_states)
389 + q, k, v, gate = self._project_qkv_gate(qkv, positions)
390 + num_tokens = hidden_states.shape[0]
391 + query = q.view(num_tokens, self.num_heads, self.head_dim)
392 + key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
393 + value = v.view(num_tokens, self.num_kv_heads, self.head_dim)
394 + attn_output = torch.empty_like(query)
395 + encoded_layer_name = _encode_layer_name(self.layer_name)
396 + if current_platform.opaque_attention_op():
397 + torch.ops.vllm.qwen4_exp_qsa_with_output(
398 + hidden_states,
399 + positions,
400 + query,
401 + key,
402 + value,
403 + attn_output,
404 + encoded_layer_name,
405 + )
406 + else:
407 + qwen4_exp_qsa_with_output(
408 + hidden_states,
409 + positions,
410 + query,
411 + key,
412 + value,
413 + attn_output,
414 + encoded_layer_name,
415 + )
416 + flat_output = attn_output.view(num_tokens, -1)
417 + if gate is not None:
418 + flat_output = flat_output * torch.sigmoid(gate)
419 + output, _ = self.o_proj(flat_output)
420 + return output
421 +
422 +
423 +def qwen4_exp_qsa_with_output(
424 + hidden_states: torch.Tensor,
425 + positions: torch.Tensor,
426 + query: torch.Tensor,
427 + key: torch.Tensor,
428 + value: torch.Tensor,
429 + output: torch.Tensor,
430 + layer_name: LayerNameType,
431 +) -> None:
432 + """Run the complete QSA state/update/attend transaction."""
433 +
434 + layer_name = _resolve_layer_name(layer_name)
435 + layer = get_forward_context().no_compile_layers[layer_name]
436 + if not isinstance(layer, Qwen4ExpQSAAttention):
437 + raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner")
438 + layer._run_qsa(
439 + hidden_states,
440 + positions,
441 + query,
442 + key,
443 + value,
444 + output,
445 + )
446 +
447 +
448 +def qwen4_exp_qsa_with_output_fake(
449 + hidden_states: torch.Tensor,
450 + positions: torch.Tensor,
451 + query: torch.Tensor,
452 + key: torch.Tensor,
453 + value: torch.Tensor,
454 + output: torch.Tensor,
455 + layer_name: LayerNameType,
456 +) -> None:
457 + del hidden_states, positions, query, key, value, output, layer_name
458 +
459 +
460 +direct_register_custom_op(
461 + op_name="qwen4_exp_qsa_with_output",
462 + op_func=qwen4_exp_qsa_with_output,
463 + mutates_args=["output"],
464 + fake_impl=qwen4_exp_qsa_with_output_fake,
465 +)
466 +
467 +
468 +__all__ = [
469 + "QSAIndexer",
470 + "Qwen4ExpQSAAttention",
471 + "Qwen4ExpQSAFlashAttentionBackend",
472 + "Qwen4ExpQSAFlashAttentionImpl",
473 + "qwen4_exp_qsa_with_output",
474 +]
vllm/models/qwen4_exp/common/__init__.py new
+17
@@ -0,0 +1,17 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Common Qwen4Exp model components."""
4 +
5 +from .hyperconnection import (
6 + GatedResidual,
7 + GroupedGemmaRMSNorm,
8 + HyperConnectionBase,
9 + HyperConnectionConfig,
10 +)
11 +
12 +__all__ = [
13 + "GatedResidual",
14 + "GroupedGemmaRMSNorm",
15 + "HyperConnectionBase",
16 + "HyperConnectionConfig",
17 +]
vllm/models/qwen4_exp/common/hyperconnection.py new
+250
@@ -0,0 +1,250 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HyperConnection (Gated Residual) utilities.
4 +
5 +Implements the HyperConnection residual scheme proposed in
6 +"HyperConnections" (https://arxiv.org/abs/2409.19606).
7 +
8 +The two concrete variants are:
9 + - ``HyperConnectionBase`` - simple average pooling across hc_count parallel
10 + streams (equivalent to hyperconnection_average).
11 + - ``GatedResidual`` - learnable low-rank gated mixing and injection
12 + (gated_residual).
13 +
14 +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner
15 +(HC outer, HS inner — checkpoint-native layout). The local torch
16 +implementation consumes the hyper input viewed as ``[..., HC, HS]``.
17 +
18 +Typical usage inside a transformer decoder layer::
19 +
20 + self.attn_hc = GatedResidual(hc_config, role="attn")
21 + self.mlp_hc = GatedResidual(hc_config, role="mlp")
22 +
23 + hidden_states, residual = self.attn_hc.mix(hidden_states)
24 + hidden_states = attention(hidden_states)
25 + hidden_states = self.attn_hc.combine(hidden_states, residual)
26 +
27 + hidden_states, residual = self.mlp_hc.mix(hidden_states)
28 + hidden_states = mlp(hidden_states)
29 + hidden_states = self.mlp_hc.combine(hidden_states, residual)
30 +"""
31 +
32 +from dataclasses import dataclass
33 +
34 +import torch
35 +import torch.nn.functional as F
36 +from torch import nn
37 +
38 +
39 +# ---------------------------------------------------------------------------
40 +# Config
41 +# ---------------------------------------------------------------------------
42 +@dataclass
43 +class HyperConnectionConfig:
44 + """Configuration shared by all HyperConnection variants."""
45 +
46 + hc_count: int = 4
47 + hidden_size: int = 64
48 + params_dtype: torch.dtype = torch.bfloat16
49 + mtp_hc: bool = False
50 + hc_lowrank: int = 16
51 + rms_norm_eps: float = 1e-6
52 + hc_per_branch_norm: bool = False
53 +
54 +
55 +class GroupedGemmaRMSNorm(nn.Module):
56 + def __init__(
57 + self,
58 + hidden_size: int,
59 + eps: float,
60 + group_size: int | None,
61 + dtype: torch.dtype | None,
62 + ) -> None:
63 + super().__init__()
64 + if group_size is not None and hidden_size % group_size:
65 + raise ValueError(
66 + f"hidden_size ({hidden_size}) must be divisible by "
67 + f"group_size ({group_size})"
68 + )
69 + self.variance_epsilon = eps
70 + self.group_size = group_size
71 + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype))
72 +
73 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
74 + input_dtype = hidden_states.dtype
75 + hidden_states = hidden_states.float()
76 + if self.group_size is None:
77 + variance = hidden_states.square().mean(dim=-1, keepdim=True)
78 + normalized = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
79 + else:
80 + grouped = hidden_states.unflatten(
81 + -1, (hidden_states.shape[-1] // self.group_size, self.group_size)
82 + )
83 + variance = grouped.square().mean(dim=-1, keepdim=True)
84 + normalized = (
85 + grouped * torch.rsqrt(variance + self.variance_epsilon)
86 + ).flatten(-2)
87 + return (normalized * (1.0 + self.weight.float())).to(input_dtype)
88 +
89 +
90 +# ---------------------------------------------------------------------------
91 +# Average-pooling variant
92 +# ---------------------------------------------------------------------------
93 +class HyperConnectionBase(nn.Module):
94 + """Average-pooling HyperConnection (``hyperconnection_average``).
95 +
96 + Splits the incoming ``[..., HC*HS]`` tensor (HC outer, HS inner) into
97 + ``HC`` parallel streams, averages them for the block input, and
98 + broadcasts the block output back to every stream.
99 + """
100 +
101 + def __init__(
102 + self,
103 + config: HyperConnectionConfig,
104 + layer_idx: int | None = None,
105 + role: str | None = None,
106 + ) -> None:
107 + super().__init__()
108 + self.config = config
109 + self.hc_count = config.hc_count
110 + self.hidden_size = config.hidden_size
111 + self.layer_idx = layer_idx
112 + self.role = role
113 +
114 + @property
115 + def hyper_hidden_size(self) -> int:
116 + return self.hc_count * self.hidden_size
117 +
118 + def mix(self, hyper_input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
119 + """Average the HC streams into a single block input."""
120 + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size
121 + # [*, HC, HS] — mean over HC (dim=-2).
122 + unflat = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size))
123 + mixed_input = unflat.mean(dim=-2)
124 + return mixed_input, hyper_input
125 +
126 + def combine(
127 + self, block_output: torch.Tensor, residual: torch.Tensor
128 + ) -> torch.Tensor:
129 + """Broadcast the block output back to every stream."""
130 + assert residual.shape[-1] == self.hc_count * self.hidden_size
131 + assert block_output.shape[-1] == self.hidden_size
132 + residual_reshaped = residual.unflatten(-1, (self.hc_count, self.hidden_size))
133 + combined = residual_reshaped + block_output.unsqueeze(-2)
134 + return combined.flatten(-2)
135 +
136 +
137 +# ---------------------------------------------------------------------------
138 +# Gated-residual variant
139 +# ---------------------------------------------------------------------------
140 +class GatedResidual(HyperConnectionBase):
141 + """Gated HyperConnection with learnable low-rank mixing and injection.
142 +
143 + ``mix()`` applies GemmaRMSNorm per HC stream and projects through a
144 + low-rank sigmoid gate to produce a single block input. ``combine()``
145 + injects the block output back into each stream through a learned
146 + per-stream injection weight.
147 +
148 + This implementation uses only PyTorch operators. Tensor-parallel
149 + collectives are supplied by its caller.
150 + """
151 +
152 + def __init__(
153 + self,
154 + config: HyperConnectionConfig,
155 + layer_idx: int | None = None,
156 + role: str | None = None,
157 + use_mix: bool = True,
158 + use_combine: bool = True,
159 + ) -> None:
160 + super().__init__(config, layer_idx, role)
161 + norm_size = (
162 + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size
163 + )
164 + group_size = config.hidden_size if config.hc_per_branch_norm else None
165 + # Normalize each H-sized HC stream independently while retaining a
166 + # separate affine weight for every element of the HC*H layout.
167 + self.hc_norm = GroupedGemmaRMSNorm(
168 + norm_size,
169 + eps=config.rms_norm_eps,
170 + group_size=group_size,
171 + dtype=config.params_dtype,
172 + )
173 +
174 + # -- raw Linear weights (checkpoint-compatible) ----------------------
175 + if use_mix:
176 + self.input_mix_weight_down = nn.Linear(
177 + self.hyper_hidden_size,
178 + config.hc_lowrank,
179 + bias=False,
180 + dtype=config.params_dtype,
181 + )
182 + self.input_mix_weight_up = nn.Linear(
183 + config.hc_lowrank,
184 + self.hyper_hidden_size,
185 + bias=False,
186 + dtype=config.params_dtype,
187 + )
188 + if use_combine:
189 + self.block_inject_weight = nn.Linear(
190 + self.hyper_hidden_size,
191 + self.hc_count,
192 + bias=False,
193 + dtype=config.params_dtype,
194 + )
195 +
196 + def _normalize(self, hyper_input: torch.Tensor) -> torch.Tensor:
197 + if self.config.hc_per_branch_norm:
198 + return self.hc_norm(hyper_input)
199 + return self.hc_norm(
200 + hyper_input.unflatten(-1, (self.hc_count, self.hidden_size))
201 + ).flatten(-2)
202 +
203 + def mix(
204 + self, hyper_input: torch.Tensor
205 + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
206 + """Mix: RMSNorm -> low-rank gate -> gated mean."""
207 + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size
208 + if not hasattr(self, "input_mix_weight_down"):
209 + raise RuntimeError("mix was disabled for this hyper-connection")
210 + hyper_input_normed = self._normalize(hyper_input)
211 + # Gate — original mix order: linear+silu then linear+sigmoid.
212 + gate = F.silu(
213 + F.linear(hyper_input_normed, self.input_mix_weight_down.weight)
214 + / self.hc_count
215 + )
216 + gate = torch.sigmoid(F.linear(gate, self.input_mix_weight_up.weight)).unflatten(
217 + -1, (self.hc_count, self.hidden_size)
218 + )
219 + mixed_input = (
220 + gate * hyper_input_normed.unflatten(-1, (self.hc_count, self.hidden_size))
221 + ).mean(dim=-2)
222 + return mixed_input.to(hyper_input.dtype), (hyper_input, hyper_input_normed)
223 +
224 + def combine(
225 + self,
226 + block_output: torch.Tensor,
227 + residuals: tuple[torch.Tensor, torch.Tensor],
228 + ) -> torch.Tensor:
229 + if not hasattr(self, "block_inject_weight"):
230 + raise RuntimeError("combine was disabled for this hyper-connection")
231 + hyper_input, hyper_input_normed = residuals
232 + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size
233 + assert block_output.shape[-1] == self.hidden_size
234 + residual = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size))
235 + # The paired mix keeps its normalized hyper input so combine uses the
236 + # same HC module's injection weight.
237 + injection_weight = 2.0 * torch.sigmoid(
238 + F.linear(hyper_input_normed, self.block_inject_weight.weight)
239 + / self.hc_count
240 + )
241 + output = residual + block_output.unsqueeze(-2) * injection_weight.unsqueeze(-1)
242 + return output.flatten(-2).to(hyper_input.dtype)
243 +
244 +
245 +__all__ = [
246 + "GatedResidual",
247 + "GroupedGemmaRMSNorm",
248 + "HyperConnectionBase",
249 + "HyperConnectionConfig",
250 +]
vllm/models/qwen4_exp/common/ple.py new
+100
@@ -0,0 +1,100 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Common Qwen4Exp PLE helpers."""
4 +
5 +from dataclasses import dataclass
6 +
7 +import torch
8 +
9 +from vllm.model_executor.layers.vocab_parallel_embedding import (
10 + VocabParallelEmbedding,
11 +)
12 +
13 +
14 +@dataclass(frozen=True)
15 +class PLEShardOverlap:
16 + """Source and destination slices for one checkpoint embedding shard."""
17 +
18 + source_start: int
19 + destination_start: int
20 + row_count: int
21 +
22 +
23 +def compute_ple_shard_overlap(
24 + *,
25 + checkpoint_start: int,
26 + checkpoint_rows: int,
27 + tp_start: int,
28 + tp_end: int,
29 +) -> PLEShardOverlap | None:
30 + """Compute the overlap of a checkpoint shard and one TP vocabulary range."""
31 +
32 + if checkpoint_start < 0 or checkpoint_rows < 0:
33 + raise ValueError("checkpoint shard bounds must be non-negative")
34 + if tp_start < 0 or tp_end < tp_start:
35 + raise ValueError("invalid TP vocabulary range")
36 + checkpoint_end = checkpoint_start + checkpoint_rows
37 + overlap_start = max(checkpoint_start, tp_start)
38 + overlap_end = min(checkpoint_end, tp_end)
39 + if overlap_start >= overlap_end:
40 + return None
41 + return PLEShardOverlap(
42 + source_start=overlap_start - checkpoint_start,
43 + destination_start=overlap_start - tp_start,
44 + row_count=overlap_end - overlap_start,
45 + )
46 +
47 +
48 +def copy_ple_embedding_shard_(
49 + destination: torch.Tensor,
50 + loaded_weight: torch.Tensor,
51 + *,
52 + checkpoint_start: int,
53 + tp_start: int,
54 + tp_end: int,
55 +) -> int:
56 + """Copy the overlapping rows of a PLE checkpoint shard into a TP table."""
57 +
58 + if destination.ndim == 0 or loaded_weight.ndim != destination.ndim:
59 + raise ValueError("destination and loaded weight must have matching ranks")
60 + if destination.shape[1:] != loaded_weight.shape[1:]:
61 + raise ValueError(
62 + "embedding shard dimensions do not match: "
63 + f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}"
64 + )
65 + if destination.shape[0] < tp_end - tp_start:
66 + raise ValueError("destination does not cover the requested TP range")
67 + overlap = compute_ple_shard_overlap(
68 + checkpoint_start=checkpoint_start,
69 + checkpoint_rows=loaded_weight.shape[0],
70 + tp_start=tp_start,
71 + tp_end=tp_end,
72 + )
73 + if overlap is None:
74 + return 0
75 + source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count)
76 + target = destination.narrow(0, overlap.destination_start, overlap.row_count)
77 + with torch.no_grad():
78 + target.copy_(source.to(device=target.device, dtype=target.dtype))
79 + return overlap.row_count
80 +
81 +
82 +class PLEVocabParallelEmbedding(VocabParallelEmbedding):
83 + """Vocab-parallel embedding that accepts checkpoint row shards."""
84 +
85 + def weight_loader(
86 + self,
87 + param: torch.Tensor,
88 + loaded_weight: torch.Tensor,
89 + checkpoint_start: int | None = None,
90 + ) -> None:
91 + if checkpoint_start is None:
92 + super().weight_loader(param, loaded_weight)
93 + return
94 + copy_ple_embedding_shard_(
95 + param,
96 + loaded_weight,
97 + checkpoint_start=checkpoint_start,
98 + tp_start=self.shard_indices.org_vocab_start_index,
99 + tp_end=self.shard_indices.org_vocab_end_index,
100 + )
vllm/models/qwen4_exp/common/qsa_cache.py new
+807
@@ -0,0 +1,807 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Paged side-cache ownership and metadata for Qwen4Exp QSA.
4 +
5 +Each QSA layer keeps a fixed circular buffer of raw index keys (the
6 +compressor state) and one compressed key. MRoPE models pack exact three-axis
7 +positions beside the raw keys; text models derive group positions from
8 +logical positions. The compressor state uses one block per request, while
9 + the compressed owner uses ``MLAAttentionSpec.tokens_per_state`` so its block
10 +table follows the main KV-cache lifecycle. Their physical tensor storage is
11 +shared by the generic cache-layout planner.
12 +"""
13 +
14 +from __future__ import annotations
15 +
16 +import math
17 +from dataclasses import dataclass
18 +from functools import cache
19 +from typing import ClassVar
20 +
21 +import torch
22 +from torch import nn
23 +
24 +from vllm.config import CacheConfig, VllmConfig
25 +from vllm.config.cache import CacheDType
26 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
27 +from vllm.platforms import current_platform
28 +from vllm.triton_utils import HAS_TRITON, tl, triton
29 +from vllm.utils.math_utils import cdiv
30 +from vllm.v1.attention.backend import (
31 + AttentionBackend,
32 + AttentionCGSupport,
33 + AttentionMetadata,
34 + AttentionMetadataBuilder,
35 + CommonAttentionMetadata,
36 +)
37 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
38 +from vllm.v1.kv_cache_interface import (
39 + AttentionSpec,
40 + CircularBufferSpec,
41 + KVCacheLayout,
42 + KVCacheSpec,
43 + MLAAttentionSpec,
44 +)
45 +
46 +
47 +def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor:
48 + """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows."""
49 +
50 + if positions.ndim == 1:
51 + positions = positions.unsqueeze(0).expand(3, -1)
52 + elif positions.ndim != 2 or positions.shape[0] not in (1, 3):
53 + raise ValueError("QSA RoPE positions must be [tokens] or [1|3, tokens]")
54 + if positions.shape[0] == 1:
55 + positions = positions.expand(3, -1)
56 + return positions.transpose(0, 1).unsqueeze(1).to(torch.int64)
57 +
58 +
59 +def _logical_positions(
60 + query_start_loc: torch.Tensor,
61 + seq_lens: torch.Tensor,
62 + token_to_req: torch.Tensor,
63 + num_tokens: int,
64 +) -> torch.Tensor:
65 + if num_tokens == 0:
66 + return seq_lens.new_empty((0,), dtype=torch.int64)
67 + arange = torch.arange(num_tokens, device=query_start_loc.device)
68 + requests = token_to_req[:num_tokens].long()
69 + query_lens = torch.diff(query_start_loc)
70 + within_query = arange - query_start_loc.index_select(0, requests)
71 + return (
72 + seq_lens.index_select(0, requests).long()
73 + - query_lens.index_select(0, requests).long()
74 + + within_query.long()
75 + )
76 +
77 +
78 +def _logical_to_physical_qsa_slots(
79 + block_table: torch.Tensor,
80 + request_indices: torch.Tensor,
81 + logical_positions: torch.Tensor,
82 + block_size: int,
83 +) -> torch.Tensor:
84 + if block_size <= 0:
85 + raise ValueError("QSA cache block size must be positive")
86 + if block_table.ndim != 2:
87 + raise ValueError("QSA block table must be two-dimensional")
88 + if request_indices.shape != logical_positions.shape:
89 + request_indices = torch.broadcast_to(request_indices, logical_positions.shape)
90 +
91 + requests = request_indices.to(device=block_table.device, dtype=torch.long)
92 + positions = logical_positions.to(device=block_table.device, dtype=torch.long)
93 + valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0)
94 + logical_blocks = torch.div(
95 + positions.clamp_min(0), block_size, rounding_mode="floor"
96 + )
97 + valid &= logical_blocks < block_table.shape[1]
98 + safe_requests = requests.clamp(0, max(block_table.shape[0] - 1, 0))
99 + safe_blocks = logical_blocks.clamp(0, max(block_table.shape[1] - 1, 0))
100 + if not all(block_table.shape):
101 + return torch.full_like(positions, PAD_SLOT_ID)
102 + physical_blocks = block_table[safe_requests, safe_blocks].long()
103 + valid &= physical_blocks >= 0
104 + slots = physical_blocks * block_size + positions.remainder(block_size)
105 + return torch.where(valid, slots, PAD_SLOT_ID)
106 +
107 +
108 +def circular_qsa_slot_mapping(
109 + block_table: torch.Tensor,
110 + token_to_req: torch.Tensor,
111 + logical_positions: torch.Tensor,
112 + compressor_state_size: int,
113 + query_start_loc: torch.Tensor | None = None,
114 + out: torch.Tensor | None = None,
115 +) -> torch.Tensor:
116 + """Map each request to its fixed physical block as a circular token ring."""
117 +
118 + if compressor_state_size <= 0:
119 + raise ValueError("QSA circular buffer size must be positive")
120 + if block_table.ndim != 2:
121 + raise ValueError("QSA block table must be two-dimensional")
122 +
123 + requests = token_to_req.to(device=block_table.device, dtype=torch.long)
124 + positions = logical_positions.to(device=block_table.device, dtype=torch.long)
125 + if not all(block_table.shape):
126 + slots = torch.full_like(positions, PAD_SLOT_ID)
127 + else:
128 + valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0)
129 + safe_requests = requests.clamp(0, block_table.shape[0] - 1)
130 + physical_blocks = block_table[safe_requests, 0].long()
131 + valid &= physical_blocks >= 0
132 + slots = physical_blocks * compressor_state_size + positions.remainder(
133 + compressor_state_size
134 + )
135 + slots = torch.where(valid, slots, PAD_SLOT_ID)
136 +
137 + if query_start_loc is not None:
138 + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2:
139 + raise ValueError("QSA query starts must contain a terminal offset")
140 + query_start_loc = query_start_loc.to(block_table.device)
141 + num_requests = query_start_loc.shape[0] - 1
142 + safe_requests = requests.clamp(0, num_requests - 1)
143 + request_ends = query_start_loc.index_select(0, safe_requests + 1)
144 + rows = torch.arange(slots.numel(), device=slots.device)
145 + keep = (
146 + (requests >= 0)
147 + & (requests < num_requests)
148 + & (rows + compressor_state_size >= request_ends)
149 + )
150 + slots = torch.where(keep, slots, PAD_SLOT_ID)
151 +
152 + slots = slots.to(torch.int64)
153 + if out is not None:
154 + out.fill_(PAD_SLOT_ID)
155 + out[: slots.numel()].copy_(slots)
156 + return out[: slots.numel()]
157 + return slots
158 +
159 +
160 +def compressed_qsa_slot_mapping(
161 + block_table: torch.Tensor,
162 + token_to_req: torch.Tensor,
163 + logical_positions: torch.Tensor,
164 + storage_block_size: int,
165 + compress_ratio: int,
166 + out: torch.Tensor | None = None,
167 +) -> torch.Tensor:
168 + """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache."""
169 +
170 + if storage_block_size <= 0 or compress_ratio <= 0:
171 + raise ValueError("QSA block size and compression ratio must be positive")
172 + compressed_positions = torch.div(
173 + logical_positions.clamp_min(0), compress_ratio, rounding_mode="floor"
174 + )
175 + slots = _logical_to_physical_qsa_slots(
176 + block_table,
177 + token_to_req,
178 + compressed_positions,
179 + storage_block_size,
180 + )
181 + valid = (logical_positions >= 0) & (
182 + (logical_positions + 1).remainder(compress_ratio) == 0
183 + )
184 + slots = torch.where(valid, slots, PAD_SLOT_ID).to(torch.int64)
185 + if out is not None:
186 + out.fill_(PAD_SLOT_ID)
187 + out[: slots.numel()].copy_(slots)
188 + return out[: slots.numel()]
189 + return slots
190 +
191 +
192 +@cache
193 +def _metadata_launch_pdl() -> bool:
194 + return current_platform.is_arch_support_pdl()
195 +
196 +
197 +@triton.jit(
198 + do_not_specialize=[
199 + "num_reqs",
200 + "num_mapped_tokens",
201 + "num_tokens",
202 + "max_num_work",
203 + "num_search_steps",
204 + "work_search_steps",
205 + ]
206 +)
207 +def _build_qsa_metadata_kernel(
208 + query_start_loc_ptr,
209 + seq_lens_ptr,
210 + common_slot_mapping_ptr,
211 + block_table_ptr,
212 + token_to_req_ptr,
213 + logical_positions_ptr,
214 + slot_mapping_ptr,
215 + k_work_metadata_ptr,
216 + block_table_stride_0: tl.constexpr,
217 + block_table_stride_1: tl.constexpr,
218 + num_reqs,
219 + num_mapped_tokens,
220 + num_tokens,
221 + max_num_work,
222 + num_search_steps,
223 + work_search_steps,
224 + storage_block_size: tl.constexpr,
225 + compress_ratio: tl.constexpr,
226 + circular_buffer_size: tl.constexpr,
227 + num_block_table_columns: tl.constexpr,
228 + launch_pdl: tl.constexpr,
229 + TOKEN_BLOCK_SIZE: tl.constexpr,
230 + REQUEST_SCAN_SIZE: tl.constexpr,
231 + WORK_BLOCK_SIZE: tl.constexpr,
232 +):
233 + if launch_pdl:
234 + tl.extra.cuda.gdc_wait()
235 +
236 + pid = tl.program_id(0)
237 + token_idx = pid * TOKEN_BLOCK_SIZE + tl.arange(0, TOKEN_BLOCK_SIZE)
238 + store_mask = token_idx < num_tokens
239 + mapped = token_idx < num_mapped_tokens
240 + search_token_idx = tl.minimum(token_idx, num_mapped_tokens - 1)
241 + request_idx = tl.zeros((TOKEN_BLOCK_SIZE,), tl.int32)
242 + # Find the last query start at or before each token. The dynamic loop avoids
243 + # compiling a kernel variant for every ceil(log2(num_reqs)).
244 + for step in tl.range(0, num_search_steps):
245 + candidate = request_idx + (1 << (num_search_steps - step - 1))
246 + valid_candidate = candidate < num_reqs
247 + candidate_start = tl.load(
248 + query_start_loc_ptr + candidate,
249 + mask=valid_candidate,
250 + other=num_mapped_tokens + 1,
251 + )
252 + advance = valid_candidate & (candidate_start <= search_token_idx)
253 + request_idx = tl.where(advance, candidate, request_idx)
254 + query_start = tl.load(query_start_loc_ptr + request_idx, mask=mapped, other=0)
255 + query_end = tl.load(query_start_loc_ptr + request_idx + 1, mask=mapped, other=0)
256 + seq_len = tl.load(seq_lens_ptr + request_idx, mask=mapped, other=0)
257 + logical_position = seq_len - (query_end - query_start) + token_idx - query_start
258 + logical_position = tl.where(mapped, logical_position, -1)
259 + tl.store(
260 + token_to_req_ptr + token_idx,
261 + tl.where(mapped, request_idx, 0),
262 + mask=store_mask,
263 + )
264 + tl.store(
265 + logical_positions_ptr + token_idx,
266 + logical_position,
267 + mask=store_mask,
268 + )
269 +
270 + # circular_buffer_size is constexpr, so each builder instance compiles out
271 + # the other QSA cache owner's slot-mapping rule.
272 + if circular_buffer_size > 0:
273 + valid = (
274 + mapped
275 + & (logical_position >= 0)
276 + & (token_idx + circular_buffer_size >= query_end)
277 + & (num_block_table_columns > 0)
278 + )
279 + physical_block = tl.load(
280 + block_table_ptr + request_idx * block_table_stride_0,
281 + mask=valid,
282 + other=-1,
283 + )
284 + valid &= physical_block >= 0
285 + slot = physical_block * circular_buffer_size + (
286 + logical_position % circular_buffer_size
287 + )
288 + elif compress_ratio != 1:
289 + compressed_position = tl.maximum(logical_position, 0) // compress_ratio
290 + logical_block = compressed_position // storage_block_size
291 + valid = (
292 + mapped
293 + & (logical_position >= 0)
294 + & ((logical_position + 1) % compress_ratio == 0)
295 + & (logical_block < num_block_table_columns)
296 + )
297 + physical_block = tl.load(
298 + block_table_ptr
299 + + request_idx * block_table_stride_0
300 + + logical_block * block_table_stride_1,
301 + mask=valid,
302 + other=-1,
303 + )
304 + valid &= physical_block >= 0
305 + valid &= (
306 + tl.load(common_slot_mapping_ptr + token_idx, mask=mapped, other=-1) >= 0
307 + )
308 + slot = physical_block * storage_block_size + (
309 + compressed_position % storage_block_size
310 + )
311 + if (circular_buffer_size > 0) or (compress_ratio != 1):
312 + tl.store(
313 + slot_mapping_ptr + token_idx,
314 + tl.where(valid, slot, -1),
315 + mask=store_mask,
316 + )
317 + work_tile_start = pid * WORK_BLOCK_SIZE
318 + has_work_tile = work_tile_start < max_num_work
319 + if k_work_metadata_ptr is not None and has_work_tile:
320 + # Every work CTA builds the request prefix in registers. Recomputing this
321 + # small vector lets CTAs write disjoint work tiles without a grid barrier.
322 + requests = tl.arange(0, REQUEST_SCAN_SIZE)
323 + valid_request = requests < num_reqs
324 + request_query_start = tl.load(
325 + query_start_loc_ptr + requests, mask=valid_request, other=0
326 + )
327 + request_query_end = tl.load(
328 + query_start_loc_ptr + requests + 1, mask=valid_request, other=0
329 + )
330 + request_seq_len = tl.load(seq_lens_ptr + requests, mask=valid_request, other=0)
331 + request_query_len = request_query_end - request_query_start
332 + chunk_start = request_seq_len - request_query_len
333 + num_groups = request_seq_len // compress_ratio - chunk_start // compress_ratio
334 + # Nonempty requests need one item even without a completed compression
335 + # group because work item zero also commits the current raw-K suffix.
336 + work_counts = tl.where(request_query_len > 0, tl.maximum(num_groups, 1), 0)
337 + work_ends = tl.cumsum(tl.where(valid_request, work_counts, 0), axis=0)
338 + total_work = tl.sum(tl.where(valid_request, work_counts, 0), axis=0)
339 + work_offsets = tl.arange(0, WORK_BLOCK_SIZE)
340 + work = work_tile_start + work_offsets
341 + in_bounds = work < max_num_work
342 + active = in_bounds & (work < total_work)
343 + request = tl.zeros((WORK_BLOCK_SIZE,), dtype=tl.int32)
344 + # Request zero starts at zero for every active work item. Descending
345 + # steps find the last request starting at or before this item.
346 + for step_idx in tl.range(0, work_search_steps):
347 + step = 1 << (work_search_steps - step_idx - 1)
348 + candidate = request + step
349 + valid_candidate = candidate < num_reqs
350 + candidate_start = tl.gather(work_ends, candidate - 1, 0)
351 + advance = active & valid_candidate & (candidate_start <= work)
352 + request = tl.where(advance, candidate, request)
353 +
354 + if launch_pdl:
355 + # Let the dependent grid start launch setup while these CTAs finish
356 + # stores; its gdc_wait still orders access to the completed metadata.
357 + tl.extra.cuda.gdc_launch_dependents()
358 +
359 + owner_work_start = tl.gather(work_ends, tl.maximum(request - 1, 0), 0)
360 + owner_work_start = tl.where(request == 0, 0, owner_work_start)
361 + work_in_request = work - owner_work_start
362 + tl.store(
363 + k_work_metadata_ptr + work * 2,
364 + tl.where(active, request, -1),
365 + mask=in_bounds,
366 + )
367 + tl.store(
368 + k_work_metadata_ptr + work * 2 + 1,
369 + tl.where(active, work_in_request, -1),
370 + mask=in_bounds,
371 + )
372 +
373 + else:
374 + if launch_pdl:
375 + # A dependent grid launches only after every CTA has signaled.
376 + tl.extra.cuda.gdc_launch_dependents()
377 +
378 +
379 +def build_qsa_metadata_triton(
380 + common_attn_metadata: CommonAttentionMetadata,
381 + token_to_req_buffer: torch.Tensor,
382 + logical_positions_buffer: torch.Tensor,
383 + slot_mapping_buffer: torch.Tensor,
384 + *,
385 + storage_block_size: int,
386 + compress_ratio: int,
387 + circular_buffer_size: int = 0,
388 + k_work_metadata_buffer: torch.Tensor | None = None,
389 + request_capacity: int | None = None,
390 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
391 + """Build QSA side-cache and optional pre-indexer work metadata."""
392 + num_tokens = common_attn_metadata.num_actual_tokens
393 + num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
394 + token_to_req = token_to_req_buffer[:num_tokens]
395 + logical_positions = logical_positions_buffer[:num_tokens]
396 + slot_mapping = slot_mapping_buffer[:num_tokens]
397 + num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1
398 + assert num_reqs > 0
399 +
400 + if k_work_metadata_buffer is not None:
401 + if request_capacity is None:
402 + request_capacity = num_reqs
403 + assert request_capacity >= num_reqs
404 + # Pad for tl.arange while keeping the scan width stable across live batches.
405 + request_scan_size = 1 << int(math.ceil(math.log2(request_capacity)))
406 + max_num_work = k_work_metadata_buffer.shape[0]
407 + else:
408 + request_scan_size = 1
409 + max_num_work = 0
410 +
411 + if num_tokens == 0 and k_work_metadata_buffer is None:
412 + return token_to_req, logical_positions, slot_mapping
413 +
414 + block_table = common_attn_metadata.block_table_tensor
415 + num_search_steps = int(math.ceil(math.log2(num_reqs)))
416 + work_search_steps = int(math.ceil(math.log2(num_reqs)))
417 + # The same grid covers token tiles and, for the compressed cache, work tiles.
418 + num_token_blocks = cdiv(num_tokens, 128)
419 + num_work_blocks = (
420 + cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0
421 + )
422 + _build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)](
423 + common_attn_metadata.query_start_loc,
424 + common_attn_metadata.seq_lens,
425 + common_attn_metadata.slot_mapping,
426 + block_table,
427 + token_to_req,
428 + logical_positions,
429 + slot_mapping,
430 + k_work_metadata_buffer,
431 + block_table.stride(0),
432 + block_table.stride(1),
433 + num_reqs,
434 + num_mapped_tokens,
435 + num_tokens,
436 + max_num_work,
437 + num_search_steps,
438 + work_search_steps,
439 + storage_block_size,
440 + compress_ratio,
441 + circular_buffer_size,
442 + block_table.shape[1],
443 + launch_pdl=_metadata_launch_pdl(),
444 + TOKEN_BLOCK_SIZE=128,
445 + REQUEST_SCAN_SIZE=request_scan_size,
446 + WORK_BLOCK_SIZE=256,
447 + num_warps=4,
448 + )
449 + if circular_buffer_size == 0 and compress_ratio == 1:
450 + slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
451 + return token_to_req, logical_positions, slot_mapping
452 +
453 +
454 +def _build_qsa_metadata_torch(
455 + common_attn_metadata: CommonAttentionMetadata,
456 + token_to_req_buffer: torch.Tensor,
457 + logical_positions_buffer: torch.Tensor,
458 + slot_mapping_buffer: torch.Tensor,
459 + *,
460 + storage_block_size: int,
461 + compress_ratio: int,
462 + circular_buffer_size: int = 0,
463 + k_work_metadata_buffer: torch.Tensor | None = None,
464 + request_capacity: int | None = None,
465 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
466 + del request_capacity
467 + num_tokens = common_attn_metadata.num_actual_tokens
468 + num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])
469 + logical_positions = logical_positions_buffer[:num_tokens]
470 +
471 + token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[
472 + :num_tokens
473 + ]
474 + logical_positions[:num_mapped_tokens].copy_(
475 + _logical_positions(
476 + common_attn_metadata.query_start_loc,
477 + common_attn_metadata.seq_lens,
478 + token_to_req[:num_mapped_tokens],
479 + num_mapped_tokens,
480 + )
481 + )
482 + if num_mapped_tokens < num_tokens:
483 + logical_positions[num_mapped_tokens:].fill_(-1)
484 + if circular_buffer_size > 0:
485 + slot_mapping = circular_qsa_slot_mapping(
486 + common_attn_metadata.block_table_tensor,
487 + token_to_req,
488 + logical_positions,
489 + circular_buffer_size,
490 + query_start_loc=common_attn_metadata.query_start_loc,
491 + out=slot_mapping_buffer,
492 + )
493 + elif compress_ratio == 1:
494 + slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]
495 + else:
496 + slot_mapping = compressed_qsa_slot_mapping(
497 + common_attn_metadata.block_table_tensor,
498 + token_to_req,
499 + logical_positions,
500 + storage_block_size,
501 + compress_ratio,
502 + slot_mapping_buffer,
503 + )
504 + slot_mapping.masked_fill_(
505 + common_attn_metadata.slot_mapping[:num_tokens] < 0, -1
506 + )
507 + if k_work_metadata_buffer is not None:
508 + query_lens = (
509 + common_attn_metadata.query_start_loc[1:]
510 + - common_attn_metadata.query_start_loc[:-1]
511 + )
512 + chunk_starts = common_attn_metadata.seq_lens - query_lens
513 + num_work_per_request = (
514 + common_attn_metadata.seq_lens // compress_ratio
515 + - chunk_starts // compress_ratio
516 + )
517 + num_work_per_request = torch.where(
518 + query_lens > 0, num_work_per_request.clamp_min(1), 0
519 + )
520 + k_start_loc = torch.empty(
521 + query_lens.shape[0] + 1,
522 + dtype=torch.int32,
523 + device=query_lens.device,
524 + )
525 + k_start_loc[0] = 0
526 + torch.cumsum(num_work_per_request, 0, out=k_start_loc[1:])
527 + work = torch.arange(
528 + k_work_metadata_buffer.shape[0],
529 + device=k_work_metadata_buffer.device,
530 + )
531 + requests = torch.searchsorted(k_start_loc[1:], work, right=True)
532 + active = work < k_start_loc[-1]
533 + work_in_request = (
534 + work - k_start_loc[requests.clamp_max(query_lens.shape[0] - 1)]
535 + )
536 + k_work_metadata_buffer[:, 0].copy_(
537 + torch.where(active, requests, -1).to(torch.int32)
538 + )
539 + k_work_metadata_buffer[:, 1].copy_(
540 + torch.where(active, work_in_request, -1).to(torch.int32)
541 + )
542 + return token_to_req, logical_positions, slot_mapping
543 +
544 +
545 +# Resolve the fallback outside the per-step metadata hot path.
546 +build_qsa_metadata = (
547 + build_qsa_metadata_triton if HAS_TRITON else _build_qsa_metadata_torch
548 +)
549 +
550 +
551 +@dataclass
552 +class QSAForwardMetadata(AttentionMetadata):
553 + """Common per-forward metadata for one QSA side cache."""
554 +
555 + block_table: torch.Tensor
556 + slot_mapping: torch.Tensor
557 + seq_lens: torch.Tensor
558 + query_start_loc: torch.Tensor
559 + token_to_req: torch.Tensor
560 + logical_positions: torch.Tensor
561 + k_work_metadata: torch.Tensor
562 + num_actual_tokens: int
563 + storage_block_size: int
564 + compress_ratio: int
565 +
566 +
567 +class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):
568 + """Build QSA metadata from vLLM's cache-group-specific common metadata."""
569 +
570 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
571 +
572 + def __init__(
573 + self,
574 + kv_cache_spec: AttentionSpec,
575 + layer_names: list[str],
576 + vllm_config: VllmConfig,
577 + device: torch.device,
578 + ) -> None:
579 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
580 + self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec)
581 + if isinstance(kv_cache_spec, MLAAttentionSpec):
582 + compress_ratio = kv_cache_spec.tokens_per_state
583 + assert isinstance(compress_ratio, int), (
584 + "QSA compression requires an integer tokens_per_state"
585 + )
586 + self.compress_ratio = compress_ratio
587 + else:
588 + self.compress_ratio = 1
589 + self.storage_block_size = kv_cache_spec.num_states
590 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
591 + self.token_to_req_buffer = torch.empty(
592 + max_tokens, dtype=torch.int32, device=device
593 + )
594 + self.slot_mapping_buffer = torch.empty(
595 + max_tokens, dtype=torch.int64, device=device
596 + )
597 + self.logical_positions_buffer = torch.empty(
598 + max_tokens, dtype=torch.int64, device=device
599 + )
600 + max_requests = vllm_config.scheduler_config.max_num_seqs
601 + self.request_capacity = max_requests
602 + if not self.is_circular_buffer and self.compress_ratio != 1:
603 + max_k_work = (
604 + max_tokens + (self.compress_ratio - 1) * max_requests
605 + ) // self.compress_ratio
606 + self.k_work_metadata_buffer = torch.empty(
607 + max_k_work, 2, dtype=torch.int32, device=device
608 + )
609 + else:
610 + self.k_work_metadata_buffer = torch.empty(
611 + 0, 2, dtype=torch.int32, device=device
612 + )
613 +
614 + def build(
615 + self,
616 + common_prefix_len: int,
617 + common_attn_metadata: CommonAttentionMetadata,
618 + fast_build: bool = False,
619 + ) -> QSAForwardMetadata:
620 + del common_prefix_len, fast_build
621 + num_tokens = common_attn_metadata.num_actual_tokens
622 + build_k_work = not self.is_circular_buffer and self.compress_ratio != 1
623 + k_work_metadata = self.k_work_metadata_buffer
624 + request_capacity = None
625 + if build_k_work:
626 + num_requests = common_attn_metadata.query_start_loc.shape[0] - 1
627 + request_capacity = self.request_capacity
628 + max_num_work = (
629 + num_tokens + (self.compress_ratio - 1) * num_requests
630 + ) // self.compress_ratio
631 + k_work_metadata = self.k_work_metadata_buffer[:max_num_work]
632 + token_to_req, logical_positions, slot_mapping = build_qsa_metadata(
633 + common_attn_metadata,
634 + self.token_to_req_buffer,
635 + self.logical_positions_buffer,
636 + self.slot_mapping_buffer,
637 + storage_block_size=self.storage_block_size,
638 + compress_ratio=self.compress_ratio,
639 + circular_buffer_size=(
640 + self.kv_cache_spec.block_size if self.is_circular_buffer else 0
641 + ),
642 + k_work_metadata_buffer=k_work_metadata if build_k_work else None,
643 + request_capacity=request_capacity,
644 + )
645 + return QSAForwardMetadata(
646 + block_table=common_attn_metadata.block_table_tensor,
647 + slot_mapping=slot_mapping,
648 + seq_lens=common_attn_metadata.seq_lens,
649 + query_start_loc=common_attn_metadata.query_start_loc,
650 + token_to_req=token_to_req,
651 + logical_positions=logical_positions,
652 + k_work_metadata=k_work_metadata,
653 + num_actual_tokens=num_tokens,
654 + storage_block_size=self.storage_block_size,
655 + compress_ratio=self.compress_ratio,
656 + )
657 +
658 +
659 +class QSAStateBackend(AttentionBackend):
660 + """Key-only dummy backend for out-of-band BF16 QSA side-cache operations."""
661 +
662 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
663 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
664 +
665 + @staticmethod
666 + def get_name() -> str:
667 + return "QWEN4_EXP_EXP_QSA_STATE"
668 +
669 + @staticmethod
670 + def get_impl_cls():
671 + raise NotImplementedError(
672 + "QSA state caches run out-of-band and have no attention impl"
673 + )
674 +
675 + @staticmethod
676 + def get_builder_cls() -> type[QSAMetadataBuilder]:
677 + return QSAMetadataBuilder
678 +
679 + @classmethod
680 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
681 + # QSA pages are packed beside the main KV pages within each block.
682 + return (KVCacheLayout.BLNHC, KVCacheLayout.BLHNC)
683 +
684 +
685 +class _QSAStateCache(nn.Module, AttentionLayerBase):
686 + supports_dcp = False
687 +
688 + def __init__(
689 + self,
690 + *,
691 + head_size: int,
692 + dtype: torch.dtype,
693 + cache_config: CacheConfig,
694 + prefix: str,
695 + vllm_config: VllmConfig,
696 + compress_ratio: int = 1,
697 + ) -> None:
698 + super().__init__()
699 + if head_size <= 0:
700 + raise ValueError("QSA cache head size must be positive")
701 + if compress_ratio <= 0:
702 + raise ValueError("QSA compression ratio must be positive")
703 + if cache_config.block_size % compress_ratio:
704 + raise ValueError(
705 + "QSA cache block size must be divisible by the compression ratio"
706 + )
707 + self.head_size = head_size
708 + self.dtype = dtype
709 + self.cache_config = cache_config
710 + self.prefix = prefix
711 + self.compress_ratio = compress_ratio
712 + self.kv_cache = torch.tensor([])
713 +
714 + static_context = vllm_config.compilation_config.static_forward_context
715 + if prefix in static_context:
716 + raise ValueError(f"Duplicate layer name: {prefix}")
717 + static_context[prefix] = self
718 +
719 + def forward(self) -> None: ...
720 +
721 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
722 + """Adapt the unified [B, H, N, C] view to QSA's [B, N, H, C]."""
723 + if kv_cache.ndim != 4 or kv_cache.shape[1] != 1:
724 + raise ValueError("QSA state cache must be [blocks, 1, states, width]")
725 + if kv_cache.dtype != torch.bfloat16 or kv_cache.shape[3] != self.head_size:
726 + raise ValueError("QSA state cache does not match its packed BF16 spec")
727 + super().bind_kv_cache(kv_cache.transpose(1, 2))
728 +
729 + def get_attn_backend(self) -> type[AttentionBackend]:
730 + return QSAStateBackend
731 +
732 +
733 +class QSAKeyStateCache(_QSAStateCache):
734 + """Raw BF16 key, optionally followed by exact int64 MRoPE positions."""
735 +
736 + _BF16_PER_INT64 = 4
737 + _NUM_ROPE_AXES = 3
738 +
739 + def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None:
740 + key_head_size = int(kwargs.pop("head_size"))
741 + self.key_head_size = key_head_size
742 + self.cache_rope_positions = bool(cache_rope_positions)
743 + self.rope_position_offset = (
744 + (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64
745 + ) * self._BF16_PER_INT64
746 + storage_head_size = key_head_size
747 + if self.cache_rope_positions:
748 + storage_head_size = self.rope_position_offset + (
749 + self._NUM_ROPE_AXES * self._BF16_PER_INT64
750 + )
751 + super().__init__(head_size=storage_head_size, **kwargs)
752 +
753 + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
754 + super().bind_kv_cache(kv_cache)
755 + qsa_cache = self.kv_cache
756 + self.key_cache = qsa_cache[..., : self.key_head_size]
757 + if self.cache_rope_positions:
758 + position_tail = qsa_cache[..., self.rope_position_offset :]
759 + self.rope_position_cache = position_tail.view(torch.int64)
760 + else:
761 + self.rope_position_cache = None
762 +
763 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
764 + # Hold the open group's committed keys plus every row a speculative
765 + # step stores before acceptance is known, rounded up to whole groups so
766 + # the ring divides the attention block size (it joins the LCM that sets
767 + # the scheduler block size). Anything narrower lets a rejected draft row
768 + # overwrite a committed key the next step needs to close the group.
769 + span = self.compress_ratio + vllm_config.num_speculative_tokens
770 + capacity = self.compress_ratio * cdiv(span, self.compress_ratio)
771 + assert self.cache_config.block_size % capacity == 0, (
772 + f"QSA ring capacity {capacity} must divide the attention block "
773 + f"size {self.cache_config.block_size}"
774 + )
775 + return CircularBufferSpec(
776 + block_size=capacity,
777 + num_kv_heads=1,
778 + head_size=self.head_size,
779 + head_size_v=0,
780 + dtype=self.dtype,
781 + )
782 +
783 +
784 +class QSACompressedKeyCache(_QSAStateCache):
785 + """Normalized, group-first-RoPE BF16 key at one row per complete group."""
786 +
787 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
788 + del vllm_config
789 + return MLAAttentionSpec(
790 + block_size=self.cache_config.block_size,
791 + num_kv_heads=1,
792 + head_size=self.head_size,
793 + dtype=self.dtype,
794 + tokens_per_state=self.compress_ratio,
795 + )
796 +
797 +
798 +__all__ = [
799 + "QSACompressedKeyCache",
800 + "QSAForwardMetadata",
801 + "QSAKeyStateCache",
802 + "QSAMetadataBuilder",
803 + "QSAStateBackend",
804 + "canonical_qsa_rope_positions",
805 + "circular_qsa_slot_mapping",
806 + "compressed_qsa_slot_mapping",
807 +]
vllm/models/qwen4_exp/config.py new
+252
@@ -0,0 +1,252 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen4Exp model configuration."""
4 +
5 +from typing import Any, ClassVar, cast
6 +
7 +from transformers import PretrainedConfig
8 +from transformers.models.qwen3_vl.configuration_qwen3_vl import (
9 + Qwen3VLVisionConfig,
10 +)
11 +
12 +from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
13 +
14 +_QSA_CONFIG_FIELDS = (
15 + "indexer_n_heads",
16 + "indexer_kv_heads",
17 + "indexer_head_dim",
18 + "indexer_budget",
19 + "indexer_compress_ratio",
20 +)
21 +
22 +
23 +class Qwen4ExpVisionConfig(Qwen3VLVisionConfig):
24 + model_type = "qwen4_exp"
25 + base_config_key = "vision_config"
26 +
27 +
28 +class Qwen4ExpTextConfig(Qwen3NextConfig):
29 + model_type = "qwen4_exp_text"
30 + base_config_key = "text_config"
31 + keys_to_ignore_at_inference = ["past_key_values"]
32 +
33 + def __init__(
34 + self,
35 + hc_count: int = 4,
36 + hc_lowrank: int = 320,
37 + ple_layer_ids: list[int] | None = None,
38 + ple_embed_dim: int | None = None,
39 + ple_conv_kernel_size: int = 4,
40 + ngram_size: int = 3,
41 + heads_per_ngram: int = 8,
42 + ngram_vocab_size_base: int = 20_000_000,
43 + make_ngram_vocab_size_divisible_by: int = 128,
44 + output_gate_type: str = "sigmoid",
45 + rope_parameters: dict[str, Any] | None = None,
46 + layer_types: list[str] | None = None,
47 + **kwargs: Any,
48 + ) -> None:
49 + if hc_count <= 1:
50 + raise ValueError(f"Qwen4Exp requires hc_count > 1, got {hc_count}.")
51 +
52 + if rope_parameters is not None:
53 + if kwargs.get("rope_scaling") is None:
54 + kwargs["rope_scaling"] = rope_parameters
55 + if kwargs.get("rope_theta") is None and "rope_theta" in rope_parameters:
56 + kwargs["rope_theta"] = rope_parameters["rope_theta"]
57 + if (
58 + kwargs.get("partial_rotary_factor") is None
59 + and "partial_rotary_factor" in rope_parameters
60 + ):
61 + kwargs["partial_rotary_factor"] = rope_parameters[
62 + "partial_rotary_factor"
63 + ]
64 +
65 + rope_scaling = kwargs.get("rope_scaling")
66 + rope_theta = kwargs.get("rope_theta", 10_000.0)
67 + super().__init__(layer_types=layer_types, **kwargs)
68 +
69 + normalized_rope_parameters = self.rope_parameters
70 + self.rope_scaling = (
71 + rope_scaling or rope_parameters or normalized_rope_parameters
72 + )
73 + self.rope_parameters = rope_parameters or normalized_rope_parameters
74 + self.rope_theta = rope_theta
75 +
76 + self.hc_count = hc_count
77 + self.hc_lowrank = hc_lowrank
78 + self.ple_layer_ids = ple_layer_ids or []
79 + self.ple_embed_dim = (
80 + self.hidden_size if ple_embed_dim is None else ple_embed_dim
81 + )
82 + self.ple_conv_kernel_size = ple_conv_kernel_size
83 + self.ngram_size = ngram_size
84 + self.heads_per_ngram = heads_per_ngram
85 + self.ngram_vocab_size_base = ngram_vocab_size_base
86 + self.make_ngram_vocab_size_divisible_by = make_ngram_vocab_size_divisible_by
87 + self.output_gate_type = output_gate_type
88 +
89 + self._validate_ple_config()
90 + self._validate_ple_layer_ids()
91 + self._validate_qsa_config()
92 +
93 + def _validate_ple_config(self) -> None:
94 + if self.hc_lowrank <= 0:
95 + raise ValueError(f"hc_lowrank must be positive, got {self.hc_lowrank}")
96 + if self.ngram_size < 2:
97 + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}")
98 + if self.heads_per_ngram <= 0:
99 + raise ValueError(
100 + f"heads_per_ngram must be positive, got {self.heads_per_ngram}"
101 + )
102 + if self.ple_embed_dim <= 0:
103 + raise ValueError(
104 + f"ple_embed_dim must be positive, got {self.ple_embed_dim}"
105 + )
106 + ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
107 + if self.ple_embed_dim % ngram_heads:
108 + raise ValueError(
109 + "ple_embed_dim must be divisible by total ngram heads: "
110 + f"{self.ple_embed_dim} % {ngram_heads} != 0"
111 + )
112 + if self.ple_conv_kernel_size <= 0:
113 + raise ValueError(
114 + "ple_conv_kernel_size must be positive, got "
115 + f"{self.ple_conv_kernel_size}"
116 + )
117 + if self.ngram_vocab_size_base <= 0:
118 + raise ValueError("ngram_vocab_size_base must be positive")
119 + if self.make_ngram_vocab_size_divisible_by <= 0:
120 + raise ValueError("make_ngram_vocab_size_divisible_by must be positive")
121 +
122 + def _validate_ple_layer_ids(self) -> None:
123 + invalid = [
124 + layer_id
125 + for layer_id in self.ple_layer_ids
126 + if not 1 <= int(layer_id) <= self.num_hidden_layers
127 + ]
128 + if invalid:
129 + raise ValueError(
130 + "ple_layer_ids are 1-based and must refer to an existing layer; "
131 + f"got {invalid} for {self.num_hidden_layers} layers"
132 + )
133 +
134 + def _validate_qsa_config(self) -> None:
135 + configured = {name: getattr(self, name, None) for name in _QSA_CONFIG_FIELDS}
136 + if all(value is None for value in configured.values()):
137 + return
138 +
139 + missing = [name for name, value in configured.items() if value is None]
140 + if missing:
141 + raise ValueError(f"QSA config is missing required fields: {missing}")
142 +
143 + values = {name: int(cast(int, value)) for name, value in configured.items()}
144 + if any(value <= 0 for value in values.values()):
145 + raise ValueError(f"QSA config values must be positive: {values}")
146 + if values["indexer_kv_heads"] != 1:
147 + raise ValueError("the QSA MQA operators require indexer_kv_heads=1")
148 + if values["indexer_budget"] % values["indexer_compress_ratio"] != 0:
149 + raise ValueError(
150 + "indexer_budget must be divisible by indexer_compress_ratio"
151 + )
152 + block_topk = values["indexer_budget"] // values["indexer_compress_ratio"]
153 + if block_topk not in (512, 2048):
154 + raise ValueError(
155 + "QSA requires indexer_budget / indexer_compress_ratio "
156 + f"to be 512 or 2048, got {block_topk}"
157 + )
158 + rotary_dim = int(self.head_dim * self.partial_rotary_factor)
159 + if rotary_dim > values["indexer_head_dim"]:
160 + raise ValueError(
161 + "QSA indexer_head_dim must cover the attention rotary "
162 + f"dimension, got {values['indexer_head_dim']} < {rotary_dim}"
163 + )
164 +
165 + @property
166 + def layers_block_type(self) -> list[str]:
167 + return [
168 + "attention" if layer_type == "full_attention" else layer_type
169 + for layer_type in self.layer_types
170 + ]
171 +
172 + @property
173 + def short_conv_layer_ids(self) -> list[int]:
174 + if not self.ple_layer_ids:
175 + return []
176 + return sorted({int(layer_id) - 1 for layer_id in self.ple_layer_ids})
177 +
178 + @property
179 + def short_conv_state_shape(self) -> tuple[int, int] | None:
180 + if not self.short_conv_layer_ids:
181 + return None
182 + ple_state_len = (self.ple_conv_kernel_size - 1) * self.ngram_size
183 + ple_channels = self.hidden_size * self.hc_count
184 + return ple_channels, ple_state_len
185 +
186 + @property
187 + def ngram_context_len(self) -> int:
188 + if not self.ple_layer_ids:
189 + return 0
190 + return max(int(self.ngram_size) - 1, 0)
191 +
192 +
193 +class Qwen4ExpConfig(PretrainedConfig):
194 + model_type = "qwen4_exp"
195 + sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = {
196 + "vision_config": Qwen4ExpVisionConfig,
197 + "text_config": Qwen4ExpTextConfig,
198 + }
199 + keys_to_ignore_at_inference = ["past_key_values"]
200 +
201 + def __init__(
202 + self,
203 + text_config: Qwen4ExpTextConfig | dict[str, Any] | None = None,
204 + vision_config: Qwen4ExpVisionConfig | dict[str, Any] | None = None,
205 + image_token_id: int = 248056,
206 + video_token_id: int = 248057,
207 + vision_start_token_id: int = 248053,
208 + vision_end_token_id: int = 248054,
209 + tie_word_embeddings: bool = False,
210 + rope_parameters: dict[str, Any] | None = None,
211 + **kwargs: Any,
212 + ) -> None:
213 + if text_config is not None:
214 + kwargs.pop("split_ngram_parts", None)
215 +
216 + text_kwargs = (
217 + dict(kwargs)
218 + if text_config is None
219 + and "hidden_size" in kwargs
220 + and "num_hidden_layers" in kwargs
221 + else {}
222 + )
223 +
224 + if isinstance(vision_config, dict):
225 + self.vision_config = self.sub_configs["vision_config"](**vision_config)
226 + elif vision_config is None:
227 + self.vision_config = self.sub_configs["vision_config"]()
228 + else:
229 + self.vision_config = vision_config
230 +
231 + if isinstance(text_config, dict):
232 + self.text_config = self.sub_configs["text_config"](**text_config)
233 + elif text_config is None:
234 + self.text_config = self.sub_configs["text_config"](**text_kwargs)
235 + else:
236 + self.text_config = text_config
237 +
238 + self.image_token_id = image_token_id
239 + self.video_token_id = video_token_id
240 + self.vision_start_token_id = vision_start_token_id
241 + self.vision_end_token_id = vision_end_token_id
242 + self.rope_parameters = rope_parameters or getattr(
243 + self.text_config, "rope_parameters", {}
244 + )
245 + super().__init__(**kwargs, tie_word_embeddings=tie_word_embeddings)
246 +
247 +
248 +__all__ = [
249 + "Qwen4ExpConfig",
250 + "Qwen4ExpTextConfig",
251 + "Qwen4ExpVisionConfig",
252 +]
vllm/models/qwen4_exp/nvidia/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/qwen4_exp/nvidia/hyperconnection.py new
+205
@@ -0,0 +1,205 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""HyperConnection (Gated Residual) utilities — NVIDIA model variant.
4 +
5 +Implements the HyperConnection residual scheme proposed in
6 +"HyperConnections" (https://arxiv.org/abs/2409.19606). This NVIDIA variant
7 +delays each HC combine to the following HC mix boundary. HC glue kernels,
8 +including fused combine+RMSNorm, live in ``ops/hc.py``; projections remain
9 +standard vLLM Linear modules.
10 +
11 +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner
12 +(HC outer, HS inner — checkpoint-native layout).
13 +
14 +Typical usage inside a transformer decoder layer::
15 +
16 + self.attn_hc = GatedResidual(hc_config)
17 +
18 + hidden_states, block_input, injection = self.attn_hc.mix(hidden_states)
19 + attention_output = attention(block_input)
20 + hidden_states, block_input, injection = self.mlp_hc.combine_and_mix(
21 + hidden_states, attention_output, injection
22 + )
23 +"""
24 +
25 +import torch
26 +from torch import nn
27 +
28 +from vllm.model_executor.layers.linear import (
29 + MergedColumnParallelLinear,
30 + ReplicatedLinear,
31 +)
32 +from vllm.model_executor.models.utils import maybe_prefix
33 +
34 +from ..common.hyperconnection import (
35 + GroupedGemmaRMSNorm,
36 + HyperConnectionConfig,
37 +)
38 +from .ops.hc import (
39 + grouped_gemma_rmsnorm,
40 + hc_combine,
41 + hc_combine_norm,
42 + hc_gate_mix,
43 + hc_silu,
44 +)
45 +
46 +
47 +# ---------------------------------------------------------------------------
48 +# Gated-residual variant
49 +# ---------------------------------------------------------------------------
50 +class GatedResidual(nn.Module):
51 + """Gated HyperConnection with learnable low-rank mixing and injection.
52 +
53 + ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged
54 + low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean
55 + over the HC streams). When passed a pending block output and an injection,
56 + it fuses their residual combine with the RMSNorm. Final mixers use
57 + ``use_combine=False`` and do not produce a new injection.
58 +
59 + Weights: the norm owns the grouped GemmaRMSNorm affine; the projections
60 + are vLLM Linear modules (merged replicated linear for down+inject), so
61 + GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the
62 + standard quant_method mechanism.
63 + """
64 +
65 + def __init__(
66 + self,
67 + config: HyperConnectionConfig,
68 + use_combine: bool = True,
69 + prefix: str = "",
70 + ) -> None:
71 + super().__init__()
72 + self.config = config
73 + self.lora_rank = config.hc_lowrank
74 + self.hc_count = config.hc_count
75 + self.hidden_size = config.hidden_size
76 + self.use_combine = use_combine
77 +
78 + norm_size = (
79 + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size
80 + )
81 + group_size = config.hidden_size if config.hc_per_branch_norm else None
82 + # Normalize each H-sized HC stream independently while retaining a
83 + # separate affine weight for every element of the HC*H layout.
84 + self.hc_norm = GroupedGemmaRMSNorm(
85 + norm_size,
86 + eps=config.rms_norm_eps,
87 + group_size=group_size,
88 + dtype=config.params_dtype,
89 + )
90 +
91 + # -- vLLM Linear weights --------------------------------------------
92 + # The merged skinny-GEMM shape is physically padded to 16 rows to ensure
93 + # good alignment and performant implementation chosen by CuBLAS heuristics.
94 + self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0
95 + if use_combine:
96 + self.input_mix_weight_down_block_inject = MergedColumnParallelLinear(
97 + self.hyper_hidden_size,
98 + [self.lora_rank, self.hc_count]
99 + + ([self.pad_size] if self.pad_size else []),
100 + bias=False,
101 + params_dtype=config.params_dtype,
102 + quant_config=None,
103 + prefix=maybe_prefix(prefix, "input_mix_weight_down_block_inject"),
104 + return_bias=False,
105 + disable_tp=True,
106 + )
107 + else:
108 + self.input_mix_weight_down = ReplicatedLinear(
109 + self.hyper_hidden_size,
110 + self.lora_rank,
111 + bias=False,
112 + params_dtype=config.params_dtype,
113 + quant_config=None,
114 + prefix=maybe_prefix(prefix, "input_mix_weight_down"),
115 + return_bias=False,
116 + )
117 + self.input_mix_weight_up = ReplicatedLinear(
118 + self.lora_rank,
119 + self.hyper_hidden_size,
120 + bias=False,
121 + params_dtype=config.params_dtype,
122 + quant_config=None,
123 + prefix=maybe_prefix(prefix, "input_mix_weight_up"),
124 + return_bias=False,
125 + )
126 +
127 + def mix(
128 + self, hidden_states: torch.Tensor
129 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
130 + xn = grouped_gemma_rmsnorm(
131 + hidden_states,
132 + self.hc_norm.weight,
133 + self.config.rms_norm_eps,
134 + self.hc_count,
135 + )
136 +
137 + if self.use_combine:
138 + # produce injection logits for combine
139 + split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
140 + down_and_injection = self.input_mix_weight_down_block_inject(xn)
141 + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
142 + else:
143 + lora = self.input_mix_weight_down(xn)
144 + injection = None
145 +
146 + lora = hc_silu(lora, self.hc_count)
147 + gate = self.input_mix_weight_up(lora) # [M, D]
148 + block_input = hc_gate_mix(xn, gate, self.hc_count)
149 +
150 + return hidden_states, block_input, injection
151 +
152 + def combine_and_mix(
153 + self,
154 + hidden_states: torch.Tensor,
155 + prev_block_output: torch.Tensor,
156 + prev_injection: torch.Tensor,
157 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
158 + """Consume a pending combine, then prepare the next block input.
159 +
160 + ``hidden_states`` is the multi-stream state from before the pending
161 + block's mix. Its combine with ``block_output`` is fused with this
162 + module's input RMSNorm.
163 + """
164 + hidden_states, xn = hc_combine_norm(
165 + hidden_states,
166 + prev_block_output,
167 + prev_injection,
168 + self.hc_norm.weight,
169 + self.config.rms_norm_eps,
170 + self.hc_count,
171 + )
172 +
173 + if self.use_combine:
174 + # produce injection logits for combine
175 + split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
176 + down_and_injection = self.input_mix_weight_down_block_inject(xn)
177 + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
178 + else:
179 + lora = self.input_mix_weight_down(xn)
180 + injection = None
181 +
182 + lora = hc_silu(lora, self.hc_count)
183 + gate = self.input_mix_weight_up(lora) # [M, D]
184 + block_input = hc_gate_mix(xn, gate, self.hc_count)
185 +
186 + return hidden_states, block_input, injection
187 +
188 + def combine(
189 + self,
190 + hidden_states: torch.Tensor,
191 + block_output: torch.Tensor,
192 + injection: torch.Tensor,
193 + ) -> torch.Tensor:
194 + return hc_combine(hidden_states, block_output, injection, self.hc_count)
195 +
196 + @property
197 + def hyper_hidden_size(self) -> int:
198 + return self.hc_count * self.hidden_size
199 +
200 +
201 +__all__ = [
202 + "GatedResidual",
203 + "GroupedGemmaRMSNorm",
204 + "HyperConnectionConfig",
205 +]
vllm/models/qwen4_exp/nvidia/indexer_qsa.py new
+356
@@ -0,0 +1,356 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen4Exp weight-free QSA indexer."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import cast
8 +
9 +import torch
10 +from torch import nn
11 +
12 +from vllm.config import VllmConfig
13 +from vllm.forward_context import get_forward_context
14 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
15 +from vllm.model_executor.layers.linear import ReplicatedLinear
16 +from vllm.model_executor.layers.quantization import QuantizationConfig
17 +from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope
18 +from vllm.transformers_utils.configs.qwen4_exp import (
19 + Qwen4ExpTextConfig,
20 +)
21 +
22 +from ..common.qsa_cache import (
23 + QSACompressedKeyCache,
24 + QSAForwardMetadata,
25 + QSAKeyStateCache,
26 + canonical_qsa_rope_positions,
27 +)
28 +from .ops.qsa_pre_indexer import qsa_pre_indexer
29 +
30 +
31 +def apply_qsa_rope(
32 + rotary_emb: nn.Module,
33 + positions: torch.Tensor,
34 + tensor: torch.Tensor,
35 +) -> torch.Tensor:
36 + """Apply the main attention's exact 1D/MRoPE composition to QSA heads."""
37 +
38 + num_tokens, _, head_dim = tensor.shape
39 + rotary_dim = rotary_emb.rotary_dim
40 + cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001
41 + cos_sin = cache[positions]
42 + cos, sin = cos_sin.chunk(2, dim=-1)
43 + if positions.ndim == 2:
44 + shape = tensor.shape
45 + tensor, _ = triton_mrope(
46 + tensor.reshape(num_tokens, -1),
47 + tensor.new_empty((num_tokens, head_dim)),
48 + cos,
49 + sin,
50 + rotary_emb.mrope_section,
51 + head_dim,
52 + rotary_dim,
53 + rotary_emb.mrope_interleaved,
54 + rotary_emb.is_neox_style,
55 + )
56 + return tensor.reshape(shape)
57 +
58 + rotated = rotary_emb.apply_rotary_emb.forward_cuda(
59 + tensor[..., :rotary_dim],
60 + cos,
61 + sin,
62 + )
63 + return torch.cat((rotated, tensor[..., rotary_dim:]), dim=-1)
64 +
65 +
66 +def _supports_fused_pre_indexer(
67 + rotary_emb: nn.Module,
68 + head_dim: int,
69 + num_kv_heads: int,
70 + compress_ratio: int,
71 +) -> bool:
72 + rotary_dim = int(rotary_emb.rotary_dim)
73 + mrope_section = getattr(rotary_emb, "mrope_section", None)
74 + return (
75 + bool(getattr(rotary_emb, "is_neox_style", False))
76 + and (
77 + not mrope_section
78 + or (
79 + len(mrope_section) == 3
80 + and sum(mrope_section) == rotary_dim // 2
81 + and bool(getattr(rotary_emb, "mrope_interleaved", False))
82 + )
83 + )
84 + and head_dim == 128
85 + and rotary_dim == 64
86 + and num_kv_heads == 1
87 + and compress_ratio > 1
88 + and compress_ratio & (compress_ratio - 1) == 0
89 + )
90 +
91 +
92 +class QSAIndexer(nn.Module):
93 + """Replicated Q/K projection plus paged, weight-free QSA selection.
94 +
95 + ``prefix`` must be the checkpoint's indexer prefix, normally
96 + ``model.layers.N.self_attn.indexer``. Consequently the trainable names are
97 + ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix.
98 + """
99 +
100 + def __init__(
101 + self,
102 + *,
103 + vllm_config: VllmConfig,
104 + config: Qwen4ExpTextConfig,
105 + layer_id: int,
106 + rotary_emb: nn.Module,
107 + quant_config: QuantizationConfig | None = None,
108 + prefix: str = "",
109 + ) -> None:
110 + super().__init__()
111 + if vllm_config.cache_config is None:
112 + raise ValueError("QSA requires a paged KV cache")
113 + if vllm_config.model_config.dtype != torch.bfloat16:
114 + raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
115 +
116 + self.layer_id = int(layer_id)
117 + self.index_n_heads = int(config.indexer_n_heads)
118 + self.index_kv_heads = int(config.indexer_kv_heads)
119 + self.index_head_dim = int(config.indexer_head_dim)
120 + self.token_topk = int(config.indexer_budget)
121 + self.compress_ratio = int(config.indexer_compress_ratio)
122 + self.rotary_emb = rotary_emb
123 + self.use_fused_pre_indexer = _supports_fused_pre_indexer(
124 + rotary_emb,
125 + self.index_head_dim,
126 + self.index_kv_heads,
127 + self.compress_ratio,
128 + )
129 + self.prefix = prefix
130 + # MTP step 0 selects the target-aligned rows; later steps reuse them
131 + # while continuing to update the QSA side cache.
132 + self.skip_topk = False
133 +
134 + self.index_qk_proj = ReplicatedLinear(
135 + int(config.hidden_size),
136 + (self.index_n_heads + self.index_kv_heads) * self.index_head_dim,
137 + bias=False,
138 + quant_config=quant_config,
139 + prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj",
140 + )
141 + self.q_layernorm = GemmaRMSNorm(
142 + self.index_head_dim,
143 + eps=float(getattr(config, "rms_norm_eps", 1e-6)),
144 + )
145 + self.k_layernorm = GemmaRMSNorm(
146 + self.index_head_dim,
147 + eps=float(getattr(config, "rms_norm_eps", 1e-6)),
148 + )
149 +
150 + cache_config = vllm_config.cache_config
151 + cache_prefix = f"{prefix}." if prefix else ""
152 + self.raw_key_cache = QSAKeyStateCache(
153 + head_size=self.index_head_dim,
154 + dtype=torch.bfloat16,
155 + cache_rope_positions=vllm_config.model_config.uses_mrope,
156 + prefix=f"{cache_prefix}raw_key_cache",
157 + cache_config=cache_config,
158 + compress_ratio=self.compress_ratio,
159 + vllm_config=vllm_config,
160 + )
161 + self.compressed_key_cache = QSACompressedKeyCache(
162 + head_size=self.index_head_dim,
163 + dtype=torch.bfloat16,
164 + compress_ratio=self.compress_ratio,
165 + prefix=f"{cache_prefix}compressed_key_cache",
166 + cache_config=cache_config,
167 + vllm_config=vllm_config,
168 + )
169 +
170 + @property
171 + def output_width(self) -> int:
172 + return self.token_topk + self.compress_ratio - 1
173 +
174 + def _metadata(
175 + self,
176 + ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None:
177 + metadata = get_forward_context().attn_metadata
178 + if isinstance(metadata, list):
179 + metadata = metadata[0]
180 + if not isinstance(metadata, dict):
181 + return None
182 + raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix])
183 + compressed = cast(
184 + QSAForwardMetadata, metadata[self.compressed_key_cache.prefix]
185 + )
186 + if raw.num_actual_tokens != compressed.num_actual_tokens:
187 + raise RuntimeError("QSA side-cache metadata token counts disagree")
188 + if not raw.logical_positions.is_cuda and (
189 + not torch.equal(raw.logical_positions, compressed.logical_positions)
190 + ):
191 + raise RuntimeError("QSA side-cache metadata positions disagree")
192 + return raw, compressed
193 +
194 + def forward(
195 + self,
196 + hidden_states: torch.Tensor,
197 + positions: torch.Tensor,
198 + out: torch.Tensor | None = None,
199 + ) -> torch.Tensor:
200 + """Return fixed-width request-relative token indices padded with ``-1``."""
201 +
202 + metadata = self._metadata()
203 + if metadata is None:
204 + # Preserve step-0 indices when later MTP steps reuse the buffer.
205 + if self.skip_topk and out is not None:
206 + return out
207 + result = torch.full(
208 + (hidden_states.shape[0], self.output_width),
209 + -1,
210 + dtype=torch.int32,
211 + device=hidden_states.device,
212 + )
213 + if out is not None:
214 + out.copy_(result)
215 + return out
216 + return result
217 +
218 + from .ops.qsa import (
219 + qsa_compress_groups_with_ratio,
220 + qsa_select_paged_tokens,
221 + qsa_store_cache_rows,
222 + )
223 +
224 + raw_metadata, compressed_metadata = metadata
225 + num_tokens = raw_metadata.num_actual_tokens
226 + hidden_states = hidden_states[:num_tokens]
227 + positions = positions[..., :num_tokens]
228 +
229 + # Q/K projection
230 + projected_qk, _ = self.index_qk_proj(hidden_states)
231 + projected_q, raw_keys = projected_qk.split(
232 + (
233 + self.index_n_heads * self.index_head_dim,
234 + self.index_kv_heads * self.index_head_dim,
235 + ),
236 + dim=-1,
237 + )
238 + raw_key_state_cache = self.raw_key_cache
239 + compressed_key_cache = self.compressed_key_cache.kv_cache
240 +
241 + if self.use_fused_pre_indexer:
242 + q = projected_q.new_empty(
243 + num_tokens,
244 + self.index_n_heads,
245 + self.index_head_dim,
246 + )
247 + qsa_pre_indexer(
248 + projected_q,
249 + raw_keys,
250 + positions,
251 + self.rotary_emb.cos_sin_cache,
252 + self.q_layernorm.weight,
253 + self.k_layernorm.weight,
254 + self.q_layernorm.variance_epsilon,
255 + q,
256 + raw_key_state_cache.kv_cache,
257 + raw_metadata.slot_mapping,
258 + raw_metadata.block_table,
259 + raw_metadata.query_start_loc,
260 + raw_metadata.logical_positions,
261 + compressed_key_cache,
262 + compressed_metadata.slot_mapping,
263 + compressed_metadata.k_work_metadata,
264 + compress_ratio=self.compress_ratio,
265 + mrope_section=getattr(self.rotary_emb, "mrope_section", None),
266 + rope_pos_offset=(
267 + raw_key_state_cache.rope_position_offset
268 + if raw_key_state_cache.rope_position_cache is not None
269 + else None
270 + ),
271 + )
272 + else:
273 + # Unfused reference path
274 + from flashinfer.norm import gemma_rmsnorm
275 +
276 + q = projected_q.reshape(-1, self.index_n_heads, self.index_head_dim)
277 + q = gemma_rmsnorm(
278 + q.reshape(-1, self.index_head_dim),
279 + self.q_layernorm.weight,
280 + self.q_layernorm.variance_epsilon,
281 + ).reshape_as(q)
282 + q = apply_qsa_rope(self.rotary_emb, positions, q)
283 +
284 + raw_key_cache = raw_key_state_cache.key_cache
285 + rope_position_cache = raw_key_state_cache.rope_position_cache
286 + if rope_position_cache is None:
287 + position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand(
288 + -1, 1, 3
289 + )
290 + else:
291 + position_rows = canonical_qsa_rope_positions(positions).to(
292 + device=raw_key_cache.device
293 + )
294 + pooled, first_positions = qsa_compress_groups_with_ratio(
295 + raw_keys.reshape(-1, 1, self.index_head_dim),
296 + position_rows,
297 + raw_key_cache,
298 + raw_metadata.block_table,
299 + raw_metadata.token_to_req,
300 + raw_metadata.query_start_loc,
301 + raw_metadata.logical_positions,
302 + compressed_metadata.slot_mapping,
303 + self.compress_ratio,
304 + rope_position_cache,
305 + )
306 + compressed_keys = gemma_rmsnorm(
307 + pooled.reshape(-1, self.index_head_dim),
308 + self.k_layernorm.weight,
309 + self.k_layernorm.variance_epsilon,
310 + ).reshape(-1, 1, self.index_head_dim)
311 + if getattr(self.rotary_emb, "mrope_section", None):
312 + first_positions = first_positions.transpose(0, 1)
313 + else:
314 + first_positions = first_positions[:, 0]
315 + compressed_keys = apply_qsa_rope(
316 + self.rotary_emb,
317 + first_positions,
318 + compressed_keys,
319 + )
320 + qsa_store_cache_rows(
321 + compressed_key_cache,
322 + compressed_metadata.slot_mapping,
323 + compressed_keys,
324 + )
325 + qsa_store_cache_rows(
326 + raw_key_cache,
327 + raw_metadata.slot_mapping,
328 + raw_keys,
329 + )
330 + if rope_position_cache is not None:
331 + qsa_store_cache_rows(
332 + rope_position_cache,
333 + raw_metadata.slot_mapping,
334 + position_rows,
335 + )
336 +
337 + if self.skip_topk:
338 + if out is None:
339 + raise RuntimeError("QSA top-k reuse requires an output buffer")
340 + return out
341 +
342 + # Score compressed keys, select blocks, then expand them to token indices.
343 + return qsa_select_paged_tokens(
344 + q,
345 + compressed_key_cache,
346 + compressed_metadata.block_table,
347 + compressed_metadata.token_to_req,
348 + compressed_metadata.logical_positions,
349 + compressed_metadata.seq_lens,
350 + self.token_topk,
351 + self.compress_ratio,
352 + out,
353 + )
354 +
355 +
356 +__all__ = ["QSAIndexer", "apply_qsa_rope"]
vllm/models/qwen4_exp/nvidia/low_latency_gemm.py new
+185
@@ -0,0 +1,185 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Qwen4Exp decode GEMM selection on Blackwell.
4 +
5 +Dispatch follows Kimi-K3 and uses the local ``(N, K)`` shape and token count.
6 +Plans contain measured CUDA graph capture sizes; other token counts use the
7 +standard linear implementation.
8 +"""
9 +
10 +import torch
11 +from torch import nn
12 +
13 +import vllm.envs as envs
14 +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
15 + SkinnyGemmConfig,
16 + shape_dynamic_skinny_gemm,
17 +)
18 +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
19 +from vllm.model_executor.layers.vocab_parallel_embedding import (
20 + ParallelLMHead,
21 + UnquantizedEmbeddingMethod,
22 +)
23 +from vllm.platforms import current_platform
24 +from vllm.utils.torch_utils import direct_register_custom_op
25 +
26 +QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
27 + # GDN fused QKVZ projection, TP=4.
28 + (4096, 2560): {
29 + 1: SkinnyGemmConfig(1, 64, 4, k_unroll=4),
30 + 2: SkinnyGemmConfig(2, 64, 4, k_unroll=4),
31 + },
32 + # GDN and QSA output projections, TP=4.
33 + (2560, 1536): {
34 + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4),
35 + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=2, vector_width=4),
36 + 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2),
37 + },
38 + # GDN fused B/A projection, TP=4.
39 + (24, 2560): {
40 + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4),
41 + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4),
42 + 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4),
43 + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
44 + 16: SkinnyGemmConfig(16, 128, 1, k_unroll=4, vector_width=4),
45 + },
46 + # QSA fused QKV/gate projection, TP=4.
47 + (3584, 2560): {
48 + 1: SkinnyGemmConfig(1, 128, 4, k_unroll=4, vector_width=4),
49 + 2: SkinnyGemmConfig(2, 64, 2, k_unroll=2),
50 + 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2),
51 + },
52 + # QSA indexer Q/K projection, replicated in a TP=4 deployment.
53 + (640, 2560): {
54 + 1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4),
55 + 2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4),
56 + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4),
57 + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
58 + },
59 + # Shared-expert fused gate/up projection, TP=4.
60 + (320, 2560): {
61 + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4),
62 + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4),
63 + 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4),
64 + 8: SkinnyGemmConfig(8, 64, 1, k_unroll=4),
65 + 16: SkinnyGemmConfig(16, 128, 2, k_unroll=4, vector_width=4),
66 + },
67 + # LM head, TP=4.
68 + (62080, 2560): {
69 + 1: SkinnyGemmConfig(1, 64, 4, k_unroll=2),
70 + 2: SkinnyGemmConfig(2, 32, 4, k_unroll=2),
71 + },
72 + # HC merged down/injection projection, replicated in a TP=4 deployment.
73 + (336, 10240): {
74 + 1: SkinnyGemmConfig(1, 128, 1, static_k=10240),
75 + 2: SkinnyGemmConfig(2, 128, 1, static_k=10240),
76 + 4: SkinnyGemmConfig(4, 128, 2, static_k=10240),
77 + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4),
78 + },
79 +}
80 +
81 +
82 +def _is_sm103() -> bool:
83 + return current_platform.is_device_capability((10, 3))
84 +
85 +
86 +def _is_packed_row_major(tensor: torch.Tensor) -> bool:
87 + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
88 +
89 +
90 +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
91 + return (
92 + not envs.VLLM_BATCH_INVARIANT
93 + and _is_packed_row_major(x)
94 + and _is_packed_row_major(weight)
95 + and x.dtype == torch.bfloat16
96 + and weight.dtype == torch.bfloat16
97 + and x.is_cuda
98 + and weight.is_cuda
99 + and x.device == weight.device
100 + and x.shape[1] == weight.shape[1]
101 + )
102 +
103 +
104 +class _Qwen4ExpLowLatencyApply:
105 + def apply(
106 + self,
107 + layer: nn.Module,
108 + x: torch.Tensor,
109 + bias: torch.Tensor | None = None,
110 + ) -> torch.Tensor:
111 + if bias is None and not envs.VLLM_BATCH_INVARIANT:
112 + return torch.ops.vllm.qwen4_exp_low_latency_gemm(x, layer.weight)
113 + return super().apply(layer, x, bias) # type: ignore[misc]
114 +
115 +
116 +class Qwen4ExpLowLatencyLinearMethod(_Qwen4ExpLowLatencyApply, UnquantizedLinearMethod):
117 + pass
118 +
119 +
120 +class Qwen4ExpLowLatencyEmbeddingMethod(
121 + _Qwen4ExpLowLatencyApply, UnquantizedEmbeddingMethod
122 +):
123 + pass
124 +
125 +
126 +def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
127 + plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1]))
128 + config = None if plan is None else plan.get(x.shape[0])
129 + if (
130 + config is not None
131 + and _runtime_ok(x, weight)
132 + and shape_dynamic_skinny_gemm.is_available()
133 + ):
134 + return shape_dynamic_skinny_gemm(x, weight, config)
135 + return torch.nn.functional.linear(x, weight)
136 +
137 +
138 +def _qwen4_exp_low_latency_gemm_fake(
139 + x: torch.Tensor, weight: torch.Tensor
140 +) -> torch.Tensor:
141 + return x.new_empty((*x.shape[:-1], weight.shape[0]))
142 +
143 +
144 +direct_register_custom_op(
145 + op_name="qwen4_exp_low_latency_gemm",
146 + op_func=_qwen4_exp_low_latency_gemm,
147 + fake_impl=_qwen4_exp_low_latency_gemm_fake,
148 +)
149 +
150 +
151 +def enable_qwen4_exp_low_latency_gemm(
152 + module: nn.Module,
153 + dtype: torch.dtype,
154 +) -> None:
155 + if dtype != torch.bfloat16 or not _is_sm103():
156 + return
157 + if not shape_dynamic_skinny_gemm.is_available():
158 + return
159 +
160 + warmup_configs: set[SkinnyGemmConfig] = set()
161 + for child in module.modules():
162 + is_linear = (
163 + isinstance(child, LinearBase)
164 + and type(child.quant_method) is UnquantizedLinearMethod
165 + )
166 + is_head = (
167 + isinstance(child, ParallelLMHead)
168 + and type(child.quant_method) is UnquantizedEmbeddingMethod
169 + )
170 + if not (is_linear or is_head):
171 + continue
172 + weight = getattr(child, "weight", None)
173 + if weight is None or weight.dim() != 2:
174 + continue
175 + plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1]))
176 + if plan is None:
177 + continue
178 + if is_linear:
179 + child.quant_method = Qwen4ExpLowLatencyLinearMethod()
180 + else:
181 + child.quant_method = Qwen4ExpLowLatencyEmbeddingMethod()
182 + warmup_configs.update(plan.values())
183 +
184 + if warmup_configs:
185 + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs)
vllm/models/qwen4_exp/nvidia/model.py new
+1065
@@ -0,0 +1,1065 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Qwen4Exp model."""
4 +
5 +from collections.abc import Iterable
6 +from itertools import islice
7 +
8 +import torch
9 +from torch import nn
10 +
11 +from vllm.compilation.decorators import support_torch_compile
12 +from vllm.config import VllmConfig
13 +from vllm.distributed import get_pp_group
14 +from vllm.model_executor.layers.fused_moe.utils import (
15 + is_model_fused_shared_expert_compatible,
16 +)
17 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
18 +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
19 + QwenGatedDeltaNetAttention,
20 +)
21 +from vllm.model_executor.layers.mamba.mamba_utils import (
22 + MambaStateCopyFunc,
23 + MambaStateCopyFuncCalculator,
24 + MambaStateCopyFuncsByType,
25 + MambaStateDtypeCalculator,
26 + MambaStateShapeCalculator,
27 +)
28 +from vllm.model_executor.layers.quantization import QuantizationConfig
29 +from vllm.model_executor.layers.vocab_parallel_embedding import (
30 + ParallelLMHead,
31 + VocabParallelEmbedding,
32 +)
33 +from vllm.model_executor.models.interfaces import (
34 + HasInnerState,
35 + IsHybrid,
36 + MixtureOfExperts,
37 + MultiModalEmbeddings,
38 + SupportsLoRA,
39 + SupportsMRoPE,
40 + SupportsPP,
41 + _require_is_multimodal,
42 +)
43 +from vllm.model_executor.models.qwen3_5 import (
44 + Qwen3_5ForConditionalGeneration,
45 + Qwen3_5Model,
46 +)
47 +from vllm.model_executor.models.qwen3_next import (
48 + Qwen3NextAttention,
49 + Qwen3NextMLP,
50 + Qwen3NextSparseMoeBlock,
51 +)
52 +from vllm.model_executor.models.qwen3_vl import (
53 + Qwen3_VisionTransformer,
54 + Qwen3VLDummyInputsBuilder,
55 + Qwen3VLMultiModalProcessor,
56 + Qwen3VLProcessingInfo,
57 +)
58 +from vllm.model_executor.models.utils import (
59 + AutoWeightsLoader,
60 + StageMissingLayer,
61 + WeightsMapper,
62 + _merge_multimodal_embeddings,
63 + extract_layer_index,
64 + make_empty_intermediate_tensors_factory,
65 + make_layers,
66 + maybe_fuse_shared_experts,
67 + maybe_prefix,
68 +)
69 +from vllm.multimodal import MULTIMODAL_REGISTRY
70 +from vllm.multimodal.inputs import MultiModalFeatureSpec
71 +from vllm.sequence import IntermediateTensors
72 +from vllm.tokenizers.registry import cached_tokenizer_from_config
73 +from vllm.transformers_utils.configs.qwen4_exp import (
74 + Qwen4ExpTextConfig,
75 +)
76 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
77 +from vllm.v1.kv_cache_interface import MambaSpec
78 +
79 +from ..config import Qwen4ExpConfig
80 +from .hyperconnection import GatedResidual, HyperConnectionConfig
81 +from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
82 +from .ple_layer import Qwen4ExpPLELayer
83 +from .qsa import Qwen4ExpQSAAttention
84 +
85 +
86 +def without_modelopt_fp4(
87 + quant_config: QuantizationConfig | None,
88 +) -> QuantizationConfig | None:
89 + """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4."""
90 +
91 + if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
92 + return None
93 + return quant_config
94 +
95 +
96 +def _remap_qsa_cache_scale_name(
97 + name: str,
98 + qsa_layer_ids: frozenset[int],
99 +) -> str:
100 + """Map serialized main-cache scales onto the merged QSA owner.
101 +
102 + Regular attention keeps cache scales below its ``attn`` child. QSA owns
103 + that cache directly, so only QSA layers need the final path component
104 + moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers.
105 + """
106 +
107 + scale_suffixes = {
108 + "k_proj.k_scale": "_k_scale",
109 + "k_proj.output_scale": "_k_scale",
110 + "attn.k_scale": "_k_scale",
111 + "attn._k_scale": "_k_scale",
112 + "k_scale": "_k_scale",
113 + "_k_scale": "_k_scale",
114 + "v_proj.v_scale": "_v_scale",
115 + "v_proj.output_scale": "_v_scale",
116 + "attn.v_scale": "_v_scale",
117 + "attn._v_scale": "_v_scale",
118 + "v_scale": "_v_scale",
119 + "_v_scale": "_v_scale",
120 + }
121 + for layer_id in qsa_layer_ids:
122 + marker = f"layers.{layer_id}.self_attn."
123 + marker_start = name.find(marker)
124 + if marker_start < 0 or (marker_start > 0 and name[marker_start - 1] != "."):
125 + continue
126 + suffix = name[marker_start + len(marker) :]
127 + mapped_suffix = scale_suffixes.get(suffix)
128 + if mapped_suffix is not None:
129 + return f"{name[: marker_start + len(marker)]}{mapped_suffix}"
130 + return name
131 +
132 +
133 +_QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [
134 + ".bias",
135 + "_bias",
136 + ".k_scale",
137 + "_k_scale",
138 + ".v_scale",
139 + "_v_scale",
140 + "_weight_scale",
141 + "_input_scale",
142 +]
143 +
144 +# The checkpoint keeps down and injection projections separate; runtime packs
145 +# them into adjacent logical shards of one MergedColumnParallelLinear.
146 +_HC_WEIGHTS_MAPPER = WeightsMapper(
147 + orig_to_new_stacked={
148 + "hyper_connection.input_mix_weight_down.weight": (
149 + "hyper_connection.input_mix_weight_down_block_inject.weight",
150 + 0,
151 + ),
152 + "hyper_connection.block_inject_weight.weight": (
153 + "hyper_connection.input_mix_weight_down_block_inject.weight",
154 + 1,
155 + ),
156 + }
157 +)
158 +
159 +
160 +class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock):
161 + """Qwen3Next MoE with Qwen4Exp HC validation."""
162 +
163 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
164 + parallel_config = vllm_config.parallel_config
165 + if parallel_config.use_sequence_parallel_moe:
166 + raise NotImplementedError(
167 + "Qwen4Exp HC does not support sequence-parallel MoE"
168 + )
169 + super().__init__(vllm_config=vllm_config, prefix=prefix)
170 + config = vllm_config.model_config.hf_text_config
171 + self.n_shared_experts = int(config.shared_expert_intermediate_size > 0)
172 +
173 +
174 +class Qwen4ExpDecoderLayer(nn.Module):
175 + def __init__(
176 + self,
177 + vllm_config: VllmConfig,
178 + layer_type: str,
179 + prefix: str = "",
180 + ) -> None:
181 + super().__init__()
182 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
183 + model_config = vllm_config.model_config
184 + cache_config = vllm_config.cache_config
185 + quant_config = vllm_config.quant_config
186 +
187 + self.config = config
188 + self.layer_type = layer_type
189 + self.layer_idx = extract_layer_index(prefix)
190 + if vllm_config.parallel_config.use_sequence_parallel_moe:
191 + raise NotImplementedError(
192 + "Qwen4Exp HC does not support sequence-parallel MoE"
193 + )
194 + self.ple: Qwen4ExpPLELayer | None = None
195 + ple_layer_ids = config.ple_layer_ids
196 + if (self.layer_idx + 1) in ple_layer_ids:
197 + ple_layer_ids_sorted = sorted(set(ple_layer_ids))
198 + ple_dense_layer_id_map = {
199 + abs_id: idx for idx, abs_id in enumerate(ple_layer_ids_sorted)
200 + }
201 + ple_dense_layer_id = ple_dense_layer_id_map[self.layer_idx + 1]
202 + self.ple = Qwen4ExpPLELayer(
203 + config,
204 + vllm_config=vllm_config,
205 + layer_idx=self.layer_idx,
206 + ple_dense_layer_id=ple_dense_layer_id,
207 + prefix=f"{prefix}.ple",
208 + )
209 +
210 + if layer_type == "linear_attention":
211 + self.linear_attn = QwenGatedDeltaNetAttention(
212 + config,
213 + vllm_config=vllm_config,
214 + prefix=f"{prefix}.linear_attn",
215 + gqa_interleaved_layout=False,
216 + )
217 + elif layer_type == "full_attention":
218 + use_qsa = getattr(config, "indexer_n_heads", None) is not None
219 + if not use_qsa:
220 + self.self_attn = Qwen3NextAttention(
221 + config,
222 + model_config=model_config,
223 + cache_config=cache_config,
224 + quant_config=quant_config,
225 + prefix=f"{prefix}.self_attn",
226 + )
227 + else:
228 + self.self_attn = Qwen4ExpQSAAttention(
229 + vllm_config=vllm_config,
230 + config=config,
231 + layer_id=self.layer_idx,
232 + quant_config=quant_config,
233 + prefix=f"{prefix}.self_attn",
234 + )
235 + else:
236 + raise ValueError(f"Invalid layer_type {layer_type}")
237 +
238 + mlp_only_layers = getattr(config, "mlp_only_layers", [])
239 + num_experts = getattr(config, "num_experts", 0) or 0
240 + absolute_layer_id = self.layer_idx + 1
241 + is_moe_layer = self.layer_idx not in mlp_only_layers and (
242 + num_experts > 0 and absolute_layer_id % config.decoder_sparse_step == 0
243 + )
244 + if is_moe_layer:
245 + self.mlp = Qwen4ExpSparseMoeBlock(
246 + vllm_config=vllm_config, prefix=f"{prefix}.mlp"
247 + )
248 + else:
249 + self.mlp = Qwen3NextMLP(
250 + hidden_size=config.hidden_size,
251 + intermediate_size=config.intermediate_size,
252 + hidden_act=config.hidden_act,
253 + quant_config=quant_config,
254 + prefix=f"{prefix}.mlp",
255 + )
256 +
257 + hc_config = HyperConnectionConfig(
258 + hc_count=config.hc_count,
259 + hidden_size=config.hidden_size,
260 + params_dtype=torch.bfloat16,
261 + hc_lowrank=config.hc_lowrank,
262 + rms_norm_eps=config.rms_norm_eps,
263 + hc_per_branch_norm=True,
264 + )
265 + self.attn_hyper_connection = GatedResidual(
266 + hc_config,
267 + prefix=maybe_prefix(prefix, "attn_hyper_connection"),
268 + )
269 + self.mlp_hyper_connection = GatedResidual(
270 + hc_config,
271 + prefix=maybe_prefix(prefix, "mlp_hyper_connection"),
272 + )
273 +
274 + def forward(
275 + self,
276 + hidden_states: torch.Tensor,
277 + prev_block_output: torch.Tensor | None,
278 + prev_injection: torch.Tensor | None,
279 + positions: torch.Tensor,
280 + *,
281 + input_ids: torch.Tensor | None,
282 + query_start_loc: torch.Tensor | None,
283 + ngram_context: torch.Tensor | None,
284 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
285 + attn_hc = self.attn_hyper_connection
286 + if self.ple is not None:
287 + # PLE adds directly to the multi-stream state, so pending HC state
288 + # must be materialized before the addition.
289 + if prev_block_output is not None and prev_injection is not None:
290 + hidden_states = attn_hc.combine(
291 + hidden_states, prev_block_output, prev_injection
292 + )
293 + prev_block_output = prev_injection = None
294 +
295 + if input_ids is None or query_start_loc is None or ngram_context is None:
296 + raise RuntimeError("PLE inputs were not prepared")
297 + hidden_states = hidden_states + self.ple(
298 + hidden_states,
299 + input_ids,
300 + query_start_loc,
301 + ngram_context,
302 + )
303 +
304 + # Fuse a pending combine with this HC module's mix when possible.
305 + if prev_block_output is not None and prev_injection is not None:
306 + hidden_states, block_input, injection = attn_hc.combine_and_mix(
307 + hidden_states, prev_block_output, prev_injection
308 + )
309 + else:
310 + hidden_states, block_input, injection = attn_hc.mix(hidden_states)
311 +
312 + if self.layer_type == "linear_attention":
313 + attn_out = self.linear_attn(hidden_states=block_input)
314 + elif self.layer_type == "full_attention":
315 + attn_out = self.self_attn(
316 + hidden_states=block_input,
317 + positions=positions,
318 + )
319 + else:
320 + raise ValueError("Invalid layer_type")
321 +
322 + mlp_hc = self.mlp_hyper_connection
323 + hidden_states, block_input, injection = mlp_hc.combine_and_mix(
324 + hidden_states, attn_out, injection
325 + )
326 + mlp_out = self.mlp(block_input)
327 + return hidden_states, mlp_out, injection
328 +
329 +
330 +class Qwen4ExpMixtureOfExperts(MixtureOfExperts):
331 + """Expose Qwen4Exp routed experts through vLLM's EPLB protocol."""
332 +
333 + def set_moe_parameters(self, layers: Iterable[nn.Module]) -> None:
334 + self.moe_layers = []
335 + self.moe_mlp_layers = []
336 + example_moe = None
337 + for layer in layers:
338 + if isinstance(layer, Qwen4ExpDecoderLayer) and isinstance(
339 + layer.mlp, Qwen4ExpSparseMoeBlock
340 + ):
341 + example_moe = layer.mlp
342 + self.moe_mlp_layers.append(layer.mlp)
343 + self.moe_layers.append(layer.mlp.experts)
344 +
345 + self.num_moe_layers = len(self.moe_layers)
346 + if example_moe is None:
347 + self.num_expert_groups = 0
348 + self.num_shared_experts = 0
349 + self.num_logical_experts = 0
350 + self.num_physical_experts = 0
351 + self.num_local_physical_experts = 0
352 + self.num_routed_experts = 0
353 + self.num_redundant_experts = 0
354 + return
355 +
356 + self.num_expert_groups = 1
357 + self.num_shared_experts = example_moe.n_shared_experts
358 + self.num_logical_experts = example_moe.n_logical_experts
359 + self.num_physical_experts = example_moe.n_physical_experts
360 + self.num_local_physical_experts = example_moe.n_local_physical_experts
361 + self.num_routed_experts = example_moe.n_routed_experts
362 + self.num_redundant_experts = example_moe.n_redundant_experts
363 +
364 + def update_physical_experts_metadata(
365 + self,
366 + num_physical_experts: int,
367 + num_local_physical_experts: int,
368 + ) -> None:
369 + assert self.num_local_physical_experts == num_local_physical_experts
370 + self.num_physical_experts = num_physical_experts
371 + self.num_local_physical_experts = num_local_physical_experts
372 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
373 + for moe in self.moe_mlp_layers:
374 + moe.n_physical_experts = num_physical_experts
375 + moe.n_local_physical_experts = num_local_physical_experts
376 + moe.n_redundant_experts = self.num_redundant_experts
377 + moe.experts.update_expert_map()
378 +
379 +
380 +@support_torch_compile(
381 + dynamic_arg_dims={
382 + "input_ids": 0,
383 + "positions": -1,
384 + "intermediate_tensors": 0,
385 + "inputs_embeds": 0,
386 + "query_start_loc": 0,
387 + "ngram_context": 0,
388 + "deepstack_input_embeds": 0,
389 + }
390 +)
391 +class Qwen4ExpModel(nn.Module):
392 + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
393 +
394 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
395 + super().__init__()
396 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
397 + self.config = config
398 + self.num_redundant_experts = (
399 + vllm_config.parallel_config.eplb_config.num_redundant_experts
400 + )
401 + self.vocab_size = config.vocab_size
402 + self._qsa_layer_ids = frozenset(
403 + layer_idx
404 + for layer_idx, layer_type in enumerate(config.layer_types)
405 + if layer_type == "full_attention"
406 + and getattr(config, "indexer_n_heads", None) is not None
407 + )
408 + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size)
409 +
410 + def get_layer(prefix: str) -> Qwen4ExpDecoderLayer:
411 + layer_idx = extract_layer_index(prefix)
412 + return Qwen4ExpDecoderLayer(
413 + vllm_config,
414 + layer_type=config.layer_types[layer_idx],
415 + prefix=prefix,
416 + )
417 +
418 + self.start_layer, self.end_layer, self.layers = make_layers(
419 + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers"
420 + )
421 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
422 + self.layers,
423 + Qwen4ExpSparseMoeBlock,
424 + "mlp",
425 + )
426 + intermediate_size = config.hidden_size * config.hc_count
427 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
428 + ["hidden_states"], intermediate_size
429 + )
430 +
431 + self.hyper_connection_mixer: GatedResidual | None
432 + if get_pp_group().is_last_rank:
433 + hc_config = HyperConnectionConfig(
434 + hc_count=config.hc_count,
435 + hidden_size=config.hidden_size,
436 + params_dtype=torch.bfloat16,
437 + hc_lowrank=config.hc_lowrank,
438 + rms_norm_eps=config.rms_norm_eps,
439 + hc_per_branch_norm=True,
440 + )
441 + self.hyper_connection_mixer = GatedResidual(
442 + hc_config,
443 + use_combine=False,
444 + prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
445 + )
446 + else:
447 + self.hyper_connection_mixer = None
448 +
449 + spec_config = vllm_config.speculative_config
450 + # MTP HC multi-stream outputs: when speculative method=="mtp" and the
451 + # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream
452 + # hidden state [T, hc_count*H] so the MTP drafter can feed a real
453 + # multi-stream backbone hidden on its first step (scheme A). Derived
454 + # purely from config (NOT node identity) so P/D nodes stay consistent.
455 + needs_mtp_hidden = (
456 + spec_config is not None
457 + and getattr(spec_config, "method", None) == "mtp"
458 + and get_pp_group().is_last_rank
459 + )
460 + if needs_mtp_hidden:
461 + self._mtp_hidden_buffer = torch.empty(
462 + vllm_config.scheduler_config.max_num_batched_tokens,
463 + config.hc_count * config.hidden_size,
464 + dtype=vllm_config.model_config.dtype,
465 + )
466 + else:
467 + self._mtp_hidden_buffer = None
468 +
469 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
470 + return self.embed_tokens(input_ids)
471 +
472 + def forward(
473 + self,
474 + input_ids: torch.Tensor | None,
475 + positions: torch.Tensor,
476 + intermediate_tensors: IntermediateTensors | None = None,
477 + inputs_embeds: torch.Tensor | None = None,
478 + query_start_loc: torch.Tensor | None = None,
479 + ngram_context: torch.Tensor | None = None,
480 + deepstack_input_embeds: IntermediateTensors | None = None,
481 + ) -> torch.Tensor | IntermediateTensors:
482 + if get_pp_group().is_first_rank:
483 + if inputs_embeds is not None:
484 + hidden_states = inputs_embeds
485 + else:
486 + if input_ids is None:
487 + raise ValueError("input_ids or inputs_embeds is required")
488 + hidden_states = self.embed_input_ids(input_ids)
489 + hidden_states = hidden_states.repeat(1, self.config.hc_count)
490 + else:
491 + if intermediate_tensors is None:
492 + raise ValueError("pipeline stage requires intermediate tensors")
493 + hidden_states = intermediate_tensors["hidden_states"]
494 +
495 + block_output = None
496 + injection = None
497 + last_layer = None
498 + for layer_idx, layer in islice(
499 + enumerate(self.layers), self.start_layer, self.end_layer
500 + ):
501 + last_layer = layer
502 + hidden_states, block_output, injection = layer(
503 + hidden_states=hidden_states,
504 + prev_block_output=block_output,
505 + prev_injection=injection,
506 + positions=positions,
507 + input_ids=input_ids,
508 + query_start_loc=query_start_loc,
509 + ngram_context=ngram_context,
510 + )
511 + if deepstack_input_embeds is not None and layer_idx < len(
512 + deepstack_input_embeds
513 + ):
514 + deepstack_embed = deepstack_input_embeds[
515 + f"deepstack_input_embeds_{layer_idx}"
516 + ]
517 + deepstack_embed = (
518 + deepstack_embed.unsqueeze(-2)
519 + .expand(
520 + *deepstack_embed.shape[:-1],
521 + self.config.hc_count,
522 + self.config.hidden_size,
523 + )
524 + .flatten(-2)
525 + )
526 + # Deepstack is an external addition to the materialized
527 + # multi-stream state and therefore terminates delayed combine.
528 + hidden_states = layer.mlp_hyper_connection.combine(
529 + hidden_states, block_output, injection
530 + )
531 + block_output = None
532 + injection = None
533 + hidden_states = hidden_states + deepstack_embed
534 +
535 + if not get_pp_group().is_last_rank:
536 + # PP transports one tensor, not the delayed HC tuple. Materialize
537 + # with the HC module that produced the pending injection.
538 + if last_layer is not None and block_output is not None:
539 + hidden_states = last_layer.mlp_hyper_connection.combine(
540 + hidden_states, block_output, injection
541 + )
542 + return IntermediateTensors({"hidden_states": hidden_states})
543 +
544 + # The final mixer consumes the last pending combine and returns both
545 + # the sampled single stream and the materialized multi-stream state.
546 + final_mixer = self.hyper_connection_mixer
547 + assert final_mixer is not None
548 + multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix(
549 + hidden_states, block_output, injection
550 + )
551 + if self._mtp_hidden_buffer is not None:
552 + # Capture the pre-final-mixer multi-stream hidden state
553 + # [T, hc_count*H] for the MTP drafter (zero extra compute:
554 + # this tensor is needed by the final mixer regardless).
555 + num_tokens = multi_hidden.shape[0]
556 + self._mtp_hidden_buffer[:num_tokens].copy_(multi_hidden)
557 + return sample_hidden_states
558 +
559 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
560 + weights = (
561 + (
562 + _remap_qsa_cache_scale_name(name, self._qsa_layer_ids),
563 + weight,
564 + )
565 + for name, weight in weights
566 + )
567 + weights = maybe_fuse_shared_experts(
568 + weights,
569 + enabled=self.is_fused_shared_expert_enabled,
570 + n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
571 + n_shared_experts=1,
572 + ckpt_prefix="mlp.shared_expert",
573 + )
574 + # Non-persistent PLE state rebuilt in __init__; skip any ckpt
575 + # column for them.
576 + skip_substrs = (
577 + "hashstats_",
578 + "token_lookup",
579 + "hyper_connection_mixer.block_inject_weight",
580 + )
581 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
582 + orig_to_new_substr={substr: None for substr in skip_substrs}
583 + )
584 + # The final HC mixer only exists on the last PP rank; earlier ranks
585 + # must drop its checkpoint weights instead of failing to place them.
586 + ignore_prefixes = (
587 + None
588 + if self.hyper_connection_mixer is not None
589 + else ["hyper_connection_mixer."]
590 + )
591 + loader = AutoWeightsLoader(
592 + self,
593 + ignore_unexpected_prefixes=ignore_prefixes,
594 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
595 + )
596 + loaded = loader.load_weights(
597 + weights,
598 + mapper=mapper,
599 + )
600 + return loaded
601 +
602 +
603 +class Qwen4ExpForCausalLM(
604 + nn.Module,
605 + HasInnerState,
606 + SupportsLoRA,
607 + SupportsMRoPE,
608 + SupportsPP,
609 + Qwen4ExpMixtureOfExperts,
610 + IsHybrid,
611 +):
612 + packed_modules_mapping = {
613 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
614 + "gate_up_proj": ["gate_proj", "up_proj"],
615 + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
616 + "in_proj_ba": ["in_proj_b", "in_proj_a"],
617 + "input_mix_weight_down_block_inject": [
618 + "input_mix_weight_down",
619 + "block_inject_weight",
620 + "_input_mix_padding",
621 + ],
622 + }
623 + hf_to_vllm_mapper = WeightsMapper(
624 + orig_to_new_prefix={"model.language_model.": "model."}
625 + )
626 + requires_raw_input_tokens = True
627 +
628 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
629 + super().__init__()
630 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
631 + self.vllm_config = vllm_config
632 + self.model_config = vllm_config.model_config
633 + self.quant_config = vllm_config.quant_config
634 + self.config = config
635 + self.scheduler_config = vllm_config.scheduler_config
636 + if vllm_config.cache_config.mamba_cache_mode == "all":
637 + raise NotImplementedError(
638 + "Qwen4Exp currently does not support 'all' prefix caching, "
639 + "please use '--mamba-cache-mode=align' instead"
640 + )
641 + self.model = Qwen4ExpModel(
642 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
643 + )
644 + self.lm_head = ParallelLMHead(
645 + config.vocab_size,
646 + config.hidden_size,
647 + prefix=maybe_prefix(prefix, "lm_head"),
648 + )
649 + self.logits_processor = LogitsProcessor(config.vocab_size)
650 + self.make_empty_intermediate_tensors = (
651 + self.model.make_empty_intermediate_tensors
652 + )
653 + self.set_moe_parameters(self.model.layers)
654 + enable_qwen4_exp_low_latency_gemm(self, self.model_config.dtype)
655 +
656 + @staticmethod
657 + def get_model_state_cls():
658 + from .model_state import Qwen4ExpModelState
659 +
660 + return Qwen4ExpModelState
661 +
662 + def forward(
663 + self,
664 + input_ids: torch.Tensor | None,
665 + positions: torch.Tensor,
666 + intermediate_tensors: IntermediateTensors | None = None,
667 + inputs_embeds: torch.Tensor | None = None,
668 + **kwargs: object,
669 + ) -> torch.Tensor | IntermediateTensors:
670 + # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs
671 + # path (query_start_loc / ngram_context) reaches Qwen4ExpModel.
672 + return self.model(
673 + input_ids,
674 + positions,
675 + intermediate_tensors,
676 + inputs_embeds,
677 + **kwargs,
678 + )
679 +
680 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
681 + return self.model.embed_input_ids(input_ids)
682 +
683 + @classmethod
684 + def get_ple_mamba_state_dtype_from_config(
685 + cls,
686 + vllm_config: VllmConfig,
687 + ) -> tuple[torch.dtype, ...]:
688 + return MambaStateDtypeCalculator.short_conv_state_dtype(
689 + vllm_config.model_config.dtype,
690 + vllm_config.cache_config.mamba_cache_dtype,
691 + )
692 +
693 + @classmethod
694 + def get_ple_mamba_state_shape_from_config(
695 + cls, vllm_config: VllmConfig
696 + ) -> tuple[tuple[int, int]]:
697 + hf_config = vllm_config.model_config.hf_text_config
698 + conv_kernel_size = hf_config.ple_conv_kernel_size
699 + short_conv_dilation = hf_config.ngram_size
700 + conv_state_len = (conv_kernel_size - 1) * short_conv_dilation
701 + num_spec = (
702 + vllm_config.speculative_config.num_speculative_tokens
703 + if vllm_config.speculative_config
704 + else 0
705 + )
706 + hc_count = hf_config.hc_count
707 + hc_hidden_size = hf_config.hidden_size * hc_count
708 + return MambaStateShapeCalculator.short_conv_state_shape(
709 + tp_world_size=1,
710 + intermediate_size=hc_hidden_size,
711 + conv_kernel=conv_state_len + 1,
712 + num_spec=num_spec,
713 + )
714 +
715 + @classmethod
716 + def get_gdn_mamba_state_dtype_from_config(
717 + cls, vllm_config: VllmConfig
718 + ) -> tuple[torch.dtype, torch.dtype]:
719 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
720 + vllm_config.model_config.dtype,
721 + vllm_config.cache_config.mamba_cache_dtype,
722 + vllm_config.cache_config.mamba_ssm_cache_dtype,
723 + )
724 +
725 + @classmethod
726 + def get_gdn_mamba_state_shape_from_config(
727 + cls, vllm_config: VllmConfig
728 + ) -> tuple[tuple[int, int], tuple[int, int]]:
729 + parallel_config = vllm_config.parallel_config
730 + hf_config = vllm_config.model_config.hf_text_config
731 + tp_size = parallel_config.tensor_parallel_size
732 + num_spec = (
733 + vllm_config.speculative_config.num_speculative_tokens
734 + if vllm_config.speculative_config
735 + else 0
736 + )
737 + return MambaStateShapeCalculator.gated_delta_net_state_shape(
738 + tp_size,
739 + hf_config.linear_num_key_heads,
740 + hf_config.linear_num_value_heads,
741 + hf_config.linear_key_head_dim,
742 + hf_config.linear_value_head_dim,
743 + hf_config.linear_conv_kernel_dim,
744 + num_spec,
745 + )
746 +
747 + @classmethod
748 + def get_mamba_state_dtype_from_config(
749 + cls,
750 + vllm_config: VllmConfig,
751 + ) -> tuple[torch.dtype, torch.dtype]:
752 + return cls.get_gdn_mamba_state_dtype_from_config(vllm_config)
753 +
754 + @classmethod
755 + def get_mamba_state_shape_from_config(
756 + cls, vllm_config: VllmConfig
757 + ) -> tuple[tuple[int, int], tuple[int, int]]:
758 + return cls.get_gdn_mamba_state_shape_from_config(vllm_config)
759 +
760 + @classmethod
761 + def get_mamba_state_copy_func(
762 + cls,
763 + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
764 + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
765 +
766 + @classmethod
767 + def get_mamba_state_copy_funcs(
768 + cls,
769 + mamba_types: set[MambaAttentionBackendEnum],
770 + ) -> MambaStateCopyFuncsByType:
771 + copy_funcs_by_type = {
772 + MambaAttentionBackendEnum.GDN_ATTN: cls.get_mamba_state_copy_func(),
773 + MambaAttentionBackendEnum.SHORT_CONV: (
774 + MambaStateCopyFuncCalculator.short_conv_state_copy_func()
775 + ),
776 + }
777 + missing_types = mamba_types - copy_funcs_by_type.keys()
778 + assert not missing_types, f"missing state copy funcs for {missing_types}"
779 + return {
780 + mamba_type: copy_funcs_by_type[mamba_type] for mamba_type in mamba_types
781 + }
782 +
783 + @classmethod
784 + def get_mamba_specs_from_config(
785 + cls, vllm_config: VllmConfig
786 + ) -> tuple[MambaSpec, ...]:
787 + """Return all MambaSpecs for this model (GDN layers + PLE layer).
788 +
789 + The PLE layer uses a separate short_conv MambaSpec whose page_size_bytes
790 + may exceed the GDN spec; callers should take the maximum.
791 + """
792 + return (
793 + MambaSpec(
794 + shapes=cls.get_gdn_mamba_state_shape_from_config(vllm_config),
795 + dtypes=cls.get_gdn_mamba_state_dtype_from_config(vllm_config),
796 + block_size=-1,
797 + ),
798 + MambaSpec(
799 + shapes=cls.get_ple_mamba_state_shape_from_config(vllm_config),
800 + dtypes=cls.get_ple_mamba_state_dtype_from_config(vllm_config),
801 + block_size=-1,
802 + tp_replicated=True,
803 + ),
804 + )
805 +
806 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
807 + return self.logits_processor(self.lm_head, hidden_states)
808 +
809 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
810 + return self.model._mtp_hidden_buffer
811 +
812 + def get_mrope_input_positions(
813 + self,
814 + input_tokens: list[int],
815 + mm_features: list[MultiModalFeatureSpec],
816 + ) -> tuple[torch.Tensor, int]:
817 + positions = torch.arange(len(input_tokens), dtype=torch.long)
818 + return positions.unsqueeze(0).expand(3, -1), 0
819 +
820 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
821 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
822 + orig_to_new_substr={"mtp.": None}
823 + )
824 + loader = AutoWeightsLoader(
825 + self,
826 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
827 + )
828 + return loader.load_weights(weights, mapper=mapper)
829 +
830 +
831 +class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo):
832 + def get_hf_config(self) -> Qwen4ExpConfig:
833 + return self.ctx.get_hf_config(Qwen4ExpConfig)
834 +
835 +
836 +@MULTIMODAL_REGISTRY.register_processor(
837 + Qwen3VLMultiModalProcessor,
838 + info=Qwen4ExpProcessingInfo,
839 + dummy_inputs=Qwen3VLDummyInputsBuilder,
840 +)
841 +class Qwen4ExpForConditionalGeneration(
842 + Qwen3_5ForConditionalGeneration,
843 + HasInnerState,
844 + Qwen4ExpMixtureOfExperts,
845 +):
846 + """Qwen3-VL vision tower backed by the Qwen4Exp language model."""
847 +
848 + requires_raw_input_tokens = True
849 +
850 + packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | {
851 + "input_mix_weight_down_block_inject": [
852 + "input_mix_weight_down",
853 + "block_inject_weight",
854 + "_input_mix_padding",
855 + ]
856 + }
857 +
858 + @staticmethod
859 + def get_model_state_cls():
860 + from .model_state import Qwen4ExpModelState
861 +
862 + return Qwen4ExpModelState
863 +
864 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None:
865 + nn.Module.__init__(self)
866 + config: Qwen4ExpConfig = vllm_config.model_config.hf_config
867 + quant_config = vllm_config.quant_config
868 + multimodal_config = vllm_config.model_config.multimodal_config
869 + if multimodal_config is None:
870 + raise ValueError(
871 + "Qwen4ExpForConditionalGeneration requires multimodal_config"
872 + )
873 +
874 + self.config = config
875 + self.model_config = vllm_config.model_config
876 + self.multimodal_config = multimodal_config
877 + self.language_model_only = multimodal_config.language_model_only
878 + if self.language_model_only:
879 + self.use_data_parallel = False
880 + self.is_multimodal_pruning_enabled = False
881 + self.video_pruning_method = None
882 + self.video_pruning_rate = 0.0
883 + self._tokenizer = None
884 + self.visual = StageMissingLayer("vision_tower")
885 + self._tower_model_names = []
886 + else:
887 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
888 + self._init_video_pruning(multimodal_config)
889 + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config)
890 +
891 + with self._mark_tower_model(vllm_config, {"image", "video"}):
892 + self.visual = Qwen3_VisionTransformer(
893 + config.vision_config,
894 + norm_eps=config.text_config.rms_norm_eps,
895 + quant_config=quant_config,
896 + prefix=maybe_prefix(prefix, "visual"),
897 + )
898 +
899 + self.use_deepstack = (
900 + not self.language_model_only
901 + and bool(config.vision_config.deepstack_visual_indexes)
902 + and not isinstance(self.visual, StageMissingLayer)
903 + )
904 + self.deepstack_num_level = (
905 + len(config.vision_config.deepstack_visual_indexes)
906 + if self.use_deepstack
907 + else 0
908 + )
909 + self.visual_dim = config.vision_config.out_hidden_size
910 + self.multiscale_dim = self.visual_dim * self.deepstack_num_level
911 +
912 + if self.use_deepstack:
913 + self.deepstack_input_embeds = [
914 + torch.zeros(
915 + vllm_config.scheduler_config.max_num_batched_tokens,
916 + config.text_config.hidden_size,
917 + )
918 + for _ in range(self.deepstack_num_level)
919 + ]
920 + self.deepstack_input_embeds_num_tokens = 0
921 +
922 + with self._mark_language_model(vllm_config):
923 + self.language_model = Qwen4ExpForCausalLM(
924 + vllm_config=vllm_config,
925 + prefix=maybe_prefix(prefix, "language_model"),
926 + )
927 +
928 + self.make_empty_intermediate_tensors = (
929 + self.language_model.make_empty_intermediate_tensors
930 + )
931 + if not get_pp_group().is_first_rank and self.use_deepstack:
932 + assert self.language_model.model.start_layer >= len(
933 + config.vision_config.deepstack_visual_indexes
934 + ), (
935 + "start_layer should be greater than or equal to "
936 + "len(deepstack_visual_indexes)"
937 + )
938 + self.set_moe_parameters(self.language_model.model.layers)
939 +
940 + def embed_input_ids(
941 + self,
942 + input_ids: torch.Tensor,
943 + multimodal_embeddings: MultiModalEmbeddings | None = None,
944 + *,
945 + is_multimodal: torch.Tensor | None = None,
946 + ) -> torch.Tensor:
947 + inputs_embeds = self._embed_text_input_ids(
948 + input_ids,
949 + self.language_model.embed_input_ids,
950 + is_multimodal=is_multimodal,
951 + )
952 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
953 + return inputs_embeds
954 + if self.language_model_only:
955 + raise ValueError(
956 + "Qwen4Exp language_model_only does not accept multimodal embeddings"
957 + )
958 +
959 + is_multimodal = _require_is_multimodal(is_multimodal)
960 + if self.use_deepstack:
961 + deepstack_input_embeds, multimodal_embeddings = (
962 + self._compute_deepstack_embeds(
963 + inputs_embeds=inputs_embeds,
964 + multimodal_embeddings=multimodal_embeddings,
965 + is_multimodal=is_multimodal,
966 + )
967 + )
968 + else:
969 + deepstack_input_embeds = None
970 +
971 + inputs_embeds = _merge_multimodal_embeddings(
972 + inputs_embeds=inputs_embeds,
973 + multimodal_embeddings=multimodal_embeddings,
974 + is_multimodal=is_multimodal,
975 + )
976 + if deepstack_input_embeds is not None:
977 + self._set_deepstack_input_embeds(deepstack_input_embeds)
978 + return inputs_embeds
979 +
980 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
981 + return self.language_model.get_mtp_target_hidden_states()
982 +
983 + def forward(
984 + self,
985 + input_ids: torch.Tensor | None,
986 + positions: torch.Tensor,
987 + intermediate_tensors: IntermediateTensors | None = None,
988 + inputs_embeds: torch.Tensor | None = None,
989 + **kwargs: object,
990 + ) -> torch.Tensor | IntermediateTensors:
991 + if intermediate_tensors is not None:
992 + inputs_embeds = None
993 + if inputs_embeds is not None and get_pp_group().is_first_rank:
994 + deepstack_input_embeds = self._get_deepstack_input_embeds(
995 + inputs_embeds.size(0)
996 + )
997 + else:
998 + deepstack_input_embeds = None
999 +
1000 + hidden_states = self.language_model.model(
1001 + input_ids=input_ids,
1002 + positions=positions,
1003 + intermediate_tensors=intermediate_tensors,
1004 + inputs_embeds=inputs_embeds,
1005 + query_start_loc=kwargs.get("query_start_loc"),
1006 + ngram_context=kwargs.get("ngram_context"),
1007 + deepstack_input_embeds=deepstack_input_embeds,
1008 + )
1009 + if inputs_embeds is not None and get_pp_group().is_first_rank:
1010 + self._clear_deepstack_input_embeds(inputs_embeds.size(0))
1011 + return hidden_states
1012 +
1013 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1014 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
1015 + orig_to_new_substr={"mtp.": None},
1016 + orig_to_new_prefix={"visual.": None} if self.language_model_only else {},
1017 + )
1018 + loader = AutoWeightsLoader(
1019 + self,
1020 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
1021 + )
1022 + return loader.load_weights(weights, mapper=mapper)
1023 +
1024 + @classmethod
1025 + def get_mamba_state_dtype_from_config(
1026 + cls,
1027 + vllm_config: VllmConfig,
1028 + ) -> tuple[torch.dtype, torch.dtype]:
1029 + return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
1030 +
1031 + @classmethod
1032 + def get_mamba_state_shape_from_config(
1033 + cls,
1034 + vllm_config: VllmConfig,
1035 + ) -> tuple[tuple[int, int], tuple[int, int]]:
1036 + return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config)
1037 +
1038 + @classmethod
1039 + def get_mamba_state_copy_func(
1040 + cls,
1041 + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
1042 + return Qwen4ExpForCausalLM.get_mamba_state_copy_func()
1043 +
1044 + @classmethod
1045 + def get_mamba_state_copy_funcs(
1046 + cls,
1047 + mamba_types: set[MambaAttentionBackendEnum],
1048 + ) -> MambaStateCopyFuncsByType:
1049 + return Qwen4ExpForCausalLM.get_mamba_state_copy_funcs(mamba_types)
1050 +
1051 + @classmethod
1052 + def get_mamba_specs_from_config(
1053 + cls, vllm_config: VllmConfig
1054 + ) -> tuple[MambaSpec, ...]:
1055 + return Qwen4ExpForCausalLM.get_mamba_specs_from_config(vllm_config)
1056 +
1057 +
1058 +__all__ = [
1059 + "Qwen4ExpDecoderLayer",
1060 + "Qwen4ExpForCausalLM",
1061 + "Qwen4ExpForConditionalGeneration",
1062 + "Qwen4ExpMixtureOfExperts",
1063 + "Qwen4ExpModel",
1064 + "Qwen4ExpSparseMoeBlock",
1065 +]
vllm/models/qwen4_exp/nvidia/model_state.py new
+143
@@ -0,0 +1,143 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Model-runner state for Qwen4Exp PLE inputs."""
4 +
5 +from typing import Any
6 +
7 +import torch
8 +import torch.nn as nn
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.v1.worker.gpu.input_batch import InputBatch
12 +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
13 +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
14 +from vllm.v1.worker.gpu.states import RequestState
15 +
16 +
17 +class Qwen4ExpModelState(MambaHybridModelState):
18 + """Add rollback-safe PLE n-gram context to the model inputs."""
19 +
20 + def __init__(
21 + self,
22 + vllm_config: VllmConfig,
23 + model: nn.Module,
24 + encoder_cache: EncoderCache | None,
25 + device: torch.device,
26 + ) -> None:
27 + super().__init__(vllm_config, model, encoder_cache, device)
28 + config = self.model_config.hf_text_config
29 + self.uses_ngram_embedding = bool(config.ple_layer_ids)
30 + if not self.uses_ngram_embedding:
31 + self.ngram_context_len = 0
32 + self.ngram_eos_token_id = 0
33 + return
34 +
35 + if vllm_config.parallel_config.pipeline_parallel_size > 1:
36 + raise RuntimeError(
37 + "N-gram PLE embedding currently requires "
38 + "pipeline_parallel_size=1 because non-first pipeline ranks do "
39 + "not receive the raw input_ids required by PLE. Please run "
40 + "with PP=1."
41 + )
42 +
43 + self.ngram_context_len = int(config.ngram_size) - 1
44 + if self.ngram_context_len <= 0:
45 + raise ValueError("N-gram embedding requires context length >= 1.")
46 + self.ngram_eos_token_id = int(config.eos_token_id)
47 + self.ngram_context = torch.full(
48 + (self.max_num_reqs, self.ngram_context_len),
49 + self.ngram_eos_token_id,
50 + dtype=torch.int32,
51 + device=self.device,
52 + )
53 + self.ngram_context_offsets = torch.arange(
54 + -self.ngram_context_len,
55 + 0,
56 + dtype=torch.int64,
57 + device=self.device,
58 + )
59 + self.ple_query_start_loc = torch.zeros(
60 + self.max_num_reqs + 1,
61 + dtype=torch.int32,
62 + device=self.device,
63 + )
64 +
65 + def _prepare_ngram_context(
66 + self,
67 + input_batch: InputBatch,
68 + req_states: RequestState,
69 + ) -> torch.Tensor:
70 + num_reqs = input_batch.num_reqs
71 + num_reqs_padded = input_batch.num_reqs_after_padding
72 + context = self.ngram_context[:num_reqs_padded]
73 + context.fill_(self.ngram_eos_token_id)
74 + if num_reqs == 0:
75 + return context
76 +
77 + request_indices = input_batch.idx_mapping[:num_reqs].long()
78 + context_end = req_states.num_computed_tokens.gpu[request_indices].long()
79 + token_indices = context_end.unsqueeze(1) + self.ngram_context_offsets
80 + valid_tokens = token_indices >= 0
81 + token_indices.clamp_min_(0)
82 + context_tokens = req_states.all_token_ids.gpu[
83 + request_indices.unsqueeze(1), token_indices
84 + ]
85 + context[:num_reqs].copy_(
86 + torch.where(
87 + valid_tokens,
88 + context_tokens,
89 + context_tokens.new_full((), self.ngram_eos_token_id),
90 + )
91 + )
92 + return context
93 +
94 + def prepare_inputs(
95 + self,
96 + input_batch: InputBatch,
97 + req_states: RequestState,
98 + ) -> dict[str, Any]:
99 + model_inputs = super().prepare_inputs(input_batch, req_states)
100 + if not self.uses_ngram_embedding:
101 + return model_inputs
102 +
103 + num_reqs_padded = input_batch.num_reqs_after_padding
104 + query_start_loc = self.ple_query_start_loc[: num_reqs_padded + 1]
105 + query_start_loc.copy_(input_batch.query_start_loc[: num_reqs_padded + 1])
106 + model_inputs.update(
107 + query_start_loc=query_start_loc,
108 + ngram_context=self._prepare_ngram_context(input_batch, req_states),
109 + )
110 + return model_inputs
111 +
112 + def prepare_dummy_inputs(
113 + self,
114 + num_reqs: int,
115 + num_tokens: int,
116 + ) -> dict[str, Any]:
117 + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens)
118 + if not self.uses_ngram_embedding:
119 + return model_inputs
120 +
121 + query_start_loc = self.ple_query_start_loc[: num_reqs + 1]
122 + query_start_loc[0] = 0
123 + tokens_per_req, num_extra_tokens = divmod(num_tokens, num_reqs)
124 + query_lens = torch.full(
125 + (num_reqs,),
126 + tokens_per_req,
127 + dtype=query_start_loc.dtype,
128 + device=query_start_loc.device,
129 + )
130 + if num_extra_tokens > 0:
131 + query_lens[-num_extra_tokens:] += 1
132 + torch.cumsum(query_lens, dim=0, out=query_start_loc[1:])
133 +
134 + ngram_context = self.ngram_context[:num_reqs]
135 + ngram_context.fill_(self.ngram_eos_token_id)
136 + model_inputs.update(
137 + query_start_loc=query_start_loc,
138 + ngram_context=ngram_context,
139 + )
140 + return model_inputs
141 +
142 +
143 +__all__ = ["Qwen4ExpModelState"]
vllm/models/qwen4_exp/nvidia/mtp.py new
+462
@@ -0,0 +1,462 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Qwen4Exp MTP (Multi-Token Predictor) model.
4 +
5 +The MTP draft model reuses the Qwen4Exp backbone (PLE/HC/MoE) but:
6 + - drops all multi-modal handling (text-only),
7 + - forces PLE off while keeping the main model's HC stream count,
8 + - fuses the backbone hidden and the new-token embedding via
9 + ``residual_linear_shared`` (fc_embedding + shared fc_hidden) instead of
10 + the ``Linear(2H, H)`` + repeat used by other MTP variants,
11 + - emits TWO hidden streams per step (scheme A): a single stream [T, H]
12 + (final-mixer collapsed, fed to the LM head) and a pre-final-mixer
13 + multi stream [T, hc_count*H] (fed to the next draft step).
14 +"""
15 +
16 +from collections.abc import Iterable
17 +
18 +import regex as re
19 +import torch
20 +from torch import nn
21 +
22 +from vllm.compilation.decorators import support_torch_compile
23 +from vllm.config import VllmConfig, replace, set_current_vllm_config
24 +from vllm.distributed import get_pp_group
25 +from vllm.model_executor.layers.fused_moe.utils import (
26 + is_model_fused_shared_expert_compatible,
27 +)
28 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
29 +from vllm.model_executor.layers.linear import ColumnParallelLinear
30 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
31 +from vllm.model_executor.layers.vocab_parallel_embedding import (
32 + ParallelLMHead,
33 + VocabParallelEmbedding,
34 +)
35 +from vllm.model_executor.model_loader.utils import configure_quant_config
36 +from vllm.model_executor.models.interfaces import SupportsPP
37 +from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
38 +from vllm.model_executor.models.utils import (
39 + AutoWeightsLoader,
40 + PPMissingLayer,
41 + WeightsMapper,
42 + get_draft_quant_config,
43 + make_empty_intermediate_tensors_factory,
44 + maybe_fuse_shared_experts,
45 + maybe_prefix,
46 +)
47 +from vllm.sequence import IntermediateTensors
48 +from vllm.transformers_utils.configs.qwen4_exp import (
49 + Qwen4ExpTextConfig,
50 +)
51 +
52 +from .hyperconnection import GatedResidual, HyperConnectionConfig
53 +from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
54 +from .model import (
55 + _HC_WEIGHTS_MAPPER,
56 + _QWEN4_EXP_IGNORED_MISSING_SUFFIXES,
57 + Qwen4ExpDecoderLayer,
58 + Qwen4ExpMixtureOfExperts,
59 + Qwen4ExpSparseMoeBlock,
60 +)
61 +
62 +
63 +def _remap_ignored_layers(
64 + ignored_layers: list[str],
65 + mtp_start_layer_idx: int,
66 +) -> list[str]:
67 + remapped: list[str] = []
68 + for name in ignored_layers:
69 + if name.startswith("mtp."):
70 + new_name = re.sub(
71 + r"(?<=\.layers\.)\d+",
72 + lambda m: str(mtp_start_layer_idx + int(m.group(0))),
73 + name,
74 + )
75 + remapped.append(new_name)
76 + else:
77 + remapped.append(name)
78 + return remapped
79 +
80 +
81 +def _remap_mtp_weight_name(name: str) -> str | None:
82 + """Map Qwen4Exp checkpoint paths into the standalone draft model."""
83 +
84 + for checkpoint_prefix in (
85 + "model.language_model.",
86 + "language_model.",
87 + ):
88 + if name.startswith(checkpoint_prefix):
89 + name = name.removeprefix(checkpoint_prefix)
90 + break
91 +
92 + if name.startswith("embed_tokens."):
93 + name = f"model.{name}"
94 + if name.startswith("model.mtp."):
95 + name = name.removeprefix("model.")
96 + if name.startswith("mtp.shared_head.head."):
97 + return name.replace("mtp.shared_head.head.", "lm_head.", 1)
98 + if name.startswith("model.shared_head.head."):
99 + return name.replace("model.shared_head.head.", "lm_head.", 1)
100 + if name.startswith("shared_head.head."):
101 + return name.replace("shared_head.head.", "lm_head.", 1)
102 + if name.startswith("model.lm_head."):
103 + return name.removeprefix("model.")
104 + if name.startswith("mtp."):
105 + return name.replace("mtp.", "model.", 1)
106 + if name.startswith("model.embed_tokens.") or name.startswith("lm_head."):
107 + return name
108 + return None
109 +
110 +
111 +def _make_draft_vllm_config(
112 + vllm_config: VllmConfig,
113 + mtp_start_layer_idx: int,
114 +) -> VllmConfig:
115 + """Ensure that the draft model config is set in the vLLM config."""
116 + speculative_config = vllm_config.speculative_config
117 + if speculative_config is None or speculative_config.draft_model_config is None:
118 + raise ValueError("speculative_config.draft_model_config must be set")
119 +
120 + draft_quant_config = get_draft_quant_config(vllm_config)
121 +
122 + # inject packed and ignored modules to the quantization config of draft model
123 + if draft_quant_config is not None:
124 + configure_quant_config(draft_quant_config, Qwen4ExpMTP)
125 + ignored_layers = getattr(draft_quant_config, "ignored_layers", None)
126 + if ignored_layers:
127 + setattr( # noqa: B010
128 + draft_quant_config,
129 + "ignored_layers",
130 + _remap_ignored_layers(ignored_layers, mtp_start_layer_idx),
131 + )
132 + exclude_modules = getattr(draft_quant_config, "exclude_modules", None)
133 + if exclude_modules:
134 + setattr( # noqa: B010
135 + draft_quant_config,
136 + "exclude_modules",
137 + _remap_ignored_layers(exclude_modules, mtp_start_layer_idx),
138 + )
139 +
140 + draft_vllm_config = replace(
141 + vllm_config,
142 + model_config=speculative_config.draft_model_config,
143 + )
144 + # VllmConfig post-init derives the target quant config, so restore the
145 + # independently resolved draft quant config after replacement.
146 + draft_vllm_config.quant_config = draft_quant_config
147 + return draft_vllm_config
148 +
149 +
150 +@support_torch_compile(
151 + dynamic_arg_dims={
152 + "input_ids": 0,
153 + "positions": -1,
154 + "intermediate_tensors": 0,
155 + "inputs_embeds": 0,
156 + "hidden_states": 0,
157 + }
158 +)
159 +class Qwen4ExpMultiTokenPredictor(nn.Module):
160 + hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _HC_WEIGHTS_MAPPER
161 +
162 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
163 + super().__init__()
164 +
165 + model_config = vllm_config.model_config
166 + config: Qwen4ExpTextConfig = model_config.hf_text_config
167 +
168 + self.config = config
169 + self.vocab_size = config.vocab_size
170 +
171 + self.mtp_start_layer_idx = config.num_hidden_layers
172 + self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1)
173 +
174 + self.hidden_size = config.hidden_size
175 + self.hc_count = config.hc_count
176 +
177 + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size)
178 + draft_vllm_config = _make_draft_vllm_config(
179 + vllm_config,
180 + self.mtp_start_layer_idx,
181 + )
182 + with set_current_vllm_config(draft_vllm_config, prefix=prefix):
183 + # residual_linear_shared fusion: fc_embedding projects the token
184 + # embedding, fc_hidden (shared across HC branches) projects the
185 + # backbone hidden; the embedding is added as a residual to every
186 + # branch (see mtp_residual_linear_shared.md).
187 + self.fc_embedding = ColumnParallelLinear(
188 + self.hidden_size,
189 + self.hidden_size,
190 + gather_output=True,
191 + bias=False,
192 + return_bias=False,
193 + quant_config=draft_vllm_config.quant_config,
194 + prefix=f"{prefix}.fc_embedding",
195 + )
196 + self.fc_hidden = ColumnParallelLinear(
197 + self.hidden_size,
198 + self.hidden_size,
199 + gather_output=True,
200 + bias=False,
201 + return_bias=False,
202 + quant_config=draft_vllm_config.quant_config,
203 + prefix=f"{prefix}.fc_hidden",
204 + )
205 + self.layers = nn.ModuleList(
206 + Qwen4ExpDecoderLayer(
207 + draft_vllm_config,
208 + layer_type="full_attention",
209 + prefix=f"{prefix}.layers.{self.mtp_start_layer_idx + idx}",
210 + )
211 + for idx in range(self.num_mtp_layers)
212 + )
213 + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
214 + self.layers,
215 + Qwen4ExpSparseMoeBlock,
216 + "mlp",
217 + )
218 +
219 + self.pre_fc_norm_embedding = GemmaRMSNorm(
220 + self.hidden_size, eps=config.rms_norm_eps
221 + )
222 + self.pre_fc_norm_hidden = GemmaRMSNorm(
223 + self.hidden_size * self.hc_count, eps=config.rms_norm_eps
224 + )
225 + # HC final mixer collapses the multi stream into [T, H] for the LM head.
226 + hc_config = HyperConnectionConfig(
227 + hc_count=config.hc_count,
228 + hidden_size=config.hidden_size,
229 + params_dtype=torch.bfloat16,
230 + hc_lowrank=config.hc_lowrank,
231 + rms_norm_eps=config.rms_norm_eps,
232 + hc_per_branch_norm=True,
233 + )
234 + self.hyper_connection_mixer = GatedResidual(
235 + hc_config,
236 + use_combine=False,
237 + prefix=maybe_prefix(prefix, "hyper_connection_mixer"),
238 + )
239 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
240 + ["hidden_states"], self.hidden_size * self.hc_count
241 + )
242 +
243 + def _iter_qsa_attentions(self):
244 + """Yield MTP attention modules that own a QSA indexer."""
245 +
246 + for layer in self.layers:
247 + attention = getattr(layer, "self_attn", None)
248 + if (
249 + attention is not None
250 + and getattr(attention, "indexer", None) is not None
251 + ):
252 + yield attention
253 +
254 + def set_skip_topk(self, skip: bool) -> None:
255 + """Select on MTP step 0 and reuse its QSA indices on later steps."""
256 +
257 + for attention in self._iter_qsa_attentions():
258 + attention.indexer.skip_topk = skip
259 +
260 + def compact_topk_indices(self, row_indices: torch.Tensor) -> None:
261 + """Keep each request's target-aligned step-0 sparse-index row."""
262 +
263 + num_rows = row_indices.numel()
264 + for attention in self._iter_qsa_attentions():
265 + buffer = attention.topk_indices_buffer
266 + selected = buffer.index_select(0, row_indices)
267 + buffer[:num_rows].copy_(selected)
268 +
269 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
270 + return self.embed_tokens(input_ids)
271 +
272 + def forward(
273 + self,
274 + input_ids: torch.Tensor | None,
275 + positions: torch.Tensor,
276 + hidden_states: torch.Tensor | None = None,
277 + intermediate_tensors: IntermediateTensors | None = None,
278 + inputs_embeds: torch.Tensor | None = None,
279 + spec_step_idx: int = 0,
280 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
281 + hc_count = self.hc_count
282 + hidden_size = self.hidden_size
283 +
284 + if get_pp_group().is_first_rank:
285 + assert hidden_states is not None
286 + if inputs_embeds is None:
287 + assert input_ids is not None
288 + inputs_embeds = self.embed_input_ids(input_ids)
289 + # Embedding branch: pre-norm -> fc_embedding -> [T, H].
290 + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
291 + inputs_embeds = self.fc_embedding(inputs_embeds)
292 +
293 + # Backbone hidden is multi-stream [T, hc_count*H] (scheme A:
294 + # the main model truly emits the pre-final-mixer multi stream
295 + # on the first step; subsequent steps reuse the prior draft
296 + # step's multi stream).
297 + num_tokens = hidden_states.shape[0]
298 + hidden_states = hidden_states.view(num_tokens, hc_count, hidden_size)
299 + hidden_states = self.pre_fc_norm_hidden(hidden_states.flatten(-2)).view(
300 + num_tokens, hc_count, hidden_size
301 + )
302 + hidden_states = self.fc_hidden(hidden_states)
303 + # Add the embedding residual to every branch, then fold back
304 + # to [T, hc_count*H] (HC outer, HS inner) for the HC decoder.
305 + hidden_states = inputs_embeds.unsqueeze(-2) + hidden_states
306 + hidden_states = hidden_states.flatten(-2)
307 + else:
308 + assert intermediate_tensors is not None
309 + hidden_states = intermediate_tensors["hidden_states"]
310 +
311 + current_step_idx = spec_step_idx % self.num_mtp_layers
312 + layer = self.layers[current_step_idx]
313 + hidden_states, block_output, injection = layer(
314 + hidden_states=hidden_states,
315 + prev_block_output=None,
316 + prev_injection=None,
317 + positions=positions,
318 + input_ids=None,
319 + query_start_loc=None,
320 + ngram_context=None,
321 + )
322 + if not get_pp_group().is_last_rank:
323 + # As in the target model, PP carries a materialized tensor rather
324 + # than the delayed hidden/output/injection tuple.
325 + hidden_states = layer.mlp_hyper_connection.combine(
326 + hidden_states, block_output, injection
327 + )
328 + return IntermediateTensors({"hidden_states": hidden_states})
329 +
330 + # Last PP rank finalize. Keep both:
331 + # (A) sample_hidden_states [T, H] -> single stream for the LM head
332 + # (B) multi_hidden [T, hc_count*H] -> pre-final-mixer multi stream
333 + # for the next draft step (zero extra compute, just kept).
334 + multi_hidden, sample_hidden_states, _ = (
335 + self.hyper_connection_mixer.combine_and_mix(
336 + hidden_states, block_output, injection
337 + )
338 + )
339 + return sample_hidden_states, multi_hidden
340 +
341 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
342 + weights = maybe_fuse_shared_experts(
343 + weights,
344 + enabled=self.is_fused_shared_expert_enabled,
345 + n_routed_experts=getattr(self.config, "num_experts", 0) or 0,
346 + n_shared_experts=1,
347 + ckpt_prefix="mlp.shared_expert",
348 + )
349 + mapper = self.hf_to_vllm_mapper | WeightsMapper(
350 + orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
351 + )
352 + loader = AutoWeightsLoader(
353 + self,
354 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
355 + )
356 + return loader.load_weights(weights, mapper=mapper)
357 +
358 +
359 +@support_torch_compile(
360 + dynamic_arg_dims={
361 + "input_ids": 0,
362 + "positions": -1,
363 + "intermediate_tensors": 0,
364 + "inputs_embeds": 0,
365 + "hidden_states": 0,
366 + }
367 +)
368 +class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts):
369 + packed_modules_mapping = {
370 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
371 + "gate_up_proj": ["gate_proj", "up_proj"],
372 + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
373 + "in_proj_ba": ["in_proj_b", "in_proj_a"],
374 + "input_mix_weight_down_block_inject": [
375 + "input_mix_weight_down",
376 + "block_inject_weight",
377 + "_input_mix_padding",
378 + ],
379 + }
380 +
381 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
382 + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
383 + self.vllm_config = vllm_config
384 + cache_config = vllm_config.cache_config
385 + if cache_config.mamba_cache_mode == "all":
386 + raise NotImplementedError(
387 + "Qwen4ExpMTP currently does not support 'all' prefix caching, "
388 + "please use '--mamba-cache-mode=align' instead"
389 + )
390 +
391 + self.quant_config = vllm_config.quant_config
392 +
393 + super().__init__()
394 + self.config = config
395 + self.model = Qwen4ExpMultiTokenPredictor(
396 + vllm_config=vllm_config,
397 + prefix=maybe_prefix(prefix, "mtp"),
398 + )
399 +
400 + if get_pp_group().is_last_rank:
401 + if config.tie_word_embeddings:
402 + self.lm_head = self.model.embed_tokens
403 + else:
404 + self.lm_head = ParallelLMHead(
405 + config.vocab_size,
406 + config.hidden_size,
407 + prefix=maybe_prefix(prefix, "lm_head"),
408 + )
409 + else:
410 + self.lm_head = PPMissingLayer()
411 +
412 + self.logits_processor = LogitsProcessor(config.vocab_size)
413 + self.make_empty_intermediate_tensors = (
414 + self.model.make_empty_intermediate_tensors
415 + )
416 + self.set_moe_parameters(self.model.layers)
417 + enable_qwen4_exp_low_latency_gemm(self, vllm_config.model_config.dtype)
418 +
419 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
420 + return self.model.embed_input_ids(input_ids)
421 +
422 + def forward(
423 + self,
424 + input_ids: torch.Tensor | None,
425 + positions: torch.Tensor,
426 + hidden_states: torch.Tensor | None = None,
427 + intermediate_tensors: IntermediateTensors | None = None,
428 + inputs_embeds: torch.Tensor | None = None,
429 + spec_step_idx: int = 0,
430 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors:
431 + return self.model(
432 + input_ids,
433 + positions,
434 + hidden_states,
435 + intermediate_tensors,
436 + inputs_embeds,
437 + spec_step_idx=spec_step_idx,
438 + )
439 +
440 + def compute_logits(
441 + self, hidden_states: torch.Tensor, spec_step_idx: int = 0
442 + ) -> torch.Tensor | None:
443 + return self.logits_processor(self.lm_head, hidden_states)
444 +
445 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
446 + def remap_weight_names():
447 + for name, weight in weights:
448 + remapped_name = _remap_mtp_weight_name(name)
449 + if remapped_name is not None:
450 + yield remapped_name, weight
451 +
452 + mapper = WeightsMapper(
453 + orig_to_new_substr={"hyper_connection_mixer.block_inject_weight": None}
454 + )
455 + loader = AutoWeightsLoader(
456 + self,
457 + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(),
458 + )
459 + return loader.load_weights(remap_weight_names(), mapper=mapper)
460 +
461 +
462 +__all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"]
vllm/models/qwen4_exp/nvidia/ops/__init__.py new
+3
@@ -0,0 +1,3 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NVIDIA-only Qwen4Exp kernels; import leaf modules directly."""
vllm/models/qwen4_exp/nvidia/ops/hc.py new
+487
@@ -0,0 +1,487 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NVIDIA HyperConnection kernels for Qwen4Exp."""
4 +
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +from vllm.triton_utils import tl, triton
9 +from vllm.utils.torch_utils import direct_register_custom_op
10 +
11 +
12 +@triton.jit
13 +def _grouped_gemma_rmsnorm_kernel(
14 + x_ptr,
15 + w_ptr,
16 + y_ptr,
17 + stride_x,
18 + stride_y,
19 + DIM: tl.constexpr,
20 + NUM_GROUPS: tl.constexpr,
21 + W_SHARED: tl.constexpr,
22 + EPS: tl.constexpr,
23 + launch_pdl: tl.constexpr,
24 +) -> None:
25 + GROUP_DIM: tl.constexpr = DIM // NUM_GROUPS
26 + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(GROUP_DIM)
27 +
28 + pid = tl.program_id(0)
29 + group_id = pid % NUM_GROUPS
30 + row = pid // NUM_GROUPS
31 +
32 + offs_g = tl.arange(0, BLOCK_SIZE)
33 + offsets = group_id * GROUP_DIM + offs_g
34 + mask = offs_g < GROUP_DIM
35 + # A [GROUP_DIM] affine is shared; a [DIM] affine follows the grouped
36 + # checkpoint layout.
37 + w_offs = offs_g if W_SHARED else offsets
38 +
39 + if launch_pdl:
40 + tl.extra.cuda.gdc_wait()
41 +
42 + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0).to(tl.float32)
43 + w = tl.load(w_ptr + w_offs, mask, other=0.0)
44 +
45 + rrms = tl.rsqrt(tl.sum(x * x) / GROUP_DIM + EPS)
46 + # Gemma's (1 + w) affine is written this way to lower to an FMA.
47 + y = x * rrms
48 + y += y * w.to(tl.float32)
49 +
50 + if launch_pdl:
51 + tl.extra.cuda.gdc_launch_dependents()
52 + tl.store(y_ptr + row * stride_y + offsets, y, mask)
53 +
54 +
55 +def _grouped_gemma_rmsnorm(
56 + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
57 +) -> torch.Tensor:
58 + N, DIM = x.shape
59 + assert x.stride(1) == 1, "grouped Gemma RMSNorm requires unit inner stride"
60 + assert weight.is_contiguous(), "grouped Gemma RMSNorm weight must be contiguous"
61 + assert DIM % num_groups == 0
62 + group_dim = DIM // num_groups
63 + assert weight.numel() in (group_dim, DIM)
64 +
65 + y = x.new_empty(x.shape)
66 + _grouped_gemma_rmsnorm_kernel[(N * num_groups,)](
67 + x,
68 + weight,
69 + y,
70 + x.stride(0),
71 + y.stride(0),
72 + DIM,
73 + num_groups,
74 + W_SHARED=weight.numel() == group_dim,
75 + EPS=eps,
76 + launch_pdl=current_platform.is_arch_support_pdl(),
77 + )
78 + return y
79 +
80 +
81 +@triton.jit
82 +def _hc_silu_kernel(
83 + x_ptr,
84 + y_ptr,
85 + stride_x,
86 + stride_y,
87 + DIM: tl.constexpr,
88 + HC: tl.constexpr,
89 + launch_pdl: tl.constexpr,
90 +) -> None:
91 + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM)
92 +
93 + row = tl.program_id(0)
94 + offs = tl.arange(0, BLOCK_SIZE)
95 + mask = offs < DIM
96 +
97 + if launch_pdl:
98 + tl.extra.cuda.gdc_wait()
99 +
100 + x = tl.load(x_ptr + row * stride_x + offs, mask).to(tl.float32) / HC
101 + y = x * tl.sigmoid(x)
102 +
103 + if launch_pdl:
104 + tl.extra.cuda.gdc_launch_dependents()
105 + tl.store(y_ptr + row * stride_y + offs, y, mask)
106 +
107 +
108 +def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
109 + num_tokens, DIM = x.shape
110 + assert x.stride(1) == 1
111 +
112 + output = x.new_empty(x.shape)
113 + _hc_silu_kernel[(num_tokens,)](
114 + x,
115 + output,
116 + x.stride(0),
117 + output.stride(0),
118 + DIM=DIM,
119 + HC=hc_count,
120 + launch_pdl=current_platform.is_arch_support_pdl(),
121 + )
122 + return output
123 +
124 +
125 +@triton.jit
126 +def _hc_gate_mix_kernel(
127 + x_ptr,
128 + g_ptr,
129 + y_ptr,
130 + stride_x,
131 + stride_g,
132 + stride_y,
133 + DIM: tl.constexpr,
134 + HC: tl.constexpr,
135 + BLOCK_SIZE: tl.constexpr,
136 + launch_pdl: tl.constexpr,
137 +) -> None:
138 + HC_DIM: tl.constexpr = DIM // HC
139 +
140 + row = tl.program_id(0)
141 + tile_id = tl.program_id(1)
142 + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
143 + mask = offs_inner < HC_DIM
144 +
145 + if launch_pdl:
146 + tl.extra.cuda.gdc_wait()
147 +
148 + # The constexpr loop is unrolled and keeps one stream live at a time.
149 + # Materializing [HC, BLOCK_SIZE] more than doubles latency at large M.
150 + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
151 + for stream in tl.static_range(HC):
152 + offsets = stream * HC_DIM + offs_inner
153 + g = tl.load(g_ptr + row * stride_g + offsets, mask, other=0.0)
154 + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0)
155 + acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32)
156 + acc /= HC
157 +
158 + if launch_pdl:
159 + tl.extra.cuda.gdc_launch_dependents()
160 + tl.store(y_ptr + row * stride_y + offs_inner, acc, mask)
161 +
162 +
163 +def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
164 + N, DIM = gate.shape
165 + assert x.shape == gate.shape
166 + assert DIM % hc_count == 0
167 + assert x.stride(1) == 1
168 + assert gate.stride(1) == 1
169 +
170 + HC_DIM = DIM // hc_count
171 + out = x.new_empty(N, HC_DIM)
172 + BLOCK_SIZE = 512
173 + _hc_gate_mix_kernel[(N, triton.cdiv(HC_DIM, BLOCK_SIZE))](
174 + x,
175 + gate,
176 + out,
177 + x.stride(0),
178 + gate.stride(0),
179 + out.stride(0),
180 + DIM,
181 + hc_count,
182 + BLOCK_SIZE,
183 + launch_pdl=current_platform.is_arch_support_pdl(),
184 + )
185 + return out
186 +
187 +
188 +@triton.jit
189 +def _hc_combine_kernel(
190 + block_ptr,
191 + res_ptr,
192 + inj_ptr,
193 + out_ptr,
194 + stride_block,
195 + stride_res,
196 + stride_inj,
197 + stride_out,
198 + HC_DIM: tl.constexpr,
199 + HC: tl.constexpr,
200 + BLOCK_SIZE: tl.constexpr,
201 + launch_pdl: tl.constexpr,
202 +) -> None:
203 + HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
204 +
205 + row = tl.program_id(0)
206 + tile_id = tl.program_id(1)
207 +
208 + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
209 + mask_inner = offs_inner < HC_DIM
210 + offs_hc = tl.arange(0, HC_PAD)
211 + mask_hc = offs_hc < HC
212 + offs = offs_hc[:, None] * HC_DIM + offs_inner[None, :]
213 + mask = mask_hc[:, None] & mask_inner[None, :]
214 +
215 + if launch_pdl:
216 + tl.extra.cuda.gdc_wait()
217 +
218 + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)
219 + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)
220 + res = tl.load(res_ptr + row * stride_res + offs, mask, other=0.0)
221 +
222 + # Keeping HC as a broadcast dimension is faster here than four separate
223 + # residual load/store sequences.
224 + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)
225 + out = res.to(tl.float32) + block.to(tl.float32)[None, :] * inj[:, None]
226 +
227 + if launch_pdl:
228 + tl.extra.cuda.gdc_launch_dependents()
229 + tl.store(out_ptr + row * stride_out + offs, out, mask=mask)
230 +
231 +
232 +def _hc_combine(
233 + residual: torch.Tensor,
234 + block_output: torch.Tensor,
235 + injection_logits: torch.Tensor,
236 + hc_count: int,
237 +) -> torch.Tensor:
238 + N, DIM = residual.shape
239 + assert DIM % hc_count == 0
240 + hc_dim = DIM // hc_count
241 + assert block_output.shape == (N, hc_dim)
242 + assert injection_logits.shape == (N, hc_count)
243 + assert residual.stride(1) == 1
244 + assert block_output.stride(1) == 1
245 + assert injection_logits.stride(1) == 1
246 +
247 + out = residual.new_empty(residual.shape)
248 + BLOCK_SIZE = 512
249 + _hc_combine_kernel[(N, triton.cdiv(hc_dim, BLOCK_SIZE))](
250 + block_output,
251 + residual,
252 + injection_logits,
253 + out,
254 + block_output.stride(0),
255 + residual.stride(0),
256 + injection_logits.stride(0),
257 + out.stride(0),
258 + hc_dim,
259 + hc_count,
260 + BLOCK_SIZE,
261 + launch_pdl=current_platform.is_arch_support_pdl(),
262 + )
263 + return out
264 +
265 +
266 +@triton.jit
267 +def _hc_combine_norm_kernel(
268 + block_ptr,
269 + res_ptr,
270 + inj_ptr,
271 + w_ptr,
272 + out_ptr,
273 + y_ptr,
274 + stride_block,
275 + stride_res,
276 + stride_inj,
277 + stride_out,
278 + stride_y,
279 + HC_DIM: tl.constexpr,
280 + HC: tl.constexpr,
281 + W_SHARED: tl.constexpr,
282 + EPS: tl.constexpr,
283 + BLOCK_SIZE: tl.constexpr,
284 + launch_pdl: tl.constexpr,
285 +) -> None:
286 + HC_PAD: tl.constexpr = triton.next_power_of_2(HC)
287 + NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE)
288 + NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES)
289 +
290 + row = tl.program_id(0)
291 + stream = tl.program_id(1)
292 + offs_hc = tl.arange(0, HC_PAD)
293 + mask_hc = offs_hc < HC
294 + tile_ids = tl.arange(0, NUM_TILES_PAD)
295 + offs_inner = tile_ids[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :]
296 + mask_inner = offs_inner < HC_DIM
297 + offs = stream * HC_DIM + offs_inner
298 + # Shared norm weights repeat across streams; per-branch weights use the
299 + # same flattened HC layout as the residual.
300 + w_offs = offs_inner if W_SHARED else offs
301 +
302 + if launch_pdl:
303 + tl.extra.cuda.gdc_wait()
304 +
305 + # Start the uncached residual load first, then issue the other combine
306 + # loads before consuming any of them.
307 + res = tl.load(res_ptr + row * stride_res + offs, mask_inner, other=0.0)
308 + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0)
309 + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0)
310 + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC)
311 + inj = tl.sum(tl.where(offs_hc == stream, inj, 0.0))
312 + # Round the materialized combine result before normalization. This matches
313 + # the unfused combine -> RMSNorm boundary.
314 + out = (res.to(tl.float32) + block.to(tl.float32) * inj).to(out_ptr.dtype.element_ty)
315 + tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner)
316 +
317 + out = out.to(tl.float32)
318 + # Keep the two-axis reduction: flattening the padded tile is ~40% slower
319 + # at decode sizes.
320 + sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0)
321 + rrms = tl.rsqrt(sum_sq / HC_DIM + EPS)
322 +
323 + if launch_pdl:
324 + tl.extra.cuda.gdc_launch_dependents()
325 +
326 + # Loading the weight earlier helps decode but keeps the tile live across
327 + # the reduction and regresses larger batches, so defer it to the norm.
328 + w = tl.load(w_ptr + w_offs, mask_inner, other=0.0)
329 + y = out * rrms
330 + y += y * w.to(tl.float32)
331 + tl.store(y_ptr + row * stride_y + offs, y, mask_inner)
332 +
333 +
334 +def _hc_combine_norm(
335 + residual: torch.Tensor,
336 + block_output: torch.Tensor,
337 + injection_logits: torch.Tensor,
338 + norm_weight: torch.Tensor,
339 + eps: float,
340 + hc_count: int,
341 +) -> tuple[torch.Tensor, torch.Tensor]:
342 + N, DIM = residual.shape
343 + assert DIM % hc_count == 0
344 + hc_dim = DIM // hc_count
345 + assert block_output.shape == (N, hc_dim)
346 + assert injection_logits.shape == (N, hc_count)
347 + assert residual.stride(1) == 1
348 + assert block_output.stride(1) == 1
349 + assert injection_logits.stride(1) == 1
350 + assert norm_weight.is_contiguous()
351 + assert norm_weight.numel() in (hc_dim, DIM)
352 +
353 + out = residual.new_empty(residual.shape)
354 + y = residual.new_empty(residual.shape)
355 + BLOCK_SIZE = 512
356 + _hc_combine_norm_kernel[(N, hc_count)](
357 + block_output,
358 + residual,
359 + injection_logits,
360 + norm_weight,
361 + out,
362 + y,
363 + block_output.stride(0),
364 + residual.stride(0),
365 + injection_logits.stride(0),
366 + out.stride(0),
367 + y.stride(0),
368 + hc_dim,
369 + hc_count,
370 + W_SHARED=norm_weight.numel() == hc_dim,
371 + EPS=eps,
372 + BLOCK_SIZE=BLOCK_SIZE,
373 + launch_pdl=current_platform.is_arch_support_pdl(),
374 + )
375 + return out, y
376 +
377 +
378 +def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor:
379 + return x.new_empty(x.shape)
380 +
381 +
382 +def _hc_gate_mix_fake(
383 + x: torch.Tensor, gate: torch.Tensor, hc_count: int
384 +) -> torch.Tensor:
385 + del gate
386 + return x.new_empty((x.shape[0], x.shape[1] // hc_count))
387 +
388 +
389 +def _hc_combine_fake(
390 + residual: torch.Tensor,
391 + block_output: torch.Tensor,
392 + injection_logits: torch.Tensor,
393 + hc_count: int,
394 +) -> torch.Tensor:
395 + del block_output, injection_logits, hc_count
396 + return residual.new_empty(residual.shape)
397 +
398 +
399 +def _hc_combine_norm_fake(
400 + residual: torch.Tensor,
401 + block_output: torch.Tensor,
402 + injection_logits: torch.Tensor,
403 + norm_weight: torch.Tensor,
404 + eps: float,
405 + hc_count: int,
406 +) -> tuple[torch.Tensor, torch.Tensor]:
407 + del block_output, injection_logits, norm_weight, eps, hc_count
408 + return residual.new_empty(residual.shape), residual.new_empty(residual.shape)
409 +
410 +
411 +direct_register_custom_op(
412 + op_name="qwen4_exp_grouped_gemma_rmsnorm",
413 + op_func=_grouped_gemma_rmsnorm,
414 + fake_impl=_same_shape_fake,
415 +)
416 +direct_register_custom_op(
417 + op_name="qwen4_exp_hc_silu",
418 + op_func=_hc_silu,
419 + fake_impl=_same_shape_fake,
420 +)
421 +direct_register_custom_op(
422 + op_name="qwen4_exp_hc_gate_mix",
423 + op_func=_hc_gate_mix,
424 + fake_impl=_hc_gate_mix_fake,
425 +)
426 +direct_register_custom_op(
427 + op_name="qwen4_exp_hc_combine",
428 + op_func=_hc_combine,
429 + fake_impl=_hc_combine_fake,
430 +)
431 +direct_register_custom_op(
432 + op_name="qwen4_exp_hc_combine_norm",
433 + op_func=_hc_combine_norm,
434 + fake_impl=_hc_combine_norm_fake,
435 +)
436 +
437 +
438 +def grouped_gemma_rmsnorm(
439 + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int
440 +) -> torch.Tensor:
441 + return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups)
442 +
443 +
444 +def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor:
445 + return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count)
446 +
447 +
448 +def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor:
449 + return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count)
450 +
451 +
452 +def hc_combine(
453 + residual: torch.Tensor,
454 + block_output: torch.Tensor,
455 + injection_logits: torch.Tensor,
456 + hc_count: int,
457 +) -> torch.Tensor:
458 + return torch.ops.vllm.qwen4_exp_hc_combine(
459 + residual, block_output, injection_logits, hc_count
460 + )
461 +
462 +
463 +def hc_combine_norm(
464 + residual: torch.Tensor,
465 + block_output: torch.Tensor,
466 + injection_logits: torch.Tensor,
467 + norm_weight: torch.Tensor,
468 + eps: float,
469 + hc_count: int,
470 +) -> tuple[torch.Tensor, torch.Tensor]:
471 + return torch.ops.vllm.qwen4_exp_hc_combine_norm(
472 + residual,
473 + block_output,
474 + injection_logits,
475 + norm_weight,
476 + eps,
477 + hc_count,
478 + )
479 +
480 +
481 +__all__ = [
482 + "grouped_gemma_rmsnorm",
483 + "hc_combine",
484 + "hc_combine_norm",
485 + "hc_gate_mix",
486 + "hc_silu",
487 +]
vllm/models/qwen4_exp/nvidia/ops/qsa.py new
+1115
@@ -0,0 +1,1115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Triton kernels for the Qwen4Exp weight-free QSA path."""
4 +
5 +from __future__ import annotations
6 +
7 +import math
8 +
9 +import torch
10 +
11 +from vllm.platforms import current_platform
12 +from vllm.triton_utils import HAS_TRITON, tl, triton
13 +
14 +_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024
15 +_TOPK_WORKSPACE_BYTES = 1024 * 1024
16 +
17 +
18 +@triton.jit
19 +def _qsa_mqa_paged_kernel(
20 + q_ptr,
21 + k_cache_ptr,
22 + page_table_ptr,
23 + token_to_req_ptr,
24 + query_positions_ptr,
25 + sequence_lengths_ptr,
26 + visible_blocks_ptr,
27 + logits_ptr,
28 + stride_q_row,
29 + stride_q_head,
30 + stride_q_dim,
31 + stride_cache_block,
32 + stride_cache_token,
33 + stride_cache_dim,
34 + stride_table_req,
35 + stride_table_page,
36 + stride_logits_row,
37 + num_rows,
38 + num_columns,
39 + num_pages,
40 + num_requests,
41 + score_divisor,
42 + PAGE_SIZE: tl.constexpr,
43 + PAGE_TABLE_WIDTH: tl.constexpr,
44 + NUM_HEADS: tl.constexpr,
45 + HEAD_DIM: tl.constexpr,
46 + BLOCK_N: tl.constexpr,
47 + BLOCK_D: tl.constexpr,
48 + TILES_PER_PROG: tl.constexpr,
49 + STAGES: tl.constexpr,
50 + MAX_N: tl.constexpr,
51 + COMPRESS_RATIO: tl.constexpr,
52 +) -> None:
53 + row = tl.program_id(0)
54 + dims = tl.arange(0, BLOCK_D)
55 + heads = tl.arange(0, MAX_N)
56 + request = tl.load(token_to_req_ptr + row)
57 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
58 + query_position = tl.load(query_positions_ptr + row)
59 + sequence_length = tl.load(
60 + sequence_lengths_ptr + safe_request,
61 + mask=(request >= 0) & (request < num_requests),
62 + other=0,
63 + )
64 + visible = tl.minimum(
65 + (query_position + 1) // COMPRESS_RATIO,
66 + sequence_length // COMPRESS_RATIO,
67 + )
68 + if tl.program_id(1) == 0:
69 + tl.store(visible_blocks_ptr + row, visible)
70 + tile_start = tl.program_id(1) * TILES_PER_PROG
71 + # Top-k is bounded by visible_blocks, so columns beyond it need no value.
72 + if tile_start * BLOCK_N >= visible:
73 + return
74 + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible, BLOCK_N))
75 + tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N))
76 +
77 + # Pad the small head axis to a tensor-core-compatible N dimension.
78 + query = tl.load(
79 + q_ptr
80 + + row * stride_q_row
81 + + heads[None, :] * stride_q_head
82 + + dims[:, None] * stride_q_dim,
83 + mask=(heads[None, :] < NUM_HEADS) & (dims[:, None] < HEAD_DIM),
84 + other=0.0,
85 + )
86 + column_offsets = tl.arange(0, BLOCK_N)
87 + for tile in tl.range(tile_start, tile_end, num_stages=STAGES):
88 + columns = tile * BLOCK_N + column_offsets
89 + live = columns < visible
90 + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1)
91 + page_offset = columns % PAGE_SIZE
92 + physical_page = tl.load(
93 + page_table_ptr
94 + + safe_request * stride_table_req
95 + + logical_page * stride_table_page,
96 + mask=live,
97 + other=-1,
98 + )
99 + page_valid = live & (physical_page >= 0) & (physical_page < num_pages)
100 + # physical_page * block stride can overflow int32 for large caches.
101 + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64)
102 + keys = tl.load(
103 + k_cache_ptr
104 + + safe_physical_page[:, None] * stride_cache_block
105 + + page_offset[:, None] * stride_cache_token
106 + + dims[None, :] * stride_cache_dim,
107 + mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM),
108 + other=0.0,
109 + eviction_policy="evict_first",
110 + )
111 + scores = tl.dot(keys, query, out_dtype=tl.float32)
112 + scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(scores, 0.0), 0.0)
113 + score = tl.sum(scores, axis=1) / score_divisor
114 + tl.store(
115 + logits_ptr + row * stride_logits_row + columns,
116 + tl.where(page_valid, score, -float("inf")),
117 + mask=live & (columns < num_columns),
118 + )
119 +
120 +
121 +@triton.jit
122 +def _expand_qsa_indices_kernel(
123 + block_indices_ptr,
124 + query_positions_ptr,
125 + sequence_lengths_ptr,
126 + token_to_req_ptr,
127 + output_ptr,
128 + stride_blocks_row,
129 + stride_blocks_column,
130 + stride_output_row,
131 + stride_output_column,
132 + rows,
133 + num_requests,
134 + BLOCK_TOPK: tl.constexpr,
135 + COMPRESS_RATIO: tl.constexpr,
136 + TOKEN_TOPK: tl.constexpr,
137 + OUTPUT_WIDTH: tl.constexpr,
138 + COLUMN_BLOCK: tl.constexpr,
139 +) -> None:
140 + row = tl.program_id(0)
141 + columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK)
142 + query_position = tl.load(query_positions_ptr + row)
143 + request = tl.load(token_to_req_ptr + row)
144 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
145 + sequence_length = tl.load(
146 + sequence_lengths_ptr + safe_request,
147 + mask=(request >= 0) & (request < num_requests),
148 + other=0,
149 + )
150 + complete_blocks = tl.minimum(
151 + tl.minimum(
152 + (query_position + 1) // COMPRESS_RATIO,
153 + sequence_length // COMPRESS_RATIO,
154 + ),
155 + BLOCK_TOPK,
156 + )
157 + expanded_count = complete_blocks * COMPRESS_RATIO
158 + tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO
159 + tail_count = (query_position + 1) - tail_start
160 +
161 + is_expanded = columns < expanded_count
162 + block_rank = columns // COMPRESS_RATIO
163 + offset = columns % COMPRESS_RATIO
164 + safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1)
165 + block = tl.load(
166 + block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column,
167 + mask=(row < rows) & is_expanded,
168 + other=-1,
169 + )
170 + expanded = block * COMPRESS_RATIO + offset
171 + tail_offset = columns - expanded_count
172 + is_tail = (
173 + (columns >= expanded_count)
174 + & (tail_offset < tail_count)
175 + & (tail_offset < COMPRESS_RATIO - 1)
176 + )
177 + token = tl.where(is_expanded, expanded, tail_start + tail_offset)
178 + valid = (
179 + (row < rows)
180 + & (columns < OUTPUT_WIDTH)
181 + & (is_expanded | is_tail)
182 + & (token >= 0)
183 + & (token < sequence_length)
184 + )
185 + tl.store(
186 + output_ptr + row * stride_output_row + columns * stride_output_column,
187 + tl.where(valid, token, -1),
188 + mask=(row < rows) & (columns < OUTPUT_WIDTH),
189 + )
190 +
191 +
192 +@triton.jit
193 +def _qsa_sparse_paged_gqa_splitk_kernel(
194 + q_ptr,
195 + k_cache_ptr,
196 + v_cache_ptr,
197 + indices_ptr,
198 + block_table_ptr,
199 + token_to_req_ptr,
200 + partial_output_ptr,
201 + partial_lse_ptr,
202 + output_ptr,
203 + stride_q_row,
204 + stride_q_head,
205 + stride_k_block,
206 + stride_k_token,
207 + stride_k_head,
208 + stride_v_block,
209 + stride_v_token,
210 + stride_v_head,
211 + stride_indices_row,
212 + stride_table_req,
213 + stride_output_row,
214 + stride_output_head,
215 + num_rows,
216 + num_cache_blocks,
217 + num_requests,
218 + TOPK: tl.constexpr,
219 + PAGE_SIZE: tl.constexpr,
220 + PAGE_TABLE_WIDTH: tl.constexpr,
221 + GROUP_SIZE: tl.constexpr,
222 + HEAD_DIM: tl.constexpr,
223 + NUM_QUERY_HEADS: tl.constexpr,
224 + NUM_SPLITS: tl.constexpr,
225 + NUM_TILES: tl.constexpr,
226 + BLOCK_M: tl.constexpr,
227 + BLOCK_N: tl.constexpr,
228 +) -> None:
229 + row = tl.program_id(0)
230 + kv_head = tl.program_id(1)
231 + split_id = tl.program_id(2)
232 + request = tl.load(token_to_req_ptr + row)
233 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
234 +
235 + head_offsets = tl.arange(0, BLOCK_M)
236 + dim_offsets = tl.arange(0, HEAD_DIM)
237 + column_offsets = tl.arange(0, BLOCK_N)
238 + first_head = kv_head * GROUP_SIZE
239 + query = tl.load(
240 + q_ptr
241 + + row * stride_q_row
242 + + (first_head + head_offsets[:, None]) * stride_q_head
243 + + dim_offsets[None, :],
244 + mask=head_offsets[:, None] < GROUP_SIZE,
245 + other=0.0,
246 + )
247 +
248 + max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32)
249 + normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32)
250 + accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)
251 + softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634
252 +
253 + # Dynamic bounds avoid padded main-loop iterations for uneven splits.
254 + split_tile_start = split_id * NUM_TILES // NUM_SPLITS
255 + split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS
256 + for tile in range(split_tile_start, split_tile_end):
257 + columns = tile * BLOCK_N + column_offsets
258 + logical_token = tl.load(
259 + indices_ptr + row * stride_indices_row + columns,
260 + mask=columns < TOPK,
261 + other=-1,
262 + )
263 + safe_token = tl.maximum(logical_token, 0)
264 + logical_page = safe_token // PAGE_SIZE
265 + page_offset = safe_token % PAGE_SIZE
266 + valid = (
267 + (request >= 0)
268 + & (request < num_requests)
269 + & (logical_token >= 0)
270 + & (logical_page < PAGE_TABLE_WIDTH)
271 + )
272 + physical_page = tl.load(
273 + block_table_ptr
274 + + safe_request * stride_table_req
275 + + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1),
276 + mask=valid,
277 + other=-1,
278 + )
279 + valid &= (physical_page >= 0) & (physical_page < num_cache_blocks)
280 + # physical_page * block stride can overflow int32 for large caches.
281 + safe_page = tl.maximum(physical_page, 0).to(tl.int64)
282 + keys = tl.load(
283 + k_cache_ptr
284 + + safe_page[None, :] * stride_k_block
285 + + page_offset[None, :] * stride_k_token
286 + + kv_head * stride_k_head
287 + + dim_offsets[:, None],
288 + mask=valid[None, :],
289 + other=0.0,
290 + )
291 + values = tl.load(
292 + v_cache_ptr
293 + + safe_page[:, None] * stride_v_block
294 + + page_offset[:, None] * stride_v_token
295 + + kv_head * stride_v_head
296 + + dim_offsets[None, :],
297 + mask=valid[:, None],
298 + other=0.0,
299 + )
300 + scores = tl.dot(query, keys)
301 + # Scaling scores avoids re-quantizing a scaled query to BF16.
302 + scores *= softmax_scale_log2
303 + scores = tl.where(valid[None, :], scores, -1.0e20)
304 + next_max = tl.maximum(max_value, tl.max(scores, axis=1))
305 + alpha = tl.math.exp2(max_value - next_max)
306 + probabilities = tl.where(
307 + valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0
308 + )
309 + accumulator = tl.dot(
310 + probabilities.to(values.dtype),
311 + values,
312 + acc=accumulator * alpha[:, None],
313 + )
314 + normalizer = normalizer * alpha + tl.sum(probabilities, axis=1)
315 + max_value = next_max
316 +
317 + has_values = normalizer > 0
318 + normalized_output = tl.where(
319 + has_values[:, None],
320 + accumulator / tl.maximum(normalizer[:, None], 1.0e-20),
321 + 0.0,
322 + )
323 + output_mask = head_offsets[:, None] < GROUP_SIZE
324 + if NUM_SPLITS == 1:
325 + tl.store(
326 + output_ptr
327 + + row * stride_output_row
328 + + (first_head + head_offsets[:, None]) * stride_output_head
329 + + dim_offsets[None, :],
330 + normalized_output,
331 + mask=output_mask,
332 + )
333 + else:
334 + partial_lse = tl.where(
335 + has_values,
336 + max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)),
337 + -float("inf"),
338 + )
339 + tl.store(
340 + partial_output_ptr
341 + + (
342 + (split_id * num_rows + row) * NUM_QUERY_HEADS
343 + + first_head
344 + + head_offsets[:, None]
345 + )
346 + * HEAD_DIM
347 + + dim_offsets[None, :],
348 + normalized_output,
349 + mask=output_mask,
350 + )
351 + tl.store(
352 + partial_lse_ptr
353 + + (split_id * num_rows + row) * NUM_QUERY_HEADS
354 + + first_head
355 + + head_offsets,
356 + partial_lse,
357 + mask=head_offsets < GROUP_SIZE,
358 + )
359 +
360 +
361 +@triton.jit
362 +def _qsa_merge_splitk_kernel(
363 + partial_output_ptr,
364 + partial_lse_ptr,
365 + output_ptr,
366 + stride_output_row,
367 + stride_output_head,
368 + num_rows,
369 + HEAD_DIM: tl.constexpr,
370 + NUM_QUERY_HEADS: tl.constexpr,
371 + NUM_SPLITS: tl.constexpr,
372 + BLOCK_SPLITS: tl.constexpr,
373 +) -> None:
374 + row = tl.program_id(0)
375 + head = tl.program_id(1)
376 + split_offsets = tl.arange(0, BLOCK_SPLITS)
377 + dim_offsets = tl.arange(0, HEAD_DIM)
378 + split_mask = split_offsets < NUM_SPLITS
379 + lse = tl.load(
380 + partial_lse_ptr + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head,
381 + mask=split_mask,
382 + other=-float("inf"),
383 + )
384 + lse_max = tl.max(lse, axis=0)
385 + has_values = lse_max > -float("inf")
386 + shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf"))
387 + weights = tl.math.exp2(shifted)
388 + denominator = tl.sum(weights, axis=0)
389 + partial_output = tl.load(
390 + partial_output_ptr
391 + + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head)
392 + * HEAD_DIM
393 + + dim_offsets[None, :],
394 + mask=split_mask[:, None],
395 + other=0.0,
396 + )
397 + merged = tl.sum(partial_output * weights[:, None], axis=0)
398 + merged = tl.where(denominator > 0, merged / denominator, 0.0)
399 + tl.store(
400 + output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets,
401 + merged,
402 + )
403 +
404 +
405 +@triton.jit
406 +def _store_qsa_rows_kernel(
407 + cache_ptr,
408 + slots_ptr,
409 + rows_ptr,
410 + stride_cache_block,
411 + stride_cache_token,
412 + stride_cache_dim,
413 + stride_rows_row,
414 + stride_rows_dim,
415 + num_rows,
416 + num_blocks,
417 + PAGE_SIZE: tl.constexpr,
418 + WIDTH: tl.constexpr,
419 + BLOCK_D: tl.constexpr,
420 +) -> None:
421 + row = tl.program_id(0)
422 + dims = tl.arange(0, BLOCK_D)
423 + slot = tl.load(slots_ptr + row)
424 + valid = (row < num_rows) & (slot >= 0) & (slot < num_blocks * PAGE_SIZE)
425 + block = tl.maximum(slot, 0) // PAGE_SIZE
426 + token = tl.maximum(slot, 0) % PAGE_SIZE
427 + values = tl.load(
428 + rows_ptr + row * stride_rows_row + dims * stride_rows_dim,
429 + mask=valid & (dims < WIDTH),
430 + other=0,
431 + )
432 + tl.store(
433 + cache_ptr
434 + + block * stride_cache_block
435 + + token * stride_cache_token
436 + + dims * stride_cache_dim,
437 + values,
438 + mask=valid & (dims < WIDTH),
439 + )
440 +
441 +
442 +@triton.jit
443 +def _compress_qsa_groups_kernel(
444 + raw_keys_ptr, # this step's raw key rows, straight from activations
445 + raw_positions_ptr, # this step's per-token positions
446 + compressor_state_cache_ptr, # per-request ring of previous raw keys
447 + rope_cache_ptr, # packed RoPE position tail of the ring
448 + compressor_state_table_ptr,
449 + token_to_req_ptr,
450 + query_start_loc_ptr,
451 + logical_positions_ptr,
452 + compressed_slots_ptr,
453 + pooled_ptr,
454 + first_positions_ptr,
455 + stride_raw_row,
456 + stride_raw_dim,
457 + stride_raw_positions_row,
458 + stride_raw_positions_dim,
459 + stride_compressor_state_block,
460 + stride_compressor_state_token,
461 + stride_compressor_state_dim,
462 + stride_rope_block,
463 + stride_rope_token,
464 + stride_rope_dim,
465 + stride_compressor_state_table_req,
466 + stride_pooled_row,
467 + stride_pooled_dim,
468 + stride_positions_row,
469 + stride_positions_dim,
470 + num_rows,
471 + num_compressor_state_blocks,
472 + num_requests,
473 + COMPRESSOR_STATE_SIZE: tl.constexpr,
474 + COMPRESS_RATIO: tl.constexpr,
475 + HEAD_DIM: tl.constexpr,
476 + BLOCK_D: tl.constexpr,
477 + LOAD_ROPE_POSITIONS: tl.constexpr,
478 +) -> None:
479 + row = tl.program_id(0)
480 + dims = tl.arange(0, BLOCK_D)
481 + request = tl.load(token_to_req_ptr + row)
482 + end_position = tl.load(logical_positions_ptr + row)
483 + compressed_slot = tl.load(compressed_slots_ptr + row)
484 + valid_request = (request >= 0) & (request < num_requests)
485 + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)
486 + query_row_start = tl.load(
487 + query_start_loc_ptr + safe_request, mask=valid_request, other=0
488 + )
489 + query_row_end = tl.load(
490 + query_start_loc_ptr + safe_request + 1, mask=valid_request, other=0
491 + )
492 + chunk_start_position = end_position - (row - query_row_start)
493 + compressor_state_block = tl.load(
494 + compressor_state_table_ptr + safe_request * stride_compressor_state_table_req,
495 + mask=valid_request,
496 + other=-1,
497 + )
498 + valid_compressor_state_block = (compressor_state_block >= 0) & (
499 + compressor_state_block < num_compressor_state_blocks
500 + )
501 + valid_row = (
502 + (row < num_rows)
503 + & valid_request
504 + & (row >= query_row_start)
505 + & (row < query_row_end)
506 + & (end_position >= COMPRESS_RATIO - 1)
507 + & (compressed_slot >= 0)
508 + )
509 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32)
510 +
511 + # A group can span the compressor-state ring (older members) and this
512 + # step's raw rows (members at positions >= chunk_start_position).
513 + for group_offset in tl.range(0, COMPRESS_RATIO):
514 + position = end_position - (COMPRESS_RATIO - 1 - group_offset)
515 + use_raw = position >= chunk_start_position
516 + raw_row = query_row_start + position - chunk_start_position
517 + raw_values = tl.load(
518 + raw_keys_ptr + raw_row * stride_raw_row + dims * stride_raw_dim,
519 + mask=valid_row
520 + & use_raw
521 + & (raw_row >= query_row_start)
522 + & (raw_row < query_row_end)
523 + & (raw_row < num_rows)
524 + & (dims < HEAD_DIM),
525 + other=0.0,
526 + ).to(tl.float32)
527 + compressor_state_values = tl.load(
528 + compressor_state_cache_ptr
529 + + tl.maximum(compressor_state_block, 0).to(tl.int64)
530 + * stride_compressor_state_block
531 + + (position % COMPRESSOR_STATE_SIZE) * stride_compressor_state_token
532 + + dims * stride_compressor_state_dim,
533 + mask=valid_row
534 + & ~use_raw
535 + & valid_compressor_state_block
536 + & (dims < HEAD_DIM),
537 + other=0.0,
538 + ).to(tl.float32)
539 + accumulator += tl.where(use_raw, raw_values, compressor_state_values)
540 +
541 + tl.store(
542 + pooled_ptr + row * stride_pooled_row + dims * stride_pooled_dim,
543 + accumulator / COMPRESS_RATIO,
544 + mask=(row < num_rows) & (dims < HEAD_DIM),
545 + )
546 +
547 + position_dims = tl.arange(0, 4)
548 + first_position = end_position - COMPRESS_RATIO + 1
549 + if LOAD_ROPE_POSITIONS:
550 + first_from_raw = first_position >= chunk_start_position
551 + raw_first_row = query_row_start + first_position - chunk_start_position
552 + raw_position_values = tl.load(
553 + raw_positions_ptr
554 + + raw_first_row * stride_raw_positions_row
555 + + position_dims * stride_raw_positions_dim,
556 + mask=valid_row
557 + & first_from_raw
558 + & (raw_first_row >= query_row_start)
559 + & (raw_first_row < query_row_end)
560 + & (raw_first_row < num_rows)
561 + & (position_dims < 3),
562 + other=0,
563 + )
564 + compressor_state_position_values = tl.load(
565 + rope_cache_ptr
566 + + tl.maximum(compressor_state_block, 0).to(tl.int64) * stride_rope_block
567 + + (first_position % COMPRESSOR_STATE_SIZE) * stride_rope_token
568 + + position_dims * stride_rope_dim,
569 + mask=valid_row
570 + & ~first_from_raw
571 + & valid_compressor_state_block
572 + & (position_dims < 3),
573 + other=0,
574 + )
575 + position_values = tl.where(
576 + first_from_raw,
577 + raw_position_values,
578 + compressor_state_position_values,
579 + )
580 + else:
581 + position_values = tl.where(valid_row, first_position, 0)
582 + tl.store(
583 + first_positions_ptr
584 + + row * stride_positions_row
585 + + position_dims * stride_positions_dim,
586 + position_values,
587 + mask=(row < num_rows) & (position_dims < 3),
588 + )
589 +
590 +
591 +def _validate_mqa(q: torch.Tensor) -> None:
592 + if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0:
593 + raise ValueError("QSA query must be [rows, heads, head_dim]")
594 +
595 +
596 +def qsa_mqa_paged(
597 + q: torch.Tensor,
598 + k_cache: torch.Tensor,
599 + page_table: torch.Tensor,
600 + token_to_req: torch.Tensor,
601 + query_positions: torch.Tensor,
602 + sequence_lengths: torch.Tensor,
603 + compress_ratio: int,
604 + num_columns: int | None = None,
605 + score_scale: float | None = None,
606 +) -> tuple[torch.Tensor, torch.Tensor]:
607 + """Compute QSA scores directly from a paged compressed-key cache."""
608 +
609 + _validate_mqa(q)
610 + if not q.is_cuda or not HAS_TRITON:
611 + raise RuntimeError("paged QSA scoring requires CUDA and Triton")
612 + if k_cache.ndim != 4 or k_cache.shape[2] != 1:
613 + raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]")
614 + if k_cache.shape[3] != q.shape[2]:
615 + raise ValueError("QSA query and cache dimensions must match")
616 + if page_table.ndim != 2:
617 + raise ValueError("QSA page table must be two-dimensional")
618 + if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)):
619 + raise ValueError("QSA paged scoring cache and page table must be nonempty")
620 + if token_to_req.shape != (q.shape[0],):
621 + raise ValueError("QSA request mapping must match query rows")
622 + if query_positions.shape != (q.shape[0],):
623 + raise ValueError("QSA query positions must match query rows")
624 + if sequence_lengths.shape != (page_table.shape[0],):
625 + raise ValueError("QSA sequence lengths must match page-table requests")
626 + if compress_ratio <= 0:
627 + raise ValueError("QSA compression ratio must be positive")
628 + score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale
629 + if score_divisor <= 0:
630 + raise ValueError("QSA score scale must be positive")
631 +
632 + capacity = page_table.shape[1] * k_cache.shape[1]
633 + columns = capacity if num_columns is None else num_columns
634 + if columns < 0:
635 + raise ValueError("QSA score width must be non-negative")
636 + logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device)
637 + visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device)
638 + if not q.shape[0] or not columns:
639 + return logits, visible_blocks
640 + BLOCK_N = 64
641 + BLOCK_D = max(16, triton.next_power_of_2(q.shape[2]))
642 + MAX_N = max(16, triton.next_power_of_2(q.shape[1]))
643 + # Tuned on GB300: larger row batches provide enough parallelism to reuse Q.
644 + tiles_per_program = 1 if q.shape[0] <= 32 else 8
645 + _qsa_mqa_paged_kernel[
646 + (q.shape[0], triton.cdiv(columns, BLOCK_N * tiles_per_program))
647 + ](
648 + q,
649 + k_cache,
650 + page_table,
651 + token_to_req,
652 + query_positions,
653 + sequence_lengths,
654 + visible_blocks,
655 + logits,
656 + q.stride(0),
657 + q.stride(1),
658 + q.stride(2),
659 + k_cache.stride(0),
660 + k_cache.stride(1),
661 + k_cache.stride(3),
662 + page_table.stride(0),
663 + page_table.stride(1),
664 + logits.stride(0),
665 + q.shape[0],
666 + columns,
667 + k_cache.shape[0],
668 + page_table.shape[0],
669 + float(score_divisor),
670 + PAGE_SIZE=k_cache.shape[1],
671 + PAGE_TABLE_WIDTH=page_table.shape[1],
672 + NUM_HEADS=q.shape[1],
673 + HEAD_DIM=q.shape[2],
674 + BLOCK_N=BLOCK_N,
675 + BLOCK_D=BLOCK_D,
676 + TILES_PER_PROG=tiles_per_program,
677 + STAGES=2,
678 + MAX_N=MAX_N,
679 + COMPRESS_RATIO=compress_ratio,
680 + num_warps=2,
681 + )
682 + return logits, visible_blocks
683 +
684 +
685 +def expand_qsa_block_indices_cuda(
686 + block_indices: torch.Tensor,
687 + query_positions: torch.Tensor,
688 + sequence_lengths: torch.Tensor,
689 + token_to_req: torch.Tensor,
690 + compress_ratio: int,
691 + token_topk: int,
692 + out: torch.Tensor | None = None,
693 +) -> torch.Tensor:
694 + """Expand compressed blocks and compact the causal tail of the open group."""
695 +
696 + if not block_indices.is_cuda or not HAS_TRITON:
697 + raise RuntimeError("QSA CUDA expansion requires Triton")
698 + if token_topk % compress_ratio:
699 + raise ValueError("QSA token top-k must be divisible by compression ratio")
700 + block_topk = token_topk // compress_ratio
701 + output_width = token_topk + compress_ratio - 1
702 + if block_indices.shape != (query_positions.numel(), block_topk):
703 + raise ValueError("QSA compressed top-k has an invalid shape")
704 + if token_to_req.shape != query_positions.shape:
705 + raise ValueError("QSA request mapping must match query positions")
706 + if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]:
707 + raise ValueError("QSA request sequence lengths must be nonempty")
708 + if out is None:
709 + out = torch.empty(
710 + (block_indices.shape[0], output_width),
711 + dtype=torch.int32,
712 + device=block_indices.device,
713 + )
714 + elif out.shape != (block_indices.shape[0], output_width):
715 + raise ValueError("QSA expansion output has an invalid shape")
716 + if not block_indices.shape[0]:
717 + return out
718 + column_block = 256
719 + _expand_qsa_indices_kernel[
720 + (block_indices.shape[0], triton.cdiv(output_width, column_block))
721 + ](
722 + block_indices,
723 + query_positions,
724 + sequence_lengths,
725 + token_to_req,
726 + out,
727 + block_indices.stride(0),
728 + block_indices.stride(1),
729 + out.stride(0),
730 + out.stride(1),
731 + block_indices.shape[0],
732 + sequence_lengths.shape[0],
733 + BLOCK_TOPK=block_topk,
734 + COMPRESS_RATIO=compress_ratio,
735 + TOKEN_TOPK=token_topk,
736 + OUTPUT_WIDTH=output_width,
737 + COLUMN_BLOCK=column_block,
738 + num_warps=4,
739 + )
740 + return out
741 +
742 +
743 +def qsa_select_paged_tokens(
744 + q: torch.Tensor,
745 + k_cache: torch.Tensor,
746 + page_table: torch.Tensor,
747 + token_to_req: torch.Tensor,
748 + query_positions: torch.Tensor,
749 + sequence_lengths: torch.Tensor,
750 + token_topk: int,
751 + compress_ratio: int,
752 + out: torch.Tensor | None = None,
753 +) -> torch.Tensor:
754 + """Score, select, and expand QSA indices without host synchronization."""
755 +
756 + rows = q.shape[0]
757 + output_width = token_topk + compress_ratio - 1
758 + if out is None:
759 + out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device)
760 + if out.shape != (rows, output_width):
761 + raise ValueError("QSA selection output has an invalid shape")
762 + if not rows:
763 + return out
764 +
765 + columns = page_table.shape[1] * k_cache.shape[1]
766 + block_topk = token_topk // compress_ratio
767 + rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1))
768 + chunk_rows = min(rows, rows_per_chunk)
769 + blocks_buffer = torch.empty(
770 + (chunk_rows, block_topk), dtype=torch.int32, device=q.device
771 + )
772 + topk_workspace = torch.empty(
773 + (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device
774 + )
775 + for row_start in range(0, rows, rows_per_chunk):
776 + row_end = min(row_start + rows_per_chunk, rows)
777 + row_slice = slice(row_start, row_end)
778 + logits, visible_blocks = qsa_mqa_paged(
779 + q[row_slice],
780 + k_cache,
781 + page_table,
782 + token_to_req[row_slice],
783 + query_positions[row_slice],
784 + sequence_lengths,
785 + compress_ratio,
786 + )
787 + blocks = blocks_buffer[: row_end - row_start]
788 + use_cooperative_topk = (
789 + blocks.shape[0] <= 32
790 + and logits.stride(0) % 4 == 0
791 + and current_platform.has_device_capability(90)
792 + and not current_platform.is_device_capability_family(120)
793 + )
794 + topk_op = (
795 + torch.ops._C.cooperative_topk
796 + if use_cooperative_topk
797 + else torch.ops._C.persistent_topk
798 + )
799 + topk_op(logits, visible_blocks, blocks, topk_workspace, block_topk, columns)
800 + expand_qsa_block_indices_cuda(
801 + blocks,
802 + query_positions[row_slice],
803 + sequence_lengths,
804 + token_to_req[row_slice],
805 + compress_ratio,
806 + token_topk,
807 + out[row_slice],
808 + )
809 + return out
810 +
811 +
812 +def qsa_sparse_paged_attention(
813 + q: torch.Tensor,
814 + k_cache: torch.Tensor,
815 + v_cache: torch.Tensor,
816 + logical_indices: torch.Tensor,
817 + block_table: torch.Tensor,
818 + token_to_req: torch.Tensor,
819 + out: torch.Tensor | None = None,
820 +) -> torch.Tensor:
821 + """Run sparse GQA directly over paged BF16 K/V caches."""
822 +
823 + if not q.is_cuda or not HAS_TRITON:
824 + raise RuntimeError("paged QSA sparse attention requires CUDA and Triton")
825 + if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape:
826 + raise ValueError("QSA sparse attention received invalid Q/K/V shapes")
827 + if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]:
828 + raise ValueError("QSA indices must have one row per query")
829 + if token_to_req.shape != (q.shape[0],) or block_table.ndim != 2:
830 + raise ValueError("QSA sparse attention metadata has invalid shapes")
831 + if not all(k_cache.shape[:3]) or not all(block_table.shape):
832 + raise ValueError("QSA sparse attention cache and block table must be nonempty")
833 + if logical_indices.shape[1] <= 0:
834 + raise ValueError("QSA sparse attention requires a positive selection width")
835 + if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]:
836 + raise ValueError("QSA sparse attention requires valid grouped-query heads")
837 + head_dim = q.shape[2]
838 + assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0
839 + assert q.dtype == k_cache.dtype == v_cache.dtype == torch.bfloat16
840 + assert logical_indices.dtype == block_table.dtype == torch.int32
841 + assert token_to_req.dtype == torch.int32
842 + assert q.device == k_cache.device == v_cache.device
843 + assert q.device == logical_indices.device == block_table.device
844 + assert q.device == token_to_req.device
845 + assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1
846 + assert logical_indices.stride(1) == block_table.stride(1) == 1
847 + assert token_to_req.stride(0) == 1
848 + if out is None:
849 + out = torch.empty_like(q)
850 + if out.shape != q.shape:
851 + raise ValueError("QSA sparse output must match its query")
852 + assert out.dtype == q.dtype and out.device == q.device
853 + assert out.stride(2) == 1
854 + if not q.shape[0]:
855 + return out
856 +
857 + group_size = q.shape[1] // k_cache.shape[2]
858 + block_m = triton.next_power_of_2(group_size)
859 + base_programs = q.shape[0] * k_cache.shape[2]
860 + small_profile_limit = 8 if block_m <= 8 else 4
861 +
862 + # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes.
863 + # Narrow tiles favor decode; wide tiles improve throughput for prefill.
864 + if base_programs <= small_profile_limit:
865 + block_n, target_splits, partial_warps = 16, 64, 4
866 + elif base_programs < 32:
867 + block_n, target_splits, partial_warps = 16, 32, 4
868 + elif base_programs <= 256:
869 + block_n, target_splits, partial_warps = 64, 8, 2
870 + elif base_programs <= 512:
871 + block_n, target_splits, partial_warps = 64, 4, 2
872 + else:
873 + block_n, target_splits, partial_warps = 64, 1, 2
874 +
875 + num_tiles = triton.cdiv(logical_indices.shape[1], block_n)
876 + # Avoid empty splits when the selection width is smaller than the profile.
877 + max_useful_splits = 1 << (num_tiles.bit_length() - 1)
878 + num_splits = min(max_useful_splits, target_splits)
879 +
880 + # Split=1 writes output directly and compiles out all workspace accesses.
881 + if num_splits == 1:
882 + partial_output = out
883 + partial_lse = out
884 + else:
885 + # FP32 partials preserve accuracy when merging independently normalized
886 + # splits.
887 + partial_output = torch.empty(
888 + (num_splits, *q.shape), dtype=torch.float32, device=q.device
889 + )
890 + partial_lse = torch.empty(
891 + (num_splits, q.shape[0], q.shape[1]),
892 + dtype=torch.float32,
893 + device=q.device,
894 + )
895 +
896 + partial_grid = (q.shape[0], k_cache.shape[2], num_splits)
897 + _qsa_sparse_paged_gqa_splitk_kernel[partial_grid](
898 + q,
899 + k_cache,
900 + v_cache,
901 + logical_indices,
902 + block_table,
903 + token_to_req,
904 + partial_output,
905 + partial_lse,
906 + out,
907 + q.stride(0),
908 + q.stride(1),
909 + k_cache.stride(0),
910 + k_cache.stride(1),
911 + k_cache.stride(2),
912 + v_cache.stride(0),
913 + v_cache.stride(1),
914 + v_cache.stride(2),
915 + logical_indices.stride(0),
916 + block_table.stride(0),
917 + out.stride(0),
918 + out.stride(1),
919 + q.shape[0],
920 + k_cache.shape[0],
921 + block_table.shape[0],
922 + TOPK=logical_indices.shape[1],
923 + PAGE_SIZE=k_cache.shape[1],
924 + PAGE_TABLE_WIDTH=block_table.shape[1],
925 + GROUP_SIZE=group_size,
926 + HEAD_DIM=q.shape[2],
927 + NUM_QUERY_HEADS=q.shape[1],
928 + NUM_SPLITS=num_splits,
929 + NUM_TILES=num_tiles,
930 + BLOCK_M=block_m,
931 + BLOCK_N=block_n,
932 + num_warps=partial_warps,
933 + num_stages=2,
934 + )
935 + if num_splits == 1:
936 + return out
937 +
938 + _qsa_merge_splitk_kernel[(q.shape[0], q.shape[1])](
939 + partial_output,
940 + partial_lse,
941 + out,
942 + out.stride(0),
943 + out.stride(1),
944 + q.shape[0],
945 + HEAD_DIM=q.shape[2],
946 + NUM_QUERY_HEADS=q.shape[1],
947 + NUM_SPLITS=num_splits,
948 + BLOCK_SPLITS=triton.next_power_of_2(num_splits),
949 + num_warps=2,
950 + num_stages=1,
951 + )
952 + return out
953 +
954 +
955 +def qsa_store_cache_rows(
956 + cache: torch.Tensor,
957 + slot_mapping: torch.Tensor,
958 + rows: torch.Tensor,
959 +) -> None:
960 + """Store fixed-width rows in a QSA cache without boolean indexing."""
961 +
962 + if not cache.is_cuda or not HAS_TRITON:
963 + raise RuntimeError("QSA CUDA cache stores require Triton")
964 + if cache.ndim != 4 or cache.shape[2] != 1:
965 + raise ValueError("QSA cache must be [pages, page_size, 1, width]")
966 + if not all(cache.shape):
967 + raise ValueError("QSA cache dimensions must be nonzero")
968 + if rows.ndim == 3:
969 + if rows.shape[1] != 1:
970 + raise ValueError("QSA cache rows must have one head")
971 + rows = rows[:, 0]
972 + if rows.shape != (slot_mapping.numel(), cache.shape[3]):
973 + raise ValueError("QSA cache rows and slots have incompatible shapes")
974 + if not rows.shape[0]:
975 + return
976 + _store_qsa_rows_kernel[(rows.shape[0],)](
977 + cache,
978 + slot_mapping,
979 + rows,
980 + cache.stride(0),
981 + cache.stride(1),
982 + cache.stride(3),
983 + rows.stride(0),
984 + rows.stride(1),
985 + rows.shape[0],
986 + cache.shape[0],
987 + PAGE_SIZE=cache.shape[1],
988 + WIDTH=cache.shape[3],
989 + BLOCK_D=triton.next_power_of_2(cache.shape[3]),
990 + num_warps=4,
991 + )
992 +
993 +
994 +def qsa_compress_groups_with_ratio(
995 + raw_keys: torch.Tensor, # this step's raw key rows [rows, 1, head_size]
996 + raw_positions: torch.Tensor, # this step's positions [rows, 1, 3] int64
997 + compressor_state_cache: torch.Tensor,
998 + compressor_state_block_table: torch.Tensor,
999 + token_to_req: torch.Tensor,
1000 + query_start_loc: torch.Tensor,
1001 + logical_positions: torch.Tensor,
1002 + compressed_slots: torch.Tensor,
1003 + compress_ratio: int,
1004 + rope_cache: torch.Tensor | None = None,
1005 +) -> tuple[torch.Tensor, torch.Tensor]:
1006 + """Pool completed groups from the compressor-state ring and raw token rows."""
1007 +
1008 + if not raw_keys.is_cuda or not HAS_TRITON:
1009 + raise RuntimeError("QSA CUDA compression requires Triton")
1010 + rows = token_to_req.numel()
1011 + if compress_ratio <= 0:
1012 + raise ValueError("QSA compression ratio must be positive")
1013 + if raw_keys.ndim != 3 or raw_keys.shape[:2] != (rows, 1):
1014 + raise ValueError("QSA raw keys must be [rows, 1, head_size]")
1015 + if raw_positions.shape != (rows, 1, 3) or raw_positions.dtype != torch.int64:
1016 + raise ValueError("QSA raw positions must be [rows, 1, 3] int64")
1017 + if logical_positions.shape != (rows,) or compressed_slots.shape != (rows,):
1018 + raise ValueError("QSA compression metadata must match token rows")
1019 + if compressor_state_cache.ndim != 4 or compressor_state_cache.shape[2] != 1:
1020 + raise ValueError("QSA compressor-state cache has an invalid shape")
1021 + if (
1022 + # The ring is wider than one group so speculative rows cannot alias
1023 + # onto the committed keys of the group still being collected.
1024 + compressor_state_cache.shape[1] < compress_ratio
1025 + or compressor_state_cache.shape[3] != raw_keys.shape[2]
1026 + or compressor_state_cache.dtype != raw_keys.dtype
1027 + ):
1028 + raise ValueError(
1029 + "QSA compressor-state cache does not match the compression layout"
1030 + )
1031 + if (
1032 + compressor_state_block_table.ndim != 2
1033 + or compressor_state_block_table.shape[1] < 1
1034 + ):
1035 + raise ValueError(
1036 + "QSA compressor-state block table must contain one block per request"
1037 + )
1038 + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2:
1039 + raise ValueError("QSA query starts must contain a terminal offset")
1040 + num_requests = query_start_loc.shape[0] - 1
1041 + if compressor_state_block_table.shape[0] < num_requests:
1042 + raise ValueError("QSA compressor-state block table has too few request rows")
1043 + if rope_cache is not None and (
1044 + rope_cache.ndim != 4
1045 + or rope_cache.shape[:3] != compressor_state_cache.shape[:3]
1046 + or rope_cache.shape[3] != 3
1047 + or rope_cache.dtype != torch.int64
1048 + ):
1049 + raise ValueError("QSA packed position view has an invalid shape or dtype")
1050 + if rows and (
1051 + not all(compressor_state_cache.shape)
1052 + or not all(compressor_state_block_table.shape)
1053 + ):
1054 + raise ValueError("QSA compressor-state cache and block table must be nonempty")
1055 + pooled = torch.empty(
1056 + (rows, 1, raw_keys.shape[2]),
1057 + dtype=raw_keys.dtype,
1058 + device=raw_keys.device,
1059 + )
1060 + first_positions = torch.empty((rows, 3), dtype=torch.int64, device=raw_keys.device)
1061 + if not rows:
1062 + return pooled, first_positions
1063 + if rope_cache is None:
1064 + rope_cache = compressor_state_cache
1065 + load_rope_positions = False
1066 + else:
1067 + load_rope_positions = True
1068 + _compress_qsa_groups_kernel[(rows,)](
1069 + raw_keys,
1070 + raw_positions,
1071 + compressor_state_cache,
1072 + rope_cache,
1073 + compressor_state_block_table,
1074 + token_to_req,
1075 + query_start_loc,
1076 + logical_positions,
1077 + compressed_slots,
1078 + pooled,
1079 + first_positions,
1080 + raw_keys.stride(0),
1081 + raw_keys.stride(2),
1082 + raw_positions.stride(0),
1083 + raw_positions.stride(2),
1084 + compressor_state_cache.stride(0),
1085 + compressor_state_cache.stride(1),
1086 + compressor_state_cache.stride(3),
1087 + rope_cache.stride(0),
1088 + rope_cache.stride(1),
1089 + rope_cache.stride(3),
1090 + compressor_state_block_table.stride(0),
1091 + pooled.stride(0),
1092 + pooled.stride(2),
1093 + first_positions.stride(0),
1094 + first_positions.stride(1),
1095 + rows,
1096 + compressor_state_cache.shape[0],
1097 + num_requests,
1098 + COMPRESSOR_STATE_SIZE=compressor_state_cache.shape[1],
1099 + COMPRESS_RATIO=compress_ratio,
1100 + HEAD_DIM=raw_keys.shape[2],
1101 + LOAD_ROPE_POSITIONS=load_rope_positions,
1102 + BLOCK_D=triton.next_power_of_2(raw_keys.shape[2]),
1103 + num_warps=4,
1104 + )
1105 + return pooled, first_positions
1106 +
1107 +
1108 +__all__ = [
1109 + "expand_qsa_block_indices_cuda",
1110 + "qsa_compress_groups_with_ratio",
1111 + "qsa_mqa_paged",
1112 + "qsa_select_paged_tokens",
1113 + "qsa_sparse_paged_attention",
1114 + "qsa_store_cache_rows",
1115 +]
vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py new
+508
@@ -0,0 +1,508 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused QSA pre-indexer kernel for Qwen4Exp."""
4 +
5 +import torch
6 +
7 +from vllm.triton_utils import tl, triton
8 +
9 +
10 +@triton.jit
11 +def _norm_rope(
12 + x,
13 + pos_t,
14 + pos_h,
15 + pos_w,
16 + cos_sin_ptr,
17 + cos_sin_stride,
18 + norm_weight_ptr,
19 + eps,
20 + IS_MROPE: tl.constexpr,
21 + MROPE_H: tl.constexpr,
22 + MROPE_W: tl.constexpr,
23 +):
24 + """Apply Gemma RMSNorm and selected-axis NeoX RoPE to register rows."""
25 + TILE_T: tl.constexpr = x.shape[0]
26 + TILE_H: tl.constexpr = x.shape[1]
27 + D: tl.constexpr = x.shape[2]
28 + ROWS: tl.constexpr = TILE_T * TILE_H
29 + HALF: tl.constexpr = D // 2
30 + QUARTER: tl.constexpr = D // 4
31 + pairs = tl.arange(0, QUARTER)
32 + if IS_MROPE:
33 + # Qwen interleaves temporal, height, and width rotary pairs. Each axis
34 + # still indexes the same position-major cos/sin table.
35 + h_mask = ((pairs % 3) == 1) & (pairs <= 3 * MROPE_H)
36 + w_mask = ((pairs % 3) == 2) & (pairs <= 3 * MROPE_W)
37 + t_mask = ~(h_mask | w_mask)
38 + base = cos_sin_ptr + pairs[None, :]
39 + pos_rows = (pos_t, pos_h, pos_w)
40 + axis_masks = (t_mask, h_mask, w_mask)
41 + cos = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty)
42 + sin = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty)
43 + for axis in tl.static_range(3):
44 + cos += tl.load(
45 + base + pos_rows[axis][:, None] * cos_sin_stride,
46 + mask=axis_masks[axis][None, :],
47 + other=0,
48 + )
49 + sin += tl.load(
50 + base + pos_rows[axis][:, None] * cos_sin_stride + QUARTER,
51 + mask=axis_masks[axis][None, :],
52 + other=0,
53 + )
54 + else:
55 + cos = tl.load(cos_sin_ptr + pos_t[:, None] * cos_sin_stride + pairs[None, :])
56 + sin = tl.load(
57 + cos_sin_ptr + pos_t[:, None] * cos_sin_stride + QUARTER + pairs[None, :]
58 + )
59 +
60 + cos = tl.reshape(
61 + tl.broadcast_to(cos[:, None, :], (TILE_T, TILE_H, QUARTER)),
62 + (ROWS, QUARTER),
63 + )
64 + sin = tl.reshape(
65 + tl.broadcast_to(sin[:, None, :], (TILE_T, TILE_H, QUARTER)),
66 + (ROWS, QUARTER),
67 + )
68 + x = tl.reshape(x, (ROWS, D)).to(tl.float32)
69 + weight = tl.load(norm_weight_ptr + tl.arange(0, D)).to(tl.float32) + 1.0
70 + rrms = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps)
71 + y = (x * rrms[:, None] * weight[None, :]).to(cos.dtype)
72 + rotated, passthrough = tl.split(
73 + tl.permute(tl.reshape(y, (ROWS, 2, HALF)), (0, 2, 1))
74 + )
75 + r0, r1 = tl.split(tl.permute(tl.reshape(rotated, (ROWS, 2, QUARTER)), (0, 2, 1)))
76 + out0 = r0 * cos - r1 * sin
77 + out1 = r1 * cos + r0 * sin
78 + rotated = tl.reshape(tl.permute(tl.join(out0, out1), (0, 2, 1)), (ROWS, HALF))
79 + result = tl.reshape(tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), (ROWS, D))
80 + return tl.reshape(result, (TILE_T, TILE_H, D))
81 +
82 +
83 +@triton.jit(
84 + do_not_specialize=[
85 + "num_tokens",
86 + "num_state_blocks",
87 + "num_compressed_blocks",
88 + "num_k_work",
89 + ]
90 +)
91 +def _qsa_pre_indexer_kernel(
92 + q_ptr,
93 + q_stride_token,
94 + k_ptr,
95 + k_stride_token,
96 + pos_ptr,
97 + pos_stride_axis,
98 + pos_stride_token,
99 + cos_sin_ptr,
100 + q_norm_weight_ptr,
101 + k_norm_weight_ptr,
102 + eps,
103 + q_out_ptr,
104 + q_out_stride_token,
105 + q_out_stride_head,
106 + state_cache_ptr,
107 + state_cache_stride_block,
108 + state_cache_stride_token,
109 + state_slots_ptr,
110 + state_table_ptr,
111 + state_table_stride_req,
112 + query_start_loc_ptr,
113 + logical_positions_ptr,
114 + compressed_slots_ptr,
115 + k_work_metadata_ptr,
116 + compressed_cache_ptr,
117 + compressed_cache_stride_block,
118 + compressed_cache_stride_token,
119 + num_tokens,
120 + num_state_blocks,
121 + num_compressed_blocks,
122 + num_k_work,
123 + HQ: tl.constexpr,
124 + D: tl.constexpr,
125 + TILE_T_Q: tl.constexpr,
126 + TILE_H_Q: tl.constexpr,
127 + COMPRESS_RATIO: tl.constexpr,
128 + STATE_SIZE: tl.constexpr,
129 + COMP_PAGE_SIZE: tl.constexpr,
130 + IS_2D_POSITIONS: tl.constexpr,
131 + IS_K_MROPE: tl.constexpr,
132 + CACHE_HAS_ROPE_POS: tl.constexpr,
133 + MROPE_H: tl.constexpr,
134 + MROPE_W: tl.constexpr,
135 +):
136 + pid = tl.program_id(0)
137 + # K work occupies the first programs; the remaining programs tile Q. This
138 + # keeps both paths in one launch while leaving their register shapes
139 + # independent.
140 + if pid >= num_k_work:
141 + q_pid = pid - num_k_work
142 + num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q)
143 + token_tile = q_pid // num_head_tiles
144 + head_tile = q_pid % num_head_tiles
145 + tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q)
146 + heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q)
147 + valid_tokens = tokens < num_tokens
148 + valid_heads = heads < HQ
149 + dims = tl.arange(0, D)
150 + mask = valid_tokens[:, None, None] & valid_heads[None, :, None]
151 + x = tl.load(
152 + q_ptr
153 + + tokens[:, None, None] * q_stride_token
154 + + heads[None, :, None] * D
155 + + dims[None, None, :],
156 + mask=mask,
157 + other=0.0,
158 + )
159 + pos_t = tl.load(pos_ptr + tokens * pos_stride_token, mask=valid_tokens, other=0)
160 + if IS_2D_POSITIONS:
161 + pos_h = tl.load(
162 + pos_ptr + pos_stride_axis + tokens * pos_stride_token,
163 + mask=valid_tokens,
164 + other=0,
165 + )
166 + pos_w = tl.load(
167 + pos_ptr + 2 * pos_stride_axis + tokens * pos_stride_token,
168 + mask=valid_tokens,
169 + other=0,
170 + )
171 + else:
172 + pos_h = pos_t
173 + pos_w = pos_t
174 + y = _norm_rope(
175 + x,
176 + pos_t,
177 + pos_h,
178 + pos_w,
179 + cos_sin_ptr,
180 + D // 2,
181 + q_norm_weight_ptr,
182 + eps,
183 + IS_2D_POSITIONS,
184 + MROPE_H,
185 + MROPE_W,
186 + )
187 + tl.store(
188 + q_out_ptr
189 + + tokens[:, None, None] * q_out_stride_token
190 + + heads[None, :, None] * q_out_stride_head
191 + + dims[None, None, :],
192 + y,
193 + mask=mask,
194 + )
195 +
196 + if pid < num_k_work:
197 + # One work item owns one completed compression group. Work item zero
198 + # additionally commits the request's current raw-K suffix below.
199 + work_metadata = tl.load(k_work_metadata_ptr + pid * 2 + tl.arange(0, 2))
200 + request, work_in_request = tl.split(work_metadata)
201 + if request < 0:
202 + return
203 +
204 + query_start = tl.load(query_start_loc_ptr + request)
205 + query_end = tl.load(query_start_loc_ptr + request + 1)
206 + query_len = query_end - query_start
207 + chunk_end = tl.load(logical_positions_ptr + query_end - 1)
208 + chunk_start = chunk_end - query_len + 1
209 + num_groups = (chunk_end + 1) // COMPRESS_RATIO - chunk_start // COMPRESS_RATIO
210 + dims = tl.arange(0, D)
211 +
212 + if work_in_request < num_groups:
213 + first_boundary = (
214 + (chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO
215 + ) * COMPRESS_RATIO - 1
216 + end_position = first_boundary + work_in_request * COMPRESS_RATIO
217 + boundary_token = query_start + end_position - chunk_start
218 + valid_token = (
219 + (boundary_token >= query_start)
220 + & (boundary_token < query_end)
221 + & (boundary_token < num_tokens)
222 + )
223 + compressed_slot = tl.load(
224 + compressed_slots_ptr + boundary_token,
225 + mask=valid_token,
226 + other=-1,
227 + )
228 + valid = (
229 + valid_token
230 + & (compressed_slot >= 0)
231 + & (compressed_slot < num_compressed_blocks * COMP_PAGE_SIZE)
232 + )
233 + state_block = tl.load(state_table_ptr + request * state_table_stride_req)
234 + state_block_valid = (state_block >= 0) & (state_block < num_state_blocks)
235 + safe_state_block = tl.maximum(state_block, 0).to(tl.int64)
236 + group_offsets = tl.arange(0, COMPRESS_RATIO)
237 + source_positions = end_position - (COMPRESS_RATIO - 1) + group_offsets
238 + source_in_chunk = source_positions >= chunk_start
239 + source_tokens = query_start + source_positions - chunk_start
240 + source_tokens_valid = (
241 + (source_tokens >= query_start)
242 + & (source_tokens < query_end)
243 + & (source_tokens < num_tokens)
244 + )
245 + current_base = (
246 + k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token
247 + )
248 + cached_base = (
249 + state_cache_ptr
250 + + safe_state_block * state_cache_stride_block
251 + + (source_positions % STATE_SIZE)[:, None] * state_cache_stride_token
252 + )
253 + # Only the first completed group can cross the chunk boundary. Select
254 + # historical rows from the ring without issuing two masked loads.
255 + source_base = tl.where(source_in_chunk[:, None], current_base, cached_base)
256 + # Pointer selection obscures alignment from Triton's analysis.
257 + source_base = tl.multiple_of(source_base, (8, 8))
258 + source_valid = tl.where(
259 + source_in_chunk, source_tokens_valid, state_block_valid
260 + )
261 + source = tl.load(
262 + source_base + dims[None, :],
263 + mask=valid & source_valid[:, None],
264 + other=0.0,
265 + ).to(tl.float32)
266 + # Match the unfused path's BF16 pooled tensor before RMSNorm.
267 + pooled = (
268 + (tl.sum(source, axis=0) / COMPRESS_RATIO).to(tl.bfloat16).to(tl.float32)
269 + )
270 +
271 + first_position = end_position - (COMPRESS_RATIO - 1)
272 + if CACHE_HAS_ROPE_POS:
273 + # RoPE uses the first token in the pooled group. Its exact MRoPE
274 + # coordinates may live in this chunk or the raw-state ring.
275 + first_in_chunk = first_position >= chunk_start
276 + first_token = query_start + first_position - chunk_start
277 + first_token_valid = (
278 + (first_token >= query_start)
279 + & (first_token < query_end)
280 + & (first_token < num_tokens)
281 + )
282 + safe_first_token = tl.maximum(first_token, 0)
283 + load_current_position = first_in_chunk & first_token_valid
284 + first_pos_t = tl.load(
285 + pos_ptr + safe_first_token * pos_stride_token,
286 + mask=load_current_position,
287 + other=0,
288 + )
289 + if IS_2D_POSITIONS:
290 + first_pos_h = tl.load(
291 + pos_ptr + pos_stride_axis + safe_first_token * pos_stride_token,
292 + mask=load_current_position,
293 + other=0,
294 + )
295 + first_pos_w = tl.load(
296 + pos_ptr
297 + + 2 * pos_stride_axis
298 + + safe_first_token * pos_stride_token,
299 + mask=load_current_position,
300 + other=0,
301 + )
302 + else:
303 + first_pos_h = first_pos_t
304 + first_pos_w = first_pos_t
305 + tail = (
306 + state_cache_ptr
307 + + safe_state_block * state_cache_stride_block
308 + + (first_position % STATE_SIZE) * state_cache_stride_token
309 + + D
310 + ).to(tl.pointer_type(tl.int64))
311 + load_cached_position = ~first_in_chunk & state_block_valid
312 + cached_pos_t = tl.load(tail, mask=load_cached_position, other=0)
313 + cached_pos_h = tl.load(tail + 1, mask=load_cached_position, other=0)
314 + cached_pos_w = tl.load(tail + 2, mask=load_cached_position, other=0)
315 + pos_t = tl.where(first_in_chunk, first_pos_t.to(tl.int64), cached_pos_t)
316 + pos_h = tl.where(first_in_chunk, first_pos_h.to(tl.int64), cached_pos_h)
317 + pos_w = tl.where(first_in_chunk, first_pos_w.to(tl.int64), cached_pos_w)
318 + else:
319 + pos_t = first_position
320 + pos_h = first_position
321 + pos_w = first_position
322 + y = _norm_rope(
323 + tl.reshape(pooled, (1, 1, D)),
324 + pos_t + tl.arange(0, 1),
325 + pos_h + tl.arange(0, 1),
326 + pos_w + tl.arange(0, 1),
327 + cos_sin_ptr,
328 + D // 2,
329 + k_norm_weight_ptr,
330 + eps,
331 + IS_K_MROPE,
332 + MROPE_H,
333 + MROPE_W,
334 + )
335 + compressed_block = (compressed_slot // COMP_PAGE_SIZE).to(tl.int64)
336 + compressed_row = compressed_slot % COMP_PAGE_SIZE
337 + tl.store(
338 + compressed_cache_ptr
339 + + compressed_block * compressed_cache_stride_block
340 + + compressed_row * compressed_cache_stride_token
341 + + dims,
342 + tl.reshape(y, (D,)),
343 + mask=valid,
344 + )
345 +
346 + if work_in_request == 0:
347 + # This CTA may have just read historical rows from the circular buffer.
348 + # Keep every lane past those loads before overwriting the same ring.
349 + tl.debug_barrier()
350 + # One CTA per request commits only the suffix retained by the ring.
351 + num_state_rows = tl.minimum(query_len, STATE_SIZE)
352 + for state_offset in tl.range(0, num_state_rows):
353 + token = query_end - num_state_rows + state_offset
354 + valid_token = (
355 + (token >= query_start) & (token < query_end) & (token < num_tokens)
356 + )
357 + slot = tl.load(state_slots_ptr + token, mask=valid_token, other=-1)
358 + valid_slot = (
359 + valid_token & (slot >= 0) & (slot < num_state_blocks * STATE_SIZE)
360 + )
361 + safe_slot = tl.maximum(slot, 0)
362 + state_row = (
363 + state_cache_ptr
364 + + (safe_slot // STATE_SIZE).to(tl.int64) * state_cache_stride_block
365 + + (safe_slot % STATE_SIZE) * state_cache_stride_token
366 + )
367 + k = tl.load(
368 + k_ptr + tl.maximum(token, 0) * k_stride_token + dims,
369 + mask=valid_slot,
370 + other=0.0,
371 + )
372 + tl.store(state_row + dims, k, mask=valid_slot)
373 + if CACHE_HAS_ROPE_POS:
374 + pos_t = tl.load(
375 + pos_ptr + tl.maximum(token, 0) * pos_stride_token,
376 + mask=valid_slot,
377 + other=0,
378 + )
379 + if IS_2D_POSITIONS:
380 + pos_h = tl.load(
381 + pos_ptr
382 + + pos_stride_axis
383 + + tl.maximum(token, 0) * pos_stride_token,
384 + mask=valid_slot,
385 + other=0,
386 + )
387 + pos_w = tl.load(
388 + pos_ptr
389 + + 2 * pos_stride_axis
390 + + tl.maximum(token, 0) * pos_stride_token,
391 + mask=valid_slot,
392 + other=0,
393 + )
394 + else:
395 + pos_h = pos_t
396 + pos_w = pos_t
397 + tail = (state_row + D).to(tl.pointer_type(tl.int64))
398 + tl.store(tail, pos_t.to(tl.int64), mask=valid_slot)
399 + tl.store(tail + 1, pos_h.to(tl.int64), mask=valid_slot)
400 + tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot)
401 +
402 +
403 +def qsa_pre_indexer(
404 + q: torch.Tensor,
405 + k: torch.Tensor,
406 + positions: torch.Tensor,
407 + cos_sin_cache: torch.Tensor,
408 + q_norm_weight: torch.Tensor,
409 + k_norm_weight: torch.Tensor,
410 + eps: float,
411 + q_out: torch.Tensor,
412 + state_cache: torch.Tensor,
413 + state_slots: torch.Tensor,
414 + state_block_table: torch.Tensor,
415 + query_start_loc: torch.Tensor,
416 + logical_positions: torch.Tensor,
417 + compressed_cache: torch.Tensor,
418 + compressed_slots: torch.Tensor,
419 + k_work_metadata: torch.Tensor,
420 + *,
421 + compress_ratio: int,
422 + mrope_section: tuple[int, int, int] | None,
423 + rope_pos_offset: int | None,
424 +) -> None:
425 + """Normalize Q, compress K, then update the circular raw state."""
426 + num_tokens = q.shape[0]
427 + if num_tokens == 0:
428 + return
429 + num_q_heads, head_dim = q_out.shape[1:]
430 + assert cos_sin_cache.shape[-1] * 2 == head_dim
431 + assert q.shape == (num_tokens, num_q_heads * head_dim)
432 + assert k.shape == (num_tokens, head_dim)
433 + assert q.stride(-1) == 1
434 + assert k.stride(-1) == 1
435 + assert q_out.stride(-1) == 1
436 + assert cos_sin_cache.is_contiguous()
437 + assert state_cache.stride(-1) == 1
438 + assert compressed_cache.stride(-1) == 1
439 + assert k_work_metadata.ndim == 2 and k_work_metadata.shape[1] == 2
440 + is_2d_positions = positions.ndim == 2
441 + is_k_mrope = bool(mrope_section)
442 + cache_has_rope_pos = rope_pos_offset is not None
443 + assert rope_pos_offset is None or rope_pos_offset == head_dim
444 + if is_2d_positions:
445 + assert positions.shape == (3, num_tokens)
446 + assert is_k_mrope
447 + pos_stride_axis, pos_stride_token = positions.stride()
448 + else:
449 + assert positions.shape == (num_tokens,)
450 + pos_stride_axis, pos_stride_token = 0, positions.stride(0)
451 + section = mrope_section if mrope_section is not None else (0, 0, 0)
452 + assert len(section) == 3
453 +
454 + if num_tokens <= 4096:
455 + TILE_T_Q, TILE_H_Q = 2, 2
456 + else:
457 + TILE_T_Q, TILE_H_Q = 2, 4
458 + num_k_work = k_work_metadata.shape[0]
459 + num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q)
460 + _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)](
461 + q,
462 + q.stride(0),
463 + k,
464 + k.stride(0),
465 + positions,
466 + pos_stride_axis,
467 + pos_stride_token,
468 + cos_sin_cache,
469 + q_norm_weight,
470 + k_norm_weight,
471 + eps,
472 + q_out,
473 + q_out.stride(0),
474 + q_out.stride(1),
475 + state_cache,
476 + state_cache.stride(0),
477 + state_cache.stride(1),
478 + state_slots,
479 + state_block_table,
480 + state_block_table.stride(0),
481 + query_start_loc,
482 + logical_positions,
483 + compressed_slots,
484 + k_work_metadata,
485 + compressed_cache,
486 + compressed_cache.stride(0),
487 + compressed_cache.stride(1),
488 + num_tokens,
489 + state_cache.shape[0],
490 + compressed_cache.shape[0],
491 + num_k_work,
492 + HQ=num_q_heads,
493 + D=head_dim,
494 + TILE_T_Q=TILE_T_Q,
495 + TILE_H_Q=TILE_H_Q,
496 + COMPRESS_RATIO=compress_ratio,
497 + STATE_SIZE=state_cache.shape[1],
498 + COMP_PAGE_SIZE=compressed_cache.shape[1],
499 + IS_2D_POSITIONS=is_2d_positions,
500 + IS_K_MROPE=is_k_mrope,
501 + CACHE_HAS_ROPE_POS=cache_has_rope_pos,
502 + MROPE_H=section[1],
503 + MROPE_W=section[2],
504 + num_warps=1,
505 + )
506 +
507 +
508 +__all__ = ["qsa_pre_indexer"]
vllm/models/qwen4_exp/nvidia/ple_layer.py new
+1256
@@ -0,0 +1,1256 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""GPU-resident Qwen4Exp position-learning enhancement layers."""
4 +
5 +import math
6 +from collections.abc import Iterable, Sequence
7 +
8 +import torch
9 +import torch.nn.functional as F
10 +from torch import nn
11 +
12 +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
13 +from vllm.forward_context import get_forward_context
14 +from vllm.model_executor.layers.linear import ReplicatedLinear
15 +from vllm.model_executor.layers.mamba.abstract import MambaBase
16 +from vllm.model_executor.layers.mamba.mamba_utils import (
17 + MambaStateDtypeCalculator,
18 + MambaStateShapeCalculator,
19 + is_conv_state_dim_first,
20 +)
21 +from vllm.model_executor.layers.quantization.base_config import (
22 + QuantizationConfig,
23 + QuantizeMethodBase,
24 +)
25 +from vllm.model_executor.layers.quantization.fp8 import Fp8Config
26 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
27 + create_fp8_scale_parameter,
28 + create_fp8_weight_parameter,
29 + is_fp8,
30 +)
31 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
32 + is_layer_skipped,
33 +)
34 +from vllm.model_executor.models.utils import AutoWeightsLoader
35 +from vllm.model_executor.parameter import PerTensorScaleParameter
36 +from vllm.transformers_utils.configs.qwen4_exp import (
37 + Qwen4ExpTextConfig,
38 +)
39 +from vllm.utils.torch_utils import direct_register_custom_op
40 +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
41 +from vllm.v1.attention.backends.short_conv_attn import (
42 + PleShortConvAttentionBackend,
43 + PleShortConvAttentionMetadata,
44 +)
45 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
46 +
47 +from ..common.ple import PLEVocabParallelEmbedding
48 +
49 +
50 +class Qwen4ExpPLEGroupedNorm(nn.Module):
51 + def __init__(
52 + self,
53 + hidden_size: int,
54 + eps: float,
55 + group_size: int | None,
56 + dtype: torch.dtype | None,
57 + ) -> None:
58 + super().__init__()
59 + if group_size is not None and hidden_size % group_size:
60 + raise ValueError(
61 + f"hidden_size ({hidden_size}) must be divisible by "
62 + f"group_size ({group_size})"
63 + )
64 + self.eps = eps
65 + self.group_size = group_size
66 + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype))
67 +
68 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
69 + input_dtype = hidden_states.dtype
70 + hidden_states = hidden_states.float()
71 + if self.group_size is None:
72 + variance = hidden_states.square().mean(dim=-1, keepdim=True)
73 + normalized = hidden_states * torch.rsqrt(variance + self.eps)
74 + else:
75 + grouped = hidden_states.unflatten(
76 + -1, (hidden_states.shape[-1] // self.group_size, self.group_size)
77 + )
78 + variance = grouped.square().mean(dim=-1, keepdim=True)
79 + normalized = (grouped * torch.rsqrt(variance + self.eps)).flatten(-2)
80 + return (normalized * (1.0 + self.weight.float())).to(input_dtype)
81 +
82 +
83 +class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase):
84 + """FP8 PLE embedding with one global checkpoint scale."""
85 +
86 + def create_weights(
87 + self,
88 + layer: nn.Module,
89 + input_size_per_partition: int,
90 + output_partition_sizes: list[int],
91 + input_size: int,
92 + output_size: int,
93 + params_dtype: torch.dtype,
94 + **extra_weight_attrs,
95 + ) -> None:
96 + del input_size, output_size, params_dtype
97 + weight_loader = extra_weight_attrs.get("weight_loader")
98 + weight = create_fp8_weight_parameter(
99 + sum(output_partition_sizes), input_size_per_partition, weight_loader
100 + )
101 + layer.register_parameter("weight", weight)
102 +
103 + weight_scale = create_fp8_scale_parameter(
104 + PerTensorScaleParameter,
105 + output_partition_sizes,
106 + input_size_per_partition,
107 + None,
108 + weight_loader,
109 + scale_dtype=torch.bfloat16,
110 + )
111 + layer.register_parameter("weight_scale", weight_scale)
112 +
113 + def apply(
114 + self,
115 + layer: nn.Module,
116 + x: torch.Tensor,
117 + bias: torch.Tensor | None = None,
118 + ) -> torch.Tensor:
119 + raise NotImplementedError("PLE FP8 weights only support embedding lookup")
120 +
121 + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor:
122 + return F.embedding(input_, layer.weight)
123 +
124 +
125 +def _get_ple_embedding_quant_method(
126 + quant_config: QuantizationConfig | None,
127 + prefix: str,
128 +) -> QuantizeMethodBase | None:
129 + """Select global-scale FP8 only for quantized PLE checkpoint shards."""
130 +
131 + if not isinstance(quant_config, Fp8Config):
132 + return None
133 + if not quant_config.is_checkpoint_fp8_serialized:
134 + return None
135 +
136 + ignored_layers = quant_config.ignored_layers
137 + if is_layer_skipped(
138 + prefix,
139 + ignored_layers,
140 + quant_config.packed_modules_mapping,
141 + match_mode=quant_config.ignored_layers_match_mode,
142 + ):
143 + return None
144 + # PLE checkpoint shards form one runtime embedding parameter.
145 + shard_prefix = f"{prefix}.shard_"
146 + if any(name.startswith(shard_prefix) for name in ignored_layers):
147 + return None
148 + return Qwen4ExpPLEFp8EmbeddingMethod()
149 +
150 +
151 +class Qwen4ExpNGramEmbedding(nn.Module):
152 + _MASK64 = (1 << 64) - 1
153 + _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
154 + _SPLITMIX_M1 = 0xBF58476D1CE4E5B9
155 + _SPLITMIX_M2 = 0x94D049BB133111EB
156 + _PLE_LAYER_PRIME = 10007
157 +
158 + @classmethod
159 + def _splitmix64(cls, value: int) -> int:
160 + """Mix an integer into a deterministic unsigned 64-bit value."""
161 + value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64
162 + value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64
163 + value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64
164 + return (value ^ (value >> 31)) & cls._MASK64
165 +
166 + @staticmethod
167 + def _is_prime_64(value: int) -> bool:
168 + """Return whether a 64-bit integer is prime."""
169 + if value < 2:
170 + return False
171 + for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
172 + if value % prime == 0:
173 + return value == prime
174 + exponent = value - 1
175 + shifts = 0
176 + while exponent % 2 == 0:
177 + exponent //= 2
178 + shifts += 1
179 + for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022):
180 + if base % value == 0:
181 + continue
182 + witness = pow(base, exponent, value)
183 + if witness in (1, value - 1):
184 + continue
185 + for _ in range(shifts - 1):
186 + witness = pow(witness, 2, value)
187 + if witness == value - 1:
188 + break
189 + else:
190 + return False
191 + return True
192 +
193 + @classmethod
194 + def _nth_prime_after(cls, start: int, count: int) -> int:
195 + """Return the ``count``-th prime strictly greater than ``start``."""
196 + prime = int(start)
197 + for _ in range(count):
198 + candidate = prime + 1
199 + if candidate <= 2:
200 + prime = 2
201 + continue
202 + if candidate % 2 == 0:
203 + candidate += 1
204 + while not cls._is_prime_64(candidate):
205 + candidate += 2
206 + prime = candidate
207 + return prime
208 +
209 + @classmethod
210 + def _make_layer_multipliers(
211 + cls,
212 + *,
213 + ngram_size: int,
214 + unigram_vocab_size: int,
215 + seed: int,
216 + ple_dense_layer_id: int,
217 + ) -> list[int]:
218 + """Build deterministic hash multipliers for one PLE layer."""
219 + max_multiplier = ((1 << 63) - 1) // unigram_vocab_size
220 + half_bound = max(1, max_multiplier // 2)
221 + base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id
222 + multipliers = []
223 + for index in range(ngram_size):
224 + value = base_seed + cls._SPLITMIX_GAMMA * (index + 1)
225 + multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1)
226 + return multipliers
227 +
228 + @classmethod
229 + def _make_vocab_layout(
230 + cls,
231 + *,
232 + ngram_vocab_size_base: int,
233 + ngram_heads: int,
234 + ple_dense_layer_id: int,
235 + ) -> tuple[list[int], list[int], int]:
236 + """Build per-head vocabulary sizes, offsets, and total row count."""
237 + sizes: list[int] = []
238 + offsets: list[int] = []
239 + offset = 0
240 + for local_head in range(ngram_heads):
241 + global_head = ple_dense_layer_id * ngram_heads + local_head
242 + size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1)
243 + sizes.append(size)
244 + offsets.append(offset)
245 + offset += size
246 + return sizes, offsets, offset
247 +
248 + def __init__(
249 + self,
250 + config: Qwen4ExpTextConfig,
251 + embedding_dim: int,
252 + ple_dense_layer_id: int,
253 + max_total_tokens: int,
254 + max_num_reqs: int,
255 + prefix: str,
256 + layer_name: str,
257 + quant_config: QuantizationConfig | None = None,
258 + params_dtype: torch.dtype | None = None,
259 + ) -> None:
260 + super().__init__()
261 + self.layer_name = layer_name
262 + self.embedding_dim = embedding_dim
263 + self.ngram_size = int(config.ngram_size)
264 + self.heads_per_ngram = int(config.heads_per_ngram)
265 + self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
266 + if self.ngram_size < 2:
267 + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}")
268 + if self.heads_per_ngram <= 0:
269 + raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}")
270 + if embedding_dim % self.ngram_heads:
271 + raise ValueError(
272 + "ple_embed_dim must be divisible by total ngram heads: "
273 + f"{embedding_dim} % {self.ngram_heads} != 0"
274 + )
275 + self.head_dim = embedding_dim // self.ngram_heads
276 + self.eos_token_id = int(config.eos_token_id)
277 + self.unigram_vocab_size = int(config.vocab_size)
278 + self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512))
279 + if self.split_ngram_parts <= 0:
280 + raise ValueError("split_ngram_parts must be positive")
281 +
282 + multipliers = self._make_layer_multipliers(
283 + ngram_size=self.ngram_size,
284 + unigram_vocab_size=self.unigram_vocab_size,
285 + seed=int(getattr(config, "seed", 1234)),
286 + ple_dense_layer_id=ple_dense_layer_id,
287 + )
288 + self.register_buffer(
289 + "layer_multipliers",
290 + torch.tensor(multipliers, dtype=torch.long),
291 + persistent=True,
292 + )
293 +
294 + sizes, offsets, total_vocab_size = self._make_vocab_layout(
295 + ngram_vocab_size_base=int(config.ngram_vocab_size_base),
296 + ngram_heads=self.ngram_heads,
297 + ple_dense_layer_id=ple_dense_layer_id,
298 + )
299 + self.register_buffer(
300 + "ngram_heads_vocab_sizes",
301 + torch.tensor(sizes, dtype=torch.long),
302 + persistent=True,
303 + )
304 + self.register_buffer(
305 + "ngram_heads_offsets",
306 + torch.tensor(offsets, dtype=torch.long),
307 + persistent=True,
308 + )
309 + divisor = int(config.make_ngram_vocab_size_divisible_by)
310 + padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor
311 + self.ngram_embedding = PLEVocabParallelEmbedding(
312 + padded_vocab_size,
313 + self.head_dim,
314 + params_dtype=params_dtype,
315 + padding_size=divisor,
316 + prefix=f"{prefix}.ngram_embedding",
317 + quant_method=_get_ple_embedding_quant_method(
318 + quant_config, f"{prefix}.ngram_embedding"
319 + ),
320 + )
321 + self.register_buffer(
322 + "positions_buffer",
323 + torch.arange(max_total_tokens, dtype=torch.int64),
324 + persistent=False,
325 + )
326 + self.register_buffer(
327 + "padded_buffer",
328 + torch.full(
329 + (max_num_reqs, max_total_tokens),
330 + self.eos_token_id,
331 + dtype=torch.int64,
332 + ),
333 + persistent=False,
334 + )
335 +
336 + @staticmethod
337 + def _shift_precompute(
338 + tokens: torch.Tensor, eos_token_id: int
339 + ) -> tuple[torch.Tensor, torch.Tensor]:
340 + if tokens.dim() != 2:
341 + raise ValueError("tokens must be a 2D tensor")
342 + batch_size, seq_len = tokens.shape
343 + positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64)
344 + eos_positions = torch.where(tokens == eos_token_id, positions, -1)
345 + previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
346 + previous_eos = torch.cat(
347 + [
348 + eos_positions.new_full((batch_size, 1), -1),
349 + previous_eos_inclusive[:, :-1],
350 + ],
351 + dim=1,
352 + )
353 + return positions, positions.unsqueeze(0) - previous_eos - 1
354 +
355 + @staticmethod
356 + def _shift_apply(
357 + tokens: torch.Tensor,
358 + positions: torch.Tensor,
359 + position_in_segment: torch.Tensor,
360 + shift: int,
361 + eos_token_id: int,
362 + ) -> torch.Tensor:
363 + if shift == 0:
364 + return tokens
365 + source = positions - shift
366 + gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1)
367 + shifted = tokens.gather(1, gather_indices)
368 + valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift)
369 + return torch.where(valid, shifted, tokens.new_full((), eos_token_id))
370 +
371 + def compute_ngram_ids(
372 + self,
373 + input_ids: torch.Tensor,
374 + query_start_loc: torch.Tensor,
375 + ngram_context: torch.Tensor,
376 + ) -> torch.Tensor:
377 + """Compute n-gram embedding indices for the current request layout."""
378 + input_ids = input_ids.reshape(-1).long()
379 + query_start_loc = query_start_loc.long()
380 + num_reqs = query_start_loc.numel() - 1
381 + num_tokens = input_ids.shape[0]
382 + if num_tokens > self.positions_buffer.numel():
383 + raise ValueError(
384 + f"PLE received {num_tokens} tokens, but its workspace supports "
385 + f"at most {self.positions_buffer.numel()}"
386 + )
387 + if num_reqs > self.padded_buffer.shape[0]:
388 + raise ValueError(
389 + f"PLE received {num_reqs} requests, but its workspace supports "
390 + f"at most {self.padded_buffer.shape[0]}"
391 + )
392 +
393 + positions = self.positions_buffer[:num_tokens]
394 + packed = self.padded_buffer[:num_reqs, :num_tokens]
395 + packed.fill_(self.eos_token_id)
396 + request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1
397 + request_indices.clamp_(max=num_reqs - 1)
398 + columns = (positions - query_start_loc[request_indices]).clamp(
399 + 0, packed.shape[1] - 1
400 + )
401 + packed[request_indices, columns] = input_ids
402 + ngram_context = ngram_context[:num_reqs].to(
403 + device=input_ids.device, dtype=torch.long
404 + )
405 +
406 + context = torch.cat([ngram_context, packed], dim=-1)
407 + positions_2d, position_in_segment = self._shift_precompute(
408 + context, self.eos_token_id
409 + )
410 + shifted = [context]
411 + for shift in range(1, self.ngram_size):
412 + shifted.append(
413 + self._shift_apply(
414 + context,
415 + positions_2d,
416 + position_in_segment,
417 + shift,
418 + self.eos_token_id,
419 + )
420 + )
421 + adjusted_columns = columns + self.ngram_size - 1
422 + id_blocks = []
423 + for ngram in range(2, self.ngram_size + 1):
424 + start = (ngram - 2) * self.heads_per_ngram
425 + end = start + self.heads_per_ngram
426 + mixed = shifted[0] * self.layer_multipliers[0]
427 + for index in range(1, ngram):
428 + mixed = torch.bitwise_xor(
429 + mixed, shifted[index] * self.layer_multipliers[index]
430 + )
431 + sizes = self.ngram_heads_vocab_sizes[start:end]
432 + offsets = self.ngram_heads_offsets[start:end]
433 + ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets
434 + id_blocks.append(ids[request_indices, adjusted_columns])
435 + return torch.cat(id_blocks, dim=-1)
436 +
437 + def forward(
438 + self,
439 + input_ids: torch.Tensor,
440 + query_start_loc: torch.Tensor,
441 + ngram_context: torch.Tensor,
442 + ) -> torch.Tensor:
443 + ngram_ids = input_ids.new_empty(
444 + (input_ids.shape[0], self.ngram_heads),
445 + dtype=torch.long,
446 + )
447 + # Keep num_reqs-dependent ID generation outside PIECEWISE CUDA graphs,
448 + # which dispatch only on the padded token count.
449 + torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids(
450 + input_ids,
451 + query_start_loc,
452 + ngram_context,
453 + ngram_ids,
454 + self.layer_name,
455 + )
456 + return self.ngram_embedding(ngram_ids).flatten(-2)
457 +
458 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
459 + """Load hash buffers and checkpoint-split embedding rows."""
460 +
461 + persistent_buffers = {
462 + "layer_multipliers": self.layer_multipliers,
463 + "ngram_heads_offsets": self.ngram_heads_offsets,
464 + "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes,
465 + }
466 + loaded: set[str] = set()
467 + regular_weights: list[tuple[str, torch.Tensor]] = []
468 + shard_prefix = "ngram_embedding.shard_"
469 +
470 + for name, loaded_weight in weights:
471 + leaf_name = name.rsplit(".", 1)[-1]
472 + if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup":
473 + continue
474 + if name in persistent_buffers:
475 + buffer = persistent_buffers[name]
476 + if buffer.shape != loaded_weight.shape:
477 + raise ValueError(
478 + f"Shape mismatch for {name}: expected "
479 + f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}"
480 + )
481 + buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype))
482 + loaded.add(name)
483 + continue
484 + if name.startswith(shard_prefix) and name.endswith(".weight"):
485 + shard_text = name[len(shard_prefix) : -len(".weight")]
486 + if not shard_text.isdigit():
487 + regular_weights.append((name, loaded_weight))
488 + continue
489 + shard_index = int(shard_text)
490 + if shard_index >= self.split_ngram_parts:
491 + raise ValueError(
492 + f"PLE embedding shard index {shard_index} exceeds "
493 + f"split_ngram_parts={self.split_ngram_parts}"
494 + )
495 + embedding = self.ngram_embedding
496 + shard_size = (
497 + embedding.org_vocab_size + self.split_ngram_parts - 1
498 + ) // self.split_ngram_parts
499 + checkpoint_start = shard_index * shard_size
500 + expected_rows = max(
501 + 0,
502 + min(shard_size, embedding.org_vocab_size - checkpoint_start),
503 + )
504 + expected_shape = (expected_rows, embedding.embedding_dim)
505 + if tuple(loaded_weight.shape) != expected_shape:
506 + raise ValueError(
507 + f"Shape mismatch for PLE embedding shard {shard_index}: "
508 + f"expected {expected_shape}, got "
509 + f"{tuple(loaded_weight.shape)}"
510 + )
511 + embedding.weight.weight_loader(
512 + embedding.weight,
513 + loaded_weight,
514 + checkpoint_start=checkpoint_start,
515 + )
516 + loaded.add("ngram_embedding.weight")
517 + continue
518 + regular_weights.append((name, loaded_weight))
519 +
520 + if regular_weights:
521 + loaded.update(AutoWeightsLoader(self).load_weights(regular_weights))
522 + return loaded
523 +
524 +
525 +class Qwen4ExpPLELayer(nn.Module, MambaBase):
526 + def __init__(
527 + self,
528 + config: Qwen4ExpTextConfig,
529 + vllm_config: VllmConfig,
530 + layer_idx: int = 0,
531 + ple_dense_layer_id: int | None = None,
532 + prefix: str = "",
533 + ) -> None:
534 + super().__init__()
535 + model_config = vllm_config.model_config
536 + cache_config = vllm_config.cache_config
537 + quant_config = vllm_config.quant_config
538 + self.model_config: ModelConfig = model_config
539 + self.cache_config: CacheConfig = cache_config
540 + self.layer_idx = layer_idx
541 + self.ple_dense_layer_id = (
542 + int(ple_dense_layer_id)
543 + if ple_dense_layer_id is not None
544 + else int(layer_idx)
545 + )
546 + self.prefix = prefix
547 + self.hidden_size = int(config.hidden_size)
548 + self.hc_count = config.hc_count
549 + self.hc_hidden_size = self.hidden_size * self.hc_count
550 + self.conv_kernel_size = int(config.ple_conv_kernel_size)
551 + self.short_conv_dilation = int(config.ngram_size)
552 + self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation
553 + self.num_spec_tokens = vllm_config.num_speculative_tokens
554 + self.activation = "silu"
555 + self.ple_embedding = Qwen4ExpNGramEmbedding(
556 + config,
557 + int(config.ple_embed_dim),
558 + self.ple_dense_layer_id,
559 + vllm_config.scheduler_config.max_num_batched_tokens,
560 + vllm_config.scheduler_config.max_num_seqs,
561 + prefix=f"{prefix}.ple_embedding",
562 + layer_name=prefix,
563 + quant_config=quant_config,
564 + params_dtype=model_config.dtype,
565 + )
566 + self.key_proj = ReplicatedLinear(
567 + int(config.ple_embed_dim),
568 + self.hc_hidden_size,
569 + bias=False,
570 + quant_config=quant_config,
571 + prefix=f"{prefix}.key_proj",
572 + )
573 + self.value_proj = ReplicatedLinear(
574 + int(config.ple_embed_dim),
575 + self.hidden_size,
576 + bias=False,
577 + quant_config=quant_config,
578 + prefix=f"{prefix}.value_proj",
579 + )
580 + norm_args = (
581 + self.hc_hidden_size,
582 + config.rms_norm_eps,
583 + self.hidden_size,
584 + model_config.dtype,
585 + )
586 + self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args)
587 + self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args)
588 + self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args)
589 + self.conv1d = nn.Conv1d(
590 + self.hc_hidden_size,
591 + self.hc_hidden_size,
592 + self.conv_kernel_size,
593 + groups=self.hc_hidden_size,
594 + padding=self.conv_state_len,
595 + dilation=self.short_conv_dilation,
596 + bias=False,
597 + dtype=model_config.dtype,
598 + )
599 + nn.init.zeros_(self.conv1d.weight)
600 + self.conv1d.weight._no_reinit = True
601 + self.kv_cache = (torch.tensor([]),)
602 + compilation_config = get_current_vllm_config().compilation_config
603 + if prefix in compilation_config.static_forward_context:
604 + raise ValueError(f"Duplicate layer name: {prefix}")
605 + compilation_config.static_forward_context[prefix] = self
606 +
607 + def _get_embedding_weight_scale(self) -> torch.Tensor | None:
608 + embedding = getattr(self.ple_embedding, "ngram_embedding", None)
609 + return getattr(embedding, "weight_scale", None)
610 +
611 + def _dequantize_embeddings(
612 + self,
613 + embeddings: torch.Tensor,
614 + output_dtype: torch.dtype,
615 + ) -> torch.Tensor:
616 + """Dequantize PLE lookup output."""
617 +
618 + if not is_fp8(embeddings):
619 + return embeddings
620 + weight_scale = self._get_embedding_weight_scale()
621 + if weight_scale is None:
622 + raise RuntimeError("FP8 PLE embedding is missing its global scale")
623 + if weight_scale.device != embeddings.device:
624 + raise RuntimeError("FP8 PLE embedding scale must be on the output device")
625 + return embeddings.to(output_dtype) * weight_scale.to(output_dtype)
626 +
627 + @property
628 + def mamba_type(self) -> MambaAttentionBackendEnum:
629 + return MambaAttentionBackendEnum.SHORT_CONV
630 +
631 + @property
632 + def is_kv_cache_tp_replicated(self) -> bool:
633 + return True
634 +
635 + def get_attn_backend(self) -> type[PleShortConvAttentionBackend]:
636 + return PleShortConvAttentionBackend
637 +
638 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:
639 + return MambaStateDtypeCalculator.short_conv_state_dtype(
640 + self.model_config.dtype, self.cache_config.mamba_cache_dtype
641 + )
642 +
643 + def get_state_shape(self) -> Sequence[tuple[int, ...]]:
644 + return MambaStateShapeCalculator.short_conv_state_shape(
645 + tp_world_size=1,
646 + intermediate_size=self.hc_hidden_size,
647 + conv_kernel=self.conv_state_len + 1,
648 + num_spec=self.num_spec_tokens,
649 + )
650 +
651 + def _apply_norm(
652 + self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor
653 + ) -> torch.Tensor:
654 + shape = hidden_states.shape
655 + return norm(hidden_states.flatten(-2)).reshape(shape)
656 +
657 + def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor:
658 + # Profiling / CUDA graph capture only; conv state is not updated.
659 + inputs_t = inputs.transpose(0, 1).unsqueeze(0)
660 + output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)]
661 + return F.silu(output).squeeze(0).transpose(0, 1)
662 +
663 + def _short_conv_dilated_decode_batched(
664 + self,
665 + x_d: torch.Tensor,
666 + conv_state: torch.Tensor,
667 + conv_weights: torch.Tensor,
668 + state_indices_tensor_d: torch.Tensor,
669 + has_initial_states_d: torch.Tensor | None,
670 + ) -> torch.Tensor:
671 + state_indices = state_indices_tensor_d.to(
672 + device=conv_state.device, dtype=torch.int64
673 + )
674 + # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to
675 + # slot 0 for a safe gather, then zero output and skip write-back.
676 + valid_state = state_indices != NULL_BLOCK_ID
677 + state_indices = torch.where(
678 + valid_state, state_indices, torch.zeros_like(state_indices)
679 + )
680 + if has_initial_states_d is None:
681 + has_initial_state = valid_state
682 + else:
683 + if has_initial_states_d.numel() < state_indices_tensor_d.numel():
684 + raise ValueError(
685 + "has_initial_states_d size mismatch: "
686 + f"got {has_initial_states_d.numel()}, "
687 + f"need >= {state_indices_tensor_d.numel()}."
688 + )
689 + has_initial_state = has_initial_states_d[
690 + : state_indices_tensor_d.numel()
691 + ].to(device=conv_state.device, dtype=torch.bool)
692 + has_initial_state = has_initial_state & valid_state
693 +
694 + cached_state = conv_state.index_select(0, state_indices)
695 + state = cached_state[..., : self.conv_state_len].to(x_d.dtype)
696 + if self.conv_state_len > 0:
697 + initial_state = torch.where(
698 + has_initial_state.view(-1, 1, 1),
699 + state,
700 + torch.zeros_like(state),
701 + )
702 + history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1)
703 + else:
704 + history = x_d.unsqueeze(-1)
705 +
706 + conv_output = F.conv1d(
707 + history,
708 + conv_weights.unsqueeze(1).contiguous(),
709 + groups=history.size(1),
710 + dilation=self.short_conv_dilation,
711 + ).squeeze(-1)
712 + output = F.silu(conv_output)
713 + output = output * valid_state.view(-1, 1).to(output.dtype)
714 +
715 + if self.conv_state_len > 0:
716 + next_state = history[..., -self.conv_state_len :]
717 + # Padded rows are remapped to the reserved null slot. Preserve its
718 + # existing value while writing the new states for valid rows.
719 + existing_base_state = cached_state[..., : self.conv_state_len]
720 + safe_next_state = torch.where(
721 + valid_state.view(-1, 1, 1),
722 + next_state.to(conv_state.dtype),
723 + existing_base_state,
724 + )
725 + cached_state[..., : self.conv_state_len] = safe_next_state
726 + conv_state.index_copy_(0, state_indices, cached_state)
727 +
728 + return output
729 +
730 + def _short_conv_dilated_prefill_batched(
731 + self,
732 + x_p: torch.Tensor,
733 + metadata: PleShortConvAttentionMetadata,
734 + conv_state: torch.Tensor,
735 + conv_weights: torch.Tensor,
736 + state_indices_tensor_p: torch.Tensor,
737 + num_prefills: int,
738 + num_decode_tokens: int,
739 + num_prefill_tokens: int,
740 + ) -> torch.Tensor:
741 + # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill)
742 + # requests and equals ``query_start_loc`` when spec-decode is inactive.
743 + non_spec_query_start_loc = metadata.non_spec_query_start_loc
744 + if non_spec_query_start_loc is None:
745 + raise ValueError("query_start_loc is required for prefill short-conv")
746 + query_start_loc_p = (
747 + non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens
748 + )
749 + # The metadata builder guarantees that the prefill query offsets start
750 + # at 0 and end at num_prefill_tokens. Avoid reading those values here,
751 + # since doing so would force a device-to-host synchronization.
752 + has_initial_states_p = metadata.has_initial_states_p
753 + if has_initial_states_p is None:
754 + raise ValueError("has_initial_states_p is required for prefill short-conv")
755 +
756 + output = torch.empty_like(x_p)
757 + q_starts = query_start_loc_p.to(torch.int64)
758 + if state_indices_tensor_p.numel() < num_prefills:
759 + raise ValueError(
760 + "state_indices_tensor_p size mismatch: "
761 + f"got {state_indices_tensor_p.numel()}, "
762 + f"need >= {num_prefills}."
763 + )
764 + if has_initial_states_p.numel() < num_prefills:
765 + raise ValueError(
766 + "has_initial_states_p size mismatch: "
767 + f"got {has_initial_states_p.numel()}, "
768 + f"need >= {num_prefills}."
769 + )
770 + if num_prefills == 0 or x_p.numel() == 0:
771 + return output
772 + lengths = q_starts[1:] - q_starts[:-1]
773 + # Use the CPU-computed packing width from the metadata builder instead
774 + # of synchronizing on lengths.max().
775 + max_len = metadata.max_prefill_query_len
776 + if max_len <= 0:
777 + return output
778 +
779 + hidden_size = x_p.shape[1]
780 + positions = torch.arange(
781 + num_prefill_tokens, device=x_p.device, dtype=torch.int64
782 + )
783 + req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
784 + col_indices = positions - q_starts[req_indices]
785 +
786 + packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size))
787 + packed_tokens[req_indices, col_indices] = x_p
788 + packed_tokens = packed_tokens.transpose(1, 2).contiguous()
789 +
790 + state_indices = state_indices_tensor_p[:num_prefills].to(
791 + device=conv_state.device, dtype=torch.int64
792 + )
793 + valid_state = state_indices != NULL_BLOCK_ID
794 + state_indices = torch.where(
795 + valid_state, state_indices, torch.zeros_like(state_indices)
796 + )
797 + has_initial = has_initial_states_p[:num_prefills].to(
798 + device=conv_state.device, dtype=torch.bool
799 + )
800 + if self.conv_state_len > 0:
801 + if conv_state.shape[0] == 0:
802 + state = conv_state.new_zeros(
803 + (num_prefills, hidden_size, self.conv_state_len),
804 + dtype=x_p.dtype,
805 + )
806 + else:
807 + state = conv_state.index_select(0, state_indices)[
808 + ..., : self.conv_state_len
809 + ].to(x_p.dtype)
810 + use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1)
811 + initial_state = torch.where(
812 + use_initial_mask,
813 + state,
814 + torch.zeros_like(state),
815 + )
816 + history = torch.cat((initial_state, packed_tokens), dim=-1)
817 + else:
818 + history = packed_tokens
819 +
820 + conv_output = F.conv1d(
821 + history,
822 + conv_weights.unsqueeze(1).contiguous(),
823 + groups=history.size(1),
824 + dilation=self.short_conv_dilation,
825 + )
826 + conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
827 +
828 + token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64)
829 + valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1)
830 + valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view(
831 + num_prefills, 1
832 + )
833 + conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0)
834 + output.copy_(conv_output[req_indices, col_indices])
835 +
836 + if self.conv_state_len > 0 and conv_state.shape[0] > 0:
837 + state_starts = lengths.to(device=history.device, dtype=torch.int64).view(
838 + num_prefills, 1, 1
839 + )
840 + state_offsets = torch.arange(
841 + self.conv_state_len, device=history.device, dtype=torch.int64
842 + ).view(1, 1, self.conv_state_len)
843 + next_state = history.gather(
844 + dim=2,
845 + index=(state_starts + state_offsets).expand(-1, history.size(1), -1),
846 + )
847 + # Write back without a host synchronization. Valid, non-empty rows
848 + # receive their new state; padding and zero-length rows keep the
849 + # current cache value.
850 + existing_state = conv_state.index_select(0, state_indices)
851 + existing_base_state = existing_state[..., : self.conv_state_len]
852 + update_mask = valid_state & (lengths.to(device=conv_state.device) > 0)
853 + safe_next_state = torch.where(
854 + update_mask.view(num_prefills, 1, 1),
855 + next_state.to(conv_state.dtype),
856 + existing_base_state,
857 + )
858 + existing_state[..., : self.conv_state_len] = safe_next_state
859 + conv_state.index_copy_(0, state_indices, existing_state)
860 + return output
861 +
862 + def _short_conv_dilated_spec_batched(
863 + self,
864 + x_spec: torch.Tensor,
865 + conv_state: torch.Tensor,
866 + conv_weights: torch.Tensor,
867 + spec_state_indices_tensor: torch.Tensor,
868 + spec_query_start_loc: torch.Tensor,
869 + num_accepted_tokens: torch.Tensor,
870 + spec_query_len: int,
871 + ) -> torch.Tensor:
872 + """Dilated short-conv for speculative-decode (MTP) requests.
873 +
874 + Each spec request feeds multiple (draft + 1) query tokens. The conv
875 + outputs are computed causally after rolling back the previous draft
876 + state by ``num_accepted_tokens - 1``. The current candidate inputs stay
877 + in the extended cache for the next forward, matching
878 + ``causal_conv1d_update``.
879 +
880 + ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query
881 + length and is a Python int, so no host synchronization is needed; this
882 + keeps the path safe for full CUDA-graph capture/replay where the buffers
883 + are padded at the request level.
884 + """
885 + num_reqs = spec_state_indices_tensor.numel()
886 + hidden_size = x_spec.size(-1)
887 + # Use a fixed packing width instead of synchronizing on lengths.max().
888 + max_len = spec_query_len
889 + # Full CUDA graphs can pad these buffers. Only the first num_reqs
890 + # accepted-token counts belong to actual speculative requests.
891 + num_accepted_tokens = num_accepted_tokens[:num_reqs]
892 + q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64)
893 + # Keep the number of real speculative tokens on the device.
894 + total_real_tokens = q_starts[num_reqs]
895 +
896 + state_indices = spec_state_indices_tensor.to(
897 + device=conv_state.device, dtype=torch.int64
898 + )
899 + valid_state = state_indices != NULL_BLOCK_ID
900 + state_indices = torch.where(
901 + valid_state, state_indices, torch.zeros_like(state_indices)
902 + )
903 + positions = torch.arange(
904 + x_spec.size(0), device=x_spec.device, dtype=torch.int64
905 + )
906 + # Route graph-padded token rows to the discarded dummy request so that
907 + # they cannot overwrite real packed data.
908 + req_indices = torch.searchsorted(q_starts[1:], positions, right=True)
909 + valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs)
910 + clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0))
911 + col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1)
912 + pack_req_indices = torch.where(
913 + valid_tokens,
914 + clamped_req_indices,
915 + torch.full_like(req_indices, num_reqs),
916 + )
917 + pack_col_indices = torch.where(
918 + valid_tokens, col_indices, torch.zeros_like(col_indices)
919 + )
920 +
921 + # The last request row is the dummy sink for graph padding.
922 + packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size))
923 + packed[pack_req_indices, pack_col_indices] = x_spec
924 + packed = packed.transpose(1, 2).contiguous()
925 +
926 + if self.conv_state_len > 0:
927 + cached_state = conv_state.index_select(0, state_indices)
928 + rollback_offsets = num_accepted_tokens.to(
929 + device=conv_state.device, dtype=torch.int64
930 + ).sub(1)
931 + rollback_offsets = torch.where(
932 + valid_state,
933 + rollback_offsets.clamp_(0, max_len - 1),
934 + torch.zeros_like(rollback_offsets),
935 + )
936 + state_offsets = torch.arange(
937 + self.conv_state_len, device=conv_state.device, dtype=torch.int64
938 + ).view(1, 1, self.conv_state_len)
939 + rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets
940 + state = cached_state.gather(
941 + 2, rollback_indices.expand(-1, hidden_size, -1)
942 + ).to(x_spec.dtype)
943 + state = torch.where(
944 + valid_state.view(num_reqs, 1, 1),
945 + state,
946 + torch.zeros_like(state),
947 + )
948 + # Append a zeroed dummy-row state to match the [num_reqs + 1] pack.
949 + dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len))
950 + state_full = torch.cat((state, dummy_state), dim=0)
951 + history = torch.cat((state_full, packed), dim=-1)
952 + else:
953 + history = packed
954 +
955 + conv_output = F.conv1d(
956 + history,
957 + conv_weights.unsqueeze(1).contiguous(),
958 + groups=history.size(1),
959 + dilation=self.short_conv_dilation,
960 + )
961 + conv_output = F.silu(conv_output).transpose(1, 2).contiguous()
962 +
963 + output = conv_output[pack_req_indices, pack_col_indices]
964 + output = output * valid_tokens.view(-1, 1).to(output.dtype)
965 +
966 + # Keep all current candidate inputs in the extended state. On the next
967 + # target forward, ``num_accepted_tokens - 1`` selects the rollback
968 + # window before processing the newly scheduled tokens.
969 + if self.conv_state_len > 0:
970 + state_capacity = self.conv_state_len + max_len - 1
971 + if conv_state.size(-1) < state_capacity:
972 + raise RuntimeError(
973 + "PLE short-conv cache cannot retain speculative tokens: "
974 + f"got {conv_state.size(-1)}, need {state_capacity}."
975 + )
976 + candidate_state = history[:num_reqs, :, 1 : state_capacity + 1]
977 + query_lengths = q_starts[1:] - q_starts[:-1]
978 + state_positions = torch.arange(
979 + state_capacity, device=history.device, dtype=torch.int64
980 + ).view(1, 1, state_capacity)
981 + update_lengths = (self.conv_state_len + query_lengths - 1).view(
982 + num_reqs, 1, 1
983 + )
984 + update_mask = valid_state.view(num_reqs, 1, 1) & (
985 + state_positions < update_lengths
986 + )
987 + existing_state = cached_state[..., :state_capacity]
988 + next_state = torch.where(
989 + update_mask,
990 + candidate_state.to(conv_state.dtype),
991 + existing_state,
992 + )
993 + cached_state[..., :state_capacity] = next_state
994 + conv_state.index_copy_(0, state_indices, cached_state)
995 +
996 + return output
997 +
998 + def _short_conv_dilated_dispatch(
999 + self,
1000 + inputs: torch.Tensor,
1001 + metadata: PleShortConvAttentionMetadata,
1002 + conv_state: torch.Tensor,
1003 + conv_weights: torch.Tensor,
1004 + ) -> torch.Tensor:
1005 + num_prefills = metadata.num_prefills
1006 + num_decodes = metadata.num_decodes
1007 + num_decode_tokens = metadata.num_decode_tokens
1008 + num_prefill_tokens = metadata.num_prefill_tokens
1009 + has_prefill = num_prefills > 0
1010 + has_decode = num_decodes > 0
1011 + has_spec = metadata.spec_sequence_masks is not None
1012 + x = inputs[: metadata.num_actual_tokens]
1013 +
1014 + # Split spec / non-spec tokens.
1015 + if has_spec:
1016 + if has_prefill or has_decode:
1017 + assert metadata.spec_token_indx is not None
1018 + assert metadata.non_spec_token_indx is not None
1019 + x_spec = x.index_select(0, metadata.spec_token_indx.long())
1020 + x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long())
1021 + else:
1022 + x_spec = x
1023 + x_non_spec = None
1024 + else:
1025 + x_spec = None
1026 + x_non_spec = x
1027 +
1028 + spec_output = None
1029 + # 1. Run the multi-query speculative-decode part.
1030 + if has_spec:
1031 + assert metadata.spec_state_indices_tensor is not None
1032 + assert metadata.spec_query_start_loc is not None
1033 + assert metadata.num_accepted_tokens is not None
1034 + spec_output = self._short_conv_dilated_spec_batched(
1035 + x_spec=x_spec,
1036 + conv_state=conv_state,
1037 + conv_weights=conv_weights,
1038 + spec_state_indices_tensor=metadata.spec_state_indices_tensor[
1039 + : metadata.num_spec_decodes
1040 + ],
1041 + spec_query_start_loc=metadata.spec_query_start_loc,
1042 + num_accepted_tokens=metadata.num_accepted_tokens,
1043 + spec_query_len=metadata.spec_query_len,
1044 + )
1045 +
1046 + # 2. Run regular decode and prefill requests.
1047 + conv_out_non_spec = None
1048 + state_indices_tensor = metadata.state_indices_tensor
1049 + if x_non_spec is not None:
1050 + assert state_indices_tensor is not None
1051 + if has_prefill:
1052 + state_indices_tensor_d, state_indices_tensor_p = torch.split(
1053 + state_indices_tensor,
1054 + [num_decodes, num_prefills],
1055 + dim=0,
1056 + )
1057 + x_d, x_p = torch.split(
1058 + x_non_spec,
1059 + [num_decode_tokens, num_prefill_tokens],
1060 + dim=0,
1061 + )
1062 + non_spec_parts: list[torch.Tensor] = []
1063 + if has_decode:
1064 + non_spec_parts.append(
1065 + self._short_conv_dilated_decode_batched(
1066 + x_d=x_d,
1067 + conv_state=conv_state,
1068 + conv_weights=conv_weights,
1069 + state_indices_tensor_d=state_indices_tensor_d,
1070 + has_initial_states_d=metadata.has_initial_states_d,
1071 + )
1072 + )
1073 + non_spec_parts.append(
1074 + self._short_conv_dilated_prefill_batched(
1075 + x_p=x_p,
1076 + metadata=metadata,
1077 + conv_state=conv_state,
1078 + conv_weights=conv_weights,
1079 + state_indices_tensor_p=state_indices_tensor_p,
1080 + num_prefills=num_prefills,
1081 + num_decode_tokens=num_decode_tokens,
1082 + num_prefill_tokens=num_prefill_tokens,
1083 + )
1084 + )
1085 + conv_out_non_spec = torch.vstack(non_spec_parts)
1086 + else:
1087 + conv_out_non_spec = self._short_conv_dilated_decode_batched(
1088 + x_d=x_non_spec,
1089 + conv_state=conv_state,
1090 + conv_weights=conv_weights,
1091 + state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)],
1092 + has_initial_states_d=metadata.has_initial_states_d,
1093 + )
1094 +
1095 + # 3. Merge both parts back into the original token order.
1096 + if has_spec and conv_out_non_spec is not None:
1097 + assert metadata.spec_token_indx is not None
1098 + assert metadata.non_spec_token_indx is not None
1099 + assert spec_output is not None
1100 + output = x.new_empty((metadata.num_actual_tokens, x.size(-1)))
1101 + output.index_copy_(0, metadata.spec_token_indx, spec_output)
1102 + output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec)
1103 + return output
1104 + elif has_spec:
1105 + assert spec_output is not None
1106 + return spec_output
1107 + if conv_out_non_spec is None:
1108 + return x
1109 + return conv_out_non_spec
1110 +
1111 + def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor:
1112 + forward_context = get_forward_context()
1113 + attn_metadata = forward_context.attn_metadata
1114 + if attn_metadata is None:
1115 + return self._short_conv_fallback(inputs)
1116 +
1117 + if not isinstance(attn_metadata, dict):
1118 + raise RuntimeError(
1119 + "PLE short-conv expects per-layer attention metadata dict "
1120 + f"during inference, got {type(attn_metadata).__name__}."
1121 + )
1122 +
1123 + layer_attn_metadata = attn_metadata.get(self.prefix)
1124 + if layer_attn_metadata is None:
1125 + # MRV2 omits Mamba-family metadata during profile warmup.
1126 + return self._short_conv_fallback(inputs)
1127 + if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata):
1128 + raise TypeError(
1129 + "Expected PleShortConvAttentionMetadata for layer "
1130 + f"'{self.prefix}', got "
1131 + f"{type(layer_attn_metadata).__name__}."
1132 + )
1133 +
1134 + conv_state = self.kv_cache[0]
1135 + if not is_conv_state_dim_first():
1136 + conv_state = conv_state.transpose(-1, -2)
1137 + conv_weights = self.conv1d.weight.squeeze(1)
1138 +
1139 + state_capacity = self.conv_state_len + self.num_spec_tokens
1140 + if state_capacity > 0:
1141 + if conv_state.size(-1) < state_capacity:
1142 + raise RuntimeError(
1143 + "PLE short-conv cache is smaller than expected for "
1144 + f"dilated convolution: got {conv_state.size(-1)}, "
1145 + f"expect at least {state_capacity}."
1146 + )
1147 + conv_state = conv_state[..., -state_capacity:]
1148 + return self._short_conv_dilated_dispatch(
1149 + inputs,
1150 + layer_attn_metadata,
1151 + conv_state,
1152 + conv_weights.to(dtype=inputs.dtype),
1153 + )
1154 +
1155 + def forward(
1156 + self,
1157 + hidden_states: torch.Tensor,
1158 + input_ids: torch.Tensor,
1159 + query_start_loc: torch.Tensor,
1160 + ngram_context: torch.Tensor,
1161 + ) -> torch.Tensor:
1162 + input_ids = input_ids.reshape(-1)
1163 + if input_ids.shape[0] != hidden_states.shape[0]:
1164 + raise ValueError(
1165 + "PLE expects input_ids and hidden_states to have the same "
1166 + f"token length, got {input_ids.shape[0]} and "
1167 + f"{hidden_states.shape[0]}"
1168 + )
1169 + embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context)
1170 + embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype)
1171 + key, _ = self.key_proj(embeddings)
1172 + value, _ = self.value_proj(embeddings)
1173 + token_count = hidden_states.shape[0]
1174 + key = key.reshape(token_count, self.hc_count, self.hidden_size)
1175 + query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size)
1176 + key = self._apply_norm(self.norm_key, key)
1177 + query = self._apply_norm(self.norm_query, query)
1178 + gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size)
1179 + gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt())
1180 + gated_value = gate * value.unsqueeze(-2)
1181 + normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2)
1182 + conv_output = torch.zeros_like(normalized)
1183 + torch.ops.vllm.qwen4_exp_ple_short_conv(
1184 + normalized,
1185 + conv_output,
1186 + self.prefix,
1187 + )
1188 + return gated_value.flatten(-2) + conv_output
1189 +
1190 +
1191 +def qwen4_exp_compute_ple_ngram_ids(
1192 + input_ids: torch.Tensor,
1193 + query_start_loc: torch.Tensor,
1194 + ngram_context: torch.Tensor,
1195 + output: torch.Tensor,
1196 + layer_name: str,
1197 +) -> None:
1198 + """Compute request-dependent PLE n-gram IDs outside piecewise graphs."""
1199 + layer = get_forward_context().no_compile_layers[layer_name]
1200 + ngram_ids = layer.ple_embedding.compute_ngram_ids(
1201 + input_ids,
1202 + query_start_loc,
1203 + ngram_context,
1204 + )
1205 + output.copy_(ngram_ids)
1206 +
1207 +
1208 +def qwen4_exp_compute_ple_ngram_ids_fake(
1209 + input_ids: torch.Tensor,
1210 + query_start_loc: torch.Tensor,
1211 + ngram_context: torch.Tensor,
1212 + output: torch.Tensor,
1213 + layer_name: str,
1214 +) -> None:
1215 + return
1216 +
1217 +
1218 +def qwen4_exp_ple_short_conv(
1219 + inputs: torch.Tensor,
1220 + output: torch.Tensor,
1221 + layer_name: str,
1222 +) -> None:
1223 + layer = get_forward_context().no_compile_layers[layer_name]
1224 + result = layer._short_conv(inputs)
1225 + output[: result.shape[0]].copy_(result)
1226 +
1227 +
1228 +def qwen4_exp_ple_short_conv_fake(
1229 + inputs: torch.Tensor,
1230 + output: torch.Tensor,
1231 + layer_name: str,
1232 +) -> None:
1233 + return
1234 +
1235 +
1236 +direct_register_custom_op(
1237 + op_name="qwen4_exp_compute_ple_ngram_ids",
1238 + op_func=qwen4_exp_compute_ple_ngram_ids,
1239 + mutates_args=["output"],
1240 + fake_impl=qwen4_exp_compute_ple_ngram_ids_fake,
1241 +)
1242 +
1243 +
1244 +direct_register_custom_op(
1245 + op_name="qwen4_exp_ple_short_conv",
1246 + op_func=qwen4_exp_ple_short_conv,
1247 + mutates_args=["output"],
1248 + fake_impl=qwen4_exp_ple_short_conv_fake,
1249 +)
1250 +
1251 +
1252 +__all__ = [
1253 + "Qwen4ExpNGramEmbedding",
1254 + "Qwen4ExpPLEGroupedNorm",
1255 + "Qwen4ExpPLELayer",
1256 +]
vllm/models/qwen4_exp/nvidia/qsa.py new
+487
@@ -0,0 +1,487 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""NVIDIA QSA owner with Triton kernels."""
4 +
5 +from __future__ import annotations
6 +
7 +from typing import ClassVar, cast
8 +
9 +import torch
10 +from torch import nn
11 +
12 +from vllm.config import VllmConfig
13 +from vllm.config.cache import CacheDType
14 +from vllm.distributed import get_tensor_model_parallel_world_size
15 +from vllm.forward_context import get_forward_context
16 +from vllm.model_executor.layers.attention.attention import (
17 + set_default_quant_scales,
18 +)
19 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
20 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm
21 +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
22 +from vllm.model_executor.layers.quantization import QuantizationConfig
23 +from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope
24 +from vllm.model_executor.models.qwen3_next import Qwen3NextAttention
25 +from vllm.platforms import current_platform
26 +from vllm.transformers_utils.configs.qwen4_exp import (
27 + Qwen4ExpTextConfig,
28 +)
29 +from vllm.utils.torch_utils import (
30 + LayerNameType,
31 + _encode_layer_name,
32 + _resolve_layer_name,
33 + canonicalize_singleton_dim_strides,
34 + direct_register_custom_op,
35 + kv_cache_dtype_str_to_dtype,
36 +)
37 +from vllm.v1.attention.backend import (
38 + AttentionBackend,
39 + AttentionCGSupport,
40 + AttentionType,
41 + MultipleOf,
42 +)
43 +from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
44 +from vllm.v1.attention.backends.flash_attn import (
45 + FlashAttentionBackend,
46 + FlashAttentionImpl,
47 + FlashAttentionMetadata,
48 + FlashAttentionMetadataBuilder,
49 +)
50 +from vllm.v1.kv_cache_interface import (
51 + FullAttentionSpec,
52 + KVCacheSpec,
53 + get_kv_quant_mode,
54 +)
55 +
56 +from ..common.qsa_cache import QSAForwardMetadata
57 +from . import model
58 +from .indexer_qsa import QSAIndexer
59 +
60 +
61 +class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
62 + """Flash metadata supporting uniform decode and target-verify graphs."""
63 +
64 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
65 +
66 +
67 +class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
68 + """FullAttentionSpec backend used by the merged QSA owner."""
69 +
70 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
71 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
72 +
73 + @staticmethod
74 + def get_name() -> str:
75 + return "QWEN4_EXP_QSA_TRITON"
76 +
77 + @staticmethod
78 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
79 + # QSA consumes manager pages directly and does not use FA4 paged attention.
80 + return [MultipleOf(16)]
81 +
82 + @staticmethod
83 + def get_impl_cls() -> type[Qwen4ExpQSAFlashAttentionImpl]:
84 + return Qwen4ExpQSAFlashAttentionImpl
85 +
86 + @staticmethod
87 + def get_builder_cls() -> type[Qwen4ExpQSAMetadataBuilder]:
88 + return Qwen4ExpQSAMetadataBuilder
89 +
90 + @classmethod
91 + def is_sparse(cls) -> bool:
92 + return True
93 +
94 + @classmethod
95 + def supports_kv_connector(cls) -> bool:
96 + return False
97 +
98 +
99 +class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl):
100 + """Run paged sparse GQA with the QSA Triton kernel."""
101 +
102 + supports_dcp: bool = False
103 + supports_pcp: bool = False
104 +
105 + def __init__(self, *args, **kwargs) -> None:
106 + super().__init__(*args, **kwargs)
107 + if not is_flash_attn_varlen_func_available():
108 + raise NotImplementedError("Qwen4Exp QSA requires FlashAttention")
109 + if self.dcp_world_size != 1:
110 + raise NotImplementedError(
111 + "Qwen4Exp QSA does not support decode context parallelism"
112 + )
113 + if self.kv_cache_dtype not in ("auto", "bfloat16"):
114 + raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
115 + self.supports_quant_query_input = False
116 +
117 + def forward_qsa(
118 + self,
119 + layer: torch.nn.Module,
120 + query: torch.Tensor,
121 + key: torch.Tensor,
122 + value: torch.Tensor,
123 + kv_cache: torch.Tensor,
124 + attn_metadata: FlashAttentionMetadata,
125 + output: torch.Tensor,
126 + token_to_req: torch.Tensor,
127 + output_scale: torch.Tensor | None = None,
128 + output_block_scale: torch.Tensor | None = None,
129 + ) -> torch.Tensor:
130 + del key, value
131 + if output_scale is not None or output_block_scale is not None:
132 + raise NotImplementedError("QSA does not support fused output quantization")
133 + if self.alibi_slopes is not None or self.sinks is not None:
134 + raise NotImplementedError("QSA does not support ALiBi or attention sinks")
135 + if self.sliding_window != (-1, -1):
136 + raise NotImplementedError("QSA does not support sliding-window attention")
137 +
138 + num_tokens = attn_metadata.num_actual_tokens
139 + output.zero_()
140 + if num_tokens == 0:
141 + return output
142 +
143 + topk_buffer = getattr(layer, "topk_indices_buffer", None)
144 + if topk_buffer is None:
145 + raise RuntimeError("QSA owner did not provide its top-k buffer")
146 + logical_indices = topk_buffer[:num_tokens]
147 + token_to_req = token_to_req[:num_tokens]
148 + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
149 + key_cache = canonicalize_singleton_dim_strides(key_cache)
150 + value_cache = canonicalize_singleton_dim_strides(value_cache)
151 + if key_cache.dtype != torch.bfloat16 or query.dtype != torch.bfloat16:
152 + raise NotImplementedError("Qwen4Exp QSA requires BF16 Q/K/V")
153 +
154 + from .ops.qsa import qsa_sparse_paged_attention
155 +
156 + qsa_sparse_paged_attention(
157 + query[:num_tokens],
158 + key_cache,
159 + value_cache,
160 + logical_indices,
161 + attn_metadata.block_table,
162 + token_to_req,
163 + output[:num_tokens],
164 + )
165 + return output
166 +
167 +
168 +class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase):
169 + """Merged Qwen full-attention owner with a QSA index side branch."""
170 +
171 + supports_dcp = False
172 +
173 + def __init__(
174 + self,
175 + *,
176 + vllm_config: VllmConfig,
177 + config: Qwen4ExpTextConfig,
178 + layer_id: int,
179 + quant_config: QuantizationConfig | None = None,
180 + reduce_results: bool = True,
181 + prefix: str = "",
182 + ) -> None:
183 + nn.Module.__init__(self)
184 + cache_config = vllm_config.cache_config
185 + model_config = vllm_config.model_config
186 + if cache_config is None:
187 + raise ValueError("Qwen4Exp QSA requires a paged KV cache")
188 + if model_config.dtype != torch.bfloat16:
189 + raise NotImplementedError("Qwen4Exp QSA currently requires BF16")
190 + if cache_config.cache_dtype not in ("auto", "bfloat16"):
191 + raise NotImplementedError("Qwen4Exp QSA requires a BF16 main KV cache")
192 + if getattr(quant_config, "kv_cache_scheme", None) is not None:
193 + raise NotImplementedError("Qwen4Exp QSA does not support KV quantization")
194 + parallel_config = vllm_config.parallel_config
195 + if (
196 + parallel_config.prefill_context_parallel_size > 1
197 + or parallel_config.decode_context_parallel_size > 1
198 + ):
199 + raise NotImplementedError(
200 + "Qwen4Exp QSA does not support context parallelism"
201 + )
202 + if not getattr(config, "is_causal", True):
203 + raise NotImplementedError("Qwen4Exp QSA requires causal decoder attention")
204 +
205 + self.config = config
206 + self.hidden_size = int(config.hidden_size)
207 + tp_size = get_tensor_model_parallel_world_size()
208 + self.total_num_heads = int(config.num_attention_heads)
209 + if self.total_num_heads % tp_size:
210 + raise ValueError("QSA attention heads must be divisible by TP size")
211 + self.num_heads = self.total_num_heads // tp_size
212 + self.total_num_kv_heads = int(config.num_key_value_heads)
213 + if self.total_num_kv_heads >= tp_size:
214 + if self.total_num_kv_heads % tp_size:
215 + raise ValueError("QSA KV heads must be divisible by TP size")
216 + elif tp_size % self.total_num_kv_heads:
217 + raise ValueError("TP size must be divisible by replicated QSA KV heads")
218 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
219 + self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads)
220 + self.q_size = self.num_heads * self.head_dim
221 + self.kv_size = self.num_kv_heads * self.head_dim
222 + self.scaling = self.head_dim**-0.5
223 + self.dual_chunk_attention_config = getattr(
224 + config, "dual_chunk_attention_config", None
225 + )
226 + if self.dual_chunk_attention_config is not None:
227 + raise NotImplementedError("Qwen4Exp QSA does not support dual-chunk RoPE")
228 + # Qwen4Exp full-attention checkpoints always pack a sigmoid output
229 + # gate next to Q, even when an inherited config default says otherwise.
230 + self.attn_output_gate = True
231 +
232 + self.qkv_proj = QKVParallelLinear(
233 + self.hidden_size,
234 + self.head_dim,
235 + self.total_num_heads * (1 + self.attn_output_gate),
236 + self.total_num_kv_heads,
237 + bias=False,
238 + quant_config=model.without_modelopt_fp4(quant_config),
239 + prefix=f"{prefix}.qkv_proj",
240 + )
241 + self.o_proj = RowParallelLinear(
242 + self.total_num_heads * self.head_dim,
243 + self.hidden_size,
244 + bias=False,
245 + reduce_results=reduce_results,
246 + quant_config=quant_config,
247 + prefix=f"{prefix}.o_proj",
248 + )
249 + self.rotary_emb = get_rope(
250 + head_size=self.head_dim,
251 + max_position=config.max_position_embeddings,
252 + rope_parameters=config.rope_parameters,
253 + )
254 + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
255 + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps)
256 +
257 + mm_config = model_config.multimodal_config
258 + text_only = mm_config is None or mm_config.language_model_only
259 + mrope_section = getattr(self.rotary_emb, "mrope_section", None)
260 + supports_mrope = bool(
261 + type(self.rotary_emb) is MRotaryEmbedding
262 + and mrope_section
263 + and len(mrope_section) == 3
264 + and sum(mrope_section) == self.rotary_emb.rotary_dim // 2
265 + and getattr(self.rotary_emb, "mrope_interleaved", False)
266 + )
267 + supports_dtype = getattr(self.rotary_emb, "dtype", None) in (
268 + torch.float16,
269 + torch.bfloat16,
270 + )
271 + self.use_fused_qk_norm_rope_gate = (
272 + self.attn_output_gate
273 + and getattr(self.rotary_emb, "is_neox_style", False)
274 + and current_platform.is_cuda()
275 + and supports_dtype
276 + and (text_only or supports_mrope)
277 + )
278 +
279 + self.layer_name = f"{prefix}.attn"
280 + self.attn_type = AttentionType.DECODER
281 + self.kv_cache_dtype = cache_config.cache_dtype
282 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
283 + self.kv_cache_dtype, model_config
284 + )
285 + if self.kv_cache_torch_dtype != torch.bfloat16:
286 + raise NotImplementedError("Qwen4Exp QSA requires BF16 cache storage")
287 + self.kv_sharing_target_layer_name = None
288 + self.kv_cache = torch.tensor([])
289 + set_default_quant_scales(self, register_buffer=True)
290 +
291 + self.attn_backend = Qwen4ExpQSAFlashAttentionBackend
292 + self.impl = Qwen4ExpQSAFlashAttentionImpl(
293 + self.num_heads,
294 + self.head_dim,
295 + self.scaling,
296 + self.num_kv_heads,
297 + None,
298 + None,
299 + self.kv_cache_dtype,
300 + None,
301 + AttentionType.DECODER,
302 + None,
303 + )
304 + self.indexer = QSAIndexer(
305 + vllm_config=vllm_config,
306 + config=config,
307 + layer_id=layer_id,
308 + rotary_emb=self.rotary_emb,
309 + quant_config=quant_config,
310 + prefix=f"{prefix}.indexer",
311 + )
312 + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
313 + self.register_buffer(
314 + "topk_indices_buffer",
315 + torch.empty(
316 + max_tokens,
317 + self.indexer.output_width,
318 + dtype=torch.int32,
319 + ),
320 + persistent=False,
321 + )
322 +
323 + static_context = vllm_config.compilation_config.static_forward_context
324 + if self.layer_name in static_context:
325 + raise ValueError(f"Duplicate layer name: {self.layer_name}")
326 + static_context[self.layer_name] = self
327 +
328 + def get_attn_backend(self) -> type[AttentionBackend]:
329 + return self.attn_backend
330 +
331 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
332 + return FullAttentionSpec(
333 + block_size=vllm_config.cache_config.block_size,
334 + num_kv_heads=self.num_kv_heads,
335 + head_size=self.head_dim,
336 + head_size_v=self.head_dim,
337 + dtype=self.kv_cache_torch_dtype,
338 + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
339 + )
340 +
341 + def _run_qsa(
342 + self,
343 + hidden_states: torch.Tensor,
344 + positions: torch.Tensor,
345 + query: torch.Tensor,
346 + key: torch.Tensor,
347 + value: torch.Tensor,
348 + output: torch.Tensor,
349 + ) -> None:
350 + metadata = get_forward_context().attn_metadata
351 + if isinstance(metadata, list):
352 + metadata = metadata[0]
353 + if not isinstance(metadata, dict):
354 + output.zero_()
355 + return
356 + main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name])
357 + if self.kv_cache.numel() == 0:
358 + raise RuntimeError("QSA main K/V cache is not bound")
359 +
360 + num_tokens = main_metadata.num_actual_tokens
361 + side_metadata = cast(
362 + QSAForwardMetadata,
363 + metadata[self.indexer.raw_key_cache.prefix],
364 + )
365 + if side_metadata.num_actual_tokens != num_tokens:
366 + raise RuntimeError("QSA main and side metadata token counts disagree")
367 + selected = self.indexer(
368 + hidden_states,
369 + positions,
370 + self.topk_indices_buffer[:num_tokens],
371 + )
372 + if selected.shape != (
373 + num_tokens,
374 + self.indexer.output_width,
375 + ):
376 + raise RuntimeError("QSA indexer returned an invalid selection shape")
377 + impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl)
378 + impl.do_kv_cache_update(
379 + self,
380 + key,
381 + value,
382 + self.kv_cache,
383 + main_metadata.slot_mapping,
384 + )
385 + impl.forward_qsa(
386 + self,
387 + query,
388 + key,
389 + value,
390 + self.kv_cache,
391 + main_metadata,
392 + output,
393 + token_to_req=side_metadata.token_to_req,
394 + )
395 +
396 + def forward(
397 + self,
398 + positions: torch.Tensor,
399 + hidden_states: torch.Tensor,
400 + ) -> torch.Tensor:
401 + qkv, _ = self.qkv_proj(hidden_states)
402 + q, k, v, gate = self._project_qkv_gate(qkv, positions)
403 + num_tokens = hidden_states.shape[0]
404 + query = q.view(num_tokens, self.num_heads, self.head_dim)
405 + key = k.view(num_tokens, self.num_kv_heads, self.head_dim)
406 + value = v.view(num_tokens, self.num_kv_heads, self.head_dim)
407 + attn_output = torch.empty_like(query)
408 + encoded_layer_name = _encode_layer_name(self.layer_name)
409 + if current_platform.opaque_attention_op():
410 + torch.ops.vllm.qwen4_exp_qsa_with_output(
411 + hidden_states,
412 + positions,
413 + query,
414 + key,
415 + value,
416 + attn_output,
417 + encoded_layer_name,
418 + )
419 + else:
420 + qwen4_exp_qsa_with_output(
421 + hidden_states,
422 + positions,
423 + query,
424 + key,
425 + value,
426 + attn_output,
427 + encoded_layer_name,
428 + )
429 + flat_output = attn_output.view(num_tokens, -1)
430 + if gate is not None:
431 + flat_output = flat_output * torch.sigmoid(gate)
432 + output, _ = self.o_proj(flat_output)
433 + return output
434 +
435 +
436 +def qwen4_exp_qsa_with_output(
437 + hidden_states: torch.Tensor,
438 + positions: torch.Tensor,
439 + query: torch.Tensor,
440 + key: torch.Tensor,
441 + value: torch.Tensor,
442 + output: torch.Tensor,
443 + layer_name: LayerNameType,
444 +) -> None:
445 + """Run the complete QSA state/update/attend transaction."""
446 +
447 + layer_name = _resolve_layer_name(layer_name)
448 + layer = get_forward_context().no_compile_layers[layer_name]
449 + if not isinstance(layer, Qwen4ExpQSAAttention):
450 + raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner")
451 + layer._run_qsa(
452 + hidden_states,
453 + positions,
454 + query,
455 + key,
456 + value,
457 + output,
458 + )
459 +
460 +
461 +def qwen4_exp_qsa_with_output_fake(
462 + hidden_states: torch.Tensor,
463 + positions: torch.Tensor,
464 + query: torch.Tensor,
465 + key: torch.Tensor,
466 + value: torch.Tensor,
467 + output: torch.Tensor,
468 + layer_name: LayerNameType,
469 +) -> None:
470 + del hidden_states, positions, query, key, value, output, layer_name
471 +
472 +
473 +direct_register_custom_op(
474 + op_name="qwen4_exp_qsa_with_output",
475 + op_func=qwen4_exp_qsa_with_output,
476 + mutates_args=["output"],
477 + fake_impl=qwen4_exp_qsa_with_output_fake,
478 +)
479 +
480 +
481 +__all__ = [
482 + "QSAIndexer",
483 + "Qwen4ExpQSAAttention",
484 + "Qwen4ExpQSAFlashAttentionBackend",
485 + "Qwen4ExpQSAFlashAttentionImpl",
486 + "qwen4_exp_qsa_with_output",
487 +]
vllm/platforms/interface.py
+12 -5
@@ -855,11 +855,18 @@ class Platform:
855 model_config.architecture,
856 model_config=model_config,
857 )
858 - mamba_page_size = MambaSpec(
859 - shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
860 - dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
861 - block_size=-1,
862 - ).page_size_bytes
858 + # Qwen4Exp has multiple Mamba state layouts with different sizes.
859 + if hasattr(model_cls, "get_mamba_specs_from_config"):
860 + mamba_page_size = max(
861 + spec.page_size_bytes
862 + for spec in model_cls.get_mamba_specs_from_config(vllm_config)
863 + )
864 + else:
865 + mamba_page_size = MambaSpec(
866 + shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
867 + dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
868 + block_size=-1,
869 + ).page_size_bytes
870
871 if mamba_page_size == 0:
872 return
vllm/transformers_utils/config.py
+2
@@ -125,6 +125,8 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
125 qianfan_ocr="QianfanOCRConfig",
126 qwen3_asr="Qwen3ASRConfig",
127 qwen3_next="Qwen3NextConfig",
128 + qwen4_exp="Qwen4ExpConfig",
129 + qwen4_exp_text="Qwen4ExpTextConfig",
130 qwen3_5="Qwen3_5Config",
131 qwen3_5_text="Qwen3_5TextConfig",
132 qwen3_5_moe="Qwen3_5MoeConfig",
vllm/transformers_utils/configs/__init__.py
+6
@@ -91,6 +91,9 @@ _CLASS_TO_MODULE: dict[str, str] = {
91 "QianfanOCRVisionConfig": "vllm.transformers_utils.configs.qianfan_ocr",
92 "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr",
93 "Qwen3NextConfig": "vllm.transformers_utils.configs.qwen3_next",
94 + "Qwen4ExpConfig": "vllm.transformers_utils.configs.qwen4_exp",
95 + "Qwen4ExpTextConfig": "vllm.transformers_utils.configs.qwen4_exp",
96 + "Qwen4ExpVisionConfig": "vllm.transformers_utils.configs.qwen4_exp",
97 "Qwen3_5Config": "vllm.transformers_utils.configs.qwen3_5",
98 "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5",
99 "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe",
@@ -176,6 +179,9 @@ __all__ = [
179 "QianfanOCRVisionConfig",
180 "Qwen3ASRConfig",
181 "Qwen3NextConfig",
182 + "Qwen4ExpConfig",
183 + "Qwen4ExpTextConfig",
184 + "Qwen4ExpVisionConfig",
185 "Qwen3_5Config",
186 "Qwen3_5TextConfig",
187 "Qwen3_5MoeConfig",
vllm/transformers_utils/configs/qwen4_exp.py new
+15
@@ -0,0 +1,15 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Exports for the package-owned Qwen4Exp configs."""
4 +
5 +from vllm.models.qwen4_exp.config import (
6 + Qwen4ExpConfig,
7 + Qwen4ExpTextConfig,
8 + Qwen4ExpVisionConfig,
9 +)
10 +
11 +__all__ = [
12 + "Qwen4ExpConfig",
13 + "Qwen4ExpTextConfig",
14 + "Qwen4ExpVisionConfig",
15 +]
vllm/transformers_utils/model_arch_config_convertor.py
+10
@@ -642,6 +642,15 @@ class Qwen3NextMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
642 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
643
644
645 +class Qwen4ExpMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
646 + def get_num_hidden_layers(self) -> int:
647 + return getattr(
648 + self.hf_text_config,
649 + "mtp_num_hidden_layers",
650 + getattr(self.hf_text_config, "num_nextn_predict_layers", 0),
651 + )
652 +
653 +
654 class BailingHybridMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
655 def get_num_hidden_layers(self) -> int:
656 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
@@ -791,6 +800,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
800 "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,
801 "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor,
802 "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,
803 + "qwen4_exp_mtp": Qwen4ExpMTPModelArchConfigConvertor,
804 "RefinedWeb": FalconModelArchConfigConvertor,
805 "RefinedWebModel": FalconModelArchConfigConvertor,
806 "step3p5_mtp": Step3p5MTPModelArchConfigConvertor,
vllm/v1/attention/backends/short_conv_attn.py
+527 -2
@@ -1,12 +1,27 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from dataclasses import dataclass
3 +from dataclasses import dataclass, replace
4 +from typing import Any
5
5 -from vllm.v1.attention.backend import AttentionBackend
6 +import torch
7 +
8 +from vllm.config import VllmConfig
9 +from vllm.utils.torch_utils import async_tensor_h2d
10 +from vllm.v1.attention.backend import (
11 + AttentionBackend,
12 + AttentionCGSupport,
13 + CommonAttentionMetadata,
14 +)
15 from vllm.v1.attention.backends.mamba_attn import (
16 BaseMambaAttentionMetadata,
17 BaseMambaAttentionMetadataBuilder,
18 )
19 +from vllm.v1.attention.backends.utils import (
20 + NULL_BLOCK_ID,
21 + compute_causal_conv1d_metadata,
22 + mamba_get_block_table_tensor,
23 +)
24 +from vllm.v1.kv_cache_interface import MambaSpec
25
26
27 class ShortConvAttentionBackend(AttentionBackend):
@@ -32,3 +47,513 @@ class ShortConvAttentionMetadataBuilder(
47 BaseMambaAttentionMetadataBuilder[ShortConvAttentionMetadata]
48 ):
49 metadata_cls = ShortConvAttentionMetadata
50 +
51 +
52 +@dataclass
53 +class PleShortConvAttentionMetadata(ShortConvAttentionMetadata):
54 + # Number of speculative-decode (multi-query / MTP) requests and the total
55 + # number of tokens they contribute. These are 0 when spec-decode is off.
56 + num_spec_decodes: int = 0
57 + num_spec_decode_tokens: int = 0
58 + num_actual_tokens: int = 0
59 +
60 + # Max query length among spec-decode requests (== num_speculative_tokens + 1).
61 + # Used as ``max_query_len`` for the varlen spec causal_conv1d_update.
62 + spec_query_len: int = 1
63 +
64 + # Max query length among the non-spec *prefill* requests, precomputed
65 + # CPU-side in the builder. The dilated PLE short-conv uses it to size its
66 + # packing buffer without a device->host sync (``lengths.max().item()``).
67 + # 0 when there are no prefill requests.
68 + max_prefill_query_len: int = 0
69 + query_start_loc: torch.Tensor | None = None
70 +
71 + # ``state_indices_tensor`` keeps the historical (non-spec) layout used by
72 + # all existing short-conv consumers: the conv-state slot for each regular
73 + # decode followed by each prefill request. When spec-decode is active this
74 + # only covers the non-spec requests.
75 + state_indices_tensor: torch.Tensor | None = None
76 + has_initial_states_d: torch.Tensor | None = None
77 +
78 + # ``non_spec_query_start_loc`` is the varlen cumulative token offset over
79 + # the non-spec requests only (decodes then prefills). It equals
80 + # ``query_start_loc`` when there are no spec-decode requests.
81 + non_spec_query_start_loc: torch.Tensor | None = None
82 +
83 + # Speculative-decode (MTP) conv metadata. Only column 0 of the block table
84 + # is needed for the convolution state, so these tensors are 1-D over the
85 + # spec-decode requests.
86 + spec_query_start_loc: torch.Tensor | None = None # [num_spec_decodes + 1]
87 + spec_state_indices_tensor: torch.Tensor | None = None # [num_spec_decodes]
88 + spec_sequence_masks: torch.Tensor | None = None # [batch]
89 + spec_token_indx: torch.Tensor | None = None
90 + non_spec_token_indx: torch.Tensor | None = None
91 + num_decode_draft_tokens_cpu: torch.Tensor | None = None
92 +
93 +
94 +class PleShortConvAttentionBackend(ShortConvAttentionBackend):
95 + @staticmethod
96 + def get_name() -> str:
97 + return "PLE_SHORT_CONV_ATTN"
98 +
99 + @staticmethod
100 + def get_builder_cls() -> type["PleShortConvAttentionMetadataBuilder"]:
101 + return PleShortConvAttentionMetadataBuilder
102 +
103 +
104 +class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
105 + metadata_cls = PleShortConvAttentionMetadata
106 + # Spec-decode requires a uniform (multi-token) decode batch for full
107 + # CUDA graph capture, matching the GDN backend.
108 + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
109 + reorder_batch_threshold: int = 1
110 + supports_update_block_table = False
111 +
112 + def __init__(
113 + self,
114 + kv_cache_spec: MambaSpec,
115 + layer_names: list[str],
116 + vllm_config: VllmConfig,
117 + device: torch.device,
118 + ) -> None:
119 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
120 + self.num_spec = self.num_spec_tokens
121 + self.use_full_cuda_graph = (
122 + self.compilation_config.cudagraph_mode.has_full_cudagraphs()
123 + )
124 +
125 + max_num_seqs = vllm_config.scheduler_config.max_num_seqs
126 + max_capture_size = self.compilation_config.max_cudagraph_capture_size
127 + self.decode_cudagraph_max_bs = max_num_seqs
128 + self.decode_cudagraph_max_tokens = max_num_seqs * (self.num_spec + 1)
129 + if max_capture_size is not None:
130 + self.decode_cudagraph_max_bs = min(
131 + self.decode_cudagraph_max_bs, max_capture_size
132 + )
133 + self.decode_cudagraph_max_tokens = min(
134 + self.decode_cudagraph_max_tokens, max_capture_size
135 + )
136 +
137 + # Persistent buffers reused during full CUDA graph capture and replay.
138 + self.spec_state_indices_tensor = torch.empty(
139 + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device
140 + )
141 + self.spec_sequence_masks = torch.empty(
142 + (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device
143 + )
144 + self.spec_token_indx = torch.empty(
145 + (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device
146 + )
147 + self.non_spec_token_indx = torch.empty(
148 + (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device
149 + )
150 + self.spec_query_start_loc = torch.empty(
151 + (self.decode_cudagraph_max_bs + 1,), dtype=torch.int32, device=device
152 + )
153 + self.num_accepted_tokens = torch.empty(
154 + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device
155 + )
156 + self.has_initial_states_d = torch.empty(
157 + (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device
158 + )
159 +
160 + def _build_non_spec_metadata(
161 + self,
162 + common_prefix_len: int,
163 + common_attn_metadata: CommonAttentionMetadata,
164 + fast_build: bool,
165 + num_decode_draft_tokens_cpu: torch.Tensor | None,
166 + **kwargs: Any,
167 + ) -> PleShortConvAttentionMetadata:
168 + metadata = super().build(
169 + common_prefix_len,
170 + common_attn_metadata,
171 + fast_build,
172 + num_accepted_tokens=None,
173 + **kwargs,
174 + )
175 + assert isinstance(metadata, PleShortConvAttentionMetadata)
176 +
177 + state_indices_d = metadata.state_indices_tensor_d
178 + if state_indices_d is not None and state_indices_d.dim() > 1:
179 + state_indices_d = state_indices_d[:, 0]
180 + state_indices_p = metadata.state_indices_tensor_p
181 + if metadata.num_prefills == 0:
182 + assert state_indices_d is not None
183 + # BaseMambaAttentionMetadataBuilder pads decode state indices into
184 + # a persistent tensor for full CUDA graphs. Keep those rows so the
185 + # PLE decode receives one cache slot per graph-padded token.
186 + state_indices_tensor = state_indices_d
187 + elif metadata.num_decodes == 0:
188 + assert state_indices_p is not None
189 + state_indices_tensor = state_indices_p[: metadata.num_prefills]
190 + else:
191 + assert state_indices_d is not None
192 + assert state_indices_p is not None
193 + state_indices_tensor = torch.cat(
194 + (state_indices_d, state_indices_p[: metadata.num_prefills])
195 + )
196 +
197 + has_initial_states_d = None
198 + if metadata.num_decodes > 0:
199 + num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
200 + has_initial_states_d = num_computed_tokens[: metadata.num_decodes] > 0
201 + if (
202 + self.use_full_cuda_graph
203 + and metadata.num_prefills == 0
204 + and metadata.num_decodes <= self.decode_cudagraph_max_bs
205 + ):
206 + assert state_indices_d is not None
207 + # Prepare tensors for CUDA graph replay. Padded rows have no
208 + # initial state and use NULL_BLOCK_ID in state_indices_d.
209 + num_decode_rows = state_indices_d.numel()
210 + self.has_initial_states_d[: metadata.num_decodes].copy_(
211 + has_initial_states_d, non_blocking=True
212 + )
213 + self.has_initial_states_d[metadata.num_decodes : num_decode_rows].fill_(
214 + False
215 + )
216 + has_initial_states_d = self.has_initial_states_d[:num_decode_rows]
217 +
218 + max_prefill_query_len = 0
219 + if metadata.num_prefills > 0:
220 + query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu)
221 + max_prefill_query_len = int(
222 + query_lens_cpu[
223 + metadata.num_decodes : (
224 + metadata.num_decodes + metadata.num_prefills
225 + )
226 + ]
227 + .max()
228 + .item()
229 + )
230 +
231 + return replace(
232 + metadata,
233 + num_actual_tokens=common_attn_metadata.num_actual_tokens,
234 + spec_query_len=self.num_spec + 1,
235 + max_prefill_query_len=max_prefill_query_len,
236 + query_start_loc=common_attn_metadata.query_start_loc,
237 + state_indices_tensor=state_indices_tensor,
238 + has_initial_states_d=has_initial_states_d,
239 + non_spec_query_start_loc=common_attn_metadata.query_start_loc,
240 + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
241 + )
242 +
243 + def build( # type: ignore[override]
244 + self,
245 + common_prefix_len: int,
246 + common_attn_metadata: CommonAttentionMetadata,
247 + fast_build: bool = False,
248 + *,
249 + num_accepted_tokens: torch.Tensor | None = None,
250 + num_decode_draft_tokens_cpu: torch.Tensor | None = None,
251 + **kwargs: Any,
252 + ) -> PleShortConvAttentionMetadata:
253 + m = common_attn_metadata
254 + spec_sequence_masks_cpu: torch.Tensor | None = None
255 + # Detect speculative-decode requests. We use -1 to mark prefill and
256 + # plain-decode requests, so any value >= 0 is a (multi-query)
257 + # spec-decode request.
258 + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None:
259 + candidate_mask = num_decode_draft_tokens_cpu[: m.num_reqs] >= 0
260 + if bool(candidate_mask.any().item()):
261 + spec_sequence_masks_cpu = candidate_mask
262 +
263 + if spec_sequence_masks_cpu is None:
264 + return self._build_non_spec_metadata(
265 + common_prefix_len,
266 + common_attn_metadata,
267 + fast_build,
268 + num_decode_draft_tokens_cpu,
269 + **kwargs,
270 + )
271 +
272 + del common_prefix_len, fast_build, kwargs
273 + query_start_loc = m.query_start_loc
274 + query_start_loc_cpu = m.query_start_loc_cpu
275 + query_lens_cpu = torch.diff(query_start_loc_cpu)
276 + block_table_tensor = mamba_get_block_table_tensor(
277 + m.block_table_tensor,
278 + m.seq_lens,
279 + self.kv_cache_spec,
280 + self.vllm_config.cache_config.mamba_cache_mode,
281 + )
282 +
283 + if query_start_loc.device.type == "cpu":
284 + spec_sequence_masks = spec_sequence_masks_cpu
285 + else:
286 + spec_sequence_masks = async_tensor_h2d(
287 + spec_sequence_masks_cpu, device=query_start_loc.device
288 + )
289 +
290 + # For causal_conv1d (non-spec prefill Triton kernel metadata).
291 + nums_dict = None
292 + batch_ptr = None
293 + token_chunk_offset_ptr = None
294 + has_initial_states_p = None
295 + has_initial_states_d = None
296 + num_computed_tokens_p = None
297 + # Original request indices of the non-spec requests, ordered
298 + # [decodes, prefills]. Used to gather per-request data consistently.
299 + non_spec_req_idx_cpu: torch.Tensor | None = None
300 +
301 + query_lens = torch.diff(query_start_loc)
302 + # Per-request classification by mask, NOT by position. With
303 + # spec-decode, the front decode group can contain both spec-decode
304 + # requests and plain non-spec single-token decodes. Spec requests are
305 + # therefore not guaranteed to occupy the first num_spec_decodes slots.
306 + non_spec_mask_cpu = ~spec_sequence_masks_cpu
307 + decode_mask_cpu = non_spec_mask_cpu & (query_lens_cpu == 1)
308 + prefill_mask_cpu = non_spec_mask_cpu & (query_lens_cpu > 1)
309 +
310 + num_spec_decodes = int(spec_sequence_masks_cpu.sum().item())
311 + num_decodes = int(decode_mask_cpu.sum().item())
312 + num_prefills = int(prefill_mask_cpu.sum().item())
313 + num_decode_tokens = num_decodes
314 + num_prefill_tokens = int(query_lens_cpu[prefill_mask_cpu].sum().item())
315 + num_spec_decode_tokens = int(
316 + query_lens_cpu[spec_sequence_masks_cpu].sum().item()
317 + )
318 + # Max prefill query length, CPU-side (no device-to-host sync) for the
319 + # PLE dilated short-conv packing buffer.
320 + max_prefill_query_len = (
321 + int(query_lens_cpu[prefill_mask_cpu].max().item())
322 + if num_prefills > 0
323 + else 0
324 + )
325 +
326 + # Original request indices grouped as
327 + # [spec | non-spec decode | non-spec prefill]; each group keeps the
328 + # original (already reordered) relative order via a stable nonzero.
329 + spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0]
330 + decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0]
331 + prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0]
332 + non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu))
333 + spec_req_idx = spec_req_idx_cpu.to(query_start_loc.device)
334 + non_spec_req_idx = non_spec_req_idx_cpu.to(query_start_loc.device)
335 +
336 + if num_decodes == 0 and num_prefills == 0:
337 + # Pure speculative-decode batch: all real tokens are spec tokens.
338 + spec_token_indx = torch.arange(
339 + num_spec_decode_tokens,
340 + dtype=torch.int32,
341 + device=query_start_loc.device,
342 + )
343 + non_spec_token_indx = torch.empty(
344 + 0, dtype=torch.int32, device=query_start_loc.device
345 + )
346 + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0]
347 + non_spec_state_indices_tensor = None
348 + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1]
349 + non_spec_query_start_loc = None
350 + non_spec_query_start_loc_cpu = None
351 + else:
352 + # Mixed batch: build a per-token group key consistent with the
353 + # request grouping above (spec=0 | decode=1 | prefill=2) and a
354 + # stable sort, so tokens of each request stay contiguous and in
355 + # request order. This yields spec tokens first, then the non-spec
356 + # [decode, prefill] tokens.
357 + req_group = torch.full(
358 + (m.num_reqs,),
359 + 2,
360 + dtype=torch.int64,
361 + device=query_start_loc.device,
362 + )
363 + req_group[spec_req_idx] = 0
364 + req_group[decode_req_idx_cpu.to(query_start_loc.device)] = 1
365 + token_group = torch.repeat_interleave(req_group, query_lens)
366 + token_perm = torch.argsort(token_group, stable=True)
367 + spec_token_indx = token_perm[:num_spec_decode_tokens]
368 + non_spec_token_indx = token_perm[num_spec_decode_tokens:]
369 +
370 + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0]
371 + non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0]
372 + spec_query_start_loc = torch.zeros(
373 + num_spec_decodes + 1,
374 + dtype=torch.int32,
375 + device=query_start_loc.device,
376 + )
377 + torch.cumsum(query_lens[spec_req_idx], dim=0, out=spec_query_start_loc[1:])
378 + non_spec_query_start_loc = torch.zeros(
379 + num_decodes + num_prefills + 1,
380 + dtype=torch.int32,
381 + device=query_start_loc.device,
382 + )
383 + torch.cumsum(
384 + query_lens[non_spec_req_idx],
385 + dim=0,
386 + out=non_spec_query_start_loc[1:],
387 + )
388 + non_spec_query_start_loc_cpu = torch.zeros(
389 + num_decodes + num_prefills + 1, dtype=torch.int32
390 + )
391 + torch.cumsum(
392 + query_lens_cpu[non_spec_req_idx_cpu],
393 + dim=0,
394 + out=non_spec_query_start_loc_cpu[1:],
395 + )
396 +
397 + assert num_accepted_tokens is not None
398 + # Accepted-token counts must follow the same request order as the
399 + # speculative state indices.
400 + num_accepted_tokens = num_accepted_tokens[
401 + spec_req_idx_cpu.to(num_accepted_tokens.device)
402 + ]
403 +
404 + # Compute the conv-state slots for the non-spec decode/prefill split,
405 + # plus the initial-state masks and Triton causal_conv1d metadata.
406 + if non_spec_state_indices_tensor is None:
407 + state_indices_tensor = block_table_tensor[:0, 0]
408 + else:
409 + state_indices_tensor = non_spec_state_indices_tensor
410 +
411 + # Build the regular decode/prefill state metadata inherited from the
412 + # generic short-conv metadata contract.
413 + query_start_loc_p = None
414 + query_start_loc_d = None
415 + state_indices_tensor_p = None
416 + state_indices_tensor_d = None
417 + if num_decodes > 0 or num_prefills > 0:
418 + num_computed_tokens = m.compute_num_computed_tokens()
419 + if non_spec_req_idx_cpu is not None:
420 + non_spec_req_idx = non_spec_req_idx_cpu.to(num_computed_tokens.device)
421 + num_computed_tokens = num_computed_tokens[non_spec_req_idx]
422 +
423 + state_indices_tensor_d = state_indices_tensor[:num_decodes]
424 + state_indices_tensor_p = state_indices_tensor[
425 + num_decodes : num_decodes + num_prefills
426 + ]
427 + if num_decodes > 0:
428 + has_initial_states_d = num_computed_tokens[:num_decodes] > 0
429 + assert non_spec_query_start_loc is not None
430 + query_start_loc_d = non_spec_query_start_loc[: num_decodes + 1]
431 + if num_prefills > 0:
432 + num_computed_tokens_p = num_computed_tokens[
433 + num_decodes : num_decodes + num_prefills
434 + ]
435 + has_initial_states_p = num_computed_tokens_p > 0
436 + assert non_spec_query_start_loc is not None
437 + assert non_spec_query_start_loc_cpu is not None
438 + query_start_loc_p = (
439 + non_spec_query_start_loc[num_decodes:] - num_decode_tokens
440 + )
441 + query_start_loc_p_cpu = (
442 + non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens
443 + )
444 + if query_start_loc.device.type != "cpu":
445 + nums_dict, batch_ptr, token_chunk_offset_ptr = (
446 + compute_causal_conv1d_metadata(
447 + query_start_loc_p_cpu,
448 + device=query_start_loc.device,
449 + )
450 + )
451 +
452 + # Prepare persistent tensors for CUDA graph capture and replay.
453 + # ``m.num_actual_tokens`` is already padded by the model runner.
454 + # Request-level buffers use ``m.num_reqs`` while token-level buffers
455 + # use their independently bounded token count.
456 + batch_size = m.num_reqs
457 + if (
458 + self.use_full_cuda_graph
459 + and num_prefills == 0
460 + and num_decodes == 0
461 + and spec_sequence_masks is not None
462 + and num_spec_decodes <= self.decode_cudagraph_max_bs
463 + and num_spec_decode_tokens <= self.decode_cudagraph_max_tokens
464 + ):
465 + assert spec_state_indices_tensor is not None
466 + self.spec_state_indices_tensor[:num_spec_decodes].copy_(
467 + spec_state_indices_tensor, non_blocking=True
468 + )
469 + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
470 + spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID)
471 +
472 + self.spec_sequence_masks[:batch_size].copy_(
473 + spec_sequence_masks[:batch_size], non_blocking=True
474 + )
475 + spec_sequence_masks = self.spec_sequence_masks[:batch_size]
476 +
477 + assert spec_query_start_loc is not None
478 + self.spec_query_start_loc[: num_spec_decodes + 1].copy_(
479 + spec_query_start_loc, non_blocking=True
480 + )
481 + spec_num_query_tokens = spec_query_start_loc[-1]
482 + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1]
483 + spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens)
484 +
485 + assert num_accepted_tokens is not None
486 + self.num_accepted_tokens[:num_spec_decodes].copy_(
487 + num_accepted_tokens, non_blocking=True
488 + )
489 + num_accepted_tokens = self.num_accepted_tokens[:batch_size]
490 + num_accepted_tokens[num_spec_decodes:].fill_(1)
491 +
492 + return PleShortConvAttentionMetadata(
493 + num_prefills=num_prefills,
494 + num_prefill_tokens=num_prefill_tokens,
495 + num_decodes=num_decodes,
496 + num_decode_tokens=num_decode_tokens,
497 + num_reqs=m.num_reqs,
498 + num_spec_decodes=num_spec_decodes,
499 + num_spec_decode_tokens=num_spec_decode_tokens,
500 + num_actual_tokens=m.num_actual_tokens,
501 + spec_query_len=self.num_spec + 1,
502 + max_prefill_query_len=max_prefill_query_len,
503 + query_start_loc=query_start_loc,
504 + state_indices_tensor=state_indices_tensor,
505 + has_initial_states_p=has_initial_states_p,
506 + has_initial_states_d=has_initial_states_d,
507 + non_spec_query_start_loc=non_spec_query_start_loc,
508 + spec_query_start_loc=spec_query_start_loc,
509 + spec_state_indices_tensor=spec_state_indices_tensor,
510 + spec_sequence_masks=spec_sequence_masks,
511 + spec_token_indx=spec_token_indx,
512 + non_spec_token_indx=non_spec_token_indx,
513 + num_accepted_tokens=num_accepted_tokens,
514 + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
515 + nums_dict=nums_dict,
516 + batch_ptr=batch_ptr,
517 + token_chunk_offset_ptr=token_chunk_offset_ptr,
518 + query_start_loc_p=query_start_loc_p,
519 + query_start_loc_d=query_start_loc_d,
520 + state_indices_tensor_p=state_indices_tensor_p,
521 + state_indices_tensor_d=state_indices_tensor_d,
522 + num_computed_tokens_p=num_computed_tokens_p,
523 + block_idx_last_scheduled_token=None,
524 + block_idx_first_scheduled_token_p=None,
525 + block_idx_last_computed_token=None,
526 + block_idx_last_scheduled_token_prev_step=None,
527 + seq_lens=m.seq_lens,
528 + )
529 +
530 + def build_for_cudagraph_capture(
531 + self, common_attn_metadata: CommonAttentionMetadata
532 + ) -> PleShortConvAttentionMetadata:
533 + """Build metadata for full CUDA graph capture.
534 +
535 + Currently, only decode is supported for full CUDA graphs with
536 + short-conv.
537 + """
538 + m = common_attn_metadata
539 + assert (
540 + m.num_reqs <= self.decode_cudagraph_max_bs
541 + and m.num_actual_tokens <= self.decode_cudagraph_max_tokens
542 + ), (
543 + "ShortConv only supports decode-only full CUDAGraph capture. "
544 + f"Make sure batch size ({m.num_reqs}) <= "
545 + f"cudagraph capture size ({self.decode_cudagraph_max_bs}) and "
546 + f"number of tokens ({m.num_actual_tokens}) <= "
547 + f"token capture size ({self.decode_cudagraph_max_tokens})."
548 + )
549 +
550 + if self.use_spec_decode:
551 + num_accepted_tokens = torch.diff(m.query_start_loc)
552 + num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu()
553 + return self.build(
554 + 0,
555 + m,
556 + num_accepted_tokens=num_accepted_tokens,
557 + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
558 + )
559 + return self.build(0, m)
vllm/v1/core/kv_cache_coordinator.py
+14 -5
@@ -599,7 +599,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
599 self.hash_block_size = hash_block_size
600 self.dcp_world_size = dcp_world_size
601 group_block_sizes = [
602 - manager.block_size for manager in self.single_type_managers
602 + manager.block_size
603 + for manager, group in zip(
604 + self.single_type_managers, kv_cache_config.kv_cache_groups
605 + )
606 + if group.kv_cache_spec.prefix_cacheable
607 ]
608 assert all(
609 block_size % hash_block_size == 0 for block_size in group_block_sizes
@@ -638,8 +642,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
642 if self.enable_partial_hash_hits:
643 unsupported_partial_hit_managers = {
644 type(manager).__name__
641 - for manager in self.single_type_managers
642 - if not manager.supports_fine_grained_hash_lookup
645 + for manager, group in zip(
646 + self.single_type_managers, kv_cache_config.kv_cache_groups
647 + )
648 + if group.kv_cache_spec.prefix_cacheable
649 + and not manager.supports_fine_grained_hash_lookup
650 and manager.block_size != hash_block_size
651 }
652 if unsupported_partial_hit_managers:
@@ -668,6 +675,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
675 """
676 self.attention_groups: list[SpecGroup] = []
677 for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
678 + if not g.kv_cache_spec.prefix_cacheable:
679 + continue
680 manager_cls = self.single_type_managers[i].__class__
681 spec = g.kv_cache_spec
682 use_eagle = i in self.eagle_group_ids
@@ -687,8 +696,8 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
696 SpecGroup(spec, [i], manager_cls, use_eagle)
697 )
698
690 - assert len(self.attention_groups) > 1, (
691 - "HybridKVCacheCoordinator requires at least two attention groups."
699 + assert self.attention_groups, (
700 + "HybridKVCacheCoordinator requires at least one cacheable group."
701 )
702
703 # Put full attention first: its efficient left-to-right scan provides
vllm/v1/core/kv_cache_manager.py
+8
@@ -779,6 +779,14 @@ class KVCacheManager:
779 self.kv_cache_config.kv_cache_groups,
780 strict=True,
781 ):
782 + if not group.kv_cache_spec.prefix_cacheable:
783 + # Scratch groups (e.g. the CSA compressor ring) hold a fixed
784 + # block covering no token range, so a lookup result never
785 + # carries computed blocks for them and their block size does
786 + # not divide the endpoint.
787 + assert not group_blocks
788 + truncated.append([])
789 + continue
790 assert num_computed_tokens % manager.block_size == 0
791 num_blocks = num_computed_tokens // manager.block_size
792 if isinstance(group.kv_cache_spec, MambaSpec):
vllm/v1/core/kv_cache_utils.py
+193 -156
@@ -16,7 +16,7 @@ from vllm import envs
16 from vllm.config import VllmConfig
17 from vllm.logger import init_logger
18 from vllm.utils.hashing import xxhash, xxhash_cbor
19 -from vllm.utils.math_utils import cdiv, round_up
19 +from vllm.utils.math_utils import cdiv
20 from vllm.utils.mem_utils import format_gib
21 from vllm.utils.torch_utils import get_dtype_size
22 from vllm.v1.kv_cache_interface import (
@@ -717,21 +717,52 @@ def resolve_kv_cache_block_sizes(
717 # scheduler block size. Read the mode from the resolved group spec because
718 # its block size may have been updated independently of cache_config.
719 if any(
720 - isinstance(g.kv_cache_spec, MambaSpec)
721 - and g.kv_cache_spec.mamba_cache_mode != "align"
722 - for g in groups
720 + isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align"
721 + for group in groups
722 + for spec in iter_layer_specs(group.kv_cache_spec)
723 ):
724 return scheduler_block_size, scheduler_block_size
725
726 + hashing_sizes = [
727 + block_size
728 + for group, block_size in zip(groups, group_block_sizes)
729 + if group.kv_cache_spec.prefix_cacheable
730 + ] or group_block_sizes
731 requested = cache_config.prefix_match_unit
727 - hash_block_size = (
728 - requested if requested is not None else math.gcd(*group_block_sizes)
732 + hash_block_size = requested if requested is not None else math.gcd(*hashing_sizes)
733 + if any(bs % hash_block_size != 0 for bs in hashing_sizes):
734 + raise ValueError(
735 + f"Invalid prefix_match_unit={hash_block_size}; prefix-cacheable "
736 + "KV cache group block sizes must be divisible by prefix_match_unit. "
737 + f"Got group block sizes={group_block_sizes}, "
738 + f"prefix-cacheable={hashing_sizes}."
739 + )
740 + prefix_alignments = {
741 + spec.tokens_per_state
742 + for group in groups
743 + for spec in iter_layer_specs(group.kv_cache_spec)
744 + if spec.prefix_cacheable
745 + and isinstance(spec.tokens_per_state, int)
746 + and spec.tokens_per_state > 1
747 + }
748 + has_partial_mamba_group = any(
749 + isinstance(spec, MambaSpec)
750 + and spec.mamba_cache_mode == "align"
751 + and (
752 + (dcp == 1 and block_size > hash_block_size)
753 + or (dcp > 1 and block_size >= hash_block_size)
754 + )
755 + for group, block_size in zip(groups, group_block_sizes)
756 + for spec in iter_layer_specs(group.kv_cache_spec)
757 + )
758 + cache_hit_alignment = (
759 + hash_block_size if has_partial_mamba_group else scheduler_block_size
760 )
730 - if any(bs % hash_block_size != 0 for bs in group_block_sizes):
761 + if any(cache_hit_alignment % alignment for alignment in prefix_alignments):
762 raise ValueError(
732 - f"Invalid prefix_match_unit={hash_block_size}; all KV cache group "
733 - f"block sizes must be divisible by prefix_match_unit. "
734 - f"Got group block sizes={group_block_sizes}."
763 + f"Invalid prefix_match_unit={hash_block_size}; prefix-cache boundaries "
764 + "must align with each spec's per-state compression. "
765 + f"Got alignments={sorted(prefix_alignments)}."
766 )
767 return scheduler_block_size, hash_block_size
768
@@ -1577,49 +1608,6 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1608 kv_cache_spec.update(_promote_local_kv_cache_specs(kv_cache_spec))
1609
1610
1580 -def group_and_unify_kv_cache_specs(
1581 - kv_cache_spec: dict[str, KVCacheSpec],
1582 -) -> list[UniformTypeKVCacheSpecs] | None:
1583 - """
1584 - Group the KV cache specs and unify each group into one UniformTypeKVCacheSpecs.
1585 - Currently, this is only used for DeepseekV4.
1586 - """
1587 - if not any(
1588 - isinstance(spec, SlidingWindowMLASpec) for spec in kv_cache_spec.values()
1589 - ):
1590 - return None
1591 -
1592 - # SlidingWindowMLASpec models with uniform page sizes don't need tuple packing.
1593 - page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()}
1594 - if len(page_sizes) <= 1:
1595 - return None
1596 -
1597 - mla_specs: dict[str, KVCacheSpec] = {}
1598 - grouped_swa_mla_specs: dict[tuple[int, int], dict[str, KVCacheSpec]] = defaultdict(
1599 - dict
1600 - )
1601 - # NOTE: Here we group SWA layers by (block_size, sliding_window), which separates
1602 - # SWA layers, C4I+C4A layers, and C128A layers into three different groups. It can
1603 - # be fragile with only block_size and sliding_window as keys, but fine for now.
1604 - for name, spec in kv_cache_spec.items():
1605 - if isinstance(spec, SlidingWindowMLASpec):
1606 - grouped_swa_mla_specs[(spec.block_size, spec.sliding_window)][name] = spec
1607 - elif isinstance(spec, MLAAttentionSpec):
1608 - mla_specs[name] = spec
1609 -
1610 - assert len(mla_specs) > 0
1611 - mla_uniform_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs)
1612 - assert mla_uniform_spec is not None
1613 -
1614 - swa_uniform_specs: list[UniformTypeKVCacheSpecs] = []
1615 - for spec_dict in grouped_swa_mla_specs.values():
1616 - uniform_spec = UniformTypeKVCacheSpecs.from_specs(spec_dict)
1617 - assert uniform_spec is not None
1618 - swa_uniform_specs.append(uniform_spec)
1619 -
1620 - return [mla_uniform_spec, *swa_uniform_specs]
1621 -
1622 -
1611 def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -> int:
1612 """Pick a chunk size that minimizes total upward padding.
1613
@@ -1655,91 +1643,145 @@ def _approximate_gcd(values: Sequence[int], *, lower_bound: int | None = None) -
1643 return best_d
1644
1645
1658 -def _get_kv_cache_groups_uniform_groups(
1659 - grouped_specs: list[UniformTypeKVCacheSpecs],
1660 -) -> list[KVCacheGroupSpec]:
1661 - """
1662 - Generate the KV cache groups from the grouped specs.
1646 +def _get_packed_kv_cache_groups(
1647 + vllm_config: VllmConfig,
1648 + kv_cache_spec: dict[str, KVCacheSpec],
1649 +) -> list[KVCacheGroupSpec] | None:
1650 + """Group mixed-page-size layers for contiguous block-outermost packing.
1651 +
1652 + Greedily buckets layers into uniform-type specs. Buckets with equal layer
1653 + counts per page size are treated as a repeating layer pattern (one layer
1654 + per page size) and split into groups covering the same number of pattern
1655 + repeats (picked by ``_approximate_gcd`` to minimize padding), so all
1656 + groups pack into the same per-block layout. Mamba buckets are additionally
1657 + split to fit the block the attention buckets already need.
1658 + Returns None when the layout is not block-outermost or all layers already
1659 + share one page size.
1660 """
1664 - assert len(grouped_specs) > 0 and all(
1665 - isinstance(spec, UniformTypeKVCacheSpecs) for spec in grouped_specs
1666 - )
1667 - # For now, we restrict the first grouped_spec to be UniformTypeKVCacheSpecs
1668 - # containing only MLAAttentionSpec.
1669 - full_mla_spec = grouped_specs[0]
1670 - assert all(
1671 - isinstance(spec, MLAAttentionSpec)
1672 - for spec in full_mla_spec.kv_cache_specs.values()
1661 + layout = vllm_config.cache_config.get_resolved_kv_cache_layout()
1662 + page_sizes = {spec.page_size_bytes for spec in kv_cache_spec.values()}
1663 + if not layout.is_block_outermost or len(page_sizes) <= 1:
1664 + return None
1665 +
1666 + buckets: list[dict[str, KVCacheSpec]] = []
1667 + for name, spec in kv_cache_spec.items():
1668 + for bucket in buckets:
1669 + candidate = {**bucket, name: spec}
1670 + if UniformTypeKVCacheSpecs.is_uniform_type(candidate):
1671 + bucket[name] = spec
1672 + break
1673 + else:
1674 + buckets.append({name: spec})
1675 +
1676 + bucketed = []
1677 + for bucket in buckets:
1678 + uniform_spec = UniformTypeKVCacheSpecs.from_specs(bucket)
1679 + assert uniform_spec is not None
1680 + page_size_layers: dict[int, list[str]] = defaultdict(list)
1681 + for layer_name, layer_spec in bucket.items():
1682 + page_size_layers[layer_spec.page_size_bytes].append(layer_name)
1683 + # Only 1:1 patterns (one layer of each page size per repeat) are
1684 + # supported; counts sharing a gcd > 1 (e.g. 2:1) could in principle
1685 + # repeat too, but such buckets are emitted whole instead.
1686 + balanced = len(set(map(len, page_size_layers.values()))) == 1
1687 + bucketed.append((uniform_spec, page_size_layers, balanced))
1688 +
1689 + # Balanced buckets that mix page sizes must stay whole, so the largest one
1690 + # sets a floor on the repeats per group; larger single-size buckets are
1691 + # split down toward it. No such bucket means nothing needs packing.
1692 + min_repeats_per_group = max(
1693 + (
1694 + spec.get_max_layers_per_page_size()
1695 + for spec, page_size_layers, balanced in bucketed
1696 + if balanced and len(page_size_layers) > 1
1697 + ),
1698 + default=0,
1699 )
1674 - full_mla_group = KVCacheGroupSpec(
1675 - layer_names=list(full_mla_spec.kv_cache_specs.keys()),
1676 - kv_cache_spec=full_mla_spec,
1700 + repeats_per_group = (
1701 + _approximate_gcd(
1702 + [
1703 + spec.get_max_layers_per_page_size()
1704 + for spec, _, balanced in bucketed
1705 + if balanced
1706 + ],
1707 + lower_bound=min_repeats_per_group,
1708 + )
1709 + if min_repeats_per_group
1710 + else None
1711 )
1712
1679 - # We define a layer tuple as a group of layers with different page sizes, and
1680 - # one UniformTypeKVCacheSpecs contains a list of layer tuples.
1681 - # For example, if we have 11 C4 layers and 10 C128 layers, we can define a layer
1682 - # tuple as [C4I, C4A, C128], and the full_mla_group will contain "11" layer tuples.
1683 - # The other uniform KV cache specs will be similarly partitioned into layer tuples.
1684 - # Say we have 21 SWA layers, all with the same page size, then we will have "21"
1685 - # layer tuples.
1686 - num_layer_tuples_per_group: list[int] = [
1687 - g_spec.get_num_layer_tuples() for g_spec in grouped_specs
1688 - ]
1689 - # Choose `num_layer_tuples` to minimize total padding across groups.
1690 - num_layer_tuples = _approximate_gcd(
1691 - num_layer_tuples_per_group, lower_bound=num_layer_tuples_per_group[0]
1692 - )
1693 - # Round up to the nearest multiple of `num_layer_tuples` (i.e., padding)
1694 - num_layer_tuples_per_group = [
1695 - round_up(x, num_layer_tuples) for x in num_layer_tuples_per_group
1696 - ]
1713 + def num_groups_for(spec: UniformTypeKVCacheSpecs, balanced: bool) -> int:
1714 + if balanced and repeats_per_group is not None:
1715 + return cdiv(spec.get_max_layers_per_page_size(), repeats_per_group)
1716 + return 1
1717
1698 - swa_mla_specs = grouped_specs[1:]
1699 - assert all(
1700 - isinstance(spec, SlidingWindowMLASpec)
1701 - for group in swa_mla_specs
1702 - for spec in group.kv_cache_specs.values()
1718 + def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int:
1719 + """Page bytes of the largest of the n groups a bucket splits into."""
1720 + return sum(
1721 + cdiv(len(names), n) * page for page, names in page_size_layers.items()
1722 + )
1723 +
1724 + # Bytes a block must hold however the mamba buckets end up split: a mamba
1725 + # bucket can go down to one state per group, every other bucket's split is
1726 + # already fixed by the repeat pattern.
1727 + anchor_bytes = max(
1728 + (
1729 + widest_group_bytes(
1730 + page_size_layers,
1731 + len(spec.kv_cache_specs)
1732 + if isinstance(spec.first_spec, MambaSpec)
1733 + else num_groups_for(spec, balanced),
1734 + )
1735 + for spec, page_size_layers, balanced in bucketed
1736 + ),
1737 + default=0,
1738 )
1739
1705 - # Split each SWA UniformKV group into smaller groups to align their
1706 - # numbers of layer tuples. The packed block planner overlays groups, so
1707 - # their page sizes do not need to match.
1708 - swa_mla_groups = []
1709 - for sm_spec in swa_mla_specs:
1710 - layers_per_size: dict[int, list[str]] = defaultdict(list)
1711 -
1712 - for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
1713 - layers_per_size[layer_spec.page_size_bytes].append(layer_name)
1714 - # NOTE(yifan): for now, inside a UniformKV group, each page_size should
1715 - # have the same number of layers. This also means we don't need to pad layers
1716 - # inside a partial-full layer tuple.
1717 - assert len(set(len(layers) for layers in layers_per_size.values())) == 1
1718 - num_layers_per_size = len(next(iter(layers_per_size.values())))
1719 -
1720 - # Split layers inside each UniformKV group for aligned #(layers).
1721 - # See `_get_kv_cache_groups_uniform_page_size` for more details.
1722 - num_tuple_groups = cdiv(num_layers_per_size, num_layer_tuples)
1723 - layer_tuples = list(zip(*layers_per_size.values()))
1724 - for i in range(num_tuple_groups):
1725 - group_layer_tuples = layer_tuples[i::num_tuple_groups]
1726 - # Flatten tuples and build dict for from_specs
1740 + groups = []
1741 + for spec, page_size_layers, balanced in bucketed:
1742 + num_groups = num_groups_for(spec, balanced)
1743 + # `_align_hybrid_block_size` pads a mamba state up to one attention
1744 + # page, so cap a mamba group at the states a block already fits rather
1745 + # than let it widen the block.
1746 + if anchor_bytes and isinstance(spec.first_spec, MambaSpec):
1747 + states_per_block = max(anchor_bytes // spec.first_spec.page_size_bytes, 1)
1748 + num_groups = max(
1749 + num_groups, cdiv(len(spec.kv_cache_specs), states_per_block)
1750 + )
1751 + if num_groups == 1:
1752 + groups.append(KVCacheGroupSpec(list(spec.kv_cache_specs), spec))
1753 + continue
1754 +
1755 + pattern_repeats = list(zip(*page_size_layers.values()))
1756 + for i in range(num_groups):
1757 group_layer_names = [
1728 - name for layer_tuple in group_layer_tuples for name in layer_tuple
1758 + name for repeat in pattern_repeats[i::num_groups] for name in repeat
1759 ]
1760 group_layer_specs = {
1731 - name: sm_spec.kv_cache_specs[name] for name in group_layer_names
1761 + name: spec.kv_cache_specs[name] for name in group_layer_names
1762 }
1733 - sub_sm_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
1734 - assert sub_sm_spec is not None
1735 - swa_mla_groups.append(
1736 - KVCacheGroupSpec(
1737 - layer_names=group_layer_names,
1738 - kv_cache_spec=sub_sm_spec,
1739 - )
1740 - )
1763 + group_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
1764 + assert group_spec is not None
1765 + groups.append(KVCacheGroupSpec(group_layer_names, group_spec))
1766 +
1767 + _annotate_eagle_groups(
1768 + vllm_config,
1769 + kv_cache_spec,
1770 + groups,
1771 + use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config),
1772 + )
1773 + _warn_if_unannotated_eagle_mamba(vllm_config, groups)
1774 + return groups
1775
1742 - return [full_mla_group, *swa_mla_groups]
1776 +
1777 +def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool:
1778 + spec_config = vllm_config.speculative_config
1779 + if spec_config is None or not spec_config.use_eagle():
1780 + return False
1781 + model_config = vllm_config.model_config
1782 + return (
1783 + model_config is not None and model_config.hf_config.model_type == "deepseek_v4"
1784 + )
1785
1786
1787 def _annotate_eagle_groups(
@@ -1763,9 +1805,10 @@ def _annotate_eagle_groups(
1805 the target's own decoder layer and so carries no spec marker. Its draft
1806 attention layer is always the last registered layer, so flag whichever
1807 group holds it. This rule is only valid where the groups partition
1766 - exactly the layers of ``kv_cache_spec``, which is true on the
1767 - group_and_unify path and not in general; other callers must leave
1768 - ``use_deepseek_v4_fallback`` False.
1808 + exactly the layers of ``kv_cache_spec``, which is true on the packed
1809 + grouping path and not in general; other callers must leave
1810 + ``use_deepseek_v4_fallback`` False. The caller gates this fallback on
1811 + the configured model type.
1812 FIXME(yifan): avoid/generalize this hacky check.
1813
1814 Args:
@@ -1773,25 +1816,21 @@ def _annotate_eagle_groups(
1816 kv_cache_spec: The kv cache spec of each attention layer, in layer
1817 registration order. Only read by rule 2.
1818 kv_cache_groups: Groups to annotate in place.
1776 - use_deepseek_v4_fallback: Enable rule 2. Only the group_and_unify path
1777 - may set this.
1819 + use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group.
1820 """
1821 spec_config = vllm_config.speculative_config
1822 if spec_config is None or not spec_config.use_eagle():
1823 return
1824
1825 for group in kv_cache_groups:
1784 - if getattr(group.kv_cache_spec, "non_causal_multi_token_decode", False):
1826 + if any(
1827 + getattr(spec, "non_causal_multi_token_decode", False)
1828 + for spec in iter_layer_specs(group.kv_cache_spec)
1829 + ):
1830 group.is_eagle_group = True
1831
1832 if not use_deepseek_v4_fallback:
1833 return
1789 - # Detection uses the merged MLA spec's model_version.
1790 - if not any(
1791 - getattr(spec, "model_version", None) == "deepseek_v4"
1792 - for spec in kv_cache_spec.values()
1793 - ):
1794 - return
1834 last_layer = next(reversed(kv_cache_spec))
1835 for group in kv_cache_groups:
1836 if last_layer in group.layer_names:
@@ -1822,7 +1861,10 @@ def _warn_if_unannotated_eagle_mamba(
1861 mamba_groups = [
1862 idx
1863 for idx, group in enumerate(kv_cache_groups)
1825 - if isinstance(group.kv_cache_spec, MambaSpec)
1864 + if any(
1865 + isinstance(spec, MambaSpec)
1866 + for spec in iter_layer_specs(group.kv_cache_spec)
1867 + )
1868 ]
1869 if not mamba_groups:
1870 return
@@ -1846,7 +1888,8 @@ def _largest_divisor_at_most(value: int, limit: int) -> int:
1888
1889
1890 def get_kv_cache_groups(
1849 - vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec]
1891 + vllm_config: VllmConfig,
1892 + kv_cache_spec: dict[str, KVCacheSpec],
1893 ) -> list[KVCacheGroupSpec]:
1894 """
1895 Split the layers in the model into groups with the same KV cache spec.
@@ -1876,22 +1919,8 @@ def get_kv_cache_groups(
1919 # full attention, or all layers are sliding window attention with the
1920 # same window size). Put all layers into one group.
1921 return _get_kv_cache_groups_uniform_type(uniform_spec)
1879 - elif grouped_specs := group_and_unify_kv_cache_specs(kv_cache_spec):
1880 - # DeepseekV4 case: All layers need the same number of token slots,
1881 - # yet some layers are full attention while others are sliding window
1882 - # attention in different sizes. Need to group layers into multiple
1883 - # UniformTypeKVCacheSpecs.
1884 - kv_cache_groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
1885 - _annotate_eagle_groups(
1886 - vllm_config,
1887 - kv_cache_spec,
1888 - kv_cache_groups,
1889 - use_deepseek_v4_fallback=True,
1890 - )
1891 - return kv_cache_groups
1892 -
1893 - # Pull HiddenStateCacheSpec layers out before the general multi-group
1894 - # path so they don't affect page-size unification or grouping.
1922 + # Hidden-state layers use their own block table and must not be absorbed
1923 + # into a compatible attention bucket.
1924 hidden_specs = {
1925 k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec)
1926 }
@@ -1901,6 +1930,14 @@ def get_kv_cache_groups(
1930 if not isinstance(v, HiddenStateCacheSpec)
1931 }
1932
1933 + if packed_groups := _get_packed_kv_cache_groups(vllm_config, filtered_spec):
1934 + # Block-outermost blocks are strided by the widest group, so hidden
1935 + # groups need no page alignment.
1936 + packed_groups += [
1937 + KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items()
1938 + ]
1939 + return packed_groups
1940 +
1941 # Prefer preserving each layer's cache semantics. If physical pages cannot
1942 # be unified, try a supported allocation-only fallback before failing.
1943 try:
vllm/v1/core/sched/scheduler.py
+13
@@ -1008,6 +1008,18 @@ class Scheduler(SchedulerInterface):
1008 )
1009 if num_new_tokens == 0:
1010 break
1011 + if (
1012 + pad_spec_decode
1013 + and num_new_tokens != 1 + self.num_spec_tokens
1014 + ):
1015 + # Alignment clipped the placeholder rows. The split
1016 + # aligns prefill chunks, but the padded tail rows are
1017 + # speculative positions, not prefill tokens. A padded
1018 + # request must keep all 1 + num_spec rows or the
1019 + # sampler's row count stops matching its query rows,
1020 + # so drop the padding instead of shortening it.
1021 + num_new_tokens = 1
1022 + pad_spec_decode = False
1023
1024 # Schedule encoder inputs.
1025 if request.has_encoder_inputs:
@@ -1170,6 +1182,7 @@ class Scheduler(SchedulerInterface):
1182 request.status = RequestStatus.RUNNING
1183 request.num_computed_tokens = num_computed_tokens
1184 if pad_spec_decode:
1185 + assert num_new_tokens == 1 + self.num_spec_tokens
1186 scheduled_spec_decode_tokens[request_id] = [
1187 -1
1188 ] * self.num_spec_tokens
vllm/v1/core/single_type_kv_cache_manager.py
+97 -2
@@ -18,6 +18,7 @@ from vllm.v1.core.kv_cache_utils import (
18 from vllm.v1.kv_cache_interface import (
19 AttentionSpec,
20 ChunkedLocalAttentionSpec,
21 + CircularBufferSpec,
22 CrossAttentionSpec,
23 FullAttentionSpec,
24 HiddenStateCacheSpec,
@@ -83,8 +84,10 @@ class SingleTypeKVCacheManager(ABC):
84 self._max_admission_blocks_per_request = max_admission_blocks_per_request
85 # Record newly allocated block ids only when worker-side zeroing will
86 # consume them and this manager holds a spec type that gets zeroed.
86 - self._record_new_block_ids = needs_kv_cache_zeroing and isinstance(
87 - kv_cache_spec, AttentionSpec
87 + self._record_new_block_ids = (
88 + needs_kv_cache_zeroing
89 + and isinstance(kv_cache_spec, AttentionSpec)
90 + and not isinstance(kv_cache_spec, CircularBufferSpec)
91 )
92 self.new_block_ids: list[int] = []
93
@@ -1109,6 +1112,93 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
1112 return 0
1113
1114
1115 +class CircularBufferManager(FullAttentionManager):
1116 + """Claims the ring's single block per request; prefix caching disabled."""
1117 +
1118 + supports_fine_grained_hash_lookup: ClassVar[bool] = False
1119 +
1120 + def _claim_ring_block(self, request_id: str) -> list[KVCacheBlock]:
1121 + req_blocks = self.req_to_blocks[request_id]
1122 + if req_blocks:
1123 + return []
1124 + new_blocks = self.block_pool.get_new_blocks(1)
1125 + req_blocks.extend(new_blocks)
1126 + if self._record_new_block_ids:
1127 + self.new_block_ids.extend(block.block_id for block in new_blocks)
1128 + return new_blocks
1129 +
1130 + def get_num_blocks_to_allocate(
1131 + self,
1132 + request_id: str,
1133 + num_tokens: int,
1134 + new_computed_blocks: Sequence[KVCacheBlock],
1135 + total_computed_tokens: int,
1136 + num_local_computed_tokens: int,
1137 + num_tokens_main_model: int,
1138 + apply_admission_cap: bool = False,
1139 + ) -> int:
1140 + return 0 if self.req_to_blocks.get(request_id) else 1
1141 +
1142 + def allocate_new_blocks(
1143 + self, request_id: str, num_tokens: int, num_tokens_main_model: int
1144 + ) -> list[KVCacheBlock]:
1145 + return self._claim_ring_block(request_id)
1146 +
1147 + def allocate_external_computed_blocks(
1148 + self,
1149 + request_id: str,
1150 + num_local_computed_tokens: int,
1151 + num_external_computed_tokens: int,
1152 + ) -> None:
1153 + self._claim_ring_block(request_id)
1154 +
1155 + @classmethod
1156 + def find_longest_cache_hit(
1157 + cls,
1158 + block_hashes: BlockHashList,
1159 + max_length: int,
1160 + kv_cache_group_ids: list[int],
1161 + block_pool: BlockPool,
1162 + kv_cache_spec: KVCacheSpec,
1163 + drop_eagle_block: bool,
1164 + alignment_tokens: int,
1165 + dcp_world_size: int = 1,
1166 + pcp_world_size: int = 1,
1167 + ) -> tuple[tuple[list[KVCacheBlock], ...], int]:
1168 + return tuple([] for _ in kv_cache_group_ids), 0
1169 +
1170 + def cache_blocks(
1171 + self,
1172 + request: Request,
1173 + num_tokens: int,
1174 + retention_interval: int | None = None,
1175 + ) -> None:
1176 + return
1177 +
1178 + def add_local_computed_blocks(
1179 + self,
1180 + request_id: str,
1181 + new_computed_blocks: Sequence[KVCacheBlock],
1182 + num_local_computed_tokens: int,
1183 + num_external_computed_tokens: int,
1184 + ) -> None:
1185 + return
1186 +
1187 + def remove_skipped_blocks(
1188 + self,
1189 + request_id: str,
1190 + processed_computed_tokens: int,
1191 + num_prompt_tokens: int | None = None,
1192 + ) -> None:
1193 + return
1194 +
1195 + def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
1196 + return 0
1197 +
1198 + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
1199 + return 0
1200 +
1201 +
1202 class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
1203 def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None:
1204 super().__init__(kv_cache_spec, **kwargs)
@@ -1953,6 +2043,11 @@ def register_all_kvcache_specs(vllm_config):
2043 SlidingWindowManager,
2044 uniform_type_base_spec=SlidingWindowSpec,
2045 )
2046 + KVCacheSpecRegistry.register(
2047 + CircularBufferSpec,
2048 + CircularBufferManager,
2049 + uniform_type_base_spec=CircularBufferSpec,
2050 + )
2051 KVCacheSpecRegistry.register(
2052 SlidingWindowMLASpec,
2053 SlidingWindowManager,
vllm/v1/kv_cache_interface.py
+59 -4
@@ -150,6 +150,10 @@ class KVCacheSpec:
150 # number of tokens in a block
151 block_size: int
152
153 + @property
154 + def prefix_cacheable(self) -> bool:
155 + return True
156 +
157 @property
158 def num_heads(self) -> int:
159 raise NotImplementedError
@@ -745,6 +749,38 @@ class SlidingWindowSpec(AttentionSpec):
749 )
750
751
752 +@dataclass(frozen=True, kw_only=True)
753 +class CircularBufferSpec(AttentionSpec):
754 + """One block per request holding the raw keys of the token group that
755 + is still being compressed.
756 +
757 + ``block_size`` is the ring capacity. It must exceed the compression ratio
758 + by the speculative lookahead: a speculative step stores all of its rows,
759 + drafts included, before acceptance is known, while the next step still
760 + reads the open group's committed keys from the ring.
761 + """
762 +
763 + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
764 + # The ring occupies one block per request for its whole lifetime.
765 + del vllm_config
766 + return self.page_size_bytes
767 +
768 + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
769 + del vllm_config, max_len
770 + return 1
771 +
772 + def is_uniform_with_collection(
773 + self, kv_cache_specs: dict[str, KVCacheSpec]
774 + ) -> bool:
775 + return all(
776 + isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values()
777 + )
778 +
779 + @property
780 + def prefix_cacheable(self) -> bool:
781 + return False
782 +
783 +
784 @dataclass(frozen=True, kw_only=True)
785 class SlidingWindowMLASpec(SlidingWindowSpec):
786 """Sliding window attention with MLA cache format."""
@@ -814,7 +850,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
850 @dataclass(frozen=True)
851 class MambaSpec(KVCacheSpec):
852 shapes: tuple[tuple[int, ...], ...]
817 - dtypes: tuple[torch.dtype]
853 + dtypes: tuple[torch.dtype, ...]
854 page_size_padded: int | None = None
855 mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2
856 mamba_cache_mode: str = "none"
@@ -822,6 +858,9 @@ class MambaSpec(KVCacheSpec):
858 num_prefill_checkpoint_blocks: int = 0
859 num_heads: int = 1
860 tokens_per_state: int = -1
861 + # False: the state is sharded across TP ranks (e.g. GDN). True: every TP
862 + # rank holds the full state (e.g. the replicated PLE conv state).
863 + tp_replicated: bool = False
864
865 @property
866 def state_content_size_bytes(self) -> int:
@@ -873,6 +912,8 @@ class MambaSpec(KVCacheSpec):
912 isinstance(spec, MambaSpec)
913 and spec.num_speculative_blocks == self.num_speculative_blocks
914 and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks
915 + and spec.page_size_bytes == self.page_size_bytes
916 + and spec.tp_replicated == self.tp_replicated
917 for spec in kv_cache_specs.values()
918 )
919
@@ -963,6 +1004,15 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1004
1005 kv_cache_specs: dict[str, KVCacheSpec]
1006
1007 + @property
1008 + def prefix_cacheable(self) -> bool:
1009 + return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values())
1010 +
1011 + @property
1012 + def first_spec(self) -> KVCacheSpec:
1013 + """Return the first spec in the group."""
1014 + return next(iter(self.kv_cache_specs.values()))
1015 +
1016 @property
1017 def page_size_bytes(self) -> int:
1018 return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values())
@@ -1015,8 +1065,9 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
1065 else:
1066 return None
1067
1018 - # NOTE: below util functions are only used by DeepseekV4 for now.
1019 - def get_num_layer_tuples(self) -> int:
1068 + def get_max_layers_per_page_size(self) -> int:
1069 + """Max number of layers sharing a page size. For a balanced bucket
1070 + this equals the number of repetitions of the layer pattern."""
1071 return Counter(
1072 spec.page_size_bytes for spec in self.kv_cache_specs.values()
1073 ).most_common(1)[0][1]
@@ -1212,7 +1263,11 @@ class KVCacheConfig:
1263
1264 @property
1265 def has_mamba_layers(self) -> bool:
1215 - return any(isinstance(g.kv_cache_spec, MambaSpec) for g in self.kv_cache_groups)
1266 + return any(
1267 + isinstance(spec, MambaSpec)
1268 + for group in self.kv_cache_groups
1269 + for spec in iter_layer_specs(group.kv_cache_spec)
1270 + )
1271
1272 @property
1273 def has_mixed_precision_kv_cache(self) -> bool:
vllm/v1/worker/gpu/block_table.py
+15 -1
@@ -26,6 +26,7 @@ class BlockTables:
26 cp_size: int = 1,
27 cp_rank: int = 0,
28 cp_interleave: int = 1,
29 + slot_mapping_enabled: list[bool] | None = None,
30 ):
31 self.block_sizes = block_sizes
32 self.kernel_block_sizes = kernel_block_sizes
@@ -39,6 +40,10 @@ class BlockTables:
40
41 self.num_kv_cache_groups = len(self.block_sizes)
42 assert len(max_num_blocks_per_group) == self.num_kv_cache_groups
43 + if slot_mapping_enabled is None:
44 + slot_mapping_enabled = [True] * self.num_kv_cache_groups
45 + assert len(slot_mapping_enabled) == self.num_kv_cache_groups
46 + self._slot_mapping_enabled = slot_mapping_enabled
47
48 self.blocks_per_kv_block = [
49 bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes)
@@ -100,6 +105,9 @@ class BlockTables:
105 self.kernel_block_sizes_tensor = torch.tensor(
106 self.kernel_block_sizes, dtype=torch.int32, device=self.device
107 )
108 + self.slot_mapping_enabled = torch.tensor(
109 + self._slot_mapping_enabled, dtype=torch.bool, device=self.device
110 + )
111 self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables)
112
113 def append_block_ids(
@@ -202,6 +210,7 @@ class BlockTables:
210 self.block_table_strides,
211 self.block_sizes_tensor,
212 self.kernel_block_sizes_tensor,
213 + self.slot_mapping_enabled,
214 slot_mappings,
215 slot_mappings.stride(0),
216 self.cp_rank,
@@ -274,6 +283,7 @@ def _compute_slot_mappings_kernel(
283 block_table_strides, # [num_kv_cache_groups]
284 block_sizes, # [num_kv_cache_groups]
285 kernel_block_sizes, # [num_kv_cache_groups]
286 + slot_mapping_enabled, # [num_kv_cache_groups]
287 slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens]
288 slot_mappings_stride,
289 cp_rank,
@@ -302,6 +312,7 @@ def _compute_slot_mappings_kernel(
312 block_table_stride = tl.load(block_table_strides + group_id)
313 kv_block_size = tl.load(block_sizes + group_id)
314 kernel_block_size = tl.load(kernel_block_sizes + group_id)
315 + mapping_enabled = tl.load(slot_mapping_enabled + group_id)
316
317 req_state_idx = tl.load(idx_mapping + batch_idx)
318 start_idx = tl.load(query_start_loc + batch_idx)
@@ -325,7 +336,9 @@ def _compute_slot_mappings_kernel(
336 local_offsets = rounds * CP_INTERLEAVE + remainder
337 local_positions = virtual_block_indices * kv_block_size + local_offsets
338
328 - block_indices = local_positions // kernel_block_size
339 + block_indices = tl.where(
340 + mapping_enabled, local_positions // kernel_block_size, 0
341 + )
342 block_offsets = local_positions % kernel_block_size
343 block_numbers = tl.load(
344 block_table_ptr + req_state_idx * block_table_stride + block_indices,
@@ -336,4 +349,5 @@ def _compute_slot_mappings_kernel(
349 if CP_SIZE != 1:
350 slot_ids = tl.where(is_local, slot_ids, PAD_ID)
351
352 + slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID)
353 tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
vllm/v1/worker/gpu/model_runner.py
+17 -2
@@ -64,7 +64,12 @@ from vllm.tasks import SupportedTask
64 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
65 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
66 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
67 -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
67 +from vllm.v1.kv_cache_interface import (
68 + CircularBufferSpec,
69 + KVCacheConfig,
70 + MambaSpec,
71 + UniformTypeKVCacheSpecs,
72 +)
73 from vllm.v1.outputs import (
74 DraftTokenIds,
75 ECConnectorOutput,
@@ -548,9 +553,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
553
554 block_sizes = []
555 max_num_blocks_per_group = []
556 + slot_mapping_enabled = []
557 for kv_cache_group in kv_cache_config.kv_cache_groups:
558 spec = kv_cache_group.kv_cache_spec
559 block_sizes.append(spec.block_size)
560 + layer_spec = (
561 + spec.first_spec if isinstance(spec, UniformTypeKVCacheSpecs) else spec
562 + )
563 + slot_mapping_enabled.append(not isinstance(layer_spec, CircularBufferSpec))
564 # Let each cache type account for CP. Attention KV is DCP-sharded,
565 # while Mamba/GDN recurrent state is replicated across DCP ranks.
566 max_num_blocks = spec.max_num_blocks_per_req(
@@ -558,7 +568,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
568 )
569 # Preserve each cache type's alignment requirements after applying
570 # its topology-aware block-table width.
561 - if isinstance(spec, MambaSpec):
571 + if isinstance(layer_spec, (MambaSpec, CircularBufferSpec)):
572 max_num_blocks = get_block_table_width(
573 max_num_blocks, spec.block_size, token_alignment=None
574 )
@@ -605,6 +615,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
615 max_num_blocks_per_group=max_num_blocks_per_group,
616 device=self.device,
617 kernel_block_sizes=self.kernel_block_sizes,
618 + slot_mapping_enabled=slot_mapping_enabled,
619 cp_size=self.dcp_size,
620 cp_rank=self.dcp_rank,
621 cp_interleave=self.cp_interleave,
@@ -1192,6 +1203,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1203 total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
1204 total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
1205 num_logits = num_draft_tokens_per_req + num_bonus_tokens
1206 + # combine_sampled_and_draft_tokens places a request's logits rows
1207 + # at [query_end - num_logits, query_end). Fewer query rows than
1208 + # that would silently select the preceding request's hidden states.
1209 + assert (num_scheduled_tokens_np >= num_logits).all()
1210 cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
1211 cu_num_logits_np[0] = 0
1212 np.cumsum(num_logits, out=cu_num_logits_np[1:])
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+29 -15
@@ -9,10 +9,14 @@ import torch.nn as nn
9
10 from vllm.config import VllmConfig
11 from vllm.config.compilation import CUDAGraphMode
12 +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType
13 from vllm.triton_utils import tl, triton
14 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
15 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
15 -from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadataBuilder
16 +from vllm.v1.attention.backends.short_conv_attn import (
17 + PleShortConvAttentionMetadataBuilder,
18 + ShortConvAttentionMetadataBuilder,
19 +)
20 from vllm.v1.core.sched.output import NewRequestData
21 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
22 from vllm.v1.utils import CpuGpuBuffer
@@ -24,7 +28,10 @@ from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
28 from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
29 from vllm.v1.worker.mamba_utils import (
30 MambaSpecDecodeGPUContext,
31 + get_mamba_group_ids,
32 + get_mamba_groups,
33 preprocess_mamba_align_fused_kernel,
34 + validate_mamba_state_copy_funcs,
35 )
36 from vllm.v1.worker.utils import AttentionGroup
37
@@ -53,6 +60,7 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
60 Mamba2AttentionMetadataBuilder,
61 GDNAttentionMetadataBuilder,
62 ShortConvAttentionMetadataBuilder,
63 + PleShortConvAttentionMetadataBuilder,
64 ),
65 ):
66 return {}
@@ -102,6 +110,7 @@ class MambaHybridModelState(DefaultModelState):
110 self._mamba_ctx: MambaSpecDecodeGPUContext | None = None
111 self._mamba_group_ids: list[int] = []
112 self._mamba_spec: MambaSpec | None = None
113 + self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None
114
115 def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
116 super().add_request(req_index, new_req_data)
@@ -117,17 +126,16 @@ class MambaHybridModelState(DefaultModelState):
126 self, kv_cache_config: KVCacheConfig
127 ) -> tuple[list[int], MambaSpec]:
128 if self._mamba_spec is None:
120 - group_ids: list[int] = []
121 - specs: list[MambaSpec] = []
122 - for i, group in enumerate(kv_cache_config.kv_cache_groups):
123 - spec = group.kv_cache_spec
124 - if isinstance(spec, MambaSpec):
125 - group_ids.append(i)
126 - specs.append(spec)
127 - assert specs, "no mamba layers in the model"
128 - assert all(specs[0] == s for s in specs)
129 - self._mamba_group_ids = group_ids
130 - self._mamba_spec = specs[0]
129 + mamba_groups = get_mamba_groups(kv_cache_config)
130 + mamba_spec = next(iter(mamba_groups))
131 + assert all(
132 + spec.block_size == mamba_spec.block_size
133 + and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks
134 + and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode
135 + for spec in mamba_groups
136 + ), "all mamba groups must share cache scheduling parameters"
137 + self._mamba_group_ids = get_mamba_group_ids(mamba_groups)
138 + self._mamba_spec = mamba_spec
139 return self._mamba_group_ids, self._mamba_spec
140
141 def _ensure_align_ctx(
@@ -136,8 +144,14 @@ class MambaHybridModelState(DefaultModelState):
144 mamba_group_ids: list[int],
145 block_tables: tuple[torch.Tensor, ...],
146 ) -> MambaSpecDecodeGPUContext:
147 + if self._mamba_state_copy_funcs is None:
148 + mamba_groups = get_mamba_groups(kv_cache_config)
149 + mamba_types = {spec.mamba_type for spec in mamba_groups}
150 + copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types)
151 + validate_mamba_state_copy_funcs(mamba_groups, copy_funcs)
152 + self._mamba_state_copy_funcs = copy_funcs
153 + copy_funcs = self._mamba_state_copy_funcs
154 if self._mamba_ctx is None:
140 - copy_funcs = self.model.get_mamba_state_copy_func()
155 # Both SD and DS conv layouts support a >0 spec-decode shift: the
156 # fused pre-copy kernel (``_copy_mamba_state_block``) applies the
157 # ``token_bias = num_accepted - 1`` window shift per conv layout
@@ -146,7 +160,7 @@ class MambaHybridModelState(DefaultModelState):
160 self._mamba_ctx = MambaSpecDecodeGPUContext.create(
161 max_num_reqs=self.max_num_reqs,
162 kv_cache_config=kv_cache_config,
149 - num_state_types=len(copy_funcs),
163 + copy_funcs=copy_funcs,
164 device=self.device,
165 make_buffer=lambda n, dtype: CpuGpuBuffer(
166 n, dtype=dtype, device=self.device
@@ -161,7 +175,7 @@ class MambaHybridModelState(DefaultModelState):
175 ctx.initialize_from_forward_context(
176 kv_cache_config,
177 forward_context,
164 - self.model.get_mamba_state_copy_func(),
178 + copy_funcs,
179 [block_tables[gid] for gid in mamba_group_ids],
180 )
181 return ctx
vllm/v1/worker/gpu/warmup.py
+12 -1
@@ -18,7 +18,13 @@ from vllm.v1.core.sched.output import (
18 NewRequestData,
19 SchedulerOutput,
20 )
21 -from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec, MambaSpec
21 +from vllm.v1.kv_cache_interface import (
22 + CircularBufferSpec,
23 + CrossAttentionSpec,
24 + KVCacheSpec,
25 + MambaSpec,
26 + UniformTypeKVCacheSpecs,
27 +)
28 from vllm.v1.request import Request
29 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
30
@@ -39,6 +45,11 @@ def _reserved_block_count(
45 `KVCacheManager.allocate_slots` reserves: the token range plus
46 `num_lookahead_tokens`, where the speculator writes the KV of its drafts.
47 """
48 + if isinstance(kvcache_spec, UniformTypeKVCacheSpecs):
49 + kvcache_spec = kvcache_spec.first_spec
50 + if isinstance(kvcache_spec, CircularBufferSpec):
51 + # Circular caches keep one physical ring block for the request lifetime.
52 + return 1
53 if isinstance(kvcache_spec, CrossAttentionSpec):
54 # Cross-attention blocks cover the encoder sequence only.
55 return cdiv(max_encoder_len, kvcache_spec.block_size)
vllm/v1/worker/gpu_model_runner.py
+15 -3
@@ -66,6 +66,7 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
66 RoutedExpertsCapturer,
67 bind_routed_experts_capturer,
68 )
69 +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType
70 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
71 initialize_mamba_ssu_backend,
72 )
@@ -997,6 +998,7 @@ class GPUModelRunner(
998 self.kv_connector_output: KVConnectorOutput | None = None
999 self.mamba_state_idx: dict[str, int] = {}
1000 self._mamba_bufs: mamba_utils.MambaBuffers | None = None
1001 + self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None
1002 self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None
1003 if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0:
1004 self.mamba_prev_last_scheduled_idx = self._make_buffer(
@@ -1053,6 +1055,15 @@ class GPUModelRunner(
1055 with_numpy=numpy,
1056 )
1057
1058 + def _get_mamba_state_copy_funcs(self) -> MambaStateCopyFuncsByType:
1059 + if self._mamba_state_copy_funcs is None:
1060 + mamba_groups = mamba_utils.get_mamba_groups(self.kv_cache_config)
1061 + mamba_types = {spec.mamba_type for spec in mamba_groups}
1062 + copy_funcs = self.model.get_mamba_state_copy_funcs(mamba_types)
1063 + mamba_utils.validate_mamba_state_copy_funcs(mamba_groups, copy_funcs)
1064 + self._mamba_state_copy_funcs = copy_funcs
1065 + return self._mamba_state_copy_funcs
1066 +
1067 def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers:
1068 # Only reachable on the ``mamba_cache_mode == "align"`` path.
1069 # The postprocess sub-object is additionally gated on spec
@@ -1062,7 +1073,7 @@ class GPUModelRunner(
1073 self._mamba_bufs = mamba_utils.MambaBuffers.create(
1074 max_num_reqs=self.max_num_reqs,
1075 kv_cache_config=self.kv_cache_config,
1065 - copy_funcs=self.model.get_mamba_state_copy_func(),
1076 + copy_funcs=self._get_mamba_state_copy_funcs(),
1077 make_buffer=self._make_buffer,
1078 device=self.device,
1079 with_postprocess_align=(
@@ -1599,7 +1610,7 @@ class GPUModelRunner(
1610 input_batch=self.input_batch,
1611 kv_cache_config=self.kv_cache_config,
1612 forward_context=self.compilation_config.static_forward_context,
1602 - mamba_state_copy_funcs=self.model.get_mamba_state_copy_func(),
1613 + mamba_state_copy_funcs=self._get_mamba_state_copy_funcs(),
1614 )
1615
1616 assert self.num_accepted_tokens_event is not None
@@ -4404,7 +4415,7 @@ class GPUModelRunner(
4415 self.input_batch,
4416 self.requests,
4417 self.compilation_config.static_forward_context,
4407 - self.model.get_mamba_state_copy_func(),
4418 + self._get_mamba_state_copy_funcs(),
4419 mamba_bufs.preprocess,
4420 align_ctx=mamba_bufs.postprocess_align,
4421 )
@@ -7488,6 +7499,7 @@ class GPUModelRunner(
7499 kv_cache_config = deepcopy(kv_cache_config)
7500 self.kv_cache_config = kv_cache_config
7501 self._mamba_bufs = None
7502 + self._mamba_state_copy_funcs = None
7503 self.may_add_encoder_only_layers_to_kv_cache_config()
7504 self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
7505 self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
vllm/v1/worker/mamba_utils.py
+143 -58
@@ -10,7 +10,7 @@ import torch
10 from vllm.config import CacheConfig
11 from vllm.logger import init_logger
12 from vllm.model_executor.layers.mamba.mamba_utils import (
13 - MambaStateCopyFunc,
13 + MambaStateCopyFuncsByType,
14 get_conv_copy_spec,
15 get_temporal_copy_spec,
16 is_conv_state_dim_first,
@@ -19,7 +19,12 @@ from vllm.triton_utils import tl, triton
19 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
20 from vllm.utils.math_utils import cdiv
21 from vllm.v1.core.sched.output import SchedulerOutput
22 -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
22 +from vllm.v1.kv_cache_interface import (
23 + KVCacheConfig,
24 + KVCacheGroupSpec,
25 + MambaSpec,
26 + UniformTypeKVCacheSpecs,
27 +)
28 from vllm.v1.utils import CpuGpuBuffer
29 from vllm.v1.worker.gpu_input_batch import CachedRequestState
30 from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
@@ -374,8 +379,8 @@ def postprocess_mamba_fused_kernel(
379 # block table.
380 block_table_ptrs_ptr,
381 block_table_stride_req: tl.int64, # stride between requests (in elements)
377 - # Mamba state metadata (per-layer, per-state-type)
378 - # These are 1D arrays indexed by (layer_idx * num_state_types + state_type_idx)
382 + # Mamba state metadata, flattened across all layer/state tensors. Different
383 + # Mamba groups may contribute different numbers of state tensors.
384 state_base_addrs_ptr, # base address of each state tensor
385 state_block_strides_ptr, # bytes per block for each state
386 state_elem_sizes_ptr, # element size for each state
@@ -415,14 +420,13 @@ def postprocess_mamba_fused_kernel(
420 Fused GPU kernel for postprocess_mamba that computes decisions AND performs
421 mamba state copies without any CPU-GPU synchronization.
422
418 - Grid: (num_reqs, num_layers * num_state_types [, TEMPORAL_TILES])
423 + Grid: (num_reqs, num_states [, TEMPORAL_TILES])
424 - program_id(0) = request/batch index
425 - program_id(1) = state_idx (flattened index into layer/state_type metadata)
426 - program_id(2) = temporal-copy tile index (0 when TEMPORAL_TILES == 1)
427
423 - Note: num_layers and num_state_types are not passed as kernel parameters
424 - because the kernel indexes directly into pre-flattened metadata arrays
425 - using program_id(1). The grid dimensions encode the total state count.
428 + The kernel indexes directly into pre-flattened metadata arrays using
429 + program_id(1). The grid dimensions encode the total state count.
430 """
431 batch_idx = tl.program_id(0)
432 state_idx = tl.program_id(1)
@@ -581,8 +585,8 @@ def precopy_mamba_align_fused_kernel(
585 copy specs), but driven by the GPU-resident src columns so it needs no
586 CPU-GPU sync (async-scheduling safe).
587
584 - Grid: (num_reqs, num_layers * num_state_types [, TEMPORAL_TILES]). V2 passes
585 - a batch-to-state idx_mapping; V1 already stores the staged arrays in batch
588 + Grid: (num_reqs, num_states [, TEMPORAL_TILES]). V2 passes a
589 + batch-to-state idx_mapping; V1 already stores the staged arrays in batch
590 order and uses HAS_IDX_MAPPING=False.
591 """
592 batch_idx = tl.program_id(0)
@@ -663,17 +667,56 @@ def batch_memcpy(src_ptrs, dst_ptrs, sizes):
667 batch_memcpy_kernel[grid](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=BLOCK_SIZE)
668
669
666 -def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSpec]:
667 - mamba_group_ids: list[int] = []
668 - mamba_specs: list[MambaSpec] = []
669 - for i in range(len(kv_cache_config.kv_cache_groups)):
670 - kv_cache_spec = kv_cache_config.kv_cache_groups[i].kv_cache_spec
671 - if isinstance(kv_cache_spec, MambaSpec):
672 - mamba_group_ids.append(i)
673 - mamba_specs.append(kv_cache_spec)
674 - assert len(mamba_group_ids) > 0, "no mamba layers in the model"
675 - assert all(mamba_specs[0] == spec for spec in mamba_specs)
676 - return mamba_group_ids, mamba_specs[0]
670 +def _get_mamba_spec_for_layer(
671 + kv_cache_group: KVCacheGroupSpec, layer_name: str
672 +) -> MambaSpec:
673 + kv_cache_spec = kv_cache_group.kv_cache_spec
674 + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
675 + kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
676 + assert isinstance(kv_cache_spec, MambaSpec)
677 + return kv_cache_spec
678 +
679 +
680 +def get_mamba_group_ids(mamba_groups: dict[MambaSpec, list[int]]) -> list[int]:
681 + """Return the sorted group ids of all mamba groups, without duplicates.
682 +
683 + A group whose layers use different ``MambaSpec``s is listed under one key
684 + per spec, so the same group id can appear more than once.
685 + """
686 + return sorted(
687 + {group_id for group_ids in mamba_groups.values() for group_id in group_ids}
688 + )
689 +
690 +
691 +def get_mamba_groups(kv_cache_config: KVCacheConfig) -> dict[MambaSpec, list[int]]:
692 + """Return a mapping from each distinct MambaSpec to its sorted group_ids."""
693 + mamba_groups: dict[MambaSpec, set[int]] = {}
694 + for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
695 + kv_cache_spec = kv_cache_group.kv_cache_spec
696 + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
697 + kv_cache_spec = kv_cache_spec.first_spec
698 + if not isinstance(kv_cache_spec, MambaSpec):
699 + continue
700 + for layer_name in kv_cache_group.layer_names:
701 + mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name)
702 + mamba_groups.setdefault(mamba_spec, set()).add(i)
703 + assert len(mamba_groups) > 0, "no mamba layers in the model"
704 + return {spec: sorted(ids) for spec, ids in mamba_groups.items()}
705 +
706 +
707 +def validate_mamba_state_copy_funcs(
708 + mamba_groups: dict[MambaSpec, list[int]],
709 + copy_funcs: MambaStateCopyFuncsByType,
710 +) -> None:
711 + for mamba_spec in mamba_groups:
712 + assert mamba_spec.mamba_type in copy_funcs, (
713 + f"missing state copy funcs for {mamba_spec.mamba_type}"
714 + )
715 + state_copy_funcs = copy_funcs[mamba_spec.mamba_type]
716 + assert len(state_copy_funcs) == len(mamba_spec.shapes), (
717 + f"{mamba_spec.mamba_type} declares {len(mamba_spec.shapes)} states, "
718 + f"but provides {len(state_copy_funcs)} state copy funcs"
719 + )
720
721
722 @dataclasses.dataclass
@@ -690,14 +733,29 @@ class MambaCopyBuffers:
733 cls,
734 max_num_reqs: int,
735 kv_cache_config: KVCacheConfig,
693 - copy_funcs: tuple[MambaStateCopyFunc, ...],
736 + copy_funcs: MambaStateCopyFuncsByType,
737 make_buffer: Callable[..., CpuGpuBuffer],
738 ) -> "MambaCopyBuffers":
696 - mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
739 + mamba_groups = get_mamba_groups(kv_cache_config)
740 + mamba_spec = next(iter(mamba_groups))
741 + assert all(
742 + spec.block_size == mamba_spec.block_size
743 + and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks
744 + and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode
745 + for spec in mamba_groups
746 + ), "all mamba groups must share cache scheduling parameters"
747 + mamba_group_ids = get_mamba_group_ids(mamba_groups)
748 entries_per_req = sum(
698 - len(kv_cache_config.kv_cache_groups[gid].layer_names)
749 + len(
750 + copy_funcs[
751 + _get_mamba_spec_for_layer(
752 + kv_cache_config.kv_cache_groups[gid], layer_name
753 + ).mamba_type
754 + ]
755 + )
756 for gid in mamba_group_ids
700 - ) * len(copy_funcs)
757 + for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names
758 + )
759 n = max_num_reqs * entries_per_req
760
761 return cls(
@@ -725,7 +783,9 @@ class MambaSpecDecodeGPUContext:
783 window with offset-based copies), 0 for temporal states (full block copies).
784 """
785
728 - # Per-state metadata tensors (shape: [num_layers * num_state_types])
786 + # Per-state metadata tensors (shape: [num_states]). Mamba groups may have
787 + # different numbers of state types (for example, GDN has two while PLE
788 + # short-conv has one).
789 # These are populated from forward_context during the first forward pass
790 state_base_addrs: torch.Tensor # int64: base address of each state tensor
791 state_block_strides: torch.Tensor # int64: bytes per block
@@ -739,8 +799,7 @@ class MambaSpecDecodeGPUContext:
799
800 # Configuration
801 block_size: int
742 - num_layers: int
743 - num_state_types: int
802 + num_states: int
803 mamba_group_ids: list[int]
804 num_groups: int
805
@@ -776,19 +835,37 @@ class MambaSpecDecodeGPUContext:
835 cls,
836 max_num_reqs: int,
837 kv_cache_config: KVCacheConfig,
779 - num_state_types: int,
838 + copy_funcs: MambaStateCopyFuncsByType,
839 device: torch.device,
840 make_buffer: Callable[..., CpuGpuBuffer],
841 ) -> "MambaSpecDecodeGPUContext":
842 """Create context with allocated buffers (metadata populated later)."""
784 - mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
785 -
786 - # Count total layers across all mamba groups
787 - num_layers = sum(
788 - len(kv_cache_config.kv_cache_groups[gid].layer_names)
843 + mamba_groups = get_mamba_groups(kv_cache_config)
844 + mamba_group_ids = get_mamba_group_ids(mamba_groups)
845 + mamba_spec = next(iter(mamba_groups))
846 + assert all(
847 + spec.block_size == mamba_spec.block_size
848 + and spec.num_speculative_blocks == mamba_spec.num_speculative_blocks
849 + and spec.mamba_cache_mode == mamba_spec.mamba_cache_mode
850 + for spec in mamba_groups
851 + ), "all mamba groups must share cache scheduling parameters"
852 + copy_funcs_by_spec = {
853 + spec: copy_funcs[spec.mamba_type] for spec in mamba_groups
854 + }
855 +
856 + # Count physical state tensors across all Mamba groups. Different
857 + # groups may expose different state specs.
858 + total_states = sum(
859 + len(
860 + copy_funcs_by_spec[
861 + _get_mamba_spec_for_layer(
862 + kv_cache_config.kv_cache_groups[gid], layer_name
863 + )
864 + ]
865 + )
866 for gid in mamba_group_ids
867 + for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names
868 )
791 - total_states = num_layers * num_state_types
869
870 return cls(
871 state_base_addrs=torch.zeros(
@@ -816,8 +893,7 @@ class MambaSpecDecodeGPUContext:
893 total_states, dtype=torch.int64, device=device
894 ),
895 block_size=mamba_spec.block_size,
819 - num_layers=num_layers,
820 - num_state_types=num_state_types,
896 + num_states=total_states,
897 mamba_group_ids=mamba_group_ids,
898 num_groups=len(mamba_group_ids),
899 num_accepted_tokens_out=torch.zeros(
@@ -848,7 +924,7 @@ class MambaSpecDecodeGPUContext:
924 self,
925 kv_cache_config: KVCacheConfig,
926 forward_context: dict[str, Any],
851 - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
927 + mamba_state_copy_funcs: MambaStateCopyFuncsByType,
928 block_tables: list[torch.Tensor],
929 ) -> None:
930 """
@@ -879,8 +955,8 @@ class MambaSpecDecodeGPUContext:
955 forward_context: Dictionary mapping layer names to attention objects,
956 populated after the model is loaded. Each attention object must
957 have a `kv_cache` attribute containing the list of state tensors.
882 - mamba_state_copy_funcs: Tuple of copy functions (one per state type)
883 - used to determine whether each state is a conv or temporal state.
958 + mamba_state_copy_funcs: Mapping from MambaAttentionBackendEnum to
959 + copy functions.
960 block_tables: per-mamba-group persistent block-table tensors, in
961 the same order as `mamba_group_ids`. Their `data_ptr()` /
962 `stride(0)` are captured once for the kernel to index into.
@@ -900,22 +976,24 @@ class MambaSpecDecodeGPUContext:
976 self,
977 kv_cache_config: KVCacheConfig,
978 forward_context: dict[str, Any],
903 - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
979 + mamba_state_copy_funcs: MambaStateCopyFuncsByType,
980 block_tables: list[torch.Tensor],
981 ) -> None:
982 idx = 0
983 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids):
908 - layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names
984 + kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id]
985 + layer_names = kv_cache_group.layer_names
986 for layer_name in layer_names:
987 + mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name)
988 + state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type]
989 attention = forward_context[layer_name]
990 kv_caches: list[torch.Tensor] = attention.kv_cache
912 -
913 - if len(kv_caches) < self.num_state_types:
991 + if len(kv_caches) < len(state_copy_funcs):
992 raise ValueError(
915 - f"Expected at least {self.num_state_types} Mamba state "
993 + f"Expected at least {len(state_copy_funcs)} Mamba state "
994 f"tensors, got {len(kv_caches)}"
995 )
918 - for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs):
996 + for state_type_idx, copy_func in enumerate(state_copy_funcs):
997 state = kv_caches[state_type_idx]
998 # Base address
999 self.state_base_addrs[idx] = _reinterpret_u64_as_i64(
@@ -994,6 +1072,8 @@ class MambaSpecDecodeGPUContext:
1072 self.state_group_indices[idx] = group_local_idx
1073 idx += 1
1074
1075 + assert idx == self.num_states
1076 +
1077 # Cache per-group block-table base addresses and per-request stride.
1078 # `block_tables[i]` is the persistent 2D int32 block-table tensor for
1079 # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the
@@ -1079,7 +1159,7 @@ class MambaSpecDecodeGPUContext:
1159 num_accepted_tokens_gpu[:num_reqs]
1160 )
1161
1082 - total_states = self.num_layers * self.num_state_types
1162 + total_states = self.num_states
1163 grid = (num_reqs, total_states, _TEMPORAL_TILES)
1164
1165 postprocess_mamba_fused_kernel[grid](
@@ -1128,7 +1208,7 @@ class MambaSpecDecodeGPUContext:
1208 """
1209 if num_reqs == 0 or not self.is_initialized:
1210 return
1131 - total_states = self.num_layers * self.num_state_types
1211 + total_states = self.num_states
1212 grid = (num_reqs, total_states, _TEMPORAL_TILES)
1213 precopy_mamba_align_fused_kernel[grid](
1214 state_idx_gpu,
@@ -1178,7 +1258,7 @@ class MambaSpecDecodeGPUContext:
1258 num_accepted_tokens_snapshot = self.num_accepted_tokens_out
1259 num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu)
1260
1181 - total_states = self.num_layers * self.num_state_types
1261 + total_states = self.num_states
1262 grid = (num_reqs, total_states, _TEMPORAL_TILES)
1263 postprocess_mamba_fused_kernel[grid](
1264 num_accepted_tokens_snapshot,
@@ -1226,7 +1306,7 @@ class MambaBuffers:
1306 cls,
1307 max_num_reqs: int,
1308 kv_cache_config: KVCacheConfig,
1229 - copy_funcs: tuple[MambaStateCopyFunc, ...],
1309 + copy_funcs: MambaStateCopyFuncsByType,
1310 make_buffer: Callable[..., CpuGpuBuffer],
1311 device: torch.device,
1312 with_postprocess_align: bool,
@@ -1239,7 +1319,7 @@ class MambaBuffers:
1319 MambaSpecDecodeGPUContext.create(
1320 max_num_reqs=max_num_reqs,
1321 kv_cache_config=kv_cache_config,
1242 - num_state_types=len(copy_funcs),
1322 + copy_funcs=copy_funcs,
1323 device=device,
1324 make_buffer=make_buffer,
1325 )
@@ -1252,7 +1332,7 @@ class MambaBuffers:
1332 def collect_mamba_copy_meta(
1333 copy_bufs: MambaCopyBuffers,
1334 kv_cache_config: KVCacheConfig,
1255 - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
1335 + mamba_state_copy_funcs: MambaStateCopyFuncsByType,
1336 mamba_group_ids: list[int],
1337 src_block_idx: int,
1338 dest_block_idx: int,
@@ -1271,11 +1351,14 @@ def collect_mamba_copy_meta(
1351 for mamba_group_id in mamba_group_ids:
1352 block_ids = req_state.block_ids[mamba_group_id]
1353 dest_block_id = block_ids[dest_block_idx]
1274 - layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names
1354 + kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id]
1355 + layer_names = kv_cache_group.layer_names
1356 for layer_name in layer_names:
1357 + mamba_spec = _get_mamba_spec_for_layer(kv_cache_group, layer_name)
1358 + state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type]
1359 attention = forward_context[layer_name]
1360 kv_caches: list[torch.Tensor] = attention.kv_cache
1278 - for state, state_copy_func in zip(kv_caches, mamba_state_copy_funcs):
1361 + for state, state_copy_func in zip(kv_caches, state_copy_funcs):
1362 copy_spec = state_copy_func(
1363 state, block_ids, src_block_idx, accept_token_bias + 1
1364 )
@@ -1355,7 +1438,7 @@ def preprocess_mamba(
1438 input_batch: GPUInputBatch,
1439 requests: dict[str, CachedRequestState],
1440 forward_context: dict[str, Any],
1358 - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
1441 + mamba_state_copy_funcs: MambaStateCopyFuncsByType,
1442 copy_bufs: MambaCopyBuffers,
1443 align_ctx: MambaSpecDecodeGPUContext | None = None,
1444 ):
@@ -1470,8 +1553,10 @@ def postprocess_mamba_all(
1553 """
1554 if num_spec_tokens <= 0:
1555 return
1473 - _, mamba_spec = get_mamba_groups(kv_cache_config)
1474 - block_size = mamba_spec.block_size
1556 + mamba_groups = get_mamba_groups(kv_cache_config)
1557 + block_sizes = {mamba_spec.block_size for mamba_spec in mamba_groups}
1558 + assert len(block_sizes) == 1, "all mamba groups must share block_size"
1559 + block_size = next(iter(block_sizes))
1560 full_decode_len = 1 + num_spec_tokens
1561 scheduled = scheduler_output.num_scheduled_tokens
1562 for req_id in input_batch.req_ids[:num_reqs]:
@@ -1508,7 +1593,7 @@ def postprocess_mamba_align_gpu(
1593 input_batch: GPUInputBatch,
1594 kv_cache_config: KVCacheConfig,
1595 forward_context: dict[str, Any],
1511 - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
1596 + mamba_state_copy_funcs: MambaStateCopyFuncsByType,
1597 ) -> None:
1598 """GPU-side mamba postprocess for spec decode + hybrid + align mode.
1599