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 5, 11:00 UTC
[CI][AMD][BugFix][Kernel] Cast induction variable to int64 on MI350 for chunk_gated_delta_rule_fwd_kernel_h_blockdim64 to avoid illegal memory access (#39087)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Apr 8, 2026 at 03:57 UTC
78434b923c80e435bcae9ad846471a48d8e3bb4e
1 file changed
+26
-6
vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+26
-6
@@ -129,22 +129,42 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
129
# main recurrence
130
for i_t in range(NT):
131
p_h1 = tl.make_block_ptr(
132
- h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)
132
+ h + i_t.to(tl.int64) * stride_h,
133
+ (V, K),
134
+ (K, 1),
135
+ (i_v * BV, 0),
136
+ (BV, 64),
137
+ (1, 0),
138
)
139
tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))
140
if K > 64:
141
p_h2 = tl.make_block_ptr(
137
- h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)
142
+ h + i_t.to(tl.int64) * stride_h,
143
+ (V, K),
144
+ (K, 1),
145
+ (i_v * BV, 64),
146
+ (BV, 64),
147
+ (1, 0),
148
)
149
tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1))
150
if K > 128:
151
p_h3 = tl.make_block_ptr(
142
- h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)
152
+ h + i_t.to(tl.int64) * stride_h,
153
+ (V, K),
154
+ (K, 1),
155
+ (i_v * BV, 128),
156
+ (BV, 64),
157
+ (1, 0),
158
)
159
tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1))
160
if K > 192:
161
p_h4 = tl.make_block_ptr(
147
- h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)
162
+ h + i_t.to(tl.int64) * stride_h,
163
+ (V, K),
164
+ (K, 1),
165
+ (i_v * BV, 192),
166
+ (BV, 64),
167
+ (1, 0),
168
)
169
tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
170
@@ -182,9 +202,9 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
202
)
203
tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1))
204
185
- last_idx = min((i_t + 1) * BT, T) - 1
205
+ last_idx = min((i_t.to(tl.int64) + 1) * BT, T) - 1
206
if USE_G:
187
- m_t = (i_t * BT + tl.arange(0, BT)) < T
207
+ m_t = (i_t.to(tl.int64) * BT + tl.arange(0, BT)) < T
208
b_g_last = tl.load(g + bos * H + last_idx * H + i_h)
209
p_g = tl.make_block_ptr(
210
g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)