{"record":{"id":"81c117705765052d","repo":"sgl-project/sglang","slug":"kda-prefill-requires-an-indexed-initial-state-pool","errorCode":null,"errorMessage":"KDA prefill requires an indexed initial-state pool","messagePattern":"KDA prefill requires an indexed initial-state pool","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_prefill.py","lineNumber":1323,"sourceCode":"    v: torch.Tensor,\n    g: torch.Tensor,\n    beta: torch.Tensor,\n    scale: float | None = None,\n    initial_state: torch.Tensor | None = None,\n    initial_state_indices: torch.Tensor | None = None,\n    use_qk_l2norm_in_kernel: bool = False,\n    cu_seqlens: torch.Tensor | None = None,\n    A_log: torch.Tensor | None = None,\n    dt_bias: torch.Tensor | None = None,\n    lower_bound: float | None = None,\n    output_intermediate_states: bool = False,\n    **kwargs: object,\n) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Match the public forward contract of SGLang's Triton ``chunk_kda``.\"\"\"\n    if scale is None:\n        scale = k.shape[-1] ** -0.5\n    if initial_state is None or initial_state_indices is None:\n        raise ValueError(\"KDA prefill requires an indexed initial-state pool\")\n\n    num_tokens = q.shape[1]\n    if g.shape[1] < num_tokens or beta.shape[1] < num_tokens:\n        raise ValueError(\"g and beta must cover every q token\")\n    g = g[:, :num_tokens]\n    beta = beta[:, :num_tokens]\n    if num_tokens == 1:\n        # Tracing constant-folds size-one dimensions, but the resulting kernel\n        # can share a cache entry with longer inputs. Keep T=1 on Triton so a\n        # short first request cannot specialize later Helion calls incorrectly.\n        return triton_chunk_kda(\n            q=q,\n            k=k,\n            v=v,\n            g=g,\n            beta=beta,\n            scale=scale,\n            initial_state=initial_state,","sourceCodeStart":1305,"sourceCodeEnd":1341,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_prefill.py#L1305-L1341","documentation":"chunk_kda (the Helion implementation matching the Triton chunked KDA prefill contract) requires an explicit initial-state pool plus per-request indices, because prefill must read and update per-sequence SSM states. If either initial_state or initial_state_indices is None it raises immediately.","triggerScenarios":"Calling chunk_kda without initial_state or without initial_state_indices — e.g. a first-prefill path that passes None for the state pool, or an integration that only wires one of the two arguments.","commonSituations":"Porting code from a reference chunk_kda that allowed a None initial state (stateless first chunk); forgetting to thread the mamba-style state pool and its index tensor through a new attention backend; short-cutting tests with initial_state=None.","solutions":["Allocate a zero state pool [num_slots, HV, V, K] and pass initial_state_indices of the request's slot ids","If no prior state exists, pass zeros(len(cu_seqlens)-1, ...) as initial_state plus valid indices — the API has no None shortcut","Wire the state pool from the model's cache manager into the prefill call site"],"exampleFix":"// before\nout = chunk_kda(q, k, v, g, beta, ..., initial_state=None)\n// after\npool = torch.zeros(num_slots, HV, V, K, device=q.device, dtype=q.dtype)\nout = chunk_kda(q, k, v, g, beta, ..., initial_state=pool, initial_state_indices=slot_ids)","handlingStrategy":"validation","validationCode":"if initial_state is None:\n    initial_state = torch.zeros(num_slots, HV, V, K, device=q.device, dtype=q.dtype)\nif initial_state_indices is None:\n    raise ValueError(\"initial_state_indices required\")\nassert initial_state is not None and initial_state_indices is not None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Thread the state pool and slot indices through every prefill call site","Wrap chunk_kda in a helper that guarantees both arguments are set"],"tags":["kda","helion","prefill","state-pool","required-argument"],"backgroundTag":"missing-required-argument","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}