1414``Int4Tensor`` weights keep falling back to torchao's default (mslk/tinygemm)
1515path.
1616
17+ Metadata encoding: the SCALE is a per-group **uint8 code** with a
18+ **per-256-super-block fp16 step** (``scale = code * step``); the ZERO now uses
19+ the SAME per-256-super-block fp16 step (per-group uint8 code + per-256 fp16
20+ step). group_size is 32, so a 256-weight super-block spans 8 groups and there
21+ are ``K/256`` scale steps and ``K/256`` zero steps per row. Dequant is
22+ unchanged: ``w = (q - zero) * scale`` with ``zero = min/scale``.
23+
24+ This mirrors GGUF Q4_K's per-super-block fp16 ``d`` for both the scale and the
25+ zero: the finer per-256 step (vs the previous per-row step) is what lifts
26+ whole-weight dequant SNR to ~45.89 dB (vs 45.15 dB for the per-row-zero-step
27+ encoding, +0.74 dB). The scale/zero codes stay single coalesced uint8s (one
28+ byte/group) so the decode kernel reads exactly one scale byte and one zero byte
29+ per group — no bit-plane reconstruct — and both per-256 steps are loaded once
30+ per super-block. Both steps MUST be fp16 (bf16 for the per-256 step costs ~0.05
31+ dB on the scale; the per-256 zero step is +0.74 dB at fp16 vs +0.52 at bf16).
32+
1733Layout difference from torchao ``Int4Tensor``:
1834 qdata : packed int4 weight (N, K/2), nibble-packed (same as Int4Tensor)
19- scale : (N, n_groups) — the *coalesced* layout, transposed from
20- torchao's documented (n_groups, N)
21- zero_point : (N, n_groups) — coalesced, transposed from (n_groups, N)
35+ scale : (N, n_groups) uint8 — per-group scale *codes*, coalesced
36+ (transposed from torchao's (n_groups, N))
37+ scale_step : (N, K/256) fp16 — per-256-super-block scale step; the real
38+ per-group scale is ``scale_code * scale_step[:, g // 8]``.
39+ zero_point : (N, n_groups) uint8 — per-group zero codes
40+ zero_point_step : (N, K/256) fp16 — per-256-super-block zero step; the real
41+ per-group zero is ``zero_code * zero_point_step[:, g // 8]``.
42+
43+ Bits-per-weight: 4.0 (qdata) + 8/32 (scale codes) + 16/256 (fp16 scale step) +
44+ 8/32 (uint8 zero codes) + 16/256 (fp16 zero step) = 4.625 bpw.
2245
2346The coalesced [N, n_groups] layout is exactly what the W4A8 dp4a matvec kernel
2447(``executorch_cuda::int4_plain_mm`` / ``int4_plain_mm.cuh``) reads row-for-row
2548with qdata, so the exported decode graph carries no per-step transpose. The
26- transpose is owned by :meth:`from_int4_tensor` so it is baked into the
27- serialized weight constant once at pack time.
49+ transpose (and the uint8 re-encoding) is owned by :meth:`from_int4_tensor` so it
50+ is baked into the serialized weight constant once at pack time.
2851"""
2952
30- from typing import List , Optional
53+ from typing import List , Optional , Tuple
3154
3255import torch
3356from torchao .quantization .quantize_ .workflows .int4 .int4_tensor import Int4Tensor
3760 "CudaCoalescedInt4Tensor" ,
3861]
3962
63+ _CODE_MAX = 255 # uint8 code range [0, 255] (both scale and zero)
64+ _SUPER_BLOCK = 256 # weights per super-block (GGUF Q4_K QK_K); scale step is per this
65+
66+
67+ def _encode_uint8_per_super (
68+ x : torch .Tensor ,
69+ group_size : int ,
70+ ) -> Tuple [torch .Tensor , torch .Tensor ]:
71+ """Encode a (n_groups, N) non-negative tensor to per-super-block uint8 codes.
72+
73+ Used for both the scale and the zero. A super-block is ``_SUPER_BLOCK``
74+ (256) weights = ``groups_per_super = 256 // group_size`` groups (8 for
75+ group_size=32). Returns ``(codes, step)`` where ``codes`` is
76+ ``(N, n_groups)`` uint8 (transposed to the coalesced layout) and ``step`` is
77+ ``(N, n_super)`` fp16 with ``n_super = n_groups // groups_per_super = K //
78+ 256``, such that ``code * step[:, g // groups_per_super] ≈ x.t()``. The step
79+ is the per-256-super-block max / 255 rounded to fp16. Rounding uses the
80+ fp16-rounded step (what the kernel reads) so encode and decode agree.
81+ """
82+ xt = x .t ().contiguous ().float () # (N, n_groups)
83+ N , n_groups = int (xt .shape [0 ]), int (xt .shape [1 ])
84+ groups_per_super = _SUPER_BLOCK // int (group_size )
85+ if groups_per_super < 1 :
86+ raise ValueError (
87+ f"group_size={ group_size } { _SUPER_BLOCK }
88+ "scale step"
89+ )
90+ if n_groups % groups_per_super != 0 :
91+ raise ValueError (
92+ f"n_groups={ n_groups } { groups_per_super }
93+ f"(K must be a multiple of { _SUPER_BLOCK } { group_size }
94+ )
95+ n_super = n_groups // groups_per_super
96+ xb = xt .reshape (N , n_super , groups_per_super ) # (N, n_super, gps)
97+ block_max = xb .amax (dim = 2 , keepdim = True ).clamp_min (1e-30 ) # (N, n_super, 1)
98+ step = (block_max / _CODE_MAX ).to (torch .float16 ) # (N, n_super, 1) fp16
99+ step_f = step .float ().clamp_min (1e-30 )
100+ codes = torch .round (xb / step_f ).clamp_ (0 , _CODE_MAX ).to (torch .uint8 )
101+ codes = codes .reshape (N , n_groups ).contiguous ()
102+ return codes , step .squeeze (2 ).contiguous ()
103+
104+
105+ def _unpack_nibble_qdata (qdata : torch .Tensor , N : int , K : int ) -> torch .Tensor :
106+ """Unpack nibble-packed int4 qdata ``(N, K/2)`` -> ``(N, K)`` uint8 [0, 15]."""
107+ qu = qdata .to (torch .uint8 )
108+ even = qu & 0xF
109+ odd = (qu >> 4 ) & 0xF
110+ return torch .stack ([even , odd ], dim = - 1 ).reshape (N , K )
111+
40112
41113class CudaCoalescedInt4Tensor (TorchAOBaseTensor ):
42- """INT4 weight with scale/zero_point in the coalesced [N, n_groups] layout .
114+ """INT4 weight, uint8 scale + per-256 fp16 step, uint8 zero + per-256 fp16 step .
43115
44116 ExecuTorch-internal; see the module docstring. Mirrors torchao
45117 ``Int4Tensor``'s data/attribute layout (so the common tensor utilities and
46- serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose
47- of scale/zero_point via :meth:`from_int4_tensor`.
118+ serialization work) but owns the [n_groups, N] -> [N, n_groups] transpose and
119+ the uint8 re-encoding via :meth:`from_int4_tensor`.
48120 """
49121
50- tensor_data_names = ["qdata" , "scale" , "zero_point" ]
122+ tensor_data_names = [
123+ "qdata" ,
124+ "scale" ,
125+ "scale_step" ,
126+ "zero_point" ,
127+ "zero_point_step" ,
128+ ]
51129 tensor_attribute_names = ["block_size" , "shape" ]
52130 optional_tensor_data_names = ["act_pre_scale" ]
53131 optional_tensor_attribute_names = ["activation_dtype" ]
@@ -56,23 +134,29 @@ def __new__(
56134 cls ,
57135 qdata : torch .Tensor ,
58136 scale : torch .Tensor ,
137+ scale_step : torch .Tensor ,
59138 zero_point : torch .Tensor ,
139+ zero_point_step : torch .Tensor ,
60140 block_size : List [int ],
61141 shape : torch .Size ,
62142 act_pre_scale : Optional [torch .Tensor ] = None ,
63143 activation_dtype : Optional [torch .dtype ] = None ,
64144 ):
65145 kwargs = {}
66146 kwargs ["device" ] = qdata .device
67- kwargs ["dtype" ] = scale .dtype
147+ kwargs ["dtype" ] = (
148+ activation_dtype if activation_dtype is not None else torch .bfloat16
149+ )
68150 kwargs ["requires_grad" ] = False
69151 return torch .Tensor ._make_wrapper_subclass (cls , shape , ** kwargs ) # type: ignore[attr-defined]
70152
71153 def __init__ (
72154 self ,
73155 qdata : torch .Tensor ,
74156 scale : torch .Tensor ,
157+ scale_step : torch .Tensor ,
75158 zero_point : torch .Tensor ,
159+ zero_point_step : torch .Tensor ,
76160 block_size : List [int ],
77161 shape : torch .Size ,
78162 act_pre_scale : Optional [torch .Tensor ] = None ,
@@ -81,7 +165,9 @@ def __init__(
81165 super ().__init__ ()
82166 self .qdata = qdata
83167 self .scale = scale
168+ self .scale_step = scale_step
84169 self .zero_point = zero_point
170+ self .zero_point_step = zero_point_step
85171 self .block_size = block_size
86172 self .activation_dtype = (
87173 activation_dtype if activation_dtype is not None else torch .bfloat16
@@ -98,21 +184,57 @@ def _quantization_type(self):
98184 def from_int4_tensor (cls , t : Int4Tensor ) -> "CudaCoalescedInt4Tensor" :
99185 """Build a coalesced tensor from a torchao ``Int4Tensor``.
100186
101- Owns the transpose: torchao stores scale/zero_point as (n_groups, N);
102- the CUDA decode kernel reads (N, n_groups). The ``.t().contiguous()``
103- here is baked into the serialized weight constant so the exported
104- decode graph has no per-step transpose/clone.
187+ Owns the transpose AND the uint8 re-encoding: torchao stores
188+ scale/zero_point as (n_groups, N) bf16. The CUDA decode kernel reads the
189+ (N, n_groups) uint8 scale/zero *codes* plus per-256-super-block
190+ (N, K/256) fp16 ``scale_step`` / ``zero_point_step`` (scale = scale_code *
191+ scale_step[:, g//8], zero = zero_code * zero_point_step[:, g//8]). The
192+ transpose + encode here is baked into the serialized weight constant so
193+ the exported decode graph has no per-step transpose/clone.
105194 """
195+ scale_codes , scale_step = _encode_uint8_per_super (t .scale , t .block_size [- 1 ])
196+ zero_codes , zero_point_step = _encode_uint8_per_super (
197+ t .zero_point , t .block_size [- 1 ]
198+ )
106199 return cls (
107200 t .qdata ,
108- t .scale .t ().contiguous (),
109- t .zero_point .t ().contiguous (),
201+ scale_codes ,
202+ scale_step ,
203+ zero_codes ,
204+ zero_point_step ,
110205 t .block_size ,
111206 t .shape ,
112207 t .act_pre_scale ,
113208 t .activation_dtype ,
114209 )
115210
211+ def dequantize (self , output_dtype : Optional [torch .dtype ] = None ) -> torch .Tensor :
212+ """Dequantize to a dense tensor: ``w = (q - zero) * scale``.
213+
214+ Reconstructs the per-group scale from the uint8 codes and the per-256
215+ fp16 step, and the per-group zero from the uint8 codes and the per-256
216+ fp16 step. Used as the numerical reference and for the tied lm_head /
217+ token embedding.
218+ """
219+ dtype = output_dtype if output_dtype is not None else torch .bfloat16
220+ N , K = int (self .shape [0 ]), int (self .shape [1 ])
221+ gs = self .block_size [- 1 ]
222+ n_groups = K // gs
223+ n_super = int (self .scale_step .shape [1 ])
224+ groups_per_super = n_groups // n_super
225+
226+ q = _unpack_nibble_qdata (self .qdata , N , K ).to (torch .float32 )
227+ scale_code = self .scale .to (torch .float32 ) # (N, n_groups)
228+ scale_step = self .scale_step .float ().repeat_interleave (groups_per_super , dim = 1 )
229+ scale = (scale_code * scale_step ).repeat_interleave (gs , dim = 1 ) # (N, K)
230+
231+ zero_code = self .zero_point .to (torch .float32 ) # (N, n_groups)
232+ zero_point_step = self .zero_point_step .float ().repeat_interleave (
233+ groups_per_super , dim = 1
234+ )
235+ zero = (zero_code * zero_point_step ).repeat_interleave (gs , dim = 1 ) # (N, K)
236+ return ((q - zero ) * scale ).to (dtype )
237+
116238
117239# Allow a model with CudaCoalescedInt4Tensor weights to be loaded with
118240# `weights_only=True` (mirrors torchao Int4Tensor).
0 commit comments