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 12, 18:30 UTC
[CI/Build][Kernel][AMD] Move extra dim to after load in _fwd_kv_parallel in lighting_attn.py (#29132)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Nov 21, 2025 at 10:53 UTC
e99e467384001e284e0722a33362866b10fed65b
1 file changed
+7
-1
vllm/model_executor/layers/lightning_attn.py
+7
-1
@@ -198,7 +198,7 @@ def _fwd_kv_parallel(
198
)
199
200
# Load the decay factors for the current head and block
201
- k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)[None, :]
201
+ k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)
202
203
kv_index = tl.arange(0, CBLOCK)
204
@@ -228,6 +228,12 @@ def _fwd_kv_parallel(
228
229
# Load decay factor and compute weighted key-value outer product
230
k_decay = tl.load(k_decay_ptr)
231
+
232
+ # NOTE: Need to add the extra dim here due to AMD MLIR lowering error.
233
+ # Please don't move it back until issue is resolved.
234
+ # Issue: https://github.com/ROCm/triton/issues/907
235
+ k_decay = k_decay[None, :]
236
+
237
kv += tl.dot(k_trans * k_decay, v)
238
239
# Move to the next sub-block