sgl-project/sglang · error · ValueError

`force_flush` must be on the same device as the inputs.

Error message

`force_flush` must be on the same device as the inputs.

What it means

When force_flush is supplied (a per-row flush control tensor), it must sit on the same device as the other ReplaySSM inputs (mixed_qkv's device). The wrapper checks force_flush.device != device separately from the cache checks because force_flush is optional; a CPU-resident index tensor with CUDA caches is the classic mistake. Launching with a cross-device index would fault inside the Triton kernel.

Source

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

    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.")
    if force_flush is not None and force_flush.device != device:
        raise ValueError("`force_flush` must be on the same device as the inputs.")

    cache_block = helion.next_power_of_2(max(16, cache_length))
    use_lower_bound = lower_bound is not None
    kernel = _select_replayssm_decode_kernel(
        is_bf16_state=initial_state.dtype is torch.bfloat16,
        num_v_heads=num_v_heads,
    )
    result = kernel(
        mixed_qkv,
        flat_a,
        b,
        A_log,
        flat_dt_bias,
        scale,
        initial_state,
        d_cache,
        k_cache,
        g_cache,

View on GitHub (pinned to 0132848349)

Solutions

1. Create force_flush with device=mixed_qkv.device (e.g. torch.zeros(B, dtype=torch.int32, device=qkv.device))
2. Or call force_flush = force_flush.to(mixed_qkv.device) before invoking the op
3. If flushing all rows, pass None and use the kernel's implicit flush path instead of a CPU tensor of ones

Example fix

# before
force_flush = torch.zeros(batch, dtype=torch.int32)
# after
force_flush = torch.zeros(batch, dtype=torch.int32, device=mixed_qkv.device)

if force_flush is not None and force_flush.device != mixed_qkv.device:\n    force_flush = force_flush.to(mixed_qkv.device)

When it happens

Trigger: Calling helion_fused_recurrent_kda_replayssm_decode with force_flush created via torch.zeros(B, dtype=torch.int32) on CPU (no device= argument) while all caches are on CUDA; commonly exercised by test_replayssm_per_row_flush_contract.

Common situations: Constructing control tensors like torch.zeros(B, dtype=torch.int32) without device=; scheduling code that builds flush flags on the scheduler (CPU) process and passes them straight to the GPU kernel wrapper.

Related errors

• ReplaySSM inputs must be on the same device.
• `d_cache` must have shape [slots, HV, L, V].
• `k_cache` must have shape [slots, H, L, K].
• `g_cache` must have shape [slots, HV, L, K].
• `g_cache` must have dtype torch.float32.

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