sgl-project/sglang · error · RuntimeError

hd256 forward non-varlen expects k rank 4 or 5, got rank {k_

Error message

hd256 forward non-varlen expects k rank 4 or 5, got rank {k_rank}

What it means

The hd256 2-CTA FMHA forward non-varlen (no cumulative-seqlen) path accepts K of rank 4 [B, S, Hk, D] or rank 5 [B, G, S, Hk, D]. Other ranks (e.g. rank-3 packed tokens without cum_seqlens) are rejected because batch/seq dims cannot be inferred.

Source

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

            if cutlass.const_expr(k_rank == 5):
                s_k = mK.shape[1]
                h_k = mK.shape[2]
            elif cutlass.const_expr(k_rank == 3):
                s_k = mK.shape[0]
                h_k = mK.shape[1]
            else:
                raise RuntimeError(
                    f"hd256 forward varlen expects k rank 3 or 5, got rank {k_rank}"
                )
        else:
            if cutlass.const_expr(k_rank == 5):
                s_k = mK.shape[1]
                h_k = mK.shape[2]
            elif cutlass.const_expr(k_rank == 4):
                s_k = mK.shape[1]
                h_k = mK.shape[2]
            else:
                raise RuntimeError(
                    f"hd256 forward non-varlen expects k rank 4 or 5, got rank {k_rank}"
                )
        if cutlass.const_expr(cum_seqlen_q is not None):
            b = mCuSeqlensQ.shape[0] - 1
        elif cutlass.const_expr(cum_seqlen_k is not None):
            b = mCuSeqlensK.shape[0] - 1
        else:
            b = mQ.shape[0]

        scale_softmax = softmax_scale
        scale_softmax_log2 = softmax_scale * math.log2(math.exp(1.0))
        scale_output = 1.0
        s_lse = s_q
        h_r = h_q // h_k
        s_q64 = Int64(s_q)
        s_k64 = Int64(s_k)
        s_lse64 = Int64(s_lse)
        d64 = cute.assume(Int64(d), divby=128)

View on GitHub (pinned to 0132848349)

Solutions

1. Pass rank 4 [B, S, Hk, D] or rank 5 [B, G, S, Hk, D] K
2. Or supply cu_seqlens to take the varlen path with rank-3 K
3. Verify upstream KV cache reshape logic

Example fix

# before
out = fmha(q, k_packed, v_packed)  # k_packed: [total_tokens, Hk, D], no cu_seqlens
# after
k = k_packed.view(B, S, Hk, D)
out = fmha(q, k, v)

assert k.ndim in (4, 5), f'non-varlen k must be rank 4 or 5, got {k.ndim}'

def is_valid_batched_k(k: torch.Tensor) -> bool:
    return k.ndim in (4, 5)

When it happens

Trigger: Calling the forward without cum_seqlens but with a rank-3 token-packed K tensor, or any rank other than 4/5.

Common situations: Migrating from varlen to batched path and forgetting to expand the packed tensor; passing a debug/random tensor with wrong ndim in unit tests.

Related errors

• hd256 forward varlen expects k rank 3 or 5, got rank {k_rank
• The layout of q is not supported
• The layout of k is not supported
• The layout of v is not supported
• Type mismatch: {self.q_dtype} != {self.k_dtype}

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