A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 20:45 UTC
[GDN] GDN Prefill kernel for SM100 (#43273)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
May 26, 2026 at 13:02 UTC
d56612c6210605aaca1ce8fa09c40bbe30ee3973
13 files changed
+3265
-160
tests/kernels/mamba/test_gdn_prefill_cutedsl.py
new
+199
@@ -0,0 +1,199 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import math
5
+
6
+import pytest
7
+import torch
8
+import torch.nn.functional as F
9
+
10
+from vllm.platforms import current_platform
11
+
12
+if not (
13
+ current_platform.is_cuda() and current_platform.is_device_capability_family(100)
14
+):
15
+ pytest.skip(
16
+ reason="GDN CuteDSL prefill requires CUDA SM10x.",
17
+ allow_module_level=True,
18
+ )
19
+
20
+from vllm.model_executor.layers.fla.ops import ( # noqa: E402
21
+ chunk_gated_delta_rule,
22
+)
23
+from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402
24
+ prepare_chunk_indices,
25
+ prepare_chunk_offsets,
26
+)
27
+from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( # noqa: E402
28
+ chunk_gated_delta_rule_cutedsl,
29
+ prepare_metadata_cutedsl,
30
+)
31
+
32
+
33
+@pytest.mark.parametrize("num_seqs", [1, 5, 257])
34
+@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
35
+def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
36
+ seq_lens = torch.randint(
37
+ 1,
38
+ 130,
39
+ (num_seqs,),
40
+ dtype=torch.int32,
41
+ )
42
+ cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
43
+ cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
44
+ total_tokens = int(cu_seqlens[-1].item())
45
+
46
+ num_k_heads = 4
47
+ num_v_heads = 8
48
+ head_k_dim = 128
49
+ head_v_dim = 128
50
+ dtype = torch.bfloat16
51
+
52
+ q = torch.randn(
53
+ 1,
54
+ total_tokens,
55
+ num_k_heads,
56
+ head_k_dim,
57
+ device="cuda",
58
+ dtype=dtype,
59
+ )
60
+ k = torch.randn_like(q)
61
+ v = torch.randn(
62
+ 1,
63
+ total_tokens,
64
+ num_v_heads,
65
+ head_v_dim,
66
+ device="cuda",
67
+ dtype=dtype,
68
+ )
69
+ q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
70
+ k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
71
+ a = torch.randn(
72
+ 1,
73
+ total_tokens,
74
+ num_v_heads,
75
+ device="cuda",
76
+ dtype=dtype,
77
+ )
78
+ b = torch.randn(
79
+ 1,
80
+ total_tokens,
81
+ num_v_heads,
82
+ device="cuda",
83
+ dtype=dtype,
84
+ )
85
+ # Match upstream FLA GatedDeltaNet synthetic initialization:
86
+ # https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
87
+ A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
88
+ A_log = torch.log(A)
89
+ dt = torch.exp(
90
+ torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
91
+ * (math.log(0.1) - math.log(0.001))
92
+ + math.log(0.001)
93
+ )
94
+ dt = torch.clamp(dt, min=1e-4)
95
+ dt_bias = dt + torch.log(-torch.expm1(-dt))
96
+ g = -A_log.exp().view(1, 1, num_v_heads) * F.softplus(
97
+ a.float() + dt_bias.view(1, 1, num_v_heads)
98
+ )
99
+ beta = torch.sigmoid(b.float())
100
+ initial_state = (
101
+ torch.randn(
102
+ num_seqs,
103
+ num_v_heads,
104
+ head_v_dim,
105
+ head_k_dim,
106
+ device="cuda",
107
+ dtype=state_dtype,
108
+ )
109
+ * 0.05
110
+ )
111
+
112
+ # check metadata kernel
113
+ chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens)
114
+ torch.accelerator.synchronize()
115
+
116
+ expected_indices = prepare_chunk_indices(cu_seqlens, 64)
117
+ expected_offsets = prepare_chunk_offsets(cu_seqlens, 64)
118
+ total_chunks = int(expected_offsets[-1].item())
119
+
120
+ torch.testing.assert_close(chunk_offsets, expected_offsets.to(torch.int32))
121
+ torch.testing.assert_close(
122
+ chunk_indices[:total_chunks],
123
+ expected_indices,
124
+ )
125
+
126
+ ref_o, ref_state = chunk_gated_delta_rule(
127
+ q=q,
128
+ k=k,
129
+ v=v,
130
+ g=g,
131
+ beta=beta,
132
+ initial_state=initial_state,
133
+ output_final_state=True,
134
+ cu_seqlens=cu_seqlens,
135
+ use_qk_l2norm_in_kernel=False,
136
+ )
137
+ actual_core_attn_out = torch.empty(
138
+ total_tokens,
139
+ num_v_heads,
140
+ head_v_dim,
141
+ device="cuda",
142
+ dtype=dtype,
143
+ )
144
+ actual_o, actual_state = chunk_gated_delta_rule_cutedsl(
145
+ q=q,
146
+ k=k,
147
+ v=v,
148
+ g=g,
149
+ beta=beta,
150
+ initial_state=initial_state,
151
+ cu_seqlens=cu_seqlens,
152
+ chunk_indices=chunk_indices,
153
+ chunk_offsets=chunk_offsets,
154
+ core_attn_out=actual_core_attn_out,
155
+ )
156
+ torch.accelerator.synchronize()
157
+
158
+ # check main kernel
159
+ o_error = (actual_o.float() - ref_o.float()).abs()
160
+ state_error = (
161
+ actual_state.float() - ref_state.to(actual_state.dtype).float()
162
+ ).abs()
163
+ assert o_error.max().item() < 2e-3
164
+ assert o_error.mean().item() < 6e-5
165
+ assert state_error.max().item() < 2e-2
166
+ assert state_error.mean().item() < 6e-4
167
+ core_attn_out_error = (
168
+ actual_core_attn_out.float() - actual_o.squeeze(0).float()
169
+ ).abs()
170
+ assert core_attn_out_error.max().item() == 0
171
+
172
+ # check main kernel when core_attn_out is not passed
173
+ no_buffer_o, no_buffer_state = chunk_gated_delta_rule_cutedsl(
174
+ q=q,
175
+ k=k,
176
+ v=v,
177
+ g=g,
178
+ beta=beta,
179
+ initial_state=initial_state,
180
+ cu_seqlens=cu_seqlens,
181
+ chunk_indices=chunk_indices,
182
+ chunk_offsets=chunk_offsets,
183
+ )
184
+ torch.accelerator.synchronize()
185
+
186
+ no_buffer_o_error = (no_buffer_o.float() - ref_o.float()).abs()
187
+ no_buffer_state_error = (
188
+ no_buffer_state.float() - ref_state.to(no_buffer_state.dtype).float()
189
+ ).abs()
190
+ buffer_o_error = (no_buffer_o.float() - actual_o.float()).abs()
191
+ buffer_state_error = (
192
+ no_buffer_state.float() - actual_state.to(no_buffer_state.dtype).float()
193
+ ).abs()
194
+ assert no_buffer_o_error.max().item() < 2e-3
195
+ assert no_buffer_o_error.mean().item() < 6e-5
196
+ assert no_buffer_state_error.max().item() < 2e-2
197
+ assert no_buffer_state_error.mean().item() < 6e-4
198
+ assert buffer_o_error.max().item() == 0
199
+ assert buffer_state_error.max().item() == 0
vllm/cute_utils/__init__.py
new
+134
@@ -0,0 +1,134 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from cutlass import BFloat16, Float32, Int64, Uint32, cute
4
+from cutlass._mlir import ir
5
+from cutlass._mlir.dialects import llvm, vector
6
+from cutlass.cute.nvgpu import cpasync
7
+from cutlass.cutlass_dsl import T, dsl_user_op
8
+
9
+# https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197
10
+EVICT_NORMAL = Int64(0x1000000000000000)
11
+EVICT_FIRST = Int64(0x12F0000000000000)
12
+EVICT_LAST = Int64(0x14F0000000000000)
13
+
14
+
15
+@dsl_user_op
16
+def recast_val(x, dtype, *, loc=None, ip=None):
17
+ return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
18
+
19
+
20
+def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
21
+ """A simple helper that wraps group_modes() and tma_partition()
22
+ NOTE: this should be called WITHOUT cute.elect_one()
23
+ """
24
+ if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
25
+ gmem = src
26
+ smem = dst
27
+ elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
28
+ smem = src
29
+ gmem = dst
30
+ else:
31
+ raise ValueError
32
+
33
+ s_part, g_part = cpasync.tma_partition(
34
+ atom,
35
+ 0,
36
+ cute.make_layout(1),
37
+ cute.group_modes(smem, 0),
38
+ cute.group_modes(gmem, 0),
39
+ )
40
+
41
+ if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
42
+ cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy)
43
+ elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
44
+ cute.copy(atom, s_part, g_part, cache_policy=cache_policy)
45
+ else:
46
+ raise ValueError
47
+
48
+
49
+# can't find the equivalent in nvvm
50
+@dsl_user_op
51
+def fence_before_tma_store(*, loc=None, ip=None):
52
+ llvm.inline_asm(
53
+ T.i32(),
54
+ [],
55
+ "mov.u32 $0, 0;\n\t"
56
+ "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;",
57
+ "=r",
58
+ has_side_effects=True,
59
+ is_align_stack=False,
60
+ loc=loc,
61
+ ip=ip,
62
+ )
63
+
64
+
65
+@dsl_user_op
66
+def mma_bf16(
67
+ a: cute.TensorSSA, b: cute.TensorSSA, c: cute.TensorSSA, *, loc=None, ip=None
68
+):
69
+ if a.element_type == BFloat16:
70
+ a = cute.recast_tensor(a, Uint32)
71
+ if b.element_type == BFloat16:
72
+ b = cute.recast_tensor(b, Uint32)
73
+
74
+ mlir_ty = Float32.mlir_type
75
+ out = llvm.inline_asm(
76
+ llvm.StructType.get_literal([mlir_ty] * 4),
77
+ [a[i].ir_value(loc=loc, ip=ip) for i in range(4)]
78
+ + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)]
79
+ + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)],
80
+ "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
81
+ "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, "
82
+ "{$10, $11, $12, $13};",
83
+ "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
84
+ has_side_effects=False,
85
+ is_align_stack=False,
86
+ loc=loc,
87
+ ip=ip,
88
+ )
89
+ vec = vector.from_elements(
90
+ ir.VectorType.get([4], mlir_ty, loc=loc),
91
+ [llvm.extractvalue(mlir_ty, out, [i], loc=loc, ip=ip) for i in range(4)],
92
+ loc=loc,
93
+ ip=ip,
94
+ )
95
+ return cute.TensorSSA(vec, 4, Float32)
96
+
97
+
98
+@dsl_user_op
99
+def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
100
+ out = llvm.inline_asm(
101
+ T.i32(),
102
+ [a.ir_value(loc=loc, ip=ip)],
103
+ "abs.bf16x2 $0, $1;",
104
+ "=r,r",
105
+ has_side_effects=False,
106
+ is_align_stack=False,
107
+ )
108
+ return Uint32(out)
109
+
110
+
111
+@dsl_user_op
112
+def _bf16x2_max(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;",
117
+ "=r,r,r",
118
+ has_side_effects=False,
119
+ is_align_stack=False,
120
+ )
121
+ return Uint32(out)
122
+
123
+
124
+@dsl_user_op
125
+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)
vllm/cute_utils/_tcgen05.py
new
+219
@@ -0,0 +1,219 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# this module is named _tcgen05 to avoid name collision with cute.nvgpu.tcgen05
4
+
5
+import cutlass
6
+from cutlass import Boolean, Float32, Int32, Uint32, Uint64, cute
7
+from cutlass._mlir import ir
8
+from cutlass._mlir.dialects import llvm, nvvm, vector
9
+from cutlass.cutlass_dsl import dsl_user_op
10
+
11
+NVVM_CTA_GROUP_MAP = [
12
+ None,
13
+ nvvm.Tcgen05GroupKind.CTA_1,
14
+ nvvm.Tcgen05GroupKind.CTA_2,
15
+]
16
+LDST_MAP = {
17
+ "32x32b": (nvvm.Tcgen05LdStShape.SHAPE_32X32B, 1),
18
+ "16x128b": (nvvm.Tcgen05LdStShape.SHAPE_16X128B, 2),
19
+ "16x256b": (nvvm.Tcgen05LdStShape.SHAPE_16X256B, 4),
20
+}
21
+
22
+
23
+def _make_tmem_llvm_ptr(addr, *, loc=None, ip=None):
24
+ ptr_ty = llvm.PointerType.get(cute.AddressSpace.tmem.value)
25
+ val = Int32(addr).ir_value(loc=loc, ip=ip)
26
+ return llvm.inttoptr(ptr_ty, val, loc=loc, ip=ip)
27
+
28
+
29
+@dsl_user_op
30
+def alloc(
31
+ taddr: cute.Pointer,
32
+ cta_group: int = 1,
33
+ *,
34
+ loc=None,
35
+ ip=None,
36
+) -> None:
37
+ nvvm.tcgen05_alloc(
38
+ taddr.to_llvm_ptr(loc=loc, ip=ip),
39
+ Uint32(512).ir_value(loc=loc, ip=ip),
40
+ group=NVVM_CTA_GROUP_MAP[cta_group],
41
+ loc=loc,
42
+ ip=ip,
43
+ )
44
+
45
+
46
+@dsl_user_op
47
+def dealloc(cta_group: int = 1, *, loc=None, ip=None) -> None:
48
+ nvvm.tcgen05_dealloc(
49
+ _make_tmem_llvm_ptr(0, loc=loc, ip=ip),
50
+ Int32(512).ir_value(loc=loc, ip=ip),
51
+ group=NVVM_CTA_GROUP_MAP[cta_group],
52
+ loc=loc,
53
+ ip=ip,
54
+ )
55
+
56
+
57
+def make_bf16_idesc(
58
+ MMA_M: int,
59
+ MMA_N: int,
60
+ *,
61
+ negate_A: bool = False,
62
+ negate_B: bool = False,
63
+ transpose_A: bool = False,
64
+ transpose_B: bool = False,
65
+):
66
+ idesc = Uint32(
67
+ (1 << 4) | (1 << 7) | (1 << 10) | ((MMA_N >> 3) << 17) | ((MMA_M >> 4) << 24)
68
+ )
69
+ idesc |= Uint32(negate_A) << 13
70
+ idesc |= Uint32(negate_B) << 14
71
+ idesc |= Uint32(transpose_A) << 15
72
+ idesc |= Uint32(transpose_B) << 16
73
+ return idesc
74
+
75
+
76
+def make_sdesc_128B_swizzle(LBO: int):
77
+ SBO = 8 * 128
78
+ return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61))
79
+
80
+
81
+@dsl_user_op
82
+def mma_f16(
83
+ d_tmem,
84
+ a_desc,
85
+ b_desc,
86
+ idesc,
87
+ enable_input_d,
88
+ cta_group: int = 1,
89
+ *,
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
+ )
104
+
105
+
106
+@dsl_user_op
107
+def mma_ts_f16(
108
+ d_tmem,
109
+ a_tmem,
110
+ b_desc,
111
+ idesc,
112
+ enable_input_d,
113
+ cta_group: int = 1,
114
+ *,
115
+ loc=None,
116
+ ip=None,
117
+) -> 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
+ )
129
+
130
+
131
+@dsl_user_op
132
+def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
133
+ mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
134
+ group = NVVM_CTA_GROUP_MAP[cta_group]
135
+ 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
+ )
143
+ else:
144
+ nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
145
+
146
+
147
+@dsl_user_op
148
+def ld(row, col, shape: str, num: int, *, loc=None, ip=None):
149
+ nvvm_shape, regs_per_num = LDST_MAP[shape]
150
+ num_regs = regs_per_num * num
151
+ tmem = (Int32(row) << Int32(16)) | Int32(col)
152
+ tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip)
153
+
154
+ if num_regs == 1:
155
+ reg = nvvm.tcgen05_ld(Int32.mlir_type, nvvm_shape, tmem_ptr, loc=loc, ip=ip)
156
+ reg_f32 = llvm.bitcast(Float32.mlir_type, reg, loc=loc, ip=ip)
157
+ return Float32(reg_f32)
158
+
159
+ else:
160
+ vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc)
161
+ vec_f32_ty = ir.VectorType.get([num_regs], Float32.mlir_type, loc=loc)
162
+ regs = nvvm.tcgen05_ld(vec_i32_ty, nvvm_shape, tmem_ptr, loc=loc, ip=ip)
163
+ regs_f32 = llvm.bitcast(vec_f32_ty, regs, loc=loc, ip=ip)
164
+ return cute.TensorSSA(regs_f32, (num_regs,), Float32)
165
+
166
+
167
+@dsl_user_op
168
+def st(row, col, shape: str, num: int, vals, *, loc=None, ip=None) -> None:
169
+ # if input is TensorSSA, convert to Tensor so we can bitcast
170
+ if isinstance(vals, cute.TensorSSA):
171
+ vals_ = cute.make_rmem_tensor_like(vals)
172
+ vals_.store(vals)
173
+ vals = vals_
174
+
175
+ # bitcast to Int32
176
+ vals = cute.recast_tensor(vals, Int32)
177
+
178
+ nvvm_shape, regs_per_num = LDST_MAP[shape]
179
+ num_regs = regs_per_num * num
180
+ tmem = (Int32(row) << Int32(16)) | Int32(col)
181
+ tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip)
182
+
183
+ if num_regs == 1:
184
+ nvvm.tcgen05_st(
185
+ nvvm_shape,
186
+ tmem_ptr,
187
+ vals[0].ir_value(loc=loc, ip=ip),
188
+ loc=loc,
189
+ ip=ip,
190
+ )
191
+ else:
192
+ vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc)
193
+ val_vec = vector.from_elements(
194
+ vec_i32_ty,
195
+ [vals[i].ir_value(loc=loc, ip=ip) for i in range(num_regs)],
196
+ loc=loc,
197
+ ip=ip,
198
+ )
199
+ nvvm.tcgen05_st(nvvm_shape, tmem_ptr, val_vec, loc=loc, ip=ip)
200
+
201
+
202
+@dsl_user_op
203
+def fence_after_thread_sync(*, loc=None, ip=None):
204
+ nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
205
+
206
+
207
+@dsl_user_op
208
+def fence_before_thread_sync(*, loc=None, ip=None):
209
+ nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
210
+
211
+
212
+@dsl_user_op
213
+def wait_ld(*, loc=None, ip=None):
214
+ nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.LOAD, loc=loc, ip=ip)
215
+
216
+
217
+@dsl_user_op
218
+def wait_st(*, loc=None, ip=None):
219
+ nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.STORE, loc=loc, ip=ip)
vllm/cute_utils/cvt.py
renamed
+33
-66
@@ -1,21 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import cutlass
5
-import cutlass.cute as cute
6
-from cutlass import Float32, Uint32
3
+from cutlass import Constexpr, Float32, Uint32, cute
4
from cutlass._mlir import ir
5
from cutlass._mlir.dialects import llvm, vector
6
from cutlass.cutlass_dsl import T, dsl_user_op
7
8
9
@dsl_user_op
13
-def _recast_val(x, dtype, *, loc=None, ip=None):
14
- return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
15
-
16
-
17
-@dsl_user_op
18
-def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
10
+def fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
11
out = llvm.inline_asm(
12
T.i32(),
13
[a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
@@ -28,62 +20,37 @@ def _fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
20
21
22
@dsl_user_op
31
-def _bf16x2_to_fp32(data: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]:
32
- out = llvm.inline_asm(
33
- llvm.StructType.get_literal([T.f32(), T.f32()]),
34
- [data.ir_value(loc=loc, ip=ip)],
35
- "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;\n",
36
- "=f,=f,r",
37
- has_side_effects=False,
38
- is_align_stack=False,
39
- )
40
- return (
41
- Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
42
- Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
43
- )
44
-
45
-
46
-@dsl_user_op
47
-def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
48
- out = llvm.inline_asm(
49
- T.i32(),
50
- [a.ir_value(loc=loc, ip=ip)],
51
- "abs.bf16x2 $0, $1;",
52
- "=r,r",
53
- has_side_effects=False,
54
- is_align_stack=False,
55
- )
56
- return Uint32(out)
57
-
58
-
59
-@dsl_user_op
60
-def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
61
- out = llvm.inline_asm(
62
- T.i32(),
63
- [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
64
- "max.bf16x2 $0, $1, $2;",
65
- "=r,r,r",
66
- has_side_effects=False,
67
- is_align_stack=False,
68
- )
69
- return Uint32(out)
70
-
71
-
72
-@dsl_user_op
73
-def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
74
- out = llvm.inline_asm(
75
- T.i32(),
76
- [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
77
- "mul.rn.bf16x2 $0, $1, $2;",
78
- "=r,r,r",
79
- has_side_effects=False,
80
- is_align_stack=False,
81
- )
82
- return Uint32(out)
23
+def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
24
+ if isinstance(data, Uint32):
25
+ out = llvm.inline_asm(
26
+ llvm.StructType.get_literal([T.f32(), T.f32()]),
27
+ [data.ir_value(loc=loc, ip=ip)],
28
+ "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;",
29
+ "=f,=f,r",
30
+ has_side_effects=False,
31
+ is_align_stack=False,
32
+ loc=loc,
33
+ ip=ip,
34
+ )
35
+ return (
36
+ Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
37
+ Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
38
+ )
39
+
40
+ elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
41
+ # NOTE: the output is always 1D
42
+ size = cute.size(data.shape)
43
+ out = cute.make_rmem_tensor(size * 2, Float32)
44
+ for i in range(size):
45
+ out[i * 2], out[i * 2 + 1] = bf16x2_to_fp32x2(data[i])
46
+ return out
47
+
48
+ else:
49
+ raise ValueError(f"Unsupported type {type(data)}")
50
51
52
@dsl_user_op
86
-def _fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
53
+def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
54
# there is only fp8->fp16 conversion, hence we need to go
55
# round trip through fp16.
56
out = llvm.inline_asm(
@@ -118,7 +85,7 @@ def _fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
85
86
87
@dsl_user_op
121
-def _fp32x4_to_fp8x4(
88
+def fp32x4_to_fp8x4(
89
a0: Float32,
90
a1: Float32,
91
a2: Float32,
@@ -151,9 +118,9 @@ def _fp32x4_to_fp8x4(
118
119
120
@dsl_user_op
154
-def _fp32x8_to_fp4x8(
121
+def fp32x8_to_fp4x8(
122
vals: cute.Tensor,
156
- offset: cutlass.Constexpr[int],
123
+ offset: Constexpr[int],
124
*,
125
loc=None,
126
ip=None,
vllm/engine/arg_utils.py
+2
-2
@@ -712,7 +712,7 @@ class EngineArgs:
712
)
713
714
fail_on_environ_validation: bool = False
715
- gdn_prefill_backend: Literal["flashinfer", "triton"] | None = None
715
+ gdn_prefill_backend: Literal["flashinfer", "triton", "cutedsl"] | None = None
716
717
def __post_init__(self):
718
# support `EngineArgs(compilation_config={...})`
@@ -1527,7 +1527,7 @@ class EngineArgs:
1527
parser.add_argument(
1528
"--gdn-prefill-backend",
1529
dest="gdn_prefill_backend",
1530
- choices=["flashinfer", "triton"],
1530
+ choices=["flashinfer", "triton", "cutedsl"],
1531
default=None,
1532
help="Select GDN prefill backend.",
1533
)
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+152
-49
@@ -3,6 +3,7 @@
3
"""Inference-only Qwen3-Next/Qwen3.5 model."""
4
5
import functools
6
+from typing import Literal
7
8
import torch
9
from einops import rearrange
@@ -83,7 +84,7 @@ logger = init_logger(__name__)
84
85
86
# TODO(arpera): remove ``_is_libs_cu13_install_intact`` and its caller in
86
-# ``_should_use_flashinfer_gdn_prefill`` once the upstream packaging bug is
87
+# ``_resolve_gdn_prefill_backend`` once the upstream packaging bug is
88
# fixed and the broken wheels are yanked / superseded on PyPI:
89
# https://github.com/NVIDIA/cutlass/issues/3170
90
# https://github.com/NVIDIA/cutlass/issues/3259
@@ -146,11 +147,12 @@ def _is_libs_cu13_install_intact() -> bool:
147
return True
148
149
149
-def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) -> bool:
150
- """Whether to use FlashInfer's GDN prefill kernel instead of the
151
- Triton/FLA fallback.
150
+def _resolve_gdn_prefill_backend(
151
+ vllm_config: VllmConfig,
152
+) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]:
153
+ """Resolve GDN prefill backend.
154
153
- Requirements:
155
+ FlashInfer's GDN prefill kernel is chosen when:
156
* ``requested in ["flashinfer", "auto"]``;
157
* ``platform == cuda``;
158
* one of the following:
@@ -158,47 +160,78 @@ def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) ->
160
- Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``,
161
and an intact ``nvidia-cutlass-dsl-libs-cu13`` install on disk
162
(see :func:`_is_libs_cu13_install_intact`).
163
+
164
+ In-tree CuteDSL GDN prefill kernel is chosen when:
165
+ * "cutedsl" is requested; (opt-in only)
166
+ * Blackwell (SM10.x) with ``head_k_dim == 128``;
167
"""
162
- if backend not in ["flashinfer", "auto"]:
163
- return False
168
+ additional_config = vllm_config.additional_config
169
+ backend_cfg = (
170
+ additional_config.get("gdn_prefill_backend", "auto")
171
+ if isinstance(additional_config, dict)
172
+ else "auto"
173
+ )
174
+ backend = str(backend_cfg).strip().lower()
175
+
176
if not current_platform.is_cuda():
165
- return False
177
+ return backend, "triton"
178
+
179
+ head_k_dim = getattr(
180
+ vllm_config.model_config.hf_config, "linear_key_head_dim", None
181
+ )
182
+
183
+ supports_flashinfer = False
184
+ supports_cutedsl = False
185
+
186
if current_platform.is_device_capability(90):
167
- return True # Hopper — no further constraints.
168
- if not current_platform.is_device_capability_family(100):
169
- return False # Neither Hopper nor Blackwell.
170
- if head_k_dim != 128:
171
- return False
172
- if current_platform.get_cuda_runtime_major() < 13:
173
- return False
174
- if not _is_libs_cu13_install_intact():
175
- logger.warning_once(
176
- "FlashInfer Blackwell GDN requires an intact nvidia-cutlass-dsl"
177
- "-libs-cu13 install, but some on-disk files do not match the "
178
- "SHA-256 declared in its RECORD (install-order race in "
179
- "nvidia-cutlass-dsl packaging — see "
180
- "https://github.com/NVIDIA/cutlass/issues/3170 and "
181
- "https://github.com/NVIDIA/cutlass/issues/3259). Falling back "
182
- "to Triton/FLA. Repair with: pip install --force-reinstall "
183
- "--no-deps nvidia-cutlass-dsl-libs-cu13"
184
- )
185
- return False
186
- return True
187
+ supports_flashinfer = True
188
+ elif (
189
+ current_platform.is_device_capability_family(100)
190
+ and head_k_dim == 128
191
+ and current_platform.get_cuda_runtime_major() >= 13
192
+ ):
193
+ supports_flashinfer = _is_libs_cu13_install_intact()
194
+ supports_cutedsl = True
195
+ if not supports_flashinfer:
196
+ logger.warning_once(
197
+ "FlashInfer Blackwell GDN requires an intact nvidia-cutlass-dsl"
198
+ "-libs-cu13 install, but some on-disk files do not match the "
199
+ "SHA-256 declared in its RECORD (install-order race in "
200
+ "nvidia-cutlass-dsl packaging -- see "
201
+ "https://github.com/NVIDIA/cutlass/issues/3170 and "
202
+ "https://github.com/NVIDIA/cutlass/issues/3259). Falling back "
203
+ "to Triton/FLA. Repair with: pip install --force-reinstall "
204
+ "--no-deps nvidia-cutlass-dsl-libs-cu13"
205
+ )
206
+
207
+ if backend in ["flashinfer", "auto"] and supports_flashinfer:
208
+ return backend, "flashinfer"
209
+ if backend == "cutedsl" and supports_cutedsl:
210
+ return backend, "cutedsl"
211
+ return backend, "triton"
212
213
214
def _log_gdn_backend_decision(
190
- backend: str, head_k_dim: int | None, use_flashinfer: bool
215
+ vllm_config: VllmConfig,
216
+ requested_backend: str,
217
+ active_backend: str,
218
) -> None:
219
"""Log the GDN prefill backend choice in the attention-selector style."""
193
- chosen = "FlashInfer" if use_flashinfer else "Triton/FLA"
220
+ head_k_dim = getattr(
221
+ vllm_config.model_config.hf_config, "linear_key_head_dim", None
222
+ )
223
+ chosen = {
224
+ "flashinfer": "FlashInfer",
225
+ "cutedsl": "CuteDSL",
226
+ "triton": "Triton/FLA",
227
+ }[active_backend]
228
logger.info_once(
229
"Using %s GDN prefill kernel (requested=%s, head_k_dim=%s).",
230
chosen,
197
- backend,
231
+ requested_backend,
232
head_k_dim,
233
)
200
- # JIT-compiled cutlass path is only used on SM90 (Hopper).
201
- if use_flashinfer and current_platform.is_device_capability(90):
234
+ if active_backend == "flashinfer" and current_platform.is_device_capability(90):
235
logger.warning_once(
236
"FlashInfer GDN prefill is JIT-compiled; first run may take a "
237
"while. Set --gdn-prefill-backend triton to skip JIT.",
@@ -256,25 +289,26 @@ def fi_chunk_gated_delta_rule(
289
290
@CustomOp.register("chunk_gated_delta_rule")
291
class ChunkGatedDeltaRule(CustomOp):
259
- def __init__(self, head_k_dim: int | None = None) -> None:
292
+ def __init__(self) -> None:
293
super().__init__()
261
- additional_config = get_current_vllm_config().additional_config
262
- assert isinstance(additional_config, dict)
263
- backend_cfg = additional_config.get("gdn_prefill_backend", "auto")
264
- backend = str(backend_cfg).strip().lower()
294
+ vllm_config = get_current_vllm_config()
295
+ backend, active_backend = _resolve_gdn_prefill_backend(vllm_config)
296
+ self.gdn_prefill_backend = active_backend
297
266
- use_flashinfer = _should_use_flashinfer_gdn_prefill(backend, head_k_dim)
267
- if backend == "flashinfer" and not use_flashinfer:
298
+ if backend in ("flashinfer", "cutedsl") and active_backend != backend:
299
logger.warning_once(
269
- "GDN prefill backend 'flashinfer' is selected but "
270
- "cannot use this kernel on the current platform. "
271
- "Falling back to Triton/FLA."
300
+ "GDN prefill backend '%s' is selected but cannot use this "
301
+ "kernel on the current platform. Falling back to Triton/FLA.",
302
+ backend,
303
)
273
- _log_gdn_backend_decision(backend, head_k_dim, use_flashinfer)
304
+ _log_gdn_backend_decision(vllm_config, backend, active_backend)
305
275
- self._forward_method = (
276
- self.forward_cuda if use_flashinfer else self.forward_native
277
- )
306
+ if active_backend == "flashinfer":
307
+ self._forward_method = self.forward_cuda
308
+ elif active_backend == "cutedsl":
309
+ self._forward_method = self.forward_cutedsl
310
+ else:
311
+ self._forward_method = self.forward_native
312
313
def forward_cuda(
314
self,
@@ -338,6 +372,49 @@ class ChunkGatedDeltaRule(CustomOp):
372
core_attn_out=core_attn_out,
373
)
374
375
+ def forward_cutedsl(
376
+ self,
377
+ q: torch.Tensor,
378
+ k: torch.Tensor,
379
+ v: torch.Tensor,
380
+ g: torch.Tensor,
381
+ beta: torch.Tensor,
382
+ initial_state: torch.Tensor,
383
+ output_final_state: bool,
384
+ cu_seqlens: torch.Tensor | None = None,
385
+ chunk_indices: torch.Tensor | None = None,
386
+ chunk_offsets: torch.Tensor | None = None,
387
+ use_qk_l2norm_in_kernel: bool = True,
388
+ core_attn_out: torch.Tensor | None = None,
389
+ ):
390
+ from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
391
+ chunk_gated_delta_rule_cutedsl,
392
+ )
393
+
394
+ if use_qk_l2norm_in_kernel:
395
+ q = l2norm_fwd(q)
396
+ k = l2norm_fwd(k)
397
+
398
+ assert cu_seqlens is not None
399
+ assert chunk_indices is not None
400
+ assert chunk_offsets is not None
401
+
402
+ o, final_state = chunk_gated_delta_rule_cutedsl(
403
+ q=q,
404
+ k=k,
405
+ v=v,
406
+ g=g,
407
+ beta=beta,
408
+ initial_state=initial_state,
409
+ cu_seqlens=cu_seqlens,
410
+ chunk_indices=chunk_indices,
411
+ chunk_offsets=chunk_offsets,
412
+ core_attn_out=core_attn_out,
413
+ )
414
+ if not output_final_state:
415
+ final_state = None
416
+ return o, final_state
417
+
418
419
@PluggableLayer.register("qwen_gated_delta_net_attention")
420
class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
@@ -474,7 +551,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
551
prefix=f"{prefix}.out_proj",
552
)
553
477
- self.chunk_gated_delta_rule = ChunkGatedDeltaRule(head_k_dim=self.head_k_dim)
554
+ self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
555
+ self.gdn_prefill_backend = self.chunk_gated_delta_rule.gdn_prefill_backend
556
self._prefill_kernels_warmed_up = False
557
self.enable_packed_recurrent_decode = (
558
envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
@@ -1060,6 +1138,16 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1138
)
1139
cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)
1140
1141
+ # CuteDSL kernels require metadata
1142
+ chunk_indices = None
1143
+ chunk_offsets = None
1144
+ if self.gdn_prefill_backend == "cutedsl":
1145
+ from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
1146
+ prepare_metadata_cutedsl,
1147
+ )
1148
+
1149
+ chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, T)
1150
+
1151
try:
1152
self.chunk_gated_delta_rule(
1153
q=q,
@@ -1070,6 +1158,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1158
initial_state=state,
1159
output_final_state=True,
1160
cu_seqlens=cu_seqlens,
1161
+ chunk_indices=chunk_indices,
1162
+ chunk_offsets=chunk_offsets,
1163
use_qk_l2norm_in_kernel=False,
1164
)
1165
except Exception:
@@ -1088,7 +1178,20 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1178
self.prefix,
1179
)
1180
finally:
1091
- del dummy_mixed_qkv, q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
1181
+ del (
1182
+ dummy_mixed_qkv,
1183
+ q,
1184
+ k,
1185
+ v,
1186
+ dummy_a,
1187
+ dummy_b,
1188
+ g,
1189
+ beta,
1190
+ state,
1191
+ cu_seqlens,
1192
+ chunk_indices,
1193
+ chunk_offsets,
1194
+ )
1195
1196
torch.accelerator.empty_cache()
1197
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/__init__.py
new
+251
@@ -0,0 +1,251 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from functools import cache
5
+
6
+import cutlass
7
+import torch
8
+from cuda.bindings.driver import CUstream
9
+from cutlass import Int32, cute
10
+from quack.compile_utils import make_fake_tensor
11
+
12
+from vllm.triton_utils import triton
13
+
14
+from .kernel_h import h_cutedsl
15
+from .kernel_kkt_inv_uw import kkt_inv_uw_cutedsl
16
+from .kernel_o import o_cutedsl
17
+
18
+
19
+class PrepMetaKernel:
20
+ def __init__(self, BT: int) -> None:
21
+ self.BT = BT
22
+ self.num_warps = 8
23
+
24
+ @cute.jit
25
+ def __call__(
26
+ self,
27
+ cu_seqlens: cute.Tensor,
28
+ chunk_indices: cute.Tensor,
29
+ chunk_offsets: cute.Tensor,
30
+ stream: CUstream,
31
+ ):
32
+ block = (self.num_warps * 32, 1, 1)
33
+ self.kernel(
34
+ cu_seqlens,
35
+ chunk_indices,
36
+ chunk_offsets,
37
+ ).launch(grid=(1, 1, 1), block=block, stream=stream)
38
+
39
+ @cute.kernel
40
+ def kernel(
41
+ self,
42
+ cu_seqlens: cute.Tensor,
43
+ chunk_indices: cute.Tensor,
44
+ chunk_offsets: cute.Tensor,
45
+ ):
46
+ tid, _, _ = cute.arch.thread_idx()
47
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
48
+ lane_id = tid % 32
49
+
50
+ num_seqs = cu_seqlens.shape[0] - 1
51
+ num_warps = self.num_warps
52
+ tb_size = num_warps * 32
53
+
54
+ if tid == 0:
55
+ chunk_offsets[0] = 0
56
+
57
+ coarsen = cute.ceil_div(num_seqs, tb_size)
58
+ seq_start = tid * coarsen
59
+ num_iters = cutlass.min(seq_start + coarsen, num_seqs) - seq_start
60
+
61
+ # First pass: compute this thread's total chunk count.
62
+ thread_sum = Int32(0)
63
+ for i in range(num_iters):
64
+ seq_id = seq_start + i
65
+ seqlen = cu_seqlens[seq_id + 1] - cu_seqlens[seq_id]
66
+ thread_sum += cute.ceil_div(seqlen, self.BT)
67
+
68
+ # warp parallel scan
69
+ cu_num_chunks = thread_sum
70
+ for i in cutlass.range_constexpr(5):
71
+ offset = cutlass.const_expr(1 << i)
72
+ lower = cute.arch.shuffle_sync_up(
73
+ cu_num_chunks, offset=offset, mask_and_clamp=0
74
+ )
75
+ if lane_id >= offset:
76
+ cu_num_chunks += lower
77
+
78
+ # cross-warp cumsum (CTA-wide)
79
+ smem = cutlass.utils.SmemAllocator()
80
+ warp_num_chunks = smem.allocate_array(Int32, num_warps)
81
+ if lane_id == 31:
82
+ warp_num_chunks[warp_id] = cu_num_chunks
83
+ cute.arch.sync_threads()
84
+
85
+ for i in cutlass.range_constexpr(1, num_warps):
86
+ if warp_id >= i:
87
+ cu_num_chunks += warp_num_chunks[i - 1]
88
+
89
+ chunk_start = cu_num_chunks - thread_sum
90
+
91
+ # Second pass: recompute per-sequence chunk counts and write results.
92
+ for i in range(num_iters):
93
+ seq_id = seq_start + i
94
+ seqlen = cu_seqlens[seq_id + 1] - cu_seqlens[seq_id]
95
+ num_chunks = cute.ceil_div(seqlen, self.BT)
96
+ chunk_end = chunk_start + num_chunks
97
+ chunk_offsets[seq_id + 1] = chunk_end
98
+
99
+ for chunk_id in range(num_chunks):
100
+ chunk_indices[chunk_start + chunk_id, 0] = seq_id
101
+ chunk_indices[chunk_start + chunk_id, 1] = chunk_id
102
+
103
+ chunk_start = chunk_end
104
+
105
+ @cache
106
+ @staticmethod
107
+ def compile(BT: int):
108
+ cu_entries = cute.sym_int()
109
+ upper_bound_chunks = cute.sym_int()
110
+
111
+ cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
112
+ chunk_indices = make_fake_tensor(Int32, (upper_bound_chunks, 2), divisibility=2)
113
+ chunk_offsets = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
114
+
115
+ kernel = PrepMetaKernel(BT)
116
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
117
+ return cute.compile(
118
+ kernel,
119
+ cu_seqlens,
120
+ chunk_indices,
121
+ chunk_offsets,
122
+ stream,
123
+ options="--enable-tvm-ffi",
124
+ )
125
+
126
+
127
+def _upper_bound_chunks(num_seqs: int, total_tokens: int, chunk_size: int) -> int:
128
+ return (num_seqs - 1) + triton.cdiv(total_tokens - (num_seqs - 1), chunk_size)
129
+
130
+
131
+def prepare_metadata_cutedsl(
132
+ cu_seqlens: torch.Tensor,
133
+ total_tokens: int,
134
+ chunk_size: int = 64,
135
+) -> tuple[torch.Tensor, torch.Tensor]:
136
+ num_seqs = cu_seqlens.numel() - 1
137
+ upper_bound_chunks = _upper_bound_chunks(num_seqs, total_tokens, chunk_size)
138
+ chunk_offsets = cu_seqlens.new_empty(num_seqs + 1, dtype=torch.int32)
139
+ chunk_indices = cu_seqlens.new_empty((upper_bound_chunks, 2), dtype=torch.int32)
140
+
141
+ PrepMetaKernel.compile(chunk_size)(cu_seqlens, chunk_indices, chunk_offsets)
142
+ return chunk_indices, chunk_offsets
143
+
144
+
145
+def chunk_gated_delta_rule_cutedsl(
146
+ q: torch.Tensor,
147
+ k: torch.Tensor,
148
+ v: torch.Tensor,
149
+ g: torch.Tensor,
150
+ beta: torch.Tensor,
151
+ initial_state: torch.Tensor,
152
+ cu_seqlens: torch.Tensor,
153
+ chunk_indices: torch.Tensor,
154
+ chunk_offsets: torch.Tensor,
155
+ core_attn_out: torch.Tensor | None = None,
156
+) -> tuple[torch.Tensor, torch.Tensor]:
157
+ """Run the GDN chunk CuteDSL prefill kernels.
158
+
159
+ Args:
160
+ q: Query tensor with shape ``[1, T, H, K]``.
161
+ k: Key tensor with shape ``[1, T, H, K]``.
162
+ v: Value tensor with shape ``[1, T, Hv, V]``.
163
+ g: Log-space decay tensor with shape ``[1, T, Hv]``.
164
+ beta: Delta-rule beta tensor with shape ``[1, T, Hv]``.
165
+ initial_state: Recurrent state with shape ``[N, Hv, V, K]``.
166
+ cu_seqlens: Cumulative sequence lengths with shape ``[N + 1]``.
167
+ chunk_indices: Chunk index metadata with shape ``[NT, 2]``.
168
+ chunk_offsets: Cumulative chunk offsets with shape ``[N + 1]``.
169
+ core_attn_out: Optional output buffer with shape ``[T, Hv, V]``.
170
+
171
+ Returns:
172
+ A tuple ``(output, final_state)`` where ``output`` has shape
173
+ ``[1, T, Hv, V]`` and ``final_state`` has shape ``[N, Hv, V, K]``.
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
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)
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,
198
+ u,
199
+ w,
200
+ g_2d,
201
+ beta_2d,
202
+ g_cu,
203
+ cu_seqlens,
204
+ chunk_indices,
205
+ total_chunks_ptr,
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)
216
+ final_state = torch.empty_like(initial_state)
217
+ h_cutedsl(
218
+ k_3d,
219
+ u,
220
+ w,
221
+ v_new,
222
+ g_cu,
223
+ h,
224
+ initial_state,
225
+ final_state,
226
+ cu_seqlens,
227
+ chunk_offsets,
228
+ )
229
+
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
232
+ o_cutedsl(
233
+ q_3d,
234
+ k_3d,
235
+ v_new.view(upper_bound_chunks, chunk_size, num_v_heads, head_v_dim),
236
+ h,
237
+ g_cu,
238
+ output,
239
+ cu_seqlens,
240
+ chunk_indices,
241
+ total_chunks_ptr,
242
+ scale,
243
+ num_sms=num_sms,
244
+ )
245
+ return output.unsqueeze(0), final_state
246
+
247
+
248
+__all__ = [
249
+ "chunk_gated_delta_rule_cutedsl",
250
+ "prepare_metadata_cutedsl",
251
+]
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_h.py
new
+753
@@ -0,0 +1,753 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from functools import cache
4
+
5
+import cutlass
6
+import torch
7
+from cuda.bindings.driver import CUstream
8
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
9
+from cutlass.cute.nvgpu import cpasync, warp
10
+from quack.compile_utils import make_fake_tensor
11
+
12
+from vllm.cute_utils import (
13
+ EVICT_FIRST,
14
+ _tcgen05,
15
+ cvt,
16
+ fence_before_tma_store,
17
+ simple_tma_copy,
18
+)
19
+
20
+
21
+class Sm100ChunkHKernel:
22
+ """For each sequence, compute the chunk recurrent update.
23
+
24
+ The input V tile is the U output from the KKT/UW kernel. For each chunk:
25
+ V_new = U - W @ H.T
26
+ (we actually do V_new.T = U.T - H @ W.T instead)
27
+
28
+ H_scaled = H * exp(g_last)
29
+ V_scaled = V_new * exp(g_last - g)
30
+ H_new = H_scaled + V_scaled.T @ K
31
+ """
32
+
33
+ def __init__(
34
+ self,
35
+ H: int,
36
+ Hv: int,
37
+ K_dim: int,
38
+ V_dim: int,
39
+ h_dtype: cutlass.Numeric = Float32,
40
+ BT: int = 64,
41
+ num_stages: int = 2,
42
+ ) -> None:
43
+ assert Hv % H == 0
44
+ assert K_dim == V_dim == 128
45
+ assert BT == 64
46
+ self.H = H
47
+ self.Hv = Hv
48
+ self.K_dim = K_dim
49
+ self.V_dim = V_dim
50
+ self.h_dtype = h_dtype
51
+ self.BT = BT
52
+ self.num_stages = num_stages
53
+ self.num_warps = 10
54
+
55
+ @cute.jit
56
+ def _make_bf16_tma_args(
57
+ self,
58
+ tensor: cute.Tensor,
59
+ dim: cutlass.Constexpr[int],
60
+ op: cpasync.TmaCopyOp,
61
+ stages: cutlass.Constexpr[int],
62
+ ):
63
+ swizzle_128B = cute.make_swizzle(3, 4, 3)
64
+ slayout = cute.make_layout(
65
+ (self.BT, 1, (64, dim // 64), stages),
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(
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
76
+
77
+ @cute.jit
78
+ def _make_h_tma_args(self, tensor: cute.Tensor, op: cpasync.TmaCopyOp):
79
+ # number of elements to fill 128B
80
+ num_elems = 128 // (tensor.element_type.width // 8)
81
+ swizzle_128B = cute.make_swizzle(3, 4, 3)
82
+ slayout = cute.make_layout(
83
+ (1, 1, self.V_dim, (num_elems, self.K_dim // num_elems)),
84
+ stride=(0, 0, num_elems, (1, self.V_dim * num_elems)),
85
+ )
86
+ slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
87
+ atom, tma_tensor = cpasync.make_tiled_tma_atom(
88
+ op,
89
+ cute.logical_divide(tensor, (None, None, None, num_elems)),
90
+ slayout,
91
+ cta_tiler=(1, 1, self.V_dim, self.K_dim),
92
+ )
93
+ return atom, tma_tensor, slayout
94
+
95
+ @cute.jit
96
+ def __call__(
97
+ self,
98
+ K: cute.Tensor,
99
+ V: cute.Tensor,
100
+ W: cute.Tensor,
101
+ V_new: cute.Tensor,
102
+ g_cu: cute.Tensor,
103
+ h: cute.Tensor,
104
+ h0: cute.Tensor,
105
+ ht: cute.Tensor,
106
+ cu_seqlens: cute.Tensor,
107
+ chunk_offsets: cute.Tensor,
108
+ stream: CUstream,
109
+ ):
110
+ tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
111
+ tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
112
+
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)
120
+
121
+ grid = (self.Hv, h0.shape[0], 1)
122
+ block = (self.num_warps * 32, 1, 1)
123
+ self.kernel(
124
+ K_args,
125
+ V_args,
126
+ W_args,
127
+ V_new_args,
128
+ H0_args,
129
+ HT_args,
130
+ H_args,
131
+ g_cu,
132
+ cu_seqlens,
133
+ chunk_offsets,
134
+ ).launch(grid=grid, block=block, stream=stream)
135
+
136
+ @cute.kernel
137
+ def kernel(
138
+ 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],
146
+ g_cu: cute.Tensor,
147
+ cu_seqlens: cute.Tensor,
148
+ chunk_offsets: cute.Tensor,
149
+ ):
150
+ tid, _, _ = cute.arch.thread_idx()
151
+ head_id, seq_id, _ = cute.arch.block_idx()
152
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
153
+ lane_id = tid % 32
154
+
155
+ BT = self.BT
156
+ V_dim = self.V_dim
157
+ K_dim = self.K_dim
158
+ num_stages = self.num_stages
159
+ is_f32 = self.h_dtype == Float32
160
+
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
+
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()
175
+
176
+ # 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]
183
+
184
+ s_v_scale = smem.allocate_array(Float32, BT)
185
+ tma_mbar = smem.allocate_array(Int64, num_stages)
186
+ wh_in_mbar = smem.allocate_array(Int64, num_stages)
187
+ wh_done_mbar = smem.allocate_array(Int64, num_stages)
188
+ vk_in_mbar = smem.allocate_array(Int64, num_stages)
189
+ vk_done_mbar = smem.allocate_array(Int64, num_stages)
190
+ h0_mbar = smem.allocate_array(Int64, 1)
191
+ taddr = smem.allocate(Int32, 4)
192
+
193
+ wh_tmem = 0
194
+ vk_tmem = wh_tmem + BT
195
+ h_tmem_base = vk_tmem + K_dim
196
+ v_tmem_base = h_tmem_base + K_dim // 2
197
+
198
+ if warp_id == 0:
199
+ with cute.arch.elect_one():
200
+ for i in cutlass.range_constexpr(num_stages):
201
+ cute.arch.mbarrier_init(tma_mbar + i, 1)
202
+ cute.arch.mbarrier_init(wh_in_mbar + i, 256)
203
+ cute.arch.mbarrier_init(wh_done_mbar + i, 1)
204
+ cute.arch.mbarrier_init(vk_in_mbar + i, 256)
205
+ cute.arch.mbarrier_init(vk_done_mbar + i, 1)
206
+ cute.arch.mbarrier_init(h0_mbar, 1)
207
+ cute.arch.mbarrier_init_fence()
208
+ 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)
216
+ cute.arch.sync_threads()
217
+
218
+ bos = cu_seqlens[seq_id]
219
+ eos = cu_seqlens[seq_id + 1]
220
+ seqlen = eos - bos
221
+ num_chunks = cute.ceil_div(seqlen, BT)
222
+
223
+ if warp_id == 9:
224
+ # TMA warp
225
+ stage_id = 0
226
+ parity = 1
227
+
228
+ k_head_id = head_id // (self.Hv // self.H)
229
+ chunk_offset = chunk_offsets[seq_id]
230
+
231
+ # load H0
232
+ with cute.arch.elect_one():
233
+ H0_size = V_dim * K_dim * self.h_dtype.width // 8
234
+ cute.arch.mbarrier_arrive_and_expect_tx(h0_mbar, H0_size)
235
+ simple_tma_copy(
236
+ H0_tma_atom, tmaH0[seq_id, head_id, None, None], sH0, h0_mbar
237
+ )
238
+
239
+ # 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))
242
+ gK_tiles = cute.logical_divide(
243
+ cute.domain_offset((bos, 0), tmaK[None, k_head_id, None]),
244
+ (BT, None),
245
+ )
246
+
247
+ for chunk_id in range(num_chunks):
248
+ mbar = tma_mbar + stage_id
249
+ gW = gW_tiles[(None, chunk_offset + chunk_id), None]
250
+ gV = gV_tiles[(None, chunk_offset + chunk_id), None]
251
+ gK = gK_tiles[(None, chunk_id), None]
252
+
253
+ # wait for MMA to release the buffer
254
+ cute.arch.mbarrier_wait(vk_done_mbar + stage_id, parity)
255
+
256
+ # load W, V (i.e. U), and K
257
+ with cute.arch.elect_one():
258
+ STAGE_SIZE = BT * (K_dim + V_dim + K_dim) * 2
259
+ cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
260
+ simple_tma_copy(
261
+ W_tma_atom, gW, sW[None, None, stage_id], mbar, EVICT_FIRST
262
+ )
263
+ simple_tma_copy(
264
+ V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
265
+ )
266
+ simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
267
+
268
+ stage_id = (stage_id + 1) % num_stages
269
+ if stage_id == 0:
270
+ parity ^= 1
271
+
272
+ elif warp_id == 8:
273
+ # MMA warp
274
+ _tcgen05.alloc(taddr)
275
+ stage_id = 0
276
+ parity = 0
277
+
278
+ wh_idesc = _tcgen05.make_bf16_idesc(V_dim, BT, negate_A=True)
279
+ vk_idesc = _tcgen05.make_bf16_idesc(V_dim, K_dim, transpose_B=True)
280
+
281
+ # LBO=BT*128 is ignored for K-major
282
+ sdesc_template = _tcgen05.make_sdesc_128B_swizzle(BT * 128)
283
+
284
+ # when using BF16 state, H is read from smem for the 1st iteration
285
+ # variable names in this conditional branch can't be the same as those
286
+ # in the mainloop below due to CuteDSL restrictions.
287
+ if cutlass.const_expr(not is_f32):
288
+ ##### 1st MMA: V_new.T = V.T - H @ W.T #####
289
+ Haddr0 = sH0[None, None].iterator.toint()
290
+ Waddr0 = sW[None, None, stage_id].iterator.toint()
291
+ hdesc0_base = sdesc_template | (Haddr0 >> 4)
292
+ wdesc0_base = sdesc_template | (Waddr0 >> 4)
293
+
294
+ cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
295
+ cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
296
+ _tcgen05.fence_after_thread_sync()
297
+
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)
305
+
306
+ ##### 2nd MMA: H_new = H + V_new.T @ K #####
307
+ Kaddr0 = sK[None, None, stage_id].iterator.toint()
308
+ kdesc0_base = sdesc_template | (Kaddr0 >> 4)
309
+
310
+ cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
311
+ _tcgen05.fence_after_thread_sync()
312
+
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)
319
+
320
+ stage_id = (stage_id + 1) % num_stages
321
+ if stage_id == 0:
322
+ parity ^= 1
323
+
324
+ num_iters = num_chunks - int(not is_f32)
325
+ for _ in range(num_iters):
326
+ ##### 1st MMA: V_new.T = V.T - H @ W.T #####
327
+ Waddr = sW[None, None, stage_id].iterator.toint()
328
+ wdesc_base = sdesc_template | (Waddr >> 4)
329
+
330
+ cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
331
+ cute.arch.mbarrier_wait(wh_in_mbar + stage_id, parity)
332
+ _tcgen05.fence_after_thread_sync()
333
+
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)
341
+
342
+ ##### 2nd MMA: H_new = H + V_new.T @ K #####
343
+ Kaddr = sK[None, None, stage_id].iterator.toint()
344
+ kdesc_base = sdesc_template | (Kaddr >> 4)
345
+
346
+ cute.arch.mbarrier_wait(vk_in_mbar + stage_id, parity)
347
+ _tcgen05.fence_after_thread_sync()
348
+
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)
355
+
356
+ stage_id = (stage_id + 1) % num_stages
357
+ if stage_id == 0:
358
+ parity ^= 1
359
+
360
+ elif warp_id >= 4:
361
+ # H warps
362
+ tid_ = tid % 128
363
+ warp_id_ = warp_id % 4
364
+ chunk_offset = chunk_offsets[seq_id]
365
+
366
+ stage_id = 0
367
+ vk_stage_id = 0
368
+ vk_parity = 0
369
+
370
+ op = cute.nvgpu.CopyUniversalOp()
371
+ cp_16B = cute.make_copy_atom(op, Float32, num_bits_per_copy=128)
372
+
373
+ ##### chunk_id = 0 #####
374
+ if True:
375
+ chunk_id = 0
376
+ end_t = min(bos + (chunk_id + 1) * BT, eos)
377
+ last_idx = end_t - 1
378
+ h_scale = cute.math.exp(g_cu[last_idx, head_id], fastmath=True)
379
+
380
+ # for 1st chunk, wait for H0 transfer from gmem
381
+ if warp_id_ == 0:
382
+ cute.arch.mbarrier_wait(h0_mbar, 0)
383
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
384
+
385
+ # when H0 is FP32, we need to pack it to BF16
386
+ # also store to smem for TMA store later.
387
+ if cutlass.const_expr(is_f32):
388
+ for i in cutlass.range_constexpr(K_dim // 32):
389
+ # H0 smem layout: (V_dim, (32, K_dim/32))
390
+ h_f32 = cute.make_rmem_tensor(32, Float32)
391
+ cute.copy(cp_16B, sH0[tid_, (None, i)], h_f32)
392
+
393
+ h_bf16 = cute.make_rmem_tensor(32, BFloat16)
394
+ h_bf16.store(h_f32.load().to(BFloat16))
395
+ _tcgen05.st(
396
+ warp_id_ * 32, h_tmem_base + i * 16, "32x32b", 16, h_bf16
397
+ )
398
+
399
+ # H smem layout: (V_dim, (64, K_dim/64))
400
+ dst = cute.local_tile(sH[tid_, None], (32,), (i,))
401
+ cute.copy(cp_16B, h_bf16, dst)
402
+
403
+ _tcgen05.wait_st()
404
+ _tcgen05.fence_before_thread_sync()
405
+ cute.arch.mbarrier_arrive(wh_in_mbar + stage_id)
406
+
407
+ # scale H for 2nd MMA
408
+ for i in cutlass.range_constexpr(K_dim // 32):
409
+ h_f32 = cute.make_rmem_tensor(32, Float32)
410
+
411
+ if cutlass.const_expr(is_f32):
412
+ cute.copy(cp_16B, sH0[tid_, (None, i)], h_f32)
413
+
414
+ else:
415
+ h_bf16 = cute.make_rmem_tensor(32, BFloat16)
416
+ sH_src = cute.local_tile(sH0[tid_, None], (32,), (i,))
417
+ cute.copy(cp_16B, sH_src, h_bf16)
418
+ h_f32.store(
419
+ cvt.bf16x2_to_fp32x2(
420
+ cute.recast_tensor(h_bf16, Uint32)
421
+ ).load()
422
+ )
423
+
424
+ for j in cutlass.range_constexpr(32):
425
+ h_f32[j] *= h_scale
426
+ _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
427
+
428
+ _tcgen05.wait_st()
429
+ _tcgen05.fence_before_thread_sync()
430
+ cute.arch.mbarrier_arrive(vk_in_mbar + stage_id)
431
+
432
+ # for BF16 H0, we issue TMA store from H0 smem
433
+ # for FP32 H0, we issue TMA store from H smem (after packing)
434
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
435
+ fence_before_tma_store()
436
+ if warp_id_ == 3:
437
+ 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)
440
+ with cute.arch.elect_one():
441
+ cute.arch.cp_async_bulk_commit_group()
442
+
443
+ # When H0 is BF16, and there is only 1 chunk, storing
444
+ # the final state to sH0 can race before this store
445
+ # has finished. hence, we need to wait here.
446
+ if cutlass.const_expr(not is_f32):
447
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
448
+
449
+ stage_id = (stage_id + 1) % num_stages
450
+
451
+ ##### subsequent chunks #####
452
+ for chunk_id in range(1, num_chunks):
453
+ end_t = min(bos + (chunk_id + 1) * BT, eos)
454
+ last_idx = end_t - 1
455
+ h_scale = cute.math.exp(g_cu[last_idx, head_id], fastmath=True)
456
+
457
+ # wait for H from previous vk MMA
458
+ if warp_id_ == 0:
459
+ cute.arch.mbarrier_wait(vk_done_mbar + vk_stage_id, vk_parity)
460
+ vk_stage_id = (vk_stage_id + 1) % num_stages
461
+ if vk_stage_id == 0:
462
+ vk_parity ^= 1
463
+ elif warp_id_ == 3:
464
+ with cute.arch.elect_one():
465
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
466
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
467
+ _tcgen05.fence_after_thread_sync()
468
+
469
+ # load FP32 H from tmem, convert to BF16, store to tmem for 1st MMA,
470
+ # store to smem for TMA store later.
471
+ for i in cutlass.range_constexpr(K_dim // 32):
472
+ h_f32 = _tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32)
473
+ h_bf16 = cute.make_rmem_tensor(32, BFloat16)
474
+ h_bf16.store(h_f32.to(BFloat16))
475
+ _tcgen05.st(
476
+ warp_id_ * 32, h_tmem_base + i * 16, "32x32b", 16, h_bf16
477
+ )
478
+
479
+ # H smem layout: (V_dim, (64, K_dim/64))
480
+ dst = cute.local_tile(sH[tid_, None], (32,), (i,))
481
+ cute.copy(cp_16B, h_bf16, dst)
482
+
483
+ _tcgen05.wait_st()
484
+ _tcgen05.fence_before_thread_sync()
485
+ cute.arch.mbarrier_arrive(wh_in_mbar + stage_id)
486
+
487
+ # scale H for 2nd MMA
488
+ for i in cutlass.range_constexpr(K_dim // 32):
489
+ h_f32 = cute.make_rmem_tensor(32, Float32)
490
+ h_f32.store(
491
+ _tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32)
492
+ )
493
+ for j in cutlass.range_constexpr(32):
494
+ h_f32[j] *= h_scale
495
+ _tcgen05.st(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32, h_f32)
496
+ _tcgen05.wait_st()
497
+ _tcgen05.fence_before_thread_sync()
498
+ cute.arch.mbarrier_arrive(vk_in_mbar + stage_id)
499
+
500
+ # issue TMA store for O kernel
501
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
502
+ fence_before_tma_store()
503
+ 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)
506
+ with cute.arch.elect_one():
507
+ cute.arch.cp_async_bulk_commit_group()
508
+
509
+ stage_id = (stage_id + 1) % num_stages
510
+
511
+ # handle final state. reuse H0 smem.
512
+ if warp_id_ == 0:
513
+ cute.arch.mbarrier_wait(vk_done_mbar + vk_stage_id, vk_parity)
514
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
515
+ _tcgen05.fence_after_thread_sync()
516
+
517
+ for i in cutlass.range_constexpr(K_dim // 32):
518
+ h_f32 = cute.make_rmem_tensor(32, Float32)
519
+ h_f32.store(_tcgen05.ld(warp_id_ * 32, vk_tmem + i * 32, "32x32b", 32))
520
+
521
+ if cutlass.const_expr(is_f32):
522
+ cute.copy(cp_16B, h_f32, sH0[tid_, (None, i)])
523
+
524
+ else:
525
+ h_bf16 = cute.make_rmem_tensor(32, BFloat16)
526
+ h_bf16.store(h_f32.load().to(BFloat16))
527
+ sH0_dst = cute.local_tile(sH0[tid_, None], (32,), (i,))
528
+ cute.copy(cp_16B, h_bf16, sH0_dst)
529
+
530
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
531
+
532
+ if warp_id_ == 0:
533
+ ht_dst = tmaHT[seq_id, head_id, None, None]
534
+ simple_tma_copy(HT_tma_atom, sH0, ht_dst)
535
+ with cute.arch.elect_one():
536
+ cute.arch.cp_async_bulk_commit_group()
537
+ if warp_id_ == 1:
538
+ _tcgen05.dealloc()
539
+
540
+ else:
541
+ # V warps
542
+ stage_id = 0
543
+ parity = 0
544
+
545
+ chunk_offset = chunk_offsets[seq_id]
546
+
547
+ ldsm_trans_op = warp.LdMatrix8x8x16bOp(num_matrices=4, transpose=True)
548
+ stsm_trans_op = warp.StMatrix8x8x16bOp(num_matrices=4, transpose=True)
549
+ ldsm_trans_atom = cute.make_copy_atom(ldsm_trans_op, BFloat16)
550
+ stsm_trans_atom = cute.make_copy_atom(stsm_trans_op, BFloat16)
551
+
552
+ # ((BT, num_BT_tiles), V_dim)
553
+ gV_new_tiles = cute.logical_divide(
554
+ tmaV_new[None, head_id, None], (BT, None)
555
+ )
556
+
557
+ # sV shape: [BT, (64, V_dim/64), num_stages]
558
+ # sV_view shape: [BT, (8, (8,2)), num_stages]
559
+ sV_view = cute.logical_divide(sV, (None, 8, None))
560
+ sV_new_view = cute.logical_divide(sV_new, (None, 8))
561
+
562
+ # [BT, 8, num_stages]
563
+ s_col = warp_id * 4 + (lane_id // 8)
564
+ sV_view = sV_view[None, (None, s_col), None]
565
+ sV_new_view = sV_new_view[None, (None, s_col)]
566
+
567
+ for chunk_id in range(num_chunks):
568
+ # wait for V to arrive
569
+ if warp_id == 0:
570
+ cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
571
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
572
+
573
+ # unpack V BF16->FP32, then store to tmem for 1st MMA
574
+ # V smem layout: [BT, (64, V_dim/64)] / [BT, V_dim]
575
+ # each iteration, CTA loads [8, V_dim] tile
576
+ # (warp loads [8, 32] tile)
577
+ for i in cutlass.range_constexpr(BT // 8):
578
+ s_row = i * 8 + (lane_id % 8)
579
+ v_bf16 = cute.make_rmem_tensor(8, BFloat16)
580
+ cute.copy(ldsm_trans_atom, sV_view[s_row, None, stage_id], v_bf16)
581
+ v_fp32 = cvt.bf16x2_to_fp32x2(cute.recast_tensor(v_bf16, Uint32))
582
+ v_fp32 = cute.logical_divide(v_fp32, 4) # (4, 2)
583
+
584
+ tcol = wh_tmem + i * 8
585
+ _tcgen05.st(warp_id * 32 + 0, tcol, "16x256b", 1, v_fp32[None, 0])
586
+ _tcgen05.st(warp_id * 32 + 16, tcol, "16x256b", 1, v_fp32[None, 1])
587
+
588
+ _tcgen05.wait_st()
589
+ _tcgen05.fence_before_thread_sync()
590
+ cute.arch.mbarrier_arrive(wh_in_mbar + stage_id)
591
+
592
+ # load g_cu for scaling
593
+ if tid < BT:
594
+ end_t = min(bos + (chunk_id + 1) * BT, eos)
595
+ last_idx = end_t - 1
596
+ t = bos + chunk_id * BT + tid
597
+ val = Float32(0.0)
598
+ if t < eos:
599
+ val = cute.math.exp(
600
+ g_cu[last_idx, head_id] - g_cu[t, head_id],
601
+ fastmath=True,
602
+ )
603
+ s_v_scale[tid] = val
604
+
605
+ # wait for 1st MMA to finish
606
+ if warp_id == 2:
607
+ cute.arch.mbarrier_wait(wh_done_mbar + stage_id, parity)
608
+ elif warp_id == 3:
609
+ with cute.arch.elect_one():
610
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
611
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
612
+ _tcgen05.fence_after_thread_sync()
613
+
614
+ for i in cutlass.range_constexpr(BT // 8):
615
+ v_new = cute.make_rmem_tensor((4, 2), Float32)
616
+ tcol = wh_tmem + i * 8
617
+ v_new[None, 0].store(
618
+ _tcgen05.ld(warp_id * 32 + 0, tcol, "16x256b", 1)
619
+ )
620
+ v_new[None, 1].store(
621
+ _tcgen05.ld(warp_id * 32 + 16, tcol, "16x256b", 1)
622
+ )
623
+ v_new_bf16 = cute.make_rmem_tensor(8, BFloat16)
624
+ v_new_bf16.store(v_new.load().to(BFloat16))
625
+
626
+ # scale V_new for 2nd MMA
627
+ scale0 = s_v_scale[i * 8 + (lane_id % 4) * 2 + 0]
628
+ scale1 = s_v_scale[i * 8 + (lane_id % 4) * 2 + 1]
629
+ v_scaled = cute.make_rmem_tensor(8, Float32)
630
+ for k in cutlass.range_constexpr(4):
631
+ v_scaled[k * 2] = v_new[k * 2] * scale0
632
+ v_scaled[k * 2 + 1] = v_new[k * 2 + 1] * scale1
633
+ v_scaled_bf16 = v_scaled.load().to(BFloat16).reshape((4, 2))
634
+
635
+ # store V_new BF16 for O kernel
636
+ s_row = i * 8 + (lane_id % 8)
637
+ cute.copy(stsm_trans_atom, v_new_bf16, sV_new_view[s_row, None])
638
+
639
+ # store to tmem
640
+ tcol = v_tmem_base + i * 4
641
+ _tcgen05.st(
642
+ warp_id * 32 + 0, tcol, "16x128b", 1, v_scaled_bf16[None, 0]
643
+ )
644
+ _tcgen05.st(
645
+ warp_id * 32 + 16, tcol, "16x128b", 1, v_scaled_bf16[None, 1]
646
+ )
647
+ _tcgen05.wait_st()
648
+ _tcgen05.fence_before_thread_sync()
649
+ cute.arch.mbarrier_arrive(vk_in_mbar + stage_id)
650
+
651
+ # issue TMA store for V_new
652
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
653
+ fence_before_tma_store()
654
+ if warp_id == 3:
655
+ gV = gV_new_tiles[(None, chunk_offset + chunk_id), None]
656
+ simple_tma_copy(V_new_tma_atom, sV_new, gV)
657
+ with cute.arch.elect_one():
658
+ cute.arch.cp_async_bulk_commit_group()
659
+
660
+ stage_id = (stage_id + 1) % num_stages
661
+ if stage_id == 0:
662
+ parity ^= 1
663
+
664
+ @cache
665
+ @staticmethod
666
+ def compile(
667
+ H: int,
668
+ Hv: int,
669
+ K_dim: int,
670
+ V_dim: int,
671
+ h_dtype: cutlass.Numeric = Float32,
672
+ BT: int = 64,
673
+ num_stages: int = 2,
674
+ ):
675
+ total_t = cute.sym_int()
676
+ pad_t = cute.sym_int()
677
+ total_chunks_n = cute.sym_int()
678
+ num_sequences = cute.sym_int()
679
+ cu_entries = cute.sym_int()
680
+
681
+ K = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
682
+ V = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
683
+ W = make_fake_tensor(BFloat16, (pad_t, Hv, K_dim), divisibility=16)
684
+ V_new = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
685
+ g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
686
+ h = make_fake_tensor(
687
+ BFloat16, (total_chunks_n, Hv, V_dim, K_dim), divisibility=16
688
+ )
689
+ h0 = make_fake_tensor(
690
+ h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
691
+ )
692
+ ht = make_fake_tensor(
693
+ h_dtype, (num_sequences, Hv, V_dim, K_dim), divisibility=16
694
+ )
695
+ cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
696
+ chunk_offsets = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
697
+
698
+ kernel = Sm100ChunkHKernel(H, Hv, K_dim, V_dim, h_dtype, BT, num_stages)
699
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
700
+ return cute.compile(
701
+ kernel,
702
+ K,
703
+ V,
704
+ W,
705
+ V_new,
706
+ g_cu,
707
+ h,
708
+ h0,
709
+ ht,
710
+ cu_seqlens,
711
+ chunk_offsets,
712
+ stream,
713
+ options="--enable-tvm-ffi",
714
+ )
715
+
716
+
717
+def h_cutedsl(
718
+ K: torch.Tensor,
719
+ V: torch.Tensor,
720
+ W: torch.Tensor,
721
+ V_new: torch.Tensor,
722
+ g_cu: torch.Tensor,
723
+ h: torch.Tensor,
724
+ h0: torch.Tensor,
725
+ ht: torch.Tensor,
726
+ cu_seqlens: torch.Tensor,
727
+ chunk_offsets: torch.Tensor,
728
+ BT: int = 64,
729
+ num_stages: int = 2,
730
+) -> None:
731
+ """Compute H/V_new with the same argument order as the CUDA wrapper."""
732
+
733
+ _, H, K_dim = K.shape
734
+ _, Hv, V_dim = V.shape
735
+ h_dtype = {
736
+ torch.bfloat16: BFloat16,
737
+ torch.float32: Float32,
738
+ }[h0.dtype]
739
+ Sm100ChunkHKernel.compile(H, Hv, K_dim, V_dim, h_dtype, BT, num_stages)(
740
+ K,
741
+ V,
742
+ W,
743
+ V_new,
744
+ g_cu,
745
+ h,
746
+ h0,
747
+ ht,
748
+ cu_seqlens,
749
+ chunk_offsets,
750
+ )
751
+
752
+
753
+h_v2b_cutedsl = h_cutedsl
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
new
+832
@@ -0,0 +1,832 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from functools import cache
4
+
5
+import cutlass
6
+import torch
7
+from cuda.bindings.driver import CUstream
8
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
9
+from cutlass.cute.nvgpu import cpasync, warp
10
+from quack.compile_utils import make_fake_tensor
11
+
12
+from vllm.cute_utils import (
13
+ EVICT_FIRST,
14
+ _tcgen05,
15
+ cvt,
16
+ fence_before_tma_store,
17
+ mma_bf16,
18
+ simple_tma_copy,
19
+)
20
+
21
+
22
+class Sm100ChunkUWKernel:
23
+ """Compute per-chunk KKT inverse preprocessing and U/W tiles.
24
+
25
+ Gamma[i,j] = exp(g_cu[i] - g_cu[j])
26
+ A = strictLower(beta * (K @ K.T) * Gamma)
27
+ Ai = inverse(I + A)
28
+ U = (Ai * beta) @ V
29
+ W = (Ai * beta * exp(g_cu)) @ K
30
+ """
31
+
32
+ def __init__(
33
+ self,
34
+ H: int,
35
+ Hv: int,
36
+ K_dim: int,
37
+ V_dim: int,
38
+ num_stages: int = 2,
39
+ ) -> None:
40
+ assert Hv % H == 0
41
+ assert K_dim == V_dim == 128
42
+ self.H = H
43
+ self.Hv = Hv
44
+ self.K_dim = K_dim
45
+ self.V_dim = V_dim
46
+ self.num_stages = num_stages
47
+
48
+ # hard-code
49
+ self.BT = 64
50
+ self.num_warps = 2 + 4 + 4
51
+
52
+ @cute.jit
53
+ def _make_tma_args(
54
+ self,
55
+ tensor: cute.Tensor,
56
+ dim: cutlass.Constexpr[int],
57
+ num_stages: int,
58
+ op: cpasync.TmaCopyOp,
59
+ ):
60
+ # logical layout: [BT, dim]
61
+ # permute for TMA: [dim/64, BT, 64] with swizzling
62
+ swizzle_128B = cute.make_swizzle(3, 4, 3)
63
+ slayout = cute.make_layout(
64
+ (self.BT, 1, (64, dim // 64), num_stages),
65
+ stride=(64, 0, (1, self.BT * 64), self.BT * dim),
66
+ )
67
+ slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
68
+
69
+ # we need to convert gmem layout to (T, H, (64, D/64)) for make_tiled_tma_atom()
70
+ # to emit a single 4D TMA. otherwise, it will emit (D/64)x 3D TMA.
71
+ atom, tma_tensor = cpasync.make_tiled_tma_atom(
72
+ op,
73
+ cute.logical_divide(tensor, (None, None, 64)),
74
+ slayout,
75
+ cta_tiler=(self.BT, 1, dim),
76
+ )
77
+ return atom, tma_tensor, slayout
78
+
79
+ @cute.jit
80
+ def __call__(
81
+ self,
82
+ K: cute.Tensor,
83
+ V: cute.Tensor,
84
+ U: cute.Tensor,
85
+ W: cute.Tensor,
86
+ g: cute.Tensor,
87
+ beta: cute.Tensor,
88
+ g_cu: cute.Tensor,
89
+ cu_seqlens: cute.Tensor,
90
+ chunk_indices: cute.Tensor,
91
+ total_chunks: cute.Tensor,
92
+ num_sms: Int32,
93
+ stream: CUstream,
94
+ ):
95
+ tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
96
+ tma_s2g = cpasync.CopyBulkTensorTileS2GOp()
97
+
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)
102
+
103
+ grid = (num_sms // self.Hv, self.Hv, 1)
104
+ block = (self.num_warps * 32, 1, 1)
105
+ self.kernel(
106
+ K_args,
107
+ V_args,
108
+ U_args,
109
+ W_args,
110
+ g,
111
+ beta,
112
+ g_cu,
113
+ cu_seqlens,
114
+ chunk_indices,
115
+ total_chunks,
116
+ ).launch(grid=grid, block=block, stream=stream)
117
+
118
+ @cute.kernel
119
+ def kernel(
120
+ 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],
125
+ g: cute.Tensor,
126
+ beta: cute.Tensor,
127
+ g_cu: cute.Tensor,
128
+ cu_seqlens: cute.Tensor,
129
+ chunk_indices: cute.Tensor,
130
+ total_chunks: cute.Tensor,
131
+ ):
132
+ tid, _, _ = cute.arch.thread_idx()
133
+ bid, head_id, _ = cute.arch.block_idx()
134
+ grid_x, _, _ = cute.arch.grid_dim()
135
+
136
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
137
+ lane_id = tid % 32
138
+ k_head_id = head_id // (self.Hv // self.H)
139
+
140
+ BT = self.BT
141
+ K_dim = self.K_dim
142
+ V_dim = self.V_dim
143
+ num_stages = self.num_stages
144
+
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
149
+
150
+ def allocate_tensor(smem, dtype, layout):
151
+ return smem.allocate_tensor(
152
+ dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
153
+ )
154
+
155
+ 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]
160
+
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)
166
+
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)
170
+
171
+ tma_mbar = smem.allocate_array(Int64, num_stages)
172
+ mma_kkt_mbar = smem.allocate_array(Int64, num_stages)
173
+ inv_mbar = smem.allocate_array(Int64, num_stages)
174
+ mma_u_mbar = smem.allocate_array(Int64, num_stages)
175
+ mma_w_mbar = smem.allocate_array(Int64, num_stages)
176
+ epi_mbar = smem.allocate_array(Int64, num_stages)
177
+ taddr = smem.allocate(Int32, 4)
178
+
179
+ kkt_tmem = 0
180
+ U_tmem_base = kkt_tmem + BT
181
+ Ab_tmem_base = U_tmem_base + V_dim * num_stages
182
+ assert Ab_tmem_base + (BT // 2) * num_stages <= 512
183
+
184
+ # prepare ldmatrix/stmatrix ops
185
+ ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4)
186
+ stsm_op = warp.StMatrix8x8x16bOp(num_matrices=4)
187
+ ldsm_trans_op = warp.LdMatrix8x8x16bOp(num_matrices=4, transpose=True)
188
+ ldsm_atom = cute.make_copy_atom(ldsm_op, BFloat16)
189
+ stsm_atom = cute.make_copy_atom(stsm_op, BFloat16)
190
+ ldsm_trans_atom = cute.make_copy_atom(ldsm_trans_op, BFloat16)
191
+
192
+ if warp_id == 0:
193
+ with cute.arch.elect_one():
194
+ for i in cutlass.range_constexpr(num_stages):
195
+ cute.arch.mbarrier_init(tma_mbar + i, 1)
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)
199
+ cute.arch.mbarrier_init(mma_w_mbar + i, 1)
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)
207
+ cute.arch.sync_threads()
208
+
209
+ num_global_chunks = total_chunks[0]
210
+ if warp_id == 9:
211
+ # TMA warp
212
+ stage_id = 0
213
+ parity = 1
214
+
215
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
216
+ seq_id = chunk_indices[global_chunk_id, 0]
217
+ chunk_id = chunk_indices[global_chunk_id, 1]
218
+ bos = cu_seqlens[seq_id]
219
+
220
+ # since off_t is not a multiple of BT, we need to use
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]),
225
+ tiler=(BT, K_dim),
226
+ coord=(chunk_id, 0),
227
+ )
228
+ gV = cute.local_tile(
229
+ cute.domain_offset((bos, 0), tmaV[None, head_id, None]),
230
+ tiler=(BT, V_dim),
231
+ coord=(chunk_id, 0),
232
+ )
233
+
234
+ # when UW MMA is done, K and V TMA buffers are released
235
+ cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity)
236
+
237
+ with cute.arch.elect_one():
238
+ STAGE_SIZE = BT * (K_dim + V_dim) * 2
239
+ cute.arch.mbarrier_arrive_and_expect_tx(mbar, STAGE_SIZE)
240
+ simple_tma_copy(K_tma_atom, gK, sK[None, None, stage_id], mbar)
241
+ simple_tma_copy(
242
+ V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
243
+ )
244
+
245
+ stage_id = (stage_id + 1) % num_stages
246
+ if stage_id == 0:
247
+ parity ^= 1
248
+
249
+ elif warp_id == 8:
250
+ # MMA warp
251
+ _tcgen05.alloc(taddr)
252
+
253
+ stage_id = 0
254
+ parity = 0
255
+
256
+ kkt_idesc = _tcgen05.make_bf16_idesc(BT, BT)
257
+ u_idesc = _tcgen05.make_bf16_idesc(BT, V_dim, transpose_B=True)
258
+ w_idesc = _tcgen05.make_bf16_idesc(BT, K_dim, transpose_B=True)
259
+
260
+ # LBO=BT*128 is ignored for K-major
261
+ sdesc_template = _tcgen05.make_sdesc_128B_swizzle(BT * 128)
262
+
263
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
264
+ U_tmem = U_tmem_base + V_dim * stage_id
265
+ W_tmem = U_tmem | (16 << 16)
266
+ Ab_tmem = Ab_tmem_base + (BT // 2) * stage_id
267
+ Abg_tmem = Ab_tmem | (16 << 16)
268
+
269
+ ##### KKT MMA: KKT = K @ K.T #####
270
+ kaddr = sK[None, None, stage_id].iterator.toint()
271
+ kdesc_base = sdesc_template | (kaddr >> 4)
272
+
273
+ # wait for TMA data to arrive
274
+ # kkt tmem is guaranteed to be free as this is issued
275
+ # after the previous kkt's consumer (inv warps)
276
+ cute.arch.mbarrier_wait(tma_mbar + stage_id, parity)
277
+ _tcgen05.fence_after_thread_sync()
278
+
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)
291
+
292
+ ##### U/W MMA: U = Ab @ V, W = Abg @ K #####
293
+ vaddr = sV[None, None, stage_id].iterator.toint()
294
+ vdesc = sdesc_template | (vaddr >> 4)
295
+ kdesc = sdesc_template | (kaddr >> 4)
296
+
297
+ # wait for epilogue to release tmem buffer
298
+ cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1)
299
+ cute.arch.mbarrier_wait(inv_mbar + stage_id, parity)
300
+ _tcgen05.fence_after_thread_sync()
301
+
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)
309
+
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)
316
+
317
+ stage_id = (stage_id + 1) % num_stages
318
+ if stage_id == 0:
319
+ parity ^= 1
320
+
321
+ cute.arch.mbarrier_wait(epi_mbar + stage_id, parity ^ 1)
322
+ _tcgen05.dealloc()
323
+
324
+ elif warp_id >= 4:
325
+ # inv warps
326
+ tid_ = tid % 128
327
+ warp_id_ = warp_id % 4
328
+
329
+ stage_id = 0
330
+ parity = 0
331
+
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)]
337
+
338
+ # 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)
342
+
343
+ # indices for ldmatrix layout later
344
+ row_indices = cute.make_rmem_tensor((1, 2, 1), Int32)
345
+ row_indices[0, 0, 0] = warp_id_ * 16 + (lane_id // 4)
346
+ row_indices[0, 1, 0] = warp_id_ * 16 + (lane_id // 4) + 8
347
+ row_indices = row_indices.load()
348
+
349
+ col_indices = cute.make_rmem_tensor((2, 1, 2), Int32)
350
+ col_indices[0, 0, 0] = (lane_id % 4) * 2 + 0
351
+ col_indices[1, 0, 0] = (lane_id % 4) * 2 + 1
352
+ col_indices[0, 0, 1] = (lane_id % 4) * 2 + 8
353
+ col_indices[1, 0, 1] = (lane_id % 4) * 2 + 9
354
+ col_indices = col_indices.load()
355
+
356
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
357
+ seq_id = chunk_indices[global_chunk_id, 0]
358
+ chunk_id = chunk_indices[global_chunk_id, 1]
359
+ bos = cu_seqlens[seq_id]
360
+ eos = cu_seqlens[seq_id + 1]
361
+ off_t = bos + chunk_id * BT
362
+
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) #####
403
+ if warp_id_ == 0:
404
+ cute.arch.mbarrier_wait(mma_kkt_mbar + stage_id, parity)
405
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
406
+ _tcgen05.fence_after_thread_sync()
407
+
408
+ # tmem 16x256b layout / ldmatrix layout
409
+ # mode0 is 8 rows together
410
+ # mode1 is top and bottom 8 rows
411
+ # 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))
418
+
419
+ # mode0 is 2 consecutive elems
420
+ # mode1 is top and bottom 8 rows
421
+ # mode2 is next 8 columns
422
+ # mode3 is repeating that 16x16 tile pattern
423
+ kkt = _tcgen05.ld(kkt_tmem, 0, "16x256b", BT // 8)
424
+ kkt = kkt.reshape((2, 2, 2, BT // 16))
425
+
426
+ for i in cutlass.range_constexpr(BT // 16):
427
+ # mode0 is 2 elems next to each other
428
+ # mode1 is 4 pairs of elems on 1 row
429
+ # mode2 is top and bottom 8 rows
430
+ # mode3 is next 16 columns
431
+ col_coord = (None, lane_id % 4, None, i)
432
+ s_g_cu_view = cute.make_tensor(s_g_cu, (2, 4, 2, BT // 16))
433
+ g_cu_col = s_g_cu_view[col_coord].load().reshape((2, 1, 2))
434
+
435
+ Gamma = cute.math.exp(g_cu_row - g_cu_col, fastmath=True)
436
+ A = kkt[None, None, None, i] * beta_row * Gamma
437
+
438
+ # strict lower mask
439
+ # NOTE: for OOB t position, s_beta is filled with zeros.
440
+ # hence, we don't need to apply bounds check for columns.
441
+ A_masked = cute.where(row_indices > col_indices + i * 16, A, 0.0)
442
+
443
+ # pack to BF16
444
+ # CuteDSL doesn't generate cvt.bf16x2.f32 here for some reasons
445
+ 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
+ )
458
+
459
+ # store to smem
460
+ cute.copy(
461
+ stsm_atom,
462
+ cute.recast_tensor(packed, BFloat16),
463
+ sA_ldsm[warp_id_, None, i],
464
+ )
465
+
466
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
467
+
468
+ ##### Phase 3: matrix inverse #####
469
+ # we use Newton-Schulz iterations to compute the inverse
470
+ # of the four 16x16 diagonal blocks.
471
+ # Ai_new = 2 Ai - Ai @ M @ Ai
472
+ # where M = I + A
473
+ #
474
+ # we do this with 2 MMAs:
475
+ # 1. -AiM = Ai @ (-M)
476
+ # 2. Ai_new = 2 Ai + (-AiM) @ Ai
477
+ zeros_f32 = cute.make_rmem_tensor(4, Float32)
478
+ zeros_f32.fill(0.0)
479
+
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
+
489
+ Ai_bf16 = cute.make_rmem_tensor(8, BFloat16)
490
+ mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16)
491
+ M_bf16 = cute.make_rmem_tensor(8, BFloat16)
492
+ acc = cute.make_rmem_tensor((4, 2), Float32)
493
+
494
+ # share the same storage
495
+ Ai = cute.recast_tensor(Ai_bf16, Uint32)
496
+ mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2)
497
+ M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2)
498
+
499
+ # initial guess: Ai = I-A
500
+ cute.copy(ldsm_atom, sA_ldsm[warp_id_, None, warp_id_], Ai_bf16)
501
+ for i in cutlass.range_constexpr(4):
502
+ Ai[i] ^= Uint32(0x80008000) # negate A
503
+ set_diagonal(Ai, lane_id)
504
+
505
+ # (4, 2)
506
+ Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4)
507
+
508
+ # 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)
511
+ for i in cutlass.range_constexpr(4):
512
+ M[i] ^= Uint32(0x80008000)
513
+
514
+ # 3 rounds of Newton-Schulz
515
+ for _ in cutlass.range_constexpr(3):
516
+ # First MMA: -AiM = Ai @ (-M)
517
+ cute.copy(stsm_atom, Ai_bf16, sA_ldsm[warp_id_, None, warp_id_])
518
+ cute.arch.sync_warp()
519
+ acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32)
520
+ acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32)
521
+ Ai_bf16.store(acc.load().to(BFloat16))
522
+
523
+ # Second MMA: Ai_new = 2Ai + (-AiM) @ Ai
524
+ for j in cutlass.range_constexpr(8):
525
+ Ai_f32[j] *= 2.0
526
+ cute.copy(
527
+ ldsm_trans_atom,
528
+ sA_ldsm[warp_id_, None, warp_id_],
529
+ mma_B_bf16,
530
+ )
531
+ Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0])
532
+ Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1])
533
+ Ai_bf16.store(Ai_f32.load().to(BFloat16))
534
+
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)
537
+
538
+ # off-diagonal by 1
539
+ # given
540
+ # [ Ai00 ]
541
+ # [ A10 Ai11 ]
542
+ # [ A20 A21 Ai22 ]
543
+ # [ A30 A31 A32 Ai33]
544
+ # warp1: Ai10 = -Ai11 @ A10 @ Ai00
545
+ # warp2: Ai21 = -Ai22 @ A21 @ Ai11
546
+ # warp3: Ai32 = -Ai33 @ A32 @ Ai22
547
+ if warp_id_ > 0:
548
+ neg_Ai = cute.make_rmem_tensor(4, Uint32)
549
+ for i in cutlass.range_constexpr(4):
550
+ neg_Ai[i] = Ai[i] ^ Uint32(0x80008000)
551
+
552
+ cute.copy(
553
+ ldsm_trans_atom,
554
+ sA_ldsm[warp_id_, None, warp_id_ - 1],
555
+ mma_B_bf16,
556
+ )
557
+ acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32)
558
+ acc[None, 1] = mma_bf16(neg_Ai, mma_B[None, 1], zeros_f32)
559
+ Ai_bf16.store(acc.load().to(BFloat16))
560
+
561
+ cute.copy(
562
+ ldsm_trans_atom,
563
+ sAi_ldsm[warp_id_ - 1, None, warp_id_ - 1],
564
+ mma_B_bf16,
565
+ )
566
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
567
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
568
+ Ai_bf16.store(acc.load().to(BFloat16))
569
+ cute.copy(
570
+ stsm_atom,
571
+ Ai_bf16,
572
+ sAi_ldsm[warp_id_, None, warp_id_ - 1],
573
+ )
574
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
575
+
576
+ # 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:
580
+ cute.copy(
581
+ ldsm_atom,
582
+ sA_ldsm[warp_id_ + 2, None, warp_id_],
583
+ Ai_bf16,
584
+ )
585
+ cute.copy(
586
+ ldsm_trans_atom,
587
+ sAi_ldsm[warp_id_, None, warp_id_],
588
+ mma_B_bf16,
589
+ )
590
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
591
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
592
+
593
+ cute.copy(
594
+ ldsm_atom,
595
+ sA_ldsm[warp_id_ + 2, None, warp_id_ + 1],
596
+ Ai_bf16,
597
+ )
598
+ cute.copy(
599
+ ldsm_trans_atom,
600
+ sAi_ldsm[warp_id_ + 1, None, warp_id_],
601
+ mma_B_bf16,
602
+ )
603
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
604
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])
605
+
606
+ tmp = cute.make_rmem_tensor(8, BFloat16)
607
+ tmp.store(acc.load().to(BFloat16))
608
+ cute.copy(stsm_atom, tmp, sAi_ldsm[warp_id_ + 2, None, warp_id_])
609
+ cute.arch.sync_warp()
610
+
611
+ cute.copy(
612
+ ldsm_atom, sAi_ldsm[warp_id_ + 2, None, warp_id_ + 2], Ai_bf16
613
+ )
614
+ for i in cutlass.range_constexpr(4):
615
+ Ai[i] ^= Uint32(0x80008000)
616
+ cute.copy(
617
+ ldsm_trans_atom,
618
+ sAi_ldsm[warp_id_ + 2, None, warp_id_],
619
+ mma_B_bf16,
620
+ )
621
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
622
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
623
+ 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)
626
+
627
+ # 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)
632
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
633
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
634
+
635
+ 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)
638
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0])
639
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1])
640
+
641
+ tmp = cute.make_rmem_tensor(8, BFloat16)
642
+ tmp.store(acc.load().to(BFloat16))
643
+ cute.copy(stsm_atom, tmp, sAi_ldsm[3, None, 0])
644
+ cute.arch.sync_warp()
645
+
646
+ cute.copy(ldsm_atom, sAi_ldsm[3, None, 3], Ai_bf16)
647
+ 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)
650
+ acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32)
651
+ acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32)
652
+ 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
+
681
+ _tcgen05.wait_st()
682
+ _tcgen05.fence_before_thread_sync()
683
+ cute.arch.mbarrier_arrive(inv_mbar + stage_id)
684
+
685
+ stage_id = (stage_id + 1) % num_stages
686
+ if stage_id == 0:
687
+ parity ^= 1
688
+
689
+ elif warp_id < 4:
690
+ # epi warps
691
+ stage_id = 0
692
+ parity = 0
693
+
694
+ # ((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))
697
+
698
+ # sW shape: [BT, (64, K_dim/64)]
699
+ # sW_view shape: [(8, 2), (4, K_dim/64)]
700
+ s_row = warp_id * 16 + lane_id % 16 # select the rows of [16,16] tile
701
+ sW_view = cute.zipped_divide(
702
+ sW[s_row, None],
703
+ tiler=cute.make_layout((8, 2)),
704
+ )
705
+ sU_view = cute.zipped_divide(
706
+ sU[s_row, None],
707
+ tiler=cute.make_layout((8, 2)),
708
+ )
709
+
710
+ # select the 8 columns within [16,16] tile
711
+ sW_view = sW_view[(None, lane_id // 16), None]
712
+ sU_view = sU_view[(None, lane_id // 16), None]
713
+
714
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
715
+ # wait for W MMA + previous TMA store to finish
716
+ U_tmem = U_tmem_base + V_dim * stage_id
717
+ if warp_id == 0:
718
+ cute.arch.mbarrier_wait(mma_w_mbar + stage_id, parity)
719
+ elif warp_id == 1:
720
+ with cute.arch.elect_one():
721
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
722
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
723
+ _tcgen05.fence_after_thread_sync()
724
+
725
+ w_f32 = _tcgen05.ld(warp_id * 32 + 16, U_tmem, "16x256b", K_dim // 8)
726
+ _tcgen05.wait_ld()
727
+ w_bf16 = cute.make_rmem_tensor((8, K_dim // 16), BFloat16)
728
+ w_bf16.store(w_f32.to(BFloat16))
729
+ cute.copy(stsm_atom, w_bf16, sW_view)
730
+
731
+ # wait for U MMA + issue W TMA store
732
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
733
+ fence_before_tma_store()
734
+ if warp_id == 0:
735
+ cute.arch.mbarrier_wait(mma_u_mbar + stage_id, parity)
736
+ elif warp_id == 1:
737
+ # don't need to commit
738
+ simple_tma_copy(
739
+ W_tma_atom, sW, gW_tiles[(None, global_chunk_id), None]
740
+ )
741
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
742
+ _tcgen05.fence_after_thread_sync()
743
+
744
+ u_f32 = _tcgen05.ld(warp_id * 32, U_tmem, "16x256b", V_dim // 8)
745
+ _tcgen05.wait_ld()
746
+ _tcgen05.fence_before_thread_sync()
747
+ cute.arch.mbarrier_arrive(epi_mbar + stage_id)
748
+ u_bf16 = cute.make_rmem_tensor((8, V_dim // 16), BFloat16)
749
+ u_bf16.store(u_f32.to(BFloat16))
750
+ cute.copy(stsm_atom, u_bf16, sU_view)
751
+
752
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
753
+ fence_before_tma_store()
754
+ if warp_id == 1:
755
+ simple_tma_copy(
756
+ U_tma_atom, sU, gU_tiles[(None, global_chunk_id), None]
757
+ )
758
+ with cute.arch.elect_one():
759
+ cute.arch.cp_async_bulk_commit_group()
760
+
761
+ stage_id = (stage_id + 1) % num_stages
762
+ if stage_id == 0:
763
+ parity ^= 1
764
+
765
+ @cache
766
+ @staticmethod
767
+ def compile(H: int, Hv: int, K_dim: int, V_dim: int, num_stages: int = 2):
768
+ total_t = cute.sym_int()
769
+ pad_t = cute.sym_int()
770
+ total_chunks_n = cute.sym_int()
771
+ num_sequences = cute.sym_int()
772
+
773
+ K = make_fake_tensor(BFloat16, (total_t, H, K_dim), divisibility=16)
774
+ V = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16)
775
+ U = make_fake_tensor(BFloat16, (pad_t, Hv, V_dim), divisibility=16)
776
+ W = make_fake_tensor(BFloat16, (pad_t, Hv, K_dim), divisibility=16)
777
+ g = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
778
+ beta = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
779
+ g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
780
+ cu_seqlens = make_fake_tensor(Int32, (num_sequences,), divisibility=1)
781
+ chunk_indices = make_fake_tensor(Int32, (total_chunks_n, 2), divisibility=2)
782
+ total_chunks = make_fake_tensor(Int32, (1,), divisibility=1)
783
+
784
+ kernel = Sm100ChunkUWKernel(H, Hv, K_dim, V_dim, num_stages)
785
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
786
+ return cute.compile(
787
+ kernel,
788
+ K,
789
+ V,
790
+ U,
791
+ W,
792
+ g,
793
+ beta,
794
+ g_cu,
795
+ cu_seqlens,
796
+ chunk_indices,
797
+ total_chunks,
798
+ Int32(148),
799
+ stream,
800
+ options="--enable-tvm-ffi",
801
+ )
802
+
803
+
804
+def kkt_inv_uw_cutedsl(
805
+ K: torch.Tensor,
806
+ V: torch.Tensor,
807
+ U: torch.Tensor,
808
+ W: torch.Tensor,
809
+ g: torch.Tensor,
810
+ beta: torch.Tensor,
811
+ g_cu: torch.Tensor,
812
+ cu_seqlens: torch.Tensor,
813
+ chunk_indices: torch.Tensor,
814
+ total_chunks: torch.Tensor,
815
+ num_sms: int = 148,
816
+) -> None:
817
+ _, Hv, V_dim = V.shape
818
+ _, H, K_dim = K.shape
819
+
820
+ Sm100ChunkUWKernel.compile(H, Hv, K_dim, V_dim)(
821
+ K,
822
+ V,
823
+ U,
824
+ W,
825
+ g,
826
+ beta,
827
+ g_cu,
828
+ cu_seqlens,
829
+ chunk_indices,
830
+ total_chunks,
831
+ num_sms,
832
+ )
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_o.py
new
+630
@@ -0,0 +1,630 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from functools import cache
4
+
5
+import cutlass
6
+import torch
7
+from cuda.bindings.driver import CUstream
8
+from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
9
+from cutlass.cute.nvgpu import cpasync, warp
10
+from quack.compile_utils import make_fake_tensor
11
+
12
+from vllm.cute_utils import (
13
+ EVICT_FIRST,
14
+ _tcgen05,
15
+ cvt,
16
+ fence_before_tma_store,
17
+ simple_tma_copy,
18
+)
19
+
20
+
21
+class Sm100ChunkOKernel:
22
+ """Compute per-token output from recurrent and intra-chunk terms.
23
+
24
+ Gamma[i,j] = exp(g_cu[i] - g_cu[j])
25
+ P = mask((Q @ K.T) * Gamma)
26
+ O = scale * (exp(g_cu) * (Q @ H.T) + P @ V)
27
+ """
28
+
29
+ def __init__(
30
+ self,
31
+ H: int,
32
+ Hv: int,
33
+ K_dim: int,
34
+ V_dim: int,
35
+ BT: int = 64,
36
+ num_stages: int = 2,
37
+ ) -> None:
38
+ assert Hv % H == 0
39
+ assert K_dim == 128
40
+ assert V_dim == 128
41
+ assert BT == 64
42
+ self.H = H
43
+ self.Hv = Hv
44
+ self.K_dim = K_dim
45
+ self.V_dim = V_dim
46
+ self.BT = BT
47
+ self.num_stages = num_stages
48
+ self.num_warps = 10
49
+
50
+ @cute.jit
51
+ def _make_bf16_tma_args(
52
+ self,
53
+ tensor: cute.Tensor,
54
+ dim: cutlass.Constexpr[int],
55
+ op: cpasync.TmaCopyOp,
56
+ stages: cutlass.Constexpr[int],
57
+ ):
58
+ swizzle_128B = cute.make_swizzle(3, 4, 3)
59
+ slayout = cute.make_layout(
60
+ (self.BT, 1, (64, dim // 64), stages),
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(
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
71
+
72
+ @cute.jit
73
+ def _make_h_tma_args(
74
+ self,
75
+ tensor: cute.Tensor,
76
+ op: cpasync.TmaCopyOp,
77
+ stages: cutlass.Constexpr[int],
78
+ ):
79
+ num_elems = 128 // (tensor.element_type.width // 8)
80
+ swizzle_128B = cute.make_swizzle(3, 4, 3)
81
+ 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),
84
+ )
85
+ slayout = cute.make_composed_layout(swizzle_128B, 0, slayout)
86
+ atom, tma_tensor = cpasync.make_tiled_tma_atom(
87
+ op,
88
+ cute.logical_divide(tensor, (None, None, num_elems)),
89
+ slayout,
90
+ cta_tiler=(1, self.V_dim, self.K_dim),
91
+ )
92
+ return atom, tma_tensor, slayout
93
+
94
+ @cute.jit
95
+ def __call__(
96
+ self,
97
+ q: cute.Tensor,
98
+ k: cute.Tensor,
99
+ v_new_chunks: cute.Tensor,
100
+ h: cute.Tensor,
101
+ g_cu: cute.Tensor,
102
+ o: cute.Tensor,
103
+ cu_seqlens: cute.Tensor,
104
+ chunk_indices: cute.Tensor,
105
+ total_chunks: cute.Tensor,
106
+ scale: Float32,
107
+ num_sms: Int32,
108
+ stream: CUstream,
109
+ ):
110
+ grid = (num_sms // self.Hv, self.Hv, 1)
111
+ block = (self.num_warps * 32, 1, 1)
112
+ tma_g2s = cpasync.CopyBulkTensorTileG2SOp()
113
+ 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(
117
+ v_new_chunks, self.V_dim, tma_g2s, self.num_stages
118
+ )
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)
121
+ self.kernel(
122
+ Q_args,
123
+ K_args,
124
+ V_args,
125
+ H_args,
126
+ O_args,
127
+ g_cu,
128
+ o,
129
+ cu_seqlens,
130
+ chunk_indices,
131
+ total_chunks,
132
+ scale,
133
+ ).launch(grid=grid, block=block, stream=stream)
134
+
135
+ @cute.kernel
136
+ def kernel(
137
+ 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],
143
+ g_cu: cute.Tensor,
144
+ o: cute.Tensor,
145
+ cu_seqlens: cute.Tensor,
146
+ chunk_indices: cute.Tensor,
147
+ total_chunks: cute.Tensor,
148
+ scale: Float32,
149
+ ):
150
+ tid, _, _ = cute.arch.thread_idx()
151
+ bid, v_head_id, _ = cute.arch.block_idx()
152
+ grid_x, _, _ = cute.arch.grid_dim()
153
+ warp_id = cute.arch.make_warp_uniform(tid // 32)
154
+ lane_id = tid % 32
155
+
156
+ BT = self.BT
157
+ K_dim = self.K_dim
158
+ V_dim = self.V_dim
159
+ num_stages = self.num_stages
160
+
161
+ heads_per_qk = self.Hv // self.H
162
+ k_head_id = v_head_id // heads_per_qk
163
+ num_global_chunks = total_chunks[0]
164
+
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
+
171
+ def allocate_tensor(smem, dtype, layout):
172
+ return smem.allocate_tensor(
173
+ dtype, layout.outer, byte_alignment=128, swizzle=layout.inner
174
+ )
175
+
176
+ 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]
182
+
183
+ s_g_cu = smem.allocate_array(Float32, BT)
184
+ qk_full_mbar = smem.allocate_array(Int64, num_stages)
185
+ hv_full_mbar = smem.allocate_array(Int64, num_stages)
186
+ qk_empty_mbar = smem.allocate_array(Int64, num_stages)
187
+ pv_mma_mbar = smem.allocate_array(Int64, num_stages)
188
+ qk_mbar = smem.allocate_array(Int64, 1)
189
+ mask_mbar = smem.allocate_array(Int64, 1)
190
+ epi_mbar = smem.allocate_array(Int64, 1)
191
+ taddr = smem.allocate(Int32, 4)
192
+
193
+ qk_tmem = 0
194
+ p_tmem = 64
195
+ out_tmem = 128
196
+ qh_tmem = 256
197
+
198
+ if warp_id == 0:
199
+ with cute.arch.elect_one():
200
+ for i in cutlass.range_constexpr(num_stages):
201
+ cute.arch.mbarrier_init(qk_full_mbar + i, 1)
202
+ cute.arch.mbarrier_init(qk_empty_mbar + i, 1)
203
+ cute.arch.mbarrier_init(hv_full_mbar + i, 1)
204
+ cute.arch.mbarrier_init(pv_mma_mbar + i, 1)
205
+ cute.arch.mbarrier_init(qk_mbar, 1)
206
+ cute.arch.mbarrier_init(mask_mbar, 128)
207
+ cute.arch.mbarrier_init(epi_mbar, 128)
208
+ cute.arch.mbarrier_init_fence()
209
+ 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)
214
+ cute.arch.sync_threads()
215
+
216
+ if warp_id == 9:
217
+ # TMA warp
218
+ stage_id = 0
219
+ parity = 1
220
+
221
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
222
+ seq_id = chunk_indices[global_chunk_id, 0]
223
+ chunk_id = chunk_indices[global_chunk_id, 1]
224
+ bos = cu_seqlens[seq_id]
225
+
226
+ # copy Q and K
227
+ q_tile = cute.local_tile(
228
+ cute.domain_offset((bos, 0), tmaQ[None, k_head_id, None]),
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]),
234
+ tiler=(BT, K_dim),
235
+ coord=(chunk_id, 0),
236
+ )
237
+ mbar = qk_full_mbar + stage_id
238
+
239
+ cute.arch.mbarrier_wait(qk_empty_mbar + stage_id, parity)
240
+
241
+ with cute.arch.elect_one():
242
+ STAGE_SIZE = BT * (K_dim + K_dim) * 2
243
+ 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
+
247
+ # copy H and V
248
+ gH = tmaH[global_chunk_id * self.Hv + v_head_id, None, None]
249
+ gV = cute.local_tile(
250
+ tmaV[None, v_head_id, None],
251
+ tiler=(BT, V_dim),
252
+ coord=(global_chunk_id, 0),
253
+ )
254
+ mbar = hv_full_mbar + stage_id
255
+
256
+ cute.arch.mbarrier_wait(pv_mma_mbar + stage_id, parity)
257
+
258
+ with cute.arch.elect_one():
259
+ H_STAGE_SIZE = V_dim * K_dim * 2
260
+ V_STAGE_SIZE = BT * V_dim * 2
261
+ cute.arch.mbarrier_arrive_and_expect_tx(
262
+ mbar, H_STAGE_SIZE + V_STAGE_SIZE
263
+ )
264
+ simple_tma_copy(
265
+ H_tma_atom, gH, sH[None, None, stage_id], mbar, EVICT_FIRST
266
+ )
267
+ simple_tma_copy(
268
+ V_tma_atom, gV, sV[None, None, stage_id], mbar, EVICT_FIRST
269
+ )
270
+
271
+ stage_id = (stage_id + 1) % num_stages
272
+ if stage_id == 0:
273
+ parity ^= 1
274
+
275
+ elif warp_id == 8:
276
+ # MMA warp
277
+ _tcgen05.alloc(taddr)
278
+
279
+ # LBO=BT*128 is ignored for K-major
280
+ sdesc_template = _tcgen05.make_sdesc_128B_swizzle(BT * 128)
281
+ qk_idesc = _tcgen05.make_bf16_idesc(BT, BT)
282
+ qh_idesc = _tcgen05.make_bf16_idesc(BT, V_dim)
283
+ pv_idesc = _tcgen05.make_bf16_idesc(BT, V_dim, transpose_B=True)
284
+
285
+ stage_id = 0
286
+ tma_parity = 0
287
+ mask_parity = 0
288
+
289
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
290
+ qaddr = sQ[None, None, stage_id].iterator.toint()
291
+ kaddr = sK[None, None, stage_id].iterator.toint()
292
+ haddr = sH[None, None, stage_id].iterator.toint()
293
+ vaddr = sV[None, None, stage_id].iterator.toint()
294
+ qdesc_base = sdesc_template | (qaddr >> 4)
295
+ kdesc_base = sdesc_template | (kaddr >> 4)
296
+ hdesc_base = sdesc_template | (haddr >> 4)
297
+ vdesc_base = sdesc_template | (vaddr >> 4)
298
+
299
+ ##### 1st MMA: Q @ K.T #####
300
+ # do this first to unblock mask(QK)
301
+ cute.arch.mbarrier_wait(epi_mbar, mask_parity ^ 1)
302
+ cute.arch.mbarrier_wait(qk_full_mbar + stage_id, tma_parity)
303
+ _tcgen05.fence_after_thread_sync()
304
+
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)
314
+
315
+ ##### 2nd MMA: Q @ H.T #####
316
+ cute.arch.mbarrier_wait(hv_full_mbar + stage_id, tma_parity)
317
+ _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)
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)
339
+
340
+ stage_id = (stage_id + 1) % num_stages
341
+ if stage_id == 0:
342
+ tma_parity ^= 1
343
+ mask_parity ^= 1
344
+
345
+ # wait for epilogue to finish for deallocation
346
+ cute.arch.mbarrier_wait(epi_mbar, mask_parity ^ 1)
347
+ _tcgen05.dealloc()
348
+
349
+ elif warp_id >= 4:
350
+ # masking warps
351
+ warp_id_ = warp_id % 4
352
+ tid_ = tid % 128
353
+ row0 = warp_id_ * 16 + lane_id // 4
354
+ row1 = row0 + 8
355
+
356
+ parity = 0
357
+
358
+ # for ldmatrix layout later
359
+ row_indices = cute.make_rmem_tensor(2, Int32)
360
+ row_indices[0] = warp_id_ * 16 + lane_id // 4
361
+ row_indices[1] = warp_id_ * 16 + lane_id // 4 + 8
362
+ row_indices = row_indices.load().reshape((1, 2))
363
+
364
+ col_indices = cute.make_rmem_tensor(2, Int32)
365
+ col_indices[0] = (lane_id % 4) * 2
366
+ col_indices[1] = (lane_id % 4) * 2 + 1
367
+ col_indices = col_indices.load().reshape((2, 1))
368
+
369
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
370
+ if tid_ < BT:
371
+ seq_id = chunk_indices[global_chunk_id, 0]
372
+ chunk_id = chunk_indices[global_chunk_id, 1]
373
+ bos = cu_seqlens[seq_id]
374
+ eos = cu_seqlens[seq_id + 1]
375
+
376
+ t_ = bos + chunk_id * BT + tid_
377
+ s_g_cu[tid_] = g_cu[t_, v_head_id] if t_ < eos else Float32(0.0)
378
+
379
+ # wait for QK MMA
380
+ if warp_id_ == 0:
381
+ cute.arch.mbarrier_wait(qk_mbar, parity)
382
+ cute.arch.barrier(barrier_id=1, number_of_threads=128)
383
+ _tcgen05.fence_after_thread_sync()
384
+ qk = _tcgen05.ld(warp_id_ * 32, qk_tmem, "16x256b", BT // 8)
385
+ qk = qk.reshape((2, 2, BT // 8))
386
+ _tcgen05.wait_ld()
387
+
388
+ g_cu_rows = cute.make_rmem_tensor(2, Float32)
389
+ g_cu_rows[0] = s_g_cu[row0]
390
+ g_cu_rows[1] = s_g_cu[row1]
391
+ g_cu_rows = g_cu_rows.load().reshape((1, 2))
392
+
393
+ for i in cutlass.range_constexpr(BT // 8):
394
+ col = i * 8 + (lane_id % 4) * 2
395
+ g_cu_cols = cute.make_rmem_tensor(2, Float32)
396
+ g_cu_cols[0] = s_g_cu[col]
397
+ g_cu_cols[1] = s_g_cu[col + 1]
398
+ g_cu_cols = g_cu_cols.load().reshape((2, 1))
399
+
400
+ # apply gamma and causal mask
401
+ Gamma = cute.math.exp(g_cu_rows - g_cu_cols, fastmath=True)
402
+ tmp = qk[None, None, i] * Gamma
403
+ tmp = cute.where(row_indices >= col_indices + i * 8, tmp, 0.0)
404
+
405
+ # CuteDSL can't emit cvt.bf16x2.f32 here
406
+ attn_lo = cute.make_rmem_tensor(2, Uint32)
407
+ attn_lo[0] = cvt.fp32x2_to_bf16x2(tmp[0, 0], tmp[1, 0])
408
+ attn_lo[1] = cvt.fp32x2_to_bf16x2(tmp[0, 1], tmp[1, 1])
409
+ _tcgen05.st(warp_id_ * 32, p_tmem + i * 4, "16x128b", 1, attn_lo)
410
+
411
+ _tcgen05.wait_st()
412
+ _tcgen05.fence_before_thread_sync()
413
+ cute.arch.mbarrier_arrive(mask_mbar)
414
+
415
+ parity ^= 1
416
+
417
+ else:
418
+ # epilogue warps
419
+ # for ldmatrix layout later
420
+ row0 = warp_id * 16 + lane_id // 4
421
+ row1 = row0 + 8
422
+
423
+ stage_id = 0
424
+ mma_parity = 0
425
+
426
+ op = cute.nvgpu.CopyUniversalOp()
427
+ cp_4B = cute.make_copy_atom(op, BFloat16, num_bits_per_copy=32)
428
+ stsm_op = warp.StMatrix8x8x16bOp(num_matrices=4, transpose=False)
429
+ stsm_atom = cute.make_copy_atom(stsm_op, BFloat16)
430
+
431
+ # ldmatrix layout
432
+ # [total_seq_len, ((2, 4, WIDTH/8), V_DIM/WIDTH)]
433
+ WIDTH = 64
434
+ o_view = cute.logical_divide(
435
+ o[None, v_head_id, None],
436
+ (None, cute.make_layout((2, 4, WIDTH // 8))),
437
+ )
438
+ # select lane: [total_seq_len, 2, WIDTH/8, V_DIM/WIDTH]
439
+ o_view = o_view[None, ((None, lane_id % 4, None), None)]
440
+
441
+ for global_chunk_id in range(bid, num_global_chunks, grid_x):
442
+ seq_id = chunk_indices[global_chunk_id, 0]
443
+ chunk_id = chunk_indices[global_chunk_id, 1]
444
+ bos = cu_seqlens[seq_id]
445
+ eos = cu_seqlens[seq_id + 1]
446
+ chunk_start = bos + chunk_id * BT
447
+ full_chunk = chunk_start + BT <= eos
448
+
449
+ g_cu_rows = cute.make_rmem_tensor(2, Float32)
450
+ g_cu_rows.fill(0.0)
451
+
452
+ # load g_cu
453
+ if chunk_start + row0 < eos:
454
+ g_cu_rows[0] = cute.math.exp(
455
+ g_cu[chunk_start + row0, v_head_id], fastmath=True
456
+ )
457
+ if chunk_start + row1 < eos:
458
+ g_cu_rows[1] = cute.math.exp(
459
+ g_cu[chunk_start + row1, v_head_id], fastmath=True
460
+ )
461
+ g_cu_rows = g_cu_rows.load().reshape((1, 2, 1))
462
+
463
+ if warp_id == 0:
464
+ cute.arch.mbarrier_wait(pv_mma_mbar + stage_id, mma_parity)
465
+ elif warp_id == 3 and full_chunk:
466
+ cute.arch.cp_async_bulk_wait_group(0, read=True)
467
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
468
+ _tcgen05.fence_after_thread_sync()
469
+
470
+ if full_chunk:
471
+ # use TMA store: tmem->rmem->smem->gmem
472
+ for i in cutlass.range_constexpr(V_dim // WIDTH):
473
+ qh = _tcgen05.ld(
474
+ warp_id * 32, qh_tmem + i * WIDTH, "16x256b", WIDTH // 8
475
+ )
476
+ pv = _tcgen05.ld(
477
+ warp_id * 32, out_tmem + i * WIDTH, "16x256b", WIDTH // 8
478
+ )
479
+ _tcgen05.wait_ld()
480
+ if i == V_dim // WIDTH - 1:
481
+ _tcgen05.fence_before_thread_sync()
482
+ cute.arch.mbarrier_arrive(epi_mbar)
483
+
484
+ qh = qh.reshape((2, 2, WIDTH // 8))
485
+ pv = pv.reshape((2, 2, WIDTH // 8))
486
+
487
+ out_f32 = scale * (g_cu_rows * qh + pv)
488
+ out_bf16 = cute.make_rmem_tensor((8, WIDTH // 16), BFloat16)
489
+ out_bf16.store(out_f32.to(BFloat16).reshape((8, WIDTH // 16)))
490
+
491
+ # TODO: issue single cute.copy()
492
+ for j in cutlass.range_constexpr(WIDTH // 16):
493
+ s_row = warp_id * 16 + lane_id % 16
494
+ s_col = i * (WIDTH // 8) + j * 2 + lane_id // 16
495
+ sO_tile = cute.local_tile(sO[s_row, None], (8,), (s_col,))
496
+ cute.copy(stsm_atom, out_bf16[None, j], sO_tile)
497
+
498
+ cute.arch.barrier(barrier_id=2, number_of_threads=128)
499
+ fence_before_tma_store()
500
+ if warp_id == 3:
501
+ gO = cute.local_tile(
502
+ cute.domain_offset((bos, 0), tmaO[None, v_head_id, None]),
503
+ tiler=(BT, V_dim),
504
+ coord=(chunk_id, 0),
505
+ )
506
+ simple_tma_copy(O_tma_atom, sO, gO)
507
+ with cute.arch.elect_one():
508
+ cute.arch.cp_async_bulk_commit_group()
509
+
510
+ else:
511
+ # direct gmem store
512
+ # TODO: explore doing multiple 1D TMAs
513
+ for i in cutlass.range_constexpr(V_dim // WIDTH):
514
+ qh = _tcgen05.ld(
515
+ warp_id * 32, qh_tmem + i * WIDTH, "16x256b", WIDTH // 8
516
+ )
517
+ pv = _tcgen05.ld(
518
+ warp_id * 32, out_tmem + i * WIDTH, "16x256b", WIDTH // 8
519
+ )
520
+ _tcgen05.wait_ld()
521
+ if i == V_dim // WIDTH - 1:
522
+ _tcgen05.fence_before_thread_sync()
523
+ cute.arch.mbarrier_arrive(epi_mbar)
524
+
525
+ qh = qh.reshape((2, 2, WIDTH // 8))
526
+ pv = pv.reshape((2, 2, WIDTH // 8))
527
+
528
+ out_f32 = scale * (g_cu_rows * qh + pv)
529
+ out_bf16 = cute.make_rmem_tensor((2, 2, WIDTH // 8), BFloat16)
530
+ out_bf16.store(out_f32.to(BFloat16))
531
+
532
+ if chunk_start + row0 < eos:
533
+ cute.copy(
534
+ cp_4B,
535
+ out_bf16[None, 0, None],
536
+ o_view[chunk_start + row0, None, None, i],
537
+ )
538
+ if chunk_start + row1 < eos:
539
+ cute.copy(
540
+ cp_4B,
541
+ out_bf16[None, 1, None],
542
+ o_view[chunk_start + row1, None, None, i],
543
+ )
544
+
545
+ stage_id = (stage_id + 1) % num_stages
546
+ if stage_id == 0:
547
+ mma_parity ^= 1
548
+
549
+ @cache
550
+ @staticmethod
551
+ def compile(
552
+ H: int,
553
+ Hv: int,
554
+ K_dim: int,
555
+ V_dim: int,
556
+ BT: int = 64,
557
+ num_stages: int = 2,
558
+ ):
559
+ total_t = cute.sym_int()
560
+ pad_t = cute.sym_int()
561
+ 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)
569
+ g_cu = make_fake_tensor(Float32, (total_t, Hv), divisibility=4)
570
+ o = make_fake_tensor(BFloat16, (total_t, Hv, V_dim), divisibility=16)
571
+ cu_seqlens = make_fake_tensor(Int32, (cu_entries,), divisibility=1)
572
+ chunk_indices = make_fake_tensor(Int32, (total_chunks_n, 2), divisibility=2)
573
+ total_chunks = make_fake_tensor(Int32, (1,), divisibility=1)
574
+
575
+ kernel = Sm100ChunkOKernel(
576
+ H,
577
+ Hv,
578
+ K_dim,
579
+ V_dim,
580
+ BT,
581
+ num_stages,
582
+ )
583
+ stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
584
+ return cute.compile(
585
+ kernel,
586
+ q,
587
+ k,
588
+ v_new,
589
+ h_flat,
590
+ g_cu,
591
+ o,
592
+ cu_seqlens,
593
+ chunk_indices,
594
+ total_chunks,
595
+ Float32(1.0),
596
+ Int32(148),
597
+ stream,
598
+ options="--enable-tvm-ffi",
599
+ )
600
+
601
+
602
+def o_cutedsl(
603
+ q: torch.Tensor,
604
+ k: torch.Tensor,
605
+ v_new_chunks: torch.Tensor,
606
+ h: torch.Tensor,
607
+ g_cu: torch.Tensor,
608
+ o: torch.Tensor,
609
+ cu_seqlens: torch.Tensor,
610
+ chunk_indices: torch.Tensor,
611
+ total_chunks: torch.Tensor,
612
+ scale: float,
613
+ num_sms: int = 148,
614
+) -> None:
615
+ _, H, K_dim = q.shape
616
+ _, Hv, V_dim = o.shape
617
+
618
+ Sm100ChunkOKernel.compile(H, Hv, K_dim, V_dim)(
619
+ q,
620
+ k,
621
+ v_new_chunks.view(-1, Hv, V_dim),
622
+ h.view(-1, V_dim, K_dim),
623
+ g_cu,
624
+ o,
625
+ cu_seqlens,
626
+ chunk_indices,
627
+ total_chunks,
628
+ float(scale),
629
+ num_sms,
630
+ )
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py
+3
-6
@@ -11,10 +11,7 @@ from cutlass import BFloat16, Int32, Uint8, Uint32
11
from cutlass.cute.nvgpu import cpasync
12
from quack.compile_utils import make_fake_tensor
13
14
-from vllm.models.deepseek_v4.nvidia.ops.cutedsl_utils import (
15
- _bf16x2_mul,
16
- _fp8x4_to_bf16x4,
17
-)
14
+from vllm.cute_utils import _bf16x2_mul, cvt
15
16
17
def dequantize_and_gather_k_cache_cutedsl(
@@ -268,8 +265,8 @@ class DequantGatherKCacheKernel:
265
dequant0 = cute.make_rmem_tensor(4, Uint32)
266
dequant1 = cute.make_rmem_tensor(4, Uint32)
267
for j in cutlass.range_constexpr(2):
271
- tmp0 = _fp8x4_to_bf16x4(data0[j])
272
- tmp1 = _fp8x4_to_bf16x4(data1[j])
268
+ tmp0 = cvt.fp8x4_to_bf16x4(data0[j])
269
+ tmp1 = cvt.fp8x4_to_bf16x4(data1[j])
270
271
# BF16 multiply is safe because the scales are exact powers of 2.
272
dequant0[j * 2] = _bf16x2_mul(tmp0[0], scale0_bf16x2)
vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py
+20
-23
@@ -9,14 +9,11 @@ from cuda.bindings.driver import CUstream
9
from cutlass import BFloat16, Float32, Int64, Uint8, Uint32, const_expr
10
from quack.compile_utils import make_fake_tensor
11
12
-from vllm.models.deepseek_v4.nvidia.ops.cutedsl_utils import (
12
+from vllm.cute_utils import (
13
_bf16x2_abs,
14
_bf16x2_max,
15
- _bf16x2_to_fp32,
16
- _fp32x2_to_bf16x2,
17
- _fp32x4_to_fp8x4,
18
- _fp32x8_to_fp4x8,
19
- _recast_val,
15
+ cvt,
16
+ recast_val,
17
)
18
from vllm.vllm_flash_attn.cute import utils as cute_utils
19
@@ -225,8 +222,8 @@ class IndexerQRopeQuantKernel:
222
cute.copy(cp_u32x4, cute.recast_tensor(sin_src, Uint32), sin_bf16x2)
223
224
for i in cutlass.range_constexpr(4):
228
- cos0, cos1 = _bf16x2_to_fp32(cos_bf16x2[i])
229
- sin0, sin1 = _bf16x2_to_fp32(sin_bf16x2[i])
225
+ cos0, cos1 = cvt.bf16x2_to_fp32x2(cos_bf16x2[i])
226
+ sin0, sin1 = cvt.bf16x2_to_fp32x2(sin_bf16x2[i])
227
cos_vals[i * 2] = cos0
228
cos_vals[i * 2 + 1] = cos1
229
sin_vals[i * 2] = sin0
@@ -234,11 +231,11 @@ class IndexerQRopeQuantKernel:
231
232
for i in cutlass.range_constexpr(self.coarsen):
233
for j in cutlass.range_constexpr(8):
237
- q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j])
234
+ q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
235
rot0 = q0 * cos_vals[j] - q1 * sin_vals[j]
236
rot1 = q0 * sin_vals[j] + q1 * cos_vals[j]
237
# convert back to BF16 to match numerics
241
- q_bf16x2[i, j] = _fp32x2_to_bf16x2(rot0, rot1)
238
+ q_bf16x2[i, j] = cvt.fp32x2_to_bf16x2(rot0, rot1)
239
240
return (
241
q_bf16x2,
@@ -327,7 +324,7 @@ class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
324
_bf16x2_max,
325
width=MXFP4_BLOCK_SIZE // 16,
326
)
330
- amax_pair = _bf16x2_to_fp32(amax_bf16x2)
327
+ amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
328
amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
329
330
if in_bounds:
@@ -336,7 +333,7 @@ class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
333
# increments the exponent whenever fp4_scale is not exactly a power of 2
334
eps = cutlass.const_expr(float.fromhex("0x6p-126"))
335
fp4_scale = cute_utils.fmax(amax, eps) * Float32(1.0 / 6.0)
339
- bits = _recast_val(fp4_scale, Uint32)
336
+ bits = recast_val(fp4_scale, Uint32)
337
ue8m0 = cute_utils.shr_u32(
338
bits + Uint32(0x7FFFFF), Uint32(23)
339
) & Uint32(0xFF)
@@ -349,18 +346,18 @@ class IndexerQMxFp4Kernel(IndexerQRopeQuantKernel):
346
# If scale = 2^A and ue8m0 = A + 127, then inverse scale has exponent
347
# -A + 127 = 254 - ue8m0.
348
inv_scale_bits = (Uint32(254) - ue8m0) << Uint32(23)
352
- inv_fp4_scale = _recast_val(inv_scale_bits, Float32)
349
+ inv_fp4_scale = recast_val(inv_scale_bits, Float32)
350
351
vals = cute.make_rmem_tensor(16, Float32)
352
for j in cutlass.range_constexpr(8):
356
- q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j])
353
+ q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j])
354
vals[j * 2] = q0 * inv_fp4_scale
355
vals[j * 2 + 1] = q1 * inv_fp4_scale
356
357
# pack to FP4
358
packed = cute.make_rmem_tensor((2,), Uint32)
362
- packed[0] = _fp32x8_to_fp4x8(vals, 0)
363
- packed[1] = _fp32x8_to_fp4x8(vals, 8)
359
+ packed[0] = cvt.fp32x8_to_fp4x8(vals, 0)
360
+ packed[1] = cvt.fp32x8_to_fp4x8(vals, 8)
361
362
dst = q_fp4_tile[i, None]
363
cp_u32x2 = cute.make_copy_atom(cp_op, Uint32, num_bits_per_copy=64)
@@ -519,24 +516,24 @@ class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
516
_bf16x2_max,
517
width=self.subwarp_size,
518
)
522
- amax_pair = _bf16x2_to_fp32(amax_bf16x2)
519
+ amax_pair = cvt.bf16x2_to_fp32x2(amax_bf16x2)
520
amax = cute_utils.fmax(amax_pair[0], amax_pair[1])
521
522
# scale = max(amax, eps) / fp8_max, then rounded UP to the next
523
# power of two. Adding the mantissa mask before shifting out the
524
# mantissa bumps the exponent whenever s isn't a pure pow2.
525
fp32_scale = cute_utils.fmax(amax, Float32(1e-4)) * Float32(1.0 / 448.0)
529
- bits = _recast_val(fp32_scale, Uint32)
526
+ bits = recast_val(fp32_scale, Uint32)
527
scale_exp = cute_utils.shr_u32(
528
bits + Uint32(0x7FFFFF), Uint32(23)
529
) & Uint32(0xFF)
530
531
# rounded scale = 2^(scale_exp - 127); bit pattern is scale_exp << 23
532
fp8_scale_bits = scale_exp << Uint32(23)
536
- fp8_scale = _recast_val(fp8_scale_bits, Float32)
533
+ fp8_scale = recast_val(fp8_scale_bits, Float32)
534
# inverse = 2^-(scale_exp - 127); bit pattern is (254 - scale_exp) << 23
535
inv_scale_bits = (Uint32(254) - scale_exp) << Uint32(23)
539
- inv_fp8_scale = _recast_val(inv_scale_bits, Float32)
536
+ inv_fp8_scale = recast_val(inv_scale_bits, Float32)
537
538
# Weight fold: weights_out = weights * q_scale * scale_combined.
539
# All threads in the subwarp share the same fp8_scale after the
@@ -553,9 +550,9 @@ class IndexerQFp8Kernel(IndexerQRopeQuantKernel):
550
# (one cp.async-shaped 128-bit store per row).
551
packed = cute.make_rmem_tensor((4,), Uint32)
552
for j in cutlass.range_constexpr(4):
556
- q0, q1 = _bf16x2_to_fp32(q_bf16x2[i, j * 2])
557
- q2, q3 = _bf16x2_to_fp32(q_bf16x2[i, j * 2 + 1])
558
- packed[j] = _fp32x4_to_fp8x4(
553
+ q0, q1 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2])
554
+ q2, q3 = cvt.bf16x2_to_fp32x2(q_bf16x2[i, j * 2 + 1])
555
+ packed[j] = cvt.fp32x4_to_fp8x4(
556
q0 * inv_fp8_scale,
557
q1 * inv_fp8_scale,
558
q2 * inv_fp8_scale,
vllm/v1/attention/backends/gdn_attn.py
+37
-14
@@ -3,6 +3,7 @@
3
"""Backend for GatedDeltaNet attention."""
4
5
from dataclasses import dataclass
6
+from typing import Literal
7
8
import torch
9
@@ -90,6 +91,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
91
self.compilation_config = vllm_config.compilation_config
92
self.speculative_config = vllm_config.speculative_config
93
self.kv_cache_spec = kv_cache_spec
94
+ from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
95
+ _resolve_gdn_prefill_backend,
96
+ )
97
+
98
+ self.gdn_prefill_backend: Literal["triton", "flashinfer", "cutedsl"]
99
+ _, self.gdn_prefill_backend = _resolve_gdn_prefill_backend(vllm_config)
100
101
if self.speculative_config:
102
assert self.speculative_config.num_speculative_tokens is not None
@@ -316,22 +323,38 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
323
chunk_indices: torch.Tensor | None = None
324
chunk_offsets: torch.Tensor | None = None
325
if num_prefills > 0:
319
- # Only prefill batches use FLA chunk ops.
320
- # Pre-compute on CPU and async-copy to GPU to avoid
321
- # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
322
- from vllm.model_executor.layers.fla.ops.index import (
323
- prepare_chunk_indices,
324
- prepare_chunk_offsets,
325
- )
326
from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
327
328
- gpu_device = query_start_loc.device
329
- chunk_indices = prepare_chunk_indices(
330
- non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
331
- ).to(device=gpu_device, non_blocking=True)
332
- chunk_offsets = prepare_chunk_offsets(
333
- non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
334
- ).to(device=gpu_device, non_blocking=True)
328
+ if self.gdn_prefill_backend == "cutedsl":
329
+ from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (
330
+ prepare_metadata_cutedsl,
331
+ )
332
+
333
+ assert non_spec_query_start_loc is not None
334
+ assert non_spec_query_start_loc_cpu is not None
335
+ total_tokens = int(non_spec_query_start_loc_cpu[-1].item())
336
+ chunk_indices, chunk_offsets = prepare_metadata_cutedsl(
337
+ non_spec_query_start_loc,
338
+ total_tokens,
339
+ FLA_CHUNK_SIZE,
340
+ )
341
+ else:
342
+ gpu_device = query_start_loc.device
343
+ # Only prefill batches use FLA chunk ops.
344
+ # Pre-compute on CPU and async-copy to GPU to avoid
345
+ # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
346
+ from vllm.model_executor.layers.fla.ops.index import (
347
+ prepare_chunk_indices,
348
+ prepare_chunk_offsets,
349
+ )
350
+
351
+ assert non_spec_query_start_loc_cpu is not None
352
+ chunk_indices = prepare_chunk_indices(
353
+ non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
354
+ ).to(device=gpu_device, non_blocking=True)
355
+ chunk_offsets = prepare_chunk_offsets(
356
+ non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
357
+ ).to(device=gpu_device, non_blocking=True)
358
359
if num_prefills > 0:
360
has_initial_state = context_lens_tensor > 0