sgl-project/sglang · error · ValueError

Invalid k_out device/dtype for fused KV materialization: got

Error message

Invalid k_out device/dtype for fused KV materialization: got device={k_out.device}, dtype={k_out.dtype}, expected device={kv.device}, dtype={kv.dtype}.

What it means

A caller-supplied k_out must live on the same device and have the same dtype as the kv input tensor, since the Triton kernel writes in-place with no conversion.

Source

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

    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:
            raise ValueError(
                "Invalid v_out device/dtype for fused KV materialization: "
                f"got device={v_out.device}, dtype={v_out.dtype}, "
                f"expected device={kv.device}, dtype={kv.dtype}."
            )

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate k_out with device=kv.device, dtype=kv.dtype or via torch.empty_like semantics.
2. If None is passed, the function allocates correctly — prefer that.
3. Move persistent buffers together with the model using register_buffer so .to() tracks them.

Example fix

// before
k_out = torch.empty(shape, dtype=torch.float16)  # kv is bf16
// after
k_out = None  # or torch.empty(shape, dtype=kv.dtype, device=kv.device)

assert k_out is None or (k_out.device == kv.device and k_out.dtype == kv.dtype)

When it happens

Trigger: Passing a CPU buffer while kv is on GPU, or a float16 buffer with bfloat16 kv.

Common situations: Preallocating output buffers at module init on the wrong device, or after moving the model with .to(device) without moving the persistent buffer.

Related errors

• Invalid v_out device/dtype for fused KV materialization: got
• v_cache must be provided
• k_cache can only be None when only_qv=True
• rope_pool_fused expects pool tensors to be 3-D
• k_pool has incompatible shape {k_pool.shape}

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