sgl-project/sglang · error · ValueError

g and beta must cover every q token

Error message

g and beta must cover every q token

What it means

chunk_kda uses q.shape[1] as the token count and requires the decay gates g and beta to have at least that many tokens along dim 1 (they are then trimmed with g[:, :num_tokens]). The error fires when g or beta is shorter than q, i.e. the gate projections do not cover all query tokens.

Source

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

    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,
            initial_state_indices=initial_state_indices,
            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
            cu_seqlens=cu_seqlens,
            A_log=A_log,

View on GitHub (pinned to 0132848349)

Solutions

  1. Recompute or re-slice g and beta so g.shape[1] >= q.shape[1] and beta.shape[1] >= q.shape[1]
  2. Check the upstream projection ran on the same token set as q/k/v (same cu_seqlens)
  3. In chunked prefill, make sure gate projections are re-run for the full extended chunk, not cached from a shorter one

Example fix

// before
g = g[:, :prev_len]  # stale chunk length
chunk_kda(q, k, v, g, beta, ...)
// after
g = compute_gates(x_full)  # covers all q tokens
chunk_kda(q, k, v, g, beta, ...)
Defensive patterns

Strategy: validation

Validate before calling

num_tokens = q.shape[1]
assert g.shape[1] >= num_tokens and beta.shape[1] >= num_tokens, (
    g.shape, beta.shape, q.shape)
g, beta = g[:, :num_tokens], beta[:, :num_tokens]

Type guard

def gates_cover_q(q, g, beta) -> bool:
    return g.shape[1] >= q.shape[1] and beta.shape[1] >= q.shape[1]

Prevention

When it happens

Trigger: Calling chunk_kda with g or beta computed over fewer tokens than q — e.g. gates sliced to a previous chunk length, a truncated gate projection, or varlen tensors packed with a different total than q.

Common situations: Chunked prefill where the gate projection lags one chunk behind q; scheduler merging requests but recomputing gates only for part of them; off-by-one slicing g[:, :T-1] when preparing inputs.

Related errors


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