sgl-project/sglang · error · ValueError

`force_flush` must be a length-B int32 tensor or None.

Error message

`force_flush` must be a length-B int32 tensor or None.

What it means

force_flush optionally forces a per-row flush of the ReplaySSM cache. When provided it must be a 1D int32 tensor of length batch (or None). The guard rejects wrong dtype, wrong rank, or wrong length so the kernel's per-row branch on force_flush[row] is valid.

Source

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

        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.")

    device = mixed_qkv.device
    if any(
        tensor.device != device for tensor in (d_cache, k_cache, g_cache, write_pos)
    ):
        raise ValueError("ReplaySSM inputs must be on the same device.")

View on GitHub (pinned to 0132848349)

Solutions

  1. Pass None when no forced flush is needed
  2. Build a per-row flag: torch.zeros/ones(batch, dtype=torch.int32, device=...) (or torch.where(cond, 1, 0).int())
  3. Never pass Python bools; the API has no scalar overload

Example fix

// before
out = helion_fused_recurrent_kda_replayssm_decode(..., force_flush=True)
// after
force = torch.ones(mixed_qkv.size(0), dtype=torch.int32, device=mixed_qkv.device)
out = helion_fused_recurrent_kda_replayssm_decode(..., force_flush=force)
Defensive patterns

Strategy: type-guard

Validate before calling

def prep_force_flush(batch, device, flag=False):
    if flag is None or flag is False:
        return None
    return torch.full((batch,), int(flag), dtype=torch.int32, device=device)

force_flush = prep_force_flush(mixed_qkv.size(0), mixed_qkv.device, want_flush)

Type guard

def valid_force_flush(t, batch) -> bool:
    return t is None or (isinstance(t, torch.Tensor) and t.ndim == 1 and t.dtype is torch.int32 and t.shape == (batch,))

Prevention

When it happens

Trigger: Passing force_flush as a Python bool, a bool tensor, an int64 tensor, or a [B,1] tensor to helion_fused_recurrent_kda_replayssm_decode.

Common situations: Calling with force_flush=True for an unconditional flush instead of a per-row tensor; broadcasting a single flag to [B] via expand (non-contiguous is fine here but dtype stays wrong); porting from a bool-mask API.

Related errors


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