sgl-project/sglang · error · ValueError

`write_pos` must have shape {(batch,)}.

Error message

`write_pos` must have shape {(batch,)}.

What it means

write_pos must have exactly one entry per batch row — shape (batch,) where batch = mixed_qkv.size(0) — because the kernel indexes it per row to place the flushed output in the replay cache. A mismatched length means rows would read out-of-bounds positions.

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_replayssm.py:726

    if dt_bias.ndim not in (1, 2) or not dt_bias.is_contiguous():
        raise ValueError("KDA `dt_bias` must be a contiguous 1D or 2D tensor.")
    flat_a = a.view(batch, -1)
    flat_dt_bias = dt_bias.view(-1)
    _, num_q_heads, num_v_heads, key_dim, value_dim = validate_packed_decode_inputs(
        mixed_qkv,
        flat_a,
        b,
        A_log,
        flat_dt_bias,
        initial_state,
        out,
        ssm_state_indices,
    )

    if write_pos.ndim != 1 or write_pos.dtype is not torch.int32:
        raise ValueError("`write_pos` must be a 1D int32 tensor.")
    if write_pos.shape != (batch,):
        raise ValueError(f"`write_pos` must have shape {(batch,)}.")
    if force_flush is not None and (
        force_flush.ndim != 1
        or force_flush.dtype is not torch.int32
        or force_flush.shape != (batch,)
    ):
        raise ValueError("`force_flush` must be a length-B int32 tensor or None.")

    cache_length = d_cache.size(2)
    if cache_length < 1:
        raise ValueError("ReplaySSM cache length must be at least 1.")
    if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):
        raise ValueError("`d_cache` must have shape [slots, HV, L, V].")
    if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):
        raise ValueError("`k_cache` must have shape [slots, H, L, K].")
    if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):
        raise ValueError("`g_cache` must have shape [slots, HV, L, K].")
    if g_cache.dtype is not torch.float32:
        raise ValueError("`g_cache` must have dtype torch.float32.")

View on GitHub (pinned to 0132848349)

Solutions

1. Regenerate write_pos per call: torch.full((batch,), pos, dtype=torch.int32, device=...)
2. Slice consistently: if you index mixed_qkv rows, index write_pos with the same mask
3. Assert write_pos.shape == (mixed_qkv.size(0),) before the call in debug builds

Example fix

// before
write_pos = cached_write_pos  # built for batch=8, now batch=5
// after
write_pos = torch.full((mixed_qkv.size(0),), pos, dtype=torch.int32, device=mixed_qkv.device)

batch = mixed_qkv.size(0)
assert write_pos.shape == (batch,), (write_pos.shape, batch)

def write_pos_matches(qkv: torch.Tensor, wp: torch.Tensor) -> bool:
    return wp.shape == (qkv.size(0),)

When it happens

Trigger: Passing write_pos sized for a different batch than mixed_qkv, e.g. after slicing qkv rows without slicing write_pos, or sharing a cached write_pos tensor across calls with varying batch.

Common situations: Continuous batching where the running batch size changes between steps but a stale write_pos is reused; tests reusing fixtures across batch sizes; prefix slicing of the batch dim on some tensors but not others.

Related errors

• Unexpected A_log shape: {A_log.shape}; expected numel={HV}
• Unexpected dt_bias shape: {dt_bias.shape}; expected numel={H
• Unexpected a shape for varlen: {a.shape}
• Unexpected a shape for dense: {a.shape}
• The number of initial states is expected to be equal to the

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