A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 03:45 UTC
[Feature]: Support NVIDIA ModelOpt HF FP8 variants FP8_PER_CHANNEL_PER_TOKEN and FP8_PB_WO in vLLM (#30957)
CedricHuang committed
Dec 22, 2025 at 11:34 UTC
19cc9468fd0fa1701e7cb74b5928b329a1d16cf1
5 files changed
+437
-15
docs/features/quantization/modelopt.md
+31
@@ -8,6 +8,16 @@ We recommend installing the library with:
8
pip install nvidia-modelopt
9
```
10
11
+## Supported ModelOpt checkpoint formats
12
+
13
+vLLM detects ModelOpt checkpoints via `hf_quant_config.json` and supports the
14
+following `quantization.quant_algo` values:
15
+
16
+- `FP8`: per-tensor weight scale (+ optional static activation scale).
17
+- `FP8_PER_CHANNEL_PER_TOKEN`: per-channel weight scale and dynamic per-token activation quantization.
18
+- `FP8_PB_WO` (ModelOpt may emit `fp8_pb_wo`): block-scaled FP8 weight-only (typically 128×128 blocks).
19
+- `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`).
20
+
21
## Quantizing HuggingFace Models with PTQ
22
23
You can quantize HuggingFace models using the example scripts provided in the Model Optimizer repository. The primary script for LLM PTQ is typically found within the `examples/llm_ptq` directory.
@@ -80,3 +90,24 @@ The quantized checkpoint can then be deployed with vLLM. As an example, the foll
90
if __name__ == "__main__":
91
main()
92
```
93
+
94
+## Running the OpenAI-compatible server
95
+
96
+To serve a local ModelOpt checkpoint via the OpenAI-compatible API:
97
+
98
+```bash
99
+vllm serve <path_to_exported_checkpoint> \
100
+ --quantization modelopt \
101
+ --host 0.0.0.0 --port 8000
102
+```
103
+
104
+## Testing (local checkpoints)
105
+
106
+vLLM's ModelOpt unit tests are gated by local checkpoint paths and are skipped
107
+by default in CI. To run the tests locally:
108
+
109
+```bash
110
+export VLLM_TEST_MODELOPT_FP8_PC_PT_MODEL_PATH=<path_to_fp8_pc_pt_checkpoint>
111
+export VLLM_TEST_MODELOPT_FP8_PB_WO_MODEL_PATH=<path_to_fp8_pb_wo_checkpoint>
112
+pytest -q tests/quantization/test_modelopt.py
113
+```
tests/quantization/test_modelopt.py
+141
@@ -6,6 +6,7 @@ Run `pytest tests/quantization/test_modelopt.py`.
6
"""
7
8
import os
9
+from typing import NoReturn
10
11
import pytest
12
import torch
@@ -19,6 +20,28 @@ def enable_pickle(monkeypatch):
20
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
21
22
23
+def _skip(msg: str) -> NoReturn:
24
+ pytest.skip(msg)
25
+ raise RuntimeError(msg)
26
+
27
+
28
+def _snapshot_download_or_skip(model_id: str) -> str:
29
+ try:
30
+ from huggingface_hub import snapshot_download
31
+ except Exception as e: # pragma: no cover
32
+ _skip(f"huggingface_hub is required to download {model_id}: {e}")
33
+
34
+ try:
35
+ return snapshot_download(
36
+ repo_id=model_id,
37
+ repo_type="model",
38
+ # These checkpoints are already small; download full repo for simplicity.
39
+ allow_patterns=["*"],
40
+ )
41
+ except Exception as e:
42
+ _skip(f"Failed to download {model_id} from the HF Hub: {e}")
43
+
44
+
45
@pytest.mark.skipif(
46
not is_quant_method_supported("modelopt"),
47
reason="ModelOpt FP8 is not supported on this GPU type.",
@@ -91,3 +114,121 @@ def test_modelopt_fp8_checkpoint_setup(vllm_runner):
114
output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
115
assert output
116
print(f"ModelOpt FP8 output: {output}")
117
+
118
+
119
+@pytest.mark.skipif(
120
+ not is_quant_method_supported("modelopt"),
121
+ reason="ModelOpt FP8 is not supported on this GPU type.",
122
+)
123
+def test_modelopt_fp8_pc_pt_checkpoint_setup(vllm_runner):
124
+ """Test ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoint setup."""
125
+ model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pc-pt"
126
+ model_path = _snapshot_download_or_skip(model_id)
127
+
128
+ with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
129
+
130
+ def check_model(model):
131
+ layer = model.model.layers[0]
132
+
133
+ qkv_proj = layer.self_attn.qkv_proj
134
+ o_proj = layer.self_attn.o_proj
135
+ gate_up_proj = layer.mlp.gate_up_proj
136
+ down_proj = layer.mlp.down_proj
137
+
138
+ from vllm.model_executor.layers.quantization.modelopt import (
139
+ ModelOptFp8PcPtLinearMethod,
140
+ )
141
+
142
+ assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod)
143
+ assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod)
144
+ assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod)
145
+ assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod)
146
+
147
+ assert qkv_proj.weight.dtype == torch.float8_e4m3fn
148
+ assert o_proj.weight.dtype == torch.float8_e4m3fn
149
+ assert gate_up_proj.weight.dtype == torch.float8_e4m3fn
150
+ assert down_proj.weight.dtype == torch.float8_e4m3fn
151
+
152
+ # Per-channel scales; activations are dynamically scaled per token.
153
+ assert hasattr(qkv_proj, "weight_scale")
154
+ assert qkv_proj.weight_scale.dtype == torch.float32
155
+ assert qkv_proj.weight_scale.dim() == 1
156
+ assert not hasattr(qkv_proj, "input_scale")
157
+
158
+ assert hasattr(o_proj, "weight_scale")
159
+ assert o_proj.weight_scale.dtype == torch.float32
160
+ assert o_proj.weight_scale.dim() == 1
161
+ assert not hasattr(o_proj, "input_scale")
162
+
163
+ assert hasattr(gate_up_proj, "weight_scale")
164
+ assert gate_up_proj.weight_scale.dtype == torch.float32
165
+ assert gate_up_proj.weight_scale.dim() == 1
166
+ assert not hasattr(gate_up_proj, "input_scale")
167
+
168
+ assert hasattr(down_proj, "weight_scale")
169
+ assert down_proj.weight_scale.dtype == torch.float32
170
+ assert down_proj.weight_scale.dim() == 1
171
+ assert not hasattr(down_proj, "input_scale")
172
+
173
+ llm.apply_model(check_model)
174
+
175
+ output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
176
+ assert output
177
+ print(f"ModelOpt FP8_PER_CHANNEL_PER_TOKEN output: {output}")
178
+
179
+
180
+@pytest.mark.skipif(
181
+ not is_quant_method_supported("modelopt"),
182
+ reason="ModelOpt FP8 is not supported on this GPU type.",
183
+)
184
+def test_modelopt_fp8_pb_wo_checkpoint_setup(vllm_runner):
185
+ """Test ModelOpt FP8_PB_WO checkpoint setup."""
186
+ model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pb-wo"
187
+ model_path = _snapshot_download_or_skip(model_id)
188
+
189
+ with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
190
+
191
+ def check_model(model):
192
+ layer = model.model.layers[0]
193
+
194
+ qkv_proj = layer.self_attn.qkv_proj
195
+ o_proj = layer.self_attn.o_proj
196
+ gate_up_proj = layer.mlp.gate_up_proj
197
+ down_proj = layer.mlp.down_proj
198
+
199
+ from vllm.model_executor.layers.quantization.modelopt import (
200
+ ModelOptFp8PbWoLinearMethod,
201
+ )
202
+
203
+ assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod)
204
+ assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod)
205
+ assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod)
206
+ assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod)
207
+
208
+ assert qkv_proj.weight.dtype == torch.float8_e4m3fn
209
+ assert o_proj.weight.dtype == torch.float8_e4m3fn
210
+ assert gate_up_proj.weight.dtype == torch.float8_e4m3fn
211
+ assert down_proj.weight.dtype == torch.float8_e4m3fn
212
+
213
+ # Block scales; should be materialized as a 2D [out_blk, in_blk] tensor.
214
+ assert hasattr(qkv_proj, "weight_scale")
215
+ assert qkv_proj.weight_scale.dtype == torch.float32
216
+ assert qkv_proj.weight_scale.dim() == 2
217
+
218
+ assert hasattr(o_proj, "weight_scale")
219
+ assert o_proj.weight_scale.dtype == torch.float32
220
+ assert o_proj.weight_scale.dim() == 2
221
+
222
+ assert hasattr(gate_up_proj, "weight_scale")
223
+ assert gate_up_proj.weight_scale.dtype == torch.float32
224
+ assert gate_up_proj.weight_scale.dim() == 2
225
+
226
+ assert hasattr(down_proj, "weight_scale")
227
+ assert down_proj.weight_scale.dtype == torch.float32
228
+ assert down_proj.weight_scale.dim() == 2
229
+
230
+ llm.apply_model(check_model)
231
+
232
+ output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
233
+ assert output
234
+ print(f"ModelOpt FP8_PB_WO output: {output}")
vllm/config/model.py
+12
-6
@@ -843,12 +843,18 @@ class ModelConfig:
843
producer_name = quant_cfg.get("producer", {}).get("name")
844
if producer_name == "modelopt":
845
quant_algo = quant_cfg.get("quantization", {}).get("quant_algo")
846
- if quant_algo == "FP8":
847
- quant_cfg["quant_method"] = "modelopt"
848
- elif quant_algo == "NVFP4":
849
- quant_cfg["quant_method"] = "modelopt_fp4"
850
- elif quant_algo is not None:
851
- raise ValueError(f"Unknown ModelOpt quant algo: {quant_algo}")
846
+ if quant_algo is not None:
847
+ quant_algo_upper = str(quant_algo).upper()
848
+ if quant_algo_upper in {
849
+ "FP8",
850
+ "FP8_PER_CHANNEL_PER_TOKEN",
851
+ "FP8_PB_WO",
852
+ }:
853
+ quant_cfg["quant_method"] = "modelopt"
854
+ elif quant_algo_upper == "NVFP4":
855
+ quant_cfg["quant_method"] = "modelopt_fp4"
856
+ else:
857
+ raise ValueError(f"Unknown ModelOpt quant algo: {quant_algo}")
858
859
return quant_cfg
860
vllm/model_executor/layers/linear.py
+2
@@ -53,6 +53,8 @@ WEIGHT_LOADER_V2_SUPPORTED = [
53
"GPTQLinearMethod",
54
"FBGEMMFp8LinearMethod",
55
"ModelOptFp8LinearMethod",
56
+ "ModelOptFp8PcPtLinearMethod",
57
+ "ModelOptFp8PbWoLinearMethod",
58
"IPEXAWQLinearMethod",
59
"IPEXGPTQLinearMethod",
60
"HQQMarlinMethod",
vllm/model_executor/layers/quantization/modelopt.py
+251
-9
@@ -55,6 +55,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
55
select_cutlass_fp8_gemm_impl,
56
swap_w13_to_w31,
57
)
58
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
59
+ W8A8BlockFp8LinearOp,
60
+)
61
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
62
get_marlin_input_dtype,
63
)
@@ -72,9 +75,15 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
75
)
76
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
77
Fp8LinearOp,
78
+ cutlass_block_fp8_supported,
79
requantize_with_max_scale,
80
)
77
-from vllm.model_executor.parameter import ModelWeightParameter, PerTensorScaleParameter
81
+from vllm.model_executor.parameter import (
82
+ BlockQuantScaleParameter,
83
+ ChannelQuantScaleParameter,
84
+ ModelWeightParameter,
85
+ PerTensorScaleParameter,
86
+)
87
from vllm.scalar_type import scalar_types
88
from vllm.utils.flashinfer import (
89
flashinfer_scaled_fp4_mm,
@@ -88,7 +97,16 @@ if TYPE_CHECKING:
97
98
logger = init_logger(__name__)
99
91
-QUANT_ALGOS = ["FP8", "NVFP4"]
100
+QUANT_ALGOS = [
101
+ # FP8 (per-tensor weight + optional static activation scale).
102
+ "FP8",
103
+ # FP8 per-channel weight scale + per-token activation scale.
104
+ "FP8_PER_CHANNEL_PER_TOKEN",
105
+ # FP8 per-block weight-only (ModelOpt may emit this as lowercase).
106
+ "FP8_PB_WO",
107
+ # FP4
108
+ "NVFP4",
109
+]
110
KV_CACHE_QUANT_ALGOS = ["FP8"]
111
112
@@ -255,6 +273,9 @@ class ModelOptQuantConfigBase(QuantizationConfig):
273
if not quant_method:
274
raise ValueError("Missing 'quant_algo' in quantization config")
275
276
+ # Normalize quant_algo for robust matching (ModelOpt may emit lowercase).
277
+ quant_method = str(quant_method).upper()
278
+
279
if kv_cache_quant_method is None:
280
# No KV cache quantization, keep this branch just to have this comment
281
pass
@@ -263,6 +284,8 @@ class ModelOptQuantConfigBase(QuantizationConfig):
284
f"kv_cache_quant_algo must be a string, got "
285
f"{type(kv_cache_quant_method)}"
286
)
287
+ else:
288
+ kv_cache_quant_method = kv_cache_quant_method.upper()
289
290
if not isinstance(exclude_modules, list):
291
raise ValueError(
@@ -302,17 +325,34 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):
325
326
def __init__(
327
self,
328
+ quant_method: str,
329
is_checkpoint_fp8_serialized: bool,
330
kv_cache_quant_method: str | None,
331
exclude_modules: list[str],
332
) -> None:
333
super().__init__(exclude_modules)
334
+ self.quant_method = quant_method
335
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
336
self.kv_cache_quant_method = kv_cache_quant_method
337
if is_checkpoint_fp8_serialized:
338
logger.warning(
314
- "Detected ModelOpt fp8 checkpoint. Please note that"
315
- " the format is experimental and could change."
339
+ "Detected ModelOpt fp8 checkpoint (quant_algo=%s). Please note "
340
+ "that the format is experimental and could change.",
341
+ quant_method,
342
+ )
343
+
344
+ # Select LinearMethod implementation based on quant_algo.
345
+ if self.quant_method == "FP8":
346
+ self.LinearMethodCls = ModelOptFp8LinearMethod
347
+ elif self.quant_method == "FP8_PER_CHANNEL_PER_TOKEN":
348
+ self.LinearMethodCls = ModelOptFp8PcPtLinearMethod
349
+ elif self.quant_method == "FP8_PB_WO":
350
+ self.LinearMethodCls = ModelOptFp8PbWoLinearMethod
351
+ else:
352
+ raise ValueError(
353
+ "Unsupported ModelOpt FP8 quant_algo for vLLM: "
354
+ f"{self.quant_method}. Supported: FP8 / "
355
+ "FP8_PER_CHANNEL_PER_TOKEN / FP8_PB_WO."
356
)
357
358
def get_name(self) -> QuantizationMethods:
@@ -346,13 +386,13 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):
386
if "quantization" in hf_quant_cfg:
387
quant_config = hf_quant_cfg["quantization"]
388
if isinstance(quant_config, dict):
349
- quant_algo = quant_config.get("quant_algo", "")
350
- if "FP8" in quant_algo:
389
+ quant_algo = str(quant_config.get("quant_algo", ""))
390
+ if "FP8" in quant_algo.upper():
391
return "modelopt"
392
else:
393
# Check for compressed-tensors style config with specific quant_algo
354
- quant_algo = hf_quant_cfg.get("quant_algo", "")
355
- if isinstance(quant_algo, str) and "FP8" in quant_algo:
394
+ quant_algo = str(hf_quant_cfg.get("quant_algo", ""))
395
+ if "FP8" in quant_algo.upper():
396
return "modelopt"
397
398
return None
@@ -369,7 +409,12 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):
409
) -> "ModelOptFp8Config":
410
is_checkpoint_fp8_serialized = "FP8" in quant_method
411
372
- return cls(is_checkpoint_fp8_serialized, kv_cache_quant_method, exclude_modules)
412
+ return cls(
413
+ quant_method,
414
+ is_checkpoint_fp8_serialized,
415
+ kv_cache_quant_method,
416
+ exclude_modules,
417
+ )
418
419
420
class ModelOptFp8LinearMethod(LinearMethodBase):
@@ -464,6 +509,203 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
509
)
510
511
512
+class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
513
+ """Linear method for ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoints.
514
+
515
+ Expected checkpoint structure (per Linear):
516
+ - weight: fp8-e4m3fn, shape [out, in]
517
+ - weight_scale: fp32, shape [out] (per-output-channel)
518
+ - no input_scale (activations are dynamically quantized per-token)
519
+ """
520
+
521
+ def __init__(self, quant_config: ModelOptFp8Config) -> None:
522
+ self.quant_config = quant_config
523
+ self.fp8_linear = Fp8LinearOp(
524
+ act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN
525
+ )
526
+
527
+ def create_weights(
528
+ self,
529
+ layer: torch.nn.Module,
530
+ input_size_per_partition: int,
531
+ output_partition_sizes: list[int],
532
+ input_size: int,
533
+ output_size: int,
534
+ params_dtype: torch.dtype,
535
+ **extra_weight_attrs,
536
+ ):
537
+ del input_size, output_size
538
+
539
+ if not self.quant_config.is_checkpoint_fp8_serialized:
540
+ raise ValueError(
541
+ "FP8_PER_CHANNEL_PER_TOKEN currently only supports "
542
+ "FP8-serialized checkpoints."
543
+ )
544
+
545
+ output_size_per_partition = sum(output_partition_sizes)
546
+ weight_loader = extra_weight_attrs.get("weight_loader")
547
+ layer.logical_widths = output_partition_sizes
548
+ layer.input_size_per_partition = input_size_per_partition
549
+ layer.output_size_per_partition = output_size_per_partition
550
+
551
+ weight = ModelWeightParameter(
552
+ data=torch.empty(
553
+ output_size_per_partition,
554
+ input_size_per_partition,
555
+ dtype=torch.float8_e4m3fn,
556
+ ),
557
+ input_dim=1,
558
+ output_dim=0,
559
+ weight_loader=weight_loader,
560
+ )
561
+ layer.register_parameter("weight", weight)
562
+
563
+ weight_scale = ChannelQuantScaleParameter(
564
+ data=torch.empty(output_size_per_partition, dtype=torch.float32),
565
+ output_dim=0,
566
+ weight_loader=weight_loader,
567
+ )
568
+ weight_scale[:] = torch.finfo(torch.float32).min
569
+ layer.register_parameter("weight_scale", weight_scale)
570
+
571
+ def process_weights_after_loading(self, layer: Module) -> None:
572
+ layer.weight = Parameter(layer.weight.t(), requires_grad=False)
573
+ layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
574
+
575
+ def apply(
576
+ self,
577
+ layer: torch.nn.Module,
578
+ x: torch.Tensor,
579
+ bias: torch.Tensor | None = None,
580
+ ) -> torch.Tensor:
581
+ return self.fp8_linear.apply(
582
+ input=x,
583
+ weight=layer.weight,
584
+ weight_scale=layer.weight_scale,
585
+ input_scale=None,
586
+ bias=bias,
587
+ )
588
+
589
+
590
+class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
591
+ """Linear method for ModelOpt FP8_PB_WO checkpoints.
592
+
593
+ ModelOpt exports `weight_scale` as a 4D tensor:
594
+ [out_blk, 1, in_blk, 1]
595
+ where block size is typically 128 for both dims.
596
+
597
+ vLLM executes it as FP8 GEMM with *dynamic per-token* activation quant.
598
+ """
599
+
600
+ _WEIGHT_BLOCK_SIZE: tuple[int, int] = (128, 128)
601
+
602
+ def __init__(self, quant_config: ModelOptFp8Config) -> None:
603
+ self.quant_config = quant_config
604
+ block_n, block_k = self._WEIGHT_BLOCK_SIZE
605
+ self.weight_block_size = list(self._WEIGHT_BLOCK_SIZE)
606
+ self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
607
+ weight_group_shape=GroupShape(block_n, block_k),
608
+ act_quant_group_shape=GroupShape(1, block_k),
609
+ cutlass_block_fp8_supported=cutlass_block_fp8_supported(),
610
+ use_aiter_and_is_supported=False,
611
+ )
612
+
613
+ def create_weights(
614
+ self,
615
+ layer: torch.nn.Module,
616
+ input_size_per_partition: int,
617
+ output_partition_sizes: list[int],
618
+ input_size: int,
619
+ output_size: int,
620
+ params_dtype: torch.dtype,
621
+ **extra_weight_attrs,
622
+ ):
623
+ del input_size, output_size
624
+
625
+ if not self.quant_config.is_checkpoint_fp8_serialized:
626
+ raise ValueError(
627
+ "FP8_PB_WO currently only supports FP8-serialized checkpoints."
628
+ )
629
+
630
+ output_size_per_partition = sum(output_partition_sizes)
631
+ weight_loader = extra_weight_attrs.get("weight_loader")
632
+ layer.logical_widths = output_partition_sizes
633
+ layer.input_size_per_partition = input_size_per_partition
634
+ layer.output_size_per_partition = output_size_per_partition
635
+
636
+ # Expose block size so the v2 weight loaders can translate offsets from
637
+ # element-space -> block-space for BlockQuantScaleParameter.
638
+ layer.weight_block_size = self.weight_block_size
639
+
640
+ weight = ModelWeightParameter(
641
+ data=torch.empty(
642
+ output_size_per_partition,
643
+ input_size_per_partition,
644
+ dtype=torch.float8_e4m3fn,
645
+ ),
646
+ input_dim=1,
647
+ output_dim=0,
648
+ weight_loader=weight_loader,
649
+ )
650
+ layer.register_parameter("weight", weight)
651
+
652
+ block_n, block_k = self._WEIGHT_BLOCK_SIZE
653
+ if output_size_per_partition % block_n != 0:
654
+ raise ValueError(
655
+ "ModelOpt FP8_PB_WO requires out_features divisible by "
656
+ f"{block_n}, got {output_size_per_partition}."
657
+ )
658
+ if input_size_per_partition % block_k != 0:
659
+ raise ValueError(
660
+ "ModelOpt FP8_PB_WO requires in_features divisible by "
661
+ f"{block_k}, got {input_size_per_partition}."
662
+ )
663
+
664
+ out_blks = output_size_per_partition // block_n
665
+ in_blks = input_size_per_partition // block_k
666
+
667
+ # Match ModelOpt's exported shape so weight loading works without a
668
+ # custom loader: [out_blk, 1, in_blk, 1]
669
+ weight_scale = BlockQuantScaleParameter(
670
+ data=torch.empty((out_blks, 1, in_blks, 1), dtype=torch.float32),
671
+ input_dim=2,
672
+ output_dim=0,
673
+ weight_loader=weight_loader,
674
+ )
675
+ weight_scale[:] = torch.finfo(torch.float32).min
676
+ layer.register_parameter("weight_scale", weight_scale)
677
+
678
+ def process_weights_after_loading(self, layer: Module) -> None:
679
+ # Keep weight in [out, in] layout for W8A8BlockFp8LinearOp.
680
+ layer.weight = Parameter(layer.weight.data, requires_grad=False)
681
+
682
+ scale = layer.weight_scale
683
+ if scale.dim() == 4:
684
+ # [out_blk, 1, in_blk, 1] -> [out_blk, in_blk]
685
+ scale = scale.squeeze(1).squeeze(-1)
686
+ elif scale.dim() != 2:
687
+ raise ValueError(
688
+ "Unexpected ModelOpt FP8_PB_WO weight_scale shape: "
689
+ f"{tuple(scale.shape)}."
690
+ )
691
+
692
+ layer.weight_scale = Parameter(scale.contiguous(), requires_grad=False)
693
+
694
+ def apply(
695
+ self,
696
+ layer: torch.nn.Module,
697
+ x: torch.Tensor,
698
+ bias: torch.Tensor | None = None,
699
+ ) -> torch.Tensor:
700
+ return self.w8a8_block_fp8_linear.apply(
701
+ input=x,
702
+ weight=layer.weight,
703
+ weight_scale=layer.weight_scale,
704
+ input_scale=None,
705
+ bias=bias,
706
+ )
707
+
708
+
709
class ModelOptFp8MoEMethod(FusedMoEMethodBase):
710
"""MoE method for ModelOpt FP8.
711
Supports loading FP8 checkpoints with static weight scale and