Skip to content

Refactor serving code #113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
superxf merged 2 commits into from
Jul 28, 2026
Merged

Refactor serving code #113

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cef4326
support kvcach dynamic shape
superxf Jul 25, 2026
a9bb0f8
supoort pypto DistributedConfig and RuntimeConfig change
superxf Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pypto-lib
Open in desktop
Submodule pypto-lib updated 7 files
+3 −2 models/deepseek/v4-flash/decode_fwd.py
+11 −3 models/deepseek/v4-flash/decode_mtp.py
+33 −31 models/deepseek/v4-flash/decode_sparse_attn.py
+42 −32 models/deepseek/v4-flash/decode_sparse_attn_hca.py
+116 −78 models/deepseek/v4-flash/lm_head.py
+3 −2 models/deepseek/v4-flash/prefill_fwd.py
+11 −3 models/deepseek/v4-flash/prefill_mtp.py
86 changes: 64 additions & 22 deletions pypto_serving/model/deepseek/npu_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
),
)
Expand All @@ -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,
),
)
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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,
),
Expand All @@ -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,
),
Expand All @@ -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,
),
Expand All @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
),
Expand All @@ -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,
),
Expand All @@ -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,
),
Expand Down Expand Up @@ -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,
Expand All @@ -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,
},
}
Expand Down
Loading
Loading