sgl-project/sglang · error · RuntimeError

head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={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.

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)

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={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE
• rope_pool_fused expects pool tensors to be 3-D
• k_pool has incompatible shape {k_pool.shape}
• v_pool shape must match k_pool shape, got {v_pool.shape} vs
• num_heads must be divisible by num_epi_subtiles

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