diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 469fca7b04d0..40695f8a800e 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -10,6 +10,8 @@ import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_fused_post_pre_tilelang, + mhc_fused_post_pre_tilelang_reuse_residual, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -1158,6 +1160,69 @@ def run_ref(): torch.testing.assert_close(x, layer_input_ref, atol=1e-2, rtol=1e-2) +def _mhc_fused_post_pre_inputs( + num_tokens: int, hidden_size: int, hc_mult: int = 4 +) -> tuple[ + torch.Tensor, + torch.Tensor, + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + float, + float, + float, + float, + int, + ], + torch.Tensor, +]: + hc_mult3 = hc_mult * 2 + hc_mult * hc_mult + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + mixes = ( + torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32), + torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32), + torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4, + torch.randn((3,), dtype=torch.float32) * 0.1, + torch.randn((hc_mult3,), dtype=torch.float32) * 0.1, + 1e-6, + 1e-6, + 1e-6, + 1.0, + 20, + ) + norm_weight = torch.randn(hidden_size, dtype=torch.bfloat16) + return x, residual, mixes, norm_weight + + +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("num_tokens", [8, 33, 1024]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +def test_mhc_fused_post_pre_reuse_residual_is_exact(num_tokens, hidden_size): + """Writing the next residual into the dead input residual changes no + output, and the unfused path then allocates no new residual.""" + torch.set_default_device(DEVICE) + set_random_seed(0) + x, residual, mixes, norm_weight = _mhc_fused_post_pre_inputs( + num_tokens, hidden_size + ) + + untouched = residual.clone() + expected = mhc_fused_post_pre_tilelang(x, residual, *mixes, norm_weight=norm_weight) + assert torch.equal(residual, untouched) + + actual = mhc_fused_post_pre_tilelang_reuse_residual( + x, residual, *mixes, norm_weight=norm_weight + ) + for got, want in zip(actual, expected): + assert torch.equal(got, want) + unfused = mhc_fused_post_pre_split_config(num_tokens, hidden_size, 4) is None + assert (actual[0].data_ptr() == residual.data_ptr()) == unfused + + def _rocm_mhc_inputs(num_tokens=2, hidden_size=256, hc_mult=4): residual = torch.randn( (num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16, device=DEVICE @@ -1676,6 +1741,58 @@ def _patch_first_rank_pp_group(monkeypatch): ) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") +@pytest.mark.parametrize("carried", [False, True]) +def test_deepseek_v4_layer_reuse_residual_keeps_caller_input(carried): + """reuse_residual changes no output and never writes the caller's x, even + when a first layer starts from ``residual = x``.""" + torch.set_default_device(DEVICE) + set_random_seed(0) + hc_mult, hidden_size, num_tokens = 4, 4096, 64 + layer = _make_mhc_decoder_layer(hc_mult=hc_mult, hidden_size=hidden_size) + _, _, mixes, norm_weight = _mhc_fused_post_pre_inputs(num_tokens, hidden_size) + _, _, fn, hc_scale, hc_base, *_ = mixes + layer.hc_attn_fn.data.copy_(fn) + layer.hc_attn_scale, layer.hc_attn_base = hc_scale, hc_base + layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base = fn, hc_scale, hc_base + norm = SimpleNamespace(weight=norm_weight, variance_epsilon=1e-6) + layer.attn_norm = layer.ffn_norm = norm + layer.rms_norm_eps, layer.hc_eps, layer.hc_post_alpha = 1e-6, 1e-6, 1.0 + layer.hc_sinkhorn_iters = 20 + layer.use_sequence_parallel = False + layer.attn = lambda positions, x, _: x * 0.5 + layer.ffn = lambda x, input_ids, metadata: x * 2.0 + positions = torch.arange(num_tokens) + + x = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + if carried: + residual = torch.randn_like(x) + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) + post_mix = torch.rand((num_tokens, hc_mult, 1), dtype=torch.float32) + res_mix = torch.rand((num_tokens, hc_mult, hc_mult), dtype=torch.float32) + else: + residual = post_mix = res_mix = None + + def run(reuse_residual): + return DeepseekV4DecoderLayer.forward( + layer, + x.clone() if carried else x, + positions, + None, + post_mix, + res_mix, + None if residual is None else residual.clone(), + reuse_residual=reuse_residual, + ) + + x_before = x.clone() + expected = run(False) + actual = run(True) + assert torch.equal(x, x_before) + for got, want in zip(actual, expected): + assert torch.equal(got, want) + + def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): """First finalize (at the end of load_weights) allocates hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 62a9a1635c1f..4ae5b436f993 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -34,6 +34,7 @@ from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer from vllm.v1.spec_decode.metadata import SpecDecodeMetadata +from vllm.v1.spec_decode.utils import extend_all_queries_by_N from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch model_dir = "meta-llama/Llama-3.1-8B-Instruct" @@ -1004,6 +1005,124 @@ def create_deterministic_logits(token_ids): assert torch.equal(result, expected_tokens) +@pytest.mark.skipif(current_platform.is_rocm(), reason="uses FLASH_ATTN") +def test_propose_refreshes_layout_caches_each_draft_step(): + """Builders called inside the draft loop must see token_to_req_indices() + and compute_num_computed_tokens() for the current layout, not values + cached during the ragged first pass or an earlier draft step.""" + device = torch.device(DEVICE_TYPE) + num_speculative_tokens = 3 + seq_lens = [5, 3] + batch_size, total_tokens, vocab_size = len(seq_lens), sum(seq_lens), 100 + + proposer = _create_proposer( + "eagle", num_speculative_tokens, attention_backend="FLASH_ATTN" + ) + hidden_size = proposer.hidden_size + model_mock = mock.MagicMock() + model_mock.side_effect = [ + (torch.zeros(n, hidden_size, device=device),) * 2 + for n in [total_tokens] + [batch_size] * (num_speculative_tokens - 1) + ] + model_mock.compute_logits.side_effect = [ + torch.zeros(batch_size, vocab_size, device=device) + ] * num_speculative_tokens + proposer.model = model_mock + proposer._draft_attn_layer_names = {"layer.0"} + + builder_cls, _ = try_get_attention_backend(AttentionBackendEnum.FLASH_ATTN) + builder = builder_cls( + kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config), + layer_names=["layer.0"], + vllm_config=proposer.vllm_config, + device=device, + ) + req_index_buffer = torch.zeros(64, dtype=torch.int32, device=device) + observed = [] + + def build_for_drafting(common_attn_metadata, draft_index): + query_lens = common_attn_metadata.query_start_loc.diff() + observed.append( + ( + common_attn_metadata.token_to_req_indices(req_index_buffer).tolist(), + torch.repeat_interleave( + torch.arange(len(query_lens), device=device), query_lens + ).tolist(), + common_attn_metadata.compute_num_computed_tokens().tolist(), + (common_attn_metadata.seq_lens - query_lens).tolist(), + ) + ) + return builder.build_for_drafting( + common_attn_metadata=common_attn_metadata, draft_index=draft_index + ) + + attn_group = mock.MagicMock() + attn_group.get_metadata_builder.return_value = builder + attn_group.build_metadata_for_drafting = build_for_drafting + attn_group.layer_names = ["layer.0"] + attn_group.kv_cache_spec = builder.kv_cache_spec + proposer.draft_attn_groups = [attn_group] + + proposer.propose( + num_speculative_tokens=num_speculative_tokens, + target_token_ids=torch.randint(0, vocab_size, (total_tokens,), device=device), + target_positions=torch.cat([torch.arange(n, device=device) for n in seq_lens]), + target_hidden_states=torch.randn(total_tokens, hidden_size, device=device), + next_token_ids=torch.randint( + 0, vocab_size, (batch_size,), dtype=torch.int32, device=device + ), + token_indices_to_sample=None, + common_attn_metadata=create_common_attn_metadata( + BatchSpec(seq_lens=seq_lens, query_lens=seq_lens), + block_size=BLOCK_SIZE, + device=device, + ), + sampling_metadata=mock.MagicMock(), + ) + + assert len(observed) == num_speculative_tokens + for step, (req, want_req, computed, want_computed) in enumerate(observed): + assert req == want_req, f"draft step {step}" + assert computed == want_computed, f"draft step {step}" + + +def test_extend_all_queries_by_n_refreshes_token_to_req_indices(): + """The extended layout must not reuse the token->request mapping cached + for the layout it was extended from.""" + device = torch.device(DEVICE_TYPE) + common_attn_metadata = create_common_attn_metadata( + BatchSpec(seq_lens=[8, 4, 4], query_lens=[3, 1, 1]), + block_size=BLOCK_SIZE, + device=device, + ) + req_index_buffer = torch.zeros(16, dtype=torch.int32, device=device) + assert common_attn_metadata.token_to_req_indices(req_index_buffer).tolist() == [ + 0, + 0, + 0, + 1, + 2, + ] + + extended = extend_all_queries_by_N( + common_attn_metadata, + N=1, + arange=torch.arange(8, dtype=torch.int32, device=device), + new_slot_mapping=torch.zeros(8, dtype=torch.int64, device=device), + ) + + assert extended.token_to_req_indices(req_index_buffer).tolist() == [ + 0, + 0, + 0, + 0, + 1, + 1, + 2, + 2, + ] + + @pytest.mark.parametrize( "attn_backend", ["ROCM_ATTN", "TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"], diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index 7cc027429819..af0a2c742ae2 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -24,6 +24,7 @@ ) from vllm.config.load import LoadConfig from vllm.model_executor.models.llama import LlamaForCausalLM +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.spec_decode.eagle import EagleProposer @@ -237,3 +238,24 @@ def create_deterministic_logits(batch_size, vocab_size, token_offset): assert model_mock.called # Verify output shape assert result.shape == (batch_size, num_speculative_tokens) + + +@pytest.mark.parametrize( + ("method", "expected"), + [ + (None, False), + ("ngram", False), + ("mtp", True), + ("eagle3", True), + ("draft_model", True), + ("dspark", False), + ], +) +def test_deepseek_v4_mtp_hidden_buffer_skips_dspark(method, expected): + """Only drafters that read DeepSeek V4's pre-hc_head residual get the + buffer; DSpark drafts from the target's aux hidden states.""" + spec_config = None + if method is not None: + spec_config = SpeculativeConfig.__new__(SpeculativeConfig) + object.__setattr__(spec_config, "method", method) + assert needs_mtp_hidden_buffer(spec_config) is expected diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 2c0f4c12a027..5a291b9c320f 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -798,7 +798,7 @@ def mhc_post_tilelang( return out -def mhc_fused_post_pre_tilelang( +def _mhc_fused_post_pre_tilelang_impl( x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, @@ -815,6 +815,7 @@ def mhc_fused_post_pre_tilelang( tile_n: int = 1, norm_weight: torch.Tensor | None = None, norm_eps: float = 1e-6, + reuse_residual: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Run one MHC post block followed by the next MHC pre block. @@ -825,6 +826,11 @@ def mhc_fused_post_pre_tilelang( ``n_splits`` and ``tile_n`` are retained for the shared MHC operator API. The TileLang path selects both values internally from the runtime shape. + With ``reuse_residual``, the unfused path writes residual_cur into + ``residual``'s storage instead of allocating a new buffer. The post kernel + reads every hc row of a hidden block before writing that block, so the + in-place update is exact. + Returns: residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size) post_mix_cur: shape (..., hc_mult, 1) @@ -927,7 +933,9 @@ def mhc_fused_post_pre_tilelang( norm_eps=norm_eps, ) else: - residual_cur = torch.empty_like(residual_flat) + residual_cur = ( + residual_flat if reuse_residual else torch.empty_like(residual_flat) + ) _MHC_POST_TILELANG_KERNEL( comb_res_mix_flat, residual_flat, @@ -969,6 +977,85 @@ def mhc_fused_post_pre_tilelang( ) +def mhc_fused_post_pre_tilelang( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """See ``_mhc_fused_post_pre_tilelang_impl``. Never writes to its inputs.""" + return _mhc_fused_post_pre_tilelang_impl( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + + +def mhc_fused_post_pre_tilelang_reuse_residual( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Like ``mhc_fused_post_pre_tilelang`` but may overwrite ``residual``. + + Not a registered custom op. Only for eager callers that own ``residual`` + and do not read it after this call. + """ + return _mhc_fused_post_pre_tilelang_impl( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + norm_weight=norm_weight, + norm_eps=norm_eps, + reuse_residual=True, + ) + + def _mhc_fused_post_pre_tilelang_fake( x: torch.Tensor, residual: torch.Tensor, diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 367c8f2cec2c..b0d798033f11 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -82,6 +82,7 @@ ) from vllm.models.deepseek_v4.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.platforms import current_platform from vllm.platforms.rocm import on_gfx950, on_gfx1250 from vllm.sequence import IntermediateTensors @@ -1191,10 +1192,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.hc_head_op = HCHeadOp() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() - ) - if get_pp_group().is_last_rank and needs_mtp_hidden_states: + if get_pp_group().is_last_rank and needs_mtp_hidden_buffer(spec_config): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, diff --git a/vllm/models/deepseek_v4/common/spec_decode.py b/vllm/models/deepseek_v4/common/spec_decode.py new file mode 100644 index 000000000000..c4b5e5a2f97f --- /dev/null +++ b/vllm/models/deepseek_v4/common/spec_decode.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from vllm.config import SpeculativeConfig + + +def needs_mtp_hidden_buffer(spec_config: SpeculativeConfig | None) -> bool: + """Whether the drafter reads the target's pre-hc_head residual buffer. + + DSpark drafts from the target's aux hidden states, so it skips the + (max_num_batched_tokens, hc_mult * hidden_size) buffer. + """ + return ( + spec_config is not None + and (spec_config.use_eagle() or spec_config.uses_draft_model()) + and not spec_config.use_dspark() + ) diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py index 71b9a9810d85..63ec337fa48b 100644 --- a/vllm/models/deepseek_v4/cpu/model.py +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -58,6 +58,7 @@ make_layers, maybe_prefix, ) +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.models.deepseek_v4.cpu.cpu_sparse import DeepseekV4CPUAttention from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -524,10 +525,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.hc_head_op = HCHeadOp() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() - ) - if get_pp_group().is_last_rank and needs_mtp_hidden_states: + if get_pp_group().is_last_rank and needs_mtp_hidden_buffer(spec_config): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index a6ec62f7aa52..04f7ae21f95d 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -21,6 +21,7 @@ from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, + mhc_fused_post_pre_tilelang_reuse_residual, mhc_post_tilelang, mhc_pre_broadcast_tilelang, mhc_pre_tilelang, @@ -76,6 +77,7 @@ ) from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.models.deepseek_v4.common.eplb_util import collect_moe_layers +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, DeepseekV4FlashInferSM120Attention, @@ -720,7 +722,16 @@ def forward( res_mix: torch.Tensor | None = None, residual: torch.Tensor | None = None, mega_gate_metadata: MegaGateRoutingMetadata | None = None, + reuse_residual: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """With ``reuse_residual``, the caller hands ``residual`` over and the + unfused mHC post/pre writes the next residual into its storage. A first + layer that starts from ``residual = x`` never overwrites ``x``.""" + fused_post_pre = ( + mhc_fused_post_pre_tilelang_reuse_residual + if reuse_residual + else mhc_fused_post_pre_tilelang + ) attn_norm_weight = self.attn_norm.weight.data attn_norm_eps = self.attn_norm.variance_epsilon if residual is None: @@ -743,6 +754,7 @@ def forward( ) else: residual = x + fused_post_pre = mhc_fused_post_pre_tilelang post_mix, res_mix, x = mhc_pre_tilelang( x, self.hc_attn_fn, @@ -757,7 +769,7 @@ def forward( norm_eps=attn_norm_eps, ) else: - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + residual, post_mix, res_mix, x = fused_post_pre( x, residual, post_mix, @@ -785,7 +797,7 @@ def forward( ffn_norm_weight = self.ffn_norm.weight.data ffn_norm_eps = self.ffn_norm.variance_epsilon - residual, post_mix, res_mix, x = mhc_fused_post_pre_tilelang( + residual, post_mix, res_mix, x = fused_post_pre( x, residual, post_mix, @@ -900,10 +912,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): requires_grad=False, ) spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() - ) - if get_pp_group().is_last_rank and needs_mtp_hidden_states: + if get_pp_group().is_last_rank and needs_mtp_hidden_buffer(spec_config): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, @@ -1002,6 +1011,7 @@ def forward( res_mix, residual, mega_gate_metadata=mega_gate_metadata, + reuse_residual=True, ) if idx + 1 in self.aux_hidden_state_layers: # Reconstruct the aux hidden state for draft models diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index 80c579abaee9..a2fb37a1c71b 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -59,6 +59,7 @@ maybe_prefix, ) from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.models.deepseek_v4.xpu.xpu_sparse import DeepseekV4XPUAttention from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -1040,10 +1041,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self.hc_head_op = HCHeadOp() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() - ) - if get_pp_group().is_last_rank and needs_mtp_hidden_states: + if get_pp_group().is_last_rank and needs_mtp_hidden_buffer(spec_config): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index de2f4e7cf9ad..e7c513047140 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -59,6 +59,7 @@ from vllm.models.deepseek_v4.amd.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.models.deepseek_v41.amd.mono_decode import MonoDecodeLayer from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterMLAAttention from vllm.models.deepseek_v41.attention import DeepseekV4Attention @@ -643,10 +644,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.mhc_post = MHCPostOp() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() - ) - if get_pp_group().is_last_rank and needs_mtp_hidden_states: + if get_pp_group().is_last_rank and needs_mtp_hidden_buffer(spec_config): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 78d3d26276d2..b71dadfe9755 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -75,6 +75,7 @@ sp_reduce_scatter, sp_shard, ) +from vllm.models.deepseek_v4.common.spec_decode import needs_mtp_hidden_buffer from vllm.models.deepseek_v4.nvidia.model import ( DeepseekV4MoE as DeepseekV4MoEBase, ) @@ -936,10 +937,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.norm = PPMissingLayer() spec_config = vllm_config.speculative_config - needs_mtp_hidden_states = spec_config is not None and ( - spec_config.use_eagle() or spec_config.uses_draft_model() - ) - if get_pp_group().is_last_rank and needs_mtp_hidden_states: + if get_pp_group().is_last_rank and needs_mtp_hidden_buffer(spec_config): self._mtp_hidden_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, self.hc_dim, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index d1bc18472ee3..5c6ce28a956d 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -486,6 +486,14 @@ def naive_query_lens(self) -> torch.Tensor: def replace(self, **kwargs) -> "CommonAttentionMetadata": return replace(self, **kwargs) + def invalidate_cached_layout(self) -> None: + """Drop values cached from query_start_loc and seq_lens. + + Call after updating either in place. + """ + self._num_computed_tokens_cache = None + self._token_to_req_indices_cache = None + def compute_num_computed_tokens(self) -> torch.Tensor: """Compute num_computed_tokens on device (seq_lens - query_lens).""" if self._num_computed_tokens_cache is None: diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index eb4db3c17521..d8e97b5beea4 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -682,6 +682,7 @@ def propose( # (i.e., not the first proposal). if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None: common_attn_metadata.seq_lens -= num_rejected_tokens_gpu + common_attn_metadata.invalidate_cached_layout() block_size = self.block_size assert block_size > 0, "block_size has not been initialized." @@ -810,6 +811,8 @@ def _update_positions_dependent_metadata( if common_attn_metadata.seq_lens_cpu_upper_bound is not None: common_attn_metadata.seq_lens_cpu_upper_bound += 1 + # seq_lens was advanced in place by the kernel above. + common_attn_metadata.invalidate_cached_layout() return positions diff --git a/vllm/v1/spec_decode/utils.py b/vllm/v1/spec_decode/utils.py index 471d51d355e5..41c91e0e474d 100644 --- a/vllm/v1/spec_decode/utils.py +++ b/vllm/v1/spec_decode/utils.py @@ -413,6 +413,8 @@ def extend_all_queries_by_N( new_cad = cad.replace( query_start_loc=new_query_start_loc, query_start_loc_cpu=new_query_start_loc_cpu, + _num_computed_tokens_cache=None, + _token_to_req_indices_cache=None, seq_lens=cad.seq_lens + N, # each request is extended by N tokens -> batch_size * N tokens are added num_actual_tokens=cad.num_actual_tokens + cad.batch_size() * N,