sgl-project/sglang · error · RuntimeError

The layout of k is not supported

Error message

The layout of k is not supported

What it means

The kernel requires K in K-major layout for the tcgen05 MMA. LayoutEnum.from_tensor(k).mma_major_mode() returned MN-major, meaning K is effectively column-major, which this kernel cannot consume.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py:429

                self.cta_tiler,
                (*self.cluster_shape_mn, 1),
            )
        else:
            self.tile_sched_params, grid = compute_grid(
                (s_q, o.shape[1], o.shape[2]) if cum_seqlen_q is not None else o.shape,
                self.cta_tiler,
                self.is_persistent,
            )

        self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode()
        self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode()
        self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode()
        self.o_layout = utils.LayoutEnum.from_tensor(o)

        if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K):
            raise RuntimeError("The layout of q is not supported")
        if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K):
            raise RuntimeError("The layout of k is not supported")
        if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.MN):
            raise RuntimeError("The layout of v is not supported")

        # check type consistency
        if cutlass.const_expr(self.q_dtype != self.k_dtype):
            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}")
        if cutlass.const_expr(self.q_dtype != self.v_dtype):
            raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}")
        self._setup_attributes()

        cta_group = tcgen05.CtaGroup.TWO
        # the intermediate tensor p is from tmem & k-major
        p_source = tcgen05.OperandSource.TMEM
        p_major_mode = tcgen05.OperandMajorMode.K
        qk_tiled_mma = sm100_utils.make_trivial_tiled_mma(
            self.q_dtype,
            self.q_major_mode,
            self.k_major_mode,

View on GitHub (pinned to 0132848349)

Solutions

  1. Pass k contiguous with unit stride in the head dim: k.contiguous()
  2. Verify the KV cache pool layout matches the kernel expectation (K-major)
  3. Fix upstream reshape/transpose of the K cache

Example fix

# before
out = fmha(q, k_cache_view.transpose(-1,-2), v)
# after
out = fmha(q, k_cache_view.contiguous(), v)
Defensive patterns

Strategy: type-guard

Validate before calling

assert k.stride(-1) == 1, 'k must be K-major'

Type guard

def k_is_k_major(k: torch.Tensor) -> bool:
    return k.stride(-1) == 1

Prevention

When it happens

Trigger: Passing a column-major (transposed) k tensor to hd256 2cta fmha forward.

Common situations: KV cache stored transposed for another backend (e.g. FA-style [B,H,S,D] vs transposed views); reusing tensors prepared for flashinfer with different layout expectations.

Related errors


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