[ROCm][Triton] Migrate Kimi-K3 kernels from make_block_ptr to tensor … (#58769)
Signed-off-by: JadenMathias <jaden.mathias@amd.com>
Jaden Mathias committed
Oct 1, 2026 at 14:41 UTC
d848c4ed3d0c576c3c24860547fa2b4ee980b37d
3 files changed
+112
-351
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk.py
+47
-175
@@ -17,7 +17,11 @@ from vllm.third_party.flash_linear_attention.ops.chunk_delta_h import (
17
from vllm.third_party.flash_linear_attention.ops.cumsum import chunk_local_cumsum
18
from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
19
from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
20
-from vllm.third_party.flash_linear_attention.ops.op import exp2, log
20
+from vllm.third_party.flash_linear_attention.ops.op import (
21
+ exp2,
22
+ log,
23
+ make_tensor_descriptor,
24
+)
25
from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE, is_amd
26
from vllm.triton_utils import tl, triton
27
from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
@@ -90,86 +94,35 @@ def recompute_w_u_fwd_kernel(
94
T = eos - bos
95
else:
96
bos, eos = i_b * T, i_b * T + T
93
- p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
94
- b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32)
97
+ o_b = i_t * BT + tl.arange(0, BT)
98
+ b_b = tl.load(beta + (bos + o_b) * H + i_h, mask=o_b < T, other=0).to(tl.float32)
99
96
- p_A = tl.make_block_ptr(
97
- A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
98
- )
99
- b_A = tl.load(p_A, boundary_check=(0, 1))
100
+ desc_A = make_tensor_descriptor(A + (bos * H + i_h) * BT, [T, BT], [H * BT, 1], [BT, BT])
101
+ b_A = desc_A.load([i_t * BT, 0])
102
103
+ desc_v = make_tensor_descriptor(v + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV])
104
+ desc_u = make_tensor_descriptor(u + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV])
105
for i_v in range(tl.cdiv(V, BV)):
102
- p_v = tl.make_block_ptr(
103
- v + (bos * H + i_h) * V,
104
- (T, V),
105
- (H * V, 1),
106
- (i_t * BT, i_v * BV),
107
- (BT, BV),
108
- (1, 0),
109
- )
110
- p_u = tl.make_block_ptr(
111
- u + (bos * H + i_h) * V,
112
- (T, V),
113
- (H * V, 1),
114
- (i_t * BT, i_v * BV),
115
- (BT, BV),
116
- (1, 0),
117
- )
118
- b_v = tl.load(p_v, boundary_check=(0, 1))
106
+ b_v = desc_v.load([i_t * BT, i_v * BV])
107
b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
108
b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)
121
- tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
109
+ desc_u.store([i_t * BT, i_v * BV], b_u.to(desc_u.dtype))
110
111
+ desc_w = make_tensor_descriptor(w + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
112
+ desc_k = make_tensor_descriptor(k + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
113
+ desc_gk = make_tensor_descriptor(gk + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
114
for i_k in range(tl.cdiv(K, BK)):
124
- p_w = tl.make_block_ptr(
125
- w + (bos * H + i_h) * K,
126
- (T, K),
127
- (H * K, 1),
128
- (i_t * BT, i_k * BK),
129
- (BT, BK),
130
- (1, 0),
131
- )
132
- p_k = tl.make_block_ptr(
133
- k + (bos * H + i_h) * K,
134
- (T, K),
135
- (H * K, 1),
136
- (i_t * BT, i_k * BK),
137
- (BT, BK),
138
- (1, 0),
139
- )
140
- b_k = tl.load(p_k, boundary_check=(0, 1))
115
+ b_k = desc_k.load([i_t * BT, i_k * BK])
116
b_kb = b_k * b_b[:, None]
117
143
- p_gk = tl.make_block_ptr(
144
- gk + (bos * H + i_h) * K,
145
- (T, K),
146
- (H * K, 1),
147
- (i_t * BT, i_k * BK),
148
- (BT, BK),
149
- (1, 0),
150
- )
151
- b_gk = tl.load(p_gk, boundary_check=(0, 1))
118
+ b_gk = desc_gk.load([i_t * BT, i_k * BK])
119
b_kb *= exp2(b_gk)
120
if STORE_QG:
154
- p_q = tl.make_block_ptr(
155
- q + (bos * H + i_h) * K,
156
- (T, K),
157
- (H * K, 1),
158
- (i_t * BT, i_k * BK),
159
- (BT, BK),
160
- (1, 0),
161
- )
162
- p_qg = tl.make_block_ptr(
163
- qg + (bos * H + i_h) * K,
164
- (T, K),
165
- (H * K, 1),
166
- (i_t * BT, i_k * BK),
167
- (BT, BK),
168
- (1, 0),
169
- )
170
- b_q = tl.load(p_q, boundary_check=(0, 1))
121
+ desc_q = make_tensor_descriptor(q + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
122
+ desc_qg = make_tensor_descriptor(qg + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
123
+ b_q = desc_q.load([i_t * BT, i_k * BK])
124
b_qg = b_q * exp2(b_gk)
172
- tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
125
+ desc_qg.store([i_t * BT, i_k * BK], b_qg.to(desc_qg.dtype))
126
if STORE_KG:
127
last_idx = min(i_t * BT + BT, T) - 1
128
@@ -180,18 +133,11 @@ def recompute_w_u_fwd_kernel(
133
)
134
b_kg = b_k * exp2(b_gn - b_gk)
135
183
- p_kg = tl.make_block_ptr(
184
- kg + (bos * H + i_h) * K,
185
- (T, K),
186
- (H * K, 1),
187
- (i_t * BT, i_k * BK),
188
- (BT, BK),
189
- (1, 0),
190
- )
191
- tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))
136
+ desc_kg = make_tensor_descriptor(kg + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
137
+ desc_kg.store([i_t * BT, i_k * BK], b_kg.to(desc_kg.dtype))
138
139
b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
194
- tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
140
+ desc_w.store([i_t * BT, i_k * BK], b_w.to(desc_w.dtype))
141
142
143
def recompute_w_u_fwd(
@@ -294,70 +240,32 @@ def chunk_gla_fwd_kernel_o(
240
m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]
241
242
b_o = tl.zeros([BT, BV], dtype=tl.float32)
243
+ desc_q = make_tensor_descriptor(q + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
244
+ desc_g = make_tensor_descriptor(g + (bos * H + i_h) * K, [T, K], [H * K, 1], [BT, BK])
245
+ desc_h = make_tensor_descriptor(h + (i_tg * H + i_h) * K * V, [V, K], [K, 1], [BV, BK])
246
for i_k in range(tl.cdiv(K, BK)):
298
- p_q = tl.make_block_ptr(
299
- q + (bos * H + i_h) * K,
300
- (T, K),
301
- (H * K, 1),
302
- (i_t * BT, i_k * BK),
303
- (BT, BK),
304
- (1, 0),
305
- )
306
- p_g = tl.make_block_ptr(
307
- g + (bos * H + i_h) * K,
308
- (T, K),
309
- (H * K, 1),
310
- (i_t * BT, i_k * BK),
311
- (BT, BK),
312
- (1, 0),
313
- )
314
- p_h = tl.make_block_ptr(
315
- h + (i_tg * H + i_h) * K * V,
316
- (V, K),
317
- (K, 1),
318
- (i_v * BV, i_k * BK),
319
- (BV, BK),
320
- (1, 0),
321
- )
322
-
247
# [BT, BK]
324
- b_q = tl.load(p_q, boundary_check=(0, 1))
248
+ b_q = desc_q.load([i_t * BT, i_k * BK])
249
b_q = (b_q * scale).to(b_q.dtype)
250
# [BT, BK]
327
- b_g = tl.load(p_g, boundary_check=(0, 1))
251
+ b_g = desc_g.load([i_t * BT, i_k * BK])
252
# [BT, BK]
253
b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
254
# [BV, BK]
331
- b_h = tl.load(p_h, boundary_check=(0, 1))
255
+ b_h = desc_h.load([i_v * BV, i_k * BK])
256
# [BT, BV]
257
if i_k >= 0:
258
b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))
335
- p_v = tl.make_block_ptr(
336
- v + (bos * H + i_h) * V,
337
- (T, V),
338
- (H * V, 1),
339
- (i_t * BT, i_v * BV),
340
- (BT, BV),
341
- (1, 0),
342
- )
343
- p_o = tl.make_block_ptr(
344
- o + (bos * H + i_h) * V,
345
- (T, V),
346
- (H * V, 1),
347
- (i_t * BT, i_v * BV),
348
- (BT, BV),
349
- (1, 0),
350
- )
351
- p_A = tl.make_block_ptr(
352
- A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
353
- )
259
+ desc_v = make_tensor_descriptor(v + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV])
260
+ desc_o = make_tensor_descriptor(o + (bos * H + i_h) * V, [T, V], [H * V, 1], [BT, BV])
261
+ desc_A = make_tensor_descriptor(A + (bos * H + i_h) * BT, [T, BT], [H * BT, 1], [BT, BT])
262
# [BT, BV]
355
- b_v = tl.load(p_v, boundary_check=(0, 1))
263
+ b_v = desc_v.load([i_t * BT, i_v * BV])
264
# [BT, BT]
357
- b_A = tl.load(p_A, boundary_check=(0, 1))
265
+ b_A = desc_A.load([i_t * BT, 0])
266
b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
267
b_o += tl.dot(b_A, b_v, allow_tf32=False)
360
- tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
268
+ desc_o.store([i_t * BT, i_v * BV], b_o.to(desc_o.dtype))
269
270
271
def chunk_gla_fwd_o_gk(
@@ -479,34 +387,13 @@ def kda_gate_chunk_cumsum_vector_kernel(
387
388
i_s -= 1
389
482
- p_s = tl.make_block_ptr(
483
- s + (bos * H + i_h) * S,
484
- (T, S),
485
- (H * S, 1),
486
- (i_t * BT, i_s * BS),
487
- (BT, BS),
488
- (1, 0),
489
- )
490
- p_o = tl.make_block_ptr(
491
- o + (bos * H + i_h) * S,
492
- (T, S),
493
- (H * S, 1),
494
- (i_t * BT, i_s * BS),
495
- (BT, BS),
496
- (1, 0),
497
- )
390
+ desc_s = make_tensor_descriptor(s + (bos * H + i_h) * S, [T, S], [H * S, 1], [BT, BS])
391
+ desc_o = make_tensor_descriptor(o + (bos * H + i_h) * S, [T, S], [H * S, 1], [BT, BS])
392
499
- b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)
393
+ b_s = desc_s.load([i_t * BT, i_s * BS]).to(tl.float32)
394
if HAS_BIAS:
501
- p_bias = tl.make_block_ptr(
502
- g_bias + i_h * S,
503
- (S,),
504
- (1,),
505
- (i_s * BS,),
506
- (BS,),
507
- (0,),
508
- )
509
- b_bias = tl.load(p_bias, boundary_check=(0,)).to(tl.float32)
395
+ desc_bias = make_tensor_descriptor(g_bias + i_h * S, [S], [1], [BS])
396
+ b_bias = desc_bias.load([i_s * BS]).to(tl.float32)
397
b_s += b_bias[None, :]
398
399
b_a = tl.exp(tl.load(A_log + i_h).to(tl.float32))
@@ -524,7 +411,7 @@ def kda_gate_chunk_cumsum_vector_kernel(
411
# Boundary loads return zero, but bias and gate activation can make padded
412
# rows nonzero. Padding trails valid rows, so it only affects masked stores.
413
b_o = tl.cumsum(b_gate, axis=0) * cumsum_scale
527
- tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
414
+ desc_o.store([i_t * BT, i_s * BS], b_o.to(desc_o.dtype))
415
416
417
def fused_kda_gate_chunk_cumsum(
@@ -862,25 +749,10 @@ def kda_gate_fwd_kernel(
749
stride_row = H * D
750
stride_col = 1
751
865
- g_ptr = tl.make_block_ptr(
866
- base=g + i_h * D,
867
- shape=(T, D),
868
- strides=(stride_row, stride_col),
869
- offsets=(n_t, 0),
870
- block_shape=(BT, BD),
871
- order=(1, 0),
872
- )
873
-
874
- y_ptr = tl.make_block_ptr(
875
- base=y + i_h * D,
876
- shape=(T, D),
877
- strides=(stride_row, stride_col),
878
- offsets=(n_t, 0),
879
- block_shape=(BT, BD),
880
- order=(1, 0),
881
- )
752
+ desc_g = make_tensor_descriptor(g + i_h * D, [T, D], [stride_row, stride_col], [BT, BD])
753
+ desc_y = make_tensor_descriptor(y + i_h * D, [T, D], [stride_row, stride_col], [BT, BD])
754
883
- b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)
755
+ b_g = desc_g.load([n_t, 0]).to(tl.float32)
756
757
if HAS_BIAS:
758
n_d = tl.arange(0, BD)
@@ -898,7 +770,7 @@ def kda_gate_fwd_kernel(
770
sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))
771
b_y = -b_a * sp
772
901
- tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))
773
+ desc_y.store([n_t, 0], b_y.to(desc_y.dtype))
774
775
776
def fused_kda_gate(
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra.py
+59
-165
@@ -12,7 +12,11 @@ import torch
12
13
from vllm.platforms import current_platform
14
from vllm.third_party.flash_linear_attention.ops.index import prepare_chunk_indices
15
-from vllm.third_party.flash_linear_attention.ops.op import exp2, gather
15
+from vllm.third_party.flash_linear_attention.ops.op import (
16
+ exp2,
17
+ gather,
18
+ make_tensor_descriptor,
19
+)
20
from vllm.third_party.flash_linear_attention.ops.utils import is_gather_supported
21
from vllm.triton_utils import tl, triton
22
@@ -125,33 +129,21 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
129
################################################################################
130
# off-diagonal blocks
131
################################################################################
132
+ desc_qk = make_tensor_descriptor(q, [T, K], [H * K, 1], [BC, BK])
133
+ desc_kk = make_tensor_descriptor(k, [T, K], [H * K, 1], [BC, BK])
134
+ desc_gk = make_tensor_descriptor(g, [T, K], [HV * K, 1], [BC, BK])
135
for i_k in range(tl.cdiv(K, BK)):
136
o_k = i_k * BK + tl.arange(0, BK)
137
m_k = o_k < K
138
132
- p_k0 = tl.make_block_ptr(
133
- k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
134
- )
135
- p_g0 = tl.make_block_ptr(
136
- g, (T, K), (HV * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
137
- )
138
- b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32)
139
- b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32)
139
+ b_k0 = desc_kk.load([i_tc0, i_k * BK]).to(tl.float32)
140
+ b_g0 = desc_gk.load([i_tc0, i_k * BK]).to(tl.float32)
141
142
if i_tc1 < T:
142
- p_q1 = tl.make_block_ptr(
143
- q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
144
- )
145
- p_k1 = tl.make_block_ptr(
146
- k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
147
- )
148
- p_g1 = tl.make_block_ptr(
149
- g, (T, K), (HV * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
150
- )
143
# [BC, BK]
152
- b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32)
153
- b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32)
154
- b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32)
144
+ b_q1 = desc_qk.load([i_tc1, i_k * BK]).to(tl.float32)
145
+ b_k1 = desc_kk.load([i_tc1, i_k * BK]).to(tl.float32)
146
+ b_g1 = desc_gk.load([i_tc1, i_k * BK]).to(tl.float32)
147
# [BK]
148
b_gn1 = tl.load(g + i_tc1 * HV * K + o_k, mask=m_k, other=0).to(tl.float32)
149
# [BC, BK]
@@ -163,19 +155,10 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
155
b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
156
157
if i_tc2 < T:
166
- p_q2 = tl.make_block_ptr(
167
- q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
168
- )
169
- p_k2 = tl.make_block_ptr(
170
- k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
171
- )
172
- p_g2 = tl.make_block_ptr(
173
- g, (T, K), (HV * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
174
- )
158
# [BC, BK]
176
- b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32)
177
- b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32)
178
- b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32)
159
+ b_q2 = desc_qk.load([i_tc2, i_k * BK]).to(tl.float32)
160
+ b_k2 = desc_kk.load([i_tc2, i_k * BK]).to(tl.float32)
161
+ b_g2 = desc_gk.load([i_tc2, i_k * BK]).to(tl.float32)
162
# [BK]
163
b_gn2 = tl.load(g + i_tc2 * HV * K + o_k, mask=m_k, other=0).to(
164
tl.float32
@@ -195,19 +178,10 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
178
b_Akk21 += tl.dot(b_kg2, b_kgt)
179
180
if i_tc3 < T:
198
- p_q3 = tl.make_block_ptr(
199
- q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
200
- )
201
- p_k3 = tl.make_block_ptr(
202
- k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
203
- )
204
- p_g3 = tl.make_block_ptr(
205
- g, (T, K), (HV * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
206
- )
181
# [BC, BK]
208
- b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32)
209
- b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32)
210
- b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32)
182
+ b_q3 = desc_qk.load([i_tc3, i_k * BK]).to(tl.float32)
183
+ b_k3 = desc_kk.load([i_tc3, i_k * BK]).to(tl.float32)
184
+ b_g3 = desc_gk.load([i_tc3, i_k * BK]).to(tl.float32)
185
# [BK]
186
b_gn3 = tl.load(g + i_tc3 * HV * K + o_k, mask=m_k, other=0).to(
187
tl.float32
@@ -235,83 +209,36 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
209
################################################################################
210
# save off-diagonal Aqk blocks and prepare Akk
211
################################################################################
212
+ desc_Aqk = make_tensor_descriptor(Aqk, [T, BT], [HV * BT, 1], [BC, BC])
213
+ p_beta = beta + bos * HV + i_hv
214
+ o_bc = tl.arange(0, BC)
215
if i_tc1 < T:
239
- p_Aqk10 = tl.make_block_ptr(
240
- Aqk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
241
- )
242
- tl.store(
243
- p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
244
- )
216
+ desc_Aqk.store([i_tc1, 0], (b_Aqk10 * scale).to(desc_Aqk.dtype))
217
246
- p_b1 = tl.make_block_ptr(
247
- beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)
248
- )
249
- b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
218
+ b_b1 = tl.load(p_beta + (i_tc1 + o_bc) * HV, mask=i_tc1 + o_bc < T, other=0).to(tl.float32)
219
b_Akk10 = b_Akk10 * b_b1[:, None]
220
if i_tc2 < T:
252
- p_Aqk20 = tl.make_block_ptr(
253
- Aqk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
254
- )
255
- p_Aqk21 = tl.make_block_ptr(
256
- Aqk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
257
- )
258
- tl.store(
259
- p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
260
- )
261
- tl.store(
262
- p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
263
- )
221
+ desc_Aqk.store([i_tc2, 0], (b_Aqk20 * scale).to(desc_Aqk.dtype))
222
+ desc_Aqk.store([i_tc2, BC], (b_Aqk21 * scale).to(desc_Aqk.dtype))
223
265
- p_b2 = tl.make_block_ptr(
266
- beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)
267
- )
268
- b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
224
+ b_b2 = tl.load(p_beta + (i_tc2 + o_bc) * HV, mask=i_tc2 + o_bc < T, other=0).to(tl.float32)
225
b_Akk20 = b_Akk20 * b_b2[:, None]
226
b_Akk21 = b_Akk21 * b_b2[:, None]
227
if i_tc3 < T:
272
- p_Aqk30 = tl.make_block_ptr(
273
- Aqk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
274
- )
275
- p_Aqk31 = tl.make_block_ptr(
276
- Aqk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
277
- )
278
- p_Aqk32 = tl.make_block_ptr(
279
- Aqk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
280
- )
281
- tl.store(
282
- p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
283
- )
284
- tl.store(
285
- p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
286
- )
287
- tl.store(
288
- p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)
289
- )
228
+ desc_Aqk.store([i_tc3, 0], (b_Aqk30 * scale).to(desc_Aqk.dtype))
229
+ desc_Aqk.store([i_tc3, BC], (b_Aqk31 * scale).to(desc_Aqk.dtype))
230
+ desc_Aqk.store([i_tc3, 2 * BC], (b_Aqk32 * scale).to(desc_Aqk.dtype))
231
291
- p_b3 = tl.make_block_ptr(
292
- beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)
293
- )
294
- b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
232
+ b_b3 = tl.load(p_beta + (i_tc3 + o_bc) * HV, mask=i_tc3 + o_bc < T, other=0).to(tl.float32)
233
b_Akk30 = b_Akk30 * b_b3[:, None]
234
b_Akk31 = b_Akk31 * b_b3[:, None]
235
b_Akk32 = b_Akk32 * b_b3[:, None]
236
299
- p_Akk00 = tl.make_block_ptr(
300
- Akkd, (T, BC), (HV * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)
301
- )
302
- p_Akk11 = tl.make_block_ptr(
303
- Akkd, (T, BC), (HV * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)
304
- )
305
- p_Akk22 = tl.make_block_ptr(
306
- Akkd, (T, BC), (HV * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)
307
- )
308
- p_Akk33 = tl.make_block_ptr(
309
- Akkd, (T, BC), (HV * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)
310
- )
311
- b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32)
312
- b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32)
313
- b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32)
314
- b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32)
237
+ desc_Akkd = make_tensor_descriptor(Akkd, [T, BC], [HV * BC, 1], [BC, BC])
238
+ b_Ai00 = desc_Akkd.load([i_tc0, 0]).to(tl.float32)
239
+ b_Ai11 = desc_Akkd.load([i_tc1, 0]).to(tl.float32)
240
+ b_Ai22 = desc_Akkd.load([i_tc2, 0]).to(tl.float32)
241
+ b_Ai33 = desc_Akkd.load([i_tc3, 0]).to(tl.float32)
242
243
################################################################################
244
# forward substitution on diagonals
@@ -397,47 +324,18 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
324
# store full Akk_inv to Akk
325
################################################################################
326
400
- p_Akk00 = tl.make_block_ptr(
401
- Akk, (T, BT), (HV * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)
402
- )
403
- p_Akk10 = tl.make_block_ptr(
404
- Akk, (T, BT), (HV * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)
405
- )
406
- p_Akk11 = tl.make_block_ptr(
407
- Akk, (T, BT), (HV * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)
408
- )
409
- p_Akk20 = tl.make_block_ptr(
410
- Akk, (T, BT), (HV * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)
411
- )
412
- p_Akk21 = tl.make_block_ptr(
413
- Akk, (T, BT), (HV * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)
414
- )
415
- p_Akk22 = tl.make_block_ptr(
416
- Akk, (T, BT), (HV * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)
417
- )
418
- p_Akk30 = tl.make_block_ptr(
419
- Akk, (T, BT), (HV * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)
420
- )
421
- p_Akk31 = tl.make_block_ptr(
422
- Akk, (T, BT), (HV * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)
423
- )
424
- p_Akk32 = tl.make_block_ptr(
425
- Akk, (T, BT), (HV * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
426
- )
427
- p_Akk33 = tl.make_block_ptr(
428
- Akk, (T, BT), (HV * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)
429
- )
327
+ desc_Akk = make_tensor_descriptor(Akk, [T, BT], [HV * BT, 1], [BC, BC])
328
431
- tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1))
432
- tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1))
433
- tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1))
434
- tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1))
435
- tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1))
436
- tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1))
437
- tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1))
438
- tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1))
439
- tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1))
440
- tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1))
329
+ desc_Akk.store([i_tc0, 0], b_Ai00.to(desc_Akk.dtype))
330
+ desc_Akk.store([i_tc1, 0], b_Ai10.to(desc_Akk.dtype))
331
+ desc_Akk.store([i_tc1, BC], b_Ai11.to(desc_Akk.dtype))
332
+ desc_Akk.store([i_tc2, 0], b_Ai20.to(desc_Akk.dtype))
333
+ desc_Akk.store([i_tc2, BC], b_Ai21.to(desc_Akk.dtype))
334
+ desc_Akk.store([i_tc2, 2 * BC], b_Ai22.to(desc_Akk.dtype))
335
+ desc_Akk.store([i_tc3, 0], b_Ai30.to(desc_Akk.dtype))
336
+ desc_Akk.store([i_tc3, BC], b_Ai31.to(desc_Akk.dtype))
337
+ desc_Akk.store([i_tc3, 2 * BC], b_Ai32.to(desc_Akk.dtype))
338
+ desc_Akk.store([i_tc3, 3 * BC], b_Ai33.to(desc_Akk.dtype))
339
340
341
@triton.heuristics(
@@ -505,16 +403,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
403
Aqk = Aqk + (bos * HV + i_hv) * BT
404
Akk = Akk + (bos * HV + i_hv) * BC
405
508
- p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
509
- p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0))
510
- p_g = tl.make_block_ptr(g, (T, K), (HV * K, 1), (i_ti, 0), (BC, BK), (1, 0))
406
+ desc_q = make_tensor_descriptor(q, [T, K], [H * K, 1], [BC, BK])
407
+ desc_k = make_tensor_descriptor(k, [T, K], [H * K, 1], [BC, BK])
408
+ desc_g = make_tensor_descriptor(g, [T, K], [HV * K, 1], [BC, BK])
409
512
- p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,))
513
-
514
- b_q = tl.load(p_q, boundary_check=(0, 1))
515
- b_k = tl.load(p_k, boundary_check=(0, 1))
516
- b_g = tl.load(p_g, boundary_check=(0, 1))
517
- b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
410
+ b_q = desc_q.load([i_ti, 0])
411
+ b_k = desc_k.load([i_ti, 0])
412
+ b_g = desc_g.load([i_ti, 0])
413
+ b_beta = tl.load(beta + o_c * HV, mask=m_c, other=0).to(tl.float32)
414
415
if USE_GATHER:
416
b_gn = gather(
@@ -546,12 +442,10 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
442
b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0)
443
b_Akk = tl.where(m_Akk, b_Akk, 0.0)
444
549
- p_Aqk = tl.make_block_ptr(
550
- Aqk, (T, BT), (HV * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)
551
- )
552
- p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV * BC, 1), (i_ti, 0), (BC, BC), (1, 0))
553
- tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))
554
- tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1))
445
+ desc_Aqk = make_tensor_descriptor(Aqk, [T, BT], [HV * BT, 1], [BC, BC])
446
+ desc_Akk = make_tensor_descriptor(Akk, [T, BC], [HV * BC, 1], [BC, BC])
447
+ desc_Aqk.store([i_ti, i_i * BC], b_Aqk.to(desc_Aqk.dtype))
448
+ desc_Akk.store([i_ti, 0], b_Akk.to(desc_Akk.dtype))
449
450
tl.debug_barrier()
451
@@ -566,7 +460,7 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
460
b_a += tl.sum(b_a[:, None] * b_Ai, 0)
461
b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai)
462
b_Ai += m_I
569
- tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1))
463
+ desc_Akk.store([i_ti, 0], b_Ai.to(desc_Akk.dtype))
464
465
466
def chunk_kda_fwd_intra(
vllm/models/kimi_k3/amd/ops/third_party/kda/chunk_intra_token_parallel.py
+6
-11
@@ -12,7 +12,7 @@
12
13
import torch
14
15
-from vllm.third_party.flash_linear_attention.ops.op import exp2
15
+from vllm.third_party.flash_linear_attention.ops.op import exp2, make_tensor_descriptor
16
from vllm.triton_utils import tl, triton
17
18
@@ -108,20 +108,15 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
108
b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
109
110
# g: [B, T, HV, K], beta: [B, T, HV]
111
- p_g = tl.make_block_ptr(
112
- g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
113
- )
114
- p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,))
115
- b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
116
- b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32)
111
+ desc_g = make_tensor_descriptor(g + i_t * HV * K, [HV, K], [K, 1], [BH, BK])
112
+ b_g = desc_g.load([i_hg * BH, 0]).to(tl.float32)
113
+ b_beta = tl.load(beta + i_t * HV + o_hv, mask=m_hv, other=0).to(tl.float32)
114
b_k *= b_beta[:, None]
115
116
for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))):
117
b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32)
121
- p_gj = tl.make_block_ptr(
122
- g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)
123
- )
124
- b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32)
118
+ desc_gj = make_tensor_descriptor(g + j * HV * K, [HV, K], [K, 1], [BH, BK])
119
+ b_gj = desc_gj.load([i_hg * BH, 0]).to(tl.float32)
120
121
b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0)
122
b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale