diff --git a/pypto-lib b/pypto-lib index ac66502f..43553649 160000 --- a/pypto-lib +++ b/pypto-lib @@ -1 +1 @@ -Subproject commit ac66502fd3bec7960d178ad537838a344511b3ec +Subproject commit 435536490d9956d37acb9f124869fc821d6b4927 diff --git a/pypto_serving/model/deepseek/npu_executor.py b/pypto_serving/model/deepseek/npu_executor.py index 0d2289fa..f52552b5 100644 --- a/pypto_serving/model/deepseek/npu_executor.py +++ b/pypto_serving/model/deepseek/npu_executor.py @@ -42,7 +42,9 @@ DeepSeekV4ModelRunner, _DECODE_FWD_TENSOR_ORDER, _PREFILL_FWD_TENSOR_ORDER, + build_deepseek_v4_cache_group_specs, build_deepseek_v4_layer_plan, + deepseek_v4_physical_cache_blocks, DEEPSEEK_V4_CSA_NUM_LAYERS, DEEPSEEK_V4_FWD_NUM_LAYERS, DEEPSEEK_V4_HCA_NUM_LAYERS, @@ -436,6 +438,8 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels: modules["prefill_mtp"].l3_mtp_prefill_fwd, self._mtp_dummy_args( modules["prefill_mtp"], + model=model, + layout=layout, num_tokens=layout.prefill_seq, ), ) @@ -444,6 +448,8 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels: modules["decode_mtp"].l3_mtp_decode_layer, self._mtp_dummy_args( modules["decode_mtp"], + model=model, + layout=layout, num_tokens=layout.decode_tokens, ), ) @@ -457,6 +463,7 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels: compress_ratios=compress_ratios, layer_plan=layer_plan, kernel_dir=str(self._kernel_dir), + runtime_model=model, prefill=prefill, decode=decode, mtp_prefill=mtp_prefill, @@ -465,6 +472,7 @@ def _compile_model(self, model: RuntimeModel) -> DeepSeekV4CompiledKernels: freqs_sin=freqs_sin, platform=self._platform, device_id=self._device_ids[0], + device_ids=self._device_ids, n_routed_experts=n_routed_experts, num_hash_layers=num_hash_layers, enable_mtp=self._enable_mtp, @@ -519,12 +527,15 @@ def _mtp_dummy_args( self, mtp_module: object, *, + model: RuntimeModel, + layout: DeepSeekV4CacheLayout, num_tokens: int, ) -> tuple[Any, ...]: """Build shape-only dummy tensors in an MTP module's tensor-spec order.""" args: list[Any] = [] pypto_root = self._kernel_dir.parents[2] is_prefill = getattr(mtp_module, "__name__", "") == "prefill_mtp" + cache_blocks = self._compile_cache_blocks(model, layout) with _deepseek_v4_import_context( self._kernel_dir, pypto_root=pypto_root, @@ -537,7 +548,16 @@ def _mtp_dummy_args( if spec.name == "num_tokens": args.append(self._int32_arg(num_tokens)) else: - args.append(torch.empty(tuple(spec.shape), dtype=spec.dtype)) + shape = tuple(spec.shape) + if spec.name == "kv_cache": + shape = ( + layout.ranks, + cache_blocks["ori"], + layout.block_size, + 1, + model.config.head_dim, + ) + args.append(torch.empty(shape, dtype=spec.dtype)) return tuple(args) def _compile_l3_callable(self, name: str, jit_fn: object, dummy_args: Sequence[Any]) -> DeepSeekV4L3Callable: @@ -572,6 +592,33 @@ def _compile_l3_callable(self, name: str, jit_fn: object, dummy_args: Sequence[A raise TypeError(f"{name} did not compile to DistributedCompiledProgram; got {type(compiled).__name__}") return DeepSeekV4L3Callable(compiled=compiled, name=name) + @staticmethod + def _compile_cache_blocks( + model: RuntimeModel, + layout: DeepSeekV4CacheLayout, + ) -> dict[str, int]: + """Return worst-case runtime pool shapes used for JIT dummy arguments. + + Actual worker allocations are HBM-sized after compilation. Like Qwen's + paged kernels, DeepSeek kernels must derive their layer/page stride from + the runtime tensor shape rather than baking these dummy dimensions into + generated code. + """ + specs = model.runtime.kv_cache_groups or build_deepseek_v4_cache_group_specs( + model.config.num_hidden_layers, + model.extra.get("compress_ratios"), + decode_batch=layout.decode_batch, + ) + local_capacity_slots = max( + 1, + (model.runtime.max_batch_size + layout.ranks - 1) // layout.ranks, + ) + return deepseek_v4_physical_cache_blocks( + specs, + local_capacity_slots, + scratch_blocks=layout.decode_batch, + ) + def _prefill_dummy_args( self, model: RuntimeModel, @@ -610,6 +657,7 @@ def _prefill_dummy_args( fwd = DEEPSEEK_V4_FWD_NUM_LAYERS csa = DEEPSEEK_V4_CSA_NUM_LAYERS hca = DEEPSEEK_V4_HCA_NUM_LAYERS + cache_blocks = self._compile_cache_blocks(model, layout) def stacked(name: str, count: int) -> torch.Tensor: base = single[name] @@ -638,7 +686,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "hca_compress_state": torch.empty( ( ranks, - hca * layout.hca_state_max_blocks, + hca * cache_blocks["hca_state"], layout.c128_state_block_size, DEEPSEEK_V4_HCA_STATE_DIM, ), @@ -652,7 +700,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "csa_compress_state": torch.empty( ( ranks, - csa * layout.csa_state_max_blocks, + csa * cache_blocks["csa_state"], layout.c4_state_block_size, DEEPSEEK_V4_CSA_STATE_DIM, ), @@ -665,7 +713,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "csa_inner_compress_state": torch.empty( ( ranks, - csa * layout.csa_inner_state_max_blocks, + csa * cache_blocks["csa_inner_state"], layout.c4_state_block_size, DEEPSEEK_V4_CSA_INNER_STATE_DIM, ), @@ -679,19 +727,19 @@ def stacked(name: str, count: int) -> torch.Tensor: # stacks across the 21 CSA layers. The kernel reshapes the fused # layer x block axis internally. "kv_cache": torch.empty( - (ranks, fwd * layout.prefill_ori_max_blocks, layout.block_size, 1, head_dim), + (ranks, fwd * cache_blocks["ori"], layout.block_size, 1, head_dim), dtype=torch.bfloat16, ), "cmp_kv": torch.empty( - (ranks, fwd * layout.prefill_cmp_block_num, layout.block_size, 1, head_dim), + (ranks, fwd * cache_blocks["cmp"], layout.block_size, 1, head_dim), dtype=torch.bfloat16, ), "idx_kv_cache": torch.empty( - (ranks, csa * layout.prefill_idx_block_num, layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), + (ranks, csa * cache_blocks["idx"], layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM), dtype=torch.int8, ), "idx_kv_scale": torch.empty( - (ranks, csa * layout.prefill_idx_block_num, layout.block_size, 1, 1), + (ranks, csa * cache_blocks["idx"], layout.block_size, 1, 1), dtype=torch.float32, ), # Shared single per-rank prefill metadata (the kernel passes each @@ -768,6 +816,7 @@ def _decode_dummy_args( fwd = DEEPSEEK_V4_FWD_NUM_LAYERS csa = DEEPSEEK_V4_CSA_NUM_LAYERS hca = DEEPSEEK_V4_HCA_NUM_LAYERS + cache_blocks = self._compile_cache_blocks(model, layout) def stacked(name: str, count: int) -> torch.Tensor: base = single[name] @@ -795,7 +844,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "kv_cache": torch.empty( ( ranks, - fwd * layout.decode_ori_max_blocks, + fwd * cache_blocks["ori"], layout.block_size, 1, model.config.head_dim, @@ -805,7 +854,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "cmp_kv": torch.empty( ( ranks, - fwd * layout.cmp_max_blocks, + fwd * cache_blocks["cmp"], layout.block_size, 1, model.config.head_dim, @@ -816,7 +865,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "idx_kv_cache": torch.empty( ( ranks, - csa * layout.idx_max_blocks, + csa * cache_blocks["idx"], layout.block_size, 1, DEEPSEEK_V4_IDX_HEAD_DIM, @@ -826,7 +875,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "idx_kv_scale": torch.empty( ( ranks, - csa * layout.idx_max_blocks, + csa * cache_blocks["idx"], layout.block_size, 1, 1, @@ -836,7 +885,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "csa_compress_state": torch.empty( ( ranks, - csa * layout.csa_state_max_blocks, + csa * cache_blocks["csa_state"], layout.c4_state_block_size, DEEPSEEK_V4_CSA_STATE_DIM, ), @@ -845,7 +894,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "csa_inner_compress_state": torch.empty( ( ranks, - csa * layout.csa_inner_state_max_blocks, + csa * cache_blocks["csa_inner_state"], layout.c4_state_block_size, DEEPSEEK_V4_CSA_INNER_STATE_DIM, ), @@ -855,7 +904,7 @@ def stacked(name: str, count: int) -> torch.Tensor: "hca_compress_state": torch.empty( ( ranks, - hca * layout.hca_state_max_blocks, + hca * cache_blocks["hca_state"], layout.c128_state_block_size, DEEPSEEK_V4_HCA_STATE_DIM, ), @@ -1064,10 +1113,6 @@ def _validate_kernel_contract(self, layout: DeepSeekV4CacheLayout) -> None: "KV_ORI_TABLE_MAX_BLOCKS": layout.ori_table_max_blocks, "KV_CMP_MAX_BLOCKS": layout.cmp_max_blocks, "IDX_CACHE_MAX_BLOCKS": layout.idx_max_blocks, - "ORI_KV_BLOCK_NUM": layout.decode_ori_max_blocks, - "HCA_STATE_PHYSICAL_BLOCKS": layout.hca_state_max_blocks, - "CSA_STATE_PHYSICAL_BLOCKS": layout.csa_state_max_blocks, - "CSA_INNER_STATE_PHYSICAL_BLOCKS": layout.csa_inner_state_max_blocks, "PREFILL_ORI_MAX_BLOCKS": layout.prefill_ori_max_blocks, "PREFILL_CMP_MAX_BLOCKS": layout.prefill_cmp_max_blocks, "PREFILL_IDX_MAX_BLOCKS": layout.prefill_idx_max_blocks, @@ -1080,13 +1125,10 @@ def _validate_kernel_contract(self, layout: DeepSeekV4CacheLayout) -> None: mismatched.append(f"{name}={actual} expected {expected}") expected_module_constants = { "prefill_attention_hca.py": { - "HCA_STATE_BLOCK_NUM": layout.hca_state_max_blocks, "HCA_STATE_MAX_BLOCKS": layout.prefill_hca_state_max_blocks, }, "prefill_attention_csa.py": { - "CSA_STATE_BLOCK_NUM": layout.csa_state_max_blocks, "CSA_STATE_MAX_BLOCKS": layout.prefill_csa_state_max_blocks, - "INNER_STATE_BLOCK_NUM": layout.csa_inner_state_max_blocks, "INNER_STATE_MAX_BLOCKS": layout.prefill_csa_inner_state_max_blocks, }, } diff --git a/pypto_serving/model/deepseek/npu_runner.py b/pypto_serving/model/deepseek/npu_runner.py index 4262e79c..102aa7b3 100644 --- a/pypto_serving/model/deepseek/npu_runner.py +++ b/pypto_serving/model/deepseek/npu_runner.py @@ -154,7 +154,14 @@ def build_deepseek_v4_cache_group_specs( *, decode_batch: int = DEEPSEEK_V4_DECODE_TOKENS, ) -> tuple[KVCacheGroupSpec, ...]: - """Describe cache namespaces and reserve padding pages for the decode specialization.""" + """Describe cache namespaces independently from their runtime pool sizes. + + ``max_blocks_per_seq`` is a per-request ring/table limit. The number of + physical blocks in each rank-local pool is intentionally left unset: the + runner fills available NPU memory after weights and persistent scratch have + been materialized, then the engine scales every group to the reported + runtime capacity. + """ decode_batch = int(decode_batch) if decode_batch <= 0: raise ValueError("decode_batch must be positive") @@ -172,6 +179,7 @@ def group( row_width: int, max_blocks: int, compress_ratio: int = 1, + extra_row_bytes: int = 0, ) -> KVCacheGroupSpec: if max_blocks <= decode_batch: raise ValueError( @@ -183,21 +191,19 @@ def group( layer_indices=layers, spec=KVCacheSpec( block_size=block_size, - page_size_bytes=block_size * row_width * element_bytes, + # One scheduler-visible block owns a page in every layer of + # this cache family. Keep the complete physical footprint here + # so all heterogeneous pools share one accurate HBM budget. + page_size_bytes=( + len(layers) + * block_size + * (row_width * element_bytes + extra_row_bytes) + ), compress_ratio=compress_ratio, ), - # One rank-local decode row owns one disjoint ring slice. This is - # the same fixed-pool partitioning used by pypto-lib's B4 host - # metadata, while allocation/release remains scheduler-owned. + # Each request may address this many pages through its logical ring; + # physical capacity across concurrent requests is runtime-sized. max_blocks_per_seq=blocks_per_request, - # pypto-lib exposes one fixed physical pool per DP rank. The - # scheduler shares that pool across the rank-local decode rows and - # applies backpressure when their combined working sets do not fit. - # Attention/compressor kernels execute every fixed B row even when - # MoE receives a smaller num_tokens prefix. Reserve one physical - # scratch page per kernel row so padding can never alias a live - # scheduler-owned page. - num_blocks=max_blocks - decode_batch, num_partitions=DEEPSEEK_V4_RANKS, ) @@ -224,10 +230,12 @@ def group( "idx", csa_layers, block_size=DEEPSEEK_V4_BLOCK_SIZE, - element_bytes=2, + element_bytes=1, row_width=DEEPSEEK_V4_IDX_HEAD_DIM, max_blocks=DEEPSEEK_V4_IDX_MAX_BLOCKS, compress_ratio=4, + # One float32 scale accompanies every int8 index-cache row. + extra_row_bytes=4, ), group( "hca_state", @@ -256,6 +264,52 @@ def group( ) +DEEPSEEK_V4_CACHE_GROUP_NAMES = ( + "ori", + "cmp", + "idx", + "hca_state", + "csa_state", + "csa_inner_state", +) + + +def deepseek_v4_cache_blocks_for_slots( + group_specs: Sequence[KVCacheGroupSpec], + capacity_slots: int, +) -> dict[str, int]: + """Return scheduler-visible blocks per rank for ``capacity_slots`` requests.""" + capacity_slots = int(capacity_slots) + if capacity_slots <= 0: + raise ValueError("DeepSeekV4 cache capacity_slots must be positive") + specs = {spec.name: spec for spec in group_specs} + missing = [name for name in DEEPSEEK_V4_CACHE_GROUP_NAMES if name not in specs] + if missing: + raise ValueError("missing DeepSeekV4 cache groups: " + ", ".join(missing)) + return { + name: capacity_slots * specs[name].max_blocks_per_seq + for name in DEEPSEEK_V4_CACHE_GROUP_NAMES + } + + +def deepseek_v4_physical_cache_blocks( + group_specs: Sequence[KVCacheGroupSpec], + capacity_slots: int, + *, + scratch_blocks: int, +) -> dict[str, int]: + """Return physical blocks per rank, including isolated padding pages.""" + if scratch_blocks <= 0: + raise ValueError("DeepSeekV4 scratch_blocks must be positive") + return { + name: num_blocks + int(scratch_blocks) + for name, num_blocks in deepseek_v4_cache_blocks_for_slots( + group_specs, + capacity_slots, + ).items() + } + + # Argument order for the packed all-43-layer ``l3_prefill_fwd`` kernel. This # mirrors pypto-lib prefill_fwd.py ``l3_prefill_fwd`` host signature: every # layer-stacked weight/state tensor in core-parameter order, followed by the @@ -510,7 +564,12 @@ def group( @dataclass(frozen=True) class DeepSeekV4CacheLayout: - """Static cache layout baked into the current DeepSeekV4 kernels.""" + """Kernel table depths and fixed execution dimensions. + + Physical cache-pool sizes are runtime state on ``DeepSeekV4ModelRunner``; + the ``*_max_blocks`` fields here only bound per-request metadata tables and + provide compatibility defaults for shape-only callers. + """ ranks: int = DEEPSEEK_V4_RANKS hc_mult: int = DEEPSEEK_V4_HC_MULT @@ -537,16 +596,6 @@ class DeepSeekV4CacheLayout: 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 - @property - def prefill_cmp_block_num(self) -> int: - """Physical cmp_kv blocks per layer in the packed prefill kernel.""" - return self.prefill_cmp_max_blocks - - @property - def prefill_idx_block_num(self) -> int: - """Physical idx_kv_cache blocks per CSA layer in the packed prefill kernel.""" - return self.prefill_idx_max_blocks - def validate_runtime(self, config: ModelConfig, runtime: RuntimeConfig, device_ids: Sequence[int]) -> None: """Validate serving/runtime options against kernel-fixed dimensions.""" if len(device_ids) != self.ranks: @@ -975,19 +1024,6 @@ class _TransientDeviceTensor: tensor: torch.Tensor -@dataclass -class DeepSeekV4LayerCache: - """Shared host backing for the rank-sharded DeepSeekV4 cache pools.""" - - kv_cache: torch.Tensor - cmp_kv: torch.Tensor - idx_kv_cache: torch.Tensor - idx_kv_scale: torch.Tensor - hca_compress_state: torch.Tensor - csa_compress_state: torch.Tensor - csa_inner_compress_state: torch.Tensor - - @dataclass class DeepSeekV4DeviceCache: """Worker-resident rank shards shared by packed prefill and decode.""" @@ -1012,6 +1048,7 @@ class DeepSeekV4CompiledKernels: compress_ratios: tuple[int, ...] layer_plan: tuple["DeepSeekV4LayerPlan", ...] kernel_dir: str + runtime_model: RuntimeModel | None = None prefill: DeepSeekV4L3Callable | None = None decode: DeepSeekV4L3Callable | None = None mtp_prefill: DeepSeekV4L3Callable | None = None @@ -1020,6 +1057,7 @@ class DeepSeekV4CompiledKernels: freqs_sin: torch.Tensor | None = None platform: str = "a2a3" device_id: int = 0 + device_ids: tuple[int, ...] = () n_routed_experts: int = 256 num_hash_layers: int = 3 embedding_weight: torch.Tensor | None = None @@ -1291,7 +1329,9 @@ def __init__(self, *, compiled: DeepSeekV4CompiledKernels) -> None: tuple[int, tuple[int, ...], torch.dtype], StackedDeviceTensor ] = {} self._l3_cache_tensor_keys: set[tuple[int, tuple[int, ...], torch.dtype]] = set() - self._decode_work_cache: DeepSeekV4LayerCache | None = None + self._cache_group_specs: tuple[KVCacheGroupSpec, ...] = () + self._cache_group_num_blocks: dict[str, int] = {} + self._cache_zero_page: torch.Tensor | None = None self._decode_device_cache: DeepSeekV4DeviceCache | None = None self._decode_cache_block_ids: dict[str, dict[str, set[int]]] = {} self._global_weights: DeepSeekV4GlobalWeights | None = None @@ -1322,21 +1362,198 @@ def __init__(self, *, compiled: DeepSeekV4CompiledKernels) -> None: self._prefill_completion = self._ignore_prefill_context def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> int: - """Initialize runner state and return scheduler-only KV block capacity. + """Allocate heterogeneous cache pools from the post-weight HBM budget. - DeepSeekV4 owns its model-specific cache tensors, while the scheduler's - grouped ``KvCacheManager`` is the sole owner of request partitions and - block IDs. No generic KV tensors are allocated here. + The returned value is the scheduler-visible ``ori`` block count in one + rank-local partition. Other group capacities are derived from the same + number of maximum-sequence cache slots by ``KvCacheManager``. """ self.input_builder = DeepSeekV4InputBuilder( layout=self._compiled.layout, hidden_size=config.hidden_size, ) self._decode_cache_block_ids.clear() + self._cache_group_specs = self._resolve_cache_group_specs(config, runtime) + + # Shape/metadata-only unit tests construct a runner without compiled + # callables. Preserve a deterministic fixed fallback for that path; the + # production executor always attaches the RuntimeModel and callables. + model = self._compiled.runtime_model + if model is None or not self._compiled.l3_callables(): + self._cache_group_num_blocks = self._fallback_cache_group_num_blocks() + return self._cache_group_num_blocks["ori"] + + logger.info("[init_kv_cache] preparing DeepSeekV4 worker and resident weights …") + self._ensure_l3_shared_buffers(model) + + requested_slots = self._compute_kv_cache_capacity_slots(runtime) + allocated_slots = self._alloc_kv_cache_with_retry(requested_slots) + self._log_kv_cache_allocation(runtime, requested_slots, allocated_slots) + return self._cache_group_num_blocks["ori"] + + def _resolve_cache_group_specs( + self, + config: ModelConfig, + runtime: RuntimeConfig, + ) -> tuple[KVCacheGroupSpec, ...]: + specs = runtime.kv_cache_groups or build_deepseek_v4_cache_group_specs( + config.num_hidden_layers, + self._compiled.compress_ratios, + decode_batch=self._compiled.layout.decode_batch, + ) + names = tuple(spec.name for spec in specs) + if names != DEEPSEEK_V4_CACHE_GROUP_NAMES: + raise ValueError( + "DeepSeekV4 KV cache groups must be ordered as " + + ", ".join(DEEPSEEK_V4_CACHE_GROUP_NAMES) + + f"; got {names}" + ) + if any(spec.num_partitions != self._compiled.layout.ranks for spec in specs): + raise ValueError( + f"DeepSeekV4 KV cache groups must use {self._compiled.layout.ranks} partitions" + ) + return tuple(specs) + + def _fallback_cache_group_num_blocks(self) -> dict[str, int]: + """Return the old fixed capacities for shape-only/non-runtime callers.""" + layout = self._compiled.layout + physical_limits = { + "ori": layout.decode_ori_max_blocks, + "cmp": layout.cmp_max_blocks, + "idx": layout.idx_max_blocks, + "hca_state": layout.hca_state_max_blocks, + "csa_state": layout.csa_state_max_blocks, + "csa_inner_state": layout.csa_inner_state_max_blocks, + } + return { + name: max(1, physical_limits[name] - layout.decode_batch) + for name in DEEPSEEK_V4_CACHE_GROUP_NAMES + } + + def _compute_kv_cache_capacity_slots(self, runtime: RuntimeConfig) -> int: + """Compute common cache slots using Qwen's utilization-budget formula.""" + ori_spec = self._cache_group_specs[0] if runtime.total_kv_pages is not None: - return int(runtime.total_kv_pages) - max_blocks_per_seq = math.ceil(runtime.max_seq_len / runtime.page_size) - return int(runtime.max_batch_size * max_blocks_per_seq) + requested_pages = int(runtime.total_kv_pages) + if requested_pages < ori_spec.max_blocks_per_seq: + raise ValueError( + "DeepSeekV4 total_kv_pages must hold at least one maximum ring: " + f"expected >= {ori_spec.max_blocks_per_seq}, got {requested_pages}" + ) + return requested_pages // ori_spec.max_blocks_per_seq + + device_ids = self._compiled.device_ids or (self._compiled.device_id,) + utilization = float(getattr(runtime, "npu_memory_utilization", 0.90)) + budgets = [] + memory_rows = [] + for device_id in device_ids: + free_bytes, total_bytes = torch.npu.mem_get_info(f"npu:{device_id}") + peak_non_kv = int(total_bytes) - int(free_bytes) + budget = int(int(total_bytes) * utilization - peak_non_kv) + budgets.append(budget) + memory_rows.append((int(device_id), int(free_bytes), int(total_bytes), budget)) + + bytes_per_slot = sum( + spec.max_blocks_per_seq * spec.spec.page_size_bytes + for spec in self._cache_group_specs + ) + scratch_bytes = sum( + self._compiled.layout.decode_batch * spec.spec.page_size_bytes + for spec in self._cache_group_specs + ) + # The optional MTP layer addresses the same ori IDs but owns a separate + # one-layer BF16 pool. + mtp_page_bytes = ( + self._compiled.layout.block_size * DEEPSEEK_V4_HEAD_DIM * torch.bfloat16.itemsize + if self._compiled.enable_mtp + else 0 + ) + bytes_per_slot += ori_spec.max_blocks_per_seq * mtp_page_bytes + scratch_bytes += self._compiled.layout.decode_batch * mtp_page_bytes + + kv_budget = min(budgets) + capacity_slots = max((kv_budget - scratch_bytes) // bytes_per_slot, 1) + logger.info( + "DeepSeekV4 KV cache sizing: utilization=%.2f, limiting_budget=%.2f GB, " + "slot=%.1f MB, scratch=%.1f MB, requested_slots=%d", + utilization, + kv_budget / 1e9, + bytes_per_slot / 1e6, + scratch_bytes / 1e6, + capacity_slots, + ) + for device_id, free_bytes, total_bytes, budget in memory_rows: + logger.info( + " npu:%d total=%.2f GB free=%.2f GB post-weight KV budget=%.2f GB", + device_id, + total_bytes / 1e9, + free_bytes / 1e9, + budget / 1e9, + ) + return int(capacity_slots) + + def _alloc_kv_cache_with_retry(self, requested_slots: int) -> int: + """Allocate every cache family atomically, halving capacity on OOM.""" + capacity_slots = max(int(requested_slots), 1) + while capacity_slots >= 1: + self._cache_group_num_blocks = deepseek_v4_cache_blocks_for_slots( + self._cache_group_specs, + capacity_slots, + ) + try: + self._materialize_decode_device_cache() + self._materialize_mtp_device_kv_cache() + self._zero_scratch_cache_blocks() + return capacity_slots + except (RuntimeError, MemoryError) as exc: + self._free_device_caches() + if capacity_slots == 1: + raise RuntimeError( + "DeepSeekV4 KV cache allocation failed at the one-slot minimum" + ) from exc + previous = capacity_slots + capacity_slots = max(capacity_slots // 2, 1) + logger.warning( + "DeepSeekV4 KV cache allocation failed (%s); retrying slots %d -> %d", + exc, + previous, + capacity_slots, + ) + raise RuntimeError("DeepSeekV4 KV cache allocation failed") + + def _log_kv_cache_allocation( + self, + runtime: RuntimeConfig, + requested_slots: int, + allocated_slots: int, + ) -> None: + main_bytes = sum( + self._physical_cache_num_blocks(spec.name) * spec.spec.page_size_bytes + for spec in self._cache_group_specs + ) + mtp_bytes = 0 + if self._compiled.enable_mtp: + mtp_bytes = ( + self._physical_cache_num_blocks("ori") + * self._compiled.layout.block_size + * DEEPSEEK_V4_HEAD_DIM + * torch.bfloat16.itemsize + ) + logger.info( + "[init_kv_cache] allocated DeepSeekV4 cache: slots=%d (requested=%d), " + "per-rank=%.2f GB, ori_blocks=%d, max_seq_len=%d", + allocated_slots, + requested_slots, + (main_bytes + mtp_bytes) / 1e9, + self._cache_group_num_blocks["ori"], + runtime.max_seq_len, + ) + + def _physical_cache_num_blocks(self, group_name: str) -> int: + try: + return self._cache_group_num_blocks[group_name] + self._compiled.layout.decode_batch + except KeyError as exc: + raise RuntimeError("DeepSeekV4 KV cache capacity is not initialized") from exc def release_finished_requests(self, request_ids: Iterable[str]) -> None: """Discard cache ownership metadata for finished or preempted requests.""" @@ -1733,15 +1950,6 @@ def _prepare_decode_inputs( "csa_inner_state_slot_mapping", ) } - group_limits = { - "ori": layout.decode_ori_max_blocks, - "cmp": layout.cmp_max_blocks, - "idx": layout.idx_max_blocks, - "hca_state": layout.hca_state_max_blocks, - "csa_state": layout.csa_state_max_blocks, - "csa_inner_state": layout.csa_inner_state_max_blocks, - } - for rank, request_indices in enumerate(indices_by_rank): if request_indices: local_positions = [positions[index] for index in request_indices] @@ -1764,16 +1972,16 @@ def _prepare_decode_inputs( padded_positions.append(local_positions[0]) padded_group_ids = {} - for name, max_blocks in group_limits.items(): + for name in DEEPSEEK_V4_CACHE_GROUP_NAMES: if local_groups: padded_group_ids[name] = self._pad_group_block_ids( [groups[name] for groups in local_groups], - max_blocks=max_blocks, + group_name=name, kernel_rows=layout.decode_batch, ) else: padded_group_ids[name] = self._scratch_group_block_ids( - max_blocks=max_blocks, + group_name=name, kernel_rows=layout.decode_batch, ) @@ -1975,14 +2183,14 @@ def _normalize_group_block_ids( ) return tuple(normalized) - @staticmethod def _pad_group_block_ids( + self, active_rows: Sequence[Sequence[int]], *, - max_blocks: int, + group_name: str, kernel_rows: int, ) -> tuple[tuple[int, ...], ...]: - """Pad inactive kernel rows without expanding pypto-lib's fixed pools.""" + """Pad inactive rows with pages at the tail of a dynamic cache pool.""" if not active_rows or len(active_rows) > kernel_rows: raise ValueError("active grouped KV rows must fit the kernel batch") normalized = [tuple(int(block_id) for block_id in row) for row in active_rows] @@ -1991,15 +2199,19 @@ def _pad_group_block_ids( used = [block_id for row in normalized for block_id in row] if len(used) != len(set(used)): raise ValueError("active grouped KV rows must not share physical blocks") - scratch = DeepSeekV4ModelRunner._scratch_group_block_ids( - max_blocks=max_blocks, + scratch = self._scratch_group_block_ids( + group_name=group_name, kernel_rows=kernel_rows, ) - allocator_blocks = max_blocks - kernel_rows + try: + allocator_blocks = self._cache_group_num_blocks[group_name] + except KeyError as exc: + raise ValueError(f"unknown DeepSeekV4 cache group: {group_name}") from exc if any(block_id < 0 or block_id >= allocator_blocks for block_id in used): raise ValueError( f"grouped KV block IDs must be in [0, {allocator_blocks}); " - f"[{allocator_blocks}, {max_blocks}) is reserved for kernel padding" + f"[{allocator_blocks}, {allocator_blocks + kernel_rows}) is reserved " + "for kernel padding" ) padded = list(normalized) # Attention and compressor cache writes are fixed-B and are not fully @@ -2008,16 +2220,21 @@ def _pad_group_block_ids( padded.extend(scratch[: kernel_rows - len(normalized)]) return tuple(padded) - @staticmethod def _scratch_group_block_ids( + self, *, - max_blocks: int, + group_name: str, kernel_rows: int, ) -> tuple[tuple[int, ...], ...]: """Return one isolated physical scratch page for every fixed kernel row.""" - if kernel_rows <= 0 or max_blocks <= kernel_rows: + if kernel_rows <= 0: + raise ValueError("kernel_rows must be positive") + try: + first = self._cache_group_num_blocks[group_name] + except KeyError as exc: + raise ValueError(f"unknown DeepSeekV4 cache group: {group_name}") from exc + if self._physical_cache_num_blocks(group_name) < first + kernel_rows: raise ValueError("cache pool must provide one scratch page per kernel row") - first = max_blocks - kernel_rows return tuple((first + row,) for row in range(kernel_rows)) def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> DeviceTensor: @@ -2299,8 +2516,7 @@ def _ensure_l3_shared_buffers(self, model: RuntimeModel) -> None: self._ensure_decode_buffers(model.config.hidden_size) with profile_span("DeepSeekV4ModelRunner.prepare.allocate_mtp_buffers", cat="executor"): self._ensure_mtp_buffers(model.config.hidden_size) - with profile_span("DeepSeekV4ModelRunner.prepare.allocate_decode_work_cache", cat="executor"): - self._ensure_decode_work_cache() + self._ensure_cache_zero_page() with profile_span("DeepSeekV4ModelRunner.prepare.allocate_prefill_outputs", cat="executor"): self._require_prefill_output_buffer(model.config.hidden_size) self._require_prefill_pre_hc_output_buffer(model.config.hidden_size) @@ -2345,7 +2561,7 @@ def _missing_l3_shared_buffers(self) -> list[str]: "freqs_sin": self._static_freqs_sin, "prefill_fwd_buffers": self._prefill_fwd_buffers, "decode_buffers": self._decode_buffers, - "decode_work_cache": self._decode_work_cache, + "cache_zero_page": self._cache_zero_page, "stacked_weights": self._stacked_host_weights or self._stacked_device_weights, "hc_head_buffers": self._hc_head_buffers, "prefill_output": self._prefill_output_buffer, @@ -2758,7 +2974,7 @@ def _advance_mtp_drafts( ] padded_blocks = self._pad_group_block_ids( active_blocks, - max_blocks=layout.prefill_ori_max_blocks, + group_name="ori", kernel_rows=layout.decode_batch, ) padded_positions = [ @@ -2769,7 +2985,7 @@ def _advance_mtp_drafts( padded_positions.append(tuple(int(value) for value in fallback_positions.tolist())) else: padded_blocks = self._scratch_group_block_ids( - max_blocks=layout.prefill_ori_max_blocks, + group_name="ori", kernel_rows=layout.decode_batch, ) padded_positions = [ @@ -3155,10 +3371,13 @@ def _ensure_mtp_buffers(self, hidden_size: int) -> _DeepSeekV4MtpSharedBuffers | hidden = int(hidden_size) loaded = self.load_mtp_weights() weights = dict(loaded.tensors) + # The real MTP pool is allocated directly on each worker after runtime + # HBM sizing. Keep only a one-page shared placeholder here so legacy + # buffer consumers retain the unified prefill/decode object contract. mtp_kv_cache = self._shared_empty( - (ranks, layout.prefill_ori_max_blocks, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), + (ranks, 1, layout.block_size, 1, DEEPSEEK_V4_HEAD_DIM), torch.bfloat16, - name="mtp_unified_kv_cache", + name="mtp_kv_cache_placeholder", ) self._mtp_buffers = _DeepSeekV4MtpSharedBuffers( weights=weights, @@ -3533,8 +3752,6 @@ def _initialize_decode_cache_blocks( raise ValueError("decode request IDs, ranks, and grouped KV rows must have the same length") if len(set(request_ids)) != len(request_ids): raise ValueError("decode request IDs must be unique") - cache = self._require_decode_work_cache() - dirty: dict[int, dict[str, set[int]]] = {} for request_id, rank, groups in zip( request_ids, ranks, @@ -3550,43 +3767,57 @@ def _initialize_decode_cache_blocks( ) for name, block_ids in groups.items() } - self._zero_decode_work_cache_blocks(cache, int(rank), new_blocks) - rank_dirty = dirty.setdefault(int(rank), {}) - for name, block_ids in new_blocks.items(): - rank_dirty.setdefault(name, set()).update(block_ids) + self._zero_device_cache_blocks(int(rank), new_blocks) for name, block_ids in groups.items(): initialized.setdefault(name, set()).update(int(value) for value in block_ids) - self._sync_decode_work_cache_blocks(cache, dirty) - def _zero_decode_work_cache_blocks( + def _zero_scratch_cache_blocks(self) -> None: + """Clear the tail pages used by inactive fixed kernel rows.""" + layout = self._compiled.layout + for rank in range(layout.ranks): + self._zero_device_cache_blocks( + rank, + { + name: tuple( + block_id + for (block_id,) in self._scratch_group_block_ids( + group_name=name, + kernel_rows=layout.decode_batch, + ) + ) + for name in DEEPSEEK_V4_CACHE_GROUP_NAMES + }, + ) + + def _zero_device_cache_blocks( self, - cache: DeepSeekV4LayerCache, rank: int, block_ids_by_group: dict[str, Sequence[int]], ) -> None: - """Clear newly assigned rank-local pages across every stacked layer.""" + """Clear selected pages in-place without a full host-side cache mirror.""" layout = self._compiled.layout if not 0 <= rank < layout.ranks: raise ValueError(f"DeepSeekV4 cache rank must be in [0, {layout.ranks - 1}]") + cache = self._decode_device_cache + worker = self._l3_worker + zero_page = self._cache_zero_page + if cache is None or worker is None or zero_page is None: + raise RuntimeError("DeepSeekV4 device cache must be materialized before clearing pages") tensors_by_group = { - "ori": ((cache.kv_cache, DEEPSEEK_V4_FWD_NUM_LAYERS, layout.decode_ori_max_blocks),), - "cmp": ((cache.cmp_kv, DEEPSEEK_V4_FWD_NUM_LAYERS, layout.cmp_max_blocks),), + "ori": ((cache.kv_cache, DEEPSEEK_V4_FWD_NUM_LAYERS),), + "cmp": ((cache.cmp_kv, DEEPSEEK_V4_FWD_NUM_LAYERS),), "idx": ( - (cache.idx_kv_cache, DEEPSEEK_V4_CSA_NUM_LAYERS, layout.idx_max_blocks), - (cache.idx_kv_scale, DEEPSEEK_V4_CSA_NUM_LAYERS, layout.idx_max_blocks), + (cache.idx_kv_cache, DEEPSEEK_V4_CSA_NUM_LAYERS), + (cache.idx_kv_scale, DEEPSEEK_V4_CSA_NUM_LAYERS), ), "hca_state": ( - (cache.hca_compress_state, DEEPSEEK_V4_HCA_NUM_LAYERS, layout.hca_state_max_blocks), + (cache.hca_compress_state, DEEPSEEK_V4_HCA_NUM_LAYERS), ), "csa_state": ( - (cache.csa_compress_state, DEEPSEEK_V4_CSA_NUM_LAYERS, layout.csa_state_max_blocks), + (cache.csa_compress_state, DEEPSEEK_V4_CSA_NUM_LAYERS), ), "csa_inner_state": ( - ( - cache.csa_inner_compress_state, - DEEPSEEK_V4_CSA_NUM_LAYERS, - layout.csa_inner_state_max_blocks, - ), + (cache.csa_inner_compress_state, DEEPSEEK_V4_CSA_NUM_LAYERS), ), } unknown = sorted(set(block_ids_by_group) - set(tensors_by_group)) @@ -3596,73 +3827,48 @@ def _zero_decode_work_cache_blocks( ids = tuple(dict.fromkeys(int(block_id) for block_id in block_ids)) if not ids: continue - for tensor, layer_count, blocks_per_layer in tensors_by_group[name]: + blocks_per_layer = self._physical_cache_num_blocks(name) + if any(block_id < 0 or block_id >= blocks_per_layer for block_id in ids): + raise ValueError( + f"DeepSeekV4 {name} block IDs must be in [0, {blocks_per_layer})" + ) + for stacked, layer_count in tensors_by_group[name]: + shard = stacked.shards[rank] + worker_id = stacked.worker_ids[rank] + page_numel = math.prod(shard.shape[1:]) + page_nbytes = page_numel * torch.empty((), dtype=shard.dtype).element_size() + if page_nbytes > zero_page.numel(): + raise RuntimeError( + f"DeepSeekV4 zero page has {zero_page.numel()} bytes, " + f"{name} needs {page_nbytes}" + ) + for layer in range(layer_count): + for block_id in ids: + flat_index = layer * blocks_per_layer + block_id + worker.copy_to( + shard.data_ptr + flat_index * page_nbytes, + zero_page.data_ptr(), + page_nbytes, + worker_id=worker_id, + ) + + if name == "ori" and self._mtp_device_kv_cache is not None: + stacked = self._mtp_device_kv_cache + shard = stacked.shards[rank] + worker_id = stacked.worker_ids[rank] + page_numel = math.prod(shard.shape[1:]) + page_nbytes = page_numel * torch.empty((), dtype=shard.dtype).element_size() if any(block_id < 0 or block_id >= blocks_per_layer for block_id in ids): raise ValueError( f"DeepSeekV4 {name} block IDs must be in [0, {blocks_per_layer})" ) - pages = tensor[rank].reshape( - layer_count, - blocks_per_layer, - *tensor.shape[2:], - ) - pages.index_fill_(1, torch.tensor(ids, dtype=torch.long), 0) - - def _sync_decode_work_cache_blocks( - self, - host: DeepSeekV4LayerCache, - dirty_by_rank: dict[int, dict[str, set[int]]], - ) -> None: - """Upload only newly cleared pages when the cache is already resident.""" - device = self._decode_device_cache - worker = self._l3_worker - if device is None: - return - if worker is None: - raise RuntimeError("DeepSeekV4 resident cache exists without an L3 worker") - layout = self._compiled.layout - tensors_by_group = { - "ori": ((host.kv_cache, device.kv_cache, DEEPSEEK_V4_FWD_NUM_LAYERS, layout.decode_ori_max_blocks),), - "cmp": ((host.cmp_kv, device.cmp_kv, DEEPSEEK_V4_FWD_NUM_LAYERS, layout.cmp_max_blocks),), - "idx": ( - (host.idx_kv_cache, device.idx_kv_cache, DEEPSEEK_V4_CSA_NUM_LAYERS, layout.idx_max_blocks), - (host.idx_kv_scale, device.idx_kv_scale, DEEPSEEK_V4_CSA_NUM_LAYERS, layout.idx_max_blocks), - ), - "hca_state": (( - host.hca_compress_state, - device.hca_compress_state, - DEEPSEEK_V4_HCA_NUM_LAYERS, - layout.hca_state_max_blocks, - ),), - "csa_state": (( - host.csa_compress_state, - device.csa_compress_state, - DEEPSEEK_V4_CSA_NUM_LAYERS, - layout.csa_state_max_blocks, - ),), - "csa_inner_state": (( - host.csa_inner_compress_state, - device.csa_inner_compress_state, - DEEPSEEK_V4_CSA_NUM_LAYERS, - layout.csa_inner_state_max_blocks, - ),), - } - for rank, groups in dirty_by_rank.items(): - for name, block_ids in groups.items(): - for host_tensor, stacked, layer_count, blocks_per_layer in tensors_by_group[name]: - shard = stacked.shards[rank] - worker_id = stacked.worker_ids[rank] - page_nbytes = host_tensor[rank, 0].numel() * host_tensor.element_size() - for layer in range(layer_count): - for block_id in block_ids: - flat_index = layer * blocks_per_layer + int(block_id) - page = host_tensor[rank, flat_index] - worker.copy_to( - shard.data_ptr + flat_index * page_nbytes, - page.data_ptr(), - page_nbytes, - worker_id=worker_id, - ) + for block_id in ids: + worker.copy_to( + shard.data_ptr + block_id * page_nbytes, + zero_page.data_ptr(), + page_nbytes, + worker_id=worker_id, + ) def _logits_for_hidden( self, @@ -4048,135 +4254,137 @@ def _inherited_host_weights(self) -> list[torch.Tensor]: tensors.extend(self._mtp_buffers.weights.values()) return tensors - def _ensure_decode_work_cache(self) -> DeepSeekV4LayerCache: - cache = self._decode_work_cache - if cache is not None: - return cache - self._ensure_shared_host_allocation_before_worker("decode work cache") - layout = self._compiled.layout - fwd_layers = DEEPSEEK_V4_FWD_NUM_LAYERS - csa_layers = DEEPSEEK_V4_CSA_NUM_LAYERS - hca_layers = DEEPSEEK_V4_HCA_NUM_LAYERS - cache = DeepSeekV4LayerCache( - kv_cache=self._shared_empty( - ( - layout.ranks, - fwd_layers * layout.decode_ori_max_blocks, - layout.block_size, - 1, - DEEPSEEK_V4_HEAD_DIM, - ), - torch.bfloat16, - name="decode_work_kv_cache", - ), - cmp_kv=self._shared_empty( - ( - layout.ranks, - fwd_layers * layout.cmp_max_blocks, - layout.block_size, - 1, - DEEPSEEK_V4_HEAD_DIM, - ), - torch.bfloat16, - name="decode_work_cmp_kv", - ), - idx_kv_cache=self._shared_empty( - ( - layout.ranks, - csa_layers * layout.idx_max_blocks, - layout.block_size, - 1, - DEEPSEEK_V4_IDX_HEAD_DIM, - ), - torch.int8, - name="decode_work_idx_kv_cache", - ), - idx_kv_scale=self._shared_empty( - ( - layout.ranks, - csa_layers * layout.idx_max_blocks, - layout.block_size, - 1, - 1, - ), - torch.float32, - name="decode_work_idx_kv_scale", - ), - hca_compress_state=self._shared_empty( - ( - layout.ranks, - hca_layers * layout.hca_state_max_blocks, - layout.c128_state_block_size, - DEEPSEEK_V4_HCA_STATE_DIM, - ), - torch.float32, - name="decode_work_hca_compress_state", - ), - csa_compress_state=self._shared_empty( - ( - layout.ranks, - csa_layers * layout.csa_state_max_blocks, - layout.c4_state_block_size, - DEEPSEEK_V4_CSA_STATE_DIM, - ), - torch.float32, - name="decode_work_csa_compress_state", - ), - csa_inner_compress_state=self._shared_empty( - ( - layout.ranks, - csa_layers * layout.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", - ), + def _ensure_cache_zero_page(self) -> torch.Tensor: + """Create one pre-fork zero source reused for device-page clearing.""" + if self._cache_zero_page is not None: + return self._cache_zero_page + self._ensure_shared_host_allocation_before_worker("cache zero page") + max_page_bytes = max( + spec.spec.page_size_bytes // max(len(spec.layer_indices), 1) + for spec in self._cache_group_specs ) - for tensor in ( - cache.kv_cache, - cache.cmp_kv, - cache.idx_kv_cache, - cache.idx_kv_scale, - cache.hca_compress_state, - cache.csa_compress_state, - cache.csa_inner_compress_state, - ): - tensor.zero_() - self._decode_work_cache = cache - return cache + max_page_bytes = max( + max_page_bytes, + self._compiled.layout.block_size + * DEEPSEEK_V4_HEAD_DIM + * torch.bfloat16.itemsize, + ) + self._cache_zero_page = self._shared_empty( + (max_page_bytes,), + torch.uint8, + name="cache_zero_page", + ) + self._cache_zero_page.zero_() + return self._cache_zero_page - 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 + def _alloc_empty_stacked_tensor( + self, + full_shape: tuple[int, ...], + dtype: torch.dtype, + ) -> StackedDeviceTensor: + """Allocate an uninitialized shard directly on every chip worker.""" + worker = self._shared_l3_worker() + worker_ids = tuple(range(self._compiled.layout.ranks)) + shards: list[DeviceTensor] = [] + try: + for worker_id in worker_ids: + shards.append( + worker.alloc_tensor( + full_shape[1:], + dtype, + worker_id=worker_id, + ) + ) + except Exception: + for shard, worker_id in zip(shards, worker_ids, strict=False): + worker.free_tensor(shard, worker_id=worker_id) + raise + return StackedDeviceTensor(shards, full_shape, worker_ids) def _materialize_decode_device_cache(self) -> DeepSeekV4DeviceCache: - """Upload one shard per rank once and keep all cache pools worker-resident.""" + """Allocate dynamically sized cache shards directly on each NPU.""" cache = self._decode_device_cache if cache is not None: return cache - host = self._require_decode_work_cache() worker = self._shared_l3_worker() allocated: list[StackedDeviceTensor] = [] + layout = self._compiled.layout - def resident(tensor: torch.Tensor) -> StackedDeviceTensor: - stacked = worker.alloc_stacked_tensor( - tensor, - worker_ids=range(self._compiled.layout.ranks), - ) + def resident(shape: tuple[int, ...], dtype: torch.dtype) -> StackedDeviceTensor: + stacked = self._alloc_empty_stacked_tensor(shape, dtype) allocated.append(stacked) return stacked try: cache = DeepSeekV4DeviceCache( - kv_cache=resident(host.kv_cache), - cmp_kv=resident(host.cmp_kv), - idx_kv_cache=resident(host.idx_kv_cache), - idx_kv_scale=resident(host.idx_kv_scale), - hca_compress_state=resident(host.hca_compress_state), - csa_compress_state=resident(host.csa_compress_state), - csa_inner_compress_state=resident(host.csa_inner_compress_state), + kv_cache=resident( + ( + layout.ranks, + DEEPSEEK_V4_FWD_NUM_LAYERS * self._physical_cache_num_blocks("ori"), + layout.block_size, + 1, + DEEPSEEK_V4_HEAD_DIM, + ), + torch.bfloat16, + ), + cmp_kv=resident( + ( + layout.ranks, + DEEPSEEK_V4_FWD_NUM_LAYERS * self._physical_cache_num_blocks("cmp"), + layout.block_size, + 1, + DEEPSEEK_V4_HEAD_DIM, + ), + torch.bfloat16, + ), + idx_kv_cache=resident( + ( + layout.ranks, + DEEPSEEK_V4_CSA_NUM_LAYERS * self._physical_cache_num_blocks("idx"), + layout.block_size, + 1, + DEEPSEEK_V4_IDX_HEAD_DIM, + ), + torch.int8, + ), + idx_kv_scale=resident( + ( + layout.ranks, + DEEPSEEK_V4_CSA_NUM_LAYERS * self._physical_cache_num_blocks("idx"), + layout.block_size, + 1, + 1, + ), + torch.float32, + ), + hca_compress_state=resident( + ( + layout.ranks, + DEEPSEEK_V4_HCA_NUM_LAYERS * self._physical_cache_num_blocks("hca_state"), + layout.c128_state_block_size, + DEEPSEEK_V4_HCA_STATE_DIM, + ), + torch.float32, + ), + csa_compress_state=resident( + ( + layout.ranks, + DEEPSEEK_V4_CSA_NUM_LAYERS * self._physical_cache_num_blocks("csa_state"), + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_STATE_DIM, + ), + torch.float32, + ), + csa_inner_compress_state=resident( + ( + layout.ranks, + DEEPSEEK_V4_CSA_NUM_LAYERS + * self._physical_cache_num_blocks("csa_inner_state"), + layout.c4_state_block_size, + DEEPSEEK_V4_CSA_INNER_STATE_DIM, + ), + torch.float32, + ), ) except Exception: for tensor in allocated: @@ -4193,13 +4401,43 @@ def _materialize_mtp_device_kv_cache(self) -> StackedDeviceTensor | None: buffers = self._mtp_buffers if buffers is None: return None - cache = self._shared_l3_worker().alloc_stacked_tensor( - buffers.prefill_kv_cache, - worker_ids=range(self._compiled.layout.ranks), + layout = self._compiled.layout + cache = self._alloc_empty_stacked_tensor( + ( + layout.ranks, + self._physical_cache_num_blocks("ori"), + layout.block_size, + 1, + DEEPSEEK_V4_HEAD_DIM, + ), + torch.bfloat16, ) self._mtp_device_kv_cache = cache return cache + def _free_device_caches(self) -> None: + worker = self._l3_worker + if worker is None: + self._decode_device_cache = None + self._mtp_device_kv_cache = None + return + cache = self._decode_device_cache + if cache is not None: + for tensor in ( + cache.kv_cache, + cache.cmp_kv, + cache.idx_kv_cache, + cache.idx_kv_scale, + cache.hca_compress_state, + cache.csa_compress_state, + cache.csa_inner_compress_state, + ): + worker.free_stacked_tensor(tensor) + if self._mtp_device_kv_cache is not None: + worker.free_stacked_tensor(self._mtp_device_kv_cache) + self._decode_device_cache = None + self._mtp_device_kv_cache = None + @staticmethod def _static_device_tensor(tensor: torch.Tensor) -> torch.Tensor: if tensor.device.type != "cpu": @@ -4226,7 +4464,8 @@ def close(self) -> None: worker.close() finally: self._l3_worker = None - self._decode_work_cache = None + self._cache_zero_page = None + self._cache_group_num_blocks.clear() self._stacked_host_weights = None self._stacked_device_weights = None self._mtp_device_weights = None diff --git a/pypto_serving/model/qwen/kernel_cache.py b/pypto_serving/model/qwen/kernel_cache.py index cdbc34d0..a5b35112 100644 --- a/pypto_serving/model/qwen/kernel_cache.py +++ b/pypto_serving/model/qwen/kernel_cache.py @@ -124,21 +124,19 @@ def compute_params_fingerprint( dummy_args: Sequence[Any], *, platform: str, - block_dim: int, ) -> str: """Fingerprint of everything that specialises this kernel's binary. The ``dummy_args`` shape+dtype signature encodes every compile-time dimension (batch, max_seq, page_size, vocab, head_dim, block-table stride, - ...); combined with the kernel name, target platform, and block dim it - uniquely identifies the compiled kernel. Two launches differing in any of - these produce different fingerprints, so a config change can never be - served a stale binary. + ...); combined with the kernel name and target platform it uniquely + identifies the compiled kernel. Two launches differing in any of these + produce different fingerprints, so a config change can never be served a + stale binary. """ digest = hashlib.sha256() digest.update(name.encode()) digest.update(platform.encode()) - digest.update(str(block_dim).encode()) for arg in dummy_args: digest.update(str(tuple(arg.shape)).encode()) digest.update(str(arg.dtype).encode()) diff --git a/pypto_serving/model/qwen/npu_executor.py b/pypto_serving/model/qwen/npu_executor.py index 63f45905..ecd46a77 100644 --- a/pypto_serving/model/qwen/npu_executor.py +++ b/pypto_serving/model/qwen/npu_executor.py @@ -38,7 +38,6 @@ _VOCAB_PAD_MULTIPLE = 512 # must be a multiple of lm_head.VOCAB_CHUNK (64) _QWEN14B_PAGE_SIZE = 128 -_QWEN14B_BLOCK_DIM = 24 @dataclass @@ -516,12 +515,9 @@ def _compile_jit_fwd_callable( distributed_config = DistributedConfig( device_ids=list(self._device_ids), num_sub_workers=0, - block_dim=_QWEN14B_BLOCK_DIM, aicpu_thread_num=4, ) - params_fingerprint = compute_params_fingerprint( - name, dummy_args, platform=self._platform, block_dim=_QWEN14B_BLOCK_DIM, - ) + params_fingerprint = compute_params_fingerprint(name, dummy_args, platform=self._platform) if self._kernel_cache is not None: cached = self._kernel_cache.load( name, @@ -536,7 +532,6 @@ def _compile_jit_fwd_callable( return _L3Callable( compiled=cached, name=name, - block_dim=_QWEN14B_BLOCK_DIM, aicpu_thread_num=4, params_fingerprint=params_fingerprint, ) @@ -562,7 +557,6 @@ def _compile_jit_fwd_callable( return _L3Callable( compiled=compiled, name=name, - block_dim=_QWEN14B_BLOCK_DIM, aicpu_thread_num=4, params_fingerprint=params_fingerprint, ) diff --git a/pypto_serving/model/qwen/npu_runner.py b/pypto_serving/model/qwen/npu_runner.py index 3d44550c..c37471bd 100644 --- a/pypto_serving/model/qwen/npu_runner.py +++ b/pypto_serving/model/qwen/npu_runner.py @@ -74,7 +74,6 @@ class _L3Callable: compiled: object name: str - block_dim: int aicpu_thread_num: int dispatch_args: tuple[Any, ...] = () #: Fingerprint of this kernel's compile parameters; written into the cache @@ -797,7 +796,6 @@ def _run_distributed_program(self, callable_spec: _L3Callable, *args: Any) -> An """Run a compiled HOST wrapper through the shared PyPTO L3 worker.""" span_args = { "kernel": callable_spec.name, - "block_dim": callable_spec.block_dim, "aicpu_thread_num": callable_spec.aicpu_thread_num, } with profile_span( diff --git a/pypto_serving/serving/engine/async_engine.py b/pypto_serving/serving/engine/async_engine.py index f16ca4a5..b0718543 100644 --- a/pypto_serving/serving/engine/async_engine.py +++ b/pypto_serving/serving/engine/async_engine.py @@ -162,10 +162,8 @@ def __init__( block_size=block_size, enable_prefix_cache=self.config.enable_prefix_cache, ) - self.kv_cache_manager.init_groups( - runtime.kv_cache_groups, - max_batch_size=runtime.max_batch_size, - ) + if runtime.kv_cache_groups and self.config.enable_prefix_cache: + raise ValueError("Prefix caching is not supported with grouped KV cache pools") self._async_scheduling = self.config.resolve_async_scheduling() scheduler_config = SchedulerConfig( @@ -232,13 +230,20 @@ async def start(self) -> None: raise logger.info("Worker ready") - # Synchronise block metadata with the actual device-side KV cache size. - actual_num_pages = num_pages_value.value - if actual_num_pages <= 0: + # Synchronise block metadata with the actual device-side KV cache + # size. Grouped runners report the first group's rank-local block + # count; generic runners report their single-pool page count. + reported_num_blocks = num_pages_value.value + if reported_num_blocks <= 0: raise RuntimeError( - f"Worker reported invalid KV cache page count: {actual_num_pages}" + f"Worker reported invalid KV cache block count: {reported_num_blocks}" + ) + if self._runtime.kv_cache_groups: + self.kv_cache_manager.init_groups( + self._runtime.kv_cache_groups, + max_batch_size=self._runtime.max_batch_size, + primary_num_blocks=reported_num_blocks, ) - if self.kv_cache_manager.has_groups: logger.info( "Grouped KV cache pools initialised: %s", ", ".join( @@ -247,10 +252,10 @@ async def start(self) -> None: ), ) else: - self.kv_cache_manager._init_blocks(actual_num_pages, self._runtime.page_size) + self.kv_cache_manager._init_blocks(reported_num_blocks, self._runtime.page_size) logger.info( "KV cache block pool initialised: num_blocks=%d, block_size=%d", - actual_num_pages, + reported_num_blocks, self._runtime.page_size, ) diff --git a/pypto_serving/serving/memory/kv_cache.py b/pypto_serving/serving/memory/kv_cache.py index 3c051498..0ee2da6d 100644 --- a/pypto_serving/serving/memory/kv_cache.py +++ b/pypto_serving/serving/memory/kv_cache.py @@ -484,8 +484,15 @@ def init_groups( group_specs: tuple[KVCacheGroupSpec, ...], *, max_batch_size: int, + primary_num_blocks: int | None = None, ) -> None: - """Initialize independent physical pools for model-specific caches.""" + """Initialize independent physical pools for model-specific caches. + + ``primary_num_blocks`` is the device-reported capacity of the first + group. Unspecified groups are scaled to the same number of + ``max_blocks_per_seq`` capacity slots, keeping heterogeneous physical + pools in lockstep with the runner's allocation. + """ if max_batch_size <= 0: raise ValueError("max_batch_size must be positive") if not group_specs: @@ -498,8 +505,25 @@ def init_groups( if len(partition_counts) != 1: raise ValueError("KV cache groups must use the same num_partitions") + capacity_slots = None + if primary_num_blocks is not None: + primary_num_blocks = int(primary_num_blocks) + primary_stride = group_specs[0].max_blocks_per_seq + if primary_num_blocks <= 0 or primary_num_blocks % primary_stride: + raise ValueError( + "primary_num_blocks must be a positive multiple of the first " + "KV cache group's max_blocks_per_seq" + ) + capacity_slots = primary_num_blocks // primary_stride requested_sizes = { - group.name: group.num_blocks or max_batch_size * group.max_blocks_per_seq + group.name: ( + group.num_blocks + or ( + capacity_slots * group.max_blocks_per_seq + if capacity_slots is not None + else max_batch_size * group.max_blocks_per_seq + ) + ) for group in group_specs } undersized = [ diff --git a/tests/test_batching.py b/tests/test_batching.py index 3dc0826d..c88466e3 100644 --- a/tests/test_batching.py +++ b/tests/test_batching.py @@ -31,6 +31,7 @@ RuntimeModel, ) from pypto_serving.model.common.executor.executor import ModelExecutor +from pypto_serving.model.qwen.kernel_cache import compute_params_fingerprint from pypto_serving.model.qwen.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor from pypto_serving.model.qwen.npu_runner import ( _CompiledKernels, @@ -356,6 +357,31 @@ def test_grouped_cache_preemption_removes_victim_from_running_queue(): assert [request.request_id for request in scheduler.waiting] == ["newer"] +def test_grouped_cache_capacity_scales_from_device_reported_primary_pool(): + manager = KvCacheManager(block_size=1, enable_prefix_cache=False) + manager.init_groups( + ( + KVCacheGroupSpec( + name="primary", + layer_indices=(0,), + spec=KVCacheSpec(block_size=1, page_size_bytes=4), + max_blocks_per_seq=3, + ), + KVCacheGroupSpec( + name="compressed", + layer_indices=(1,), + spec=KVCacheSpec(block_size=1, page_size_bytes=2), + max_blocks_per_seq=2, + ), + ), + max_batch_size=8, + primary_num_blocks=6, + ) + + assert manager.group_num_blocks("primary") == 6 + assert manager.group_num_blocks("compressed") == 4 + + def test_step_command_preserves_grouped_cache_metadata_on_preempted_restart(): core = ReplicaEngineCore.__new__(ReplicaEngineCore) core._worker_known_req_ids = {"req"} @@ -532,7 +558,6 @@ def _compiled_kernels( callable_ = _L3Callable( compiled=object(), name="fake", - block_dim=1, aicpu_thread_num=1, ) if decode_weights is None: @@ -1588,7 +1613,6 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc fake_callable = _L3Callable( compiled=fake_kernel, name="fake", - block_dim=1, aicpu_thread_num=1, ) compiled = _compiled_kernels( @@ -1647,6 +1671,82 @@ def test_pypto_executor_uses_cached_kernel_weights_after_registration(monkeypatc manager.free(decode_alloc) +def test_qwen_compile_uses_current_distributed_config_interface(monkeypatch): + import pypto.ir.distributed_compiled_program as distributed_program + + class StrictDistributedConfig: + """Mirror runtimes that retain AICPU tuning but remove block_dim.""" + + def __init__(self, *, device_ids, num_sub_workers, aicpu_thread_num): + self.device_ids = device_ids + self.num_sub_workers = num_sub_workers + self.aicpu_thread_num = aicpu_thread_num + + captured = {} + + class FakeJitFunction: + def compile(self, *args, config): + captured["config"] = config + return distributed_program.DistributedCompiledProgram.__new__( + distributed_program.DistributedCompiledProgram + ) + + monkeypatch.setattr(distributed_program, "DistributedConfig", StrictDistributedConfig) + executor = PyptoExecutor(device_ids=(3,)) + + callable_spec = executor._compile_jit_fwd_callable( + "fake", + FakeJitFunction(), + [], + ) + + run_config = captured["config"] + assert not hasattr(run_config, "block_dim") + assert run_config.distributed_config.device_ids == [3] + assert run_config.distributed_config.num_sub_workers == 0 + assert run_config.distributed_config.aicpu_thread_num == 4 + assert callable_spec.aicpu_thread_num == 4 + + +def test_qwen_kernel_cache_hit_uses_current_callable_interface(): + cached_program = object() + captured = {} + + class FakeKernelCache: + def load(self, name, params_fingerprint, *, platform, distributed_config): + captured["name"] = name + captured["params_fingerprint"] = params_fingerprint + captured["platform"] = platform + captured["distributed_config"] = distributed_config + return cached_program + + class CompileMustNotRun: + def compile(self, *args, config): + raise AssertionError("a kernel-cache hit must skip JIT compilation") + + executor = PyptoExecutor(device_ids=(3,)) + executor._kernel_cache = FakeKernelCache() + + callable_spec = executor._compile_jit_fwd_callable( + "fake", + CompileMustNotRun(), + [torch.empty((2, 4), dtype=torch.bfloat16)], + ) + + assert callable_spec.compiled is cached_program + assert callable_spec.name == "fake" + assert callable_spec.aicpu_thread_num == 4 + assert not hasattr(callable_spec, "block_dim") + assert captured["name"] == "fake" + assert captured["params_fingerprint"] == compute_params_fingerprint( + "fake", + [torch.empty((2, 4), dtype=torch.bfloat16)], + platform=executor._platform, + ) + assert captured["platform"] == executor._platform + assert captured["distributed_config"].device_ids == [3] + + def test_pypto_executor_preserves_device_group(): executor = PyptoExecutor(device_ids=[3, 4]) diff --git a/tests/test_deepseek_v4.py b/tests/test_deepseek_v4.py index 020c1219..6f4e315b 100644 --- a/tests/test_deepseek_v4.py +++ b/tests/test_deepseek_v4.py @@ -16,6 +16,7 @@ import pytest import torch +from pypto.runtime import DeviceTensor import pypto_serving.cli.main as cli from pypto_serving.config.types import DecodeBatch, PrefillBatch, RuntimeConfig @@ -30,7 +31,9 @@ DeepSeekV4L3Callable, DeepSeekV4ModelRunner, accept_mtp_tokens, + build_deepseek_v4_cache_group_specs, build_deepseek_v4_layer_plan, + deepseek_v4_cache_blocks_for_slots, ) from pypto_serving.model.deepseek.weight_loader import ( DeepSeekV4WeightStore, @@ -317,6 +320,10 @@ def _fake_compile(self, name, jit_fn, dummy_args): assert compiled_args["deepseek_v4_decode"][0].shape == (8, 8, 4, 4096) assert compiled_args["deepseek_v4_prefill"][0].dtype == torch.float32 assert compiled_args["deepseek_v4_decode"][0].dtype == torch.float32 + cache_blocks = executor._compile_cache_blocks( + loaded.runtime_model, + DeepSeekV4CacheLayout(decode_batch=8, decode_seq=1, decode_tokens=8), + ) prefill_order = npu_executor._PREFILL_FWD_TENSOR_ORDER # Packed prefill flattens the FWD work caches to 5-D (kv_cache/cmp_kv stack x43, # idx_kv_cache stacks x21 across the CSA group) and stacks the compress-state @@ -325,17 +332,31 @@ def _fake_compile(self, name, jit_fn, dummy_args): # per-rank copies (the kernel slices them per layer). The kernel emits # final-normalized hidden rows. prefill_args = compiled_args["deepseek_v4_prefill"] - assert prefill_args[prefill_order.index("kv_cache")].shape == (8, 43 * 128, 128, 1, 512) - assert prefill_args[prefill_order.index("cmp_kv")].shape == (8, 43 * 32, 128, 1, 512) - assert prefill_args[prefill_order.index("idx_kv_cache")].shape == (8, 21 * 64, 128, 1, 128) + assert prefill_args[prefill_order.index("kv_cache")].shape == ( + 8, 43 * cache_blocks["ori"], 128, 1, 512 + ) + assert prefill_args[prefill_order.index("cmp_kv")].shape == ( + 8, 43 * cache_blocks["cmp"], 128, 1, 512 + ) + assert prefill_args[prefill_order.index("idx_kv_cache")].shape == ( + 8, 21 * cache_blocks["idx"], 128, 1, 128 + ) assert prefill_args[prefill_order.index("idx_kv_cache")].dtype == torch.int8 - assert prefill_args[prefill_order.index("idx_kv_scale")].shape == (8, 21 * 64, 128, 1, 1) + assert prefill_args[prefill_order.index("idx_kv_scale")].shape == ( + 8, 21 * cache_blocks["idx"], 128, 1, 1 + ) assert prefill_args[prefill_order.index("hca_cmp_wkv")].shape == (8, 20 * 512, 4096) assert prefill_args[prefill_order.index("csa_cmp_wkv")].shape == (8, 21 * 1024, 4096) assert prefill_args[prefill_order.index("csa_inner_wkv")].shape == (8, 21 * 256, 4096) - assert prefill_args[prefill_order.index("hca_compress_state")].shape == (8, 20 * 64, 8, 1024) - assert prefill_args[prefill_order.index("csa_compress_state")].shape == (8, 21 * 65, 4, 2048) - assert prefill_args[prefill_order.index("csa_inner_compress_state")].shape == (8, 21 * 65, 4, 512) + assert prefill_args[prefill_order.index("hca_compress_state")].shape == ( + 8, 20 * cache_blocks["hca_state"], 8, 1024 + ) + assert prefill_args[prefill_order.index("csa_compress_state")].shape == ( + 8, 21 * cache_blocks["csa_state"], 4, 2048 + ) + assert prefill_args[prefill_order.index("csa_inner_compress_state")].shape == ( + 8, 21 * cache_blocks["csa_inner_state"], 4, 512 + ) assert prefill_args[prefill_order.index("hca_compress_state_block_table")].shape == (8, 2048) assert prefill_args[prefill_order.index("csa_compress_state_block_table")].shape == (8, 4096) assert prefill_args[prefill_order.index("csa_inner_compress_state_block_table")].shape == (8, 4096) @@ -357,13 +378,17 @@ def _fake_compile(self, name, jit_fn, dummy_args): assert prefill_args[prefill_order.index("num_tokens_per_owner")].shape == (8,) assert prefill_args[prefill_order.index("logit_row_indices")].shape == (8, 8) decode_order = npu_executor._DECODE_FWD_TENSOR_ORDER - # Compress-state work caches are stacked across the CSA (x21) and HCA (x20) layer - # groups, each layer holding decode_batch (4) x state_max_blocks rows. - assert compiled_args["deepseek_v4_decode"][decode_order.index("hca_compress_state")].shape == (8, 20 * 64, 8, 1024) - assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_compress_state")].shape == (8, 21 * 65, 4, 2048) + # Cache tensor dimensions follow the runtime pool capacity; table depths + # remain the fixed kernel ABI below. + assert compiled_args["deepseek_v4_decode"][decode_order.index("hca_compress_state")].shape == ( + 8, 20 * cache_blocks["hca_state"], 8, 1024 + ) + assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_compress_state")].shape == ( + 8, 21 * cache_blocks["csa_state"], 4, 2048 + ) assert compiled_args["deepseek_v4_decode"][decode_order.index("csa_inner_compress_state")].shape == ( 8, - 21 * 65, + 21 * cache_blocks["csa_inner_state"], 4, 512, ) @@ -383,12 +408,17 @@ def _fake_compile(self, name, jit_fn, dummy_args): assert compiled_args["deepseek_v4_decode"][decode_order.index("logits")].shape == (8, 8, 129280) assert compiled_args["deepseek_v4_decode"][decode_order.index("num_tokens_per_owner")].shape == (8,) assert compiled_args["deepseek_v4_decode"][decode_order.index("logit_row_indices")].shape == (8, 8) - # Decode ori-KV uses the same fixed 128-block physical pool as prefill. + assert compiled_args["deepseek_v4_decode"][decode_order.index("num_logit_rows")].shape == (8,) + # Prefill and decode share the same runtime-sized physical pool. decode_args = compiled_args["deepseek_v4_decode"] - assert decode_args[decode_order.index("kv_cache")].shape == (8, 43 * 128, 128, 1, 512) + assert decode_args[decode_order.index("kv_cache")].shape == ( + 8, 43 * cache_blocks["ori"], 128, 1, 512 + ) assert decode_args[decode_order.index("block_table")].shape == (8, 8, 128) assert decode_args[decode_order.index("idx_kv_cache")].dtype == torch.int8 - assert decode_args[decode_order.index("idx_kv_scale")].shape == (8, 21 * 64, 128, 1, 1) + assert decode_args[decode_order.index("idx_kv_scale")].shape == ( + 8, 21 * cache_blocks["idx"], 128, 1, 1 + ) # SWA/HCA/CSA metadata all use the cache-first full window, plus the paged # write slot mapping. assert decode_args[decode_order.index("swa_slot_mapping")].shape == (8, 8) @@ -784,6 +814,169 @@ def test_deepseek_cache_metadata_maps_scheduler_block_ids(): assert hca_state_mapping.tolist() == [[64 * 8, 64 * 8 + 1, 64 * 8 + 2]] +def test_deepseek_cache_group_specs_leave_physical_capacity_for_runtime_sizing(): + compress_ratios = (0, 0, *([4] * 21), *([128] * 20)) + specs = build_deepseek_v4_cache_group_specs(43, compress_ratios, decode_batch=8) + by_name = {spec.name: spec for spec in specs} + + assert all(spec.num_blocks is None for spec in specs) + assert by_name["ori"].spec.page_size_bytes == 43 * 128 * 512 * 2 + assert by_name["idx"].spec.page_size_bytes == 21 * 128 * (128 + 4) + assert by_name["hca_state"].spec.page_size_bytes == 20 * 8 * 1024 * 4 + assert deepseek_v4_cache_blocks_for_slots(specs, 3) == { + "ori": 96, + "cmp": 24, + "idx": 48, + "hca_state": 48, + "csa_state": 48, + "csa_inner_state": 48, + } + + +def test_deepseek_cache_sizing_uses_limiting_rank_post_weight_budget(monkeypatch): + layout = DeepSeekV4CacheLayout(decode_batch=8, decode_seq=1, decode_tokens=8) + runner = DeepSeekV4ModelRunner( + compiled=DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=tuple([0] * 43), + layer_plan=(), + kernel_dir="", + device_id=2, + device_ids=(2, 5), + ) + ) + runner._cache_group_specs = build_deepseek_v4_cache_group_specs( + 43, + runner._compiled.compress_ratios, + decode_batch=layout.decode_batch, + ) + memory = { + "npu:2": (5_000_000_000, 10_000_000_000), + "npu:5": (4_000_000_000, 10_000_000_000), + } + monkeypatch.setattr(torch.npu, "mem_get_info", lambda device: memory[device]) + runtime = RuntimeConfig(npu_memory_utilization=0.8) + + bytes_per_slot = sum( + spec.max_blocks_per_seq * spec.spec.page_size_bytes + for spec in runner._cache_group_specs + ) + scratch_bytes = sum( + layout.decode_batch * spec.spec.page_size_bytes + for spec in runner._cache_group_specs + ) + expected = max((2_000_000_000 - scratch_bytes) // bytes_per_slot, 1) + + assert runner._compute_kv_cache_capacity_slots(runtime) == expected + + +def test_deepseek_cache_allocation_halves_all_groups_together_on_oom(monkeypatch): + layout = DeepSeekV4CacheLayout(decode_batch=8, decode_seq=1, decode_tokens=8) + runner = DeepSeekV4ModelRunner( + compiled=DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=tuple([0] * 43), + layer_plan=(), + kernel_dir="", + ) + ) + runner._cache_group_specs = build_deepseek_v4_cache_group_specs( + 43, + runner._compiled.compress_ratios, + decode_batch=layout.decode_batch, + ) + attempts = [] + + def allocate_main_cache(): + slots = runner._cache_group_num_blocks["ori"] // 32 + attempts.append(slots) + if slots > 2: + raise MemoryError("synthetic OOM") + return object() + + monkeypatch.setattr(runner, "_materialize_decode_device_cache", allocate_main_cache) + monkeypatch.setattr(runner, "_materialize_mtp_device_kv_cache", lambda: None) + monkeypatch.setattr(runner, "_zero_scratch_cache_blocks", lambda: None) + + assert runner._alloc_kv_cache_with_retry(8) == 2 + assert attempts == [8, 4, 2] + assert runner._cache_group_num_blocks == deepseek_v4_cache_blocks_for_slots( + runner._cache_group_specs, + 2, + ) + + +def test_deepseek_device_cache_allocates_runtime_sized_rank_shards(): + layout = DeepSeekV4CacheLayout( + ranks=2, + block_size=1, + decode_batch=1, + decode_seq=1, + decode_tokens=1, + ) + runner = DeepSeekV4ModelRunner( + compiled=DeepSeekV4CompiledKernels( + layout=layout, + model_dir="", + weight_map={}, + weight_store=None, + compress_ratios=tuple([0] * 43), + layer_plan=(), + kernel_dir="", + ) + ) + runner._cache_group_num_blocks = { + name: 2 + for name in ("ori", "cmp", "idx", "hca_state", "csa_state", "csa_inner_state") + } + + class FakeWorker: + def __init__(self): + self.allocations = [] + self.frees = [] + + def alloc_tensor(self, shape, dtype, *, worker_id=0): + tensor = DeviceTensor( + 0x1000 + len(self.allocations) * 0x100, + tuple(shape), + dtype, + ) + self.allocations.append((worker_id, tensor)) + return tensor + + def free_tensor(self, tensor, *, worker_id=0): + self.frees.append((worker_id, tensor)) + + def free_stacked_tensor(self, stacked): + for tensor, worker_id in zip( + stacked.shards, + stacked.worker_ids, + strict=True, + ): + self.free_tensor(tensor, worker_id=worker_id) + + worker = FakeWorker() + runner._l3_worker = worker + + cache = runner._materialize_decode_device_cache() + + assert cache.kv_cache.full_shape == (2, 43 * 3, 1, 1, 512) + assert cache.cmp_kv.full_shape == (2, 43 * 3, 1, 1, 512) + assert cache.idx_kv_cache.full_shape == (2, 21 * 3, 1, 1, 128) + assert cache.hca_compress_state.full_shape == (2, 20 * 3, 8, 1024) + assert len(worker.allocations) == 14 + assert {worker_id for worker_id, _tensor in worker.allocations} == {0, 1} + + runner._free_device_caches() + assert len(worker.frees) == 14 + + def _grouped_cache_rows(count: int) -> list[dict[str, list[int]]]: names = ("ori", "cmp", "idx", "hca_state", "csa_state", "csa_inner_state") return [ @@ -1100,10 +1293,9 @@ def test_deepseek_prefill_staging_keeps_worker_resident_cache_tensors_out(): freqs_sin=torch.empty((1, 1), dtype=torch.bfloat16), ) ) - cache = runner._ensure_decode_work_cache() prefill = runner._ensure_prefill_fwd_buffers(hidden_size=1) - assert cache is runner._decode_work_cache + assert not hasattr(runner, "_decode_work_cache") for name in ( "kv_cache", "cmp_kv", @@ -1116,6 +1308,37 @@ def test_deepseek_prefill_staging_keeps_worker_resident_cache_tensors_out(): assert name not in prefill.tensors +def test_deepseek_shared_buffer_setup_does_not_require_decode_work_cache(monkeypatch): + runner, model = _runner_for_prepared_inputs() + runner._stacked_device_weights = object() + zero_page_calls = [] + + for method_name in ( + "load_packed_global_weights", + "_static_freqs_cos_tensor", + "_static_freqs_sin_tensor", + "_ensure_decode_buffers", + "_ensure_mtp_buffers", + "_require_prefill_output_buffer", + "_require_prefill_pre_hc_output_buffer", + "_require_prefill_logits_buffer", + "_static_final_norm_weight_tensor", + "_static_lm_head_weight_tensor", + "_require_decode_logits_buffer", + "_hc_head_tensors", + "_ensure_prefill_fwd_buffers", + "_assert_l3_shared_buffers_preallocated", + "_materialize_resident_weights", + ): + monkeypatch.setattr(runner, method_name, lambda *args, **kwargs: None) + monkeypatch.setattr(runner, "_ensure_cache_zero_page", lambda: zero_page_calls.append(True)) + + runner._ensure_l3_shared_buffers(model) + + assert zero_page_calls == [True] + assert not hasattr(runner, "_decode_work_cache") + + def test_deepseek_mtp_prefill_and_decode_reuse_same_kv_cache(): layout = DeepSeekV4CacheLayout( ranks=1, diff --git a/tests/test_kernel_cache.py b/tests/test_kernel_cache.py index 259a3805..e2ad2c6f 100644 --- a/tests/test_kernel_cache.py +++ b/tests/test_kernel_cache.py @@ -90,8 +90,8 @@ def _cache(tmp_path, code_fingerprint="v1+abcd"): # --- params fingerprint ---------------------------------------------------- -def _pf(name, args, platform="a2a3", block_dim=24): - return compute_params_fingerprint(name, args, platform=platform, block_dim=block_dim) +def _pf(name, args, platform="a2a3"): + return compute_params_fingerprint(name, args, platform=platform) def test_params_fingerprint_is_stable(): @@ -105,10 +105,9 @@ def test_params_fingerprint_tracks_every_distinguishing_dimension(): base = _pf("decode_fwd", args) # max_seq (shape) change -> the exact bug the name-only key missed assert base != _pf("decode_fwd", [_FakeTensor((16, 2048), "bfloat16")]) - # dtype, platform, block_dim, and kernel name all distinguish a binary + # dtype, platform, and kernel name all distinguish a binary assert base != _pf("decode_fwd", [_FakeTensor((16, 512), "float32")]) assert base != _pf("decode_fwd", args, platform="a2a3sim") - assert base != _pf("decode_fwd", args, block_dim=48) assert base != _pf("prefill_fwd", args)