diff --git a/docs/dev/model/deepseek-v4.md b/docs/dev/model/deepseek-v4.md new file mode 100644 index 0000000..ee4dc35 --- /dev/null +++ b/docs/dev/model/deepseek-v4.md @@ -0,0 +1,28 @@ +# DeepSeek V4 NPU Serving Dev Notes + +These commands are for DeepSeek V4 Flash W8A8 serving checks on shared Ascend +development machines with `task-submit`. Run them from the pypto-serving +checkout. + +## 8-Device TP Serving + +Use the quantized checkpoint under `/data/models/dsv4-flash-w8a8` and run with +TP=8 on devices 8-15: + +```bash +task-submit --device 8,9,10,11,12,13,14,15 --max-time 0 --timeout 0 --ptoas 0.46 --run "PYPTO_RUNTIME_LOG=error PTO2_RING_DEP_POOL=131072 PTO2_RING_TASK_WINDOW=131072 PTO2_RING_HEAP=2147483648 PTO2_OP_EXECUTE_TIMEOUT_US=400000000 PTO2_STREAM_SYNC_TIMEOUT_MS=440000 PTO2_SCHEDULER_TIMEOUT_MS=320000 SERVING_WORKER_STEP_TIMEOUT=1800 python python/cli/main.py --model /data/models/dsv4-flash-w8a8 --served-model-name dsv4-flash-w8a8 --backend npu --platform a2a3 --devices 8,9,10,11,12,13,14,15 --dp 1 --tp 8 --block-size 128 --max-model-len 260 --max-num-seqs 4 --max-num-batched-tokens 512 --long-prefill-token-threshold 2048 --no-enable-prefix-caching --port 8225 --show-startup-logs" +``` + +## Completion Check + +Check server health first: + +```bash +curl --noproxy "*" http://127.0.0.1:8225/health +``` + +Then send a deterministic completion request: + +```bash +curl --noproxy "*" -s http://127.0.0.1:8225/v1/completions -H "Content-Type: application/json" -d '{"model":"dsv4-flash-w8a8","prompt":"Huawei is","max_tokens":25,"temperature":0.0}' +``` diff --git a/examples/model/deepseek_v4/__init__.py b/examples/model/deepseek_v4/__init__.py new file mode 100644 index 0000000..29b1abc --- /dev/null +++ b/examples/model/deepseek_v4/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""DeepSeekV4 serving integration package.""" diff --git a/examples/model/deepseek_v4/runner/__init__.py b/examples/model/deepseek_v4/runner/__init__.py new file mode 100644 index 0000000..9202a54 --- /dev/null +++ b/examples/model/deepseek_v4/runner/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""DeepSeekV4 runner components.""" diff --git a/examples/model/deepseek_v4/runner/npu_executor.py b/examples/model/deepseek_v4/runner/npu_executor.py new file mode 100644 index 0000000..3edf8f0 --- /dev/null +++ b/examples/model/deepseek_v4/runner/npu_executor.py @@ -0,0 +1,948 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import ast +import contextlib +import importlib +import importlib.util +import operator +import os +import sys +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import Any + +import torch + +from examples.model.deepseek_v4.runner.npu_runner import ( + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + DEEPSEEK_V4_CSA_INNER_STATE_DIM, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + DEEPSEEK_V4_CSA_STATE_DIM, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + DEEPSEEK_V4_HCA_STATE_DIM, + DEEPSEEK_V4_HC_MULT, + DEEPSEEK_V4_IDX_HEAD_DIM, + DeepSeekV4CacheLayout, + DeepSeekV4CompiledKernels, + DeepSeekV4L3Callable, + DeepSeekV4ModelRunner, + _DECODE_FWD_TENSOR_ORDER, + _PREFILL_FWD_TENSOR_ORDER, + build_deepseek_v4_layer_plan, + DEEPSEEK_V4_CSA_NUM_LAYERS, + DEEPSEEK_V4_FWD_NUM_LAYERS, + DEEPSEEK_V4_HCA_NUM_LAYERS, +) +from examples.model.deepseek_v4.runner.weight_loader import DeepSeekV4WeightStore +from python.core.model_runner import ModelRunner +from python.core.pypto_executor import PyptoExecutor as CorePyptoExecutor +from python.core.types import RuntimeModel + + +_AST_INT_OPERATORS = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.FloorDiv: operator.floordiv, +} +# CSA-group (x21) and HCA-group (x20) layer-stacked weight names emitted by the +# per-layer common dummy builder. Everything else there is a FWD weight (x43). +# Shared single-copy inputs (freqs/input_ids) are handled explicitly, not stacked. +_DECODE_FWD_CSA_STACKED_NAMES = frozenset( + { + "csa_cmp_wkv", + "csa_cmp_wgate", + "csa_cmp_ape", + "csa_cmp_norm_w", + "csa_idx_wq_b", + "csa_idx_wq_b_scale", + "csa_weights_proj", + "csa_hadamard_idx", + "csa_inner_wkv", + "csa_inner_wgate", + "csa_inner_ape", + "csa_inner_norm_w", + } +) +_DECODE_FWD_HCA_STACKED_NAMES = frozenset( + { + "hca_cmp_wkv", + "hca_cmp_wgate", + "hca_cmp_ape", + "hca_cmp_norm_w", + } +) +_DECODE_FWD_SHARED_COMMON_NAMES = frozenset({"freqs_cos", "freqs_sin", "input_ids"}) +# Packed prefill now mirrors decode: the RoPE tables and input ids are passed as a +# single per-rank copy (the kernel slices them per layer internally), not stacked +# across the 43 forward layers. +_PREFILL_FWD_SHARED_COMMON_NAMES = frozenset({"freqs_cos", "freqs_sin", "input_ids"}) +_DEEPSEEK_V4_IMPORT_MODULES = ( + "config", + "moe", + "combine", + "decode_attention_csa", + "decode_attention_hca", + "decode_attention_swa", + "decode_fwd", + "decode_indexer", + "decode_indexer_compressor", + "decode_layer", + "decode_sparse_attn", + "decode_sparse_attn_csa", + "decode_sparse_attn_hca", + "decode_sparse_attn_swa", + "dispatch", + "expert_routed", + "expert_shared", + "gate", + "hc_post", + "hc_pre", + "prefill_attention_csa", + "prefill_attention_hca", + "prefill_attention_swa", + "prefill_indexer_compressor", + "prefill_layer", + "prefill_fwd", + "prefill_sparse_attn", + "qkv_proj_rope", + "rmsnorm", + "rope_tables", +) + + +def _find_pypto_lib_deepseek_v4_dir(pypto_root: str | None = None) -> Path: + """Find the DeepSeekV4 kernel directory.""" + if pypto_root is None: + pypto_root = os.environ.get("PYPTO_ROOT") + if pypto_root: + root = Path(pypto_root) + candidate = root / "models" / "deepseek" / "v4" + if candidate.is_dir(): + return candidate + raise FileNotFoundError(f"DeepSeekV4 kernel directory not found under PYPTO_ROOT={pypto_root!r}") + + start_dir = Path(__file__).resolve().parent + for directory in (start_dir, *start_dir.parents): + pypto_lib_dir = directory / "pypto-lib" + candidate = pypto_lib_dir / "models" / "deepseek" / "v4" + if candidate.is_dir(): + return candidate + + raise FileNotFoundError( + "Cannot locate DeepSeekV4 kernels. Run from a checkout with pypto-lib available " + "or set PYPTO_ROOT to a pypto-lib checkout." + ) + + +def _int_constant_from_file(path: Path, name: str) -> int | None: + """Read a simple integer module constant without importing kernel code.""" + tree = ast.parse(path.read_text(), filename=str(path)) + assignments = { + target.id: node.value + for node in tree.body + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) + } + config_assignments = None + + def _eval_int(node: ast.AST) -> int | None: + nonlocal config_assignments + if isinstance(node, ast.Constant) and isinstance(node.value, int): + return int(node.value) + if isinstance(node, ast.Name): + if node.id in assignments: + return _eval_int(assignments[node.id]) + if config_assignments is None: + config_path = path.parent / "config.py" + if config_path == path or not config_path.exists(): + config_assignments = {} + else: + config_tree = ast.parse(config_path.read_text(), filename=str(config_path)) + config_assignments = { + target.id: cfg_node.value + for cfg_node in config_tree.body + if isinstance(cfg_node, ast.Assign) + for target in cfg_node.targets + if isinstance(target, ast.Name) + } + config_node = config_assignments.get(node.id) + return _eval_int(config_node) if config_node is not None else None + if isinstance(node, ast.BinOp): + left = _eval_int(node.left) + right = _eval_int(node.right) + op = _AST_INT_OPERATORS.get(type(node.op)) + if left is None or right is None or op is None: + return None + return int(op(left, right)) + return None + + for node in tree.body: + if not isinstance(node, ast.Assign): + continue + if not any(isinstance(target, ast.Name) and target.id == name for target in node.targets): + continue + return _eval_int(node.value) + return None + + +def _is_deepseek_v4_module_file(path: Path, kernel_dir: Path) -> bool: + """Return whether ``path`` is one of the top-level DeepSeekV4 kernel modules.""" + resolved = path.resolve() + if resolved.is_relative_to(kernel_dir): + return True + parts = resolved.parts + return len(parts) >= 4 and parts[-4:-1] == ("models", "deepseek", "v4") + + +@contextlib.contextmanager +def _deepseek_v4_import_context( + kernel_dir: Path, + *, + pypto_root: Path, + ep: int, + moe_shape: str | None = None, + num_layers: int | None = None, +): + """Temporarily import DeepSeekV4 pypto-lib modules with a fixed EP argv.""" + old_argv = list(sys.argv) + old_path = list(sys.path) + missing = object() + old_modules = { + module_name: sys.modules.get(module_name, missing) + for module_name in _DEEPSEEK_V4_IMPORT_MODULES + } + for module_name in _DEEPSEEK_V4_IMPORT_MODULES: + module = sys.modules.get(module_name) + module_file = getattr(module, "__file__", None) + if module_file is not None and _is_deepseek_v4_module_file(Path(module_file), kernel_dir): + sys.modules.pop(module_name, None) + sys.argv = ["pypto-serving-deepseek-v4", "--ep", str(int(ep))] + if moe_shape is not None: + sys.argv.extend(["--moe-shape", moe_shape]) + if num_layers is not None: + # prefill_fwd freezes its layer-stack span from ``--num-layers`` at import; + # serving always packs the full 43-layer forward. + sys.argv.extend(["--num-layers", str(int(num_layers))]) + sys.path.insert(0, str(kernel_dir)) + sys.path.insert(0, str(pypto_root)) + try: + yield + finally: + sys.argv = old_argv + sys.path[:] = old_path + for module_name, module in old_modules.items(): + if module is missing: + sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = module + + +class DeepSeekV4PyptoExecutor(CorePyptoExecutor): + """PyPTO executor boundary for DeepSeekV4 Flash W8A8 serving.""" + + def __init__( + self, + kv_cache_manager=None, + *, + platform: str = "a2a3sim", + device_id: int = 0, + device_ids: Sequence[int] | None = None, + save_kernels_dir: str | None = None, + pypto_root: str | None = None, + compile_kernels: bool = False, + l3_trace: bool = False, + ) -> None: + worker_device_ids = tuple(device_ids) if device_ids is not None else (int(device_id),) + super().__init__( + kv_cache_manager, + platform=platform, + device_ids=worker_device_ids, + save_kernels_dir=save_kernels_dir, + ) + self._pypto_root = pypto_root + self._kernel_dir = _find_pypto_lib_deepseek_v4_dir(pypto_root) + self._compile_kernels = bool(compile_kernels) + self._l3_trace = l3_trace + self._embedding_cache: dict[str, torch.Tensor] = {} + + @property + def profile_verbose(self) -> bool: + """Return whether compile and L3 execution timing logs are enabled.""" + return self._l3_trace + + 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) + if not isinstance(compiled, DeepSeekV4CompiledKernels): + raise RuntimeError(f"DeepSeekV4 model {model.config.model_id!r} is not registered") + embed_weight = self._embedding_cache.get(model.config.model_id) + if embed_weight is None: + embed_weight = compiled.weight_store.load_tensor("embed.weight").contiguous() + if embed_weight.ndim != 2: + raise ValueError(f"embed.weight must be rank-2, got shape={tuple(embed_weight.shape)}") + if int(embed_weight.shape[0]) != model.config.vocab_size: + raise ValueError( + f"embed.weight vocab size must be {model.config.vocab_size}, " + f"got {int(embed_weight.shape[0])}" + ) + if int(embed_weight.shape[1]) != model.config.hidden_size: + raise ValueError( + f"embed.weight hidden size must be {model.config.hidden_size}, " + f"got {int(embed_weight.shape[1])}" + ) + self._embedding_cache[model.config.model_id] = embed_weight + + flat_ids = token_ids.detach().to(device="cpu", dtype=torch.long).reshape(-1) + embeddings = embed_weight.index_select(0, flat_ids) + return embeddings.reshape(*token_ids.shape, model.config.hidden_size).to(device=token_ids.device) + + def release_finished_requests(self, request_ids: Iterable[str]) -> None: + """Release runner-owned DeepSeekV4 cache slots for finished requests.""" + for runner in self._runners.values(): + release = getattr(runner, "release_finished_requests", None) + if callable(release): + release(request_ids) + + def _create_runner(self, model_id: str, compiled: object) -> ModelRunner: + """Create the DeepSeekV4 runtime runner.""" + if not isinstance(compiled, DeepSeekV4CompiledKernels): + raise TypeError("DeepSeekV4PyptoExecutor requires DeepSeekV4 compiled metadata.") + return DeepSeekV4ModelRunner(compiled=compiled) + + def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels: + """Validate DeepSeekV4 W8A8 metadata and return runner artifacts. + + The current pypto-lib DeepSeekV4 programs are single-layer kernels. This + method intentionally validates and packages the serving contract without + pretending those kernels are already a full-model generator. + """ + metadata = model.extra + if metadata.get("family") != "deepseek_v4": + raise ValueError("DeepSeekV4PyptoExecutor received a non-DeepSeekV4 model") + if metadata.get("checkpoint_format") != "w8a8-compressed-tensors": + raise ValueError("DeepSeekV4PyptoExecutor requires the W8A8 compressed-tensors checkpoint") + + layout = DeepSeekV4CacheLayout() + layout.validate_runtime(model.config, model.runtime, self._device_ids) + self._validate_kernel_contract(layout) + compress_ratios = tuple(int(ratio) for ratio in metadata["compress_ratios"]) + if len(compress_ratios) != model.config.num_hidden_layers + 1: + raise ValueError("DeepSeekV4 compress_ratios must include hidden layers plus MTP/final entry") + config_data = metadata.get("config_data", {}) + n_routed_experts = int(config_data.get("n_routed_experts", 256)) if isinstance(config_data, dict) else 256 + num_hash_layers = int(config_data.get("num_hash_layers", 3)) if isinstance(config_data, dict) else 3 + layer_plan = build_deepseek_v4_layer_plan( + compress_ratios=compress_ratios, + num_hidden_layers=model.config.num_hidden_layers, + num_hash_layers=num_hash_layers, + ) + weight_map = dict(metadata["weight_map"]) + weight_store = DeepSeekV4WeightStore(model_dir=str(metadata["model_dir"]), weight_map=weight_map) + weight_store.validate_startup_contract( + num_hidden_layers=model.config.num_hidden_layers, + n_routed_experts=n_routed_experts, + compress_ratios=compress_ratios, + num_hash_layers=num_hash_layers, + ) + + prefill = None + decode = None + freqs_cos = freqs_sin = None + if self._compile_kernels: + modules = self._load_kernel_modules(layout) + prefill = self._compile_l3_callable( + "deepseek_v4_prefill", + modules["prefill_fwd"].l3_prefill_fwd, + self._prefill_dummy_args(model, layout, modules["config"]), + ) + decode = self._compile_l3_callable( + "deepseek_v4_decode", + modules["decode_fwd"].l3_decode_fwd, + self._decode_dummy_args(model, layout, modules["config"]), + ) + freqs_cos, freqs_sin = self._build_rope_tables(modules["rope_tables"], modules["config"]) + + return DeepSeekV4CompiledKernels( + layout=layout, + model_dir=str(metadata["model_dir"]), + weight_map=weight_map, + weight_store=weight_store, + compress_ratios=compress_ratios, + layer_plan=layer_plan, + kernel_dir=str(self._kernel_dir), + prefill=prefill, + decode=decode, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + platform=self._platform, + device_id=self._device_ids[0], + n_routed_experts=n_routed_experts, + num_hash_layers=num_hash_layers, + ) + + def _load_kernel_modules(self, layout: DeepSeekV4CacheLayout) -> dict[str, object]: + """Import DeepSeekV4 pypto-lib modules with EP fixed to the serving world size.""" + pypto_root = self._kernel_dir.parents[2] + ranks = layout.ranks + fwd_layers = DEEPSEEK_V4_FWD_NUM_LAYERS + with _deepseek_v4_import_context( + self._kernel_dir, + pypto_root=pypto_root, + ep=ranks, + moe_shape="prefill", + num_layers=fwd_layers, + ): + prefill_layer = importlib.import_module("prefill_layer") + prefill_fwd = importlib.import_module("prefill_fwd") + with _deepseek_v4_import_context(self._kernel_dir, pypto_root=pypto_root, ep=ranks, moe_shape="decode"): + modules = { + name: importlib.import_module(name) + for name in ("config", "decode_layer", "decode_fwd", "rope_tables") + } + modules["prefill_layer"] = prefill_layer + modules["prefill_fwd"] = prefill_fwd + return modules + + def _compile_l3_callable(self, name: str, jit_fn: object, dummy_args: Sequence[Any]) -> DeepSeekV4L3Callable: + """Compile one DeepSeekV4 HOST wrapper into a distributed program.""" + from pypto.ir.distributed_compiled_program import DistributedCompiledProgram # noqa: PLC0415 + from pypto.ir.distributed_compiled_program import DistributedConfig # noqa: PLC0415 + from pypto.runtime import RunConfig # noqa: PLC0415 + + config = self._run_config(codegen_only=True) + distributed_config = DistributedConfig( + device_ids=list(self._device_ids), + num_sub_workers=0, + ) + run_config = RunConfig( + platform=config.platform, + device_id=config.device_id, + backend_type=config.backend_type, + strategy=config.strategy, + dump_passes=config.dump_passes, + save_kernels=config.save_kernels, + save_kernels_dir=config.save_kernels_dir, + codegen_only=True, + pto_isa_commit=config.pto_isa_commit, + diagnostic_phase=config.diagnostic_phase, + disabled_diagnostics=config.disabled_diagnostics, + compile_profiling=config.compile_profiling, + enable_scope_stats=True, + distributed_config=distributed_config, + ) + compiled = jit_fn.compile(*dummy_args, config=run_config) + if not isinstance(compiled, DistributedCompiledProgram): + raise TypeError(f"{name} did not compile to DistributedCompiledProgram; got {type(compiled).__name__}") + return DeepSeekV4L3Callable(compiled=compiled, name=name) + + def _prefill_dummy_args( + self, + model: RuntimeModel, + layout: DeepSeekV4CacheLayout, + config_module: object, + ) -> tuple[Any, ...]: + """Return explicit serving dummy args for the packed ``l3_prefill_fwd``. + + Like the packed decode_fwd kernel, every weight is layer-stacked on dim 1: + FWD weights stack across all 43 hidden layers, CSA-group weights across the + 21 compress_ratio==4 layers, HCA-group weights across the 20 + compress_ratio==128 layers. The work caches (kv_cache/cmp_kv stack x43, + idx_kv_cache stacks x21) and compressor-state kv/score caches are stacked on + the layer axis. The per-step metadata (slot mappings, block tables, sparse + tables, position ids, input ids), the RoPE tables and the compressor-state + block tables are shared single per-rank copies, matching decode -- the kernel + slices them per layer internally. Prefill runs final RMSNorm and emits + normalized ``x_out`` hidden rows, so host-side LM-head can project only the + rows selected for sampling. It takes a trailing ``num_tokens`` scalar. + """ + cfg = config_module.FLASH + single = self._layer_common_dummy_tensors( + model, + layout, + cfg, + tokens=layout.prefill_seq, + include_decode_indexer=True, + include_prefill_temporaries=False, + ) + ranks = layout.ranks + seq = layout.prefill_seq + hidden = model.config.hidden_size + head_dim = model.config.head_dim + hc_dim = int(cfg.hc_dim) + + fwd = DEEPSEEK_V4_FWD_NUM_LAYERS + csa = DEEPSEEK_V4_CSA_NUM_LAYERS + hca = DEEPSEEK_V4_HCA_NUM_LAYERS + + def stacked(name: str, count: int) -> torch.Tensor: + base = single[name] + shape = (base.shape[0], count * base.shape[1], *base.shape[2:]) + return torch.empty(shape, dtype=base.dtype) + + values: dict[str, torch.Tensor] = {} + # CSA-group weights stack x21; HCA-group weights stack x20; everything else + # in the per-layer common tensors is a FWD weight and stacks x43. The RoPE + # tables and input ids are shared single per-rank copies (the kernel slices + # them per layer internally), matching decode. + for name, base in single.items(): + if name in _PREFILL_FWD_SHARED_COMMON_NAMES: + values[name] = base + elif name in _DECODE_FWD_CSA_STACKED_NAMES: + values[name] = stacked(name, csa) + elif name in _DECODE_FWD_HCA_STACKED_NAMES: + values[name] = stacked(name, hca) + else: + values[name] = stacked(name, fwd) + + values.update( + { + "x_hc": torch.empty((ranks, seq, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + # HCA-group prefill compressor state (x20). + "hca_cmp_kv_state": torch.empty( + ( + ranks, + hca * layout.prefill_hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ), + dtype=torch.float32, + ), + "hca_cmp_score_state": torch.empty( + ( + ranks, + hca * layout.prefill_hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ), + dtype=torch.float32, + ), + "hca_compress_state_block_table": torch.empty( + (ranks, layout.prefill_hca_state_max_blocks), + dtype=torch.int32, + ), + # CSA-group prefill compressor state (x21). + "csa_cmp_kv_state": torch.empty( + ( + ranks, + csa * layout.prefill_csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ), + dtype=torch.float32, + ), + "csa_cmp_score_state": torch.empty( + ( + ranks, + csa * layout.prefill_csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ), + dtype=torch.float32, + ), + "csa_compress_state_block_table": torch.empty( + (ranks, layout.prefill_csa_state_max_blocks), + dtype=torch.int32, + ), + "csa_inner_kv_state": torch.empty( + ( + ranks, + csa * layout.prefill_csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ), + dtype=torch.float32, + ), + "csa_inner_score_state": torch.empty( + ( + ranks, + csa * layout.prefill_csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ), + dtype=torch.float32, + ), + "csa_inner_compress_state_block_table": torch.empty( + (ranks, layout.prefill_csa_inner_state_max_blocks), + dtype=torch.int32, + ), + # FWD-stacked prefill work caches (x43, flattened 5-D); idx_kv_cache + # stacks across the 21 CSA layers. The kernel reshapes the fused + # layer x block axis internally. + "kv_cache": torch.empty( + (ranks, fwd * layout.ori_max_blocks, layout.block_size, 1, head_dim), + dtype=torch.bfloat16, + ), + "cmp_kv": torch.empty( + (ranks, fwd * layout.prefill_cmp_block_num, layout.block_size, 1, head_dim), + dtype=torch.bfloat16, + ), + "idx_kv_cache": torch.empty( + (ranks, csa * layout.prefill_idx_block_num, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + dtype=torch.bfloat16, + ), + # Shared single per-rank prefill metadata (the kernel passes each + # whole tensor to every layer). + "ori_block_table": torch.empty((ranks, layout.ori_max_blocks), dtype=torch.int32), + "ori_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + "cmp_block_table": torch.empty((ranks, layout.prefill_cmp_max_blocks), dtype=torch.int32), + "cmp_sparse_indices": torch.empty( + (ranks, seq, layout.prefill_sparse_topk), + dtype=torch.int32, + ), + "cmp_sparse_lens": torch.empty((ranks, seq), dtype=torch.int32), + "idx_block_table": torch.empty((ranks, layout.prefill_idx_max_blocks), dtype=torch.int32), + "position_ids": torch.empty((ranks, seq), dtype=torch.int32), + "hca_cmp_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + "hca_state_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + "csa_cmp_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + "csa_idx_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + "csa_state_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + "csa_inner_state_slot_mapping": torch.empty((ranks, seq), dtype=torch.long), + # 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), + "hc_head_base": torch.empty((ranks, DEEPSEEK_V4_HC_MULT), dtype=torch.float32), + # Final RMSNorm in-kernel; host-side LM-head consumes selected + # normalized rows from x_out. + "final_norm_w": torch.empty((ranks, hidden), dtype=torch.bfloat16), + "x_out": torch.empty((ranks, seq, hidden), dtype=torch.bfloat16), + } + ) + # The packed prefill kernel emits normalized hidden rows and takes a + # trailing INT32 ``num_tokens`` scalar. + return (*self._ordered_dummy_args(values, _PREFILL_FWD_TENSOR_ORDER), self._int32_arg(seq)) + + def _decode_dummy_args( + self, + model: RuntimeModel, + layout: DeepSeekV4CacheLayout, + config_module: object, + ) -> tuple[Any, ...]: + """Return explicit serving dummy args for the packed ``l3_decode_fwd``. + + Every weight/state argument is layer-stacked on dim 1: FWD weights and + the kv/cmp work caches stack across all 43 hidden layers; CSA-group + weights and state stack across the 21 compress_ratio==4 layers; HCA-group + weights and state stack across the 20 compress_ratio==128 layers. + """ + cfg = config_module.FLASH + single = self._layer_common_dummy_tensors( + model, + layout, + cfg, + tokens=layout.decode_tokens, + include_decode_indexer=True, + include_prefill_temporaries=False, + ) + ranks = layout.ranks + batch = layout.decode_batch + tokens = layout.decode_tokens + hidden = model.config.hidden_size + hc_dim = int(cfg.hc_dim) + + fwd = DEEPSEEK_V4_FWD_NUM_LAYERS + csa = DEEPSEEK_V4_CSA_NUM_LAYERS + hca = DEEPSEEK_V4_HCA_NUM_LAYERS + + def stacked(name: str, count: int) -> torch.Tensor: + base = single[name] + shape = (base.shape[0], count * base.shape[1], *base.shape[2:]) + return torch.empty(shape, dtype=base.dtype) + + values: dict[str, torch.Tensor] = {} + # CSA-group weights stack x21; HCA-group weights stack x20; everything + # else in the per-layer common tensors is a FWD weight and stacks x43. + # Shared single-copy inputs (freqs/input_ids) are populated explicitly. + for name, base in single.items(): + if name in _DECODE_FWD_SHARED_COMMON_NAMES: + values[name] = base + elif name in _DECODE_FWD_CSA_STACKED_NAMES: + values[name] = stacked(name, csa) + elif name in _DECODE_FWD_HCA_STACKED_NAMES: + values[name] = stacked(name, hca) + else: + values[name] = stacked(name, fwd) + + values.update( + { + "x_hc": torch.empty((ranks, tokens, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + # FWD-stacked work caches (x43). + "kv_cache": torch.empty( + ( + ranks, + fwd * batch * layout.ori_max_blocks, + layout.block_size, + 1, + model.config.head_dim, + ), + dtype=torch.bfloat16, + ), + "cmp_kv": torch.empty( + ( + ranks, + fwd * batch * layout.cmp_max_blocks, + layout.block_size, + 1, + model.config.head_dim, + ), + dtype=torch.bfloat16, + ), + # CSA-group state (x21). + "idx_kv_cache": torch.empty( + ( + ranks, + csa * batch * layout.idx_max_blocks, + layout.block_size, + 1, + DEEPSEEK_V4_IDX_HEAD_DIM, + ), + dtype=torch.bfloat16, + ), + "csa_compress_state": torch.empty( + ( + ranks, + csa * batch * layout.csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_STATE_DIM, + ), + dtype=torch.float32, + ), + "csa_inner_compress_state": torch.empty( + ( + ranks, + csa * batch * layout.csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_STATE_DIM, + ), + dtype=torch.float32, + ), + # HCA-group state (x20). + "hca_compress_state": torch.empty( + ( + ranks, + hca * batch * layout.hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_STATE_DIM, + ), + dtype=torch.float32, + ), + # Shared single-copy per-step inputs. + "block_table": torch.empty((ranks, batch, layout.ori_max_blocks), dtype=torch.int32), + "ori_slot_mapping": torch.empty((ranks, tokens), dtype=torch.long), + "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( + (ranks, batch, layout.hca_state_max_blocks), + dtype=torch.int32, + ), + "csa_compress_state_block_table": torch.empty( + (ranks, batch, layout.csa_state_max_blocks), + dtype=torch.int32, + ), + "csa_inner_compress_state_block_table": torch.empty( + (ranks, batch, layout.csa_inner_state_max_blocks), + dtype=torch.int32, + ), + "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), + # 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), + "hc_head_base": torch.empty((ranks, DEEPSEEK_V4_HC_MULT), dtype=torch.float32), + # Decode writes final-normalized hidden rows; host-side LM-head + # turns the selected rows into logits. + "final_norm_w": torch.empty((ranks, hidden), dtype=torch.bfloat16), + "x_out": torch.empty((ranks, tokens, hidden), dtype=torch.bfloat16), + } + ) + # The packed decode kernel takes a trailing INT32 ``num_tokens`` scalar + # (the real active token count), mirroring prefill. + return (*self._ordered_dummy_args(values, _DECODE_FWD_TENSOR_ORDER), self._int32_arg(tokens)) + + def _layer_common_dummy_tensors( + self, + model: RuntimeModel, + layout: DeepSeekV4CacheLayout, + cfg: object, + *, + tokens: int, + include_decode_indexer: bool, + include_prefill_temporaries: bool, + ) -> dict[str, torch.Tensor]: + """Return explicit dummy tensors shared by prefill and decode layer kernels.""" + del include_prefill_temporaries + ranks = layout.ranks + hidden = model.config.hidden_size + heads = model.config.num_attention_heads + head_dim = model.config.head_dim + q_lora = int(cfg.q_lora_rank) + o_lora = int(cfg.o_lora_rank) + o_groups = int(cfg.o_groups) + o_group_in = heads * head_dim // o_groups + mix_hc = int(cfg.mix_hc) + hc_dim = int(cfg.hc_dim) + max_seq_len = int(cfg.max_position_embeddings) + rope_dim = int(cfg.qk_rope_head_dim) + moe_inter = int(cfg.moe_intermediate_size) + n_routed_experts = int(cfg.n_routed_experts) + n_local = n_routed_experts // ranks + topk = int(cfg.num_experts_per_tok) + index_heads = int(cfg.index_n_heads) + index_dim = int(cfg.index_head_dim) + values = { + "hc_attn_fn": torch.empty((ranks, mix_hc, hc_dim), dtype=torch.float32), + "hc_attn_scale": torch.empty((ranks, 3), dtype=torch.float32), + "hc_attn_base": torch.empty((ranks, mix_hc), dtype=torch.float32), + "attn_norm_w": torch.empty((ranks, hidden), dtype=torch.bfloat16), + "wq_a": torch.empty((ranks, hidden, q_lora), dtype=torch.bfloat16), + "wq_b": torch.empty((ranks, q_lora, heads * head_dim), dtype=torch.int8), + "wq_b_scale": torch.empty((ranks, heads * head_dim), dtype=torch.float32), + "wkv": torch.empty((ranks, hidden, head_dim), dtype=torch.bfloat16), + "gamma_cq": torch.empty((ranks, q_lora), dtype=torch.bfloat16), + "gamma_ckv": torch.empty((ranks, head_dim), dtype=torch.bfloat16), + "freqs_cos": torch.empty((ranks, max_seq_len, rope_dim), dtype=torch.bfloat16), + "freqs_sin": torch.empty((ranks, max_seq_len, rope_dim), dtype=torch.bfloat16), + "hca_cmp_wkv": torch.empty((ranks, DEEPSEEK_V4_HCA_MAIN_OUT_DIM, hidden), dtype=torch.bfloat16), + "hca_cmp_wgate": torch.empty((ranks, DEEPSEEK_V4_HCA_MAIN_OUT_DIM, hidden), dtype=torch.bfloat16), + "hca_cmp_ape": torch.empty((ranks, 128, DEEPSEEK_V4_HCA_MAIN_OUT_DIM), dtype=torch.float32), + "hca_cmp_norm_w": torch.empty((ranks, head_dim), dtype=torch.bfloat16), + "csa_cmp_wkv": torch.empty((ranks, DEEPSEEK_V4_CSA_MAIN_OUT_DIM, hidden), dtype=torch.bfloat16), + "csa_cmp_wgate": torch.empty((ranks, DEEPSEEK_V4_CSA_MAIN_OUT_DIM, hidden), dtype=torch.bfloat16), + "csa_cmp_ape": torch.empty((ranks, 4, DEEPSEEK_V4_CSA_MAIN_OUT_DIM), dtype=torch.float32), + "csa_cmp_norm_w": torch.empty((ranks, head_dim), dtype=torch.bfloat16), + "csa_hadamard_idx": torch.empty((ranks, index_dim, index_dim), dtype=torch.bfloat16), + "csa_inner_wkv": torch.empty((ranks, DEEPSEEK_V4_CSA_INNER_OUT_DIM, hidden), dtype=torch.bfloat16), + "csa_inner_wgate": torch.empty((ranks, DEEPSEEK_V4_CSA_INNER_OUT_DIM, hidden), dtype=torch.bfloat16), + "csa_inner_ape": torch.empty((ranks, 4, DEEPSEEK_V4_CSA_INNER_OUT_DIM), dtype=torch.float32), + "csa_inner_norm_w": torch.empty((ranks, index_dim), dtype=torch.bfloat16), + "attn_sink": torch.empty((ranks, heads), dtype=torch.float32), + "wo_a": torch.empty((ranks, o_groups, o_lora, o_group_in), dtype=torch.bfloat16), + "wo_b": torch.empty((ranks, hidden, o_groups * o_lora), dtype=torch.int8), + "wo_b_scale": torch.empty((ranks, hidden), dtype=torch.float32), + "hc_ffn_fn": torch.empty((ranks, mix_hc, hc_dim), dtype=torch.float32), + "hc_ffn_scale": torch.empty((ranks, 3), dtype=torch.float32), + "hc_ffn_base": torch.empty((ranks, mix_hc), dtype=torch.float32), + "norm_w": torch.empty((ranks, hidden), dtype=torch.bfloat16), + "gate_w": torch.empty((ranks, n_routed_experts, hidden), dtype=torch.float32), + "gate_bias": torch.empty((ranks, n_routed_experts), dtype=torch.float32), + "tid2eid": torch.empty((ranks, model.config.vocab_size, topk), dtype=torch.int32), + "input_ids": torch.empty((ranks, tokens), dtype=torch.long), + "routed_w1": torch.empty((ranks, n_local, moe_inter, hidden), dtype=torch.int8), + "routed_w1_scale": torch.empty((ranks, n_local, moe_inter), dtype=torch.float32), + "routed_w3": torch.empty((ranks, n_local, moe_inter, hidden), dtype=torch.int8), + "routed_w3_scale": torch.empty((ranks, n_local, moe_inter), dtype=torch.float32), + "routed_w2": torch.empty((ranks, n_local, hidden, moe_inter), dtype=torch.int8), + "routed_w2_scale": torch.empty((ranks, n_local, hidden), dtype=torch.float32), + "shared_w1": torch.empty((ranks, moe_inter, hidden), dtype=torch.int8), + "shared_w1_scale": torch.empty((ranks, moe_inter), dtype=torch.float32), + "shared_w3": torch.empty((ranks, moe_inter, hidden), dtype=torch.int8), + "shared_w3_scale": torch.empty((ranks, moe_inter), dtype=torch.float32), + "shared_w2": torch.empty((ranks, hidden, moe_inter), dtype=torch.int8), + "shared_w2_scale": torch.empty((ranks, hidden), dtype=torch.float32), + } + if include_decode_indexer: + values.update( + { + "csa_idx_wq_b": torch.empty((ranks, q_lora, index_heads * index_dim), dtype=torch.int8), + "csa_idx_wq_b_scale": torch.empty((ranks, index_heads * index_dim), dtype=torch.float32), + "csa_weights_proj": torch.empty((ranks, hidden, index_heads), dtype=torch.bfloat16), + } + ) + return values + + @staticmethod + def _ordered_dummy_args(values: dict[str, torch.Tensor], names: Sequence[str]) -> tuple[torch.Tensor, ...]: + missing = [name for name in names if name not in values] + if missing: + raise KeyError(f"DeepSeekV4 compile dummy args missing tensors: {', '.join(missing)}") + return tuple(values[name] for name in names) + + @staticmethod + def _int32_arg(value: int) -> Any: + import ctypes + + return ctypes.c_int32(int(value)) + + def _build_rope_tables(self, rope_tables_module: object, config_module: object) -> tuple[torch.Tensor, torch.Tensor]: + """Build full-sequence DeepSeekV4 RoPE tables using pypto-lib's helper.""" + freqs_cos, freqs_sin = rope_tables_module.build_deepseek_v4_rope_tables( + config_module.FLASH, + 0, + dtype=torch.bfloat16, + ) + return freqs_cos.contiguous().cpu(), freqs_sin.contiguous().cpu() + + def _validate_kernel_contract(self, layout: DeepSeekV4CacheLayout) -> None: + """Fail fast when the checked-out pypto-lib kernels do not match serving topology.""" + required_modules = ( + "config.py", + "prefill_attention_hca.py", + "prefill_attention_csa.py", + "prefill_layer.py", + "prefill_fwd.py", + "decode_layer.py", + "decode_fwd.py", + ) + missing = [name for name in required_modules if not (self._kernel_dir / name).is_file()] + if missing: + raise FileNotFoundError( + "DeepSeekV4 kernel directory is missing required modules: " + ", ".join(missing) + ) + + config_path = self._kernel_dir / "config.py" + expected_config = { + "BLOCK_SIZE": layout.block_size, + "DECODE_BATCH": layout.decode_batch, + "DECODE_SEQ": layout.decode_seq, + "DECODE_TOKENS": layout.decode_tokens, + "PREFILL_BATCH": layout.prefill_batch, + "PREFILL_SEQ": layout.prefill_seq, + "KV_ORI_MAX_BLOCKS": layout.ori_max_blocks, + "KV_CMP_MAX_BLOCKS": layout.cmp_max_blocks, + "IDX_CACHE_MAX_BLOCKS": layout.idx_max_blocks, + "PREFILL_CMP_MAX_BLOCKS": layout.prefill_cmp_max_blocks, + "PREFILL_IDX_MAX_BLOCKS": layout.prefill_idx_max_blocks, + "EP_WORLD_SIZE": layout.ranks, + } + mismatched = [] + for name, expected in expected_config.items(): + actual = _int_constant_from_file(config_path, name) + if actual is not None and actual != expected: + mismatched.append(f"{name}={actual} expected {expected}") + expected_module_constants = { + "prefill_attention_hca.py": { + "HCA_STATE_BLOCK_NUM": layout.prefill_hca_state_max_blocks, + "HCA_STATE_MAX_BLOCKS": layout.prefill_hca_state_max_blocks, + }, + "prefill_attention_csa.py": { + "CSA_STATE_BLOCK_NUM": layout.prefill_csa_state_max_blocks, + "CSA_STATE_MAX_BLOCKS": layout.prefill_csa_state_max_blocks, + "INNER_STATE_BLOCK_NUM": layout.prefill_csa_inner_state_max_blocks, + "INNER_STATE_MAX_BLOCKS": layout.prefill_csa_inner_state_max_blocks, + }, + } + for filename, expected_constants in expected_module_constants.items(): + module_path = self._kernel_dir / filename + for name, expected in expected_constants.items(): + actual = _int_constant_from_file(module_path, name) + if actual is not None and actual != expected: + mismatched.append(f"{filename}:{name}={actual} expected {expected}") + if mismatched: + raise ValueError("DeepSeekV4 kernel config does not match serving layout: " + ", ".join(mismatched)) diff --git a/examples/model/deepseek_v4/runner/npu_runner.py b/examples/model/deepseek_v4/runner/npu_runner.py new file mode 100644 index 0000000..8fb3d22 --- /dev/null +++ b/examples/model/deepseek_v4/runner/npu_runner.py @@ -0,0 +1,2963 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import math +import logging +import os +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field, replace +from typing import Any + +import torch +from pypto.runtime import DeviceTensor + +from examples.model.deepseek_v4.runner.weight_loader import DeepSeekV4WeightStore +from examples.model.deepseek_v4.runner.weight_loader import DeepSeekV4GlobalWeights +from examples.model.deepseek_v4.runner.weight_loader import DeepSeekV4StackedLayerWeights +from python.core.model_runner import ModelRunner +from python.core.types import ( + DecodeBatch, + DecodeResult, + ModelConfig, + ModelRecord, + PrefillBatch, + PrefillResult, + RuntimeConfig, + RuntimeModel, +) + +logger = logging.getLogger(__name__) + + +DEEPSEEK_V4_RANKS = 8 +DEEPSEEK_V4_HC_MULT = 4 +DEEPSEEK_V4_BLOCK_SIZE = 128 +DEEPSEEK_V4_DECODE_BATCH = 8 +DEEPSEEK_V4_DECODE_SEQ = 1 +DEEPSEEK_V4_DECODE_TOKENS = DEEPSEEK_V4_DECODE_BATCH * DEEPSEEK_V4_DECODE_SEQ +DEEPSEEK_V4_PREFILL_BATCH = 1 +DEEPSEEK_V4_PREFILL_SEQ = 128 +DEEPSEEK_V4_ORI_MAX_BLOCKS = 1 +DEEPSEEK_V4_CMP_MAX_BLOCKS = 32 +DEEPSEEK_V4_IDX_MAX_BLOCKS = 64 +DEEPSEEK_V4_HCA_STATE_MAX_BLOCKS = 64 +DEEPSEEK_V4_CSA_STATE_MAX_BLOCKS = 65 +DEEPSEEK_V4_CSA_INNER_STATE_MAX_BLOCKS = 65 +DEEPSEEK_V4_C128_STATE_BLOCK_SIZE = 8 +DEEPSEEK_V4_C4_STATE_BLOCK_SIZE = 4 +DEEPSEEK_V4_PREFILL_CMP_MAX_BLOCKS = DEEPSEEK_V4_CMP_MAX_BLOCKS +DEEPSEEK_V4_PREFILL_IDX_MAX_BLOCKS = DEEPSEEK_V4_IDX_MAX_BLOCKS +DEEPSEEK_V4_PREFILL_HCA_STATE_MAX_BLOCKS = 2048 +DEEPSEEK_V4_PREFILL_CSA_STATE_MAX_BLOCKS = 4096 +DEEPSEEK_V4_PREFILL_CSA_INNER_STATE_MAX_BLOCKS = 4096 +DEEPSEEK_V4_INDEX_TOPK = 512 +DEEPSEEK_V4_PREFILL_SPARSE_TOPK = DEEPSEEK_V4_BLOCK_SIZE + DEEPSEEK_V4_INDEX_TOPK +DEEPSEEK_V4_HEAD_DIM = 512 +DEEPSEEK_V4_IDX_HEAD_DIM = 128 +DEEPSEEK_V4_HCA_MAIN_OUT_DIM = 512 +DEEPSEEK_V4_CSA_MAIN_OUT_DIM = 1024 +DEEPSEEK_V4_CSA_INNER_OUT_DIM = 256 +DEEPSEEK_V4_HCA_STATE_DIM = 2 * DEEPSEEK_V4_HCA_MAIN_OUT_DIM +DEEPSEEK_V4_CSA_STATE_DIM = 2 * DEEPSEEK_V4_CSA_MAIN_OUT_DIM +DEEPSEEK_V4_CSA_INNER_STATE_DIM = 2 * DEEPSEEK_V4_CSA_INNER_OUT_DIM +DEEPSEEK_V4_RMS_NORM_EPS = 1e-6 +DEEPSEEK_V4_HC_EPS = 1e-6 +# Layer-stacking counts for the packed all-layer decode_fwd kernel. +DEEPSEEK_V4_FWD_NUM_LAYERS = 43 +DEEPSEEK_V4_CSA_NUM_LAYERS = 21 +DEEPSEEK_V4_HCA_NUM_LAYERS = 20 + + +# Argument order for the packed all-43-layer ``l3_prefill_fwd`` kernel. This +# mirrors pypto-lib prefill_fwd.py ``l3_prefill_fwd`` host signature: every +# layer-stacked weight/state tensor in core-parameter order, followed by the +# ``hc_head`` collapse weights, final RMSNorm input and an ``x_out`` output. The +# kernel stops before LM-head; logits are computed on the host from selected +# normalized hidden rows. A trailing ``num_tokens`` scalar is appended at dispatch. +# The work caches +# (kv_cache/cmp_kv/idx_kv_cache) are kernel ``pl.Out`` tensors; weights and +# metadata are inputs. +_PREFILL_FWD_TENSOR_ORDER = ( + "x_hc", + "hc_attn_fn", + "hc_attn_scale", + "hc_attn_base", + "attn_norm_w", + "wq_a", + "wq_b", + "wq_b_scale", + "wkv", + "gamma_cq", + "gamma_ckv", + "kv_cache", + "attn_sink", + "wo_a", + "wo_b", + "wo_b_scale", + "cmp_kv", + "hca_cmp_wkv", + "hca_cmp_wgate", + "hca_cmp_ape", + "hca_cmp_norm_w", + "hca_cmp_kv_state", + "hca_cmp_score_state", + "csa_cmp_wkv", + "csa_cmp_wgate", + "csa_cmp_ape", + "csa_cmp_norm_w", + "csa_cmp_kv_state", + "csa_cmp_score_state", + "csa_hadamard_idx", + "csa_idx_wq_b", + "csa_idx_wq_b_scale", + "csa_weights_proj", + "csa_inner_wkv", + "csa_inner_wgate", + "csa_inner_ape", + "csa_inner_norm_w", + "csa_inner_kv_state", + "csa_inner_score_state", + "idx_kv_cache", + "hca_compress_state_block_table", + "csa_compress_state_block_table", + "csa_inner_compress_state_block_table", + "freqs_cos", + "freqs_sin", + "ori_block_table", + "cmp_block_table", + "idx_block_table", + "ori_slot_mapping", + "position_ids", + "input_ids", + "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", + "cmp_sparse_indices", + "cmp_sparse_lens", + "hc_ffn_fn", + "hc_ffn_scale", + "hc_ffn_base", + "norm_w", + "gate_w", + "gate_bias", + "tid2eid", + "routed_w1", + "routed_w1_scale", + "routed_w3", + "routed_w3_scale", + "routed_w2", + "routed_w2_scale", + "shared_w1", + "shared_w1_scale", + "shared_w3", + "shared_w3_scale", + "shared_w2", + "shared_w2_scale", + "hc_head_fn", + "hc_head_scale", + "hc_head_base", + "final_norm_w", + "x_out", +) + +# Argument order for the packed all-43-layer ``l3_decode_fwd`` kernel. This +# mirrors pypto-lib decode_fwd.py ``l3_decode_fwd`` host signature: after the +# ``hc_head`` collapse weights the kernel performs final RMSNorm and writes +# normalized ``x_out``. LM-head is computed on the host side. +_DECODE_FWD_TENSOR_ORDER = ( + "x_hc", + "hc_attn_fn", + "hc_attn_scale", + "hc_attn_base", + "attn_norm_w", + "wq_a", + "wq_b", + "wq_b_scale", + "wkv", + "gamma_cq", + "gamma_ckv", + "kv_cache", + "attn_sink", + "wo_a", + "wo_b", + "wo_b_scale", + "hca_cmp_wkv", + "hca_cmp_wgate", + "hca_cmp_ape", + "hca_cmp_norm_w", + "hca_compress_state", + "csa_cmp_wkv", + "csa_cmp_wgate", + "csa_cmp_ape", + "csa_cmp_norm_w", + "csa_compress_state", + "csa_idx_wq_b", + "csa_idx_wq_b_scale", + "csa_weights_proj", + "csa_hadamard_idx", + "csa_inner_wkv", + "csa_inner_wgate", + "csa_inner_ape", + "csa_inner_norm_w", + "csa_inner_compress_state", + "cmp_kv", + "idx_kv_cache", + "hc_ffn_fn", + "hc_ffn_scale", + "hc_ffn_base", + "norm_w", + "gate_w", + "gate_bias", + "tid2eid", + "routed_w1", + "routed_w1_scale", + "routed_w3", + "routed_w3_scale", + "routed_w2", + "routed_w2_scale", + "shared_w1", + "shared_w1_scale", + "shared_w3", + "shared_w3_scale", + "shared_w2", + "shared_w2_scale", + "freqs_cos", + "freqs_sin", + "block_table", + "ori_slot_mapping", + "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", + "csa_compress_state_block_table", + "csa_inner_compress_state_block_table", + "cmp_block_table", + "idx_block_table", + "input_ids", + "hc_head_fn", + "hc_head_scale", + "hc_head_base", + "final_norm_w", + "x_out", +) + +_DECODE_INPUT_TENSOR_FIELDS = ( + "input_ids", + "position_ids", + "kv_seq_lens", + "block_table", + "ori_slot_mapping", + "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", +) + + +@dataclass(frozen=True) +class DeepSeekV4CacheLayout: + """Static cache layout baked into the current DeepSeekV4 kernels.""" + + ranks: int = DEEPSEEK_V4_RANKS + hc_mult: int = DEEPSEEK_V4_HC_MULT + block_size: int = DEEPSEEK_V4_BLOCK_SIZE + decode_batch: int = DEEPSEEK_V4_DECODE_BATCH + decode_seq: int = DEEPSEEK_V4_DECODE_SEQ + decode_tokens: int = DEEPSEEK_V4_DECODE_TOKENS + prefill_batch: int = DEEPSEEK_V4_PREFILL_BATCH + prefill_seq: int = DEEPSEEK_V4_PREFILL_SEQ + ori_max_blocks: int = DEEPSEEK_V4_ORI_MAX_BLOCKS + cmp_max_blocks: int = DEEPSEEK_V4_CMP_MAX_BLOCKS + idx_max_blocks: int = DEEPSEEK_V4_IDX_MAX_BLOCKS + hca_state_max_blocks: int = DEEPSEEK_V4_HCA_STATE_MAX_BLOCKS + csa_state_max_blocks: int = DEEPSEEK_V4_CSA_STATE_MAX_BLOCKS + csa_inner_state_max_blocks: int = DEEPSEEK_V4_CSA_INNER_STATE_MAX_BLOCKS + c128_state_block_size: int = DEEPSEEK_V4_C128_STATE_BLOCK_SIZE + c4_state_block_size: int = DEEPSEEK_V4_C4_STATE_BLOCK_SIZE + prefill_cmp_max_blocks: int = DEEPSEEK_V4_PREFILL_CMP_MAX_BLOCKS + prefill_idx_max_blocks: int = DEEPSEEK_V4_PREFILL_IDX_MAX_BLOCKS + prefill_hca_state_max_blocks: int = DEEPSEEK_V4_PREFILL_HCA_STATE_MAX_BLOCKS + prefill_csa_state_max_blocks: int = DEEPSEEK_V4_PREFILL_CSA_STATE_MAX_BLOCKS + prefill_csa_inner_state_max_blocks: int = DEEPSEEK_V4_PREFILL_CSA_INNER_STATE_MAX_BLOCKS + prefill_sparse_topk: int = DEEPSEEK_V4_PREFILL_SPARSE_TOPK + + @property + def prefill_cmp_block_num(self) -> int: + """Physical cmp_kv blocks per layer in the packed prefill kernel.""" + return self.decode_batch * self.prefill_cmp_max_blocks + + @property + def prefill_idx_block_num(self) -> int: + """Physical idx_kv_cache blocks per CSA layer in the packed prefill kernel.""" + return self.decode_batch * self.prefill_idx_max_blocks + + def validate_runtime(self, config: ModelConfig, runtime: RuntimeConfig, device_ids: Sequence[int]) -> None: + """Validate serving/runtime options against kernel-fixed dimensions.""" + if len(device_ids) != self.ranks: + raise ValueError(f"DeepSeekV4 requires exactly {self.ranks} devices, got {len(device_ids)}") + if runtime.page_size != self.block_size: + raise ValueError(f"DeepSeekV4 kernels require page_size={self.block_size}, got {runtime.page_size}") + if runtime.max_batch_size > self.decode_batch: + raise ValueError( + f"DeepSeekV4 decode kernels support at most {self.decode_batch} active rows, " + f"got max_batch_size={runtime.max_batch_size}" + ) + decode_state_capacity = self.csa_state_max_blocks * self.c4_state_block_size + if runtime.max_seq_len > decode_state_capacity: + raise ValueError( + "DeepSeekV4 pypto-lib decode CSA state tables currently support at most " + f"max_seq_len={decode_state_capacity}, got {runtime.max_seq_len}. " + "Increase the decode CSA state table depth in pypto-lib before serving longer contexts." + ) + if self.decode_tokens != self.decode_batch * self.decode_seq: + raise ValueError("DeepSeekV4 layout decode_tokens must equal decode_batch * decode_seq") + expected = { + "hidden_size": 4096, + "num_hidden_layers": 43, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "vocab_size": 129280, + } + actual = { + "hidden_size": config.hidden_size, + "num_hidden_layers": config.num_hidden_layers, + "num_attention_heads": config.num_attention_heads, + "num_key_value_heads": config.num_key_value_heads, + "head_dim": config.head_dim, + "vocab_size": config.vocab_size, + } + if actual != expected: + mismatch = ", ".join(f"{name}={actual[name]} expected {value}" for name, value in expected.items()) + raise ValueError("DeepSeekV4 W8A8 kernels require Flash shape: " + mismatch) + + +@dataclass +class DeepSeekV4CacheManager: + """Request-to-cache-slot mapping and table builders for DeepSeekV4 kernels.""" + + layout: DeepSeekV4CacheLayout = field(default_factory=DeepSeekV4CacheLayout) + _request_to_slot: dict[str, int] = field(default_factory=dict) + _free_slots: list[int] = field(default_factory=list) + + def __post_init__(self) -> None: + if not self._free_slots: + self._free_slots = list(range(self.layout.decode_batch)) + + @property + def active_slots(self) -> dict[str, int]: + """Return a copy of currently assigned request slots.""" + return dict(self._request_to_slot) + + @property + def free_count(self) -> int: + """Return the number of unassigned decode slots.""" + return len(self._free_slots) + + def allocate(self, request_id: str) -> int | None: + """Assign a stable decode slot to ``request_id``.""" + if request_id in self._request_to_slot: + return self._request_to_slot[request_id] + if not self._free_slots: + return None + slot = self._free_slots.pop(0) + self._request_to_slot[request_id] = slot + return slot + + def release(self, request_ids: Iterable[str]) -> None: + """Release slots held by finished or aborted requests.""" + for request_id in request_ids: + slot = self._request_to_slot.pop(request_id, None) + if slot is not None and slot not in self._free_slots: + self._free_slots.append(slot) + self._free_slots.sort() + + def slots_for_request_ids(self, request_ids: Sequence[str]) -> list[int]: + """Return assigned slots for request ids, allocating missing slots.""" + slots = [] + for request_id in request_ids: + slot = self.allocate(request_id) + if slot is None: + raise RuntimeError("DeepSeekV4 cache slots exhausted") + slots.append(slot) + return slots + + def block_table(self, slots: Sequence[int], *, max_blocks: int) -> torch.Tensor: + """Build a row-major block table for request-owned physical block ranges.""" + table = torch.empty((len(slots), max_blocks), dtype=torch.int32) + for row, slot in enumerate(slots): + start = int(slot) * max_blocks + table[row].copy_(torch.arange(start, start + max_blocks, dtype=torch.int32)) + return table + + def slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[Sequence[int]], + *, + max_blocks: int, + block_size: int | None = None, + compress_ratio: int = 1, + ) -> torch.Tensor: + """Map logical token positions to physical cache rows for each request slot.""" + block_size = self.layout.block_size if block_size is None else int(block_size) + if compress_ratio <= 0: + raise ValueError("compress_ratio must be positive") + capacity = max_blocks * block_size + max_tokens = max((len(row) for row in positions), default=0) + mapping = torch.full((len(slots), max_tokens), -1, dtype=torch.int64) + for row, (slot, row_positions) in enumerate(zip(slots, positions, strict=True)): + base = int(slot) * capacity + for col, position in enumerate(row_positions): + logical = int(position) // compress_ratio + if logical >= capacity: + raise ValueError( + f"position {position} maps to logical cache row {logical}, " + f"but capacity is {capacity}" + ) + mapping[row, col] = base + logical + return mapping + + def block_table_for_kernel_rows( + self, + slots: Sequence[int], + *, + max_blocks: int, + kernel_rows: int, + ) -> torch.Tensor: + """Build a fixed-row block table, replicating row 0 into inactive rows.""" + if not slots: + raise ValueError("slots must not be empty") + active = self.block_table(slots, max_blocks=max_blocks) + return self.replicate_first_row(active, actual_rows=len(slots), kernel_rows=kernel_rows) + + def sliding_window_slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[Sequence[int]], + *, + kernel_rows: int, + ) -> torch.Tensor: + """Map absolute positions into the 128-token ori sliding-window cache.""" + rows = self._replicated_slots_and_positions(slots, positions, kernel_rows=kernel_rows) + mapping = torch.full((kernel_rows, max((len(row) for _, row in rows), default=0)), -1, dtype=torch.int64) + for row_idx, (slot, row_positions) in enumerate(rows): + base = int(slot) * self.layout.ori_max_blocks * self.layout.block_size + for col, position in enumerate(row_positions): + window_slot = int(position) % self.layout.block_size + mapping[row_idx, col] = base + window_slot + return mapping + + def compressed_slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[Sequence[int]], + *, + max_blocks: int, + compress_ratio: int, + kernel_rows: int, + ) -> torch.Tensor: + """Map compression-boundary positions into a compressed KV cache.""" + rows = self._replicated_slots_and_positions(slots, positions, kernel_rows=kernel_rows) + mapping = torch.full((kernel_rows, max((len(row) for _, row in rows), default=0)), -1, dtype=torch.int64) + capacity = max_blocks * self.layout.block_size + for row_idx, (slot, row_positions) in enumerate(rows): + base = int(slot) * capacity + for col, position in enumerate(row_positions): + position = int(position) + if (position + 1) % compress_ratio != 0: + continue + logical = position // compress_ratio + if logical >= capacity: + raise ValueError( + f"position {position} maps to compressed row {logical}, " + f"but capacity is {capacity}" + ) + mapping[row_idx, col] = base + logical + return mapping + + def state_slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[Sequence[int]], + *, + max_blocks: int, + state_block_size: int, + kernel_rows: int, + ) -> torch.Tensor: + """Map absolute token positions into a compressor-state cache.""" + rows = self._replicated_slots_and_positions(slots, positions, kernel_rows=kernel_rows) + mapping = torch.full((kernel_rows, max((len(row) for _, row in rows), default=0)), -1, dtype=torch.int64) + capacity = max_blocks * state_block_size + for row_idx, (slot, row_positions) in enumerate(rows): + base = int(slot) * capacity + for col, position in enumerate(row_positions): + position = int(position) + if position >= capacity: + raise ValueError( + f"position {position} exceeds compressor-state capacity {capacity} " + f"(max_blocks={max_blocks}, state_block_size={state_block_size})" + ) + mapping[row_idx, col] = base + position + return mapping + + @staticmethod + def _replicated_slots_and_positions( + slots: Sequence[int], + positions: Sequence[Sequence[int]], + *, + kernel_rows: int, + ) -> list[tuple[int, Sequence[int]]]: + if not slots: + raise ValueError("slots must not be empty") + if len(slots) != len(positions): + raise ValueError("slots and positions must have the same active row count") + if len(slots) > kernel_rows: + raise ValueError("active rows exceed kernel_rows") + rows = [(int(slot), tuple(int(pos) for pos in row)) for slot, row in zip(slots, positions, strict=True)] + rows.extend((rows[0][0], rows[0][1]) for _ in range(kernel_rows - len(rows))) + return rows + + @staticmethod + def replicate_first_row(tensor: torch.Tensor, *, actual_rows: int, kernel_rows: int) -> torch.Tensor: + """Pad kernel inputs by replicating row 0 into inactive rows.""" + if actual_rows <= 0: + raise ValueError("actual_rows must be positive") + if kernel_rows < actual_rows: + raise ValueError("kernel_rows must be >= actual_rows") + if tensor.shape[0] < actual_rows: + raise ValueError("tensor has fewer rows than actual_rows") + out = torch.empty((kernel_rows, *tensor.shape[1:]), dtype=tensor.dtype) + out[:actual_rows].copy_(tensor[:actual_rows]) + if actual_rows < kernel_rows: + out[actual_rows:].copy_(tensor[0:1].expand(kernel_rows - actual_rows, *tensor.shape[1:])) + return out + + +class DeepSeekV4InputBuilder: + """Build fixed-shape host inputs for DeepSeekV4 HC-stack kernels.""" + + def __init__(self, *, layout: DeepSeekV4CacheLayout, hidden_size: int) -> None: + self.layout = layout + self.hidden_size = int(hidden_size) + + def prefill_x_hc(self, embeddings: torch.Tensor, *, actual_tokens: int) -> torch.Tensor: + """Build ``[ranks, 128, hc_mult, hidden]`` prefill HC input.""" + if embeddings.ndim != 2: + raise ValueError(f"prefill embeddings must be rank-2, got shape={tuple(embeddings.shape)}") + return self._x_hc_from_rows( + embeddings, + actual_tokens=actual_tokens, + token_rows=self.layout.prefill_seq, + ) + + def decode_x_hc( + self, + embeddings: torch.Tensor, + *, + actual_batch: int, + prev_embeddings: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build ``[ranks, 128, hc_mult, hidden]`` decode HC input. + + Current DeepSeekV4 decode kernels use a fixed ``decode_tokens`` contract + with ``decode_seq`` token slots per request. If ``decode_seq`` is greater + than one and ``prev_embeddings`` is provided, earlier slots carry the + previous token and the final slot carries the last token. Padding rows + still replicate row 0 / their own embedding to keep the fixed rows valid. + """ + if embeddings.ndim != 2: + raise ValueError(f"decode embeddings must be rank-2, got shape={tuple(embeddings.shape)}") + if actual_batch <= 0: + raise ValueError("actual_batch must be positive") + if actual_batch > self.layout.decode_batch: + raise ValueError( + f"actual_batch={actual_batch} exceeds decode batch capacity {self.layout.decode_batch}" + ) + if embeddings.shape[0] < actual_batch: + raise ValueError("decode embeddings has fewer rows than actual_batch") + if prev_embeddings is not None and prev_embeddings.shape[0] < actual_batch: + raise ValueError("decode prev_embeddings has fewer rows than actual_batch") + rows = torch.zeros( + (self.layout.decode_tokens, self.hidden_size), + dtype=embeddings.dtype, + device=embeddings.device, + ) + decode_seq = self.layout.decode_seq + # When the caller supplies a full per-row embedding tensor (one row per + # decode-batch slot), use each row's own embedding so the MoE gate routes + # the 128 tokens across many experts. Otherwise replicate slot 0 into the + # padding rows as before. + per_row = embeddings.shape[0] >= self.layout.decode_batch + for row in range(self.layout.decode_batch): + source_row = row if per_row else (row if row < actual_batch else 0) + start = row * decode_seq + if prev_embeddings is not None and row < actual_batch: + # Fill every slot with prev, then overwrite the final slot with + # the last token. + rows[start : start + decode_seq].copy_( + prev_embeddings[row : row + 1].expand(decode_seq, self.hidden_size) + ) + rows[start + decode_seq - 1].copy_(embeddings[row]) + else: + rows[start : start + decode_seq].copy_( + embeddings[source_row : source_row + 1].expand(decode_seq, self.hidden_size) + ) + return self._expand_hc_and_ranks(rows) + + def _x_hc_from_rows( + self, + embeddings: torch.Tensor, + *, + actual_tokens: int, + token_rows: int, + ) -> torch.Tensor: + if actual_tokens <= 0: + raise ValueError("actual_tokens must be positive") + if actual_tokens > token_rows: + raise ValueError(f"actual_tokens={actual_tokens} exceeds token row capacity {token_rows}") + if embeddings.shape[0] < actual_tokens: + raise ValueError("embeddings has fewer rows than actual_tokens") + if int(embeddings.shape[1]) != self.hidden_size: + raise ValueError(f"embedding hidden size must be {self.hidden_size}, got {int(embeddings.shape[1])}") + rows = torch.zeros((token_rows, self.hidden_size), dtype=embeddings.dtype, device=embeddings.device) + rows[:actual_tokens].copy_(embeddings[:actual_tokens]) + return self._expand_hc_and_ranks(rows) + + def _expand_hc_and_ranks(self, rows: torch.Tensor) -> torch.Tensor: + return ( + rows.unsqueeze(1) + .expand(rows.shape[0], self.layout.hc_mult, self.hidden_size) + .unsqueeze(0) + .expand(self.layout.ranks, rows.shape[0], self.layout.hc_mult, self.hidden_size) + .contiguous() + ) + + +@dataclass +class DeepSeekV4L3Callable: + """Compiled HOST-dispatched DeepSeekV4 program.""" + + compiled: object + name: str + + +@dataclass +class _StaticDeviceTensor: + """CPU tensor marker uploaded to the shared worker once.""" + + tensor: torch.Tensor + + +@dataclass +class _TransientDeviceTensor: + """CPU tensor marker uploaded for one layer dispatch and then freed.""" + + tensor: torch.Tensor + + +@dataclass +class DeepSeekV4LayerCache: + """Shared decode work-cache tensors for one DeepSeekV4 layer dispatch.""" + + kv_cache: torch.Tensor + cmp_kv: torch.Tensor + idx_kv_cache: torch.Tensor + hca_compress_state: torch.Tensor + csa_compress_state: torch.Tensor + csa_inner_compress_state: torch.Tensor + + +@dataclass +class DeepSeekV4LayerCacheSnapshot: + """Compact parent-side cache snapshot captured after prefill for one layer.""" + + tensors: dict[str, torch.Tensor] + + +@dataclass +class DeepSeekV4CompiledKernels: + """Compiled-kernel placeholder and immutable DeepSeekV4 runtime metadata.""" + + layout: DeepSeekV4CacheLayout + model_dir: str + weight_map: dict[str, str] + weight_store: DeepSeekV4WeightStore + compress_ratios: tuple[int, ...] + layer_plan: tuple["DeepSeekV4LayerPlan", ...] + kernel_dir: str + prefill: DeepSeekV4L3Callable | None = None + decode: DeepSeekV4L3Callable | None = None + freqs_cos: torch.Tensor | None = None + freqs_sin: torch.Tensor | None = None + platform: str = "a2a3" + device_id: int = 0 + n_routed_experts: int = 256 + num_hash_layers: int = 3 + + def l3_callables(self) -> tuple[DeepSeekV4L3Callable, ...]: + """Return every compiled L3 program that the shared worker may run.""" + callables: list[DeepSeekV4L3Callable] = [] + if self.prefill is not None: + callables.append(self.prefill) + if self.decode is not None: + callables.append(self.decode) + return tuple(callables) + + +@dataclass(frozen=True) +class DeepSeekV4PreparedPrefillInputs: + """Fixed-shape host tensors derived from one serving prefill chunk.""" + + request_id: str + slot: int + actual_tokens: int + x_hc: torch.Tensor + input_ids: torch.Tensor + position_ids: torch.Tensor + ori_block_table: torch.Tensor + ori_slot_mapping: 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 + cmp_sparse_indices_by_ratio: dict[int, torch.Tensor] + cmp_sparse_lens_by_ratio: dict[int, torch.Tensor] + + def sparse_inputs_for_ratio(self, compress_ratio: int) -> tuple[torch.Tensor, torch.Tensor]: + """Return prefill sparse-attention inputs for one layer compression ratio.""" + ratio = int(compress_ratio) + return self.cmp_sparse_indices_by_ratio[ratio], self.cmp_sparse_lens_by_ratio[ratio] + + +@dataclass(frozen=True) +class DeepSeekV4PreparedDecodeInputs: + """Fixed-shape host tensors derived from one decode scheduler batch.""" + + request_ids: tuple[str, ...] + slots: tuple[int, ...] + kernel_slots: tuple[int, ...] + actual_batch: int + x_hc: torch.Tensor + input_ids: torch.Tensor + position_ids: torch.Tensor + kv_seq_lens: torch.Tensor + block_table: torch.Tensor + ori_slot_mapping: 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 + + +@dataclass +class _DeepSeekV4DecodeSharedBuffers: + """Reusable decode shared-memory buffers inherited by the L3 chip workers.""" + + x_hc_a: torch.Tensor + x_hc_b: torch.Tensor + x_out: torch.Tensor + tensors: dict[str, torch.Tensor] + + +@dataclass +class _DeepSeekV4PrefillFwdSharedBuffers: + """Reusable packed-prefill shared buffers inherited by the L3 chip workers. + + For the single ``l3_prefill_fwd`` dispatch the work caches are flattened 5-D + (kv_cache/cmp_kv stack across all 43 hidden layers, idx_kv_cache across the 21 + compress_ratio==4 layers) and the compress-state kv/score caches stack across + the CSA (x21) and HCA (x20) groups. The per-step metadata, RoPE tables and + compress-state block tables are shared single per-rank copies (the kernel + slices them per layer). ``tensors`` is keyed by ``_PREFILL_FWD_TENSOR_ORDER`` + name (excluding the stacked weights, which live in ``_stacked_weight_buffers``, + and ``freqs_*``/``x_hc`` which are tracked explicitly). The final normalized + hidden output is held separately in ``_prefill_output_buffer``. + """ + + x_hc: torch.Tensor + freqs_cos: torch.Tensor + freqs_sin: torch.Tensor + tensors: dict[str, torch.Tensor] + + +@dataclass(frozen=True) +class DeepSeekV4LayerPlan: + """Per-layer execution metadata for DeepSeekV4 serving.""" + + layer_id: int + compress_ratio: int + attention_kind: str + include_tid2eid: bool + include_gate_bias: bool + + +def deepseek_v4_attention_kind(compress_ratio: int) -> str: + """Return the DeepSeekV4 attention family for a compression ratio.""" + if compress_ratio == 0: + return "swa" + if compress_ratio == 128: + return "hca" + if compress_ratio == 4: + return "csa" + raise ValueError(f"unsupported DeepSeekV4 attention compress ratio: {compress_ratio}") + + +def build_deepseek_v4_layer_plan( + *, + compress_ratios: Sequence[int], + num_hidden_layers: int, + num_hash_layers: int, +) -> tuple[DeepSeekV4LayerPlan, ...]: + """Build the per-layer serving plan from config metadata.""" + if len(compress_ratios) < num_hidden_layers: + raise ValueError("compress_ratios must include at least one entry per hidden layer") + return tuple( + DeepSeekV4LayerPlan( + layer_id=layer_id, + compress_ratio=int(compress_ratios[layer_id]), + attention_kind=deepseek_v4_attention_kind(int(compress_ratios[layer_id])), + include_tid2eid=layer_id < num_hash_layers, + include_gate_bias=layer_id >= num_hash_layers, + ) + for layer_id in range(num_hidden_layers) + ) + + +class DeepSeekV4ModelRunner(ModelRunner): + """Runner boundary for DeepSeekV4 W8A8 kernels and model-specific caches.""" + + def __init__(self, *, compiled: DeepSeekV4CompiledKernels) -> None: + super().__init__() + self._compiled = compiled + self.cache_manager = DeepSeekV4CacheManager(layout=compiled.layout) + self.input_builder: DeepSeekV4InputBuilder | None = None + self._l3_worker: Any | None = None + self._l3_static_tensors: dict[tuple[int, tuple[int, ...], torch.dtype], DeviceTensor] = {} + self._decode_work_cache: DeepSeekV4LayerCache | None = None + self._decode_cache_seeded_slots: set[int] = set() + self._prefill_cache_snapshots: dict[int, DeepSeekV4LayerCacheSnapshot] = {} + self._global_weights: DeepSeekV4GlobalWeights | None = None + self._static_final_norm_weight: torch.Tensor | None = None + self._static_freqs_cos: torch.Tensor | None = None + self._static_freqs_sin: torch.Tensor | None = None + self._prefill_fwd_buffers: _DeepSeekV4PrefillFwdSharedBuffers | None = None + self._decode_buffers: _DeepSeekV4DecodeSharedBuffers | None = None + self._stacked_weight_buffers: dict[str, torch.Tensor] | None = None + self._hc_head_buffers: dict[str, torch.Tensor] | None = None + self._decode_logits_buffer: torch.Tensor | None = None + self._prefill_output_buffer: torch.Tensor | None = None + + def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> int: + """Initialize runner state and return scheduler-only KV block capacity. + + DeepSeekV4 owns its NPU cache tensors and fixed slot mapping internally, + so no generic KV tensors are allocated here. The scheduler still needs a + positive block pool for host-side request budgeting and preemption. + """ + self.input_builder = DeepSeekV4InputBuilder( + layout=self._compiled.layout, + hidden_size=config.hidden_size, + ) + self._decode_cache_seeded_slots.clear() + if runtime.total_kv_pages is not None: + return int(runtime.total_kv_pages) + max_blocks_per_seq = math.ceil(runtime.max_seq_len / runtime.page_size) + return int(runtime.max_batch_size * max_blocks_per_seq) + + def release_finished_requests(self, request_ids: Iterable[str]) -> None: + """Release runner-owned cache slots for finished requests.""" + request_ids = tuple(request_ids) + self.cache_manager.release(request_ids) + if request_ids: + self._prefill_cache_snapshots.clear() + self._decode_cache_seeded_slots.clear() + + def preflight(self, record: ModelRecord) -> None: + """Eagerly load all W8A8 weights and stage shared buffers before ready. + + DeepSeekV4 defers safetensors reads and layer-stacking to first inference + by design (the model loader only parses config/index). Move that work + here so the serving worker only signals ready once weights are fully + materialized, matching the Qwen eager-load contract. ``_ensure_l3_shared_buffers`` + is idempotent, so the first real request still takes the same fast path. + """ + self._ensure_l3_shared_buffers(record.runtime_model) + + def load_packed_global_weights(self) -> DeepSeekV4GlobalWeights: + """Load global tensors and pack the LM head for host-side projection.""" + if self._global_weights is None: + self._global_weights = self._compiled.weight_store.load_packed_global_weights( + ranks=self._compiled.layout.ranks + ) + return self._global_weights + + def load_stacked_layer_weights(self) -> DeepSeekV4StackedLayerWeights: + """Load and stack all hidden-layer weights for the packed decode_fwd kernel.""" + compress_ratios = tuple(int(layer.compress_ratio) for layer in self._compiled.layer_plan) + return self._compiled.weight_store.load_stacked_layer_weights( + ranks=self._compiled.layout.ranks, + n_routed_experts=self._compiled.n_routed_experts, + compress_ratios=compress_ratios, + num_hash_layers=self._compiled.num_hash_layers, + ) + + def prepare_prefill_inputs(self, model: RuntimeModel, batch: PrefillBatch) -> DeepSeekV4PreparedPrefillInputs: + """Build DeepSeekV4 prefill host inputs for the current scheduler chunk.""" + builder = self._require_input_builder() + layout = self._compiled.layout + if len(batch.request_ids) != layout.prefill_batch: + raise ValueError( + f"DeepSeekV4 prefill kernels support exactly {layout.prefill_batch} request per dispatch, " + f"got {len(batch.request_ids)}" + ) + request_id = batch.request_ids[0] + slot = self.cache_manager.allocate(request_id) + if slot is None: + raise RuntimeError("DeepSeekV4 cache slots exhausted") + + actual_tokens = self._prefill_actual_tokens(batch) + positions = self._prefill_positions(batch, actual_tokens) + if positions[-1] >= model.runtime.max_seq_len: + raise ValueError( + f"prefill position {positions[-1]} exceeds max_seq_len={model.runtime.max_seq_len}" + ) + embeddings = batch.input_embeddings[0, :actual_tokens].to(torch.bfloat16).cpu() + token_ids = batch.token_ids[0, :actual_tokens].detach().cpu().to(torch.long) + kernel_tokens = self._prefill_kernel_tokens(actual_tokens) + kernel_positions = self._prefill_kernel_positions( + positions, + kernel_tokens=kernel_tokens, + max_seq_len=model.runtime.max_seq_len, + ) + kernel_slots = self._prefill_kernel_slots( + slot, + actual_tokens=actual_tokens, + kernel_tokens=kernel_tokens, + ) + kernel_embeddings = self._padded_rows(embeddings, kernel_tokens) + kernel_token_ids = self._padded_vector(token_ids, kernel_tokens, dtype=torch.long) + sparse_by_ratio = self._prefill_sparse_by_ratio(kernel_positions, kernel_tokens) + + return DeepSeekV4PreparedPrefillInputs( + request_id=request_id, + slot=slot, + actual_tokens=actual_tokens, + x_hc=builder.prefill_x_hc(kernel_embeddings, actual_tokens=kernel_tokens), + input_ids=self._rank_stack(self._padded_vector(kernel_token_ids, layout.prefill_seq, dtype=torch.long)), + position_ids=self._rank_stack(self._prefill_position_ids(kernel_positions, layout.prefill_seq)), + ori_block_table=self._rank_stack( + self.cache_manager.block_table([slot], max_blocks=layout.ori_max_blocks)[0] + ), + ori_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_sliding_window_slot_mapping(kernel_slots, kernel_positions), + layout.prefill_seq, + ) + ), + cmp_block_table=self._rank_stack( + self.cache_manager.block_table([slot], max_blocks=layout.prefill_cmp_max_blocks)[0] + ), + idx_block_table=self._rank_stack( + self.cache_manager.block_table([slot], max_blocks=layout.prefill_idx_max_blocks)[0] + ), + hca_compress_state_block_table=self._rank_stack( + self.cache_manager.block_table([slot], max_blocks=layout.prefill_hca_state_max_blocks)[0] + ), + csa_compress_state_block_table=self._rank_stack( + self.cache_manager.block_table([slot], max_blocks=layout.prefill_csa_state_max_blocks)[0] + ), + csa_inner_compress_state_block_table=self._rank_stack( + self.cache_manager.block_table([slot], max_blocks=layout.prefill_csa_inner_state_max_blocks)[0] + ), + hca_cmp_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_compressed_slot_mapping( + kernel_slots, + kernel_positions, + max_blocks=layout.prefill_cmp_max_blocks, + compress_ratio=128, + ), + layout.prefill_seq, + ) + ), + hca_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_state_slot_mapping( + kernel_slots, + kernel_positions, + max_blocks=layout.prefill_hca_state_max_blocks, + state_block_size=layout.c128_state_block_size, + ), + layout.prefill_seq, + ) + ), + csa_cmp_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_compressed_slot_mapping( + kernel_slots, + kernel_positions, + max_blocks=layout.prefill_cmp_max_blocks, + compress_ratio=4, + ), + layout.prefill_seq, + ) + ), + csa_idx_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_compressed_slot_mapping( + kernel_slots, + kernel_positions, + max_blocks=layout.prefill_idx_max_blocks, + compress_ratio=4, + ), + layout.prefill_seq, + ) + ), + csa_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_state_slot_mapping( + kernel_slots, + kernel_positions, + max_blocks=layout.prefill_csa_state_max_blocks, + state_block_size=layout.c4_state_block_size, + ), + layout.prefill_seq, + ) + ), + csa_inner_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self._prefill_state_slot_mapping( + kernel_slots, + kernel_positions, + max_blocks=layout.prefill_csa_inner_state_max_blocks, + state_block_size=layout.c4_state_block_size, + ), + layout.prefill_seq, + ) + ), + cmp_sparse_indices_by_ratio={ + ratio: self._rank_stack(indices) + for ratio, (indices, _) in sparse_by_ratio.items() + }, + cmp_sparse_lens_by_ratio={ + ratio: self._rank_stack(lens) + for ratio, (_, lens) in sparse_by_ratio.items() + }, + ) + + def prepare_decode_inputs(self, model: RuntimeModel, batch: DecodeBatch) -> DeepSeekV4PreparedDecodeInputs: + """Build DeepSeekV4 decode host inputs for the current scheduler batch.""" + builder = self._require_input_builder() + layout = self._compiled.layout + actual_batch = len(batch.request_ids) + if actual_batch <= 0: + raise ValueError("decode batch must contain at least one request") + if actual_batch > layout.decode_batch: + raise ValueError(f"decode batch {actual_batch} exceeds kernel batch {layout.decode_batch}") + slots = self.cache_manager.slots_for_request_ids(batch.request_ids) + positions = self._decode_positions(batch, actual_batch) + max_position = max(max(row) for row in positions) + if max_position >= model.runtime.max_seq_len: + raise ValueError(f"decode position {max_position} exceeds max_seq_len={model.runtime.max_seq_len}") + + prev_token_ids = ( + batch.prev_token_ids.detach().cpu().to(torch.long) + if batch.prev_token_ids is not None + else None + ) + token_ids = self._decode_token_rows( + batch.token_ids.detach().cpu().to(torch.long), + actual_batch, + vocab_size=model.config.vocab_size, + prev_token_ids=prev_token_ids, + ) + decode_embeds = batch.hidden_states.to(torch.bfloat16).cpu() + prev_embeds = ( + batch.prev_hidden_states.to(torch.bfloat16).cpu() + if batch.prev_hidden_states is not None + else None + ) + if os.environ.get("PYPTO_DSV4_DIVERSE_DECODE_PAD") == "1" and actual_batch < layout.decode_batch: + decode_embeds = self._diverse_decode_pad_embeddings(model, decode_embeds, actual_batch) + x_hc = builder.decode_x_hc(decode_embeds, actual_batch=actual_batch, prev_embeddings=prev_embeds) + decode_slots = self._decode_kernel_slots(slots) + decode_positions = (*positions, *((positions[0],) * (layout.decode_batch - actual_batch))) + ori_slot_mapping = self.cache_manager.sliding_window_slot_mapping( + decode_slots, + decode_positions, + kernel_rows=layout.decode_batch, + ) + hca_cmp_slot_mapping = self.cache_manager.compressed_slot_mapping( + decode_slots, + decode_positions, + max_blocks=layout.cmp_max_blocks, + compress_ratio=128, + kernel_rows=layout.decode_batch, + ) + hca_state_slot_mapping = self.cache_manager.state_slot_mapping( + decode_slots, + decode_positions, + max_blocks=layout.hca_state_max_blocks, + state_block_size=layout.c128_state_block_size, + kernel_rows=layout.decode_batch, + ) + csa_cmp_slot_mapping = self.cache_manager.compressed_slot_mapping( + decode_slots, + decode_positions, + max_blocks=layout.cmp_max_blocks, + compress_ratio=4, + kernel_rows=layout.decode_batch, + ) + csa_idx_slot_mapping = self.cache_manager.compressed_slot_mapping( + decode_slots, + decode_positions, + max_blocks=layout.idx_max_blocks, + compress_ratio=4, + kernel_rows=layout.decode_batch, + ) + csa_state_slot_mapping = self.cache_manager.state_slot_mapping( + decode_slots, + decode_positions, + max_blocks=layout.csa_state_max_blocks, + state_block_size=layout.c4_state_block_size, + kernel_rows=layout.decode_batch, + ) + csa_inner_state_slot_mapping = self.cache_manager.state_slot_mapping( + decode_slots, + decode_positions, + max_blocks=layout.csa_inner_state_max_blocks, + state_block_size=layout.c4_state_block_size, + kernel_rows=layout.decode_batch, + ) + + return DeepSeekV4PreparedDecodeInputs( + request_ids=tuple(batch.request_ids), + slots=tuple(slots), + kernel_slots=decode_slots, + actual_batch=actual_batch, + x_hc=x_hc, + input_ids=self._rank_stack(token_ids), + position_ids=self._rank_stack(torch.tensor(decode_positions, dtype=torch.int32).reshape(-1)), + kv_seq_lens=self._rank_stack(self._decode_kv_seq_lens(batch.seq_lens, actual_batch)), + block_table=self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + decode_slots, + max_blocks=layout.ori_max_blocks, + kernel_rows=layout.decode_batch, + ) + ), + ori_slot_mapping=self._rank_stack(ori_slot_mapping.reshape(-1)), + cmp_block_table=self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + decode_slots, + max_blocks=layout.cmp_max_blocks, + kernel_rows=layout.decode_batch, + ) + ), + idx_block_table=self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + decode_slots, + max_blocks=layout.idx_max_blocks, + kernel_rows=layout.decode_batch, + ) + ), + hca_compress_state_block_table=self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + decode_slots, + max_blocks=layout.hca_state_max_blocks, + kernel_rows=layout.decode_batch, + ) + ), + csa_compress_state_block_table=self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + decode_slots, + max_blocks=layout.csa_state_max_blocks, + kernel_rows=layout.decode_batch, + ) + ), + csa_inner_compress_state_block_table=self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + decode_slots, + max_blocks=layout.csa_inner_state_max_blocks, + kernel_rows=layout.decode_batch, + ) + ), + hca_cmp_slot_mapping=self._rank_stack(hca_cmp_slot_mapping.reshape(-1)), + hca_state_slot_mapping=self._rank_stack(hca_state_slot_mapping.reshape(-1)), + csa_cmp_slot_mapping=self._rank_stack(csa_cmp_slot_mapping.reshape(-1)), + csa_idx_slot_mapping=self._rank_stack(csa_idx_slot_mapping.reshape(-1)), + csa_state_slot_mapping=self._rank_stack(csa_state_slot_mapping.reshape(-1)), + csa_inner_state_slot_mapping=self._rank_stack(csa_inner_state_slot_mapping.reshape(-1)), + ) + + def _diverse_decode_pad_embeddings( + self, model, active_embeds: torch.Tensor, actual_batch: int + ) -> torch.Tensor: + """Diagnostic: build a full [decode_batch, hidden] embedding tensor whose + padding rows carry distinct real token embeddings, so the decode MoE gate + routes the 128 tokens across many experts instead of all padding rows + mirroring slot 0. Active rows keep their real embeddings.""" + layout = self._compiled.layout + embed = getattr(self, "_diverse_embed_cache", None) + if embed is None: + embed = self._compiled.weight_store.load_tensor("embed.weight").contiguous() + self._diverse_embed_cache = embed + vocab = int(model.config.vocab_size) + hidden = int(embed.shape[1]) + full = torch.zeros((layout.decode_batch, hidden), dtype=active_embeds.dtype) + full[:actual_batch].copy_(active_embeds[:actual_batch].to(full.dtype)) + # Distinct, spread-out, non-special token ids for the padding rows. + pad_ids = [ + max(100, (1000 + row * 2659) % vocab) + for row in range(actual_batch, layout.decode_batch) + ] + pad_embed = embed.index_select(0, torch.tensor(pad_ids, dtype=torch.long)).to(full.dtype) + full[actual_batch:].copy_(pad_embed) + return full + + def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> DeviceTensor: + raise NotImplementedError("DeepSeekV4 uses model-specific cache pools, not generic KV tensors") + + def _free_kv_cache_tensor(self, tensor: DeviceTensor) -> None: + return None + + def run_prefill(self, model, batch: PrefillBatch) -> PrefillResult: + """Run all DeepSeekV4 hidden layers for one prefill chunk in a single packed call.""" + if self._compiled.prefill is None: + raise RuntimeError("DeepSeekV4 kernels were not compiled for this runner") + self._ensure_l3_shared_buffers(model) + inputs = self.prepare_prefill_inputs(model, batch) + if inputs.slot != 0: + raise RuntimeError( + "DeepSeekV4 prefill currently supports the first active serving slot only. " + "Run with one concurrent request until pypto-lib exposes a 64-slot prefill kernel." + ) + self._stage_prefill_fwd_inputs(inputs) + hidden_buffer = self._require_prefill_output_buffer(model.config.hidden_size) + hidden_buffer.zero_() + args = self._prefill_fwd_args(hidden_buffer) + self._debug_prefill_dispatch(inputs, args) + if os.environ.get("PYPTO_DSV4_SKIP_PREFILL_KERNEL") == "1": + # Diagnostic: skip the prefill kernel dispatch to isolate the decode + # deadlock from prefill device/ring state. All host-side prep above + # ran; we only skip the device kernel. Snapshot the (un-run) packed + # caches so decode can proceed, and force the first token to " a" + # (id 260) so decode runs on a realistic input. + self._snapshot_prefill_fwd_caches(inputs.slot) + forced = torch.zeros((1, int(model.config.vocab_size)), dtype=torch.float32) + forced[0, 260] = 1.0e4 + return PrefillResult(last_hidden=None, logits=forced) + try: + self._run_l3( + self._require_prefill_callable(), + *args, + self._int32_scalar(self._prefill_kernel_tokens(inputs.actual_tokens)), + ) + except RuntimeError as exc: + raise RuntimeError( + "DeepSeekV4 packed prefill dispatch failed " + f"(tokens={inputs.actual_tokens}, slot={inputs.slot})" + ) from exc + self._snapshot_prefill_fwd_caches(inputs.slot) + self._decode_cache_seeded_slots.clear() + + active_hidden = hidden_buffer[:, : inputs.actual_tokens, :] + self._debug_tensor_stats("prefill.output.hidden.active", active_hidden, per_rank=True) + if self._debug_tensor_stats_enabled() and not self._tensor_is_finite(active_hidden): + raise RuntimeError("DeepSeekV4 packed prefill produced non-finite active hidden rows") + + # Sample the last real prompt row (``actual_tokens - 1``) from host-side + # LM-head logits, mirroring the decode path. + last_row = inputs.actual_tokens - 1 + logits = self._logits_for_hidden(hidden_buffer, active_rows=(last_row,), label="prefill").float() + return PrefillResult(last_hidden=None, logits=logits) + + def run_decode(self, model, batch: DecodeBatch) -> DecodeResult: + """Run all DeepSeekV4 hidden layers for one decode batch in a single packed call.""" + if self._compiled.decode is None: + raise RuntimeError("DeepSeekV4 kernels were not compiled for this runner") + self._ensure_l3_shared_buffers(model) + inputs = self._stage_decode_inputs(self.prepare_decode_inputs(model, batch)) + if inputs.actual_batch != 1 or inputs.slots != (0,): + raise RuntimeError( + "DeepSeekV4 decode currently supports the first active serving slot only. " + "Run with one concurrent request until the compact cache handoff supports multiple slots." + ) + self._require_prefill_cache_snapshots() + self._seed_decode_work_cache(inputs.kernel_slots) + decode_buffers = self._require_decode_buffers() + x_hc = decode_buffers.x_hc_a + active_decode_tokens = inputs.actual_batch * self._compiled.layout.decode_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) + hidden_buffer.zero_() + # ``num_tokens`` is the real active token count. PR 677 restores the + # gate norm/quant scopes for this num_tokens-aware path; the fixed padding + # rows remain valid metadata for attention but must not be routed by MoE. + num_tokens = active_decode_tokens + args = self._decode_fwd_args(inputs, x_hc, hidden_buffer) + self._debug_decode_dispatch(inputs, args) + try: + self._run_l3( + self._require_decode_callable(), + *args, + self._int32_scalar(num_tokens), + ) + except RuntimeError as exc: + raise RuntimeError( + "DeepSeekV4 packed decode dispatch failed " + f"(actual_batch={inputs.actual_batch}, slots={inputs.slots})" + ) from exc + active_hidden = hidden_buffer[:, :active_decode_tokens, :] + self._debug_tensor_stats("decode.output.hidden.active", active_hidden, per_rank=True) + if self._debug_tensor_stats_enabled() and not self._tensor_is_finite(active_hidden): + raise RuntimeError("DeepSeekV4 packed decode produced non-finite active hidden rows") + + # Sample the final MTP slot (position seq_len-1), which predicts the next + # token: row r's sampled slot is ``r * decode_seq + (decode_seq - 1)``. + decode_seq = self._compiled.layout.decode_seq + active_rows = tuple(row * decode_seq + (decode_seq - 1) for row in range(inputs.actual_batch)) + logits = self._logits_for_hidden(hidden_buffer, active_rows=active_rows, label="decode").float() + return DecodeResult(hidden_states=None, logits=logits) + + def _require_prefill_callable(self) -> DeepSeekV4L3Callable: + if self._compiled.prefill is None: + raise RuntimeError("DeepSeekV4 prefill kernel is not compiled") + return self._compiled.prefill + + def _require_decode_callable(self) -> DeepSeekV4L3Callable: + if self._compiled.decode is None: + raise RuntimeError("DeepSeekV4 decode kernel is not compiled") + return self._compiled.decode + + def _ensure_l3_shared_buffers(self, model: RuntimeModel) -> None: + """Allocate every CPU tensor visible to the L3 worker before it forks. + + ``DistributedWorker`` creates per-chip children on first use. Any CPU + tensor argument those children access must already live in shared memory + at that point, so this method stages all packed prefill/decode input and + output buffers before the first ``_run_l3`` call. + """ + self.load_packed_global_weights() + self._static_freqs_cos_tensor() + self._static_freqs_sin_tensor() + self._ensure_decode_buffers(model.config.hidden_size) + self._ensure_decode_work_cache() + self._require_prefill_output_buffer(model.config.hidden_size) + self._static_final_norm_weight_tensor() + if self._stacked_weight_buffers is None: + self._stage_stacked_weights(self.load_stacked_layer_weights()) + self._hc_head_tensors() + self._ensure_prefill_fwd_buffers(model.config.hidden_size) + self._assert_l3_shared_buffers_preallocated() + + def _assert_l3_shared_buffers_preallocated(self) -> None: + missing = self._missing_l3_shared_buffers() + if missing: + raise RuntimeError( + "DeepSeekV4 L3 worker cannot start before all shared host buffers are preallocated; " + "missing: " + ", ".join(missing) + ) + + def _missing_l3_shared_buffers(self) -> list[str]: + missing: list[str] = [] + expected = { + "final_norm_w": self._static_final_norm_weight, + "freqs_cos": self._static_freqs_cos, + "freqs_sin": self._static_freqs_sin, + "prefill_fwd_buffers": self._prefill_fwd_buffers, + "decode_buffers": self._decode_buffers, + "decode_work_cache": self._decode_work_cache, + "stacked_weight_buffers": self._stacked_weight_buffers, + "hc_head_buffers": self._hc_head_buffers, + "prefill_output": self._prefill_output_buffer, + } + for name, value in expected.items(): + if value is None: + missing.append(name) + if self._stacked_weight_buffers is not None and not self._stacked_weight_buffers: + missing.append("stacked_weight_buffers") + if self._hc_head_buffers is not None and not self._hc_head_buffers: + missing.append("hc_head_buffers") + return missing + + def _prefill_fwd_args(self, x_out: torch.Tensor) -> tuple[Any, ...]: + """Build the single packed ``l3_prefill_fwd`` argument tuple. + + The kernel runs final RMSNorm and emits normalized hidden rows. LM-head is + computed on the host from the selected rows. + """ + buffers = self._require_prefill_fwd_buffers() + stacked = self._require_stacked_weights() + hc_head = self._hc_head_tensors() + values = dict(stacked.tensors) + values.update( + { + "x_hc": buffers.x_hc, + "freqs_cos": buffers.freqs_cos, + "freqs_sin": buffers.freqs_sin, + "hc_head_fn": hc_head["hc_head_fn"], + "hc_head_scale": hc_head["hc_head_scale"], + "hc_head_base": hc_head["hc_head_base"], + "final_norm_w": self._static_final_norm_weight_tensor(), + "x_out": x_out, + } + ) + values.update(buffers.tensors) + return self._ordered_layer_args(values, _PREFILL_FWD_TENSOR_ORDER) + + def _decode_fwd_args( + self, + inputs: DeepSeekV4PreparedDecodeInputs, + x_hc: torch.Tensor, + x_out: torch.Tensor, + ) -> tuple[Any, ...]: + """Build the single packed ``l3_decode_fwd`` argument tuple.""" + cache = self._require_decode_work_cache() + stacked = self._require_stacked_weights() + hc_head = self._hc_head_tensors() + values = dict(stacked.tensors) + values.update( + { + "x_hc": x_hc, + "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, + "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, + "hca_compress_state_block_table": inputs.hca_compress_state_block_table, + "csa_compress_state": cache.csa_compress_state, + "csa_compress_state_block_table": inputs.csa_compress_state_block_table, + "csa_inner_compress_state": cache.csa_inner_compress_state, + "csa_inner_compress_state_block_table": inputs.csa_inner_compress_state_block_table, + "cmp_kv": cache.cmp_kv, + "cmp_block_table": inputs.cmp_block_table, + "idx_kv_cache": cache.idx_kv_cache, + "idx_block_table": inputs.idx_block_table, + "input_ids": inputs.input_ids, + "hc_head_fn": hc_head["hc_head_fn"], + "hc_head_scale": hc_head["hc_head_scale"], + "hc_head_base": hc_head["hc_head_base"], + "final_norm_w": self._static_final_norm_weight_tensor(), + "x_out": x_out, + } + ) + return self._ordered_layer_args(values, _DECODE_FWD_TENSOR_ORDER) + + def _require_stacked_weights(self) -> DeepSeekV4StackedLayerWeights: + if self._stacked_weight_buffers is None: + raise RuntimeError("DeepSeekV4 stacked decode weights were not staged") + return DeepSeekV4StackedLayerWeights(tensors=self._stacked_weight_buffers) + + def _ordered_layer_args(self, values: dict[str, Any], names: Sequence[str]) -> tuple[Any, ...]: + missing = [name for name in names if name not in values] + if missing: + raise KeyError(f"DeepSeekV4 layer dispatch is missing tensors: {', '.join(missing)}") + return tuple(values[name] for name in names) + + def _debug_prefill_dispatch( + self, + inputs: DeepSeekV4PreparedPrefillInputs, + args: Sequence[Any], + ) -> None: + if os.getenv("PYPTO_DSV4_DEBUG") != "1": + return + named_args = dict(zip(_PREFILL_FWD_TENSOR_ORDER, args, strict=True)) + interesting = ( + "x_hc", + "kv_cache", + "cmp_kv", + "idx_kv_cache", + "ori_block_table", + "cmp_block_table", + "idx_block_table", + "cmp_sparse_indices", + "cmp_sparse_lens", + "input_ids", + "x_out", + ) + tensor_names = [ + name + for name, tensor in named_args.items() + if isinstance(tensor, torch.Tensor) and tensor.device.type == "cpu" + ] + non_shared = [name for name in tensor_names if not named_args[name].is_shared()] + parts = [] + for name in interesting: + tensor = named_args[name] + if isinstance(tensor, torch.Tensor): + parts.append(f"{name}={tuple(tensor.shape)}/{tensor.dtype}/shared={tensor.is_shared()}") + elif isinstance(tensor, DeviceTensor): + parts.append(f"{name}=DeviceTensor") + else: + parts.append(f"{name}={type(tensor).__name__}") + print( + "DeepSeekV4 packed prefill dispatch " + f"tokens={inputs.actual_tokens} slot={inputs.slot} " + f"worker_started={self._l3_worker is not None} " + f"cpu_tensor_args={len(tensor_names)} non_shared={non_shared} " + + " ".join(parts), + flush=True, + ) + if os.getenv("PYPTO_DSV4_DEBUG_ARGS") == "1": + for name in _PREFILL_FWD_TENSOR_ORDER: + tensor = named_args[name] + if isinstance(tensor, torch.Tensor): + print( + "DeepSeekV4 prefill arg " + f"{name}: shape={tuple(tensor.shape)} dtype={tensor.dtype} " + f"device={tensor.device} shared={tensor.is_shared()}", + flush=True, + ) + + def _debug_decode_dispatch( + self, + inputs: DeepSeekV4PreparedDecodeInputs, + args: Sequence[Any], + ) -> None: + if os.getenv("PYPTO_DSV4_DEBUG") != "1": + return + named_args = dict(zip(_DECODE_FWD_TENSOR_ORDER, args, strict=True)) + interesting = ( + "x_hc", + "kv_cache", + "block_table", + "ori_slot_mapping", + "cmp_kv", + "cmp_block_table", + "idx_kv_cache", + "idx_block_table", + "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", + "x_out", + ) + tensor_names = [ + name + for name, tensor in named_args.items() + if isinstance(tensor, torch.Tensor) and tensor.device.type == "cpu" + ] + non_shared = [name for name in tensor_names if not named_args[name].is_shared()] + parts = [] + for name in interesting: + tensor = named_args[name] + if isinstance(tensor, torch.Tensor): + parts.append(f"{name}={tuple(tensor.shape)}/{tensor.dtype}/shared={tensor.is_shared()}") + elif isinstance(tensor, DeviceTensor): + parts.append(f"{name}=DeviceTensor") + else: + parts.append(f"{name}={type(tensor).__name__}") + print( + "DeepSeekV4 packed decode dispatch " + f"actual_batch={inputs.actual_batch} active_tokens={inputs.actual_batch * self._compiled.layout.decode_seq} " + f"slots={inputs.slots} " + f"worker_started={self._l3_worker is not None} " + f"cpu_tensor_args={len(tensor_names)} non_shared={non_shared} " + + " ".join(parts), + flush=True, + ) + if os.getenv("PYPTO_DSV4_DEBUG_ARGS") == "1": + for name in _DECODE_FWD_TENSOR_ORDER: + tensor = named_args[name] + if isinstance(tensor, torch.Tensor): + print( + "DeepSeekV4 decode arg " + f"{name}: shape={tuple(tensor.shape)} dtype={tensor.dtype} " + f"device={tensor.device} shared={tensor.is_shared()}", + flush=True, + ) + self._debug_tensor_stats(f"dispatch.fwd.{name}", tensor) + + @staticmethod + def _is_layer_weight_name(name: str) -> bool: + runtime_names = { + "x_hc", + "freqs_cos", + "freqs_sin", + "hca_cmp_kv_state", + "hca_cmp_score_state", + "hca_compress_state_block_table", + "csa_cmp_kv_state", + "csa_cmp_score_state", + "csa_compress_state_block_table", + "csa_inner_kv_state", + "csa_inner_score_state", + "csa_inner_compress_state_block_table", + "kv_cache", + "ori_block_table", + "block_table", + "ori_slot_mapping", + "cmp_kv", + "cmp_block_table", + "cmp_sparse_indices", + "cmp_sparse_lens", + "idx_kv_cache", + "idx_block_table", + "position_ids", + "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", + "kv_seq_lens", + "input_ids", + "x_next", + } + return name not in runtime_names + + def _ensure_decode_buffers(self, hidden_size: int) -> _DeepSeekV4DecodeSharedBuffers: + buffers = self._decode_buffers + if buffers is None: + self._ensure_shared_host_allocation_before_worker("decode inputs") + layout = self._compiled.layout + ranks = layout.ranks + 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.bfloat16, + name="decode_x_hc", + ), + x_hc_b=self._shared_empty( + (ranks, tokens, layout.hc_mult, int(hidden_size)), + torch.bfloat16, + name="decode_x_hc_next", + ), + x_out=self._shared_empty( + (ranks, tokens, int(hidden_size)), + torch.bfloat16, + name="decode_x_out", + ), + 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"), + "kv_seq_lens": self._shared_empty((ranks, batch), torch.int32, name="decode_kv_seq_lens"), + "block_table": self._shared_empty( + (ranks, batch, layout.ori_max_blocks), + torch.int32, + name="decode_block_table", + ), + "ori_slot_mapping": self._shared_empty( + (ranks, tokens), + torch.long, + name="decode_ori_slot_mapping", + ), + "cmp_block_table": self._shared_empty( + (ranks, batch, layout.cmp_max_blocks), + torch.int32, + name="decode_cmp_block_table", + ), + "idx_block_table": self._shared_empty( + (ranks, batch, layout.idx_max_blocks), + torch.int32, + name="decode_idx_block_table", + ), + "hca_compress_state_block_table": self._shared_empty( + (ranks, batch, layout.hca_state_max_blocks), + torch.int32, + name="decode_hca_compress_state_block_table", + ), + "csa_compress_state_block_table": self._shared_empty( + (ranks, batch, layout.csa_state_max_blocks), + torch.int32, + name="decode_csa_compress_state_block_table", + ), + "csa_inner_compress_state_block_table": self._shared_empty( + (ranks, batch, layout.csa_inner_state_max_blocks), + 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", + ), + }, + ) + self._decode_buffers = buffers + return buffers + + def _stage_decode_inputs(self, inputs: DeepSeekV4PreparedDecodeInputs) -> DeepSeekV4PreparedDecodeInputs: + buffers = self._ensure_decode_buffers(inputs.x_hc.shape[-1]) + self._copy_shared(buffers.x_hc_a, inputs.x_hc, name="decode_x_hc") + staged_values: dict[str, torch.Tensor] = {} + 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) + + def _ensure_prefill_fwd_buffers(self, hidden_size: int) -> _DeepSeekV4PrefillFwdSharedBuffers: + """Allocate the layer-stacked shared buffers for the packed prefill dispatch.""" + buffers = self._prefill_fwd_buffers + if buffers is not None: + return buffers + self._ensure_shared_host_allocation_before_worker("prefill_fwd buffers") + layout = self._compiled.layout + ranks = layout.ranks + seq = layout.prefill_seq + hidden = int(hidden_size) + fwd = DEEPSEEK_V4_FWD_NUM_LAYERS + csa = DEEPSEEK_V4_CSA_NUM_LAYERS + hca = DEEPSEEK_V4_HCA_NUM_LAYERS + rope_dim = self._compiled.freqs_cos.shape[-1] if self._compiled.freqs_cos is not None else 0 + max_seq_len = self._compiled.freqs_cos.shape[0] if self._compiled.freqs_cos is not None else 0 + + def shared(shape, dtype, name): + return self._shared_empty(shape, dtype, name=name) + + tensors: dict[str, torch.Tensor] = { + # HCA-group prefill compressor state (x20). + "hca_cmp_kv_state": shared( + (ranks, hca * layout.prefill_hca_state_max_blocks, layout.c128_state_block_size, DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + torch.float32, + "prefill_fwd_hca_cmp_kv_state", + ), + "hca_cmp_score_state": shared( + (ranks, hca * layout.prefill_hca_state_max_blocks, layout.c128_state_block_size, DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + torch.float32, + "prefill_fwd_hca_cmp_score_state", + ), + "hca_compress_state_block_table": shared( + (ranks, layout.prefill_hca_state_max_blocks), torch.int32, "prefill_fwd_hca_state_block_table" + ), + # CSA-group prefill compressor state (x21). + "csa_cmp_kv_state": shared( + (ranks, csa * layout.prefill_csa_state_max_blocks, layout.c4_state_block_size, DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + torch.float32, + "prefill_fwd_csa_cmp_kv_state", + ), + "csa_cmp_score_state": shared( + (ranks, csa * layout.prefill_csa_state_max_blocks, layout.c4_state_block_size, DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + torch.float32, + "prefill_fwd_csa_cmp_score_state", + ), + "csa_compress_state_block_table": shared( + (ranks, layout.prefill_csa_state_max_blocks), torch.int32, "prefill_fwd_csa_state_block_table" + ), + "csa_inner_kv_state": shared( + (ranks, csa * layout.prefill_csa_inner_state_max_blocks, layout.c4_state_block_size, DEEPSEEK_V4_CSA_INNER_OUT_DIM), + torch.float32, + "prefill_fwd_csa_inner_kv_state", + ), + "csa_inner_score_state": shared( + (ranks, csa * layout.prefill_csa_inner_state_max_blocks, layout.c4_state_block_size, DEEPSEEK_V4_CSA_INNER_OUT_DIM), + torch.float32, + "prefill_fwd_csa_inner_score_state", + ), + "csa_inner_compress_state_block_table": shared( + (ranks, layout.prefill_csa_inner_state_max_blocks), torch.int32, "prefill_fwd_csa_inner_state_block_table" + ), + # Work caches: kv_cache/cmp_kv stack x43, idx_kv_cache stacks x21 (CSA), + # all flattened 5-D (the kernel reshapes the fused layer x block axis). + "kv_cache": shared( + (ranks, fwd * layout.ori_max_blocks, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), + torch.bfloat16, + "prefill_fwd_kv_cache", + ), + "cmp_kv": shared( + (ranks, fwd * layout.prefill_cmp_block_num, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), + torch.bfloat16, + "prefill_fwd_cmp_kv", + ), + "idx_kv_cache": shared( + (ranks, csa * layout.prefill_idx_block_num, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + torch.bfloat16, + "prefill_fwd_idx_kv_cache", + ), + # Shared single per-rank metadata (the kernel passes each whole tensor + # to every layer). + "ori_block_table": shared((ranks, layout.ori_max_blocks), torch.int32, "prefill_fwd_ori_block_table"), + "ori_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_ori_slot_mapping"), + "cmp_block_table": shared((ranks, layout.prefill_cmp_max_blocks), torch.int32, "prefill_fwd_cmp_block_table"), + "cmp_sparse_indices": shared((ranks, seq, layout.prefill_sparse_topk), torch.int32, "prefill_fwd_cmp_sparse_indices"), + "cmp_sparse_lens": shared((ranks, seq), torch.int32, "prefill_fwd_cmp_sparse_lens"), + "idx_block_table": shared((ranks, layout.prefill_idx_max_blocks), torch.int32, "prefill_fwd_idx_block_table"), + "position_ids": shared((ranks, seq), torch.int32, "prefill_fwd_position_ids"), + "hca_cmp_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_hca_cmp_slot_mapping"), + "hca_state_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_hca_state_slot_mapping"), + "csa_cmp_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_csa_cmp_slot_mapping"), + "csa_idx_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_csa_idx_slot_mapping"), + "csa_state_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_csa_state_slot_mapping"), + "csa_inner_state_slot_mapping": shared((ranks, seq), torch.long, "prefill_fwd_csa_inner_state_slot_mapping"), + "input_ids": shared((ranks, seq), torch.long, "prefill_fwd_input_ids"), + } + buffers = _DeepSeekV4PrefillFwdSharedBuffers( + x_hc=shared((ranks, seq, layout.hc_mult, hidden), torch.bfloat16, "prefill_fwd_x_hc"), + freqs_cos=shared((ranks, max_seq_len, rope_dim), torch.bfloat16, "prefill_fwd_freqs_cos"), + freqs_sin=shared((ranks, max_seq_len, rope_dim), torch.bfloat16, "prefill_fwd_freqs_sin"), + tensors=tensors, + ) + self._prefill_fwd_buffers = buffers + return buffers + + def _require_prefill_fwd_buffers(self) -> _DeepSeekV4PrefillFwdSharedBuffers: + if self._prefill_fwd_buffers is None: + raise RuntimeError("DeepSeekV4 packed prefill shared buffers were not staged") + return self._prefill_fwd_buffers + + def _stage_prefill_fwd_inputs(self, inputs: DeepSeekV4PreparedPrefillInputs) -> None: + """Copy one prefill chunk's metadata/state into the packed buffers. + + The per-request metadata (slot mappings, block tables, sparse tables, + position/input ids), the RoPE tables and the compressor-state block tables + are shared single per-rank copies (the kernel slices them per layer + internally). The compressor-state and work caches start zeroed and are + produced by the kernel. + """ + buffers = self._require_prefill_fwd_buffers() + + # x_hc / output collapse weights. + self._copy_shared(buffers.x_hc, inputs.x_hc, name="prefill_fwd_x_hc") + self._copy_shared( + buffers.freqs_cos, + self._static_freqs_cos_table(), + name="prefill_fwd_freqs_cos", + ) + self._copy_shared( + buffers.freqs_sin, + self._static_freqs_sin_table(), + name="prefill_fwd_freqs_sin", + ) + + # Shared single per-rank metadata (the kernel slices it per layer). + shared_metadata = { + "ori_block_table": inputs.ori_block_table, + "ori_slot_mapping": inputs.ori_slot_mapping, + "cmp_block_table": inputs.cmp_block_table, + "idx_block_table": inputs.idx_block_table, + "position_ids": inputs.position_ids, + "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, + "input_ids": inputs.input_ids, + "hca_compress_state_block_table": inputs.hca_compress_state_block_table, + "csa_compress_state_block_table": inputs.csa_compress_state_block_table, + "csa_inner_compress_state_block_table": inputs.csa_inner_compress_state_block_table, + } + for name, tensor in shared_metadata.items(): + self._copy_shared(buffers.tensors[name], tensor, name=f"prefill_fwd_{name}") + + # Sparse tables are now a single shared copy used by every layer. The kernel + # consumes the sliding-window index set (the ratio-0 view), so pass that copy. + sparse_indices, sparse_lens = inputs.sparse_inputs_for_ratio(0) + self._copy_shared( + buffers.tensors["cmp_sparse_indices"], + sparse_indices, + name="prefill_fwd_cmp_sparse_indices", + ) + self._copy_shared( + buffers.tensors["cmp_sparse_lens"], + sparse_lens, + name="prefill_fwd_cmp_sparse_lens", + ) + + # Compressor-state and work caches start zeroed; the kernel populates them. + for name in ( + "hca_cmp_kv_state", + "hca_cmp_score_state", + "csa_cmp_kv_state", + "csa_cmp_score_state", + "csa_inner_kv_state", + "csa_inner_score_state", + "kv_cache", + "cmp_kv", + "idx_kv_cache", + ): + buffers.tensors[name].zero_() + + def _static_freqs_cos_table(self) -> torch.Tensor: + if self._compiled.freqs_cos is None: + raise RuntimeError("DeepSeekV4 RoPE cosine table is not initialized") + return self._rank_stack(self._compiled.freqs_cos) + + def _static_freqs_sin_table(self) -> torch.Tensor: + if self._compiled.freqs_sin is None: + raise RuntimeError("DeepSeekV4 RoPE sine table is not initialized") + return self._rank_stack(self._compiled.freqs_sin) + + def _snapshot_prefill_fwd_caches(self, slot: int) -> None: + """Capture per-layer cache slices from the packed prefill Out caches.""" + buffers = self._require_prefill_fwd_buffers() + csa_order = 0 + hca_order = 0 + layout = self._compiled.layout + for layer in self._compiled.layer_plan: + fwd = int(layer.layer_id) + # The flattened 5-D work caches fuse the layer axis with the per-layer + # block axis; slice the contiguous per-layer block span back out. + snapshot: dict[str, torch.Tensor] = { + "kv_cache": self._slice_layer_state(buffers.tensors["kv_cache"], fwd, layout.ori_max_blocks), + "cmp_kv": self._slice_layer_slot_state( + buffers.tensors["cmp_kv"], + fwd, + layout.prefill_cmp_block_num, + slot, + layout.prefill_cmp_max_blocks, + ), + } + if layer.compress_ratio == 4: + # idx_kv_cache stacks across the CSA layers (x21), so index by csa_order. + snapshot["idx_kv_cache"] = self._slice_layer_slot_state( + buffers.tensors["idx_kv_cache"], + csa_order, + layout.prefill_idx_block_num, + slot, + layout.prefill_idx_max_blocks, + ) + snapshot["csa_cmp_kv_state"] = self._slice_layer_state( + buffers.tensors["csa_cmp_kv_state"], csa_order, self._compiled.layout.prefill_csa_state_max_blocks + ) + snapshot["csa_cmp_score_state"] = self._slice_layer_state( + buffers.tensors["csa_cmp_score_state"], csa_order, self._compiled.layout.prefill_csa_state_max_blocks + ) + snapshot["csa_inner_kv_state"] = self._slice_layer_state( + buffers.tensors["csa_inner_kv_state"], csa_order, self._compiled.layout.prefill_csa_inner_state_max_blocks + ) + snapshot["csa_inner_score_state"] = self._slice_layer_state( + buffers.tensors["csa_inner_score_state"], csa_order, self._compiled.layout.prefill_csa_inner_state_max_blocks + ) + csa_order += 1 + elif layer.compress_ratio == 128: + snapshot["hca_cmp_kv_state"] = self._slice_layer_state( + buffers.tensors["hca_cmp_kv_state"], hca_order, self._compiled.layout.prefill_hca_state_max_blocks + ) + snapshot["hca_cmp_score_state"] = self._slice_layer_state( + buffers.tensors["hca_cmp_score_state"], hca_order, self._compiled.layout.prefill_hca_state_max_blocks + ) + hca_order += 1 + self._prefill_cache_snapshots[layer.layer_id] = DeepSeekV4LayerCacheSnapshot(snapshot) + + @staticmethod + def _slice_layer_state(stacked: torch.Tensor, order: int, blocks_per_layer: int) -> torch.Tensor: + start = int(order) * int(blocks_per_layer) + return stacked[:, start : start + int(blocks_per_layer)].detach().cpu().contiguous().clone() + + @staticmethod + def _slice_layer_slot_state( + stacked: torch.Tensor, + order: int, + blocks_per_layer: int, + slot: int, + blocks_per_slot: int, + ) -> torch.Tensor: + start = int(order) * int(blocks_per_layer) + layer = stacked[:, start : start + int(blocks_per_layer)] + slot_slice = DeepSeekV4ModelRunner._slot_block_slice(int(slot), int(blocks_per_slot)) + return layer[:, slot_slice].detach().cpu().contiguous().clone() + + def _stage_stacked_weights(self, weights: DeepSeekV4StackedLayerWeights) -> DeepSeekV4StackedLayerWeights: + """Copy the layer-stacked decode_fwd weights into shared buffers once.""" + buffers = self._stacked_weight_buffers + if buffers is None: + self._ensure_shared_host_allocation_before_worker("stacked layer weights") + buffers = { + name: self._new_shared_like(tensor, name=f"stacked_weight[{name}]") + for name, tensor in weights.tensors.items() + } + self._stacked_weight_buffers = buffers + + missing = sorted(set(weights.tensors) - set(buffers)) + if missing: + raise KeyError(f"DeepSeekV4 shared stacked-weight buffers are missing: {', '.join(missing)}") + + for name, tensor in weights.tensors.items(): + self._copy_shared(buffers[name], tensor, name=f"stacked_weight[{name}]") + return DeepSeekV4StackedLayerWeights(tensors=buffers) + + def _hc_head_tensors(self) -> dict[str, torch.Tensor]: + """Return rank-replicated hc_head weights for the decode_fwd output collapse.""" + buffers = self._hc_head_buffers + if buffers is not None: + return buffers + self._ensure_shared_host_allocation_before_worker("hc_head weights") + global_weights = self.load_packed_global_weights() + ranks = self._compiled.layout.ranks + # The kernel hc_head_fn is [HC_MULT, HC_DIM]; the checkpoint stores it as + # [HC_MULT, hidden*HC_MULT] (== [HC_MULT, HC_DIM]). Scale/base are scalars + # per HC_MULT row, rank-replicated. + hc_head_fn = global_weights.hc_head_fn.to(torch.float32).contiguous().cpu() + hc_head_scale = global_weights.hc_head_scale.to(torch.float32).contiguous().cpu() + hc_head_base = global_weights.hc_head_base.to(torch.float32).contiguous().cpu() + buffers = { + "hc_head_fn": self._static_device_tensor(self._rank_stack(hc_head_fn)), + "hc_head_scale": self._static_device_tensor(self._rank_stack(hc_head_scale)), + "hc_head_base": self._static_device_tensor(self._rank_stack(hc_head_base)), + } + self._hc_head_buffers = buffers + return buffers + + def _require_decode_buffers(self) -> _DeepSeekV4DecodeSharedBuffers: + if self._decode_buffers is None: + raise RuntimeError("DeepSeekV4 decode shared buffers were not staged") + return self._decode_buffers + + def _require_decode_output_buffer(self, hidden_size: int) -> torch.Tensor: + return self._ensure_decode_buffers(int(hidden_size)).x_out + + def _require_decode_logits_buffer(self, vocab_size: int) -> torch.Tensor: + """Return a legacy shared ``[ranks, decode_tokens, vocab]`` logits buffer.""" + layout = self._compiled.layout + logits_shape = (layout.ranks, layout.decode_tokens, int(vocab_size)) + if self._decode_logits_buffer is None: + self._ensure_shared_host_allocation_before_worker("decode_logits") + self._decode_logits_buffer = self._shared_empty(logits_shape, torch.float32, name="decode_logits") + return self._decode_logits_buffer + + def _require_prefill_output_buffer(self, hidden_size: int) -> torch.Tensor: + """Return the shared ``[ranks, prefill_seq, hidden]`` prefill hidden output.""" + layout = self._compiled.layout + output_shape = (layout.ranks, layout.prefill_seq, int(hidden_size)) + if self._prefill_output_buffer is None: + self._ensure_shared_host_allocation_before_worker("prefill_output") + self._prefill_output_buffer = self._shared_empty(output_shape, torch.bfloat16, name="prefill_output") + return self._prefill_output_buffer + + def _static_final_norm_weight_tensor(self) -> torch.Tensor: + """Return the worker-resident per-rank final RMSNorm weight ``[ranks, D]``. + + Reuses the same ``final_norm_weight`` already loaded for the host-side + ``_final_norm`` collapse, rank-replicated and cast to bf16 for the kernel. + """ + if self._static_final_norm_weight is None: + global_weights = self.load_packed_global_weights() + self._ensure_shared_host_allocation_before_worker("final_norm_w") + final_norm_w = global_weights.final_norm_weight.to(torch.bfloat16).contiguous().cpu() + self._static_final_norm_weight = self._static_device_tensor(self._rank_stack(final_norm_w)) + return self._static_final_norm_weight + + def _static_freqs_cos_tensor(self) -> torch.Tensor: + if self._static_freqs_cos is None: + if self._compiled.freqs_cos is None: + raise RuntimeError("DeepSeekV4 RoPE cosine table is not initialized") + self._ensure_shared_host_allocation_before_worker("freqs_cos") + self._static_freqs_cos = self._static_device_tensor(self._rank_stack(self._compiled.freqs_cos)) + return self._static_freqs_cos + + def _static_freqs_sin_tensor(self) -> torch.Tensor: + if self._static_freqs_sin is None: + if self._compiled.freqs_sin is None: + raise RuntimeError("DeepSeekV4 RoPE sine table is not initialized") + self._ensure_shared_host_allocation_before_worker("freqs_sin") + self._static_freqs_sin = self._static_device_tensor(self._rank_stack(self._compiled.freqs_sin)) + return self._static_freqs_sin + + def _require_prefill_cache_snapshots(self) -> None: + missing = [ + str(layer.layer_id) + for layer in self._compiled.layer_plan + if layer.layer_id not in self._prefill_cache_snapshots + ] + if missing: + raise RuntimeError( + "DeepSeekV4 decode requires prefill cache snapshots before decode; " + "missing layers: " + ", ".join(missing) + ) + + def _seed_decode_work_cache(self, kernel_slots: Sequence[int]) -> None: + """Seed uninitialized decode cache slots from every prefill snapshot. + + The packed ``l3_decode_fwd`` kernel reads all 43 layers' KV/compress state + in one call, and then mutates those same buffers with generated-token + state. Seed each slot from the prefill snapshot once, then preserve the + decode-produced state across later decode steps. Each layer's blocks live + at a stacked offset on dim 1: FWD layers use the layer id (0..42), + CSA-group state uses the CSA order index (0..20), and HCA-group state uses + the HCA order index (0..19). Within a layer the slot offset is + ``layer_offset * decode_batch + slot``. + """ + slots_to_seed = tuple( + int(slot) + for slot in kernel_slots + if int(slot) not in self._decode_cache_seeded_slots + ) + if not slots_to_seed: + return + + cache = self._require_decode_work_cache() + layout = self._compiled.layout + batch = layout.decode_batch + + csa_order = 0 + hca_order = 0 + for layer in self._compiled.layer_plan: + snapshot = self._prefill_cache_snapshots.get(layer.layer_id) + if snapshot is None: + raise RuntimeError(f"DeepSeekV4 decode cache snapshot missing for layer {layer.layer_id}") + tensors = snapshot.tensors + fwd_offset = int(layer.layer_id) + for slot in slots_to_seed: + self._copy_snapshot_blocks_to_work( + tensors["kv_cache"], + cache.kv_cache, + fwd_offset * batch + int(slot), + layout.ori_max_blocks, + ) + self._copy_snapshot_blocks_to_work( + tensors["cmp_kv"], + cache.cmp_kv, + fwd_offset * batch + int(slot), + layout.cmp_max_blocks, + ) + if layer.compress_ratio == 4: + for slot in slots_to_seed: + self._copy_snapshot_blocks_to_work( + tensors["idx_kv_cache"], + cache.idx_kv_cache, + csa_order * batch + int(slot), + layout.idx_max_blocks, + ) + self._copy_split_state_to_work( + tensors["csa_cmp_kv_state"], + tensors["csa_cmp_score_state"], + cache.csa_compress_state, + csa_order * batch + int(slot), + layout.csa_state_max_blocks, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ) + self._copy_split_state_to_work( + tensors["csa_inner_kv_state"], + tensors["csa_inner_score_state"], + cache.csa_inner_compress_state, + csa_order * batch + int(slot), + layout.csa_inner_state_max_blocks, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ) + csa_order += 1 + elif layer.compress_ratio == 128: + for slot in slots_to_seed: + self._copy_split_state_to_work( + tensors["hca_cmp_kv_state"], + tensors["hca_cmp_score_state"], + cache.hca_compress_state, + hca_order * batch + int(slot), + layout.hca_state_max_blocks, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ) + hca_order += 1 + + if csa_order != DEEPSEEK_V4_CSA_NUM_LAYERS: + raise RuntimeError( + f"DeepSeekV4 decode expected {DEEPSEEK_V4_CSA_NUM_LAYERS} CSA layers, found {csa_order}" + ) + if hca_order != DEEPSEEK_V4_HCA_NUM_LAYERS: + raise RuntimeError( + f"DeepSeekV4 decode expected {DEEPSEEK_V4_HCA_NUM_LAYERS} HCA layers, found {hca_order}" + ) + self._decode_cache_seeded_slots.update(slots_to_seed) + + @staticmethod + def _slot_block_slice(slot: int, blocks_per_slot: int) -> slice: + if slot < 0: + raise ValueError("slot must be non-negative") + start = int(slot) * int(blocks_per_slot) + return slice(start, start + int(blocks_per_slot)) + + def _copy_snapshot_blocks_to_work( + self, + snapshot: torch.Tensor, + work: torch.Tensor, + slot: int, + blocks_per_slot: int, + ) -> None: + del self + slot_slice = DeepSeekV4ModelRunner._slot_block_slice(slot, blocks_per_slot) + dst = work[:, slot_slice] + dst.zero_() + blocks = min(snapshot.shape[1], int(blocks_per_slot)) + dst[:, :blocks].copy_(snapshot[:, :blocks]) + + def _copy_split_state_to_work( + self, + kv_state: torch.Tensor, + score_state: torch.Tensor, + work: torch.Tensor, + slot: int, + blocks_per_slot: int, + out_dim: int, + ) -> None: + del self + slot_slice = DeepSeekV4ModelRunner._slot_block_slice(slot, blocks_per_slot) + dst = work[:, slot_slice] + dst.zero_() + blocks = min(kv_state.shape[1], score_state.shape[1], int(blocks_per_slot)) + dst[:, :blocks, ..., :out_dim].copy_(kv_state[:, :blocks]) + dst[:, :blocks, ..., out_dim : 2 * out_dim].copy_(score_state[:, :blocks]) + + def _logits_for_hidden( + self, + x_hc: torch.Tensor, + *, + active_rows: Sequence[int], + label: str = "unknown", + ) -> torch.Tensor: + global_weights = self.load_packed_global_weights() + if x_hc.ndim == 3: + # Decode output is already collapsed and final-normalized by + # ``l3_decode_fwd``; host LM-head consumes it directly. + hidden = x_hc + else: + hidden = self._final_hidden(x_hc) + rows = tuple(int(row) for row in active_rows) + if not rows: + raise ValueError("DeepSeekV4 LM-head requires at least one active row") + if min(rows) < 0 or max(rows) >= hidden.shape[1]: + raise ValueError( + f"DeepSeekV4 LM-head active rows {rows} exceed hidden rows={hidden.shape[1]}" + ) + row_list = list(rows) + if self._debug_tensor_stats_enabled(): + print(f"DSV4_DEBUG lm_head.label={label} active_rows={rows}", flush=True) + if x_hc.ndim == 4: + self._debug_tensor_stats("lm_head.x_hc.active", x_hc[:, row_list, :, :]) + self._debug_tensor_stats("lm_head.hidden.active", hidden[:, row_list, :]) + + layout = global_weights.lm_head_layout + if global_weights.lm_head_weight.shape[0] != layout.ranks: + raise ValueError( + "DeepSeekV4 packed LM-head rank count mismatch: " + f"weight ranks={global_weights.lm_head_weight.shape[0]} layout ranks={layout.ranks}" + ) + if global_weights.lm_head_weight.shape[1] < layout.vocab_per_rank: + raise ValueError( + "DeepSeekV4 packed LM-head shard is smaller than the real vocab shard: " + f"shape={tuple(global_weights.lm_head_weight.shape)} vocab_per_rank={layout.vocab_per_rank}" + ) + + selected = hidden[0, row_list, :].detach().cpu().to(torch.float32).contiguous() + logits_parts = [] + for rank in range(layout.ranks): + shard = global_weights.lm_head_weight[rank, : layout.vocab_per_rank, :] + shard = shard.detach().cpu().to(torch.float32).contiguous() + logits_parts.append(torch.matmul(selected, shard.t())) + logits = torch.cat(logits_parts, dim=-1) + if logits.shape[-1] != layout.vocab_size: + logits = logits[:, : layout.vocab_size].contiguous() + else: + logits = logits.contiguous() + self._debug_tensor_stats("lm_head.logits.returned", logits) + return logits + + @staticmethod + def _debug_tensor_stats_enabled() -> bool: + return os.getenv("PYPTO_DSV4_LOGIT_DEBUG") == "1" + + @staticmethod + def _debug_tensor_stats(name: str, tensor: torch.Tensor, *, per_rank: bool = False) -> None: + if not DeepSeekV4ModelRunner._debug_tensor_stats_enabled(): + return + data = tensor.detach().cpu().to(torch.float32) + finite = torch.isfinite(data) + finite_count = int(finite.sum().item()) + total = data.numel() + nan_count = int(torch.isnan(data).sum().item()) + pos_inf_count = int(torch.isposinf(data).sum().item()) + neg_inf_count = int(torch.isneginf(data).sum().item()) + if finite_count: + finite_values = data[finite] + min_value = float(finite_values.min().item()) + max_value = float(finite_values.max().item()) + absmax_value = float(finite_values.abs().max().item()) + else: + min_value = float("nan") + max_value = float("nan") + absmax_value = float("nan") + print( + "DSV4_DEBUG " + f"{name} shape={tuple(tensor.shape)} dtype={tensor.dtype} " + f"finite={finite_count}/{total} nan={nan_count} " + f"+inf={pos_inf_count} -inf={neg_inf_count} " + f"min={min_value:.6g} max={max_value:.6g} absmax={absmax_value:.6g}", + flush=True, + ) + if per_rank and data.ndim >= 1: + rank_view = data.reshape(data.shape[0], -1) + rank_finite = torch.isfinite(rank_view) + rank_counts = (rank_view.shape[1] - rank_finite.sum(dim=1)).tolist() + print(f"DSV4_DEBUG {name} nonfinite_by_rank={rank_counts}", flush=True) + + @staticmethod + def _tensor_is_finite(tensor: torch.Tensor) -> bool: + return bool(torch.isfinite(tensor.detach().cpu().to(torch.float32)).all().item()) + + def _final_hidden(self, x_hc: torch.Tensor) -> torch.Tensor: + """Collapse a ``[ranks, T, HC_MULT, D]`` HC stack and apply the final norm.""" + weights = self.load_packed_global_weights() + x_hc = x_hc.to(torch.bfloat16).cpu() + x_float = x_hc.float() + flat = x_float.flatten(2) + rms = torch.sqrt(flat.double().square().mean(dim=-1, keepdim=True) + DEEPSEEK_V4_RMS_NORM_EPS) + normed_flat = flat / rms.to(torch.float32) + mixes = torch.matmul(normed_flat, weights.hc_head_fn.t()) + pre = torch.sigmoid(mixes * weights.hc_head_scale + weights.hc_head_base) + DEEPSEEK_V4_HC_EPS + collapsed = torch.sum(pre.unsqueeze(-1).double() * x_float.double(), dim=2) + return self._final_norm(collapsed) + + def _final_norm(self, collapsed: torch.Tensor) -> torch.Tensor: + """Apply the final RMS norm to an already-collapsed ``[ranks, T, D]`` hidden. + + The packed ``l3_decode_fwd`` kernel collapses HC_MULT in-kernel via + ``hc_head`` and returns the collapsed (pre-final-norm) hidden, so decode + only needs the model's final RMS norm before the LM head. + """ + collapsed = collapsed.cpu().double() + weights = self.load_packed_global_weights() + norm_inv = torch.rsqrt(collapsed.square().mean(dim=-1, keepdim=True) + DEEPSEEK_V4_RMS_NORM_EPS) + normed = collapsed * norm_inv * weights.final_norm_weight.double() + return normed.to(torch.float32).to(torch.bfloat16).contiguous() + + def _scope_stats_run_config(self) -> Any: + """Optional per-dispatch RunConfig that captures device scope stats. + + Enabled with ``PYPTO_DSV4_SCOPE_STATS=1`` to dump per-scope + heap / task_window / tensormap peaks under ``/dfx_outputs/``. + """ + if os.getenv("PYPTO_DSV4_SCOPE_STATS") != "1": + return None + from pypto.runtime import RunConfig # noqa: PLC0415 + + out_dir = os.getenv("PYPTO_DSV4_SCOPE_STATS_DIR", "/data/liuxu/pypto-serving/dsv4_scope_stats") + return RunConfig( + platform=self._compiled.platform, + device_id=self._compiled.device_id, + enable_scope_stats=True, + save_kernels=True, + save_kernels_dir=out_dir, + ) + + def _run_l3(self, callable_spec: DeepSeekV4L3Callable, *args: Any) -> Any: + if self._l3_worker is None: + self._assert_l3_args_shared_before_worker(callable_spec, args) + worker = self._shared_l3_worker() + run_config = self._scope_stats_run_config() + uploaded: list[DeviceTensor] = [] + try: + l3_args = tuple(self._coerce_l3_arg(worker, arg, uploaded) for arg in args) + if run_config is not None: + return worker.run(callable_spec.compiled, *l3_args, config=run_config) + return worker.run(callable_spec.compiled, *l3_args) + finally: + for tensor in uploaded: + worker.free_tensor(tensor) + + @staticmethod + def _share_cpu_tensor(tensor: torch.Tensor) -> torch.Tensor: + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + if not tensor.is_shared(): + tensor = tensor.share_memory_() + return tensor + + @staticmethod + def _shared_empty(shape: Sequence[int], dtype: torch.dtype, *, name: str) -> torch.Tensor: + del name + return torch.empty(tuple(int(dim) for dim in shape), dtype=dtype).share_memory_() + + @staticmethod + def _new_shared_like(tensor: torch.Tensor, *, name: str) -> torch.Tensor: + if tensor.device.type != "cpu": + raise ValueError(f"{name} must be a CPU tensor") + return torch.empty_like(tensor.contiguous(), memory_format=torch.contiguous_format).share_memory_() + + @staticmethod + def _copy_shared(dst: torch.Tensor, src: torch.Tensor, *, name: str) -> None: + if src.device.type != "cpu": + src = src.cpu() + if not src.is_contiguous(): + src = src.contiguous() + if tuple(dst.shape) != tuple(src.shape) or dst.dtype != src.dtype: + raise ValueError( + f"{name} shared buffer shape/dtype mismatch: " + f"buffer shape={tuple(dst.shape)} dtype={dst.dtype}, " + f"source shape={tuple(src.shape)} dtype={src.dtype}" + ) + dst.copy_(src) + + @staticmethod + def _int32_scalar(value: int) -> int: + return int(value) + + def _ensure_shared_host_allocation_before_worker(self, name: str) -> None: + if self._l3_worker is not None: + raise RuntimeError( + f"DeepSeekV4 shared host buffer '{name}' must be allocated before the L3 worker starts" + ) + + def _assert_l3_args_shared_before_worker( + self, + callable_spec: DeepSeekV4L3Callable, + args: Sequence[Any], + ) -> None: + for index, arg in enumerate(args): + self._assert_l3_arg_shared(arg, name=f"{callable_spec.name}[{index}]") + + def _assert_l3_arg_shared(self, arg: Any, *, name: str) -> None: + if isinstance(arg, (_StaticDeviceTensor, _TransientDeviceTensor)): + self._assert_l3_arg_shared(arg.tensor, name=f"{name}.tensor") + return + if isinstance(arg, torch.Tensor) and arg.device.type == "cpu" and not arg.is_shared(): + raise TypeError( + "DeepSeekV4 L3 dispatch requires shared-memory CPU tensors allocated before " + f"the L3 worker starts; got {name} shape={tuple(arg.shape)} dtype={arg.dtype}" + ) + if isinstance(arg, Sequence) and not isinstance(arg, (str, bytes, bytearray)): + for index, item in enumerate(arg): + self._assert_l3_arg_shared(item, name=f"{name}[{index}]") + return + if isinstance(arg, dict): + for key, item in arg.items(): + self._assert_l3_arg_shared(item, name=f"{name}[{key!r}]") + + def _coerce_l3_arg(self, worker: Any, arg: Any, uploaded: list[DeviceTensor]) -> Any: + if isinstance(arg, _StaticDeviceTensor): + self._assert_l3_arg_shared(arg, name="static") + return arg.tensor + if isinstance(arg, _TransientDeviceTensor): + tensor = arg.tensor + self._assert_l3_arg_shared(arg, name="transient") + dev = worker.alloc_tensor(tensor.shape, tensor.dtype, init=tensor) + uploaded.append(dev) + return dev + if isinstance(arg, torch.Tensor) and arg.device.type == "cpu" and not arg.is_shared(): + raise TypeError( + "DeepSeekV4 L3 dispatch requires shared-memory CPU tensors allocated before " + f"the worker starts; got non-shared tensor shape={tuple(arg.shape)} dtype={arg.dtype}" + ) + return arg + + def _shared_l3_worker(self) -> Any: + worker = self._l3_worker + if worker is None: + self._assert_l3_shared_buffers_preallocated() + compiled_callables = self._compiled.l3_callables() + if not compiled_callables: + raise RuntimeError("DeepSeekV4 L3 callables are not compiled") + from pypto.runtime import DistributedWorker # noqa: PLC0415 + + worker = DistributedWorker([callable_spec.compiled for callable_spec in compiled_callables]) + self._l3_worker = worker + return worker + + def _ensure_decode_work_cache(self) -> DeepSeekV4LayerCache: + cache = self._decode_work_cache + if cache is not None: + return cache + self._ensure_shared_host_allocation_before_worker("decode work cache") + layout = self._compiled.layout + fwd_layers = DEEPSEEK_V4_FWD_NUM_LAYERS + csa_layers = DEEPSEEK_V4_CSA_NUM_LAYERS + hca_layers = DEEPSEEK_V4_HCA_NUM_LAYERS + cache = DeepSeekV4LayerCache( + kv_cache=self._shared_empty( + ( + layout.ranks, + fwd_layers * layout.decode_batch * layout.ori_max_blocks, + layout.block_size, + 1, + DEEPSEEK_V4_HEAD_DIM, + ), + torch.bfloat16, + name="decode_work_kv_cache", + ), + cmp_kv=self._shared_empty( + ( + layout.ranks, + fwd_layers * layout.decode_batch * layout.cmp_max_blocks, + layout.block_size, + 1, + DEEPSEEK_V4_HEAD_DIM, + ), + torch.bfloat16, + name="decode_work_cmp_kv", + ), + idx_kv_cache=self._shared_empty( + ( + layout.ranks, + csa_layers * layout.decode_batch * layout.idx_max_blocks, + layout.block_size, + 1, + DEEPSEEK_V4_IDX_HEAD_DIM, + ), + torch.bfloat16, + name="decode_work_idx_kv_cache", + ), + hca_compress_state=self._shared_empty( + ( + layout.ranks, + hca_layers * layout.decode_batch * layout.hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_STATE_DIM, + ), + torch.float32, + name="decode_work_hca_compress_state", + ), + csa_compress_state=self._shared_empty( + ( + layout.ranks, + csa_layers * layout.decode_batch * layout.csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_STATE_DIM, + ), + torch.float32, + name="decode_work_csa_compress_state", + ), + csa_inner_compress_state=self._shared_empty( + ( + layout.ranks, + csa_layers * layout.decode_batch * layout.csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_STATE_DIM, + ), + torch.float32, + name="decode_work_csa_inner_compress_state", + ), + ) + self._decode_work_cache = cache + return cache + + def _require_decode_work_cache(self) -> DeepSeekV4LayerCache: + if self._decode_work_cache is None: + raise RuntimeError("DeepSeekV4 decode work cache was not allocated before the L3 worker started") + return self._decode_work_cache + + @staticmethod + def _static_device_tensor(tensor: torch.Tensor) -> torch.Tensor: + if tensor.device.type != "cpu": + raise ValueError("worker-resident tensor must be on CPU") + if not tensor.is_contiguous(): + raise ValueError("worker-resident tensor must be contiguous") + return DeepSeekV4ModelRunner._share_cpu_tensor(tensor) + + def _reset_l3_worker(self) -> None: + worker = self._l3_worker + if worker is None: + return + try: + for tensor in self._l3_static_tensors.values(): + worker.free_tensor(tensor) + worker.close() + finally: + self._l3_worker = None + self._l3_static_tensors.clear() + + def close(self) -> None: + worker = self._l3_worker + try: + if worker is not None: + for tensor in self._l3_static_tensors.values(): + worker.free_tensor(tensor) + worker.close() + finally: + self._l3_worker = None + self._decode_work_cache = None + self._decode_cache_seeded_slots.clear() + self._prefill_cache_snapshots.clear() + self._l3_static_tensors.clear() + + def _require_input_builder(self) -> DeepSeekV4InputBuilder: + if self.input_builder is None: + raise RuntimeError("DeepSeekV4 input builder is not initialized") + return self.input_builder + + def _rank_stack(self, tensor: torch.Tensor) -> torch.Tensor: + return tensor.unsqueeze(0).expand(self._compiled.layout.ranks, *tensor.shape).contiguous() + + def _prefill_kernel_tokens(self, actual_tokens: int) -> int: + # The packed prefill kernel currently uses its static 128-row contract. + if actual_tokens <= 0: + raise ValueError("actual_tokens must be positive") + return self._compiled.layout.prefill_seq + + @staticmethod + def _prefill_kernel_positions( + positions: Sequence[int], + *, + kernel_tokens: int, + max_seq_len: int, + ) -> list[int]: + if len(positions) <= 0: + raise ValueError("positions must not be empty") + if kernel_tokens < len(positions): + raise ValueError("kernel_tokens must cover all active positions") + start = int(positions[0]) + kernel_positions = list(range(start, start + kernel_tokens)) + if kernel_positions[-1] >= max_seq_len: + raise ValueError( + f"prefill static kernel position {kernel_positions[-1]} exceeds max_seq_len={max_seq_len}" + ) + return kernel_positions + + def _prefill_kernel_slots(self, slot: int, *, actual_tokens: int, kernel_tokens: int) -> list[int]: + if actual_tokens <= 0: + raise ValueError("actual_tokens must be positive") + if kernel_tokens < actual_tokens: + raise ValueError("kernel_tokens must cover all active tokens") + slot = int(slot) + scratch_slot = slot + if kernel_tokens > actual_tokens and self._compiled.layout.decode_batch > 1: + scratch_slot = (slot + self._compiled.layout.decode_batch - 1) % self._compiled.layout.decode_batch + return [slot] * actual_tokens + [scratch_slot] * (kernel_tokens - actual_tokens) + + def _prefill_sliding_window_slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[int], + ) -> torch.Tensor: + layout = self._compiled.layout + if len(slots) != len(positions): + raise ValueError("prefill slots and positions must have the same length") + mapping = torch.empty((len(positions),), dtype=torch.long) + capacity = layout.ori_max_blocks * layout.block_size + for row, (slot, position) in enumerate(zip(slots, positions, strict=True)): + mapping[row] = int(slot) * capacity + (int(position) % layout.block_size) + return mapping + + def _prefill_compressed_slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[int], + *, + max_blocks: int, + compress_ratio: int, + ) -> torch.Tensor: + if len(slots) != len(positions): + raise ValueError("prefill slots and positions must have the same length") + if compress_ratio <= 0: + raise ValueError("compress_ratio must be positive") + capacity = int(max_blocks) * self._compiled.layout.block_size + mapping = torch.full((len(positions),), -1, dtype=torch.long) + for row, (slot, position) in enumerate(zip(slots, positions, strict=True)): + position = int(position) + if (position + 1) % compress_ratio != 0: + continue + logical = position // compress_ratio + if logical >= capacity: + raise ValueError( + f"position {position} maps to compressed row {logical}, but capacity is {capacity}" + ) + mapping[row] = int(slot) * capacity + logical + return mapping + + def _prefill_state_slot_mapping( + self, + slots: Sequence[int], + positions: Sequence[int], + *, + max_blocks: int, + state_block_size: int, + ) -> torch.Tensor: + if len(slots) != len(positions): + raise ValueError("prefill slots and positions must have the same length") + capacity = int(max_blocks) * int(state_block_size) + mapping = torch.empty((len(positions),), dtype=torch.long) + for row, (slot, position) in enumerate(zip(slots, positions, strict=True)): + position = int(position) + if position >= capacity: + raise ValueError( + f"position {position} exceeds compressor-state capacity {capacity} " + f"(max_blocks={max_blocks}, state_block_size={state_block_size})" + ) + mapping[row] = int(slot) * capacity + position + return mapping + + @staticmethod + def _padded_rows(values: torch.Tensor, length: int) -> torch.Tensor: + if values.ndim != 2: + raise ValueError(f"values must be rank-2, got shape={tuple(values.shape)}") + if values.shape[0] <= 0: + raise ValueError("values must not be empty") + if values.shape[0] > length: + raise ValueError(f"values rows {values.shape[0]} exceed padded length {length}") + out = torch.empty((length, values.shape[1]), dtype=values.dtype, device=values.device) + out[: values.shape[0]].copy_(values) + if values.shape[0] < length: + pad_rows = torch.arange(values.shape[0], length, device=values.device) % values.shape[0] + out[values.shape[0] :].copy_(values.index_select(0, pad_rows)) + return out + + @staticmethod + def _padded_vector(values: torch.Tensor, length: int, *, dtype: torch.dtype) -> torch.Tensor: + if values.numel() <= 0: + raise ValueError("values must not be empty") + if values.numel() > length: + raise ValueError(f"values length {values.numel()} exceeds padded length {length}") + out = torch.empty((length,), dtype=dtype) + out[: values.numel()] = values.to(dtype=dtype) + if values.numel() < length: + pad_rows = torch.arange(values.numel(), length) % values.numel() + out[values.numel() :] = values.to(dtype=dtype).index_select(0, pad_rows) + return out + + @staticmethod + def _prefill_position_ids(positions: Sequence[int], length: int) -> torch.Tensor: + if len(positions) <= 0: + raise ValueError("positions must not be empty") + if len(positions) > length: + raise ValueError(f"positions length {len(positions)} exceeds padded length {length}") + out = torch.arange(length, dtype=torch.int32) + out[: len(positions)] = torch.tensor(tuple(int(pos) for pos in positions), dtype=torch.int32) + return out + + @staticmethod + def _pad_prefill_mapping(mapping: torch.Tensor, length: int) -> torch.Tensor: + if mapping.ndim != 1: + raise ValueError(f"prefill mapping must be rank-1, got shape={tuple(mapping.shape)}") + if mapping.numel() > length: + raise ValueError(f"prefill mapping length {mapping.numel()} exceeds padded length {length}") + out = torch.full((length,), -1, dtype=mapping.dtype) + out[: mapping.numel()].copy_(mapping.to(dtype=mapping.dtype)) + return out + + @staticmethod + def _prefill_actual_tokens(batch: PrefillBatch) -> int: + if batch.positions is not None: + valid = batch.positions[0].detach().cpu() + valid = valid[valid >= 0] + if valid.numel() <= 0: + raise ValueError("prefill positions must include at least one token") + return int(valid.numel()) + seq_len = int(batch.seq_lens[0].item()) + if seq_len <= 0: + raise ValueError("prefill seq_len must be positive") + return seq_len + + @staticmethod + def _prefill_positions(batch: PrefillBatch, actual_tokens: int) -> list[int]: + if batch.positions is None: + positions = list(range(actual_tokens)) + else: + raw = batch.positions[0, :actual_tokens].detach().cpu().to(torch.long) + positions = [int(pos) for pos in raw.tolist()] + if any(pos < 0 for pos in positions): + raise ValueError("prefill positions must be non-negative") + expected = list(range(positions[0], positions[0] + actual_tokens)) + if positions != expected: + raise ValueError( + "prefill positions must form one contiguous chunk: " + f"positions={positions[:8]}{'...' if len(positions) > 8 else ''}" + ) + return positions + + def _prefill_sparse_by_ratio( + self, + positions: Sequence[int], + actual_tokens: int, + ) -> dict[int, tuple[torch.Tensor, torch.Tensor]]: + return { + ratio: self._prefill_sparse_indices(positions, actual_tokens, compress_ratio=ratio) + for ratio in (0, 4, 128) + } + + def _prefill_sparse_indices( + self, + positions: Sequence[int], + actual_tokens: int, + *, + compress_ratio: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + layout = self._compiled.layout + indices = torch.full( + (layout.prefill_seq, layout.prefill_sparse_topk), + -1, + dtype=torch.int32, + ) + lens = torch.zeros((layout.prefill_seq,), dtype=torch.int32) + current = {int(pos): row for row, pos in enumerate(positions[:actual_tokens])} + for row in range(actual_tokens): + abs_pos = int(positions[row]) + window_valid = min(layout.block_size, abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + cursor = 0 + for key_abs in range(key_start_abs, abs_pos + 1): + overlay_row = current.get(key_abs) + if overlay_row is not None and overlay_row <= row: + indices[row, cursor] = layout.block_size + overlay_row + else: + indices[row, cursor] = key_abs % layout.block_size + cursor += 1 + if compress_ratio > 0: + compressed_visible = min( + (abs_pos + 1) // compress_ratio, + layout.cmp_max_blocks * layout.block_size, + layout.prefill_sparse_topk - layout.block_size, + ) + for cmp_slot in range(compressed_visible): + if cursor >= layout.prefill_sparse_topk: + break + indices[row, cursor] = layout.block_size + layout.prefill_seq + cmp_slot + cursor += 1 + lens[row] = cursor + return indices, lens + + def _decode_positions(self, batch: DecodeBatch, actual_batch: int) -> tuple[tuple[int, ...], ...]: + decode_seq = self._compiled.layout.decode_seq + positions = [] + for row in range(actual_batch): + seq_len = int(batch.seq_lens[row].item()) + if seq_len < decode_seq: + raise ValueError( + f"decode seq_lens must be >= decode_seq ({decode_seq}), got {seq_len}" + ) + # MTP feeds ``decode_seq`` real trailing tokens ending at the last real + # position ``seq_len-1`` (so positions are ``seq_len-decode_seq .. seq_len-1``). + first_position = seq_len - decode_seq + positions.append(tuple(first_position + offset for offset in range(decode_seq))) + return tuple(positions) + + def _decode_token_rows( + self, + token_ids: torch.Tensor, + actual_batch: int, + *, + vocab_size: int, + prev_token_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + layout = self._compiled.layout + if token_ids.ndim == 1: + active = token_ids[:actual_batch].reshape(actual_batch, 1) + else: + active = token_ids[:actual_batch, :1] + prev_active = None + if prev_token_ids is not None: + prev_active = prev_token_ids[:actual_batch].reshape(actual_batch, 1) + if vocab_size <= 0: + raise ValueError("vocab_size must be positive") + rows = torch.empty(layout.decode_tokens, dtype=torch.long).reshape( + layout.decode_batch, + layout.decode_seq, + ) + if prev_active is not None: + rows.copy_(prev_active[0, 0].expand(layout.decode_batch, layout.decode_seq)) + rows[:, layout.decode_seq - 1].copy_(active[0, 0]) + else: + rows.copy_(active[0, 0].expand(layout.decode_batch, layout.decode_seq)) + for row in range(actual_batch): + if prev_active is not None: + # Earlier slots use prev token; final slot uses last token. + rows[row].copy_(prev_active[row, 0].expand(layout.decode_seq)) + rows[row, layout.decode_seq - 1].copy_(active[row, 0]) + else: + rows[row].copy_(active[row, 0].expand(layout.decode_seq)) + return rows.reshape(layout.decode_tokens) + + def _decode_kernel_slots(self, active_slots: Sequence[int]) -> tuple[int, ...]: + """Route padded fixed decode rows into scratch cache slots.""" + layout = self._compiled.layout + slots = [int(slot) for slot in active_slots] + if not slots: + raise ValueError("decode must include at least one active slot") + if len(set(slots)) != len(slots): + raise ValueError(f"decode slots must be unique, got {slots}") + if len(slots) > layout.decode_batch: + raise ValueError(f"decode slots exceed kernel batch {layout.decode_batch}: {slots}") + active_set = set(slots) + for scratch_slot in range(layout.decode_batch): + if len(slots) >= layout.decode_batch: + break + if scratch_slot not in active_set: + slots.append(scratch_slot) + if len(slots) != layout.decode_batch: + raise RuntimeError( + f"DeepSeekV4 decode needs {layout.decode_batch} kernel slots, built {len(slots)}" + ) + return tuple(slots) + + def _decode_kv_seq_lens(self, seq_lens: torch.Tensor, actual_batch: int) -> torch.Tensor: + layout = self._compiled.layout + # The last written KV position is ``seq_len-1``, so the valid KV history + # is exactly ``seq_len`` entries. (yangyaodong's "seq_len+1" was relative + # to a seq_len = prompt length, which does not count the prefill token; + # our seq_len already does.) + active = seq_lens[:actual_batch].detach().cpu().to(torch.int32) + return DeepSeekV4CacheManager.replicate_first_row( + active.reshape(actual_batch, 1), + actual_rows=actual_batch, + kernel_rows=layout.decode_batch, + ).reshape(layout.decode_batch) diff --git a/examples/model/deepseek_v4/runner/weight_loader.py b/examples/model/deepseek_v4/runner/weight_loader.py new file mode 100644 index 0000000..95bce9a --- /dev/null +++ b/examples/model/deepseek_v4/runner/weight_loader.py @@ -0,0 +1,977 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import logging +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import ContextManager, Protocol + +import torch + +logger = logging.getLogger(__name__) + + +class _SafeTensorReader(Protocol): + """Minimal safetensors reader protocol used by the lazy weight store.""" + + def get_tensor(self, name: str) -> torch.Tensor: + """Return one tensor by name.""" + raise NotImplementedError + + +class _SafeOpenFn(Protocol): + """Callable shape for injectable safetensors openers.""" + + def __call__(self, path: Path, device: str) -> ContextManager[_SafeTensorReader]: + """Open one safetensors shard.""" + raise NotImplementedError + + +_GLOBAL_WEIGHT_NAMES = ( + "embed.weight", + "norm.weight", + "head.weight", + "hc_head_fn", + "hc_head_scale", + "hc_head_base", +) +_LM_HEAD_VOCAB_CHUNK = 512 +_LAYER_COMMON_SUFFIXES = ( + "attn.attn_sink", + "attn.kv_norm.weight", + "attn.q_norm.weight", + "attn.wkv.weight", + "attn.wo_a.weight", + "attn.wo_b.weight", + "attn.wo_b.scale", + "attn.wq_a.weight", + "attn.wq_b.weight", + "attn.wq_b.scale", + "attn_norm.weight", + "ffn.gate.weight", + "ffn.shared_experts.w1.weight", + "ffn.shared_experts.w1.scale", + "ffn.shared_experts.w2.weight", + "ffn.shared_experts.w2.scale", + "ffn.shared_experts.w3.weight", + "ffn.shared_experts.w3.scale", + "ffn_norm.weight", + "hc_attn_base", + "hc_attn_fn", + "hc_attn_scale", + "hc_ffn_base", + "hc_ffn_fn", + "hc_ffn_scale", +) +_LAYER_COMPRESSOR_SUFFIXES = ( + "attn.compressor.ape", + "attn.compressor.norm.weight", + "attn.compressor.wgate.weight", + "attn.compressor.wkv.weight", +) +_LAYER_INDEXER_SUFFIXES = ( + "attn.indexer.compressor.ape", + "attn.indexer.compressor.norm.weight", + "attn.indexer.compressor.wgate.weight", + "attn.indexer.compressor.wkv.weight", + "attn.indexer.weights_proj.weight", + "attn.indexer.wq_b.weight", + "attn.indexer.wq_b.scale", +) +_EXPERT_SUFFIXES = ("w1.weight", "w1.scale", "w2.weight", "w2.scale", "w3.weight", "w3.scale") +_DEEPSEEK_V4_O_GROUPS = 8 +_DEEPSEEK_V4_HADAMARD_IDX_DIM = 128 +_DEEPSEEK_V4_HCA_COMPRESS_RATIO = 128 +_DEEPSEEK_V4_CSA_COMPRESS_RATIO = 4 +_DEEPSEEK_V4_HCA_MAIN_OUT_DIM = 512 +_DEEPSEEK_V4_CSA_MAIN_OUT_DIM = 1024 +_DEEPSEEK_V4_CSA_INNER_OUT_DIM = 256 +_DEEPSEEK_V4_HIDDEN_SIZE = 4096 +_DEEPSEEK_V4_Q_LORA = 1024 +_DEEPSEEK_V4_HEAD_DIM = 512 +_DEEPSEEK_V4_ATTENTION_OUT = 64 * 512 +_DEEPSEEK_V4_N_ROUTED_EXPERTS = 256 +_DEEPSEEK_V4_TOPK = 6 +_DEEPSEEK_V4_VOCAB_SIZE = 129280 + + +def _default_safe_open(path: Path, device: str) -> ContextManager[_SafeTensorReader]: + """Open a safetensors shard without loading unrelated tensors.""" + try: + from safetensors import safe_open + except ImportError as exc: + raise RuntimeError("safetensors is required to read DeepSeekV4 W8A8 weights.") from exc + + return safe_open(str(path), framework="pt", device=device) + + +def deepseek_v4_global_weight_names() -> tuple[str, ...]: + """Return global DeepSeekV4 tensor names needed outside the layer stack.""" + return _GLOBAL_WEIGHT_NAMES + + +@dataclass(frozen=True) +class DeepSeekV4LmHeadLayout: + """8-way tensor-parallel LM-head shard layout.""" + + ranks: int + vocab_size: int + hidden_size: int + vocab_per_rank: int + padded_vocab_per_rank: int + + +@dataclass(frozen=True) +class DeepSeekV4GlobalWeights: + """Global DeepSeekV4 weights packed for serving kernels.""" + + embed_weight: torch.Tensor + final_norm_weight: torch.Tensor + lm_head_weight: torch.Tensor + lm_head_layout: DeepSeekV4LmHeadLayout + hc_head_fn: torch.Tensor + hc_head_scale: torch.Tensor + hc_head_base: torch.Tensor + + +@dataclass(frozen=True) +class DeepSeekV4PackedLayerWeights: + """One DeepSeekV4 layer's tensors packed in pypto-lib host argument names.""" + + layer_id: int + tensors: Mapping[str, torch.Tensor] + + def args(self, names: Sequence[str]) -> tuple[torch.Tensor, ...]: + """Return tensors in a kernel host order.""" + missing = [name for name in names if name not in self.tensors] + if missing: + raise KeyError(f"Packed DeepSeekV4 layer is missing tensors: {', '.join(missing)}") + return tuple(self.tensors[name] for name in names) + + +# Layer-stacking groups for the packed all-layer ``l3_decode_fwd`` kernel. These +# mirror the name groups in pypto-lib decode_fwd.py, but only cover *loaded* +# weights -- the per-layer work-cache/state tensors (kv_cache, cmp_kv, +# idx_kv_cache, *_compress_state) are owned by the runner work cache and are not +# emitted by the weight loader. +DEEPSEEK_V4_CSA_STACKED_WEIGHT_NAMES = ( + "csa_cmp_wkv", + "csa_cmp_wgate", + "csa_cmp_ape", + "csa_cmp_norm_w", + "csa_idx_wq_b", + "csa_idx_wq_b_scale", + "csa_weights_proj", + "csa_hadamard_idx", + "csa_inner_wkv", + "csa_inner_wgate", + "csa_inner_ape", + "csa_inner_norm_w", +) +DEEPSEEK_V4_HCA_STACKED_WEIGHT_NAMES = ( + "hca_cmp_wkv", + "hca_cmp_wgate", + "hca_cmp_ape", + "hca_cmp_norm_w", +) +_DEEPSEEK_V4_CSA_COMPRESS_RATIO_VALUE = 4 +_DEEPSEEK_V4_HCA_COMPRESS_RATIO_VALUE = 128 + + +@dataclass(frozen=True) +class DeepSeekV4StackedLayerWeights: + """All hidden-layer weights stacked on the layer axis for ``l3_decode_fwd``. + + Each tensor fuses its layer axis into dim 1: ``[ranks, layer_count*d1, ...]``. + FWD weights stack across all 43 hidden layers; CSA-group weights stack across + the 21 compress_ratio==4 layers in order; HCA-group weights stack across the + 20 compress_ratio==128 layers in order. + """ + + tensors: Mapping[str, torch.Tensor] + + def args(self, names: Sequence[str]) -> tuple[torch.Tensor, ...]: + """Return stacked tensors in a kernel host order.""" + missing = [name for name in names if name not in self.tensors] + if missing: + raise KeyError(f"Stacked DeepSeekV4 weights are missing tensors: {', '.join(missing)}") + return tuple(self.tensors[name] for name in names) + + +def deepseek_v4_lm_head_layout( + *, + vocab_size: int, + hidden_size: int, + ranks: int, + vocab_chunk: int = _LM_HEAD_VOCAB_CHUNK, +) -> DeepSeekV4LmHeadLayout: + """Return the LM-head shard shape expected by ``lm_head.py``.""" + if ranks <= 0: + raise ValueError("ranks must be positive") + if vocab_chunk <= 0: + raise ValueError("vocab_chunk must be positive") + if vocab_size % ranks != 0: + raise ValueError(f"vocab_size={vocab_size} must divide evenly across ranks={ranks}") + vocab_per_rank = vocab_size // ranks + padded_vocab_per_rank = ((vocab_per_rank + vocab_chunk - 1) // vocab_chunk) * vocab_chunk + return DeepSeekV4LmHeadLayout( + ranks=ranks, + vocab_size=vocab_size, + hidden_size=hidden_size, + vocab_per_rank=vocab_per_rank, + padded_vocab_per_rank=padded_vocab_per_rank, + ) + + +def pack_deepseek_v4_lm_head_weight( + weight: torch.Tensor, + *, + ranks: int, + vocab_chunk: int = _LM_HEAD_VOCAB_CHUNK, +) -> tuple[torch.Tensor, DeepSeekV4LmHeadLayout]: + """Pack flat ``head.weight`` into contiguous TP vocab shards.""" + if weight.ndim != 2: + raise ValueError(f"lm_head weight must be rank-2, got shape={tuple(weight.shape)}") + vocab_size, hidden_size = (int(dim) for dim in weight.shape) + layout = deepseek_v4_lm_head_layout( + vocab_size=vocab_size, + hidden_size=hidden_size, + ranks=ranks, + vocab_chunk=vocab_chunk, + ) + packed = torch.zeros( + (layout.ranks, layout.padded_vocab_per_rank, layout.hidden_size), + dtype=weight.dtype, + device=weight.device, + ) + for rank in range(layout.ranks): + start = rank * layout.vocab_per_rank + end = start + layout.vocab_per_rank + packed[rank, : layout.vocab_per_rank].copy_(weight[start:end]) + return packed.contiguous(), layout + + +def _attention_suffixes_for_compress_ratio(compress_ratio: int) -> tuple[str, ...]: + """Return attention parameter suffixes required by one layer attention mode.""" + if compress_ratio == 0: + return () + if compress_ratio == 128: + return _LAYER_COMPRESSOR_SUFFIXES + if compress_ratio == 4: + return (*_LAYER_COMPRESSOR_SUFFIXES, *_LAYER_INDEXER_SUFFIXES) + raise ValueError(f"unsupported DeepSeekV4 attention compress ratio: {compress_ratio}") + + +def deepseek_v4_layer_core_weight_names( + layer_id: int, + *, + compress_ratio: int = 0, + include_tid2eid: bool = False, + include_gate_bias: bool = False, +) -> tuple[str, ...]: + """Return non-routed-expert tensor names for one DeepSeekV4 layer.""" + prefix = f"layers.{int(layer_id)}" + suffixes = [*_LAYER_COMMON_SUFFIXES, *_attention_suffixes_for_compress_ratio(compress_ratio)] + if include_tid2eid: + suffixes.append("ffn.gate.tid2eid") + if include_gate_bias: + suffixes.append("ffn.gate.bias") + return tuple(f"{prefix}.{suffix}" for suffix in suffixes) + + +def deepseek_v4_routed_expert_weight_names(layer_id: int, expert_ids: Iterable[int]) -> tuple[str, ...]: + """Return routed expert tensor names for one DeepSeekV4 layer.""" + names: list[str] = [] + for expert_id in expert_ids: + prefix = f"layers.{int(layer_id)}.ffn.experts.{int(expert_id)}" + names.extend(f"{prefix}.{suffix}" for suffix in _EXPERT_SUFFIXES) + return tuple(names) + + +def deepseek_v4_local_expert_ids(*, rank: int, ranks: int, n_routed_experts: int) -> tuple[int, ...]: + """Return the contiguous routed-expert ids owned by one EP rank.""" + if ranks <= 0: + raise ValueError("ranks must be positive") + if not 0 <= rank < ranks: + raise ValueError(f"rank must be in [0, {ranks - 1}], got {rank}") + if n_routed_experts <= 0: + raise ValueError("n_routed_experts must be positive") + if n_routed_experts % ranks != 0: + raise ValueError(f"n_routed_experts={n_routed_experts} must divide evenly across ranks={ranks}") + local_count = n_routed_experts // ranks + start = rank * local_count + return tuple(range(start, start + local_count)) + + +def deepseek_v4_hadamard_idx(dim: int = _DEEPSEEK_V4_HADAMARD_IDX_DIM) -> torch.Tensor: + """Return the normalized Hadamard matrix used by the CSA indexer.""" + if dim <= 0 or dim & (dim - 1) != 0: + raise ValueError("Hadamard dimension must be a positive power of two") + h = torch.ones((1, 1), dtype=torch.bfloat16) + while h.shape[0] < dim: + h = torch.cat( + [torch.cat([h, h], dim=1), torch.cat([h, -h], dim=1)], + dim=0, + ) + return (h * (dim**-0.5)).contiguous() + + +def deepseek_v4_layer_weight_names( + layer_id: int, + *, + n_routed_experts: int, + compress_ratio: int = 0, + include_tid2eid: bool = False, + include_gate_bias: bool = False, + expert_ids: Iterable[int] | None = None, +) -> tuple[str, ...]: + """Return all tensor names needed to execute one DeepSeekV4 layer.""" + if n_routed_experts <= 0: + raise ValueError("n_routed_experts must be positive") + expert_ids = range(n_routed_experts) if expert_ids is None else tuple(expert_ids) + return ( + *deepseek_v4_layer_core_weight_names( + layer_id, + compress_ratio=compress_ratio, + include_tid2eid=include_tid2eid, + include_gate_bias=include_gate_bias, + ), + *deepseek_v4_routed_expert_weight_names(layer_id, expert_ids), + ) + + +def deepseek_v4_startup_weight_names( + num_hidden_layers: int, + *, + n_routed_experts: int, + compress_ratios: Sequence[int] | None = None, + num_hash_layers: int = 3, +) -> tuple[str, ...]: + """Return tensor names used for metadata-only checkpoint contract validation. + + Startup checks every layer's core tensors plus the first and last routed + expert in each layer. Full expert materialization remains an explicit + per-layer load so serving startup does not read shard payloads. + """ + if num_hidden_layers <= 0: + raise ValueError("num_hidden_layers must be positive") + if n_routed_experts <= 0: + raise ValueError("n_routed_experts must be positive") + if compress_ratios is None: + compress_ratios = (0,) * num_hidden_layers + if len(compress_ratios) < num_hidden_layers: + raise ValueError("compress_ratios must include at least one entry per hidden layer") + + edge_experts = tuple(dict.fromkeys((0, n_routed_experts - 1))) + names = list(_GLOBAL_WEIGHT_NAMES) + for layer_id in range(num_hidden_layers): + names.extend( + deepseek_v4_layer_core_weight_names( + layer_id, + compress_ratio=int(compress_ratios[layer_id]), + include_tid2eid=layer_id < num_hash_layers, + include_gate_bias=layer_id >= num_hash_layers, + ) + ) + names.extend(deepseek_v4_routed_expert_weight_names(layer_id, edge_experts)) + return tuple(dict.fromkeys(names)) + + +class DeepSeekV4WeightStore: + """Lazy name-based safetensors access for DeepSeekV4 W8A8 checkpoints.""" + + def __init__( + self, + *, + model_dir: str | Path, + weight_map: Mapping[str, str], + device: str = "cpu", + safe_open_fn: _SafeOpenFn | None = None, + ) -> None: + """Create a store from the Hugging Face safetensors index.""" + self.model_dir = Path(model_dir) + self.weight_map = dict(weight_map) + self.device = device + self._safe_open_fn = _default_safe_open if safe_open_fn is None else safe_open_fn + + def __contains__(self, name: object) -> bool: + """Return whether the checkpoint index exposes ``name``.""" + return isinstance(name, str) and name in self.weight_map + + def filename_for(self, name: str) -> str: + """Return the safetensors shard filename for ``name``.""" + try: + return self.weight_map[name] + except KeyError as exc: + raise KeyError(f"Missing DeepSeekV4 weight tensor in index: {name}") from exc + + def path_for(self, name: str) -> Path: + """Return the shard path containing ``name``.""" + return self.model_dir / self.filename_for(name) + + def require(self, names: Iterable[str]) -> None: + """Validate that all tensor names are present in the checkpoint index.""" + missing = [name for name in names if name not in self.weight_map] + if missing: + preview = ", ".join(missing[:8]) + suffix = "" if len(missing) <= 8 else f", ... ({len(missing)} total)" + raise KeyError(f"DeepSeekV4 W8A8 checkpoint is missing required tensors: {preview}{suffix}") + + def validate_startup_contract( + self, + *, + num_hidden_layers: int, + n_routed_experts: int, + compress_ratios: Sequence[int] | None = None, + num_hash_layers: int = 3, + ) -> None: + """Validate the startup-visible checkpoint contract without opening shards.""" + self.require( + deepseek_v4_startup_weight_names( + num_hidden_layers, + n_routed_experts=n_routed_experts, + compress_ratios=compress_ratios, + num_hash_layers=num_hash_layers, + ) + ) + + def load_tensor(self, name: str) -> torch.Tensor: + """Load one tensor by name, leaving all unrelated shard tensors untouched.""" + return self.load_many([name])[name] + + def load_many(self, names: Sequence[str]) -> dict[str, torch.Tensor]: + """Load a set of named tensors grouped by shard file.""" + unique_names = tuple(dict.fromkeys(names)) + self.require(unique_names) + + groups: dict[str, list[str]] = {} + for name in unique_names: + groups.setdefault(self.filename_for(name), []).append(name) + + loaded: dict[str, torch.Tensor] = {} + for filename, shard_names in groups.items(): + path = self.model_dir / filename + if not path.exists(): + raise FileNotFoundError(f"Missing safetensors shard for DeepSeekV4 weight load: {path}") + with self._safe_open_fn(path, self.device) as reader: + for name in shard_names: + loaded[name] = reader.get_tensor(name) + + return {name: loaded[name] for name in unique_names} + + def load_global_weights(self) -> dict[str, torch.Tensor]: + """Load embedding, final norm, and LM head tensors.""" + return self.load_many(deepseek_v4_global_weight_names()) + + def load_packed_global_weights(self, *, ranks: int) -> DeepSeekV4GlobalWeights: + """Load and pack global tensors for the DeepSeekV4 serving kernels.""" + weights = self.load_global_weights() + packed_lm_head, layout = pack_deepseek_v4_lm_head_weight(weights["head.weight"], ranks=ranks) + if weights["embed.weight"].ndim != 2: + raise ValueError(f"embed.weight must be rank-2, got shape={tuple(weights['embed.weight'].shape)}") + if weights["norm.weight"].ndim != 1: + raise ValueError(f"norm.weight must be rank-1, got shape={tuple(weights['norm.weight'].shape)}") + if tuple(weights["embed.weight"].shape) != (layout.vocab_size, layout.hidden_size): + raise ValueError( + "embed.weight shape must match head.weight shape, " + f"got embed={tuple(weights['embed.weight'].shape)}, head={tuple(weights['head.weight'].shape)}" + ) + if int(weights["norm.weight"].shape[0]) != layout.hidden_size: + raise ValueError( + f"norm.weight hidden size must be {layout.hidden_size}, " + f"got {int(weights['norm.weight'].shape[0])}" + ) + if tuple(weights["hc_head_fn"].shape) != (4, layout.hidden_size * 4): + raise ValueError(f"hc_head_fn has unsupported shape {tuple(weights['hc_head_fn'].shape)}") + if tuple(weights["hc_head_scale"].shape) != (1,): + raise ValueError(f"hc_head_scale has unsupported shape {tuple(weights['hc_head_scale'].shape)}") + if tuple(weights["hc_head_base"].shape) != (4,): + raise ValueError(f"hc_head_base has unsupported shape {tuple(weights['hc_head_base'].shape)}") + return DeepSeekV4GlobalWeights( + embed_weight=weights["embed.weight"], + final_norm_weight=weights["norm.weight"], + lm_head_weight=packed_lm_head, + lm_head_layout=layout, + hc_head_fn=weights["hc_head_fn"].to(torch.float32).contiguous().cpu(), + hc_head_scale=weights["hc_head_scale"].to(torch.float32).contiguous().cpu(), + hc_head_base=weights["hc_head_base"].to(torch.float32).contiguous().cpu(), + ) + + def load_layer_weights( + self, + layer_id: int, + *, + n_routed_experts: int, + compress_ratio: int = 0, + include_tid2eid: bool = False, + include_gate_bias: bool = False, + expert_ids: Iterable[int] | None = None, + ) -> dict[str, torch.Tensor]: + """Load all tensors needed for one DeepSeekV4 layer.""" + return self.load_many( + deepseek_v4_layer_weight_names( + layer_id, + n_routed_experts=n_routed_experts, + compress_ratio=compress_ratio, + include_tid2eid=include_tid2eid, + include_gate_bias=include_gate_bias, + expert_ids=expert_ids, + ) + ) + + def load_rank_layer_weights( + self, + layer_id: int, + *, + rank: int, + ranks: int, + n_routed_experts: int, + compress_ratio: int = 0, + include_tid2eid: bool = False, + include_gate_bias: bool = False, + ) -> dict[str, torch.Tensor]: + """Load common layer tensors plus the routed experts owned by one rank.""" + local_experts = deepseek_v4_local_expert_ids( + rank=rank, + ranks=ranks, + n_routed_experts=n_routed_experts, + ) + return self.load_layer_weights( + layer_id, + n_routed_experts=n_routed_experts, + compress_ratio=compress_ratio, + include_tid2eid=include_tid2eid, + include_gate_bias=include_gate_bias, + expert_ids=local_experts, + ) + + def load_packed_layer_weights( + self, + layer_id: int, + *, + ranks: int, + n_routed_experts: int, + compress_ratio: int = 0, + include_tid2eid: bool = False, + include_gate_bias: bool = False, + ) -> DeepSeekV4PackedLayerWeights: + """Load and pack one layer into the tensor names expected by pypto-lib kernels.""" + all_experts = range(n_routed_experts) + raw = self.load_layer_weights( + layer_id, + n_routed_experts=n_routed_experts, + compress_ratio=compress_ratio, + include_tid2eid=include_tid2eid, + include_gate_bias=include_gate_bias, + expert_ids=all_experts, + ) + return pack_deepseek_v4_layer_weights( + layer_id, + raw, + ranks=ranks, + n_routed_experts=n_routed_experts, + compress_ratio=compress_ratio, + include_tid2eid=include_tid2eid, + include_gate_bias=include_gate_bias, + ) + + def load_stacked_layer_weights( + self, + *, + ranks: int, + n_routed_experts: int, + compress_ratios: Sequence[int], + num_hash_layers: int, + ) -> DeepSeekV4StackedLayerWeights: + """Load every hidden layer once and stack weights on the layer axis. + + FWD weights are concatenated across all hidden layers in order; CSA-group + weights across the compress_ratio==4 layers in order; HCA-group weights + across the compress_ratio==128 layers in order. Each per-layer tensor is + ``[ranks, d1, ...]`` and stacking concatenates on dim 1. + + Layers are packed serially. A thread pool was tried but regressed: pack + is a mixed IO+CPU workload and per-layer packing allocates ~8 GB of + intermediate tensors (256 routed experts each ``torch.stack``-ed and + rank-replicated), so N parallel layers multiply peak allocation and + contend on CPU memory bandwidth; the GIL-switch cost also exceeded the + parallel gain when workers <= layer count. Serial packing keeps the + working set to one layer at a time and lets the disk prefetcher run. + """ + num_hidden_layers = len(compress_ratios) + if num_hidden_layers <= 0: + raise ValueError("compress_ratios must include at least one entry per hidden layer") + + per_layer: list[DeepSeekV4PackedLayerWeights] = [] + import time + + pack_t0 = time.perf_counter() + for layer_id in range(num_hidden_layers): + per_layer.append( + self.load_packed_layer_weights( + layer_id, + ranks=ranks, + n_routed_experts=n_routed_experts, + compress_ratio=int(compress_ratios[layer_id]), + include_tid2eid=layer_id < num_hash_layers, + include_gate_bias=layer_id >= num_hash_layers, + ) + ) + if layer_id % 5 == 0 or layer_id == num_hidden_layers - 1: + logger.info( + "DeepSeekV4 weight load progress: layer %d/%d", + layer_id + 1, + num_hidden_layers, + ) + return stack_deepseek_v4_layer_weights(per_layer, compress_ratios=compress_ratios) + + +def stack_deepseek_v4_layer_weights( + per_layer: Sequence[DeepSeekV4PackedLayerWeights], + *, + compress_ratios: Sequence[int], +) -> DeepSeekV4StackedLayerWeights: + """Concatenate per-layer packed weights into the layer-stacked decode_fwd groups.""" + num_hidden_layers = len(per_layer) + if num_hidden_layers != len(compress_ratios): + raise ValueError("per_layer count must match compress_ratios length") + if num_hidden_layers <= 0: + raise ValueError("per_layer must include at least one layer") + + csa_layers = [i for i in range(num_hidden_layers) if int(compress_ratios[i]) == _DEEPSEEK_V4_CSA_COMPRESS_RATIO_VALUE] + hca_layers = [i for i in range(num_hidden_layers) if int(compress_ratios[i]) == _DEEPSEEK_V4_HCA_COMPRESS_RATIO_VALUE] + + csa_grouped = set(DEEPSEEK_V4_CSA_STACKED_WEIGHT_NAMES) + hca_grouped = set(DEEPSEEK_V4_HCA_STACKED_WEIGHT_NAMES) + fwd_names = [ + name + for name in per_layer[0].tensors + if name not in csa_grouped and name not in hca_grouped + ] + + def cat(names: Sequence[str], layer_ids: Sequence[int]) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for name in names: + tensors = [] + for layer_id in layer_ids: + tensor = per_layer[layer_id].tensors[name] + tensors.append(tensor.contiguous().cpu()) + out[name] = torch.cat(tensors, dim=1).contiguous() + return out + + stacked: dict[str, torch.Tensor] = {} + stacked.update(cat(fwd_names, range(num_hidden_layers))) + stacked.update(cat(DEEPSEEK_V4_CSA_STACKED_WEIGHT_NAMES, csa_layers)) + stacked.update(cat(DEEPSEEK_V4_HCA_STACKED_WEIGHT_NAMES, hca_layers)) + return DeepSeekV4StackedLayerWeights(tensors=stacked) + + +def pack_deepseek_v4_layer_weights( + layer_id: int, + raw: Mapping[str, torch.Tensor], + *, + ranks: int, + n_routed_experts: int, + compress_ratio: int, + include_tid2eid: bool, + include_gate_bias: bool, +) -> DeepSeekV4PackedLayerWeights: + """Pack raw checkpoint tensors for one layer into rank-stacked kernel tensors.""" + prefix = f"layers.{int(layer_id)}" + + def get(suffix: str) -> torch.Tensor: + name = f"{prefix}.{suffix}" + try: + return raw[name] + except KeyError as exc: + raise KeyError(f"missing raw DeepSeekV4 layer tensor: {name}") from exc + + def maybe(suffix: str) -> torch.Tensor | None: + return raw.get(f"{prefix}.{suffix}") + + def replicated(tensor: torch.Tensor, *, dtype: torch.dtype | None = None) -> torch.Tensor: + if dtype is not None: + tensor = tensor.to(dtype=dtype) + tensor = tensor.contiguous().cpu() + return tensor.unsqueeze(0).expand(ranks, *tensor.shape).contiguous() + + def transposed(tensor: torch.Tensor, *, dtype: torch.dtype | None = None) -> torch.Tensor: + out = tensor.transpose(0, 1).contiguous().cpu() + return out.to(dtype=dtype) if dtype is not None else out + + def replicated_transposed(tensor: torch.Tensor, *, dtype: torch.dtype | None = None) -> torch.Tensor: + return replicated(transposed(tensor, dtype=dtype)) + + tensors: dict[str, torch.Tensor] = { + "hc_attn_fn": replicated(get("hc_attn_fn"), dtype=torch.float32), + "hc_attn_scale": replicated(get("hc_attn_scale"), dtype=torch.float32), + "hc_attn_base": replicated(get("hc_attn_base"), dtype=torch.float32), + "attn_norm_w": replicated(get("attn_norm.weight"), dtype=torch.bfloat16), + "wq_a": replicated_transposed(get("attn.wq_a.weight"), dtype=torch.bfloat16), + "wq_b": replicated_transposed(get("attn.wq_b.weight"), dtype=torch.int8), + "wq_b_scale": replicated(get("attn.wq_b.scale"), dtype=torch.float32), + "wkv": replicated_transposed(get("attn.wkv.weight"), dtype=torch.bfloat16), + "gamma_cq": replicated(get("attn.q_norm.weight"), dtype=torch.bfloat16), + "gamma_ckv": replicated(get("attn.kv_norm.weight"), dtype=torch.bfloat16), + "attn_sink": replicated(get("attn.attn_sink"), dtype=torch.float32), + "wo_a": replicated(_pack_wo_a(get("attn.wo_a.weight")), dtype=torch.bfloat16), + "wo_b": replicated(get("attn.wo_b.weight"), dtype=torch.int8), + "wo_b_scale": replicated(get("attn.wo_b.scale"), dtype=torch.float32), + "hc_ffn_fn": replicated(get("hc_ffn_fn"), dtype=torch.float32), + "hc_ffn_scale": replicated(get("hc_ffn_scale"), dtype=torch.float32), + "hc_ffn_base": replicated(get("hc_ffn_base"), dtype=torch.float32), + "norm_w": replicated(get("ffn_norm.weight"), dtype=torch.bfloat16), + "gate_w": replicated(get("ffn.gate.weight"), dtype=torch.float32), + "shared_w1": replicated(get("ffn.shared_experts.w1.weight"), dtype=torch.int8), + "shared_w1_scale": replicated(get("ffn.shared_experts.w1.scale"), dtype=torch.float32), + "shared_w3": replicated(get("ffn.shared_experts.w3.weight"), dtype=torch.int8), + "shared_w3_scale": replicated(get("ffn.shared_experts.w3.scale"), dtype=torch.float32), + "shared_w2": replicated(get("ffn.shared_experts.w2.weight"), dtype=torch.int8), + "shared_w2_scale": replicated(get("ffn.shared_experts.w2.scale"), dtype=torch.float32), + } + + tensors.update(_pack_deepseek_v4_optional_attention(prefix, raw, ranks, compress_ratio=compress_ratio)) + tensors.update( + _pack_deepseek_v4_router( + prefix, + raw, + ranks=ranks, + n_routed_experts=n_routed_experts, + include_tid2eid=include_tid2eid, + include_gate_bias=include_gate_bias, + ) + ) + tensors.update( + _pack_deepseek_v4_routed_experts( + prefix, + raw, + ranks=ranks, + n_routed_experts=n_routed_experts, + ) + ) + return DeepSeekV4PackedLayerWeights(layer_id=layer_id, tensors=tensors) + + +def _pack_wo_a(weight: torch.Tensor) -> torch.Tensor: + """Pack flattened output-LoRA A projection to ``[o_groups, o_lora, group_in]``.""" + if weight.ndim != 2: + raise ValueError(f"wo_a weight must be rank-2, got shape={tuple(weight.shape)}") + if int(weight.shape[0]) % _DEEPSEEK_V4_O_GROUPS != 0: + raise ValueError( + f"wo_a first dimension {int(weight.shape[0])} must divide by {_DEEPSEEK_V4_O_GROUPS}" + ) + return weight.reshape(_DEEPSEEK_V4_O_GROUPS, int(weight.shape[0]) // _DEEPSEEK_V4_O_GROUPS, int(weight.shape[1])) + + +def _pack_deepseek_v4_optional_attention( + prefix: str, + raw: Mapping[str, torch.Tensor], + ranks: int, + *, + compress_ratio: int, +) -> dict[str, torch.Tensor]: + """Pack compressor/indexer tensors, filling inactive branch placeholders.""" + + def raw_tensor(suffix: str) -> torch.Tensor | None: + return raw.get(f"{prefix}.{suffix}") + + def zeros(shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + return torch.zeros((ranks, *shape), dtype=dtype) + + def replicated(tensor: torch.Tensor, *, dtype: torch.dtype | None = None) -> torch.Tensor: + if dtype is not None: + tensor = tensor.to(dtype=dtype) + tensor = tensor.contiguous().cpu() + return tensor.unsqueeze(0).expand(ranks, *tensor.shape).contiguous() + + def replicated_transposed( + suffix: str, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + enabled: bool, + ) -> torch.Tensor: + tensor = raw_tensor(suffix) if enabled else None + if tensor is None: + return zeros(shape, dtype) + return replicated(tensor.transpose(0, 1).contiguous(), dtype=dtype) + + def replicated_plain( + suffix: str, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + enabled: bool, + ) -> torch.Tensor: + tensor = raw_tensor(suffix) if enabled else None + if tensor is None: + return zeros(shape, dtype) + return replicated(tensor, dtype=dtype) + + is_hca = int(compress_ratio) == _DEEPSEEK_V4_HCA_COMPRESS_RATIO + is_csa = int(compress_ratio) == _DEEPSEEK_V4_CSA_COMPRESS_RATIO + return { + "hca_cmp_wkv": replicated_plain( + "attn.compressor.wkv.weight", + (_DEEPSEEK_V4_HCA_MAIN_OUT_DIM, _DEEPSEEK_V4_HIDDEN_SIZE), + torch.bfloat16, + enabled=is_hca, + ), + "hca_cmp_wgate": replicated_plain( + "attn.compressor.wgate.weight", + (_DEEPSEEK_V4_HCA_MAIN_OUT_DIM, _DEEPSEEK_V4_HIDDEN_SIZE), + torch.bfloat16, + enabled=is_hca, + ), + "hca_cmp_ape": replicated_plain( + "attn.compressor.ape", + (_DEEPSEEK_V4_HCA_COMPRESS_RATIO, _DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + torch.float32, + enabled=is_hca, + ), + "hca_cmp_norm_w": replicated_plain( + "attn.compressor.norm.weight", + (_DEEPSEEK_V4_HEAD_DIM,), + torch.bfloat16, + enabled=is_hca, + ), + "csa_cmp_wkv": replicated_plain( + "attn.compressor.wkv.weight", + (_DEEPSEEK_V4_CSA_MAIN_OUT_DIM, _DEEPSEEK_V4_HIDDEN_SIZE), + torch.bfloat16, + enabled=is_csa, + ), + "csa_cmp_wgate": replicated_plain( + "attn.compressor.wgate.weight", + (_DEEPSEEK_V4_CSA_MAIN_OUT_DIM, _DEEPSEEK_V4_HIDDEN_SIZE), + torch.bfloat16, + enabled=is_csa, + ), + "csa_cmp_ape": replicated_plain( + "attn.compressor.ape", + (_DEEPSEEK_V4_CSA_COMPRESS_RATIO, _DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + torch.float32, + enabled=is_csa, + ), + "csa_cmp_norm_w": replicated_plain( + "attn.compressor.norm.weight", + (_DEEPSEEK_V4_HEAD_DIM,), + torch.bfloat16, + enabled=is_csa, + ), + "csa_idx_wq_b": replicated_transposed( + "attn.indexer.wq_b.weight", + (_DEEPSEEK_V4_Q_LORA, _DEEPSEEK_V4_ATTENTION_OUT // 4), + torch.int8, + enabled=is_csa, + ), + "csa_idx_wq_b_scale": replicated_plain( + "attn.indexer.wq_b.scale", + (_DEEPSEEK_V4_ATTENTION_OUT // 4,), + torch.float32, + enabled=is_csa, + ), + "csa_weights_proj": replicated_transposed( + "attn.indexer.weights_proj.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, 64), + torch.bfloat16, + enabled=is_csa, + ), + "csa_hadamard_idx": replicated(deepseek_v4_hadamard_idx(), dtype=torch.bfloat16), + "csa_inner_wkv": replicated_plain( + "attn.indexer.compressor.wkv.weight", + (_DEEPSEEK_V4_CSA_INNER_OUT_DIM, _DEEPSEEK_V4_HIDDEN_SIZE), + torch.bfloat16, + enabled=is_csa, + ), + "csa_inner_wgate": replicated_plain( + "attn.indexer.compressor.wgate.weight", + (_DEEPSEEK_V4_CSA_INNER_OUT_DIM, _DEEPSEEK_V4_HIDDEN_SIZE), + torch.bfloat16, + enabled=is_csa, + ), + "csa_inner_ape": replicated_plain( + "attn.indexer.compressor.ape", + (_DEEPSEEK_V4_CSA_COMPRESS_RATIO, _DEEPSEEK_V4_CSA_INNER_OUT_DIM), + torch.float32, + enabled=is_csa, + ), + "csa_inner_norm_w": replicated_plain( + "attn.indexer.compressor.norm.weight", + (_DEEPSEEK_V4_HADAMARD_IDX_DIM,), + torch.bfloat16, + enabled=is_csa, + ), + } + + +def _pack_deepseek_v4_router( + prefix: str, + raw: Mapping[str, torch.Tensor], + *, + ranks: int, + n_routed_experts: int, + include_tid2eid: bool, + include_gate_bias: bool, +) -> dict[str, torch.Tensor]: + """Pack router-only tensors and placeholders for inactive router modes.""" + gate_bias = raw.get(f"{prefix}.ffn.gate.bias") + if gate_bias is None: + if include_gate_bias: + raise KeyError(f"missing raw DeepSeekV4 layer tensor: {prefix}.ffn.gate.bias") + gate_bias = torch.zeros((n_routed_experts,), dtype=torch.float32) + tid2eid = raw.get(f"{prefix}.ffn.gate.tid2eid") + if tid2eid is None: + if include_tid2eid: + raise KeyError(f"missing raw DeepSeekV4 layer tensor: {prefix}.ffn.gate.tid2eid") + tid2eid = torch.zeros((_DEEPSEEK_V4_VOCAB_SIZE, _DEEPSEEK_V4_TOPK), dtype=torch.int32) + return { + "gate_bias": gate_bias.to(torch.float32).contiguous().cpu().unsqueeze(0).expand(ranks, -1).contiguous(), + "tid2eid": tid2eid.to(torch.int32).contiguous().cpu().unsqueeze(0).expand(ranks, *tid2eid.shape).contiguous(), + } + + +def _pack_deepseek_v4_routed_experts( + prefix: str, + raw: Mapping[str, torch.Tensor], + *, + ranks: int, + n_routed_experts: int, +) -> dict[str, torch.Tensor]: + """Stack rank-local routed experts into EP-rank-major tensors.""" + + def expert(expert_id: int, suffix: str) -> torch.Tensor: + name = f"{prefix}.ffn.experts.{expert_id}.{suffix}" + try: + return raw[name].contiguous().cpu() + except KeyError as exc: + raise KeyError(f"missing raw DeepSeekV4 expert tensor: {name}") from exc + + def stack(suffix: str, dtype: torch.dtype) -> torch.Tensor: + per_rank = [] + for rank in range(ranks): + ids = deepseek_v4_local_expert_ids( + rank=rank, + ranks=ranks, + n_routed_experts=n_routed_experts, + ) + per_rank.append(torch.stack([expert(expert_id, suffix).to(dtype=dtype) for expert_id in ids], dim=0)) + return torch.stack(per_rank, dim=0).contiguous() + + return { + "routed_w1": stack("w1.weight", torch.int8), + "routed_w1_scale": stack("w1.scale", torch.float32), + "routed_w3": stack("w3.weight", torch.int8), + "routed_w3_scale": stack("w3.scale", torch.float32), + "routed_w2": stack("w2.weight", torch.int8), + "routed_w2_scale": stack("w2.scale", torch.float32), + } diff --git a/pypto-lib b/pypto-lib index 57772f3..c159c32 160000 --- a/pypto-lib +++ b/pypto-lib @@ -1 +1 @@ -Subproject commit 57772f304bbcaee927b51227f6aa495a5591debf +Subproject commit c159c325a8d4279d0af65fbd528c9096bfb6a57a diff --git a/python/cli/main.py b/python/cli/main.py index a6f0067..b648742 100644 --- a/python/cli/main.py +++ b/python/cli/main.py @@ -11,6 +11,7 @@ import argparse import contextlib +import json import os import sys from collections.abc import Iterator, Sequence @@ -134,27 +135,36 @@ def build_serving_engine_config(args: argparse.Namespace) -> EngineConfig: model_dir = str(Path(args.model).resolve()) executor_kwargs = _build_executor_kwargs() devices = parse_device_ids(args.devices, default_device=args.device) + model_family = _detect_model_family(Path(model_dir)) + if model_family == "deepseek_v4": + executor_kwargs["compile_kernels"] = True parallel_config = ParallelConfig( data_parallel_size=args.data_parallel_size, tensor_parallel_size=args.tensor_parallel_size, devices=devices, data_parallel_routing=args.data_parallel_routing, ) + _validate_model_topology(model_family, args, parallel_config) first_group = parallel_config.replica_device_groups[0] + worker_device_ids = first_group if parallel_config.data_parallel_size == 1 else () + enable_prefix_cache = args.enable_prefix_caching + if model_family == "deepseek_v4": + enable_prefix_cache = False return EngineConfig( model_id=args.served_model_name or Path(args.model).name, model_dir=model_dir, platform=args.platform, device_id=first_group[0], + device_ids=worker_device_ids, parallel_config=parallel_config, - executor_cls="PyptoQwen14BExecutor", + executor_cls=_executor_cls_for_model_family(model_family), executor_kwargs=executor_kwargs, runtime_config=_build_runtime_config(args), max_num_running_reqs=args.max_num_seqs, max_num_scheduled_tokens=args.max_num_batched_tokens, long_prefill_token_threshold=args.long_prefill_token_threshold, - enable_prefix_cache=args.enable_prefix_caching, + enable_prefix_cache=enable_prefix_cache, enable_chunk_prefill=args.enable_chunked_prefill, ) @@ -187,6 +197,60 @@ def _build_executor_kwargs() -> dict[str, object]: return executor_kwargs +def _detect_model_family(model_dir: Path) -> str: + """Return the serving model family inferred from config.json.""" + config_path = model_dir / "config.json" + if not config_path.exists(): + return "qwen" + try: + config_data = json.loads(config_path.read_text()) + except json.JSONDecodeError: + return "qwen" + model_type = str(config_data.get("model_type") or "").lower() + architectures = {str(item).lower() for item in (config_data.get("architectures") or [])} + if model_type == "deepseek_v4" or "deepseekv4forcausallm" in architectures: + return "deepseek_v4" + return "qwen" + + +def _executor_cls_for_model_family(model_family: str) -> str: + """Map model family metadata to the worker executor class id.""" + if model_family == "deepseek_v4": + return "PyptoDeepSeekV4Executor" + return "PyptoQwen14BExecutor" + + +def _validate_model_topology( + model_family: str, + args: argparse.Namespace, + parallel_config, +) -> None: + """Validate model-specific serving topology constraints.""" + if model_family != "deepseek_v4": + return + config_data = json.loads((Path(args.model).resolve() / "config.json").read_text()) + quantization = config_data.get("quantization_config") or {} + if quantization.get("quant_method") != "compressed-tensors": + raise ValueError( + "DeepSeekV4 serving requires the quantized W8A8 compressed-tensors checkpoint " + "such as /data/models/dsv4-flash-w8a8; the original checkpoint is too large for 8 NPUs." + ) + if parallel_config.data_parallel_size != 1 or parallel_config.tensor_parallel_size != 8: + raise ValueError("DeepSeekV4 serving requires --dp 1 --tp 8") + if len(parallel_config.devices) != 8: + raise ValueError("DeepSeekV4 serving requires exactly 8 NPU device ids") + if args.block_size != 128: + raise ValueError("DeepSeekV4 kernels require --block-size 128") + if args.max_num_seqs > 64: + raise ValueError("DeepSeekV4 decode kernels support at most --max-num-seqs 64") + if args.max_model_len > 260: + raise ValueError( + "DeepSeekV4 pypto-lib decode CSA state tables currently support at most " + "--max-model-len 260. Increase the decode CSA state table depth in pypto-lib " + "before serving longer contexts." + ) + + def run_serve( config: EngineConfig, *, @@ -281,7 +345,7 @@ def _ensure_core_imports() -> None: try: from ..core.parallel import ParallelConfig as imported_parallel_config from ..core.parallel import parse_device_ids as imported_parse_device_ids - except ImportError: + except (ImportError, ValueError): from python.core.parallel import ParallelConfig as imported_parallel_config from python.core.parallel import parse_device_ids as imported_parse_device_ids diff --git a/python/core/async_engine.py b/python/core/async_engine.py index 077dd60..7ae3f15 100644 --- a/python/core/async_engine.py +++ b/python/core/async_engine.py @@ -12,6 +12,7 @@ import asyncio import contextlib import logging +import os import queue import time from collections.abc import AsyncGenerator, Sequence @@ -26,6 +27,33 @@ from .serving_worker import spawn_worker logger = logging.getLogger(__name__) +_DEFAULT_WORKER_INIT_TIMEOUT_SECONDS = 1800.0 +_DEFAULT_WORKER_STEP_TIMEOUT_SECONDS = 300.0 +_DEFAULT_DEEPSEEK_V4_WORKER_STEP_TIMEOUT_SECONDS = 1200.0 + + +def _positive_env_timeout_seconds(name: str, default: float) -> float: + raw = os.environ.get(name) + if raw is None or raw.strip() == "": + return default + try: + timeout = float(raw) + except ValueError as exc: + raise ValueError(f"{name} must be a positive number of seconds") from exc + if timeout <= 0: + raise ValueError(f"{name} must be a positive number of seconds") + return timeout + + +def _worker_init_timeout_seconds() -> float: + return _positive_env_timeout_seconds("PYPTO_WORKER_INIT_TIMEOUT", _DEFAULT_WORKER_INIT_TIMEOUT_SECONDS) + + +def _worker_step_timeout_seconds(executor_cls: str = "") -> float: + default = _DEFAULT_WORKER_STEP_TIMEOUT_SECONDS + if executor_cls == "PyptoDeepSeekV4Executor": + default = _DEFAULT_DEEPSEEK_V4_WORKER_STEP_TIMEOUT_SECONDS + return _positive_env_timeout_seconds("SERVING_WORKER_STEP_TIMEOUT", default) @dataclass @@ -150,11 +178,15 @@ async def start(self) -> None: logger.info("Waiting for worker to initialize model...") try: - ready = await asyncio.to_thread(ready_event.wait, timeout=600) + init_timeout = _worker_init_timeout_seconds() + ready = await asyncio.to_thread(ready_event.wait, timeout=init_timeout) if not ready: - raise RuntimeError("Worker failed to initialize within timeout") + raise RuntimeError( + f"Worker failed to initialize within {init_timeout:g}s timeout; " + "set PYPTO_WORKER_INIT_TIMEOUT to allow more time for large checkpoints" + ) except BaseException: - self._shutdown_worker(timeout=5) + await asyncio.to_thread(self._shutdown_worker, timeout=5) raise logger.info("Worker ready") @@ -182,7 +214,7 @@ async def stop(self) -> None: await self._loop_task self._loop_task = None - self._shutdown_worker(timeout=30) + await asyncio.to_thread(self._shutdown_worker, timeout=30) logger.info("ReplicaEngineCore stopped") def generate_request_id(self) -> str: @@ -308,11 +340,12 @@ async def _engine_loop(self) -> None: try: with profile_span("scheduler.wait_worker_output", cat="scheduler"): + step_timeout = _worker_step_timeout_seconds(self.config.executor_cls) step_output: StepOutput = await asyncio.to_thread( - self._output_queue.get, timeout=300 + self._output_queue.get, timeout=step_timeout ) except queue.Empty: - logger.error("Worker response timed out (300s)") + logger.error(f"Worker response timed out ({step_timeout:g}s)") self._handle_step_error(scheduler_output) continue @@ -371,12 +404,15 @@ def _process_step_output( def _handle_step_error(self, scheduler_output: SchedulerOutput) -> None: """On worker error, abort all requests in the failed batch.""" for sr in scheduler_output.scheduled_requests: - ctx = self._request_contexts.get(sr.request.request_id) + request_id = sr.request.request_id + ctx = self._request_contexts.get(request_id) if ctx is not None: ctx.queue.put_nowait( TokenOutput(finished=True, finish_reason="error") ) - self.scheduler.abort_request(sr.request.request_id) + if request_id not in self._pending_free_ids: + self._pending_free_ids.append(request_id) + self.scheduler.abort_request(request_id) def _shutdown_worker(self, *, timeout: float) -> None: input_q = self._input_queue @@ -470,8 +506,7 @@ async def start(self) -> None: async def stop(self) -> None: """Stop all DP engine cores.""" - for core in reversed(self._cores): - await core.stop() + await asyncio.gather(*(core.stop() for core in reversed(self._cores))) def generate_request_id(self) -> str: self._request_counter += 1 diff --git a/python/core/model_loader.py b/python/core/model_loader.py index 8db02a1..1a3e662 100644 --- a/python/core/model_loader.py +++ b/python/core/model_loader.py @@ -66,6 +66,41 @@ def load(self, request: ModelLoadRequest) -> LoadedModel: raise NotImplementedError +def _load_safetensors_weight_map(model_dir: Path) -> dict[str, str]: + """Return the ``{tensor_name: shard_filename}`` map from a safetensors index. + + Reads ``model.safetensors.index.json`` when present; otherwise synthesizes a + flat map from every ``*.safetensors`` shard in the directory (each shard + owning all of its tensors, which is only correct for the single-shard case). + Shared by every format loader that needs checkpoint layout without reading + tensor payloads. + """ + index_path = model_dir / "model.safetensors.index.json" + if index_path.exists(): + index_data = json.loads(index_path.read_text()) + return dict(index_data.get("weight_map", {})) + shards = sorted(path.name for path in model_dir.glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"No .safetensors files found in {model_dir}") + if len(shards) == 1: + try: + from safetensors import safe_open + + with safe_open(str(model_dir / shards[0]), framework="pt", device="cpu") as reader: + return {name: shards[0] for name in reader.keys()} + except ImportError: + raise RuntimeError("safetensors is required to read weight names from a single-shard checkpoint.") + raise FileNotFoundError( + f"{model_dir} has multiple safetensors shards but no model.safetensors.index.json; " + "the index is required to map tensor names to shards." + ) + + +def _safetensors_shard_filenames(weight_map: dict[str, str]) -> list[str]: + """Return the sorted unique shard filenames referenced by a weight map.""" + return sorted(set(weight_map.values())) + + def _load_safetensors_dir(model_dir: Path) -> dict[str, torch.Tensor]: """Load all safetensors shards from a local Hugging Face directory.""" try: @@ -73,15 +108,8 @@ def _load_safetensors_dir(model_dir: Path) -> dict[str, torch.Tensor]: except ImportError as exc: raise RuntimeError("safetensors is required to load weights from a local model directory.") from exc - index_path = model_dir / "model.safetensors.index.json" - if index_path.exists(): - index_data = json.loads(index_path.read_text()) - filenames = sorted(set(index_data["weight_map"].values())) - else: - filenames = sorted(path.name for path in model_dir.glob("*.safetensors")) - if not filenames: - raise FileNotFoundError(f"No .safetensors files found in {model_dir}") - + weight_map = _load_safetensors_weight_map(model_dir) + filenames = _safetensors_shard_filenames(weight_map) state_dict: dict[str, torch.Tensor] = {} for filename in filenames: state_dict.update(load_file(str(model_dir / filename))) @@ -103,6 +131,20 @@ def _optional_tensor(state_dict: dict[str, torch.Tensor], names: list[str]) -> t return None +def _build_tokenizer(model_path: Path, *, trust_remote_code: bool = False) -> TokenizerAdapter: + """Build a tokenizer adapter, preferring a fast ``tokenizer.json`` when present. + + Shared by every format loader so new models do not re-implement the + ``tokenizer.json``-first, ``from_pretrained``-fallback heuristic. + """ + if (model_path / "tokenizer.json").exists(): + return TransformersTokenizerAdapter.from_tokenizer_file(str(model_path)) + return TransformersTokenizerAdapter.from_pretrained( + str(model_path), + trust_remote_code=trust_remote_code, + ) + + def _build_model_config(model_id: str, config_data: dict, tokenizer: TokenizerAdapter) -> ModelConfig: """Build internal model metadata from Hugging Face config JSON.""" hidden_size = int(config_data["hidden_size"]) @@ -188,10 +230,7 @@ def _mark(label: str) -> None: raise FileNotFoundError(f"Missing config.json in {model_path}") trust_remote_code = bool(request.loader_options.get("trust_remote_code", False)) - tokenizer = TransformersTokenizerAdapter.from_pretrained( - str(model_path), - trust_remote_code=trust_remote_code, - ) + tokenizer = _build_tokenizer(model_path, trust_remote_code=trust_remote_code) _mark("load_tokenizer") config_data = json.loads(config_path.read_text()) config = _build_model_config(request.model_id, config_data, tokenizer) @@ -276,12 +315,146 @@ def _mark(label: str) -> None: ) +class DeepSeekV4W8A8DirectoryLoader: + """Lazy loader for the local DeepSeekV4 Flash W8A8 checkpoint.""" + + format_names = ("deepseek_v4_w8a8", "deepseek-v4-w8a8", "dsv4-w8a8") + + def supports_format(self, model_format: str) -> bool: + """Return whether ``model_format`` names the DeepSeekV4 W8A8 loader.""" + return model_format.lower() in self.format_names + + def can_load(self, model_path: Path) -> bool: + """Detect a DeepSeekV4 compressed-tensors checkpoint directory.""" + config_path = model_path / "config.json" + index_path = model_path / "model.safetensors.index.json" + if not config_path.exists() or not index_path.exists(): + return False + try: + config_data = json.loads(config_path.read_text()) + except json.JSONDecodeError: + return False + return _is_deepseek_v4_config(config_data) + + def load(self, request: ModelLoadRequest) -> LoadedModel: + """Load tokenizer and metadata without materializing all quantized weights.""" + model_path = Path(request.model_dir) + config_path = model_path / "config.json" + index_path = model_path / "model.safetensors.index.json" + if not config_path.exists(): + raise FileNotFoundError(f"Missing config.json in {model_path}") + if not index_path.exists(): + raise FileNotFoundError(f"Missing model.safetensors.index.json in {model_path}") + + config_data = json.loads(config_path.read_text()) + if not _is_deepseek_v4_config(config_data): + raise ValueError(f"{model_path} is not a DeepSeekV4 checkpoint") + quantization = config_data.get("quantization_config", {}) + if quantization.get("quant_method") != "compressed-tensors": + raise ValueError( + "DeepSeekV4 serving requires the W8A8 compressed-tensors checkpoint; " + f"got quant_method={quantization.get('quant_method')!r}" + ) + + trust_remote_code = bool(request.loader_options.get("trust_remote_code", False)) + tokenizer = _build_tokenizer(model_path, trust_remote_code=trust_remote_code) + config = _build_deepseek_v4_model_config(request.model_id, config_data, tokenizer) + runtime = request.runtime_config or RuntimeConfig(max_seq_len=min(config.max_position_embeddings, 8192)) + layer_specs = _build_layer_specs(config) + weight_map = _load_safetensors_weight_map(model_path) + _validate_deepseek_v4_weight_index(weight_map, config_data) + + placeholder = torch.empty(0, config.hidden_size, dtype=torch.bfloat16) + runtime_model = RuntimeModel( + config=config, + runtime=runtime, + embed_tokens=placeholder, + final_norm_weight=torch.empty(0, dtype=torch.bfloat16), + lm_head=placeholder, + layers=[], + extra={ + "family": "deepseek_v4", + "checkpoint_format": "w8a8-compressed-tensors", + "config_data": config_data, + "quantization_config": quantization, + "weight_map": weight_map, + "model_dir": str(model_path), + "compress_ratios": tuple(int(ratio) for ratio in config_data["compress_ratios"]), + }, + ) + + return LoadedModel( + model_id=request.model_id, + model_dir=str(model_path), + config=config, + tokenizer=tokenizer, + layer_specs=layer_specs, + runtime_model=runtime_model, + ) + + +def _is_deepseek_v4_config(config_data: dict) -> bool: + """Return whether config metadata names DeepSeekV4.""" + model_type = str(config_data.get("model_type", "")).lower() + architectures = {str(item).lower() for item in config_data.get("architectures", [])} + return model_type == "deepseek_v4" or "deepseekv4forcausallm" in architectures + + +def _build_deepseek_v4_model_config( + model_id: str, + config_data: dict, + tokenizer: TokenizerAdapter, +) -> ModelConfig: + """Build internal metadata for DeepSeekV4 Flash.""" + return ModelConfig( + model_id=model_id, + architecture=str(config_data.get("architectures", ["DeepseekV4ForCausalLM"])[0]), + vocab_size=int(config_data["vocab_size"]), + hidden_size=int(config_data["hidden_size"]), + intermediate_size=int(config_data["moe_intermediate_size"]), + num_hidden_layers=int(config_data["num_hidden_layers"]), + num_attention_heads=int(config_data["num_attention_heads"]), + num_key_value_heads=int(config_data.get("num_key_value_heads", 1)), + head_dim=int(config_data["head_dim"]), + max_position_embeddings=int(config_data["max_position_embeddings"]), + rms_norm_eps=float(config_data["rms_norm_eps"]), + rope_theta=float(config_data["rope_theta"]), + bos_token_id=config_data.get("bos_token_id", tokenizer.bos_token_id), + eos_token_id=config_data.get("eos_token_id", tokenizer.eos_token_id), + pad_token_id=config_data.get("pad_token_id", tokenizer.pad_token_id), + torch_dtype=str(config_data.get("torch_dtype", "bfloat16")), + ) + + +def _validate_deepseek_v4_weight_index(weight_map: dict[str, str], config_data: dict) -> None: + """Fail early if the W8A8 checkpoint does not expose required tensor names.""" + required = [ + "embed.weight", + "norm.weight", + "head.weight", + "layers.0.attn.wq_b.weight", + "layers.0.attn.wq_b.scale", + "layers.0.attn.wo_b.weight", + "layers.0.attn.wo_b.scale", + "layers.0.ffn.experts.0.w1.weight", + "layers.0.ffn.experts.0.w1.scale", + ] + missing = [name for name in required if name not in weight_map] + if missing: + raise KeyError(f"DeepSeekV4 W8A8 checkpoint is missing required tensors: {', '.join(missing)}") + ratios = config_data.get("compress_ratios") + if not isinstance(ratios, list) or len(ratios) != int(config_data["num_hidden_layers"]) + 1: + raise ValueError( + "DeepSeekV4 config compress_ratios must include one entry per hidden layer plus MTP/final entry" + ) + + class ModelLoader: """Registry that selects a model-format loader and loads models.""" def __init__(self, format_loaders: list[ModelFormatLoader] | None = None) -> None: """Create a loader registry with optional custom format loaders.""" - self._format_loaders = format_loaders or [HuggingFaceDirectoryLoader()] + self._format_loaders = format_loaders or [DeepSeekV4W8A8DirectoryLoader(), HuggingFaceDirectoryLoader()] def register(self, format_loader: ModelFormatLoader) -> None: """Register an additional model format loader.""" diff --git a/python/core/model_runner.py b/python/core/model_runner.py index 819c059..c602366 100644 --- a/python/core/model_runner.py +++ b/python/core/model_runner.py @@ -21,6 +21,7 @@ DecodeBatch, DecodeResult, ModelConfig, + ModelRecord, PrefillBatch, PrefillResult, RuntimeConfig, @@ -91,6 +92,17 @@ def close_kv_cache(self) -> None: self._free_kv_cache_tensor(pool.value_pages) self._kv_caches.clear() + def preflight(self, record: ModelRecord) -> None: + """Eagerly materialize weights and shared buffers before the worker signals ready. + + Called once by ``PyptoExecutor.register_model`` after ``init_kv_cache``, + before the serving worker sets its ready event. Runners whose weights are + already loaded by the model loader (e.g. HuggingFace eager load) inherit + this no-op; runners that defer weight reads to first inference override + it so the serving "ready" contract means weights are fully loaded. + """ + del record + @abstractmethod def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> DeviceTensor: """Allocate one worker-resident KV cache tensor.""" diff --git a/python/core/pypto_executor.py b/python/core/pypto_executor.py index 026d8a1..3b4edc4 100644 --- a/python/core/pypto_executor.py +++ b/python/core/pypto_executor.py @@ -53,16 +53,22 @@ def __init__( self._runners: dict[str, ModelRunner] = {} self._compiled: dict[str, object] = {} - def register_model(self, model_id: str, record: ModelRecord) -> int: + def register_model(self, model_id: str, record: ModelRecord) -> None: """Compile kernels for ``record`` and attach a runner to ``model_id``. Returns the number of KV cache pages allocated on the device so the caller can synchronise host-side block metadata. """ print("[register_model] compiling kernels …", flush=True) + def register_model(self, model_id: str, record: ModelRecord) -> None: + """Compile kernels for ``record`` and attach a runner to ``model_id``.""" + import time + with profile_span("PyptoExecutor.register_model", cat="executor", args={"model_id": model_id}): + start_t0 = time.perf_counter() compiled = self._compile_model(record.runtime_model) runner = self._create_runner(model_id, compiled) + try: num_pages = runner.init_kv_cache(model_id, record.config, record.runtime) except Exception: @@ -70,7 +76,14 @@ def register_model(self, model_id: str, record: ModelRecord) -> int: if callable(close): close() raise - self._compiled[model_id] = compiled + + with profile_span("PyptoExecutor.preflight", cat="executor", args={"model_id": model_id}): + runner.preflight(record) + logger.info( + "PyptoExecutor %s: model loaded (%.1fs total)", + model_id, + time.perf_counter() - start_t0, + ) self._runners[model_id] = runner return num_pages diff --git a/python/core/serving_worker.py b/python/core/serving_worker.py index 9576d3b..93b470e 100644 --- a/python/core/serving_worker.py +++ b/python/core/serving_worker.py @@ -12,6 +12,7 @@ import logging import multiprocessing as mp import os +import sys from pathlib import Path import torch @@ -120,6 +121,9 @@ def _resolve_executor_cls(self): if self.config.executor_cls == "PyptoQwen14BExecutor": from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor return Qwen314BPyptoExecutor + if self.config.executor_cls == "PyptoDeepSeekV4Executor": + from examples.model.deepseek_v4.runner.npu_executor import DeepSeekV4PyptoExecutor + return DeepSeekV4PyptoExecutor from .executor import ModelExecutor return ModelExecutor @@ -144,7 +148,9 @@ def busy_loop(self) -> None: break elif cmd.type == "step": if cmd.finished_request_ids: - pass # No allocation cleanup needed + release_finished = getattr(self.executor, "release_finished_requests", None) + if callable(release_finished): + release_finished(cmd.finished_request_ids) try: result = self._execute_step(cmd.scheduler_output) @@ -288,6 +294,7 @@ def _batch_decode( device = runtime_model.runtime.device decode_tokens = [] + prev_tokens = [] block_ids_list = [] seq_lens = [] allow_device_greedy_sampling = ( @@ -297,12 +304,19 @@ def _batch_decode( for sr in scheduled: request = sr.request - last_token = ( - request.output_token_ids[-1] - if request.output_token_ids - else request.prompt_token_ids[-1] - ) + output_ids = request.output_token_ids + prompt_ids = request.prompt_token_ids + last_token = output_ids[-1] if output_ids else prompt_ids[-1] + # Token at absolute position ``seq_len-2``; guard the single-token + # edge so we never index out of range. + if len(output_ids) >= 2: + prev_token = output_ids[-2] + elif output_ids and prompt_ids: + prev_token = prompt_ids[-1] + else: + prev_token = last_token decode_tokens.append(last_token) + prev_tokens.append(prev_token) block_ids_list.append(sr.block_ids) seq_lens.append(request.num_tokens) @@ -313,8 +327,14 @@ def _batch_decode( dtype=runtime_model.embed_tokens.dtype, device=device, ) + prev_embeddings = torch.zeros_like(decode_embeddings) else: decode_embeddings = self.executor.lookup_embeddings(runtime_model, decode_token_tensor) + 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: + prev_token_tensor = torch.tensor(prev_tokens, dtype=torch.long, device=device) decode_result = self.executor.run_decode( runtime_model, @@ -325,6 +345,8 @@ def _batch_decode( seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), allow_device_greedy_sampling=allow_device_greedy_sampling, block_ids=block_ids_list, + prev_token_ids=prev_token_tensor, + prev_hidden_states=prev_embeddings, ), ) @@ -382,6 +404,17 @@ def _worker_entry( for _n in ("simpler_setup", "pypto", "simpler"): logging.getLogger(_n).setLevel(logging.WARNING) + # Spawned workers do not inherit the parent's logging config; configure a + # stderr handler so per-stage progress logs (weight load, preflight) are + # visible alongside kernel/perf output. + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s.%(msecs)03d %(levelname)s | %(message)s", + datefmt="%H:%M:%S", + stream=sys.stderr, + force=True, + ) + worker = WorkerProcess(config, input_queue, output_queue) try: num_pages = worker.init_device_and_model() diff --git a/python/core/tokenizer.py b/python/core/tokenizer.py index 5abb7b2..353ccf9 100644 --- a/python/core/tokenizer.py +++ b/python/core/tokenizer.py @@ -9,10 +9,15 @@ from __future__ import annotations +import json +import logging from dataclasses import dataclass from pathlib import Path +logger = logging.getLogger(__name__) + + class TokenizerAdapter: """Minimal tokenizer interface required by the generation engine.""" @@ -50,20 +55,41 @@ class TransformersTokenizerAdapter(TokenizerAdapter): def from_pretrained(cls, model_dir: str, trust_remote_code: bool = False) -> "TransformersTokenizerAdapter": """Load a local Hugging Face tokenizer directory.""" try: - from transformers import AutoTokenizer + from transformers import AutoTokenizer, PreTrainedTokenizerFast except ImportError as exc: raise RuntimeError( "transformers is required for the current local Hugging Face tokenizer adapter." ) from exc - tokenizer = AutoTokenizer.from_pretrained( - str(Path(model_dir)), - local_files_only=True, - trust_remote_code=trust_remote_code, - use_fast=True, - ) + model_path = Path(model_dir) + try: + tokenizer = AutoTokenizer.from_pretrained( + str(model_path), + local_files_only=True, + trust_remote_code=trust_remote_code, + use_fast=True, + ) + except (OSError, ValueError, AttributeError) as exc: + logger.warning( + "AutoTokenizer.from_pretrained failed for %s: %s; falling back to local tokenizer.json", + model_path, + exc, + ) + tokenizer = _load_fast_tokenizer_from_file(model_path, PreTrainedTokenizerFast) return cls(tokenizer=tokenizer) + @classmethod + def from_tokenizer_file(cls, model_dir: str) -> "TransformersTokenizerAdapter": + """Load ``tokenizer.json`` directly without consulting model config.""" + try: + from transformers import PreTrainedTokenizerFast + except ImportError as exc: + raise RuntimeError( + "transformers is required for the current local Hugging Face tokenizer adapter." + ) from exc + + return cls(tokenizer=_load_fast_tokenizer_from_file(Path(model_dir), PreTrainedTokenizerFast)) + def encode(self, text: str) -> list[int]: """Encode text using the wrapped Hugging Face tokenizer.""" return list(self.tokenizer.encode(text, add_special_tokens=False)) @@ -86,3 +112,26 @@ def eos_token_id(self) -> int | None: def pad_token_id(self) -> int | None: """Return the wrapped tokenizer PAD token ID.""" return self.tokenizer.pad_token_id + + +def _token_content(value: object) -> str | None: + """Extract a special token string from tokenizer_config JSON.""" + if isinstance(value, dict): + content = value.get("content") + return content if isinstance(content, str) else None + return value if isinstance(value, str) else None + + +def _load_fast_tokenizer_from_file(model_path: Path, tokenizer_cls: type) -> object: + """Load a local tokenizer.json with special tokens from tokenizer_config.""" + tokenizer_file = model_path / "tokenizer.json" + if not tokenizer_file.exists(): + raise FileNotFoundError(f"Missing tokenizer.json in {model_path}") + config_path = model_path / "tokenizer_config.json" + tokenizer_config = json.loads(config_path.read_text()) if config_path.exists() else {} + special_tokens = { + name: _token_content(tokenizer_config.get(name)) + for name in ("bos_token", "eos_token", "pad_token", "unk_token") + if _token_content(tokenizer_config.get(name)) is not None + } + return tokenizer_cls(tokenizer_file=str(tokenizer_file), **special_tokens) diff --git a/python/core/types.py b/python/core/types.py index e261c46..9f0dc78 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -111,6 +111,7 @@ class RuntimeModel: final_norm_weight: torch.Tensor lm_head: torch.Tensor layers: list[LayerWeights] + extra: dict[str, object] = field(default_factory=dict) @dataclass @@ -211,6 +212,12 @@ class DecodeBatch: allow_device_greedy_sampling: bool = False kv_allocations: list[KvAllocation] = field(default_factory=list) block_ids: list[list[int]] = field(default_factory=list) + # Optional MTP context for models (e.g. DeepSeek V4) that decode two real + # trailing tokens per step. ``prev_token_ids`` holds the token id at absolute + # position ``seq_len-2`` per request (shape ``[B]``) and ``prev_hidden_states`` + # its embedding (shape ``[B, hidden]``). Left ``None`` for single-token decoders. + prev_token_ids: torch.Tensor | None = None + prev_hidden_states: torch.Tensor | None = None @dataclass diff --git a/tests/test_batching.py b/tests/test_batching.py index 43c5154..af2b8ee 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -7,6 +7,7 @@ # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +import asyncio from pathlib import Path from types import SimpleNamespace @@ -14,6 +15,7 @@ import torch from simpler.task_interface import DataType +from python.core.async_engine import ReplicaEngineCore, TokenOutput from python.core.engine import LLMEngine from python.core.executor import ModelExecutor from python.core.kv_cache import KvCacheManager @@ -54,6 +56,33 @@ def decode(self, token_ids: list[int]) -> str: return " ".join(str(token_id) for token_id in token_ids) +def test_worker_step_error_queues_finished_ids_for_executor_release(): + aborted: list[str] = [] + core = ReplicaEngineCore.__new__(ReplicaEngineCore) + core.scheduler = SimpleNamespace(abort_request=aborted.append) + core._pending_free_ids = [] + core._request_contexts = { + "req-a": SimpleNamespace(queue=asyncio.Queue()), + "req-b": SimpleNamespace(queue=asyncio.Queue()), + } + scheduler_output = SimpleNamespace( + scheduled_requests=[ + SimpleNamespace(request=SimpleNamespace(request_id="req-a")), + SimpleNamespace(request=SimpleNamespace(request_id="req-b")), + ] + ) + + core._handle_step_error(scheduler_output) + + assert aborted == ["req-a", "req-b"] + assert core._pending_free_ids == ["req-a", "req-b"] + for request_id in ("req-a", "req-b"): + token = core._request_contexts[request_id].queue.get_nowait() + assert isinstance(token, TokenOutput) + assert token.finished is True + assert token.finish_reason == "error" + + def _model( max_batch_size: int, max_seq_len: int = 128, @@ -547,7 +576,7 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc compiled=compiled, ) monkeypatch.setattr(runner, "_shared_l3_worker", lambda: _FakeWorker()) - monkeypatch.setattr(runner, "_compute_kv_cache_pages", lambda config, runtime: 1) + monkeypatch.setattr(runner, "_compute_kv_cache_pages", lambda config, runtime, device_id=0: 1) monkeypatch.setattr(runner, "_print_memory_breakdown", lambda *a, **kw: None) runner.init_kv_cache(model.config.model_id, model.config, model.runtime) monkeypatch.setattr(runner, "_static_device_tensor", lambda tensor: tensor) @@ -626,7 +655,10 @@ def test_decode_host_inlines_embedding_and_sampling_into_decode_fwd(): if not QWEN3_KERNEL_DIR.is_dir(): pytest.skip("pypto-lib submodule is not checked out") - decode_source = (QWEN3_KERNEL_DIR / "decode_layer.py").read_text(encoding="utf-8") + decode_path = QWEN3_KERNEL_DIR / "decode_layer.py" + if not decode_path.is_file(): + decode_path = QWEN3_KERNEL_DIR / "decode_fwd.py" + decode_source = decode_path.read_text(encoding="utf-8") assert 'name_hint="token_embed"' in decode_source assert 'name_hint="greedy_sample"' in decode_source diff --git a/tests/test_deepseek_v4.py b/tests/test_deepseek_v4.py new file mode 100644 index 0000000..74275e2 --- /dev/null +++ b/tests/test_deepseek_v4.py @@ -0,0 +1,1698 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import ctypes +import json +import sys +from pathlib import Path + +import pytest +import torch + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +import python.cli.main as cli +from python.core import async_engine +from python.core import tokenizer as tokenizer_module +from examples.model.deepseek_v4.runner import npu_executor, npu_runner, weight_loader +from examples.model.deepseek_v4.runner.npu_runner import ( + DeepSeekV4CacheLayout, + DeepSeekV4CacheManager, + DeepSeekV4CompiledKernels, + DeepSeekV4InputBuilder, + DeepSeekV4L3Callable, + DeepSeekV4LayerCache, + DeepSeekV4LayerCacheSnapshot, + DeepSeekV4LayerPlan, + DeepSeekV4ModelRunner, + build_deepseek_v4_layer_plan, + deepseek_v4_attention_kind, +) +from examples.model.deepseek_v4.runner.weight_loader import ( + DeepSeekV4WeightStore, + deepseek_v4_layer_core_weight_names, + deepseek_v4_hadamard_idx, + deepseek_v4_local_expert_ids, + deepseek_v4_routed_expert_weight_names, + deepseek_v4_startup_weight_names, + pack_deepseek_v4_lm_head_weight, + pack_deepseek_v4_layer_weights, +) +from python.core import model_loader +from python.core.model_loader import ModelLoader +from python.core.types import DecodeBatch, PrefillBatch, RuntimeConfig + + +def test_cli_selects_deepseek_executor_and_forces_prefix_cache_off(tmp_path): + model_dir = _write_deepseek_model_dir(tmp_path) + args = cli.build_parser().parse_args( + [ + "--model", str(model_dir), + "--devices", "0,1,2,3,4,5,6,7", + "--dp", "1", + "--tp", "8", + "--block-size", "128", + "--max-model-len", "260", + "--dtype", "int8", + ] + ) + + config = cli.build_serving_engine_config(args) + + assert config.executor_cls == "PyptoDeepSeekV4Executor" + assert config.device_ids == (0, 1, 2, 3, 4, 5, 6, 7) + assert config.parallel_config.replica_device_groups == ((0, 1, 2, 3, 4, 5, 6, 7),) + assert config.runtime_config.page_size == 128 + assert config.runtime_config.weight_dtype == "int8" + assert config.enable_prefix_cache is False + + +def test_tokenizer_falls_back_when_deepseek_config_raises_attribute_error(tmp_path, monkeypatch): + class AutoTokenizer: + @staticmethod + def from_pretrained(*args, **kwargs): + raise AttributeError("'PreTrainedConfig' object has no attribute 'max_position_embeddings'") + + sentinel = object() + fake_transformers = type( + "FakeTransformers", + (), + { + "AutoTokenizer": AutoTokenizer, + "PreTrainedTokenizerFast": object, + }, + ) + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + monkeypatch.setattr( + tokenizer_module, + "_load_fast_tokenizer_from_file", + lambda model_path, tokenizer_cls: sentinel, + ) + + adapter = tokenizer_module.TransformersTokenizerAdapter.from_pretrained(str(tmp_path)) + + assert adapter.tokenizer is sentinel + + +def test_cli_rejects_deepseek_non_w8a8_checkpoint(tmp_path): + model_dir = _write_deepseek_model_dir(tmp_path, quant_method="fp8") + args = cli.build_parser().parse_args( + [ + "--model", str(model_dir), + "--devices", "0,1,2,3,4,5,6,7", + "--dp", "1", + "--tp", "8", + "--block-size", "128", + ] + ) + + with pytest.raises(ValueError, match="compressed-tensors"): + cli.build_serving_engine_config(args) + + +def test_cli_rejects_deepseek_non_8_way_topology(tmp_path): + model_dir = _write_deepseek_model_dir(tmp_path) + args = cli.build_parser().parse_args( + [ + "--model", str(model_dir), + "--devices", "0,1,2,3", + "--dp", "1", + "--tp", "4", + "--block-size", "128", + ] + ) + + with pytest.raises(ValueError, match="--dp 1 --tp 8"): + cli.build_serving_engine_config(args) + + +def test_cli_rejects_deepseek_context_beyond_decode_state_capacity(tmp_path): + model_dir = _write_deepseek_model_dir(tmp_path) + args = cli.build_parser().parse_args( + [ + "--model", str(model_dir), + "--devices", "0,1,2,3,4,5,6,7", + "--dp", "1", + "--tp", "8", + "--block-size", "128", + "--max-model-len", "512", + ] + ) + + with pytest.raises(ValueError, match="--max-model-len 260"): + cli.build_serving_engine_config(args) + + +def test_deepseek_worker_step_timeout_default_allows_lazy_first_step(monkeypatch): + monkeypatch.delenv("SERVING_WORKER_STEP_TIMEOUT", raising=False) + + assert async_engine._worker_step_timeout_seconds("PyptoDeepSeekV4Executor") == 1200.0 + assert async_engine._worker_step_timeout_seconds("PyptoQwen14BExecutor") == 300.0 + + monkeypatch.setenv("SERVING_WORKER_STEP_TIMEOUT", "42") + assert async_engine._worker_step_timeout_seconds("PyptoDeepSeekV4Executor") == 42.0 + + +def test_deepseek_loader_keeps_w8a8_weights_lazy(tmp_path, monkeypatch): + model_dir = _write_deepseek_model_dir(tmp_path) + monkeypatch.setattr( + model_loader.TransformersTokenizerAdapter, + "from_pretrained", + lambda *args, **kwargs: _Tokenizer(), + ) + + loaded = ModelLoader().load( + model_id="dsv4", + model_dir=str(model_dir), + runtime_config=RuntimeConfig(page_size=128, max_batch_size=4, max_seq_len=256, weight_dtype="int8"), + ) + + assert loaded.config.architecture == "DeepseekV4ForCausalLM" + assert loaded.config.head_dim == 512 + assert loaded.runtime_model.layers == [] + assert loaded.runtime_model.embed_tokens.numel() == 0 + assert loaded.runtime_model.extra["family"] == "deepseek_v4" + assert loaded.runtime_model.extra["checkpoint_format"] == "w8a8-compressed-tensors" + assert "layers.0.attn.wq_b.scale" in loaded.runtime_model.extra["weight_map"] + assert "layers.2.attn.indexer.wq_b.scale" in loaded.runtime_model.extra["weight_map"] + assert "layers.3.attn.compressor.wkv.weight" in loaded.runtime_model.extra["weight_map"] + assert "layers.3.ffn.gate.bias" in loaded.runtime_model.extra["weight_map"] + + +def test_deepseek_compile_attaches_lazy_weight_store_without_opening_shards(tmp_path, monkeypatch): + model_dir = _write_deepseek_model_dir(tmp_path) + kernel_dir = _write_deepseek_kernel_dir(tmp_path, lm_head_tp_size=8) + monkeypatch.setattr( + model_loader.TransformersTokenizerAdapter, + "from_pretrained", + lambda *args, **kwargs: _Tokenizer(), + ) + opened: list[Path] = [] + + def _fail_open(path: Path, device: str): + opened.append(path) + raise AssertionError(f"unexpected safetensors open on {device}: {path}") + + monkeypatch.setattr(weight_loader, "_default_safe_open", _fail_open) + monkeypatch.setattr(npu_executor, "_find_pypto_lib_deepseek_v4_dir", lambda *args, **kwargs: kernel_dir) + loaded = ModelLoader().load( + model_id="dsv4", + model_dir=str(model_dir), + runtime_config=RuntimeConfig(page_size=128, max_batch_size=4, max_seq_len=256, weight_dtype="int8"), + ) + executor = npu_executor.DeepSeekV4PyptoExecutor(platform="a2a3sim", device_ids=tuple(range(8))) + + compiled = executor._compile_model(loaded.runtime_model) + + assert opened == [] + assert isinstance(compiled.weight_store, DeepSeekV4WeightStore) + assert compiled.weight_store.filename_for("head.weight") == "model-00001-of-00001.safetensors" + assert compiled.weight_store.device == "cpu" + assert compiled.layer_plan[0].attention_kind == "swa" + assert compiled.layer_plan[2].attention_kind == "csa" + assert compiled.layer_plan[2].include_tid2eid is True + assert compiled.layer_plan[3].attention_kind == "hca" + assert compiled.layer_plan[3].include_gate_bias is True + + +def test_deepseek_compile_builds_one_runtime_scalar_layer_callable(tmp_path, monkeypatch): + model_dir = _write_deepseek_model_dir(tmp_path) + kernel_dir = _write_deepseek_kernel_dir(tmp_path, lm_head_tp_size=8) + monkeypatch.setattr( + model_loader.TransformersTokenizerAdapter, + "from_pretrained", + lambda *args, **kwargs: _Tokenizer(), + ) + monkeypatch.setattr(npu_executor, "_find_pypto_lib_deepseek_v4_dir", lambda *args, **kwargs: kernel_dir) + loaded = ModelLoader().load( + model_id="dsv4", + model_dir=str(model_dir), + runtime_config=RuntimeConfig(page_size=128, max_batch_size=4, max_seq_len=256, weight_dtype="int8"), + ) + compiled_args: dict[str, tuple[object, ...]] = {} + + class _PrefillModule: + l3_prefill_layer = object() + + class _PrefillFwdModule: + l3_prefill_fwd = object() + + class _DecodeModule: + l3_decode_layer = object() + + class _DecodeFwdModule: + l3_decode_fwd = object() + + class _FlashConfig: + hidden_size = 4096 + num_attention_heads = 64 + head_dim = 512 + qk_rope_head_dim = 64 + q_lora_rank = 1024 + o_lora_rank = 1024 + o_groups = 8 + mix_hc = 24 + hc_dim = 16384 + max_position_embeddings = 8192 + moe_intermediate_size = 2048 + n_routed_experts = 256 + num_experts_per_tok = 6 + index_n_heads = 64 + index_head_dim = 128 + + class _ConfigModule: + FLASH = _FlashConfig + + compiled_names: list[str] = [] + + def _fake_compile(self, name, jit_fn, dummy_args): + compiled_names.append(name) + compiled_args[name] = tuple(dummy_args) + return DeepSeekV4L3Callable(compiled=object(), name=name) + + monkeypatch.setattr( + npu_executor.DeepSeekV4PyptoExecutor, + "_load_kernel_modules", + lambda self, layout: { + "config": _ConfigModule, + "prefill_layer": _PrefillModule, + "prefill_fwd": _PrefillFwdModule, + "decode_layer": _DecodeModule, + "decode_fwd": _DecodeFwdModule, + "rope_tables": object(), + }, + ) + monkeypatch.setattr(npu_executor.DeepSeekV4PyptoExecutor, "_compile_l3_callable", _fake_compile) + monkeypatch.setattr( + npu_executor.DeepSeekV4PyptoExecutor, + "_build_rope_tables", + lambda self, rope_tables_module, config_module: (torch.empty(1), torch.empty(1)), + ) + executor = npu_executor.DeepSeekV4PyptoExecutor( + platform="a2a3sim", + device_ids=tuple(range(8)), + compile_kernels=True, + ) + + executor._compile_model(loaded.runtime_model) + + assert compiled_names == ["deepseek_v4_prefill", "deepseek_v4_decode"] + # The packed l3_prefill_fwd emits final-normalized x_out and carries a trailing + # num_tokens scalar. LM-head is computed on the host side. + assert len(compiled_args["deepseek_v4_prefill"]) == 84 + # The packed l3_decode_fwd emits final-normalized x_out and carries a trailing + # num_tokens scalar. LM-head is computed on the host side. + assert len(compiled_args["deepseek_v4_decode"]) == 80 + # Both packed kernels carry a trailing num_tokens scalar. + assert isinstance(compiled_args["deepseek_v4_prefill"][-1], ctypes.c_int32) + assert isinstance(compiled_args["deepseek_v4_decode"][-1], ctypes.c_int32) + assert compiled_args["deepseek_v4_prefill"][0].shape == (8, 128, 4, 4096) + assert compiled_args["deepseek_v4_decode"][0].shape == (8, 8, 4, 4096) + prefill_order = npu_executor._PREFILL_FWD_TENSOR_ORDER + # Packed prefill flattens the FWD work caches to 5-D (kv_cache/cmp_kv stack x43, + # idx_kv_cache stacks x21 across the CSA group) and stacks the compress-state + # kv/score caches across the CSA (x21) and HCA (x20) groups. The per-step + # metadata, RoPE tables and compress-state block tables are shared single + # per-rank copies (the kernel slices them per layer). The kernel emits + # final-normalized hidden rows. + prefill_args = compiled_args["deepseek_v4_prefill"] + assert prefill_args[prefill_order.index("kv_cache")].shape == (8, 43 * 1, 128, 1, 512) + assert prefill_args[prefill_order.index("cmp_kv")].shape == (8, 43 * 256, 128, 1, 512) + assert prefill_args[prefill_order.index("idx_kv_cache")].shape == (8, 21 * 512, 128, 1, 128) + assert prefill_args[prefill_order.index("hca_cmp_wkv")].shape == (8, 20 * 512, 4096) + assert prefill_args[prefill_order.index("csa_cmp_wkv")].shape == (8, 21 * 1024, 4096) + assert prefill_args[prefill_order.index("csa_inner_wkv")].shape == (8, 21 * 256, 4096) + assert prefill_args[prefill_order.index("hca_cmp_kv_state")].shape == (8, 20 * 2048, 8, 512) + assert prefill_args[prefill_order.index("csa_cmp_kv_state")].shape == (8, 21 * 4096, 4, 1024) + assert prefill_args[prefill_order.index("csa_inner_kv_state")].shape == (8, 21 * 4096, 4, 256) + assert prefill_args[prefill_order.index("hca_compress_state_block_table")].shape == (8, 2048) + assert prefill_args[prefill_order.index("csa_compress_state_block_table")].shape == (8, 4096) + assert prefill_args[prefill_order.index("csa_inner_compress_state_block_table")].shape == (8, 4096) + assert prefill_args[prefill_order.index("ori_block_table")].shape == (8, 1) + assert prefill_args[prefill_order.index("cmp_block_table")].shape == (8, 32) + assert prefill_args[prefill_order.index("idx_block_table")].shape == (8, 64) + assert prefill_args[prefill_order.index("ori_slot_mapping")].shape == (8, 128) + assert prefill_args[prefill_order.index("position_ids")].shape == (8, 128) + assert prefill_args[prefill_order.index("input_ids")].shape == (8, 128) + assert prefill_args[prefill_order.index("cmp_sparse_indices")].shape == (8, 128, 640) + assert prefill_args[prefill_order.index("cmp_sparse_lens")].shape == (8, 128) + assert prefill_args[prefill_order.index("freqs_cos")].shape == (8, 8192, 64) + # In-kernel final RMSNorm only; host-side LM-head consumes selected rows. + assert prefill_args[prefill_order.index("final_norm_w")].shape == (8, 4096) + assert prefill_args[prefill_order.index("x_out")].shape == (8, 128, 4096) + decode_order = npu_executor._DECODE_FWD_TENSOR_ORDER + # Compress-state work caches are stacked across the CSA (x21) and HCA (x20) layer + # groups, each layer holding decode_batch (8) x state_max_blocks rows. + assert compiled_args["deepseek_v4_decode"][decode_order.index("hca_compress_state")].shape == (8, 20 * 8 * 64, 8, 1024) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_compress_state")].shape == (8, 21 * 8 * 65, 4, 2048) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_inner_compress_state")].shape == ( + 8, + 21 * 8 * 65, + 4, + 512, + ) + assert compiled_args["deepseek_v4_decode"][decode_order.index("hca_cmp_wkv")].shape == (8, 20 * 512, 4096) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_cmp_wkv")].shape == (8, 21 * 1024, 4096) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_inner_wkv")].shape == (8, 21 * 256, 4096) + # Decode emits final-normalized hidden rows; host-side LM-head consumes those + # rows and the TP vocab shards from the packed checkpoint weights. + assert compiled_args["deepseek_v4_decode"][decode_order.index("final_norm_w")].shape == (8, 4096) + assert compiled_args["deepseek_v4_decode"][decode_order.index("x_out")].shape == (8, 8, 4096) + + +def test_deepseek_layer_plan_tracks_attention_and_router_metadata(): + plan = build_deepseek_v4_layer_plan( + compress_ratios=_deepseek_flash_compress_ratios(), + num_hidden_layers=43, + num_hash_layers=3, + ) + + assert [(layer.attention_kind, layer.include_tid2eid) for layer in plan[:5]] == [ + ("swa", True), + ("swa", True), + ("csa", True), + ("hca", False), + ("csa", False), + ] + + +def test_deepseek_kernel_contract_does_not_require_device_lm_head(tmp_path): + kernel_dir = _write_deepseek_kernel_dir(tmp_path, lm_head_tp_size=2) + executor = npu_executor.DeepSeekV4PyptoExecutor.__new__(npu_executor.DeepSeekV4PyptoExecutor) + executor._kernel_dir = kernel_dir + + executor._validate_kernel_contract(DeepSeekV4CacheLayout()) + + +def test_deepseek_kernel_contract_accepts_config_named_tp_size(tmp_path): + kernel_dir = _write_deepseek_kernel_dir(tmp_path, lm_head_tp_size=8, use_config_constant=True) + executor = npu_executor.DeepSeekV4PyptoExecutor.__new__(npu_executor.DeepSeekV4PyptoExecutor) + executor._kernel_dir = kernel_dir + + executor._validate_kernel_contract(DeepSeekV4CacheLayout()) + + +def test_deepseek_kernel_contract_rejects_config_dimension_mismatch(tmp_path): + kernel_dir = _write_deepseek_kernel_dir(tmp_path, lm_head_tp_size=8, block_size=64) + executor = npu_executor.DeepSeekV4PyptoExecutor.__new__(npu_executor.DeepSeekV4PyptoExecutor) + executor._kernel_dir = kernel_dir + + with pytest.raises(ValueError, match="BLOCK_SIZE=64 expected 128"): + executor._validate_kernel_contract(DeepSeekV4CacheLayout()) + + +def test_deepseek_kernel_contract_rejects_prefill_state_mismatch(tmp_path): + kernel_dir = _write_deepseek_kernel_dir( + tmp_path, + lm_head_tp_size=8, + hca_state_blocks=1024, + csa_state_blocks=2048, + csa_inner_state_blocks=2048, + ) + executor = npu_executor.DeepSeekV4PyptoExecutor.__new__(npu_executor.DeepSeekV4PyptoExecutor) + executor._kernel_dir = kernel_dir + + with pytest.raises( + ValueError, + match=( + r"prefill_attention_hca.py:HCA_STATE_BLOCK_NUM=1024 expected 2048" + r".*prefill_attention_csa.py:CSA_STATE_BLOCK_NUM=2048 expected 4096" + r".*prefill_attention_csa.py:INNER_STATE_BLOCK_NUM=2048 expected 4096" + ), + ): + executor._validate_kernel_contract(DeepSeekV4CacheLayout()) + + +def test_deepseek_hc_input_builder_shapes_prefill_and_decode(): + # Exercise a wide, batch-agnostic input layout independent of production B=8/S=1. + builder = DeepSeekV4InputBuilder( + layout=DeepSeekV4CacheLayout(decode_batch=32, decode_seq=2, decode_tokens=64), hidden_size=4 + ) + + prefill = builder.prefill_x_hc(torch.arange(12, dtype=torch.bfloat16).reshape(3, 4), actual_tokens=3) + decode = builder.decode_x_hc(torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), actual_batch=2) + + assert prefill.shape == (8, 128, 4, 4) + assert prefill[0, 0, 0].tolist() == [0, 1, 2, 3] + assert prefill[7, 2, 3].tolist() == [8, 9, 10, 11] + assert torch.count_nonzero(prefill[:, 3:]) == 0 + assert decode.shape == (8, 64, 4, 4) + assert decode[0, 0, 0].tolist() == [0, 1, 2, 3] + assert decode[0, 1, 3].tolist() == [0, 1, 2, 3] + assert decode[7, 2, 0].tolist() == [4, 5, 6, 7] + assert decode[7, 3, 3].tolist() == [4, 5, 6, 7] + assert decode[0, 4, 0].tolist() == [0, 1, 2, 3] + assert decode[7, 5, 3].tolist() == [0, 1, 2, 3] + assert torch.equal(decode[:, 4:], decode[:, 0:2].repeat(1, 30, 1, 1)) + + +def test_deepseek_layout_rejects_context_beyond_decode_state_capacity(): + model = _runtime_model_for_embeddings() + + with pytest.raises(ValueError, match="max_seq_len=260"): + DeepSeekV4CacheLayout().validate_runtime( + model.config, + RuntimeConfig(page_size=128, max_batch_size=1, max_seq_len=261, weight_dtype="int8"), + tuple(range(8)), + ) + + +def test_deepseek_layer_plan_tracks_attention_and_gate_modes(): + plan = build_deepseek_v4_layer_plan( + compress_ratios=[0, 0, 4, 128, 4], + num_hidden_layers=5, + num_hash_layers=3, + ) + + assert [layer.attention_kind for layer in plan] == ["swa", "swa", "csa", "hca", "csa"] + assert [layer.include_tid2eid for layer in plan] == [True, True, True, False, False] + assert [layer.include_gate_bias for layer in plan] == [False, False, False, True, True] + + +def test_deepseek_weight_store_groups_requested_reads_by_shard(tmp_path): + weight_map = { + "a": "one.safetensors", + "b": "one.safetensors", + "c": "two.safetensors", + } + for filename in set(weight_map.values()): + (tmp_path / filename).touch() + opened: list[tuple[str, str]] = [] + reads: list[tuple[str, str]] = [] + tensors = { + "a": torch.tensor([1]), + "c": torch.tensor([3]), + } + + class _Reader: + def __init__(self, filename: str) -> None: + self.filename = filename + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, name: str) -> torch.Tensor: + reads.append((self.filename, name)) + return tensors[name] + + def _open(path: Path, device: str): + opened.append((path.name, device)) + return _Reader(path.name) + + store = DeepSeekV4WeightStore(model_dir=tmp_path, weight_map=weight_map, safe_open_fn=_open) + + loaded = store.load_many(["c", "a"]) + + assert list(loaded) == ["c", "a"] + assert loaded["c"].item() == 3 + assert loaded["a"].item() == 1 + assert opened == [("two.safetensors", "cpu"), ("one.safetensors", "cpu")] + assert reads == [("two.safetensors", "c"), ("one.safetensors", "a")] + + +def test_deepseek_weight_store_reads_real_safetensors_by_name(tmp_path): + from safetensors.torch import save_file + + save_file( + { + "embed.weight": torch.arange(4, dtype=torch.float32).reshape(2, 2), + "head.weight": torch.ones(2, 2), + }, + str(tmp_path / "global.safetensors"), + ) + store = DeepSeekV4WeightStore( + model_dir=tmp_path, + weight_map={ + "embed.weight": "global.safetensors", + "head.weight": "global.safetensors", + }, + ) + + loaded = store.load_tensor("embed.weight") + + assert loaded.tolist() == [[0.0, 1.0], [2.0, 3.0]] + + +def test_deepseek_executor_lazily_loads_and_caches_embeddings(tmp_path): + from safetensors.torch import save_file + + save_file( + {"embed.weight": torch.arange(24, dtype=torch.float32).reshape(6, 4)}, + str(tmp_path / "embed.safetensors"), + ) + open_count = 0 + store = DeepSeekV4WeightStore( + model_dir=tmp_path, + weight_map={"embed.weight": "embed.safetensors"}, + ) + original_open = store._safe_open_fn + + def _counting_open(path: Path, device: str): + nonlocal open_count + open_count += 1 + return original_open(path, device) + + store._safe_open_fn = _counting_open + executor = npu_executor.DeepSeekV4PyptoExecutor.__new__(npu_executor.DeepSeekV4PyptoExecutor) + executor._compiled = { + "dsv4": DeepSeekV4CompiledKernels( + layout=DeepSeekV4CacheLayout(), + model_dir=str(tmp_path), + weight_map=store.weight_map, + weight_store=store, + compress_ratios=tuple([0] * 44), + layer_plan=build_deepseek_v4_layer_plan( + compress_ratios=tuple([0] * 44), + num_hidden_layers=43, + num_hash_layers=3, + ), + kernel_dir=str(tmp_path), + ) + } + executor._embedding_cache = {} + model = _runtime_model_for_embeddings() + + first = executor.lookup_embeddings(model, torch.tensor([1, 3], dtype=torch.long)) + second = executor.lookup_embeddings(model, torch.tensor([[2, 4]], dtype=torch.long)) + + assert first.tolist() == [[4.0, 5.0, 6.0, 7.0], [12.0, 13.0, 14.0, 15.0]] + assert second.shape == (1, 2, 4) + assert second[0, 1].tolist() == [16.0, 17.0, 18.0, 19.0] + assert open_count == 1 + + +def test_deepseek_weight_store_loads_rank_local_experts(tmp_path): + core_names = deepseek_v4_layer_core_weight_names(0, include_tid2eid=True) + local_experts = deepseek_v4_local_expert_ids(rank=1, ranks=4, n_routed_experts=8) + expert_names = deepseek_v4_routed_expert_weight_names(0, local_experts) + weight_map = {name: "layer.safetensors" for name in (*core_names, *expert_names)} + (tmp_path / "layer.safetensors").touch() + reads: list[str] = [] + + class _Reader: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, name: str) -> torch.Tensor: + reads.append(name) + return torch.tensor([len(reads)]) + + store = DeepSeekV4WeightStore(model_dir=tmp_path, weight_map=weight_map, safe_open_fn=lambda path, device: _Reader()) + + loaded = store.load_rank_layer_weights( + 0, + rank=1, + ranks=4, + n_routed_experts=8, + include_tid2eid=True, + ) + + assert local_experts == (2, 3) + assert set(loaded) == set(weight_map) + assert all(".experts.2." in name or ".experts.3." in name for name in expert_names) + assert not any(".experts.0." in name or ".experts.1." in name for name in loaded) + + +def test_deepseek_weight_store_packs_lm_head_into_8_tp_shards(tmp_path): + from safetensors.torch import save_file + + save_file( + { + "embed.weight": torch.arange(64, dtype=torch.float32).reshape(16, 4), + "norm.weight": torch.arange(4, dtype=torch.float32), + "head.weight": torch.arange(64, dtype=torch.float32).reshape(16, 4) + 100, + "hc_head_fn": torch.zeros((4, 16), dtype=torch.float32), + "hc_head_scale": torch.ones((1,), dtype=torch.float32), + "hc_head_base": torch.zeros((4,), dtype=torch.float32), + }, + str(tmp_path / "global.safetensors"), + ) + store = DeepSeekV4WeightStore( + model_dir=tmp_path, + weight_map={ + "embed.weight": "global.safetensors", + "norm.weight": "global.safetensors", + "head.weight": "global.safetensors", + "hc_head_fn": "global.safetensors", + "hc_head_scale": "global.safetensors", + "hc_head_base": "global.safetensors", + }, + ) + + global_weights = store.load_packed_global_weights(ranks=8) + + assert global_weights.lm_head_layout.vocab_per_rank == 2 + assert global_weights.lm_head_layout.padded_vocab_per_rank == 512 + assert global_weights.lm_head_weight.shape == (8, 512, 4) + assert global_weights.lm_head_weight[0, :2].tolist() == [[100.0, 101.0, 102.0, 103.0], [104.0, 105.0, 106.0, 107.0]] + assert global_weights.lm_head_weight[1, :2].tolist() == [[108.0, 109.0, 110.0, 111.0], [112.0, 113.0, 114.0, 115.0]] + assert torch.count_nonzero(global_weights.lm_head_weight[:, 2:]) == 0 + + +def test_deepseek_lm_head_packer_rejects_uneven_vocab(): + with pytest.raises(ValueError, match="divide evenly"): + pack_deepseek_v4_lm_head_weight(torch.zeros((17, 4)), ranks=8) + + +def test_deepseek_layer_packer_transposes_and_stacks_rank_local_experts(): + raw = _synthetic_layer_raw(layer_id=0, n_experts=4) + + packed = pack_deepseek_v4_layer_weights( + 0, + raw, + ranks=2, + n_routed_experts=4, + compress_ratio=4, + include_tid2eid=False, + include_gate_bias=True, + ) + + assert packed.tensors["wq_a"].shape == (2, 4, 2) + assert packed.tensors["wq_a"][0].tolist() == raw["layers.0.attn.wq_a.weight"].t().tolist() + assert packed.tensors["wo_a"].shape == (2, 8, 2, 4) + assert packed.tensors["csa_cmp_wkv"].shape == (2, 2, 4) + assert packed.tensors["csa_cmp_wkv"][0].tolist() == raw["layers.0.attn.compressor.wkv.weight"].tolist() + assert packed.tensors["csa_inner_wkv"].shape == (2, 2, 4) + assert packed.tensors["csa_inner_wkv"][0].tolist() == raw["layers.0.attn.indexer.compressor.wkv.weight"].tolist() + assert packed.tensors["hca_cmp_wkv"].shape == (2, 512, 4096) + assert torch.count_nonzero(packed.tensors["hca_cmp_wkv"]) == 0 + assert packed.tensors["gate_bias"].shape == (2, 4) + assert packed.tensors["tid2eid"].shape == (2, 129280, 6) + assert packed.tensors["routed_w1"].shape == (2, 2, 2, 4) + assert packed.tensors["routed_w1"][0, 0].tolist() == raw["layers.0.ffn.experts.0.w1.weight"].tolist() + assert packed.tensors["routed_w1"][1, 0].tolist() == raw["layers.0.ffn.experts.2.w1.weight"].tolist() + assert torch.equal(packed.tensors["csa_hadamard_idx"][0], deepseek_v4_hadamard_idx()) + + +def test_deepseek_cache_slots_tables_and_mappings(): + manager = DeepSeekV4CacheManager(layout=DeepSeekV4CacheLayout()) + + assert manager.allocate("req-a") == 0 + assert manager.allocate("req-b") == 1 + assert manager.allocate("req-a") == 0 + + table = manager.block_table([1], max_blocks=64) + assert table.shape == (1, 64) + assert table[0, 0].item() == 64 + assert table[0, 63].item() == 127 + + cmp_mapping = manager.slot_mapping([1], [[0, 4, 256]], max_blocks=64, compress_ratio=4) + base = 1 * 64 * 128 + assert cmp_mapping.tolist() == [[base, base + 1, base + 64]] + + hca_state_mapping = manager.slot_mapping( + [1], + [[0, 128, 256]], + max_blocks=64, + block_size=8, + compress_ratio=128, + ) + assert hca_state_mapping.tolist() == [[1 * 64 * 8, 1 * 64 * 8 + 1, 1 * 64 * 8 + 2]] + + manager.release(["req-a"]) + assert manager.allocate("req-c") == 0 + + +def test_deepseek_prepare_prefill_inputs_maps_chunk_and_sparse_tables(): + runner, model = _runner_for_prepared_inputs() + layout = runner._compiled.layout + scratch_slot = layout.decode_batch - 1 + embeddings = torch.arange(12, dtype=torch.bfloat16).reshape(1, 3, 4) + + prepared = runner.prepare_prefill_inputs( + model, + PrefillBatch( + request_ids=["req-a"], + token_ids=torch.tensor([[10, 11, 12]], dtype=torch.long), + input_embeddings=embeddings, + seq_lens=torch.tensor([129], dtype=torch.int32), + positions=torch.tensor([[126, 127, 128]], dtype=torch.long), + ), + ) + + assert prepared.request_id == "req-a" + assert prepared.slot == 0 + assert prepared.actual_tokens == 3 + assert prepared.x_hc.shape == (8, 128, 4, 4) + assert prepared.cmp_block_table.shape == (8, 32) + assert prepared.idx_block_table.shape == (8, 64) + assert prepared.position_ids.shape == (8, 128) + assert prepared.position_ids[0, :4].tolist() == [126, 127, 128, 129] + assert prepared.input_ids[0, :4].tolist() == [10, 11, 12, 10] + assert prepared.ori_slot_mapping.shape == (8, 128) + assert prepared.ori_slot_mapping[0, :4].tolist() == [126, 127, 0, scratch_slot * 128 + 1] + assert prepared.hca_cmp_slot_mapping.shape == (8, 128) + assert prepared.hca_cmp_slot_mapping[0, :3].tolist() == [-1, 0, -1] + assert prepared.hca_cmp_slot_mapping[0, 3].item() == -1 + assert prepared.csa_cmp_slot_mapping.shape == (8, 128) + assert prepared.csa_cmp_slot_mapping[0, :3].tolist() == [-1, 31, -1] + assert prepared.csa_cmp_slot_mapping[0, 3].item() == -1 + assert prepared.csa_idx_slot_mapping.shape == (8, 128) + assert prepared.csa_idx_slot_mapping[0, :3].tolist() == [-1, 31, -1] + assert prepared.csa_idx_slot_mapping[0, 3].item() == -1 + assert prepared.hca_state_slot_mapping.shape == (8, 128) + assert prepared.hca_state_slot_mapping[0, :4].tolist() == [ + 126, + 127, + 128, + scratch_slot * layout.prefill_hca_state_max_blocks * layout.c128_state_block_size + 129, + ] + assert prepared.csa_state_slot_mapping.shape == (8, 128) + assert prepared.csa_state_slot_mapping[0, :4].tolist() == [ + 126, + 127, + 128, + scratch_slot * layout.prefill_csa_state_max_blocks * layout.c4_state_block_size + 129, + ] + assert prepared.csa_inner_state_slot_mapping.shape == (8, 128) + assert prepared.csa_inner_state_slot_mapping[0, :4].tolist() == [ + 126, + 127, + 128, + scratch_slot * layout.prefill_csa_inner_state_max_blocks * layout.c4_state_block_size + 129, + ] + sparse4, lens4 = prepared.sparse_inputs_for_ratio(4) + sparse0, lens0 = prepared.sparse_inputs_for_ratio(0) + assert sparse4.shape == (8, 128, 640) + assert lens4[0, 1].item() > lens0[0, 1].item() + assert sparse4[0, 1, 0].item() == 0 + assert sparse4[0, 1, 127].item() == 129 + assert sparse4[0, 1, 128].item() == 256 + + +def test_deepseek_prepare_decode_inputs_uses_scratch_slots_for_fixed_rows(): + runner, model = _runner_for_prepared_inputs() + + prepared = runner.prepare_decode_inputs( + model, + DecodeBatch( + request_ids=["req-a", "req-b"], + token_ids=torch.tensor([[5], [9]], dtype=torch.long), + hidden_states=torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + seq_lens=torch.tensor([128, 5], dtype=torch.int32), + ), + ) + + assert prepared.actual_batch == 2 + assert prepared.slots == (0, 1) + assert prepared.kernel_slots[:4] == (0, 1, 2, 3) + assert len(prepared.kernel_slots) == 32 + assert prepared.x_hc.shape == (8, 64, 4, 4) + assert prepared.x_hc[0, 0, 0].tolist() == [0, 1, 2, 3] + assert prepared.x_hc[0, 1, 0].tolist() == [0, 1, 2, 3] + assert prepared.x_hc[0, 2, 0].tolist() == [4, 5, 6, 7] + assert prepared.x_hc[0, 3, 0].tolist() == [4, 5, 6, 7] + assert prepared.x_hc[0, 4, 0].tolist() == [0, 1, 2, 3] + # No prev_token_ids supplied: inactive fixed rows mirror row 0 so the packed + # decode tile can execute all rows without arbitrary routing metadata. + assert prepared.input_ids[0, :6].tolist() == [5, 5, 9, 9, 5, 5] + # Positions are the two real trailing slots (seq_len-2, seq_len-1). + assert prepared.position_ids[0, :6].tolist() == [126, 127, 3, 4, 126, 127] + # kv_seq_lens = seq_len: last written position is seq_len-1 and seq_len already + # counts the prefill-generated last token, so the KV history is seq_len entries. + assert prepared.kv_seq_lens[0, :4].tolist() == [128, 5, 128, 128] + assert prepared.block_table.shape == (8, 32, 1) + assert prepared.cmp_block_table.shape == (8, 32, 32) + assert prepared.ori_slot_mapping[0, :6].tolist() == [126, 127, 131, 132, 382, 383] + assert prepared.hca_cmp_slot_mapping[0, :6].tolist() == [-1, 0, -1, -1, -1, 8192] + assert prepared.csa_cmp_slot_mapping[0, :6].tolist() == [-1, 31, 4096, -1, -1, 8223] + assert prepared.csa_idx_slot_mapping[0, :6].tolist() == [-1, 31, 8192, -1, -1, 16415] + assert prepared.csa_state_slot_mapping[0, :6].tolist() == [126, 127, 263, 264, 646, 647] + + +def test_deepseek_prepare_decode_inputs_feeds_two_real_tokens(): + runner, model = _runner_for_prepared_inputs() + + prepared = runner.prepare_decode_inputs( + model, + DecodeBatch( + request_ids=["req-a", "req-b"], + token_ids=torch.tensor([[5], [9]], dtype=torch.long), + hidden_states=torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + seq_lens=torch.tensor([128, 5], dtype=torch.int32), + prev_token_ids=torch.tensor([3, 7], dtype=torch.long), + prev_hidden_states=torch.arange(8, 16, dtype=torch.bfloat16).reshape(2, 4), + ), + ) + + # Active rows get [prev_token, last_token]; positions are (seq_len-2, seq_len-1). + assert prepared.input_ids[0, :6].tolist() == [3, 5, 7, 9, 3, 5] + assert prepared.position_ids[0, :6].tolist() == [126, 127, 3, 4, 126, 127] + assert prepared.kv_seq_lens[0, :4].tolist() == [128, 5, 128, 128] + # 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[0, 2, 0].tolist() == [12, 13, 14, 15] + assert prepared.x_hc[0, 3, 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] + + +def test_deepseek_decode_x_hc_prev_last_two_token_slots(): + builder = DeepSeekV4InputBuilder( + layout=DeepSeekV4CacheLayout(decode_batch=32, decode_seq=2, decode_tokens=64), hidden_size=4 + ) + + decode = builder.decode_x_hc( + torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + actual_batch=2, + prev_embeddings=torch.arange(8, 16, dtype=torch.bfloat16).reshape(2, 4), + ) + + assert decode.shape == (8, 64, 4, 4) + # Active row 0: slot 0 = prev, slot 1 = last. + assert decode[0, 0, 0].tolist() == [8, 9, 10, 11] + assert decode[0, 1, 3].tolist() == [0, 1, 2, 3] + # Active row 1: slot 0 = prev, slot 1 = last. + assert decode[7, 2, 0].tolist() == [12, 13, 14, 15] + assert decode[7, 3, 3].tolist() == [4, 5, 6, 7] + # Padding rows replicate active row 0's last embedding (both slots). + assert decode[0, 4, 0].tolist() == [0, 1, 2, 3] + assert decode[0, 5, 3].tolist() == [0, 1, 2, 3] + + +def test_deepseek_stage_decode_inputs_uses_shared_buffers(): + runner, model = _runner_for_prepared_inputs() + prepared = runner.prepare_decode_inputs( + model, + DecodeBatch( + request_ids=["req-a"], + token_ids=torch.tensor([[5]], dtype=torch.long), + hidden_states=torch.arange(4, dtype=torch.bfloat16).reshape(1, 4), + seq_lens=torch.tensor([128], dtype=torch.int32), + ), + ) + + staged = runner._stage_decode_inputs(prepared) + + assert staged.x_hc.is_shared() + 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", + "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", + ): + assert getattr(staged, name).is_shared() + + +def test_deepseek_run_decode_dispatches_active_token_count(): + from types import SimpleNamespace + + runner, model = _runner_for_prepared_inputs() + runner._compiled.decode = DeepSeekV4L3Callable(compiled=object(), name="decode") + captured: dict[str, object] = {} + + def fake_stage(inputs): + captured["prepared"] = inputs + return inputs + + def fake_decode_fwd_args(inputs, x_hc, x_out): + captured["x_hc_shape"] = tuple(x_hc.shape) + return (x_hc, x_out) + + def fake_run_l3(_callable, *args): + captured["num_tokens"] = args[-1] + args[-2].fill_(1) + + def fake_logits(_hidden, *, active_rows, label): + captured["active_rows"] = active_rows + captured["label"] = label + return torch.zeros((len(active_rows), model.config.vocab_size), dtype=torch.float32) + + hidden_out = torch.empty( + runner._compiled.layout.ranks, + runner._compiled.layout.decode_tokens, + model.config.hidden_size, + dtype=torch.bfloat16, + ) + runner._ensure_l3_shared_buffers = lambda _model: None + runner._stage_decode_inputs = fake_stage + runner._require_prefill_cache_snapshots = lambda: None + runner._seed_decode_work_cache = lambda _slots: None + runner._require_decode_buffers = lambda: SimpleNamespace(x_hc_a=captured["prepared"].x_hc) + runner._require_decode_output_buffer = lambda _hidden_size: hidden_out + runner._decode_fwd_args = fake_decode_fwd_args + runner._run_l3 = fake_run_l3 + runner._logits_for_hidden = fake_logits + + result = runner.run_decode( + model, + DecodeBatch( + request_ids=["req-a"], + token_ids=torch.tensor([[5]], dtype=torch.long), + hidden_states=torch.arange(4, dtype=torch.bfloat16).reshape(1, 4), + seq_lens=torch.tensor([128], dtype=torch.int32), + ), + ) + + assert captured["num_tokens"] == runner._compiled.layout.decode_seq + assert captured["active_rows"] == (runner._compiled.layout.decode_seq - 1,) + assert captured["label"] == "decode" + 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) + + +def test_deepseek_run_prefill_dispatches_static_prefill_token_count_temporarily(): + runner, model = _runner_for_prepared_inputs() + runner._compiled.prefill = DeepSeekV4L3Callable(compiled=object(), name="prefill") + captured: dict[str, object] = {} + + def fake_stage(inputs): + captured["prepared"] = inputs + return inputs + + def fake_prefill_fwd_args(x_out): + captured["x_out_shape"] = tuple(x_out.shape) + return (x_out,) + + def fake_run_l3(_callable, *args): + captured["num_tokens"] = args[-1] + args[-2].fill_(1) + + def fake_logits(_hidden, *, active_rows, label): + captured["active_rows"] = active_rows + captured["label"] = label + return torch.zeros((len(active_rows), model.config.vocab_size), dtype=torch.float32) + + runner._ensure_l3_shared_buffers = lambda _model: None + runner._stage_prefill_fwd_inputs = fake_stage + runner._prefill_fwd_args = fake_prefill_fwd_args + runner._run_l3 = fake_run_l3 + runner._snapshot_prefill_fwd_caches = lambda _slot: None + runner._logits_for_hidden = fake_logits + + result = runner.run_prefill( + model, + PrefillBatch( + request_ids=["req-a"], + token_ids=torch.tensor([[10, 11, 12]], dtype=torch.long), + input_embeddings=torch.arange(12, dtype=torch.bfloat16).reshape(1, 3, 4), + seq_lens=torch.tensor([3], dtype=torch.int32), + positions=torch.tensor([[0, 1, 2]], dtype=torch.long), + ), + ) + + assert captured["prepared"].actual_tokens == 3 + assert captured["num_tokens"] == runner._compiled.layout.prefill_seq + assert captured["active_rows"] == (2,) + assert captured["label"] == "prefill" + assert captured["x_out_shape"] == ( + runner._compiled.layout.ranks, + runner._compiled.layout.prefill_seq, + model.config.hidden_size, + ) + assert result.logits.shape == (1, model.config.vocab_size) + + +def test_deepseek_l3_dispatch_rejects_non_shared_tensor_before_worker_start(): + runner, _model = _runner_for_prepared_inputs() + + with pytest.raises(TypeError, match="before the L3 worker starts"): + runner._run_l3(DeepSeekV4L3Callable(compiled=object(), name="fake"), torch.zeros(1)) + + +def test_deepseek_l3_worker_requires_full_shared_preallocation_before_start(): + runner, _model = _runner_for_prepared_inputs() + + with pytest.raises(RuntimeError, match="shared host buffers are preallocated"): + runner._run_l3( + DeepSeekV4L3Callable(compiled=object(), name="fake"), + torch.zeros(1).share_memory_(), + ) + + +def test_deepseek_l3_scalars_are_runtime_python_ints(): + runner, _model = _runner_for_prepared_inputs() + + value = runner._int32_scalar(7) + + assert isinstance(value, int) + assert value == 7 + + +def test_deepseek_cache_replicates_decode_padding_rows(): + active = torch.tensor([[10, 11], [20, 21]], dtype=torch.int32) + + padded = DeepSeekV4CacheManager.replicate_first_row(active, actual_rows=2, kernel_rows=4) + + assert padded.tolist() == [[10, 11], [20, 21], [10, 11], [10, 11]] + + +def test_deepseek_decode_work_cache_loads_snapshot_into_kernel_slots(): + layout = DeepSeekV4CacheLayout( + ranks=1, + decode_batch=3, + decode_seq=2, + decode_tokens=6, + ori_max_blocks=1, + cmp_max_blocks=1, + ) + layer = DeepSeekV4LayerPlan( + layer_id=0, + compress_ratio=0, + attention_kind="swa", + include_tid2eid=True, + include_gate_bias=False, + ) + compiled = DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=(), + layer_plan=(layer,), + kernel_dir="", + ) + runner = DeepSeekV4ModelRunner(compiled=compiled) + + # ``_populate_decode_work_cache`` zeroes the stacked cache, then copies each + # layer's prefill snapshot into its kernel slots at stacked offset + # ``fwd_offset * decode_batch + slot``. Exercise that per-slot copy primitive + # directly for a single swa layer (fwd_offset 0) into kernel slots 0 and 2. + work_kv = torch.full((1, 3, 1, 1, 1), -1.0, dtype=torch.bfloat16) + work_cmp = torch.full((1, 3, 1, 1, 1), -2.0, dtype=torch.bfloat16) + snap_kv = torch.tensor([[[[[7.0]]]]], dtype=torch.bfloat16) + snap_cmp = torch.tensor([[[[[8.0]]]]], dtype=torch.bfloat16) + work_kv.zero_() + work_cmp.zero_() + for slot in (0, 2): + runner._copy_snapshot_blocks_to_work(snap_kv, work_kv, slot, layout.ori_max_blocks) + runner._copy_snapshot_blocks_to_work(snap_cmp, work_cmp, slot, layout.cmp_max_blocks) + + # The unused slot 1 is left at zero. + assert work_kv.flatten().tolist() == [7.0, 0.0, 7.0] + assert work_cmp.flatten().tolist() == [8.0, 0.0, 8.0] + + +def test_deepseek_prefill_snapshot_slices_physical_slot_pool(): + layout = DeepSeekV4CacheLayout( + ranks=1, + decode_batch=2, + decode_seq=1, + decode_tokens=2, + block_size=1, + ori_max_blocks=1, + prefill_cmp_max_blocks=2, + prefill_idx_max_blocks=3, + prefill_hca_state_max_blocks=1, + prefill_csa_state_max_blocks=1, + prefill_csa_inner_state_max_blocks=1, + ) + layer_plan = ( + DeepSeekV4LayerPlan( + layer_id=0, + compress_ratio=4, + attention_kind="csa", + include_tid2eid=True, + include_gate_bias=False, + ), + DeepSeekV4LayerPlan( + layer_id=1, + compress_ratio=128, + attention_kind="hca", + include_tid2eid=False, + include_gate_bias=True, + ), + ) + runner = DeepSeekV4ModelRunner( + compiled=DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=(4, 128), + layer_plan=layer_plan, + kernel_dir="", + ) + ) + runner._prefill_fwd_buffers = npu_runner._DeepSeekV4PrefillFwdSharedBuffers( + x_hc=torch.empty(0), + freqs_cos=torch.empty(0), + freqs_sin=torch.empty(0), + tensors={ + "kv_cache": torch.tensor([100.0, 200.0], dtype=torch.bfloat16).reshape(1, 2, 1, 1, 1), + "cmp_kv": torch.tensor( + [10.0, 11.0, 12.0, 13.0, 20.0, 21.0, 22.0, 23.0], + dtype=torch.bfloat16, + ).reshape(1, 8, 1, 1, 1), + "idx_kv_cache": torch.tensor([30.0, 31.0, 32.0, 33.0, 34.0, 35.0], dtype=torch.bfloat16).reshape( + 1, 6, 1, 1, 1 + ), + "csa_cmp_kv_state": torch.ones((1, 1, 1, 1, 1), dtype=torch.float32), + "csa_cmp_score_state": torch.ones((1, 1, 1, 1, 1), dtype=torch.float32) * 2, + "csa_inner_kv_state": torch.ones((1, 1, 1, 1, 1), dtype=torch.float32) * 3, + "csa_inner_score_state": torch.ones((1, 1, 1, 1, 1), dtype=torch.float32) * 4, + "hca_cmp_kv_state": torch.ones((1, 1, 1, 1, 1), dtype=torch.float32) * 5, + "hca_cmp_score_state": torch.ones((1, 1, 1, 1, 1), dtype=torch.float32) * 6, + }, + ) + + runner._snapshot_prefill_fwd_caches(slot=1) + + csa_snapshot = runner._prefill_cache_snapshots[0].tensors + hca_snapshot = runner._prefill_cache_snapshots[1].tensors + assert csa_snapshot["kv_cache"].flatten().tolist() == [100.0] + assert csa_snapshot["cmp_kv"].flatten().tolist() == [12.0, 13.0] + assert csa_snapshot["idx_kv_cache"].flatten().tolist() == [33.0, 34.0, 35.0] + assert hca_snapshot["kv_cache"].flatten().tolist() == [200.0] + assert hca_snapshot["cmp_kv"].flatten().tolist() == [22.0, 23.0] + + +def test_deepseek_decode_work_cache_preserves_decode_state_after_initial_seed(): + layout = DeepSeekV4CacheLayout( + ranks=1, + decode_batch=3, + decode_seq=1, + decode_tokens=3, + ori_max_blocks=1, + cmp_max_blocks=1, + idx_max_blocks=1, + hca_state_max_blocks=1, + csa_state_max_blocks=1, + csa_inner_state_max_blocks=1, + block_size=1, + c128_state_block_size=1, + c4_state_block_size=1, + ) + ratios = _deepseek_flash_compress_ratios() + layer_plan = tuple( + DeepSeekV4LayerPlan( + layer_id=layer_id, + compress_ratio=ratio, + attention_kind=deepseek_v4_attention_kind(ratio), + include_tid2eid=layer_id < 3, + include_gate_bias=layer_id >= 3, + ) + for layer_id, ratio in enumerate(ratios) + ) + runner = DeepSeekV4ModelRunner( + compiled=DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=tuple(ratios), + layer_plan=layer_plan, + kernel_dir="", + ) + ) + + hca_dim = 2 * npu_runner.DEEPSEEK_V4_HCA_MAIN_OUT_DIM + csa_dim = 2 * npu_runner.DEEPSEEK_V4_CSA_MAIN_OUT_DIM + csa_inner_dim = 2 * npu_runner.DEEPSEEK_V4_CSA_INNER_OUT_DIM + runner._decode_work_cache = DeepSeekV4LayerCache( + kv_cache=torch.zeros((1, 43 * layout.decode_batch, 1, 1, 1), dtype=torch.bfloat16), + cmp_kv=torch.zeros((1, 43 * layout.decode_batch, 1, 1, 1), dtype=torch.bfloat16), + idx_kv_cache=torch.zeros((1, 21 * layout.decode_batch, 1, 1, 1), dtype=torch.bfloat16), + hca_compress_state=torch.zeros((1, 20 * layout.decode_batch, 1, 1, hca_dim), dtype=torch.float32), + csa_compress_state=torch.zeros((1, 21 * layout.decode_batch, 1, 1, csa_dim), dtype=torch.float32), + csa_inner_compress_state=torch.zeros((1, 21 * layout.decode_batch, 1, 1, csa_inner_dim), dtype=torch.float32), + ) + + def snapshot_for( + layer_id: int, + ratio: int, + *, + value_offset: float = 0.0, + ) -> DeepSeekV4LayerCacheSnapshot: + value = float(layer_id + 1) + value_offset + tensors = { + "kv_cache": torch.full((1, 1, 1, 1, 1), value, dtype=torch.bfloat16), + "cmp_kv": torch.full((1, 1, 1, 1, 1), value + 0.25, dtype=torch.bfloat16), + } + if ratio == 4: + tensors.update( + { + "idx_kv_cache": torch.full((1, 1, 1, 1, 1), value + 0.5, dtype=torch.bfloat16), + "csa_cmp_kv_state": torch.full( + (1, 1, 1, 1, npu_runner.DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + value, + dtype=torch.float32, + ), + "csa_cmp_score_state": torch.full( + (1, 1, 1, 1, npu_runner.DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + value + 1.0, + dtype=torch.float32, + ), + "csa_inner_kv_state": torch.full( + (1, 1, 1, 1, npu_runner.DEEPSEEK_V4_CSA_INNER_OUT_DIM), + value, + dtype=torch.float32, + ), + "csa_inner_score_state": torch.full( + (1, 1, 1, 1, npu_runner.DEEPSEEK_V4_CSA_INNER_OUT_DIM), + value + 1.0, + dtype=torch.float32, + ), + } + ) + elif ratio == 128: + tensors.update( + { + "hca_cmp_kv_state": torch.full( + (1, 1, 1, 1, npu_runner.DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + value, + dtype=torch.float32, + ), + "hca_cmp_score_state": torch.full( + (1, 1, 1, 1, npu_runner.DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + value + 1.0, + dtype=torch.float32, + ), + } + ) + return DeepSeekV4LayerCacheSnapshot(tensors) + + runner._prefill_cache_snapshots = { + layer_id: snapshot_for(layer_id, ratio) + for layer_id, ratio in enumerate(ratios) + } + + runner._seed_decode_work_cache((0, 2)) + assert runner._decode_work_cache.kv_cache[0, 0, 0, 0, 0].item() == 1.0 + assert runner._decode_work_cache.kv_cache[0, 2, 0, 0, 0].item() == 1.0 + + runner._decode_work_cache.kv_cache[0, 0, 0, 0, 0] = 99.0 + runner._seed_decode_work_cache((0, 2)) + assert runner._decode_work_cache.kv_cache[0, 0, 0, 0, 0].item() == 99.0 + + runner._seed_decode_work_cache((1,)) + assert runner._decode_work_cache.kv_cache[0, 1, 0, 0, 0].item() == 1.0 + + runner._prefill_cache_snapshots = { + layer_id: snapshot_for(layer_id, ratio, value_offset=10.0) + for layer_id, ratio in enumerate(ratios) + } + runner._decode_work_cache.kv_cache[0, 2, 0, 0, 0] = 77.0 + runner._decode_cache_seeded_slots.clear() + + runner._seed_decode_work_cache((0, 1, 2)) + + assert runner._decode_work_cache.kv_cache[0, 0, 0, 0, 0].item() == 11.0 + assert runner._decode_work_cache.kv_cache[0, 1, 0, 0, 0].item() == 11.0 + assert runner._decode_work_cache.kv_cache[0, 2, 0, 0, 0].item() == 11.0 + + +def test_deepseek_release_invalidates_all_decode_kernel_slots(): + layout = DeepSeekV4CacheLayout(decode_batch=3) + runner = DeepSeekV4ModelRunner( + compiled=DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=(), + layer_plan=(), + kernel_dir="", + ) + ) + assert runner.cache_manager.allocate("request-a") == 0 + runner._decode_cache_seeded_slots.update({0, 1, 2}) + runner._prefill_cache_snapshots[0] = DeepSeekV4LayerCacheSnapshot({}) + + runner.release_finished_requests(["request-a"]) + + assert runner._decode_cache_seeded_slots == set() + assert runner._prefill_cache_snapshots == {} + + +def _write_deepseek_model_dir(tmp_path: Path, *, quant_method: str = "compressed-tensors") -> Path: + model_dir = tmp_path / "dsv4-flash-w8a8" + model_dir.mkdir() + compress_ratios = _deepseek_flash_compress_ratios() + config = { + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4", + "vocab_size": 129280, + "hidden_size": 4096, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "num_hidden_layers": 43, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "max_position_embeddings": 1048576, + "rms_norm_eps": 1e-6, + "rope_theta": 10000, + "bos_token_id": 0, + "eos_token_id": 1, + "torch_dtype": "bfloat16", + "compress_ratios": compress_ratios, + "quantization_config": { + "quant_method": quant_method, + "format": "int-quantized", + "quantization_status": "compressed", + }, + } + (model_dir / "config.json").write_text(json.dumps(config)) + weight_names = deepseek_v4_startup_weight_names( + 43, + n_routed_experts=256, + compress_ratios=compress_ratios, + num_hash_layers=3, + ) + index = {"weight_map": {name: "model-00001-of-00001.safetensors" for name in weight_names}} + (model_dir / "model.safetensors.index.json").write_text(json.dumps(index)) + return model_dir + + +def _deepseek_flash_compress_ratios() -> list[int]: + return [0, 0, *(4 if layer_id % 2 == 0 else 128 for layer_id in range(2, 43)), 0] + + +def _write_deepseek_kernel_dir( + tmp_path: Path, + *, + lm_head_tp_size: int, + use_config_constant: bool = False, + block_size: int = 128, + hca_state_blocks: int = 2048, + csa_state_blocks: int = 4096, + csa_inner_state_blocks: int = 4096, +) -> Path: + kernel_dir = tmp_path / f"deepseek-v4-kernels-tp{lm_head_tp_size}" + kernel_dir.mkdir() + (kernel_dir / "prefill_attention_hca.py").write_text( + "\n".join( + [ + f"HCA_STATE_BLOCK_NUM = {hca_state_blocks}", + "HCA_STATE_MAX_BLOCKS = HCA_STATE_BLOCK_NUM", + "", + ] + ) + ) + (kernel_dir / "prefill_attention_csa.py").write_text( + "\n".join( + [ + f"CSA_STATE_BLOCK_NUM = {csa_state_blocks}", + "CSA_STATE_MAX_BLOCKS = CSA_STATE_BLOCK_NUM", + f"INNER_STATE_BLOCK_NUM = {csa_inner_state_blocks}", + "INNER_STATE_MAX_BLOCKS = INNER_STATE_BLOCK_NUM", + "", + ] + ) + ) + (kernel_dir / "prefill_layer.py").write_text("") + (kernel_dir / "prefill_fwd.py").write_text("") + (kernel_dir / "decode_layer.py").write_text("") + (kernel_dir / "decode_fwd.py").write_text("") + (kernel_dir / "config.py").write_text( + "\n".join( + [ + f"BLOCK_SIZE = {block_size}", + "DECODE_BATCH = 8", + "DECODE_SEQ = 1", + "DECODE_TOKENS = DECODE_BATCH * DECODE_SEQ", + "PREFILL_BATCH = 1", + "PREFILL_SEQ = 128", + "KV_ORI_MAX_BLOCKS = 1", + "KV_CMP_MAX_BLOCKS = 32", + "IDX_CACHE_MAX_BLOCKS = 64", + "PREFILL_CMP_MAX_BLOCKS = KV_CMP_MAX_BLOCKS", + "PREFILL_IDX_MAX_BLOCKS = IDX_CACHE_MAX_BLOCKS", + "EP_WORLD_SIZE = 8", + f"LM_HEAD_TP_SIZE = {lm_head_tp_size}", + "", + ] + ) + ) + if use_config_constant: + (kernel_dir / "lm_head.py").write_text("TP_SIZE = LM_HEAD_TP_SIZE\n") + else: + (kernel_dir / "lm_head.py").write_text(f"TP_SIZE = {lm_head_tp_size}\n") + return kernel_dir + + +def _synthetic_layer_raw(*, layer_id: int, n_experts: int) -> dict[str, torch.Tensor]: + prefix = f"layers.{layer_id}" + raw = { + f"{prefix}.hc_attn_fn": torch.arange(4, dtype=torch.float32).reshape(1, 4), + f"{prefix}.hc_attn_scale": torch.arange(3, dtype=torch.float32), + f"{prefix}.hc_attn_base": torch.arange(1, dtype=torch.float32), + f"{prefix}.attn_norm.weight": torch.arange(4, dtype=torch.bfloat16), + f"{prefix}.attn.wq_a.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + f"{prefix}.attn.wq_b.weight": torch.arange(12, dtype=torch.int8).reshape(6, 2), + f"{prefix}.attn.wq_b.scale": torch.arange(6, dtype=torch.float32), + f"{prefix}.attn.wkv.weight": torch.arange(12, dtype=torch.bfloat16).reshape(3, 4), + f"{prefix}.attn.q_norm.weight": torch.arange(2, dtype=torch.bfloat16), + f"{prefix}.attn.kv_norm.weight": torch.arange(3, dtype=torch.bfloat16), + f"{prefix}.attn.attn_sink": torch.arange(2, dtype=torch.float32), + f"{prefix}.attn.wo_a.weight": torch.arange(64, dtype=torch.bfloat16).reshape(16, 4), + f"{prefix}.attn.wo_b.weight": torch.arange(64, dtype=torch.int8).reshape(4, 16), + f"{prefix}.attn.wo_b.scale": torch.arange(4, dtype=torch.float32), + f"{prefix}.hc_ffn_fn": torch.arange(4, dtype=torch.float32).reshape(1, 4), + f"{prefix}.hc_ffn_scale": torch.arange(3, dtype=torch.float32), + f"{prefix}.hc_ffn_base": torch.arange(1, dtype=torch.float32), + f"{prefix}.ffn_norm.weight": torch.arange(4, dtype=torch.bfloat16), + f"{prefix}.ffn.gate.weight": torch.arange(16, dtype=torch.bfloat16).reshape(4, 4), + f"{prefix}.ffn.gate.bias": torch.arange(4, dtype=torch.float32), + f"{prefix}.ffn.shared_experts.w1.weight": torch.arange(8, dtype=torch.int8).reshape(2, 4), + f"{prefix}.ffn.shared_experts.w1.scale": torch.arange(2, dtype=torch.float32), + f"{prefix}.ffn.shared_experts.w2.weight": torch.arange(8, dtype=torch.int8).reshape(4, 2), + f"{prefix}.ffn.shared_experts.w2.scale": torch.arange(4, dtype=torch.float32), + f"{prefix}.ffn.shared_experts.w3.weight": torch.arange(8, dtype=torch.int8).reshape(2, 4), + f"{prefix}.ffn.shared_experts.w3.scale": torch.arange(2, dtype=torch.float32), + f"{prefix}.attn.compressor.wkv.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + f"{prefix}.attn.compressor.wgate.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + f"{prefix}.attn.compressor.ape": torch.arange(8, dtype=torch.float32).reshape(4, 2), + f"{prefix}.attn.compressor.norm.weight": torch.arange(3, dtype=torch.bfloat16), + f"{prefix}.attn.indexer.wq_b.weight": torch.arange(12, dtype=torch.int8).reshape(6, 2), + f"{prefix}.attn.indexer.wq_b.scale": torch.arange(6, dtype=torch.float32), + f"{prefix}.attn.indexer.weights_proj.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + f"{prefix}.attn.indexer.compressor.wkv.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + f"{prefix}.attn.indexer.compressor.wgate.weight": torch.arange(8, dtype=torch.bfloat16).reshape(2, 4), + f"{prefix}.attn.indexer.compressor.ape": torch.arange(8, dtype=torch.float32).reshape(4, 2), + f"{prefix}.attn.indexer.compressor.norm.weight": torch.arange(2, dtype=torch.bfloat16), + } + for expert_id in range(n_experts): + base = expert_id * 10 + raw.update( + { + f"{prefix}.ffn.experts.{expert_id}.w1.weight": torch.full((2, 4), base, dtype=torch.int8), + f"{prefix}.ffn.experts.{expert_id}.w1.scale": torch.full((2,), base + 1, dtype=torch.float32), + f"{prefix}.ffn.experts.{expert_id}.w2.weight": torch.full((4, 2), base + 2, dtype=torch.int8), + f"{prefix}.ffn.experts.{expert_id}.w2.scale": torch.full((4,), base + 3, dtype=torch.float32), + f"{prefix}.ffn.experts.{expert_id}.w3.weight": torch.full((2, 4), base + 4, dtype=torch.int8), + f"{prefix}.ffn.experts.{expert_id}.w3.scale": torch.full((2,), base + 5, dtype=torch.float32), + } + ) + return raw + + +class _Tokenizer: + bos_token_id = 0 + eos_token_id = 1 + pad_token_id = None + + def encode(self, text: str) -> list[int]: + return [1] + + def decode(self, token_ids: list[int]) -> str: + return "" + + +def _runtime_model_for_embeddings(): + from python.core.types import ModelConfig, RuntimeModel + + config = ModelConfig( + model_id="dsv4", + architecture="DeepseekV4ForCausalLM", + vocab_size=6, + hidden_size=4, + intermediate_size=8, + num_hidden_layers=43, + num_attention_heads=64, + num_key_value_heads=1, + head_dim=512, + max_position_embeddings=8192, + rms_norm_eps=1e-6, + rope_theta=10000.0, + bos_token_id=0, + eos_token_id=1, + pad_token_id=1, + torch_dtype="bfloat16", + ) + runtime = RuntimeConfig(page_size=128, max_batch_size=1, max_seq_len=260, weight_dtype="int8") + placeholder = torch.empty(0, config.hidden_size) + return RuntimeModel( + config=config, + runtime=runtime, + embed_tokens=placeholder, + final_norm_weight=torch.empty(0), + lm_head=placeholder, + layers=[], + ) + + +def _runner_for_prepared_inputs() -> tuple[DeepSeekV4ModelRunner, object]: + model = _runtime_model_for_embeddings() + compiled = DeepSeekV4CompiledKernels( + # Exercise a wide, batch-agnostic input layout independent of production B=8/S=1. + layout=DeepSeekV4CacheLayout(decode_batch=32, decode_seq=2, decode_tokens=64), + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=tuple([0] * 44), + layer_plan=build_deepseek_v4_layer_plan( + compress_ratios=tuple([0] * 44), + num_hidden_layers=43, + num_hash_layers=3, + ), + kernel_dir="", + ) + runner = DeepSeekV4ModelRunner(compiled=compiled) + runner.init_kv_cache("dsv4", model.config, model.runtime) + return runner, model + + +def test_deepseek_init_kv_cache_returns_scheduler_block_capacity(): + model = _runtime_model_for_embeddings() + compiled = DeepSeekV4CompiledKernels( + layout=DeepSeekV4CacheLayout(), + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=(), + layer_plan=(), + kernel_dir="", + ) + runner = DeepSeekV4ModelRunner(compiled=compiled) + + assert runner.init_kv_cache("dsv4", model.config, model.runtime) == 3 + + runtime = RuntimeConfig( + page_size=model.runtime.page_size, + max_batch_size=model.runtime.max_batch_size, + max_seq_len=model.runtime.max_seq_len, + total_kv_pages=17, + weight_dtype=model.runtime.weight_dtype, + ) + + assert runner.init_kv_cache("dsv4", model.config, runtime) == 17 + + +def test_deepseek_lm_head_computes_selected_rows_on_host_without_padded_vocab(): + layout = DeepSeekV4CacheLayout(ranks=2, decode_batch=2, decode_seq=2, decode_tokens=4) + compiled = DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=(), + layer_plan=(), + kernel_dir="", + ) + runner = DeepSeekV4ModelRunner(compiled=compiled) + lm_head_weight = torch.zeros((layout.ranks, 4, 3), dtype=torch.bfloat16) + lm_head_weight[0, 0] = torch.tensor([1.0, 0.0, 0.0]) + lm_head_weight[0, 1] = torch.tensor([0.0, 1.0, 0.0]) + lm_head_weight[0, 2] = torch.tensor([0.0, 0.0, 1.0]) + lm_head_weight[1, 0] = torch.tensor([1.0, 1.0, 0.0]) + lm_head_weight[1, 1] = torch.tensor([0.0, 1.0, 1.0]) + runner._global_weights = weight_loader.DeepSeekV4GlobalWeights( + embed_weight=torch.empty(0), + final_norm_weight=torch.empty(0), + lm_head_weight=lm_head_weight, + lm_head_layout=weight_loader.DeepSeekV4LmHeadLayout( + ranks=layout.ranks, + vocab_size=5, + hidden_size=3, + vocab_per_rank=3, + padded_vocab_per_rank=4, + ), + hc_head_fn=torch.empty(0), + hc_head_scale=torch.empty(0), + hc_head_base=torch.empty(0), + ) + hidden = torch.arange(layout.ranks * 6 * 3, dtype=torch.float32).reshape(layout.ranks, 6, 3).to(torch.bfloat16) + + def fail_run_l3(*args): + raise AssertionError("host LM-head must not dispatch an L3 program") + + runner._run_l3 = fail_run_l3 + logits = runner._logits_for_hidden(hidden, active_rows=(5, 2)) + + assert logits.shape == (2, 5) + assert logits[0].tolist() == [15, 16, 17, 31, 33] + assert logits[1].tolist() == [6, 7, 8, 13, 15] + + +def test_deepseek_final_hidden_normalizes_before_hc_head_projection_overflows(): + compiled = DeepSeekV4CompiledKernels( + layout=DeepSeekV4CacheLayout(), + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=(), + layer_plan=(), + kernel_dir="", + ) + runner = DeepSeekV4ModelRunner(compiled=compiled) + hidden_size = 3 + runner._global_weights = weight_loader.DeepSeekV4GlobalWeights( + embed_weight=torch.empty(0), + final_norm_weight=torch.ones(hidden_size), + lm_head_weight=torch.empty(0), + lm_head_layout=weight_loader.DeepSeekV4LmHeadLayout( + ranks=1, + vocab_size=1, + hidden_size=hidden_size, + vocab_per_rank=1, + padded_vocab_per_rank=1, + ), + hc_head_fn=torch.ones((4, hidden_size * 4), dtype=torch.float32), + hc_head_scale=torch.ones((1,), dtype=torch.float32), + hc_head_base=torch.zeros((4,), dtype=torch.float32), + ) + x_hc = torch.full( + (1, 2, 4, hidden_size), + torch.finfo(torch.bfloat16).max, + dtype=torch.bfloat16, + ) + + flat = x_hc.flatten(2).float() + inv_rms = torch.rsqrt(flat.square().mean(dim=-1, keepdim=True) + 1e-6) + unstable_mixes = torch.matmul(flat, runner._global_weights.hc_head_fn.t()) * inv_rms + assert not torch.isfinite(unstable_mixes).all() + + hidden = runner._final_hidden(x_hc) + + assert hidden.shape == (1, 2, hidden_size) + assert torch.isfinite(hidden.float()).all() diff --git a/tests/test_parallel.py b/tests/test_parallel.py index 6c59b46..3a48bbe 100644 --- a/tests/test_parallel.py +++ b/tests/test_parallel.py @@ -29,7 +29,6 @@ def _parse_cli_args(argv: list[str]): return cli.build_parser().parse_args(argv) - def test_parallel_config_groups_dp_replicas_into_tp_groups(): config = ParallelConfig( data_parallel_size=2,