sgl-project/sglang · error · ValueError

varlen KDA requires batch size 1

Error message

varlen KDA requires batch size 1

What it means

In varlen mode (cu_seqlens provided), the Helion KDA prefill chunked gate/cumsum path only supports a single logical batch: g.size(0) must be 1. gate_chunk_cumsum_operands raises this when it sees a varlen request whose gate tensor still carries a batch dimension greater than 1.

Source

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

    flat_a_log = (
        a_log.reshape(-1)
        if a_log is not None
        else torch.empty(1, device=g.device, dtype=torch.float32)
    )
    flat_bias = (
        dt_bias.reshape(-1)
        if dt_bias is not None
        else torch.empty(1, device=g.device, dtype=torch.float32)
    )
    activate = a_log is not None
    has_bias = dt_bias is not None
    use_lower_bound = lower_bound is not None
    lower_bound_value = 0.0 if lower_bound is None else lower_bound

    is_varlen = cu_seqlens is not None
    if is_varlen:
        if g.size(0) != 1:
            raise ValueError("varlen KDA requires batch size 1")
        if chunk_indices is None:
            chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)
        metadata = cu_seqlens
        gate_kernel = _gate_cumsum_operands_varlen
    else:
        metadata = torch.empty(0, device=g.device, dtype=torch.int32)
        chunk_indices = torch.empty(0, 2, device=g.device, dtype=torch.long)
        gate_kernel = _gate_cumsum_operands

    return gate_kernel(
        g,
        q,
        k,
        beta,
        flat_a_log,
        flat_bias,
        metadata,
        chunk_indices,

View on GitHub (pinned to 0132848349)

Solutions

1. Flatten the batch into [1, total_tokens] and pass matching cu_seqlens covering all packed sequences
2. If you truly need batch > 1, don't use varlen — drop cu_seqlens and pad instead
3. Check the caller (e.g. attention backend prefill) flattens requests before invoking chunk_kda

Example fix

// before
g = g  # [2, T, HV, K] with cu_seqlens
chunk_kda(..., g=g, cu_seqlens=cu)
// after
g = g.reshape(1, -1, *g.shape[2:])  # [1, 2T, HV, K]
chunk_kda(..., g=g, cu_seqlens=cu)  # cu covers both seqs

if cu_seqlens is not None:
    assert g.size(0) == 1, "flatten varlen inputs to [1, total_tokens]"
    g = g.reshape(1, -1, *g.shape[2:])

def valid_varlen_gates(g: torch.Tensor, cu: torch.Tensor | None) -> bool:
    return cu is None or g.size(0) == 1

When it happens

Trigger: Calling chunk_kda with cu_seqlens set (varlen batch) but g shaped [B>1, T, ...] — e.g. feeding a padded batch tensor together with cumulative-sequence-length metadata.

Common situations: Mixing padded-batch tensors with varlen metadata after a scheduler change; tests that pass a batch of 2 with cu_seqlens; leftover batching from a prefill aggregation path that should have flattened to [1, total_tokens].

Related errors

• The batch size is expected to be 1 rather than {q.shape[0]}
• Unexpected a shape for varlen: {a.shape}
• `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=

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