sgl-project/sglang · error · ValueError

{name} must be on device {device}, got {t.device}

Error message

{name} must be on device {device}, got {t.device}

What it means

All output buffers passed to sparse_mla_q8kv8_prefill_fwd must live on the same CUDA device the kernel launches on. _check_out_buffer compares t.device against the expected device and raises ValueError on mismatch.

Source

Thrown at python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py:118

# id and must not be used here.
_get_current_stream_raw = torch._C._cuda_getCurrentRawStream


# Module-level cache for kernel-write-only output tensors. The active s_q rows
# are overwritten every call; buffers grow monotonically by device/head shape.
def _check_out_buffer(
    t: torch.Tensor,
    name: str,
    shape: tuple,
    dtype: torch.dtype,
    device: torch.device,
) -> None:
    if tuple(t.shape) != tuple(shape):
        raise ValueError(f"{name} must have shape {tuple(shape)}, got {tuple(t.shape)}")
    if t.dtype != dtype:
        raise ValueError(f"{name} must have dtype {dtype}, got {t.dtype}")
    if t.device != device:
        raise ValueError(f"{name} must be on device {device}, got {t.device}")
    if not t.is_contiguous():
        raise ValueError(f"{name} must be contiguous")


# Internal custom-op wrappers so the JIT kernel calls participate in
# torch.library / torch.compile tracing and kernel-API debug logging.
# The dispatch_full variant carries the optional attn_sink / topk_length
# tensors as required args; the public API chooses which op to call.
@register_custom_op(
    op_name="sparse_mla_q8kv8_prefill",
    mutates_args=["out", "max_logits", "lse"],
)
def _sparse_mla_q8kv8_prefill_op(
    q: torch.Tensor,
    kv: torch.Tensor,
    indices: torch.Tensor,
    q_scale: torch.Tensor,
    kv_scale: torch.Tensor,

View on GitHub (pinned to 0132848349)

Solutions

  1. Allocate output buffers with device=q.device (derive from the inputs)
  2. Never hardcode 'cuda' or 'cuda:0' in TP/EP rank code; use the rank's device or set_device

Example fix

# before
out = torch.empty(s_q, h_q, d_v, dtype=torch.bfloat16, device='cuda')  # maybe cuda:0
# after
out = torch.empty(s_q, h_q, d_v, dtype=torch.bfloat16, device=q.device)
Defensive patterns

Strategy: validation

Validate before calling

for buf in (out, max_logits, lse):
    assert buf.device == q.device, (buf.device, q.device)

Prevention

When it happens

Trigger: Passing out tensors allocated on a different GPU (or on CPU) than the one q/kv/indices reside on — common with cuda:0 vs cuda:1 mismatches in multi-GPU (TP/EP) inference.

Common situations: Tensor-parallel serving where buffers are allocated with a device string instead of the input tensors' device; multi-process rank code assuming device 0.

Related errors


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