diff --git a/python/tvm/dlight/gpu/matmul.py b/python/tvm/dlight/gpu/matmul.py index a5759941caf5..25cc649b44dd 100644 --- a/python/tvm/dlight/gpu/matmul.py +++ b/python/tvm/dlight/gpu/matmul.py @@ -27,7 +27,7 @@ from tvm.tir.analysis import undefined_vars from tvm.tir.schedule.schedule import BlockRV -from ..base import analysis +from ..base import analysis, BlockInfo, IterInfo from .base import GPUScheduleRule @@ -273,6 +273,32 @@ def get_index_map(block: tir.Block) -> Optional[Tuple[tir.IndexMap, ...]]: ) +def get_block_info(sch: tir.Schedule, block: tir.schedule.BlockRV) -> BlockInfo: + def _iter_kind(loop: tir.IterVar) -> str: + return {tir.IterVar.DataPar: "S", tir.IterVar.CommReduce: "R"}.get(loop.iter_type, "O") + + def _is_reduction_block(block: tir.schedule.BlockRV): + for iter_var in sch.get(block).iter_vars: + if _iter_kind(iter_var) == "R": + return True + return False + + return BlockInfo( + name=sch.get(block).name_hint, + iters=[ + IterInfo( + kind=_iter_kind(iter_var), + var=iter_var.var, + dom=iter_var.dom.extent, + loop_rv=loop_rv, + ) + for loop_rv, iter_var in zip(sch.get_loops(block), sch.get(block).iter_vars) + ], + block_rv=block, + reduction_block=_is_reduction_block(block), + ) + + def get_reduction_blocks(sch, blocks) -> bool: # Get the main computation block def is_reduction(block: BlockRV) -> bool: @@ -914,17 +940,19 @@ def get_configs(self, target: Target) -> Config: storage_align=True, inner_x=False, ) - elif target.kind.name == "opencl" and "android" in str(target.host): + elif target.kind.name == "opencl" and ( + ("android" in str(target.host)) or ("windows" in str(target.host)) + ): return Matmul.Config( - block_size_x=8, - block_size_y=16, + block_size_x=32, + block_size_y=8, vthread_x=1, vthread_y=1, micro_size_x=8, micro_size_y=2, micro_size_k=16, vector_size=8, - unroll=64, + unroll=4, use_shared=False, storage_align=False, inner_x=True, @@ -941,6 +969,7 @@ def apply( # pylint: disable=too-many-locals,missing-docstring if not isinstance(func, tir.PrimFunc) or not self.is_target_available(target): return None sch = tir.Schedule(func) + config = self.get_configs(target) root_block = analysis.get_root_block(sch) blocks = sch.get_child_blocks(root_block) @@ -953,9 +982,22 @@ def apply( # pylint: disable=too-many-locals,missing-docstring index_maps = get_index_map(block_stmt) if index_maps is None: return None - matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps + + main_block_info = get_block_info(sch, main_block) + iter_infos = main_block_info.iters + + # Checks if it's a inner reduction by getting the last matrix's inner Index + def is_inner_reduction(block_stmt, iter_infos): + end_it = block_stmt.reads[-1].region[-1].min + return {it.var: it.kind for it in iter_infos}.get(end_it, "O") == "R" + + if target.kind.name == "opencl" and not is_inner_reduction(block_stmt, iter_infos): + ret = self.sch_outer_reduction(sch, config, main_block, blocks) + if ret is not None: + return ret # Step 0. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K] + matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps block = sch.reindex(main_block, ("read", 0)) sch.transform_layout(block, ("write", 0), a_index_map) block = sch.reindex(main_block, ("read", 1)) @@ -994,10 +1036,7 @@ def apply( # pylint: disable=too-many-locals,missing-docstring except: # pylint: disable=bare-except pass - # Step 2. Get schedule config. - config = self.get_configs(target) - - # Step 3. Schedule matmul + # Step 2. Schedule matmul y_kernel_size = config.vthread_y * config.block_size_y * config.micro_size_y x_kernel_size = config.vthread_x * config.block_size_x * config.micro_size_x if config.inner_x: @@ -1075,3 +1114,88 @@ def _cooperative_fetch(index, vec_len): sch.decompose_reduction(main_block, ko) return sch + + def sch_outer_reduction( + self, + sch: tir.Schedule, + config: Config, + reduction_block: tir.schedule.BlockRV, + blocks: List[tir.schedule.BlockRV], + ) -> Optional[tir.Schedule]: + reduction_loops = sch.get_loops(reduction_block) + if not len(reduction_loops) == 4: + return None + + mb, ms, n, k = reduction_loops + if not ( + isinstance(sch.get(n).extent, tir.IntImm) + and isinstance(sch.get(mb).extent, tir.IntImm) + and isinstance(sch.get(ms).extent, tir.Var) + ): + return None + + Threads_X, Threads_Y, VecSize, Unroll_M = ( + config.block_size_x, + config.block_size_y, + config.vector_size, + config.unroll, + ) + + is_dequant_block = len(blocks) > 1 + if is_dequant_block: + compute_block, dequant_block, matmul_block = blocks + sch.compute_inline(compute_block) + else: + (matmul_block,) = blocks + + m = sch.fuse(mb, ms) + + sch.pad_einsum(matmul_block, [1, Threads_Y * Unroll_M, Threads_X * VecSize, 1]) + + rmat_block, wmat_block = ( + sch.get_producers(matmul_block)[0], + sch.get_consumers(matmul_block)[0], + ) + mo, mi, mu = sch.split(m, [None, Threads_Y, Unroll_M]) + no, ni, nv = sch.split(n, [None, Threads_X, VecSize]) + k0, k1, k2, k3 = sch.split(k, [None, (Threads_X * VecSize) // 32, 4, 8]) + sch.reorder(no, mo, ni, mi, k0, k1, k2, k3, mu, nv) + + sch.compute_at(rmat_block, k0) + if is_dequant_block: + sch.compute_at(dequant_block, k3) + sch.reverse_compute_at(wmat_block, mi) + sch.set_scope(rmat_block, 0, "shared") + sch.set_scope(matmul_block, 0, "local") + if is_dequant_block: + sch.set_scope(dequant_block, 0, "local") + + sch.bind(mo, "blockIdx.y") + sch.bind(no, "blockIdx.x") + sch.bind(mi, "threadIdx.y") + sch.bind(ni, "threadIdx.x") + sch.vectorize(sch.get_loops(matmul_block)[-1]) + if is_dequant_block: + sch.vectorize(sch.get_loops(dequant_block)[-1]) + + # Co-operative Memory Fetch + ro, rv = sch.split(sch.get_loops(rmat_block)[-1], [None, VecSize]) + sch.bind(ro, "threadIdx.x") + sch.vectorize(rv) + + wv = sch.get_loops(wmat_block)[-1] + sch.vectorize(wv) + + # Scale and Quant Cache + if is_dequant_block: + qb = sch.cache_read(dequant_block, 0, "local") + sb = sch.cache_read(dequant_block, 1, "local") + sch.compute_at(sb, k1) + sch.compute_at(qb, k2) + sch.set_scope(sb, 0, "local") + sch.set_scope(qb, 0, "local") + sch.vectorize(sch.get_loops(qb)[-1]) + sch.vectorize(sch.get_loops(sb)[-1]) + + sch.decompose_reduction(matmul_block, k0) + return sch diff --git a/tests/python/dlight/test_gpu_matmul.py b/tests/python/dlight/test_gpu_matmul.py index ca32c286abfe..4cef7f1c27c3 100644 --- a/tests/python/dlight/test_gpu_matmul.py +++ b/tests/python/dlight/test_gpu_matmul.py @@ -634,42 +634,166 @@ def expected(var_inp0: T.handle, inp1: T.Buffer((T.int64(4096), T.int64(4096)), inp0 = T.match_buffer(var_inp0, (T.int64(1), m, T.int64(4096))) matmul = T.match_buffer(var_matmul, (T.int64(1), m, T.int64(4096))) # with T.block("root"): - matmul_reindex_pad_local = T.alloc_buffer((T.int64(1), (m + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), scope="local") - for ax0_ax1_0_fused in T.thread_binding((m + T.int64(31)) // T.int64(32), thread="blockIdx.y"): - for ax2_0 in T.thread_binding(T.int64(64), thread="blockIdx.x"): - for ax1_1 in T.thread_binding(T.int64(1), thread="vthread.y"): - for ax2_1 in T.thread_binding(T.int64(1), thread="vthread.x"): - for ax1_2 in T.thread_binding(T.int64(16), thread="threadIdx.y"): - for ax2_2 in T.thread_binding(T.int64(8), thread="threadIdx.x", annotations={"pragma_auto_unroll_max_step": 64, "pragma_unroll_explicit": 1}): - for ax1_3_init, ax2_3_0_init in T.grid(T.int64(2), T.int64(1)): - for ax2_3_1_init in T.vectorized(T.int64(8)): - with T.block("matmul_init"): + inp0_pad_shared = T.alloc_buffer((T.int64(1), (m + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), scope="shared") + matmul_pad_local = T.alloc_buffer((T.int64(1), (m + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), scope="local") + for i2_0 in T.thread_binding(T.int64(16), thread="blockIdx.x"): + for i0_i1_fused_0 in T.thread_binding((m + T.int64(31)) // T.int64(32), thread="blockIdx.y"): + for i2_1 in T.thread_binding(T.int64(32), thread="threadIdx.x"): + for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"): + for i0_i1_fused_2_init in range(T.int64(4)): + for i2_2_init in T.vectorized(T.int64(8)): + with T.block("matmul_init"): + v_i0 = T.axis.spatial(T.int64(1), T.int64(0)) + v_i1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + i0_i1_fused_2_init) + v_i2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2_init) + T.reads() + T.writes(matmul_pad_local[v_i0, v_i1, v_i2]) + matmul_pad_local[v_i0, v_i1, v_i2] = T.float32(0) + for k_0 in range(T.int64(16)): + for ax0 in range(T.int64(4)): + for ax1_0 in T.thread_binding(T.int64(32), thread="threadIdx.x"): + for ax1_1 in T.vectorized(T.int64(8)): + with T.block("inp0_pad"): v0 = T.axis.spatial(T.int64(1), T.int64(0)) - v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax0_ax1_0_fused * T.int64(32) + ax1_1 * T.int64(32) + ax1_2 * T.int64(2) + ax1_3_init) - v2 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(8) + ax2_3_0_init * T.int64(8) + ax2_3_1_init) - T.reads() - T.writes(matmul_reindex_pad_local[T.int64(0), v1, v2]) - matmul_reindex_pad_local[T.int64(0), v1, v2] = T.float32(0) - for ax3_0, ax3_1, ax1_3, ax2_3_0 in T.grid(T.int64(256), T.int64(16), T.int64(2), T.int64(1)): - for ax2_3_1 in T.vectorized(T.int64(8)): - with T.block("matmul_update"): + v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + ax0) + v2 = T.axis.spatial(T.int64(4096), k_0 * T.int64(256) + ax1_0 * T.int64(8) + ax1_1) + T.reads(inp0[v0, v1, v2]) + T.writes(inp0_pad_shared[v0, v1, v2]) + inp0_pad_shared[v0, v1, v2] = T.if_then_else(v1 < m, inp0[v0, v1, v2], T.float32(0)) + for k_1, k_2, k_3, i0_i1_fused_2 in T.grid(T.int64(8), T.int64(4), T.int64(8), T.int64(4)): + for i2_2 in T.vectorized(T.int64(8)): + with T.block("matmul_update"): + v_i0 = T.axis.spatial(T.int64(1), T.int64(0)) + v_i1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + i0_i1_fused_2) + v_i2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2) + v_k = T.axis.reduce(T.int64(4096), k_0 * T.int64(256) + k_1 * T.int64(32) + k_2 * T.int64(8) + k_3) + T.reads(matmul_pad_local[v_i0, v_i1, v_i2], inp0_pad_shared[v_i0, v_i1, v_k], inp1[v_k, v_i2]) + T.writes(matmul_pad_local[v_i0, v_i1, v_i2]) + matmul_pad_local[v_i0, v_i1, v_i2] = matmul_pad_local[v_i0, v_i1, v_i2] + inp0_pad_shared[v_i0, v_i1, v_k] * inp1[v_k, v_i2] + for ax0 in range(T.int64(4)): + for ax1 in T.vectorized(T.int64(8)): + with T.block("matmul_pad"): + v0 = T.axis.spatial(T.int64(1), T.int64(0)) + v1 = T.axis.spatial(m, i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + ax0) + v2 = T.axis.spatial(T.int64(4096), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax1) + T.where((i0_i1_fused_0 - (m + T.int64(31)) // T.int64(32) < T.int64(0) or i0_i1_fused_0 == T.int64(0)) and i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + ax0 < m) + T.reads(matmul_pad_local[v0, v1, v2]) + T.writes(matmul[v0, v1, v2]) + matmul[v0, v1, v2] = matmul_pad_local[v0, v1, v2] + + +class TestFusedDequantMatmulAndroid(AndroidBeforeAfter): + # fmt: off + @T.prim_func + def before(lv840: T.Buffer((T.int64(512), T.int64(12288)), "uint32"), lv841: T.Buffer((T.int64(128), T.int64(12288)), "float16"), p_rms_norm260: T.handle, p_output0: T.handle): + T.func_attr({"tir.noalias": T.bool(True)}) + seq_len = T.int64() + rms_norm260 = T.match_buffer(p_rms_norm260, (T.int64(1), seq_len, T.int64(4096)), "float16") + matmul_intermediate = T.match_buffer(p_output0, (T.int64(1), seq_len, T.int64(12288)), "float16") + # with T.block("root"): + compute = T.alloc_buffer((T.int64(4096), T.int64(12288)), "float16") + dequantize_intermediate_intermediate = T.alloc_buffer((T.int64(4096), T.int64(12288)), "float16") + for i0, i1 in T.grid(T.int64(4096), T.int64(12288)): + with T.block("compute"): + v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) + T.reads(lv840[v_i0 // T.int64(8), v_i1]) + T.writes(compute[v_i0, v_i1]) + compute[v_i0, v_i1] = T.Cast("float16", T.bitwise_and(T.shift_right(lv840[v_i0 // T.int64(8), v_i1], T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4))), T.uint32(15))) + for i0, i1 in T.grid(T.int64(4096), T.int64(12288)): + with T.block("dequantize"): + v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) + T.reads(compute[v_i0, v_i1], lv841[v_i0 // T.int64(32), v_i1]) + T.writes(dequantize_intermediate_intermediate[v_i0, v_i1]) + dequantize_intermediate_intermediate[v_i0, v_i1] = (compute[v_i0, v_i1] - T.float16(7)) * lv841[v_i0 // T.int64(32), v_i1] + for i0, i1, i2, k in T.grid(T.int64(1), seq_len, T.int64(12288), T.int64(4096)): + with T.block("matmul"): + v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k]) + T.reads(rms_norm260[v_i0, v_i1, v_k], dequantize_intermediate_intermediate[v_k, v_i2]) + T.writes(matmul_intermediate[v_i0, v_i1, v_i2]) + with T.init(): + matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0) + matmul_intermediate[v_i0, v_i1, v_i2] = matmul_intermediate[v_i0, v_i1, v_i2] + rms_norm260[v_i0, v_i1, v_k] * dequantize_intermediate_intermediate[v_k, v_i2] + + @T.prim_func + def expected(lv840: T.Buffer((T.int64(512), T.int64(12288)), "uint32"), lv841: T.Buffer((T.int64(128), T.int64(12288)), "float16"), p_rms_norm260: T.handle, p_output0: T.handle): + T.func_attr({"global_symbol": "main", "tir.is_scheduled": 1, "tir.noalias": T.bool(True)}) + seq_len = T.int64() + rms_norm260 = T.match_buffer(p_rms_norm260, (T.int64(1), seq_len, T.int64(4096)), "float16") + matmul_intermediate = T.match_buffer(p_output0, (T.int64(1), seq_len, T.int64(12288)), "float16") + # with T.block("root"): + dequantize_intermediate_intermediate_local = T.alloc_buffer((T.int64(4096), T.int64(12288)), "float16", scope="local") + rms_norm260_pad_shared = T.alloc_buffer((T.int64(1), (seq_len + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(4096)), "float16", scope="shared") + matmul_intermediate_pad_local = T.alloc_buffer((T.int64(1), (seq_len + T.int64(31)) // T.int64(32) * T.int64(32), T.int64(12288)), "float16", scope="local") + lv840_local = T.alloc_buffer((T.int64(512), T.int64(12288)), "uint32", scope="local") + lv841_local = T.alloc_buffer((T.int64(128), T.int64(12288)), "float16", scope="local") + for i2_0 in T.thread_binding(T.int64(48), thread="blockIdx.x"): + for i0_i1_fused_0 in T.thread_binding((seq_len + T.int64(31)) // T.int64(32), thread="blockIdx.y"): + for i2_1 in T.thread_binding(T.int64(32), thread="threadIdx.x"): + for i0_i1_fused_1 in T.thread_binding(T.int64(8), thread="threadIdx.y"): + for i0_i1_fused_2_init in range(T.int64(4)): + for i2_2_init in T.vectorized(T.int64(8)): + with T.block("matmul_init"): + v_i0 = T.axis.spatial(T.int64(1), T.int64(0)) + v_i1 = T.axis.spatial((seq_len + T.int64(31)) // T.int64(32) * T.int64(32), i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + i0_i1_fused_2_init) + v_i2 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2_init) + T.reads() + T.writes(matmul_intermediate_pad_local[v_i0, v_i1, v_i2]) + matmul_intermediate_pad_local[v_i0, v_i1, v_i2] = T.float16(0) + for k_0 in range(T.int64(16)): + for ax0 in range(T.int64(4)): + for ax1_0 in T.thread_binding(T.int64(32), thread="threadIdx.x"): + for ax1_1 in T.vectorized(T.int64(8)): + with T.block("rms_norm260_pad"): v0 = T.axis.spatial(T.int64(1), T.int64(0)) - v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax0_ax1_0_fused * T.int64(32) + ax1_1 * T.int64(32) + ax1_2 * T.int64(2) + ax1_3) - v2 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + ax2_1 * T.int64(64) + ax2_2 * T.int64(8) + ax2_3_0 * T.int64(8) + ax2_3_1) - v3 = T.axis.reduce(T.int64(4096), ax3_0 * T.int64(16) + ax3_1) - T.reads(matmul_reindex_pad_local[T.int64(0), v1, v2], inp0[T.int64(0), v1, v3], inp1[v3, v2]) - T.writes(matmul_reindex_pad_local[T.int64(0), v1, v2]) - matmul_reindex_pad_local[T.int64(0), v1, v2] = matmul_reindex_pad_local[T.int64(0), v1, v2] + T.if_then_else(v1 < m, inp0[T.int64(0), v1, v3], T.float32(0)) * inp1[v3, v2] - for ax0, ax1, ax2_0_1 in T.grid(T.int64(1), T.int64(2), T.int64(1)): - for ax2_1_1 in T.vectorized(T.int64(8)): - with T.block("matmul_reindex_pad_local"): - v0 = T.axis.spatial(T.int64(1), ax0) - v1 = T.axis.spatial((m + T.int64(31)) // T.int64(32) * T.int64(32), ax0_ax1_0_fused * T.int64(32) + ax1_2 * T.int64(2) + ax1) - v2 = T.axis.spatial(T.int64(4096), ax2_0 * T.int64(64) + ax2_2 * T.int64(8) + ax2_0_1 * T.int64(8) + ax2_1_1) - T.where(ax0_ax1_0_fused * T.int64(32) + ax1_2 * T.int64(2) + ax1 < m) - T.reads(matmul_reindex_pad_local[v0, v1, v2]) - T.writes(matmul[T.int64(0), v1, v2]) - matmul[T.int64(0), v1, v2] = matmul_reindex_pad_local[v0, v1, v2] + v1 = T.axis.spatial((seq_len + T.int64(31)) // T.int64(32) * T.int64(32), i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + ax0) + v2 = T.axis.spatial(T.int64(4096), k_0 * T.int64(256) + ax1_0 * T.int64(8) + ax1_1) + T.reads(rms_norm260[v0, v1, v2]) + T.writes(rms_norm260_pad_shared[v0, v1, v2]) + rms_norm260_pad_shared[v0, v1, v2] = T.if_then_else(v1 < seq_len, rms_norm260[v0, v1, v2], T.float16(0)) + for k_1 in range(T.int64(8)): + for ax0 in T.vectorized(T.int64(8)): + with T.block("lv841_local"): + v0 = T.axis.spatial(T.int64(128), k_0 * T.int64(8) + k_1) + v1 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax0) + T.reads(lv841[v0, v1]) + T.writes(lv841_local[v0, v1]) + lv841_local[v0, v1] = lv841[v0, v1] + for k_2 in range(T.int64(4)): + for ax0 in T.vectorized(T.int64(8)): + with T.block("lv840_local"): + v0 = T.axis.spatial(T.int64(512), k_0 * T.int64(32) + k_1 * T.int64(4) + k_2) + v1 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax0) + T.reads(lv840[v0, v1]) + T.writes(lv840_local[v0, v1]) + lv840_local[v0, v1] = lv840[v0, v1] + for k_3 in range(T.int64(8)): + for ax0 in T.vectorized(T.int64(8)): + with T.block("dequantize"): + v_i0 = T.axis.spatial(T.int64(4096), k_0 * T.int64(256) + k_1 * T.int64(32) + k_2 * T.int64(8) + k_3) + v_i1 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax0) + T.reads(lv840_local[v_i0 // T.int64(8), v_i1], lv841_local[v_i0 // T.int64(32), v_i1]) + T.writes(dequantize_intermediate_intermediate_local[v_i0, v_i1]) + dequantize_intermediate_intermediate_local[v_i0, v_i1] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv840_local[v_i0 // T.int64(8), v_i1], T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4))), T.uint32(15))) - T.float16(7)) * lv841_local[v_i0 // T.int64(32), v_i1] + for i0_i1_fused_2 in range(T.int64(4)): + for i2_2 in T.vectorized(T.int64(8)): + with T.block("matmul_update"): + v_i0 = T.axis.spatial(T.int64(1), T.int64(0)) + v_i1 = T.axis.spatial((seq_len + T.int64(31)) // T.int64(32) * T.int64(32), i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + i0_i1_fused_2) + v_i2 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + i2_2) + v_k = T.axis.reduce(T.int64(4096), k_0 * T.int64(256) + k_1 * T.int64(32) + k_2 * T.int64(8) + k_3) + T.reads(matmul_intermediate_pad_local[v_i0, v_i1, v_i2], rms_norm260_pad_shared[v_i0, v_i1, v_k], dequantize_intermediate_intermediate_local[v_k, v_i2]) + T.writes(matmul_intermediate_pad_local[v_i0, v_i1, v_i2]) + matmul_intermediate_pad_local[v_i0, v_i1, v_i2] = matmul_intermediate_pad_local[v_i0, v_i1, v_i2] + rms_norm260_pad_shared[v_i0, v_i1, v_k] * dequantize_intermediate_intermediate_local[v_k, v_i2] + for ax0 in range(T.int64(4)): + for ax1 in T.vectorized(T.int64(8)): + with T.block("matmul_intermediate_pad"): + v0 = T.axis.spatial(T.int64(1), T.int64(0)) + v1 = T.axis.spatial(seq_len, i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + ax0) + v2 = T.axis.spatial(T.int64(12288), i2_0 * T.int64(256) + i2_1 * T.int64(8) + ax1) + T.where((i0_i1_fused_0 - (seq_len + T.int64(31)) // T.int64(32) < T.int64(0) or i0_i1_fused_0 == T.int64(0)) and i0_i1_fused_0 * T.int64(32) + i0_i1_fused_1 * T.int64(4) + ax0 < seq_len) + T.reads(matmul_intermediate_pad_local[v0, v1, v2]) + T.writes(matmul_intermediate[v0, v1, v2]) + matmul_intermediate[v0, v1, v2] = matmul_intermediate_pad_local[v0, v1, v2] # fmt: on