mirrors/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/frontend/nn/llm/_decode_kernels.py
+526
View File
@@ -0,0 +1,526 @@
# 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
# fmt: off
"""Single-token decode attention kernels and attention-state merge helpers.
Contents:
- ``_attention_decode_cpu`` / ``_attention_decode`` — paged-KV decode (one Q token
per sequence), CPU scalar and GPU allreduce variants.
- ``_merge_state_inplace_cpu`` / ``_merge_state_inplace`` — combine two
log-sum-exp attention outputs in place. Used by multi-stage decoding and by
the distributed KV-transfer path.
"""
# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
import math
from typing import Any
from tvm.script import tirx as T
from tvm.target import Target
from ._kernel_common import (
_declare_length_info,
_get_kv_chunk_len,
_get_seq_offset,
_rope,
_var,
_var_cpu,
check_thread_limits,
get_max_num_threads_per_block,
)
def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], page_size: int = 16):
H_qo = num_qo_heads
H_kv = num_kv_heads
D = head_dim
group_size = num_qo_heads // num_kv_heads
global_symbol = "batch_decode_paged_kv_cpu"
if sliding_window:
global_symbol += "_sliding_window"
@T.prim_func(s_tir=True)
def batch_decode_paged_kv(
Q_handle: T.handle,
pages_handle: T.handle,
page_table_indptr_handle: T.handle,
page_table_values_handle: T.handle,
var_length_info: T.handle, # [b] when sliding window = False, or otherwise [3, b]
k_rope_pos_offset_handle: T.handle,
q_rope_position_handle: T.handle,
output_handle: T.handle,
lse_handle: T.handle,
rotary_mode: T.int32,
rope_scale: T.float32,
rope_theta: T.float32,
sm_scale: T.float32,
):
T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
B = T.int32()
nnz_pages = T.int32()
max_num_pages = T.int32()
page_indptr_elem_offset = T.int32()
page_values_elem_offset = T.int32()
k_rope_pos_offset_elem_offset = T.int32()
q_rope_position_elem_offset = T.int32()
length_info_elem_offset = T.int32()
Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype)
page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
lse = T.match_buffer(lse_handle, (B, H_qo), "float32") # pylint: disable=unused-variable
# The length information of the sequences.
# - It is in shape `(3, batch_size)` when sliding window is enabled.
# For a sequence "i", location
# - "(0, i)" is the number of KV slots used in the last page of the seq ("last_page_len"),
# - "(1, i)" is the starting offset of the sliding window in the seq,
# - "(2, i)" is the attn sink length of the sequence.
# - It is in shape `(batch_size,)` when sliding window is disabled,
# denoting the "last_page_len".
length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
for b in T.serial(B):
with T.sblock("attn"):
O_local = T.sblock_alloc_buffer((D,), "float32")
Q_local = T.sblock_alloc_buffer((D,), "float32")
K_local = T.sblock_alloc_buffer((D,), "float32")
V_local = T.sblock_alloc_buffer((D,), "float32")
kv_chunk_len = T.sblock_alloc_buffer((1,), "int32")
m_val = T.sblock_alloc_buffer((1,), "float32")
new_m = T.sblock_alloc_buffer((1,), "float32")
d_val = T.sblock_alloc_buffer((1,), "float32")
S_val = T.sblock_alloc_buffer((1,), "float32")
scale_O = T.sblock_alloc_buffer((1,), "float32")
factor = T.sblock_alloc_buffer((1,), "float32")
cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[b]
cur_page_indptr_end: T.let[T.int32] = page_table_indptr[b + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b, length_info, sliding_window),
0,
)
for h_qo in T.serial(H_qo):
m_val[0] = -5e4
d_val[0] = 1.0
for d in T.serial(D):
O_local[d] = 0.0
for d in T.serial(D):
Q_local[d] = T.if_then_else(
rotary_mode == 1,
_rope(Q, q_rope_position[b], head_dim, rope_theta, rope_scale, (b, h_qo, d), qkv_dtype, rope_scaling),
Q[b, h_qo, d],
)
for row_idx in T.serial(kv_chunk_len[0]):
seq_offset: T.let[T.int32()] = _get_seq_offset(row_idx, b, length_info, sliding_window)
page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + (seq_offset // page_size)]
page_offset: T.let[T.int32()] = seq_offset % page_size
for d in T.serial(D):
K_local[d] = T.if_then_else(
rotary_mode == 1,
_rope(pages, k_rope_pos_offset[b] + row_idx, head_dim, rope_theta, rope_scale, (page_no, 0, h_qo // group_size, page_offset, d), qkv_dtype, rope_scaling),
pages[page_no, 0, h_qo // group_size, page_offset, d],
)
S_val[0] = 0.0
for d in T.serial(D):
S_val[0] += Q_local[d] * K_local[d]
S_val[0] *= sm_scale * math.log2(math.exp(1))
new_m[0] = T.max(m_val[0], S_val[0])
d_val[0] = (d_val[0] * T.exp2(m_val[0] - new_m[0])) + T.exp2(S_val[0] - new_m[0])
scale_O[0] = T.exp2(m_val[0] - new_m[0])
for d in T.serial(D):
O_local[d] = O_local[d] * scale_O[0]
m_val[0] = new_m[0]
for d in T.serial(D):
V_local[d] = pages[page_no, 1, h_qo // group_size, page_offset, d]
factor[0] = T.exp2(S_val[0] - m_val[0])
for d in T.serial(D):
O_local[d] = O_local[d] + V_local[d] * factor[0]
for d in T.serial(D):
O_local[d] = O_local[d] / d_val[0]
output[b, h_qo, d] = O_local[d]
lse[b, h_qo] = m_val[0] + T.log2(d_val[0])
return batch_decode_paged_kv
def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], target: Target, page_size: int = 16):
qkv_dtype_bytes = 2
H_qo = num_qo_heads
H_kv = num_kv_heads
D = head_dim
THREAD_LIMIT = 512
TILE_SIZE_PER_BDX = 2
if target.kind.name == "opencl" and (("android" in str(target.host)) or ("adreno" in str(target.attrs))):
# Keeping lower thread limit for this kernel on adreno target
# to avoid register spill
THREAD_LIMIT = 256
TILE_SIZE_PER_BDX = 1
max_num_threads_per_block = get_max_num_threads_per_block(target)
thread_limit = min(max_num_threads_per_block, THREAD_LIMIT)
GROUP_SIZE = H_qo // H_kv
VEC_SIZE = min(max(8 // qkv_dtype_bytes, D // 32), 4)
bdx = D // VEC_SIZE
bdy = GROUP_SIZE
while bdx * bdy > thread_limit and bdy > 1:
bdy //= 2
gdz = GROUP_SIZE // bdy
threads_per_CTA = max(thread_limit, bdx * bdy)
bdz = threads_per_CTA // (bdx * bdy)
tile_size_per_bdx = TILE_SIZE_PER_BDX if GROUP_SIZE == 1 else 1
check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=bdz, gdz=1)
global_symbol = "batch_decode_paged_kv"
if sliding_window:
global_symbol += "_sliding_window"
# pylint: disable=too-many-branches
@T.prim_func(s_tir=True)
def batch_decode_paged_kv(
Q_handle: T.handle,
pages_handle: T.handle,
page_table_indptr_handle: T.handle,
page_table_values_handle: T.handle,
var_length_info: T.handle, # [b] when sliding window = False, or otherwise [3, b]
k_rope_pos_offset_handle: T.handle,
q_rope_position_handle: T.handle,
output_handle: T.handle,
lse_handle: T.handle,
rotary_mode: T.int32,
rope_scale: T.float32,
rope_theta: T.float32,
sm_scale: T.float32,
):
T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
B = T.int32()
nnz_pages = T.int32()
max_num_pages = T.int32()
pages_elem_offset = T.int64()
page_indptr_elem_offset = T.int32()
page_values_elem_offset = T.int32()
k_rope_pos_offset_elem_offset = T.int32()
q_rope_position_elem_offset = T.int32()
length_info_elem_offset = T.int32()
Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype, elem_offset=pages_elem_offset)
page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
lse = T.match_buffer(lse_handle, (B, H_qo), "float32") # pylint: disable=unused-variable
length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
for bx in T.thread_binding(B, thread="blockIdx.x"):
for fused_by_bz in T.thread_binding(H_kv * gdz, thread="blockIdx.y"):
for ty in T.thread_binding(bdy, thread="threadIdx.y"):
for tx in T.thread_binding(bdx, thread="threadIdx.x"):
for tz in T.thread_binding(bdz, thread="threadIdx.z"):
with T.sblock("attn"):
Q_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
kv_chunk_len = T.sblock_alloc_buffer((1,), "int32", scope="local")
K_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
V_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
O_allreduce = T.sblock_alloc_buffer((bdz, bdy, D), "float32", scope="shared")
md_allreduce = T.sblock_alloc_buffer((bdz, bdy, 2), "float32", scope="shared")
S_reduce_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
S_local = T.sblock_alloc_buffer((bdy * tile_size_per_bdx), "float32", scope="local")
QK_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
V_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
m_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
d_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
other_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
other_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
exp_mprev = T.sblock_alloc_buffer((1,), "float32", scope="local")
exp_otherm = T.sblock_alloc_buffer((1,), "float32", scope="local")
other_o = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
st_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
st_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
O_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
by: T.let[T.int32] = fused_by_bz % H_kv
bz: T.let[T.int32] = fused_by_bz // H_kv
batch_idx: T.let[T.int32] = bx
cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[batch_idx]
cur_page_indptr_end: T.let[T.int32] = page_table_indptr[batch_idx + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, batch_idx, length_info, sliding_window),
0
)
# init states
st_m[0] = -5e4
st_d[0] = 1.0
for vec in T.vectorized(VEC_SIZE):
O_local[vec] = 0.0
# load q
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_scaling),
Q[bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec]
)
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_size_per_bdx * bdy * bdz)):
tile_start_s: T.let[T.int32()] = (tz * bdy + ty) * tile_size_per_bdx # type: ignore
tile_start_g: T.let[T.int32()] = ((iterator * bdz + tz) * bdy + ty) * tile_size_per_bdx # type: ignore
# load KV from global memory to shared memory
for j in T.serial(tile_size_per_bdx):
with T.sblock("KV_load"):
T.reads()
T.writes()
row_g: T.let[T.int32()] = tile_start_g + j # type: ignore
if row_g < kv_chunk_len[0]:
seq_offset: T.let[T.int32()] = _get_seq_offset(row_g, batch_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.let[T.int32()] = T.floormod(seq_offset, page_size) # type: ignore
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_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]
else:
for vec in T.vectorized(VEC_SIZE):
K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
T.tvm_storage_sync("shared")
# compute QK
m_prev[0] = st_m[0]
for j in T.serial(bdy * tile_size_per_bdx):
# compute S = Q * K * sm_scale
for vec in T.vectorized(VEC_SIZE):
QK_local[vec] = T.cast(Q_local[vec], "float32") * T.cast(K_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec], "float32") * sm_scale * math.log2(math.exp(1))
S_reduce_local[0] = 0
for vec in T.unroll(VEC_SIZE):
S_reduce_local[0] += QK_local[vec]
with T.sblock("block_cross_thread"):
T.reads(S_reduce_local[0])
T.writes(t0[0])
T.attr(
T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
"reduce_scope",
T.int32(0),
)
T.tvm_thread_allreduce(T.uint32(1), S_reduce_local[0], True, t0[0], tx, dtype="void")
S_local[j] = -5e4
if (iterator * bdz + tz) * bdy * tile_size_per_bdx + j < kv_chunk_len[0]:
S_local[j] = t0[0]
# update st_m
st_m[0] = T.max(st_m[0], S_local[j])
# update st_d, st_O
o_scale: T.let[T.float32] = T.exp2(m_prev[0] - st_m[0])
st_d[0] *= o_scale
for j in T.serial(bdy * tile_size_per_bdx):
S_local[j] = T.exp2(S_local[j] - st_m[0])
st_d[0] += S_local[j]
for j in T.vectorized(VEC_SIZE):
O_local[j] *= o_scale
# load V from shared memory to local memory
# compute O
for j in T.serial(bdy * tile_size_per_bdx):
for vec in T.vectorized(VEC_SIZE):
V_local[vec] = V_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec]
for vec in T.vectorized(VEC_SIZE):
O_local[vec] += T.cast(V_local[vec], "float32") * S_local[j]
if bdz > 1:
# allreduce over bdz
for vec in T.vectorized(VEC_SIZE):
O_allreduce[tz, ty, tx * VEC_SIZE + vec] = O_local[vec]
md_allreduce[tz, ty, 0] = st_m[0]
md_allreduce[tz, ty, 1] = st_d[0]
T.tvm_storage_sync("shared")
st_m[0] = -5e4
st_d[0] = 1.0
for vec in T.vectorized(VEC_SIZE):
O_local[vec] = 0.0
for j in T.serial(bdz):
m_prev[0] = st_m[0]
d_prev[0] = st_d[0]
other_m[0] = md_allreduce[j, ty, 0]
other_d[0] = md_allreduce[j, ty, 1]
for vec in T.vectorized(VEC_SIZE):
other_o[vec] = O_allreduce[j, ty, tx * VEC_SIZE + vec]
st_m[0] = T.max(st_m[0], other_m[0])
st_d[0] = d_prev[0] * T.exp2(m_prev[0] - st_m[0]) + other_d[0] * T.exp2(other_m[0] - st_m[0])
exp_mprev[0] = T.exp2(m_prev[0] - st_m[0])
exp_otherm[0] = T.exp2(other_m[0] - st_m[0])
for vec in T.vectorized(VEC_SIZE):
O_local[vec] = O_local[vec] * exp_mprev[0] + other_o[vec] * exp_otherm[0]
# normalize O
for vec in T.vectorized(VEC_SIZE):
O_local[vec] /= st_d[0]
# store O to global memory
for vec in T.vectorized(VEC_SIZE):
output[batch_idx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec] = O_local[vec]
# store lse to global memory
lse[batch_idx, by * GROUP_SIZE + bz * bdy + ty] = st_m[0] + T.log2(st_d[0])
# pylint: enable=too-many-branches
return batch_decode_paged_kv
def _merge_state_inplace_cpu(v_dtype):
@T.prim_func(s_tir=True)
def merge_state_inplace_cpu(
v: T.handle,
s: T.handle,
v_other: T.handle,
s_other: T.handle,
):
T.func_attr({"tirx.is_scheduled": True})
N = T.int32()
H = T.int32()
D = T.int32()
V = T.match_buffer(v, (N, H, D), v_dtype)
S = T.match_buffer(s, (N, H), "float32")
V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
S_other = T.match_buffer(s_other, (N, H), "float32")
for n in T.serial(N):
for h in T.serial(H):
with T.sblock("merge"):
s_val = _var_cpu("float32")
s_other_val = _var_cpu("float32")
s_max = _var_cpu("float32")
scale = _var_cpu("float32")
other_scale = _var_cpu("float32")
s_val[0] = S[n, h]
s_other_val[0] = S_other[n, h]
s_max[0] = T.max(s_val[0], s_other_val[0])
s_val[0] = T.exp2(s_val[0] - s_max[0])
s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
for d in T.serial(D):
V[n, h, d] = V[n, h, d] * scale[0] + V_other[n, h, d] * other_scale[0]
S[n, h] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
return merge_state_inplace_cpu
def _merge_state_inplace(num_heads, head_dim, v_dtype, target: Target, global_symbol: str | None = None):
v_dtype_bytes = 2
VEC_SIZE = min(max(8 // v_dtype_bytes, head_dim // 32), 4)
bdx = head_dim // VEC_SIZE
bdy = num_heads
max_num_threads_per_block = get_max_num_threads_per_block(target)
while bdx * bdy > max_num_threads_per_block and bdy > 1:
bdy //= 2
gdy = num_heads // bdy
check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=1, gdz=1)
@T.prim_func(s_tir=True)
def merge_state_inplace(
v: T.handle,
s: T.handle,
v_other: T.handle,
s_other: T.handle,
):
T.func_attr({"tirx.is_scheduled": True})
N = T.int32()
H = T.int32()
D = T.int32()
V = T.match_buffer(v, (N, H, D), v_dtype)
S = T.match_buffer(s, (N, H), "float32")
V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
S_other = T.match_buffer(s_other, (N, H), "float32")
for bx in T.thread_binding(N, thread="blockIdx.x"):
for by in T.thread_binding(gdy, thread="blockIdx.y"):
for ty in T.thread_binding(bdy, thread="threadIdx.y"):
for tx in T.thread_binding(bdx, thread="threadIdx.x"):
with T.sblock("merge"):
s_val = _var("float32")
s_other_val = _var("float32")
s_max = _var("float32")
scale = _var("float32")
other_scale = _var("float32")
v_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
v_other_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
s_val[0] = S[bx, ty + by * bdy]
s_other_val[0] = S_other[bx, ty + by * bdy]
s_max[0] = T.max(s_val[0], s_other_val[0])
s_val[0] = T.exp2(s_val[0] - s_max[0])
s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
# load v
for vec in T.vectorized(VEC_SIZE):
v_vec[vec] = V[bx, ty + by * bdy, tx * VEC_SIZE + vec]
# load v_other
for vec in T.vectorized(VEC_SIZE):
v_other_vec[vec] = V_other[bx, ty + by * bdy, tx * VEC_SIZE + vec]
# merge
for vec in T.serial(VEC_SIZE):
v_vec[vec] = v_vec[vec] * scale[0] + v_other_vec[vec] * other_scale[0]
# store v
for vec in T.vectorized(VEC_SIZE):
V[bx, ty + by * bdy, tx * VEC_SIZE + vec] = v_vec[vec]
# store s
S[bx, ty + by * bdy] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
func = merge_state_inplace
if global_symbol:
func = func.with_attr("global_symbol", global_symbol)
return func