lesj0610andClaude Opus 5
3b1911c0ca
Three initialization and config defects found serving Qwen4Exp, each with a regression test. **The MTP hidden buffer lands on the ambient device.** torch.empty for the spec-decode hidden buffer takes no device, so it follows whatever device is current rather than vllm_config.device_config.device. The allocation was duplicated verbatim in the NVIDIA and AMD backends; it moves to common/mtp.py so the two cannot drift, and the caller passes is_last_rank rather than reaching for the process group from shared code. **M-RoPE keys are reported as unrecognized.** Qwen4Exp drives M-RoPE through mrope_section and mrope_interleaved, which transformers' rope validator does not know for rope_type="default", so every load warns about a configuration the model genuinely uses. Declare them at rope validation on the text config and on the outer config that mirrors it; every other key still warns. **transformers' layer_types spelling is unread.** Its Qwen4Exp config canonicalizes "full_attention" entries to "qwen_sparse_attention" in memory, so checkpoints re-exported through it serialize that name. vLLM keys QSA off "full_attention" plus indexer_n_heads, so map the spelling back. **Qwen4Exp is missing from the Triton warmup set.** Its GDN prefill and QSA run on other backends, but the kernels this warmup covers -- the gated RMS norm, the causal conv, the post-conv preparation, the decode state update -- are Triton whichever prefill kernel is chosen, so they compile on the first prompt instead. qwen4_exp_mtp stays out: kernel_warmup passes the worker's own model_config, which is the target model's. tests/models/qwen4_exp/test_config.py: 15 passed. tests/model_executor/test_qwen_triton_warmup.py: 6 passed. ruff check, ruff format --check and mypy (Python 3.10) report what they already report on the base commit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Shared MTP drafter state for Qwen4Exp."""
|
|
|
|
import torch
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.transformers_utils.configs.qwen4_exp import Qwen4ExpTextConfig
|
|
|
|
|
|
def make_mtp_hidden_buffer(
|
|
vllm_config: VllmConfig,
|
|
config: Qwen4ExpTextConfig,
|
|
*,
|
|
is_last_rank: bool,
|
|
) -> torch.Tensor | None:
|
|
"""Buffer holding the pre-final-mixer multi-stream hidden state, or None.
|
|
|
|
When speculative ``method == "mtp"`` the drafter feeds a real multi-stream
|
|
backbone hidden on its first step, so the last stage retains
|
|
``[T, hc_count * H]``. Whether it is needed is derived from the config
|
|
alone, not from node identity, so prefill and decode nodes agree.
|
|
"""
|
|
spec_config = vllm_config.speculative_config
|
|
needs_mtp_hidden = (
|
|
spec_config is not None
|
|
and getattr(spec_config, "method", None) == "mtp"
|
|
and is_last_rank
|
|
)
|
|
if not needs_mtp_hidden:
|
|
return None
|
|
return torch.empty(
|
|
vllm_config.scheduler_config.max_num_batched_tokens,
|
|
config.hc_count * config.hidden_size,
|
|
dtype=vllm_config.model_config.dtype,
|
|
device=vllm_config.device_config.device,
|
|
)
|