{"record":{"id":"594788c0f8ce4ccd","repo":"sgl-project/sglang","slug":"invalid-k-out-shape-for-fused-kv-materialization","errorCode":null,"errorMessage":"Invalid k_out shape for fused KV materialization: got {tuple(k_out.shape)}, expected {expected_shape}.","messagePattern":"Invalid k_out shape for fused KV materialization: got (.+?), expected (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/speculative/fused_kv_materialize.py","lineNumber":187,"sourceCode":"        raise ValueError(\n            \"Invalid stacked eps shape for fused KV materialization: \"\n            f\"got {tuple(eps.shape)}, expected {(n_layers,)}.\"\n        )\n\n    half_rotary_dim = rotary_dim // 2\n    BLOCK_HD = triton.next_power_of_2(head_dim)\n\n    if positions.device != kv.device:\n        positions = positions.to(device=kv.device, dtype=torch.int64)\n    elif positions.dtype != torch.int64:\n        positions = positions.to(torch.int64)\n\n    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)\n    if k_out is None:\n        k_out = torch.empty(expected_shape, dtype=kv.dtype, device=kv.device)\n    else:\n        if k_out.shape != expected_shape:\n            raise ValueError(\n                \"Invalid k_out shape for fused KV materialization: \"\n                f\"got {tuple(k_out.shape)}, expected {expected_shape}.\"\n            )\n        if k_out.device != kv.device or k_out.dtype != kv.dtype:\n            raise ValueError(\n                \"Invalid k_out device/dtype for fused KV materialization: \"\n                f\"got device={k_out.device}, dtype={k_out.dtype}, \"\n                f\"expected device={kv.device}, dtype={kv.dtype}.\"\n            )\n    if v_out is None:\n        v_out = torch.empty_like(k_out)\n    else:\n        if v_out.shape != expected_shape:\n            raise ValueError(\n                \"Invalid v_out shape for fused KV materialization: \"\n                f\"got {tuple(v_out.shape)}, expected {expected_shape}.\"\n            )\n        if v_out.device != kv.device or v_out.dtype != kv.dtype:","sourceCodeStart":169,"sourceCodeEnd":205,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/speculative/fused_kv_materialize.py#L169-L205","documentation":"When a preallocated k_out buffer is supplied, its shape must exactly equal [n_layers, total_ctx, num_kv_heads, head_dim]. Otherwise the kernel would write out of bounds or produce a wrongly laid-out cache.","triggerScenarios":"Passing k_out with wrong layer/tokens dimension order (e.g. [total_ctx, n_layers, ...]) or sized for a different token count.","commonSituations":"Reusing a KV cache slab across requests without resizing, or permuting dimensions when adapting an existing cache layout to the fused path.","solutions":["Allocate k_out as torch.empty((n_layers, total_ctx, num_kv_heads, head_dim), ...) or pass None to let it be allocated.","Double-check total_ctx equals positions.numel() and the kv tensor's dim 0.","Match n_layers to kv.shape[1]."],"exampleFix":"// before\nk_out = torch.empty(total_ctx, n_layers, H, D)\n// after\nk_out = torch.empty(n_layers, total_ctx, H, D, dtype=kv.dtype, device=kv.device)","handlingStrategy":"validation","validationCode":"expected = (n_layers, total_ctx, num_kv_heads, head_dim)\nassert k_out is None or k_out.shape == expected","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer passing None and letting the function allocate.","Never reuse cache buffers across differently-sized batches."],"tags":["shape-validation","kv-cache","speculative-decoding"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}