sgl-project/sglang · error · ValueError

Invalid k_out shape for fused KV materialization: got {tuple

Error message

Invalid k_out shape for fused KV materialization: got {tuple(k_out.shape)}, expected {expected_shape}.

What it means

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.

Source

Thrown at python/sglang/kernels/ops/speculative/fused_kv_materialize.py:187

        raise ValueError(
            "Invalid stacked eps shape for fused KV materialization: "
            f"got {tuple(eps.shape)}, expected {(n_layers,)}."
        )

    half_rotary_dim = rotary_dim // 2
    BLOCK_HD = triton.next_power_of_2(head_dim)

    if positions.device != kv.device:
        positions = positions.to(device=kv.device, dtype=torch.int64)
    elif positions.dtype != torch.int64:
        positions = positions.to(torch.int64)

    expected_shape = (n_layers, total_ctx, num_kv_heads, head_dim)
    if k_out is None:
        k_out = torch.empty(expected_shape, dtype=kv.dtype, device=kv.device)
    else:
        if k_out.shape != expected_shape:
            raise ValueError(
                "Invalid k_out shape for fused KV materialization: "
                f"got {tuple(k_out.shape)}, expected {expected_shape}."
            )
        if k_out.device != kv.device or k_out.dtype != kv.dtype:
            raise ValueError(
                "Invalid k_out device/dtype for fused KV materialization: "
                f"got device={k_out.device}, dtype={k_out.dtype}, "
                f"expected device={kv.device}, dtype={kv.dtype}."
            )
    if v_out is None:
        v_out = torch.empty_like(k_out)
    else:
        if v_out.shape != expected_shape:
            raise ValueError(
                "Invalid v_out shape for fused KV materialization: "
                f"got {tuple(v_out.shape)}, expected {expected_shape}."
            )
        if v_out.device != kv.device or v_out.dtype != kv.dtype:

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate k_out as torch.empty((n_layers, total_ctx, num_kv_heads, head_dim), ...) or pass None to let it be allocated.
2. Double-check total_ctx equals positions.numel() and the kv tensor's dim 0.
3. Match n_layers to kv.shape[1].

Example fix

// before
k_out = torch.empty(total_ctx, n_layers, H, D)
// after
k_out = torch.empty(n_layers, total_ctx, H, D, dtype=kv.dtype, device=kv.device)

expected = (n_layers, total_ctx, num_kv_heads, head_dim)
assert k_out is None or k_out.shape == expected

When it happens

Trigger: Passing k_out with wrong layer/tokens dimension order (e.g. [total_ctx, n_layers, ...]) or sized for a different token count.

Common situations: Reusing a KV cache slab across requests without resizing, or permuting dimensions when adapting an existing cache layout to the fused path.

Related errors

• 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
• D={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE
• kv_scales shape {tuple(kv_scales.shape)} does not match expe

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