TensorRT-LLM 源码完全拆解（十四）：量化全链路——checkpoint 怎么读、权重怎么量化、Linear 怎么选方法

容量：权重和 KV 缓存是两笔独立的账

以一个 70B 稠密模型为例（$P \approx 70\times 10^9$ 个参数）：

• BF16 权重：$2 \times 70\text{G} = 140$ GB；
• FP8 权重：约 70 GB；
• NVFP4（4-bit）权重：约 35 GB，再加缩放因子开销。

权重是静态的，第二笔账 KV 缓存则是随序列长度线性增长的动态开销。对于采用 MLA 或者 GQA 的模型，单个 token 的 KV 缓存大小可以写成：

其中 $L$ 是层数，$H_{kv}$ 是 KV 头数（GQA 下远小于注意力头数），$d_h$ 是每头维度，$b$ 是每个元素的字节数，系数 2 来自 K 和 V 两份。把 $b$ 从 2（BF16）换成 1（FP8）或 0.5（NVFP4），KV 缓存的容量上限直接翻倍或翻四倍。这两笔账都直接决定「这张卡能不能装下这个模型、这个上下文」。

带宽：解码阶段的真正瓶颈

自回归解码每生成一个 token 都要把全部权重从 HBM 读进 SM 一次。于是单 token 的解码时间是：

$BW_{\text{HBM}}$ 是显存带宽。这个式子说明：在 batch 不够大、算力没被喂满的解码阶段，推理速度与每次前向要搬运的字节数成正比。权重从 FP16 压到 FP8，理论解码吞吐直接翻倍；压到 FP4 再翻一倍。这就是量化在推理侧最直接的收益，与「算力省不省」无关。

Kernel：量化把「乘加」换成了「解码 + 乘加」

代价在第三层。Tensor Core 的 FP8/FP4 指令不是通用的：mma 系列在 FP8 下要求 e4m3 输入、要求缩放因子以特定布局铺在内存里，Blackwell 的 FP4 指令更要求把 e2m1 的两两打包和 E4M3 的 block scale 交错（interleave）成固定块结构。于是「权重怎么摆、scale 怎么存、激活什么时候量化」全部变成 kernel 级的设计问题——这部分正是 TensorRT-LLM 量化子系统在源码层面要解决的事。

一个真实的权衡表：

这张表里最值得记住的是最后一行与「内核复杂度」这一列的关系：收益越大，需要写的 kernel 越多。

QuantAlgo：一串把算法名与位宽绑在一起的枚举

读取路径的终点是一个 QuantConfig，它的 quant_algo 字段取自 tensorrt_llm/quantization/mode.py 里的 QuantAlgo 枚举。这个枚举不只是名字列表，每一项都隐含了「权重位宽、激活位宽、缩放粒度」三元组：

1
class QuantAlgo(StrEnum, metaclass=BaseEnumMeta):
2
    W8A16 = auto()
3
    W4A16 = auto()
4
    W4A16_AWQ = auto()
5
    W4A8_AWQ = auto()
6
    W8A16_GPTQ = auto()
7
    W4A16_GPTQ = auto()
8
    W8A8_SQ_PER_CHANNEL = auto()
9
    W8A8_SQ_PER_TENSOR_PLUGIN = auto()
10
    W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN = auto()
11
    W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN = auto()
12
    W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN = auto()
13
    W4A8_QSERVE_PER_GROUP = auto()
14
    W4A8_QSERVE_PER_CHANNEL = auto()
15
    FP8 = auto()
16
    FP8_PER_CHANNEL_PER_TOKEN = auto()
17
    FP8_BLOCK_SCALES = auto()
18
    NVFP4 = auto()
19
    W4A8_NVFP4_FP8 = auto()
20
    W4A8_MXFP4_FP8 = auto()
21
    W4A8_MXFP4_MXFP8 = auto()
22
    W4A16_MXFP4 = auto()
23
    MXFP8 = auto()
24
    W4A16_NVFP4 = auto()
25
    NVFP4_AWQ = auto()
26
    NVFP4_ARC = auto()
27
    NO_QUANT = auto()

命名规则是 W<bits>A<bits> 加算法后缀：NVFP4 表示权重与激活都是 FP4，W4A16_NVFP4 表示权重 FP4、激活不量化（也就是常说的「weight-only」），W4A8_NVFP4_FP8 表示权重 FP4、激活 FP8。后缀是量化算法的来源：_AWQ 表示用 AWQ 的按通道缩放选显著权重，_GPTQ 表示 GPTQ 的二阶误差补偿，_SQ 表示 SmoothQuant 的等价平滑迁移，_MXFP4 / _MXFP8 表示 OCP microscaling 格式。这些算法的原理在本站已有专门文章拆解，例如 AWQ 完全拆解、GPTQ 完全拆解 与 SmoothQuant 完全拆解。

kv_cache_quant_algo 是另一条独立的枚举，只有 FP8、INT8、NVFP4 三个可选值——KV 缓存的量化粒度与权重完全不同，它是 per-head 或 per-tensor 的，不参与权重那套 per-group 讨论。

QuantMode：位标志如何把枚举翻译成执行条件

QuantAlgo 是给人看的名字，QuantMode 则是给代码看的位掩码。它是一个 IntFlag，每一位代表一个执行条件：

1
class QuantMode(IntFlag):
2
    INT4_WEIGHTS = auto()
3
    INT8_WEIGHTS = auto()
4
    ACTIVATIONS = auto()
5
    PER_CHANNEL = auto()
6
    PER_TOKEN = auto()
7
    PER_GROUP = auto()
8
    INT8_KV_CACHE = auto()
9
    FP8_KV_CACHE = auto()
10
    FP8_QDQ = auto()
11
    FP8_ROWWISE = auto()
12
    FP8_1x128_128x128 = auto()
13
    NVFP4 = auto()
14
    NVFP4_KV_CACHE = auto()
15
    W4A8_NVFP4_FP8 = auto()
16
    W4A8_MXFP4_FP8 = auto()
17
    W4A8_MXFP4_MXFP8 = auto()
18
    W4A16_MXFP4 = auto()
19
    MXFP8 = auto()
20
    COUNT = auto()

注意注释里那句 KEEP BELOW DEFINITION IN SYNC WITH cpp/tensorrt_llm/common/quantization.h：这个位掩码同时被 Python 与 C++ 两侧解读，因为最终执行 GEMM 的是 C++/CUDA 侧。这就是量化子系统最容易出问题的地方——一份解释在两处实现，任何一边加位都要同步。

QuantMode.from_quant_algo() 负责把枚举翻成掩码，它的逻辑本身就说明了每类算法的确切含义。摘几段关键分支：

1
if quant_algo == QuantAlgo.FP8:
2
    mode = QuantMode.use_fp8_qdq()
3
elif quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN:
4
    mode = QuantMode.use_fp8_rowwise()
5
elif quant_algo == QuantAlgo.FP8_BLOCK_SCALES:
6
    mode = QuantMode.use_fp8_block_scales()
7
elif quant_algo == QuantAlgo.NVFP4:
8
    mode = QuantMode.use_nvfp4()
9
elif quant_algo == QuantAlgo.W4A16_NVFP4:
10
    mode = QuantMode.use_w4a16_nvfp4()
11
elif quant_algo == QuantAlgo.W4A8_NVFP4_FP8:
12
    mode = QuantMode.use_w4a8_nvfp4_fp8()
13
elif quant_algo == QuantAlgo.MXFP8:
14
    mode = QuantMode.use_mxfp8()

再看 from_description 的约束，它把「哪些组合在数学上讲不通」写成了断言式的早期报错：

1
# We must quantize weights when we quantize activations.
2
if quantize_activations and not quantize_weights:
3
    raise_error()
4

5
# If we set per_token or per_channel, we must quantize both weights and activations.
6
if (per_token or per_channel) and not (quantize_weights and quantize_activations):
7
    raise_error()

第一条约束看起来是约定，其实是实现限制：没有哪块 Tensor Core 支持「权重 FP16、激活 FP8」的混合输入，所以「激活量化必然伴随权重量化」。第二条则来自 per-channel 与 per-token 缩放因子的定义——它们是绑在 GEMM 的某一维上的，只有当权重和激活都被量化成同一种定点/浮点格式时才有意义。

值得单独看的是 FP8 rowwise 这一支：

1
if use_fp8_rowwise:
2
    mode = mode | QuantMode.FP8_ROWWISE | QuantMode.PER_TOKEN | QuantMode.PER_CHANNEL

一行代码同时点亮三个位：FP8 rowwise 的本质就是「权重按输出通道给 scale（per-channel）、激活按 token 给 scale（per-token）」，也就是 DeepSeek 那套 1×128 与 128×128 缩放之外的另一种常见形态。位掩码在这里的好处是可以直接用 has_fp8_rowwise() 这种查询函数做判断，而不用去记忆枚举之间的映射。

三类 checkpoint 描述文件

读了什么文件决定走哪条解析分支。TensorRT-LLM 至少支持四种来源：

1. ModelOpt 的 hf_quant_config.json（或 legacy 的 quant_cfg.json），里面是一个 quantization 字典，字段包括 quant_algo、kv_cache_quant_algo、group_size、exclude_modules，AWQ 系还额外带 has_zero_point 与 pre_quant_scale。
2. llm-compressor / compressed-tensors 的 config_groups 格式，用 weights.strategy 与 input_activations.strategy 描述粒度。
3. HuggingFace 原生 FP8 checkpoint，靠 quant_method: "fp8" 与 weight_block_size 识别。
4. MXFP4 / MXFP8 的 quant_method 标记。

1
quant_config.quant_algo = (
2
    QuantAlgo(canonicalize_quant_algo(json_quant_configs['quant_algo']))
3
    if json_quant_configs.get('quant_algo') is not None else None)
4
quant_config.kv_cache_quant_algo = (
5
    QuantAlgo(json_quant_configs['kv_cache_quant_algo']) if
6
    json_quant_configs.get('kv_cache_quant_algo') is not None else None)
7
quant_config.group_size = json_quant_configs.get('group_size', None)
8
quant_config.exclude_modules = json_quant_configs.get('exclude_modules', None)
9
if 'has_zero_point' in json_quant_configs:
10
    quant_config.has_zero_point = json_quant_configs['has_zero_point']
11
if 'pre_quant_scale' in json_quant_configs:
12
    quant_config.pre_quant_scale = json_quant_configs['pre_quant_scale']

1
if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION:
2
    ...
3
    mixed_quant_configs = json_quant_configs.get('quantized_layers', None)
4
    ...
5
    for layer in mixed_quant_configs:
6
        layer_cfg = mixed_quant_configs[layer]
7
        config = QuantConfig()
8
        config.kv_cache_quant_algo = kv_cache_quant_algo
9
        config.quant_algo = QuantAlgo(canonicalize_quant_algo(layer_cfg['quant_algo']))
10
        config.group_size = layer_cfg.get('group_size', None)
11
        ModelConfig._validate_nvfp4_group_size(config, layer)
12
        ...
13
    layer_quant_config = mixed_quant_configs

1
if kv_quant_lhs != kv_quant_rhs:
2
    raise RuntimeError(
3
        f"The kvcache config in 'quant_cfg.json', {kv_quant_lhs},"
4
        f"is different from 'hf_quant_config.json', {kv_quant_rhs}!"
5
    )

1
# Set default exclude_modules for FP8_BLOCK_SCALES
2
# kv_b_proj must always be excluded: FP8 128x128 block boundaries
3
# don't necessarily align with per-head dim boundaries (e.g. GLM-5
4
# has qk_nope_head_dim=192), so the scale tensor cannot be cleanly
5
# reshaped per-head. The dequant path handles this correctly.
6
default_exclude = ["*kv_b_proj*", "*k_b_proj*", "*eh_proj"]

权重张量名与 scale 名：约定在哪

读取的最后一步是「按名字找张量」。这一步依赖约定，而约定分散在几处：

• params_map = dict(FUSED_MODULE_COMPONENTS)：把 checkpoint 里的独立张量名（如 q_proj、k_proj、v_proj）映射到运行时融合模块的组件名，用于 FUSED_QKV 模式。
• _get_scale_name：FP8 block scales 的 scale 张量在不同 recipe 下可能叫 weight_scale_inv（DeepSeek recipe）或 weight_scale（ModelOpt recipe）：

1
def _get_scale_name(self, weights: List[Dict]):
2
    # `weight_scale_inv` for DS recipe and `weight_scale` for ModelOpt recipe.
3
    # Actually they hold identical values of data_amax / 448.
4
    for w in weights:
5
        if "weight_scale_inv" in w:
6
            return "weight_scale_inv"
7
    return "weight_scale"

注释指出了关键事实：两个名字的数值含义相同（都是 data_amax / 448，即「原始值除以最大值」，反量化时相乘），只是历史命名不同。理解这一点，读其他框架的 FP8 checkpoint 时就不会被名字误导。

• tie_word_embeddings 与 lm_head 的跳过逻辑：共享词表时 lm_head 与 embed_tokens 是同一份权重，只能加载一次。_resolve_lm_head_quant_config 则处理「lm_head 是否量化」这个细节问题——很多 checkpoint 的 lm_head 保持高精度，因为它直接决定输出分布。

NVFP4：两层 scale 的设计与那个 448x6

NVFP4 是 Blackwell 引入的 4-bit 浮点格式，它的量化方案不是「一个 scale 走天下」，而是两层：

1. 每 16 个元素一个 FP8 (e4m3) 的 block scale（$g=16$）。之所以是 16 而不是 32，来自硬件指令 block scaling 的粒度选择；
2. 整个张量再加一个 FP32 的 per-tensor global scale。

代码里 NVFP4LinearMethod.create_weights 的注释把这两个数字的含义直接写了出来：

1
# FP32 per-tensor global scaling factor = 448*6/amax_input
2
module.input_scale = Parameter(torch.empty([1], dtype=torch.float32), ...)
3

4
# (amax_input * amax_weight) / (448*6 * 448*6)
5
module.alpha = Parameter(torch.empty([1], dtype=torch.float32), ...)
6

7
# Global weight scale: amax_weight / (448*6)
8
module.weight_scale_2 = Parameter(torch.empty([1], dtype=torch.float32), ...)

三个常数要一起看：

• 448 是 e4m3 的最大值（FP8 block scale 的编码范围）；
• 6 是 e2m1 的最大值（FP4 权重/激活能表示的绝对值上限：e2m1 的 $2^2 \times (1 + 0.5) = 6$）；
• $448 \times 6 = 2688$ 就是「FP8 scale 乘 FP4 元素」这个二级结构能覆盖的整体动态范围上限。

alpha 这个名字容易让人困惑，它的物理含义是「这次 GEMM 输出相对于输入的总放缩倍数」。前向计算时激活被量化成 FP4，乘上 input_scale = 2688/amax_input（把激活归一化到 FP4 范围），权重同样按 weight_scale_2 = amax_weight/2688 归一化，两者相乘后必须再乘回 alpha 才能还原数值：

1
FP8_MAX, E2M1_MAX = 448.0, 6.0
2
amax_input = torch.amax(torch.abs(input)).float()
3
input_scale = FP8_MAX * E2M1_MAX / amax_input
4
alpha = (amax_input / (FP8_MAX * E2M1_MAX)) * module.weight_scale_2

动态模式（force_dynamic_quantization 为真或静态 scale 缺失）下，input_scale 与 alpha 在每次前向时现算，因为激活的最大值随 batch 与 token 变化；静态模式下二者在加载完成后由 _finalize_nvfp4_scales 固化下来，省掉一次全张量 amax 归约。这就是 NVFP4 的「静态/动态激活量化」开关，也是实测中精度与延迟的调节旋钮。

还有一个容易踩的约束：input_scale 与 weight_scale_2 是标量（shape 为 [1] 的 FP32），在 TP 切分后各 rank 的这两个值必须一致。_finalize_nvfp4_scales 里为此做了断言：

1
for s in input_scale_list[1:]:
2
    assert torch.allclose(input_scale_list[0], s), \
3
        f"input_scale mismatch across shards: {input_scale_list}"

如果某个 rank 的激活 amax 统计不一致，这里会直接报错而不是悄悄给出偏差结果。

weight_scale 的物理布局：为什么必须 interleave

NVFP4 的 block scale 在内存里不是简单地按行铺开。load_weight_scales 在 vanilla 路径下做了一次转换：

1
ws = module.load_shard(w["weight_scale"], device=device, scale_span=scale_span).contiguous()
2
ws = ws.view(fp4_utils.float4_sf_dtype)
3
ws = torch.ops.trtllm.block_scale_interleave(ws)
4
copy_weight(module.weight_scale, ws)

block_scale_interleave 是硬件要求：FP4 GEMM 的 Tensor Core 指令以固定的 tile（例如 128 行 x 4 列 scale）为单位读取 scale 张量，如果不做交错，指令取到的 scale 与它正在计算的元素块就对不上。同时 create_weights 里对 scale 张量的尺寸做了补零：

1
nrows = fp4_utils.pad_up(out_features, 128)
2
ncols = fp4_utils.pad_up(in_features // module.scaling_vector_size, 4)
3
module.weight_scale = Parameter(torch.empty([nrows * ncols], dtype=fp4_utils.float4_sf_dtype), ...)

pad_up(x, y) 是把 x 向上取整到 y 的倍数。补零的原因同样来自硬件：tile 尺寸是 128x4，若张量尺寸不是它的整数倍，最后一块会读越界，所以必须补齐。get_fp4_shape 里的 scale_shape = pad_up(m, 128) * pad_up(K // sf_vec_size, 4) 是同一个公式的另一种写法。这类「先补零再交错」的操作在 FP4 kernel 里是标配，也是为什么 FP4 权重的显存占用不是精确的 0.5 byte/元素。

fp4_utils 里还有一个值得看的函数——reorder_rows_for_gated_act_gemm：

1
def get_reorder_rows_for_gated_act_gemm_row_indices(x) -> torch.Tensor:
2
    # Reorders rows in the gemm/MOE_gemm weight matrix for min-latency
3
    # [r0, r1, r2, ..., rN/2, r(N/2+1), .. r(N-1)]
4
    # to
5
    # [r0, rN/2, r1, rN/2+1, ..., r(N/2-1), r(N-1)]

这是 SwiGLU 融合 GEMM 的需求：gate 与 up 两个投影的权重在 checkpoint 里是「gate 的全部行、up 的全部行」顺序拼接，但融合 kernel 希望两者交错排列，这样同一个线程块可以同时算出 gate(x) 与 up(x) 并立刻做逐元素相乘。这种「加载时重排权重行」的做法，把布局转换的开销从每次前向挪到了加载阶段——量化子系统里反复出现这个思路。

一个 Linear 模块的三段生命周期

TensorRT-LLM 的 Linear 模块不是简单地持有 nn.Parameter 然后做矩阵乘。它的生命周期分成三段，每段都可能被量化方法改写：

1. 构造（create_weights_impl）：调 self.quant_method = self.get_quant_method(self.quant_config)，再让方法分配权重张量。此时张量是空的，只是按量化格式决定了形状与 dtype；
2. 加载（load_weights 与 post_load_weights）：从 checkpoint 读进数据，可能做重排与交错；
3. 前向（apply_linear）：调 quant_method.apply(module, input, bias)。

这个设计的好处是：量化差异被压缩进一个方法对象里，Linear 本身不用知道自己是 FP8 还是 FP4。代价是方法的构造必须发生在权重分配之前，因为张量的形状由它决定——这也解释了为什么 create_weights 里会有断言：

1
assert in_features % module.scaling_vector_size == 0, \
2
    f"in_features {in_features} must be divisible by scaling_vector_size {module.scaling_vector_size}"

如果输入维度不是 16 的倍数，NVFP4 的块划分就无法整齐覆盖整个 K 维，直接拒绝加载比运行时报错好。

get_quant_method：一张决策表

方法选择的全部逻辑在一个函数里：

1
def get_quant_method(quant_config: Optional[QuantConfig] = None):
2
    if quant_config is None or not quant_config.layer_quant_mode.has_any_quant(exclude_kv_cache=True):
3
        return UnquantizedLinearMethod()
4
    if quant_config.layer_quant_mode.has_fp8_qdq():
5
        return FP8QDQLinearMethod()
6
    if quant_config.layer_quant_mode.has_fp8_rowwise():
7
        return FP8RowwiseLinearMethod()
8
    if quant_config.layer_quant_mode.has_fp8_block_scales():
9
        return FP8BlockScalesLinearMethod()
10
    if quant_config.quant_algo == QuantAlgo.W4A16_NVFP4:
11
        return W4A16NVFP4LinearMethod()
12
    if quant_config.layer_quant_mode.has_nvfp4():
13
        if quant_config.quant_algo == QuantAlgo.NVFP4_ARC:
14
            return NVFP4ARCLinearMethod()
15
        return NVFP4LinearMethod()
16
    if quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8():
17
        return W4A8NVFP4FP8LinearMethod()
18
    if quant_config.layer_quant_mode.has_w4a8_mxfp4_fp8():
19
        return W4A8MXFP4FP8LinearMethod()
20
    if quant_config.layer_quant_mode.has_mxfp8():
21
        return MXFP8LinearMethod()
22
    if quant_config.layer_quant_mode.is_weight_only() and not quant_config.layer_quant_mode.has_per_group_scaling():
23
        return WeightOnlyQuantLinearMethod()
24
    if quant_config.layer_quant_mode.is_int4_weight_only_per_group() and quant_config.quant_algo == QuantAlgo.W4A16_AWQ:
25
        return W4A16_AWQ_LinearMethod()
26
    if quant_config.layer_quant_mode.is_int4_weight_only_per_group() and quant_config.quant_algo == QuantAlgo.W4A8_AWQ:
27
        return W4A8_AWQ_LinearMethod()
28
    if quant_config.layer_quant_mode.has_w4a8_mxfp4_mxfp8():
29
        return W4A8MXFP4MXFP8LinearMethod()
30
    raise ValueError(f'unsupported quant mode: {quant_config.quant_mode}')

三点需要留意：

第一，KV 缓存的量化被排除在方法选择之外（exclude_kv_cache=True）。KV 量化不影响 Linear 的张量布局，它属于注意力后端的事，本站已有 TensorRT-LLM 源码完全拆解（七）：KV Cache 的内存布局与量化存储 专门讲这条线。量化子系统与 KV 子系统在这里正交分离。

第二，判断顺序是「格式优先、特例随后」。 FP8 的三种粒度（QDQ、rowwise、block scales）先判，因为它们互斥且由 QuantMode 位直接区分；NVFP4 之后则是按 quant_algo 再细分出 W4A16_NVFP4、NVFP4_ARC 等变体。最后一行 raise ValueError 而不是返回默认方法，是因为「没匹配上」意味着 QuantMode 与 QuantAlgo 不一致，属于配置错误，静默降级会掩盖问题。

第三，INT4 权重量化的分支结构说明：方法选择先看「是否 per-group」，per-group 之下再按 quant_algo 区分 AWQ 与非 AWQ。 这是因为 AWQ 的 checkpoint 多带一个 pre_quant_scale，需要专门的加载与激活预处理路径。

硬件相关的二次选择

选出方法之后还有一层硬件相关覆盖，在 Linear.get_quant_method 里：

1
def get_quant_method(self, quant_config: Optional[QuantConfig] = None):
2
    quant_method = get_quant_method(quant_config)
3
    method_type = type(quant_method)
4
    if method_type is W4A16NVFP4LinearMethod:
5
        # Marlin by default on SM120/121, explicit opt-in elsewhere.
6
        if MarlinNVFP4LinearMethod.is_enabled(self):
7
            return MarlinNVFP4LinearMethod()
8
    elif method_type is NVFP4LinearMethod:
9
        if _uses_marlin_nvfp4_backend(self):
10
            return MarlinNVFP4LinearMethod()
11
    return quant_method

这段代码回答了一个自然的问题：同一份 checkpoint 在不同 GPU 上要用不同的 kernel 吗？ 答案是会。Marlin 是一套 W4A16 的权重专用 kernel（原理见 Marlin 完全拆解），它在 SM120/121（消费级 Blackwell）上是默认选择，在其他架构上需要显式开启。注释里特别强调方法选择与 uses_marlin_nvfp4 属性必须用同一个判断函数，否则「选出来的方法」与「对外声明的能力」可能不一致，上游就会按 A4 的假设去调一个 A16 的 kernel。

还有一个细节：create_weights_impl 开头会根据是否存在 NVFP4 激活量化来决定是否关闭融合路径：

1
if not quant_config_has_nvfp4_activation_quantization(self.quant_config):
2
    self.use_fused_gemm_allreduce = False

GEMM 加 AllReduce 的融合要求输出 buffer 是一个特定的对称内存窗口，而只有激活也被量化的路径才支持这种写法。这类「量化方式决定融合能不能做」的耦合，是量化子系统与并行通信子系统之间的接口，本站 TensorRT-LLM 源码完全拆解（十一）：并行拓扑与通信 讲了另一半。

WeightMode：同一份权重的三种摆放

load_weights 会按 WeightMode 分派到三个不同的加载函数：

1
if weight_mode == WeightMode.VANILLA:
2
    load_fn = self.load_weights_vanilla
3
elif weight_mode == WeightMode.FUSED_QKV_LINEAR:
4
    load_fn = self.load_weights_fused_qkv_linear
5
elif weight_mode == WeightMode.FUSED_GATE_UP_LINEAR:
6
    load_fn = self.load_weights_fused_gate_up_linear

• VANILLA：checkpoint 里 q/k/v 是三个独立张量，运行时也保留三块权重；
• FUSED_QKV_LINEAR：q/k/v 融合成一块，由一个 GEMM 一次算出，减少 kernel 启动次数与重复读取激活；
• FUSED_GATE_UP_LINEAR：gate/up 融合，对应前面 reorder_rows_for_gated_act_gemm 的行重排。

融合带来的收益是访存与启动开销，代价是量化的 shard 对齐变得复杂。NVFP4 的 load_weight_scales 就写明了这一点：

1
scale_span = (module.scaling_vector_size
2
              if module.tp_mode == TensorParallelMode.ROW else 1)

在行并行（ROW）下，权重沿 K 维切分，而 NVFP4 的 scale 是每 16 个 K 元素一个，所以切分时必须按 16 个元素为最小单位，否则每个 rank 都会拿到半截 block。scale_span 这个变量就是在告诉 load_shard：「你在 K 维每移动一个单位，scale 张量走多远」。TP 切分与量化粒度的这个交点，在任何框架里都是 bug 高发区。

前向：一次 FP8 QDQ GEMM 的完整调用

以 FP8QDQLinearMethod.apply 为例，看一次量化 GEMM 在运行时做了什么：

1
static_input_scale = self.get_static_input_scale(module)
2
cur_input_scale = module.input_scale
3
if input.dtype != torch.float8_e4m3fn:
4
    if static_input_scale is not None:
5
        # Static quantization
6
        qinput, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(
7
            input, static_input_scale)
8
    else:
9
        # Dynamic quantization
10
        qinput, cur_input_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(input)
11
        cur_input_scale = cur_input_scale.to(torch.float32)
12
else:
13
    qinput = input

三件事依次发生：

1. 判断输入是否已经被量化。如果上一个算子已经输出了 FP8（前一层的量化输出被直接复用，或者注意力输出已经是 FP8），就跳过量化。这是消除冗余量化的关键优化——if input.dtype != torch.float8_e4m3fn 这一句省掉的是整个张量的一次读写。
2. 静态与动态量化的分流。静态用 checkpoint 里存的 input_scale，动态则调 quantize_e4m3_per_tensor 现算 amax 得到 scale。
3. 选 GEMM 后端：

1
if _fp8_per_tensor_uses_cute_dsl_sm107():
2
    output = torch.ops.trtllm.cute_dsl_fp8_per_tensor_gemm_rubin(...)
3
elif module.enable_cuda_core and qinput.shape[0] <= 8:
4
    # use cuda core for small m dimension
5
    output = torch.ops.trtllm.cuda_scaled_mm(...)
6
else:
7
    output = torch.ops.trtllm.cublas_scaled_mm(...)

第二支是按 M 维大小分流：解码阶段 M 很小（batch 乘 1 个 token），Tensor Core 的 tile 利用率低，改用 CUDA Core 的 cuda_scaled_mm 反而更快。这条经验在很多框架的量化 kernel 里都能看到。

FP8 block scales：布局转换发生在加载期

FP8BlockScalesLinearMethod 展示了 block 量化的完整形态：

1
scale_shape = (math.ceil(out_features / 128), math.ceil(in_features / 128))
2
module.weight_scale = Parameter(torch.empty(scale_shape, dtype=torch.float32), ...)

scale 是二维的 128x128 块网格，每个块一个 fp32 值。前向时激活的量化与 GEMM 后端也按 SM 版本分流：fp8_swap_ab_gemm（DeepGEMM 路线）、fp8_block_scaling_gemm（trtllm-gen 路线）、cute_dsl_fp8_gemm_blackwell（SM100/103 的 CuTe DSL 路线），以及 SM107 上用 UE8M0 编码 scale 的 cute_dsl_mxfp8_gemm_rubin。

其中一个容易被忽略的类型细节是 scale 的类型：激活 scale 的 dtype 决定了走哪条路：

1
uses_cute_dsl_rubin = (activation_scale.dtype == torch.uint8
2
                       and sm_version == 107
3
                       and IS_CUTLASS_DSL_RUBIN_AVAILABLE
4
                       and (...))
5
elif (activation_scale.dtype == torch.int32 and is_sm_100f()
6
      and not module.disable_deep_gemm):
7
    output = torch.ops.trtllm.fp8_prequantized_swap_ab_gemm(...)

uint8 意味着 scale 被编码成 UE8M0（无符号 8 位纯指数），int32 意味着 scale 被打包成 int32 以便被 DeepGEMM 风格的内核读取。同一份数学量、不同的内存编码，对应的后端完全不同。这解释了为什么 transform_weights 里有一句注释说「把 weight_scale 重新铺成 UE8M0 K32 R128c4」——布局转换被挪到加载阶段完成，前向时只需按 dtype 分派。

MXFP8 的参考实现：十几行讲完一种格式

mxfp8_utils.py 提供了 MXFP8 的 PyTorch 参考路径，注释特别说明它不是生产路径，而是作为 CUTLASS kernel 的数值对照：

1
def quant_bf16_to_mxfp8(w: torch.Tensor, block_size: int = 32):
2
    # [O, K] bf16/fp32 -> (e4m3 [O,K], ue8m0 uint8 [O, K//block])
3
    wb = w.float().view(o, k // block_size, block_size)
4
    amax = wb.abs().amax(dim=-1).clamp_min(1e-12)
5
    E4M3_MAX = 448.0
6
    exp = torch.ceil(torch.log2(amax / E4M3_MAX))
7
    scale_ue8m0 = (exp + UE8M0_BIAS).clamp(0, 255).to(torch.uint8)
8
    scale = torch.exp2(exp).unsqueeze(-1)
9
    w_e4m3 = (wb / scale).to(torch.float8_e4m3fn).view(o, k)
10
    return w_e4m3, scale_ue8m0.view(o, k // block_size)

这段十几行的代码把 MXFP8 的全部数学讲完了：每 32 个元素取 amax，用 $2^e \ge \text{amax}/448$ 求出最近的 2 的幂指数 $e$，把 $e$ 加上偏置 127 存成 uint8（这就是 UE8M0：只存指数、尾数固定为 1），然后元素除以 $2^e$ 再转成 e4m3。MX 格式的核心是「scale 必须是 2 的幂」，好处是缩放可以用指数加减实现、无需乘法器，代价是 scale 只能近似到最近的 2 的幂，精度略低于自由的 FP32 scale。对照 NVFP4 用 e4m3（有 3 位尾数）做 block scale，可以看到两条格式路线在「scale 表达能力与硬件开销」上的不同取舍。