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 4, 17:30 UTC
[Bugfix] Fix IMA in DSA + MTP (#40772)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Apr 24, 2026 at 01:40 UTC
7d3195ea9fc88e31131099d2d2122fe38558a87a
1 file changed
+14
-7
csrc/cache_kernels.cu
+14
-7
@@ -599,6 +599,11 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
599
const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
600
// Find batch index within a block
601
__shared__ int batch_idx[BLOCK_Y_SIZE];
602
+ if (threadIdx.x == 0) {
603
+ batch_idx[threadIdx.y] = -1;
604
+ }
605
+ __syncthreads();
606
+
607
for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
608
iter++) {
609
int tid = iter * blockDim.x + threadIdx.x;
@@ -611,16 +616,18 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
616
}
617
}
618
614
-#ifndef USE_ROCM
615
- __syncwarp();
616
-#endif
619
+ __syncthreads();
620
618
- if (head_idx >= head_dim || token_idx >= num_tokens) {
621
+ // num_tokens may be an allocation upper bound when Python avoids a D2H sync.
622
+ // Only tokens covered by the exact device-side cu_seq_lens are valid to
623
+ // gather.
624
+ const int batch = batch_idx[threadIdx.y];
625
+ if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
626
return;
627
}
621
- const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
622
- const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
623
- inbatch_seq_idx / cache_block_size];
628
+ const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
629
+ const int block_idx =
630
+ block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
631
const int64_t src_block_offset = block_idx * block_stride;
632
const int64_t cache_inblock_offset =
633
(inbatch_seq_idx % cache_block_size) * head_dim + head_idx;