Prhub

#36061 [Diffusion] Dispatch mixed Comfy NVFP4 and INT8 layers

原始 PR 作者 mickqian 合并时间 2026-08-25 09:19 文件变更 4 提交数 18 评论 1 代码增减 +110 / -2

执行摘要

实现 Comfy 混合 NVFP4/INT8 检查点按层分发量化

PR body 明确指出:社区导出越来越多地把 NVFP4 MLPs 与 INT8 ConvRot attention 混在同一个 safetensors 文件中;父 PR(#36044)已经能识别 NVFP4 布局,但遇到混合标记集时会直接拒绝,而不是尊重每个序列化层自己的格式。本 PR 的目标是让加载器按层标记分发量化实现,复用现有 ModelOpt NVFP4 与 Kitchen INT8 代码,用户无需显式选择量化方式。

值得精读。set_comfy_layer_markers 展示了如何在一份量化配置内按层组合两种量化方法,get_quant_method 的 prefix 分发模式也可复用于其他混合量化检查点。注意本 PR 栈式叠加在 #36044 之上,建议结合父 PR 与 #36066 一起阅读,理解 Diffusion 量化加载的完整演进方向。

讨论亮点

该 PR 没有 reviewer 评论,唯一评论来自 mintlify bot 的文档预览部署通知,指向 MiniMax-H3 cookbook 的 Mintlify Preview,无实质技术讨论。技术决策主要体现在 commit 历史中:例如“Fix mixed NVFP4 dispatch test capability”与“Mock optional INT8 kernel in mixed quant test”说明作者处理了设备能力与可选 kernel 的 mock 问题;多次“merge encoder quant marker fix”提交显示编码器量化标记的合并逻辑经过了多轮修复。

实现拆解

  1. 扩展 NVFP4 配置以承载 INT8 子配置(modelopt_quant.py):ModelOptFp4Config 新增 _comfy_int8_config 属性与 set_comfy_layer_markers 方法,校验 layer_markers 中只允许 nvfp4 与 int8_tensorwise 两种格式,并收集 int8 标记构造 KitchenInt8Config。get_quant_method 对命中 INT8 标记的 prefix 委托给 KitchenInt8Config,其余层走原 NVFP4 路径,从而在一份配置内实现跨量化方法的分发。
  2. 放宽检查点量化解析分支(minimax_h3_weights.py):resolve_minimax_h3_checkpoint_quantization 将进入 NVFP4 路径的条件从 formats == 与“nvfp4” 改为 “nvfp4” in formats,使混合标记集也能进入;对白名单外的伴随格式提前抛出 NotImplementedError;用 isinstance(config, ModelOptFp4Config) 替代原 None 检查,并在返回前调用 set_comfy_layer_markers 注入逐层标记。
  3. 新增逐层分发测试(test_transformer_quant.py):test_minimax_h3_mixed_nvfp4_int8_dispatches_each_layer 构造同时含 nvfp4(qkv_proj)与 int8_tensorwise(out_proj,convrot_groupsize=256)标记的 safetensors,在 mock 设备能力后断言 get_quant_method 分别返回 ModelOptFp4LinearMethod 与 KitchenInt8LinearMethod。
  4. 文档补充(MiniMax-H3.mdx):在 cookbook 的量化小节说明混合文件中标记为 int8_tensorwise 的层会自动分发到 Kitchen INT8 ConvRot 路径。

测试配套方面,单元测试覆盖了标记解析、格式白名单校验与逐层分发;测试中 mock 了 current_platform.get_device_capability 与 kitchen_int8._load_comfy_kitchen,未覆盖真实 kernel 加载。本次无配置项、部署或 schema 变更。

文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py 量化配置 modified 7.24
python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py 权重加载 modified 6.27
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py 量化测试 modified 6.16
docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx 文档 modified 2.22

关键符号

set_comfy_layer_markers ModelOptFp4Config.get_quant_method resolve_minimax_h3_checkpoint_quantization test_minimax_h3_mixed_nvfp4_int8_dispatches_each_layer

关键源码片段

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py data-contract

核心变更文件。在 ModelOptFp4Config 中新增 set_comfy_layer_markers 与 _comfy_int8_config,并在 get_quant_method 中实现按 prefix 向 Kitchen INT8 量化方法分发,是混合检查点支持的数据契约核心。

# modelopt_quant.py 中的关键片段(节选)
class ModelOptFp4Config(ModelOptQuantConfig):
    '''NVFP4 量化配置类。    在 Comfy 混合检查点场景下,一份 safetensors 里可能同时包含
    NVFP4 MLP 与 INT8 ConvRot attention 两种格式,因此该类额外
    持有一个 Kitchen INT8 子配置,并按 layer prefix 将层分发给
    对应的量化实现。
    '''
​
    def __init__(self, ...):
        ...
        # 新增:保存从 layer_markers 提取的 INT8 子配置;
        # 为 None 时表示纯 NVFP4 检查点,行为与旧版本一致
        self._comfy_int8_config: KitchenInt8Config | None = None
​
    def set_comfy_layer_markers(self, layer_markers: dict[str, dict[str, Any]]) -> None:
        '''接收 checkpoint 的按层格式标记,并校验支持的格式集合。'''
        # 白名单只有 nvfp4 与 int8_tensorwise,出现其他格式直接报错,
        # 避免带着未处理的格式进入后续权重加载
        unsupported = {
            str(marker.get('format')) for marker in layer_markers.values()
        } - {'nvfp4', 'int8_tensorwise'}
        if unsupported:
            raise ValueError(
                'NVFP4 checkpoints cannot dispatch companion Comfy formats: '
                + ', '.join(sorted(unsupported))
            )
        # 收集所有 int8_tensorwise 层并构造 Kitchen INT8 子配置,
        # 供 get_quant_method 按 prefix 精确分发
        int8_markers = {
            prefix: marker
            for prefix, marker in layer_markers.items()
            if marker.get('format') == 'int8_tensorwise'
        }
        self._comfy_int8_config = (
            KitchenInt8Config(layer_markers=int8_markers) if int8_markers else None
        )
​
    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
        # INT8 标记层优先委托给 Kitchen INT8 实现,其余层仍走 NVFP4 路径
        if (
            self._comfy_int8_config is not None
            and prefix in self._comfy_int8_config.layer_markers
        ):
            return self._comfy_int8_config.get_quant_method(layer, prefix)
        return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod)
python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py dependency-wiring

修改 MiniMax-H3 检查点量化解析分支,从只接受纯 NVFP4 标记集扩展为接受含 nvfp4 的混合标记集,并在构建配置后注入逐层标记,是整个功能的接线点。

# minimax_h3_weights.py:解析 MiniMax-H3 检查点的量化配置(节选)
def resolve_minimax_h3_checkpoint_quantization(
    layer_markers: dict[str, dict[str, Any]],
    safetensors_list: list[str] | None = None,
    param_names_mapping: dict | None = None,
    reverse_param_names_mapping: dict | None = None,
) -> QuantizationConfig | None:
    # 只要包含 nvfp4 标记就进入 NVFP4 构建路径,
    # 这样混合标记集(nvfp4 + int8_tensorwise)也能被解析
    formats = {str(marker.get('format')) for marker in layer_markers.values()}
    if 'nvfp4' in formats:
        # 白名单外的伴随格式直接拒绝,宁可报错也不要静默误解析
        unsupported = formats - {'nvfp4', 'int8_tensorwise'}
        if unsupported:
            raise NotImplementedError(
                'Unsupported Comfy NVFP4 companion format(s): '
                + ', '.join(sorted(unsupported))
            )
        if safetensors_list is None:
            raise ValueError('MiniMax-H3 NVFP4 metadata requires checkpoint files')
        config = build_nvfp4_config_from_safetensors_list(
            safetensors_list,
            param_names_mapping,
            reverse_param_names_mapping,
        )
        # 用 isinstance 校验返回类型,防止将来返回其他配置类型时静默出错
        if not isinstance(config, ModelOptFp4Config):
            raise ValueError('Could not resolve MiniMax-H3 NVFP4 checkpoint layout')
        # 把逐层标记注入配置,使 NVFP4 config 内部能够按层分发 INT8 实现
        config.set_comfy_layer_markers(layer_markers)
        config.checkpoint_uses_comfy_quantization = True
        config.checkpoint_uses_native_qkv_layout = True
        config.checkpoint_weight_scale_layout = 'swizzled'
        config.swap_weight_nibbles = True
        return config
    # 非 NVFP4 场景继续走通用 Comfy 量化解析
    return resolve_comfy_checkpoint_quantization(layer_markers)

评论区精华

Mintlify 文档预览部署 other

mintlify bot 在 PR 上发布了文档预览链接,指向 MiniMax-H3 cookbook 的预览站点,无实质技术讨论。

结论:无需要处理的技术结论;PR 无 reviewer 评论,最终由作者直接合并。 · 无实质技术讨论

风险与影响

核心风险集中在分发与平台依赖两方面:get_quant_method 依赖 prefix 在 _comfy_int8_config.layer_markers 中的精确匹配,若 checkpoint 标记前缀与模型内部 prefix 不一致,INT8 层会被静默降级到 NVFP4 路径,权重解析可能错位;Kitchen INT8 ConvRot kernel 并非所有平台可用,测试中 mock 了 _load_comfy_kitchen,真实加载失败场景未被 CI 覆盖;NVFP4 路径本身要求 SM100+。另外 resolve 逻辑收紧后,未来出现白名单外的伴随格式会直接 NotImplementedError,需要同步扩展白名单。影响面局限在 Diffusion 运行时 MiniMax-H3 加载路径,不涉及 SRT 核心调度与推理路径。

用户可直接加载社区导出的混合量化 MiniMax-H3 safetensors 而无需指定 --quantization;纯 NVFP4 检查点无 INT8 标记时 _comfy_int8_config 为 None,行为与旧版本一致。系统侧新增一层按层分发逻辑,但只影响含 int8_tensorwise 标记的检查点。团队侧为后续支持更多 Comfy 伴随格式(如 FP8,见 #36066)提供了可复用的分发框架。

量化配置契约变更 平台与 kernel 依赖未全覆盖 prefix 精确匹配分发 混合格式白名单较窄

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论