diff --git a/.gitignore b/.gitignore index b664ac6a..281796a3 100644 --- a/.gitignore +++ b/.gitignore @@ -21,6 +21,7 @@ env/ build_output/ outputs/ profile_out/ +artifacts/ *.log # Local editor and agent metadata diff --git a/pypto-lib b/pypto-lib index 7e7d4cce..95206fc4 160000 --- a/pypto-lib +++ b/pypto-lib @@ -1 +1 @@ -Subproject commit 7e7d4cce8f3199f27cfefb818332af5c6b95a06e +Subproject commit 95206fc4017577c7c975b67f128478eea946f060 diff --git a/pypto_serving/model/common/executor/executor.py b/pypto_serving/model/common/executor/executor.py index ad226af2..518a7495 100644 --- a/pypto_serving/model/common/executor/executor.py +++ b/pypto_serving/model/common/executor/executor.py @@ -54,6 +54,11 @@ def supports_device_embedding(self) -> bool: """ return False + @property + def supports_device_decode_embedding(self) -> bool: + """Return whether decode kernels gather embeddings from token IDs.""" + return self.supports_device_embedding + @property def max_prefill_batch_size(self) -> int | None: """Return an executor-specific prefill dispatch limit, if any.""" diff --git a/pypto_serving/model/deepseek/npu_executor.py b/pypto_serving/model/deepseek/npu_executor.py index f52552b5..79de74c0 100644 --- a/pypto_serving/model/deepseek/npu_executor.py +++ b/pypto_serving/model/deepseek/npu_executor.py @@ -30,12 +30,14 @@ DEEPSEEK_V4_CSA_INNER_STATE_DIM, DEEPSEEK_V4_CSA_MAIN_OUT_DIM, DEEPSEEK_V4_CSA_STATE_DIM, + DEEPSEEK_V4_CACHE_GROUP_NAMES, DEEPSEEK_V4_HCA_MAIN_OUT_DIM, DEEPSEEK_V4_HCA_STATE_DIM, DEEPSEEK_V4_HC_MULT, DEEPSEEK_V4_IDX_HEAD_DIM, DEEPSEEK_V4_LM_HEAD_TP_SIZE, DEEPSEEK_V4_MAX_LOGIT_ROWS, + DEEPSEEK_V4_SAMPLED_IDS_PAD, DeepSeekV4CacheLayout, DeepSeekV4CompiledKernels, DeepSeekV4L3Callable, @@ -100,10 +102,13 @@ "decode_attention_hca", "decode_attention_swa", "decode_fwd", + "decode_input_pack", "decode_indexer", "decode_indexer_compressor", "decode_layer", + "decode_metadata_device", "decode_mtp", + "lookup_embedding", "decode_sparse_attn", "decode_sparse_attn_csa", "decode_sparse_attn_hca", @@ -327,6 +332,11 @@ def supports_device_sampling(self) -> bool: """Enable executor-provided greedy token acceptance for MTP only.""" return self._enable_mtp + @property + def supports_device_decode_embedding(self) -> bool: + """Use token IDs directly in the packed DeepSeek decode kernels.""" + return True + def lookup_embeddings(self, model: RuntimeModel, token_ids: torch.Tensor) -> torch.Tensor: """Lookup token embeddings from the lazily loaded DeepSeekV4 embedding table.""" compiled = self._compiled.get(model.config.model_id) @@ -839,7 +849,9 @@ def stacked(name: str, count: int) -> torch.Tensor: values.update( { - "x_hc": torch.empty((ranks, tokens, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.float32), + "embed_weight": torch.empty( + (ranks, model.config.vocab_size, hidden), dtype=torch.bfloat16 + ), # FWD-stacked work caches (x43). "kv_cache": torch.empty( ( @@ -912,18 +924,6 @@ def stacked(name: str, count: int) -> torch.Tensor: ), # Shared single-copy per-step inputs. "block_table": torch.empty((ranks, batch, layout.ori_table_max_blocks), dtype=torch.int32), - "ori_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "window_swa_indices": torch.empty((ranks, tokens, layout.sliding_window), dtype=torch.int32), - "window_swa_lens": torch.empty((ranks, tokens), dtype=torch.int32), - "swa_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "swa_indices": torch.empty((ranks, tokens, layout.sliding_window), dtype=torch.int32), - "swa_lens": torch.empty((ranks, tokens), dtype=torch.int32), - "hca_cmp_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "hca_state_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "csa_cmp_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "csa_idx_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "csa_state_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), - "csa_inner_state_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), "position_ids": torch.empty((ranks, tokens), dtype=torch.int32), "kv_seq_lens": torch.empty((ranks, batch), dtype=torch.int32), "hca_compress_state_block_table": torch.empty( @@ -940,6 +940,10 @@ def stacked(name: str, count: int) -> torch.Tensor: ), "cmp_block_table": torch.empty((ranks, batch, layout.cmp_max_blocks), dtype=torch.int32), "idx_block_table": torch.empty((ranks, batch, layout.idx_max_blocks), dtype=torch.int32), + "block_counts": torch.empty( + (ranks, batch, len(DEEPSEEK_V4_CACHE_GROUP_NAMES)), + dtype=torch.int32, + ), # hc_head output-collapse weights (single copy per rank). "hc_head_fn": torch.empty((ranks, DEEPSEEK_V4_HC_MULT, hc_dim), dtype=torch.float32), "hc_head_scale": torch.empty((ranks, 1), dtype=torch.float32), @@ -961,6 +965,10 @@ def stacked(name: str, count: int) -> torch.Tensor: ), "hidden_out": torch.empty((ranks, tokens, hidden), dtype=torch.bfloat16), "logits": torch.empty((ranks, tokens, model.config.vocab_size), dtype=torch.float32), + "sampled_ids": torch.empty( + (ranks, DEEPSEEK_V4_MAX_LOGIT_ROWS, DEEPSEEK_V4_SAMPLED_IDS_PAD), + dtype=torch.int32, + ), "num_tokens_per_owner": torch.full((ranks,), tokens, dtype=torch.int32), "logit_row_indices": _dummy_logit_row_indices(ranks, tokens), } diff --git a/pypto_serving/model/deepseek/npu_runner.py b/pypto_serving/model/deepseek/npu_runner.py index 2947b2f7..6b81001a 100644 --- a/pypto_serving/model/deepseek/npu_runner.py +++ b/pypto_serving/model/deepseek/npu_runner.py @@ -110,6 +110,7 @@ def _add_run_timing_args(args: dict[str, Any], timing: Any) -> None: DEEPSEEK_V4_HCA_NUM_LAYERS = 20 DEEPSEEK_V4_LM_HEAD_TP_SIZE = 4 DEEPSEEK_V4_MAX_LOGIT_ROWS = 8 +DEEPSEEK_V4_SAMPLED_IDS_PAD = 8 # Policy values indicate whether a resident argument contains mutable request # cache state and therefore must be invalidated before its Host backing is @@ -139,6 +140,7 @@ def _add_run_timing_args(args: dict[str, Any], timing: Any) -> None: _DECODE_RESIDENT_POLICY = { **_MAIN_STATIC_RESIDENT_POLICY, **_MAIN_CACHE_RESIDENT_POLICY, + "embed_weight": False, } _MTP_RESIDENT_POLICY = { "freqs_cos": False, @@ -146,6 +148,10 @@ def _add_run_timing_args(args: dict[str, Any], timing: Any) -> None: "lm_head_weight": False, "kv_cache": True, } +_MTP_DECODE_RESIDENT_POLICY = { + **_MTP_RESIDENT_POLICY, + "embed_weight": False, +} def build_deepseek_v4_cache_group_specs( @@ -409,7 +415,7 @@ def deepseek_v4_physical_cache_blocks( # ``hc_head`` collapse weights the kernel performs final RMSNorm and device # LM-head projection. _DECODE_FWD_TENSOR_ORDER = ( - "x_hc", + "embed_weight", "hc_attn_fn", "hc_attn_scale", "hc_attn_base", @@ -469,18 +475,6 @@ def deepseek_v4_physical_cache_blocks( "freqs_cos", "freqs_sin", "block_table", - "ori_slot_mapping", - "window_swa_indices", - "window_swa_lens", - "swa_slot_mapping", - "swa_indices", - "swa_lens", - "hca_cmp_slot_mapping", - "hca_state_slot_mapping", - "csa_cmp_slot_mapping", - "csa_idx_slot_mapping", - "csa_state_slot_mapping", - "csa_inner_state_slot_mapping", "position_ids", "kv_seq_lens", "hca_compress_state_block_table", @@ -488,6 +482,7 @@ def deepseek_v4_physical_cache_blocks( "csa_inner_compress_state_block_table", "cmp_block_table", "idx_block_table", + "block_counts", "input_ids", "hc_head_fn", "hc_head_scale", @@ -497,6 +492,7 @@ def deepseek_v4_physical_cache_blocks( "lm_head_weight", "hidden_out", "logits", + "sampled_ids", "num_tokens_per_owner", "logit_row_indices", ) @@ -519,12 +515,13 @@ def deepseek_v4_physical_cache_blocks( ) _MTP_DECODE_TENSOR_ORDER = ( - "hidden_states", "prev_pre_hc_hidden", "position_ids", + "embed_weight", "main_pre_hc_hidden", "tail_pre_hc_pool", + "accepted_counts", "tail_slot_ids", "position_ids", "enorm_w", "hnorm_w", "e_proj_w", "e_proj_w_scale", "e_proj_smooth", "h_proj_w", "h_proj_w_scale", "h_proj_smooth", "hc_attn_fn", "hc_attn_scale", "hc_attn_base", "attn_norm_w", "wq_a", "wq_b", "wq_b_scale", "wkv", "gamma_cq", "gamma_ckv", - "freqs_cos", "freqs_sin", "kv_cache", "swa_slot_mapping", "swa_indices", "swa_lens", + "freqs_cos", "freqs_sin", "kv_cache", "ori_block_table", "attn_sink", "wo_a", "wo_b", "wo_b_scale", "hc_ffn_fn", "hc_ffn_scale", "hc_ffn_base", "norm_w", "gate_w", "gate_bias", "tid2eid", "input_ids", @@ -532,7 +529,8 @@ def deepseek_v4_physical_cache_blocks( "routed_w2", "routed_w2_scale", "shared_w1", "shared_w1_scale", "shared_w3", "shared_w3_scale", "shared_w2", "shared_w2_scale", "mtp_hc_head_fn", "mtp_hc_head_scale", "mtp_hc_head_base", "mtp_norm_w", - "lm_head_weight", "hidden_out", "next_pre_hc_hidden", "logits", "logit_row_indices", + "lm_head_weight", "hidden_out", "next_pre_hc_hidden", "logits", "sampled_ids", + "logit_row_indices", ) _DECODE_INPUT_TENSOR_FIELDS = ( @@ -540,23 +538,12 @@ def deepseek_v4_physical_cache_blocks( "position_ids", "kv_seq_lens", "block_table", - "ori_slot_mapping", - "window_swa_indices", - "window_swa_lens", - "swa_slot_mapping", - "swa_indices", - "swa_lens", "cmp_block_table", "idx_block_table", "hca_compress_state_block_table", "csa_compress_state_block_table", "csa_inner_compress_state_block_table", - "hca_cmp_slot_mapping", - "hca_state_slot_mapping", - "csa_cmp_slot_mapping", - "csa_idx_slot_mapping", - "csa_state_slot_mapping", - "csa_inner_state_slot_mapping", + "block_counts", "num_tokens_per_owner", "logit_row_indices", ) @@ -1178,28 +1165,17 @@ class DeepSeekV4PreparedDecodeInputs: local_rows: tuple[int, ...] per_rank_counts: tuple[int, ...] actual_batch: int - x_hc: torch.Tensor + x_hc: torch.Tensor | None input_ids: torch.Tensor position_ids: torch.Tensor kv_seq_lens: torch.Tensor block_table: torch.Tensor - ori_slot_mapping: torch.Tensor - window_swa_indices: torch.Tensor - window_swa_lens: torch.Tensor - swa_slot_mapping: torch.Tensor - swa_indices: torch.Tensor - swa_lens: torch.Tensor cmp_block_table: torch.Tensor idx_block_table: torch.Tensor hca_compress_state_block_table: torch.Tensor csa_compress_state_block_table: torch.Tensor csa_inner_compress_state_block_table: torch.Tensor - hca_cmp_slot_mapping: torch.Tensor - hca_state_slot_mapping: torch.Tensor - csa_cmp_slot_mapping: torch.Tensor - csa_idx_slot_mapping: torch.Tensor - csa_state_slot_mapping: torch.Tensor - csa_inner_state_slot_mapping: torch.Tensor + block_counts: torch.Tensor block_ids_by_group: tuple[dict[str, tuple[int, ...]], ...] num_tokens_per_owner: torch.Tensor logit_row_indices: torch.Tensor @@ -1221,18 +1197,17 @@ class _DeepSeekV4MainDecodeOutput: inputs: DeepSeekV4PreparedDecodeInputs hidden: torch.Tensor - pre_hc_hidden: torch.Tensor + pre_hc_hidden: StackedDeviceTensor logits: torch.Tensor + sampled_ids: torch.Tensor @dataclass class _DeepSeekV4DecodeSharedBuffers: """Reusable decode shared-memory buffers inherited by the L3 chip workers.""" - x_hc_a: torch.Tensor - x_hc_b: torch.Tensor - pre_hc_hidden_out: torch.Tensor x_out: torch.Tensor + sampled_ids: torch.Tensor tensors: dict[str, torch.Tensor] @@ -1269,13 +1244,11 @@ class _DeepSeekV4MtpSharedBuffers: prefill_block_table: torch.Tensor prefill_slot_mapping: torch.Tensor prefill_kv_cache: torch.Tensor - decode_hidden_in: torch.Tensor - decode_prev_hidden_in: torch.Tensor decode_input_ids: torch.Tensor decode_position_ids: torch.Tensor - decode_slot_mapping: torch.Tensor - decode_swa_indices: torch.Tensor - decode_swa_lens: torch.Tensor + decode_accepted_counts: torch.Tensor + decode_tail_slot_ids: torch.Tensor + tail_init_hidden: torch.Tensor decode_kv_cache: torch.Tensor prefill_hidden_out: torch.Tensor prefill_pre_hc_out: torch.Tensor @@ -1284,6 +1257,7 @@ class _DeepSeekV4MtpSharedBuffers: decode_hidden_out: torch.Tensor decode_pre_hc_out: torch.Tensor decode_logits: torch.Tensor + decode_sampled_ids: torch.Tensor decode_logit_row_indices: torch.Tensor @@ -1308,7 +1282,8 @@ class _DeepSeekV4MtpRequestState: prefill_context: _DeepSeekV4MtpPrefillContext | None = None draft_token_id: int | None = None tail_token_id: int | None = None - tail_pre_hc_hidden: torch.Tensor | None = None + tail_rank: int | None = None + tail_slot_id: int | None = None tail_position: int | None = None proposed_tokens: int = 0 accepted_tokens: int = 0 @@ -1404,12 +1379,14 @@ def __init__(self, *, compiled: DeepSeekV4CompiledKernels) -> None: self._static_freqs_sin: torch.Tensor | None = None self._prefill_fwd_buffers: _DeepSeekV4PrefillFwdSharedBuffers | None = None self._decode_buffers: _DeepSeekV4DecodeSharedBuffers | None = None - self._decode_x_hc_initialized = False self._decode_static_metadata_keys: list[tuple[object, ...] | None] = [ None ] * compiled.layout.ranks self._stacked_host_weights: dict[str, torch.Tensor] | None = None self._stacked_device_weights: dict[str, StackedDeviceTensor] | None = None + self._embedding_device_weight: StackedDeviceTensor | None = None + self._main_pre_hc_device: StackedDeviceTensor | None = None + self._mtp_tail_pre_hc_pool: StackedDeviceTensor | None = None self._mtp_device_weights: dict[str, StackedDeviceTensor] | None = None self._hc_head_buffers: dict[str, torch.Tensor] | None = None self._prefill_output_buffer: torch.Tensor | None = None @@ -1420,6 +1397,10 @@ def __init__(self, *, compiled: DeepSeekV4CompiledKernels) -> None: self._mtp_decode_inputs_initialized = False self._mtp_device_kv_cache: StackedDeviceTensor | None = None self._mtp_request_states: dict[str, _DeepSeekV4MtpRequestState] = {} + self._mtp_free_tail_slots: list[list[int]] = [ + list(range(compiled.layout.decode_batch - 1, -1, -1)) + for _ in range(compiled.layout.ranks) + ] self._mtp_proposed_tokens = 0 self._mtp_accepted_tokens = 0 if compiled.enable_mtp: @@ -1651,7 +1632,11 @@ def release_finished_requests(self, request_ids: Iterable[str]) -> None: request_ids = tuple(request_ids) for request_id in request_ids: state = self._mtp_request_states.pop(request_id, None) - if state is not None and state.proposed_tokens: + if state is None: + continue + if state.tail_rank is not None and state.tail_slot_id is not None: + self._mtp_free_tail_slots[state.tail_rank].append(state.tail_slot_id) + if state.proposed_tokens: logger.info( "DeepSeekV4 MTP acceptance for %s: accepted=%d proposed=%d rate=%.2f%%", request_id, @@ -1672,8 +1657,17 @@ def load_packed_global_weights(self) -> DeepSeekV4GlobalWeights: loaded = self._compiled.weight_store.load_packed_global_weights( ranks=DEEPSEEK_V4_LM_HEAD_TP_SIZE ) + embed_weight = loaded.embed_weight.to( + device="cpu", + dtype=torch.bfloat16, + ).contiguous() exact_weight = loaded.lm_head_weight[:, : loaded.lm_head_layout.vocab_per_rank, :].contiguous() - self._global_weights = replace(loaded, lm_head_weight=exact_weight) + self._global_weights = replace( + loaded, + embed_weight=embed_weight, + lm_head_weight=exact_weight, + ) + self._compiled.embedding_weight = embed_weight return self._global_weights def load_stacked_layer_weights(self) -> DeepSeekV4StackedLayerWeights: @@ -1915,44 +1909,22 @@ def prepare_prefill_inputs(self, model: RuntimeModel, batch: PrefillBatch) -> De logit_row_indices=logit_row_indices, ) - @staticmethod - def _require_decode_hidden_states(batch: DecodeBatch) -> torch.Tensor: - """Return DeepSeek decode hidden states or reject a device-embedding batch.""" - hidden_states = batch.hidden_states - if hidden_states is None: - raise ValueError("DeepSeek V4 decode requires host hidden states") - return hidden_states - def prepare_decode_inputs( self, model: RuntimeModel, batch: DecodeBatch, ) -> DeepSeekV4PreparedDecodeInputs: """Build inputs for the single-token autoregressive decode flow.""" - hidden_states = self._require_decode_hidden_states(batch) assignment = self._decode_assignment(batch) if self._compiled.layout.decode_seq != 1 and max(assignment.per_rank_counts) > 1: raise ValueError( "DeepSeekV4 non-MTP decode supports at most one request per DP rank; " "the fixed S=2 kernel can expose only one cache-safe active token per rank" ) - builder = self._require_input_builder() actual_batch = len(batch.request_ids) positions = self._autoregressive_decode_positions(batch, actual_batch) - buffers = self._ensure_decode_buffers(model.config.hidden_size) - with profile_span("DeepSeekV4ModelRunner.decode.prepare_embeddings", cat="executor"): - embeddings = hidden_states.to(torch.float32).cpu() - with profile_span("DeepSeekV4ModelRunner.decode.pack_x_hc", cat="executor"): - token_rows = self._autoregressive_decode_token_rows(batch.token_ids, actual_batch) - x_hc = builder.decode_x_hc( - embeddings, - ranks=assignment.ranks, - local_rows=assignment.local_rows, - out=buffers.x_hc_a, - fill_all=not self._decode_x_hc_initialized, - ) - self._decode_x_hc_initialized = True - with profile_span("DeepSeekV4ModelRunner.decode.build_metadata", cat="executor"): + token_rows = self._autoregressive_decode_token_rows(batch.token_ids, actual_batch) + with profile_span("DeepSeekV4ModelRunner.decode.prepare_metadata_sources", cat="executor"): return self._prepare_decode_inputs( model, batch, @@ -1960,7 +1932,7 @@ def prepare_decode_inputs( active_seq=1, positions=positions, token_rows=token_rows, - x_hc=x_hc, + x_hc=None, ) def prepare_mtp_decode_inputs( @@ -1969,33 +1941,17 @@ def prepare_mtp_decode_inputs( batch: DecodeBatch, ) -> DeepSeekV4PreparedDecodeInputs: """Build paired main-model verification inputs for the MTP flow.""" - hidden_states = self._require_decode_hidden_states(batch) - if batch.prev_token_ids is None or batch.prev_hidden_states is None: - raise ValueError("DeepSeekV4 MTP decode requires previous token IDs and embeddings") + if batch.prev_token_ids is None: + raise ValueError("DeepSeekV4 MTP decode requires previous token IDs") assignment = self._decode_assignment(batch) - builder = self._require_input_builder() actual_batch = len(batch.request_ids) positions = self._mtp_decode_positions(batch, actual_batch) - buffers = self._ensure_decode_buffers(model.config.hidden_size) - with profile_span("DeepSeekV4ModelRunner.decode.prepare_embeddings", cat="executor"): - embeddings = hidden_states.to(torch.float32).cpu() - previous_embeddings = batch.prev_hidden_states.to(torch.float32).cpu() - with profile_span("DeepSeekV4ModelRunner.decode.pack_x_hc", cat="executor"): - token_rows = self._mtp_decode_token_rows( - batch.token_ids, - batch.prev_token_ids, - actual_batch, - ) - x_hc = builder.mtp_decode_x_hc( - embeddings, - prev_embeddings=previous_embeddings, - ranks=assignment.ranks, - local_rows=assignment.local_rows, - out=buffers.x_hc_a, - fill_all=not self._decode_x_hc_initialized, - ) - self._decode_x_hc_initialized = True - with profile_span("DeepSeekV4ModelRunner.decode.build_metadata", cat="executor"): + token_rows = self._mtp_decode_token_rows( + batch.token_ids, + batch.prev_token_ids, + actual_batch, + ) + with profile_span("DeepSeekV4ModelRunner.decode.prepare_metadata_sources", cat="executor"): return self._prepare_decode_inputs( model, batch, @@ -2003,7 +1959,7 @@ def prepare_mtp_decode_inputs( active_seq=self._compiled.layout.decode_seq, positions=positions, token_rows=token_rows, - x_hc=x_hc, + x_hc=None, ) def _prepare_decode_inputs( @@ -2015,7 +1971,7 @@ def _prepare_decode_inputs( active_seq: int, positions: tuple[tuple[int, ...], ...], token_rows: torch.Tensor, - x_hc: torch.Tensor, + x_hc: torch.Tensor | None, ) -> DeepSeekV4PreparedDecodeInputs: """Build mode-independent cache metadata around explicit token rows.""" layout = self._compiled.layout @@ -2033,8 +1989,6 @@ def _prepare_decode_inputs( raise ValueError(f"decode position {max_position} exceeds max_seq_len={model.runtime.max_seq_len}") buffers = self._ensure_decode_buffers(model.config.hidden_size) - with profile_span("DeepSeekV4ModelRunner.decode.stage_x_hc", cat="executor"): - self._copy_shared(buffers.x_hc_a, x_hc, name="decode_x_hc") staged = buffers.tensors staged["num_tokens_per_owner"].zero_() staged["logit_row_indices"].fill_(-1) @@ -2125,6 +2079,16 @@ def _prepare_decode_inputs( max_blocks=layout.prefill_csa_inner_state_max_blocks, ) ), + "block_counts": torch.tensor( + [ + [ + len(padded_group_ids[name][row]) + for name in DEEPSEEK_V4_CACHE_GROUP_NAMES + ] + for row in range(layout.decode_batch) + ], + dtype=torch.int32, + ), } for name, value in static_values.items(): self._copy_shared( @@ -2134,90 +2098,6 @@ def _prepare_decode_inputs( ) self._decode_static_metadata_keys[rank] = static_key - self._copy_shared( - staged["ori_slot_mapping"][rank], - self.cache_metadata.sliding_window_slot_mapping_from_ids( - padded_group_ids["ori"], - padded_positions, - ).reshape(-1), - name=f"decode_ori_slot_mapping_rank{rank}", - ) - self._copy_shared( - staged["swa_slot_mapping"][rank], - self.cache_metadata.paged_decode_slot_mapping_from_ids( - padded_group_ids["ori"], - padded_positions, - ).reshape(-1), - name=f"decode_swa_slot_mapping_rank{rank}", - ) - swa_indices, swa_lens = self.cache_metadata.swa_window_indices_and_lens_from_ids( - padded_group_ids["ori"], padded_positions - ) - for name, value in ( - ("swa_indices", swa_indices), - ("window_swa_indices", swa_indices), - ("swa_lens", swa_lens), - ("window_swa_lens", swa_lens), - ): - self._copy_shared( - staged[name][rank], - value, - name=f"decode_{name}_rank{rank}", - ) - dynamic_mappings = { - "hca_cmp_slot_mapping": ( - self.cache_metadata.compressed_slot_mapping_from_ids( - padded_group_ids["cmp"], - padded_positions, - block_size=layout.block_size, - compress_ratio=128, - ).reshape(-1) - ), - "hca_state_slot_mapping": ( - self.cache_metadata.state_slot_mapping_from_ids( - padded_group_ids["hca_state"], - padded_positions, - state_block_size=layout.c128_state_block_size, - ).reshape(-1) - ), - "csa_cmp_slot_mapping": ( - self.cache_metadata.compressed_slot_mapping_from_ids( - padded_group_ids["cmp"], - padded_positions, - block_size=layout.block_size, - compress_ratio=4, - ).reshape(-1) - ), - "csa_idx_slot_mapping": ( - self.cache_metadata.compressed_slot_mapping_from_ids( - padded_group_ids["idx"], - padded_positions, - block_size=layout.block_size, - compress_ratio=4, - ).reshape(-1) - ), - "csa_state_slot_mapping": ( - self.cache_metadata.state_slot_mapping_from_ids( - padded_group_ids["csa_state"], - padded_positions, - state_block_size=layout.c4_state_block_size, - ).reshape(-1) - ), - "csa_inner_state_slot_mapping": ( - self.cache_metadata.state_slot_mapping_from_ids( - padded_group_ids["csa_inner_state"], - padded_positions, - state_block_size=layout.c4_state_block_size, - ).reshape(-1) - ), - } - for name, value in dynamic_mappings.items(): - self._copy_shared( - staged[name][rank], - value, - name=f"decode_{name}_rank{rank}", - ) - for rank, count in enumerate(per_rank_counts): row_count = count * active_seq if row_count > DEEPSEEK_V4_MAX_LOGIT_ROWS: @@ -2238,17 +2118,11 @@ def _prepare_decode_inputs( local_rows=tuple(local_rows), per_rank_counts=per_rank_counts, actual_batch=actual_batch, - x_hc=buffers.x_hc_a, + x_hc=x_hc, input_ids=staged["input_ids"], position_ids=staged["position_ids"], kv_seq_lens=staged["kv_seq_lens"], block_table=staged["block_table"], - ori_slot_mapping=staged["ori_slot_mapping"], - window_swa_indices=staged["window_swa_indices"], - window_swa_lens=staged["window_swa_lens"], - swa_slot_mapping=staged["swa_slot_mapping"], - swa_indices=staged["swa_indices"], - swa_lens=staged["swa_lens"], cmp_block_table=staged["cmp_block_table"], idx_block_table=staged["idx_block_table"], hca_compress_state_block_table=staged["hca_compress_state_block_table"], @@ -2256,12 +2130,7 @@ def _prepare_decode_inputs( csa_inner_compress_state_block_table=staged[ "csa_inner_compress_state_block_table" ], - hca_cmp_slot_mapping=staged["hca_cmp_slot_mapping"], - hca_state_slot_mapping=staged["hca_state_slot_mapping"], - csa_cmp_slot_mapping=staged["csa_cmp_slot_mapping"], - csa_idx_slot_mapping=staged["csa_idx_slot_mapping"], - csa_state_slot_mapping=staged["csa_state_slot_mapping"], - csa_inner_state_slot_mapping=staged["csa_inner_state_slot_mapping"], + block_counts=staged["block_counts"], block_ids_by_group=active_group_ids, num_tokens_per_owner=staged["num_tokens_per_owner"], logit_row_indices=staged["logit_row_indices"], @@ -2446,14 +2315,13 @@ def _run_mtp_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResu inputs = output.inputs decode_seq = self._compiled.layout.decode_seq with profile_span("DeepSeekV4ModelRunner.decode.mtp_accept", cat="executor"): - pair_logits = torch.stack( + main_ids = torch.stack( tuple( - output.logits[rank, local_row * decode_seq + offset] + output.sampled_ids[rank, local_row * decode_seq + offset, 0] for rank, local_row in zip(inputs.ranks, inputs.local_rows, strict=True) for offset in range(decode_seq) ) - ).float() - main_ids = pair_logits.argmax(dim=-1).reshape(inputs.actual_batch, decode_seq) + ).to(torch.long).reshape(inputs.actual_batch, decode_seq) accepted = accept_mtp_tokens(main_ids, draft_token_ids) self._mtp_proposed_tokens += inputs.actual_batch self._mtp_accepted_tokens += sum(len(tokens) == decode_seq for tokens in accepted) @@ -2483,12 +2351,11 @@ def _run_mtp_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResu self._advance_mtp_drafts( inputs, main_ids, - output.pre_hc_hidden, accepted_counts=tuple(len(tokens) for tokens in accepted), ) return DecodeResult( hidden_states=None, - logits=pair_logits[::decode_seq], + logits=None, accepted_token_ids=accepted, ) @@ -2503,13 +2370,12 @@ def _execute_main_decode( with profile_span("DeepSeekV4ModelRunner.decode.prepare_inputs", cat="executor"): inputs = self._stage_decode_inputs(prepared) decode_buffers = self._require_decode_buffers() - x_hc = decode_buffers.x_hc_a active_decode_tokens = max(inputs.per_rank_counts) * active_seq - self._debug_tensor_stats("decode.input.initial.active", x_hc[:, :active_decode_tokens, :, :]) hidden_buffer = self._require_decode_output_buffer(model.config.hidden_size) - pre_hc_hidden_buffer = decode_buffers.pre_hc_hidden_out + pre_hc_hidden_buffer = self._materialize_main_pre_hc_device(model.config.hidden_size) logits_buffer = self._require_decode_logits_buffer(model.config.vocab_size) + sampled_ids_buffer = decode_buffers.sampled_ids # The active hidden/pre-HC rows are pl.Out tensors and the grouped LM # head overwrites every logits row, including rows selected with -1. # Pre-clearing these reusable buffers only adds host memory bandwidth. @@ -2521,10 +2387,10 @@ def _execute_main_decode( ): args = self._decode_fwd_args( inputs, - x_hc, pre_hc_hidden_buffer, hidden_buffer, logits_buffer, + sampled_ids_buffer, ) self._debug_decode_dispatch(inputs, args) try: @@ -2552,6 +2418,7 @@ def _execute_main_decode( hidden=hidden_buffer, pre_hc_hidden=pre_hc_hidden_buffer, logits=logits_buffer, + sampled_ids=sampled_ids_buffer, ) @staticmethod @@ -2715,10 +2582,10 @@ def _prefill_fwd_args( def _decode_fwd_args( self, inputs: DeepSeekV4PreparedDecodeInputs, - x_hc: torch.Tensor, - pre_hc_hidden_out: torch.Tensor, + pre_hc_hidden_out: StackedDeviceTensor, hidden_out: torch.Tensor, logits: torch.Tensor, + sampled_ids: torch.Tensor, ) -> tuple[Any, ...]: """Build the single packed ``l3_decode_fwd`` argument tuple.""" cache = self._materialize_decode_device_cache() @@ -2727,23 +2594,11 @@ def _decode_fwd_args( values = dict(stacked.tensors) values.update( { - "x_hc": x_hc, + "embed_weight": self._materialize_embedding_device_weight(), "freqs_cos": self._static_freqs_cos_tensor(), "freqs_sin": self._static_freqs_sin_tensor(), "kv_cache": cache.kv_cache, "block_table": inputs.block_table, - "ori_slot_mapping": inputs.ori_slot_mapping, - "window_swa_indices": inputs.window_swa_indices, - "window_swa_lens": inputs.window_swa_lens, - "swa_slot_mapping": inputs.swa_slot_mapping, - "swa_indices": inputs.swa_indices, - "swa_lens": inputs.swa_lens, - "hca_cmp_slot_mapping": inputs.hca_cmp_slot_mapping, - "hca_state_slot_mapping": inputs.hca_state_slot_mapping, - "csa_cmp_slot_mapping": inputs.csa_cmp_slot_mapping, - "csa_idx_slot_mapping": inputs.csa_idx_slot_mapping, - "csa_state_slot_mapping": inputs.csa_state_slot_mapping, - "csa_inner_state_slot_mapping": inputs.csa_inner_state_slot_mapping, "position_ids": inputs.position_ids, "kv_seq_lens": inputs.kv_seq_lens, "hca_compress_state": cache.hca_compress_state, @@ -2757,6 +2612,7 @@ def _decode_fwd_args( "idx_kv_cache": cache.idx_kv_cache, "idx_kv_scale": cache.idx_kv_scale, "idx_block_table": inputs.idx_block_table, + "block_counts": inputs.block_counts, "input_ids": inputs.input_ids, "hc_head_fn": hc_head["hc_head_fn"], "hc_head_scale": hc_head["hc_head_scale"], @@ -2766,6 +2622,7 @@ def _decode_fwd_args( "lm_head_weight": self._static_lm_head_weight_tensor(), "hidden_out": hidden_out, "logits": logits, + "sampled_ids": sampled_ids, "num_tokens_per_owner": inputs.num_tokens_per_owner, "logit_row_indices": inputs.logit_row_indices, } @@ -2821,24 +2678,30 @@ def _mtp_decode_args(self) -> tuple[Any, ...]: values = dict(self._mtp_device_weights or buffers.weights) values.update( { - "hidden_states": buffers.decode_hidden_in, - "prev_pre_hc_hidden": buffers.decode_prev_hidden_in, + "embed_weight": self._materialize_embedding_device_weight(), + "main_pre_hc_hidden": self._materialize_main_pre_hc_device( + int(self.load_packed_global_weights().embed_weight.shape[1]) + ), + "tail_pre_hc_pool": self._materialize_mtp_tail_pre_hc_pool( + int(self.load_packed_global_weights().embed_weight.shape[1]) + ), + "accepted_counts": buffers.decode_accepted_counts, + "tail_slot_ids": buffers.decode_tail_slot_ids, "position_ids": buffers.decode_position_ids, "freqs_cos": self._static_freqs_cos_tensor(), "freqs_sin": self._static_freqs_sin_tensor(), "kv_cache": kv_cache, - "swa_slot_mapping": buffers.decode_slot_mapping, - "swa_indices": buffers.decode_swa_indices, - "swa_lens": buffers.decode_swa_lens, + "ori_block_table": self._require_decode_buffers().tensors["block_table"], "input_ids": buffers.decode_input_ids, "lm_head_weight": self._static_lm_head_weight_tensor(), "hidden_out": buffers.decode_hidden_out, "next_pre_hc_hidden": buffers.decode_pre_hc_out, "logits": buffers.decode_logits, + "sampled_ids": buffers.decode_sampled_ids, "logit_row_indices": buffers.decode_logit_row_indices, } ) - values = self._mark_resident_args(values, _MTP_RESIDENT_POLICY) + values = self._mark_resident_args(values, _MTP_DECODE_RESIDENT_POLICY) return self._ordered_layer_args(values, _MTP_DECODE_TENSOR_ORDER) def _require_mtp_buffers(self) -> _DeepSeekV4MtpSharedBuffers: @@ -2872,14 +2735,12 @@ def _main_speculative_batch( actual_batch = len(batch.request_ids) draft = draft_token_ids[:actual_batch].detach().cpu().to(torch.long) current = batch.token_ids[:actual_batch].detach().cpu().to(torch.long).reshape(-1) - current_hidden = self._require_decode_hidden_states(batch)[:actual_batch].detach().cpu() - draft_hidden = self._embedding_rows(draft, current_hidden.dtype) return replace( batch, token_ids=draft.reshape(actual_batch, 1), - hidden_states=draft_hidden, + hidden_states=None, prev_token_ids=current, - prev_hidden_states=current_hidden, + prev_hidden_states=None, seq_lens=batch.seq_lens.detach().cpu().to(torch.int32) + 1, ) @@ -2961,19 +2822,55 @@ def _initialize_mtp_draft( ) state.draft_token_id = int(buffers.prefill_logits[owner_rank, 0].argmax().item()) state.tail_token_id = int(first_token[0].item()) - state.tail_pre_hc_hidden = context.prev_hidden_states[n - 1].clone() + self._initialize_mtp_tail_slot( + state, + owner_rank, + context.prev_hidden_states[n - 1], + ) state.tail_position = int(context.position_ids[n - 1].item()) state.prefill_context = None + def _initialize_mtp_tail_slot( + self, + state: _DeepSeekV4MtpRequestState, + rank: int, + hidden: torch.Tensor, + ) -> None: + """Allocate and initialize one persistent device tail slot.""" + if state.tail_slot_id is not None: + return + free_slots = self._mtp_free_tail_slots[rank] + if not free_slots: + raise RuntimeError(f"DeepSeekV4 MTP tail slots are exhausted on rank {rank}") + slot = free_slots.pop() + try: + buffers = self._require_mtp_buffers() + staged = buffers.tail_init_hidden[rank, slot] + staged.copy_(hidden.to(dtype=torch.float32, device="cpu")) + pool = self._materialize_mtp_tail_pre_hc_pool(int(staged.shape[-1])) + shard = pool.shards[rank] + row_nbytes = staged.numel() * staged.element_size() + dst = shard.data_ptr + slot * row_nbytes + self._shared_l3_worker().copy_to( + dst, + staged.data_ptr(), + row_nbytes, + worker_id=rank, + ) + except Exception: + free_slots.append(slot) + raise + state.tail_rank = rank + state.tail_slot_id = slot + def _mtp_committed_window( self, inputs: DeepSeekV4PreparedDecodeInputs, main_ids: torch.Tensor, - main_pre_hc: torch.Tensor, *, request_index: int, accepted_count: int, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: """Build one request's fixed MTP window from committed outputs only.""" layout = self._compiled.layout request_id = inputs.request_ids[request_index] @@ -2984,7 +2881,6 @@ def _mtp_committed_window( if accepted_count == layout.decode_seq: return ( main_ids[request_index, :layout.decode_seq].detach().cpu().to(torch.long), - main_pre_hc[owner_rank, row_slice].detach().cpu(), inputs.position_ids[owner_rank, row_slice].detach().cpu().to(torch.int32), ) if accepted_count != 1 or layout.decode_seq != 2: @@ -2994,7 +2890,7 @@ def _mtp_committed_window( ) if ( state.tail_token_id is None - or state.tail_pre_hc_hidden is None + or state.tail_slot_id is None or state.tail_position is None ): raise RuntimeError(f"DeepSeekV4 MTP committed tail is not initialized for {request_id!r}") @@ -3003,12 +2899,6 @@ def _mtp_committed_window( (state.tail_token_id, int(main_ids[request_index, 0].item())), dtype=torch.long, ), - torch.stack( - ( - state.tail_pre_hc_hidden, - main_pre_hc[owner_rank, row_start].detach().cpu(), - ) - ), torch.tensor( (state.tail_position, int(inputs.position_ids[owner_rank, row_start].item())), dtype=torch.int32, @@ -3019,7 +2909,6 @@ def _advance_mtp_drafts( self, inputs: DeepSeekV4PreparedDecodeInputs, main_ids: torch.Tensor, - main_pre_hc: torch.Tensor, *, accepted_counts: Sequence[int], ) -> None: @@ -3033,13 +2922,12 @@ def _advance_mtp_drafts( self._mtp_committed_window( inputs, main_ids, - main_pre_hc, request_index=index, accepted_count=int(accepted_counts[index]), ) for index in range(inputs.actual_batch) ] - active_tokens = self._stage_mtp_decode_inputs(inputs, committed) + active_tokens = self._stage_mtp_decode_inputs(inputs, committed, accepted_counts) # The MTP pl.Out tensors follow the same overwrite contract as the main # decode outputs, so no host-side clear is required before dispatch. with profile_span( @@ -3054,23 +2942,27 @@ def _advance_mtp_drafts( ) for request_index, request_id in enumerate(inputs.request_ids): state = self._require_mtp_request_state(request_id) - _, committed_hidden, committed_positions = committed[request_index] + committed_ids, committed_positions = committed[request_index] state.draft_token_id = int( - buffers.decode_logits[inputs.ranks[request_index], inputs.local_rows[request_index]].argmax().item() + buffers.decode_sampled_ids[ + inputs.ranks[request_index], + inputs.local_rows[request_index], + 0, + ].item() ) - state.tail_token_id = int(committed[request_index][0][-1].item()) - state.tail_pre_hc_hidden = committed_hidden[-1].clone() + state.tail_token_id = int(committed_ids[-1].item()) state.tail_position = int(committed_positions[-1].item()) def _stage_mtp_decode_inputs( self, inputs: DeepSeekV4PreparedDecodeInputs, - committed: Sequence[tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + committed: Sequence[tuple[torch.Tensor, torch.Tensor]], + accepted_counts: Sequence[int], ) -> int: """Stage compact committed windows into persistent MTP decode buffers.""" buffers = self._require_mtp_buffers() layout = self._compiled.layout - fallback_ids, fallback_hidden, fallback_positions = committed[0] + fallback_ids, fallback_positions = committed[0] active_tokens = max(inputs.per_rank_counts) * layout.decode_seq fill_tokens = ( layout.decode_tokens @@ -3082,88 +2974,31 @@ def _stage_mtp_decode_inputs( with profile_span("DeepSeekV4ModelRunner.mtp.pack_inputs", cat="executor"): fallback_id_rows = fallback_ids.repeat(fill_rows) fallback_position_rows = fallback_positions.repeat(fill_rows) - fallback_hidden_rows = fallback_hidden.repeat(fill_rows, 1, 1) buffers.decode_input_ids[:, :fill_tokens].copy_( fallback_id_rows.unsqueeze(0).expand(layout.ranks, -1) ) buffers.decode_position_ids[:, :fill_tokens].copy_( fallback_position_rows.unsqueeze(0).expand(layout.ranks, -1) ) - buffers.decode_prev_hidden_in[:, :fill_tokens].copy_( - fallback_hidden_rows.unsqueeze(0).expand(layout.ranks, -1, -1, -1) - ) + buffers.decode_accepted_counts.fill_(1) + buffers.decode_tail_slot_ids.fill_(-1) for request_index, (rank, local_row) in enumerate( zip(inputs.ranks, inputs.local_rows, strict=True) ): row_start = local_row * layout.decode_seq row_slice = slice(row_start, row_start + layout.decode_seq) - ids, hidden, positions = committed[request_index] + ids, positions = committed[request_index] buffers.decode_input_ids[rank, row_slice].copy_(ids) buffers.decode_position_ids[rank, row_slice].copy_(positions) - buffers.decode_prev_hidden_in[rank, row_slice].copy_(hidden) - - slot_mappings = [] - swa_indices_by_rank = [] - swa_lens_by_rank = [] - with profile_span("DeepSeekV4ModelRunner.mtp.build_metadata", cat="executor"): - for rank in range(layout.ranks): - request_indices = [ - index for index, owner_rank in enumerate(inputs.ranks) if owner_rank == rank - ] - if request_indices: - active_blocks = [ - inputs.block_ids_by_group[index]["ori"] for index in request_indices - ] - padded_blocks = self._pad_group_block_ids( - active_blocks, - group_name="ori", - kernel_rows=layout.decode_batch, - ) - padded_positions = [ - tuple(int(value) for value in committed[index][2].tolist()) - for index in request_indices - ] - while len(padded_positions) < layout.decode_batch: - padded_positions.append( - tuple(int(value) for value in fallback_positions.tolist()) - ) - else: - padded_blocks = self._scratch_group_block_ids( - group_name="ori", - kernel_rows=layout.decode_batch, - ) - padded_positions = [ - tuple(int(value) for value in fallback_positions.tolist()) - for _ in range(layout.decode_batch) - ] - slot_mappings.append( - self.cache_metadata.paged_decode_slot_mapping_from_ids( - padded_blocks, - padded_positions, - ).reshape(-1) - ) - rank_swa_indices, rank_swa_lens = ( - self.cache_metadata.swa_window_indices_and_lens_from_ids( - padded_blocks, - padded_positions, + state = self._require_mtp_request_state(inputs.request_ids[request_index]) + if state.tail_slot_id is None: + raise RuntimeError( + f"DeepSeekV4 MTP tail slot is missing for {inputs.request_ids[request_index]!r}" ) + buffers.decode_accepted_counts[rank, local_row] = int( + accepted_counts[request_index] ) - swa_indices_by_rank.append(rank_swa_indices) - swa_lens_by_rank.append(rank_swa_lens) - - buffers.decode_slot_mapping.copy_(torch.stack(slot_mappings)) - buffers.decode_swa_indices.copy_(torch.stack(swa_indices_by_rank)) - buffers.decode_swa_lens.copy_(torch.stack(swa_lens_by_rank)) - - with profile_span("DeepSeekV4ModelRunner.mtp.lookup_embeddings", cat="executor"): - staged_ids = buffers.decode_input_ids[:, :fill_tokens] - buffers.decode_hidden_in[:, :fill_tokens].copy_( - self._embedding_rows(staged_ids.reshape(-1), torch.bfloat16).reshape( - layout.ranks, - fill_tokens, - -1, - ) - ) + buffers.decode_tail_slot_ids[rank, local_row] = state.tail_slot_id buffers.decode_logit_row_indices.fill_(-1) for rank, local_row in zip(inputs.ranks, inputs.local_rows, strict=True): buffers.decode_logit_row_indices[rank, local_row] = ( @@ -3250,17 +3085,14 @@ def _debug_decode_dispatch( "x_hc", "kv_cache", "block_table", - "ori_slot_mapping", "cmp_kv", "cmp_block_table", "idx_kv_cache", "idx_block_table", + "block_counts", "hca_compress_state", - "hca_state_slot_mapping", "csa_compress_state", - "csa_state_slot_mapping", "csa_inner_compress_state", - "csa_inner_state_slot_mapping", "position_ids", "kv_seq_lens", "input_ids", @@ -3316,24 +3148,13 @@ def _is_layer_weight_name(name: str) -> bool: "kv_cache", "ori_block_table", "block_table", - "ori_slot_mapping", "cmp_kv", "cmp_block_table", "idx_kv_cache", "idx_kv_scale", "idx_block_table", + "block_counts", "position_ids", - "window_swa_indices", - "window_swa_lens", - "swa_slot_mapping", - "swa_indices", - "swa_lens", - "hca_cmp_slot_mapping", - "hca_state_slot_mapping", - "csa_cmp_slot_mapping", - "csa_idx_slot_mapping", - "csa_state_slot_mapping", - "csa_inner_state_slot_mapping", "hca_compress_state", "csa_compress_state", "csa_inner_compress_state", @@ -3352,26 +3173,16 @@ def _ensure_decode_buffers(self, hidden_size: int) -> _DeepSeekV4DecodeSharedBuf batch = layout.decode_batch tokens = layout.decode_tokens buffers = _DeepSeekV4DecodeSharedBuffers( - x_hc_a=self._shared_empty( - (ranks, tokens, layout.hc_mult, int(hidden_size)), - torch.float32, - name="decode_x_hc", - ), - x_hc_b=self._shared_empty( - (ranks, tokens, layout.hc_mult, int(hidden_size)), - torch.float32, - name="decode_x_hc_next", - ), - pre_hc_hidden_out=self._shared_empty( - (ranks, tokens, layout.hc_mult, int(hidden_size)), - torch.float32, - name="decode_pre_hc_hidden_out", - ), x_out=self._shared_empty( (ranks, tokens, int(hidden_size)), torch.bfloat16, name="decode_x_out", ), + sampled_ids=self._shared_empty( + (ranks, DEEPSEEK_V4_MAX_LOGIT_ROWS, DEEPSEEK_V4_SAMPLED_IDS_PAD), + torch.int32, + name="decode_sampled_ids", + ), tensors={ "input_ids": self._shared_empty((ranks, tokens), torch.long, name="decode_input_ids"), "position_ids": self._shared_empty((ranks, tokens), torch.int32, name="decode_position_ids"), @@ -3381,36 +3192,6 @@ def _ensure_decode_buffers(self, hidden_size: int) -> _DeepSeekV4DecodeSharedBuf torch.int32, name="decode_block_table", ), - "ori_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_ori_slot_mapping", - ), - "window_swa_indices": self._shared_empty( - (ranks, tokens, layout.sliding_window), - torch.int32, - name="decode_window_swa_indices", - ), - "window_swa_lens": self._shared_empty( - (ranks, tokens), - torch.int32, - name="decode_window_swa_lens", - ), - "swa_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_swa_slot_mapping", - ), - "swa_indices": self._shared_empty( - (ranks, tokens, layout.sliding_window), - torch.int32, - name="decode_swa_indices", - ), - "swa_lens": self._shared_empty( - (ranks, tokens), - torch.int32, - name="decode_swa_lens", - ), "cmp_block_table": self._shared_empty( (ranks, batch, layout.cmp_max_blocks), torch.int32, @@ -3436,35 +3217,10 @@ def _ensure_decode_buffers(self, hidden_size: int) -> _DeepSeekV4DecodeSharedBuf torch.int32, name="decode_csa_inner_compress_state_block_table", ), - "hca_cmp_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_hca_cmp_slot_mapping", - ), - "hca_state_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_hca_state_slot_mapping", - ), - "csa_cmp_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_csa_cmp_slot_mapping", - ), - "csa_idx_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_csa_idx_slot_mapping", - ), - "csa_state_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_csa_state_slot_mapping", - ), - "csa_inner_state_slot_mapping": self._shared_empty( - (ranks, tokens), - torch.long, - name="decode_csa_inner_state_slot_mapping", + "block_counts": self._shared_empty( + (ranks, batch, len(DEEPSEEK_V4_CACHE_GROUP_NAMES)), + torch.int32, + name="decode_block_counts", ), "num_tokens_per_owner": self._shared_empty( (ranks,), @@ -3479,7 +3235,6 @@ def _ensure_decode_buffers(self, hidden_size: int) -> _DeepSeekV4DecodeSharedBuf }, ) self._decode_buffers = buffers - self._decode_x_hc_initialized = False self._decode_static_metadata_keys = [None] * layout.ranks return buffers @@ -3527,28 +3282,22 @@ def _ensure_mtp_buffers(self, hidden_size: int) -> _DeepSeekV4MtpSharedBuffers | (ranks, layout.prefill_seq), torch.long, name="mtp_prefill_slot_mapping" ), prefill_kv_cache=mtp_kv_cache, - decode_hidden_in=self._shared_empty( - (ranks, tokens, hidden), torch.bfloat16, name="mtp_decode_hidden_in" - ), - decode_prev_hidden_in=self._shared_empty( - (ranks, tokens, layout.hc_mult, hidden), - torch.float32, - name="mtp_decode_prev_hidden_in", - ), decode_input_ids=self._shared_empty( (ranks, tokens), torch.long, name="mtp_decode_input_ids" ), decode_position_ids=self._shared_empty( (ranks, tokens), torch.int32, name="mtp_decode_position_ids" ), - decode_slot_mapping=self._shared_empty( - (ranks, tokens), torch.long, name="mtp_decode_slot_mapping" + decode_accepted_counts=self._shared_empty( + (ranks, layout.decode_batch), torch.int32, name="mtp_decode_accepted_counts" ), - decode_swa_indices=self._shared_empty( - (ranks, tokens, layout.sliding_window), torch.int32, name="mtp_decode_swa_indices" + decode_tail_slot_ids=self._shared_empty( + (ranks, layout.decode_batch), torch.int32, name="mtp_decode_tail_slot_ids" ), - decode_swa_lens=self._shared_empty( - (ranks, tokens), torch.int32, name="mtp_decode_swa_lens" + tail_init_hidden=self._shared_empty( + (ranks, layout.decode_batch, layout.hc_mult, hidden), + torch.float32, + name="mtp_tail_init_hidden", ), decode_kv_cache=mtp_kv_cache, prefill_hidden_out=self._shared_empty( @@ -3580,6 +3329,11 @@ def _ensure_mtp_buffers(self, hidden_size: int) -> _DeepSeekV4MtpSharedBuffers | torch.float32, name="mtp_decode_logits", ), + decode_sampled_ids=self._shared_empty( + (ranks, DEEPSEEK_V4_MAX_LOGIT_ROWS, DEEPSEEK_V4_SAMPLED_IDS_PAD), + torch.int32, + name="mtp_decode_sampled_ids", + ), decode_logit_row_indices=self._shared_empty( (ranks, DEEPSEEK_V4_MAX_LOGIT_ROWS), torch.int32, name="mtp_decode_logit_row_indices" ), @@ -3589,16 +3343,14 @@ def _ensure_mtp_buffers(self, hidden_size: int) -> _DeepSeekV4MtpSharedBuffers | return self._mtp_buffers def _stage_decode_inputs(self, inputs: DeepSeekV4PreparedDecodeInputs) -> DeepSeekV4PreparedDecodeInputs: - buffers = self._ensure_decode_buffers(inputs.x_hc.shape[-1]) - with profile_span("DeepSeekV4ModelRunner.decode.stage_x_hc", cat="executor"): - self._copy_shared(buffers.x_hc_a, inputs.x_hc, name="decode_x_hc") + buffers = self._require_decode_buffers() staged_values: dict[str, torch.Tensor] = {} with profile_span("DeepSeekV4ModelRunner.decode.stage_metadata", cat="executor"): for name in _DECODE_INPUT_TENSOR_FIELDS: dst = buffers.tensors[name] self._copy_shared(dst, getattr(inputs, name), name=f"decode_{name}") staged_values[name] = dst - return replace(inputs, x_hc=buffers.x_hc_a, **staged_values) + return replace(inputs, **staged_values) def _ensure_prefill_fwd_buffers(self, hidden_size: int) -> _DeepSeekV4PrefillFwdSharedBuffers: """Allocate the layer-stacked shared buffers for the packed prefill dispatch.""" @@ -4177,6 +3929,12 @@ def _upload_weight_group( def _materialize_resident_weights(self) -> None: """Upload inherited weights once and release their parent-process Host references.""" worker = self._shared_l3_worker() + if self._global_weights is not None: + hidden_size = int(self.load_packed_global_weights().embed_weight.shape[1]) + self._materialize_embedding_device_weight() + self._materialize_main_pre_hc_device(hidden_size) + if self._compiled.enable_mtp: + self._materialize_mtp_tail_pre_hc_pool(hidden_size) if self._stacked_device_weights is None: host_weights = self._stacked_host_weights if not host_weights: @@ -4243,6 +4001,9 @@ def _shared_l3_worker(self) -> Any: def _inherited_host_weights(self) -> list[torch.Tensor]: """Return immutable main and MTP weights that must be visible at worker fork.""" tensors = list(self._stacked_host_weights.values()) if self._stacked_host_weights else [] + global_weights = getattr(self, "_global_weights", None) + if global_weights is not None: + tensors.append(global_weights.embed_weight) if self._mtp_buffers is not None: tensors.extend(self._mtp_buffers.weights.values()) return tensors @@ -4271,6 +4032,56 @@ def _alloc_empty_stacked_tensor( raise return StackedDeviceTensor(shards, full_shape, worker_ids) + def _materialize_embedding_device_weight(self) -> StackedDeviceTensor: + """Upload one full embedding table to every decode rank.""" + stacked = self._embedding_device_weight + if stacked is not None: + return stacked + source = self.load_packed_global_weights().embed_weight + if source.device.type != "cpu" or source.dtype != torch.bfloat16 or not source.is_contiguous(): + raise ValueError("DeepSeekV4 embedding weight must be contiguous BF16 CPU storage before worker fork") + worker = self._shared_l3_worker() + worker_ids = tuple(range(self._compiled.layout.ranks)) + shards: list[DeviceTensor] = [] + try: + for worker_id in worker_ids: + shards.append(worker.alloc_tensor(source.shape, source.dtype, init=source, worker_id=worker_id)) + except Exception: + for shard, worker_id in zip(shards, worker_ids, strict=False): + worker.free_tensor(shard, worker_id=worker_id) + raise + stacked = StackedDeviceTensor( + shards, + (self._compiled.layout.ranks, *source.shape), + worker_ids, + ) + self._embedding_device_weight = stacked + return stacked + + def _materialize_main_pre_hc_device(self, hidden_size: int) -> StackedDeviceTensor: + """Allocate the main decode pre-HC output used directly by MTP.""" + stacked = self._main_pre_hc_device + if stacked is None: + layout = self._compiled.layout + stacked = self._alloc_empty_stacked_tensor( + (layout.ranks, layout.decode_tokens, layout.hc_mult, int(hidden_size)), + torch.float32, + ) + self._main_pre_hc_device = stacked + return stacked + + def _materialize_mtp_tail_pre_hc_pool(self, hidden_size: int) -> StackedDeviceTensor: + """Allocate persistent request tail hidden slots on every rank.""" + stacked = self._mtp_tail_pre_hc_pool + if stacked is None: + layout = self._compiled.layout + stacked = self._alloc_empty_stacked_tensor( + (layout.ranks, layout.decode_batch, layout.hc_mult, int(hidden_size)), + torch.float32, + ) + self._mtp_tail_pre_hc_pool = stacked + return stacked + def _materialize_decode_device_cache(self) -> DeepSeekV4DeviceCache: """Allocate dynamically sized cache shards directly on each NPU.""" cache = self._decode_device_cache @@ -4426,6 +4237,9 @@ def _reset_l3_worker(self) -> None: self._l3_worker = None self._l3_static_tensors.clear() self._l3_cache_tensor_keys.clear() + self._embedding_device_weight = None + self._main_pre_hc_device = None + self._mtp_tail_pre_hc_pool = None def close(self) -> None: worker = self._l3_worker @@ -4437,6 +4251,9 @@ def close(self) -> None: self._cache_group_num_blocks.clear() self._stacked_host_weights = None self._stacked_device_weights = None + self._embedding_device_weight = None + self._main_pre_hc_device = None + self._mtp_tail_pre_hc_pool = None self._mtp_device_weights = None self._mtp_buffers = None self._mtp_decode_inputs_initialized = False @@ -4444,7 +4261,6 @@ def close(self) -> None: self._decode_device_cache = None self._mtp_device_kv_cache = None self._mtp_request_states.clear() - self._decode_x_hc_initialized = False self._decode_static_metadata_keys = [None] * self._compiled.layout.ranks self._l3_static_tensors.clear() self._l3_cache_tensor_keys.clear() diff --git a/pypto_serving/serving/engine/engine.py b/pypto_serving/serving/engine/engine.py index fba09950..e3bf90a1 100644 --- a/pypto_serving/serving/engine/engine.py +++ b/pypto_serving/serving/engine/engine.py @@ -519,7 +519,7 @@ def _decode_embeddings_from_cache_or_lookup( decode_token_tensor: torch.Tensor, ) -> torch.Tensor | None: """Build decode hidden states only when the executor consumes them.""" - if self._executor.supports_device_embedding: + if self._executor.supports_device_decode_embedding: return None return self._executor.lookup_embeddings(runtime_model, decode_token_tensor) diff --git a/pypto_serving/serving/server/serving_worker.py b/pypto_serving/serving/server/serving_worker.py index 07dae884..9301f9ee 100644 --- a/pypto_serving/serving/server/serving_worker.py +++ b/pypto_serving/serving/server/serving_worker.py @@ -395,7 +395,7 @@ def _batch_decode( seq_lens = [dr.seq_len for dr in scheduled] decode_token_tensor = torch.tensor(decode_tokens, dtype=torch.long, device=device) - if self.executor.supports_device_embedding: + if self.executor.supports_device_decode_embedding: # Device kernel embeds directly from token ids — do not build # host-side embedding tensors. decode_embeddings = None @@ -405,7 +405,7 @@ def _batch_decode( prev_token_tensor = torch.tensor(prev_tokens, dtype=torch.long, device=device) prev_embeddings = self.executor.lookup_embeddings(runtime_model, prev_token_tensor) - if self.executor.supports_device_embedding: + if self.executor.supports_device_decode_embedding: prev_token_tensor = torch.tensor(prev_tokens, dtype=torch.long, device=device) decode_result = self.executor.run_decode( diff --git a/tests/test_deepseek_v4.py b/tests/test_deepseek_v4.py index 3df8457f..6c5848d5 100644 --- a/tests/test_deepseek_v4.py +++ b/tests/test_deepseek_v4.py @@ -87,8 +87,8 @@ def test_deepseek_mtp_speculative_batch_feeds_current_then_draft(): assert speculative.prev_token_ids.tolist() == [2] assert speculative.token_ids.tolist() == [[5]] assert speculative.seq_lens.tolist() == [130] - torch.testing.assert_close(speculative.prev_hidden_states, embed[2:3].to(torch.bfloat16)) - torch.testing.assert_close(speculative.hidden_states, embed[5:6].to(torch.bfloat16)) + assert speculative.prev_hidden_states is None + assert speculative.hidden_states is None def test_deepseek_mtp_committed_window_handles_rejection_and_acceptance(): @@ -103,41 +103,31 @@ def test_deepseek_mtp_committed_window_handles_rejection_and_acceptance(): ), ) main_ids = torch.tensor([[11, 12]], dtype=torch.long) - main_pre_hc = torch.arange( - runner._compiled.layout.ranks * runner._compiled.layout.decode_seq * 4, - dtype=torch.float32, - ).reshape(runner._compiled.layout.ranks, runner._compiled.layout.decode_seq, 4, 1) - tail_hidden = torch.full((4, 1), 99.0, dtype=torch.float32) runner._mtp_request_states["req-a"] = SimpleNamespace( tail_token_id=9, - tail_pre_hc_hidden=tail_hidden, + tail_slot_id=0, tail_position=6, ) - committed_ids, committed_hidden, committed_positions = runner._mtp_committed_window( + committed_ids, committed_positions = runner._mtp_committed_window( inputs, main_ids, - main_pre_hc, request_index=0, accepted_count=1, ) assert committed_ids.tolist() == [9, 11] assert committed_positions.tolist() == [6, 7] - torch.testing.assert_close(committed_hidden[0], tail_hidden) - torch.testing.assert_close(committed_hidden[1], main_pre_hc[0, 0]) - committed_ids, committed_hidden, committed_positions = runner._mtp_committed_window( + committed_ids, committed_positions = runner._mtp_committed_window( inputs, main_ids, - main_pre_hc, request_index=0, accepted_count=2, ) assert committed_ids.tolist() == [11, 12] assert committed_positions.tolist() == [7, 8] - torch.testing.assert_close(committed_hidden, main_pre_hc[0]) def test_cli_selects_deepseek_executor_and_forces_prefix_cache_off(tmp_path): @@ -317,9 +307,9 @@ def _fake_compile(self, name, jit_fn, dummy_args): assert len(compiled_args["deepseek_v4_prefill"]) == len(npu_executor._PREFILL_FWD_TENSOR_ORDER) assert len(compiled_args["deepseek_v4_decode"]) == len(npu_executor._DECODE_FWD_TENSOR_ORDER) assert compiled_args["deepseek_v4_prefill"][0].shape == (8, 128, 4, 4096) - assert compiled_args["deepseek_v4_decode"][0].shape == (8, 8, 4, 4096) + assert compiled_args["deepseek_v4_decode"][0].shape == (8, 129280, 4096) assert compiled_args["deepseek_v4_prefill"][0].dtype == torch.float32 - assert compiled_args["deepseek_v4_decode"][0].dtype == torch.float32 + assert compiled_args["deepseek_v4_decode"][0].dtype == torch.bfloat16 cache_blocks = executor._compile_cache_blocks( loaded.runtime_model, DeepSeekV4CacheLayout(decode_batch=8, decode_seq=1, decode_tokens=8), @@ -406,9 +396,9 @@ def _fake_compile(self, name, jit_fn, dummy_args): assert compiled_args["deepseek_v4_decode"][decode_order.index("lm_head_weight")].shape == (8, 32320, 4096) assert compiled_args["deepseek_v4_decode"][decode_order.index("hidden_out")].shape == (8, 8, 4096) assert compiled_args["deepseek_v4_decode"][decode_order.index("logits")].shape == (8, 8, 129280) + assert compiled_args["deepseek_v4_decode"][decode_order.index("sampled_ids")].shape == (8, 8, 8) assert compiled_args["deepseek_v4_decode"][decode_order.index("num_tokens_per_owner")].shape == (8,) assert compiled_args["deepseek_v4_decode"][decode_order.index("logit_row_indices")].shape == (8, 8) - assert compiled_args["deepseek_v4_decode"][decode_order.index("num_logit_rows")].shape == (8,) # Prefill and decode share the same runtime-sized physical pool. decode_args = compiled_args["deepseek_v4_decode"] assert decode_args[decode_order.index("kv_cache")].shape == ( @@ -419,13 +409,12 @@ def _fake_compile(self, name, jit_fn, dummy_args): assert decode_args[decode_order.index("idx_kv_scale")].shape == ( 8, 21 * cache_blocks["idx"], 128, 1, 1 ) - # SWA/HCA/CSA metadata all use the cache-first full window, plus the paged - # write slot mapping. - assert decode_args[decode_order.index("swa_slot_mapping")].shape == (8, 8) - assert decode_args[decode_order.index("swa_indices")].shape == (8, 8, 128) - assert decode_args[decode_order.index("swa_lens")].shape == (8, 8) - assert decode_args[decode_order.index("window_swa_indices")].shape == (8, 8, 128) - assert decode_args[decode_order.index("window_swa_lens")].shape == (8, 8) + # Position-dependent SWA/HCA/CSA metadata is lowered on device. The host + # ABI retains only allocator-owned tables and their active block counts. + assert decode_args[decode_order.index("block_counts")].shape == (8, 8, 6) + assert "swa_slot_mapping" not in decode_order + assert "swa_indices" not in decode_order + assert "hca_state_slot_mapping" not in decode_order def test_deepseek_kernel_contract_rejects_config_dimension_mismatch(tmp_path): @@ -779,6 +768,7 @@ def release_inherited_host_tensor_refs(self): runner = DeepSeekV4ModelRunner.__new__(DeepSeekV4ModelRunner) runner._stacked_host_weights = {"main": main_weight} runner._stacked_device_weights = None + runner._global_weights = None runner._mtp_buffers = type("MtpBuffers", (), {"weights": {"mtp": mtp_weight}})() runner._mtp_device_weights = None runner._shared_l3_worker = lambda: worker @@ -1052,22 +1042,25 @@ def test_deepseek_prepare_prefill_inputs_maps_chunk_metadata(): assert prepared.logit_row_indices[0].tolist() == [2, -1, -1, -1, -1, -1, -1, -1] -def test_deepseek_prepare_decode_inputs_requires_hidden_states(): +def test_deepseek_prepare_decode_inputs_accepts_device_embedding_batch(): runner, model = _runner_for_prepared_inputs() - with pytest.raises(ValueError, match="DeepSeek V4 decode requires host hidden states"): - runner.prepare_decode_inputs( - model, - DecodeBatch( - request_ids=["req-a"], - token_ids=torch.tensor([[5]], dtype=torch.long), - hidden_states=None, - seq_lens=torch.tensor([128], dtype=torch.int32), - ), - ) + prepared = runner.prepare_decode_inputs( + model, + DecodeBatch( + request_ids=["req-a"], + token_ids=torch.tensor([[5]], dtype=torch.long), + hidden_states=None, + seq_lens=torch.tensor([128], dtype=torch.int32), + block_ids_by_group=_grouped_cache_rows(1), + cache_partitions=[0], + ), + ) + assert prepared.x_hc is None + assert prepared.input_ids[0, :2].tolist() == [5, 5] -def test_deepseek_prepare_mtp_decode_inputs_builds_sliding_window_metadata(): +def test_deepseek_prepare_mtp_decode_inputs_stages_device_metadata_sources(): runner, model = _runner_for_prepared_inputs() prepared = runner.prepare_mtp_decode_inputs( @@ -1088,22 +1081,11 @@ def test_deepseek_prepare_mtp_decode_inputs_builds_sliding_window_metadata(): assert prepared.block_table.shape == (8, 4, 128) assert prepared.block_table[0, 0, :4].tolist() == [0, 0, 0, 0] assert prepared.block_table[1, 0, :4].tolist() == [6, 6, 6, 6] - # SWA full window includes the current position and lowers through the - # request's physical ori page. - assert prepared.swa_lens[0, 0].item() == 127 - assert prepared.swa_lens[0, 1].item() == 128 - assert prepared.swa_lens[1, 1].item() == 5 - assert prepared.swa_indices[0, 1, :4].tolist() == [0, 1, 2, 3] - assert prepared.swa_indices[1, 1, :5].tolist() == [768, 769, 770, 771, 772] - # Paged write slot for the current token. - assert prepared.swa_slot_mapping[0, 1].item() == 127 - assert prepared.swa_slot_mapping[1, 1].item() == 772 - # HCA/CSA follow pypto-lib's cache-first contract: their window contains - # historical rows and the current decode position after KV writeback. - assert prepared.window_swa_lens[0, 0].item() == 127 - assert prepared.window_swa_lens[0, 1].item() == 128 - assert prepared.window_swa_lens[1, 1].item() == 5 - assert prepared.window_swa_indices[1, 1, :5].tolist() == [768, 769, 770, 771, 772] + assert prepared.position_ids[0, :2].tolist() == [126, 127] + assert prepared.position_ids[1, :2].tolist() == [3, 4] + assert prepared.block_counts.shape == (8, 4, 6) + assert prepared.block_counts[0, 0].tolist() == [1, 1, 1, 1, 1, 1] + assert prepared.block_counts[1, 0].tolist() == [1, 1, 1, 1, 1, 1] assert prepared.num_tokens_per_owner.tolist() == [2, 2, 0, 0, 0, 0, 0, 0] assert prepared.logit_row_indices[0].tolist() == [0, 1, -1, -1, -1, -1, -1, -1] @@ -1132,13 +1114,7 @@ def test_deepseek_prepare_mtp_decode_inputs_feeds_two_real_tokens(): assert prepared.position_ids[1, :2].tolist() == [3, 4] assert prepared.kv_seq_lens[0, :4].tolist() == [128, 128, 128, 128] assert prepared.kv_seq_lens[1, :4].tolist() == [5, 5, 5, 5] - # slot 0 carries the prev-token embedding, slot 1 the last-token embedding. - assert prepared.x_hc[0, 0, 0].tolist() == [8, 9, 10, 11] - assert prepared.x_hc[0, 1, 0].tolist() == [0, 1, 2, 3] - assert prepared.x_hc[1, 0, 0].tolist() == [12, 13, 14, 15] - assert prepared.x_hc[1, 1, 0].tolist() == [4, 5, 6, 7] - # Padding row keeps replicating row 0's last embedding. - assert prepared.x_hc[0, 4, 0].tolist() == [0, 1, 2, 3] + assert prepared.x_hc is None def test_deepseek_stage_decode_inputs_uses_shared_buffers(): @@ -1157,32 +1133,19 @@ def test_deepseek_stage_decode_inputs_uses_shared_buffers(): staged = runner._stage_decode_inputs(prepared) - assert staged.x_hc.is_shared() - assert staged.x_hc.dtype == torch.float32 + assert staged.x_hc is None assert runner._decode_buffers is not None - assert runner._decode_buffers.x_hc_b.is_shared() for name in ( "input_ids", "position_ids", "kv_seq_lens", "block_table", - "ori_slot_mapping", - "window_swa_indices", - "window_swa_lens", - "swa_slot_mapping", - "swa_indices", - "swa_lens", "cmp_block_table", "idx_block_table", "hca_compress_state_block_table", "csa_compress_state_block_table", "csa_inner_compress_state_block_table", - "hca_cmp_slot_mapping", - "hca_state_slot_mapping", - "csa_cmp_slot_mapping", - "csa_idx_slot_mapping", - "csa_state_slot_mapping", - "csa_inner_state_slot_mapping", + "block_counts", ): assert getattr(staged, name).is_shared() @@ -1237,24 +1200,7 @@ def prepare(seq_len, grouped_rows): def test_deepseek_stage_mtp_decode_inputs_updates_only_active_prefix_after_first_step(): runner, model = _runner_for_prepared_inputs() layout = runner._compiled.layout - runner._compiled.embedding_weight = torch.arange( - model.config.vocab_size * model.config.hidden_size, - dtype=torch.float32, - ).reshape(model.config.vocab_size, model.config.hidden_size) runner._mtp_buffers = SimpleNamespace( - decode_hidden_in=torch.empty( - layout.ranks, - layout.decode_tokens, - model.config.hidden_size, - dtype=torch.bfloat16, - ), - decode_prev_hidden_in=torch.empty( - layout.ranks, - layout.decode_tokens, - layout.hc_mult, - model.config.hidden_size, - dtype=torch.float32, - ), decode_input_ids=torch.empty( layout.ranks, layout.decode_tokens, @@ -1265,20 +1211,14 @@ def test_deepseek_stage_mtp_decode_inputs_updates_only_active_prefix_after_first layout.decode_tokens, dtype=torch.int32, ), - decode_slot_mapping=torch.empty( - layout.ranks, - layout.decode_tokens, - dtype=torch.long, - ), - decode_swa_indices=torch.empty( + decode_accepted_counts=torch.empty( layout.ranks, - layout.decode_tokens, - layout.sliding_window, + layout.decode_batch, dtype=torch.int32, ), - decode_swa_lens=torch.empty( + decode_tail_slot_ids=torch.empty( layout.ranks, - layout.decode_tokens, + layout.decode_batch, dtype=torch.int32, ), decode_logit_row_indices=torch.empty( @@ -1300,29 +1240,20 @@ def test_deepseek_stage_mtp_decode_inputs_updates_only_active_prefix_after_first prev_hidden_states=torch.arange(4, 8, dtype=torch.bfloat16).reshape(1, 4), ), ) - committed_hidden = torch.arange( - layout.decode_seq * layout.hc_mult * model.config.hidden_size, - dtype=torch.float32, - ).reshape(layout.decode_seq, layout.hc_mult, model.config.hidden_size) committed = [ ( torch.tensor([3, 5], dtype=torch.long), - committed_hidden, torch.tensor([126, 127], dtype=torch.int32), ) ] + runner._mtp_request_states["req-a"] = SimpleNamespace(tail_slot_id=2) - assert runner._stage_mtp_decode_inputs(inputs, committed) == 2 + assert runner._stage_mtp_decode_inputs(inputs, committed, [1]) == 2 assert runner._mtp_buffers.decode_input_ids[0, :2].tolist() == [3, 5] assert runner._mtp_buffers.decode_input_ids[1, :2].tolist() == [3, 5] - torch.testing.assert_close( - runner._mtp_buffers.decode_hidden_in[0, :2], - runner._compiled.embedding_weight[torch.tensor([3, 5])].to(torch.bfloat16), - ) - - runner._mtp_buffers.decode_prev_hidden_in[:, 2:].fill_(99) - runner._stage_mtp_decode_inputs(inputs, committed) - assert torch.all(runner._mtp_buffers.decode_prev_hidden_in[:, 2:] == 99) + assert runner._mtp_buffers.decode_accepted_counts[0, 0].item() == 1 + assert runner._mtp_buffers.decode_tail_slot_ids[0, 0].item() == 2 + assert runner._mtp_buffers.decode_tail_slot_ids[1, 0].item() == -1 def test_deepseek_run_decode_dispatches_active_token_count(): @@ -1334,13 +1265,12 @@ def fake_stage(inputs): captured["prepared"] = inputs return inputs - def fake_decode_fwd_args(inputs, x_hc, pre_hc_hidden_out, hidden_out, logits): - captured["x_hc_shape"] = tuple(x_hc.shape) + def fake_decode_fwd_args(inputs, pre_hc_hidden_out, hidden_out, logits, sampled_ids): captured["num_tokens_per_owner"] = inputs.num_tokens_per_owner - return (x_hc, pre_hc_hidden_out, hidden_out, logits) + return (pre_hc_hidden_out, hidden_out, logits, sampled_ids) def fake_run_l3(_callable, *args): - args[-1].fill_(1) + args[-2].fill_(1) hidden_out = torch.empty( runner._compiled.layout.ranks, @@ -1351,15 +1281,20 @@ def fake_run_l3(_callable, *args): runner._ensure_l3_shared_buffers = lambda _model: None runner._stage_decode_inputs = fake_stage runner._require_decode_buffers = lambda: SimpleNamespace( - x_hc_a=captured["prepared"].x_hc, - pre_hc_hidden_out=torch.empty( + sampled_ids=torch.empty( runner._compiled.layout.ranks, runner._compiled.layout.decode_tokens, - runner._compiled.layout.hc_mult, - model.config.hidden_size, - dtype=torch.float32, + 8, + dtype=torch.int32, ), ) + runner._materialize_main_pre_hc_device = lambda _hidden_size: torch.empty( + runner._compiled.layout.ranks, + runner._compiled.layout.decode_tokens, + runner._compiled.layout.hc_mult, + model.config.hidden_size, + dtype=torch.float32, + ) runner._require_decode_output_buffer = lambda _hidden_size: hidden_out runner._require_decode_logits_buffer = lambda _vocab_size: torch.empty( runner._compiled.layout.ranks, @@ -1383,12 +1318,6 @@ def fake_run_l3(_callable, *args): ) assert captured["num_tokens_per_owner"].tolist() == [1, 0, 0, 0, 0, 0, 0, 0] - assert captured["x_hc_shape"] == ( - runner._compiled.layout.ranks, - runner._compiled.layout.decode_tokens, - runner._compiled.layout.hc_mult, - model.config.hidden_size, - ) assert result.logits.shape == (1, model.config.vocab_size) @@ -1512,15 +1441,17 @@ def test_deepseek_mtp_prefill_and_decode_reuse_same_kv_cache(): def test_deepseek_release_finished_requests_discards_mtp_state(): runner, _model = _runner_for_prepared_inputs() runner._mtp_request_states = { - "req-a": SimpleNamespace(proposed_tokens=0), - "req-b": SimpleNamespace(proposed_tokens=0), + "req-a": SimpleNamespace(proposed_tokens=0, tail_rank=1, tail_slot_id=2), + "req-b": SimpleNamespace(proposed_tokens=0, tail_rank=None, tail_slot_id=None), } + runner._mtp_free_tail_slots = [[] for _ in range(runner._compiled.layout.ranks)] runner.release_finished_requests(["req-a"]) assert runner._mtp_request_states == { - "req-b": SimpleNamespace(proposed_tokens=0), + "req-b": SimpleNamespace(proposed_tokens=0, tail_rank=None, tail_slot_id=None), } + assert runner._mtp_free_tail_slots[1] == [2] def _write_deepseek_model_dir(tmp_path: Path, *, quant_method: str = "compressed-tensors") -> Path: