From 0abb60565cf170afadca36c55724a4b68bf08147 Mon Sep 17 00:00:00 2001 From: JCNTH Date: Thu, 24 Sep 2026 10:10:05 -0400 Subject: [PATCH] [Relax][WebGPU] Sequence-sharded paged decode attention for long contexts The paged decode attention kernel spreads threads across the head dimension (bdx = head_dim / vec), so each q.k dot product ends in a cross-thread allreduce and one trip of the serial loop over the context covers only bdz * bdy positions. On WebGPU, with at most 256 invocations per workgroup and every workgroupBarrier a real sync, that loop dominates decode time at long context. Add a second thread nest for WebGPU that cuts the context into shards owned by a few lanes each (64 dims of q and O per lane in registers), exchanges partial scores through a double-buffered workgroup array so one sync per position suffices, and merges the per-shard online-softmax states in two phases that fit the default 16 KiB of workgroup storage. Wide heads (128, 256) split across threadIdx.y lanes so the workgroup stays at 256 invocations. The existing tiled nest is wrapped in a T.macro unchanged and remains in charge of short contexts; the choice is made on the host from the symbolic page count via gpu_thread_invariant. Other targets are unaffected. Attention time per decode step at 3060 tokens of context, Apple M4 Max / Chrome, sharded vs tiled nest: 2.9x / 4.0x / 3.4x faster for head_dim 64 / 128 / 256 (Llama-3.1-8B shape: 36.6 -> 9.0 ms). Output matches an f64 reference to f16 rounding (0.02% relative L2) at 1..3060 positions. --- .../relax/frontend/nn/llm/_decode_kernels.py | 222 ++++++++++++++++-- .../test_frontend_nn_llm_kernel_config.py | 43 +++- 2 files changed, 244 insertions(+), 21 deletions(-) diff --git a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py index 165ad11915fc..2104edf0a1da 100644 --- a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py +++ b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py @@ -202,6 +202,11 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w if sliding_window: global_symbol += "_sliding_window" + # webgpu: long contexts (picked on the host from the page count) use a sequence-sharded nest + use_sharded = target.kind.name == "webgpu" and D % 4 == 0 + LONG_CONTEXT_PAGES = 256 // page_size + sharded_nest = _make_seq_sharded_decode_nest(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window, rope_scaling, target, page_size) if use_sharded else None + # pylint: disable=too-many-branches B = T.dynamic("B", "int32") nnz_pages = T.dynamic("nnz_pages", "int32") @@ -212,26 +217,8 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w k_rope_pos_offset_elem_offset = T.dynamic("k_rope_pos_offset_elem_offset", "int32") q_rope_position_elem_offset = T.dynamic("q_rope_position_elem_offset", "int32") length_info_elem_offset = T.dynamic("length_info_elem_offset", "int32") - @Ts.function - def batch_decode_paged_kv( - Q: T.Tensor((B, H_qo, D), qkv_dtype), - pages: T.Tensor((max_num_pages, 2, H_kv, page_size, D), qkv_dtype, elem_offset=pages_elem_offset), - page_table_indptr: T.Tensor((B + 1,), 'int32', elem_offset=page_indptr_elem_offset), - page_table_values: T.Tensor((nnz_pages,), 'int32', elem_offset=page_values_elem_offset), - length_info: _length_info_buffer(B, sliding_window, length_info_elem_offset), # [b] when sliding window = False, or otherwise [3, b] - k_rope_pos_offset: T.Tensor((B,), 'int32', elem_offset=k_rope_pos_offset_elem_offset), - q_rope_position: T.Tensor((B,), 'int32', elem_offset=q_rope_position_elem_offset), - output: T.Tensor((B, H_qo, D), qkv_dtype), - lse: T.Tensor((B, H_qo), 'float32'), - 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}) - - # pylint: disable=unused-variable - + @T.macro + def tiled_nest(Q, pages, page_table_indptr, page_table_values, length_info, k_rope_pos_offset, q_rope_position, output, lse, B, rotary_mode, rope_scale, rope_theta, sm_scale): for bx in T.parallel(B, annotations={"thread_binding": "blockIdx.x"}): for fused_by_bz in T.parallel(H_kv * gdz, annotations={"thread_binding": "blockIdx.y"}): for ty in T.parallel(bdy, annotations={"thread_binding": "threadIdx.y"}): @@ -386,9 +373,204 @@ def batch_decode_paged_kv( # store lse to global memory lse[batch_idx, by * GROUP_SIZE + bz * bdy + ty] = st_m[0] + T.log2(st_d[0]) + @Ts.function + def batch_decode_paged_kv( + Q: T.Tensor((B, H_qo, D), qkv_dtype), + pages: T.Tensor((max_num_pages, 2, H_kv, page_size, D), qkv_dtype, elem_offset=pages_elem_offset), + page_table_indptr: T.Tensor((B + 1,), 'int32', elem_offset=page_indptr_elem_offset), + page_table_values: T.Tensor((nnz_pages,), 'int32', elem_offset=page_values_elem_offset), + length_info: _length_info_buffer(B, sliding_window, length_info_elem_offset), # [b] when sliding window = False, or otherwise [3, b] + k_rope_pos_offset: T.Tensor((B,), 'int32', elem_offset=k_rope_pos_offset_elem_offset), + q_rope_position: T.Tensor((B,), 'int32', elem_offset=q_rope_position_elem_offset), + output: T.Tensor((B, H_qo, D), qkv_dtype), + lse: T.Tensor((B, H_qo), 'float32'), + 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}) + + # pylint: disable=unused-variable + + if T.constexpr(use_sharded): + if T.gpu_thread_invariant(nnz_pages > B * LONG_CONTEXT_PAGES): + sharded_nest(Q, pages, page_table_indptr, page_table_values, length_info, k_rope_pos_offset, q_rope_position, output, lse, B, rotary_mode, rope_scale, rope_theta, sm_scale) + else: + tiled_nest(Q, pages, page_table_indptr, page_table_values, length_info, k_rope_pos_offset, q_rope_position, output, lse, B, rotary_mode, rope_scale, rope_theta, sm_scale) + else: + tiled_nest(Q, pages, page_table_indptr, page_table_values, length_info, k_rope_pos_offset, q_rope_position, output, lse, B, rotary_mode, rope_scale, rope_theta, sm_scale) # pylint: enable=too-many-branches return batch_decode_paged_kv +def _make_seq_sharded_decode_nest(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], target: Target, page_size: int = 16): + """Build a `@T.macro` decode nest that shards the KV sequence across threads (WebGPU long context).""" + H_qo = num_qo_heads + H_kv = num_kv_heads + D = head_dim + GROUP_SIZE = H_qo // H_kv + VEC_SIZE = 4 + + # a lane keeps 64 dims of q and O in registers; wider heads use several lanes per shard + MAX_DIMS_PER_LANE = 64 + MAX_SHARDS = 128 + MIN_SHARD_LEN = 8 + SMEM_BYTES = 16384 # WebGPU default maxComputeWorkgroupStorageSize + lanes_per_shard = D // MAX_DIMS_PER_LANE if (D > MAX_DIMS_PER_LANE and D % MAX_DIMS_PER_LANE == 0) else 1 + dims_per_lane = D // lanes_per_shard + vec_per_lane = dims_per_lane // VEC_SIZE + gdz = GROUP_SIZE # one workgroup per query head + max_threads = get_max_num_threads_per_block(target) + num_shards = 1 + while num_shards * 2 * lanes_per_shard <= max_threads and num_shards * 2 <= MAX_SHARDS: + num_shards *= 2 + # merge O in chunks so the per-thread workgroup memory (chunk + m, d, two scores) fits + chunk = dims_per_lane + while chunk > VEC_SIZE and num_shards * lanes_per_shard * (chunk + 4) * 4 > SMEM_BYTES: + chunk //= 2 + chunk_vecs = chunk // VEC_SIZE + num_chunks = dims_per_lane // chunk + check_thread_limits(target, bdx=num_shards, bdy=lanes_per_shard, bdz=1, gdz=1) + + # pylint: disable=too-many-branches,too-many-statements + @T.macro + def sharded_nest(Q, pages, page_table_indptr, page_table_values, length_info, k_rope_pos_offset, q_rope_position, output, lse, B, rotary_mode, rope_scale, rope_theta, sm_scale): + for bx in T.parallel(B, annotations={"thread_binding": "blockIdx.x"}): + for fused_by_bz in T.parallel(H_kv * gdz, annotations={"thread_binding": "blockIdx.y"}): + for ty in T.parallel(lanes_per_shard, annotations={"thread_binding": "threadIdx.y"}): + for tx in T.parallel(num_shards, annotations={"thread_binding": "threadIdx.x"}): + with Ts.sblock("attn"): + Q_local = Ts.sblock_alloc_tensor((dims_per_lane,), "float32", scope="local") + O_local = Ts.sblock_alloc_tensor((dims_per_lane,), "float32", scope="local") + QK_local = Ts.sblock_alloc_tensor((VEC_SIZE,), "float32", scope="local") + kv_chunk_len = Ts.sblock_alloc_tensor((1,), "int32", scope="local") + shard_len = Ts.sblock_alloc_tensor((1,), "int32", scope="local") + num_active = Ts.sblock_alloc_tensor((1,), "int32", scope="local") + # per-lane partial scores; two buffers so the sync pass emits one barrier per position + score_a = Ts.sblock_alloc_tensor((num_shards, lanes_per_shard), "float32", scope="shared") + score_b = Ts.sblock_alloc_tensor((num_shards, lanes_per_shard), "float32", scope="shared") + O_chunk = Ts.sblock_alloc_tensor((num_shards, lanes_per_shard, chunk), "float32", scope="shared") + md_shards = Ts.sblock_alloc_tensor((num_shards, 2), "float32", scope="shared") + acc = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + own_m = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + weight = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + s_acc = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + st_m = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + st_d = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + m_prev = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + exp_mprev = Ts.sblock_alloc_tensor((1,), "float32", scope="local") + exp_otherm = Ts.sblock_alloc_tensor((1,), "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 + q_head: T.let[T.int32] = by * GROUP_SIZE + bz + 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 + ) + shard_len[0] = T.max(T.ceildiv(kv_chunk_len[0], num_shards), MIN_SHARD_LEN) + num_active[0] = T.ceildiv(kv_chunk_len[0], shard_len[0]) + + # init states; load q (idle shards beyond the context only join the syncs) + st_m[0] = -5e4 + st_d[0] = 1.0 + for lane in T.serial(vec_per_lane): + for vec in T.vectorized(VEC_SIZE): + O_local[lane * VEC_SIZE + vec] = 0.0 + if tx < num_active[0]: + for lane in T.serial(vec_per_lane): + for vec in T.vectorized(VEC_SIZE): + Q_local[lane * VEC_SIZE + vec] = T.cast(T.if_then_else( + rotary_mode == 1, + _rope(Q, q_rope_position[batch_idx], head_dim, rope_theta, rope_scale, (bx, q_head, ty * dims_per_lane + lane * VEC_SIZE + vec), qkv_dtype, rope_scaling), + Q[bx, q_head, ty * dims_per_lane + lane * VEC_SIZE + vec] + ), "float32") + + # shard tx owns rows [tx * shard_len, (tx + 1) * shard_len); the loop is + # uniform across threads, invalid rows score 0 and skip the update + for it2 in T.serial(T.ceildiv(shard_len[0], 2)): + for half in T.unroll(2): + row_g: T.let[T.int32()] = tx * shard_len[0] + it2 * 2 + half # type: ignore + row_c: T.let[T.int32()] = T.max(T.min(row_g, kv_chunk_len[0] - 1), 0) # type: ignore + seq_offset: T.let[T.int32()] = _get_seq_offset(row_c, 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 + # an odd shard_len's extra row belongs to the next shard + valid: T.let[T.bool] = T.And(it2 * 2 + half < shard_len[0], row_g < kv_chunk_len[0]) # type: ignore + for vec in T.vectorized(VEC_SIZE): + QK_local[vec] = 0.0 + if valid: + if rotary_mode == 1: + for lane in T.serial(vec_per_lane): + for vec in T.vectorized(VEC_SIZE): + QK_local[vec] += Q_local[lane * VEC_SIZE + vec] * T.cast(_rope(pages, k_rope_pos_offset[batch_idx] + row_g, head_dim, rope_theta, rope_scale, (page_no, 0, by, page_offset, ty * dims_per_lane + lane * VEC_SIZE + vec), qkv_dtype, rope_scaling), "float32") + else: + for lane in T.serial(vec_per_lane): + for vec in T.vectorized(VEC_SIZE): + QK_local[vec] += Q_local[lane * VEC_SIZE + vec] * T.cast(pages[page_no, 0, by, page_offset, ty * dims_per_lane + lane * VEC_SIZE + vec], "float32") + s_acc[0] = 0.0 + for vec in T.unroll(VEC_SIZE): + s_acc[0] += QK_local[vec] + if half == 0: + score_a[tx, ty] = s_acc[0] + else: + score_b[tx, ty] = s_acc[0] + T.gpu_storage_sync("shared") + s_acc[0] = 0.0 + for lane_i in T.unroll(lanes_per_shard): + if half == 0: + s_acc[0] += score_a[tx, lane_i] + else: + s_acc[0] += score_b[tx, lane_i] + if valid: + # online softmax update; (m, d) is identical on every lane + s_acc[0] = s_acc[0] * sm_scale * math.log2(math.exp(1)) + m_prev[0] = st_m[0] + st_m[0] = T.max(st_m[0], s_acc[0]) + exp_mprev[0] = T.exp2(m_prev[0] - st_m[0]) + exp_otherm[0] = T.exp2(s_acc[0] - st_m[0]) + st_d[0] = st_d[0] * exp_mprev[0] + exp_otherm[0] + for lane in T.serial(vec_per_lane): + for vec in T.vectorized(VEC_SIZE): + O_local[lane * VEC_SIZE + vec] = O_local[lane * VEC_SIZE + vec] * exp_mprev[0] + T.cast(pages[page_no, 1, by, page_offset, ty * dims_per_lane + lane * VEC_SIZE + vec], "float32") * exp_otherm[0] + + # merge (m, d) across shards + if ty == 0: + md_shards[tx, 0] = st_m[0] + md_shards[tx, 1] = st_d[0] + own_m[0] = st_m[0] + T.gpu_storage_sync("shared") + st_m[0] = -5e4 + st_d[0] = 1.0 + for j in T.serial(num_active[0]): + m_prev[0] = st_m[0] + st_m[0] = T.max(st_m[0], md_shards[j, 0]) + st_d[0] = st_d[0] * T.exp2(m_prev[0] - st_m[0]) + md_shards[j, 1] * T.exp2(md_shards[j, 0] - st_m[0]) + weight[0] = T.exp2(own_m[0] - st_m[0]) / st_d[0] + + # merge O in chunks: thread (tx < chunk, ty) sums one dim over the shards + for c in T.serial(num_chunks): + if tx < num_active[0]: + for lane in T.serial(chunk_vecs): + for vec in T.vectorized(VEC_SIZE): + O_chunk[tx, ty, lane * VEC_SIZE + vec] = O_local[c * chunk + lane * VEC_SIZE + vec] * weight[0] + T.gpu_storage_sync("shared") + if tx < chunk: + acc[0] = 0.0 + for j in T.serial(num_active[0]): + acc[0] += O_chunk[j, ty, tx] + output[batch_idx, q_head, ty * dims_per_lane + c * chunk + tx] = T.cast(acc[0], qkv_dtype) + T.gpu_storage_sync("shared") + if tx == 0 and ty == 0: + lse[batch_idx, q_head] = st_m[0] + T.log2(st_d[0]) + # pylint: enable=too-many-branches,too-many-statements + return sharded_nest + + def _merge_state_inplace_cpu(v_dtype): N = T.dynamic("N", "int32") H = T.dynamic("H", "int32") diff --git a/tests/python/relax/test_frontend_nn_llm_kernel_config.py b/tests/python/relax/test_frontend_nn_llm_kernel_config.py index 85fb823f9014..a14e8dd56797 100644 --- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py +++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py @@ -15,11 +15,13 @@ # specific language governing permissions and limitations # under the License. +import re + import pytest import tvm import tvm.testing -from tvm.relax.frontend.nn.llm._decode_kernels import _merge_state_inplace +from tvm.relax.frontend.nn.llm._decode_kernels import _attention_decode, _merge_state_inplace from tvm.relax.frontend.nn.llm._kernel_common import ( _get_prefill_kernel_config, _get_prefill_shared_memory_usage, @@ -274,5 +276,44 @@ def test_merge_state_waits_before_overwriting_the_shared_lse(head_dim): ) +def _webgpu_kernels(func, target): + wgsl = tvm.tirx.compile(func, target=target).mod.imports[0].inspect_source("wgsl") + return [k for k in re.split(r"(?=// Function: )", wgsl) if "@compute" in k] + + +def _workgroup_size(kernel): + return re.search(r"@workgroup_size\(([^)]*)\)", kernel).group(1) + + +@pytest.mark.parametrize("sliding_window", [False, True]) +def test_webgpu_decode_attention_dispatches_on_context_length(sliding_window): + target = tvm.target.Target("webgpu", host="llvm") + func = _attention_decode(8, 32, 64, "float16", sliding_window, {}, target) + kernels = _webgpu_kernels(func, target) + # sharded nest (one workgroup per query head, 128 shards), then the tiled nest unchanged + assert [_workgroup_size(k) for k in kernels] == ["128, 1, 1", "16, 4, 4"] + assert kernels[0].count("workgroupBarrier") < kernels[1].count("workgroupBarrier") + + +@pytest.mark.parametrize( + "head_dim, size", [(128, "128, 2, 1"), (256, "64, 4, 1"), (96, "128, 1, 1")] +) +def test_webgpu_decode_attention_lanes_follow_head_dim(head_dim, size): + target = tvm.target.Target("webgpu", host="llvm") + func = _attention_decode(4, 8, head_dim, "float16", False, {}, target) + kernel = _webgpu_kernels(func, target)[0] + assert _workgroup_size(kernel) == size + smem = sum( + int(n) * 4 for n in re.findall(r"var\s+\w+\s*:\s*array", kernel) + ) + assert smem <= 16384 + + +@pytest.mark.parametrize("kind", ["cuda", "metal"]) +def test_decode_attention_unchanged_off_webgpu(kind): + func = _attention_decode(8, 32, 64, "float16", False, {}, tvm.target.Target(kind)) + assert "gpu_thread_invariant" not in func.script() + + if __name__ == "__main__": tvm.testing.main()