Skip to content

vllm.model_executor.layers.quantization.online.lm_head

Online-quantized copies of an already-loaded lm_head.

Functions:

_LMHeadNvfp4Method

Bases: OnlineLinearBase

Weight-only NVFP4 lm_head copy on the platform's W4A16 kernel; activations stay in the model dtype.

Source code in vllm/model_executor/layers/quantization/online/lm_head.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class _LMHeadNvfp4Method(OnlineLinearBase):
    """Weight-only NVFP4 lm_head copy on the platform's W4A16 kernel;
    activations stay in the model dtype."""

    def process_weights_after_loading(self, layer: Module) -> None:
        packed, scales, global_scale = _quantize_nvfp4(layer.weight)
        replace_parameter(layer, "weight", packed)
        layer.weight_scale = Parameter(scales, requires_grad=False)
        layer.weight_global_scale = Parameter(global_scale, requires_grad=False)
        self.kernel = init_nvfp4_linear_kernel(use_a16=True)
        self.kernel.process_weights_after_loading(layer)

    def apply(
        self, layer: Module, x: torch.Tensor, bias: torch.Tensor | None = None
    ) -> torch.Tensor:
        return self.kernel.apply_weights(layer, x, bias)

_quantize_nvfp4(weight, chunk_rows=8192)

NVFP4-quantize the rows of weight without the SM100+ quant kernel.

Returns packed E2M1 codes (low nibble first), E4M3 scales per 16 values and the fp32 dequantization global scale, in the checkpoint layout.

Source code in vllm/model_executor/layers/quantization/online/lm_head.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def _quantize_nvfp4(
    weight: torch.Tensor, chunk_rows: int = 8192
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """NVFP4-quantize the rows of `weight` without the SM100+ quant kernel.

    Returns packed E2M1 codes (low nibble first), E4M3 scales per 16 values and
    the fp32 dequantization global scale, in the checkpoint layout.
    """
    fp8_max = torch.finfo(torch.float8_e4m3fn).max
    global_scale = (weight_amax(weight).float() / (FLOAT4_E2M1_MAX * fp8_max)).clamp(
        min=1e-12
    )
    inv_global_scale = (1.0 / global_scale).reshape(())
    # E2M1 magnitudes 0, 0.5, ..., 6 doubled to integers, mapped to their codes.
    code_of = torch.zeros(13, dtype=torch.uint8, device=weight.device)
    code_of[torch.tensor([0, 1, 2, 3, 4, 6, 8, 12])] = torch.arange(
        8, dtype=torch.uint8
    ).to(weight.device)
    n, k = weight.shape
    packed = torch.empty(n, k // 2, dtype=torch.uint8, device=weight.device)
    scales = torch.empty(n, k // 16, dtype=torch.float8_e4m3fn, device=weight.device)
    for start in range(0, n, chunk_rows):
        rows = slice(start, start + chunk_rows)
        values, block_scales = ref_nvfp4_quant(weight[rows], inv_global_scale, 16)
        codes = code_of[(values.abs() * 2).long()] | ((values < 0).to(torch.uint8) << 3)
        packed[rows] = codes[:, 0::2] | (codes[:, 1::2] << 4)
        scales[rows] = block_scales.to(torch.float8_e4m3fn)
    return packed, scales, global_scale

quantized_lm_head_copy(lm_head, quantization)

A ParallelLMHead holding this rank's shard of lm_head, quantized.

"fp8" uses per-row weight and per-token activation scales on the platform's FP8 linear kernel; "nvfp4" is weight-only on the platform's W4A16 kernel. lm_head is left unchanged and no bf16 copy of its weight is made. The copy is derived state: refresh_quantized_lm_heads re-derives it after lm_head changes.

Source code in vllm/model_executor/layers/quantization/online/lm_head.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def quantized_lm_head_copy(
    lm_head: ParallelLMHead, quantization: Literal["fp8", "nvfp4"]
) -> ParallelLMHead:
    """A ParallelLMHead holding this rank's shard of `lm_head`, quantized.

    "fp8" uses per-row weight and per-token activation scales on the platform's
    FP8 linear kernel; "nvfp4" is weight-only on the
    platform's W4A16 kernel. `lm_head` is left
    unchanged and no bf16 copy of its weight is made. The copy is derived state:
    `refresh_quantized_lm_heads` re-derives it after `lm_head` changes.
    """
    head = _quantize_copy(lm_head, quantization)
    # Keep it out of layerwise reload, which would re-run
    # process_weights_after_loading on the already-quantized weight.
    get_layerwise_info(head).reset()
    _SOURCES[head] = (lm_head, quantization)
    logger.info(
        "Quantized a %s copy of the lm_head (%d x %d per rank).",
        quantization,
        head.output_size_per_partition,
        head.input_size_per_partition,
    )
    return head

refresh_quantized_lm_heads(model)

Re-derive the quantized lm_head copies in model from their source heads.

Copies in place, so CUDA graphs keep pointing at valid tensors.

Source code in vllm/model_executor/layers/quantization/online/lm_head.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
@torch.no_grad()
def refresh_quantized_lm_heads(model: Module) -> None:
    """Re-derive the quantized lm_head copies in `model` from their source heads.

    Copies in place, so CUDA graphs keep pointing at valid tensors.
    """
    for module in model.modules():
        if module not in _SOURCES:
            continue
        fresh = _quantize_copy(*_SOURCES[module])
        for name, tensor in chain(
            fresh.named_parameters(recurse=False), fresh.named_buffers(recurse=False)
        ):
            getattr(module, name).copy_(tensor)

restore_quantized_lm_heads(model, refresh=True)

Undo share_source_lm_heads, then re-derive the copies in model.

Parameters:

  • model

    (Module) –

    The model passed to share_source_lm_heads, if any.

  • refresh

    (bool, default: True ) –

    Whether to re-derive the copies; False when the source heads may be mid-reload.

Source code in vllm/model_executor/layers/quantization/online/lm_head.py 
192
193
194
195
196
197
198
199
200
201
202
203
204
def restore_quantized_lm_heads(model: Module, refresh: bool = True) -> None:
    """Undo `share_source_lm_heads`, then re-derive the copies in `model`.

    Args:
        model: The model passed to `share_source_lm_heads`, if any.
        refresh: Whether to re-derive the copies; False when the source heads
            may be mid-reload.

    """
    for module, name, head in _SWAPPED.pop(model, ()):
        setattr(module, name, head)
    if refresh:
        refresh_quantized_lm_heads(model)

share_source_lm_heads(model)

Replace each quantized copy in model with its source head.

model then loads weights like a drafter sharing the unquantized head, until restore_quantized_lm_heads swaps the copies back in.

Source code in vllm/model_executor/layers/quantization/online/lm_head.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def share_source_lm_heads(model: Module) -> None:
    """Replace each quantized copy in `model` with its source head.

    `model` then loads weights like a drafter sharing the unquantized head,
    until `restore_quantized_lm_heads` swaps the copies back in.
    """
    swapped = [
        (module, name, child)
        for module in model.modules()
        for name, child in module.named_children()
        if child in _SOURCES
    ]
    for module, name, child in swapped:
        setattr(module, name, _SOURCES[child][0])
    _SWAPPED[model] = swapped