This is an automated email from the ASF dual-hosted git repository.
tlopex 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 993863d083 [WebGPU] Parallelize prefill attention for 512-wide FP32
heads (#20700)
993863d083 is described below
commit 993863d0831c15d31b5da37e609fc175791be718
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Oct 13 19:59:15 2026 -0700
[WebGPU] Parallelize prefill attention for 512-wide FP32 heads (#20700)
WebGPU prefill with 512-wide FP32 heads spends most of each Q/K dot
product in a serial loop. This affects Gemma 4 E2B `q4f32_1`; the
small-tile fallback in #20658 makes the kernel fit in shared memory but
leaves this bottleneck.
Distribute each dot product across 32 threads in a 128-thread workgroup.
For paged global attention, split the KV range into four dispatches and
merge their outputs using log-sum-exp weights, so one dispatch does not
process the entire context. Ragged and sliding attention use a single
dispatch.
Select this path only for WebGPU FP32 attention with 512-wide heads.
---
.../tvm/relax/frontend/nn/llm/_prefill_kernels.py | 10 +
.../tvm/relax/frontend/nn/llm/_prefill_parallel.py | 486 +++++++++++++++++++++
tests/python/relax/test_nn_llm_prefill_parallel.py | 151 +++++++
3 files changed, 647 insertions(+)
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
index fd98fe83b3..dc5e77576a 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
@@ -50,6 +50,10 @@ from ._kernel_common import (
_rope,
_schedule_prefill_kernel,
)
+from ._prefill_parallel import (
+ _attention_prefill_paged_parallel,
+ _attention_prefill_ragged_parallel,
+)
def _attention_prefill_cpu(
@@ -213,6 +217,10 @@ def _attention_prefill(
page_size: int = 16,
sliding_window_size: int = 1024,
):
+ if target.kind.name == "webgpu" and dtype == "float32" and d == 512:
+ return _attention_prefill_paged_parallel(
+ h_kv, h_q, rope_scaling, sliding_window, page_size,
sliding_window_size
+ )
NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z =
_get_prefill_kernel_config(h_kv, h_q, d, dtype, target)
global_symbol = "batch_prefill_paged_kv"
@@ -747,6 +755,8 @@ def _attention_prefill_ragged_cpu(h_kv, h_q, d_qk, d_v,
dtype, rope_scaling: dic
return batch_prefill_ragged_kv
def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling:
dict[str, Any], target: Target):
+ if target.kind.name == "webgpu" and dtype == "float32" and d_qk == d_v ==
512:
+ return _attention_prefill_ragged_parallel(h_kv, h_q, rope_scaling)
NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z =
_get_prefill_kernel_config(h_kv, h_q, d_qk, dtype, target, d_v=d_v)
init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, _,
advance_tile_batch, paged_store_output_lse, *_ = _make_prefill_macros(tile_x,
tile_y, tile_z, d_v, bdx, num_warps, group_size)
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
new file mode 100644
index 0000000000..73b7da44f9
--- /dev/null
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_parallel.py
@@ -0,0 +1,486 @@
+# 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.
+"""Parallel dot-product prefill for wide FP32 WebGPU attention heads."""
+
+from tvm.script import s_tir as Ts
+from tvm.script import tirx as T
+
+from ._kernel_common import _get_kv_chunk_len, _get_seq_offset,
_length_info_buffer, _rope
+
+
+def _parallel_prefill_body(
+ h_kv,
+ h_q,
+ rope_scaling,
+ paged=False,
+ sliding_window=False,
+ page_size=16,
+ sliding_window_size=1024,
+ partitions=1,
+):
+ # One workgroup handles one query/head. Four groups of 32 threads each
+ # reduce a 512-element Q/K dot product, then update the online softmax.
+ dtype = "float32"
+
+ def kv_length(length, batch, length_info):
+ return (
+ _get_kv_chunk_len(length, page_size, batch, length_info,
sliding_window)
+ if paged
+ else length
+ )
+
+ def kv_page(indptr, indices, length_info, batch, key):
+ if paged:
+ slot = _get_seq_offset(key, batch, length_info, sliding_window)
+ return indices[indptr[batch] + slot // page_size]
+ return indptr[batch] + key
+
+ def kv_offset(length_info, batch, key):
+ return _get_seq_offset(key, batch, length_info, sliding_window) %
page_size if paged else 0
+
+ def kv_index(page, offset, head, feature, kind):
+ return (page, kind, head, offset, feature) if paged else (page, head,
feature)
+
+ def load_key(
+ source, page, offset, position, head, feature, rotary_mode,
rope_scale, rope_theta
+ ):
+ index = kv_index(page, offset, head, feature, 0)
+ return T.if_then_else(
+ rotary_mode == 1,
+ _rope(source, position, 512, rope_theta, rope_scale, index, dtype,
rope_scaling),
+ source[index],
+ )
+
+ def load_value(source, page, offset, head, feature):
+ return source[kv_index(page, offset, head, feature, 1)]
+
+ @T.macro
+ def body(
+ q,
+ q_indptr,
+ k,
+ v,
+ kv_indptr,
+ q_rope_position,
+ k_rope_pos_offset,
+ output,
+ lse,
+ causal,
+ rotary_mode,
+ rope_scale,
+ rope_theta,
+ sm_scale,
+ page_values,
+ length_info,
+ part_idx=0,
+ ):
+ for row in T.parallel(q.shape[0], annotations={"thread_binding":
"blockIdx.x"}):
+ for head in T.parallel(h_q, annotations={"thread_binding":
"blockIdx.y"}):
+ for lane in T.parallel(128, annotations={"thread_binding":
"threadIdx.x"}):
+ with Ts.sblock("attention"):
+ Ts.reads()
+ Ts.writes()
+ partial = Ts.sblock_alloc_tensor((128,), "float32",
scope="shared")
+ probabilities = Ts.sblock_alloc_tensor((4,),
"float32", scope="shared")
+ state = Ts.sblock_alloc_tensor((3,), "float32",
scope="shared")
+ query = Ts.sblock_alloc_tensor((16,), "float32",
scope="local")
+ accum = Ts.sblock_alloc_tensor((4,), "float32",
scope="local")
+ score = Ts.sblock_alloc_tensor((1,), "float32",
scope="local")
+ weights = Ts.sblock_alloc_tensor((4,), "float32",
scope="local")
+ denominator = Ts.sblock_alloc_tensor((1,), "float32",
scope="local")
+ rescale = Ts.sblock_alloc_tensor((1,), "float32",
scope="local")
+ b = Ts.sblock_alloc_tensor((1,), "int32",
scope="local")
+ b[0] = 0
+ while row >= q_indptr[b[0] + 1]:
+ b[0] += 1
+ begin: T.let[T.int32] = kv_indptr[b[0]]
+ count: T.let[T.int32] = kv_length(
+ kv_indptr[b[0] + 1] - begin, b[0], length_info
+ )
+ limit: T.let[T.int32] = T.if_then_else(
+ causal != 0 and not sliding_window,
+ T.min(
+ count,
+ row
+ - q_indptr[b[0]]
+ + count
+ - (q_indptr[b[0] + 1] - q_indptr[b[0]])
+ + 1,
+ ),
+ count,
+ )
+ first: T.let[T.int32] = T.if_then_else(
+ causal != 0 and sliding_window,
+ T.max(
+ count - T.max(sliding_window_size - (row -
q_indptr[b[0]]) - 1, 0),
+ 0,
+ ),
+ 0,
+ )
+ width: T.let[T.int32] = T.ceildiv(T.max(limit - first,
0), partitions)
+ partition_first: T.let[T.int32] = T.min(limit, first +
part_idx * width)
+ partition_limit: T.let[T.int32] = T.min(
+ limit, first + (part_idx + 1) * width
+ )
+ kh: T.let[T.int32] = head // (h_q // h_kv)
+ for j in T.unroll(16):
+ feature: T.let[T.int32] = j * 32 + lane % 32
+ query[j] = T.if_then_else(
+ rotary_mode == 1,
+ _rope(
+ q,
+ q_rope_position[row],
+ 512,
+ rope_theta,
+ rope_scale,
+ (row, head, feature),
+ dtype,
+ rope_scaling,
+ ),
+ q[row, head, feature],
+ )
+ for j in T.unroll(4):
+ accum[j] = 0
+ if lane == 0:
+ state[0] = -50000
+ state[1] = 0
+ T.gpu_storage_sync("shared")
+ for tile in T.serial(T.ceildiv(partition_limit -
partition_first, 4)):
+ key: T.let[T.int32] = partition_first + tile * 4 +
lane // 32
+ score[0] = 0
+ if key < partition_limit:
+ page: T.let[T.int32] = kv_page(
+ kv_indptr, page_values, length_info, b[0],
key
+ )
+ offset: T.let[T.int32] =
kv_offset(length_info, b[0], key)
+ for j in T.unroll(16):
+ feature: T.let[T.int32] = j * 32 + lane %
32
+ score[0] += query[j] * load_key(
+ k,
+ page,
+ offset,
+ k_rope_pos_offset[b[0]] + key,
+ kh,
+ feature,
+ rotary_mode,
+ rope_scale,
+ rope_theta,
+ )
+ partial[lane] = score[0]
+ T.gpu_storage_sync("shared")
+ for reduction in T.unroll(5):
+ width: T.let[T.int32] = 16 >> reduction
+ if lane % 32 < width:
+ partial[lane] += partial[lane + width]
+ T.gpu_storage_sync("shared")
+ if lane == 0:
+ score[0] = state[0]
+ for j in T.unroll(4):
+ weights[j] = T.if_then_else(
+ partition_first + tile * 4 + j <
partition_limit,
+ partial[j * 32] * sm_scale *
1.4426950408889634,
+ T.float32(-50000),
+ )
+ score[0] = T.max(score[0], weights[j])
+ rescale[0] = T.exp2(state[0] - score[0])
+ denominator[0] = state[1] * rescale[0]
+ for j in T.unroll(4):
+ weights[j] = T.if_then_else(
+ partition_first + tile * 4 + j <
partition_limit,
+ T.exp2(weights[j] - score[0]),
+ T.float32(0),
+ )
+ denominator[0] += weights[j]
+ T.gpu_storage_sync("shared")
+ if lane == 0:
+ for j in T.unroll(4):
+ probabilities[j] = weights[j]
+ state[0] = score[0]
+ state[1] = denominator[0]
+ state[2] = rescale[0]
+ T.gpu_storage_sync("shared")
+ for feature in T.unroll(4):
+ accum[feature] *= state[2]
+ for j in T.unroll(4):
+ if partition_first + tile * 4 + j <
partition_limit:
+ page: T.let[T.int32] = kv_page(
+ kv_indptr,
+ page_values,
+ length_info,
+ b[0],
+ partition_first + tile * 4 + j,
+ )
+ offset: T.let[T.int32] = kv_offset(
+ length_info, b[0], partition_first
+ tile * 4 + j
+ )
+ accum[feature] += probabilities[j] *
load_value(
+ v, page, offset, kh, feature * 128
+ lane
+ )
+ T.gpu_storage_sync("shared")
+ for feature in T.unroll(4):
+ output[row, head, feature * 128 + lane] =
T.if_then_else(
+ state[1] > 0, accum[feature] / state[1],
T.float32(0)
+ )
+ if lane == 0:
+ lse[row, head] = T.if_then_else(
+ state[1] > 0, state[0] + T.log2(state[1]),
T.float32(-50000)
+ )
+
+ return body
+
+
+def _attention_prefill_ragged_parallel(h_kv, h_q, rope_scaling):
+ batch, qlen, klen = [T.dynamic(x, "int32") for x in ["batch_size",
"qo_len", "kv_len"]]
+ qi, ki, qp, kp = [
+ T.dynamic(x, "int32")
+ for x in [
+ "q_indptr_elem_offset",
+ "kv_indptr_elem_offset",
+ "q_rope_position_elem_offset",
+ "k_rope_pos_offset_elem_offset",
+ ]
+ ]
+ body = _parallel_prefill_body(h_kv, h_q, rope_scaling)
+
+ @Ts.function
+ def batch_prefill_ragged_kv(
+ q: T.Tensor((qlen, h_q, 512), "float32"),
+ q_indptr: T.Tensor((batch + 1,), "int32", elem_offset=qi),
+ k: T.Tensor((klen, h_kv, 512), "float32"),
+ v: T.Tensor((klen, h_kv, 512), "float32"),
+ kv_indptr: T.Tensor((batch + 1,), "int32", elem_offset=ki),
+ q_rope_position: T.Tensor((qlen,), "int32", elem_offset=qp),
+ k_rope_pos_offset: T.Tensor((batch,), "int32", elem_offset=kp),
+ output: T.Tensor((qlen, h_q, 512), "float32"),
+ lse: T.Tensor((qlen, h_q), "float32"),
+ causal: T.int32,
+ rotary_mode: T.int32,
+ rope_scale: T.float32,
+ rope_theta: T.float32,
+ sm_scale: T.float32,
+ ):
+ body(
+ q,
+ q_indptr,
+ k,
+ v,
+ kv_indptr,
+ q_rope_position,
+ k_rope_pos_offset,
+ output,
+ lse,
+ causal,
+ rotary_mode,
+ rope_scale,
+ rope_theta,
+ sm_scale,
+ kv_indptr,
+ kv_indptr,
+ )
+
+ return batch_prefill_ragged_kv.with_attr("tirx.is_scheduled", True)
+
+
+def _attention_prefill_paged_parallel(
+ h_kv, h_q, rope_scaling, sliding_window, page_size=16,
sliding_window_size=1024
+):
+ if not sliding_window:
+ return _attention_prefill_paged_split(
+ h_kv, h_q, rope_scaling, sliding_window, page_size,
sliding_window_size
+ )
+ batch, qlen, num_pages, nnz_pages = [
+ T.dynamic(x, "int32") for x in ["batch_size", "qo_len", "num_pages",
"nnz_pages"]
+ ]
+ qi, pi, pv, li, qp, kp, po = [
+ T.dynamic(x, "int32")
+ for x in [
+ "q_indptr_elem_offset",
+ "page_indptr_elem_offset",
+ "page_values_elem_offset",
+ "length_info_elem_offset",
+ "q_rope_position_elem_offset",
+ "k_rope_pos_offset_elem_offset",
+ "pages_elem_offset",
+ ]
+ ]
+ body = _parallel_prefill_body(
+ h_kv, h_q, rope_scaling, True, sliding_window, page_size,
sliding_window_size
+ )
+ symbol = "batch_prefill_paged_kv_sliding_window" if sliding_window else
"batch_prefill_paged_kv"
+
+ @Ts.function
+ def batch_prefill_paged_kv(
+ q: T.Tensor((qlen, h_q, 512), "float32"),
+ q_indptr: T.Tensor((batch + 1,), "int32", elem_offset=qi),
+ pages: T.Tensor((num_pages, 2, h_kv, page_size, 512), "float32",
elem_offset=po),
+ page_indptr: T.Tensor((batch + 1,), "int32", elem_offset=pi),
+ page_values: T.Tensor((nnz_pages,), "int32", elem_offset=pv),
+ length_info: _length_info_buffer(batch, sliding_window, li),
+ k_rope_pos_offset: T.Tensor((batch,), "int32", elem_offset=kp),
+ q_rope_position: T.Tensor((qlen,), "int32", elem_offset=qp),
+ output: T.Tensor((qlen, h_q, 512), "float32"),
+ lse: T.Tensor((qlen, h_q), "float32"),
+ causal: T.int32,
+ rotary_mode: T.int32,
+ rope_scale: T.float32,
+ rope_theta: T.float32,
+ sm_scale: T.float32,
+ ):
+ T.func_attr({"global_symbol": symbol})
+ body(
+ q,
+ q_indptr,
+ pages,
+ pages,
+ page_indptr,
+ q_rope_position,
+ k_rope_pos_offset,
+ output,
+ lse,
+ causal,
+ rotary_mode,
+ rope_scale,
+ rope_theta,
+ sm_scale,
+ page_values,
+ length_info,
+ )
+
+ return batch_prefill_paged_kv.with_attr("tirx.is_scheduled", True)
+
+
+def _attention_prefill_paged_split(
+ h_kv, h_q, rope_scaling, sliding_window, page_size=16,
sliding_window_size=1024
+):
+ batch, qlen, num_pages, nnz_pages = [
+ T.dynamic(x, "int32") for x in ["batch_size", "qo_len", "num_pages",
"nnz_pages"]
+ ]
+ qi, pi, pv, li, qp, kp, po = [
+ T.dynamic(x, "int32")
+ for x in [
+ "q_indptr_elem_offset",
+ "page_indptr_elem_offset",
+ "page_values_elem_offset",
+ "length_info_elem_offset",
+ "q_rope_position_elem_offset",
+ "k_rope_pos_offset_elem_offset",
+ "pages_elem_offset",
+ ]
+ ]
+ body = _parallel_prefill_body(
+ h_kv, h_q, rope_scaling, True, False, page_size, sliding_window_size,
partitions=4
+ )
+ symbol = "batch_prefill_paged_kv_sliding_window" if sliding_window else
"batch_prefill_paged_kv"
+
+ @Ts.function
+ def batch_prefill_paged_kv(
+ q: T.Tensor((qlen, h_q, 512), "float32"),
+ q_indptr: T.Tensor((batch + 1,), "int32", elem_offset=qi),
+ pages: T.Tensor((num_pages, 2, h_kv, page_size, 512), "float32",
elem_offset=po),
+ page_indptr: T.Tensor((batch + 1,), "int32", elem_offset=pi),
+ page_values: T.Tensor((nnz_pages,), "int32", elem_offset=pv),
+ length_info: _length_info_buffer(batch, sliding_window, li),
+ k_rope_pos_offset: T.Tensor((batch,), "int32", elem_offset=kp),
+ q_rope_position: T.Tensor((qlen,), "int32", elem_offset=qp),
+ output: T.Tensor((qlen, h_q, 512), "float32"),
+ lse: T.Tensor((qlen, h_q), "float32"),
+ causal: T.int32,
+ rotary_mode: T.int32,
+ rope_scale: T.float32,
+ rope_theta: T.float32,
+ sm_scale: T.float32,
+ ):
+ T.func_attr({"global_symbol": symbol})
+ # Launch each KV partition separately so long contexts do not produce
+ # one long-running dispatch. Merge their softmax-weighted outputs
below.
+ temporary = Ts.sblock_alloc_tensor(
+ (T.int64(4), T.Cast("int64", qlen), T.int64(h_q), T.int64(512)),
"float32"
+ )
+ temporary_lse = Ts.sblock_alloc_tensor(
+ (T.int64(4), T.Cast("int64", qlen), T.int64(h_q)), "float32"
+ )
+ for part_idx in T.serial(4):
+ with Ts.sblock("part_idx"):
+ part_output = Ts.match_tensor(
+ temporary[part_idx, 0:qlen, 0:h_q, 0:512], (qlen, h_q,
512), offset_factor=1
+ )
+ part_lse = Ts.match_tensor(
+ temporary_lse[part_idx, 0:qlen, 0:h_q], (qlen, h_q),
offset_factor=1
+ )
+ body(
+ q,
+ q_indptr,
+ pages,
+ pages,
+ page_indptr,
+ q_rope_position,
+ k_rope_pos_offset,
+ part_output,
+ part_lse,
+ causal,
+ rotary_mode,
+ rope_scale,
+ rope_theta,
+ sm_scale,
+ page_values,
+ length_info,
+ part_idx,
+ )
+ for row in T.parallel(qlen, annotations={"thread_binding":
"blockIdx.x"}):
+ for head in T.parallel(h_q, annotations={"thread_binding":
"blockIdx.y"}):
+ for lane in T.parallel(128, annotations={"thread_binding":
"threadIdx.x"}):
+ with Ts.sblock("merge"):
+ Ts.reads()
+ Ts.writes()
+ weights = Ts.sblock_alloc_tensor((4,), "float32",
scope="shared")
+ local_weights = Ts.sblock_alloc_tensor((4,),
"float32", scope="local")
+ maximum = Ts.sblock_alloc_tensor((1,), "float32",
scope="local")
+ denominator = Ts.sblock_alloc_tensor((1,), "float32",
scope="local")
+ value = Ts.sblock_alloc_tensor((1,), "float32",
scope="local")
+ if lane == 0:
+ maximum[0] = T.float32(-50000)
+ for part in T.unroll(4):
+ local_weights[part] = temporary_lse[part, row,
head]
+ maximum[0] = T.max(maximum[0],
local_weights[part])
+ denominator[0] = 0
+ for part in T.unroll(4):
+ local_weights[part] = T.if_then_else(
+ local_weights[part] > T.float32(-49999),
+ T.exp2(local_weights[part] - maximum[0]),
+ T.float32(0),
+ )
+ denominator[0] += local_weights[part]
+ for part in T.unroll(4):
+ weights[part] = local_weights[part] / T.max(
+ denominator[0], T.float32(1)
+ )
+ lse[row, head] = T.if_then_else(
+ denominator[0] > 0,
+ maximum[0] + T.log2(denominator[0]),
+ T.float32(-50000),
+ )
+ T.gpu_storage_sync("shared")
+ for feature in T.unroll(4):
+ value[0] = 0
+ for part in T.unroll(4):
+ value[0] += (
+ weights[part] * temporary[part, row, head,
feature * 128 + lane]
+ )
+ output[row, head, feature * 128 + lane] = value[0]
+
+ # Let the caller assign a symbol when scratch buffers are lifted.
+ return
batch_prefill_paged_kv.without_attr("global_symbol").with_attr("tirx.is_scheduled",
True)
diff --git a/tests/python/relax/test_nn_llm_prefill_parallel.py
b/tests/python/relax/test_nn_llm_prefill_parallel.py
new file mode 100644
index 0000000000..7009302b5c
--- /dev/null
+++ b/tests/python/relax/test_nn_llm_prefill_parallel.py
@@ -0,0 +1,151 @@
+# 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.
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.relax.frontend.nn.llm._prefill_kernels import (
+ _attention_prefill,
+ _attention_prefill_ragged,
+)
+from tvm.relax.frontend.nn.llm._prefill_parallel import (
+ _attention_prefill_paged_parallel,
+ _attention_prefill_ragged_parallel,
+)
+
+
[email protected]("layout", ["ragged", "paged", "sliding"])
+def test_webgpu_parallel_prefill_compiles(layout):
+ target = tvm.target.Target("webgpu")
+ if layout == "ragged":
+ func = _attention_prefill_ragged(1, 8, 512, 512, "float32", {}, target)
+ else:
+ func = _attention_prefill(1, 8, 512, "float32", layout == "sliding",
{}, target)
+ compiled = tvm.tirx.compile(func.with_attr("global_symbol", "main"),
target=target)
+ assert compiled.mod.imports[0].inspect_source()
+
+
+def rope(values, positions):
+ frequency = positions[:, None] / 10000 ** (np.arange(256, dtype="float64")
/ 256)
+ angles = np.concatenate([frequency, frequency], axis=-1)[:, None, :]
+ rotated = np.concatenate([-values[..., 256:], values[..., :256]], axis=-1)
+ return values * np.cos(angles) + rotated * np.sin(angles)
+
+
[email protected]
[email protected]("device_type", ["cuda", "metal"])
[email protected]("layout", ["ragged", "paged", "sliding"])
[email protected]("rotary", [0, 1])
[email protected]("causal", [0, 1])
[email protected]("first_kv_length", [1, 23])
+def test_parallel_prefill(device_type, layout, rotary, causal,
first_kv_length):
+ if not tvm.testing.device_enabled(device_type):
+ pytest.skip(f"{device_type} not enabled")
+ device = tvm.device(device_type)
+ target = tvm.target.Target.from_device(device)
+ rng = np.random.default_rng(127)
+ lengths = [(7, first_kv_length), (17, 513)]
+ qi = np.array([0, 7, 24], dtype="int32")
+ ki = np.array([0, first_kv_length, first_kv_length + 513], dtype="int32")
+ q = rng.normal(0, 0.3, (24, 8, 512)).astype("float32")
+ k = rng.normal(0, 0.3, (first_kv_length + 513, 2, 512)).astype("float32")
+ v = rng.normal(size=(first_kv_length + 513, 2, 512)).astype("float32")
+ kp = np.array([32, 1234], dtype="int32")
+ qp = np.concatenate(
+ [np.arange(nk - nq, nk) + kp[b] for b, (nq, nk) in enumerate(lengths)]
+ ).astype("int32")
+ reference = np.zeros_like(q)
+ reference_lse = np.full((24, 8), -50000, dtype="float32")
+ scale = 512**-0.5
+ for b, (nq, nk) in enumerate(lengths):
+ query = q[qi[b] : qi[b + 1]].astype("float64")
+ keys = k[ki[b] : ki[b + 1]].astype("float64")
+ if rotary:
+ query = rope(query, qp[qi[b] : qi[b + 1]])
+ keys = rope(keys, np.arange(nk) + kp[b])
+ keys = np.repeat(keys, 4, axis=1)
+ values = np.repeat(v[ki[b] : ki[b + 1]], 4, axis=1)
+ scores = np.einsum("qhd,khd->hqk", query, keys) * scale
+ mask = np.ones((nq, nk), dtype=bool)
+ if causal:
+ if layout == "sliding":
+ mask = np.arange(nk)[None, :] >= np.maximum(
+ nk - np.maximum(15 - np.arange(nq)[:, None], 0), 0
+ )
+ else:
+ mask = np.arange(nk)[None, :] <= np.arange(nq)[:, None] + nk -
nq
+ for row in range(nq):
+ if not mask[row].any():
+ continue
+ scores_row = scores[:, row, mask[row]]
+ maximum = scores_row.max(axis=-1, keepdims=True)
+ weights = np.exp(scores_row - maximum)
+ denominator = weights.sum(axis=-1, keepdims=True)
+ reference[qi[b] + row] = np.einsum(
+ "hk,khd->hd", weights / denominator, values[mask[row]]
+ )
+ reference_lse[qi[b] + row] = ((maximum + np.log(denominator)) /
np.log(2))[:, 0]
+ output = tvm.runtime.empty(q.shape, "float32", device)
+ lse = tvm.runtime.empty(reference_lse.shape, "float32", device)
+
+ def tensor(value):
+ return tvm.runtime.tensor(value, device)
+
+ if layout == "ragged":
+ func = _attention_prefill_ragged_parallel(2, 8, {})
+ args = [tensor(q), tensor(qi), tensor(k), tensor(v), tensor(ki),
tensor(qp), tensor(kp)]
+ else:
+ sliding = layout == "sliding"
+ offsets = [3, 17] if sliding else [0, 0]
+ sinks = [0, 2] if sliding else [0, 0]
+ physical_lengths = [
+ nk + offset - sink for (_, nk), offset, sink in zip(lengths,
offsets, sinks)
+ ]
+ counts = [(n + 15) // 16 for n in physical_lengths]
+ pi = np.array([0, counts[0], sum(counts)], dtype="int32")
+ indices = rng.permutation(sum(counts)).astype("int32")
+ pages = np.zeros((sum(counts), 2, 2, 16, 512), dtype="float32")
+ for b, (_, nk) in enumerate(lengths):
+ for row in range(nk):
+ slot = row if row < sinks[b] else row - sinks[b] + offsets[b]
+ page = indices[pi[b] + slot // 16]
+ pages[page, 0, :, slot % 16] = k[ki[b] + row]
+ pages[page, 1, :, slot % 16] = v[ki[b] + row]
+ last = [(n - 1) % 16 + 1 for n in physical_lengths]
+ info = np.array([last, offsets, sinks] if sliding else last,
dtype="int32")
+ func = _attention_prefill_paged_parallel(2, 8, {}, sliding,
sliding_window_size=16)
+ args = [
+ tensor(q),
+ tensor(qi),
+ tensor(pages),
+ tensor(pi),
+ tensor(indices),
+ tensor(info),
+ tensor(kp),
+ tensor(qp),
+ ]
+ module = tvm.compile(
+ tvm.IRModule({"main": func.with_attr("global_symbol", "main")}),
target=target
+ )
+ module(*args, output, lse, causal, rotary, 1.0, 10000.0, scale)
+ tvm.testing.assert_allclose(output.numpy(), reference, rtol=2e-4,
atol=2e-5)
+ tvm.testing.assert_allclose(lse.numpy(), reference_lse, rtol=2e-5,
atol=2e-5)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
