sgl-project/sglang · error · RuntimeError

head_dim mismatch: unified_kv=

Error message

head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.size(-1)}

What it means

sparse_attn_v4_paged_prefill requires the paged unified_kv cache and the non-paged extend kv tensor to have the same head_dim (last dimension). A mismatch means the two KV sources describe different model geometries and cannot be attended over together.

Solutions

1. Verify both the cache allocation and the extend kv projection use the same head_dim from model config
2. Rebuild/reallocate the unified KV cache after any head_dim or model change
3. Print both .shape[-1] values at the call site to find which producer is wrong

assert unified_kv.size(-1) == kv.size(-1) == q.size(-1), (
    unified_kv.size(-1), kv.size(-1), q.size(-1))

When it happens

Trigger: Calling sparse_attn_v4_paged_prefill with unified_kv.shape[-1] != kv.shape[-1], e.g. cache allocated for head_dim 128 while the extend projection emits 64 (or a typo'd config).

Common situations: Changing head_dim in model config without resizing the paged cache; MLA-style models where KV compression dim differs from Q head dim and the wrong dim was cached; leftover cache from a previous model.

Related errors

• D= must be divisible by GROUP_SIZE=
• head_dim must be positive, even, and <= 128
• Invalid k_out shape for fused KV materialization: got
• Invalid v_out shape for fused KV materialization: got
• k_pool has incompatible shape

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

Appendix: source

Thrown at python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_prefill.py:243

    attn_sink: torch.Tensor,
    softmax_scale: float,
) -> torch.Tensor:
    if not q.is_cuda:
        raise RuntimeError(
            "Triton sparse_attn_v4_paged_prefill requires CUDA/HIP tensors"
        )
    if q.dtype not in (torch.bfloat16, torch.float16):
        raise RuntimeError(
            f"sparse_attn_v4_paged_prefill expects fp16/bf16 q, got {q.dtype}"
        )
    if unified_kv.dtype != q.dtype:
        raise RuntimeError(
            f"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}"
        )
    if kv.dtype != q.dtype:
        raise RuntimeError(f"kv dtype mismatch: kv={kv.dtype}, q={q.dtype}")
    if unified_kv.size(-1) != kv.size(-1):
        raise RuntimeError(
            f"head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.size(-1)}"
        )

    T, H, D = q.shape
    out = torch.empty_like(q)
    kv_indices_prefix = kv_indices_prefix.to(torch.int32).contiguous()
    kv_indptr_prefix = kv_indptr_prefix.to(torch.int32).contiguous()
    kv_indices_extend = kv_indices_extend.to(torch.int32).contiguous()
    kv_indptr_extend = kv_indptr_extend.to(torch.int32).contiguous()

    block_h = 16  # AMD MFMA min tile
    block_d = triton.next_power_of_2(D)
    block_k = 16 if D >= 256 else 32
    _sparse_attn_v4_paged_prefill_kernel[(T, triton.cdiv(H, block_h))](
        q,
        unified_kv,
        kv_indices_prefix,
        kv_indptr_prefix,

View on GitHub (pinned to 0132848349)