sgl-project/sglang · error · ValueError

attn_sink must be float32 with shape

Error message

attn_sink must be float32 with shape ({h_q},), got {tuple(attn_sink.shape)}/{attn_sink.dtype}

What it means

sparse_mla_q8kv8_prefill_fwd validates the optional attn_sink tensor before launching the sparse MLA q8kv8 prefill kernel. attn_sink must be a 1-D float32 tensor of length h_q (the number of query heads); any other shape or dtype is rejected because the kernel indexes it per head. The error message reports both the offending shape and dtype.

Solutions

Build the sink as torch.zeros(h_q, dtype=torch.float32, device=q.device) and fill per-head values
Check h_q from q.shape[1] before constructing attn_sink so lengths always match
Cast an existing sink with attn_sink.float() if the values are right but the dtype is wrong

Example fix

// before
attn_sink = torch.zeros(num_kv_heads, dtype=torch.bfloat16, device=q.device)
sparse_mla_q8kv8_prefill_fwd(q, ..., attn_sink=attn_sink)
// after
h_q = q.shape[1]
attn_sink = torch.zeros(h_q, dtype=torch.float32, device=q.device)
sparse_mla_q8kv8_prefill_fwd(q, ..., attn_sink=attn_sink)

Defensive patterns

Strategy: validation

Validate before calling

assert attn_sink is None or (attn_sink.dtype == torch.float32 and attn_sink.shape == (q.shape[1],)), 'attn_sink must be float32 shape (h_q,)'

Type guard

def valid_attn_sink(t: torch.Tensor, h_q: int) -> bool:
    return t.dtype == torch.float32 and tuple(t.shape) == (h_q,)

Prevention

Derive h_q from q.shape[1] at the call site instead of hardcoding head counts
Centralize sink construction in one helper that always uses float32/(h_q,)

When it happens

Trigger: Calling sparse_mla_q8kv8_prefill_fwd(..., attn_sink=tensor) where tensor has shape (h_kv,), a scalar shape (), (1, h_q), or dtype torch.bfloat16/float16 instead of torch.float32.

Common situations: Passing a KV-head-shaped sink (forgetting it is per query head), reusing a bf16 model parameter as the sink, or passing a sink tensor that was indexed/squeezed incorrectly.

Related errors

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

Appendix: source

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

            )
        if not topk_length.is_contiguous():
            raise ValueError("topk_length must be contiguous")
        if torch.any(topk_length < 0).item() or torch.any(topk_length > topk).item():
            raise ValueError(
                "topk_length values must satisfy " f"0 <= topk_length <= topk ({topk})"
            )

    if d_v != 512:
        raise ValueError(
            f"sparse_mla_q8kv8_prefill_fwd only supports d_v=512, got {d_v}"
        )

    if attn_sink is not None and topk_length is None:
        raise ValueError("attn_sink requires topk_length to be provided as well")

    if attn_sink is not None:
        if attn_sink.shape != (h_q,) or attn_sink.dtype != torch.float32:
            raise ValueError(
                f"attn_sink must be float32 with shape ({h_q},), got "
                f"{tuple(attn_sink.shape)}/{attn_sink.dtype}"
            )
        if not attn_sink.is_cuda:
            raise ValueError("attn_sink must be a CUDA tensor")
        if attn_sink.device != device:
            raise ValueError(
                f"attn_sink must be on q's device {device}, got {attn_sink.device}"
            )
        if not attn_sink.is_contiguous():
            raise ValueError("attn_sink must be contiguous")

    for name, scale in (("q_scale", q_scale), ("kv_scale", kv_scale)):
        if not isinstance(scale, torch.Tensor):
            raise ValueError(f"{name} must be a torch.Tensor")
        if not scale.is_cuda:
            raise ValueError(f"{name} must be a CUDA tensor")
        if scale.device != device:

View on GitHub (pinned to 0132848349)