sgl-project/sglang · error · ValueError

{name} must have shape {tuple(shape)}, got {tuple(t.shape)}

Error message

{name} must have shape {tuple(shape)}, got {tuple(t.shape)}

What it means

sparse_mla_q8kv8_prefill_fwd validates caller-provided output buffers through _check_out_buffer; each buffer must exactly match the expected shape for the kernel launch. A mismatch raises ValueError naming the buffer, expected shape, and actual shape.

Source

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

# ---------------------------------------------------------------------------

# torch._C._cuda_getCurrentRawStream returns the cudaStream_t pointer expected
# by the JIT wrapper. torch._C._cuda_getCurrentStream returns a packed stream
# 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,

View on GitHub (pinned to 0132848349)

Solutions

  1. Allocate output buffers from the current inputs' shapes at call time (or use the API variant that allocates outputs)
  2. Key any output-buffer cache on (s_q, h_q, topk, dtype) exactly
  3. Re-check the documented shapes in the docstring (e.g. out [s_q,h_q,d_v], max_logits/lse [s_q,h_q])

Example fix

# before
out = torch.empty(prev_s_q, h_q, d_v, ...)  # stale s_q
sparse_mla_q8kv8_prefill_fwd(q, kv, indices, out=out, ...)
# after
out = torch.empty(q.shape[0], q.shape[1], d_v, dtype=torch.bfloat16, device=q.device)
sparse_mla_q8kv8_prefill_fwd(q, kv, indices, out=out, ...)
Defensive patterns

Strategy: validation

Validate before calling

expected = (s_q, h_q, d_v)
assert tuple(out.shape) == expected, (out.shape, expected)
sparse_mla_q8kv8_prefill_fwd(q, kv, idx, out=out, ...)

Prevention

When it happens

Trigger: Calling sparse_mla_q8kv8_prefill_fwd with an out= buffer (e.g. out/max_logits/lse) allocated with wrong dimensions — typically a stale shape from a different batch size, num heads, or sequence length.

Common situations: Reusing preallocated output buffers across requests with changing s_q or h_q; buffer caches keyed incompletely (missing a dim); after changing topk or head config without reallocating outputs.

Related errors

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