[Kimi Bug] Fix kimi k3 startup cuda graph issue with recoverSSM

[yewentao256]

Purpose

vllm serve moonshotai/Kimi-K3   --trust-remote-code   --tensor-parallel-size 8   --load-format fastsafetensors   --gpu-memory-utilization 0.9   --enable-prefix-caching   --reasoning-parser kimi_k3   --enable-auto-tool-choice   --tool-call-parser kimi_k3   --host 0.0.0.0   --port 30000   --max-model-len auto   --max-num-seqs 48   --speculative-config '{"model":"RedHatAI/Kimi-K3-speculator.dspark","method":"dspark","num_speculative_tokens":8,"draft_sample_method":"probabilistic","rejection_sample_method":"standard"}'   --use-replayssm 

Will raise error

(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]   File "/home/yewentao256/vllm-source/vllm/v1/worker/gpu/model_states/mamba_hybrid.py", line 313, in prepare_attn
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]     attn_metadata = build_attn_metadata(
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]                     ^^^^^^^^^^^^^^^^^^^^
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]   File "/home/yewentao256/vllm-source/vllm/v1/worker/gpu/attn_utils.py", line 343, in build_attn_metadata
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]     metadata = attn_metadata_builder.build(
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]   File "/home/yewentao256/vllm-source/vllm/models/kimi_k3/nvidia/kda_metadata.py", line 442, in build
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055]     raise ValueError(
(Worker_TP5 pid=1272371) ERROR 09-07 17:14:52 [multiproc_executor.py:1055] ValueError: KDA RecoverSSM speculative decode query length exceeds its activation capacity (9)

This PR fixes the issue by skipping the cudagraph capture outside of workspace

Test

(APIServer pid=2088238) INFO:     Started server process [2088238]
(APIServer pid=2088238) INFO:     Waiting for application startup.
(APIServer pid=2088238) INFO:     Application startup complete.

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87593 for commit 76e5267433ac.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #87689 for commit 69998fce9e74.