diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 552910e85a..045e04533a 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -82,18 +82,20 @@ steps: - image-build-amd source_file_dependencies: - vllm/models/deepseek_v4/ + - vllm/models/deepseek_v41/ - tests/models/test_deepseek_v4_vl_rocm.py - tests/models/test_hyv4_rocm.py - tests/models/transformers/test_layer_registry_e2e.py - tests/models/test_deepseek_v41_replay_start.py - tests/models/test_deepseek_v41_decoder_replay_layers.py - tests/models/test_deepseek_v41_replay_batch.py + - tests/models/test_deepseek_v41_amd_replay_forward.py - tests/models/test_deepseek_v4_dspark_rocm.py - tests/models/test_deepseek_v4_rocm_wo_a.py - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py - tests/models/test_dsv41_mono_decode.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - label: ":nvidia: (B200) Inkling" diff --git a/tests/models/test_deepseek_v41_amd_replay_forward.py b/tests/models/test_deepseek_v41_amd_replay_forward.py new file mode 100644 index 0000000000..fe22bc8025 --- /dev/null +++ b/tests/models/test_deepseek_v41_amd_replay_forward.py @@ -0,0 +1,329 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""What the AMD replay forward does with a prefill longer than the window. + +The cut in this fixture is layer 1, so replay starts at layer 1. That is the +same boundary as production layer 20, on a 4-layer stand-in. The final +``mhc_post`` and the aux capture at that cut both run on the window rows. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm.forward_context import ForwardContext, override_forward_context +from vllm.models.deepseek_v41.decoder_replay_layers import ReplayBatch +from vllm.sequence import IntermediateTensors + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +WINDOW = 128 +FULL = 300 +# Trailing window of a 300-token prefill. +ROWS = list(range(FULL - WINDOW, FULL)) +HC, HIDDEN = 4, 8 +DEVICE = torch.device("cuda") +# gfx942 writes the V4 record (584 B). Family-100 writes the V4.1 record (528 B). +# The SWA layer uses block size 32, which is the size the KV-only insert checks. +_BLOCK = 32 +_HEAD_DIM = 512 + + +class _RocmKvInsert: + """SWA cache from ``forward_kv``, checked against the KV-only insert.""" + + def __init__(self, num_tokens: int): + from vllm.models.deepseek_v41.attention import _use_v41_mxfp8_kv_record + + if not hasattr(torch.ops._C, "fused_deepseek_v4_kv_rope_insert"): + pytest.skip("fused DeepseekV4 KV insert op is not built") + + self.num_tokens = num_tokens + self.kv_mxfp8 = _use_v41_mxfp8_kv_record() + bytes_per = 528 if self.kv_mxfp8 else 584 + num_blocks = (num_tokens + _BLOCK - 1) // _BLOCK + 1 + cache = torch.full( + (num_blocks, _BLOCK, bytes_per), + 0xA5, + dtype=torch.uint8, + device=DEVICE, + ) + self.poison = cache.clone() + self.kv = torch.randn( + num_tokens, _HEAD_DIM, dtype=torch.bfloat16, device=DEVICE + ) + prefix = "replay.swa" + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=DEVICE) + self.metadata = { + prefix: SimpleNamespace(slot_mapping=slot_mapping, block_size=_BLOCK) + } + self.context = ForwardContext({}, self.metadata, {}) + inv_freq = 1.0 / ( + 10000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device=DEVICE) / 64) + ) + freqs = torch.outer( + torch.arange(num_tokens + 8, dtype=torch.float32, device=DEVICE), inv_freq + ) + self.attn = SimpleNamespace( + kv_mxfp8=self.kv_mxfp8, + compressor=None, + indexer=None, + rotary_emb=SimpleNamespace( + cos_sin_cache=torch.cat((freqs.cos(), freqs.sin()), dim=-1) + ), + swa_cache_layer=SimpleNamespace( + prefix=prefix, kv_cache=cache, block_size=_BLOCK + ), + _run_parallel_input_projections=lambda hidden: (hidden, None, None), + _split_qkv_and_norm=lambda qr_kv: (None, None, self.kv[: qr_kv.shape[0]]), + ) + self.last_rows: int | None = None + self.last_positions: torch.Tensor | None = None + + def write( + self, + x, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + ): + from vllm.models.deepseek_v41.amd.rocm import ( + DeepseekV41ROCMAiterMLAAttention, + ) + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + assert ( + DeepseekV41ROCMAiterMLAAttention.forward_kv + is DeepseekV4Attention.forward_kv + ) + self.last_rows = x.shape[0] + self.last_positions = positions + DeepseekV41ROCMAiterMLAAttention.forward_kv(self.attn, positions, x) + + def assert_matches_forward(self): + assert self.last_rows == self.num_tokens + assert self.last_positions is not None + written = self.attn.swa_cache_layer.kv_cache + assert not torch.equal(written, self.poison) + ref = self.poison.clone() + meta = self.metadata[self.attn.swa_cache_layer.prefix] + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + self.kv, + ref, + meta.slot_mapping, + self.last_positions, + self.attn.rotary_emb.cos_sin_cache, + _BLOCK, + None, + self.kv_mxfp8, + ) + assert torch.equal(written, ref) + assert torch.equal(written[-1], self.poison[-1]) + + +class _Layer(nn.Module): + def __init__( + self, + idx: int, + compressed: str | None = None, + indexer: str | None = None, + ): + super().__init__() + self.idx = idx + self.seen: list[torch.Tensor] = [] + self.attn = SimpleNamespace( + swa_cache_layer=SimpleNamespace(prefix=f"layers.{idx}.swa"), + compressed_cache_prefix=compressed, + indexer=( + None + if indexer is None + else SimpleNamespace(k_cache=SimpleNamespace(prefix=indexer)) + ), + compress_ratio=1, + ) + + def forward( + self, + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + ): + self.seen.append(hidden_states[:, 0, 0].detach().clone()) + if residual is None: + residual = hidden_states + if pre_mix is None: + pre_mix = hidden_states[:, :, 0] + if post_mix is None: + post_mix = pre_mix + if res_mix is None: + res_mix = pre_mix + return hidden_states, residual, post_mix, res_mix, pre_mix + + +def _model(monkeypatch): + import vllm.models.deepseek_v41.amd.model as amd + + monkeypatch.setattr( + amd, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=False, is_last_rank=False), + ) + monkeypatch.setattr(amd.EngramLayout, "from_config", lambda config: None) + monkeypatch.setattr(amd, "MHCPostOp", lambda: SimpleNamespace()) + + # Layer 0 sits before the cut, so its keys are not collected. Layer 2 + # shares layer 1's compressed cache. Layer 3 has neither. + layers = [ + _Layer(0, "layers.0.compressed", "layers.0.indexer"), + _Layer(1, "layers.1.compressed", "layers.1.indexer"), + _Layer(2, "layers.1.compressed", "layers.2.indexer"), + _Layer(3), + ] + # Layer 1 is the KV source, so the replay callback is this layer's write. + insert = _RocmKvInsert(FULL) + layers[1].write_kv = insert.write + monkeypatch.setattr(amd, "make_layers", lambda n, factory, prefix: (0, n, layers)) + + config = SimpleNamespace( + vocab_size=32, + hc_eps=1e-6, + hc_mult=HC, + hidden_size=HIDDEN, + rms_norm_eps=1e-6, + index_topk=4, + kv_source_layer_ids=[1], + num_hidden_layers=4, + sliding_window=WINDOW, + engram_layer_ids=(), + candidate_source_layer_id=-1, + candidate_topk_blocks=0, + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=config, dtype=torch.bfloat16), + quant_config=None, + cache_config=SimpleNamespace(swa_bounded_replay=True), + use_v2_model_runner=True, + parallel_config=SimpleNamespace( + prefill_context_parallel_size=1, + use_ubatching=False, + enable_expert_parallel=False, + ), + speculative_config=None, + scheduler_config=SimpleNamespace(max_num_batched_tokens=FULL), + ) + model = amd.DeepseekV4Model(vllm_config=vllm_config) + # Production requests aux at the first replay layer. Here that id is the cut. + model.aux_hidden_state_layers = (model.decoder_replay_start,) + model.topk_indices_buffer = model.topk_indices_buffer.to(DEVICE) + + calls: list[torch.Tensor] = [] + + def mhc_post(hidden, residual, post_mix, res_mix): + calls.append(hidden[:, 0, 0].detach().clone()) + return hidden + + model.mhc_post = mhc_post + saved_aux: list[list[torch.Tensor]] = [] + run_layers = model._run_layers + + def wrapped(*args, **kwargs): + result = run_layers(*args, **kwargs) + # Snapshot: forward() extends the early list with the scattered aux. + saved_aux.append(list(result[5])) + return result + + model._run_layers = wrapped + scattered: list[tuple[torch.Tensor, ...]] = [] + replay = model.decoder_replay_layers + run = replay._run + + def wrapped_run(*args, **kwargs): + result = run(*args, **kwargs) + scattered.append(result) + return result + + replay._run = wrapped_run + return model, layers, calls, saved_aux, scattered, insert + + +def _forward(model, insert: _RocmKvInsert): + token_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32) + hidden = token_ids[:, None, None].expand(FULL, HC, HIDDEN).contiguous() + pre_mix = torch.ones(FULL, HC, device=DEVICE) + rows = torch.tensor(ROWS, device=DEVICE) + model.decoder_replay_layers.replay_batch = ReplayBatch( + rows, + ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={}), + ) + positions = torch.arange(FULL, device=DEVICE) + with override_forward_context(insert.context): + out = model( + input_ids=positions.to(torch.int32), + positions=positions, + intermediate_tensors=IntermediateTensors( + {"hidden_states": hidden, "pre_mix": pre_mix} + ), + ) + return out + + +def test_amd_replay_forward_trims_to_the_window(monkeypatch): + model, layers, calls, saved_aux, scattered, insert = _model(monkeypatch) + replay_layers = model.decoder_replay_layers + assert model.decoder_replay_start == 1 + assert replay_layers.first_swa_prefix == "layers.1.swa" + assert replay_layers.metadata_prefixes == { + "layers.1.swa", + "layers.1.compressed", + "layers.1.indexer", + "layers.2.swa", + "layers.2.indexer", + "layers.3.swa", + } + assert replay_layers.replay_batch is None + out = _forward(model, insert) + insert.assert_matches_forward() + replay = model.decoder_replay_layers.replay_batch + assert replay is not None and replay.rows.tolist() == ROWS + + early = [layer.seen[0] for layer in layers[: model.decoder_replay_start]] + late = [layer.seen[0] for layer in layers[model.decoder_replay_start :]] + full_ids = torch.arange(FULL, device=DEVICE, dtype=torch.float32) + window_ids = full_ids[ROWS] + assert [t.tolist() for t in early] == [full_ids.tolist()] + assert [t.tolist() for t in late] == [ + window_ids.tolist(), + window_ids.tolist(), + window_ids.tolist(), + ] + + # Scatter writes the replay rows back and leaves the trimmed rows at zero. + hidden = out["hidden_states"][:, 0, 0] + assert torch.equal(hidden[ROWS], window_ids) + assert hidden[: FULL - WINDOW].eq(0).all() + + # calls[0] is the cut-aux capture. calls[1] is the final mhc_post. + # Both run inside the replay batch, before the scatter back to FULL. + assert len(calls) == 2 + assert [t.shape[0] for t in calls] == [WINDOW, WINDOW] + assert torch.equal(calls[0], window_ids) + assert torch.equal(calls[1], window_ids) + + # The early loop skips this id. The replay loop does not see it either. + assert saved_aux[0] == [] + assert saved_aux[1] == [] + _hidden, _pre_mix, aux = scattered[0] + assert aux.shape[0] == FULL + assert torch.equal(aux[ROWS, 0], window_ids) + assert aux[: FULL - WINDOW].eq(0).all() diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index bcd4db64b2..5a3625dcbf 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -290,6 +290,81 @@ def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state): assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [1, 129] +def test_zero_prefix_window_warms_gather_and_skips_pad(monkeypatch): + """ROCm leaves the prefix-cache window at 0. Replay still needs replay_start + and a warmed gather kernel. The pad kernel stays on the prefix-cache path.""" + import vllm.models.deepseek_v41.nvidia.model_state as model_state + + calls = {"pad_warmup": 0, "gather_warmup": 0, "pad_launch": 0} + + def pad_warmup(*args, **kwargs): + calls["pad_warmup"] += 1 + + def gather_warmup(*args, **kwargs): + calls["gather_warmup"] += 1 + + class _Pad: + warmup = staticmethod(pad_warmup) + + def __getitem__(self, grid): + def launch(*args, **kwargs): + calls["pad_launch"] += 1 + + return launch + + monkeypatch.setattr(model_state, "_pad_replayed_slots_kernel", _Pad()) + monkeypatch.setattr( + model_state._gather_replay_batch_kernel, "warmup", gather_warmup + ) + + cfg = MagicMock() + cfg.model_config.enable_prompt_embeds = False + cfg.model_config.uses_mrope = False + cfg.model_config.is_multimodal_model = False + cfg.scheduler_config.max_num_seqs = 16 + cfg.scheduler_config.max_num_batched_tokens = 1024 + cfg.parallel_config.data_parallel_size = 1 + cfg.compilation_config.fast_moe_cold_start = False + cfg.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + layers = SimpleNamespace(window=WINDOW, replay_batch=None, metadata_prefixes=set()) + model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) + groups = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=["swa"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=False, prefix_replay_tokens=0 + ), + ), + SimpleNamespace( + layer_names=["mla"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=True, prefix_replay_tokens=0 + ), + ), + ] + ) + + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw): + return {"swa": kw["model_specific_attn_metadata"]} + + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) + state = DeepseekV41ModelState(cfg, model, None, DEVICE) + state._replay_start_np[0] = 16 + batch = _input_batch([8], [24], [True]) + state.prepare_attn( + batch, + CUDAGraphMode.NONE, + BLOCK_TABLES, + _slot_mappings(batch.num_tokens_after_padding), + [], + groups, + ) + assert state._replay[0] == 0 + assert calls == {"pad_warmup": 0, "gather_warmup": 1, "pad_launch": 0} + assert state._replay_start[:1].tolist() == [16] + + def test_dp_ranks_pad_to_one_replay_graph(dp_state): """Every rank pads its replay batch to the graph fitting the largest rank's.""" dp_state.replay_cudagraphs = MagicMock() diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 4c5f158587..a52a408c74 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -62,6 +62,7 @@ from vllm.models.deepseek_v4.amd.model import ( 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 +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -421,6 +422,42 @@ class DeepseekV4DecoderLayer(nn.Module): x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre + def write_kv( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> None: + """Write the KV ``forward`` would for the replay layers' inputs, skipping + the rest of the layer; the replay batch then reruns its rows.""" + # The replay batch rules out Engram and sequence parallel here. + assert self.engram is None and not self.use_sequence_parallel + fuse_attn_norm = self.fuse_seam_norm + residual, post_mix, res_mix, x, _attn_pre = self.mhc_pre_delayed( + residual, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + pre_mix=pre_mix, + sublayer_out=x, + post_layer_mix=post_mix, + comb_res_mix=res_mix, + norm_weight=self.attn_norm.weight if fuse_attn_norm else None, + norm_eps=self.attn_norm.variance_epsilon, + ) + if not fuse_attn_norm: + x = self.attn_norm(x) + self.attn.forward_kv(positions, x) + class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -498,6 +535,45 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): prefix=f"{prefix}.layers", ) + # Config flag is the switch: ROCm forces swa_cache_layer.bounded_replay off. + self.decoder_replay_layers: DecoderReplayLayers | None = None + self.decoder_replay_start = self.end_layer + cut = max(config.kv_source_layer_ids) + if ( + vllm_config.cache_config.swa_bounded_replay + and vllm_config.use_v2_model_runner + and cut < self.end_layer - 1 + and self._decoder_replay_supported(vllm_config, cut) + ): + assert self.layers[cut].attn.compress_ratio <= 1 + # The last KV source writes every row's KV, then replays. + self.decoder_replay_start = cut + # The attention metadata keys the replay layers read. + metadata_prefixes: set[str] = set() + for layer in islice(self.layers, cut, self.end_layer): + attn = typing.cast(DeepseekV4DecoderLayer, layer).attn + metadata_prefixes.add(attn.swa_cache_layer.prefix) + if attn.compressed_cache_prefix is not None: + metadata_prefixes.add(attn.compressed_cache_prefix) + if attn.indexer is not None: + metadata_prefixes.add(attn.indexer.k_cache.prefix) + self.decoder_replay_layers = DecoderReplayLayers( + config.sliding_window, + self._run_replay_layers, + self.layers[cut].write_kv, + metadata_prefixes, + self.layers[cut].attn.swa_cache_layer.prefix, + ) + logger.info_once( + "Decoder SWA bounded replay: in eager prefill steps, layer %d " + "writes its KV for every token, and the rest of it and layers " + "%d-%d run on each request's last %d tokens only.", + cut, + cut + 1, + self.end_layer - 1, + config.sliding_window, + ) + # The n-gram hash needs a slot-keyed rolling store of compressed ids # (chunked prefill / decode lookback); key it off the first local # layer's sliding-window KV cache. Only PP ranks owning an engram @@ -657,13 +733,33 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - aux_hidden_states: list[torch.Tensor] = [] - final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call - for idx, layer in enumerate( - islice(self.layers, self.start_layer, self.end_layer), - start=self.start_layer, - ): - hidden_states, residual, post_mix, res_mix, pre_mix = layer( + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.start_layer, self.decoder_replay_start), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + full_num_tokens, + ) + # Replay collapses inside its own batch. The early loop still sets + # ran_layers, so that flag alone must not post again on the full batch. + collapsed = False + if self.decoder_replay_layers is not None: + hidden_states, pre_mix, *late_aux = self.decoder_replay_layers( hidden_states, positions, input_ids, @@ -671,20 +767,15 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): post_mix, res_mix, residual, - engram_hashes, - engram_mask, ) - if idx + 1 in self.aux_hidden_state_layers: - # Reconstruct the aux hidden state for draft models - aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) - aux_hidden_state = aux_recon.mean(dim=1) - if self.use_sequence_parallel: - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] - aux_hidden_states.append(aux_hidden_state) - final_aux_recon = aux_recon - if layer is not None: - # Reuse if the last layer was captured as an aux hidden state - if self.end_layer in self.aux_hidden_state_layers: + aux_hidden_states.extend(late_aux) + final_aux_recon = None + collapsed = True + if ran_layers and not collapsed: + if ( + self.end_layer in self.aux_hidden_state_layers + and final_aux_recon is not None + ): hidden_states = final_aux_recon else: hidden_states = self.mhc_post( @@ -714,6 +805,162 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): return hidden_states, aux_hidden_states return hidden_states + def _run_layers( + self, + layer_ids: range, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor | None, + post_mix: torch.Tensor | None, + res_mix: torch.Tensor | None, + residual: torch.Tensor | None, + engram_hashes: torch.Tensor | None, + engram_mask: torch.Tensor | None, + full_num_tokens: int, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + list[torch.Tensor], + torch.Tensor | None, + bool, + ]: + aux_hidden_states: list[torch.Tensor] = [] + final_aux_recon: torch.Tensor | None = None + ran_layers = False + for idx in layer_ids: + ran_layers = True + hidden_states, residual, post_mix, res_mix, pre_mix = self.layers[idx]( + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + ) + # The aux id of the first replay layer is the post of the last + # full-batch layer. Capture it on the trimmed rows instead. + if idx + 1 in self.aux_hidden_state_layers and not ( + self.decoder_replay_layers is not None + and idx + 1 == self.decoder_replay_start + ): + aux_recon = self.mhc_post(hidden_states, residual, post_mix, res_mix) + aux_hidden_state = aux_recon.mean(dim=1) + if self.use_sequence_parallel: + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) + final_aux_recon = aux_recon + return ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux_hidden_states, + final_aux_recon, + ran_layers, + ) + + def _run_replay_layers( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + pre_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Layers past the last KV source, on the rows they are given. + + The final ``mhc_post`` runs here, on this batch. The return is + ``(hidden_states, pre_mix, *aux)``. + """ + # Post of the layer before the cut. These rows are already trimmed, + # and sequence parallel is unsupported while replay is on. + cut_aux: list[torch.Tensor] = [] + if self.decoder_replay_start in self.aux_hidden_state_layers: + cut_aux.append( + self.mhc_post(hidden_states, residual, post_mix, res_mix).mean(dim=1) + ) + ( + hidden_states, + residual, + post_mix, + res_mix, + pre_mix, + aux, + final_aux_recon, + ran_layers, + ) = self._run_layers( + range(self.decoder_replay_start, self.end_layer), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + None, + None, + positions.shape[0], + ) + assert residual is not None and post_mix is not None + assert res_mix is not None and pre_mix is not None + if ( + self.end_layer in self.aux_hidden_state_layers + and final_aux_recon is not None + ): + hidden_states = final_aux_recon + elif ran_layers: + hidden_states = self.mhc_post(hidden_states, residual, post_mix, res_mix) + return (hidden_states, pre_mix, *cut_aux, *aux) + + def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: + """Whether this rank may replay from layer ``cut`` on; warns when not.""" + parallel_config = vllm_config.parallel_config + spec_config = vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + draft_hf_config = getattr(draft_config, "hf_config", None) + draft_window = getattr(draft_hf_config, "sliding_window", None) + draft_layer_types = getattr(draft_hf_config, "layer_types", None) or () + window = self.config.sliding_window + if self.start_layer > cut or self.end_layer < self.config.num_hidden_layers: + reason = ( + "the pipeline stage holding the last KV source layer must also " + "hold every layer after it" + ) + elif ( + self.use_sequence_parallel + or parallel_config.prefill_context_parallel_size > 1 + or parallel_config.use_ubatching + ): + reason = ( + "the replay-layer batch shrinks per rank, which sequence and " + "prefill-context parallelism and microbatching cannot follow" + ) + elif any(i >= cut for i in getattr(self.config, "engram_layer_ids", ())): + reason = "an Engram layer sits at or after the last KV source layer" + elif draft_config is not None and ( + draft_window is None + or draft_window > window + or any(t != "sliding_attention" for t in draft_layer_types) + ): + reason = ( + f"the drafter (sliding window {draft_window}) reads hidden states " + f"outside the target's {window}-token window" + ) + else: + return True + logger.warning_once("Decoder SWA bounded replay is off: %s.", reason) + return False + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: stacked_params_mapping = [ # (param_name, shard_name, shard_id) @@ -1128,6 +1375,10 @@ class DeepseekV41LLMForCausalLM( engram_hash = self.model.engram_hash return engram_hash.lookback_depth if engram_hash is not None else 0 + @property + def decoder_replay_layers(self) -> DecoderReplayLayers | None: + return self.model.decoder_replay_layers + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v41/amd/vl_model.py b/vllm/models/deepseek_v41/amd/vl_model.py index 5ed045029b..5b8f7b8262 100644 --- a/vllm/models/deepseek_v41/amd/vl_model.py +++ b/vllm/models/deepseek_v41/amd/vl_model.py @@ -51,6 +51,7 @@ from ..common.mm_preprocess import ( DeepseekV4VLProcessingInfo, ) from ..common.vl_cudagraph import DeepseekV4VLEncoderCudaGraphMixin +from ..decoder_replay_layers import DecoderReplayLayers from .model import ( DeepseekV41LLMForCausalLM, _linear_scale_param_name, @@ -299,6 +300,10 @@ class DeepseekV41ForCausalLM( def token_lookback_depth(self) -> int: return self.language_model.token_lookback_depth + @property + def decoder_replay_layers(self) -> DecoderReplayLayers | None: + return self.language_model.decoder_replay_layers + def forward( self, input_ids: torch.Tensor, diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 559d874cc4..4d460d89b8 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -547,11 +547,11 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): swa_bounded_replay = False if swa_bounded_replay and current_platform.is_rocm(): logger.warning_once( - "SWA bounded replay is off on ROCm (the sparse SWA metadata " - "builders forward replay_start, but the window clamp it relies " - "on lives in the FlashInfer and FlashMLA prefill kernels, so " - "the padded slots fault); the sliding-window cache takes part " - "in prefix caching instead." + "Prefix-cache SWA replay is off on ROCm (the sparse SWA " + "metadata builders forward replay_start, but the window clamp " + "it relies on lives in the FlashInfer and FlashMLA prefill " + "kernels, so the padded slots fault); the sliding-window cache " + "takes part in prefix caching instead." ) swa_bounded_replay = False self.swa_cache_layer = DeepseekV4SWACache( diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index 7fc0bd25ff..e8bbaef64d 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -337,7 +337,7 @@ class DeepseekV41ModelState(DefaultModelState): self._warm_up_replay_kernels(input_batch, slot_mappings) window, cacheable_groups = self._replay replay_start: torch.Tensor | None = None - if window: + if window or self.decoder_replay_layers is not None: num_reqs = input_batch.num_reqs # Decode rows sit above the hit, so only prefills carry a replay # start; dummy batches (captures, profiling) carry none. @@ -349,7 +349,7 @@ class DeepseekV41ModelState(DefaultModelState): replay_start = self._replay_start_staging.copy_to_gpu( replay_start_np, out=self._replay_start[:num_reqs] ) - if replay_start_np.any(): + if window and replay_start_np.any(): # The replayed tokens rebuild window KV only: their slots in the # prefix-cacheable groups are padded so the cached KV stays as is. _pad_replayed_slots_kernel[(num_reqs,)]( @@ -397,21 +397,21 @@ class DeepseekV41ModelState(DefaultModelState): dummy batches never replay or trim, so they would not launch them.""" assert self._replay is not None window, cacheable_groups = self._replay - if not window: - return - _pad_replayed_slots_kernel.warmup( - slot_mappings, - slot_mappings.stride(0), - cacheable_groups, - input_batch.query_start_loc, - input_batch.positions, - self._replay_start, - window, - PAD_SLOT_ID, - NUM_GROUPS=cacheable_groups.numel(), - BLOCK=1024, - grid=(1,), - ) + # Pad warmup is prefix-cache SWA. A zero window must not skip the gather warmup. + if window: + _pad_replayed_slots_kernel.warmup( + slot_mappings, + slot_mappings.stride(0), + cacheable_groups, + input_batch.query_start_loc, + input_batch.positions, + self._replay_start, + window, + PAD_SLOT_ID, + NUM_GROUPS=cacheable_groups.numel(), + BLOCK=1024, + grid=(1,), + ) if self.decoder_replay_layers is None: return assert self._kept_slot_mappings is not None