Lower DeepSeek decode inputs to device

.gitignore

@@ -21,6 +21,7 @@ env/
 build_output/
 outputs/
 profile_out/
+artifacts/
 *.log
 
 # Local editor and agent metadata

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."""

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),
             }