__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; // 已存在
            }
        }
    }
}

// 寄存器大户：CUB 需要一个数组来存每个元素的状态
int unique_flags[kItemsPerThread]; // Flag: 1=Unique, 0=Duplicate
// 1. 排序
BlockRadixSortT(temp_storage).Sort(values);
// 2. 找不同 (FlagHeads)
// CUB 自动处理了跨线程的边界比较，结果存入 unique_flags 数组
BlockDiscontinuityT(temp_storage).FlagHeads(unique_flags, values, NotEqual());
// 3. 计数
int unique_count = 0;
#pragma unroll
for (int i = 0; i < kItemsPerThread; ++i) unique_count += unique_flags[i];
// 4. 前缀和计算写入偏移 (Offset)
BlockScanT(temp_storage).ExclusiveSum(unique_count, thread_offset, total_unique);
// 5. 写回 (Scatter)
#pragma unroll
for (int i = 0; i < kItemsPerThread; ++i) {
    if (unique_flags[i]) {
        // 依然依赖 flags 数组做判断
        output[thread_offset + local_offset++] = values[i];
    }
}

// Before: 标量读取 (Scalar Load)
// 产生大量 LDG.E.32 指令，带宽利用率低
int values[kItemsPerThread];
for (int i = 0; i < kItemsPerThread; ++i) {
    values[i] = input[tid + i * blockDim.x];
}
// After: 向量化读取 (Vectorized Load)
// 产生 LDG.E.128 指令，一次搬运 16 字节
const int4* vec_ptr = reinterpret_cast<const int4*>(input);
int4 loaded = vec_ptr[tid + i * blockDim.x]; 
values[0] = loaded.x; values[1] = loaded.y; ...

// Before: 直接离散写回 GMEM (低效)
// 线程写的位置跳跃，无法合并，带宽浪费严重
if (is_unique) {
    gmem_output[global_offset + local_idx] = val; 
}
// After: SMEM 中转 + 向量化写回 (高效)
// Step 1: 先写到 SMEM (片上极速)
if (is_unique) {
    smem_buffer[local_idx] = val;
}
__syncthreads();
// Step 2: 连续、向量化写回 GMEM
// 此时 idx 是连续的 (tid * 4)，可以生成高效的 STG.128 指令
int4* out_vec = reinterpret_cast<int4*>(gmem_output);
const int4* smem_vec = reinterpret_cast<const int4*>(smem_buffer);
out_vec[tid] = smem_vec[tid]; // 完美合并写

// 【优化点1】砍掉数组，只用 1 个寄存器
uint32_t unique_mask = 0; 
// 手写 Discontinuity 逻辑，替代 CUB
// ... (省略跨线程边界交换代码) ...
// 生成掩码：如果当前元素 != 上一个，则置位
#pragma unroll
for (int i = 0; i < kItemsPerThread; ++i) {
    if (val != prev && val != sentinel) {
        unique_mask |= (1u << i); // 设置第 i 位为 1
    }
}
// 【优化点2】用 __popc 替代循环计数
int unique_count = __popc(unique_mask);
BlockScanT(temp_storage).ExclusiveSum(unique_count, thread_offset, total_unique);
// 【优化点3】极速 Scatter
#pragma unroll
for (int i = 0; i < kItemsPerThread; ++i) {
    // 检查第 i 位是否为 1
    if (unique_mask & (1u << i)) {
        // Magic: 计算低 i 位有多少个 1，直接得到局部偏移
        // (1u << i) - 1 生成形如 000...011111 的掩码
        int local_idx = __popc(unique_mask & ((1u << i) - 1));
        // 直接写入 Shared Memory，无循环依赖
        output_buffer[thread_offset + local_idx] = values[i];
    }
}

#!/usr/bin/env python3
import argparse
import sys
from pathlib import Path
def _add_repo_path() -> None:
    script_dir = Path(__file__).resolve()
    repo_root = script_dir.parents[2]
    python_root = repo_root / "BAIDU_REPO" / "aiak_sglang_offload" / "python"
    sys.path.insert(0, str(python_root))
def main() -> None:
    parser = argparse.ArgumentParser(description="Benchmark deduplicate_topk CUDA op.")
    parser.add_argument("--bs", type=int, default=115)
    parser.add_argument("--mtp-step", type=int, choices=[1, 2, 3, 4], default=2)
    parser.add_argument("--k", type=int, default=2048)
    parser.add_argument("--iters", type=int, default=10)
    parser.add_argument("--warmup", type=int, default=1)
    parser.add_argument("--vocab-size", type=int, default=1 << 31)
    parser.add_argument(
        "--ncu",
        action="store_true",
        help="Run a single kernel for ncu capture",
    )
    args = parser.parse_args()
    _add_repo_path()
    import torch
    from sglang.srt.utils_op.kv_offload.deduplicate_topk import DeduplicateTopk
    total = args.bs * args.mtp_step
    x = torch.randint(
        0,
        args.vocab_size,
        (total, args.k),
        device="cuda",
        dtype=torch.int32,
    )
    op = DeduplicateTopk()
    if args.ncu:
        _ = op.cuda_impl(x, args.mtp_step)
        torch.cuda.synchronize()
        print(f"bs={args.bs} mtp_step={args.mtp_step} k={args.k}")
        return
    for _ in range(args.warmup):
        _ = op.cuda_impl(x, args.mtp_step)
    torch.cuda.synchronize()
    y_cuda = op.cuda_impl(x, args.mtp_step)
    y_aten = op.aten_impl(x, args.mtp_step)
    torch.cuda.synchronize()
    is_match = bool(torch.equal(y_cuda, y_aten))
    print(f"check_equal={is_match}")
    torch.cuda.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        _ = op.cuda_impl(x, args.mtp_step)
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(args.iters):
        graph.replay()
    end.record()
    torch.cuda.synchronize()
    avg_ms = start.elapsed_time(end) / args.iters
    print(f"bs={args.bs} mtp_step={args.mtp_step} k={args.k} total={total}")
    print(f"avg_ms_cuda_graph={avg_ms:.4f}")
if __name__ == "__main__":
    main()