[Kimi Bug] Fix kimi k3 startup cuda graph issue with recoverSSM (#55774)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 8, 2026 at 10:29 UTC
07950d47347f7fff41f98616096fd10349171e3d
2 files changed
+15
-2
tests/v1/cudagraph/test_cudagraph_manager.py
+10
-1
@@ -172,6 +172,8 @@ def _create_decode_vllm_config(
172
capture_sizes: list[int],
173
num_speculative_tokens: int = 0,
174
dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
175
+ max_num_seqs: int = 8,
176
+ use_kda_recoverssm: bool = False,
177
) -> MagicMock:
178
compilation_config = CompilationConfig(
179
cudagraph_mode="FULL_AND_PIECEWISE",
@@ -182,8 +184,11 @@ def _create_decode_vllm_config(
184
185
vllm_config = MagicMock(spec=VllmConfig)
186
vllm_config.compilation_config = compilation_config
185
- vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=8)
187
+ vllm_config.scheduler_config = SchedulerConfig.default_factory(
188
+ max_num_seqs=max_num_seqs
189
+ )
190
vllm_config.parallel_config = ParallelConfig()
191
+ vllm_config.cache_config = SimpleNamespace(use_kda_recoverssm=use_kda_recoverssm)
192
vllm_config.num_speculative_tokens = num_speculative_tokens
193
if dynamic_spec_schedule is None:
194
vllm_config.speculative_config = None
@@ -201,6 +206,8 @@ def _make_spec_decode_manager(
206
capture_sizes: list[int] | None = None,
207
num_speculative_tokens: int = 0,
208
dynamic_spec_schedule: list[tuple[int, int, int]] | None = None,
209
+ max_num_seqs: int = 8,
210
+ use_kda_recoverssm: bool = False,
211
) -> gpu_cudagraph_utils.CudaGraphManager:
212
monkeypatch.setattr(
213
gpu_cudagraph_utils,
@@ -217,6 +224,8 @@ def _make_spec_decode_manager(
224
capture_sizes or [1, 2, 4, 8, 16, 24],
225
num_speculative_tokens=num_speculative_tokens,
226
dynamic_spec_schedule=dynamic_spec_schedule,
227
+ max_num_seqs=max_num_seqs,
228
+ use_kda_recoverssm=use_kda_recoverssm,
229
),
230
device=torch.device("cpu"),
231
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
vllm/v1/worker/gpu/cudagraph_utils.py
+5
-1
@@ -299,7 +299,11 @@ class CudaGraphManager:
299
if desc not in descs_by_mode[decode_mode]:
300
descs_by_mode[decode_mode].append(desc)
301
302
- if mixed_mode:
302
+ # recoverSSM cannot capture a dummy query wider than its workspace.
303
+ if mixed_mode and (
304
+ not self.vllm_config.cache_config.use_kda_recoverssm
305
+ or num_tokens <= max_decode_tokens
306
+ ):
307
# for PIECEWISE graphs there is no limit on requests when replaying
308
# i.e. no request padding is needed, so we leave it as None.
309
# For breakable PW graphs, break-point kernels read the real batch