diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 3fc63a26ad..5bb1aa0279 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -27,14 +27,19 @@ from vllm.model_executor.kernels.linear import ( from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( + LINEAR_ALGOS, ModelOptFp8Config, - ModelOptFp8LinearMethod, - ModelOptFp8PbWoLinearMethod, + ModelOptLinearMethod, ModelOptMixedPrecisionConfig, ModelOptMxFp8Config, ModelOptNvFp4Config, - ModelOptNvFp4LinearMethod, - ModelOptNvFp4W4A16LinearMethod, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8StaticTensorSym, + kMxfp8Dynamic, + kMxfp8Static, + kNvfp4Dynamic, + kNvfp4Static, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -114,15 +119,25 @@ def test_modelopt_nvfp4_quantizes_parallel_lm_head(): exclude_modules=[], ) - with patch( - "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel" - ): - method = config.get_quant_method(_mock_lm_head(), prefix="lm_head") + method = config.get_quant_method(_mock_lm_head(), prefix="lm_head") - assert isinstance(method, ModelOptNvFp4LinearMethod) + assert isinstance(method, ModelOptLinearMethod) + assert method.spec.weight is kNvfp4Static + assert method.spec.activation is kNvfp4Dynamic def test_modelopt_fp8_updates_weight_dims_after_transpose(): + """Humming reads weight.input_dim/output_dim. Swapping the + ModelWeightParameter for a plain Parameter drops them, so the per-tensor + FP8 scheme must restore them for the transposed [in, out] layout. + """ + from vllm.config.quantization import QuantSpec + from vllm.model_executor.layers.quantization.modelopt import ( + SCHEME_FOR, + CkptCtx, + FormatScheme, + ) + layer = torch.nn.Module() layer.register_parameter( "weight", torch.nn.Parameter(torch.empty(3, 2), requires_grad=False) @@ -133,73 +148,80 @@ def test_modelopt_fp8_updates_weight_dims_after_transpose(): layer.register_parameter( "input_scale", torch.nn.Parameter(torch.ones(1), requires_grad=False) ) + layer.logical_widths = [3] - method = ModelOptFp8LinearMethod.__new__(ModelOptFp8LinearMethod) - method.fp8_linear = Mock() + method = ModelOptLinearMethod.__new__(ModelOptLinearMethod) + method.spec = QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym) + method.ctx = CkptCtx() + method.fmt = FormatScheme() + method.wkey = SCHEME_FOR[kFp8StaticTensorSym] + method.akey = SCHEME_FOR[kFp8StaticTensorSym] + method.kernel = Mock() method.process_weights_after_loading(layer) assert layer.weight.shape == (2, 3) assert layer.weight.input_dim == 0 assert layer.weight.output_dim == 1 - method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer) + method.kernel.process_weights_after_loading.assert_called_once_with(layer) -def test_modelopt_fp8_pb_wo_hides_output_padding(default_vllm_config): - config = ModelOptFp8Config( - quant_method="FP8_PB_WO", - is_checkpoint_fp8_serialized=True, - kv_cache_quant_method=None, - exclude_modules=[], +def test_modelopt_linear_algos_table_matches_resolve(): + """LINEAR_ALGOS is the single source of truth for supported linear algos. + + Every entry must be dispatchable by resolve(), and every config's + validation list must be derived from it -- so adding a format is one row + here plus one row in resolve(), with nothing else to keep in sync. + """ + from vllm.model_executor.layers.quantization.modelopt import ( + QUANT_ALGOS, + algos_owned_by, + resolve, ) - default_vllm_config.model_config = Mock(dtype=torch.bfloat16) - kernel = Mock() - with ( - set_current_vllm_config(default_vllm_config), - patch( - "vllm.model_executor.layers.quantization.modelopt.init_fp8_linear_kernel", - return_value=kernel, - ), - patch( - "vllm.model_executor.parameter.get_tensor_model_parallel_rank", - return_value=0, - ), - patch( - "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", - return_value=1, - ), - ): - method = ModelOptFp8PbWoLinearMethod(config) - layer = torch.nn.Module() - method.create_weights( - layer, - input_size_per_partition=128, - output_partition_sizes=[128, 65], - input_size=128, - output_size=193, - params_dtype=torch.bfloat16, - ) + class _Cfg: + group_size = 16 - assert layer.weight.shape == (193, 128) - assert layer.weight_scale.shape == (2, 1, 1, 1) + for algo in LINEAR_ALGOS: + spec, _, _ = resolve(algo, _Cfg(), "layer") + assert spec.weight is not None, algo - layer.weight.data.fill_(1) - method.process_weights_after_loading(layer) + assert list(QUANT_ALGOS) == [*LINEAR_ALGOS, "MIXED_PRECISION"] + assert algos_owned_by("modelopt") == ( + "FP8", + "FP8_PER_CHANNEL_PER_TOKEN", + "FP8_PB_WO", + ) + assert algos_owned_by("modelopt_fp4") == ("NVFP4", "W4A16_NVFP4") + assert algos_owned_by("modelopt_mxfp8") == ("MXFP8",) - assert layer.weight.shape == (256, 128) - assert torch.count_nonzero(layer.weight[193:].float()) == 0 - kernel.process_weights_after_loading.assert_called_once_with(layer) - physical_output = torch.randn(2, 256, dtype=torch.bfloat16) - kernel.apply_weights.return_value = physical_output - bias = torch.randn(193, dtype=torch.bfloat16) - output = method.apply(layer, torch.randn(2, 128), bias) +@pytest.mark.parametrize("algo", list(LINEAR_ALGOS)) +def test_modelopt_mixed_precision_dispatches_every_linear_algo(algo): + """Mixed precision must route every algo in LINEAR_ALGOS through the + generic method. FP8_PER_CHANNEL_PER_TOKEN and FP8_PB_WO used to fall + through to UnquantizedLinearMethod, which loses the checkpoint's scales. + """ + from vllm.model_executor.layers.linear import LinearBase + from vllm.model_executor.layers.quantization import modelopt as m - torch.testing.assert_close(output, physical_output[:, :193] + bias) - assert output.shape == (2, 193) - assert output.is_contiguous() - kernel.apply_weights.assert_called_once() - assert kernel.apply_weights.call_args.args[2] is None + config = m.ModelOptMixedPrecisionConfig.from_config( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "exclude_modules": [], + "group_size": 16, + "quantized_layers": { + "model.layers.0.mlp.down_proj": {"quant_algo": algo} + }, + } + } + ) + method = config.get_quant_method( + MagicMock(spec=LinearBase), "model.layers.0.mlp.down_proj" + ) + + assert isinstance(method, ModelOptLinearMethod), (algo, type(method).__name__) def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized(): @@ -219,12 +241,11 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head(): {"lm_head": {"quant_algo": "NVFP4", "group_size": 16}} ) - with patch( - "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel" - ): - method = config.get_quant_method(_mock_lm_head(), prefix="lm_head") + method = config.get_quant_method(_mock_lm_head(), prefix="lm_head") - assert isinstance(method, ModelOptNvFp4LinearMethod) + assert isinstance(method, ModelOptLinearMethod) + assert method.spec.weight is kNvfp4Static + assert method.spec.activation is kNvfp4Dynamic def test_modelopt_mixed_precision_resolves_declared_packed_projection(): @@ -302,12 +323,11 @@ def test_modelopt_mixed_precision_infers_fused_gate_up_projection(): ) fake_layer = MagicMock(spec=LinearBase) - with patch( - "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel" - ): - method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj") + method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj") - assert isinstance(method, ModelOptNvFp4LinearMethod) + assert isinstance(method, ModelOptLinearMethod) + assert method.spec.weight is kNvfp4Static + assert method.spec.activation is kNvfp4Dynamic @pytest.mark.parametrize( @@ -363,7 +383,7 @@ def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner): "This test requires a local ModelOpt FP8 checkpoint." ) - # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod. + # Set model config as model_config.dtype is required in ModelOptLinearMethod. default_vllm_config.model_config = ModelConfig() with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm: @@ -376,14 +396,10 @@ def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner): down_proj = layer.mlp.down_proj # Check that ModelOpt quantization method is properly applied - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptFp8LinearMethod, - ) - - assert isinstance(qkv_proj.quant_method, ModelOptFp8LinearMethod) - assert isinstance(o_proj.quant_method, ModelOptFp8LinearMethod) - assert isinstance(gate_up_proj.quant_method, ModelOptFp8LinearMethod) - assert isinstance(down_proj.quant_method, ModelOptFp8LinearMethod) + assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod) + assert isinstance(o_proj.quant_method, ModelOptLinearMethod) + assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod) + assert isinstance(down_proj.quant_method, ModelOptLinearMethod) # Check weight dtype is FP8 assert qkv_proj.weight.dtype == torch.float8_e4m3fn @@ -440,14 +456,10 @@ def test_modelopt_fp8_pc_pt_checkpoint_setup(monkeypatch, dist_init, workspace_i gate_up_proj = layer.mlp.gate_up_proj down_proj = layer.mlp.down_proj - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptFp8PcPtLinearMethod, - ) - - assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod) - assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod) - assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod) - assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod) + assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod) + assert isinstance(o_proj.quant_method, ModelOptLinearMethod) + assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod) + assert isinstance(down_proj.quant_method, ModelOptLinearMethod) fp8_dtype = current_platform.fp8_dtype() assert qkv_proj.weight.dtype == fp8_dtype @@ -494,14 +506,10 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_i gate_up_proj = layer.mlp.gate_up_proj down_proj = layer.mlp.down_proj - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptFp8PbWoLinearMethod, - ) - - assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod) - assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod) - assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod) - assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod) + assert isinstance(qkv_proj.quant_method, ModelOptLinearMethod) + assert isinstance(o_proj.quant_method, ModelOptLinearMethod) + assert isinstance(gate_up_proj.quant_method, ModelOptLinearMethod) + assert isinstance(down_proj.quant_method, ModelOptLinearMethod) fp8_dtype = current_platform.fp8_dtype() assert qkv_proj.weight.dtype == fp8_dtype @@ -528,12 +536,10 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_i def test_modelopt_nvfp4_config_dispatches_w4a4_method(): - """``quant_method="NVFP4"`` (W4A4 default) routes to the existing - ``ModelOptNvFp4LinearMethod``.""" - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ModelOptNvFp4LinearMethod, - ) + """``quant_method="NVFP4"`` (W4A4) resolves to a + ``(kNvfp4Static, kNvfp4Dynamic)`` QuantSpec under the generic + ``ModelOptLinearMethod``.""" + from vllm.model_executor.layers.linear import LinearBase config = ModelOptNvFp4Config( quant_method="NVFP4", @@ -541,24 +547,25 @@ def test_modelopt_nvfp4_config_dispatches_w4a4_method(): kv_cache_quant_algo=None, exclude_modules=[], ) - assert config.LinearMethodCls is ModelOptNvFp4LinearMethod assert config.quant_method == "NVFP4" + method = config.get_quant_method( + MagicMock(spec=LinearBase), "model.layers.0.fake_proj" + ) + assert isinstance(method, ModelOptLinearMethod) + assert method.spec.weight is kNvfp4Static + assert method.spec.activation is kNvfp4Dynamic + def test_modelopt_nvfp4_config_dispatches_w4a16_method(): - """``quant_method="W4A16_NVFP4"`` routes to - ``ModelOptNvFp4W4A16LinearMethod`` instead of the W4A4 sibling. + """``quant_method="W4A16_NVFP4"`` resolves to a weight-only QuantSpec + (``activation=None``) — distinct from the W4A4 sibling. - Mirrors the FP8 dispatch precedent (``ModelOptFp8Config`` selects - one of three FP8 LinearMethods on ``quant_method``); a regression - here would mean a W4A16 NVFP4 checkpoint silently loaded under the - W4A4 activation-quantization path. + A regression here would mean a W4A16 NVFP4 checkpoint silently loaded + with a dynamic fp4 activation key, registering an ``input_scale`` and + routing to the cutlass W4A4 NVFP4 GEMM instead of FP4 Marlin. """ - from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptNvFp4Config, - ModelOptNvFp4LinearMethod, - ModelOptNvFp4W4A16LinearMethod, - ) + from vllm.model_executor.layers.linear import LinearBase config = ModelOptNvFp4Config( quant_method="W4A16_NVFP4", @@ -566,10 +573,38 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method(): kv_cache_quant_algo=None, exclude_modules=[], ) - assert config.LinearMethodCls is ModelOptNvFp4W4A16LinearMethod - assert config.LinearMethodCls is not ModelOptNvFp4LinearMethod assert config.quant_method == "W4A16_NVFP4" + method = config.get_quant_method( + MagicMock(spec=LinearBase), "model.layers.0.fake_proj" + ) + assert isinstance(method, ModelOptLinearMethod) + assert method.spec.weight is kNvfp4Static + assert method.spec.activation is None + + +def test_modelopt_linear_method_builder_registry_override(monkeypatch): + """The bespoke-method escape hatch: a format registered in + ``LINEAR_METHOD_BUILDERS`` routes that algo to its own method instead of the + generic ``ModelOptLinearMethod``. This is how a format that cannot be a + ``(weight, activation)`` key pair plugs into dispatch.""" + from vllm.model_executor.layers.linear import LinearBase + from vllm.model_executor.layers.quantization import modelopt as m + + sentinel = object() + monkeypatch.setitem(m.LINEAR_METHOD_BUILDERS, "NVFP4", lambda cfg, prefix: sentinel) + + config = ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + ) + method = config.get_quant_method( + MagicMock(spec=LinearBase), "model.layers.0.fake_proj" + ) + assert method is sentinel # bespoke builder wins over the generic path + @pytest.mark.parametrize( ("linear_backend", "kernel_cls"), @@ -581,6 +616,15 @@ def test_modelopt_nvfp4_config_dispatches_w4a16_method(): ) @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): + """W4A16 (`activation=None`) kernel selection honors ``--linear-backend``: + ``use_a16=True`` defaults to Marlin, but an explicit backend wins. The + generic method routes this through ``select_linear_kernel``.""" + from vllm.config.quantization import QuantSpec + from vllm.model_executor.layers.quantization.modelopt import ( + RuntimeDtypes, + select_linear_kernel, + ) + if linear_backend != "auto": is_supported, reason = kernel_cls.is_supported() if not is_supported: @@ -588,11 +632,52 @@ def test_modelopt_w4a16_respects_linear_backend(linear_backend, kernel_cls): vllm_config = VllmConfig() vllm_config.kernel_config.linear_backend = linear_backend + spec = QuantSpec(weight=kNvfp4Static, activation=None) + rt = RuntimeDtypes(torch.bfloat16, torch.bfloat16) with set_current_vllm_config(vllm_config): - method = ModelOptNvFp4W4A16LinearMethod( - ModelOptNvFp4Config(quant_method="W4A16_NVFP4") + kernel = select_linear_kernel(spec, MagicMock(), rt) + assert isinstance(kernel, kernel_cls) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +def test_modelopt_linear_exposes_humming_layer_attrs(dist_init, monkeypatch): + """``prepare_humming_linear_layer_config`` reads ``output_partition_sizes`` + and ``has_bias`` straight off the layer, so ``--linear-backend=humming`` + needs create_weights to leave both there. Nothing else sets + ``output_partition_sizes``; ``LinearBase`` sets ``has_bias`` but + ``ParallelLMHead`` does not. + """ + from vllm.config.quantization import QuantSpec + from vllm.model_executor.layers.quantization import modelopt as mo + + monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: Mock()) + monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, kernel: None) + + def build(layer): + method = ModelOptLinearMethod.__new__(ModelOptLinearMethod) + method.spec = QuantSpec(weight=kNvfp4Static, activation=None) + method.ctx = mo.CkptCtx(group_size=16) + method.fmt = mo.FormatScheme() + method.wkey = mo.SCHEME_FOR[kNvfp4Static] + method.akey = None + method.input_dtype = method.out_dtype = torch.bfloat16 + method.marlin_input_dtype = None + method.create_weights( + layer, 64, [32, 32], 64, 64, torch.bfloat16, weight_loader=Mock() ) - assert isinstance(method.kernel, kernel_cls) + + # ParallelLMHead-style: a bias slot but no has_bias attribute. + lm_head = torch.nn.Module() + lm_head.register_parameter("bias", None) + build(lm_head) + assert lm_head.output_partition_sizes == [32, 32] + assert lm_head.has_bias is False + + # LinearBase already decided has_bias; we must not overwrite it. + linear = torch.nn.Module() + linear.has_bias = True + build(linear) + assert linear.has_bias is True @pytest.mark.parametrize( @@ -654,40 +739,28 @@ def test_modelopt_nvfp4_moe_dispatches_to_marlin_when_w4a16( @pytest.mark.parametrize( - "per_layer_algo, expected_linear_cls_name", + "per_layer_algo, expected_weight, expected_activation", [ - ("NVFP4", "ModelOptNvFp4LinearMethod"), - ("W4A16_NVFP4", "ModelOptNvFp4W4A16LinearMethod"), + ("NVFP4", kNvfp4Static, kNvfp4Dynamic), + ("W4A16_NVFP4", kNvfp4Static, None), + ("FP8", kFp8StaticTensorSym, kFp8StaticTensorSym), + ("MXFP8", kMxfp8Static, kMxfp8Dynamic), ], ) -def test_modelopt_mixed_precision_dispatches_w4a16_layer( - per_layer_algo, expected_linear_cls_name +def test_modelopt_mixed_precision_dispatches_linear_layer( + per_layer_algo, expected_weight, expected_activation ): - """``ModelOptMixedPrecisionConfig.get_quant_method`` must route a Linear - layer to the right LinearMethod based on its per-layer ``quant_algo`` - entry in ``quantized_layers``. Verifies the new ``W4A16_NVFP4`` branch - coexists with the existing ``NVFP4`` branch without regression. A - regression here would mean a W4A16 layer in a mixed-precision ckpt - silently fell through to ``UnquantizedLinearMethod``. - - NOTE: FP8 dispatch (the third branch of get_quant_method) is not - covered here because ``ModelOptFp8LinearMethod.__init__`` reads - ``get_current_vllm_config().model_config.dtype``, which requires a - fully constructed ``ModelConfig`` (real model path). FP8 routing in - mixed-precision is exercised by the existing integration tests - above that use the ``vllm_runner`` fixture (e.g. - ``test_modelopt_fp8_checkpoint_setup``). Our PR doesn't change the - FP8 branch, so this isn't a coverage gap. + """``ModelOptMixedPrecisionConfig.get_quant_method`` routes a Linear layer + to the generic ``ModelOptLinearMethod`` with the ``QuantSpec`` resolved + from its per-layer ``quant_algo`` entry in ``quantized_layers``. A + regression here would mean a layer got the wrong ``(weight, activation)`` + key pair or fell through to ``UnquantizedLinearMethod`` — e.g. a W4A16 + layer picking up a dynamic fp4 activation key (cutlass W4A4 path) instead + of the weight-only Marlin path. """ from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization import modelopt as m - if ( - expected_linear_cls_name == "ModelOptNvFp4W4A16LinearMethod" - and current_platform.is_rocm() - ): - pytest.skip("ModelOptNvFp4W4A16LinearMethod is not supported with rocm") - hf_quant_config: dict[str, Any] = { "quantization": { "quant_algo": "MIXED_PRECISION", @@ -704,10 +777,9 @@ def test_modelopt_mixed_precision_dispatches_w4a16_layer( fake_layer = MagicMock(spec=LinearBase) method = config.get_quant_method(fake_layer, "model.layers.0.fake_proj") - expected_cls = getattr(m, expected_linear_cls_name) - assert isinstance(method, expected_cls), ( - f"Expected {expected_linear_cls_name}, got {type(method).__name__}" - ) + assert isinstance(method, m.ModelOptLinearMethod) + assert method.spec.weight is expected_weight + assert method.spec.activation is expected_activation def test_modelopt_mixed_precision_builds_w4a16_sibling_config(): @@ -734,6 +806,94 @@ def test_modelopt_mixed_precision_builds_w4a16_sibling_config(): config = m.ModelOptMixedPrecisionConfig.from_config(hf_quant_config) assert config.nvfp4_config.quant_method == "NVFP4" - assert config.nvfp4_config.LinearMethodCls is m.ModelOptNvFp4LinearMethod assert config.w4a16_nvfp4_config.quant_method == "W4A16_NVFP4" - assert config.w4a16_nvfp4_config.LinearMethodCls is m.ModelOptNvFp4W4A16LinearMethod + + +def test_modelopt_fp8_pb_wo_hides_output_padding(monkeypatch): + """FP8_PB_WO output width that is not a multiple of 128 (a partial trailing + block) is padded up to a block boundary before the kernel post-load, the + GEMM runs on the padded weight, and the output is trimmed back to the + logical width with bias added after. Faithful port of wei-zhao #53132's + test_modelopt_fp8_pb_wo_hides_output_padding for the generic method + + _Fp8PbWoPartialBlock FormatScheme. + + Width is the real motivating case -- GLM's fused qkv_a_proj, q_a 2048 + + kv_a 576 = 2624 (replicated, so no TP degree makes it a 128-multiple), + padded to 2688 = 21 * 128. + """ + from vllm.config.quantization import QuantSpec + from vllm.model_executor.layers.quantization import modelopt as mo + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Dynamic128Sym, + kFp8Static128BlockSym, + ) + + kernel = Mock() + monkeypatch.setattr(mo, "select_linear_kernel", lambda spec, layer, rt: kernel) + monkeypatch.setattr(mo, "expose_input_quant_key", lambda layer, k: None) + + method = ModelOptLinearMethod.__new__(ModelOptLinearMethod) + method.spec = QuantSpec(weight=kFp8Static128BlockSym, activation=kFp8Dynamic128Sym) + method.ctx = mo.CkptCtx() + method.fmt = mo._PB_WO_PARTIAL_BLOCK + method.wkey = mo.SCHEME_FOR[kFp8Static128BlockSym] + method.akey = mo.SCHEME_FOR[kFp8Dynamic128Sym] + method.input_dtype = method.out_dtype = torch.bfloat16 + method.marlin_input_dtype = None + + layer = torch.nn.Module() + with ( + patch( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", + return_value=0, + ), + patch( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", + return_value=1, + ), + ): + # output 2624 = 2048 + 576, input 128 + method.create_weights( + layer, 128, [2048, 576], 128, 2624, torch.bfloat16, weight_loader=Mock() + ) + + # loaded at logical size; scale is cdiv(2624, 128) = 21 block rows + assert layer.weight.shape == (2624, 128) + assert layer.weight_scale.shape == (21, 1, 1, 1) + + layer.weight.data.fill_(1) + method.process_weights_after_loading(layer) + + # weight padded to the block boundary; the pad rows are zero + assert layer.weight.shape == (2688, 128) + assert torch.count_nonzero(layer.weight[2624:].float()) == 0 + kernel.process_weights_after_loading.assert_called_once_with(layer) + + # apply: GEMM on padded weight (bias=None), output trimmed + bias added + physical_output = torch.randn(4, 2688, dtype=torch.bfloat16) + kernel.apply_weights.return_value = physical_output + bias = torch.randn(2624, dtype=torch.bfloat16) + output = method.apply(layer, torch.randn(4, 128), bias) + + torch.testing.assert_close(output, physical_output[:, :2624] + bias) + assert output.shape == (4, 2624) + assert output.is_contiguous() + kernel.apply_weights.assert_called_once() + assert kernel.apply_weights.call_args.kwargs["bias"] is None + + +def test_modelopt_fp8_pb_wo_rejects_non_128_input(): + """Input width must still be a multiple of 128 (same as #53132, which only + pads the output). A partial input block is refused loudly rather than + silently loading wrong scales.""" + from vllm.model_executor.layers.quantization import modelopt as mo + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Static128BlockSym, + ) + + scheme = mo.SCHEME_FOR[kFp8Static128BlockSym] + shapes = mo.Shapes([128], 100, torch.bfloat16) # input 100 not divisible by 128 + with pytest.raises(ValueError, match="in divisible by 128"): + scheme.create_weights( + torch.nn.Module(), mo.WEIGHT, mo.CkptCtx(), shapes, Mock() + ) diff --git a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py index 28c958d4fd..c8b8582c0e 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py +++ b/vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py @@ -6,6 +6,8 @@ from dataclasses import dataclass import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + @dataclass class Mxfp8LinearLayerConfig: @@ -30,6 +32,12 @@ class Mxfp8LinearKernel(ABC): assert self.is_supported()[0] self.config = c + def input_quant_key(self) -> QuantKey | None: + """Return the input quantization key supported by this kernel. If the kernel + does not support input quantization outside of the kernel, return None. + """ + return None + @classmethod @abstractmethod def is_supported( diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 479da9721c..55f7d82882 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -377,12 +377,18 @@ def _get_kv_b_proj_input_dtype( return None if weight_dtype == current_platform.fp8_dtype(): from vllm.model_executor.layers.quantization.modelopt import ( - ModelOptFp8PbWoLinearMethod, + ModelOptLinearMethod, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kFp8Static128BlockSym, ) quant_method = getattr(kv_b_proj, "quant_method", None) # FP8_PB_WO dynamically quantizes BF16/FP16 inputs in the linear method. - if isinstance(quant_method, ModelOptFp8PbWoLinearMethod): + if ( + isinstance(quant_method, ModelOptLinearMethod) + and quant_method.spec.weight is kFp8Static128BlockSym + ): return quant_method.input_dtype if not use_fp8_prefill: return None diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 817409c0c4..a33f3a1166 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -55,13 +55,10 @@ WEIGHT_LOADER_V2_SUPPORTED = [ "AutoGPTQLinearMethod", "Fp8LinearMethod", "FBGEMMFp8LinearMethod", - "ModelOptFp8LinearMethod", - "ModelOptFp8PcPtLinearMethod", - "ModelOptFp8PbWoLinearMethod", "QuarkLinearMethod", - "ModelOptNvFp4LinearMethod", - "ModelOptNvFp4W4A16LinearMethod", "HummingLinearMethod", + # ModelOptLinearMethod self-registers via + # register_weight_loader_v2_supported_method (see modelopt.py). ] diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index aa11f9269e..465e8bd573 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum from fnmatch import fnmatch from typing import TYPE_CHECKING, Any, cast @@ -8,7 +11,8 @@ import torch from torch.nn.parameter import Parameter import vllm.envs as envs -from vllm.config import get_current_vllm_config +from vllm.config import get_current_vllm_config_or_none +from vllm.config.quantization import QuantSpec from vllm.logger import init_logger from vllm.model_executor.kernels.linear import ( init_fp8_linear_kernel, @@ -49,6 +53,7 @@ from vllm.model_executor.layers.linear import ( LinearBase, LinearMethodBase, UnquantizedLinearMethod, + register_weight_loader_v2_supported_method, ) from vllm.model_executor.layers.quantization import QuantizationMethods from vllm.model_executor.layers.quantization.base_config import ( @@ -57,6 +62,7 @@ from vllm.model_executor.layers.quantization.base_config import ( ) from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + FP8_SCALE_SENTINEL, process_fp8_input_tensor_strategy_moe, process_fp8_weight_channel_strategy, process_fp8_weight_tensor_strategy_moe, @@ -70,12 +76,16 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( MXFP8_VALUE_DTYPE, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( - GroupShape, - create_fp8_quant_key, + FP4_DTYPE, + QuantKey, is_layer_skipped, + kFp8Dynamic128Sym, kFp8DynamicTokenSym, + kFp8Static128BlockSym, kFp8StaticTensorSym, kFp8StaticTokenSym, + kMxfp8Dynamic, + kMxfp8Static, kNvfp4Dynamic, kNvfp4Static, ) @@ -86,33 +96,49 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.parameter import ( BlockQuantScaleParameter, ChannelQuantScaleParameter, - GroupQuantScaleParameter, ModelWeightParameter, PerTensorScaleParameter, ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.utils.math_utils import cdiv if TYPE_CHECKING: from vllm.model_executor.models.utils import WeightsMapper logger = init_logger(__name__) -QUANT_ALGOS = [ - # FP8 (per-tensor weight + optional static activation scale). - "FP8", +# Single source of truth for the ModelOpt linear algos. +# +# Each entry is (owning config's name, mixed-precision sub-config attribute). +# ``resolve`` turns the algo into a QuantSpec; this table only records which +# config object holds the checkpoint parameters resolve may read -- today just +# group_size, and only for the NVFP4 family. +# +# To add a format: one row here plus one row in resolve(). Every consumer below +# is derived, so nothing else needs editing. +LINEAR_ALGOS: dict[str, tuple[str, str]] = { + # FP8 per-tensor weight + optional static activation scale. + "FP8": ("modelopt", "fp8_config"), # FP8 per-channel weight scale + per-token activation scale. - "FP8_PER_CHANNEL_PER_TOKEN", - # FP8 per-block weight-only (ModelOpt may emit this as lowercase). - "FP8_PB_WO", - # NVFP4 W4A4 (4-bit float weights AND 4-bit float activations). - "NVFP4", - # W4A16 NVFP4 (4-bit float weights, fp16/bf16 activations). - "W4A16_NVFP4", - # MXFP8 - "MXFP8", - # MIXED_PRECISION, - "MIXED_PRECISION", -] + "FP8_PER_CHANNEL_PER_TOKEN": ("modelopt", "fp8_config"), + # FP8 128x128 per-block weight (ModelOpt may emit this lowercase). + "FP8_PB_WO": ("modelopt", "fp8_config"), + # NVFP4 W4A4: 4-bit float weights AND 4-bit float activations. + "NVFP4": ("modelopt_fp4", "nvfp4_config"), + # W4A16 NVFP4: 4-bit float weights, fp16/bf16 activations. + "W4A16_NVFP4": ("modelopt_fp4", "w4a16_nvfp4_config"), + # MXFP8: e4m3 weights with per-32-block e8m0 scales. + "MXFP8": ("modelopt_mxfp8", "mxfp8_config"), +} + +# MIXED_PRECISION is not a linear algo; it selects a per-layer algo from the +# table above. +QUANT_ALGOS = [*LINEAR_ALGOS, "MIXED_PRECISION"] + + +def algos_owned_by(config_name: str) -> tuple[str, ...]: + """The linear algos a given config accepts, for its quant_algo validation.""" + return tuple(a for a, (owner, _) in LINEAR_ALGOS.items() if owner == config_name) class ModelOptKVCacheMethod(BaseKVCacheMethod): @@ -125,7 +151,11 @@ class ModelOptKVCacheMethod(BaseKVCacheMethod): class ModelOptQuantConfigBase(QuantizationConfig): - LinearMethodCls: type = LinearMethodBase + # ModelOpt quant-algo string, set by each subclass. Fed to resolve() to + # build the QuantSpec for the generic ModelOptLinearMethod. The mixed + # config resolves the algo per-prefix instead and does not set this. + quant_method: str + FusedMoEMethodCls: type = FusedMoEMethodBase KVCacheMethodCls: type = BaseKVCacheMethod @@ -201,10 +231,7 @@ class ModelOptQuantConfigBase(QuantizationConfig): # now, the layer is quantized, handle it here if isinstance(layer, (LinearBase, ParallelLMHead)): - quant_method = self.LinearMethodCls(self) - if getattr(quant_method, "backend", "") == "marlin": - quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix) - return quant_method + return build_linear_method(self, self.quant_method, prefix) elif isinstance(layer, RoutedExperts): quant_method = self.FusedMoEMethodCls( quant_config=self, moe_config=layer.moe_config @@ -375,27 +402,22 @@ class ModelOptFp8Config(ModelOptQuantConfigBase): self.quant_method = quant_method self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized self.kv_cache_quant_method = kv_cache_quant_method - if is_checkpoint_fp8_serialized: - logger.warning( - "Detected ModelOpt fp8 checkpoint (quant_algo=%s). Please note " - "that the format is experimental and could change.", - quant_method, - ) - # Select LinearMethod implementation based on quant_algo. - if self.quant_method == "FP8": - self.LinearMethodCls = ModelOptFp8LinearMethod - elif self.quant_method == "FP8_PER_CHANNEL_PER_TOKEN": - self.LinearMethodCls = ModelOptFp8PcPtLinearMethod - elif self.quant_method == "FP8_PB_WO": - self.LinearMethodCls = ModelOptFp8PbWoLinearMethod - else: + # Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch. + supported = algos_owned_by("modelopt") + if self.quant_method not in supported: raise ValueError( "Unsupported ModelOpt FP8 quant_algo for vLLM: " - f"{self.quant_method}. Supported: FP8 / " - "FP8_PER_CHANNEL_PER_TOKEN / FP8_PB_WO." + f"{self.quant_method}. Supported: {' / '.join(supported)}." ) + def has_blocked_weights(self) -> bool: + # Gates "+quant_fp8" in VllmConfig. FP8_PB_WO runs the block-scaled + # kernels, whose DeepGEMM path needs UE8M0-packed group scales; without + # the op enabled QuantFP8 falls back to forward_native, which emits + # unpacked fp32 scales and the GEMM fails to launch. + return self.quant_method == "FP8_PB_WO" + def get_name(self) -> QuantizationMethods: return "modelopt" @@ -435,334 +457,6 @@ class ModelOptFp8Config(ModelOptQuantConfigBase): ) -class ModelOptFp8LinearMethod(LinearMethodBase): - """Linear method for Model Optimizer static quantization. - Supports loading FP8 checkpoints with static weight scale and - activation scale. Future support might be added for dynamic - scales. - - Limitations: - 1. Only support per-tensor quantization due to torch._scaled_mm support. - 2. Only support float8_e4m3fn datatype - Args: quant_config: The ModelOpt quantization config. - """ - - def __init__(self, quant_config: ModelOptFp8Config) -> None: - self.quant_config = quant_config - self.out_dtype = get_current_vllm_config().model_config.dtype - self.input_dtype = get_current_vllm_config().model_config.dtype - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - layer.orig_dtype = params_dtype - weight_dtype = ( - torch.float8_e4m3fn - if self.quant_config.is_checkpoint_fp8_serialized - else params_dtype - ) - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, input_size_per_partition, dtype=weight_dtype - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - if self.quant_config.is_checkpoint_fp8_serialized: - # WEIGHT SCALE - weight_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - weight_scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("weight_scale", weight_scale) - # INPUT SCALE - scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - - scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("input_scale", scale) - - self.fp8_linear = init_fp8_linear_kernel( - activation_quant_key=kFp8StaticTensorSym, - weight_quant_key=kFp8StaticTensorSym, - weight_shape=layer.weight.shape, - input_dtype=self.input_dtype, - out_dtype=self.out_dtype, - module_name=self.__class__.__name__, - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - weight = layer.weight - max_w_scale = layer.weight_scale.max() - if not (layer.weight_scale == layer.weight_scale[0]).all(): - max_w_scale, weight = requantize_with_max_scale( - layer.weight, layer.weight_scale, layer.logical_widths - ) - layer.weight = Parameter(weight.t(), requires_grad=False) - layer.weight.input_dim = 0 - layer.weight.output_dim = 1 - layer.weight_scale = Parameter(max_w_scale, requires_grad=False) - layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) - self.fp8_linear.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - return self.fp8_linear.apply_weights(layer, x, bias) - - -class ModelOptFp8PcPtLinearMethod(LinearMethodBase): - """Linear method for ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoints. - - Expected checkpoint structure (per Linear): - - weight: fp8-e4m3fn, shape [out, in] - - weight_scale: fp32, shape [out] (per-output-channel) - - no input_scale (activations are dynamically quantized per-token) - """ - - def __init__(self, quant_config: ModelOptFp8Config) -> None: - self.quant_config = quant_config - self.out_dtype = get_current_vllm_config().model_config.dtype - self.input_dtype = get_current_vllm_config().model_config.dtype - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - - if not self.quant_config.is_checkpoint_fp8_serialized: - raise ValueError( - "FP8_PER_CHANNEL_PER_TOKEN currently only supports " - "FP8-serialized checkpoints." - ) - - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition, - dtype=torch.float8_e4m3fn, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - weight_scale = ChannelQuantScaleParameter( - data=torch.empty(output_size_per_partition, dtype=torch.float32), - output_dim=0, - weight_loader=weight_loader, - ) - weight_scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("weight_scale", weight_scale) - - self.fp8_linear = init_fp8_linear_kernel( - activation_quant_key=kFp8DynamicTokenSym, - weight_quant_key=kFp8StaticTokenSym, - weight_shape=layer.weight.shape, - input_dtype=self.input_dtype, - out_dtype=self.out_dtype, - module_name=self.__class__.__name__, - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - weight, weight_scale, _ = process_fp8_weight_channel_strategy( - layer.weight, layer.weight_scale.data - ) - layer.weight = Parameter(weight.t(), requires_grad=False) - layer.weight_scale = Parameter(weight_scale, requires_grad=False) - self.fp8_linear.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - return self.fp8_linear.apply_weights(layer, x, bias) - - -class ModelOptFp8PbWoLinearMethod(LinearMethodBase): - """Linear method for ModelOpt FP8_PB_WO checkpoints. - - ModelOpt exports `weight_scale` as a 4D tensor: - [out_blk, 1, in_blk, 1] - where block size is typically 128 for both dims. - - vLLM executes it as FP8 GEMM with *dynamic per-token* activation quant. - Output widths that are not block-aligned are padded and restored - to their logical width before returning to the model. - """ - - _WEIGHT_BLOCK_SIZE: tuple[int, int] = (128, 128) - - def __init__(self, quant_config: ModelOptFp8Config) -> None: - self.quant_config = quant_config - block_n, block_k = self._WEIGHT_BLOCK_SIZE - self.weight_block_size = list(self._WEIGHT_BLOCK_SIZE) - - self.activation_quant_key = create_fp8_quant_key( - static=False, group_shape=GroupShape(1, block_k) - ) - self.weight_quant_key = create_fp8_quant_key( - static=True, group_shape=GroupShape(block_n, block_k) - ) - - self.out_dtype = get_current_vllm_config().model_config.dtype - self.input_dtype = get_current_vllm_config().model_config.dtype - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - - if not self.quant_config.is_checkpoint_fp8_serialized: - raise ValueError( - "FP8_PB_WO currently only supports FP8-serialized checkpoints." - ) - - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - - # Expose block size so the v2 weight loaders can translate offsets from - # element-space -> block-space for BlockQuantScaleParameter. - layer.weight_block_size = self.weight_block_size - - block_n, block_k = self._WEIGHT_BLOCK_SIZE - remainder = output_size_per_partition % block_n - self.output_padding = 0 if remainder == 0 else block_n - remainder - self.logical_output_size = output_size_per_partition - physical_output_size = output_size_per_partition + self.output_padding - - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition, - dtype=torch.float8_e4m3fn, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - if input_size_per_partition % block_k != 0: - raise ValueError( - "ModelOpt FP8_PB_WO requires in_features divisible by " - f"{block_k}, got {input_size_per_partition}." - ) - - out_blks = physical_output_size // block_n - in_blks = input_size_per_partition // block_k - - # Match ModelOpt's exported shape so weight loading works without a - # custom loader: [out_blk, 1, in_blk, 1] - weight_scale = BlockQuantScaleParameter( - data=torch.empty((out_blks, 1, in_blks, 1), dtype=torch.float32), - input_dim=2, - output_dim=0, - weight_loader=weight_loader, - ) - weight_scale[:] = torch.finfo(torch.float32).min - layer.register_parameter("weight_scale", weight_scale) - - self.w8a8_block_fp8_linear = init_fp8_linear_kernel( - activation_quant_key=self.activation_quant_key, - weight_quant_key=self.weight_quant_key, - weight_shape=(physical_output_size, input_size_per_partition), - input_dtype=self.input_dtype, - out_dtype=self.out_dtype, - module_name=self.__class__.__name__, - ) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Keep weight in [out, in] layout for Fp8BlockScaledMMLinearKernel. - weight = layer.weight.data - if self.output_padding: - padded_weight = weight.new_zeros( - self.logical_output_size + self.output_padding, - weight.shape[1], - ) - padded_weight[: self.logical_output_size].copy_(weight) - weight = padded_weight - layer.weight = Parameter(weight, requires_grad=False) - - scale = layer.weight_scale - if scale.dim() == 4: - # [out_blk, 1, in_blk, 1] -> [out_blk, in_blk] - scale = scale.squeeze(1).squeeze(-1) - elif scale.dim() != 2: - raise ValueError( - "Unexpected ModelOpt FP8_PB_WO weight_scale shape: " - f"{tuple(scale.shape)}." - ) - - layer.weight_scale = Parameter(scale.contiguous(), requires_grad=False) - - self.w8a8_block_fp8_linear.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - kernel_bias = None if self.output_padding else bias - output = self.w8a8_block_fp8_linear.apply_weights(layer, x, kernel_bias) - if not self.output_padding: - return output - - output = output[..., : self.logical_output_size].contiguous() - if bias is not None: - output.add_(bias) - return output - - class ModelOptFp8MoEMethod(FusedMoEMethodBase): """MoE method for ModelOpt FP8. Supports loading FP8 checkpoints with static weight scale and @@ -1004,7 +698,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): ) -ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod ModelOptFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod @@ -1026,27 +719,17 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase): self.quant_method = quant_method self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized if is_checkpoint_nvfp4_serialized: - logger.warning( - "Detected ModelOpt NVFP4 checkpoint (quant_algo=%s). Please " - "note that the format is experimental and could change in " - "future.", - quant_method, - ) - self.group_size = group_size self.kv_cache_quant_algo = kv_cache_quant_algo - # Select LinearMethod implementation based on quant_algo (FP8 pattern). - # NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization - # W4A16_NVFP4 -> W4A16: flashinfer cute-dsl or Marlin with 16-bit inputs - if quant_method == "NVFP4": - self.LinearMethodCls = ModelOptNvFp4LinearMethod - elif quant_method == "W4A16_NVFP4": - self.LinearMethodCls = ModelOptNvFp4W4A16LinearMethod - else: + # Validate the quant_algo; resolve() maps it to a QuantSpec at dispatch. + # NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization + # W4A16_NVFP4 -> W4A16: FP4 Marlin GEMM with bf16/fp16 activations + supported = algos_owned_by("modelopt_fp4") + if quant_method not in supported: raise ValueError( f"Unsupported ModelOpt NVFP4 quant_algo: {quant_method}. " - "Supported: NVFP4 / W4A16_NVFP4." + f"Supported: {' / '.join(supported)}." ) def get_name(self) -> QuantizationMethods: @@ -1123,285 +806,6 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase): ) -class ModelOptNvFp4LinearMethod(LinearMethodBase): - """Linear method for Model Optimizer NVFP4. - Supports loading NVFP4 checkpoints with the following structure: - - input_scale: torch.float32, scalar , - weight: NVFP4(represented as byte) Shape: [1, X, y/2] - weight_scale: FP8-E4M3, Shape: [X, Y], aka per block scale, - weight_scale_2: torch.float32, scalar, - Args: quant_config: The ModelOpt quantization config. - """ - - def __init__(self, quant_config: ModelOptNvFp4Config) -> None: - self.quant_config = quant_config - self.marlin_input_dtype = None - self.kernel = init_nvfp4_linear_kernel() - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - if not self.quant_config.is_checkpoint_nvfp4_serialized: - raise ValueError( - "NVFP4 quantization was selected, " - " dynamic quantization is not supported." - ) - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - - if input_size_per_partition % 16 != 0: - raise ValueError( - "Unsupported model when in features size is not multiple of 16" - ) - # The nvfp4 weight is still represented as - weight_dtype = ( - torch.float8_e4m3fn - if self.quant_config.is_checkpoint_nvfp4_serialized - else params_dtype - ) - # Weight - weight = ModelWeightParameter( - data=torch.empty( - # 2 fp4 items are packed in the input dimension - layer.output_size_per_partition, - layer.input_size_per_partition // 2, - dtype=torch.uint8, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - # Input Global Scale - input_global_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - layer.register_parameter("input_scale", input_global_scale) - - # Weight Global Scale - weight_global_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - layer.register_parameter("weight_scale_2", weight_global_scale) - - # Per Block Weight Scale - weight_scale = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // self.quant_config.group_size, - dtype=weight_dtype, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - - layer.register_parameter("weight_scale", weight_scale) - - expose_input_quant_key(layer, self.kernel) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if ( - torch.unique(layer.input_scale).numel() != 1 - or torch.unique(layer.weight_scale_2).numel() != 1 - ): - logger.warning_once( - "In NVFP4 linear, the global scale for input or weight are different" - " for parallel layers (e.g. q_proj, k_proj, v_proj). This " - " will likely results in reduce accuracy. Please verify the model" - " accuracy. Consider using a checkpoint with a shared global NVFP4" - " scale for parallel layers." - ) - - # Rename ModelOpt checkpoint names to standardized names - input_global_scale = layer.input_scale.max().to(torch.float32) - layer.input_global_scale = Parameter(input_global_scale, requires_grad=False) - del layer.input_scale - - weight_global_scale = layer.weight_scale_2.max().to(torch.float32) - layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False) - del layer.weight_scale_2 - - # Pre-compute alpha and inverse for runtime quantization - layer.alpha = Parameter( - layer.input_global_scale * layer.weight_global_scale, requires_grad=False - ) - layer.input_global_scale_inv = Parameter( - (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False - ) - - # Convert layer to NVFP4 linear kernel format - self.kernel.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - return self.kernel.apply_weights(layer=layer, x=x, bias=bias) - - -class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase): - """Linear method for ModelOpt NVFP4 W4A16. - - 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names - directly (no on-disk conversion) and dispatches to a W4A16 GEMM: - - weight uint8 packed NVFP4 (2 nibbles/byte along input dim) - weight_scale fp8-e4m3 per 16-elem group along input dim - weight_scale_2 fp32 per-tensor global scale = amax / (6.0 * 448.0) - - No activation quantization. ModelOpt stores the global scale as - amax/2688, so we rename weight_scale_2 -> weight_global_scale without - reciprocation. The selected kernel converts it to its runtime format. - The CT W4A16 path reciprocates because CT stores the inverse on disk. - - We also register a placeholder input_scale parameter so that W4A4-shaped - checkpoints (which contain *_proj.input_scale tensors) can be loaded - under this method without the per-shard loader hitting a KeyError on - the merged-name lookup. The placeholder is discarded in - process_weights_after_loading -- its value is never used. - """ - - def __init__(self, quant_config: ModelOptNvFp4Config) -> None: - self.quant_config = quant_config - self.marlin_input_dtype = None - # `init_nvfp4_linear_kernel(use_a16=True)` is best of both worlds: - # 1. `use_a16=True` forces flashinfer cutedsl for sm100 and `Marlin` for others. - # for `--linear-backend=auto`, avoiding a W4A4 kernel that requires input_scale. - # 2. Specifying e.g. `--linear-backend=humming` will override. - self.kernel = init_nvfp4_linear_kernel(use_a16=True) - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - if not self.quant_config.is_checkpoint_nvfp4_serialized: - raise ValueError( - "W4A16_NVFP4 quantization was selected; " - "dynamic quantization is not supported." - ) - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - layer.output_partition_sizes = output_partition_sizes - - if input_size_per_partition % 16 != 0: - raise ValueError( - "Unsupported model: input feature size is not a multiple of 16." - ) - - # Packed NVFP4 weights: uint8, 2 nibbles per byte along the input dim. - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // 2, - dtype=torch.uint8, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - # Per-tensor global weight scale (fp32). ModelOpt stores - # amax / (NVFP4_max * fp8_e4m3_max) = amax / 2688. PerTensorScaleParameter - # holds one entry per fused output partition (e.g. q/k/v in a fused QKV). - weight_scale_2 = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - layer.register_parameter("weight_scale_2", weight_scale_2) - - # Per-group fp8 weight scale. - weight_scale = GroupQuantScaleParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // self.quant_config.group_size, - dtype=torch.float8_e4m3fn, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight_scale", weight_scale) - - # Placeholder input_scale param so W4A4-shaped checkpoints can be - # loaded under this method without KeyError on the merged-name - # lookup (qwen2-style stacked-loader path renames *_proj.input_scale - # to e.g. qkv_proj.input_scale and looks it up unconditionally). - # Discarded in process_weights_after_loading; never read by the kernel. - # For native W4A16 checkpoints (no input_scale on disk) the param - # stays uninitialized and is simply deleted. - input_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - layer.register_parameter("input_scale", input_scale) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if not hasattr(layer, "has_bias"): - layer.has_bias = getattr(layer, "bias", None) is not None - - # Discard the input_scale placeholder. Whether it carries values - # (W4A4 ckpt loaded as W4A16) or is uninitialized (native W4A16 - # ckpt), W4A16 mode does not quantize activations, so this is unused. - if hasattr(layer, "input_scale"): - del layer.input_scale - - if torch.unique(layer.weight_scale_2).numel() != 1: - logger.warning_once( - "In W4A16_NVFP4 linear, the global weight scale " - "(weight_scale_2) differs across fused parallel layers " - "(e.g. q/k/v_proj). This will likely reduce accuracy. " - "Consider a checkpoint with a shared global scale." - ) - - # Rename weight_scale_2 -> weight_global_scale. NO reciprocation: - # ModelOpt already stores amax/2688, which is exactly what Marlin - # consumes via nvfp4_marlin_process_global_scale (called inside the - # Marlin adapter's process_weights_after_loading). - layer.weight_global_scale = Parameter( - layer.weight_scale_2.max().to(torch.float32), requires_grad=False - ) - del layer.weight_scale_2 - - self.kernel.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - return self.kernel.apply_weights(layer=layer, x=x, bias=bias) - - class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): """ MoE Method for FP4 Quantization. @@ -1685,7 +1089,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase): ) -ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE ModelOptNvFp4Config.KVCacheMethodCls = ModelOptKVCacheMethod @@ -1700,6 +1103,7 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase): exclude_modules: list[str], ) -> None: super().__init__(exclude_modules) + self.quant_method = "MXFP8" self.is_checkpoint_mxfp8_serialized = is_checkpoint_mxfp8_serialized if not is_checkpoint_mxfp8_serialized: @@ -1708,11 +1112,6 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase): "Dynamic quantization is not supported." ) - logger.warning( - "Detected ModelOpt MXFP8 checkpoint. Please note that " - "the format is experimental and could change in future." - ) - self.kv_cache_quant_algo = kv_cache_quant_algo def get_name(self) -> QuantizationMethods: @@ -1783,117 +1182,6 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase): ) -class ModelOptMxFp8LinearMethod(LinearMethodBase): - """Linear method for ModelOpt MXFP8 quantization.""" - - def __init__(self, quant_config: ModelOptMxFp8Config) -> None: - self.quant_config = quant_config - - if not self.quant_config.is_checkpoint_mxfp8_serialized: - raise ValueError( - "MXFP8 currently only supports serialized checkpoints. " - "Dynamic quantization is not supported." - ) - - self.kernel = init_mxfp8_linear_kernel() - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - - if not self.quant_config.is_checkpoint_mxfp8_serialized: - raise ValueError( - "MXFP8 quantization was selected, but checkpoint is not " - "MXFP8 serialized. Dynamic quantization is not supported." - ) - - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - - if input_size_per_partition % MXFP8_BLOCK_SIZE != 0: - raise ValueError( - f"MXFP8 requires input dimension to be divisible by " - f"{MXFP8_BLOCK_SIZE}, got {input_size_per_partition}" - ) - - # Weight tensor: FP8 E4M3 format - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition, - dtype=MXFP8_VALUE_DTYPE, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - # Weight scale tensor (E8M0 encoded as uint8), one scale per block of 32 along K - weight_scale = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // MXFP8_BLOCK_SIZE, - dtype=MXFP8_SCALE_DTYPE, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight_scale", weight_scale) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Idempotent: the emulation kernel may dequant the weight to BF16 at load - # time (>=2-byte). If already converted, there is nothing left to do -- - # avoid re-running the MXFP8-only validation/conversion below. - if layer.weight.element_size() >= 2: - return - - # Validate weight tensor - if layer.weight.ndim != 2: - raise ValueError( - f"MXFP8 weight must be 2D tensor [N, K], got {layer.weight.ndim}D " - f"with shape {tuple(layer.weight.shape)}" - ) - - if layer.weight.dtype != MXFP8_VALUE_DTYPE: - raise ValueError( - f"MXFP8 weight must be {MXFP8_VALUE_DTYPE} (FP8 E4M3), " - f"got {layer.weight.dtype}. The checkpoint may not be properly " - f"quantized with MXFP8." - ) - - # Validate weight scale tensor (should be 2D, not swizzled) - assert layer.weight_scale.ndim == 2, ( - f"MXFP8 weight scale must be 2D, got {layer.weight_scale.ndim}D" - ) - assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE, ( - f"MXFP8 weight scale must be {MXFP8_SCALE_DTYPE}," - f" got {layer.weight_scale.dtype}" - ) - - self.kernel.process_weights_after_loading(layer) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - return self.kernel.apply_weights(layer, x, bias) - - class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): """FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints.""" @@ -2171,7 +1459,6 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase): # Register the method classes for ModelOptMxFp8Config -ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod ModelOptMxFp8Config.FusedMoEMethodCls = ModelOptMxFp8FusedMoE ModelOptMxFp8Config.KVCacheMethodCls = ModelOptKVCacheMethod @@ -2204,6 +1491,14 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): self.w4a16_nvfp4_config = w4a16_nvfp4_config self.mxfp8_config = mxfp8_config + def has_blocked_weights(self) -> bool: + # Same gate as ModelOptFp8Config.has_blocked_weights, resolved per + # layer: "+quant_fp8" must be on as soon as any layer is block-scaled. + return any( + info.get("quant_algo", "").upper() == "FP8_PB_WO" + for info in self.quantized_layers.values() + ) + def get_name(self) -> QuantizationMethods: return "modelopt_mixed" @@ -2281,10 +1576,10 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): group_size=group_size, ) # Sibling config for layers that declare quant_algo: "W4A16_NVFP4". - # ModelOptNvFp4Config.__init__ keys LinearMethodCls off quant_method, - # so this instance auto-selects ModelOptNvFp4W4A16LinearMethod. The - # MoE side reads quant_config.quant_method == "W4A16_NVFP4" to set - # use_a16 → Marlin backend in ModelOptNvFp4FusedMoE.__init__. + # get_quant_method resolves this sub-config to the (kNvfp4Static, None) + # QuantSpec (W4A16) for the generic ModelOptLinearMethod. The MoE side + # reads quant_config.quant_method == "W4A16_NVFP4" to set use_a16 → + # Marlin backend in ModelOptNvFp4FusedMoE.__init__. w4a16_nvfp4_config = ModelOptNvFp4Config( quant_method="W4A16_NVFP4", is_checkpoint_nvfp4_serialized=True, @@ -2424,18 +1719,15 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): quant_algo = self._resolve_quant_algo(prefix) if isinstance(layer, (LinearBase, ParallelLMHead)): - if quant_algo == "FP8": - return ModelOptFp8LinearMethod(self.fp8_config) - if quant_algo == "FP8_PB_WO": - return ModelOptFp8PbWoLinearMethod(self.fp8_config) - if quant_algo == "NVFP4": - return ModelOptNvFp4LinearMethod(self.nvfp4_config) - if quant_algo == "W4A16_NVFP4": - return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config) - if quant_algo == "MXFP8": - return ModelOptMxFp8LinearMethod(self.mxfp8_config) - # Layer not in quantized_layers — leave unquantized - return UnquantizedLinearMethod() + # Per-prefix algo -> its sub-config, then the generic linear method. + # resolve() reads group_size off whichever sub-config it is handed. + # Read-only — never write back, or a linear-only change leaks into + # the MoE method that shares the sub-config. + if quant_algo is None or quant_algo not in LINEAR_ALGOS: + # Layer not in quantized_layers — leave unquantized + return UnquantizedLinearMethod() + _, subcfg_attr = LINEAR_ALGOS[quant_algo] + return build_linear_method(getattr(self, subcfg_attr), quant_algo, prefix) if isinstance(layer, RoutedExperts): if quant_algo == "FP8": @@ -2466,3 +1758,752 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): super().apply_vllm_mapper(hf_to_vllm_mapper) if self.quantized_layers: self.quantized_layers = hf_to_vllm_mapper.apply_dict(self.quantized_layers) + + +# =========================================================================== +# Generic QuantKey-driven linear method +# +# ``ModelOptLinearMethod`` serves every ModelOpt linear format. It composes a +# per-QuantKey weight scheme + activation scheme (the ``QuantSpec`` pair that +# ``resolve`` produces from the checkpoint's quant_algo), runs a fixed +# create/process lifecycle, selects the kernel from the pair, and applies. +# +# Adding a format: +# * Composes as a (weight, activation) key pair -> add a QuantKeyScheme per +# new key to SCHEME_FOR, plus a resolve() row returning the QuantSpec. No +# new method class; this is how every current format is built. +# * Needs format-wide state but the same lifecycle -> also return a +# FormatScheme subclass from that resolve() row (extra_weights / pre_process +# / post_process hooks). +# * Needs a different lifecycle entirely -> write a LinearMethodBase and +# register it in LINEAR_METHOD_BUILDERS by algo. +# In all cases accept the algo in the owning config's quant_method validation, +# so get_quant_method can dispatch it. +# =========================================================================== + + +class Role(Enum): + WEIGHT = "weight" + ACT = "activation" + + +WEIGHT = Role.WEIGHT +ACT = Role.ACT + + +@dataclass(frozen=True) +class CkptCtx: + """Per-checkpoint facts a QuantKey cannot carry.""" + + group_size: int | None = None + + +@dataclass(frozen=True) +class RuntimeDtypes: + """Runtime/model dtypes kernels need — format-agnostic.""" + + input_dtype: torch.dtype + out_dtype: torch.dtype + marlin_input_dtype: torch.dtype | None = None + + +@dataclass(frozen=True) +class Shapes: + """Layer geometry from create_weights args.""" + + output_partition_sizes: list[int] + input_size_per_partition: int + params_dtype: torch.dtype + + @property + def output_size_per_partition(self) -> int: + return sum(self.output_partition_sizes) + + @property + def num_partitions(self) -> int: + return len(self.output_partition_sizes) + + +class QuantKeyScheme: + """One scheme per QuantKey. Selected by key *content*; the base supplies the + ``role`` from the slot it fills the key into (wkey->WEIGHT, akey->ACT). + + Every scheme branches on ``role`` explicitly and ``reject``s any role it has + not validated — never falling through to the wrong role's registration, + which would allocate the wrong parameters and produce garbage silently. + """ + + key: QuantKey + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + raise NotImplementedError + + def process(self, layer, role) -> None: + pass + + @staticmethod + def reject(role) -> None: + raise NotImplementedError( + f"role {role!r} not validated for this QuantKey scheme" + ) + + @staticmethod + def register_params(layer, name, shape, dtype, cls, wl, *, init=None, **dims): + p = cls(data=torch.empty(shape, dtype=dtype), weight_loader=wl, **dims) + if init is not None: + p.data.fill_(init) + layer.register_parameter(name, p) + + +class KNvfp4Static(QuantKeyScheme): + """NVFP4 weight scheme (W4A4 and W4A16 share it). Weight-role only today.""" + + key = kNvfp4Static + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is not WEIGHT: + self.reject(role) + if shapes.input_size_per_partition % 16 != 0: + raise ValueError( + "Unsupported model when in features size is not multiple of 16" + ) + # Packed NVFP4 weight: 2 fp4 items per byte along the input dim. + self.register_params( + layer, + "weight", + (shapes.output_size_per_partition, shapes.input_size_per_partition // 2), + torch.uint8, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + # Per-tensor global weight scale. + self.register_params( + layer, + "weight_scale_2", + (shapes.num_partitions,), + torch.float32, + PerTensorScaleParameter, + wl, + ) + # Per-block (group_size) weight scale. + self.register_params( + layer, + "weight_scale", + ( + shapes.output_size_per_partition, + shapes.input_size_per_partition // ctx.group_size, + ), + torch.float8_e4m3fn, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + + def process(self, layer, role) -> None: + if role is not WEIGHT: + self.reject(role) + if torch.unique(layer.weight_scale_2).numel() != 1: + logger.warning_once( + "In NVFP4 linear, the global weight scale differs across " + "parallel layers (e.g. q_proj, k_proj, v_proj). This will " + "likely reduce accuracy. Consider a checkpoint with the same " + "global NVFP4 scale for fused weights." + ) + # Raw max, no reciprocation — Marlin/cutlass want ModelOpt's amax/2688. + weight_global_scale = layer.weight_scale_2.max().to(torch.float32) + layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False) + del layer.weight_scale_2 + + +class KNvfp4Dynamic(QuantKeyScheme): + """NVFP4 activation scheme (W4A4). Has a static global input scale on disk; + the per-group scale is computed at runtime inside the kernel.""" + + key = kNvfp4Dynamic + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is not ACT: + self.reject(role) + self.register_params( + layer, + "input_scale", + (shapes.num_partitions,), + torch.float32, + PerTensorScaleParameter, + wl, + ) + + def process(self, layer, role) -> None: + if role is not ACT: + self.reject(role) + if torch.unique(layer.input_scale).numel() != 1: + logger.warning_once( + "In NVFP4 linear, the global input scale differs across " + "parallel layers (e.g. q_proj, k_proj, v_proj). This will " + "likely reduce accuracy. Consider a checkpoint with the same " + "global NVFP4 scale for fused weights." + ) + input_global_scale = layer.input_scale.max().to(torch.float32) + layer.input_global_scale = Parameter(input_global_scale, requires_grad=False) + layer.input_global_scale_inv = Parameter( + (1.0 / layer.input_global_scale).to(torch.float32), + requires_grad=False, + ) + del layer.input_scale + + +class KFp8StaticTensor(QuantKeyScheme): + """Plain per-tensor static FP8 — bivalent: serves BOTH the weight slot and + the activation slot (W8A8). One key in both QuantSpec slots.""" + + key = kFp8StaticTensorSym + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is WEIGHT: + self.register_params( + layer, + "weight", + (shapes.output_size_per_partition, shapes.input_size_per_partition), + torch.float8_e4m3fn, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + layer.orig_dtype = shapes.params_dtype + self.register_params( + layer, + "weight_scale", + (shapes.num_partitions,), + torch.float32, + PerTensorScaleParameter, + wl, + init=FP8_SCALE_SENTINEL, + ) + elif role is ACT: + self.register_params( + layer, + "input_scale", + (shapes.num_partitions,), + torch.float32, + PerTensorScaleParameter, + wl, + init=FP8_SCALE_SENTINEL, + ) + else: + self.reject(role) + + def process(self, layer, role) -> None: + if role is WEIGHT: + weight = layer.weight + max_w_scale = layer.weight_scale.max() + if not (layer.weight_scale == layer.weight_scale[0]).all(): + max_w_scale, weight = requantize_with_max_scale( + layer.weight, layer.weight_scale, layer.logical_widths + ) + # Transposed here rather than in the kernel: ModelOpt maps each + # fp8 key to exactly one kernel, so the layout is key-determined. + layer.weight = Parameter(weight.t(), requires_grad=False) + # The plain Parameter drops ModelWeightParameter's dim attributes; + # restore them for the transposed layout, which Humming reads. + layer.weight.input_dim = 0 + layer.weight.output_dim = 1 + layer.weight_scale = Parameter(max_w_scale, requires_grad=False) + elif role is ACT: + if torch.unique(layer.input_scale).numel() != 1: + logger.warning_once( + "In FP8 linear, the static input scale differs across " + "parallel layers (e.g. q_proj, k_proj, v_proj). Collapsing " + "them to the max will likely reduce accuracy. Consider a " + "checkpoint with the same input scale for fused weights." + ) + layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False) + else: + self.reject(role) + + +class KFp8StaticChannel(QuantKeyScheme): + """Per-channel static FP8 weight (the 'PcPt' weight). Weight-role only — + there is no static per-channel *activation* today.""" + + key = kFp8StaticTokenSym + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is not WEIGHT: + self.reject(role) + self.register_params( + layer, + "weight", + (shapes.output_size_per_partition, shapes.input_size_per_partition), + torch.float8_e4m3fn, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + self.register_params( + layer, + "weight_scale", + (shapes.output_size_per_partition,), + torch.float32, + ChannelQuantScaleParameter, + wl, + output_dim=0, + init=FP8_SCALE_SENTINEL, + ) + + def process(self, layer, role) -> None: + if role is not WEIGHT: + self.reject(role) + weight, weight_scale, _ = process_fp8_weight_channel_strategy( + layer.weight, layer.weight_scale.data + ) + layer.weight = Parameter(weight.t(), requires_grad=False) + layer.weight_scale = Parameter(weight_scale, requires_grad=False) + + +class KFp8Block128(QuantKeyScheme): + """128x128 block-static FP8 weight ('PbWo'). Weight-role only. ModelOpt + exports the scale 4-D [out_blk,1,in_blk,1]; process squeezes to 2-D. + No transpose (block kernel keeps [out,in]).""" + + key = kFp8Static128BlockSym + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is not WEIGHT: + self.reject(role) + out = shapes.output_size_per_partition + in_ = shapes.input_size_per_partition + # Scale sized by cdiv to match the checkpoint's block grid (the last + # output block may be partial). Weight is registered at its logical + # size; a non-128 output width is padded to a block boundary by the + # _Fp8PbWoPartialBlock format scheme. Input must be 128-aligned (same as + # #53132, which only pads the output). + if in_ % 128 != 0: + raise ValueError(f"FP8_PB_WO requires in divisible by 128, got {in_}") + self.register_params( + layer, + "weight", + (out, in_), + torch.float8_e4m3fn, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + ob, ib = cdiv(out, 128), in_ // 128 + self.register_params( + layer, + "weight_scale", + (ob, 1, ib, 1), + torch.float32, + BlockQuantScaleParameter, + wl, + input_dim=2, + output_dim=0, + init=FP8_SCALE_SENTINEL, + ) + layer.weight_block_size = [128, 128] + + def process(self, layer, role) -> None: + if role is not WEIGHT: + self.reject(role) + layer.weight = Parameter(layer.weight.data, requires_grad=False) + s = layer.weight_scale + if s.dim() == 4: + s = s.squeeze(1).squeeze(-1) # [ob,1,ib,1] -> [ob,ib] + elif s.dim() != 2: + raise ValueError( + f"Unexpected FP8_PB_WO weight_scale shape {tuple(s.shape)}" + ) + layer.weight_scale = Parameter(s.contiguous(), requires_grad=False) + + +class KMxfp8Static(QuantKeyScheme): + """MXFP8 weight: fp8-e4m3 values + per-32-block e8m0 (uint8) scale. + Weight-role only. process is validate-only plus an idempotency guard.""" + + key = kMxfp8Static + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is not WEIGHT: + self.reject(role) + if shapes.input_size_per_partition % MXFP8_BLOCK_SIZE != 0: + raise ValueError( + f"MXFP8 requires input size divisible by {MXFP8_BLOCK_SIZE}, " + f"got {shapes.input_size_per_partition}" + ) + self.register_params( + layer, + "weight", + (shapes.output_size_per_partition, shapes.input_size_per_partition), + MXFP8_VALUE_DTYPE, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + self.register_params( + layer, + "weight_scale", + ( + shapes.output_size_per_partition, + shapes.input_size_per_partition // MXFP8_BLOCK_SIZE, + ), + MXFP8_SCALE_DTYPE, + ModelWeightParameter, + wl, + input_dim=1, + output_dim=0, + ) + + def process(self, layer, role) -> None: + if role is not WEIGHT: + self.reject(role) + # Idempotency: the emulation kernel's post-load may replace the 1-byte + # MXFP8 weight with BF16 (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD, default + # on). On a weight reload process runs again after that swap, so skip + # the fp8 asserts when the weight is already >=2-byte. + if layer.weight.element_size() >= 2: + return + assert layer.weight.ndim == 2 and layer.weight.dtype == MXFP8_VALUE_DTYPE + assert layer.weight_scale.ndim == 2 + assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE + + +class KDynamicNoParam(QuantKeyScheme): + """Dynamic activation with no stored scale (W8A8): quantized at runtime in + the kernel. NOT the same as activation=None (weight-only) — init_fp8 needs a + non-None activation key. Activation-role only. Serves the fp8 per-token, fp8 + per-block, and mxfp8 dynamic activation keys.""" + + def create_weights(self, layer, role, ctx, shapes, wl) -> None: + if role is not ACT: + self.reject(role) + # dynamic -> nothing stored + + def process(self, layer, role) -> None: + if role is not ACT: + self.reject(role) + + +SCHEME_FOR: dict[QuantKey | None, QuantKeyScheme] = { + kNvfp4Static: KNvfp4Static(), + kNvfp4Dynamic: KNvfp4Dynamic(), + kFp8StaticTensorSym: KFp8StaticTensor(), + kFp8StaticTokenSym: KFp8StaticChannel(), + kFp8Static128BlockSym: KFp8Block128(), + kMxfp8Static: KMxfp8Static(), + kFp8DynamicTokenSym: KDynamicNoParam(), + kFp8Dynamic128Sym: KDynamicNoParam(), + kMxfp8Dynamic: KDynamicNoParam(), +} + + +def maybe_fuse_global_scales(layer) -> None: + """alpha = input_global_scale * weight_global_scale, presence-gated. + + W4A4 has both -> computed; W4A16 has no input_global_scale -> skipped. + """ + if hasattr(layer, "weight_global_scale") and hasattr(layer, "input_global_scale"): + layer.alpha = Parameter( + layer.input_global_scale * layer.weight_global_scale, + requires_grad=False, + ) + + +def select_linear_kernel(spec: QuantSpec, layer, rt: RuntimeDtypes): + """Thin family dispatcher on the weight key: nvfp4 / mxfp8 / fp8.""" + w = spec.weight + assert isinstance(w, QuantKey), f"resolve() must supply a weight key, got {w!r}" + if w.dtype == FP4_DTYPE: + # W4A16 (activation is None) → use_a16=True defaults to Marlin *and* + # honors --linear-backend (matches upstream ModelOptNvFp4W4A16LinearMethod + # after #50273); W4A4 → use_a16=False. + return init_nvfp4_linear_kernel(use_a16=spec.activation is None) + if w.scale.dtype == MXFP8_SCALE_DTYPE: + return init_mxfp8_linear_kernel() + # fp8 family: init_fp8 routes block-vs-plain itself off the activation key, + # and needs a real key -- weight-only fp8 is not a ModelOpt format. + act = spec.activation + assert isinstance(act, QuantKey), f"fp8 needs an activation key, got {act!r}" + return init_fp8_linear_kernel( + activation_quant_key=act, + weight_quant_key=w, + input_dtype=rt.input_dtype, + out_dtype=rt.out_dtype, + weight_shape=layer.weight.shape, + module_name=type(layer).__name__, + ) + + +class FormatScheme: + """Optional per-format hooks that compose *around* the QuantKey schemes. + + Extension seam for residue that belongs to a format as a whole rather than a + single QuantKey — an extra parameter the weight/activation schemes don't + cover, or a tweak before/after their ``process``. All hooks default to + no-ops; most formats need none (they are a pure ``(weight, activation)`` key + pair). To add one: subclass, override the hooks you need, and return an + instance from the format's ``resolve()`` branch — no change to + ``ModelOptLinearMethod`` itself. + """ + + def extra_weights(self, layer, shapes: "Shapes", ctx: CkptCtx, wl) -> None: + """Register format-level params (after the key schemes' weights).""" + + def pre_process(self, layer) -> None: + """Run before the key schemes' ``process``.""" + + def post_process(self, layer) -> None: + """Run after the key schemes' ``process``, before the kernel's.""" + + def apply(self, layer, x, bias, kernel_apply): + """Wrap the kernel's forward. ``kernel_apply(layer, x, bias) -> Tensor``. + + Default: delegate unchanged. A format whose residue lives at compute + time (e.g. run the GEMM on a padded weight then slice the output back + to its logical width, #53132-style) overrides this and calls + ``kernel_apply`` in the middle. + """ + return kernel_apply(layer, x, bias) + + +class _Fp8PbWoPartialBlock(FormatScheme): + """FP8_PB_WO output width that is not a multiple of 128 (a partial trailing + block; e.g. GLM's replicated ``fused_qkv_a_proj`` = 2048 + 576 = 2624). + + The block kernel needs full 128-blocks, so pad the weight up to the block + boundary with zeros before the kernel's post-load, then slice the GEMM + output back to the logical width. This is #53132's approach, expressed as + the format's compute-time residue on the generic method. The cdiv-sized + scale from ``KFp8Block128`` already matches the padded block count, so only + the weight rows and the output need adjusting. + """ + + def post_process(self, layer) -> None: + w = layer.weight + out = w.shape[0] + pad = (-out) % 128 + if not pad: + return + layer._pbwo_logical_out = out + padded = w.data.new_zeros(out + pad, w.shape[1]) + padded[:out].copy_(w.data) + layer.weight = Parameter(padded, requires_grad=False) + + def apply(self, layer, x, bias, kernel_apply): + logical = getattr(layer, "_pbwo_logical_out", None) + if logical is None: # width was block-aligned: nothing to trim + return kernel_apply(layer, x, bias) + # Run the GEMM on the padded weight without bias, then trim to the + # logical width and add bias (the padded rows produce zeros we drop). + out = kernel_apply(layer, x, None) + out = out[..., :logical].contiguous() + return out if bias is None else out.add_(bias) + + +_PB_WO_PARTIAL_BLOCK = _Fp8PbWoPartialBlock() + + +class _DropInputScale(FormatScheme): + """Interim: register then drop a W4A16 checkpoint's on-disk input_scale. + + Kept for backward compat with previously-exported ModelOpt checkpoints. + """ + + def extra_weights(self, layer, shapes, ctx, wl) -> None: + data = torch.full((shapes.num_partitions,), torch.nan) + layer.register_parameter( + "input_scale", PerTensorScaleParameter(data=data, weight_loader=wl) + ) + + def post_process(self, layer) -> None: + scale = getattr(layer, "input_scale", None) + if scale is None: + return + if not torch.isnan(scale).all(): + logger.warning_once( + "Dropping input_scale from a weight-only W4A16 NVFP4 " + "checkpoint (deprecated)." + ) + del layer.input_scale + + +_DROP_INPUT_SCALE = _DropInputScale() + + +@register_weight_loader_v2_supported_method +class ModelOptLinearMethod(LinearMethodBase): + """Generic, format-agnostic ModelOpt linear method. Holds a weight scheme + + an activation scheme (from the QuantSpec pair) and an optional + ``FormatScheme``, runs a fixed lifecycle, selects the kernel from the pair, + and applies. + + Registered for weight_loader_v2 (like every ModelOpt linear method) so the + ``BasevLLMParameter`` params route through the v2 fused loader, not the + legacy shape-assert path. + """ + + def __init__( + self, + spec: QuantSpec, + ctx: CkptCtx, + format_scheme=None, + ) -> None: + self.spec = spec + self.ctx = ctx + self.fmt = format_scheme or FormatScheme() + self.wkey = SCHEME_FOR[spec.weight] + self.akey = None if spec.activation is None else SCHEME_FOR[spec.activation] + # Only the fp8/mxfp8 kernels read input_dtype; nvfp4 ignores it. During + # real serving model_config is always set; fall back defensively when + # there is no config context (bare unit-test dispatch). + model_config = getattr(get_current_vllm_config_or_none(), "model_config", None) + dtype = ( + model_config.dtype + if model_config is not None + else torch.get_default_dtype() + ) + self.out_dtype = dtype + self.input_dtype = dtype + self.marlin_input_dtype = None + # Kernel/backend are chosen in create_weights (after get_quant_method), + # so the front-end marlin poke stays dormant here — same as the old + # NVFP4 methods. + self.kernel: Any = None + + def create_weights( + self, + layer, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + del input_size, output_size + weight_loader = extra_weight_attrs.get("weight_loader") + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = sum(output_partition_sizes) + # Humming reads both off the layer in + # prepare_humming_linear_layer_config. LinearBase sets them itself; + # ParallelLMHead does not, so supply them here. + layer.output_partition_sizes = output_partition_sizes + if not hasattr(layer, "has_bias"): + layer.has_bias = getattr(layer, "bias", None) is not None + shapes = Shapes(output_partition_sizes, input_size_per_partition, params_dtype) + + self.wkey.create_weights(layer, WEIGHT, self.ctx, shapes, weight_loader) + if self.akey: + self.akey.create_weights(layer, ACT, self.ctx, shapes, weight_loader) + self.fmt.extra_weights(layer, shapes, self.ctx, weight_loader) + + rt = RuntimeDtypes(self.input_dtype, self.out_dtype, self.marlin_input_dtype) + self.kernel = select_linear_kernel(self.spec, layer, rt) + expose_input_quant_key(layer, self.kernel) + + def process_weights_after_loading(self, layer) -> None: + self.fmt.pre_process(layer) + self.wkey.process(layer, WEIGHT) + if self.akey: + self.akey.process(layer, ACT) + maybe_fuse_global_scales(layer) + self.fmt.post_process(layer) + self.kernel.process_weights_after_loading(layer) + + def apply(self, layer, x, bias=None): + return self.fmt.apply( + layer, + x, + bias, + lambda lyr, inp, b: self.kernel.apply_weights(layer=lyr, x=inp, bias=b), + ) + + +def resolve(algo: str, subcfg, prefix: str): + """Turn a sub-config into (QuantSpec, CkptCtx, format_scheme). + + Strictly read-only over ``subcfg``: in mixed precision the same sub-config + is shared with the MoE method, so writing back would leak a linear-only + change into MoE. + """ + if algo == "FP8": + # Plain per-tensor static FP8 (W8A8): same key in both slots (bivalent). + ctx = CkptCtx() + return ( + QuantSpec(weight=kFp8StaticTensorSym, activation=kFp8StaticTensorSym), + ctx, + None, + ) + if algo == "FP8_PER_CHANNEL_PER_TOKEN": + # PcPt: per-channel static weight, dynamic per-token activation (W8A8). + ctx = CkptCtx() + return ( + QuantSpec(weight=kFp8StaticTokenSym, activation=kFp8DynamicTokenSym), + ctx, + None, + ) + if algo == "FP8_PB_WO": + # PbWo: 128x128 block-static weight, dynamic per-block activation (W8A8). + # The block kernel's post-load runs here; CompressedTensors block-FP8 + # is the reference for this path. The FormatScheme pads a non-128 output + # width to a block boundary and trims the output back (#53132) -- a no-op + # for the common block-aligned case. + ctx = CkptCtx() + return ( + QuantSpec(weight=kFp8Static128BlockSym, activation=kFp8Dynamic128Sym), + ctx, + _PB_WO_PARTIAL_BLOCK, + ) + if algo == "MXFP8": + # MXFP8: block(32) e4m3 weight + e8m0 scale, dynamic activation. + ctx = CkptCtx() + return QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic), ctx, None + + # NVFP4 family (W4A4 / W4A16). + ctx = CkptCtx(group_size=subcfg.group_size) + if algo == "NVFP4": + # W4A4: static fp4 weight + dynamic fp4 activation (has a static global + # input scale). alpha = weight_gs * input_gs. + return QuantSpec(weight=kNvfp4Static, activation=kNvfp4Dynamic), ctx, None + if algo == "W4A16_NVFP4": + # W4A16: fp4 weight, no activation quant (-> Marlin). _DROP_INPUT_SCALE + # is a no-op for genuine exports; keeps legacy input_scale ckpts loading. + return QuantSpec(weight=kNvfp4Static, activation=None), ctx, _DROP_INPUT_SCALE + raise NotImplementedError(f"resolve: unsupported ModelOpt linear algo {algo!r}") + + +# Bespoke-method escape hatch. Almost every ModelOpt linear format is a +# (weight, activation) QuantKey pair handled by the generic ModelOptLinearMethod +# (optionally + a FormatScheme), so this stays empty. A format that genuinely +# cannot be expressed that way — a fundamentally different create/process/apply +# lifecycle — registers its own LinearMethodBase here, keyed by algo: +# +# LINEAR_METHOD_BUILDERS["FOO"] = lambda cfg, prefix: ModelOptFooLinearMethod(cfg) +# +# Keep the ModelOpt* class-name prefix and decorate the class with +# @register_weight_loader_v2_supported_method if it uses BasevLLMParameter +# params. Also add the algo to the owning config's quant_method validation. +LINEAR_METHOD_BUILDERS: dict[str, Callable[..., LinearMethodBase]] = {} + + +def build_linear_method(config, algo: str, prefix: str) -> LinearMethodBase: + """Construct the linear method for ``algo``. + + Returns a bespoke method if one is registered in ``LINEAR_METHOD_BUILDERS``, + else the generic ``ModelOptLinearMethod`` built from ``resolve(algo, config, + prefix)``. Single indirection point shared by the homogeneous and + mixed-precision dispatch, so a new format plugs in without editing either + ``get_quant_method``. + """ + builder = LINEAR_METHOD_BUILDERS.get(algo) + if builder is not None: + return builder(config, prefix) + spec, ctx, format_scheme = resolve(algo, config, prefix) + return ModelOptLinearMethod(spec, ctx, format_scheme) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 5df6bcb734..0d49e3510f 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -38,6 +38,11 @@ from vllm.utils.platform_utils import get_device_name_as_file_name logger = init_logger(__name__) +# Pre-fill value for scale parameters whose shards load independently. The +# shards are combined with .max(), so an unloaded shard must never win; the +# smallest representable float32 guarantees that. +FP8_SCALE_SENTINEL = torch.finfo(torch.float32).min + def is_fp8(x: torch.dtype | torch.Tensor) -> bool: if isinstance(x, torch.Tensor):