@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 43916891b2

[GDN] Improve kkt kernel of CuteDSL prefill backend (#46346)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Jun 30, 2026 at 09:34 UTC 43916891b222c4ebe85e4c1f4da08202373ede30
6 files changed +554 -474
vllm/cute_utils/__init__.py
+29 -15
@@ -95,40 +95,54 @@ def mma_bf16(
95 return cute.TensorSSA(vec, 4, Float32)
96
97
98 -@dsl_user_op
99 -def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
98 +def _bf16x2_unary(asm: str, a: Uint32, *, loc=None, ip=None) -> Uint32:
99 out = llvm.inline_asm(
100 T.i32(),
101 [a.ir_value(loc=loc, ip=ip)],
103 - "abs.bf16x2 $0, $1;",
102 + f"{asm}.bf16x2 $0, $1;",
103 "=r,r",
104 has_side_effects=False,
105 is_align_stack=False,
106 + loc=loc,
107 + ip=ip,
108 )
109 return Uint32(out)
110
111
111 -@dsl_user_op
112 -def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
112 +def _bf16x2_binary(asm: str, a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
113 out = llvm.inline_asm(
114 T.i32(),
115 [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
116 - "max.bf16x2 $0, $1, $2;",
116 + f"{asm}.bf16x2 $0, $1, $2;",
117 "=r,r,r",
118 has_side_effects=False,
119 is_align_stack=False,
120 + loc=loc,
121 + ip=ip,
122 )
123 return Uint32(out)
124
125
126 +@dsl_user_op
127 +def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
128 + return _bf16x2_unary("abs", a, loc=loc, ip=ip)
129 +
130 +
131 +@dsl_user_op
132 +def _bf16x2_neg(a: Uint32, *, loc=None, ip=None) -> Uint32:
133 + return _bf16x2_unary("neg", a, loc=loc, ip=ip)
134 +
135 +
136 +@dsl_user_op
137 +def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
138 + return _bf16x2_binary("max", a, b, loc=loc, ip=ip)
139 +
140 +
141 @dsl_user_op
142 def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
126 - out = llvm.inline_asm(
127 - T.i32(),
128 - [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
129 - "mul.rn.bf16x2 $0, $1, $2;",
130 - "=r,r,r",
131 - has_side_effects=False,
132 - is_align_stack=False,
133 - )
134 - return Uint32(out)
143 + return _bf16x2_binary("mul.rn", a, b, loc=loc, ip=ip)
144 +
145 +
146 +@dsl_user_op
147 +def _bf16x2_sub(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
148 + return _bf16x2_binary("sub.rn", a, b, loc=loc, ip=ip)
vllm/cute_utils/_tcgen05.py
+34 -30
@@ -90,17 +90,18 @@ def mma_f16(
90 loc=None,
91 ip=None,
92 ) -> None:
93 - nvvm.tcgen05_mma(
94 - nvvm.Tcgen05MMAKind.F16,
95 - NVVM_CTA_GROUP_MAP[cta_group],
96 - _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
97 - Uint64(a_desc).ir_value(loc=loc, ip=ip),
98 - Uint64(b_desc).ir_value(loc=loc, ip=ip),
99 - Int32(idesc).ir_value(loc=loc, ip=ip),
100 - Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
101 - loc=loc,
102 - ip=ip,
103 - )
93 + with cute.arch.elect_one():
94 + nvvm.tcgen05_mma(
95 + nvvm.Tcgen05MMAKind.F16,
96 + NVVM_CTA_GROUP_MAP[cta_group],
97 + _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
98 + Uint64(a_desc).ir_value(loc=loc, ip=ip),
99 + Uint64(b_desc).ir_value(loc=loc, ip=ip),
100 + Int32(idesc).ir_value(loc=loc, ip=ip),
101 + Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
102 + loc=loc,
103 + ip=ip,
104 + )
105
106
107 @dsl_user_op
@@ -115,17 +116,18 @@ def mma_ts_f16(
116 loc=None,
117 ip=None,
118 ) -> None:
118 - nvvm.tcgen05_mma(
119 - nvvm.Tcgen05MMAKind.F16,
120 - NVVM_CTA_GROUP_MAP[cta_group],
121 - _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
122 - _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip),
123 - Uint64(b_desc).ir_value(loc=loc, ip=ip),
124 - Int32(idesc).ir_value(loc=loc, ip=ip),
125 - Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
126 - loc=loc,
127 - ip=ip,
128 - )
119 + with cute.arch.elect_one():
120 + nvvm.tcgen05_mma(
121 + nvvm.Tcgen05MMAKind.F16,
122 + NVVM_CTA_GROUP_MAP[cta_group],
123 + _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
124 + _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip),
125 + Uint64(b_desc).ir_value(loc=loc, ip=ip),
126 + Int32(idesc).ir_value(loc=loc, ip=ip),
127 + Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
128 + loc=loc,
129 + ip=ip,
130 + )
131
132
133 @dsl_user_op
@@ -133,15 +135,17 @@ def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
135 mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
136 group = NVVM_CTA_GROUP_MAP[cta_group]
137 if cutlass.const_expr(cta_mask is not None):
136 - nvvm.tcgen05_commit_arrive(
137 - mbar_llvm,
138 - multicast_mask=cta_mask.ir_value(loc=loc, ip=ip),
139 - group=group,
140 - loc=loc,
141 - ip=ip,
142 - )
138 + with cute.arch.elect_one():
139 + nvvm.tcgen05_commit_arrive(
140 + mbar_llvm,
141 + multicast_mask=cta_mask.ir_value(loc=loc, ip=ip),
142 + group=group,
143 + loc=loc,
144 + ip=ip,
145 + )
146 else:
144 - nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
147 + with cute.arch.elect_one():
148 + nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
149
150
151 @dsl_user_op
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py
+23 -28
@@ -174,31 +174,31 @@ def chunk_gated_delta_rule_cutedsl(
174 When ``core_attn_out`` is provided, ``output`` is an unsqueezed view of
175 that buffer.
176 """
177 - q_3d = q.squeeze(0)
178 - k_3d = k.squeeze(0)
179 - v_3d = v.squeeze(0)
180 - g_2d = g.squeeze(0)
181 - beta_2d = beta.squeeze(0)
182 -
183 - _, _, head_k_dim = k_3d.shape
184 - _, num_v_heads, head_v_dim = v_3d.shape
177 + q = q.squeeze(0)
178 + k = k.squeeze(0)
179 + v = v.squeeze(0)
180 + g = g.squeeze(0)
181 + beta = beta.squeeze(0)
182 +
183 + _, _, K_dim = k.shape
184 + _, num_v_heads, V_dim = v.shape
185 chunk_size = 64
186 upper_bound_chunks = chunk_indices.shape[0]
187 pad_t = upper_bound_chunks * chunk_size
188 total_chunks_ptr = chunk_offsets[-1:]
189
190 - g_cu = torch.empty_like(g_2d, dtype=torch.float32)
191 - u = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
192 - w = q_3d.new_empty(pad_t, num_v_heads, head_k_dim)
190 + g_cu = torch.empty_like(g, dtype=torch.float32)
191 + u = q.new_empty(pad_t, num_v_heads, V_dim)
192 + w = q.new_empty(pad_t, num_v_heads, K_dim)
193
194 num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
195 kkt_inv_uw_cutedsl(
196 - k_3d,
197 - v_3d,
196 + k,
197 + v,
198 u,
199 w,
200 - g_2d,
201 - beta_2d,
200 + g,
201 + beta,
202 g_cu,
203 cu_seqlens,
204 chunk_indices,
@@ -206,16 +206,11 @@ def chunk_gated_delta_rule_cutedsl(
206 num_sms=num_sms,
207 )
208
209 - h = k_3d.new_empty(
210 - upper_bound_chunks,
211 - num_v_heads,
212 - head_v_dim,
213 - head_k_dim,
214 - )
215 - v_new = q_3d.new_empty(pad_t, num_v_heads, head_v_dim)
209 + h = k.new_empty(upper_bound_chunks, num_v_heads, V_dim, K_dim)
210 + v_new = q.new_empty(pad_t, num_v_heads, V_dim)
211 final_state = torch.empty_like(initial_state)
212 h_cutedsl(
218 - k_3d,
213 + k,
214 u,
215 w,
216 v_new,
@@ -227,12 +222,12 @@ def chunk_gated_delta_rule_cutedsl(
222 chunk_offsets,
223 )
224
230 - output = core_attn_out if core_attn_out is not None else torch.empty_like(v_3d)
231 - scale = head_k_dim**-0.5
225 + output = core_attn_out if core_attn_out is not None else torch.empty_like(v)
226 + scale = K_dim**-0.5
227 o_cutedsl(
233 - q_3d,
234 - k_3d,
235 - v_new.view(upper_bound_chunks, chunk_size, num_v_heads, head_v_dim),
228 + q,
229 + k,
230 + v_new,
231 h,
232 g_cu,
233 output,
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py
+89 -90
@@ -66,13 +66,12 @@ class Sm100ChunkHKernel:
66 stride=(64, 0, (1, self.BT * 64), self.BT * dim),
67 )
68 slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
69 - atom, tma_tensor = cpasync.make_tiled_tma_atom(
69 + return cpasync.make_tiled_tma_atom(
70 op,
71 cute.logical_divide(tensor, (None, None, 64)),
72 slayout,
73 cta_tiler=(self.BT, 1, dim),
74 )
75 - return atom, tma_tensor, slayout
75
76 @cute.jit
77 def _make_h_tma_args(self, tensor: cute.Tensor, op: cpasync.TmaCopyOp):
@@ -84,13 +83,12 @@ class Sm100ChunkHKernel:
83 stride=(0, 0, num_elems, (1, self.V_dim * num_elems)),
84 )
85 slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
87 - atom, tma_tensor = cpasync.make_tiled_tma_atom(
86 + return cpasync.make_tiled_tma_atom(
87 op,
88 cute.logical_divide(tensor, (None, None, None, num_elems)),
89 slayout,
90 cta_tiler=(1, 1, self.V_dim, self.K_dim),
91 )
93 - return atom, tma_tensor, slayout
92
93 @cute.jit
94 def __call__(
@@ -110,39 +108,39 @@ class Sm100ChunkHKernel:
108 tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
109 tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
110
113 - K_args = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages)
114 - V_args = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages)
115 - W_args = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages)
116 - V_new_args = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1)
117 - H0_args = self._make_h_tma_args(h0, tma_g2s)
118 - HT_args = self._make_h_tma_args(ht, tma_s2g)
119 - H_args = self._make_h_tma_args(h, tma_s2g)
111 + K_tma = self._make_bf16_tma_args(K, self.K_dim, tma_g2s, self.num_stages)
112 + V_tma = self._make_bf16_tma_args(V, self.V_dim, tma_g2s, self.num_stages)
113 + W_tma = self._make_bf16_tma_args(W, self.K_dim, tma_g2s, self.num_stages)
114 + V_new_tma = self._make_bf16_tma_args(V_new, self.V_dim, tma_s2g, 1)
115 + H0_tma = self._make_h_tma_args(h0, tma_g2s)
116 + HT_tma = self._make_h_tma_args(ht, tma_s2g)
117 + H_tma = self._make_h_tma_args(h, tma_s2g)
118
119 grid = (self.Hv, h0.shape[0], 1)
120 block = (self.num_warps * 32, 1, 1)
121 self.kernel(
124 - K_args,
125 - V_args,
126 - W_args,
127 - V_new_args,
128 - H0_args,
129 - HT_args,
130 - H_args,
122 + K_tma,
123 + V_tma,
124 + W_tma,
125 + V_new_tma,
126 + H0_tma,
127 + HT_tma,
128 + H_tma,
129 g_cu,
130 cu_seqlens,
131 chunk_offsets,
134 - ).launch(grid=grid, block=block, stream=stream)
132 + ).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream)
133
134 @cute.kernel
135 def kernel(
136 self,
139 - K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
140 - V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
141 - W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
142 - V_new_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
143 - H0_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
144 - HT_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
145 - H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
137 + K_tma: cpasync.TmaInfo,
138 + V_tma: cpasync.TmaInfo,
139 + W_tma: cpasync.TmaInfo,
140 + V_new_tma: cpasync.TmaInfo,
141 + H0_tma: cpasync.TmaInfo,
142 + HT_tma: cpasync.TmaInfo,
143 + H_tma: cpasync.TmaInfo,
144 g_cu: cute.Tensor,
145 cu_seqlens: cute.Tensor,
146 chunk_offsets: cute.Tensor,
@@ -158,14 +156,6 @@ class Sm100ChunkHKernel:
156 num_stages = self.num_stages
157 is_f32 = self.h_dtype == Float32
158
161 - K_tma_atom, tmaK, sK_layout = K_args
162 - V_tma_atom, tmaV, sV_layout = V_args
163 - W_tma_atom, tmaW, sW_layout = W_args
164 - V_new_tma_atom, tmaV_new, sV_new_layout = V_new_args
165 - H0_tma_atom, tmaH0, sH0_layout = H0_args
166 - HT_tma_atom, tmaHT, _ = HT_args
167 - H_tma_atom, tmaH, sH_layout = H_args
168 -
159 def allocate_tensor(smem, dtype, layout):
160 return smem.allocate_tensor(
161 dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
@@ -174,12 +164,14 @@ class Sm100ChunkHKernel:
164 smem = cutlass.utils.SmemAllocator()
165
166 # remove size=1 modes
177 - sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, None]
178 - sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
179 - sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
180 - sH0 = allocate_tensor(smem, self.h_dtype, sH0_layout)[0, 0, None, None]
181 - sH = allocate_tensor(smem, BFloat16, sH_layout)[0, 0, None, None]
182 - sV_new = allocate_tensor(smem, BFloat16, sV_new_layout)[None, 0, None, 0]
167 + sW = allocate_tensor(smem, BFloat16, W_tma.smem_layout)[None, 0, None, None]
168 + sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None]
169 + sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None]
170 + sH0 = allocate_tensor(smem, self.h_dtype, H0_tma.smem_layout)[0, 0, None, None]
171 + sH = allocate_tensor(smem, BFloat16, H_tma.smem_layout)[0, 0, None, None]
172 + sV_new = allocate_tensor(smem, BFloat16, V_new_tma.smem_layout)[
173 + None, 0, None, 0
174 + ]
175
176 s_v_scale = smem.allocate_array(Float32, BT)
177 tma_mbar = smem.allocate_array(Int64, num_stages)
@@ -206,13 +198,13 @@ class Sm100ChunkHKernel:
198 cute.arch.mbarrier_init(h0_mbar, 1)
199 cute.arch.mbarrier_init_fence()
200 elif warp_id == 1:
209 - cpasync.prefetch_descriptor(H0_tma_atom)
210 - cpasync.prefetch_descriptor(W_tma_atom)
211 - cpasync.prefetch_descriptor(V_tma_atom)
212 - cpasync.prefetch_descriptor(K_tma_atom)
213 - cpasync.prefetch_descriptor(HT_tma_atom)
214 - cpasync.prefetch_descriptor(H_tma_atom)
215 - cpasync.prefetch_descriptor(V_new_tma_atom)
201 + cpasync.prefetch_descriptor(H0_tma.atom)
202 + cpasync.prefetch_descriptor(W_tma.atom)
203 + cpasync.prefetch_descriptor(V_tma.atom)
204 + cpasync.prefetch_descriptor(K_tma.atom)
205 + cpasync.prefetch_descriptor(HT_tma.atom)
206 + cpasync.prefetch_descriptor(H_tma.atom)
207 + cpasync.prefetch_descriptor(V_new_tma.atom)
208 cute.arch.sync_threads()
209
210 bos = cu_seqlens[seq_id]
@@ -233,14 +225,21 @@ class Sm100ChunkHKernel:
225 H0_size = V_dim * K_dim * self.h_dtype.width // 8
226 cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size)
227 simple_tma_copy(
236 - H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar
228 + H0_tma.atom,
229 + H0_tma.tma_tensor[seq_id, head_id, None, None],
230 + sH0,
231 + h0_mbar,
232 )
233
234 # shape: ((BT, num_BT_tiles), (64, 2))
240 - gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
241 - gV_tiles = cute.logical_divide(tmaV[None, head_id, None], (BT, None))
235 + gW_tiles = cute.logical_divide(
236 + W_tma.tma_tensor[None, head_id, None], (BT, None)
237 + )
238 + gV_tiles = cute.logical_divide(
239 + V_tma.tma_tensor[None, head_id, None], (BT, None)
240 + )
241 gK_tiles = cute.logical_divide(
243 - cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
242 + cute.domain_offset((bos, 0), K_tma.tma_tensor[None, k_head_id, None]),
243 (BT, None),
244 )
245
@@ -258,12 +257,12 @@ class Sm100ChunkHKernel:
257 STAGE_SIZE = BT * (K_dim + V_dim + K_dim) * 2
258 cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
259 simple_tma_copy(
261 - W_tma_atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST
260 + W_tma.atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST
261 )
262 simple_tma_copy(
264 - V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
263 + V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
264 )
266 - simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
265 + simple_tma_copy(K_tma.atom, gK, sK[None, None, stage_id], mbar)
266
267 stage_id = (stage_id + 1) % num_stages
268 if stage_id == 0:
@@ -295,13 +294,12 @@ class Sm100ChunkHKernel:
294 cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
295 _tcgen05.fence_after_thread_sync()
296
298 - with cute.arch.elect_one():
299 - for i in cutlass.range_constexpr(K_dim // 64):
300 - for j in cutlass.range_constexpr(64 // 16):
301 - hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4)
302 - wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4)
303 - _tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True)
304 - _tcgen05.commit(wh_done_mbar + stage_id)
297 + for i in cutlass.range_constexpr(K_dim // 64):
298 + for j in cutlass.range_constexpr(64 // 16):
299 + hdesc0 = hdesc0_base | ((i * V_dim * 128 + j * 32) >> 4)
300 + wdesc0 = wdesc0_base | ((i * BT * 128 + j * 32) >> 4)
301 + _tcgen05.mma_f16(wh_tmem, hdesc0, wdesc0, wh_idesc, True)
302 + _tcgen05.commit(wh_done_mbar + stage_id)
303
304 ##### 2nd MMA: H_new = H + V_new.T @ K #####
305 Kaddr0 = sK[None, None, stage_id].iterator.toint()
@@ -310,12 +308,11 @@ class Sm100ChunkHKernel:
308 cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
309 _tcgen05.fence_after_thread_sync()
310
313 - with cute.arch.elect_one():
314 - for k in cutlass.range_constexpr(BT // 16):
315 - vtmem0 = v_tmem_base + k * 8
316 - kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4)
317 - _tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True)
318 - _tcgen05.commit(vk_done_mbar + stage_id)
311 + for k in cutlass.range_constexpr(BT // 16):
312 + vtmem0 = v_tmem_base + k * 8
313 + kdesc0 = kdesc0_base | ((k * 16 * 128) >> 4)
314 + _tcgen05.mma_ts_f16(vk_tmem, vtmem0, kdesc0, vk_idesc, True)
315 + _tcgen05.commit(vk_done_mbar + stage_id)
316
317 stage_id = (stage_id + 1) % num_stages
318 if stage_id == 0:
@@ -331,13 +328,12 @@ class Sm100ChunkHKernel:
328 cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
329 _tcgen05.fence_after_thread_sync()
330
334 - with cute.arch.elect_one():
335 - for i in cutlass.range_constexpr(K_dim // 64):
336 - for j in cutlass.range_constexpr(64 // 16):
337 - htmem = h_tmem_base + i * 32 + j * 8
338 - wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4)
339 - _tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True)
340 - _tcgen05.commit(wh_done_mbar + stage_id)
331 + for i in cutlass.range_constexpr(K_dim // 64):
332 + for j in cutlass.range_constexpr(64 // 16):
333 + htmem = h_tmem_base + i * 32 + j * 8
334 + wdesc = wdesc_base | ((i * BT * 128 + j * 32) >> 4)
335 + _tcgen05.mma_ts_f16(wh_tmem, htmem, wdesc, wh_idesc, True)
336 + _tcgen05.commit(wh_done_mbar + stage_id)
337
338 ##### 2nd MMA: H_new = H + V_new.T @ K #####
339 Kaddr = sK[None, None, stage_id].iterator.toint()
@@ -346,12 +342,11 @@ class Sm100ChunkHKernel:
342 cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
343 _tcgen05.fence_after_thread_sync()
344
349 - with cute.arch.elect_one():
350 - for k in cutlass.range_constexpr(BT // 16):
351 - vtmem = v_tmem_base + k * 8
352 - kdesc = kdesc_base | ((k * 16 * 128) >> 4)
353 - _tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True)
354 - _tcgen05.commit(vk_done_mbar + stage_id)
345 + for k in cutlass.range_constexpr(BT // 16):
346 + vtmem = v_tmem_base + k * 8
347 + kdesc = kdesc_base | ((k * 16 * 128) >> 4)
348 + _tcgen05.mma_ts_f16(vk_tmem, vtmem, kdesc, vk_idesc, True)
349 + _tcgen05.commit(vk_done_mbar + stage_id)
350
351 stage_id = (stage_id + 1) % num_stages
352 if stage_id == 0:
@@ -421,7 +416,7 @@ class Sm100ChunkHKernel:
416 ).load()
417 )
418
424 - for j in cutlass.range_constexpr(32):
419 + for j in cutlass.range(32, vectorize=True):
420 h_f32[j] *= h_scale
421 _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
422
@@ -435,8 +430,10 @@ class Sm100ChunkHKernel:
430 fence_before_tma_store()
431 if warp_id_ == 3:
432 h_src = sH if cutlass.const_expr(is_f32) else sH0
438 - h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None]
439 - simple_tma_copy(H_tma_atom, h_src, h_dst)
433 + h_dst = H_tma.tma_tensor[
434 + chunk_offset + chunk_id, head_id, None, None
435 + ]
436 + simple_tma_copy(H_tma.atom, h_src, h_dst)
437 with cute.arch.elect_one():
438 cute.arch.cp_async_bulk_commit_group()
439
@@ -490,7 +487,7 @@ class Sm100ChunkHKernel:
487 h_f32.store(
488 _tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32)
489 )
493 - for j in cutlass.range_constexpr(32):
490 + for j in cutlass.range(32, vectorize=True):
491 h_f32[j] *= h_scale
492 _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
493 _tcgen05.wait_st()
@@ -501,8 +498,10 @@ class Sm100ChunkHKernel:
498 cute.arch.barrier(barrier_id=1, number_of_threads=128)
499 fence_before_tma_store()
500 if warp_id_ == 3:
504 - h_dst = tmaH[chunk_offset + chunk_id, head_id, None, None]
505 - simple_tma_copy(H_tma_atom, sH, h_dst)
501 + h_dst = H_tma.tma_tensor[
502 + chunk_offset + chunk_id, head_id, None, None
503 + ]
504 + simple_tma_copy(H_tma.atom, sH, h_dst)
505 with cute.arch.elect_one():
506 cute.arch.cp_async_bulk_commit_group()
507
@@ -530,8 +529,8 @@ class Sm100ChunkHKernel:
529 cute.arch.barrier(barrier_id=1, number_of_threads=128)
530
531 if warp_id_ == 0:
533 - ht_dst = tmaHT[seq_id, head_id, None, None]
534 - simple_tma_copy(HT_tma_atom, sH0, ht_dst)
532 + ht_dst = HT_tma.tma_tensor[seq_id, head_id, None, None]
533 + simple_tma_copy(HT_tma.atom, sH0, ht_dst)
534 with cute.arch.elect_one():
535 cute.arch.cp_async_bulk_commit_group()
536 if warp_id_ == 1:
@@ -551,7 +550,7 @@ class Sm100ChunkHKernel:
550
551 # ((BT, num_BT_tiles), V_dim)
552 gV_new_tiles = cute.logical_divide(
554 - tmaV_new[None, head_id, None], (BT, None)
553 + V_new_tma.tma_tensor[None, head_id, None], (BT, None)
554 )
555
556 # sV shape: [BT, (64, V_dim/64), num_stages]
@@ -653,7 +652,7 @@ class Sm100ChunkHKernel:
652 fence_before_tma_store()
653 if warp_id == 3:
654 gV = gV_new_tiles[(None, chunk_offset + chunk_id), None]
656 - simple_tma_copy(V_new_tma_atom, sV_new, gV)
655 + simple_tma_copy(V_new_tma.atom, sV_new, gV)
656 with cute.arch.elect_one():
657 cute.arch.cp_async_bulk_commit_group()
658
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
+297 -231
@@ -11,6 +11,8 @@ from quack.compile_utils import make_fake_tensor
11
12 from vllm.cute_utils import (
13 EVICT_FIRST,
14 + _bf16x2_neg,
15 + _bf16x2_sub,
16 _tcgen05,
17 cvt,
18 fence_before_tma_store,
@@ -47,7 +49,7 @@ class Sm100ChunkUWKernel:
49
50 # hard-code
51 self.BT = 64
50 - self.num_warps = 2 + 4 + 4
52 + self.num_warps = 4 + 4 + 4
53
54 @cute.jit
55 def _make_tma_args(
@@ -68,13 +70,12 @@ class Sm100ChunkUWKernel:
70
71 # we need to convert gmem layout to (T, H, (64, D/64)) for make_tiled_tma_atom()
72 # to emit a single 4D TMA. otherwise, it will emit (D/64)x 3D TMA.
71 - atom, tma_tensor = cpasync.make_tiled_tma_atom(
73 + return cpasync.make_tiled_tma_atom(
74 op,
75 cute.logical_divide(tensor, (None, None, 64)),
76 slayout,
77 cta_tiler=(self.BT, 1, dim),
78 )
77 - return atom, tma_tensor, slayout
79
80 @cute.jit
81 def __call__(
@@ -95,33 +96,33 @@ class Sm100ChunkUWKernel:
96 tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
97 tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
98
98 - K_args = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s)
99 - V_args = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s)
100 - U_args = self._make_tma_args(U, self.V_dim, 1, tma_s2g)
101 - W_args = self._make_tma_args(W, self.K_dim, 1, tma_s2g)
99 + K_tma = self._make_tma_args(K, self.K_dim, self.num_stages, tma_g2s)
100 + V_tma = self._make_tma_args(V, self.V_dim, self.num_stages, tma_g2s)
101 + U_tma = self._make_tma_args(U, self.V_dim, 1, tma_s2g)
102 + W_tma = self._make_tma_args(W, self.K_dim, 1, tma_s2g)
103
104 grid = (num_sms // self.Hv, self.Hv, 1)
105 block = (self.num_warps * 32, 1, 1)
106 self.kernel(
106 - K_args,
107 - V_args,
108 - U_args,
109 - W_args,
107 + K_tma,
108 + V_tma,
109 + U_tma,
110 + W_tma,
111 g,
112 beta,
113 g_cu,
114 cu_seqlens,
115 chunk_indices,
116 total_chunks,
116 - ).launch(grid=grid, block=block, stream=stream)
117 + ).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream)
118
119 @cute.kernel
120 def kernel(
121 self,
121 - K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
122 - V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
123 - U_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
124 - W_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
122 + K_tma: cpasync.TmaInfo,
123 + V_tma: cpasync.TmaInfo,
124 + U_tma: cpasync.TmaInfo,
125 + W_tma: cpasync.TmaInfo,
126 g: cute.Tensor,
127 beta: cute.Tensor,
128 g_cu: cute.Tensor,
@@ -142,10 +143,10 @@ class Sm100ChunkUWKernel:
143 V_dim = self.V_dim
144 num_stages = self.num_stages
145
145 - K_tma_atom, tmaK, sK_layout = K_args
146 - V_tma_atom, tmaV, sV_layout = V_args
147 - U_tma_atom, tmaU, sU_layout = U_args
148 - W_tma_atom, tmaW, sW_layout = W_args
146 + INV_BAR = 1
147 + EPI_BAR = 2
148 + SCAN_BAR = 3
149 + TMEM_ALLOC_BAR = 4
150
151 def allocate_tensor(smem, dtype, layout):
152 return smem.allocate_tensor(
@@ -153,22 +154,20 @@ class Sm100ChunkUWKernel:
154 )
155
156 smem = cutlass.utils.SmemAllocator()
156 - sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
157 - sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
158 - sU = allocate_tensor(smem, BFloat16, sU_layout)[None, 0, None, 0]
159 - sW = allocate_tensor(smem, BFloat16, sW_layout)[None, 0, None, 0]
157 + sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None]
158 + sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None]
159 + sU = allocate_tensor(smem, BFloat16, U_tma.smem_layout)[None, 0, None, 0]
160 + sW = allocate_tensor(smem, BFloat16, W_tma.smem_layout)[None, 0, None, 0]
161
161 - swizzle_128B = cute.make_swizzle(3, 4, 3)
162 - sA_layout = cute.make_layout((BT, (64, 1)), stride=(64, (1, BT * 64)))
163 - sA_layout = cute.make_composed_layout(swizzle_128B, 0, sA_layout)
164 - sA = allocate_tensor(smem, BFloat16, sA_layout)
165 - sAi = allocate_tensor(smem, BFloat16, sA_layout)
162 + sA_ptr = smem.allocate_array(BFloat16, BT * BT, byte_alignment=16)
163 + sAi_ptr = smem.allocate_array(BFloat16, BT * BT, byte_alignment=16)
164
167 - s_beta = smem.allocate_array(Float32, BT)
168 - s_g_cu_exp = smem.allocate_array(Float32, BT)
169 - s_g_cu = smem.allocate_array(Float32, BT)
165 + s_beta = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages)))
166 + s_g_cu = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages)))
167 + s_beta_g = smem.allocate_tensor(Float32, cute.make_layout((BT, num_stages)))
168
169 tma_mbar = smem.allocate_array(Int64, num_stages)
170 + prep_mbar = smem.allocate_array(Int64, num_stages)
171 mma_kkt_mbar = smem.allocate_array(Int64, num_stages)
172 inv_mbar = smem.allocate_array(Int64, num_stages)
173 mma_u_mbar = smem.allocate_array(Int64, num_stages)
@@ -179,7 +178,7 @@ class Sm100ChunkUWKernel:
178 kkt_tmem = 0
179 U_tmem_base = kkt_tmem + BT
180 Ab_tmem_base = U_tmem_base + V_dim * num_stages
182 - assert Ab_tmem_base + (BT // 2) * num_stages <= 512
181 + assert Ab_tmem_base + (BT // 2) <= 512
182
183 # prepare ldmatrix/stmatrix ops
184 ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4)
@@ -193,6 +192,7 @@ class Sm100ChunkUWKernel:
192 with cute.arch.elect_one():
193 for i in cutlass.range_constexpr(num_stages):
194 cute.arch.mbarrier_init(tma_mbar + i, 1)
195 + cute.arch.mbarrier_init(prep_mbar + i, 64)
196 cute.arch.mbarrier_init(mma_kkt_mbar + i, 1)
197 cute.arch.mbarrier_init(inv_mbar + i, 128)
198 cute.arch.mbarrier_init(mma_u_mbar + i, 1)
@@ -200,14 +200,14 @@ class Sm100ChunkUWKernel:
200 cute.arch.mbarrier_init(epi_mbar + i, 128)
201 cute.arch.mbarrier_init_fence()
202 elif warp_id == 1:
203 - cpasync.prefetch_descriptor(K_tma_atom)
204 - cpasync.prefetch_descriptor(V_tma_atom)
205 - cpasync.prefetch_descriptor(U_tma_atom)
206 - cpasync.prefetch_descriptor(W_tma_atom)
203 + cpasync.prefetch_descriptor(K_tma.atom)
204 + cpasync.prefetch_descriptor(V_tma.atom)
205 + cpasync.prefetch_descriptor(U_tma.atom)
206 + cpasync.prefetch_descriptor(W_tma.atom)
207 cute.arch.sync_threads()
208
209 num_global_chunks = total_chunks[0]
210 - if warp_id == 9:
210 + if warp_id == 11:
211 # TMA warp
212 stage_id = 0
213 parity = 1
@@ -221,12 +221,14 @@ class Sm100ChunkUWKernel:
221 # domain_offset() to shift the pointer first.
222 mbar = tma_mbar + stage_id
223 gK = cute.local_tile(
224 - cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
224 + cute.domain_offset(
225 + (bos, 0), K_tma.tma_tensor[None, k_head_id, None]
226 + ),
227 tiler=(BT, K_dim),
228 coord=(chunk_id, 0),
229 )
230 gV = cute.local_tile(
229 - cute.domain_offset((bos, 0), tmaV[None, head_id, None]),
231 + cute.domain_offset((bos, 0), V_tma.tma_tensor[None, head_id, None]),
232 tiler=(BT, V_dim),
233 coord=(chunk_id, 0),
234 )
@@ -237,18 +239,19 @@ class Sm100ChunkUWKernel:
239 with cute.arch.elect_one():
240 STAGE_SIZE = BT * (K_dim + V_dim) * 2
241 cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
240 - simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
242 + simple_tma_copy(K_tma.atom, gK, sK[None, None, stage_id], mbar)
243 simple_tma_copy(
242 - V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
244 + V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
245 )
246
247 stage_id = (stage_id + 1) % num_stages
248 if stage_id == 0:
249 parity ^= 1
250
249 - elif warp_id == 8:
251 + elif warp_id == 10:
252 # MMA warp
253 _tcgen05.alloc(taddr)
254 + cute.arch.barrier(barrier_id=TMEM_ALLOC_BAR, number_of_threads=160)
255
256 stage_id = 0
257 parity = 0
@@ -263,8 +266,8 @@ class Sm100ChunkUWKernel:
266 for global_chunk_id in range(bid, num_global_chunks, grid_x):
267 U_tmem = U_tmem_base + V_dim * stage_id
268 W_tmem = U_tmem | (16 << 16)
266 - Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id
267 - Abg_tmem = Ab_tmem | (16 << 16)
269 + Ab_tmem = Ab_tmem_base
270 + Abg_tmem = Ab_tmem_base | (16 << 16)
271
272 ##### KKT MMA: KKT = K @ K.T #####
273 kaddr = sK[None, None, stage_id].iterator.toint()
@@ -276,18 +279,12 @@ class Sm100ChunkUWKernel:
279 cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
280 _tcgen05.fence_after_thread_sync()
281
279 - with cute.arch.elect_one():
280 - for i in cutlass.range_constexpr(K_dim // 64):
281 - for j in cutlass.range_constexpr(64 // 16):
282 - kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
283 - _tcgen05.mma_f16(
284 - kkt_tmem,
285 - kdesc,
286 - kdesc,
287 - kkt_idesc,
288 - (i > 0) or (j > 0),
289 - )
290 - _tcgen05.commit(mma_kkt_mbar + stage_id)
282 + for i in cutlass.range_constexpr(K_dim // 64):
283 + for j in cutlass.range_constexpr(64 // 16):
284 + kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
285 + enable_d = (i > 0) or (j > 0)
286 + _tcgen05.mma_f16(kkt_tmem, kdesc, kdesc, kkt_idesc, enable_d)
287 + _tcgen05.commit(mma_kkt_mbar + stage_id)
288
289 ##### U/W MMA: U = Ab @ V, W = Abg @ K #####
290 vaddr = sV[None, None, stage_id].iterator.toint()
@@ -299,20 +296,15 @@ class Sm100ChunkUWKernel:
296 cute.arch.mbarrier_wait(inv_mbar + stage_id, parity)
297 _tcgen05.fence_after_thread_sync()
298
302 - with cute.arch.elect_one():
303 - for i in cutlass.range_constexpr(BT // 16):
304 - _tcgen05.mma_ts_f16(
305 - W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0
306 - )
307 - kdesc += (16 * 128) >> 4
308 - _tcgen05.commit(mma_w_mbar + stage_id)
299 + for i in cutlass.range_constexpr(BT // 16):
300 + _tcgen05.mma_ts_f16(W_tmem, Abg_tmem + i * 8, kdesc, w_idesc, i > 0)
301 + kdesc += (16 * 128) >> 4
302 + _tcgen05.commit(mma_w_mbar + stage_id)
303
310 - for i in cutlass.range_constexpr(BT // 16):
311 - _tcgen05.mma_ts_f16(
312 - U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0
313 - )
314 - vdesc += (16 * 128) >> 4
315 - _tcgen05.commit(mma_u_mbar + stage_id)
304 + for i in cutlass.range_constexpr(BT // 16):
305 + _tcgen05.mma_ts_f16(U_tmem, Ab_tmem + i * 8, vdesc, u_idesc, i > 0)
306 + vdesc += (16 * 128) >> 4
307 + _tcgen05.commit(mma_u_mbar + stage_id)
308
309 stage_id = (stage_id + 1) % num_stages
310 if stage_id == 0:
@@ -321,24 +313,124 @@ class Sm100ChunkUWKernel:
313 cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1)
314 _tcgen05.dealloc()
315
316 + elif warp_id >= 8:
317 + # dedicated prep warps for beta and gate, consumed by INV warps
318 + stage_id = 0
319 + parity = 0
320 + tid_ = tid % 128
321 + warp_id_ = warp_id % 4
322 +
323 + for global_chunk_id in range(bid, num_global_chunks, grid_x):
324 + seq_id = chunk_indices[global_chunk_id, 0]
325 + chunk_id = chunk_indices[global_chunk_id, 1]
326 + bos = cu_seqlens[seq_id]
327 + eos = cu_seqlens[seq_id + 1]
328 + off_t = bos + chunk_id * BT
329 + t = off_t + tid_
330 +
331 + in_bounds = t < eos
332 + beta_val = beta[t, head_id] if in_bounds else Float32(0.0)
333 + g_val = g[t, head_id] if in_bounds else Float32(0.0)
334 +
335 + # warp-local prefix scan
336 + for i in cutlass.range_constexpr(5):
337 + offset = cutlass.const_expr(1 << i)
338 + lower = cute.arch.shuffle_sync_up(g_val, offset, mask_and_clamp=0)
339 + if lane_id >= offset:
340 + g_val += lower
341 +
342 + # Delay the stage-reuse wait until just before touching smem:
343 + # global loads and the warp-local scan do not use staged buffers.
344 + if warp_id_ == 0:
345 + cute.arch.mbarrier_wait(inv_mbar + stage_id, parity ^ 1)
346 + cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT)
347 +
348 + # Store beta and the per-warp scan totals for the cross-warp fixup.
349 + s_beta[tid_, stage_id] = beta_val
350 + if lane_id == 31:
351 + s_g_cu[warp_id_, stage_id] = g_val
352 + cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT)
353 +
354 + # Add the sum from the lower prep warp.
355 + if warp_id_ == 1:
356 + g_val += s_g_cu[0, stage_id]
357 + cute.arch.barrier(barrier_id=SCAN_BAR, number_of_threads=BT)
358 +
359 + if in_bounds:
360 + g_cu[t, head_id] = g_val
361 +
362 + s_g_cu[tid_, stage_id] = g_val
363 + s_beta_g[tid_, stage_id] = beta_val * cute.math.exp(g_val)
364 + cute.arch.mbarrier_arrive(prep_mbar + stage_id)
365 +
366 + stage_id = (stage_id + 1) % num_stages
367 + if stage_id == 0:
368 + parity ^= 1
369 +
370 elif warp_id >= 4:
371 # inv warps
372 tid_ = tid % 128
373 warp_id_ = warp_id % 4
374
375 + def store_ab_abg(
376 + Ai_f32, s_beta, s_beta_g, warp_id_, lane_id, tile_col, Ab_tmem_base
377 + ):
378 + # compute Ab and Abg from Ai, then store to tmem
379 + beta_col = cute.make_rmem_tensor((2, 2), Float32)
380 + beta_g_col = cute.make_rmem_tensor((2, 2), Float32)
381 +
382 + for i in cutlass.range_constexpr(2):
383 + base = tile_col * 16 + i * 8 + (lane_id % 4) * 2
384 + for j in cutlass.range_constexpr(2):
385 + beta_col[j, i] = s_beta[base + j]
386 + beta_g_col[j, i] = s_beta_g[base + j]
387 +
388 + # without conversion to TensorSSA, cutlass.range(vectorize=True) fails
389 + beta_col = beta_col.load()
390 + beta_g_col = beta_g_col.load()
391 +
392 + Ab_f32 = cute.make_rmem_tensor(8, Float32)
393 + Abg_f32 = cute.make_rmem_tensor(8, Float32)
394 + for i in cutlass.range(8, vectorize=True):
395 + scale_idx = (i // 4) * 2 + (i % 2)
396 + Ab_f32[i] = Ai_f32[i] * beta_col[scale_idx]
397 + Abg_f32[i] = Ai_f32[i] * beta_g_col[scale_idx]
398 +
399 + Ab = Ab_f32.load().to(BFloat16)
400 + Abg = Abg_f32.load().to(BFloat16)
401 + Ab_tmem = Ab_tmem_base + tile_col * 8
402 + _tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab)
403 + _tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg)
404 +
405 + # clear the Ab/Abg tmem buffer once before mainloop
406 + # this is to keep the upper triangular tiles zeros
407 + cute.arch.barrier(barrier_id=TMEM_ALLOC_BAR, number_of_threads=160)
408 + zeros_ = cute.make_rmem_tensor(BT // 2, Float32)
409 + zeros_.fill(0.0)
410 + _tcgen05.st(warp_id_ * 32, Ab_tmem_base, "32x32b", BT // 2, zeros_)
411 +
412 stage_id = 0
413 parity = 0
414
332 - # view into (16,16) sub-tiles, then ldmatrix layout
333 - sA_ldsm = cute.logical_divide(sA, (16, cute.make_layout((8, 2))))
334 - sAi_ldsm = cute.logical_divide(sAi, (16, cute.make_layout((8, 2))))
335 - sA_ldsm = sA_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)]
336 - sAi_ldsm = sAi_ldsm[(lane_id % 16, None), ((None, lane_id // 16), None)]
415 + # for sA, we can avoid bank conflict without using swizzling because
416 + # this is only used as tmp buffer between rmem<->smem, no gmem interactions.
417 + # to do so, we can put (8,8) tile contiguous in memory. logically, we are
418 + # partitioning (64,64) tile into 4x (16,16) tiles. to make indexing easier
419 + # later, we view (16,16) tile as (32,8) tile here.
420 + # (4,4) is (row_tile, col_tile).
421 + sA_layout = cute.make_layout(((8, 32), (4, 4)))
422 + sA = cute.make_tensor(sA_ptr, sA_layout)
423 + sAi = cute.make_tensor(sAi_ptr, sA_layout)
424 +
425 + # pre-compute ldmatrix addresses
426 + sA_ldsm = sA[(None, lane_id), None]
427 + sAi_ldsm = sAi[(None, lane_id), None]
428
429 # init Ai smem buffer with zeros (only the first 48 rows)
339 - for i in cutlass.range_constexpr((BT // 4 * 3) * BT // 128):
340 - idx = i * 128 + tid_
341 - sAi[idx // BT, idx % BT] = BFloat16(0.0)
430 + zeros_bf16 = cute.make_rmem_tensor(8, BFloat16)
431 + zeros_bf16.fill(0.0)
432 + for i in cutlass.range_constexpr(3):
433 + cute.copy(stsm_atom, zeros_bf16, sAi_ldsm[None, (i, warp_id_)])
434
435 # indices for ldmatrix layout later
436 row_indices = cute.make_rmem_tensor((1, 2, 1), Int32)
@@ -360,61 +452,31 @@ class Sm100ChunkUWKernel:
452 eos = cu_seqlens[seq_id + 1]
453 off_t = bos + chunk_id * BT
454
363 - t = off_t + tid_
364 -
365 - ##### Phase 1: load g and beta #####
366 - if tid_ < BT:
367 - in_bounds = t < eos
368 - beta_val = beta[t, head_id] if in_bounds else Float32(0.0)
369 - g_val = g[t, head_id] if in_bounds else Float32(0.0)
370 -
371 - s_beta[tid_] = beta_val
372 -
373 - # compute cumsum(g)
374 - # parallel scan within a warp
375 - for i in cutlass.range_constexpr(5):
376 - offset = cutlass.const_expr(1 << i)
377 - lower = cute.arch.shuffle_sync_up(
378 - g_val, offset, mask_and_clamp=0
379 - )
380 - if lane_id >= offset:
381 - g_val += lower
382 -
383 - # store warp sum
384 - if lane_id == 31:
385 - s_g_cu[warp_id_] = g_val
386 - cute.arch.barrier(barrier_id=3, number_of_threads=BT)
387 -
388 - # add warp sum from lower warps
389 - for i in cutlass.range_constexpr(1, BT // 32):
390 - if warp_id_ >= i:
391 - g_val += s_g_cu[i - 1]
392 - cute.arch.barrier(barrier_id=3, number_of_threads=BT)
393 -
394 - # store g_cu to gmem for H and O kernels
395 - if in_bounds:
396 - g_cu[t, head_id] = g_val
397 -
398 - # store g and g_cu to smem for later
399 - s_g_cu[tid_] = g_val
400 - s_g_cu_exp[tid_] = cute.math.exp(g_val) if in_bounds else 0.0
401 -
402 - ##### Phase 2: A = strictLower(beta * kkt * Gamma) #####
455 + ##### Phase 1: A = strictLower(beta * kkt * Gamma) #####
456 + # Ab/Abg share one tmem slot across stages. The MMA warp commits
457 + # tcgen05 groups in program order: KKT_i, W_i, U_i, KKT_{i+1}.
458 + # Waiting for KKT_i means the previous W/U commits have completed,
459 + # so the INV warps can safely overwrite Ab/Abg for this iteration.
460 + # Wait for prep warps to publish beta/gate and for KKT MMA.
461 if warp_id_ == 0:
462 + cute.arch.mbarrier_wait(prep_mbar + stage_id, parity)
463 cute.arch.mbarrier_wait(mma_kkt_mbar + stage_id, parity)
405 - cute.arch.barrier(barrier_id=1, number_of_threads=128)
464 + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
465 _tcgen05.fence_after_thread_sync()
466
467 # tmem 16x256b layout / ldmatrix layout
468 + beta_row = cute.make_rmem_tensor(2, Float32)
469 + g_cu_row = cute.make_rmem_tensor(2, Float32)
470 + for i in cutlass.range_constexpr(2):
471 + idx = warp_id_ * 16 + i * 8 + (lane_id // 4)
472 + beta_row[i] = s_beta[idx, stage_id]
473 + g_cu_row[i] = s_g_cu[idx, stage_id]
474 +
475 # mode0 is 8 rows together
476 # mode1 is top and bottom 8 rows
477 # mode2 is groups of 16 rows
412 - row_coord = (lane_id // 4, None, warp_id_)
413 - s_beta_view = cute.make_tensor(s_beta, (8, 2, 4))
414 - beta_row = s_beta_view[row_coord].load().reshape((1, 2, 1))
415 -
416 - s_g_cu_view = cute.make_tensor(s_g_cu, (8, 2, 4))
417 - g_cu_row = s_g_cu_view[row_coord].load().reshape((1, 2, 1))
478 + beta_row = beta_row.load().reshape((1, 2, 1))
479 + g_cu_row = g_cu_row.load().reshape((1, 2, 1))
480
481 # mode0 is 2 consecutive elems
482 # mode1 is top and bottom 8 rows
@@ -429,7 +491,9 @@ class Sm100ChunkUWKernel:
491 # mode2 is top and bottom 8 rows
492 # mode3 is next 16 columns
493 col_coord = (None, lane_id % 4, None, i)
432 - s_g_cu_view = cute.make_tensor(s_g_cu, (2, 4, 2, BT // 16))
494 + s_g_cu_view = cute.make_tensor(
495 + s_g_cu[None, stage_id].iterator, (2, 4, 2, BT // 16)
496 + )
497 g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
498
499 Gamma = cute.math.exp(g_cu_row - g_cu_col, fastmath=True)
@@ -443,29 +507,23 @@ class Sm100ChunkUWKernel:
507 # pack to BF16
508 # CuteDSL doesn't generate cvt.bf16x2.f32 here for some reasons
509 packed = cute.make_rmem_tensor(4, Uint32)
446 - packed[0] = cvt.fp32x2_to_bf16x2(
447 - A_masked[0, 0, 0], A_masked[1, 0, 0]
448 - )
449 - packed[1] = cvt.fp32x2_to_bf16x2(
450 - A_masked[0, 1, 0], A_masked[1, 1, 0]
451 - )
452 - packed[2] = cvt.fp32x2_to_bf16x2(
453 - A_masked[0, 0, 1], A_masked[1, 0, 1]
454 - )
455 - packed[3] = cvt.fp32x2_to_bf16x2(
456 - A_masked[0, 1, 1], A_masked[1, 1, 1]
457 - )
510 + for j in cutlass.range_constexpr(4):
511 + packed[j] = cvt.fp32x2_to_bf16x2(
512 + A_masked[j * 2], A_masked[j * 2 + 1]
513 + )
514
515 # store to smem
516 cute.copy(
517 stsm_atom,
518 cute.recast_tensor(packed, BFloat16),
463 - sA_ldsm[warp_id_, None, i],
519 + sA_ldsm[None, (warp_id_, i)],
520 )
521
466 - cute.arch.barrier(barrier_id=1, number_of_threads=128)
522 + # use sync warp instead of bar.sync because for block-diagonal inverse,
523 + # each warp reads its own private smem memory.
524 + cute.arch.sync_warp()
525
468 - ##### Phase 3: matrix inverse #####
526 + ##### Phase 2: matrix inverse #####
527 # we use Newton-Schulz iterations to compute the inverse
528 # of the four 16x16 diagonal blocks.
529 # Ai_new = 2 Ai - Ai @ M @ Ai
@@ -477,15 +535,6 @@ class Sm100ChunkUWKernel:
535 zeros_f32 = cute.make_rmem_tensor(4, Float32)
536 zeros_f32.fill(0.0)
537
480 - def set_diagonal(A: cute.Tensor, lane_id: Int32):
481 - "Set the diagonal to 1s"
482 - if lane_id % 9 == 0:
483 - A[0] = (A[0] & Uint32(0xFFFF0000)) | Uint32(0x00003F80)
484 - A[3] = (A[3] & Uint32(0xFFFF0000)) | Uint32(0x00003F80)
485 - elif lane_id % 9 == 4:
486 - A[0] = (A[0] & Uint32(0x0000FFFF)) | Uint32(0x3F800000)
487 - A[3] = (A[3] & Uint32(0x0000FFFF)) | Uint32(0x3F800000)
488 -
538 Ai_bf16 = cute.make_rmem_tensor(8, BFloat16)
539 mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16)
540 M_bf16 = cute.make_rmem_tensor(8, BFloat16)
@@ -496,44 +545,60 @@ class Sm100ChunkUWKernel:
545 mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2)
546 M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2)
547
548 + # construct rmem-backed identity matrix
549 + eye = cute.make_rmem_tensor(4, Uint32)
550 + eye[0] = Uint32(lane_id % 9 == 0) * Uint32(0x00003F80) + Uint32(
551 + lane_id % 9 == 4
552 + ) * Uint32(0x3F800000)
553 + eye[1] = 0
554 + eye[2] = 0
555 + eye[3] = eye[0]
556 +
557 # initial guess: Ai = I-A
500 - cute.copy(ldsm_atom, sA_ldsm[warp_id_, None, warp_id_], Ai_bf16)
558 + cute.copy(ldsm_atom, sA_ldsm[None, (warp_id_, warp_id_)], Ai_bf16)
559 for i in cutlass.range_constexpr(4):
502 - Ai[i] ^= Uint32(0x80008000) # negate A
503 - set_diagonal(Ai, lane_id)
560 + Ai[i] = _bf16x2_sub(eye[i], Ai[i])
561
562 # (4, 2)
563 Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4)
564
565 # M is holding -(I+A), stay constant throughout the iterations
509 - cute.copy(ldsm_trans_atom, sA_ldsm[warp_id_, None, warp_id_], M_bf16)
510 - set_diagonal(M, lane_id)
566 + cute.copy(ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_)], M_bf16)
567 for i in cutlass.range_constexpr(4):
512 - M[i] ^= Uint32(0x80008000)
568 + M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i])
569
570 # 3 rounds of Newton-Schulz
571 for _ in cutlass.range_constexpr(3):
572 # First MMA: -AiM = Ai @ (-M)
517 - cute.copy(stsm_atom, Ai_bf16, sA_ldsm[warp_id_, None, warp_id_])
573 + cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)])
574 cute.arch.sync_warp()
575 acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32)
576 acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32)
577 Ai_bf16.store(acc.load().to(BFloat16))
578
579 # Second MMA: Ai_new = 2Ai + (-AiM) @ Ai
524 - for j in cutlass.range_constexpr(8):
580 + for j in cutlass.range(8, vectorize=True):
581 Ai_f32[j] *= 2.0
582 cute.copy(
583 ldsm_trans_atom,
528 - sA_ldsm[warp_id_, None, warp_id_],
584 + sA_ldsm[None, (warp_id_, warp_id_)],
585 mma_B_bf16,
586 )
587 Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0])
588 Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1])
589 Ai_bf16.store(Ai_f32.load().to(BFloat16))
590
535 - cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[warp_id_, None, warp_id_])
536 - cute.arch.barrier(barrier_id=1, number_of_threads=128)
591 + cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[None, (warp_id_, warp_id_)])
592 + store_ab_abg(
593 + Ai_f32,
594 + s_beta[None, stage_id],
595 + s_beta_g[None, stage_id],
596 + warp_id_,
597 + lane_id,
598 + warp_id_,
599 + Ab_tmem_base,
600 + )
601 + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
602
603 # off-diagonal by 1
604 # given
@@ -544,14 +609,14 @@ class Sm100ChunkUWKernel:
609 # warp1: Ai10 = -Ai11 @ A10 @ Ai00
610 # warp2: Ai21 = -Ai22 @ A21 @ Ai11
611 # warp3: Ai32 = -Ai33 @ A32 @ Ai22
547 - if warp_id_ > 0:
612 + if warp_id_ >= 1:
613 neg_Ai = cute.make_rmem_tensor(4, Uint32)
614 for i in cutlass.range_constexpr(4):
550 - neg_Ai[i] = Ai[i] ^ Uint32(0x80008000)
615 + neg_Ai[i] = _bf16x2_neg(Ai[i])
616
617 cute.copy(
618 ldsm_trans_atom,
554 - sA_ldsm[warp_id_, None, warp_id_ - 1],
619 + sA_ldsm[None, (warp_id_, warp_id_ - 1)],
620 mma_B_bf16,
621 )
622 acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32)
@@ -560,44 +625,52 @@ class Sm100ChunkUWKernel:
625
626 cute.copy(
627 ldsm_trans_atom,
563 - sAi_ldsm[warp_id_ - 1, None, warp_id_ - 1],
628 + sAi_ldsm[None, (warp_id_ - 1, warp_id_ - 1)],
629 mma_B_bf16,
630 )
631 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
632 acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
633 Ai_bf16.store(acc.load().to(BFloat16))
634 + store_ab_abg(
635 + acc,
636 + s_beta[None, stage_id],
637 + s_beta_g[None, stage_id],
638 + warp_id_,
639 + lane_id,
640 + warp_id_ - 1,
641 + Ab_tmem_base,
642 + )
643 cute.copy(
644 stsm_atom,
645 Ai_bf16,
572 - sAi_ldsm[warp_id_, None, warp_id_ - 1],
646 + sAi_ldsm[None, (warp_id_, warp_id_ - 1)],
647 )
574 - cute.arch.barrier(barrier_id=1, number_of_threads=128)
648 + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
649
650 # off-diagonal by 2
577 - # warp0: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10)
578 - # warp1: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21)
579 - if warp_id_ < 2:
651 + # warp2: Ai20 = -Ai22 @ (A20 @ Ai00 + A21 @ Ai10)
652 + # warp3: Ai31 = -Ai33 @ (A31 @ Ai11 + A32 @ Ai21)
653 + if warp_id_ >= 2:
654 + tile_col = warp_id_ - 2
655 cute.copy(
656 ldsm_atom,
582 - sA_ldsm[warp_id_ + 2, None, warp_id_],
657 + sA_ldsm[None, (warp_id_, tile_col)],
658 Ai_bf16,
659 )
660 cute.copy(
661 ldsm_trans_atom,
587 - sAi_ldsm[warp_id_, None, warp_id_],
662 + sAi_ldsm[None, (tile_col, tile_col)],
663 mma_B_bf16,
664 )
665 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
666 acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
667
668 cute.copy(
594 - ldsm_atom,
595 - sA_ldsm[warp_id_ + 2, None, warp_id_ + 1],
596 - Ai_bf16,
669 + ldsm_atom, sA_ldsm[None, (warp_id_, tile_col + 1)], Ai_bf16
670 )
671 cute.copy(
672 ldsm_trans_atom,
600 - sAi_ldsm[warp_id_ + 1, None, warp_id_],
673 + sAi_ldsm[None, (tile_col + 1, tile_col)],
674 mma_B_bf16,
675 )
676 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
@@ -605,79 +678,68 @@ class Sm100ChunkUWKernel:
678
679 tmp = cute.make_rmem_tensor(8, BFloat16)
680 tmp.store(acc.load().to(BFloat16))
608 - cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
681 + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)])
682 cute.arch.sync_warp()
683
611 - cute.copy(
612 - ldsm_atom, sAi_ldsm[warp_id_ + 2, None, warp_id_ + 2], Ai_bf16
613 - )
684 + cute.copy(ldsm_atom, sAi_ldsm[None, (warp_id_, warp_id_)], Ai_bf16)
685 for i in cutlass.range_constexpr(4):
615 - Ai[i] ^= Uint32(0x80008000)
686 + Ai[i] = _bf16x2_neg(Ai[i])
687 cute.copy(
688 ldsm_trans_atom,
618 - sAi_ldsm[warp_id_ + 2, None, warp_id_],
689 + sAi_ldsm[None, (warp_id_, tile_col)],
690 mma_B_bf16,
691 )
692 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
693 acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
694 tmp.store(acc.load().to(BFloat16))
624 - cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
625 - cute.arch.barrier(barrier_id=1, number_of_threads=128)
695 + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)])
696 + store_ab_abg(
697 + acc,
698 + s_beta[None, stage_id],
699 + s_beta_g[None, stage_id],
700 + warp_id_,
701 + lane_id,
702 + tile_col,
703 + Ab_tmem_base,
704 + )
705 + cute.arch.barrier(barrier_id=INV_BAR, number_of_threads=128)
706
707 # off-diagonal by 3
628 - # warp0: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20)
629 - if warp_id_ == 0:
630 - cute.copy(ldsm_atom, sA_ldsm[3, None, 0], Ai_bf16)
631 - cute.copy(ldsm_trans_atom, sAi_ldsm[0, None, 0], mma_B_bf16)
708 + # warp3: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20)
709 + if warp_id_ == 3:
710 + cute.copy(ldsm_atom, sA_ldsm[None, (3, 0)], Ai_bf16)
711 + cute.copy(ldsm_trans_atom, sAi_ldsm[None, (0, 0)], mma_B_bf16)
712 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
713 acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
714
715 for i in cutlass.range_constexpr(1, 3):
636 - cute.copy(ldsm_atom, sA_ldsm[3, None, i], Ai_bf16)
637 - cute.copy(ldsm_trans_atom, sAi_ldsm[i, None, 0], mma_B_bf16)
716 + cute.copy(ldsm_atom, sA_ldsm[None, (3, i)], Ai_bf16)
717 + cute.copy(ldsm_trans_atom, sAi_ldsm[None, (i, 0)], mma_B_bf16)
718 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
719 acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])
720
721 tmp = cute.make_rmem_tensor(8, BFloat16)
722 tmp.store(acc.load().to(BFloat16))
643 - cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
723 + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)])
724 cute.arch.sync_warp()
725
646 - cute.copy(ldsm_atom, sAi_ldsm[3, None, 3], Ai_bf16)
726 + cute.copy(ldsm_atom, sAi_ldsm[None, (3, 3)], Ai_bf16)
727 for i in cutlass.range_constexpr(4):
648 - Ai[i] ^= Uint32(0x80008000)
649 - cute.copy(ldsm_trans_atom, sAi_ldsm[3, None, 0], mma_B_bf16)
728 + Ai[i] = _bf16x2_neg(Ai[i])
729 + cute.copy(ldsm_trans_atom, sAi_ldsm[None, (3, 0)], mma_B_bf16)
730 acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
731 acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
732 tmp.store(acc.load().to(BFloat16))
653 - cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
654 -
655 - ##### Phase 4: compute Ab, Abg #####
656 - if warp_id_ == 3:
657 - cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity ^ 1)
658 - cute.arch.barrier(barrier_id=1, number_of_threads=128)
659 -
660 - for i in cutlass.range_constexpr(BT // 16):
661 - cute.copy(ldsm_atom, sAi_ldsm[warp_id_, None, i], Ai_bf16)
662 -
663 - col_coord = (None, lane_id % 4, None, i)
664 - s_beta_view = cute.make_tensor(s_beta, (2, 4, 2, BT // 16))
665 - beta_col = s_beta_view[col_coord].load().reshape((2, 1, 2))
666 -
667 - s_g_cu_view = cute.make_tensor(s_g_cu_exp, (2, 4, 2, BT // 16))
668 - g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
669 -
670 - Ai_f32 = cvt.bf16x2_to_fp32x2(Ai).load().reshape((2, 2, 2))
671 -
672 - Ab_f32 = Ai_f32 * beta_col
673 - Ab = Ab_f32.to(BFloat16)
674 - Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id + i * 8
675 - _tcgen05.st(warp_id_ * 32, Ab_tmem, "16x128b", 2, Ab)
676 -
677 - Abg_f32 = Ab_f32 * g_cu_col
678 - Abg = Abg_f32.to(BFloat16)
679 - _tcgen05.st(warp_id_ * 32 + 16, Ab_tmem, "16x128b", 2, Abg)
680 -
733 + cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)])
734 + store_ab_abg(
735 + acc,
736 + s_beta[None, stage_id],
737 + s_beta_g[None, stage_id],
738 + warp_id_,
739 + lane_id,
740 + 0,
741 + Ab_tmem_base,
742 + )
743 _tcgen05.wait_st()
744 _tcgen05.fence_before_thread_sync()
745 cute.arch.mbarrier_arrive(inv_mbar + stage_id)
@@ -686,14 +748,18 @@ class Sm100ChunkUWKernel:
748 if stage_id == 0:
749 parity ^= 1
750
689 - elif warp_id < 4:
751 + else:
752 # epi warps
753 stage_id = 0
754 parity = 0
755
756 # ((BT, num_global_chunks), V_dim)
695 - gU_tiles = cute.logical_divide(tmaU[None, head_id, None], (BT, None))
696 - gW_tiles = cute.logical_divide(tmaW[None, head_id, None], (BT, None))
757 + gU_tiles = cute.logical_divide(
758 + U_tma.tma_tensor[None, head_id, None], (BT, None)
759 + )
760 + gW_tiles = cute.logical_divide(
761 + W_tma.tma_tensor[None, head_id, None], (BT, None)
762 + )
763
764 # sW shape: [BT, (64, K_dim/64)]
765 # sW_view shape: [(8, 2), (4, K_dim/64)]
@@ -719,7 +785,7 @@ class Sm100ChunkUWKernel:
785 elif warp_id == 1:
786 with cute.arch.elect_one():
787 cute.arch.cp_async_bulk_wait_group(0, read=True)
722 - cute.arch.barrier(barrier_id=2, number_of_threads=128)
788 + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
789 _tcgen05.fence_after_thread_sync()
790
791 w_f32 = _tcgen05.ld(warp_id * 32 + 16, U_tmem, "16x256b", K_dim // 8)
@@ -729,16 +795,16 @@ class Sm100ChunkUWKernel:
795 cute.copy(stsm_atom, w_bf16, sW_view)
796
797 # wait for U MMA + issue W TMA store
732 - cute.arch.barrier(barrier_id=2, number_of_threads=128)
798 + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
799 fence_before_tma_store()
800 if warp_id == 0:
801 cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity)
802 elif warp_id == 1:
803 # don't need to commit
804 simple_tma_copy(
739 - W_tma_atom, sW, gW_tiles[(None, global_chunk_id), None]
805 + W_tma.atom, sW, gW_tiles[(None, global_chunk_id), None]
806 )
741 - cute.arch.barrier(barrier_id=2, number_of_threads=128)
807 + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
808 _tcgen05.fence_after_thread_sync()
809
810 u_f32 = _tcgen05.ld(warp_id * 32, U_tmem, "16x256b", V_dim // 8)
@@ -749,11 +815,11 @@ class Sm100ChunkUWKernel:
815 u_bf16.store(u_f32.to(BFloat16))
816 cute.copy(stsm_atom, u_bf16, sU_view)
817
752 - cute.arch.barrier(barrier_id=2, number_of_threads=128)
818 + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=128)
819 fence_before_tma_store()
820 if warp_id == 1:
821 simple_tma_copy(
756 - U_tma_atom, sU, gU_tiles[(None, global_chunk_id), None]
822 + U_tma.atom, sU, gU_tiles[(None, global_chunk_id), None]
823 )
824 with cute.arch.elect_one():
825 cute.arch.cp_async_bulk_commit_group()
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py
+82 -80
@@ -61,13 +61,12 @@ class Sm100ChunkOKernel:
61 stride=(64, 0, (1, self.BT * 64), self.BT * dim),
62 )
63 slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
64 - atom, tma_tensor = cpasync.make_tiled_tma_atom(
64 + return cpasync.make_tiled_tma_atom(
65 op,
66 cute.logical_divide(tensor, (None, None, 64)),
67 slayout,
68 cta_tiler=(self.BT, 1, dim),
69 )
70 - return atom, tma_tensor, slayout
70
71 @cute.jit
72 def _make_h_tma_args(
@@ -79,17 +78,22 @@ class Sm100ChunkOKernel:
78 num_elems = 128 // (tensor.element_type.width // 8)
79 swizzle_128B = cute.make_swizzle(3, 4, 3)
80 slayout = cute.make_layout(
82 - (1, self.V_dim, (num_elems, self.K_dim // num_elems), stages),
83 - stride=(0, num_elems, (1, self.V_dim * num_elems), self.V_dim * self.K_dim),
81 + (1, 1, self.V_dim, (num_elems, self.K_dim // num_elems), stages),
82 + stride=(
83 + 0,
84 + 0,
85 + num_elems,
86 + (1, self.V_dim * num_elems),
87 + self.V_dim * self.K_dim,
88 + ),
89 )
90 slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
86 - atom, tma_tensor = cpasync.make_tiled_tma_atom(
91 + return cpasync.make_tiled_tma_atom(
92 op,
88 - cute.logical_divide(tensor, (None, None, num_elems)),
93 + cute.logical_divide(tensor, (None, None, None, num_elems)),
94 slayout,
90 - cta_tiler=(1, self.V_dim, self.K_dim),
95 + cta_tiler=(1, 1, self.V_dim, self.K_dim),
96 )
92 - return atom, tma_tensor, slayout
97
98 @cute.jit
99 def __call__(
@@ -111,35 +115,35 @@ class Sm100ChunkOKernel:
115 block = (self.num_warps * 32, 1, 1)
116 tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
117 tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
114 - Q_args = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages)
115 - K_args = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages)
116 - V_args = self._make_bf16_tma_args(
118 + Q_tma = self._make_bf16_tma_args(q, self.K_dim, tma_g2s, self.num_stages)
119 + K_tma = self._make_bf16_tma_args(k, self.K_dim, tma_g2s, self.num_stages)
120 + V_tma = self._make_bf16_tma_args(
121 v_new_chunks, self.V_dim, tma_g2s, self.num_stages
122 )
119 - H_args = self._make_h_tma_args(h, tma_g2s, self.num_stages)
120 - O_args = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1)
123 + H_tma = self._make_h_tma_args(h, tma_g2s, self.num_stages)
124 + O_tma = self._make_bf16_tma_args(o, self.V_dim, tma_s2g, 1)
125 self.kernel(
122 - Q_args,
123 - K_args,
124 - V_args,
125 - H_args,
126 - O_args,
126 + Q_tma,
127 + K_tma,
128 + V_tma,
129 + H_tma,
130 + O_tma,
131 g_cu,
132 o,
133 cu_seqlens,
134 chunk_indices,
135 total_chunks,
136 scale,
133 - ).launch(grid=grid, block=block, stream=stream)
137 + ).launch(grid=grid, block=block, min_blocks_per_mp=1, stream=stream)
138
139 @cute.kernel
140 def kernel(
141 self,
138 - Q_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
139 - K_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
140 - V_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
141 - H_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
142 - O_args: tuple[cute.CopyAtom, cute.Tensor, cute.ComposedLayout],
142 + Q_tma: cpasync.TmaInfo,
143 + K_tma: cpasync.TmaInfo,
144 + V_tma: cpasync.TmaInfo,
145 + H_tma: cpasync.TmaInfo,
146 + O_tma: cpasync.TmaInfo,
147 g_cu: cute.Tensor,
148 o: cute.Tensor,
149 cu_seqlens: cute.Tensor,
@@ -162,23 +166,17 @@ class Sm100ChunkOKernel:
166 k_head_id = v_head_id // heads_per_qk
167 num_global_chunks = total_chunks[0]
168
165 - Q_tma_atom, tmaQ, sQ_layout = Q_args
166 - K_tma_atom, tmaK, sK_layout = K_args
167 - V_tma_atom, tmaV, sV_layout = V_args
168 - H_tma_atom, tmaH, sH_layout = H_args
169 - O_tma_atom, tmaO, sO_layout = O_args
170 -
169 def allocate_tensor(smem, dtype, layout):
170 return smem.allocate_tensor(
171 dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
172 )
173
174 smem = cutlass.utils.SmemAllocator()
177 - sQ = allocate_tensor(smem, BFloat16, sQ_layout)[None, 0, None, None]
178 - sK = allocate_tensor(smem, BFloat16, sK_layout)[None, 0, None, None]
179 - sV = allocate_tensor(smem, BFloat16, sV_layout)[None, 0, None, None]
180 - sH = allocate_tensor(smem, BFloat16, sH_layout)[0, None, None, None]
181 - sO = allocate_tensor(smem, BFloat16, sO_layout)[None, 0, None, 0]
175 + sQ = allocate_tensor(smem, BFloat16, Q_tma.smem_layout)[None, 0, None, None]
176 + sK = allocate_tensor(smem, BFloat16, K_tma.smem_layout)[None, 0, None, None]
177 + sV = allocate_tensor(smem, BFloat16, V_tma.smem_layout)[None, 0, None, None]
178 + sH = allocate_tensor(smem, BFloat16, H_tma.smem_layout)[0, 0, None, None, None]
179 + sO = allocate_tensor(smem, BFloat16, O_tma.smem_layout)[None, 0, None, 0]
180
181 s_g_cu = smem.allocate_array(Float32, BT)
182 qk_full_mbar = smem.allocate_array(Int64, num_stages)
@@ -207,10 +205,10 @@ class Sm100ChunkOKernel:
205 cute.arch.mbarrier_init(epi_mbar, 128)
206 cute.arch.mbarrier_init_fence()
207 elif warp_id == 9:
210 - cpasync.prefetch_descriptor(Q_tma_atom)
211 - cpasync.prefetch_descriptor(K_tma_atom)
212 - cpasync.prefetch_descriptor(V_tma_atom)
213 - cpasync.prefetch_descriptor(H_tma_atom)
208 + cpasync.prefetch_descriptor(Q_tma.atom)
209 + cpasync.prefetch_descriptor(K_tma.atom)
210 + cpasync.prefetch_descriptor(V_tma.atom)
211 + cpasync.prefetch_descriptor(H_tma.atom)
212 cute.arch.sync_threads()
213
214 if warp_id == 9:
@@ -225,12 +223,16 @@ class Sm100ChunkOKernel:
223
224 # copy Q and K
225 q_tile = cute.local_tile(
228 - cute.domain_offset((bos, 0), tmaQ[None, k_head_id, None]),
226 + cute.domain_offset(
227 + (bos, 0), Q_tma.tma_tensor[None, k_head_id, None]
228 + ),
229 tiler=(BT, K_dim),
230 coord=(chunk_id, 0),
231 )
232 k_tile = cute.local_tile(
233 - cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
233 + cute.domain_offset(
234 + (bos, 0), K_tma.tma_tensor[None, k_head_id, None]
235 + ),
236 tiler=(BT, K_dim),
237 coord=(chunk_id, 0),
238 )
@@ -241,13 +243,13 @@ class Sm100ChunkOKernel:
243 with cute.arch.elect_one():
244 STAGE_SIZE = BT * (K_dim + K_dim) * 2
245 cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
244 - simple_tma_copy(Q_tma_atom, q_tile, sQ[None, None, stage_id], mbar)
245 - simple_tma_copy(K_tma_atom, k_tile, sK[None, None, stage_id], mbar)
246 + simple_tma_copy(Q_tma.atom, q_tile, sQ[None, None, stage_id], mbar)
247 + simple_tma_copy(K_tma.atom, k_tile, sK[None, None, stage_id], mbar)
248
249 # copy H and V
248 - gH = tmaH[global_chunk_id * self.Hv + v_head_id, None, None]
250 + gH = H_tma.tma_tensor[global_chunk_id, v_head_id, None, None]
251 gV = cute.local_tile(
250 - tmaV[None, v_head_id, None],
252 + V_tma.tma_tensor[None, v_head_id, None],
253 tiler=(BT, V_dim),
254 coord=(global_chunk_id, 0),
255 )
@@ -262,10 +264,10 @@ class Sm100ChunkOKernel:
264 mbar, H_STAGE_SIZE + V_STAGE_SIZE
265 )
266 simple_tma_copy(
265 - H_tma_atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST
267 + H_tma.atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST
268 )
269 simple_tma_copy(
268 - V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
270 + V_tma.atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
271 )
272
273 stage_id = (stage_id + 1) % num_stages
@@ -302,40 +304,37 @@ class Sm100ChunkOKernel:
304 cute.arch.mbarrier_wait(qk_full_mbar + stage_id, tma_parity)
305 _tcgen05.fence_after_thread_sync()
306
305 - with cute.arch.elect_one():
306 - for i in cutlass.range_constexpr(K_dim // BT):
307 - for j in cutlass.range_constexpr(BT // 16):
308 - qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
309 - kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
310 - _tcgen05.mma_f16(
311 - qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0)
312 - )
313 - _tcgen05.commit(qk_mbar)
307 + for i in cutlass.range_constexpr(K_dim // 64):
308 + for j in cutlass.range_constexpr(64 // 16):
309 + qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
310 + kdesc = kdesc_base | ((i * BT * 128 + j * 32) >> 4)
311 + _tcgen05.mma_f16(
312 + qk_tmem, qdesc, kdesc, qk_idesc, (i > 0) or (j > 0)
313 + )
314 + _tcgen05.commit(qk_mbar)
315
316 ##### 2nd MMA: Q @ H.T #####
317 cute.arch.mbarrier_wait(hv_full_mbar + stage_id, tma_parity)
318 _tcgen05.fence_after_thread_sync()
318 - with cute.arch.elect_one():
319 - for i in cutlass.range_constexpr(K_dim // BT):
320 - for j in cutlass.range_constexpr(BT // 16):
321 - qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
322 - hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4)
323 - _tcgen05.mma_f16(
324 - qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0)
325 - )
326 - _tcgen05.commit(qk_empty_mbar + stage_id)
319 + for i in cutlass.range_constexpr(K_dim // 64):
320 + for j in cutlass.range_constexpr(64 // 16):
321 + qdesc = qdesc_base | ((i * BT * 128 + j * 32) >> 4)
322 + hdesc = hdesc_base | ((i * V_dim * 128 + j * 32) >> 4)
323 + _tcgen05.mma_f16(
324 + qh_tmem, qdesc, hdesc, qh_idesc, (i > 0) or (j > 0)
325 + )
326 + _tcgen05.commit(qk_empty_mbar + stage_id)
327
328 ##### 3rd MMA: P @ V #####
329 # stalled by mask(QK)
330 cute.arch.mbarrier_wait(mask_mbar, mask_parity)
331 _tcgen05.fence_after_thread_sync()
332 - with cute.arch.elect_one():
333 - for i in cutlass.range_constexpr(BT // 16):
334 - vdesc = vdesc_base | ((i * 16 * 128) >> 4)
335 - _tcgen05.mma_ts_f16(
336 - out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0
337 - )
338 - _tcgen05.commit(pv_mma_mbar + stage_id)
332 + for i in cutlass.range_constexpr(BT // 16):
333 + vdesc = vdesc_base | ((i * 16 * 128) >> 4)
334 + _tcgen05.mma_ts_f16(
335 + out_tmem, p_tmem + i * 8, vdesc, pv_idesc, i > 0
336 + )
337 + _tcgen05.commit(pv_mma_mbar + stage_id)
338
339 stage_id = (stage_id + 1) % num_stages
340 if stage_id == 0:
@@ -499,11 +498,13 @@ class Sm100ChunkOKernel:
498 fence_before_tma_store()
499 if warp_id == 3:
500 gO = cute.local_tile(
502 - cute.domain_offset((bos, 0), tmaO[None, v_head_id, None]),
501 + cute.domain_offset(
502 + (bos, 0), O_tma.tma_tensor[None, v_head_id, None]
503 + ),
504 tiler=(BT, V_dim),
505 coord=(chunk_id, 0),
506 )
506 - simple_tma_copy(O_tma_atom, sO, gO)
507 + simple_tma_copy(O_tma.atom, sO, gO)
508 with cute.arch.elect_one():
509 cute.arch.cp_async_bulk_commit_group()
510
@@ -559,13 +560,14 @@ class Sm100ChunkOKernel:
560 total_t = cute.sym_int()
561 pad_t = cute.sym_int()
562 total_chunks_n = cute.sym_int()
562 - h_outer_n = cute.sym_int()
563 cu_entries = cute.sym_int()
564
565 q = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
566 k = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
567 v_new = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
568 - h_flat = make_fake_tensor(BFloat16, (h_outer_n, V_dim, K_dim), divisibility=16)
568 + h = make_fake_tensor(
569 + BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16
570 + )
571 g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
572 o = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16)
573 cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
@@ -586,7 +588,7 @@ class Sm100ChunkOKernel:
588 q,
589 k,
590 v_new,
589 - h_flat,
591 + h,
592 g_cu,
593 o,
594 cu_seqlens,
@@ -602,7 +604,7 @@ class Sm100ChunkOKernel:
604 def o_cutedsl(
605 q: torch.Tensor,
606 k: torch.Tensor,
605 - v_new_chunks: torch.Tensor,
607 + v_new: torch.Tensor,
608 h: torch.Tensor,
609 g_cu: torch.Tensor,
610 o: torch.Tensor,
@@ -618,8 +620,8 @@ def o_cutedsl(
620 Sm100ChunkOKernel.compile(H, Hv, K_dim, V_dim)(
621 q,
622 k,
621 - v_new_chunks.view(-1, Hv, V_dim),
622 - h.view(-1, V_dim, K_dim),
623 + v_new,
624 + h,
625 g_cu,
626 o,
627 cu_seqlens,