sgl-project/sglang · error · RuntimeError

kv_scales shape {tuple(kv_scales.shape)} does not match expe

Error message

kv_scales shape {tuple(kv_scales.shape)} does not match expected ({unified_kv.shape[0]}, {expected_g})

What it means

Given unified_kv of shape (num_pages, ..., D), the kernel requires kv_scales to be exactly (unified_kv.shape[0], D // _FP8_GROUP_SIZE): one row of per-group scales per page. Any other shape (wrong page count, wrong group count, flat vector) is rejected.

Source

Thrown at python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_decode.py:676

        )

    quant_kv = kv_scales is not None
    if quant_kv:
        if unified_kv.dtype != _FP8_DTYPE:
            raise RuntimeError(
                f"kv_scales supplied but unified_kv is {unified_kv.dtype}, "
                f"expected {_FP8_DTYPE}"
            )
        if kv_scales.dtype != torch.float32:
            raise RuntimeError(f"kv_scales must be fp32, got {kv_scales.dtype}")
        D_check = unified_kv.shape[-1]
        if D_check % _FP8_GROUP_SIZE != 0:
            raise RuntimeError(
                f"D={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE}"
            )
        expected_g = D_check // _FP8_GROUP_SIZE
        if kv_scales.shape != (unified_kv.shape[0], expected_g):
            raise RuntimeError(
                f"kv_scales shape {tuple(kv_scales.shape)} does not match "
                f"expected ({unified_kv.shape[0]}, {expected_g})"
            )
        if kv_scales.stride(-1) != 1:
            kv_scales = kv_scales.contiguous()
    else:
        if unified_kv.dtype != q.dtype:
            raise RuntimeError(
                f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}"
            )

    T, H, D = q.shape
    out = torch.empty_like(q)

    if block_h is None:
        block_h = triton.next_power_of_2(min(H, 64))
    else:
        block_h = triton.next_power_of_2(block_h)

View on GitHub (pinned to 0132848349)

Solutions

  1. Recompute/requantize kv_scales against the current unified_kv so its shape is (num_pages, D // _FP8_GROUP_SIZE)
  2. Check that the page dimension used when quantizing matches unified_kv.shape[0]
  3. Ensure the cache and scales are produced by the same allocation/quantization step, not independently
Defensive patterns

Strategy: validation

Validate before calling

if kv_scales is not None:
    expected = (unified_kv.shape[0], unified_kv.shape[-1] // FP8_GROUP_SIZE)
    assert tuple(kv_scales.shape) == expected, (kv_scales.shape, expected)

Prevention

When it happens

Trigger: Passing kv_scales with shape (num_tokens, groups) instead of (num_pages, groups), or scales computed for a different cache size / head_dim than the current unified_kv.

Common situations: Reusing a scales tensor after the paged cache grew (pages appended); mismatch between the quantization routine's page layout and the kernel's expectation; off-by-one in group count after a head_dim change.

Related errors


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