diff --git a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py index 751ef411e340..96bafd381bf9 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, layout=T.TileLayout(T.S[(max_num_pages, 2, H_kv, page_size, D)] + (pages_elem_offset))), - page_table_indptr: T.Tensor((B + 1,), 'int32', layout=T.TileLayout(T.S[(B + 1,)] + (page_indptr_elem_offset))), - page_table_values: T.Tensor((nnz_pages,), 'int32', layout=T.TileLayout(T.S[(nnz_pages,)] + (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', layout=T.TileLayout(T.S[(B,)] + (k_rope_pos_offset_elem_offset))), - q_rope_position: T.Tensor((B,), 'int32', layout=T.TileLayout(T.S[(B,)] + (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, layout=T.TileLayout(T.S[(max_num_pages, 2, H_kv, page_size, D)] + (pages_elem_offset))), + page_table_indptr: T.Tensor((B + 1,), 'int32', layout=T.TileLayout(T.S[(B + 1,)] + (page_indptr_elem_offset))), + page_table_values: T.Tensor((nnz_pages,), 'int32', layout=T.TileLayout(T.S[(nnz_pages,)] + (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', layout=T.TileLayout(T.S[(B,)] + (k_rope_pos_offset_elem_offset))), + q_rope_position: T.Tensor((B,), 'int32', layout=T.TileLayout(T.S[(B,)] + (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 61cad58b4c24..52404399cddb 100644 --- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py +++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py @@ -15,12 +15,14 @@ # specific language governing permissions and limitations # under the License. +import re + import pytest import tvm import tvm_tirx import tvm_tirx.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, @@ -287,5 +289,44 @@ def test_merge_state_waits_before_overwriting_the_shared_lse(head_dim): ) +def _webgpu_kernels(func, target): + wgsl = tvm_tirx.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_tirx.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_tirx.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_tirx.target.Target(kind)) + assert "gpu_thread_invariant" not in func.script() + + if __name__ == "__main__": tvm_tirx.testing.main()