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 13, 05:45 UTC
[Bugfix][CI/Test][Spec Decode] Fix illegal memory access in offline_inference/spec_decode.py (Issue 27619) (#28432)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
rasmith committed
Nov 14, 2025 at 00:34 UTC
15ae8e0784d3889c6aa2c487ca00df4e3fde6f44
1 file changed
+4
-2
vllm/attention/ops/triton_reshape_and_cache_flash.py
+4
-2
@@ -97,7 +97,6 @@ def triton_reshape_and_cache_flash(
97
k_scale: torch.Tensor, # float32
98
v_scale: torch.Tensor, # float32
99
):
100
- num_tokens = key.shape[0]
100
num_heads = key.shape[1]
101
head_size = key.shape[2]
102
block_size = key_cache.shape[1]
@@ -155,7 +154,10 @@ def triton_reshape_and_cache_flash(
154
155
# TODO(ngl): maybe replace with static launch grid to avoid overhead if
156
# using cudagraphs
158
- grid = lambda meta: (int(num_tokens), triton.cdiv(n, meta["TILE_SIZE"]))
157
+ grid = lambda meta: (
158
+ slot_mapping.shape[0],
159
+ triton.cdiv(n, meta["TILE_SIZE"]),
160
+ )
161
162
reshape_and_cache_kernel_flash[grid](
163
key_ptr=key,