第 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）：

flowchart TD
    subgraph TupleReduction["ArgMax 二元元组仲裁"]
        P1["元组 1: (val_1, idx_1)"] --> Cmp{"val_1 > val_2 ?"}
        P2["元组 2: (val_2, idx_2)"] --> Cmp
        Cmp -->|是| Win1["胜出: (val_1, idx_1)"]
        Cmp -->|否: val_2 > val_1| Win2["胜出: (val_2, idx_2)"]
        Cmp -->|等值: val_1 == val_2| Tie{"min(idx_1, idx_2)"}
        Tie --> Deterministic["严格保留最小索引 (确定性 Tie-Breaking)"]
    end

图 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 配置的 100% 确定性。

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）：

sequenceDiagram
    autonumber
    participant Lane as Warp 内部 32 线程
    participant Shfl as __shfl_down_sync 洗牌
    participant Arb as 严格仲裁算子

    Note over Lane: 步进 16 -> 8 -> 4 -> 2 -> 1 折半归约
    Lane->>Shfl: 同时洗牌传输 val 与 idx (或 64 位整型打包)
    Shfl->>Arb: 比较 pair 与 other_pair
    Arb->>Lane: 更新局部胜出元组
    Note over Lane: 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 为非负数，可利用位运算将浮点数的有序编码与 32-bit 整数打包进单个 64-bit 整数（uint64_t），单条洗牌指令同时传输数值与下标。
两阶段 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) {
        if (partial_val[i] > local.val) {
            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) {
        if (partial_val[i] > local.val) {
            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 参考答案

输入阶段（Tuple State Initialization & Grid-Stride Ingestion）：
- 算法分配 num_blocks 个 Block 启动 Pass 1。每个线程在本地寄存器中初始化候选双元组 ValIdx local = {-FLT_MAX, -1}；
- 线程以 Grid-Stride 跨步遍历其负责的一组连续或步进元素。每次遇到新输入 $x_i$ ，将其与当前本地 local 进行字典序比较：若 $x_i > \text{local.val}$ ，更新 (x_i, i)。连续的线程读取连续全局显存，保持合并访存；
状态变化阶段（Two-Pass Hierarchical Reduction & Strict Tie-breaking）：
- Warp 级双元组洗牌：Warp 内 32 个线程调用 __shfl_down_sync 同时交换 val 与 idx。两两比较时执行确定性仲裁规则： $(b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$ ，确保数值相等时物理全局下标较小者胜出；
- Block 级与中间发布：各 Warp 的 Lane 0 将部分结果写入 Shared Memory，二次归约后，每个 Block 的 Thread 0 将该 Block 胜出的 (val, idx) 分别写入显存中间数组 partial_val[blockIdx.x] 与 partial_idx[blockIdx.x]；
输出阶段（Final Single-Block Reduction & Scalar Store）：
- 启动仅有 1 个 ThreadBlock 的 Pass 2 Kernel，并行读取 num_blocks 个候选元组并在 Block 内执行终局树状归约；
- Thread 0 将取得全局绝对最大值、且下标最小的合法索引单个整数写入输出指针 *output。

Q2 参考答案

若比较逻辑仅简单使用单向严格大于 >（即 if (b.val > a.val) a = b;），当输入数据中存在多个数值完全相同的全局最大值时，不同的 Block 划分与调度顺序将导致计算结果失去确定性（Non-deterministic Outcome）：

树状对偶归约的物理非对称性：
- 在折半归约步进中（如 offset = 16, 8, 4, 2, 1），当前线程持有的状态作为左操作数 $a$ ，接收洗牌传来的数据作为右操作数 $b$ ；
- 若 $a.val == b.val$ ，条件 b.val > a.val 返回 false，左操作数 $a$ 维持不变保留下来。这表面上看似乎“保留了左边”，但左操作数对应的全局索引并不保证小于右操作数！
Block 划分改变比较拓扑：
- 在 Block 内部，跨步循环 col += BLOCK_SIZE 会使同一个线程处理跨度很大的不同索引（如索引 0 和 256）；
- 当改变 num_blocks（例如从 64 改为 128）或改变 BLOCK_SIZE 时，相同的两个相等元素（比如下标 10 和下标 500）可能由同一个线程串行比较，也可能被切分到两个不同的 Warp 在不同树枝上并行比较；
- 若没有 (b.val == a.val && b.idx < a.idx) 这一强制下标仲裁，哪一个下标能存活到最终阶段完全取决于物理折半树在何时何地折叠了它们，导致输出索引在 10 与 500 之间随机震荡。

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 --- ::: {.callout-note appearance="simple"} **参考答案版**：包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。 ::: ## 本课目标与完成标准学完后你应能：归约 `(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）**： ```{mermaid} flowchart TD subgraph TupleReduction["ArgMax 二元元组仲裁"] P1["元组 1: (val_1, idx_1)"] --> Cmp{"val_1 > val_2 ?"} P2["元组 2: (val_2, idx_2)"] --> Cmp Cmp -->|是| Win1["胜出: (val_1, idx_1)"] Cmp -->|否: val_2 > val_1| Win2["胜出: (val_2, idx_2)"] Cmp -->|等值: val_1 == val_2| Tie{"min(idx_1, idx_2)"} Tie --> Deterministic["严格保留最小索引 (确定性 Tie-Breaking)"] end ``` 图 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 配置的 100% 确定性。 ### 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）**： ```{mermaid} sequenceDiagram autonumber participant Lane as Warp 内部 32 线程 participant Shfl as __shfl_down_sync 洗牌 participant Arb as 严格仲裁算子 Note over Lane: 步进 16 -> 8 -> 4 -> 2 -> 1 折半归约 Lane->>Shfl: 同时洗牌传输 val 与 idx (或 64 位整型打包) Shfl->>Arb: 比较 pair 与 other_pair Arb->>Lane: 更新局部胜出元组 Note over Lane: 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]`。 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` 为非负数，可利用位运算将浮点数的有序编码与 32-bit 整数打包进单个 64-bit 整数（`uint64_t`），单条洗牌指令同时传输数值与下标。 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) { if (partial_val[i] > local.val) { 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) { if (partial_val[i] > local.val) { 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 参考答案 1. **输入阶段（Tuple State Initialization & Grid-Stride Ingestion）**： - 算法分配 `num_blocks` 个 Block 启动 Pass 1。每个线程在本地寄存器中初始化候选双元组 `ValIdx local = {-FLT_MAX, -1}`； - 线程以 Grid-Stride 跨步遍历其负责的一组连续或步进元素。每次遇到新输入 $x_i$，将其与当前本地 `local` 进行字典序比较：若 $x_i > \text{local.val}$，更新 `(x_i, i)`。连续的线程读取连续全局显存，保持合并访存； 2. **状态变化阶段（Two-Pass Hierarchical Reduction & Strict Tie-breaking）**： - **Warp 级双元组洗牌**：Warp 内 32 个线程调用 `__shfl_down_sync` 同时交换 `val` 与 `idx`。两两比较时执行确定性仲裁规则：$(b.val > a.val) \lor (b.val == a.val \land b.idx < a.idx)$，确保数值相等时物理全局下标较小者胜出； - **Block 级与中间发布**：各 Warp 的 Lane 0 将部分结果写入 Shared Memory，二次归约后，每个 Block 的 Thread 0 将该 Block 胜出的 `(val, idx)` 分别写入显存中间数组 `partial_val[blockIdx.x]` 与 `partial_idx[blockIdx.x]`； 3. **输出阶段（Final Single-Block Reduction & Scalar Store）**： - 启动仅有 1 个 ThreadBlock 的 Pass 2 Kernel，并行读取 `num_blocks` 个候选元组并在 Block 内执行终局树状归约； - Thread 0 将取得全局绝对最大值、且下标最小的合法索引单个整数写入输出指针 `*output`。 ### Q2 参考答案若比较逻辑仅简单使用单向严格大于 `>`（即 `if (b.val > a.val) a = b;`），当输入数据中存在多个数值完全相同的全局最大值时，不同的 Block 划分与调度顺序将导致**计算结果失去确定性（Non-deterministic Outcome）**： 1. **树状对偶归约的物理非对称性**： - 在折半归约步进中（如 `offset = 16, 8, 4, 2, 1`），当前线程持有的状态作为左操作数 $a$，接收洗牌传来的数据作为右操作数 $b$； - 若 $a.val == b.val$，条件 `b.val > a.val` 返回 `false`，左操作数 $a$ 维持不变保留下来。这表面上看似乎“保留了左边”，但**左操作数对应的全局索引并不保证小于右操作数**！ 2. **Block 划分改变比较拓扑**： - 在 Block 内部，跨步循环 `col += BLOCK_SIZE` 会使同一个线程处理跨度很大的不同索引（如索引 0 和 256）； - 当改变 `num_blocks`（例如从 64 改为 128）或改变 `BLOCK_SIZE` 时，相同的两个相等元素（比如下标 10 和下标 500）可能由同一个线程串行比较，也可能被切分到两个不同的 Warp 在不同树枝上并行比较； - 若没有 `(b.val == a.val && b.idx < a.idx)` 这一强制下标仲裁，哪一个下标能存活到最终阶段完全取决于物理折半树在何时何地折叠了它们，导致输出索引在 10 与 500 之间随机震荡。 ### 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} [**← 上一课** 第 11 课：SiLU 与算子融合](lesson11_silu.qmd){.pagination-link .pagination-prev} [**下一课 →** 第 13 课：DeepSeek DeepGEM...](lesson13_deepgemm.qmd){.pagination-link .pagination-next} :::