sgl-project/sglang · error · NotImplementedError

MHATokenToKOnlyPool: use set_index_k_buffer on the parent Mi

Error message

MHATokenToKOnlyPool: use set_index_k_buffer on the parent MiniMaxSparseKVPool — this pool does not store V

What it means

MHATokenToKOnlyPool.set_kv_buffer is intentionally disabled: writing K (and V) for K-only sparse layers must go through MiniMaxSparseKVPool.set_index_k_buffer, which handles dtype conversion and the layer-id remapping. Calling the generic setter directly is a misuse of the sub-pool.

Source

Thrown at python/sglang/srt/mem_cache/memory_pool.py:4799

    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
        raise NotImplementedError("MHATokenToKOnlyPool does not allocate V")

    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
        raise NotImplementedError("MHATokenToKOnlyPool does not allocate V")

    def set_kv_buffer(
        self,
        layer: RadixAttention,
        loc: torch.Tensor,
        cache_k: torch.Tensor,
        cache_v: torch.Tensor,
        k_scale: Optional[float] = None,
        v_scale: Optional[float] = None,
        layer_id_override: Optional[int] = None,
    ) -> None:
        # Routed through MiniMaxSparseKVPool.set_index_k_buffer instead.
        raise NotImplementedError(
            "MHATokenToKOnlyPool: use set_index_k_buffer on the parent "
            "MiniMaxSparseKVPool — this pool does not store V"
        )

    def get_kv_size_bytes(self):
        k_size_bytes = sum(get_tensor_size_bytes(k) for k in self.k_buffer)
        return k_size_bytes, 0


class MiniMaxSparseKVPool(KVCache):
    def __init__(
        self,
        size: int,
        page_size: int,
        dtype: torch.dtype,
        head_num: int,
        head_dim: int,
        idx_head_dim: int,

View on GitHub (pinned to 0132848349)

Solutions

  1. Call MiniMaxSparseKVPool.set_index_k_buffer(layer, loc, cache_idx_k, k_scale) for K-only layers
  2. Ensure the model's attention layers call the parent pool's set_kv_buffer dispatcher, which routes K-only layers correctly

Example fix

// before
k_only_pool.set_kv_buffer(layer, loc, cache_k, cache_v)
// after
sparse_pool.set_index_k_buffer(layer, loc, cache_k, k_scale=layer.k_scale)
Defensive patterns

Strategy: validation

Validate before calling

# always write through the parent dispatcher
sparse_pool.set_kv_buffer(layer, loc, cache_k, cache_v, k_scale=layer.k_scale, v_scale=layer.v_scale)

Try / catch

try:
    sub_pool.set_kv_buffer(layer, loc, cache_k, cache_v)
except NotImplementedError as e:
    raise RuntimeError('use MiniMaxSparseKVPool.set_index_k_buffer for K-only layers') from e

Prevention

When it happens

Trigger: Calling set_kv_buffer(layer, loc, cache_k, cache_v, ...) directly on an MHATokenToKOnlyPool instead of on the parent MiniMaxSparseKVPool or its set_index_k_buffer entry point.

Common situations: Custom model code or refactored attention layers that grab the sub-pool directly; version changes that moved the set_index_k_buffer entry point.

Related errors


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