diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index 5f5926917e..1eee51019f 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -95,40 +95,54 @@ def mma_bf16( return cute.TensorSSA(vec, 4, Float32) -@dsl_user_op -def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32: +def _bf16x2_unary(asm: str, a: Uint32, *, loc=None, ip=None) -> Uint32: out = llvm.inline_asm( T.i32(), [a.ir_value(loc=loc, ip=ip)], - "abs.bf16x2 $0, $1;", + f"{asm}.bf16x2 $0, $1;", "=r,r", has_side_effects=False, is_align_stack=False, + loc=loc, + ip=ip, ) return Uint32(out) +def _bf16x2_binary(asm: str, a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: + out = llvm.inline_asm( + T.i32(), + [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], + f"{asm}.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + loc=loc, + ip=ip, + ) + return Uint32(out) + + +@dsl_user_op +def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32: + return _bf16x2_unary("abs", a, loc=loc, ip=ip) + + +@dsl_user_op +def _bf16x2_neg(a: Uint32, *, loc=None, ip=None) -> Uint32: + return _bf16x2_unary("neg", a, loc=loc, ip=ip) + + @dsl_user_op def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: - out = llvm.inline_asm( - T.i32(), - [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], - "max.bf16x2 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - ) - return Uint32(out) + return _bf16x2_binary("max", a, b, loc=loc, ip=ip) @dsl_user_op def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: - out = llvm.inline_asm( - T.i32(), - [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)], - "mul.rn.bf16x2 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - ) - return Uint32(out) + return _bf16x2_binary("mul.rn", a, b, loc=loc, ip=ip) + + +@dsl_user_op +def _bf16x2_sub(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: + return _bf16x2_binary("sub.rn", a, b, loc=loc, ip=ip) diff --git a/vllm/cute_utils/_tcgen05.py b/vllm/cute_utils/_tcgen05.py index 01dc998fd0..9367fa12a1 100644 --- a/vllm/cute_utils/_tcgen05.py +++ b/vllm/cute_utils/_tcgen05.py @@ -90,17 +90,18 @@ def mma_f16( loc=None, ip=None, ) -> None: - nvvm.tcgen05_mma( - nvvm.Tcgen05MMAKind.F16, - NVVM_CTA_GROUP_MAP[cta_group], - _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), - Uint64(a_desc).ir_value(loc=loc, ip=ip), - Uint64(b_desc).ir_value(loc=loc, ip=ip), - Int32(idesc).ir_value(loc=loc, ip=ip), - Boolean(enable_input_d).ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) + with cute.arch.elect_one(): + nvvm.tcgen05_mma( + nvvm.Tcgen05MMAKind.F16, + NVVM_CTA_GROUP_MAP[cta_group], + _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), + Uint64(a_desc).ir_value(loc=loc, ip=ip), + Uint64(b_desc).ir_value(loc=loc, ip=ip), + Int32(idesc).ir_value(loc=loc, ip=ip), + Boolean(enable_input_d).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) @dsl_user_op @@ -115,17 +116,18 @@ def mma_ts_f16( loc=None, ip=None, ) -> None: - nvvm.tcgen05_mma( - nvvm.Tcgen05MMAKind.F16, - NVVM_CTA_GROUP_MAP[cta_group], - _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), - _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip), - Uint64(b_desc).ir_value(loc=loc, ip=ip), - Int32(idesc).ir_value(loc=loc, ip=ip), - Boolean(enable_input_d).ir_value(loc=loc, ip=ip), - loc=loc, - ip=ip, - ) + with cute.arch.elect_one(): + nvvm.tcgen05_mma( + nvvm.Tcgen05MMAKind.F16, + NVVM_CTA_GROUP_MAP[cta_group], + _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), + _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip), + Uint64(b_desc).ir_value(loc=loc, ip=ip), + Int32(idesc).ir_value(loc=loc, ip=ip), + Boolean(enable_input_d).ir_value(loc=loc, ip=ip), + loc=loc, + ip=ip, + ) @dsl_user_op @@ -133,15 +135,17 @@ def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None): mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip) group = NVVM_CTA_GROUP_MAP[cta_group] if cutlass.const_expr(cta_mask is not None): - nvvm.tcgen05_commit_arrive( - mbar_llvm, - multicast_mask=cta_mask.ir_value(loc=loc, ip=ip), - group=group, - loc=loc, - ip=ip, - ) + with cute.arch.elect_one(): + nvvm.tcgen05_commit_arrive( + mbar_llvm, + multicast_mask=cta_mask.ir_value(loc=loc, ip=ip), + group=group, + loc=loc, + ip=ip, + ) else: - nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip) + with cute.arch.elect_one(): + nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip) @dsl_user_op diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py index 040788da3a..992f4ff68b 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py @@ -174,31 +174,31 @@ def chunk_gated_delta_rule_cutedsl( When ``core_attn_out`` is provided, ``output`` is an unsqueezed view of that buffer. """ - q_3d = q.squeeze(0) - k_3d = k.squeeze(0) - v_3d = v.squeeze(0) - g_2d = g.squeeze(0) - beta_2d = beta.squeeze(0) + q = q.squeeze(0) + k = k.squeeze(0) + v = v.squeeze(0) + g = g.squeeze(0) + beta = beta.squeeze(0) - _, _, head_k_dim = k_3d.shape - _, num_v_heads, head_v_dim = v_3d.shape + _, _, K_dim = k.shape + _, num_v_heads, V_dim = v.shape chunk_size = 64 upper_bound_chunks = chunk_indices.shape[0] pad_t = upper_bound_chunks * chunk_size total_chunks_ptr = chunk_offsets[-1:] - g_cu = torch.empty_like(g_2d, dtype=torch.float32) - u = q_3d.new_empty(pad_t, num_v_heads, head_v_dim) - w = q_3d.new_empty(pad_t, num_v_heads, head_k_dim) + g_cu = torch.empty_like(g, dtype=torch.float32) + u = q.new_empty(pad_t, num_v_heads, V_dim) + w = q.new_empty(pad_t, num_v_heads, K_dim) num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count kkt_inv_uw_cutedsl( - k_3d, - v_3d, + k, + v, u, w, - g_2d, - beta_2d, + g, + beta, g_cu, cu_seqlens, chunk_indices, @@ -206,16 +206,11 @@ def chunk_gated_delta_rule_cutedsl( num_sms=num_sms, ) - h = k_3d.new_empty( - upper_bound_chunks, - num_v_heads, - head_v_dim, - head_k_dim, - ) - v_new = q_3d.new_empty(pad_t, num_v_heads, head_v_dim) + h = k.new_empty(upper_bound_chunks, num_v_heads, V_dim, K_dim) + v_new = q.new_empty(pad_t, num_v_heads, V_dim) final_state = torch.empty_like(initial_state) h_cutedsl( - k_3d, + k, u, w, v_new, @@ -227,12 +222,12 @@ def chunk_gated_delta_rule_cutedsl( chunk_offsets, ) - output = core_attn_out if core_attn_out is not None else torch.empty_like(v_3d) - scale = head_k_dim**-0.5 + output = core_attn_out if core_attn_out is not None else torch.empty_like(v) + scale = K_dim**-0.5 o_cutedsl( - q_3d, - k_3d, - v_new.view(upper_bound_chunks, chunk_size, num_v_heads, head_v_dim), + q, + k, + v_new, h, g_cu, output, diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py index 5dff6b6dfd..7cdd652e7d 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py @@ -66,13 +66,12 @@ class Sm100ChunkHKernel: stride=(64, 0, (1, self.BT * 64), self.BT * dim), ) slayout = cute.make_composed_layout(swizzle_128B, 0, slayout) - atom, tma_tensor = cpasync.make_tiled_tma_atom( + return cpasync.make_tiled_tma_atom( op, cute.logical_divide(tensor, (None, None, 64)), slayout, cta_tiler=(self.BT, 1, dim), ) - return atom, tma_tensor, slayout @cute.jit def _make_h_tma_args(self, tensor: cute.Tensor, op: cpasync.TmaCopyOp): @@ -84,13 +83,12 @@ class Sm100ChunkHKernel: stride=(0, 0, num_elems, (1, self.V_dim * num_elems)), ) slayout = cute.make_composed_layout(swizzle_128B, 0, slayout) - atom, tma_tensor = cpasync.make_tiled_tma_atom( + return cpasync.make_tiled_tma_atom( op, cute.logical_divide(tensor, (None, None, None, num_elems)), slayout, cta_tiler=(1, 1, self.V_dim, self.K_dim), ) - return atom, tma_tensor, slayout @cute.jit def __call__( @@ -110,39 +108,39 @@ class Sm100ChunkHKernel: tma_g2s = cpasync.CopyBulkTensorTileG2SOp() tma_s2g = cpasync.CopyBulkTensorTileS2GOp() - K_args = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages) - V_args = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages) - W_args = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages) - V_new_args = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1) - H0_args = self._make_h_tma_args(h0, tma_g2s) - HT_args = self._make_h_tma_args(ht, tma_s2g) - H_args = self._make_h_tma_args(h, tma_s2g) + K_tma = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages) + V_tma = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages) + W_tma = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages) + V_new_tma = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1) + H0_tma = self._make_h_tma_args(h0, tma_g2s) + HT_tma = self._make_h_tma_args(ht, tma_s2g) + H_tma = self._make_h_tma_args(h, tma_s2g) grid = (self.Hv, h0.shape[0], 1) block = (self.num_warps * 32, 1, 1) self.kernel( - K_args, - V_args, - W_args, - V_new_args, - H0_args, - HT_args, - H_args, + K_tma, + V_tma, + W_tma, + V_new_tma, + H0_tma, + HT_tma, + H_tma, g_cu, cu_seqlens, chunk_offsets, - ).launch(grid=grid, block=block, stream=stream) + ).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream) @cute.kernel def kernel( self, - K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - V_new_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - H0_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - HT_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], + K_tma: cpasync.TmaInfo, + V_tma: cpasync.TmaInfo, + W_tma: cpasync.TmaInfo, + V_new_tma: cpasync.TmaInfo, + H0_tma: cpasync.TmaInfo, + HT_tma: cpasync.TmaInfo, + H_tma: cpasync.TmaInfo, g_cu: cute.Tensor, cu_seqlens: cute.Tensor, chunk_offsets: cute.Tensor, @@ -158,14 +156,6 @@ class Sm100ChunkHKernel: num_stages = self.num_stages is_f32 = self.h_dtype == Float32 - K_tma_atom, tmaK, sK_layout = K_args - V_tma_atom, tmaV, sV_layout = V_args - W_tma_atom, tmaW, sW_layout = W_args - V_new_tma_atom, tmaV_new, sV_new_layout = V_new_args - H0_tma_atom, tmaH0, sH0_layout = H0_args - HT_tma_atom, tmaHT, _ = HT_args - H_tma_atom, tmaH, sH_layout = H_args - def allocate_tensor(smem, dtype, layout): return smem.allocate_tensor( dtype, layout.outer, byte_alignment=128, swizzle=layout.inner @@ -174,12 +164,14 @@ class Sm100ChunkHKernel: smem = cutlass.utils.SmemAllocator() # remove size=1 modes - sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, None] - sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None] - sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None] - sH0 = allocate_tensor(smem, self.h_dtype, sH0_layout)[0, 0, None, None] - sH = allocate_tensor(smem, BFloat16, sH_layout)[0, 0, None, None] - sV_new = allocate_tensor(smem, BFloat16, sV_new_layout)[None, 0, None, 0] + sW = allocate_tensor(smem, BFloat16, W_tma.smem_layout)[None, 0, None, None] + sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None] + sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None] + sH0 = allocate_tensor(smem, self.h_dtype, H0_tma.smem_layout)[0, 0, None, None] + sH = allocate_tensor(smem, BFloat16, H_tma.smem_layout)[0, 0, None, None] + sV_new = allocate_tensor(smem, BFloat16, V_new_tma.smem_layout)[ + None, 0, None, 0 + ] s_v_scale = smem.allocate_array(Float32, BT) tma_mbar = smem.allocate_array(Int64, num_stages) @@ -206,13 +198,13 @@ class Sm100ChunkHKernel: cute.arch.mbarrier_init(h0_mbar, 1) cute.arch.mbarrier_init_fence() elif warp_id == 1: - cpasync.prefetch_descriptor(H0_tma_atom) - cpasync.prefetch_descriptor(W_tma_atom) - cpasync.prefetch_descriptor(V_tma_atom) - cpasync.prefetch_descriptor(K_tma_atom) - cpasync.prefetch_descriptor(HT_tma_atom) - cpasync.prefetch_descriptor(H_tma_atom) - cpasync.prefetch_descriptor(V_new_tma_atom) + cpasync.prefetch_descriptor(H0_tma.atom) + cpasync.prefetch_descriptor(W_tma.atom) + cpasync.prefetch_descriptor(V_tma.atom) + cpasync.prefetch_descriptor(K_tma.atom) + cpasync.prefetch_descriptor(HT_tma.atom) + cpasync.prefetch_descriptor(H_tma.atom) + cpasync.prefetch_descriptor(V_new_tma.atom) cute.arch.sync_threads() bos = cu_seqlens[seq_id] @@ -233,14 +225,21 @@ class Sm100ChunkHKernel: H0_size = V_dim * K_dim * self.h_dtype.width // 8 cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size) simple_tma_copy( - H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar + H0_tma.atom, + H0_tma.tma_tensor[seq_id, head_id, None, None], + sH0, + h0_mbar, ) # shape: ((BT, num_BT_tiles), (64, 2)) - gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None)) - gV_tiles = cute.logical_divide(tmaV[None, head_id, None], (BT, None)) + gW_tiles = cute.logical_divide( + W_tma.tma_tensor[None, head_id, None], (BT, None) + ) + gV_tiles = cute.logical_divide( + V_tma.tma_tensor[None, head_id, None], (BT, None) + ) gK_tiles = cute.logical_divide( - cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]), + cute.domain_offset((bos, 0), K_tma.tma_tensor[None, k_head_id, None]), (BT, None), ) @@ -258,12 +257,12 @@ class Sm100ChunkHKernel: STAGE_SIZE = BT * (K_dim + V_dim + K_dim) * 2 cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE) simple_tma_copy( - W_tma_atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST + W_tma.atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST ) simple_tma_copy( - V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST + V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST ) - simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar) + simple_tma_copy(K_tma.atom, gK, sK[None, None, stage_id], mbar) stage_id = (stage_id + 1) % num_stages if stage_id == 0: @@ -295,13 +294,12 @@ class Sm100ChunkHKernel: cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(K_dim // 64): - for j in cutlass.range_constexpr(64 // 16): - hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4) - wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4) - _tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True) - _tcgen05.commit(wh_done_mbar + stage_id) + for i in cutlass.range_constexpr(K_dim // 64): + for j in cutlass.range_constexpr(64 // 16): + hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4) + wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4) + _tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True) + _tcgen05.commit(wh_done_mbar + stage_id) ##### 2nd MMA: H_new = H + V_new.T @ K ##### Kaddr0 = sK[None, None, stage_id].iterator.toint() @@ -310,12 +308,11 @@ class Sm100ChunkHKernel: cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for k in cutlass.range_constexpr(BT // 16): - vtmem0 = v_tmem_base + k * 8 - kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4) - _tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True) - _tcgen05.commit(vk_done_mbar + stage_id) + for k in cutlass.range_constexpr(BT // 16): + vtmem0 = v_tmem_base + k * 8 + kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4) + _tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True) + _tcgen05.commit(vk_done_mbar + stage_id) stage_id = (stage_id + 1) % num_stages if stage_id == 0: @@ -331,13 +328,12 @@ class Sm100ChunkHKernel: cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(K_dim // 64): - for j in cutlass.range_constexpr(64 // 16): - htmem = h_tmem_base + i * 32 + j * 8 - wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4) - _tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True) - _tcgen05.commit(wh_done_mbar + stage_id) + for i in cutlass.range_constexpr(K_dim // 64): + for j in cutlass.range_constexpr(64 // 16): + htmem = h_tmem_base + i * 32 + j * 8 + wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4) + _tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True) + _tcgen05.commit(wh_done_mbar + stage_id) ##### 2nd MMA: H_new = H + V_new.T @ K ##### Kaddr = sK[None, None, stage_id].iterator.toint() @@ -346,12 +342,11 @@ class Sm100ChunkHKernel: cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for k in cutlass.range_constexpr(BT // 16): - vtmem = v_tmem_base + k * 8 - kdesc = kdesc_base | ((k * 16 * 128) >> 4) - _tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True) - _tcgen05.commit(vk_done_mbar + stage_id) + for k in cutlass.range_constexpr(BT // 16): + vtmem = v_tmem_base + k * 8 + kdesc = kdesc_base | ((k * 16 * 128) >> 4) + _tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True) + _tcgen05.commit(vk_done_mbar + stage_id) stage_id = (stage_id + 1) % num_stages if stage_id == 0: @@ -421,7 +416,7 @@ class Sm100ChunkHKernel: ).load() ) - for j in cutlass.range_constexpr(32): + for j in cutlass.range(32, vectorize=True): h_f32[j] *= h_scale _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32) @@ -435,8 +430,10 @@ class Sm100ChunkHKernel: fence_before_tma_store() if warp_id_ == 3: h_src = sH if cutlass.const_expr(is_f32) else sH0 - h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None] - simple_tma_copy(H_tma_atom, h_src, h_dst) + h_dst = H_tma.tma_tensor[ + chunk_offset + chunk_id, head_id, None, None + ] + simple_tma_copy(H_tma.atom, h_src, h_dst) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() @@ -490,7 +487,7 @@ class Sm100ChunkHKernel: h_f32.store( _tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32) ) - for j in cutlass.range_constexpr(32): + for j in cutlass.range(32, vectorize=True): h_f32[j] *= h_scale _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32) _tcgen05.wait_st() @@ -501,8 +498,10 @@ class Sm100ChunkHKernel: cute.arch.barrier(barrier_id=1, number_of_threads=128) fence_before_tma_store() if warp_id_ == 3: - h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None] - simple_tma_copy(H_tma_atom, sH, h_dst) + h_dst = H_tma.tma_tensor[ + chunk_offset + chunk_id, head_id, None, None + ] + simple_tma_copy(H_tma.atom, sH, h_dst) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() @@ -530,8 +529,8 @@ class Sm100ChunkHKernel: cute.arch.barrier(barrier_id=1, number_of_threads=128) if warp_id_ == 0: - ht_dst = tmaHT[seq_id, head_id, None, None] - simple_tma_copy(HT_tma_atom, sH0, ht_dst) + ht_dst = HT_tma.tma_tensor[seq_id, head_id, None, None] + simple_tma_copy(HT_tma.atom, sH0, ht_dst) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() if warp_id_ == 1: @@ -551,7 +550,7 @@ class Sm100ChunkHKernel: # ((BT, num_BT_tiles), V_dim) gV_new_tiles = cute.logical_divide( - tmaV_new[None, head_id, None], (BT, None) + V_new_tma.tma_tensor[None, head_id, None], (BT, None) ) # sV shape: [BT, (64, V_dim/64), num_stages] @@ -653,7 +652,7 @@ class Sm100ChunkHKernel: fence_before_tma_store() if warp_id == 3: gV = gV_new_tiles[(None, chunk_offset + chunk_id), None] - simple_tma_copy(V_new_tma_atom, sV_new, gV) + simple_tma_copy(V_new_tma.atom, sV_new, gV) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py index 7e42d80b64..2066c04f52 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py @@ -11,6 +11,8 @@ from quack.compile_utils import make_fake_tensor from vllm.cute_utils import ( EVICT_FIRST, + _bf16x2_neg, + _bf16x2_sub, _tcgen05, cvt, fence_before_tma_store, @@ -47,7 +49,7 @@ class Sm100ChunkUWKernel: # hard-code self.BT = 64 - self.num_warps = 2 + 4 + 4 + self.num_warps = 4 + 4 + 4 @cute.jit def _make_tma_args( @@ -68,13 +70,12 @@ class Sm100ChunkUWKernel: # we need to convert gmem layout to (T, H, (64, D/64)) for make_tiled_tma_atom() # to emit a single 4D TMA. otherwise, it will emit (D/64)x 3D TMA. - atom, tma_tensor = cpasync.make_tiled_tma_atom( + return cpasync.make_tiled_tma_atom( op, cute.logical_divide(tensor, (None, None, 64)), slayout, cta_tiler=(self.BT, 1, dim), ) - return atom, tma_tensor, slayout @cute.jit def __call__( @@ -95,33 +96,33 @@ class Sm100ChunkUWKernel: tma_g2s = cpasync.CopyBulkTensorTileG2SOp() tma_s2g = cpasync.CopyBulkTensorTileS2GOp() - K_args = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s) - V_args = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s) - U_args = self._make_tma_args(U, self.V_dim, 1, tma_s2g) - W_args = self._make_tma_args(W, self.K_dim, 1, tma_s2g) + K_tma = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s) + V_tma = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s) + U_tma = self._make_tma_args(U, self.V_dim, 1, tma_s2g) + W_tma = self._make_tma_args(W, self.K_dim, 1, tma_s2g) grid = (num_sms // self.Hv, self.Hv, 1) block = (self.num_warps * 32, 1, 1) self.kernel( - K_args, - V_args, - U_args, - W_args, + K_tma, + V_tma, + U_tma, + W_tma, g, beta, g_cu, cu_seqlens, chunk_indices, total_chunks, - ).launch(grid=grid, block=block, stream=stream) + ).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream) @cute.kernel def kernel( self, - K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - U_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], + K_tma: cpasync.TmaInfo, + V_tma: cpasync.TmaInfo, + U_tma: cpasync.TmaInfo, + W_tma: cpasync.TmaInfo, g: cute.Tensor, beta: cute.Tensor, g_cu: cute.Tensor, @@ -142,10 +143,10 @@ class Sm100ChunkUWKernel: V_dim = self.V_dim num_stages = self.num_stages - K_tma_atom, tmaK, sK_layout = K_args - V_tma_atom, tmaV, sV_layout = V_args - U_tma_atom, tmaU, sU_layout = U_args - W_tma_atom, tmaW, sW_layout = W_args + INV_BAR = 1 + EPI_BAR = 2 + SCAN_BAR = 3 + TMEM_ALLOC_BAR = 4 def allocate_tensor(smem, dtype, layout): return smem.allocate_tensor( @@ -153,22 +154,20 @@ class Sm100ChunkUWKernel: ) smem = cutlass.utils.SmemAllocator() - sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None] - sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None] - sU = allocate_tensor(smem, BFloat16, sU_layout)[None, 0, None, 0] - sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, 0] + sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None] + sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None] + sU = allocate_tensor(smem, BFloat16, U_tma.smem_layout)[None, 0, None, 0] + sW = allocate_tensor(smem, BFloat16, W_tma.smem_layout)[None, 0, None, 0] - swizzle_128B = cute.make_swizzle(3, 4, 3) - sA_layout = cute.make_layout((BT, (64, 1)), stride=(64, (1, BT * 64))) - sA_layout = cute.make_composed_layout(swizzle_128B, 0, sA_layout) - sA = allocate_tensor(smem, BFloat16, sA_layout) - sAi = allocate_tensor(smem, BFloat16, sA_layout) + sA_ptr = smem.allocate_array(BFloat16, BT * BT, byte_alignment=16) + sAi_ptr = smem.allocate_array(BFloat16, BT * BT, byte_alignment=16) - s_beta = smem.allocate_array(Float32, BT) - s_g_cu_exp = smem.allocate_array(Float32, BT) - s_g_cu = smem.allocate_array(Float32, BT) + s_beta = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages))) + s_g_cu = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages))) + s_beta_g = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages))) tma_mbar = smem.allocate_array(Int64, num_stages) + prep_mbar = smem.allocate_array(Int64, num_stages) mma_kkt_mbar = smem.allocate_array(Int64, num_stages) inv_mbar = smem.allocate_array(Int64, num_stages) mma_u_mbar = smem.allocate_array(Int64, num_stages) @@ -179,7 +178,7 @@ class Sm100ChunkUWKernel: kkt_tmem = 0 U_tmem_base = kkt_tmem + BT Ab_tmem_base = U_tmem_base + V_dim * num_stages - assert Ab_tmem_base + (BT // 2) * num_stages <= 512 + assert Ab_tmem_base + (BT // 2) <= 512 # prepare ldmatrix/stmatrix ops ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) @@ -193,6 +192,7 @@ class Sm100ChunkUWKernel: with cute.arch.elect_one(): for i in cutlass.range_constexpr(num_stages): cute.arch.mbarrier_init(tma_mbar + i, 1) + cute.arch.mbarrier_init(prep_mbar + i, 64) cute.arch.mbarrier_init(mma_kkt_mbar + i, 1) cute.arch.mbarrier_init(inv_mbar + i, 128) cute.arch.mbarrier_init(mma_u_mbar + i, 1) @@ -200,14 +200,14 @@ class Sm100ChunkUWKernel: cute.arch.mbarrier_init(epi_mbar + i, 128) cute.arch.mbarrier_init_fence() elif warp_id == 1: - cpasync.prefetch_descriptor(K_tma_atom) - cpasync.prefetch_descriptor(V_tma_atom) - cpasync.prefetch_descriptor(U_tma_atom) - cpasync.prefetch_descriptor(W_tma_atom) + cpasync.prefetch_descriptor(K_tma.atom) + cpasync.prefetch_descriptor(V_tma.atom) + cpasync.prefetch_descriptor(U_tma.atom) + cpasync.prefetch_descriptor(W_tma.atom) cute.arch.sync_threads() num_global_chunks = total_chunks[0] - if warp_id == 9: + if warp_id == 11: # TMA warp stage_id = 0 parity = 1 @@ -221,12 +221,14 @@ class Sm100ChunkUWKernel: # domain_offset() to shift the pointer first. mbar = tma_mbar + stage_id gK = cute.local_tile( - cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]), + cute.domain_offset( + (bos, 0), K_tma.tma_tensor[None, k_head_id, None] + ), tiler=(BT, K_dim), coord=(chunk_id, 0), ) gV = cute.local_tile( - cute.domain_offset((bos, 0), tmaV[None, head_id, None]), + cute.domain_offset((bos, 0), V_tma.tma_tensor[None, head_id, None]), tiler=(BT, V_dim), coord=(chunk_id, 0), ) @@ -237,18 +239,19 @@ class Sm100ChunkUWKernel: with cute.arch.elect_one(): STAGE_SIZE = BT * (K_dim + V_dim) * 2 cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE) - simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar) + simple_tma_copy(K_tma.atom, gK, sK[None, None, stage_id], mbar) simple_tma_copy( - V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST + V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST ) stage_id = (stage_id + 1) % num_stages if stage_id == 0: parity ^= 1 - elif warp_id == 8: + elif warp_id == 10: # MMA warp _tcgen05.alloc(taddr) + cute.arch.barrier(barrier_id=TMEM_ALLOC_BAR, number_of_threads=160) stage_id = 0 parity = 0 @@ -263,8 +266,8 @@ class Sm100ChunkUWKernel: for global_chunk_id in range(bid, num_global_chunks, grid_x): U_tmem = U_tmem_base + V_dim * stage_id W_tmem = U_tmem | (16 << 16) - Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id - Abg_tmem = Ab_tmem | (16 << 16) + Ab_tmem = Ab_tmem_base + Abg_tmem = Ab_tmem_base | (16 << 16) ##### KKT MMA: KKT = K @ K.T ##### kaddr = sK[None, None, stage_id].iterator.toint() @@ -276,18 +279,12 @@ class Sm100ChunkUWKernel: cute.arch.mbarrier_wait(tma_mbar + stage_id, parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(K_dim // 64): - for j in cutlass.range_constexpr(64 // 16): - kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4) - _tcgen05.mma_f16( - kkt_tmem, - kdesc, - kdesc, - kkt_idesc, - (i > 0) or (j > 0), - ) - _tcgen05.commit(mma_kkt_mbar + stage_id) + for i in cutlass.range_constexpr(K_dim // 64): + for j in cutlass.range_constexpr(64 // 16): + kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4) + enable_d = (i > 0) or (j > 0) + _tcgen05.mma_f16(kkt_tmem, kdesc, kdesc, kkt_idesc, enable_d) + _tcgen05.commit(mma_kkt_mbar + stage_id) ##### U/W MMA: U = Ab @ V, W = Abg @ K ##### vaddr = sV[None, None, stage_id].iterator.toint() @@ -299,20 +296,15 @@ class Sm100ChunkUWKernel: cute.arch.mbarrier_wait(inv_mbar + stage_id, parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(BT // 16): - _tcgen05.mma_ts_f16( - W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0 - ) - kdesc += (16 * 128) >> 4 - _tcgen05.commit(mma_w_mbar + stage_id) + for i in cutlass.range_constexpr(BT // 16): + _tcgen05.mma_ts_f16(W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0) + kdesc += (16 * 128) >> 4 + _tcgen05.commit(mma_w_mbar + stage_id) - for i in cutlass.range_constexpr(BT // 16): - _tcgen05.mma_ts_f16( - U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0 - ) - vdesc += (16 * 128) >> 4 - _tcgen05.commit(mma_u_mbar + stage_id) + for i in cutlass.range_constexpr(BT // 16): + _tcgen05.mma_ts_f16(U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0) + vdesc += (16 * 128) >> 4 + _tcgen05.commit(mma_u_mbar + stage_id) stage_id = (stage_id + 1) % num_stages if stage_id == 0: @@ -321,24 +313,124 @@ class Sm100ChunkUWKernel: cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1) _tcgen05.dealloc() + elif warp_id >= 8: + # dedicated prep warps for beta and gate, consumed by INV warps + stage_id = 0 + parity = 0 + tid_ = tid % 128 + warp_id_ = warp_id % 4 + + for global_chunk_id in range(bid, num_global_chunks, grid_x): + seq_id = chunk_indices[global_chunk_id, 0] + chunk_id = chunk_indices[global_chunk_id, 1] + bos = cu_seqlens[seq_id] + eos = cu_seqlens[seq_id + 1] + off_t = bos + chunk_id * BT + t = off_t + tid_ + + in_bounds = t < eos + beta_val = beta[t, head_id] if in_bounds else Float32(0.0) + g_val = g[t, head_id] if in_bounds else Float32(0.0) + + # warp-local prefix scan + for i in cutlass.range_constexpr(5): + offset = cutlass.const_expr(1 << i) + lower = cute.arch.shuffle_sync_up(g_val, offset, mask_and_clamp=0) + if lane_id >= offset: + g_val += lower + + # Delay the stage-reuse wait until just before touching smem: + # global loads and the warp-local scan do not use staged buffers. + if warp_id_ == 0: + cute.arch.mbarrier_wait(inv_mbar + stage_id, parity ^ 1) + cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT) + + # Store beta and the per-warp scan totals for the cross-warp fixup. + s_beta[tid_, stage_id] = beta_val + if lane_id == 31: + s_g_cu[warp_id_, stage_id] = g_val + cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT) + + # Add the sum from the lower prep warp. + if warp_id_ == 1: + g_val += s_g_cu[0, stage_id] + cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT) + + if in_bounds: + g_cu[t, head_id] = g_val + + s_g_cu[tid_, stage_id] = g_val + s_beta_g[tid_, stage_id] = beta_val * cute.math.exp(g_val) + cute.arch.mbarrier_arrive(prep_mbar + stage_id) + + stage_id = (stage_id + 1) % num_stages + if stage_id == 0: + parity ^= 1 + elif warp_id >= 4: # inv warps tid_ = tid % 128 warp_id_ = warp_id % 4 + def store_ab_abg( + Ai_f32, s_beta, s_beta_g, warp_id_, lane_id, tile_col, Ab_tmem_base + ): + # compute Ab and Abg from Ai, then store to tmem + beta_col = cute.make_rmem_tensor((2, 2), Float32) + beta_g_col = cute.make_rmem_tensor((2, 2), Float32) + + for i in cutlass.range_constexpr(2): + base = tile_col * 16 + i * 8 + (lane_id % 4) * 2 + for j in cutlass.range_constexpr(2): + beta_col[j, i] = s_beta[base + j] + beta_g_col[j, i] = s_beta_g[base + j] + + # without conversion to TensorSSA, cutlass.range(vectorize=True) fails + beta_col = beta_col.load() + beta_g_col = beta_g_col.load() + + Ab_f32 = cute.make_rmem_tensor(8, Float32) + Abg_f32 = cute.make_rmem_tensor(8, Float32) + for i in cutlass.range(8, vectorize=True): + scale_idx = (i // 4) * 2 + (i % 2) + Ab_f32[i] = Ai_f32[i] * beta_col[scale_idx] + Abg_f32[i] = Ai_f32[i] * beta_g_col[scale_idx] + + Ab = Ab_f32.load().to(BFloat16) + Abg = Abg_f32.load().to(BFloat16) + Ab_tmem = Ab_tmem_base + tile_col * 8 + _tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab) + _tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg) + + # clear the Ab/Abg tmem buffer once before mainloop + # this is to keep the upper triangular tiles zeros + cute.arch.barrier(barrier_id=TMEM_ALLOC_BAR, number_of_threads=160) + zeros_ = cute.make_rmem_tensor(BT // 2, Float32) + zeros_.fill(0.0) + _tcgen05.st(warp_id_ * 32, Ab_tmem_base, "32x32b", BT // 2, zeros_) + stage_id = 0 parity = 0 - # view into (16,16) sub-tiles, then ldmatrix layout - sA_ldsm = cute.logical_divide(sA, (16, cute.make_layout((8, 2)))) - sAi_ldsm = cute.logical_divide(sAi, (16, cute.make_layout((8, 2)))) - sA_ldsm = sA_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)] - sAi_ldsm = sAi_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)] + # for sA, we can avoid bank conflict without using swizzling because + # this is only used as tmp buffer between rmem<->smem, no gmem interactions. + # to do so, we can put (8,8) tile contiguous in memory. logically, we are + # partitioning (64,64) tile into 4x (16,16) tiles. to make indexing easier + # later, we view (16,16) tile as (32,8) tile here. + # (4,4) is (row_tile, col_tile). + sA_layout = cute.make_layout(((8, 32), (4, 4))) + sA = cute.make_tensor(sA_ptr, sA_layout) + sAi = cute.make_tensor(sAi_ptr, sA_layout) + + # pre-compute ldmatrix addresses + sA_ldsm = sA[(None, lane_id), None] + sAi_ldsm = sAi[(None, lane_id), None] # init Ai smem buffer with zeros (only the first 48 rows) - for i in cutlass.range_constexpr((BT // 4 * 3) * BT // 128): - idx = i * 128 + tid_ - sAi[idx // BT, idx % BT] = BFloat16(0.0) + zeros_bf16 = cute.make_rmem_tensor(8, BFloat16) + zeros_bf16.fill(0.0) + for i in cutlass.range_constexpr(3): + cute.copy(stsm_atom, zeros_bf16, sAi_ldsm[None, (i, warp_id_)]) # indices for ldmatrix layout later row_indices = cute.make_rmem_tensor((1, 2, 1), Int32) @@ -360,61 +452,31 @@ class Sm100ChunkUWKernel: eos = cu_seqlens[seq_id + 1] off_t = bos + chunk_id * BT - t = off_t + tid_ - - ##### Phase 1: load g and beta ##### - if tid_ < BT: - in_bounds = t < eos - beta_val = beta[t, head_id] if in_bounds else Float32(0.0) - g_val = g[t, head_id] if in_bounds else Float32(0.0) - - s_beta[tid_] = beta_val - - # compute cumsum(g) - # parallel scan within a warp - for i in cutlass.range_constexpr(5): - offset = cutlass.const_expr(1 << i) - lower = cute.arch.shuffle_sync_up( - g_val, offset, mask_and_clamp=0 - ) - if lane_id >= offset: - g_val += lower - - # store warp sum - if lane_id == 31: - s_g_cu[warp_id_] = g_val - cute.arch.barrier(barrier_id=3, number_of_threads=BT) - - # add warp sum from lower warps - for i in cutlass.range_constexpr(1, BT // 32): - if warp_id_ >= i: - g_val += s_g_cu[i - 1] - cute.arch.barrier(barrier_id=3, number_of_threads=BT) - - # store g_cu to gmem for H and O kernels - if in_bounds: - g_cu[t, head_id] = g_val - - # store g and g_cu to smem for later - s_g_cu[tid_] = g_val - s_g_cu_exp[tid_] = cute.math.exp(g_val) if in_bounds else 0.0 - - ##### Phase 2: A = strictLower(beta * kkt * Gamma) ##### + ##### Phase 1: A = strictLower(beta * kkt * Gamma) ##### + # Ab/Abg share one tmem slot across stages. The MMA warp commits + # tcgen05 groups in program order: KKT_i, W_i, U_i, KKT_{i+1}. + # Waiting for KKT_i means the previous W/U commits have completed, + # so the INV warps can safely overwrite Ab/Abg for this iteration. + # Wait for prep warps to publish beta/gate and for KKT MMA. if warp_id_ == 0: + cute.arch.mbarrier_wait(prep_mbar + stage_id, parity) cute.arch.mbarrier_wait(mma_kkt_mbar + stage_id, parity) - cute.arch.barrier(barrier_id=1, number_of_threads=128) + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128) _tcgen05.fence_after_thread_sync() # tmem 16x256b layout / ldmatrix layout + beta_row = cute.make_rmem_tensor(2, Float32) + g_cu_row = cute.make_rmem_tensor(2, Float32) + for i in cutlass.range_constexpr(2): + idx = warp_id_ * 16 + i * 8 + (lane_id // 4) + beta_row[i] = s_beta[idx, stage_id] + g_cu_row[i] = s_g_cu[idx, stage_id] + # mode0 is 8 rows together # mode1 is top and bottom 8 rows # mode2 is groups of 16 rows - row_coord = (lane_id // 4, None, warp_id_) - s_beta_view = cute.make_tensor(s_beta, (8, 2, 4)) - beta_row = s_beta_view[row_coord].load().reshape((1, 2, 1)) - - s_g_cu_view = cute.make_tensor(s_g_cu, (8, 2, 4)) - g_cu_row = s_g_cu_view[row_coord].load().reshape((1, 2, 1)) + beta_row = beta_row.load().reshape((1, 2, 1)) + g_cu_row = g_cu_row.load().reshape((1, 2, 1)) # mode0 is 2 consecutive elems # mode1 is top and bottom 8 rows @@ -429,7 +491,9 @@ class Sm100ChunkUWKernel: # mode2 is top and bottom 8 rows # mode3 is next 16 columns col_coord = (None, lane_id % 4, None, i) - s_g_cu_view = cute.make_tensor(s_g_cu, (2, 4, 2, BT // 16)) + s_g_cu_view = cute.make_tensor( + s_g_cu[None, stage_id].iterator, (2, 4, 2, BT // 16) + ) g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2)) Gamma = cute.math.exp(g_cu_row - g_cu_col, fastmath=True) @@ -443,29 +507,23 @@ class Sm100ChunkUWKernel: # pack to BF16 # CuteDSL doesn't generate cvt.bf16x2.f32 here for some reasons packed = cute.make_rmem_tensor(4, Uint32) - packed[0] = cvt.fp32x2_to_bf16x2( - A_masked[0, 0, 0], A_masked[1, 0, 0] - ) - packed[1] = cvt.fp32x2_to_bf16x2( - A_masked[0, 1, 0], A_masked[1, 1, 0] - ) - packed[2] = cvt.fp32x2_to_bf16x2( - A_masked[0, 0, 1], A_masked[1, 0, 1] - ) - packed[3] = cvt.fp32x2_to_bf16x2( - A_masked[0, 1, 1], A_masked[1, 1, 1] - ) + for j in cutlass.range_constexpr(4): + packed[j] = cvt.fp32x2_to_bf16x2( + A_masked[j * 2], A_masked[j * 2 + 1] + ) # store to smem cute.copy( stsm_atom, cute.recast_tensor(packed, BFloat16), - sA_ldsm[warp_id_, None, i], + sA_ldsm[None, (warp_id_, i)], ) - cute.arch.barrier(barrier_id=1, number_of_threads=128) + # use sync warp instead of bar.sync because for block-diagonal inverse, + # each warp reads its own private smem memory. + cute.arch.sync_warp() - ##### Phase 3: matrix inverse ##### + ##### Phase 2: matrix inverse ##### # we use Newton-Schulz iterations to compute the inverse # of the four 16x16 diagonal blocks. # Ai_new = 2 Ai - Ai @ M @ Ai @@ -477,15 +535,6 @@ class Sm100ChunkUWKernel: zeros_f32 = cute.make_rmem_tensor(4, Float32) zeros_f32.fill(0.0) - def set_diagonal(A: cute.Tensor, lane_id: Int32): - "Set the diagonal to 1s" - if lane_id % 9 == 0: - A[0] = (A[0] & Uint32(0xFFFF0000)) | Uint32(0x00003F80) - A[3] = (A[3] & Uint32(0xFFFF0000)) | Uint32(0x00003F80) - elif lane_id % 9 == 4: - A[0] = (A[0] & Uint32(0x0000FFFF)) | Uint32(0x3F800000) - A[3] = (A[3] & Uint32(0x0000FFFF)) | Uint32(0x3F800000) - Ai_bf16 = cute.make_rmem_tensor(8, BFloat16) mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16) M_bf16 = cute.make_rmem_tensor(8, BFloat16) @@ -496,44 +545,60 @@ class Sm100ChunkUWKernel: mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2) M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2) + # construct rmem-backed identity matrix + eye = cute.make_rmem_tensor(4, Uint32) + eye[0] = Uint32(lane_id % 9 == 0) * Uint32(0x00003F80) + Uint32( + lane_id % 9 == 4 + ) * Uint32(0x3F800000) + eye[1] = 0 + eye[2] = 0 + eye[3] = eye[0] + # initial guess: Ai = I-A - cute.copy(ldsm_atom, sA_ldsm[warp_id_, None, warp_id_], Ai_bf16) + cute.copy(ldsm_atom, sA_ldsm[None, (warp_id_, warp_id_)], Ai_bf16) for i in cutlass.range_constexpr(4): - Ai[i] ^= Uint32(0x80008000) # negate A - set_diagonal(Ai, lane_id) + Ai[i] = _bf16x2_sub(eye[i], Ai[i]) # (4, 2) Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4) # M is holding -(I+A), stay constant throughout the iterations - cute.copy(ldsm_trans_atom, sA_ldsm[warp_id_, None, warp_id_], M_bf16) - set_diagonal(M, lane_id) + cute.copy(ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_)], M_bf16) for i in cutlass.range_constexpr(4): - M[i] ^= Uint32(0x80008000) + M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i]) # 3 rounds of Newton-Schulz for _ in cutlass.range_constexpr(3): # First MMA: -AiM = Ai @ (-M) - cute.copy(stsm_atom, Ai_bf16, sA_ldsm[warp_id_, None, warp_id_]) + cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)]) cute.arch.sync_warp() acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32) acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32) Ai_bf16.store(acc.load().to(BFloat16)) # Second MMA: Ai_new = 2Ai + (-AiM) @ Ai - for j in cutlass.range_constexpr(8): + for j in cutlass.range(8, vectorize=True): Ai_f32[j] *= 2.0 cute.copy( ldsm_trans_atom, - sA_ldsm[warp_id_, None, warp_id_], + sA_ldsm[None, (warp_id_, warp_id_)], mma_B_bf16, ) Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0]) Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1]) Ai_bf16.store(Ai_f32.load().to(BFloat16)) - cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[warp_id_, None, warp_id_]) - cute.arch.barrier(barrier_id=1, number_of_threads=128) + cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[None, (warp_id_, warp_id_)]) + store_ab_abg( + Ai_f32, + s_beta[None, stage_id], + s_beta_g[None, stage_id], + warp_id_, + lane_id, + warp_id_, + Ab_tmem_base, + ) + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128) # off-diagonal by 1 # given @@ -544,14 +609,14 @@ class Sm100ChunkUWKernel: # warp1: Ai10 = -Ai11 @ A10 @ Ai00 # warp2: Ai21 = -Ai22 @ A21 @ Ai11 # warp3: Ai32 = -Ai33 @ A32 @ Ai22 - if warp_id_ > 0: + if warp_id_ >= 1: neg_Ai = cute.make_rmem_tensor(4, Uint32) for i in cutlass.range_constexpr(4): - neg_Ai[i] = Ai[i] ^ Uint32(0x80008000) + neg_Ai[i] = _bf16x2_neg(Ai[i]) cute.copy( ldsm_trans_atom, - sA_ldsm[warp_id_, None, warp_id_ - 1], + sA_ldsm[None, (warp_id_, warp_id_ - 1)], mma_B_bf16, ) acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32) @@ -560,44 +625,52 @@ class Sm100ChunkUWKernel: cute.copy( ldsm_trans_atom, - sAi_ldsm[warp_id_ - 1, None, warp_id_ - 1], + sAi_ldsm[None, (warp_id_ - 1, warp_id_ - 1)], mma_B_bf16, ) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) Ai_bf16.store(acc.load().to(BFloat16)) + store_ab_abg( + acc, + s_beta[None, stage_id], + s_beta_g[None, stage_id], + warp_id_, + lane_id, + warp_id_ - 1, + Ab_tmem_base, + ) cute.copy( stsm_atom, Ai_bf16, - sAi_ldsm[warp_id_, None, warp_id_ - 1], + sAi_ldsm[None, (warp_id_, warp_id_ - 1)], ) - cute.arch.barrier(barrier_id=1, number_of_threads=128) + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128) # off-diagonal by 2 - # warp0: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10) - # warp1: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21) - if warp_id_ < 2: + # warp2: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10) + # warp3: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21) + if warp_id_ >= 2: + tile_col = warp_id_ - 2 cute.copy( ldsm_atom, - sA_ldsm[warp_id_ + 2, None, warp_id_], + sA_ldsm[None, (warp_id_, tile_col)], Ai_bf16, ) cute.copy( ldsm_trans_atom, - sAi_ldsm[warp_id_, None, warp_id_], + sAi_ldsm[None, (tile_col, tile_col)], mma_B_bf16, ) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) cute.copy( - ldsm_atom, - sA_ldsm[warp_id_ + 2, None, warp_id_ + 1], - Ai_bf16, + ldsm_atom, sA_ldsm[None, (warp_id_, tile_col + 1)], Ai_bf16 ) cute.copy( ldsm_trans_atom, - sAi_ldsm[warp_id_ + 1, None, warp_id_], + sAi_ldsm[None, (tile_col + 1, tile_col)], mma_B_bf16, ) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0]) @@ -605,79 +678,68 @@ class Sm100ChunkUWKernel: tmp = cute.make_rmem_tensor(8, BFloat16) tmp.store(acc.load().to(BFloat16)) - cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_]) + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)]) cute.arch.sync_warp() - cute.copy( - ldsm_atom, sAi_ldsm[warp_id_ + 2, None, warp_id_ + 2], Ai_bf16 - ) + cute.copy(ldsm_atom, sAi_ldsm[None, (warp_id_, warp_id_)], Ai_bf16) for i in cutlass.range_constexpr(4): - Ai[i] ^= Uint32(0x80008000) + Ai[i] = _bf16x2_neg(Ai[i]) cute.copy( ldsm_trans_atom, - sAi_ldsm[warp_id_ + 2, None, warp_id_], + sAi_ldsm[None, (warp_id_, tile_col)], mma_B_bf16, ) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) tmp.store(acc.load().to(BFloat16)) - cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_]) - cute.arch.barrier(barrier_id=1, number_of_threads=128) + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)]) + store_ab_abg( + acc, + s_beta[None, stage_id], + s_beta_g[None, stage_id], + warp_id_, + lane_id, + tile_col, + Ab_tmem_base, + ) + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128) # off-diagonal by 3 - # warp0: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20) - if warp_id_ == 0: - cute.copy(ldsm_atom, sA_ldsm[3, None, 0], Ai_bf16) - cute.copy(ldsm_trans_atom, sAi_ldsm[0, None, 0], mma_B_bf16) + # warp3: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20) + if warp_id_ == 3: + cute.copy(ldsm_atom, sA_ldsm[None, (3, 0)], Ai_bf16) + cute.copy(ldsm_trans_atom, sAi_ldsm[None, (0, 0)], mma_B_bf16) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) for i in cutlass.range_constexpr(1, 3): - cute.copy(ldsm_atom, sA_ldsm[3, None, i], Ai_bf16) - cute.copy(ldsm_trans_atom, sAi_ldsm[i, None, 0], mma_B_bf16) + cute.copy(ldsm_atom, sA_ldsm[None, (3, i)], Ai_bf16) + cute.copy(ldsm_trans_atom, sAi_ldsm[None, (i, 0)], mma_B_bf16) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0]) acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1]) tmp = cute.make_rmem_tensor(8, BFloat16) tmp.store(acc.load().to(BFloat16)) - cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0]) + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)]) cute.arch.sync_warp() - cute.copy(ldsm_atom, sAi_ldsm[3, None, 3], Ai_bf16) + cute.copy(ldsm_atom, sAi_ldsm[None, (3, 3)], Ai_bf16) for i in cutlass.range_constexpr(4): - Ai[i] ^= Uint32(0x80008000) - cute.copy(ldsm_trans_atom, sAi_ldsm[3, None, 0], mma_B_bf16) + Ai[i] = _bf16x2_neg(Ai[i]) + cute.copy(ldsm_trans_atom, sAi_ldsm[None, (3, 0)], mma_B_bf16) acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) tmp.store(acc.load().to(BFloat16)) - cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0]) - - ##### Phase 4: compute Ab, Abg ##### - if warp_id_ == 3: - cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity ^ 1) - cute.arch.barrier(barrier_id=1, number_of_threads=128) - - for i in cutlass.range_constexpr(BT // 16): - cute.copy(ldsm_atom, sAi_ldsm[warp_id_, None, i], Ai_bf16) - - col_coord = (None, lane_id % 4, None, i) - s_beta_view = cute.make_tensor(s_beta, (2, 4, 2, BT // 16)) - beta_col = s_beta_view[col_coord].load().reshape((2, 1, 2)) - - s_g_cu_view = cute.make_tensor(s_g_cu_exp, (2, 4, 2, BT // 16)) - g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2)) - - Ai_f32 = cvt.bf16x2_to_fp32x2(Ai).load().reshape((2, 2, 2)) - - Ab_f32 = Ai_f32 * beta_col - Ab = Ab_f32.to(BFloat16) - Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id + i * 8 - _tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab) - - Abg_f32 = Ab_f32 * g_cu_col - Abg = Abg_f32.to(BFloat16) - _tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg) - + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)]) + store_ab_abg( + acc, + s_beta[None, stage_id], + s_beta_g[None, stage_id], + warp_id_, + lane_id, + 0, + Ab_tmem_base, + ) _tcgen05.wait_st() _tcgen05.fence_before_thread_sync() cute.arch.mbarrier_arrive(inv_mbar + stage_id) @@ -686,14 +748,18 @@ class Sm100ChunkUWKernel: if stage_id == 0: parity ^= 1 - elif warp_id < 4: + else: # epi warps stage_id = 0 parity = 0 # ((BT, num_global_chunks), V_dim) - gU_tiles = cute.logical_divide(tmaU[None, head_id, None], (BT, None)) - gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None)) + gU_tiles = cute.logical_divide( + U_tma.tma_tensor[None, head_id, None], (BT, None) + ) + gW_tiles = cute.logical_divide( + W_tma.tma_tensor[None, head_id, None], (BT, None) + ) # sW shape: [BT, (64, K_dim/64)] # sW_view shape: [(8, 2), (4, K_dim/64)] @@ -719,7 +785,7 @@ class Sm100ChunkUWKernel: elif warp_id == 1: with cute.arch.elect_one(): cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.barrier(barrier_id=2, number_of_threads=128) + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128) _tcgen05.fence_after_thread_sync() w_f32 = _tcgen05.ld(warp_id * 32 + 16, U_tmem, "16x256b", K_dim // 8) @@ -729,16 +795,16 @@ class Sm100ChunkUWKernel: cute.copy(stsm_atom, w_bf16, sW_view) # wait for U MMA + issue W TMA store - cute.arch.barrier(barrier_id=2, number_of_threads=128) + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128) fence_before_tma_store() if warp_id == 0: cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity) elif warp_id == 1: # don't need to commit simple_tma_copy( - W_tma_atom, sW, gW_tiles[(None, global_chunk_id), None] + W_tma.atom, sW, gW_tiles[(None, global_chunk_id), None] ) - cute.arch.barrier(barrier_id=2, number_of_threads=128) + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128) _tcgen05.fence_after_thread_sync() u_f32 = _tcgen05.ld(warp_id * 32, U_tmem, "16x256b", V_dim // 8) @@ -749,11 +815,11 @@ class Sm100ChunkUWKernel: u_bf16.store(u_f32.to(BFloat16)) cute.copy(stsm_atom, u_bf16, sU_view) - cute.arch.barrier(barrier_id=2, number_of_threads=128) + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128) fence_before_tma_store() if warp_id == 1: simple_tma_copy( - U_tma_atom, sU, gU_tiles[(None, global_chunk_id), None] + U_tma.atom, sU, gU_tiles[(None, global_chunk_id), None] ) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py index f132ebaa6f..4be42f620a 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py @@ -61,13 +61,12 @@ class Sm100ChunkOKernel: stride=(64, 0, (1, self.BT * 64), self.BT * dim), ) slayout = cute.make_composed_layout(swizzle_128B, 0, slayout) - atom, tma_tensor = cpasync.make_tiled_tma_atom( + return cpasync.make_tiled_tma_atom( op, cute.logical_divide(tensor, (None, None, 64)), slayout, cta_tiler=(self.BT, 1, dim), ) - return atom, tma_tensor, slayout @cute.jit def _make_h_tma_args( @@ -79,17 +78,22 @@ class Sm100ChunkOKernel: num_elems = 128 // (tensor.element_type.width // 8) swizzle_128B = cute.make_swizzle(3, 4, 3) slayout = cute.make_layout( - (1, self.V_dim, (num_elems, self.K_dim // num_elems), stages), - stride=(0, num_elems, (1, self.V_dim * num_elems), self.V_dim * self.K_dim), + (1, 1, self.V_dim, (num_elems, self.K_dim // num_elems), stages), + stride=( + 0, + 0, + num_elems, + (1, self.V_dim * num_elems), + self.V_dim * self.K_dim, + ), ) slayout = cute.make_composed_layout(swizzle_128B, 0, slayout) - atom, tma_tensor = cpasync.make_tiled_tma_atom( + return cpasync.make_tiled_tma_atom( op, - cute.logical_divide(tensor, (None, None, num_elems)), + cute.logical_divide(tensor, (None, None, None, num_elems)), slayout, - cta_tiler=(1, self.V_dim, self.K_dim), + cta_tiler=(1, 1, self.V_dim, self.K_dim), ) - return atom, tma_tensor, slayout @cute.jit def __call__( @@ -111,35 +115,35 @@ class Sm100ChunkOKernel: block = (self.num_warps * 32, 1, 1) tma_g2s = cpasync.CopyBulkTensorTileG2SOp() tma_s2g = cpasync.CopyBulkTensorTileS2GOp() - Q_args = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages) - K_args = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages) - V_args = self._make_bf16_tma_args( + Q_tma = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages) + K_tma = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages) + V_tma = self._make_bf16_tma_args( v_new_chunks, self.V_dim, tma_g2s, self.num_stages ) - H_args = self._make_h_tma_args(h, tma_g2s, self.num_stages) - O_args = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1) + H_tma = self._make_h_tma_args(h, tma_g2s, self.num_stages) + O_tma = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1) self.kernel( - Q_args, - K_args, - V_args, - H_args, - O_args, + Q_tma, + K_tma, + V_tma, + H_tma, + O_tma, g_cu, o, cu_seqlens, chunk_indices, total_chunks, scale, - ).launch(grid=grid, block=block, stream=stream) + ).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream) @cute.kernel def kernel( self, - Q_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], - O_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout], + Q_tma: cpasync.TmaInfo, + K_tma: cpasync.TmaInfo, + V_tma: cpasync.TmaInfo, + H_tma: cpasync.TmaInfo, + O_tma: cpasync.TmaInfo, g_cu: cute.Tensor, o: cute.Tensor, cu_seqlens: cute.Tensor, @@ -162,23 +166,17 @@ class Sm100ChunkOKernel: k_head_id = v_head_id // heads_per_qk num_global_chunks = total_chunks[0] - Q_tma_atom, tmaQ, sQ_layout = Q_args - K_tma_atom, tmaK, sK_layout = K_args - V_tma_atom, tmaV, sV_layout = V_args - H_tma_atom, tmaH, sH_layout = H_args - O_tma_atom, tmaO, sO_layout = O_args - def allocate_tensor(smem, dtype, layout): return smem.allocate_tensor( dtype, layout.outer, byte_alignment=128, swizzle=layout.inner ) smem = cutlass.utils.SmemAllocator() - sQ = allocate_tensor(smem, BFloat16, sQ_layout)[None, 0, None, None] - sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None] - sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None] - sH = allocate_tensor(smem, BFloat16, sH_layout)[0, None, None, None] - sO = allocate_tensor(smem, BFloat16, sO_layout)[None, 0, None, 0] + sQ = allocate_tensor(smem, BFloat16, Q_tma.smem_layout)[None, 0, None, None] + sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None] + sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None] + sH = allocate_tensor(smem, BFloat16, H_tma.smem_layout)[0, 0, None, None, None] + sO = allocate_tensor(smem, BFloat16, O_tma.smem_layout)[None, 0, None, 0] s_g_cu = smem.allocate_array(Float32, BT) qk_full_mbar = smem.allocate_array(Int64, num_stages) @@ -207,10 +205,10 @@ class Sm100ChunkOKernel: cute.arch.mbarrier_init(epi_mbar, 128) cute.arch.mbarrier_init_fence() elif warp_id == 9: - cpasync.prefetch_descriptor(Q_tma_atom) - cpasync.prefetch_descriptor(K_tma_atom) - cpasync.prefetch_descriptor(V_tma_atom) - cpasync.prefetch_descriptor(H_tma_atom) + cpasync.prefetch_descriptor(Q_tma.atom) + cpasync.prefetch_descriptor(K_tma.atom) + cpasync.prefetch_descriptor(V_tma.atom) + cpasync.prefetch_descriptor(H_tma.atom) cute.arch.sync_threads() if warp_id == 9: @@ -225,12 +223,16 @@ class Sm100ChunkOKernel: # copy Q and K q_tile = cute.local_tile( - cute.domain_offset((bos, 0), tmaQ[None, k_head_id, None]), + cute.domain_offset( + (bos, 0), Q_tma.tma_tensor[None, k_head_id, None] + ), tiler=(BT, K_dim), coord=(chunk_id, 0), ) k_tile = cute.local_tile( - cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]), + cute.domain_offset( + (bos, 0), K_tma.tma_tensor[None, k_head_id, None] + ), tiler=(BT, K_dim), coord=(chunk_id, 0), ) @@ -241,13 +243,13 @@ class Sm100ChunkOKernel: with cute.arch.elect_one(): STAGE_SIZE = BT * (K_dim + K_dim) * 2 cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE) - simple_tma_copy(Q_tma_atom, q_tile, sQ[None, None, stage_id], mbar) - simple_tma_copy(K_tma_atom, k_tile, sK[None, None, stage_id], mbar) + simple_tma_copy(Q_tma.atom, q_tile, sQ[None, None, stage_id], mbar) + simple_tma_copy(K_tma.atom, k_tile, sK[None, None, stage_id], mbar) # copy H and V - gH = tmaH[global_chunk_id * self.Hv + v_head_id, None, None] + gH = H_tma.tma_tensor[global_chunk_id, v_head_id, None, None] gV = cute.local_tile( - tmaV[None, v_head_id, None], + V_tma.tma_tensor[None, v_head_id, None], tiler=(BT, V_dim), coord=(global_chunk_id, 0), ) @@ -262,10 +264,10 @@ class Sm100ChunkOKernel: mbar, H_STAGE_SIZE + V_STAGE_SIZE ) simple_tma_copy( - H_tma_atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST + H_tma.atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST ) simple_tma_copy( - V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST + V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST ) stage_id = (stage_id + 1) % num_stages @@ -302,40 +304,37 @@ class Sm100ChunkOKernel: cute.arch.mbarrier_wait(qk_full_mbar + stage_id, tma_parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(K_dim // BT): - for j in cutlass.range_constexpr(BT // 16): - qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4) - kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4) - _tcgen05.mma_f16( - qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0) - ) - _tcgen05.commit(qk_mbar) + for i in cutlass.range_constexpr(K_dim // 64): + for j in cutlass.range_constexpr(64 // 16): + qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4) + kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4) + _tcgen05.mma_f16( + qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0) + ) + _tcgen05.commit(qk_mbar) ##### 2nd MMA: Q @ H.T ##### cute.arch.mbarrier_wait(hv_full_mbar + stage_id, tma_parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(K_dim // BT): - for j in cutlass.range_constexpr(BT // 16): - qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4) - hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4) - _tcgen05.mma_f16( - qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0) - ) - _tcgen05.commit(qk_empty_mbar + stage_id) + for i in cutlass.range_constexpr(K_dim // 64): + for j in cutlass.range_constexpr(64 // 16): + qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4) + hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4) + _tcgen05.mma_f16( + qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0) + ) + _tcgen05.commit(qk_empty_mbar + stage_id) ##### 3rd MMA: P @ V ##### # stalled by mask(QK) cute.arch.mbarrier_wait(mask_mbar, mask_parity) _tcgen05.fence_after_thread_sync() - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(BT // 16): - vdesc = vdesc_base | ((i * 16 * 128) >> 4) - _tcgen05.mma_ts_f16( - out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0 - ) - _tcgen05.commit(pv_mma_mbar + stage_id) + for i in cutlass.range_constexpr(BT // 16): + vdesc = vdesc_base | ((i * 16 * 128) >> 4) + _tcgen05.mma_ts_f16( + out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0 + ) + _tcgen05.commit(pv_mma_mbar + stage_id) stage_id = (stage_id + 1) % num_stages if stage_id == 0: @@ -499,11 +498,13 @@ class Sm100ChunkOKernel: fence_before_tma_store() if warp_id == 3: gO = cute.local_tile( - cute.domain_offset((bos, 0), tmaO[None, v_head_id, None]), + cute.domain_offset( + (bos, 0), O_tma.tma_tensor[None, v_head_id, None] + ), tiler=(BT, V_dim), coord=(chunk_id, 0), ) - simple_tma_copy(O_tma_atom, sO, gO) + simple_tma_copy(O_tma.atom, sO, gO) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() @@ -559,13 +560,14 @@ class Sm100ChunkOKernel: total_t = cute.sym_int() pad_t = cute.sym_int() total_chunks_n = cute.sym_int() - h_outer_n = cute.sym_int() cu_entries = cute.sym_int() q = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16) k = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16) v_new = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16) - h_flat = make_fake_tensor(BFloat16, (h_outer_n, V_dim, K_dim), divisibility=16) + h = make_fake_tensor( + BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16 + ) g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4) o = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16) cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1) @@ -586,7 +588,7 @@ class Sm100ChunkOKernel: q, k, v_new, - h_flat, + h, g_cu, o, cu_seqlens, @@ -602,7 +604,7 @@ class Sm100ChunkOKernel: def o_cutedsl( q: torch.Tensor, k: torch.Tensor, - v_new_chunks: torch.Tensor, + v_new: torch.Tensor, h: torch.Tensor, g_cu: torch.Tensor, o: torch.Tensor, @@ -618,8 +620,8 @@ def o_cutedsl( Sm100ChunkOKernel.compile(H, Hv, K_dim, V_dim)( q, k, - v_new_chunks.view(-1, Hv, V_dim), - h.view(-1, V_dim, K_dim), + v_new, + h, g_cu, o, cu_seqlens,