资讯详情

SGLang 量化架构深度解析:以 W4AFp8 为例理解 create_weights → process_weights_after_loading → apply 三阶段设计

📅 2026/10/9 1:37:55 | 华诺云谱 👁 阅读
SGLang 量化架构深度解析:以 W4AFp8 为例理解 create_weights → process_weights_after_loading → apply 三阶段设计
文档教程人工智能大模型RLHF【免费下载链接】Awesome-ML-SYS-TutorialMy learning notes for ML SYS.项目地址https://gitcode.com/gh_mirrors/aw/Awesome-ML-SYS-Tutorial点击查看免费下载导读本文以 W4AFp8INT4 权重 FP8 激活量化方案为切入点系统拆解 SGLang 推理引擎的量化设计与实现思路。SGLang 将全部量化实现收纳在python/sglang/srt/layers/quantization/目录中用一套恰到好处的抽象基类与钩子函数Hook Function把模型构建、权重加载、推理执行三个生命周期阶段有机连接起来。读完本文你将掌握 SGLang 量化的核心抽象模型配置类与量化方法类的继承层次、三条关键调用链路的完整走向、W4AFp8 在 MoE 层与线性层的落地细节以及如何不改动框架核心代码、仅通过三步扩展接入一种全新的量化方案如 W2A8。一、量化三阶段抽象一条贯穿模型全生命周期的水管SGLang 把量化拆解为三个核心阶段对应量化方法类上的三个钩子函数阶段钩子函数职责类比模型构建create_weights预分配量化权重、Scale 因子等参数张量的内存铺设水管尚未通水权重加载完成后process_weights_after_loading把权重和 Scale 转换成计算内核如 CUTLASS最优的格式与布局重排、降精度等数据转换、蓄水推理执行apply指挥底层计算内核如 FP8 GEMM让激活与权重真正流动完成计算开水放行这套抽象的精妙之处在于模型代码完全不需要感知量化细节。线性层与 MoE 层在初始化时只负责调用self.quant_method上的钩子而quant_method由配置类按层类型动态分发。对于不量化的模型quant_method会被设置为UnquantizedLinearMethod其create_weights直接创建指定形状和数据类型的普通权重参数apply退化为常规 GEMM——这保证了量化路径与非量化路径的代码完全统一。从仓库笔记 SGLang 模型加载过程 可以印证这一抽象在实际代码中的位置Linear层如ColumnParallelLinear等在__init__中调用self.quant_method.create_weights创建权重参数在forward函数中调用quant_method.apply完成实际计算权重加载完成后模型加载器会遍历每个子模块执行量化方法中定义的process_weights_after_loading且该阶段要求参数位于全局目标设备上在 CPU 卸载场景下SGLang 会将参数先移到设备、执行后处理、再移回原位置。二、整体流程从配置解析到量化 GEMM 的完整链路SGLang 量化启动的完整调用链如下ModelConfig._parse_quant_hf_config → 判定 quant_method如 w4afp8 ↓ weight_utils.get_quant_config → 构造对应 QuantizationConfig 实例 ↓ _initialize_model(...) → 把 quant_config 传入模型/各层 ↓ LinearBase.quant_method.create_weights → 注册量化权重占位 ↓ DefaultModelLoader.load_weights_and_postprocess → 先 load_weights 加载权重 再逐层调用 quant_method.process_weights_after_loading ↓ 推理时 LinearBase.forward → quant_method.apply 执行量化 GEMM链路两端各有一个关键的桥配置侧ModelConfig._parse_quant_hf_config从模型目录的量化配置如hf_quant_config.json中判定出字符串标识quant_method随后weight_utils.get_quant_config依据注册表把字符串解析为具体的QuantizationConfig实例并在_initialize_model阶段随模型初始化传入各层。执行侧权重加载由DefaultModelLoader统一负责——先通过load_weights把 checkpoint 数据写入create_weights注册的参数占位符再逐层触发process_weights_after_loading做格式转换推理时则由各层forward中的quant_method.apply完成最终计算。三、类继承层次配置类与量化方法类的双向解耦SGLang 的量化抽象由两条继承链组成职责划分清晰配置类继承关系定义于 base_config.py QuantizationConfig抽象基类 职责解析量化配置、硬件校验与激活数据类型校验并按层返回正确的量化方法实例 ├─ from_config()从配置字典解析并实例化 ├─ get_quant_method()根据层类型返回量化方法 ├─ get_min_capability()校验硬件兼容性 │ 硬件校验包括 │ • NVIDIA GPU通过 CUDA capability计算能力检查 │ 如 70 (Volta)、75 (Turing)、80 (Ampere)、90 (Hopper) 等 │ • AMD GPU通过 ROCm/HIP 平台检测某些方案会检查特定 GCN 架构如 gfx94 └─ get_supported_act_dtypes()返回支持的激活数据类型 ↓ W4AFp8Config具体配置类 └─→ get_quant_method() 根据层类型返回 ├─→ LinearBase → Fp8LinearMethod └─→ FusedMoE → W4AFp8MoEMethod 量化方法类继承关系定义于 base_config.py QuantizeMethodBase抽象基类 职责权重注册、权重加载完成后的处理以及前向执行 ├─ create_weights()注册量化权重占位符模型构建阶段 ├─ process_weights_after_loading()权重后处理权重加载完成后 └─ apply()前向传播执行量化计算推理阶段 ├─→ LinearMethodBase │ └─→ Fp8LinearMethod用于普通线性层 └─→ FusedMoEMethodBase └─→ W4AFp8MoEMethod用于 MoE 层 量化方法注册表位于 __init__.py 将原生方案AWQ、GPTQ、FP8、W4AFp8、ModelOpt 等映射成字符串标识从源码结构可以清晰看出两条设计原则配置与实现分离QuantizationConfig只负责配置解析 硬件校验 方法分发不触碰任何张量逻辑QuantizeMethodBase只负责建参、转格式、算前向不关心配置从哪来。同方案、异实现一个配置类可以通过get_quant_method按层类型分发到不同的方法实例如 W4AFp8 对普通线性层用Fp8LinearMethod、对 MoE 层用W4AFp8MoEMethod复用同一份配置语义。四、三条调用链路的完整走向4.1 create_weights模型构建阶段以 DeepSeekV2 的 MLA 注意力 MoE 结构为例DeepseekV2DecoderLayer.__init__() ↓ DeepseekV2AttentionMLA.__init__() ↓ RowParallelLinear.__init__() ↓ LinearBase.__init__() (super().__init__()) # self.quant_method 在父类初始化函数里赋值 ↓ quant_config.get_quant_method() → 返回 Fp8LinearMethod ↓ RowParallelLinear.__init__() 继续执行 ↓ Fp8LinearMethod.create_weights() ↓ 注册 weight、weight_scale、input_scale 等参数占位符关键点在于self.quant_method是在LinearBase的父类初始化函数中被赋值的各具体线性层无需关心量化方案的差异。4.2 process_weights_after_loading权重加载完成后Scheduler.__init__() ↓ TpModelWorker.__init__() ↓ ModelRunner.__init__() ↓ ModelRunner.initialize() ↓ get_model() ↓ DefaultModelLoader.load_model() ↓ DefaultModelLoader.load_weights_and_postprocess() ↓ model.load_weights() → 加载权重数据 ↓ 逐层遍历调用 quant_method.process_weights_after_loading() ↓ Fp8LinearMethod.process_weights_after_loading() └─→ 或 W4AFp8MoEMethod.process_weights_after_loading()这一阶段在仓库笔记中同样有据可查SGLang 模型加载过程 明确记载了权重加载完成后遍历子模块、执行process_weights_after_loading的行为并且该钩子在 在线权重更新 场景下也被复用——在线更新流程在model.load_weights(iter)之后会遍历self.model.named_modules()对每个带quant_method的模块在device_loading_context作用域内重新调用process_weights_after_loading(module)。这说明三阶段抽象不仅是模型冷启动的流程也是热更新权重替换后重新做格式转换的公共基础设施。4.3 apply推理阶段DeepseekV2DecoderLayer.forward() ↓ DeepseekV2AttentionMLA.forward() ↓ RowParallelLinear.forward() ↓ self.quant_method.apply() ↓ Fp8LinearMethod.apply() └─→ 或 W4AFp8MoEMethod.apply() ↓ 调用底层内核CUTLASS/Marlin/torch执行量化 GEMM五、W4AFp8 量化方案深入解析W4AFp8 即权重 INT4 激活 FP8的混合精度方案SGLang 对其的实现包含配置类W4AFp8Config与两个量化方法类W4AFp8MoEMethod、Fp8LinearMethod。5.1 W4AFp8Config配置 → 具体量化方法的映射W4AFp8Config继承自QuantizationConfig负责描述清楚配置 → 具体量化方法这一映射关系。配置识别当hf_quant_config.json中quant_algo MIXED_PRECISION时ModelConfig会把量化方案映射为w4afp8并校验硬件兼容性# ModelConfig._parse_modelopt_quant_config if quant_algo MIXED_PRECISION: return {quant_method: w4afp8}对象构造weight_utils.get_quant_config获取W4AFp8Config类然后调用from_config方法进行实例化。关键方法W4AFp8Config.from_config()从配置字典解析并实例化配置对象。W4AFp8Config.get_quant_method(layer, prefix)核心方法根据层类型返回对应的量化方法实例if isinstance(layer, LinearBase): return Fp8LinearMethod(self) # 普通层用 Fp8LinearMethod elif isinstance(layer, FusedMoE): return W4AFp8MoEMethod(self) # MoE 层用 W4AFp8MoEMethod5.2 W4AFp8MoEMethodMoE 层的三步落地W4AFp8MoEMethod是 W4AFp8 在 MoE 层上的具体实现严格遵循三阶段抽象。第一步 create_weights参数预分配在FusedMoE模块初始化时为 MoE 层创建量化所需的参数容器。主要工作包括创建量化权重张量w13_weightgate 和 up projection和w2_weightdown projection注意类型是int8分配权重缩放因子w13_weight_scale_inv和w2_weight_scale_inv每组 128 个元素共享一个 scalegroup-wise 量化准备激活缩放因子w13_input_scale和w2_input_scale初始化计算所需的元数据如 stride、expert offsets 等。注意此时参数为空使用torch.empty创建仅完成内存布局的初始化尚未填充实际数据def create_weights(self, layer, num_experts, hidden_size, ...): # 创建量化权重容器INT8 类型 layer.register_parameter(w13_weight, torch.empty(..., dtypetorch.int8)) layer.register_parameter(w2_weight, torch.empty(..., dtypetorch.int8)) # 创建权重缩放因子group-wise每组 128 元素 layer.register_parameter(w13_weight_scale_inv, torch.zeros(...)) layer.register_parameter(w2_weight_scale_inv, torch.zeros(...)) # 创建输入缩放因子静态量化时使用 layer.register_parameter(w13_input_scale, torch.ones(..., dtypetorch.bfloat16)) # 初始化 stride 等计算元数据 self.a_strides1 torch.full((num_experts, 3), hidden_size, ...)第二步 process_weights_after_loading格式转换与布局优化权重数据从 Checkpoint 加载后需要进行格式转换以适配底层计算内核权重 scale 的格式优化将 float32 格式的 scale 转换为 bfloat16减少 50% 内存占用并调用interleave_scales函数对 scale 进行交错重排。重排的目的是匹配 CUTLASS 内核的内存访问模式参考了 TRT-LLM 的实现重排后内核在计算时能够更高效地访问数据提升缓存命中率输入 scale 的聚合在静态量化模式下把每个专家的输入 scale 聚合为单一标量减少推理计算量。def process_weights_after_loading(self, layer: Module) - None: # 将权重 scale 转换为 bfloat16 并重新排列以匹配 CUTLASS 布局 w13_weight_scale layer.w13_weight_scale_inv.to(torch.bfloat16) w13_weight_scale interleave_scales(w13_weight_scale) layer.w13_weight_scale_inv Parameter(w13_weight_scale, requires_gradFalse) # 将输入 scale 聚合为单一标量静态量化模式 w13_input_scale_max layer.w13_input_scale.max().to(torch.bfloat16).item() layer.w13_input_scale Parameter(torch.tensor([w13_input_scale_max], dtypetorch.bfloat16), requires_gradFalse)第三步 apply混合精度 GEMM 执行在前向传播阶段apply收集所有预处理完成的数据激活、重排后的权重和 scale、路由结果等然后调用cutlass_w4a8_moe底层内核执行两个 GEMM 操作GEMM1w13_weightgate 和 upGEMM2w2_weightdown。cutlass_w4a8_moe是封装了 CUTLASS 库的底层函数实现了 INT4 权重与 FP8 激活的混合精度矩阵乘法充分利用硬件的量化计算能力def apply(self, layer, dispatch_output) - CombineInput: from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe x dispatch_output.hidden_states topk_weights, topk_ids, _ dispatch_output.topk_output # 调用 CUTLASS 内核执行混合精度 MoE 计算 output cutlass_w4a8_moe( x, layer.w13_weight, layer.w2_weight, layer.w13_weight_scale_inv, layer.w2_weight_scale_inv, topk_weights, topk_ids, self.a_strides1, self.b_strides1, self.c_strides1, # GEMM1 的 stride self.a_strides2, self.b_strides2, self.c_strides2, # GEMM2 的 stride self.s_strides13, self.s_strides2, # Scale 的 stride self.expert_offsets, self.problem_sizes1, self.problem_sizes2, layer.w13_input_scale, layer.w2_input_scale, ) # 应用路由缩放因子 if self.moe_runner_config.routed_scaling_factor is not None: output * self.moe_runner_config.routed_scaling_factor return StandardCombineInput(hidden_statesoutput)5.3 Fp8LinearMethod线性层的简化实现对于普通线性层W4AFp8Config分配Fp8LinearMethod逻辑与 MoE 版本类似但更简单create_weights注册weight、weight_scale和input_scale占位符process_weights_after_loading根据硬件Marlin、CUTLASS 等要求对权重和 scale 进行格式转换apply调用合适的内核Marlin、CUTLASS 等执行 FP8 GEMM 计算。这里也体现了 SGLang 量化后端的多内核策略同一个 FP8 线性层可根据硬件能力和部署偏好选择 CUTLASS 或 Marlin 内核内核选择与量化语义三阶段钩子解耦。六、如何扩展一种全新的量化方案SGLang 的可扩展性在于恰到好处、不多不少的抽象。要接入新的量化方案例如 W2A8无需修改框架核心代码只需按照以下三步实现实现配置类继承QuantizationConfig解析自定义参数并实现get_quant_method方法以及必要的get_min_capability、get_supported_act_dtypes校验。实现量化方法类继承LinearMethodBase、FusedMoEMethodBase实现create_weights、process_weights_after_loading和apply三个方法。注册方案在__init__.py的BASE_QUANTIZATION_METHODS中注册建立字符串标识与配置类的映射关系。接入之后模型构建、权重加载、推理执行的既有链路会自动复用新方案——get_quant_method的分发、DefaultModelLoader的逐层后处理、各层forward的apply调用都不需要任何改动。这也是 在线权重更新 能直接复用process_weights_after_loading的原因钩子抽象天然支持换权重后重新转换格式的场景。七、附录SGLang 已支持的量化方法一览以下表格汇总了 SGLang 量化体系中已覆盖的方案类别便于对照选择与扩展CategoryRepresentative ConfigurationsDescriptionFP8 Seriesfp8,w8a8_fp8,modelopt_fp8,fbgemm_fp8Native FP8、W8A8-FP8 混合、ModelOpt/FBGEMM 扩展INT8 Seriesw8a8_int8,blockwise_int8经典 8bit 权重/激活、blockwise INT8INT4/Mixed Precisionw4afp8,qoq,moe_wna164bit 权重 FP8 激活、QoQ、WNA16W4A16/W8A16FP4 / MXFP4modelopt_fp4,petit_nvfp4,mxfp4,quarkFP4 / MXFP4 方案quark为 ROCm 专属Pre-quantized Formatsawq,awq_marlin,gptq,gptq_marlin,gguf,compressed-tensors,auto-round,modelopt与外部工具链或压缩张量框架的集成modelopt可自动识别 FP8/FP4KV Cache QuantizationBaseKVCacheMethod及其在kv_cache.py中的子类为注意力缓存提供 scale 与 zero-point 管理从该表格可以看出SGLang 量化体系的覆盖面分为原生计算方案FP8/INT8/INT4/FP4直接驱动 GEMM 内核与预量化格式兼容层AWQ/GPTQ/GGUF 等负责解析外部工具链产出的权重两类而它们最终都收敛到本文所述的同一套三阶段抽象之上。仓库中还提供了配套的量化背景资料见 general_background.md与 BNB/Unsloth 量化模型接入案例见 unsloth.md可以作为理解 SGLang 量化生态的延伸阅读。结语SGLang 的量化设计可以概括为一句话用一对抽象基类配置类 方法类和三个钩子create_weights / process_weights_after_loading / apply解耦了量化语义与计算内核。模型代码只认quant_method接口配置解析只负责分发格式转换与内核调用完全收敛在量化方法类内部。这种设计既让 W4AFp8 这类混合精度方案可以优雅落地也让社区扩展新量化方案的成本降到了新增两个文件、注册一行映射的量级是值得借鉴的推理框架插件化架构范本。赞分享文档教程人工智能大模型RLHF【免费下载链接】Awesome-ML-SYS-TutorialMy learning notes for ML SYS.项目地址https://gitcode.com/gh_mirrors/aw/Awesome-ML-SYS-Tutorial点击查看免费下载相关推荐Soldier76安全使用指南避免检测的终极策略Soldier76安全使用指南避免检测的终极策略 Soldier76是一款功能强大的PUBG罗技鼠标宏工具能够帮助玩家实现自动压枪支持12种枪械无后坐力效游戏开发如何构建C语言解释器深入理解编译器三阶段处理模型如何构建C语言解释器深入理解编译器三阶段处理模型 想要学习编译原理却不知从何入手今天我将为你揭秘C语言解释器的 架构设计 让你轻松掌握 编译器三阶段处理模编译器/解释器文档/教程SGLang 多模态请求生命周期深度解析以 Qwen2.5-VL 为例的端到端架构详解SGLang 多模态请求生命周期深度解析以 Qwen2.5 VL 为例的端到端架构详解 导读本文以 Qwen2.5 VL 为参考模型逐层拆解 SGLang文档教程人工智能大模型RLHF上一篇3分钟掌握SiYuan笔记终极特殊符号输入技巧指南下一篇10倍速浏览数据库DBeaver元数据缓存机制的设计与实现创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑