diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 62805e90f6..54130d0d02 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -66,8 +66,10 @@ steps: - tests/models/transformers/test_backend.py - tests/models/test_registry.py - tests/models/test_deepseek_v4_vl_rocm.py + - tests/models/test_deepseek_v41_decoder_replay_layers.py + - tests/models/test_deepseek_v41_replay_batch.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 + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py mirror: amd: label: ":amd: (MI355 DPX) Basic Models (Other)" diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index ba03387aac..84dd62851e 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/tests/distributed/test_engram_dp_shard.py @@ -507,20 +507,22 @@ def _check_dummy_hash_model_forward( assert torch.all(hashes == engram_ops.DEAD_ID) output = engram.embed(hashes[:, 0]) # Trailing None is previous_aux; this stub captures no aux states. - return output, None, None, None, None, None + return output, hidden, hidden, hidden, hidden, None - model = SimpleNamespace( - use_mega_moe=False, - use_sequence_parallel=False, - fuse_mhc_all_reduce=False, - engram_hash=state, - engram_swa_prefix="swa", - engram_dp_shared_memory=dp_shared_memory, - layers=[Decoder(engram=engram)], - start_layer=0, - end_layer=1, - aux_hidden_state_layers=(), - ) + model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model) + torch.nn.Module.__init__(model) + model.use_mega_moe = False + model.use_sequence_parallel = False + model.fuse_mhc_all_reduce = False + model.engram_hash = state + model.engram_swa_prefix = "swa" + model.engram_dp_shared_memory = dp_shared_memory + model.layers = [Decoder(engram=engram)] + model.start_layer = 0 + model.end_layer = 1 + model.decoder_replay_start = model.end_layer + model.decoder_replay_layers = None + model.aux_hidden_state_layers = () metadata = ( None if dp_rank == 1 @@ -550,8 +552,7 @@ def _check_dummy_hash_model_forward( dtype=torch.int32, ), ): - output = model_ops.DeepseekV4Model.forward( - model, + output = model( torch.arange(tokens, device="cuda"), torch.arange(tokens, device="cuda"), intermediate_tensors=None, diff --git a/tests/models/test_deepseek_v41_decoder_replay_layers.py b/tests/models/test_deepseek_v41_decoder_replay_layers.py new file mode 100644 index 0000000000..4af797a1e6 --- /dev/null +++ b/tests/models/test_deepseek_v41_decoder_replay_layers.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DecoderReplayLayers runs the replay layers on the forward's replay batch.""" + +import pytest +import torch + +from vllm.forward_context import ( + ForwardContext, + get_forward_context, + override_forward_context, +) +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +DEVICE = torch.device("cuda") +NUM_TOKENS = 401 +HC, HIDDEN = 2, 3 +# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300) +WINDOW = 128 +REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)] + + +def _context(metadata): + return ForwardContext( + no_compile_layers={}, + attn_metadata=metadata, + slot_mapping={}, + is_padding=torch.zeros(NUM_TOKENS, dtype=torch.bool, device=DEVICE), + ) + + +def _set_replay_batch(layers, rows, metadata): + layers.rows = torch.tensor(rows, device=DEVICE) + layers.forward_context = ForwardContext( + no_compile_layers={}, attn_metadata=metadata, slot_mapping={} + ) + + +def _states(num_tokens: int): + """(hidden_states, positions, input_ids, pre_mix, post_mix, res_mix, residual), + shaped as the model's: hidden states [T, H], mixes [T, HC], residual [T, HC, H].""" + hidden = torch.arange(num_tokens, device=DEVICE, dtype=torch.float32) + hidden = hidden[:, None].expand(num_tokens, HIDDEN).contiguous() + mix = hidden[:, :HC].contiguous() + residual = hidden[:, None, :].expand(num_tokens, HC, HIDDEN).contiguous() + return ( + hidden, + torch.arange(num_tokens, device=DEVICE), + None, + mix + 1, + mix + 2, + mix + 3, + residual, + ) + + +def test_run_gathers_states_and_realigns_shared_indexer_buffers(): + topk = torch.arange(NUM_TOKENS * 4, device=DEVICE).view(NUM_TOKENS, 4).int() + candidates = torch.arange(NUM_TOKENS * 3, device=DEVICE).view(NUM_TOKENS, 3).int() + topk_before, candidates_before = topk.clone(), candidates.clone() + seen = {} + + def run_layers(hidden_states, positions, input_ids, pre_mix, post, res, residual): + seen["hidden_states"] = hidden_states.clone() + replay_context = get_forward_context() + seen["attn_metadata"] = replay_context.attn_metadata + seen["is_padding"] = replay_context.is_padding + return residual, pre_mix + + layers = DecoderReplayLayers(WINDOW, run_layers, [topk, candidates]) + states = _states(NUM_TOKENS) + hidden, residual = states[0], states[-1] + full, sub = object(), object() + _set_replay_batch(layers, REPLAY_ROWS, sub) + context = _context(full) + with override_forward_context(context): + outputs = layers(*states) + assert get_forward_context() is context + + rows = torch.tensor(REPLAY_ROWS, device=DEVICE) + assert torch.equal(seen["hidden_states"], hidden[rows]) + assert seen["attn_metadata"] is sub and seen["is_padding"] is None + assert torch.equal(topk[: len(REPLAY_ROWS)], topk_before[rows]) + assert torch.equal(candidates[: len(REPLAY_ROWS)], candidates_before[rows]) + assert len(outputs) == 2 and outputs[0].shape == residual.shape + assert torch.equal(outputs[0][rows], residual[rows]) + assert outputs[0][1:173].abs().sum() == 0 + + +def test_no_replay_batch_runs_the_whole_batch(): + seen = {} + + def run_layers(hidden_states, *rest): + seen["attn_metadata"] = get_forward_context().attn_metadata + return (hidden_states,) + + layers = DecoderReplayLayers(WINDOW, run_layers, []) + states = _states(NUM_TOKENS) + full = object() + with override_forward_context(_context(full)): + (output,) = layers(*states) + assert output is states[0] and seen["attn_metadata"] is full diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py new file mode 100644 index 0000000000..f6efe5c516 --- /dev/null +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The decoder-side SWA bounded replay batch DeepseekV41ModelState prepares.""" + +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.model_states.default import DefaultModelState + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="the replay buffers live on the GPU" +) + +WINDOW = 128 +DEVICE = torch.device("cuda") +GROUPS = SimpleNamespace( + kv_cache_groups=[ + SimpleNamespace( + layer_names=["swa"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=False, prefix_replay_tokens=WINDOW + ), + ), + SimpleNamespace( + layer_names=["mla"], + kv_cache_spec=SimpleNamespace( + prefix_cacheable=True, prefix_replay_tokens=0 + ), + ), + ] +) +BLOCK_TABLES = (torch.zeros(4, 4, device=DEVICE),) * 2 + +# decode (1 token), trimmed prefill (300 of 300), untrimmed prefill (100 of 300) +QUERY_LENS = [1, 300, 100] +SEQ_LENS = [500, 300, 300] +PREFILLING = [False, True, True] +REPLAY_ROWS = [0, *range(301 - WINDOW, 301), *range(301, 401)] + + +@pytest.fixture +def state(monkeypatch): + """A DeepseekV41ModelState whose attention builds are recorded, not run.""" + 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 + layers = SimpleNamespace(window=WINDOW, rows=None, forward_context=None) + model = SimpleNamespace(token_lookback_depth=0, decoder_replay_layers=layers) + builds: list = [] + + def prepare_attn(self, input_batch, cg_mode, block_tables, slot_mappings, *a, **kw): + build = SimpleNamespace( + batch=input_batch, + cg_mode=cg_mode, + slot_mappings=slot_mappings.clone(), + replay_start=kw["model_specific_attn_metadata"].replay_start, + attn_metadata=object(), + ) + builds.append(build) + return {"swa": build.attn_metadata} + + monkeypatch.setattr(DefaultModelState, "prepare_attn", prepare_attn) + state = DeepseekV41ModelState(cfg, model, None, DEVICE) + state.builds = builds # type: ignore[attr-defined] + return state + + +def _input_batch( + query_lens: list[int], + seq_lens: list[int], + is_prefilling: list[bool], + device_query_lens: list[int] | None = None, + num_tokens_after_padding: int | None = None, + max_query_len: int | None = None, +) -> InputBatch: + num_reqs, num_tokens = len(query_lens), sum(query_lens) + num_padded = num_tokens_after_padding or num_tokens + batch = InputBatch.make_dummy(num_reqs, num_tokens, InputBuffers(16, 1024, DEVICE)) + return replace( + batch, + num_tokens=num_tokens, + num_tokens_after_padding=num_padded, + query_start_loc=torch.tensor( + [0, *np.cumsum(device_query_lens or query_lens)], + dtype=torch.int32, + device=DEVICE, + ), + query_start_loc_np=np.array([0, *np.cumsum(query_lens)], dtype=np.int32), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=DEVICE), + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32), + is_prefilling_np=np.array(is_prefilling, dtype=np.bool_), + positions=torch.cat( + [torch.arange(s - q, s) for q, s in zip(query_lens, seq_lens)] + + [torch.zeros(num_padded - num_tokens, dtype=torch.int64)] + ).to(DEVICE), + is_padding=torch.arange(num_padded, device=DEVICE) >= num_tokens, + max_query_len=max_query_len, + ) + + +def _slot_mappings(num_tokens: int) -> torch.Tensor: + return torch.stack([torch.arange(num_tokens), torch.arange(num_tokens) * 10]).to( + DEVICE + ) + + +def _prepare(state, batch, cg_mode): + """Runs prepare_attn; returns the replay batch and the replay build's inputs.""" + state.prepare_attn( + batch, + cg_mode, + BLOCK_TABLES, + _slot_mappings(batch.num_tokens_after_padding), + [], + GROUPS, + ) + layers = state.decoder_replay_layers + if layers.rows is None: + assert layers.forward_context is None + return None, None + return SimpleNamespace( + rows=layers.rows, forward_context=layers.forward_context + ), state.builds[-1] + + +def test_replay_batch_keeps_each_request_window(state): + state._replay_start_np[2] = 50 # the encoder-side replay start of request 2 + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING) + replay, build = _prepare(state, batch, CUDAGraphMode.NONE) + assert replay is not None + rows = torch.tensor(REPLAY_ROWS, device=DEVICE) + assert torch.equal(replay.rows, rows) + sub = build.batch + assert build.cg_mode == CUDAGraphMode.NONE + assert sub.query_start_loc.tolist() == [0, 1, 129, 229] + assert sub.query_start_loc_np.tolist() == [0, 1, 129, 229] + assert sub.num_tokens == sub.num_tokens_after_padding == 229 + assert sub.max_query_len == WINDOW + # The trimmed request's window starts at its replay window; a higher + # encoder-side replay start stands. + assert build.replay_start.tolist() == [0, 300 - WINDOW, 50] + assert torch.equal(sub.positions, batch.positions[rows]) + assert torch.equal(build.slot_mappings, _slot_mappings(401)[:, rows]) + context = replay.forward_context + assert context.attn_metadata == {"swa": build.attn_metadata} + assert torch.equal(context.slot_mapping["mla"], _slot_mappings(401)[1, rows]) + assert context.dp_metadata is None and context.is_padding is None + + +def test_replay_batch_keeps_device_decode_boundaries(state): + """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device + [1, 3]): their rows stay and the boundaries shift only by trimmed rows.""" + batch = _input_batch( + [2, 2, 300], [10, 10, 300], PREFILLING, device_query_lens=[1, 3, 300] + ) + replay, build = _prepare(state, batch, CUDAGraphMode.NONE) + assert replay is not None + assert replay.rows.tolist() == [0, 1, 2, 3, *range(304 - WINDOW, 304)] + assert build.batch.query_start_loc.tolist() == [0, 1, 4, 4 + WINDOW] + + +def test_prompt_logprobs_requests_keep_their_rows(state): + """Every prompt row of a prompt-logprobs request is read, so it never + trims; the other prefills still do.""" + state._keeps_rows_np[1] = True + batch = _input_batch([1, 300, 300], [500, 300, 300], PREFILLING) + replay, build = _prepare(state, batch, CUDAGraphMode.NONE) + assert replay is not None + assert replay.rows.tolist() == [0, *range(1, 301), *range(601 - WINDOW, 601)] + assert build.batch.max_query_len == 300 + + +def test_replay_batch_keeps_adaptive_verification_query_bound(state): + """Adaptive verification bounds the decodes' device-side query lengths + above their CPU lengths; the replay batch keeps that bound.""" + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, max_query_len=200) + replay, build = _prepare(state, batch, CUDAGraphMode.NONE) + assert replay is not None and build.batch.max_query_len == 200 + + +def test_only_eager_steps_trim(state): + """A CUDA graph keeps the layers on its whole batch; an eager step trims a + padded batch to its real rows.""" + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING, num_tokens_after_padding=512) + for cg_mode in (CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL): + assert _prepare(state, batch, cg_mode) == (None, None) + replay, build = _prepare(state, batch, CUDAGraphMode.NONE) + assert replay is not None + assert replay.rows.tolist() == REPLAY_ROWS + assert build.cg_mode == CUDAGraphMode.NONE + assert build.batch.num_tokens_after_padding == len(REPLAY_ROWS) + + +def test_whole_batch_forwards_get_no_replay_batch(state): + short = _input_batch([100, 100], [100, 100], [True, True]) + for cg_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL): + assert _prepare(state, short, cg_mode) == (None, None) + assert len(state.builds) == 3 # the batch's own metadata only + # Dummy batches, captures included (prepared as eager), are not prefills and + # keep their rows. + dummy = _input_batch( + QUERY_LENS, SEQ_LENS, [False] * 3, num_tokens_after_padding=512 + ) + assert _prepare(state, dummy, CUDAGraphMode.NONE) == (None, None) + + +@pytest.fixture +def dp_state(state, monkeypatch): + """`state` on DP rank 0 of 2; `state.other` sets what rank 1 reports + (trims, replay tokens).""" + import vllm.models.deepseek_v41.nvidia.model_state as module + + state.vllm_config.parallel_config.data_parallel_size = 2 + state.vllm_config.parallel_config.data_parallel_rank = 0 + state.vllm_config.parallel_config.is_moe_model = True + state.other = (False, 0) # type: ignore[attr-defined] + + def all_reduce(tensor, group): + tensor[:, 1] = torch.tensor(state.other, dtype=torch.int32) + + monkeypatch.setattr(module.dist, "all_reduce", all_reduce) + monkeypatch.setattr(module, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)) + return state + + +def test_dp_ranks_replay_together(dp_state): + """A trimming rank makes the others replay too, on every rank's count.""" + short = _input_batch([100, 100], [100, 100], [True, True]) + dp_state.other = (False, 200) + assert _prepare(dp_state, short, CUDAGraphMode.NONE) == (None, None) + + dp_state.other = (True, 300) # rank 1 trims to 300 rows + replay, build = _prepare(dp_state, short, CUDAGraphMode.NONE) + assert replay is not None + assert replay.rows.tolist() == list(range(200)) # this rank keeps every row + dp_metadata = replay.forward_context.dp_metadata + assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [200, 300] + + +def test_idle_dp_rank_dummy_trims_with_its_peers(dp_state): + dummy = _input_batch([512], [512], [False]) + dp_state.other = (True, 129) + replay, _ = _prepare(dp_state, dummy, CUDAGraphMode.NONE) + assert replay is not None and replay.rows.shape[0] == WINDOW + dp_metadata = replay.forward_context.dp_metadata + assert dp_metadata.num_tokens_across_dp_cpu.tolist() == [WINDOW, 129] diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 3121f49532..eeff06933e 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -242,7 +242,9 @@ class CacheConfig: swa_bounded_replay: bool = True """Keep the sliding-window KV of models that support it (DeepSeek-V4.1) out of prefix caching and rebuild it after a prefix hit by recomputing the - hit's last window. Requires model runner V2.""" + hit's last window. The layers past the last KV-source layer then also run + eager prefill steps on each request's trailing window only. Requires model + runner V2.""" kv_cache_memory_bytes: int | None = None """Size of KV Cache per GPU in bytes. By default, this is set to None diff --git a/vllm/models/deepseek_v41/decoder_replay_layers.py b/vllm/models/deepseek_v41/decoder_replay_layers.py new file mode 100644 index 0000000000..869029aa74 --- /dev/null +++ b/vllm/models/deepseek_v41/decoder_replay_layers.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Decoder-side SWA bounded replay: running the replay layers on their batch. + +The layers past the last KV-source layer own nothing but sliding-window KV, so +in eager prefill steps they run on each request's last ``window`` rows only. +``DeepseekV41ModelState`` prepares those rows as a sub-batch with attention +metadata and a forward context of its own, like a microbatch; +``DecoderReplayLayers`` gathers the layer inputs by its rows, runs the layers +under that context and scatters the outputs back to batch rows. Steps that run +in a CUDA graph keep the layers on the whole batch, inside the graph. +""" + +from collections.abc import Callable + +import torch + +from vllm.forward_context import ForwardContext, override_forward_context +from vllm.model_executor.layers.fused_moe.moe_output import MoEOutput + + +class DecoderReplayLayers: + """Runs the replay layers on the step's replay batch. + + ``run_layers`` takes a batch's layer inputs and returns its per-row + outputs. ``row_buffers`` hold per-row results the source layer's indexer + publishes for the layers after it; they are compacted to the replay rows + in place. + """ + + def __init__( + self, + window: int, + run_layers: Callable[..., tuple[torch.Tensor, ...]], + row_buffers: list[torch.Tensor], + ) -> None: + self.window = window + self.run_layers = run_layers + self.row_buffers = row_buffers + # The replay batch, set by the model state every step: its rows of the + # batch and its forward context. None runs the layers on the batch. + self.rows: torch.Tensor | None = None + self.forward_context: ForwardContext | None = None + + def __call__( + self, hidden_states: torch.Tensor | MoEOutput, *states: torch.Tensor | None + ) -> tuple[torch.Tensor, ...]: + rows = self.rows + if rows is None: + return self.run_layers(hidden_states, *states) + # A trimming step holds a prefill longer than the window, more tokens + # than any step whose MoE leaves its finalize to the next layer. + assert isinstance(hidden_states, torch.Tensor) + num_rows = rows.shape[0] + for buf in self.row_buffers: + buf[:num_rows].copy_(buf.index_select(0, rows)) + with override_forward_context(self.forward_context): + row_outputs = self.run_layers( + hidden_states.index_select(0, rows), + *(None if t is None else t.index_select(0, rows) for t in states), + ) + # The trimmed rows' outputs stay zero; nothing reads them. + num_tokens = hidden_states.shape[0] + return tuple( + out.new_zeros((num_tokens, *out.shape[1:])).index_copy_(0, rows, out) + for out in row_outputs + ) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 9e7a35f0fc..f7647ec81d 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -83,6 +83,7 @@ from vllm.models.deepseek_v4.nvidia.model import ( prepare_mega_gate_routing_metadata, ) from vllm.models.deepseek_v41.attention import DeepseekV4Attention +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import ( DeepseekV4MegaAttnAttention, ) @@ -741,6 +742,35 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): if successor.engram is None: layer.ffn.defer_finalize() + # Decoder-side SWA bounded replay: in eager prefill steps the layers past + # the last KV source run on each request's trailing window only + # (decoder_replay_layers.py). + self.decoder_replay_layers: DecoderReplayLayers | None = None + self.decoder_replay_start = self.end_layer + cut = max(config.kv_source_layer_ids) + if ( + cut < self.end_layer - 1 + and self._decoder_replay_supported(vllm_config, cut) + and self.layers[cut].attn.swa_cache_layer.bounded_replay + ): + self.decoder_replay_start = cut + 1 + self.decoder_replay_layers = DecoderReplayLayers( + config.sliding_window, + self._run_replay_layers, + [ + buf + for buf in (self.topk_indices_buffer, self.candidate_block_buffer) + if buf is not None + ], + ) + logger.info_once( + "Decoder SWA bounded replay: in eager prefill steps, layers " + "%d-%d run on each request's last %d tokens only.", + 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 @@ -904,29 +934,28 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): hidden_states = sp_shard(hidden_states) input_ids = sp_shard(input_ids) - mega_gate_metadata = None - if self.use_mega_moe: - mega_gate_metadata = prepare_mega_gate_routing_metadata( - input_ids, - has_hash_routing=False, - image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID - if getattr(self.config, "vision_n_layers", 0) > 0 - else None, - ) - residual, post_mix, res_mix = None, None, None pre_mix: torch.Tensor | None = None if not get_pp_group().is_first_rank: assert intermediate_tensors is not None pre_mix = intermediate_tensors["pre_mix"] - # Every layer's post runs inside the next layer's fused pre, so aux - # hidden states are read back from there instead of recomputed. aux_hidden_by_layer: dict[int, torch.Tensor] = {} - 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, previous_aux = layer( + hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers( + range(self.start_layer, self.decoder_replay_start), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_by_layer, + engram_hashes, + engram_mask, + ) + late_aux: list[torch.Tensor] = [] + if self.decoder_replay_layers is not None: + hidden_states, pre_mix, *late_aux = self.decoder_replay_layers( hidden_states, positions, input_ids, @@ -934,36 +963,22 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): post_mix, res_mix, residual, - engram_hashes, - engram_mask, - capture_previous_aux=idx in self.aux_hidden_state_layers, - mega_gate_metadata=mega_gate_metadata, ) - if previous_aux is not None: - # idx is the one-based id of the layer whose post this is. - if self.use_sequence_parallel: - previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] - aux_hidden_by_layer[idx] = previous_aux - # Without a successor boundary, the last layer finalized its own MoE. - assert isinstance(hidden_states, torch.Tensor) - if layer is not None: - # The last layer has no successor to fold its post into. - if self.fuse_mhc_all_reduce: - hidden_states = tensor_model_parallel_all_reduce(hidden_states) - hidden_states = mhc_post_tilelang( - hidden_states, residual, post_mix, res_mix + else: + hidden_states = self._collapse( + hidden_states, + residual, + post_mix, + res_mix, + aux_hidden_by_layer, + full_num_tokens, ) - if self.end_layer in self.aux_hidden_state_layers: - final_aux = hidden_states.mean(dim=1) - if self.use_sequence_parallel: - final_aux = sp_all_gather(final_aux)[:full_num_tokens] - aux_hidden_by_layer[self.end_layer] = final_aux aux_hidden_states = [ aux_hidden_by_layer[layer_id] for layer_id in self.aux_hidden_state_layers if layer_id in aux_hidden_by_layer - ] + ] + late_aux if not get_pp_group().is_last_rank: return IntermediateTensors( @@ -992,6 +1007,169 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): return hidden_states, aux_hidden_states return hidden_states + def _mega_gate_metadata( + self, input_ids: torch.Tensor | None + ) -> MegaGateRoutingMetadata | None: + if not self.use_mega_moe: + return None + assert input_ids is not None + return prepare_mega_gate_routing_metadata( + input_ids, + has_hash_routing=False, + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID + if getattr(self.config, "vision_n_layers", 0) > 0 + else None, + ) + + def _run_layers( + self, + layer_ids: range, + hidden_states: torch.Tensor | MoEOutput, + 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, + aux_hidden_by_layer: dict[int, torch.Tensor], + engram_hashes: torch.Tensor | None = None, + engram_mask: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor | MoEOutput, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + ]: + # Every layer's post runs inside the next layer's fused pre, so aux + # hidden states are read back from there instead of recomputed. + full_num_tokens = positions.shape[0] + mega_gate_metadata = self._mega_gate_metadata(input_ids) + for idx in layer_ids: + hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = ( + self.layers[idx]( + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + engram_hashes, + engram_mask, + capture_previous_aux=idx in self.aux_hidden_state_layers, + mega_gate_metadata=mega_gate_metadata, + ) + ) + if previous_aux is not None: + # idx is the one-based id of the layer whose post this is. + if self.use_sequence_parallel: + previous_aux = sp_all_gather(previous_aux)[:full_num_tokens] + aux_hidden_by_layer[idx] = previous_aux + assert residual is not None and post_mix is not None + assert res_mix is not None and pre_mix is not None + return hidden_states, residual, post_mix, res_mix, pre_mix + + def _collapse( + self, + hidden_states: torch.Tensor | MoEOutput, + residual: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + aux_hidden_by_layer: dict[int, torch.Tensor], + full_num_tokens: int, + ) -> torch.Tensor: + # Without a successor boundary, the last layer finalized its own MoE. + assert isinstance(hidden_states, torch.Tensor) + # The last layer has no successor to fold its post into. + if self.fuse_mhc_all_reduce: + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) + if self.end_layer in self.aux_hidden_state_layers: + final_aux = hidden_states.mean(dim=1) + if self.use_sequence_parallel: + final_aux = sp_all_gather(final_aux)[:full_num_tokens] + aux_hidden_by_layer[self.end_layer] = final_aux + return hidden_states + + def _run_replay_layers( + self, + hidden_states: torch.Tensor | MoEOutput, + 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, ...]: + """The layers past the last KV source, on whatever rows they are given; + returns their output, the last FFN's pre-mix and the aux hidden states + they capture.""" + aux_hidden_by_layer: dict[int, torch.Tensor] = {} + hidden_states, residual, post_mix, res_mix, pre_mix = self._run_layers( + range(self.decoder_replay_start, self.end_layer), + hidden_states, + positions, + input_ids, + pre_mix, + post_mix, + res_mix, + residual, + aux_hidden_by_layer, + ) + hidden_states = self._collapse( + hidden_states, + residual, + post_mix, + res_mix, + aux_hidden_by_layer, + positions.shape[0], + ) + return ( + hidden_states, + pre_mix, + *( + aux_hidden_by_layer[layer_id] + for layer_id in self.aux_hidden_state_layers + if layer_id in aux_hidden_by_layer + ), + ) + + def _decoder_replay_supported(self, vllm_config: VllmConfig, cut: int) -> bool: + """Whether this rank may trim the layers after ``cut``; 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 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) @@ -1397,6 +1575,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/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index af20976ded..30b287db22 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -1,17 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import replace from typing import Any import numpy as np import torch +import torch.distributed as dist import torch.nn as nn from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed.parallel_state import get_dp_group +from vllm.forward_context import DPMetadata, create_forward_context +from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadataBuilder from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.core.sched.output import NewRequestData from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.dp_utils import should_skip_dp_coordination +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -81,6 +88,61 @@ def _pad_replayed_slots_kernel( ) +@triton.jit +def _gather_replay_batch_kernel( + query_start_loc_ptr, # [num_reqs + 1] + dropped_before_ptr, # [num_reqs + 1] rows trimmed before each boundary + seq_lens_ptr, # [num_reqs] + replay_start_ptr, # [num_reqs] the encoder-side replay start + positions_ptr, + slot_mappings_ptr, # [num_groups, num_tokens] + slot_mappings_stride, + rows_ptr, # out: [num_kept] batch row of every kept row + kept_query_start_loc_ptr, # out: [num_reqs + 1] + kept_positions_ptr, # out: [num_kept] + kept_slot_mappings_ptr, # out: [num_groups, num_kept] + kept_slot_mappings_stride, + kept_kv_start_ptr, # out: [num_reqs] + window, + NUM_GROUPS: tl.constexpr, + BLOCK: tl.constexpr, +): + """One program per request: its kept rows are the last ones before its + boundary, and land at that boundary minus the rows trimmed before it.""" + req = tl.program_id(0) + begin = tl.load(query_start_loc_ptr + req) + end = tl.load(query_start_loc_ptr + req + 1) + kept_begin = begin - tl.load(dropped_before_ptr + req) + kept_end = end - tl.load(dropped_before_ptr + req + 1) + if req == 0: + tl.store(kept_query_start_loc_ptr, kept_begin) + tl.store(kept_query_start_loc_ptr + req + 1, kept_end) + + # A trimmed request holds no replay-layer window KV below its kept rows. + kv_start = tl.load(replay_start_ptr + req) + if kept_end - kept_begin < end - begin: + kv_start = tl.maximum(kv_start, tl.load(seq_lens_ptr + req) - window) + tl.store(kept_kv_start_ptr + req, kv_start) + + first_row = end - (kept_end - kept_begin) + for tok in range(kept_begin, kept_end, BLOCK): + offs = tok + tl.arange(0, BLOCK) + mask = offs < kept_end + row = first_row + (offs - kept_begin) + tl.store(rows_ptr + offs, row.to(tl.int64), mask=mask) + tl.store( + kept_positions_ptr + offs, + tl.load(positions_ptr + row, mask=mask), + mask=mask, + ) + for g in tl.static_range(NUM_GROUPS): + tl.store( + kept_slot_mappings_ptr + g * kept_slot_mappings_stride + offs, + tl.load(slot_mappings_ptr + g * slot_mappings_stride + row, mask=mask), + mask=mask, + ) + + class ReplayAttnMetadata(ModelSpecificAttnMetadata): """Hands the batch's replay starts to the sliding-window builders.""" @@ -111,7 +173,10 @@ class DeepseekV41ModelState(DefaultModelState): ``prepare_attn`` gathers those per batch, pads the replayed tokens' slots in the prefix-cacheable groups so the cached KV stays as is, and hands the starts to the sliding-window metadata builders, whose kernels read no - window KV below them. + window KV below them. With the decoder side on + (``model.decoder_replay_layers``) it also prepares the replay layers' + batch in eager steps: each prefill's last ``window`` rows as a sub-batch + with metadata and a forward context of its own, like a microbatch. """ def __init__( @@ -140,9 +205,40 @@ class DeepseekV41ModelState(DefaultModelState): # the KV cache config on first use; None until then. self._replay: tuple[int, torch.Tensor] | None = None + self.decoder_replay_layers: DecoderReplayLayers | None = getattr( + model, "decoder_replay_layers", None + ) + if self.decoder_replay_layers is not None: + # Requests whose every prompt row is read (prompt logprobs) never trim. + self._keeps_rows_np = np.zeros(self.max_num_reqs, dtype=np.bool_) + self._dropped_before = UvaBufferPool(self.max_num_reqs + 1, torch.int32) + # The replay batch's inputs: its rows of the batch, query offsets, + # positions, slot mappings (sized with the KV cache groups on first + # use) and the position from which each request holds replay-layer + # window KV. + self._kept_rows = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=device + ) + self._kept_query_start_loc = torch.zeros( + self.max_num_reqs + 1, dtype=torch.int32, device=device + ) + self._kept_positions = torch.zeros( + self.max_num_tokens, dtype=torch.int64, device=device + ) + self._kept_slot_mappings: torch.Tensor | None = None + self._kept_kv_start = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + self._replay_attn_groups: list[list[AttentionGroup]] | None = None + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) self._replay_start_np[req_index] = new_req_data.replay_start + if self.decoder_replay_layers is not None: + params = new_req_data.sampling_params + self._keeps_rows_np[req_index] = ( + params is not None and params.prompt_logprobs is not None + ) def prepare_inputs( self, input_batch: InputBatch, req_states: RequestState @@ -196,7 +292,16 @@ class DeepseekV41ModelState(DefaultModelState): device=self.device, ), ) + if self.decoder_replay_layers is not None: + self._kept_slot_mappings = torch.zeros( + len(specs), + self.max_num_tokens, + dtype=torch.int64, + device=self.device, + ) + self._warm_up_replay_kernels(input_batch, slot_mappings) window, cacheable_groups = self._replay + replay_start: torch.Tensor | None = None if window: num_reqs = input_batch.num_reqs # Decode rows sit above the hit, so only prefills carry a replay @@ -226,7 +331,7 @@ class DeepseekV41ModelState(DefaultModelState): ) assert model_specific_attn_metadata is None model_specific_attn_metadata = ReplayAttnMetadata(replay_start) - return super().prepare_attn( + attn_metadata = super().prepare_attn( input_batch, cudagraph_mode, block_tables, @@ -237,3 +342,217 @@ class DeepseekV41ModelState(DefaultModelState): ubatch_idx=ubatch_idx, model_specific_attn_metadata=model_specific_attn_metadata, ) + if self.decoder_replay_layers is not None: + assert replay_start is not None + self._prepare_replay_batch( + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + replay_start, + ) + return attn_metadata + + def _warm_up_replay_kernels( + self, input_batch: InputBatch, slot_mappings: torch.Tensor + ) -> None: + """Compile the replay kernels on the buffers they run on: the startup + 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,), + ) + if self.decoder_replay_layers is None: + return + assert self._kept_slot_mappings is not None + _gather_replay_batch_kernel.warmup( + input_batch.query_start_loc, + self._dropped_before.copy_to_uva(np.zeros(1, dtype=np.int32)), + input_batch.seq_lens, + self._replay_start, + input_batch.positions, + slot_mappings, + slot_mappings.stride(0), + self._kept_rows, + self._kept_query_start_loc, + self._kept_positions, + self._kept_slot_mappings, + self._kept_slot_mappings.stride(0), + self._kept_kv_start, + self.decoder_replay_layers.window, + NUM_GROUPS=slot_mappings.shape[0], + BLOCK=1024, + grid=(1,), + ) + + def _prepare_replay_batch( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + replay_start: torch.Tensor, + ) -> None: + """Set the replay layers' batch for this forward, or none when nothing + trims. Only eager steps trim: a CUDA graph keeps the layers on the whole + batch, and ranks share the graph mode. Under data parallelism every rank + replays if any does.""" + layers = self.decoder_replay_layers + assert layers is not None + layers.rows = layers.forward_context = None + if cudagraph_mode != CUDAGraphMode.NONE: + return + + num_reqs = input_batch.num_reqs + query_lens = np.diff(input_batch.query_start_loc_np[: num_reqs + 1]) + kept_lens = np.where( + self._keeps_rows_np[input_batch.idx_mapping_np[:num_reqs]], + query_lens, + np.minimum(query_lens, layers.window), + ) + # Dummy batches (captures, an idle DP rank's) are not prefills: they + # keep their rows unless a rank trims. + trims = bool( + (kept_lens < query_lens)[input_batch.is_prefilling_np[:num_reqs]].any() + ) + num_tokens = int(kept_lens.sum()) + dp_metadata = None + if self.vllm_config.parallel_config.data_parallel_size > 1: + trims, dp_metadata = self._agree_across_dp(trims, num_tokens) + if not trims: + return + + kept_batch, kept_slot_mappings = self._kept_input_batch( + input_batch, slot_mappings, replay_start, kept_lens + ) + attn_metadata = super().prepare_attn( + kept_batch, + CUDAGraphMode.NONE, + block_tables, + kept_slot_mappings, + self._replay_groups(attn_groups), + kv_cache_config, + model_specific_attn_metadata=ReplayAttnMetadata( + self._kept_kv_start[:num_reqs] + ), + ) + layers.rows = self._kept_rows[:num_tokens] + layers.forward_context = create_forward_context( + attn_metadata, + self.vllm_config, + dp_metadata=dp_metadata, + slot_mapping=build_slot_mappings_by_layer( + kept_slot_mappings, kv_cache_config + ), + ) + + def _kept_input_batch( + self, + input_batch: InputBatch, + slot_mappings: torch.Tensor, + replay_start: torch.Tensor, + kept_lens: np.ndarray, + ) -> tuple[InputBatch, torch.Tensor]: + """Build the sub-``InputBatch`` of each request's last ``kept_lens`` + rows, and its slot mappings. + + The kept rows follow the device boundaries minus the rows trimmed before + them: adaptive verification resizes the decodes on the GPU alone, and + decodes never trim. Like a microbatch's, the sub-batch describes the + forward only; its sampling fields are carried over and must not be read. + """ + num_reqs = input_batch.num_reqs + num_tokens = int(kept_lens.sum()) + kept_query_start_loc_np = np.zeros(num_reqs + 1, dtype=np.int32) + np.cumsum(kept_lens, out=kept_query_start_loc_np[1:]) + dropped_before = self._dropped_before.copy_to_uva( + input_batch.query_start_loc_np[: num_reqs + 1] - kept_query_start_loc_np + ) + assert self._kept_slot_mappings is not None + kept_slot_mappings = self._kept_slot_mappings[:, :num_tokens] + _gather_replay_batch_kernel[(num_reqs,)]( + input_batch.query_start_loc, + dropped_before, + input_batch.seq_lens, + replay_start, + input_batch.positions, + slot_mappings, + slot_mappings.stride(0), + self._kept_rows, + self._kept_query_start_loc, + self._kept_positions, + kept_slot_mappings, + kept_slot_mappings.stride(0), + self._kept_kv_start, + self.decoder_replay_layers.window, # type: ignore[union-attr] + NUM_GROUPS=slot_mappings.shape[0], + BLOCK=1024, + ) + # Adaptive verification may lengthen the decodes on the device, up to + # the batch's bound. + max_query_len = max(int(kept_lens.max()), input_batch.max_query_len or 0) + kept_batch = replace( + input_batch, + num_tokens=num_tokens, + num_tokens_after_padding=num_tokens, + query_start_loc=self._kept_query_start_loc[: num_reqs + 1], + query_start_loc_np=kept_query_start_loc_np, + max_query_len=max_query_len, + positions=self._kept_positions[:num_tokens], + fast_prefill=None, + ) + return kept_batch, kept_slot_mappings + + def _agree_across_dp( + self, trims: bool, num_tokens: int + ) -> tuple[bool, DPMetadata | None]: + """Whether any rank trims and, then, the replay layers' DP metadata from + every rank's replay token count.""" + parallel_config = self.vllm_config.parallel_config + agreed = torch.zeros(2, parallel_config.data_parallel_size, dtype=torch.int32) + agreed[:, parallel_config.data_parallel_rank] = torch.tensor( + [trims, num_tokens] + ) + if not should_skip_dp_coordination(): + dist.all_reduce(agreed, group=get_dp_group().cpu_group) + if not agreed[0].any(): + return False, None + return True, DPMetadata.make(parallel_config, num_tokens, agreed[1]) + + def _replay_groups( + self, attn_groups: list[list[AttentionGroup]] + ) -> list[list[AttentionGroup]]: + """The attention groups with metadata builders of their own, like each + microbatch's: a builder keeps the metadata it built, and the runner's + hold the batch's.""" + if self._replay_attn_groups is None: + self._replay_attn_groups = [] + for groups in attn_groups: + replay_groups = [] + for group in groups: + replay_group = replace(group, metadata_builders=[]) + replay_group.create_metadata_builders( + self.vllm_config, + self.device, + group.metadata_builders[0].kernel_block_size, + ) + replay_groups.append(replay_group) + self._replay_attn_groups.append(replay_groups) + return self._replay_attn_groups diff --git a/vllm/models/deepseek_v41/nvidia/vl_model.py b/vllm/models/deepseek_v41/nvidia/vl_model.py index 901985bdd7..d67753405a 100644 --- a/vllm/models/deepseek_v41/nvidia/vl_model.py +++ b/vllm/models/deepseek_v41/nvidia/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, @@ -281,6 +282,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,