Add DASH-Q remote-code inference (Triton decode kernel)

README.md

@@ -9,7 +9,6 @@ tags:
 - post-training-quantization
 - int4
 ---
-
 ![DASH-Q](https://raw.githubusercontent.com/JaeminK/dashq/main/assets/dashq_banner.png)
 
 # Qwen3.5-35B-A3B-DASHQ-INT4-g64
@@ -17,21 +16,41 @@ tags:
 > **DASH-Q** — Diagonal-Aware Shrinkage for Robust PTQ.
 > `INT4` · group size 64 · **21.8928 GB** (from 71.9039 GB — **3.3x smaller**)
 
-DASH-Q checkpoints load with the lightweight DASH-Q runtime — linear layers are packed `PackedQuantizedLinear` modules, not plain Transformers weights.
+This checkpoint runs directly with Transformers: the packed quantized layers and a Triton decode kernel are bundled in the repository, so no additional package is needed.
+
+## Usage
 
-## Install
+```python
+from transformers import AutoModelForImageTextToText, AutoTokenizer
 
-```bash
-pip install git+https://github.com/JaeminK/dashq.git
+model = AutoModelForImageTextToText.from_pretrained(
+    "jkim96/Qwen3.5-35B-A3B-DASHQ-INT4-g64", trust_remote_code=True, device_map="cuda", dtype="auto"
+)
+tokenizer = AutoTokenizer.from_pretrained("jkim96/Qwen3.5-35B-A3B-DASHQ-INT4-g64")
+
+messages = [{"role": "user", "content": "Explain 2-bit quantization in one sentence."}]
+text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+inputs = tokenizer(text, return_tensors="pt").to(model.device)
+print(tokenizer.decode(model.generate(**inputs, max_new_tokens=256)[0]))
 ```
 
-## Load
+No extra package is required: this repository carries its own inference code
+(`modeling_dashq.py`) and a Triton weight-only GEMV decode kernel
+(`dashq_kernel.py`). Quantized layers are rebuilt from `dashq_config.json` at load
+time and converted to the kernel automatically; on CPU or without Triton the model
+falls back to a PyTorch dequantize-and-matmul path.
 
-```python
-from dashq import load_quantized
+Requirements: `transformers`, `torch`, and `triton` (bundled with CUDA builds of
+PyTorch). The [DASH-Q repository](https://github.com/JaeminK/dashq) is only needed
+to quantize your own models.
 
-model, tokenizer = load_quantized("jkim96/Qwen3.5-35B-A3B-DASHQ-INT4-g64", device_map="auto")
-```
+### Runtime format
+
+| Field | Value |
+| --- | --- |
+| Weights | group-wise asymmetric integers (4-bit, group size 64) packed into int32 words |
+| Decode kernel | Triton K-major GEMV (supports 2/3/4/8-bit at any group size) |
+| Prefill | unpack-and-matmul on the same buffers |
 
 ## Quantization
 

config.json

@@ -1,7 +1,30 @@
 {
   "architectures": [
-    "Qwen3_5MoeForConditionalGeneration"
+    "DashQQwen3_5MoeForConditionalGeneration"
   ],
+  "auto_map": {
+    "AutoModelForCausalLM": "modeling_dashq.DashQQwen3_5MoeForConditionalGeneration",
+    "AutoModelForImageTextToText": "modeling_dashq.DashQQwen3_5MoeForConditionalGeneration"
+  },
+  "dashq": {
+    "format": "dashq-packed-linear",
+    "format_version": 1,
+    "layer_metadata": "dashq_config.json",
+    "method": "dashq",
+    "n_quantized_modules": 20830,
+    "params": {
+      "bits": 4,
+      "group_size": 64,
+      "low_memory_optimization": false,
+      "moe_hessian_scope": "shared",
+      "n_samples": 128,
+      "scale_zero_dtype": "float16",
+      "symmetric": false,
+      "use_error_compensation": true,
+      "use_optimal_shrinkage": true,
+      "use_weighted_quantization": true
+    }
+  },
   "dtype": "bfloat16",
   "image_token_id": 248056,
   "model_type": "qwen3_5_moe",
@@ -120,4 +143,4 @@
   },
   "vision_end_token_id": 248054,
   "vision_start_token_id": 248053
-}
+}

dashq_kernel.py

@@ -0,0 +1,236 @@
+"""Triton weight-only GEMV backend for DASH-Q packed checkpoints.
+
+Group-wise asymmetric integer weights (the format DASH-Q emits) are stored
+K-major so that a decode-time GEMV reads each packed word exactly once with
+fully coalesced loads:
+
+    W_q : (K // elements_per_word, N) int32   (packed along K, N contiguous)
+    s,z : (K // group_size, N)                (one group per program)
+
+Supported bit widths: 2, 3, 4, 8 (and 1). 3-bit uses two bit-planes -- a
+2-bit plane plus a 1-bit plane -- which is exactly 3 bits per weight and is
+not covered by existing kernel libraries.
+
+Batched inputs (prefill) fall back to an unpack-and-matmul path that uses the
+same K-major buffers, so the original torch buffers can be released.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+try:
+    import triton
+    import triton.language as tl
+
+    TRITON_AVAILABLE = True
+except Exception:  # pragma: no cover - triton is an optional dependency
+    TRITON_AVAILABLE = False
+
+SUPPORTED_NBITS = (1, 2, 3, 4, 8)
+
+
+if TRITON_AVAILABLE:
+
+    @triton.jit
+    def _dashq_gemv_kernel(
+        x_ptr, w_ptr, lo_ptr, s_ptr, z_ptr, y_ptr,
+        N, K,
+        NBITS: tl.constexpr, EPS: tl.constexpr, GS: tl.constexpr,
+        BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    ):
+        pid_n = tl.program_id(0)
+        pid_k = tl.program_id(1) * 2
+        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_N), BLOCK_N)
+
+        # one scale/zero group per program (2 * BLOCK_K == GS)
+        k_m = (pid_k * BLOCK_K) // GS
+        scales = tl.load(s_ptr + k_m * N + offs_n).to(tl.float32)
+        zeros = tl.load(z_ptr + k_m * N + offs_n).to(tl.float32)
+
+        acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
+        for _ in tl.static_range(2):
+            a = tl.load(x_ptr + offs_k, eviction_policy="evict_last").to(tl.float32)
+            if NBITS == 3:
+                hw = tl.load(
+                    w_ptr + (offs_k // 16)[:, None] * N + offs_n[None, :],
+                    eviction_policy="evict_first",
+                )
+                lw = tl.load(
+                    lo_ptr + (offs_k // 32)[:, None] * N + offs_n[None, :],
+                    eviction_policy="evict_first",
+                )
+                q = (((hw >> (((offs_k % 16) * 2)[:, None])) & 3) << 1) | (
+                    (lw >> ((offs_k % 32)[:, None])) & 1
+                )
+            else:
+                wv = tl.load(
+                    w_ptr + (offs_k // EPS)[:, None] * N + offs_n[None, :],
+                    eviction_policy="evict_first",
+                )
+                q = (wv >> (((offs_k % EPS) * NBITS)[:, None])) & ((1 << NBITS) - 1)
+            b = (q.to(tl.float32) - zeros[None, :]) * scales[None, :]
+            acc += tl.sum(a[:, None] * b, axis=0)
+            offs_k += BLOCK_K
+
+        tl.atomic_add(y_ptr + offs_n, acc, sem="relaxed")
+
+
+def _pack_kmajor(q_kn: torch.Tensor, bits: int) -> torch.Tensor:
+    """(K, N) uint8 codes -> (K // eps, N) int32, value k in word k // eps."""
+    K, N = q_kn.shape
+    eps = 32 // bits
+    v = q_kn.to(torch.int32).reshape(K // eps, eps, N)
+    words = torch.zeros(K // eps, N, dtype=torch.int32, device=q_kn.device)
+    for j in range(eps):
+        words |= v[:, j, :] << (bits * j)
+    return words
+
+
+def _unpack_kmajor(words: torch.Tensor, bits: int, K: int) -> torch.Tensor:
+    eps = 32 // bits
+    WK, N = words.shape
+    shifts = (torch.arange(eps, device=words.device, dtype=torch.int32) * bits).view(1, eps, 1)
+    q = (words.view(WK, 1, N) >> shifts) & ((1 << bits) - 1)
+    return q.reshape(WK * eps, N)[:K]
+
+
+class TritonQuantLinear(nn.Module):
+    """Decode-optimized replacement for a DASH-Q PackedQuantizedLinear."""
+
+    def __init__(
+        self,
+        W_int: torch.Tensor,          # (out_features, in_features) integer codes
+        scale: torch.Tensor,          # (out_features, num_groups)
+        zero: torch.Tensor,           # (out_features, num_groups)
+        nbits: int,
+        group_size: int,
+        bias: Optional[torch.Tensor] = None,
+        out_dtype: torch.dtype = torch.float16,
+        block_n: int = 128,
+        num_warps: int = 1,
+    ) -> None:
+        super().__init__()
+        if not TRITON_AVAILABLE:
+            raise RuntimeError("Triton is not available.")
+        if nbits not in SUPPORTED_NBITS:
+            raise ValueError(f"Unsupported nbits for the Triton backend: {nbits}")
+
+        out_features, in_features = W_int.shape
+        if in_features % group_size != 0:
+            raise ValueError("in_features must be divisible by group_size.")
+        if group_size % 2 != 0:
+            raise ValueError("group_size must be even.")
+
+        self.out_features = out_features
+        self.in_features = in_features
+        self.nbits = int(nbits)
+        self.group_size = int(group_size)
+        self.out_dtype = out_dtype
+        self.block_n = int(block_n)
+        self.num_warps = int(num_warps)
+        self.block_k = self.group_size // 2
+
+        q_kn = W_int.t().contiguous().to(torch.uint8)
+        if nbits == 3:
+            self.register_buffer("W_q", _pack_kmajor(q_kn >> 1, 2))
+            self.register_buffer("W_lo", _pack_kmajor(q_kn & 1, 1))
+            self.eps = 16
+        else:
+            self.register_buffer("W_q", _pack_kmajor(q_kn, nbits))
+            self.register_buffer("W_lo", torch.zeros(1, dtype=torch.int32, device=q_kn.device))
+            self.eps = 32 // nbits
+        del q_kn
+
+        self.register_buffer("scale", scale.t().contiguous().to(out_dtype))
+        self.register_buffer("zero", zero.t().contiguous().to(out_dtype))
+        if bias is not None:
+            self.register_buffer("bias", bias.detach().clone().to(out_dtype))
+        else:
+            self.bias = None
+
+        # accumulator is seeded with the bias, so the kernel never adds it
+        # (each K-split program contributes once via atomic_add)
+        acc_init = torch.zeros(out_features, dtype=torch.float32, device=self.W_q.device)
+        if bias is not None:
+            acc_init.copy_(self.bias.float())
+        self.register_buffer("_acc_init", acc_init)
+        self.register_buffer("_acc", acc_init.clone())
+        self._grid = (
+            (out_features + self.block_n - 1) // self.block_n,
+            in_features // self.group_size,
+        )
+
+    @classmethod
+    def from_packed(cls, module: nn.Module, **kwargs) -> "TritonQuantLinear":
+        """Build from a dashq.quantization.PackedQuantizedLinear instance."""
+        from dashq.quantization import _unpack_int_values
+
+        K = int(getattr(module, "quant_in_features", module.in_features))
+        N = int(module.out_features)
+        W_int = _unpack_int_values(module.W_q_packed, module.nbits, module.numel).view(N, K)
+        num_groups = K // int(module.group_size)
+        scale = module.scale.view(N, num_groups)
+        zero = module.zero.view(N, num_groups)
+        bias = module.bias if getattr(module, "bias", None) is not None else None
+        return cls(
+            W_int,
+            scale,
+            zero,
+            int(module.nbits),
+            int(module.group_size),
+            bias=bias,
+            out_dtype=getattr(module, "linear_dtype", torch.float16),
+            **kwargs,
+        )
+
+    def dequantize_weight(self, dtype: torch.dtype) -> torch.Tensor:
+        """Returns W^T as (in_features, out_features), matching the K-major layout."""
+        if self.nbits == 3:
+            q = (_unpack_kmajor(self.W_q, 2, self.in_features).to(torch.int32) << 1) | (
+                _unpack_kmajor(self.W_lo, 1, self.in_features).to(torch.int32)
+            )
+        else:
+            q = _unpack_kmajor(self.W_q, self.nbits, self.in_features)
+        s = self.scale.repeat_interleave(self.group_size, dim=0).to(dtype)
+        z = self.zero.repeat_interleave(self.group_size, dim=0).to(dtype)
+        return (q.to(dtype) - z) * s
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = x.shape
+        tokens = x.numel() // shape[-1]
+        if tokens == 1 and x.is_cuda:
+            self._acc.copy_(self._acc_init)
+            _dashq_gemv_kernel[self._grid](
+                x.reshape(-1),
+                self.W_q,
+                self.W_lo,
+                self.scale,
+                self.zero,
+                self._acc,
+                self.out_features,
+                self.in_features,
+                self.nbits,
+                self.eps,
+                self.group_size,
+                self.block_n,
+                self.block_k,
+                num_warps=self.num_warps,
+            )
+            return self._acc.to(x.dtype).reshape(*shape[:-1], self.out_features)
+
+        w_t = self.dequantize_weight(x.dtype)
+        out = x.reshape(tokens, -1) @ w_t
+        if self.bias is not None:
+            out = out + self.bias.to(x.dtype)
+        return out.reshape(*shape[:-1], self.out_features)
+
+    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}, backend=triton"
+        )

modeling_dashq.py

@@ -0,0 +1,229 @@
+"""Self-contained DASH-Q inference code for this checkpoint.
+
+Generated by export_hf_repo.py -- do not edit by hand.
+
+Weights are stored as group-wise asymmetric integers packed into int32 words.
+The layout of every quantized layer is described by `dashq_config.json` in this
+repository. At load time each quantized module is converted to the K-major
+layout used by the DASH-Q Triton decode kernel (`dashq_kernel.py`); when Triton
+is unavailable (e.g. CPU) the modules fall back to an unpack-and-matmul path in
+pure PyTorch.
+"""
+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, Qwen3_5MoeForConditionalGeneration
+
+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):
+    """Holds the checkpoint buffers; converted to the Triton kernel after load."""
+
+    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 DashQQwen3_5MoeForConditionalGeneration(Qwen3_5MoeForConditionalGeneration):
+    """Qwen3_5MoeForConditionalGeneration 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) -> "DashQQwen3_5MoeForConditionalGeneration":
+        """Convert packed buffers to the Triton decode kernel (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)