first commit

dev/functions/config_builder/__init__.py

@@ -0,0 +1,151 @@
+import torch
+
+from .qwen3moe import Qwen3MoEQuantConfigBuilder
+from .default import DefaultQuantConfigBuilder
+from .kimik2 import KimiK2QuantConfigBuilder
+from .glm4_7 import GLM4_7QuantConfigBuilder
+from ..specs import QuantArtifacts, ResolvedQuantConfig, DequantSpec
+
+
+def create_quant_config_builder(qconfig, custom_dequant_config, model_type):
+
+    if model_type in ["qwen3_moe", "minimax_m2", "minimax", "qwen3_5_moe", "deepseek_v3", "glm_moe_dsa"]:
+        return Qwen3MoEQuantConfigBuilder(qconfig, custom_dequant_config)
+    elif model_type == "kimik2":
+        return KimiK2QuantConfigBuilder(qconfig, custom_dequant_config)
+    elif model_type == "glm4_moe":
+        return GLM4_7QuantConfigBuilder(qconfig, custom_dequant_config)
+
+    return DefaultQuantConfigBuilder(
+        qconfig,
+        custom_dequant_config
+    )
+
+
+def derive_quant_key(weight_name: str, key_type: str, tensors: dict[str, torch.Tensor]) -> str | None:
+    """
+    根据 weight_name 推导量化辅助 tensor 名称。
+    """
+    # 定义所有可能的原始 weight 后缀
+    ORIG_WEIGHT_SUFFIXES = QWEIGHT_SUFFIXES
+
+    # 选择目标后缀表
+    if key_type == "scale":
+        target_suffixes = SCALE_SUFFIXES
+    elif key_type == "zero_point":
+        target_suffixes = ZERO_POINT_SUFFIXES
+    elif key_type == "qweight":
+        target_suffixes = QWEIGHT_SUFFIXES
+    elif key_type == "shape":
+        target_suffixes = SHAPE_SUFFIXES
+    else:
+        raise ValueError(f"Unknown key_type {key_type}")
+
+    candidates = []
+    for orig_suffix in ORIG_WEIGHT_SUFFIXES:
+        if weight_name.endswith(orig_suffix):
+            prefix = weight_name[: -len(orig_suffix)]
+            candidates.extend([prefix + s for s in target_suffixes])
+            break
+    else:
+        # weight_name 不匹配任何已知 weight 后缀，直接加 target_suffix
+        candidates.extend([weight_name + s for s in target_suffixes])
+
+    # 返回第一个存在于 tensors 的 key
+    for k in candidates:
+        if k in tensors:
+            return k
+    return None
+
+
+SCALE_SUFFIXES = [
+    ".weight_scale",
+    ".weight_scale_inv",
+    ".weight_block_scale",
+    "_scale",
+    "_scale_inv",
+    "_block_scale",
+]
+
+ZERO_POINT_SUFFIXES = [
+    ".weight_zero_point",
+    ".weight_zp",
+    "_zero_point",
+    "_zp",
+]
+
+QWEIGHT_SUFFIXES = [
+    ".weight",
+    ".weight_packed",
+    ".weight_q",
+    "_q",
+    "_packed",
+]
+
+SHAPE_SUFFIXES = [
+    ".weight_shape",
+    ".weight_block_shape",
+    "_shape",
+    "_block_shape",
+]
+
+
+def resolve_quant_artifacts(
+    weight_name: str,
+    tensors: dict[str, torch.Tensor],
+    rq: ResolvedQuantConfig,
+) -> QuantArtifacts:
+
+    if weight_name not in tensors:
+        raise KeyError(f"Missing quantized weight: {weight_name}")
+
+    if any(
+        weight_name.endswith(suffix)
+        for suffix in SCALE_SUFFIXES + ZERO_POINT_SUFFIXES + SHAPE_SUFFIXES
+    ):
+        return None
+
+    qweight_key = derive_quant_key(weight_name, "qweight", tensors)
+    scale_key = derive_quant_key(weight_name, "scale", tensors)
+    zero_point_key = derive_quant_key(weight_name, "zero_point", tensors)
+    shape_key = derive_quant_key(weight_name, "shape", tensors)
+
+    try:
+        qweight = tensors[qweight_key]
+    except KeyError:
+        raise KeyError(f"Missing quantized weight: {weight_name}")
+
+    # if qweight_key is None:
+    #     raise KeyError(f"Cannot find qweight for {weight_name}")
+
+    qweight = tensors[qweight_key]
+
+    if not scale_key:
+        raise KeyError(f"Missing scale tensor: {weight_name}")
+
+    scale = tensors[scale_key]
+
+    zero_point = None
+    if not rq.symmetric:
+        if zero_point_key is None:
+            raise KeyError(f"Cannot find zero_point tensor for {weight_name}")
+        zero_point = tensors[zero_point_key]
+
+    # shape = tensors.get(shape_key, None)
+    shape = None
+    if shape_key is not None:
+        shape = tensors[shape_key]
+
+    return QuantArtifacts(
+        qweight=qweight, scale=scale, zero_point=zero_point, shape=shape
+    )
+
+
+def to_dequant_spec(rq: ResolvedQuantConfig) -> DequantSpec:
+    return DequantSpec(
+        dtype=rq.dtype,
+        num_bits=rq.num_bits,
+        packed=rq.packed,
+        strategy=rq.strategy,
+        symmetric=rq.symmetric,
+    )

dev/functions/config_builder/base.py

@@ -0,0 +1,36 @@
+# config_builder/base.py
+import re
+
+from typing import Optional
+from ..specs import ResolvedQuantConfig
+
+
+class QuantConfigBuilder:
+
+    def resolve(
+        self,
+        weight_name: str,
+    ) -> Optional[ResolvedQuantConfig]:
+        pass
+    
+    def get_ignore(self, quantization_config: dict):
+        ignore = []
+
+        candidate_suffixes = ["modules_to_not_convert", "ignore"]
+        for cs in candidate_suffixes:
+            ig = quantization_config.get(cs)
+            if ig:
+                ignore.extend(ig)
+
+        return ignore
+    
+    def compile_ignore(self, ignore_list):
+        exact = set()
+        regex = []
+
+        for x in (ignore_list or []):
+            if x.startswith("re:"):
+                regex.append(re.compile(x[3:]))
+            else:
+                exact.add(x)
+        return list(exact), regex

dev/functions/config_builder/default.py

@@ -0,0 +1,56 @@
+import re
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class DefaultQuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        self.qconfig = qconfig
+        self.groups = qconfig["config_groups"]
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.custom_ignore)
+    
+    def _ignored(self, weight_name: str) -> bool:
+        if any(weight_name.startswith(ig) for ig in self.exact_ignore):
+            return True
+        if any(re.match(p, weight_name) for p in self.regex_ignore):
+            return True
+        
+        return False
+
+    def _match_group(self, weight_name: str):
+        for g in self.groups.values():
+            for p in g.get("weight_patterns", []):
+                if p.startswith("re:") and re.match(p[3:], weight_name):
+                    return g
+        return None
+    
+    def build_weight_quant_config(self, qconfig: dict, group_cfg: dict) -> dict:
+        wcfg = group_cfg["weights"]
+
+        return {
+            "dtype": wcfg["type"],  # int / float
+            "num_bits": wcfg["num_bits"],
+            "packed": qconfig["format"] == "pack-quantized",
+            "strategy": wcfg["strategy"],
+            "group_size": wcfg.get("group_size"),
+            "block_structure": wcfg.get("block_structure"),
+            "symmetric": wcfg["symmetric"],
+            "dynamic": wcfg["dynamic"],
+        }
+    
+    def resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        if len(self.groups) == 1:
+            group = next(iter(self.groups.values()))
+        else:
+            group = self._match_group(weight_name)
+            if group is None:
+                return None
+
+        return ResolvedQuantConfig(
+            **self.build_weight_quant_config(self.qconfig, group)
+        )

dev/functions/config_builder/glm4_7.py

@@ -0,0 +1,39 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class GLM4_7QuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        # get ignore from qconfig
+        self.original_ignore = self.get_ignore(qc)
+        # get ignore from custom_dequant_config
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_ignore+self.custom_ignore)
+
+    def _ignored(self, weight_name: str) -> bool:
+        if any(weight_name.startswith(ig) for ig in self.exact_ignore):
+            return True
+        if any(re.match(p, weight_name) for p in self.regex_ignore):
+            return True
+        
+        return False
+
+    def resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="float",
+            num_bits=8,
+            packed=False,
+            strategy="channel",
+            group_size=None,
+            block_structure=(0,0),
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/config_builder/kimik2.py

@@ -0,0 +1,38 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class KimiK2QuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        self.group_size = qconfig['config_groups']['group_0']['weights']['group_size']
+        self.original_ignore = self.get_ignore(qc)
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_ignore+self.custom_ignore)
+
+    def _ignored(self, weight_name: str) -> bool:
+        if any(weight_name.startswith(ig) for ig in self.exact_ignore):
+            return True
+        if any(re.match(p, weight_name) for p in self.regex_ignore):
+            return True
+        
+        return False
+
+    def resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="int",
+            num_bits=4,
+            packed=True,
+            strategy="group",
+            group_size=self.group_size,
+            block_structure=None,
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/config_builder/kimik25.py

@@ -0,0 +1,38 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class KimiK25QuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        self.group_size = qconfig['config_groups']['group_0']['weights']['group_size']
+        self.original_ignore = self.get_ignore(qc)
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_ignore+self.custom_ignore)
+
+    def _ignored(self, weight_name: str) -> bool:
+        if any(weight_name.startswith(ig) for ig in self.exact_ignore):
+            return True
+        if any(re.match(p, weight_name) for p in self.regex_ignore):
+            return True
+        
+        return False
+
+    def resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="int",
+            num_bits=4,
+            packed=True,
+            strategy="group",
+            group_size=self.group_size,
+            block_structure=None,
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/config_builder/qwen3moe.py

@@ -0,0 +1,40 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class Qwen3MoEQuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        self.block_h, self.block_w = qc["weight_block_size"]
+        # get ignore from qconfig
+        self.original_ignore = self.get_ignore(qc)
+        # get ignore from custom_dequant_config
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_ignore+self.custom_ignore)
+
+    def _ignored(self, weight_name: str) -> bool:
+        if any(weight_name.startswith(ig) for ig in self.exact_ignore):
+            return True
+        if any(re.match(p, weight_name) for p in self.regex_ignore):
+            return True
+        
+        return False
+
+    def resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="float",
+            num_bits=8,
+            packed=False,
+            strategy="block",
+            group_size=None,
+            block_structure=(self.block_h, self.block_w),
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/dequant.py

@@ -0,0 +1,156 @@
+import re
+import os
+import json
+import torch
+
+from safetensors.torch import load_file
+
+from .specs import *
+from .utils import compile_ignore, get_ignore, estimate_gpu_bytes
+from .dequant_register import DEQUANT_KERNEL_TABLE
+from .config_builder import create_quant_config_builder, resolve_quant_artifacts, to_dequant_spec
+
+
+
+from safetensors import safe_open
+
+
+class LazyTensorDict(dict):
+
+    def __init__(self, base_tensors, tensor_index, input_dir):
+        super().__init__(base_tensors)
+        self.tensor_index = tensor_index
+        self.input_dir = input_dir
+        self.file_cache = {}
+
+    def _load_tensor(self, name):
+
+        file_name = self.tensor_index[name]
+        path = os.path.join(self.input_dir, file_name)
+
+        if file_name not in self.file_cache:
+            self.file_cache[file_name] = safe_open(path, framework="pt")
+        return self.file_cache[file_name].get_tensor(name)
+
+    def __getitem__(self, key):
+        try:
+            return super().__getitem__(key)
+        except KeyError:
+            # 键不在字典中，检查是否在索引中
+            if self.tensor_index and key in self.tensor_index:
+                tensor = self._load_tensor(key)
+                self[key] = tensor  # 现在可以安全地添加
+                return tensor
+            else:
+                raise KeyError(key)
+    
+    def __contains__(self, key):
+
+        if super().__contains__(key):
+            return True
+
+        if self.tensor_index and key in self.tensor_index:
+            return True
+
+        return False
+
+
+
+def dequantize_(
+    tensors: dict[str, torch.Tensor],
+    qconfig: dict,
+    custom_dequant_config,
+    output_dtype=torch.bfloat16,
+    device: torch.device = "cpu",
+    **kwargs
+):
+    ignore_exact, ignore_regex = compile_ignore(
+        get_ignore(qconfig),
+        get_ignore(custom_dequant_config),
+    )
+
+    builder = create_quant_config_builder(
+        qconfig,
+        custom_dequant_config,
+        model_type=kwargs.get("model_type")
+    )
+
+    outputs = {}
+    gpu_cache = {}
+    gpu_cache_size = 0
+    max_gpu_cache_size = 16 * (1024**3)
+    weight_names = list(tensors.keys())
+    # for weight_name in tensors.keys():
+    for weight_name in weight_names:
+        rq = builder.resolve(weight_name)
+
+        if rq is None:  # 包含scale的一定不会在这里
+            assert "scale" not in weight_name, f"{weight_name}"
+            outputs[weight_name] = tensors[weight_name]
+            continue
+
+        assert not rq.dynamic, "Dynamic quantization not supported in offline dequant"
+
+        artifacts = resolve_quant_artifacts(weight_name, tensors, rq)
+
+        if artifacts is None:
+            # 跳过scale, shape等
+            continue
+
+        spec = to_dequant_spec(rq)
+        if spec not in DEQUANT_KERNEL_TABLE:
+            raise RuntimeError(f"No dequant kernel registered for {spec}")
+
+        kernel = DEQUANT_KERNEL_TABLE[spec]
+
+        gpu_cache_size += estimate_gpu_bytes(artifacts, rq)
+
+        try:
+            output = kernel(
+                artifacts=artifacts,
+                rq=rq,
+                output_dtype=output_dtype,
+                device=device,
+            )
+        except Exception as e:
+            print(weight_name)
+            raise e
+        # 处理名称
+        weight_name = weight_name.replace("weight_packed", "weight").replace("qweight", "weight")
+        gpu_cache[weight_name] = output
+
+        if gpu_cache_size >= max_gpu_cache_size:
+            for k, v in gpu_cache.items():
+                outputs[k] = v.cpu()
+            gpu_cache.clear()
+            gpu_cache_size = 0
+    
+    for k, v in gpu_cache.items():
+        outputs[k] = v.cpu()
+        
+    gpu_cache.clear()
+    gpu_cache_size = 0
+
+    return outputs
+
+
+def dequant(file_name: str,
+            quant_config,
+            custom_dequant_config,
+            device: torch.device,
+            model_type,
+            tensor_index=None,
+            input_dir=None):
+
+    base_tensors = load_file(file_name)
+
+    if tensor_index is not None:
+        tensors = LazyTensorDict(base_tensors, tensor_index, input_dir)
+    else:
+        tensors = base_tensors
+
+    return dequantize_(tensors, quant_config, custom_dequant_config, device=device, model_type=model_type)
+
+
+if __name__ == "__main__":
+    ...

dev/functions/dequant_kernel_dispatch.py

@@ -0,0 +1,16 @@
+from core import weight_dequant_fp8e4m3_to_bf16
+
+import torch
+
+
+def fake_kernel(*args, **kwargs):
+    return args
+
+
+def dispatch_kernel(dtype):
+    if dtype in [torch.bfloat16, torch.float16, torch.half, torch.float, torch.float32]:
+        return fake_kernel
+    elif dtype == torch.float8_e4m3fn:
+        return weight_dequant_fp8e4m3_to_bf16
+    else:
+        raise NotImplemented

dev/functions/dequant_register.py

@@ -0,0 +1,20 @@
+from .specs import DequantSpec
+from core.kernels import (dequant_fp8_block_sym, 
+                          dequant_int4_group_sym,
+                          dequant_fp8_channel_sym)
+
+DEQUANT_KERNEL_TABLE = {}
+
+DEQUANT_KERNEL_TABLE[
+    DequantSpec("float", 8, False, "block", True)
+] = dequant_fp8_block_sym
+
+
+DEQUANT_KERNEL_TABLE[
+    DequantSpec("int", 4, True, "group", True)
+] = dequant_int4_group_sym
+
+
+DEQUANT_KERNEL_TABLE[
+    DequantSpec("float", 8, False, "channel", True)
+] = dequant_fp8_channel_sym

dev/functions/processor.py

@@ -0,0 +1,45 @@
+
+def quant_post_process(artifacts, version, model_type):
+    quant_dtype = artifacts.quant_type
+    if artifacts.extra is not None:
+        i8scales = artifacts.extra.get("i8scales")
+        i8zeros = artifacts.extra.get("i8zeros")
+    else:
+        i8scales = None
+        i8zeros = None
+
+    for idx, name in enumerate(artifacts.names):
+        weight = artifacts.qweight[idx]  # [#experts, m, n]
+        scale = artifacts.scale[idx]  # [#experts, m, 1]
+        i8scale = None if i8scales is None else i8scales[idx]
+        i8zero = None if i8zeros is None else i8zeros[idx]
+
+        num_experts = weight.shape[0]
+        
+        if "expert" in name and "shared" not in name:
+            scale = scale.contiguous().view(num_experts, 1, -1) if version == 2 else scale.contiguous()
+            if model_type == "qwen3_vl_moe":
+                weight = weight.transpose(-2, -1).contiguous()
+                if quant_dtype == "int4":
+                    scale = scale.transpose(-2, -1).contiguous()
+                elif quant_dtype == "int8":
+                    scale = scale.transpose(-2, -1).contiguous()
+            
+            if quant_dtype == "int8":
+                scale = scale.transpose(-2, -1).contiguous()
+
+            if i8scale is not None:
+                i8scale.squeeze_(0)
+                assert i8scale.dim() == 2
+            
+            if i8zero is not None:
+                i8zero.squeeze_(0)
+                assert i8zero.dim() == 2
+        
+        if num_experts == 1:
+            weight = weight.squeeze(0)
+            scale = scale.squeeze(0)
+        
+        # TODO: i8scale, i8zero
+        
+        yield name, weight, scale, i8scale, i8zero

dev/functions/quant.py

@@ -0,0 +1,159 @@
+import torch
+
+from collections import defaultdict
+from safetensors.torch import load_file
+
+from .processor import quant_post_process
+from .specs import QuantBatch, QuantBatchArtifacts, QuantSpec
+from .utils import compile_ignore as compile_pattern, should_ignore, should_quant, get_format
+from core.kernels import weight_quant_bf16_to_int8, weight_quant_bf16_to_int4pack8
+
+
+def quant(
+    file_name: str | None = None,
+    state_dict: dict[str, torch.Tensor] | None = None,
+    custom_quant_config: dict | None = None,  # from model.yaml
+    device: torch.device = "cpu",
+    model_type: str | None = None
+):
+    if file_name:
+        assert state_dict is None
+        state_dict = load_file(file_name)
+    else:
+        assert state_dict is not None
+    
+    exact_ignore, regex_ignore = compile_pattern(custom_quant_config['ignore'] or [])
+
+    int4_exact_pattern, int4_regex_pattern = compile_pattern(custom_quant_config['class']['int4'])
+    int8_exact_pattern, int8_regex_pattern = compile_pattern(custom_quant_config['class']['int8'])
+
+    batch_size = custom_quant_config['batch_size']
+
+    buffer = defaultdict(list)  # key = (quant_type, shape, symmetric, version, format)
+    outputs = {}
+    quant_map = {}
+
+    for weight_name, weight in state_dict.items():
+        # 不量化
+        if should_ignore(weight_name, exact_ignore, regex_ignore):
+            outputs[weight_name] = weight.cpu()
+            quant_map[weight_name] = str(weight.dtype)
+            continue
+        
+        if should_quant(weight_name, int4_exact_pattern, int4_regex_pattern):
+            qs = QuantSpec(
+                quant_type="int4",
+                shape=tuple(weight.shape),
+                symmetric=custom_quant_config['symmetric'],
+                version=custom_quant_config['version'],
+                format=get_format(weight_name, custom_quant_config['format'])
+            )
+            quant_map[weight_name] = f"int4-{get_format(weight_name, custom_quant_config['format'])}"
+        elif should_quant(weight_name, int8_exact_pattern, int8_regex_pattern):
+            qs = QuantSpec(
+                quant_type="int8",
+                shape=tuple(weight.shape),
+                symmetric=custom_quant_config['symmetric'],
+                version=custom_quant_config['version'],
+                format=get_format(weight_name, custom_quant_config['format'])
+            )
+            quant_map[weight_name] = f"int8-{get_format(weight_name, custom_quant_config['format'])}"
+        else:
+            raise ValueError(f"{weight_name} missing in ignore or quant pattern!!!")
+
+        buffer[qs].append(weight_name) 
+
+    # 构建batch
+    for qs, weight_names in buffer.items():
+        for i in range(0, len(weight_names), batch_size):
+            chunk_names = weight_names[i:i+batch_size]
+            tensors = []
+            for name in chunk_names:
+                w = state_dict[name]
+                if w.device != device:
+                    w = w.to(device, non_blocking=True)
+                tensors.append(w)
+            
+            batch_tensor = torch.stack(tensors, dim=0)  # [batch_size, num_experts, o, i] / [batch_size, o, i]
+
+            batch = QuantBatch(
+                names=chunk_names,
+                tensor=batch_tensor.transpose(-2,-1) if model_type == "qwen3_vl_moe" and "expert" in chunk_names[0] else batch_tensor,
+                device=device,
+                symmetric=qs.symmetric,
+                version=qs.version,
+                format=qs.format
+            )
+
+            # if model_type is qwen3-vl-moe, the order of input and output (moe) need to be transposed.
+            # and the batch.tensor.shape = (batch_size, #experts, input, output)
+            if batch.tensor.ndim == 3: # (bs, m, n)
+                batch.tensor.unsqueeze_(1)
+            
+            assert batch.tensor.ndim == 4, f"The len of batch.tensor should be 4, but got {batch.tensor.ndim}"
+
+            if qs.quant_type == "int8":
+                temp_qw, temp_s = [], []
+                for idx in range(batch.tensor.shape[1]):
+                    quant_weight, scale = weight_quant_bf16_to_int8(batch.tensor[:, idx])  # (batch_size, m, n) , (batch_size, m)
+                    temp_qw.append(quant_weight.unsqueeze(1))
+                    temp_s.append(scale.unsqueeze(1))
+                
+                quant_weight = torch.cat(temp_qw, dim=1)  # (batch_size, #experts, m, n)
+                scale = torch.cat(temp_s, dim=1)  # (bach_size, #experts, m)
+
+                artifacts = QuantBatchArtifacts(
+                    names=batch.names,
+                    qweight=quant_weight,
+                    scale=scale,
+                    quant_type=qs.quant_type
+                )
+            elif qs.quant_type == "int4":
+                temp_qw, temp_s, temp_i8scales, temp_i8zeros = [], [], [], []
+                i8scales, i8zeros = None, None
+                for idx in range(batch.tensor.shape[1]):
+                    quant_weight, scale, i8scales, i8zeros = weight_quant_bf16_to_int4pack8(
+                        batch.tensor[:, idx],
+                        block_size=custom_quant_config.get("block_size", 128),
+                        group_size=custom_quant_config.get("group_size", -1),
+                        format=qs.format,
+                        symmetric=qs.symmetric,
+                        version=qs.version
+                    )
+                    temp_qw.append(quant_weight.unsqueeze(1))
+                    temp_s.append(scale.unsqueeze(1))
+                    if i8scales is not None:
+                        temp_i8scales.append(i8scales.unsqueeze(1))
+                    if i8zeros is not None:
+                        temp_i8zeros.append(i8zeros.unsqueeze(1))
+                    
+                quant_weight = torch.cat(temp_qw, dim=1)  # (batch_size, #experts, m, n)
+                scale = torch.cat(temp_s, dim=1)  # (bach_size, #experts, m, n)
+
+                if temp_i8scales:
+                    i8scales = torch.cat(temp_i8scales, dim=1)
+                if temp_i8zeros:
+                    i8zeros = torch.cat(temp_i8zeros, dim=1)
+
+                artifacts = QuantBatchArtifacts(
+                    names=batch.names,
+                    qweight=quant_weight,
+                    scale=scale,
+                    extra={
+                        "i8scales": i8scales,
+                        "i8zeros": i8zeros,
+                    },
+                    quant_type=qs.quant_type
+                )
+            else:
+                raise RuntimeError(f"Unsupported quant type {qs.quant_type}")
+            
+            for name, weight, scale, i8scale, i8zero in quant_post_process(artifacts, qs.version, model_type):
+                outputs[name] = weight.cpu()
+                outputs[name+"_scale"] = scale.cpu()
+                if i8scale is not None:
+                    outputs[name.removesuffix(".weight")+".i8_weight_scale"] = i8scale.cpu()
+                if i8zero is not None:
+                    outputs[name.removesuffix(".weight")+".i8_weight_zero"] = i8zero.cpu()
+    
+    return outputs, quant_map

dev/functions/quant.py.pre_glm5_group64.bak

@@ -0,0 +1,157 @@
+import torch
+
+from collections import defaultdict
+from safetensors.torch import load_file
+
+from .processor import quant_post_process
+from .specs import QuantBatch, QuantBatchArtifacts, QuantSpec
+from .utils import compile_ignore as compile_pattern, should_ignore, should_quant, get_format
+from core.kernels import weight_quant_bf16_to_int8, weight_quant_bf16_to_int4pack8
+
+
+def quant(
+    file_name: str | None = None,
+    state_dict: dict[str, torch.Tensor] | None = None,
+    custom_quant_config: dict | None = None,  # from model.yaml
+    device: torch.device = "cpu",
+    model_type: str | None = None
+):
+    if file_name:
+        assert state_dict is None
+        state_dict = load_file(file_name)
+    else:
+        assert state_dict is not None
+    
+    exact_ignore, regex_ignore = compile_pattern(custom_quant_config['ignore'] or [])
+
+    int4_exact_pattern, int4_regex_pattern = compile_pattern(custom_quant_config['class']['int4'])
+    int8_exact_pattern, int8_regex_pattern = compile_pattern(custom_quant_config['class']['int8'])
+
+    batch_size = custom_quant_config['batch_size']
+
+    buffer = defaultdict(list)  # key = (quant_type, shape, symmetric, version, format)
+    outputs = {}
+    quant_map = {}
+
+    for weight_name, weight in state_dict.items():
+        # 不量化
+        if should_ignore(weight_name, exact_ignore, regex_ignore):
+            outputs[weight_name] = weight.cpu()
+            quant_map[weight_name] = str(weight.dtype)
+            continue
+        
+        if should_quant(weight_name, int4_exact_pattern, int4_regex_pattern):
+            qs = QuantSpec(
+                quant_type="int4",
+                shape=tuple(weight.shape),
+                symmetric=custom_quant_config['symmetric'],
+                version=custom_quant_config['version'],
+                format=get_format(weight_name, custom_quant_config['format'])
+            )
+            quant_map[weight_name] = f"int4-{get_format(weight_name, custom_quant_config['format'])}"
+        elif should_quant(weight_name, int8_exact_pattern, int8_regex_pattern):
+            qs = QuantSpec(
+                quant_type="int8",
+                shape=tuple(weight.shape),
+                symmetric=custom_quant_config['symmetric'],
+                version=custom_quant_config['version'],
+                format=get_format(weight_name, custom_quant_config['format'])
+            )
+            quant_map[weight_name] = f"int8-{get_format(weight_name, custom_quant_config['format'])}"
+        else:
+            raise ValueError(f"{weight_name} missing in ignore or quant pattern!!!")
+
+        buffer[qs].append(weight_name) 
+
+    # 构建batch
+    for qs, weight_names in buffer.items():
+        for i in range(0, len(weight_names), batch_size):
+            chunk_names = weight_names[i:i+batch_size]
+            tensors = []
+            for name in chunk_names:
+                w = state_dict[name]
+                if w.device != device:
+                    w = w.to(device, non_blocking=True)
+                tensors.append(w)
+            
+            batch_tensor = torch.stack(tensors, dim=0)  # [batch_size, num_experts, o, i] / [batch_size, o, i]
+
+            batch = QuantBatch(
+                names=chunk_names,
+                tensor=batch_tensor.transpose(-2,-1) if model_type == "qwen3_vl_moe" and "expert" in chunk_names[0] else batch_tensor,
+                device=device,
+                symmetric=qs.symmetric,
+                version=qs.version,
+                format=qs.format
+            )
+
+            # if model_type is qwen3-vl-moe, the order of input and output (moe) need to be transposed.
+            # and the batch.tensor.shape = (batch_size, #experts, input, output)
+            if batch.tensor.ndim == 3: # (bs, m, n)
+                batch.tensor.unsqueeze_(1)
+            
+            assert batch.tensor.ndim == 4, f"The len of batch.tensor should be 4, but got {batch.tensor.ndim}"
+
+            if qs.quant_type == "int8":
+                temp_qw, temp_s = [], []
+                for idx in range(batch.tensor.shape[1]):
+                    quant_weight, scale = weight_quant_bf16_to_int8(batch.tensor[:, idx])  # (batch_size, m, n) , (batch_size, m)
+                    temp_qw.append(quant_weight.unsqueeze(1))
+                    temp_s.append(scale.unsqueeze(1))
+                
+                quant_weight = torch.cat(temp_qw, dim=1)  # (batch_size, #experts, m, n)
+                scale = torch.cat(temp_s, dim=1)  # (bach_size, #experts, m)
+
+                artifacts = QuantBatchArtifacts(
+                    names=batch.names,
+                    qweight=quant_weight,
+                    scale=scale,
+                    quant_type=qs.quant_type
+                )
+            elif qs.quant_type == "int4":
+                temp_qw, temp_s, temp_i8scales, temp_i8zeros = [], [], [], []
+                i8scales, i8zeros = None, None
+                for idx in range(batch.tensor.shape[1]):
+                    quant_weight, scale, i8scales, i8zeros = weight_quant_bf16_to_int4pack8(
+                        batch.tensor[:, idx],
+                        block_size=custom_quant_config.get("block_size", 128),
+                        group_size=custom_quant_config.get("group_size", -1),
+                        format=qs.format,
+                        symmetric=qs.symmetric,
+                        version=qs.version
+                    )
+                    temp_qw.append(quant_weight.unsqueeze(1))
+                    temp_s.append(scale.unsqueeze(1))
+                    if i8scales is not None:
+                        temp_i8scales.append(i8scales.unsqueeze(1))
+                        temp_i8zeros.append(i8zeros.unsqueeze(1))
+                    
+                quant_weight = torch.cat(temp_qw, dim=1)  # (batch_size, #experts, m, n)
+                scale = torch.cat(temp_s, dim=1)  # (bach_size, #experts, m, n)
+
+                if len(temp_i8scales) > 0:
+                    i8scales = torch.cat(temp_i8scales, dim=1)
+                    i8zeros = torch.cat(temp_i8zeros, dim=1)
+
+                artifacts = QuantBatchArtifacts(
+                    names=batch.names,
+                    qweight=quant_weight,
+                    scale=scale,
+                    extra={
+                        "i8scales": i8scales,
+                        "i8zeros": i8zeros,
+                    },
+                    quant_type=qs.quant_type
+                )
+            else:
+                raise RuntimeError(f"Unsupported quant type {qs.quant_type}")
+            
+            for name, weight, scale, i8scale, i8zero in quant_post_process(artifacts, qs.version, model_type):
+                outputs[name] = weight.cpu()
+                outputs[name+"_scale"] = scale.cpu()
+                if i8scale is not None:
+                    outputs[name+"_i8scales"] = i8scale.cpu()
+                if i8zero is not None:
+                    outputs[name+"_i8zeros"] = i8zero.cpu()
+    
+    return outputs, quant_map

dev/functions/rules.py

@@ -0,0 +1 @@
+

dev/functions/saver/base.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+from typing import Dict
+import torch
+
+
+class BaseSaver(ABC):
+    def __init__(self, output_dir: str):
+        self.output_dir = output_dir
+
+    @abstractmethod
+    def save(
+        self,
+        state_dict: Dict[str, torch.Tensor],
+        src_file: str,
+    ):
+        pass

dev/functions/specs.py

@@ -0,0 +1,60 @@
+import torch
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class DequantSpec:
+    dtype: str  # "int" | "float"
+    num_bits: int  # 4 / 8
+    packed: bool  # True / False
+    strategy: str  # "group" | "block" | "channel" | "tensor"
+    symmetric: bool  # True / False
+
+
+@dataclass
+class ResolvedQuantConfig:
+    dtype: str
+    num_bits: int
+    packed: bool
+    strategy: str
+    group_size: int | None
+    block_structure: tuple | None
+    symmetric: bool
+    dynamic: bool
+
+
+@dataclass
+class QuantArtifacts:
+    qweight: torch.Tensor
+    scale: torch.Tensor
+    zero_point: torch.Tensor | None
+    shape: torch.Tensor | None
+
+
+@dataclass(frozen=True)
+class QuantSpec:
+    quant_type: str
+    shape: tuple
+    symmetric: bool = True
+    version: int = 2
+    format: str = "TN"
+
+
+@dataclass
+class QuantBatch:
+    names: list[str]
+    tensor: torch.Tensor
+    device: torch.device
+    symmetric: bool = True
+    version: int = 2
+    format: str = "TN"
+
+
+@dataclass
+class QuantBatchArtifacts:
+    names: list[str]
+    qweight: torch.Tensor
+    scale: torch.Tensor
+    quant_type: str | None = None
+    extra: dict | None = None

dev/functions/utils.py

@@ -0,0 +1,84 @@
+import re
+import torch
+
+
+def compile_ignore(ignore_list: list | None = None, 
+                   extra_ignore_list: list | None = None):
+    exact = set()
+    regex = []
+
+    for x in (ignore_list or []) + (extra_ignore_list or []):
+        if x.startswith("re:"):
+            regex.append(re.compile(x[3:]))
+        else:
+            exact.add(x)
+    return exact, regex
+
+
+def should_ignore(weight_name, exact_ignore: list, regex_ignore: list):
+    if any(weight_name.startswith(ig) for ig in exact_ignore):
+        return True
+    if any(re.match(p, weight_name) for p in regex_ignore):
+        return True
+    
+    return False
+
+
+def should_quant(weight_name, exact_pattern: list, regex_pattern: list):
+    return should_ignore(weight_name, exact_pattern, regex_pattern)
+
+    
+def get_ignore(quantization_config: dict):
+    ignore = []
+
+    candidate_suffixes = ["modules_to_not_convert", "ignore"]
+    for cs in candidate_suffixes:
+        ig = quantization_config.get(cs)
+        if ig:
+            ignore.extend(ig)
+
+    return ignore
+
+
+def get_format(weight_name, format):
+    tn_group = format['TN']
+    if tn_group:
+        exact, regex = compile_ignore(tn_group)
+        if should_ignore(weight_name, exact, regex):
+            return "TN"
+    return "NN"
+
+
+def estimate_gpu_bytes(
+    artifacts,
+    rq,
+    output_dtype=torch.bfloat16,
+) -> int:
+    """
+    Conservative estimation of peak GPU memory usage (bytes)
+    during dequantization of ONE weight.
+
+    Used for batching / flushing decision only.
+    """
+
+    bytes_total = 0
+
+    def tensor_bytes(t: torch.Tensor) -> int:
+        return t.numel() * t.element_size()
+
+    qweight = artifacts.qweight
+    scale = artifacts.scale
+
+    bytes_total += tensor_bytes(qweight)
+    bytes_total += tensor_bytes(scale)
+
+    if artifacts.zero_point is not None:
+        bytes_total += tensor_bytes(artifacts.zero_point)
+
+    bytes_total += qweight.numel() * 4  # fp32
+
+    out_elem_size = torch.tensor([], dtype=output_dtype).element_size()
+    bytes_total += qweight.numel() * out_elem_size
+
+    SAFETY_FACTOR = 1.3
+    return int(bytes_total * SAFETY_FACTOR)