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

README.md

@@ -32,10 +32,9 @@ inputs = tokenizer(text, return_tensors="pt").to(model.device)
 print(tokenizer.decode(model.generate(**inputs, max_new_tokens=256)[0]))
 ```
 
-`trust_remote_code=True` is required: the checkpoint ships the quantized-layer
-implementation (`modeling_dashq.py`) and a Triton weight-only GEMV decode kernel
-(`dashq_kernel.py`). Without Triton or on CPU the model falls back to a PyTorch
-dequantize-and-matmul path.
+`trust_remote_code=True` is required: the checkpoint ships its quantized-layer
+implementation (`modeling_dashq.py`) and Triton kernels (`dashq_kernel.py`).
+Without Triton, or on CPU, it falls back to dequantize-and-matmul in PyTorch.
 
 ### Requirements
 

dashq_kernel.py

@@ -1,18 +1,16 @@
-"""Triton weight-only GEMV backend for DASH-Q packed checkpoints.
+"""Triton kernels for group-wise asymmetric integer weights.
 
-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:
+Weights are stored K-major: W_q has shape (K // elements_per_word, N) with
+values packed along K, and scale/zero have shape (K // group_size, N).
 
-    W_q : (K // elements_per_word, N) int32   (packed along K, N contiguous)
-    s,z : (K // group_size, N)                (one group per program)
+Three kernels are selected by the number of input rows M:
 
-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.
+    M == 1        GEMV
+    2 <= M <= 32  fused dequantize-GEMM with split-K
+    M > 32        fused dequantize-GEMM, accumulator kept in registers
 
-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.
+Supported bit widths are 1, 2, 3, 4 and 8. 3-bit is stored as a 2-bit plane
+plus a 1-bit plane, so it occupies exactly 3 bits per weight.
 """
 from __future__ import annotations
 
@@ -26,11 +24,13 @@ try:
     import triton.language as tl
 
     TRITON_AVAILABLE = True
-except Exception:  # pragma: no cover - triton is an optional dependency
+except Exception:  # triton is optional
     TRITON_AVAILABLE = False
 
 SUPPORTED_NBITS = (1, 2, 3, 4, 8)
 
+_GEMM_CONFIG_CACHE = {}
+
 
 if TRITON_AVAILABLE:
 
@@ -46,7 +46,7 @@ if TRITON_AVAILABLE:
         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)
+        # 2 * BLOCK_K == GS, so a program covers exactly one scale group.
         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)
@@ -76,12 +76,69 @@ if TRITON_AVAILABLE:
             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")
 
+    @triton.jit
+    def _dashq_gemm_kernel(
+        x_ptr, w_ptr, lo_ptr, s_ptr, z_ptr, y_ptr,
+        M, N, K,
+        NBITS: tl.constexpr, EPS: tl.constexpr, GS: tl.constexpr,
+        BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+        SPLIT_K: tl.constexpr, OUT_DTYPE: tl.constexpr,
+    ):
+        """y[M, N] = x[M, K] @ dequantize(w)[K, N]
+
+        BLOCK_K divides the group size, so a K-tile lies inside one group and the
+        scale/zero load is a single (1, BLOCK_N) vector.
+        """
+        pid_m = tl.program_id(0)
+        pid_n = tl.program_id(1)
+        pid_k = tl.program_id(2)
+
+        offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        mask_m = offs_m < M
+        mask_n = offs_n < N
+
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+        for t in range(pid_k, tl.cdiv(K, BLOCK_K), SPLIT_K):
+            k0 = t * BLOCK_K
+            offs_k = k0 + tl.arange(0, BLOCK_K)
+            mask_k = offs_k < K
+
+            x = tl.load(x_ptr + offs_m[:, None] * K + offs_k[None, :],
+                        mask=mask_m[:, None] & mask_k[None, :], other=0.0)
+
+            if NBITS == 3:
+                hw = tl.load(w_ptr + (offs_k // 16)[:, None] * N + offs_n[None, :],
+                             mask=mask_k[:, None] & mask_n[None, :], other=0)
+                lw = tl.load(lo_ptr + (offs_k // 32)[:, None] * N + offs_n[None, :],
+                             mask=mask_k[:, None] & mask_n[None, :], other=0)
+                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, :],
+                             mask=mask_k[:, None] & mask_n[None, :], other=0)
+                q = (wv >> (((offs_k % EPS) * NBITS)[:, None])) & ((1 << NBITS) - 1)
+
+            g = k0 // GS
+            s = tl.load(s_ptr + g * N + offs_n, mask=mask_n, other=0.0).to(tl.float32)
+            z = tl.load(z_ptr + g * N + offs_n, mask=mask_n, other=0.0).to(tl.float32)
+            w = (q.to(tl.float32) - z[None, :]) * s[None, :]
+
+            acc += tl.dot(x, w.to(x.dtype), out_dtype=tl.float32)
+
+        out = acc.to(OUT_DTYPE)
+        y_ptrs = y_ptr + offs_m[:, None] * N + offs_n[None, :]
+        if SPLIT_K == 1:
+            tl.store(y_ptrs, out, mask=mask_m[:, None] & mask_n[None, :])
+        else:
+            tl.atomic_add(y_ptrs, out, mask=mask_m[:, None] & mask_n[None, :], 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) 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)
@@ -100,7 +157,7 @@ def _unpack_kmajor(words: torch.Tensor, bits: int, K: int) -> torch.Tensor:
 
 
 class TritonQuantLinear(nn.Module):
-    """Decode-optimized replacement for a DASH-Q PackedQuantizedLinear."""
+    """Linear layer over group-wise asymmetric integer weights."""
 
     def __init__(
         self,
@@ -118,7 +175,7 @@ class TritonQuantLinear(nn.Module):
         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}")
+            raise ValueError(f"Unsupported nbits: {nbits}")
 
         out_features, in_features = W_int.shape
         if in_features % group_size != 0:
@@ -153,8 +210,7 @@ class TritonQuantLinear(nn.Module):
         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)
+        # The GEMV accumulates with atomics, so it starts from the bias.
         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())
@@ -166,7 +222,7 @@ class TritonQuantLinear(nn.Module):
         )
 
     def dequantize_weight(self, dtype: torch.dtype) -> torch.Tensor:
-        """Returns W^T as (in_features, out_features), matching the K-major layout."""
+        """Returns W^T with shape (in_features, out_features)."""
         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)
@@ -177,37 +233,69 @@ class TritonQuantLinear(nn.Module):
         z = self.zero.repeat_interleave(self.group_size, dim=0).to(dtype)
         return (q.to(dtype) - z) * s
 
+    # (BLOCK_M, BLOCK_N, SPLIT_K, num_warps, num_stages), largest tile first;
+    # the first entry that fits in shared memory is cached per shape.
+    _SMALL_M_CONFIGS = ((16, 64, 8, 4, 2), (16, 64, 4, 4, 1))
+    _LARGE_M_CONFIGS = ((128, 128, 1, 8, 4), (128, 128, 1, 8, 3),
+                        (128, 64, 1, 4, 3), (64, 64, 1, 4, 2))
+
+    def _gemm(self, x2d: torch.Tensor) -> torch.Tensor:
+        M = x2d.shape[0]
+        N, K, gs = self.out_features, self.in_features, self.group_size
+        block_k = min(gs, 32)
+        configs = self._SMALL_M_CONFIGS if M <= 32 else self._LARGE_M_CONFIGS
+        cache_key = (M <= 32, N, K, gs, self.nbits)
+        if cache_key in _GEMM_CONFIG_CACHE:
+            configs = (_GEMM_CONFIG_CACHE[cache_key],)
+
+        tl_dtype = tl.float16 if self.out_dtype == torch.float16 else tl.bfloat16
+        last_err = None
+        for cfg in configs:
+            block_m, block_n, split_k, warps, stages = cfg
+            split_k = min(split_k, max(1, K // block_k))
+            alloc = torch.empty if split_k == 1 else torch.zeros
+            y = alloc(M, N, dtype=self.out_dtype, device=x2d.device)
+            grid = (triton.cdiv(M, block_m), triton.cdiv(N, block_n), split_k)
+            try:
+                _dashq_gemm_kernel[grid](
+                    x2d, self.W_q, self.W_lo, self.scale, self.zero, y,
+                    M, N, K,
+                    self.nbits, self.eps, gs,
+                    block_m, block_n, block_k, split_k, tl_dtype,
+                    num_warps=warps, num_stages=stages,
+                )
+            except triton.runtime.errors.OutOfResources as exc:
+                last_err = exc
+                continue
+            _GEMM_CONFIG_CACHE[cache_key] = cfg
+            return y
+        raise last_err
+
     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,
+                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
+        x2d = x.reshape(tokens, -1)
+        if x.is_cuda and TRITON_AVAILABLE:
+            out = self._gemm(x2d)
+        else:
+            out = x2d @ self.dequantize_weight(x.dtype)
         if self.bias is not None:
-            out = out + self.bias.to(x.dtype)
-        return out.reshape(*shape[:-1], self.out_features)
+            out = out + self.bias.to(out.dtype)
+        return out.to(x.dtype).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"
+            f"nbits={self.nbits}, group_size={self.group_size}"
         )

modeling_dashq.py

@@ -1,13 +1,11 @@
-"""Self-contained DASH-Q inference code for this checkpoint.
+"""Inference code for this DASH-Q 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.
+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
 
@@ -38,7 +36,7 @@ def _unpack_int_values(packed: torch.Tensor, nbits: int, numel: int) -> torch.Te
 
 
 class DashQPackedLinear(nn.Module):
-    """Holds the checkpoint buffers; converted to the Triton kernel after load."""
+    """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:
@@ -203,7 +201,7 @@ class DashQQwen3_5MoeForConditionalGeneration(Qwen3_5MoeForConditionalGeneration
         return model
 
     def build_dashq_kernels(self, verbose: bool = True) -> "DashQQwen3_5MoeForConditionalGeneration":
-        """Convert packed buffers to the Triton decode kernel (no-op off CUDA)."""
+        """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):