wehub-resource-sync
94 lines
3.6 KiB
Python
94 lines
3.6 KiB
Python
"""Attention KV cache modeling."""
|
|
|
|
import json
|
|
from typing import Any, Dict, List, Literal, Optional, Union # noqa: UP035
|
|
|
|
import numpy as np
|
|
from tvm import relax as rx
|
|
from tvm import tirx
|
|
from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache as TVMPagedKVCache
|
|
from tvm.relax.frontend.nn.llm.kv_cache import RopeMode
|
|
|
|
|
|
class PagedKVCache(TVMPagedKVCache):
|
|
"""The Paged KV Cache used in LLM batching for efficient attention computation."""
|
|
|
|
@staticmethod
|
|
def create_generic(
|
|
attn_kind: Union[Literal["mha", "mla"], List[Literal["mha", "mla", "mha_sliding"]]], # noqa: UP006
|
|
max_batch_size: tirx.Var,
|
|
max_total_seq_len: tirx.Var,
|
|
prefill_chunk_size: tirx.Var,
|
|
page_size: tirx.Var,
|
|
support_sliding_window: tirx.Var,
|
|
num_hidden_layers: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
qk_head_dim: int,
|
|
v_head_dim: int,
|
|
rope_mode: RopeMode,
|
|
rope_scale: int,
|
|
rope_theta: int,
|
|
dtype: str,
|
|
mla_original_qk_head_dim: int = 0,
|
|
mla_original_v_head_dim: int = 0,
|
|
rotary_dim: Optional[int] = None,
|
|
rope_scaling: Optional[Dict[str, Any]] = None, # noqa: UP006
|
|
rope_ext_factors: Optional[List[int]] = None, # noqa: UP006
|
|
layer_partition: Optional[List[int]] = None, # noqa: UP006
|
|
enable_disaggregation: bool = False,
|
|
name: str = "paged_kv_cache",
|
|
) -> "PagedKVCache":
|
|
"""The generic function of creating a multi-head attention PagedKVCache,
|
|
which will be rewritten by functions in compilation pipeline.
|
|
"""
|
|
if rotary_dim is None:
|
|
rotary_dim = qk_head_dim
|
|
if rope_scaling is None:
|
|
rope_scaling = {}
|
|
if layer_partition is None:
|
|
layer_partition = [0, num_hidden_layers]
|
|
if isinstance(attn_kind, List): # noqa: UP006
|
|
rx_attn_kind = [rx.StringImm(layer_kind) for layer_kind in attn_kind]
|
|
else:
|
|
rx_attn_kind = rx.StringImm(attn_kind)
|
|
return PagedKVCache(
|
|
_expr=rx.call_pure_packed(
|
|
"mlc.create_paged_kv_cache_generic",
|
|
rx_attn_kind,
|
|
rx.ShapeExpr(
|
|
[
|
|
max_batch_size,
|
|
max_total_seq_len,
|
|
prefill_chunk_size,
|
|
page_size,
|
|
support_sliding_window,
|
|
]
|
|
),
|
|
rx.ShapeExpr(layer_partition),
|
|
rx.prim_value(num_hidden_layers),
|
|
rx.prim_value(num_attention_heads),
|
|
rx.prim_value(num_key_value_heads),
|
|
rx.prim_value(qk_head_dim),
|
|
rx.prim_value(v_head_dim),
|
|
rx.prim_value(mla_original_qk_head_dim),
|
|
rx.prim_value(mla_original_v_head_dim),
|
|
rx.prim_value(rope_mode),
|
|
rx.prim_value(rope_scale),
|
|
rx.prim_value(rope_theta),
|
|
rx.StringImm(json.dumps(rope_scaling)),
|
|
(
|
|
rx.const(np.array(rope_ext_factors, "float32"))
|
|
if rope_ext_factors is not None
|
|
else rx.prim_value(0)
|
|
# NOTE: since relax does not have "Optional" type, we use prim_value(0)
|
|
# to represent "undefined".
|
|
),
|
|
rx.prim_value(rotary_dim),
|
|
rx.prim_value(int(enable_disaggregation)),
|
|
rx.DataTypeImm(dtype),
|
|
ty_args=rx.ObjectType(),
|
|
),
|
|
_name=name,
|
|
)
|