Skip to content

[not ready for review] extend fp8 online quant with blockwise scaling #32485

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vkuzo wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[not ready for review] extend fp8 online quant with blockwise scaling #32485

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions tests/quantization/test_fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,9 +126,21 @@ def check_model(model):
not is_quant_method_supported("fp8"),
reason="FP8 is not supported on this GPU type.",
)
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
# @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
@pytest.mark.parametrize(
"force_marlin", [False] if current_platform.is_rocm() else [False, True]
"kv_cache_dtype",
[
"auto",
],
)
@pytest.mark.parametrize(
# "force_marlin", [False] if current_platform.is_rocm() else [False, True]
"force_marlin",
[False]
if current_platform.is_rocm()
else [
False,
],
)
@pytest.mark.parametrize(
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
Expand All @@ -150,7 +162,8 @@ def test_load_fp16_model(
monkeypatch.setenv("VLLM_TEST_FORCE_FP8_MARLIN", "1")

with vllm_runner(
"facebook/opt-125m",
# "facebook/opt-125m",
"Qwen/Qwen1.5-MoE-A2.7B",
Copy link

@cursor cursor bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Test validation disabled with debug model change

Medium Severity

The test model was changed from facebook/opt-125m to Qwen/Qwen1.5-MoE-A2.7B, but the check_model validation function (which references opt-125m-specific layer paths like model.model.decoder.layers[0].fc1) was commented out instead of updated. The test now only runs inference without validating quantization was applied correctly. Additionally, test parameterization was reduced, decreasing coverage.

Additional Locations (2)

Fix in Cursor Fix in Web

quantization="fp8",
enforce_eager=True,
kv_cache_dtype=kv_cache_dtype,
Expand Down Expand Up @@ -189,7 +202,10 @@ def check_model(model):
"It only runs on CUDA and ROCm platform."
)

llm.apply_model(check_model)
# below currently hardcodes opt-125m layers, skip for now
# llm.apply_model(check_model)
outputs = llm.generate_greedy(["Hello my name is"], max_tokens=20)
print(outputs[0][1])


@pytest.mark.skipif(
Expand Down
197 changes: 163 additions & 34 deletions vllm/model_executor/layers/quantization/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from enum import Enum
from typing import TYPE_CHECKING, Any, Optional

import torch
Expand Down Expand Up @@ -103,6 +104,13 @@
logger = init_logger(__name__)


class OnlineQuantScalingType(Enum):
# TODO(before land): align on naming and add descriptive comments
# to each enum value
TENSORWISE = "tensorwise"
BLOCKWISE = "blockwise"


class Fp8Config(QuantizationConfig):
"""Config class for FP8."""

Expand Down Expand Up @@ -140,6 +148,11 @@
)
self.weight_block_size = weight_block_size

# TODO(before land): hook this up to user UI, for now hardcode it here
self.online_quant_scaling_type = OnlineQuantScalingType.BLOCKWISE
# self.online_quant_scaling_type = OnlineQuantScalingType.TENSORWISE
self.online_block_size = [128, 128] # [block_n, block_k]

@classmethod
def get_name(cls) -> QuantizationMethods:
return "fp8"
Expand Down Expand Up @@ -328,20 +341,33 @@
self.weight_block_size = self.quant_config.weight_block_size
self.block_quant = self.weight_block_size is not None
self.act_q_static = self.quant_config.activation_scheme == "static"

if self.weight_block_size:
self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])
elif (
self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
self.act_q_group_shape = GroupShape(
1, self.quant_config.online_block_size[0]
)
else:
# Use per-token quantization for better perf if dynamic and cutlass
if not self.act_q_static and cutlass_fp8_supported():
self.act_q_group_shape = GroupShape.PER_TOKEN
else:
self.act_q_group_shape = GroupShape.PER_TENSOR

if self.block_quant:
if (
self.block_quant
or self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
Comment on lines +361 to +365
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The condition self.block_quant or self.quant_config.online_quant_scaling_type is OnlineQuantScalingType.BLOCKWISE is repeated in several places within this class (e.g., in process_weights_after_loading and apply). To improve maintainability and reduce code duplication, consider creating a helper property within the Fp8LinearMethod class to encapsulate this logic. For example:

@property
def _is_blockwise_quant(self):
    return (self.block_quant or
            self.quant_config.online_quant_scaling_type is
            OnlineQuantScalingType.BLOCKWISE)

You could then use if self._is_blockwise_quant: in this and other locations.

block_size = self.weight_block_size or self.quant_config.online_block_size
assert block_size is not None
assert not self.act_q_static
assert self.weight_block_size is not None
self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
weight_group_shape=GroupShape(*self.weight_block_size),
weight_group_shape=GroupShape(*block_size),
act_quant_group_shape=self.act_q_group_shape,
cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
use_aiter_and_is_supported=self.use_aiter_and_is_supported,
Expand Down Expand Up @@ -487,8 +513,38 @@
# If checkpoint not serialized fp8, quantize the weights.
else:
if not self.quant_config.is_checkpoint_fp8_serialized:
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
weight = qweight.t()
# Online quantization
if (
self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
# blockwise
from vllm.utils.deep_gemm import per_block_cast_to_fp8
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Local imports should generally be at the top of the file. Please move from vllm.utils.deep_gemm import per_block_cast_to_fp8 to the top to follow standard Python style guidelines and improve maintainability.


block_size = self.quant_config.online_block_size
# layer.weight is [N, K] where N=output_size, K=input_size
qweight, weight_scale_inv = per_block_cast_to_fp8(
layer.weight, block_size=block_size
)
# qweight: [N, K] in FP8
# weight_scale_inv: [N/block_n, K/block_k] - inverse scales
# Note: block ops expect [N, K] format (no transpose)
replace_parameter(layer, "weight", qweight.data)
replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
layer.weight_block_size = block_size
size_k_first = False
else:
# tensorwise
assert (
self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.TENSORWISE
)
qweight, weight_scale = ops.scaled_fp8_quant(
layer.weight, scale=None
)
weight = qweight.t()
replace_parameter(layer, "weight", weight.data)
replace_parameter(layer, "weight_scale", weight_scale.data)

# If checkpoint is fp8 per-tensor, handle that there are N scales for N
# shards in a fused module
Expand All @@ -512,9 +568,9 @@
input_scale = input_scale.max()
weight = weight.t()

# Update layer with new values.
replace_parameter(layer, "weight", weight.data)
replace_parameter(layer, "weight_scale", weight_scale.data)
# Update layer with new values.
replace_parameter(layer, "weight", weight.data)
replace_parameter(layer, "weight_scale", weight_scale.data)

if input_scale is not None:
replace_parameter(layer, "input_scale", input_scale)
Expand All @@ -529,7 +585,11 @@
del layer.input_scale
return

if self.block_quant:
if (
self.block_quant
or self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
maybe_post_process_fp8_weight_block(layer)

def apply(
Expand All @@ -541,8 +601,11 @@
# if batch invariant mode is enabled, prefer DeepGEMM FP8 path
# we will use BF16 dequant when DeepGEMM is not supported.
if vllm_is_batch_invariant():
if self.block_quant:
assert self.weight_block_size is not None
if (
self.block_quant
or self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
Copy link

@cursor cursor bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blockwise condition breaks serialized per-tensor FP8 checkpoints

High Severity

The conditions checking online_quant_scaling_type is BLOCKWISE in the non-Marlin apply paths don't account for serialized checkpoints. For pre-quantized per-tensor FP8 checkpoints, create_weights creates weight_scale (not weight_scale_inv), but the hardcoded BLOCKWISE setting causes both the batch-invariant path and the main w8a8_block_fp8_linear.apply path to access layer.weight_scale_inv, which doesn't exist. This causes AttributeError when loading any serialized per-tensor FP8 checkpoint on FP8-capable hardware.

Additional Locations (1)

Fix in Cursor Fix in Web

return self.w8a8_block_fp8_linear.apply(
input=x,
weight=layer.weight,
Expand Down Expand Up @@ -591,9 +654,11 @@
bias=bias,
)

if self.block_quant:
assert self.weight_block_size is not None

if (
Copy link

@cursor cursor bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Marlin path accesses non-existent weight scale attribute

High Severity

The Marlin code path in apply checks only self.block_quant to decide between layer.weight_scale_inv and layer.weight_scale, but block_quant is False when doing online quantization (it only reflects checkpoint-provided weight_block_size). Since online_quant_scaling_type is hardcoded to BLOCKWISE, process_weights_after_loading creates weight_scale_inv, but the Marlin path tries to access weight_scale which doesn't exist, causing an AttributeError on GPUs without FP8 hardware support.

Additional Locations (1)

Fix in Cursor Fix in Web

self.block_quant
or self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
return self.w8a8_block_fp8_linear.apply(
input=x,
weight=layer.weight,
Expand Down Expand Up @@ -1089,7 +1154,21 @@
super().__init__(quant_config, layer)
assert not quant_config.is_checkpoint_fp8_serialized
assert quant_config.activation_scheme == "dynamic"
assert quant_config.weight_block_size is None

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The assertion assert quant_config.weight_block_size is None was removed. Fp8OnlineMoEMethod is used for online quantization (is_checkpoint_fp8_serialized=False), where weight_block_size in Fp8Config is expected to be None. This assertion is a crucial sanity check to prevent misconfiguration. Please restore it.

Suggested change
assert self.quant_config.weight_block_size is None

# Override parent class attributes for online blockwise quantization
if (
self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
self.weight_block_size = self.quant_config.online_block_size
self.block_quant = True
self.weight_scale_name = "weight_scale_inv"
# Re-select backend with correct block_quant flag
self.fp8_backend = select_fp8_moe_backend(
block_quant=self.block_quant,

Check failure on line 1168 in vllm/model_executor/layers/quantization/fp8.py

View workflow job for this annotation

GitHub Actions / pre-commit 

Missing positional argument "is_act_and_mul" in call to "select_fp8_moe_backend"  [call-arg]

Check failure on line 1168 in vllm/model_executor/layers/quantization/fp8.py

View workflow job for this annotation

GitHub Actions / pre-commit 

Missing positional argument "is_act_and_mul" in call to "select_fp8_moe_backend"  [call-arg]

Check failure on line 1168 in vllm/model_executor/layers/quantization/fp8.py

View workflow job for this annotation

GitHub Actions / pre-commit 

Missing positional argument "is_act_and_mul" in call to "select_fp8_moe_backend"  [call-arg]

Check failure on line 1168 in vllm/model_executor/layers/quantization/fp8.py

View workflow job for this annotation

GitHub Actions / pre-commit 

Missing positional argument "is_act_and_mul" in call to "select_fp8_moe_backend"  [call-arg]
tp_size=layer.moe_parallel_config.tp_size,
with_lora_support=self.moe.is_lora_enabled,
)
Copy link

@cursor cursor bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parent validation rejects online blockwise before child overrides

Medium Severity

Fp8OnlineMoEMethod.__init__ calls super().__init__() before setting block_quant=True for online blockwise quantization. The parent class validates using its initial block_quant=False and activation_scheme="dynamic", computing dynamic_per_token=True. On SM90/SM100 GPUs with FlashInfer enabled, the parent selects a FlashInfer backend and raises NotImplementedError about "dynamic per token activation quantization" before the child can override block_quant=True. This causes online blockwise MoE quantization to fail on H100/Blackwell GPUs with a misleading error.

Additional Locations (1)

Fix in Cursor Fix in Web


def create_weights(
self,
Expand Down Expand Up @@ -1168,16 +1247,43 @@
set_weight_attrs(w2_weight, extra_weight_attrs)

# WEIGHT_SCALES
# Allocate 2 scales for w1 and w3 respectively.
# They will be combined to a single scale after weight loading.
w13_weight_scale = torch.nn.Parameter(
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
)
w2_weight_scale = torch.nn.Parameter(
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
if (
self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
# For blockwise, scales are per block (typically 128x128)
block_size = self.quant_config.online_block_size
block_n, block_k = block_size[0], block_size[1]
w13_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
2 * ((intermediate_size_per_partition + block_n - 1) // block_n),
(hidden_size + block_k - 1) // block_k,
dtype=torch.float32,
),
requires_grad=False,
)
w2_weight_scale = torch.nn.Parameter(
torch.ones(
num_experts,
(hidden_size + block_n - 1) // block_n,
(intermediate_size_per_partition + block_k - 1) // block_k,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
else:
# For tensorwise, scales are per expert
w13_weight_scale = torch.nn.Parameter(
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
)
w2_weight_scale = torch.nn.Parameter(
torch.ones(num_experts, dtype=torch.float32), requires_grad=False
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)

Expand All @@ -1192,16 +1298,39 @@
fp8_dtype = current_platform.fp8_dtype()
w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
w13_scale = layer.w13_weight_scale
w2_scale = layer.w2_weight_scale

for expert in range(layer.local_num_experts):
w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
layer.w13_weight[expert, :, :]
)
w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
layer.w2_weight[expert, :, :]
)
if (
self.quant_config.online_quant_scaling_type
is OnlineQuantScalingType.BLOCKWISE
):
# Blockwise quantization
from vllm.utils.deep_gemm import per_block_cast_to_fp8
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This local import should be moved to the top of the file. This improves code readability and adheres to standard Python style guides.


block_size = self.quant_config.online_block_size
w13_scale = layer.w13_weight_scale_inv
w2_scale = layer.w2_weight_scale_inv

for expert in range(layer.local_num_experts):
w13[expert, :, :], w13_scale[expert, :, :] = per_block_cast_to_fp8(
layer.w13_weight[expert, :, :], block_size=block_size
)
w2[expert, :, :], w2_scale[expert, :, :] = per_block_cast_to_fp8(
layer.w2_weight[expert, :, :], block_size=block_size
)

layer.weight_block_size = block_size
else:
# Tensorwise quantization
w13_scale = layer.w13_weight_scale
w2_scale = layer.w2_weight_scale

for expert in range(layer.local_num_experts):
w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
layer.w13_weight[expert, :, :]
)
w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
layer.w2_weight[expert, :, :]
)

# Shuffle weights to runtime format and setup kernel.
self._setup_kernel(
Expand Down
Loading