sgl-project/sglang · error · RuntimeError

hd256 forward varlen expects k rank 3 or 5, got rank {k_rank

Error message

hd256 forward varlen expects k rank 3 or 5, got rank {k_rank}

What it means

The hd256 (head-dim 256) 2-CTA FMHA forward varlen kernel accepts K tensors of rank 3 (token-major [total_tokens, Hk, D]) or rank 5 (batched [B, G, S, Hk, D]) only. Any other rank is rejected at compile-time of the CUTLASS JitExecutor call. It means the caller passed a K layout the varlen kernel cannot interpret.

Source

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

                d = mQ.shape[4]
            elif cutlass.const_expr(q_rank == 4):
                s_q = mQ.shape[1]
                h_q = mQ.shape[2]
                d = mQ.shape[3]
            else:
                raise RuntimeError(
                    f"hd256 forward non-varlen expects q rank 4 or 5, got rank {q_rank}"
                )

        if cutlass.const_expr(cum_seqlen_k is not None):
            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:

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape k to rank 3 [total_tokens, Hk, D] when using varlen (cum_seqlen) path
2. Or pass rank 5 [B, G, S, Hk, D] layout
3. Drop cum_seqlens to use the non-varlen path which accepts rank 4/5
4. Check the caller in the attention backend that packs KV for hd256 models

Example fix

# before
out = fmha(q, k, k, cu_seqlens_q, cu_seqlens_k)  # k is [B, S, Hk, D]
# after
k = k.reshape(-1, k.shape[-2], k.shape[-1])  # [total_tokens, Hk, D]
out = fmha(q, k, k, cu_seqlens_q, cu_seqlens_k)

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

def is_valid_varlen_k(k: torch.Tensor, varlen: bool) -> bool:
    return k.ndim in (3, 5) if varlen else k.ndim in (4, 5)

When it happens

Trigger: Calling the sm100 hd256 2cta fmha forward with cum_seqlens present and a k tensor whose ndim is not 3 or 5 (e.g. rank-4 [B, S, Hk, D] with varlen path, or a squeezed/expanded tensor).

Common situations: Passing a non-varlen shaped [B,S,H,D] K while still supplying cum_seqlens; upstream attention backend reshaping KV cache incorrectly; head-dim-256 model (e.g. GQA-256) misconfigured.

Related errors

• hd256 forward non-varlen expects k rank 4 or 5, got rank {k_
• SplitKV partial output (mO) must be Float32
• All tensors must have the same data type
• Only Float16 or BFloat16 is supported
• LSE tensor must be Float32

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