Bias-Add 与 Quant/Dequant 辅助函数

本页中的所有地址都适用于来自 libtpu-0.0.40-cp314 wheel 的 libtpu.so（build-id 89edbbe81c5b328a958fe628a9f2207d，781,691,048 字节，未 strip）。.text 和 .rodata VMA 等于其文件偏移；.data.rel.ro VMA 减去 0x200000 等于其文件偏移。其他 libtpu 构建会有所不同。

摘要

这些是围绕 TensorCore 上 systolic-array matmul/conv 的数值辅助函数：位于 MXU 输出的原始 int32/f32 accumulator 与最终存储值之间的层。具体包括：bias-add、activation clamp，以及 quantize/dequantize 序列（应用 scale、窄浮点和整数转换、闭合往返的饱和 clamp）。寻找 EmitBias 或 QuantizeOp 的重新实现者不会找到它们，而这种缺席是第一个发现。

这里没有专用 bias op。“matmul + bias” HLO 是一个输出 fusion，其 root 是 conv，body 是普通的逐元素 add(conv, broadcast(bias))，由通用 fusion vectorizer 在 MatrixMultiplyAccumulateFunctor::DoOutputFusion 内按输出 chunk 发出。bias 被广播到 tile 形状（HandleBroadcast），以 f32 相加（HandleAdd → VaddF32），再执行 activation（HandleMaximum/HandleClamp），最后结果被重新量化（HandleConvert）并存储。accumulator dtype 按 dtype 取键：int8 matmul 累加到 int32（VaddS32），并在浮点 epilogue 之前 dequantize 到 f32；bf16/fp32 matmul 累加到 f32（VaddF32），并有两个专用混合精度 add（0x11e/0x11f）把 bf16 high/low pass 折叠进 f32 accumulator。

quant/dequant 叶子是由 LloRegionBuilder::Vcvt*/Vpack*/Vunpack* 构建的单个 LLO vector op。round-to-nearest 窄浮点 quantize 走 EXMY 转换（0x6e/0x6f）或 bf16→fp8 pack；随机舍入走四个 VcvtSrF32To* binop（0x70-0x73），其第二个操作数是逐 lane 随机 dither；int8 quantize 是 round → VclampSymmetric → VpackcB8。在设备快速路径上，量化是对称的（zero-point 隐式为 0）；非对称 / zero-point 模型只存在于 host 侧 MLIR quant dialect 中。对称 absmax 8-bit 量化在二进制中最干净的参考实现是 collective pincer_utils 家族，本文按 op 记录它。

对于重新实现，契约是：

Bias-add 是 fusion，不是 op： AccumulateOutputChunk → DoOutputFusion 主干，以及 VectorizingEmitter Handle* dispatch（HandleBroadcast = bias splat，HandleAdd = bias add）。
Accumulator dtype 分裂： int32 VaddS32（int8 matmul）对 f32 VaddF32（bf16/fp32），再加上 bf16-in-f32 混合 add 0x11e/0x11f。
Convert 叶子： EXMY 窄浮点（0x6e/0x6f）、四个随机舍入 binop（0x70-0x73）、VcvtF32ToS32 舍入方向，以及 lane pack/unpack 辅助函数。
Saturate 层： VclampSymmetric*（原生 + min/max fallback 0x15f）对 VclampAsymmetric*（binop 0x4c + min/max），以及 VclampGez*（ReLU）。
对称 absmax quantizer（pincer_utils）： absmax → scale（qmax/absmax）→ quantize → 8-bit wire → f32 requant-reduce → dequant，并带三个按 dtype 的 max 常量。


Bias-add 路径	`MatrixMultiplyAccumulateFunctor::DoOutputFusion` (`0x13124360`); `HandleBroadcast` (`0x136f1d60`) + `HandleAdd` (`0x13e22e80`)
Accumulator	int8 → int32 `VaddS32` (`0x1d51e9c0`); bf16/fp32 → f32 `VaddF32` (`0x1d525160`); mixed `0x11e`/`0x11f`
RNE 窄浮点	`VcvtEXMYToE4M3`/`E5M2` (`0x6e`/`0x6f`); `VpackCBf16ToE5M2`/`E4M3Fn` (`0x126`)
随机舍入	`VcvtSrF32To{E5M2,E4M3,If8,Bf16}` (`0x70`-`0x73`), binop, rng-bits 第二操作数
f32→int32	`VcvtF32ToS32` (`0x1d524760`): `RoundNearest`→`VroundF32`, 否则 `VtruncF32`，然后 `0x5e`
Clamp	`VclampSymmetricF32`/`Bf16`（原生 + `0x15f` fallback）；`VclampAsymmetric`（`0x4c`）；`VclampGez`（ReLU）
Collective quantizer	`pincer_utils`: `UpdateScale` (`0x137b75c0`), `GetUnpackFunction` (`0x137b8720`); types `{S8, F8E5M2, F8E4M3B11FNUZ}`

Bias-Add 是输出 Fusion Epilogue

目的

重新实现者需要知道，bias、activation 和 output-requantize 不是对物化 matmul 输出的独立 pass。它们按输出 chunk 内联发出，直接从 MXU result register 读取 matmul 结果，这正是“output fusion”的意义。

算法

规范化后的 “matmul + bias + activation” HLO 是一个输出 fusion：

text

fusion {
  conv  = convolution(activations, weights)   // MXU
  bcast = broadcast(bias)                      // bias to output tile shape
  add   = add(conv, bcast)                     // bias add
  [opt] = maximum(add, 0) / clamp(...) / tanh  // activation
  [opt] = convert(...) -> out_dtype            // (re)quantize
}
```text

由按 chunk 的 MXU emission 主干降低：

```c
MatrixMultiplyAccumulateFunctor::operator()                       // 0x1310cd80
  -> AccumOutputChunkAndPossiblyDoOutputFusion                    // 0x13124100
       -> AccumulateOutputChunk                                   // 0x13123a80 : matres -> accumulator
       -> RegisterAccumOutput                                     // 0x131196a0
       -> PossiblyDoOutputFusion                                  // 0x13122d80
            -> DoOutputFusion                                     // 0x13124360

DoOutputFusion（0x13124360）根据 ShapeUtil::ElementIsIntegral/ElementIsFloating(out) 分支，在浮点 epilogue 之前用 VcvtS32ToF32 对 int32 accumulator 做 dequantize（因此即使是 int8 matmul，bias-add 和 activation 也在 f32 中运行），然后通过 VectorizingEmitter Handle* dispatch 逐 op 遍历 fusion body，最后执行 VreducePrecision（bf16 舍入，除非 AllowExcessPrecision）+ Vst。

逐元素 op 链由这些 handler JIT 发出：

Handler	地址	角色
`HandleBroadcast`	`0x136f1d60`	Bias broadcast（splat 到 tile）
`HandleAdd`	`0x13e22e80`	Bias add（`VaddF32`/`VaddS32`，按 dtype 取键）
`HandleConvert`	`0x13e22040`	Fusion 内的 (re)quant/dequant
`HandleMaximum`	`0x13e21f80`	ReLU / clip-low
`HandleClamp`	`0x13e21f20`	两侧 clamp
`HandleMultiply`	`0x13e228c0`	Scale multiply
`HandleReducePrecision`	`0x13e23a40`	bf16 round-to-nearest-even
`HandleStochasticConvert`	`0x13e22520`	HLO `kStochasticConvert`

HandleBroadcast 读取 broadcast 操作数，把 bias splat 到输出 tile（PackOrMakeTuple），并通过 OpaqueCopy 物化。bias 操作数经由 FusionEmitter::AddInput（0x136d5f20）进入。

注意 —— 没有 BiasAddEmitter / EmitBias。Bias-add 是通用 fusion vectorizer 发出 Add(value, Broadcast(bias))。应把“bias”视为普通 fused input edge，而不是特殊 op。

Accumulator Dtype 和 bf16-in-f32 累加

目的

accumulate op 由输出元素类型选择，而高精度 bf16 matmul 模式需要专用混合精度 add。重新实现者必须按 dtype 选择正确的 accumulate op，否则会静默损失精度。

算法

AccumulateOutputChunk（0x13123a80）选择：

int8 / int4 matmul → int32 accumulator → VaddS32（0x1d51e9c0）
bf16 / fp32 matmul → f32 accumulator → VaddF32（0x1d525160）

bf16-in-f32 多 pass 累加（bf16x2/bf16x3 高精度模式）使用两个专用混合 add，把 bf16 结果半部折叠进 f32 accumulator：

Builder method	地址	LLO opcode	角色
`VaddF32MixedBF16HighInst`	`0x1d55dd40`	`0x11e`	把 bf16 HIGH pass 累加进 f32
`VaddF32MixedBF16LowInst`	`0x1d55dd80`	`0x11f`	把 bf16 LOW pass 累加进 f32
`VaddF32`	`0x1d525160`	(binop)	普通 f32 accumulate
`VaddS32`	`0x1d51e9c0`	(binop)	普通 int32 accumulate

反编译确认 VaddF32MixedBF16HighInst 调用 CreateVectorBinop(0x11E, …)。accumulate 始终是 f32 宽度（systolic array 乘 bf16，但 partial product 累加到 f32 寄存器文件），这就是 fp32 和 bf16 共享相同逐 matmul MXU latency 的原因。

Convert 叶子：RNE 和随机舍入

目的

每个 fp8/int8 quantize 及其反向 dequant 都是一个单独 LLO op。重新实现者需要每个转换的 opcode 和创建器，才能让转换成为一条指令，而不是仿真序列（在有能力的 target 上）。

编码 —— round-to-nearest 窄浮点

function VcvtEXMYToE4M3(in, exp_bits):                  // 0x1d52b440
    assert target().SupportsFloat8EXMY()
    inst = CreateVectorEXMYConversion(110 /*0x6e*/, …)
    return region->AppendInstruction(inst)
```text

| Method | 地址 | LLO opcode | Creator |
|---|---|---|---|
| `VcvtEXMYToE4M3` | `0x1d52b440` | `0x6e` | `CreateVectorEXMYConversion`（门控 `SupportsFloat8EXMY`） |
| `VcvtEXMYToE5M2` | `0x1d52b520` | `0x6f` | `CreateVectorEXMYConversion` |
| `VpackCBf16ToE5M2` | `0x1d567a40` | `0x126` | `CreateVectorPack`（VpackFormat） |
| `VpackCBf16ToE4M3Fn` | `0x1d567b20` | `0x126` | `CreateVectorPack`（VpackFormat） |
| `VcvtF32ToNarrowFloat` | `0x1d560960` | (composite) | `VimmIf`+`VabsF32`+`SimplifyPackF16`（软件 RNE） |

### 编码 —— 随机舍入

每个 Sr convert 都是一个 binop，第二个操作数是逐 lane 随机 dither，因此转换会以等于小数部分的概率向上舍入：

```c
function VcvtSrF32ToE5M2(src_f32, rng_bits):            // 0x1d52b600
    inst = CreateVectorBinop(0x70, src_f32, rng_bits, region)
    return region->AppendInstruction(inst)

Method	地址	LLO opcode
`VcvtSrF32ToE5M2`	`0x1d52b600`	`0x70`
`VcvtSrF32ToE4M3`	`0x1d52b640`	`0x71`
`VcvtSrF32ToIf8`	`0x1d52b680`	`0x72`
`VcvtSrF32ToBf16`	`0x1d52b6c0`	`0x73`
`VcvtF32ToS32Stochastic`	`0x1d52ade0`	`0x5e`（fast）/ `0x15c`（sw fallback）

反编译确认 VcvtSrF32ToE5M2 → CreateVectorBinop(0x70u, …)，VcvtSrF32ToE4M3 → CreateVectorBinop(0x71u, …)。在有能力的 target 上，fp8/bf16 Sr convert 没有软件仿真：它们各自都是一个 op。HLO kStochasticConvert（HandleStochasticConvert 0x13e22520）会对 fp8/bf16 输出 dispatch 到这个家族，对整数输出 dispatch 到 VcvtF32ToS32Stochastic。

按代际的 stochastic 支持：Target::SupportsVectorConvertF32Stochastic::kSupportedTypes = {10, 16, 19, 23} = {F16, BF16, F8E5M2, F8E4M3B11FNUZ}，在 Ghostlite（0xb530678）和 Viperfish（0xb530784）上均如此。F8E4M3Fn（20）不在 Sr 支持集合中：只有 B11FNUZ fp8 variant 有原生 f32→fp8 随机舍入。

编码 —— f32 → int32 舍入方向

VcvtF32ToS32(value, RoundDirection)（0x1d524760）：RoundNearest（1）会前置 VroundF32，否则（TowardsZero）前置 VtruncF32，然后发出 opcode 0x5e。具名 pseudo-variant：VcvtF32ToS32TiesToEven（0x1d52ab00）、...TowardsZeroPseudo（0x1d52ad20）、...KnownIntegral（0x1d52ad80）、...Stochastic（0x1d52ade0）。

饱和 Clamp

目的

quantize 往返和 activation 以饱和 clamp 到目标范围结束。对称与非对称的分裂，是以两个 builder 家族呈现的 device-symmetric 与 host-zero-point 区分。

编码

function VclampSymmetricF32(value, bound):              // 0x1d527fc0
    if target supports native clamp:
        inst = CreateVectorClampSymmetric(value, bound, …)
    else:                                                // fallback
        v = CreateVectorMinimumF32(value, +bound)        // opcode 0x15f path
        inst = CreateVectorMaximumF32(v, -bound)
    return region->AppendInstruction(inst)
```text

| Method | 地址 | Body |
|---|---|---|
| `VclampSymmetricF32` | `0x1d527fc0` | `CreateVectorClampSymmetric` + min/max（`0x15f`）fallback |
| `VclampSymmetricBf16` | `0x1d5283a0` | 同上 |
| `VclampAsymmetricF32` | `0x1d528480` | binop `0x4c` + `Minimum`/`Maximum`（lo ≠ −hi） |
| `VclampAsymmetricBF16` | `0x1d528740` | 同上 |
| `VclampGezF32` / `Bf16` | `0x1d527ea0` / `0x1d527ee0` | clamp ≥ 0（ReLU） |
| `VclampS32` / `U32` / `Float` | `0x1d528820` / `0x1d528c40` / `0x1d528fe0` | integer / generic clamp |

`VclampSymmetricF32` 的反编译同时显示原生 `CreateVectorClampSymmetric` 分支，以及带 opcode `0x15f` 的 `CreateVectorMinimumF32`/`CreateVectorMaximumF32` fallback。**Symmetric** 饱和到 `[-bound, +bound]`（zero-point 0），即 int8/fp8 设备 saturate。**Asymmetric** 饱和到 `[lo, hi]`，其中 `lo ≠ -hi`，即带 zero-point / unsigned 的 host 路径。

`VreducePrecision`（`0x1d5697c0`，从 `DoOutputFusion` 和 `HandleReducePrecision` 调用）是把 f32 mantissa 截断为 bf16 的 bit-manipulation round-to-nearest-even：向 mantissa 加入舍入 bias，mask 掉低位，修正 tie-to-even 和 NaN/Inf 指数。LLO op 是 `llo.vround.nearest_even.f32`（`VroundF32` `0x1d525480`）。

---

## Collective 对称 Quantizer（`pincer_utils`）

### 目的

quantized all-reduce 在 wire 写入前把每个 shard 量化到 8 bit，用正在运行的 scale 以 8-bit-on-wire 形式 reduce，并在结束时 dequantize。这些 `xla::jellyfish::pincer_utils` 辅助函数是对称 absmax 8-bit 量化在二进制中最干净的参考，重新实现者可以直接抽取其 scale 公式。

### 算法

```c
function UpdateScale(dtype, builder, max_abs_addr, scale_addr):   // 0x137b75c0
    switch dtype:
      case S8(0x2):           dtype_max = 127.0        // dword_84A2A28
      case F8E4M3B11FNUZ(0x17): dtype_max = 30.0       // dword_84A27FC
      case F8E5M2(0x13):      dtype_max = 57344.0      // dword_84A2530
      default: LogFatal("quantized_type_bound != nullptr") + PrimitiveType_Name(dtype)
    m       = VimmF32(dtype_max)
    max_abs = Vld(max_abs_addr)                         // expected in VMEM (checked)
    scale   = VdivF32(m, max_abs)                       // scale_factor = qmax / absmax
    Vst(scale_addr, scale)

反编译精确固定了 dtype switch 和三个 .rodata float，并且重要的是固定了除法操作数顺序为 VdivF32(dtype_max, max_abs)，即存储的 scale factor 是 qmax / absmax（quantize 乘数）。然后 Quantize 为 q = round(x * scale_factor)；dequantize 为 x = q / scale_factor。按 dtype 的 max 常量：

`.rodata` 地址	值	dtype	PrimitiveType
`0x84a2a28`	`127.0`	int8 max	`0x2`
`0x84a27fc`	`30.0`	F8E4M3B11FNUZ max finite	`0x17`
`0x84a2530`	`57344.0`	F8E5M2 max finite	`0x13`

五阶段往返：

Helper	地址	角色
`UpdateMaxLocalChunk`	`0x137b73a0`	运行中的 absmax：`Vld` → `VunpackCF32`（若 bf16）→ `VabsF32` → `VmaxF32` → `Vst`
`UpdateScale`	`0x137b75c0`	absmax → scale（`qmax/absmax`）；3 个 dtype-max 常量
`SymmetricallyQuantizeShardInPlaceTo8Bits`	`0x137b7740`	`round(x*scale)` → `VpackcB16` → `VpackcB8`（×4 展开）
`SymmetricallyDequantizeShardInPlace8Bit`	`0x137b7fc0`	`unpack` → `VcvtS32ToF32` → `/scale` → `VpackCBf16`
`ReduceSymmetricallyQuantized8BitShardInPlace`	`0x137b8880`	f32 requant-reduce（dequant → merge → 重新跟踪 absmax）
`GetUnpackFunction`	`0x137b8720`	dtype → unpack fn pointer

GetUnpackFunction（0x137b8720）映射 S8(0x2)→VunpackCS8、F8E5M2(0x13)→VunpackCF8E5M2、F8E4M3B11FNUZ(0x17)→VunpackCF8E4M3B11，否则返回携带 PrimitiveType 名称的 StatusOr 错误。ring reduction（ReduceSymmetricallyQuantized8BitShardInPlace）在 f32 中完成：dequant → merge functor → 重新跟踪 absmax，因此 8-bit 格式只是 wire 表示，而不是算术格式。RotatedPincerQuantizedEmitter::kSupportedQuantizationTypes（0xae5b90c）= {2, 19, 23} = {S8, F8E5M2, F8E4M3B11FNUZ}。

注意 —— 设备侧对称 quantizer（pincer_utils + Vcvt*）没有 zero-point：on-device fast path 是对称的（zero_point == 0）。非对称 / zero-point 量化只存在于 host 侧 MLIR quant-dialect 层。

Per-Tensor vs Per-Channel vs Sub-Channel Scale（Host）

在设备 codegen 路径上，量化是对称的，粒度是 shard / vector tile（一个 scale，隐式 zero-point 0）。更丰富的 scale + zero-point 模型存在于 host 侧 MLIR quant dialect 中，并完整链接进 libtpu：

MLIR type	Scale	Zero-point	Extra fields	粒度
`UniformQuantizedType`	`getScale()`	`getZeroPoint()`	—	Per-tensor
`UniformQuantizedPerAxisType`	`getScales()[]`	`getZeroPoints()[]`	`getQuantizedDimension()`	Per-channel
`UniformQuantizedSubChannelType`	`getScales()[]`	`getZeroPoints()[]`	`getBlockSizes()[]`, `getQuantizedDimensions()[]`	Block-wise
`CalibratedQuantizedType`	`getMin()`/`getMax()`	—	(no scale yet)	Calibration

这些会通过 mlir::quant::stablehlo::ConvertUniform{Quantize,Dequantize,Requantize,QuantizedDotGeneral,QuantizedConvolution,QuantizedAdd,QuantizedClipByValue,...}Op（十个不同的 ConvertUniform*Op 家族，暴露为 40 个 QuantizedStablehloOpConversion::matchAndRewrite 实例）降低到 integer dot/conv + affine requant round((scale_in/scale_out)·(acc - zp_in)) + zp_out。guard 字符串 "Cannot requantize while changing quantization_axis" 确认 requantize 不能移动 per-channel axis。TPU dynamic per-column int8 quantizer（convert_dynamic_quantize_ops → damax_output → per-column scale）由 xla_tpu_experimental_enable_dynamic_int8_quantization flag 门控。

示例：int8 Matmul + Bias + ReLU → int8

text

HLO (post DotCanonicalizer + output fusion):
  fusion {
    conv = convolution(int8 act, int8 weights) : s32   // MXU, int32 accumulate
    bc   = broadcast(f32 bias)
    deq  = convert(conv) : f32
    add  = add(deq, bc)
    relu = maximum(add, 0)
    req  = convert(relu) : s8
  }

Emission per output chunk (inside MatrixMultiplyAccumulateFunctor):
  AccumulateOutputChunk:  int32 accumulator (VaddS32 over K-tiles)
  DoOutputFusion:
    VcvtS32ToF32(acc)                       // int32 -> f32 dequant
    HandleBroadcast(bias) -> OpaqueCopy     // bias splat to tile
    HandleAdd            -> VaddF32         // + bias
    HandleMaximum        -> VclampGezF32    // ReLU
    HandleConvert        -> VmulF32(.,scale) + VcvtF32ToS32(RNE)
                            + VclampSymmetricF32(127) + VpackcB16/VpackcB8   // requant s8
    Vst                                     // store int8 output
```text

---

## 未解码内容

- host 侧 affine-requant 常量（`M0`、shift），位于 `ConvertUniformRequantizeOp` / `ConvertUniformQuantizedDotOp` 内：已恢复 pattern symbol 和结构公式；`matchAndRewrite` 未逐 op 展开（upstream open-source legalization）。
- TPU `DynamicQuantize` custom-op emitter body（per-column damax → scale → int8）和 `xla::jellyfish::QuantizationConfig` proto 字段布局。
- 在 `0x126 CreateVectorPack` 内区分 E5M2 / E4M3Fn / E4M3B11 的 `VpackFormat` ordinal（三个 bf16→fp8 pack 方法已具名；enum 值未在此绑定，见 [Pack/Unpack Precision](pack-unpack-precision.md)）。
- `VcvtF32ToNarrowFloat` 软件 RNE composite 对每个 fp8 格式的精确 bit 布局（exponent bias、subnormal flush、overflow-to-Inf vs saturate）。
- `VclampAsymmetric` lo/hi 操作数来源（quant-type zero-point 还是显式 fusion clip 常量）：已确认是 asymmetric（`lo ≠ -hi`），但未追踪操作数来源。

---

## 交叉引用

- [Pack/Unpack Precision](pack-unpack-precision.md) —— `VpackFormat` 枚举、`set_pack_format_sublane`，以及 quant pack/unpack 辅助函数共享的 bf16↔f32 lane pack/unpack
- [VPU（Vector-ALU）Slot](slot-vpu.md) —— 承载 `vcvt`/`vclamp`/`vadd` opcode immediate 和 `.sr` stochastic-round convert 家族的 VALU slot
- [EUP / Transcendental Slot](slot-eup-transcendental.md) —— `SupportsBf16AluInstructions` 以及 bf16 accumulate 和 unpack 背后的 lane-width 模型
- [XLU Op 名册](xlu-op-roster.md) —— 此 convert/clamp 家族所在的更广泛 XLU op-to-factory 表
- [Bundle 模型](bundle-model-overview.md) —— convert/clamp/add op 按代际调度进入的 VLIW bundle

Bias-Add 与 Quant/Dequant 辅助函数 ​

摘要 ​

Bias-Add 是输出 Fusion Epilogue ​

目的 ​

算法 ​

Accumulator Dtype 和 bf16-in-f32 累加 ​

目的 ​

算法 ​

Convert 叶子：RNE 和随机舍入 ​

目的 ​

编码 —— round-to-nearest 窄浮点 ​

编码 —— f32 → int32 舍入方向 ​

饱和 Clamp ​

目的 ​

编码 ​

Per-Tensor vs Per-Channel vs Sub-Channel Scale（Host） ​

示例：int8 Matmul + Bias + ReLU → int8 ​

Bias-Add 与 Quant/Dequant 辅助函数

摘要

Bias-Add 是输出 Fusion Epilogue

目的

算法

Accumulator Dtype 和 bf16-in-f32 累加

目的

算法

Convert 叶子：RNE 和随机舍入

目的

编码 —— round-to-nearest 窄浮点

编码 —— f32 → int32 舍入方向

饱和 Clamp

目的

编码

Per-Tensor vs Per-Channel vs Sub-Channel Scale（Host）

示例：int8 Matmul + Bias + ReLU → int8