sgl-project/sglang · error · ValueError

`write_pos` must be a 1D int32 tensor.

Error message

`write_pos` must be a 1D int32 tensor.

What it means

ReplaySSM decode writes each row's recomputed output into a per-request ring cache at position write_pos[row]; the kernel requires write_pos to be a 1D int32 tensor so its indexing is well-defined. The guard rejects other dtypes (int64, bool) and multi-dimensional tensors before the launch.

Source

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

    if a.ndim not in (2, 3) or not a.is_contiguous():
        raise ValueError("KDA `a` must be a contiguous 2D or 3D tensor.")
    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].")

View on GitHub (pinned to 0132848349)

Solutions

1. Cast explicitly: write_pos = write_pos.to(torch.int32)
2. Construct with the right dtype: torch.arange(B, dtype=torch.int32, device=...)
3. Squeeze stray dims: write_pos = write_pos.squeeze(-1) if it came in as [B,1]

Example fix

// before
write_pos = torch.arange(batch, device=dev)  # int64
// after
write_pos = torch.arange(batch, device=dev, dtype=torch.int32)

if write_pos.dtype is not torch.int32 or write_pos.ndim != 1:
    write_pos = write_pos.reshape(-1).to(torch.int32)

def valid_write_pos(t: torch.Tensor) -> bool:
    return isinstance(t, torch.Tensor) and t.ndim == 1 and t.dtype is torch.int32

When it happens

Trigger: Passing write_pos as torch.int64 (the default from torch.arange / tensor(...)), a 0-dim scalar tensor, or a [B,1] tensor to helion_fused_recurrent_kda_replayssm_decode.

Common situations: Creating write_pos with torch.arange without dtype=torch.int32; index arithmetic promoting to long; a caller porting from an API that accepted int64 positions.

Related errors

• `force_flush` must be a length-B int32 tensor or None.
• KDA `a` must be a contiguous 2D or 3D tensor.
• `g_cache` must have dtype torch.float32.
• `write_pos` must be a 1D int32 tensor.
• {name} must have dtype torch.int32

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