sgl-project/sglang · error · ValueError

O partial tensor must have 4 or 5 dimensions: (num_splits, b

Error message

O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)

What it means

The combine kernel accepts O partial tensors only in 4D (num_splits, batch, seqlen, nheads, headdim) or 5D varlen forms; other ranks cannot be mapped to the kernel's memory layout.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py:231

        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
        varlen_batch_idx: Optional[cute.Tensor] = None,
        semaphore_to_reset: Optional[cute.Tensor] = None,
        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
        stream: cuda.CUstream = None,
    ):
        # Type checking
        if const_expr(not (mO_partial.element_type == self.dtype_partial)):
            raise TypeError("O partial tensor must match dtype_partial")
        if const_expr(not (mO.element_type == self.dtype)):
            raise TypeError("O tensor must match dtype")
        if const_expr(mLSE_partial.element_type not in [Float32]):
            raise TypeError("LSE partial tensor must be Float32")
        if const_expr(mLSE is not None and mLSE.element_type not in [Float32]):
            raise TypeError("LSE tensor must be Float32")

        # Shape validation - input tensors are in user format, need to be converted to kernel format
        if const_expr(len(mO_partial.shape) not in [4, 5]):
            raise ValueError(
                "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)"
            )
        if const_expr(len(mLSE_partial.shape) not in [3, 4]):
            raise ValueError(
                "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)"
            )
        if const_expr(len(mO.shape) not in [3, 4]):
            raise ValueError(
                "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)"
            )
        if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]):
            raise ValueError(
                "LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)"
            )

        mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)]
        # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b)
        # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h)

View on GitHub (pinned to 0132848349)

Solutions

1. Ensure mO_partial keeps the leading num_splits dimension (4D batched or 5D varlen)
2. If you squeezed it, reshape back: o_partial.unsqueeze(0) when num_splits==1

Example fix

// before
combine(o_3d, lse_partial, mO)  # o_3d: (b, s, h, d)
// after
combine(o_4d, lse_partial, mO)  # o_4d: (num_splits, b, s, h, d)

assert mO_partial.dim() in (4, 5), mO_partial.shape

def valid_o_partial(t) -> bool: return t.dim() in (4, 5)

When it happens

Trigger: Passing mO_partial with a rank other than 4 or 5, e.g. a 3D (batch, seqlen, dim) tensor without the num_splits leading axis.

Common situations: Feeding raw attention output (3D) instead of per-split partials; dropping the num_splits dimension via squeeze/reshape before combining.

Related errors

• LSE partial tensor must have 3 or 4 dimensions: (num_splits,
• O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads
• LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nhea
• v_cache must be provided
• q can only be None when only_qv=True

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