sgl-project/sglang · error · RuntimeError

The layout of mBias is wrong

Error message

The layout of mBias is wrong

What it means

The attention bias tensor passed to the SM90/SM100 kernel must be in K-major layout (row/col-major matching tcgen05 OperandMajorMode.K) so TMA can stream it into shared memory for the MMA. A transposed/mixed bias layout is rejected.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm100.py:1301

                num_bits_per_copy=universal_copy_bits,
            )
            tO_shape_dim_1 = sO_layout.outer.shape[1][0] // async_copy_elems
            tO_layout = cute.make_ordered_layout(
                (self.num_epilogue_threads // tO_shape_dim_1, tO_shape_dim_1),
                order=(1, 0),
            )
            # So that we don't have to check if we overshoot kBlockM when we store O
            assert self.m_block_size % tO_layout.shape[0] == 0
            vO_layout = cute.make_layout((1, async_copy_elems))
            gmem_tiled_copy_O = cute.make_tiled_copy_tv(
                atom_universal_copy, tO_layout, vO_layout
            )

        if const_expr(mBias is not None):
            bias_layout_enum = cutlass.utils.LayoutEnum.from_tensor(mBias)
            self.bias_major_mode = bias_layout_enum.mma_major_mode()
            if const_expr(self.bias_major_mode != tcgen05.OperandMajorMode.K):
                raise RuntimeError("The layout of mBias is wrong")
            # (bias_block_size, n_block_size, bias_stage)
            sBias_layout = sm100_utils_basic.make_smem_layout_epi(
                self.bias_dtype,
                bias_layout_enum,
                (self.bias_block_size, self.n_block_size),
                self.bias_stage,
            )
            sBias_size = cute.cosize(sBias_layout)
            # Set after the Q/K/V cta_group_size scaling loop above so bias (non-multicast) isn't double-scaled.
            self.tma_copy_bytes["bias"] = cute.size_in_bytes(
                self.bias_dtype, cute.select(sBias_layout, mode=[0, 1])
            )
            tma_atom_bias, mBias = cpasync.make_tiled_tma_atom(
                cpasync.CopyBulkTensorTileG2SOp(),
                mBias,
                cute.select(sBias_layout, mode=[0, 1]),
                (self.bias_block_size, self.n_block_size),
                1,  # no mcast

View on GitHub (pinned to 0132848349)

Solutions

  1. Materialize the bias contiguous in K-major (row-major over the K dimension): ensure last-dim stride 1 along the K/head-k axis
  2. Avoid .transpose()/.T on the bias; build it directly in the correct orientation
  3. Call .contiguous() after any transpose

Example fix

// before
bias = raw_bias.T  # q-major after transpose
attn(..., bias=bias)
// after
bias = raw_bias.T.contiguous()  # or build K-major directly
attn(..., bias=bias)
Defensive patterns

Strategy: validation

Validate before calling

assert bias.stride(-1) == 1 and is_bias_k_major(bias)  # K axis last & contiguous

Type guard

def bias_k_major(bias) -> bool:
    # K-major: last dim (K) has stride 1 and tensor is contiguous
    return bias.is_contiguous() and bias.stride(-1) == 1

Prevention

When it happens

Trigger: Passing an mBias tensor whose memory layout is not K-major (e.g. transposed bias or a non-standard stride order) with attention bias / ALiBi-style score bias enabled.

Common situations: Creating bias tensors via .transpose() or .permute() (producing non-contiguous layouts), or materializing bias in (d_k, q) orientation instead of (q, k).

Related errors


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