sgl-project/sglang · error · ValueError

All inputs must be on the same device.

Error message

All inputs must be on the same device.

What it means

The packed decode kernel launches on mixed_qkv.device and reads/writes a, b, A_log, dt_bias, initial_state, out, and ssm_state_indices with raw pointers, so every tensor must live on the same CUDA device. The wrapper iterates over all inputs and raises on the first mismatch.

Solutions

1. Move everything to one device: tensors = [t.to(dev, non_blocking=True) for t in tensors] before the call
2. In TP setups, pass each rank its own slice already resident on that rank's device; verify with assert all(t.device == dev ...)

Example fix

# before
out, s = ...(initial_state=cpu_state, ...)  # cpu_state on 'cpu'
# after
out, s = ...(initial_state=cpu_state.to(mixed_qkv.device), ...)

dev = mixed_qkv.device
for t in (a, b, A_log, dt_bias, initial_state, out, ssm_state_indices):
    assert t.device == dev, (t.shape, t.device, dev)

def all_on_device(dev: torch.device, *ts: torch.Tensor) -> bool:
    return all(t.device == dev for t in ts)

When it happens

Trigger: Keeping initial_state/out on cuda:0 while the token inputs are on cuda:1 in a TP>=2 setup; parameters (A_log/dt_bias) left on CPU after a partial .to(device); mixed-precision pipelines where some buffers were moved and others weren't.

Common situations: Multi-GPU tensor parallel runs with per-rank device ids; loading model weights to meta/cpu and forgetting SSM state caches; writing a single-device benchmark then running under device_map=auto.

Related errors

• `a` and `b` must be 2D tensors
• `a`/`b` must be contiguous in the last dim.
• `a`/`b` must have shape [B, HV] with HV=
• `A_log` and `dt_bias` must have
• `A_log`/`dt_bias` must be 1D tensors.

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

Appendix: source

Thrown at python/sglang/kernels/ops/attention/fla/fused_recurrent.py:308

    if a.stride(-1) != 1 or b.stride(-1) != 1:
        raise ValueError("`a`/`b` must be contiguous in the last dim.")
    if A_log.ndim != 1 or dt_bias.ndim != 1:
        raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
        raise ValueError("`A_log`/`dt_bias` must be contiguous.")
    if ssm_state_indices.ndim != 1:
        raise ValueError(
            f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
        )
    if not out.is_contiguous():
        raise ValueError("`out` must be contiguous.")

    dev = mixed_qkv.device
    if any(
        t.device != dev
        for t in (a, b, A_log, dt_bias, initial_state, out, ssm_state_indices)
    ):
        raise ValueError("All inputs must be on the same device.")

    B = mixed_qkv.shape[0]
    if a.shape[0] != B or b.shape[0] != B:
        raise ValueError(
            "Mismatched batch sizes: "
            f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}."
        )
    if ssm_state_indices.shape[0] != B:
        raise ValueError(
            f"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},))."
        )

    if initial_state.ndim != 4:
        raise ValueError(
            f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
        )
    if initial_state.stride(-1) != 1:
        raise ValueError("`initial_state` must be contiguous in the last dim.")

View on GitHub (pinned to 0132848349)