| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """A QuantizedActivation is a pre-quantized activation produced by a fused kernel |
| 4 | and consumed directly by a linear layer, letting the layer skip its own input |
| 5 | quantization. A linear advertises the key its kernel can consume via |
| 6 | expose_input_quant_key; the kernel validates and reads the activation via |
| 7 | as_quantized_activation. Producers query the effective capability through |
| 8 | get_input_quant_key, which hides the key when another consumer branch needs the |
| 9 | original activation. |
| 10 | """ |
| 11 | |
| 12 | from dataclasses import dataclass, replace |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class QuantizedActivation: |
| 21 | """A quantized activation paired with its scale and original metadata. |
| 22 | |
| 23 | The quant_key describes how data and scale are to be interpreted (dtype, |
| 24 | scale granularity, value packing). Details the key does not capture, such |
| 25 | as blockscale layout or activation padding, must follow the consumer |
| 26 | kernel's convention. |
| 27 | |
| 28 | TODO(mgoin): Encode layout and padding requirements in the contract so |
| 29 | producers can match consumer kernels without relying on convention. |
| 30 | """ |
| 31 | |
| 32 | data: torch.Tensor |
| 33 | scale: torch.Tensor |
| 34 | orig_dtype: torch.dtype |
| 35 | orig_shape: torch.Size |
| 36 | quant_key: QuantKey |
| 37 | |
| 38 | def weak_ref(self) -> "QuantizedActivation": |
| 39 | """Return a copy with non-owning tensor references for CUDA graph replay.""" |
| 40 | from vllm.utils.torch_utils import weak_ref_tensor |
| 41 | |
| 42 | return replace( |
| 43 | self, |
| 44 | data=weak_ref_tensor(self.data), |
| 45 | scale=weak_ref_tensor(self.scale), |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | def expose_input_quant_key(layer: torch.nn.Module, kernel) -> None: |
| 50 | """Store the kernel's pre-quantized input key on the layer, if any. |
| 51 | |
| 52 | This is the bridge from a kernel's input_quant_key() to the layer capability |
| 53 | that fusion call sites read through get_input_quant_key. The raw key is left |
| 54 | unset when the kernel quantizes its own input. |
| 55 | |
| 56 | TODO(mgoin): Producers also need the consumer's quantization scales (e.g. |
| 57 | static input scale, global scale). Expose those here as well so producers |
| 58 | do not reach into kernel-specific layer attributes. |
| 59 | """ |
| 60 | key = kernel.input_quant_key() |
| 61 | if key is not None: |
| 62 | layer._input_quant_key = key |
| 63 | |
| 64 | |
| 65 | def get_input_quant_key(layer: torch.nn.Module) -> QuantKey | None: |
| 66 | """Return the pre-quantized input key the complete layer can consume.""" |
| 67 | if getattr(layer, "requires_unquantized_input", False): |
| 68 | return None |
| 69 | return getattr(layer, "_input_quant_key", None) |
| 70 | |
| 71 | |
| 72 | def as_quantized_activation( |
| 73 | x: "torch.Tensor | QuantizedActivation", expected_key: QuantKey | None |
| 74 | ) -> "QuantizedActivation | None": |
| 75 | """Validate and narrow a pre-quantized activation for a consumer kernel. |
| 76 | |
| 77 | Returns the QuantizedActivation when x is one whose key matches the |
| 78 | kernel's declared expected_key, and None when x is a plain tensor (the |
| 79 | caller quantizes in-kernel). Raises on a key mismatch so a wrongly routed |
| 80 | activation fails loudly instead of being silently re-quantized. |
| 81 | """ |
| 82 | if not isinstance(x, QuantizedActivation): |
| 83 | return None |
| 84 | assert x.quant_key == expected_key, ( |
| 85 | f"QuantizedActivation key {x.quant_key} != consumer kernel " |
| 86 | f"input_quant_key {expected_key}" |
| 87 | ) |
| 88 | return x |