# 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