This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new bf7bbefd36 [Python][Relax] Rotary positional embedding scaling (#17305)
bf7bbefd36 is described below
commit bf7bbefd36ac91242496d533d2bfff71570bf04a
Author: Ruihang Lai <[email protected]>
AuthorDate: Tue Aug 27 10:19:28 2024 -0400
[Python][Relax] Rotary positional embedding scaling (#17305)
This PR introduces two styles of RoPE scaling: the llama3 style
and the longrope scale.
---
python/tvm/relax/frontend/nn/llm/kv_cache.py | 396 +++++++++++++++++++--
.../relax/frontend/nn/llm/position_embedding.py | 191 +++++++++-
python/tvm/relax/frontend/nn/llm/tree_attn.py | 26 +-
...runtime_builtin_paged_attention_kv_cache_tir.py | 19 +-
4 files changed, 579 insertions(+), 53 deletions(-)
diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py
b/python/tvm/relax/frontend/nn/llm/kv_cache.py
index 25a3a1a00d..5ddce76eab 100644
--- a/python/tvm/relax/frontend/nn/llm/kv_cache.py
+++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py
@@ -20,7 +20,7 @@
# pylint:
disable=too-many-statements,too-many-lines,too-many-arguments,invalid-name
import enum
import math
-from typing import Tuple
+from typing import Any, Dict, Tuple
from tvm import relax as rx
from tvm import tir
@@ -29,7 +29,7 @@ from tvm.runtime import DataType
from tvm.script import tir as T
from tvm.target import Target
-from .position_embedding import llama_rope_with_position_map, rope_freq
+from .position_embedding import llama_rope_with_position_map,
switch_rope_freq_func
from .tree_attn import tree_attn
@@ -166,6 +166,8 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
rope_mode: RopeMode,
rope_scale: int,
rope_theta: int,
+ rope_scaling: Dict[str, Any],
+ rope_ext_factors: rx.Expr,
rotary_dim: int,
dtype: str,
target: Target,
@@ -195,6 +197,9 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
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.
@@ -205,6 +210,8 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
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.
"""
@@ -235,8 +242,8 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
bb.add_func(_kv_cache_transpose_append(num_key_value_heads,
head_dim, dtype), "kv_cache_transpose_append"),
rx.extern("flashinfer.attention_kernel_prefill_with_paged_kv_cache"),
rx.extern("flashinfer.attention_kernel_decode_with_paged_kv_cache"),
- bb.add_func(_attention_prefill(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, target),
"tir_attention_prefill_sliding_window"),
- bb.add_func(_attention_decode(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, target),
"tir_attention_decode_sliding_window"),
+ bb.add_func(_attention_prefill(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, rope_scaling, target),
"tir_attention_prefill_sliding_window"),
+ bb.add_func(_attention_decode(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, rope_scaling, target),
"tir_attention_decode_sliding_window"),
rx.extern("flashinfer.attention_kernel_prefill_with_ragged_kv_cache"),
rx.extern("flashinfer.attention_kernel_prefill_with_ragged_kv_cache_begin_forward"),
rx.extern("flashinfer.attention_kernel_prefill_with_ragged_kv_cache_end_forward"),
@@ -245,11 +252,12 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
rx.extern("flashinfer.attention_kernel_decode_with_paged_kv_cache_begin_forward"),
rx.extern("flashinfer.attention_kernel_decode_with_paged_kv_cache_end_forward"),
rx.extern("flashinfer.merge_state_in_place"),
- bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale,
head_dim, num_attention_heads, num_key_value_heads, dtype, rotary_dim),
"tir_split_rotary"),
+ bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale,
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,
head_dim, dtype, target), "kv_cache_copy_single_page"),
bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers,
num_key_value_heads, head_dim, dtype), "kv_cache_debug_get_kv"),
bb.add_func(_compact_kv_copy(num_key_value_heads, head_dim, dtype,
target), "kv_cache_compact_kv_copy"),
- bb.add_func(tree_attn(num_key_value_heads, num_attention_heads,
head_dim, dtype, target), "tir_attention_prefill_with_tree_mask"),
+ bb.add_func(tree_attn(num_key_value_heads, num_attention_heads,
head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask"),
+ rope_ext_factors,
# fmt: on
# pylint: enable=line-too-long
]
@@ -281,6 +289,8 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
head_dim: int,
rope_scale: int,
rope_theta: int,
+ rope_scaling: Dict[str, Any],
+ rope_ext_factors: rx.Expr,
rotary_dim: int,
dtype: str,
target: Target,
@@ -321,6 +331,10 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
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.
target : Target
@@ -349,17 +363,18 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
# pylint: disable=line-too-long
# fmt: off
bb.add_func(_kv_cache_transpose_append(num_key_value_heads,
head_dim, dtype), "kv_cache_transpose_append"),
- bb.add_func(_attention_prefill(num_key_value_heads,
num_attention_heads, head_dim, dtype, False, target), "tir_attention_prefill"),
- bb.add_func(_attention_decode(num_key_value_heads,
num_attention_heads, head_dim, dtype, False, target), "tir_attention_decode"),
- bb.add_func(_attention_prefill(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, target),
"tir_attention_prefill_sliding_window"),
- bb.add_func(_attention_decode(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, target),
"tir_attention_decode_sliding_window"),
- bb.add_func(_attention_prefill_ragged(num_key_value_heads,
num_attention_heads, head_dim, dtype, target), "tir_attention_prefill_ragged"),
+ bb.add_func(_attention_prefill(num_key_value_heads,
num_attention_heads, head_dim, dtype, False, rope_scaling, target),
"tir_attention_prefill"),
+ bb.add_func(_attention_decode(num_key_value_heads,
num_attention_heads, head_dim, dtype, False, rope_scaling, target),
"tir_attention_decode"),
+ bb.add_func(_attention_prefill(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, rope_scaling, target),
"tir_attention_prefill_sliding_window"),
+ bb.add_func(_attention_decode(num_key_value_heads,
num_attention_heads, head_dim, dtype, True, rope_scaling, target),
"tir_attention_decode_sliding_window"),
+ bb.add_func(_attention_prefill_ragged(num_key_value_heads,
num_attention_heads, head_dim, dtype, rope_scaling, target),
"tir_attention_prefill_ragged"),
bb.add_func(_merge_state_inplace(num_attention_heads, head_dim,
dtype, target), "tir_attention_merge_state"),
- bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale,
head_dim, num_attention_heads, num_key_value_heads, dtype, rotary_dim),
"tir_split_rotary"),
+ bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale,
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,
head_dim, dtype, target), "kv_cache_copy_single_page"),
bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers,
num_key_value_heads, head_dim, dtype), "kv_cache_debug_get_kv"),
bb.add_func(_compact_kv_copy(num_key_value_heads, head_dim, dtype,
target), "kv_cache_compact_kv_copy"),
- bb.add_func(tree_attn(num_key_value_heads, num_attention_heads,
head_dim, dtype, target), "tir_attention_prefill_with_tree_mask"),
+ bb.add_func(tree_attn(num_key_value_heads, num_attention_heads,
head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask"),
+ rope_ext_factors,
# fmt: on
# pylint: enable=line-too-long
]
@@ -464,17 +479,23 @@ def _rope(
theta: tir.Var,
scale: tir.Var,
indices: Tuple[tir.Var, ...],
- qkv_dtype="float16",
+ qkv_dtype: str,
+ rope_scaling: Dict[str, Any],
):
d = indices[-1]
- cos_freq, sin_freq = rope_freq(offset * scale, d, rotary_dim, theta,
"float32")
+ cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+ offset * scale, d, rotary_dim, theta, "float32"
+ )
cos = cos_freq * buffer[indices].astype("float32")
sin = sin_freq * tir.if_then_else(
d < rotary_dim // 2,
-buffer[indices[:-1] + (d + rotary_dim // 2,)],
buffer[indices[:-1] + (d - rotary_dim // 2,)],
).astype("float32")
- return (cos + sin).astype(qkv_dtype)
+ expr = (cos + sin).astype(qkv_dtype)
+ for var, value in var_map.items():
+ expr = tir.Let(var, value, expr)
+ return expr
def _var(dtype):
@@ -520,7 +541,9 @@ def _get_seq_offset(pos, seq_id, length_info,
sliding_window):
)
-def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, target:
Target):
+def _attention_prefill(
+ h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling: Dict[str, Any],
target: Target
+):
NUM_BLKS = 16
LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8) # 8 bytes
group_size = h_q // h_kv
@@ -680,7 +703,7 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window:
bool, target: Target
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
- _rope(q,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype),
+ _rope(q,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype,
rope_scaling),
q[cur_L, cur_H_qo, j]
)
else:
@@ -701,7 +724,7 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window:
bool, target: Target
page_offset:
T.int32(is_size_var=True) = T.floormod(seq_offset, 16) # type: ignore
K_smem[i, j] =
T.if_then_else(
rotary_mode == 1,
- _rope(pages,
k_rope_pos_offset[b_idx] + cur_L, d, rope_theta, rope_scale, (page_no, 0, by,
page_offset, j), dtype),
+ _rope(pages,
k_rope_pos_offset[b_idx] + cur_L, d, rope_theta, rope_scale, (page_no, 0, by,
page_offset, j), dtype, rope_scaling),
pages[page_no, 0, by,
page_offset, j]
)
else:
@@ -890,6 +913,7 @@ def _attention_decode(
head_dim,
qkv_dtype,
sliding_window: bool,
+ rope_scaling: Dict[str, Any],
target: Target,
):
qkv_dtype_bytes = 2
@@ -1023,7 +1047,7 @@ def _attention_decode(
for vec in T.vectorized(VEC_SIZE):
Q_local[vec] = T.if_then_else(
rotary_mode == 1,
- _rope(Q, q_rope_position[batch_idx],
head_dim, rope_theta, rope_scale, (bx, by * GROUP_SIZE + bz * bdy + ty, tx *
VEC_SIZE + vec), qkv_dtype),
+ _rope(Q, q_rope_position[batch_idx],
head_dim, rope_theta, rope_scale, (bx, by * GROUP_SIZE + bz * bdy + ty, tx *
VEC_SIZE + vec), qkv_dtype, rope_scaling),
Q[bx, by * GROUP_SIZE + bz * bdy + ty,
tx * VEC_SIZE + vec]
)
@@ -1043,7 +1067,7 @@ def _attention_decode(
for vec in
T.vectorized(VEC_SIZE):
K_smem[tile_start_s + j,
tx * VEC_SIZE + vec] = T.if_then_else(
rotary_mode == 1,
- _rope(pages,
k_rope_pos_offset[batch_idx] + row_g, head_dim, rope_theta, rope_scale,
(page_no, 0, by, page_offset, tx * VEC_SIZE + vec), qkv_dtype),
+ _rope(pages,
k_rope_pos_offset[batch_idx] + row_g, head_dim, rope_theta, rope_scale,
(page_no, 0, by, page_offset, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
pages[page_no, 0, by,
page_offset, tx * VEC_SIZE + vec]
)
V_smem[tile_start_s + j,
tx * VEC_SIZE + vec] = pages[page_no, 1, by, page_offset, tx * VEC_SIZE + vec]
@@ -1210,7 +1234,331 @@ def _merge_state_inplace(num_heads, head_dim, v_dtype,
target: Target):
return merge_state_inplace
-def _attention_prefill_ragged(h_kv, h_q, d, dtype, target: Target):
+def _attention_sequence_prefill(
+ batch_size, h_kv, h_q, d, dtype, target: Target, causal=0,
attn_score_scaling_factor=1.0
+): # pylint: disable=line-too-long
+ LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8) # 8 bytes
+ group_size = h_q // h_kv
+ sm_scale = 1.0 / math.sqrt(float(d)) * math.log2(math.exp(1))
+
+ bdx = 32
+ num_warps = 4
+ tile_x, tile_y, tile_z = 64 // ((DataType(dtype).bits + 7) // 8) // max(d
// 128, 1), d, 16
+
+ # Otherwise we would exceed maxComputeWorkgroupStorageSize
+ if (
+ str(target.kind) == "webgpu"
+ and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
+ ):
+ tile_z = 8
+ num_warps = 2
+
+ # fmt: off
+ @T.prim_func
+ def batch_sequence_prefill_kv( # pylint: disable=too-many-branches
+ var_q: T.handle, # [total_len, h_q, d]
+ var_k: T.handle, # [total_len, h_kv, d]
+ var_v: T.handle, # [total_len, h_kv, d]
+ var_output: T.handle, # [total_len, h_q, d]
+ var_lse: T.handle # [total_len, h_q]
+ ):
+ qo_len = T.int32(is_size_var=True)
+ kv_len = T.int32(is_size_var=True)
+ q = T.match_buffer(var_q, (batch_size, qo_len, h_q, d), dtype)
+ k = T.match_buffer(var_k, (batch_size, kv_len, h_kv, d), dtype)
+ v = T.match_buffer(var_v, (batch_size, kv_len, h_kv, d), dtype)
+ output = T.match_buffer(var_output, (batch_size, qo_len, h_q, d),
dtype)
+ lse = T.match_buffer(var_lse, (batch_size, qo_len, h_q), dtype) #
pylint: disable=unused-variable
+
+ batch_tiles: T.int32 = T.ceildiv(qo_len * group_size, tile_x)
+
+ # kernel code
+ for lbx in T.thread_binding(T.cast(batch_size, "int32") * batch_tiles,
thread="blockIdx.x"):
+ for lby in T.thread_binding(h_kv, thread="blockIdx.y"):
+ for lty in T.thread_binding(num_warps, thread="threadIdx.y"):
+ for ltx in T.thread_binding(bdx, thread="threadIdx.x"):
+ with T.block("attn"):
+ vbx, by, ty, tx = T.axis.remap("SSSS", [lbx, lby,
lty, ltx])
+ T.reads()
+ T.writes()
+
+ Q_smem = T.alloc_buffer((tile_x, d), dtype,
scope="shared")
+ K_smem = T.alloc_buffer((tile_z, d), dtype,
scope="shared")
+ V_smem = T.alloc_buffer((tile_z, d), dtype,
scope="shared")
+ S_smem = T.alloc_buffer((tile_x, tile_z),
"float32", scope="shared")
+
+ S_local = T.alloc_buffer((tile_x, tile_z),
"float32", scope="local")
+ O_local = T.alloc_buffer((tile_x, d), "float32",
scope="local")
+
+ m_smem = T.alloc_buffer((tile_x,), "float32",
scope="shared")
+ m_prev_smem = T.alloc_buffer((tile_x,), "float32",
scope="shared")
+ d_smem = T.alloc_buffer((tile_x,), "float32",
scope="shared")
+
+ m_new = T.alloc_buffer(
+ (math.ceil(tile_x / (bdx * num_warps)),),
"float32", scope="local"
+ )
+ m_prev = T.alloc_buffer(
+ (math.ceil(tile_x / (bdx * num_warps)),),
"float32", scope="local"
+ )
+ d_new = T.alloc_buffer(
+ (math.ceil(tile_x / (bdx * num_warps)),),
"float32", scope="local"
+ )
+
+ b_idx: T.int32 = vbx // batch_tiles
+ tile_id: T.int32 = vbx % batch_tiles
+ LH_start: T.int32 = tile_id * tile_x
+ T.tvm_storage_sync("shared")
+
+ # init states
+ for i in T.serial(T.ceildiv(tile_x, bdx *
num_warps)):
+ row: T.int32 = i * bdx * num_warps + ty * bdx
+ tx
+ if row < tile_x:
+ m_smem[row] = -5e4
+ d_smem[row] = 1.0
+
+ for li, lj in T.grid(tile_x, tile_y):
+ with T.block("O_init"):
+ i, j = T.axis.remap("SS", [li, lj])
+ O_local[i, j] = 0.0
+ T.tvm_storage_sync("shared")
+
+ # Load Q from gmem to smem
+ for li, lj in T.grid(tile_x, tile_y):
+ with T.block("Q_load"):
+ i, j = T.axis.remap("SS", [li, lj])
+ T.reads()
+ T.writes()
+ cur_L = (LH_start + i) // group_size
+ cur_H_qo = by * group_size + (LH_start +
i) % group_size
+ if cur_L < qo_len:
+ Q_smem[i, j] = q[b_idx, cur_L,
cur_H_qo, j]
+ else:
+ Q_smem[i, j] = 0.0
+ T.tvm_storage_sync("shared")
+
+ for iterator in T.serial(T.ceildiv(kv_len,
tile_z)):
+ L_kv_start: T.int32 = iterator * tile_z
+ L_kv_base: T.int32 = 0
+ for lz, ly in T.grid(tile_z, tile_y):
+ with T.block("K_load"):
+ i, j = T.axis.remap("SS", [lz, ly])
+ T.reads()
+ T.writes()
+ cur_L = L_kv_start + i
+ if cur_L < kv_len:
+ K_smem[i, j] = k[
+ b_idx, L_kv_base + cur_L, by, j
+ ]
+ else:
+ K_smem[i, j] = 0.0
+ T.tvm_storage_sync("shared")
+ for lz, ly in T.grid(tile_z, tile_y):
+ with T.block("V_load"):
+ i, j = T.axis.remap("SS", [lz, ly])
+ T.reads()
+ T.writes()
+ cur_L = L_kv_start + i
+ if cur_L < kv_len:
+ V_smem[i, j] = v[
+ b_idx, L_kv_base + cur_L, by, j
+ ]
+ else:
+ V_smem[i, j] = 0.0
+ T.tvm_storage_sync("shared")
+
+ # Compute S
+ with T.block():
+ for li, lj, lk in T.grid(tile_x, tile_z,
tile_y):
+ with T.block("S_gemm"):
+ i, j, k = T.axis.remap("SSR", [li,
lj, lk])
+ with T.init():
+ S_local[i, j] = 0.0
+ S_local[i, j] += (
+ T.cast(Q_smem[i, k], "float32")
+ * T.cast(K_smem[j, k],
"float32")
+ * attn_score_scaling_factor
+ * sm_scale
+ )
+ T.tvm_storage_sync("shared")
+ for li, lj in T.grid(tile_x, tile_z):
+ with T.block("S_store"):
+ i, j = T.axis.remap("SS", [li, lj])
+ S_smem[i, j] = S_local[i, j]
+ T.tvm_storage_sync("shared")
+
+ # Update S, m, d
+ for i in T.serial(T.ceildiv(tile_x, bdx *
num_warps)):
+ row: T.int32 = i * bdx * num_warps + ty *
bdx + tx
+ if row < tile_x:
+ with T.block("update1"):
+ m_prev[i] = m_smem[row]
+ m_new[i] = m_smem[row]
+ # mask out of kv_chunk_len S
+ row_: T.int32 = (LH_start + row)
// group_size
+ for j in T.serial(tile_z):
+ if _causal_mask(
+ causal,
+ row=row_,
+ col=L_kv_start + j,
+ kv_len=kv_len,
+ qo_len=qo_len,
+ ):
+ m_new[i] = T.max(
+ m_new[i], S_smem[row,
j]
+ )
+ d_new[i] = d_smem[row] * T.exp2(
+ m_prev[i] - m_new[i]
+ )
+
+ for i in T.serial(T.ceildiv(tile_x, bdx *
num_warps)):
+ row: T.int32 = i * bdx * num_warps + ty *
bdx + tx
+ with T.block("update"):
+ for j in T.serial(tile_z):
+ # this is to avoid sync inside
condition branch
+ if row < tile_x:
+ row_: T.int32 = (
+ LH_start + row
+ ) // group_size
+ if _causal_mask(
+ causal,
+ row=row_,
+ col=L_kv_start + j,
+ kv_len=kv_len,
+ qo_len=qo_len,
+ ):
+ S_smem[row, j] = T.exp2(
+ S_smem[row, j] -
m_new[i]
+ )
+ else:
+ S_smem[row, j] =
T.exp2(-5e4 - m_new[i])
+
+ for i in T.serial(T.ceildiv(tile_x, bdx *
num_warps)):
+ row: T.int32 = i * bdx * num_warps + ty *
bdx + tx
+ if row < tile_x:
+ with T.block("update"):
+ for j in T.serial(tile_z):
+ d_new[i] += S_smem[row, j]
+ m_smem[row] = m_new[i]
+ d_smem[row] = d_new[i]
+ m_prev_smem[row] = m_prev[i]
+ T.tvm_storage_sync("shared")
+
+ # Update O
+ with T.block():
+ for li, lj, lk in T.grid(tile_x, tile_y,
tile_z):
+ with T.block("O_gemm"):
+ i, j, k = T.axis.remap("SSR", [li,
lj, lk])
+ with T.init():
+ O_local[i, j] *= T.exp2(
+ m_prev_smem[i] - m_smem[i]
+ )
+ O_local[i, j] += S_smem[i, k] *
T.cast(
+ V_smem[k, j], "float32"
+ )
+
+ # Store O from smem to gmem
+ for li, lj in T.grid(tile_x, tile_y):
+ with T.block("O_store"):
+ i, j = T.axis.remap("SS", [li, lj])
+ cur_L: T.int32 = 0 + (LH_start + i) //
group_size
+ cur_H_qo: T.int32 = (
+ by * group_size + (LH_start + i) %
group_size
+ )
+ if cur_L < qo_len:
+ output[b_idx, cur_L, cur_H_qo, j] = (
+ O_local[i, j] / d_smem[i]
+ )
+
+ # Store LSE to gmem
+ for li in T.grid(tile_x):
+ with T.block("lse_store"):
+ i = T.axis.remap("S", [li])
+ cur_L: T.int32 = 0 + (LH_start + i) //
group_size
+ cur_H_qo: T.int32 = (
+ by * group_size + (LH_start + i) %
group_size
+ )
+ if cur_L < qo_len:
+ lse[b_idx, cur_L, cur_H_qo] =
m_smem[i] + T.log2(
+ d_smem[i]
+ )
+
+ # fmt: on
+ # pylint: enable=line-too-long,too-many-branches
+ sch = tir.Schedule(batch_sequence_prefill_kv)
+
+ def get_tile_size(x, y, t):
+ cnt = (x * y) // t
+ assert (x * y) % t == 0
+ tile_y = (int)(math.ceil(math.sqrt(cnt)))
+ while (cnt % tile_y != 0 or y % tile_y != 0) and tile_y <= cnt:
+ tile_y += 1
+ assert tile_y <= cnt
+ tile_x = cnt // tile_y
+ return tile_x, tile_y
+
+ def apply_to_qkv_load(sch: tir.Schedule, block):
+ loop_x, loop_y = sch.get_loops(block)[-2:]
+ loop = sch.fuse(loop_x, loop_y)
+ _, ty, tx, vec = sch.split(
+ loop, factors=[None, num_warps, bdx, LOAD_VEC],
preserve_unit_iters=True
+ )
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.vectorize(vec)
+
+ def apply_to_so_ewise(sch: tir.Schedule, block, tile):
+ loop_x, loop_y = sch.get_loops(block)[-2:]
+ xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+ yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+ sch.reorder(xo, yo, xi, yi)
+ t = sch.fuse(xo, yo)
+ ty, tx = sch.split(t, factors=[None, bdx])
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+
+ def apply_to_gemm( # pylint: disable=unused-argument
+ sch: tir.Schedule, block, tile, read_0, read_1, r_len=8, k_major=False
+ ):
+ loop_x, loop_y, loop_z = sch.get_loops(block)[-3:]
+ xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+ yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+ sch.reorder(xo, yo, xi, yi)
+ t = sch.fuse(xo, yo)
+ ty, tx = sch.split(t, factors=[None, bdx])
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+
+ ko, ki = sch.split(loop_z, factors=[None, r_len])
+ if k_major:
+ sch.reorder(ko, xi, yi, ki)
+ else:
+ sch.reorder(ko, ki, xi, yi)
+ sch.decompose_reduction(block, ty)
+
+ def apply_to_md(sch, block):
+ loop = sch.get_loops(block)[-1]
+ _, ty, tx = sch.split(loop, factors=[None, num_warps, bdx])
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+
+ def apply_schedule(sch):
+ tile_s = get_tile_size(tile_x, tile_z, bdx * num_warps)
+ tile_o = get_tile_size(tile_x, tile_y, bdx * num_warps)
+ apply_to_gemm(sch, sch.get_block("S_gemm"), tile_s, 0, 1, k_major=True)
+ apply_to_gemm(sch, sch.get_block("O_gemm"), tile_o, 2, 3,
k_major=False)
+ apply_to_so_ewise(sch, sch.get_block("S_store"), tile_s)
+ apply_to_so_ewise(sch, sch.get_block("O_init"), tile_o)
+ apply_to_so_ewise(sch, sch.get_block("O_store"), tile_o)
+ apply_to_qkv_load(sch, sch.get_block("Q_load"))
+ apply_to_qkv_load(sch, sch.get_block("K_load"))
+ apply_to_qkv_load(sch, sch.get_block("V_load"))
+
+ apply_schedule(sch)
+ apply_to_md(sch, sch.get_block("lse_store"))
+ return sch.mod["main"].with_attr("tir.is_scheduled", 1)
+
+
+def _attention_prefill_ragged(h_kv, h_q, d, dtype, rope_scaling: Dict[str,
Any], target: Target):
# pylint: disable=line-too-long
NUM_BLKS = 16
LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8) # 8 bytes
@@ -1344,7 +1692,7 @@ def _attention_prefill_ragged(h_kv, h_q, d, dtype,
target: Target):
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
- _rope(q,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype),
+ _rope(q,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype,
rope_scaling),
q[cur_L, cur_H_qo, j]
)
else:
@@ -1363,7 +1711,7 @@ def _attention_prefill_ragged(h_kv, h_q, d, dtype,
target: Target):
if cur_L < kv_chunk_len[0]:
K_smem[i, j] =
T.if_then_else(
rotary_mode == 1,
- _rope(k,
k_rope_pos_offset[b_idx] + cur_L, d, rope_theta, rope_scale, (L_kv_base +
cur_L, by, j), dtype),
+ _rope(k,
k_rope_pos_offset[b_idx] + cur_L, d, rope_theta, rope_scale, (L_kv_base +
cur_L, by, j), dtype, rope_scaling),
k[L_kv_base + cur_L,
by, j]
)
else:
diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py
b/python/tvm/relax/frontend/nn/llm/position_embedding.py
index b224ce04c5..4373395e32 100644
--- a/python/tvm/relax/frontend/nn/llm/position_embedding.py
+++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py
@@ -17,7 +17,9 @@
"""Operators for positional embeddings, e.g. RoPE."""
-from typing import Optional, Tuple
+import math
+from functools import partial
+from typing import Any, Callable, Dict, Optional, Tuple
from tvm import tir
from tvm.relax.frontend.nn import Tensor, op
@@ -26,7 +28,7 @@ from tvm.script import tir as T
# pylint: disable=invalid-name
-def rope_freq(s: tir.Var, d: tir.Var, d_range: int, theta: float, dtype: str):
+def rope_freq_default(s: tir.Var, d: tir.Var, d_range: int, theta: float,
dtype: str):
"""Compute the inverse frequency of RoPE and then return the cosine and
sine of it.
Parameters
@@ -53,11 +55,95 @@ def rope_freq(s: tir.Var, d: tir.Var, d_range: int, theta:
float, dtype: str):
sin_freq : Tensor
The sine of the inverse frequency.
+
+ var_map: Dict[tir.Var, tir.PrimExpr]
+ The common expression map.
"""
freq = s / tir.power(theta, d * 2 % d_range / tir.const(d_range,
"float32"))
- cos_freq = tir.cos(freq).astype(dtype)
- sin_freq = tir.sin(freq).astype(dtype)
- return cos_freq, sin_freq
+ freq_var = tir.Var("freq", "float32")
+ cos_freq = tir.cos(freq_var).astype(dtype)
+ sin_freq = tir.sin(freq_var).astype(dtype)
+ return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_llama3( # pylint: disable=too-many-arguments,too-many-locals
+ s: tir.Var,
+ d: tir.Var,
+ d_range: int,
+ theta: float,
+ dtype: str,
+ factor: float,
+ low_freq_factor: float,
+ high_freq_factor: float,
+ original_max_position_embeddings: float,
+):
+ """Compute the inverse frequency of RoPE for llama3 RoPE scaling."""
+ orig_freq = tir.const(1, "float32") / tir.power(
+ theta, d * 2 % d_range / tir.const(d_range, "float32")
+ )
+ orig_freq_var = tir.Var("orig_freq", "float32")
+ inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+ llama3_inv_scaling_factor = 1.0 / factor
+ llama3_alpha = original_max_position_embeddings / (2 * math.pi) *
inv_diff_freq_factor
+ llama3_beta = low_freq_factor * inv_diff_freq_factor
+ smooth = tir.max(0.0, tir.min(1.0, llama3_alpha * orig_freq_var -
llama3_beta))
+ smoothed_freq = s * (
+ (1.0 - smooth) * orig_freq_var * llama3_inv_scaling_factor + smooth *
orig_freq_var
+ )
+ smoothed_freq_var = tir.Var("smoothed_freq", "float32")
+ cos_freq = tir.cos(smoothed_freq_var).astype(dtype)
+ sin_freq = tir.sin(smoothed_freq_var).astype(dtype)
+ return cos_freq, sin_freq, {smoothed_freq_var: smoothed_freq,
orig_freq_var: orig_freq}
+
+
+def rope_freq_longrope( # pylint: disable=too-many-arguments
+ s: tir.Var,
+ d: tir.Var,
+ d_range: int,
+ theta: float,
+ dtype: str,
+ max_position_embeddings: int,
+ original_max_position_embeddings: int,
+ ext_factors: Optional[T.Buffer] = None,
+):
+ """Compute the inverse frequency of RoPE for longrope scaling."""
+ scale = max_position_embeddings / original_max_position_embeddings
+ scaling_factor = (
+ math.sqrt(1 + math.log(scale) /
math.log(original_max_position_embeddings))
+ if scale > 1.0
+ else 1.0
+ )
+ divisor = tir.power(theta, d * 2 % d_range / tir.const(d_range, "float32"))
+ if ext_factors is not None:
+ divisor = ext_factors[d % (d_range // 2)] * divisor
+ freq = s / divisor
+ freq_var = tir.Var("freq", "float32")
+ cos_freq = (tir.cos(freq_var) * scaling_factor).astype(dtype)
+ sin_freq = (tir.sin(freq_var) * scaling_factor).astype(dtype)
+ return cos_freq, sin_freq, {freq_var: freq}
+
+
+def switch_rope_freq_func(rope_scaling: Dict[str, Any]) -> Callable:
+ """Return the RoPE inverse frequency computation function based
+ on the given RoPE scaling.
+ """
+ if "rope_type" not in rope_scaling:
+ return rope_freq_default
+ if rope_scaling["rope_type"] == "llama3":
+ return partial(
+ rope_freq_llama3,
+ factor=rope_scaling["factor"],
+ low_freq_factor=rope_scaling["low_freq_factor"],
+ high_freq_factor=rope_scaling["high_freq_factor"],
+
original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+ )
+ if rope_scaling["rope_type"] == "longrope":
+ return partial(
+ rope_freq_longrope,
+ max_position_embeddings=rope_scaling["max_position_embeddings"],
+
original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+ )
+ raise ValueError(f'Unsupported RoPE scaling type:
{rope_scaling["rope_type"]}')
# mypy: disable-error-code="attr-defined"
@@ -67,9 +153,10 @@ def llama_rope( # pylint: disable=too-many-arguments
qkv: Tensor,
total_seq_len: tir.Var,
theta: float,
+ scale: float,
num_q_heads: int,
num_kv_heads: int,
- scale: float = 1.0,
+ rope_scaling: Dict[str, Any],
rotary_dim: Optional[int] = None,
) -> Tuple[Tensor, Tensor, Tensor]:
"""Llama-style RoPE. Given a fused QKV tensor, it returns three tensors,
Q, K, and V, where Q
@@ -96,6 +183,9 @@ def llama_rope( # pylint: disable=too-many-arguments
num_kv_heads : int
The number of key/value heads. It differs from `num_q_heads` in
group-query attention.
+ rope_scaling : Dict
+ The configuration of RoPE scaling.
+
rotary_dim : Optional[int]
The number of dimensions in the embedding that RoPE is applied to. By
default, the
rotary_dim is the same as head_dim.
@@ -126,14 +216,19 @@ def llama_rope( # pylint: disable=too-many-arguments
d: tir.Var,
offset: tir.Var,
):
- cos_freq, sin_freq = rope_freq((s + offset) * scale, d, rotary_dim,
theta, dtype)
+ cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+ (s + offset) * scale, d, rotary_dim, theta, dtype
+ )
cos = cos_freq * x[b, s, h, d]
sin = sin_freq * tir.if_then_else(
d < rotary_dim // 2,
-x[b, s, h, d + rotary_dim // 2],
x[b, s, h, d - rotary_dim // 2],
)
- return cos + sin
+ expr = cos + sin
+ for var, value in var_map.items():
+ expr = tir.Let(var, value, expr)
+ return expr
@T.prim_func(private=True)
def fused_rope( # pylint: disable=too-many-locals
@@ -193,6 +288,7 @@ def llama_rope_with_position_map( # pylint:
disable=too-many-arguments
num_q_heads: int,
num_kv_heads: int,
dtype: str,
+ rope_scaling: Dict[str, Any],
rotary_dim: Optional[int] = None,
):
"""Return the TIR function that computes Llama-style RoPE with q position
map.
@@ -217,6 +313,9 @@ def llama_rope_with_position_map( # pylint:
disable=too-many-arguments
dtype : str
The dtype of qkv data.
+ rope_scaling : Dict
+ The configuration of RoPE scaling.
+
rotary_dim : int
The number of dimensions in the embedding that RoPE is applied to. By
default, the
rotary_dim is the same as head_dim.
@@ -225,6 +324,7 @@ def llama_rope_with_position_map( # pylint:
disable=too-many-arguments
if rotary_dim is None:
rotary_dim = head_dim
scale = tir.const(scale, "float32")
+ is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
def _rope( # pylint: disable=too-many-arguments
x: T.Buffer,
@@ -232,15 +332,24 @@ def llama_rope_with_position_map( # pylint:
disable=too-many-arguments
h: tir.Var,
d: tir.Var,
pos: tir.Var,
+ ext_factors: Optional[T.Buffer] = None,
):
- cos_freq, sin_freq = rope_freq(pos * scale, d, rotary_dim, theta,
"float32")
+ kwargs = {}
+ if ext_factors:
+ kwargs["ext_factors"] = ext_factors
+ cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+ pos * scale, d, rotary_dim, theta, "float32", **kwargs
+ )
cos = cos_freq * x[s, h, d].astype("float32")
sin = sin_freq * tir.if_then_else(
d < rotary_dim // 2,
-x[s, h, d + rotary_dim // 2],
x[s, h, d - rotary_dim // 2],
).astype("float32")
- return (cos + sin).astype(dtype)
+ expr = (cos + sin).astype(dtype)
+ for var, value in var_map.items():
+ expr = tir.Let(var, value, expr)
+ return expr
@T.prim_func
def fused_rope( # pylint: disable=too-many-locals
@@ -257,8 +366,8 @@ def llama_rope_with_position_map( # pylint:
disable=too-many-arguments
"tir.noalias": T.bool(True),
}
)
- seq_len = T.int64()
- position_map_elem_offset = T.int64()
+ seq_len = T.int32()
+ position_map_elem_offset = T.int32()
qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
@@ -284,4 +393,62 @@ def llama_rope_with_position_map( # pylint:
disable=too-many-arguments
else:
v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+ @T.prim_func
+ def fused_rope_longrope_scaling( # pylint: disable=too-many-locals
+ var_qkv: T.handle,
+ var_position_map: T.handle,
+ var_q: T.handle,
+ var_k: T.handle,
+ var_v: T.handle,
+ ext_factors: T.Buffer((head_dim // 2,), "float32"), # type: ignore
+ ):
+ T.func_attr(
+ {
+ "op_pattern": 8, # 2 means injective, 8 means opaque
+ "tir.noalias": T.bool(True),
+ }
+ )
+ seq_len = T.int64()
+ position_map_elem_offset = T.int64()
+ qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+ q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+ k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+ v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+ position_map = T.match_buffer(
+ var_position_map, (seq_len,), "int32",
elem_offset=position_map_elem_offset
+ )
+ for iters in T.grid(seq_len, fused_heads, head_dim):
+ with T.block("llama_fused_rope"):
+ s, h, d = T.axis.remap("SSS", iters)
+ if h < num_q_heads:
+ q[s, h, d] = T.if_then_else(
+ d < rotary_dim,
+ _rope(
+ qkv,
+ s,
+ h,
+ d,
+ position_map[s],
+ ext_factors if is_longrope_scaling else None,
+ ),
+ qkv[s, h, d],
+ )
+ elif h < num_q_heads + num_kv_heads:
+ k[s, h - num_q_heads, d] = T.if_then_else(
+ d < rotary_dim,
+ _rope(
+ qkv,
+ s,
+ h,
+ d,
+ position_map[s],
+ ext_factors if is_longrope_scaling else None,
+ ),
+ qkv[s, h, d],
+ )
+ else:
+ v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+ if is_longrope_scaling:
+ return fused_rope_longrope_scaling
return fused_rope
diff --git a/python/tvm/relax/frontend/nn/llm/tree_attn.py
b/python/tvm/relax/frontend/nn/llm/tree_attn.py
index 486491dbf2..069eb48923 100644
--- a/python/tvm/relax/frontend/nn/llm/tree_attn.py
+++ b/python/tvm/relax/frontend/nn/llm/tree_attn.py
@@ -19,14 +19,14 @@
"""Operators for tree attention."""
import math
-from typing import Tuple
+from typing import Any, Dict, Tuple
from tvm import tir
from tvm.runtime import DataType
from tvm.script import tir as T
from tvm.target import Target
-from .position_embedding import rope_freq
+from .position_embedding import switch_rope_freq_func
# mypy: disable-error-code="attr-defined,valid-type,no-redef"
# pylint: disable=too-many-statements,too-many-locals,too-many-arguments
@@ -43,24 +43,30 @@ def _rope(
theta: tir.Var,
scale: tir.Var,
indices: Tuple[tir.Var, ...],
- qkv_dtype="float16",
+ qkv_dtype: str,
+ rope_scaling: Dict[str, Any],
):
d = indices[-1]
- cos_freq, sin_freq = rope_freq(offset * scale, d, rotary_dim, theta,
qkv_dtype)
- cos = cos_freq * buffer[indices]
+ cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+ offset * scale, d, rotary_dim, theta, "float32"
+ )
+ cos = cos_freq * buffer[indices].astype("float32")
sin = sin_freq * tir.if_then_else(
d < rotary_dim // 2,
-buffer[indices[:-1] + (d + rotary_dim // 2,)],
buffer[indices[:-1] + (d - rotary_dim // 2,)],
- )
- return cos + sin
+ ).astype("float32")
+ expr = (cos + sin).astype(qkv_dtype)
+ for var, value in var_map.items():
+ expr = tir.Let(var, value, expr)
+ return expr
def _tree_mask(row, col, mask_ptr, offset, stride, kv_len):
return tir.all(col < kv_len, mask_ptr[offset + row * stride + col] == 1)
-def tree_attn(h_kv, h_q, d, dtype, target: Target): # pylint:
disable=unused-argument
+def tree_attn(h_kv, h_q, d, dtype, rope_scaling: Dict[str, Any], target:
Target):
"""Generate tree attention kernel for batched tree attention.
Parameters
@@ -217,7 +223,7 @@ def tree_attn(h_kv, h_q, d, dtype, target: Target): #
pylint: disable=unused-ar
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
- _rope(q,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype),
+ _rope(q,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, cur_H_qo, j), dtype,
rope_scaling),
q[cur_L, cur_H_qo, j]
)
else:
@@ -236,7 +242,7 @@ def tree_attn(h_kv, h_q, d, dtype, target: Target): #
pylint: disable=unused-ar
if L_kv_start + i <
kv_chunk_len[0]:
K_smem[i, j] =
T.if_then_else(
rotary_mode == 1,
- _rope(k,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, by, j), dtype),
+ _rope(k,
q_rope_position[cur_L], d, rope_theta, rope_scale, (cur_L, by, j), dtype,
rope_scaling),
k[cur_L, by, j]
)
V_smem[i, j] = v[cur_L,
by, j]
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index ff655e141b..c35b7062cd 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -49,6 +49,7 @@ num_kv_heads = 4
head_dim = None
rope_scale = 1.0
rope_theta = 1e4
+rope_scaling = {}
dtype = None
device = tvm.cuda()
@@ -113,15 +114,19 @@ def set_global_func(head_dim, dtype):
for tir_func in [
_kv_cache_transpose_append(num_kv_heads, head_dim, dtype),
_kv_cache_debug_get_kv(num_layers, num_kv_heads, head_dim, dtype),
- _attention_prefill(num_kv_heads, num_qo_heads, head_dim, dtype, False,
target),
- _attention_decode(num_kv_heads, num_qo_heads, head_dim, dtype, False,
target),
- _attention_prefill(num_kv_heads, num_qo_heads, head_dim, dtype, True,
target),
- _attention_decode(num_kv_heads, num_qo_heads, head_dim, dtype, True,
target),
- _attention_prefill_ragged(num_kv_heads, num_qo_heads, head_dim, dtype,
target),
- tree_attn(num_kv_heads, num_qo_heads, head_dim, dtype, target),
+ _attention_prefill(
+ num_kv_heads, num_qo_heads, head_dim, dtype, False, rope_scaling,
target
+ ),
+ _attention_decode(num_kv_heads, num_qo_heads, head_dim, dtype, False,
rope_scaling, target),
+ _attention_prefill(num_kv_heads, num_qo_heads, head_dim, dtype, True,
rope_scaling, target),
+ _attention_decode(num_kv_heads, num_qo_heads, head_dim, dtype, True,
rope_scaling, target),
+ _attention_prefill_ragged(
+ num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling, target
+ ),
+ tree_attn(num_kv_heads, num_qo_heads, head_dim, dtype, rope_scaling,
target),
_merge_state_inplace(num_qo_heads, head_dim, dtype, target),
llama_rope_with_position_map(
- rope_theta, rope_scale, head_dim, num_qo_heads, num_kv_heads, dtype
+ rope_theta, rope_scale, head_dim, num_qo_heads, num_kv_heads,
dtype, rope_scaling
),
_copy_single_page(num_kv_heads, page_size, head_dim, dtype, target),
_compact_kv_copy(num_kv_heads, head_dim, dtype, target),
