sgl-project/sglang · error · ValueError

O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads

Error message

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

What it means

The final output tensor mO must be 3D (batch, seqlen, nheads, headdim) or 4D varlen (total_q, nheads, headdim) — i.e. one rank less than the partials, without the num_splits axis.

Source

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

            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)
        O_partial_layout_transpose = (
            [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2]
        )
        # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h)
        mO_partial = cute.make_tensor(
            mO_partial.iterator,
            cute.select(mO_partial.layout, mode=O_partial_layout_transpose),
        )

View on GitHub (pinned to 0132848349)

Solutions

Allocate mO without the leading num_splits dimension
For batched input use (b, s, h, d); for varlen use (total_q, h, d)

Example fix

// before
mO = torch.empty(num_splits, b, s, h, d, dtype=dt)
// after
mO = torch.empty(b, s, h, d, dtype=dt)

Defensive patterns

Strategy: validation

Validate before calling

assert mO.dim() in (3, 4) and mO.dim() == mO_partial.dim() - 1

Type guard

def final_o_ok(o, o_partial) -> bool:
    return o.dim() in (3, 4) and o.dim() == o_partial.dim() - 1

Prevention

Allocate output as o_partial[0].new_empty(o_partial.shape[1:], ...) to inherit correct rank

When it happens

Trigger: Passing an mO that still contains the num_splits axis, or a 2D folded output.

Common situations: Reusing the partial tensor's shape when allocating the output buffer; forgetting the combine reduces over splits.

Related errors

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