Step 0: 存储格式 → 逻辑值#

FP4 的元素格式是 E2M1（1-bit sign + 2-bit exponent + 1-bit mantissa），可表示的规格化范围为 ±[0.5, 1.75]，再乘以 2^sf 的 block scale 恢复为实际值。需要特别说明的是，SM100 对 MXFP4 (Microscaling FP4) 的支持是一种混合精度方案：

• 存储位宽：4-bit（E2M1），这是带宽和显存占用的决定因素
• Block scale 位宽：8-bit（UE8M0），每 128 元素共享一个 scale
• Tensor core 内部算术精度：乘法和累加在更高精度下完成

这里的 “FP4 native support” 指的是硬件理解 E2M1 编码和 UE8M0 block scale 的语义，无需软件显式反量化——tensor core 从 SMEM 读到 FP4 字节流后，在内部流水线中自动完成 scale 解码与精度提升，然后执行乘法。中间的反量化值自始至终不存在于通用寄存器（GPREG）中，整个过程是 tensor core 指令（tcgen05.mma.kind::mxf4.block_scale）内部封闭的硬件行为。

至于 tensor core 内部的乘法精度，这里有一个容易忽略的细节：UE8M0 提供的是 8-bit unsigned exponent（范围 2^0 ~ 2^255），而 FP16 仅有 5-bit exponent（范围 2^-14 ~ 2^15）。如果内部乘法使用 FP16，则当 block scale 较大时，反量化后的值将立即溢出——这显然不可接受。因此 Blackwell tensor core 实际使用的内部格式是 FP19（1-bit sign + 8-bit exponent + 10-bit mantissa），其 8-bit exponent 恰好覆盖 UE8M0 的全范围而无需额外钳位。累加器则继续使用 FP32。

Comment: FP19 之于 MXFP4/MXFP6/MXFP8，类似于 FP32 之于 FP16/BF16/FP8——它是 tensor core 内部的”通用运算精度”。8-bit exponent 的设计使得它既能无损容纳 UE8M0 的全范围 scale，又比 FP32 节省了 13-bit mantissa（反正源头只有 4-bit，精度已定）。此外，“中间值不在寄存器上”这一点也对性能有直接影响：tensor core 的输入直接来自 SMEM 和 TMEM，无需将反量化结果回写 GPREG 再读入，省去了一个完整的数据往返。Q 的 128 元素 block size 恰好等于 UMMA_N = BLOCK_Q × num_heads，这意味着每个 UMMA tile 内部共享唯一的 SF_Q 和 SF_KV 组合，避免了 tile 内 SF 切换的开销。

Step 1: TMA 搬运 → SMEM#

SMEM 中的数据保持其原始存储类型。Q 和 KV 的 64B swizzle 布局是 UMMA 指令的对齐要求，数据本身不做类型转换。

Step 2: UMMA 发射 → TMEM 累加器#

UMMA 指令描述符中声明了 A/B 矩阵为 float_e2m1_t，累加器类型为 float。该指令在 SM100 tensor core 上的执行流程为：

1. 从 SMEM 读取 packed FP4（SWIZZLE_64B 解交织后得到 E2M1 元素字节流）
2. 从 TMEM 读取对应的 SF 值（UE8M0，已由 UTCCP 搬运到位）
3. Tensor core 内部自动将 E2M1 解码并与 block scale 组合，提升到内部运算精度后执行乘法——整个过程在 tensor core 流水线内完成，无显式反量化步骤
4. 矩阵乘累加：A(K-major) × B(K-major) → C，累加器为 float（FP32）
5. 结果写入 TMEM 的累加器列

Comment: MXFP4 的关键特性在于 SF 存储在 TMEM 而非 SMEM 中。与 FP8 不同（SF 可作为普通的 SMEM scale vector 传入），FP4 的 kind::mxf4.block_scale 指令要求 SF 必须驻留在 TMEM 内——这解释了为什么 kernel 需要先通过 UTCCP 将 SF_KV 搬到 TMEM 再做 UMMA。TMEM 作为 tensor core 的”近存”，SF 访问延迟极低，使得 block scale 能在每个 UMMA K-step 无开销地切换。

Step 3: TMEM → Math Registers#

Math warp 通过 SM100_TMEM_LOAD_32dp32b32x 指令将累加器从 TMEM 读入寄存器。对于 num_heads = 64，每次加载 32 个 float：

1
tmem_load(Int<32>{}, tmem_addr, accum);
2
tmem_load(Int<32>{}, tmem_addr + 32, accum + 32);

加载后需要 fence_view_async_tmem_load() 保证 TMEM 读可见。

Step 4: ReLU + Weighted Sum（Math 阶段）#

这是唯一在常规寄存器中完成全部计算的步骤：

1
// Type of each operand:
2
sum_0, sum_1 → float2  (initialized with {0.0f, 0.0f})
3
accum[j]     → float   // ReLU input
4
fmaxf(...)   → float   // ReLU output
5
weights[i][j]→ float   // from SMEM, via ld_shared
6
__ffma2_rn   → float2  // FMA on two packed floats at once
7
__fadd2_rn   → float2  // pairwise combine
8
sum.x + sum.y → float  // final scalar reduction

__ffma2_rn 是一条 PTX 级指令，在一个 cycle 内完成两组 a.x*b.x + c.x 和 a.y*b.y + c.y，吞吐是两次标量 FMA 的两倍。四个 head 为一轮（两个 head 归入 sum_0，两个 head 归入 sum_1），最终 pair-wise 合并为单个标量。

Step 5: 写回 Global Memory#

1
auto result = static_cast<logits_dtype_t>(sum.x + sum.y);
2
logits[q_offset + kv_offset] = result;

当 logits_dtype_t 为 BF16 时，static_cast 将 FP32 的低 16 位丢弃（round-to-nearest-even 由硬件/编译器实现）。这是整个流程中唯一的精度损失点——从 FP32 累加器截断到目标输出类型。

数据类型流转总图#

1
Q (packed FP4, ½B/el) ──┐                        Weights (FP32, 4B/el)
2
SF_Q (UE8M0, ¼B/el) ────┤                                  │
3
                          ├─ UMMA: E2M1×2^UE8M0 ── FP19 mul (tensor core 内部, 不经过 GPREG) ── float acc ── TMEM
4
KV (packed FP4, ½B/el) ──┤                                                     │
5
SF_KV (UE8M0, ¼B/el) ────┘                                                     │
6
                                                                                ▼
7
                                                    TMEM float ── load ── registers (float)
8
                                                                                │
9
                                                           ┌────────────────────┤
10
                                                           │ fmaxf (ReLU)       │ ld_shared weights
11
                                                           ▼ float              ▼ float
12
                                                     __ffma2_rn ── float2 ── fadd2 ── float
13
                                                                                           │
14
                                                                            static_cast ──▼
15
                                                                         BF16 / FP16 / FP32 → Global Memory

Comment: 需要特别指出的是，MXFP4 的 block quantization 并非完全无损——128 个元素共享同一个 SF，意味着该 block 内的动态范围被限制为 ±1.75 × 2^sf。精度损失发生在量化阶段（FP32 → E2M1 + SF），而非 tensor core 内部运算。对于 MQA logits 这种 scoring 任务，量化损失通常不敏感于最终排序结果。Weights 保持 FP32 则是因为它在 ReLU 后直接参与乘法，若降低精度会引入额外的量化误差。

5.4.1 Tile 尺寸#

5.4.2 SF 搬运到 TMEM：UTCCP 转置#

Q Scale Factor 在 SMEM 中按 [num_heads, block_q] 布局存放（TMA 写入时 transpose）。在 UMMA 发射前，需要通过 UTCCP (Unified Tensor Core Copy) 将其搬运到 TMEM（Tensor Memory）：

1
// Warp-level transpose via shared memory (UTCCP requires K-major layout)
2
utccp_required_smem_warp_transpose(smem_ptr);
3
cute::SM100_UTCCP_4x32dp128bit_1cta::copy(sf_desc, tmem_col);

UTCCP 要求数据在 SMEM 中为 K-major layout，而 TMA 写入的是 MN-major，因此需要一个 warp-level 的 SMEM 转置步骤（通过 ld_shared / st_shared + XOR lane swapping 实现）。

SF_Q 搬入 TMEM 后位于 tmem_col = kTmemStartColOfSFQ；SF_KV 搬入后位于 kTmemStartColOfSFKV。

5.4.3 发射 UMMA#

1
auto runtime_instr_desc = make_runtime_instr_desc_with_sf_id(
2
    instr_desc, k * 2, k * 2);  // SF ids for A and B
3
auto a_desc = make_smem_desc(SWIZZLE_64B,
4
    smem_kv + i * UMMA_M * (head_dim/2) + k * UMMA_K/2,
5
    8 * (head_dim/2), 0);
6
auto b_desc = make_smem_desc(SWIZZLE_64B,
7
    smem_q + k * UMMA_K/2,
8
    8 * (head_dim/2), 0);
9
SM100_MMA_MXF4_SS::fma(a_desc, b_desc, tmem_addr, k,
10
    runtime_instr_desc,
11
    kTmemStartColOfSFKV + i * 4, kTmemStartColOfSFQ);

其中：

• SM100_MMA_MXF4_SS::fma 封装了 PTX 指令 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32（CUDA ≥ 12.9）或 kind::mxf4.block_scale.scale_vec::2X（旧版本）。
• A 矩阵是 KV（K-major, SWIZZLE_64B），B 矩阵是 Q（K-major, SWIZZLE_64B）。
• runtime_instr_desc 中的 a_sf_id / b_sf_id 分别指向 TMEM 中对应 K 步的 SF 列位置（A→SF_KV, B→SF_Q）。
• 累加结果写入 TMEM 地址 tmem_addr = tmem_stage_idx * UMMA_N。

UMMA 完成后通过 tcgen05.commit + mbarrier 通知 math warps。

5.4.4 TMEM 三级流水线#

与 Q/KV 类似，TMEM 也使用 3 级流水线（kNumTmemStages = 3）。UMMA warp 循环填充 TMEM stage，math warp groups 轮流消费。每个 stage 的容量为 UMMA_N = BLOCK_Q * num_heads = 128 列。

TMEM 总分配列数包括：

1
kNumTmemCols = align(
2
    kNumAccumTmemCols                // BLOCK_Q * num_heads * 3 stages
3
    + kNumSFQ / 32                   // SF_Q columns
4
    + kNumSFKV / 32,                 // SF_KV columns
5
    到最近的 2 的幂
6
)

Comment: 在 Blackwell 上，每个 SM 只有 512 列 TMEM（约 64KB），3 级流水线 + 两份 SF 存储，256 列左右，这在容量预算内。