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.

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)

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), ...)
Defensive patterns

Strategy: validation

Validate before calling

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)

Type guard

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

Prevention

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

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