{"record":{"id":"32a9bdda3b4426d7","repo":"sgl-project/sglang","slug":"head-dim-mismatch-unified-kv-unified-kv-size-1","errorCode":null,"errorMessage":"head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.size(-1)}","messagePattern":"head_dim mismatch: unified_kv=(.+?), kv=(.+?)","errorType":"validation","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_prefill.py","lineNumber":243,"sourceCode":"    attn_sink: torch.Tensor,\n    softmax_scale: float,\n) -> torch.Tensor:\n    if not q.is_cuda:\n        raise RuntimeError(\n            \"Triton sparse_attn_v4_paged_prefill requires CUDA/HIP tensors\"\n        )\n    if q.dtype not in (torch.bfloat16, torch.float16):\n        raise RuntimeError(\n            f\"sparse_attn_v4_paged_prefill expects fp16/bf16 q, got {q.dtype}\"\n        )\n    if unified_kv.dtype != q.dtype:\n        raise RuntimeError(\n            f\"unified_kv dtype mismatch: kv={unified_kv.dtype}, q={q.dtype}\"\n        )\n    if kv.dtype != q.dtype:\n        raise RuntimeError(f\"kv dtype mismatch: kv={kv.dtype}, q={q.dtype}\")\n    if unified_kv.size(-1) != kv.size(-1):\n        raise RuntimeError(\n            f\"head_dim mismatch: unified_kv={unified_kv.size(-1)}, kv={kv.size(-1)}\"\n        )\n\n    T, H, D = q.shape\n    out = torch.empty_like(q)\n    kv_indices_prefix = kv_indices_prefix.to(torch.int32).contiguous()\n    kv_indptr_prefix = kv_indptr_prefix.to(torch.int32).contiguous()\n    kv_indices_extend = kv_indices_extend.to(torch.int32).contiguous()\n    kv_indptr_extend = kv_indptr_extend.to(torch.int32).contiguous()\n\n    block_h = 16  # AMD MFMA min tile\n    block_d = triton.next_power_of_2(D)\n    block_k = 16 if D >= 256 else 32\n    _sparse_attn_v4_paged_prefill_kernel[(T, triton.cdiv(H, block_h))](\n        q,\n        unified_kv,\n        kv_indices_prefix,\n        kv_indptr_prefix,","sourceCodeStart":225,"sourceCodeEnd":261,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/dsv4/unified_kv_kernels/paged_prefill.py#L225-L261","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Verify both the cache allocation and the extend kv projection use the same head_dim from model config","Rebuild/reallocate the unified KV cache after any head_dim or model change","Print both .shape[-1] values at the call site to find which producer is wrong"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert unified_kv.size(-1) == kv.size(-1) == q.size(-1), (\n    unified_kv.size(-1), kv.size(-1), q.size(-1))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Validate head_dim consistency once at model load, covering cache allocation","Reallocate caches after any model geometry change"],"tags":["attention","head-dim","shape-validation","kv-cache"],"backgroundTag":"shape-validation-failed","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}