diff --git a/tests/model_executor/layers/quantization/test_cpu_int4.py b/tests/model_executor/layers/quantization/test_cpu_int4.py new file mode 100644 index 000000000..1899df63e --- /dev/null +++ b/tests/model_executor/layers/quantization/test_cpu_int4.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the cpu_int4 W4A16 method.""" + +import pytest +import torch + +from vllm_omni.model_executor.layers.quantization.cpu_int4 import ( + CPUInt4Config, + _to_tinygemm, + dequantize, + pack_nibbles, + quantize_groupwise, + to_split_halves, + unpack_nibbles, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _weights(n=64, k=256, seed=0): + torch.manual_seed(seed) + return torch.randn(n, k) + + +def test_nibble_pack_round_trips(): + q = torch.randint(0, 16, (32, 128), dtype=torch.int32) + assert torch.equal(unpack_nibbles(pack_nibbles(q)), q) + + +def test_quantized_weights_stay_on_the_four_bit_grid(): + w = _weights() + packed, scale, zero = quantize_groupwise(w, 64) + codes = unpack_nibbles(packed) + assert codes.min() >= 0 and codes.max() <= 15 + assert packed.dtype == torch.uint8 + assert packed.shape == (64, 128) + assert scale.shape == zero.shape == (64, 4) + + +def test_group_extremes_are_reproduced_exactly(): + """min/max grid: the two ends of each group must dequantize back to + themselves, which is what makes the offset form worth its extra bytes.""" + w = _weights(n=16, k=128) + group = 64 + packed, scale, zero = quantize_groupwise(w, group) + codes = unpack_nibbles(packed).float().reshape(16, 128 // group, group) + deq = (codes - 8.0) * scale.unsqueeze(-1) + zero.unsqueeze(-1) + wg = w.reshape(16, 128 // group, group) + torch.testing.assert_close(deq.amax(-1), wg.amax(-1), rtol=1e-5, atol=1e-5) + torch.testing.assert_close(deq.amin(-1), wg.amin(-1), rtol=1e-5, atol=1e-5) + + +def test_split_halves_is_a_relabelling_of_the_same_codes(): + w = _weights() + packed, _, _ = quantize_groupwise(w, 64) + halves = to_split_halves(packed) + codes = unpack_nibbles(packed) + k = codes.shape[1] + assert torch.equal((halves & 0x0F).to(torch.int32), codes[:, : k // 2]) + assert torch.equal((halves >> 4).to(torch.int32), codes[:, k // 2 :]) + + +@pytest.mark.parametrize("group", [32, 64, 128]) +def test_dequantize_path_agrees_with_the_tinygemm_kernel(group): + """The two GEMM paths must be the same function of the weights. + + Prefill takes the dequantize-then-AMX path and decode the tinygemm one; + a mismatch here would show up as a model that generates fine but answers + differently depending on prompt length, which no throughput benchmark + would catch. + """ + w = _weights(n=64, k=256) + packed, scale, zero = quantize_groupwise(w, group) + inner, sz = _to_tinygemm(packed, scale, zero) + halves = to_split_halves(packed) + sc, ze = scale.to(torch.bfloat16), zero.to(torch.bfloat16) + + x = torch.randn(8, 256, dtype=torch.bfloat16) + via_kernel = torch.ops.aten._weight_int4pack_mm_for_cpu(x, inner, group, sz) + via_dequant = torch.nn.functional.linear(x, dequantize(halves, sc, ze, group)) + + # Scale-free: both sides accumulate in bf16, so the gap is bf16 rounding + # on outputs whose magnitude grows with K, not a layout disagreement. + a, b = via_kernel.float(), via_dequant.float() + snr = 20 * torch.log10(a.norm() / (a - b).norm()) + assert snr > 40.0, f'paths disagree beyond bf16 rounding: {snr:.1f} dB' + + + +def test_four_bit_error_is_small_enough_to_be_useful(): + w = _weights(n=256, k=1024) + packed, scale, zero = quantize_groupwise(w, 64) + codes = unpack_nibbles(packed).float().reshape(256, 1024 // 64, 64) + deq = ((codes - 8.0) * scale.unsqueeze(-1) + zero.unsqueeze(-1)).reshape(256, 1024) + snr = 20 * torch.log10(w.norm() / (w - deq).norm()) + assert snr > 19.0, snr + + +def test_config_round_trips_and_reports_bf16_only(): + cfg = CPUInt4Config.from_config( + {"quant_method": "cpu_int4", "group_size": 128, "prefill_dequant": False} + ) + assert cfg.group_size == 128 + assert cfg.prefill_dequant is False + assert cfg.get_supported_act_dtypes() == [torch.bfloat16] + assert cfg.get_name() == "cpu_int4" + + +def test_registered_into_vllms_quantization_registry(): + from vllm.model_executor.layers.quantization import ( + QUANTIZATION_METHODS, + get_quantization_config, + ) + + import vllm_omni.model_executor.layers.quantization # noqa: F401 + + assert "cpu_int4" in QUANTIZATION_METHODS + assert get_quantization_config("cpu_int4") is CPUInt4Config + + +def test_ignore_list_falls_back_to_full_precision(): + cfg = CPUInt4Config(group_size=64, ignore=["lm_head"]) + from vllm.model_executor.layers.linear import ( + ReplicatedLinear, + UnquantizedLinearMethod, + ) + + layer = ReplicatedLinear.__new__(ReplicatedLinear) + assert isinstance(cfg.get_quant_method(layer, "lm_head"), UnquantizedLinearMethod) + assert not isinstance( + cfg.get_quant_method(layer, "model.layers.0.mlp.down_proj"), + UnquantizedLinearMethod, + ) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 3fc39dc28..6b9ae78d4 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -140,6 +140,10 @@ def register_omni_models_to_vllm(): # Register omni-specific reasoning parsers (e.g., step_audio). import vllm_omni.reasoning # noqa: F401 + # Register out-of-tree quantization methods (e.g. cpu_int4, which gives + # the CPU backend the 4-bit weight path vLLM 0.28 does not ship). + import vllm_omni.model_executor.layers.quantization # noqa: F401 + @dataclass class OmniEngineArgs(EngineArgs): diff --git a/vllm_omni/model_executor/layers/quantization/__init__.py b/vllm_omni/model_executor/layers/quantization/__init__.py new file mode 100644 index 000000000..9840a73d4 --- /dev/null +++ b/vllm_omni/model_executor/layers/quantization/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Out-of-tree quantization methods registered into vLLM.""" + +from vllm_omni.model_executor.layers.quantization.cpu_int4 import ( # noqa: F401 + CPUInt4Config, +) + +__all__ = ["CPUInt4Config"] diff --git a/vllm_omni/model_executor/layers/quantization/cpu_int4.py b/vllm_omni/model_executor/layers/quantization/cpu_int4.py new file mode 100644 index 000000000..2b2ee6470 --- /dev/null +++ b/vllm_omni/model_executor/layers/quantization/cpu_int4.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""W4A16 group-wise weights for vLLM's CPU backend. + +Single-stream decode on a CPU is a bandwidth problem: the whole weight matrix +is read to produce one token, so the bits per weight, not the FLOPs, set the +token rate. Four bits is how llama.cpp's Q4_K_M builds win single-stream +decode against an engine that beats them everywhere else. + +vLLM 0.28 does reach four bits on CPU, through the GPTQ/AWQ kernel selector +(`int4_scaled_mm_cpu`, AMX W4A8 -- see `gptq_cpu_export`). That path pins the +zero point at 8, so its grid is symmetric and its activations are int8. This +one is the other trade: an asymmetric min/max grid and bf16 activations, which +costs some speed and buys back fidelity. + +Weights are quantized to 4 bits on groups of ``group_size`` columns with a +bf16 scale and offset per group; at the default group of 64 that is 4.5 bits +per weight, the same budget Q4_K_M spends. The GEMM is PyTorch's ``_weight_int4pack_mm_for_cpu`` tinygemm +kernel, which has AVX-512 paths and, unlike the fp32 entry point, keeps +bf16 activations in bf16 end to end. + +Checkpoints carry, per quantized linear: + + .weight_packed uint8 [N, K // 2] two 4-bit codes per byte, + byte j = w[2j] | w[2j+1] << 4 + .weight_scale bf16 [N, K // group] + .weight_zero bf16 [N, K // group] + +and dequantize as ``w = (q - 8) * scale + zero``. The row-major nibble layout +is the portable one; the tinygemm layout the kernel wants blends rows within a +block of 16 and is built in ``process_weights_after_loading``. + +The tinygemm kernel dequantizes once per row of activations, so its cost is +linear in the batch and it loses to AMX bf16 above about eight rows -- fine +for decode, ruinous for prefill. Large batches therefore take a second path: +the weight is unpacked to bf16 once per call and handed to the ordinary +oneDNN/AMX GEMM, which costs one extra pass over the weights and wins back +prefill. That path reads a row-major copy, so with ``prefill_dequant`` on +(the default) a layer holds both layouts -- 8 bits per weight resident, but +still 4.5 bits per weight *read* in a decode step, which is what sets the +token rate. Set ``prefill_dequant: false`` to trade prefill for the smaller +footprint. + +``dequant_threshold`` defaults to 2, i.e. only a single-row step takes the +tinygemm path. That is not just a speed choice. On a Sapphire Rapids host, +running the tinygemm kernel on a multi-row batch and then vLLM's +``cpu_attention_with_kv_cache`` in the same forward kills the process with +SIGILL inside the attention kernel -- the signature of an AMX instruction +issued with no valid tile configuration loaded. The same prompt is fine in +bf16, and fine at any batch size once every multi-row GEMM goes through the +dequantize path, so the two kernels disagree about who owns the tile state. +Single-row steps never trip it, which is why the crossover sits at 2 rather +than at the ~8 rows where the dequantize path starts winning on speed. +Batched decode therefore pays the dequantize path; single-stream decode, the +case 4-bit weights exist for, does not. +""" + +from typing import Any + +import torch + +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.model_executor.layers.quantization import register_quantization_config +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.parameter import ( + GroupQuantScaleParameter, + PackedvLLMParameter, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.utils.torch_utils import direct_register_custom_op + +# The tinygemm packer refuses anything else. +PACK_N_ALIGN = 16 +INNER_K_TILES = 1 + + +def pack_nibbles(q: torch.Tensor) -> torch.Tensor: + """[N, K] codes in [0, 15] -> [N, K // 2] uint8, low nibble first.""" + q = q.to(torch.uint8) + lo, hi = q[:, 0::2], q[:, 1::2] + return (lo | (hi << 4)).contiguous() + + +def unpack_nibbles(packed: torch.Tensor) -> torch.Tensor: + """[N, K // 2] uint8 -> [N, K] int32 codes in [0, 15].""" + n, half = packed.shape + out = torch.empty(n, half * 2, dtype=torch.int32) + out[:, 0::2] = (packed & 0x0F).to(torch.int32) + out[:, 1::2] = (packed >> 4).to(torch.int32) + return out + + +def quantize_groupwise(w: torch.Tensor, group: int): + """Min/max 4-bit quantization of [N, K] -> (packed, scale, zero). + + ``scale = (max - min) / 15`` and ``zero = min + 8 * scale`` make the + kernel's ``(q - 8) * scale + zero`` reproduce the group's extremes + exactly, which is the same asymmetric grid llama.cpp's K-quants use. + """ + n, k = w.shape + if k % group: + raise ValueError(f"input dim {k} is not a multiple of group size {group}") + wf = w.float().reshape(n, k // group, group) + wmin, wmax = wf.amin(dim=-1), wf.amax(dim=-1) + scale = (wmax - wmin).clamp(min=1e-9) / 15.0 + zero = wmin + 8.0 * scale + q = ((wf - wmin.unsqueeze(-1)) / scale.unsqueeze(-1)).round_().clamp_(0, 15) + return pack_nibbles(q.reshape(n, k)), scale, zero + + +def _to_tinygemm(packed: torch.Tensor, scale: torch.Tensor, zero: torch.Tensor): + """Row-major nibbles + [N, G] scales -> kernel layout + [G, N, 2] scales.""" + codes = unpack_nibbles(packed) + inner = torch.ops.aten._convert_weight_to_int4pack_for_cpu(codes, INNER_K_TILES) + sz = torch.stack([scale.float(), zero.float()], dim=-1) # [N, G, 2] + sz = sz.transpose(0, 1).contiguous().to(torch.bfloat16) # [G, N, 2] + return inner, sz + + +def to_split_halves(packed: torch.Tensor) -> torch.Tensor: + """Interleaved nibbles -> byte j holding w[j] and w[j + K/2]. + + The on-disk layout interleaves so that a tensor-parallel split along K + stays a plain slice. Once the shard is local, splitting the halves makes + both nibble planes contiguous runs, and the dequantize loop drops its + lane shuffle: measured 3.1 ms per layer instead of 4.5. + """ + codes = unpack_nibbles(packed) + n, k = codes.shape + lo, hi = codes[:, : k // 2], codes[:, k // 2 :] + return (lo.to(torch.uint8) | (hi.to(torch.uint8) << 4)).contiguous() + + +def _dequantize_halves( + halves: torch.Tensor, + scale: torch.Tensor, + zero: torch.Tensor, + group: int, +) -> torch.Tensor: + n, half = halves.shape + g2 = half // group + lo = (halves & 0x0F).to(torch.bfloat16).reshape(n, g2, group) + hi = (halves >> 4).to(torch.bfloat16).reshape(n, g2, group) + a = (lo - 8.0) * scale[:, :g2].unsqueeze(-1) + zero[:, :g2].unsqueeze(-1) + b = (hi - 8.0) * scale[:, g2:].unsqueeze(-1) + zero[:, g2:].unsqueeze(-1) + return torch.cat([a.reshape(n, half), b.reshape(n, half)], dim=1) + + +_compiled_dequant = None + + +def dequantize(halves, scale, zero, group): + """bf16 weight from the row-major halves; compiled, so it is one pass.""" + global _compiled_dequant + if _compiled_dequant is None: + # Compiled outside any vLLM graph: this runs inside an opaque custom + # op, so dynamo never traces it and never guards on the batch size. + _compiled_dequant = torch.compile(_dequantize_halves, dynamic=False) + return _compiled_dequant(halves, scale, zero, group) + + +def _cpu_int4_linear( + x: torch.Tensor, + int4pack: torch.Tensor, + scales_and_zeros: torch.Tensor, + halves: torch.Tensor | None, + scale: torch.Tensor | None, + zero: torch.Tensor | None, + group: int, + threshold: int, +) -> torch.Tensor: + if x.dtype != torch.bfloat16: + x = x.to(torch.bfloat16) + if halves is not None and x.shape[0] >= threshold: + weight = dequantize(halves, scale, zero, group) + return torch.nn.functional.linear(x, weight) + return torch.ops.aten._weight_int4pack_mm_for_cpu( + x, int4pack, group, scales_and_zeros + ) + + +def _cpu_int4_linear_fake( + x: torch.Tensor, + int4pack: torch.Tensor, + scales_and_zeros: torch.Tensor, + halves: torch.Tensor | None, + scale: torch.Tensor | None, + zero: torch.Tensor | None, + group: int, + threshold: int, +) -> torch.Tensor: + return x.new_empty((x.shape[0], int4pack.shape[0]), dtype=torch.bfloat16) + + +direct_register_custom_op( + op_name="cpu_int4_linear", + op_func=_cpu_int4_linear, + mutates_args=[], + fake_impl=_cpu_int4_linear_fake, + dispatch_key="CPU", +) + + +@register_quantization_config("cpu_int4") +class CPUInt4Config(QuantizationConfig): + """Group-wise 4-bit weight-only quantization for the CPU backend.""" + + def __init__( + self, + group_size: int = 64, + ignore: list[str] | None = None, + quantize_embedding: bool = True, + prefill_dequant: bool = True, + dequant_threshold: int = 2, + use_custom_op: bool = True, + ) -> None: + super().__init__() + self.group_size = group_size + self.ignore = list(ignore or []) + self.quantize_embedding = quantize_embedding + self.prefill_dequant = prefill_dequant + self.dequant_threshold = dequant_threshold + self.use_custom_op = use_custom_op + + def __repr__(self) -> str: + return ( + f"CPUInt4Config(group_size={self.group_size}, " + f"ignore={self.ignore}, quantize_embedding={self.quantize_embedding}, " + f"prefill_dequant={self.prefill_dequant}, " + f"dequant_threshold={self.dequant_threshold}, " + f"use_custom_op={self.use_custom_op})" + ) + + def get_name(self) -> str: + return "cpu_int4" + + def get_supported_act_dtypes(self) -> list[torch.dtype]: + # The fp32 entry point of the tinygemm kernel has no vectorized path + # and runs ~30x slower, so bf16 is the only activation type offered. + return [torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 0 + + @staticmethod + def get_config_filenames() -> list[str]: + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "CPUInt4Config": + return cls( + group_size=cls.get_from_keys_or(config, ["group_size"], 64), + ignore=cls.get_from_keys_or(config, ["ignore"], []), + quantize_embedding=cls.get_from_keys_or( + config, ["quantize_embedding"], True + ), + prefill_dequant=cls.get_from_keys_or(config, ["prefill_dequant"], True), + dequant_threshold=cls.get_from_keys_or(config, ["dequant_threshold"], 2), + use_custom_op=cls.get_from_keys_or(config, ["use_custom_op"], True), + ) + + def _is_ignored(self, prefix: str) -> bool: + return any(prefix == p or prefix.endswith("." + p) for p in self.ignore) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> QuantizeMethodBase | None: + from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, + ) + + if isinstance(layer, VocabParallelEmbedding): + if not self.quantize_embedding or self._is_ignored(prefix): + return None + return CPUInt4EmbeddingMethod(self) + if isinstance(layer, LinearBase): + if self._is_ignored(prefix): + return UnquantizedLinearMethod() + return CPUInt4LinearMethod(self) + return None + + +def _create_int4_weights( + method: "CPUInt4LinearMethod", + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + params_dtype: torch.dtype, + extra_weight_attrs: dict, +) -> None: + group = method.quant_config.group_size + out_size = sum(output_partition_sizes) + if input_size_per_partition % group: + raise ValueError( + f"input dim {input_size_per_partition} is not a multiple of " + f"group size {group}" + ) + if out_size % PACK_N_ALIGN: + raise ValueError( + f"output dim {out_size} is not a multiple of {PACK_N_ALIGN}, which " + "the tinygemm packer requires" + ) + weight_loader = extra_weight_attrs.pop("weight_loader") + + packed = PackedvLLMParameter( + data=torch.empty(out_size, input_size_per_partition // 2, dtype=torch.uint8), + input_dim=1, + output_dim=0, + packed_dim=1, + packed_factor=2, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_packed", packed) + for name in ("weight_scale", "weight_zero"): + p = GroupQuantScaleParameter( + data=torch.empty( + out_size, input_size_per_partition // group, dtype=torch.bfloat16 + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter(name, p) + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = out_size + set_weight_attrs(packed, extra_weight_attrs) + + +def _finalize_int4( + layer: torch.nn.Module, keep_row_major: bool, prefill_dequant: bool +) -> None: + """Repack into the layouts the two GEMM paths want; once per layer.""" + if getattr(layer, "_cpu_int4_ready", False): + return + packed = layer.weight_packed.data + scale = layer.weight_scale.data + zero = layer.weight_zero.data + inner, sz = _to_tinygemm(packed, scale, zero) + layer.weight_int4pack = torch.nn.Parameter(inner, requires_grad=False) + layer.weight_scales_and_zeros = torch.nn.Parameter(sz, requires_grad=False) + if prefill_dequant: + layer.weight_halves = torch.nn.Parameter( + to_split_halves(packed), requires_grad=False + ) + layer.weight_group_scale = torch.nn.Parameter(scale.clone(), requires_grad=False) + layer.weight_group_zero = torch.nn.Parameter(zero.clone(), requires_grad=False) + else: + layer.weight_halves = None + layer.weight_group_scale = None + layer.weight_group_zero = None + if keep_row_major: + # Only the vocab table needs this: the tinygemm layout interleaves + # rows inside blocks of 16, so a single-token embedding lookup cannot + # address it. One row is a few hundred bytes, so keeping the portable + # copy costs no bandwidth at decode. + layer.weight_packed = torch.nn.Parameter(packed, requires_grad=False) + layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + layer.weight_zero = torch.nn.Parameter(zero, requires_grad=False) + else: + del layer.weight_packed + del layer.weight_scale + del layer.weight_zero + layer._cpu_int4_ready = True + + +def _int4_mm( + layer: torch.nn.Module, + x: torch.Tensor, + group: int, + threshold: int, + custom_op: bool = True, +) -> torch.Tensor: + orig_shape = x.shape + x2d = x.reshape(-1, orig_shape[-1]) + if x2d.dtype != torch.bfloat16: + # Static at trace time: the config advertises bf16 activations only, + # and the fp32 entry point of the tinygemm kernel has no vectorized + # path (~30x slower), so never fall into it by accident. + x2d = x2d.to(torch.bfloat16) + if custom_op: + # Opaque to dynamo, so the batch-size branch never becomes a guard on + # vLLM's single traced graph. Costs ~5 us a call in dispatch. + out = torch.ops.vllm.cpu_int4_linear( + x2d, + layer.weight_int4pack, + layer.weight_scales_and_zeros, + layer.weight_halves, + layer.weight_group_scale, + layer.weight_group_zero, + group, + threshold, + ) + elif layer.weight_halves is not None and x2d.shape[0] >= threshold: + out = torch.nn.functional.linear( + x2d, + _dequantize_halves( + layer.weight_halves, + layer.weight_group_scale, + layer.weight_group_zero, + group, + ), + ) + else: + out = torch.ops.aten._weight_int4pack_mm_for_cpu( + x2d, layer.weight_int4pack, group, layer.weight_scales_and_zeros + ) + return out.reshape(*orig_shape[:-1], out.shape[-1]) + + +class CPUInt4LinearMethod(QuantizeMethodBase): + """Group-wise 4-bit linear.""" + + def __init__(self, quant_config: CPUInt4Config) -> None: + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + _create_int4_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + params_dtype, + extra_weight_attrs, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + _finalize_int4( + layer, + keep_row_major=False, + prefill_dequant=self.quant_config.prefill_dequant, + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + out = _int4_mm( + layer, + x, + self.quant_config.group_size, + self.quant_config.dequant_threshold, + self.quant_config.use_custom_op, + ) + if bias is not None: + out = out + bias + return out + + +class CPUInt4EmbeddingMethod(CPUInt4LinearMethod): + """The same weights, used as a vocab table and as the output head. + + With tied embeddings these are one tensor, so quantizing it pays twice: + the lm_head GEMM is the single largest read in a decode step, and the + table itself is 16% of the model's bytes. + """ + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + _create_int4_weights( + self, + layer, + input_size_per_partition, + output_partition_sizes, + params_dtype, + extra_weight_attrs, + ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # No dequantize path for the vocab table: the head only ever sees the + # tokens being sampled (one row per sequence), and unpacking a + # 131072 x 2048 matrix to bf16 would cost more than the GEMM saves. + _finalize_int4(layer, keep_row_major=True, prefill_dequant=False) + + def tie_weights(self, layer: torch.nn.Module, embed_tokens: torch.nn.Module): + """Share every quantized tensor, not just ``weight``.""" + for name in ( + "weight_packed", + "weight_scale", + "weight_zero", + "weight_int4pack", + "weight_scales_and_zeros", + "weight_halves", + "weight_group_scale", + "weight_group_zero", + ): + src = getattr(embed_tokens, name, None) + if src is not None: + setattr(layer, name, src) + layer._cpu_int4_tied_to = embed_tokens + return layer + + def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: + group = self.quant_config.group_size + rows = layer.weight_packed.data.index_select(0, input_.reshape(-1)) + codes = unpack_nibbles(rows).float() + n, k = codes.shape + scale = layer.weight_scale.data.index_select(0, input_.reshape(-1)).float() + zero = layer.weight_zero.data.index_select(0, input_.reshape(-1)).float() + deq = (codes.reshape(n, k // group, group) - 8.0) * scale.unsqueeze( + -1 + ) + zero.unsqueeze(-1) + out = deq.reshape(n, k).to(torch.bfloat16) + return out.reshape(*input_.shape, k)