sgl-project/sglang · error · ValueError

Unsupported layout for models with head_dim != v_head_dim: {

Error message

Unsupported layout for models with head_dim != v_head_dim: {self.layout}; expected 'page_first' or 'page_first_direct'.

What it means

In the MLA variant with head_dim != v_head_dim, K and V buffers need different shapes/strides; init_kv_buffer only knows how to allocate them for 'page_first' and 'page_first_direct'. Any other layout (notably 'layer_first') cannot represent the asymmetric K/V dims.

Source

Thrown at python/sglang/srt/mem_cache/pool_host/mha.py:1103

            k_dims = (self.size, self.layer_num, self.head_num, self.head_dim)
            v_dims = (self.size, self.layer_num, self.head_num, self.v_head_dim)
        elif self.layout == "page_first_direct":
            k_dims = (
                self.page_num,
                self.layer_num,
                self.page_size,
                self.head_num,
                self.head_dim,
            )
            v_dims = (
                self.page_num,
                self.layer_num,
                self.page_size,
                self.head_num,
                self.v_head_dim,
            )
        else:
            raise ValueError(
                f"Unsupported layout for models with head_dim != v_head_dim: "
                f"{self.layout}; expected 'page_first' or 'page_first_direct'."
            )

        # token_stride_size / layout_dim are intentionally NOT set: K and V
        # have different strides, so any caller that reaches for a single
        # shared stride is a bug. Such callers will fail loudly with
        # AttributeError rather than silently use the K stride for V copies.

        alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device]
        k_buffer = alloc_func(
            k_dims,
            dtype=self.dtype,
            device=self.device,
            pin_memory=self.pin_memory,
            allocator=self.allocator,
        )
        v_buffer = alloc_func(

View on GitHub (pinned to 0132848349)

Solutions

  1. Use layout='page_first' (recommended, works with kernel backend) or 'page_first_direct'
  2. Remove any explicit layer_first host-layout override in server args / hicache config
  3. Check model config: if head_dim == v_head_dim you can use the standard MHA pool with more layouts

Example fix

# before
pool = MLATokenToKVPoolHost(..., layout="layer_first")
# after
pool = MLATokenToKVPoolHost(..., layout="page_first")
Defensive patterns

Strategy: validation

Validate before calling

if getattr(pool, "head_dim", None) != getattr(pool, "v_head_dim", None):
    assert pool.layout in ("page_first", "page_first_direct"), "asymmetric head dims require page-first layouts"

Type guard

def mla_pool_layout_ok(pool) -> bool:
    if pool.head_dim == pool.v_head_dim:
        return True
    return pool.layout in ("page_first", "page_first_direct")

Prevention

When it happens

Trigger: Constructing the head_dim!=v_head_dim host pool class with layout='layer_first' or an unknown value; init_kv_buffer then refuses to allocate k_buffer/v_buffer.

Common situations: Running MLA-style or hybrid models with a legacy layer-first host layout config; carrying over configs from MHA models to MLA models.

Related errors


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