"""Inference code for this DASH-Q checkpoint. Generated by export_hf_repo.py -- do not edit by hand. Weights are group-wise asymmetric integers packed into int32 words; the layout of each quantized layer is described by `dashq_config.json`. At load time the layers are converted to the format used by the Triton kernels in `dashq_kernel.py`, with a PyTorch dequantize-and-matmul fallback when Triton is unavailable. """ from __future__ import annotations import json import os from typing import Any, Dict, Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoConfig, MuseGlimmerForConditionalGeneration try: from .dashq_kernel import TRITON_AVAILABLE, SUPPORTED_NBITS, TritonQuantLinear except ImportError: # loaded as a flat module by trust_remote_code from dashq_kernel import TRITON_AVAILABLE, SUPPORTED_NBITS, TritonQuantLinear DASHQ_CONFIG_FILE = "dashq_config.json" def _unpack_int_values(packed: torch.Tensor, nbits: int, numel: int) -> torch.Tensor: values_per_word = max(1, 32 // nbits) mask = (1 << nbits) - 1 shifts = torch.arange(values_per_word, device=packed.device, dtype=torch.int32) * nbits out = (packed.view(-1, 1) >> shifts.view(1, -1)) & mask return out.reshape(-1)[:numel] class DashQPackedLinear(nn.Module): """Checkpoint buffers for one quantized layer.""" def __init__(self, in_features: int, out_features: int, nbits: int, group_size: int, bias: bool, dtype: torch.dtype, quant_in_features: Optional[int] = None) -> None: super().__init__() self.in_features = int(in_features) self.quant_in_features = int(quant_in_features or in_features) self.out_features = int(out_features) self.nbits = int(nbits) self.group_size = int(group_size) self.linear_dtype = dtype self.numel = self.out_features * self.quant_in_features self.num_groups = self.numel // self.group_size values_per_word = max(1, 32 // self.nbits) n_words = (self.numel + values_per_word - 1) // values_per_word self.register_buffer("W_q_packed", torch.zeros(n_words, dtype=torch.int32)) self.register_buffer("scale", torch.zeros(self.num_groups, 1, dtype=torch.float16)) self.register_buffer("zero", torch.zeros(self.num_groups, 1, dtype=torch.float16)) if bias: self.bias = nn.Parameter(torch.zeros(self.out_features, dtype=dtype), requires_grad=False) else: self.register_parameter("bias", None) self.kernel: Optional[nn.Module] = None def dequantize_weight(self, dtype: torch.dtype) -> torch.Tensor: W_int = _unpack_int_values(self.W_q_packed, self.nbits, self.numel) W_int = W_int.view(self.num_groups, self.group_size).to(dtype) W = (W_int - self.zero.to(dtype)) * self.scale.to(dtype) return W.view(self.out_features, self.quant_in_features) @torch.no_grad() def build_kernel(self) -> bool: if self.kernel is not None: return True if not TRITON_AVAILABLE or self.nbits not in SUPPORTED_NBITS: return False if self.quant_in_features % self.group_size or self.group_size % 2: return False if self.W_q_packed is None or not self.W_q_packed.is_cuda: return False W_int = _unpack_int_values(self.W_q_packed, self.nbits, self.numel) W_int = W_int.view(self.out_features, self.quant_in_features) ng = self.quant_in_features // self.group_size self.kernel = TritonQuantLinear( W_int, self.scale.view(self.out_features, ng), self.zero.view(self.out_features, ng), self.nbits, self.group_size, bias=self.bias.data if self.bias is not None else None, out_dtype=self.linear_dtype, ) del W_int self._buffers["W_q_packed"] = None self._buffers["scale"] = None self._buffers["zero"] = None return True def forward(self, x: torch.Tensor) -> torch.Tensor: if self.kernel is not None: return self.kernel(x) weight = self.dequantize_weight(x.dtype) bias = self.bias.to(x.dtype) if self.bias is not None else None return F.linear(x, weight, bias) def extra_repr(self) -> str: return (f"in_features={self.in_features}, out_features={self.out_features}, " f"nbits={self.nbits}, group_size={self.group_size}") def _set_module(root: nn.Module, name: str, new_module: nn.Module) -> None: parts = name.split(".") parent = root for part in parts[:-1]: parent = getattr(parent, part) setattr(parent, parts[-1], new_module) def _get_module(root: nn.Module, name: str) -> Optional[nn.Module]: obj = root for part in name.split("."): if not hasattr(obj, part): return None obj = getattr(obj, part) return obj _DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32} def _load_dashq_spec(model_id_or_path, **kwargs) -> Dict[str, Any]: """Read dashq_config.json from a local dir or the Hub.""" path = None if model_id_or_path is not None: local = os.path.join(str(model_id_or_path), DASHQ_CONFIG_FILE) if os.path.isfile(local): path = local if path is None and model_id_or_path is not None: try: from huggingface_hub import hf_hub_download path = hf_hub_download( repo_id=str(model_id_or_path), filename=DASHQ_CONFIG_FILE, revision=kwargs.get("revision"), token=kwargs.get("token"), cache_dir=kwargs.get("cache_dir"), ) except Exception: return {} if path is None: return {} with open(path, encoding="utf-8") as f: return json.load(f) def _swap_quantized_modules(model: nn.Module, modules: Dict[str, Any]) -> int: count = 0 for name, meta in modules.items(): target = _get_module(model, name) if target is None or isinstance(target, DashQPackedLinear): continue module = DashQPackedLinear( in_features=meta["in_features"], out_features=meta["out_features"], nbits=meta["nbits"], group_size=meta["group_size"], bias=getattr(target, "bias", None) is not None, dtype=_DTYPES.get(meta.get("linear_dtype", "float16"), torch.float16), quant_in_features=meta.get("quant_in_features"), ) _set_module(model, name, module) count += 1 return count class DashQMuseGlimmerForConditionalGeneration(MuseGlimmerForConditionalGeneration): """MuseGlimmerForConditionalGeneration whose linear layers hold DASH-Q packed quantized weights.""" def __init__(self, config): super().__init__(config) modules = (getattr(config, "dashq_modules", None) or {}) if modules: _swap_quantized_modules(self, modules) @classmethod def from_pretrained(cls, pretrained_model_name_or_path=None, *args, **kwargs): config = kwargs.pop("config", None) if config is None and pretrained_model_name_or_path is not None: config = AutoConfig.from_pretrained( pretrained_model_name_or_path, trust_remote_code=True, revision=kwargs.get("revision"), token=kwargs.get("token"), cache_dir=kwargs.get("cache_dir"), ) if config is not None and not getattr(config, "dashq_modules", None): spec = _load_dashq_spec(pretrained_model_name_or_path, **kwargs) config.dashq_modules = spec.get("quantized_modules", {}) model = super().from_pretrained(pretrained_model_name_or_path, *args, config=config, **kwargs) model.build_dashq_kernels() return model def build_dashq_kernels(self, verbose: bool = True) -> "DashQMuseGlimmerForConditionalGeneration": """Move the packed buffers to the Triton kernel layout (no-op off CUDA).""" total = built = 0 for module in self.modules(): if isinstance(module, DashQPackedLinear): total += 1 built += int(module.build_kernel()) if verbose and total: if built: print(f">> DASH-Q: {built}/{total} linear layers using the Triton decode kernel.") else: print(f">> DASH-Q: {total} quantized layers using the PyTorch fallback path.") return self def save_pretrained(self, *args, **kwargs): released = any( isinstance(m, DashQPackedLinear) and m.kernel is not None for m in self.modules() ) if released: raise RuntimeError( "This model has been converted to the DASH-Q Triton kernel, so the packed " "buffers are no longer materialized and saving would produce an incomplete " "checkpoint. Reload with build_dashq_kernels() skipped if you need to re-save." ) return super().save_pretrained(*args, **kwargs)