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..2ff4d8a --- /dev/null +++ b/examples/model/deepseek_v4/runner/npu_executor.py @@ -0,0 +1,777 @@ +# 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_TENSOR_ORDER, + _PREFILL_TENSOR_ORDER, + build_deepseek_v4_cache_group_specs, + build_deepseek_v4_layer_plan, +) +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, +} +_DEEPSEEK_V4_IMPORT_MODULES = ( + "config", + "moe", + "combine", + "decode_attention_csa", + "decode_attention_hca", + "decode_attention_swa", + "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_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): + """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]) + 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_layer"].l3_prefill_layer, + self._prefill_dummy_args(model, layout, modules["config"]), + ) + decode = self._compile_l3_callable( + "deepseek_v4_decode", + modules["decode_layer"].l3_decode_layer, + 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 + with _deepseek_v4_import_context(self._kernel_dir, pypto_root=pypto_root, ep=ranks, moe_shape="prefill"): + prefill = importlib.import_module("prefill_layer") + 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", "lm_head", "rope_tables") + } + modules["prefill_layer"] = prefill + 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 ``l3_prefill_layer``.""" + cfg = config_module.FLASH + values = self._layer_common_dummy_tensors( + model, + layout, + cfg, + tokens=layout.prefill_seq, + include_decode_indexer=True, + include_prefill_temporaries=True, + ) + ranks = layout.ranks + hidden = model.config.hidden_size + values.update( + { + "x_hc": torch.empty((ranks, layout.prefill_seq, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + "hca_cmp_kv_state": torch.empty( + ( + ranks, + 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, + layout.prefill_hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ), + dtype=torch.float32, + ), + "hca_compress_state_block_table": torch.empty( + (ranks, layout.prefill_hca_state_max_blocks), + dtype=torch.int32, + ), + "csa_cmp_kv_state": torch.empty( + ( + ranks, + 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, + layout.prefill_csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ), + dtype=torch.float32, + ), + "csa_compress_state_block_table": torch.empty( + (ranks, layout.prefill_csa_state_max_blocks), + dtype=torch.int32, + ), + "csa_inner_kv_state": torch.empty( + ( + ranks, + 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, + layout.prefill_csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ), + dtype=torch.float32, + ), + "csa_inner_compress_state_block_table": torch.empty( + (ranks, layout.prefill_csa_inner_state_max_blocks), + dtype=torch.int32, + ), + "kv_cache": torch.empty( + (ranks, layout.ori_max_blocks, layout.block_size, 1, model.config.head_dim), + dtype=torch.bfloat16, + ), + "ori_block_table": torch.empty((ranks, layout.ori_max_blocks), dtype=torch.int32), + "ori_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "cmp_kv": torch.empty( + (ranks, layout.prefill_cmp_max_blocks, layout.block_size, 1, model.config.head_dim), + dtype=torch.bfloat16, + ), + "cmp_block_table": torch.empty((ranks, layout.prefill_cmp_max_blocks), dtype=torch.int32), + "cmp_sparse_indices": torch.empty( + (ranks, layout.prefill_seq, layout.prefill_sparse_topk), + dtype=torch.int32, + ), + "cmp_sparse_lens": torch.empty((ranks, layout.prefill_seq), dtype=torch.int32), + "idx_kv_cache": torch.empty( + (ranks, layout.prefill_cmp_max_blocks, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + dtype=torch.bfloat16, + ), + "idx_block_table": torch.empty((ranks, layout.prefill_idx_max_blocks), dtype=torch.int32), + "position_ids": torch.empty((ranks, layout.prefill_seq), dtype=torch.int32), + "hca_cmp_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "hca_state_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "csa_cmp_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "csa_idx_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "csa_state_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "csa_inner_state_slot_mapping": torch.empty((ranks, layout.prefill_seq), dtype=torch.long), + "x_next": torch.empty((ranks, layout.prefill_seq, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + } + ) + return (*self._ordered_dummy_args(values, _PREFILL_TENSOR_ORDER), self._int32_arg(layout.prefill_seq), self._int32_arg(0)) + + def _decode_dummy_args( + self, + model: RuntimeModel, + layout: DeepSeekV4CacheLayout, + config_module: object, + ) -> tuple[Any, ...]: + """Return explicit serving dummy args for ``l3_decode_layer``.""" + cfg = config_module.FLASH + values = 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 + values.update( + { + "x_hc": torch.empty((ranks, tokens, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + "kv_cache": torch.empty( + (ranks, batch * layout.ori_max_blocks, layout.block_size, 1, model.config.head_dim), + dtype=torch.bfloat16, + ), + "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": torch.empty( + ( + ranks, + batch * layout.hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_STATE_DIM, + ), + dtype=torch.float32, + ), + "hca_compress_state_block_table": torch.empty( + (ranks, batch, layout.hca_state_max_blocks), + dtype=torch.int32, + ), + "csa_compress_state": torch.empty( + ( + ranks, + batch * layout.csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_STATE_DIM, + ), + dtype=torch.float32, + ), + "csa_compress_state_block_table": torch.empty( + (ranks, batch, layout.csa_state_max_blocks), + dtype=torch.int32, + ), + "csa_inner_compress_state": torch.empty( + ( + ranks, + batch * layout.csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_STATE_DIM, + ), + dtype=torch.float32, + ), + "csa_inner_compress_state_block_table": torch.empty( + (ranks, batch, layout.csa_inner_state_max_blocks), + dtype=torch.int32, + ), + "cmp_kv": torch.empty( + (ranks, batch * layout.cmp_max_blocks, layout.block_size, 1, model.config.head_dim), + dtype=torch.bfloat16, + ), + "cmp_block_table": torch.empty((ranks, batch, layout.cmp_max_blocks), dtype=torch.int32), + "idx_kv_cache": torch.empty( + (ranks, batch * layout.idx_max_blocks, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + dtype=torch.bfloat16, + ), + "idx_block_table": torch.empty((ranks, batch, layout.idx_max_blocks), dtype=torch.int32), + "x_next": torch.empty((ranks, tokens, DEEPSEEK_V4_HC_MULT, hidden), dtype=torch.bfloat16), + } + ) + return ( + *self._ordered_dummy_args(values, _DECODE_TENSOR_ORDER), + self._int32_arg(0), + self._int32_arg(layout.decode_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", "decode_layer.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..3605d1e --- /dev/null +++ b/examples/model/deepseek_v4/runner/npu_runner.py @@ -0,0 +1,2746 @@ +# 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 DeepSeekV4PackedLayerWeights +from python.core.model_runner import ModelRunner +from python.core.types import ( + DecodeBatch, + DecodeResult, + KVCacheGroupSpec, + KVCacheSpec, + ModelConfig, + PrefillBatch, + PrefillResult, + RuntimeConfig, + RuntimeModel, +) + + +DEEPSEEK_V4_RANKS = 8 +DEEPSEEK_V4_HC_MULT = 4 +DEEPSEEK_V4_BLOCK_SIZE = 128 +DEEPSEEK_V4_DECODE_BATCH = 32 +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 = 64 +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 + + +def build_deepseek_v4_cache_group_specs( + num_hidden_layers: int, +) -> tuple[KVCacheGroupSpec, ...]: + """Build KVCacheGroupSpecs for all 6 DeepSeekV4 cache families.""" + all_layers = tuple(range(num_hidden_layers)) + block_size = DEEPSEEK_V4_BLOCK_SIZE + head_dim = DEEPSEEK_V4_HEAD_DIM + idx_head_dim = DEEPSEEK_V4_IDX_HEAD_DIM + return ( + KVCacheGroupSpec( + name="ori", + layer_indices=all_layers, + spec=KVCacheSpec(block_size=block_size, page_size_bytes=block_size * head_dim * 2), + max_blocks_per_seq=DEEPSEEK_V4_ORI_MAX_BLOCKS, + ), + KVCacheGroupSpec( + name="cmp", + layer_indices=all_layers, + spec=KVCacheSpec(block_size=block_size, page_size_bytes=block_size * head_dim * 2), + max_blocks_per_seq=DEEPSEEK_V4_CMP_MAX_BLOCKS, + ), + KVCacheGroupSpec( + name="idx", + layer_indices=all_layers, + spec=KVCacheSpec(block_size=block_size, page_size_bytes=block_size * idx_head_dim * 2), + max_blocks_per_seq=DEEPSEEK_V4_IDX_MAX_BLOCKS, + ), + KVCacheGroupSpec( + name="hca_state", + layer_indices=all_layers, + spec=KVCacheSpec( + block_size=DEEPSEEK_V4_C128_STATE_BLOCK_SIZE, + page_size_bytes=DEEPSEEK_V4_C128_STATE_BLOCK_SIZE * DEEPSEEK_V4_HCA_STATE_DIM * 2, + compress_ratio=1, + ), + max_blocks_per_seq=DEEPSEEK_V4_HCA_STATE_MAX_BLOCKS, + ), + KVCacheGroupSpec( + name="csa_state", + layer_indices=all_layers, + spec=KVCacheSpec( + block_size=DEEPSEEK_V4_C4_STATE_BLOCK_SIZE, + page_size_bytes=DEEPSEEK_V4_C4_STATE_BLOCK_SIZE * DEEPSEEK_V4_CSA_STATE_DIM * 2, + compress_ratio=1, + ), + max_blocks_per_seq=DEEPSEEK_V4_CSA_STATE_MAX_BLOCKS, + ), + KVCacheGroupSpec( + name="csa_inner_state", + layer_indices=all_layers, + spec=KVCacheSpec( + block_size=DEEPSEEK_V4_C4_STATE_BLOCK_SIZE, + page_size_bytes=DEEPSEEK_V4_C4_STATE_BLOCK_SIZE * DEEPSEEK_V4_CSA_INNER_STATE_DIM * 2, + compress_ratio=1, + ), + max_blocks_per_seq=DEEPSEEK_V4_CSA_INNER_STATE_MAX_BLOCKS, + ), + ) + + +_PREFILL_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_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_next", +) + +_DECODE_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", + "kv_cache", + "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", + "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", + "hca_compress_state_block_table", + "csa_cmp_wkv", + "csa_cmp_wgate", + "csa_cmp_ape", + "csa_cmp_norm_w", + "csa_compress_state", + "csa_compress_state_block_table", + "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", + "csa_inner_compress_state_block_table", + "cmp_kv", + "cmp_block_table", + "idx_kv_cache", + "idx_block_table", + "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_next", +) + +_PREFILL_INPUT_TENSOR_FIELDS = ( + "input_ids", + "position_ids", + "ori_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", +) + +_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. + + When no external block IDs are provided, falls back to contiguous + slot-based addressing for backwards compatibility. + """ + 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 block_table_from_ids( + self, + per_request_block_ids: Sequence[list[int]], + *, + max_blocks: int, + ) -> torch.Tensor: + """Build a block table from scheduler-allocated block IDs.""" + table = torch.zeros((len(per_request_block_ids), max_blocks), dtype=torch.int32) + for row, block_ids in enumerate(per_request_block_ids): + n = min(len(block_ids), max_blocks) + table[row, :n] = torch.tensor(block_ids[:n], 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 slot_mapping_from_ids( + self, + per_request_block_ids: Sequence[list[int]], + positions: Sequence[Sequence[int]], + *, + block_size: int | None = None, + compress_ratio: int = 1, + ) -> torch.Tensor: + """Map token positions to physical slots using scheduler block IDs.""" + block_size = self.layout.block_size if block_size is None else int(block_size) + max_tokens = max((len(row) for row in positions), default=0) + mapping = torch.full((len(per_request_block_ids), max_tokens), -1, dtype=torch.int64) + for row, (block_ids, row_positions) in enumerate( + zip(per_request_block_ids, positions, strict=True) + ): + for col, position in enumerate(row_positions): + logical = int(position) // compress_ratio + block_idx = logical // block_size + offset = logical % block_size + if block_idx < len(block_ids): + mapping[row, col] = block_ids[block_idx] * block_size + offset + return mapping + + def sliding_window_slot_mapping_from_ids( + self, + per_request_block_ids: Sequence[list[int]], + positions: Sequence[Sequence[int]], + *, + kernel_rows: int, + ) -> torch.Tensor: + """Map positions into sliding-window cache using scheduler block IDs.""" + active = list(zip(per_request_block_ids, positions, strict=True)) + if not active: + raise ValueError("block_ids must not be empty") + while len(active) < kernel_rows: + active.append(active[0]) + mapping = torch.full( + (kernel_rows, max((len(pos) for _, pos in active), default=0)), + -1, dtype=torch.int64, + ) + for row_idx, (block_ids, row_positions) in enumerate(active): + for col, position in enumerate(row_positions): + window_slot = int(position) % self.layout.block_size + block_idx = window_slot // self.layout.block_size + offset = window_slot % self.layout.block_size + if block_idx < len(block_ids): + mapping[row_idx, col] = block_ids[block_idx] * self.layout.block_size + offset + elif block_ids: + mapping[row_idx, col] = block_ids[0] * self.layout.block_size + window_slot + return mapping + + def compressed_slot_mapping_from_ids( + self, + per_request_block_ids: Sequence[list[int]], + positions: Sequence[Sequence[int]], + *, + compress_ratio: int, + block_size: int | None = None, + kernel_rows: int, + ) -> torch.Tensor: + """Map compression-boundary positions using scheduler block IDs.""" + block_size = self.layout.block_size if block_size is None else int(block_size) + active: list[tuple[list[int], Sequence[int]]] = list( + zip(per_request_block_ids, positions, strict=True) + ) + if not active: + raise ValueError("block_ids must not be empty") + while len(active) < kernel_rows: + active.append(active[0]) + mapping = torch.full( + (kernel_rows, max((len(pos) for _, pos in active), default=0)), + -1, dtype=torch.int64, + ) + for row_idx, (block_ids, row_positions) in enumerate(active): + for col, position in enumerate(row_positions): + position = int(position) + if (position + 1) % compress_ratio != 0: + continue + logical = position // compress_ratio + block_idx = logical // block_size + offset = logical % block_size + if block_idx < len(block_ids): + mapping[row_idx, col] = block_ids[block_idx] * block_size + offset + return mapping + + def state_slot_mapping_from_ids( + self, + per_request_block_ids: Sequence[list[int]], + positions: Sequence[Sequence[int]], + *, + state_block_size: int, + kernel_rows: int, + ) -> torch.Tensor: + """Map positions into compressor-state cache using scheduler block IDs.""" + active: list[tuple[list[int], Sequence[int]]] = list( + zip(per_request_block_ids, positions, strict=True) + ) + if not active: + raise ValueError("block_ids must not be empty") + while len(active) < kernel_rows: + active.append(active[0]) + mapping = torch.full( + (kernel_rows, max((len(pos) for _, pos in active), default=0)), + -1, dtype=torch.int64, + ) + for row_idx, (block_ids, row_positions) in enumerate(active): + for col, position in enumerate(row_positions): + position = int(position) + block_idx = position // state_block_size + offset = position % state_block_size + if block_idx < len(block_ids): + mapping[row_idx, col] = block_ids[block_idx] * state_block_size + offset + 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 block_table_for_kernel_rows_from_ids( + self, + per_request_block_ids: Sequence[list[int]], + *, + max_blocks: int, + kernel_rows: int, + ) -> torch.Tensor: + """Build a fixed-row block table from scheduler block IDs.""" + if not per_request_block_ids: + raise ValueError("block_ids must not be empty") + active = self.block_table_from_ids(per_request_block_ids, max_blocks=max_blocks) + return self.replicate_first_row( + active, actual_rows=len(per_request_block_ids), 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 reserve two token rows per request for + MTP. Serving generation feeds the current token into the first row and + duplicates it into the second row until MTP dispatch is wired. + """ + 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, + ) + for row in range(self.layout.decode_batch): + source_row = 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, ...] + 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 _DeepSeekV4PrefillSharedBuffers: + """Reusable shared-memory buffers inherited by the L3 chip workers.""" + + x_hc_a: torch.Tensor + x_hc_b: torch.Tensor + tensors: dict[str, torch.Tensor] + cmp_sparse_indices_by_ratio: dict[int, torch.Tensor] + cmp_sparse_lens_by_ratio: dict[int, torch.Tensor] + temporaries: dict[str, 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 + 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_freqs_cos: torch.Tensor | None = None + self._static_freqs_sin: torch.Tensor | None = None + self._prefill_buffers: _DeepSeekV4PrefillSharedBuffers | None = None + self._decode_buffers: _DeepSeekV4DecodeSharedBuffers | None = None + self._layer_weight_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 + + 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_packed_layer_weights(self, layer: "DeepSeekV4LayerPlan") -> DeepSeekV4PackedLayerWeights: + """Load and pack one DeepSeekV4 layer for the per-layer PyPTO kernels.""" + return self._compiled.weight_store.load_packed_layer_weights( + layer.layer_id, + ranks=self._compiled.layout.ranks, + n_routed_experts=self._compiled.n_routed_experts, + compress_ratio=layer.compress_ratio, + include_tid2eid=layer.include_tid2eid, + include_gate_bias=layer.include_gate_bias, + ) + + def _has_group_block_ids(self, batch_group_ids: list[dict[str, list[int]]]) -> bool: + """Check if scheduler-provided group block IDs are available.""" + return bool(batch_group_ids) and any(g for g in batch_group_ids) + + def _get_group_ids( + self, batch_group_ids: list[dict[str, list[int]]], group_name: str, request_idx: int + ) -> list[int]: + """Extract block IDs for a specific group and request from batch group IDs.""" + if request_idx < len(batch_group_ids) and batch_group_ids[request_idx]: + return batch_group_ids[request_idx].get(group_name, []) + return [] + + 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) + + use_group_ids = self._has_group_block_ids(batch.block_ids_by_group) + + if use_group_ids: + ori_ids = [self._get_group_ids(batch.block_ids_by_group, "ori", 0)] + cmp_ids = [self._get_group_ids(batch.block_ids_by_group, "cmp", 0)] + idx_ids = [self._get_group_ids(batch.block_ids_by_group, "idx", 0)] + hca_state_ids = [self._get_group_ids(batch.block_ids_by_group, "hca_state", 0)] + csa_state_ids = [self._get_group_ids(batch.block_ids_by_group, "csa_state", 0)] + csa_inner_ids = [self._get_group_ids(batch.block_ids_by_group, "csa_inner_state", 0)] + + ori_block_table = self.cache_manager.block_table_from_ids( + ori_ids, max_blocks=layout.ori_max_blocks + )[0] + ori_slot_mapping = self.cache_manager.sliding_window_slot_mapping_from_ids( + ori_ids, [positions], kernel_rows=layout.prefill_batch + )[0] + cmp_block_table = self.cache_manager.block_table_from_ids( + cmp_ids, max_blocks=layout.prefill_cmp_max_blocks + )[0] + idx_block_table = self.cache_manager.block_table_from_ids( + idx_ids, max_blocks=layout.prefill_idx_max_blocks + )[0] + hca_state_bt = self.cache_manager.block_table_from_ids( + hca_state_ids, max_blocks=layout.prefill_hca_state_max_blocks + )[0] + csa_state_bt = self.cache_manager.block_table_from_ids( + csa_state_ids, max_blocks=layout.prefill_csa_state_max_blocks + )[0] + csa_inner_bt = self.cache_manager.block_table_from_ids( + csa_inner_ids, max_blocks=layout.prefill_csa_inner_state_max_blocks + )[0] + + hca_cmp_sm = self.cache_manager.compressed_slot_mapping_from_ids( + cmp_ids, [positions], compress_ratio=128, kernel_rows=layout.prefill_batch + )[0] + hca_state_sm = self.cache_manager.state_slot_mapping_from_ids( + hca_state_ids, [positions], + state_block_size=layout.c128_state_block_size, kernel_rows=layout.prefill_batch + )[0] + csa_cmp_sm = self.cache_manager.compressed_slot_mapping_from_ids( + cmp_ids, [positions], compress_ratio=4, kernel_rows=layout.prefill_batch + )[0] + csa_idx_sm = self.cache_manager.compressed_slot_mapping_from_ids( + idx_ids, [positions], compress_ratio=4, kernel_rows=layout.prefill_batch + )[0] + csa_state_sm = self.cache_manager.state_slot_mapping_from_ids( + csa_state_ids, [positions], + state_block_size=layout.c4_state_block_size, kernel_rows=layout.prefill_batch + )[0] + csa_inner_sm = self.cache_manager.state_slot_mapping_from_ids( + csa_inner_ids, [positions], + state_block_size=layout.c4_state_block_size, kernel_rows=layout.prefill_batch + )[0] + else: + ori_block_table = self.cache_manager.block_table( + [slot], max_blocks=layout.ori_max_blocks + )[0] + ori_slot_mapping = self.cache_manager.sliding_window_slot_mapping( + [slot], [positions], kernel_rows=layout.prefill_batch + )[0] + cmp_block_table = self.cache_manager.block_table( + [slot], max_blocks=layout.prefill_cmp_max_blocks + )[0] + idx_block_table = self.cache_manager.block_table( + [slot], max_blocks=layout.prefill_idx_max_blocks + )[0] + hca_state_bt = self.cache_manager.block_table( + [slot], max_blocks=layout.prefill_hca_state_max_blocks + )[0] + csa_state_bt = self.cache_manager.block_table( + [slot], max_blocks=layout.prefill_csa_state_max_blocks + )[0] + csa_inner_bt = self.cache_manager.block_table( + [slot], max_blocks=layout.prefill_csa_inner_state_max_blocks + )[0] + + hca_cmp_sm = self.cache_manager.compressed_slot_mapping( + [slot], [positions], max_blocks=layout.prefill_cmp_max_blocks, + compress_ratio=128, kernel_rows=layout.prefill_batch + )[0] + hca_state_sm = 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] + csa_cmp_sm = self.cache_manager.compressed_slot_mapping( + [slot], [positions], max_blocks=layout.prefill_cmp_max_blocks, + compress_ratio=4, kernel_rows=layout.prefill_batch + )[0] + csa_idx_sm = self.cache_manager.compressed_slot_mapping( + [slot], [positions], max_blocks=layout.prefill_idx_max_blocks, + compress_ratio=4, kernel_rows=layout.prefill_batch + )[0] + csa_state_sm = 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] + csa_inner_sm = 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] + + 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(ori_block_table), + ori_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(ori_slot_mapping, layout.prefill_seq) + ), + cmp_block_table=self._rank_stack(cmp_block_table), + idx_block_table=self._rank_stack(idx_block_table), + hca_compress_state_block_table=self._rank_stack(hca_state_bt), + csa_compress_state_block_table=self._rank_stack(csa_state_bt), + csa_inner_compress_state_block_table=self._rank_stack(csa_inner_bt), + hca_cmp_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(hca_cmp_sm, layout.prefill_seq) + ), + hca_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(hca_state_sm, layout.prefill_seq) + ), + csa_cmp_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(csa_cmp_sm, layout.prefill_seq) + ), + csa_idx_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(csa_idx_sm, layout.prefill_seq) + ), + csa_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(csa_state_sm, layout.prefill_seq) + ), + csa_inner_state_slot_mapping=self._rank_stack( + self._pad_prefill_mapping(csa_inner_sm, 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, + ) + x_hc = builder.decode_x_hc(batch.hidden_states.to(torch.bfloat16).cpu(), actual_batch=actual_batch) + decode_slots = (*slots, *((slots[0],) * (layout.decode_batch - actual_batch))) + decode_positions = (*positions, *((positions[0],) * (layout.decode_batch - actual_batch))) + + use_group_ids = self._has_group_block_ids(batch.block_ids_by_group) + + if use_group_ids: + ori_ids = [self._get_group_ids(batch.block_ids_by_group, "ori", i) for i in range(actual_batch)] + cmp_ids = [self._get_group_ids(batch.block_ids_by_group, "cmp", i) for i in range(actual_batch)] + idx_ids = [self._get_group_ids(batch.block_ids_by_group, "idx", i) for i in range(actual_batch)] + hca_state_ids = [self._get_group_ids(batch.block_ids_by_group, "hca_state", i) for i in range(actual_batch)] + csa_state_ids = [self._get_group_ids(batch.block_ids_by_group, "csa_state", i) for i in range(actual_batch)] + csa_inner_ids = [self._get_group_ids(batch.block_ids_by_group, "csa_inner_state", i) for i in range(actual_batch)] + + ori_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.sliding_window_slot_mapping_from_ids( + ori_ids, positions, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + hca_cmp_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.compressed_slot_mapping_from_ids( + cmp_ids, positions, compress_ratio=128, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + hca_state_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.state_slot_mapping_from_ids( + hca_state_ids, positions, + state_block_size=layout.c128_state_block_size, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_cmp_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.compressed_slot_mapping_from_ids( + cmp_ids, positions, compress_ratio=4, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_idx_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.compressed_slot_mapping_from_ids( + idx_ids, positions, compress_ratio=4, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_state_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.state_slot_mapping_from_ids( + csa_state_ids, positions, + state_block_size=layout.c4_state_block_size, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_inner_state_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.state_slot_mapping_from_ids( + csa_inner_ids, positions, + state_block_size=layout.c4_state_block_size, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows_from_ids( + ori_ids, max_blocks=layout.ori_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + cmp_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows_from_ids( + cmp_ids, 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_from_ids( + idx_ids, max_blocks=layout.idx_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + hca_state_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows_from_ids( + hca_state_ids, max_blocks=layout.hca_state_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + csa_state_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows_from_ids( + csa_state_ids, max_blocks=layout.csa_state_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + csa_inner_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows_from_ids( + csa_inner_ids, max_blocks=layout.csa_inner_state_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + else: + ori_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.sliding_window_slot_mapping( + slots, positions, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + hca_cmp_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.compressed_slot_mapping( + slots, positions, max_blocks=layout.cmp_max_blocks, + compress_ratio=128, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + hca_state_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.state_slot_mapping( + slots, positions, max_blocks=layout.hca_state_max_blocks, + state_block_size=layout.c128_state_block_size, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_cmp_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.compressed_slot_mapping( + slots, positions, max_blocks=layout.cmp_max_blocks, + compress_ratio=4, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_idx_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.compressed_slot_mapping( + slots, positions, max_blocks=layout.idx_max_blocks, + compress_ratio=4, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_state_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.state_slot_mapping( + slots, positions, max_blocks=layout.csa_state_max_blocks, + state_block_size=layout.c4_state_block_size, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + csa_inner_state_slot_mapping = self._mask_inactive_decode_slots( + self.cache_manager.state_slot_mapping( + slots, positions, max_blocks=layout.csa_inner_state_max_blocks, + state_block_size=layout.c4_state_block_size, kernel_rows=layout.decode_batch, + ), actual_batch, + ) + block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + slots, max_blocks=layout.ori_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + cmp_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + 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( + slots, max_blocks=layout.idx_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + hca_state_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + slots, max_blocks=layout.hca_state_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + csa_state_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + slots, max_blocks=layout.csa_state_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + csa_inner_block_table = self._rank_stack( + self.cache_manager.block_table_for_kernel_rows( + slots, max_blocks=layout.csa_inner_state_max_blocks, kernel_rows=layout.decode_batch, + ) + ) + + return DeepSeekV4PreparedDecodeInputs( + request_ids=tuple(batch.request_ids), + slots=tuple(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=block_table, + ori_slot_mapping=self._rank_stack(ori_slot_mapping.reshape(-1)), + cmp_block_table=cmp_block_table, + idx_block_table=idx_block_table, + hca_compress_state_block_table=hca_state_block_table, + csa_compress_state_block_table=csa_state_block_table, + csa_inner_compress_state_block_table=csa_inner_block_table, + 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 _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 and return logits.""" + 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() + inputs = self._stage_prefill_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." + ) + prefill_buffers = self._require_prefill_buffers() + x_hc = prefill_buffers.x_hc_a + x_next = prefill_buffers.x_hc_b + for layer in self._compiled.layer_plan: + weights = self._stage_layer_weights(self.load_packed_layer_weights(layer)) + temp = self._alloc_prefill_temporaries() + self._reset_prefill_temporaries(temp) + try: + args = self._prefill_layer_args(layer, weights, inputs, x_hc, x_next, temp) + self._debug_prefill_dispatch(layer, inputs, args) + try: + self._run_l3( + self._require_prefill_callable(), + *args, + self._int32_scalar(inputs.actual_tokens), + self._int32_scalar(layer.layer_id), + ) + except RuntimeError as exc: + raise RuntimeError( + "DeepSeekV4 prefill dispatch failed " + f"at layer {layer.layer_id} ({layer.attention_kind}, ratio={layer.compress_ratio}, " + f"tokens={inputs.actual_tokens})" + ) from exc + self._snapshot_prefill_cache(layer, temp) + finally: + self._free_prefill_temporaries(temp) + x_hc, x_next = x_next, x_hc + + logits = self._logits_for_hidden(x_hc, active_rows=(inputs.actual_tokens - 1,)) + return PrefillResult(last_hidden=None, logits=logits) + + def run_decode(self, model, batch: DecodeBatch) -> DecodeResult: + """Run all DeepSeekV4 hidden layers for one decode batch and return logits.""" + 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() + 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." + ) + if self._layer_weight_buffers is None: + self._stage_layer_weights(self.load_packed_layer_weights(self._compiled.layer_plan[0])) + self._require_prefill_cache_snapshots() + decode_buffers = self._require_decode_buffers() + x_hc = decode_buffers.x_hc_a + x_next = decode_buffers.x_hc_b + active_decode_tokens = inputs.actual_batch * self._compiled.layout.decode_seq + for layer in self._compiled.layer_plan: + weights = self._stage_layer_weights(self.load_packed_layer_weights(layer)) + self._load_decode_work_cache(layer, inputs.slots[0]) + args = self._decode_layer_args(layer, weights, inputs, x_hc, x_next) + self._debug_decode_dispatch(layer, inputs, args) + try: + self._run_l3( + self._require_decode_callable(), + *args, + self._int32_scalar(layer.layer_id), + self._int32_scalar(active_decode_tokens), + ) + except RuntimeError as exc: + raise RuntimeError( + "DeepSeekV4 decode dispatch failed " + f"at layer {layer.layer_id} ({layer.attention_kind}, ratio={layer.compress_ratio}, " + f"actual_batch={inputs.actual_batch}, slots={inputs.slots})" + ) from exc + self._snapshot_decode_work_cache(layer, inputs.slots[0]) + x_hc, x_next = x_next, x_hc + + active_rows = tuple(row * self._compiled.layout.decode_seq for row in range(inputs.actual_batch)) + logits = self._logits_for_hidden(x_hc, active_rows=active_rows) + hidden = self._final_hidden(x_hc)[0, list(active_rows)].float() + return DecodeResult(hidden_states=hidden, 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_layer_args( + self, + layer: DeepSeekV4LayerPlan, + weights: DeepSeekV4PackedLayerWeights, + inputs: DeepSeekV4PreparedPrefillInputs, + x_hc: torch.Tensor, + x_next: torch.Tensor, + temp: dict[str, Any], + ) -> tuple[Any, ...]: + sparse_indices, sparse_lens = inputs.sparse_inputs_for_ratio(layer.compress_ratio) + values = dict(weights.tensors) + values.update( + { + "x_hc": x_hc, + "freqs_cos": self._static_freqs_cos_tensor(), + "freqs_sin": self._static_freqs_sin_tensor(), + "hca_cmp_kv_state": temp["hca_cmp_kv_state"], + "hca_cmp_score_state": temp["hca_cmp_score_state"], + "csa_cmp_kv_state": temp["csa_cmp_kv_state"], + "csa_cmp_score_state": temp["csa_cmp_score_state"], + "csa_inner_kv_state": temp["csa_inner_kv_state"], + "csa_inner_score_state": temp["csa_inner_score_state"], + "kv_cache": temp["kv_cache"], + "cmp_kv": temp["cmp_kv"], + "idx_kv_cache": temp["idx_kv_cache"], + "ori_block_table": inputs.ori_block_table, + "ori_slot_mapping": inputs.ori_slot_mapping, + "cmp_block_table": inputs.cmp_block_table, + "cmp_sparse_indices": sparse_indices, + "cmp_sparse_lens": sparse_lens, + "idx_block_table": inputs.idx_block_table, + "position_ids": inputs.position_ids, + "hca_compress_state_block_table": inputs.hca_compress_state_block_table, + "csa_compress_state_block_table": inputs.csa_compress_state_block_table, + "csa_inner_compress_state_block_table": inputs.csa_inner_compress_state_block_table, + "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, + "x_next": x_next, + } + ) + return self._ordered_layer_args(values, _PREFILL_TENSOR_ORDER) + + def _decode_layer_args( + self, + layer: DeepSeekV4LayerPlan, + weights: DeepSeekV4PackedLayerWeights, + inputs: DeepSeekV4PreparedDecodeInputs, + x_hc: torch.Tensor, + x_next: torch.Tensor, + ) -> tuple[Any, ...]: + cache = self._require_decode_work_cache() + values = dict(weights.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, + "x_next": x_next, + } + ) + return self._ordered_layer_args(values, _DECODE_TENSOR_ORDER) + + 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, + layer: DeepSeekV4LayerPlan, + inputs: DeepSeekV4PreparedPrefillInputs, + args: Sequence[Any], + ) -> None: + if os.getenv("PYPTO_DSV4_DEBUG") != "1": + return + named_args = dict(zip(_PREFILL_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_next", + ) + 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 prefill dispatch " + f"layer={layer.layer_id} kind={layer.attention_kind} ratio={layer.compress_ratio} " + 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_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, + layer: DeepSeekV4LayerPlan, + inputs: DeepSeekV4PreparedDecodeInputs, + args: Sequence[Any], + ) -> None: + if os.getenv("PYPTO_DSV4_DEBUG") != "1": + return + named_args = dict(zip(_DECODE_TENSOR_ORDER, args, strict=True)) + interesting = ( + "x_hc", + "kv_cache", + "block_table", + "ori_slot_mapping", + "cmp_kv", + "cmp_block_table", + "idx_kv_cache", + "idx_block_table", + "hca_compress_state", + "hca_state_slot_mapping", + "csa_compress_state", + "csa_state_slot_mapping", + "csa_inner_compress_state", + "csa_inner_state_slot_mapping", + "position_ids", + "kv_seq_lens", + "input_ids", + "x_next", + ) + 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 decode dispatch " + f"layer={layer.layer_id} kind={layer.attention_kind} ratio={layer.compress_ratio} " + 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_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, + ) + + @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 _stage_prefill_inputs(self, inputs: DeepSeekV4PreparedPrefillInputs) -> DeepSeekV4PreparedPrefillInputs: + buffers = self._prefill_buffers + if buffers is None: + self._ensure_shared_host_allocation_before_worker("prefill inputs") + buffers = _DeepSeekV4PrefillSharedBuffers( + x_hc_a=self._new_shared_like(inputs.x_hc, name="x_hc"), + x_hc_b=self._new_shared_like(inputs.x_hc, name="x_hc_next"), + tensors={ + name: self._new_shared_like(getattr(inputs, name), name=name) + for name in _PREFILL_INPUT_TENSOR_FIELDS + }, + cmp_sparse_indices_by_ratio={ + ratio: self._new_shared_like(tensor, name=f"cmp_sparse_indices[{ratio}]") + for ratio, tensor in inputs.cmp_sparse_indices_by_ratio.items() + }, + cmp_sparse_lens_by_ratio={ + ratio: self._new_shared_like(tensor, name=f"cmp_sparse_lens[{ratio}]") + for ratio, tensor in inputs.cmp_sparse_lens_by_ratio.items() + }, + temporaries=self._new_prefill_temporaries(), + ) + self._prefill_buffers = buffers + + self._copy_shared(buffers.x_hc_a, inputs.x_hc, name="x_hc") + staged_values: dict[str, torch.Tensor] = {} + for name in _PREFILL_INPUT_TENSOR_FIELDS: + dst = buffers.tensors[name] + self._copy_shared(dst, getattr(inputs, name), name=name) + staged_values[name] = dst + + staged_sparse_indices: dict[int, torch.Tensor] = {} + for ratio, src in inputs.cmp_sparse_indices_by_ratio.items(): + dst = buffers.cmp_sparse_indices_by_ratio[ratio] + self._copy_shared(dst, src, name=f"cmp_sparse_indices[{ratio}]") + staged_sparse_indices[ratio] = dst + + staged_sparse_lens: dict[int, torch.Tensor] = {} + for ratio, src in inputs.cmp_sparse_lens_by_ratio.items(): + dst = buffers.cmp_sparse_lens_by_ratio[ratio] + self._copy_shared(dst, src, name=f"cmp_sparse_lens[{ratio}]") + staged_sparse_lens[ratio] = dst + + return replace( + inputs, + x_hc=buffers.x_hc_a, + cmp_sparse_indices_by_ratio=staged_sparse_indices, + cmp_sparse_lens_by_ratio=staged_sparse_lens, + **staged_values, + ) + + 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", + ), + 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 _stage_layer_weights(self, weights: DeepSeekV4PackedLayerWeights) -> DeepSeekV4PackedLayerWeights: + buffers = self._layer_weight_buffers + if buffers is None: + self._ensure_shared_host_allocation_before_worker("layer weights") + buffers = { + name: self._new_shared_like(tensor, name=f"layer_weight[{name}]") + for name, tensor in weights.tensors.items() + } + self._layer_weight_buffers = buffers + + missing = sorted(set(weights.tensors) - set(buffers)) + if missing: + raise KeyError(f"DeepSeekV4 shared layer-weight buffers are missing: {', '.join(missing)}") + + for name, tensor in weights.tensors.items(): + self._copy_shared(buffers[name], tensor, name=f"layer_weight[{name}]") + return DeepSeekV4PackedLayerWeights(layer_id=weights.layer_id, tensors=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_prefill_buffers(self) -> _DeepSeekV4PrefillSharedBuffers: + if self._prefill_buffers is None: + raise RuntimeError("DeepSeekV4 prefill shared buffers were not staged") + return self._prefill_buffers + + def _require_decode_buffers(self) -> _DeepSeekV4DecodeSharedBuffers: + if self._decode_buffers is None: + raise RuntimeError("DeepSeekV4 decode shared buffers were not staged") + return self._decode_buffers + + def _new_prefill_temporaries(self) -> dict[str, torch.Tensor]: + layout = self._compiled.layout + ranks = layout.ranks + return { + "kv_cache": self._shared_empty( + (ranks, layout.ori_max_blocks, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), + torch.bfloat16, + name="prefill_kv_cache", + ), + "cmp_kv": self._shared_empty( + (ranks, layout.prefill_cmp_max_blocks, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), + torch.bfloat16, + name="prefill_cmp_kv", + ), + "idx_kv_cache": self._shared_empty( + (ranks, layout.prefill_cmp_max_blocks, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + torch.bfloat16, + name="prefill_idx_kv_cache", + ), + "hca_cmp_kv_state": self._shared_empty( + ( + ranks, + layout.prefill_hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ), + torch.float32, + name="prefill_hca_cmp_kv_state", + ), + "hca_cmp_score_state": self._shared_empty( + ( + ranks, + layout.prefill_hca_state_max_blocks, + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ), + torch.float32, + name="prefill_hca_cmp_score_state", + ), + "csa_cmp_kv_state": self._shared_empty( + ( + ranks, + layout.prefill_csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ), + torch.float32, + name="prefill_csa_cmp_kv_state", + ), + "csa_cmp_score_state": self._shared_empty( + ( + ranks, + layout.prefill_csa_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ), + torch.float32, + name="prefill_csa_cmp_score_state", + ), + "csa_inner_kv_state": self._shared_empty( + ( + ranks, + layout.prefill_csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ), + torch.float32, + name="prefill_csa_inner_kv_state", + ), + "csa_inner_score_state": self._shared_empty( + ( + ranks, + layout.prefill_csa_inner_state_max_blocks, + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ), + torch.float32, + name="prefill_csa_inner_score_state", + ), + } + + 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 _alloc_prefill_temporaries(self) -> dict[str, Any]: + return self._require_prefill_buffers().temporaries + + @staticmethod + def _reset_prefill_temporaries(temp: dict[str, Any]) -> None: + for tensor in temp.values(): + if isinstance(tensor, torch.Tensor): + tensor.zero_() + + def _free_prefill_temporaries(self, temp: dict[str, Any]) -> None: + worker = self._l3_worker + if worker is None: + return + for tensor in temp.values(): + if isinstance(tensor, DeviceTensor): + worker.free_tensor(tensor) + + def _snapshot_prefill_cache(self, layer: DeepSeekV4LayerPlan, temp: dict[str, Any]) -> None: + """Capture the compact cache fields produced by one prefill layer.""" + snapshot: dict[str, torch.Tensor] = {} + for name in self._prefill_cache_snapshot_fields(layer): + tensor = temp[name] + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"DeepSeekV4 prefill cache field {name} must be a CPU tensor") + snapshot[name] = tensor.detach().cpu().contiguous().clone() + self._prefill_cache_snapshots[layer.layer_id] = DeepSeekV4LayerCacheSnapshot(snapshot) + + @staticmethod + def _prefill_cache_snapshot_fields(layer: DeepSeekV4LayerPlan) -> tuple[str, ...]: + fields = ["kv_cache", "cmp_kv"] + if layer.compress_ratio == 4: + fields.extend( + ( + "idx_kv_cache", + "csa_cmp_kv_state", + "csa_cmp_score_state", + "csa_inner_kv_state", + "csa_inner_score_state", + ) + ) + elif layer.compress_ratio == 128: + fields.extend(("hca_cmp_kv_state", "hca_cmp_score_state")) + return tuple(fields) + + 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 _load_decode_work_cache(self, layer: DeepSeekV4LayerPlan, slot: int) -> None: + 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}") + cache = self._require_decode_work_cache() + layout = self._compiled.layout + tensors = snapshot.tensors + + self._copy_snapshot_blocks_to_work(tensors["kv_cache"], cache.kv_cache, slot, layout.ori_max_blocks) + self._copy_snapshot_blocks_to_work(tensors["cmp_kv"], cache.cmp_kv, slot, layout.cmp_max_blocks) + if layer.compress_ratio == 4: + self._copy_snapshot_blocks_to_work(tensors["idx_kv_cache"], cache.idx_kv_cache, 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, + 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, + slot, + layout.csa_inner_state_max_blocks, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ) + elif layer.compress_ratio == 128: + self._copy_split_state_to_work( + tensors["hca_cmp_kv_state"], + tensors["hca_cmp_score_state"], + cache.hca_compress_state, + slot, + layout.hca_state_max_blocks, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ) + + def _snapshot_decode_work_cache(self, layer: DeepSeekV4LayerPlan, slot: int) -> None: + snapshot = self._prefill_cache_snapshots.get(layer.layer_id) + if snapshot is None: + return + cache = self._require_decode_work_cache() + layout = self._compiled.layout + tensors = snapshot.tensors + + self._copy_work_blocks_to_snapshot(cache.kv_cache, tensors["kv_cache"], slot, layout.ori_max_blocks) + self._copy_work_blocks_to_snapshot(cache.cmp_kv, tensors["cmp_kv"], slot, layout.cmp_max_blocks) + if layer.compress_ratio == 4: + self._copy_work_blocks_to_snapshot(cache.idx_kv_cache, tensors["idx_kv_cache"], slot, layout.idx_max_blocks) + self._copy_split_state_to_snapshot( + cache.csa_compress_state, + tensors["csa_cmp_kv_state"], + tensors["csa_cmp_score_state"], + slot, + layout.csa_state_max_blocks, + DEEPSEEK_V4_CSA_MAIN_OUT_DIM, + ) + self._copy_split_state_to_snapshot( + cache.csa_inner_compress_state, + tensors["csa_inner_kv_state"], + tensors["csa_inner_score_state"], + slot, + layout.csa_inner_state_max_blocks, + DEEPSEEK_V4_CSA_INNER_OUT_DIM, + ) + elif layer.compress_ratio == 128: + self._copy_split_state_to_snapshot( + cache.hca_compress_state, + tensors["hca_cmp_kv_state"], + tensors["hca_cmp_score_state"], + slot, + layout.hca_state_max_blocks, + DEEPSEEK_V4_HCA_MAIN_OUT_DIM, + ) + + @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_work_blocks_to_snapshot( + self, + work: torch.Tensor, + snapshot: torch.Tensor, + slot: int, + blocks_per_slot: int, + ) -> None: + del self + slot_slice = DeepSeekV4ModelRunner._slot_block_slice(slot, blocks_per_slot) + blocks = min(snapshot.shape[1], int(blocks_per_slot)) + snapshot[:, :blocks].copy_(work[:, slot_slice][:, :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 _copy_split_state_to_snapshot( + self, + work: torch.Tensor, + kv_state: torch.Tensor, + score_state: torch.Tensor, + slot: int, + blocks_per_slot: int, + out_dim: int, + ) -> None: + del self + slot_slice = DeepSeekV4ModelRunner._slot_block_slice(slot, blocks_per_slot) + src = work[:, slot_slice] + blocks = min(kv_state.shape[1], score_state.shape[1], int(blocks_per_slot)) + kv_state[:, :blocks].copy_(src[:, :blocks, ..., :out_dim]) + score_state[:, :blocks].copy_(src[:, :blocks, ..., out_dim : 2 * out_dim]) + + def _logits_for_hidden(self, x_hc: torch.Tensor, *, active_rows: Sequence[int]) -> 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) + 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]}" + ) + self._lm_head_hidden_buffer.zero_() + self._lm_head_hidden_buffer[:, : len(rows)].copy_(hidden[:, list(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, + ) + return self._lm_head_logits_buffer[0, : len(rows), :].contiguous() + + def _final_hidden(self, x_hc: torch.Tensor) -> torch.Tensor: + 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) + 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 _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() + uploaded: list[DeviceTensor] = [] + try: + l3_args = tuple(self._coerce_l3_arg(worker, arg, uploaded) for arg in args) + 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 + cache = DeepSeekV4LayerCache( + kv_cache=self._shared_empty( + ( + layout.ranks, + 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, + 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, + 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, + 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, + 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, + 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) + + @staticmethod + def _mask_inactive_decode_slots(mapping: torch.Tensor, actual_batch: int) -> torch.Tensor: + """Prevent fixed-batch padding rows from writing into the active request cache.""" + if mapping.ndim != 2: + raise ValueError(f"decode slot mapping must be rank-2, got shape={tuple(mapping.shape)}") + if actual_batch <= 0 or actual_batch > mapping.shape[0]: + raise ValueError("actual_batch must be within the decode mapping row count") + masked = mapping.clone() + if actual_batch < masked.shape[0]: + masked[actual_batch:].fill_(-1) + return masked + + 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..e283f8a --- /dev/null +++ b/examples/model/deepseek_v4/runner/weight_loader.py @@ -0,0 +1,835 @@ +# 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) + + +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 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/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..d01235c 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__( @@ -94,6 +132,8 @@ def __init__( block_size=block_size, enable_prefix_cache=self.config.enable_prefix_cache, ) + if runtime.kv_cache_groups: + self.kv_cache_manager.init_groups(runtime.kv_cache_groups) scheduler_config = SchedulerConfig( max_num_running_reqs=self.config.max_num_running_reqs, @@ -117,21 +157,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 +193,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 +247,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 +305,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 +375,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/kv_cache.py b/python/core/kv_cache.py index 9d0096e..0885c13 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -14,7 +14,7 @@ import torch -from .types import KvAllocation, ModelConfig, RuntimeConfig +from .types import KvAllocation, KVCacheGroupSpec, ModelConfig, RuntimeConfig NONE_HASH = hash(("__none__",)) @@ -114,6 +114,25 @@ class _CachePool: value_pages: torch.Tensor +@dataclass +class _GroupBlockPool: + """Per-group block pool with independent free queue and capacity.""" + + name: str + spec: KVCacheGroupSpec + blocks: list[KVCacheBlock] + free_queue: FreeKVCacheBlockQueue + request_blocks: dict[str, list[KVCacheBlock]] + + @property + def num_free_blocks(self) -> int: + return self.free_queue.count + + @property + def num_blocks(self) -> int: + return len(self.blocks) + + class KvCacheManager: """Unified KV block metadata and paged KV tensor storage manager.""" @@ -132,6 +151,7 @@ def __init__( self.free_queue = FreeKVCacheBlockQueue() self.hash_to_block: dict[int, KVCacheBlock] = {} self.request_blocks: dict[str, list[KVCacheBlock]] = {} + self._group_pools: dict[str, _GroupBlockPool] = {} if num_blocks is not None: self._init_blocks(num_blocks, block_size) @@ -408,6 +428,152 @@ def free(self, alloc: KvAllocation) -> None: alloc.tokens_capacity = 0 alloc.tokens_used = 0 + # ------------------------------------------------------------------ + # Multi-group KV cache support (for models like DeepSeekV4 with + # multiple cache families: ori, cmp, idx, states, etc.) + # ------------------------------------------------------------------ + + @property + def has_groups(self) -> bool: + """True when multi-group cache pools have been initialized.""" + return len(self._group_pools) > 0 + + @property + def group_names(self) -> list[str]: + """Return registered group names.""" + return list(self._group_pools.keys()) + + def init_groups(self, group_specs: tuple[KVCacheGroupSpec, ...]) -> None: + """Initialize per-group block pools from group specifications. + + Each group gets its own independent block pool. Groups are used by + models like DeepSeekV4 that have multiple cache families with + different block sizes and compression ratios. + """ + if self._group_pools: + existing = set(self._group_pools.keys()) + incoming = {gs.name for gs in group_specs} + if existing != incoming: + raise ValueError( + f"Group pools already initialized with {existing}, got {incoming}" + ) + return + for gs in group_specs: + num_blocks = gs.max_blocks_per_seq * ( + self._max_batch_size_from_default_pool() or 32 + ) + blocks = [KVCacheBlock(block_id=i) for i in range(num_blocks)] + free_queue = FreeKVCacheBlockQueue() + for block in blocks: + free_queue.append(block) + self._group_pools[gs.name] = _GroupBlockPool( + name=gs.name, + spec=gs, + blocks=blocks, + free_queue=free_queue, + request_blocks={}, + ) + + def _max_batch_size_from_default_pool(self) -> int | None: + """Infer max batch from the default pool if available.""" + if self.blocks: + return max(1, len(self.blocks) // max(1, self.block_size)) + return None + + def allocate_group_blocks( + self, group_name: str, num_blocks: int + ) -> list[int] | None: + """Allocate blocks from a named group pool. Returns block IDs or None.""" + pool = self._group_pool(group_name) + if num_blocks <= 0: + return [] + if pool.num_free_blocks < num_blocks: + return None + allocated: list[int] = [] + for _ in range(num_blocks): + block = pool.free_queue.popleft() + if block is None: + for bid in allocated: + self._release_group_block(pool, pool.blocks[bid]) + return None + block.ref_cnt = 1 + allocated.append(block.block_id) + return allocated + + def allocate_for_groups( + self, + request_id: str, + token_count: int, + ) -> dict[str, list[int]]: + """Allocate blocks across all groups for one request. + + Returns a dict mapping group_name -> list of block IDs. + Raises RuntimeError if any group cannot satisfy the allocation. + """ + result: dict[str, list[int]] = {} + allocated_groups: list[str] = [] + try: + for name, pool in self._group_pools.items(): + spec = pool.spec + num_blocks_needed = math.ceil( + token_count / (spec.spec.block_size // max(1, spec.spec.compress_ratio)) + ) + num_blocks_needed = min(num_blocks_needed, spec.max_blocks_per_seq) + block_ids = self.allocate_group_blocks(name, num_blocks_needed) + if block_ids is None: + raise RuntimeError( + f"Insufficient blocks in group {name!r}: " + f"need {num_blocks_needed}, have {pool.num_free_blocks}" + ) + result[name] = block_ids + pool.request_blocks.setdefault(request_id, []).extend( + pool.blocks[bid] for bid in block_ids + ) + allocated_groups.append(name) + except RuntimeError: + for gname in allocated_groups: + self.release_group_request(gname, request_id) + raise + return result + + def release_group_request(self, group_name: str, request_id: str) -> None: + """Release all blocks held by a request in one group.""" + pool = self._group_pool(group_name) + blocks = pool.request_blocks.pop(request_id, []) + for block in blocks: + self._release_group_block(pool, block) + + def release_all_group_requests(self, request_id: str) -> None: + """Release blocks held by a request across all groups.""" + for pool in self._group_pools.values(): + blocks = pool.request_blocks.pop(request_id, []) + for block in blocks: + self._release_group_block(pool, block) + + def release_group_blocks_by_ids( + self, group_name: str, block_ids: list[int] + ) -> None: + """Release specific block IDs in a named group.""" + pool = self._group_pool(group_name) + for bid in block_ids: + self._release_group_block(pool, pool.blocks[bid]) + + def group_free_count(self, group_name: str) -> int: + """Return free block count for a named group.""" + return self._group_pool(group_name).num_free_blocks + + def _release_group_block(self, pool: _GroupBlockPool, block: KVCacheBlock) -> None: + if block.ref_cnt <= 0: + return + block.ref_cnt -= 1 + if block.ref_cnt == 0: + pool.free_queue.append(block) + + def _group_pool(self, group_name: str) -> _GroupBlockPool: + if group_name not in self._group_pools: + raise KeyError(f"Group {group_name!r} is not registered") + return self._group_pools[group_name] + def _pool(self, model_id: str) -> _CachePool: """Return the registered cache pool for a model.""" if model_id not in self._pools: 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/scheduler.py b/python/core/scheduler.py index 6f9305c..7511f9d 100644 --- a/python/core/scheduler.py +++ b/python/core/scheduler.py @@ -64,6 +64,7 @@ class Request: cached_block_ids: list[int] = field(default_factory=list) allocated_block_ids: list[int] = field(default_factory=list) block_hashes: list[int] = field(default_factory=list) + allocated_group_block_ids: dict[str, list[int]] = field(default_factory=dict) @property def num_prompt_tokens(self) -> int: @@ -93,6 +94,7 @@ class ScheduledRequest: is_prefill: bool num_computed_tokens: int = 0 block_ids: list[int] = field(default_factory=list) + block_ids_by_group: dict[str, list[int]] = field(default_factory=dict) resumed_from_preemption: bool = False @@ -195,6 +197,7 @@ def schedule(self) -> SchedulerOutput: is_prefill = request.is_prefill all_block_ids = request.cached_block_ids + request.allocated_block_ids + self._try_allocate_group_blocks(request, request.num_computed_tokens + num_new) output.scheduled_requests.append( ScheduledRequest( request=request, @@ -202,6 +205,7 @@ def schedule(self) -> SchedulerOutput: is_prefill=is_prefill, num_computed_tokens=request.num_computed_tokens, block_ids=list(all_block_ids), + block_ids_by_group=dict(request.allocated_group_block_ids), ) ) scheduled_req_ids.add(request.request_id) @@ -252,6 +256,7 @@ def schedule(self) -> SchedulerOutput: request.status = RequestStatus.RUNNING self.running.append(request) all_block_ids = request.cached_block_ids + request.allocated_block_ids + self._try_allocate_group_blocks(request, request.num_computed_tokens + num_new) output.scheduled_requests.append( ScheduledRequest( request=request, @@ -259,6 +264,7 @@ def schedule(self) -> SchedulerOutput: is_prefill=True, num_computed_tokens=request.num_computed_tokens, block_ids=list(all_block_ids), + block_ids_by_group=dict(request.allocated_group_block_ids), ) ) output.num_prefill_tokens += num_new @@ -386,6 +392,7 @@ def _preempt_lowest_priority( victim.num_computed_tokens = 0 victim.cached_block_ids = [] victim.allocated_block_ids = [] + victim.allocated_group_block_ids = {} self.running = [r for r in self.running if r.request_id != victim.request_id] self.waiting.appendleft(victim) return {"request": victim, "returned_tokens": returned_tokens} @@ -397,6 +404,22 @@ def _free_request_blocks(self, request: Request) -> None: ) request.cached_block_ids = [] request.allocated_block_ids = [] + if request.allocated_group_block_ids: + self.kv_cache_manager.release_all_group_requests(request.request_id) + request.allocated_group_block_ids = {} + + def _try_allocate_group_blocks(self, request: Request, total_tokens: int) -> None: + """Allocate multi-group blocks if groups are configured.""" + if not self.kv_cache_manager.has_groups: + return + if request.allocated_group_block_ids: + return + try: + request.allocated_group_block_ids = self.kv_cache_manager.allocate_for_groups( + request.request_id, total_tokens + ) + except RuntimeError: + request.allocated_group_block_ids = {} def _cache_completed_blocks(self, request: Request) -> None: """Register completed blocks in the prefix cache.""" 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..b4b325d 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) @@ -211,6 +224,7 @@ def _batch_prefill( list(positions_list[i]), dtype=torch.long, device=device ) + group_block_ids_list = [sr.block_ids_by_group for sr in scheduled] prefill_result = self.executor.run_prefill( runtime_model, PrefillBatch( @@ -220,6 +234,7 @@ def _batch_prefill( seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), positions=positions_tensor, block_ids=block_ids_list, + block_ids_by_group=group_block_ids_list, ), ) @@ -271,6 +286,7 @@ def _batch_decode( decode_token_tensor = torch.tensor(decode_tokens, dtype=torch.long, device=device) decode_embeddings = self.executor.lookup_embeddings(runtime_model, decode_token_tensor) + group_block_ids_list = [sr.block_ids_by_group for sr in scheduled] decode_result = self.executor.run_decode( runtime_model, DecodeBatch( @@ -279,6 +295,7 @@ def _batch_decode( hidden_states=decode_embeddings, seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), block_ids=block_ids_list, + block_ids_by_group=group_block_ids_list, ), ) 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..39dcc18 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -51,6 +51,29 @@ class ModelConfig: torch_dtype: str +@dataclass(frozen=True) +class KVCacheSpec: + """Describes one KV cache family's shape and memory budget per block.""" + + block_size: int + page_size_bytes: int + compress_ratio: int = 1 + + @property + def storage_block_size(self) -> int: + return self.block_size // self.compress_ratio + + +@dataclass(frozen=True) +class KVCacheGroupSpec: + """Maps a named cache group to a KVCacheSpec and the layers that use it.""" + + name: str + layer_indices: tuple[int, ...] + spec: KVCacheSpec + max_blocks_per_seq: int + + @dataclass(frozen=True) class RuntimeConfig: """Runtime limits and device placement for one loaded model.""" @@ -64,6 +87,7 @@ class RuntimeConfig: total_kv_pages: int | None = None # Compile-time generation limit used by model-specific runners. max_new_tokens: int = 256 + kv_cache_groups: tuple[KVCacheGroupSpec, ...] = () @dataclass(frozen=True) @@ -105,6 +129,7 @@ class RuntimeModel: final_norm_weight: torch.Tensor lm_head: torch.Tensor layers: list[LayerWeights] + extra: dict[str, object] = field(default_factory=dict) @dataclass @@ -181,6 +206,7 @@ class PrefillBatch: kv_allocations: list[KvAllocation] = field(default_factory=list) positions: torch.Tensor | None = None block_ids: list[list[int]] = field(default_factory=list) + block_ids_by_group: list[dict[str, list[int]]] = field(default_factory=list) @dataclass @@ -201,6 +227,7 @@ class DecodeBatch: seq_lens: torch.Tensor kv_allocations: list[KvAllocation] = field(default_factory=list) block_ids: list[list[int]] = field(default_factory=list) + block_ids_by_group: list[dict[str, list[int]]] = field(default_factory=list) @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..860a8b0 --- /dev/null +++ b/tests/test_deepseek_v4.py @@ -0,0 +1,1067 @@ +# 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, + 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 _DecodeModule: + l3_decode_layer = 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, + "decode_layer": _DecodeModule, + "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"] + assert len(compiled_args["deepseek_v4_prefill"]) == 78 + assert len(compiled_args["deepseek_v4_decode"]) == 77 + assert len(compiled_args["deepseek_v4_lm_head"]) == 3 + assert isinstance(compiled_args["deepseek_v4_prefill"][-2], ctypes.c_int32) + assert isinstance(compiled_args["deepseek_v4_prefill"][-1], ctypes.c_int32) + assert isinstance(compiled_args["deepseek_v4_decode"][-2], 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, 64, 4, 4096) + decode_order = npu_executor._DECODE_TENSOR_ORDER + assert compiled_args["deepseek_v4_decode"][decode_order.index("hca_compress_state")].shape == (8, 2048, 8, 1024) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_compress_state")].shape == (8, 2080, 4, 2048) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_inner_compress_state")].shape == ( + 8, + 2080, + 4, + 512, + ) + 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(): + builder = DeepSeekV4InputBuilder(layout=DeepSeekV4CacheLayout(), 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[:, 6:], decode[:, 0:1].expand_as(decode[:, 6:])) + + +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_masks_inactive_write_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.x_hc.shape == (8, 64, 4, 4) + assert prepared.x_hc[0, 0, 0].tolist() == [0, 1, 2, 3] + assert prepared.x_hc[0, 2, 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, 64) + assert prepared.ori_slot_mapping[0, :6].tolist() == [127, 0, 132, 133, -1, -1] + assert prepared.hca_cmp_slot_mapping[0, :6].tolist() == [0, -1, -1, -1, -1, -1] + assert prepared.csa_cmp_slot_mapping[0, :6].tolist() == [31, -1, -1, -1, -1, -1] + assert prepared.csa_idx_slot_mapping[0, :6].tolist() == [31, -1, -1, -1, -1, -1] + assert prepared.csa_state_slot_mapping[0, :6].tolist() == [127, 128, 264, 265, -1, -1] + + +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 _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 / "decode_layer.py").write_text("") + (kernel_dir / "config.py").write_text( + "\n".join( + [ + f"BLOCK_SIZE = {block_size}", + "DECODE_BATCH = 32", + "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( + layout=DeepSeekV4CacheLayout(), + 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