第 12 课：ArgMax 对归约与确定性

本课目标与完成标准

学完后你应能：归约 (value,index) 对，处理相等值、NaN 和两阶段输出。

通过必须同时满足：

独立补齐本课唯一的代码填空题，并通过给定检查；
三个问答题均说明因果链，而不是只报术语；
能指出至少一个正确性边界和一个性能取舍；
总分不低于 8/10，且没有一票否决级概念错误。

前置关系

课程前置：C/C++、线性代数、基本并行编程
本课在路线中的作用：ArgMax 不只归约值，还要同步携带索引；这是自定义归约算子的典型模式。

核心心智模型

1. 它是什么，解决什么问题

ArgMax 算子（最大值索引定位） 是模型解码贪婪采样（Greedy Decoding）、Top-1 类别预测以及强化学习离散动作选择的核心基础。不同于普通的 ReduceMax 只关心最大数值本身，ArgMax 要求输出取得最大值的物理全局下标（Index）： $\text{ArgMax}(\mathbf{x}) = \arg\max_{0 \le i < N} x_i = \{i \mid \forall j, \ x_i \ge x_j\}$

并行归约中的确定性挑战（Tie-Breaking 不变量，图 12-1）：

ArgMax 二元元组仲裁

元组 1: (val_1, idx_1)

元组 2: (val_2, idx_2)

val_1 > val_2 ?

胜出: (val_1, idx_1)

胜出: (val_2, idx_2)

min(idx_1, idx_2)

严格保留最小索引 (确定性 Tie-Breaking)

是否: val_2 > val_1等值: val_1 == val_2

图 12-1：ArgMax 键值对元组结构与确定性仲裁规则

在串行 CPU 代码中，通常用 if (x[i] > max_val) { max_val = x[i]; best_idx = i; }，天然具有“遇到相等最大值保留最左侧（最小下标）”的确定性；
但在 GPU 上，成千上万个线程以树状结构并行对偶比较。当输入数组中存在多个相等的最大值（例如分词 Logits 中有多个未见过的 Padding Token 得分相同）时，由于 Warp 调度和 Block 树状归约的无序性，若仅仅通过 val_b > val_a 进行筛选，谁会胜出完全取决于线程树的配对顺序；
确定性仲裁规则（Tie-breaking Policy）：必须将数据结构扩充为双元组 ValIdx { float val; int idx; }，并在比较逻辑中严格引入字典序比较：数值大者胜出；若数值相等，则全局下标更小者胜出。这样结果与比较顺序无关，在不同架构和 Grid 配置下都相同。

2. 它如何工作（结合图 12-1 与图 12-2）

由于需要支持百万到数十亿规模的超大张量，算法采用经典的 两阶段分层归约架构（Two-Stage Grid Reduction）：

阶段一：网格局部归约（Pass 1: Block-level ArgMax）
- 启动包含 num_blocks 个 Block 的网格，各线程以 Grid-Stride 方式遍历整个长数组；
- 每个线程在本地寄存器中维护当前的局部最佳二元状态 ValIdx local = {-FLT_MAX, -1}；
- Warp 级元组洗牌（图 12-2）：

Warp 内部 32 线程

__shfl_down_sync 洗牌

严格仲裁算子

步进 16 -> 8 -> 4 -> 2 -> 1 折半归约

1同时洗牌传输 val 与 idx (或 64 位整型打包)

2比较 pair 与 other_pair

3更新局部胜出元组

Lane 0 获得全局极值与准确最小全局索引

图 12-2：Warp 级确定性 ArgMax 归约时序流

利用 __shfl_down_sync 同时在寄存器间交换 val 与 idx（或打包传输），在 5 步折半归约中通过确定性仲裁逻辑选出当前 Warp 的胜者； - Block 级发布：各 Warp 的 Lane 0 将局部的 ValIdx 写入共享内存 smem，由 Warp 0 再次归约出当前 Block 的最佳结果，最终由 Thread 0 将该 Block 的局部极值写入全局临时数组 partial_val[blockIdx.x] 与 partial_idx[blockIdx.x]。

阶段二：单块全局终局归约（Pass 2: Single-Block Final Reduce）
- 启动仅有 1 个 ThreadBlock 的 Kernel，Block 内线程并行读取 num_blocks 个局部极值候选者；
- 在 Block 内部再次执行树状元组归约，Thread 0 将最终胜者的全局索引单次写回输出指针 *output。

3. 正确性条件与常见误区

相等最大值的非确定性漏洞（Tie-breaking 缺失）：
- 如果比较逻辑仅写为 if (b.val > a.val) a = b;：当 b.val == a.val 时，条件为假，保留 a。但是，在树状归约中，谁作为参数 a、谁作为参数 b 完全受折半步进的硬件偏移量（Offset）支配。例如在 Offset=16 的比较中，Lane 0 比较 Lane 16；在 Offset=8 时，Lane 0 比较 Lane 8。这会导致最终保留下来的不一定是全局最小下标，甚至不同 Grid 大小跑出来的下标完全不同；
- 强制不变量：必须严格书写确定性判据： $b \text{ 胜出} \iff (b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$
NaN 污染策略（NaN Handling Policy）：
- 根据 IEEE 754 规范，任何包含 NaN 的浮点比较（NaN > x、x > NaN、NaN == x）在硬件中全部返回 false；
- 若未对 NaN 进行特判，NaN 永远无法在 > 比较中胜出，可能导致真正的最大值被静默吞掉，或者导致首个遇到的正常数被当作结果；工业级实现（如 PyTorch）通常定义“遇到 NaN 则直接传播 NaN 或优先返回首个 NaN 的索引”。

4. 性能与工程取舍

元组洗牌带宽开销 vs. 整数压缩：
- 普通 ReduceMax 每次洗牌仅需传输 1 个 32-bit 的 float；
- ArgMax 每次洗牌必须传递 val（4 字节）和 idx（4 字节），洗牌指令数翻倍；
- 工程优化：若浮点数为 FP16/BF16（2 字节），或者能保证 val 为非负数，可把浮点数的有序编码与下标打包进一个 64-bit 整数（uint64_t），一次比较同时完成数值与下标的 tie-break。注意 64-bit __shfl 会被编译成两条 32-bit SHFL；若要真正单条，需要打包进 32 位（如 FP16 值 + 16 位下标）。
两阶段 Kernel vs. 原子操作/跨 Block 栅障：
- 采用两阶段独立 Kernel Launch 拥有绝对的硬件可移植性与安全性，中间通过 Global Memory 传递 partial 结果，天然具备 Stream 内存可见性；
- 单 Kernel 协作（使用 Cooperative Groups 网格级同步）可省去二次 Launch 开销，但对 GPU 架构与 Occupancy 调度有严格的硬件配额约束。

具体演示：平局仲裁确定性演算

考虑极小输入数组： $x = [5.0, 7.0, 7.0, 3.0]$ ，其中最大值 $7.0$ 在下标 1 和 2 处重复出现：

无下标仲裁（仅凭 val_b > val_a）：
- 线程 1 拥有 $(val=7.0, idx=1)$ ，线程 2 拥有 $(val=7.0, idx=2)$ ；
- 若在归约折半中由线程 2 比较线程 1（b = 线程1, a = 线程2）： $7.0 > 7.0$ 为假，保留线程 2 的结果，最终输出 $\text{Index} = 2$ ；
- 若换一种网格拓扑或线程配对顺序，由线程 1 比较线程 2（b = 线程2, a = 线程1）： $7.0 > 7.0$ 为假，保留线程 1 的结果，最终输出 $\text{Index} = 1$ ；
- 严重后果：相同的模型权重和输入，在不同的 Batch 大小或不同的 GPU 架构上推理，生成的下一个 Token 竟然完全不同。
严格确定性字典序仲裁：
- 比较逻辑： $(b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$ ；
- 无论二者谁作为参数 $a$ 谁作为参数 $b$ ，下标较小的 $idx=1$ 恒定被判定为胜者；
- 无论如何调度，输出绝对恒为 $\text{Index} = 1$ ，完美捍卫了模型贪婪采样的可复现性（Reproducibility）。

实践任务：唯一代码填空题

补齐确定性比较条件；注意初始化 idx=-1。

规则：只能修改 TODO/______ 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。

Code

%%writefile /tmp/12_argmax.cu
#include <cuda_runtime.h>
#include <float.h>
#include <stdio.h>
#include <stdlib.h>

namespace {

// ArgMax: 找到 input[0:n] 中最大值的索引。
//
// 面试重点：
// 1. reduce 时同时传递 (value, index) 对。
// 2. warp shuffle 传 value 比较，再传 index 选择。
// 3. 两轮 kernel：第一轮每个 block 出一个 (val, idx)，第二轮归并。

struct ValIdx {
    float val;
    int idx;
};

template<int BLOCK_SIZE>
__device__ __forceinline__ ValIdx warp_reduce_max(ValIdx vi) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other_val = __shfl_down_sync(0xffffffff, vi.val, offset);
        int other_idx = __shfl_down_sync(0xffffffff, vi.idx, offset);
        if (______) {  // TODO: 选择更大值；相等时选择更小有效索引
            vi.val = other_val;
            vi.idx = other_idx;
        }
    }
    return vi;
}

template<int BLOCK_SIZE>
__device__ __forceinline__ ValIdx block_reduce_max(ValIdx vi) {
    constexpr int NUM_WARPS = (BLOCK_SIZE + 31) / 32;
    __shared__ float smem_val[NUM_WARPS];
    __shared__ int smem_idx[NUM_WARPS];

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    vi = warp_reduce_max<BLOCK_SIZE>(vi);

    if (lane == 0) {
        smem_val[warp] = vi.val;
        smem_idx[warp] = vi.idx;
    }
    __syncthreads();

    if (warp == 0) {
        vi.val = (lane < NUM_WARPS) ? smem_val[lane] : -FLT_MAX;
        vi.idx = (lane < NUM_WARPS) ? smem_idx[lane] : -1;
        vi = warp_reduce_max<BLOCK_SIZE>(vi);
    }

    return vi;
}

// 第一轮：每个 block 用 grid-stride loop 找到局部最大值及其索引
template<int BLOCK_SIZE>
__global__ void argmax_kernel(
    const float* __restrict__ input,
    float* __restrict__ partial_val,
    int* __restrict__ partial_idx,
    int n
) {
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + tid;
    int stride = blockDim.x * gridDim.x;

    // 每个线程先通过 grid-stride loop 找到自己负责的元素中的最大值
    ValIdx local = {-FLT_MAX, -1};
    for (int i = idx; i < n; i += stride) {
        if (input[i] > local.val) {
            local.val = input[i];
            local.idx = i;
        }
    }

    // block 级 reduce
    ValIdx result = block_reduce_max<BLOCK_SIZE>(local);

    if (tid == 0) {
        partial_val[blockIdx.x] = result.val;
        partial_idx[blockIdx.x] = result.idx;
    }
}

// 第二轮：把所有 block 的 partial 结果归并，只需一个 block
template<int BLOCK_SIZE>
__global__ void argmax_final_kernel(
    const float* __restrict__ partial_val,
    const int* __restrict__ partial_idx,
    int* __restrict__ output,
    int num_blocks
) {
    int tid = threadIdx.x;

    ValIdx local = {-FLT_MAX, -1};
    // 一个 block 处理所有 partial，stride = BLOCK_SIZE
    for (int i = tid; i < num_blocks; i += BLOCK_SIZE) {
        // 与 warp/block 归约使用同一 tie-break：值更大，或值相等且下标更小
        if (partial_val[i] > local.val ||
            (partial_val[i] == local.val && (local.idx < 0 || partial_idx[i] < local.idx))) {
            local.val = partial_val[i];
            local.idx = partial_idx[i];
        }
    }

    ValIdx result = block_reduce_max<BLOCK_SIZE>(local);

    if (tid == 0) {
        *output = result.idx;
    }
}

} // namespace

// ===================== Host 调用接口 =====================

void launch_argmax(
    const float* d_input,
    int* d_output,
    float* d_partial_val,
    int* d_partial_idx,
    int n,
    int num_blocks,
    cudaStream_t stream
) {
    constexpr int BLOCK_SIZE = 256;
    argmax_kernel<BLOCK_SIZE><<<num_blocks, BLOCK_SIZE, 0, stream>>>(
        d_input, d_partial_val, d_partial_idx, n
    );
    argmax_final_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE, 0, stream>>>(
        d_partial_val, d_partial_idx, d_output, num_blocks
    );
}

// ===================== 测试 =====================

int main() {
    const int N = 1 << 20;  // 1M 元素
    const int NUM_BLOCKS = 128;

    // Host 数据
    float* h_input = (float*)malloc(N * sizeof(float));
    srand(42);
    int expected_idx = 0;
    float expected_val = -FLT_MAX;
    for (int i = 0; i < N; i++) {
        h_input[i] = (float)rand() / RAND_MAX;
        if (h_input[i] > expected_val) {
            expected_val = h_input[i];
            expected_idx = i;
        }
    }

    // Device 数据
    float *d_input, *d_partial_val;
    int *d_partial_idx, *d_output;
    cudaMalloc(&d_input, N * sizeof(float));
    cudaMalloc(&d_partial_val, NUM_BLOCKS * sizeof(float));
    cudaMalloc(&d_partial_idx, NUM_BLOCKS * sizeof(int));
    cudaMalloc(&d_output, sizeof(int));

    cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice);

    // 执行
    launch_argmax(d_input, d_output, d_partial_val, d_partial_idx, N, NUM_BLOCKS, 0);

    // 取结果
    int h_output;
    cudaMemcpy(&h_output, d_output, sizeof(int), cudaMemcpyDeviceToHost);

    printf("ArgMax result: index = %d, value = %f\n", h_output, h_input[h_output]);
    printf("Expected:      index = %d, value = %f\n", expected_idx, expected_val);
    printf("%s\n", (h_output == expected_idx) ? "PASS" : "FAIL");

    // 清理
    free(h_input);
    cudaFree(d_input);
    cudaFree(d_partial_val);
    cudaFree(d_partial_idx);
    cudaFree(d_output);

    return 0;
}

检查方法

有 CUDA 环境时执行 nvcc -std=c++17 -c /tmp/12_argmax.cu -o /tmp/12_argmax.cu.o；无 CUDA 环境时只做静态审查并登记待验证。

提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“ArgMax 对归约与确定性”的工作机制。

你的答案：

Q2

只用 > 时为什么不同 block 划分可能返回不同的相等最大值索引？

你的答案：

Q3

若输入包含 NaN，NumPy、PyTorch 与你的 kernel 应采用哪种约定？

你的答案：

评分与通过规则

代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分；
Q1～Q3 各 2 分；
一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。

💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。

Code

%%writefile /tmp/12_argmax.cu
#include <cuda_runtime.h>
#include <float.h>
#include <stdio.h>
#include <stdlib.h>

namespace {

// ArgMax: 找到 input[0:n] 中最大值的索引。
//
// 面试重点：
// 1. reduce 时同时传递 (value, index) 对。
// 2. warp shuffle 传 value 比较，再传 index 选择。
// 3. 两轮 kernel：第一轮每个 block 出一个 (val, idx)，第二轮归并。

struct ValIdx {
    float val;
    int idx;
};

template<int BLOCK_SIZE>
__device__ __forceinline__ ValIdx warp_reduce_max(ValIdx vi) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other_val = __shfl_down_sync(0xffffffff, vi.val, offset);
        int other_idx = __shfl_down_sync(0xffffffff, vi.idx, offset);
        if (other_val > vi.val || (other_val == vi.val && other_idx >= 0 && (vi.idx < 0 || other_idx < vi.idx))) {
            vi.val = other_val;
            vi.idx = other_idx;
        }
    }
    return vi;
}

template<int BLOCK_SIZE>
__device__ __forceinline__ ValIdx block_reduce_max(ValIdx vi) {
    constexpr int NUM_WARPS = (BLOCK_SIZE + 31) / 32;
    __shared__ float smem_val[NUM_WARPS];
    __shared__ int smem_idx[NUM_WARPS];

    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;

    vi = warp_reduce_max<BLOCK_SIZE>(vi);

    if (lane == 0) {
        smem_val[warp] = vi.val;
        smem_idx[warp] = vi.idx;
    }
    __syncthreads();

    if (warp == 0) {
        vi.val = (lane < NUM_WARPS) ? smem_val[lane] : -FLT_MAX;
        vi.idx = (lane < NUM_WARPS) ? smem_idx[lane] : -1;
        vi = warp_reduce_max<BLOCK_SIZE>(vi);
    }

    return vi;
}

// 第一轮：每个 block 用 grid-stride loop 找到局部最大值及其索引
template<int BLOCK_SIZE>
__global__ void argmax_kernel(
    const float* __restrict__ input,
    float* __restrict__ partial_val,
    int* __restrict__ partial_idx,
    int n
) {
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + tid;
    int stride = blockDim.x * gridDim.x;

    // 每个线程先通过 grid-stride loop 找到自己负责的元素中的最大值
    ValIdx local = {-FLT_MAX, -1};
    for (int i = idx; i < n; i += stride) {
        if (input[i] > local.val) {
            local.val = input[i];
            local.idx = i;
        }
    }

    // block 级 reduce
    ValIdx result = block_reduce_max<BLOCK_SIZE>(local);

    if (tid == 0) {
        partial_val[blockIdx.x] = result.val;
        partial_idx[blockIdx.x] = result.idx;
    }
}

// 第二轮：把所有 block 的 partial 结果归并，只需一个 block
template<int BLOCK_SIZE>
__global__ void argmax_final_kernel(
    const float* __restrict__ partial_val,
    const int* __restrict__ partial_idx,
    int* __restrict__ output,
    int num_blocks
) {
    int tid = threadIdx.x;

    ValIdx local = {-FLT_MAX, -1};
    // 一个 block 处理所有 partial，stride = BLOCK_SIZE
    for (int i = tid; i < num_blocks; i += BLOCK_SIZE) {
        // 与 warp/block 归约使用同一 tie-break：值更大，或值相等且下标更小
        if (partial_val[i] > local.val ||
            (partial_val[i] == local.val && (local.idx < 0 || partial_idx[i] < local.idx))) {
            local.val = partial_val[i];
            local.idx = partial_idx[i];
        }
    }

    ValIdx result = block_reduce_max<BLOCK_SIZE>(local);

    if (tid == 0) {
        *output = result.idx;
    }
}

} // namespace

// ===================== Host 调用接口 =====================

void launch_argmax(
    const float* d_input,
    int* d_output,
    float* d_partial_val,
    int* d_partial_idx,
    int n,
    int num_blocks,
    cudaStream_t stream
) {
    constexpr int BLOCK_SIZE = 256;
    argmax_kernel<BLOCK_SIZE><<<num_blocks, BLOCK_SIZE, 0, stream>>>(
        d_input, d_partial_val, d_partial_idx, n
    );
    argmax_final_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE, 0, stream>>>(
        d_partial_val, d_partial_idx, d_output, num_blocks
    );
}

// ===================== 测试 =====================

int main() {
    const int N = 1 << 20;  // 1M 元素
    const int NUM_BLOCKS = 128;

    // Host 数据
    float* h_input = (float*)malloc(N * sizeof(float));
    srand(42);
    int expected_idx = 0;
    float expected_val = -FLT_MAX;
    for (int i = 0; i < N; i++) {
        h_input[i] = (float)rand() / RAND_MAX;
        if (h_input[i] > expected_val) {
            expected_val = h_input[i];
            expected_idx = i;
        }
    }

    // Device 数据
    float *d_input, *d_partial_val;
    int *d_partial_idx, *d_output;
    cudaMalloc(&d_input, N * sizeof(float));
    cudaMalloc(&d_partial_val, NUM_BLOCKS * sizeof(float));
    cudaMalloc(&d_partial_idx, NUM_BLOCKS * sizeof(int));
    cudaMalloc(&d_output, sizeof(int));

    cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice);

    // 执行
    launch_argmax(d_input, d_output, d_partial_val, d_partial_idx, N, NUM_BLOCKS, 0);

    // 取结果
    int h_output;
    cudaMemcpy(&h_output, d_output, sizeof(int), cudaMemcpyDeviceToHost);

    printf("ArgMax result: index = %d, value = %f\n", h_output, h_input[h_output]);
    printf("Expected:      index = %d, value = %f\n", expected_idx, expected_val);
    printf("%s\n", (h_output == expected_idx) ? "PASS" : "FAIL");

    // 清理
    free(h_input);
    cudaFree(d_input);
    cudaFree(d_partial_val);
    cudaFree(d_partial_idx);
    cudaFree(d_output);

    return 0;
}

Q1 参考答案

输入：Pass 1 的每个线程把 ValIdx local = {-FLT_MAX, -1} 放在寄存器里，以 Grid-Stride 合并读取输入；同一线程按下标递增遍历，用 x_i > local.val 更新即可保留最小下标。
状态变化：Warp 内用 __shfl_down_sync 同时交换 val 与 idx，按 $(b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$ 仲裁；各 Warp 结果经 Shared Memory 再归约，每个 Block 的 Thread 0 写出 partial_val/partial_idx[blockIdx.x]（详见核心概念第 2 节）。
输出：Pass 2 用单个 Block 对 num_blocks 个候选按同一规则归约，Thread 0 把最大值的最小下标写入 *output。

Q2 参考答案

只用 > 时，相等的两个值比较后保留的是左操作数 $a$ ，即洗牌接收方当前持有的元组，而树状归约中 $a$ 的下标不一定更小。改变 num_blocks 或 BLOCK_SIZE 会改变谁和谁在哪一步比较：例如下标 10 和 500 可能由同一线程串行比较，也可能分在不同 Warp 或不同 Block 里，在树的不同位置相遇，最终留下哪个下标取决于这棵比较树的形状。加上 b.val == a.val && b.idx < a.idx 后，仲裁结果与比较顺序无关，始终是最小下标（详见核心概念第 3 节）。

Q3 参考答案

当输入数据中混入不可比的 NaN（Not a Number）异常值时，各生态系统与高性能算子库的约定及底层逻辑如下：

NumPy 约定（np.argmax）：
- 策略：首个 NaN 优先捕获。若输入包含 NaN，NumPy 会将输出索引锁定为数组中第一次出现 NaN 的物理下标。这有助于开发者迅速定位哪一个特征发生了数值污染；
PyTorch 约定（torch.argmax）：
- 策略：在早期版本或某些后端上，torch.argmax 在遇到 NaN 时的行为可能未定义（受底层 cuBLAS/CUB 实现影响）；但在标准行为中，若存在 NaN，通常会优先传播该异常，直接返回包含 NaN 的首个位置下标。若用户希望忽略 NaN，则需显式调用 torch.nanargmax（跳过 NaN 仅在合法实数中找最大值）；
本课 CUDA Kernel 的工程与数值约定建议：
- 朴素浮点比较的静默崩溃缺陷：在 IEEE 754 标准下，NaN > x 与 x > NaN 恒为 false。若最大候选值为 NaN，它无法战胜任何数，导致 NaN 被悄无声息地吞噬掉，掩盖了严重的数值 Bug；
- 生产级确定性约定：在 Kernel 内部定义三元仲裁谓词：
```
bool b_wins = isnan(b.val) ? (!isnan(a.val) || b.idx < a.idx)
                           : (isnan(a.val) ? false : (b.val > a.val || (b.val == a.val && b.idx < a.idx)));
```
  即：若发生 NaN 污染，将 NaN 视为具有最高传播优先级的异常，并严格返回出现 NaN 的最小下标；或者如果按 nanargmax 规则，则在读取阶段直接将 NaN 过滤为 -FLT_MAX 并记录告警。保持对 NaN 的显式定义是不变量工程的核心防线。

参考资料

资料用于建立事实基线；面试回答仍需用自己的语言组织。

← 上一课
第 11 课：SiLU 与算子融合 下一课 →
第 13 课：DeepSeek DeepGEM…

--- title: "第 12 课：ArgMax 对归约与确定性" jupyter: python3 --- ## 本课目标与完成标准学完后你应能：归约 `(value,index)` 对，处理相等值、NaN 和两阶段输出。通过必须同时满足： - 独立补齐本课唯一的代码填空题，并通过给定检查； - 三个问答题均说明因果链，而不是只报术语； - 能指出至少一个正确性边界和一个性能取舍； - 总分不低于 8/10，且没有一票否决级概念错误。 ## 前置关系 - 课程前置：C/C++、线性代数、基本并行编程 - 本课在路线中的作用：ArgMax 不只归约值，还要同步携带索引；这是自定义归约算子的典型模式。 ## 核心心智模型 ### 1. 它是什么，解决什么问题 **ArgMax 算子（最大值索引定位）** 是模型解码贪婪采样（Greedy Decoding）、Top-1 类别预测以及强化学习离散动作选择的核心基础。不同于普通的 ReduceMax 只关心最大数值本身，ArgMax 要求输出取得最大值的**物理全局下标（Index）**： $$\text{ArgMax}(\mathbf{x}) = \arg\max_{0 \le i < N} x_i = \{i \mid \forall j, \ x_i \ge x_j\}$$ **并行归约中的确定性挑战（Tie-Breaking 不变量，图 12-1）**： ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-sg" data-nid="TupleReduction" data-depth="0"><div class="fd-sg-title">ArgMax 二元元组仲裁</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank fd-few"><div class="fd-node fd-box" data-nid="P1">元组 1: (val_1, idx_1)</div><div class="fd-node fd-box" data-nid="P2">元组 2: (val_2, idx_2)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-dec" data-nid="Cmp">val_1 > val_2 ?</div></div><div class="fd-gap" style="--gl:171px;--gt:53px"></div><div class="fd-rank fd-few"><div class="fd-node fd-box" data-nid="Win1">胜出: (val_1, idx_1)</div><div class="fd-node fd-box" data-nid="Win2">胜出: (val_2, idx_2)</div><div class="fd-node fd-dec" data-nid="Tie">min(idx_1, idx_2)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Deterministic">严格保留最小索引 (确定性 Tie-Breaking)</div></div></div></div></div></div></div><i class="fd-e" data-s="P1" data-t="Cmp" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="P2" data-t="Cmp" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Cmp" data-t="Win1" data-k="solid" data-h="1" data-b="0">是</i><i class="fd-e" data-s="Cmp" data-t="Win2" data-k="solid" data-h="1" data-b="0">否: val_2 > val_1</i><i class="fd-e" data-s="Cmp" data-t="Tie" data-k="solid" data-h="1" data-b="0">等值: val_1 == val_2</i><i class="fd-e" data-s="Tie" data-t="Deterministic" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 12-1：ArgMax 键值对元组结构与确定性仲裁规则</em></p> - 在串行 CPU 代码中，通常用 `if (x[i] > max_val) { max_val = x[i]; best_idx = i; }`，天然具有“遇到相等最大值保留最左侧（最小下标）”的确定性； - 但在 GPU 上，成千上万个线程以树状结构并行对偶比较。当输入数组中存在多个相等的最大值（例如分词 Logits 中有多个未见过的 Padding Token 得分相同）时，由于 Warp 调度和 Block 树状归约的无序性，若仅仅通过 `val_b > val_a` 进行筛选，谁会胜出完全取决于线程树的配对顺序； - **确定性仲裁规则（Tie-breaking Policy）**：必须将数据结构扩充为双元组 `ValIdx { float val; int idx; }`，并在比较逻辑中严格引入字典序比较：数值大者胜出；若数值相等，则**全局下标更小者胜出**。这样结果与比较顺序无关，在不同架构和 Grid 配置下都相同。 ### 2. 它如何工作（结合图 12-1 与图 12-2）由于需要支持百万到数十亿规模的超大张量，算法采用经典的 **两阶段分层归约架构（Two-Stage Grid Reduction）**： 1. **阶段一：网格局部归约（Pass 1: Block-level ArgMax）** - 启动包含 `num_blocks` 个 Block 的网格，各线程以 Grid-Stride 方式遍历整个长数组； - 每个线程在本地寄存器中维护当前的局部最佳二元状态 `ValIdx local = {-FLT_MAX, -1}`； - **Warp 级元组洗牌（图 12-2）**： ```{=html} <div class="fd fd-seq" role="img"><div class="sd-grid" style="--n:3"><div class="sd-life" style="grid-column:1;grid-row:1/7"></div><div class="sd-life" style="grid-column:2;grid-row:1/7"></div><div class="sd-life" style="grid-column:3;grid-row:1/7"></div><div class="sd-actor" style="grid-column:1;grid-row:1">Warp 内部 32 线程</div><div class="sd-actor" style="grid-column:2;grid-row:1">__shfl_down_sync 洗牌</div><div class="sd-actor" style="grid-column:3;grid-row:1">严格仲裁算子</div><div class="sd-note sd-note-over" style="grid-column:1/2;grid-row:2;--k:1">步进 16 -> 8 -> 4 -> 2 -> 1 折半归约</div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:1/3;grid-row:3;--k:2"><div class="sd-text"><span class="sd-num">1</span>同时洗牌传输 val 与 idx (或 64 位整型打包)</div><div class="sd-line"></div></div><div class="sd-msg sd-r sd-h-arrow" style="grid-column:2/4;grid-row:4;--k:2"><div class="sd-text"><span class="sd-num">2</span>比较 pair 与 other_pair</div><div class="sd-line"></div></div><div class="sd-msg sd-l sd-h-arrow" style="grid-column:1/4;grid-row:5;--k:3"><div class="sd-text"><span class="sd-num">3</span>更新局部胜出元组</div><div class="sd-line"></div></div><div class="sd-note sd-note-over" style="grid-column:1/2;grid-row:6;--k:1">Lane 0 获得全局极值与准确最小全局索引</div></div></div> ``` <p class="caption" align="center"><em>图 12-2：Warp 级确定性 ArgMax 归约时序流</em></p> 利用 `__shfl_down_sync` 同时在寄存器间交换 `val` 与 `idx`（或打包传输），在 5 步折半归约中通过确定性仲裁逻辑选出当前 Warp 的胜者； - **Block 级发布**：各 Warp 的 Lane 0 将局部的 `ValIdx` 写入共享内存 `smem`，由 Warp 0 再次归约出当前 Block 的最佳结果，最终由 Thread 0 将该 Block 的局部极值写入全局临时数组 `partial_val[blockIdx.x]` 与 `partial_idx[blockIdx.x]`。 2. **阶段二：单块全局终局归约（Pass 2: Single-Block Final Reduce）** - 启动仅有 1 个 ThreadBlock 的 Kernel，Block 内线程并行读取 `num_blocks` 个局部极值候选者； - 在 Block 内部再次执行树状元组归约，Thread 0 将最终胜者的全局索引单次写回输出指针 `*output`。 ### 3. 正确性条件与常见误区 1. **相等最大值的非确定性漏洞（Tie-breaking 缺失）**： - 如果比较逻辑仅写为 `if (b.val > a.val) a = b;`：当 `b.val == a.val` 时，条件为假，保留 `a`。但是，在树状归约中，谁作为参数 `a`、谁作为参数 `b` 完全受折半步进的硬件偏移量（Offset）支配。例如在 Offset=16 的比较中，Lane 0 比较 Lane 16；在 Offset=8 时，Lane 0 比较 Lane 8。这会导致最终保留下来的不一定是全局最小下标，甚至不同 Grid 大小跑出来的下标完全不同； - **强制不变量**：必须严格书写确定性判据： $$b \text{ 胜出} \iff (b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$$ 2. **NaN 污染策略（NaN Handling Policy）**： - 根据 IEEE 754 规范，任何包含 NaN 的浮点比较（`NaN > x`、`x > NaN`、`NaN == x`）在硬件中全部返回 `false`； - 若未对 NaN 进行特判，NaN 永远无法在 `>` 比较中胜出，可能导致真正的最大值被静默吞掉，或者导致首个遇到的正常数被当作结果；工业级实现（如 PyTorch）通常定义“遇到 NaN 则直接传播 NaN 或优先返回首个 NaN 的索引”。 ### 4. 性能与工程取舍 1. **元组洗牌带宽开销 vs. 整数压缩**： - 普通 ReduceMax 每次洗牌仅需传输 1 个 32-bit 的 `float`； - ArgMax 每次洗牌必须传递 `val`（4 字节）和 `idx`（4 字节），洗牌指令数翻倍； - **工程优化**：若浮点数为 FP16/BF16（2 字节），或者能保证 `val` 为非负数，可把浮点数的有序编码与下标打包进一个 64-bit 整数（`uint64_t`），一次比较同时完成数值与下标的 tie-break。注意 64-bit `__shfl` 会被编译成两条 32-bit SHFL；若要真正单条，需要打包进 32 位（如 FP16 值 + 16 位下标）。 2. **两阶段 Kernel vs. 原子操作/跨 Block 栅障**： - 采用两阶段独立 Kernel Launch 拥有绝对的硬件可移植性与安全性，中间通过 Global Memory 传递 partial 结果，天然具备 Stream 内存可见性； - 单 Kernel 协作（使用 Cooperative Groups 网格级同步）可省去二次 Launch 开销，但对 GPU 架构与 Occupancy 调度有严格的硬件配额约束。 ## 具体演示：平局仲裁确定性演算考虑极小输入数组：$x = [5.0, 7.0, 7.0, 3.0]$，其中最大值 $7.0$ 在下标 1 和 2 处重复出现： - **无下标仲裁（仅凭 `val_b > val_a`）**： - 线程 1 拥有 $(val=7.0, idx=1)$，线程 2 拥有 $(val=7.0, idx=2)$； - 若在归约折半中由线程 2 比较线程 1（`b = 线程1, a = 线程2`）：$7.0 > 7.0$ 为假，保留线程 2 的结果，最终输出 $\text{Index} = 2$； - 若换一种网格拓扑或线程配对顺序，由线程 1 比较线程 2（`b = 线程2, a = 线程1`）：$7.0 > 7.0$ 为假，保留线程 1 的结果，最终输出 $\text{Index} = 1$； - **严重后果**：相同的模型权重和输入，在不同的 Batch 大小或不同的 GPU 架构上推理，生成的下一个 Token 竟然完全不同。 - **严格确定性字典序仲裁**： - 比较逻辑：$(b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$； - 无论二者谁作为参数 $a$ 谁作为参数 $b$，下标较小的 $idx=1$ 恒定被判定为胜者； - 无论如何调度，输出绝对恒为 $\text{Index} = 1$，完美捍卫了模型贪婪采样的可复现性（Reproducibility）。 ## 实践任务：唯一代码填空题补齐确定性比较条件；注意初始化 idx=-1。规则：只能修改 `TODO`/`______` 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。 ```{python} #| course_role: exercise %%writefile /tmp/12_argmax.cu #include <cuda_runtime.h> #include <float.h> #include <stdio.h> #include <stdlib.h> namespace { // ArgMax: 找到 input[0:n] 中最大值的索引。 // // 面试重点： // 1. reduce 时同时传递 (value, index) 对。 // 2. warp shuffle 传 value 比较，再传 index 选择。 // 3. 两轮 kernel：第一轮每个 block 出一个 (val, idx)，第二轮归并。 struct ValIdx { float val; int idx; }; template<int BLOCK_SIZE> __device__ __forceinline__ ValIdx warp_reduce_max(ValIdx vi) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_val = __shfl_down_sync(0xffffffff, vi.val, offset); int other_idx = __shfl_down_sync(0xffffffff, vi.idx, offset); if (______) { // TODO: 选择更大值；相等时选择更小有效索引 vi.val = other_val; vi.idx = other_idx; } } return vi; } template<int BLOCK_SIZE> __device__ __forceinline__ ValIdx block_reduce_max(ValIdx vi) { constexpr int NUM_WARPS = (BLOCK_SIZE + 31) / 32; __shared__ float smem_val[NUM_WARPS]; __shared__ int smem_idx[NUM_WARPS]; int lane = threadIdx.x & 31; int warp = threadIdx.x >> 5; vi = warp_reduce_max<BLOCK_SIZE>(vi); if (lane == 0) { smem_val[warp] = vi.val; smem_idx[warp] = vi.idx; } __syncthreads(); if (warp == 0) { vi.val = (lane < NUM_WARPS) ? smem_val[lane] : -FLT_MAX; vi.idx = (lane < NUM_WARPS) ? smem_idx[lane] : -1; vi = warp_reduce_max<BLOCK_SIZE>(vi); } return vi; } // 第一轮：每个 block 用 grid-stride loop 找到局部最大值及其索引 template<int BLOCK_SIZE> __global__ void argmax_kernel( const float* __restrict__ input, float* __restrict__ partial_val, int* __restrict__ partial_idx, int n ) { int tid = threadIdx.x; int idx = blockIdx.x * blockDim.x + tid; int stride = blockDim.x * gridDim.x; // 每个线程先通过 grid-stride loop 找到自己负责的元素中的最大值 ValIdx local = {-FLT_MAX, -1}; for (int i = idx; i < n; i += stride) { if (input[i] > local.val) { local.val = input[i]; local.idx = i; } } // block 级 reduce ValIdx result = block_reduce_max<BLOCK_SIZE>(local); if (tid == 0) { partial_val[blockIdx.x] = result.val; partial_idx[blockIdx.x] = result.idx; } } // 第二轮：把所有 block 的 partial 结果归并，只需一个 block template<int BLOCK_SIZE> __global__ void argmax_final_kernel( const float* __restrict__ partial_val, const int* __restrict__ partial_idx, int* __restrict__ output, int num_blocks ) { int tid = threadIdx.x; ValIdx local = {-FLT_MAX, -1}; // 一个 block 处理所有 partial，stride = BLOCK_SIZE for (int i = tid; i < num_blocks; i += BLOCK_SIZE) { // 与 warp/block 归约使用同一 tie-break：值更大，或值相等且下标更小 if (partial_val[i] > local.val || (partial_val[i] == local.val && (local.idx < 0 || partial_idx[i] < local.idx))) { local.val = partial_val[i]; local.idx = partial_idx[i]; } } ValIdx result = block_reduce_max<BLOCK_SIZE>(local); if (tid == 0) { *output = result.idx; } } } // namespace // ===================== Host 调用接口 ===================== void launch_argmax( const float* d_input, int* d_output, float* d_partial_val, int* d_partial_idx, int n, int num_blocks, cudaStream_t stream ) { constexpr int BLOCK_SIZE = 256; argmax_kernel<BLOCK_SIZE><<<num_blocks, BLOCK_SIZE, 0, stream>>>( d_input, d_partial_val, d_partial_idx, n ); argmax_final_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE, 0, stream>>>( d_partial_val, d_partial_idx, d_output, num_blocks ); } // ===================== 测试 ===================== int main() { const int N = 1 << 20; // 1M 元素 const int NUM_BLOCKS = 128; // Host 数据 float* h_input = (float*)malloc(N * sizeof(float)); srand(42); int expected_idx = 0; float expected_val = -FLT_MAX; for (int i = 0; i < N; i++) { h_input[i] = (float)rand() / RAND_MAX; if (h_input[i] > expected_val) { expected_val = h_input[i]; expected_idx = i; } } // Device 数据 float *d_input, *d_partial_val; int *d_partial_idx, *d_output; cudaMalloc(&d_input, N * sizeof(float)); cudaMalloc(&d_partial_val, NUM_BLOCKS * sizeof(float)); cudaMalloc(&d_partial_idx, NUM_BLOCKS * sizeof(int)); cudaMalloc(&d_output, sizeof(int)); cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice); // 执行 launch_argmax(d_input, d_output, d_partial_val, d_partial_idx, N, NUM_BLOCKS, 0); // 取结果 int h_output; cudaMemcpy(&h_output, d_output, sizeof(int), cudaMemcpyDeviceToHost); printf("ArgMax result: index = %d, value = %f\n", h_output, h_input[h_output]); printf("Expected: index = %d, value = %f\n", expected_idx, expected_val); printf("%s\n", (h_output == expected_idx) ? "PASS" : "FAIL"); // 清理 free(h_input); cudaFree(d_input); cudaFree(d_partial_val); cudaFree(d_partial_idx); cudaFree(d_output); return 0; } ``` ### 检查方法有 CUDA 环境时执行 `nvcc -std=c++17 -c /tmp/12_argmax.cu -o /tmp/12_argmax.cu.o`；无 CUDA 环境时只做静态审查并登记待验证。提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。 ### Q1 不要背定义：请从输入、状态变化和输出三个阶段解释“ArgMax 对归约与确定性”的工作机制。 **你的答案：** ### Q2 只用 `>` 时为什么不同 block 划分可能返回不同的相等最大值索引？ **你的答案：** ### Q3 若输入包含 NaN，NumPy、PyTorch 与你的 kernel 应采用哪种约定？ **你的答案：** ## 评分与通过规则 - 代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分； - Q1～Q3 各 2 分； - 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。 ```{python} #| course_role: solution %%writefile /tmp/12_argmax.cu #include <cuda_runtime.h> #include <float.h> #include <stdio.h> #include <stdlib.h> namespace { // ArgMax: 找到 input[0:n] 中最大值的索引。 // // 面试重点： // 1. reduce 时同时传递 (value, index) 对。 // 2. warp shuffle 传 value 比较，再传 index 选择。 // 3. 两轮 kernel：第一轮每个 block 出一个 (val, idx)，第二轮归并。 struct ValIdx { float val; int idx; }; template<int BLOCK_SIZE> __device__ __forceinline__ ValIdx warp_reduce_max(ValIdx vi) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_val = __shfl_down_sync(0xffffffff, vi.val, offset); int other_idx = __shfl_down_sync(0xffffffff, vi.idx, offset); if (other_val > vi.val || (other_val == vi.val && other_idx >= 0 && (vi.idx < 0 || other_idx < vi.idx))) { vi.val = other_val; vi.idx = other_idx; } } return vi; } template<int BLOCK_SIZE> __device__ __forceinline__ ValIdx block_reduce_max(ValIdx vi) { constexpr int NUM_WARPS = (BLOCK_SIZE + 31) / 32; __shared__ float smem_val[NUM_WARPS]; __shared__ int smem_idx[NUM_WARPS]; int lane = threadIdx.x & 31; int warp = threadIdx.x >> 5; vi = warp_reduce_max<BLOCK_SIZE>(vi); if (lane == 0) { smem_val[warp] = vi.val; smem_idx[warp] = vi.idx; } __syncthreads(); if (warp == 0) { vi.val = (lane < NUM_WARPS) ? smem_val[lane] : -FLT_MAX; vi.idx = (lane < NUM_WARPS) ? smem_idx[lane] : -1; vi = warp_reduce_max<BLOCK_SIZE>(vi); } return vi; } // 第一轮：每个 block 用 grid-stride loop 找到局部最大值及其索引 template<int BLOCK_SIZE> __global__ void argmax_kernel( const float* __restrict__ input, float* __restrict__ partial_val, int* __restrict__ partial_idx, int n ) { int tid = threadIdx.x; int idx = blockIdx.x * blockDim.x + tid; int stride = blockDim.x * gridDim.x; // 每个线程先通过 grid-stride loop 找到自己负责的元素中的最大值 ValIdx local = {-FLT_MAX, -1}; for (int i = idx; i < n; i += stride) { if (input[i] > local.val) { local.val = input[i]; local.idx = i; } } // block 级 reduce ValIdx result = block_reduce_max<BLOCK_SIZE>(local); if (tid == 0) { partial_val[blockIdx.x] = result.val; partial_idx[blockIdx.x] = result.idx; } } // 第二轮：把所有 block 的 partial 结果归并，只需一个 block template<int BLOCK_SIZE> __global__ void argmax_final_kernel( const float* __restrict__ partial_val, const int* __restrict__ partial_idx, int* __restrict__ output, int num_blocks ) { int tid = threadIdx.x; ValIdx local = {-FLT_MAX, -1}; // 一个 block 处理所有 partial，stride = BLOCK_SIZE for (int i = tid; i < num_blocks; i += BLOCK_SIZE) { // 与 warp/block 归约使用同一 tie-break：值更大，或值相等且下标更小 if (partial_val[i] > local.val || (partial_val[i] == local.val && (local.idx < 0 || partial_idx[i] < local.idx))) { local.val = partial_val[i]; local.idx = partial_idx[i]; } } ValIdx result = block_reduce_max<BLOCK_SIZE>(local); if (tid == 0) { *output = result.idx; } } } // namespace // ===================== Host 调用接口 ===================== void launch_argmax( const float* d_input, int* d_output, float* d_partial_val, int* d_partial_idx, int n, int num_blocks, cudaStream_t stream ) { constexpr int BLOCK_SIZE = 256; argmax_kernel<BLOCK_SIZE><<<num_blocks, BLOCK_SIZE, 0, stream>>>( d_input, d_partial_val, d_partial_idx, n ); argmax_final_kernel<BLOCK_SIZE><<<1, BLOCK_SIZE, 0, stream>>>( d_partial_val, d_partial_idx, d_output, num_blocks ); } // ===================== 测试 ===================== int main() { const int N = 1 << 20; // 1M 元素 const int NUM_BLOCKS = 128; // Host 数据 float* h_input = (float*)malloc(N * sizeof(float)); srand(42); int expected_idx = 0; float expected_val = -FLT_MAX; for (int i = 0; i < N; i++) { h_input[i] = (float)rand() / RAND_MAX; if (h_input[i] > expected_val) { expected_val = h_input[i]; expected_idx = i; } } // Device 数据 float *d_input, *d_partial_val; int *d_partial_idx, *d_output; cudaMalloc(&d_input, N * sizeof(float)); cudaMalloc(&d_partial_val, NUM_BLOCKS * sizeof(float)); cudaMalloc(&d_partial_idx, NUM_BLOCKS * sizeof(int)); cudaMalloc(&d_output, sizeof(int)); cudaMemcpy(d_input, h_input, N * sizeof(float), cudaMemcpyHostToDevice); // 执行 launch_argmax(d_input, d_output, d_partial_val, d_partial_idx, N, NUM_BLOCKS, 0); // 取结果 int h_output; cudaMemcpy(&h_output, d_output, sizeof(int), cudaMemcpyDeviceToHost); printf("ArgMax result: index = %d, value = %f\n", h_output, h_input[h_output]); printf("Expected: index = %d, value = %f\n", expected_idx, expected_val); printf("%s\n", (h_output == expected_idx) ? "PASS" : "FAIL"); // 清理 free(h_input); cudaFree(d_input); cudaFree(d_partial_val); cudaFree(d_partial_idx); cudaFree(d_output); return 0; } ``` ### Q1 参考答案 - **输入**：Pass 1 的每个线程把 `ValIdx local = {-FLT_MAX, -1}` 放在寄存器里，以 Grid-Stride 合并读取输入；同一线程按下标递增遍历，用 `x_i > local.val` 更新即可保留最小下标。 - **状态变化**：Warp 内用 `__shfl_down_sync` 同时交换 `val` 与 `idx`，按 $(b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$ 仲裁；各 Warp 结果经 Shared Memory 再归约，每个 Block 的 Thread 0 写出 `partial_val/partial_idx[blockIdx.x]`（详见核心概念第 2 节）。 - **输出**：Pass 2 用单个 Block 对 `num_blocks` 个候选按同一规则归约，Thread 0 把最大值的最小下标写入 `*output`。 ### Q2 参考答案只用 `>` 时，相等的两个值比较后保留的是左操作数 $a$，即洗牌接收方当前持有的元组，而树状归约中 $a$ 的下标不一定更小。改变 `num_blocks` 或 `BLOCK_SIZE` 会改变谁和谁在哪一步比较：例如下标 10 和 500 可能由同一线程串行比较，也可能分在不同 Warp 或不同 Block 里，在树的不同位置相遇，最终留下哪个下标取决于这棵比较树的形状。加上 `b.val == a.val && b.idx < a.idx` 后，仲裁结果与比较顺序无关，始终是最小下标（详见核心概念第 3 节）。 ### Q3 参考答案当输入数据中混入不可比的 `NaN`（Not a Number）异常值时，各生态系统与高性能算子库的约定及底层逻辑如下： 1. **NumPy 约定（`np.argmax`）**： - 策略：**首个 NaN 优先捕获**。若输入包含 NaN，NumPy 会将输出索引锁定为数组中**第一次出现 NaN 的物理下标**。这有助于开发者迅速定位哪一个特征发生了数值污染； 2. **PyTorch 约定（`torch.argmax`）**： - 策略：在早期版本或某些后端上，`torch.argmax` 在遇到 NaN 时的行为可能未定义（受底层 cuBLAS/CUB 实现影响）；但在标准行为中，若存在 NaN，通常会优先传播该异常，直接返回包含 NaN 的首个位置下标。若用户希望忽略 NaN，则需显式调用 `torch.nanargmax`（跳过 NaN 仅在合法实数中找最大值）； 3. **本课 CUDA Kernel 的工程与数值约定建议**： - **朴素浮点比较的静默崩溃缺陷**：在 IEEE 754 标准下，`NaN > x` 与 `x > NaN` 恒为 `false`。若最大候选值为 NaN，它无法战胜任何数，导致 NaN 被悄无声息地吞噬掉，掩盖了严重的数值 Bug； - **生产级确定性约定**：在 Kernel 内部定义三元仲裁谓词： ```cpp bool b_wins = isnan(b.val) ? (!isnan(a.val) || b.idx < a.idx) : (isnan(a.val) ? false : (b.val > a.val || (b.val == a.val && b.idx < a.idx))); ``` 即：若发生 NaN 污染，将 NaN 视为具有最高传播优先级的异常，并严格返回出现 NaN 的最小下标；或者如果按 `nanargmax` 规则，则在读取阶段直接将 NaN 过滤为 `-FLT_MAX` 并记录告警。保持对 NaN 的显式定义是不变量工程的核心防线。 ## 参考资料 - [CUDA Programming Guide](https://docs.nvidia.com/cuda/cuda-programming-guide/) - [CUDA Best Practices Guide](https://docs.nvidia.com/cuda/cuda-c-best-practices-guide/) 资料用于建立事实基线；面试回答仍需用自己的语言组织。 ::: ::: {.course-pagination} [**← 上一课**<br>第 11 课：SiLU 与算子融合](lesson11_silu.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 13 课：DeepSeek DeepGEM...](lesson13_deepgemm.qmd){.pagination-link .pagination-next} :::