sgl-project/sglang · error · ValueError

ReplaySSM inputs must be on the same device.

Error message

ReplaySSM inputs must be on the same device.

What it means

All ReplaySSM decode inputs (d_cache, k_cache, g_cache, write_pos) must live on the same CUDA device as mixed_qkv. The wrapper compares each tensor's .device against mixed_qkv.device because the Helion/Triton kernel cannot follow pointers across devices and would silently corrupt or crash. This fires when caches are on cuda:0 but the activations are on cuda:1, or caches are left on CPU.

Source

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

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

View on GitHub (pinned to 0132848349)

Solutions

1. Move all caches and write_pos to mixed_qkv.device before the call (or allocate them on mixed_qkv.device in the first place)
2. In TP runs, allocate per-rank state pools on torch.cuda.current_device() rather than a fixed cuda:0
3. Add a debug assert: assert all(t.device == mixed_qkv.device for t in (...)) in your wrapper

Example fix

# before
out = helion_fused_recurrent_kda_replayssm_decode(mixed_qkv, d_cache_cpu, k_cache_cpu, g_cache_cpu, write_pos, ...)
# after
dev = mixed_qkv.device
out = helion_fused_recurrent_kda_replayssm_decode(mixed_qkv, d_cache_cpu.to(dev), k_cache_cpu.to(dev), g_cache_cpu.to(dev), write_pos.to(dev), ...)

dev = mixed_qkv.device
assert all(t.device == dev for t in (d_cache, k_cache, g_cache, write_pos))

When it happens

Trigger: Calling helion_fused_recurrent_kda_replayssm_decode in a multi-GPU (TP/EP) setup where the state pool was allocated on a different rank/device than mixed_qkv, or caches created on CPU ('cpu' device) for testing and never moved.

Common situations: Tensor-parallel inference where the memory pool is allocated once on the primary device; unit tests that build caches with device='cpu' by default; device migration code that moves qkv but not the state pool.

Related errors

• `force_flush` must be on the same device as the inputs.
• All 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].

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