[Kimi Bug] Fix kda ima `Triton Error [CUDA]: an illegal memory access was encountered` (#55924)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 8, 2026 at 17:07 UTC
bfb443a6b6f670e68e211112a141d089f1cf956f
1 file changed
+1
-1
vllm/models/kimi_k3/nvidia/kda.py
+1
-1
@@ -415,7 +415,7 @@ def _store_cache_checkpoints_kernel(
415
# store checkpoints to cache
416
seq_idx = tl.program_id(0)
417
cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
418
- state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx)
418
+ state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx).to(tl.int64)
419
checkpoint_offset = tl.load(
420
checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride
421
)