执行摘要

该PR修复了在DSA（Direct Sparse Attention）与MTP（Multi-Token Prediction）同时启用时出现的IMA（Invalid Memory Access）崩溃。根本原因是上一轮性能优化（#40654）为了消除GPU→CPU同步，使kernel中的num_tokens参数变为内存分配上界（可能大于实际有效token数），导致kernel中访问无效的批次索引。

功能与动机

PR #40654 引入了一个重要的性能优化：避免GPU→CPU同步获取精确的seq_lens，而是使用一个上界seq_lens_cpu_upper_bound。然而，这在DSA+MTP场景下导致cp_gather_indexer_k_quant_cache_kernel中的某些线程访问了无效的batch_idx，因为该上界可能大于实际批次中的有效序列数。

实现拆解

修复仅涉及一个文件：csrc/cache_kernels.cu，核心内核函数cp_gather_indexer_k_quant_cache_kernel。

1. 共享数组初始化：将batch_idx共享内存数组初始化为-1，表示无效批次。此操作由每个warp的第一个线程（threadIdx.x == 0）执行。
2. 全局同步增强：将原有的条件__syncwarp()（仅在非ROCm环境生效）替换为全块同步__syncthreads()，确保所有线程的batch_idx写入完成后再进行后续判断。
3. 边界检查：在访问batch_idx后，增加batch < 0的检查，如果批次索引无效则直接返回，避免后续的cu_seq_lens[batch]越界访问。

csrc/cache_kernels.cu

核心内核函数 cp_gather_indexer_k_quant_cache_kernel 的修复，解决了DSA+MTP场景下因num_tokens上界导致的越界访问。

说明: 该kernel原先假设所有线程都能找到有效的batch索引，但num_tokens上界可能导致部分线程访问无效batch。通过初始化共享数组为-1，并在访问前检查batch<0，优雅地跳过无效线程。

关键源码片段

csrc/cache_kernels.cu

核心内核函数 cp_gather_indexer_k_quant_cache_kernel 的修复，解决了DSA+MTP场景下因num_tokens上界导致的越界访问。

__global__ void cp_gather_indexer_k_quant_cache_kernel(...) {
    // ...
    __shared__ int batch_idx[BLOCK_Y_SIZE];
    if (threadIdx.x == 0) {
        batch_idx[threadIdx.y] = -1; // 初始化为无效值，防止未更新时使用
    }
    __syncthreads();
​
    for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x)); iter++) {
        int tid = iter * blockDim.x + threadIdx.x;
        if (tid < batch_size) {
            // 某个线程负责写入 batch_idx
            batch_idx[threadIdx.y] = /* 计算 */ ;
        }
    }
    __syncthreads(); // 确保所有线程的 batch_idx 已更新
​
    // num_tokens 可能为分配上界，需校验 batch 有效性
    const int batch = batch_idx[threadIdx.y];
    if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
        return; // batch<0 表示该线程负责的批次索引未初始化，跳过
    }
    // 使用安全的 batch 访问后续数据
    const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
    // ...
}

说明: 该kernel原先假设所有线程都能找到有效的batch索引，但num_tokens上界可能导致部分线程访问无效batch。通过初始化共享数组为-1，并在访问前检查batch<0，优雅地跳过无效线程。

评论区精华

无讨论。

风险与影响

• 回归风险：替换__syncwarp()为__syncthreads()可能对非ROCm平台的性能有轻微影响，但逻辑更安全。
• 影响范围：仅影响DSA+MTP用户，修复后此类场景将不再触发IMA崩溃。其他场景无影响。

关联脉络

该PR是PR #40654（[Core] Avoid seq_lens_cpu GPU->CPU sync）的补丁，展示了性能优化引入的副作用如何被修复。它属于speculative decoding功能线的持续演进。