modelopt.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,661 lines 100 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	from collections.abc import Callable
5	from dataclasses import dataclass
6	from enum import Enum
7	from fnmatch import fnmatch
8	from typing import TYPE_CHECKING, Any, cast
9
10	import torch
11	from torch.nn.parameter import Parameter
12
13	import vllm.envs as envs
14	from vllm.config import get_current_vllm_config_or_none
15	from vllm.config.quantization import QuantSpec
16	from vllm.logger import init_logger
17	from vllm.model_executor.kernels.linear import (
18	init_fp8_linear_kernel,
19	init_mxfp8_linear_kernel,
20	init_nvfp4_linear_kernel,
21	)
22	from vllm.model_executor.layers.attention import Attention, MLAAttention
23	from vllm.model_executor.layers.fused_moe import (
24	FusedMoEConfig,
25	FusedMoEMethodBase,
26	FusedMoEQuantConfig,
27	FusedMoeWeightScaleSupported,
28	RoutedExperts,
29	SharedExperts,
30	)
31	from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
32	from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
33	Fp8MoeBackend,
34	convert_to_fp8_moe_kernel_format,
35	make_fp8_moe_kernel,
36	make_fp8_moe_quant_config,
37	select_fp8_moe_backend,
38	)
39	from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
40	select_mxfp8_moe_backend,
41	)
42	from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
43	NvFp4MoeBackend,
44	convert_to_nvfp4_moe_kernel_format,
45	is_global_sf_supported_for_nvfp4_backend,
46	make_nvfp4_moe_kernel,
47	make_nvfp4_moe_quant_config,
48	select_nvfp4_moe_backend,
49	)
50	from vllm.model_executor.layers.fusion.quant_activation import (
51	expose_input_quant_key,
52	)
53	from vllm.model_executor.layers.linear import (
54	LinearBase,
55	LinearMethodBase,
56	UnquantizedLinearMethod,
57	register_weight_loader_v2_supported_method,
58	)
59	from vllm.model_executor.layers.quantization import QuantizationMethods
60	from vllm.model_executor.layers.quantization.base_config import (
61	QuantizationConfig,
62	QuantizeMethodBase,
63	)
64	from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
65	from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
66	from vllm.model_executor.layers.quantization.utils.fp8_utils import (
67	FP8_SCALE_SENTINEL,
68	process_fp8_input_tensor_strategy_moe,
69	process_fp8_weight_channel_strategy,
70	process_fp8_weight_tensor_strategy_moe,
71	)
72	from vllm.model_executor.layers.quantization.utils.marlin_utils import (
73	get_marlin_input_dtype,
74	)
75	from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
76	MXFP8_BLOCK_SIZE,
77	MXFP8_SCALE_DTYPE,
78	MXFP8_VALUE_DTYPE,
79	)
80	from vllm.model_executor.layers.quantization.utils.quant_utils import (
81	FP4_DTYPE,
82	QuantKey,
83	is_layer_skipped,
84	kFp8Dynamic128Sym,
85	kFp8DynamicTokenSym,
86	kFp8Static128BlockSym,
87	kFp8StaticTensorSym,
88	kFp8StaticTokenSym,
89	kMxfp8Dynamic,
90	kMxfp8Static,
91	kNvfp4Dynamic,
92	kNvfp4Static,
93	)
94	from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
95	requantize_with_max_scale,
96	)
97	from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
98	from vllm.model_executor.parameter import (
99	BlockQuantScaleParameter,
100	ChannelQuantScaleParameter,
101	ModelWeightParameter,
102	PerTensorScaleParameter,
103	)
104	from vllm.model_executor.utils import (
105	is_weights_pre_processed,
106	replace_parameter,
107	set_weight_attrs,
108	)
109	from vllm.utils.math_utils import cdiv
110
111	if TYPE_CHECKING:
112	from vllm.model_executor.models.utils import WeightsMapper
113
114	logger = init_logger(__name__)
115
116	# ``FP8_PB_WO`` is ModelOpt's canonical 2D block-FP8 name. Early composed
117	# Qwen3.8-Flash-Next checkpoints used ``FP8_BLOCK_SCALES`` for the same tensor
118	# layout, so retain it as a checkpoint-compatibility alias.
119	_BLOCK_FP8_MOE_ALGOS = ("FP8_PB_WO", "FP8_BLOCK_SCALES")
120
121	# Single source of truth for the ModelOpt linear algos.
122	#
123	# Each entry is (owning config's name, mixed-precision sub-config attribute).
124	# ``resolve`` turns the algo into a QuantSpec; this table only records which
125	# config object holds the checkpoint parameters resolve may read -- today just
126	# group_size, and only for the NVFP4 family.
127	#
128	# To add a format: one row here plus one row in resolve(). Every consumer below
129	# is derived, so nothing else needs editing.
130	LINEAR_ALGOS: dict[str, tuple[str, str]] = {
131	# FP8 per-tensor weight + optional static activation scale.
132	"FP8": ("modelopt", "fp8_config"),
133	# FP8 per-channel weight scale + per-token activation scale.
134	"FP8_PER_CHANNEL_PER_TOKEN": ("modelopt", "fp8_config"),
135	# FP8 128x128 per-block weight (ModelOpt may emit this lowercase).
136	"FP8_PB_WO": ("modelopt", "fp8_config"),
137	# NVFP4 W4A4: 4-bit float weights AND 4-bit float activations.
138	"NVFP4": ("modelopt_fp4", "nvfp4_config"),
139	# W4A16 NVFP4: 4-bit float weights, fp16/bf16 activations.
140	"W4A16_NVFP4": ("modelopt_fp4", "w4a16_nvfp4_config"),
141	# MXFP8: e4m3 weights with per-32-block e8m0 scales.
142	"MXFP8": ("modelopt_mxfp8", "mxfp8_config"),
143	}
144
145	# MIXED_PRECISION is not a linear algo; it selects a per-layer algo from the
146	# table above.
147	QUANT_ALGOS = [*LINEAR_ALGOS, "MIXED_PRECISION"]
148
149
150	def algos_owned_by(config_name: str) -> tuple[str, ...]:
151	"""The linear algos a given config accepts, for its quant_algo validation."""
152	return tuple(a for a, (owner, _) in LINEAR_ALGOS.items() if owner == config_name)
153
154
155	class ModelOptKVCacheMethod(BaseKVCacheMethod):
156	"""Supports loading kv-cache scaling factors from FP8 or NVFP4 checkpoints."""
157
158	def __init__(self, quant_config: "ModelOptQuantConfigBase"):
159	super().__init__(quant_config)
160
161
162	class ModelOptQuantConfigBase(QuantizationConfig):
163	# ModelOpt quant-algo string, set by each subclass. Fed to resolve() to
164	# build the QuantSpec for the generic ModelOptLinearMethod. The mixed
165	# config resolves the algo per-prefix instead and does not set this.
166	quant_method: str
167
168	FusedMoEMethodCls: type = FusedMoEMethodBase
169	KVCacheMethodCls: type = BaseKVCacheMethod
170
171	def __init__(
172	self,
173	exclude_modules: list[str],
174	):
175	super().__init__()
176	self.exclude_modules: list[str] = exclude_modules
177
178	def is_layer_excluded(self, prefix: str) -> bool:
179	"""Check if a layer should be excluded from quantization.
180
181	Handles both exact matching (for fused layers) and ModelOpt wildcard matching.
182
183	The ModelOpt exclude_modules list is a list of wildcards.
184	"""
185	if len(self.exclude_modules) == 0:
186	return False
187
188	# First check exact matching with fused layer support
189	if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):
190	return True
191
192	# TODO: This special hard coded logic is not needed for quantized checkpoints
193	# generated by ModelOpt >= 0.39.0 where they are handled natually by the
194	# exclude_modules config. But need to keep them for loading quantized
195	# checkpoints generated by older versions. Then check substring matching
196	# for patterns not caught by exact match
197	for exclude_module in self.exclude_modules:
198	# Skip exact matches already handled above
199	if exclude_module != prefix and (
200	exclude_module in prefix
201	or (
202	prefix.startswith("language_model.")
203	and exclude_module in prefix.removeprefix("language_model.")
204	)
205	):
206	return True
207
208	# modelopt exclude modules are not simple strings, they are wildcards
209	for wildcard_pattern in self.exclude_modules:
210	if fnmatch(prefix, wildcard_pattern):
211	return True
212
213	return False
214
215	def get_quant_method(
216	self, layer: torch.nn.Module, prefix: str
217	) -> "QuantizeMethodBase \| None":
218	# handle kv-cache first so we can focus only on weight quantization thereafter
219	if isinstance(layer, (Attention, MLAAttention)):
220	return self.KVCacheMethodCls(self)
221
222	# handle exclusion
223	if self.is_layer_excluded(prefix):
224	if isinstance(layer, (LinearBase, ParallelLMHead)):
225	return UnquantizedLinearMethod()
226	return None
227
228	# TODO: This special hard coded logic is not needed for quantized checkpoints
229	# generated by ModelOpt >= 0.39.0 where they are handled natually by the
230	# exclude_modules config. But need to keep them for loading quantized
231	# checkpoints generated by older versions. Then check substring matching
232	# for patterns not caught by exact match
233	if (
234	"vision_tower" in prefix
235	or "vision_model" in prefix
236	or "vit_large_projector" in prefix
237	):
238	return UnquantizedLinearMethod()
239
240	# now, the layer is quantized, handle it here
241	if isinstance(layer, (LinearBase, ParallelLMHead)):
242	return build_linear_method(self, self.quant_method, prefix)
243	elif isinstance(layer, RoutedExperts):
244	quant_method = self.FusedMoEMethodCls(
245	quant_config=self, moe_config=layer.moe_config
246	)
247	if getattr(quant_method, "backend", "") == "marlin":
248	quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
249	return quant_method
250
251	return None
252
253	def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
254	if len(self.exclude_modules) > 0:
255	# This is a workaround for the weights remapping issue:
256	# https://github.com/vllm-project/vllm/issues/28072
257	# Right now, the Nvidia ModelOpt library use just one wildcard pattern:
258	# module_path*
259	# It gets applied if the whole tree of modules rooted at module_path
260	# is not quantized. Here we replace such pattern by 2 patterns that are
261	# collectively equivalent to the original pattern:
262	# module_path
263	# module_path.*
264	new_exclude_modules = []
265	for exclude in self.exclude_modules:
266	if len(exclude) >= 2 and exclude[-1] == "*" and exclude[-2] != ".":
267	new_exclude_modules.append(exclude[:-1])
268	new_exclude_modules.append(exclude[:-1] + ".*")
269	else:
270	new_exclude_modules.append(exclude)
271
272	self.exclude_modules = hf_to_vllm_mapper.apply_list(new_exclude_modules)
273
274	@staticmethod
275	def _extract_modelopt_quant_algo(
276	hf_quant_cfg: dict[str, Any] \| None,
277	) -> str \| None:
278	"""Extract upper-cased quant_algo from a modelopt config.
279
280	Returns the quant_algo string (upper-cased), or None if the config
281	is not a modelopt config.
282	"""
283	if hf_quant_cfg is None:
284	return None
285	if not hf_quant_cfg.get("quant_method", "").lower().startswith("modelopt"):
286	return None
287	if "quantization" in hf_quant_cfg:
288	quant_config = hf_quant_cfg["quantization"]
289	if isinstance(quant_config, dict):
290	return str(quant_config.get("quant_algo", "")).upper()
291	return None
292	return str(hf_quant_cfg.get("quant_algo", "")).upper()
293
294	@staticmethod
295	def get_config_filenames() -> list[str]:
296	return ["hf_quant_config.json"]
297
298	@classmethod
299	def _from_config(
300	cls,
301	*,
302	quant_method: str,
303	kv_cache_quant_method: str \| None,
304	exclude_modules: list[str],
305	original_config: dict[str, Any],
306	group_size: int \| None,
307	) -> "ModelOptQuantConfigBase":
308	raise NotImplementedError("Please implement this function in sub classes")
309
310	@classmethod
311	def from_config(cls, config: dict[str, Any]) -> "ModelOptQuantConfigBase":
312	# Handle both ModelOpt format and compressed-tensors style format
313	if "quantization" in config:
314	# Traditional ModelOpt format:
315	# {"quantization": {"quant_algo": "..."}}
316	quant_config = cls.get_from_keys(config, ["quantization"])
317	if not isinstance(quant_config, dict):
318	raise ValueError("Expected 'quantization' to be a dictionary in config")
319
320	quant_method = quant_config.get("quant_algo")
321
322	# Handle kv_cache_quant_algo with proper type validation
323	kv_cache_quant_method = quant_config.get("kv_cache_quant_algo")
324
325	# Handle group_size with proper type validation
326	group_size_raw = quant_config.get("group_size")
327
328	# "exclude_modules" is the key in the legacy hf_quant_config.json
329	exclude_modules = quant_config.get("exclude_modules", [])
330	else:
331	# Compressed-tensors style format (config.json quantization_config):
332	# {"quant_algo": "...", "quant_method": "modelopt"}
333	quant_method = config.get("quant_algo")
334
335	# "kv_cache_scheme" (a dict) instead of "kv_cache_quant_algo" (a string).
336	kv_cache_scheme = config.get("kv_cache_scheme")
337	if isinstance(kv_cache_scheme, dict) and (
338	kv_cache_scheme.get("type") == "float"
339	and kv_cache_scheme.get("num_bits") == 8
340	):
341	kv_cache_quant_method = "FP8"
342	else:
343	kv_cache_quant_method = None
344
345	# "ignore" is the key in config.json
346	exclude_modules = config.get("ignore", [])
347	group_size_raw = config.get("group_size")
348
349	if not quant_method:
350	raise ValueError("Missing 'quant_algo' in quantization config")
351
352	# Normalize quant_algo for robust matching (ModelOpt may emit lowercase).
353	quant_method = str(quant_method).upper()
354
355	if kv_cache_quant_method is None:
356	# No KV cache quantization, keep this branch just to have this comment
357	pass
358	elif not isinstance(kv_cache_quant_method, str):
359	raise ValueError(
360	f"kv_cache_quant_algo must be a string, got "
361	f"{type(kv_cache_quant_method)}"
362	)
363	else:
364	kv_cache_quant_method = kv_cache_quant_method.upper()
365
366	if not isinstance(exclude_modules, list):
367	raise ValueError(
368	f"exclude_modules must be a list, got {type(exclude_modules)}"
369	)
370
371	if group_size_raw is None:
372	group_size = None
373	elif isinstance(group_size_raw, int):
374	group_size = group_size_raw
375	else:
376	try:
377	group_size = int(group_size_raw)
378	except (ValueError, TypeError):
379	raise ValueError(
380	f"group_size must be an integer, got {type(group_size_raw)}"
381	) from None
382
383	if quant_method not in QUANT_ALGOS:
384	raise ValueError(
385	f"ModelOpt currently only supports: {QUANT_ALGOS} "
386	"quantizations in vLLM. Please check the "
387	"`hf_quant_config.json` file for your model's "
388	"quant configuration."
389	)
390	return cls._from_config(
391	quant_method=quant_method,
392	kv_cache_quant_method=kv_cache_quant_method,
393	exclude_modules=exclude_modules,
394	group_size=group_size,
395	original_config=config,
396	)
397
398
399	class ModelOptFp8Config(ModelOptQuantConfigBase):
400	"""Config class for ModelOpt FP8."""
401
402	def __init__(
403	self,
404	quant_method: str,
405	is_checkpoint_fp8_serialized: bool,
406	kv_cache_quant_method: str \| None,
407	exclude_modules: list[str],
408	) -> None:
409	super().__init__(exclude_modules)
410	self.quant_method = quant_method
411	self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
412	self.kv_cache_quant_method = kv_cache_quant_method
413
414	# Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch.
415	supported = algos_owned_by("modelopt")
416	if self.quant_method not in supported:
417	raise ValueError(
418	"Unsupported ModelOpt FP8 quant_algo for vLLM: "
419	f"{self.quant_method}. Supported: {' / '.join(supported)}."
420	)
421
422	def has_blocked_weights(self) -> bool:
423	# Gates "+quant_fp8" in VllmConfig. FP8_PB_WO runs the block-scaled
424	# kernels, whose DeepGEMM path needs UE8M0-packed group scales; without
425	# the op enabled QuantFP8 falls back to forward_native, which emits
426	# unpacked fp32 scales and the GEMM fails to launch.
427	return self.quant_method == "FP8_PB_WO"
428
429	def get_name(self) -> QuantizationMethods:
430	return "modelopt"
431
432	def get_supported_act_dtypes(self) -> list[torch.dtype]:
433	return [torch.bfloat16, torch.half]
434
435	@classmethod
436	def get_min_capability(cls) -> int:
437	return 80
438
439	@classmethod
440	def override_quantization_method(
441	cls, hf_quant_cfg, user_quant, hf_config=None
442	) -> QuantizationMethods \| None:
443	algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)
444	if algo is not None and algo == "FP8":
445	return "modelopt"
446	return None
447
448	@classmethod
449	def _from_config(
450	cls,
451	*,
452	quant_method: str,
453	kv_cache_quant_method: str \| None,
454	exclude_modules: list[str],
455	original_config: dict[str, Any],
456	**kwargs: Any,
457	) -> "ModelOptFp8Config":
458	is_checkpoint_fp8_serialized = "FP8" in quant_method
459
460	return cls(
461	quant_method,
462	is_checkpoint_fp8_serialized,
463	kv_cache_quant_method,
464	exclude_modules,
465	)
466
467
468	class ModelOptFp8MoEMethod(FusedMoEMethodBase):
469	"""MoE method for ModelOpt FP8.
470	Supports loading FP8 checkpoints with static weight scale and
471	activation scale.
472
473	Args:
474	quant_config: The ModelOpt quantization config.
475
476	"""
477
478	def __init__(
479	self,
480	quant_config: ModelOptFp8Config,
481	moe_config: FusedMoEConfig,
482	) -> None:
483	super().__init__(moe_config)
484	self.quant_config = quant_config
485	assert self.quant_config.is_checkpoint_fp8_serialized
486
487	# Select Fp8 MoE backend
488	self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
489	config=self.moe,
490	weight_key=kFp8StaticTensorSym,
491	activation_key=kFp8StaticTensorSym,
492	)
493
494	def create_weights(
495	self,
496	layer: RoutedExperts,
497	num_experts: int,
498	hidden_size: int,
499	intermediate_size_per_partition: int,
500	params_dtype: torch.dtype,
501	**extra_weight_attrs,
502	):
503	layer.orig_dtype = params_dtype
504	layer.num_experts = num_experts
505
506	# Use FP8 dtype if checkpoint is serialized
507	weight_dtype = (
508	torch.float8_e4m3fn
509	if self.quant_config.is_checkpoint_fp8_serialized
510	else params_dtype
511	)
512	weight_loader = extra_weight_attrs.get("weight_loader")
513
514	w13_num_shards = 2 if self.moe.is_act_and_mul else 1
515
516	w13_weight = ModelWeightParameter(
517	data=torch.empty(
518	num_experts,
519	w13_num_shards * intermediate_size_per_partition,
520	hidden_size,
521	dtype=weight_dtype,
522	),
523	input_dim=2,
524	output_dim=1,
525	weight_loader=weight_loader,
526	)
527	layer.register_parameter("w13_weight", w13_weight)
528
529	w2_weight = ModelWeightParameter(
530	data=torch.empty(
531	num_experts,
532	hidden_size,
533	intermediate_size_per_partition,
534	dtype=weight_dtype,
535	),
536	input_dim=2,
537	output_dim=1,
538	weight_loader=weight_loader,
539	)
540	layer.register_parameter("w2_weight", w2_weight)
541
542	# WEIGHT SCALES - Per-tensor scaling for ModelOpts
543	# For gated MoE, allocate 2 scales for w1 and w3 respectively.
544	# They will be combined to a single scale after weight loading.
545	# For non-gated MoE, allocate 1 scale for w13.
546	w13_weight_scale = PerTensorScaleParameter(
547	data=torch.full(
548	(num_experts, w13_num_shards),
549	1.0,
550	dtype=torch.float32,
551	),
552	weight_loader=weight_loader,
553	)
554	w2_weight_scale = PerTensorScaleParameter(
555	data=torch.full((num_experts,), 1.0, dtype=torch.float32),
556	weight_loader=weight_loader,
557	)
558	layer.register_parameter("w13_weight_scale", w13_weight_scale)
559	layer.register_parameter("w2_weight_scale", w2_weight_scale)
560
561	# INPUT SCALES - Per-tensor scaling for ModelOpt
562	w13_input_scale = PerTensorScaleParameter(
563	data=torch.full((num_experts,), 1.0, dtype=torch.float32),
564	weight_loader=weight_loader,
565	)
566	w2_input_scale = PerTensorScaleParameter(
567	data=torch.full((num_experts,), 1.0, dtype=torch.float32),
568	weight_loader=weight_loader,
569	)
570	layer.register_parameter("w13_input_scale", w13_input_scale)
571	layer.register_parameter("w2_input_scale", w2_input_scale)
572
573	def _setup_kernel(
574	self,
575	layer: RoutedExperts,
576	w13: torch.Tensor,
577	w2: torch.Tensor,
578	w13_scale: torch.Tensor,
579	w2_scale: torch.Tensor,
580	w13_input_scale: torch.Tensor,
581	w2_input_scale: torch.Tensor,
582	):
583	w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
584	fp8_backend=self.fp8_backend,
585	layer=layer,
586	w13=w13,
587	w2=w2,
588	w13_scale=w13_scale,
589	w2_scale=w2_scale,
590	w13_input_scale=w13_input_scale,
591	w2_input_scale=w2_input_scale,
592	)
593
594	# Replace parameters with updated versions. Note that this helper
595	# function ensures the replacement is compatible with RL weight reloads.
596	replace_parameter(layer, "w13_weight", w13)
597	replace_parameter(layer, "w2_weight", w2)
598	replace_parameter(layer, "w13_weight_scale", w13_scale)
599	replace_parameter(layer, "w2_weight_scale", w2_scale)
600
601	# Setup modular kernel.
602	self.moe_quant_config = self.get_fused_moe_quant_config(layer)
603	assert self.experts_cls is not None
604	self.moe_kernel = make_fp8_moe_kernel(
605	moe_quant_config=self.moe_quant_config,
606	moe_config=self.moe,
607	fp8_backend=self.fp8_backend,
608	experts_cls=self.experts_cls,
609	routing_tables=layer._expert_routing_tables(),
610	)
611
612	def process_weights_after_loading(self, layer: RoutedExperts) -> None:
613	w13 = layer.w13_weight
614	w2 = layer.w2_weight
615	w13_scale = layer.w13_weight_scale
616	w2_scale = layer.w2_weight_scale
617	w13_input_scale = layer.w13_input_scale
618	w2_input_scale = layer.w2_input_scale
619
620	# Per tensor kernels require single activation scale. Use the max.
621	w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
622	w13_input_scale,
623	w2_input_scale,
624	layer.moe_config.moe_parallel_config.enable_eplb,
625	)
626	replace_parameter(layer, "w13_input_scale", w13_input_scale)
627	replace_parameter(layer, "w2_input_scale", w2_input_scale)
628
629	# Per tensor kernels require single weight scale for w13 per expert, but
630	# on disk there is a scale for w1 and w3. Use the max to requantize.
631	shard_size = layer.intermediate_size_per_partition
632	w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
633	w13,
634	w13_scale,
635	shard_size,
636	num_experts=layer.w13_weight.shape[0],
637	is_act_and_mul=self.moe.is_act_and_mul,
638	)
639
640	# Shuffle weights to runtime format and setup kernel.
641	self._setup_kernel(
642	layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
643	)
644
645	def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
646	w1_scale = layer.w13_weight_scale
647	w2_scale = layer.w2_weight_scale
648	a1_scale = layer.w13_input_scale
649	a2_scale = layer.w2_input_scale
650
651	return make_fp8_moe_quant_config(
652	fp8_backend=self.fp8_backend,
653	w1_scale=w1_scale,
654	w2_scale=w2_scale,
655	a1_scale=a1_scale,
656	a2_scale=a2_scale,
657	swiglu_limit=getattr(layer, "swiglu_limit", None),
658	layer=layer,
659	)
660
661	def apply_monolithic(
662	self,
663	layer: RoutedExperts,
664	x: torch.Tensor,
665	router_logits: torch.Tensor,
666	input_ids: torch.Tensor \| None = None,
667	) -> torch.Tensor:
668	assert self.is_monolithic
669	assert self.moe_kernel is not None
670	return self.moe_kernel.apply_monolithic(
671	x,
672	layer.w13_weight,
673	layer.w2_weight,
674	router_logits,
675	activation=layer.activation,
676	global_num_experts=layer.global_num_experts,
677	expert_map=layer.expert_map,
678	apply_router_weight_on_input=layer.apply_router_weight_on_input,
679	num_expert_group=layer.num_expert_group,
680	topk_group=layer.topk_group,
681	e_score_correction_bias=layer.e_score_correction_bias,
682	routed_scaling_factor=layer.routed_scaling_factor,
683	)
684
685	def apply(
686	self,
687	layer: RoutedExperts,
688	x: torch.Tensor,
689	topk_weights: torch.Tensor,
690	topk_ids: torch.Tensor,
691	shared_experts: SharedExperts \| None,
692	shared_experts_input: torch.Tensor \| None,
693	) -> torch.Tensor \| UnfinalizedMoEOutput:
694	assert not self.is_monolithic
695	assert self.moe_kernel is not None
696	return self.moe_kernel.apply(
697	x,
698	layer.w13_weight,
699	layer.w2_weight,
700	topk_weights,
701	topk_ids,
702	activation=layer.activation,
703	global_num_experts=layer.global_num_experts,
704	expert_map=layer.expert_map,
705	apply_router_weight_on_input=layer.apply_router_weight_on_input,
706	shared_experts=shared_experts,
707	shared_experts_input=shared_experts_input,
708	)
709
710
711	ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod
712	ModelOptFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod
713
714
715	class ModelOptNvFp4Config(ModelOptQuantConfigBase):
716	"""Config class for ModelOpt FP4."""
717
718	def __init__(
719	self,
720	quant_method: str = "NVFP4",
721	is_checkpoint_nvfp4_serialized: bool = False,
722	kv_cache_quant_algo: str \| None = None,
723	exclude_modules: list[str] \| None = None,
724	group_size: int = 16,
725	) -> None:
726	if exclude_modules is None:
727	exclude_modules = []
728	super().__init__(exclude_modules)
729	self.quant_method = quant_method
730	self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
731	if is_checkpoint_nvfp4_serialized:
732	self.group_size = group_size
733	self.kv_cache_quant_algo = kv_cache_quant_algo
734
735	# Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch.
736	# NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization
737	# W4A16_NVFP4 -> W4A16: FP4 Marlin GEMM with bf16/fp16 activations
738	supported = algos_owned_by("modelopt_fp4")
739	if quant_method not in supported:
740	raise ValueError(
741	f"Unsupported ModelOpt NVFP4 quant_algo: {quant_method}. "
742	f"Supported: {' / '.join(supported)}."
743	)
744
745	def get_name(self) -> QuantizationMethods:
746	return "modelopt_fp4"
747
748	def get_supported_act_dtypes(self) -> list[torch.dtype]:
749	return [torch.bfloat16, torch.half, torch.float8_e4m3fn]
750
751	@classmethod
752	def get_min_capability(cls) -> int:
753	return 75
754
755	@classmethod
756	def override_quantization_method(
757	cls, hf_quant_cfg, user_quant, hf_config=None
758	) -> QuantizationMethods \| None:
759	algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)
760	if algo is not None and ("NVFP4" in algo or "FP4" in algo):
761	return "modelopt_fp4"
762	return None
763
764	@classmethod
765	def _from_config(
766	cls,
767	*,
768	quant_method: str,
769	kv_cache_quant_method: str \| None,
770	exclude_modules: list[str],
771	original_config: dict[str, Any],
772	group_size: int \| None,
773	**kwargs: Any,
774	) -> "ModelOptNvFp4Config":
775	is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method
776
777	# A checkpoint can declare quant_algo "NVFP4" yet quantize weights
778	# only (input_activations is null in every config group). The W4A4
779	# MoE path then folds an uninitialized input_scale that was never
780	# loaded and can zero every expert, so route it through W4A16.
781	if quant_method == "NVFP4":
782	config_groups = original_config.get("config_groups")
783	if (
784	isinstance(config_groups, dict)
785	and config_groups
786	and all(
787	isinstance(group, dict) and group.get("input_activations") is None
788	for group in config_groups.values()
789	)
790	):
791	quant_method = "W4A16_NVFP4"
792
793	if group_size is None:
794	group_size = 16 # Default value
795
796	# For FP4, these fields are required
797	if is_checkpoint_nvfp4_serialized and "quantization" in original_config:
798	# Check if required fields are present in the quantization config
799	quant_config = original_config["quantization"]
800	required_fields = ["group_size", "kv_cache_quant_algo", "exclude_modules"]
801	missing_fields = [
802	field for field in required_fields if field not in quant_config
803	]
804	if missing_fields:
805	raise ValueError(
806	f"NVFP4 quantization requires the following fields in "
807	f"hf_quant_config.json: {missing_fields}"
808	)
809
810	return cls(
811	quant_method,
812	is_checkpoint_nvfp4_serialized,
813	kv_cache_quant_method,
814	exclude_modules,
815	group_size,
816	)
817
818
819	class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
820	"""MoE Method for FP4 Quantization.
821
822	Args:
823	quant_config: NVFP4 Quant Config
824
825	"""
826
827	supports_pre_processed_weights = True
828
829	def __init__(
830	self,
831	quant_config: ModelOptNvFp4Config,
832	moe_config: FusedMoEConfig,
833	) -> None:
834	super().__init__(moe_config)
835	self.quant_config = quant_config
836	# W4A16 mode fires for W4A16_NVFP4 on-disk checkpoints. With
837	# activation_key=None every W4A4 backend's _supports_quant_scheme
838	# rejects itself (they all require (kNvfp4Static, kNvfp4Dynamic)
839	# exactly); only Marlin survives. Marlin's MoE path drops
840	# activation scales in convert_to_nvfp4_moe_kernel_format, so no
841	# other change is needed.
842	self.use_a16 = quant_config.quant_method == "W4A16_NVFP4"
843	self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
844	config=self.moe,
845	weight_key=kNvfp4Static,
846	activation_key=None if self.use_a16 else kNvfp4Dynamic,
847	)
848
849	self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
850	self.nvfp4_backend
851	)
852
853	def uses_weight_scale_2_pattern(self) -> bool:
854	"""FP4 variants use 'weight_scale_2' pattern for per-tensor weight scales."""
855	return True
856
857	def create_weights(
858	self,
859	layer: RoutedExperts,
860	num_experts: int,
861	hidden_size: int,
862	intermediate_size_per_partition: int,
863	params_dtype: torch.dtype,
864	**extra_weight_attrs,
865	):
866	assert self.quant_config.is_checkpoint_nvfp4_serialized
867
868	layer.num_experts = num_experts
869	layer.params_dtype = params_dtype
870	layer.quant_config = self.quant_config
871	weight_dtype = torch.uint8
872	weight_scale_dtype = torch.float8_e4m3fn
873	weight_loader = extra_weight_attrs.get("weight_loader")
874	global_num_experts = extra_weight_attrs.get("global_num_experts")
875	w13_num_shards = 2 if self.moe.is_act_and_mul else 1
876	# GEMM 1
877	w13_weight = ModelWeightParameter(
878	data=torch.empty(
879	num_experts,
880	w13_num_shards * intermediate_size_per_partition,
881	# 2 fp4 items are packed in the input dimension
882	hidden_size // 2,
883	dtype=weight_dtype,
884	),
885	input_dim=1,
886	output_dim=2,
887	weight_loader=weight_loader,
888	)
889	layer.register_parameter("w13_weight", w13_weight)
890
891	# GEMM 2
892	w2_weight = ModelWeightParameter(
893	data=torch.empty(
894	num_experts,
895	hidden_size,
896	# 2 fp4 items are packed in the input dimension
897	intermediate_size_per_partition // 2,
898	dtype=weight_dtype,
899	),
900	input_dim=1,
901	output_dim=2,
902	weight_loader=weight_loader,
903	)
904	layer.register_parameter("w2_weight", w2_weight)
905
906	w13_weight_scale = ModelWeightParameter(
907	data=torch.empty(
908	num_experts,
909	w13_num_shards * intermediate_size_per_partition,
910	# 2 fp4 items are packed in the input dimension
911	hidden_size // self.quant_config.group_size,
912	dtype=weight_scale_dtype,
913	),
914	input_dim=1,
915	output_dim=2,
916	weight_loader=weight_loader,
917	)
918	layer.register_parameter("w13_weight_scale", w13_weight_scale)
919
920	w2_weight_scale = ModelWeightParameter(
921	data=torch.empty(
922	num_experts,
923	hidden_size,
924	# 2 fp4 items are packed in the input dimension
925	intermediate_size_per_partition // self.quant_config.group_size,
926	dtype=weight_scale_dtype,
927	),
928	input_dim=1,
929	output_dim=2,
930	weight_loader=weight_loader,
931	)
932	layer.register_parameter("w2_weight_scale", w2_weight_scale)
933
934	extra_weight_attrs.update(
935	{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
936	)
937
938	w13_weight_scale_2 = PerTensorScaleParameter(
939	data=torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
940	weight_loader=weight_loader,
941	)
942	layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)
943
944	w2_weight_scale_2 = PerTensorScaleParameter(
945	data=torch.empty(num_experts, dtype=torch.float32),
946	weight_loader=weight_loader,
947	)
948	layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2)
949
950	extra_weight_attrs.update(
951	{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
952	)
953
954	global_sf_num_experts = (
955	global_num_experts if self.use_global_sf else num_experts
956	)
957	w13_input_scale = PerTensorScaleParameter(
958	data=torch.empty(
959	global_sf_num_experts,
960	w13_num_shards,
961	dtype=torch.float32,
962	),
963	weight_loader=weight_loader,
964	)
965	layer.register_parameter("w13_input_scale", w13_input_scale)
966
967	w2_input_scale = PerTensorScaleParameter(
968	data=torch.empty(global_sf_num_experts, dtype=torch.float32),
969	weight_loader=weight_loader,
970	)
971	layer.register_parameter("w2_input_scale", w2_input_scale)
972
973	def process_weights_after_loading(self, layer: RoutedExperts) -> None:
974	"""Convert NVFP4 MoE weights into kernel format and setup the kernel."""
975	if is_weights_pre_processed():
976	if self.nvfp4_backend != NvFp4MoeBackend.FLASHINFER_TRTLLM:
977	raise RuntimeError(
978	"pre-processed weights require FLASHINFER_TRTLLM backend, "
979	f"moe backend, got {self.nvfp4_backend}"
980	)
981	self._restore_padded_moe_dims(layer)
982	self._build_moe_kernel(layer)
983	return
984
985	# Use a single gscale for w13.
986	if self.moe.is_act_and_mul and not torch.allclose(
987	layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]
988	):
989	logger.warning_once(
990	"w1_weight_scale_2 must match w3_weight_scale_2. "
991	"Accuracy may be affected."
992	)
993	w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0].contiguous()
994
995	(
996	w13,
997	w13_scale,
998	w13_scale_2,
999	a13_scale,
1000	w2,
1001	w2_scale,
1002	w2_scale_2,
1003	a2_scale,
1004	) = convert_to_nvfp4_moe_kernel_format(
1005	nvfp4_backend=self.nvfp4_backend,
1006	layer=layer,
1007	w13=layer.w13_weight,
1008	w13_scale=layer.w13_weight_scale,
1009	w13_scale_2=w13_weight_scale_2,
1010	a13_scale=layer.w13_input_scale,
1011	w2=layer.w2_weight,
1012	w2_scale=layer.w2_weight_scale,
1013	w2_scale_2=layer.w2_weight_scale_2,
1014	a2_scale=layer.w2_input_scale,
1015	is_act_and_mul=self.moe.is_act_and_mul,
1016	use_a16=self.use_a16,
1017	)
1018
1019	replace_parameter(layer, "w13_weight", w13)
1020	replace_parameter(layer, "w13_weight_scale", w13_scale)
1021	replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
1022	replace_parameter(layer, "w13_input_scale", a13_scale)
1023	replace_parameter(layer, "w2_weight", w2)
1024	replace_parameter(layer, "w2_weight_scale", w2_scale)
1025	replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
1026	replace_parameter(layer, "w2_input_scale", a2_scale)
1027
1028	self._build_moe_kernel(layer)
1029
1030	def _build_moe_kernel(self, layer: RoutedExperts) -> None:
1031	"""Build the modular MoE kernel from the (already in-format) weights."""
1032	self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1033	assert self.experts_cls is not None
1034	self.moe_kernel = make_nvfp4_moe_kernel(
1035	moe_quant_config=self.moe_quant_config,
1036	moe_config=self.moe,
1037	experts_cls=self.experts_cls,
1038	backend=self.nvfp4_backend,
1039	routing_tables=layer._expert_routing_tables(),
1040	)
1041	self.moe_kernel.fused_experts.process_weights_after_loading(layer)
1042
1043	def _restore_padded_moe_dims(self, layer: RoutedExperts) -> None:
1044	"""Recover the padded ``moe_config`` dims from the exported weights."""
1045	mc = layer.moe_config
1046	padded_hidden = layer.w2_weight.shape[1]
1047	if padded_hidden != mc.hidden_dim:
1048	if mc.hidden_dim_unpadded is None:
1049	mc.hidden_dim_unpadded = mc.hidden_dim
1050	mc.hidden_dim = padded_hidden
1051	mc.intermediate_size_per_partition = layer.w2_weight.shape[2] * 2
1052
1053	def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
1054	return make_nvfp4_moe_quant_config(
1055	backend=self.nvfp4_backend,
1056	w13_scale=layer.w13_weight_scale,
1057	w2_scale=layer.w2_weight_scale,
1058	w13_scale_2=layer.w13_weight_scale_2,
1059	w2_scale_2=layer.w2_weight_scale_2,
1060	a13_scale=layer.w13_input_scale,
1061	a2_scale=layer.w2_input_scale,
1062	swiglu_limit=getattr(layer, "swiglu_limit", None),
1063	swiglu_alpha=getattr(layer, "swiglu_alpha", None),
1064	swiglu_beta=getattr(layer, "swiglu_beta", None),
1065	layer=layer,
1066	use_a16=self.use_a16,
1067	)
1068
1069	@property
1070	def supports_eplb(self) -> bool:
1071	return True
1072
1073	def apply_monolithic(
1074	self,
1075	layer: RoutedExperts,
1076	x: torch.Tensor,
1077	router_logits: torch.Tensor,
1078	input_ids: torch.Tensor \| None = None,
1079	) -> torch.Tensor \| UnfinalizedMoEOutput:
1080	assert self.is_monolithic
1081	assert self.moe_kernel is not None
1082	return self.moe_kernel.apply_monolithic(
1083	x,
1084	layer.w13_weight,
1085	layer.w2_weight,
1086	router_logits,
1087	activation=layer.activation,
1088	global_num_experts=layer.global_num_experts,
1089	expert_map=layer.expert_map,
1090	apply_router_weight_on_input=layer.apply_router_weight_on_input,
1091	num_expert_group=layer.num_expert_group,
1092	topk_group=layer.topk_group,
1093	e_score_correction_bias=layer.e_score_correction_bias,
1094	routed_scaling_factor=layer.routed_scaling_factor,
1095	)
1096
1097	def apply(
1098	self,
1099	layer: RoutedExperts,
1100	x: torch.Tensor,
1101	topk_weights: torch.Tensor,
1102	topk_ids: torch.Tensor,
1103	shared_experts: SharedExperts \| None,
1104	shared_experts_input: torch.Tensor \| None,
1105	) -> torch.Tensor \| UnfinalizedMoEOutput:
1106	assert not self.is_monolithic
1107	assert self.moe_kernel is not None
1108	return self.moe_kernel.apply(
1109	x,
1110	layer.w13_weight,
1111	layer.w2_weight,
1112	topk_weights,
1113	topk_ids,
1114	activation=layer.activation,
1115	global_num_experts=layer.global_num_experts,
1116	expert_map=layer.expert_map,
1117	apply_router_weight_on_input=layer.apply_router_weight_on_input,
1118	shared_experts=shared_experts,
1119	shared_experts_input=shared_experts_input,
1120	)
1121
1122
1123	ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE
1124	ModelOptNvFp4Config.KVCacheMethodCls = ModelOptKVCacheMethod
1125
1126
1127	class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1128	"""Config class for ModelOpt MXFP8."""
1129
1130	def __init__(
1131	self,
1132	is_checkpoint_mxfp8_serialized: bool,
1133	kv_cache_quant_algo: str \| None,
1134	exclude_modules: list[str],
1135	) -> None:
1136	super().__init__(exclude_modules)
1137	self.quant_method = "MXFP8"
1138	self.is_checkpoint_mxfp8_serialized = is_checkpoint_mxfp8_serialized
1139
1140	if not is_checkpoint_mxfp8_serialized:
1141	raise ValueError(
1142	"MXFP8 quantization requires a serialized checkpoint. "
1143	"Dynamic quantization is not supported."
1144	)
1145
1146	self.kv_cache_quant_algo = kv_cache_quant_algo
1147
1148	def get_name(self) -> QuantizationMethods:
1149	return "modelopt_mxfp8"
1150
1151	def get_supported_act_dtypes(self) -> list[torch.dtype]:
1152	return [torch.bfloat16]
1153
1154	@classmethod
1155	def get_min_capability(cls) -> int:
1156	# Marlin kernel supports MXFP8 on SM80+
1157	return 80
1158
1159	@classmethod
1160	def override_quantization_method(
1161	cls, hf_quant_cfg, user_quant, hf_config=None
1162	) -> QuantizationMethods \| None:
1163	algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)
1164	if algo is not None and "MXFP8" in algo:
1165	return "modelopt_mxfp8"
1166	return None
1167
1168	@classmethod
1169	def from_config(cls, config: dict[str, Any]) -> "ModelOptMxFp8Config":
1170	# MiniMax-style checkpoints tag `quant_method: "mxfp8"` + `ignored_layers`
1171	# (same on-disk format as ModelOpt MXFP8); normalize to the ModelOpt
1172	# schema and reuse the shared parser.
1173	if "quantization" not in config and not config.get("quant_algo"):
1174	config = {
1175	"quant_method": "modelopt",
1176	"quantization": {
1177	"quant_algo": "MXFP8",
1178	"kv_cache_quant_algo": config.get("kv_cache_quant_algo"),
1179	"exclude_modules": config.get("ignored_layers", []) or [],
1180	},
1181	}
1182	return cast("ModelOptMxFp8Config", super().from_config(config))
1183
1184	@classmethod
1185	def _from_config(
1186	cls,
1187	*,
1188	quant_method: str,
1189	kv_cache_quant_method: str \| None,
1190	exclude_modules: list[str],
1191	original_config: dict[str, Any],
1192	**kwargs: Any,
1193	) -> "ModelOptMxFp8Config":
1194	is_checkpoint_mxfp8_serialized = "MXFP8" in quant_method.upper()
1195
1196	# For MXFP8, validate required fields in the config
1197	if is_checkpoint_mxfp8_serialized and "quantization" in original_config:
1198	quant_config = original_config["quantization"]
1199	required_fields = ["kv_cache_quant_algo", "exclude_modules"]
1200	missing_fields = [
1201	field for field in required_fields if field not in quant_config
1202	]
1203	if missing_fields:
1204	raise ValueError(
1205	f"MXFP8 quantization requires the following fields in "
1206	f"hf_quant_config.json: {missing_fields}"
1207	)
1208
1209	return cls(
1210	is_checkpoint_mxfp8_serialized,
1211	kv_cache_quant_method,
1212	exclude_modules,
1213	)
1214
1215
1216	class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1217	"""FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints."""
1218
1219	def __init__(
1220	self,
1221	quant_config: ModelOptMxFp8Config,
1222	moe_config: FusedMoEConfig,
1223	) -> None:
1224	super().__init__(moe_config)
1225	self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
1226	self.quant_config = quant_config
1227	assert self.quant_config.is_checkpoint_mxfp8_serialized
1228
1229	self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
1230
1231	def create_weights(
1232	self,
1233	layer: RoutedExperts,
1234	num_experts: int,
1235	hidden_size: int,
1236	intermediate_size_per_partition: int,
1237	params_dtype: torch.dtype,
1238	**extra_weight_attrs,
1239	):
1240	assert layer.intermediate_size_per_partition == intermediate_size_per_partition
1241	assert layer.hidden_size == hidden_size
1242	layer.orig_dtype = params_dtype
1243
1244	if hidden_size % MXFP8_BLOCK_SIZE != 0:
1245	raise ValueError(
1246	f"MXFP8 MoE requires hidden_size divisible by {MXFP8_BLOCK_SIZE}, "
1247	f"got {hidden_size}."
1248	)
1249	if intermediate_size_per_partition % MXFP8_BLOCK_SIZE != 0:
1250	raise ValueError(
1251	"MXFP8 MoE requires intermediate_size_per_partition divisible by "
1252	f"{MXFP8_BLOCK_SIZE}, got {intermediate_size_per_partition}."
1253	)
1254
1255	layer.num_experts = num_experts
1256	weight_loader = extra_weight_attrs.get("weight_loader")
1257	w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1258
1259	# GEMM 1 weights: [E, (2I or I), H]
1260	w13_weight = ModelWeightParameter(
1261	data=torch.empty(
1262	num_experts,
1263	w13_num_shards * intermediate_size_per_partition,
1264	hidden_size,
1265	dtype=MXFP8_VALUE_DTYPE,
1266	),
1267	input_dim=2,
1268	output_dim=1,
1269	weight_loader=weight_loader,
1270	)
1271	layer.register_parameter("w13_weight", w13_weight)
1272
1273	# GEMM 2 weights: [E, H, I]
1274	w2_weight = ModelWeightParameter(
1275	data=torch.empty(
1276	num_experts,
1277	hidden_size,
1278	intermediate_size_per_partition,
1279	dtype=MXFP8_VALUE_DTYPE,
1280	),
1281	input_dim=2,
1282	output_dim=1,
1283	weight_loader=weight_loader,
1284	)
1285	layer.register_parameter("w2_weight", w2_weight)
1286
1287	# Per-block (K=32) E8M0 scales.
1288	w13_weight_scale = ModelWeightParameter(
1289	data=torch.empty(
1290	num_experts,
1291	w13_num_shards * intermediate_size_per_partition,
1292	hidden_size // MXFP8_BLOCK_SIZE,
1293	dtype=MXFP8_SCALE_DTYPE,
1294	),
1295	input_dim=2,
1296	output_dim=1,
1297	weight_loader=weight_loader,
1298	)
1299	layer.register_parameter("w13_weight_scale", w13_weight_scale)
1300
1301	w2_weight_scale = ModelWeightParameter(
1302	data=torch.empty(
1303	num_experts,
1304	hidden_size,
1305	intermediate_size_per_partition // MXFP8_BLOCK_SIZE,
1306	dtype=MXFP8_SCALE_DTYPE,
1307	),
1308	input_dim=2,
1309	output_dim=1,
1310	weight_loader=weight_loader,
1311	)
1312	layer.register_parameter("w2_weight_scale", w2_weight_scale)
1313
1314	# Ensure the generic MoE weight-loader treats these as block scales.
1315	set_weight_attrs(
1316	layer.w13_weight_scale,
1317	{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value},
1318	)
1319	set_weight_attrs(
1320	layer.w2_weight_scale,
1321	{"quant_method": FusedMoeWeightScaleSupported.BLOCK.value},
1322	)
1323
1324	@staticmethod
1325	def _check_weight_dtypes(layer: torch.nn.Module) -> None:
1326	"""Validate weight and scale dtypes before processing."""
1327	expected = {
1328	"w13_weight": MXFP8_VALUE_DTYPE,
1329	"w2_weight": MXFP8_VALUE_DTYPE,
1330	"w13_weight_scale": MXFP8_SCALE_DTYPE,
1331	"w2_weight_scale": MXFP8_SCALE_DTYPE,
1332	}
1333	for name, expected_dtype in expected.items():
1334	actual = getattr(layer, name).dtype
1335	if actual != expected_dtype:
1336	raise ValueError(
1337	f"Expected {name} dtype {expected_dtype}, got {actual}."
1338	)
1339
1340	def _dequant_mxfp8_weights_to_bf16(self, layer: RoutedExperts) -> None:
1341	"""One-time MXFP8->BF16 weight dequant for the emulation path.
1342
1343	On devices without a native MXFP8 MoE kernel (e.g. gfx942 / MI300),
1344	``Mxfp8EmulationTritonExperts`` otherwise dequantizes every expert
1345	weight to BF16 on every forward step -- the dominant cost (conc1
1346	~1.3 tok/s). Doing the dequant once here and replacing the MXFP8
1347	parameters with BF16 makes the MoE run exactly like a plain BF16
1348	checkpoint (full precision, no per-step dequant); SwiGLU-OAI is still
1349	applied by the experts' ``activation()`` override. The MXFP8 weights
1350	are freed by ``replace_parameter`` (BF16 is 2x their size; the small
1351	E8M0 scale tensors are left in place, unused).
1352	"""
1353	from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
1354	dequant_mxfp8_to_bf16,
1355	)
1356
1357	target_dtype = getattr(layer, "orig_dtype", torch.bfloat16)
1358	num_experts = layer.w13_weight.shape[0]
1359
1360	# dequant_mxfp8_to_bf16 handles arbitrary leading dims (*x.shape[:-1]),
1361	# so dequant the whole [E, N, K] weight in one vectorized call.
1362	w13_bf16 = dequant_mxfp8_to_bf16(layer.w13_weight, layer.w13_weight_scale).to(
1363	target_dtype
1364	)
1365	w2_bf16 = dequant_mxfp8_to_bf16(layer.w2_weight, layer.w2_weight_scale).to(
1366	target_dtype
1367	)
1368
1369	replace_parameter(layer, "w13_weight", w13_bf16)
1370	replace_parameter(layer, "w2_weight", w2_bf16)
1371
1372	logger.info_once(
1373	"MXFP8->BF16 load-time dequant complete (%d experts/layer); MoE "
1374	"now runs in BF16 with no per-step dequant.",
1375	num_experts,
1376	)
1377
1378	def process_weights_after_loading(self, layer: RoutedExperts) -> None:
1379	# TODO(bnell): why is this required only for mxfp8?
1380	if getattr(layer, "_already_called_process_weights_after_loading", False):
1381	return
1382	layer._already_called_process_weights_after_loading = True
1383
1384	self._check_weight_dtypes(layer)
1385
1386	layer.weight_block_size = self.weight_block_size
1387
1388	w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
1389	fp8_backend=self.mxfp8_backend,
1390	layer=layer,
1391	w13=layer.w13_weight,
1392	w2=layer.w2_weight,
1393	w13_scale=layer.w13_weight_scale,
1394	w2_scale=layer.w2_weight_scale,
1395	w13_input_scale=None,
1396	w2_input_scale=None,
1397	)
1398
1399	replace_parameter(layer, "w13_weight", w13)
1400	replace_parameter(layer, "w2_weight", w2)
1401	replace_parameter(layer, "w13_weight_scale", w13_scale)
1402	replace_parameter(layer, "w2_weight_scale", w2_scale)
1403
1404	self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1405	assert self.moe_quant_config is not None
1406	assert self.experts_cls is not None
1407	self.moe_kernel = make_fp8_moe_kernel(
1408	moe_quant_config=self.moe_quant_config,
1409	moe_config=self.moe,
1410	fp8_backend=self.mxfp8_backend,
1411	experts_cls=self.experts_cls,
1412	routing_tables=layer._expert_routing_tables(),
1413	)
1414
1415	# No native MXFP8 MoE kernel on this device (e.g. gfx942): the emulation
1416	# experts would dequant MXFP8->BF16 every forward step. Convert the
1417	# weights to BF16 once, here, so the MoE runs like a BF16 checkpoint.
1418	# Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the 1-byte
1419	# MXFP8 weights and dequant per-step (~half the memory, much slower).
1420	if (
1421	self.mxfp8_backend == Fp8MoeBackend.EMULATION
1422	and envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD
1423	):
1424	self._dequant_mxfp8_weights_to_bf16(layer)
1425
1426	def get_fused_moe_quant_config(
1427	self, layer: RoutedExperts
1428	) -> FusedMoEQuantConfig \| None:
1429	return make_fp8_moe_quant_config(
1430	fp8_backend=self.mxfp8_backend,
1431	w1_scale=layer.w13_weight_scale,
1432	w2_scale=layer.w2_weight_scale,
1433	a1_scale=None,
1434	a2_scale=None,
1435	block_shape=self.weight_block_size,
1436	swiglu_limit=getattr(layer, "swiglu_limit", None),
1437	gemm1_alpha=getattr(layer, "swiglu_alpha", None),
1438	gemm1_beta=getattr(layer, "swiglu_beta", None),
1439	layer=layer,
1440	)
1441
1442	def apply_monolithic(
1443	self,
1444	layer: RoutedExperts,
1445	x: torch.Tensor,
1446	router_logits: torch.Tensor,
1447	input_ids: torch.Tensor \| None = None,
1448	) -> torch.Tensor:
1449	assert self.is_monolithic
1450	assert self.moe_kernel is not None
1451	return self.moe_kernel.apply_monolithic(
1452	x,
1453	layer.w13_weight,
1454	layer.w2_weight,
1455	router_logits,
1456	activation=layer.activation,
1457	global_num_experts=layer.global_num_experts,
1458	expert_map=layer.expert_map,
1459	apply_router_weight_on_input=layer.apply_router_weight_on_input,
1460	num_expert_group=layer.num_expert_group,
1461	topk_group=layer.topk_group,
1462	e_score_correction_bias=layer.e_score_correction_bias,
1463	routed_scaling_factor=layer.routed_scaling_factor,
1464	)
1465
1466	def apply(
1467	self,
1468	layer: RoutedExperts,
1469	x: torch.Tensor,
1470	topk_weights: torch.Tensor,
1471	topk_ids: torch.Tensor,
1472	shared_experts: SharedExperts \| None,
1473	shared_experts_input: torch.Tensor \| None,
1474	) -> torch.Tensor \| UnfinalizedMoEOutput:
1475	assert not self.is_monolithic
1476	assert self.moe_kernel is not None
1477	return self.moe_kernel.apply(
1478	x,
1479	layer.w13_weight,
1480	layer.w2_weight,
1481	topk_weights,
1482	topk_ids,
1483	activation=layer.activation,
1484	global_num_experts=layer.global_num_experts,
1485	expert_map=layer.expert_map,
1486	apply_router_weight_on_input=layer.apply_router_weight_on_input,
1487	shared_experts=shared_experts,
1488	shared_experts_input=shared_experts_input,
1489	)
1490
1491
1492	# Register the method classes for ModelOptMxFp8Config
1493	ModelOptMxFp8Config.FusedMoEMethodCls = ModelOptMxFp8FusedMoE
1494	ModelOptMxFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod
1495
1496
1497	class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1498	"""Config class for ModelOpt MIXED_PRECISION.
1499
1500	Supports checkpoints where different layers use different quantization
1501	algorithms (e.g., FP8 for dense layers and NVFP4 for MoE experts).
1502	The per-layer algorithm is specified in the ``quantized_layers`` dict
1503	inside ``config.json``'s ``quantization_config`` (preferred) or the
1504	legacy ``hf_quant_config.json``.
1505	"""
1506
1507	def __init__(
1508	self,
1509	kv_cache_quant_method: str \| None,
1510	exclude_modules: list[str],
1511	quantized_layers: dict[str, dict[str, Any]],
1512	fp8_config: ModelOptFp8Config,
1513	nvfp4_config: ModelOptNvFp4Config,
1514	w4a16_nvfp4_config: ModelOptNvFp4Config,
1515	mxfp8_config: ModelOptMxFp8Config,
1516	) -> None:
1517	super().__init__(exclude_modules)
1518	self.kv_cache_quant_method = kv_cache_quant_method
1519	self.quantized_layers = quantized_layers
1520	self.fp8_config = fp8_config
1521	self.nvfp4_config = nvfp4_config
1522	self.w4a16_nvfp4_config = w4a16_nvfp4_config
1523	self.mxfp8_config = mxfp8_config
1524
1525	block_sizes = {
1526	int(layer_info.get("group_size", 128))
1527	for layer_info in quantized_layers.values()
1528	if layer_info.get("quant_algo", "").upper() in _BLOCK_FP8_MOE_ALGOS
1529	}
1530	if len(block_sizes) > 1:
1531	raise ValueError(
1532	"MIXED_PRECISION currently requires all block-FP8 MoE layers "
1533	f"to use one group_size, got {sorted(block_sizes)}."
1534	)
1535	block_size = next(iter(block_sizes), 128)
1536	self.fp8_block_config = Fp8Config(
1537	is_checkpoint_fp8_serialized=True,
1538	activation_scheme="dynamic",
1539	weight_block_size=[block_size, block_size],
1540	)
1541
1542	def has_blocked_weights(self) -> bool:
1543	# Same gate as ModelOptFp8Config.has_blocked_weights, resolved per
1544	# layer: "+quant_fp8" must be on as soon as any layer is block-scaled.
1545	return any(
1546	info.get("quant_algo", "").upper() in _BLOCK_FP8_MOE_ALGOS
1547	for info in self.quantized_layers.values()
1548	)
1549
1550	def get_name(self) -> QuantizationMethods:
1551	return "modelopt_mixed"
1552
1553	def get_supported_act_dtypes(self) -> list[torch.dtype]:
1554	return [torch.bfloat16, torch.half]
1555
1556	@classmethod
1557	def get_min_capability(cls) -> int:
1558	# Turing and up (SM75+): NVFP4 routed experts run via Marlin W4A16
1559	# (SM75+), FP8 weight-only dense via MarlinFP8 (cc>=7.5), and FP8 MoE,
1560	# if present, via Marlin (TritonExperts gates its FP8 schemes behind
1561	# supports_fp8(), cc>=89). None of these paths require native FP8 tensor
1562	# cores, so SM75 is sufficient. Validated end-to-end on a Tesla T4
1563	# (SM75) and A100 (SM80). Pairs with the FlashInfer attention SM80
1564	# lower bound so SM75 auto-selects a supported attention backend.
1565	return 75
1566
1567	@classmethod
1568	def override_quantization_method(
1569	cls, hf_quant_cfg, user_quant, hf_config=None
1570	) -> QuantizationMethods \| None:
1571	algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)
1572	if algo is not None and algo == "MIXED_PRECISION":
1573	return "modelopt_mixed"
1574	return None
1575
1576	@classmethod
1577	def _from_config(
1578	cls,
1579	*,
1580	quant_method: str,
1581	kv_cache_quant_method: str \| None,
1582	exclude_modules: list[str],
1583	original_config: dict[str, Any],
1584	group_size: int \| None,
1585	**kwargs: Any,
1586	) -> "ModelOptMixedPrecisionConfig":
1587	if "quantization" in original_config:
1588	quantized_layers = original_config["quantization"].get(
1589	"quantized_layers", {}
1590	)
1591	else:
1592	quantized_layers = original_config.get("quantized_layers", {})
1593
1594	if not quantized_layers:
1595	raise ValueError(
1596	"MIXED_PRECISION quant_algo requires a non-empty "
1597	"'quantized_layers' mapping in the quantization config."
1598	)
1599
1600	# Determine group_size from the first NVFP4-family entry if not
1601	# provided. Both NVFP4 (W4A4) and W4A16_NVFP4 share the same packing
1602	# + group-size convention; either entry resolves the value.
1603	if group_size is None:
1604	for layer_info in quantized_layers.values():
1605	if layer_info.get("quant_algo", "").upper() in (
1606	"NVFP4",
1607	"W4A16_NVFP4",
1608	):
1609	group_size = layer_info.get("group_size", 16)
1610	break
1611	if group_size is None:
1612	group_size = 16
1613
1614	fp8_config = ModelOptFp8Config(
1615	quant_method="FP8",
1616	is_checkpoint_fp8_serialized=True,
1617	kv_cache_quant_method=kv_cache_quant_method,
1618	exclude_modules=[],
1619	)
1620	nvfp4_config = ModelOptNvFp4Config(
1621	is_checkpoint_nvfp4_serialized=True,
1622	kv_cache_quant_algo=kv_cache_quant_method,
1623	exclude_modules=[],
1624	group_size=group_size,
1625	)
1626	# Sibling config for layers that declare quant_algo: "W4A16_NVFP4".
1627	# get_quant_method resolves this sub-config to the (kNvfp4Static, None)
1628	# QuantSpec (W4A16) for the generic ModelOptLinearMethod. The MoE side
1629	# reads quant_config.quant_method == "W4A16_NVFP4" to set use_a16 →
1630	# Marlin backend in ModelOptNvFp4FusedMoE.__init__.
1631	w4a16_nvfp4_config = ModelOptNvFp4Config(
1632	quant_method="W4A16_NVFP4",
1633	is_checkpoint_nvfp4_serialized=True,
1634	kv_cache_quant_algo=kv_cache_quant_method,
1635	exclude_modules=[],
1636	group_size=group_size,
1637	)
1638
1639	mxfp8_config = ModelOptMxFp8Config(
1640	is_checkpoint_mxfp8_serialized=True,
1641	kv_cache_quant_algo=kv_cache_quant_method,
1642	exclude_modules=[],
1643	)
1644
1645	return cls(
1646	kv_cache_quant_method=kv_cache_quant_method,
1647	exclude_modules=exclude_modules,
1648	quantized_layers=quantized_layers,
1649	fp8_config=fp8_config,
1650	nvfp4_config=nvfp4_config,
1651	w4a16_nvfp4_config=w4a16_nvfp4_config,
1652	mxfp8_config=mxfp8_config,
1653	)
1654
1655	def _resolve_quant_algo(self, prefix: str) -> str \| None:
1656	"""Look up the quant_algo for a vLLM-side layer prefix.
1657
1658	Tries three strategies in order:
1659	1. Direct lookup in ``quantized_layers``.
1660	2. Packed/fused-layer lookup (unfuse via ``packed_modules_mapping``).
1661	3. Prefix-based lookup for RoutedExperts (any child key starts with
1662	``prefix + "."``).
1663
1664	Returns the upper-cased quant_algo string, or None if the prefix
1665	is not found.
1666	"""
1667	# 1. Direct lookup
1668	for candidate in self._quantized_layer_prefix_candidates(prefix):
1669	if candidate in self.quantized_layers:
1670	return self.quantized_layers[candidate]["quant_algo"].upper()
1671
1672	# 2. Packed / fused layer lookup
1673	proj_name = prefix.rsplit(".", 1)[-1]
1674	if self.packed_modules_mapping and proj_name in self.packed_modules_mapping:
1675	algos: set[str] = set()
1676	base = prefix.rsplit(".", 1)[0]
1677	for base_candidate in self._quantized_layer_prefix_candidates(base):
1678	for shard_name in self.packed_modules_mapping[proj_name]:
1679	shard_prefix = f"{base_candidate}.{shard_name}"
1680	if shard_prefix in self.quantized_layers:
1681	algos.add(
1682	self.quantized_layers[shard_prefix]["quant_algo"].upper()
1683	)
1684	if len(algos) == 1:
1685	return algos.pop()
1686	if len(algos) > 1:
1687	raise ValueError(
1688	f"Mixed quant_algo within fused layer {prefix}: "
1689	f"{algos}. All shards must use the same quantization."
1690	)
1691
1692	# 3. Prefix-based lookup (for RoutedExperts / parent modules)
1693	for candidate in self._quantized_layer_prefix_candidates(prefix):
1694	prefix_dot = candidate + "."
1695	for key, info in self.quantized_layers.items():
1696	if key.startswith(prefix_dot):
1697	return info["quant_algo"].upper()
1698
1699	# RoutedExperts expert prefix is e.g. "...moe.experts", while ModelOpt's
1700	# quantized_layers entries use "...moe.gate_proj" / "...moe.up_proj".
1701	if prefix.endswith(".experts"):
1702	parent_dot = prefix.rsplit(".experts", 1)[0] + "."
1703	for key, info in self.quantized_layers.items():
1704	if key.startswith(parent_dot):
1705	return info["quant_algo"].upper()
1706
1707	# 4. Parent-prefix fallback for fused projections whose config lists
1708	# shard names instead of vLLM's packed module name.
1709	fused_projection_shards = {
1710	"qkv_proj": ("q_proj", "k_proj", "v_proj"),
1711	"gate_up_proj": ("gate_proj", "up_proj"),
1712	}
1713	shard_names = fused_projection_shards.get(proj_name)
1714	if shard_names is not None:
1715	for candidate in self._quantized_layer_prefix_candidates(prefix):
1716	parent_dot = candidate.rsplit(".", 1)[0] + "."
1717	shard_algos: set[str] = set()
1718	for shard_name in shard_names:
1719	shard_prefix = f"{parent_dot}{shard_name}"
1720	if shard_prefix in self.quantized_layers:
1721	algo = self.quantized_layers[shard_prefix]["quant_algo"].upper()
1722	shard_algos.add(algo)
1723	if len(shard_algos) == 1:
1724	return shard_algos.pop()
1725	if len(shard_algos) > 1:
1726	raise ValueError(
1727	f"Mixed quant_algo within fused layer {prefix}: "
1728	f"{shard_algos}. All shards must use the same quantization."
1729	)
1730
1731	return None
1732
1733	@staticmethod
1734	def _quantized_layer_prefix_candidates(prefix: str) -> tuple[str, ...]:
1735	candidates = [prefix]
1736
1737	if prefix.endswith(".lm_head"):
1738	candidates.append("lm_head")
1739
1740	if prefix.startswith("language_model.model."):
1741	candidates.append(
1742	"model.language_model." + prefix[len("language_model.model.") :]
1743	)
1744	elif prefix.startswith("model.language_model."):
1745	candidates.append(
1746	"language_model.model." + prefix[len("model.language_model.") :]
1747	)
1748
1749	return tuple(dict.fromkeys(candidates))
1750
1751	def get_quant_method(
1752	self, layer: torch.nn.Module, prefix: str
1753	) -> "QuantizeMethodBase \| None":
1754	"""Return quantize-method based on layer."""
1755	# KV-cache quantization
1756	if isinstance(layer, Attention):
1757	if self.kv_cache_quant_method:
1758	return ModelOptKVCacheMethod(self)
1759	return None
1760
1761	# Excluded layers
1762	if self.is_layer_excluded(prefix):
1763	if isinstance(layer, (LinearBase, ParallelLMHead)):
1764	return UnquantizedLinearMethod()
1765	return None
1766
1767	quant_algo = self._resolve_quant_algo(prefix)
1768
1769	if isinstance(layer, (LinearBase, ParallelLMHead)):
1770	# Per-prefix algo -> its sub-config, then the generic linear method.
1771	# resolve() reads group_size off whichever sub-config it is handed.
1772	# Read-only — never write back, or a linear-only change leaks into
1773	# the MoE method that shares the sub-config.
1774	if quant_algo is None or quant_algo not in LINEAR_ALGOS:
1775	# Layer not in quantized_layers — leave unquantized
1776	return UnquantizedLinearMethod()
1777	_, subcfg_attr = LINEAR_ALGOS[quant_algo]
1778	return build_linear_method(getattr(self, subcfg_attr), quant_algo, prefix)
1779
1780	if isinstance(layer, RoutedExperts):
1781	if quant_algo in _BLOCK_FP8_MOE_ALGOS:
1782	return Fp8MoEMethod(self.fp8_block_config, layer)
1783	if quant_algo == "FP8":
1784	return ModelOptFp8MoEMethod(
1785	quant_config=self.fp8_config,
1786	moe_config=layer.moe_config,
1787	)
1788	if quant_algo == "NVFP4":
1789	return ModelOptNvFp4FusedMoE(
1790	quant_config=self.nvfp4_config,
1791	moe_config=layer.moe_config,
1792	)
1793	if quant_algo == "W4A16_NVFP4":
1794	return ModelOptNvFp4FusedMoE(
1795	quant_config=self.w4a16_nvfp4_config,
1796	moe_config=layer.moe_config,
1797	)
1798	if quant_algo == "MXFP8":
1799	return ModelOptMxFp8FusedMoE(
1800	quant_config=self.mxfp8_config,
1801	moe_config=layer.moe_config,
1802	)
1803	return None
1804
1805	return None
1806
1807	def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
1808	super().apply_vllm_mapper(hf_to_vllm_mapper)
1809	if self.quantized_layers:
1810	self.quantized_layers = hf_to_vllm_mapper.apply_dict(self.quantized_layers)
1811
1812
1813	# ===========================================================================
1814	# Generic QuantKey-driven linear method
1815	#
1816	# ``ModelOptLinearMethod`` serves every ModelOpt linear format. It composes a
1817	# per-QuantKey weight scheme + activation scheme (the ``QuantSpec`` pair that
1818	# ``resolve`` produces from the checkpoint's quant_algo), runs a fixed
1819	# create/process lifecycle, selects the kernel from the pair, and applies.
1820	#
1821	# Adding a format:
1822	# * Composes as a (weight, activation) key pair -> add a QuantKeyScheme per
1823	# new key to SCHEME_FOR, plus a resolve() row returning the QuantSpec. No
1824	# new method class; this is how every current format is built.
1825	# * Needs format-wide state but the same lifecycle -> also return a
1826	# FormatScheme subclass from that resolve() row (extra_weights / pre_process
1827	# / post_process hooks).
1828	# * Needs a different lifecycle entirely -> write a LinearMethodBase and
1829	# register it in LINEAR_METHOD_BUILDERS by algo.
1830	# In all cases accept the algo in the owning config's quant_method validation,
1831	# so get_quant_method can dispatch it.
1832	# ===========================================================================
1833
1834
1835	class Role(Enum):
1836	WEIGHT = "weight"
1837	ACT = "activation"
1838
1839
1840	WEIGHT = Role.WEIGHT
1841	ACT = Role.ACT
1842
1843
1844	@dataclass(frozen=True)
1845	class CkptCtx:
1846	"""Per-checkpoint facts a QuantKey cannot carry."""
1847
1848	group_size: int \| None = None
1849	scale_block_size: tuple[int, int] \| None = None
1850
1851
1852	@dataclass(frozen=True)
1853	class RuntimeDtypes:
1854	"""Runtime/model dtypes kernels need — format-agnostic."""
1855
1856	input_dtype: torch.dtype
1857	out_dtype: torch.dtype
1858	marlin_input_dtype: torch.dtype \| None = None
1859
1860
1861	@dataclass(frozen=True)
1862	class Shapes:
1863	"""Layer geometry from create_weights args."""
1864
1865	output_partition_sizes: list[int]
1866	input_size_per_partition: int
1867	params_dtype: torch.dtype
1868
1869	@property
1870	def output_size_per_partition(self) -> int:
1871	return sum(self.output_partition_sizes)
1872
1873	@property
1874	def num_partitions(self) -> int:
1875	return len(self.output_partition_sizes)
1876
1877
1878	class QuantKeyScheme:
1879	"""One scheme per QuantKey. Selected by key content; the base supplies the
1880	``role`` from the slot it fills the key into (wkey->WEIGHT, akey->ACT).
1881
1882	Every scheme branches on ``role`` explicitly and ``reject``s any role it has
1883	not validated — never falling through to the wrong role's registration,
1884	which would allocate the wrong parameters and produce garbage silently.
1885	"""
1886
1887	key: QuantKey
1888
1889	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1890	raise NotImplementedError
1891
1892	def process(self, layer, role) -> None:
1893	pass
1894
1895	@staticmethod
1896	def reject(role) -> None:
1897	raise NotImplementedError(
1898	f"role {role!r} not validated for this QuantKey scheme"
1899	)
1900
1901	@staticmethod
1902	def register_params(layer, name, shape, dtype, cls, wl, , init=None, *dims):
1903	p = cls(data=torch.empty(shape, dtype=dtype), weight_loader=wl, **dims)
1904	if init is not None:
1905	p.data.fill_(init)
1906	layer.register_parameter(name, p)
1907
1908
1909	class KNvfp4Static(QuantKeyScheme):
1910	"""NVFP4 weight scheme (W4A4 and W4A16 share it). Weight-role only today."""
1911
1912	key = kNvfp4Static
1913
1914	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1915	if role is not WEIGHT:
1916	self.reject(role)
1917	if shapes.input_size_per_partition % 16 != 0:
1918	raise ValueError(
1919	"Unsupported model when in features size is not multiple of 16"
1920	)
1921	# Packed NVFP4 weight: 2 fp4 items per byte along the input dim.
1922	self.register_params(
1923	layer,
1924	"weight",
1925	(shapes.output_size_per_partition, shapes.input_size_per_partition // 2),
1926	torch.uint8,
1927	ModelWeightParameter,
1928	wl,
1929	input_dim=1,
1930	output_dim=0,
1931	)
1932	# Per-tensor global weight scale.
1933	self.register_params(
1934	layer,
1935	"weight_scale_2",
1936	(shapes.num_partitions,),
1937	torch.float32,
1938	PerTensorScaleParameter,
1939	wl,
1940	)
1941	# Per-block (group_size) weight scale. Initialized to NaN (rather than
1942	# left as torch.empty garbage) so an unloaded scale is unambiguously
1943	# detectable in process() below, matching the sentinel pattern in #45320.
1944	self.register_params(
1945	layer,
1946	"weight_scale",
1947	(
1948	shapes.output_size_per_partition,
1949	shapes.input_size_per_partition // ctx.group_size,
1950	),
1951	torch.float8_e4m3fn,
1952	ModelWeightParameter,
1953	wl,
1954	input_dim=1,
1955	output_dim=0,
1956	init=float("nan"),
1957	)
1958
1959	def process(self, layer, role) -> None:
1960	if role is not WEIGHT:
1961	self.reject(role)
1962	# Sanity-check: weight_scale must have been overwritten by the weight
1963	# loader. A value still equal to the NaN sentinel it was created with
1964	# means the FP4 weights were never actually loaded (e.g. the
1965	# checkpoint stores this layer as BF16 and the weight loader silently
1966	# skipped it).
1967	if torch.isnan(layer.weight_scale.float()).any():
1968	raise RuntimeError(
1969	f"NVFP4 weight_scale for layer "
1970	f"{getattr(layer, 'name', repr(layer))!r} was never loaded "
1971	"(still NaN). The checkpoint likely stores this layer as "
1972	"BF16 (not FP4). Fix: pass quant_config=None when "
1973	"constructing this layer, or add it to the quantization "
1974	"ignore list."
1975	)
1976	if torch.unique(layer.weight_scale_2).numel() != 1:
1977	logger.warning_once(
1978	"In NVFP4 linear, the global weight scale differs across "
1979	"parallel layers (e.g. q_proj, k_proj, v_proj). This will "
1980	"likely reduce accuracy. Consider a checkpoint with the same "
1981	"global NVFP4 scale for fused weights."
1982	)
1983	# Raw max, no reciprocation — Marlin/cutlass want ModelOpt's amax/2688.
1984	weight_global_scale = layer.weight_scale_2.max().to(torch.float32)
1985	layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
1986	del layer.weight_scale_2
1987
1988
1989	class KNvfp4Dynamic(QuantKeyScheme):
1990	"""NVFP4 activation scheme (W4A4). Has a static global input scale on disk;
1991	the per-group scale is computed at runtime inside the kernel."""
1992
1993	key = kNvfp4Dynamic
1994
1995	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
1996	if role is not ACT:
1997	self.reject(role)
1998	self.register_params(
1999	layer,
2000	"input_scale",
2001	(shapes.num_partitions,),
2002	torch.float32,
2003	PerTensorScaleParameter,
2004	wl,
2005	)
2006
2007	def process(self, layer, role) -> None:
2008	if role is not ACT:
2009	self.reject(role)
2010	if torch.unique(layer.input_scale).numel() != 1:
2011	logger.warning_once(
2012	"In NVFP4 linear, the global input scale differs across "
2013	"parallel layers (e.g. q_proj, k_proj, v_proj). This will "
2014	"likely reduce accuracy. Consider a checkpoint with the same "
2015	"global NVFP4 scale for fused weights."
2016	)
2017	input_global_scale = layer.input_scale.max().to(torch.float32)
2018	layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
2019	layer.input_global_scale_inv = Parameter(
2020	(1.0 / layer.input_global_scale).to(torch.float32),
2021	requires_grad=False,
2022	)
2023	del layer.input_scale
2024
2025
2026	class KFp8StaticTensor(QuantKeyScheme):
2027	"""Plain per-tensor static FP8 — bivalent: serves BOTH the weight slot and
2028	the activation slot (W8A8). One key in both QuantSpec slots."""
2029
2030	key = kFp8StaticTensorSym
2031
2032	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2033	if role is WEIGHT:
2034	self.register_params(
2035	layer,
2036	"weight",
2037	(shapes.output_size_per_partition, shapes.input_size_per_partition),
2038	torch.float8_e4m3fn,
2039	ModelWeightParameter,
2040	wl,
2041	input_dim=1,
2042	output_dim=0,
2043	)
2044	layer.orig_dtype = shapes.params_dtype
2045	self.register_params(
2046	layer,
2047	"weight_scale",
2048	(shapes.num_partitions,),
2049	torch.float32,
2050	PerTensorScaleParameter,
2051	wl,
2052	init=FP8_SCALE_SENTINEL,
2053	)
2054	elif role is ACT:
2055	self.register_params(
2056	layer,
2057	"input_scale",
2058	(shapes.num_partitions,),
2059	torch.float32,
2060	PerTensorScaleParameter,
2061	wl,
2062	init=FP8_SCALE_SENTINEL,
2063	)
2064	else:
2065	self.reject(role)
2066
2067	def process(self, layer, role) -> None:
2068	if role is WEIGHT:
2069	weight = layer.weight
2070	max_w_scale = layer.weight_scale.max()
2071	if not (layer.weight_scale == layer.weight_scale[0]).all():
2072	max_w_scale, weight = requantize_with_max_scale(
2073	layer.weight, layer.weight_scale, layer.logical_widths
2074	)
2075	# Transposed here rather than in the kernel: ModelOpt maps each
2076	# fp8 key to exactly one kernel, so the layout is key-determined.
2077	layer.weight = Parameter(weight.t(), requires_grad=False)
2078	# The plain Parameter drops ModelWeightParameter's dim attributes;
2079	# restore them for the transposed layout, which Humming reads.
2080	layer.weight.input_dim = 0
2081	layer.weight.output_dim = 1
2082	layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
2083	elif role is ACT:
2084	if torch.unique(layer.input_scale).numel() != 1:
2085	logger.warning_once(
2086	"In FP8 linear, the static input scale differs across "
2087	"parallel layers (e.g. q_proj, k_proj, v_proj). Collapsing "
2088	"them to the max will likely reduce accuracy. Consider a "
2089	"checkpoint with the same input scale for fused weights."
2090	)
2091	layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
2092	else:
2093	self.reject(role)
2094
2095
2096	class KFp8StaticChannel(QuantKeyScheme):
2097	"""Per-channel static FP8 weight (the 'PcPt' weight). Weight-role only —
2098	there is no static per-channel activation today."""
2099
2100	key = kFp8StaticTokenSym
2101
2102	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2103	if role is not WEIGHT:
2104	self.reject(role)
2105	self.register_params(
2106	layer,
2107	"weight",
2108	(shapes.output_size_per_partition, shapes.input_size_per_partition),
2109	torch.float8_e4m3fn,
2110	ModelWeightParameter,
2111	wl,
2112	input_dim=1,
2113	output_dim=0,
2114	)
2115	self.register_params(
2116	layer,
2117	"weight_scale",
2118	(shapes.output_size_per_partition,),
2119	torch.float32,
2120	ChannelQuantScaleParameter,
2121	wl,
2122	output_dim=0,
2123	init=FP8_SCALE_SENTINEL,
2124	)
2125
2126	def process(self, layer, role) -> None:
2127	if role is not WEIGHT:
2128	self.reject(role)
2129	weight, weight_scale, _ = process_fp8_weight_channel_strategy(
2130	layer.weight, layer.weight_scale.data
2131	)
2132	layer.weight = Parameter(weight, requires_grad=False)
2133	layer.weight_scale = Parameter(weight_scale, requires_grad=False)
2134
2135
2136	class KFp8Block128(QuantKeyScheme):
2137	"""128x128 block-static FP8 weight ('PbWo'). Weight-role only. ModelOpt
2138	exports the scale 4-D [out_blk,1,in_blk,1]; process squeezes to 2-D.
2139	No transpose (block kernel keeps [out,in])."""
2140
2141	key = kFp8Static128BlockSym
2142
2143	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2144	if role is not WEIGHT:
2145	self.reject(role)
2146	out = shapes.output_size_per_partition
2147	in_ = shapes.input_size_per_partition
2148	# Scale sized by cdiv to match the checkpoint's block grid (the last
2149	# output block may be partial). Weight is registered at its logical
2150	# size; a non-128 output width is padded to a block boundary by the
2151	# _Fp8PbWoPartialBlock format scheme. Input must be 128-aligned (same as
2152	# #53132, which only pads the output).
2153	if in_ % 128 != 0:
2154	raise ValueError(f"FP8_PB_WO requires in divisible by 128, got {in_}")
2155	self.register_params(
2156	layer,
2157	"weight",
2158	(out, in_),
2159	torch.float8_e4m3fn,
2160	ModelWeightParameter,
2161	wl,
2162	input_dim=1,
2163	output_dim=0,
2164	)
2165	ob, ib = cdiv(out, 128), in_ // 128
2166	self.register_params(
2167	layer,
2168	"weight_scale",
2169	(ob, 1, ib, 1),
2170	torch.float32,
2171	BlockQuantScaleParameter,
2172	wl,
2173	input_dim=2,
2174	output_dim=0,
2175	init=FP8_SCALE_SENTINEL,
2176	)
2177	layer.weight_block_size = [128, 128]
2178
2179	def process(self, layer, role) -> None:
2180	if role is not WEIGHT:
2181	self.reject(role)
2182	layer.weight = Parameter(layer.weight.data, requires_grad=False)
2183	s = layer.weight_scale
2184	if s.dim() == 4:
2185	s = s.squeeze(1).squeeze(-1) # [ob,1,ib,1] -> [ob,ib]
2186	elif s.dim() != 2:
2187	raise ValueError(
2188	f"Unexpected FP8_PB_WO weight_scale shape {tuple(s.shape)}"
2189	)
2190	layer.weight_scale = Parameter(s.contiguous(), requires_grad=False)
2191
2192
2193	class KMxfp8Static(QuantKeyScheme):
2194	"""MXFP8 weight: fp8-e4m3 values + per-32-block e8m0 (uint8) scale.
2195	Weight-role only. process is validate-only plus an idempotency guard."""
2196
2197	key = kMxfp8Static
2198
2199	@staticmethod
2200	def get_scale_weight_loader(weight_loader: Callable, ctx: CkptCtx) -> Callable:
2201	block_rows, block_cols = ctx.scale_block_size or (1, MXFP8_BLOCK_SIZE)
2202	if block_rows < 1 or block_cols != MXFP8_BLOCK_SIZE:
2203	raise NotImplementedError(
2204	f"MXFP8 checkpoint scale block {ctx.scale_block_size} is unsupported"
2205	)
2206
2207	def scaled_loader(param, loaded_weight, args, *kwargs):
2208	assert loaded_weight.dtype in (torch.uint8, torch.float8_e8m0fnu)
2209	loaded_weight = loaded_weight.view(torch.uint8).repeat_interleave(
2210	block_rows, dim=0
2211	)
2212	return weight_loader(param, loaded_weight, args, *kwargs)
2213
2214	return scaled_loader if block_rows > 1 else weight_loader
2215
2216	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2217	if role is not WEIGHT:
2218	self.reject(role)
2219	if shapes.input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
2220	raise ValueError(
2221	f"MXFP8 requires input size divisible by {MXFP8_BLOCK_SIZE}, "
2222	f"got {shapes.input_size_per_partition}"
2223	)
2224	self.register_params(
2225	layer,
2226	"weight",
2227	(shapes.output_size_per_partition, shapes.input_size_per_partition),
2228	MXFP8_VALUE_DTYPE,
2229	ModelWeightParameter,
2230	wl,
2231	input_dim=1,
2232	output_dim=0,
2233	)
2234	scale_loader = self.get_scale_weight_loader(wl, ctx)
2235	self.register_params(
2236	layer,
2237	"weight_scale",
2238	(
2239	shapes.output_size_per_partition,
2240	shapes.input_size_per_partition // MXFP8_BLOCK_SIZE,
2241	),
2242	MXFP8_SCALE_DTYPE,
2243	ModelWeightParameter,
2244	scale_loader,
2245	input_dim=1,
2246	output_dim=0,
2247	)
2248	layer.weight_block_size = [1, MXFP8_BLOCK_SIZE]
2249
2250	def process(self, layer, role) -> None:
2251	if role is not WEIGHT:
2252	self.reject(role)
2253	# Idempotency: the emulation kernel's post-load may replace the 1-byte
2254	# MXFP8 weight with BF16 (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD, default
2255	# on). On a weight reload process runs again after that swap, so skip
2256	# the fp8 asserts when the weight is already >=2-byte.
2257	if layer.weight.element_size() >= 2:
2258	return
2259	if getattr(layer, "is_bmm", False) and layer.weight.ndim == 3:
2260	return
2261	assert layer.weight.ndim == 2 and layer.weight.dtype == MXFP8_VALUE_DTYPE
2262	assert layer.weight_scale.ndim == 2
2263	assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE
2264
2265
2266	class KDynamicNoParam(QuantKeyScheme):
2267	"""Dynamic activation with no stored scale (W8A8): quantized at runtime in
2268	the kernel. NOT the same as activation=None (weight-only) — init_fp8 needs a
2269	non-None activation key. Activation-role only. Serves the fp8 per-token, fp8
2270	per-block, and mxfp8 dynamic activation keys."""
2271
2272	def create_weights(self, layer, role, ctx, shapes, wl) -> None:
2273	if role is not ACT:
2274	self.reject(role)
2275	# dynamic -> nothing stored
2276
2277	def process(self, layer, role) -> None:
2278	if role is not ACT:
2279	self.reject(role)
2280
2281
2282	SCHEME_FOR: dict[QuantKey \| None, QuantKeyScheme] = {
2283	kNvfp4Static: KNvfp4Static(),
2284	kNvfp4Dynamic: KNvfp4Dynamic(),
2285	kFp8StaticTensorSym: KFp8StaticTensor(),
2286	kFp8StaticTokenSym: KFp8StaticChannel(),
2287	kFp8Static128BlockSym: KFp8Block128(),
2288	kMxfp8Static: KMxfp8Static(),
2289	kFp8DynamicTokenSym: KDynamicNoParam(),
2290	kFp8Dynamic128Sym: KDynamicNoParam(),
2291	kMxfp8Dynamic: KDynamicNoParam(),
2292	}
2293
2294
2295	def maybe_fuse_global_scales(layer) -> None:
2296	"""Alpha = input_global_scale * weight_global_scale, presence-gated.
2297
2298	W4A4 has both -> computed; W4A16 has no input_global_scale -> skipped.
2299	"""
2300	if hasattr(layer, "weight_global_scale") and hasattr(layer, "input_global_scale"):
2301	layer.alpha = Parameter(
2302	layer.input_global_scale * layer.weight_global_scale,
2303	requires_grad=False,
2304	)
2305
2306
2307	def select_linear_kernel(
2308	spec: QuantSpec,
2309	layer,
2310	rt: RuntimeDtypes,
2311	weight_shape: tuple[int, int] \| None = None,
2312	):
2313	"""Thin family dispatcher on the weight key: nvfp4 / mxfp8 / fp8."""
2314	w = spec.weight
2315	assert isinstance(w, QuantKey), f"resolve() must supply a weight key, got {w!r}"
2316	if w.dtype == FP4_DTYPE:
2317	# W4A16 (activation is None) → use_a16=True, auto prefers FlashInfer
2318	# CuTe-DSL backend on SM100/103 when available and Marlin otherwise,
2319	# and honors --linear-backend (#50273); W4A4 → use_a16=False.
2320	return init_nvfp4_linear_kernel(use_a16=spec.activation is None)
2321	if w.scale.dtype == MXFP8_SCALE_DTYPE:
2322	return init_mxfp8_linear_kernel(
2323	weight_shape=weight_shape or layer.weight.shape,
2324	bmm_batch_size=getattr(layer, "bmm_batch_size", None),
2325	)
2326	# fp8 family: init_fp8 routes block-vs-plain itself off the activation key,
2327	# and needs a real key -- weight-only fp8 is not a ModelOpt format.
2328	act = spec.activation
2329	assert isinstance(act, QuantKey), f"fp8 needs an activation key, got {act!r}"
2330	return init_fp8_linear_kernel(
2331	activation_quant_key=act,
2332	weight_quant_key=w,
2333	input_dtype=rt.input_dtype,
2334	out_dtype=rt.out_dtype,
2335	weight_shape=weight_shape or layer.weight.shape,
2336	module_name=type(layer).__name__,
2337	)
2338
2339
2340	class FormatScheme:
2341	"""Optional per-format hooks that compose around the QuantKey schemes.
2342
2343	Extension seam for residue that belongs to a format as a whole rather than a
2344	single QuantKey — an extra parameter the weight/activation schemes don't
2345	cover, or a tweak before/after their ``process``. All hooks default to
2346	no-ops; most formats need none (they are a pure ``(weight, activation)`` key
2347	pair). To add one: subclass, override the hooks you need, and return an
2348	instance from the format's ``resolve()`` branch — no change to
2349	``ModelOptLinearMethod`` itself.
2350	"""
2351
2352	def extra_weights(self, layer, shapes: "Shapes", ctx: CkptCtx, wl) -> None:
2353	"""Register format-level params (after the key schemes' weights)."""
2354
2355	def kernel_weight_shape(self, layer) -> tuple[int, int]:
2356	"""Return the weight shape the selected kernel will receive."""
2357	return tuple(layer.weight.shape)
2358
2359	def pre_process(self, layer) -> None:
2360	"""Run before the key schemes' ``process``."""
2361
2362	def post_process(self, layer) -> None:
2363	"""Run after the key schemes' ``process``, before the kernel's."""
2364
2365	def apply(self, layer, x, bias, kernel_apply):
2366	"""Wrap the kernel's forward. ``kernel_apply(layer, x, bias) -> Tensor``.
2367
2368	Default: delegate unchanged. A format whose residue lives at compute
2369	time (e.g. run the GEMM on a padded weight then slice the output back
2370	to its logical width, #53132-style) overrides this and calls
2371	``kernel_apply`` in the middle.
2372	"""
2373	return kernel_apply(layer, x, bias)
2374
2375
2376	class _Fp8PbWoPartialBlock(FormatScheme):
2377	"""FP8_PB_WO output width that is not a multiple of 128 (a partial trailing
2378	block; e.g. GLM's replicated ``fused_qkv_a_proj`` = 2048 + 576 = 2624).
2379
2380	The block kernel needs full 128-blocks, so pad the weight up to the block
2381	boundary with zeros before the kernel's post-load, then slice the GEMM
2382	output back to the logical width. This is #53132's approach, expressed as
2383	the format's compute-time residue on the generic method. The cdiv-sized
2384	scale from ``KFp8Block128`` already matches the padded block count, so only
2385	the weight rows and the output need adjusting.
2386	"""
2387
2388	def kernel_weight_shape(self, layer) -> tuple[int, int]:
2389	out, in_ = layer.weight.shape
2390	return (cdiv(out, 128) * 128, in_)
2391
2392	def post_process(self, layer) -> None:
2393	w = layer.weight
2394	out = w.shape[0]
2395	pad = (-out) % 128
2396	if not pad:
2397	return
2398	layer._pbwo_logical_out = out
2399	padded = w.data.new_zeros(out + pad, w.shape[1])
2400	padded[:out].copy_(w.data)
2401	layer.weight = Parameter(padded, requires_grad=False)
2402
2403	def apply(self, layer, x, bias, kernel_apply):
2404	logical = getattr(layer, "_pbwo_logical_out", None)
2405	if logical is None: # width was block-aligned: nothing to trim
2406	return kernel_apply(layer, x, bias)
2407	# Run the GEMM on the padded weight without bias, then trim to the
2408	# logical width and add bias (the padded rows produce zeros we drop).
2409	out = kernel_apply(layer, x, None)
2410	out = out[..., :logical].contiguous()
2411	return out if bias is None else out.add_(bias)
2412
2413
2414	_PB_WO_PARTIAL_BLOCK = _Fp8PbWoPartialBlock()
2415
2416
2417	class _DropInputScale(FormatScheme):
2418	"""Interim: register then drop a W4A16 checkpoint's on-disk input_scale.
2419
2420	Kept for backward compat with previously-exported ModelOpt checkpoints.
2421	"""
2422
2423	def extra_weights(self, layer, shapes, ctx, wl) -> None:
2424	data = torch.full((shapes.num_partitions,), torch.nan)
2425	layer.register_parameter(
2426	"input_scale", PerTensorScaleParameter(data=data, weight_loader=wl)
2427	)
2428
2429	def post_process(self, layer) -> None:
2430	scale = getattr(layer, "input_scale", None)
2431	if scale is None:
2432	return
2433	if not torch.isnan(scale).all():
2434	logger.warning_once(
2435	"Dropping input_scale from a weight-only W4A16 NVFP4 "
2436	"checkpoint (deprecated)."
2437	)
2438	del layer.input_scale
2439
2440
2441	_DROP_INPUT_SCALE = _DropInputScale()
2442
2443
2444	@register_weight_loader_v2_supported_method
2445	class ModelOptLinearMethod(LinearMethodBase):
2446	"""Generic, format-agnostic ModelOpt linear method. Holds a weight scheme +
2447	an activation scheme (from the QuantSpec pair) and an optional
2448	``FormatScheme``, runs a fixed lifecycle, selects the kernel from the pair,
2449	and applies.
2450
2451	Registered for weight_loader_v2 (like every ModelOpt linear method) so the
2452	``BasevLLMParameter`` params route through the v2 fused loader, not the
2453	legacy shape-assert path.
2454	"""
2455
2456	def __init__(
2457	self,
2458	spec: QuantSpec,
2459	ctx: CkptCtx,
2460	format_scheme=None,
2461	) -> None:
2462	self.spec = spec
2463	self.ctx = ctx
2464	self.fmt = format_scheme or FormatScheme()
2465	self.wkey = SCHEME_FOR[spec.weight]
2466	self.akey = None if spec.activation is None else SCHEME_FOR[spec.activation]
2467	# Only the fp8/mxfp8 kernels read input_dtype; nvfp4 ignores it. During
2468	# real serving model_config is always set; fall back defensively when
2469	# there is no config context (bare unit-test dispatch).
2470	model_config = getattr(get_current_vllm_config_or_none(), "model_config", None)
2471	dtype = (
2472	model_config.dtype
2473	if model_config is not None
2474	else torch.get_default_dtype()
2475	)
2476	self.out_dtype = dtype
2477	self.input_dtype = dtype
2478	self.marlin_input_dtype = None
2479	# Kernel/backend are chosen in create_weights (after get_quant_method),
2480	# so the front-end marlin poke stays dormant here — same as the old
2481	# NVFP4 methods.
2482	self.kernel: Any = None
2483
2484	@property
2485	def supports_pre_processed_weights(self) -> bool: # type: ignore[override]
2486	# TODO(Isotr0py): support fp8 ModelOpt kernels transpose/repack.
2487	w = self.spec.weight
2488	return isinstance(w, QuantKey) and (
2489	w.dtype == FP4_DTYPE
2490	or (
2491	w == kMxfp8Static
2492	and self.kernel is not None
2493	and self.kernel.supports_pre_processed_weights
2494	)
2495	)
2496
2497	def create_weights(
2498	self,
2499	layer,
2500	input_size_per_partition: int,
2501	output_partition_sizes: list[int],
2502	input_size: int,
2503	output_size: int,
2504	params_dtype: torch.dtype,
2505	**extra_weight_attrs,
2506	) -> None:
2507	del input_size, output_size
2508	weight_loader = extra_weight_attrs.get("weight_loader")
2509	layer.logical_widths = output_partition_sizes
2510	layer.input_size_per_partition = input_size_per_partition
2511	layer.output_size_per_partition = sum(output_partition_sizes)
2512	shapes = Shapes(output_partition_sizes, input_size_per_partition, params_dtype)
2513
2514	self.wkey.create_weights(layer, WEIGHT, self.ctx, shapes, weight_loader)
2515	if self.akey:
2516	self.akey.create_weights(layer, ACT, self.ctx, shapes, weight_loader)
2517	self.fmt.extra_weights(layer, shapes, self.ctx, weight_loader)
2518
2519	rt = RuntimeDtypes(self.input_dtype, self.out_dtype, self.marlin_input_dtype)
2520	self.kernel = select_linear_kernel(
2521	self.spec,
2522	layer,
2523	rt,
2524	weight_shape=self.fmt.kernel_weight_shape(layer),
2525	)
2526	expose_input_quant_key(layer, self.kernel)
2527
2528	def process_weights_after_loading(self, layer) -> None:
2529	if self.spec.weight == kMxfp8Static and getattr(layer, "is_bmm", False):
2530	self.kernel = init_mxfp8_linear_kernel(
2531	weight_shape=(layer.weight.shape[-2], layer.weight.shape[-1]),
2532	bmm_batch_size=layer.bmm_batch_size,
2533	)
2534	if is_weights_pre_processed():
2535	if not self.supports_pre_processed_weights:
2536	raise RuntimeError(
2537	f"{type(self.kernel).__name__} cannot use pre-processed weights"
2538	)
2539	return
2540	self.fmt.pre_process(layer)
2541	self.wkey.process(layer, WEIGHT)
2542	if self.akey:
2543	self.akey.process(layer, ACT)
2544	maybe_fuse_global_scales(layer)
2545	self.fmt.post_process(layer)
2546	layer.is_w4a16_nvfp4 = (
2547	self.spec.weight == kNvfp4Static and self.spec.activation is None
2548	)
2549	if getattr(layer, "_retain_weight_for_gather", False):
2550	if not layer.is_w4a16_nvfp4:
2551	raise NotImplementedError(
2552	"Gathered projection is only supported for ModelOpt W4A16 "
2553	"NVFP4 weights."
2554	)
2555	assert self.ctx.group_size is not None
2556	layer.register_buffer(
2557	"_nvfp4_weight_for_gather", layer.weight.detach(), persistent=False
2558	)
2559	layer.register_buffer(
2560	"_nvfp4_weight_scale_for_gather",
2561	layer.weight_scale.detach(),
2562	persistent=False,
2563	)
2564	layer.register_buffer(
2565	"_nvfp4_weight_global_scale_for_gather",
2566	layer.weight_global_scale.detach(),
2567	persistent=False,
2568	)
2569	layer._nvfp4_group_size_for_gather = self.ctx.group_size
2570	self.kernel.process_weights_after_loading(layer)
2571
2572	def apply(self, layer, x, bias=None):
2573	return self.fmt.apply(
2574	layer,
2575	x,
2576	bias,
2577	lambda lyr, inp, b: self.kernel.apply_weights(layer=lyr, x=inp, bias=b),
2578	)
2579
2580
2581	def resolve(algo: str, subcfg, prefix: str):
2582	"""Turn a sub-config into (QuantSpec, CkptCtx, format_scheme).
2583
2584	Strictly read-only over ``subcfg``: in mixed precision the same sub-config
2585	is shared with the MoE method, so writing back would leak a linear-only
2586	change into MoE.
2587	"""
2588	if algo == "FP8":
2589	# Plain per-tensor static FP8 (W8A8): same key in both slots (bivalent).
2590	ctx = CkptCtx()
2591	return (
2592	QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym),
2593	ctx,
2594	None,
2595	)
2596	if algo == "FP8_PER_CHANNEL_PER_TOKEN":
2597	# PcPt: per-channel static weight, dynamic per-token activation (W8A8).
2598	ctx = CkptCtx()
2599	return (
2600	QuantSpec(weight=kFp8StaticTokenSym, activation=kFp8DynamicTokenSym),
2601	ctx,
2602	None,
2603	)
2604	if algo == "FP8_PB_WO":
2605	# PbWo: 128x128 block-static weight, dynamic per-block activation (W8A8).
2606	# The block kernel's post-load runs here; CompressedTensors block-FP8
2607	# is the reference for this path. The FormatScheme pads a non-128 output
2608	# width to a block boundary and trims the output back (#53132) -- a no-op
2609	# for the common block-aligned case.
2610	ctx = CkptCtx()
2611	return (
2612	QuantSpec(weight=kFp8Static128BlockSym, activation=kFp8Dynamic128Sym),
2613	ctx,
2614	_PB_WO_PARTIAL_BLOCK,
2615	)
2616	if algo == "MXFP8":
2617	# MXFP8: block(32) e4m3 weight + e8m0 scale, dynamic activation.
2618	ctx = CkptCtx()
2619	return QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic), ctx, None
2620
2621	# NVFP4 family (W4A4 / W4A16).
2622	ctx = CkptCtx(group_size=subcfg.group_size)
2623	if algo == "NVFP4":
2624	# W4A4: static fp4 weight + dynamic fp4 activation (has a static global
2625	# input scale). alpha = weight_gs * input_gs.
2626	return QuantSpec(weight=kNvfp4Static, activation=kNvfp4Dynamic), ctx, None
2627	if algo == "W4A16_NVFP4":
2628	# W4A16: fp4 weight, no activation quant (-> Marlin). _DROP_INPUT_SCALE
2629	# is a no-op for genuine exports; keeps legacy input_scale ckpts loading.
2630	return QuantSpec(weight=kNvfp4Static, activation=None), ctx, _DROP_INPUT_SCALE
2631	raise NotImplementedError(f"resolve: unsupported ModelOpt linear algo {algo!r}")
2632
2633
2634	# Bespoke-method escape hatch. Almost every ModelOpt linear format is a
2635	# (weight, activation) QuantKey pair handled by the generic ModelOptLinearMethod
2636	# (optionally + a FormatScheme), so this stays empty. A format that genuinely
2637	# cannot be expressed that way — a fundamentally different create/process/apply
2638	# lifecycle — registers its own LinearMethodBase here, keyed by algo:
2639	#
2640	# LINEAR_METHOD_BUILDERS["FOO"] = lambda cfg, prefix: ModelOptFooLinearMethod(cfg)
2641	#
2642	# Keep the ModelOpt* class-name prefix and decorate the class with
2643	# @register_weight_loader_v2_supported_method if it uses BasevLLMParameter
2644	# params. Also add the algo to the owning config's quant_method validation.
2645	LINEAR_METHOD_BUILDERS: dict[str, Callable[..., LinearMethodBase]] = {}
2646
2647
2648	def build_linear_method(config, algo: str, prefix: str) -> LinearMethodBase:
2649	"""Construct the linear method for ``algo``.
2650
2651	Returns a bespoke method if one is registered in ``LINEAR_METHOD_BUILDERS``,
2652	else the generic ``ModelOptLinearMethod`` built from ``resolve(algo, config,
2653	prefix)``. Single indirection point shared by the homogeneous and
2654	mixed-precision dispatch, so a new format plugs in without editing either
2655	``get_quant_method``.
2656	"""
2657	builder = LINEAR_METHOD_BUILDERS.get(algo)
2658	if builder is not None:
2659	return builder(config, prefix)
2660	spec, ctx, format_scheme = resolve(algo, config, prefix)
2661	return ModelOptLinearMethod(spec, ctx, format_scheme)