背景：

在 kvoffload 这个版本的 sglang 的推理的投机采样（Speculative Sampling）阶段，引擎需要对生成的 TopK 候选 Token 索引进行去重，这就引入了本次的优化算子。

算子语义：

• 对每个 batch 将 mtp_step 个 top‑k index 行做去重，输出长度固定为 mtp_step * k，不足用 -1 填充

场景特征：

• 小 Batch：bs=115，GPU 并行度（Occupancy）天然不足。
• 延迟敏感：这是在线服务链路的一环，对 Latency 极度敏感。
• 输入规模：固定几档（Total elements = 2048, 4096...），适合特化。

我们经历了一个从“Naive 实现”到“指令级 Hacking”的完整过程，最终将耗时压缩了 120倍 ，从 2ms 压缩至 16us。

基线：直觉的陷阱 (Shared Hash Table)

当我们拿到“去重”这个需求时，第一反应通常是：哈希表。这在 CPU 上是标准答案（$O(N)$）。于是前人就写出了第一版 Kernel，在 Shared Memory 中维护一个哈希表。

代码片段 (Naive Implementation)

__global__ void deduplicate_topk_kernel_v0(
    const int* topk_indices, 
    int* topk_indices_spec, 
    int total_elements, 
    int k) 
{
    // ... setup shared memory ...
    extern __shared__ int shared_mem[];
    int* seen = shared_mem; // 哈希表
    // 初始化哈希表
    for (int i = tid; i < total_elements; i += blockDim.x) seen[i] = -1;
    __syncthreads();
    // 循环处理每个元素
    for (int i = tid; i < total_elements; i += blockDim.x) {
        int val = topk_indices[row * k + col];
        // 简单的线性探测哈希
        int hash_idx = val % total_elements;
        // 【致命点1】探测循环：可能 1 次成功，也可能 10 次
        // 导致 Warp 内线程严重发散 (Divergence)
        for (int probe = 0; probe < total_elements; probe++) {
            int idx = (hash_idx + probe) % total_elements;
            // 【致命点2】原子操作竞争：所有线程抢着写
            int old = atomicCAS(&seen[idx], -1, val);
            if (old == -1) { // 抢到了，是新元素
                int pos = atomicAdd(unique_count, 1);
                unique_list[pos] = val; // 【致命点3】随机写，非合并访存
                break;
            } else if (old == val) {
                break; // 已存在
            }
        }
    }
}

为什么慢？

NCU 的报告显示：

1. Warp Divergence (发散)：哈希冲突是随机的。Warp 里有的线程 1 次 probe 就结束了，有的线程要 probe 10 次。根据 SIMT 机制，整个 Warp 必须陪着最慢的那个线程空转。
2. Memory Coalescing (非合并)：哈希表的读写是完全随机的，显存带宽利用率（Memory Throughput）只有可怜的 4%。
3. Atomic Contention (原子竞争)：数百个线程同时轰炸 atomicCAS，硬件序列化严重。

初步结论：在 GPU 上，“确定的执行路径” 往往比理论上的算法复杂度更重要。哈希表在 GPU 上通常不是好选择。同时，考虑到推理阶段的特性，输入的 2048 个索引中，很可能重复的索引并不多，导致了哈希表的冲突比理想情况糟糕得多。