sgl-project/sglang · error · RuntimeError

FlashInfer KDA verify requires an identity intermediate row-

Error message

FlashInfer KDA verify requires an identity intermediate row-map (verify_intermediate_state_indices must be arange).

What it means

The FlashInfer verify kernel's fast seed-copy assumes row n of the scratch buffer belongs to request n; when no explicit ssm_state_indices are provided it requires verify_intermediate_state_indices to equal arange(batch_size). A permuted row-map breaks the assumption, so it fails fast.

Source

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

                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):
                raise RuntimeError(
                    "FlashInfer KDA verify requires an identity intermediate row-map "
                    "(verify_intermediate_state_indices must be arange)."
                )
            step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32)
            ssm_state_indices = (
                base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :]
            ).contiguous()  # [N, T]
            self._verify_idx_cache[cache_key] = ssm_state_indices

        # Seed step 0 from committed state, then recurrent_kda overwrites it with
        # token-0 post-state. Padded graph rows clamp to slot 0; their output is ignored.
        base_state = ssm_states.index_select(
            0, cache_indices[:batch_size].clamp(min=0).to(torch.int64)
        )
        scratch[:batch_size, 0].copy_(base_state)

        # Same storage as scratch, flattened over the allocated step stride.
        state_pool = scratch.view(

View on GitHub (pinned to 0132848349)

Solutions

Ensure verify_intermediate_state_indices is the identity arange for the verify batch
Or pass explicit ssm_state_indices covering every (row, step) so the kernel doesn't need the assumption
Avoid custom remapping of the intermediate SSM buffer between decode and verify

Example fix

# before
verify_intermediate_state_indices = permuted_rows  # custom remap
# after
verify_intermediate_state_indices = torch.arange(batch_size, dtype=torch.int32)

Defensive patterns

Strategy: validation

Validate before calling

expected = torch.arange(batch_size, device=idx.device, dtype=idx.dtype)
if ssm_state_indices is None:
    assert torch.equal(verify_intermediate_state_indices, expected), 'row-map must be identity arange'

Type guard

def is_identity_row_map(idx: torch.Tensor, batch_size: int) -> bool:
    return idx.numel() == batch_size and bool((idx == torch.arange(batch_size, device=idx.device, dtype=idx.dtype)).all())

Prevention

Pass explicit ssm_state_indices when rows may be remapped
Never permute intermediate state rows between decode and verify

When it happens

Trigger: Calling target_verify with ssm_state_indices=None while intermediate_state_indices[:batch_size] is not arange(batch_size) — e.g. a remapped or compacted scratch pool after request eviction/reordering.

Common situations: Custom radix-cache or scheduling logic that reuses/permutes intermediate state rows; partial-batch reordering that reassigns scratch rows non-identically.

Related errors

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