sgl-project/sglang · error · ValueError

KDA prefill requires an indexed initial-state pool

Error message

KDA prefill requires an indexed initial-state pool

What it means

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.

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_prefill.py:1323

    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float | None = None,
    initial_state: torch.Tensor | None = None,
    initial_state_indices: torch.Tensor | None = None,
    use_qk_l2norm_in_kernel: bool = False,
    cu_seqlens: torch.Tensor | None = None,
    A_log: torch.Tensor | None = None,
    dt_bias: torch.Tensor | None = None,
    lower_bound: float | None = None,
    output_intermediate_states: bool = False,
    **kwargs: object,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Match the public forward contract of SGLang's Triton ``chunk_kda``."""
    if scale is None:
        scale = k.shape[-1] ** -0.5
    if initial_state is None or initial_state_indices is None:
        raise ValueError("KDA prefill requires an indexed initial-state pool")

    num_tokens = q.shape[1]
    if g.shape[1] < num_tokens or beta.shape[1] < num_tokens:
        raise ValueError("g and beta must cover every q token")
    g = g[:, :num_tokens]
    beta = beta[:, :num_tokens]
    if num_tokens == 1:
        # Tracing constant-folds size-one dimensions, but the resulting kernel
        # can share a cache entry with longer inputs. Keep T=1 on Triton so a
        # short first request cannot specialize later Helion calls incorrectly.
        return triton_chunk_kda(
            q=q,
            k=k,
            v=v,
            g=g,
            beta=beta,
            scale=scale,
            initial_state=initial_state,

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate a zero state pool [num_slots, HV, V, K] and pass initial_state_indices of the request's slot ids
2. If no prior state exists, pass zeros(len(cu_seqlens)-1, ...) as initial_state plus valid indices — the API has no None shortcut
3. Wire the state pool from the model's cache manager into the prefill call site

Example fix

// before
out = chunk_kda(q, k, v, g, beta, ..., initial_state=None)
// after
pool = torch.zeros(num_slots, HV, V, K, device=q.device, dtype=q.dtype)
out = chunk_kda(q, k, v, g, beta, ..., initial_state=pool, initial_state_indices=slot_ids)

if initial_state is None:
    initial_state = torch.zeros(num_slots, HV, V, K, device=q.device, dtype=q.dtype)
if initial_state_indices is None:
    raise ValueError("initial_state_indices required")
assert initial_state is not None and initial_state_indices is not None

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

• g and beta must cover every q token
• `mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim}).
• `mixed_qkv` must be contiguous in the last dim.
• `a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim=
• `a`/`b` must be contiguous in the last dim.

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