wehub-resource-sync
687 lines
32 KiB
Python
687 lines
32 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, RUF012
|
|
# fmt: off
|
|
|
|
"""Attention KV cache modeling.
|
|
|
|
This module exposes the public ``PagedKVCache`` classes (``FlashInferPagedKVCache``
|
|
and ``TIRPagedKVCache``). The kernel factories that build the underlying TIR
|
|
functions are split across sibling private modules:
|
|
|
|
- ``_kernel_common``: shared helpers (enums, RoPE, mask, tile allocators,
|
|
``@T.macro`` bundle, tiling config, scheduling).
|
|
- ``_page_kernels``: page management (append, debug, copy, compact).
|
|
- ``_prefill_kernels``: prefill attention kernels (paged/ragged/MLA/dense).
|
|
- ``_decode_kernels``: decode attention kernels and state-merge helpers.
|
|
|
|
The private-named kernel factories are re-exported from this module so the
|
|
test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache``.
|
|
"""
|
|
|
|
# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
|
|
import math
|
|
from typing import Any, Literal
|
|
|
|
import tvm
|
|
from tvm import relax as rx
|
|
from tvm import tirx
|
|
from tvm.relax.frontend.nn import Object, Tensor
|
|
from tvm.target import Target
|
|
|
|
# Re-export enums + kernel factories so existing ``from kv_cache import ...``
|
|
# users (test suite, tree_attn.py, mlc-llm, etc.) continue to work after the
|
|
# split. These names are referenced in ``__all__`` below to signal to linters
|
|
# that the imports are intentional public API (not dead code).
|
|
from ._decode_kernels import (
|
|
_attention_decode,
|
|
_attention_decode_cpu,
|
|
_merge_state_inplace,
|
|
_merge_state_inplace_cpu,
|
|
)
|
|
from ._kernel_common import AttnKind, RopeMode
|
|
from ._page_kernels import (
|
|
_compact_kv_copy,
|
|
_compact_kv_copy_cpu,
|
|
_copy_single_page,
|
|
_copy_single_page_cpu,
|
|
_copy_single_page_mla,
|
|
_kv_cache_debug_get_kv,
|
|
_kv_cache_debug_get_kv_mla,
|
|
_kv_cache_transpose_append,
|
|
_kv_cache_transpose_append_mla,
|
|
)
|
|
from ._prefill_kernels import (
|
|
_attention_prefill,
|
|
_attention_prefill_cpu,
|
|
_attention_prefill_mla,
|
|
_attention_prefill_ragged,
|
|
_attention_prefill_ragged_cpu,
|
|
_attention_sequence_prefill,
|
|
_attention_sequence_prefill_with_mask,
|
|
)
|
|
from .position_embedding import llama_rope_with_position_map
|
|
from .tree_attn import (
|
|
tree_attn,
|
|
tree_attn_cpu,
|
|
tree_attn_with_paged_kv_cache,
|
|
tree_attn_with_paged_kv_cache_cpu,
|
|
)
|
|
|
|
__all__ = [
|
|
"AttnKind",
|
|
"FlashInferPagedKVCache",
|
|
"PagedKVCache",
|
|
"RopeMode",
|
|
"TIRPagedKVCache",
|
|
"_attention_decode",
|
|
"_attention_decode_cpu",
|
|
"_attention_prefill",
|
|
"_attention_prefill_cpu",
|
|
"_attention_prefill_mla",
|
|
"_attention_prefill_ragged",
|
|
"_attention_prefill_ragged_cpu",
|
|
"_attention_sequence_prefill",
|
|
"_attention_sequence_prefill_with_mask",
|
|
"_compact_kv_copy",
|
|
"_compact_kv_copy_cpu",
|
|
"_copy_single_page",
|
|
"_copy_single_page_cpu",
|
|
"_copy_single_page_mla",
|
|
"_kv_cache_debug_get_kv",
|
|
"_kv_cache_debug_get_kv_mla",
|
|
"_kv_cache_transpose_append",
|
|
"_kv_cache_transpose_append_mla",
|
|
"_merge_state_inplace",
|
|
"_merge_state_inplace_cpu",
|
|
"llama_rope_with_position_map",
|
|
"tree_attn",
|
|
"tree_attn_cpu",
|
|
"tree_attn_with_paged_kv_cache",
|
|
"tree_attn_with_paged_kv_cache_cpu",
|
|
]
|
|
|
|
|
|
class PagedKVCache(Object): # pylint: disable=too-few-public-methods
|
|
"""The Paged KV Cache used in LLM batching for efficient attention computation."""
|
|
|
|
extern_mods: list[tvm.runtime.Module] = []
|
|
|
|
def attention_with_fused_qkv(
|
|
self,
|
|
layer_id: int,
|
|
qkv: Tensor,
|
|
num_qo_heads: int,
|
|
sm_scale: float,
|
|
) -> Tensor:
|
|
"""Compute attention with the given fused q/k/v data and in-cache k/v data
|
|
on the specified layer. Rotary position embeddings are applied to k/v
|
|
within this function.
|
|
|
|
- For prefill, the input qkv and output tensor have shape
|
|
(1, total_seq_len) for the first two dimensions.
|
|
- For decode, the input qkv and output tensor have shape
|
|
(batch_size, 1) for the first two dimensions.
|
|
- The input qkv have `2 * num_qo_heads + num_kv_heads` at the third dim.
|
|
- The output tensor have `num_qo_heads` at the third dim.
|
|
- The input qkv and output tensor have `head_dim` at the last dim.
|
|
"""
|
|
# pylint: disable=protected-access
|
|
b, s, _, d = qkv._expr.ty.shape
|
|
qkv = qkv.reshape(b * s, qkv.shape[2], d)
|
|
return Tensor(
|
|
_expr=rx.BlockBuilder.current().emit(
|
|
rx.call_dps_packed(
|
|
"vm.builtin.attention_kv_cache_attention_with_fused_qkv",
|
|
[
|
|
self._expr,
|
|
rx.prim_value(layer_id), # type: ignore[arg-type]
|
|
rx.prim_value(sm_scale),
|
|
qkv._expr,
|
|
],
|
|
out_ty=rx.TensorType((b * s, num_qo_heads, d), qkv.dtype),
|
|
)
|
|
)
|
|
).reshape(b, s, num_qo_heads, d)
|
|
|
|
def self_attention( # pylint: disable=too-many-locals
|
|
self,
|
|
layer_id: int,
|
|
q: Tensor,
|
|
k: Tensor,
|
|
v: Tensor,
|
|
sm_scale: float,
|
|
) -> tuple[Tensor, Tensor]:
|
|
"""Fine-grained API that computes ragged self attention with Q/K/V data."""
|
|
# pylint: disable=protected-access
|
|
b, s, h_qo, d_qk = q._expr.ty.shape
|
|
_, _, h_kv, d_v = v._expr.ty.shape
|
|
q = q.reshape(b * s, h_qo, d_qk)
|
|
k = k.reshape(b * s, h_kv, d_qk)
|
|
v = v.reshape(b * s, h_kv, d_v)
|
|
bb = rx.BlockBuilder.current()
|
|
attn_results = bb.emit(
|
|
rx.call_dps_packed(
|
|
"vm.builtin.attention_kv_cache_self_attention",
|
|
[
|
|
self._expr,
|
|
rx.prim_value(layer_id), # type: ignore[arg-type]
|
|
rx.prim_value(sm_scale),
|
|
q._expr,
|
|
k._expr,
|
|
v._expr,
|
|
],
|
|
out_ty=[
|
|
rx.TensorType((b * s, h_qo, d_v), q.dtype),
|
|
rx.TensorType((b * s, h_qo), "float32"),
|
|
],
|
|
)
|
|
)
|
|
assert isinstance(attn_results.ty, rx.TupleType)
|
|
assert len(attn_results.ty.fields) == 2
|
|
o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, d_v)
|
|
lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
|
|
return o, lse
|
|
|
|
def cross_attention(
|
|
self,
|
|
layer_id: int,
|
|
q: Tensor,
|
|
v_head_dim: int,
|
|
sm_scale: float,
|
|
) -> tuple[Tensor, Tensor]:
|
|
"""Fine-grained API that computes paged cross attention with Q and in-cache KV data."""
|
|
# pylint: disable=protected-access
|
|
b, s, h_qo, d_qk = q._expr.ty.shape
|
|
q = q.reshape(b * s, h_qo, d_qk)
|
|
bb = rx.BlockBuilder.current()
|
|
attn_results = bb.emit(
|
|
rx.call_dps_packed(
|
|
"vm.builtin.attention_kv_cache_cross_attention",
|
|
[
|
|
self._expr,
|
|
rx.prim_value(layer_id), # type: ignore[arg-type]
|
|
rx.prim_value(sm_scale),
|
|
q._expr,
|
|
],
|
|
out_ty=[
|
|
rx.TensorType((b * s, h_qo, v_head_dim), q.dtype),
|
|
rx.TensorType((b * s, h_qo), "float32"),
|
|
],
|
|
)
|
|
)
|
|
assert isinstance(attn_results.ty, rx.TupleType)
|
|
assert len(attn_results.ty.fields) == 2
|
|
o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, v_head_dim)
|
|
lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
|
|
return o, lse
|
|
|
|
def append_mla_kv(self, layer_id: int, kv: Tensor) -> "PagedKVCache":
|
|
"""Fine-grained API that appends the MLA K/V data to KV cache."""
|
|
# pylint: disable=protected-access
|
|
b, s, _, d_qk = kv._expr.ty.shape
|
|
kv = kv.reshape(b * s, d_qk)
|
|
return PagedKVCache(
|
|
_expr=rx.call_pure_packed(
|
|
"vm.builtin.attention_kv_cache_append_mla_kv",
|
|
self._expr,
|
|
rx.prim_value(layer_id), # type: ignore[arg-type]
|
|
kv._expr,
|
|
ty_args=rx.AnyType(),
|
|
),
|
|
_name="paged_kv_cache",
|
|
)
|
|
|
|
def merge_attn_output_inplace(
|
|
self,
|
|
o_self_attn: Tensor,
|
|
lse_self_attn: Tensor,
|
|
o_cross_attn: Tensor,
|
|
lse_cross_attn: Tensor,
|
|
) -> tuple[Tensor, Tensor]:
|
|
"""Fine-grained API that merges the attention output from two sources.
|
|
The first two tensors will be inplace updated.
|
|
"""
|
|
# pylint: disable=protected-access
|
|
b, s, h_qo, d_v = o_self_attn._expr.ty.shape
|
|
o_self_attn = o_self_attn.reshape(b * s, h_qo, d_v)
|
|
lse_self_attn = lse_self_attn.reshape(b * s, h_qo)
|
|
o_cross_attn = o_cross_attn.reshape(b * s, h_qo, d_v)
|
|
lse_cross_attn = lse_cross_attn.reshape(b * s, h_qo)
|
|
bb = rx.BlockBuilder.current()
|
|
merge_results = bb.emit(
|
|
rx.call_pure_packed(
|
|
"vm.builtin.attention_kv_cache_merge_attn_output_inplace",
|
|
self._expr,
|
|
o_self_attn._expr,
|
|
lse_self_attn._expr,
|
|
o_cross_attn._expr,
|
|
lse_cross_attn._expr,
|
|
ty_args=rx.TupleType(
|
|
[o_self_attn._expr.ty, lse_self_attn._expr.ty]
|
|
),
|
|
)
|
|
)
|
|
assert isinstance(merge_results.ty, rx.TupleType)
|
|
assert len(merge_results.ty.fields) == 2
|
|
o_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 0))).reshape(
|
|
b, s, h_qo, d_v
|
|
)
|
|
lse_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 1))).reshape(b, s, h_qo)
|
|
return o_self_attn, lse_self_attn
|
|
|
|
def get_query_positions(self, total_length: tirx.Expr) -> Tensor:
|
|
"""Get the in-sequence positions of each slot in the query,
|
|
which are needed for applying positional embeddings in some models.
|
|
|
|
Parameters
|
|
----------
|
|
total_length : tirx.Expr
|
|
The summed-up total sequence length of queries in
|
|
the batch being forwarded.
|
|
|
|
Returns
|
|
-------
|
|
q_positions : Tensor
|
|
The in-sequence query positions, in shape `(total_length,)`
|
|
"""
|
|
return Tensor(
|
|
_expr=rx.BlockBuilder.current().emit(
|
|
rx.call_pure_packed(
|
|
"vm.builtin.attention_kv_cache_get_query_positions",
|
|
self._expr,
|
|
ty_args=rx.TensorType((total_length,), "int32"),
|
|
)
|
|
)
|
|
)
|
|
|
|
# pylint: enable=protected-access
|
|
|
|
|
|
def _prepare_yarn_rope_scaling(rope_scaling: dict[str, Any] | None, rope_theta: float | None) -> dict[str, Any] | None:
|
|
"""Ensure Yarn-specific scaling configs include the theta metadata."""
|
|
if rope_scaling is None:
|
|
return None
|
|
if rope_scaling.get("rope_type") != "yarn":
|
|
return rope_scaling
|
|
|
|
rope_scaling_updated = dict(rope_scaling)
|
|
if "inv_theta_log_scale" not in rope_scaling_updated and rope_theta is not None:
|
|
theta_value = float(rope_theta)
|
|
rope_scaling_updated["inv_theta_log_scale"] = 1.0 / (2 * math.log(theta_value))
|
|
return rope_scaling_updated
|
|
|
|
|
|
class FlashInferPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-methods
|
|
"""Paged KV cache using FlashInfer (CUDA) kernels."""
|
|
|
|
def __init__( # pylint: disable=too-many-locals
|
|
self,
|
|
attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
|
|
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,
|
|
layer_partition: rx.ShapeExpr,
|
|
num_hidden_layers: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
qk_head_dim: int,
|
|
v_head_dim: int,
|
|
mla_original_qk_head_dim: int,
|
|
mla_original_v_head_dim: int,
|
|
rope_mode: RopeMode,
|
|
rope_scale: int,
|
|
rope_theta: int,
|
|
rope_scaling: dict[str, Any],
|
|
rope_ext_factors: rx.Expr,
|
|
rotary_dim: int,
|
|
enable_disaggregation: bool,
|
|
dtype: str,
|
|
target: Target,
|
|
name: str = "paged_kv_cache",
|
|
) -> None:
|
|
"""Create a paged KV cache object with FlashInfer kernels.
|
|
|
|
Parameters
|
|
----------
|
|
max_batch_size : tirx.Var
|
|
The maximum allowed batch size of the KV cache.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
max_total_seq_len : tirx.Var
|
|
The maximum allowed total sequence length of the KV cache.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
prefill_chunk_size : tirx.Var
|
|
The maximum total sequence length in a prefill.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
page_size : tirx.Var
|
|
The size (a.k.a. number of tokens) of each page.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
support_sliding_window : tirx.Var
|
|
0 or 1, denoting whether the KV cache supports sliding window.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
layer_partition : rx.ShapeExpr
|
|
The KV cache layer partition for pipeline stages.
|
|
It is an indptr array, denoting the starting layer of each pipeline stage.
|
|
rope_mode : RopeMode
|
|
The RoPE mode of the Paged KV cache.
|
|
If it is normal, RoPE will be applied to k before adding k to cache.
|
|
Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
|
|
rope_scale : int
|
|
The scale of rotary position embedding.
|
|
rope_theta : int
|
|
The base of rotary position embedding.
|
|
rope_scaling: Dict[str, Any]
|
|
The RoPE scaling information dict.
|
|
rope_ext_factors: rx.Expr
|
|
The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
|
|
rotary_dim : int
|
|
The number of dimensions in the embedding that RoPE is applied to.
|
|
enable_disaggregation : bool
|
|
Whether to enable disaggregation in the KV cache.
|
|
"""
|
|
assert rope_mode != RopeMode.INLINE, "FlashInfer RoPE does not support inline mode."
|
|
rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
|
|
|
|
attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
|
|
if attn_kind_single == "mha_sliding":
|
|
attn_kind_single = "mha"
|
|
flashinfer_prefill_mods = rx.backend.cuda.flashinfer.gen_flashinfer_prefill_module(
|
|
dtype_q=dtype,
|
|
dtype_kv=dtype,
|
|
dtype_o=dtype,
|
|
qk_head_dim=(qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim),
|
|
v_head_dim=(v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim),
|
|
enable_inline_rope=False,
|
|
return_static_libs=True,
|
|
)
|
|
flashinfer_decode_mods = (
|
|
rx.backend.cuda.flashinfer.gen_flashinfer_decode_module(
|
|
dtype_q=dtype,
|
|
dtype_kv=dtype,
|
|
dtype_o=dtype,
|
|
qk_head_dim=qk_head_dim,
|
|
v_head_dim=v_head_dim,
|
|
enable_inline_rope=False,
|
|
return_static_libs=True,
|
|
)
|
|
if attn_kind_single == "mha"
|
|
else []
|
|
)
|
|
flashinfer_mla_mods = (
|
|
rx.backend.cuda.flashinfer.gen_flashinfer_mla_module(
|
|
dtype_q=dtype,
|
|
dtype_kv=dtype,
|
|
dtype_o=dtype,
|
|
head_dim_ckv=v_head_dim,
|
|
head_dim_kpe=qk_head_dim - v_head_dim,
|
|
return_static_libs=True,
|
|
)
|
|
if attn_kind_single == "mla"
|
|
else []
|
|
)
|
|
self.extern_mods = flashinfer_prefill_mods + flashinfer_decode_mods + flashinfer_mla_mods
|
|
|
|
bb = rx.BlockBuilder.current()
|
|
mha_functions = (
|
|
[
|
|
rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
|
|
rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
|
|
]
|
|
if attn_kind_single == "mha"
|
|
else [rx.Tuple([]) for _ in range(6)]
|
|
)
|
|
ragged_prefill_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan")]) if attn_kind_single == "mha" else rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan"), rx.prim_value(mla_original_qk_head_dim), rx.prim_value(mla_original_v_head_dim)])
|
|
mla_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_mla_run"), rx.ExternFunc("batch_mla_plan")] if attn_kind_single == "mla" else [])
|
|
attn_merge_functions = [
|
|
bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
|
|
]
|
|
if attn_kind_single == "mla":
|
|
attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
|
|
|
|
if isinstance(attn_kind, list):
|
|
attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
|
|
else:
|
|
attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
|
|
|
|
args = [
|
|
rx.ShapeExpr(
|
|
[
|
|
max_batch_size,
|
|
max_total_seq_len,
|
|
prefill_chunk_size,
|
|
page_size,
|
|
support_sliding_window,
|
|
]
|
|
),
|
|
layer_partition,
|
|
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.ShapeExpr(attn_kind),
|
|
rx.prim_value(enable_disaggregation),
|
|
rx.prim_value(rope_mode),
|
|
rx.prim_value(rope_scale),
|
|
rx.prim_value(rope_theta),
|
|
rope_ext_factors,
|
|
rx.op.zeros((), dtype),
|
|
bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
|
|
bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
|
|
ragged_prefill_function,
|
|
*mha_functions,
|
|
mla_function,
|
|
rx.Tuple(attn_merge_functions),
|
|
bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
|
|
bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
|
|
bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
|
|
bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
|
|
]
|
|
super().__init__(
|
|
_expr=rx.call_pure_packed(
|
|
"vm.builtin.paged_attention_kv_cache_create",
|
|
*args,
|
|
ty_args=rx.AnyType(),
|
|
),
|
|
_name=name,
|
|
)
|
|
|
|
|
|
class TIRPagedKVCache(PagedKVCache): # pylint: disable=too-few-public-methods
|
|
"""Paged KV cache using TIR kernels."""
|
|
|
|
def __init__( # pylint: disable=too-many-locals
|
|
self,
|
|
attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
|
|
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,
|
|
layer_partition: rx.ShapeExpr,
|
|
num_hidden_layers: int,
|
|
num_attention_heads: int,
|
|
num_key_value_heads: int,
|
|
qk_head_dim: int,
|
|
v_head_dim: int,
|
|
mla_original_qk_head_dim: int,
|
|
mla_original_v_head_dim: int,
|
|
rope_mode: RopeMode,
|
|
rope_scale: int,
|
|
rope_theta: int,
|
|
rope_scaling: dict[str, Any],
|
|
rope_ext_factors: rx.Expr,
|
|
rotary_dim: int,
|
|
enable_disaggregation: bool,
|
|
dtype: str,
|
|
target: Target,
|
|
name: str = "paged_kv_cache",
|
|
) -> None:
|
|
"""Create a paged KV cache object with TIR kernels.
|
|
|
|
Parameters
|
|
----------
|
|
max_batch_size : tirx.Var
|
|
The maximum allowed batch size of the KV cache.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
max_total_seq_len : tirx.Var
|
|
The maximum allowed total sequence length of the KV cache.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
prefill_chunk_size : tirx.Var
|
|
The maximum total sequence length in a prefill.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
page_size : tirx.Var
|
|
The size (a.k.a. number of tokens) of each page.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
support_sliding_window : tirx.Var
|
|
0 or 1, denoting whether the KV cache supports sliding window.
|
|
It is a symbolic variable whose concrete value is specified
|
|
at runtime.
|
|
layer_partition : rx.ShapeExpr
|
|
The KV cache layer partition for pipeline stages.
|
|
It is an indptr array, denoting the starting layer of each pipeline stage.
|
|
rope_mode : RopeMode
|
|
The RoPE mode of the Paged KV cache.
|
|
If it is normal, RoPE will be applied to k before adding k to cache.
|
|
Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
|
|
rope_scale : int
|
|
The scale of rotary position embedding.
|
|
rope_theta : int
|
|
The base of rotary position embedding.
|
|
rope_scaling: Dict[str, Any]
|
|
The RoPE scaling information dict.
|
|
rope_ext_factors: rx.Expr
|
|
The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
|
|
rotary_dim : int
|
|
The number of dimensions in the embedding that RoPE is applied to.
|
|
enable_disaggregation : bool
|
|
Whether to enable disaggregation in the KV cache.
|
|
target : Target
|
|
The target to build the model to.
|
|
"""
|
|
rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
|
|
attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
|
|
if attn_kind_single == "mha_sliding":
|
|
attn_kind_single = "mha"
|
|
if isinstance(attn_kind, list):
|
|
attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
|
|
else:
|
|
attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
|
|
bb = rx.BlockBuilder.current()
|
|
args = [
|
|
rx.ShapeExpr(
|
|
[
|
|
max_batch_size,
|
|
max_total_seq_len,
|
|
prefill_chunk_size,
|
|
page_size,
|
|
support_sliding_window,
|
|
]
|
|
),
|
|
layer_partition,
|
|
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.ShapeExpr(attn_kind),
|
|
rx.prim_value(enable_disaggregation),
|
|
rx.prim_value(rope_mode),
|
|
rx.prim_value(rope_scale),
|
|
rx.prim_value(rope_theta),
|
|
rope_ext_factors,
|
|
rx.op.zeros((), dtype),
|
|
bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
|
|
bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
|
|
]
|
|
|
|
if target.kind.name == "llvm":
|
|
if attn_kind_single == "mla":
|
|
raise ValueError("MLA is not supported in TIR kernels for now.")
|
|
args.extend(
|
|
[
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_prefill_cpu_sliding_window")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
|
|
rx.Tuple([]), # f_mla_prefill
|
|
rx.Tuple([bb.add_func(_merge_state_inplace_cpu(dtype), "tir_attention_merge_state_cpu")]),
|
|
bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
|
|
bb.add_func(_copy_single_page_cpu(num_key_value_heads, page_size, qk_head_dim, dtype), "kv_cache_copy_single_page_cpu"),
|
|
bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
|
|
bb.add_func(_compact_kv_copy_cpu(num_key_value_heads, qk_head_dim, dtype), "kv_cache_compact_kv_copy_cpu"),
|
|
]
|
|
)
|
|
else:
|
|
ragged_qk_head_dim = qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim
|
|
ragged_v_head_dim = v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim
|
|
args.append(rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged(num_key_value_heads if attn_kind_single == "mha" else num_attention_heads, num_attention_heads, ragged_qk_head_dim, ragged_v_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_ragged")]))
|
|
mha_functions = (
|
|
[
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
|
|
rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
|
|
]
|
|
if attn_kind_single == "mha"
|
|
else [rx.Tuple([]) for _ in range(6)]
|
|
)
|
|
mla_function = rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_mla(num_attention_heads, v_head_dim, qk_head_dim - v_head_dim, dtype, False, target), "tir_attention_prefill_mla")] if attn_kind_single == "mla" else [])
|
|
attn_merge_functions = [
|
|
bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
|
|
]
|
|
if attn_kind_single == "mla":
|
|
attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
|
|
args.extend(mha_functions)
|
|
args.append(mla_function)
|
|
args.extend(
|
|
[
|
|
rx.Tuple(attn_merge_functions),
|
|
bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
|
|
bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
|
|
bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
|
|
bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
|
|
]
|
|
)
|
|
|
|
super().__init__(
|
|
_expr=rx.call_pure_packed(
|
|
"vm.builtin.paged_attention_kv_cache_create",
|
|
*args,
|
|
ty_args=rx.AnyType(),
|
|
),
|
|
_name=name,
|
|
)
|