diff --git a/.gitignore b/.gitignore index 3b512e5..41df846 100644 --- a/.gitignore +++ b/.gitignore @@ -27,3 +27,6 @@ profile_out/ # OS files .DS_Store Thumbs.db + +# Personal Claude Code guidance +CLAUDE.local.md diff --git a/README.md b/README.md index 998e613..b02ea38 100644 --- a/README.md +++ b/README.md @@ -80,37 +80,42 @@ task-submit --device auto --max-time 0 --run \ ## HTTP Serving (OpenAI-compatible API) Start the serving server with a multiprocess worker. When launching through -`task-submit`, keep `--device {}` so the selected NPU ID is substituted into the -server command: +`task-submit`, use single quotes around the `--run` payload so `$TASK_DEVICE` +expands inside the task: ```bash -task-submit --device auto --run \ - "python -m python.cli.main \ - --model /path/to/Qwen3-14B \ +task-submit --device auto --max-time 1200 --run \ + 'export PTO2_RING_HEAP=4294967296 PTO2_RING_TASK_WINDOW=131072 PTO2_RING_DEP_POOL=131072; \ + python python/cli/main.py \ + --model /data/linyifan/models/Qwen3-14B \ --backend npu \ --platform a2a3 \ - --device {} \ - --port 8899" + --device "$TASK_DEVICE" \ + --dp 1 \ + --tp 1 \ + --max-model-len 512 \ + --max-new-tokens 16 \ + --port 19340' ``` Send a generation request after the server logs `Application startup complete`: ```bash # Health check -curl --noproxy "*" http://127.0.0.1:8899/health +curl --noproxy "*" http://127.0.0.1:19340/health # Completion -curl --noproxy "*" http://127.0.0.1:8899/v1/completions \ +curl --noproxy "*" http://127.0.0.1:19340/v1/completions \ -H "Content-Type: application/json" \ -d '{"prompt": "Huawei is", "max_tokens": 32, "temperature": 0.0}' # Streaming -curl --noproxy "*" http://127.0.0.1:8899/v1/completions \ +curl --noproxy "*" http://127.0.0.1:19340/v1/completions \ -H "Content-Type: application/json" \ -d '{"prompt": "Huawei is", "max_tokens": 32, "stream": true}' # Chat completion -curl --noproxy "*" http://127.0.0.1:8899/v1/chat/completions \ +curl --noproxy "*" http://127.0.0.1:19340/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{"messages": [{"role": "user", "content": "What is 1+1?"}], "max_tokens": 32}' ``` @@ -118,31 +123,83 @@ curl --noproxy "*" http://127.0.0.1:8899/v1/chat/completions \ Run the serving benchmark: ```bash -python tests/bench_serving.py --port 8899 --stream -n 8 -c 4 --max-tokens 16 +python tests/bench_serving.py --port 19340 --stream -n 8 -c 4 --max-tokens 16 +``` + +### Parallel Strategy V1 + +Serving supports a v1 `DP x TP` device topology. Data parallelism creates one +independent serving engine per replica, and tensor parallelism passes one device +group to the PyPTO L3 distributed worker for that replica. Single-device serving +remains the default. Pipeline parallelism and expert parallelism are accepted in +the config surface but rejected until the model kernels support them. + +For example, run two data-parallel replicas on the two devices selected by +`task-submit`: + +```bash +task-submit --device auto --device-num 2 --max-time 1800 --run \ + 'export PTO2_RING_HEAP=4294967296 PTO2_RING_TASK_WINDOW=131072 PTO2_RING_DEP_POOL=131072; \ + python python/cli/main.py \ + --model /data/linyifan/models/Qwen3-14B \ + --backend npu \ + --platform a2a3 \ + --devices "$TASK_DEVICE" \ + --dp 2 \ + --tp 1 \ + --max-model-len 512 \ + --max-new-tokens 16 \ + --port 19339' +``` + +Send a completion request to the DP=2 server: + +```bash +curl --noproxy "*" http://127.0.0.1:19339/v1/completions \ + -H "Content-Type: application/json" \ + -d '{"prompt":"Huawei is","max_tokens":16,"temperature":0.0}' ``` +Offline `npu_generate.py` supports `--devices` and `--tp` for one logical TP +replica. It intentionally rejects `--dp > 1`; launch separate offline jobs if +data-parallel offline generation is needed. + Single-request HTTP serving does not require the larger PTO2 ring settings. For -concurrent NPU serving, start the server with the larger PTO2 ring settings: +concurrent NPU serving, use topology-specific ring settings. + +Single-replica concurrent serving uses the larger task window and dependency +pool: ```bash task-submit --device auto --run \ - "PTO2_RING_HEAP=4294967296 PTO2_RING_TASK_WINDOW=1048576 PTO2_RING_DEP_POOL=1048576 \ - python -m python.cli.main \ + 'export PTO2_RING_HEAP=4294967296 PTO2_RING_TASK_WINDOW=1048576 PTO2_RING_DEP_POOL=1048576; \ + python python/cli/main.py \ --model /path/to/Qwen3-14B \ --backend npu \ --platform a2a3 \ - --device {} \ - --port 8899" + --device "$TASK_DEVICE" \ + --port 8899' +``` + +DP=2+ concurrent serving should keep the smaller task window and dependency +pool used by the DP=2 command above: + +```bash +PTO2_RING_HEAP=4294967296 PTO2_RING_TASK_WINDOW=131072 PTO2_RING_DEP_POOL=131072 ``` Without these settings, multi-request serving may return HTTP 200 while generating no tokens and logging worker runtime failures such as `rtMalloc failed: 207001`, `507018`, or `507046`. +For DP=2+, setting `PTO2_RING_TASK_WINDOW` and `PTO2_RING_DEP_POOL` to `1048576` +with a 4 GiB heap can reserve about 19 GiB of runtime arena per replica and fail +with `rtMalloc failed: 207001`. + ## Notes - All model/device/runtime options are passed via CLI arguments. Run - `python -m python.cli.main --help` for the full list. + `python python/cli/main.py --help` for the full list. - Generated kernel artifacts are written under `build_output/` and are ignored by git. - This repository expects PyPTO, CANN, torch, safetensors, transformers, and the 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..3a3c79e --- /dev/null +++ b/examples/model/deepseek_v4/runner/npu_executor.py @@ -0,0 +1,942 @@ +# 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 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 deepseek_v4_lm_head_layout +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 + + +_FALLBACK_PYPTO_LIB_ROOT = Path("/data/liuxu/pypto-lib") +_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"}) +_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", + "lm_head", + "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: + 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 + + candidate = _FALLBACK_PYPTO_LIB_ROOT / "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) + 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 + + +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: + super().__init__( + kv_cache_manager, + platform=platform, + device_id=device_id, + device_ids=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) + + 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 + lm_head = None + freqs_cos = freqs_sin = None + if self._compile_kernels: + modules = self._load_kernel_modules() + 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"]), + ) + lm_head = self._compile_l3_callable( + "deepseek_v4_lm_head", + modules["lm_head"].l3_lm_head, + self._lm_head_dummy_args(model, layout), + ) + 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, + lm_head=lm_head, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + platform=self._platform, + device_id=self._device_id, + n_routed_experts=n_routed_experts, + num_hash_layers=num_hash_layers, + ) + + def _load_kernel_modules(self) -> dict[str, object]: + """Import DeepSeekV4 pypto-lib modules with EP fixed to the serving world size.""" + pypto_root = self._kernel_dir.parents[2] + ranks = DeepSeekV4CacheLayout().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", "lm_head", "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, + 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. Prefill additionally packs the per-layer + forward metadata (slot mappings, block tables, sparse tables, position ids, + input ids) and the prefill compressor-state / work caches on the layer axis, + and -- unlike decode -- carries the RoPE tables per-layer (x43). The single + trailing ``num_tokens`` scalar replaces the per-layer ``layer_id`` scalars. + """ + 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 + + 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. Prefill + # packs the RoPE tables per-layer, so they are FWD-stacked rather than shared. + for name, base in single.items(): + if 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, hca * 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, csa * 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, csa * layout.prefill_csa_inner_state_max_blocks), + dtype=torch.int32, + ), + # FWD-stacked prefill work caches (x43); each carries an explicit + # per-layer block dim that the kernel reshapes internally. + "kv_cache": torch.empty( + (ranks, fwd, layout.ori_max_blocks, layout.block_size, 1, head_dim), + dtype=torch.bfloat16, + ), + "ori_block_table": torch.empty((ranks, fwd * layout.ori_max_blocks), dtype=torch.int32), + "ori_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "cmp_kv": torch.empty( + (ranks, fwd, layout.prefill_cmp_max_blocks, layout.block_size, 1, head_dim), + dtype=torch.bfloat16, + ), + "cmp_block_table": torch.empty((ranks, fwd * layout.prefill_cmp_max_blocks), dtype=torch.int32), + "cmp_sparse_indices": torch.empty( + (ranks, fwd * seq, layout.prefill_sparse_topk), + dtype=torch.int32, + ), + "cmp_sparse_lens": torch.empty((ranks, fwd * seq), dtype=torch.int32), + "idx_kv_cache": torch.empty( + (ranks, fwd, layout.prefill_cmp_max_blocks, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + dtype=torch.bfloat16, + ), + "idx_block_table": torch.empty((ranks, fwd * layout.prefill_idx_max_blocks), dtype=torch.int32), + "position_ids": torch.empty((ranks, fwd * seq), dtype=torch.int32), + "hca_cmp_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "hca_state_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "csa_cmp_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "csa_idx_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "csa_state_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "csa_inner_state_slot_mapping": torch.empty((ranks, fwd * seq), dtype=torch.long), + "input_ids": torch.empty((ranks, fwd * seq), dtype=torch.long), + "x_out": torch.empty((ranks, seq, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + } + ) + 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) + lm_layout = deepseek_v4_lm_head_layout( + vocab_size=model.config.vocab_size, + hidden_size=hidden, + ranks=ranks, + ) + + 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), + # In-kernel final RMSNorm + LM head: per-rank norm weight and TP + # vocab shards in, per-rank logits out (the kernel replaced x_out). + "final_norm_w": torch.empty((ranks, hidden), dtype=torch.bfloat16), + "lm_head_weight": torch.empty( + (ranks, lm_layout.padded_vocab_per_rank, hidden), + dtype=torch.bfloat16, + ), + "logits": torch.empty((ranks, tokens, model.config.vocab_size), dtype=torch.float32), + } + ) + # 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 _lm_head_dummy_args(self, model: RuntimeModel, layout: DeepSeekV4CacheLayout) -> tuple[torch.Tensor, ...]: + """Return explicit serving dummy args for ``l3_lm_head``.""" + lm_layout = deepseek_v4_lm_head_layout( + vocab_size=model.config.vocab_size, + hidden_size=model.config.hidden_size, + ranks=layout.ranks, + ) + return ( + torch.empty( + (layout.ranks, layout.decode_tokens, model.config.hidden_size), + dtype=torch.bfloat16, + ), + torch.empty( + (layout.ranks, lm_layout.padded_vocab_per_rank, model.config.hidden_size), + dtype=torch.bfloat16, + ), + torch.empty( + (layout.ranks, layout.decode_tokens, model.config.vocab_size), + dtype=torch.float32, + ), + ) + + 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, hidden, DEEPSEEK_V4_HCA_MAIN_OUT_DIM), dtype=torch.bfloat16), + "hca_cmp_wgate": torch.empty((ranks, hidden, DEEPSEEK_V4_HCA_MAIN_OUT_DIM), 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, hidden, DEEPSEEK_V4_CSA_MAIN_OUT_DIM), dtype=torch.bfloat16), + "csa_cmp_wgate": torch.empty((ranks, hidden, DEEPSEEK_V4_CSA_MAIN_OUT_DIM), 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, hidden, DEEPSEEK_V4_CSA_INNER_OUT_DIM), dtype=torch.bfloat16), + "csa_inner_wgate": torch.empty((ranks, hidden, DEEPSEEK_V4_CSA_INNER_OUT_DIM), 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_layer.py", + "prefill_fwd.py", + "decode_layer.py", + "decode_fwd.py", + "lm_head.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, + "EP_WORLD_SIZE": layout.ranks, + "LM_HEAD_TP_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}") + if mismatched: + raise ValueError("DeepSeekV4 kernel config does not match serving layout: " + ", ".join(mismatched)) + + lm_head_path = self._kernel_dir / "lm_head.py" + tp_size = _int_constant_from_file(lm_head_path, "TP_SIZE") + if tp_size is not None and tp_size != layout.ranks: + raise ValueError( + f"DeepSeekV4 serving requires lm_head.py TP_SIZE={layout.ranks}, " + f"but {lm_head_path} has TP_SIZE={tp_size}. Update the pypto-lib " + "DeepSeekV4 LM head kernel before running 8-NPU generation serving." + ) 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..a6aecbc --- /dev/null +++ b/examples/model/deepseek_v4/runner/npu_runner.py @@ -0,0 +1,2734 @@ +# 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 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, + PrefillBatch, + PrefillResult, + RuntimeConfig, + RuntimeModel, +) + + +DEEPSEEK_V4_RANKS = 8 +DEEPSEEK_V4_HC_MULT = 4 +DEEPSEEK_V4_BLOCK_SIZE = 128 +DEEPSEEK_V4_DECODE_BATCH = 64 +DEEPSEEK_V4_DECODE_SEQ = 2 +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 = 8 +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 = 2 +DEEPSEEK_V4_PREFILL_IDX_MAX_BLOCKS = 2 +DEEPSEEK_V4_PREFILL_HCA_STATE_MAX_BLOCKS = 1024 +DEEPSEEK_V4_PREFILL_CSA_STATE_MAX_BLOCKS = 2048 +DEEPSEEK_V4_PREFILL_CSA_INNER_STATE_MAX_BLOCKS = 2048 +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 ``prefill_fwd_rank_core`` / ``l3_prefill_fwd`` +# host signature: every layer-stacked weight/state tensor in core-parameter order, +# then ``x_out`` and a trailing ``num_tokens`` scalar. 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", + "freqs_cos", + "freqs_sin", + "hca_cmp_wkv", + "hca_cmp_wgate", + "hca_cmp_ape", + "hca_cmp_norm_w", + "hca_cmp_kv_state", + "hca_cmp_score_state", + "hca_compress_state_block_table", + "csa_cmp_wkv", + "csa_cmp_wgate", + "csa_cmp_ape", + "csa_cmp_norm_w", + "csa_cmp_kv_state", + "csa_cmp_score_state", + "csa_compress_state_block_table", + "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", + "csa_inner_compress_state_block_table", + "kv_cache", + "ori_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", + "attn_sink", + "wo_a", + "wo_b", + "wo_b_scale", + "hc_ffn_fn", + "hc_ffn_scale", + "hc_ffn_base", + "norm_w", + "gate_w", + "gate_bias", + "tid2eid", + "input_ids", + "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", + "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 now performs the final RMSNorm and LM +# head in-kernel, so it takes ``final_norm_w`` + ``lm_head_weight`` inputs and a +# ``logits`` output (the trailing ``num_tokens`` scalar is appended at dispatch, +# like prefill, and is not part of this tensor tuple). +_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", + "lm_head_weight", + "logits", +) + +_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 + + 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) -> torch.Tensor: + """Build ``[ranks, 128, hc_mult, hidden]`` decode HC input. + + Current DeepSeekV4 decode kernels use a fixed 64-token contract. Serving + generation only samples from active request rows, but the fused decode + layer still expects every fixed row to carry valid token data. + """ + 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") + rows = torch.zeros( + (self.layout.decode_tokens, self.hidden_size), + dtype=embeddings.dtype, + device=embeddings.device, + ) + # 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 * self.layout.decode_seq + rows[start : start + self.layout.decode_seq].copy_( + embeddings[source_row : source_row + 1].expand(self.layout.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 + lm_head: 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) + if self.lm_head is not None: + callables.append(self.lm_head) + 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. + + Every tensor is layer-stacked for the single ``l3_prefill_fwd`` dispatch: FWD + metadata/work caches stack across all 43 hidden layers, CSA-group state across + the 21 compress_ratio==4 layers, HCA-group state across the 20 + compress_ratio==128 layers. ``tensors`` is keyed by ``_PREFILL_FWD_TENSOR_ORDER`` + name (excluding the stacked weights, which live in ``_stacked_weight_buffers``, + and ``freqs_*``/``x_hc``/``x_out`` which are tracked explicitly). + """ + + x_hc: torch.Tensor + x_out: 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._prefill_cache_snapshots: dict[int, DeepSeekV4LayerCacheSnapshot] = {} + self._global_weights: DeepSeekV4GlobalWeights | None = None + self._static_lm_head_weight: torch.Tensor | 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._lm_head_hidden_buffer: torch.Tensor | None = None + self._lm_head_logits_buffer: torch.Tensor | None = None + self._decode_logits_buffer: torch.Tensor | None = None + + def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> None: + """DeepSeekV4 owns several cache families, so generic KV allocation is bypassed.""" + self.input_builder = DeepSeekV4InputBuilder( + layout=self._compiled.layout, + hidden_size=config.hidden_size, + ) + return None + + 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() + + def load_packed_global_weights(self) -> DeepSeekV4GlobalWeights: + """Load global tensors and pack the LM head for the 8-way TP kernel.""" + if self._global_weights is None: + self._global_weights = self._compiled.weight_store.load_packed_global_weights( + ranks=self._compiled.layout.ranks + ) + self._ensure_shared_host_allocation_before_worker("lm_head_weight") + self._static_lm_head_weight = self._static_device_tensor(self._global_weights.lm_head_weight) + 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) + sparse_by_ratio = self._prefill_sparse_by_ratio(positions, actual_tokens) + + return DeepSeekV4PreparedPrefillInputs( + request_id=request_id, + slot=slot, + actual_tokens=actual_tokens, + x_hc=builder.prefill_x_hc(embeddings, actual_tokens=actual_tokens), + input_ids=self._rank_stack(self._padded_vector(token_ids, layout.prefill_seq, dtype=torch.long)), + position_ids=self._rank_stack(self._prefill_position_ids(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.cache_manager.sliding_window_slot_mapping( + [slot], + [positions], + kernel_rows=layout.prefill_batch, + )[0], + 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.cache_manager.compressed_slot_mapping( + [slot], + [positions], + max_blocks=layout.prefill_cmp_max_blocks, + compress_ratio=128, + kernel_rows=layout.prefill_batch, + )[0], + layout.prefill_seq, + ) + ), + hca_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self.cache_manager.state_slot_mapping( + [slot], + [positions], + max_blocks=layout.prefill_hca_state_max_blocks, + state_block_size=layout.c128_state_block_size, + kernel_rows=layout.prefill_batch, + )[0], + layout.prefill_seq, + ) + ), + csa_cmp_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self.cache_manager.compressed_slot_mapping( + [slot], + [positions], + max_blocks=layout.prefill_cmp_max_blocks, + compress_ratio=4, + kernel_rows=layout.prefill_batch, + )[0], + layout.prefill_seq, + ) + ), + csa_idx_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self.cache_manager.compressed_slot_mapping( + [slot], + [positions], + max_blocks=layout.prefill_idx_max_blocks, + compress_ratio=4, + kernel_rows=layout.prefill_batch, + )[0], + layout.prefill_seq, + ) + ), + csa_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self.cache_manager.state_slot_mapping( + [slot], + [positions], + max_blocks=layout.prefill_csa_state_max_blocks, + state_block_size=layout.c4_state_block_size, + kernel_rows=layout.prefill_batch, + )[0], + layout.prefill_seq, + ) + ), + csa_inner_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping( + self.cache_manager.state_slot_mapping( + [slot], + [positions], + max_blocks=layout.prefill_csa_inner_state_max_blocks, + state_block_size=layout.c4_state_block_size, + kernel_rows=layout.prefill_batch, + )[0], + 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}") + + token_ids = self._decode_token_rows( + batch.token_ids.detach().cpu().to(torch.long), + actual_batch, + vocab_size=model.config.vocab_size, + ) + decode_embeds = batch.hidden_states.to(torch.bfloat16).cpu() + 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) + 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.load_packed_global_weights() + self._static_freqs_cos_tensor() + self._static_freqs_sin_tensor() + self._ensure_lm_head_buffers() + self._ensure_decode_buffers(model.config.hidden_size) + self._ensure_decode_work_cache() + # The packed decode kernel emits logits directly and runs the final + # RMSNorm + LM head in-kernel; stage its output buffer and the + # final_norm_w / lm_head_weight inputs before the worker forks. + self._require_decode_logits_buffer(model.config.vocab_size) + self._static_final_norm_weight_tensor() + self._require_static_lm_head_weight() + # Pre-allocate the packed shared buffers (stacked layer weights, hc_head and + # packed prefill caches/metadata) before the prefill dispatch forks the L3 + # worker, so the child inherits them. + 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) + 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) + x_out = self._require_prefill_fwd_buffers().x_out + x_out.zero_() + args = self._prefill_fwd_args(x_out) + 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() + 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(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() + + logits = self._logits_for_hidden(x_out, active_rows=(inputs.actual_tokens - 1,), label="prefill") + 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.load_packed_global_weights() + self._static_freqs_cos_tensor() + self._static_freqs_sin_tensor() + self._ensure_lm_head_buffers() + self._require_decode_logits_buffer(model.config.vocab_size) + self._static_final_norm_weight_tensor() + self._require_static_lm_head_weight() + self._ensure_decode_work_cache() + 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() + if self._stacked_weight_buffers is None: + self._stage_stacked_weights(self.load_stacked_layer_weights()) + self._populate_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._debug_tensor_stats("decode.input.initial.active", x_hc[:, :active_decode_tokens, :, :]) + + logits_buffer = self._require_decode_logits_buffer(model.config.vocab_size) + logits_buffer.zero_() + # The packed decode kernel now runs the final RMSNorm + LM head in-kernel, + # emitting per-rank logits directly. ``num_tokens`` is the real active token + # count (actual_batch decode rows x decode_seq), mirroring prefill's + # actual_tokens scalar. + num_tokens = inputs.actual_batch * self._compiled.layout.decode_seq + args = self._decode_fwd_args(inputs, x_hc, logits_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 + self._debug_tensor_stats( + "decode.output.logits.active", + logits_buffer[:, :active_decode_tokens, :], + per_rank=True, + ) + if self._debug_tensor_stats_enabled() and not self._tensor_is_finite( + logits_buffer[:, :active_decode_tokens, :] + ): + raise RuntimeError("DeepSeekV4 packed decode produced non-finite active logits rows") + + active_rows = tuple(row * self._compiled.layout.decode_seq for row in range(inputs.actual_batch)) + logits = logits_buffer[0, list(active_rows), :].contiguous().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 _require_lm_head_callable(self) -> DeepSeekV4L3Callable: + if self._compiled.lm_head is None: + raise RuntimeError("DeepSeekV4 LM-head kernel is not compiled") + return self._compiled.lm_head + + def _prefill_fwd_args(self, x_out: torch.Tensor) -> tuple[Any, ...]: + """Build the single packed ``l3_prefill_fwd`` argument tuple.""" + buffers = self._require_prefill_fwd_buffers() + stacked = self._require_stacked_weights() + values = dict(stacked.tensors) + values.update( + { + "x_hc": buffers.x_hc, + "freqs_cos": buffers.freqs_cos, + "freqs_sin": buffers.freqs_sin, + "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, + logits: 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(), + "lm_head_weight": self._require_static_lm_head_weight(), + "logits": logits, + } + ) + 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", + "logits", + ) + 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, hca * 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, csa * 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, csa * layout.prefill_csa_inner_state_max_blocks), torch.int32, "prefill_fwd_csa_inner_state_block_table" + ), + # FWD-stacked work caches (x43, explicit per-layer block dim). + "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_max_blocks, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), + torch.bfloat16, + "prefill_fwd_cmp_kv", + ), + "idx_kv_cache": shared( + (ranks, fwd, layout.prefill_cmp_max_blocks, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + torch.bfloat16, + "prefill_fwd_idx_kv_cache", + ), + # FWD-stacked per-layer metadata (x43). + "ori_block_table": shared((ranks, fwd * layout.ori_max_blocks), torch.int32, "prefill_fwd_ori_block_table"), + "ori_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_ori_slot_mapping"), + "cmp_block_table": shared((ranks, fwd * layout.prefill_cmp_max_blocks), torch.int32, "prefill_fwd_cmp_block_table"), + "cmp_sparse_indices": shared((ranks, fwd * seq, layout.prefill_sparse_topk), torch.int32, "prefill_fwd_cmp_sparse_indices"), + "cmp_sparse_lens": shared((ranks, fwd * seq), torch.int32, "prefill_fwd_cmp_sparse_lens"), + "idx_block_table": shared((ranks, fwd * layout.prefill_idx_max_blocks), torch.int32, "prefill_fwd_idx_block_table"), + "position_ids": shared((ranks, fwd * seq), torch.int32, "prefill_fwd_position_ids"), + "hca_cmp_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_hca_cmp_slot_mapping"), + "hca_state_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_hca_state_slot_mapping"), + "csa_cmp_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_csa_cmp_slot_mapping"), + "csa_idx_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_csa_idx_slot_mapping"), + "csa_state_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_csa_state_slot_mapping"), + "csa_inner_state_slot_mapping": shared((ranks, fwd * seq), torch.long, "prefill_fwd_csa_inner_state_slot_mapping"), + "input_ids": shared((ranks, fwd * seq), torch.long, "prefill_fwd_input_ids"), + } + buffers = _DeepSeekV4PrefillFwdSharedBuffers( + x_hc=shared((ranks, seq, layout.hc_mult, hidden), torch.bfloat16, "prefill_fwd_x_hc"), + x_out=shared((ranks, seq, layout.hc_mult, hidden), torch.bfloat16, "prefill_fwd_x_out"), + freqs_cos=shared((ranks, fwd * max_seq_len, rope_dim), torch.bfloat16, "prefill_fwd_freqs_cos"), + freqs_sin=shared((ranks, fwd * 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 stacked metadata/state into the packed buffers. + + The per-request metadata is layer-independent (one slot, one token chunk), + so each FWD-stacked field tiles the single ``[ranks, dim]`` tensor x43 on a + fused layer axis; the sparse tables are selected per layer by compress + ratio. Compressor-state and work caches start zeroed and are produced by the + kernel. Compacted block tables tile across the CSA (x21) / HCA (x20) groups. + """ + buffers = self._require_prefill_fwd_buffers() + layout = self._compiled.layout + fwd = DEEPSEEK_V4_FWD_NUM_LAYERS + plan = self._compiled.layer_plan + + # 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._tile_layers(self._static_freqs_cos_table(), fwd), + name="prefill_fwd_freqs_cos", + ) + self._copy_shared( + buffers.freqs_sin, + self._tile_layers(self._static_freqs_sin_table(), fwd), + name="prefill_fwd_freqs_sin", + ) + + # Layer-independent FWD metadata tiled x43 along the fused layer axis. + fwd_tiled = { + "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, + } + for name, tensor in fwd_tiled.items(): + self._copy_shared(buffers.tensors[name], self._tile_layers(tensor, fwd), name=f"prefill_fwd_{name}") + + # Sparse tables vary by per-layer compression ratio. + sparse_indices_layers = [] + sparse_lens_layers = [] + for layer in plan: + indices, lens = inputs.sparse_inputs_for_ratio(layer.compress_ratio) + sparse_indices_layers.append(indices) + sparse_lens_layers.append(lens) + self._copy_shared( + buffers.tensors["cmp_sparse_indices"], + self._concat_layers(sparse_indices_layers, axis=1), + name="prefill_fwd_cmp_sparse_indices", + ) + self._copy_shared( + buffers.tensors["cmp_sparse_lens"], + self._concat_layers(sparse_lens_layers, axis=1), + name="prefill_fwd_cmp_sparse_lens", + ) + + # Compacted state block tables (CSA x21 / HCA x20). + self._copy_shared( + buffers.tensors["hca_compress_state_block_table"], + self._tile_layers(inputs.hca_compress_state_block_table, DEEPSEEK_V4_HCA_NUM_LAYERS), + name="prefill_fwd_hca_state_block_table", + ) + self._copy_shared( + buffers.tensors["csa_compress_state_block_table"], + self._tile_layers(inputs.csa_compress_state_block_table, DEEPSEEK_V4_CSA_NUM_LAYERS), + name="prefill_fwd_csa_state_block_table", + ) + self._copy_shared( + buffers.tensors["csa_inner_compress_state_block_table"], + self._tile_layers(inputs.csa_inner_compress_state_block_table, DEEPSEEK_V4_CSA_NUM_LAYERS), + name="prefill_fwd_csa_inner_state_block_table", + ) + + # 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_() + + @staticmethod + def _tile_layers(tensor: torch.Tensor, count: int) -> torch.Tensor: + """Tile a ``[ranks, d1, ...]`` tensor ``count`` times on a fused layer axis.""" + if tensor.ndim < 2: + raise ValueError(f"layer-tiled tensor must be rank>=2, got shape={tuple(tensor.shape)}") + expanded = tensor.unsqueeze(1).expand(tensor.shape[0], count, *tensor.shape[1:]) + return expanded.reshape(tensor.shape[0], count * tensor.shape[1], *tensor.shape[2:]).contiguous() + + @staticmethod + def _concat_layers(per_layer: Sequence[torch.Tensor], *, axis: int) -> torch.Tensor: + """Concatenate per-layer ``[ranks, d1, ...]`` tensors on the fused layer axis.""" + return torch.cat([tensor.contiguous() for tensor in per_layer], dim=axis).contiguous() + + 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) -> None: + """Capture per-layer cache slices from the packed prefill Out caches.""" + buffers = self._require_prefill_fwd_buffers() + csa_order = 0 + hca_order = 0 + for layer in self._compiled.layer_plan: + fwd = int(layer.layer_id) + snapshot: dict[str, torch.Tensor] = { + "kv_cache": buffers.tensors["kv_cache"][:, fwd].detach().cpu().contiguous().clone(), + "cmp_kv": buffers.tensors["cmp_kv"][:, fwd].detach().cpu().contiguous().clone(), + } + if layer.compress_ratio == 4: + snapshot["idx_kv_cache"] = buffers.tensors["idx_kv_cache"][:, fwd].detach().cpu().contiguous().clone() + 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() + + 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 _ensure_lm_head_buffers(self) -> None: + weights = self.load_packed_global_weights() + layout = self._compiled.layout + hidden_shape = (layout.ranks, layout.decode_tokens, weights.lm_head_layout.hidden_size) + logits_shape = (layout.ranks, layout.decode_tokens, weights.lm_head_layout.vocab_size) + if self._lm_head_hidden_buffer is None: + self._ensure_shared_host_allocation_before_worker("lm_head_hidden") + self._lm_head_hidden_buffer = self._shared_empty(hidden_shape, torch.bfloat16, name="lm_head_hidden") + if self._lm_head_logits_buffer is None: + self._ensure_shared_host_allocation_before_worker("lm_head_logits") + self._lm_head_logits_buffer = self._shared_empty(logits_shape, torch.float32, name="lm_head_logits") + + 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 the shared ``[ranks, decode_tokens, vocab]`` decode logits output. + + The packed ``l3_decode_fwd`` kernel emits logits directly, so decode no + longer reuses the separate-lm_head 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_static_lm_head_weight(self) -> torch.Tensor: + """Return the worker-resident per-rank LM-head weight shards. + + Reuses the exact tensor staged for the separate-lm_head step + (``[ranks, VOCAB_PER_TP, hidden]`` packed shards). + """ + if self._static_lm_head_weight is None: + global_weights = self.load_packed_global_weights() + self._static_lm_head_weight = self._static_device_tensor(global_weights.lm_head_weight) + return self._static_lm_head_weight + + 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 _populate_decode_work_cache(self, kernel_slots: Sequence[int]) -> None: + """Fill the layer-stacked decode work cache from every prefill snapshot. + + The packed ``l3_decode_fwd`` kernel reads all 43 layers' KV/compress state + in one call, so the full work cache is populated once before dispatch. 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``. + """ + cache = self._require_decode_work_cache() + layout = self._compiled.layout + batch = layout.decode_batch + cache.kv_cache.zero_() + cache.cmp_kv.zero_() + cache.idx_kv_cache.zero_() + cache.hca_compress_state.zero_() + cache.csa_compress_state.zero_() + cache.csa_inner_compress_state.zero_() + + 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 kernel_slots: + 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 kernel_slots: + 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 kernel_slots: + 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}" + ) + + @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() + self._ensure_lm_head_buffers() + if self._static_lm_head_weight is None: + self._static_lm_head_weight = self._static_device_tensor(global_weights.lm_head_weight) + if x_hc.ndim == 3: + # Decode output: the in-kernel hc_head already collapsed HC_MULT, so + # only the final RMS norm remains before the LM head. + hidden = self._final_norm(x_hc) + else: + hidden = self._final_hidden(x_hc) + if self._lm_head_hidden_buffer is None or self._lm_head_logits_buffer is None: + raise RuntimeError("DeepSeekV4 LM-head shared buffers were not initialized") + 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]}" + ) + if len(rows) > self._lm_head_hidden_buffer.shape[1]: + raise ValueError( + "DeepSeekV4 LM-head active row count exceeds kernel capacity: " + f"rows={len(rows)} capacity={self._lm_head_hidden_buffer.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, :]) + self._lm_head_hidden_buffer.zero_() + self._lm_head_hidden_buffer[:, : len(rows)].copy_(hidden[:, row_list, :]) + if self._debug_tensor_stats_enabled(): + self._debug_tensor_stats("lm_head.hidden_buffer.active", self._lm_head_hidden_buffer[:, : len(rows), :]) + self._lm_head_logits_buffer.zero_() + self._run_l3( + self._require_lm_head_callable(), + self._lm_head_hidden_buffer, + self._static_lm_head_weight, + self._lm_head_logits_buffer, + ) + active_logits = self._lm_head_logits_buffer[:, : len(rows), :] + self._debug_tensor_stats("lm_head.logits.active", active_logits, per_rank=True) + logits = self._lm_head_logits_buffer[0, : len(rows), :].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: + 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._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() + + @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: + out[values.numel() :] = out[0] + 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, ...], ...]: + positions = [] + for row in range(actual_batch): + seq_len = int(batch.seq_lens[row].item()) + if seq_len <= 0: + raise ValueError("decode seq_lens must be positive") + first_position = seq_len - 1 + positions.append(tuple(first_position + offset for offset in range(self._compiled.layout.decode_seq))) + return tuple(positions) + + def _decode_token_rows(self, token_ids: torch.Tensor, actual_batch: int, *, vocab_size: int) -> 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] + if vocab_size <= 0: + raise ValueError("vocab_size must be positive") + rows = torch.arange(layout.decode_tokens, dtype=torch.long).reshape( + layout.decode_batch, + layout.decode_seq, + ) + rows.remainder_(int(vocab_size)) + for row in range(actual_batch): + 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 + active = seq_lens[:actual_batch].detach().cpu().to(torch.int32) + (layout.decode_seq - 1) + 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..3d820a0 --- /dev/null +++ b/examples/model/deepseek_v4/runner/weight_loader.py @@ -0,0 +1,956 @@ +# 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 + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import ContextManager, Protocol + +import torch + + +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. + """ + 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] = [] + 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, + ) + ) + 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_transposed( + "attn.compressor.wkv.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, _DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + torch.bfloat16, + enabled=is_hca, + ), + "hca_cmp_wgate": replicated_transposed( + "attn.compressor.wgate.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, _DEEPSEEK_V4_HCA_MAIN_OUT_DIM), + 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_transposed( + "attn.compressor.wkv.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, _DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + torch.bfloat16, + enabled=is_csa, + ), + "csa_cmp_wgate": replicated_transposed( + "attn.compressor.wgate.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, _DEEPSEEK_V4_CSA_MAIN_OUT_DIM), + 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_transposed( + "attn.indexer.compressor.wkv.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, _DEEPSEEK_V4_CSA_INNER_OUT_DIM), + torch.bfloat16, + enabled=is_csa, + ), + "csa_inner_wgate": replicated_transposed( + "attn.indexer.compressor.wgate.weight", + (_DEEPSEEK_V4_HIDDEN_SIZE, _DEEPSEEK_V4_CSA_INNER_OUT_DIM), + 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/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index 526acc4..fcf194b 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -33,6 +33,7 @@ def _bootstrap_package_root() -> None: from python.core import GenerateConfig, LLMEngine, RuntimeConfig from python.core.kv_cache import KvCacheManager +from python.core.parallel import ParallelConfig, parse_device_ids from python.profile import get_profiler, merge_profile, profile_span from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor from python.core.types import LoadedModel @@ -346,6 +347,20 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--model-id", default="qwen3-14b-local") parser.add_argument("--platform", default="a2a3", choices=["a2a3sim", "a2a3", "a5sim", "a5"]) parser.add_argument("--device-id", type=int, default=0) + parser.add_argument( + "--devices", + default=None, + help="Comma-separated NPU device ids for one tensor-parallel L3 worker group.", + ) + parser.add_argument("--tp", "--tensor-parallel-size", dest="tensor_parallel_size", type=int, default=1) + parser.add_argument( + "--dp", + "--data-parallel-size", + dest="data_parallel_size", + type=int, + default=1, + help="Offline generation does not launch DP replicas; values > 1 fail fast.", + ) parser.add_argument("--max-seq-len", type=int, default=4096) parser.add_argument("--max-new-tokens", type=int, default=32) parser.add_argument("--temperature", type=float, default=0.0) @@ -384,12 +399,21 @@ def main() -> None: profile_enabled = args.profile or args.profile_verbose collector = _TimingCollector() if profile_enabled else None + parallel_config = ParallelConfig( + data_parallel_size=args.data_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, + devices=parse_device_ids(args.devices, default_device=args.device_id), + ) + if parallel_config.data_parallel_size != 1: + raise ValueError("offline npu_generate.py supports tensor parallelism only; data_parallel_size must be 1") + device_ids = parallel_config.replica_device_groups[0] kv_cache_manager = KvCacheManager() executor = PyptoExecutor( kv_cache_manager, platform=args.platform, - device_id=args.device_id, + device_id=device_ids[0], + device_ids=device_ids, save_kernels_dir=args.save_kernels_dir, l3_trace=args.profile_verbose, ) diff --git a/examples/model/qwen3_14b/npu_serving.json b/examples/model/qwen3_14b/npu_serving.json index 4659ec4..c4f021f 100644 --- a/examples/model/qwen3_14b/npu_serving.json +++ b/examples/model/qwen3_14b/npu_serving.json @@ -25,5 +25,15 @@ "device_id": 0, "save_kernels_dir": null, "l3": false + }, + "parallel": { + "devices": [0], + "data_parallel_size": 1, + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "enable_expert_parallel": false, + "expert_placement_strategy": "linear", + "all2all_backend": "none", + "data_parallel_routing": "least_pending_tokens" } } diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index c9b0ba5..e66c885 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -11,6 +11,7 @@ import importlib.util import sys +from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path from typing import Any @@ -105,6 +106,7 @@ def __init__( *, platform: str = "a2a3sim", device_id: int = 0, + device_ids: Sequence[int] | None = None, save_kernels_dir: str | None = None, l3_trace: bool = False, ) -> None: @@ -112,6 +114,7 @@ def __init__( kv_cache_manager, platform=platform, device_id=device_id, + device_ids=device_ids, save_kernels_dir=save_kernels_dir, ) self._l3_trace = l3_trace @@ -428,7 +431,7 @@ def _compile_jit_fwd_callable( config = self._run_config(codegen_only=True) distributed_config = DistributedConfig( - device_ids=[self._device_id], + device_ids=list(self._device_ids), num_sub_workers=0, block_dim=_QWEN14B_BLOCK_DIM, aicpu_thread_num=4, diff --git a/pypto-lib b/pypto-lib index 65871d8..ca927d2 160000 --- a/pypto-lib +++ b/pypto-lib @@ -1 +1 @@ -Subproject commit 65871d8d2e50e6c943e029272249aef933c2bbba +Subproject commit ca927d2c487dd7c7d5ffc6cc173482376371ce84 diff --git a/python/cli/main.py b/python/cli/main.py index 32342b9..c0ea9ba 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 @@ -23,6 +24,8 @@ from python.profile import get_profiler, merge_profile RuntimeConfig = None +ParallelConfig = None +parse_device_ids = None _VALID_BACKENDS = {"npu"} @@ -42,6 +45,29 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--backend", default="npu", choices=sorted(_VALID_BACKENDS), help="Inference backend (default: npu).") parser.add_argument("--platform", default="a2a3", help="NPU platform (default: a2a3).") parser.add_argument("--device", type=int, default=0, help="NPU device ID (default: 0).") + parser.add_argument( + "--devices", + default=None, + help="Comma-separated NPU device ids for DP x TP placement, for example 0,1,2,3.", + ) + parser.add_argument("--dp", "--data-parallel-size", dest="data_parallel_size", type=int, default=1, help="Data-parallel replica count.") + parser.add_argument("--tp", "--tensor-parallel-size", dest="tensor_parallel_size", type=int, default=1, help="Tensor-parallel group size.") + parser.add_argument("--pp", "--pipeline-parallel-size", dest="pipeline_parallel_size", type=int, default=1, help="Pipeline parallel size.") + parser.add_argument( + "--ep", + "--enable-expert-parallel", + dest="enable_expert_parallel", + action="store_true", + help="Enable expert parallelism. Not supported yet; accepted for config compatibility.", + ) + parser.add_argument("--expert-placement-strategy", default="linear", help="Expert placement policy placeholder.") + parser.add_argument("--all2all-backend", default="none", help="Expert-parallel all2all backend placeholder.") + parser.add_argument( + "--data-parallel-routing", + default="least_pending_tokens", + choices=["least_pending_tokens"], + help="Data-parallel request routing policy.", + ) # Dtype parser.add_argument("--dtype", default="float32", help="Weight data type (default: float32).") parser.add_argument("--kv-cache-dtype", default="bfloat16", help="KV cache data type. 'auto' follows --dtype (default: bfloat16).") @@ -100,19 +126,42 @@ 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, + pipeline_parallel_size=args.pipeline_parallel_size, + enable_expert_parallel=args.enable_expert_parallel, + expert_placement_strategy=args.expert_placement_strategy, + all2all_backend=args.all2all_backend, + devices=devices, + data_parallel_routing=args.data_parallel_routing, + ) + _validate_model_topology(model_family, args, parallel_config) + device_groups = parallel_config.replica_device_groups + first_group = device_groups[0] + worker_device_ids = first_group if parallel_config.data_parallel_size == 1 else devices + 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=args.device, - executor_cls="PyptoQwen14BExecutor", + device_id=first_group[0], + device_ids=worker_device_ids, + parallel_config=parallel_config, + 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, ) @@ -144,6 +193,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", "")).lower() + architectures = {str(item).lower() for item in config_data.get("architectures", [])} + 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", {}) + 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, *, @@ -155,15 +258,18 @@ def run_serve( except ImportError as e: raise ImportError("Serving mode requires uvicorn. Install with: pip install uvicorn") from e - from ..core.async_engine import AsyncLLMEngine - from ..core.server import create_serving_app - from ..core.tokenizer import TransformersTokenizerAdapter + from python.core.async_engine import AsyncLLMEngineClient, DPEngineCore + from python.core.server import create_serving_app + from python.core.tokenizer import TransformersTokenizerAdapter model_id = config.model_id get_profiler(process_name="pypto-serving-api") tokenizer = TransformersTokenizerAdapter.from_pretrained(config.model_dir) + engine_cls = DPEngineCore + if config.parallel_config is not None and config.parallel_config.data_parallel_size > 1: + engine_cls = AsyncLLMEngineClient - async_engine = AsyncLLMEngine( + async_engine = engine_cls( config=config, tokenizer=tokenizer, eos_token_id=tokenizer.eos_token_id, @@ -183,7 +289,7 @@ async def shutdown(): print(f"Starting PyPTO serving on {host}:{port}") print(f" Model: {model_id} (loaded in worker process)") - print(f" Platform: {config.platform}, Device: {config.device_id}") + print(f" Platform: {config.platform}, Device groups: {_format_device_groups(config)}") print(f" Max running requests: {config.max_num_running_reqs}") print(f" Max scheduled tokens/iter: {config.max_num_scheduled_tokens}") print(f" Chunked prefill threshold: {config.long_prefill_token_threshold}") @@ -194,6 +300,13 @@ async def shutdown(): uvicorn.run(app, host=host, port=port, log_level="info") +def _format_device_groups(config: EngineConfig) -> str: + parallel_config = config.parallel_config + if parallel_config is None: + return str(list(config.worker_device_ids())) + return str([list(group) for group in parallel_config.replica_device_groups]) + + def _validate_backend(backend: str) -> None: if backend != "npu": raise ValueError(f"Only NPU backend is supported, got: {backend}") @@ -215,7 +328,7 @@ def main(argv: Sequence[str] | None = None) -> int: def _ensure_core_imports() -> None: - global RuntimeConfig + global ParallelConfig, RuntimeConfig, parse_device_ids if RuntimeConfig is None: try: @@ -224,6 +337,16 @@ def _ensure_core_imports() -> None: from python.core import RuntimeConfig as imported_runtime_config RuntimeConfig = imported_runtime_config + if ParallelConfig is None or parse_device_ids is 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: + from python.core.parallel import ParallelConfig as imported_parallel_config + from python.core.parallel import parse_device_ids as imported_parse_device_ids + + ParallelConfig = imported_parallel_config + parse_device_ids = imported_parse_device_ids @contextlib.contextmanager diff --git a/python/core/__init__.py b/python/core/__init__.py index 5e39670..2a2786c 100644 --- a/python/core/__init__.py +++ b/python/core/__init__.py @@ -7,8 +7,19 @@ # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +from .async_engine import AsyncLLMEngineClient, DPEngineCore, EngineConfig from .engine import LLMEngine from .model_loader import ModelLoader +from .parallel import ParallelConfig from .types import GenerateConfig, RuntimeConfig -__all__ = ["GenerateConfig", "LLMEngine", "ModelLoader", "RuntimeConfig"] +__all__ = [ + "AsyncLLMEngineClient", + "DPEngineCore", + "EngineConfig", + "GenerateConfig", + "LLMEngine", + "ModelLoader", + "ParallelConfig", + "RuntimeConfig", +] diff --git a/python/core/api.py b/python/core/api.py index 5e39670..2a2786c 100644 --- a/python/core/api.py +++ b/python/core/api.py @@ -7,8 +7,19 @@ # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +from .async_engine import AsyncLLMEngineClient, DPEngineCore, EngineConfig from .engine import LLMEngine from .model_loader import ModelLoader +from .parallel import ParallelConfig from .types import GenerateConfig, RuntimeConfig -__all__ = ["GenerateConfig", "LLMEngine", "ModelLoader", "RuntimeConfig"] +__all__ = [ + "AsyncLLMEngineClient", + "DPEngineCore", + "EngineConfig", + "GenerateConfig", + "LLMEngine", + "ModelLoader", + "ParallelConfig", + "RuntimeConfig", +] diff --git a/python/core/async_engine.py b/python/core/async_engine.py index 5374389..0a20ef8 100644 --- a/python/core/async_engine.py +++ b/python/core/async_engine.py @@ -11,18 +11,44 @@ import asyncio import logging +import os import queue import time from collections.abc import AsyncGenerator -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +from typing import Callable from .kv_cache import KvCacheManager from python.profile import profile_instant, profile_span +from .parallel import ParallelConfig from .scheduler import Request, RequestStatus, Scheduler, SchedulerConfig, SchedulerOutput from .types import RuntimeConfig, StepOutput, WorkerCommand from .serving_worker import spawn_worker logger = logging.getLogger(__name__) +_DEFAULT_WORKER_INIT_TIMEOUT_SECONDS = 600.0 +_DEFAULT_WORKER_STEP_TIMEOUT_SECONDS = 300.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() -> float: + return _positive_env_timeout_seconds("PYPTO_WORKER_STEP_TIMEOUT", _DEFAULT_WORKER_STEP_TIMEOUT_SECONDS) @dataclass @@ -34,6 +60,9 @@ class EngineConfig: # Device / executor platform: str = "a2a3" device_id: int = 0 + device_ids: tuple[int, ...] = () + parallel_config: ParallelConfig | None = None + dp_rank: int = 0 executor_cls: str = "PyptoQwen14BExecutor" executor_kwargs: dict = field(default_factory=dict) @@ -50,6 +79,14 @@ class EngineConfig: enable_prefix_cache: bool = True enable_chunk_prefill: bool = True + def worker_device_ids(self) -> tuple[int, ...]: + """Return the device ids this engine worker should own.""" + if self.device_ids: + return tuple(int(device) for device in self.device_ids) + if self.parallel_config is not None and self.parallel_config.data_parallel_size == 1: + return self.parallel_config.replica_device_groups[0] + return (int(self.device_id),) + @dataclass class _RequestContext: @@ -65,12 +102,13 @@ class TokenOutput: finish_reason: str = "" -class AsyncLLMEngine: - """Async engine with multiprocess worker for NPU execution. +class DPEngineCore: + """Engine core for one data-parallel replica. - Architecture: - Main process: scheduler + API serving + output processing - Worker process: NPU device + model execution (single-card, extensible to multi-card) + A core owns all mutable serving state for one replica: one scheduler, one + KV cache manager, one worker process, one executor/model runtime, and one + tensor-parallel device group. Requests assigned to this core are scheduled + only against this core's local KV cache and worker state. """ def __init__( @@ -117,21 +155,26 @@ def __init__( async def start(self) -> None: """Start worker process and engine loop.""" - with profile_span("AsyncLLMEngine.start", cat="serving"): - process, input_q, output_q, ready_event = spawn_worker(self.config) - self._worker_process = process - self._input_queue = input_q - self._output_queue = output_q - - logger.info("Waiting for worker to initialize model...") - await asyncio.to_thread(ready_event.wait, timeout=600) - if not ready_event.is_set(): - raise RuntimeError("Worker failed to initialize within timeout") - logger.info("Worker ready") - - self._running = True - self._loop_task = asyncio.create_task(self._engine_loop()) - logger.info("AsyncLLMEngine started") + try: + with profile_span("DPEngineCore.start", cat="serving"): + process, input_q, output_q, ready_event = spawn_worker(self.config) + self._worker_process = process + self._input_queue = input_q + self._output_queue = output_q + + logger.info("Waiting for worker to initialize model...") + init_timeout = _worker_init_timeout_seconds() + await asyncio.to_thread(ready_event.wait, timeout=init_timeout) + if not ready_event.is_set(): + raise RuntimeError(f"Worker failed to initialize within {init_timeout:g}s") + logger.info("Worker ready") + + self._running = True + self._loop_task = asyncio.create_task(self._engine_loop()) + logger.info("DPEngineCore started") + except BaseException: + await self.stop() + raise async def stop(self) -> None: """Stop engine loop and worker process.""" @@ -148,23 +191,36 @@ async def stop(self) -> None: if self._worker_process.is_alive(): self._worker_process.terminate() self._worker_process = None - logger.info("AsyncLLMEngine stopped") + logger.info("DPEngineCore stopped") def generate_request_id(self) -> str: self._request_counter += 1 return f"serving-req-{self._request_counter}" + def pending_token_load(self) -> int: + """Estimate unfinished work for routing new data-parallel requests.""" + load = 0 + for request in self.scheduler.requests.values(): + if request.status.is_finished: + continue + prompt_remaining = max(0, request.num_prompt_tokens - request.num_computed_tokens) + generation_remaining = max(0, request.max_new_tokens - len(request.output_token_ids)) + load += prompt_remaining + generation_remaining + return load + async def add_request( self, request_id: str, prompt: str, config, + *, + on_queued: Callable[[], None] | None = None, ) -> AsyncGenerator[TokenOutput, None]: """Add a request and yield token outputs as they are generated.""" if self.tokenizer is None: raise RuntimeError("Tokenizer is required for request processing") with profile_span( - "AsyncLLMEngine.add_request", + "DPEngineCore.add_request", cat="serving", args={"request_id": request_id, "max_new_tokens": config.max_new_tokens}, ): @@ -189,6 +245,8 @@ async def add_request( ctx = _RequestContext(request=request) self._request_contexts[request_id] = ctx self.scheduler.add_request(request) + if on_queued is not None: + on_queued() profile_instant( "request.queued", cat="serving", @@ -245,11 +303,12 @@ async def _engine_loop(self) -> None: try: with profile_span("scheduler.wait_worker_output", cat="scheduler"): + step_timeout = _worker_step_timeout_seconds() 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 @@ -314,3 +373,146 @@ def _handle_step_error(self, scheduler_output: SchedulerOutput) -> None: TokenOutput(finished=True, finish_reason="error") ) self.scheduler.abort_request(sr.request.request_id) + + +class AsyncLLMEngineClient: + """Async serving client that routes requests across DP engine cores. + + The client owns one or more ``DPEngineCore`` instances and exposes the + server-facing async API: ``start``, ``stop``, ``add_request``, + ``abort_request``, and ``generate_request_id``. For ``data_parallel_size=1`` + it can still wrap a single core, while for DP>1 it selects a core for each + request and records request placement so aborts are sent to the correct + replica. It does not own scheduler, KV cache, model, or worker state itself. + """ + + def __init__( + self, + config: EngineConfig, + tokenizer=None, + eos_token_id: int | None = None, + bos_token_id: int | None = None, + *, + core_factory: Callable[..., DPEngineCore] = DPEngineCore, + ) -> None: + parallel = config.parallel_config + if parallel is None: + parallel = ParallelConfig(devices=config.worker_device_ids()) + config = replace(config, parallel_config=parallel) + + self.config = config + self.tokenizer = tokenizer + self.eos_token_id = eos_token_id + self.bos_token_id = bos_token_id + self.parallel_config = parallel + self._request_counter = 0 + self._route_counter = 0 + self._request_to_replica: dict[str, int] = {} + self._route_extra_load = [0 for _ in parallel.replica_device_groups] + self._cores: list[DPEngineCore] = [] + + for dp_rank, device_group in enumerate(parallel.replica_device_groups): + replica_parallel = parallel.for_replica(device_group) + replica_config = replace( + config, + device_id=device_group[0], + device_ids=device_group, + parallel_config=replica_parallel, + dp_rank=dp_rank, + ) + self._cores.append( + core_factory( + config=replica_config, + tokenizer=tokenizer, + eos_token_id=eos_token_id, + bos_token_id=bos_token_id, + ) + ) + + async def start(self) -> None: + """Start all DP engine cores in parallel.""" + tasks = [asyncio.create_task(core.start()) for core in self._cores] + try: + await asyncio.gather(*tasks) + except BaseException: + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self.stop() + raise + + async def stop(self) -> None: + """Stop all DP engine cores.""" + for core in reversed(self._cores): + await core.stop() + + def generate_request_id(self) -> str: + self._request_counter += 1 + return f"serving-req-{self._request_counter}" + + def pending_token_load(self) -> int: + return sum(core.pending_token_load() for core in self._cores) + + async def add_request( + self, + request_id: str, + prompt: str, + config, + ) -> AsyncGenerator[TokenOutput, None]: + replica_idx = self._select_replica() + request_load = self._estimate_request_load(prompt, config) + self._route_extra_load[replica_idx] += request_load + self._request_to_replica[request_id] = replica_idx + route_extra_active = True + + def clear_route_extra_load() -> None: + nonlocal route_extra_active + if not route_extra_active: + return + self._route_extra_load[replica_idx] = max( + 0, + self._route_extra_load[replica_idx] - request_load, + ) + route_extra_active = False + + try: + core = self._cores[replica_idx] + async for output in core.add_request( + request_id, + prompt, + config, + on_queued=clear_route_extra_load, + ): + yield output + finally: + self._request_to_replica.pop(request_id, None) + clear_route_extra_load() + + async def abort_request(self, request_id: str) -> None: + replica_idx = self._request_to_replica.get(request_id) + if replica_idx is not None: + await self._cores[replica_idx].abort_request(request_id) + return + for core in self._cores: + await core.abort_request(request_id) + + def _select_replica(self) -> int: + loads = [ + core.pending_token_load() + self._route_extra_load[idx] + for idx, core in enumerate(self._cores) + ] + replica_count = len(self._cores) + ordered = [ + (loads[idx], (idx - self._route_counter) % replica_count, idx) + for idx in range(replica_count) + ] + replica_idx = min(ordered)[2] + self._route_counter = (replica_idx + 1) % replica_count + return replica_idx + + def _estimate_request_load(self, prompt: str, config) -> int: + prompt_tokens = 0 + if self.tokenizer is not None: + prompt_tokens = len(self.tokenizer.encode(prompt)) + return prompt_tokens + int(getattr(config, "max_new_tokens", 0)) diff --git a/python/core/model_loader.py b/python/core/model_loader.py index 8db02a1..c6c0404 100644 --- a/python/core/model_loader.py +++ b/python/core/model_loader.py @@ -276,12 +276,153 @@ 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)) + if (model_path / "tokenizer.json").exists(): + tokenizer = TransformersTokenizerAdapter.from_tokenizer_file(str(model_path)) + else: + tokenizer = TransformersTokenizerAdapter.from_pretrained( + str(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) + index_data = json.loads(index_path.read_text()) + weight_map = dict(index_data.get("weight_map", {})) + _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/parallel.py b/python/core/parallel.py new file mode 100644 index 0000000..d6eca2d --- /dev/null +++ b/python/core/parallel.py @@ -0,0 +1,98 @@ +# 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 + +from dataclasses import dataclass + + +_SUPPORTED_ROUTING_POLICIES = {"least_pending_tokens"} + + +@dataclass(frozen=True) +class ParallelConfig: + """Serving parallelism contract for logical model replicas.""" + + data_parallel_size: int = 1 + tensor_parallel_size: int = 1 + pipeline_parallel_size: int = 1 + enable_expert_parallel: bool = False + expert_placement_strategy: str = "linear" + all2all_backend: str = "none" + devices: tuple[int, ...] = (0,) + data_parallel_routing: str = "least_pending_tokens" + + def __post_init__(self) -> None: + devices = tuple(int(device) for device in self.devices) + object.__setattr__(self, "devices", devices) + + if self.data_parallel_size < 1: + raise ValueError("data_parallel_size must be >= 1") + if self.tensor_parallel_size < 1: + raise ValueError("tensor_parallel_size must be >= 1") + if self.pipeline_parallel_size < 1: + raise ValueError("pipeline_parallel_size must be >= 1") + if self.pipeline_parallel_size != 1: + raise ValueError("pipeline_parallel_size > 1 is not supported yet") + if self.enable_expert_parallel: + raise ValueError("expert parallel is not supported yet") + if self.data_parallel_routing not in _SUPPORTED_ROUTING_POLICIES: + supported = ", ".join(sorted(_SUPPORTED_ROUTING_POLICIES)) + raise ValueError( + f"unsupported data_parallel_routing={self.data_parallel_routing!r}; " + f"supported policies: {supported}" + ) + if not devices: + raise ValueError("devices must contain at least one device id") + if len(set(devices)) != len(devices): + raise ValueError(f"devices must not contain duplicates: {devices}") + + expected_devices = self.data_parallel_size * self.tensor_parallel_size + if len(devices) != expected_devices: + raise ValueError( + "number of devices must equal data_parallel_size * tensor_parallel_size: " + f"devices={len(devices)}, data_parallel_size={self.data_parallel_size}, " + f"tensor_parallel_size={self.tensor_parallel_size}" + ) + + @property + def replica_device_groups(self) -> tuple[tuple[int, ...], ...]: + """Return devices grouped by DP replica, each group being one TP group.""" + groups = [] + for dp_rank in range(self.data_parallel_size): + start = dp_rank * self.tensor_parallel_size + end = start + self.tensor_parallel_size + groups.append(self.devices[start:end]) + return tuple(groups) + + def for_replica(self, device_group: tuple[int, ...]) -> "ParallelConfig": + """Return a single-DP-replica view for one worker process.""" + return ParallelConfig( + data_parallel_size=1, + tensor_parallel_size=len(device_group), + pipeline_parallel_size=self.pipeline_parallel_size, + enable_expert_parallel=self.enable_expert_parallel, + expert_placement_strategy=self.expert_placement_strategy, + all2all_backend=self.all2all_backend, + devices=device_group, + data_parallel_routing=self.data_parallel_routing, + ) + + +def parse_device_ids(value: str | None, *, default_device: int = 0) -> tuple[int, ...]: + """Parse a comma-separated device list, falling back to one default device.""" + if value is None or not value.strip(): + return (int(default_device),) + parts = [part.strip() for part in value.split(",")] + if any(not part for part in parts): + raise ValueError(f"invalid devices list: {value!r}") + try: + return tuple(int(part) for part in parts) + except ValueError as exc: + raise ValueError(f"devices must be a comma-separated list of integers: {value!r}") from exc diff --git a/python/core/pypto_executor.py b/python/core/pypto_executor.py index 6785d39..052b775 100644 --- a/python/core/pypto_executor.py +++ b/python/core/pypto_executor.py @@ -12,6 +12,7 @@ import contextlib import logging from abc import ABC, abstractmethod +from collections.abc import Sequence from .executor import ModelExecutor from .kv_cache import KvCacheManager @@ -40,12 +41,16 @@ def __init__( *, platform: str = "a2a3sim", device_id: int = 0, + device_ids: Sequence[int] | None = None, save_kernels_dir: str | None = None, ) -> None: """Initialize common PyPTO runtime options and model registries.""" super().__init__(kv_cache_manager) self._platform = platform - self._device_id = device_id + self._device_ids = tuple(int(device) for device in (device_ids if device_ids is not None else (device_id,))) + if not self._device_ids: + raise ValueError("device_ids must contain at least one device id") + self._device_id = self._device_ids[0] self._save_kernels_dir = save_kernels_dir self._runners: dict[str, ModelRunner] = {} self._compiled: dict[str, object] = {} diff --git a/python/core/server.py b/python/core/server.py index 7e70fd5..42253ab 100644 --- a/python/core/server.py +++ b/python/core/server.py @@ -14,7 +14,7 @@ import time import uuid -from .async_engine import AsyncLLMEngine +from .async_engine import AsyncLLMEngineClient, DPEngineCore from python.profile import profile_instant, profile_span from .types import GenerateConfig @@ -91,7 +91,7 @@ class ChatCompletionResponse(BaseModel): # --- Server --- class ServingServer: - def __init__(self, async_engine: AsyncLLMEngine, model_id: str) -> None: + def __init__(self, async_engine: AsyncLLMEngineClient | DPEngineCore, model_id: str) -> None: self.engine = async_engine self.model_id = model_id self.app = FastAPI(title="PyPTO Serving") @@ -275,6 +275,6 @@ def _map_finish_reason(reason: str) -> str: return mapping.get(reason, "stop") -def create_serving_app(async_engine: AsyncLLMEngine, model_id: str) -> FastAPI: +def create_serving_app(async_engine: AsyncLLMEngineClient | DPEngineCore, model_id: str) -> FastAPI: server = ServingServer(async_engine, model_id) return server.app diff --git a/python/core/serving_worker.py b/python/core/serving_worker.py index 0a7662f..69fa132 100644 --- a/python/core/serving_worker.py +++ b/python/core/serving_worker.py @@ -58,16 +58,23 @@ def init_device_and_model(self) -> None: from .sampler import Sampler from .types import ModelRecord + device_ids = self.config.worker_device_ids() + device_label = ",".join(str(device_id) for device_id in device_ids) if mp.current_process().name != "MainProcess": - get_profiler(process_name=f"serving-worker-{self.config.device_id}") + get_profiler(process_name=f"serving-worker-{device_label}") with profile_span( "WorkerProcess.init_device_and_model", cat="worker", - args={"model_id": self.config.model_id, "device_id": self.config.device_id}, + args={ + "model_id": self.config.model_id, + "device_id": self.config.device_id, + "device_ids": list(device_ids), + "dp_rank": self.config.dp_rank, + }, ): logger.info( f"Worker initializing: platform={self.config.platform}, " - f"device={self.config.device_id}" + f"devices={list(device_ids)}, dp_rank={self.config.dp_rank}" ) self.sampler = Sampler() @@ -76,6 +83,7 @@ def init_device_and_model(self) -> None: self.executor = executor_cls( platform=self.config.platform, device_id=self.config.device_id, + device_ids=device_ids, **self.config.executor_kwargs, ) @@ -105,6 +113,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 @@ -121,7 +132,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) diff --git a/python/core/tokenizer.py b/python/core/tokenizer.py index 5d06b0d..938e93a 100644 --- a/python/core/tokenizer.py +++ b/python/core/tokenizer.py @@ -9,6 +9,7 @@ from __future__ import annotations +import json from dataclasses import dataclass from pathlib import Path @@ -50,20 +51,36 @@ 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 Exception: + 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 +103,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 bb3343e..aae584a 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -105,6 +105,7 @@ class RuntimeModel: final_norm_weight: torch.Tensor lm_head: torch.Tensor layers: list[LayerWeights] + extra: dict[str, object] = field(default_factory=dict) @dataclass diff --git a/tests/test_batching.py b/tests/test_batching.py index 639860b..1d53c21 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -380,6 +380,13 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc manager.free(decode_alloc) +def test_pypto_executor_preserves_device_group(): + executor = PyptoExecutor(device_id=3, device_ids=[3, 4]) + + assert executor._device_id == 3 + assert executor._device_ids == (3, 4) + + def test_kernel_profile_helpers_emit_kernel_name_and_runtime_timing(): args = {"runtime": "tensormap_and_ringbuffer"} host_wall_us, device_wall_us = _run_timing_us( diff --git a/tests/test_cli.py b/tests/test_cli.py index f88e5cb..41be9aa 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -55,6 +55,9 @@ def test_build_serving_engine_config_uses_cli_args(tmp_path, monkeypatch): assert config.model_dir == str(model_dir.resolve()) assert config.platform == "a5" assert config.device_id == 2 + assert config.device_ids == (2,) + assert config.parallel_config.devices == (2,) + assert config.parallel_config.replica_device_groups == ((2,),) assert config.executor_cls == "PyptoQwen14BExecutor" assert config.executor_kwargs == { "pypto_root": "/tmp/pypto", @@ -73,6 +76,39 @@ def test_build_serving_engine_config_uses_cli_args(tmp_path, monkeypatch): assert config.enable_chunk_prefill is False +def test_build_serving_engine_config_accepts_parallel_topology(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + args = _parse_args([ + "--model", str(model_dir), + "--devices", "0,1,2,3", + "--dp", "2", + "--tp", "2", + ]) + + config = cli.build_serving_engine_config(args) + + assert config.device_id == 0 + assert config.device_ids == (0, 1, 2, 3) + assert config.parallel_config.data_parallel_size == 2 + assert config.parallel_config.tensor_parallel_size == 2 + assert config.parallel_config.replica_device_groups == ((0, 1), (2, 3)) + + +def test_build_serving_engine_config_rejects_invalid_parallel_topology(tmp_path): + model_dir = tmp_path / "model" + model_dir.mkdir() + args = _parse_args([ + "--model", str(model_dir), + "--devices", "0,1,2", + "--dp", "2", + "--tp", "2", + ]) + + with pytest.raises(ValueError, match="number of devices"): + cli.build_serving_engine_config(args) + + def test_parser_rejects_invalid_backend(tmp_path): model_dir = tmp_path / "model" model_dir.mkdir() diff --git a/tests/test_deepseek_v4.py b/tests/test_deepseek_v4.py new file mode 100644 index 0000000..66e993c --- /dev/null +++ b/tests/test_deepseek_v4.py @@ -0,0 +1,1155 @@ +# 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 examples.model.deepseek_v4.runner import npu_executor, weight_loader +from examples.model.deepseek_v4.runner.npu_runner import ( + DeepSeekV4CacheLayout, + DeepSeekV4CacheManager, + DeepSeekV4CompiledKernels, + DeepSeekV4InputBuilder, + DeepSeekV4L3Callable, + DeepSeekV4LayerPlan, + DeepSeekV4ModelRunner, + build_deepseek_v4_layer_plan, +) +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_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_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=16, 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=16, 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=16, 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 _LmHeadModule: + l3_lm_head = 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: { + "config": _ConfigModule, + "prefill_layer": _PrefillModule, + "prefill_fwd": _PrefillFwdModule, + "decode_layer": _DecodeModule, + "decode_fwd": _DecodeFwdModule, + "lm_head": _LmHeadModule, + "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", "deepseek_v4_lm_head"] + # The packed l3_prefill_fwd takes 79 layer-stacked tensors (incl. the CSA + # indexer wq_b/scale/weights_proj) and a single trailing num_tokens scalar. + assert len(compiled_args["deepseek_v4_prefill"]) == 80 + # The packed l3_decode_fwd takes 80 layer-stacked tensors (the in-kernel final + # RMSNorm + LM head replaced x_out with final_norm_w + lm_head_weight inputs and + # a logits output) plus a single trailing num_tokens scalar. + assert len(compiled_args["deepseek_v4_decode"]) == 81 + assert len(compiled_args["deepseek_v4_lm_head"]) == 3 + 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, 128, 4, 4096) + prefill_order = npu_executor._PREFILL_FWD_TENSOR_ORDER + # Packed prefill stacks FWD work caches x43 and the compress-state caches across + # the CSA (x21) and HCA (x20) groups; x_out collapses back to one HC stack. + 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, 2, 128, 1, 512) + assert prefill_args[prefill_order.index("idx_kv_cache")].shape == (8, 43, 2, 128, 1, 128) + assert prefill_args[prefill_order.index("hca_cmp_kv_state")].shape == (8, 20 * 1024, 8, 512) + assert prefill_args[prefill_order.index("csa_cmp_kv_state")].shape == (8, 21 * 2048, 4, 1024) + assert prefill_args[prefill_order.index("csa_inner_kv_state")].shape == (8, 21 * 2048, 4, 256) + assert prefill_args[prefill_order.index("cmp_sparse_indices")].shape == (8, 43 * 128, 640) + assert prefill_args[prefill_order.index("freqs_cos")].shape == (8, 43 * 8192, 64) + assert prefill_args[prefill_order.index("x_out")].shape == (8, 128, 4, 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 (64) x state_max_blocks rows. + assert compiled_args["deepseek_v4_decode"][decode_order.index("hca_compress_state")].shape == (8, 81920, 8, 1024) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_compress_state")].shape == (8, 87360, 4, 2048) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_inner_compress_state")].shape == ( + 8, + 87360, + 4, + 512, + ) + # In-kernel final RMSNorm + LM head: per-rank norm weight + TP vocab shards in, + # per-rank [ranks, decode_tokens, vocab] logits out. + assert compiled_args["deepseek_v4_decode"][decode_order.index("final_norm_w")].shape == (8, 4096) + assert compiled_args["deepseek_v4_decode"][decode_order.index("lm_head_weight")].shape == (8, 16384, 4096) + assert compiled_args["deepseek_v4_decode"][decode_order.index("logits")].shape == (8, 128, 129280) + assert compiled_args["deepseek_v4_lm_head"][1].shape == (8, 16384, 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_compile_rejects_non_8_way_lm_head_kernel(tmp_path, monkeypatch): + model_dir = _write_deepseek_model_dir(tmp_path) + kernel_dir = _write_deepseek_kernel_dir(tmp_path, lm_head_tp_size=2) + 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=16, max_seq_len=256, weight_dtype="int8"), + ) + executor = npu_executor.DeepSeekV4PyptoExecutor(platform="a2a3sim", device_ids=tuple(range(8))) + + with pytest.raises(ValueError, match="LM_HEAD_TP_SIZE=2 expected 8"): + executor._compile_model(loaded.runtime_model) + + +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_hc_input_builder_shapes_prefill_and_decode(): + # Exercise the batch-agnostic input layout at decode_batch=32 (production uses 64). + 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, 4, 2) + assert packed.tensors["hca_cmp_wkv"].shape == (2, 4096, 512) + 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() + 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, 2) + assert prepared.idx_block_table.shape == (8, 2) + assert prepared.position_ids.shape == (8, 128) + assert prepared.position_ids[0, :4].tolist() == [126, 127, 128, 3] + 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, :3].tolist() == [126, 127, 0] + assert prepared.ori_slot_mapping[0, 3].item() == -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, :3].tolist() == [126, 127, 128] + assert prepared.hca_state_slot_mapping[0, 3].item() == -1 + assert prepared.csa_state_slot_mapping.shape == (8, 128) + assert prepared.csa_state_slot_mapping[0, :3].tolist() == [126, 127, 128] + assert prepared.csa_state_slot_mapping[0, 3].item() == -1 + assert prepared.csa_inner_state_slot_mapping.shape == (8, 128) + assert prepared.csa_inner_state_slot_mapping[0, :3].tolist() == [126, 127, 128] + assert prepared.csa_inner_state_slot_mapping[0, 3].item() == -1 + 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] + assert prepared.input_ids[0, :6].tolist() == [5, 5, 9, 9, 4, 5] + assert prepared.position_ids[0, :6].tolist() == [127, 128, 4, 5, 127, 128] + assert prepared.kv_seq_lens[0, :4].tolist() == [129, 6, 129, 129] + assert prepared.block_table.shape == (8, 32, 1) + assert prepared.cmp_block_table.shape == (8, 32, 8) + assert prepared.ori_slot_mapping[0, :6].tolist() == [127, 0, 132, 133, 383, 256] + assert prepared.hca_cmp_slot_mapping[0, :6].tolist() == [0, -1, -1, -1, 2048, -1] + assert prepared.csa_cmp_slot_mapping[0, :6].tolist() == [31, -1, -1, -1, 2079, -1] + assert prepared.csa_idx_slot_mapping[0, :6].tolist() == [31, -1, -1, -1, 16415, -1] + assert prepared.csa_state_slot_mapping[0, :6].tolist() == [127, 128, 264, 265, 647, 648] + + +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_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_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 _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, +) -> Path: + kernel_dir = tmp_path / f"deepseek-v4-kernels-tp{lm_head_tp_size}" + kernel_dir.mkdir() + (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 = 64", + "DECODE_SEQ = 2", + "DECODE_TOKENS = DECODE_BATCH * DECODE_SEQ", + "PREFILL_BATCH = 1", + "PREFILL_SEQ = 128", + "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 the batch-agnostic input layout at decode_batch=32 (production uses 64). + 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_lm_head_packs_requested_rows_into_kernel_capacity(): + 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="", + lm_head=DeepSeekV4L3Callable(compiled=object(), name="lm_head"), + ) + runner = DeepSeekV4ModelRunner(compiled=compiled) + runner._global_weights = weight_loader.DeepSeekV4GlobalWeights( + embed_weight=torch.empty(0), + final_norm_weight=torch.empty(0), + lm_head_weight=torch.zeros((layout.ranks, 4, 3), dtype=torch.bfloat16).share_memory_(), + 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), + ) + runner._static_lm_head_weight = runner._global_weights.lm_head_weight + hidden = torch.arange(layout.ranks * 6 * 3, dtype=torch.float32).reshape(layout.ranks, 6, 3).to(torch.bfloat16) + captured = {} + + runner._final_hidden = lambda _: hidden + + def fake_run_l3(callable_spec, hidden_states, lm_head_weight, logits): + del callable_spec, lm_head_weight + captured["hidden_states"] = hidden_states.clone() + logits.zero_() + logits[0, 0].copy_(torch.arange(logits.shape[-1], dtype=torch.float32)) + logits[0, 1].copy_(torch.arange(logits.shape[-1], dtype=torch.float32) + 10) + + runner._run_l3 = fake_run_l3 + + logits = runner._logits_for_hidden(torch.empty(0), active_rows=(5, 2)) + + assert torch.equal(captured["hidden_states"][:, :2], hidden[:, [5, 2]]) + assert torch.count_nonzero(captured["hidden_states"][:, 2:]) == 0 + assert logits.shape == (2, 5) + assert logits[0].tolist() == [0, 1, 2, 3, 4] + assert logits[1].tolist() == [10, 11, 12, 13, 14] + + +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_npu_prefix_chunk.py b/tests/test_npu_prefix_chunk.py index 7a80f5e..c3d8519 100644 --- a/tests/test_npu_prefix_chunk.py +++ b/tests/test_npu_prefix_chunk.py @@ -16,7 +16,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from python.core.async_engine import AsyncLLMEngine, EngineConfig +from python.core.async_engine import DPEngineCore, EngineConfig from python.core.tokenizer import TransformersTokenizerAdapter from python.core.types import GenerateConfig, RuntimeConfig @@ -71,7 +71,7 @@ async def main(): long_prefill_token_threshold=128, ) - engine = AsyncLLMEngine( + engine = DPEngineCore( config=engine_config, tokenizer=tokenizer, eos_token_id=tokenizer.eos_token_id, diff --git a/tests/test_parallel.py b/tests/test_parallel.py new file mode 100644 index 0000000..f7d8829 --- /dev/null +++ b/tests/test_parallel.py @@ -0,0 +1,270 @@ +# 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 asyncio +import sys +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from python.core.async_engine import AsyncLLMEngineClient, EngineConfig, TokenOutput +from python.core.parallel import ParallelConfig, parse_device_ids +from python.core.types import GenerateConfig + + +def test_parallel_config_groups_dp_replicas_into_tp_groups(): + config = ParallelConfig( + data_parallel_size=2, + tensor_parallel_size=2, + devices=(0, 1, 2, 3), + ) + + assert config.replica_device_groups == ((0, 1), (2, 3)) + assert config.for_replica((2, 3)).data_parallel_size == 1 + assert config.for_replica((2, 3)).devices == (2, 3) + + +def test_parallel_config_rejects_unsupported_modes(): + with pytest.raises(ValueError, match="pipeline_parallel_size"): + ParallelConfig(pipeline_parallel_size=2) + + with pytest.raises(ValueError, match="expert parallel"): + ParallelConfig(enable_expert_parallel=True) + + with pytest.raises(ValueError, match="duplicates"): + ParallelConfig(data_parallel_size=1, tensor_parallel_size=2, devices=(0, 0)) + + +def test_parse_device_ids_uses_default_device(): + assert parse_device_ids(None, default_device=3) == (3,) + assert parse_device_ids("0, 2,4", default_device=3) == (0, 2, 4) + + +def test_async_llm_engine_client_routes_to_least_pending_tokens(): + asyncio.run(_check_async_llm_engine_client_routes_to_least_pending_tokens()) + + +def test_async_llm_engine_client_does_not_double_count_admitted_requests(): + asyncio.run(_check_async_llm_engine_client_does_not_double_count_admitted_requests()) + + +def test_async_llm_engine_client_cleans_up_failed_startup(): + asyncio.run(_check_async_llm_engine_client_cleans_up_failed_startup()) + + +async def _check_async_llm_engine_client_routes_to_least_pending_tokens(): + created: list[_FakeCore] = [] + + def factory(**kwargs): + core = _FakeCore(**kwargs) + created.append(core) + return core + + engine = AsyncLLMEngineClient( + EngineConfig( + model_id="model", + model_dir="/tmp/model", + parallel_config=ParallelConfig( + data_parallel_size=2, + tensor_parallel_size=1, + devices=(0, 1), + ), + ), + tokenizer=_Tokenizer(), + core_factory=factory, + ) + created[0].load = 10 + created[1].load = 1 + + outputs = [ + output + async for output in engine.add_request( + "req-0", + "hello world", + GenerateConfig(max_new_tokens=4), + ) + ] + + assert [fake.config.device_ids for fake in created] == [(0,), (1,)] + assert created[0].requests == [] + assert created[1].requests == ["req-0"] + assert outputs == [TokenOutput(finished=True, finish_reason="FINISHED_LENGTH")] + + +async def _check_async_llm_engine_client_does_not_double_count_admitted_requests(): + created: list[_BlockingCore] = [] + + def factory(**kwargs): + core = _BlockingCore(**kwargs) + created.append(core) + return core + + engine = AsyncLLMEngineClient( + EngineConfig( + model_id="model", + model_dir="/tmp/model", + parallel_config=ParallelConfig( + data_parallel_size=2, + tensor_parallel_size=1, + devices=(0, 1), + ), + ), + tokenizer=_Tokenizer(), + core_factory=factory, + ) + created[1].load = 5 + + first_task = asyncio.create_task( + _collect_outputs( + engine.add_request( + "req-0", + "long prompt that should not remain route extra load", + GenerateConfig(max_new_tokens=4), + ) + ) + ) + await created[0].admitted.wait() + + outputs = [ + output + async for output in engine.add_request( + "req-1", + "short prompt", + GenerateConfig(max_new_tokens=1), + ) + ] + + created[0].release.set() + await first_task + + assert engine._route_extra_load == [0, 0] + assert created[0].requests == ["req-0", "req-1"] + assert created[1].requests == [] + assert outputs == [TokenOutput(finished=True, finish_reason="FINISHED_LENGTH")] + + +async def _check_async_llm_engine_client_cleans_up_failed_startup(): + created: list[_StartupCore] = [] + + def factory(**kwargs): + core = _StartupCore(should_fail=len(created) == 1, **kwargs) + created.append(core) + return core + + engine = AsyncLLMEngineClient( + EngineConfig( + model_id="model", + model_dir="/tmp/model", + parallel_config=ParallelConfig( + data_parallel_size=2, + tensor_parallel_size=1, + devices=(0, 1), + ), + ), + tokenizer=_Tokenizer(), + core_factory=factory, + ) + + with pytest.raises(RuntimeError, match="startup failed"): + await engine.start() + + assert [core.started for core in created] == [True, True] + assert [core.stopped for core in created] == [True, True] + + +async def _collect_outputs(outputs): + return [output async for output in outputs] + + +class _Tokenizer: + def encode(self, text: str) -> list[int]: + return [1 for _ in text.split()] + + +class _FakeCore: + def __init__(self, *, config, tokenizer=None, eos_token_id=None, bos_token_id=None) -> None: + self.config = config + self.load = 0 + self.requests: list[str] = [] + + async def start(self) -> None: + return None + + async def stop(self) -> None: + return None + + def pending_token_load(self) -> int: + return self.load + + async def add_request(self, request_id: str, prompt: str, config, *, on_queued=None): + self.requests.append(request_id) + if on_queued is not None: + on_queued() + yield TokenOutput(finished=True, finish_reason="FINISHED_LENGTH") + + async def abort_request(self, request_id: str) -> None: + return None + + +class _BlockingCore(_FakeCore): + def __init__(self, *, config, tokenizer=None, eos_token_id=None, bos_token_id=None) -> None: + super().__init__( + config=config, + tokenizer=tokenizer, + eos_token_id=eos_token_id, + bos_token_id=bos_token_id, + ) + self.admitted = asyncio.Event() + self.release = asyncio.Event() + + async def add_request(self, request_id: str, prompt: str, config, *, on_queued=None): + self.requests.append(request_id) + self.load = 1 + if on_queued is not None: + on_queued() + self.admitted.set() + if request_id == "req-0": + await self.release.wait() + yield TokenOutput(finished=True, finish_reason="FINISHED_LENGTH") + + +class _StartupCore(_FakeCore): + def __init__( + self, + *, + config, + tokenizer=None, + eos_token_id=None, + bos_token_id=None, + should_fail: bool = False, + ) -> None: + super().__init__( + config=config, + tokenizer=tokenizer, + eos_token_id=eos_token_id, + bos_token_id=bos_token_id, + ) + self.should_fail = should_fail + self.started = False + self.stopped = False + + async def start(self) -> None: + self.started = True + if self.should_fail: + raise RuntimeError("startup failed") + + async def stop(self) -> None: + self.stopped = True