modeling_dashq.py

"""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, Phi3ForCausalLM

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 DashQPhi3ForCausalLM(Phi3ForCausalLM):
    """Phi3ForCausalLM 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) -> "DashQPhi3ForCausalLM":
        """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)