Fix KV-canary workspace accounting after graph capture (#38596)

python/sglang/kernels/ops/kv_canary/verify.py

@@ -197,6 +197,10 @@ class VerifyPlan:
     verify_num_valid: torch.Tensor
     enable: torch.Tensor
 
+    @staticmethod
+    def allocation_bytes(verify_capacity: int) -> int:
+        return 4 * verify_capacity * torch.int64.itemsize + 2 * torch.int32.itemsize
+
     @classmethod
     def allocate(cls, *, verify_capacity: int, device: torch.device) -> VerifyPlan:
         if verify_capacity <= 0:

python/sglang/kernels/ops/kv_canary/write.py

@@ -46,6 +46,12 @@ class WritePlan:
     write_seed_slot_indices: torch.Tensor
     write_num_valid_reqs: torch.Tensor
 
+    @staticmethod
+    def allocation_bytes(write_req_capacity: int) -> int:
+        return (
+            2 * write_req_capacity + 1
+        ) * torch.int64.itemsize + torch.int32.itemsize
+
     @classmethod
     def allocate(
         cls,

python/sglang/srt/kv_canary/capacities.py

@@ -3,6 +3,10 @@ from __future__ import annotations
 import math
 from dataclasses import dataclass
 
+from sglang.kernels.ops.kv_canary.verify import VerifyPlan
+from sglang.kernels.ops.kv_canary.write import WritePlan
+from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
+from sglang.srt.kv_canary.plan_input import PlanInput
 from sglang.srt.runtime_context import (
     get_exec,
     get_schedule,
@@ -116,3 +120,14 @@ class CanaryLaunchCapacities:
             per_forward_write_req_capacity=max_bs,
             per_forward_write_entry_capacity=write_entry_capacity,
         )
+
+    def per_forward_workspace_bytes(self, *, num_buffer_groups: int) -> int:
+        return (
+            num_buffer_groups
+            * (
+                VerifyPlan.allocation_bytes(self.per_forward_verify_capacity)
+                + WritePlan.allocation_bytes(self.per_forward_write_req_capacity)
+            )
+            + ExpectedInputs.allocation_bytes(self.per_forward_write_entry_capacity)
+            + PlanInput.allocation_bytes(self.per_forward_write_req_capacity)
+        )

python/sglang/srt/kv_canary/expected_inputs.py

@@ -10,6 +10,10 @@ class ExpectedInputs:
     tokens: torch.Tensor
     positions: torch.Tensor
 
+    @staticmethod
+    def allocation_bytes(capacity: int) -> int:
+        return 2 * capacity * torch.int64.itemsize
+
     @classmethod
     def allocate(cls, *, capacity: int, device: torch.device) -> ExpectedInputs:
         return cls(

python/sglang/srt/kv_canary/plan_input.py

@@ -41,6 +41,10 @@ class PlanInput:
     extend_seq_lens: torch.Tensor
     req_to_verify_expected_tokens_valid_lens: torch.Tensor
 
+    @staticmethod
+    def allocation_bytes(bs_capacity: int) -> int:
+        return 4 * bs_capacity * torch.int64.itemsize
+
     def zero_(self) -> None:
         self.req_pool_indices.zero_()
         self.prefix_lens.zero_()

python/sglang/srt/kv_canary/runner/canary_manager.py

@@ -64,6 +64,7 @@ class CanaryManager:
         self._model_forward_bracket_depth: int = 0
 
         self._buffer_groups: tuple[CanaryBufferGroup, ...] = tuple(buffer_groups)
+        self._launch_capacities = launch_capacities
 
         self._device_state = CanaryDeviceState.allocate(
             config=config,
@@ -167,6 +168,11 @@ class CanaryManager:
             for _ in range(num_sfms)
         )
 
+    def per_forward_workspace_bytes(self) -> int:
+        return self._launch_capacities.per_forward_workspace_bytes(
+            num_buffer_groups=len(self._buffer_groups)
+        )
+
     @contextlib.contextmanager
     def with_active_single_forward_manager(self, index: int) -> Iterator[None]:
         assert self._active_single_forward_manager_index is None, (

python/sglang/srt/managers/scheduler.py

@@ -1115,7 +1115,13 @@ class Scheduler(
                 self.draft_worker.prewarm_sampling()
         if model_runner.token_to_kv_pool.post_capture_active:
             tic = time.perf_counter()
-            model_runner.post_capture_resize_kv_pool()
+            model_runner.post_capture_resize_kv_pool(
+                draft_runners=(
+                    self.draft_worker._draft_model_runners()
+                    if self.draft_worker is not None
+                    else ()
+                )
+            )
             self.kv_cache_allocation_time += time.perf_counter() - tic
 
         if get_model().is_startup_weight_load_overlap:

python/sglang/srt/model_executor/model_runner.py

@@ -966,8 +966,8 @@ class ModelRunner:
             ),
         )
 
-    def post_capture_resize_kv_pool(self):
-        resize = compute_post_capture_kv_resize(self)
+    def post_capture_resize_kv_pool(self, *, draft_runners=()):
+        resize = compute_post_capture_kv_resize(self, draft_runners=draft_runners)
         self.max_total_num_tokens = resize.max_total_num_tokens
         if self.is_hybrid_swa:
             self.full_max_total_num_tokens = resize.full_max_total_num_tokens

python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py

@@ -48,6 +48,8 @@ class PostCaptureKVResize(msgspec.Struct, frozen=True, kw_only=True):
 
 def compute_post_capture_kv_resize(
     model_runner: ModelRunner,
+    *,
+    draft_runners: tuple[ModelRunner, ...] = (),
 ) -> PostCaptureKVResize:
     """Resize the KV pool after capture and return the new sizes for the
     orchestrator to assign. Takes the live ModelRunner because it reads
@@ -97,9 +99,19 @@ def compute_post_capture_kv_resize(
         is_multimodal=model_runner.model_config.is_multimodal,
         mm_feature_transport=get_mm().mm_feature_transport,
     )
+    # Sequential target/draft forwards reuse workspace at unchanged capacities.
+    canary_workspace_bytes = max(
+        (
+            runner.canary_manager.per_forward_workspace_bytes()
+            for runner in (model_runner, *draft_runners)
+            if runner.canary_manager is not None
+        ),
+        default=0,
+    )
     budget_bytes = (
         int(max(0.0, free_gb - headroom_gb - mm_reservation_gb) * (1 << 30))
         + pool.post_capture_backed_bytes
+        - canary_workspace_bytes
     )
     config = model_runner.kv_cache_configurator.config_from_budget(
         budget_bytes, cap_tokens=model_runner.max_total_num_tokens
@@ -107,6 +119,11 @@ def compute_post_capture_kv_resize(
     pool.finalize_backing(config)
     model_runner.token_to_kv_pool_allocator.resize(config)
     model_runner.req_to_token_pool.reset_aux_cache_allocator()
+    if canary_workspace_bytes:
+        logger.info(
+            "Post-capture KV sizing: KV-canary per-forward workspace %.2f GB",
+            canary_workspace_bytes / (1 << 30),
+        )
 
     capped_max_running_requests = None
     if model_runner.max_running_requests is not None: