[Bugfix][Kimi-K3] Fix two crashes that make --use-replayssm unusable

[mispa-ms]

Purpose

--use-replayssm selects the RecoverSSM path on Kimi-K3. It cannot start: it
raises during cudagraph capture, and after that is fixed it raises again on the
first execute_model.

1. GDNAttentionMetadataBuilder never calls its base __init__.

AttributeError: 'KimiK3KDAMetadataBuilder' object has no attribute 'layer_names'
  kda_metadata.py, in _get_recoverssm_context
    layers = [forward_context[layer_name] for layer_name in self.layer_names]

AttentionMetadataBuilder.__init__ stores kv_cache_spec, layer_names,
vllm_config and device. This builder calls none of it and assigns four
fields by hand, leaving out layer_names and device. layer_names has one
reader, _get_recoverssm_context; device has none yet. Calling super()
restores both and matches flash_attn, mamba_attn and triton_attn.

2. The checkpoint block's guard does not match the one that binds its
variable.

UnboundLocalError: cannot access local variable 'active_non_spec_mask_cpu'
  kda_metadata.py, in build
    request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist()

active_non_spec_mask_cpu is bound only under num_spec_decodes > 0, which is
what the sibling has_initial_state use guards on. The checkpoint block guards
on spec_sequence_masks_cpu is not None.

The reachable case is an all-False mask, not a zero draft sum.
num_decode_draft_tokens is filled with -1 and written only on spec-decode
rows, so on a prefill-only step num_decode_draft_tokens_cpu >= 0 is entirely
False. Without RecoverSSM the classification clears that mask to None; the
clearing carries not self.use_recoverssm, so under RecoverSSM a non-None
all-False mask survives with num_spec_decodes == 0. This is therefore every
prefill step under align with checkpoint blocks, not a rare corner, which is
why it fires on the first execute_model.

The root cause is a dropped normalization clause. This file states it
"intentionally mirror[s] GDNAttentionMetadataBuilder", and the shared builder
keeps spec_sequence_masks_cpu is not None equivalent to num_spec_decodes > 0
at the source:

# vllm/v1/attention/backends/gdn_attn.py
spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
num_spec_decodes = spec_sequence_masks_cpu.sum().item()
if (
    num_spec_decodes == 0
    or num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
    == 0
):
    num_spec_decodes = 0
    spec_sequence_masks = None
    spec_sequence_masks_cpu = None

That invariant is why guarding on the mask is safe there. The Kimi copy dropped
the num_spec_decodes == 0 disjunct and moved the count into the else, and
adding not self.use_recoverssm then made the gap reachable. Restoring the
clause is the other possible fix; it is left as a follow-up because it would
make the mask None again and so stop the regression test below from
distinguishing the two guards.

num_spec_decodes > 0 is also the correct guard, not just the bound one:
checkpoint_offsets is consumed by _store_cache_checkpoints_kernel on a grid
of its own numel(), with seq_idx indexing non_spec_query_start_loc. At
num_spec_decodes == 0 that tensor is the full query_start_loc, so
request_rows must be every row — the fallback value. Masking there would have
shortened request_rows against a full-length start-loc array.

Not a duplicate

active_non_spec_mask_cpu returns no results in this repo's issues or PRs. The
RecoverSSM hits are #51855 (closed, added this path), #52506 and #54255
(FlashInfer backends), #54103 (ROCm). None touches either line.

Test

Two regression tests added. The existing RecoverSSM tests set
builder.recoverssm_context = Mock(), so _get_recoverssm_context returns
early and the layer_names read never executes -- which is why CI is green on
a path that cannot start.

Kimi-K3-MXFP4, 16×GB300, aggregated, TP16 × DCP8 × PP1, DSpark
num_speculative_tokens=4, mamba_cache_mode=align, mamba-backend=triton,
VLLM_USE_V2_MODEL_RUNNER=1, 3600 s per point. Three arms one flag apart.

	base	+triton	+--use-replayssm
c64, tok/s/GPU	3,760.3	3,816.9	3,940.8
c64, ITL p90 ms	166.94	170.48	165.28
c128, tok/s/GPU	2,749.9	2,706.1	2,828.5
c128, ITL p90 ms	205.23	204.00	192.51

+3.2% at c64 and +4.5% at c128 over the triton control, in ITL.

Limitations

• The throughput arms use synthetic acceptance with ignore_eos, so a stale
  recurrent state would not show in them. GSM8K with
  rejection_sample_method=standard, 1319 questions, 5 shots, same config one
  flag apart: 0.9530 off / 0.9545 on, invalid 0.08% / 0.15%. Two questions
  on 1319, and in the wrong direction for a staleness bug.
• Kimi-K3 / NVIDIA KDA only. Fix 1 is in shared GDN code; the AMD KDA builder
  has no RecoverSSM or checkpoint path and does not read either field, so it is
  unaffected.
• Single configuration, no repeats, so no error bars.

AI assistance

Written with Claude Code. Both crashes were reproduced on hardware and the
tracebacks are verbatim.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.