sgl-project/sglang · error · RuntimeError

f"KDA verify needs {draft_token_num} scratch steps, but inte

Error message

f"KDA verify needs {draft_token_num} scratch steps, but intermediate_ssm only has {scratch_steps}."

What it means

During target_verify, each of the draft_token_num speculative steps needs one scratch row in the intermediate SSM buffer; if the buffer's step dimension T is smaller than draft_token_num, the kernel cannot index the flat state pool and raises RuntimeError.

Source

Thrown at python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py:270

        num_v_heads = v.shape[2]
        head_v_dim = v.shape[3]

        # Packed [1, N*T, ...] inputs, cu_seqlens = query_start_loc (draft stride).
        # recurrent_kda is bf16-only (see decode), so cast every input to bf16.
        q_fi = q.reshape(1, seq_len, num_heads, head_k_dim).to(torch.bfloat16)
        k_fi = k.reshape(1, seq_len, num_heads, head_k_dim).to(torch.bfloat16)
        v_fi = v.reshape(1, seq_len, num_v_heads, head_v_dim).to(torch.bfloat16)
        g_fi = a.reshape(1, seq_len, num_v_heads, head_k_dim).to(torch.bfloat16)
        beta_fi = self._beta_logit_to_prob(b).reshape(1, seq_len, num_v_heads)

        A_log_fi, dt_bias_fi = self._prep_gate_params(A_log, dt_bias)

        # recurrent_kda indexes a flat state pool. Map each request/step to the
        # matching slot in SGLang's [scratch_row, allocated_step, HV, V, K] buffer.
        scratch = intermediate_states_buffer  # [N_scratch, T, HV, V, K]
        scratch_steps = scratch.shape[1]
        if draft_token_num > scratch_steps:
            raise RuntimeError(
                f"KDA verify needs {draft_token_num} scratch steps, "
                f"but intermediate_ssm only has {scratch_steps}."
            )

        base_rows = intermediate_state_indices[:batch_size]
        cache_key = (
            id(intermediate_state_indices),
            batch_size,
            draft_token_num,
            scratch_steps,
        )
        ssm_state_indices = self._verify_idx_cache.get(cache_key)
        if ssm_state_indices is None:
            # The fast seed copy below assumes row n in scratch belongs to request n.
            expected = torch.arange(
                batch_size, device=base_rows.device, dtype=base_rows.dtype
            )
            if not torch.equal(base_rows, expected):

View on GitHub (pinned to 0132848349)

Solutions

  1. Increase the intermediate SSM buffer's step dimension to at least the max draft_token_num (align speculative_num_steps with cache sizing)
  2. Lower --speculative-num-steps so draft_token_num fits the buffer
  3. Re-run with the standard speculative config used when the cache was allocated

Example fix

# before
--speculative-num-steps 8   # buffer T was sized for 4
# after
--speculative-num-steps 4   # or resize intermediate_ssm to T>=8
Defensive patterns

Strategy: validation

Validate before calling

draft_token_num = speculative_num_steps  # per request
assert draft_token_num <= intermediate_ssm.shape[1], (
    f'need {draft_token_num} scratch steps, buffer has {intermediate_ssm.shape[1]}')

Type guard

null

Prevention

When it happens

Trigger: Speculative config produces more draft tokens per request than the intermediate_ssm buffer was allocated for (e.g. speculative_num_steps increased after cache sizing, or draft_token_num > buffer T).

Common situations: Raising --speculative-num-steps / --speculative-eagle-topk without resizing the mamba intermediate buffer; mismatched server args between cache allocation and verify path.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/54919db54a9458bf. Report an issue: GitHub.