From 99d65b886cdbe7675705082d5007c6d900e08cf2 Mon Sep 17 00:00:00 2001 From: Tom Date: Thu, 10 Sep 2026 14:25:38 +0800 Subject: [PATCH] Report prompt KV weight versions as meta_info.prefill_weight_versions Add optional per-slot KV weight-version tracking and expose run-length encoded prompt provenance through generation metadata, including cache reuse and abort responses. Capture target write locations and serving versions alongside routed-expert capture. Carry records through the same result paths and finalize them beside routed-expert and indexer outputs after their existing synchronization. Centralize optional tracker creation and reject encoder-decoder, embedding/reward, and pipeline-parallel configurations when enabled. Include unit coverage and speculative multi-turn cache-reuse assertions. Align regression tests with the existing merged IPC role field and session-end version publication contract. Keep strict schema and pre-publication assertions. Limit long test generations to 3072 tokens so prompt plus completion fits the 4096-token CI model context; retain all boundary and abort assertions. Validated on two H200 GPUs: 263 unit tests and 24 RL tests passed, including EAGLE, DP Attention, and multi-turn prefix reuse. --- docs/docs/advanced_features/sglang_for_rl.mdx | 1 + python/sglang/srt/arg_groups/field_order.py | 1 + python/sglang/srt/arg_groups/fields/exec_.py | 4 + python/sglang/srt/arg_groups/kv_cache_hook.py | 14 + python/sglang/srt/disaggregation/prefill.py | 3 + .../srt/managers/detokenizer_manager.py | 1 + python/sglang/srt/managers/io_struct.py | 3 + .../srt/managers/multi_tokenizer_mixin.py | 6 + python/sglang/srt/managers/schedule_batch.py | 3 + python/sglang/srt/managers/scheduler.py | 13 + .../batch_result_processor.py | 12 + .../scheduler_components/output_streamer.py | 8 + .../sglang/srt/managers/tokenizer_manager.py | 15 +- python/sglang/srt/managers/tp_worker.py | 1 + python/sglang/srt/managers/utils.py | 3 + .../mem_cache/kv_weight_version_tracker.py | 159 +++++++++ .../sglang/srt/model_executor/model_runner.py | 6 + .../srt/speculative/dflash_worker_v2.py | 1 + .../srt/speculative/eagle_worker_common.py | 1 + python/sglang/srt/utils/weight_versions.py | 23 +- .../rl/test_weight_version_spans.py | 303 ++++++++++++++-- .../mlx/test_attention_patching.py | 1 + ...st_batch_result_processor_hidden_states.py | 3 + ...t_batch_result_processor_mamba_boundary.py | 1 + ...est_batch_result_processor_spec_grammar.py | 37 ++ .../managers/test_msgpack_ipc_roundtrip.py | 27 +- .../managers/test_multi_tokenizer_mixin.py | 23 ++ .../test_output_streamer_customized_info.py | 50 +++ .../test_scheduler_weight_version_tracking.py | 256 ++++++++++++- .../test_kv_weight_version_tracker.py | 336 ++++++++++++++++++ .../test_model_runner_kv_weight_versions.py | 103 ++++++ .../unit/server_args/test_server_args.py | 26 ++ .../unit/utils/test_weight_versions.py | 46 +++ 33 files changed, 1441 insertions(+), 49 deletions(-) create mode 100644 python/sglang/srt/mem_cache/kv_weight_version_tracker.py create mode 100644 test/registered/unit/mem_cache/test_kv_weight_version_tracker.py create mode 100644 test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py diff --git a/docs/docs/advanced_features/sglang_for_rl.mdx b/docs/docs/advanced_features/sglang_for_rl.mdx index 1c2c663dd6a3..e86875a1ff64 100644 --- a/docs/docs/advanced_features/sglang_for_rl.mdx +++ b/docs/docs/advanced_features/sglang_for_rl.mdx @@ -608,6 +608,7 @@ A request can outlive a version change: the RL flow retracts it with `pause_gene ``` - Ranges are half-open `[start, end)` over output-token indices, prompt excluded. The usual case is a single span. +- `--enable-prefill-weight-versions` adds `prefill_weight_versions`, the same span list over the prompt tokens, so KV reused from the prefix cache under older weights is visible. ## Easy To Postpone Generation diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index fd126f2a4061..464fd4ef5223 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/python/sglang/srt/arg_groups/field_order.py @@ -520,5 +520,6 @@ "enable_decoder_swa_bounded_replay", "enable_response_store", "disaggregation_decode_host_receive_threshold", + "enable_prefill_weight_versions", ) # fmt: on diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index 159d4b0a0b1c..bb3db33894d0 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -98,6 +98,10 @@ class ExecFeatures(msgspec.Struct): bool, "DeepSeek-V4.1 decoder SWA bounded replay: after the last kv_source layer, run the remaining layers over only the last window_size tokens of a prefill. Main and indexer KV stay exact; nothing is replayed. Deterministic for a fixed prompt and chunk size.", ] = False + enable_prefill_weight_versions: A[ + bool, + "Enable returning the weight version that computed each prompt token's KV cache with responses, as meta_info['prefill_weight_versions'].", + ] = False sampling_mask_max_tokens: A[ int, "The maximum number of token IDs in a returned sampling mask. Requests " diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 5db8194b9dda..94d47365322a 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -396,6 +396,20 @@ def handle_cache_compatibility(server_args: Any) -> None: "and cannot be used at the same time. Please use only one of them." ) + if cfg.enable_prefill_weight_versions: + if cfg.enable_hierarchical_cache: + raise ValueError( + "The argument enable-prefill-weight-versions cannot be combined with " + "enable-hierarchical-cache: KV restored from the host or storage tiers carries " + "no weight version, so the reported prompt versions would be wrong." + ) + if cfg.disaggregation_mode != "null": + raise ValueError( + "The argument enable-prefill-weight-versions is not supported in disaggregation " + f"mode {cfg.disaggregation_mode!r}: prefill workers do not report prompt weight " + "versions yet." + ) + if cfg.disaggregation_decode_enable_offload_kvcache: if cfg.disaggregation_mode != "decode": raise ValueError( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9383cfb679e9..64b23b65e939 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -795,6 +795,9 @@ def process_batch_result_disagg_prefill( if result.indexer_topk_output is not None: result.indexer_topk_output.finalize() result.indexer_topk_output = None + if (record := result.kv_weight_version_record) is not None: + record.finalize(tracker=self.kv_weight_version_tracker) + result.kv_weight_version_record = None logprob_pt = 0 aborted_reqs: List[Req] = [] diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index ef253a939838..01058f95a3de 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -504,6 +504,7 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput): placeholder_tokens_val=None, retraction_counts=recv_obj.retraction_counts, weight_versions=recv_obj.weight_versions, + prefill_weight_versions=recv_obj.prefill_weight_versions, token_steps=recv_obj.token_steps, beam_search_output=recv_obj.beam_search_output, dp_ranks=recv_obj.dp_ranks, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 4ebce9cb53d7..a5044cc114e0 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1588,6 +1588,7 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): beam_search_output: Optional[List[Optional[BeamSearchOutput]]] = None weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None + prefill_weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None # The trainer step id. Used to know which step's weights are used for sampling. token_steps: Optional[List[List[int]]] = None @@ -1686,6 +1687,7 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): beam_search_output: Optional[List[Optional[BeamSearchOutput]]] = None weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None + prefill_weight_versions: Optional[List[Optional[WeightVersionSpans]]] = None # The trainer step id. Used to know which step's weights are used for sampling. token_steps: Optional[List[List[int]]] = None @@ -2232,6 +2234,7 @@ class AbortReq(BaseReq, kw_only=True): # because batch requests derive child rids as ``f"{rid}_{i}"`` and an # abort for the parent rid must cover them. prefix: bool = False + prefill_weight_versions: Optional[WeightVersionSpans] = None def __post_init__(self): # FIXME: This is a hack to keep the same with the old code diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py index 150ee49055b5..0021979fa948 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -259,6 +259,9 @@ def _handle_output_by_index(output, i): indexer_topk_num_layers=getattr(output, "indexer_topk_num_layers", None), retraction_counts=_extract_field_by_index(output, "retraction_counts", i), weight_versions=_extract_field_by_index(output, "weight_versions", i), + prefill_weight_versions=_extract_field_by_index( + output, "prefill_weight_versions", i + ), placeholder_tokens_idx=None, placeholder_tokens_val=None, beam_search_output=_extract_field_by_index( @@ -387,6 +390,9 @@ def _handle_output_by_index(output, i): output, "beam_search_output", i, check_length=True ), weight_versions=_extract_field_by_index(output, "weight_versions", i), + prefill_weight_versions=_extract_field_by_index( + output, "prefill_weight_versions", i + ), token_steps=_extract_field_by_index( output, "token_steps", i, check_length=False ), diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index a94a1df7ae7a..74bed1fcb048 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -22,6 +22,7 @@ ) from sglang.srt.utils.weight_versions import ( WeightVersionEvent, + WeightVersionSpans, truncate_weight_version_events, ) @@ -1230,6 +1231,7 @@ def __init__( self.retracted_stain = False self.weight_version_events: List[WeightVersionEvent] = [] + self.prefill_weight_versions: Optional[WeightVersionSpans] = None # Incremental streamining self.send_token_offset: int = 0 @@ -1994,6 +1996,7 @@ def reset_for_retract(self): self.already_computed = 0 assert not self.kv.holds_kv, "expect it is already released" self.kv.kv_committed_len = 0 + self.prefill_weight_versions = None self.extend_batch_idx = 0 self.decode_batch_idx = 0 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8bd6402da8e0..a74a095e3d0c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -297,6 +297,7 @@ maybe_cache_unfinished_req, release_kv_cache, ) +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionTracker from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.model_executor.runner_utils.pool import prewarm_graph_pool_borrow from sglang.srt.model_loader.utils import get_resolved_model_impl @@ -710,6 +711,8 @@ def __init__( self.init_load_inquirer() + self.init_kv_weight_version_tracker() + self.init_output_streamer() self.init_batch_result_processor() @@ -2466,6 +2469,14 @@ def init_load_inquirer(self) -> None: get_decode_moment_totals=lambda: self.decode_moment_totals, ) + def init_kv_weight_version_tracker(self) -> None: + self.kv_weight_version_tracker = KvWeightVersionTracker.maybe_create( + server_args=self.server_args, + model_config=self.model_config, + allocator=self.token_to_kv_pool_allocator, + req_to_token_pool=self.req_to_token_pool, + ) + def init_output_streamer(self) -> None: self.output_streamer = self.get_output_streamer_class()( send_to_detokenizer=self.ipc_channels.send_to_detokenizer, @@ -2508,6 +2519,7 @@ def init_batch_result_processor(self) -> None: tree_cache=self.tree_cache, hisparse_coordinator=self.hisparse_coordinator, req_to_token_pool=self.req_to_token_pool, + kv_weight_version_tracker=self.kv_weight_version_tracker, decode_offload_manager=self.decode_offload_manager, metrics_collector=self.metrics_collector, metrics_reporter=self.metrics_reporter, @@ -6136,4 +6148,5 @@ def _make_abort_req( current_version=get_serving().weight_version, num_output_tokens=len(req.output_ids), ), + prefill_weight_versions=req.prefill_weight_versions, ) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index e0f922fe8df6..9ebcfbcd064f 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -76,6 +76,7 @@ ) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionTracker from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector @@ -103,6 +104,7 @@ class SchedulerBatchResultProcessor: tree_cache: BasePrefixCache hisparse_coordinator: Optional[HiSparseCoordinator] req_to_token_pool: ReqToTokenPool + kv_weight_version_tracker: Optional[KvWeightVersionTracker] decode_offload_manager: Optional[DecodeKVCacheOffloadManager] metrics_collector: SchedulerMetricsCollector metrics_reporter: SchedulerMetricsReporter @@ -275,6 +277,9 @@ def process_batch_result_prefill( if result.indexer_topk_output is not None: result.indexer_topk_output.finalize() result.indexer_topk_output = None + if (record := result.kv_weight_version_record) is not None: + record.finalize(tracker=self.kv_weight_version_tracker) + result.kv_weight_version_record = None ( logits_output, @@ -343,6 +348,8 @@ def process_batch_result_prefill( if req.inflight_middle_chunks <= 0: req.time_stats.set_prefill_finished_time() + if (x := self.kv_weight_version_tracker) is not None: + x.fill_req_prefill_weight_versions(req) if sampling_mask_finish_reason is not None: req.to_finish = sampling_mask_finish_reason @@ -463,6 +470,8 @@ def process_batch_result_prefill( req.pooled_hidden_state = phs[i] if req.inflight_middle_chunks <= 0: req.time_stats.set_prefill_finished_time() + if (x := self.kv_weight_version_tracker) is not None: + x.fill_req_prefill_weight_versions(req) # Dummy output token for embedding models req.output_ids.append(0) req.update_finish_state() @@ -933,6 +942,9 @@ def process_batch_result_decode( if result.indexer_topk_output is not None: result.indexer_topk_output.finalize() result.indexer_topk_output = None + if (record := result.kv_weight_version_record) is not None: + record.finalize(tracker=self.kv_weight_version_tracker) + result.kv_weight_version_record = None logits_output, next_token_ids, can_run_cuda_graph = ( result.logits_output, diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index c394a9f64979..9594e3c2230f 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py @@ -356,6 +356,7 @@ class _GenerationStreamAccumulator: spec_cap_lens_histogram: list = field(default_factory=list) retraction_counts: list = field(default_factory=list) weight_versions: list = field(default_factory=list) + prefill_weight_versions: list = field(default_factory=list) output_hidden_states: Optional[list] = None routed_experts: Optional[list] = None indexer_topk: Optional[list] = None @@ -524,8 +525,10 @@ def accept(self, *, req: Req) -> None: num_output_tokens=len(output_ids_), ) ) + self.prefill_weight_versions.append(req.prefill_weight_versions) else: self.weight_versions.append(None) + self.prefill_weight_versions.append(None) self.time_stats.append(req.time_stats) @@ -768,5 +771,10 @@ def to_payload( weight_versions=( self.weight_versions if any(self.weight_versions) else None ), + prefill_weight_versions=( + self.prefill_weight_versions + if any(self.prefill_weight_versions) + else None + ), dp_ranks=dp_ranks, ) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 110c1d9fd573..81cd975306d9 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -182,7 +182,10 @@ from sglang.srt.utils.network import get_zmq_socket from sglang.srt.utils.request_logger import RequestLogger from sglang.srt.utils.watchdog import Watchdog -from sglang.srt.utils.weight_versions import add_weight_versions_to_meta_info +from sglang.srt.utils.weight_versions import ( + add_weight_versions_to_meta_info, + weight_version_spans_to_json, +) from sglang.utils import TypeBasedDispatcher, get_exception_traceback asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) @@ -2524,6 +2527,12 @@ async def _handle_batch_output( spans, num_output_tokens=recv_obj.completion_tokens[i], ) + if recv_obj.prefill_weight_versions and ( + prefill_spans := recv_obj.prefill_weight_versions[i] + ): + meta_info["prefill_weight_versions"] = weight_version_spans_to_json( + prefill_spans + ) # Add detailed cache breakdown if available if ( hasattr(recv_obj, "cached_tokens_details") @@ -3501,6 +3510,10 @@ def _handle_abort_req(self, recv_obj: AbortReq): recv_obj.weight_versions, num_output_tokens=len(state.output_ids), ) + if prefill_spans := recv_obj.prefill_weight_versions: + meta_info["prefill_weight_versions"] = weight_version_spans_to_json( + prefill_spans + ) is_stream = getattr(state.obj, "stream", False) if getattr(state.obj, "return_logprob", False): self.add_logprob_to_meta_info( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index f38f6e6b57d6..90c1b0ce6653 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -600,6 +600,7 @@ def forward_batch_generation( expert_distribution_metrics=out.expert_distribution_metrics, routed_experts_output=out.routed_experts_output, indexer_topk_output=out.indexer_topk_output, + kv_weight_version_record=out.kv_weight_version_record, ) capture_pre_sample_logits(batch, forward_batch, logits_output) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index ed40ce3694a5..7a6a49f0f7c8 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -17,6 +17,7 @@ ) from sglang.srt.managers import io_struct from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.runtime_context import get_spec, max_speculative_num_draft_tokens from sglang.srt.speculative.spec_info import SpeculativeAlgorithm @@ -131,6 +132,7 @@ class GenerationBatchResult: # Routed experts: pending async D2H for overlap scheduling routed_experts_output: Optional[TopkCaptureOutput] = None indexer_topk_output: Optional[TopkCaptureOutput] = None + kv_weight_version_record: Optional[KvWeightVersionRecord] = None # metrics expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None @@ -207,6 +209,7 @@ def copy_to_cpu(self, return_logprob: bool, return_hidden_states: bool = True): sampling_mask_output, self.routed_experts_output, self.indexer_topk_output, + self.kv_weight_version_record, self.expert_distribution_metrics, ): if holder is not None: diff --git a/python/sglang/srt/mem_cache/kv_weight_version_tracker.py b/python/sglang/srt/mem_cache/kv_weight_version_tracker.py new file mode 100644 index 000000000000..3633f48e4214 --- /dev/null +++ b/python/sglang/srt/mem_cache/kv_weight_version_tracker.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Dict, List, Optional + +import msgspec +import torch + +from sglang.srt.runtime_context import get_serving +from sglang.srt.utils.weight_versions import WeightVersionSpan, WeightVersionSpans + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.server_args import ServerArgs + + +_UNWRITTEN_VERSION_ID = -1 + + +class KvWeightVersionRecord(msgspec.Struct): + slot_indices: torch.Tensor + version: str + + @classmethod + def maybe_capture( + cls, model_runner: ModelRunner, forward_batch: ForwardBatch + ) -> Optional[KvWeightVersionRecord]: + if ( + not model_runner.is_draft_worker + and model_runner.server_args.enable_prefill_weight_versions + and forward_batch.out_cache_loc is not None + ): + return cls.capture( + slot_indices=forward_batch.out_cache_loc, + version=get_serving().weight_version, + ) + return None + + @classmethod + def capture( + cls, *, slot_indices: torch.Tensor, version: Optional[str] + ) -> KvWeightVersionRecord: + assert version is not None + return cls(slot_indices=slot_indices.clone(), version=version) + + def map_device_tensors(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> None: + self.slot_indices = fn(self.slot_indices) + + def finalize(self, *, tracker: Optional[KvWeightVersionTracker]) -> None: + assert tracker is not None + tracker.record(slot_indices=self.slot_indices, version=self.version) + + +class KvWeightVersionTracker: + @classmethod + def maybe_create( + cls, + *, + server_args: ServerArgs, + model_config: ModelConfig, + allocator: BaseTokenToKVPoolAllocator, + req_to_token_pool: ReqToTokenPool, + ) -> Optional[KvWeightVersionTracker]: + if not server_args.enable_prefill_weight_versions: + return None + + assert not model_config.is_encoder_decoder, ( + "--enable-prefill-weight-versions does not support encoder-decoder models" + ) + assert model_config.is_generation, ( + "--enable-prefill-weight-versions does not support embedding or reward models" + ) + assert server_args.pp_size == 1, ( + "--enable-prefill-weight-versions does not support pipeline parallelism" + ) + + from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator + + if isinstance(allocator, SWATokenToKVPoolAllocator): + allocator = allocator.full_attn_allocator + + return cls( + num_slots=allocator.size_full + allocator.page_size, + device=allocator.device, + req_to_token_pool=req_to_token_pool, + ) + + def __init__( + self, *, num_slots: int, device: str, req_to_token_pool: ReqToTokenPool + ): + self._slot_version_ids = torch.full( + (num_slots,), _UNWRITTEN_VERSION_ID, dtype=torch.int32, device=device + ) + self._req_to_token_pool = req_to_token_pool + self._versions = _StringInterner() + + def record(self, *, slot_indices: torch.Tensor, version: Optional[str]) -> None: + assert version is not None + self._slot_version_ids[slot_indices] = self._versions.intern(version) + + def fill_req_prefill_weight_versions(self, req: Req) -> None: + num_prompt_tokens = len(req.origin_input_ids) + assert req.kv.kv_committed_len >= num_prompt_tokens, ( + f"prefill finished with {req.kv.kv_committed_len} committed KV tokens " + f"for a {num_prompt_tokens}-token prompt" + ) + req.prefill_weight_versions = self._lookup_spans( + self._req_to_token_pool.req_to_token[ + req.kv.req_pool_idx, :num_prompt_tokens + ] + ) + + def _lookup_spans(self, slot_indices: torch.Tensor) -> WeightVersionSpans: + version_ids = self._slot_version_ids[slot_indices] + if len(version_ids) == 0: + return [] + if (is_unwritten := version_ids == _UNWRITTEN_VERSION_ID).any(): + raise ValueError( + "KV slots without a recorded weight version were looked up: " + f"{slot_indices[is_unwritten].tolist()}" + ) + + version_changes_at: List[int] = ( + (version_ids[1:] != version_ids[:-1]).nonzero().flatten().tolist() + ) + run_starts = [0] + [position + 1 for position in version_changes_at] + run_ends = run_starts[1:] + [len(version_ids)] + run_version_ids: List[int] = version_ids[run_starts].tolist() + + return [ + WeightVersionSpan( + version=self._versions.lookup(version_id), start=start, end=end + ) + for version_id, start, end in zip( + run_version_ids, run_starts, run_ends, strict=True + ) + ] + + +class _StringInterner: + def __init__(self): + self._str_by_id: List[str] = [] + self._id_by_str: Dict[str, int] = {} + + def intern(self, value: str) -> int: + if (value_id := self._id_by_str.get(value)) is not None: + return value_id + + value_id = len(self._str_by_id) + self._str_by_id.append(value) + self._id_by_str[value] = value_id + return value_id + + def lookup(self, value_id: int) -> str: + return self._str_by_id[value_id] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 9f7f8636d226..437a632441ab 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -87,6 +87,7 @@ KVCacheConfigurator, ) from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.model_executor.cuda_graph_config import ( Backend, @@ -286,6 +287,7 @@ class ModelRunnerOutput: expert_distribution_metrics: Optional[ExpertDistributionMetrics] = None routed_experts_output: Optional[TopkCaptureOutput] = None indexer_topk_output: Optional[TopkCaptureOutput] = None + kv_weight_version_record: Optional[KvWeightVersionRecord] = None def resolve_draft_attention_backend( @@ -1806,6 +1808,10 @@ def forward( no_copy_to_cpu=no_copy_to_cpu, ) + output.kv_weight_version_record = KvWeightVersionRecord.maybe_capture( + model_runner=self, forward_batch=forward_batch + ) + # should_advance_eplb_counter is always True on non-hip platform if self.eplb_manager is not None and should_advance_eplb_counter(forward_batch): self.eplb_manager.on_forward_pass_end() diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 35361a58a518..f570abe4b789 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -2896,4 +2896,5 @@ def forward_batch_generation( new_seq_lens=new_seq_lens, routed_experts_output=target_out.routed_experts_output, indexer_topk_output=target_out.indexer_topk_output, + kv_weight_version_record=target_out.kv_weight_version_record, ) diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py index 0c2c53086020..f390e7d6cf4b 100644 --- a/python/sglang/srt/speculative/eagle_worker_common.py +++ b/python/sglang/srt/speculative/eagle_worker_common.py @@ -671,5 +671,6 @@ def run_eagle_verify( new_seq_lens=new_seq_lens, routed_experts_output=forward_batch_output.routed_experts_output, indexer_topk_output=forward_batch_output.indexer_topk_output, + kv_weight_version_record=forward_batch_output.kv_weight_version_record, extra_keep_alive_refs=[verify_forward_batch], ) diff --git a/python/sglang/srt/utils/weight_versions.py b/python/sglang/srt/utils/weight_versions.py index 6ea1db20a91b..c64cd3479b2c 100644 --- a/python/sglang/srt/utils/weight_versions.py +++ b/python/sglang/srt/utils/weight_versions.py @@ -1,6 +1,7 @@ from __future__ import annotations import dataclasses +import sys from typing import TYPE_CHECKING, Any, Dict, Iterable, List import msgspec @@ -96,17 +97,21 @@ def add_weight_versions_to_meta_info( span for span in spans if span.start < num_output_tokens or span.start == 0 ] - meta_info["weight_versions"] = [ - { - "version": span.version, - "start": span.start, - "end": min(span.end, num_output_tokens), - } - for span in visible - ] + meta_info["weight_versions"] = weight_version_spans_to_json( + visible, end_limit=num_output_tokens + ) meta_info["weight_version"] = visible[-1].version +def weight_version_spans_to_json( + spans: WeightVersionSpans, end_limit: int = sys.maxsize +) -> List[Dict[str, Any]]: + return [ + {"version": span.version, "start": span.start, "end": min(span.end, end_limit)} + for span in spans + ] + + # ====================================================================== # OpenAI-compatible endpoints # ====================================================================== @@ -114,4 +119,6 @@ def build_endpoint_weight_version_metadata(meta_info: Dict[str, Any]) -> Dict[st metadata = {"weight_version": meta_info["weight_version"]} if "weight_versions" in meta_info: metadata["weight_versions"] = meta_info["weight_versions"] + if "prefill_weight_versions" in meta_info: + metadata["prefill_weight_versions"] = meta_info["prefill_weight_versions"] return metadata diff --git a/test/registered/rl/test_weight_version_spans.py b/test/registered/rl/test_weight_version_spans.py index f5fa8d4fe3f8..4ded05a71b6e 100644 --- a/test/registered/rl/test_weight_version_spans.py +++ b/test/registered/rl/test_weight_version_spans.py @@ -2,6 +2,7 @@ import time import unittest from concurrent.futures import ThreadPoolExecutor +from typing import Optional import requests @@ -15,21 +16,37 @@ popen_launch_server, ) -register_cuda_ci( - est_time=180, stage="nightly", runner_config="2-gpu-large", nightly=True -) +register_cuda_ci(est_time=180, stage="nightly", runner_config="2-gpu-large") _REQUEST_TIMEOUT = 180 +_LONG_GENERATION_TOKENS = 3072 + +_SHARED_PREFIX = ( + "You are a meticulous assistant. Follow every instruction to the letter, " + "explain your reasoning step by step, and never skip an intermediate result. " + "Here is the background material you must rely on for the whole conversation. " +) def _assert_spans_contiguous(test, meta_info): - spans = meta_info["weight_versions"] + spans = _assert_contiguous(test, meta_info["weight_versions"]) + test.assertEqual(meta_info["weight_version"], spans[-1]["version"]) + return spans + + +def _assert_prefill_spans_contiguous(test, meta_info, prompt_tokens=None): + spans = _assert_contiguous(test, meta_info["prefill_weight_versions"]) + if prompt_tokens is not None: + test.assertEqual(spans[-1]["end"], prompt_tokens) + return spans + + +def _assert_contiguous(test, spans): test.assertGreater(len(spans), 0) test.assertEqual(spans[0]["start"], 0) for prev, cur in zip(spans, spans[1:]): test.assertEqual(prev["end"], cur["start"]) test.assertNotEqual(prev["version"], cur["version"]) - test.assertEqual(meta_info["weight_version"], spans[-1]["version"]) return spans @@ -65,6 +82,7 @@ def setUpClass(cls): "32", "--max-running-requests", "8", + "--enable-prefill-weight-versions", ], ) @@ -72,7 +90,13 @@ def setUpClass(cls): def tearDownClass(cls): kill_process_tree(cls.process.pid) - def _generate(self, max_new_tokens: int, prompt: str = "The capital of France is"): + def _generate( + self, + max_new_tokens: int, + prompt: str = "The capital of France is", + dp_rank: Optional[int] = None, + extra_key: Optional[str] = None, + ): response = requests.post( f"{self.base_url}/generate", json={ @@ -82,6 +106,8 @@ def _generate(self, max_new_tokens: int, prompt: str = "The capital of France is "max_new_tokens": max_new_tokens, "ignore_eos": True, }, + "routed_dp_rank": dp_rank, + "extra_key": extra_key, }, timeout=_REQUEST_TIMEOUT, ) @@ -107,6 +133,14 @@ def _continue(self): timeout=30, ).raise_for_status() + def _abort_all(self): + requests.post( + f"{self.base_url}/abort_request", json={"abort_all": True}, timeout=30 + ).raise_for_status() + + def _flush_cache(self): + requests.post(f"{self.base_url}/flush_cache", timeout=30).raise_for_status() + def _set_weight_version(self, new_version: str, abort_all_requests: bool = False): response = requests.post( f"{self.base_url}/update_weight_version", @@ -133,7 +167,7 @@ def _run_while_paused( num_requests: int, while_paused, mode: str = "retract", - max_new_tokens: int = 1024, + max_new_tokens: int = _LONG_GENERATION_TOKENS, ): with ThreadPoolExecutor(max_workers=num_requests) as executor: futures = [ @@ -164,6 +198,11 @@ def test_01_single_span_without_update(self): self.assertEqual(spans[0]["version"], "base-v0") self.assertEqual(spans[0]["end"], meta_info["completion_tokens"]) self.assertEqual(meta_info["weight_version"], "base-v0") + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual(len(prefill_spans), 1) + self.assertEqual(prefill_spans[0]["version"], "base-v0") def test_02_update_weight_version_endpoint_applies_to_new_requests(self): """The endpoint returns only once every scheduler stamps new requests with the new version.""" @@ -181,6 +220,10 @@ def test_02_update_weight_version_endpoint_applies_to_new_requests(self): spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(len(spans), 1) self.assertEqual(spans[0]["version"], "endpoint-v1") + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertIn(prefill_span["version"], ("base-v0", "endpoint-v1")) def test_03_spans_split_across_pause_update_continue(self): """Requests spanning pause -> update_weights_from_disk -> continue report one span per version.""" @@ -201,6 +244,9 @@ def test_03_spans_split_across_pause_update_continue(self): versions = [span["version"] for span in spans] self.assertEqual(versions[0], base_version) self.assertIn(versions[-1], (base_version, "disk-v2")) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) if len(spans) > 1: multi_span_count += 1 self.assertEqual(versions, [base_version, "disk-v2"]) @@ -235,6 +281,9 @@ def test_04_openai_metadata_contains_weight_versions(self): self.assertEqual(metadata["weight_version"], self._current_version()) self.assertEqual(spans[0]["start"], 0) self.assertEqual(spans[0]["end"], data["usage"]["completion_tokens"]) + _assert_prefill_spans_contiguous( + self, metadata, prompt_tokens=data["usage"]["prompt_tokens"] + ) def test_04b_openai_metadata_reports_the_first_choice_when_n_is_greater_than_one( self, @@ -261,18 +310,14 @@ def test_04b_openai_metadata_reports_the_first_choice_when_n_is_greater_than_one self.assertEqual(spans[0]["version"], metadata["weight_version"]) self.assertEqual(metadata["weight_version"], self._current_version()) self.assertEqual(spans[0]["start"], 0) + _assert_prefill_spans_contiguous( + self, metadata, prompt_tokens=data["usage"]["prompt_tokens"] + ) def test_05_aborted_retracted_requests_report_spans(self): """Requests aborted while retracted in the waiting queue still report their spans.""" - def abort_all(): - requests.post( - f"{self.base_url}/abort_request", - json={"abort_all": True}, - timeout=30, - ).raise_for_status() - - results = self._run_while_paused(num_requests=4, while_paused=abort_all) + results = self._run_while_paused(num_requests=4, while_paused=self._abort_all) aborted_with_spans = 0 for data in results: @@ -299,7 +344,6 @@ def test_06_running_requests_split_without_abort(self): num_requests=4, while_paused=lambda: self._set_weight_version("inplace-v3"), mode="in_place", - max_new_tokens=1024, ) self.assertEqual(self._current_version(), "inplace-v3") @@ -310,6 +354,10 @@ def test_06_running_requests_split_without_abort(self): spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) self.assertEqual(spans[0]["version"], previous_version) + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertNotEqual(prefill_span["version"], "inplace-v3") if len(spans) > 1: split_count += 1 self.assertEqual( @@ -339,9 +387,13 @@ def test_07_update_without_weight_version_does_not_split(self): self.assertEqual(len(spans), 1) self.assertEqual(spans[0]["version"], version) self.assertEqual(spans[0]["end"], meta_info["completion_tokens"]) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) def test_08_reannouncing_current_version_is_a_noop(self): """Re-announcing the version the server already has must not split anything.""" + self._flush_cache() version = self._current_version() results = self._run_while_paused( @@ -353,9 +405,14 @@ def test_08_reannouncing_current_version_is_a_noop(self): self.assertEqual(self._current_version(), version) for data in results: - spans = _assert_spans_contiguous(self, data["meta_info"]) + meta_info = data["meta_info"] + spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(len(spans), 1) self.assertEqual(spans[0]["version"], version) + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertEqual(prefill_span["version"], version) def test_09_three_spans_across_two_updates(self): """Two updates during one request produce three ordered, non-empty spans.""" @@ -365,7 +422,7 @@ def test_09_three_spans_across_two_updates(self): futures = [ executor.submit( self._generate, - max_new_tokens=2048, + max_new_tokens=_LONG_GENERATION_TOKENS, prompt=f"Write a long story about the number {i}.", ) for i in range(4) @@ -391,6 +448,10 @@ def test_09_three_spans_across_two_updates(self): for span in spans: self.assertGreater(span["end"], span["start"]) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) + for prefill_span in _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ): + self.assertNotIn(prefill_span["version"], ("multi-a", "multi-b")) if len(spans) == 3: three_span_count += 1 @@ -404,16 +465,9 @@ def test_10_abort_before_first_token_reports_empty_span(self): """A request aborted before producing a token still reports a well-formed span.""" version = self._current_version() - def abort_all(): - requests.post( - f"{self.base_url}/abort_request", - json={"abort_all": True}, - timeout=30, - ).raise_for_status() - results = self._run_while_paused( num_requests=16, - while_paused=abort_all, + while_paused=self._abort_all, ) empty_aborts = [ @@ -429,6 +483,7 @@ def abort_all(): [{"version": version, "start": 0, "end": 0}], ) self.assertEqual(meta_info["weight_version"], version) + self.assertNotIn("prefill_weight_versions", meta_info) def test_11_streaming_reports_spans_only_on_the_final_chunk(self): """Intermediate stream chunks carry no spans; the finishing chunk carries them all.""" @@ -462,12 +517,16 @@ def test_11_streaming_reports_spans_only_on_the_final_chunk(self): self.assertGreater(len(chunks), 1) for chunk in chunks[:-1]: self.assertNotIn("weight_versions", chunk["meta_info"]) + self.assertNotIn("prefill_weight_versions", chunk["meta_info"]) self.assertEqual(chunk["meta_info"]["weight_version"], version) meta_info = chunks[-1]["meta_info"] spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) self.assertEqual(spans[-1]["end"], max_new_tokens) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) def test_12_completions_endpoint_reports_metadata(self): """/v1/completions surfaces the spans the same way /v1/chat/completions does.""" @@ -491,6 +550,9 @@ def test_12_completions_endpoint_reports_metadata(self): self.assertEqual(spans[0]["version"], metadata["weight_version"]) self.assertEqual(spans[0]["start"], 0) self.assertEqual(spans[0]["end"], data["usage"]["completion_tokens"]) + _assert_prefill_spans_contiguous( + self, metadata, prompt_tokens=data["usage"]["prompt_tokens"] + ) def test_13_new_requests_after_the_endpoint_returns_see_the_new_version(self): """Once the endpoint returns, every concurrently admitted request stamps the new version.""" @@ -570,6 +632,9 @@ def test_16_retract_update_continue_keeps_exact_boundaries(self): spans = _assert_spans_contiguous(self, meta_info) self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) self.assertEqual(spans[0]["version"], previous_version) + _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) if len(spans) > 1: split_count += 1 self.assertEqual( @@ -585,6 +650,192 @@ def test_16_retract_update_continue_keeps_exact_boundaries(self): "speculative decoding was not recorded.", ) + def test_17_reused_prefix_keeps_the_version_that_computed_it(self): + """A radix hit on KV computed before a relabel shows up as a stale leading prefill span.""" + self._flush_cache() + self._generate( + max_new_tokens=8, prompt=_SHARED_PREFIX + "Name three colors.", dp_rank=0 + ) + version_before = self._current_version() + + self._set_weight_version("reuse-v4") + data = self._generate( + max_new_tokens=8, + prompt=_SHARED_PREFIX + "Name three colors. Then name three shapes.", + dp_rank=0, + ) + + meta_info = data["meta_info"] + self.assertGreater(meta_info["cached_tokens"], 0) + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual( + [span["version"] for span in prefill_spans], [version_before, "reuse-v4"] + ) + self.assertEqual(prefill_spans[0]["end"], meta_info["cached_tokens"]) + spans = _assert_spans_contiguous(self, meta_info) + self.assertEqual([span["version"] for span in spans], ["reuse-v4"]) + + def test_18_flushed_cache_recomputes_the_prefix(self): + """After a flush the whole prompt is recomputed under the current version.""" + self._flush_cache() + data = self._generate(max_new_tokens=8, prompt=_SHARED_PREFIX + "Say hello.") + + meta_info = data["meta_info"] + self.assertEqual(meta_info["cached_tokens"], 0) + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual(len(prefill_spans), 1) + self.assertEqual(prefill_spans[0]["version"], self._current_version()) + + def test_19_aborted_after_prefill_reports_both_spans(self): + """A request aborted mid-decode keeps its prompt spans and reports output spans.""" + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit( + self._generate, + max_new_tokens=_LONG_GENERATION_TOKENS, + prompt=_SHARED_PREFIX + "Write an even longer story.", + ) + time.sleep(1) + self._pause("in_place") + try: + self._abort_all() + finally: + self._continue() + data = future.result() + + meta_info = data["meta_info"] + self.assertEqual(meta_info["finish_reason"]["type"], "abort") + self.assertGreater(meta_info["completion_tokens"], 0) + spans = _assert_spans_contiguous(self, meta_info) + self.assertEqual(spans[-1]["end"], meta_info["completion_tokens"]) + _assert_prefill_spans_contiguous(self, meta_info) + + def test_20_prefix_grown_across_versions_reports_every_version(self): + """A prompt whose prefix was extended under three successive versions reports three prefill spans.""" + self._flush_cache() + self._set_weight_version("grow-v1") + self._generate(max_new_tokens=8, prompt=_SHARED_PREFIX + "Step one.", dp_rank=0) + self._set_weight_version("grow-v2") + self._generate( + max_new_tokens=8, + prompt=_SHARED_PREFIX + "Step one. Step two follows the first one.", + dp_rank=0, + ) + self._set_weight_version("grow-v3") + + data = self._generate( + max_new_tokens=8, + prompt=_SHARED_PREFIX + + "Step one. Step two follows the first one. Step three ends it.", + dp_rank=0, + ) + + meta_info = data["meta_info"] + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual( + [span["version"] for span in prefill_spans], + ["grow-v1", "grow-v2", "grow-v3"], + ) + self.assertEqual(prefill_spans[1]["end"], meta_info["cached_tokens"]) + self.assertEqual( + [span["version"] for span in _assert_spans_contiguous(self, meta_info)], + ["grow-v3"], + ) + + def test_21_multi_turn_reuses_the_previous_turn_under_its_version(self): + """The second turn of a conversation reports the first turn's prompt and reply as a stale span.""" + self._flush_cache() + self._set_weight_version("turn-v1") + first_prompt = _SHARED_PREFIX + "User: Name a planet.\nAssistant:" + first = self._generate(max_new_tokens=16, prompt=first_prompt, dp_rank=0) + self.assertGreater(first["meta_info"]["spec_verify_ct"], 0) + self.assertEqual( + [span["version"] for span in first["meta_info"]["prefill_weight_versions"]], + ["turn-v1"], + ) + + self._set_weight_version("turn-v2") + second = self._generate( + max_new_tokens=16, + prompt=first_prompt + + first["text"] + + "\nUser: Name another one.\nAssistant:", + dp_rank=0, + ) + + meta_info = second["meta_info"] + self.assertGreater(meta_info["spec_verify_ct"], 0) + self.assertGreater( + meta_info["cached_tokens"], first["meta_info"]["prompt_tokens"] + ) + prefill_spans = _assert_prefill_spans_contiguous( + self, meta_info, prompt_tokens=meta_info["prompt_tokens"] + ) + self.assertEqual( + [span["version"] for span in prefill_spans], ["turn-v1", "turn-v2"] + ) + self.assertEqual(prefill_spans[0]["end"], meta_info["cached_tokens"]) + self.assertEqual( + [span["version"] for span in _assert_spans_contiguous(self, meta_info)], + ["turn-v2"], + ) + + def test_22_the_other_dp_rank_recomputes_the_prefix(self): + """A prefix cached on one DP rank is not shared, so the other rank prefills it fresh.""" + self._flush_cache() + self._set_weight_version("rank-v1") + prompt = _SHARED_PREFIX + "Which rank served me?" + self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0) + self._set_weight_version("rank-v2") + + cached = self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0) + fresh = self._generate(max_new_tokens=8, prompt=prompt, dp_rank=1) + + self.assertEqual( + [ + span["version"] + for span in cached["meta_info"]["prefill_weight_versions"] + ], + ["rank-v1", "rank-v2"], + ) + self.assertEqual(fresh["meta_info"]["cached_tokens"], 0) + self.assertEqual( + [span["version"] for span in fresh["meta_info"]["prefill_weight_versions"]], + ["rank-v2"], + ) + + def test_23_extra_key_partitions_stay_apart(self): + """A prefix cached under one extra_key is invisible to another, so only the same key sees the stale span.""" + self._flush_cache() + self._set_weight_version("key-v1") + prompt = _SHARED_PREFIX + "Which partition served me?" + self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0, extra_key="alpha") + self._set_weight_version("key-v2") + + same_key = self._generate( + max_new_tokens=8, prompt=prompt, dp_rank=0, extra_key="alpha" + ) + other_key = self._generate( + max_new_tokens=8, prompt=prompt, dp_rank=0, extra_key="beta" + ) + no_key = self._generate(max_new_tokens=8, prompt=prompt, dp_rank=0) + + self.assertEqual( + [s["version"] for s in same_key["meta_info"]["prefill_weight_versions"]], + ["key-v1", "key-v2"], + ) + for data in (other_key, no_key): + self.assertEqual(data["meta_info"]["cached_tokens"], 0) + self.assertEqual( + [s["version"] for s in data["meta_info"]["prefill_weight_versions"]], + ["key-v2"], + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 38964a655d9d..1a673cb62849 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1203,6 +1203,7 @@ def test_finished_request_snapshots_before_release(self): tree_cache=tree_cache, hisparse_coordinator=SimpleNamespace(request_finished=lambda req: None), req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=None, diff --git a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py index 103432afbd62..e5d159b02b4d 100644 --- a/test/registered/unit/managers/test_batch_result_processor_hidden_states.py +++ b/test/registered/unit/managers/test_batch_result_processor_hidden_states.py @@ -39,6 +39,7 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc tree_cache=None, hisparse_coordinator=None, req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=metrics_reporter, @@ -236,6 +237,7 @@ def test_active_middle_chunk_advances_before_new_last_request(self): auxiliary_host_output=None, routed_experts_output=None, indexer_topk_output=None, + kv_weight_version_record=None, logits_output=SimpleNamespace( hidden_states=hidden_states, customized_info=None, @@ -291,6 +293,7 @@ def test_sampling_mask_middle_chunk_does_not_require_logits_output(self): routed_experts_output=None, indexer_topk_output=None, logits_output=None, + kv_weight_version_record=None, next_token_ids=torch.zeros(1, dtype=torch.int64), extend_input_len_per_req=None, extend_logprob_start_len_per_req=None, diff --git a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py index 97eb071fb1f0..ffce27861120 100644 --- a/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py +++ b/test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py @@ -67,6 +67,7 @@ def _make_processor() -> SchedulerBatchResultProcessor: hisparse_coordinator=None, beam_coordinator=MagicMock(), req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=metrics_reporter, diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index 6fc1a88c424e..0d69add93a1f 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -5,6 +5,7 @@ """ import unittest +from dataclasses import replace from types import SimpleNamespace import torch @@ -15,6 +16,7 @@ SchedulerBatchResultProcessor, ) from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionTracker from sglang.srt.sampling.sampling_params import ( REQUEST_REASONING_END_TOKEN_IDS_KEY, SamplingParams, @@ -75,6 +77,7 @@ def _make_processor() -> SchedulerBatchResultProcessor: tree_cache=None, hisparse_coordinator=None, req_to_token_pool=None, + kv_weight_version_tracker=None, decode_offload_manager=None, metrics_collector=None, metrics_reporter=SimpleNamespace(), @@ -154,6 +157,40 @@ def _commit_disagg_handoff( class TestSpecV2GrammarTruncation(CustomTestCase): + def test_grammar_commit_does_not_relabel_forward_recorded_kv(self) -> None: + """Grammar commitment leaves forward-recorded provenance unchanged.""" + for overlap in (False, True): + with self.subTest(overlap=overlap): + pool = SimpleNamespace(req_to_token=torch.tensor([[1, 2, 3, 7, 5, 6]])) + tracker = KvWeightVersionTracker( + num_slots=10, device="cpu", req_to_token_pool=pool + ) + tracker.record( + slot_indices=torch.tensor([1, 2, 3, 7, 5, 6]), version="v0" + ) + tracker.record(slot_indices=torch.tensor([7, 5, 6]), version="v1") + req = _make_req(terminate_after=2) + req.kv.req_pool_idx = 0 + req.kv.kv_committed_len = 3 + processor = replace( + _make_processor(), + enable_overlap=overlap, + req_to_token_pool=pool, + kv_weight_version_tracker=tracker, + ) + batch = _FakeBatch([req]) + batch.out_cache_loc = torch.tensor([1, 2, 3]) + result = _make_result(4, [3], [101, 102, 103, 0]) + + processor._resolve_spec_v2_tokens(result, batch) + + spans = tracker._lookup_spans(torch.tensor([1, 2, 3, 7, 5, 6])) + self.assertEqual( + [(span.version, span.start, span.end) for span in spans], + [("v0", 0, 3), ("v1", 3, 6)], + ) + self.assertEqual(req.kv.kv_committed_len, 5) + def test_resolve_truncates_after_grammar_completion(self): req = _make_req(terminate_after=2) proc = _make_processor() diff --git a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py index 65a28ebb9660..99f5ebd8304d 100644 --- a/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py +++ b/test/registered/unit/managers/test_msgpack_ipc_roundtrip.py @@ -178,12 +178,13 @@ def test_backup_dram_req_is_narrowed(self): pointer = decoded.weight_pointer_map["experts.0.gate_proj"] self.assertEqual((pointer.weight_ptr, pointer.byte_size), (8, 4)) - def test_check_weights_mirrors_match_pydantic_models(self): + def test_check_weights_mirrors_match_pydantic_models(self) -> None: + """Wire fields preserve checker metadata and add the merged runner role.""" # Field-parity guard: the msgspec wire structs must not drift from the # pydantic source of truth in weight_checker. self.assertEqual( {f.name for f in msgspec.structs.fields(ParallelismInfo)}, - set(PydanticParallelismInfo.model_fields), + set(PydanticParallelismInfo.model_fields) | {"role"}, ) self.assertEqual( {f.name for f in msgspec.structs.fields(ChecksumInfo)}, @@ -282,6 +283,28 @@ def test_abort_req_carries_spans(self): def test_abort_req_defaults_to_no_spans(self): """The field is optional on the wire, so an abort without spans decodes to None.""" self.assertIsNone(_round_trip(AbortReq(rid="r0")).weight_versions) + self.assertIsNone(_round_trip(AbortReq(rid="r0")).prefill_weight_versions) + + def test_abort_req_carries_prefill_spans(self): + """Prompt-side spans survive the wire independently of the sampling spans.""" + obj = AbortReq( + rid="r0", + prefill_weight_versions=[ + WeightVersionSpan(version="unknown", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=9), + ], + ) + + decoded = _double_hop(obj) + + self.assertIsNone(decoded.weight_versions) + self.assertEqual( + decoded.prefill_weight_versions, + [ + WeightVersionSpan(version="unknown", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=9), + ], + ) if __name__ == "__main__": diff --git a/test/registered/unit/managers/test_multi_tokenizer_mixin.py b/test/registered/unit/managers/test_multi_tokenizer_mixin.py index 02a05ca209e5..b73251bc9a09 100644 --- a/test/registered/unit/managers/test_multi_tokenizer_mixin.py +++ b/test/registered/unit/managers/test_multi_tokenizer_mixin.py @@ -97,6 +97,10 @@ def _make_batch_str_output() -> BatchStrOutput: ], [WeightVersionSpan(version="v2", start=0, end=2)], ], + prefill_weight_versions=[ + [WeightVersionSpan(version="v0", start=0, end=7)], + None, + ], ) @@ -148,6 +152,25 @@ def test_batch_str_output_without_weight_versions_stays_none(self): self.assertIsNone(_handle_output_by_index(output, 0).weight_versions) + def test_batch_str_output_keeps_prefill_weight_versions_nested_per_request(self): + """Prefill spans split per request the same way the sampling spans do.""" + output = _make_batch_str_output() + + self.assertEqual( + _handle_output_by_index(output, 0).prefill_weight_versions, + [[WeightVersionSpan(version="v0", start=0, end=7)]], + ) + self.assertEqual( + _handle_output_by_index(output, 1).prefill_weight_versions, [None] + ) + + def test_batch_str_output_without_prefill_weight_versions_stays_none(self): + """An output produced without the flag splits into None.""" + output = _make_batch_str_output() + output.prefill_weight_versions = None + + self.assertIsNone(_handle_output_by_index(output, 0).prefill_weight_versions) + def test_get_tokenizer_worker_class_uses_default(self): self.assertIs(get_tokenizer_worker_class(DefaultServerArgs()), TokenizerWorker) diff --git a/test/registered/unit/managers/test_output_streamer_customized_info.py b/test/registered/unit/managers/test_output_streamer_customized_info.py index 8817dc762561..d02e58dd8e1c 100644 --- a/test/registered/unit/managers/test_output_streamer_customized_info.py +++ b/test/registered/unit/managers/test_output_streamer_customized_info.py @@ -83,6 +83,7 @@ def __init__( self.multimodal_inputs = None self.customized_info = customized_info self.weight_version_events = [] + self.prefill_weight_versions = None def finished(self): return self._finished @@ -537,5 +538,54 @@ def test_clients_receive_the_sampler_rows_unchanged(self): self.assertEqual(received, expected) +class TestOutputStreamerPrefillWeightVersions(unittest.TestCase): + def test_payload_carries_prefill_spans_for_finished_requests(self): + """Prefill spans computed before the KV release ride out with the finished request.""" + streaming_req = _FakeReq("r0", [10, 11]) + finished_req = _FakeReq("r1", [20], finished=True) + finished_req.prefill_weight_versions = [ + WeightVersionSpan(version="v0", start=0, end=4), + WeightVersionSpan(version="v1", start=4, end=6), + ] + + accumulator = _accumulator(current_weight_version="v1") + accumulator.accept(req=streaming_req) + accumulator.accept(req=finished_req) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + + self.assertEqual( + payload.prefill_weight_versions, + [ + None, + [ + WeightVersionSpan(version="v0", start=0, end=4), + WeightVersionSpan(version="v1", start=4, end=6), + ], + ], + ) + + def test_an_unfinished_request_withholds_its_prefill_spans(self): + """Spans of a request still generating must not ship before its finishing chunk.""" + streaming_req = _FakeReq("r0", [10, 11]) + streaming_req.prefill_weight_versions = [ + WeightVersionSpan(version="v0", start=0, end=4) + ] + + accumulator = _accumulator(current_weight_version="v1") + accumulator.accept(req=streaming_req) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + + self.assertIsNone(payload.prefill_weight_versions) + + def test_payload_omits_prefill_spans_when_the_flag_is_off(self): + """With tracking disabled every request contributes None, so nothing goes on the wire.""" + accumulator = _accumulator(current_weight_version="v1") + accumulator.accept(req=_FakeReq("r0", [10], finished=True)) + payload = accumulator.to_payload(dp_rank=0, is_idle_batch=False) + + self.assertIsNone(payload.prefill_weight_versions) + self.assertIsNotNone(payload.weight_versions) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py index 8eb58463dd79..7e3c79651eb2 100644 --- a/test/registered/unit/managers/test_scheduler_weight_version_tracking.py +++ b/test/registered/unit/managers/test_scheduler_weight_version_tracking.py @@ -1,16 +1,31 @@ import unittest from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch +import torch + +from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.managers.io_struct import ( BeginWeightUpdateReqInput, EndWeightUpdateReqInput, ) +from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.scheduler_components.batch_result_processor import ( + SchedulerBatchResultProcessor, +) from sglang.srt.managers.scheduler_components.weight_updater import ( SchedulerWeightUpdaterManager, _WeightUpdateSession, ) +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.kv_weight_version_tracker import ( + KvWeightVersionRecord, + KvWeightVersionTracker, +) +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode +from sglang.srt.model_executor.model_runner import ModelRunner +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 @@ -101,6 +116,203 @@ def test_every_source_of_live_requests_is_stamped(self): ) +class TestSchedulerBatchWeightVersion(CustomTestCase): + def test_embedding_rejects_enabled_prefill_tracking(self) -> None: + """Embedding and reward models reject enabled prefill tracking before allocation.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=True) + scheduler.model_config = SimpleNamespace( + is_encoder_decoder=False, is_generation=False + ) + + with self.assertRaisesRegex( + AssertionError, "does not support embedding or reward" + ): + scheduler.init_kv_weight_version_tracker() + + def test_embedding_without_prefill_tracking_remains_supported(self) -> None: + """Embedding initialization remains unchanged when prefill tracking is disabled.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=False) + scheduler.model_config = SimpleNamespace( + is_encoder_decoder=False, is_generation=False + ) + + scheduler.init_kv_weight_version_tracker() + + self.assertIsNone(scheduler.kv_weight_version_tracker) + + def test_encoder_decoder_rejects_enabled_prefill_tracking(self) -> None: + """Encoder-decoder models reject prefill tracking before allocating a tracker.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=True) + scheduler.model_config = SimpleNamespace(is_encoder_decoder=True) + + with self.assertRaisesRegex(AssertionError, "does not support encoder-decoder"): + scheduler.init_kv_weight_version_tracker() + + def test_encoder_decoder_without_prefill_tracking_remains_supported(self) -> None: + """Disabled tracking leaves encoder-decoder initialization unchanged.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.server_args = SimpleNamespace(enable_prefill_weight_versions=False) + scheduler.model_config = SimpleNamespace(is_encoder_decoder=True) + + scheduler.init_kv_weight_version_tracker() + + self.assertIsNone(scheduler.kv_weight_version_tracker) + + def _scheduler(self) -> Scheduler: + scheduler = Scheduler.__new__(Scheduler) + scheduler.token_to_kv_pool_allocator = None + scheduler.req_to_token_pool = None + scheduler.publish_load_snapshot = MagicMock() + scheduler.scheduler_stage_metrics = None + scheduler.load_publisher = MagicMock() + scheduler.load_inquirer = MagicMock() + scheduler.tree_cache = SimpleNamespace(flush_pending_backups=MagicMock()) + scheduler.kv_weight_version_tracker = KvWeightVersionTracker( + num_slots=8, + device="cpu", + req_to_token_pool=SimpleNamespace(req_to_token=torch.tensor([[3]])), + ) + scheduler.batch_result_processor = SchedulerBatchResultProcessor( + is_generation=True, + disaggregation_mode=DisaggregationMode.NULL, + enable_overlap=False, + enable_overlap_mlx=False, + model_config=SimpleNamespace(think_end_ids=None), + token_to_kv_pool_allocator=MagicMock(), + tree_cache=scheduler.tree_cache, + hisparse_coordinator=None, + req_to_token_pool=None, + kv_weight_version_tracker=scheduler.kv_weight_version_tracker, + decode_offload_manager=None, + metrics_collector=None, + metrics_reporter=MagicMock(), + draft_worker=None, + model_worker=MagicMock(), + logprob_result_processor=None, + output_streamer=MagicMock(), + beam_coordinator=MagicMock(), + abort_request=lambda *args, **kwargs: None, + ) + for name, value in ( + ("get_observability", SimpleNamespace(enable_metrics=False)), + ("get_server_return_hidden_states_mode", CaptureHiddenMode.NULL), + ): + patcher = patch( + f"sglang.srt.managers.scheduler_components.batch_result_processor.{name}", + return_value=value, + ) + patcher.start() + self.addCleanup(patcher.stop) + scheduler.disaggregation_mode = DisaggregationMode.NULL + scheduler._record_step_counters = MagicMock() + scheduler.metrics_reporter = MagicMock() + scheduler.enable_fpm = False + scheduler._maybe_clear_mm_inputs = MagicMock() + scheduler.maybe_send_health_check_signal = MagicMock() + return scheduler + + def _batch(self) -> ScheduleBatch: + return ScheduleBatch( + reqs=[], + forward_mode=ForwardMode.EXTEND, + out_cache_loc=torch.tensor([3], dtype=torch.int64), + ) + + def test_pending_result_uses_version_captured_when_batch_ran(self) -> None: + """A pending result keeps the forward-time version across a serving update.""" + serving = _ServingStub("v0") + scheduler = self._scheduler() + batch = self._batch() + + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = False + runner.server_args = SimpleNamespace(enable_prefill_weight_versions=True) + with patch( + "sglang.srt.mem_cache.kv_weight_version_tracker.get_serving", + return_value=serving, + ): + result = GenerationBatchResult( + next_token_ids=torch.tensor([], dtype=torch.int64), + kv_weight_version_record=KvWeightVersionRecord.maybe_capture( + model_runner=runner, forward_batch=batch + ), + ) + queued_batch = batch.copy() + serving.weight_version = "v1" + Scheduler.process_batch_result(scheduler, queued_batch, result) + + spans = scheduler.kv_weight_version_tracker._lookup_spans( + queued_batch.out_cache_loc + ) + self.assertEqual( + [(span.version, span.start, span.end) for span in spans], [("v0", 0, 1)] + ) + + def test_spec_decode_does_not_restamp_restored_prefill_slots(self) -> None: + """Restoring the batch after speculative forward must not relabel old prompt KV.""" + for overlap in (False, True): + with self.subTest(overlap=overlap): + scheduler = self._scheduler() + scheduler.record_batch_in_overlap = lambda batch: None + batch = self._batch() + batch.forward_mode = ForwardMode.DECODE + batch.spec_algorithm = SpeculativeAlgorithm.EAGLE + tracker = scheduler.kv_weight_version_tracker + tracker.record(slot_indices=batch.out_cache_loc, version="v0") + + with scheduler._forward_isolation(batch, overlap=overlap): + forward_slots = torch.tensor([4, 5, 6]) + batch.out_cache_loc = forward_slots + result = GenerationBatchResult( + next_token_ids=torch.tensor([], dtype=torch.int64), + accept_lens=torch.tensor([], dtype=torch.int32), + speculative_num_draft_tokens=3, + kv_weight_version_record=KvWeightVersionRecord.capture( + slot_indices=forward_slots, version="v1" + ), + ) + forward_slots.fill_(7) + scheduler.process_batch_result(batch.copy(), result) + + spans = tracker._lookup_spans(torch.tensor([3, 4, 5, 6])) + self.assertEqual( + [(span.version, span.start, span.end) for span in spans], + [("v0", 0, 1), ("v1", 1, 4)], + ) + + def test_result_copy_finishes_before_slot_indices_are_consumed(self) -> None: + """Pending slot copies complete before the tracker reads their contents.""" + scheduler = self._scheduler() + batch = self._batch() + record = KvWeightVersionRecord.capture( + slot_indices=torch.tensor([0]), version="v0" + ) + synchronize = MagicMock(side_effect=lambda: record.slot_indices.fill_(3)) + result = GenerationBatchResult( + next_token_ids=torch.tensor([], dtype=torch.int64), + kv_weight_version_record=record, + copy_done=SimpleNamespace(synchronize=synchronize), + ) + + scheduler.process_batch_result(batch, result) + + spans = scheduler.kv_weight_version_tracker._lookup_spans(torch.tensor([3])) + self.assertEqual([span.version for span in spans], ["v0"]) + self.assertIsNone(result.kv_weight_version_record) + synchronize.assert_called_once_with() + + def _runner(result=(True, "ok")): runner = Mock() for method in ( @@ -110,6 +322,9 @@ def _runner(result=(True, "ok")): "load_weights_from_distributed", ): getattr(runner.weight_updater, method).return_value = result + runner.weight_updater.receive_weights_from_distributed.return_value = [ + ("weight", torch.tensor([1.0])) + ] return runner @@ -171,23 +386,27 @@ def _manager(self, target, draft=None, *, session=True): class TestRecordWeightVersionAfterUpdate(_WeightUpdaterManagerTestBase): def test_successful_update_records_the_version(self): """A refit that reports success advances the scheduler-side version.""" - output = self._manager(_runner()).update_weights_from_disk(_request()) + output = self._manager(_runner(), session=False).update_weights_from_disk( + _request() + ) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) def test_failed_update_does_not_record_the_version(self): """A refit that fails must leave the version alone, or later tokens are mislabelled.""" - output = self._manager(_runner((False, "boom"))).update_weights_from_disk( - _request() - ) + output = self._manager( + _runner((False, "boom")), session=False + ).update_weights_from_disk(_request()) self.assertFalse(output.success) self.assertEqual(self.recorded, []) def test_draft_failure_does_not_record_the_version(self): """The target succeeding is not enough: a failed draft refit leaves the engine mixed.""" - manager = self._manager(_runner(), draft=_runner((False, "draft boom"))) + manager = self._manager( + _runner(), draft=_runner((False, "draft boom")), session=False + ) output = manager.update_weights_from_disk(_request()) @@ -196,7 +415,12 @@ def test_draft_failure_does_not_record_the_version(self): def test_successful_distributed_update_records_the_version(self): """The distributed refit is the path an RL trainer actually drives, so it must record too.""" - output = self._manager(_runner()).update_weights_from_distributed(_request()) + manager = self._manager(_runner()) + output = manager.update_weights_from_distributed(_request()) + + self.assertTrue(output.success) + self.assertEqual(self.recorded, []) + output = manager.end_weight_update(EndWeightUpdateReqInput()) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) @@ -212,23 +436,31 @@ def test_failed_distributed_update_does_not_record_the_version(self): def test_successful_tensor_update_records_the_version(self): """The tensor refit records the version once the load reports success.""" - output = self._manager(_runner()).update_weights_from_tensor(_request()) + updater = self._manager(_runner()) + output = updater.update_weights_from_tensor(_request()) + + self.assertTrue(output.success) + self.assertEqual(self.recorded, []) + with patch("torch.distributed.barrier"): + output = updater.end_weight_update(EndWeightUpdateReqInput()) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) def test_successful_ipc_update_records_the_version(self): """The checkpoint-engine IPC refit records the version like every other path.""" - output = self._manager(_runner()).update_weights_from_ipc(_request()) + output = self._manager(_runner(), session=False).update_weights_from_ipc( + _request() + ) self.assertTrue(output.success) self.assertEqual(self.recorded, ["v2"]) def test_failed_ipc_update_does_not_record_the_version(self): """The IPC path branches on success separately from the cache flush, so failure must record nothing.""" - output = self._manager(_runner((False, "boom"))).update_weights_from_ipc( - _request() - ) + output = self._manager( + _runner((False, "boom")), session=False + ).update_weights_from_ipc(_request()) self.assertFalse(output.success) self.assertEqual(self.recorded, []) diff --git a/test/registered/unit/mem_cache/test_kv_weight_version_tracker.py b/test/registered/unit/mem_cache/test_kv_weight_version_tracker.py new file mode 100644 index 000000000000..4b90a25865e7 --- /dev/null +++ b/test/registered/unit/mem_cache/test_kv_weight_version_tracker.py @@ -0,0 +1,336 @@ +import unittest +from types import SimpleNamespace +from typing import List + +import torch + +from sglang.srt.mem_cache.kv_weight_version_tracker import ( + KvWeightVersionRecord, + KvWeightVersionTracker, + _StringInterner, +) +from sglang.srt.utils.weight_versions import WeightVersionSpan +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") + + +class _ReqStub: + def __init__(self, num_prompt_tokens: int, kv_committed_len: int): + self.origin_input_ids = [0] * num_prompt_tokens + self.kv = SimpleNamespace(kv_committed_len=kv_committed_len, req_pool_idx=1) + self.prefill_weight_versions = None + + +class _ReqToTokenPoolStub: + def __init__(self, slots_of_req: List[int]): + self.req_to_token = torch.zeros((2, 32), dtype=torch.int32) + self.req_to_token[1, : len(slots_of_req)] = torch.tensor( + slots_of_req, dtype=torch.int32 + ) + + +def _table(slots_of_req: List[int] = ()) -> KvWeightVersionTracker: + return KvWeightVersionTracker( + num_slots=16, + device="cpu", + req_to_token_pool=_ReqToTokenPoolStub(list(slots_of_req)), + ) + + +def _slots(*indices: int) -> torch.Tensor: + return torch.tensor(indices, dtype=torch.int64) + + +class TestKvWeightVersionRecord(CustomTestCase): + def test_capture_without_a_version_fails(self) -> None: + """A forward cannot publish KV provenance without a weight version.""" + with self.assertRaises(AssertionError): + KvWeightVersionRecord.capture(slot_indices=_slots(1), version=None) + + def test_capture_survives_source_reuse_and_tensor_mapping(self) -> None: + """Reusing forward buffers cannot change a pending record's slots.""" + slots = _slots(4, 5, 6) + record = KvWeightVersionRecord.capture(slot_indices=slots, version="v0") + slots.fill_(9) + record.map_device_tensors(lambda tensor: tensor.to(dtype=torch.int32)) + table = _table() + + table.record(slot_indices=record.slot_indices, version=record.version) + + self.assertEqual( + table._lookup_spans(_slots(4, 5, 6)), + [WeightVersionSpan(version="v0", start=0, end=3)], + ) + self.assertEqual(record.slot_indices.dtype, torch.int32) + + +class TestKvWeightVersionTracker(CustomTestCase): + def test_pipeline_parallelism_is_rejected_only_when_tracking_is_enabled( + self, + ) -> None: + """The optional tracker rejects PP before allocation without restricting disabled runs.""" + for enabled in (False, True): + with self.subTest(enabled=enabled): + kwargs = dict( + server_args=SimpleNamespace( + enable_prefill_weight_versions=enabled, pp_size=2 + ), + model_config=SimpleNamespace( + is_encoder_decoder=False, is_generation=True + ), + allocator=None, + req_to_token_pool=None, + ) + if enabled: + with self.assertRaisesRegex(AssertionError, "pipeline parallelism"): + KvWeightVersionTracker.maybe_create(**kwargs) + else: + self.assertIsNone(KvWeightVersionTracker.maybe_create(**kwargs)) + + def test_factory_includes_allocator_padding_and_uses_the_request_pool(self) -> None: + """Factory sizing includes the allocator's padding page and preserves its mapping.""" + tracker = KvWeightVersionTracker.maybe_create( + server_args=SimpleNamespace(enable_prefill_weight_versions=True, pp_size=1), + model_config=SimpleNamespace(is_encoder_decoder=False, is_generation=True), + allocator=SimpleNamespace(size_full=12, page_size=4, device="cpu"), + req_to_token_pool=_ReqToTokenPoolStub([15]), + ) + req = _ReqStub(num_prompt_tokens=1, kv_committed_len=1) + + tracker.record(slot_indices=_slots(15), version="v0") + tracker.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [WeightVersionSpan(version="v0", start=0, end=1)], + ) + + def test_unified_swa_virtual_slots_above_the_static_quota_keep_their_versions( + self, + ) -> None: + """Legal virtual slots above the static token quota remain addressable by the tracker.""" + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + for page_size in (1, 4): + with self.subTest(page_size=page_size): + allocator = UnifiedSWATokenToKVPoolAllocator.__new__( + UnifiedSWATokenToKVPoolAllocator + ) + allocator._size_full = 8 + allocator.page_size = page_size + allocator.device = "cpu" + allocator.full_attn_allocator = SimpleNamespace( + size_full=32, page_size=page_size, device="cpu" + ) + tracker = KvWeightVersionTracker.maybe_create( + server_args=SimpleNamespace( + enable_prefill_weight_versions=True, pp_size=1 + ), + model_config=SimpleNamespace( + is_encoder_decoder=False, is_generation=True + ), + allocator=allocator, + req_to_token_pool=_ReqToTokenPoolStub([31]), + ) + req = _ReqStub(num_prompt_tokens=1, kv_committed_len=1) + + tracker.record(slot_indices=_slots(31), version="v2") + tracker.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [WeightVersionSpan(version="v2", start=0, end=1)], + ) + + def test_equal_neighbours_merge_while_a_returning_version_starts_a_new_span(self): + """Run-length compression merges adjacent equal versions and never merges across a change.""" + table = _table() + table.record(slot_indices=_slots(4, 5), version="v0") + table.record(slot_indices=_slots(6), version="v1") + table.record(slot_indices=_slots(7), version="v0") + + self.assertEqual( + table._lookup_spans(_slots(4, 5, 6, 7)), + [ + WeightVersionSpan(version="v0", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=3), + WeightVersionSpan(version="v0", start=3, end=4), + ], + ) + + def test_a_single_slot_lookup_yields_one_unit_span(self): + """A one-token prompt maps to exactly one [0, 1) span.""" + table = _table() + table.record(slot_indices=_slots(7), version="v0") + + self.assertEqual( + table._lookup_spans(_slots(7)), + [WeightVersionSpan(version="v0", start=0, end=1)], + ) + + def test_an_empty_lookup_returns_no_spans(self): + """Looking up an empty prompt yields an empty span list.""" + self.assertEqual(_table()._lookup_spans(_slots()), []) + + def test_unwritten_slots_are_reported_by_index_in_lookup_order(self): + """The error names exactly the never-stamped slots, ordered by their place in the lookup.""" + table = _table() + table.record(slot_indices=_slots(1, 3), version="v0") + + with self.assertRaisesRegex(ValueError, r": \[2, 0\]$"): + table._lookup_spans(_slots(2, 3, 0, 1)) + + def test_a_reused_slot_reports_the_version_that_last_wrote_it(self): + """The table is keyed by slot, so a slot refilled under a new version forgets the old one.""" + table = _table() + table.record(slot_indices=_slots(3), version="v0") + table.record(slot_indices=_slots(3), version="v1") + + self.assertEqual( + table._lookup_spans(_slots(3)), + [WeightVersionSpan(version="v1", start=0, end=1)], + ) + + def test_int32_slot_indices_are_accepted(self): + """out_cache_loc arrives as int32 and must index the table like int64 does.""" + table = _table() + table.record(slot_indices=torch.tensor([2, 3], dtype=torch.int32), version="v0") + + self.assertEqual( + table._lookup_spans(torch.tensor([2, 3], dtype=torch.int32)), + [WeightVersionSpan(version="v0", start=0, end=2)], + ) + + def test_the_last_slot_is_addressable(self): + """The table covers num_slots entries, so index num_slots - 1 is valid.""" + table = _table() + table.record(slot_indices=_slots(15), version="v0") + + self.assertEqual( + table._lookup_spans(_slots(15)), + [WeightVersionSpan(version="v0", start=0, end=1)], + ) + + def test_version_ids_are_dense_and_reused(self): + """Every new version gets the next id and already seen ones are reused.""" + table = _table() + for index, version in enumerate(["v0", "v1", "v2", "v1", "v0"]): + table.record(slot_indices=_slots(index), version=version) + + self.assertEqual(table._slot_version_ids[:5].tolist(), [0, 1, 2, 1, 0]) + + +class TestStringInterner(CustomTestCase): + def test_first_value_gets_id_zero(self): + """A fresh interner starts numbering at zero.""" + self.assertEqual(_StringInterner().intern("v0"), 0) + + def test_new_values_get_consecutive_ids_in_first_seen_order(self): + """Ids are dense and follow the order in which values were first interned.""" + interner = _StringInterner() + + self.assertEqual([interner.intern(v) for v in ["b", "a", "c"]], [0, 1, 2]) + + def test_a_seen_value_reuses_its_id(self): + """Interning a value again returns the id it already has instead of a new one.""" + interner = _StringInterner() + first = interner.intern("v0") + interner.intern("v1") + + self.assertEqual(interner.intern("v0"), first) + self.assertEqual(interner.intern("v2"), 2) + + def test_lookup_returns_the_interned_value(self): + """lookup is the inverse of intern for every id handed out.""" + interner = _StringInterner() + values = ["v0", "v1", "v2"] + ids = [interner.intern(v) for v in values] + + self.assertEqual([interner.lookup(i) for i in ids], values) + + def test_lookup_of_an_unknown_id_fails(self): + """An id that was never handed out cannot be resolved.""" + interner = _StringInterner() + interner.intern("v0") + + with self.assertRaises(IndexError): + interner.lookup(1) + + def test_distinct_strings_never_share_an_id(self): + """Values that differ only slightly are still distinct entries.""" + interner = _StringInterner() + + ids = {interner.intern(v) for v in ["1", "01", "1 ", "default", ""]} + + self.assertEqual(len(ids), 5) + + +class TestFillReqPrefillWeightVersions(CustomTestCase): + def test_prompt_slots_resolve_to_the_versions_that_computed_them(self): + """Only the prompt's own slots are read, even when more KV is committed than there are prompt tokens.""" + table = _table(slots_of_req=[4, 5, 6, 7, 8]) + table.record(slot_indices=_slots(4, 5), version="v0") + table.record(slot_indices=_slots(6), version="v1") + req = _ReqStub(num_prompt_tokens=3, kv_committed_len=5) + + table.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [ + WeightVersionSpan(version="v0", start=0, end=2), + WeightVersionSpan(version="v1", start=2, end=3), + ], + ) + + def test_a_partially_committed_prompt_fails_the_fill(self): + """Filling before the whole prompt is committed is a scheduling bug and must not yield short spans.""" + table = _table(slots_of_req=[4, 5]) + table.record(slot_indices=_slots(4, 5), version="v0") + req = _ReqStub(num_prompt_tokens=5, kv_committed_len=2) + + with self.assertRaisesRegex( + AssertionError, "2 committed KV tokens for a 5-token prompt" + ): + table.fill_req_prefill_weight_versions(req) + + def test_an_empty_prompt_yields_no_spans(self): + """A request without prompt tokens reports an empty span list, not an error.""" + table = _table(slots_of_req=[]) + req = _ReqStub(num_prompt_tokens=0, kv_committed_len=0) + + table.fill_req_prefill_weight_versions(req) + + self.assertEqual(req.prefill_weight_versions, []) + + def test_an_unstamped_prompt_slot_fails_the_fill(self): + """A prompt slot the tracker never saw surfaces as an error on the request path too.""" + table = _table(slots_of_req=[4, 5, 6]) + table.record(slot_indices=_slots(4, 6), version="v0") + req = _ReqStub(num_prompt_tokens=3, kv_committed_len=3) + + with self.assertRaisesRegex(ValueError, r"\[5\]"): + table.fill_req_prefill_weight_versions(req) + + def test_a_second_fill_replaces_the_previous_spans(self): + """Filling again after a re-prefill replaces the stale spans instead of appending.""" + table = _table(slots_of_req=[4, 5]) + table.record(slot_indices=_slots(4, 5), version="v0") + req = _ReqStub(num_prompt_tokens=2, kv_committed_len=2) + table.fill_req_prefill_weight_versions(req) + table.record(slot_indices=_slots(4, 5), version="v1") + + table.fill_req_prefill_weight_versions(req) + + self.assertEqual( + req.prefill_weight_versions, + [WeightVersionSpan(version="v1", start=0, end=2)], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py b/test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py new file mode 100644 index 000000000000..8caa2294bc57 --- /dev/null +++ b/test/registered/unit/model_executor/test_model_runner_kv_weight_versions.py @@ -0,0 +1,103 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.mem_cache.kv_weight_version_tracker import KvWeightVersionRecord +from sglang.srt.model_executor.model_runner import ModelRunner, ModelRunnerOutput +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestModelRunnerKvWeightVersions(CustomTestCase): + def test_only_enabled_target_forwards_capture_slots(self) -> None: + """Draft forwards and disabled tracking never publish target provenance.""" + for enabled, draft, has_slots in ( + (False, False, True), + (True, True, True), + (True, False, False), + (True, False, True), + ): + with self.subTest(enabled=enabled, draft=draft, has_slots=has_slots): + runner = ModelRunner.__new__(ModelRunner) + runner.is_draft_worker = draft + runner.server_args = SimpleNamespace( + enable_prefill_weight_versions=enabled + ) + batch = SimpleNamespace( + out_cache_loc=torch.tensor([4, 5]) if has_slots else None + ) + with patch( + "sglang.srt.mem_cache.kv_weight_version_tracker.get_serving", + return_value=SimpleNamespace(weight_version="v0"), + ): + record = KvWeightVersionRecord.maybe_capture( + model_runner=runner, forward_batch=batch + ) + + if enabled and not draft and has_slots: + batch.out_cache_loc.fill_(9) + self.assertEqual(record.slot_indices.tolist(), [4, 5]) + self.assertEqual(record.version, "v0") + else: + self.assertIsNone(record) + + def test_target_worker_preserves_forward_provenance(self) -> None: + """TP result construction retains the target forward record.""" + record = KvWeightVersionRecord.capture( + slot_indices=torch.tensor([4, 5]), version="v0" + ) + output = ModelRunnerOutput( + logits_output=None, + can_run_graph=False, + kv_weight_version_record=record, + ) + worker = SimpleNamespace( + is_dllm=lambda: False, + pp_group=SimpleNamespace(is_last_rank=True), + model_runner=SimpleNamespace(forward=lambda *args, **kwargs: output), + ) + batch = SimpleNamespace( + apply_deprecated_skip_attn_backend_init=lambda value: None + ) + + result = TpModelWorker.forward_batch_generation( + worker, batch=None, forward_batch=batch, is_verify=True + ) + + self.assertIs(result.kv_weight_version_record, record) + + def test_result_copies_record_before_publishing_completion(self) -> None: + """The shared result-copy path includes KV slots before its completion event.""" + record = KvWeightVersionRecord.capture( + slot_indices=torch.tensor([4, 5]), version="v0" + ) + source = record.slot_indices + completed: list[bool] = [] + + def publish_completion() -> None: + self.assertIsNot(record.slot_indices, source) + completed.append(True) + + result = GenerationBatchResult( + next_token_ids=torch.tensor([1]), + kv_weight_version_record=record, + copy_done=SimpleNamespace(record=publish_completion), + ) + with patch( + "sglang.srt.managers.utils._async_d2h", + side_effect=lambda tensor: tensor.clone(), + ): + result.copy_to_cpu(return_logprob=False, return_hidden_states=False) + + self.assertEqual(record.slot_indices.tolist(), [4, 5]) + self.assertEqual(completed, [True]) + + +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 8aff0ab199c2..e379b04869bb 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2244,6 +2244,32 @@ def test_hicache_kernel_keeps_implicit_fa3_decode_backend(self): self.assertEqual(resolution_result(args, "hicache_mem_layout"), "page_first") self.assertIsNone(resolution_result(args, "decode_attention_backend")) + def test_prefill_weight_versions_rejects_hierarchical_cache(self): + """KV restored from a host or storage tier carries no weight version, so the two cannot combine.""" + args = self._make_args( + enable_prefill_weight_versions=True, + enable_hierarchical_cache=True, + ) + + with self.assertRaisesRegex(ValueError, "enable-hierarchical-cache"): + handle_cache_compatibility(args) + + def test_prefill_weight_versions_rejects_disaggregation(self): + """PD prefill workers do not stamp prompt tokens yet, so the flag is refused there.""" + args = self._make_args( + enable_prefill_weight_versions=True, + disaggregation_mode="prefill", + ) + + with self.assertRaisesRegex(ValueError, "disaggregation"): + handle_cache_compatibility(args) + + def test_prefill_weight_versions_alone_is_accepted(self): + """A plain engine may report prompt weight versions.""" + args = self._make_args(enable_prefill_weight_versions=True) + + handle_cache_compatibility(args) + def test_decode_offload_rejects_host_pool_retraction(self): args = self._make_args( disaggregation_mode="decode", diff --git a/test/registered/unit/utils/test_weight_versions.py b/test/registered/unit/utils/test_weight_versions.py index 7c2508b4b799..556407028b70 100644 --- a/test/registered/unit/utils/test_weight_versions.py +++ b/test/registered/unit/utils/test_weight_versions.py @@ -13,6 +13,7 @@ compute_weight_version_spans, record_weight_version_events, truncate_weight_version_events, + weight_version_spans_to_json, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import ( @@ -28,6 +29,7 @@ class _ReqStub: def __init__(self, output_len: int): self.output_ids = [0] * output_len self.weight_version_events = [] + self.prefill_weight_versions = None def record_weight_version_change(self, old_version): record_weight_version_events([self], old_version=old_version) @@ -540,6 +542,50 @@ def test_metadata_omits_spans_when_absent(self): metadata = build_endpoint_weight_version_metadata({"weight_version": "v2"}) self.assertEqual(metadata, {"weight_version": "v2"}) + def test_metadata_carries_prefill_spans_when_present(self): + """prefill_weight_versions rides along into endpoint metadata next to the output spans.""" + prefill_spans = [{"version": "v1", "start": 0, "end": 5}] + metadata = build_endpoint_weight_version_metadata( + { + "weight_version": "v2", + "weight_versions": [{"version": "v2", "start": 0, "end": 3}], + "prefill_weight_versions": prefill_spans, + "prompt_tokens": 5, + } + ) + self.assertEqual(metadata["prefill_weight_versions"], prefill_spans) + self.assertNotIn("prompt_tokens", metadata) + + +class TestWeightVersionSpansToJson(CustomTestCase): + def test_spans_pass_through_unclamped_without_a_limit(self): + """Prefill spans keep their exact ends because no output-token limit applies.""" + spans = [ + WeightVersionSpan(version="v1", start=0, end=5), + WeightVersionSpan(version="v2", start=5, end=9), + ] + self.assertEqual( + weight_version_spans_to_json(spans), + [ + {"version": "v1", "start": 0, "end": 5}, + {"version": "v2", "start": 5, "end": 9}, + ], + ) + + def test_ends_are_clamped_to_the_limit(self): + """With a limit every end is capped, matching the output-span clamp.""" + spans = [ + WeightVersionSpan(version="v1", start=0, end=5), + WeightVersionSpan(version="v2", start=5, end=9), + ] + self.assertEqual( + weight_version_spans_to_json(spans, end_limit=7), + [ + {"version": "v1", "start": 0, "end": 5}, + {"version": "v2", "start": 5, "end": 7}, + ], + ) + class TestAddWeightVersionsToMetaInfo(CustomTestCase): def test_meta_info_gets_dict_spans_and_legacy_field(self):