From a3cc771526dd19b07e230f0762e463add6092d67 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:04:50 +0800 Subject: [PATCH 01/28] [Model] Start Qwen3.8 Flash Next SM70 adaptation Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 90 +++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 docs/design/sm70_qwen38_flash_next_nvfp4.md diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md new file mode 100644 index 0000000000..b76c30c842 --- /dev/null +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -0,0 +1,90 @@ +# Qwen3.8 Flash Next NVFP4 on SM70 + +## Status and ownership + +- Status: bring-up in progress; no route, quality, memory, or speed claim yet. +- Integration line: `private/main`. +- Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. +- Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision + `7b719225242aacd3dbd3f9407468c2ee9a9d2594`. +- Model download: `/data/models/RadixArk/Qwen3.8-Flash-Next-NVFP4`. +- Upstream references: + [vLLM PR 53896](https://github.com/vllm-project/vllm/pull/53896) and + [SGLang PR 36497](https://github.com/sgl-project/sglang/pull/36497). + +The upstream PRs are implementation references, not acceptance evidence. Both +were still open when this work started, so imports must be narrowed to the +Qwen4Exp route and validated against this branch. + +## Frozen first-pass contract + +The first correctness route deliberately excludes speculative decoding. + +- Hardware: four NVIDIA V100-SXM2-32GB GPUs (SM70). +- Parallelism: TP4, PP1, no expert parallelism. +- Compute dtype: FP16; no BF16 or native FP8/NVFP4 tensor-core assumptions. +- Checkpoint: ModelOpt NVFP4 routed-expert weights, consumed as an SM70 + weight-only W4A16 route. Ignored dense, attention, GDN, shared-expert, GR, + PLE, and LM-head modules stay in their checkpoint dtypes and execute in + FP16 where required. +- PLE/N-gram table: allocate and load each TP shard directly in pinned host + memory. It must never be materialized on a GPU before being moved to the + host. Gathered rows are transferred asynchronously and converted to FP16. +- Initial KV cache: FP16. FP8 KV cache is a separate, quality-gated follow-up. +- Initial decoding: MTP disabled. MTP may be enabled only after the no-MTP + route is correct and its emitted-token baseline is recorded. + +## Architecture facts that affect the port + +The text stack has 48 layers: 36 gated-delta-net layers and 12 QSA layers in a +3:1 pattern. Hidden size is 2560. QSA uses 24 query heads, two KV heads, +head-dimension 256, index dimension 128, compression ratio four, and a 2048 +token sparse budget. The MoE has 512 routed experts, top-10 routing, a 640-wide +routed expert, and one 640-wide shared expert. General residual connections +use four streams and rank 320. + +PLE is a learned trigram embedding, not prompt-ngram speculative decoding. It +uses 16 heads (`ngram_size=3`, eight heads per n-gram order), embedding width +2560, and FP8 E4M3 storage. Native speculative decoding is the separate +one-layer MTP head. + +## Memory budget hypothesis + +Safetensor payloads total about 125.910 GiB. The sharded PLE payload is about +47.684 GiB, or about 11.921 GiB of pinned host memory per TP rank. Removing PLE +from device residency leaves an idealized 19.556 GiB of checkpoint payload per +GPU before replicated tensors, KV/index caches, CUDA graphs, and workspaces. + +This is a planning bound, not a measured peak. Startup must record host RSS, +pinned memory, per-rank device peak, post-load device residency, and whether a +loader creates duplicate staging buffers. A 262144-token context is admitted +only after the measured peak leaves a safe margin on every 32GB GPU. + +## Acceptance gates + +1. Static route: Transformers config, model registry, multimodal processor, + QSA/GDN/GR/PLE modules, and ModelOpt NVFP4 mapping load without importing an + Ampere-only backend. +2. Loader route: TP4 expert shards select TurboMind SM70 W4A16; PLE shards are + born on pinned CPU memory and do not consume persistent device memory. +3. Numerical route: focused operator comparisons against FP32/FP16 references, + followed by deterministic token-ID and output checks on the full model. +4. Memory route: 32K bring-up first, then 128K and the exact 262144 boundary; + report controlled OOM separately from corrupted output. +5. Performance route: one request, TP4, PP1, MTP off, FP16 activations, 8192 + input tokens and 512 output tokens. Report TTFT/prefill separately and + calculate steady pure decode from emitted tokens 33-512. The target is at + least 100 emitted tokens/s (at most 10 ms/token) with CUDA graphs enabled. + Record an otherwise identical eager control. +6. MTP follow-up: report accepted length, target passes, emitted tokens/s, and + output quality separately; do not compare accepted candidates with emitted + tokens. + +## Initial implementation boundary + +Reuse the upstream Qwen4Exp Python structure and tests where they match this +tree. Do not import unrelated AMD, SM90, build-system, or broad engine changes. +The first SM70-specific changes are limited to genericizing the existing +TurboMind NVFP4 MoE shape contract, adding the QSA/indexer route, and adding a +pinned-host PLE loader/gather path. Optimize GDN, GR, sparse attention, and MTP +only after profiles identify them as measured decode bottlenecks. From a90419ba0b11b0f72c2e6ac9b4a138792c1a79db Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:06:40 +0800 Subject: [PATCH 02/28] [Model] Add Qwen3.8 Flash Next SM70 V2 route Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 97 +- .../test_weights_mapper_stacked.py | 76 + tests/models/qwen4_exp/__init__.py | 2 + tests/models/qwen4_exp/test_config.py | 271 ++++ tests/models/qwen4_exp/test_ple.py | 478 ++++++ tests/models/qwen4_exp/test_weight_loading.py | 131 ++ .../test_sm70_modelopt_mixed_nvfp4.py | 43 +- tests/v1/core/test_qwen4_exp_kv_cache.py | 198 +++ tests/v1/worker/test_gpu_warmup_blocks.py | 23 +- tests/v1/worker/test_qwen4_exp_ngram.py | 80 + tests/v1/worker/test_qwen4_exp_v2.py | 134 ++ vllm/config/compilation.py | 2 + vllm/config/vllm.py | 19 +- vllm/model_executor/layers/linear.py | 24 + vllm/model_executor/layers/mamba/abstract.py | 5 + .../layers/mamba/mamba_utils.py | 8 +- .../layers/quantization/modelopt.py | 35 + .../layers/quantization/nvfp4_sm70_moe.py | 116 +- .../layers/vocab_parallel_embedding.py | 18 +- vllm/model_executor/model_loader/utils.py | 3 +- vllm/model_executor/models/config.py | 73 + vllm/model_executor/models/interfaces.py | 14 +- vllm/model_executor/models/registry.py | 8 + vllm/model_executor/models/utils.py | 91 +- vllm/models/qwen4_exp/__init__.py | 61 + vllm/models/qwen4_exp/common/__init__.py | 17 + .../qwen4_exp/common/hyperconnection.py | 250 ++++ vllm/models/qwen4_exp/common/ple.py | 75 + vllm/models/qwen4_exp/common/qsa_cache.py | 824 +++++++++++ vllm/models/qwen4_exp/config.py | 252 ++++ vllm/models/qwen4_exp/nvidia/__init__.py | 2 + .../qwen4_exp/nvidia/hyperconnection.py | 205 +++ vllm/models/qwen4_exp/nvidia/indexer_qsa.py | 355 +++++ .../qwen4_exp/nvidia/low_latency_gemm.py | 185 +++ vllm/models/qwen4_exp/nvidia/model.py | 1070 ++++++++++++++ vllm/models/qwen4_exp/nvidia/model_state.py | 143 ++ vllm/models/qwen4_exp/nvidia/mtp.py | 460 ++++++ vllm/models/qwen4_exp/nvidia/ops/__init__.py | 3 + vllm/models/qwen4_exp/nvidia/ops/hc.py | 487 ++++++ vllm/models/qwen4_exp/nvidia/ops/qsa.py | 1116 ++++++++++++++ .../qwen4_exp/nvidia/ops/qsa_pre_indexer.py | 516 +++++++ vllm/models/qwen4_exp/nvidia/ple_layer.py | 1310 +++++++++++++++++ vllm/models/qwen4_exp/nvidia/qsa.py | 518 +++++++ vllm/transformers_utils/config.py | 2 + vllm/transformers_utils/configs/__init__.py | 6 + vllm/transformers_utils/configs/qwen4_exp.py | 15 + vllm/v1/attention/backends/short_conv_attn.py | 529 ++++++- vllm/v1/core/kv_cache_utils.py | 371 ++++- vllm/v1/core/single_type_kv_cache_manager.py | 94 ++ vllm/v1/kv_cache_interface.py | 26 + vllm/v1/worker/gpu/attn_utils.py | 4 +- vllm/v1/worker/gpu/block_table.py | 16 +- vllm/v1/worker/gpu/model_runner.py | 27 +- .../worker/gpu/model_states/mamba_hybrid.py | 9 +- vllm/v1/worker/gpu/warmup.py | 14 +- vllm/v1/worker/gpu_model_runner.py | 162 +- 56 files changed, 10965 insertions(+), 108 deletions(-) create mode 100644 tests/model_executor/test_weights_mapper_stacked.py create mode 100644 tests/models/qwen4_exp/__init__.py create mode 100644 tests/models/qwen4_exp/test_config.py create mode 100644 tests/models/qwen4_exp/test_ple.py create mode 100644 tests/models/qwen4_exp/test_weight_loading.py create mode 100644 tests/v1/core/test_qwen4_exp_kv_cache.py create mode 100644 tests/v1/worker/test_qwen4_exp_ngram.py create mode 100644 tests/v1/worker/test_qwen4_exp_v2.py create mode 100644 vllm/models/qwen4_exp/__init__.py create mode 100644 vllm/models/qwen4_exp/common/__init__.py create mode 100644 vllm/models/qwen4_exp/common/hyperconnection.py create mode 100644 vllm/models/qwen4_exp/common/ple.py create mode 100644 vllm/models/qwen4_exp/common/qsa_cache.py create mode 100644 vllm/models/qwen4_exp/config.py create mode 100644 vllm/models/qwen4_exp/nvidia/__init__.py create mode 100644 vllm/models/qwen4_exp/nvidia/hyperconnection.py create mode 100644 vllm/models/qwen4_exp/nvidia/indexer_qsa.py create mode 100644 vllm/models/qwen4_exp/nvidia/low_latency_gemm.py create mode 100644 vllm/models/qwen4_exp/nvidia/model.py create mode 100644 vllm/models/qwen4_exp/nvidia/model_state.py create mode 100644 vllm/models/qwen4_exp/nvidia/mtp.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/__init__.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/hc.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/qsa.py create mode 100644 vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py create mode 100644 vllm/models/qwen4_exp/nvidia/ple_layer.py create mode 100644 vllm/models/qwen4_exp/nvidia/qsa.py create mode 100644 vllm/transformers_utils/configs/qwen4_exp.py diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index b76c30c842..fcf159bddc 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,12 +2,16 @@ ## Status and ownership -- Status: bring-up in progress; no route, quality, memory, or speed claim yet. +- Status: source bring-up implemented; CPU/configuration gates pass. Full-model + load, output quality, measured memory, and speed are not claimed yet. - Integration line: `private/main`. - Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. - Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision `7b719225242aacd3dbd3f9407468c2ee9a9d2594`. - Model download: `/data/models/RadixArk/Qwen3.8-Flash-Next-NVFP4`. +- Download source: ModelScope `master`, verified against the fixed Hugging Face + revision above: all 419 file sizes match and all 208 comparable LFS SHA-256 + values match. - Upstream references: [vLLM PR 53896](https://github.com/vllm-project/vllm/pull/53896) and [SGLang PR 36497](https://github.com/sgl-project/sglang/pull/36497). @@ -22,6 +26,10 @@ The first correctness route deliberately excludes speculative decoding. - Hardware: four NVIDIA V100-SXM2-32GB GPUs (SM70). - Parallelism: TP4, PP1, no expert parallelism. +- Model mode: `--language-model-only`. The initial SM70 route deliberately + excludes the vision tower from its memory, quality, and performance gates; + omitting the flag fails during configuration instead of reaching a private + Qwen3.5 multimodal API mismatch at model construction. - Compute dtype: FP16; no BF16 or native FP8/NVFP4 tensor-core assumptions. - Checkpoint: ModelOpt NVFP4 routed-expert weights, consumed as an SM70 weight-only W4A16 route. Ignored dense, attention, GDN, shared-expert, GR, @@ -33,6 +41,12 @@ The first correctness route deliberately excludes speculative decoding. - Initial KV cache: FP16. FP8 KV cache is a separate, quality-gated follow-up. - Initial decoding: MTP disabled. MTP may be enabled only after the no-MTP route is correct and its emitted-token baseline is recorded. +- Model runner: V2 is the default initial route. Its Qwen4Exp model state keeps + raw token IDs and builds the PLE context from committed tokens, so rejected + speculative candidates cannot leak into the next trigram. V1 remains a + correctness control, not the primary performance route. +- Prefix caching: disabled for the first route. The fixed QSA ring manager does + not yet implement reusable prefix blocks. ## Architecture facts that affect the port @@ -60,11 +74,41 @@ pinned memory, per-rank device peak, post-load device residency, and whether a loader creates duplicate staging buffers. A 262144-token context is admitted only after the measured peak leaves a safe margin on every 32GB GPU. +The SM70 TurboMind repack changes routed-expert FP4 scales from FP8 to FP16. +For TP4, routed experts are estimated at about 15.82 GiB/rank in the source +checkpoint and 17.57 GiB/rank after repack. This puts the idealized final +device weights near 21.3 GiB/rank. Because layers are repacked sequentially, +the estimated transient weight peak is about 22.1 GiB/rank before runtime +buffers, KV/index caches, NCCL, and CUDA graphs. These are storage calculations, +not `torch.cuda.max_memory_allocated` measurements. + +The loader marks the PLE parameter as permanently host-resident so generic +quantization post-processing cannot stage the entire 11.921 GiB TP shard on a +GPU. Only its small scale parameter resides on device; lookup reads selected +FP8 rows through a stable UVA view and converts the gathered output to FP16. + +With the real SM70 platform alignment and the QSA attention backend selected, +the V2 scheduler block is 784 tokens, the recurrent-state block is 32768 +tokens at a 32768-token initial maximum length, and each padded recurrent page +is 802816 bytes. The exact synthetic model layout has one 24-layer uniform QSA +main/compressed group, one 12-layer fixed circular QSA ring group, three +12-layer GDN state groups, and one PLE short-convolution state group. It +allocates 24 physical cache tensors. The aligned pool cost is 10235904 bytes +(9.762 MiB) per shared block per TP rank. + +For one request, the resulting cache-pool planning estimates are about 0.448 +GiB/rank at 32K (47 shared blocks), 1.649 GiB/rank at 128K (173 blocks), and +3.241 GiB/rank at 262144 tokens (340 blocks). Combining the last figure with +the estimated 21.3 GiB final weights gives about 24.54 GiB/rank before CUDA +graphs, workspaces, NCCL, allocator fragmentation, and loader transients. This +explains why TP4 is plausible, but it is not evidence that the maximum context +will load safely. + ## Acceptance gates -1. Static route: Transformers config, model registry, multimodal processor, +1. Static route: Transformers config, model registry/processor registration, QSA/GDN/GR/PLE modules, and ModelOpt NVFP4 mapping load without importing an - Ampere-only backend. + Ampere-only backend. Multimodal execution is outside this first route. 2. Loader route: TP4 expert shards select TurboMind SM70 W4A16; PLE shards are born on pinned CPU memory and do not consume persistent device memory. 3. Numerical route: focused operator comparisons against FP32/FP16 references, @@ -88,3 +132,50 @@ The first SM70-specific changes are limited to genericizing the existing TurboMind NVFP4 MoE shape contract, adding the QSA/indexer route, and adding a pinned-host PLE loader/gather path. Optimize GDN, GR, sparse attention, and MTP only after profiles identify them as measured decode bottlenecks. + +## Source validation snapshot + +- The real downloaded `config.json` resolves without remote model code as + `Qwen4ExpConfig` / `Qwen4ExpTextConfig`: 48 layers, 36 GDN, 12 QSA, 512 + experts, top-10, HC count four/rank 320, and one trigram PLE layer. +- Exact-SM70 configuration construction with FP16, TP4, prefix caching off, + language-model-only mode, and V2 selects `ModelOptNvFp4Config`, the + pinned-host PLE default, and the Qwen4Exp PLE/QSA compilation split + operators. The same real configuration rejects the unvalidated multimodal + route with an actionable `--language-model-only` error. +- Full 48-layer meta construction from the real checkpoint config succeeds in + language-model-only mode. It instantiates QSA, GDN, HC, PLE, and all 512 + experts without materializing weights; the routed experts select + `ModelOptNvFp4SM70MoEMethod(use_a16=True)` and the PLE table has shape + `(320001536, 160)` with FP8 E4M3 storage. This constructor probe used TP1; + TP4 selection and expert geometry are covered separately and full TP4 load + remains pending. +- Focused CPU tests cover PLE shard loading and hashing, `seed=None`, permanent + host residency during post-load processing, QSA cache grouping, V1 and V2 + n-gram inputs, V2 circular block-table sizing, scheduler-manager conversion, + official checkpoint weight mappings, and Qwen3.6/Qwen3.8 NVFP4 route + selection. The current CPU-only focused run is 76 passed and 7 CUDA skips; + all 55 changed Python files pass Ruff, format, and compileall checks. +- In the pre-final real V100-SXM2-32GB snapshot, 63 focused tests pass. They + include the Triton V2 slot-mapping kernel with its QSA circular group + disabled, pinned-host FP8 lookup through a CUDA UVA view, the compressed QSA + storage-page reshape, V2 committed-token PLE state, and the SM70 ModelOpt + NVFP4 selection gates. A final V100 rerun is still required after the current + GPU owners release a device. +- The upstream QSA fused pre-indexer executes on SM70 for both ordinary RoPE + and MRoPE inputs and matches a PyTorch normalization reference. This also + exposed and fixed two private-tree API differences: QKV projection is local + because the branch's `Qwen3NextAttention` has no `_project_qkv_gate`, and its + `triton_mrope` accepts eight rather than nine arguments. +- Actual SM70 platform alignment produces a 784-token attention block and an + 802816-byte padded recurrent page; the exact synthetic 48-layer cache layout + validates successfully after that alignment. +- The HC grouped norm/gate/combine kernels pass FP16 reference checks on a real + V100. A captured pinned-host FP8 embedding probe (228.9 MiB synthetic table, + 16 rows by 160 elements per replay) measured 95.81 microseconds/replay, + including the input-ID copy. This only demonstrates that the isolated UVA + lookup can be captured and is not by itself an end-to-end throughput result + or a measurement of the full 11.921 GiB TP shard. +- The existing general KV-cache utility/manager suites pass 69 tests; one + unrelated DeepSeek-v4 fixture failure is unchanged from the integration + base because its `SimpleNamespace` omits `max_in_flight_tokens`. diff --git a/tests/model_executor/test_weights_mapper_stacked.py b/tests/model_executor/test_weights_mapper_stacked.py new file mode 100644 index 0000000000..07f45a55dd --- /dev/null +++ b/tests/model_executor/test_weights_mapper_stacked.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.model_executor.layers.linear import MergedColumnParallelLinear +from vllm.model_executor.models.utils import WeightsMapper + + +def test_weights_mapper_preserves_stacked_shard_id() -> None: + mapper = WeightsMapper( + orig_to_new_stacked={ + ".input_mix_weight_down.weight": ( + ".input_mix_weight_down_block_inject.weight", + 0, + ), + ".block_inject_weight.weight": ( + ".input_mix_weight_down_block_inject.weight", + 1, + ), + } + ) + down = torch.ones(2, 4) + injection = torch.full((3, 4), 2.0) + + mapped = list( + mapper.apply( + [ + ("layer.input_mix_weight_down.weight", down), + ("layer.block_inject_weight.weight", injection), + ] + ) + ) + + assert [name for name, _ in mapped] == [ + "layer.input_mix_weight_down_block_inject.weight", + "layer.input_mix_weight_down_block_inject.weight", + ] + assert down.shard_id == 0 + assert injection.shard_id == 1 + + +def test_merged_column_load_weights_forwards_stacked_shards() -> None: + layer = object.__new__(MergedColumnParallelLinear) + torch.nn.Module.__init__(layer) + layer.output_sizes = [2, 3] + layer.tp_size = 1 + layer.tp_rank = 0 + layer.prefix = "hc.input_mix_weight_down_block_inject" + weight = torch.nn.Parameter(torch.zeros(5, 4)) + calls = [] + + def weight_loader(param, loaded_weight, shard_id) -> None: + calls.append((param, loaded_weight, shard_id)) + + weight.weight_loader = weight_loader + layer.register_parameter("weight", weight) + down = torch.ones(2, 4) + down.shard_id = 0 + injection = torch.full((3, 4), 2.0) + injection.shard_id = 1 + + loaded = list( + layer.load_weights( + [ + ("weight", down), + ("weight", injection), + ] + ) + ) + + assert loaded == ["weight", "weight"] + assert calls == [ + (weight, down, 0), + (weight, injection, 1), + ] diff --git a/tests/models/qwen4_exp/__init__.py b/tests/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/models/qwen4_exp/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py new file mode 100644 index 0000000000..f7abe35682 --- /dev/null +++ b/tests/models/qwen4_exp/test_config.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.models.config import ( + Qwen3_5ForConditionalGenerationConfig, + Qwen4ExpForConditionalGenerationConfig, +) +from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionMetadataBuilder, +) +from vllm.v1.worker.gpu.model_states.mamba_hybrid import ( + MambaHybridAttnMetadata, + MambaHybridModelState, +) + + +def test_initial_sm70_route_accepts_v2_runner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + text_config = SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + rope_parameters={"mrope_section": [11, 11, 10]}, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(rope_parameters={"mrope_interleaved": True}), + hf_text_config=text_config, + multimodal_config=SimpleNamespace(language_model_only=True), + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=None, + use_v2_model_runner=True, + ) + + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + assert "mrope_section" not in text_config.rope_parameters + assert "mrope_interleaved" not in vllm_config.model_config.hf_config.rope_parameters + + +def test_initial_sm70_route_rejects_multimodal_tower( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + text_config = SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + rope_parameters={"mrope_section": [11, 11, 10]}, + ) + multimodal_config = SimpleNamespace(language_model_only=False) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(rope_parameters={}), + hf_text_config=text_config, + multimodal_config=multimodal_config, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=None, + ) + + with pytest.raises(NotImplementedError, match="--language-model-only"): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + +@pytest.mark.parametrize( + "architecture", + ["Qwen4ExpForCausalLM", "Qwen4ExpForConditionalGeneration"], +) +def test_qwen4_exp_defaults_to_v2_even_when_quantized_moe( + architecture: str, +) -> None: + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + runner_type="generate", + architectures=[architecture], + is_moe=True, + is_quantized=True, + ) + ) + + assert VllmConfig._is_default_v2_model_runner_model(vllm_config) + + +def test_initial_sm70_v2_route_rejects_speculative_decode( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + ), + multimodal_config=None, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=SimpleNamespace(method="mtp"), + ) + + with pytest.raises(NotImplementedError, match="initial SM70 V2 route"): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + +def test_qwen4_exp_registers_v2_model_state() -> None: + from vllm.models.qwen4_exp.nvidia.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + + assert Qwen4ExpForCausalLM.get_model_state_cls() is Qwen4ExpModelState + assert Qwen4ExpForConditionalGeneration.get_model_state_cls() is Qwen4ExpModelState + + +def test_qwen4_exp_v2_model_state_uses_committed_ngram_context() -> None: + model_state = object.__new__(Qwen4ExpModelState) + model_state.uses_ngram_embedding = True + model_state.ngram_context_len = 3 + model_state.ngram_eos_token_id = 99 + model_state.ngram_context = torch.empty((4, 3), dtype=torch.int32) + model_state.ngram_context_offsets = torch.arange(-3, 0, dtype=torch.int64) + model_state.ple_query_start_loc = torch.empty(5, dtype=torch.int32) + + input_batch = SimpleNamespace( + num_reqs=2, + num_reqs_after_padding=3, + idx_mapping=torch.tensor([1, 0]), + query_start_loc=torch.tensor([0, 2, 3, 3], dtype=torch.int32), + ) + req_states = SimpleNamespace( + num_computed_tokens=SimpleNamespace(gpu=torch.tensor([3, 1])), + all_token_ids=SimpleNamespace( + gpu=torch.tensor([[1, 2, 3, 4], [20, 21, 22, 23]], dtype=torch.int32) + ), + ) + + with patch.object(MambaHybridModelState, "prepare_inputs", return_value={}): + model_inputs = model_state.prepare_inputs(input_batch, req_states) + + torch.testing.assert_close( + model_inputs["query_start_loc"], + torch.tensor([0, 2, 3, 3], dtype=torch.int32), + ) + torch.testing.assert_close( + model_inputs["ngram_context"], + torch.tensor([[99, 99, 20], [1, 2, 3], [99, 99, 99]], dtype=torch.int32), + ) + + +def test_qwen4_exp_ple_builder_receives_v2_decode_metadata() -> None: + num_accepted_tokens = torch.tensor([1, 2], dtype=torch.int32) + num_decode_draft_tokens_cpu = torch.tensor([-1, 2], dtype=torch.int32) + metadata = MambaHybridAttnMetadata( + is_prefilling=torch.tensor([False, False]), + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + builder = PleShortConvAttentionMetadataBuilder.__new__( + PleShortConvAttentionMetadataBuilder + ) + + kwargs = metadata.get_extra_attn_kwargs(builder, num_reqs=2) + + torch.testing.assert_close(kwargs["num_accepted_tokens"], num_accepted_tokens) + torch.testing.assert_close( + kwargs["num_decode_draft_tokens_cpu"], num_decode_draft_tokens_cpu + ) + + +def test_qwen4_exp_qsa_owns_qkv_projection_for_private_qwen3_api() -> None: + from vllm.models.qwen4_exp.nvidia.qsa import Qwen4ExpQSAAttention + + class AddRotaryOffset(torch.nn.Module): + def forward(self, positions, query, key): + del positions + return query + 1, key + 2 + + layer = object.__new__(Qwen4ExpQSAAttention) + torch.nn.Module.__init__(layer) + layer.q_size = 4 + layer.kv_size = 2 + layer.num_heads = 2 + layer.num_kv_heads = 1 + layer.head_dim = 2 + layer.q_norm = torch.nn.Identity() + layer.k_norm = torch.nn.Identity() + layer.rotary_emb = AddRotaryOffset() + qkv = torch.arange(12, dtype=torch.float32).reshape(1, 12) + + query, key, value, gate = layer._project_qkv_gate( + qkv, torch.tensor([0], dtype=torch.int64) + ) + + torch.testing.assert_close(query, torch.tensor([[1, 2, 5, 6.0]])) + torch.testing.assert_close(gate, torch.tensor([[2, 3, 6, 7.0]])) + torch.testing.assert_close(key, torch.tensor([[10, 11.0]])) + torch.testing.assert_close(value, torch.tensor([[10, 11.0]])) + + +def test_qwen4_exp_qsa_uses_private_mrope_signature( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.models.qwen4_exp.nvidia import indexer_qsa + + calls = [] + + def private_triton_mrope( + query, + key, + cos, + sin, + mrope_section, + head_size, + rotary_dim, + mrope_interleaved, + ): + calls.append( + ( + cos.shape, + sin.shape, + mrope_section, + head_size, + rotary_dim, + mrope_interleaved, + ) + ) + return query, key + + monkeypatch.setattr(indexer_qsa, "triton_mrope", private_triton_mrope) + cache = torch.zeros(16, 4) + rotary_emb = SimpleNamespace( + rotary_dim=2, + mrope_section=[1, 0, 0], + mrope_interleaved=True, + _match_cos_sin_cache_dtype=lambda _tensor: cache, + ) + tensor = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4) + positions = torch.tensor([[0, 1], [2, 3], [4, 5]]) + + output = indexer_qsa.apply_qsa_rope(rotary_emb, positions, tensor) + + torch.testing.assert_close(output, tensor) + assert calls == [ + (torch.Size([3, 2, 2]), torch.Size([3, 2, 2]), [1, 0, 0], 4, 2, True) + ] diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py new file mode 100644 index 0000000000..fcb081596e --- /dev/null +++ b/tests/models/qwen4_exp/test_ple.py @@ -0,0 +1,478 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.nn import functional as F + +import vllm.model_executor.layers.vocab_parallel_embedding as embedding_module +import vllm.model_executor.parameter as parameter_module +import vllm.models.qwen4_exp.nvidia.ple_layer as ple_module +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.model_loader.utils import device_loading_context +from vllm.models.qwen4_exp.common.ple import ( + PLEShardOverlap, + compute_ple_shard_overlap, + copy_ple_embedding_shard_, +) +from vllm.models.qwen4_exp.nvidia.ple_layer import ( + Qwen4ExpNGramEmbedding, + Qwen4ExpPinnedHostEmbedding, + Qwen4ExpPLEFp8EmbeddingMethod, + Qwen4ExpPLELayer, + _get_ple_embedding_quant_method, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA pinned memory") +def test_pinned_host_ple_allocates_tp_shard_without_device_table( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(ple_module, "is_pin_memory_available", lambda: True) + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 2) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 2) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + + layer = Qwen4ExpPinnedHostEmbedding( + num_embeddings=32, + embedding_dim=8, + params_dtype=torch.float16, + padding_size=8, + prefix="model.layers.2.ple.ngram_embedding", + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), + ) + + assert layer.tp_size == 4 + assert layer.weight.shape == (8, 8) + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight.device.type == "cpu" + assert layer.weight.is_pinned() + assert layer.weight._vllm_keep_on_cpu + assert not layer.weight.is_meta + assert not layer.weight_scale.is_meta + assert layer._accelerator_weight_views == {} + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires an exact SM70 CUDA device", +) +def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(ple_module, "is_pin_memory_available", lambda: True) + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + layer = Qwen4ExpPinnedHostEmbedding( + num_embeddings=8, + embedding_dim=8, + params_dtype=torch.float16, + padding_size=8, + prefix="model.layers.2.ple.ngram_embedding", + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), + ) + expected = torch.arange(64, dtype=torch.float32).reshape(8, 8) % 16 + layer.weight.data.copy_(expected.to(torch.float8_e4m3fn)) + + accelerator_weight = layer.get_accelerator_weight(torch.device("cuda")) + output = F.embedding( + torch.tensor([0, 7], dtype=torch.int64, device="cuda"), accelerator_weight + ) + torch.cuda.synchronize() + + assert output.dtype == torch.float8_e4m3fn + torch.testing.assert_close(output.float().cpu(), expected[[0, 7]]) + + +def test_post_load_context_keeps_marked_parameter_on_cpu() -> None: + module = nn.Module() + host_weight = nn.Parameter(torch.ones(2)) + host_weight._vllm_keep_on_cpu = True + module.register_parameter("weight", host_weight) + + with device_loading_context(module, torch.device("meta")): + assert module.weight.device.type == "cpu" + + assert module.weight.device.type == "cpu" + + +def test_ngram_embedding_accepts_checkpoint_seed_none( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 4 + ) + config = SimpleNamespace( + ngram_size=3, + heads_per_ngram=8, + eos_token_id=2, + vocab_size=64, + split_ngram_parts=2, + seed=None, + ngram_vocab_size_base=101, + make_ngram_vocab_size_divisible_by=128, + ple_embedding_dtype="float8_e4m3fn", + ple_offload_embedding=False, + ) + + with torch.device("meta"): + layer = Qwen4ExpNGramEmbedding( + config, + embedding_dim=256, + ple_dense_layer_id=0, + max_total_tokens=8, + max_num_reqs=2, + prefix="model.layers.2.ple.ple_embedding", + params_dtype=torch.float16, + ) + + assert layer.ngram_heads == 16 + assert layer.head_dim == 16 + assert layer.ngram_embedding.weight.dtype == torch.float8_e4m3fn + assert layer.ngram_embedding.weight.is_meta + + +def _make_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: + module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding) + nn.Module.__init__(module) + module.split_ngram_parts = 2 + module.register_buffer("layer_multipliers", torch.zeros(1, dtype=torch.long)) + module.register_buffer("ngram_heads_offsets", torch.zeros(1, dtype=torch.long)) + module.register_buffer("ngram_heads_vocab_sizes", torch.zeros(1, dtype=torch.long)) + module.ngram_embedding = SimpleNamespace( + org_vocab_size=8, + embedding_dim=2, + weight=nn.Parameter(torch.full((4, 2), -1.0)), + shard_indices=SimpleNamespace( + org_vocab_start_index=2, + org_vocab_end_index=6, + ), + ) + return module + + +def _make_fp8_ngram_embedding_for_load_test() -> Qwen4ExpNGramEmbedding: + module = _make_ngram_embedding_for_load_test() + embedding = nn.Module() + embedding.org_vocab_size = 8 + embedding.embedding_dim = 2 + embedding.shard_indices = SimpleNamespace( + org_vocab_start_index=2, + org_vocab_end_index=6, + ) + embedding.register_parameter( + "weight", + nn.Parameter( + torch.full((4, 2), -1.0).to(torch.float8_e4m3fn), + requires_grad=False, + ), + ) + embedding.register_parameter( + "weight_scale", + nn.Parameter(torch.zeros(1, dtype=torch.bfloat16), requires_grad=False), + ) + module.ngram_embedding = embedding + return module + + +def test_ple_shard_overlap_and_copy() -> None: + overlap = compute_ple_shard_overlap( + checkpoint_start=2, checkpoint_rows=5, tp_start=4, tp_end=8 + ) + assert overlap == PLEShardOverlap(source_start=2, destination_start=0, row_count=3) + + destination = torch.full((4, 2), -1.0) + loaded = torch.arange(10, dtype=torch.float64).reshape(5, 2) + copied = copy_ple_embedding_shard_( + destination, + loaded, + checkpoint_start=2, + tp_start=4, + tp_end=8, + ) + + assert copied == 3 + torch.testing.assert_close(destination[:3], loaded[2:5].float()) + torch.testing.assert_close(destination[3], torch.tensor([-1.0, -1.0])) + + +def test_ple_shard_copy_is_a_noop_without_overlap() -> None: + destination = torch.ones(4, 2) + copied = copy_ple_embedding_shard_( + destination, + torch.zeros(2, 2), + checkpoint_start=10, + tp_start=4, + tp_end=8, + ) + + assert copied == 0 + assert torch.equal(destination, torch.ones_like(destination)) + + +def test_ngram_embedding_loads_shards_and_ignores_legacy_token_lookup() -> None: + module = _make_ngram_embedding_for_load_test() + shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2) + shard_1 = torch.arange(8, 16, dtype=torch.float32).reshape(4, 2) + + loaded = module.load_weights( + [ + ("ngram_embedding.shard_0.weight", shard_0), + ("ngram_embedding.shard_1.weight", shard_1), + ("token_lookup", torch.tensor([2, 1, 0])), + ] + ) + + assert loaded == {"ngram_embedding.weight"} + torch.testing.assert_close( + module.ngram_embedding.weight, + torch.cat((shard_0[2:4], shard_1[0:2])), + ) + + +def test_ngram_embedding_rejects_mismatched_checkpoint_shard() -> None: + module = _make_ngram_embedding_for_load_test() + + with pytest.raises( + ValueError, + match=r"Shape mismatch for PLE embedding shard 0", + ): + module.load_weights([("ngram_embedding.shard_0.weight", torch.zeros(3, 2))]) + + +def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None: + module = _make_fp8_ngram_embedding_for_load_test() + shard_0 = torch.arange(8, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn) + shard_1 = ( + torch.arange(8, 16, dtype=torch.float32).reshape(4, 2).to(torch.float8_e4m3fn) + ) + weight_scale = torch.tensor([0.25], dtype=torch.bfloat16) + + loaded = module.load_weights( + [ + ("ngram_embedding.shard_0.weight", shard_0), + ("ngram_embedding.shard_1.weight", shard_1), + ("ngram_embedding.weight_scale", weight_scale), + ] + ) + + assert loaded == {"ngram_embedding.weight", "ngram_embedding.weight_scale"} + assert module.ngram_embedding.weight.dtype == torch.float8_e4m3fn + assert torch.equal( + module.ngram_embedding.weight.float(), + torch.cat((shard_0[2:4], shard_1[0:2])).float(), + ) + assert torch.equal(module.ngram_embedding.weight_scale, weight_scale) + + +def _make_fp8_embedding_layer( + monkeypatch: pytest.MonkeyPatch, +) -> embedding_module.VocabParallelEmbedding: + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + embedding_module, + "tensor_model_parallel_all_reduce", + lambda tensor: tensor, + ) + method = Qwen4ExpPLEFp8EmbeddingMethod() + layer = embedding_module.VocabParallelEmbedding( + 3, + 2, + params_dtype=torch.bfloat16, + padding_size=1, + quant_method=method, + ) + weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]]) + layer.weight.data.copy_(weight.to(torch.float8_e4m3fn)) + layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=torch.bfloat16)) + return layer + + +def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None: + layer = _make_fp8_embedding_layer(monkeypatch) + quantized_output = layer(torch.tensor([2, 0])) + ple_layer = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + nn.Module.__init__(ple_layer) + ple_layer.ple_embedding = nn.Module() + ple_layer.ple_embedding.ngram_embedding = layer + + output = ple_layer._dequantize_embeddings( + quantized_output, + torch.bfloat16, + ) + + assert layer.weight.dtype == torch.float8_e4m3fn + assert layer.weight_scale.dtype == torch.bfloat16 + assert quantized_output.dtype == torch.float8_e4m3fn + assert output.dtype == torch.bfloat16 + weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]]) + torch.testing.assert_close(output, (weight[[2, 0]] * 0.25).bfloat16()) + + +def test_ple_fp8_embedding_uses_int8_for_tp_reduce(monkeypatch) -> None: + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + embedding_module, "get_tensor_model_parallel_world_size", lambda: 2 + ) + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + parameter_module, "get_tensor_model_parallel_world_size", lambda: 2 + ) + monkeypatch.setattr( + embedding_module, + "get_masked_input_and_mask", + lambda *args: ( + torch.tensor([0, 0]), + torch.tensor([False, True]), + ), + ) + reduced_dtypes = [] + + def all_reduce(tensor: torch.Tensor) -> torch.Tensor: + reduced_dtypes.append(tensor.dtype) + return tensor.clone() + + monkeypatch.setattr( + embedding_module, + "tensor_model_parallel_all_reduce", + all_reduce, + ) + layer = embedding_module.VocabParallelEmbedding( + 4, + 2, + params_dtype=torch.bfloat16, + padding_size=1, + quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), + ) + layer.weight.data.copy_( + torch.tensor([[1.0, 2.0], [4.0, 8.0]]).to(torch.float8_e4m3fn) + ) + + output = layer(torch.tensor([0, 2])) + + assert reduced_dtypes == [torch.int8] + assert output.dtype == torch.float8_e4m3fn + torch.testing.assert_close(output[0].float(), layer.weight[0].float()) + assert torch.count_nonzero(output[1].float()) == 0 + + +def test_ple_fp8_embedding_respects_checkpoint_shard_exclusions() -> None: + prefix = "model.layers.1.ple.ple_embedding.ngram_embedding" + quant_config = Fp8Config( + is_checkpoint_fp8_serialized=True, + ignored_layers=[], + weight_block_size=[128, 128], + ) + assert isinstance( + _get_ple_embedding_quant_method(quant_config, prefix), + Qwen4ExpPLEFp8EmbeddingMethod, + ) + + quant_config.ignored_layers = [f"{prefix}.shard_0"] + assert _get_ple_embedding_quant_method(quant_config, prefix) is None + + +def test_dilated_ple_spec_state_rolls_back_before_next_forward() -> None: + module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + nn.Module.__init__(module) + module.conv_state_len = 6 + module.short_conv_dilation = 2 + + conv_weights = torch.tensor([[0.25, -0.5, 0.75, 1.0]]) + conv_state = torch.zeros(2, 1, 9) + conv_state[1] = torch.arange(1, 10, dtype=torch.float32).reshape(1, 9) + first_inputs = torch.tensor([[10.0], [20.0], [30.0], [40.0]]) + initial_state = conv_state[1:].clone() + first_history = torch.cat( + (initial_state[..., : module.conv_state_len], first_inputs.T.unsqueeze(0)), + dim=-1, + ) + + graph_padded_inputs = F.pad(first_inputs, (0, 0, 0, 4)) + first_output = module._short_conv_dilated_spec_batched( + graph_padded_inputs, + conv_state, + conv_weights, + torch.tensor([1, 0]), + torch.tensor([0, 4, 4]), + torch.tensor([1, 0]), + spec_query_len=4, + ) + + expected_first_output = F.silu( + F.conv1d( + first_history, + conv_weights.unsqueeze(1), + groups=1, + dilation=module.short_conv_dilation, + ) + ).transpose(1, 2)[0] + expected_first_state = first_history[..., 1:10] + torch.testing.assert_close(first_output[:4], expected_first_output) + assert torch.count_nonzero(first_output[4:]) == 0 + assert torch.count_nonzero(conv_state[0]) == 0 + torch.testing.assert_close(conv_state[1:], expected_first_state) + + second_inputs = torch.tensor([[50.0], [60.0]]) + rollback_state = expected_first_state[..., 1:7] + padded_second_inputs = F.pad(second_inputs.T.unsqueeze(0), (0, 2)) + second_history = torch.cat((rollback_state, padded_second_inputs), dim=-1) + expected_second_state = expected_first_state.clone() + expected_second_state[..., :7] = second_history[..., 1:8] + + second_output = module._short_conv_dilated_spec_batched( + second_inputs, + conv_state, + conv_weights, + torch.tensor([1]), + torch.tensor([0, 2]), + torch.tensor([2]), + spec_query_len=4, + ) + + expected_second_output = F.silu( + F.conv1d( + second_history, + conv_weights.unsqueeze(1), + groups=1, + dilation=module.short_conv_dilation, + ) + ).transpose(1, 2)[0, :2] + torch.testing.assert_close(second_output, expected_second_output) + torch.testing.assert_close(conv_state[1:], expected_second_state) + + +def test_ple_state_shape_reserves_speculative_tokens() -> None: + module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer) + nn.Module.__init__(module) + module.hc_hidden_size = 32 + module.conv_state_len = 9 + module.num_spec_tokens = 3 + + assert module.get_state_shape()[0] in ((32, 12), (12, 32)) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py new file mode 100644 index 0000000000..81df1c5a9c --- /dev/null +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.models.qwen4_exp.nvidia.model import ( + Qwen4ExpForConditionalGeneration, + Qwen4ExpModel, + _remap_qsa_cache_scale_name, +) + + +@pytest.mark.parametrize( + ("checkpoint_name", "model_name", "shard_id"), + [ + ( + "layers.0.self_attn.q_proj.weight", + "layers.0.self_attn.qkv_proj.weight", + "q", + ), + ( + "layers.0.self_attn.k_proj.weight", + "layers.0.self_attn.qkv_proj.weight", + "k", + ), + ( + "layers.1.linear_attn.in_proj_qkv.weight", + "layers.1.linear_attn.in_proj_qkvz.weight", + (0, 1, 2), + ), + ( + "layers.1.linear_attn.in_proj_z.weight", + "layers.1.linear_attn.in_proj_qkvz.weight", + 3, + ), + ( + "layers.1.linear_attn.in_proj_b.weight", + "layers.1.linear_attn.in_proj_ba.weight", + 0, + ), + ( + "layers.1.mlp.gate_proj.weight", + "layers.1.mlp.gate_up_proj.weight", + 0, + ), + ( + "layers.1.mlp.experts.0.gate_proj.weight", + "layers.1.mlp.experts.0.gate_proj.weight", + None, + ), + ( + "layers.0.self_attn.indexer.index_qk_proj.weight", + "layers.0.self_attn.indexer.index_qk_proj.weight", + None, + ), + ( + "layers.0.attn_hyper_connection.input_mix_weight_down.weight", + "layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + 0, + ), + ( + "layers.0.attn_hyper_connection.block_inject_weight.weight", + "layers.0.attn_hyper_connection.input_mix_weight_down_block_inject.weight", + 1, + ), + ( + "hyper_connection_mixer.input_mix_weight_down.weight", + "hyper_connection_mixer.input_mix_weight_down.weight", + None, + ), + ( + "layers.1.ple.ple_embedding.layer_multipliers", + "layers.1.ple.ple_embedding.layer_multipliers", + None, + ), + ], +) +def test_text_checkpoint_mapper_preserves_qwen4_exp_specific_weights( + checkpoint_name: str, + model_name: str, + shard_id: str | int | tuple[int, ...] | None, +) -> None: + assert Qwen4ExpModel.hf_to_vllm_mapper._map_name_with_shard(checkpoint_name) == ( + model_name, + shard_id, + ) + + +def test_outer_checkpoint_mapper_selects_language_model_only_paths() -> None: + mapper = Qwen4ExpForConditionalGeneration.hf_to_vllm_mapper + + assert ( + mapper._map_name("model.language_model.layers.0.ple.key_proj.weight") + == "language_model.model.layers.0.ple.key_proj.weight" + ) + assert mapper._map_name("lm_head.weight") == "language_model.lm_head.weight" + assert mapper._map_name("model.visual.blocks.0.attn.qkv.weight") == ( + "visual.blocks.0.attn.qkv.weight" + ) + + +@pytest.mark.parametrize( + ("checkpoint_name", "model_name"), + [ + ( + "layers.0.self_attn.k_proj.k_scale", + "layers.0.self_attn._k_scale", + ), + ( + "layers.0.self_attn.v_proj.output_scale", + "layers.0.self_attn._v_scale", + ), + ( + "language_model.model.layers.0.self_attn.attn.k_scale", + "language_model.model.layers.0.self_attn._k_scale", + ), + ( + "layers.0.self_attn.indexer.index_qk_proj.weight_scale", + "layers.0.self_attn.indexer.index_qk_proj.weight_scale", + ), + ( + "layers.1.self_attn.k_proj.k_scale", + "layers.1.self_attn.k_proj.k_scale", + ), + ], +) +def test_only_qsa_main_cache_scales_move_to_the_merged_owner( + checkpoint_name: str, + model_name: str, +) -> None: + assert _remap_qsa_cache_scale_name(checkpoint_name, frozenset({0})) == model_name diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 065b447cc9..35982e851f 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -57,6 +57,19 @@ def _moe_contract(**overrides): return SimpleNamespace(**values) +def _qwen4_moe_contract(**overrides): + values = { + "num_experts": 512, + "experts_per_token": 10, + "hidden_dim": 2560, + "intermediate_size_per_partition": 160, + "tp_size": 4, + "moe_parallel_config": SimpleNamespace(use_all2all_kernels=False), + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_mixed_min_capability_requires_exact_sm70_and_both_turbomind_routes(): with ( patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), @@ -92,12 +105,16 @@ def test_nvfp4_grouped_prefill_defaults_on_and_can_be_disabled(monkeypatch): ("tp_size", 2), ], ) -def test_nvfp4_moe_contract_rejects_non_qwen36_shapes(field, value): +def test_nvfp4_moe_contract_rejects_unvalidated_shapes(field, value): validate_nvfp4_sm70_moe_contract(_moe_contract()) with pytest.raises(NotImplementedError): validate_nvfp4_sm70_moe_contract(_moe_contract(**{field: value})) +def test_nvfp4_moe_contract_accepts_qwen4_exp_tp4(): + validate_nvfp4_sm70_moe_contract(_qwen4_moe_contract()) + + def test_nvfp4_moe_contract_rejects_shape_consistent_unvalidated_tp8(): with pytest.raises(NotImplementedError, match="tensor parallel"): validate_nvfp4_sm70_moe_contract( @@ -160,7 +177,7 @@ def test_nvfp4_sm70_moe_owns_routing_without_generic_modular_wrapper(): not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires an exact SM70 CUDA device", ) -@pytest.mark.parametrize("total_slots", (8, 72, 80)) +@pytest.mark.parametrize("total_slots", (8, 72, 80, 100)) def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slots): sorted_expert_ids = ( torch.arange(total_slots, dtype=torch.int32, device="cuda") // 3 @@ -190,3 +207,25 @@ class FakeRoutedExperts: pytest.raises(NotImplementedError, match="TurboMind"), ): config.get_quant_method(FakeRoutedExperts(), "model.layers.0.mlp.experts") + + +def test_pure_nvfp4_qwen4_moe_uses_turbomind_w4a16_on_sm70(): + config = ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + ) + + class FakeRoutedExperts: + moe_config = _qwen4_moe_contract() + + with ( + patch.object(modelopt, "RoutedExperts", FakeRoutedExperts), + patch.object(sm70_tm, "is_exact_sm70_cuda_platform", return_value=True), + patch.object(sm70_tm, "should_use_nvfp4_moe_turbomind", return_value=True), + ): + method = config.get_quant_method( + FakeRoutedExperts(), "model.layers.0.mlp.experts" + ) + + assert isinstance(method, ModelOptNvFp4SM70MoEMethod) + assert method.use_a16 diff --git a/tests/v1/core/test_qwen4_exp_kv_cache.py b/tests/v1/core/test_qwen4_exp_kv_cache.py new file mode 100644 index 0000000000..1763bb1dda --- /dev/null +++ b/tests/v1/core/test_qwen4_exp_kv_cache.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.models.qwen4_exp.common.qsa_cache import QSAStateBackend +from vllm.v1.core.block_pool import BlockPool +from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator +from vllm.v1.core.kv_cache_utils import ( + _get_csa_linear_tensor_layout, + generate_scheduler_kv_cache_config, + get_kv_cache_config_from_groups, + get_kv_cache_groups, +) +from vllm.v1.core.single_type_kv_cache_manager import CircularBufferManager +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + MambaSpec, + MLAAttentionSpec, +) +from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache +from vllm.v1.worker.utils import AttentionGroup + + +class _ModelConfig: + max_model_len = 8192 + + def get_num_kv_heads(self, parallel_config) -> int: + del parallel_config + return 1 + + def get_total_num_hidden_layers(self) -> int: + return 8 + + +def _vllm_config(): + return SimpleNamespace( + model_config=_ModelConfig(), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + cache_config=SimpleNamespace( + num_gpu_blocks_override=None, + mamba_cache_mode="none", + ), + ) + + +def _qwen4_exp_cache_specs(): + specs = {} + for layer in (3, 7): + prefix = f"model.layers.{layer}.self_attn" + specs[prefix] = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=256, + head_size_v=256, + dtype=torch.float16, + ) + specs[f"{prefix}.compressed"] = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + compress_ratio=4, + ) + specs[f"{prefix}.compressor_state"] = CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.float16, + ) + + for layer in (0, 1, 2, 4, 5, 6): + specs[f"model.layers.{layer}.linear_attn"] = MambaSpec( + block_size=16, + shapes=((1, 64),), + dtypes=(torch.float16,), + ) + specs["model.layers.2.ple"] = MambaSpec( + block_size=16, + shapes=((1, 64),), + dtypes=(torch.float16,), + tp_replicated=True, + ) + return specs + + +def test_qwen4_exp_csa_linear_cache_layout() -> None: + groups = get_kv_cache_groups(_vllm_config(), _qwen4_exp_cache_specs()) + layout = _get_csa_linear_tensor_layout(groups) + + assert layout is not None + assert [len(group.layer_names) for group in groups] == [4, 2, 2, 2, 2, 1] + assert len(layout.main_kv_names) == 2 + assert len(layout.compressed_names) == 2 + assert len(layout.compressor_state_names) == 2 + assert len(layout.mamba_groups) == 4 + + cache_config = get_kv_cache_config_from_groups( + _vllm_config(), groups, available_memory=1 << 30 + ) + assert len(cache_config.kv_cache_tensors) == 4 + assert all(len(tensor.shared_by) >= 2 for tensor in cache_config.kv_cache_tensors) + + scheduler_config = generate_scheduler_kv_cache_config([cache_config]) + scheduler_config.num_blocks = 32 + assert isinstance( + scheduler_config.kv_cache_groups[0].kv_cache_spec, FullAttentionSpec + ) + assert isinstance( + scheduler_config.kv_cache_groups[1].kv_cache_spec, CircularBufferSpec + ) + coordinator = get_kv_cache_coordinator( + scheduler_config, + max_model_len=8192, + max_in_flight_tokens=128, + use_eagle=False, + enable_caching=False, + enable_kv_cache_events=False, + dcp_world_size=1, + pcp_world_size=1, + hash_block_size=4, + ) + assert isinstance(coordinator.single_type_managers[1], CircularBufferManager) + + +def test_qwen4_exp_circular_cache_stores_keys_without_unused_values() -> None: + spec = CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.float16, + ) + + assert spec.real_page_size_bytes == 4 * 128 * 2 + assert spec.max_memory_usage_bytes(_vllm_config()) == spec.page_size_bytes + + +def test_qwen4_exp_circular_manager_owns_one_block_per_request() -> None: + spec = CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.float16, + ) + block_pool = BlockPool( + num_gpu_blocks=8, + enable_caching=False, + hash_block_size=spec.block_size, + ) + manager = CircularBufferManager( + spec, + block_pool=block_pool, + enable_caching=False, + kv_cache_group_id=0, + ) + + assert manager.get_num_blocks_to_allocate("req", 4096, (), 0, 4096) == 1 + blocks = manager.allocate_new_blocks("req", 4096, 4096) + assert len(blocks) == 1 + assert manager.req_to_blocks["req"] == blocks + assert manager.get_num_blocks_to_allocate("req", 8192, (), 4096, 8192) == 0 + assert manager.allocate_new_blocks("req", 8192, 8192) == [] + + +def test_qwen4_exp_compressed_qsa_reshape_uses_storage_block_size() -> None: + spec = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + compress_ratio=4, + ) + num_blocks = 3 + raw = torch.empty(num_blocks * spec.page_size_bytes, dtype=torch.int8) + group = AttentionGroup( + QSAStateBackend, + ["compressed"], + spec, + kv_cache_group_id=0, + ) + + caches = _reshape_kv_cache( + attn_groups=[group], + kv_cache_raw_tensors={"compressed": raw}, + cache_dtype="auto", + kernel_block_sizes=[16], + shared_kv_cache_layers={}, + ) + + assert caches["compressed"].shape == (num_blocks, 4, 1, 128) + assert caches["compressed"].untyped_storage().data_ptr() == raw.data_ptr() diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index b4fcb8bd0a..4beda08b46 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -7,7 +7,12 @@ from vllm.config.speculative import SpeculativeConfig from vllm.config.vllm import VllmConfig -from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + MambaSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.worker.gpu.warmup import _reserved_block_count BLOCK_SIZE = 16 @@ -73,10 +78,26 @@ def _mamba_spec(mode: str) -> MambaSpec: ) +def _circular_spec() -> CircularBufferSpec: + return CircularBufferSpec( + block_size=4, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + ) + + @pytest.mark.parametrize( ("spec", "expected"), [ (_full_attention_spec(), 2), + (_circular_spec(), 1), + ( + UniformTypeKVCacheSpecs( + block_size=4, kv_cache_specs={"compressor_state": _circular_spec()} + ), + 1, + ), (_mamba_spec("align"), 8), (_mamba_spec("none"), 9), (_mamba_spec("all"), 9), diff --git a/tests/v1/worker/test_qwen4_exp_ngram.py b/tests/v1/worker/test_qwen4_exp_ngram.py new file mode 100644 index 0000000000..c3904e1b65 --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_ngram.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import torch + +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + +class _Buffer: + def __init__(self, shape: tuple[int, ...]) -> None: + self.np = np.zeros(shape, dtype=np.int32) + self.gpu = torch.zeros(shape, dtype=torch.int32) + + def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: + if n is None: + self.gpu.copy_(torch.from_numpy(self.np)) + else: + self.gpu[:n].copy_(torch.from_numpy(self.np[:n])) + return self.gpu + + +def _runner() -> GPUModelRunner: + runner = object.__new__(GPUModelRunner) + runner.uses_ngram_embedding = True + runner.ngram_context_len = 2 + runner.ngram_eos_token_id = 99 + runner.enable_prompt_embeds = False + runner._sm70_async_staged_input_prep_active = False + runner.ngram_context = _Buffer((4, 2)) + runner.query_start_loc = _Buffer((5,)) + runner.input_batch = SimpleNamespace( + num_computed_tokens_cpu=np.array([3, 1, 0, 0], dtype=np.int32), + token_ids_cpu=np.array( + [ + [1, 2, 3, 4], + [20, 21, 22, 23], + [0, 0, 0, 0], + [0, 0, 0, 0], + ], + dtype=np.int32, + ), + is_token_ids=np.ones((4, 4), dtype=bool), + ) + return runner + + +def test_v1_runner_prepares_committed_ngram_context() -> None: + context = _runner()._prepare_ngram_context(num_reqs=2, num_reqs_padded=3) + + torch.testing.assert_close( + context, + torch.tensor([[2, 3], [99, 20], [99, 99]], dtype=torch.int32), + ) + + +def test_v1_runner_uses_stable_dummy_ngram_buffers() -> None: + runner = _runner() + model_kwargs = {} + + runner._maybe_add_ngram_kwargs( + model_kwargs, + num_reqs=2, + num_reqs_padded=3, + is_first_rank=True, + is_encoder_decoder=False, + use_dummy_context=True, + num_scheduled_tokens=[2, 1], + ) + + torch.testing.assert_close( + model_kwargs["query_start_loc"], + torch.tensor([0, 2, 3, 3], dtype=torch.int32), + ) + torch.testing.assert_close( + model_kwargs["ngram_context"], + torch.full((3, 2), 99, dtype=torch.int32), + ) diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py new file mode 100644 index 0000000000..13e9ee95bb --- /dev/null +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.worker.gpu import model_runner as mrv2 +from vllm.v1.worker.gpu.block_table import BlockTables + + +def test_qsa_circular_group_uses_one_block_and_custom_slot_mapping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + runner = mrv2.GPUModelRunner.__new__(mrv2.GPUModelRunner) + runner.max_model_len = 262144 + runner.is_encoder_decoder = False + runner.dcp_size = 1 + runner.dcp_rank = 0 + runner.cp_interleave = 1 + runner.cache_config = SimpleNamespace(enable_prefix_caching=False) + runner.vllm_config = SimpleNamespace() + runner.max_num_reqs = 1 + runner.max_num_tokens = 2 + runner.device = torch.device("cpu") + + circular_spec = CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + ) + attention_spec = FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + ) + compressed_spec = MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.float16, + ) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + layer_names=["compressor_state"], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=8, kv_cache_specs={"compressor_state": circular_spec} + ), + ), + KVCacheGroupSpec( + layer_names=["attention", "compressed"], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={ + "attention": attention_spec, + "compressed": compressed_spec, + }, + ), + ), + ], + ) + + monkeypatch.setattr( + mrv2, + "init_attn_backend", + lambda *args: ([], SimpleNamespace(), [8, 16]), + ) + captured: dict[str, object] = {} + + class BlockTablesCaptured(Exception): + pass + + def capture_block_tables(**kwargs): + captured.update(kwargs) + raise BlockTablesCaptured + + monkeypatch.setattr(mrv2, "BlockTables", capture_block_tables) + + with pytest.raises(BlockTablesCaptured): + runner.initialize_kv_cache(kv_cache_config) + + assert captured["max_num_blocks_per_group"] == [1, 16384] + assert captured["slot_mapping_enabled"] == [False, True] + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires an exact SM70 CUDA device", +) +def test_qsa_circular_group_emits_no_generic_slots_on_sm70() -> None: + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[8, 262144], + max_num_reqs=1, + max_num_batched_tokens=4, + max_num_blocks_per_group=[1, 1], + device=device, + kernel_block_sizes=[8, 262144], + slot_mapping_enabled=[False, True], + ) + block_tables.append_block_ids( + req_index=0, + new_block_ids=([7], [12]), + overwrite=True, + ) + block_tables.apply_staged_writes() + + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping=torch.tensor([0], dtype=torch.int32, device=device), + query_start_loc=torch.tensor([0, 2], dtype=torch.int32, device=device), + positions=torch.tensor([153797, 165757], dtype=torch.int64, device=device), + num_tokens_padded=2, + ) + torch.cuda.synchronize() + + assert slot_mappings[0].tolist() == [-1, -1] + assert slot_mappings[1].tolist() == [ + 12 * 262144 + 153797, + 12 * 262144 + 165757, + ] diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 2b0833e545..3a11face66 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -759,6 +759,8 @@ class CompilationConfig: "vllm::mamba_mixer2", "vllm::mamba_mixer", "vllm::short_conv", + "vllm::qwen4_exp_ple_short_conv", + "vllm::qwen4_exp_qsa_with_output", "vllm::linear_attention", "vllm::plamo2_mamba_mixer", "vllm::qwen_gdn_attention_core", diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 961b979c12..a24c2ac59b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -66,7 +66,13 @@ logger = init_logger(__name__) -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset({"Qwen3ForCausalLM"}) +DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( + { + "Qwen3ForCausalLM", + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + } +) _SM70_NOMTP_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 16) _SM70_MTP_CUDAGRAPH_REQUEST_SIZES = (1, 2, 4, 6, 8, 12, 16) @@ -693,11 +699,18 @@ def _is_default_v2_model_runner_model(self) -> bool: return False architectures = getattr(model_config, "architectures", []) - if not any( + is_default_v2_architecture = any( arch in DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES for arch in architectures - ): + ) + if not is_default_v2_architecture: return False + # Qwen4Exp is a hybrid MoE model whose QSA ring cache and PLE raw-token + # inputs are implemented by Model Runner V2. Its quantized checkpoint + # is therefore an explicit V2 route rather than a generic fallback. + if any(arch.startswith("Qwen4ExpFor") for arch in architectures): + return True + return not model_config.is_moe and not model_config.is_quantized @property diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 68f99ee3eb..47477033f3 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -3,6 +3,7 @@ import itertools from abc import abstractmethod +from collections.abc import Iterable import torch from torch.nn.parameter import Parameter, UninitializedParameter @@ -1215,6 +1216,29 @@ def weight_loader_v2( tp_rank=self.tp_rank, ) + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[str]: + for name, loaded_weight in weights: + shard_id = getattr(loaded_weight, "shard_id", None) + self.validate_shard_id(shard_id) + if "." in name: + submodule, _, attr = name.rpartition(".") + param = getattr(self.get_submodule(submodule), attr, self) + else: + param = getattr(self, name, self) + if param is None and name == "bias": + continue + param.weight_loader(param, loaded_weight, shard_id) + logger.debug( + "Loaded shard %s with shape %s into %s.%s", + shard_id, + loaded_weight.shape, + self.prefix, + name, + ) + yield name + class QKVParallelLinear(ColumnParallelLinear): """Linear layers for the attention's QKV transformation. diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index f8b88a5e8d..13aaf64b3a 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -37,6 +37,10 @@ def get_state_shape(self) -> Iterable[tuple[int, ...]]: def mamba_type(self) -> MambaAttentionBackendEnum: pass + @property + def is_kv_cache_tp_replicated(self) -> bool: + return False + @abstractmethod def get_state_dtype(self) -> tuple[torch.dtype, ...]: pass @@ -51,6 +55,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: block_size=mamba_block_size, page_size_padded=page_size_padded, mamba_type=self.mamba_type, + tp_replicated=self.is_kv_cache_tp_replicated, mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, num_speculative_blocks=( vllm_config.speculative_config.num_speculative_state_tokens() diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 935aa700f1..9344943de6 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -17,6 +17,7 @@ STR_DTYPE_TO_TORCH_DTYPE, get_kv_cache_torch_dtype, ) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum logger = init_logger(__name__) @@ -199,9 +200,10 @@ def short_conv_state_shape( tp_world_size: int, intermediate_size: int, conv_kernel: int, + num_spec: int = 0, ) -> tuple[tuple[int, int]]: conv_dim = divide(intermediate_size, tp_world_size) - conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1) + conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1 + num_spec) return (conv_state_shape,) @classmethod @@ -301,6 +303,10 @@ class MambaCopySpec: num_accepted_tokens: int - number of accepted tokens used to compute the copy offset. Range: 1 .. 1 + num_speculative_tokens (inclusive). """ +MambaStateCopyFuncs: TypeAlias = tuple[MambaStateCopyFunc, ...] +MambaStateCopyFuncsByType: TypeAlias = dict[ + MambaAttentionBackendEnum, MambaStateCopyFuncs +] def get_conv_copy_spec( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 09b3049800..35384dbdf0 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1161,6 +1161,28 @@ def get_name(self) -> QuantizationMethods: def get_supported_act_dtypes(self) -> list[torch.dtype]: return [torch.bfloat16, torch.half, torch.float8_e4m3fn] + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + if ( + isinstance(layer, RoutedExperts) + and not self.is_layer_excluded(prefix) + and sm70_tm.is_exact_sm70_cuda_platform() + ): + if not sm70_tm.should_use_nvfp4_moe_turbomind(): + raise NotImplementedError( + "ModelOpt NVFP4 MoE on SM70 requires the TurboMind backend." + ) + from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + ) + + return ModelOptNvFp4SM70MoEMethod( + quant_config=self, + moe_config=layer.moe_config, + ) + return super().get_quant_method(layer, prefix) + @classmethod def get_min_capability(cls) -> int: # Do not unconditionally lower the class-wide gate to 70. W4A4 @@ -2548,6 +2570,19 @@ def get_quant_method( moe_config=layer.moe_config, ) if quant_algo == "NVFP4": + if sm70_tm.is_exact_sm70_cuda_platform(): + if not sm70_tm.should_use_nvfp4_moe_turbomind(): + raise NotImplementedError( + "ModelOpt NVFP4 MoE on SM70 requires the TurboMind backend." + ) + from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( + ModelOptNvFp4SM70MoEMethod, + ) + + return ModelOptNvFp4SM70MoEMethod( + quant_config=self.nvfp4_config, + moe_config=layer.moe_config, + ) return ModelOptNvFp4FusedMoE( quant_config=self.nvfp4_config, moe_config=layer.moe_config, diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 7eb31d9ce3..ccc73f908c 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Native SM70 TurboMind NVFP4 MoE for Qwen3.6-35B-A3B. +"""Native SM70 TurboMind NVFP4 MoE for validated Qwen expert shapes. The route keeps ModelOpt W4A16_NVFP4 expert weights packed. It combines the checkpoint's FP8 block scales with its explicit ModelOpt global scales once at @@ -38,13 +38,15 @@ logger = init_logger(__name__) -_QWEN36_HIDDEN_SIZE: Final = 2048 -_QWEN36_INTERMEDIATE_SIZE: Final = 512 -_QWEN36_NUM_EXPERTS: Final = 256 -_QWEN36_TOP_K: Final = 8 -_QWEN36_SUPPORTED_TP_SIZES: Final = (1, 2, 4) +_SUPPORTED_CONTRACTS: Final = { + # (hidden size, global expert intermediate size, experts, top-k) + (2048, 512, 256, 8), # Qwen3.6-35B-A3B + (2560, 640, 512, 10), # Qwen3.8-Flash-Next +} +_SUPPORTED_TP_SIZES: Final = (1, 2, 4) _GRAPH_SAFE_MAX_TOKENS: Final = 18 _COMPACT_GROUPED_MAX_TOKENS: Final = 10 +_MAX_SUPPORTED_TOP_K: Final = max(contract[3] for contract in _SUPPORTED_CONTRACTS) @triton.jit @@ -80,7 +82,7 @@ def _prepare_compact_slot_groups( active_expert_ids: torch.Tensor, ) -> None: total_slots = sorted_expert_ids.numel() - max_slots = _COMPACT_GROUPED_MAX_TOKENS * _QWEN36_TOP_K + max_slots = _COMPACT_GROUPED_MAX_TOKENS * _MAX_SUPPORTED_TOP_K if not (0 < total_slots <= max_slots): raise ValueError(f"Unsupported SM70 NVFP4 active-expert slots: {total_slots}") block = triton.next_power_of_2(total_slots + 1) @@ -100,25 +102,10 @@ def _prepare_compact_slot_groups( def validate_nvfp4_sm70_moe_contract(moe: FusedMoEConfig) -> None: """Reject every topology outside the validated SM70 NVFP4 contract.""" - if moe.num_experts != _QWEN36_NUM_EXPERTS: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports Qwen3.6-35B-A3B " - f"with {_QWEN36_NUM_EXPERTS} experts, got {moe.num_experts}." - ) - if moe.experts_per_token != _QWEN36_TOP_K: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports top-k=" - f"{_QWEN36_TOP_K}, got {moe.experts_per_token}." - ) - if moe.hidden_dim != _QWEN36_HIDDEN_SIZE: - raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports hidden size " - f"{_QWEN36_HIDDEN_SIZE}, got {moe.hidden_dim}." - ) - if moe.tp_size not in _QWEN36_SUPPORTED_TP_SIZES: + if moe.tp_size not in _SUPPORTED_TP_SIZES: raise NotImplementedError( "SM70 TurboMind NVFP4 MoE currently supports tensor parallel " - f"sizes {_QWEN36_SUPPORTED_TP_SIZES}, got {moe.tp_size}." + f"sizes {_SUPPORTED_TP_SIZES}, got {moe.tp_size}." ) local_intermediate = moe.intermediate_size_per_partition if local_intermediate <= 0 or local_intermediate % NVFP4_GROUP_SIZE: @@ -126,11 +113,19 @@ def validate_nvfp4_sm70_moe_contract(moe: FusedMoEConfig) -> None: "SM70 TurboMind NVFP4 MoE requires a positive local intermediate " f"size divisible by {NVFP4_GROUP_SIZE}, got {local_intermediate}." ) - if local_intermediate * max(moe.tp_size, 1) != _QWEN36_INTERMEDIATE_SIZE: + global_intermediate = local_intermediate * max(moe.tp_size, 1) + contract = ( + moe.hidden_dim, + global_intermediate, + moe.num_experts, + moe.experts_per_token, + ) + if contract not in _SUPPORTED_CONTRACTS: raise NotImplementedError( - "SM70 TurboMind NVFP4 MoE currently supports Qwen3.6 expert " - f"intermediate size {_QWEN36_INTERMEDIATE_SIZE}; got local=" - f"{local_intermediate}, tp_size={moe.tp_size}." + "SM70 TurboMind NVFP4 MoE shape is not validated: " + f"hidden={moe.hidden_dim}, intermediate={global_intermediate}, " + f"experts={moe.num_experts}, top_k={moe.experts_per_token}. " + f"Validated contracts: {sorted(_SUPPORTED_CONTRACTS)}." ) if moe.moe_parallel_config.use_all2all_kernels: raise NotImplementedError( @@ -175,7 +170,7 @@ def _validate_weight_layout(layer: RoutedExperts) -> None: class ModelOptNvFp4SM70MoEMethod(ModelOptNvFp4FusedMoE): - """Qwen3.6 ModelOpt W4A16_NVFP4 experts on native TurboMind SM70.""" + """ModelOpt NVFP4 experts with FP16 activations on native SM70.""" def __init__( self, @@ -183,10 +178,10 @@ def __init__( moe_config: FusedMoEConfig, ) -> None: FusedMoEMethodBase.__init__(self, moe_config) - if quant_config.quant_method != "W4A16_NVFP4": + if quant_config.quant_method not in {"NVFP4", "W4A16_NVFP4"}: raise NotImplementedError( - "SM70 TurboMind ModelOpt NVFP4 MoE currently requires " - "W4A16_NVFP4 checkpoint weights." + "SM70 TurboMind ModelOpt NVFP4 MoE requires NVFP4-family " + f"checkpoint weights, got {quant_config.quant_method}." ) self.quant_config = quant_config self.use_a16 = True @@ -215,13 +210,11 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: missing = [name for name in required_ops if not hasattr(torch.ops._C, name)] if missing: raise RuntimeError( - "Qwen3.6 NVFP4 MoE on SM70 requires the TurboMind extension " + "SM70 NVFP4 MoE requires the TurboMind extension " "with " + ", ".join(missing) + "." ) if not hasattr(torch.ops._moe_C, "moe_permute_with_scratch"): - raise RuntimeError( - "Qwen3.6 NVFP4 MoE on SM70 requires graph-safe MoE permute ops." - ) + raise RuntimeError("SM70 NVFP4 MoE requires graph-safe MoE permute ops.") if self.moe.has_bias: raise NotImplementedError("SM70 NVFP4 MoE does not support expert bias.") if layer.activation != MoEActivation.SILU: @@ -316,6 +309,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.sm70_nvfp4_num_experts = num_experts layer.sm70_nvfp4_hidden_size = hidden layer.sm70_nvfp4_intermediate_size = intermediate + layer.sm70_nvfp4_top_k = int(layer.moe_config.experts_per_token) layer.sm70_nvfp4_w13_k_dim = hidden layer.sm70_nvfp4_w13_n_dim = 2 * intermediate layer.sm70_nvfp4_w2_k_dim = intermediate @@ -334,17 +328,21 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: del layer.w2_weight_scale_2 del layer.w2_input_scale logger.info_once( - "SM70 ModelOpt NVFP4 TurboMind MoE path enabled for " - "Qwen3.6-35B-A3B (local_experts=%d, graph_safe_decode=B1-B%d, " - "compact_grouped_decode=B1-B%d).", + "SM70 ModelOpt NVFP4 TurboMind MoE path enabled " + "(hidden=%d, local_intermediate=%d, local_experts=%d, top_k=%d, " + "graph_safe_decode=B1-B%d, compact_grouped_decode=B1-B%d).", + hidden, + intermediate, num_experts, + layer.sm70_nvfp4_top_k, _GRAPH_SAFE_MAX_TOKENS, _COMPACT_GROUPED_MAX_TOKENS, ) def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: device = layer.w13_tm_weight.device - max_slots = _GRAPH_SAFE_MAX_TOKENS * _QWEN36_TOP_K + top_k = int(layer.sm70_nvfp4_top_k) + max_slots = _GRAPH_SAFE_MAX_TOKENS * top_k experts = int(layer.sm70_nvfp4_num_experts) hidden = int(layer.sm70_nvfp4_hidden_size) intermediate = int(layer.sm70_nvfp4_intermediate_size) @@ -372,19 +370,19 @@ def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: ) layer._nvfp4_sm70_inv_permuted_idx = torch.empty( _GRAPH_SAFE_MAX_TOKENS, - _QWEN36_TOP_K, + top_k, dtype=torch.int32, device=device, ) layer._nvfp4_sm70_topk_ids = torch.empty( _GRAPH_SAFE_MAX_TOKENS, - _QWEN36_TOP_K, + top_k, dtype=torch.int32, device=device, ) layer._nvfp4_sm70_token_expert_indices = torch.arange( max_slots, dtype=torch.int32, device=device - ).view(_GRAPH_SAFE_MAX_TOKENS, _QWEN36_TOP_K) + ).view(_GRAPH_SAFE_MAX_TOKENS, top_k) layer._nvfp4_sm70_permuted_idx = torch.empty( max_slots, dtype=torch.int32, device=device ) @@ -417,7 +415,7 @@ def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: def _persistent_buffers( layer: RoutedExperts, num_tokens: int ) -> dict[str, torch.Tensor]: - slots = num_tokens * _QWEN36_TOP_K + slots = num_tokens * int(layer.sm70_nvfp4_top_k) return { "output": layer._nvfp4_sm70_output[:num_tokens], "permuted_input": layer._nvfp4_sm70_permuted_input[:slots], @@ -446,7 +444,8 @@ def _eager_buffers( layer: RoutedExperts, num_tokens: int ) -> dict[str, torch.Tensor]: device = layer.w13_tm_weight.device - slots = num_tokens * _QWEN36_TOP_K + top_k = int(layer.sm70_nvfp4_top_k) + slots = num_tokens * top_k experts = int(layer.sm70_nvfp4_num_experts) hidden = int(layer.sm70_nvfp4_hidden_size) intermediate = int(layer.sm70_nvfp4_intermediate_size) @@ -476,14 +475,14 @@ def _eager_buffers( experts + 1, dtype=torch.int64, device=device ), "inv_permuted_idx": torch.empty( - num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device + num_tokens, top_k, dtype=torch.int32, device=device ), "topk_ids": torch.empty( - num_tokens, _QWEN36_TOP_K, dtype=torch.int32, device=device + num_tokens, top_k, dtype=torch.int32, device=device ), "token_expert_indices": torch.arange( slots, dtype=torch.int32, device=device - ).view(num_tokens, _QWEN36_TOP_K), + ).view(num_tokens, top_k), "permuted_idx": torch.empty(slots, dtype=torch.int32, device=device), "sort_workspace": torch.empty( workspace_size, dtype=torch.int8, device=device @@ -528,13 +527,18 @@ def apply( raise TypeError("SM70 NVFP4 MoE requires CUDA FP16 activations [M, H].") if not is_exact_sm70_cuda(x, enabled=True): raise RuntimeError("SM70 NVFP4 MoE dispatch is restricted to CUDA SM70.") - if x.shape[1] != _QWEN36_HIDDEN_SIZE: + hidden = int(layer.sm70_nvfp4_hidden_size) + top_k = int(layer.sm70_nvfp4_top_k) + if x.shape[1] != hidden: raise ValueError( "SM70 NVFP4 MoE activation hidden size mismatch: expected " - f"{_QWEN36_HIDDEN_SIZE}, got {x.shape[1]}." + f"{hidden}, got {x.shape[1]}." + ) + if tuple(topk_ids.shape) != (x.shape[0], top_k): + raise ValueError( + "SM70 NVFP4 MoE top-k ID shape mismatch: expected " + f"{(x.shape[0], top_k)}, got {tuple(topk_ids.shape)}." ) - if tuple(topk_ids.shape) != (x.shape[0], _QWEN36_TOP_K): - raise ValueError("SM70 NVFP4 MoE requires top-k IDs with shape [M, 8].") if tuple(topk_weights.shape) != tuple(topk_ids.shape): raise ValueError("SM70 NVFP4 MoE top-k weights and IDs must share shape.") if topk_weights.dtype != torch.float32: @@ -542,11 +546,11 @@ def apply( num_tokens = x.shape[0] if num_tokens == 0: - return x.new_empty((0, _QWEN36_HIDDEN_SIZE)) + return x.new_empty((0, hidden)) buffers = self._get_buffers(layer, num_tokens) output = buffers["output"] output.zero_() - slots = num_tokens * _QWEN36_TOP_K + slots = num_tokens * top_k topk_ids_i32 = buffers["topk_ids"] topk_ids_i32.copy_(topk_ids, non_blocking=True) buffers["permuted_idx"].fill_(slots) @@ -557,7 +561,7 @@ def apply( layer.expert_map, layer.global_num_experts, layer.local_num_experts, - _QWEN36_TOP_K, + top_k, buffers["permuted_input"], buffers["expert_offsets64"], buffers["inv_permuted_idx"], @@ -613,7 +617,7 @@ def apply( topk_weights, buffers["inv_permuted_idx"], buffers["expert_offsets64"], - _QWEN36_TOP_K, + top_k, output, ) return output diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 9f69084c62..7bec6af393 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -824,6 +824,8 @@ def __init__( padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + quant_method: QuantizeMethodBase | None = None, ): super().__init__() self.prefix = prefix @@ -853,8 +855,10 @@ def __init__( ) self.embedding_dim = embedding_dim - quant_method = None - if quant_config is not None: + # Model-specific embeddings can preselect a storage method. This is + # required by Qwen4Exp PLE, whose table remains FP8 even when the + # routed experts use a different checkpoint quantization config. + if quant_method is None and quant_config is not None: quant_method = quant_config.get_quant_method(self, prefix=prefix) if quant_method is None: quant_method = UnquantizedEmbeddingMethod() @@ -1077,6 +1081,16 @@ def forward(self, input_): output_parallel = self.quant_method.embedding(self, masked_input.long()) # Mask the output embedding. if self.tp_size > 1: + if output_parallel.dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ): + # Each token row has exactly one TP owner. Communicate the raw + # FP8 bytes as int8 because NCCL does not reduce FP8 directly. + comm_output = output_parallel.view(torch.int8) + comm_output.masked_fill_(input_mask.unsqueeze(-1), 0) + output = tensor_model_parallel_all_reduce(comm_output) + return output.view(output_parallel.dtype) output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) # Reduce across all the model parallel GPUs. output = tensor_model_parallel_all_reduce(output_parallel) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 2a5f746d78..3b4fcfa042 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -139,7 +139,8 @@ def device_loading_context(module: torch.nn.Module, target_device: torch.device) # Store original device states and move parameters to GPU if they're on CPU for name, p in module.named_parameters(): - if p.device.type == "cpu": + keep_on_cpu = getattr(p, "_vllm_keep_on_cpu", False) + if p.device.type == "cpu" and not keep_on_cpu: original_device_states[name] = p.device p.data = p.data.to(target_device) if getattr(p, "_vllm_is_uva_offloaded", False): diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index de9786c777..7674b2c539 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -568,6 +568,77 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: ) +def _strip_qwen4_exp_mrope(model_config: "ModelConfig") -> None: + configs = { + id(config): config + for config in ( + getattr(model_config, "hf_config", None), + model_config.hf_text_config, + ) + if config is not None + } + for config in configs.values(): + rope_parameters = getattr(config, "rope_parameters", None) + if rope_parameters is not None: + rope_parameters.pop("mrope_section", None) + rope_parameters.pop("mrope_interleaved", None) + + +class Qwen4ExpForConditionalGenerationConfig(Qwen3_5ForConditionalGenerationConfig): + """Apply the Qwen3.5 hybrid-cache contract to Qwen4Exp.""" + + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen3_5ForConditionalGenerationConfig.verify_and_update_config(vllm_config) + text_config = vllm_config.model_config.hf_text_config + if text_config.hc_count <= 1: + raise ValueError("Qwen4Exp requires hc_count > 1") + + if vllm_config.cache_config.enable_prefix_caching: + raise NotImplementedError( + "Qwen4Exp prefix caching is not enabled in the initial SM70 " + "route; disable it while the QSA ring cache is in use" + ) + + parallel_config = vllm_config.parallel_config + uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( + getattr(text_config, "indexer_n_heads", None) is not None + ) + if uses_ple_or_qsa and ( + parallel_config.enable_dbo or parallel_config.ubatch_size > 1 + ): + raise NotImplementedError( + "Qwen4Exp PLE/QSA does not support dual-batch overlap or microbatching" + ) + + model_config = vllm_config.model_config + multimodal_config = model_config.multimodal_config + if multimodal_config is not None: + if not multimodal_config.language_model_only: + raise NotImplementedError( + "Qwen4Exp multimodal inference is not enabled in the initial " + "SM70 route; pass --language-model-only while the vision " + "tower remains outside the TP4 memory and quality gates" + ) + _strip_qwen4_exp_mrope(model_config) + + spec_config = vllm_config.speculative_config + if spec_config is not None: + raise NotImplementedError( + "Qwen4Exp speculative decoding is disabled for the initial " + "SM70 V2 route. The checkpoint's PLE n-gram embedding remains " + "enabled; disable --speculative-config until the native MTP " + "follow-up is quality-gated." + ) + + +class Qwen4ExpForCausalLMConfig(Qwen4ExpForConditionalGenerationConfig): + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + _strip_qwen4_exp_mrope(vllm_config.model_config) + + class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -628,6 +699,8 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig, "Qwen3_5ForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig, + "Qwen4ExpForConditionalGeneration": Qwen4ExpForConditionalGenerationConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, } diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index d6011a3e58..741b34fd05 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -32,11 +32,15 @@ from vllm.config import ModelConfig, SpeechToTextConfig, SpeechToTextParams from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncsByType, +) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.tasks import ScoreType from vllm.utils.collection_utils import common_prefix from vllm.utils.func_utils import supports_kw +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from .interfaces_base import VllmModel @@ -837,6 +841,14 @@ def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]: """ ... + @classmethod + def get_mamba_state_copy_funcs( + cls, + mamba_types: set[MambaAttentionBackendEnum], + ) -> MambaStateCopyFuncsByType: + copy_funcs = cls.get_mamba_state_copy_func() + return {mamba_type: copy_funcs for mamba_type in mamba_types} + @overload def is_hybrid(model: object) -> TypeIs[IsHybrid]: ... diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5723dbe285..e053e2def8 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -117,6 +117,10 @@ "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"), "Gemma4ForCausalLM": ("gemma4", "Gemma4ForCausalLM"), "Qwen3NextForCausalLM": ("qwen3_next", "Qwen3NextForCausalLM"), + "Qwen4ExpForCausalLM": ( + "vllm.models.qwen4_exp", + "Qwen4ExpForCausalLM", + ), "GlmForCausalLM": ("glm", "GlmForCausalLM"), "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"), "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"), @@ -562,6 +566,10 @@ "qwen3_5", "Qwen3_5MoeForConditionalGeneration", ), + "Qwen4ExpForConditionalGeneration": ( + "vllm.models.qwen4_exp", + "Qwen4ExpForConditionalGeneration", + ), "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"), "SkyworkR1VChatModel": ("skyworkr1v", "SkyworkR1VChatModel"), "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 095d0e363d..302e71dbd7 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -46,18 +46,32 @@ class WeightsMapper: orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict) orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict) + orig_to_new_stacked: Mapping[str, tuple[str, str | int | tuple[int, ...]]] = field( + default_factory=dict + ) orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict) orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict) def __or__(self, other: "WeightsMapper") -> "WeightsMapper": """Combine two `WeightsMapper`s by merging their mappings.""" return WeightsMapper( + orig_to_new_regex={**self.orig_to_new_regex, **other.orig_to_new_regex}, orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr}, + orig_to_new_stacked={ + **self.orig_to_new_stacked, + **other.orig_to_new_stacked, + }, orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix}, orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix}, ) def _map_name(self, key: str) -> str | None: + result = self._map_name_with_shard(key) + return result[0] if result is not None else None + + def _map_name_with_shard( + self, key: str + ) -> tuple[str, str | int | tuple[int, ...] | None] | None: for pattern, new_key in self.orig_to_new_regex.items(): if pattern.search(key): if new_key is None: @@ -72,6 +86,12 @@ def _map_name(self, key: str) -> str | None: key = key.replace(substr, new_key, 1) + shard_id: str | int | tuple[int, ...] | None = None + for substr, (new_key, new_shard_id) in self.orig_to_new_stacked.items(): + if substr in key: + key = key.replace(substr, new_key, 1) + shard_id = new_shard_id + for prefix, new_key in self.orig_to_new_prefix.items(): if key.startswith(prefix): if new_key is None: @@ -86,16 +106,19 @@ def _map_name(self, key: str) -> str | None: key = new_key.join(key.rsplit(suffix, 1)) - return key + return key, shard_id def apply( self, weights: Iterable[tuple[str, torch.Tensor]] ) -> Iterable[tuple[str, torch.Tensor]]: - return ( - (out_name, data) - for name, data in weights - if (out_name := self._map_name(name)) is not None - ) + for name, data in weights: + result = self._map_name_with_shard(name) + if result is None: + continue + out_name, shard_id = result + if shard_id is not None: + data.shard_id = shard_id + yield out_name, data def apply_list(self, values: list[str]) -> list[str]: return [ @@ -354,6 +377,62 @@ def load_weights( return autoloaded_weights +def maybe_fuse_shared_experts( + weights: Iterable[tuple[str, torch.Tensor]], + *, + n_routed_experts: int, + n_shared_experts: int, + ckpt_prefix: str = "mlp.shared_experts", + enabled: bool | None = None, +) -> Iterable[tuple[str, torch.Tensor]]: + """Route AITER fused-shared-expert weights into routed-expert slots. + + CUDA leaves the input stream unchanged. On ROCm, AITER can fuse shared + experts into the routed expert tensor; in that case split the checkpoint + tensor into the extra expert slots expected by ``RoutedExperts``. + """ + if enabled is None: + from vllm._aiter_ops import rocm_aiter_ops + + enabled = rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() + if not enabled: + yield from weights + return + + prefix = f"{ckpt_prefix}." + for name, loaded_weight in weights: + if prefix not in name: + yield name, loaded_weight + continue + + split_dim = 1 if ("down_proj.weight" in name and loaded_weight.ndim > 1) else 0 + total = loaded_weight.shape[split_dim] + if total % n_shared_experts != 0: + raise ValueError( + f"FSE shared-expert weight {name!r} has size {total} along axis " + f"{split_dim}, not divisible by " + f"n_shared_experts={n_shared_experts}." + ) + chunk = total // n_shared_experts + for shared_expert_idx in range(n_shared_experts): + chunk_slice = slice( + shared_expert_idx * chunk, (shared_expert_idx + 1) * chunk + ) + if loaded_weight.ndim == 1: + chunk_weight = loaded_weight[chunk_slice] + elif split_dim == 0: + chunk_weight = loaded_weight[chunk_slice, :] + else: + chunk_weight = loaded_weight[:, chunk_slice] + yield ( + name.replace( + prefix, + f"mlp.experts.{n_routed_experts + shared_expert_idx}.", + ), + chunk_weight, + ) + + def init_vllm_registered_model( vllm_config: VllmConfig, *, diff --git a/vllm/models/qwen4_exp/__init__.py b/vllm/models/qwen4_exp/__init__.py new file mode 100644 index 0000000000..9e533e28b4 --- /dev/null +++ b/vllm/models/qwen4_exp/__init__.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp model package.""" + +from typing import TYPE_CHECKING, Any + +from .common.hyperconnection import ( + GatedResidual, + GroupedGemmaRMSNorm, + HyperConnectionBase, + HyperConnectionConfig, +) + +if TYPE_CHECKING: + from .nvidia.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + from .nvidia.mtp import Qwen4ExpMTP + + +def __getattr__(name: str) -> Any: + if name in { + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMTP", + }: + from vllm.platforms import current_platform + + if current_platform.is_xpu() or current_platform.is_tpu(): + raise NotImplementedError("Qwen4Exp currently supports CUDA and ROCm only") + if current_platform.is_rocm(): + from .amd.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + from .amd.mtp import Qwen4ExpMTP + else: + from .nvidia.model import ( + Qwen4ExpForCausalLM, + Qwen4ExpForConditionalGeneration, + ) + from .nvidia.mtp import Qwen4ExpMTP + + return { + "Qwen4ExpForCausalLM": Qwen4ExpForCausalLM, + "Qwen4ExpForConditionalGeneration": (Qwen4ExpForConditionalGeneration), + "Qwen4ExpMTP": Qwen4ExpMTP, + }[name] + raise AttributeError(name) + + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionBase", + "HyperConnectionConfig", + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMTP", +] diff --git a/vllm/models/qwen4_exp/common/__init__.py b/vllm/models/qwen4_exp/common/__init__.py new file mode 100644 index 0000000000..a89d77b976 --- /dev/null +++ b/vllm/models/qwen4_exp/common/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Common Qwen4Exp model components.""" + +from .hyperconnection import ( + GatedResidual, + GroupedGemmaRMSNorm, + HyperConnectionBase, + HyperConnectionConfig, +) + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionBase", + "HyperConnectionConfig", +] diff --git a/vllm/models/qwen4_exp/common/hyperconnection.py b/vllm/models/qwen4_exp/common/hyperconnection.py new file mode 100644 index 0000000000..5f8bde0479 --- /dev/null +++ b/vllm/models/qwen4_exp/common/hyperconnection.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HyperConnection (Gated Residual) utilities. + +Implements the HyperConnection residual scheme proposed in +"HyperConnections" (https://arxiv.org/abs/2409.19606). + +The two concrete variants are: + - ``HyperConnectionBase`` - simple average pooling across hc_count parallel + streams (equivalent to hyperconnection_average). + - ``GatedResidual`` - learnable low-rank gated mixing and injection + (gated_residual). + +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner +(HC outer, HS inner — checkpoint-native layout). The local torch +implementation consumes the hyper input viewed as ``[..., HC, HS]``. + +Typical usage inside a transformer decoder layer:: + + self.attn_hc = GatedResidual(hc_config, role="attn") + self.mlp_hc = GatedResidual(hc_config, role="mlp") + + hidden_states, residual = self.attn_hc.mix(hidden_states) + hidden_states = attention(hidden_states) + hidden_states = self.attn_hc.combine(hidden_states, residual) + + hidden_states, residual = self.mlp_hc.mix(hidden_states) + hidden_states = mlp(hidden_states) + hidden_states = self.mlp_hc.combine(hidden_states, residual) +""" + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch import nn + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- +@dataclass +class HyperConnectionConfig: + """Configuration shared by all HyperConnection variants.""" + + hc_count: int = 4 + hidden_size: int = 64 + params_dtype: torch.dtype = torch.bfloat16 + mtp_hc: bool = False + hc_lowrank: int = 16 + rms_norm_eps: float = 1e-6 + hc_per_branch_norm: bool = False + + +class GroupedGemmaRMSNorm(nn.Module): + def __init__( + self, + hidden_size: int, + eps: float, + group_size: int | None, + dtype: torch.dtype | None, + ) -> None: + super().__init__() + if group_size is not None and hidden_size % group_size: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by " + f"group_size ({group_size})" + ) + self.variance_epsilon = eps + self.group_size = group_size + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + if self.group_size is None: + variance = hidden_states.square().mean(dim=-1, keepdim=True) + normalized = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + else: + grouped = hidden_states.unflatten( + -1, (hidden_states.shape[-1] // self.group_size, self.group_size) + ) + variance = grouped.square().mean(dim=-1, keepdim=True) + normalized = ( + grouped * torch.rsqrt(variance + self.variance_epsilon) + ).flatten(-2) + return (normalized * (1.0 + self.weight.float())).to(input_dtype) + + +# --------------------------------------------------------------------------- +# Average-pooling variant +# --------------------------------------------------------------------------- +class HyperConnectionBase(nn.Module): + """Average-pooling HyperConnection (``hyperconnection_average``). + + Splits the incoming ``[..., HC*HS]`` tensor (HC outer, HS inner) into + ``HC`` parallel streams, averages them for the block input, and + broadcasts the block output back to every stream. + """ + + def __init__( + self, + config: HyperConnectionConfig, + layer_idx: int | None = None, + role: str | None = None, + ) -> None: + super().__init__() + self.config = config + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.role = role + + @property + def hyper_hidden_size(self) -> int: + return self.hc_count * self.hidden_size + + def mix(self, hyper_input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Average the HC streams into a single block input.""" + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size + # [*, HC, HS] — mean over HC (dim=-2). + unflat = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) + mixed_input = unflat.mean(dim=-2) + return mixed_input, hyper_input + + def combine( + self, block_output: torch.Tensor, residual: torch.Tensor + ) -> torch.Tensor: + """Broadcast the block output back to every stream.""" + assert residual.shape[-1] == self.hc_count * self.hidden_size + assert block_output.shape[-1] == self.hidden_size + residual_reshaped = residual.unflatten(-1, (self.hc_count, self.hidden_size)) + combined = residual_reshaped + block_output.unsqueeze(-2) + return combined.flatten(-2) + + +# --------------------------------------------------------------------------- +# Gated-residual variant +# --------------------------------------------------------------------------- +class GatedResidual(HyperConnectionBase): + """Gated HyperConnection with learnable low-rank mixing and injection. + + ``mix()`` applies GemmaRMSNorm per HC stream and projects through a + low-rank sigmoid gate to produce a single block input. ``combine()`` + injects the block output back into each stream through a learned + per-stream injection weight. + + This implementation uses only PyTorch operators. Tensor-parallel + collectives are supplied by its caller. + """ + + def __init__( + self, + config: HyperConnectionConfig, + layer_idx: int | None = None, + role: str | None = None, + use_mix: bool = True, + use_combine: bool = True, + ) -> None: + super().__init__(config, layer_idx, role) + norm_size = ( + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size + ) + group_size = config.hidden_size if config.hc_per_branch_norm else None + # Normalize each H-sized HC stream independently while retaining a + # separate affine weight for every element of the HC*H layout. + self.hc_norm = GroupedGemmaRMSNorm( + norm_size, + eps=config.rms_norm_eps, + group_size=group_size, + dtype=config.params_dtype, + ) + + # -- raw Linear weights (checkpoint-compatible) ---------------------- + if use_mix: + self.input_mix_weight_down = nn.Linear( + self.hyper_hidden_size, + config.hc_lowrank, + bias=False, + dtype=config.params_dtype, + ) + self.input_mix_weight_up = nn.Linear( + config.hc_lowrank, + self.hyper_hidden_size, + bias=False, + dtype=config.params_dtype, + ) + if use_combine: + self.block_inject_weight = nn.Linear( + self.hyper_hidden_size, + self.hc_count, + bias=False, + dtype=config.params_dtype, + ) + + def _normalize(self, hyper_input: torch.Tensor) -> torch.Tensor: + if self.config.hc_per_branch_norm: + return self.hc_norm(hyper_input) + return self.hc_norm( + hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) + ).flatten(-2) + + def mix( + self, hyper_input: torch.Tensor + ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Mix: RMSNorm -> low-rank gate -> gated mean.""" + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size + if not hasattr(self, "input_mix_weight_down"): + raise RuntimeError("mix was disabled for this hyper-connection") + hyper_input_normed = self._normalize(hyper_input) + # Gate — original mix order: linear+silu then linear+sigmoid. + gate = F.silu( + F.linear(hyper_input_normed, self.input_mix_weight_down.weight) + / self.hc_count + ) + gate = torch.sigmoid(F.linear(gate, self.input_mix_weight_up.weight)).unflatten( + -1, (self.hc_count, self.hidden_size) + ) + mixed_input = ( + gate * hyper_input_normed.unflatten(-1, (self.hc_count, self.hidden_size)) + ).mean(dim=-2) + return mixed_input.to(hyper_input.dtype), (hyper_input, hyper_input_normed) + + def combine( + self, + block_output: torch.Tensor, + residuals: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + if not hasattr(self, "block_inject_weight"): + raise RuntimeError("combine was disabled for this hyper-connection") + hyper_input, hyper_input_normed = residuals + assert hyper_input.shape[-1] == self.hc_count * self.hidden_size + assert block_output.shape[-1] == self.hidden_size + residual = hyper_input.unflatten(-1, (self.hc_count, self.hidden_size)) + # The paired mix keeps its normalized hyper input so combine uses the + # same HC module's injection weight. + injection_weight = 2.0 * torch.sigmoid( + F.linear(hyper_input_normed, self.block_inject_weight.weight) + / self.hc_count + ) + output = residual + block_output.unsqueeze(-2) * injection_weight.unsqueeze(-1) + return output.flatten(-2).to(hyper_input.dtype) + + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionBase", + "HyperConnectionConfig", +] diff --git a/vllm/models/qwen4_exp/common/ple.py b/vllm/models/qwen4_exp/common/ple.py new file mode 100644 index 0000000000..8047f1ad10 --- /dev/null +++ b/vllm/models/qwen4_exp/common/ple.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Common Qwen4Exp PLE helpers.""" + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class PLEShardOverlap: + """Source and destination slices for one checkpoint embedding shard.""" + + source_start: int + destination_start: int + row_count: int + + +def compute_ple_shard_overlap( + *, + checkpoint_start: int, + checkpoint_rows: int, + tp_start: int, + tp_end: int, +) -> PLEShardOverlap | None: + """Compute the overlap of a checkpoint shard and one TP vocabulary range.""" + + if checkpoint_start < 0 or checkpoint_rows < 0: + raise ValueError("checkpoint shard bounds must be non-negative") + if tp_start < 0 or tp_end < tp_start: + raise ValueError("invalid TP vocabulary range") + checkpoint_end = checkpoint_start + checkpoint_rows + overlap_start = max(checkpoint_start, tp_start) + overlap_end = min(checkpoint_end, tp_end) + if overlap_start >= overlap_end: + return None + return PLEShardOverlap( + source_start=overlap_start - checkpoint_start, + destination_start=overlap_start - tp_start, + row_count=overlap_end - overlap_start, + ) + + +def copy_ple_embedding_shard_( + destination: torch.Tensor, + loaded_weight: torch.Tensor, + *, + checkpoint_start: int, + tp_start: int, + tp_end: int, +) -> int: + """Copy the overlapping rows of a PLE checkpoint shard into a TP table.""" + + if destination.ndim == 0 or loaded_weight.ndim != destination.ndim: + raise ValueError("destination and loaded weight must have matching ranks") + if destination.shape[1:] != loaded_weight.shape[1:]: + raise ValueError( + "embedding shard dimensions do not match: " + f"{tuple(destination.shape[1:])} != {tuple(loaded_weight.shape[1:])}" + ) + if destination.shape[0] < tp_end - tp_start: + raise ValueError("destination does not cover the requested TP range") + overlap = compute_ple_shard_overlap( + checkpoint_start=checkpoint_start, + checkpoint_rows=loaded_weight.shape[0], + tp_start=tp_start, + tp_end=tp_end, + ) + if overlap is None: + return 0 + source = loaded_weight.narrow(0, overlap.source_start, overlap.row_count) + target = destination.narrow(0, overlap.destination_start, overlap.row_count) + with torch.no_grad(): + target.copy_(source.to(device=target.device, dtype=target.dtype)) + return overlap.row_count diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py new file mode 100644 index 0000000000..bb5f890cef --- /dev/null +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -0,0 +1,824 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Paged side-cache ownership and metadata for Qwen4Exp QSA. + +Each QSA layer keeps a fixed circular buffer of raw index keys (the +compressor state) and one compressed key. MRoPE models pack exact three-axis +positions beside the raw keys; text models derive group positions from +logical positions. The compressor state uses one block per request, while +the compressed owner uses ``MLAAttentionSpec.compress_ratio`` so its block +table follows the main KV-cache lifecycle. Their physical tensor storage is +shared by the generic cache-layout planner. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from functools import cache +from typing import ClassVar + +import torch +from torch import nn + +from vllm.config import CacheConfig, VllmConfig +from vllm.config.cache import CacheDType +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + CircularBufferSpec, + KVCacheSpec, + MLAAttentionSpec, +) + + +def canonical_qsa_rope_positions(positions: torch.Tensor) -> torch.Tensor: + """Return exact per-token positions as ``[tokens, 1, 3]`` int64 rows.""" + + if positions.ndim == 1: + positions = positions.unsqueeze(0).expand(3, -1) + elif positions.ndim != 2 or positions.shape[0] not in (1, 3): + raise ValueError("QSA RoPE positions must be [tokens] or [1|3, tokens]") + if positions.shape[0] == 1: + positions = positions.expand(3, -1) + return positions.transpose(0, 1).unsqueeze(1).to(torch.int64) + + +def _logical_positions( + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + token_to_req: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + if num_tokens == 0: + return seq_lens.new_empty((0,), dtype=torch.int64) + arange = torch.arange(num_tokens, device=query_start_loc.device) + requests = token_to_req[:num_tokens].long() + query_lens = torch.diff(query_start_loc) + within_query = arange - query_start_loc.index_select(0, requests) + return ( + seq_lens.index_select(0, requests).long() + - query_lens.index_select(0, requests).long() + + within_query.long() + ) + + +def _logical_to_physical_qsa_slots( + block_table: torch.Tensor, + request_indices: torch.Tensor, + logical_positions: torch.Tensor, + block_size: int, +) -> torch.Tensor: + if block_size <= 0: + raise ValueError("QSA cache block size must be positive") + if block_table.ndim != 2: + raise ValueError("QSA block table must be two-dimensional") + if request_indices.shape != logical_positions.shape: + request_indices = torch.broadcast_to(request_indices, logical_positions.shape) + + requests = request_indices.to(device=block_table.device, dtype=torch.long) + positions = logical_positions.to(device=block_table.device, dtype=torch.long) + valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) + logical_blocks = torch.div( + positions.clamp_min(0), block_size, rounding_mode="floor" + ) + valid &= logical_blocks < block_table.shape[1] + safe_requests = requests.clamp(0, max(block_table.shape[0] - 1, 0)) + safe_blocks = logical_blocks.clamp(0, max(block_table.shape[1] - 1, 0)) + if not all(block_table.shape): + return torch.full_like(positions, PAD_SLOT_ID) + physical_blocks = block_table[safe_requests, safe_blocks].long() + valid &= physical_blocks >= 0 + slots = physical_blocks * block_size + positions.remainder(block_size) + return torch.where(valid, slots, PAD_SLOT_ID) + + +def circular_qsa_slot_mapping( + block_table: torch.Tensor, + token_to_req: torch.Tensor, + logical_positions: torch.Tensor, + compressor_state_size: int, + query_start_loc: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Map each request to its fixed physical block as a circular token ring.""" + + if compressor_state_size <= 0: + raise ValueError("QSA circular buffer size must be positive") + if block_table.ndim != 2: + raise ValueError("QSA block table must be two-dimensional") + + requests = token_to_req.to(device=block_table.device, dtype=torch.long) + positions = logical_positions.to(device=block_table.device, dtype=torch.long) + if not all(block_table.shape): + slots = torch.full_like(positions, PAD_SLOT_ID) + else: + valid = (requests >= 0) & (requests < block_table.shape[0]) & (positions >= 0) + safe_requests = requests.clamp(0, block_table.shape[0] - 1) + physical_blocks = block_table[safe_requests, 0].long() + valid &= physical_blocks >= 0 + slots = physical_blocks * compressor_state_size + positions.remainder( + compressor_state_size + ) + slots = torch.where(valid, slots, PAD_SLOT_ID) + + if query_start_loc is not None: + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: + raise ValueError("QSA query starts must contain a terminal offset") + query_start_loc = query_start_loc.to(block_table.device) + num_requests = query_start_loc.shape[0] - 1 + safe_requests = requests.clamp(0, num_requests - 1) + request_ends = query_start_loc.index_select(0, safe_requests + 1) + rows = torch.arange(slots.numel(), device=slots.device) + keep = ( + (requests >= 0) + & (requests < num_requests) + & (rows + compressor_state_size >= request_ends) + ) + slots = torch.where(keep, slots, PAD_SLOT_ID) + + slots = slots.to(torch.int64) + if out is not None: + out.fill_(PAD_SLOT_ID) + out[: slots.numel()].copy_(slots) + return out[: slots.numel()] + return slots + + +def compressed_qsa_slot_mapping( + block_table: torch.Tensor, + token_to_req: torch.Tensor, + logical_positions: torch.Tensor, + storage_block_size: int, + compress_ratio: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Build boundary-only slots for an ``MLAAttentionSpec`` QSA cache.""" + + if storage_block_size <= 0 or compress_ratio <= 0: + raise ValueError("QSA block size and compression ratio must be positive") + compressed_positions = torch.div( + logical_positions.clamp_min(0), compress_ratio, rounding_mode="floor" + ) + slots = _logical_to_physical_qsa_slots( + block_table, + token_to_req, + compressed_positions, + storage_block_size, + ) + valid = (logical_positions >= 0) & ( + (logical_positions + 1).remainder(compress_ratio) == 0 + ) + slots = torch.where(valid, slots, PAD_SLOT_ID).to(torch.int64) + if out is not None: + out.fill_(PAD_SLOT_ID) + out[: slots.numel()].copy_(slots) + return out[: slots.numel()] + return slots + + +@cache +def _metadata_launch_pdl() -> bool: + return current_platform.is_arch_support_pdl() + + +@triton.jit( + do_not_specialize=[ + "num_reqs", + "num_mapped_tokens", + "num_tokens", + "max_num_work", + "num_search_steps", + "work_search_steps", + ] +) +def _build_qsa_metadata_kernel( + query_start_loc_ptr, + seq_lens_ptr, + common_slot_mapping_ptr, + block_table_ptr, + token_to_req_ptr, + logical_positions_ptr, + slot_mapping_ptr, + k_work_metadata_ptr, + block_table_stride_0: tl.constexpr, + block_table_stride_1: tl.constexpr, + num_reqs, + num_mapped_tokens, + num_tokens, + max_num_work, + num_search_steps, + work_search_steps, + storage_block_size: tl.constexpr, + compress_ratio: tl.constexpr, + circular_buffer_size: tl.constexpr, + num_block_table_columns: tl.constexpr, + launch_pdl: tl.constexpr, + TOKEN_BLOCK_SIZE: tl.constexpr, + REQUEST_SCAN_SIZE: tl.constexpr, + WORK_BLOCK_SIZE: tl.constexpr, +): + if launch_pdl: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + token_idx = pid * TOKEN_BLOCK_SIZE + tl.arange(0, TOKEN_BLOCK_SIZE) + store_mask = token_idx < num_tokens + mapped = token_idx < num_mapped_tokens + search_token_idx = tl.minimum(token_idx, num_mapped_tokens - 1) + request_idx = tl.zeros((TOKEN_BLOCK_SIZE,), tl.int32) + # Find the last query start at or before each token. The dynamic loop avoids + # compiling a kernel variant for every ceil(log2(num_reqs)). + for step in tl.range(0, num_search_steps): + candidate = request_idx + (1 << (num_search_steps - step - 1)) + valid_candidate = candidate < num_reqs + candidate_start = tl.load( + query_start_loc_ptr + candidate, + mask=valid_candidate, + other=num_mapped_tokens + 1, + ) + advance = valid_candidate & (candidate_start <= search_token_idx) + request_idx = tl.where(advance, candidate, request_idx) + query_start = tl.load(query_start_loc_ptr + request_idx, mask=mapped, other=0) + query_end = tl.load(query_start_loc_ptr + request_idx + 1, mask=mapped, other=0) + seq_len = tl.load(seq_lens_ptr + request_idx, mask=mapped, other=0) + logical_position = seq_len - (query_end - query_start) + token_idx - query_start + logical_position = tl.where(mapped, logical_position, -1) + tl.store( + token_to_req_ptr + token_idx, + tl.where(mapped, request_idx, 0), + mask=store_mask, + ) + tl.store( + logical_positions_ptr + token_idx, + logical_position, + mask=store_mask, + ) + + # circular_buffer_size is constexpr, so each builder instance compiles out + # the other QSA cache owner's slot-mapping rule. + if circular_buffer_size > 0: + valid = ( + mapped + & (logical_position >= 0) + & (token_idx + circular_buffer_size >= query_end) + & (num_block_table_columns > 0) + ) + physical_block = tl.load( + block_table_ptr + request_idx * block_table_stride_0, + mask=valid, + other=-1, + ) + valid &= physical_block >= 0 + slot = physical_block * circular_buffer_size + ( + logical_position % circular_buffer_size + ) + elif compress_ratio != 1: + compressed_position = tl.maximum(logical_position, 0) // compress_ratio + logical_block = compressed_position // storage_block_size + valid = ( + mapped + & (logical_position >= 0) + & ((logical_position + 1) % compress_ratio == 0) + & (logical_block < num_block_table_columns) + ) + physical_block = tl.load( + block_table_ptr + + request_idx * block_table_stride_0 + + logical_block * block_table_stride_1, + mask=valid, + other=-1, + ) + valid &= physical_block >= 0 + valid &= ( + tl.load(common_slot_mapping_ptr + token_idx, mask=mapped, other=-1) >= 0 + ) + slot = physical_block * storage_block_size + ( + compressed_position % storage_block_size + ) + if (circular_buffer_size > 0) or (compress_ratio != 1): + tl.store( + slot_mapping_ptr + token_idx, + tl.where(valid, slot, -1), + mask=store_mask, + ) + work_tile_start = pid * WORK_BLOCK_SIZE + has_work_tile = work_tile_start < max_num_work + if k_work_metadata_ptr is not None and has_work_tile: + # Every work CTA builds the request prefix in registers. Recomputing this + # small vector lets CTAs write disjoint work tiles without a grid barrier. + requests = tl.arange(0, REQUEST_SCAN_SIZE) + valid_request = requests < num_reqs + request_query_start = tl.load( + query_start_loc_ptr + requests, mask=valid_request, other=0 + ) + request_query_end = tl.load( + query_start_loc_ptr + requests + 1, mask=valid_request, other=0 + ) + request_seq_len = tl.load(seq_lens_ptr + requests, mask=valid_request, other=0) + request_query_len = request_query_end - request_query_start + chunk_start = request_seq_len - request_query_len + num_groups = request_seq_len // compress_ratio - chunk_start // compress_ratio + # Nonempty requests need one item even without a completed compression + # group because work item zero also commits the current raw-K suffix. + work_counts = tl.where(request_query_len > 0, tl.maximum(num_groups, 1), 0) + work_ends = tl.cumsum(tl.where(valid_request, work_counts, 0), axis=0) + total_work = tl.sum(tl.where(valid_request, work_counts, 0), axis=0) + work_offsets = tl.arange(0, WORK_BLOCK_SIZE) + work = work_tile_start + work_offsets + in_bounds = work < max_num_work + active = in_bounds & (work < total_work) + request = tl.zeros((WORK_BLOCK_SIZE,), dtype=tl.int32) + # Request zero starts at zero for every active work item. Descending + # steps find the last request starting at or before this item. + for step_idx in tl.range(0, work_search_steps): + step = 1 << (work_search_steps - step_idx - 1) + candidate = request + step + valid_candidate = candidate < num_reqs + candidate_start = tl.gather(work_ends, candidate - 1, 0) + advance = active & valid_candidate & (candidate_start <= work) + request = tl.where(advance, candidate, request) + + if launch_pdl: + # Let the dependent grid start launch setup while these CTAs finish + # stores; its gdc_wait still orders access to the completed metadata. + tl.extra.cuda.gdc_launch_dependents() + + owner_work_start = tl.gather(work_ends, tl.maximum(request - 1, 0), 0) + owner_work_start = tl.where(request == 0, 0, owner_work_start) + work_in_request = work - owner_work_start + tl.store( + k_work_metadata_ptr + work * 2, + tl.where(active, request, -1), + mask=in_bounds, + ) + tl.store( + k_work_metadata_ptr + work * 2 + 1, + tl.where(active, work_in_request, -1), + mask=in_bounds, + ) + + else: + if launch_pdl: + # A dependent grid launches only after every CTA has signaled. + tl.extra.cuda.gdc_launch_dependents() + + +def build_qsa_metadata_triton( + common_attn_metadata: CommonAttentionMetadata, + token_to_req_buffer: torch.Tensor, + logical_positions_buffer: torch.Tensor, + slot_mapping_buffer: torch.Tensor, + *, + storage_block_size: int, + compress_ratio: int, + circular_buffer_size: int = 0, + k_work_metadata_buffer: torch.Tensor | None = None, + request_capacity: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build QSA side-cache and optional pre-indexer work metadata.""" + num_tokens = common_attn_metadata.num_actual_tokens + num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + token_to_req = token_to_req_buffer[:num_tokens] + logical_positions = logical_positions_buffer[:num_tokens] + slot_mapping = slot_mapping_buffer[:num_tokens] + num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1 + assert num_reqs > 0 + + if k_work_metadata_buffer is not None: + if request_capacity is None: + request_capacity = num_reqs + assert request_capacity >= num_reqs + # Pad for tl.arange while keeping the scan width stable across live batches. + request_scan_size = 1 << int(math.ceil(math.log2(request_capacity))) + max_num_work = k_work_metadata_buffer.shape[0] + else: + request_scan_size = 1 + max_num_work = 0 + + if num_tokens == 0 and k_work_metadata_buffer is None: + return token_to_req, logical_positions, slot_mapping + + block_table = common_attn_metadata.block_table_tensor + num_search_steps = int(math.ceil(math.log2(num_reqs))) + work_search_steps = int(math.ceil(math.log2(num_reqs))) + # The same grid covers token tiles and, for the compressed cache, work tiles. + num_token_blocks = cdiv(num_tokens, 128) + num_work_blocks = ( + cdiv(max_num_work, 256) if k_work_metadata_buffer is not None else 0 + ) + _build_qsa_metadata_kernel[(max(num_token_blocks, num_work_blocks, 1),)]( + common_attn_metadata.query_start_loc, + common_attn_metadata.seq_lens, + common_attn_metadata.slot_mapping, + block_table, + token_to_req, + logical_positions, + slot_mapping, + k_work_metadata_buffer, + block_table.stride(0), + block_table.stride(1), + num_reqs, + num_mapped_tokens, + num_tokens, + max_num_work, + num_search_steps, + work_search_steps, + storage_block_size, + compress_ratio, + circular_buffer_size, + block_table.shape[1], + launch_pdl=_metadata_launch_pdl(), + TOKEN_BLOCK_SIZE=128, + REQUEST_SCAN_SIZE=request_scan_size, + WORK_BLOCK_SIZE=256, + num_warps=4, + ) + if circular_buffer_size == 0 and compress_ratio == 1: + slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] + return token_to_req, logical_positions, slot_mapping + + +def _build_qsa_metadata_torch( + common_attn_metadata: CommonAttentionMetadata, + token_to_req_buffer: torch.Tensor, + logical_positions_buffer: torch.Tensor, + slot_mapping_buffer: torch.Tensor, + *, + storage_block_size: int, + compress_ratio: int, + circular_buffer_size: int = 0, + k_work_metadata_buffer: torch.Tensor | None = None, + request_capacity: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del request_capacity + num_tokens = common_attn_metadata.num_actual_tokens + num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + logical_positions = logical_positions_buffer[:num_tokens] + + token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[ + :num_tokens + ] + logical_positions[:num_mapped_tokens].copy_( + _logical_positions( + common_attn_metadata.query_start_loc, + common_attn_metadata.seq_lens, + token_to_req[:num_mapped_tokens], + num_mapped_tokens, + ) + ) + if num_mapped_tokens < num_tokens: + logical_positions[num_mapped_tokens:].fill_(-1) + if circular_buffer_size > 0: + slot_mapping = circular_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + circular_buffer_size, + query_start_loc=common_attn_metadata.query_start_loc, + out=slot_mapping_buffer, + ) + elif compress_ratio == 1: + slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] + else: + slot_mapping = compressed_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + storage_block_size, + compress_ratio, + slot_mapping_buffer, + ) + slot_mapping.masked_fill_( + common_attn_metadata.slot_mapping[:num_tokens] < 0, -1 + ) + if k_work_metadata_buffer is not None: + query_lens = ( + common_attn_metadata.query_start_loc[1:] + - common_attn_metadata.query_start_loc[:-1] + ) + chunk_starts = common_attn_metadata.seq_lens - query_lens + num_work_per_request = ( + common_attn_metadata.seq_lens // compress_ratio + - chunk_starts // compress_ratio + ) + num_work_per_request = torch.where( + query_lens > 0, num_work_per_request.clamp_min(1), 0 + ) + k_start_loc = torch.empty( + query_lens.shape[0] + 1, + dtype=torch.int32, + device=query_lens.device, + ) + k_start_loc[0] = 0 + torch.cumsum(num_work_per_request, 0, out=k_start_loc[1:]) + work = torch.arange( + k_work_metadata_buffer.shape[0], + device=k_work_metadata_buffer.device, + ) + requests = torch.searchsorted(k_start_loc[1:], work, right=True) + active = work < k_start_loc[-1] + work_in_request = ( + work - k_start_loc[requests.clamp_max(query_lens.shape[0] - 1)] + ) + k_work_metadata_buffer[:, 0].copy_( + torch.where(active, requests, -1).to(torch.int32) + ) + k_work_metadata_buffer[:, 1].copy_( + torch.where(active, work_in_request, -1).to(torch.int32) + ) + return token_to_req, logical_positions, slot_mapping + + +# Resolve the fallback outside the per-step metadata hot path. +build_qsa_metadata = ( + build_qsa_metadata_triton if HAS_TRITON else _build_qsa_metadata_torch +) + + +@dataclass +class QSAForwardMetadata(AttentionMetadata): + """Common per-forward metadata for one QSA side cache.""" + + block_table: torch.Tensor + slot_mapping: torch.Tensor + seq_lens: torch.Tensor + query_start_loc: torch.Tensor + token_to_req: torch.Tensor + logical_positions: torch.Tensor + k_work_metadata: torch.Tensor + num_actual_tokens: int + storage_block_size: int + compress_ratio: int + + +class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]): + """Build QSA metadata from vLLM's cache-group-specific common metadata.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec) + if isinstance(kv_cache_spec, MLAAttentionSpec): + self.compress_ratio = kv_cache_spec.compress_ratio + else: + self.compress_ratio = 1 + self.storage_block_size = kv_cache_spec.storage_block_size + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.token_to_req_buffer = torch.empty( + max_tokens, dtype=torch.int32, device=device + ) + self.slot_mapping_buffer = torch.empty( + max_tokens, dtype=torch.int64, device=device + ) + self.logical_positions_buffer = torch.empty( + max_tokens, dtype=torch.int64, device=device + ) + max_requests = vllm_config.scheduler_config.max_num_seqs + self.request_capacity = max_requests + if not self.is_circular_buffer and self.compress_ratio != 1: + max_k_work = ( + max_tokens + (self.compress_ratio - 1) * max_requests + ) // self.compress_ratio + self.k_work_metadata_buffer = torch.empty( + max_k_work, 2, dtype=torch.int32, device=device + ) + else: + self.k_work_metadata_buffer = torch.empty( + 0, 2, dtype=torch.int32, device=device + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> QSAForwardMetadata: + del common_prefix_len, fast_build + num_tokens = common_attn_metadata.num_actual_tokens + build_k_work = not self.is_circular_buffer and self.compress_ratio != 1 + k_work_metadata = self.k_work_metadata_buffer + request_capacity = None + if build_k_work: + num_requests = common_attn_metadata.query_start_loc.shape[0] - 1 + request_capacity = self.request_capacity + max_num_work = ( + num_tokens + (self.compress_ratio - 1) * num_requests + ) // self.compress_ratio + k_work_metadata = self.k_work_metadata_buffer[:max_num_work] + token_to_req, logical_positions, slot_mapping = build_qsa_metadata( + common_attn_metadata, + self.token_to_req_buffer, + self.logical_positions_buffer, + self.slot_mapping_buffer, + storage_block_size=self.storage_block_size, + compress_ratio=self.compress_ratio, + circular_buffer_size=( + self.kv_cache_spec.block_size if self.is_circular_buffer else 0 + ), + k_work_metadata_buffer=k_work_metadata if build_k_work else None, + request_capacity=request_capacity, + ) + return QSAForwardMetadata( + block_table=common_attn_metadata.block_table_tensor, + slot_mapping=slot_mapping, + seq_lens=common_attn_metadata.seq_lens, + query_start_loc=common_attn_metadata.query_start_loc, + token_to_req=token_to_req, + logical_positions=logical_positions, + k_work_metadata=k_work_metadata, + num_actual_tokens=num_tokens, + storage_block_size=self.storage_block_size, + compress_ratio=self.compress_ratio, + ) + + +class QSAStateBackend(AttentionBackend): + """Key-only dummy backend for out-of-band QSA side-cache operations.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + ] + + @staticmethod + def get_name() -> str: + return "QWEN4_EXP_EXP_QSA_STATE" + + @staticmethod + def get_impl_cls(): + raise NotImplementedError( + "QSA state caches run out-of-band and have no attention impl" + ) + + @staticmethod + def get_builder_cls() -> type[QSAMetadataBuilder]: + return QSAMetadataBuilder + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + del cache_dtype_str + if num_kv_heads != 1: + raise ValueError("QSA side caches require exactly one KV head") + return (num_blocks, block_size, num_kv_heads, head_size) + + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + if include_num_layers_dimension: + return (0, 1, 2, 3, 4) + return (0, 1, 2, 3) + + +class _QSAStateCache(nn.Module, AttentionLayerBase): + supports_dcp = False + + def __init__( + self, + *, + head_size: int, + dtype: torch.dtype, + cache_config: CacheConfig, + prefix: str, + vllm_config: VllmConfig, + compress_ratio: int = 1, + ) -> None: + super().__init__() + if head_size <= 0: + raise ValueError("QSA cache head size must be positive") + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + if cache_config.block_size % compress_ratio: + raise ValueError( + "QSA cache block size must be divisible by the compression ratio" + ) + self.head_size = head_size + self.dtype = dtype + self.cache_config = cache_config + self.prefix = prefix + self.compress_ratio = compress_ratio + self.kv_cache = torch.tensor([]) + + static_context = vllm_config.compilation_config.static_forward_context + if prefix in static_context: + raise ValueError(f"Duplicate layer name: {prefix}") + static_context[prefix] = self + + def forward(self) -> None: ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return QSAStateBackend + + +class QSAKeyStateCache(_QSAStateCache): + """Raw 16-bit key, optionally followed by exact int64 MRoPE positions.""" + + _BF16_PER_INT64 = 4 + _NUM_ROPE_AXES = 3 + + def __init__(self, *, cache_rope_positions: bool = False, **kwargs) -> None: + key_head_size = int(kwargs.pop("head_size")) + self.key_head_size = key_head_size + self.cache_rope_positions = bool(cache_rope_positions) + self.rope_position_offset = ( + (key_head_size + self._BF16_PER_INT64 - 1) // self._BF16_PER_INT64 + ) * self._BF16_PER_INT64 + storage_head_size = key_head_size + if self.cache_rope_positions: + storage_head_size = self.rope_position_offset + ( + self._NUM_ROPE_AXES * self._BF16_PER_INT64 + ) + super().__init__(head_size=storage_head_size, **kwargs) + + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + if kv_cache.ndim != 4 or kv_cache.shape[2] != 1: + raise ValueError("QSA raw cache must be [blocks, block_size, 1, width]") + if kv_cache.dtype != self.dtype or kv_cache.shape[3] != self.head_size: + raise ValueError("QSA raw cache does not match its packed cache spec") + super().bind_kv_cache(kv_cache) + self.key_cache = kv_cache[..., : self.key_head_size] + if self.cache_rope_positions: + position_tail = kv_cache[..., self.rope_position_offset :] + self.rope_position_cache = position_tail.view(torch.int64) + else: + self.rope_position_cache = None + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # Hold the open group's committed keys plus every row a speculative + # step stores before acceptance is known, rounded up to whole groups so + # the ring divides the attention block size (it joins the LCM that sets + # the scheduler block size). Anything narrower lets a rejected draft row + # overwrite a committed key the next step needs to close the group. + span = self.compress_ratio + vllm_config.num_speculative_tokens + capacity = self.compress_ratio * cdiv(span, self.compress_ratio) + assert self.cache_config.block_size % capacity == 0, ( + f"QSA ring capacity {capacity} must divide the attention block " + f"size {self.cache_config.block_size}" + ) + return CircularBufferSpec( + block_size=capacity, + num_kv_heads=1, + head_size=self.head_size, + head_size_v=0, + dtype=self.dtype, + ) + + +class QSACompressedKeyCache(_QSAStateCache): + """Normalized, group-first-RoPE 16-bit key per complete group.""" + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + del vllm_config + return MLAAttentionSpec( + block_size=self.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_size, + dtype=self.dtype, + compress_ratio=self.compress_ratio, + ) + + +__all__ = [ + "QSACompressedKeyCache", + "QSAForwardMetadata", + "QSAKeyStateCache", + "QSAMetadataBuilder", + "QSAStateBackend", + "canonical_qsa_rope_positions", + "circular_qsa_slot_mapping", + "compressed_qsa_slot_mapping", +] diff --git a/vllm/models/qwen4_exp/config.py b/vllm/models/qwen4_exp/config.py new file mode 100644 index 0000000000..a8eedab83f --- /dev/null +++ b/vllm/models/qwen4_exp/config.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp model configuration.""" + +from typing import Any, ClassVar, cast + +from transformers import PretrainedConfig +from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLVisionConfig, +) + +from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig + +_QSA_CONFIG_FIELDS = ( + "indexer_n_heads", + "indexer_kv_heads", + "indexer_head_dim", + "indexer_budget", + "indexer_compress_ratio", +) + + +class Qwen4ExpVisionConfig(Qwen3VLVisionConfig): + model_type = "qwen4_exp" + base_config_key = "vision_config" + + +class Qwen4ExpTextConfig(Qwen3NextConfig): + model_type = "qwen4_exp_text" + base_config_key = "text_config" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + hc_count: int = 4, + hc_lowrank: int = 320, + ple_layer_ids: list[int] | None = None, + ple_embed_dim: int | None = None, + ple_conv_kernel_size: int = 4, + ngram_size: int = 3, + heads_per_ngram: int = 8, + ngram_vocab_size_base: int = 20_000_000, + make_ngram_vocab_size_divisible_by: int = 128, + output_gate_type: str = "sigmoid", + rope_parameters: dict[str, Any] | None = None, + layer_types: list[str] | None = None, + **kwargs: Any, + ) -> None: + if hc_count <= 1: + raise ValueError(f"Qwen4Exp requires hc_count > 1, got {hc_count}.") + + if rope_parameters is not None: + if kwargs.get("rope_scaling") is None: + kwargs["rope_scaling"] = rope_parameters + if kwargs.get("rope_theta") is None and "rope_theta" in rope_parameters: + kwargs["rope_theta"] = rope_parameters["rope_theta"] + if ( + kwargs.get("partial_rotary_factor") is None + and "partial_rotary_factor" in rope_parameters + ): + kwargs["partial_rotary_factor"] = rope_parameters[ + "partial_rotary_factor" + ] + + rope_scaling = kwargs.get("rope_scaling") + rope_theta = kwargs.get("rope_theta", 10_000.0) + super().__init__(layer_types=layer_types, **kwargs) + + normalized_rope_parameters = self.rope_parameters + self.rope_scaling = ( + rope_scaling or rope_parameters or normalized_rope_parameters + ) + self.rope_parameters = rope_parameters or normalized_rope_parameters + self.rope_theta = rope_theta + + self.hc_count = hc_count + self.hc_lowrank = hc_lowrank + self.ple_layer_ids = ple_layer_ids or [] + self.ple_embed_dim = ( + self.hidden_size if ple_embed_dim is None else ple_embed_dim + ) + self.ple_conv_kernel_size = ple_conv_kernel_size + self.ngram_size = ngram_size + self.heads_per_ngram = heads_per_ngram + self.ngram_vocab_size_base = ngram_vocab_size_base + self.make_ngram_vocab_size_divisible_by = make_ngram_vocab_size_divisible_by + self.output_gate_type = output_gate_type + + self._validate_ple_config() + self._validate_ple_layer_ids() + self._validate_qsa_config() + + def _validate_ple_config(self) -> None: + if self.hc_lowrank <= 0: + raise ValueError(f"hc_lowrank must be positive, got {self.hc_lowrank}") + if self.ngram_size < 2: + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") + if self.heads_per_ngram <= 0: + raise ValueError( + f"heads_per_ngram must be positive, got {self.heads_per_ngram}" + ) + if self.ple_embed_dim <= 0: + raise ValueError( + f"ple_embed_dim must be positive, got {self.ple_embed_dim}" + ) + ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram + if self.ple_embed_dim % ngram_heads: + raise ValueError( + "ple_embed_dim must be divisible by total ngram heads: " + f"{self.ple_embed_dim} % {ngram_heads} != 0" + ) + if self.ple_conv_kernel_size <= 0: + raise ValueError( + "ple_conv_kernel_size must be positive, got " + f"{self.ple_conv_kernel_size}" + ) + if self.ngram_vocab_size_base <= 0: + raise ValueError("ngram_vocab_size_base must be positive") + if self.make_ngram_vocab_size_divisible_by <= 0: + raise ValueError("make_ngram_vocab_size_divisible_by must be positive") + + def _validate_ple_layer_ids(self) -> None: + invalid = [ + layer_id + for layer_id in self.ple_layer_ids + if not 1 <= int(layer_id) <= self.num_hidden_layers + ] + if invalid: + raise ValueError( + "ple_layer_ids are 1-based and must refer to an existing layer; " + f"got {invalid} for {self.num_hidden_layers} layers" + ) + + def _validate_qsa_config(self) -> None: + configured = {name: getattr(self, name, None) for name in _QSA_CONFIG_FIELDS} + if all(value is None for value in configured.values()): + return + + missing = [name for name, value in configured.items() if value is None] + if missing: + raise ValueError(f"QSA config is missing required fields: {missing}") + + values = {name: int(cast(int, value)) for name, value in configured.items()} + if any(value <= 0 for value in values.values()): + raise ValueError(f"QSA config values must be positive: {values}") + if values["indexer_kv_heads"] != 1: + raise ValueError("the QSA MQA operators require indexer_kv_heads=1") + if values["indexer_budget"] % values["indexer_compress_ratio"] != 0: + raise ValueError( + "indexer_budget must be divisible by indexer_compress_ratio" + ) + block_topk = values["indexer_budget"] // values["indexer_compress_ratio"] + if block_topk not in (512, 2048): + raise ValueError( + "QSA requires indexer_budget / indexer_compress_ratio " + f"to be 512 or 2048, got {block_topk}" + ) + rotary_dim = int(self.head_dim * self.partial_rotary_factor) + if rotary_dim > values["indexer_head_dim"]: + raise ValueError( + "QSA indexer_head_dim must cover the attention rotary " + f"dimension, got {values['indexer_head_dim']} < {rotary_dim}" + ) + + @property + def layers_block_type(self) -> list[str]: + return [ + "attention" if layer_type == "full_attention" else layer_type + for layer_type in self.layer_types + ] + + @property + def short_conv_layer_ids(self) -> list[int]: + if not self.ple_layer_ids: + return [] + return sorted({int(layer_id) - 1 for layer_id in self.ple_layer_ids}) + + @property + def short_conv_state_shape(self) -> tuple[int, int] | None: + if not self.short_conv_layer_ids: + return None + ple_state_len = (self.ple_conv_kernel_size - 1) * self.ngram_size + ple_channels = self.hidden_size * self.hc_count + return ple_channels, ple_state_len + + @property + def ngram_context_len(self) -> int: + if not self.ple_layer_ids: + return 0 + return max(int(self.ngram_size) - 1, 0) + + +class Qwen4ExpConfig(PretrainedConfig): + model_type = "qwen4_exp" + sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = { + "vision_config": Qwen4ExpVisionConfig, + "text_config": Qwen4ExpTextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config: Qwen4ExpTextConfig | dict[str, Any] | None = None, + vision_config: Qwen4ExpVisionConfig | dict[str, Any] | None = None, + image_token_id: int = 248056, + video_token_id: int = 248057, + vision_start_token_id: int = 248053, + vision_end_token_id: int = 248054, + tie_word_embeddings: bool = False, + rope_parameters: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + if text_config is not None: + kwargs.pop("split_ngram_parts", None) + + text_kwargs = ( + dict(kwargs) + if text_config is None + and "hidden_size" in kwargs + and "num_hidden_layers" in kwargs + else {} + ) + + if isinstance(vision_config, dict): + self.vision_config = self.sub_configs["vision_config"](**vision_config) + elif vision_config is None: + self.vision_config = self.sub_configs["vision_config"]() + else: + self.vision_config = vision_config + + if isinstance(text_config, dict): + self.text_config = self.sub_configs["text_config"](**text_config) + elif text_config is None: + self.text_config = self.sub_configs["text_config"](**text_kwargs) + else: + self.text_config = text_config + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + self.rope_parameters = rope_parameters or getattr( + self.text_config, "rope_parameters", {} + ) + super().__init__(**kwargs, tie_word_embeddings=tie_word_embeddings) + + +__all__ = [ + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", +] diff --git a/vllm/models/qwen4_exp/nvidia/__init__.py b/vllm/models/qwen4_exp/nvidia/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/qwen4_exp/nvidia/hyperconnection.py b/vllm/models/qwen4_exp/nvidia/hyperconnection.py new file mode 100644 index 0000000000..8ca503d214 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/hyperconnection.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HyperConnection (Gated Residual) utilities — NVIDIA model variant. + +Implements the HyperConnection residual scheme proposed in +"HyperConnections" (https://arxiv.org/abs/2409.19606). This NVIDIA variant +delays each HC combine to the following HC mix boundary. HC glue kernels, +including fused combine+RMSNorm, live in ``ops/hc.py``; projections remain +standard vLLM Linear modules. + +Hidden states between layers have shape ``[..., HC*HS]`` with HS inner +(HC outer, HS inner — checkpoint-native layout). + +Typical usage inside a transformer decoder layer:: + + self.attn_hc = GatedResidual(hc_config) + + hidden_states, block_input, injection = self.attn_hc.mix(hidden_states) + attention_output = attention(block_input) + hidden_states, block_input, injection = self.mlp_hc.combine_and_mix( + hidden_states, attention_output, injection + ) +""" + +import torch +from torch import nn + +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, +) +from vllm.model_executor.models.utils import maybe_prefix + +from ..common.hyperconnection import ( + GroupedGemmaRMSNorm, + HyperConnectionConfig, +) +from .ops.hc import ( + grouped_gemma_rmsnorm, + hc_combine, + hc_combine_norm, + hc_gate_mix, + hc_silu, +) + + +# --------------------------------------------------------------------------- +# Gated-residual variant +# --------------------------------------------------------------------------- +class GatedResidual(nn.Module): + """Gated HyperConnection with learnable low-rank mixing and injection. + + ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged + low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean + over the HC streams). When passed a pending block output and an injection, + it fuses their residual combine with the RMSNorm. Final mixers use + ``use_combine=False`` and do not produce a new injection. + + Weights: the norm owns the grouped GemmaRMSNorm affine; the projections + are vLLM Linear modules (merged replicated linear for down+inject), so + GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the + standard quant_method mechanism. + """ + + def __init__( + self, + config: HyperConnectionConfig, + use_combine: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.lora_rank = config.hc_lowrank + self.hc_count = config.hc_count + self.hidden_size = config.hidden_size + self.use_combine = use_combine + + norm_size = ( + self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size + ) + group_size = config.hidden_size if config.hc_per_branch_norm else None + # Normalize each H-sized HC stream independently while retaining a + # separate affine weight for every element of the HC*H layout. + self.hc_norm = GroupedGemmaRMSNorm( + norm_size, + eps=config.rms_norm_eps, + group_size=group_size, + dtype=config.params_dtype, + ) + + # -- vLLM Linear weights -------------------------------------------- + # The merged skinny-GEMM shape is physically padded to 16 rows to ensure + # good alignment and performant implementation chosen by CuBLAS heuristics. + self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0 + if use_combine: + self.input_mix_weight_down_block_inject = MergedColumnParallelLinear( + self.hyper_hidden_size, + [self.lora_rank, self.hc_count] + + ([self.pad_size] if self.pad_size else []), + bias=False, + params_dtype=config.params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "input_mix_weight_down_block_inject"), + return_bias=False, + disable_tp=True, + ) + else: + self.input_mix_weight_down = ReplicatedLinear( + self.hyper_hidden_size, + self.lora_rank, + bias=False, + params_dtype=config.params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "input_mix_weight_down"), + return_bias=False, + ) + self.input_mix_weight_up = ReplicatedLinear( + self.lora_rank, + self.hyper_hidden_size, + bias=False, + params_dtype=config.params_dtype, + quant_config=None, + prefix=maybe_prefix(prefix, "input_mix_weight_up"), + return_bias=False, + ) + + def mix( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + xn = grouped_gemma_rmsnorm( + hidden_states, + self.hc_norm.weight, + self.config.rms_norm_eps, + self.hc_count, + ) + + if self.use_combine: + # produce injection logits for combine + split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + down_and_injection = self.input_mix_weight_down_block_inject(xn) + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + else: + lora = self.input_mix_weight_down(xn) + injection = None + + lora = hc_silu(lora, self.hc_count) + gate = self.input_mix_weight_up(lora) # [M, D] + block_input = hc_gate_mix(xn, gate, self.hc_count) + + return hidden_states, block_input, injection + + def combine_and_mix( + self, + hidden_states: torch.Tensor, + prev_block_output: torch.Tensor, + prev_injection: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Consume a pending combine, then prepare the next block input. + + ``hidden_states`` is the multi-stream state from before the pending + block's mix. Its combine with ``block_output`` is fused with this + module's input RMSNorm. + """ + hidden_states, xn = hc_combine_norm( + hidden_states, + prev_block_output, + prev_injection, + self.hc_norm.weight, + self.config.rms_norm_eps, + self.hc_count, + ) + + if self.use_combine: + # produce injection logits for combine + split_sizes = [self.lora_rank, self.hc_count, self.pad_size] + down_and_injection = self.input_mix_weight_down_block_inject(xn) + lora, injection, _ = down_and_injection.split(split_sizes, dim=-1) + else: + lora = self.input_mix_weight_down(xn) + injection = None + + lora = hc_silu(lora, self.hc_count) + gate = self.input_mix_weight_up(lora) # [M, D] + block_input = hc_gate_mix(xn, gate, self.hc_count) + + return hidden_states, block_input, injection + + def combine( + self, + hidden_states: torch.Tensor, + block_output: torch.Tensor, + injection: torch.Tensor, + ) -> torch.Tensor: + return hc_combine(hidden_states, block_output, injection, self.hc_count) + + @property + def hyper_hidden_size(self) -> int: + return self.hc_count * self.hidden_size + + +__all__ = [ + "GatedResidual", + "GroupedGemmaRMSNorm", + "HyperConnectionConfig", +] diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py new file mode 100644 index 0000000000..e704624146 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -0,0 +1,355 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp weight-free QSA indexer.""" + +from __future__ import annotations + +from typing import cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) + +from ..common.qsa_cache import ( + QSACompressedKeyCache, + QSAForwardMetadata, + QSAKeyStateCache, + canonical_qsa_rope_positions, +) +from .ops.qsa_pre_indexer import qsa_pre_indexer + + +def apply_qsa_rope( + rotary_emb: nn.Module, + positions: torch.Tensor, + tensor: torch.Tensor, +) -> torch.Tensor: + """Apply the main attention's exact 1D/MRoPE composition to QSA heads.""" + + num_tokens, _, head_dim = tensor.shape + rotary_dim = rotary_emb.rotary_dim + cache = rotary_emb._match_cos_sin_cache_dtype(tensor) # noqa: SLF001 + cos_sin = cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if positions.ndim == 2: + shape = tensor.shape + tensor, _ = triton_mrope( + tensor.reshape(num_tokens, -1), + tensor.new_empty((num_tokens, head_dim)), + cos, + sin, + rotary_emb.mrope_section, + head_dim, + rotary_dim, + rotary_emb.mrope_interleaved, + ) + return tensor.reshape(shape) + + rotated = rotary_emb.apply_rotary_emb.forward_cuda( + tensor[..., :rotary_dim], + cos, + sin, + ) + return torch.cat((rotated, tensor[..., rotary_dim:]), dim=-1) + + +def _supports_fused_pre_indexer( + rotary_emb: nn.Module, + head_dim: int, + num_kv_heads: int, + compress_ratio: int, +) -> bool: + rotary_dim = int(rotary_emb.rotary_dim) + mrope_section = getattr(rotary_emb, "mrope_section", None) + return ( + bool(getattr(rotary_emb, "is_neox_style", False)) + and ( + not mrope_section + or ( + len(mrope_section) == 3 + and sum(mrope_section) == rotary_dim // 2 + and bool(getattr(rotary_emb, "mrope_interleaved", False)) + ) + ) + and head_dim == 128 + and rotary_dim == 64 + and num_kv_heads == 1 + and compress_ratio > 1 + and compress_ratio & (compress_ratio - 1) == 0 + ) + + +class QSAIndexer(nn.Module): + """Replicated Q/K projection plus paged, weight-free QSA selection. + + ``prefix`` must be the checkpoint's indexer prefix, normally + ``model.layers.N.self_attn.indexer``. Consequently the trainable names are + ``index_qk_proj``, ``q_layernorm`` and ``k_layernorm`` under that prefix. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: Qwen4ExpTextConfig, + layer_id: int, + rotary_emb: nn.Module, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if vllm_config.cache_config is None: + raise ValueError("QSA requires a paged KV cache") + if vllm_config.model_config.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError("Qwen4Exp QSA requires FP16 or BF16") + + self.layer_id = int(layer_id) + self.index_n_heads = int(config.indexer_n_heads) + self.index_kv_heads = int(config.indexer_kv_heads) + self.index_head_dim = int(config.indexer_head_dim) + self.token_topk = int(config.indexer_budget) + self.compress_ratio = int(config.indexer_compress_ratio) + self.rotary_emb = rotary_emb + self.use_fused_pre_indexer = _supports_fused_pre_indexer( + rotary_emb, + self.index_head_dim, + self.index_kv_heads, + self.compress_ratio, + ) + self.prefix = prefix + # MTP step 0 selects the target-aligned rows; later steps reuse them + # while continuing to update the QSA side cache. + self.skip_topk = False + + self.index_qk_proj = ReplicatedLinear( + int(config.hidden_size), + (self.index_n_heads + self.index_kv_heads) * self.index_head_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.index_qk_proj" if prefix else "index_qk_proj", + ) + self.q_layernorm = GemmaRMSNorm( + self.index_head_dim, + eps=float(getattr(config, "rms_norm_eps", 1e-6)), + ) + self.k_layernorm = GemmaRMSNorm( + self.index_head_dim, + eps=float(getattr(config, "rms_norm_eps", 1e-6)), + ) + + cache_config = vllm_config.cache_config + cache_prefix = f"{prefix}." if prefix else "" + self.raw_key_cache = QSAKeyStateCache( + head_size=self.index_head_dim, + dtype=vllm_config.model_config.dtype, + cache_rope_positions=vllm_config.model_config.uses_mrope, + prefix=f"{cache_prefix}raw_key_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + vllm_config=vllm_config, + ) + self.compressed_key_cache = QSACompressedKeyCache( + head_size=self.index_head_dim, + dtype=vllm_config.model_config.dtype, + compress_ratio=self.compress_ratio, + prefix=f"{cache_prefix}compressed_key_cache", + cache_config=cache_config, + vllm_config=vllm_config, + ) + + @property + def output_width(self) -> int: + return self.token_topk + self.compress_ratio - 1 + + def _metadata( + self, + ) -> tuple[QSAForwardMetadata, QSAForwardMetadata] | None: + metadata = get_forward_context().attn_metadata + if isinstance(metadata, list): + metadata = metadata[0] + if not isinstance(metadata, dict): + return None + raw = cast(QSAForwardMetadata, metadata[self.raw_key_cache.prefix]) + compressed = cast( + QSAForwardMetadata, metadata[self.compressed_key_cache.prefix] + ) + if raw.num_actual_tokens != compressed.num_actual_tokens: + raise RuntimeError("QSA side-cache metadata token counts disagree") + if not raw.logical_positions.is_cuda and ( + not torch.equal(raw.logical_positions, compressed.logical_positions) + ): + raise RuntimeError("QSA side-cache metadata positions disagree") + return raw, compressed + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return fixed-width request-relative token indices padded with ``-1``.""" + + metadata = self._metadata() + if metadata is None: + # Preserve step-0 indices when later MTP steps reuse the buffer. + if self.skip_topk and out is not None: + return out + result = torch.full( + (hidden_states.shape[0], self.output_width), + -1, + dtype=torch.int32, + device=hidden_states.device, + ) + if out is not None: + out.copy_(result) + return out + return result + + from .ops.qsa import ( + qsa_compress_groups_with_ratio, + qsa_select_paged_tokens, + qsa_store_cache_rows, + ) + + raw_metadata, compressed_metadata = metadata + num_tokens = raw_metadata.num_actual_tokens + hidden_states = hidden_states[:num_tokens] + positions = positions[..., :num_tokens] + + # Q/K projection + projected_qk, _ = self.index_qk_proj(hidden_states) + projected_q, raw_keys = projected_qk.split( + ( + self.index_n_heads * self.index_head_dim, + self.index_kv_heads * self.index_head_dim, + ), + dim=-1, + ) + raw_key_state_cache = self.raw_key_cache + compressed_key_cache = self.compressed_key_cache.kv_cache + + if self.use_fused_pre_indexer: + q = projected_q.new_empty( + num_tokens, + self.index_n_heads, + self.index_head_dim, + ) + qsa_pre_indexer( + projected_q, + raw_keys, + positions, + self.rotary_emb.cos_sin_cache, + self.q_layernorm.weight, + self.k_layernorm.weight, + self.q_layernorm.variance_epsilon, + q, + raw_key_state_cache.kv_cache, + raw_metadata.slot_mapping, + raw_metadata.block_table, + raw_metadata.query_start_loc, + raw_metadata.logical_positions, + compressed_key_cache, + compressed_metadata.slot_mapping, + compressed_metadata.k_work_metadata, + compress_ratio=self.compress_ratio, + mrope_section=getattr(self.rotary_emb, "mrope_section", None), + rope_pos_offset=( + raw_key_state_cache.rope_position_offset + if raw_key_state_cache.rope_position_cache is not None + else None + ), + ) + else: + # Unfused reference path + from flashinfer.norm import gemma_rmsnorm + + q = projected_q.reshape(-1, self.index_n_heads, self.index_head_dim) + q = gemma_rmsnorm( + q.reshape(-1, self.index_head_dim), + self.q_layernorm.weight, + self.q_layernorm.variance_epsilon, + ).reshape_as(q) + q = apply_qsa_rope(self.rotary_emb, positions, q) + + raw_key_cache = raw_key_state_cache.key_cache + rope_position_cache = raw_key_state_cache.rope_position_cache + if rope_position_cache is None: + position_rows = raw_metadata.logical_positions.view(-1, 1, 1).expand( + -1, 1, 3 + ) + else: + position_rows = canonical_qsa_rope_positions(positions).to( + device=raw_key_cache.device + ) + pooled, first_positions = qsa_compress_groups_with_ratio( + raw_keys.reshape(-1, 1, self.index_head_dim), + position_rows, + raw_key_cache, + raw_metadata.block_table, + raw_metadata.token_to_req, + raw_metadata.query_start_loc, + raw_metadata.logical_positions, + compressed_metadata.slot_mapping, + self.compress_ratio, + rope_position_cache, + ) + compressed_keys = gemma_rmsnorm( + pooled.reshape(-1, self.index_head_dim), + self.k_layernorm.weight, + self.k_layernorm.variance_epsilon, + ).reshape(-1, 1, self.index_head_dim) + if getattr(self.rotary_emb, "mrope_section", None): + first_positions = first_positions.transpose(0, 1) + else: + first_positions = first_positions[:, 0] + compressed_keys = apply_qsa_rope( + self.rotary_emb, + first_positions, + compressed_keys, + ) + qsa_store_cache_rows( + compressed_key_cache, + compressed_metadata.slot_mapping, + compressed_keys, + ) + qsa_store_cache_rows( + raw_key_cache, + raw_metadata.slot_mapping, + raw_keys, + ) + if rope_position_cache is not None: + qsa_store_cache_rows( + rope_position_cache, + raw_metadata.slot_mapping, + position_rows, + ) + + if self.skip_topk: + if out is None: + raise RuntimeError("QSA top-k reuse requires an output buffer") + return out + + # Score compressed keys, select blocks, then expand them to token indices. + return qsa_select_paged_tokens( + q, + compressed_key_cache, + compressed_metadata.block_table, + compressed_metadata.token_to_req, + compressed_metadata.logical_positions, + compressed_metadata.seq_lens, + self.token_topk, + self.compress_ratio, + out, + ) + + +__all__ = ["QSAIndexer", "apply_qsa_rope"] diff --git a/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py new file mode 100644 index 0000000000..43b7a76bc5 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/low_latency_gemm.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp decode GEMM selection on Blackwell. + +Dispatch follows Kimi-K3 and uses the local ``(N, K)`` shape and token count. +Plans contain measured CUDA graph capture sizes; other token counts use the +standard linear implementation. +""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( + SkinnyGemmConfig, + shape_dynamic_skinny_gemm, +) +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + UnquantizedEmbeddingMethod, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = { + # GDN fused QKVZ projection, TP=4. + (4096, 2560): { + 1: SkinnyGemmConfig(1, 64, 4, k_unroll=4), + 2: SkinnyGemmConfig(2, 64, 4, k_unroll=4), + }, + # GDN and QSA output projections, TP=4. + (2560, 1536): { + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=2, vector_width=4), + 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2), + }, + # GDN fused B/A projection, TP=4. + (24, 2560): { + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4), + 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4), + 16: SkinnyGemmConfig(16, 128, 1, k_unroll=4, vector_width=4), + }, + # QSA fused QKV/gate projection, TP=4. + (3584, 2560): { + 1: SkinnyGemmConfig(1, 128, 4, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 64, 2, k_unroll=2), + 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2), + }, + # QSA indexer Q/K projection, replicated in a TP=4 deployment. + (640, 2560): { + 1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4), + 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4), + }, + # Shared-expert fused gate/up projection, TP=4. + (320, 2560): { + 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4), + 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4), + 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4), + 8: SkinnyGemmConfig(8, 64, 1, k_unroll=4), + 16: SkinnyGemmConfig(16, 128, 2, k_unroll=4, vector_width=4), + }, + # LM head, TP=4. + (62080, 2560): { + 1: SkinnyGemmConfig(1, 64, 4, k_unroll=2), + 2: SkinnyGemmConfig(2, 32, 4, k_unroll=2), + }, + # HC merged down/injection projection, replicated in a TP=4 deployment. + (336, 10240): { + 1: SkinnyGemmConfig(1, 128, 1, static_k=10240), + 2: SkinnyGemmConfig(2, 128, 1, static_k=10240), + 4: SkinnyGemmConfig(4, 128, 2, static_k=10240), + 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4), + }, +} + + +def _is_sm103() -> bool: + return current_platform.is_device_capability((10, 3)) + + +def _is_packed_row_major(tensor: torch.Tensor) -> bool: + return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1) + + +def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + not envs.VLLM_BATCH_INVARIANT + and _is_packed_row_major(x) + and _is_packed_row_major(weight) + and x.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and x.is_cuda + and weight.is_cuda + and x.device == weight.device + and x.shape[1] == weight.shape[1] + ) + + +class _Qwen4ExpLowLatencyApply: + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if bias is None and not envs.VLLM_BATCH_INVARIANT: + return torch.ops.vllm.qwen4_exp_low_latency_gemm(x, layer.weight) + return super().apply(layer, x, bias) # type: ignore[misc] + + +class Qwen4ExpLowLatencyLinearMethod(_Qwen4ExpLowLatencyApply, UnquantizedLinearMethod): + pass + + +class Qwen4ExpLowLatencyEmbeddingMethod( + _Qwen4ExpLowLatencyApply, UnquantizedEmbeddingMethod +): + pass + + +def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1])) + config = None if plan is None else plan.get(x.shape[0]) + if ( + config is not None + and _runtime_ok(x, weight) + and shape_dynamic_skinny_gemm.is_available() + ): + return shape_dynamic_skinny_gemm(x, weight, config) + return torch.nn.functional.linear(x, weight) + + +def _qwen4_exp_low_latency_gemm_fake( + x: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + return x.new_empty((*x.shape[:-1], weight.shape[0])) + + +direct_register_custom_op( + op_name="qwen4_exp_low_latency_gemm", + op_func=_qwen4_exp_low_latency_gemm, + fake_impl=_qwen4_exp_low_latency_gemm_fake, +) + + +def enable_qwen4_exp_low_latency_gemm( + module: nn.Module, + dtype: torch.dtype, +) -> None: + if dtype != torch.bfloat16 or not _is_sm103(): + return + if not shape_dynamic_skinny_gemm.is_available(): + return + + warmup_configs: set[SkinnyGemmConfig] = set() + for child in module.modules(): + is_linear = ( + isinstance(child, LinearBase) + and type(child.quant_method) is UnquantizedLinearMethod + ) + is_head = ( + isinstance(child, ParallelLMHead) + and type(child.quant_method) is UnquantizedEmbeddingMethod + ) + if not (is_linear or is_head): + continue + weight = getattr(child, "weight", None) + if weight is None or weight.dim() != 2: + continue + plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1])) + if plan is None: + continue + if is_linear: + child.quant_method = Qwen4ExpLowLatencyLinearMethod() + else: + child.quant_method = Qwen4ExpLowLatencyEmbeddingMethod() + warmup_configs.update(plan.values()) + + if warmup_configs: + shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py new file mode 100644 index 0000000000..5f32b8ee6e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -0,0 +1,1070 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen4Exp model.""" + +from collections.abc import Iterable +from itertools import islice + +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, +) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, + MambaStateCopyFuncsByType, + MambaStateDtypeCalculator, + MambaStateShapeCalculator, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import ( + HasInnerState, + IsHybrid, + MixtureOfExperts, + MultiModalEmbeddings, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, + _require_is_multimodal, +) +from vllm.model_executor.models.qwen3_5 import ( + Qwen3_5ForConditionalGeneration, +) +from vllm.model_executor.models.qwen3_next import ( + Qwen3NextAttention, + Qwen3NextMLP, + Qwen3NextSparseMoeBlock, +) +from vllm.model_executor.models.qwen3_vl import ( + Qwen3_VisionTransformer, + Qwen3VLDummyInputsBuilder, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + StageMissingLayer, + WeightsMapper, + _merge_multimodal_embeddings, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_fuse_shared_experts, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sequence import IntermediateTensors +from vllm.tokenizers.registry import cached_tokenizer_from_config +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import MambaSpec + +from ..config import Qwen4ExpConfig +from .hyperconnection import GatedResidual, HyperConnectionConfig + +try: + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +except ModuleNotFoundError as exc: + # The Blackwell-only CuTe DSL helper is absent from the 1Cat SM70 tree. + if exc.name != "vllm.model_executor.kernels.linear.cute_dsl": + raise + + def enable_qwen4_exp_low_latency_gemm( + module: nn.Module, dtype: torch.dtype + ) -> None: + del module, dtype + + +from .ple_layer import Qwen4ExpPLELayer +from .qsa import Qwen4ExpQSAAttention + + +def without_modelopt_fp4( + quant_config: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Return ``None`` for weights excluded from Qwen4Exp ModelOpt-FP4.""" + + if quant_config is not None and quant_config.get_name() == "modelopt_fp4": + return None + return quant_config + + +def _remap_qsa_cache_scale_name( + name: str, + qsa_layer_ids: frozenset[int], +) -> str: + """Map serialized main-cache scales onto the merged QSA owner. + + Regular attention keeps cache scales below its ``attn`` child. QSA owns + that cache directly, so only QSA layers need the final path component + moved to the owner's persistent ``_k_scale``/``_v_scale`` buffers. + """ + + scale_suffixes = { + "k_proj.k_scale": "_k_scale", + "k_proj.output_scale": "_k_scale", + "attn.k_scale": "_k_scale", + "attn._k_scale": "_k_scale", + "k_scale": "_k_scale", + "_k_scale": "_k_scale", + "v_proj.v_scale": "_v_scale", + "v_proj.output_scale": "_v_scale", + "attn.v_scale": "_v_scale", + "attn._v_scale": "_v_scale", + "v_scale": "_v_scale", + "_v_scale": "_v_scale", + } + for layer_id in qsa_layer_ids: + marker = f"layers.{layer_id}.self_attn." + marker_start = name.find(marker) + if marker_start < 0 or (marker_start > 0 and name[marker_start - 1] != "."): + continue + suffix = name[marker_start + len(marker) :] + mapped_suffix = scale_suffixes.get(suffix) + if mapped_suffix is not None: + return f"{name[: marker_start + len(marker)]}{mapped_suffix}" + return name + + +_QWEN4_EXP_IGNORED_MISSING_SUFFIXES = [ + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + "_weight_scale", + "_input_scale", +] + +# The checkpoint keeps down and injection projections separate; runtime packs +# them into adjacent logical shards of one MergedColumnParallelLinear. +_HC_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_stacked={ + "hyper_connection.input_mix_weight_down.weight": ( + "hyper_connection.input_mix_weight_down_block_inject.weight", + 0, + ), + "hyper_connection.block_inject_weight.weight": ( + "hyper_connection.input_mix_weight_down_block_inject.weight", + 1, + ), + } +) + +_QWEN3_5_WEIGHTS_MAPPER = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".mlp.gate_proj": (".mlp.gate_up_proj", 0), + ".mlp.up_proj": (".mlp.gate_up_proj", 1), + ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0), + ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1), + ".in_proj_qkv": (".in_proj_qkvz", (0, 1, 2)), + ".in_proj_z": (".in_proj_qkvz", 3), + ".in_proj_b": (".in_proj_ba", 0), + ".in_proj_a": (".in_proj_ba", 1), + } +) + + +class Qwen4ExpSparseMoeBlock(Qwen3NextSparseMoeBlock): + """Qwen3Next MoE with Qwen4Exp HC validation.""" + + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + parallel_config = vllm_config.parallel_config + if parallel_config.use_sequence_parallel_moe: + raise NotImplementedError( + "Qwen4Exp HC does not support sequence-parallel MoE" + ) + super().__init__(vllm_config=vllm_config, prefix=prefix) + config = vllm_config.model_config.hf_text_config + self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) + + +class Qwen4ExpDecoderLayer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + layer_type: str, + prefix: str = "", + ) -> None: + super().__init__() + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + self.config = config + self.layer_type = layer_type + self.layer_idx = extract_layer_index(prefix) + if vllm_config.parallel_config.use_sequence_parallel_moe: + raise NotImplementedError( + "Qwen4Exp HC does not support sequence-parallel MoE" + ) + self.ple: Qwen4ExpPLELayer | None = None + ple_layer_ids = config.ple_layer_ids + if (self.layer_idx + 1) in ple_layer_ids: + ple_layer_ids_sorted = sorted(set(ple_layer_ids)) + ple_dense_layer_id_map = { + abs_id: idx for idx, abs_id in enumerate(ple_layer_ids_sorted) + } + ple_dense_layer_id = ple_dense_layer_id_map[self.layer_idx + 1] + self.ple = Qwen4ExpPLELayer( + config, + vllm_config=vllm_config, + layer_idx=self.layer_idx, + ple_dense_layer_id=ple_dense_layer_id, + prefix=f"{prefix}.ple", + ) + + if layer_type == "linear_attention": + self.linear_attn = QwenGatedDeltaNetAttention( + config, + vllm_config=vllm_config, + prefix=f"{prefix}.linear_attn", + gqa_interleaved_layout=False, + ) + elif layer_type == "full_attention": + use_qsa = getattr(config, "indexer_n_heads", None) is not None + if not use_qsa: + self.self_attn = Qwen3NextAttention( + config, + model_config=model_config, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + else: + self.self_attn = Qwen4ExpQSAAttention( + vllm_config=vllm_config, + config=config, + layer_id=self.layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + else: + raise ValueError(f"Invalid layer_type {layer_type}") + + mlp_only_layers = getattr(config, "mlp_only_layers", []) + num_experts = getattr(config, "num_experts", 0) or 0 + absolute_layer_id = self.layer_idx + 1 + is_moe_layer = self.layer_idx not in mlp_only_layers and ( + num_experts > 0 and absolute_layer_id % config.decoder_sparse_step == 0 + ) + if is_moe_layer: + self.mlp = Qwen4ExpSparseMoeBlock( + vllm_config=vllm_config, prefix=f"{prefix}.mlp" + ) + else: + self.mlp = Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + hc_config = HyperConnectionConfig( + hc_count=config.hc_count, + hidden_size=config.hidden_size, + params_dtype=model_config.dtype, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.attn_hyper_connection = GatedResidual( + hc_config, + prefix=maybe_prefix(prefix, "attn_hyper_connection"), + ) + self.mlp_hyper_connection = GatedResidual( + hc_config, + prefix=maybe_prefix(prefix, "mlp_hyper_connection"), + ) + + def forward( + self, + hidden_states: torch.Tensor, + prev_block_output: torch.Tensor | None, + prev_injection: torch.Tensor | None, + positions: torch.Tensor, + *, + input_ids: torch.Tensor | None, + query_start_loc: torch.Tensor | None, + ngram_context: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + attn_hc = self.attn_hyper_connection + if self.ple is not None: + # PLE adds directly to the multi-stream state, so pending HC state + # must be materialized before the addition. + if prev_block_output is not None and prev_injection is not None: + hidden_states = attn_hc.combine( + hidden_states, prev_block_output, prev_injection + ) + prev_block_output = prev_injection = None + + if input_ids is None or query_start_loc is None or ngram_context is None: + raise RuntimeError("PLE inputs were not prepared") + hidden_states = hidden_states + self.ple( + hidden_states, + input_ids, + query_start_loc, + ngram_context, + ) + + # Fuse a pending combine with this HC module's mix when possible. + if prev_block_output is not None and prev_injection is not None: + hidden_states, block_input, injection = attn_hc.combine_and_mix( + hidden_states, prev_block_output, prev_injection + ) + else: + hidden_states, block_input, injection = attn_hc.mix(hidden_states) + + if self.layer_type == "linear_attention": + attn_out = self.linear_attn(hidden_states=block_input) + elif self.layer_type == "full_attention": + attn_out = self.self_attn( + hidden_states=block_input, + positions=positions, + ) + else: + raise ValueError("Invalid layer_type") + + mlp_hc = self.mlp_hyper_connection + hidden_states, block_input, injection = mlp_hc.combine_and_mix( + hidden_states, attn_out, injection + ) + mlp_out = self.mlp(block_input) + return hidden_states, mlp_out, injection + + +class Qwen4ExpMixtureOfExperts(MixtureOfExperts): + """Expose Qwen4Exp routed experts through vLLM's EPLB protocol.""" + + def set_moe_parameters(self, layers: Iterable[nn.Module]) -> None: + self.moe_layers = [] + self.moe_mlp_layers = [] + example_moe = None + for layer in layers: + if isinstance(layer, Qwen4ExpDecoderLayer) and isinstance( + layer.mlp, Qwen4ExpSparseMoeBlock + ): + example_moe = layer.mlp + self.moe_mlp_layers.append(layer.mlp) + self.moe_layers.append(layer.mlp.experts) + + self.num_moe_layers = len(self.moe_layers) + if example_moe is None: + self.num_expert_groups = 0 + self.num_shared_experts = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_redundant_experts = 0 + return + + self.num_expert_groups = 1 + self.num_shared_experts = example_moe.n_shared_experts + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for moe in self.moe_mlp_layers: + moe.n_physical_experts = num_physical_experts + moe.n_local_physical_experts = num_local_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "query_start_loc": 0, + "ngram_context": 0, + "deepstack_input_embeds": 0, + } +) +class Qwen4ExpModel(nn.Module): + hf_to_vllm_mapper = _QWEN3_5_WEIGHTS_MAPPER | _HC_WEIGHTS_MAPPER + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.config = config + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + self.vocab_size = config.vocab_size + self._qsa_layer_ids = frozenset( + layer_idx + for layer_idx, layer_type in enumerate(config.layer_types) + if layer_type == "full_attention" + and getattr(config, "indexer_n_heads", None) is not None + ) + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, config.hidden_size) + + def get_layer(prefix: str) -> Qwen4ExpDecoderLayer: + layer_idx = extract_layer_index(prefix) + return Qwen4ExpDecoderLayer( + vllm_config, + layer_type=config.layer_types[layer_idx], + prefix=prefix, + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + intermediate_size = config.hidden_size * config.hc_count + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], intermediate_size + ) + + self.hyper_connection_mixer: GatedResidual | None + if get_pp_group().is_last_rank: + hc_config = HyperConnectionConfig( + hc_count=config.hc_count, + hidden_size=config.hidden_size, + params_dtype=vllm_config.model_config.dtype, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.hyper_connection_mixer = GatedResidual( + hc_config, + use_combine=False, + prefix=maybe_prefix(prefix, "hyper_connection_mixer"), + ) + else: + self.hyper_connection_mixer = None + + spec_config = vllm_config.speculative_config + # MTP HC multi-stream outputs: when speculative method=="mtp" and the + # model uses HC with hc_count>1, retain the pre-final-mixer multi-stream + # hidden state [T, hc_count*H] so the MTP drafter can feed a real + # multi-stream backbone hidden on its first step (scheme A). Derived + # purely from config (NOT node identity) so P/D nodes stay consistent. + needs_mtp_hidden = ( + spec_config is not None + and getattr(spec_config, "method", None) == "mtp" + and get_pp_group().is_last_rank + ) + if needs_mtp_hidden: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.hc_count * config.hidden_size, + dtype=vllm_config.model_config.dtype, + ) + else: + self._mtp_hidden_buffer = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + ngram_context: torch.Tensor | None = None, + deepstack_input_embeds: IntermediateTensors | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + if input_ids is None: + raise ValueError("input_ids or inputs_embeds is required") + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.repeat(1, self.config.hc_count) + else: + if intermediate_tensors is None: + raise ValueError("pipeline stage requires intermediate tensors") + hidden_states = intermediate_tensors["hidden_states"] + + block_output = None + injection = None + last_layer = None + for layer_idx, layer in islice( + enumerate(self.layers), self.start_layer, self.end_layer + ): + last_layer = layer + hidden_states, block_output, injection = layer( + hidden_states=hidden_states, + prev_block_output=block_output, + prev_injection=injection, + positions=positions, + input_ids=input_ids, + query_start_loc=query_start_loc, + ngram_context=ngram_context, + ) + if deepstack_input_embeds is not None and layer_idx < len( + deepstack_input_embeds + ): + deepstack_embed = deepstack_input_embeds[ + f"deepstack_input_embeds_{layer_idx}" + ] + deepstack_embed = ( + deepstack_embed.unsqueeze(-2) + .expand( + *deepstack_embed.shape[:-1], + self.config.hc_count, + self.config.hidden_size, + ) + .flatten(-2) + ) + # Deepstack is an external addition to the materialized + # multi-stream state and therefore terminates delayed combine. + hidden_states = layer.mlp_hyper_connection.combine( + hidden_states, block_output, injection + ) + block_output = None + injection = None + hidden_states = hidden_states + deepstack_embed + + if not get_pp_group().is_last_rank: + # PP transports one tensor, not the delayed HC tuple. Materialize + # with the HC module that produced the pending injection. + if last_layer is not None and block_output is not None: + hidden_states = last_layer.mlp_hyper_connection.combine( + hidden_states, block_output, injection + ) + return IntermediateTensors({"hidden_states": hidden_states}) + + # The final mixer consumes the last pending combine and returns both + # the sampled single stream and the materialized multi-stream state. + final_mixer = self.hyper_connection_mixer + assert final_mixer is not None + multi_hidden, sample_hidden_states, _ = final_mixer.combine_and_mix( + hidden_states, block_output, injection + ) + if self._mtp_hidden_buffer is not None: + # Capture the pre-final-mixer multi-stream hidden state + # [T, hc_count*H] for the MTP drafter (zero extra compute: + # this tensor is needed by the final mixer regardless). + num_tokens = multi_hidden.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(multi_hidden) + return sample_hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = ( + ( + _remap_qsa_cache_scale_name(name, self._qsa_layer_ids), + weight, + ) + for name, weight in weights + ) + weights = maybe_fuse_shared_experts( + weights, + n_routed_experts=getattr(self.config, "num_experts", 0) or 0, + n_shared_experts=1, + ckpt_prefix="mlp.shared_expert", + ) + # Non-persistent PLE state rebuilt in __init__; skip any ckpt + # column for them. + skip_substrs = [ + "hashstats_", + "token_lookup", + "hyper_connection_mixer.block_inject_weight", + ] + loader = AutoWeightsLoader( + self, + skip_substrs=skip_substrs, + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + loaded = loader.load_weights( + weights, + mapper=self.hf_to_vllm_mapper, + ) + return loaded + + +class Qwen4ExpForCausalLM( + nn.Module, + HasInnerState, + SupportsLoRA, + SupportsMRoPE, + SupportsPP, + Qwen4ExpMixtureOfExperts, + IsHybrid, +): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + "input_mix_weight_down_block_inject": [ + "input_mix_weight_down", + "block_inject_weight", + "_input_mix_padding", + ], + } + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"model.language_model.": "model."} + ) + requires_raw_input_tokens = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.quant_config = vllm_config.quant_config + self.config = config + self.scheduler_config = vllm_config.scheduler_config + if vllm_config.cache_config.mamba_cache_mode == "all": + raise NotImplementedError( + "Qwen4Exp currently does not support 'all' prefix caching, " + "please use '--mamba-cache-mode=align' instead" + ) + self.model = Qwen4ExpModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + self.set_moe_parameters(self.model.layers) + enable_qwen4_exp_low_latency_gemm(self, self.model_config.dtype) + + @staticmethod + def get_model_state_cls(): + from .model_state import Qwen4ExpModelState + + return Qwen4ExpModelState + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + # Forward kwargs unchanged so the runner's _maybe_add_ngram_kwargs + # path (query_start_loc / ngram_context) reaches Qwen4ExpModel. + return self.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + **kwargs, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + @classmethod + def get_ple_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[torch.dtype, ...]: + return MambaStateDtypeCalculator.short_conv_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + ) + + @classmethod + def get_ple_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int]]: + hf_config = vllm_config.model_config.hf_text_config + conv_kernel_size = hf_config.ple_conv_kernel_size + short_conv_dilation = hf_config.ngram_size + conv_state_len = (conv_kernel_size - 1) * short_conv_dilation + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + hc_count = hf_config.hc_count + hc_hidden_size = hf_config.hidden_size * hc_count + return MambaStateShapeCalculator.short_conv_state_shape( + tp_world_size=1, + intermediate_size=hc_hidden_size, + conv_kernel=conv_state_len + 1, + num_spec=num_spec, + ) + + @classmethod + def get_gdn_mamba_state_dtype_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[torch.dtype, torch.dtype]: + return MambaStateDtypeCalculator.gated_delta_net_state_dtype( + vllm_config.model_config.dtype, + vllm_config.cache_config.mamba_cache_dtype, + vllm_config.cache_config.mamba_ssm_cache_dtype, + ) + + @classmethod + def get_gdn_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int], tuple[int, int]]: + parallel_config = vllm_config.parallel_config + hf_config = vllm_config.model_config.hf_text_config + tp_size = parallel_config.tensor_parallel_size + num_spec = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config + else 0 + ) + return MambaStateShapeCalculator.gated_delta_net_state_shape( + tp_size, + hf_config.linear_num_key_heads, + hf_config.linear_num_value_heads, + hf_config.linear_key_head_dim, + hf_config.linear_value_head_dim, + hf_config.linear_conv_kernel_dim, + num_spec, + ) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[torch.dtype, torch.dtype]: + return cls.get_gdn_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[tuple[int, int], tuple[int, int]]: + return cls.get_gdn_mamba_state_shape_from_config(vllm_config) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func() + + @classmethod + def get_mamba_state_copy_funcs( + cls, + mamba_types: set[MambaAttentionBackendEnum], + ) -> MambaStateCopyFuncsByType: + copy_funcs_by_type = { + MambaAttentionBackendEnum.GDN_ATTN: cls.get_mamba_state_copy_func(), + MambaAttentionBackendEnum.SHORT_CONV: ( + MambaStateCopyFuncCalculator.short_conv_state_copy_func() + ), + } + missing_types = mamba_types - copy_funcs_by_type.keys() + assert not missing_types, f"missing state copy funcs for {missing_types}" + return { + mamba_type: copy_funcs_by_type[mamba_type] for mamba_type in mamba_types + } + + @classmethod + def get_mamba_specs_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[MambaSpec, ...]: + """Return all MambaSpecs for this model (GDN layers + PLE layer). + + The PLE layer uses a separate short_conv MambaSpec whose page_size_bytes + may exceed the GDN spec; callers should take the maximum. + """ + return ( + MambaSpec( + shapes=cls.get_gdn_mamba_state_shape_from_config(vllm_config), + dtypes=cls.get_gdn_mamba_state_dtype_from_config(vllm_config), + block_size=-1, + ), + MambaSpec( + shapes=cls.get_ple_mamba_state_shape_from_config(vllm_config), + dtypes=cls.get_ple_mamba_state_dtype_from_config(vllm_config), + block_size=-1, + tp_replicated=True, + ), + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + return self.model._mtp_hidden_buffer + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + positions = torch.arange(len(input_tokens), dtype=torch.long) + return positions.unsqueeze(0).expand(3, -1), 0 + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_substrs=["mtp."], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +class Qwen4ExpProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self) -> Qwen4ExpConfig: + return self.ctx.get_hf_config(Qwen4ExpConfig) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen4ExpProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class Qwen4ExpForConditionalGeneration( + Qwen3_5ForConditionalGeneration, + HasInnerState, + Qwen4ExpMixtureOfExperts, +): + """Qwen3-VL vision tower backed by the Qwen4Exp language model.""" + + requires_raw_input_tokens = True + + packed_modules_mapping = Qwen3_5ForConditionalGeneration.packed_modules_mapping | { + "input_mix_weight_down_block_inject": [ + "input_mix_weight_down", + "block_inject_weight", + "_input_mix_padding", + ] + } + + @staticmethod + def get_model_state_cls(): + from .model_state import Qwen4ExpModelState + + return Qwen4ExpModelState + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None: + nn.Module.__init__(self) + config: Qwen4ExpConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + if multimodal_config is None: + raise ValueError( + "Qwen4ExpForConditionalGeneration requires multimodal_config" + ) + + self.config = config + self.model_config = vllm_config.model_config + self.multimodal_config = multimodal_config + self.language_model_only = multimodal_config.language_model_only + if self.language_model_only: + self.use_data_parallel = False + self.is_multimodal_pruning_enabled = False + self.video_pruning_method = None + self.video_pruning_rate = 0.0 + self._tokenizer = None + self.visual = StageMissingLayer("vision_tower") + self._tower_model_names = [] + else: + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" + self._init_video_pruning(multimodal_config) + self._tokenizer = cached_tokenizer_from_config(vllm_config.model_config) + + with self._mark_tower_model(vllm_config, {"image", "video"}): + self.visual = Qwen3_VisionTransformer( + config.vision_config, + norm_eps=config.text_config.rms_norm_eps, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + + self.use_deepstack = ( + not self.language_model_only + and bool(config.vision_config.deepstack_visual_indexes) + and not isinstance(self.visual, StageMissingLayer) + ) + self.deepstack_num_level = ( + len(config.vision_config.deepstack_visual_indexes) + if self.use_deepstack + else 0 + ) + self.visual_dim = config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + if self.use_deepstack: + self.deepstack_input_embeds = [ + torch.zeros( + vllm_config.scheduler_config.max_num_batched_tokens, + config.text_config.hidden_size, + ) + for _ in range(self.deepstack_num_level) + ] + self.deepstack_input_embeds_num_tokens = 0 + + with self._mark_language_model(vllm_config): + self.language_model = Qwen4ExpForCausalLM( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + if not get_pp_group().is_first_rank and self.use_deepstack: + assert self.language_model.model.start_layer >= len( + config.vision_config.deepstack_visual_indexes + ), ( + "start_layer should be greater than or equal to " + "len(deepstack_visual_indexes)" + ) + self.set_moe_parameters(self.language_model.model.layers) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.language_model.embed_input_ids, + is_multimodal=is_multimodal, + ) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + if self.language_model_only: + raise ValueError( + "Qwen4Exp language_model_only does not accept multimodal embeddings" + ) + + is_multimodal = _require_is_multimodal(is_multimodal) + if self.use_deepstack: + deepstack_input_embeds, multimodal_embeddings = ( + self._compute_deepstack_embeds( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + ) + else: + deepstack_input_embeds = None + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + if deepstack_input_embeds is not None: + self._set_deepstack_input_embeds(deepstack_input_embeds) + return inputs_embeds + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + return self.language_model.get_mtp_target_hidden_states() + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + if inputs_embeds is not None and get_pp_group().is_first_rank: + deepstack_input_embeds = self._get_deepstack_input_embeds( + inputs_embeds.size(0) + ) + else: + deepstack_input_embeds = None + + hidden_states = self.language_model.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + query_start_loc=kwargs.get("query_start_loc"), + ngram_context=kwargs.get("ngram_context"), + deepstack_input_embeds=deepstack_input_embeds, + ) + if inputs_embeds is not None and get_pp_group().is_first_rank: + self._clear_deepstack_input_embeds(inputs_embeds.size(0)) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=["visual."] if self.language_model_only else None, + skip_substrs=["mtp."], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + @classmethod + def get_mamba_state_dtype_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[torch.dtype, torch.dtype]: + return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config) + + @classmethod + def get_mamba_state_shape_from_config( + cls, + vllm_config: VllmConfig, + ) -> tuple[tuple[int, int], tuple[int, int]]: + return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config) + + @classmethod + def get_mamba_state_copy_func( + cls, + ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: + return Qwen4ExpForCausalLM.get_mamba_state_copy_func() + + @classmethod + def get_mamba_state_copy_funcs( + cls, + mamba_types: set[MambaAttentionBackendEnum], + ) -> MambaStateCopyFuncsByType: + return Qwen4ExpForCausalLM.get_mamba_state_copy_funcs(mamba_types) + + @classmethod + def get_mamba_specs_from_config( + cls, vllm_config: VllmConfig + ) -> tuple[MambaSpec, ...]: + return Qwen4ExpForCausalLM.get_mamba_specs_from_config(vllm_config) + + +__all__ = [ + "Qwen4ExpDecoderLayer", + "Qwen4ExpForCausalLM", + "Qwen4ExpForConditionalGeneration", + "Qwen4ExpMixtureOfExperts", + "Qwen4ExpModel", + "Qwen4ExpSparseMoeBlock", +] diff --git a/vllm/models/qwen4_exp/nvidia/model_state.py b/vllm/models/qwen4_exp/nvidia/model_state.py new file mode 100644 index 0000000000..ec2a0fef8b --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/model_state.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-runner state for Qwen4Exp PLE inputs.""" + +from typing import Any + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache +from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState +from vllm.v1.worker.gpu.states import RequestState + + +class Qwen4ExpModelState(MambaHybridModelState): + """Add rollback-safe PLE n-gram context to the model inputs.""" + + def __init__( + self, + vllm_config: VllmConfig, + model: nn.Module, + encoder_cache: EncoderCache | None, + device: torch.device, + ) -> None: + super().__init__(vllm_config, model, encoder_cache, device) + config = self.model_config.hf_text_config + self.uses_ngram_embedding = bool(config.ple_layer_ids) + if not self.uses_ngram_embedding: + self.ngram_context_len = 0 + self.ngram_eos_token_id = 0 + return + + if vllm_config.parallel_config.pipeline_parallel_size > 1: + raise RuntimeError( + "N-gram PLE embedding currently requires " + "pipeline_parallel_size=1 because non-first pipeline ranks do " + "not receive the raw input_ids required by PLE. Please run " + "with PP=1." + ) + + self.ngram_context_len = int(config.ngram_size) - 1 + if self.ngram_context_len <= 0: + raise ValueError("N-gram embedding requires context length >= 1.") + self.ngram_eos_token_id = int(config.eos_token_id) + self.ngram_context = torch.full( + (self.max_num_reqs, self.ngram_context_len), + self.ngram_eos_token_id, + dtype=torch.int32, + device=self.device, + ) + self.ngram_context_offsets = torch.arange( + -self.ngram_context_len, + 0, + dtype=torch.int64, + device=self.device, + ) + self.ple_query_start_loc = torch.zeros( + self.max_num_reqs + 1, + dtype=torch.int32, + device=self.device, + ) + + def _prepare_ngram_context( + self, + input_batch: InputBatch, + req_states: RequestState, + ) -> torch.Tensor: + num_reqs = input_batch.num_reqs + num_reqs_padded = input_batch.num_reqs_after_padding + context = self.ngram_context[:num_reqs_padded] + context.fill_(self.ngram_eos_token_id) + if num_reqs == 0: + return context + + request_indices = input_batch.idx_mapping[:num_reqs].long() + context_end = req_states.num_computed_tokens.gpu[request_indices].long() + token_indices = context_end.unsqueeze(1) + self.ngram_context_offsets + valid_tokens = token_indices >= 0 + token_indices.clamp_min_(0) + context_tokens = req_states.all_token_ids.gpu[ + request_indices.unsqueeze(1), token_indices + ] + context[:num_reqs].copy_( + torch.where( + valid_tokens, + context_tokens, + context_tokens.new_full((), self.ngram_eos_token_id), + ) + ) + return context + + def prepare_inputs( + self, + input_batch: InputBatch, + req_states: RequestState, + ) -> dict[str, Any]: + model_inputs = super().prepare_inputs(input_batch, req_states) + if not self.uses_ngram_embedding: + return model_inputs + + num_reqs_padded = input_batch.num_reqs_after_padding + query_start_loc = self.ple_query_start_loc[: num_reqs_padded + 1] + query_start_loc.copy_(input_batch.query_start_loc[: num_reqs_padded + 1]) + model_inputs.update( + query_start_loc=query_start_loc, + ngram_context=self._prepare_ngram_context(input_batch, req_states), + ) + return model_inputs + + def prepare_dummy_inputs( + self, + num_reqs: int, + num_tokens: int, + ) -> dict[str, Any]: + model_inputs = super().prepare_dummy_inputs(num_reqs, num_tokens) + if not self.uses_ngram_embedding: + return model_inputs + + query_start_loc = self.ple_query_start_loc[: num_reqs + 1] + query_start_loc[0] = 0 + tokens_per_req, num_extra_tokens = divmod(num_tokens, num_reqs) + query_lens = torch.full( + (num_reqs,), + tokens_per_req, + dtype=query_start_loc.dtype, + device=query_start_loc.device, + ) + if num_extra_tokens > 0: + query_lens[-num_extra_tokens:] += 1 + torch.cumsum(query_lens, dim=0, out=query_start_loc[1:]) + + ngram_context = self.ngram_context[:num_reqs] + ngram_context.fill_(self.ngram_eos_token_id) + model_inputs.update( + query_start_loc=query_start_loc, + ngram_context=ngram_context, + ) + return model_inputs + + +__all__ = ["Qwen4ExpModelState"] diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py new file mode 100644 index 0000000000..87599d99f8 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -0,0 +1,460 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Qwen4Exp MTP (Multi-Token Predictor) model. + +The MTP draft model reuses the Qwen4Exp backbone (PLE/HC/MoE) but: + - drops all multi-modal handling (text-only), + - forces PLE off while keeping the main model's HC stream count, + - fuses the backbone hidden and the new-token embedding via + ``residual_linear_shared`` (fc_embedding + shared fc_hidden) instead of + the ``Linear(2H, H)`` + repeat used by other MTP variants, + - emits TWO hidden streams per step (scheme A): a single stream [T, H] + (final-mixer collapsed, fed to the LM head) and a pre-final-mixer + multi stream [T, hc_count*H] (fed to the next draft step). +""" + +from collections.abc import Iterable + +import regex as re +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig, replace, set_current_vllm_config +from vllm.distributed import get_pp_group +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.utils import configure_quant_config +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + get_draft_quant_config, + make_empty_intermediate_tensors_factory, + maybe_fuse_shared_experts, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) + +from .hyperconnection import GatedResidual, HyperConnectionConfig + +try: + from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm +except ModuleNotFoundError as exc: + # The Blackwell-only CuTe DSL helper is absent from the 1Cat SM70 tree. + if exc.name != "vllm.model_executor.kernels.linear.cute_dsl": + raise + + def enable_qwen4_exp_low_latency_gemm( + module: nn.Module, dtype: torch.dtype + ) -> None: + del module, dtype + + +from .model import ( + _HC_WEIGHTS_MAPPER, + _QWEN3_5_WEIGHTS_MAPPER, + _QWEN4_EXP_IGNORED_MISSING_SUFFIXES, + Qwen4ExpDecoderLayer, + Qwen4ExpMixtureOfExperts, +) + + +def _remap_ignored_layers( + ignored_layers: list[str], + mtp_start_layer_idx: int, +) -> list[str]: + remapped: list[str] = [] + for name in ignored_layers: + if name.startswith("mtp."): + new_name = re.sub( + r"(?<=\.layers\.)\d+", + lambda m: str(mtp_start_layer_idx + int(m.group(0))), + name, + ) + remapped.append(new_name) + else: + remapped.append(name) + return remapped + + +def _remap_mtp_weight_name(name: str) -> str | None: + """Map Qwen4Exp checkpoint paths into the standalone draft model.""" + + for checkpoint_prefix in ( + "model.language_model.", + "language_model.", + ): + if name.startswith(checkpoint_prefix): + name = name.removeprefix(checkpoint_prefix) + break + + if name.startswith("embed_tokens."): + name = f"model.{name}" + if name.startswith("model.mtp."): + name = name.removeprefix("model.") + if name.startswith("mtp.shared_head.head."): + return name.replace("mtp.shared_head.head.", "lm_head.", 1) + if name.startswith("model.shared_head.head."): + return name.replace("model.shared_head.head.", "lm_head.", 1) + if name.startswith("shared_head.head."): + return name.replace("shared_head.head.", "lm_head.", 1) + if name.startswith("model.lm_head."): + return name.removeprefix("model.") + if name.startswith("mtp."): + return name.replace("mtp.", "model.", 1) + if name.startswith("model.embed_tokens.") or name.startswith("lm_head."): + return name + return None + + +def _make_draft_vllm_config( + vllm_config: VllmConfig, + mtp_start_layer_idx: int, +) -> VllmConfig: + """Ensure that the draft model config is set in the vLLM config.""" + speculative_config = vllm_config.speculative_config + if speculative_config is None or speculative_config.draft_model_config is None: + raise ValueError("speculative_config.draft_model_config must be set") + + draft_quant_config = get_draft_quant_config(vllm_config) + + # inject packed and ignored modules to the quantization config of draft model + if draft_quant_config is not None: + configure_quant_config(draft_quant_config, Qwen4ExpMTP) + ignored_layers = getattr(draft_quant_config, "ignored_layers", None) + if ignored_layers: + setattr( # noqa: B010 + draft_quant_config, + "ignored_layers", + _remap_ignored_layers(ignored_layers, mtp_start_layer_idx), + ) + exclude_modules = getattr(draft_quant_config, "exclude_modules", None) + if exclude_modules: + setattr( # noqa: B010 + draft_quant_config, + "exclude_modules", + _remap_ignored_layers(exclude_modules, mtp_start_layer_idx), + ) + + draft_vllm_config = replace( + vllm_config, + model_config=speculative_config.draft_model_config, + ) + # VllmConfig post-init derives the target quant config, so restore the + # independently resolved draft quant config after replacement. + draft_vllm_config.quant_config = draft_quant_config + return draft_vllm_config + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "hidden_states": 0, + } +) +class Qwen4ExpMultiTokenPredictor(nn.Module): + hf_to_vllm_mapper = _QWEN3_5_WEIGHTS_MAPPER | _HC_WEIGHTS_MAPPER + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + model_config = vllm_config.model_config + config: Qwen4ExpTextConfig = model_config.hf_text_config + + self.config = config + self.vocab_size = config.vocab_size + + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1) + + self.hidden_size = config.hidden_size + self.hc_count = config.hc_count + + self.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size) + draft_vllm_config = _make_draft_vllm_config( + vllm_config, + self.mtp_start_layer_idx, + ) + with set_current_vllm_config(draft_vllm_config, prefix=prefix): + # residual_linear_shared fusion: fc_embedding projects the token + # embedding, fc_hidden (shared across HC branches) projects the + # backbone hidden; the embedding is added as a residual to every + # branch (see mtp_residual_linear_shared.md). + self.fc_embedding = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + gather_output=True, + bias=False, + return_bias=False, + quant_config=draft_vllm_config.quant_config, + prefix=f"{prefix}.fc_embedding", + ) + self.fc_hidden = ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + gather_output=True, + bias=False, + return_bias=False, + quant_config=draft_vllm_config.quant_config, + prefix=f"{prefix}.fc_hidden", + ) + self.layers = nn.ModuleList( + Qwen4ExpDecoderLayer( + draft_vllm_config, + layer_type="full_attention", + prefix=f"{prefix}.layers.{self.mtp_start_layer_idx + idx}", + ) + for idx in range(self.num_mtp_layers) + ) + + self.pre_fc_norm_embedding = GemmaRMSNorm( + self.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_hidden = GemmaRMSNorm( + self.hidden_size * self.hc_count, eps=config.rms_norm_eps + ) + # HC final mixer collapses the multi stream into [T, H] for the LM head. + hc_config = HyperConnectionConfig( + hc_count=config.hc_count, + hidden_size=config.hidden_size, + params_dtype=model_config.dtype, + hc_lowrank=config.hc_lowrank, + rms_norm_eps=config.rms_norm_eps, + hc_per_branch_norm=True, + ) + self.hyper_connection_mixer = GatedResidual( + hc_config, + use_combine=False, + prefix=maybe_prefix(prefix, "hyper_connection_mixer"), + ) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], self.hidden_size * self.hc_count + ) + + def _iter_qsa_attentions(self): + """Yield MTP attention modules that own a QSA indexer.""" + + for layer in self.layers: + attention = getattr(layer, "self_attn", None) + if ( + attention is not None + and getattr(attention, "indexer", None) is not None + ): + yield attention + + def set_skip_topk(self, skip: bool) -> None: + """Select on MTP step 0 and reuse its QSA indices on later steps.""" + + for attention in self._iter_qsa_attentions(): + attention.indexer.skip_topk = skip + + def compact_topk_indices(self, row_indices: torch.Tensor) -> None: + """Keep each request's target-aligned step-0 sparse-index row.""" + + num_rows = row_indices.numel() + for attention in self._iter_qsa_attentions(): + buffer = attention.topk_indices_buffer + selected = buffer.index_select(0, row_indices) + buffer[:num_rows].copy_(selected) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors: + hc_count = self.hc_count + hidden_size = self.hidden_size + + if get_pp_group().is_first_rank: + assert hidden_states is not None + if inputs_embeds is None: + assert input_ids is not None + inputs_embeds = self.embed_input_ids(input_ids) + # Embedding branch: pre-norm -> fc_embedding -> [T, H]. + inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds) + inputs_embeds = self.fc_embedding(inputs_embeds) + + # Backbone hidden is multi-stream [T, hc_count*H] (scheme A: + # the main model truly emits the pre-final-mixer multi stream + # on the first step; subsequent steps reuse the prior draft + # step's multi stream). + num_tokens = hidden_states.shape[0] + hidden_states = hidden_states.view(num_tokens, hc_count, hidden_size) + hidden_states = self.pre_fc_norm_hidden(hidden_states.flatten(-2)).view( + num_tokens, hc_count, hidden_size + ) + hidden_states = self.fc_hidden(hidden_states) + # Add the embedding residual to every branch, then fold back + # to [T, hc_count*H] (HC outer, HS inner) for the HC decoder. + hidden_states = inputs_embeds.unsqueeze(-2) + hidden_states + hidden_states = hidden_states.flatten(-2) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + current_step_idx = spec_step_idx % self.num_mtp_layers + layer = self.layers[current_step_idx] + hidden_states, block_output, injection = layer( + hidden_states=hidden_states, + prev_block_output=None, + prev_injection=None, + positions=positions, + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + if not get_pp_group().is_last_rank: + # As in the target model, PP carries a materialized tensor rather + # than the delayed hidden/output/injection tuple. + hidden_states = layer.mlp_hyper_connection.combine( + hidden_states, block_output, injection + ) + return IntermediateTensors({"hidden_states": hidden_states}) + + # Last PP rank finalize. Keep both: + # (A) sample_hidden_states [T, H] -> single stream for the LM head + # (B) multi_hidden [T, hc_count*H] -> pre-final-mixer multi stream + # for the next draft step (zero extra compute, just kept). + multi_hidden, sample_hidden_states, _ = ( + self.hyper_connection_mixer.combine_and_mix( + hidden_states, block_output, injection + ) + ) + return sample_hidden_states, multi_hidden + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + weights = maybe_fuse_shared_experts( + weights, + n_routed_experts=getattr(self.config, "num_experts", 0) or 0, + n_shared_experts=1, + ckpt_prefix="mlp.shared_expert", + ) + loader = AutoWeightsLoader( + self, + skip_substrs=["hyper_connection_mixer.block_inject_weight"], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "hidden_states": 0, + } +) +class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"], + "in_proj_ba": ["in_proj_b", "in_proj_a"], + "input_mix_weight_down_block_inject": [ + "input_mix_weight_down", + "block_inject_weight", + "_input_mix_padding", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config + self.vllm_config = vllm_config + cache_config = vllm_config.cache_config + if cache_config.mamba_cache_mode == "all": + raise NotImplementedError( + "Qwen4ExpMTP currently does not support 'all' prefix caching, " + "please use '--mamba-cache-mode=align' instead" + ) + + self.quant_config = vllm_config.quant_config + + super().__init__() + self.config = config + self.model = Qwen4ExpMultiTokenPredictor( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "mtp"), + ) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + self.set_moe_parameters(self.model.layers) + enable_qwen4_exp_low_latency_gemm(self, vllm_config.model_config.dtype) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | IntermediateTensors: + return self.model( + input_ids, + positions, + hidden_states, + intermediate_tensors, + inputs_embeds, + spec_step_idx=spec_step_idx, + ) + + def compute_logits( + self, hidden_states: torch.Tensor, spec_step_idx: int = 0 + ) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + def remap_weight_names(): + for name, weight in weights: + remapped_name = _remap_mtp_weight_name(name) + if remapped_name is not None: + yield remapped_name, weight + + loader = AutoWeightsLoader( + self, + skip_substrs=["hyper_connection_mixer.block_inject_weight"], + ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), + ) + return loader.load_weights(remap_weight_names()) + + +__all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"] diff --git a/vllm/models/qwen4_exp/nvidia/ops/__init__.py b/vllm/models/qwen4_exp/nvidia/ops/__init__.py new file mode 100644 index 0000000000..779c57191d --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA-only Qwen4Exp kernels; import leaf modules directly.""" diff --git a/vllm/models/qwen4_exp/nvidia/ops/hc.py b/vllm/models/qwen4_exp/nvidia/ops/hc.py new file mode 100644 index 0000000000..a3487b916e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/hc.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA HyperConnection kernels for Qwen4Exp.""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + + +@triton.jit +def _grouped_gemma_rmsnorm_kernel( + x_ptr, + w_ptr, + y_ptr, + stride_x, + stride_y, + DIM: tl.constexpr, + NUM_GROUPS: tl.constexpr, + W_SHARED: tl.constexpr, + EPS: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + GROUP_DIM: tl.constexpr = DIM // NUM_GROUPS + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(GROUP_DIM) + + pid = tl.program_id(0) + group_id = pid % NUM_GROUPS + row = pid // NUM_GROUPS + + offs_g = tl.arange(0, BLOCK_SIZE) + offsets = group_id * GROUP_DIM + offs_g + mask = offs_g < GROUP_DIM + # A [GROUP_DIM] affine is shared; a [DIM] affine follows the grouped + # checkpoint layout. + w_offs = offs_g if W_SHARED else offsets + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + w_offs, mask, other=0.0) + + rrms = tl.rsqrt(tl.sum(x * x) / GROUP_DIM + EPS) + # Gemma's (1 + w) affine is written this way to lower to an FMA. + y = x * rrms + y += y * w.to(tl.float32) + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(y_ptr + row * stride_y + offsets, y, mask) + + +def _grouped_gemma_rmsnorm( + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int +) -> torch.Tensor: + N, DIM = x.shape + assert x.stride(1) == 1, "grouped Gemma RMSNorm requires unit inner stride" + assert weight.is_contiguous(), "grouped Gemma RMSNorm weight must be contiguous" + assert DIM % num_groups == 0 + group_dim = DIM // num_groups + assert weight.numel() in (group_dim, DIM) + + y = x.new_empty(x.shape) + _grouped_gemma_rmsnorm_kernel[(N * num_groups,)]( + x, + weight, + y, + x.stride(0), + y.stride(0), + DIM, + num_groups, + W_SHARED=weight.numel() == group_dim, + EPS=eps, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return y + + +@triton.jit +def _hc_silu_kernel( + x_ptr, + y_ptr, + stride_x, + stride_y, + DIM: tl.constexpr, + HC: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + BLOCK_SIZE: tl.constexpr = triton.next_power_of_2(DIM) + + row = tl.program_id(0) + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < DIM + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + x = tl.load(x_ptr + row * stride_x + offs, mask).to(tl.float32) / HC + y = x * tl.sigmoid(x) + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(y_ptr + row * stride_y + offs, y, mask) + + +def _hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: + num_tokens, DIM = x.shape + assert x.stride(1) == 1 + + output = x.new_empty(x.shape) + _hc_silu_kernel[(num_tokens,)]( + x, + output, + x.stride(0), + output.stride(0), + DIM=DIM, + HC=hc_count, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return output + + +@triton.jit +def _hc_gate_mix_kernel( + x_ptr, + g_ptr, + y_ptr, + stride_x, + stride_g, + stride_y, + DIM: tl.constexpr, + HC: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + HC_DIM: tl.constexpr = DIM // HC + + row = tl.program_id(0) + tile_id = tl.program_id(1) + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs_inner < HC_DIM + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + # The constexpr loop is unrolled and keeps one stream live at a time. + # Materializing [HC, BLOCK_SIZE] more than doubles latency at large M. + acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for stream in tl.static_range(HC): + offsets = stream * HC_DIM + offs_inner + g = tl.load(g_ptr + row * stride_g + offsets, mask, other=0.0) + x = tl.load(x_ptr + row * stride_x + offsets, mask, other=0.0) + acc += tl.sigmoid(g.to(tl.float32)) * x.to(tl.float32) + acc /= HC + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(y_ptr + row * stride_y + offs_inner, acc, mask) + + +def _hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: + N, DIM = gate.shape + assert x.shape == gate.shape + assert DIM % hc_count == 0 + assert x.stride(1) == 1 + assert gate.stride(1) == 1 + + HC_DIM = DIM // hc_count + out = x.new_empty(N, HC_DIM) + BLOCK_SIZE = 512 + _hc_gate_mix_kernel[(N, triton.cdiv(HC_DIM, BLOCK_SIZE))]( + x, + gate, + out, + x.stride(0), + gate.stride(0), + out.stride(0), + DIM, + hc_count, + BLOCK_SIZE, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out + + +@triton.jit +def _hc_combine_kernel( + block_ptr, + res_ptr, + inj_ptr, + out_ptr, + stride_block, + stride_res, + stride_inj, + stride_out, + HC_DIM: tl.constexpr, + HC: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + HC_PAD: tl.constexpr = triton.next_power_of_2(HC) + + row = tl.program_id(0) + tile_id = tl.program_id(1) + + offs_inner = tile_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask_inner = offs_inner < HC_DIM + offs_hc = tl.arange(0, HC_PAD) + mask_hc = offs_hc < HC + offs = offs_hc[:, None] * HC_DIM + offs_inner[None, :] + mask = mask_hc[:, None] & mask_inner[None, :] + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0) + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0) + res = tl.load(res_ptr + row * stride_res + offs, mask, other=0.0) + + # Keeping HC as a broadcast dimension is faster here than four separate + # residual load/store sequences. + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC) + out = res.to(tl.float32) + block.to(tl.float32)[None, :] * inj[:, None] + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + tl.store(out_ptr + row * stride_out + offs, out, mask=mask) + + +def _hc_combine( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + hc_count: int, +) -> torch.Tensor: + N, DIM = residual.shape + assert DIM % hc_count == 0 + hc_dim = DIM // hc_count + assert block_output.shape == (N, hc_dim) + assert injection_logits.shape == (N, hc_count) + assert residual.stride(1) == 1 + assert block_output.stride(1) == 1 + assert injection_logits.stride(1) == 1 + + out = residual.new_empty(residual.shape) + BLOCK_SIZE = 512 + _hc_combine_kernel[(N, triton.cdiv(hc_dim, BLOCK_SIZE))]( + block_output, + residual, + injection_logits, + out, + block_output.stride(0), + residual.stride(0), + injection_logits.stride(0), + out.stride(0), + hc_dim, + hc_count, + BLOCK_SIZE, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out + + +@triton.jit +def _hc_combine_norm_kernel( + block_ptr, + res_ptr, + inj_ptr, + w_ptr, + out_ptr, + y_ptr, + stride_block, + stride_res, + stride_inj, + stride_out, + stride_y, + HC_DIM: tl.constexpr, + HC: tl.constexpr, + W_SHARED: tl.constexpr, + EPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + launch_pdl: tl.constexpr, +) -> None: + HC_PAD: tl.constexpr = triton.next_power_of_2(HC) + NUM_TILES: tl.constexpr = triton.cdiv(HC_DIM, BLOCK_SIZE) + NUM_TILES_PAD: tl.constexpr = triton.next_power_of_2(NUM_TILES) + + row = tl.program_id(0) + stream = tl.program_id(1) + offs_hc = tl.arange(0, HC_PAD) + mask_hc = offs_hc < HC + tile_ids = tl.arange(0, NUM_TILES_PAD) + offs_inner = tile_ids[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] + mask_inner = offs_inner < HC_DIM + offs = stream * HC_DIM + offs_inner + # Shared norm weights repeat across streams; per-branch weights use the + # same flattened HC layout as the residual. + w_offs = offs_inner if W_SHARED else offs + + if launch_pdl: + tl.extra.cuda.gdc_wait() + + # Start the uncached residual load first, then issue the other combine + # loads before consuming any of them. + res = tl.load(res_ptr + row * stride_res + offs, mask_inner, other=0.0) + inj = tl.load(inj_ptr + row * stride_inj + offs_hc, mask_hc, other=0.0) + block = tl.load(block_ptr + row * stride_block + offs_inner, mask_inner, other=0.0) + inj = 2.0 * tl.sigmoid(inj.to(tl.float32) / HC) + inj = tl.sum(tl.where(offs_hc == stream, inj, 0.0)) + # Round the materialized combine result before normalization. This matches + # the unfused combine -> RMSNorm boundary. + out = (res.to(tl.float32) + block.to(tl.float32) * inj).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) + + out = out.to(tl.float32) + # Keep the two-axis reduction: flattening the padded tile is ~40% slower + # at decode sizes. + sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) + + if launch_pdl: + tl.extra.cuda.gdc_launch_dependents() + + # Loading the weight earlier helps decode but keeps the tile live across + # the reduction and regresses larger batches, so defer it to the norm. + w = tl.load(w_ptr + w_offs, mask_inner, other=0.0) + y = out * rrms + y += y * w.to(tl.float32) + tl.store(y_ptr + row * stride_y + offs, y, mask_inner) + + +def _hc_combine_norm( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + N, DIM = residual.shape + assert DIM % hc_count == 0 + hc_dim = DIM // hc_count + assert block_output.shape == (N, hc_dim) + assert injection_logits.shape == (N, hc_count) + assert residual.stride(1) == 1 + assert block_output.stride(1) == 1 + assert injection_logits.stride(1) == 1 + assert norm_weight.is_contiguous() + assert norm_weight.numel() in (hc_dim, DIM) + + out = residual.new_empty(residual.shape) + y = residual.new_empty(residual.shape) + BLOCK_SIZE = 512 + _hc_combine_norm_kernel[(N, hc_count)]( + block_output, + residual, + injection_logits, + norm_weight, + out, + y, + block_output.stride(0), + residual.stride(0), + injection_logits.stride(0), + out.stride(0), + y.stride(0), + hc_dim, + hc_count, + W_SHARED=norm_weight.numel() == hc_dim, + EPS=eps, + BLOCK_SIZE=BLOCK_SIZE, + launch_pdl=current_platform.is_arch_support_pdl(), + ) + return out, y + + +def _same_shape_fake(x: torch.Tensor, *args) -> torch.Tensor: + return x.new_empty(x.shape) + + +def _hc_gate_mix_fake( + x: torch.Tensor, gate: torch.Tensor, hc_count: int +) -> torch.Tensor: + del gate + return x.new_empty((x.shape[0], x.shape[1] // hc_count)) + + +def _hc_combine_fake( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + hc_count: int, +) -> torch.Tensor: + del block_output, injection_logits, hc_count + return residual.new_empty(residual.shape) + + +def _hc_combine_norm_fake( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del block_output, injection_logits, norm_weight, eps, hc_count + return residual.new_empty(residual.shape), residual.new_empty(residual.shape) + + +direct_register_custom_op( + op_name="qwen4_exp_grouped_gemma_rmsnorm", + op_func=_grouped_gemma_rmsnorm, + fake_impl=_same_shape_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_silu", + op_func=_hc_silu, + fake_impl=_same_shape_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_gate_mix", + op_func=_hc_gate_mix, + fake_impl=_hc_gate_mix_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_combine", + op_func=_hc_combine, + fake_impl=_hc_combine_fake, +) +direct_register_custom_op( + op_name="qwen4_exp_hc_combine_norm", + op_func=_hc_combine_norm, + fake_impl=_hc_combine_norm_fake, +) + + +def grouped_gemma_rmsnorm( + x: torch.Tensor, weight: torch.Tensor, eps: float, num_groups: int +) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_grouped_gemma_rmsnorm(x, weight, eps, num_groups) + + +def hc_silu(x: torch.Tensor, hc_count: int) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_silu(x, hc_count) + + +def hc_gate_mix(x: torch.Tensor, gate: torch.Tensor, hc_count: int) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_gate_mix(x, gate, hc_count) + + +def hc_combine( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + hc_count: int, +) -> torch.Tensor: + return torch.ops.vllm.qwen4_exp_hc_combine( + residual, block_output, injection_logits, hc_count + ) + + +def hc_combine_norm( + residual: torch.Tensor, + block_output: torch.Tensor, + injection_logits: torch.Tensor, + norm_weight: torch.Tensor, + eps: float, + hc_count: int, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops.vllm.qwen4_exp_hc_combine_norm( + residual, + block_output, + injection_logits, + norm_weight, + eps, + hc_count, + ) + + +__all__ = [ + "grouped_gemma_rmsnorm", + "hc_combine", + "hc_combine_norm", + "hc_gate_mix", + "hc_silu", +] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py new file mode 100644 index 0000000000..b1fda40d18 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -0,0 +1,1116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernels for the Qwen4Exp weight-free QSA path.""" + +from __future__ import annotations + +import math + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON, tl, triton + +_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 +_TOPK_WORKSPACE_BYTES = 1024 * 1024 + + +@triton.jit +def _qsa_mqa_paged_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + token_to_req_ptr, + query_positions_ptr, + sequence_lengths_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_q_dim, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_table_req, + stride_table_page, + stride_logits_row, + num_rows, + num_columns, + num_pages, + num_requests, + score_divisor, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + TILES_PER_PROG: tl.constexpr, + STAGES: tl.constexpr, + MAX_N: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + heads = tl.arange(0, MAX_N) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_position = tl.load(query_positions_ptr + row) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + visible = tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ) + if tl.program_id(1) == 0: + tl.store(visible_blocks_ptr + row, visible) + tile_start = tl.program_id(1) * TILES_PER_PROG + # Top-k is bounded by visible_blocks, so columns beyond it need no value. + if tile_start * BLOCK_N >= visible: + return + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible, BLOCK_N)) + tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N)) + + # Pad the small head axis to a tensor-core-compatible N dimension. + query = tl.load( + q_ptr + + row * stride_q_row + + heads[None, :] * stride_q_head + + dims[:, None] * stride_q_dim, + mask=(heads[None, :] < NUM_HEADS) & (dims[:, None] < HEAD_DIM), + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < visible + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + + safe_request * stride_table_req + + logical_page * stride_table_page, + mask=live, + other=-1, + ) + page_valid = live & (physical_page >= 0) & (physical_page < num_pages) + # physical_page * block stride can overflow int32 for large caches. + safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_physical_page[:, None] * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :] * stride_cache_dim, + mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + scores = tl.dot(keys, query, out_dtype=tl.float32) + scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(scores, 0.0), 0.0) + score = tl.sum(scores, axis=1) / score_divisor + tl.store( + logits_ptr + row * stride_logits_row + columns, + tl.where(page_valid, score, -float("inf")), + mask=live & (columns < num_columns), + ) + + +@triton.jit +def _expand_qsa_indices_kernel( + block_indices_ptr, + query_positions_ptr, + sequence_lengths_ptr, + token_to_req_ptr, + output_ptr, + stride_blocks_row, + stride_blocks_column, + stride_output_row, + stride_output_column, + rows, + num_requests, + BLOCK_TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + TOKEN_TOPK: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + COLUMN_BLOCK: tl.constexpr, +) -> None: + row = tl.program_id(0) + columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK) + query_position = tl.load(query_positions_ptr + row) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + sequence_length = tl.load( + sequence_lengths_ptr + safe_request, + mask=(request >= 0) & (request < num_requests), + other=0, + ) + complete_blocks = tl.minimum( + tl.minimum( + (query_position + 1) // COMPRESS_RATIO, + sequence_length // COMPRESS_RATIO, + ), + BLOCK_TOPK, + ) + expanded_count = complete_blocks * COMPRESS_RATIO + tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO + tail_count = (query_position + 1) - tail_start + + is_expanded = columns < expanded_count + block_rank = columns // COMPRESS_RATIO + offset = columns % COMPRESS_RATIO + safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1) + block = tl.load( + block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column, + mask=(row < rows) & is_expanded, + other=-1, + ) + expanded = block * COMPRESS_RATIO + offset + tail_offset = columns - expanded_count + is_tail = ( + (columns >= expanded_count) + & (tail_offset < tail_count) + & (tail_offset < COMPRESS_RATIO - 1) + ) + token = tl.where(is_expanded, expanded, tail_start + tail_offset) + valid = ( + (row < rows) + & (columns < OUTPUT_WIDTH) + & (is_expanded | is_tail) + & (token >= 0) + & (token < sequence_length) + ) + tl.store( + output_ptr + row * stride_output_row + columns * stride_output_column, + tl.where(valid, token, -1), + mask=(row < rows) & (columns < OUTPUT_WIDTH), + ) + + +@triton.jit +def _qsa_sparse_paged_gqa_splitk_kernel( + q_ptr, + k_cache_ptr, + v_cache_ptr, + indices_ptr, + block_table_ptr, + token_to_req_ptr, + partial_output_ptr, + partial_lse_ptr, + output_ptr, + stride_q_row, + stride_q_head, + stride_k_block, + stride_k_token, + stride_k_head, + stride_v_block, + stride_v_token, + stride_v_head, + stride_indices_row, + stride_table_req, + stride_output_row, + stride_output_head, + num_rows, + num_cache_blocks, + num_requests, + TOPK: tl.constexpr, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + GROUP_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + NUM_TILES: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +) -> None: + row = tl.program_id(0) + kv_head = tl.program_id(1) + split_id = tl.program_id(2) + request = tl.load(token_to_req_ptr + row) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + + head_offsets = tl.arange(0, BLOCK_M) + dim_offsets = tl.arange(0, HEAD_DIM) + column_offsets = tl.arange(0, BLOCK_N) + first_head = kv_head * GROUP_SIZE + query = tl.load( + q_ptr + + row * stride_q_row + + (first_head + head_offsets[:, None]) * stride_q_head + + dim_offsets[None, :], + mask=head_offsets[:, None] < GROUP_SIZE, + other=0.0, + ) + + max_value = tl.full((BLOCK_M,), -1.0e20, dtype=tl.float32) + normalizer = tl.zeros((BLOCK_M,), dtype=tl.float32) + accumulator = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32) + softmax_scale_log2: tl.constexpr = (HEAD_DIM**-0.5) * 1.4426950408889634 + + # Dynamic bounds avoid padded main-loop iterations for uneven splits. + split_tile_start = split_id * NUM_TILES // NUM_SPLITS + split_tile_end = (split_id + 1) * NUM_TILES // NUM_SPLITS + for tile in range(split_tile_start, split_tile_end): + columns = tile * BLOCK_N + column_offsets + logical_token = tl.load( + indices_ptr + row * stride_indices_row + columns, + mask=columns < TOPK, + other=-1, + ) + safe_token = tl.maximum(logical_token, 0) + logical_page = safe_token // PAGE_SIZE + page_offset = safe_token % PAGE_SIZE + valid = ( + (request >= 0) + & (request < num_requests) + & (logical_token >= 0) + & (logical_page < PAGE_TABLE_WIDTH) + ) + physical_page = tl.load( + block_table_ptr + + safe_request * stride_table_req + + tl.minimum(logical_page, PAGE_TABLE_WIDTH - 1), + mask=valid, + other=-1, + ) + valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) + # physical_page * block stride can overflow int32 for large caches. + safe_page = tl.maximum(physical_page, 0).to(tl.int64) + keys = tl.load( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + mask=valid[None, :], + other=0.0, + ) + values = tl.load( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) + scores = tl.dot(query, keys) + # Scaling scores avoids re-quantizing a scaled query to BF16. + scores *= softmax_scale_log2 + scores = tl.where(valid[None, :], scores, -1.0e20) + next_max = tl.maximum(max_value, tl.max(scores, axis=1)) + alpha = tl.math.exp2(max_value - next_max) + probabilities = tl.where( + valid[None, :], tl.math.exp2(scores - next_max[:, None]), 0.0 + ) + accumulator = tl.dot( + probabilities.to(values.dtype), + values, + acc=accumulator * alpha[:, None], + ) + normalizer = normalizer * alpha + tl.sum(probabilities, axis=1) + max_value = next_max + + has_values = normalizer > 0 + normalized_output = tl.where( + has_values[:, None], + accumulator / tl.maximum(normalizer[:, None], 1.0e-20), + 0.0, + ) + output_mask = head_offsets[:, None] < GROUP_SIZE + if NUM_SPLITS == 1: + tl.store( + output_ptr + + row * stride_output_row + + (first_head + head_offsets[:, None]) * stride_output_head + + dim_offsets[None, :], + normalized_output, + mask=output_mask, + ) + else: + partial_lse = tl.where( + has_values, + max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), + -float("inf"), + ) + tl.store( + partial_output_ptr + + ( + (split_id * num_rows + row) * NUM_QUERY_HEADS + + first_head + + head_offsets[:, None] + ) + * HEAD_DIM + + dim_offsets[None, :], + normalized_output, + mask=output_mask, + ) + tl.store( + partial_lse_ptr + + (split_id * num_rows + row) * NUM_QUERY_HEADS + + first_head + + head_offsets, + partial_lse, + mask=head_offsets < GROUP_SIZE, + ) + + +@triton.jit +def _qsa_merge_splitk_kernel( + partial_output_ptr, + partial_lse_ptr, + output_ptr, + stride_output_row, + stride_output_head, + num_rows, + HEAD_DIM: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + NUM_SPLITS: tl.constexpr, + BLOCK_SPLITS: tl.constexpr, +) -> None: + row = tl.program_id(0) + head = tl.program_id(1) + split_offsets = tl.arange(0, BLOCK_SPLITS) + dim_offsets = tl.arange(0, HEAD_DIM) + split_mask = split_offsets < NUM_SPLITS + lse = tl.load( + partial_lse_ptr + (split_offsets * num_rows + row) * NUM_QUERY_HEADS + head, + mask=split_mask, + other=-float("inf"), + ) + lse_max = tl.max(lse, axis=0) + has_values = lse_max > -float("inf") + shifted = tl.where(split_mask & has_values, lse - lse_max, -float("inf")) + weights = tl.math.exp2(shifted) + denominator = tl.sum(weights, axis=0) + partial_output = tl.load( + partial_output_ptr + + ((split_offsets[:, None] * num_rows + row) * NUM_QUERY_HEADS + head) + * HEAD_DIM + + dim_offsets[None, :], + mask=split_mask[:, None], + other=0.0, + ) + merged = tl.sum(partial_output * weights[:, None], axis=0) + merged = tl.where(denominator > 0, merged / denominator, 0.0) + tl.store( + output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, + merged, + ) + + +@triton.jit +def _store_qsa_rows_kernel( + cache_ptr, + slots_ptr, + rows_ptr, + stride_cache_block, + stride_cache_token, + stride_cache_dim, + stride_rows_row, + stride_rows_dim, + num_rows, + num_blocks, + PAGE_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_D: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + slot = tl.load(slots_ptr + row) + valid = (row < num_rows) & (slot >= 0) & (slot < num_blocks * PAGE_SIZE) + block = tl.maximum(slot, 0) // PAGE_SIZE + token = tl.maximum(slot, 0) % PAGE_SIZE + values = tl.load( + rows_ptr + row * stride_rows_row + dims * stride_rows_dim, + mask=valid & (dims < WIDTH), + other=0, + ) + tl.store( + cache_ptr + + block * stride_cache_block + + token * stride_cache_token + + dims * stride_cache_dim, + values, + mask=valid & (dims < WIDTH), + ) + + +@triton.jit +def _compress_qsa_groups_kernel( + raw_keys_ptr, # this step's raw key rows, straight from activations + raw_positions_ptr, # this step's per-token positions + compressor_state_cache_ptr, # per-request ring of previous raw keys + rope_cache_ptr, # packed RoPE position tail of the ring + compressor_state_table_ptr, + token_to_req_ptr, + query_start_loc_ptr, + logical_positions_ptr, + compressed_slots_ptr, + pooled_ptr, + first_positions_ptr, + stride_raw_row, + stride_raw_dim, + stride_raw_positions_row, + stride_raw_positions_dim, + stride_compressor_state_block, + stride_compressor_state_token, + stride_compressor_state_dim, + stride_rope_block, + stride_rope_token, + stride_rope_dim, + stride_compressor_state_table_req, + stride_pooled_row, + stride_pooled_dim, + stride_positions_row, + stride_positions_dim, + num_rows, + num_compressor_state_blocks, + num_requests, + COMPRESSOR_STATE_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + LOAD_ROPE_POSITIONS: tl.constexpr, +) -> None: + row = tl.program_id(0) + dims = tl.arange(0, BLOCK_D) + request = tl.load(token_to_req_ptr + row) + end_position = tl.load(logical_positions_ptr + row) + compressed_slot = tl.load(compressed_slots_ptr + row) + valid_request = (request >= 0) & (request < num_requests) + safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) + query_row_start = tl.load( + query_start_loc_ptr + safe_request, mask=valid_request, other=0 + ) + query_row_end = tl.load( + query_start_loc_ptr + safe_request + 1, mask=valid_request, other=0 + ) + chunk_start_position = end_position - (row - query_row_start) + compressor_state_block = tl.load( + compressor_state_table_ptr + safe_request * stride_compressor_state_table_req, + mask=valid_request, + other=-1, + ) + valid_compressor_state_block = (compressor_state_block >= 0) & ( + compressor_state_block < num_compressor_state_blocks + ) + valid_row = ( + (row < num_rows) + & valid_request + & (row >= query_row_start) + & (row < query_row_end) + & (end_position >= COMPRESS_RATIO - 1) + & (compressed_slot >= 0) + ) + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + + # A group can span the compressor-state ring (older members) and this + # step's raw rows (members at positions >= chunk_start_position). + for group_offset in tl.range(0, COMPRESS_RATIO): + position = end_position - (COMPRESS_RATIO - 1 - group_offset) + use_raw = position >= chunk_start_position + raw_row = query_row_start + position - chunk_start_position + raw_values = tl.load( + raw_keys_ptr + raw_row * stride_raw_row + dims * stride_raw_dim, + mask=valid_row + & use_raw + & (raw_row >= query_row_start) + & (raw_row < query_row_end) + & (raw_row < num_rows) + & (dims < HEAD_DIM), + other=0.0, + ).to(tl.float32) + compressor_state_values = tl.load( + compressor_state_cache_ptr + + tl.maximum(compressor_state_block, 0).to(tl.int64) + * stride_compressor_state_block + + (position % COMPRESSOR_STATE_SIZE) * stride_compressor_state_token + + dims * stride_compressor_state_dim, + mask=valid_row + & ~use_raw + & valid_compressor_state_block + & (dims < HEAD_DIM), + other=0.0, + ).to(tl.float32) + accumulator += tl.where(use_raw, raw_values, compressor_state_values) + + tl.store( + pooled_ptr + row * stride_pooled_row + dims * stride_pooled_dim, + accumulator / COMPRESS_RATIO, + mask=(row < num_rows) & (dims < HEAD_DIM), + ) + + position_dims = tl.arange(0, 4) + first_position = end_position - COMPRESS_RATIO + 1 + if LOAD_ROPE_POSITIONS: + first_from_raw = first_position >= chunk_start_position + raw_first_row = query_row_start + first_position - chunk_start_position + raw_position_values = tl.load( + raw_positions_ptr + + raw_first_row * stride_raw_positions_row + + position_dims * stride_raw_positions_dim, + mask=valid_row + & first_from_raw + & (raw_first_row >= query_row_start) + & (raw_first_row < query_row_end) + & (raw_first_row < num_rows) + & (position_dims < 3), + other=0, + ) + compressor_state_position_values = tl.load( + rope_cache_ptr + + tl.maximum(compressor_state_block, 0).to(tl.int64) * stride_rope_block + + (first_position % COMPRESSOR_STATE_SIZE) * stride_rope_token + + position_dims * stride_rope_dim, + mask=valid_row + & ~first_from_raw + & valid_compressor_state_block + & (position_dims < 3), + other=0, + ) + position_values = tl.where( + first_from_raw, + raw_position_values, + compressor_state_position_values, + ) + else: + position_values = tl.where(valid_row, first_position, 0) + tl.store( + first_positions_ptr + + row * stride_positions_row + + position_dims * stride_positions_dim, + position_values, + mask=(row < num_rows) & (position_dims < 3), + ) + + +def _validate_mqa(q: torch.Tensor) -> None: + if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0: + raise ValueError("QSA query must be [rows, heads, head_dim]") + + +def qsa_mqa_paged( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + compress_ratio: int, + num_columns: int | None = None, + score_scale: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute QSA scores directly from a paged compressed-key cache.""" + + _validate_mqa(q) + if not q.is_cuda or not HAS_TRITON: + raise RuntimeError("paged QSA scoring requires CUDA and Triton") + if k_cache.ndim != 4 or k_cache.shape[2] != 1: + raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]") + if k_cache.shape[3] != q.shape[2]: + raise ValueError("QSA query and cache dimensions must match") + if page_table.ndim != 2: + raise ValueError("QSA page table must be two-dimensional") + if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)): + raise ValueError("QSA paged scoring cache and page table must be nonempty") + if token_to_req.shape != (q.shape[0],): + raise ValueError("QSA request mapping must match query rows") + if query_positions.shape != (q.shape[0],): + raise ValueError("QSA query positions must match query rows") + if sequence_lengths.shape != (page_table.shape[0],): + raise ValueError("QSA sequence lengths must match page-table requests") + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale + if score_divisor <= 0: + raise ValueError("QSA score scale must be positive") + + capacity = page_table.shape[1] * k_cache.shape[1] + columns = capacity if num_columns is None else num_columns + if columns < 0: + raise ValueError("QSA score width must be non-negative") + logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) + visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) + if not q.shape[0] or not columns: + return logits, visible_blocks + BLOCK_N = 64 + BLOCK_D = max(16, triton.next_power_of_2(q.shape[2])) + MAX_N = max(16, triton.next_power_of_2(q.shape[1])) + # Tuned on GB300: larger row batches provide enough parallelism to reuse Q. + tiles_per_program = 1 if q.shape[0] <= 32 else 8 + _qsa_mqa_paged_kernel[ + (q.shape[0], triton.cdiv(columns, BLOCK_N * tiles_per_program)) + ]( + q, + k_cache, + page_table, + token_to_req, + query_positions, + sequence_lengths, + visible_blocks, + logits, + q.stride(0), + q.stride(1), + q.stride(2), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(3), + page_table.stride(0), + page_table.stride(1), + logits.stride(0), + q.shape[0], + columns, + k_cache.shape[0], + page_table.shape[0], + float(score_divisor), + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + NUM_HEADS=q.shape[1], + HEAD_DIM=q.shape[2], + BLOCK_N=BLOCK_N, + BLOCK_D=BLOCK_D, + TILES_PER_PROG=tiles_per_program, + STAGES=2, + MAX_N=MAX_N, + COMPRESS_RATIO=compress_ratio, + num_warps=2, + ) + return logits, visible_blocks + + +def expand_qsa_block_indices_cuda( + block_indices: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + token_to_req: torch.Tensor, + compress_ratio: int, + token_topk: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Expand compressed blocks and compact the causal tail of the open group.""" + + if not block_indices.is_cuda or not HAS_TRITON: + raise RuntimeError("QSA CUDA expansion requires Triton") + if token_topk % compress_ratio: + raise ValueError("QSA token top-k must be divisible by compression ratio") + block_topk = token_topk // compress_ratio + output_width = token_topk + compress_ratio - 1 + if block_indices.shape != (query_positions.numel(), block_topk): + raise ValueError("QSA compressed top-k has an invalid shape") + if token_to_req.shape != query_positions.shape: + raise ValueError("QSA request mapping must match query positions") + if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]: + raise ValueError("QSA request sequence lengths must be nonempty") + if out is None: + out = torch.empty( + (block_indices.shape[0], output_width), + dtype=torch.int32, + device=block_indices.device, + ) + elif out.shape != (block_indices.shape[0], output_width): + raise ValueError("QSA expansion output has an invalid shape") + if not block_indices.shape[0]: + return out + column_block = 256 + _expand_qsa_indices_kernel[ + (block_indices.shape[0], triton.cdiv(output_width, column_block)) + ]( + block_indices, + query_positions, + sequence_lengths, + token_to_req, + out, + block_indices.stride(0), + block_indices.stride(1), + out.stride(0), + out.stride(1), + block_indices.shape[0], + sequence_lengths.shape[0], + BLOCK_TOPK=block_topk, + COMPRESS_RATIO=compress_ratio, + TOKEN_TOPK=token_topk, + OUTPUT_WIDTH=output_width, + COLUMN_BLOCK=column_block, + num_warps=4, + ) + return out + + +def qsa_select_paged_tokens( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + token_to_req: torch.Tensor, + query_positions: torch.Tensor, + sequence_lengths: torch.Tensor, + token_topk: int, + compress_ratio: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Score, select, and expand QSA indices without host synchronization.""" + + rows = q.shape[0] + output_width = token_topk + compress_ratio - 1 + if out is None: + out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device) + if out.shape != (rows, output_width): + raise ValueError("QSA selection output has an invalid shape") + if not rows: + return out + + columns = page_table.shape[1] * k_cache.shape[1] + block_topk = token_topk // compress_ratio + rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)) + chunk_rows = min(rows, rows_per_chunk) + blocks_buffer = torch.empty( + (chunk_rows, block_topk), dtype=torch.int32, device=q.device + ) + topk_workspace = torch.empty( + (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device + ) + for row_start in range(0, rows, rows_per_chunk): + row_end = min(row_start + rows_per_chunk, rows) + row_slice = slice(row_start, row_end) + logits, visible_blocks = qsa_mqa_paged( + q[row_slice], + k_cache, + page_table, + token_to_req[row_slice], + query_positions[row_slice], + sequence_lengths, + compress_ratio, + ) + blocks = blocks_buffer[: row_end - row_start] + use_cooperative_topk = ( + blocks.shape[0] <= 32 + and logits.stride(0) % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + topk_op = ( + torch.ops._C.cooperative_topk + if use_cooperative_topk + else torch.ops._C.persistent_topk + ) + topk_op(logits, visible_blocks, blocks, topk_workspace, block_topk, columns) + expand_qsa_block_indices_cuda( + blocks, + query_positions[row_slice], + sequence_lengths, + token_to_req[row_slice], + compress_ratio, + token_topk, + out[row_slice], + ) + return out + + +def qsa_sparse_paged_attention( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + logical_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_req: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Run sparse GQA directly over paged FP16/BF16 K/V caches.""" + + if not q.is_cuda or not HAS_TRITON: + raise RuntimeError("paged QSA sparse attention requires CUDA and Triton") + if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("QSA sparse attention received invalid Q/K/V shapes") + if logical_indices.ndim != 2 or logical_indices.shape[0] != q.shape[0]: + raise ValueError("QSA indices must have one row per query") + if token_to_req.shape != (q.shape[0],) or block_table.ndim != 2: + raise ValueError("QSA sparse attention metadata has invalid shapes") + if not all(k_cache.shape[:3]) or not all(block_table.shape): + raise ValueError("QSA sparse attention cache and block table must be nonempty") + if logical_indices.shape[1] <= 0: + raise ValueError("QSA sparse attention requires a positive selection width") + if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: + raise ValueError("QSA sparse attention requires valid grouped-query heads") + head_dim = q.shape[2] + assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 + assert q.dtype == k_cache.dtype == v_cache.dtype + assert q.dtype in (torch.float16, torch.bfloat16) + assert logical_indices.dtype == block_table.dtype == torch.int32 + assert token_to_req.dtype == torch.int32 + assert q.device == k_cache.device == v_cache.device + assert q.device == logical_indices.device == block_table.device + assert q.device == token_to_req.device + assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 + assert logical_indices.stride(1) == block_table.stride(1) == 1 + assert token_to_req.stride(0) == 1 + if out is None: + out = torch.empty_like(q) + if out.shape != q.shape: + raise ValueError("QSA sparse output must match its query") + assert out.dtype == q.dtype and out.device == q.device + assert out.stride(2) == 1 + if not q.shape[0]: + return out + + group_size = q.shape[1] // k_cache.shape[2] + block_m = triton.next_power_of_2(group_size) + base_programs = q.shape[0] * k_cache.shape[2] + small_profile_limit = 8 if block_m <= 8 else 4 + + # Tuned on GB300 for the Qwen-Air TP1, TP2, and TP4 attention shapes. + # Narrow tiles favor decode; wide tiles improve throughput for prefill. + if base_programs <= small_profile_limit: + block_n, target_splits, partial_warps = 16, 64, 4 + elif base_programs < 32: + block_n, target_splits, partial_warps = 16, 32, 4 + elif base_programs <= 256: + block_n, target_splits, partial_warps = 64, 8, 2 + elif base_programs <= 512: + block_n, target_splits, partial_warps = 64, 4, 2 + else: + block_n, target_splits, partial_warps = 64, 1, 2 + + num_tiles = triton.cdiv(logical_indices.shape[1], block_n) + # Avoid empty splits when the selection width is smaller than the profile. + max_useful_splits = 1 << (num_tiles.bit_length() - 1) + num_splits = min(max_useful_splits, target_splits) + + # Split=1 writes output directly and compiles out all workspace accesses. + if num_splits == 1: + partial_output = out + partial_lse = out + else: + # FP32 partials preserve accuracy when merging independently normalized + # splits. + partial_output = torch.empty( + (num_splits, *q.shape), dtype=torch.float32, device=q.device + ) + partial_lse = torch.empty( + (num_splits, q.shape[0], q.shape[1]), + dtype=torch.float32, + device=q.device, + ) + + partial_grid = (q.shape[0], k_cache.shape[2], num_splits) + _qsa_sparse_paged_gqa_splitk_kernel[partial_grid]( + q, + k_cache, + v_cache, + logical_indices, + block_table, + token_to_req, + partial_output, + partial_lse, + out, + q.stride(0), + q.stride(1), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + logical_indices.stride(0), + block_table.stride(0), + out.stride(0), + out.stride(1), + q.shape[0], + k_cache.shape[0], + block_table.shape[0], + TOPK=logical_indices.shape[1], + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=block_table.shape[1], + GROUP_SIZE=group_size, + HEAD_DIM=q.shape[2], + NUM_QUERY_HEADS=q.shape[1], + NUM_SPLITS=num_splits, + NUM_TILES=num_tiles, + BLOCK_M=block_m, + BLOCK_N=block_n, + num_warps=partial_warps, + num_stages=2, + ) + if num_splits == 1: + return out + + _qsa_merge_splitk_kernel[(q.shape[0], q.shape[1])]( + partial_output, + partial_lse, + out, + out.stride(0), + out.stride(1), + q.shape[0], + HEAD_DIM=q.shape[2], + NUM_QUERY_HEADS=q.shape[1], + NUM_SPLITS=num_splits, + BLOCK_SPLITS=triton.next_power_of_2(num_splits), + num_warps=2, + num_stages=1, + ) + return out + + +def qsa_store_cache_rows( + cache: torch.Tensor, + slot_mapping: torch.Tensor, + rows: torch.Tensor, +) -> None: + """Store fixed-width rows in a QSA cache without boolean indexing.""" + + if not cache.is_cuda or not HAS_TRITON: + raise RuntimeError("QSA CUDA cache stores require Triton") + if cache.ndim != 4 or cache.shape[2] != 1: + raise ValueError("QSA cache must be [pages, page_size, 1, width]") + if not all(cache.shape): + raise ValueError("QSA cache dimensions must be nonzero") + if rows.ndim == 3: + if rows.shape[1] != 1: + raise ValueError("QSA cache rows must have one head") + rows = rows[:, 0] + if rows.shape != (slot_mapping.numel(), cache.shape[3]): + raise ValueError("QSA cache rows and slots have incompatible shapes") + if not rows.shape[0]: + return + _store_qsa_rows_kernel[(rows.shape[0],)]( + cache, + slot_mapping, + rows, + cache.stride(0), + cache.stride(1), + cache.stride(3), + rows.stride(0), + rows.stride(1), + rows.shape[0], + cache.shape[0], + PAGE_SIZE=cache.shape[1], + WIDTH=cache.shape[3], + BLOCK_D=triton.next_power_of_2(cache.shape[3]), + num_warps=4, + ) + + +def qsa_compress_groups_with_ratio( + raw_keys: torch.Tensor, # this step's raw key rows [rows, 1, head_size] + raw_positions: torch.Tensor, # this step's positions [rows, 1, 3] int64 + compressor_state_cache: torch.Tensor, + compressor_state_block_table: torch.Tensor, + token_to_req: torch.Tensor, + query_start_loc: torch.Tensor, + logical_positions: torch.Tensor, + compressed_slots: torch.Tensor, + compress_ratio: int, + rope_cache: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pool completed groups from the compressor-state ring and raw token rows.""" + + if not raw_keys.is_cuda or not HAS_TRITON: + raise RuntimeError("QSA CUDA compression requires Triton") + rows = token_to_req.numel() + if compress_ratio <= 0: + raise ValueError("QSA compression ratio must be positive") + if raw_keys.ndim != 3 or raw_keys.shape[:2] != (rows, 1): + raise ValueError("QSA raw keys must be [rows, 1, head_size]") + if raw_positions.shape != (rows, 1, 3) or raw_positions.dtype != torch.int64: + raise ValueError("QSA raw positions must be [rows, 1, 3] int64") + if logical_positions.shape != (rows,) or compressed_slots.shape != (rows,): + raise ValueError("QSA compression metadata must match token rows") + if compressor_state_cache.ndim != 4 or compressor_state_cache.shape[2] != 1: + raise ValueError("QSA compressor-state cache has an invalid shape") + if ( + # The ring is wider than one group so speculative rows cannot alias + # onto the committed keys of the group still being collected. + compressor_state_cache.shape[1] < compress_ratio + or compressor_state_cache.shape[3] != raw_keys.shape[2] + or compressor_state_cache.dtype != raw_keys.dtype + ): + raise ValueError( + "QSA compressor-state cache does not match the compression layout" + ) + if ( + compressor_state_block_table.ndim != 2 + or compressor_state_block_table.shape[1] < 1 + ): + raise ValueError( + "QSA compressor-state block table must contain one block per request" + ) + if query_start_loc.ndim != 1 or query_start_loc.shape[0] < 2: + raise ValueError("QSA query starts must contain a terminal offset") + num_requests = query_start_loc.shape[0] - 1 + if compressor_state_block_table.shape[0] < num_requests: + raise ValueError("QSA compressor-state block table has too few request rows") + if rope_cache is not None and ( + rope_cache.ndim != 4 + or rope_cache.shape[:3] != compressor_state_cache.shape[:3] + or rope_cache.shape[3] != 3 + or rope_cache.dtype != torch.int64 + ): + raise ValueError("QSA packed position view has an invalid shape or dtype") + if rows and ( + not all(compressor_state_cache.shape) + or not all(compressor_state_block_table.shape) + ): + raise ValueError("QSA compressor-state cache and block table must be nonempty") + pooled = torch.empty( + (rows, 1, raw_keys.shape[2]), + dtype=raw_keys.dtype, + device=raw_keys.device, + ) + first_positions = torch.empty((rows, 3), dtype=torch.int64, device=raw_keys.device) + if not rows: + return pooled, first_positions + if rope_cache is None: + rope_cache = compressor_state_cache + load_rope_positions = False + else: + load_rope_positions = True + _compress_qsa_groups_kernel[(rows,)]( + raw_keys, + raw_positions, + compressor_state_cache, + rope_cache, + compressor_state_block_table, + token_to_req, + query_start_loc, + logical_positions, + compressed_slots, + pooled, + first_positions, + raw_keys.stride(0), + raw_keys.stride(2), + raw_positions.stride(0), + raw_positions.stride(2), + compressor_state_cache.stride(0), + compressor_state_cache.stride(1), + compressor_state_cache.stride(3), + rope_cache.stride(0), + rope_cache.stride(1), + rope_cache.stride(3), + compressor_state_block_table.stride(0), + pooled.stride(0), + pooled.stride(2), + first_positions.stride(0), + first_positions.stride(1), + rows, + compressor_state_cache.shape[0], + num_requests, + COMPRESSOR_STATE_SIZE=compressor_state_cache.shape[1], + COMPRESS_RATIO=compress_ratio, + HEAD_DIM=raw_keys.shape[2], + LOAD_ROPE_POSITIONS=load_rope_positions, + BLOCK_D=triton.next_power_of_2(raw_keys.shape[2]), + num_warps=4, + ) + return pooled, first_positions + + +__all__ = [ + "expand_qsa_block_indices_cuda", + "qsa_compress_groups_with_ratio", + "qsa_mqa_paged", + "qsa_select_paged_tokens", + "qsa_sparse_paged_attention", + "qsa_store_cache_rows", +] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py new file mode 100644 index 0000000000..3c1695548e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_pre_indexer.py @@ -0,0 +1,516 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused QSA pre-indexer kernel for Qwen4Exp.""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _norm_rope( + x, + pos_t, + pos_h, + pos_w, + cos_sin_ptr, + cos_sin_stride, + norm_weight_ptr, + eps, + IS_MROPE: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, +): + """Apply Gemma RMSNorm and selected-axis NeoX RoPE to register rows.""" + TILE_T: tl.constexpr = x.shape[0] + TILE_H: tl.constexpr = x.shape[1] + D: tl.constexpr = x.shape[2] + ROWS: tl.constexpr = TILE_T * TILE_H + HALF: tl.constexpr = D // 2 + QUARTER: tl.constexpr = D // 4 + pairs = tl.arange(0, QUARTER) + if IS_MROPE: + # Qwen interleaves temporal, height, and width rotary pairs. Each axis + # still indexes the same position-major cos/sin table. + h_mask = ((pairs % 3) == 1) & (pairs <= 3 * MROPE_H) + w_mask = ((pairs % 3) == 2) & (pairs <= 3 * MROPE_W) + t_mask = ~(h_mask | w_mask) + base = cos_sin_ptr + pairs[None, :] + pos_rows = (pos_t, pos_h, pos_w) + axis_masks = (t_mask, h_mask, w_mask) + cos = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) + sin = tl.zeros((TILE_T, QUARTER), dtype=cos_sin_ptr.dtype.element_ty) + for axis in tl.static_range(3): + cos += tl.load( + base + pos_rows[axis][:, None] * cos_sin_stride, + mask=axis_masks[axis][None, :], + other=0, + ) + sin += tl.load( + base + pos_rows[axis][:, None] * cos_sin_stride + QUARTER, + mask=axis_masks[axis][None, :], + other=0, + ) + else: + cos = tl.load(cos_sin_ptr + pos_t[:, None] * cos_sin_stride + pairs[None, :]) + sin = tl.load( + cos_sin_ptr + pos_t[:, None] * cos_sin_stride + QUARTER + pairs[None, :] + ) + + cos = tl.reshape( + tl.broadcast_to(cos[:, None, :], (TILE_T, TILE_H, QUARTER)), + (ROWS, QUARTER), + ) + sin = tl.reshape( + tl.broadcast_to(sin[:, None, :], (TILE_T, TILE_H, QUARTER)), + (ROWS, QUARTER), + ) + x = tl.reshape(x, (ROWS, D)).to(tl.float32) + weight = tl.load(norm_weight_ptr + tl.arange(0, D)).to(tl.float32) + 1.0 + rrms = tl.rsqrt(tl.sum(x * x, axis=1) / D + eps) + y = (x * rrms[:, None] * weight[None, :]).to(cos.dtype) + rotated, passthrough = tl.split( + tl.permute(tl.reshape(y, (ROWS, 2, HALF)), (0, 2, 1)) + ) + r0, r1 = tl.split(tl.permute(tl.reshape(rotated, (ROWS, 2, QUARTER)), (0, 2, 1))) + out0 = r0 * cos - r1 * sin + out1 = r1 * cos + r0 * sin + rotated = tl.reshape(tl.permute(tl.join(out0, out1), (0, 2, 1)), (ROWS, HALF)) + result = tl.reshape(tl.permute(tl.join(rotated, passthrough), (0, 2, 1)), (ROWS, D)) + return tl.reshape(result, (TILE_T, TILE_H, D)) + + +@triton.jit( + do_not_specialize=[ + "num_tokens", + "num_state_blocks", + "num_compressed_blocks", + "num_k_work", + ] +) +def _qsa_pre_indexer_kernel( + q_ptr, + q_stride_token, + k_ptr, + k_stride_token, + pos_ptr, + pos_stride_axis, + pos_stride_token, + cos_sin_ptr, + q_norm_weight_ptr, + k_norm_weight_ptr, + eps, + q_out_ptr, + q_out_stride_token, + q_out_stride_head, + state_cache_ptr, + state_cache_stride_block, + state_cache_stride_token, + state_slots_ptr, + state_table_ptr, + state_table_stride_req, + query_start_loc_ptr, + logical_positions_ptr, + compressed_slots_ptr, + k_work_metadata_ptr, + compressed_cache_ptr, + compressed_cache_stride_block, + compressed_cache_stride_token, + num_tokens, + num_state_blocks, + num_compressed_blocks, + num_k_work, + HQ: tl.constexpr, + D: tl.constexpr, + TILE_T_Q: tl.constexpr, + TILE_H_Q: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + STATE_SIZE: tl.constexpr, + COMP_PAGE_SIZE: tl.constexpr, + IS_2D_POSITIONS: tl.constexpr, + IS_K_MROPE: tl.constexpr, + CACHE_HAS_ROPE_POS: tl.constexpr, + CACHE_IS_FP16: tl.constexpr, + MROPE_H: tl.constexpr, + MROPE_W: tl.constexpr, +): + pid = tl.program_id(0) + # K work occupies the first programs; the remaining programs tile Q. This + # keeps both paths in one launch while leaving their register shapes + # independent. + if pid >= num_k_work: + q_pid = pid - num_k_work + num_head_tiles: tl.constexpr = tl.cdiv(HQ, TILE_H_Q) + token_tile = q_pid // num_head_tiles + head_tile = q_pid % num_head_tiles + tokens = token_tile * TILE_T_Q + tl.arange(0, TILE_T_Q) + heads = head_tile * TILE_H_Q + tl.arange(0, TILE_H_Q) + valid_tokens = tokens < num_tokens + valid_heads = heads < HQ + dims = tl.arange(0, D) + mask = valid_tokens[:, None, None] & valid_heads[None, :, None] + x = tl.load( + q_ptr + + tokens[:, None, None] * q_stride_token + + heads[None, :, None] * D + + dims[None, None, :], + mask=mask, + other=0.0, + ) + pos_t = tl.load(pos_ptr + tokens * pos_stride_token, mask=valid_tokens, other=0) + if IS_2D_POSITIONS: + pos_h = tl.load( + pos_ptr + pos_stride_axis + tokens * pos_stride_token, + mask=valid_tokens, + other=0, + ) + pos_w = tl.load( + pos_ptr + 2 * pos_stride_axis + tokens * pos_stride_token, + mask=valid_tokens, + other=0, + ) + else: + pos_h = pos_t + pos_w = pos_t + y = _norm_rope( + x, + pos_t, + pos_h, + pos_w, + cos_sin_ptr, + D // 2, + q_norm_weight_ptr, + eps, + IS_2D_POSITIONS, + MROPE_H, + MROPE_W, + ) + tl.store( + q_out_ptr + + tokens[:, None, None] * q_out_stride_token + + heads[None, :, None] * q_out_stride_head + + dims[None, None, :], + y, + mask=mask, + ) + + if pid < num_k_work: + # One work item owns one completed compression group. Work item zero + # additionally commits the request's current raw-K suffix below. + work_metadata = tl.load(k_work_metadata_ptr + pid * 2 + tl.arange(0, 2)) + request, work_in_request = tl.split(work_metadata) + if request < 0: + return + + query_start = tl.load(query_start_loc_ptr + request) + query_end = tl.load(query_start_loc_ptr + request + 1) + query_len = query_end - query_start + chunk_end = tl.load(logical_positions_ptr + query_end - 1) + chunk_start = chunk_end - query_len + 1 + num_groups = (chunk_end + 1) // COMPRESS_RATIO - chunk_start // COMPRESS_RATIO + dims = tl.arange(0, D) + + if work_in_request < num_groups: + first_boundary = ( + (chunk_start + COMPRESS_RATIO) // COMPRESS_RATIO + ) * COMPRESS_RATIO - 1 + end_position = first_boundary + work_in_request * COMPRESS_RATIO + boundary_token = query_start + end_position - chunk_start + valid_token = ( + (boundary_token >= query_start) + & (boundary_token < query_end) + & (boundary_token < num_tokens) + ) + compressed_slot = tl.load( + compressed_slots_ptr + boundary_token, + mask=valid_token, + other=-1, + ) + valid = ( + valid_token + & (compressed_slot >= 0) + & (compressed_slot < num_compressed_blocks * COMP_PAGE_SIZE) + ) + state_block = tl.load(state_table_ptr + request * state_table_stride_req) + state_block_valid = (state_block >= 0) & (state_block < num_state_blocks) + safe_state_block = tl.maximum(state_block, 0).to(tl.int64) + group_offsets = tl.arange(0, COMPRESS_RATIO) + source_positions = end_position - (COMPRESS_RATIO - 1) + group_offsets + source_in_chunk = source_positions >= chunk_start + source_tokens = query_start + source_positions - chunk_start + source_tokens_valid = ( + (source_tokens >= query_start) + & (source_tokens < query_end) + & (source_tokens < num_tokens) + ) + current_base = ( + k_ptr + tl.maximum(source_tokens, 0)[:, None] * k_stride_token + ) + cached_base = ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (source_positions % STATE_SIZE)[:, None] * state_cache_stride_token + ) + # Only the first completed group can cross the chunk boundary. Select + # historical rows from the ring without issuing two masked loads. + source_base = tl.where(source_in_chunk[:, None], current_base, cached_base) + # Pointer selection obscures alignment from Triton's analysis. + source_base = tl.multiple_of(source_base, (8, 8)) + source_valid = tl.where( + source_in_chunk, source_tokens_valid, state_block_valid + ) + source = tl.load( + source_base + dims[None, :], + mask=valid & source_valid[:, None], + other=0.0, + ).to(tl.float32) + # Match the unfused path's storage-dtype pooled tensor before + # RMSNorm. V100 executes this route in FP16. + pooled = tl.sum(source, axis=0) / COMPRESS_RATIO + if CACHE_IS_FP16: + pooled = pooled.to(tl.float16).to(tl.float32) + else: + pooled = pooled.to(tl.bfloat16).to(tl.float32) + + first_position = end_position - (COMPRESS_RATIO - 1) + if CACHE_HAS_ROPE_POS: + # RoPE uses the first token in the pooled group. Its exact MRoPE + # coordinates may live in this chunk or the raw-state ring. + first_in_chunk = first_position >= chunk_start + first_token = query_start + first_position - chunk_start + first_token_valid = ( + (first_token >= query_start) + & (first_token < query_end) + & (first_token < num_tokens) + ) + safe_first_token = tl.maximum(first_token, 0) + load_current_position = first_in_chunk & first_token_valid + first_pos_t = tl.load( + pos_ptr + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + if IS_2D_POSITIONS: + first_pos_h = tl.load( + pos_ptr + pos_stride_axis + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + first_pos_w = tl.load( + pos_ptr + + 2 * pos_stride_axis + + safe_first_token * pos_stride_token, + mask=load_current_position, + other=0, + ) + else: + first_pos_h = first_pos_t + first_pos_w = first_pos_t + tail = ( + state_cache_ptr + + safe_state_block * state_cache_stride_block + + (first_position % STATE_SIZE) * state_cache_stride_token + + D + ).to(tl.pointer_type(tl.int64)) + load_cached_position = ~first_in_chunk & state_block_valid + cached_pos_t = tl.load(tail, mask=load_cached_position, other=0) + cached_pos_h = tl.load(tail + 1, mask=load_cached_position, other=0) + cached_pos_w = tl.load(tail + 2, mask=load_cached_position, other=0) + pos_t = tl.where(first_in_chunk, first_pos_t.to(tl.int64), cached_pos_t) + pos_h = tl.where(first_in_chunk, first_pos_h.to(tl.int64), cached_pos_h) + pos_w = tl.where(first_in_chunk, first_pos_w.to(tl.int64), cached_pos_w) + else: + pos_t = first_position + pos_h = first_position + pos_w = first_position + y = _norm_rope( + tl.reshape(pooled, (1, 1, D)), + pos_t + tl.arange(0, 1), + pos_h + tl.arange(0, 1), + pos_w + tl.arange(0, 1), + cos_sin_ptr, + D // 2, + k_norm_weight_ptr, + eps, + IS_K_MROPE, + MROPE_H, + MROPE_W, + ) + compressed_block = (compressed_slot // COMP_PAGE_SIZE).to(tl.int64) + compressed_row = compressed_slot % COMP_PAGE_SIZE + tl.store( + compressed_cache_ptr + + compressed_block * compressed_cache_stride_block + + compressed_row * compressed_cache_stride_token + + dims, + tl.reshape(y, (D,)), + mask=valid, + ) + + if work_in_request == 0: + # This CTA may have just read historical rows from the circular buffer. + # Keep every lane past those loads before overwriting the same ring. + tl.debug_barrier() + # One CTA per request commits only the suffix retained by the ring. + num_state_rows = tl.minimum(query_len, STATE_SIZE) + for state_offset in tl.range(0, num_state_rows): + token = query_end - num_state_rows + state_offset + valid_token = ( + (token >= query_start) & (token < query_end) & (token < num_tokens) + ) + slot = tl.load(state_slots_ptr + token, mask=valid_token, other=-1) + valid_slot = ( + valid_token & (slot >= 0) & (slot < num_state_blocks * STATE_SIZE) + ) + safe_slot = tl.maximum(slot, 0) + state_row = ( + state_cache_ptr + + (safe_slot // STATE_SIZE).to(tl.int64) * state_cache_stride_block + + (safe_slot % STATE_SIZE) * state_cache_stride_token + ) + k = tl.load( + k_ptr + tl.maximum(token, 0) * k_stride_token + dims, + mask=valid_slot, + other=0.0, + ) + tl.store(state_row + dims, k, mask=valid_slot) + if CACHE_HAS_ROPE_POS: + pos_t = tl.load( + pos_ptr + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + if IS_2D_POSITIONS: + pos_h = tl.load( + pos_ptr + + pos_stride_axis + + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + pos_w = tl.load( + pos_ptr + + 2 * pos_stride_axis + + tl.maximum(token, 0) * pos_stride_token, + mask=valid_slot, + other=0, + ) + else: + pos_h = pos_t + pos_w = pos_t + tail = (state_row + D).to(tl.pointer_type(tl.int64)) + tl.store(tail, pos_t.to(tl.int64), mask=valid_slot) + tl.store(tail + 1, pos_h.to(tl.int64), mask=valid_slot) + tl.store(tail + 2, pos_w.to(tl.int64), mask=valid_slot) + + +def qsa_pre_indexer( + q: torch.Tensor, + k: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + eps: float, + q_out: torch.Tensor, + state_cache: torch.Tensor, + state_slots: torch.Tensor, + state_block_table: torch.Tensor, + query_start_loc: torch.Tensor, + logical_positions: torch.Tensor, + compressed_cache: torch.Tensor, + compressed_slots: torch.Tensor, + k_work_metadata: torch.Tensor, + *, + compress_ratio: int, + mrope_section: tuple[int, int, int] | None, + rope_pos_offset: int | None, +) -> None: + """Normalize Q, compress K, then update the circular raw state.""" + num_tokens = q.shape[0] + if num_tokens == 0: + return + num_q_heads, head_dim = q_out.shape[1:] + assert cos_sin_cache.shape[-1] * 2 == head_dim + assert q.shape == (num_tokens, num_q_heads * head_dim) + assert k.shape == (num_tokens, head_dim) + assert q.stride(-1) == 1 + assert k.stride(-1) == 1 + assert q.dtype == k.dtype == q_out.dtype + assert q.dtype == state_cache.dtype == compressed_cache.dtype + assert q.dtype in (torch.float16, torch.bfloat16) + assert q_out.stride(-1) == 1 + assert cos_sin_cache.is_contiguous() + assert state_cache.stride(-1) == 1 + assert compressed_cache.stride(-1) == 1 + assert k_work_metadata.ndim == 2 and k_work_metadata.shape[1] == 2 + is_2d_positions = positions.ndim == 2 + is_k_mrope = bool(mrope_section) + cache_has_rope_pos = rope_pos_offset is not None + assert rope_pos_offset is None or rope_pos_offset == head_dim + if is_2d_positions: + assert positions.shape == (3, num_tokens) + assert is_k_mrope + pos_stride_axis, pos_stride_token = positions.stride() + else: + assert positions.shape == (num_tokens,) + pos_stride_axis, pos_stride_token = 0, positions.stride(0) + section = mrope_section if mrope_section is not None else (0, 0, 0) + assert len(section) == 3 + + if num_tokens <= 4096: + TILE_T_Q, TILE_H_Q = 2, 2 + else: + TILE_T_Q, TILE_H_Q = 2, 4 + num_k_work = k_work_metadata.shape[0] + num_q_work = triton.cdiv(num_tokens, TILE_T_Q) * triton.cdiv(num_q_heads, TILE_H_Q) + _qsa_pre_indexer_kernel[(num_k_work + num_q_work,)]( + q, + q.stride(0), + k, + k.stride(0), + positions, + pos_stride_axis, + pos_stride_token, + cos_sin_cache, + q_norm_weight, + k_norm_weight, + eps, + q_out, + q_out.stride(0), + q_out.stride(1), + state_cache, + state_cache.stride(0), + state_cache.stride(1), + state_slots, + state_block_table, + state_block_table.stride(0), + query_start_loc, + logical_positions, + compressed_slots, + k_work_metadata, + compressed_cache, + compressed_cache.stride(0), + compressed_cache.stride(1), + num_tokens, + state_cache.shape[0], + compressed_cache.shape[0], + num_k_work, + HQ=num_q_heads, + D=head_dim, + TILE_T_Q=TILE_T_Q, + TILE_H_Q=TILE_H_Q, + COMPRESS_RATIO=compress_ratio, + STATE_SIZE=state_cache.shape[1], + COMP_PAGE_SIZE=compressed_cache.shape[1], + IS_2D_POSITIONS=is_2d_positions, + IS_K_MROPE=is_k_mrope, + CACHE_HAS_ROPE_POS=cache_has_rope_pos, + CACHE_IS_FP16=q.dtype == torch.float16, + MROPE_H=section[1], + MROPE_W=section[2], + num_warps=1, + ) + + +__all__ = ["qsa_pre_indexer"] diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py new file mode 100644 index 0000000000..b4549add56 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -0,0 +1,1310 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen4Exp position-learning enhancement layers.""" + +import math +from collections.abc import Iterable, Sequence + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.mamba.abstract import MambaBase +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateDtypeCalculator, + MambaStateShapeCalculator, + is_conv_state_dim_first, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + create_fp8_scale_parameter, + create_fp8_weight_parameter, + is_fp8, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.parameter import ( + ModelWeightParameter, + PerTensorScaleParameter, +) +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) +from vllm.utils.mem_utils import format_gib +from vllm.utils.platform_utils import is_pin_memory_available +from vllm.utils.torch_utils import ( + direct_register_custom_op, + get_accelerator_view_from_cpu_tensor, +) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionBackend, + PleShortConvAttentionMetadata, +) +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +from ..common.ple import copy_ple_embedding_shard_ + +_MASK64 = (1 << 64) - 1 +_SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 +_SPLITMIX_M1 = 0xBF58476D1CE4E5B9 +_SPLITMIX_M2 = 0x94D049BB133111EB +_PLE_LAYER_PRIME = 10007 + +logger = init_logger(__name__) + + +def _splitmix64(value: int) -> int: + value = (value + _SPLITMIX_GAMMA) & _MASK64 + value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64 + value = ((value ^ (value >> 27)) * _SPLITMIX_M2) & _MASK64 + return (value ^ (value >> 31)) & _MASK64 + + +def _is_prime_64(value: int) -> bool: + if value < 2: + return False + for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): + if value % prime == 0: + return value == prime + exponent = value - 1 + shifts = 0 + while exponent % 2 == 0: + exponent //= 2 + shifts += 1 + for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): + if base % value == 0: + continue + witness = pow(base, exponent, value) + if witness in (1, value - 1): + continue + for _ in range(shifts - 1): + witness = pow(witness, 2, value) + if witness == value - 1: + break + else: + return False + return True + + +def _nth_prime_after(start: int, count: int) -> int: + prime = int(start) + for _ in range(count): + candidate = prime + 1 + if candidate <= 2: + prime = 2 + continue + if candidate % 2 == 0: + candidate += 1 + while not _is_prime_64(candidate): + candidate += 2 + prime = candidate + return prime + + +class Qwen4ExpPLEGroupedNorm(nn.Module): + def __init__( + self, + hidden_size: int, + eps: float, + group_size: int | None, + dtype: torch.dtype | None, + ) -> None: + super().__init__() + if group_size is not None and hidden_size % group_size: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by " + f"group_size ({group_size})" + ) + self.eps = eps + self.group_size = group_size + self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.float() + if self.group_size is None: + variance = hidden_states.square().mean(dim=-1, keepdim=True) + normalized = hidden_states * torch.rsqrt(variance + self.eps) + else: + grouped = hidden_states.unflatten( + -1, (hidden_states.shape[-1] // self.group_size, self.group_size) + ) + variance = grouped.square().mean(dim=-1, keepdim=True) + normalized = (grouped * torch.rsqrt(variance + self.eps)).flatten(-2) + return (normalized * (1.0 + self.weight.float())).to(input_dtype) + + +class Qwen4ExpPLEFp8EmbeddingMethod(QuantizeMethodBase): + """FP8 PLE embedding with one global checkpoint scale.""" + + def create_weights( + self, + layer: nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size, params_dtype + weight_loader = extra_weight_attrs.get("weight_loader") + weight = create_fp8_weight_parameter( + sum(output_partition_sizes), input_size_per_partition, weight_loader + ) + layer.register_parameter("weight", weight) + + weight_scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + scale_dtype=torch.bfloat16, + ) + layer.register_parameter("weight_scale", weight_scale) + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + raise NotImplementedError("PLE FP8 weights only support embedding lookup") + + def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: + get_accelerator_weight = getattr(layer, "get_accelerator_weight", None) + weight = ( + get_accelerator_weight(input_.device) + if get_accelerator_weight is not None + else layer.weight + ) + return F.embedding(input_, weight) + + def process_weights_after_loading(self, layer: nn.Module) -> None: + prepare_accelerator_weight = getattr(layer, "prepare_accelerator_weight", None) + if prepare_accelerator_weight is not None: + prepare_accelerator_weight() + + +def _get_ple_embedding_quant_method( + quant_config: QuantizationConfig | None, + prefix: str, + *, + force_fp8_storage: bool = False, +) -> QuantizeMethodBase | None: + """Select global-scale FP8 only for quantized PLE checkpoint shards.""" + + if force_fp8_storage: + return Qwen4ExpPLEFp8EmbeddingMethod() + if not isinstance(quant_config, Fp8Config): + return None + if not quant_config.is_checkpoint_fp8_serialized: + return None + + ignored_layers = quant_config.ignored_layers + if is_layer_skipped( + prefix, + ignored_layers, + quant_config.packed_modules_mapping, + ): + return None + # PLE checkpoint shards form one runtime embedding parameter. + shard_prefix = f"{prefix}.shard_" + if any(name.startswith(shard_prefix) for name in ignored_layers): + return None + return Qwen4ExpPLEFp8EmbeddingMethod() + + +def _should_use_pinned_host_ple(config: Qwen4ExpTextConfig) -> bool: + explicit = getattr(config, "ple_offload_embedding", None) + if explicit is not None: + return bool(explicit) + capability = current_platform.get_device_capability() + return capability is not None and capability.to_int() == 70 + + +class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding): + """TP-sharded FP8 PLE table backed directly by pinned host memory. + + The base embedding is constructed on the meta device, so the full shard is + never allocated on a GPU. Checkpoint shards copy directly into the pinned + CPU parameter. A stable UVA view is created after loading and used only for + embedding gathers. + """ + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + params_dtype: torch.dtype | None, + padding_size: int, + prefix: str, + quant_method: QuantizeMethodBase, + ) -> None: + if not is_pin_memory_available(): + raise RuntimeError("Qwen4Exp PLE host offload requires pinned host memory") + if not isinstance(quant_method, Qwen4ExpPLEFp8EmbeddingMethod): + raise NotImplementedError( + "Qwen4Exp pinned-host PLE currently requires FP8 checkpoint storage" + ) + + with torch.device("meta"): + super().__init__( + num_embeddings, + embedding_dim, + params_dtype=params_dtype, + padding_size=padding_size, + prefix=prefix, + quant_method=quant_method, + ) + + host_weight = ModelWeightParameter( + data=torch.empty( + tuple(self.weight.shape), + dtype=self.weight.dtype, + device="cpu", + pin_memory=True, + ), + input_dim=1, + output_dim=0, + weight_loader=self.weight_loader, + ) + host_weight._vllm_keep_on_cpu = True + self.weight = host_weight + self.weight_scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + [self.num_embeddings_per_partition], + self.embedding_dim, + None, + self.weight_loader, + scale_dtype=torch.bfloat16, + ) + self._accelerator_weight_views: dict[int, torch.Tensor] = {} + logger.info( + "Qwen4Exp PLE shard allocated in pinned host memory: %s", + format_gib(self.weight.numel() * self.weight.element_size()), + ) + + def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: + if device.type != "cuda": + raise RuntimeError( + f"Qwen4Exp pinned-host PLE requires a CUDA input, got {device}" + ) + device_index = ( + torch.cuda.current_device() if device.index is None else device.index + ) + view = self._accelerator_weight_views.get(device_index) + if view is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Qwen4Exp PLE UVA view must be prepared before CUDA graph capture" + ) + with torch.cuda.device(device_index): + view = get_accelerator_view_from_cpu_tensor(self.weight) + self._accelerator_weight_views[device_index] = view + return view + + def prepare_accelerator_weight(self) -> None: + self.get_accelerator_weight(torch.device("cuda", torch.cuda.current_device())) + + +class Qwen4ExpNGramEmbedding(nn.Module): + def __init__( + self, + config: Qwen4ExpTextConfig, + embedding_dim: int, + ple_dense_layer_id: int, + max_total_tokens: int, + max_num_reqs: int, + prefix: str, + quant_config: QuantizationConfig | None = None, + params_dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + self.embedding_dim = embedding_dim + self.ngram_size = int(config.ngram_size) + self.heads_per_ngram = int(config.heads_per_ngram) + self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram + if self.ngram_size < 2: + raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") + if self.heads_per_ngram <= 0: + raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") + if embedding_dim % self.ngram_heads: + raise ValueError( + "ple_embed_dim must be divisible by total ngram heads: " + f"{embedding_dim} % {self.ngram_heads} != 0" + ) + self.head_dim = embedding_dim // self.ngram_heads + self.eos_token_id = int(config.eos_token_id) + self.unigram_vocab_size = int(config.vocab_size) + self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) + if self.split_ngram_parts <= 0: + raise ValueError("split_ngram_parts must be positive") + + max_multiplier = ((1 << 63) - 1) // self.unigram_vocab_size + half_bound = max(1, max_multiplier // 2) + seed = int(getattr(config, "seed", None) or 1234) + base_seed = seed + _PLE_LAYER_PRIME * ple_dense_layer_id + multipliers = [] + for index in range(self.ngram_size): + value = base_seed + _SPLITMIX_GAMMA * (index + 1) + multipliers.append(2 * (_splitmix64(value) % half_bound) + 1) + self.register_buffer( + "layer_multipliers", + torch.tensor(multipliers, dtype=torch.long), + persistent=True, + ) + + ngram_vocab_size_base = int(config.ngram_vocab_size_base) + sizes: list[int] = [] + offsets: list[int] = [] + offset = 0 + for local_head in range(self.ngram_heads): + global_head = ple_dense_layer_id * self.ngram_heads + local_head + size = _nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) + sizes.append(size) + offsets.append(offset) + offset += size + self.register_buffer( + "ngram_heads_vocab_sizes", + torch.tensor(sizes, dtype=torch.long), + persistent=True, + ) + self.register_buffer( + "ngram_heads_offsets", + torch.tensor(offsets, dtype=torch.long), + persistent=True, + ) + divisor = int(config.make_ngram_vocab_size_divisible_by) + padded_vocab_size = ((offset + divisor - 1) // divisor) * divisor + embedding_prefix = f"{prefix}.ngram_embedding" + ple_storage_dtype = str( + getattr(config, "ple_embedding_dtype", "") + ).removeprefix("torch.") + quant_method = _get_ple_embedding_quant_method( + quant_config, + embedding_prefix, + force_fp8_storage=ple_storage_dtype == "float8_e4m3fn", + ) + if _should_use_pinned_host_ple(config): + if quant_method is None: + raise NotImplementedError( + "Qwen4Exp pinned-host PLE requires FP8 checkpoint storage" + ) + self.ngram_embedding = Qwen4ExpPinnedHostEmbedding( + padded_vocab_size, + self.head_dim, + params_dtype=params_dtype, + padding_size=divisor, + prefix=embedding_prefix, + quant_method=quant_method, + ) + else: + self.ngram_embedding = VocabParallelEmbedding( + padded_vocab_size, + self.head_dim, + params_dtype=params_dtype, + padding_size=divisor, + prefix=embedding_prefix, + quant_method=quant_method, + ) + self.register_buffer( + "positions_buffer", + torch.arange(max_total_tokens, dtype=torch.int64), + persistent=False, + ) + self.register_buffer( + "padded_buffer", + torch.full( + (max_num_reqs, max_total_tokens), + self.eos_token_id, + dtype=torch.int64, + ), + persistent=False, + ) + + @staticmethod + def _shift_precompute( + tokens: torch.Tensor, eos_token_id: int + ) -> tuple[torch.Tensor, torch.Tensor]: + if tokens.dim() != 2: + raise ValueError("tokens must be a 2D tensor") + batch_size, seq_len = tokens.shape + positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) + eos_positions = torch.where(tokens == eos_token_id, positions, -1) + previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values + previous_eos = torch.cat( + [ + eos_positions.new_full((batch_size, 1), -1), + previous_eos_inclusive[:, :-1], + ], + dim=1, + ) + return positions, positions.unsqueeze(0) - previous_eos - 1 + + @staticmethod + def _shift_apply( + tokens: torch.Tensor, + positions: torch.Tensor, + position_in_segment: torch.Tensor, + shift: int, + eos_token_id: int, + ) -> torch.Tensor: + if shift == 0: + return tokens + source = positions - shift + gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) + shifted = tokens.gather(1, gather_indices) + valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) + return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) + + def forward( + self, + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + ) -> torch.Tensor: + input_ids = input_ids.reshape(-1).long() + query_start_loc = query_start_loc.long() + num_reqs = query_start_loc.numel() - 1 + num_tokens = input_ids.shape[0] + if num_tokens > self.positions_buffer.numel(): + raise ValueError( + f"PLE received {num_tokens} tokens, but its workspace supports " + f"at most {self.positions_buffer.numel()}" + ) + if num_reqs > self.padded_buffer.shape[0]: + raise ValueError( + f"PLE received {num_reqs} requests, but its workspace supports " + f"at most {self.padded_buffer.shape[0]}" + ) + + positions = self.positions_buffer[:num_tokens] + packed = self.padded_buffer[:num_reqs] + packed.fill_(self.eos_token_id) + request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 + request_indices.clamp_(max=num_reqs - 1) + columns = (positions - query_start_loc[request_indices]).clamp( + 0, packed.shape[1] - 1 + ) + packed[request_indices, columns] = input_ids + ngram_context = ngram_context[:num_reqs].to( + device=input_ids.device, dtype=torch.long + ) + + context = torch.cat([ngram_context, packed], dim=-1) + positions_2d, position_in_segment = self._shift_precompute( + context, self.eos_token_id + ) + shifted = [context] + for shift in range(1, self.ngram_size): + shifted.append( + self._shift_apply( + context, + positions_2d, + position_in_segment, + shift, + self.eos_token_id, + ) + ) + adjusted_columns = columns + self.ngram_size - 1 + id_blocks = [] + for ngram in range(2, self.ngram_size + 1): + start = (ngram - 2) * self.heads_per_ngram + end = start + self.heads_per_ngram + mixed = shifted[0] * self.layer_multipliers[0] + for index in range(1, ngram): + mixed = torch.bitwise_xor( + mixed, shifted[index] * self.layer_multipliers[index] + ) + sizes = self.ngram_heads_vocab_sizes[start:end] + offsets = self.ngram_heads_offsets[start:end] + ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets + id_blocks.append(ids[request_indices, adjusted_columns]) + ngram_ids = torch.cat(id_blocks, dim=-1) + return self.ngram_embedding(ngram_ids).flatten(-2) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load hash buffers and checkpoint-split embedding rows.""" + + persistent_buffers = { + "layer_multipliers": self.layer_multipliers, + "ngram_heads_offsets": self.ngram_heads_offsets, + "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, + } + loaded: set[str] = set() + regular_weights: list[tuple[str, torch.Tensor]] = [] + shard_prefix = "ngram_embedding.shard_" + + for name, loaded_weight in weights: + leaf_name = name.rsplit(".", 1)[-1] + if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": + continue + if name in persistent_buffers: + buffer = persistent_buffers[name] + if buffer.shape != loaded_weight.shape: + raise ValueError( + f"Shape mismatch for {name}: expected " + f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" + ) + buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) + loaded.add(name) + continue + if name.startswith(shard_prefix) and name.endswith(".weight"): + shard_text = name[len(shard_prefix) : -len(".weight")] + if not shard_text.isdigit(): + regular_weights.append((name, loaded_weight)) + continue + shard_index = int(shard_text) + if shard_index >= self.split_ngram_parts: + raise ValueError( + f"PLE embedding shard index {shard_index} exceeds " + f"split_ngram_parts={self.split_ngram_parts}" + ) + embedding = self.ngram_embedding + shard_size = ( + embedding.org_vocab_size + self.split_ngram_parts - 1 + ) // self.split_ngram_parts + checkpoint_start = shard_index * shard_size + expected_rows = max( + 0, + min(shard_size, embedding.org_vocab_size - checkpoint_start), + ) + expected_shape = (expected_rows, embedding.embedding_dim) + if tuple(loaded_weight.shape) != expected_shape: + raise ValueError( + f"Shape mismatch for PLE embedding shard {shard_index}: " + f"expected {expected_shape}, got " + f"{tuple(loaded_weight.shape)}" + ) + copy_ple_embedding_shard_( + embedding.weight.data, + loaded_weight, + checkpoint_start=checkpoint_start, + tp_start=embedding.shard_indices.org_vocab_start_index, + tp_end=embedding.shard_indices.org_vocab_end_index, + ) + loaded.add("ngram_embedding.weight") + continue + regular_weights.append((name, loaded_weight)) + + if regular_weights: + loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) + return loaded + + +class Qwen4ExpPLELayer(nn.Module, MambaBase): + def __init__( + self, + config: Qwen4ExpTextConfig, + vllm_config: VllmConfig, + layer_idx: int = 0, + ple_dense_layer_id: int | None = None, + prefix: str = "", + ) -> None: + super().__init__() + model_config = vllm_config.model_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.model_config: ModelConfig = model_config + self.cache_config: CacheConfig = cache_config + self.layer_idx = layer_idx + self.ple_dense_layer_id = ( + int(ple_dense_layer_id) + if ple_dense_layer_id is not None + else int(layer_idx) + ) + self.prefix = prefix + self.hidden_size = int(config.hidden_size) + self.hc_count = config.hc_count + self.hc_hidden_size = self.hidden_size * self.hc_count + self.conv_kernel_size = int(config.ple_conv_kernel_size) + self.short_conv_dilation = int(config.ngram_size) + self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation + self.num_spec_tokens = vllm_config.num_speculative_tokens + self.activation = "silu" + self.ple_embedding: nn.Module = Qwen4ExpNGramEmbedding( + config, + int(config.ple_embed_dim), + self.ple_dense_layer_id, + vllm_config.scheduler_config.max_num_batched_tokens, + vllm_config.scheduler_config.max_num_seqs, + f"{prefix}.ple_embedding", + quant_config=quant_config, + params_dtype=model_config.dtype, + ) + self.key_proj = ReplicatedLinear( + int(config.ple_embed_dim), + self.hc_hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.key_proj", + ) + self.value_proj = ReplicatedLinear( + int(config.ple_embed_dim), + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.value_proj", + ) + norm_args = ( + self.hc_hidden_size, + config.rms_norm_eps, + self.hidden_size, + model_config.dtype, + ) + self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args) + self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args) + self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args) + self.conv1d = nn.Conv1d( + self.hc_hidden_size, + self.hc_hidden_size, + self.conv_kernel_size, + groups=self.hc_hidden_size, + padding=self.conv_state_len, + dilation=self.short_conv_dilation, + bias=False, + dtype=model_config.dtype, + ) + nn.init.zeros_(self.conv1d.weight) + self.conv1d.weight._no_reinit = True + self.kv_cache = (torch.tensor([]),) + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _get_embedding_weight_scale(self) -> torch.Tensor | None: + embedding = getattr(self.ple_embedding, "ngram_embedding", None) + return getattr(embedding, "weight_scale", None) + + def _dequantize_embeddings( + self, + embeddings: torch.Tensor, + output_dtype: torch.dtype, + ) -> torch.Tensor: + """Dequantize PLE lookup output.""" + + if not is_fp8(embeddings): + return embeddings + weight_scale = self._get_embedding_weight_scale() + if weight_scale is None: + raise RuntimeError("FP8 PLE embedding is missing its global scale") + if weight_scale.device != embeddings.device: + raise RuntimeError("FP8 PLE embedding scale must be on the output device") + return embeddings.to(output_dtype) * weight_scale.to(output_dtype) + + @property + def mamba_type(self) -> MambaAttentionBackendEnum: + return MambaAttentionBackendEnum.SHORT_CONV + + @property + def is_kv_cache_tp_replicated(self) -> bool: + return True + + def get_attn_backend(self) -> type[PleShortConvAttentionBackend]: + return PleShortConvAttentionBackend + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: + return MambaStateDtypeCalculator.short_conv_state_dtype( + self.model_config.dtype, self.cache_config.mamba_cache_dtype + ) + + def get_state_shape(self) -> Sequence[tuple[int, ...]]: + return MambaStateShapeCalculator.short_conv_state_shape( + tp_world_size=1, + intermediate_size=self.hc_hidden_size, + conv_kernel=self.conv_state_len + 1, + num_spec=self.num_spec_tokens, + ) + + def _apply_norm( + self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor + ) -> torch.Tensor: + shape = hidden_states.shape + return norm(hidden_states.flatten(-2)).reshape(shape) + + def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor: + # Profiling / CUDA graph capture only; conv state is not updated. + inputs_t = inputs.transpose(0, 1).unsqueeze(0) + output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)] + return F.silu(output).squeeze(0).transpose(0, 1) + + def _short_conv_dilated_decode_batched( + self, + x_d: torch.Tensor, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + state_indices_tensor_d: torch.Tensor, + has_initial_states_d: torch.Tensor | None, + ) -> torch.Tensor: + state_indices = state_indices_tensor_d.to( + device=conv_state.device, dtype=torch.int64 + ) + # TODO: need double-check + # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to + # slot 0 for a safe gather, then zero output and skip write-back. + valid_state = state_indices != NULL_BLOCK_ID + state_indices = torch.where( + valid_state, state_indices, torch.zeros_like(state_indices) + ) + if has_initial_states_d is None: + has_initial_state = valid_state + else: + if has_initial_states_d.numel() < state_indices_tensor_d.numel(): + raise ValueError( + "has_initial_states_d size mismatch: " + f"got {has_initial_states_d.numel()}, " + f"need >= {state_indices_tensor_d.numel()}." + ) + has_initial_state = has_initial_states_d[ + : state_indices_tensor_d.numel() + ].to(device=conv_state.device, dtype=torch.bool) + has_initial_state = has_initial_state & valid_state + + cached_state = conv_state.index_select(0, state_indices) + state = cached_state[..., : self.conv_state_len].to(x_d.dtype) + if self.conv_state_len > 0: + initial_state = torch.where( + has_initial_state.view(-1, 1, 1), + state, + torch.zeros_like(state), + ) + history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1) + else: + history = x_d.unsqueeze(-1) + + conv_output = F.conv1d( + history, + conv_weights.unsqueeze(1).contiguous(), + groups=history.size(1), + dilation=self.short_conv_dilation, + ).squeeze(-1) + output = F.silu(conv_output) + output = output * valid_state.view(-1, 1).to(output.dtype) + + if self.conv_state_len > 0: + next_state = history[..., -self.conv_state_len :] + # Padded rows are remapped to the reserved null slot. Preserve its + # existing value while writing the new states for valid rows. + existing_base_state = cached_state[..., : self.conv_state_len] + safe_next_state = torch.where( + valid_state.view(-1, 1, 1), + next_state.to(conv_state.dtype), + existing_base_state, + ) + cached_state[..., : self.conv_state_len] = safe_next_state + conv_state.index_copy_(0, state_indices, cached_state) + + return output + + def _short_conv_dilated_prefill_batched( + self, + x_p: torch.Tensor, + metadata: PleShortConvAttentionMetadata, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + state_indices_tensor_p: torch.Tensor, + num_prefills: int, + num_decode_tokens: int, + num_prefill_tokens: int, + ) -> torch.Tensor: + # ``non_spec_query_start_loc`` covers the non-spec (decode + prefill) + # requests and equals ``query_start_loc`` when spec-decode is inactive. + non_spec_query_start_loc = metadata.non_spec_query_start_loc + if non_spec_query_start_loc is None: + raise ValueError("query_start_loc is required for prefill short-conv") + query_start_loc_p = ( + non_spec_query_start_loc[-num_prefills - 1 :] - num_decode_tokens + ) + # The metadata builder guarantees that the prefill query offsets start + # at 0 and end at num_prefill_tokens. Avoid reading those values here, + # since doing so would force a device-to-host synchronization. + has_initial_states_p = metadata.has_initial_states_p + if has_initial_states_p is None: + raise ValueError("has_initial_states_p is required for prefill short-conv") + + output = torch.empty_like(x_p) + q_starts = query_start_loc_p.to(torch.int64) + if state_indices_tensor_p.numel() < num_prefills: + raise ValueError( + "state_indices_tensor_p size mismatch: " + f"got {state_indices_tensor_p.numel()}, " + f"need >= {num_prefills}." + ) + if has_initial_states_p.numel() < num_prefills: + raise ValueError( + "has_initial_states_p size mismatch: " + f"got {has_initial_states_p.numel()}, " + f"need >= {num_prefills}." + ) + if num_prefills == 0 or x_p.numel() == 0: + return output + lengths = q_starts[1:] - q_starts[:-1] + # Use the CPU-computed packing width from the metadata builder instead + # of synchronizing on lengths.max(). + max_len = metadata.max_prefill_query_len + if max_len <= 0: + return output + + hidden_size = x_p.shape[1] + positions = torch.arange( + num_prefill_tokens, device=x_p.device, dtype=torch.int64 + ) + req_indices = torch.searchsorted(q_starts[1:], positions, right=True) + col_indices = positions - q_starts[req_indices] + + packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size)) + packed_tokens[req_indices, col_indices] = x_p + packed_tokens = packed_tokens.transpose(1, 2).contiguous() + + state_indices = state_indices_tensor_p[:num_prefills].to( + device=conv_state.device, dtype=torch.int64 + ) + valid_state = state_indices != NULL_BLOCK_ID + state_indices = torch.where( + valid_state, state_indices, torch.zeros_like(state_indices) + ) + has_initial = has_initial_states_p[:num_prefills].to( + device=conv_state.device, dtype=torch.bool + ) + if self.conv_state_len > 0: + if conv_state.shape[0] == 0: + state = conv_state.new_zeros( + (num_prefills, hidden_size, self.conv_state_len), + dtype=x_p.dtype, + ) + else: + state = conv_state.index_select(0, state_indices)[ + ..., : self.conv_state_len + ].to(x_p.dtype) + use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1) + initial_state = torch.where( + use_initial_mask, + state, + torch.zeros_like(state), + ) + history = torch.cat((initial_state, packed_tokens), dim=-1) + else: + history = packed_tokens + + conv_output = F.conv1d( + history, + conv_weights.unsqueeze(1).contiguous(), + groups=history.size(1), + dilation=self.short_conv_dilation, + ) + conv_output = F.silu(conv_output).transpose(1, 2).contiguous() + + token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64) + valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1) + valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view( + num_prefills, 1 + ) + conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0) + output.copy_(conv_output[req_indices, col_indices]) + + if self.conv_state_len > 0 and conv_state.shape[0] > 0: + state_starts = lengths.to(device=history.device, dtype=torch.int64).view( + num_prefills, 1, 1 + ) + state_offsets = torch.arange( + self.conv_state_len, device=history.device, dtype=torch.int64 + ).view(1, 1, self.conv_state_len) + next_state = history.gather( + dim=2, + index=(state_starts + state_offsets).expand(-1, history.size(1), -1), + ) + # Write back without a host synchronization. Valid, non-empty rows + # receive their new state; padding and zero-length rows keep the + # current cache value. + existing_state = conv_state.index_select(0, state_indices) + existing_base_state = existing_state[..., : self.conv_state_len] + update_mask = valid_state & (lengths.to(device=conv_state.device) > 0) + safe_next_state = torch.where( + update_mask.view(num_prefills, 1, 1), + next_state.to(conv_state.dtype), + existing_base_state, + ) + existing_state[..., : self.conv_state_len] = safe_next_state + conv_state.index_copy_(0, state_indices, existing_state) + return output + + def _short_conv_dilated_spec_batched( + self, + x_spec: torch.Tensor, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + spec_state_indices_tensor: torch.Tensor, + spec_query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + spec_query_len: int, + ) -> torch.Tensor: + """Dilated short-conv for speculative-decode (MTP) requests. + + Each spec request feeds multiple (draft + 1) query tokens. The conv + outputs are computed causally after rolling back the previous draft + state by ``num_accepted_tokens - 1``. The current candidate inputs stay + in the extended cache for the next forward, matching + ``causal_conv1d_update``. + + ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query + length and is a Python int, so no host synchronization is needed; this + keeps the path safe for full CUDA-graph capture/replay where the buffers + are padded at the request level. + """ + num_reqs = spec_state_indices_tensor.numel() + hidden_size = x_spec.size(-1) + # Use a fixed packing width instead of synchronizing on lengths.max(). + max_len = spec_query_len + # Full CUDA graphs can pad these buffers. Only the first num_reqs + # accepted-token counts belong to actual speculative requests. + num_accepted_tokens = num_accepted_tokens[:num_reqs] + q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64) + # Keep the number of real speculative tokens on the device. + total_real_tokens = q_starts[num_reqs] + + state_indices = spec_state_indices_tensor.to( + device=conv_state.device, dtype=torch.int64 + ) + valid_state = state_indices != NULL_BLOCK_ID + state_indices = torch.where( + valid_state, state_indices, torch.zeros_like(state_indices) + ) + positions = torch.arange( + x_spec.size(0), device=x_spec.device, dtype=torch.int64 + ) + # Route graph-padded token rows to the discarded dummy request so that + # they cannot overwrite real packed data. + req_indices = torch.searchsorted(q_starts[1:], positions, right=True) + valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs) + clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0)) + col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1) + pack_req_indices = torch.where( + valid_tokens, + clamped_req_indices, + torch.full_like(req_indices, num_reqs), + ) + pack_col_indices = torch.where( + valid_tokens, col_indices, torch.zeros_like(col_indices) + ) + + # The last request row is the dummy sink for graph padding. + packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size)) + packed[pack_req_indices, pack_col_indices] = x_spec + packed = packed.transpose(1, 2).contiguous() + + if self.conv_state_len > 0: + cached_state = conv_state.index_select(0, state_indices) + rollback_offsets = num_accepted_tokens.to( + device=conv_state.device, dtype=torch.int64 + ).sub(1) + rollback_offsets = torch.where( + valid_state, + rollback_offsets.clamp_(0, max_len - 1), + torch.zeros_like(rollback_offsets), + ) + state_offsets = torch.arange( + self.conv_state_len, device=conv_state.device, dtype=torch.int64 + ).view(1, 1, self.conv_state_len) + rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets + state = cached_state.gather( + 2, rollback_indices.expand(-1, hidden_size, -1) + ).to(x_spec.dtype) + state = torch.where( + valid_state.view(num_reqs, 1, 1), + state, + torch.zeros_like(state), + ) + # Append a zeroed dummy-row state to match the [num_reqs + 1] pack. + dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len)) + state_full = torch.cat((state, dummy_state), dim=0) + history = torch.cat((state_full, packed), dim=-1) + else: + history = packed + + conv_output = F.conv1d( + history, + conv_weights.unsqueeze(1).contiguous(), + groups=history.size(1), + dilation=self.short_conv_dilation, + ) + conv_output = F.silu(conv_output).transpose(1, 2).contiguous() + + output = conv_output[pack_req_indices, pack_col_indices] + output = output * valid_tokens.view(-1, 1).to(output.dtype) + + # Keep all current candidate inputs in the extended state. On the next + # target forward, ``num_accepted_tokens - 1`` selects the rollback + # window before processing the newly scheduled tokens. + if self.conv_state_len > 0: + state_capacity = self.conv_state_len + max_len - 1 + if conv_state.size(-1) < state_capacity: + raise RuntimeError( + "PLE short-conv cache cannot retain speculative tokens: " + f"got {conv_state.size(-1)}, need {state_capacity}." + ) + candidate_state = history[:num_reqs, :, 1 : state_capacity + 1] + query_lengths = q_starts[1:] - q_starts[:-1] + state_positions = torch.arange( + state_capacity, device=history.device, dtype=torch.int64 + ).view(1, 1, state_capacity) + update_lengths = (self.conv_state_len + query_lengths - 1).view( + num_reqs, 1, 1 + ) + update_mask = valid_state.view(num_reqs, 1, 1) & ( + state_positions < update_lengths + ) + existing_state = cached_state[..., :state_capacity] + next_state = torch.where( + update_mask, + candidate_state.to(conv_state.dtype), + existing_state, + ) + cached_state[..., :state_capacity] = next_state + conv_state.index_copy_(0, state_indices, cached_state) + + return output + + def _short_conv_dilated_dispatch( + self, + inputs: torch.Tensor, + metadata: PleShortConvAttentionMetadata, + conv_state: torch.Tensor, + conv_weights: torch.Tensor, + ) -> torch.Tensor: + num_prefills = metadata.num_prefills + num_decodes = metadata.num_decodes + num_decode_tokens = metadata.num_decode_tokens + num_prefill_tokens = metadata.num_prefill_tokens + has_prefill = num_prefills > 0 + has_decode = num_decodes > 0 + has_spec = metadata.spec_sequence_masks is not None + x = inputs[: metadata.num_actual_tokens] + + # Split spec / non-spec tokens. + if has_spec: + if has_prefill or has_decode: + assert metadata.spec_token_indx is not None + assert metadata.non_spec_token_indx is not None + x_spec = x.index_select(0, metadata.spec_token_indx.long()) + x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long()) + else: + x_spec = x + x_non_spec = None + else: + x_spec = None + x_non_spec = x + + spec_output = None + # 1. Run the multi-query speculative-decode part. + if has_spec: + assert metadata.spec_state_indices_tensor is not None + assert metadata.spec_query_start_loc is not None + assert metadata.num_accepted_tokens is not None + spec_output = self._short_conv_dilated_spec_batched( + x_spec=x_spec, + conv_state=conv_state, + conv_weights=conv_weights, + spec_state_indices_tensor=metadata.spec_state_indices_tensor[ + : metadata.num_spec_decodes + ], + spec_query_start_loc=metadata.spec_query_start_loc, + num_accepted_tokens=metadata.num_accepted_tokens, + spec_query_len=metadata.spec_query_len, + ) + + # 2. Run regular decode and prefill requests. + conv_out_non_spec = None + state_indices_tensor = metadata.state_indices_tensor + if x_non_spec is not None: + assert state_indices_tensor is not None + if has_prefill: + state_indices_tensor_d, state_indices_tensor_p = torch.split( + state_indices_tensor, + [num_decodes, num_prefills], + dim=0, + ) + x_d, x_p = torch.split( + x_non_spec, + [num_decode_tokens, num_prefill_tokens], + dim=0, + ) + non_spec_parts: list[torch.Tensor] = [] + if has_decode: + non_spec_parts.append( + self._short_conv_dilated_decode_batched( + x_d=x_d, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices_tensor_d=state_indices_tensor_d, + has_initial_states_d=metadata.has_initial_states_d, + ) + ) + non_spec_parts.append( + self._short_conv_dilated_prefill_batched( + x_p=x_p, + metadata=metadata, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices_tensor_p=state_indices_tensor_p, + num_prefills=num_prefills, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_prefill_tokens, + ) + ) + conv_out_non_spec = torch.vstack(non_spec_parts) + else: + conv_out_non_spec = self._short_conv_dilated_decode_batched( + x_d=x_non_spec, + conv_state=conv_state, + conv_weights=conv_weights, + state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)], + has_initial_states_d=metadata.has_initial_states_d, + ) + + # 3. Merge both parts back into the original token order. + if has_spec and conv_out_non_spec is not None: + assert metadata.spec_token_indx is not None + assert metadata.non_spec_token_indx is not None + assert spec_output is not None + output = x.new_empty((metadata.num_actual_tokens, x.size(-1))) + output.index_copy_(0, metadata.spec_token_indx, spec_output) + output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec) + return output + elif has_spec: + assert spec_output is not None + return spec_output + if conv_out_non_spec is None: + return x + return conv_out_non_spec + + def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + return self._short_conv_fallback(inputs) + + if not isinstance(attn_metadata, dict): + raise RuntimeError( + "PLE short-conv expects per-layer attention metadata dict " + f"during inference, got {type(attn_metadata).__name__}." + ) + + layer_attn_metadata = attn_metadata.get(self.prefix) + if layer_attn_metadata is None: + raise RuntimeError( + f"Missing short-conv metadata for layer '{self.prefix}'. " + "This would bypass conv-state updates and is not allowed." + ) + if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): + raise TypeError( + "Expected PleShortConvAttentionMetadata for layer " + f"'{self.prefix}', got " + f"{type(layer_attn_metadata).__name__}." + ) + + conv_state = self.kv_cache[0] + if not is_conv_state_dim_first(): + conv_state = conv_state.transpose(-1, -2) + conv_weights = self.conv1d.weight.squeeze(1) + + state_capacity = self.conv_state_len + self.num_spec_tokens + if state_capacity > 0: + if conv_state.size(-1) < state_capacity: + raise RuntimeError( + "PLE short-conv cache is smaller than expected for " + f"dilated convolution: got {conv_state.size(-1)}, " + f"expect at least {state_capacity}." + ) + conv_state = conv_state[..., -state_capacity:] + return self._short_conv_dilated_dispatch( + inputs, + layer_attn_metadata, + conv_state, + conv_weights.to(dtype=inputs.dtype), + ) + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + query_start_loc: torch.Tensor, + ngram_context: torch.Tensor, + ) -> torch.Tensor: + input_ids = input_ids.reshape(-1) + if input_ids.shape[0] != hidden_states.shape[0]: + raise ValueError( + "PLE expects input_ids and hidden_states to have the same " + f"token length, got {input_ids.shape[0]} and " + f"{hidden_states.shape[0]}" + ) + embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context) + embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype) + key, _ = self.key_proj(embeddings) + value, _ = self.value_proj(embeddings) + token_count = hidden_states.shape[0] + key = key.reshape(token_count, self.hc_count, self.hidden_size) + query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size) + key = self._apply_norm(self.norm_key, key) + query = self._apply_norm(self.norm_query, query) + gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) + gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt()) + gated_value = gate * value.unsqueeze(-2) + normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) + conv_output = torch.zeros_like(normalized) + torch.ops.vllm.qwen4_exp_ple_short_conv( + normalized, + conv_output, + self.prefix, + ) + return gated_value.flatten(-2) + conv_output + + +def qwen4_exp_ple_short_conv( + inputs: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + layer = get_forward_context().no_compile_layers[layer_name] + result = layer._short_conv(inputs) + output[: result.shape[0]].copy_(result) + + +def qwen4_exp_ple_short_conv_fake( + inputs: torch.Tensor, + output: torch.Tensor, + layer_name: str, +) -> None: + return + + +direct_register_custom_op( + op_name="qwen4_exp_ple_short_conv", + op_func=qwen4_exp_ple_short_conv, + mutates_args=["output"], + fake_impl=qwen4_exp_ple_short_conv_fake, +) + + +__all__ = [ + "Qwen4ExpNGramEmbedding", + "Qwen4ExpPLEGroupedNorm", + "Qwen4ExpPLELayer", +] diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py new file mode 100644 index 0000000000..0bab984eac --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA QSA owner with Triton kernels.""" + +from __future__ import annotations + +from typing import ClassVar, cast + +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.config.cache import CacheDType +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention.attention import ( + set_default_quant_scales, +) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import GemmaRMSNorm +from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding, get_rope +from vllm.model_executor.models.qwen3_next import Qwen3NextAttention +from vllm.platforms import current_platform +from vllm.transformers_utils.configs.qwen4_exp import ( + Qwen4ExpTextConfig, +) +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + canonicalize_singleton_dim_strides, + direct_register_custom_op, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionType, +) +from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available +from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionImpl, + FlashAttentionMetadata, + FlashAttentionMetadataBuilder, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + get_kv_quant_mode, +) + +from ..common.qsa_cache import QSAForwardMetadata +from .indexer_qsa import QSAIndexer + + +class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder): + """Flash metadata supporting uniform decode and target-verify graphs.""" + + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + +class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend): + """FullAttentionSpec backend used by the merged QSA owner.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [ + torch.float16, + torch.bfloat16, + ] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + ] + + @staticmethod + def get_name() -> str: + return "QWEN4_EXP_QSA_TRITON" + + @staticmethod + def get_impl_cls() -> type[Qwen4ExpQSAFlashAttentionImpl]: + return Qwen4ExpQSAFlashAttentionImpl + + @staticmethod + def get_builder_cls() -> type[Qwen4ExpQSAMetadataBuilder]: + return Qwen4ExpQSAMetadataBuilder + + @classmethod + def is_sparse(cls) -> bool: + return True + + @classmethod + def supports_kv_connector(cls) -> bool: + return False + + +class Qwen4ExpQSAFlashAttentionImpl(FlashAttentionImpl): + """Run paged sparse GQA with the QSA Triton kernel.""" + + supports_dcp: bool = False + supports_pcp: bool = False + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + if not is_flash_attn_varlen_func_available(): + raise NotImplementedError("Qwen4Exp QSA requires FlashAttention") + if self.dcp_world_size != 1: + raise NotImplementedError( + "Qwen4Exp QSA does not support decode context parallelism" + ) + if self.kv_cache_dtype not in ("auto", "float16", "bfloat16"): + raise NotImplementedError( + "Qwen4Exp QSA requires an FP16/BF16 main KV cache" + ) + self.supports_quant_query_input = False + + def forward_qsa( + self, + layer: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + kv_cache: torch.Tensor, + attn_metadata: FlashAttentionMetadata, + output: torch.Tensor, + token_to_req: torch.Tensor, + output_scale: torch.Tensor | None = None, + output_block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + del key, value + if output_scale is not None or output_block_scale is not None: + raise NotImplementedError("QSA does not support fused output quantization") + if self.alibi_slopes is not None or self.sinks is not None: + raise NotImplementedError("QSA does not support ALiBi or attention sinks") + if self.sliding_window != (-1, -1): + raise NotImplementedError("QSA does not support sliding-window attention") + + num_tokens = attn_metadata.num_actual_tokens + output.zero_() + if num_tokens == 0: + return output + + topk_buffer = getattr(layer, "topk_indices_buffer", None) + if topk_buffer is None: + raise RuntimeError("QSA owner did not provide its top-k buffer") + logical_indices = topk_buffer[:num_tokens] + token_to_req = token_to_req[:num_tokens] + key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + key_cache = canonicalize_singleton_dim_strides(key_cache) + value_cache = canonicalize_singleton_dim_strides(value_cache) + if key_cache.dtype != query.dtype or query.dtype not in ( + torch.float16, + torch.bfloat16, + ): + raise NotImplementedError("Qwen4Exp QSA requires FP16/BF16 Q/K/V") + + from .ops.qsa import qsa_sparse_paged_attention + + qsa_sparse_paged_attention( + query[:num_tokens], + key_cache, + value_cache, + logical_indices, + attn_metadata.block_table, + token_to_req, + output[:num_tokens], + ) + return output + + +class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): + """Merged Qwen full-attention owner with a QSA index side branch.""" + + supports_dcp = False + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: Qwen4ExpTextConfig, + layer_id: int, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + cache_config = vllm_config.cache_config + model_config = vllm_config.model_config + if cache_config is None: + raise ValueError("Qwen4Exp QSA requires a paged KV cache") + if model_config.dtype not in (torch.float16, torch.bfloat16): + raise NotImplementedError("Qwen4Exp QSA requires FP16 or BF16") + if cache_config.cache_dtype not in ("auto", "float16", "bfloat16"): + raise NotImplementedError( + "Qwen4Exp QSA requires an FP16/BF16 main KV cache" + ) + if getattr(quant_config, "kv_cache_scheme", None) is not None: + raise NotImplementedError("Qwen4Exp QSA does not support KV quantization") + parallel_config = vllm_config.parallel_config + if ( + parallel_config.prefill_context_parallel_size > 1 + or parallel_config.decode_context_parallel_size > 1 + ): + raise NotImplementedError( + "Qwen4Exp QSA does not support context parallelism" + ) + if not getattr(config, "is_causal", True): + raise NotImplementedError("Qwen4Exp QSA requires causal decoder attention") + + self.config = config + self.hidden_size = int(config.hidden_size) + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = int(config.num_attention_heads) + if self.total_num_heads % tp_size: + raise ValueError("QSA attention heads must be divisible by TP size") + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = int(config.num_key_value_heads) + if self.total_num_kv_heads >= tp_size: + if self.total_num_kv_heads % tp_size: + raise ValueError("QSA KV heads must be divisible by TP size") + elif tp_size % self.total_num_kv_heads: + raise ValueError("TP size must be divisible by replicated QSA KV heads") + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = int(config.head_dim or self.hidden_size // self.num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.dual_chunk_attention_config = getattr( + config, "dual_chunk_attention_config", None + ) + if self.dual_chunk_attention_config is not None: + raise NotImplementedError("Qwen4Exp QSA does not support dual-chunk RoPE") + # Qwen4Exp full-attention checkpoints always pack a sigmoid output + # gate next to Q, even when an inherited config default says otherwise. + self.attn_output_gate = True + qkv_quant_config = quant_config + if quant_config is not None and quant_config.get_name() == "modelopt_fp4": + qkv_quant_config = None + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads * (1 + self.attn_output_gate), + self.total_num_kv_heads, + bias=False, + quant_config=qkv_quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + reduce_results=reduce_results, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + self.rotary_emb = get_rope( + head_size=self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters, + ) + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + mm_config = model_config.multimodal_config + text_only = mm_config is None or mm_config.language_model_only + mrope_section = getattr(self.rotary_emb, "mrope_section", None) + supports_mrope = bool( + type(self.rotary_emb) is MRotaryEmbedding + and mrope_section + and len(mrope_section) == 3 + and sum(mrope_section) == self.rotary_emb.rotary_dim // 2 + and getattr(self.rotary_emb, "mrope_interleaved", False) + ) + supports_dtype = getattr(self.rotary_emb, "dtype", None) in ( + torch.float16, + torch.bfloat16, + ) + self.use_fused_qk_norm_rope_gate = ( + self.attn_output_gate + and getattr(self.rotary_emb, "is_neox_style", False) + and current_platform.is_cuda() + and supports_dtype + and (text_only or supports_mrope) + ) + + self.layer_name = f"{prefix}.attn" + self.attn_type = AttentionType.DECODER + self.kv_cache_dtype = cache_config.cache_dtype + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, model_config + ) + if self.kv_cache_torch_dtype != model_config.dtype: + raise NotImplementedError( + "Qwen4Exp QSA main cache dtype must match the model dtype" + ) + self.kv_sharing_target_layer_name = None + self.kv_cache = torch.tensor([]) + set_default_quant_scales(self, register_buffer=True) + + self.attn_backend = Qwen4ExpQSAFlashAttentionBackend + self.impl = Qwen4ExpQSAFlashAttentionImpl( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + None, + None, + self.kv_cache_dtype, + None, + AttentionType.DECODER, + None, + ) + self.indexer = QSAIndexer( + vllm_config=vllm_config, + config=config, + layer_id=layer_id, + rotary_emb=self.rotary_emb, + quant_config=quant_config, + prefix=f"{prefix}.indexer", + ) + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.register_buffer( + "topk_indices_buffer", + torch.empty( + max_tokens, + self.indexer.output_width, + dtype=torch.int32, + ), + persistent=False, + ) + + static_context = vllm_config.compilation_config.static_forward_context + if self.layer_name in static_context: + raise ValueError(f"Duplicate layer name: {self.layer_name}") + static_context[self.layer_name] = self + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.attn_backend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return FullAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + head_size_v=self.head_dim, + dtype=self.kv_cache_torch_dtype, + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + ) + + def _project_qkv_gate( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Split, normalize, and rotate Q/K using this tree's Qwen3 API.""" + q_gate, key, value = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + token_shape = q_gate.shape[:-1] + q_gate = q_gate.view(*token_shape, self.num_heads, 2 * self.head_dim) + query, gate = torch.chunk(q_gate, 2, dim=-1) + query = self.q_norm(query).reshape(*token_shape, self.q_size) + key = self.k_norm( + key.view(*token_shape, self.num_kv_heads, self.head_dim) + ).reshape(*token_shape, self.kv_size) + query, key = self.rotary_emb(positions, query, key) + return query, key, value, gate.reshape(*token_shape, self.q_size) + + def _run_qsa( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + ) -> None: + metadata = get_forward_context().attn_metadata + if isinstance(metadata, list): + metadata = metadata[0] + if not isinstance(metadata, dict): + output.zero_() + return + main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name]) + if self.kv_cache.numel() == 0: + raise RuntimeError("QSA main K/V cache is not bound") + + num_tokens = main_metadata.num_actual_tokens + side_metadata = cast( + QSAForwardMetadata, + metadata[self.indexer.raw_key_cache.prefix], + ) + if side_metadata.num_actual_tokens != num_tokens: + raise RuntimeError("QSA main and side metadata token counts disagree") + selected = self.indexer( + hidden_states, + positions, + self.topk_indices_buffer[:num_tokens], + ) + if selected.shape != ( + num_tokens, + self.indexer.output_width, + ): + raise RuntimeError("QSA indexer returned an invalid selection shape") + impl = cast(Qwen4ExpQSAFlashAttentionImpl, self.impl) + impl.do_kv_cache_update( + self, + key, + value, + self.kv_cache, + main_metadata.slot_mapping, + ) + impl.forward_qsa( + self, + query, + key, + value, + self.kv_cache, + main_metadata, + output, + token_to_req=side_metadata.token_to_req, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v, gate = self._project_qkv_gate(qkv, positions) + num_tokens = hidden_states.shape[0] + query = q.view(num_tokens, self.num_heads, self.head_dim) + key = k.view(num_tokens, self.num_kv_heads, self.head_dim) + value = v.view(num_tokens, self.num_kv_heads, self.head_dim) + attn_output = torch.empty_like(query) + encoded_layer_name = _encode_layer_name(self.layer_name) + if current_platform.opaque_attention_op(): + torch.ops.vllm.qwen4_exp_qsa_with_output( + hidden_states, + positions, + query, + key, + value, + attn_output, + encoded_layer_name, + ) + else: + qwen4_exp_qsa_with_output( + hidden_states, + positions, + query, + key, + value, + attn_output, + encoded_layer_name, + ) + flat_output = attn_output.view(num_tokens, -1) + if gate is not None: + flat_output = flat_output * torch.sigmoid(gate) + output, _ = self.o_proj(flat_output) + return output + + +def qwen4_exp_qsa_with_output( + hidden_states: torch.Tensor, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + layer_name: LayerNameType, +) -> None: + """Run the complete QSA state/update/attend transaction.""" + + layer_name = _resolve_layer_name(layer_name) + layer = get_forward_context().no_compile_layers[layer_name] + if not isinstance(layer, Qwen4ExpQSAAttention): + raise TypeError(f"{layer_name} is not a Qwen4Exp QSA owner") + layer._run_qsa( + hidden_states, + positions, + query, + key, + value, + output, + ) + + +def qwen4_exp_qsa_with_output_fake( + hidden_states: torch.Tensor, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + layer_name: LayerNameType, +) -> None: + del hidden_states, positions, query, key, value, output, layer_name + + +direct_register_custom_op( + op_name="qwen4_exp_qsa_with_output", + op_func=qwen4_exp_qsa_with_output, + mutates_args=["output"], + fake_impl=qwen4_exp_qsa_with_output_fake, +) + + +__all__ = [ + "QSAIndexer", + "Qwen4ExpQSAAttention", + "Qwen4ExpQSAFlashAttentionBackend", + "Qwen4ExpQSAFlashAttentionImpl", + "qwen4_exp_qsa_with_output", +] diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index f940739a96..23d833ac67 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -131,6 +131,8 @@ def __getitem__(self, key): qianfan_ocr="QianfanOCRConfig", qwen3_asr="Qwen3ASRConfig", qwen3_next="Qwen3NextConfig", + qwen4_exp="Qwen4ExpConfig", + qwen4_exp_text="Qwen4ExpTextConfig", qwen3_5="Qwen3_5Config", qwen3_5_moe="Qwen3_5MoeConfig", laguna="LagunaConfig", diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 5998e61dfd..700d9cda73 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -75,6 +75,9 @@ "QianfanOCRVisionConfig": "vllm.transformers_utils.configs.qianfan_ocr", "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr", "Qwen3NextConfig": "vllm.transformers_utils.configs.qwen3_next", + "Qwen4ExpConfig": "vllm.transformers_utils.configs.qwen4_exp", + "Qwen4ExpTextConfig": "vllm.transformers_utils.configs.qwen4_exp", + "Qwen4ExpVisionConfig": "vllm.transformers_utils.configs.qwen4_exp", "Qwen3_5Config": "vllm.transformers_utils.configs.qwen3_5", "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5", "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe", @@ -143,6 +146,9 @@ "QianfanOCRVisionConfig", "Qwen3ASRConfig", "Qwen3NextConfig", + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", "Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5MoeConfig", diff --git a/vllm/transformers_utils/configs/qwen4_exp.py b/vllm/transformers_utils/configs/qwen4_exp.py new file mode 100644 index 0000000000..ed4c6cca8e --- /dev/null +++ b/vllm/transformers_utils/configs/qwen4_exp.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exports for the package-owned Qwen4Exp configs.""" + +from vllm.models.qwen4_exp.config import ( + Qwen4ExpConfig, + Qwen4ExpTextConfig, + Qwen4ExpVisionConfig, +) + +__all__ = [ + "Qwen4ExpConfig", + "Qwen4ExpTextConfig", + "Qwen4ExpVisionConfig", +] diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 9c85ec5efb..3e888e69b7 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -1,12 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace +from typing import Any -from vllm.v1.attention.backend import AttentionBackend +import torch + +from vllm.config import VllmConfig +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + CommonAttentionMetadata, +) from vllm.v1.attention.backends.mamba_attn import ( BaseMambaAttentionMetadata, BaseMambaAttentionMetadataBuilder, ) +from vllm.v1.attention.backends.utils import ( + NULL_BLOCK_ID, + compute_causal_conv1d_metadata, + mamba_get_block_table_tensor, +) +from vllm.v1.kv_cache_interface import MambaSpec class ShortConvAttentionBackend(AttentionBackend): @@ -32,3 +47,513 @@ class ShortConvAttentionMetadataBuilder( BaseMambaAttentionMetadataBuilder[ShortConvAttentionMetadata] ): metadata_cls = ShortConvAttentionMetadata + + +@dataclass +class PleShortConvAttentionMetadata(ShortConvAttentionMetadata): + # Number of speculative-decode (multi-query / MTP) requests and the total + # number of tokens they contribute. These are 0 when spec-decode is off. + num_spec_decodes: int = 0 + num_spec_decode_tokens: int = 0 + num_actual_tokens: int = 0 + + # Max query length among spec-decode requests (== num_speculative_tokens + 1). + # Used as ``max_query_len`` for the varlen spec causal_conv1d_update. + spec_query_len: int = 1 + + # Max query length among the non-spec *prefill* requests, precomputed + # CPU-side in the builder. The dilated PLE short-conv uses it to size its + # packing buffer without a device->host sync (``lengths.max().item()``). + # 0 when there are no prefill requests. + max_prefill_query_len: int = 0 + query_start_loc: torch.Tensor | None = None + + # ``state_indices_tensor`` keeps the historical (non-spec) layout used by + # all existing short-conv consumers: the conv-state slot for each regular + # decode followed by each prefill request. When spec-decode is active this + # only covers the non-spec requests. + state_indices_tensor: torch.Tensor | None = None + has_initial_states_d: torch.Tensor | None = None + + # ``non_spec_query_start_loc`` is the varlen cumulative token offset over + # the non-spec requests only (decodes then prefills). It equals + # ``query_start_loc`` when there are no spec-decode requests. + non_spec_query_start_loc: torch.Tensor | None = None + + # Speculative-decode (MTP) conv metadata. Only column 0 of the block table + # is needed for the convolution state, so these tensors are 1-D over the + # spec-decode requests. + spec_query_start_loc: torch.Tensor | None = None # [num_spec_decodes + 1] + spec_state_indices_tensor: torch.Tensor | None = None # [num_spec_decodes] + spec_sequence_masks: torch.Tensor | None = None # [batch] + spec_token_indx: torch.Tensor | None = None + non_spec_token_indx: torch.Tensor | None = None + num_decode_draft_tokens_cpu: torch.Tensor | None = None + + +class PleShortConvAttentionBackend(ShortConvAttentionBackend): + @staticmethod + def get_name() -> str: + return "PLE_SHORT_CONV_ATTN" + + @staticmethod + def get_builder_cls() -> type["PleShortConvAttentionMetadataBuilder"]: + return PleShortConvAttentionMetadataBuilder + + +class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): + metadata_cls = PleShortConvAttentionMetadata + # Spec-decode requires a uniform (multi-token) decode batch for full + # CUDA graph capture, matching the GDN backend. + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH + reorder_batch_threshold: int = 1 + supports_update_block_table = False + + def __init__( + self, + kv_cache_spec: MambaSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.num_spec = self.num_spec_tokens + self.use_full_cuda_graph = ( + self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ) + + max_num_seqs = vllm_config.scheduler_config.max_num_seqs + max_capture_size = self.compilation_config.max_cudagraph_capture_size + self.decode_cudagraph_max_bs = max_num_seqs + self.decode_cudagraph_max_tokens = max_num_seqs * (self.num_spec + 1) + if max_capture_size is not None: + self.decode_cudagraph_max_bs = min( + self.decode_cudagraph_max_bs, max_capture_size + ) + self.decode_cudagraph_max_tokens = min( + self.decode_cudagraph_max_tokens, max_capture_size + ) + + # Persistent buffers reused during full CUDA graph capture and replay. + self.spec_state_indices_tensor = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device + ) + self.spec_sequence_masks = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device + ) + self.spec_token_indx = torch.empty( + (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device + ) + self.non_spec_token_indx = torch.empty( + (self.decode_cudagraph_max_tokens,), dtype=torch.int32, device=device + ) + self.spec_query_start_loc = torch.empty( + (self.decode_cudagraph_max_bs + 1,), dtype=torch.int32, device=device + ) + self.num_accepted_tokens = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device + ) + self.has_initial_states_d = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device + ) + + def _build_non_spec_metadata( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool, + num_decode_draft_tokens_cpu: torch.Tensor | None, + **kwargs: Any, + ) -> PleShortConvAttentionMetadata: + metadata = super().build( + common_prefix_len, + common_attn_metadata, + fast_build, + num_accepted_tokens=None, + **kwargs, + ) + assert isinstance(metadata, PleShortConvAttentionMetadata) + + state_indices_d = metadata.state_indices_tensor_d + if state_indices_d is not None and state_indices_d.dim() > 1: + state_indices_d = state_indices_d[:, 0] + state_indices_p = metadata.state_indices_tensor_p + if metadata.num_prefills == 0: + assert state_indices_d is not None + # BaseMambaAttentionMetadataBuilder pads decode state indices into + # a persistent tensor for full CUDA graphs. Keep those rows so the + # PLE decode receives one cache slot per graph-padded token. + state_indices_tensor = state_indices_d + elif metadata.num_decodes == 0: + assert state_indices_p is not None + state_indices_tensor = state_indices_p[: metadata.num_prefills] + else: + assert state_indices_d is not None + assert state_indices_p is not None + state_indices_tensor = torch.cat( + (state_indices_d, state_indices_p[: metadata.num_prefills]) + ) + + has_initial_states_d = None + if metadata.num_decodes > 0: + num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() + has_initial_states_d = num_computed_tokens[: metadata.num_decodes] > 0 + if ( + self.use_full_cuda_graph + and metadata.num_prefills == 0 + and metadata.num_decodes <= self.decode_cudagraph_max_bs + ): + assert state_indices_d is not None + # Prepare tensors for CUDA graph replay. Padded rows have no + # initial state and use NULL_BLOCK_ID in state_indices_d. + num_decode_rows = state_indices_d.numel() + self.has_initial_states_d[: metadata.num_decodes].copy_( + has_initial_states_d, non_blocking=True + ) + self.has_initial_states_d[metadata.num_decodes : num_decode_rows].fill_( + False + ) + has_initial_states_d = self.has_initial_states_d[:num_decode_rows] + + max_prefill_query_len = 0 + if metadata.num_prefills > 0: + query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu) + max_prefill_query_len = int( + query_lens_cpu[ + metadata.num_decodes : ( + metadata.num_decodes + metadata.num_prefills + ) + ] + .max() + .item() + ) + + return replace( + metadata, + num_actual_tokens=common_attn_metadata.num_actual_tokens, + spec_query_len=self.num_spec + 1, + max_prefill_query_len=max_prefill_query_len, + query_start_loc=common_attn_metadata.query_start_loc, + state_indices_tensor=state_indices_tensor, + has_initial_states_d=has_initial_states_d, + non_spec_query_start_loc=common_attn_metadata.query_start_loc, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + + def build( # type: ignore[override] + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + *, + num_accepted_tokens: torch.Tensor | None = None, + num_decode_draft_tokens_cpu: torch.Tensor | None = None, + **kwargs: Any, + ) -> PleShortConvAttentionMetadata: + m = common_attn_metadata + spec_sequence_masks_cpu: torch.Tensor | None = None + # Detect speculative-decode requests. We use -1 to mark prefill and + # plain-decode requests, so any value >= 0 is a (multi-query) + # spec-decode request. + if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: + candidate_mask = num_decode_draft_tokens_cpu[: m.num_reqs] >= 0 + if bool(candidate_mask.any().item()): + spec_sequence_masks_cpu = candidate_mask + + if spec_sequence_masks_cpu is None: + return self._build_non_spec_metadata( + common_prefix_len, + common_attn_metadata, + fast_build, + num_decode_draft_tokens_cpu, + **kwargs, + ) + + del common_prefix_len, fast_build, kwargs + query_start_loc = m.query_start_loc + query_start_loc_cpu = m.query_start_loc_cpu + query_lens_cpu = torch.diff(query_start_loc_cpu) + block_table_tensor = mamba_get_block_table_tensor( + m.block_table_tensor, + m.seq_lens, + self.kv_cache_spec, + self.vllm_config.cache_config.mamba_cache_mode, + ) + + if query_start_loc.device.type == "cpu": + spec_sequence_masks = spec_sequence_masks_cpu + else: + spec_sequence_masks = async_tensor_h2d( + spec_sequence_masks_cpu, device=query_start_loc.device + ) + + # For causal_conv1d (non-spec prefill Triton kernel metadata). + nums_dict = None + batch_ptr = None + token_chunk_offset_ptr = None + has_initial_states_p = None + has_initial_states_d = None + num_computed_tokens_p = None + # Original request indices of the non-spec requests, ordered + # [decodes, prefills]. Used to gather per-request data consistently. + non_spec_req_idx_cpu: torch.Tensor | None = None + + query_lens = torch.diff(query_start_loc) + # Per-request classification by mask, NOT by position. With + # spec-decode, the front decode group can contain both spec-decode + # requests and plain non-spec single-token decodes. Spec requests are + # therefore not guaranteed to occupy the first num_spec_decodes slots. + non_spec_mask_cpu = ~spec_sequence_masks_cpu + decode_mask_cpu = non_spec_mask_cpu & (query_lens_cpu == 1) + prefill_mask_cpu = non_spec_mask_cpu & (query_lens_cpu > 1) + + num_spec_decodes = int(spec_sequence_masks_cpu.sum().item()) + num_decodes = int(decode_mask_cpu.sum().item()) + num_prefills = int(prefill_mask_cpu.sum().item()) + num_decode_tokens = num_decodes + num_prefill_tokens = int(query_lens_cpu[prefill_mask_cpu].sum().item()) + num_spec_decode_tokens = int( + query_lens_cpu[spec_sequence_masks_cpu].sum().item() + ) + # Max prefill query length, CPU-side (no device-to-host sync) for the + # PLE dilated short-conv packing buffer. + max_prefill_query_len = ( + int(query_lens_cpu[prefill_mask_cpu].max().item()) + if num_prefills > 0 + else 0 + ) + + # Original request indices grouped as + # [spec | non-spec decode | non-spec prefill]; each group keeps the + # original (already reordered) relative order via a stable nonzero. + spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] + decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] + prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] + non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) + spec_req_idx = spec_req_idx_cpu.to(query_start_loc.device) + non_spec_req_idx = non_spec_req_idx_cpu.to(query_start_loc.device) + + if num_decodes == 0 and num_prefills == 0: + # Pure speculative-decode batch: all real tokens are spec tokens. + spec_token_indx = torch.arange( + num_spec_decode_tokens, + dtype=torch.int32, + device=query_start_loc.device, + ) + non_spec_token_indx = torch.empty( + 0, dtype=torch.int32, device=query_start_loc.device + ) + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] + non_spec_state_indices_tensor = None + spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] + non_spec_query_start_loc = None + non_spec_query_start_loc_cpu = None + else: + # Mixed batch: build a per-token group key consistent with the + # request grouping above (spec=0 | decode=1 | prefill=2) and a + # stable sort, so tokens of each request stay contiguous and in + # request order. This yields spec tokens first, then the non-spec + # [decode, prefill] tokens. + req_group = torch.full( + (m.num_reqs,), + 2, + dtype=torch.int64, + device=query_start_loc.device, + ) + req_group[spec_req_idx] = 0 + req_group[decode_req_idx_cpu.to(query_start_loc.device)] = 1 + token_group = torch.repeat_interleave(req_group, query_lens) + token_perm = torch.argsort(token_group, stable=True) + spec_token_indx = token_perm[:num_spec_decode_tokens] + non_spec_token_indx = token_perm[num_spec_decode_tokens:] + + spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] + non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0] + spec_query_start_loc = torch.zeros( + num_spec_decodes + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum(query_lens[spec_req_idx], dim=0, out=spec_query_start_loc[1:]) + non_spec_query_start_loc = torch.zeros( + num_decodes + num_prefills + 1, + dtype=torch.int32, + device=query_start_loc.device, + ) + torch.cumsum( + query_lens[non_spec_req_idx], + dim=0, + out=non_spec_query_start_loc[1:], + ) + non_spec_query_start_loc_cpu = torch.zeros( + num_decodes + num_prefills + 1, dtype=torch.int32 + ) + torch.cumsum( + query_lens_cpu[non_spec_req_idx_cpu], + dim=0, + out=non_spec_query_start_loc_cpu[1:], + ) + + assert num_accepted_tokens is not None + # Accepted-token counts must follow the same request order as the + # speculative state indices. + num_accepted_tokens = num_accepted_tokens[ + spec_req_idx_cpu.to(num_accepted_tokens.device) + ] + + # Compute the conv-state slots for the non-spec decode/prefill split, + # plus the initial-state masks and Triton causal_conv1d metadata. + if non_spec_state_indices_tensor is None: + state_indices_tensor = block_table_tensor[:0, 0] + else: + state_indices_tensor = non_spec_state_indices_tensor + + # Build the regular decode/prefill state metadata inherited from the + # generic short-conv metadata contract. + query_start_loc_p = None + query_start_loc_d = None + state_indices_tensor_p = None + state_indices_tensor_d = None + if num_decodes > 0 or num_prefills > 0: + num_computed_tokens = m.compute_num_computed_tokens() + if non_spec_req_idx_cpu is not None: + non_spec_req_idx = non_spec_req_idx_cpu.to(num_computed_tokens.device) + num_computed_tokens = num_computed_tokens[non_spec_req_idx] + + state_indices_tensor_d = state_indices_tensor[:num_decodes] + state_indices_tensor_p = state_indices_tensor[ + num_decodes : num_decodes + num_prefills + ] + if num_decodes > 0: + has_initial_states_d = num_computed_tokens[:num_decodes] > 0 + assert non_spec_query_start_loc is not None + query_start_loc_d = non_spec_query_start_loc[: num_decodes + 1] + if num_prefills > 0: + num_computed_tokens_p = num_computed_tokens[ + num_decodes : num_decodes + num_prefills + ] + has_initial_states_p = num_computed_tokens_p > 0 + assert non_spec_query_start_loc is not None + assert non_spec_query_start_loc_cpu is not None + query_start_loc_p = ( + non_spec_query_start_loc[num_decodes:] - num_decode_tokens + ) + query_start_loc_p_cpu = ( + non_spec_query_start_loc_cpu[num_decodes:] - num_decode_tokens + ) + if query_start_loc.device.type != "cpu": + nums_dict, batch_ptr, token_chunk_offset_ptr = ( + compute_causal_conv1d_metadata( + query_start_loc_p_cpu, + device=query_start_loc.device, + ) + ) + + # Prepare persistent tensors for CUDA graph capture and replay. + # ``m.num_actual_tokens`` is already padded by the model runner. + # Request-level buffers use ``m.num_reqs`` while token-level buffers + # use their independently bounded token count. + batch_size = m.num_reqs + if ( + self.use_full_cuda_graph + and num_prefills == 0 + and num_decodes == 0 + and spec_sequence_masks is not None + and num_spec_decodes <= self.decode_cudagraph_max_bs + and num_spec_decode_tokens <= self.decode_cudagraph_max_tokens + ): + assert spec_state_indices_tensor is not None + self.spec_state_indices_tensor[:num_spec_decodes].copy_( + spec_state_indices_tensor, non_blocking=True + ) + spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] + spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID) + + self.spec_sequence_masks[:batch_size].copy_( + spec_sequence_masks[:batch_size], non_blocking=True + ) + spec_sequence_masks = self.spec_sequence_masks[:batch_size] + + assert spec_query_start_loc is not None + self.spec_query_start_loc[: num_spec_decodes + 1].copy_( + spec_query_start_loc, non_blocking=True + ) + spec_num_query_tokens = spec_query_start_loc[-1] + spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] + spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens) + + assert num_accepted_tokens is not None + self.num_accepted_tokens[:num_spec_decodes].copy_( + num_accepted_tokens, non_blocking=True + ) + num_accepted_tokens = self.num_accepted_tokens[:batch_size] + num_accepted_tokens[num_spec_decodes:].fill_(1) + + return PleShortConvAttentionMetadata( + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_reqs=m.num_reqs, + num_spec_decodes=num_spec_decodes, + num_spec_decode_tokens=num_spec_decode_tokens, + num_actual_tokens=m.num_actual_tokens, + spec_query_len=self.num_spec + 1, + max_prefill_query_len=max_prefill_query_len, + query_start_loc=query_start_loc, + state_indices_tensor=state_indices_tensor, + has_initial_states_p=has_initial_states_p, + has_initial_states_d=has_initial_states_d, + non_spec_query_start_loc=non_spec_query_start_loc, + spec_query_start_loc=spec_query_start_loc, + spec_state_indices_tensor=spec_state_indices_tensor, + spec_sequence_masks=spec_sequence_masks, + spec_token_indx=spec_token_indx, + non_spec_token_indx=non_spec_token_indx, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + nums_dict=nums_dict, + batch_ptr=batch_ptr, + token_chunk_offset_ptr=token_chunk_offset_ptr, + query_start_loc_p=query_start_loc_p, + query_start_loc_d=query_start_loc_d, + state_indices_tensor_p=state_indices_tensor_p, + state_indices_tensor_d=state_indices_tensor_d, + num_computed_tokens_p=num_computed_tokens_p, + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + seq_lens=m.seq_lens, + ) + + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ) -> PleShortConvAttentionMetadata: + """Build metadata for full CUDA graph capture. + + Currently, only decode is supported for full CUDA graphs with + short-conv. + """ + m = common_attn_metadata + assert ( + m.num_reqs <= self.decode_cudagraph_max_bs + and m.num_actual_tokens <= self.decode_cudagraph_max_tokens + ), ( + "ShortConv only supports decode-only full CUDAGraph capture. " + f"Make sure batch size ({m.num_reqs}) <= " + f"cudagraph capture size ({self.decode_cudagraph_max_bs}) and " + f"number of tokens ({m.num_actual_tokens}) <= " + f"token capture size ({self.decode_cudagraph_max_tokens})." + ) + + if self.use_spec_decode: + num_accepted_tokens = torch.diff(m.query_start_loc) + num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu() + return self.build( + 0, + m, + num_accepted_tokens=num_accepted_tokens, + num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + ) + return self.build(0, m) diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 9f1f447340..256443f819 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -7,10 +7,10 @@ import math import os from collections import defaultdict -from collections.abc import Callable, Iterable, Iterator, Sequence +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, replace from functools import partial -from typing import Any, NewType, TypeAlias, cast, overload +from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload from vllm import envs from vllm.config import VllmConfig @@ -21,6 +21,7 @@ from vllm.utils.torch_utils import get_dtype_size from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, FullAttentionSpec, HiddenStateCacheSpec, KVCacheConfig, @@ -880,18 +881,14 @@ def get_max_concurrency_for_kv_cache_config( """ Get the maximum concurrency for the given KV cache configuration. """ - num_layer_per_group = max( - len(group.layer_names) for group in kv_cache_config.kv_cache_groups - ) - max_memory_usage_per_request = num_layer_per_group * max_memory_usage_bytes( - vllm_config, (group.kv_cache_spec for group in kv_cache_config.kv_cache_groups) - ) - memory_per_block = ( - kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes - * num_layer_per_group + num_blocks_per_request = sum( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in kv_cache_config.kv_cache_groups ) - num_block_per_request = cdiv(max_memory_usage_per_request, memory_per_block) - max_concurrency = kv_cache_config.num_blocks / num_block_per_request + max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request return max_concurrency @@ -916,6 +913,8 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): return kv_cache_groups[0].kv_cache_spec.page_size_bytes + if layout := _get_csa_linear_tensor_layout(kv_cache_groups): + return layout.bytes_per_block if all( isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups ): @@ -1248,6 +1247,39 @@ def _get_kv_cache_config_deepseek_v4( return num_blocks, kv_cache_tensors +def _get_kv_cache_config_csa_linear( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + available_memory: int, +) -> tuple[int, list[KVCacheTensor]] | None: + layout = _get_csa_linear_tensor_layout(kv_cache_groups) + if layout is None: + return None + + num_blocks = available_memory // layout.bytes_per_block + num_blocks = may_override_num_blocks(vllm_config, num_blocks) + kv_cache_tensors = [ + KVCacheTensor( + size=layout.main_kv_page_size * num_blocks, + shared_by=[main_kv_name] + + [ + group.layer_names[index] + for group in layout.mamba_groups + if index < len(group.layer_names) + ], + ) + for index, main_kv_name in enumerate(layout.main_kv_names) + ] + kv_cache_tensors.extend( + KVCacheTensor( + size=layout.compressed_page_size * num_blocks, + shared_by=[compressed_name, layout.compressor_state_names[index]], + ) + for index, compressed_name in enumerate(layout.compressed_names) + ) + return num_blocks, kv_cache_tensors + + def get_kv_cache_config_from_groups( vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], @@ -1292,6 +1324,10 @@ def get_kv_cache_config_from_groups( ) for layer_name in kv_cache_groups[0].layer_names ] + elif csa_config := _get_kv_cache_config_csa_linear( + vllm_config, kv_cache_groups, available_memory + ): + num_blocks, kv_cache_tensors = csa_config elif all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups @@ -1470,6 +1506,303 @@ def group_and_unify_kv_cache_specs( return [mla_uniform_spec, *swa_uniform_specs] +class _CSALinearCacheTuple(NamedTuple): + """Three cache owners attached to one compressed-sparse layer.""" + + layer_index: int + main_kv: tuple[str, FullAttentionSpec] + compressed: tuple[str, MLAAttentionSpec] + compressor_state: tuple[str, CircularBufferSpec] + + +@dataclass(frozen=True) +class _CSALinearTensorLayout: + main_kv_names: list[str] + compressed_names: list[str] + compressor_state_names: list[str] + mamba_groups: list[KVCacheGroupSpec] + main_kv_page_size: int + compressed_page_size: int + + @property + def bytes_per_block(self) -> int: + return len(self.main_kv_names) * ( + self.main_kv_page_size + self.compressed_page_size + ) + + +class _CSALinearRoles(NamedTuple): + main_kv: dict[str, FullAttentionSpec] + compressed: dict[str, MLAAttentionSpec] + compressor_state: dict[str, CircularBufferSpec] + mamba: dict[str, MambaSpec] + + +def _classify_csa_linear_specs( + kv_cache_spec: Mapping[str, KVCacheSpec], +) -> _CSALinearRoles | None: + if not any(type(spec) is CircularBufferSpec for spec in kv_cache_spec.values()): + return None + + roles = _CSALinearRoles(main_kv={}, compressed={}, compressor_state={}, mamba={}) + unsupported: list[str] = [] + for name, spec in kv_cache_spec.items(): + if type(spec) is FullAttentionSpec: + roles.main_kv[name] = spec + elif type(spec) is MLAAttentionSpec and spec.compress_ratio > 1: + roles.compressed[name] = spec + elif type(spec) is CircularBufferSpec: + roles.compressor_state[name] = spec + elif type(spec) is MambaSpec: + roles.mamba[name] = spec + else: + unsupported.append(name) + if unsupported: + raise ValueError(f"CSA+linear unsupported cache owners: {unsupported}.") + if not roles.main_kv or not roles.compressed or not roles.mamba: + raise ValueError( + "CSA+linear requires main KV, compressed, compressor-state, and " + "Mamba cache owners." + ) + return roles + + +def _get_csa_linear_cache_tuples( + roles: _CSALinearRoles, +) -> list[_CSALinearCacheTuple]: + from vllm.model_executor.models.utils import extract_layer_index + + def by_layer_index( + specs: Mapping[str, KVCacheSpec], role: str + ) -> dict[int, tuple[str, KVCacheSpec]]: + indexed: dict[int, tuple[str, KVCacheSpec]] = {} + for name, spec in specs.items(): + try: + layer_index = extract_layer_index(name) + except (AssertionError, ValueError) as exc: + raise ValueError( + f"CSA+linear {role} owner {name!r} does not identify " + "exactly one transformer layer." + ) from exc + if layer_index in indexed: + other_name = indexed[layer_index][0] + raise ValueError( + f"CSA+linear layer {layer_index} has duplicate {role} " + f"owners: {other_name!r} and {name!r}." + ) + indexed[layer_index] = (name, spec) + return indexed + + main_kv = by_layer_index(roles.main_kv, "main-KV") + compressed = by_layer_index(roles.compressed, "compressed") + compressor_state = by_layer_index(roles.compressor_state, "compressor-state") + layer_indices = set(main_kv) + if set(compressed) != layer_indices or set(compressor_state) != layer_indices: + raise ValueError( + "CSA+linear main-KV, compressed, and compressor-state owners must " + "have matching transformer-layer indices." + ) + + return [ + _CSALinearCacheTuple( + layer_index, + cast(tuple[str, FullAttentionSpec], main_kv[layer_index]), + cast(tuple[str, MLAAttentionSpec], compressed[layer_index]), + cast(tuple[str, CircularBufferSpec], compressor_state[layer_index]), + ) + for layer_index in sorted(layer_indices) + ] + + +def _get_csa_linear_mamba_group_count( + vllm_config: VllmConfig, + mamba_names: list[str], + main_kv_names: list[str], +) -> int | None: + num_groups = cdiv(len(mamba_names), len(main_kv_names)) + pp_size = vllm_config.parallel_config.pipeline_parallel_size + if pp_size == 1: + return num_groups + + from vllm.distributed.utils import get_pp_indices + from vllm.model_executor.models.utils import extract_layer_index + + total_layers = vllm_config.model_config.get_total_num_hidden_layers() + mamba_indices = [extract_layer_index(name) for name in mamba_names] + main_kv_indices = [extract_layer_index(name) for name in main_kv_names] + for rank in range(pp_size): + start, end = get_pp_indices(total_layers, rank, pp_size) + num_mamba = sum(start <= index < end for index in mamba_indices) + num_main_kv = sum(start <= index < end for index in main_kv_indices) + if not num_mamba: + continue + if not num_main_kv: + return None + num_groups = max(num_groups, cdiv(num_mamba, num_main_kv)) + return num_groups + + +def _get_kv_cache_groups_csa_linear( + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec] | None: + """Build Qwen4Exp's main/compressed/ring/recurrent cache groups.""" + + roles = _classify_csa_linear_specs(kv_cache_spec) + if roles is None: + return None + tuples = _get_csa_linear_cache_tuples(roles) + + expected_local_kv_heads = vllm_config.model_config.get_num_kv_heads( + vllm_config.parallel_config + ) + if any( + cache.main_kv[1].num_kv_heads != expected_local_kv_heads for cache in tuples + ): + raise ValueError( + "CSA+linear main-KV specs do not match the TP-local KV-head geometry." + ) + + for cache in tuples: + _, main_kv = cache.main_kv + _, compressed = cache.compressed + _, compressor_state = cache.compressor_state + if not ( + main_kv.block_size == compressed.block_size + and compressor_state.real_page_size_bytes <= compressed.page_size_bytes + and all( + spec.page_size_padded is None + for spec in (main_kv, compressed, compressor_state) + ) + ): + raise ValueError( + f"CSA+linear layer {cache.layer_index} violates cache geometry." + ) + + shapes = { + ( + cache.compressed[1].compress_ratio, + cache.main_kv[1].block_size, + cache.main_kv[1].page_size_bytes, + cache.compressed[1].page_size_bytes, + ) + for cache in tuples + } + if len(shapes) != 1: + raise ValueError( + "CSA+linear layers must share one block size, compression ratio, " + "and main/compressed page geometry." + ) + _, _, main_kv_page, compressed_page = next(iter(shapes)) + + padded_compressor_specs: dict[str, KVCacheSpec] = { + cache.compressor_state[0]: replace( + cache.compressor_state[1], page_size_padded=compressed_page + ) + for cache in tuples + } + compressed_sparse_specs: dict[str, KVCacheSpec] = { + name: spec + for cache in tuples + for name, spec in (cache.main_kv, cache.compressed) + } + compressed_sparse_uniform = UniformTypeKVCacheSpecs.from_specs( + compressed_sparse_specs + ) + compressor_uniform = UniformTypeKVCacheSpecs.from_specs(padded_compressor_specs) + if compressed_sparse_uniform is None or compressor_uniform is None: + raise ValueError("CSA+linear cache owners have incompatible lifetimes.") + + groups = [ + KVCacheGroupSpec(list(compressed_sparse_specs), compressed_sparse_uniform), + KVCacheGroupSpec(list(padded_compressor_specs), compressor_uniform), + ] + main_kv_names = [cache.main_kv[0] for cache in tuples] + for tp_replicated in (False, True): + names = [ + name + for name, spec in roles.mamba.items() + if spec.tp_replicated is tp_replicated + ] + if not names: + continue + representative = roles.mamba[names[0]] + if any(roles.mamba[name] != representative for name in names[1:]): + policy = "replicated" if tp_replicated else "sharded" + raise ValueError( + f"CSA+linear {policy} Mamba owners must use one cache spec." + ) + unpadded_page = replace(representative, page_size_padded=None).page_size_bytes + if unpadded_page > main_kv_page: + raise ValueError( + f"CSA+linear Mamba owner {names[0]!r} needs {unpadded_page} " + f"bytes, but a main-KV page has {main_kv_page} bytes." + ) + num_groups = _get_csa_linear_mamba_group_count( + vllm_config, names, main_kv_names + ) + if num_groups is None: + raise ValueError( + "CSA+linear pipeline stage has Mamba owners but no main-KV slots." + ) + padded_spec = replace(representative, page_size_padded=main_kv_page) + grouped_names: list[list[str]] = [[] for _ in range(num_groups)] + for index, name in enumerate(names): + grouped_names[index % num_groups].append(name) + groups.extend( + KVCacheGroupSpec(group_names, padded_spec) for group_names in grouped_names + ) + return groups + + +def _get_csa_linear_tensor_layout( + kv_cache_groups: list[KVCacheGroupSpec], +) -> _CSALinearTensorLayout | None: + compressed_sparse: Mapping[str, KVCacheSpec] | None = None + compressor_state: Mapping[str, KVCacheSpec] | None = None + mamba_groups: list[KVCacheGroupSpec] = [] + for group in kv_cache_groups: + if not group.layer_names: + continue + spec = group.kv_cache_spec + if isinstance(spec, MambaSpec): + mamba_groups.append(group) + continue + if not isinstance(spec, UniformTypeKVCacheSpecs): + return None + member = next(iter(spec.kv_cache_specs.values())) + if type(member) is CircularBufferSpec: + compressor_state = spec.kv_cache_specs + elif type(member) is FullAttentionSpec: + compressed_sparse = spec.kv_cache_specs + else: + return None + if compressed_sparse is None or compressor_state is None: + return None + + main_kv_names = [ + name + for name, spec in compressed_sparse.items() + if type(spec) is FullAttentionSpec + ] + compressed_names = [ + name + for name, spec in compressed_sparse.items() + if type(spec) is MLAAttentionSpec + ] + if not main_kv_names or not compressed_names: + return None + + return _CSALinearTensorLayout( + main_kv_names=main_kv_names, + compressed_names=compressed_names, + compressor_state_names=list(compressor_state), + mamba_groups=mamba_groups, + main_kv_page_size=compressed_sparse[main_kv_names[0]].page_size_bytes, + compressed_page_size=compressed_sparse[compressed_names[0]].page_size_bytes, + ) + + def _select_deepseek_v4_tuple_width( vllm_config: VllmConfig, grouped_specs: list[UniformTypeKVCacheSpecs], @@ -1663,6 +1996,8 @@ def get_kv_cache_groups( ) _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups) return kv_cache_groups + elif csa_groups := _get_kv_cache_groups_csa_linear(vllm_config, kv_cache_spec): + return csa_groups # Pull HiddenStateCacheSpec layers out before the general multi-group # path so they don't affect page-size unification or grouping. @@ -1770,6 +2105,16 @@ def _max_memory_usage_bytes_from_groups( spec.max_memory_usage_bytes(vllm_config) for spec in per_layer_specs.values() ) + elif layout := _get_csa_linear_tensor_layout(kv_cache_groups): + blocks_needed = sum( + cdiv( + group.kv_cache_spec.max_memory_usage_bytes(vllm_config), + group.kv_cache_spec.page_size_bytes, + ) + for group in kv_cache_groups + if group.layer_names + ) + return layout.bytes_per_block * blocks_needed elif all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index b3980c592f..69f4717792 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -14,6 +14,7 @@ ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, + CircularBufferSpec, CrossAttentionSpec, FullAttentionSpec, HiddenStateCacheSpec, @@ -592,6 +593,98 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: return num_common_blocks +class CircularBufferManager(FullAttentionManager): + """Own exactly one QSA compressor-ring block per live request.""" + + def _claim_ring_block(self, request_id: str) -> list[KVCacheBlock]: + req_blocks = self.req_to_blocks[request_id] + if req_blocks: + return [] + new_blocks = self.block_pool.get_new_blocks(1) + req_blocks.extend(new_blocks) + self.new_block_ids.extend(block.block_id for block in new_blocks) + return new_blocks + + def get_num_blocks_to_allocate( + self, + request_id: str, + num_tokens: int, + new_computed_blocks: Sequence[KVCacheBlock], + total_computed_tokens: int, + num_tokens_main_model: int, + apply_admission_cap: bool = False, + ) -> int: + del ( + num_tokens, + new_computed_blocks, + total_computed_tokens, + num_tokens_main_model, + apply_admission_cap, + ) + return 0 if self.req_to_blocks.get(request_id) else 1 + + def allocate_new_computed_blocks( + self, + request_id: str, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + num_external_computed_tokens: int, + ) -> None: + del new_computed_blocks, num_local_computed_tokens, num_external_computed_tokens + self._claim_ring_block(request_id) + + def allocate_new_blocks( + self, request_id: str, num_tokens: int, num_tokens_main_model: int + ) -> list[KVCacheBlock]: + del num_tokens, num_tokens_main_model + return self._claim_ring_block(request_id) + + @classmethod + def find_longest_cache_hit( + cls, + block_hashes: BlockHashList, + max_length: int, + kv_cache_group_ids: list[int], + block_pool: BlockPool, + kv_cache_spec: KVCacheSpec, + use_eagle: bool, + alignment_tokens: int, + dcp_world_size: int = 1, + pcp_world_size: int = 1, + ) -> tuple[list[KVCacheBlock], ...]: + del ( + block_hashes, + max_length, + block_pool, + kv_cache_spec, + use_eagle, + alignment_tokens, + dcp_world_size, + pcp_world_size, + ) + return tuple([] for _ in kv_cache_group_ids) + + def cache_blocks( + self, + request: Request, + num_tokens: int, + alignment_tokens: int | None = None, + ) -> None: + del request, num_tokens, alignment_tokens + + def remove_skipped_blocks( + self, + request_id: str, + processed_computed_tokens: int, + num_prompt_tokens: int | None = None, + ) -> None: + del request_id, processed_computed_tokens, num_prompt_tokens + + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: + del running_request_id + return 0 + + class PrefixAnchoredSWAManager(FullAttentionManager): """KV cache manager for prefix-anchored sliding-window attention. @@ -1326,6 +1419,7 @@ def __init__( MambaSpec: MambaManager, CrossAttentionSpec: CrossAttentionManager, SinkFullAttentionSpec: SinkFullAttentionManager, + CircularBufferSpec: CircularBufferManager, } diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 83f32dfccf..d55f5a9e41 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -538,6 +538,26 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: return max_blocks * self.page_size_bytes +@dataclass(frozen=True, kw_only=True) +class CircularBufferSpec(AttentionSpec): + """One fixed block per request for QSA's uncompressed key ring.""" + + head_size_v: int = 0 + + @property + def real_page_size_bytes(self) -> int: + return ( + self.block_size + * self.num_kv_heads + * (self.head_size + self.head_size_v) + * get_dtype_size(self.dtype) + ) + + def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: + del vllm_config + return self.page_size_bytes + + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): """Sliding window attention with MLA cache format.""" @@ -611,6 +631,8 @@ class MambaSpec(KVCacheSpec): mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 mamba_cache_mode: str = "none" num_speculative_blocks: int = 0 + # PLE short-conv state is replicated; GDN state is TP-sharded. + tp_replicated: bool = False @property def page_size_bytes(self) -> int: @@ -749,6 +771,10 @@ def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: and spec.sliding_window == one_spec.sliding_window for spec in kv_cache_specs.values() ) + elif isinstance(one_spec, CircularBufferSpec): + return all( + isinstance(spec, CircularBufferSpec) for spec in kv_cache_specs.values() + ) elif isinstance(one_spec, FullAttentionSpec): return all( isinstance(spec, FullAttentionSpec) for spec in kv_cache_specs.values() diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 771f06bff5..87fcf73d1b 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -199,7 +199,9 @@ def _reshape_kv_cache( if isinstance(kv_cache_spec, AttentionSpec): has_attn = True - num_blocks_per_kv_block = kv_cache_spec.block_size // kernel_block_size + num_blocks_per_kv_block = ( + kv_cache_spec.storage_block_size // kernel_block_size + ) kernel_num_blocks = num_blocks * num_blocks_per_kv_block kv_cache_shape = group.backend.get_kv_cache_shape( kernel_num_blocks, diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 41692f58e7..d403019400 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -21,6 +21,7 @@ def __init__( cp_size: int = 1, cp_rank: int = 0, cp_interleave: int = 1, + slot_mapping_enabled: list[bool] | None = None, ): self.block_sizes = block_sizes self.kernel_block_sizes = kernel_block_sizes @@ -34,6 +35,10 @@ def __init__( self.num_kv_cache_groups = len(self.block_sizes) assert len(max_num_blocks_per_group) == self.num_kv_cache_groups + if slot_mapping_enabled is None: + slot_mapping_enabled = [True] * self.num_kv_cache_groups + assert len(slot_mapping_enabled) == self.num_kv_cache_groups + self._slot_mapping_enabled = slot_mapping_enabled self.blocks_per_kv_block = [ bs // kbs for bs, kbs in zip(block_sizes, kernel_block_sizes) @@ -91,6 +96,9 @@ def init_block_table_layout_tensors(self) -> None: self.block_sizes_tensor = torch.tensor( self.kernel_block_sizes, dtype=torch.int32, device=self.device ) + self.slot_mapping_enabled = torch.tensor( + self._slot_mapping_enabled, dtype=torch.bool, device=self.device + ) self.input_block_table_ptrs = self._make_ptr_tensor(self.input_block_tables) def append_block_ids( @@ -158,6 +166,7 @@ def compute_slot_mappings( self.block_table_ptrs, self.block_table_strides, self.block_sizes_tensor, + self.slot_mapping_enabled, self.slot_mappings, self.slot_mappings.stride(0), self.cp_rank, @@ -229,6 +238,7 @@ def _compute_slot_mappings_kernel( block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] block_sizes, # [num_kv_cache_groups] + slot_mapping_enabled, # [num_kv_cache_groups] slot_mappings_ptr, # [num_kv_cache_groups, max_num_tokens] slot_mappings_stride, cp_rank, @@ -256,6 +266,7 @@ def _compute_slot_mappings_kernel( block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32) block_table_stride = tl.load(block_table_strides + group_id) block_size = tl.load(block_sizes + group_id) + mapping_enabled = tl.load(slot_mapping_enabled + group_id) req_state_idx = tl.load(idx_mapping + batch_idx) start_idx = tl.load(query_start_loc + batch_idx) @@ -264,7 +275,9 @@ def _compute_slot_mappings_kernel( offset = i + tl.arange(0, TRITON_BLOCK_SIZE) positions = tl.load(pos + offset, mask=offset < end_idx, other=0) - block_indices = positions // (block_size * CP_SIZE) + block_indices = tl.where( + mapping_enabled, positions // (block_size * CP_SIZE), 0 + ) block_offsets = positions % (block_size * CP_SIZE) block_numbers = tl.load( block_table_ptr + req_state_idx * block_table_stride + block_indices @@ -282,6 +295,7 @@ def _compute_slot_mappings_kernel( slot_ids = block_numbers * block_size + local_offsets slot_ids = tl.where(is_local, slot_ids, PAD_ID) + slot_ids = tl.where(mapping_enabled, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 47d1da107e..b85b6d8e9b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -50,7 +50,12 @@ from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheConfig, + MambaSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput @@ -395,18 +400,31 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: block_sizes = [] max_num_blocks_per_group = [] + slot_mapping_enabled = [] for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec + if isinstance(spec, UniformTypeKVCacheSpecs): + specs = tuple(spec.kv_cache_specs.values()) + assert specs + is_circular = all( + isinstance(member, CircularBufferSpec) for member in specs + ) + spec = specs[0] + else: + is_circular = isinstance(spec, CircularBufferSpec) block_sizes.append(spec.block_size) + slot_mapping_enabled.append(not is_circular) # When using DCP, each request's KV cache is sharded among different ranks. # As a result, one block on the current rank covers `block_size * cp_size` # tokens in the full, global (unsharded) sequence. - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.dcp_size + max_num_blocks = ( + 1 + if is_circular + else cdiv(block_table_max_model_len, spec.block_size * self.dcp_size) ) # Align to a multiple of (128 / block_size) as required by some attention # backends such as TRTLLM (#39324) - if spec.block_size <= 128: + if not is_circular and spec.block_size <= 128: alignment = 128 // spec.block_size max_num_blocks = cdiv(max_num_blocks, alignment) * alignment # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens @@ -430,6 +448,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cp_size=self.dcp_size, cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, + slot_mapping_enabled=slot_mapping_enabled, ) initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 140cf97413..2c55f6ac35 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -31,6 +31,9 @@ prepare_dflash2_gdn_group_metadata, ) from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder +from vllm.v1.attention.backends.short_conv_attn import ( + PleShortConvAttentionMetadataBuilder, +) from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer @@ -90,7 +93,11 @@ def get_extra_attn_kwargs( } if not isinstance( attn_metadata_builder, - (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder), + ( + Mamba2AttentionMetadataBuilder, + GDNAttentionMetadataBuilder, + PleShortConvAttentionMetadataBuilder, + ), ): return {} kwargs = { diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 1cdcf567d6..9bcafe5210 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -15,7 +15,13 @@ NewRequestData, SchedulerOutput, ) -from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec, MambaSpec +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + CrossAttentionSpec, + KVCacheSpec, + MambaSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner @@ -29,6 +35,12 @@ def _reserved_block_count( max_encoder_len: int, ) -> int: """Match the scheduler's block reservation in hand-built warmup batches.""" + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + specs = tuple(kv_cache_spec.kv_cache_specs.values()) + assert specs + kv_cache_spec = specs[0] + if isinstance(kv_cache_spec, CircularBufferSpec): + return 1 if isinstance(kv_cache_spec, CrossAttentionSpec): return cdiv(max_encoder_len, kv_cache_spec.block_size) num_speculative_blocks = 0 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c150cf791e..53efe8b579 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -153,6 +153,7 @@ from vllm.v1.kv_cache_interface import ( AttentionSpec, ChunkedLocalAttentionSpec, + CircularBufferSpec, CrossAttentionSpec, EncoderOnlyAttentionSpec, FullAttentionSpec, @@ -1380,6 +1381,19 @@ def __init__( # Only relevant for models using ALiBi (e.g, MPT) self.use_alibi = model_config.uses_alibi + ple_layer_ids = getattr(model_config.hf_text_config, "ple_layer_ids", ()) + self.uses_ngram_embedding = bool(ple_layer_ids) + if self.uses_ngram_embedding: + self.ngram_context_len = int(model_config.hf_text_config.ngram_size) - 1 + self.ngram_eos_token_id = int(model_config.hf_text_config.eos_token_id) + else: + self.ngram_context_len = 0 + self.ngram_eos_token_id = 0 + if self.uses_ngram_embedding and self.ngram_context_len <= 0: + raise ValueError("N-gram embedding requires context length >= 1") + if self.uses_ngram_embedding and parallel_config.pipeline_parallel_size > 1: + raise RuntimeError("N-gram PLE embedding requires pipeline_parallel_size=1") + self.cascade_attn_enabled = not self.model_config.disable_cascade_attn self.is_mm_prefix_lm = self.model_config.is_mm_prefix_lm @@ -1713,6 +1727,12 @@ def __init__( self.inputs_embeds = self._make_buffer( self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False ) + if self.uses_ngram_embedding: + self.ngram_context = self._make_buffer( + self.max_num_reqs, + self.ngram_context_len, + dtype=torch.int32, + ) self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool) self.discard_request_mask = self._make_buffer( self.max_num_reqs, dtype=torch.bool @@ -5088,11 +5108,26 @@ def _prepare_inputs( profile_positions_ms = (time.perf_counter() - profile_inner_t0) * 1000.0 profile_inner_t0 = time.perf_counter() if profile_inputs else 0.0 - self.input_batch.block_table.compute_slot_mapping( - num_reqs, - self.query_start_loc.gpu[: num_reqs + 1], - self.positions[:total_num_scheduled_tokens], - ) + for group_id, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups): + kv_cache_spec = kv_cache_group.kv_cache_spec + is_circular = isinstance(kv_cache_spec, CircularBufferSpec) + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs): + is_circular = all( + isinstance(spec, CircularBufferSpec) + for spec in kv_cache_spec.kv_cache_specs.values() + ) + block_table = self.input_batch.block_table[group_id] + if is_circular: + # QSA derives ring slots from logical positions in its own + # metadata builder. Generic mapping would index beyond the + # ring's single block-table column. + block_table.slot_mapping.gpu.fill_(PAD_SLOT_ID) + continue + block_table.compute_slot_mapping( + num_reqs, + self.query_start_loc.gpu[: num_reqs + 1], + self.positions[:total_num_scheduled_tokens], + ) if profile_inputs: profile_slot_mapping_inner_ms = ( time.perf_counter() - profile_inner_t0 @@ -6593,10 +6628,83 @@ def _prepare_mm_inputs( inputs_embeds = self.inputs_embeds.gpu[:num_tokens] return input_ids, inputs_embeds + def _prepare_ngram_context( + self, + num_reqs: int, + num_reqs_padded: int, + ) -> torch.Tensor: + """Copy committed per-request token history into the PLE context.""" + if not self.uses_ngram_embedding: + raise RuntimeError("N-gram context requested for a non-PLE model") + + eos_token_id = int(self.ngram_eos_token_id) + context_cpu = self.ngram_context.np[:num_reqs_padded] + context_cpu.fill(eos_token_id) + num_computed = self.input_batch.num_computed_tokens_cpu + token_ids = self.input_batch.token_ids_cpu + is_token_ids = self.input_batch.is_token_ids + + for req_idx in range(num_reqs): + end = int(num_computed[req_idx]) + if end <= 0: + continue + start = max(0, end - self.ngram_context_len) + context_tokens = token_ids[req_idx, start:end] + if context_tokens.size == 0: + continue + if self.enable_prompt_embeds and not is_token_ids[req_idx, start:end].all(): + context_tokens = context_tokens.copy() + context_tokens[~is_token_ids[req_idx, start:end]] = eos_token_id + context_cpu[req_idx, -context_tokens.size :] = context_tokens + + self._copy_buffer_to_gpu(self.ngram_context, num_reqs_padded) + return self.ngram_context.gpu[:num_reqs_padded] + + def _maybe_add_ngram_kwargs( + self, + model_kwargs: dict[str, Any], + *, + num_reqs: int, + num_reqs_padded: int, + is_first_rank: bool, + is_encoder_decoder: bool, + use_dummy_context: bool, + query_start_loc: torch.Tensor | None = None, + num_scheduled_tokens: Sequence[int] | np.ndarray | None = None, + ) -> None: + if not self.uses_ngram_embedding or not is_first_rank or is_encoder_decoder: + return + + if query_start_loc is None: + if num_scheduled_tokens is None: + raise RuntimeError("query_start_loc is required for N-gram PLE") + scheduled = np.asarray(num_scheduled_tokens, dtype=np.int32) + cu_num_tokens = np.cumsum(scheduled, dtype=np.int32) + last = int(cu_num_tokens[-1]) if num_reqs > 0 else 0 + self.query_start_loc.np[0] = 0 + if num_reqs > 0: + self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens + self.query_start_loc.np[num_reqs + 1 :].fill(last) + self._copy_buffer_to_gpu(self.query_start_loc) + query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1] + model_kwargs["query_start_loc"] = query_start_loc + + if use_dummy_context: + self.ngram_context.np[:num_reqs_padded].fill(int(self.ngram_eos_token_id)) + self._copy_buffer_to_gpu(self.ngram_context, num_reqs_padded) + model_kwargs["ngram_context"] = self.ngram_context.gpu[:num_reqs_padded] + else: + model_kwargs["ngram_context"] = self._prepare_ngram_context( + num_reqs, + num_reqs_padded, + ) + def _preprocess( self, scheduler_output: "SchedulerOutput", num_input_tokens: int, # Padded + num_reqs: int, + num_reqs_padded: int, intermediate_tensors: IntermediateTensors | None = None, ) -> tuple[ torch.Tensor | None, @@ -6675,6 +6783,25 @@ def _preprocess( inputs_embeds = None model_kwargs = self._init_model_kwargs() + if ( + self.uses_ngram_embedding + and is_first_rank + and not is_encoder_decoder + and input_ids is None + ): + raise RuntimeError( + "N-gram PLE requires token IDs on the first pipeline rank" + ) + self._maybe_add_ngram_kwargs( + model_kwargs, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + is_first_rank=is_first_rank, + is_encoder_decoder=is_encoder_decoder, + use_dummy_context=False, + query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], + ) + if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_input_tokens] elif self.uses_xdrope_dim > 0: @@ -8540,7 +8667,11 @@ def execute_model( model_kwargs, ec_connector_output, ) = self._preprocess( - scheduler_output, num_tokens_padded, intermediate_tensors + scheduler_output, + num_tokens_padded, + num_reqs, + num_reqs_padded, + intermediate_tensors, ) if trace_log: trace_model_preprocess_ms = ( @@ -10832,6 +10963,16 @@ def _dummy_run( input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None + self._maybe_add_ngram_kwargs( + model_kwargs, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + is_first_rank=get_pp_group().is_first_rank, + is_encoder_decoder=self.model_config.is_encoder_decoder, + use_dummy_context=True, + num_scheduled_tokens=num_scheduled_tokens, + ) + if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] elif self.uses_xdrope_dim > 0: @@ -12049,9 +12190,12 @@ def may_reinitialize_input_batch( continue block_size = kv_cache_group.kv_cache_spec.block_size block_sizes.append(block_size) - max_num_blocks_per_req = cdiv( - max_model_len, block_size * get_total_cp_world_size() - ) + if isinstance(kv_cache_group.kv_cache_spec, CircularBufferSpec): + max_num_blocks_per_req = 1 + else: + max_num_blocks_per_req = cdiv( + max_model_len, block_size * get_total_cp_world_size() + ) if isinstance(kv_cache_group.kv_cache_spec, MambaSpec): max_num_blocks_per_req = ( max_num_blocks_per_req From 6b23eea07b01f7ebccf93991503ae33d706a9cee Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:10:37 +0800 Subject: [PATCH 03/28] [Doc] Record verified ModelScope snapshot Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index fcf159bddc..54191ae5c1 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,8 +2,9 @@ ## Status and ownership -- Status: source bring-up implemented; CPU/configuration gates pass. Full-model - load, output quality, measured memory, and speed are not claimed yet. +- Status: source bring-up implemented; CPU/configuration gates pass and the + local ModelScope snapshot is fully verified. Full-model load, output quality, + measured memory, and speed are not claimed yet. - Integration line: `private/main`. - Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. - Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision @@ -11,7 +12,9 @@ - Model download: `/data/models/RadixArk/Qwen3.8-Flash-Next-NVFP4`. - Download source: ModelScope `master`, verified against the fixed Hugging Face revision above: all 419 file sizes match and all 208 comparable LFS SHA-256 - values match. + values match. After download, all 419 local files were independently hashed + against the ModelScope manifest with zero missing, size-mismatched, or + SHA-256-mismatched files. - Upstream references: [vLLM PR 53896](https://github.com/vllm-project/vllm/pull/53896) and [SGLang PR 36497](https://github.com/sgl-project/sglang/pull/36497). @@ -135,6 +138,10 @@ only after profiles identify them as measured decode bottlenecks. ## Source validation snapshot +- The ModelScope download completed successfully at the path above. A full + post-download verification checked 419 files and about 125.910 GiB of + safetensor payload: zero files were missing and zero size or SHA-256 values + differed from the remote manifest. - The real downloaded `config.json` resolves without remote model code as `Qwen4ExpConfig` / `Qwen4ExpTextConfig`: 48 layers, 36 GDN, 12 QSA, 512 experts, top-10, HC count four/rank 320, and one trigram PLE layer. From 3fc247d1e9fedbef0743e07d6c64df898bbb9f1f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:49:09 +0800 Subject: [PATCH 04/28] [Bugfix] Attach Qwen4Exp expert mapping Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_weight_loading.py | 54 +++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model.py | 14 +++++ 2 files changed, 68 insertions(+) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py index 81df1c5a9c..fb5d9903c7 100644 --- a/tests/models/qwen4_exp/test_weight_loading.py +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +from torch import nn +from vllm.model_executor.models.qwen3_next import Qwen3NextSparseMoeBlock from vllm.models.qwen4_exp.nvidia.model import ( Qwen4ExpForConditionalGeneration, Qwen4ExpModel, + Qwen4ExpSparseMoeBlock, _remap_qsa_cache_scale_name, ) @@ -99,6 +102,57 @@ def test_outer_checkpoint_mapper_selects_language_model_only_paths() -> None: ) +def test_sparse_moe_attaches_private_recursive_loader_mapping(monkeypatch) -> None: + class FakeExperts(nn.Module): + def __init__(self) -> None: + super().__init__() + self.expert_mapping = None + + def fake_qwen3_next_init(self, vllm_config, prefix="") -> None: + del vllm_config, prefix + nn.Module.__init__(self) + self.n_routed_experts = 2 + self.n_redundant_experts = 0 + self.experts = FakeExperts() + + monkeypatch.setattr( + Qwen3NextSparseMoeBlock, + "__init__", + fake_qwen3_next_init, + ) + vllm_config = type( + "VllmConfigStub", + (), + { + "parallel_config": type( + "ParallelConfigStub", (), {"use_sequence_parallel_moe": False} + )(), + "model_config": type( + "ModelConfigStub", + (), + { + "hf_text_config": type( + "TextConfigStub", + (), + {"shared_expert_intermediate_size": 640}, + )() + }, + )(), + }, + )() + + block = Qwen4ExpSparseMoeBlock(vllm_config, prefix="model.layers.0.mlp") + + assert block.experts.expert_mapping == [ + ("experts.w13_", "experts.0.gate_proj.", 0, "w1"), + ("experts.w2_", "experts.0.down_proj.", 0, "w2"), + ("experts.w13_", "experts.0.up_proj.", 0, "w3"), + ("experts.w13_", "experts.1.gate_proj.", 1, "w1"), + ("experts.w2_", "experts.1.down_proj.", 1, "w2"), + ("experts.w13_", "experts.1.up_proj.", 1, "w3"), + ] + + @pytest.mark.parametrize( ("checkpoint_name", "model_name"), [ diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 5f32b8ee6e..08520d3be4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -11,6 +11,9 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_pp_group +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( QwenGatedDeltaNetAttention, @@ -194,6 +197,17 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__(vllm_config=vllm_config, prefix=prefix) config = vllm_config.model_config.hf_text_config self.n_shared_experts = int(config.shared_expert_intermediate_size > 0) + # Qwen3Next loads expert tensors manually in this private tree, so its + # FusedMoE does not receive the mapping required by AutoWeightsLoader. + # Qwen4Exp uses recursive loading and must attach that mapping here. + self.experts.expert_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.n_routed_experts, + num_redundant_experts=self.n_redundant_experts, + ) class Qwen4ExpDecoderLayer(nn.Module): From 09f68f836116c57bdd19f62c74b38ade6899a226 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:59:55 +0800 Subject: [PATCH 05/28] [Core] Report checkpoint name on weight-load failure Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/model_executor/models/utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 302e71dbd7..561bd419a4 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -243,7 +243,14 @@ def _load_param( ) weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, weight_data) + try: + weight_loader(param, weight_data) + except Exception as exc: + raise RuntimeError( + f"Error loading weight {weight_qualname!r} " + f"with checkpoint shape {tuple(weight_data.shape)} into " + f"parameter shape {tuple(param.shape)}" + ) from exc logger.debug("Loaded weight %s with shape %s", weight_qualname, param.shape) From 86b19124e53b8b45b745334075d8e62872b5fab1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:08:17 +0800 Subject: [PATCH 06/28] [Bugfix] Preserve QKV shard metadata in auto loading Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_weights_mapper_stacked.py | 42 ++++++++++++++++++- vllm/model_executor/layers/linear.py | 24 +++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/tests/model_executor/test_weights_mapper_stacked.py b/tests/model_executor/test_weights_mapper_stacked.py index 07f45a55dd..2cbe9b13fd 100644 --- a/tests/model_executor/test_weights_mapper_stacked.py +++ b/tests/model_executor/test_weights_mapper_stacked.py @@ -3,7 +3,10 @@ import torch -from vllm.model_executor.layers.linear import MergedColumnParallelLinear +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, +) from vllm.model_executor.models.utils import WeightsMapper @@ -74,3 +77,40 @@ def weight_loader(param, loaded_weight, shard_id) -> None: (weight, down, 0), (weight, injection, 1), ] + + +def test_qkv_load_weights_forwards_stacked_shards() -> None: + layer = object.__new__(QKVParallelLinear) + torch.nn.Module.__init__(layer) + layer.prefix = "attn.qkv_proj" + weight = torch.nn.Parameter(torch.zeros(7, 4)) + calls = [] + + def weight_loader(param, loaded_weight, shard_id) -> None: + calls.append((param, loaded_weight, shard_id)) + + weight.weight_loader = weight_loader + layer.register_parameter("weight", weight) + query = torch.ones(3, 4) + query.shard_id = "q" + key = torch.full((2, 4), 2.0) + key.shard_id = "k" + value = torch.full((2, 4), 3.0) + value.shard_id = "v" + + loaded = list( + layer.load_weights( + [ + ("weight", query), + ("weight", key), + ("weight", value), + ] + ) + ) + + assert loaded == ["weight", "weight", "weight"] + assert calls == [ + (weight, query, "q"), + (weight, key, "k"), + (weight, value, "v"), + ] diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 47477033f3..1d1f726b45 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1337,6 +1337,30 @@ def validate_shard_id(self, loaded_shard_id: str | None): return raise ValueError("This line should not be reached") + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[str]: + """Load separately serialized Q/K/V shards through AutoWeightsLoader.""" + for name, loaded_weight in weights: + shard_id = getattr(loaded_weight, "shard_id", None) + self.validate_shard_id(shard_id) + if "." in name: + submodule, _, attr = name.rpartition(".") + param = getattr(self.get_submodule(submodule), attr, self) + else: + param = getattr(self, name, self) + if param is None and name == "bias": + continue + param.weight_loader(param, loaded_weight, shard_id) + logger.debug( + "Loaded shard %s with shape %s into %s.%s", + shard_id, + loaded_weight.shape, + self.prefix, + name, + ) + yield name + def _get_shard_offset_mapping(self, loaded_shard_id: str): shard_offset_mapping = { "q": 0, From 78225e67fa22003f6864e61528bda1935734d36c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:16:52 +0800 Subject: [PATCH 07/28] [Model] Pass 1Cat GDN output buffer in Qwen4Exp Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_model_forward.py | 63 ++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/model.py | 22 ++++++- 2 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 tests/models/qwen4_exp/test_model_forward.py diff --git a/tests/models/qwen4_exp/test_model_forward.py b/tests/models/qwen4_exp/test_model_forward.py new file mode 100644 index 0000000000..9a34e3cf31 --- /dev/null +++ b/tests/models/qwen4_exp/test_model_forward.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.models.qwen4_exp.nvidia.model import Qwen4ExpDecoderLayer + + +class _AttentionHyperConnection: + def mix(self, hidden_states: torch.Tensor): + return hidden_states, hidden_states, torch.zeros_like(hidden_states) + + +class _MlpHyperConnection: + def combine_and_mix( + self, + hidden_states: torch.Tensor, + block_output: torch.Tensor, + injection: torch.Tensor, + ): + return hidden_states, block_output, injection + + +class _OutputOnlyGdn: + def __init__(self) -> None: + self.output: torch.Tensor | None = None + + def __call__( + self, + *, + hidden_states: torch.Tensor, + output: torch.Tensor | None, + ) -> torch.Tensor: + assert output is not None + self.output = output + output.copy_(hidden_states + 1) + return hidden_states + 100 + + +def test_linear_attention_forwards_preallocated_output_buffer() -> None: + layer = object.__new__(Qwen4ExpDecoderLayer) + object.__setattr__(layer, "ple", None) + object.__setattr__(layer, "layer_type", "linear_attention") + object.__setattr__(layer, "attn_hyper_connection", _AttentionHyperConnection()) + object.__setattr__(layer, "mlp_hyper_connection", _MlpHyperConnection()) + gdn = _OutputOnlyGdn() + object.__setattr__(layer, "linear_attn", gdn) + object.__setattr__(layer, "mlp", lambda hidden_states: hidden_states) + hidden_states = torch.arange(6, dtype=torch.float32).view(2, 3) + + _, mlp_out, _ = Qwen4ExpDecoderLayer.forward( + layer, + hidden_states, + None, + None, + torch.arange(2), + input_ids=None, + query_start_loc=None, + ngram_context=None, + ) + + assert gdn.output is not None + torch.testing.assert_close(mlp_out, hidden_states + 1) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 08520d3be4..7eddc1bb35 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -8,6 +8,7 @@ import torch from torch import nn +from vllm import envs from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_pp_group @@ -349,7 +350,26 @@ def forward( hidden_states, block_input, injection = attn_hc.mix(hidden_states) if self.layer_type == "linear_attention": - attn_out = self.linear_attn(hidden_states=block_input) + use_direct_attention_output = ( + envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM + and torch.compiler.is_compiling() + ) + attn_buffer = ( + None if use_direct_attention_output else torch.empty_like(block_input) + ) + projected_attn_out = self.linear_attn( + hidden_states=block_input, + output=attn_buffer, + ) + if use_direct_attention_output: + if projected_attn_out is None: + raise RuntimeError( + "SM70 TP4 fused prefill requires a direct GDN output" + ) + attn_out = projected_attn_out + else: + assert attn_buffer is not None + attn_out = attn_buffer elif self.layer_type == "full_attention": attn_out = self.self_attn( hidden_states=block_input, From d174f3ff0ed1df9d9dedeeb9bb073d47b5f8cffe Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:25:34 +0800 Subject: [PATCH 08/28] [Model] Honor custom V2 model state providers Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_model_state_init.py | 41 +++++++++++++++++++++ vllm/v1/worker/gpu/model_states/__init__.py | 5 +++ 2 files changed, 46 insertions(+) create mode 100644 tests/v1/worker/test_model_state_init.py diff --git a/tests/v1/worker/test_model_state_init.py b/tests/v1/worker/test_model_state_init.py new file mode 100644 index 0000000000..7535a70b6c --- /dev/null +++ b/tests/v1/worker/test_model_state_init.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch +from torch import nn + +from vllm.v1.worker.gpu.model_states import init_model_state + + +def test_model_can_select_custom_model_state() -> None: + captured = {} + + class CustomModelState: + def __init__(self, vllm_config, model, encoder_cache, device) -> None: + captured.update( + vllm_config=vllm_config, + model=model, + encoder_cache=encoder_cache, + device=device, + ) + + class CustomModel(nn.Module): + @staticmethod + def get_model_state_cls(): + return CustomModelState + + model = CustomModel() + vllm_config = SimpleNamespace() + device = torch.device("cpu") + + model_state = init_model_state(vllm_config, model, None, device) + + assert isinstance(model_state, CustomModelState) + assert captured == { + "vllm_config": vllm_config, + "model": model, + "encoder_cache": None, + "device": device, + } diff --git a/vllm/v1/worker/gpu/model_states/__init__.py b/vllm/v1/worker/gpu/model_states/__init__.py index 06b5a92c39..d7f9e6f83e 100644 --- a/vllm/v1/worker/gpu/model_states/__init__.py +++ b/vllm/v1/worker/gpu/model_states/__init__.py @@ -13,6 +13,11 @@ def init_model_state( encoder_cache: EncoderCache | None, device: torch.device, ): + get_model_state_cls = getattr(model, "get_model_state_cls", None) + if callable(get_model_state_cls): + model_state_cls = get_model_state_cls() + return model_state_cls(vllm_config, model, encoder_cache, device) + if "WhisperForConditionalGeneration" in vllm_config.model_config.architectures: from vllm.v1.worker.gpu.model_states.whisper import WhisperModelState From 7fb606850bae15770716c81fe663b091bb67a180 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 02:34:12 +0800 Subject: [PATCH 09/28] [Bugfix] Zero heterogeneous hybrid KV pages Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_kv_block_zeroer.py | 47 ++++++++++ vllm/v1/worker/utils.py | 112 +++++++++++++----------- 2 files changed, 107 insertions(+), 52 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 332111b7f4..4e31e11bde 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -159,3 +159,50 @@ def test_zero_block_ids_multiple_interleaved_pools() -> None: index for index in range(cells.shape[0]) if bool((cells[index] == 0).all()) } assert zeroed == expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_zero_block_ids_nonuniform_page_sizes() -> None: + device = torch.device("cuda") + page_sizes = [8, 16] + buffers = [ + torch.ones(N_BLOCKS * page_size, dtype=torch.int32, device=device) + for page_size in page_sizes + ] + groups = [] + static_forward_context = {} + for index, (page_size, buffer) in enumerate(zip(page_sizes, buffers)): + spec = MambaSpec( + block_size=16, + shapes=((1,),), + dtypes=(torch.int32,), + page_size_padded=page_size * 4, + ) + layer_name = f"mamba.nonuniform.{index}" + groups.append( + SimpleNamespace( + kv_cache_spec=spec, + kv_cache_group_id=index, + layer_names=[layer_name], + backend=None, + ) + ) + static_forward_context[layer_name] = SimpleNamespace(kv_cache=[buffer]) + + zeroer = KVBlockZeroer(device, pin_memory=False) + zeroer.init_meta( + groups, + kernel_block_sizes=[16, 16], + cache_dtype="auto", + runner_only_attn_layers=set(), + static_forward_context=static_forward_context, + ) + + target = 2 + zeroer.zero_block_ids([target]) + torch.accelerator.synchronize() + + for page_size, buffer in zip(page_sizes, buffers): + pages = buffer.view(N_BLOCKS, page_size) + zeroed = {index for index in range(N_BLOCKS) if bool((pages[index] == 0).all())} + assert zeroed == {target} diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 6b634557a5..8f011bc02a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -16,7 +16,6 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.platforms import current_platform from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import largest_power_of_2_divisor from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.v1.attention.backend import ( AttentionBackend, @@ -75,14 +74,12 @@ def _infer_segment_block_strides( return block_strides -@triton.jit(do_not_specialize=["n_blocks"]) +@triton.jit def _zero_kv_blocks_kernel( seg_addrs_ptr, seg_block_strides_ptr, + seg_page_sizes_ptr, block_ids_ptr, - n_blocks, - N_SEGS: tl.constexpr, - PAGE_SIZE_EL: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): """Zero KV cache blocks across all segments in a single launch. @@ -98,27 +95,27 @@ def _zero_kv_blocks_kernel( elements, kept separate from the PAGE_SIZE_EL span each block zeros (they differ for interleaved layouts, see init_meta). - Programs are mapped as (block_index, seg_index, chunk_index). + Programs are mapped as (block_index, seg_index, chunk_index). Segment page + sizes may differ across hybrid cache groups. """ - pid = tl.program_id(0) - chunks = PAGE_SIZE_EL // BLOCK_SIZE - work_per_block = N_SEGS * chunks - block_index = pid // work_per_block - if block_index >= n_blocks: + block_index = tl.program_id(0) + seg_index = tl.program_id(1) + chunk_index = tl.program_id(2) + page_size_el = tl.load(seg_page_sizes_ptr + seg_index) + chunk_offset = chunk_index.to(tl.int64) * BLOCK_SIZE + if chunk_offset >= page_size_el: return - remainder = pid % work_per_block - seg_index = remainder // chunks - chunk_index = remainder % chunks block_id = tl.load(block_ids_ptr + block_index) seg_addr = tl.load(seg_addrs_ptr + seg_index) block_stride_el = tl.load(seg_block_strides_ptr + seg_index) ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32)) - offset = ( - block_id.to(tl.int64) * block_stride_el.to(tl.int64) - + chunk_index.to(tl.int64) * BLOCK_SIZE + block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64) + cols = chunk_offset + tl.arange(0, BLOCK_SIZE).to(tl.int64) + tl.store( + ptr + block_offset + cols, + tl.zeros([BLOCK_SIZE], dtype=tl.int32), + mask=cols < page_size_el, ) - cols = tl.arange(0, BLOCK_SIZE).to(tl.int64) - tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32)) class KVBlockZeroer: @@ -132,7 +129,9 @@ class KVBlockZeroer: def __init__(self, device: torch.device, pin_memory: bool): self.device = device self.pin_memory = pin_memory - self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None + self._meta: ( + tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None + ) = None self._id_cap: int = 0 self._ids_pinned: torch.Tensor | None = None self._ids_gpu: torch.Tensor | None = None @@ -161,9 +160,17 @@ def init_meta( state tensor views over one raw page per block; the first state tensor starts at the page base, so one segment per layer zeros the whole page. """ - seen_ptrs: set[int] = set() + seen_ptrs: dict[int, int] = {} seg_addrs: list[int] = [] - page_size_el: int | None = None + seg_page_sizes: list[int] = [] + + def add_segment(address: int, page_size_el: int) -> None: + if (index := seen_ptrs.get(address)) is not None: + seg_page_sizes[index] = max(seg_page_sizes[index], page_size_el) + return + seen_ptrs[address] = len(seg_addrs) + seg_addrs.append(address) + seg_page_sizes.append(page_size_el) for group in attn_groups_iter: spec = group.kv_cache_spec @@ -186,21 +193,12 @@ def init_meta( cache_dtype_str=cache_dtype, ) dp = kv.data_ptr() - if dp in seen_ptrs: - continue - seen_ptrs.add(dp) el = kv.element_size() cur_bytes = kv.stride(block_dim) * el assert cur_bytes % 4 == 0 kernel_block_el = cur_bytes // 4 cur_page_el = kernel_block_el * ratio - if page_size_el is None: - page_size_el = cur_page_el - else: - assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" - ) block_stride_bytes = cur_bytes outer_dims = [ @@ -211,7 +209,7 @@ def init_meta( outer_strides = [kv.stride(d) * el for d in outer_dims] for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)): off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) - seg_addrs.append(dp + off_bytes) + add_segment(dp + off_bytes, cur_page_el) elif isinstance(spec, MambaSpec): if not isinstance(kv, (list, tuple)) or not kv: continue @@ -219,34 +217,38 @@ def init_meta( if not isinstance(first_state, torch.Tensor): continue dp = first_state.data_ptr() - if dp in seen_ptrs: - continue - seen_ptrs.add(dp) page_bytes = spec.page_size_bytes assert page_bytes % 4 == 0 cur_page_el = page_bytes // 4 - if page_size_el is None: - page_size_el = cur_page_el - else: - assert page_size_el == cur_page_el, ( - f"Non-uniform page sizes: {page_size_el} vs {cur_page_el}" - ) - seg_addrs.append(dp) - - if not seg_addrs or page_size_el is None: + add_segment(dp, cur_page_el) + + if not seg_addrs: self._meta = None return - blk_size = min(largest_power_of_2_divisor(page_size_el), 1024) + max_page_size_el = max(seg_page_sizes) + blk_size = min(1 << (max_page_size_el - 1).bit_length(), 1024) # Dense layouts space blocks page_size_el apart. Block-major # interleaved pools expose segment starts one cell apart and advance by # the number of segments in that contiguous address run. Infer each # run separately because a process may own more than one KV pool. n_segs = len(seg_addrs) + block_strides = [0] * n_segs + for page_size_el in set(seg_page_sizes): + indices = [ + index + for index, size in enumerate(seg_page_sizes) + if size == page_size_el + ] + inferred = _infer_segment_block_strides( + [seg_addrs[index] for index in indices], page_size_el + ) + for index, stride in zip(indices, inferred): + block_strides[index] = stride seg_block_strides = torch.tensor( - _infer_segment_block_strides(seg_addrs, page_size_el), + block_strides, dtype=torch.int64, device=self.device, ) @@ -261,7 +263,8 @@ def init_meta( self._meta = ( torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), seg_block_strides, - page_size_el, + torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device), + (max_page_size_el + blk_size - 1) // blk_size, blk_size, n_segs, ) @@ -270,7 +273,14 @@ def zero_block_ids(self, block_ids: list[int]) -> None: """Zero the KV cache memory for the given block IDs.""" if not block_ids or self._meta is None: return - seg_addrs, seg_block_strides, page_size_el, blk_size, n_segs = self._meta + ( + seg_addrs, + seg_block_strides, + seg_page_sizes, + max_chunks, + blk_size, + n_segs, + ) = self._meta n_blocks = len(block_ids) if n_blocks > self._id_cap: self._id_cap = n_blocks * 2 @@ -286,14 +296,12 @@ def zero_block_ids(self, block_ids: list[int]) -> None: self._ids_pinned[:n_blocks].numpy()[:] = block_ids idx = self._ids_gpu[:n_blocks] idx.copy_(self._ids_pinned[:n_blocks], non_blocking=True) - grid = (n_blocks * n_segs * (page_size_el // blk_size),) + grid = (n_blocks, n_segs, max_chunks) _zero_kv_blocks_kernel[grid]( seg_addrs, seg_block_strides, + seg_page_sizes, idx, - n_blocks, - N_SEGS=n_segs, - PAGE_SIZE_EL=page_size_el, BLOCK_SIZE=blk_size, ) From f73cedb3ec20b864049dc2e411851fa1ef049b5e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:03:56 +0800 Subject: [PATCH 10/28] [Bugfix] Fall back for unsupported custom AR dtypes Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_custom_all_reduce_dispatch.py | 35 +++++++++++++++++++ .../device_communicators/custom_all_reduce.py | 2 ++ 2 files changed, 37 insertions(+) create mode 100644 tests/distributed/test_custom_all_reduce_dispatch.py diff --git a/tests/distributed/test_custom_all_reduce_dispatch.py b/tests/distributed/test_custom_all_reduce_dispatch.py new file mode 100644 index 0000000000..3703931fd3 --- /dev/null +++ b/tests/distributed/test_custom_all_reduce_dispatch.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def _mock_communicator() -> CustomAllreduce: + communicator = object.__new__(CustomAllreduce) + communicator.disabled = False + communicator._ptr = 0 + communicator.world_size = 4 + communicator.fully_connected = True + communicator.tp8_hierarchical = False + communicator.dispatch_max_size = 1024 * 1024 + return communicator + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_should_custom_ar_accepts_supported_dtype(dtype: torch.dtype) -> None: + communicator = _mock_communicator() + + assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) + + +@pytest.mark.parametrize( + "dtype", + [torch.float64, torch.int64, torch.int32, torch.int8, torch.uint8, torch.bool], +) +def test_should_custom_ar_rejects_unsupported_dtype(dtype: torch.dtype) -> None: + communicator = _mock_communicator() + + assert not communicator.should_custom_ar(torch.empty(16, dtype=dtype)) diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 8d55ceeaff..790922fd6f 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -380,6 +380,8 @@ def register_graph_buffers(self): def should_custom_ar(self, inp: torch.Tensor): if self.disabled: return False + if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16): + return False inp_size = inp.numel() * inp.element_size() # custom allreduce requires input byte size to be multiples of 16 if inp_size % 16 != 0: From d64150fb9bd60eaa3bbac0db074e06652c0de7bc Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:18:08 +0800 Subject: [PATCH 11/28] [Core] Invoke custom KV cache bind hooks Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_qsa_cache.py | 38 +++++++++++++++++++ .../layers/attention_layer_base.py | 6 +++ vllm/v1/worker/utils.py | 2 +- 3 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 tests/models/qwen4_exp/test_qsa_cache.py diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py new file mode 100644 index 0000000000..494e9b5ae7 --- /dev/null +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.models.qwen4_exp.common.qsa_cache import QSAKeyStateCache +from vllm.v1.worker.utils import bind_kv_cache + + +def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: + prefix = "model.layers.0.self_attn.raw_key_cache" + static_forward_context = {} + layer = QSAKeyStateCache( + head_size=128, + dtype=torch.float16, + cache_rope_positions=True, + prefix=prefix, + cache_config=SimpleNamespace(block_size=16), + compress_ratio=4, + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace( + static_forward_context=static_forward_context + ) + ), + ) + cache = torch.empty(2, 8, 1, layer.head_size, dtype=torch.float16) + runner_kv_caches = [] + + bind_kv_cache({prefix: cache}, static_forward_context, runner_kv_caches) + + assert layer.kv_cache is cache + assert layer.key_cache.shape == (2, 8, 1, 128) + assert layer.key_cache.untyped_storage().data_ptr() == cache.data_ptr() + assert layer.rope_position_cache.shape == (2, 8, 1, 3) + assert layer.rope_position_cache.dtype == torch.int64 + assert runner_kv_caches == [cache] diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index 97395b6414..512ffd182c 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -4,6 +4,8 @@ from abc import ABC, abstractmethod +import torch + from vllm.config import VllmConfig from vllm.v1.attention.backend import AttentionBackend, AttentionImpl from vllm.v1.kv_cache_interface import KVCacheSpec @@ -20,6 +22,10 @@ class AttentionLayerBase(ABC): impl: "AttentionImpl" + def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: + """Bind an allocated KV cache view to this layer.""" + self.kv_cache = kv_cache + @abstractmethod def get_attn_backend(self) -> type[AttentionBackend]: """Get the attention backend class for this layer.""" diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 8f011bc02a..a9e1bd2501 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -609,7 +609,7 @@ def bind_kv_cache( # Bind kv_caches to forward context for layer_name, kv_cache in kv_caches.items(): - forward_context[layer_name].kv_cache = kv_cache + forward_context[layer_name].bind_kv_cache(kv_cache) def is_residual_scattered_for_sp( From b37cd74c94f471d15d146ea3d8bd8ac316e1891b Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:48:20 +0800 Subject: [PATCH 12/28] [Bugfix] Match QSA to local Flash KV layout Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_qsa_cache.py | 59 ++++++++++++++++++++++++ vllm/models/qwen4_exp/nvidia/qsa.py | 4 +- 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py index 494e9b5ae7..9a5ed1daca 100644 --- a/tests/models/qwen4_exp/test_qsa_cache.py +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -6,6 +6,7 @@ import torch from vllm.models.qwen4_exp.common.qsa_cache import QSAKeyStateCache +from vllm.models.qwen4_exp.nvidia.qsa import Qwen4ExpQSAFlashAttentionImpl from vllm.v1.worker.utils import bind_kv_cache @@ -36,3 +37,61 @@ def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: assert layer.rope_position_cache.shape == (2, 8, 1, 3) assert layer.rope_position_cache.dtype == torch.int64 assert runner_kv_caches == [cache] + + +def test_qsa_forward_splits_local_flash_cache_layout(monkeypatch) -> None: + from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops + + num_blocks, block_size, head_size = 2, 16, 8 + kv_cache = torch.arange( + num_blocks * 2 * block_size * head_size, + dtype=torch.float16, + ).view(num_blocks, 2, block_size, 1, head_size) + query = torch.zeros(1, 2, head_size, dtype=torch.float16) + output = torch.empty_like(query) + logical_indices = torch.zeros(1, 4, dtype=torch.int32) + block_table = torch.zeros(1, 1, dtype=torch.int32) + token_to_req = torch.zeros(1, dtype=torch.int32) + captured = {} + + def fake_sparse_attention( + query_arg, + key_cache_arg, + value_cache_arg, + logical_indices_arg, + block_table_arg, + token_to_req_arg, + output_arg, + ): + captured["key_cache"] = key_cache_arg + captured["value_cache"] = value_cache_arg + assert torch.equal(query_arg, query) + assert torch.equal(logical_indices_arg, logical_indices) + assert torch.equal(block_table_arg, block_table) + assert torch.equal(token_to_req_arg, token_to_req) + output_arg.fill_(1) + return output_arg + + monkeypatch.setattr(qsa_ops, "qsa_sparse_paged_attention", fake_sparse_attention) + impl = object.__new__(Qwen4ExpQSAFlashAttentionImpl) + impl.head_size = head_size + impl.alibi_slopes = None + impl.sinks = None + impl.sliding_window = (-1, -1) + + result = impl.forward_qsa( + SimpleNamespace(topk_indices_buffer=logical_indices), + query, + query[:, :1], + query[:, :1], + kv_cache, + SimpleNamespace(num_actual_tokens=1, block_table=block_table), + output, + token_to_req, + ) + + expected_key, expected_value = kv_cache.unbind(1) + assert torch.equal(captured["key_cache"], expected_key) + assert torch.equal(captured["value_cache"], expected_value) + assert result is output + assert torch.equal(output, torch.ones_like(output)) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 0bab984eac..b6b68b45cb 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -147,7 +147,9 @@ def forward_qsa( raise RuntimeError("QSA owner did not provide its top-k buffer") logical_indices = topk_buffer[:num_tokens] token_to_req = token_to_req[:num_tokens] - key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + # This tree's FlashAttention cache ABI keeps K/V on dimension 1: + # [num_blocks, 2, block_size, num_kv_heads, head_size]. + key_cache, value_cache = kv_cache.unbind(1) key_cache = canonicalize_singleton_dim_strides(key_cache) value_cache = canonicalize_singleton_dim_strides(value_cache) if key_cache.dtype != query.dtype or query.dtype not in ( From 7349fe259df3a0eefc274626a75e02abc4b34ed1 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 06:20:56 +0800 Subject: [PATCH 13/28] [Bugfix] Keep PLE scales in model dtype Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_ple.py | 12 ++++++++++-- vllm/models/qwen4_exp/nvidia/ple_layer.py | 9 ++++++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index fcb081596e..14d8206fe7 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -58,6 +58,7 @@ def test_pinned_host_ple_allocates_tp_shard_without_device_table( assert layer.weight._vllm_keep_on_cpu assert not layer.weight.is_meta assert not layer.weight_scale.is_meta + assert layer.weight_scale.dtype == torch.float16 assert layer._accelerator_weight_views == {} @@ -286,6 +287,7 @@ def test_ngram_embedding_loads_fp8_shards_and_global_scale() -> None: def _make_fp8_embedding_layer( monkeypatch: pytest.MonkeyPatch, + params_dtype: torch.dtype = torch.bfloat16, ) -> embedding_module.VocabParallelEmbedding: monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr( @@ -304,16 +306,22 @@ def _make_fp8_embedding_layer( layer = embedding_module.VocabParallelEmbedding( 3, 2, - params_dtype=torch.bfloat16, + params_dtype=params_dtype, padding_size=1, quant_method=method, ) weight = torch.tensor([[1.0, 2.0], [4.0, 8.0], [16.0, 32.0]]) layer.weight.data.copy_(weight.to(torch.float8_e4m3fn)) - layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=torch.bfloat16)) + layer.weight_scale.data.copy_(torch.tensor([0.25], dtype=params_dtype)) return layer +def test_ple_fp8_embedding_scale_matches_model_dtype(monkeypatch) -> None: + layer = _make_fp8_embedding_layer(monkeypatch, params_dtype=torch.float16) + + assert layer.weight_scale.dtype == torch.float16 + + def test_ple_fp8_embedding_dequantizes_in_ple_layer(monkeypatch) -> None: layer = _make_fp8_embedding_layer(monkeypatch) quantized_output = layer(torch.tensor([2, 0])) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index b4549add56..4f38d0dc61 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -162,7 +162,7 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: - del input_size, output_size, params_dtype + del input_size, output_size weight_loader = extra_weight_attrs.get("weight_loader") weight = create_fp8_weight_parameter( sum(output_partition_sizes), input_size_per_partition, weight_loader @@ -175,7 +175,10 @@ def create_weights( input_size_per_partition, None, weight_loader, - scale_dtype=torch.bfloat16, + # Keep graph inputs in the requested model dtype. In particular, + # an otherwise-FP16 graph cannot retain a BF16 scale parameter on + # SM70 because Inductor rejects BF16 graph inputs there. + scale_dtype=params_dtype, ) layer.register_parameter("weight_scale", weight_scale) @@ -293,7 +296,7 @@ def __init__( self.embedding_dim, None, self.weight_loader, - scale_dtype=torch.bfloat16, + scale_dtype=params_dtype, ) self._accelerator_weight_views: dict[int, torch.Tensor] = {} logger.info( From c228412e47aff9862841ca152f79c751fed8b781 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 06:53:46 +0800 Subject: [PATCH 14/28] [Kernel] Decode offloaded PLE FP8 rows on SM70 Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_ple.py | 23 ++-- vllm/models/qwen4_exp/nvidia/ple_layer.py | 125 ++++++++++++++++++++++ 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 14d8206fe7..102d1045da 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -78,6 +78,9 @@ def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( monkeypatch.setattr( parameter_module, "get_tensor_model_parallel_world_size", lambda: 1 ) + monkeypatch.setattr( + embedding_module, "tensor_model_parallel_all_reduce", lambda tensor: tensor + ) layer = Qwen4ExpPinnedHostEmbedding( num_embeddings=8, embedding_dim=8, @@ -86,16 +89,22 @@ def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( prefix="model.layers.2.ple.ngram_embedding", quant_method=Qwen4ExpPLEFp8EmbeddingMethod(), ) - expected = torch.arange(64, dtype=torch.float32).reshape(8, 8) % 16 - layer.weight.data.copy_(expected.to(torch.float8_e4m3fn)) - - accelerator_weight = layer.get_accelerator_weight(torch.device("cuda")) - output = F.embedding( - torch.tensor([0, 7], dtype=torch.int64, device="cuda"), accelerator_weight + raw = torch.tensor( + [0x00, 0x01, 0x08, 0x38, 0x7E, 0x80, 0xB8, 0xFE], + dtype=torch.uint8, + ).repeat(8, 1) + layer.weight.data.copy_(raw.view(torch.float8_e4m3fn)) + layer.weight_scale = nn.Parameter( + torch.tensor([0.25], dtype=torch.float16, device="cuda"), + requires_grad=False, ) + layer.prepare_accelerator_weight() + + output = layer(torch.tensor([0, 7], dtype=torch.int64, device="cuda")) torch.cuda.synchronize() - assert output.dtype == torch.float8_e4m3fn + assert output.dtype == torch.float16 + expected = raw.view(torch.float8_e4m3fn).float() * 0.25 torch.testing.assert_close(output.float().cpu(), expected[[0, 7]]) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 4f38d0dc61..0f071ba965 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -44,6 +44,7 @@ from vllm.transformers_utils.configs.qwen4_exp import ( Qwen4ExpTextConfig, ) +from vllm.triton_utils import tl, triton from vllm.utils.mem_utils import format_gib from vllm.utils.platform_utils import is_pin_memory_available from vllm.utils.torch_utils import ( @@ -68,6 +69,59 @@ logger = init_logger(__name__) +@triton.jit +def _apply_float32_sign_bit(value, sign_bit): + """Apply an FP8 sign bit without canonicalizing negative zero.""" + + value_bits = tl.cast(value, tl.uint32, bitcast=True) + signed_bits = value_bits | (sign_bit.to(tl.uint32) << 31) + return tl.cast(signed_bits, tl.float32, bitcast=True) + + +@triton.jit +def _e4m3fn_byte_to_float(raw): + """Decode E4M3FN bytes without requiring native FP8 on SM70.""" + + raw_i32 = raw.to(tl.int32) + sign_bit = (raw_i32 >> 7) & 1 + exponent = (raw_i32 >> 3) & 0x0F + mantissa = raw_i32 & 0x07 + mantissa_f32 = mantissa.to(tl.float32) + normal = (1.0 + mantissa_f32 * 0.125) * tl.exp2(exponent.to(tl.float32) - 7.0) + subnormal = mantissa_f32 * 0.001953125 # 2**-9 + value = _apply_float32_sign_bit( + tl.where(exponent == 0, subnormal, normal), sign_bit + ) + is_nan = (exponent == 0x0F) & (mantissa == 0x07) + return tl.where(is_nan, float("nan"), value) + + +@triton.jit +def _gather_ple_fp8_from_pinned_kernel( + weight_ptr, + ids_ptr, + scale_ptr, + output_ptr, + embedding_dim: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Gather and dequantize one PLE row per program from pinned host memory.""" + + row_idx = tl.program_id(0) + local_idx = tl.load(ids_ptr + row_idx) + offsets = tl.arange(0, BLOCK_D) + mask = offsets < embedding_dim + byte_ptr = weight_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) + raw = tl.load( + byte_ptr + local_idx * embedding_dim + offsets, + mask=mask, + other=0, + ) + scale = tl.load(scale_ptr).to(tl.float32) + values = _e4m3fn_byte_to_float(raw) * scale + tl.store(output_ptr + row_idx * embedding_dim + offsets, values, mask=mask) + + def _splitmix64(value: int) -> int: value = (value + _SPLITMIX_GAMMA) & _MASK64 value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64 @@ -191,6 +245,9 @@ def apply( raise NotImplementedError("PLE FP8 weights only support embedding lookup") def embedding(self, layer: nn.Module, input_: torch.Tensor) -> torch.Tensor: + pinned_lookup = getattr(layer, "embedding_lookup", None) + if pinned_lookup is not None: + return pinned_lookup(input_) get_accelerator_weight = getattr(layer, "get_accelerator_weight", None) weight = ( get_accelerator_weight(input_.device) @@ -299,6 +356,8 @@ def __init__( scale_dtype=params_dtype, ) self._accelerator_weight_views: dict[int, torch.Tensor] = {} + self._accelerator_weight_ptrs: dict[int, int] = {} + self._output_dtype = self.weight_scale.dtype logger.info( "Qwen4Exp PLE shard allocated in pinned host memory: %s", format_gib(self.weight.numel() * self.weight.element_size()), @@ -321,11 +380,38 @@ def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: with torch.cuda.device(device_index): view = get_accelerator_view_from_cpu_tensor(self.weight) self._accelerator_weight_views[device_index] = view + self._accelerator_weight_ptrs[device_index] = view.data_ptr() return view def prepare_accelerator_weight(self) -> None: self.get_accelerator_weight(torch.device("cuda", torch.cuda.current_device())) + def embedding_lookup(self, input_: torch.Tensor) -> torch.Tensor: + """Gather FP8 UVA rows and emit scaled model-dtype values.""" + + device_index = ( + torch.cuda.current_device() + if input_.device.index is None + else input_.device.index + ) + weight_ptr = self._accelerator_weight_ptrs.get(device_index) + if weight_ptr is None: + self.get_accelerator_weight(input_.device) + weight_ptr = self._accelerator_weight_ptrs[device_index] + output = torch.empty( + (*input_.shape, self.embedding_dim), + dtype=self._output_dtype, + device=input_.device, + ) + torch.ops.vllm.qwen4_exp_ple_pinned_gather( + input_.reshape(-1), + output.reshape(-1, self.embedding_dim), + self.weight_scale, + weight_ptr, + self.embedding_dim, + ) + return output + class Qwen4ExpNGramEmbedding(nn.Module): def __init__( @@ -1298,6 +1384,45 @@ def qwen4_exp_ple_short_conv_fake( return +def qwen4_exp_ple_pinned_gather( + input_ids: torch.Tensor, + output: torch.Tensor, + weight_scale: torch.Tensor, + weight_ptr: int, + embedding_dim: int, +) -> None: + if input_ids.numel() == 0: + return + block_d = triton.next_power_of_2(embedding_dim) + _gather_ple_fp8_from_pinned_kernel[(input_ids.numel(),)]( + weight_ptr, + input_ids, + weight_scale, + output, + embedding_dim=embedding_dim, + BLOCK_D=block_d, + num_warps=4, + ) + + +def qwen4_exp_ple_pinned_gather_fake( + input_ids: torch.Tensor, + output: torch.Tensor, + weight_scale: torch.Tensor, + weight_ptr: int, + embedding_dim: int, +) -> None: + return + + +direct_register_custom_op( + op_name="qwen4_exp_ple_pinned_gather", + op_func=qwen4_exp_ple_pinned_gather, + mutates_args=["output"], + fake_impl=qwen4_exp_ple_pinned_gather_fake, +) + + direct_register_custom_op( op_name="qwen4_exp_ple_short_conv", op_func=qwen4_exp_ple_short_conv, From de10e2ea4c2f667354c5559c60e9d3cf5ca2508e Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 07:02:15 +0800 Subject: [PATCH 15/28] [Bugfix] Preserve dynamic PLE request shapes Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/models/qwen4_exp/nvidia/ple_layer.py | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 0f071ba965..1ccad20a4b 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -571,18 +571,11 @@ def forward( ) -> torch.Tensor: input_ids = input_ids.reshape(-1).long() query_start_loc = query_start_loc.long() - num_reqs = query_start_loc.numel() - 1 + num_reqs = ngram_context.shape[0] num_tokens = input_ids.shape[0] - if num_tokens > self.positions_buffer.numel(): - raise ValueError( - f"PLE received {num_tokens} tokens, but its workspace supports " - f"at most {self.positions_buffer.numel()}" - ) - if num_reqs > self.padded_buffer.shape[0]: - raise ValueError( - f"PLE received {num_reqs} requests, but its workspace supports " - f"at most {self.padded_buffer.shape[0]}" - ) + # The scheduler already enforces max_num_batched_tokens and max_num_seqs. + # Python comparisons here would specialize these symbolic dimensions and + # violate V2's dynamic-shape contract during AOT compilation. positions = self.positions_buffer[:num_tokens] packed = self.padded_buffer[:num_reqs] @@ -593,9 +586,7 @@ def forward( 0, packed.shape[1] - 1 ) packed[request_indices, columns] = input_ids - ngram_context = ngram_context[:num_reqs].to( - device=input_ids.device, dtype=torch.long - ) + ngram_context = ngram_context.to(device=input_ids.device, dtype=torch.long) context = torch.cat([ngram_context, packed], dim=-1) positions_2d, position_in_segment = self._shift_precompute( From 383ff458dbc8af1dda79a1a29bf700019030f054 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:03:49 +0800 Subject: [PATCH 16/28] [Kernel][SM70] Bypass single-token NVFP4 MoE sorting Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- .../test_sm70_modelopt_mixed_nvfp4.py | 52 ++++++ .../layers/quantization/nvfp4_sm70_moe.py | 175 ++++++++++++++---- 2 files changed, 193 insertions(+), 34 deletions(-) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 35982e851f..469e8aadb9 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -19,6 +19,8 @@ from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( ModelOptNvFp4SM70MoEMethod, _prepare_compact_slot_groups, + _prepare_single_token_slots, + _single_token_weighted_reduce, _validate_weight_layout, validate_nvfp4_sm70_moe_contract, ) @@ -194,6 +196,56 @@ def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slot assert torch.equal(active_expert_ids.cpu(), sorted_expert_ids.cpu()) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), + reason="requires an exact SM70 CUDA device", +) +def test_nvfp4_single_token_direct_routing_is_exact_and_graph_dynamic(): + top_k = 10 + hidden = 2560 + x = torch.randn(1, hidden, dtype=torch.float16, device="cuda") + topk_ids = torch.tensor( + [[401, 7, 310, 99, 211, 3, 470, 120, 55, 256]], + dtype=torch.int64, + device="cuda", + ) + weights = torch.softmax(torch.randn(1, top_k, device="cuda"), dim=-1) + expert_output = torch.randn(top_k, hidden, dtype=torch.float16, device="cuda") + expanded = torch.empty_like(expert_output) + active_ids = torch.empty(top_k, dtype=torch.int32, device="cuda") + output = torch.empty(1, hidden, dtype=torch.float16, device="cuda") + reference = torch.empty_like(output) + identity = torch.arange(top_k, dtype=torch.int32, device="cuda").view(1, -1) + + _prepare_single_token_slots(x, topk_ids, expanded, active_ids) + _single_token_weighted_reduce(expert_output, weights, output) + torch.ops._moe_C.moe_unpermute( + expert_output, weights, identity, None, top_k, reference + ) + torch.cuda.synchronize() + assert torch.equal(expanded, x.expand(top_k, -1)) + assert torch.equal(active_ids, topk_ids[0].int()) + assert torch.equal(output, reference) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _prepare_single_token_slots(x, topk_ids, expanded, active_ids) + _single_token_weighted_reduce(expert_output, weights, output) + + x.add_(0.25) + topk_ids.copy_(topk_ids.roll(1, dims=1)) + weights.copy_(weights.roll(2, dims=1)) + expert_output.mul_(0.5) + graph.replay() + torch.ops._moe_C.moe_unpermute( + expert_output, weights, identity, None, top_k, reference + ) + torch.cuda.synchronize() + assert torch.equal(expanded, x.expand(top_k, -1)) + assert torch.equal(active_ids, topk_ids[0].int()) + assert torch.equal(output, reference) + + def test_mixed_w4a16_moe_requires_turbomind_on_sm70(): config = _mixed_config() diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index ccc73f908c..938439ddb9 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -49,6 +49,92 @@ _MAX_SUPPORTED_TOP_K: Final = max(contract[3] for contract in _SUPPORTED_CONTRACTS) +@triton.jit +def _prepare_single_token_slots_kernel( + input_ptr, + topk_ids_ptr, + expanded_input_ptr, + active_expert_ids_ptr, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + slot = tl.program_id(0) + offsets = tl.arange(0, BLOCK) + mask = offsets < HIDDEN + values = tl.load(input_ptr + offsets, mask=mask, other=0.0) + tl.store(expanded_input_ptr + slot * HIDDEN + offsets, values, mask=mask) + expert_id = tl.load(topk_ids_ptr + slot) + tl.store(active_expert_ids_ptr + slot, expert_id.to(tl.int32)) + + +def _prepare_single_token_slots( + x: torch.Tensor, + topk_ids: torch.Tensor, + expanded_input: torch.Tensor, + active_expert_ids: torch.Tensor, +) -> None: + top_k = topk_ids.numel() + hidden = x.shape[1] + if x.shape[0] != 1 or tuple(topk_ids.shape) != (1, top_k): + raise ValueError("SM70 NVFP4 direct routing requires one input token.") + if tuple(expanded_input.shape) != (top_k, hidden): + raise ValueError("SM70 NVFP4 direct routing buffer shape mismatch.") + if active_expert_ids.numel() != top_k: + raise ValueError("SM70 NVFP4 direct expert-ID buffer shape mismatch.") + _prepare_single_token_slots_kernel[(top_k,)]( + x, + topk_ids, + expanded_input, + active_expert_ids, + HIDDEN=hidden, + BLOCK=triton.next_power_of_2(hidden), + num_warps=8, + ) + + +@triton.jit +def _single_token_weighted_reduce_kernel( + expert_output_ptr, + topk_weights_ptr, + output_ptr, + HIDDEN: tl.constexpr, + TOP_K: tl.constexpr, + BLOCK: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < HIDDEN + acc = tl.zeros((BLOCK,), tl.float32) + for slot in tl.static_range(0, TOP_K): + values = tl.load( + expert_output_ptr + slot * HIDDEN + offsets, + mask=mask, + other=0.0, + ) + weight = tl.load(topk_weights_ptr + slot) + acc += values.to(tl.float32) * weight + tl.store(output_ptr + offsets, acc, mask=mask) + + +def _single_token_weighted_reduce( + expert_output: torch.Tensor, + topk_weights: torch.Tensor, + output: torch.Tensor, +) -> None: + top_k, hidden = expert_output.shape + if tuple(topk_weights.shape) != (1, top_k) or tuple(output.shape) != (1, hidden): + raise ValueError("SM70 NVFP4 direct weighted-reduce shape mismatch.") + block = 256 + _single_token_weighted_reduce_kernel[(triton.cdiv(hidden, block),)]( + expert_output, + topk_weights, + output, + HIDDEN=hidden, + TOP_K=top_k, + BLOCK=block, + num_warps=4, + ) + + @triton.jit def _prepare_compact_slot_groups_kernel( sorted_expert_ids_ptr, @@ -404,7 +490,7 @@ def _allocate_graph_safe_decode_buffers(self, layer: RoutedExperts) -> None: layer._nvfp4_sm70_dense_expert_ids = torch.arange( experts, dtype=torch.int32, device=device ) - layer._nvfp4_sm70_compact_offsets = torch.empty( + layer._nvfp4_sm70_compact_offsets = torch.arange( max_slots + 1, dtype=torch.int32, device=device ) layer._nvfp4_sm70_active_expert_ids = torch.empty( @@ -491,7 +577,9 @@ def _eager_buffers( "sorted_row_idx": torch.empty(slots, dtype=torch.int32, device=device), "topk_ids_for_sort": torch.empty(slots, dtype=torch.int32, device=device), "dense_expert_ids": layer._nvfp4_sm70_dense_expert_ids, - "compact_offsets": torch.empty(slots + 1, dtype=torch.int32, device=device), + "compact_offsets": torch.arange( + slots + 1, dtype=torch.int32, device=device + ), "active_expert_ids": torch.empty(slots, dtype=torch.int32, device=device), } @@ -549,31 +637,45 @@ def apply( return x.new_empty((0, hidden)) buffers = self._get_buffers(layer, num_tokens) output = buffers["output"] - output.zero_() slots = num_tokens * top_k - topk_ids_i32 = buffers["topk_ids"] - topk_ids_i32.copy_(topk_ids, non_blocking=True) - buffers["permuted_idx"].fill_(slots) - torch.ops._moe_C.moe_permute_with_scratch( - x, - topk_ids_i32, - buffers["token_expert_indices"], - layer.expert_map, - layer.global_num_experts, - layer.local_num_experts, - top_k, - buffers["permuted_input"], - buffers["expert_offsets64"], - buffers["inv_permuted_idx"], - buffers["permuted_idx"], - buffers["sort_workspace"], - buffers["permuted_experts_id"], - buffers["sorted_row_idx"], - buffers["topk_ids_for_sort"], - ) - buffers["expert_offsets"].copy_(buffers["expert_offsets64"], non_blocking=True) + direct_single_token = num_tokens == 1 + if direct_single_token: + _prepare_single_token_slots( + x, + topk_ids, + buffers["permuted_input"], + buffers["active_expert_ids"], + ) + stage_offsets = buffers["compact_offsets"] + stage_expert_ids = buffers["active_expert_ids"] + stage_experts = top_k + else: + output.zero_() + topk_ids_i32 = buffers["topk_ids"] + topk_ids_i32.copy_(topk_ids, non_blocking=True) + buffers["permuted_idx"].fill_(slots) + torch.ops._moe_C.moe_permute_with_scratch( + x, + topk_ids_i32, + buffers["token_expert_indices"], + layer.expert_map, + layer.global_num_experts, + layer.local_num_experts, + top_k, + buffers["permuted_input"], + buffers["expert_offsets64"], + buffers["inv_permuted_idx"], + buffers["permuted_idx"], + buffers["sort_workspace"], + buffers["permuted_experts_id"], + buffers["sorted_row_idx"], + buffers["topk_ids_for_sort"], + ) + buffers["expert_offsets"].copy_( + buffers["expert_offsets64"], non_blocking=True + ) - if num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: + if not direct_single_token and num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: _prepare_compact_slot_groups( buffers["permuted_experts_id"], buffers["compact_offsets"], @@ -582,7 +684,7 @@ def apply( stage_offsets = buffers["compact_offsets"] stage_expert_ids = buffers["active_expert_ids"] stage_experts = slots - else: + elif not direct_single_token: stage_offsets = buffers["expert_offsets"] stage_expert_ids = buffers["dense_expert_ids"] stage_experts = int(layer.sm70_nvfp4_num_experts) @@ -612,14 +714,19 @@ def apply( layer.sm70_nvfp4_w2_n_dim, layer.sm70_nvfp4_group_size, ) - torch.ops._moe_C.moe_unpermute( - buffers["sorted_output"], - topk_weights, - buffers["inv_permuted_idx"], - buffers["expert_offsets64"], - top_k, - output, - ) + if direct_single_token: + _single_token_weighted_reduce( + buffers["sorted_output"], topk_weights, output + ) + else: + torch.ops._moe_C.moe_unpermute( + buffers["sorted_output"], + topk_weights, + buffers["inv_permuted_idx"], + buffers["expert_offsets64"], + top_k, + output, + ) return output def apply_monolithic( From da7beb2c72fb6df44409aa05cf50630a9d553fdc Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:29:25 +0800 Subject: [PATCH 17/28] [Model][SM70] Enable Qwen4Exp native MTP Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_config.py | 29 ++- tests/v1/spec_decode/test_qwen4_exp.py | 245 +++++++++++++++++++++++ vllm/config/speculative.py | 27 +++ vllm/model_executor/models/config.py | 24 ++- vllm/model_executor/models/registry.py | 1 + vllm/v1/spec_decode/llm_base_proposer.py | 6 +- vllm/v1/spec_decode/qwen4_exp.py | 178 ++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 16 +- 8 files changed, 517 insertions(+), 9 deletions(-) create mode 100644 tests/v1/spec_decode/test_qwen4_exp.py create mode 100644 vllm/v1/spec_decode/qwen4_exp.py diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index f7abe35682..ba72c90eb8 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -102,7 +102,7 @@ def test_qwen4_exp_defaults_to_v2_even_when_quantized_moe( assert VllmConfig._is_default_v2_model_runner_model(vllm_config) -def test_initial_sm70_v2_route_rejects_speculative_decode( +def test_initial_sm70_v2_route_accepts_native_mtp( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( @@ -124,7 +124,32 @@ def test_initial_sm70_v2_route_rejects_speculative_decode( speculative_config=SimpleNamespace(method="mtp"), ) - with pytest.raises(NotImplementedError, match="initial SM70 V2 route"): + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + + +def test_initial_sm70_v2_route_rejects_unvalidated_speculator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + Qwen3_5ForConditionalGenerationConfig, + "verify_and_update_config", + lambda _config: None, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace( + hc_count=4, + ple_layer_ids=[2], + indexer_n_heads=4, + ), + multimodal_config=None, + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), + speculative_config=SimpleNamespace(method="dflash"), + ) + + with pytest.raises(NotImplementedError, match="supports only its native MTP"): Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) diff --git a/tests/v1/spec_decode/test_qwen4_exp.py b/tests/v1/spec_decode/test_qwen4_exp.py new file mode 100644 index 0000000000..c430b22199 --- /dev/null +++ b/tests/v1/spec_decode/test_qwen4_exp.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the Qwen4Exp MTP proposer cache topology.""" + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.spec_decode.qwen4_exp as qwen_proposer +from vllm.config.speculative import SpeculativeConfig +from vllm.models.qwen4_exp.common.qsa_cache import ( + circular_qsa_slot_mapping, + compressed_qsa_slot_mapping, +) +from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + FullAttentionSpec, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.spec_decode.qwen4_exp import Qwen4ExpMTPProposer + +SCHEDULER_BLOCK_SIZE = 256 +KERNEL_BLOCK_SIZE = 64 +RAW_CAPACITY = 128 +MAIN_LAYER = "draft.mtp.layers.0.self_attn.attn" +RAW_LAYER = "draft.mtp.layers.0.self_attn.indexer.raw_key_cache" +COMPRESSED_LAYER = "draft.mtp.layers.0.self_attn.indexer.compressed_key_cache" + + +class _FakeBackend: + def __init__(self, name: str) -> None: + self.name = name + + def full_cls_name(self) -> tuple[str, str]: + return (__name__, self.name) + + +class _FakeAttentionGroup: + def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id): + self.backend = backend + self.layer_names = list(layer_names) + self.kv_cache_spec = kv_cache_spec + self.kv_cache_group_id = kv_cache_group_id + self.builder = SimpleNamespace(kv_cache_spec=kv_cache_spec) + + def create_metadata_builders(self, vllm_config, device, kernel_block_size=None): + if kernel_block_size is not None: + self.builder.kv_cache_spec = self.kv_cache_spec.copy_with_new_block_size( + kernel_block_size + ) + + def get_metadata_builder(self): + return self.builder + + +def _make_specs(): + main_spec = FullAttentionSpec( + block_size=SCHEDULER_BLOCK_SIZE, + num_kv_heads=1, + head_size=128, + head_size_v=128, + dtype=torch.bfloat16, + ) + compressed_spec = MLAAttentionSpec( + block_size=SCHEDULER_BLOCK_SIZE, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + compress_ratio=64, + ) + raw_spec = CircularBufferSpec( + block_size=RAW_CAPACITY, + num_kv_heads=1, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + ) + return main_spec, compressed_spec, raw_spec + + +def _make_proposer_and_config( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[Qwen4ExpMTPProposer, SimpleNamespace]: + main_spec, compressed_spec, raw_spec = _make_specs() + backends = { + MAIN_LAYER: _FakeBackend("MainBackend"), + RAW_LAYER: _FakeBackend("QSAStateBackend"), + COMPRESSED_LAYER: _FakeBackend("QSAStateBackend"), + } + fake_layers = { + name: SimpleNamespace( + get_attn_backend=lambda backend=backend: backend, + num_heads=1, + ) + for name, backend in backends.items() + } + monkeypatch.setattr( + qwen_proposer, + "get_layers_from_vllm_config", + lambda *args, **kwargs: fake_layers, + ) + monkeypatch.setattr(qwen_proposer, "AttentionGroup", _FakeAttentionGroup) + + proposer = Qwen4ExpMTPProposer.__new__(Qwen4ExpMTPProposer) + proposer.vllm_config = None + proposer.draft_model_config = SimpleNamespace( + hf_text_config=SimpleNamespace(mtp_num_hidden_layers=1) + ) + proposer.device = torch.device("cpu") + proposer._draft_attn_layer_names = {MAIN_LAYER, RAW_LAYER, COMPRESSED_LAYER} + proposer.kv_cache_gid = -1 + proposer.draft_attn_groups = [] + proposer.block_size = -1 + proposer._per_group_block_tables = {} + + config = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=[MAIN_LAYER, COMPRESSED_LAYER], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=SCHEDULER_BLOCK_SIZE, + kv_cache_specs={ + MAIN_LAYER: main_spec, + COMPRESSED_LAYER: compressed_spec, + }, + ), + ), + SimpleNamespace( + layer_names=[RAW_LAYER], + kv_cache_spec=UniformTypeKVCacheSpecs( + block_size=RAW_CAPACITY, + kv_cache_specs={RAW_LAYER: raw_spec}, + ), + ), + ] + ) + return proposer, config + + +def test_initializes_and_builds_current_qwen_cache_topology( + monkeypatch: pytest.MonkeyPatch, +) -> None: + proposer, config = _make_proposer_and_config(monkeypatch) + proposer.initialize_attn_backend( + config, + kernel_block_sizes=[KERNEL_BLOCK_SIZE, RAW_CAPACITY], + ) + + assert proposer.kv_cache_gid == 0 + assert proposer.block_size == KERNEL_BLOCK_SIZE + assert [group.kv_cache_group_id for group in proposer.draft_attn_groups] == [ + 0, + 0, + 1, + ] + + logical_positions = torch.tensor([62, 63, 64, 65], dtype=torch.int64) + token_to_req = torch.zeros(4, dtype=torch.int32) + + for group in proposer.draft_attn_groups: + + def build_for_drafting(*, common_attn_metadata, draft_index, group=group): + spec = group.builder.kv_cache_spec + if isinstance(spec, CircularBufferSpec): + slots = circular_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + spec.block_size, + query_start_loc=common_attn_metadata.query_start_loc, + ) + elif isinstance(spec, MLAAttentionSpec): + slots = compressed_qsa_slot_mapping( + common_attn_metadata.block_table_tensor, + token_to_req, + logical_positions, + spec.storage_block_size, + spec.compress_ratio, + ) + else: + slots = common_attn_metadata.slot_mapping + return SimpleNamespace( + common=common_attn_metadata, + draft_index=draft_index, + slot_mapping=slots, + ) + + group.builder.build_for_drafting = build_for_drafting + + main_table = torch.tensor([[10, 11, 12, 13]], dtype=torch.int32) + raw_table = torch.tensor([[20]], dtype=torch.int32) + main_slots = torch.tensor([702, 703, 704, 705], dtype=torch.int64) + proposer.set_per_group_block_table(1, raw_table) + common = SimpleNamespace( + num_reqs=1, + query_start_loc=torch.tensor([0, 4], dtype=torch.int32), + block_table_tensor=main_table, + slot_mapping=main_slots, + ) + + _, per_layer = proposer.build_per_group_and_layer_attn_metadata( + common, + draft_index=2, + ) + + assert torch.equal(per_layer[MAIN_LAYER].slot_mapping, main_slots) + assert torch.equal( + per_layer[COMPRESSED_LAYER].slot_mapping, + torch.tensor([-1, 10, -1, -1], dtype=torch.int64), + ) + assert torch.equal( + per_layer[RAW_LAYER].slot_mapping, + torch.tensor([2622, 2623, 2624, 2625], dtype=torch.int64), + ) + + +def test_rejects_multiple_mtp_layers(monkeypatch: pytest.MonkeyPatch) -> None: + proposer, config = _make_proposer_and_config(monkeypatch) + proposer.draft_model_config.hf_text_config.mtp_num_hidden_layers = 2 + + with pytest.raises(NotImplementedError, match="only supports one MTP layer"): + proposer.initialize_attn_backend( + config, + kernel_block_sizes=[KERNEL_BLOCK_SIZE, RAW_CAPACITY], + ) + + +def test_qwen_proposer_hidden_and_config_contract() -> None: + proposer = Qwen4ExpMTPProposer.__new__(Qwen4ExpMTPProposer) + proposer.draft_model_config = SimpleNamespace( + hf_config=SimpleNamespace(hc_mult=4), + get_hidden_size=lambda: 1024, + ) + config = SimpleNamespace( + method="mtp", + draft_model_config=SimpleNamespace( + hf_config=SimpleNamespace(model_type="qwen4_exp_mtp") + ), + ) + + assert proposer._get_hidden_size() == 4096 + assert proposer.model_returns_tuple() + assert SpeculativeConfig.use_qwen4_exp_mtp(config) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index c74aeb882e..07507b9be4 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -45,6 +45,7 @@ "exaone_moe_mtp", "exaone4_5_mtp", "qwen3_next_mtp", + "qwen4_exp_mtp", "qwen3_5_mtp", "longcat_flash_mtp", "mtp", @@ -516,6 +517,23 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: {"n_predict": n_predict, "architectures": ["Qwen3NextMTP"]} ) + if hf_config.model_type in {"qwen4_exp", "qwen4_exp_text"}: + hf_config.model_type = "qwen4_exp_mtp" + if hf_config.model_type == "qwen4_exp_mtp": + text_config = get_hf_text_config(hf_config) + n_predict = getattr( + text_config, + "mtp_num_hidden_layers", + getattr(text_config, "num_nextn_predict_layers", None), + ) + hf_config.update( + { + "hc_mult": int(text_config.hc_count), + "n_predict": n_predict, + "architectures": ["Qwen4ExpMTP"], + } + ) + if hf_config.model_type == "exaone_moe": hf_config.model_type = "exaone_moe_mtp" if hf_config.model_type == "exaone_moe_mtp": @@ -1284,6 +1302,15 @@ def use_step3p5_mtp(self) -> bool: == "step3p5_mtp" ) + def use_qwen4_exp_mtp(self) -> bool: + """Return whether Qwen4Exp needs its dedicated proposer.""" + return ( + self.method == "mtp" + and self.draft_model_config is not None + and getattr(self.draft_model_config.hf_config, "model_type", None) + == "qwen4_exp_mtp" + ) + def use_eagle(self) -> bool: return ( self.method in ("eagle", "eagle3", "mtp") diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 7674b2c539..dc0444b74c 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -623,12 +623,14 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: _strip_qwen4_exp_mrope(model_config) spec_config = vllm_config.speculative_config - if spec_config is not None: + if spec_config is not None and spec_config.method not in { + "mtp", + "ngram", + "ngram_gpu", + }: raise NotImplementedError( - "Qwen4Exp speculative decoding is disabled for the initial " - "SM70 V2 route. The checkpoint's PLE n-gram embedding remains " - "enabled; disable --speculative-config until the native MTP " - "follow-up is quality-gated." + "Qwen4Exp speculative decoding supports only its native MTP " + "checkpoint and linear n-gram proposers" ) @@ -639,6 +641,17 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: _strip_qwen4_exp_mrope(vllm_config.model_config) +class Qwen4ExpMTPConfig(Qwen4ExpForConditionalGenerationConfig): + """Preserve MRoPE for a VL target and use 1D RoPE for a text target.""" + + @staticmethod + def verify_and_update_config(vllm_config: "VllmConfig") -> None: + Qwen4ExpForConditionalGenerationConfig.verify_and_update_config(vllm_config) + if hasattr(vllm_config.model_config.hf_config, "vision_config"): + return + _strip_qwen4_exp_mrope(vllm_config.model_config) + + class SnowflakeGteNewModelConfig(VerifyAndUpdateConfig): @staticmethod def verify_and_update_model_config(model_config: "ModelConfig") -> None: @@ -701,6 +714,7 @@ def verify_and_update_model_config(model_config: "ModelConfig") -> None: "Qwen3_5MoeForConditionalGeneration": Qwen3_5ForConditionalGenerationConfig, "Qwen4ExpForCausalLM": Qwen4ExpForCausalLMConfig, "Qwen4ExpForConditionalGeneration": Qwen4ExpForConditionalGenerationConfig, + "Qwen4ExpMTP": Qwen4ExpMTPConfig, "VoyageQwen3BidirectionalEmbedModel": VoyageQwen3BidirectionalEmbedModelConfig, "XLMRobertaModel": JinaRobertaModelConfig, } diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index e053e2def8..80beb6e8fe 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -640,6 +640,7 @@ "MedusaModel": ("medusa", "Medusa"), "OpenPanguMTPModel": ("openpangu_mtp", "OpenPanguMTP"), "Qwen3NextMTP": ("qwen3_next_mtp", "Qwen3NextMTP"), + "Qwen4ExpMTP": ("vllm.models.qwen4_exp", "Qwen4ExpMTP"), "Step3p5MTP": ("step3p5_mtp", "Step3p5MTP"), "Qwen3_5MTP": ("qwen3_5_mtp", "Qwen3_5MTP"), "Qwen3_5MoeMTP": ("qwen3_5_mtp", "Qwen3_5MoeMTP"), diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 98bed3ac9d..a643262407 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -231,7 +231,7 @@ def __init__( # We need to get the hidden size from the draft model config because # the draft model's hidden size can be different from the target model's # hidden size (e.g., Llama 3.3 70B). - self.hidden_size = self.draft_model_config.get_hidden_size() + self.hidden_size = self._get_hidden_size() self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size() # DeepSeek V4 MTP consumes the target's pre-hc_head residual stream, @@ -461,6 +461,10 @@ def __init__( self.allowed_attn_types = tuple(rocm_types) + def _get_hidden_size(self) -> int: + """Return the hidden width consumed by the draft model.""" + return self.draft_model_config.get_hidden_size() + def _raise_if_padded_drafter_batch_disabled(self): if self.speculative_config.disable_padded_drafter_batch: raise NotImplementedError( diff --git a/vllm/v1/spec_decode/qwen4_exp.py b/vllm/v1/spec_decode/qwen4_exp.py new file mode 100644 index 0000000000..a1792b2c27 --- /dev/null +++ b/vllm/v1/spec_decode/qwen4_exp.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from copy import copy + +import torch + +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.worker.utils import AttentionGroup + + +class Qwen4ExpMTPProposer(EagleProposer): + """Speculative decoding proposer for Qwen4Exp MTP.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + runner=None, + ) -> None: + super().__init__(vllm_config, device, runner) + self._per_group_block_tables: dict[int, torch.Tensor] = {} + + def _get_hidden_size(self) -> int: + """Return the multi-stream feedback width consumed by Qwen MTP.""" + return int( + self.draft_model_config.get_hidden_size() + * self.draft_model_config.hf_config.hc_mult + ) + + def model_returns_tuple(self) -> bool: + """Qwen MTP returns separate logits and feedback hidden states.""" + return True + + def set_per_group_block_table( + self, + gid: int, + block_table: torch.Tensor, + ) -> None: + """Stage one scheduler group's block table for drafting.""" + self._per_group_block_tables[gid] = block_table + + def build_per_group_and_layer_attn_metadata( + self, + common_attn_metadata: CommonAttentionMetadata, + draft_index: int = 0, + ) -> tuple[list[object], dict[str, object]]: + """Build each Qwen cache owner's metadata from its scheduler group.""" + per_group_attn_metadata: list[object] = [] + per_layer_attn_metadata: dict[str, object] = {} + common_by_gid: dict[int, CommonAttentionMetadata] = {} + num_reqs = common_attn_metadata.num_reqs + + for attn_group in self.draft_attn_groups: + gid = attn_group.kv_cache_group_id + group_common = common_by_gid.get(gid) + if group_common is None: + if gid == self.kv_cache_gid: + group_common = common_attn_metadata + else: + block_table = self._per_group_block_tables.get(gid) + assert block_table is not None, ( + f"Missing Qwen draft block table for KV cache group {gid}" + ) + group_common = copy(common_attn_metadata) + group_common.block_table_tensor = block_table[:num_reqs] + common_by_gid[gid] = group_common + + attn_metadata = attn_group.get_metadata_builder().build_for_drafting( + common_attn_metadata=group_common, + draft_index=draft_index, + ) + per_group_attn_metadata.append(attn_metadata) + for layer_name in attn_group.layer_names: + per_layer_attn_metadata[layer_name] = attn_metadata + + return per_group_attn_metadata, per_layer_attn_metadata + + def initialize_attn_backend( + self, + kv_cache_config: KVCacheConfig, + kernel_block_sizes: list[int] | None = None, + ) -> None: + """Initialize Qwen main, compressed, and circular cache owners.""" + num_mtp_layers = self.draft_model_config.hf_text_config.mtp_num_hidden_layers + if num_mtp_layers != 1: + raise NotImplementedError( + "Qwen4Exp MTP proposer only supports one MTP layer" + ) + assert kernel_block_sizes is not None, ( + "Qwen MTP requires resolved kernel block sizes" + ) + assert len(kernel_block_sizes) == len(kv_cache_config.kv_cache_groups), ( + "Qwen MTP requires one kernel block size per KV cache group" + ) + + all_attn_layers = get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + ) + layer_to_gid, layer_to_spec = self._map_draft_layers_to_groups(kv_cache_config) + main_layers = [ + name + for name, spec in layer_to_spec.items() + if type(spec) is FullAttentionSpec + ] + assert len(main_layers) == 1, ( + "Qwen4Exp MTP requires exactly one main cache owner" + ) + self.kv_cache_gid = layer_to_gid[main_layers[0]] + + attention_groups: list[AttentionGroup] = [] + for layer_name in sorted(self._draft_attn_layer_names): + attn_layer = all_attn_layers[layer_name] + gid = layer_to_gid[layer_name] + attn_group = AttentionGroup( + backend=attn_layer.get_attn_backend(), + layer_names=[layer_name], + kv_cache_spec=layer_to_spec[layer_name], + kv_cache_group_id=gid, + ) + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_size=kernel_block_sizes[gid], + ) + attention_groups.append(attn_group) + + self.draft_attn_groups = sorted( + attention_groups, + key=lambda group: ( + group.kv_cache_group_id != self.kv_cache_gid, + group.kv_cache_group_id, + group.backend.full_cls_name(), + group.layer_names[0], + ), + ) + self.block_size = kernel_block_sizes[self.kv_cache_gid] + + def _map_draft_layers_to_groups( + self, + kv_cache_config: KVCacheConfig, + ) -> tuple[dict[str, int], dict[str, KVCacheSpec]]: + """Map Qwen draft cache owners to scheduler groups and concrete specs.""" + layer_to_gid: dict[str, int] = {} + layer_to_spec: dict[str, KVCacheSpec] = {} + for gid, group in enumerate(kv_cache_config.kv_cache_groups): + group_spec = group.kv_cache_spec + for layer_name in group.layer_names: + if layer_name not in self._draft_attn_layer_names: + continue + assert isinstance(group_spec, UniformTypeKVCacheSpecs), ( + "Qwen draft cache owners require packed KV cache groups" + ) + spec = group_spec.kv_cache_specs.get(layer_name) + assert spec is not None, ( + f"Qwen draft cache group {gid} has no spec for {layer_name}" + ) + layer_to_gid[layer_name] = gid + layer_to_spec[layer_name] = spec + + assert layer_to_spec.keys() == self._draft_attn_layer_names, ( + "Qwen draft KV cache configuration is missing layers: " + f"{sorted(self._draft_attn_layer_names - layer_to_spec.keys())}" + ) + return layer_to_gid, layer_to_spec + + +__all__ = ["Qwen4ExpMTPProposer"] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 53efe8b579..8ada303e4d 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -219,6 +219,7 @@ update_ngram_gpu_tensors_incremental, update_scheduler_for_invalid_drafts, ) +from vllm.v1.spec_decode.qwen4_exp import Qwen4ExpMTPProposer from vllm.v1.spec_decode.static_draft_vocab import ( DynamicDraftVocabPrefillBootstrapState, resolve_mtp_draft_vocab_config, @@ -1518,6 +1519,7 @@ def __init__( | ExtractHiddenStatesProposer | Gemma4Proposer | Step3p5MTPProposer + | Qwen4ExpMTPProposer ) if self.speculative_config.method == "custom_class": self.drafter = create_custom_proposer( # type: ignore[assignment] @@ -1554,6 +1556,10 @@ def __init__( self.drafter = Gemma4Proposer(self.vllm_config, self.device, self) elif self.speculative_config.use_step3p5_mtp(): self.drafter = Step3p5MTPProposer(self.vllm_config, self.device, self) + elif self.speculative_config.use_qwen4_exp_mtp(): + self.drafter = Qwen4ExpMTPProposer( + self.vllm_config, self.device, self + ) elif self.speculative_config.use_dspark(): self.drafter = DSparkProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = True @@ -5664,7 +5670,15 @@ def _build_attn_group_metadata( else: spec_decode_common_attn_metadata = cm # Capture per-group block tables for multi-group proposers. - if self.speculative_config and isinstance(self.drafter, Step3p5MTPProposer): + if self.speculative_config and isinstance( + self.drafter, Qwen4ExpMTPProposer + ): + self.drafter.set_per_group_block_table( + kv_cache_gid, cm.block_table_tensor + ) + elif self.speculative_config and isinstance( + self.drafter, Step3p5MTPProposer + ): self.drafter.set_per_group_attn_metadata( kv_cache_gid, cm.block_table_tensor, cm.slot_mapping ) From aa5384d298d15ac1f0c29e8ab48eeecf305120f0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:41:56 +0800 Subject: [PATCH 18/28] [Core][SM70] Enable Qwen4Exp prefix caching Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_config.py | 4 ++-- tests/v1/core/test_qwen4_exp_kv_cache.py | 16 ++++++++++++++++ vllm/model_executor/models/config.py | 6 ------ vllm/v1/core/kv_cache_coordinator.py | 6 ++++-- vllm/v1/kv_cache_interface.py | 12 ++++++++++++ 5 files changed, 34 insertions(+), 10 deletions(-) diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index ba72c90eb8..c85602a401 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -22,7 +22,7 @@ ) -def test_initial_sm70_route_accepts_v2_runner( +def test_sm70_v2_route_accepts_prefix_caching( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( @@ -42,7 +42,7 @@ def test_initial_sm70_route_accepts_v2_runner( hf_text_config=text_config, multimodal_config=SimpleNamespace(language_model_only=True), ), - cache_config=SimpleNamespace(enable_prefix_caching=False), + cache_config=SimpleNamespace(enable_prefix_caching=True), parallel_config=SimpleNamespace(enable_dbo=False, ubatch_size=1), speculative_config=None, use_v2_model_runner=True, diff --git a/tests/v1/core/test_qwen4_exp_kv_cache.py b/tests/v1/core/test_qwen4_exp_kv_cache.py index 1763bb1dda..4b777dc2e5 100644 --- a/tests/v1/core/test_qwen4_exp_kv_cache.py +++ b/tests/v1/core/test_qwen4_exp_kv_cache.py @@ -127,6 +127,22 @@ def test_qwen4_exp_csa_linear_cache_layout() -> None: ) assert isinstance(coordinator.single_type_managers[1], CircularBufferManager) + prefix_coordinator = get_kv_cache_coordinator( + scheduler_config, + max_model_len=8192, + max_in_flight_tokens=128, + use_eagle=False, + enable_caching=True, + enable_kv_cache_events=False, + dcp_world_size=1, + pcp_world_size=1, + hash_block_size=4, + ) + assert all( + not isinstance(spec, CircularBufferSpec) + for spec, _, _ in prefix_coordinator.attention_groups + ) + def test_qwen4_exp_circular_cache_stores_keys_without_unused_values() -> None: spec = CircularBufferSpec( diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index dc0444b74c..1cb631a8af 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -594,12 +594,6 @@ def verify_and_update_config(vllm_config: "VllmConfig") -> None: if text_config.hc_count <= 1: raise ValueError("Qwen4Exp requires hc_count > 1") - if vllm_config.cache_config.enable_prefix_caching: - raise NotImplementedError( - "Qwen4Exp prefix caching is not enabled in the initial SM70 " - "route; disable it while the QSA ring cache is in use" - ) - parallel_config = vllm_config.parallel_config uses_ple_or_qsa = bool(text_config.ple_layer_ids) or ( getattr(text_config, "indexer_n_heads", None) is not None diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 224b4d98d7..eb0eec34dc 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -451,6 +451,8 @@ def verify_and_split_kv_cache_groups(self) -> None: ] = [] for i, g in enumerate(self.kv_cache_config.kv_cache_groups): + if not g.kv_cache_spec.prefix_cacheable: + continue manager_cls = self.single_type_managers[i].__class__ spec = g.kv_cache_spec @@ -465,8 +467,8 @@ def verify_and_split_kv_cache_groups(self) -> None: else: attention_groups.append((spec, [i], manager_cls)) - assert len(attention_groups) > 1, ( - "HybridKVCacheCoordinator requires at least two attention groups." + assert attention_groups, ( + "HybridKVCacheCoordinator requires at least one cacheable group." ) # Put full attention first: its efficient left-to-right scan provides diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d55f5a9e41..d6ae8e0fc2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -100,6 +100,10 @@ class KVCacheSpec: # number of tokens in a block block_size: int + @property + def prefix_cacheable(self) -> bool: + return True + @property def page_size_bytes(self) -> int: """ @@ -557,6 +561,10 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: del vllm_config return self.page_size_bytes + @property + def prefix_cacheable(self) -> bool: + return False + @dataclass(frozen=True, kw_only=True) class SlidingWindowMLASpec(SlidingWindowSpec): @@ -740,6 +748,10 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): kv_cache_specs: dict[str, KVCacheSpec] + @property + def prefix_cacheable(self) -> bool: + return all(spec.prefix_cacheable for spec in self.kv_cache_specs.values()) + @property def page_size_bytes(self) -> int: return sum(spec.page_size_bytes for spec in self.kv_cache_specs.values()) From a21d15e88092a4678aff62605828df63551eb69c Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:37:09 +0800 Subject: [PATCH 19/28] [Spec Decode][SM70] Handle Qwen4Exp MTP feedback states Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_qwen4_exp_v2.py | 37 +++++++++++++++++++ .../gpu/spec_decode/eagle/speculator.py | 8 ++-- 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 13e9ee95bb..29ab4fe9bb 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -16,6 +17,42 @@ ) from vllm.v1.worker.gpu import model_runner as mrv2 from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.spec_decode.eagle import speculator as eagle_speculator + + +def test_qwen4_exp_mtp_v2_unpacks_logits_and_feedback_hidden_states( + monkeypatch: pytest.MonkeyPatch, +) -> None: + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.device = torch.device("cpu") + speculator.vllm_config = SimpleNamespace() + speculator.supports_mm_inputs = False + speculator.input_buffers = SimpleNamespace( + input_ids=torch.zeros(3, dtype=torch.int64), + positions=torch.arange(3, dtype=torch.int64), + ) + speculator.hidden_states = torch.zeros(3, 16) + + logits_hidden = torch.ones(3, 4) + feedback_hidden = torch.ones(3, 16) + speculator.model = lambda **_kwargs: (logits_hidden, feedback_hidden) + monkeypatch.setattr( + eagle_speculator, + "set_forward_context", + lambda *_args, **_kwargs: nullcontext(), + ) + + actual_logits_hidden, actual_feedback_hidden = speculator.run_model( + num_tokens=3, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + ) + + assert actual_logits_hidden is logits_hidden + assert actual_feedback_hidden is feedback_hidden def test_qsa_circular_group_uses_one_block_and_custom_slot_mapping( diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 8376a51558..3c4609d357 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -222,11 +222,13 @@ def run_model( hidden_states=self.hidden_states[:num_tokens], inputs_embeds=inputs_embeds, ) - if self.method == "mtp": + # Some MTP models declare a single-tensor contract but return + # (logits_hidden, feedback_hidden) for final-norm correctness. + if isinstance(ret_hidden_states, tuple): + last_hidden_states, hidden_states = ret_hidden_states + else: last_hidden_states = ret_hidden_states hidden_states = ret_hidden_states - else: - last_hidden_states, hidden_states = ret_hidden_states return last_hidden_states, hidden_states def _sample_draft( From eb9547ee4c397d33586bdb8fd82cff6f6e1131bf Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:49:22 +0800 Subject: [PATCH 20/28] [Core] Align Qwen4Exp short-conv metadata Port the applicable runtime hunk from origin/pr-338 commit 4de4737ad3 so CUDA graph capture passes the required bool dtype to async_tensor_h2d and uses the current AttentionSpec interface. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/v1/attention/backends/short_conv_attn.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 3e888e69b7..91e3e7bcbc 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -21,7 +21,7 @@ compute_causal_conv1d_metadata, mamba_get_block_table_tensor, ) -from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.kv_cache_interface import AttentionSpec class ShortConvAttentionBackend(AttentionBackend): @@ -111,7 +111,7 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): def __init__( self, - kv_cache_spec: MambaSpec, + kv_cache_spec: AttentionSpec, layer_names: list[str], vllm_config: VllmConfig, device: torch.device, @@ -284,7 +284,9 @@ def build( # type: ignore[override] spec_sequence_masks = spec_sequence_masks_cpu else: spec_sequence_masks = async_tensor_h2d( - spec_sequence_masks_cpu, device=query_start_loc.device + spec_sequence_masks_cpu, + dtype=torch.bool, + device=query_start_loc.device, ) # For causal_conv1d (non-spec prefill Triton kernel metadata). From 6097689ea3f4ef004c36ec8a66ed456f0fa5431f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:58:32 +0800 Subject: [PATCH 21/28] [Bugfix] Stage Qwen4Exp spec masks from lists Convert the CPU bool tensor to the list contract required by async_tensor_h2d. This preserves the pinned non-blocking copy and fixes full CUDA graph capture after the PR 338 API alignment. Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- vllm/v1/attention/backends/short_conv_attn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/attention/backends/short_conv_attn.py b/vllm/v1/attention/backends/short_conv_attn.py index 91e3e7bcbc..efbbd1e075 100644 --- a/vllm/v1/attention/backends/short_conv_attn.py +++ b/vllm/v1/attention/backends/short_conv_attn.py @@ -284,7 +284,7 @@ def build( # type: ignore[override] spec_sequence_masks = spec_sequence_masks_cpu else: spec_sequence_masks = async_tensor_h2d( - spec_sequence_masks_cpu, + spec_sequence_masks_cpu.tolist(), dtype=torch.bool, device=query_start_loc.device, ) From d9a39ea43477e092a44cd7e7b4e21d5a25cf6cb6 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:28:50 +0800 Subject: [PATCH 22/28] [Core] Support heterogeneous Mamba state groups Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/v1/worker/test_mamba_utils.py | 94 ++++++++++- vllm/v1/worker/gpu/mamba_align.py | 6 +- .../worker/gpu/model_states/mamba_hybrid.py | 35 ++-- vllm/v1/worker/mamba_utils.py | 155 ++++++++++++------ 4 files changed, 223 insertions(+), 67 deletions(-) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index fc2ecea746..7c1d40500d 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -13,6 +13,7 @@ get_conv_copy_spec, get_temporal_copy_spec, ) +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec from vllm.v1.worker import mamba_utils as worker_mamba_utils @@ -469,12 +470,103 @@ def make_buffer(n, dtype): return MambaSpecDecodeGPUContext.create( max_num_reqs=cfg.max_num_reqs, kv_cache_config=kv_cache_config, - num_state_types=2, + mamba_state_copy_funcs=_COPY_FUNCS, device=device, make_buffer=make_buffer, ) +def test_mamba_context_supports_heterogeneous_state_groups(): + device = torch.device("cpu") + block_size = 16 + num_speculative_blocks = 4 + gdn_spec = MambaSpec( + block_size=block_size, + shapes=((4, 8), (16,)), + dtypes=(torch.float16, torch.float32), + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + ) + ple_spec = MambaSpec( + block_size=block_size, + shapes=((3, 8),), + dtypes=(torch.float16,), + mamba_type=MambaAttentionBackendEnum.SHORT_CONV, + mamba_cache_mode="align", + num_speculative_blocks=num_speculative_blocks, + tp_replicated=True, + ) + kv_cache_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=["gdn_0", "gdn_1"], kv_cache_spec=gdn_spec), + KVCacheGroupSpec(layer_names=["ple"], kv_cache_spec=ple_spec), + ], + ) + copy_funcs_by_type = { + MambaAttentionBackendEnum.GDN_ATTN: ( + get_conv_copy_spec, + get_temporal_copy_spec, + ), + MambaAttentionBackendEnum.SHORT_CONV: (get_conv_copy_spec,), + } + + def make_buffer(n, dtype): + return _MockCpuGpuBuffer(n, dtype, device) + + ctx = MambaSpecDecodeGPUContext.create( + max_num_reqs=2, + kv_cache_config=kv_cache_config, + mamba_state_copy_funcs=copy_funcs_by_type, + device=device, + make_buffer=make_buffer, + ) + copy_bufs = MambaCopyBuffers.create( + max_num_reqs=2, + kv_cache_config=kv_cache_config, + copy_funcs=copy_funcs_by_type, + make_buffer=make_buffer, + ) + + assert ctx.num_layers == 3 + assert ctx.num_states == 5 + assert ctx.state_base_addrs.numel() == 5 + assert copy_bufs.src_ptrs.cpu.numel() == 10 + + def attention(*states: torch.Tensor) -> MagicMock: + mock = MagicMock() + mock.kv_cache = list(states) + return mock + + forward_context = { + "gdn_0": attention( + torch.zeros(4, 4, 8, dtype=torch.float16), + torch.zeros(4, 16, dtype=torch.float32), + ), + "gdn_1": attention( + torch.zeros(4, 4, 8, dtype=torch.float16), + torch.zeros(4, 16, dtype=torch.float32), + ), + "ple": attention(torch.zeros(4, 3, 8, dtype=torch.float16)), + } + block_tables = [ + torch.zeros(2, 8, dtype=torch.int32), + torch.zeros(2, 8, dtype=torch.int32), + ] + ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + copy_funcs_by_type, + block_tables, + ) + + assert ctx.state_group_indices.tolist() == [0, 0, 0, 0, 1] + assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 3] + assert ctx.state_inner_sizes.tolist() == [8, 16, 8, 16, 8] + + def _run_gpu_postprocess( gpu_ctx: MambaSpecDecodeGPUContext, *, diff --git a/vllm/v1/worker/gpu/mamba_align.py b/vllm/v1/worker/gpu/mamba_align.py index 52a03e83c3..c24ce798a3 100644 --- a/vllm/v1/worker/gpu/mamba_align.py +++ b/vllm/v1/worker/gpu/mamba_align.py @@ -261,8 +261,7 @@ def run_mamba_align_precopy( ) -> None: if num_reqs == 0 or not ctx.is_initialized: return - total_states = ctx.num_layers * ctx.num_state_types - grid = (num_reqs, total_states, _TEMPORAL_TILES) + grid = (num_reqs, ctx.num_states, _TEMPORAL_TILES) _precopy_mamba_align_kernel[grid]( state_idx, src_col, @@ -296,8 +295,7 @@ def run_mamba_align_postprocess( # output count while programs copying other state tensors still read it. snapshot = ctx.num_accepted_tokens_out snapshot.copy_(num_accepted_tokens) - total_states = ctx.num_layers * ctx.num_state_types - grid = (num_reqs, total_states, _TEMPORAL_TILES) + grid = (num_reqs, ctx.num_states, _TEMPORAL_TILES) _postprocess_mamba_align_kernel[grid]( snapshot, num_accepted_tokens, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 2c55f6ac35..f38703a49c 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -52,7 +52,11 @@ from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata from vllm.v1.worker.gpu.spec_decode import uses_dflash_selector_engine -from vllm.v1.worker.mamba_utils import MambaSpecDecodeGPUContext +from vllm.v1.worker.mamba_utils import ( + MambaSpecDecodeGPUContext, + get_mamba_groups, + get_mamba_types, +) from vllm.v1.worker.utils import AttentionGroup logger = init_logger(__name__) @@ -194,16 +198,7 @@ def _get_mamba_group_info( self, kv_cache_config: KVCacheConfig ) -> tuple[list[int], MambaSpec]: if self._mamba_spec is None: - group_ids: list[int] = [] - specs: list[MambaSpec] = [] - for group_id, group in enumerate(kv_cache_config.kv_cache_groups): - if isinstance(group.kv_cache_spec, MambaSpec): - group_ids.append(group_id) - specs.append(group.kv_cache_spec) - assert specs, "no mamba layers in the model" - assert all(specs[0] == spec for spec in specs) - self._mamba_group_ids = group_ids - self._mamba_spec = specs[0] + self._mamba_group_ids, self._mamba_spec = get_mamba_groups(kv_cache_config) return self._mamba_group_ids, self._mamba_spec def _ensure_align_ctx( @@ -212,12 +207,18 @@ def _ensure_align_ctx( mamba_group_ids: list[int], block_tables: tuple[torch.Tensor, ...], ) -> MambaSpecDecodeGPUContext: + copy_funcs = None if self._mamba_ctx is None: - copy_funcs = self.model.get_mamba_state_copy_func() + copy_funcs = self.model.get_mamba_state_copy_funcs( + get_mamba_types(kv_cache_config) + ) # This V100 closure intentionally retains the tree's default SD # layout. The official DS-row extension is unrelated to Qwen3.8's # configured path and would expand the DDTree-sensitive closure. - if get_conv_copy_spec in copy_funcs and is_conv_state_dim_first(): + if ( + any(get_conv_copy_spec in funcs for funcs in copy_funcs.values()) + and is_conv_state_dim_first() + ): raise ValueError( "MRV2 align prefix caching with speculative decoding requires " "the default SD Mamba conv-state layout on this V100 path." @@ -225,7 +226,7 @@ def _ensure_align_ctx( self._mamba_ctx = MambaSpecDecodeGPUContext.create( max_num_reqs=self.max_num_reqs, kv_cache_config=kv_cache_config, - num_state_types=len(copy_funcs), + mamba_state_copy_funcs=copy_funcs, device=self.device, make_buffer=lambda n, dtype: CpuGpuBuffer( n, @@ -236,11 +237,15 @@ def _ensure_align_ctx( ) ctx = self._mamba_ctx if not ctx.is_initialized: + if copy_funcs is None: + copy_funcs = self.model.get_mamba_state_copy_funcs( + get_mamba_types(kv_cache_config) + ) forward_context = self.vllm_config.compilation_config.static_forward_context ctx.initialize_from_forward_context( kv_cache_config, forward_context, - self.model.get_mamba_state_copy_func(), + copy_funcs, [block_tables[group_id] for group_id in mamba_group_ids], ) return ctx diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7ffb0c28f9..ef363d610a 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -13,12 +13,14 @@ from vllm.config import CacheConfig from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, + MambaStateCopyFuncs, + MambaStateCopyFuncsByType, get_conv_copy_spec, get_temporal_copy_spec, ) from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.utils import CpuGpuBuffer @@ -77,8 +79,7 @@ def postprocess_mamba_fused_kernel( # block table. block_table_ptrs_ptr, block_table_stride_req: tl.int64, # stride between requests (in elements) - # Mamba state metadata (per-layer, per-state-type) - # These are 1D arrays indexed by (layer_idx * num_state_types + state_type_idx) + # Mamba state metadata, flattened across all layers and state types. state_base_addrs_ptr, # base address of each state tensor state_block_strides_ptr, # bytes per block for each state state_elem_sizes_ptr, # element size for each state @@ -101,13 +102,13 @@ def postprocess_mamba_fused_kernel( Fused GPU kernel for postprocess_mamba that computes decisions AND performs mamba state copies without any CPU-GPU synchronization. - Grid: (num_reqs, num_layers * num_state_types) + Grid: (num_reqs, num_states) - program_id(0) = request index - program_id(1) = state_idx (flattened index into layer/state_type metadata) - Note: num_layers and num_state_types are not passed as kernel parameters - because the kernel indexes directly into pre-flattened metadata arrays - using program_id(1). The grid dimensions encode the total state count. + The kernel indexes directly into pre-flattened metadata arrays using + program_id(1). The grid dimensions encode the total state count, allowing + different Mamba layer types to own different numbers of state tensors. """ req_idx = tl.program_id(0) state_idx = tl.program_id(1) @@ -307,6 +308,9 @@ def warmup_batch_memcpy_kernel(device: torch.device) -> bool: return True +MambaStateCopyFuncsInput = MambaStateCopyFuncs | MambaStateCopyFuncsByType + + def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSpec]: mamba_group_ids: list[int] = [] mamba_specs: list[MambaSpec] = [] @@ -316,10 +320,52 @@ def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSp mamba_group_ids.append(i) mamba_specs.append(kv_cache_spec) assert len(mamba_group_ids) > 0, "no mamba layers in the model" - assert all(mamba_specs[0] == spec for spec in mamba_specs) + # State shapes, dtypes, backend types, and TP replication may differ across + # groups (for example Qwen4Exp GDN versus its PLE short-conv state). The + # align scheduler and copy kernels only require shared block semantics. + for field in ("block_size", "mamba_cache_mode", "num_speculative_blocks"): + values = {getattr(spec, field) for spec in mamba_specs} + assert len(values) == 1, f"all mamba groups must share {field}, got {values}" return mamba_group_ids, mamba_specs[0] +def get_mamba_types( + kv_cache_config: KVCacheConfig, +) -> set[MambaAttentionBackendEnum]: + mamba_group_ids, _ = get_mamba_groups(kv_cache_config) + mamba_types: set[MambaAttentionBackendEnum] = set() + for group_id in mamba_group_ids: + spec = kv_cache_config.kv_cache_groups[group_id].kv_cache_spec + assert isinstance(spec, MambaSpec) + mamba_types.add(spec.mamba_type) + return mamba_types + + +def _get_copy_funcs_for_group( + kv_cache_config: KVCacheConfig, + mamba_group_id: int, + copy_funcs: MambaStateCopyFuncsInput, +) -> MambaStateCopyFuncs: + spec = kv_cache_config.kv_cache_groups[mamba_group_id].kv_cache_spec + assert isinstance(spec, MambaSpec) + if isinstance(copy_funcs, dict): + group_copy_funcs = copy_funcs[spec.mamba_type] + assert len(group_copy_funcs) == len(spec.shapes), ( + f"{spec.mamba_type} exposes {len(spec.shapes)} states but has " + f"{len(group_copy_funcs)} copy functions" + ) + return group_copy_funcs + + # Backward compatibility for homogeneous models and legacy callers. Some + # mixed models append a one-state short-conv group after a two-state GDN + # group, so selecting the matching prefix is safe for their old tuple. + assert len(copy_funcs) >= len(spec.shapes), ( + f"{spec.mamba_type} exposes {len(spec.shapes)} states but has only " + f"{len(copy_funcs)} copy functions" + ) + return copy_funcs[: len(spec.shapes)] + + @dataclasses.dataclass class MambaCopyBuffers: src_ptrs: CpuGpuBuffer @@ -334,15 +380,16 @@ def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, - copy_funcs: tuple[MambaStateCopyFunc, ...], + copy_funcs: MambaStateCopyFuncsInput, make_buffer: Callable[..., CpuGpuBuffer], copies_per_req: int = 1, ) -> "MambaCopyBuffers": mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config) entries_per_req = sum( len(kv_cache_config.kv_cache_groups[gid].layer_names) + * len(_get_copy_funcs_for_group(kv_cache_config, gid, copy_funcs)) for gid in mamba_group_ids - ) * len(copy_funcs) + ) n = max_num_reqs * entries_per_req * copies_per_req return cls( src_ptrs=make_buffer(n, dtype=torch.int64), @@ -369,7 +416,7 @@ class MambaSpecDecodeGPUContext: window with offset-based copies), 0 for temporal states (full block copies). """ - # Per-state metadata tensors (shape: [num_layers * num_state_types]) + # Per-state metadata tensors (shape: [num_states]) # These are populated from forward_context during the first forward pass state_base_addrs: torch.Tensor # int64: base address of each state tensor state_block_strides: torch.Tensor # int64: bytes per block @@ -381,7 +428,7 @@ class MambaSpecDecodeGPUContext: # Configuration block_size: int num_layers: int - num_state_types: int + num_states: int mamba_group_ids: list[int] num_groups: int @@ -412,7 +459,7 @@ def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, - num_state_types: int, + mamba_state_copy_funcs: MambaStateCopyFuncsInput, device: torch.device, make_buffer: Callable[..., CpuGpuBuffer], ) -> "MambaSpecDecodeGPUContext": @@ -424,30 +471,28 @@ def create( len(kv_cache_config.kv_cache_groups[gid].layer_names) for gid in mamba_group_ids ) - total_states = num_layers * num_state_types + num_states = sum( + len(kv_cache_config.kv_cache_groups[gid].layer_names) + * len( + _get_copy_funcs_for_group(kv_cache_config, gid, mamba_state_copy_funcs) + ) + for gid in mamba_group_ids + ) return cls( - state_base_addrs=torch.zeros( - total_states, dtype=torch.int64, device=device - ), + state_base_addrs=torch.zeros(num_states, dtype=torch.int64, device=device), state_block_strides=torch.zeros( - total_states, dtype=torch.int64, device=device - ), - state_elem_sizes=torch.zeros( - total_states, dtype=torch.int32, device=device - ), - state_inner_sizes=torch.zeros( - total_states, dtype=torch.int64, device=device - ), - state_conv_widths=torch.zeros( - total_states, dtype=torch.int32, device=device + num_states, dtype=torch.int64, device=device ), + state_elem_sizes=torch.zeros(num_states, dtype=torch.int32, device=device), + state_inner_sizes=torch.zeros(num_states, dtype=torch.int64, device=device), + state_conv_widths=torch.zeros(num_states, dtype=torch.int32, device=device), state_group_indices=torch.zeros( - total_states, dtype=torch.int32, device=device + num_states, dtype=torch.int32, device=device ), block_size=mamba_spec.block_size, num_layers=num_layers, - num_state_types=num_state_types, + num_states=num_states, mamba_group_ids=mamba_group_ids, num_groups=len(mamba_group_ids), num_accepted_tokens_out=torch.zeros( @@ -470,7 +515,7 @@ def initialize_from_forward_context( self, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, block_tables: list[torch.Tensor], ) -> None: """ @@ -501,8 +546,9 @@ def initialize_from_forward_context( forward_context: Dictionary mapping layer names to attention objects, populated after the model is loaded. Each attention object must have a `kv_cache` attribute containing the list of state tensors. - mamba_state_copy_funcs: Tuple of copy functions (one per state type) - used to determine whether each state is a conv or temporal state. + mamba_state_copy_funcs: Copy functions keyed by Mamba backend type, + or a legacy homogeneous tuple. They determine whether each state + is a conv or temporal state. block_tables: per-mamba-group persistent block-table tensors, in the same order as `mamba_group_ids`. Their `data_ptr()` / `stride(0)` are captured once for the kernel to index into. @@ -513,11 +559,18 @@ def initialize_from_forward_context( idx = 0 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names + group_copy_funcs = _get_copy_funcs_for_group( + kv_cache_config, mamba_group_id, mamba_state_copy_funcs + ) for layer_name in layer_names: attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache + assert len(kv_caches) == len(group_copy_funcs), ( + f"{layer_name} has {len(kv_caches)} state tensors but " + f"{len(group_copy_funcs)} copy functions" + ) - for state_type_idx, state in enumerate(kv_caches): + for state, copy_func in zip(kv_caches, group_copy_funcs): # Base address self.state_base_addrs[idx] = state.data_ptr() @@ -534,7 +587,6 @@ def initialize_from_forward_context( # Element size self.state_elem_sizes[idx] = state.element_size() - copy_func = mamba_state_copy_funcs[state_type_idx] assert ( copy_func is get_conv_copy_spec or copy_func is get_temporal_copy_spec @@ -566,6 +618,10 @@ def initialize_from_forward_context( self.state_group_indices[idx] = group_local_idx idx += 1 + assert idx == self.num_states, ( + f"initialized {idx} mamba states, expected {self.num_states}" + ) + # Cache per-group block-table base addresses and per-request stride. # `block_tables[i]` is the persistent 2D int32 block-table tensor for # `mamba_group_ids[i]`; `data_ptr()` / `stride(0)` are stable for the @@ -586,7 +642,7 @@ def initialize_from_forward_context( "postprocess_init", num_groups=self.num_groups, num_layers=self.num_layers, - num_state_types=self.num_state_types, + num_states=self.num_states, block_size=self.block_size, block_table_stride_req=self.block_table_stride_req, ) @@ -628,8 +684,7 @@ def run_fused_postprocess( spec_state_slot_selectors_gpu[:num_reqs] ) - total_states = self.num_layers * self.num_state_types - grid = (num_reqs, total_states) + grid = (num_reqs, self.num_states) has_ddtree = ddtree_accepted_node_indices is not None if ddtree_accepted_node_indices is None: ddtree_accepted_node_indices = num_accepted_tokens_gpu.new_empty((1, 1)) @@ -665,7 +720,6 @@ def warmup_fused_postprocess(self) -> bool: return False device = self.state_base_addrs.device - total_states = self.num_layers * self.num_state_types warmup_sizes = {1} if envs.VLLM_SM70_MTP_CONCURRENCY_WARMUP: max_num_reqs = int(self.num_accepted_tokens_out.shape[0]) @@ -689,7 +743,7 @@ def warmup_fused_postprocess(self) -> bool: (1, 1), dtype=torch.int32, device=device ) - postprocess_mamba_fused_kernel[(num_reqs, total_states)]( + postprocess_mamba_fused_kernel[(num_reqs, self.num_states)]( num_accepted_tokens, spec_state_slot_selectors, ddtree_accepted_node_indices, @@ -735,7 +789,7 @@ def create( cls, max_num_reqs: int, kv_cache_config: KVCacheConfig, - copy_funcs: tuple[MambaStateCopyFunc, ...], + copy_funcs: MambaStateCopyFuncsInput, make_buffer: Callable[..., CpuGpuBuffer], device: torch.device, with_postprocess_align: bool, @@ -748,7 +802,7 @@ def create( MambaSpecDecodeGPUContext.create( max_num_reqs=max_num_reqs, kv_cache_config=kv_cache_config, - num_state_types=len(copy_funcs), + mamba_state_copy_funcs=copy_funcs, device=device, make_buffer=make_buffer, ) @@ -761,7 +815,7 @@ def create( def collect_mamba_copy_meta( copy_bufs: MambaCopyBuffers, kv_cache_config: KVCacheConfig, - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, mamba_group_ids: list[int], src_block_idx: int, dest_block_idx: int, @@ -781,13 +835,20 @@ def collect_mamba_copy_meta( offset = copy_bufs.offset for mamba_group_id in mamba_group_ids: + group_copy_funcs = _get_copy_funcs_for_group( + kv_cache_config, mamba_group_id, mamba_state_copy_funcs + ) block_ids = req_state.block_ids[mamba_group_id] dest_block_id = block_ids[dest_block_idx] layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names for layer_name in layer_names: attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache - for state, state_copy_func in zip(kv_caches, mamba_state_copy_funcs): + assert len(kv_caches) == len(group_copy_funcs), ( + f"{layer_name} has {len(kv_caches)} state tensors but " + f"{len(group_copy_funcs)} copy functions" + ) + for state, state_copy_func in zip(kv_caches, group_copy_funcs): copy_src_block_idx = src_block_idx copy_num_accepted_tokens = accept_token_bias + 1 if state_slot_bias != accept_token_bias: @@ -844,7 +905,7 @@ def preprocess_mamba( input_batch: GPUInputBatch, requests: dict[str, CachedRequestState], forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, copy_bufs: MambaCopyBuffers, ): """ @@ -968,7 +1029,7 @@ def postprocess_mamba( requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, copy_bufs: MambaCopyBuffers, ddtree_accepted_node_indices: torch.Tensor | None = None, ): @@ -1089,7 +1150,7 @@ def postprocess_mamba_align_gpu( input_batch: GPUInputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], - mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + mamba_state_copy_funcs: MambaStateCopyFuncsInput, ddtree_accepted_node_indices: torch.Tensor | None = None, ) -> None: """GPU-side mamba postprocess for spec decode + hybrid + align mode. @@ -1123,7 +1184,7 @@ def postprocess_mamba_align_gpu( "postprocess_align_gpu", num_reqs=num_reqs, num_layers=ctx.num_layers, - num_state_types=ctx.num_state_types, + num_states=ctx.num_states, block_size=ctx.block_size, ) ctx.run_fused_postprocess( From b6ded5b8d4b701d9cb71fbfb2c864b8038e342e7 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:00:20 +0800 Subject: [PATCH 23/28] [Doc] Record Qwen4Exp MTP4 bring-up Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- docs/design/sm70_qwen38_flash_next_nvfp4.md | 75 +++++++++++++++++++-- 1 file changed, 68 insertions(+), 7 deletions(-) diff --git a/docs/design/sm70_qwen38_flash_next_nvfp4.md b/docs/design/sm70_qwen38_flash_next_nvfp4.md index 54191ae5c1..169425e151 100644 --- a/docs/design/sm70_qwen38_flash_next_nvfp4.md +++ b/docs/design/sm70_qwen38_flash_next_nvfp4.md @@ -2,9 +2,12 @@ ## Status and ownership -- Status: source bring-up implemented; CPU/configuration gates pass and the - local ModelScope snapshot is fully verified. Full-model load, output quality, - measured memory, and speed are not claimed yet. +- Status: source bring-up, native Qwen4Exp MTP, and prefix-cache configuration + are implemented; focused CPU/configuration gates pass and the local + ModelScope snapshot is fully verified. Native MTP4 now completes TP4 model + load, graph capture, warmup, and two 1024x256 requests. Acceptance, repeated + token equality, memory, and steady decode are measured below; matched no-MTP + token equality and verifier cost remain pending. - Integration line: `private/main`. - Base SHA: `d63e9490f65f9e01f6649053c1ab72922034b931`. - Model: `RadixArk/Qwen3.8-Flash-Next-NVFP4` at revision @@ -48,8 +51,18 @@ The first correctness route deliberately excludes speculative decoding. raw token IDs and builds the PLE context from committed tokens, so rejected speculative candidates cannot leak into the next trigram. V1 remains a correctness control, not the primary performance route. -- Prefix caching: disabled for the first route. The fixed QSA ring manager does - not yet implement reusable prefix blocks. +- Prefix caching: enabled for the MTP validation route. Hybrid recurrent state + uses `mamba_cache_mode=align` with chunked prefill. The fixed QSA compressor + ring is explicitly non-cacheable and is excluded from prefix-hit + reconciliation; a clean ring block is allocated after a hit. Main QSA KV, + compressed QSA KV, and aligned GDN/Mamba state remain cacheable. + +The native-MTP validation route keeps the same TP4/PP1, FP16 activation, +ModelOpt NVFP4, language-model-only, and FlashAttention-V100 contract. It uses +four speculative tokens, V2, CUDA graphs, prefix caching, a deterministic +greedy prompt, and two identical requests. The second request is both the hot +speed sample and the prefix-cache reuse check. A matched no-MTP run is still +required for exact token-ID quality comparison and incremental verifier cost. ## Architecture facts that affect the port @@ -107,6 +120,39 @@ graphs, workspaces, NCCL, allocator fragmentation, and loader transients. This explains why TP4 is plausible, but it is not evidence that the maximum context will load safely. +## Native MTP4 TP4 validation snapshot + +The first complete native-MTP run uses V2, TP4, FP16 activations, ModelOpt +NVFP4 weights, `FLASH_ATTN_V100` for target and draft attention, four draft +tokens, `mamba_cache_mode=align`, prefix caching, chunked prefill, and +FULL+PIECEWISE CUDA graphs. It runs two identical deterministic 1024-token +prompts with 256 forced output tokens each. The artifact is +`.artifacts/qwen4_exp_mtp_tp4_20260827/mtp4_prefix_graph_i1024_o256_r2_hetero_v2.json`. + +- Source HEAD is `d9a39ea434` on the Qwen3.8 worktree branch. The measurement + also sees the worktree's separately owned, uncommitted SM70 kernel changes; + it is bring-up evidence and must be repeated from a clean, pinned source + before becoming release-baseline evidence. +- Target plus MTP weights use 23.16 GiB/rank. The aligned MTP attention block + is 816 tokens with 1.62% recurrent-page padding. Available KV-cache memory + is 4.77 GiB/rank, or 219,942 tokens. Observed peak device memory is + 32,330 MiB on every 32,768-MiB V100; the run therefore has only 438 MiB of + peak device headroom at `gpu_memory_utilization=0.90`. +- PLE remains host-resident at 11.92 GiB/rank. The sampled minimum host + `MemAvailable` is 48.107 GiB and minimum free swap is 247.994 GiB. +- Both repeats emit 256 tokens and are exactly equal token-for-token. The first + request has 9.750-second TTFT and 53.125 steady decode tokens/s. The repeated + prefix has 2.667-second TTFT and 52.187 steady decode tokens/s. Mean steady + decode is 52.656 tokens/s; the first end-to-end request includes prefill and + JIT and is not a decode baseline. +- Across 256 speculative steps, the MTP head proposes 1,024 draft tokens and + 254 are accepted. Mean acceptance length including the target bonus is + 1.9921875. Draft acceptance is 24.8047%; per-position acceptance is + 54.6875%, 26.5625%, 13.28125%, and 4.6875%. +- A strictly matched no-MTP run is required before claiming target-token + equality or a verifier-cost ratio. Older 1024x256 artifacts use the distinct + Qwen3.8-27B checkpoint and are not valid controls for Flash Next. + ## Acceptance gates 1. Static route: Transformers config, model registry/processor registration, @@ -150,13 +196,28 @@ only after profiles identify them as measured decode bottlenecks. pinned-host PLE default, and the Qwen4Exp PLE/QSA compilation split operators. The same real configuration rejects the unvalidated multimodal route with an actionable `--language-model-only` error. +- Exact real-checkpoint construction with native MTP4 and prefix caching on + resolves the target as `Qwen4ExpForConditionalGeneration`, the draft as + `Qwen4ExpMTP`, target and draft attention as `FLASH_ATTN_V100`, and recurrent + caching as `align` with 16-token configured blocks and chunked prefill. + Focused prefix-cache/QSA coordinator and model-config tests pass. The + non-cacheable QSA compressor ring is skipped during prefix-hit matching, + matching the current upstream Qwen4Exp contract. +- TP4 loaded all 206 checkpoint shards in 109.27 seconds and then correctly + rejected an incomplete runtime extension set: `_C` and + `_C_stable_libtorch` were present, but `_moe_C` was omitted. A complete + runtime must carry all three. The matching `_moe_C` artifact has SHA-256 + `a14eeb4fa06947e335cf69ee188e23509fc294da61cada786df09888ca5b4469`; + its graph-safe permute, workspace-size, unpermute schemas, and SM70 support + probe all pass before the next full-model attempt. - Full 48-layer meta construction from the real checkpoint config succeeds in language-model-only mode. It instantiates QSA, GDN, HC, PLE, and all 512 experts without materializing weights; the routed experts select `ModelOptNvFp4SM70MoEMethod(use_a16=True)` and the PLE table has shape `(320001536, 160)` with FP8 E4M3 storage. This constructor probe used TP1; - TP4 selection and expert geometry are covered separately and full TP4 load - remains pending. + TP4 selection and expert geometry are covered separately. Full TP4 target + plus native-MTP loading now completes with 23.16 GiB/rank of loaded model + state before cache and graph allocation. - Focused CPU tests cover PLE shard loading and hashing, `seed=None`, permanent host residency during post-load processing, QSA cache grouping, V1 and V2 n-gram inputs, V2 circular block-table sizing, scheduler-manager conversion, From 70b63a1e4502627285cb436aa71ca459da5b14c4 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:32:06 +0800 Subject: [PATCH 24/28] [Spec Decode] Reduce Qwen4Exp MTP cost Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/engine/test_arg_utils.py | 4 +- tests/v1/worker/test_qwen4_exp_v2.py | 57 ++++++++++ vllm/config/speculative.py | 24 ++++ vllm/engine/arg_utils.py | 7 +- vllm/models/qwen4_exp/nvidia/mtp.py | 5 + .../gpu/spec_decode/eagle/speculator.py | 107 ++++++++++++++---- 6 files changed, 176 insertions(+), 28 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 541f3d5d42..4cf07e5599 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -293,7 +293,7 @@ def test_sm70_mtp_defaults_require_env_opt_in(monkeypatch): "method": "mtp", "num_speculative_tokens": 4, "use_local_argmax_reduction": True, - "draft_sample_method": "probabilistic", + "draft_sample_method": "greedy", "attention_backend": "FLASH_ATTN_V100", } assert args.enable_prefix_caching is True @@ -379,7 +379,7 @@ def test_sm70_explicit_mtp_still_gets_safe_defaults(monkeypatch): assert args.speculative_config == { "method": "mtp", "num_speculative_tokens": 2, - "draft_sample_method": "probabilistic", + "draft_sample_method": "greedy", "use_local_argmax_reduction": True, "attention_backend": "FLASH_ATTN_V100", } diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 29ab4fe9bb..1e6d9d5af2 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -55,6 +55,63 @@ def test_qwen4_exp_mtp_v2_unpacks_logits_and_feedback_hidden_states( assert actual_feedback_hidden is feedback_hidden +def test_qwen4_exp_mtp_v2_uses_local_argmax_without_full_logits() -> None: + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.use_local_argmax_reduction = True + + class DraftModel: + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.shape == (2, 4) + return torch.tensor([7, 11]) + + def compute_logits(self, _hidden_states: torch.Tensor) -> torch.Tensor: + raise AssertionError("local argmax must not materialize full logits") + + speculator.model = DraftModel() + top_tokens = speculator._sample_draft( + hidden_states=torch.zeros(2, 4), + idx_mapping=torch.arange(2, dtype=torch.int32), + pos=torch.arange(2), + draft_step=torch.tensor(0), + draft_logits=None, + ) + + torch.testing.assert_close(top_tokens, torch.tensor([7, 11])) + + +def test_qwen4_exp_mtp_v2_reuses_step_zero_qsa_indices() -> None: + calls: list[tuple[str, object]] = [] + + class MTPBackbone: + def set_skip_topk(self, skip: bool) -> None: + calls.append(("skip", skip)) + + def compact_topk_indices(self, row_indices: torch.Tensor) -> None: + calls.append(("compact", row_indices.tolist())) + + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.share_mtp_topk_indices = True + speculator.num_speculative_steps = 3 + speculator.last_token_indices = torch.tensor([4, 9, 12]) + speculator.model = SimpleNamespace(model=MTPBackbone()) + + speculator._mtp_prefill_begin() + speculator._mtp_prefill_end(num_reqs=2) + speculator._mtp_decode_begin() + speculator._mtp_decode_end() + + assert calls == [ + ("skip", False), + ("compact", [4, 9]), + ("skip", True), + ("skip", False), + ] + + def test_qsa_circular_group_uses_one_block_and_custom_slot_mapping( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 07507b9be4..9d18c98185 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -139,6 +139,9 @@ class SpeculativeConfig: """The specific revision to use for the draft model code on Hugging Face Hub. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version.""" + index_share_for_mtp_iteration: bool | None = None + """Override whether MTP iterations reuse the first step's sparse indices. + If ``None``, use the value from the draft model's Hugging Face config.""" # Advanced control disable_padded_drafter_batch: bool = False @@ -368,6 +371,15 @@ def compute_hash(self) -> str: ) ) + if self.method == "mtp" and self.draft_model_config is not None: + factors.append( + getattr( + self.draft_model_config.hf_config, + "index_share_for_mtp_iteration", + False, + ) + ) + hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() return hash_str @@ -531,6 +543,9 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: "hc_mult": int(text_config.hc_count), "n_predict": n_predict, "architectures": ["Qwen4ExpMTP"], + "index_share_for_mtp_iteration": getattr( + text_config, "index_share_for_mtp_iteration", True + ), } ) @@ -986,6 +1001,15 @@ def __post_init__(self): ), ) ) + + if self.index_share_for_mtp_iteration is not None: + if self.method != "mtp" or self.draft_model_config is None: + raise ValueError( + "index_share_for_mtp_iteration is only supported with method='mtp'" + ) + self.draft_model_config.hf_config.index_share_for_mtp_iteration = ( + self.index_share_for_mtp_iteration + ) return self def _verify_dspark_final_stage_ownership(self) -> None: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 31850497d3..4dc0954a72 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -1819,10 +1819,9 @@ def _maybe_apply_sm70_mtp_defaults( return if "draft_sample_method" not in self.speculative_config: - # Preserve the validated probabilistic default for MTP and the - # DFlash family. DFlash2's greedy performance lane remains - # available when explicitly requested, without changing DFlash1. - draft_sample_method = "probabilistic" + # MTP can use the vocab-parallel local-argmax fast path only in + # greedy mode. DFlash keeps probabilistic drafting by default. + draft_sample_method = "greedy" if spec_method == "mtp" else "probabilistic" self.speculative_config["draft_sample_method"] = draft_sample_method profile_updates.append( f"speculative_config.draft_sample_method={draft_sample_method}" diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 87599d99f8..3945ebc663 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -442,6 +442,11 @@ def compute_logits( ) -> torch.Tensor | None: return self.logits_processor(self.lm_head, hidden_states) + def get_top_tokens( + self, hidden_states: torch.Tensor, spec_step_idx: int = 0 + ) -> torch.Tensor: + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def remap_weight_names(): for name, weight in weights: diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 3c4609d357..07afbd9c04 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -65,6 +65,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vocab_size = self.draft_model_config.get_vocab_size() self.dtype = vllm_config.model_config.dtype self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel + self.use_local_argmax_reduction = ( + self.speculative_config.use_local_argmax_reduction + ) + self.share_mtp_topk_indices = False # DP configuration self.dp_size = vllm_config.parallel_config.data_parallel_size @@ -157,6 +161,17 @@ def load_model(self, target_model: nn.Module) -> None: ).keys() self.model = load_eagle_model(target_model, self.vllm_config) + self._validate_local_argmax_reduction() + + draft_hf_config = self.draft_model_config.hf_config + self.share_mtp_topk_indices = ( + self.method == "mtp" + and getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + and hasattr(self.model.model, "set_skip_topk") + and hasattr(self.model.model, "compact_topk_indices") + ) + if self.share_mtp_topk_indices: + logger.info("Reusing target-aligned step-0 QSA indices for MTP steps 1+.") all_attn_layers = get_layers_from_vllm_config( self.vllm_config, @@ -231,15 +246,35 @@ def run_model( hidden_states = ret_hidden_states return last_hidden_states, hidden_states + def _validate_local_argmax_reduction(self) -> None: + if not self.use_local_argmax_reduction: + return + if self.speculative_config.draft_sample_method == "probabilistic": + raise ValueError( + "use_local_argmax_reduction is not compatible with " + "draft_sample_method='probabilistic'." + ) + if not hasattr(self.model, "get_top_tokens"): + raise ValueError( + "use_local_argmax_reduction is enabled but draft model " + f"{self.model.__class__.__name__} does not implement " + "get_top_tokens()." + ) + logger.info( + "Using local argmax reduction for draft token generation " + "(communication: O(2*tp_size) vs O(vocab_size))." + ) + def _sample_draft( self, - logits: torch.Tensor, + hidden_states: torch.Tensor, idx_mapping: torch.Tensor, pos: torch.Tensor, draft_step: torch.Tensor, draft_logits: torch.Tensor | None, ) -> torch.Tensor: if draft_logits is not None: + logits = self.model.compute_logits(hidden_states) # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise # used for draft and target sampling. return gumbel_sample( @@ -253,8 +288,26 @@ def _sample_draft( output_processed_logits_col=draft_step, use_fp64=self.use_fp64_gumbel, ) - else: - return logits.argmax(dim=-1) + if self.use_local_argmax_reduction: + return self.model.get_top_tokens(hidden_states) + logits = self.model.compute_logits(hidden_states) + return logits.argmax(dim=-1) + + def _mtp_prefill_begin(self) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) + + def _mtp_prefill_end(self, num_reqs: int) -> None: + if self.share_mtp_topk_indices and self.num_speculative_steps > 1: + self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs]) + + def _mtp_decode_begin(self) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(True) + + def _mtp_decode_end(self) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) def prefill( self, @@ -279,10 +332,9 @@ def prefill( mm_inputs=mm_inputs, ) sample_hidden_states = last_hidden_states[last_token_indices] - logits = self.model.compute_logits(sample_hidden_states) self.draft_tokens[:num_reqs, 0] = self._sample_draft( - logits, + sample_hidden_states, idx_mapping, pos, self.current_draft_step, @@ -364,9 +416,8 @@ def generate_draft( last_hidden_states = last_hidden_states[:num_reqs] # Sample the draft tokens. - logits = self.model.compute_logits(last_hidden_states) draft_tokens = self._sample_draft( - logits, + last_hidden_states, idx_mapping, positions, self.current_draft_step, @@ -434,11 +485,13 @@ def capture( # For PIECEWISE, only the model's compiled regions are captured # and the rest (compute_logits, gumbel_sample) runs eagerly. assert self.prefill_cudagraph_manager is not None + self._mtp_prefill_begin() self.prefill_cudagraph_manager.capture( self.prefill, attn_states, progress_bar_desc="Capturing eagle prefill CUDA graphs", ) + self._mtp_prefill_end(self.max_num_reqs) if self.num_speculative_steps == 1: return @@ -447,15 +500,19 @@ def capture( # compute_logits + sample + update_eagle_inputs) for a single # step. assert self.decode_cudagraph_manager is not None - self.decode_cudagraph_manager.capture( - self.generate_draft, - self.model_state, - self.input_buffers, - self.block_tables, - self.attn_groups, - self.kv_cache_config, - progress_bar_desc="Capturing eagle decode CUDA graphs", - ) + self._mtp_decode_begin() + try: + self.decode_cudagraph_manager.capture( + self.generate_draft, + self.model_state, + self.input_buffers, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + progress_bar_desc="Capturing eagle decode CUDA graphs", + ) + finally: + self._mtp_decode_end() @torch.inference_mode() def propose( @@ -550,6 +607,7 @@ def propose( need_eager=is_profile, ) + self._mtp_prefill_begin() if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. assert self.prefill_cudagraph_manager is not None @@ -567,6 +625,7 @@ def propose( cudagraph_runtime_mode=prefill_batch_desc.cg_mode, mm_inputs=mm_inputs, ) + self._mtp_prefill_end(num_reqs) if self.num_speculative_steps == 1: # Early exit. @@ -595,12 +654,16 @@ def propose( ) # Generate the remaining num_speculative_steps - 1 draft tokens. - self.multi_step_decode( - num_reqs, - dummy_run and skip_attn_for_dummy_run, - decode_batch_desc, - num_tokens_across_dp, - ) + self._mtp_decode_begin() + try: + self.multi_step_decode( + num_reqs, + dummy_run and skip_attn_for_dummy_run, + decode_batch_desc, + num_tokens_across_dp, + ) + finally: + self._mtp_decode_end() return self.draft_tokens[:num_reqs] From 5abeaa617754090ae20dc225134e310f729f0d2f Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:16:54 +0800 Subject: [PATCH 25/28] [Bench] Record per-prompt speculative acceptance Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- benchmarks/benchmark_sm70_model_tokens.py | 80 ++++++++++++++++++- .../benchmarks/test_benchmark_sm70_decode.py | 32 +++++++- 2 files changed, 110 insertions(+), 2 deletions(-) diff --git a/benchmarks/benchmark_sm70_model_tokens.py b/benchmarks/benchmark_sm70_model_tokens.py index 2f0d580bcf..75112d51fd 100644 --- a/benchmarks/benchmark_sm70_model_tokens.py +++ b/benchmarks/benchmark_sm70_model_tokens.py @@ -240,6 +240,66 @@ def _spec_decoding_summary( } +def _spec_decoding_delta( + before: list[dict[str, Any]], + after: list[dict[str, Any]], +) -> dict[str, Any] | None: + """Return speculative-decoding counters added between two snapshots.""" + + before_summary = _spec_decoding_summary(before) + after_summary = _spec_decoding_summary(after) + if after_summary is None: + return None + + before_drafts = int(before_summary["num_drafts"]) if before_summary else 0 + before_draft_tokens = ( + int(before_summary["num_draft_tokens"]) if before_summary else 0 + ) + before_accepted = ( + int(before_summary["num_accepted_tokens"]) if before_summary else 0 + ) + after_drafts = int(after_summary["num_drafts"]) + after_draft_tokens = int(after_summary["num_draft_tokens"]) + after_accepted = int(after_summary["num_accepted_tokens"]) + + num_drafts = after_drafts - before_drafts + num_draft_tokens = after_draft_tokens - before_draft_tokens + num_accepted_tokens = after_accepted - before_accepted + if num_drafts <= 0: + return None + if num_draft_tokens < 0 or num_accepted_tokens < 0: + raise RuntimeError("Speculative-decoding counters moved backwards") + + before_per_pos = ( + list(before_summary["accepted_tokens_per_pos"]) if before_summary else [] + ) + after_per_pos = list(after_summary["accepted_tokens_per_pos"]) + width = max(len(before_per_pos), len(after_per_pos)) + before_per_pos.extend([0] * (width - len(before_per_pos))) + after_per_pos.extend([0] * (width - len(after_per_pos))) + accepted_tokens_per_pos = [ + end - start for start, end in zip(before_per_pos, after_per_pos, strict=True) + ] + if any(value < 0 for value in accepted_tokens_per_pos): + raise RuntimeError("Per-position speculative counters moved backwards") + + avg_accepted_tokens = num_accepted_tokens / num_drafts + return { + "num_drafts": num_drafts, + "num_draft_tokens": num_draft_tokens, + "num_accepted_tokens": num_accepted_tokens, + "avg_accepted_tokens_no_bonus": avg_accepted_tokens, + "mean_acceptance_length": 1 + avg_accepted_tokens, + "draft_acceptance_rate": ( + num_accepted_tokens / num_draft_tokens if num_draft_tokens else None + ), + "accepted_tokens_per_pos": accepted_tokens_per_pos, + "per_position_acceptance_rate": [ + value / num_drafts for value in accepted_tokens_per_pos + ], + } + + def _load_prompts(args: argparse.Namespace) -> list[Any]: if args.input_lens is not None: if args.input_len is not None or args.prompt or args.prompts_json is not None: @@ -1713,6 +1773,7 @@ def _dump(args: argparse.Namespace) -> int: torch.accelerator.synchronize() torch.cuda.cudart().cudaProfilerStart() generate_seconds_by_repeat: list[float] = [] + sequential_prompt_metrics: list[dict[str, Any]] = [] outputs = [] try: for repeat_index in range(args.repeat_count): @@ -1724,8 +1785,24 @@ def _dump(args: argparse.Namespace) -> int: raise RuntimeError("Failed to reset the idle prefix cache") generate_start = time.perf_counter() if args.sequential_prompts: - for prompt in prompts: + metrics_before_prompt = _metric_snapshot(llm) + for prompt_index, prompt in enumerate(prompts): + prompt_start = time.perf_counter() outputs.extend(llm.generate([prompt], sampling_params)) + prompt_seconds = time.perf_counter() - prompt_start + metrics_after_prompt = _metric_snapshot(llm) + sequential_prompt_metrics.append( + { + "repeat_index": repeat_index, + "prompt_index": prompt_index, + "generate_seconds": prompt_seconds, + "spec_decoding_metrics": _spec_decoding_delta( + metrics_before_prompt, + metrics_after_prompt, + ), + } + ) + metrics_before_prompt = metrics_after_prompt else: outputs.extend(llm.generate(prompts, sampling_params)) generate_seconds_by_repeat.append(time.perf_counter() - generate_start) @@ -1831,6 +1908,7 @@ def _dump(args: argparse.Namespace) -> int: "eos_token_ids": eos_token_ids, "metrics_snapshot": metrics_snapshot, "spec_decoding_metrics": _spec_decoding_summary(metrics_snapshot), + "sequential_prompt_metrics": sequential_prompt_metrics, "sampling_params": { "max_tokens": args.max_tokens, "temperature": args.temperature, diff --git a/tests/benchmarks/test_benchmark_sm70_decode.py b/tests/benchmarks/test_benchmark_sm70_decode.py index b9fe6ebb78..0895a8ab17 100644 --- a/tests/benchmarks/test_benchmark_sm70_decode.py +++ b/tests/benchmarks/test_benchmark_sm70_decode.py @@ -3,7 +3,7 @@ import types -from benchmarks import benchmark_sm70_decode +from benchmarks import benchmark_sm70_decode, benchmark_sm70_model_tokens def test_sm70_fa2_d256_prefill_status_reports_import_error(monkeypatch): @@ -46,3 +46,33 @@ def test_sm70_fa2_d256_prefill_status_requires_dense_and_paged_ops(monkeypatch): assert status["error"] is None assert all(status["required_ops"].values()) assert all(status["optional_ops"].values()) + + +def test_spec_decoding_delta_reports_one_sequential_prompt(): + before = [ + {"name": "vllm:spec_decode_num_drafts", "value": 10}, + {"name": "vllm:spec_decode_num_draft_tokens", "value": 30}, + {"name": "vllm:spec_decode_num_accepted_tokens", "value": 20}, + { + "name": "vllm:spec_decode_num_accepted_tokens_per_pos", + "values": [8, 7, 5], + }, + ] + after = [ + {"name": "vllm:spec_decode_num_drafts", "value": 14}, + {"name": "vllm:spec_decode_num_draft_tokens", "value": 42}, + {"name": "vllm:spec_decode_num_accepted_tokens", "value": 27}, + { + "name": "vllm:spec_decode_num_accepted_tokens_per_pos", + "values": [11, 9, 7], + }, + ] + + metrics = benchmark_sm70_model_tokens._spec_decoding_delta(before, after) + + assert metrics is not None + assert metrics["num_drafts"] == 4 + assert metrics["num_draft_tokens"] == 12 + assert metrics["num_accepted_tokens"] == 7 + assert metrics["accepted_tokens_per_pos"] == [3, 2, 2] + assert metrics["mean_acceptance_length"] == 2.75 From cf5c44a1aa65bd8556ff9050bf9cb19d1c5f439d Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:48:32 +0800 Subject: [PATCH 26/28] [Bugfix] Load fused Qwen4Exp MTP experts Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_weight_loading.py | 71 +++++++++++++++++++ vllm/model_executor/layers/fused_moe/layer.py | 47 +++++++++++- vllm/models/qwen4_exp/nvidia/model.py | 1 + 3 files changed, 118 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py index fb5d9903c7..09a975f196 100644 --- a/tests/models/qwen4_exp/test_weight_loading.py +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from torch import nn +from vllm.model_executor.layers.fused_moe.layer import FusedMoE from vllm.model_executor.models.qwen3_next import Qwen3NextSparseMoeBlock from vllm.models.qwen4_exp.nvidia.model import ( Qwen4ExpForConditionalGeneration, @@ -144,6 +146,9 @@ def fake_qwen3_next_init(self, vllm_config, prefix="") -> None: block = Qwen4ExpSparseMoeBlock(vllm_config, prefix="model.layers.0.mlp") assert block.experts.expert_mapping == [ + ("experts.w13_weight", "experts.gate_up_proj", 0, "w1"), + ("experts.w13_weight", "experts.gate_up_proj", 1, "w3"), + ("experts.w2_weight", "experts.down_proj", 0, "w2"), ("experts.w13_", "experts.0.gate_proj.", 0, "w1"), ("experts.w2_", "experts.0.down_proj.", 0, "w2"), ("experts.w13_", "experts.0.up_proj.", 0, "w3"), @@ -153,6 +158,72 @@ def fake_qwen3_next_init(self, vllm_config, prefix="") -> None: ] +def test_fused_mtp_expert_checkpoint_loads_every_expert() -> None: + class FakeFusedExperts(nn.Module): + load_weights = FusedMoE.load_weights + + def __init__(self) -> None: + super().__init__() + self.layer_name = "model.layers.0.mlp.experts" + self.w13_weight = nn.Parameter(torch.empty(1)) + self.w2_weight = nn.Parameter(torch.empty(1)) + self.calls: list[tuple[str, str, int, torch.Tensor]] = [] + self.expert_mapping = FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=2, + include_fused=True, + ) + + def weight_loader( + self, + *, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + return_success: bool, + ) -> bool: + assert return_success + param_name = "w13" if param is self.w13_weight else "w2" + self.calls.append((param_name, shard_id, expert_id, loaded_weight.clone())) + return True + + experts = FakeFusedExperts() + gate_up = torch.arange(2 * 6 * 2).reshape(2, 6, 2) + down = torch.arange(2 * 2 * 3).reshape(2, 2, 3) + + loaded = list( + experts.load_weights([("gate_up_proj", gate_up), ("down_proj", down)]) + ) + + assert loaded == [ + "w13_weight", + "w13_weight", + "w13_weight", + "w13_weight", + "w2_weight", + "w2_weight", + ] + assert [(name, shard, expert) for name, shard, expert, _ in experts.calls] == [ + ("w13", "w1", 0), + ("w13", "w1", 1), + ("w13", "w3", 0), + ("w13", "w3", 1), + ("w2", "w2", 0), + ("w2", "w2", 1), + ] + torch.testing.assert_close(experts.calls[0][3], gate_up[0, :3]) + torch.testing.assert_close(experts.calls[1][3], gate_up[1, :3]) + torch.testing.assert_close(experts.calls[2][3], gate_up[0, 3:]) + torch.testing.assert_close(experts.calls[3][3], gate_up[1, 3:]) + torch.testing.assert_close(experts.calls[4][3], down[0]) + torch.testing.assert_close(experts.calls[5][3], down[1]) + + @pytest.mark.parametrize( ("checkpoint_name", "model_name"), [ diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 4ff43ce21b..beb1f6be7c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1338,6 +1338,7 @@ def make_expert_params_mapping( ckpt_up_proj_name: str, num_experts: int, num_redundant_experts: int = 0, + include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: num_physical_experts = num_experts + num_redundant_experts @@ -1357,7 +1358,48 @@ def make_expert_params_mapping( else "" ) - return [ + fused_mapping: list[tuple[str, str, int, str]] = [] + if include_fused: + gate_up_name = None + if ckpt_gate_proj_name == "gate_proj" and ckpt_up_proj_name == "up_proj": + gate_up_name = "gate_up_proj" + elif ckpt_gate_proj_name == "w1" and ckpt_up_proj_name == "w3": + gate_up_name = "w13" + else: + logger.warning( + "Unexpected gate/up projection names: %s, %s. " + "Fused gate/up mapping will be skipped.", + ckpt_gate_proj_name, + ckpt_up_proj_name, + ) + + if gate_up_name is not None: + # Some checkpoints store every expert in one 3D tensor and + # concatenate gate/up along dim 1. FusedMoE.load_weights + # already knows how to split and TP-shard that layout; these + # aliases make the checkpoint names reachable. + fused_mapping = [ + ( + f"experts.{base_layer}w13_weight", + f"experts.{gate_up_name}", + 0, + "w1", + ), + ( + f"experts.{base_layer}w13_weight", + f"experts.{gate_up_name}", + 1, + "w3", + ), + ( + f"experts.{base_layer}w2_weight", + f"experts.{ckpt_down_proj_name}", + 0, + "w2", + ), + ] + + per_expert_mapping = [ # (param_name, weight_name, expert_id, shard_id) ( f"experts.{base_layer}w13_" @@ -1374,6 +1416,7 @@ def make_expert_params_mapping( ("w3", ckpt_up_proj_name), ] ] + return fused_mapping + per_expert_mapping @property def hidden_size(self) -> int: @@ -1404,6 +1447,7 @@ def fused_moe_make_expert_params_mapping( ckpt_up_proj_name: str, num_experts: int, num_redundant_experts: int = 0, + include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: return FusedMoE.make_expert_params_mapping( model, @@ -1412,6 +1456,7 @@ def fused_moe_make_expert_params_mapping( ckpt_up_proj_name, num_experts, num_redundant_experts, + include_fused, ) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 7eddc1bb35..06ce82e8cf 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -208,6 +208,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: ckpt_up_proj_name="up_proj", num_experts=self.n_routed_experts, num_redundant_experts=self.n_redundant_experts, + include_fused=True, ) From ceb543c05573c9bcbf1cc9563bab0c2e8c746ae3 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Thu, 27 Aug 2026 22:56:37 +0800 Subject: [PATCH 27/28] [Bugfix] Reject incomplete Qwen4Exp MTP loads Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_weight_loading.py | 19 ++++++++++++++ vllm/models/qwen4_exp/nvidia/mtp.py | 26 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/tests/models/qwen4_exp/test_weight_loading.py b/tests/models/qwen4_exp/test_weight_loading.py index 09a975f196..ce95ac3441 100644 --- a/tests/models/qwen4_exp/test_weight_loading.py +++ b/tests/models/qwen4_exp/test_weight_loading.py @@ -13,6 +13,7 @@ Qwen4ExpSparseMoeBlock, _remap_qsa_cache_scale_name, ) +from vllm.models.qwen4_exp.nvidia.mtp import _validate_mtp_expert_weights_loaded @pytest.mark.parametrize( @@ -224,6 +225,24 @@ def weight_loader( torch.testing.assert_close(experts.calls[5][3], down[1]) +def test_mtp_expert_loading_fails_closed() -> None: + model = nn.Module() + model.model = nn.Module() + model.model.layers = nn.ModuleList([nn.Module()]) + model.model.layers[0].mlp = nn.Module() + model.model.layers[0].mlp.experts = nn.Module() + experts = model.model.layers[0].mlp.experts + experts.w13_weight = nn.Parameter(torch.empty(1)) + experts.w2_weight = nn.Parameter(torch.empty(1)) + + w13_name = "model.layers.0.mlp.experts.w13_weight" + w2_name = "model.layers.0.mlp.experts.w2_weight" + _validate_mtp_expert_weights_loaded(model, {w13_name, w2_name}) + + with pytest.raises(ValueError, match="w2_weight"): + _validate_mtp_expert_weights_loaded(model, {w13_name}) + + @pytest.mark.parametrize( ("checkpoint_name", "model_name"), [ diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 3945ebc663..1b1df75d37 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -116,6 +116,28 @@ def _remap_mtp_weight_name(name: str) -> str | None: return None +def _validate_mtp_expert_weights_loaded( + model: nn.Module, + loaded_weights: set[str], +) -> None: + """Reject a silently incomplete Qwen4Exp MTP routed-expert load.""" + + required_suffixes = ( + ".mlp.experts.w13_weight", + ".mlp.experts.w2_weight", + ) + required = { + name for name, _ in model.named_parameters() if name.endswith(required_suffixes) + } + missing = required - loaded_weights + if missing: + missing_names = ", ".join(sorted(missing)) + raise ValueError( + "Qwen4Exp MTP routed-expert checkpoint weights were not loaded: " + f"{missing_names}. Check fused/per-expert checkpoint mappings." + ) + + def _make_draft_vllm_config( vllm_config: VllmConfig, mtp_start_layer_idx: int, @@ -459,7 +481,9 @@ def remap_weight_names(): skip_substrs=["hyper_connection_mixer.block_inject_weight"], ignore_unexpected_suffixes=_QWEN4_EXP_IGNORED_MISSING_SUFFIXES.copy(), ) - return loader.load_weights(remap_weight_names()) + loaded_weights = loader.load_weights(remap_weight_names()) + _validate_mtp_expert_weights_loaded(self, loaded_weights) + return loaded_weights __all__ = ["Qwen4ExpMTP", "Qwen4ExpMultiTokenPredictor"] From 95fdf660b325ef1627a01b414518d1c45f7e67d0 Mon Sep 17 00:00:00 2001 From: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:38:05 +0800 Subject: [PATCH 28/28] [Bugfix][SM70] Repair Qwen4Exp integration gates Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com> --- tests/models/qwen4_exp/test_ple.py | 2 +- tests/models/qwen4_exp/test_qsa_cache.py | 5 +++-- .../test_sm70_modelopt_mixed_nvfp4.py | 4 ++-- tests/v1/worker/test_model_state_init.py | 2 +- tests/v1/worker/test_qwen4_exp_ngram.py | 3 ++- tests/v1/worker/test_qwen4_exp_v2.py | 2 +- vllm/model_executor/layers/linear.py | 12 ++++++++---- vllm/models/qwen4_exp/nvidia/model.py | 13 +++++++------ vllm/models/qwen4_exp/nvidia/mtp.py | 4 ++-- vllm/models/qwen4_exp/nvidia/ple_layer.py | 19 +++++++++++++------ vllm/models/qwen4_exp/nvidia/qsa.py | 7 +++++-- vllm/v1/kv_cache_interface.py | 2 +- .../gpu/spec_decode/eagle/speculator.py | 7 ++++--- vllm/v1/worker/gpu_model_runner.py | 4 +--- 14 files changed, 51 insertions(+), 35 deletions(-) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 102d1045da..6777828c76 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -101,7 +101,7 @@ def test_pinned_host_ple_fp8_rows_are_gatherable_on_sm70( layer.prepare_accelerator_weight() output = layer(torch.tensor([0, 7], dtype=torch.int64, device="cuda")) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert output.dtype == torch.float16 expected = raw.view(torch.float8_e4m3fn).float() * 0.25 diff --git a/tests/models/qwen4_exp/test_qsa_cache.py b/tests/models/qwen4_exp/test_qsa_cache.py index 9a5ed1daca..d3cfea8d5f 100644 --- a/tests/models/qwen4_exp/test_qsa_cache.py +++ b/tests/models/qwen4_exp/test_qsa_cache.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any import torch @@ -12,7 +13,7 @@ def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: prefix = "model.layers.0.self_attn.raw_key_cache" - static_forward_context = {} + static_forward_context: dict[str, Any] = {} layer = QSAKeyStateCache( head_size=128, dtype=torch.float16, @@ -27,7 +28,7 @@ def test_bind_qsa_key_cache_builds_key_and_mrope_views() -> None: ), ) cache = torch.empty(2, 8, 1, layer.head_size, dtype=torch.float16) - runner_kv_caches = [] + runner_kv_caches: list[torch.Tensor] = [] bind_kv_cache({prefix: cache}, static_forward_context, runner_kv_caches) diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index 469e8aadb9..da2cd7dc1f 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -222,7 +222,7 @@ def test_nvfp4_single_token_direct_routing_is_exact_and_graph_dynamic(): torch.ops._moe_C.moe_unpermute( expert_output, weights, identity, None, top_k, reference ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert torch.equal(expanded, x.expand(top_k, -1)) assert torch.equal(active_ids, topk_ids[0].int()) assert torch.equal(output, reference) @@ -240,7 +240,7 @@ def test_nvfp4_single_token_direct_routing_is_exact_and_graph_dynamic(): torch.ops._moe_C.moe_unpermute( expert_output, weights, identity, None, top_k, reference ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert torch.equal(expanded, x.expand(top_k, -1)) assert torch.equal(active_ids, topk_ids[0].int()) assert torch.equal(output, reference) diff --git a/tests/v1/worker/test_model_state_init.py b/tests/v1/worker/test_model_state_init.py index 7535a70b6c..44200b5806 100644 --- a/tests/v1/worker/test_model_state_init.py +++ b/tests/v1/worker/test_model_state_init.py @@ -10,7 +10,7 @@ def test_model_can_select_custom_model_state() -> None: - captured = {} + captured: dict[str, object] = {} class CustomModelState: def __init__(self, vllm_config, model, encoder_cache, device) -> None: diff --git a/tests/v1/worker/test_qwen4_exp_ngram.py b/tests/v1/worker/test_qwen4_exp_ngram.py index c3904e1b65..6c2459c897 100644 --- a/tests/v1/worker/test_qwen4_exp_ngram.py +++ b/tests/v1/worker/test_qwen4_exp_ngram.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace +from typing import Any import numpy as np import torch @@ -58,7 +59,7 @@ def test_v1_runner_prepares_committed_ngram_context() -> None: def test_v1_runner_uses_stable_dummy_ngram_buffers() -> None: runner = _runner() - model_kwargs = {} + model_kwargs: dict[str, Any] = {} runner._maybe_add_ngram_kwargs( model_kwargs, diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 1e6d9d5af2..954779157f 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -219,7 +219,7 @@ def test_qsa_circular_group_emits_no_generic_slots_on_sm70() -> None: positions=torch.tensor([153797, 165757], dtype=torch.int64, device=device), num_tokens_padded=2, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert slot_mappings[0].tolist() == [-1, -1] assert slot_mappings[1].tolist() == [ diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 1d1f726b45..7f549fed61 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -1224,11 +1224,13 @@ def load_weights( self.validate_shard_id(shard_id) if "." in name: submodule, _, attr = name.rpartition(".") - param = getattr(self.get_submodule(submodule), attr, self) + param = getattr(self.get_submodule(submodule), attr, None) else: - param = getattr(self, name, self) + param = getattr(self, name, None) if param is None and name == "bias": continue + if param is None: + raise ValueError(f"Unknown parameter {name!r} for {self.prefix}") param.weight_loader(param, loaded_weight, shard_id) logger.debug( "Loaded shard %s with shape %s into %s.%s", @@ -1346,11 +1348,13 @@ def load_weights( self.validate_shard_id(shard_id) if "." in name: submodule, _, attr = name.rpartition(".") - param = getattr(self.get_submodule(submodule), attr, self) + param = getattr(self.get_submodule(submodule), attr, None) else: - param = getattr(self, name, self) + param = getattr(self, name, None) if param is None and name == "bias": continue + if param is None: + raise ValueError(f"Unknown parameter {name!r} for {self.prefix}") param.weight_loader(param, loaded_weight, shard_id) logger.debug( "Loaded shard %s with shape %s into %s.%s", diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 06ce82e8cf..526acbe0a4 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/vllm/models/qwen4_exp/nvidia/model.py @@ -375,6 +375,7 @@ def forward( attn_out = self.self_attn( hidden_states=block_input, positions=positions, + output=None, ) else: raise ValueError("Invalid layer_type") @@ -691,7 +692,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) self.set_moe_parameters(self.model.layers) @@ -769,7 +770,7 @@ def get_gdn_mamba_state_dtype_from_config( @classmethod def get_gdn_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: parallel_config = vllm_config.parallel_config hf_config = vllm_config.model_config.hf_text_config tp_size = parallel_config.tensor_parallel_size @@ -798,7 +799,7 @@ def get_mamba_state_dtype_from_config( @classmethod def get_mamba_state_shape_from_config( cls, vllm_config: VllmConfig - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: return cls.get_gdn_mamba_state_shape_from_config(vllm_config) @classmethod @@ -967,7 +968,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model") -> None: prefix=maybe_prefix(prefix, "language_model"), ) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.language_model.make_empty_intermediate_tensors ) if not get_pp_group().is_first_rank and self.use_deepstack: @@ -1069,10 +1070,10 @@ def get_mamba_state_dtype_from_config( return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config) @classmethod - def get_mamba_state_shape_from_config( + def get_mamba_state_shape_from_config( # type: ignore[override] cls, vllm_config: VllmConfig, - ) -> tuple[tuple[int, int], tuple[int, int]]: + ) -> tuple[tuple[int, int], tuple[int, int, int]]: return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config) @classmethod diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 1b1df75d37..18f4575f54 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -432,7 +432,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] self.model.make_empty_intermediate_tensors ) self.set_moe_parameters(self.model.layers) @@ -441,7 +441,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) - def forward( + def forward( # type: ignore[override] self, input_ids: torch.Tensor | None, positions: torch.Tensor, diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index 1ccad20a4b..b4902d6c46 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py @@ -334,10 +334,13 @@ def __init__( quant_method=quant_method, ) + meta_weight = self._parameters.get("weight") + if not isinstance(meta_weight, torch.Tensor): + raise RuntimeError("Qwen4Exp PLE meta weight was not initialized") host_weight = ModelWeightParameter( data=torch.empty( - tuple(self.weight.shape), - dtype=self.weight.dtype, + tuple(meta_weight.shape), + dtype=meta_weight.dtype, device="cpu", pin_memory=True, ), @@ -369,7 +372,9 @@ def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: f"Qwen4Exp pinned-host PLE requires a CUDA input, got {device}" ) device_index = ( - torch.cuda.current_device() if device.index is None else device.index + torch.accelerator.current_device_index() + if device.index is None + else device.index ) view = self._accelerator_weight_views.get(device_index) if view is None: @@ -377,20 +382,22 @@ def get_accelerator_weight(self, device: torch.device) -> torch.Tensor: raise RuntimeError( "Qwen4Exp PLE UVA view must be prepared before CUDA graph capture" ) - with torch.cuda.device(device_index): + with torch.accelerator.device_index(device_index): view = get_accelerator_view_from_cpu_tensor(self.weight) self._accelerator_weight_views[device_index] = view self._accelerator_weight_ptrs[device_index] = view.data_ptr() return view def prepare_accelerator_weight(self) -> None: - self.get_accelerator_weight(torch.device("cuda", torch.cuda.current_device())) + self.get_accelerator_weight( + torch.device("cuda", torch.accelerator.current_device_index()) + ) def embedding_lookup(self, input_: torch.Tensor) -> torch.Tensor: """Gather FP8 UVA rows and emit scaled model-dtype values.""" device_index = ( - torch.cuda.current_device() + torch.accelerator.current_device_index() if input_.device.index is None else input_.device.index ) diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index b6b68b45cb..25e9ab967d 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -429,6 +429,7 @@ def _run_qsa( def forward( self, positions: torch.Tensor, + output: torch.Tensor | None, hidden_states: torch.Tensor, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) @@ -462,8 +463,10 @@ def forward( flat_output = attn_output.view(num_tokens, -1) if gate is not None: flat_output = flat_output * torch.sigmoid(gate) - output, _ = self.o_proj(flat_output) - return output + projected_output, _ = self.o_proj(flat_output) + if output is not None: + output.copy_(projected_output) + return projected_output def qwen4_exp_qsa_with_output( diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d6ae8e0fc2..10edbf4eb5 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -634,7 +634,7 @@ def merge(cls, specs: list[Self]) -> Self: @dataclass(frozen=True) class MambaSpec(KVCacheSpec): shapes: tuple[tuple[int, ...], ...] - dtypes: tuple[torch.dtype] + dtypes: tuple[torch.dtype, ...] page_size_padded: int | None = None mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 mamba_cache_mode: str = "none" diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index 07afbd9c04..c09186ceba 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -6,7 +6,7 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config import SpeculativeConfig, VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger @@ -43,8 +43,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.device = device - self.speculative_config = vllm_config.speculative_config - assert self.speculative_config is not None + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + self.speculative_config: SpeculativeConfig = speculative_config self.method = self.speculative_config.method self.num_speculative_steps = self.speculative_config.num_speculative_tokens self.draft_model_config = self.speculative_config.draft_model_config diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 8ada303e4d..237861e3d7 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1557,9 +1557,7 @@ def __init__( elif self.speculative_config.use_step3p5_mtp(): self.drafter = Step3p5MTPProposer(self.vllm_config, self.device, self) elif self.speculative_config.use_qwen4_exp_mtp(): - self.drafter = Qwen4ExpMTPProposer( - self.vllm_config, self.device, self - ) + self.drafter = Qwen4ExpMTPProposer(self.vllm_config, self.device, self) elif self.speculative_config.use_dspark(): self.drafter = DSparkProposer(self.vllm_config, self.device, self) self.use_aux_hidden_state_outputs = True