Report prompt KV weight versions in generation metadata

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
 

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

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 "

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(

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] = []

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,

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

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
             ),

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
 

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,
     )

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,

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,
         )

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(

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)

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: