sgl-project/sglang · error · ValueError

kv d_qk must match q d_qk={d_qk}, got {kv_d_qk}

Error message

kv d_qk must match q d_qk={d_qk}, got {kv_d_qk}

What it means

The QK head dimension of the kv cache (kv_d_qk, derived from kv's shape) must equal the q tensor's head dimension d_qk (512 or 576). The kernel assumes q and kv share the same per-head QK width for its FP8 dot products.

Source

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

    if indices.device != device:
        raise ValueError(
            f"indices must be on q's device {device}, got {indices.device}"
        )

    if q.dtype != torch.float8_e4m3fn:
        raise ValueError(f"q must be torch.float8_e4m3fn, got {q.dtype}")
    if kv.dtype != torch.float8_e4m3fn:
        raise ValueError(f"kv must be torch.float8_e4m3fn, got {kv.dtype}")

    if not q.is_contiguous():
        raise ValueError("q must be contiguous")
    if not kv.is_contiguous():
        raise ValueError("kv must be contiguous")
    if not indices.is_contiguous():
        raise ValueError("indices must be contiguous")

    if kv_d_qk != d_qk:
        raise ValueError(f"kv d_qk must match q d_qk={d_qk}, got {kv_d_qk}")

    # The CUDA implementation uses B_H=64 and launches h_q / B_H CTAs.
    # Reject unpadded TP-local head counts instead of launching zero CTAs and
    # returning uninitialized outputs, which can appear to callers as a hang or
    # a later collective failure.
    if h_q == 0 or h_q % 64 != 0:
        raise ValueError(
            "sparse_mla_q8kv8_prefill_fwd requires h_q padded to a positive "
            f"multiple of 64, got {h_q}"
        )

    if h_kv != 1:
        raise ValueError(f"sparse_mla_q8kv8_prefill_fwd requires h_kv=1, got {h_kv}")

    if d_qk not in (512, 576):
        raise ValueError(
            f"sparse_mla_q8kv8_prefill_fwd supports d_qk=512/576, got {d_qk}"
        )

View on GitHub (pinned to 0132848349)

Solutions

  1. Make the kv cache's QK head dim match q's (reallocate the cache with the correct d_qk)
  2. Verify the model config's qk_head_dim / kv_lora_rank settings flow into both the projection and the cache allocator
  3. If intentional mismatch (nope-only kv), this kernel does not support it; use a different backend

Example fix

// before
kv = torch.empty(..., 512, ..., dtype=torch.float8_e4m3fn)  # q is 576
// after
kv = torch.empty(..., 576, ..., dtype=torch.float8_e4m3fn)  # matches q d_qk
Defensive patterns

Strategy: validation

Validate before calling

d_qk = q.shape[-1]
assert kv.shape[-1] == d_qk, f"kv d_qk {kv.shape[-1]} != q d_qk {d_qk}"

Type guard

def qkv_dims_match(q: torch.Tensor, kv: torch.Tensor) -> bool:
    return q.shape[-1] == kv.shape[-1]

Prevention

When it happens

Trigger: q has d_qk=576 (e.g. 512+64 RoPE carriers) while the kv cache was allocated with d_qk=512, or vice versa.

Common situations: Model variants (DeepSeek MLA with/without rope delta), TP sharding that halves head dims inconsistently, cache allocation using a different config than the attention backend reads.

Related errors


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