minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

📝 [NPU] Use vendor-neutral wording in quantization comments (#36768)

Browse Source
This commit is contained in:
Junlin Wu
2026-08-29 23:05:35 +03:00
committed by GitHub
parent 6afb5e1771
commit 09ecb9aaaa
8 changed files with 34 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py
+1 -7
View File
@@ -1,4 +1,4 @@
"""ModelSlim MXFP4 scheme for pre-quantized weight inference on Ascend NPU. """ModelSlim MXFP4 scheme for pre-quantized weight inference on NPU.
Loads weights pre-quantized by msmodelslim and runs MXFP4 dual-level Loads weights pre-quantized by msmodelslim and runs MXFP4 dual-level
matmul at inference via npu_dual_level_quant_matmul. matmul at inference via npu_dual_level_quant_matmul.
@@ -8,9 +8,6 @@ Checkpoint tensor formats (verified from msmodelslim export):
weight_scale: [out, in/32] uint8 (L1 block scales, e8m0+127) weight_scale: [out, in/32] uint8 (L1 block scales, e8m0+127)
weight_dual_scale:[out, in/512, 1] float32 (L0 coarse scales) weight_dual_scale:[out, in/512, 1] float32 (L0 coarse scales)
mul_scale: [in] float32 (smooth quant activation scale) mul_scale: [in] float32 (smooth quant activation scale)
Reference: MindIE-SD W4A4MXFP4DualQuantLinear
(MindIE-SD/mindiesd/quantization/layer.py)
""" """
from typing import List, Optional from typing import List, Optional
@@ -98,7 +95,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
# After repack, it becomes `<prefix>.mul_scale`. # After repack, it becomes `<prefix>.mul_scale`.
# This is CRITICAL: the offline-quantized weights were calibrated with # This is CRITICAL: the offline-quantized weights were calibrated with
# x * mul_scale applied to the activation. Omitting it causes mosaic output. # x * mul_scale applied to the activation. Omitting it causes mosaic output.
# Ref: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul lines 385-386.
mul_scale = BasevLLMParameter( mul_scale = BasevLLMParameter(
data=torch.empty( data=torch.empty(
(input_size_per_partition,), (input_size_per_partition,),
@@ -119,7 +115,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
weight = weight.to(f"npu:{torch.npu.current_device()}") weight = weight.to(f"npu:{torch.npu.current_device()}")
weight = torch_npu.npu_dtype_cast(weight, torch_npu.float4_e2m1fn_x2) weight = torch_npu.npu_dtype_cast(weight, torch_npu.float4_e2m1fn_x2)
# npu_dual_level_quant_matmul requires x2 in FRACTAL_NZ format (format 29). # npu_dual_level_quant_matmul requires x2 in FRACTAL_NZ format (format 29).
# Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
weight = torch_npu.npu_format_cast( weight = torch_npu.npu_format_cast(
weight.view(torch.int8), 29, customize_dtype=torch.int8 weight.view(torch.int8), 29, customize_dtype=torch.int8
) )
@@ -170,7 +165,6 @@ class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
# Apply smooth quant scale before activation quantization. # Apply smooth quant scale before activation quantization.
# The offline-quantized weights were calibrated under x * mul_scale, # The offline-quantized weights were calibrated under x * mul_scale,
# so we MUST apply it here for scale alignment. # so we MUST apply it here for scale alignment.
# Reference: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul
mul_scale = layer.mul_scale mul_scale = layer.mul_scale
if getattr(layer, "use_mul_scale", True): if getattr(layer, "use_mul_scale", True):
x_2d = x_2d * mul_scale.to(x_2d.dtype) x_2d = x_2d * mul_scale.to(x_2d.dtype)
python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py
+1 -1
View File
@@ -1,4 +1,4 @@
"""ModelSlim MXFP8 scheme for pre-quantized weight inference on Ascend NPU. """ModelSlim MXFP8 scheme for pre-quantized weight inference on NPU.
Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights, Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights,
uint8 scales) and runs MXFP8 matmul at inference. uint8 scales) and runs MXFP8 matmul at inference.
python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py
+10 -13
View File
@@ -1,4 +1,4 @@
"""Online MXFP4 quantization for Diffusion models on Ascend NPU. """Online MXFP4 quantization for Diffusion models on NPU.
Provides ``NPUMXFP4Config`` (registered as ``"mxfp4_npu"``) and Provides ``NPUMXFP4Config`` (registered as ``"mxfp4_npu"``) and
``NPUMXFP4DiffusionLinearMethod`` which quantises FP16/BF16 weights to MXFP4 ``NPUMXFP4DiffusionLinearMethod`` which quantises FP16/BF16 weights to MXFP4
@@ -10,9 +10,10 @@ The ``"mxfp4_npu"`` key is distinct from upstream's ROCm ``"mxfp4"``
(``Mxfp4Config`` in ``mxfp4.py``) which targets AMD MI350+ via aiter kernels. (``Mxfp4Config`` in ``mxfp4.py``) which targets AMD MI350+ via aiter kernels.
NOTE: Online weight quantization via ``npu_dynamic_dual_level_mx_quant`` is NOTE: Online weight quantization via ``npu_dynamic_dual_level_mx_quant`` is
experimental. MindIE-SD only uses an offline (pre-quantized) path for MXFP4 experimental; the established MXFP4 path for these models is offline
weights. The online path quantizes FP16/BF16 weights at load time, which may (pre-quantized) only. The online path quantizes FP16/BF16 weights at load
produce different numerical results than the offline calibrated path. time, which may produce different numerical results than the offline
calibrated path.
""" """
from __future__ import annotations from __future__ import annotations
@@ -41,7 +42,7 @@ logger = init_logger(__name__)
class NPUMXFP4Config(QuantizationConfig): class NPUMXFP4Config(QuantizationConfig):
"""Config for online MXFP4 quantization on Ascend NPU (Diffusion).""" """Config for online MXFP4 quantization on NPU (Diffusion)."""
def __init__(self) -> None: def __init__(self) -> None:
super().__init__() super().__init__()
@@ -78,13 +79,11 @@ class NPUMXFP4Config(QuantizationConfig):
class NPUMXFP4DiffusionLinearMethod(LinearMethodBase): class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
"""Ascend NPU MXFP4 linear method for Diffusion models (dual-level). """NPU MXFP4 linear method for Diffusion models (dual-level).
Online mode: loads FP16/BF16 weights → quantises to MXFP4 at load time Online mode: loads FP16/BF16 weights → quantises to MXFP4 at load time
via ``npu_dynamic_dual_level_mx_quant``. via ``npu_dynamic_dual_level_mx_quant``.
Inference: dynamic dual-level MXFP4 activation quant + dual-level matmul. Inference: dynamic dual-level MXFP4 activation quant + dual-level matmul.
Reference: MindIE-SD ``W4A4MXFP4DualQuantLinear`` (offline path only).
""" """
def __init__(self, quant_config: NPUMXFP4Config): def __init__(self, quant_config: NPUMXFP4Config):
@@ -133,9 +132,9 @@ class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}") weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
# Online dual-level MXFP4 weight quantisation. # Online dual-level MXFP4 weight quantisation.
# NOTE: This is experimental — MindIE-SD only has an offline path for # NOTE: This is experimental — the established MXFP4 path for these
# MXFP4 weights. We assume npu_dynamic_dual_level_mx_quant can also # models is offline only. We assume npu_dynamic_dual_level_mx_quant can
# quantise weights (not just activations). # also quantise weights (not just activations).
# Returns: (qw, w_dual_scale, w_scale) # Returns: (qw, w_dual_scale, w_scale)
# qw — quantized weight in float4_e2m1fn_x2 (2 FP4 packed/byte) # qw — quantized weight in float4_e2m1fn_x2 (2 FP4 packed/byte)
# w_dual_scale — L0-level scale (goes to pos 3 in npu_dual_level_quant_matmul) # w_dual_scale — L0-level scale (goes to pos 3 in npu_dual_level_quant_matmul)
@@ -145,14 +144,12 @@ class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
) )
# npu_dual_level_quant_matmul requires x2 (weight) in FRACTAL_NZ format. # npu_dual_level_quant_matmul requires x2 (weight) in FRACTAL_NZ format.
# Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
qw = torch_npu.npu_format_cast( qw = torch_npu.npu_format_cast(
qw.view(torch.int8), 29, customize_dtype=torch.int8 qw.view(torch.int8), 29, customize_dtype=torch.int8
) )
# x2Level0Scale must be [in/level0_block_size, out] — transpose from # x2Level0Scale must be [in/level0_block_size, out] — transpose from
# the [out, in/level0_block_size] shape returned by the quant op. # the [out, in/level0_block_size] shape returned by the quant op.
# Reference: MindIE-SD layer.py:409
w_dual_scale = w_dual_scale.squeeze(-1).transpose(0, 1).contiguous() w_dual_scale = w_dual_scale.squeeze(-1).transpose(0, 1).contiguous()
layer.weight = Parameter(qw, requires_grad=False) layer.weight = Parameter(qw, requires_grad=False)
python/sglang/srt/hardware_backend/npu/moe/init_routing.py
+1 -2
View File
@@ -22,8 +22,7 @@ def _normalize_mxfp_scale(scale: torch.Tensor) -> torch.Tensor:
``npu_moe_init_routing_v2(quant_mode=3)`` emits the scale flat, while the ``npu_moe_init_routing_v2(quant_mode=3)`` emits the scale flat, while the
grouped matmul wants the pair-split view. Already-3D scales (what grouped matmul wants the pair-split view. Already-3D scales (what
``npu_dynamic_mx_quant`` returns) pass through untouched. Mirrors ``npu_dynamic_mx_quant`` returns) pass through untouched.
vllm-ascend's ``maybe_normalize_mxfp_scale_layout``.
""" """
if scale is None or scale.ndim != 2: if scale is None or scale.ndim != 2:
return scale return scale
python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py
+11 -12
View File
@@ -156,7 +156,7 @@ class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
class NPUMXFP8LinearMethod(_NPULinearMethodBase): class NPUMXFP8LinearMethod(_NPULinearMethodBase):
"""Ascend NPU MXFP8 linear method for LLM (SRT) models. """NPU MXFP8 linear method for LLM (SRT) models.
Shared kernel for both the online config path (``--quantization mxfp8``) and Shared kernel for both the online config path (``--quantization mxfp8``) and
the offline ModelSlimMXFP8Scheme (which delegates to this as ``self.kernel``). the offline ModelSlimMXFP8Scheme (which delegates to this as ``self.kernel``).
@@ -343,7 +343,7 @@ class NPU_W4A4DynamicLinearMethod(_NPULinearMethodBase):
class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase): class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
"""Ascend NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations. """NPU W4A8 online quantization: MXFP4 weights + MXFP8 activations.
This is a *true* W4(weight) A8(activation) path: it mirrors the offline This is a *true* W4(weight) A8(activation) path: it mirrors the offline
``W4A8_MXFP`` kernel (``NPUMXFP4W4A8OfflineLinearMethod``) exactly — the only ``W4A8_MXFP`` kernel (``NPUMXFP4W4A8OfflineLinearMethod``) exactly — the only
@@ -363,7 +363,7 @@ class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn) (A8, FP8) BF16/FP16 activation → npu_dynamic_mx_quant(dst=float8_e4m3fn) (A8, FP8)
→ npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block]) → npu_quant_matmul(x2_dtype=float4_e2m1fn_x2, group_sizes=[0, 0, block])
Hardware: Ascend 950 (A5) + a recent torch_npu with the FP4 npu_quant_matmul Hardware: A5 NPU + a recent torch_npu with the FP4 npu_quant_matmul
(same requirement as the offline W4A8 path — see that class's docstring). (same requirement as the offline W4A8 path — see that class's docstring).
""" """
@@ -513,7 +513,7 @@ class NPUMXFP4W4A8LinearMethod(_NPULinearMethodBase):
class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
"""Ascend NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations. """NPU offline W4A8 (ModelSlim ``W4A8_MXFP``): packed-FP4 weights + MXFP8 activations.
Kernel for the offline ModelSlimMXFP4W4A8Scheme (delegated as ``self.kernel``). Kernel for the offline ModelSlimMXFP4W4A8Scheme (delegated as ``self.kernel``).
The msmodelslim ``W4A8_MXFP`` checkpoint stores weights as *packed FP4* The msmodelslim ``W4A8_MXFP`` checkpoint stores weights as *packed FP4*
@@ -626,7 +626,7 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase):
class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):
"""Ascend NPU W4A4 online quantization: single-level MXFP4. """NPU W4A4 online quantization: single-level MXFP4.
True W4(weight) A4(activation): both weights and activations are quantised to True W4(weight) A4(activation): both weights and activations are quantised to
single-level MXFP4 (``float4_e2m1fn_x2``), unlike the W4A8 path which keeps FP8 single-level MXFP4 (``float4_e2m1fn_x2``), unlike the W4A8 path which keeps FP8
@@ -642,8 +642,8 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):
→ npu_quant_matmul(x1_dtype = x2_dtype = float4_e2m1fn_x2, → npu_quant_matmul(x1_dtype = x2_dtype = float4_e2m1fn_x2,
group_sizes=[1, 1, MXFP4_BLOCK_SIZE]) group_sizes=[1, 1, MXFP4_BLOCK_SIZE])
Triggered by ``--quantization mxfp4`` on Ascend NPU. Hardware: Ascend 950 (A5) Triggered by ``--quantization mxfp4`` on NPU. Hardware: A5 NPU with a recent
with a recent torch_npu exposing ``float4_e2m1fn_x2``. torch_npu exposing ``float4_e2m1fn_x2``.
""" """
def create_weights( def create_weights(
@@ -765,7 +765,7 @@ class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase):
class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod):
"""Ascend NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): packed FP4 weights. """NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): packed FP4 weights.
Kernel for the offline ``ModelSlimMXFP4Scheme`` (delegated as ``self.kernel``). Kernel for the offline ``ModelSlimMXFP4Scheme`` (delegated as ``self.kernel``).
The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as packed ``uint8`` The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as packed ``uint8``
@@ -796,7 +796,7 @@ class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod):
class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod): class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod):
"""Ascend NPU W4A4 online quantization: dual-level MXFP4 (higher accuracy). """NPU W4A4 online quantization: dual-level MXFP4 (higher accuracy).
This is the sole online ``--quantization mxfp4`` linear path. Instead of a single This is the sole online ``--quantization mxfp4`` linear path. Instead of a single
UE8M0 (power-of-2) block scale, dual-level MX quant produces a finer L0 (FP8 E4M3) UE8M0 (power-of-2) block scale, dual-level MX quant produces a finer L0 (FP8 E4M3)
@@ -819,9 +819,8 @@ class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod):
BF16/FP16 activation → npu_dynamic_dual_level_mx_quant (A4, dual-level) BF16/FP16 activation → npu_dynamic_dual_level_mx_quant (A4, dual-level)
→ npu_dual_level_quant_matmul(act, weight, act_l0, w_l0, act_l1, w_l1) → npu_dual_level_quant_matmul(act, weight, act_l0, w_l0, act_l1, w_l1)
Reference: Diffusion ``NPUMXFP4DiffusionLinearMethod`` / MindIE-SD Reference: Diffusion ``NPUMXFP4DiffusionLinearMethod``. Hardware: A5 NPU
``W4A4MXFP4DualQuantLinear``. Hardware: Ascend 950 (A5) only — the only — the ``DualLevelQuantBatchMatmul`` op is unavailable on A2/A3.
``DualLevelQuantBatchMatmul`` op is unavailable on A2/A3.
""" """
def process_weights_after_loading(self, layer: torch.nn.Module) -> None: def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py
+5 -5
View File
@@ -223,7 +223,7 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase):
).transpose(1, 2) ).transpose(1, 2)
weight_scale.data = scale weight_scale.data = scale
# The refactored Ascend dispatchers currently support BF16 and INT8. # The refactored NPU dispatchers currently support BF16 and INT8.
# Keep dispatch in BF16 and quantize to MXFP8 immediately before GMM. # Keep dispatch in BF16 and quantize to MXFP8 immediately before GMM.
if weight_prefix == "w13": if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "bf16") self._set_dispatcher_output_dtype(layer, "bf16")
@@ -304,7 +304,7 @@ class NPUW4A4MXFP4MoEMethod(_NPUMoEMethodBase):
).transpose(1, 2) ).transpose(1, 2)
weight_scale.data = scale weight_scale.data = scale
# The refactored Ascend dispatchers currently support BF16 and INT8. # The refactored NPU dispatchers currently support BF16 and INT8.
# Keep dispatch in BF16 and quantize immediately before each GMM. # Keep dispatch in BF16 and quantize immediately before each GMM.
if weight_prefix == "w13": if weight_prefix == "w13":
self._set_dispatcher_output_dtype(layer, "bf16") self._set_dispatcher_output_dtype(layer, "bf16")
@@ -929,7 +929,7 @@ class NPUUnquantMoEMethod(_NPUMoEMethodBase):
# NPUMXFP8MoEMethod # NPUMXFP8MoEMethod
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class NPUMXFP8MoEMethod(_NPUMoEMethodBase): class NPUMXFP8MoEMethod(_NPUMoEMethodBase):
"""MXFP8 MoE on Ascend A5 – float8_e4m3fn weights with e8m0 block scales. """MXFP8 MoE on A5 NPU – float8_e4m3fn weights with e8m0 block scales.
Serves both the online config path (``--quantization mxfp8``, weights Serves both the online config path (``--quantization mxfp8``, weights
quantised at load time) and the offline ModelSlim ``W8A8_MXFP8`` scheme quantised at load time) and the offline ModelSlim ``W8A8_MXFP8`` scheme
@@ -1016,8 +1016,8 @@ class NPUMXFP8MoEMethod(_NPUMoEMethodBase):
# [E, K//64, N, 2] as strided transpose views — DO NOT call # [E, K//64, N, 2] as strided transpose views — DO NOT call
# .contiguous(). Beyond breaking the transpose-flag match above, it # .contiguous(). Beyond breaking the transpose-flag match above, it
# measures slower on the same probe: making both sides contiguous costs # measures slower on the same probe: making both sides contiguous costs
# 6.2% on decode. This matches NPUMXFP8LinearMethod, msmodelslim's # 6.2% on decode. This matches NPUMXFP8LinearMethod and msmodelslim's
# offline layout and vllm-ascend's AscendW8A8MXFP8DynamicFusedMoEMethod. # offline layout.
setattr( setattr(
layer, layer,
f"{weight_prefix}_weight", f"{weight_prefix}_weight",
python/sglang/srt/hardware_backend/npu/quantization/online_moe_methods.py
+2 -2
View File
@@ -1,4 +1,4 @@
"""Online (config-driven) quantized FusedMoE methods for Ascend NPU. """Online (config-driven) quantized FusedMoE methods for NPU.
These are the ``--quantization <scheme>`` entry points: the checkpoint holds These are the ``--quantization <scheme>`` entry points: the checkpoint holds
BF16/FP16 expert weights and the per-gmm kernels quantize them at load time. BF16/FP16 expert weights and the per-gmm kernels quantize them at load time.
@@ -28,7 +28,7 @@ class NPUMXFP8OnlineMoEMethod(UnquantizedFusedMoEMethod):
"""Online MXFP8 FusedMoE entry point (``--quantization mxfp8`` on A5). """Online MXFP8 FusedMoE entry point (``--quantization mxfp8`` on A5).
Weight creation, weight post-processing and the forward pass are identical Weight creation, weight post-processing and the forward pass are identical
to the unquantized Ascend path — the only difference is which per-gmm kernel to the unquantized NPU path — the only difference is which per-gmm kernel
the layer gets, so everything but ``create_moe_runner`` is inherited. the layer gets, so everything but ``create_moe_runner`` is inherited.
``NPUMXFP8MoEMethod`` then quantizes the BF16 expert weights to MXFP8 in ``NPUMXFP8MoEMethod`` then quantizes the BF16 expert weights to MXFP8 in
``process_weights_after_loading``. ``process_weights_after_loading``.
python/sglang/srt/models/qwen3_moe.py
+3 -3
View File
@@ -267,9 +267,9 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
routing_method_type=RoutingMethodType.Renormalize, routing_method_type=RoutingMethodType.Renormalize,
) )
# Router gate: description-driven quant, mirroring vllm-ascend. Only the # Router gate: description-driven quant. Only the offline ModelSlim path
# offline ModelSlim path (which carries a per-layer quant_model_description) # (which carries a per-layer quant_model_description) may quantise the
# may quantise the gate — if the checkpoint stored it as MXFP8 it is loaded # gate — if the checkpoint stored it as MXFP8 it is loaded
# and dequantised correctly instead of cast to bf16 without its block scale. # and dequantised correctly instead of cast to bf16 without its block scale.
# The online Fp8/mxfp8 path keeps the gate in bf16 (unchanged, verified). # The online Fp8/mxfp8 path keeps the gate in bf16 (unchanged, verified).
gate_quant_config = ( gate_quant_config = (