{"record":{"id":"b7767f739460ac43","repo":"sgl-project/sglang","slug":"flashinfer-kda-verify-requires-an-identity-interme","errorCode":null,"errorMessage":"FlashInfer KDA verify requires an identity intermediate row-map (verify_intermediate_state_indices must be arange).","messagePattern":"FlashInfer KDA verify requires an identity intermediate row-map \\(verify_intermediate_state_indices must be arange\\)\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py","lineNumber":289,"sourceCode":"                f\"KDA verify needs {draft_token_num} scratch steps, \"\n                f\"but intermediate_ssm only has {scratch_steps}.\"\n            )\n\n        base_rows = intermediate_state_indices[:batch_size]\n        cache_key = (\n            id(intermediate_state_indices),\n            batch_size,\n            draft_token_num,\n            scratch_steps,\n        )\n        ssm_state_indices = self._verify_idx_cache.get(cache_key)\n        if ssm_state_indices is None:\n            # The fast seed copy below assumes row n in scratch belongs to request n.\n            expected = torch.arange(\n                batch_size, device=base_rows.device, dtype=base_rows.dtype\n            )\n            if not torch.equal(base_rows, expected):\n                raise RuntimeError(\n                    \"FlashInfer KDA verify requires an identity intermediate row-map \"\n                    \"(verify_intermediate_state_indices must be arange).\"\n                )\n            step = torch.arange(draft_token_num, device=q.device, dtype=torch.int32)\n            ssm_state_indices = (\n                base_rows.to(torch.int32)[:, None] * scratch_steps + step[None, :]\n            ).contiguous()  # [N, T]\n            self._verify_idx_cache[cache_key] = ssm_state_indices\n\n        # Seed step 0 from committed state, then recurrent_kda overwrites it with\n        # token-0 post-state. Padded graph rows clamp to slot 0; their output is ignored.\n        base_state = ssm_states.index_select(\n            0, cache_indices[:batch_size].clamp(min=0).to(torch.int64)\n        )\n        scratch[:batch_size, 0].copy_(base_state)\n\n        # Same storage as scratch, flattened over the allocated step stride.\n        state_pool = scratch.view(","sourceCodeStart":271,"sourceCodeEnd":307,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py#L271-L307","documentation":"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.","triggerScenarios":"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.","commonSituations":"Custom radix-cache or scheduling logic that reuses/permutes intermediate state rows; partial-batch reordering that reassigns scratch rows non-identically.","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"],"exampleFix":"# before\nverify_intermediate_state_indices = permuted_rows  # custom remap\n# after\nverify_intermediate_state_indices = torch.arange(batch_size, dtype=torch.int32)","handlingStrategy":"validation","validationCode":"expected = torch.arange(batch_size, device=idx.device, dtype=idx.dtype)\nif ssm_state_indices is None:\n    assert torch.equal(verify_intermediate_state_indices, expected), 'row-map must be identity arange'","typeGuard":"def is_identity_row_map(idx: torch.Tensor, batch_size: int) -> bool:\n    return idx.numel() == batch_size and bool((idx == torch.arange(batch_size, device=idx.device, dtype=idx.dtype)).all())","tryCatchPattern":null,"preventionTips":["Pass explicit ssm_state_indices when rows may be remapped","Never permute intermediate state rows between decode and verify"],"tags":["sglang","kda","speculative-decoding","index-mapping","invariant"],"backgroundTag":"internal-invariant-violation","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}