From 835c583dfbfceb5129e408b2920bc8d1aec88eb9 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Thu, 17 Sep 2026 05:18:30 -0700 Subject: [PATCH 01/11] [kv-shard 3c] Validate and enable the sharded control plane Expose the flag only after scratch and scheduler contracts exist. Resolve supported configurations, reserve scratch HBM, select sharded pools and allocator, and require the rotation-aware cache together. (cherry picked from commit 40fa5266f3d41167163a4ce2e4029526a1d803dc) --- python/sglang/srt/arg_groups/field_order.py | 1 + .../sglang/srt/arg_groups/fields/parallel.py | 8 + python/sglang/srt/arg_groups/kv_shard_hook.py | 287 ++++++++++++++++++ python/sglang/srt/arg_groups/overrides.py | 6 + python/sglang/srt/arg_groups/pipeline.py | 10 + .../sglang/srt/arg_groups/resolution_hooks.py | 1 + .../srt/mem_cache/kv_cache_configurator.py | 75 ++++- python/sglang/srt/mem_cache/registry.py | 11 +- .../srt/model_executor/pool_configurator.py | 12 +- test/registered/dcp/test_dcp_layout_unit.py | 2 + .../test_mooncake_efa_allocator.py | 1 + .../unit/mem_cache/test_registry.py | 32 ++ .../model_executor/test_pool_configurator.py | 74 +++++ .../unit/models/test_deepseek_mla_dispatch.py | 108 ++++++- .../unit/server_args/test_server_args.py | 213 +++++++++++++ 15 files changed, 836 insertions(+), 5 deletions(-) create mode 100644 python/sglang/srt/arg_groups/kv_shard_hook.py diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index cd77b9ebb305..8c63306092f6 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/python/sglang/srt/arg_groups/field_order.py @@ -90,6 +90,7 @@ "cp_strategy", "enable_dsa_cache_layer_split", "enable_cp_decode_attn_tp", + "enable_kv_cache_sharding", "enable_dp_attention", "enable_dp_attention_local_control_broadcast", "enable_dp_lm_head", diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index cca9bc325fa4..caa549ae5a02 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -170,6 +170,14 @@ class Parallel(msgspec.Struct): bool, "Enable attention tensor-parallel weight slicing during decode under context parallel (cp_size>1). Slices the replicated attention linears to the local CP partition, eliminating redundant decode GEMMs.", ] = False + enable_kv_cache_sharding: A[ + bool, + "Shard at-rest KV cache storage at page granularity across the group " + "that currently replicates it (attention-CP ranks for GQA models " + "under prefill CP; attention-TP ranks for MLA models), multiplying " + "the group's unique KV capacity. Prefill workers of PD " + "disaggregation only.", + ] = False # Deprecated spelling of `attn_dp_size`: `--dp-size N --enable-dp-attention` # resolves to `attn_dp_size = N`, `dp_size = 1`. TODO: remove after 2026-12-31. enable_dp_attention: A[bool, Arg(no_cli=True, resolvable=True)] = False diff --git a/python/sglang/srt/arg_groups/kv_shard_hook.py b/python/sglang/srt/arg_groups/kv_shard_hook.py new file mode 100644 index 000000000000..d9c7747dfe25 --- /dev/null +++ b/python/sglang/srt/arg_groups/kv_shard_hook.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Resolution and validation for logical-page KV cache sharding.""" + +from __future__ import annotations + +import logging +from typing import Any, Optional + +from sglang.srt.arg_groups.overrides import ( + attention_backends_of, + declare_resolution, + model_config_of, + resolved_view, + resolving_view, + use_mla_backend, +) +from sglang.srt.configs.model_config import is_deepseek_dsa +from sglang.srt.environ import envs +from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase +from sglang.srt.utils.common import get_device_memory_capacity + +logger = logging.getLogger(__name__) + + +def kv_shard_group_size(server_args: Any) -> int: + """Return the active KV-shard group size, or one when disabled.""" + cfg = resolving_view(server_args) + if not cfg.enable_kv_cache_sharding: + return 1 + if cfg.attn_cp_size > 1: + return cfg.attn_cp_size + if use_mla_backend(server_args): + # Plain-TP MLA is validated below, so attn_tp_size == tp_size here. + return cfg.tp_size + return 1 + + +def validate_kv_shard_attention_backend(server_args: Any, cfg: Any = None) -> None: + cfg = cfg or resolving_view(server_args) + prefill_backend, _ = attention_backends_of(resolved_view(server_args)) + if prefill_backend == "fa3": + return + if prefill_backend != "trtllm_mla": + raise ValueError( + "--enable-kv-cache-sharding currently requires the fa3 prefill " + "attention backend, or trtllm_mla for plain-TP MLA; " + f"got {prefill_backend!r}." + ) + if not use_mla_backend(server_args): + raise ValueError("KV sharding with trtllm_mla requires an MLA model.") + if cfg.attn_cp_size > 1: + raise ValueError( + "KV sharding with trtllm_mla only supports plain-TP MLA " + "(attn_cp_size must be 1)." + ) + if cfg.disable_chunked_prefix_cache: + raise ValueError("KV sharding with trtllm_mla requires chunked prefix caching.") + + +def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) -> None: + """Validate KV sharding and declare its derived eager-prefill settings.""" + cfg = resolving_view(server_args) + if not cfg.enable_kv_cache_sharding: + return + + if gpu_mem is None: + # `handle_gpu_memory_settings` resolves capacity internally now, so the + # pipeline no longer threads it through. Only the chunk-size rounding + # below reads it, and only when mem_fraction_static was left unset. + gpu_mem = get_device_memory_capacity(cfg.device) + + if cfg.disaggregation_mode != "prefill": + raise ValueError( + "--enable-kv-cache-sharding is only supported on PD prefill " + "workers: a rank's pool holds only its stripe of every cached " + "page, which is incompatible with local decode." + ) + if cfg.speculative_algorithm is not None: + raise ValueError( + "--enable-kv-cache-sharding does not support speculative decoding." + ) + if cfg.dllm_algorithm is not None: + raise ValueError( + "--enable-kv-cache-sharding does not support diffusion language " + "models: DLLM admission bypasses the shard-scratch budget and " + "resolves out of PD-prefill mode." + ) + if cfg.radix_cache_backend is not None: + raise ValueError( + "--enable-kv-cache-sharding requires the built-in " + "UnifiedRadixCache with a rotation-aware UnifiedTreeCore and does " + "not support " + f"--radix-cache-backend={cfg.radix_cache_backend!r}." + ) + tree_core_backend = envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get() + if tree_core_backend != "python": + raise ValueError( + "--enable-kv-cache-sharding requires the Python UnifiedTreeCore: " + f"SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND={tree_core_backend!r} is " + "not supported for sharded KV." + ) + if cfg.enable_unified_cache_external_linker: + raise ValueError( + "--enable-kv-cache-sharding does not support " + "--enable-unified-cache-external-linker: load-back allocates " + "through the flat allocator API." + ) + if cfg.enable_streaming_session: + raise ValueError( + "--enable-kv-cache-sharding does not support streaming sessions: " + "their virtual nodes do not carry KV-shard rotation metadata." + ) + if cfg.enable_session_radix_cache: + raise ValueError( + "--enable-kv-cache-sharding does not support session radix cache yet." + ) + if cfg.enable_flexkv: + raise ValueError( + "--enable-kv-cache-sharding does not support --enable-flexkv: " + "FlexKV load-back allocates through the stock flat allocator API." + ) + if cfg.enable_unified_memory: + raise ValueError( + "--enable-kv-cache-sharding does not support --enable-unified-memory: " + "the unified pool uses a different virtual-id allocator and layout." + ) + if cfg.dcp_size > 1: + raise ValueError( + "--enable-kv-cache-sharding and decode context parallelism " + "(--dcp-size) are mutually exclusive users of the widened " + "allocator index space." + ) + if cfg.enable_hierarchical_cache: + raise ValueError( + "--enable-kv-cache-sharding does not support " + "--enable-hierarchical-cache yet: HiCache backup/load assumes " + "each rank holds full pages." + ) + if cfg.enable_lmcache: + raise ValueError( + "--enable-kv-cache-sharding does not support --enable-lmcache." + ) + if cfg.enable_hisparse: + raise ValueError( + "--enable-kv-cache-sharding does not support --enable-hisparse: " + "HiSparse selects a non-sharded allocator and assumes each rank " + "owns complete KV pages." + ) + if cfg.enable_dynamic_chunking: + raise ValueError( + "--enable-kv-cache-sharding requires fixed chunk sizes; disable " + "--enable-dynamic-chunking." + ) + if cfg.enable_two_batch_overlap: + raise ValueError( + "--enable-kv-cache-sharding does not support " + "--enable-two-batch-overlap: the sharded pool's per-batch gather " + "plan is a process-wide singleton, but TBO rebuilds attention " + "metadata multiple times per extend forward and clobbers the plan." + ) + + model_config = model_config_of(server_args) + if model_config.is_encoder_decoder: + raise ValueError( + "--enable-kv-cache-sharding does not support encoder-decoder " + "models: cross-attention uses a separate encoder page table that " + "is not part of the sharded gather plan." + ) + if model_config.attention_chunk_size is not None: + raise ValueError( + "--enable-kv-cache-sharding does not support chunked/local-attention " + "models yet: local-attention metadata is built from the page table " + "before the scratch translation." + ) + + validate_kv_shard_attention_backend(server_args, cfg) + if cfg.disaggregation_transfer_backend not in ("mooncake", "nixl"): + raise ValueError( + "--enable-kv-cache-sharding currently only supports the mooncake " + "and nixl transfer backends: both route sends through " + "CommonKVSender._prepare_send_indices and pair destinations by " + "fancy index. Got --disaggregation-transfer-backend " + f"{cfg.disaggregation_transfer_backend!r}." + ) + if envs.SGLANG_DISAGG_STAGING_BUFFER.get(): + raise ValueError( + "--enable-kv-cache-sharding is incompatible with " + "SGLANG_DISAGG_STAGING_BUFFER (the staging path assumes " + "contiguous send slices)." + ) + + if use_mla_backend(server_args): + if is_deepseek_dsa(model_config.hf_config): + raise ValueError( + "--enable-kv-cache-sharding does not support DSA models yet " + "(indexer buffers are not striped)." + ) + if cfg.attn_cp_size <= 1: + if cfg.enable_dp_attention: + raise ValueError( + "--enable-kv-cache-sharding for MLA models without " + "prefill CP shards across attention-TP and requires plain " + "TP attention; disable --enable-dp-attention." + ) + if cfg.tp_size <= 1: + raise ValueError( + "--enable-kv-cache-sharding for MLA models without " + "prefill CP shards across attention-TP ranks and needs " + "tp_size > 1." + ) + elif not cfg.enable_prefill_cp or cfg.attn_cp_size <= 1: + raise ValueError( + "--enable-kv-cache-sharding for GQA models shards across " + "attention-CP ranks and requires --enable-prefill-cp with " + "attn_cp_size > 1." + ) + + page_size = cfg.page_size if cfg.page_size is not None else 1 + if page_size <= 1: + raise ValueError( + "--enable-kv-cache-sharding requires a real page size " + "(--page-size 64 recommended): the physical page is the P/D " + "transfer descriptor and gather-copy unit; page_size=1 " + "degenerates both into per-token operations." + ) + granule = kv_shard_group_size(server_args) * page_size + chunk_size = cfg.chunked_prefill_size + if chunk_size is None or chunk_size <= 0: + raise ValueError( + "--enable-kv-cache-sharding requires chunked prefill " + "(--chunked-prefill-size > 0)." + ) + + declarations = {} + if chunk_size % granule != 0: + rounded = (chunk_size + granule - 1) // granule * granule + declarations["chunked_prefill_size"] = rounded + + # The raw input, deliberately not the resolving view: the question is + # whether the user pinned mem_fraction_static, and the resolved value + # would answer yes for every run. + raw_mem_fraction = (getattr(server_args, "_raw_input", None) or {}).get( + "mem_fraction_static" + ) + activation_delta_mb = 1.5 * (max(rounded, 2048) - max(chunk_size, 2048)) + if raw_mem_fraction is None and activation_delta_mb > 0: + if not gpu_mem or cfg.mem_fraction_static is None: + raise ValueError( + "Cannot safely auto-round --chunked-prefill-size for KV " + "sharding without a known device-memory capacity; pass a " + f"multiple of shard_size * page_size ({granule})." + ) + corrected_mem_fraction = ( + float(cfg.mem_fraction_static) - activation_delta_mb / gpu_mem + ) + if corrected_mem_fraction <= 0: + raise ValueError( + "Auto-rounding --chunked-prefill-size for KV sharding " + "would consume all automatically reserved GPU memory; " + f"pass a smaller multiple of {granule}." + ) + declarations["mem_fraction_static"] = corrected_mem_fraction + + logger.warning( + "Rounding --chunked-prefill-size up from %d to %d " + "(a multiple of shard_size * page_size = %d) for KV sharding.", + chunk_size, + rounded, + granule, + ) + + if cfg.cuda_graph_config.prefill.backend != Backend.DISABLED: + logger.warning( + "Prefill CUDA graph (%s) is incompatible with " + "--enable-kv-cache-sharding; disabling it.", + cfg.cuda_graph_config.prefill.backend, + ) + declarations["cuda_graph_config"] = with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ) + + if declarations: + declare_resolution( + server_args, + "handle_kv_cache_sharding", + **declarations, + ) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index bb7641350252..ddc9e372d8cc 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1710,6 +1710,12 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool: False for any config the runtime won't post-capture-size, else it gets an under-reserved fraction.""" cfg = resolving_view(server_args) + # Logical-page sharding always runs prefill eagerly. This predicate runs + # before the sharding hook disables prefill capture, so gate on the feature + # itself to retain eager activation headroom in the memory heuristic. + if cfg.enable_kv_cache_sharding: + return False + mla_enabled = use_mla_backend(server_args) if not envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get(): return False diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 0ff314930f04..cb3ac9025d5e 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -6,6 +6,7 @@ from typing import Any from sglang.srt.arg_groups.arg_utils import record_fields +from sglang.srt.arg_groups.kv_shard_hook import handle_kv_cache_sharding from sglang.srt.arg_groups.overrides import ( _page_size_default, _pipeline_parallel_overlap_disable, @@ -266,6 +267,10 @@ def run_resolution_pipeline(server_args: Any) -> None: run_hook(handle_context_parallelism, server_args) + # Validate logical-page KV cache sharding after its page size, attention + # backend, memory budget, and shard topology have all resolved. + run_hook(handle_kv_cache_sharding, server_args) + from sglang.srt.arg_groups.moe_hook import ( handle_a2a_moe, handle_moe_kernel_config, @@ -336,6 +341,11 @@ def run_resolution_pipeline(server_args: Any) -> None: finalize_cuda_graph_prefill_max_context(server_args) + # Capability adjustments can late-force a different attention backend or + # disable chunked prefill. Re-run the idempotent sharding gate over the + # final resolving view. + run_hook(handle_kv_cache_sharding, server_args) + # Validate after all batch-size declarations are visible. run_hook(validate_deepep_v2_speculative_draft, server_args) run_hook(validate_deepep_v2_dispatch_token_budget, server_args) diff --git a/python/sglang/srt/arg_groups/resolution_hooks.py b/python/sglang/srt/arg_groups/resolution_hooks.py index cc401b45caef..68897a2c811c 100644 --- a/python/sglang/srt/arg_groups/resolution_hooks.py +++ b/python/sglang/srt/arg_groups/resolution_hooks.py @@ -68,6 +68,7 @@ "handle_data_parallelism", "handle_load_balance_method", "handle_context_parallelism", + "handle_kv_cache_sharding", "handle_moe_kernel_config", "handle_a2a_moe", "handle_eplb_and_dispatch", diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1e1667f20bf5..ec72f5236cae 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -47,6 +47,7 @@ DeepSeekV4HiSparseTokenToKVPoolAllocator, HiSparseTokenToKVPoolAllocator, ) +from sglang.srt.mem_cache.allocator.page_interleave import PageInterleavePoolAllocator from sglang.srt.mem_cache.allocator.swa import ( PureSWATokenToKVPoolAllocator, SWATokenToKVPoolAllocator, @@ -78,6 +79,16 @@ ReqToTokenPool, get_minimax_sparse_index_dtype, ) +from sglang.srt.mem_cache.page_interleave import ( + PageShardSpec, + get_kv_shard_group, + make_page_shard_spec, +) +from sglang.srt.mem_cache.page_interleave_pool import ( + PageInterleaveKVPoolMixin, + PageInterleaveMHATokenToKVPool, + PageInterleaveMLATokenToKVPool, +) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import ( @@ -219,6 +230,7 @@ def _pp_local_per_request_bytes( if TYPE_CHECKING: + from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.mem_cache.unified_memory_pool import ( UnifiedKVPool, UnifiedPoolBundle, @@ -301,6 +313,10 @@ class KVCacheConfigurator: mambaish_config: Optional[Any] = field(init=False) hybrid_gdn_config: Optional[Any] = field(init=False) hybrid_kda_config: Optional[Any] = field(init=False) + kv_shard_rank: Optional[int] = field(init=False) + kv_shard_size: int = field(init=False) + kv_shard_spec: Optional[PageShardSpec] = field(init=False) + kv_shard_group: Optional[GroupCoordinator] = field(init=False) is_hybrid_swa_mtp_draft: bool = field(init=False) draft_swa_full_capacity: bool = field(init=False) @@ -308,6 +324,7 @@ def __post_init__(self) -> None: self.mambaish_config = mambaish_config(self.model_config) self.hybrid_gdn_config = hybrid_gdn_config(self.model_config) self.hybrid_kda_config = hybrid_kda_config(self.model_config) + self._resolve_kv_shard() self.is_hybrid_swa_mtp_draft = self.layer_info.is_hybrid_swa_mtp_draft self.draft_swa_full_capacity = self.is_hybrid_swa_mtp_draft and bool( self.layer_info.swa_attention_layer_ids @@ -341,6 +358,21 @@ def logical_token_capacity(self, *, max_total_num_tokens: int) -> int: max_total_num_tokens * get_parallel().attn_dcp_size // self.loc_space_scale ) + def _resolve_kv_shard(self) -> None: + """Resolve the logical-page KV sharding identity: ``(rank, size)`` of + the shard group, or ``(None, 1)`` when sharding is off. Consumed by the + pool builders and by the widened allocator branch.""" + self.kv_shard_spec = make_page_shard_spec(self) + self.kv_shard_rank = None + self.kv_shard_size = 1 + self.kv_shard_group = None + if self.kv_shard_spec is None: + return + + self.kv_shard_rank = self.kv_shard_spec.shard_rank + self.kv_shard_size = self.kv_shard_spec.shard_size + self.kv_shard_group = get_kv_shard_group(self.use_mla_backend) + def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): return None @@ -632,6 +664,13 @@ def _init_pools( req_to_token_pool=req_to_token_pool, ) + if self.kv_shard_rank is not None: + if not isinstance(token_to_kv_pool, PageInterleaveKVPoolMixin): + raise ValueError( + "--enable-kv-cache-sharding does not support this model " + f"family (pool {type(token_to_kv_pool).__name__})." + ) + if draft_virtual_id_space is not None: assert token_to_kv_pool.size >= draft_virtual_id_space, ( "draft token_to_kv_pool smaller than the shared unified " @@ -1770,7 +1809,13 @@ def _build_mla_fp4_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: return token_to_kv_pool def _build_mla_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: - token_to_kv_pool = MLATokenToKVPool( + mla_pool_cls = MLATokenToKVPool + mla_pool_kwargs = {} + if self.kv_shard_rank is not None: + mla_pool_cls = PageInterleaveMLATokenToKVPool + mla_pool_kwargs["shard_spec"] = self.kv_shard_spec + mla_pool_kwargs["shard_group"] = self.kv_shard_group + token_to_kv_pool = mla_pool_cls( max_total_num_tokens, page_size=self.pool_page_size, dtype=self.kv_cache_dtype, @@ -1781,6 +1826,7 @@ def _build_mla_kv_pool(self, *, max_total_num_tokens: int) -> KVCache: enable_memory_saver=get_exec().features.enable_memory_saver, start_layer=self.layer_info.start_layer, end_layer=self.layer_info.end_layer, + **mla_pool_kwargs, ) return token_to_kv_pool @@ -2006,6 +2052,20 @@ def _build_mha_kv_pool( pool_kwargs["quant_method"] = quant_method else: pool_kwargs["post_capture_active"] = self.post_capture_kv_active + if self.kv_shard_rank is not None: + # The striped pool subclasses the plain per-layer MHA pool; the + # mxfp8 / NoOp / page-major variants have incompatible buffer + # layouts, and a quantized KV cache has no striped store path. + assert pool_cls is MHATokenToKVPool, ( + "--enable-kv-cache-sharding is incompatible with " + f"the {pool_cls.__name__} pool" + ) + assert quant_method is None, ( + "--enable-kv-cache-sharding is incompatible with a quantized KV cache" + ) + pool_cls = PageInterleaveMHATokenToKVPool + pool_kwargs["shard_spec"] = self.kv_shard_spec + pool_kwargs["shard_group"] = self.kv_shard_group token_to_kv_pool = pool_cls( max_total_num_tokens, page_size=self.pool_page_size, @@ -2140,6 +2200,19 @@ def _build_token_to_kv_pool_allocator( need_sort=need_sort, host_to_device_ratio=hisparse_cfg.host_to_device_ratio, ) + elif self.kv_shard_size > 1: + # Logical-page sharding: index space widened xN over + # the stock 1x pool; the allocator page stays physical. + token_to_kv_pool_allocator = PageInterleavePoolAllocator( + sizes.max_total_num_tokens, + physical_page_size=self.page_size, + shard_size=self.kv_shard_size, + dtype=self.kv_cache_dtype, + device=self.device, + kvcache=token_to_kv_pool, + need_sort=need_sort, + shard_spec=self.kv_shard_spec, + ) elif ( get_schedule().page_size == 1 and not get_parallel().dcp_enabled ): diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 638db7fb85ae..b5ca9a9b4161 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -17,7 +17,12 @@ from sglang.srt.hardware_backend.mlx.runtime import use_mlx from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.runtime_context import get_disagg, get_memory, get_serving +from sglang.srt.runtime_context import ( + get_disagg, + get_memory, + get_parallel, + get_serving, +) if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig @@ -106,6 +111,10 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: return SWAChunkCache(params) + if get_parallel().enable_kv_cache_sharding: + logger.info("Using UnifiedRadixCache for logical-page KV sharding.") + return create_unified_radix_cache(ctx) + if get_memory().enable_lmcache: from sglang.srt.mem_cache.storage.lmcache.lmcache_unified_radix_cache import ( LMCacheUnifiedRadixCache, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 3986089f9840..62e59e2e931f 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -42,6 +42,7 @@ DSATokenToKVPool, get_minimax_sparse_index_dtype, ) +from sglang.srt.mem_cache.page_interleave import compute_page_shard_scratch_bytes from sglang.srt.runtime_context import ( get_disagg, get_exec, @@ -213,7 +214,10 @@ def validate_swa_pool_size( class DefaultPoolConfigurator(MemoryPoolConfigurator): - """Standard models (MHA, MLA, DSA, FP4): coeff = cell_size, bias = 0.""" + """Standard models (MHA, MLA, DSA, FP4). + + coeff = cell_size; bias = double-buffered KV-shard scratch, otherwise zero. + """ def __init__(self, kvc: KVCacheConfigurator): self.kv_cache_dtype_str = kvc.kv_cache_dtype_str @@ -242,6 +246,10 @@ def __init__(self, kvc: KVCacheConfigurator): else get_schedule().max_total_tokens or kvc.model_config.context_len ) + # Logical-page KV sharding reserves a fixed double-buffered assembly + # scratch next to the pool; charge it before token sizing. + self._fixed_overhead_bytes = compute_page_shard_scratch_bytes(kvc) + # EAGLE/STANDALONE: assumes the draft shares the target's per-layer KV size, # which holds for EAGLE/MTP drafts that reuse the target's attention config. if ( @@ -592,7 +600,7 @@ def _compute_dsa_indexer_cell_size( def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - available_bytes = max(available_bytes, 0) + available_bytes = max(available_bytes - self._fixed_overhead_bytes, 0) max_total_num_tokens = ( available_bytes // self._cell_size if self._cell_size diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index a9c4948d9bef..37a16bc7f9d0 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/test/registered/dcp/test_dcp_layout_unit.py @@ -442,6 +442,8 @@ def test_configurator_scales_only_the_virtual_dcp_allocator(self): kv_cache_dtype=torch.bfloat16, device="cpu", is_draft_worker=False, + kv_shard_size=1, + kv_shard_spec=None, ) # The allocator widens from get_parallel(), not from the injected # server_args stand-in -- drive the cause, not the effect. diff --git a/test/registered/unit/disaggregation/test_mooncake_efa_allocator.py b/test/registered/unit/disaggregation/test_mooncake_efa_allocator.py index 2ed4793bad56..4d2b17e60489 100644 --- a/test/registered/unit/disaggregation/test_mooncake_efa_allocator.py +++ b/test/registered/unit/disaggregation/test_mooncake_efa_allocator.py @@ -75,6 +75,7 @@ def test_non_efa_keeps_custom_memory_pool_behavior(self): def test_efa_disables_post_capture_kv_sizing(self): cfg = SimpleNamespace( + enable_kv_cache_sharding=False, enable_unified_memory=False, device="cuda", dcp_size=1, diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 39d434be01b4..1ed0609333a6 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -45,6 +45,7 @@ def _make_ctx( disable_radix_cache=False, effective_chunked_prefill_size=None, full_tokens_per_layer=None, + enable_kv_cache_sharding=False, ): # The factory reads the published bags for the cache-backend leaves, so the # fixture publishes them; the instance stays for the whole-object contract @@ -56,6 +57,7 @@ def _make_ctx( enable_lmcache=enable_lmcache, enable_flexkv=False, enable_unified_cache_external_linker=False, + enable_kv_cache_sharding=enable_kv_cache_sharding, ) return TreeCacheBuildContext( server_args=server_args, @@ -158,6 +160,20 @@ def test_chunk_cache_when_chunked_prefill_and_disable_radix(self): ChunkCache.assert_called_once_with(ctx.params) self.assertIs(result, ChunkCache.return_value) + def test_kv_sharding_keeps_chunk_cache_when_radix_is_disabled(self): + ctx = _make_ctx( + self, + effective_chunked_prefill_size=512, + disable_radix_cache=True, + enable_kv_cache_sharding=True, + ) + with patch("sglang.srt.mem_cache.chunk_cache.ChunkCache") as ChunkCache: + ChunkCache.return_value = MagicMock() + result = default_radix_cache_factory(ctx) + + ChunkCache.assert_called_once_with(ctx.params) + self.assertIs(result, ChunkCache.return_value) + def test_swa_chunk_cache_when_chunked_prefill_disable_and_hybrid_swa(self): ctx = _make_ctx( self, @@ -230,6 +246,22 @@ def test_mamba_rejected_on_cache_without_mamba(self): with self.assertRaisesRegex(NotImplementedError, "not verified"): create_tree_cache(_make_ctx(self, backend="nomamba", is_hybrid_ssm=True)) + def test_kv_sharding_uses_unified_radix_cache(self): + ctx = _make_ctx(self, enable_kv_cache_sharding=True) + fake_components = MagicMock() + fake_radix = MagicMock() + with patch.dict( + "sys.modules", + { + "sglang.srt.mem_cache.unified_cache.components": fake_components, + "sglang.srt.mem_cache.unified_radix_cache": fake_radix, + }, + ): + result = default_radix_cache_factory(ctx) + + fake_radix.UnifiedRadixCache.assert_called_once_with(ctx.params) + self.assertIs(result, fake_radix.UnifiedRadixCache.return_value) + def test_unified_radix_cache_is_the_default(self): ctx = _make_ctx( self, diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3fa29e5ce0c3..9dea37c10631 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -6,7 +6,12 @@ from unittest.mock import MagicMock, patch from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.mem_cache.page_interleave import ( + compute_page_shard_scratch_bytes, + make_page_shard_spec, +) from sglang.srt.runtime_context import ( + get_context, get_memory, get_parallel, get_schedule, @@ -246,6 +251,75 @@ def test_page_alignment(self): _, _, config = self._run(available, page_size=128) self.assertEqual(config.max_total_num_tokens % 128, 0) + def test_eight_context_scratch_is_reserved_before_kv_pool_sizing(self): + page_size = 16 + mr = _make_model_runner( + self, + use_mla_backend=True, + num_layers=3, + page_size=page_size, + chunked_prefill_size=1024, + max_running_requests=1, + ) + mr.model_config.context_len = 8193 + with ( + mock_cpu_env(), + patch( + "sglang.srt.mem_cache.page_interleave.get_kv_shard_group_info", + return_value=(0, 4), + ), + get_context().override_server_args( + prefill_max_requests=1, + chunked_prefill_size=1024, + max_running_requests=1, + ), + ): + spec = make_page_shard_spec(mr) + # Each context rounds to 8256 tokens, a 64-token gather multiple, + # even with request limits of one. Only the prefix is multiplied. + self.assertEqual(spec.max_prefix_tokens, 8 * 8256) + self.assertEqual(spec.chunk_tokens, 1024) + row_bytes = (512 + 64) * KV_SIZE + scratch_bytes = 2 * (8 * 8256 + 1024 + page_size) * row_bytes + self.assertEqual(compute_page_shard_scratch_bytes(mr), scratch_bytes) + + # Scratch has two one-layer slots; persistent KV spans all three + # layers. Charge the former before dividing by the latter's cost. + cell_bytes = row_bytes * mr.num_effective_layers + page_bytes = page_size * cell_bytes + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + configurator = create_memory_pool_configurator(mr) + for desired_tokens in (page_size, 128): + for slack_bytes in (0, page_bytes - 1): + with self.subTest( + desired_tokens=desired_tokens, slack_bytes=slack_bytes + ): + budget = ( + scratch_bytes + desired_tokens * cell_bytes + slack_bytes + ) + config = configurator.calculate_pool_sizes(budget, page_size) + self.assertEqual(config.max_total_num_tokens, desired_tokens) + self.assertEqual( + budget + - scratch_bytes + - config.max_total_num_tokens * cell_bytes, + slack_bytes, + ) + + # The fixed eight-context scratch is never silently reduced to + # leave room for persistent KV, including a sub-page remainder. + for budget in ( + scratch_bytes - 1, + scratch_bytes, + scratch_bytes + page_bytes - 1, + ): + with self.subTest(insufficient_budget=budget): + with self.assertRaisesRegex(RuntimeError, "Not enough memory"): + configurator.calculate_pool_sizes(budget, page_size) + def test_constraint_respected(self): """calculate_pool_sizes_from_max_tokens respects the limit.""" mr, cfg, config = self._run(10_000_000) diff --git a/test/registered/unit/models/test_deepseek_mla_dispatch.py b/test/registered/unit/models/test_deepseek_mla_dispatch.py index dbd144857485..91173aa933bf 100644 --- a/test/registered/unit/models/test_deepseek_mla_dispatch.py +++ b/test/registered/unit/models/test_deepseek_mla_dispatch.py @@ -1,4 +1,4 @@ -"""Hermetic unit tests for DeepSeek MLA attention-method dispatch on ROCm. +"""Hermetic unit tests for DeepSeek MLA attention-method dispatch. `_dispatch_mla_subtype` picks the forward method for MLA attention. On ROCm the fused-decode-MLA + fused-RoPE fast path (`MLA_FUSED_ROPE_ROCM`) is only correct @@ -142,5 +142,111 @@ def test_strategy_cp_uses_absorbed_mla_without_legacy_flags(self): ) +class TestKVShardedMLADispatch(CustomTestCase): + def setUp(self): + super().setUp() + self.attn = SimpleNamespace( + chunked_prefix_cache_threshold=0, + disable_chunked_prefix_cache=False, + flashinfer_mla_disable_ragged=False, + ) + self.exec_config = SimpleNamespace( + deterministic=SimpleNamespace(enable_deterministic_inference=False) + ) + for patcher in ( + mock.patch.object(abh, "_is_hip", False), + mock.patch.object( + abh, + "get_parallel", + return_value=SimpleNamespace(enable_kv_cache_sharding=True), + ), + mock.patch.object(abh, "get_exec", return_value=self.exec_config), + mock.patch.object(abh, "is_in_tc_piecewise_cuda_graph", return_value=False), + mock.patch.object(abh, "is_in_breakable_cuda_graph", return_value=False), + mock.patch.object(cp_base, "_STRATEGY", None), + ): + patcher.start() + self.addCleanup(patcher.stop) + + def _batch(self, prefix=64, capacity=8192): + num_tokens = 3952 + return SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + input_ids=range(num_tokens), + attn_cp_metadata=None, + extend_prefix_lens_cpu=[prefix], + extend_seq_lens_cpu=[num_tokens], + seq_lens_cpu=[prefix + num_tokens], + get_max_chunk_capacity=lambda: capacity, + ) + + def test_fa3_disabled_chunked_prefix_cache_uses_absorbed_mla(self): + # A cached prefix or the second prompt chunk must not enter FA3's + # chunked-prefix path, which asserts prefix caching is enabled. + self.attn.disable_chunked_prefix_cache = True + self.assertEqual( + abh.handle_attention_fa3(self.attn, self._batch()), + AttnForwardMethod.MLA, + ) + + def test_fa3_deterministic_inference_uses_absorbed_mla(self): + self.exec_config.deterministic.enable_deterministic_inference = True + for prefix in (0, 64): + with self.subTest(prefix=prefix): + self.assertEqual( + abh.handle_attention_fa3(self.attn, self._batch(prefix=prefix)), + AttnForwardMethod.MLA, + ) + + def test_fa3_preserves_capacity_based_mha_dispatch(self): + # Both readers can use sharded scratch. Retain one-shot attention + # when the complete sequence fits, and chunked MHA otherwise. + for prefix in (0, 64): + for capacity, expected in ( + (8192, AttnForwardMethod.MHA_ONE_SHOT), + (0, AttnForwardMethod.MHA_CHUNKED_KV), + ): + with self.subTest(prefix=prefix, capacity=capacity): + self.assertEqual( + abh.handle_attention_fa3( + self.attn, self._batch(prefix=prefix, capacity=capacity) + ), + expected, + ) + + def test_fa3_short_prefix_uses_absorbed_mla(self): + self.attn.chunked_prefix_cache_threshold = 4096 + self.assertEqual( + abh.handle_attention_fa3(self.attn, self._batch()), + AttnForwardMethod.MLA, + ) + + def test_fa3_active_cp_uses_absorbed_mla(self): + # CP must gather rank-local latent activations before the pool write. + with mock.patch.object(cp_base, "_STRATEGY", ZigzagCPStrategy(cp_size=4)): + for prefix in (0, 64): + for capacity in (0, 8192): + with self.subTest(prefix=prefix, capacity=capacity): + self.assertEqual( + abh.handle_attention_fa3( + self.attn, + self._batch(prefix=prefix, capacity=capacity), + ), + AttnForwardMethod.MLA, + ) + + def test_trtllm_mla_keeps_chunked_mha_for_prefill(self): + # TRTLLM has a separate dispatcher and requires prefix caching for + # sharding; changing FA3 dispatch must not change its supported path. + for prefix in (0, 64): + with self.subTest(prefix=prefix): + self.assertEqual( + abh.handle_attention_trtllm_mla( + self.attn, self._batch(prefix=prefix) + ), + AttnForwardMethod.MHA_CHUNKED_KV, + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 0a4be745f22b..beb9bcd29398 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -43,6 +43,10 @@ handle_nvfp4_prefill_kv_dequant_dtype, validate_prefill_only_disable_kv_cache_args, ) +from sglang.srt.arg_groups.kv_shard_hook import ( + handle_kv_cache_sharding, + validate_kv_shard_attention_backend, +) from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend from sglang.srt.arg_groups.memory_hook import handle_gpu_memory_settings from sglang.srt.arg_groups.model_path_hook import handle_load_format @@ -52,7 +56,9 @@ validate_deepep_v2_speculative_draft, ) from sglang.srt.arg_groups.overrides import ( + declare_resolution, max_speculative_num_draft_tokens, + post_capture_kv_sizing_planned, resolution_result, ) from sglang.srt.arg_groups.parallel_hook import ( @@ -76,6 +82,7 @@ check_pipeline_parallel_compat, check_two_batch_overlap, ) +from sglang.srt.configs.model_config import AttentionArch from sglang.srt.entrypoints.sidecar import ( SGLANG_GRPC_ENDPOINT_ENV, Sidecar, @@ -3749,6 +3756,212 @@ def test_ep_a2a_backend_ok_without_dp_attention(self): check_two_batch_overlap(args) +class TestKvCacheShardingCompatibility(CustomTestCase): + def _args(self, **overrides): + model_config = overrides.pop("model_config", None) + args = ServerArgs( + model_path="dummy", + enable_kv_cache_sharding=True, + disaggregation_mode="prefill", + ) + for key, value in overrides.items(): + setattr(args, key, value) + args._model_config = model_config or SimpleNamespace( + is_encoder_decoder=False, + attention_chunk_size=None, + attention_arch=AttentionArch.MHA, + hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]), + ) + return args + + def _rounding_args(self, *, raw_mem_fraction, resolved_mem_fraction): + args = self._args( + tp_size=2, + page_size=64, + chunked_prefill_size=2050, + attention_backend="fa3", + disaggregation_transfer_backend="mooncake", + mem_fraction_static=resolved_mem_fraction, + cuda_graph_config=CudaGraphConfig( + prefill=PhaseConfig(backend=Backend.DISABLED) + ), + ) + args._model_config = SimpleNamespace( + is_encoder_decoder=False, + attention_chunk_size=None, + attention_arch=AttentionArch.MLA, + hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]), + ) + args._raw_input = {"mem_fraction_static": raw_mem_fraction} + args._resolved_overrides = [ + ( + "_handle_gpu_memory_settings", + {"mem_fraction_static": resolved_mem_fraction}, + ) + ] + return args + + def test_encoder_decoder_is_rejected_before_backend_setup(self): + args = self._args(model_config=SimpleNamespace(is_encoder_decoder=True)) + + with self.assertRaisesRegex(ValueError, "does not support encoder-decoder"): + handle_kv_cache_sharding(args) + + def test_trtllm_mla_accepts_plain_tp_with_chunked_prefix_cache(self): + args = self._rounding_args(raw_mem_fraction=0.8, resolved_mem_fraction=0.8) + args.attention_backend = "trtllm_mla" + + with ( + patch.object(envs.SGLANG_DISAGG_STAGING_BUFFER, "get", return_value=False), + ): + handle_kv_cache_sharding(args, 80 * 1024) + + self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2176) + + def test_trtllm_mla_rejects_unsupported_shard_topologies(self): + cases = ( + (False, 1, False, "requires an MLA model"), + (True, 2, False, "only supports plain-TP MLA"), + (True, 1, True, "requires chunked prefix caching"), + ) + for is_mla, attn_cp_size, disable_chunked_prefix_cache, error in cases: + with self.subTest(error=error): + args = self._args() + args.attention_backend = "trtllm_mla" + args._model_config.attention_arch = ( + AttentionArch.MLA if is_mla else AttentionArch.MHA + ) + view = SimpleNamespace( + attn_cp_size=attn_cp_size, + disable_chunked_prefix_cache=disable_chunked_prefix_cache, + ) + + with self.assertRaisesRegex(ValueError, error): + validate_kv_shard_attention_backend(args, view) + + def test_hisparse_allocator_is_rejected(self): + args = self._args(enable_hisparse=True) + + with self.assertRaisesRegex(ValueError, "does not support --enable-hisparse"): + handle_kv_cache_sharding(args) + + def test_incompatible_cache_tree_modes_are_rejected(self): + cases = ( + ({"radix_cache_backend": "custom"}, "built-in UnifiedRadixCache"), + ( + {"enable_unified_cache_external_linker": True}, + "does not support --enable-unified-cache-external-linker", + ), + ( + {"enable_streaming_session": True}, + "does not support streaming sessions", + ), + ( + {"enable_session_radix_cache": True}, + "does not support session radix cache yet", + ), + ) + for overrides, error in cases: + with self.subTest(overrides=overrides): + with self.assertRaisesRegex(ValueError, error): + handle_kv_cache_sharding(self._args(**overrides)) + + def test_rust_unified_tree_core_is_rejected(self): + with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("rust"): + with self.assertRaisesRegex(ValueError, "requires the Python"): + handle_kv_cache_sharding(self._args()) + + def test_post_capture_kv_sizing_keeps_eager_prefill_headroom(self): + prefill_graph = SimpleNamespace(backend=Backend.BREAKABLE, bs=[4096]) + args = ServerArgs( + model_path="dummy", + enable_kv_cache_sharding=True, + device="cuda", + dcp_size=1, + kv_cache_dtype="auto", + prefill_only_disable_kv_cache=False, + enable_memory_saver=False, + disaggregation_mode="prefill", + cuda_graph_config=CudaGraphConfig(prefill=prefill_graph), + chunked_prefill_size=4096, + ) + with ( + patch( + "sglang.srt.arg_groups.overrides.use_mla_backend", + return_value=False, + ), + patch( + "sglang.srt.arg_groups.overrides.max_prefill_buffer_tokens", + return_value=4096, + ), + patch( + "sglang.srt.arg_groups.overrides.model_config_of", + return_value=SimpleNamespace( + hf_config=SimpleNamespace(architectures=["LlamaForCausalLM"]) + ), + ), + patch.object( + envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING, + "get", + return_value=True, + ), + patch.object( + envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL, + "get", + return_value=None, + ), + ): + self.assertFalse(post_capture_kv_sizing_planned(args)) + args.enable_kv_cache_sharding = False + self.assertTrue(post_capture_kv_sizing_planned(args)) + + def test_chunk_rounding_expands_automatic_activation_headroom(self): + args = self._rounding_args(raw_mem_fraction=None, resolved_mem_fraction=0.9) + gpu_mem = 80 * 1024 + + with ( + patch.object(envs.SGLANG_DISAGG_STAGING_BUFFER, "get", return_value=False), + ): + handle_kv_cache_sharding(args, gpu_mem) + + rounded = 2176 # ceil(2050 / (tp_size * page_size)) * 128 + expected_fraction = 0.9 - 1.5 * (rounded - 2050) / gpu_mem + self.assertEqual(resolution_result(args, "chunked_prefill_size"), rounded) + self.assertAlmostEqual( + resolution_result(args, "mem_fraction_static"), expected_fraction + ) + + def test_chunk_rounding_preserves_explicit_mem_fraction(self): + args = self._rounding_args(raw_mem_fraction=0.8, resolved_mem_fraction=0.8) + + with ( + patch.object(envs.SGLANG_DISAGG_STAGING_BUFFER, "get", return_value=False), + ): + handle_kv_cache_sharding(args, 80 * 1024) + + self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2176) + self.assertEqual(resolution_result(args, "mem_fraction_static"), 0.8) + + def test_final_pass_rejects_late_model_capability_overrides(self): + cases = ( + ({"attention_backend": "triton"}, "requires the fa3"), + ({"chunked_prefill_size": -1}, "requires chunked prefill"), + ) + with ( + patch.object(envs.SGLANG_DISAGG_STAGING_BUFFER, "get", return_value=False), + ): + for declarations, error in cases: + with self.subTest(declarations=declarations): + args = self._rounding_args( + raw_mem_fraction=0.8, resolved_mem_fraction=0.8 + ) + handle_kv_cache_sharding(args, 80 * 1024) + declare_resolution(args, "late_model_capability", **declarations) + + with self.assertRaisesRegex(ValueError, error): + handle_kv_cache_sharding(args, 80 * 1024) + + class TestDcpKvEventContract(CustomTestCase): """DCP widens the radix-tree page to page_size * dcp_size, which the advertised KV-event block size must reflect.""" From 6d6a01988ea49c34ac40ef1bdc3e2027c91f7d4b Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Sun, 27 Sep 2026 23:50:17 -0700 Subject: [PATCH 02/11] Narrow KV-shard configuration guards (cherry picked from commit f025d206f3567cad2c7c04b014fc538a7afa2707) --- python/sglang/srt/arg_groups/kv_shard_hook.py | 40 +++++++--------- .../srt/mem_cache/kv_cache_configurator.py | 7 ++- .../sglang/srt/mem_cache/page_interleave.py | 8 ++-- .../unit/mem_cache/test_registry.py | 6 ++- .../unit/server_args/test_server_args.py | 46 +++++++++++++++++-- 5 files changed, 75 insertions(+), 32 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_shard_hook.py b/python/sglang/srt/arg_groups/kv_shard_hook.py index d9c7747dfe25..b063e4f8afaa 100644 --- a/python/sglang/srt/arg_groups/kv_shard_hook.py +++ b/python/sglang/srt/arg_groups/kv_shard_hook.py @@ -22,19 +22,6 @@ logger = logging.getLogger(__name__) -def kv_shard_group_size(server_args: Any) -> int: - """Return the active KV-shard group size, or one when disabled.""" - cfg = resolving_view(server_args) - if not cfg.enable_kv_cache_sharding: - return 1 - if cfg.attn_cp_size > 1: - return cfg.attn_cp_size - if use_mla_backend(server_args): - # Plain-TP MLA is validated below, so attn_tp_size == tp_size here. - return cfg.tp_size - return 1 - - def validate_kv_shard_attention_backend(server_args: Any, cfg: Any = None) -> None: cfg = cfg or resolving_view(server_args) prefill_backend, _ = attention_backends_of(resolved_view(server_args)) @@ -93,7 +80,9 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) f"--radix-cache-backend={cfg.radix_cache_backend!r}." ) tree_core_backend = envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get() - if tree_core_backend != "python": + # With radix caching disabled, chunked prefill uses ChunkCache and keeps + # the rotation base on each request; it never constructs a tree core. + if not cfg.disable_radix_cache and tree_core_backend != "python": raise ValueError( "--enable-kv-cache-sharding requires the Python UnifiedTreeCore: " f"SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND={tree_core_backend!r} is " @@ -146,9 +135,13 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) "HiSparse selects a non-sharded allocator and assumes each rank " "owns complete KV pages." ) - if cfg.enable_dynamic_chunking: + # The dynamic sizer only runs with pipeline parallelism. Its profiling + # and predicted chunks can exceed the fixed shard-scratch capacity. + if cfg.enable_dynamic_chunking and cfg.pp_size > 1: raise ValueError( - "--enable-kv-cache-sharding requires fixed chunk sizes; disable " + "--enable-kv-cache-sharding does not support dynamic chunk sizing " + "with pipeline parallelism: profiling and predicted chunks can " + "exceed the fixed sharded scratch capacity. Disable " "--enable-dynamic-chunking." ) if cfg.enable_two_batch_overlap: @@ -223,7 +216,6 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) "transfer descriptor and gather-copy unit; page_size=1 " "degenerates both into per-token operations." ) - granule = kv_shard_group_size(server_args) * page_size chunk_size = cfg.chunked_prefill_size if chunk_size is None or chunk_size <= 0: raise ValueError( @@ -232,8 +224,10 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) ) declarations = {} - if chunk_size % granule != 0: - rounded = (chunk_size + granule - 1) // granule * granule + # Chunk boundaries must stay on physical pages. Uneven ownership across + # ranks is supported; only the prefix gather pads to a full-group span. + if chunk_size % page_size != 0: + rounded = (chunk_size + page_size - 1) // page_size * page_size declarations["chunked_prefill_size"] = rounded # The raw input, deliberately not the resolving view: the question is @@ -248,7 +242,7 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) raise ValueError( "Cannot safely auto-round --chunked-prefill-size for KV " "sharding without a known device-memory capacity; pass a " - f"multiple of shard_size * page_size ({granule})." + f"multiple of page_size ({page_size})." ) corrected_mem_fraction = ( float(cfg.mem_fraction_static) - activation_delta_mb / gpu_mem @@ -257,16 +251,16 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None) raise ValueError( "Auto-rounding --chunked-prefill-size for KV sharding " "would consume all automatically reserved GPU memory; " - f"pass a smaller multiple of {granule}." + f"pass a smaller multiple of {page_size}." ) declarations["mem_fraction_static"] = corrected_mem_fraction logger.warning( "Rounding --chunked-prefill-size up from %d to %d " - "(a multiple of shard_size * page_size = %d) for KV sharding.", + "(a multiple of page_size = %d) for KV sharding.", chunk_size, rounded, - granule, + page_size, ) if cfg.cuda_graph_config.prefill.backend != Backend.DISABLED: diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index ec72f5236cae..6e938f3c01ef 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2055,13 +2055,16 @@ def _build_mha_kv_pool( if self.kv_shard_rank is not None: # The striped pool subclasses the plain per-layer MHA pool; the # mxfp8 / NoOp / page-major variants have incompatible buffer - # layouts, and a quantized KV cache has no striped store path. + # layouts. Specialized quantization methods need packed-data and + # scale-aware sharded writes/gathers; ordinary CUDA FP8 uses the + # plain pool and has no specialized quant_method. assert pool_cls is MHATokenToKVPool, ( "--enable-kv-cache-sharding is incompatible with " f"the {pool_cls.__name__} pool" ) assert quant_method is None, ( - "--enable-kv-cache-sharding is incompatible with a quantized KV cache" + "--enable-kv-cache-sharding does not support specialized KV-cache " + "quantization methods (e.g. FP4); ordinary CUDA FP8 is supported." ) pool_cls = PageInterleaveMHATokenToKVPool pool_kwargs["shard_spec"] = self.kv_shard_spec diff --git a/python/sglang/srt/mem_cache/page_interleave.py b/python/sglang/srt/mem_cache/page_interleave.py index a9e5c837ddac..3764cca21548 100644 --- a/python/sglang/srt/mem_cache/page_interleave.py +++ b/python/sglang/srt/mem_cache/page_interleave.py @@ -68,9 +68,11 @@ def scratch_rows(self) -> int: @property def logical_page_size(self) -> int: - """The N*page_size span of one logical page. It bounds the assembly - scratch and rounds chunked_prefill_size; the allocator and the tree - both keep drawing and matching at the physical ``page_size``.""" + """The N*page_size span used to align prefix-gather scratch. + + Allocation, prefix matching, and chunk boundaries use the physical + ``page_size``. + """ return self.shard_size * self.page_size diff --git a/test/registered/unit/mem_cache/test_registry.py b/test/registered/unit/mem_cache/test_registry.py index 1ed0609333a6..5a8c5f1cf905 100644 --- a/test/registered/unit/mem_cache/test_registry.py +++ b/test/registered/unit/mem_cache/test_registry.py @@ -1,5 +1,6 @@ """Unit tests for the radix-cache registry, routing, and selection chain.""" +from sglang.srt.environ import envs from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci @@ -167,7 +168,10 @@ def test_kv_sharding_keeps_chunk_cache_when_radix_is_disabled(self): disable_radix_cache=True, enable_kv_cache_sharding=True, ) - with patch("sglang.srt.mem_cache.chunk_cache.ChunkCache") as ChunkCache: + with ( + envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("rust"), + patch("sglang.srt.mem_cache.chunk_cache.ChunkCache") as ChunkCache, + ): ChunkCache.return_value = MagicMock() result = default_radix_cache_factory(ctx) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index beb9bcd29398..0a5965eacf6b 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -3816,7 +3816,7 @@ def test_trtllm_mla_accepts_plain_tp_with_chunked_prefix_cache(self): ): handle_kv_cache_sharding(args, 80 * 1024) - self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2176) + self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2112) def test_trtllm_mla_rejects_unsupported_shard_topologies(self): cases = ( @@ -3871,6 +3871,37 @@ def test_rust_unified_tree_core_is_rejected(self): with self.assertRaisesRegex(ValueError, "requires the Python"): handle_kv_cache_sharding(self._args()) + def test_unused_rust_tree_core_is_allowed_with_radix_disabled(self): + args = self._rounding_args(raw_mem_fraction=0.8, resolved_mem_fraction=0.8) + args.disable_radix_cache = True + with ( + envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("rust"), + envs.SGLANG_DISAGG_STAGING_BUFFER.override(False), + ): + handle_kv_cache_sharding(args, 80 * 1024) + + self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2112) + + def test_dynamic_chunking_is_rejected_only_when_pipeline_parallelism_uses_it(self): + for pp_size in (1, 2): + with self.subTest(pp_size=pp_size): + args = self._rounding_args( + raw_mem_fraction=0.8, resolved_mem_fraction=0.8 + ) + args.enable_dynamic_chunking = True + args.pp_size = pp_size + with envs.SGLANG_DISAGG_STAGING_BUFFER.override(False): + if pp_size > 1: + with self.assertRaisesRegex( + ValueError, "dynamic chunk sizing with pipeline parallelism" + ): + handle_kv_cache_sharding(args, 80 * 1024) + else: + handle_kv_cache_sharding(args, 80 * 1024) + self.assertEqual( + resolution_result(args, "chunked_prefill_size"), 2112 + ) + def test_post_capture_kv_sizing_keeps_eager_prefill_headroom(self): prefill_graph = SimpleNamespace(backend=Backend.BREAKABLE, bs=[4096]) args = ServerArgs( @@ -3924,13 +3955,22 @@ def test_chunk_rounding_expands_automatic_activation_headroom(self): ): handle_kv_cache_sharding(args, gpu_mem) - rounded = 2176 # ceil(2050 / (tp_size * page_size)) * 128 + rounded = 2112 # ceil(2050 / page_size) * 64; 33 pages across two shards expected_fraction = 0.9 - 1.5 * (rounded - 2050) / gpu_mem self.assertEqual(resolution_result(args, "chunked_prefill_size"), rounded) self.assertAlmostEqual( resolution_result(args, "mem_fraction_static"), expected_fraction ) + def test_page_aligned_chunk_need_not_span_the_whole_shard_group(self): + args = self._rounding_args(raw_mem_fraction=None, resolved_mem_fraction=0.9) + args.chunked_prefill_size = 2112 # 33 pages, not divisible by TP size 2. + with envs.SGLANG_DISAGG_STAGING_BUFFER.override(False): + handle_kv_cache_sharding(args, 80 * 1024) + + self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2112) + self.assertEqual(resolution_result(args, "mem_fraction_static"), 0.9) + def test_chunk_rounding_preserves_explicit_mem_fraction(self): args = self._rounding_args(raw_mem_fraction=0.8, resolved_mem_fraction=0.8) @@ -3939,7 +3979,7 @@ def test_chunk_rounding_preserves_explicit_mem_fraction(self): ): handle_kv_cache_sharding(args, 80 * 1024) - self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2176) + self.assertEqual(resolution_result(args, "chunked_prefill_size"), 2112) self.assertEqual(resolution_result(args, "mem_fraction_static"), 0.8) def test_final_pass_rejects_late_model_capability_overrides(self): From 4b9939e3ab063eff0f1f2fc531fcffa9f8dffa8d Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 29 Sep 2026 19:43:01 -0700 Subject: [PATCH 03/11] Fix KV-sharding tests under the Rust tree-core default (cherry picked from commit 93c9783afae497e2cfc6de73f09eb56738d9485e) --- python/sglang/srt/mem_cache/registry.py | 5 ----- test/registered/unit/server_args/test_server_args.py | 8 ++++++++ 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index b5ca9a9b4161..24611808f45b 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -20,7 +20,6 @@ from sglang.srt.runtime_context import ( get_disagg, get_memory, - get_parallel, get_serving, ) @@ -111,10 +110,6 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: return SWAChunkCache(params) - if get_parallel().enable_kv_cache_sharding: - logger.info("Using UnifiedRadixCache for logical-page KV sharding.") - return create_unified_radix_cache(ctx) - if get_memory().enable_lmcache: from sglang.srt.mem_cache.storage.lmcache.lmcache_unified_radix_cache import ( LMCacheUnifiedRadixCache, diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 0a5965eacf6b..ad1572935269 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -114,6 +114,7 @@ from sglang.test.test_utils import ( DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, CustomTestCase, + enter_scope, ) register_cpu_ci(est_time=14, suite="base-a-test-cpu") @@ -3757,6 +3758,13 @@ def test_ep_a2a_backend_ok_without_dp_attention(self): class TestKvCacheShardingCompatibility(CustomTestCase): + def setUp(self): + super().setUp() + # Sharding currently requires Python, independent of the default backend. + enter_scope( + self, envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("python") + ) + def _args(self, **overrides): model_config = overrides.pop("model_config", None) args = ServerArgs( From 4c489523127cf88617c88a55d5572c39f68dfdd3 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:23:42 -0700 Subject: [PATCH 04/11] [kv-shard 4/4] P/D transfer: owned-page send filter, bootstrap, E2E tests Last of four PRs for logical-page KV cache sharding; stacked on "[kv-shard 3/4] Enable the feature". A sharded prefill rank holds only its stripe of every cached page, so each prefill rank ships the pages it owns and decode pulls from every rank in the shard group. UnifiedRadixCache preserves the logical Full-KV page ids and rotation invariant consumed by this tree-independent transfer path. Contents: - `disaggregation/utils.py`: add `filter_kv_indices_for_shard_rank`. Ownership comes from the value (`owner = logical_page % N`), independent of placement order. Owned entries go on the wire as `logical_page // N`, the owner's local page id. The contiguous destination slice becomes an explicit positions array for positional pairing. - `disaggregation/common/conn.py`: resolve and advertise the prefill shard identity; make decode pull from all CP ranks or register every prefill TP rank as a real sender; reject decode layouts that would require DCP relayout. - `disaggregation/nixl/conn.py`: emit a standalone notification when a rank owns no page in a non-final chunk, preventing a missing chunk id from stalling the receiver. - `disaggregation/prefill.py`: widen the bootstrap capacity bound, require a page-aligned decode-cached prefix, retain scheduler-stage metrics, and run busy invariant checks in both disaggregated prefill loops. - `disaggregation/mooncake/conn.py`: keep transfer metrics aligned with bytes actually sent when replicated state transfer is suppressed. Tests cover the send partition and wire ids, Mooncake transfer accounting, NIXL empty-chunk notification, and bootstrap flag propagation. The new 4-GPU MLA and 8-GPU GQA/CP disaggregated E2Es exercise both shard axes; the GQA case uses the canonical `--enable-prefill-cp --cp-strategy zigzag` configuration. (cherry picked from commit 9f6b07266f22a72c0ebe7d202168c87697e5e472) --- .../sglang/srt/disaggregation/common/conn.py | 99 ++++++++++++++++++- .../srt/disaggregation/mooncake/conn.py | 11 ++- python/sglang/srt/disaggregation/nixl/conn.py | 37 +++++-- python/sglang/srt/disaggregation/prefill.py | 31 +++++- python/sglang/srt/disaggregation/utils.py | 30 ++++++ .../models/test_kv_cache_sharding_gqa_cp.py | 93 +++++++++++++++++ .../e2e/models/test_kv_cache_sharding_mla.py | 81 +++++++++++++++ .../test_kv_transfer_replica_metric.py | 46 +++++++++ .../disaggregation/test_nixl_backend_basic.py | 12 +++ .../test_register_to_bootstrap.py | 12 ++- .../mem_cache/test_page_interleave_shard.py | 56 +++++++++++ 11 files changed, 486 insertions(+), 22 deletions(-) create mode 100644 test/registered/e2e/models/test_kv_cache_sharding_gqa_cp.py create mode 100644 test/registered/e2e/models/test_kv_cache_sharding_mla.py diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 371bf2d291f0..67243a5bdacb 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -36,6 +36,7 @@ from sglang.srt.disaggregation.utils import ( DisaggregationMode, filter_kv_indices_for_cp_rank, + filter_kv_indices_for_shard_rank, get_dsv41_spec_layout, ) from sglang.srt.environ import envs @@ -113,6 +114,9 @@ class PrefillServerInfo: enable_dsa_cache_layer_split: bool = False dsv41_spec_layout: Optional[dict] = None decode_allocation_policy: str = "early" + # Prefill stores each KV page on exactly one shard-group rank + # (logical-page KV sharding); decode must pull from every shard rank. + enable_kv_cache_sharding: bool = False # PD true-retraction rebootstrap: the prefill's HTTP API port. The decode # already knows the prefill host (the bootstrap_addr host), so it can POST @@ -139,6 +143,7 @@ def __post_init__(self): ) self.follow_bootstrap_room = bool(self.follow_bootstrap_room) self.enable_dsa_cache_layer_split = bool(self.enable_dsa_cache_layer_split) + self.enable_kv_cache_sharding = bool(self.enable_kv_cache_sharding) self.prefill_http_port = ( int(self.prefill_http_port) if self.prefill_http_port is not None else None ) @@ -330,6 +335,27 @@ def __init__( self.pp_size = get_parallel().pp_size self.pp_rank = self.kv_args.pp_rank self.local_ip = get_local_ip_auto() + # Logical-page KV sharding (prefill side): the group whose ranks each + # hold a disjoint 1/N of every page, chosen by topology (mirrors + # mem_cache.page_interleave.get_kv_shard_group): an active attn-CP + # group takes precedence; without CP, MLA shards across attn-TP. + # (0, 1) when sharding is off (including on decode workers). + if ( + get_parallel().enable_kv_cache_sharding + and disaggregation_mode == DisaggregationMode.PREFILL + ): + if self.attn_cp_size > 1: + self.kv_shard_rank = self.attn_cp_rank + self.kv_shard_size = self.attn_cp_size + elif self.is_mla_backend: + self.kv_shard_rank = self.attn_tp_rank + self.kv_shard_size = self.attn_tp_size + else: + self.kv_shard_rank = 0 + self.kv_shard_size = 1 + else: + self.kv_shard_rank = 0 + self.kv_shard_size = 1 cp_sharded_prefill = self.attn_cp_size > 1 and ( self.is_hybrid_mla_backend or get_parallel().enable_dsa_cache_layer_split ) @@ -342,6 +368,9 @@ def __init__( envs.SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER.get() or cp_sharded_prefill or hybrid_decode_pulls_all_ranks + # Under CP-axis sharding no rank holds the full KV, so every CP + # rank must participate in the transfer. + or (get_parallel().enable_kv_cache_sharding and self.attn_cp_size > 1) ) # bind zmq socket @@ -482,6 +511,16 @@ def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool: and dst_dcp_size > 1 and (self.is_mla_backend or self.is_hybrid_mla_backend) ): + if self.kv_shard_size > 1: + # The relayout path indexes chunks by kv_chunk.index_slice.start, + # but the shard send filter replaces that slice with a fancy + # index of owned positions. Fail here instead of deep in the + # transfer worker. + raise RuntimeError( + "KV cache sharding cannot serve a decode instance that " + f"requires DCP relayout (decode dcp_size={dst_dcp_size}). " + "Disable --enable-kv-cache-sharding or --dcp-size." + ) return True raise RuntimeError( @@ -1224,7 +1263,26 @@ def _resolve_rank_mapping(self, info: PrefillServerInfo) -> None: """Compute TP/CP/PP rank mapping and store on the PrefillServerInfo object. Deterministic for a given (bootstrap_addr, decode engine) pair.""" # TP rank mapping - if self.attn_tp_size == info.attn_tp_size: + if ( + info.enable_kv_cache_sharding + and self.is_mla_backend + and info.attn_cp_size == 1 + ): + # Logical-page KV sharding on the attn-TP axis (MLA prefill + # without CP — with CP active the shard axis is the CP group and + # the CP fan-in below handles it): every decode rank pulls its + # full latent KV from ALL prefill TP ranks (each sends the pages + # it owns, paired positionally), which needs equal TP sizes. + assert self.attn_tp_size == info.attn_tp_size, ( + "MLA KV cache sharding requires equal prefill/decode attn TP " + f"sizes, got prefill={info.attn_tp_size} decode={self.attn_tp_size}" + ) + target_tp_rank = self.kv_args.engine_rank % self.attn_tp_size + target_tp_ranks = list(range(info.attn_tp_size)) + # Every decode TP rank registers with every prefill TP rank. + required_dst_info_num = self.attn_tp_size + required_prefill_response_num = info.attn_tp_size + elif self.attn_tp_size == info.attn_tp_size: target_tp_rank = self.kv_args.engine_rank % self.attn_tp_size required_dst_info_num = 1 required_prefill_response_num = 1 @@ -1278,6 +1336,9 @@ def _resolve_rank_mapping(self, info: PrefillServerInfo) -> None: pull_from_all_cp_ranks = ( self.enable_all_cp_ranks_for_transfer or info.enable_dsa_cache_layer_split + # CP-axis KV sharding: each prefill CP rank holds a disjoint + # 1/N of every page, so decode must pull from all of them. + or info.enable_kv_cache_sharding ) if not pull_from_all_cp_ranks: # Only retrieve from prefill CP rank 0 when not using all ranks @@ -1368,6 +1429,7 @@ def register_to_bootstrap(self): "dsv41_spec_layout": self.dsv41_spec_layout, "load_balance_method": get_parallel().load_balance_method, "enable_dsa_cache_layer_split": get_parallel().enable_dsa_cache_layer_split, + "enable_kv_cache_sharding": get_parallel().enable_kv_cache_sharding, # Self-register the HTTP API port so the decode can derive the PD # retract rebootstrap /generate URL from bootstrap info instead of a # router-injected pd_rebootstrap_prefill_url. @@ -1910,7 +1972,17 @@ def _prepare_send_indices( self.curr_idx += len(kv_indices) is_last_chunk = self.curr_idx == self.num_kv_indices - if ( + if self.kv_mgr.kv_shard_size > 1: + # Logical-page KV sharding: send only the pages this rank owns + # (ownership derived from the logical page ids in kv_indices), + # paired with their canonical positions (an index array replaces + # the contiguous slice). + kv_indices, index_slice = filter_kv_indices_for_shard_rank( + self.kv_mgr, + kv_indices, + index_slice, + ) + elif ( self.kv_mgr.enable_all_cp_ranks_for_transfer and not get_parallel().enable_dsa_cache_layer_split ): @@ -2162,14 +2234,26 @@ def _setup_bootstrap_infos(self): target_pp_rank, ) if bootstrap_info is not None: - if self.kv_mgr.is_mla_backend: + # TP-axis KV sharding (MLA without CP) needs every + # prefill TP rank to be a real sender; CP-axis + # sharding keeps MLA's one-real-TP-rank rule (the + # pages fan in across CP ranks instead). Keep the + # prefill_info read behind the MLA test: only that + # branch consults it. + if self.kv_mgr.is_mla_backend and not ( + self.prefill_info.enable_kv_cache_sharding + and self.prefill_info.attn_cp_size == 1 + ): # For MLA: target_tp_rank is the selected real rank, others are dummy ranks bootstrap_info["is_dummy"] = not bool( target_tp_rank == self.target_tp_rank or self.target_tp_rank is None ) else: - # For non-MLA: all target_tp_ranks are selected real ranks + # For non-MLA — and for MLA under logical-page + # KV sharding, where every prefill TP rank + # sends the pages it owns — all + # target_tp_ranks are selected real ranks bootstrap_info["is_dummy"] = False logger.debug( f"Fetched bootstrap info: {bootstrap_info} for DP {self.prefill_dp_rank} CP {target_cp_rank} TP {target_tp_rank} PP {target_pp_rank}" @@ -2475,6 +2559,7 @@ def __init__(self, host: str, port: int): self.dsv41_spec_layout: Optional[dict] = None self.follow_bootstrap_room: Optional[bool] = None self.enable_dsa_cache_layer_split: Optional[bool] = None + self.enable_kv_cache_sharding: Optional[bool] = None self.prefill_http_port: Optional[int] = None self.decode_allocation_policy = "early" self.prefill_port_table: Dict[ @@ -2591,6 +2676,11 @@ async def _handle_route_put(self, request: web.Request): data.get("enable_dsa_cache_layer_split", False) ) + if self.enable_kv_cache_sharding is None: + self.enable_kv_cache_sharding = bool( + data.get("enable_kv_cache_sharding", False) + ) + if system_dp_size == 1: dp_group = attn_dp_rank else: @@ -2656,6 +2746,7 @@ async def _handle_route_get(self, request: web.Request): else True ), enable_dsa_cache_layer_split=bool(self.enable_dsa_cache_layer_split), + enable_kv_cache_sharding=bool(self.enable_kv_cache_sharding), prefill_http_port=self.prefill_http_port, decode_allocation_policy=self.decode_allocation_policy, ) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 8924a03569e4..664edb7ec4f6 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -3061,7 +3061,16 @@ def send( trace_ctx=self.trace_ctx.copy_for_thread(), wait_event=wait_event, ) - self._record_transfer_indices(kv_indices, state_indices) + # Nonzero CP ranks still enqueue the final chunk so their completion + # notification is emitted, but the worker suppresses replicated state + # from those ranks. Keep the transfer metric aligned with the bytes + # that are actually sent. + metric_state_indices = ( + None + if self.kv_mgr._should_skip_cp_replicated_state_transfer() + else state_indices + ) + self._record_transfer_indices(kv_indices, metric_state_indices) def poll(self) -> KVPoll: if self.conclude_state is None: diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 0612bed8829c..bff3eb211121 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -1302,9 +1302,17 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) dst_info = self.decode_kv_args_table[req.agent_name] decode_tp_size = dst_info.decode_tp_size + notif = ( + f"{req.room}_kv_{kv_chunk.chunk_id}" + f"_{int(kv_chunk.is_last_chunk)}_{self.transfer_source_rank}" + ) + # Skip KV RDMA transfer when there are no pages to send - # (e.g., decode-side radix cache matched the entire prefix). - # Aux data is still sent below when is_last_chunk=True. + # (decode-side radix cache matched the entire prefix, this + # rank owns no page of the chunk under KV cache sharding, or + # this rank holds no attention KV at all). The chunk id is + # still accounted for below, and aux data is still sent when + # is_last_chunk=True. if ( len(kv_chunk.prefill_kv_indices) > 0 and self.kv_args.kv_data_ptrs @@ -1333,11 +1341,6 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) src_prefill_kv_indices = kv_chunk.prefill_kv_indices - notif = ( - f"{req.room}_kv_{kv_chunk.chunk_id}" - f"_{int(kv_chunk.is_last_chunk)}_{self.transfer_source_rank}" - ) - # Decide which kv send path to use: # 1. Staging (heterogeneous TP, both sides have # registered staging, watermark/alloc ready) @@ -1449,6 +1452,18 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) if kv_xfer_handle is not None: handles.append(kv_xfer_handle) + elif not kv_chunk.is_last_chunk: + # Zero pages for this rank on a non-final chunk: deliver + # the chunk notif standalone. chunk_id advances whenever + # the request as a whole has pages to send, so a rank + # that happens to own none of this chunk's pages (KV + # cache sharding assigns them round-robin) still has to + # claim the id -- the receiver requires every id in + # [0, expected) from every source rank, and a silently + # skipped id hangs it until the transfer timeout. The + # final chunk instead reports its count through the aux + # "nokv" marker below. + self.agent.send_notif(req.agent_name, notif.encode("ascii")) if kv_chunk.is_last_chunk: dst_info = self.decode_kv_args_table[req.agent_name] @@ -1472,9 +1487,11 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) if kv_chunk.prefill_aux_index is None: raise RuntimeError("Missing aux index for last chunk") - # A no-KV notification still identifies its PP source. - # Empty non-final chunks do not consume chunk IDs, so a - # final no-KV chunk_id equals the prior KV chunk count. + # A no-KV notification still identifies its PP source and + # reports chunk_id as the expected KV-chunk count: the + # final chunk carried nothing, and every earlier id was + # claimed either by its RDMA transfer or by the + # standalone notif above. aux_notif = f"{req.room}_aux" if ( len(kv_chunk.prefill_kv_indices) == 0 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index df2974c4853b..a725341e90ec 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -190,6 +190,8 @@ def __init__( self.queue: List[Req] = [] self.scheduler = scheduler self.scheduler_stage_metrics = scheduler_stage_metrics + # Capacity bound in the allocator's (logical) token units — widened + # by the shard-group size under logical-page KV sharding. self.max_total_num_tokens = ( self.scheduler.tp_worker.model_runner.effective_logical_max_total_num_tokens ) @@ -416,6 +418,21 @@ def finalize_bootstrap(self, req: Req) -> bool: req.disagg_decode_prefix_len = decode_prefix_len num_pages = self._num_pages_to_send(req, decode_prefix_len) req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index) + if ( + get_parallel().enable_kv_cache_sharding + and self.kv_manager.kv_shard_size > 1 + ): + # Logical-page KV sharding: chunk starts must stay page-aligned — + # the sender samples one wire entry per physical page (stride + # ps), so an off-page start would misalign every sample. + # Ownership itself is value-derived (owner = logical page % N in + # filter_kv_indices_for_shard_rank), independent of where send + # position 0 sits. + physical_page_size = self.token_to_kv_pool.page_size + assert decode_prefix_len % physical_page_size == 0, ( + f"decode-cached prefix ({decode_prefix_len}) must be " + f"page-aligned under KV sharding" + ) req.pending_bootstrap = False return True @@ -726,6 +743,8 @@ def event_loop_normal_disagg_prefill(self: Scheduler) -> None: # Update last_batch self.last_batch = batch + if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): + self.invariant_checker.self_check_during_busy() def _is_continuous_input_polling_enabled(self: Scheduler) -> bool: """Check support once on entry to the prefill overlap loop.""" @@ -855,6 +874,8 @@ def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self._sched_idled = batch is None self.last_batch = batch waiting_for_batch_result = False + if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): + self.invariant_checker.self_check_during_busy() def is_disagg_prefill_batch_result_ready( self: Scheduler, @@ -1442,9 +1463,9 @@ def maybe_send_cached_prefix_chunk(self: Scheduler, req: Req) -> None: if cached_end <= req.start_send_idx: return if cached_end % self.token_to_kv_pool_allocator.page_size != 0: - # DCP radix hits can end on a logical cache-page boundary that is - # not a complete physical DCP page. The regular final send covers - # the full range; only skip this optional early-send optimization. + # Under DCP the allocator page is wider than the kernel page and + # a radix hit may end inside one. The regular final send covers the + # full range; only this optional early send is skipped. return # Early-send issues the KV read before this step's forward is enqueued, # but under overlap scheduling the PRIOR step's prefill forward may still @@ -1658,6 +1679,10 @@ def _request_state_payload(): raw_kv_indices ) ) + # Under logical-page KV sharding these remain logical page ids: + # PageInterleavePoolAllocator intentionally inherits the identity + # transfer translation. The sender derives ownership from each id + # and emits the owning rank's local page id. page_indices = kv_to_page_indices(kv_indices, page_size) segment_is_last = last_chunk and is_final_segment if not req.disagg_kv_sender.should_send_kv_chunk( diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index db8cb7a27fae..66c3ba743a87 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -694,6 +694,36 @@ def _get_cp_rank_page_bounds( return local_start, local_start + n_pages +def filter_kv_indices_for_shard_rank( + kv_mgr: CommonKVManager, + kv_indices: np.ndarray, + index_slice: slice, +) -> Tuple[np.ndarray, np.ndarray]: + """Ownership-driven page filter for logical-page KV cache sharding. + + ``kv_indices`` holds one entry per PHYSICAL page position of this chunk; + the entry value is the LOGICAL page id ``l = loc // ps``, identical on + every rank (mirrored ``req_to_token``). Ownership is derived from the + value — owner rank ``l % shard_size`` — so the filter is independent of + the allocator's placement policy. The owned entries are emitted as wire + values ``l // shard_size``: the owner's local page id in its own pool, so + the decode side needs no change. + + Returns the owned wire values and an int64 positions array (in the + request's canonical send-position space) replacing the contiguous + ``index_slice`` contract — decode-side dst pairing is positional and + consumes a fancy index the same way it consumes a slice. + """ + shard_rank = kv_mgr.kv_shard_rank + shard_size = kv_mgr.kv_shard_size + chunk_start = index_slice.start if index_slice.start is not None else 0 + chunk_end = index_slice.stop if index_slice.stop is not None else chunk_start + positions = np.arange(chunk_start, chunk_end, dtype=np.int64) + logical_pages = np.asarray(kv_indices) + owned = logical_pages % shard_size == shard_rank + return logical_pages[owned] // shard_size, positions[owned] + + def filter_kv_indices_for_cp_rank( kv_mgr: CommonKVManager, kv_indices: np.ndarray, diff --git a/test/registered/e2e/models/test_kv_cache_sharding_gqa_cp.py b/test/registered/e2e/models/test_kv_cache_sharding_gqa_cp.py new file mode 100644 index 000000000000..7f6f4685b9c5 --- /dev/null +++ b/test/registered/e2e/models/test_kv_cache_sharding_gqa_cp.py @@ -0,0 +1,93 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""E2E test for logical-page KV cache sharding on a GQA model under prefill CP. + +Prefill CP replicates GQA KV storage across attention-CP ranks; with +``--enable-kv-cache-sharding`` each CP rank stores only its page-interleaved +stripe (~CPx unique KV capacity for the group), extend attention reads the +assembled scratch through the translated FA3 page table, and each CP rank +sends only its owned pages to decode. + +PD-disaggregation fixture, 8 GPUs: prefill TP4 (attn-CP 2) on GPUs 0-3, +decode TP2 on GPUs 4-5. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) + +register_cuda_ci(est_time=900, stage="base-c", runner_config="8-gpu-h200") + + +class TestGQACPKVCacheSharding(PDDisaggregationServerBase, GSM8KMixin): + model = "Qwen/Qwen3-30B-A3B-FP8" + + gsm8k_accuracy_thres = 0.90 + gsm8k_num_questions = 200 + gsm8k_num_threads = 128 + + # Canonical zigzag prefill CP plus KV sharding across the attn-CP group. + extra_prefill_args = [ + "--tp", + "4", + "--moe-dp-size", + "2", + "--ep-size", + "2", + "--attn-cp-size", + "2", + "--enable-prefill-cp", + "--cp-strategy", + "zigzag", + "--attention-backend", + "fa3", + "--enable-kv-cache-sharding", + "--page-size", + "64", + "--chunked-prefill-size", + "4096", + "--max-prefill-tokens", + "4096", + "--mem-fraction-static", + "0.8", + ] + # Decode attn-TP matches the prefill attn-TP (tp 4 / cp 2 -> attn_tp 2): + # the owner-strided send filter pairs pages positionally per (tp, cp) + # sender; a TP-size mismatch would add non-MLA head slicing on top and + # fragment every page transfer. + extra_decode_args = [ + "--tp", + "2", + "--attention-backend", + "fa3", + "--page-size", + "64", + "--mem-fraction-static", + "0.8", + "--base-gpu-id", + "4", + ] + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.launch_all() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/e2e/models/test_kv_cache_sharding_mla.py b/test/registered/e2e/models/test_kv_cache_sharding_mla.py new file mode 100644 index 000000000000..baff92f41b3d --- /dev/null +++ b/test/registered/e2e/models/test_kv_cache_sharding_mla.py @@ -0,0 +1,81 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""E2E test for logical-page KV cache sharding on an MLA model. + +MLA latent KV is replicated across attention-TP ranks; with +``--enable-kv-cache-sharding`` each prefill TP rank stores only its +page-interleaved stripe (~TPx unique KV capacity for the group), reads the +prefix through the assembled scratch (chunked-prefix MHA path), and sends only +its owned pages to decode (which pulls from every prefill TP rank). + +Uses the PD-disaggregation fixture because KV sharding is a PD-prefill-only +feature. 4 GPUs: prefill TP2 on GPUs 0-1, decode TP2 on GPUs 2-3. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import DEFAULT_MLA_MODEL_NAME_FOR_TEST + +register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-h100") + + +class TestMLAKVCacheSharding(PDDisaggregationServerBase, GSM8KMixin): + model = DEFAULT_MLA_MODEL_NAME_FOR_TEST + + gsm8k_accuracy_thres = 0.70 + gsm8k_num_questions = 200 + gsm8k_num_threads = 128 + + extra_prefill_args = [ + "--tp", + "2", + "--attention-backend", + "fa3", + "--enable-kv-cache-sharding", + "--page-size", + "64", + "--chunked-prefill-size", + "2048", + "--max-prefill-tokens", + "2048", + "--mem-fraction-static", + "0.8", + ] + + extra_decode_args = [ + "--tp", + "2", + "--attention-backend", + "fa3", + "--page-size", + "64", + "--mem-fraction-static", + "0.8", + "--base-gpu-id", + "2", + ] + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.launch_all() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py b/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py index 75baa3f52abf..b962f8bbd38c 100644 --- a/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py +++ b/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py @@ -9,11 +9,13 @@ import unittest from types import SimpleNamespace +from unittest.mock import Mock import numpy as np from sglang.srt.disaggregation.base.conn import KVTransferMetric from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVSender +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVSender from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -70,6 +72,50 @@ def _make_sender(kv_mgr): class TestKVTransferReplicaMetric(CustomTestCase): + def test_mooncake_does_not_count_cp_replicated_state_that_is_skipped(self): + state_indices = [np.array([9], dtype=np.int32)] + + for cp_rank, expected_kv_count, expected_state_count in ( + (0, 2, 1), + (1, 0, 0), + ): + with self.subTest(cp_rank=cp_rank): + mgr = SimpleNamespace( + kv_shard_rank=cp_rank, + kv_shard_size=2, + add_transfer_request=Mock(), + _should_skip_cp_replicated_state_transfer=( + lambda cp_rank=cp_rank: cp_rank != 0 + ), + ) + sender = MooncakeKVSender.__new__(MooncakeKVSender) + sender.kv_mgr = mgr + sender.bootstrap_room = 17 + sender.aux_index = 3 + sender.curr_idx = 0 + sender.num_kv_indices = 2 + sender.trace_ctx = Mock() + sender._transfer_metric = KVTransferMetric() + sender._transfer_num_kv_indices = 0 + sender._transfer_num_state_indices = 0 + + # Both logical pages are owned by rank 0, so rank 1 must still + # enqueue an empty final chunk but must not count the replicated + # state that Mooncake's worker skips for nonzero CP ranks. + sender.send( + np.array([0, 2], dtype=np.int32), + state_indices=state_indices, + ) + + self.assertEqual(sender._transfer_num_kv_indices, expected_kv_count) + self.assertEqual( + sender._transfer_num_state_indices, expected_state_count + ) + mgr.add_transfer_request.assert_called_once() + call = mgr.add_transfer_request.call_args + self.assertTrue(call.args[3]) + self.assertEqual(len(call.args[1]), expected_kv_count) + def test_mla_scales_kv_and_state_bytes_by_fan_out(self): # CP rank 0 replicates to 4 decode TP ranks; both kv and state scale. mgr = _make_kv_mgr(is_mla_backend=True) diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 668237a44d31..aa568887acb2 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -791,6 +791,18 @@ def test_given_non_last_chunk_aborts_mid_transfer_when_worker_finishes_then_fail self.assertIn(room, mgr.req_to_decode_prefix_len) mgr.send_kvcache.assert_called_once() + def test_empty_non_last_chunk_sends_standalone_notification(self): + room = 24 + mgr = self._make_manager(room) + mgr.agent.send_notif = MagicMock() + mgr.send_kvcache = MagicMock() + chunk = self._make_chunk(room, [], is_last_chunk=False) + + self._run_worker_once(mgr, chunk) + + mgr.agent.send_notif.assert_called_once_with("agent", b"24_kv_0_0_0") + mgr.send_kvcache.assert_not_called() + def test_dcp_destinations_use_disjoint_pack_regions_before_chunk_barrier(self): room = 23 mgr = self._make_manager(room) diff --git a/test/registered/unit/disaggregation/test_register_to_bootstrap.py b/test/registered/unit/disaggregation/test_register_to_bootstrap.py index 322bca9e6659..5db76bc7a045 100644 --- a/test/registered/unit/disaggregation/test_register_to_bootstrap.py +++ b/test/registered/unit/disaggregation/test_register_to_bootstrap.py @@ -17,11 +17,13 @@ class TestRegisterToBootstrap(CustomTestCase): """Tests for CommonKVManager.register_to_bootstrap retry/backoff behavior.""" def setUp(self): - # register_to_bootstrap reads get_parallel().load_balance_method / - # .enable_dsa_cache_layer_split and get_serving().port from the - # published config. + # register_to_bootstrap reads get_parallel().load_balance_method, + # .enable_dsa_cache_layer_split, and .enable_kv_cache_sharding plus + # get_serving().port from the published config. override = get_context().override_server_args( - load_balance_method="follow_bootstrap_room", port=30000 + load_balance_method="follow_bootstrap_room", + port=30000, + enable_kv_cache_sharding=True, ) override.install() self.addCleanup(override.restore) @@ -231,12 +233,14 @@ def test_payload_contains_required_fields(self, mock_put, mock_time): "rank_port", "page_size", "kv_cache_dtype", + "enable_kv_cache_sharding", # Self-registered HTTP API port used to derive the PD retract # rebootstrap /generate URL on the decode side. "prefill_http_port", ] for field in required_fields: self.assertIn(field, payload) + self.assertIs(payload["enable_kv_cache_sharding"], True) self.assertEqual(payload["prefill_http_port"], 30000) @patch("sglang.srt.disaggregation.common.conn.time") diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index e0a8592b0267..2ae799d664bf 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -31,6 +31,8 @@ class draws (owners exactly cyclic along a chain), least-full root 5. ``begin_shard_extend`` plan capture (page positions, padded send rows, owner-congruence guard) with the gather stubbed out, following the SimpleNamespace binding pattern of ``test_dsa_layer_shard_utils.py``. +6. The value-derived P/D ownership send filter + ``filter_kv_indices_for_shard_rank``. The second section (``TestPageInterleaveGatherMultiGpu``, at the bottom) drives real pools over a real 2-rank process group. It is the only check that the plan @@ -47,10 +49,12 @@ class draws (owners exactly cyclic along a chain), least-full root from contextlib import contextmanager from types import SimpleNamespace +import numpy as np import torch import torch.multiprocessing as mp from parameterized import parameterized_class +from sglang.srt.disaggregation.utils import filter_kv_indices_for_shard_rank from sglang.srt.distributed import ( init_distributed_environment, initialize_model_parallel, @@ -1342,6 +1346,58 @@ def test_partial_tail_page_may_own_nothing(self): self.assertEqual(local_rows.numel(), 0) +class TestShardSendFilter(CustomTestCase): + """Value-derived P/D ownership partition: owner = logical page % N, wire + value = logical page // N — independent of positions, so it stays + correct under any placement policy.""" + + def test_partition_and_positional_pairing(self): + mgr = SimpleNamespace(kv_shard_rank=0, kv_shard_size=N) + # A base-2 rotated chain: owners are NOT position-congruent. + logical_pages = np.array(_chain_pages(base=2, n_pages=9), dtype=np.int64) + sl = slice(0, 9) + got = {} + for r in range(N): + mgr.kv_shard_rank = r + wire, pos = filter_kv_indices_for_shard_rank(mgr, logical_pages, sl) + got[r] = (wire, pos) + # Ownership by VALUE (l % N == r), wire id = owner's local page. + np.testing.assert_array_equal(logical_pages[pos] % N, r) + np.testing.assert_array_equal(wire, logical_pages[pos] // N) + # Disjoint cover of all positions. + all_pos = np.sort(np.concatenate([got[r][1] for r in range(N)])) + np.testing.assert_array_equal(all_pos, np.arange(9)) + + def test_wire_value_is_the_owner_local_page(self): + """The wire byte contract: l // N == loc // (N * ps) — the owner-local + page id, so the decode side needs no change.""" + mgr = SimpleNamespace(kv_shard_rank=1, kv_shard_size=N) + logical_pages = np.array(_chain_pages(base=0, n_pages=8), dtype=np.int64) + locs = logical_pages * PS # page-head loc of each entry + wire, pos = filter_kv_indices_for_shard_rank(mgr, logical_pages, slice(0, 8)) + np.testing.assert_array_equal(wire, locs[pos] // (N * PS)) + + def test_mid_request_chunk_positions_canonical(self): + mgr = SimpleNamespace(kv_shard_rank=2, kv_shard_size=N) + logical_pages = np.array( + _chain_pages(base=1, n_pages=32)[20:32], dtype=np.int64 + ) + wire, pos = filter_kv_indices_for_shard_rank(mgr, logical_pages, slice(20, 32)) + # Positions are canonical (chunk-global), values from this chunk. + self.assertTrue(((pos >= 20) & (pos < 32)).all()) + np.testing.assert_array_equal(logical_pages[pos - 20] % N, 2) + np.testing.assert_array_equal(wire, logical_pages[pos - 20] // N) + + def test_rank_owning_nothing_in_chunk(self): + """A short chunk can leave a rank with zero pages; the filter must + return empty arrays, not error.""" + mgr = SimpleNamespace(kv_shard_rank=3, kv_shard_size=N) + pages = np.array(_chain_pages(base=0, n_pages=2), dtype=np.int64) # owners 0,1 + wire, pos = filter_kv_indices_for_shard_rank(mgr, pages, slice(0, 2)) + self.assertEqual(len(wire), 0) + self.assertEqual(len(pos), 0) + + # ============================================================================= # Multi-GPU: the real NCCL layer-ahead gather (2 GPUs). # From 2d78f5c4991569bf6074da5f40ad1f94924b6919 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Wed, 23 Sep 2026 01:44:16 -0700 Subject: [PATCH 05/11] feat: support page-sharded DSA indexer KV cache (cherry picked from commit ee639a5a2b0eb4dd6e274e6b21dd7a645480fe41) --- .../srt/mem_cache/page_interleave_pool.py | 248 +++++++- .../mem_cache/test_page_interleave_shard.py | 542 +++++++++++++++++- 2 files changed, 768 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index 637a94bfb74e..b0a859536606 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -56,9 +56,15 @@ import torch +from sglang.kernels.ops.attention.dsa import index_buf_accessor +from sglang.kernels.ops.attention.dsa.dequant_k_cache import ( + dequantize_k_cache_paged, +) from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator +from sglang.srt.mem_cache.index_key_cache import IndexKeyCache from sglang.srt.mem_cache.memory_pool import ( GPU_MEMORY_TYPE_KV_CACHE, + DSATokenToKVPool, MHATokenToKVPool, MLATokenToKVPool, RadixAttention, @@ -68,10 +74,7 @@ PageInterleavePlacement, PageShardSpec, ) -from sglang.srt.mem_cache.utils import ( - get_mla_kv_buffer_triton, - set_mla_kv_buffer_triton, -) +from sglang.srt.mem_cache.utils import get_mla_kv_buffer_triton from sglang.srt.utils import ceil_div, get_bool_env_var if TYPE_CHECKING: @@ -166,6 +169,7 @@ def _init_page_shard_state( # sort key of logical page l is (l % N) * stride + l // N. self._local_page_stride = self.size // spec.page_size + 2 self._debug_plan_checks = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL") + self._send_pages: Optional[torch.Tensor] = None self._send_rows: Optional[torch.Tensor] = None self._write_plan_key = None self._write_plan: Optional[Tuple[torch.Tensor, torch.Tensor]] = None @@ -227,6 +231,10 @@ def _gather_pairs(self, local_layer: int) -> List[Tuple[torch.Tensor, str]]: """(pool buffer of ``local_layer``, scratch tensor name) pairs.""" raise NotImplementedError + def _gather_send_rows(self, name: str) -> torch.Tensor: + """Source indices in a gather pair's first dimension (tokens by default).""" + return self._send_rows + # ---- per-batch plan ------------------------------------------------------- def begin_shard_extend( @@ -355,12 +363,14 @@ def begin_shard_extend( torch.zeros((n_pad,), dtype=torch.int64, device=self.device), ] ) + self._send_pages = own_local self._send_rows = ( own_local[:, None] * ps + torch.arange(ps, dtype=torch.int64, device=self.device) ).reshape(-1) self._prefetch_layer(self.start_layer) else: + self._send_pages = None self._send_rows = None def end_shard_extend(self) -> None: @@ -395,16 +405,17 @@ def _prefetch_layer(self, layer_id: int) -> None: key = (layer_id, self._epoch) if slot.resident_key == key: return - block = self._block_pages * self.shard_spec.page_size # Order the gather after all prior compute-stream work: the # previous tenant's reads (attention of layer_id - 2) and the pool # writes that produced the prefix rows. self.kv_gather_stream.wait_stream(self.device_module.current_stream()) with self.device_module.stream(self.kv_gather_stream): for pool_buf, name in self._gather_pairs(local_layer): + send_rows = self._gather_send_rows(name) + block = send_rows.numel() scratch = slot.tensors[name] send = scratch[self.shard_rank * block : (self.shard_rank + 1) * block] - torch.index_select(pool_buf, 0, self._send_rows, out=send) + torch.index_select(pool_buf, 0, send_rows, out=send) # In-place regular allgather: send is exactly the rank's # block of the output, every rank contributes `block` rows. with self.kv_gather_comm.change_state(enable=True): @@ -668,14 +679,9 @@ def set_mla_kv_buffer( if self._shard_extend_active: slot = self._slots[layer_id % 2] rows = self._translate_loc_cached(loc) - staged_nope, staged_rope = cache_k_nope, cache_k_rope - if staged_nope.dtype != self.dtype: - staged_nope = staged_nope.to(self.dtype) - staged_rope = staged_rope.to(self.dtype) - if self.store_dtype != self.dtype: - staged_nope = staged_nope.view(self.store_dtype) - staged_rope = staged_rope.view(self.store_dtype) - set_mla_kv_buffer_triton(slot.tensors["kv"], rows, staged_nope, staged_rope) + self._write_mla_kv_buffer( + slot.tensors["kv"], rows, cache_k_nope, cache_k_rope + ) owned_idx, local_rows = self._get_write_plan(loc) if owned_idx.numel() == 0: return @@ -734,3 +740,217 @@ def get_mla_kv_buffer( ) get_mla_kv_buffer_triton(kv_buffer, rows, cache_k_nope, cache_k_rope) return cache_k_nope, cache_k_rope + + +class _PageInterleaveIndexKeyCache(IndexKeyCache): + """Packed indexer pages using the latent pool's ownership and scratch plan.""" + + def get_buffer(self, layer_id: int) -> torch.Tensor: + pool = self.pool + local = self.get_local_buffer(layer_id) + if pool._shard_extend_active: + slot = pool._acquire_slot_for_read(layer_id) + # Raw getters preserve the empty skip-topk placeholder, but still + # advance layer prefetch rather than exposing stale scratch bytes. + return slot.tensors["index_k"] if local.shape[0] else local + return local + + def _scratch_pages(self, layer_id: int, page_indices: torch.Tensor) -> torch.Tensor: + assert self.pool._shard_extend_active, ( + "logical indexer page reads require begin_shard_extend" + ) + assert self.buffer[layer_id - self.pool.start_layer].shape[0], ( + "cannot read indexer keys for a skip-topk layer" + ) + return self.pool._page_pos[page_indices.long()] + + def get_k_continuous(self, layer_id: int, seq_len: int, page_indices: torch.Tensor): + return super().get_k_continuous( + layer_id, seq_len, self._scratch_pages(layer_id, page_indices) + ) + + def get_k_scale_continuous( + self, layer_id: int, seq_len: int, page_indices: torch.Tensor + ): + return super().get_k_scale_continuous( + layer_id, seq_len, self._scratch_pages(layer_id, page_indices) + ) + + def get_k_and_scale( + self, + layer_id: int, + seq_len_tensor: torch.Tensor, + page_indices: torch.Tensor, + seq_len_sum: int, + max_seq_len: int, + ): + return super().get_k_and_scale( + layer_id, + seq_len_tensor, + self._scratch_pages(layer_id, page_indices), + seq_len_sum, + max_seq_len, + ) + + def store_quantized( + self, + layer_id: int, + loc: torch.Tensor, + index_k: torch.Tensor, + index_k_scale: torch.Tensor, + ) -> None: + if loc.numel() == 0: + return + pool = self.pool + assert self.buffer[layer_id - pool.start_layer].shape[0], ( + "cannot write indexer keys for a skip-topk layer" + ) + if pool._shard_extend_active: + index_buf_accessor.SetKAndS.execute( + pool=pool, + buf=pool._slots[layer_id % 2].tensors["index_k"], + loc=pool.translate_loc_to_scratch(loc), + index_k=index_k, + index_k_scale=index_k_scale, + ) + # Indexer callers may reuse or mutate their location tensors. Compute + # fresh indices instead of memoizing only their pointer and length. + owned_idx = torch.nonzero( + pool.placement.local_mask(loc, pool.shard_rank) + ).squeeze(1) + local_rows = pool.placement.local_index(loc[owned_idx]) + if owned_idx.numel(): + super().store_quantized( + layer_id, + local_rows, + index_k.index_select(0, owned_idx), + index_k_scale.index_select(0, owned_idx), + ) + + +class PageInterleaveDSATokenToKVPool(PageInterleaveMLATokenToKVPool, DSATokenToKVPool): + """DSA latent and indexer caches sharing one logical-page shard plan. + + The indexer keeps its native page layout: all FP8 keys followed by all + FP32 scales in each uint8 page. Prefix gathers copy complete packed pages; + chunk stores stage every token and persist only the owning rank's tokens. + + ``get_index_k_continuous``, ``get_index_k_scale_continuous``, and + ``get_index_k_scale_buffer`` accept logical page IDs during an active + sharded extend. Consumers of the raw ``get_index_k_with_scale_buffer`` + must translate their page IDs with the batch's scratch plan themselves, + just as consumers of ``get_key_buffer`` do. Outside an extend, raw getters + expose owner-local storage; ``index_key_cache.get_local_buffer`` always + exposes that storage. + + This is a pool-level implementation. DSA backend metadata, fused indexer + stores, and model-runner selection still require separate integration. + """ + + def __init__( + self, + *args, + shard_spec: PageShardSpec, + shard_group: GroupCoordinator, + **kwargs, + ): + # Initialize both DSA buffers before allocating their shared scratch; + # the ordinary sharded MLA constructor deliberately rejects use_dsa. + DSATokenToKVPool.__init__(self, *args, **kwargs) + assert self.index_buf_size == self.size, ( + "page-interleave DSA requires equal latent and indexer capacities" + ) + self._init_page_shard_state(shard_spec, shard_group) + + def _create_index_key_cache(self) -> IndexKeyCache: + return _PageInterleaveIndexKeyCache(self, self.index_buf_size) + + def _scratch_tensor_specs(self, rows: int) -> Dict[str, torch.Tensor]: + tensors = super()._scratch_tensor_specs(rows) + tensors["index_k"] = torch.zeros( + self.index_key_cache._buffer_shape(rows // self.page_size), + dtype=self.index_k_with_scale_buffer_dtype, + device=self.device, + ) + return tensors + + def _gather_pairs(self, local_layer: int): + pairs = super()._gather_pairs(local_layer) + index_k = self.index_k_with_scale_buffer[local_layer] + if index_k.shape[0]: + pairs.append((index_k, "index_k")) + return pairs + + def _gather_send_rows(self, name: str) -> torch.Tensor: + # The indexer's first dimension counts complete pages, while latent + # storage counts tokens. Both use the exact same page ownership order. + if name == "index_k": + return self._send_pages + return super()._gather_send_rows(name) + + def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: + return self.index_key_cache.get_buffer(layer_id) + + def set_kv_buffer( + self, + layer: RadixAttention, + loc_info, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + layer_id_override: Optional[int] = None, + ): + if self.dsa_kv_cache_store_fp8: + raise NotImplementedError( + "packed DSA latent writes require set_mla_kv_buffer" + ) + return super().set_kv_buffer( + layer, loc_info, cache_k, cache_v, layer_id_override + ) + + def get_mla_kv_buffer( + self, + layer: RadixAttention, + loc: torch.Tensor, + dst_dtype: Optional[torch.dtype] = None, + ): + """Read logical rows as separate NoPE and RoPE tensors. + + Packed FP8 rows contain scales and BF16 RoPE bytes, so decode them + after gathering instead of using the ordinary MLA row-copy kernel. + The DSA decoder produces BF16 by default; an explicit ``dst_dtype`` + converts those decoded values to the requested type. + """ + if not self.dsa_kv_cache_store_fp8: + return super().get_mla_kv_buffer(layer, loc, dst_dtype) + + slot = self._acquire_slot_for_read(layer.layer_id) + rows = self.translate_loc_to_scratch(loc) + kv = dequantize_k_cache_paged(self._scratch_kv(slot), rows) + if dst_dtype is not None: + kv = kv.to(dst_dtype) + return ( + kv[..., : self.kv_lora_rank].contiguous(), + kv[..., self.kv_lora_rank :].contiguous(), + ) + + def get_value_buffer(self, layer_id: int): + if self.dsa_kv_cache_store_fp8: + raise NotImplementedError( + "packed DSA latent reads require the raw key buffer and a DSA backend" + ) + return super().get_value_buffer(layer_id) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + raise NotImplementedError( + "DSA cache relocation requires a transfer between logical-page owners" + ) + + def get_cpu_copy(self, indices, mamba_indices=None): + raise NotImplementedError( + "DSA CPU offload is unsupported under logical-page KV sharding" + ) + + def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None): + raise NotImplementedError( + "DSA CPU restore is unsupported under logical-page KV sharding" + ) diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index 2ae799d664bf..9ddbb89c1b2e 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -13,8 +13,8 @@ # ============================================================================== """Tests for logical-page KV cache sharding. -Two sections. The first (CPU only, what the CPU CI job runs) pins the pure -arithmetic that rotated owner-classed allocation hangs on: +The CPU tests pin the pure arithmetic that rotated owner-classed +allocation hangs on: 1. The placement bijection ``loc = Q*(N*ps) + r*ps + o`` — owner / local-row round-trip, disjoint equal partition across ranks. @@ -34,15 +34,17 @@ class draws (owners exactly cyclic along a chain), least-full root 6. The value-derived P/D ownership send filter ``filter_kv_indices_for_shard_rank``. -The second section (``TestPageInterleaveGatherMultiGpu``, at the bottom) drives -real pools over a real 2-rank process group. It is the only check that the plan -the CPU stub validates actually addresses the bytes NCCL delivers, so it is -skipped rather than dropped when fewer than 2 CUDA devices are visible — which -is every run of the CPU suite this file is registered to. +The GPU tests exercise real pools and NCCL collectives: MLA and MHA gathers +on 2 ranks (``TestPageInterleaveGatherMultiGpu``), and byte-exact DSA indexer +keys, scales, and latent KV on 2 and 4 ranks (``TestPageInterleaveIndexer``), +including dequantized MLA reads from packed FP8 scratch. +GPU tests skip when too few CUDA devices are visible, so the same file runs +in both CPU and CUDA CI suites. """ import os import shutil +import tempfile import unittest import unittest.mock from array import array @@ -93,11 +95,12 @@ class draws (owners exactly cyclic along a chain), least-full root from sglang.srt.runtime_context import get_parallel, publish from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ceil_div -from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci from sglang.test.mem_cache_utils import finish_req from sglang.test.test_utils import CustomTestCase, publish_build_topology register_cpu_ci(est_time=30, suite="base-a-test-cpu") +register_cuda_ci(est_time=120, stage="base-c", runner_config="4-gpu-b200") N = 4 # shard size PS = 16 # physical page size @@ -1739,5 +1742,528 @@ def test_mha_shard_over_attention_cp(self): mp.spawn(_run_mha, args=(_GATHER_WORLD, 29812), nprocs=_GATHER_WORLD, join=True) +# ============================================================================= +# Indexer: packed FP8 keys and FP32 scales over real NCCL (2 and 4 GPUs). +# ============================================================================= + +_INDEXER_PAGE_SIZE = 64 +_INDEXER_SIZE = _INDEXER_PAGE_SIZE * 32 +_INDEXER_HEAD_DIM = 128 +_INDEXER_KV_LORA_RANK = 512 +_INDEXER_QK_ROPE = 64 +_INDEXER_START_LAYER = 3 +_INDEXER_SKIP_TOPK = [False, True, False, False, False] + + +def _assert_bytes(got, expected, label): + got = got.contiguous().view(torch.uint8) + expected = expected.contiguous().view(torch.uint8) + assert got.shape == expected.shape, label + assert torch.equal(got, expected), f"{label}: cache bytes differ" + + +def _index_values(locs, layer_id, generation): + dims = torch.arange(_INDEXER_HEAD_DIM, device=locs.device) + keys = ( + (locs[:, None] * 7 + dims * 3 + layer_id * 11 + generation * 13) % 31 - 15 + ).to(torch.float8_e4m3fn) + scales = ( + (locs % 29 + layer_id * 31 + generation * 17 + 1).float() / 128 + ).unsqueeze(1) + return keys, scales + + +def _latent_values(locs, layer_id, generation): + dims = torch.arange(_INDEXER_KV_LORA_RANK + _INDEXER_QK_ROPE, device=locs.device) + values = ((locs[:, None] * 3 + dims + layer_id * 5 + generation * 7) % 113 - 56).to( + torch.bfloat16 + ) / 16 + # Vary the quantization scale across tokens and 128-value NoPE groups. + amplitude = torch.exp2( + ((locs[:, None] + dims // 128 + layer_id + generation) % 5).float() - 2 + ).to(torch.bfloat16) + return (values * amplitude).unsqueeze(1) + + +def _locs(pages, seq_len, device): + return ( + torch.tensor(pages, dtype=torch.int64, device=device)[:, None] + * _INDEXER_PAGE_SIZE + + torch.arange(_INDEXER_PAGE_SIZE, device=device) + ).flatten()[:seq_len] + + +def _unpack_index(buf, pages, seq_len): + """Independent unpacking of [page FP8 keys | page FP32 scales].""" + selected = buf[pages.long()] + split = _INDEXER_PAGE_SIZE * _INDEXER_HEAD_DIM + keys = selected[:, :split].reshape(-1, _INDEXER_HEAD_DIM)[:seq_len] + scales = selected[:, split:].reshape(-1, 4)[:seq_len] + return keys, scales + + +def _unpack_latent(buf, locs, packed): + """Decode reference rows independently of the production DSA kernels.""" + selected = buf[locs] + if not packed: + return ( + selected[..., :_INDEXER_KV_LORA_RANK], + selected[..., _INDEXER_KV_LORA_RANK:], + ) + raw = selected.contiguous().view(torch.uint8) + scale_end = _INDEXER_KV_LORA_RANK + _INDEXER_KV_LORA_RANK // 128 * 4 + values = raw[..., :_INDEXER_KV_LORA_RANK].view(torch.float8_e4m3fn).float() + scales = ( + raw[..., _INDEXER_KV_LORA_RANK:scale_end] + .contiguous() + .view(torch.float32) + .repeat_interleave(128, dim=-1) + ) + return ( + (values * scales).to(torch.bfloat16), + raw[..., scale_end:].contiguous().view(torch.bfloat16), + ) + + +class _IndexerCacheCheck: + def __init__(self, group, packed): + from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool + from sglang.srt.mem_cache.page_interleave import PageShardSpec + from sglang.srt.mem_cache.page_interleave_pool import ( + PageInterleaveDSATokenToKVPool, + ) + + self.world = group.world_size + self.rank = group.rank_in_group + self.device = group.device + self.packed = packed + kwargs = dict( + page_size=_INDEXER_PAGE_SIZE, + kv_lora_rank=_INDEXER_KV_LORA_RANK, + dtype=torch.float8_e4m3fn if packed else torch.bfloat16, + qk_rope_head_dim=_INDEXER_QK_ROPE, + layer_num=len(_INDEXER_SKIP_TOPK), + device=str(self.device), + index_head_dim=_INDEXER_HEAD_DIM, + enable_memory_saver=False, + kv_cache_dim=656 if packed else _INDEXER_KV_LORA_RANK + _INDEXER_QK_ROPE, + start_layer=_INDEXER_START_LAYER, + end_layer=_INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) - 1, + skip_topk_layers=_INDEXER_SKIP_TOPK, + ) + self.pool = PageInterleaveDSATokenToKVPool( + _INDEXER_SIZE, + **kwargs, + shard_spec=PageShardSpec( + shard_rank=self.rank, + shard_size=self.world, + page_size=_INDEXER_PAGE_SIZE, + max_prefix_tokens=8 * self.world * _INDEXER_PAGE_SIZE, + chunk_tokens=8 * _INDEXER_PAGE_SIZE, + ), + shard_group=group, + ) + # Include each rank's reserved/padded physical page in the reference. + self.reference = DSATokenToKVPool( + (_INDEXER_SIZE + _INDEXER_PAGE_SIZE) * self.world, **kwargs + ) + + def store(self, locs, layer_id, generation): + values = _latent_values(locs, layer_id, generation) + for pool in (self.pool, self.reference): + pool.set_mla_kv_buffer( + SimpleNamespace(layer_id=layer_id), + locs, + values[..., :_INDEXER_KV_LORA_RANK].contiguous(), + values[..., _INDEXER_KV_LORA_RANK:].contiguous(), + ) + if _INDEXER_SKIP_TOPK[layer_id - _INDEXER_START_LAYER]: + return + keys, scales = _index_values(locs, layer_id, generation) + write_scales = scales[:, 0] if generation % 2 else scales + for pool in (self.pool, self.reference): + pool.set_index_k_scale_buffer(layer_id, locs, keys, write_scales) + + # Check the reference's packing independently, including both scales + # shapes accepted by SetKAndS and every valid token in partial pages. + buf = self.reference.index_key_cache.get_local_buffer(layer_id) + pages, offsets = locs // _INDEXER_PAGE_SIZE, locs % _INDEXER_PAGE_SIZE + unpacked_keys = buf[:, : _INDEXER_PAGE_SIZE * _INDEXER_HEAD_DIM].reshape( + -1, _INDEXER_PAGE_SIZE, _INDEXER_HEAD_DIM + ) + unpacked_scales = buf[:, _INDEXER_PAGE_SIZE * _INDEXER_HEAD_DIM :].reshape( + -1, _INDEXER_PAGE_SIZE, 4 + ) + _assert_bytes(unpacked_keys[pages, offsets], keys, "reference FP8 keys") + _assert_bytes(unpacked_scales[pages, offsets], scales, "reference FP32 scales") + + def assert_owned(self, layer_id): + local_layer = layer_id - _INDEXER_START_LAYER + owned = self.pool.index_key_cache.get_local_buffer(layer_id) + assert ( + owned.data_ptr() + == self.pool.index_k_with_scale_buffer[local_layer].data_ptr() + ) + if _INDEXER_SKIP_TOPK[local_layer]: + assert owned.shape[0] == 0 + return + logical_pages = ( + torch.arange(owned.shape[0], device=self.device) * self.world + self.rank + ) + expected = self.reference.index_k_with_scale_buffer[local_layer][logical_pages] + # Comparing the whole allocation catches writes to foreign pages and + # overwrites of untouched bytes in a partially filled page. + _assert_bytes(owned, expected, f"owned index layer {layer_id}") + + def seed(self, pages, generation): + self.pool.end_shard_extend() + locs = _locs(pages, len(pages) * _INDEXER_PAGE_SIZE, self.device) + for layer_id in range( + _INDEXER_START_LAYER, _INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) + ): + self.store(locs, layer_id, generation) + self.assert_owned(layer_id) + + def assert_latent_reads(self, layer_id, rows): + # Cover every gathered prefix/chunk row in reverse order, repeated + # locations, partial pages, and a noncontiguous logical-location tensor. + all_locs = torch.cat(rows) + selected = torch.cat((all_locs.flip(0), all_locs[1:2].expand(2))) + storage = torch.empty( + selected.numel() * 2, dtype=selected.dtype, device=self.device + ) + read_locs = storage[::2] + read_locs.copy_(selected) + reference = self.reference.kv_buffer[layer_id - _INDEXER_START_LAYER] + layer = SimpleNamespace(layer_id=layer_id) + for locs in (read_locs, read_locs[:0]): + expected = _unpack_latent(reference, locs, self.packed) + for dtype in (None, torch.bfloat16, torch.float16, torch.float32): + actual = self.pool.get_mla_kv_buffer(layer, locs, dst_dtype=dtype) + expected_dtype = dtype or torch.bfloat16 + for got, value in zip(actual, expected): + assert got.shape == value.shape + assert got.dtype == expected_dtype + assert got.is_contiguous() + _assert_bytes(got, value.to(expected_dtype), "decoded latent KV") + + def batch(self, requests, generation): + """requests contains (logical pages, prefix token count, sequence length).""" + seq_lens = [seq_len for _, _, seq_len in requests] + prefix_lens = [prefix_len for _, prefix_len, _ in requests] + rows = [_locs(pages, seq_len, self.device) for pages, _, seq_len in requests] + # Nonconsecutive request IDs ensure the plan indexes req_pool_indices. + req_ids = torch.arange(len(rows), device=self.device) * 2 + 1 + req_to_token = torch.zeros( + (2 * len(rows) + 1, max(seq_lens)), dtype=torch.int64, device=self.device + ) + page_table = torch.zeros( + (len(rows), max(len(pages) for pages, _, _ in requests)), + dtype=torch.int32, + device=self.device, + ) + for i, ((pages, _, seq_len), row) in enumerate(zip(requests, rows)): + req_to_token[req_ids[i], :seq_len] = row + page_table[i, : len(pages)] = torch.tensor(pages, device=self.device) + chunk_locs = torch.cat([row[p:] for row, p in zip(rows, prefix_lens)]) + self.pool.begin_shard_extend(req_to_token, req_ids, prefix_lens, seq_lens) + + for layer_id in range( + _INDEXER_START_LAYER, _INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) + ): + if chunk_locs.numel(): + self.store(chunk_locs, layer_id, generation) + # Read before the raw-key accessor so MLA reads must themselves + # wait for the prefix allgather and include the locally staged chunk. + self.assert_latent_reads(layer_id, rows) + if _INDEXER_SKIP_TOPK[layer_id - _INDEXER_START_LAYER]: + assert self.pool.get_index_k_with_scale_buffer(layer_id).shape[0] == 0 + # Every layer advances prefetch, including layers reusing top-k. + latent = self.pool.get_key_buffer(layer_id) + for row in rows: + scratch_rows = self.pool.translate_loc_to_scratch(row) + expected = self.reference.kv_buffer[layer_id - _INDEXER_START_LAYER][ + row + ] + _assert_bytes(latent[scratch_rows], expected, "latent KV") + + self.assert_owned(layer_id) + if _INDEXER_SKIP_TOPK[layer_id - _INDEXER_START_LAYER]: + with unittest.TestCase().assertRaisesRegex(AssertionError, "skip-topk"): + self.pool.get_index_k_continuous(layer_id, 1, page_table[0]) + continue + + # Repeated public reads must preserve the resident layer while + # the next layer's NCCL allgather is in flight on the other slot. + for _ in range(2): + for i, (_, _, seq_len) in enumerate(requests): + pages = page_table[i].contiguous() + reference_buf = self.reference.get_index_k_with_scale_buffer( + layer_id + ) + expected_k, expected_s = _unpack_index( + reference_buf, pages, seq_len + ) + _assert_bytes( + self.pool.get_index_k_continuous(layer_id, seq_len, pages), + expected_k, + "continuous FP8 keys", + ) + _assert_bytes( + self.pool.get_index_k_scale_continuous( + layer_id, seq_len, pages + ), + expected_s, + "continuous FP32 scales", + ) + scratch_pages = ( + self.pool.translate_loc_to_scratch( + pages.long() * _INDEXER_PAGE_SIZE + ) + // _INDEXER_PAGE_SIZE + ) + raw = self.pool.get_index_k_with_scale_buffer(layer_id) + raw_k, raw_s = _unpack_index(raw, scratch_pages, seq_len) + _assert_bytes(raw_k, expected_k, "raw scratch FP8 keys") + _assert_bytes(raw_s, expected_s, "raw scratch FP32 scales") + + lengths = torch.tensor(seq_lens, dtype=torch.int32, device=self.device) + got = self.pool.get_index_k_scale_buffer( + layer_id, lengths, page_table, sum(seq_lens), max(seq_lens) + ) + expected = self.reference.get_index_k_scale_buffer( + layer_id, lengths, page_table, sum(seq_lens), max(seq_lens) + ) + for actual, reference in zip(got, expected): + _assert_bytes(actual, reference, "batched index gather") + + self.pool.end_shard_extend() + for layer_id in range( + _INDEXER_START_LAYER, _INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) + ): + assert self.pool.get_index_k_with_scale_buffer(layer_id).data_ptr() == ( + self.pool.index_key_cache.get_local_buffer(layer_id).data_ptr() + ) + torch.cuda.synchronize() + torch.distributed.barrier() + + def mutated_index_locations(self): + # Indexer callers may reuse the same storage with new logical locations; + # a write plan cached only by pointer/length must not survive that change. + mutable_loc = torch.empty(1, dtype=torch.int64, device=self.device) + for active in (False, True): + pages = torch.tensor( + [28 * self.world + 1, 29 * self.world], device=self.device + ) + if active: + self.pool.begin_shard_extend( + pages[:, None] * _INDEXER_PAGE_SIZE, + torch.arange(2, device=self.device), + [0, 0], + [1, 1], + ) + else: + with unittest.TestCase().assertRaisesRegex( + AssertionError, "begin_shard_extend" + ): + self.pool.get_index_k_continuous(_INDEXER_START_LAYER, 1, pages[:1]) + + for layer_id in range( + _INDEXER_START_LAYER, _INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) + ): + # An empty store must also be harmless on a skip-topk layer. + empty_loc = mutable_loc[:0] + empty_k, empty_s = _index_values(empty_loc, layer_id, 1) + self.pool.set_index_k_scale_buffer( + layer_id, empty_loc, empty_k, empty_s + ) + if _INDEXER_SKIP_TOPK[layer_id - _INDEXER_START_LAYER]: + continue + + for i in range(2): + mutable_loc.copy_(pages[i : i + 1] * _INDEXER_PAGE_SIZE) + # Both an in-place mutation and another tensor object + # aliasing the old allocation retain the same data pointer. + loc = mutable_loc if i == 0 else mutable_loc.view_as(mutable_loc) + keys, scales = _index_values(loc, layer_id, 7 + i + int(active)) + for pool in (self.pool, self.reference): + pool.set_index_k_scale_buffer(layer_id, loc, keys, scales) + self.assert_owned(layer_id) + + if active: + for i in range(2): + expected_k, expected_s = _index_values( + pages[i : i + 1] * _INDEXER_PAGE_SIZE, layer_id, 8 + i + ) + _assert_bytes( + self.pool.get_index_k_continuous( + layer_id, 1, pages[i : i + 1] + ), + expected_k, + "mutated location FP8 keys", + ) + _assert_bytes( + self.pool.get_index_k_scale_continuous( + layer_id, 1, pages[i : i + 1] + ), + expected_s, + "mutated location FP32 scales", + ) + self.pool.end_shard_extend() + torch.cuda.synchronize() + torch.distributed.barrier() + + def async_prefix(self, pages): + """Queue all layers' reads before synchronizing to exercise slot reuse.""" + locs = _locs(pages, len(pages) * _INDEXER_PAGE_SIZE, self.device) + page_ids = torch.tensor(pages, dtype=torch.int32, device=self.device) + self.pool.begin_shard_extend( + locs[None, :], + torch.tensor([0], device=self.device), + [locs.numel()], + [locs.numel()], + ) + scratch_pages = ( + self.pool.translate_loc_to_scratch(page_ids.long() * _INDEXER_PAGE_SIZE) + // _INDEXER_PAGE_SIZE + ) + snapshots = [] + for layer_id in range( + _INDEXER_START_LAYER, _INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) + ): + decoded = self.pool.get_mla_kv_buffer( + SimpleNamespace(layer_id=layer_id), locs + ) + expected_latent = _unpack_latent( + self.reference.kv_buffer[layer_id - _INDEXER_START_LAYER], + locs, + self.packed, + ) + snapshots.extend(zip(decoded, expected_latent)) + raw = self.pool.get_index_k_with_scale_buffer(layer_id) + if _INDEXER_SKIP_TOPK[layer_id - _INDEXER_START_LAYER]: + assert raw.shape[0] == 0 + continue + reference = self.reference.get_index_k_with_scale_buffer(layer_id) + snapshots.append((raw[scratch_pages].clone(), reference[page_ids])) + expected_k, expected_s = _unpack_index(reference, page_ids, locs.numel()) + snapshots.append( + ( + self.pool.get_index_k_continuous(layer_id, locs.numel(), page_ids), + expected_k, + ) + ) + snapshots.append( + ( + self.pool.get_index_k_scale_continuous( + layer_id, locs.numel(), page_ids + ), + expected_s, + ) + ) + torch.cuda.synchronize() + for actual, expected in snapshots: + _assert_bytes(actual, expected, "asynchronous layer pipeline") + self.pool.end_shard_extend() + torch.distributed.barrier() + + +def _run_indexer(rank, world, init_method): + os.environ.update( + RANK=str(rank), + WORLD_SIZE=str(world), + ) + os.environ.setdefault("no_proxy", "127.0.0.1,localhost") + os.environ["SGLANG_DEBUG_MEMORY_POOL"] = "1" + torch.cuda.set_device(rank) + from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, + ) + from sglang.srt.runtime_context import get_parallel + + init_distributed_environment( + world_size=world, + rank=rank, + local_rank=rank, + distributed_init_method=init_method, + backend="nccl", + ) + publish_build_topology(tp_size=world, attn_cp_size=world, world_rank=rank) + initialize_model_parallel() + with get_parallel().override(attn_cp_size=world): + group = get_parallel().attn_cp_group + for packed in (False, True): + check = _IndexerCacheCheck(group, packed) + # Fragmented physical pages with a rotated cyclic ownership run. + prefix = [ + physical * world + (i + 1) % world + for i, physical in enumerate([5, 2, 9, 4, 7][: world + 1]) + ] + check.seed(prefix, generation=1) + chunk = [12 * world + (world + 2) % world, 11 * world + (world + 3) % world] + check.batch( + [ + ( + prefix + chunk, + len(prefix) * _INDEXER_PAGE_SIZE, + (len(prefix) + 1) * _INDEXER_PAGE_SIZE + 13, + ), + ( + [prefix[0], 16 * world + 2 % world], + _INDEXER_PAGE_SIZE, + _INDEXER_PAGE_SIZE + 7, + ), + ], + generation=2, + ) + # Prefix-only batch: exactly one rank owns a page, all peers send + # padding. Also checks double-buffer epoch invalidation. + check.batch( + [([prefix[0]], _INDEXER_PAGE_SIZE, _INDEXER_PAGE_SIZE)], generation=3 + ) + fresh = [20 * world + world - 1, 21 * world] + check.batch([(fresh, 0, _INDEXER_PAGE_SIZE + 9)], generation=4) + # Same logical prefix and layer IDs, new bytes after page reuse. + check.seed([prefix[0]], generation=5) + check.batch( + [ + ( + [prefix[0], 23 * world + 2 % world], + _INDEXER_PAGE_SIZE, + _INDEXER_PAGE_SIZE + 19, + ) + ], + generation=6, + ) + check.async_prefix(prefix) + check.mutated_index_locations() + print(f"rank {rank}: {'packed FP8' if packed else 'BF16'} DSA cache OK") + del check + destroy_model_parallel() + destroy_distributed_environment() + + +class TestPageInterleaveIndexer(CustomTestCase): + def _test_world(self, world): + if torch.cuda.device_count() < world: + self.skipTest(f"indexer KV sharding test requires {world} GPUs") + with tempfile.TemporaryDirectory(prefix="sglang-indexer-shard-") as tmp: + mp.spawn( + _run_indexer, + args=(world, f"file://{tmp}/rendezvous"), + nprocs=world, + join=True, + ) + + def test_two_gpu(self): + self._test_world(2) + + def test_four_gpu(self): + self._test_world(4) + + if __name__ == "__main__": unittest.main() From 8ed0347e1704897a36a1c519ace9dee7d7bf849b Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Mon, 28 Sep 2026 00:58:45 -0700 Subject: [PATCH 06/11] docs: explain DSA shard buffer layouts (cherry picked from commit f176b897a4fe170b0437054e5342e68f28144a8f) --- .../sglang/srt/mem_cache/page_interleave_pool.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index b0a859536606..ffaae4a44c45 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -831,9 +831,19 @@ def store_quantized( class PageInterleaveDSATokenToKVPool(PageInterleaveMLATokenToKVPool, DSATokenToKVPool): """DSA latent and indexer caches sharing one logical-page shard plan. - The indexer keeps its native page layout: all FP8 keys followed by all - FP32 scales in each uint8 page. Prefix gathers copy complete packed pages; - chunk stores stage every token and persist only the owning rank's tokens. + The buffers use different units along their first dimension: + - Latent KV: [num_tokens, 1, kv_cache_dim], one token per row. + - Indexer: [num_pages, page_size * (128 + 4)] in uint8, one page per row. + Each page packs all 128-element FP8 keys first, followed by one FP32 + scale per token. With page_size=64, one indexer row contains 8448 bytes. + + For example, local page 2 corresponds to latent rows 128..191 but indexer + row 2 when page_size=64. The ``_gather_send_rows`` hook therefore selects + ``_send_rows`` for latent KV and ``_send_pages`` for the indexer, letting + the shared gather compute rank offsets and block sizes in the correct + units while retaining the same page ownership order and scratch plan. + Prefix gathers copy complete packed indexer pages; chunk stores stage + every token and persist only the owning rank's tokens. ``get_index_k_continuous``, ``get_index_k_scale_continuous``, and ``get_index_k_scale_buffer`` accept logical page IDs during an active From d4c90961fb6ca73fca4dfca9665e8f755cbdc90d Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 6 Oct 2026 01:17:02 -0700 Subject: [PATCH 07/11] fix: read sharded DSA indexer keys through the batched getter Upstream #41469 removed get_index_k_continuous and get_index_k_scale_continuous (and IndexKeyCache.get_k_continuous / get_k_scale_continuous) as orphaned. The sharded index cache still overrode the IndexKeyCache pair, and the indexer test called the pool pair, so CI on the main-merged branch failed with AttributeError. Drop the two overrides and read through get_index_k_scale_buffer, which still translates logical page IDs into the scratch plan and keeps the skip-topk and begin_shard_extend assertions. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit e33582fe5b790d1905ed1b227b6017e745b4df87) --- .../srt/mem_cache/page_interleave_pool.py | 15 +---- .../mem_cache/test_page_interleave_shard.py | 59 ++++++------------- 2 files changed, 18 insertions(+), 56 deletions(-) diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index ffaae4a44c45..c18165cf964d 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -764,18 +764,6 @@ def _scratch_pages(self, layer_id: int, page_indices: torch.Tensor) -> torch.Ten ) return self.pool._page_pos[page_indices.long()] - def get_k_continuous(self, layer_id: int, seq_len: int, page_indices: torch.Tensor): - return super().get_k_continuous( - layer_id, seq_len, self._scratch_pages(layer_id, page_indices) - ) - - def get_k_scale_continuous( - self, layer_id: int, seq_len: int, page_indices: torch.Tensor - ): - return super().get_k_scale_continuous( - layer_id, seq_len, self._scratch_pages(layer_id, page_indices) - ) - def get_k_and_scale( self, layer_id: int, @@ -845,8 +833,7 @@ class PageInterleaveDSATokenToKVPool(PageInterleaveMLATokenToKVPool, DSATokenToK Prefix gathers copy complete packed indexer pages; chunk stores stage every token and persist only the owning rank's tokens. - ``get_index_k_continuous``, ``get_index_k_scale_continuous``, and - ``get_index_k_scale_buffer`` accept logical page IDs during an active + ``get_index_k_scale_buffer`` accepts logical page IDs during an active sharded extend. Consumers of the raw ``get_index_k_with_scale_buffer`` must translate their page IDs with the batch's scratch plan themselves, just as consumers of ``get_key_buffer`` do. Outside an extend, raw getters diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index 9ddbb89c1b2e..bf349b0249f8 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -1947,6 +1947,13 @@ def assert_latent_reads(self, layer_id, rows): assert got.is_contiguous() _assert_bytes(got, value.to(expected_dtype), "decoded latent KV") + def read_index(self, layer_id, pages, seq_len): + """Gather one request's keys and scales by logical page IDs.""" + lengths = torch.tensor([seq_len], dtype=torch.int32, device=self.device) + return self.pool.get_index_k_scale_buffer( + layer_id, lengths, pages[None], seq_len, seq_len + ) + def batch(self, requests, generation): """requests contains (logical pages, prefix token count, sequence length).""" seq_lens = [seq_len for _, _, seq_len in requests] @@ -1990,7 +1997,7 @@ def batch(self, requests, generation): self.assert_owned(layer_id) if _INDEXER_SKIP_TOPK[layer_id - _INDEXER_START_LAYER]: with unittest.TestCase().assertRaisesRegex(AssertionError, "skip-topk"): - self.pool.get_index_k_continuous(layer_id, 1, page_table[0]) + self.read_index(layer_id, page_table[0], 1) continue # Repeated public reads must preserve the resident layer while @@ -2004,18 +2011,9 @@ def batch(self, requests, generation): expected_k, expected_s = _unpack_index( reference_buf, pages, seq_len ) - _assert_bytes( - self.pool.get_index_k_continuous(layer_id, seq_len, pages), - expected_k, - "continuous FP8 keys", - ) - _assert_bytes( - self.pool.get_index_k_scale_continuous( - layer_id, seq_len, pages - ), - expected_s, - "continuous FP32 scales", - ) + got_k, got_s = self.read_index(layer_id, pages, seq_len) + _assert_bytes(got_k, expected_k, "per-request FP8 keys") + _assert_bytes(got_s, expected_s, "per-request FP32 scales") scratch_pages = ( self.pool.translate_loc_to_scratch( pages.long() * _INDEXER_PAGE_SIZE @@ -2066,7 +2064,7 @@ def mutated_index_locations(self): with unittest.TestCase().assertRaisesRegex( AssertionError, "begin_shard_extend" ): - self.pool.get_index_k_continuous(_INDEXER_START_LAYER, 1, pages[:1]) + self.read_index(_INDEXER_START_LAYER, pages[:1], 1) for layer_id in range( _INDEXER_START_LAYER, _INDEXER_START_LAYER + len(_INDEXER_SKIP_TOPK) @@ -2095,20 +2093,9 @@ def mutated_index_locations(self): expected_k, expected_s = _index_values( pages[i : i + 1] * _INDEXER_PAGE_SIZE, layer_id, 8 + i ) - _assert_bytes( - self.pool.get_index_k_continuous( - layer_id, 1, pages[i : i + 1] - ), - expected_k, - "mutated location FP8 keys", - ) - _assert_bytes( - self.pool.get_index_k_scale_continuous( - layer_id, 1, pages[i : i + 1] - ), - expected_s, - "mutated location FP32 scales", - ) + got_k, got_s = self.read_index(layer_id, pages[i : i + 1], 1) + _assert_bytes(got_k, expected_k, "mutated location FP8 keys") + _assert_bytes(got_s, expected_s, "mutated location FP32 scales") self.pool.end_shard_extend() torch.cuda.synchronize() torch.distributed.barrier() @@ -2147,20 +2134,8 @@ def async_prefix(self, pages): reference = self.reference.get_index_k_with_scale_buffer(layer_id) snapshots.append((raw[scratch_pages].clone(), reference[page_ids])) expected_k, expected_s = _unpack_index(reference, page_ids, locs.numel()) - snapshots.append( - ( - self.pool.get_index_k_continuous(layer_id, locs.numel(), page_ids), - expected_k, - ) - ) - snapshots.append( - ( - self.pool.get_index_k_scale_continuous( - layer_id, locs.numel(), page_ids - ), - expected_s, - ) - ) + got_k, got_s = self.read_index(layer_id, page_ids, locs.numel()) + snapshots.extend(((got_k, expected_k), (got_s, expected_s))) torch.cuda.synchronize() for actual, expected in snapshots: _assert_bytes(actual, expected, "asynchronous layer pipeline") From 8cfbeac22482f6d95099cd8ac266ce82a44c26e8 Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Tue, 6 Oct 2026 03:55:25 -0700 Subject: [PATCH 08/11] fix: reject the k-pool indexer layout in the sharded DSA pool Since #42178, slots_per_page = page_size // index_kpool, but the sharded pool translates indexer locations in token units. With index_kpool > 1 that would address the wrong rows, so fail at construction instead. Co-Authored-By: Claude Opus 5.5 (cherry picked from commit 5a82ac9f311a2d558ecf2ccf57c0f2d2eda1fc81) --- .../srt/mem_cache/page_interleave_pool.py | 5 ++++ .../mem_cache/test_page_interleave_shard.py | 29 ++++++++++++++++++- 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index c18165cf964d..2a932b34a8f1 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -857,6 +857,11 @@ def __init__( assert self.index_buf_size == self.size, ( "page-interleave DSA requires equal latent and indexer capacities" ) + # Indexer locations are translated in token units, which equal k-pool + # slots only when every page keeps all of its tokens. + assert self.index_kpool == 1, ( + f"page-interleave DSA does not support index_kpool={self.index_kpool}" + ) self._init_page_shard_state(shard_spec, shard_group) def _create_index_key_cache(self) -> IndexKeyCache: diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index bf349b0249f8..c3007f11aab6 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -1833,11 +1833,12 @@ def __init__(self, group, packed): PageInterleaveDSATokenToKVPool, ) + self.group = group self.world = group.world_size self.rank = group.rank_in_group self.device = group.device self.packed = packed - kwargs = dict( + self.kwargs = kwargs = dict( page_size=_INDEXER_PAGE_SIZE, kv_lora_rank=_INDEXER_KV_LORA_RANK, dtype=torch.float8_e4m3fn if packed else torch.bfloat16, @@ -1868,6 +1869,31 @@ def __init__(self, group, packed): (_INDEXER_SIZE + _INDEXER_PAGE_SIZE) * self.world, **kwargs ) + def rejects_kpool(self): + from sglang.srt.mem_cache.page_interleave import PageShardSpec + from sglang.srt.mem_cache.page_interleave_pool import ( + PageInterleaveDSATokenToKVPool, + ) + + # A 128-token page with index_kpool=2 holds 64 indexer slots, so the + # token-unit location translation would address the wrong rows. + page_size = 2 * _INDEXER_PAGE_SIZE + with unittest.TestCase().assertRaisesRegex( + AssertionError, "does not support index_kpool=2" + ): + PageInterleaveDSATokenToKVPool( + _INDEXER_SIZE, + **{**self.kwargs, "page_size": page_size, "index_kpool": 2}, + shard_spec=PageShardSpec( + shard_rank=self.rank, + shard_size=self.world, + page_size=page_size, + max_prefix_tokens=8 * self.world * page_size, + chunk_tokens=8 * page_size, + ), + shard_group=self.group, + ) + def store(self, locs, layer_id, generation): values = _latent_values(locs, layer_id, generation) for pool in (self.pool, self.reference): @@ -2172,6 +2198,7 @@ def _run_indexer(rank, world, init_method): group = get_parallel().attn_cp_group for packed in (False, True): check = _IndexerCacheCheck(group, packed) + check.rejects_kpool() # Fragmented physical pages with a rotated cyclic ownership run. prefix = [ physical * world + (i + 1) % world From adf94562b7c118d1e2e7272f03f644d1e5404c7a Mon Sep 17 00:00:00 2001 From: shunkangz <182541032+Shunkangz@users.noreply.github.com> Date: Wed, 23 Sep 2026 07:41:23 -0700 Subject: [PATCH 09/11] Add MTP KV pool sharding primitives and correctness tests (cherry picked from commit 01fba662537d5b26c99d706ff63c20a085ecc883) --- .../srt/mem_cache/kv_cache_configurator.py | 122 ++++++ .../sglang/srt/mem_cache/page_interleave.py | 76 +++- .../srt/mem_cache/page_interleave_pool.py | 1 + test/manual/test_page_interleave_mtp.py | 389 ++++++++++++++++++ .../mem_cache/test_mtp_kv_shard_config.py | 365 ++++++++++++++++ .../mem_cache/test_page_interleave_shard.py | 1 + 6 files changed, 950 insertions(+), 4 deletions(-) create mode 100644 test/manual/test_page_interleave_mtp.py create mode 100644 test/registered/unit/mem_cache/test_mtp_kv_shard_config.py diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 6e938f3c01ef..b4fadd759e82 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1259,6 +1259,12 @@ def _build_token_to_kv_pool( is_dsv4_model: bool, req_to_token_pool: ReqToTokenPool, ) -> KVCache: + mtp_pool = self._build_mtp_kv_shard_pool( + sizes=sizes, is_dsa_model=is_dsa_model, is_dsv4_model=is_dsv4_model + ) + if mtp_pool is not None: + return mtp_pool + # Page-granularity envelope layout for the MHA-shaped (full / SWA) pools, # selected by swapping in the PageMajorMHATokenToKVPool subclass. The # default keeps upstream's per-layer layout. The Mamba state pool is routed @@ -1375,6 +1381,122 @@ def _build_token_to_kv_pool( ) return token_to_kv_pool + def _build_mtp_kv_shard_pool( + self, *, sizes: _PoolSizes, is_dsa_model: bool, is_dsv4_model: bool + ) -> Optional[KVCache]: + """Build a draft pool in its target's logical space, with its own KV. + + This is a pool-level opt-in through a shared sharded allocator. Target + pool creation and speculative serving remain separate integration work. + """ + if not self.is_draft_worker: + return None + + from sglang.srt.mem_cache.page_interleave import get_shared_kv_shard_pool + + target_pool = get_shared_kv_shard_pool(self) + if target_pool is None: + return None + + from sglang.srt.mem_cache.page_interleave_pool import ( + PageInterleaveMHATokenToKVPool, + PageInterleaveMLATokenToKVPool, + ) + + if ( + self.spec_algorithm != SpeculativeAlgorithm.EAGLE + or (self.model_config.num_nextn_predict_layers or 0) <= 0 + or get_spec().enable_multi_layer_eagle + ): + raise ValueError( + "Shared KV sharding supports only an EAGLE MTP draft runner." + ) + if get_parallel().attn_dcp_size > 1: + raise ValueError("MTP KV sharding is incompatible with DCP.") + if ( + is_dsa_model + or is_dsv4_model + or self.is_hybrid_swa + or self.mambaish_config is not None + or self.sliding_window_size is not None + or is_minimax_sparse(self.model_config.hf_config) + or current_platform.is_out_of_tree() + or _is_npu + ): + raise ValueError("MTP KV sharding supports only dense MLA and MHA pools.") + if ( + self.post_capture_kv_active + or get_memory().enable_page_major_kv_layout + or self.kv_cache_dtype_str == "mxfp8" + or is_float4_e2m1fn_x2(self.kv_cache_dtype) + ): + raise ValueError("MTP KV sharding requires the plain per-layer KV layout.") + if ( + self.page_size != target_pool.page_size + or self.pool_page_size != target_pool.page_size + ): + raise ValueError( + "MTP and target KV pools must use the same physical page size." + ) + if sizes.max_total_num_tokens != target_pool.size: + raise ValueError( + "MTP and target sharded KV pools must have the same per-rank capacity." + ) + target_is_mla = isinstance(target_pool, PageInterleaveMLATokenToKVPool) + if self.use_mla_backend != target_is_mla or not isinstance( + target_pool, + (PageInterleaveMLATokenToKVPool, PageInterleaveMHATokenToKVPool), + ): + raise ValueError( + "MTP and target KV pools must use the same attention geometry." + ) + fields = ( + ("kv_lora_rank", "qk_rope_head_dim") + if self.use_mla_backend + else ("head_dim", "v_head_dim") + ) + if ( + self.kv_cache_dtype != target_pool.dtype + or any( + getattr(self.model_config, name) != getattr(target_pool, name) + for name in fields + ) + or ( + not self.use_mla_backend + and self.model_config.get_num_kv_heads(get_parallel().attn_tp_size) + != target_pool.head_num + ) + ): + raise ValueError( + "MTP and target KV pools must use the same KV geometry and dtype." + ) + + kwargs = dict( + size=target_pool.size, + page_size=target_pool.page_size, + dtype=self.kv_cache_dtype, + layer_num=self.layer_info.num_effective_layers, + device=self.device, + enable_memory_saver=get_exec().features.enable_memory_saver, + start_layer=self.layer_info.start_layer, + end_layer=self.layer_info.end_layer, + shard_spec=target_pool.shard_spec, + shard_group=target_pool.shard_group, + ) + if self.use_mla_backend: + return PageInterleaveMLATokenToKVPool( + **kwargs, + kv_lora_rank=self.model_config.kv_lora_rank, + qk_rope_head_dim=self.model_config.qk_rope_head_dim, + ) + return PageInterleaveMHATokenToKVPool( + **kwargs, + head_num=target_pool.head_num, + head_dim=self.model_config.head_dim, + v_head_dim=self.model_config.v_head_dim, + enable_alt_stream=not get_disagg().enable_pdmux, + ) + def _build_dsv4_kv_pool( self, *, diff --git a/python/sglang/srt/mem_cache/page_interleave.py b/python/sglang/srt/mem_cache/page_interleave.py index 3764cca21548..ae967e41341c 100644 --- a/python/sglang/srt/mem_cache/page_interleave.py +++ b/python/sglang/srt/mem_cache/page_interleave.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator + from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin logger = logging.getLogger(__name__) @@ -127,11 +128,48 @@ def get_kv_shard_group(use_mla_backend: bool) -> GroupCoordinator: return cp_group +def get_shared_kv_shard_pool( + kvc: KVCacheConfigurator, +) -> Optional[PageInterleaveKVPoolMixin]: + """Find a draft's sharded target through their shared logical allocator.""" + if not kvc.is_draft_worker: + return None + + from sglang.srt.mem_cache.allocator.page_interleave import ( + PageInterleavePoolAllocator, + ) + + allocator = kvc.token_to_kv_pool_allocator + if not isinstance(allocator, PageInterleavePoolAllocator): + return None + + from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveKVPoolMixin + + pool = allocator.get_kvcache() + if not isinstance(pool, PageInterleaveKVPoolMixin): + raise ValueError("A sharded MTP draft requires the target's sharded KV pool.") + if ( + allocator.shard_spec != pool.shard_spec + or allocator.shard_size != pool.shard_spec.shard_size + or allocator.page_size != pool.page_size + or allocator.size != pool.size * allocator.shard_size + ): + raise ValueError( + "The shared allocator and target KV pool have different shard layouts." + ) + return pool + + def get_kv_shard_group_info( kvc: KVCacheConfigurator, ) -> Tuple[Optional[int], int]: """``(shard_rank, shard_size)`` for the KV pool; ``(None, 1)`` disables.""" - if kvc.is_draft_worker or not get_parallel().enable_kv_cache_sharding: + if kvc.is_draft_worker: + pool = get_shared_kv_shard_pool(kvc) + if pool is None: + return None, 1 + return pool.shard_spec.shard_rank, pool.shard_spec.shard_size + if not get_parallel().enable_kv_cache_sharding: return None, 1 group = get_kv_shard_group(kvc.use_mla_backend) if group.world_size <= 1: @@ -163,7 +201,14 @@ def make_page_shard_spec(kvc: KVCacheConfigurator) -> Optional[PageShardSpec]: or chunk size. Both scratch slots are charged against HBM before sizing persistent KV; large model contexts therefore reduce persistent capacity. Admission still checks the actual batch's padded prefix and chunk spans. + + MTP drafts inherit the target's exact bounds: both pools serve the same + scheduled batch and logical locations, even if their model configs differ + in context length. """ + if kvc.is_draft_worker: + pool = get_shared_kv_shard_pool(kvc) + return pool.shard_spec if pool is not None else None shard_rank, shard_size = get_kv_shard_group_info(kvc) if shard_rank is None: return None @@ -181,11 +226,34 @@ def make_page_shard_spec(kvc: KVCacheConfigurator) -> Optional[PageShardSpec]: ) -def compute_page_shard_scratch_bytes(kvc: KVCacheConfigurator) -> int: +def compute_page_shard_scratch_bytes( + kvc: KVCacheConfigurator, *, include_mtp: bool = False +) -> int: """Fixed HBM cost of the double-buffered assembly scratch, charged against the KV budget before pool sizing. Use the same spec as the pool builders: - two slots, each holding ONE layer's ``[prefix | chunk | trash page]``.""" + two slots, each holding ONE layer's ``[prefix | chunk | trash page]``. + + ``include_mtp`` additionally budgets a target's single EAGLE MTP runner + at the same KV geometry, as required by the draft pool builder. Its two + scratch slots are reused across draft layers. Target sizing integration + can opt into this estimate without changing the per-pool default. + """ spec = make_page_shard_spec(kvc) if spec is None: return 0 - return 2 * spec.scratch_rows * _page_shard_row_bytes(kvc) + scratch_bytes = 2 * spec.scratch_rows * _page_shard_row_bytes(kvc) + if include_mtp: + from sglang.srt.runtime_context import get_spec + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + if ( + kvc.is_draft_worker + or kvc.spec_algorithm != SpeculativeAlgorithm.EAGLE + or (kvc.spec_aux_config.eagle_draft_num_layers or 0) <= 0 + or get_spec().enable_multi_layer_eagle + ): + raise ValueError( + "MTP scratch budgeting requires a target with one EAGLE MTP runner." + ) + scratch_bytes *= 2 + return scratch_bytes diff --git a/python/sglang/srt/mem_cache/page_interleave_pool.py b/python/sglang/srt/mem_cache/page_interleave_pool.py index 2a932b34a8f1..c97773d5774a 100644 --- a/python/sglang/srt/mem_cache/page_interleave_pool.py +++ b/python/sglang/srt/mem_cache/page_interleave_pool.py @@ -123,6 +123,7 @@ def _init_page_shard_state( assert spec.chunk_tokens % spec.page_size == 0 self.shard_spec = spec + self.shard_group = shard_group self.placement = PageInterleavePlacement(spec) self.shard_rank = spec.shard_rank self.shard_size = spec.shard_size diff --git a/test/manual/test_page_interleave_mtp.py b/test/manual/test_page_interleave_mtp.py new file mode 100644 index 000000000000..2499aaf96833 --- /dev/null +++ b/test/manual/test_page_interleave_mtp.py @@ -0,0 +1,389 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Two-GPU MTP KV-pool correctness, without model weights or a serving stack. + +Target and single-layer MTP pools share logical token locations but own their +storage, prefix gathers, and scratch plans. Check every persisted physical row +and every assembled prefix/chunk row against canonical, exactly representable +BF16 values. Exercise MLA over TP and MHA over CP, local and nonzero draft layer +IDs, uneven owner counts, fragmented pages, and a partially filled final page. +Both MLA write APIs, repeated reads, and plan changes must preserve independent +target/draft state across batches. Draft pools use the real configurator and +share the target's allocator. This does not run a model forward or a P/D transfer. + +Run: CUDA_VISIBLE_DEVICES=0,1 python test/manual/test_page_interleave_mtp.py +""" + +import argparse +import os +from types import SimpleNamespace + +import torch +import torch.multiprocessing as mp + +WORLD = 2 +PAGE_SIZE = 16 +SIZE = 32 * PAGE_SIZE # Physical token capacity of each pool on each rank. +TARGET_LAYERS = 4 +KV_LORA_RANK = 128 +QK_ROPE = 32 +HEAD_NUM = 2 +HEAD_DIM = 32 +DTYPE = torch.bfloat16 +UNWRITTEN = -256 + + +def _dist_init(rank, port, kind): + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + RANK=str(rank), + WORLD_SIZE=str(WORLD), + ) + os.environ.setdefault("no_proxy", "127.0.0.1,localhost") + torch.cuda.set_device(rank) + + from sglang.srt.distributed import ( + init_distributed_environment, + initialize_model_parallel, + ) + from sglang.srt.runtime_context import publish + from sglang.srt.server_args import ServerArgs + from sglang.test.test_utils import publish_build_topology + + init_distributed_environment( + world_size=WORLD, + rank=rank, + local_rank=rank, + distributed_init_method=f"tcp://127.0.0.1:{port}", + backend="nccl", + ) + cp_size = WORLD if kind == "mha" else 1 + publish( + ServerArgs(model_path="dummy", tp_size=WORLD, attn_cp_size=cp_size), + role="scheduler", + ) + publish_build_topology(tp_size=WORLD, attn_cp_size=cp_size, world_rank=rank) + initialize_model_parallel() + + +def _make_pool(rank, group, kind, start_layer, layer_num): + from sglang.srt.mem_cache.page_interleave import PageShardSpec + from sglang.srt.mem_cache.page_interleave_pool import ( + PageInterleaveMHATokenToKVPool, + PageInterleaveMLATokenToKVPool, + ) + + kwargs = dict( + size=SIZE, + page_size=PAGE_SIZE, + dtype=DTYPE, + layer_num=layer_num, + device=f"cuda:{rank}", + enable_memory_saver=False, + start_layer=start_layer, + end_layer=start_layer + layer_num - 1, + shard_spec=PageShardSpec( + shard_rank=group.rank_in_group, + shard_size=WORLD, + page_size=PAGE_SIZE, + max_prefix_tokens=16 * WORLD * PAGE_SIZE, + chunk_tokens=4 * PAGE_SIZE, + ), + shard_group=group, + ) + if kind == "mla": + pool = PageInterleaveMLATokenToKVPool( + **kwargs, kv_lora_rank=KV_LORA_RANK, qk_rope_head_dim=QK_ROPE + ) + else: + pool = PageInterleaveMHATokenToKVPool( + **kwargs, + head_num=HEAD_NUM, + head_dim=HEAD_DIM, + enable_alt_stream=False, + ) + for local_layer in range(layer_num): + for buffer in _buffers(pool, kind, local_layer): + buffer.fill_(UNWRITTEN) + return pool + + +def _buffers(pool, kind, local_layer): + if kind == "mla": + return (pool.kv_buffer[local_layer],) + return (pool.k_buffer[local_layer], pool.v_buffer[local_layer]) + + +def _make_draft_pool(target, allocator, kind, start_layer): + from sglang.srt.mem_cache.kv_cache_configurator import ( + KVCacheConfigurator, + _PoolSizes, + ) + from sglang.srt.runtime_context import get_context + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + # Skip model loading while using the production pool dispatch and sharing + # the target allocator. A replicated draft pool must fail this test. + kvc = KVCacheConfigurator.__new__(KVCacheConfigurator) + kvc.device = target.device + kvc.is_draft_worker = True + kvc.page_size = PAGE_SIZE + kvc.use_mla_backend = kind == "mla" + kvc.kv_cache_dtype = DTYPE + kvc.kv_cache_dtype_str = "bfloat16" + kvc.post_capture_kv_active = False + kvc.is_hybrid_swa = False + kvc.sliding_window_size = None + kvc.mambaish_config = None + kvc.token_to_kv_pool_allocator = allocator + kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE + kvc.model_config = SimpleNamespace( + num_nextn_predict_layers=1, + hf_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]), + kv_lora_rank=KV_LORA_RANK, + qk_rope_head_dim=QK_ROPE, + head_dim=HEAD_DIM, + v_head_dim=HEAD_DIM, + get_num_kv_heads=lambda *_args: HEAD_NUM, + ) + kvc.layer_info = SimpleNamespace( + start_layer=start_layer, end_layer=start_layer + 1, num_effective_layers=1 + ) + sizes = _PoolSizes( + max_total_num_tokens=target.size, + max_running_requests=1, + full_max_total_num_tokens=None, + swa_max_total_num_tokens=None, + c4_max_total_num_tokens=0, + c128_max_total_num_tokens=0, + c4_state_pool_size=0, + c128_state_pool_size=0, + c4_state_dtype=None, + c128_state_dtype=None, + ) + with get_context().override_server_args( + page_size=PAGE_SIZE, + tp_size=WORLD, + attn_cp_size=WORLD if kind == "mha" else 1, + speculative_algorithm="EAGLE", + enable_multi_layer_eagle=False, + ): + pool = kvc._build_token_to_kv_pool( + sizes=sizes, + is_dsa_model=False, + is_dsv4_model=False, + req_to_token_pool=None, + ) + assert pool.shard_spec is target.shard_spec + assert pool.shard_group is target.shard_group + for buffer in _buffers(pool, kind, 0): + buffer.fill_(UNWRITTEN) + return pool + + +def _values(locs, width, tag, local_layer, is_v=False): + """Each location, feature, layer, and pool has distinguishable BF16 bytes. + + Small integers survive BF16 exactly; high-valued logical slot IDs plus tiny + feature offsets would round away the differences this test must detect. + The first two columns encode the full logical slot, avoiding hash aliases. + """ + features = torch.arange(width, device=locs.device) + values = ( + locs[:, None] * 37 + + features[None, :] * 17 + + tag * 73 + + local_layer * 11 + + int(is_v) * 97 + ) % 251 - 125 + values[:, 0] = locs % 128 + values[:, 1] = locs // 128 + values[:, 2] = tag + values[:, 3] = local_layer + values[:, 4] = int(is_v) + return values.to(DTYPE) + + +def _canonical(kind, locs, tag, local_layer): + if kind == "mla": + return (_values(locs, KV_LORA_RANK + QK_ROPE, tag, local_layer).unsqueeze(1),) + return tuple( + _values(locs, HEAD_NUM * HEAD_DIM, tag, local_layer, is_v).reshape( + -1, HEAD_NUM, HEAD_DIM + ) + for is_v in (False, True) + ) + + +def _assert_bytes(got, expected, label): + assert got.shape == expected.shape, label + assert got.dtype == expected.dtype == DTYPE, label + assert torch.equal( + got.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8) + ), f"{label}: KV bytes differ" + + +def _logical_chain(device): + # Owners alternate from rank 1, while local pages run out of order. Some + # logical slots exceed SIZE, catching a replicated draft indexing them raw. + local_pages = torch.tensor([25, 3, 20, 6, 28, 4, 30, 8, 22], device=device) + owners = (1 + torch.arange(local_pages.numel(), device=device)) % WORLD + pages = local_pages * WORLD + owners + return ( + pages[:, None] * PAGE_SIZE + torch.arange(PAGE_SIZE, device=device) + ).reshape(-1) + + +def _write(pool, kind, locs, tag, local_layer, epoch): + layer = SimpleNamespace(layer_id=pool.start_layer + local_layer) + values = _canonical(kind, locs, tag, local_layer) + if kind == "mha": + pool.set_kv_buffer(layer, locs, *values) + elif epoch % 2: + # Both MLA write entry points must owner-filter and stage MTP rows. + pool.set_mla_kv_buffer( + layer, locs, values[0][..., :KV_LORA_RANK], values[0][..., KV_LORA_RANK:] + ) + else: + pool.set_kv_buffer(layer, locs, values[0], values[0][..., :KV_LORA_RANK]) + + +def _check_persisted(pool, kind, locs, tag, rank): + owned_locs = locs[(locs // PAGE_SIZE) % WORLD == rank] + local_rows = owned_locs // (WORLD * PAGE_SIZE) * PAGE_SIZE + owned_locs % PAGE_SIZE + for local_layer in range(pool.layer_num): + canonical = _canonical(kind, owned_locs, tag, local_layer) + for buffer, values in zip(_buffers(pool, kind, local_layer), canonical): + expected = torch.full_like(buffer, UNWRITTEN) + expected[local_rows] = values + _assert_bytes( + buffer, expected, f"{kind} pool {tag} owned layer {local_layer}" + ) + + +def _check_scratch(pool, kind, locs, prefix_len, tag, local_layer): + layer_id = pool.start_layer + local_layer + rows = pool.translate_loc_to_scratch(locs) + expected = _canonical(kind, locs, tag, local_layer) + _assert_bytes( + pool.get_key_buffer(layer_id)[rows], expected[0], f"{kind} pool {tag} key" + ) + _assert_bytes( + pool.get_value_buffer(layer_id)[rows], + expected[0][..., :KV_LORA_RANK] if kind == "mla" else expected[1], + f"{kind} pool {tag} value", + ) + if kind == "mla" and prefix_len: + # Arbitrary prefix indices also exercise the chunked MLA consumer. + subset = locs[3:prefix_len:7].clone() + k_nope, k_rope = pool.get_mla_kv_buffer( + SimpleNamespace(layer_id=layer_id), subset, DTYPE + ) + reference = _canonical(kind, subset, tag, local_layer)[0] + _assert_bytes(k_nope, reference[..., :KV_LORA_RANK], f"pool {tag} latent") + _assert_bytes(k_rope, reference[..., KV_LORA_RANK:], f"pool {tag} rope") + + +def _run(rank, port, kind): + _dist_init(rank, port, kind) + from sglang.srt.mem_cache.allocator.page_interleave import ( + PageInterleavePoolAllocator, + ) + from sglang.srt.mem_cache.page_interleave import get_kv_shard_group + + group = get_kv_shard_group(use_mla_backend=kind == "mla") + assert group.world_size == WORLD + target = _make_pool(rank, group, kind, start_layer=0, layer_num=TARGET_LAYERS) + allocator = PageInterleavePoolAllocator( + size=target.size, + physical_page_size=PAGE_SIZE, + shard_size=WORLD, + dtype=DTYPE, + device=target.device, + kvcache=target, + need_sort=False, + shard_spec=target.shard_spec, + ) + # DeepSeek's NextN uses local layer 0. Also cover an absolute/offset layer ID + # so the one-layer gather terminates correctly on either scratch-slot parity. + drafts = [ + _make_draft_pool(target, allocator, kind, start_layer=start) + for start in (0, 61) + ] + pools = [target, *drafts] + for other in drafts: + assert other._page_pos.data_ptr() != target._page_pos.data_ptr() + assert other.kv_gather_comm is not target.kv_gather_comm + for slot, target_slot in zip(other._slots, target._slots): + for name in slot.tensors: + assert ( + slot.tensors[name].data_ptr() + != target_slot.tensors[name].data_ptr() + ) + + locs = _logical_chain(target.device) + req_to_token = locs.to(torch.int32).unsqueeze(0) + req_indices = torch.tensor([0], dtype=torch.int64, device=locs.device) + # Each cached prefix is page aligned. The third batch writes only five + # tokens of its tail page; the fourth recomputes and completes that page. + batches = [(0, 48), (48, 96), (96, 133), (128, 144)] + for epoch, (prefix_len, seq_len) in enumerate(batches, start=1): + active_locs = locs[:seq_len] + chunk_locs = locs[prefix_len:seq_len] + target_epoch = target._epoch + for tag, pool in enumerate(pools, start=1): + pool.begin_shard_extend(req_to_token, req_indices, [prefix_len], [seq_len]) + assert pool._epoch == epoch + for local_layer in range(pool.layer_num): + _write(pool, kind, chunk_locs, tag, local_layer, epoch) + _check_scratch(pool, kind, active_locs, prefix_len, tag, local_layer) + _check_persisted(pool, kind, active_locs, tag, rank) + # MTP metadata/writes/reads must leave the target's last resident + # layer and gather plan usable, despite identical logical locations. + assert target._epoch == target_epoch + 1 + _check_scratch(target, kind, active_locs, prefix_len, 1, TARGET_LAYERS - 1) + if pool.layer_num == 1 and prefix_len: + resident = pool._slots[pool.start_layer % 2] + assert resident.resident_key == (pool.start_layer, epoch) + assert pool._slots[(pool.start_layer + 1) % 2].resident_key is None + if rank == 0: + print(f"PASS: {kind.upper()} target/MTP prefix={prefix_len}, seq={seq_len}") + + # Deactivating one draft's plan must not deactivate the other pools. + drafts[0].end_shard_extend() + assert not drafts[0]._shard_extend_active + assert target._shard_extend_active and drafts[1]._shard_extend_active + _check_scratch(drafts[1], kind, locs, 128, 3, 0) + _check_scratch(target, kind, locs, 128, 1, TARGET_LAYERS - 1) + torch.cuda.synchronize() + torch.distributed.barrier() + torch.distributed.destroy_process_group() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pool", choices=("all", "mla", "mha"), default="all") + parser.add_argument("--port", type=int, default=29821) + args = parser.parse_args() + if torch.cuda.device_count() < WORLD: + raise RuntimeError("MTP KV sharding correctness requires two visible GPUs") + kinds = ("mla", "mha") if args.pool == "all" else (args.pool,) + for offset, kind in enumerate(kinds): + mp.spawn(_run, args=(args.port + offset, kind), nprocs=WORLD, join=True) + print("PASS: MTP page-interleave KV cache correctness") + + +if __name__ == "__main__": + main() diff --git a/test/registered/unit/mem_cache/test_mtp_kv_shard_config.py b/test/registered/unit/mem_cache/test_mtp_kv_shard_config.py new file mode 100644 index 000000000000..6700d607b0d3 --- /dev/null +++ b/test/registered/unit/mem_cache/test_mtp_kv_shard_config.py @@ -0,0 +1,365 @@ +"""CPU coverage of MTP pool construction against a shared sharded allocator.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.mem_cache import page_interleave +from sglang.srt.mem_cache.allocator.page_interleave import PageInterleavePoolAllocator +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator +from sglang.srt.mem_cache.page_interleave import ( + PageShardSpec, + compute_page_shard_scratch_bytes, + get_kv_shard_group_info, + get_shared_kv_shard_pool, + make_page_shard_spec, +) +from sglang.srt.mem_cache.page_interleave_pool import ( + PageInterleaveMHATokenToKVPool, + PageInterleaveMLATokenToKVPool, +) +from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig +from sglang.srt.runtime_context import ( + get_context, + get_memory, + get_parallel, + get_schedule, + get_spec, +) +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _make_draft_configurator(*, use_mla=True): + # Keep the production pool/allocator types, but allocate only the CPU + # allocator metadata. Pool construction itself is mocked in these tests. + pool_class = ( + PageInterleaveMLATokenToKVPool if use_mla else PageInterleaveMHATokenToKVPool + ) + target_pool = pool_class.__new__(pool_class) + target_pool.shard_spec = PageShardSpec( + shard_rank=1, + shard_size=2, + page_size=16, + max_prefix_tokens=128, + chunk_tokens=64, + ) + target_pool.shard_group = SimpleNamespace(rank_in_group=1, world_size=2) + target_pool.size = 128 + target_pool.page_size = 16 + target_pool.dtype = torch.bfloat16 + target_pool.kv_lora_rank = 8 + target_pool.qk_rope_head_dim = 4 + target_pool.head_num = 2 + target_pool.head_dim = 4 + target_pool.v_head_dim = 6 + allocator = PageInterleavePoolAllocator( + size=target_pool.size, + physical_page_size=16, + shard_size=2, + dtype=torch.bfloat16, + device="cpu", + kvcache=target_pool, + need_sort=False, + shard_spec=target_pool.shard_spec, + ) + kvc = KVCacheConfigurator.__new__(KVCacheConfigurator) + kvc.device = "cpu" + kvc.gpu_id = 0 + kvc.is_draft_worker = True + kvc.page_size = 16 + kvc.use_mla_backend = use_mla + kvc.kv_cache_dtype = torch.bfloat16 + kvc.kv_cache_dtype_str = "bfloat16" + kvc.token_to_kv_pool_allocator = allocator + kvc.req_to_token_pool = None + kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE + kvc.model_config = SimpleNamespace( + num_nextn_predict_layers=1, + context_len=100, + kv_lora_rank=8, + qk_rope_head_dim=4, + head_dim=4, + v_head_dim=6, + get_num_kv_heads=lambda tp, dcp=1: 4 // tp, + hf_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]), + ) + kvc.layer_info = SimpleNamespace(start_layer=5, end_layer=6, num_effective_layers=1) + kvc.spec_aux_config = SimpleNamespace(eagle_draft_num_layers=1) + kvc.is_hybrid_swa = False + kvc.is_hybrid_swa_compress = False + kvc.sliding_window_size = None + kvc.mambaish_config = None + kvc.hybrid_gdn_config = None + kvc.hybrid_kda_config = None + kvc.post_capture_kv_active = False + kvc.memory_pool_config = MemoryPoolConfig( + max_total_num_tokens=128, max_running_requests=4 + ) + return kvc, target_pool + + +def _build_draft_pool(kvc, *, is_dsa_model=False, is_dsv4_model=False): + sizes = kvc._derive_pool_sizes(config=kvc.memory_pool_config) + return kvc._build_token_to_kv_pool( + sizes=sizes, + is_dsa_model=is_dsa_model, + is_dsv4_model=is_dsv4_model, + req_to_token_pool=None, + ) + + +class TestMTPKVShardConfig(CustomTestCase): + def setUp(self): + super().setUp() + self.enterContext( + get_context().override_server_args( + page_size=16, + chunked_prefill_size=64, + enable_multi_layer_eagle=False, + ) + ) + self.enterContext( + get_parallel().override( + tp_size=2, attn_tp_size=2, moe_tp_size=2, attn_dcp_size=1 + ) + ) + + def test_inherits_exact_target_spec_without_resolving_current_group(self): + for use_mla in (True, False): + with self.subTest(use_mla=use_mla): + kvc, target = _make_draft_configurator(use_mla=use_mla) + with patch.object( + page_interleave, "get_kv_shard_group", side_effect=AssertionError + ): + self.assertIs(get_shared_kv_shard_pool(kvc), target) + self.assertEqual(get_kv_shard_group_info(kvc), (1, 2)) + self.assertIs(make_page_shard_spec(kvc), target.shard_spec) + + def test_builds_separate_mtp_pool_with_target_layout_and_draft_layer_offsets(self): + for use_mla in (True, False): + with self.subTest(use_mla=use_mla): + kvc, target = _make_draft_configurator(use_mla=use_mla) + pool_class = type(target) + # Patch __init__, not the class: target validation still sees + # the real type, and construction returns a distinct instance. + with patch.object(pool_class, "__init__", return_value=None) as init: + pool = _build_draft_pool(kvc) + self.assertIsInstance(pool, pool_class) + self.assertIsNot(pool, target) + kwargs = init.call_args.kwargs + self.assertEqual(kwargs["size"], target.size) + self.assertIs(kwargs["shard_spec"], target.shard_spec) + self.assertIs(kwargs["shard_group"], target.shard_group) + self.assertEqual(kwargs["page_size"], 16) + self.assertEqual(kwargs["layer_num"], 1) + self.assertEqual(kwargs["start_layer"], 5) + self.assertEqual(kwargs["end_layer"], 6) + self.assertEqual(kwargs["dtype"], torch.bfloat16) + sizes = kvc._derive_pool_sizes(config=kvc.memory_pool_config) + self.assertIs( + kvc._build_token_to_kv_pool_allocator( + sizes=sizes, + token_to_kv_pool=pool, + is_dsv4_model=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=kvc.token_to_kv_pool_allocator, + ), + kvc.token_to_kv_pool_allocator, + ) + + def test_draft_without_shared_sharded_allocator_keeps_existing_pool(self): + for allocator in (None, object()): + with self.subTest(allocator=allocator): + kvc, _ = _make_draft_configurator() + kvc.token_to_kv_pool_allocator = allocator + self.assertIsNone(get_shared_kv_shard_pool(kvc)) + self.assertEqual(get_kv_shard_group_info(kvc), (None, 1)) + self.assertIsNone(make_page_shard_spec(kvc)) + self.assertEqual(compute_page_shard_scratch_bytes(kvc), 0) + with patch.object(KVCacheConfigurator, "_build_mla_kv_pool") as build: + self.assertIs(_build_draft_pool(kvc), build.return_value) + build.assert_called_once_with(max_total_num_tokens=128) + + def test_target_without_shared_allocator_keeps_existing_pool(self): + kvc, _ = _make_draft_configurator() + kvc.is_draft_worker = False + kvc.token_to_kv_pool_allocator = None + # Ordinary target construction does not require the global sharding + # flag that has not yet been integrated into server arguments. + with patch.object(KVCacheConfigurator, "_build_mla_kv_pool") as build: + self.assertIs(_build_draft_pool(kvc), build.return_value) + build.assert_called_once_with(max_total_num_tokens=128) + + def test_rejects_invalid_shared_target_pool_and_spec(self): + kvc, _ = _make_draft_configurator() + kvc.token_to_kv_pool_allocator._kvcache = object() + with self.assertRaisesRegex(ValueError, "target's sharded KV pool"): + get_shared_kv_shard_pool(kvc) + mismatched_spec = PageShardSpec( + shard_rank=0, + shard_size=2, + page_size=16, + max_prefix_tokens=128, + chunk_tokens=64, + ) + for attribute, value in ( + ("shard_spec", mismatched_spec), + ("shard_size", 4), + ("page_size", 32), + ("size", 128), + ): + with self.subTest(attribute=attribute): + kvc, _ = _make_draft_configurator() + setattr(kvc.token_to_kv_pool_allocator, attribute, value) + with self.assertRaisesRegex(ValueError, "different shard layouts"): + get_shared_kv_shard_pool(kvc) + + def test_rejects_page_dtype_and_attention_mode_mismatches(self): + for attribute, value, message in ( + ("page_size", 32, "physical page size"), + ("kv_cache_dtype", torch.float16, "geometry and dtype"), + ("use_mla_backend", False, "attention geometry"), + ): + with self.subTest(attribute=attribute): + kvc, _ = _make_draft_configurator() + setattr(kvc, attribute, value) + with self.assertRaisesRegex(ValueError, message): + _build_draft_pool(kvc) + + def test_rejects_mla_and_mha_geometry_mismatches(self): + for use_mla, attribute, value in ( + (True, "kv_lora_rank", 16), + (True, "qk_rope_head_dim", 8), + (False, "head_dim", 8), + (False, "v_head_dim", 8), + (False, "get_num_kv_heads", lambda tp, dcp=1: 4), + ): + with self.subTest(use_mla=use_mla, attribute=attribute): + kvc, _ = _make_draft_configurator(use_mla=use_mla) + setattr(kvc.model_config, attribute, value) + with self.assertRaisesRegex(ValueError, "geometry and dtype"): + _build_draft_pool(kvc) + + def test_rejects_draft_capacity_different_from_target(self): + for capacity in (64, 256): + with self.subTest(capacity=capacity): + kvc, _ = _make_draft_configurator() + kvc.memory_pool_config = MemoryPoolConfig( + max_total_num_tokens=capacity, max_running_requests=4 + ) + with self.assertRaisesRegex(ValueError, "same per-rank capacity"): + _build_draft_pool(kvc) + + def test_rejects_other_draft_algorithms_and_multi_runner_mode(self): + for algorithm in ( + SpeculativeAlgorithm.EAGLE3, + SpeculativeAlgorithm.DFLASH, + SpeculativeAlgorithm.FROZEN_KV_MTP, + ): + with self.subTest(algorithm=algorithm): + kvc, _ = _make_draft_configurator() + kvc.spec_algorithm = algorithm + with self.assertRaisesRegex(ValueError, "EAGLE MTP"): + _build_draft_pool(kvc) + kvc, _ = _make_draft_configurator() + with get_spec().override(enable_multi_layer_eagle=True): + with self.assertRaisesRegex(ValueError, "EAGLE MTP"): + _build_draft_pool(kvc) + for layers in (None, 0): + with self.subTest(layers=layers): + kvc.model_config.num_nextn_predict_layers = layers + with self.assertRaisesRegex(ValueError, "EAGLE MTP"): + _build_draft_pool(kvc) + + def test_rejects_unsupported_model_pool_layouts(self): + for options in ({"is_dsa_model": True}, {"is_dsv4_model": True}): + with self.subTest(options=options): + kvc, _ = _make_draft_configurator() + with self.assertRaisesRegex(ValueError, "dense MLA and MHA"): + _build_draft_pool(kvc, **options) + for attribute, value in ( + ("is_hybrid_swa", True), + ("mambaish_config", object()), + ("sliding_window_size", 128), + ): + with self.subTest(attribute=attribute): + kvc, _ = _make_draft_configurator() + setattr(kvc, attribute, value) + with self.assertRaisesRegex(ValueError, "dense MLA and MHA"): + _build_draft_pool(kvc) + + def test_rejects_dcp(self): + kvc, _ = _make_draft_configurator() + with get_parallel().override(attn_dcp_size=2): + with self.assertRaisesRegex(ValueError, "incompatible with DCP"): + _build_draft_pool(kvc) + + def test_rejects_special_kv_storage_layouts(self): + for attribute, value in ( + ("post_capture_kv_active", True), + ("kv_cache_dtype_str", "mxfp8"), + ): + with self.subTest(attribute=attribute): + kvc, _ = _make_draft_configurator() + setattr(kvc, attribute, value) + with self.assertRaisesRegex(ValueError, "plain per-layer KV layout"): + _build_draft_pool(kvc) + kvc, _ = _make_draft_configurator() + with get_memory().override(enable_page_major_kv_layout=True): + with self.assertRaisesRegex(ValueError, "plain per-layer KV layout"): + _build_draft_pool(kvc) + + def test_draft_scratch_uses_inherited_bounds_and_has_one_buffer_pair(self): + for use_mla in (True, False): + with self.subTest(use_mla=use_mla): + kvc, _ = _make_draft_configurator(use_mla=use_mla) + kvc.model_config.context_len = 9999 + row_bytes = (8 + 4) * 2 if use_mla else 2 * (4 + 6) * 2 + expected = 2 * (128 + 64 + 16) * row_bytes + with get_schedule().override(chunked_prefill_size=1024): + self.assertEqual(compute_page_shard_scratch_bytes(kvc), expected) + with self.assertRaisesRegex(ValueError, "target"): + compute_page_shard_scratch_bytes(kvc, include_mtp=True) + + def test_target_plus_mtp_scratch_preserves_eight_context_budget(self): + kvc, target = _make_draft_configurator() + kvc.is_draft_worker = False + kvc.token_to_kv_pool_allocator = None + # Some architectures populate num_nextn_predict_layers only in their + # draft config; target sizing uses the resolved draft layer count. + kvc.model_config.num_nextn_predict_layers = None + kvc.layer_info = SimpleNamespace( + start_layer=0, end_layer=8, num_effective_layers=8 + ) + parallel = SimpleNamespace( + enable_kv_cache_sharding=True, + attn_cp_group=SimpleNamespace(world_size=1), + attn_tp_group=target.shard_group, + attn_tp_size=2, + ) + # 100 tokens rounds to 128 per context before multiplying by eight. + single_pool_scratch = 2 * (8 * 128 + 64 + 16) * (8 + 4) * 2 + with patch.object(page_interleave, "get_parallel", return_value=parallel): + for layers in (1, 3): + with self.subTest(layers=layers): + kvc.spec_aux_config.eagle_draft_num_layers = layers + self.assertEqual( + compute_page_shard_scratch_bytes(kvc), single_pool_scratch + ) + self.assertEqual( + compute_page_shard_scratch_bytes(kvc, include_mtp=True), + 2 * single_pool_scratch, + ) + kvc.spec_aux_config.eagle_draft_num_layers = None + with self.assertRaisesRegex(ValueError, "EAGLE MTP"): + compute_page_shard_scratch_bytes(kvc, include_mtp=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index c3007f11aab6..ad5fea9b0212 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -141,6 +141,7 @@ def _fixture( ) kvc = SimpleNamespace( is_draft_worker=draft, + token_to_kv_pool_allocator=None, use_mla_backend=use_mla, page_size=16, kv_cache_dtype=torch.bfloat16, From 52219b057392b6628e5aee4dc04e5448dc55aa4a Mon Sep 17 00:00:00 2001 From: "daowu.hzy" Date: Wed, 7 Oct 2026 21:04:26 +0800 Subject: [PATCH 10/11] test(kv-cache): reconcile upstream MTP tests with current main Preserve the final test coverage from Shunkangz/sglang PR #40929 commits ea0efc16ea, 16c7162899 and 156b0db949 while retaining current-main, PD and DSA tests. Keep current-main import boundaries in the pool implementations. Co-authored-by: shunkangz <182541032+Shunkangz@users.noreply.github.com> --- test/manual/test_page_interleave_mtp.py | 389 ------------ .../unit/managers/test_prefill_adder.py | 1 + .../mem_cache/test_mtp_kv_shard_config.py | 365 ----------- .../mem_cache/test_page_interleave_shard.py | 585 +++++++++++++++++- 4 files changed, 583 insertions(+), 757 deletions(-) delete mode 100644 test/manual/test_page_interleave_mtp.py delete mode 100644 test/registered/unit/mem_cache/test_mtp_kv_shard_config.py diff --git a/test/manual/test_page_interleave_mtp.py b/test/manual/test_page_interleave_mtp.py deleted file mode 100644 index 2499aaf96833..000000000000 --- a/test/manual/test_page_interleave_mtp.py +++ /dev/null @@ -1,389 +0,0 @@ -# Copyright 2023-2026 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Two-GPU MTP KV-pool correctness, without model weights or a serving stack. - -Target and single-layer MTP pools share logical token locations but own their -storage, prefix gathers, and scratch plans. Check every persisted physical row -and every assembled prefix/chunk row against canonical, exactly representable -BF16 values. Exercise MLA over TP and MHA over CP, local and nonzero draft layer -IDs, uneven owner counts, fragmented pages, and a partially filled final page. -Both MLA write APIs, repeated reads, and plan changes must preserve independent -target/draft state across batches. Draft pools use the real configurator and -share the target's allocator. This does not run a model forward or a P/D transfer. - -Run: CUDA_VISIBLE_DEVICES=0,1 python test/manual/test_page_interleave_mtp.py -""" - -import argparse -import os -from types import SimpleNamespace - -import torch -import torch.multiprocessing as mp - -WORLD = 2 -PAGE_SIZE = 16 -SIZE = 32 * PAGE_SIZE # Physical token capacity of each pool on each rank. -TARGET_LAYERS = 4 -KV_LORA_RANK = 128 -QK_ROPE = 32 -HEAD_NUM = 2 -HEAD_DIM = 32 -DTYPE = torch.bfloat16 -UNWRITTEN = -256 - - -def _dist_init(rank, port, kind): - os.environ.update( - MASTER_ADDR="127.0.0.1", - MASTER_PORT=str(port), - RANK=str(rank), - WORLD_SIZE=str(WORLD), - ) - os.environ.setdefault("no_proxy", "127.0.0.1,localhost") - torch.cuda.set_device(rank) - - from sglang.srt.distributed import ( - init_distributed_environment, - initialize_model_parallel, - ) - from sglang.srt.runtime_context import publish - from sglang.srt.server_args import ServerArgs - from sglang.test.test_utils import publish_build_topology - - init_distributed_environment( - world_size=WORLD, - rank=rank, - local_rank=rank, - distributed_init_method=f"tcp://127.0.0.1:{port}", - backend="nccl", - ) - cp_size = WORLD if kind == "mha" else 1 - publish( - ServerArgs(model_path="dummy", tp_size=WORLD, attn_cp_size=cp_size), - role="scheduler", - ) - publish_build_topology(tp_size=WORLD, attn_cp_size=cp_size, world_rank=rank) - initialize_model_parallel() - - -def _make_pool(rank, group, kind, start_layer, layer_num): - from sglang.srt.mem_cache.page_interleave import PageShardSpec - from sglang.srt.mem_cache.page_interleave_pool import ( - PageInterleaveMHATokenToKVPool, - PageInterleaveMLATokenToKVPool, - ) - - kwargs = dict( - size=SIZE, - page_size=PAGE_SIZE, - dtype=DTYPE, - layer_num=layer_num, - device=f"cuda:{rank}", - enable_memory_saver=False, - start_layer=start_layer, - end_layer=start_layer + layer_num - 1, - shard_spec=PageShardSpec( - shard_rank=group.rank_in_group, - shard_size=WORLD, - page_size=PAGE_SIZE, - max_prefix_tokens=16 * WORLD * PAGE_SIZE, - chunk_tokens=4 * PAGE_SIZE, - ), - shard_group=group, - ) - if kind == "mla": - pool = PageInterleaveMLATokenToKVPool( - **kwargs, kv_lora_rank=KV_LORA_RANK, qk_rope_head_dim=QK_ROPE - ) - else: - pool = PageInterleaveMHATokenToKVPool( - **kwargs, - head_num=HEAD_NUM, - head_dim=HEAD_DIM, - enable_alt_stream=False, - ) - for local_layer in range(layer_num): - for buffer in _buffers(pool, kind, local_layer): - buffer.fill_(UNWRITTEN) - return pool - - -def _buffers(pool, kind, local_layer): - if kind == "mla": - return (pool.kv_buffer[local_layer],) - return (pool.k_buffer[local_layer], pool.v_buffer[local_layer]) - - -def _make_draft_pool(target, allocator, kind, start_layer): - from sglang.srt.mem_cache.kv_cache_configurator import ( - KVCacheConfigurator, - _PoolSizes, - ) - from sglang.srt.runtime_context import get_context - from sglang.srt.speculative.spec_info import SpeculativeAlgorithm - - # Skip model loading while using the production pool dispatch and sharing - # the target allocator. A replicated draft pool must fail this test. - kvc = KVCacheConfigurator.__new__(KVCacheConfigurator) - kvc.device = target.device - kvc.is_draft_worker = True - kvc.page_size = PAGE_SIZE - kvc.use_mla_backend = kind == "mla" - kvc.kv_cache_dtype = DTYPE - kvc.kv_cache_dtype_str = "bfloat16" - kvc.post_capture_kv_active = False - kvc.is_hybrid_swa = False - kvc.sliding_window_size = None - kvc.mambaish_config = None - kvc.token_to_kv_pool_allocator = allocator - kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE - kvc.model_config = SimpleNamespace( - num_nextn_predict_layers=1, - hf_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]), - kv_lora_rank=KV_LORA_RANK, - qk_rope_head_dim=QK_ROPE, - head_dim=HEAD_DIM, - v_head_dim=HEAD_DIM, - get_num_kv_heads=lambda *_args: HEAD_NUM, - ) - kvc.layer_info = SimpleNamespace( - start_layer=start_layer, end_layer=start_layer + 1, num_effective_layers=1 - ) - sizes = _PoolSizes( - max_total_num_tokens=target.size, - max_running_requests=1, - full_max_total_num_tokens=None, - swa_max_total_num_tokens=None, - c4_max_total_num_tokens=0, - c128_max_total_num_tokens=0, - c4_state_pool_size=0, - c128_state_pool_size=0, - c4_state_dtype=None, - c128_state_dtype=None, - ) - with get_context().override_server_args( - page_size=PAGE_SIZE, - tp_size=WORLD, - attn_cp_size=WORLD if kind == "mha" else 1, - speculative_algorithm="EAGLE", - enable_multi_layer_eagle=False, - ): - pool = kvc._build_token_to_kv_pool( - sizes=sizes, - is_dsa_model=False, - is_dsv4_model=False, - req_to_token_pool=None, - ) - assert pool.shard_spec is target.shard_spec - assert pool.shard_group is target.shard_group - for buffer in _buffers(pool, kind, 0): - buffer.fill_(UNWRITTEN) - return pool - - -def _values(locs, width, tag, local_layer, is_v=False): - """Each location, feature, layer, and pool has distinguishable BF16 bytes. - - Small integers survive BF16 exactly; high-valued logical slot IDs plus tiny - feature offsets would round away the differences this test must detect. - The first two columns encode the full logical slot, avoiding hash aliases. - """ - features = torch.arange(width, device=locs.device) - values = ( - locs[:, None] * 37 - + features[None, :] * 17 - + tag * 73 - + local_layer * 11 - + int(is_v) * 97 - ) % 251 - 125 - values[:, 0] = locs % 128 - values[:, 1] = locs // 128 - values[:, 2] = tag - values[:, 3] = local_layer - values[:, 4] = int(is_v) - return values.to(DTYPE) - - -def _canonical(kind, locs, tag, local_layer): - if kind == "mla": - return (_values(locs, KV_LORA_RANK + QK_ROPE, tag, local_layer).unsqueeze(1),) - return tuple( - _values(locs, HEAD_NUM * HEAD_DIM, tag, local_layer, is_v).reshape( - -1, HEAD_NUM, HEAD_DIM - ) - for is_v in (False, True) - ) - - -def _assert_bytes(got, expected, label): - assert got.shape == expected.shape, label - assert got.dtype == expected.dtype == DTYPE, label - assert torch.equal( - got.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8) - ), f"{label}: KV bytes differ" - - -def _logical_chain(device): - # Owners alternate from rank 1, while local pages run out of order. Some - # logical slots exceed SIZE, catching a replicated draft indexing them raw. - local_pages = torch.tensor([25, 3, 20, 6, 28, 4, 30, 8, 22], device=device) - owners = (1 + torch.arange(local_pages.numel(), device=device)) % WORLD - pages = local_pages * WORLD + owners - return ( - pages[:, None] * PAGE_SIZE + torch.arange(PAGE_SIZE, device=device) - ).reshape(-1) - - -def _write(pool, kind, locs, tag, local_layer, epoch): - layer = SimpleNamespace(layer_id=pool.start_layer + local_layer) - values = _canonical(kind, locs, tag, local_layer) - if kind == "mha": - pool.set_kv_buffer(layer, locs, *values) - elif epoch % 2: - # Both MLA write entry points must owner-filter and stage MTP rows. - pool.set_mla_kv_buffer( - layer, locs, values[0][..., :KV_LORA_RANK], values[0][..., KV_LORA_RANK:] - ) - else: - pool.set_kv_buffer(layer, locs, values[0], values[0][..., :KV_LORA_RANK]) - - -def _check_persisted(pool, kind, locs, tag, rank): - owned_locs = locs[(locs // PAGE_SIZE) % WORLD == rank] - local_rows = owned_locs // (WORLD * PAGE_SIZE) * PAGE_SIZE + owned_locs % PAGE_SIZE - for local_layer in range(pool.layer_num): - canonical = _canonical(kind, owned_locs, tag, local_layer) - for buffer, values in zip(_buffers(pool, kind, local_layer), canonical): - expected = torch.full_like(buffer, UNWRITTEN) - expected[local_rows] = values - _assert_bytes( - buffer, expected, f"{kind} pool {tag} owned layer {local_layer}" - ) - - -def _check_scratch(pool, kind, locs, prefix_len, tag, local_layer): - layer_id = pool.start_layer + local_layer - rows = pool.translate_loc_to_scratch(locs) - expected = _canonical(kind, locs, tag, local_layer) - _assert_bytes( - pool.get_key_buffer(layer_id)[rows], expected[0], f"{kind} pool {tag} key" - ) - _assert_bytes( - pool.get_value_buffer(layer_id)[rows], - expected[0][..., :KV_LORA_RANK] if kind == "mla" else expected[1], - f"{kind} pool {tag} value", - ) - if kind == "mla" and prefix_len: - # Arbitrary prefix indices also exercise the chunked MLA consumer. - subset = locs[3:prefix_len:7].clone() - k_nope, k_rope = pool.get_mla_kv_buffer( - SimpleNamespace(layer_id=layer_id), subset, DTYPE - ) - reference = _canonical(kind, subset, tag, local_layer)[0] - _assert_bytes(k_nope, reference[..., :KV_LORA_RANK], f"pool {tag} latent") - _assert_bytes(k_rope, reference[..., KV_LORA_RANK:], f"pool {tag} rope") - - -def _run(rank, port, kind): - _dist_init(rank, port, kind) - from sglang.srt.mem_cache.allocator.page_interleave import ( - PageInterleavePoolAllocator, - ) - from sglang.srt.mem_cache.page_interleave import get_kv_shard_group - - group = get_kv_shard_group(use_mla_backend=kind == "mla") - assert group.world_size == WORLD - target = _make_pool(rank, group, kind, start_layer=0, layer_num=TARGET_LAYERS) - allocator = PageInterleavePoolAllocator( - size=target.size, - physical_page_size=PAGE_SIZE, - shard_size=WORLD, - dtype=DTYPE, - device=target.device, - kvcache=target, - need_sort=False, - shard_spec=target.shard_spec, - ) - # DeepSeek's NextN uses local layer 0. Also cover an absolute/offset layer ID - # so the one-layer gather terminates correctly on either scratch-slot parity. - drafts = [ - _make_draft_pool(target, allocator, kind, start_layer=start) - for start in (0, 61) - ] - pools = [target, *drafts] - for other in drafts: - assert other._page_pos.data_ptr() != target._page_pos.data_ptr() - assert other.kv_gather_comm is not target.kv_gather_comm - for slot, target_slot in zip(other._slots, target._slots): - for name in slot.tensors: - assert ( - slot.tensors[name].data_ptr() - != target_slot.tensors[name].data_ptr() - ) - - locs = _logical_chain(target.device) - req_to_token = locs.to(torch.int32).unsqueeze(0) - req_indices = torch.tensor([0], dtype=torch.int64, device=locs.device) - # Each cached prefix is page aligned. The third batch writes only five - # tokens of its tail page; the fourth recomputes and completes that page. - batches = [(0, 48), (48, 96), (96, 133), (128, 144)] - for epoch, (prefix_len, seq_len) in enumerate(batches, start=1): - active_locs = locs[:seq_len] - chunk_locs = locs[prefix_len:seq_len] - target_epoch = target._epoch - for tag, pool in enumerate(pools, start=1): - pool.begin_shard_extend(req_to_token, req_indices, [prefix_len], [seq_len]) - assert pool._epoch == epoch - for local_layer in range(pool.layer_num): - _write(pool, kind, chunk_locs, tag, local_layer, epoch) - _check_scratch(pool, kind, active_locs, prefix_len, tag, local_layer) - _check_persisted(pool, kind, active_locs, tag, rank) - # MTP metadata/writes/reads must leave the target's last resident - # layer and gather plan usable, despite identical logical locations. - assert target._epoch == target_epoch + 1 - _check_scratch(target, kind, active_locs, prefix_len, 1, TARGET_LAYERS - 1) - if pool.layer_num == 1 and prefix_len: - resident = pool._slots[pool.start_layer % 2] - assert resident.resident_key == (pool.start_layer, epoch) - assert pool._slots[(pool.start_layer + 1) % 2].resident_key is None - if rank == 0: - print(f"PASS: {kind.upper()} target/MTP prefix={prefix_len}, seq={seq_len}") - - # Deactivating one draft's plan must not deactivate the other pools. - drafts[0].end_shard_extend() - assert not drafts[0]._shard_extend_active - assert target._shard_extend_active and drafts[1]._shard_extend_active - _check_scratch(drafts[1], kind, locs, 128, 3, 0) - _check_scratch(target, kind, locs, 128, 1, TARGET_LAYERS - 1) - torch.cuda.synchronize() - torch.distributed.barrier() - torch.distributed.destroy_process_group() - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--pool", choices=("all", "mla", "mha"), default="all") - parser.add_argument("--port", type=int, default=29821) - args = parser.parse_args() - if torch.cuda.device_count() < WORLD: - raise RuntimeError("MTP KV sharding correctness requires two visible GPUs") - kinds = ("mla", "mha") if args.pool == "all" else (args.pool,) - for offset, kind in enumerate(kinds): - mp.spawn(_run, args=(args.port + offset, kind), nprocs=WORLD, join=True) - print("PASS: MTP page-interleave KV cache correctness") - - -if __name__ == "__main__": - main() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 44a5d95b833b..b49afeb86fa3 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -882,6 +882,7 @@ def test_eight_long_prefixes_fit_an_eight_k_chunk(self): extend_len = 1024 prefix_len = context_len - extend_len kvc = SimpleNamespace( + is_draft_worker=False, page_size=4, model_config=SimpleNamespace(context_len=context_len), ) diff --git a/test/registered/unit/mem_cache/test_mtp_kv_shard_config.py b/test/registered/unit/mem_cache/test_mtp_kv_shard_config.py deleted file mode 100644 index 6700d607b0d3..000000000000 --- a/test/registered/unit/mem_cache/test_mtp_kv_shard_config.py +++ /dev/null @@ -1,365 +0,0 @@ -"""CPU coverage of MTP pool construction against a shared sharded allocator.""" - -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -import torch - -from sglang.srt.mem_cache import page_interleave -from sglang.srt.mem_cache.allocator.page_interleave import PageInterleavePoolAllocator -from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator -from sglang.srt.mem_cache.page_interleave import ( - PageShardSpec, - compute_page_shard_scratch_bytes, - get_kv_shard_group_info, - get_shared_kv_shard_pool, - make_page_shard_spec, -) -from sglang.srt.mem_cache.page_interleave_pool import ( - PageInterleaveMHATokenToKVPool, - PageInterleaveMLATokenToKVPool, -) -from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig -from sglang.srt.runtime_context import ( - get_context, - get_memory, - get_parallel, - get_schedule, - get_spec, -) -from sglang.srt.speculative.spec_info import SpeculativeAlgorithm -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - - -def _make_draft_configurator(*, use_mla=True): - # Keep the production pool/allocator types, but allocate only the CPU - # allocator metadata. Pool construction itself is mocked in these tests. - pool_class = ( - PageInterleaveMLATokenToKVPool if use_mla else PageInterleaveMHATokenToKVPool - ) - target_pool = pool_class.__new__(pool_class) - target_pool.shard_spec = PageShardSpec( - shard_rank=1, - shard_size=2, - page_size=16, - max_prefix_tokens=128, - chunk_tokens=64, - ) - target_pool.shard_group = SimpleNamespace(rank_in_group=1, world_size=2) - target_pool.size = 128 - target_pool.page_size = 16 - target_pool.dtype = torch.bfloat16 - target_pool.kv_lora_rank = 8 - target_pool.qk_rope_head_dim = 4 - target_pool.head_num = 2 - target_pool.head_dim = 4 - target_pool.v_head_dim = 6 - allocator = PageInterleavePoolAllocator( - size=target_pool.size, - physical_page_size=16, - shard_size=2, - dtype=torch.bfloat16, - device="cpu", - kvcache=target_pool, - need_sort=False, - shard_spec=target_pool.shard_spec, - ) - kvc = KVCacheConfigurator.__new__(KVCacheConfigurator) - kvc.device = "cpu" - kvc.gpu_id = 0 - kvc.is_draft_worker = True - kvc.page_size = 16 - kvc.use_mla_backend = use_mla - kvc.kv_cache_dtype = torch.bfloat16 - kvc.kv_cache_dtype_str = "bfloat16" - kvc.token_to_kv_pool_allocator = allocator - kvc.req_to_token_pool = None - kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE - kvc.model_config = SimpleNamespace( - num_nextn_predict_layers=1, - context_len=100, - kv_lora_rank=8, - qk_rope_head_dim=4, - head_dim=4, - v_head_dim=6, - get_num_kv_heads=lambda tp, dcp=1: 4 // tp, - hf_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]), - ) - kvc.layer_info = SimpleNamespace(start_layer=5, end_layer=6, num_effective_layers=1) - kvc.spec_aux_config = SimpleNamespace(eagle_draft_num_layers=1) - kvc.is_hybrid_swa = False - kvc.is_hybrid_swa_compress = False - kvc.sliding_window_size = None - kvc.mambaish_config = None - kvc.hybrid_gdn_config = None - kvc.hybrid_kda_config = None - kvc.post_capture_kv_active = False - kvc.memory_pool_config = MemoryPoolConfig( - max_total_num_tokens=128, max_running_requests=4 - ) - return kvc, target_pool - - -def _build_draft_pool(kvc, *, is_dsa_model=False, is_dsv4_model=False): - sizes = kvc._derive_pool_sizes(config=kvc.memory_pool_config) - return kvc._build_token_to_kv_pool( - sizes=sizes, - is_dsa_model=is_dsa_model, - is_dsv4_model=is_dsv4_model, - req_to_token_pool=None, - ) - - -class TestMTPKVShardConfig(CustomTestCase): - def setUp(self): - super().setUp() - self.enterContext( - get_context().override_server_args( - page_size=16, - chunked_prefill_size=64, - enable_multi_layer_eagle=False, - ) - ) - self.enterContext( - get_parallel().override( - tp_size=2, attn_tp_size=2, moe_tp_size=2, attn_dcp_size=1 - ) - ) - - def test_inherits_exact_target_spec_without_resolving_current_group(self): - for use_mla in (True, False): - with self.subTest(use_mla=use_mla): - kvc, target = _make_draft_configurator(use_mla=use_mla) - with patch.object( - page_interleave, "get_kv_shard_group", side_effect=AssertionError - ): - self.assertIs(get_shared_kv_shard_pool(kvc), target) - self.assertEqual(get_kv_shard_group_info(kvc), (1, 2)) - self.assertIs(make_page_shard_spec(kvc), target.shard_spec) - - def test_builds_separate_mtp_pool_with_target_layout_and_draft_layer_offsets(self): - for use_mla in (True, False): - with self.subTest(use_mla=use_mla): - kvc, target = _make_draft_configurator(use_mla=use_mla) - pool_class = type(target) - # Patch __init__, not the class: target validation still sees - # the real type, and construction returns a distinct instance. - with patch.object(pool_class, "__init__", return_value=None) as init: - pool = _build_draft_pool(kvc) - self.assertIsInstance(pool, pool_class) - self.assertIsNot(pool, target) - kwargs = init.call_args.kwargs - self.assertEqual(kwargs["size"], target.size) - self.assertIs(kwargs["shard_spec"], target.shard_spec) - self.assertIs(kwargs["shard_group"], target.shard_group) - self.assertEqual(kwargs["page_size"], 16) - self.assertEqual(kwargs["layer_num"], 1) - self.assertEqual(kwargs["start_layer"], 5) - self.assertEqual(kwargs["end_layer"], 6) - self.assertEqual(kwargs["dtype"], torch.bfloat16) - sizes = kvc._derive_pool_sizes(config=kvc.memory_pool_config) - self.assertIs( - kvc._build_token_to_kv_pool_allocator( - sizes=sizes, - token_to_kv_pool=pool, - is_dsv4_model=False, - req_to_token_pool=None, - token_to_kv_pool_allocator=kvc.token_to_kv_pool_allocator, - ), - kvc.token_to_kv_pool_allocator, - ) - - def test_draft_without_shared_sharded_allocator_keeps_existing_pool(self): - for allocator in (None, object()): - with self.subTest(allocator=allocator): - kvc, _ = _make_draft_configurator() - kvc.token_to_kv_pool_allocator = allocator - self.assertIsNone(get_shared_kv_shard_pool(kvc)) - self.assertEqual(get_kv_shard_group_info(kvc), (None, 1)) - self.assertIsNone(make_page_shard_spec(kvc)) - self.assertEqual(compute_page_shard_scratch_bytes(kvc), 0) - with patch.object(KVCacheConfigurator, "_build_mla_kv_pool") as build: - self.assertIs(_build_draft_pool(kvc), build.return_value) - build.assert_called_once_with(max_total_num_tokens=128) - - def test_target_without_shared_allocator_keeps_existing_pool(self): - kvc, _ = _make_draft_configurator() - kvc.is_draft_worker = False - kvc.token_to_kv_pool_allocator = None - # Ordinary target construction does not require the global sharding - # flag that has not yet been integrated into server arguments. - with patch.object(KVCacheConfigurator, "_build_mla_kv_pool") as build: - self.assertIs(_build_draft_pool(kvc), build.return_value) - build.assert_called_once_with(max_total_num_tokens=128) - - def test_rejects_invalid_shared_target_pool_and_spec(self): - kvc, _ = _make_draft_configurator() - kvc.token_to_kv_pool_allocator._kvcache = object() - with self.assertRaisesRegex(ValueError, "target's sharded KV pool"): - get_shared_kv_shard_pool(kvc) - mismatched_spec = PageShardSpec( - shard_rank=0, - shard_size=2, - page_size=16, - max_prefix_tokens=128, - chunk_tokens=64, - ) - for attribute, value in ( - ("shard_spec", mismatched_spec), - ("shard_size", 4), - ("page_size", 32), - ("size", 128), - ): - with self.subTest(attribute=attribute): - kvc, _ = _make_draft_configurator() - setattr(kvc.token_to_kv_pool_allocator, attribute, value) - with self.assertRaisesRegex(ValueError, "different shard layouts"): - get_shared_kv_shard_pool(kvc) - - def test_rejects_page_dtype_and_attention_mode_mismatches(self): - for attribute, value, message in ( - ("page_size", 32, "physical page size"), - ("kv_cache_dtype", torch.float16, "geometry and dtype"), - ("use_mla_backend", False, "attention geometry"), - ): - with self.subTest(attribute=attribute): - kvc, _ = _make_draft_configurator() - setattr(kvc, attribute, value) - with self.assertRaisesRegex(ValueError, message): - _build_draft_pool(kvc) - - def test_rejects_mla_and_mha_geometry_mismatches(self): - for use_mla, attribute, value in ( - (True, "kv_lora_rank", 16), - (True, "qk_rope_head_dim", 8), - (False, "head_dim", 8), - (False, "v_head_dim", 8), - (False, "get_num_kv_heads", lambda tp, dcp=1: 4), - ): - with self.subTest(use_mla=use_mla, attribute=attribute): - kvc, _ = _make_draft_configurator(use_mla=use_mla) - setattr(kvc.model_config, attribute, value) - with self.assertRaisesRegex(ValueError, "geometry and dtype"): - _build_draft_pool(kvc) - - def test_rejects_draft_capacity_different_from_target(self): - for capacity in (64, 256): - with self.subTest(capacity=capacity): - kvc, _ = _make_draft_configurator() - kvc.memory_pool_config = MemoryPoolConfig( - max_total_num_tokens=capacity, max_running_requests=4 - ) - with self.assertRaisesRegex(ValueError, "same per-rank capacity"): - _build_draft_pool(kvc) - - def test_rejects_other_draft_algorithms_and_multi_runner_mode(self): - for algorithm in ( - SpeculativeAlgorithm.EAGLE3, - SpeculativeAlgorithm.DFLASH, - SpeculativeAlgorithm.FROZEN_KV_MTP, - ): - with self.subTest(algorithm=algorithm): - kvc, _ = _make_draft_configurator() - kvc.spec_algorithm = algorithm - with self.assertRaisesRegex(ValueError, "EAGLE MTP"): - _build_draft_pool(kvc) - kvc, _ = _make_draft_configurator() - with get_spec().override(enable_multi_layer_eagle=True): - with self.assertRaisesRegex(ValueError, "EAGLE MTP"): - _build_draft_pool(kvc) - for layers in (None, 0): - with self.subTest(layers=layers): - kvc.model_config.num_nextn_predict_layers = layers - with self.assertRaisesRegex(ValueError, "EAGLE MTP"): - _build_draft_pool(kvc) - - def test_rejects_unsupported_model_pool_layouts(self): - for options in ({"is_dsa_model": True}, {"is_dsv4_model": True}): - with self.subTest(options=options): - kvc, _ = _make_draft_configurator() - with self.assertRaisesRegex(ValueError, "dense MLA and MHA"): - _build_draft_pool(kvc, **options) - for attribute, value in ( - ("is_hybrid_swa", True), - ("mambaish_config", object()), - ("sliding_window_size", 128), - ): - with self.subTest(attribute=attribute): - kvc, _ = _make_draft_configurator() - setattr(kvc, attribute, value) - with self.assertRaisesRegex(ValueError, "dense MLA and MHA"): - _build_draft_pool(kvc) - - def test_rejects_dcp(self): - kvc, _ = _make_draft_configurator() - with get_parallel().override(attn_dcp_size=2): - with self.assertRaisesRegex(ValueError, "incompatible with DCP"): - _build_draft_pool(kvc) - - def test_rejects_special_kv_storage_layouts(self): - for attribute, value in ( - ("post_capture_kv_active", True), - ("kv_cache_dtype_str", "mxfp8"), - ): - with self.subTest(attribute=attribute): - kvc, _ = _make_draft_configurator() - setattr(kvc, attribute, value) - with self.assertRaisesRegex(ValueError, "plain per-layer KV layout"): - _build_draft_pool(kvc) - kvc, _ = _make_draft_configurator() - with get_memory().override(enable_page_major_kv_layout=True): - with self.assertRaisesRegex(ValueError, "plain per-layer KV layout"): - _build_draft_pool(kvc) - - def test_draft_scratch_uses_inherited_bounds_and_has_one_buffer_pair(self): - for use_mla in (True, False): - with self.subTest(use_mla=use_mla): - kvc, _ = _make_draft_configurator(use_mla=use_mla) - kvc.model_config.context_len = 9999 - row_bytes = (8 + 4) * 2 if use_mla else 2 * (4 + 6) * 2 - expected = 2 * (128 + 64 + 16) * row_bytes - with get_schedule().override(chunked_prefill_size=1024): - self.assertEqual(compute_page_shard_scratch_bytes(kvc), expected) - with self.assertRaisesRegex(ValueError, "target"): - compute_page_shard_scratch_bytes(kvc, include_mtp=True) - - def test_target_plus_mtp_scratch_preserves_eight_context_budget(self): - kvc, target = _make_draft_configurator() - kvc.is_draft_worker = False - kvc.token_to_kv_pool_allocator = None - # Some architectures populate num_nextn_predict_layers only in their - # draft config; target sizing uses the resolved draft layer count. - kvc.model_config.num_nextn_predict_layers = None - kvc.layer_info = SimpleNamespace( - start_layer=0, end_layer=8, num_effective_layers=8 - ) - parallel = SimpleNamespace( - enable_kv_cache_sharding=True, - attn_cp_group=SimpleNamespace(world_size=1), - attn_tp_group=target.shard_group, - attn_tp_size=2, - ) - # 100 tokens rounds to 128 per context before multiplying by eight. - single_pool_scratch = 2 * (8 * 128 + 64 + 16) * (8 + 4) * 2 - with patch.object(page_interleave, "get_parallel", return_value=parallel): - for layers in (1, 3): - with self.subTest(layers=layers): - kvc.spec_aux_config.eagle_draft_num_layers = layers - self.assertEqual( - compute_page_shard_scratch_bytes(kvc), single_pool_scratch - ) - self.assertEqual( - compute_page_shard_scratch_bytes(kvc, include_mtp=True), - 2 * single_pool_scratch, - ) - kvc.spec_aux_config.eagle_draft_num_layers = None - with self.assertRaisesRegex(ValueError, "EAGLE MTP"): - compute_page_shard_scratch_bytes(kvc, include_mtp=True) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index ad5fea9b0212..32a48c974dd4 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -33,11 +33,15 @@ class draws (owners exactly cyclic along a chain), least-full root SimpleNamespace binding pattern of ``test_dsa_layer_shard_utils.py``. 6. The value-derived P/D ownership send filter ``filter_kv_indices_for_shard_rank``. +7. MTP draft pool construction, shared allocator/layout validation, and scratch + sizing for separate target and draft pools. The GPU tests exercise real pools and NCCL collectives: MLA and MHA gathers on 2 ranks (``TestPageInterleaveGatherMultiGpu``), and byte-exact DSA indexer keys, scales, and latent KV on 2 and 4 ranks (``TestPageInterleaveIndexer``), -including dequantized MLA reads from packed FP8 scratch. +including dequantized MLA reads from packed FP8 scratch. MTP cases cover exact +KV bytes, separate target/draft storage and scratch plans, fragmented and partial +pages, both MLA write APIs, and local/nonzero draft layer IDs across batches. GPU tests skip when too few CUDA devices are visible, so the same file runs in both CPU and CUDA CI suites. """ @@ -48,7 +52,7 @@ class draws (owners exactly cyclic along a chain), least-full root import unittest import unittest.mock from array import array -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from types import SimpleNamespace import numpy as np @@ -76,12 +80,15 @@ class draws (owners exactly cyclic along a chain), least-full root ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import _evict_until_allocatable +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool from sglang.srt.mem_cache.page_interleave import ( PageInterleavePlacement, PageShardSpec, compute_page_shard_scratch_bytes, get_kv_shard_group, + get_kv_shard_group_info, + get_shared_kv_shard_pool, make_page_shard_spec, ) from sglang.srt.mem_cache.page_interleave_pool import ( @@ -92,8 +99,15 @@ class draws (owners exactly cyclic along a chain), least-full root from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.components import ComponentType from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache -from sglang.srt.runtime_context import get_parallel, publish +from sglang.srt.runtime_context import ( + get_context, + get_memory, + get_parallel, + get_spec, + publish, +) from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import ceil_div from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci from sglang.test.mem_cache_utils import finish_req @@ -248,6 +262,275 @@ def test_estimate_matches_actual_mla_and_gqa_scratch_tensors(self): ) self.assertEqual(estimated, actual) + def test_target_with_mtp_budgets_one_extra_scratch_pair(self): + with ( + self._fixture(context_len=100, chunk_tokens=64, shard_size=2) as kvc, + get_context().override_server_args(enable_multi_layer_eagle=False), + ): + kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE + # Target configs can omit next-N layers; use the resolved draft count. + kvc.model_config.num_nextn_predict_layers = None + kvc.spec_aux_config = SimpleNamespace(eagle_draft_num_layers=1) + expected = 2 * (8 * 128 + 64 + 16) * (16 + 8) * 2 + for layers in (1, 3): + with self.subTest(layers=layers): + kvc.spec_aux_config.eagle_draft_num_layers = layers + self.assertEqual(compute_page_shard_scratch_bytes(kvc), expected) + self.assertEqual( + compute_page_shard_scratch_bytes(kvc, include_mtp=True), + 2 * expected, + ) + kvc.spec_aux_config.eagle_draft_num_layers = None + with self.assertRaisesRegex(ValueError, "EAGLE MTP"): + compute_page_shard_scratch_bytes(kvc, include_mtp=True) + + +def _make_mtp_configurator(use_mla=True): + # Keep real pool/allocator types for validation; allocate only CPU metadata. + pool_class = ( + PageInterleaveMLATokenToKVPool if use_mla else PageInterleaveMHATokenToKVPool + ) + target = pool_class.__new__(pool_class) + target.shard_spec = PageShardSpec(1, 2, 16, 128, 64) + target.shard_group = SimpleNamespace(rank_in_group=1, world_size=2) + target.size, target.page_size, target.dtype = 128, 16, torch.bfloat16 + target.kv_lora_rank, target.qk_rope_head_dim = 8, 4 + target.head_num, target.head_dim, target.v_head_dim = 2, 4, 6 + kvc = KVCacheConfigurator.__new__(KVCacheConfigurator) + kvc.token_to_kv_pool_allocator = PageInterleavePoolAllocator( + size=target.size, + physical_page_size=target.page_size, + shard_size=target.shard_spec.shard_size, + dtype=target.dtype, + device="cpu", + kvcache=target, + need_sort=False, + shard_spec=target.shard_spec, + ) + kvc.device = "cpu" + kvc.is_draft_worker = True + kvc.page_size = target.page_size + kvc.use_mla_backend = use_mla + kvc.kv_cache_dtype = target.dtype + kvc.kv_cache_dtype_str = "bfloat16" + kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE + kvc.model_config = SimpleNamespace( + num_nextn_predict_layers=1, + context_len=9999, + kv_lora_rank=8, + qk_rope_head_dim=4, + head_dim=4, + v_head_dim=6, + get_num_kv_heads=lambda tp: 4 // tp, + hf_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]), + ) + kvc.layer_info = SimpleNamespace(start_layer=5, end_layer=6, num_effective_layers=1) + kvc.is_hybrid_swa = False + kvc.sliding_window_size = None + kvc.mambaish_config = None + kvc.post_capture_kv_active = False + return kvc, target + + +def _build_mtp_pool(kvc, *, capacity=128, is_dsa_model=False, is_dsv4_model=False): + return kvc._build_token_to_kv_pool( + sizes=SimpleNamespace(max_total_num_tokens=capacity), + is_dsa_model=is_dsa_model, + is_dsv4_model=is_dsv4_model, + req_to_token_pool=None, + ) + + +class TestMTPKVShardConfig(CustomTestCase): + def setUp(self): + super().setUp() + stack = ExitStack() + self.addCleanup(stack.close) + stack.enter_context( + get_context().override_server_args( + page_size=16, chunked_prefill_size=1024, enable_multi_layer_eagle=False + ) + ) + stack.enter_context( + get_parallel().override( + tp_size=2, attn_tp_size=2, moe_tp_size=2, attn_dcp_size=1 + ) + ) + + def test_draft_shares_target_layout_but_has_its_own_pool(self): + for use_mla in (True, False): + with self.subTest(use_mla=use_mla): + kvc, target = _make_mtp_configurator(use_mla) + pool_class = type(target) + # Mock allocation only; preserve real types and distinct instances. + with ( + unittest.mock.patch.object( + pool_class, "__init__", return_value=None + ) as init, + unittest.mock.patch.object( + page_interleave, + "get_kv_shard_group", + side_effect=AssertionError, + ), + ): + self.assertIs(get_shared_kv_shard_pool(kvc), target) + self.assertEqual(get_kv_shard_group_info(kvc), (1, 2)) + self.assertIs(make_page_shard_spec(kvc), target.shard_spec) + pool = _build_mtp_pool(kvc) + self.assertIsInstance(pool, pool_class) + self.assertIsNot(pool, target) + kwargs = init.call_args.kwargs + self.assertIs(kwargs["shard_spec"], target.shard_spec) + self.assertIs(kwargs["shard_group"], target.shard_group) + for name, expected in dict( + size=128, + page_size=16, + dtype=torch.bfloat16, + layer_num=1, + start_layer=5, + end_layer=6, + ).items(): + self.assertEqual(kwargs[name], expected, name) + allocator = kvc.token_to_kv_pool_allocator + self.assertIs( + kvc._build_token_to_kv_pool_allocator( + sizes=SimpleNamespace(max_total_num_tokens=128), + token_to_kv_pool=pool, + is_dsv4_model=False, + req_to_token_pool=None, + token_to_kv_pool_allocator=allocator, + ), + allocator, + ) + # Scratch uses inherited bounds despite different draft context/chunk sizes. + row_bytes = (8 + 4) * 2 if use_mla else 2 * (4 + 6) * 2 + self.assertEqual( + compute_page_shard_scratch_bytes(kvc), + 2 * (128 + 64 + 16) * row_bytes, + ) + with self.assertRaisesRegex(ValueError, "target"): + compute_page_shard_scratch_bytes(kvc, include_mtp=True) + + def test_without_shared_allocator_keeps_existing_pool(self): + for draft, allocator in ((True, None), (True, object()), (False, None)): + with self.subTest(draft=draft, allocator=allocator): + kvc, _ = _make_mtp_configurator() + kvc.is_draft_worker = draft + kvc.token_to_kv_pool_allocator = allocator + self.assertIsNone(get_shared_kv_shard_pool(kvc)) + if draft: + self.assertEqual(get_kv_shard_group_info(kvc), (None, 1)) + self.assertIsNone(make_page_shard_spec(kvc)) + self.assertEqual(compute_page_shard_scratch_bytes(kvc), 0) + with unittest.mock.patch.object( + KVCacheConfigurator, "_build_mla_kv_pool" + ) as build: + self.assertIs(_build_mtp_pool(kvc), build.return_value) + build.assert_called_once_with(max_total_num_tokens=128) + + def test_rejects_invalid_shared_allocator(self): + kvc, _ = _make_mtp_configurator() + for name, value, message in ( + ("_kvcache", object(), "target's sharded KV pool"), + ("shard_spec", PageShardSpec(0, 2, 16, 128, 64), "different shard layouts"), + ("shard_size", 4, "different shard layouts"), + ("page_size", 32, "different shard layouts"), + ("size", 128, "different shard layouts"), + ): + with ( + self.subTest(attribute=name), + unittest.mock.patch.object(kvc.token_to_kv_pool_allocator, name, value), + self.assertRaisesRegex(ValueError, message), + ): + get_shared_kv_shard_pool(kvc) + + def test_rejects_incompatible_pool_geometry(self): + for use_mla in (True, False): + kvc, _ = _make_mtp_configurator(use_mla) + geometry = ( + ("kv_lora_rank", "qk_rope_head_dim") + if use_mla + else ("head_dim", "v_head_dim") + ) + cases = [ + (kvc, "page_size", 32, "physical page size"), + (kvc, "kv_cache_dtype", torch.float16, "geometry and dtype"), + (kvc, "use_mla_backend", not use_mla, "attention geometry"), + *( + (kvc.model_config, field, 16, "geometry and dtype") + for field in geometry + ), + ] + if not use_mla: + cases.append( + ( + kvc.model_config, + "get_num_kv_heads", + lambda tp: 4, + "geometry and dtype", + ) + ) + for obj, name, value, message in cases: + with ( + self.subTest(use_mla=use_mla, attribute=name), + unittest.mock.patch.object(obj, name, value), + self.assertRaisesRegex(ValueError, message), + ): + _build_mtp_pool(kvc) + for capacity in (64, 256): + with ( + self.subTest(use_mla=use_mla, capacity=capacity), + self.assertRaisesRegex(ValueError, "same per-rank capacity"), + ): + _build_mtp_pool(kvc, capacity=capacity) + + def test_rejects_unsupported_draft_config(self): + kvc, _ = _make_mtp_configurator() + cases = [ + (kvc, "spec_algorithm", algorithm, "EAGLE MTP") + for algorithm in ( + SpeculativeAlgorithm.EAGLE3, + SpeculativeAlgorithm.DFLASH, + SpeculativeAlgorithm.FROZEN_KV_MTP, + ) + ] + [ + (kvc.model_config, "num_nextn_predict_layers", None, "EAGLE MTP"), + (kvc.model_config, "num_nextn_predict_layers", 0, "EAGLE MTP"), + (kvc, "is_hybrid_swa", True, "dense MLA and MHA"), + (kvc, "mambaish_config", object(), "dense MLA and MHA"), + (kvc, "sliding_window_size", 128, "dense MLA and MHA"), + (kvc, "post_capture_kv_active", True, "plain per-layer KV layout"), + (kvc, "kv_cache_dtype_str", "mxfp8", "plain per-layer KV layout"), + ] + for obj, name, value, message in cases: + with ( + self.subTest(attribute=name, value=value), + unittest.mock.patch.object(obj, name, value), + self.assertRaisesRegex(ValueError, message), + ): + _build_mtp_pool(kvc) + for options in ({"is_dsa_model": True}, {"is_dsv4_model": True}): + with ( + self.subTest(options=options), + self.assertRaisesRegex(ValueError, "dense MLA and MHA"), + ): + _build_mtp_pool(kvc, **options) + for context, options, message in ( + (get_spec(), {"enable_multi_layer_eagle": True}, "EAGLE MTP"), + (get_parallel(), {"attn_dcp_size": 2}, "incompatible with DCP"), + ( + get_memory(), + {"enable_page_major_kv_layout": True}, + "plain per-layer KV layout", + ), + ): + with ( + self.subTest(options=options), + context.override(**options), + self.assertRaisesRegex(ValueError, message), + ): + _build_mtp_pool(kvc) + def _make_spec(shard_rank=0, max_prefix_groups=64, chunk_pages=32): return PageShardSpec( @@ -1434,6 +1717,8 @@ def test_rank_owning_nothing_in_chunk(self): _GATHER_HEAD_NUM = 2 _GATHER_HEAD_DIM = 32 _GATHER_DTYPE = torch.bfloat16 +_MTP_SIZE = 32 * _GATHER_PAGE_SIZE # Physical token capacity per rank. +_MTP_UNWRITTEN = -256 def _mla_value(loc, dim): @@ -1726,6 +2011,294 @@ def kv_value(locs, layer_id, is_v): print("PASS: MHA page-interleave shard (attn-CP axis)") +def _mtp_make_pool(rank, group, kind, start_layer, layer_num): + kwargs = dict( + size=_MTP_SIZE, + page_size=_GATHER_PAGE_SIZE, + dtype=_GATHER_DTYPE, + layer_num=layer_num, + device=f"cuda:{rank}", + enable_memory_saver=False, + start_layer=start_layer, + end_layer=start_layer + layer_num - 1, + shard_spec=PageShardSpec( + shard_rank=group.rank_in_group, + shard_size=_GATHER_WORLD, + page_size=_GATHER_PAGE_SIZE, + max_prefix_tokens=16 * _GATHER_WORLD * _GATHER_PAGE_SIZE, + chunk_tokens=4 * _GATHER_PAGE_SIZE, + ), + shard_group=group, + ) + if kind == "mla": + pool = PageInterleaveMLATokenToKVPool( + **kwargs, + kv_lora_rank=_GATHER_KV_LORA_RANK, + qk_rope_head_dim=_GATHER_QK_ROPE, + ) + else: + pool = PageInterleaveMHATokenToKVPool( + **kwargs, + head_num=_GATHER_HEAD_NUM, + head_dim=_GATHER_HEAD_DIM, + enable_alt_stream=False, + ) + for local_layer in range(layer_num): + for buffer in _mtp_buffers(pool, kind, local_layer): + buffer.fill_(_MTP_UNWRITTEN) + return pool + + +def _mtp_buffers(pool, kind, local_layer): + if kind == "mla": + return (pool.kv_buffer[local_layer],) + return (pool.k_buffer[local_layer], pool.v_buffer[local_layer]) + + +def _mtp_make_draft_pool(target, allocator, kind, start_layer): + # Skip model loading while using the production pool dispatch and sharing + # the target allocator. A replicated draft pool must fail this test. + kvc = KVCacheConfigurator.__new__(KVCacheConfigurator) + kvc.device = target.device + kvc.is_draft_worker = True + kvc.page_size = _GATHER_PAGE_SIZE + kvc.use_mla_backend = kind == "mla" + kvc.kv_cache_dtype = _GATHER_DTYPE + kvc.kv_cache_dtype_str = "bfloat16" + kvc.post_capture_kv_active = False + kvc.is_hybrid_swa = False + kvc.sliding_window_size = None + kvc.mambaish_config = None + kvc.token_to_kv_pool_allocator = allocator + kvc.spec_algorithm = SpeculativeAlgorithm.EAGLE + kvc.model_config = SimpleNamespace( + num_nextn_predict_layers=1, + hf_config=SimpleNamespace(architectures=["DeepseekV3ForCausalLM"]), + kv_lora_rank=_GATHER_KV_LORA_RANK, + qk_rope_head_dim=_GATHER_QK_ROPE, + head_dim=_GATHER_HEAD_DIM, + v_head_dim=_GATHER_HEAD_DIM, + get_num_kv_heads=lambda *_args: _GATHER_HEAD_NUM, + ) + kvc.layer_info = SimpleNamespace( + start_layer=start_layer, end_layer=start_layer + 1, num_effective_layers=1 + ) + with get_context().override_server_args( + page_size=_GATHER_PAGE_SIZE, + tp_size=_GATHER_WORLD, + attn_cp_size=_GATHER_WORLD if kind == "mha" else 1, + speculative_algorithm="EAGLE", + enable_multi_layer_eagle=False, + ): + pool = _build_mtp_pool(kvc, capacity=target.size) + assert pool.shard_spec is target.shard_spec + assert pool.shard_group is target.shard_group + for buffer in _mtp_buffers(pool, kind, 0): + buffer.fill_(_MTP_UNWRITTEN) + return pool + + +def _mtp_values(locs, width, tag, local_layer, is_v=False): + """Each location, feature, layer, and pool has distinguishable BF16 bytes. + + Small integers survive BF16 exactly; high-valued logical slot IDs plus tiny + feature offsets would round away the differences this test must detect. + The first two columns encode the full logical slot, avoiding hash aliases. + """ + features = torch.arange(width, device=locs.device) + values = ( + locs[:, None] * 37 + + features[None, :] * 17 + + tag * 73 + + local_layer * 11 + + int(is_v) * 97 + ) % 251 - 125 + values[:, 0] = locs % 128 + values[:, 1] = locs // 128 + values[:, 2] = tag + values[:, 3] = local_layer + values[:, 4] = int(is_v) + return values.to(_GATHER_DTYPE) + + +def _mtp_canonical(kind, locs, tag, local_layer): + if kind == "mla": + return ( + _mtp_values( + locs, _GATHER_KV_LORA_RANK + _GATHER_QK_ROPE, tag, local_layer + ).unsqueeze(1), + ) + return tuple( + _mtp_values( + locs, _GATHER_HEAD_NUM * _GATHER_HEAD_DIM, tag, local_layer, is_v + ).reshape(-1, _GATHER_HEAD_NUM, _GATHER_HEAD_DIM) + for is_v in (False, True) + ) + + +def _mtp_assert_bytes(got, expected, label): + assert got.shape == expected.shape, label + assert got.dtype == expected.dtype == _GATHER_DTYPE, label + assert torch.equal( + got.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8) + ), f"{label}: KV bytes differ" + + +def _mtp_logical_chain(device): + # Owners alternate from rank 1, while local pages run out of order. Some + # logical slots exceed the physical pool size, catching a replicated draft + # indexing them raw. + local_pages = torch.tensor([25, 3, 20, 6, 28, 4, 30, 8, 22], device=device) + owners = (1 + torch.arange(local_pages.numel(), device=device)) % _GATHER_WORLD + pages = local_pages * _GATHER_WORLD + owners + return ( + pages[:, None] * _GATHER_PAGE_SIZE + + torch.arange(_GATHER_PAGE_SIZE, device=device) + ).reshape(-1) + + +def _mtp_write(pool, kind, locs, tag, local_layer, epoch): + layer = SimpleNamespace(layer_id=pool.start_layer + local_layer) + values = _mtp_canonical(kind, locs, tag, local_layer) + if kind == "mha": + pool.set_kv_buffer(layer, locs, *values) + elif epoch % 2: + # Both MLA write entry points must owner-filter and stage MTP rows. + pool.set_mla_kv_buffer( + layer, + locs, + values[0][..., :_GATHER_KV_LORA_RANK], + values[0][..., _GATHER_KV_LORA_RANK:], + ) + else: + pool.set_kv_buffer( + layer, locs, values[0], values[0][..., :_GATHER_KV_LORA_RANK] + ) + + +def _mtp_check_persisted(pool, kind, locs, tag, rank): + owned_locs = locs[(locs // _GATHER_PAGE_SIZE) % _GATHER_WORLD == rank] + local_rows = ( + owned_locs // (_GATHER_WORLD * _GATHER_PAGE_SIZE) * _GATHER_PAGE_SIZE + + owned_locs % _GATHER_PAGE_SIZE + ) + for local_layer in range(pool.layer_num): + canonical = _mtp_canonical(kind, owned_locs, tag, local_layer) + for buffer, values in zip(_mtp_buffers(pool, kind, local_layer), canonical): + expected = torch.full_like(buffer, _MTP_UNWRITTEN) + expected[local_rows] = values + _mtp_assert_bytes( + buffer, expected, f"{kind} pool {tag} owned layer {local_layer}" + ) + + +def _mtp_check_scratch(pool, kind, locs, prefix_len, tag, local_layer): + layer_id = pool.start_layer + local_layer + rows = pool.translate_loc_to_scratch(locs) + expected = _mtp_canonical(kind, locs, tag, local_layer) + _mtp_assert_bytes( + pool.get_key_buffer(layer_id)[rows], expected[0], f"{kind} pool {tag} key" + ) + _mtp_assert_bytes( + pool.get_value_buffer(layer_id)[rows], + expected[0][..., :_GATHER_KV_LORA_RANK] if kind == "mla" else expected[1], + f"{kind} pool {tag} value", + ) + if kind == "mla" and prefix_len: + # Arbitrary prefix indices also exercise the chunked MLA consumer. + subset = locs[3:prefix_len:7].clone() + k_nope, k_rope = pool.get_mla_kv_buffer( + SimpleNamespace(layer_id=layer_id), subset, _GATHER_DTYPE + ) + reference = _mtp_canonical(kind, subset, tag, local_layer)[0] + _mtp_assert_bytes( + k_nope, reference[..., :_GATHER_KV_LORA_RANK], f"pool {tag} latent" + ) + _mtp_assert_bytes( + k_rope, reference[..., _GATHER_KV_LORA_RANK:], f"pool {tag} rope" + ) + + +def _run_mtp(rank, port, kind): + """Check exact KV bytes and independent target/draft state across batches.""" + _dist_init( + rank, _GATHER_WORLD, port, attn_cp_size=_GATHER_WORLD if kind == "mha" else 1 + ) + group = get_kv_shard_group(use_mla_backend=kind == "mla") + assert group.world_size == _GATHER_WORLD + target = _mtp_make_pool( + rank, group, kind, start_layer=0, layer_num=_GATHER_LAYER_NUM + ) + allocator = PageInterleavePoolAllocator( + size=target.size, + physical_page_size=_GATHER_PAGE_SIZE, + shard_size=_GATHER_WORLD, + dtype=_GATHER_DTYPE, + device=target.device, + kvcache=target, + need_sort=False, + shard_spec=target.shard_spec, + ) + # DeepSeek's NextN uses local layer 0. Also cover an absolute/offset layer ID + # so the one-layer gather terminates correctly on either scratch-slot parity. + drafts = [ + _mtp_make_draft_pool(target, allocator, kind, start_layer=start) + for start in (0, 61) + ] + pools = [target, *drafts] + for other in drafts: + assert other._page_pos.data_ptr() != target._page_pos.data_ptr() + assert other.kv_gather_comm is not target.kv_gather_comm + for slot, target_slot in zip(other._slots, target._slots): + for name in slot.tensors: + assert ( + slot.tensors[name].data_ptr() + != target_slot.tensors[name].data_ptr() + ) + + locs = _mtp_logical_chain(target.device) + req_to_token = locs.to(torch.int32).unsqueeze(0) + req_indices = torch.tensor([0], dtype=torch.int64, device=locs.device) + # Each cached prefix is page aligned. The third batch writes only five + # tokens of its tail page; the fourth recomputes and completes that page. + batches = [(0, 48), (48, 96), (96, 133), (128, 144)] + for epoch, (prefix_len, seq_len) in enumerate(batches, start=1): + active_locs = locs[:seq_len] + chunk_locs = locs[prefix_len:seq_len] + target_epoch = target._epoch + for tag, pool in enumerate(pools, start=1): + pool.begin_shard_extend(req_to_token, req_indices, [prefix_len], [seq_len]) + assert pool._epoch == epoch + for local_layer in range(pool.layer_num): + _mtp_write(pool, kind, chunk_locs, tag, local_layer, epoch) + _mtp_check_scratch( + pool, kind, active_locs, prefix_len, tag, local_layer + ) + _mtp_check_persisted(pool, kind, active_locs, tag, rank) + # MTP metadata/writes/reads must leave the target's last resident + # layer and gather plan usable, despite identical logical locations. + assert target._epoch == target_epoch + 1 + _mtp_check_scratch( + target, kind, active_locs, prefix_len, 1, _GATHER_LAYER_NUM - 1 + ) + if pool.layer_num == 1 and prefix_len: + resident = pool._slots[pool.start_layer % 2] + assert resident.resident_key == (pool.start_layer, epoch) + assert pool._slots[(pool.start_layer + 1) % 2].resident_key is None + if rank == 0: + print(f"PASS: {kind.upper()} target/MTP prefix={prefix_len}, seq={seq_len}") + + # Deactivating one draft's plan must not deactivate the other pools. + drafts[0].end_shard_extend() + assert not drafts[0]._shard_extend_active + assert target._shard_extend_active and drafts[1]._shard_extend_active + _mtp_check_scratch(drafts[1], kind, locs, 128, 3, 0) + _mtp_check_scratch(target, kind, locs, 128, 1, _GATHER_LAYER_NUM - 1) + torch.cuda.synchronize() + torch.distributed.barrier() + torch.distributed.destroy_process_group() + + @unittest.skipIf( torch.cuda.device_count() < 2, "page-interleave gather needs 2 CUDA devices" ) @@ -1742,6 +2315,12 @@ def test_mla_shard_over_attention_tp(self): def test_mha_shard_over_attention_cp(self): mp.spawn(_run_mha, args=(_GATHER_WORLD, 29812), nprocs=_GATHER_WORLD, join=True) + def test_mtp_mla_shard_over_attention_tp(self): + mp.spawn(_run_mtp, args=(29813, "mla"), nprocs=_GATHER_WORLD, join=True) + + def test_mtp_mha_shard_over_attention_cp(self): + mp.spawn(_run_mtp, args=(29814, "mha"), nprocs=_GATHER_WORLD, join=True) + # ============================================================================= # Indexer: packed FP8 keys and FP32 scales over real NCCL (2 and 4 GPUs). From a22db8cd580112174f70b69663c3a4e38961c4e1 Mon Sep 17 00:00:00 2001 From: "daowu.hzy" Date: Wed, 7 Oct 2026 21:09:10 +0800 Subject: [PATCH 11/11] feat(disagg): transfer page-sharded DSA target and draft indexers Extend the existing sharded Mooncake protocol to route packed DSA indexer pages by owner. Preserve source/destination positions, reject mismatched page counts, and exclude empty indexer layers from RDMA descriptors. Count each sender only for its owned state pages. Add CPU byte-exact descriptor tests for CP4/CP8, target plus draft buffers, empty owners, prefix-hit final chunks, layout rejection and unsharded compatibility. --- .../sglang/srt/disaggregation/common/conn.py | 5 +- .../srt/disaggregation/mooncake/conn.py | 49 ++- .../test_dsa_sharded_pd_transfer.py | 296 ++++++++++++++++++ 3 files changed, 347 insertions(+), 3 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_dsa_sharded_pd_transfer.py diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 67243a5bdacb..ba15f8deaa6e 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -489,12 +489,15 @@ def _should_skip_cp_replicated_state_transfer(self) -> bool: Prefill CP materializes global token order before writing state pools, so every CP rank holds the same state. When all CP ranks transfer their KV shards, only rank 0 needs to send that state. Cache layer split is the - exception because each CP rank owns different state layers. + exception because each CP rank owns different state layers. Page-sharded + DSA indexer state likewise needs every owner; Mooncake filters its page + lists before transfer instead of suppressing whole ranks. """ return ( self.attn_cp_size > 1 and self.attn_cp_rank != 0 and not get_parallel().enable_dsa_cache_layer_split + and getattr(self, "kv_shard_size", 1) <= 1 ) def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 664edb7ec4f6..0ae2ab54f350 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -54,6 +54,7 @@ build_dsa_tail_transfer_blocks, build_transfer_entry_pairs, compute_mamba_state_slice_byte_blocks, + filter_kv_indices_for_shard_rank, resolve_dcp_dst_entry_indices, should_send_replicated_state, slice_dsa_tail_dst_ptrs_for_pp, @@ -1082,6 +1083,13 @@ def _send_kvcache_generic( for layer_id in range(layers_current_pp_stage) ] assert layers_params is not None + if state_type == StateType.DSA and getattr(self, "kv_shard_size", 1) > 1: + # Skip-topk layers retain zero-size indexer entries to keep target + # and draft layer numbering aligned. Layout checking above still + # compares both peers; only omit their empty RDMA descriptors. + layers_params = [entry for entry in layers_params if entry[2] > 0] + if not layers_params: + return 0 def set_transfer_blocks( src_ptr: int, dst_ptr: int, item_len: int @@ -1784,8 +1792,8 @@ def _get_dsa_cache_transfer_skip_flags( # Prefill CP all-gathers before writing the pool, so every CP rank holds # the full state regardless of whether the pool is hybrid. We assume no # structure about the state rows, so we don't split them across CP ranks - # -- just let rank 0 send the whole thing (unless layer split already - # shards it per rank). + # -- just let rank 0 send the whole thing. Layer-split and page-sharded + # DSA pools are exceptions: every owner must send its distinct state. if self._should_skip_cp_replicated_state_transfer(): skip_state = True @@ -2015,6 +2023,26 @@ def maybe_send_extra( continue src_indices = list(indices) dst_indices_local = list(dst_indices) + if st == StateType.DSA and getattr(self, "kv_shard_size", 1) > 1: + # Target and draft indexer buffers share the latent KV page + # allocator. Pair logical source pages with destinations + # BEFORE converting owned sources to rank-local page IDs. + # Truncating a mismatched list would silently shift tokens. + if len(src_indices) != len(dst_indices_local): + raise RuntimeError( + "Sharded DSA state page count mismatch: " + f"prefill={len(src_indices)}, dst={len(dst_indices_local)}" + ) + src_indices, positions = filter_kv_indices_for_shard_rank( + self, + np.asarray(src_indices, dtype=np.int32), + slice(0, len(src_indices)), + ) + dst_indices_local = np.asarray(dst_indices_local, dtype=np.int32)[ + positions + ] + if len(src_indices) == 0: + continue if ( st == StateType.DSV4_REQUEST_STATE and len(src_indices) == 0 @@ -3070,6 +3098,23 @@ def send( if self.kv_mgr._should_skip_cp_replicated_state_transfer() else state_indices ) + if metric_state_indices and self.kv_mgr.kv_shard_size > 1: + # The queued state remains logical and is filtered by the worker. + # Account only for this sender's pages, including all target/draft + # indexer buffers represented by the DSA component's item lengths. + state_types = getattr( + getattr(self.kv_mgr, "kv_args", None), "state_types", [] + ) + metric_state_indices = list(metric_state_indices) + for i, st in enumerate(state_types): + if st == StateType.DSA and i < len(metric_state_indices): + indices = metric_state_indices[i] + if indices is not None: + metric_state_indices[i], _ = filter_kv_indices_for_shard_rank( + self.kv_mgr, + np.asarray(indices, dtype=np.int32), + slice(0, len(indices)), + ) self._record_transfer_indices(kv_indices, metric_state_indices) def poll(self) -> KVPoll: diff --git a/test/registered/unit/disaggregation/test_dsa_sharded_pd_transfer.py b/test/registered/unit/disaggregation/test_dsa_sharded_pd_transfer.py new file mode 100644 index 000000000000..923a991e5c22 --- /dev/null +++ b/test/registered/unit/disaggregation/test_dsa_sharded_pd_transfer.py @@ -0,0 +1,296 @@ +"""CPU byte-exact tests of page-sharded target/draft DSA PD transfer. + +Exercise the production owner filter and Mooncake descriptor builder with a +memory-copy transport. No GPU, RDMA engine or model weights are required. +""" + +import ctypes +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np + +from sglang.srt.disaggregation.base.conn import KVPoll, KVTransferMetric, StateType +from sglang.srt.disaggregation.common.conn import CommonKVManager +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager, MooncakeKVSender +from sglang.srt.disaggregation.utils import setup_state_kv_args +from sglang.srt.mem_cache.page_interleave_pool import PageInterleaveDSATokenToKVPool +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + +_ROOM = 7 +_INDEX_PAGE_BYTES = 64 * (128 + 4) + + +def _manager(rank, size, sources, destinations): + mgr = MooncakeKVManager.__new__(MooncakeKVManager) + mgr.attn_cp_rank, mgr.attn_cp_size = rank, size + mgr.kv_shard_rank, mgr.kv_shard_size = rank, size + mgr.attn_tp_size = mgr.pp_size = 1 + mgr.is_mla_backend = True + mgr.is_hybrid_mla_backend = False + mgr.enable_custom_mem_pool = False + mgr.max_transfer_batch_indices = 0 + mgr.state_strides_validated = set() + mgr.request_status = {_ROOM: KVPoll.Transferring} + mgr.kv_args = SimpleNamespace( + state_types=[StateType.DSA], + state_data_ptrs=[[x.ctypes.data for x in sources]], + state_item_lens=[[_INDEX_PAGE_BYTES if len(x) else 0 for x in sources]], + state_dim_per_tensor=[[]], + state_layer_ids=[[]], + prefill_start_layer=0, + ) + mgr.blocks = [] + + def copy_blocks(session, blocks): + mgr.blocks.extend(blocks) + for src, dst, length in blocks: + ctypes.memmove(dst, src, length) + return 0 + + mgr._transfer_data = copy_blocks + peer = SimpleNamespace( + dst_state_data_ptrs=[[x.ctypes.data for x in destinations]], + dst_state_item_lens=[ + [_INDEX_PAGE_BYTES if len(x) else 0 for x in destinations] + ], + dst_state_dim_per_tensor=[[]], + dst_state_layer_ids=[[]], + dst_attn_tp_size=1, + ) + return mgr, peer + + +def _send_state(mgr, peer, src_pages, dst_pages): + with patch( + "sglang.srt.disaggregation.mooncake.conn.get_memory", + return_value=SimpleNamespace(enable_unified_memory=False), + ): + return mgr.maybe_send_extra( + SimpleNamespace( + room=_ROOM, + mooncake_session_id="cpu", + dst_state_indices=[dst_pages], + ), + [src_pages], + executor=None, + target_rank_registration_info=peer, + ) + + +class TestDSAShardedPDTransfer(unittest.TestCase): + def test_target_and_draft_indexers_register_as_one_page_component(self): + pools = [] + for pointers in ([1000, 2000, 3000], [4000]): + pool = PageInterleaveDSATokenToKVPool.__new__( + PageInterleaveDSATokenToKVPool + ) + pool.get_state_buf_infos = lambda pointers=pointers: ( + pointers, + [8 * _INDEX_PAGE_BYTES] * len(pointers), + [_INDEX_PAGE_BYTES] * len(pointers), + ) + pool.get_compress_tail_buf_infos = lambda: ([], [], []) + pools.append(pool) + args = SimpleNamespace() + setup_state_kv_args(args, pools[0], draft_token_to_kv_pool=pools[1]) + self.assertEqual(args.state_types, [StateType.DSA]) + self.assertEqual(args.state_data_ptrs, [[1000, 2000, 3000, 4000]]) + self.assertEqual(args.state_item_lens, [[_INDEX_PAGE_BYTES] * 4]) + + def test_fragmented_pages_preserve_every_target_and_draft_byte(self): + # A full indexer page includes its packed K region AND scale region. + # Random bytes encode layer, logical page and byte offset independently. + # The fourth layer represents draft, preventing target/draft aliasing + # from passing as an apparently correct copy. + for size in (4, 8): + with self.subTest(cp_size=size): + logical_pages = np.array( + [15, 3, 24, 9, 18, 7, 28, 12, 22, 1, 8, 30, 4], + dtype=np.int32, + ) + dst_pages = np.array( + [29, 5, 16, 8, 1, 27, 3, 20, 7, 15, 10, 31, 0], + dtype=np.int32, + ) + rng = np.random.default_rng(92) + canonical = rng.integers( + 0, 256, (4, 32, _INDEX_PAGE_BYTES), dtype=np.uint8 + ) + destinations = [np.full_like(x, 0xA5) for x in canonical] + expected = [x.copy() for x in destinations] + for layer in range(4): + expected[layer][dst_pages] = canonical[layer][logical_pages] + byte_count = 0 + for rank in range(size): + sources = [np.ascontiguousarray(x[rank::size]) for x in canonical] + mgr, peer = _manager(rank, size, sources, destinations) + with patch( + "sglang.srt.disaggregation.common.conn.get_parallel", + return_value=SimpleNamespace( + enable_dsa_cache_layer_split=False + ), + ): + self.assertEqual( + mgr._get_dsa_cache_transfer_skip_flags(peer), (False, False) + ) + self.assertEqual( + _send_state(mgr, peer, logical_pages, dst_pages), 0 + ) + byte_count += sum(length for _, _, length in mgr.blocks) + for actual, oracle in zip(destinations, expected): + np.testing.assert_array_equal(actual, oracle) + self.assertEqual(byte_count, len(logical_pages) * 4 * _INDEX_PAGE_BYTES) + + def test_mismatch_is_rejected_before_any_write(self): + src = [np.zeros((8, _INDEX_PAGE_BYTES), dtype=np.uint8)] + dst = [np.full((32, _INDEX_PAGE_BYTES), 0xA5, dtype=np.uint8)] + for src_pages, dst_pages in (([4, 8], [0]), ([4], [0, 1]), ([], [0])): + with self.subTest(src=src_pages, dst=dst_pages): + mgr, peer = _manager(0, 4, src, dst) + with self.assertRaisesRegex(RuntimeError, "state page count mismatch"): + _send_state(mgr, peer, src_pages, dst_pages) + self.assertEqual(mgr.blocks, []) + self.assertTrue(np.all(dst[0] == 0xA5)) + + def test_rank_without_owned_pages_writes_nothing(self): + src = [np.zeros((8, _INDEX_PAGE_BYTES), dtype=np.uint8)] + dst = [np.full((32, _INDEX_PAGE_BYTES), 0xA5, dtype=np.uint8)] + mgr, peer = _manager(3, 4, src, dst) + self.assertEqual(_send_state(mgr, peer, [4, 8], [7, 2]), 0) + self.assertEqual(mgr.blocks, []) + self.assertTrue(np.all(dst[0] == 0xA5)) + + def test_packed_page_stride_mismatch_is_rejected_before_any_write(self): + src = [np.zeros((8, _INDEX_PAGE_BYTES), dtype=np.uint8)] + dst = [np.full((32, _INDEX_PAGE_BYTES), 0xA5, dtype=np.uint8)] + mgr, peer = _manager(0, 4, src, dst) + mgr.conclude_failure = Mock() + peer.dst_state_item_lens = [[_INDEX_PAGE_BYTES - 64 * 4]] + self.assertEqual(_send_state(mgr, peer, [4, 8], [7, 2]), -1) + mgr.conclude_failure.assert_called_once() + self.assertEqual(mgr.blocks, []) + self.assertTrue(np.all(dst[0] == 0xA5)) + + def test_skip_topk_layers_keep_pairing_without_empty_rdma_descriptors(self): + for skip_layers in ([False, True, False, True], [True] * 4): + with self.subTest(skip_layers=skip_layers): + sources = [ + np.full( + (0 if skip else 8, _INDEX_PAGE_BYTES), layer, dtype=np.uint8 + ) + for layer, skip in enumerate(skip_layers) + ] + destinations = [ + np.full( + (0 if skip else 32, _INDEX_PAGE_BYTES), 0xA5, dtype=np.uint8 + ) + for skip in skip_layers + ] + mgr, peer = _manager(0, 4, sources, destinations) + self.assertEqual(_send_state(mgr, peer, [4, 8], [7, 2]), 0) + self.assertTrue(all(length > 0 for _, _, length in mgr.blocks)) + for src, dst, skip in zip(sources, destinations, skip_layers): + if not skip: + np.testing.assert_array_equal(dst[[7, 2]], src[[1, 2]]) + self.assertEqual( + sum(length for _, _, length in mgr.blocks), + skip_layers.count(False) * 2 * _INDEX_PAGE_BYTES, + ) + + def test_unsharded_cp_policy_and_state_copy_are_unchanged(self): + src = [np.arange(8 * _INDEX_PAGE_BYTES, dtype=np.uint8).reshape(8, -1)] + dst = [np.full_like(src[0], 0xA5)] + mgr, peer = _manager(0, 1, src, dst) + self.assertEqual(_send_state(mgr, peer, [5, 2], [0, 7]), 0) + np.testing.assert_array_equal(dst[0][[0, 7]], src[0][[5, 2]]) + mgr.attn_cp_size, mgr.attn_cp_rank = 8, 3 + with patch( + "sglang.srt.disaggregation.common.conn.get_parallel", + return_value=SimpleNamespace(enable_dsa_cache_layer_split=False), + ): + self.assertTrue(mgr._should_skip_cp_replicated_state_transfer()) + + def test_empty_final_chunk_is_enqueued_and_metric_counts_only_owned_state(self): + for rank, expected_count in ((0, 2), (1, 1), (2, 0), (3, 0)): + with self.subTest(rank=rank): + mgr = CommonKVManager.__new__(CommonKVManager) + mgr.attn_cp_size = mgr.kv_shard_size = 4 + mgr.attn_cp_rank = mgr.kv_shard_rank = rank + mgr.add_transfer_request = Mock() + mgr.kv_args = SimpleNamespace(state_types=[StateType.DSA]) + mgr.kv_item_lens_sum = 64 * 576 * 4 # 3 target + 1 draft + mgr.state_item_lens_sum = _INDEX_PAGE_BYTES * 4 + mgr._kv_replica_factor = 1 + sender = MooncakeKVSender.__new__(MooncakeKVSender) + sender.kv_mgr = mgr + sender.bootstrap_room = _ROOM + sender.aux_index = 0 + sender.curr_idx = 0 + sender.num_kv_indices = 3 + sender.trace_ctx = Mock() + sender._transfer_metric = KVTransferMetric() + sender._transfer_num_kv_indices = 0 + sender._transfer_num_state_indices = 0 + pages = np.array([4, 9, 12], dtype=np.int32) + with patch( + "sglang.srt.disaggregation.common.conn.get_parallel", + return_value=SimpleNamespace(enable_dsa_cache_layer_split=False), + ): + sender.send(pages, [pages]) + mgr.add_transfer_request.assert_called_once() + call = mgr.add_transfer_request.call_args + self.assertTrue(call.args[3]) # final chunk, even with 0 pages + self.assertEqual(len(call.args[1]), expected_count) + np.testing.assert_array_equal(call.kwargs["state_indices"][0], pages) + self.assertEqual(sender._transfer_num_state_indices, expected_count) + self.assertEqual( + sender.get_transfer_metric().transfer_total_bytes, + expected_count * (mgr.kv_item_lens_sum + mgr.state_item_lens_sum), + ) + + def test_chunked_kv_delta_does_not_filter_away_full_prefix_indexer_state(self): + # Decode cached two prefix pages. Latent transfer contains only the + # suffix; DSA state still addresses the full sequence. Rank 3 owns one + # prefix indexer page but no suffix KV page. Its empty final KV chunk + # must retain that state, with independently filtered state metrics. + src = [np.full((8, _INDEX_PAGE_BYTES), 0x19, dtype=np.uint8)] + dst = [np.full((32, _INDEX_PAGE_BYTES), 0xA5, dtype=np.uint8)] + mgr, peer = _manager(3, 4, src, dst) + mgr.add_transfer_request = Mock() + sender = MooncakeKVSender.__new__(MooncakeKVSender) + sender.kv_mgr = mgr + sender.bootstrap_room = _ROOM + sender.aux_index = 0 + sender.curr_idx = 0 + sender.num_kv_indices = 3 + sender.trace_ctx = Mock() + sender._transfer_metric = KVTransferMetric() + sender._transfer_num_kv_indices = 0 + sender._transfer_num_state_indices = 0 + full_pages = np.array([3, 6, 4, 9, 12], dtype=np.int32) + with patch( + "sglang.srt.disaggregation.common.conn.get_parallel", + return_value=SimpleNamespace(enable_dsa_cache_layer_split=False), + ): + sender.send(full_pages[2:4]) + sender.send(full_pages[4:], [full_pages]) + calls = mgr.add_transfer_request.call_args_list + self.assertEqual(len(calls), 2) + self.assertFalse(calls[0].args[3]) + self.assertTrue(calls[1].args[3]) + self.assertEqual(len(calls[1].args[1]), 0) + self.assertEqual(sender._transfer_num_kv_indices, 0) + self.assertEqual(sender._transfer_num_state_indices, 1) + np.testing.assert_array_equal(calls[1].kwargs["state_indices"][0], full_pages) + self.assertEqual(_send_state(mgr, peer, full_pages, [7, 2, 4, 8, 11]), 0) + expected = np.full_like(dst[0], 0xA5) + expected[7] = src[0][0] # logical page 3 on rank 3 is local page 0 + np.testing.assert_array_equal(dst[0], expected) + + +if __name__ == "__main__": + unittest.main()