[Bugfix][Kimi-K3] Reserve KDA prefill workspace before profiling (#63989)

Signed-off-by: jbernloehr <jbernloehr@nvidia.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Co-authored-by: vllm-agent <claw@inferact.ai>

vllm/models/kimi_k3/nvidia/kda.py

@@ -73,7 +73,10 @@ from vllm.utils.flashinfer import (
 )
 from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.kv_cache_interface import MambaSpec
-from vllm.v1.worker.workspace import current_workspace_manager
+from vllm.v1.worker.workspace import (
+    current_workspace_manager,
+    is_workspace_manager_initialized,
+)
 
 logger = init_logger(__name__)
 
@@ -857,6 +860,18 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                 self.model_config.dtype,
             )
 
+        # Profiling skips attention, so reserve its runtime layout before the
+        # KV cache is sized. Layers share the manager's grow-only allocation.
+        if is_workspace_manager_initialized():
+            if self._flashkda_buffer_specs is not None:
+                current_workspace_manager().get_simultaneous(
+                    *self._flashkda_buffer_specs
+                )
+            elif self._flashinfer_kda_output_spec is not None:
+                current_workspace_manager().get_simultaneous(
+                    self._flashinfer_kda_output_spec
+                )
+
         self.o_norm = FusedRMSNormGated(self.head_dim, activation="sigmoid")
         decode_norm_weight = None
         if decode_conv1d_weight is not None: