@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / debd6bbf09
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 19:45 UTC

[Kernel] Add ModelOpt FP4 Checkpoint Support (#12520)

Signed-off-by: Pavani Majety <pmajety@nvidia.com>

Pavani Majety committed Mar 11, 2025 at 22:13 UTC debd6bbf0951222c9ad7a1a91b958027f2ed0782
10 files changed +388 -30
csrc/ops.h
+5 -3
@@ -160,14 +160,16 @@ torch::Tensor ggml_moe_a8(torch::Tensor X, torch::Tensor W,
160 int64_t ggml_moe_get_block_size(int64_t type);
161
162 #ifndef USE_ROCM
163 +
164 +bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
165 +bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
166 +bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
167 +
168 void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
169 torch::Tensor const& B, torch::Tensor const& A_sf,
170 torch::Tensor const& B_sf,
171 torch::Tensor const& alpha);
172
168 -bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
169 -bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
170 -
173 void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a,
174 torch::Tensor const& b, torch::Tensor const& a_scales,
175 torch::Tensor const& b_scales,
csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu
+6
@@ -36,3 +36,9 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
36 "be compiled using CUDA 12.8 and target "
37 "compute capability 100 or above.");
38 }
39 +
40 +bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability) {
41 + int runtimeVersion;
42 + cudaRuntimeGetVersion(&runtimeVersion);
43 + return cuda_device_capability >= 100 && runtimeVersion >= 12080;
44 +}
\ No newline at end of file
csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu
+4 -3
@@ -201,10 +201,11 @@ void runGemm(at::Tensor& D, at::Tensor const& A, at::Tensor const& B,
201 #endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
202
203 #define CHECK_TYPE(x, st, m) \
204 - TORCH_CHECK(x.scalar_type() == st, "Inconsistency of Tensor type:", m)
205 -#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
204 + TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
205 +#define CHECK_TH_CUDA(x, m) \
206 + TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor")
207 #define CHECK_CONTIGUOUS(x, m) \
207 - TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
208 + TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous")
209 #define CHECK_INPUT(x, st, m) \
210 CHECK_TH_CUDA(x, m); \
211 CHECK_CONTIGUOUS(x, m); \
csrc/torch_bindings.cpp
+4
@@ -434,6 +434,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
434 " Tensor! output_scale, Tensor input_scale) -> ()");
435 ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
436
437 + // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
438 + // of the given capability
439 + ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
440 + ops.impl("cutlass_scaled_mm_supports_fp4", &cutlass_scaled_mm_supports_fp4);
441 #endif
442
443 // Quantized GEMM for GPTQ.
tests/models/decoder_only/language/test_nvfp4.py new
+82
@@ -0,0 +1,82 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# flake8: noqa
3 +"""Tests Model Optimizer nvfp4 models against ground truth generation
4 +Note: these tests will only pass on B200
5 +"""
6 +import os
7 +from typing import List
8 +
9 +import pytest
10 +from transformers import AutoTokenizer
11 +
12 +from tests.quantization.utils import is_quant_method_supported
13 +from vllm import LLM, SamplingParams
14 +
15 +os.environ["TOKENIZERS_PARALLELISM"] = "true"
16 +
17 +MAX_MODEL_LEN = 1024
18 +
19 +MODELS = ["nvidia/Llama-3.3-70B-Instruct-FP4"]
20 +
21 +EXPECTED_STRS_MAP = {
22 + "nvidia/Llama-3.3-70B-Instruct-FP4": [
23 + 'vLLM (Vectorized Large Language Model) is indeed a high-throughput and memory-efficient inference',
24 + 'Here are the major milestones in the development of artificial intelligence (AI) from 1950 to ',
25 + 'Artificial intelligence (AI) and human intelligence (HI) are two distinct forms of intelligence that process',
26 + 'A neural network is a type of machine learning model inspired by the structure and function of the human brain',
27 + 'In the heart of a cutting-edge robotics lab, a team of engineers had been working tirelessly to push',
28 + 'The COVID-19 pandemic has had a profound impact on global economic structures and future business models, leading',
29 + 'The Mona Lisa, painted by Leonardo da Vinci in the early 16th century, is one of',
30 + 'Here are the translations:\n\n* Japanese: (Sasuga no tori ga miwa o ts'
31 + ]
32 +}
33 +
34 +
35 +# This test compares against golden strings for exact match since
36 +# there is no baseline implementation to compare against
37 +# and is unstable w.r.t specifics of the fp4 implementation or
38 +# the hardware being run on.
39 +# Disabled to prevent it from breaking the build
40 +@pytest.mark.skip(
41 + reason=
42 + "Prevent unstable test based on golden strings from breaking the build "
43 + " and test input model being too large and hanging the system.")
44 +@pytest.mark.quant_model
45 +@pytest.mark.skipif(not is_quant_method_supported("nvfp4"),
46 + reason="nvfp4 is not supported on this GPU type.")
47 +@pytest.mark.parametrize("model_name", MODELS)
48 +def test_models(example_prompts, model_name) -> None:
49 + model = LLM(
50 + model=model_name,
51 + max_model_len=MAX_MODEL_LEN,
52 + trust_remote_code=True,
53 + enforce_eager=True,
54 + quantization="nvfp4",
55 + )
56 +
57 + tokenizer = AutoTokenizer.from_pretrained(model_name)
58 + formatted_prompts = [
59 + tokenizer.apply_chat_template([{
60 + "role": "user",
61 + "content": prompt
62 + }],
63 + tokenize=False,
64 + add_generation_prompt=True)
65 + for prompt in example_prompts
66 + ]
67 + params = SamplingParams(max_tokens=20, temperature=0)
68 + generations: List[str] = []
69 + # Note: these need to be run 1 at a time due to numerical precision,
70 + # since the expected strs were generated this way.
71 + for prompt in formatted_prompts:
72 + outputs = model.generate(prompt, params)
73 + generations.append(outputs[0].outputs[0].text)
74 + del model
75 +
76 + print(model_name, generations)
77 + expected_strs = EXPECTED_STRS_MAP[model_name]
78 + for i in range(len(example_prompts)):
79 + generated_str = generations[i]
80 + expected_str = expected_strs[i]
81 + assert expected_str == generated_str, (
82 + f"Test{i}:\nExpected: {expected_str!r}\nvLLM: {generated_str!r}")
vllm/_custom_ops.py
+4
@@ -467,6 +467,10 @@ if hasattr(torch.ops._C, "ggml_dequantize"):
467
468
469 # cutlass
470 +def cutlass_scaled_mm_supports_fp4(cuda_device_capability: int) -> bool:
471 + return torch.ops._C.cutlass_scaled_mm_supports_fp4(cuda_device_capability)
472 +
473 +
474 def cutlass_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,
475 block_scale_a: torch.Tensor,
476 block_scale_b: torch.Tensor, alpha: torch.Tensor,
vllm/config.py
+1 -1
@@ -613,7 +613,7 @@ class ModelConfig:
613 optimized_quantization_methods = [
614 "fp8", "marlin", "modelopt", "gptq_marlin_24", "gptq_marlin",
615 "awq_marlin", "fbgemm_fp8", "compressed_tensors",
616 - "compressed-tensors", "experts_int8", "quark"
616 + "compressed-tensors", "experts_int8", "quark", "nvfp4"
617 ]
618 if self.quantization is not None:
619 self.quantization = self.quantization.lower()
vllm/model_executor/layers/linear.py
+17 -6
@@ -30,12 +30,23 @@ from vllm.model_executor.utils import set_weight_attrs
30 logger = init_logger(__name__)
31
32 WEIGHT_LOADER_V2_SUPPORTED = [
33 - "CompressedTensorsLinearMethod", "AWQMarlinLinearMethod",
34 - "AWQLinearMethod", "GPTQMarlinLinearMethod", "Fp8LinearMethod",
35 - "MarlinLinearMethod", "QQQLinearMethod", "GPTQMarlin24LinearMethod",
36 - "TPUInt8LinearMethod", "GPTQLinearMethod", "FBGEMMFp8LinearMethod",
37 - "ModelOptFp8LinearMethod", "IPEXAWQLinearMethod", "IPEXGPTQLinearMethod",
38 - "HQQMarlinMethod", "QuarkLinearMethod"
33 + "CompressedTensorsLinearMethod",
34 + "AWQMarlinLinearMethod",
35 + "AWQLinearMethod",
36 + "GPTQMarlinLinearMethod",
37 + "Fp8LinearMethod",
38 + "MarlinLinearMethod",
39 + "QQQLinearMethod",
40 + "GPTQMarlin24LinearMethod",
41 + "TPUInt8LinearMethod",
42 + "GPTQLinearMethod",
43 + "FBGEMMFp8LinearMethod",
44 + "ModelOptFp8LinearMethod",
45 + "IPEXAWQLinearMethod",
46 + "IPEXGPTQLinearMethod",
47 + "HQQMarlinMethod",
48 + "QuarkLinearMethod",
49 + "ModelOptNvFp4LinearMethod",
50 ]
51
52
vllm/model_executor/layers/quantization/__init__.py
+3 -1
@@ -14,6 +14,7 @@ QUANTIZATION_METHODS: List[str] = [
14 "ptpc_fp8",
15 "fbgemm_fp8",
16 "modelopt",
17 + "nvfp4",
18 # The order of gptq methods is important for config.py iteration over
19 # override_quantization_method(..)
20 "marlin",
@@ -97,7 +98,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
98 from .hqq_marlin import HQQMarlinConfig
99 from .ipex_quant import IPEXConfig
100 from .marlin import MarlinConfig
100 - from .modelopt import ModelOptFp8Config
101 + from .modelopt import ModelOptFp8Config, ModelOptNvFp4Config
102 from .moe_wna16 import MoeWNA16Config
103 from .neuron_quant import NeuronQuantConfig
104 from .ptpc_fp8 import PTPCFp8Config
@@ -112,6 +113,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
113 "fp8": Fp8Config,
114 "fbgemm_fp8": FBGEMMFp8Config,
115 "modelopt": ModelOptFp8Config,
116 + "nvfp4": ModelOptNvFp4Config,
117 # The order of gptq methods is important for config.py iteration over
118 # override_quantization_method(..)
119 "marlin": MarlinConfig,
vllm/model_executor/layers/quantization/modelopt.py
+262 -16
@@ -1,24 +1,31 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from typing import Any, Dict, List, Optional
3 +from typing import Any, Dict, List, Optional, Union
4
5 import torch
6 from torch.nn import Module
7 from torch.nn.parameter import Parameter
8
9 +from vllm._custom_ops import (cutlass_scaled_fp4_mm,
10 + cutlass_scaled_mm_supports_fp4, scaled_fp4_quant)
11 from vllm.logger import init_logger
10 -from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
12 +from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
13 + UnquantizedLinearMethod)
14 from vllm.model_executor.layers.quantization.base_config import (
15 QuantizationConfig, QuantizeMethodBase)
16 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
17 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 + is_layer_skipped)
19 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
20 Fp8LinearOp, requantize_with_max_scale)
21 from vllm.model_executor.parameter import (ModelWeightParameter,
22 PerTensorScaleParameter)
23 +from vllm.platforms import current_platform
24
25 logger = init_logger(__name__)
26
21 -ACTIVATION_SCHEMES = ["static"]
27 +QUANT_ALGOS = ["FP8", "NVFP4"]
28 +KV_CACHE_QUANT_ALGOS = ["FP8"]
29
30
31 class ModelOptFp8Config(QuantizationConfig):
@@ -54,12 +61,13 @@ class ModelOptFp8Config(QuantizationConfig):
61 def from_config(cls, config: Dict[str, Any]) -> "ModelOptFp8Config":
62 quant_config = cls.get_from_keys(config, ["quantization"])
63 quant_method = quant_config["quant_algo"]
57 - is_checkpoint_fp8_serialized = ("FP8" in quant_method)
58 - if not is_checkpoint_fp8_serialized:
59 - raise ValueError("ModelOpt currently only supports static FP8 "
60 - "quantization in vLLM. Please check the "
64 + if quant_method not in QUANT_ALGOS:
65 + raise ValueError(f"ModelOpt currently only supports: {QUANT_ALGOS}"
66 + " quantizations in vLLM. Please check the "
67 "`hf_quant_config.json` file for your model's "
68 "quant configuration.")
69 + is_checkpoint_fp8_serialized = ("FP8" in quant_method)
70 +
71 return cls(is_checkpoint_fp8_serialized)
72
73 def get_quant_method(self, layer: torch.nn.Module,
@@ -72,15 +80,6 @@ class ModelOptFp8Config(QuantizationConfig):
80 return None
81
82
75 -class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
76 - """
77 - Supports loading kv-cache scaling factors from FP8 checkpoints.
78 - """
79 -
80 - def __init__(self, quant_config: ModelOptFp8Config):
81 - super().__init__(quant_config)
82 -
83 -
83 class ModelOptFp8LinearMethod(LinearMethodBase):
84 """Linear method for Model Optimizer static quantization.
85 Supports loading FP8 checkpoints with static weight scale and
@@ -162,3 +161,250 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
161 weight_scale=layer.weight_scale,
162 input_scale=layer.input_scale,
163 bias=bias)
164 +
165 +
166 +class ModelOptNvFp4Config(QuantizationConfig):
167 + """Config class for ModelOpt FP4."""
168 +
169 + def __init__(
170 + self,
171 + is_checkpoint_nvfp4_serialized: bool,
172 + kv_cache_quant_algo: str,
173 + exclude_modules: List[str],
174 + group_size: int = 16,
175 + ) -> None:
176 + self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
177 + if is_checkpoint_nvfp4_serialized:
178 + logger.warning(
179 + "Detected ModelOpt NVFP4 checkpoint. Please note that"
180 + " the format is experimental and could change in future.")
181 +
182 + self.group_size = group_size
183 + self.kv_cache_quant_algo = kv_cache_quant_algo
184 + self.exclude_modules = exclude_modules
185 +
186 + @classmethod
187 + def get_name(cls) -> str:
188 + return "modelopt_nvfp4"
189 +
190 + @classmethod
191 + def get_supported_act_dtypes(cls) -> List[torch.dtype]:
192 + return [torch.bfloat16, torch.half, torch.float8_e4m3fn]
193 +
194 + @classmethod
195 + def get_min_capability(cls) -> int:
196 + return 100
197 +
198 + @classmethod
199 + def get_config_filenames(cls) -> List[str]:
200 + return ["hf_quant_config.json"]
201 +
202 + @classmethod
203 + def from_config(cls, config: Dict[str, Any]) -> "ModelOptNvFp4Config":
204 + quant_config = cls.get_from_keys(config, ["quantization"])
205 + quant_method = quant_config["quant_algo"]
206 + if quant_method not in QUANT_ALGOS:
207 + raise ValueError(f"ModelOpt currently only supports: {QUANT_ALGOS}"
208 + " quantizations in vLLM. Please check the "
209 + "`hf_quant_config.json` file for your model's "
210 + "quant configuration.")
211 + is_checkpoint_nvfp4_serialized = ("NVFP4" in quant_method)
212 + kv_cache_quant_algo = quant_config["kv_cache_quant_algo"]
213 + group_size = quant_config["group_size"]
214 + exclude_modules = quant_config["exclude_modules"]
215 + if not (group_size and kv_cache_quant_algo and exclude_modules):
216 + raise ValueError("NVFP4 quantization requires group size and "
217 + "kv_cache_quant_algo specified in "
218 + "hf_quant_config.json")
219 + return cls(is_checkpoint_nvfp4_serialized, kv_cache_quant_algo,
220 + exclude_modules, group_size)
221 +
222 + def get_quant_method(self, layer: torch.nn.Module,
223 + prefix: str) -> Optional["QuantizeMethodBase"]:
224 + from vllm.attention.layer import Attention # Avoid circular import
225 + if isinstance(layer, LinearBase):
226 + if is_layer_skipped(prefix, self.exclude_modules):
227 + return UnquantizedLinearMethod()
228 + return ModelOptNvFp4LinearMethod(self)
229 + elif isinstance(layer, Attention):
230 + return ModelOptFp8KVCacheMethod(self)
231 + return None
232 +
233 +
234 +def cutlass_fp4_supported() -> bool:
235 + if not current_platform.is_cuda():
236 + return False
237 + capability_tuple = current_platform.get_device_capability()
238 + capability = -1 if capability_tuple is None else capability_tuple.to_int()
239 + return cutlass_scaled_mm_supports_fp4(capability)
240 +
241 +
242 +class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):
243 + """
244 + Supports loading kv-cache scaling factors from FP8 checkpoints.
245 + """
246 +
247 + def __init__(self, quant_config: Union[ModelOptFp8Config,
248 + ModelOptNvFp4Config]):
249 + super().__init__(quant_config)
250 +
251 +
252 +class ModelOptNvFp4LinearMethod(LinearMethodBase):
253 + """Linear method for Model Optimizer NVFP4.
254 + Supports loading NVFP4 checkpoints with the following structure:
255 +
256 + input_scale: torch.float32, scalar ,
257 + weight: NVFP4(represented as byte) Shape: [1, X, y/2]
258 + weight_scale: FP8-E4M3, Shape: [X, Y], aka per block scale,
259 + weight_scale_2: torch.float32, scalar,
260 + Args: quant_config: The ModelOpt quantization config.
261 + """
262 +
263 + def __init__(self, quant_config: ModelOptNvFp4Config):
264 + self.quant_config = quant_config
265 + self.cutlass_nvfp4_supported = cutlass_fp4_supported()
266 + if not self.cutlass_nvfp4_supported:
267 + raise ValueError("Current platform does not support NVFP4"
268 + " quantization. Please use Blackwell and above.")
269 +
270 + def create_weights(
271 + self,
272 + layer: torch.nn.Module,
273 + input_size_per_partition: int,
274 + output_partition_sizes: List[int],
275 + input_size: int,
276 + output_size: int,
277 + params_dtype: torch.dtype,
278 + **extra_weight_attrs,
279 + ):
280 + del input_size, output_size
281 + if not self.quant_config.is_checkpoint_nvfp4_serialized:
282 + raise ValueError("NVFP4 quantization was selected, "
283 + " dynamic quantization is not supported.")
284 + output_size_per_partition = sum(output_partition_sizes)
285 + weight_loader = extra_weight_attrs.get("weight_loader")
286 + layer.logical_widths = output_partition_sizes
287 + layer.input_size_per_partition = input_size_per_partition
288 + layer.output_size_per_partition = output_size_per_partition
289 +
290 + if (input_size_per_partition % 16 != 0):
291 + raise ValueError("Unsupported model when in features size is "
292 + "not multiple of 16")
293 + # The nvfp4 weight is still represented as
294 + weight_dtype = (torch.float8_e4m3fn
295 + if self.quant_config.is_checkpoint_nvfp4_serialized
296 + else params_dtype)
297 + # Weight
298 + weight = ModelWeightParameter(
299 + data=torch.empty(
300 + # 2 fp4 items are packed in the input dimension
301 + layer.output_size_per_partition,
302 + layer.input_size_per_partition // 2,
303 + dtype=torch.uint8),
304 + input_dim=1,
305 + output_dim=0,
306 + weight_loader=weight_loader)
307 + layer.register_parameter("weight", weight)
308 +
309 + # Input Weight Scale
310 + input_scale = PerTensorScaleParameter(data=torch.empty(
311 + len(output_partition_sizes), dtype=torch.float32),
312 + weight_loader=weight_loader)
313 + layer.register_parameter("input_scale", input_scale)
314 +
315 + # Global Weight Scale
316 + weight_scale_2 = PerTensorScaleParameter(data=torch.empty(
317 + len(output_partition_sizes), dtype=torch.float32),
318 + weight_loader=weight_loader)
319 + layer.register_parameter("weight_scale_2", weight_scale_2)
320 +
321 + # Per Block Weight Scale
322 + weight_scale = ModelWeightParameter(data=torch.empty(
323 + output_size_per_partition,
324 + input_size_per_partition // self.quant_config.group_size,
325 + dtype=weight_dtype,
326 + ),
327 + input_dim=1,
328 + output_dim=0,
329 + weight_loader=weight_loader)
330 +
331 + layer.register_parameter("weight_scale", weight_scale)
332 +
333 + def swizzle_blockscale(self, scale: torch.tensor):
334 + assert (scale.dtype == torch.float8_e4m3fn)
335 + # Pad and blockwise interleave weight_scale
336 + scale_ndim = scale.ndim
337 + if scale.ndim == 2:
338 + scale = scale.unsqueeze(0)
339 + assert scale.ndim == 3
340 + B, M, K = scale.shape
341 + round_up_multiple = lambda x, m: (x + m - 1) // m * m
342 + M_padded = round_up_multiple(M, 128)
343 + K_padded = round_up_multiple(K, 4)
344 + padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
345 + padded_scale[:B, :M, :K] = scale
346 + batches, rows, cols = padded_scale.shape
347 + assert rows % 128 == 0
348 + assert cols % 4 == 0
349 + padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
350 + cols // 4, 4)
351 + swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
352 + swizzled_scale = swizzled_scale.contiguous().cuda()
353 + return (swizzled_scale.reshape(M, K)
354 + if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
355 +
356 + def process_weights_after_loading(self, layer: Module) -> None:
357 +
358 + # global scales:
359 + input_scale_2 = layer.input_scale.max().to(torch.float32)
360 + layer.input_scale = Parameter(input_scale_2, requires_grad=False)
361 +
362 + weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
363 + layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False)
364 +
365 + layer.alpha = Parameter(layer.input_scale * layer.weight_scale_2,
366 + requires_grad=False)
367 +
368 + # Swizzle the weight blockscale.
369 + # contracting dimension is input dimension
370 + # block_size = 16;
371 + assert (layer.weight_scale.shape[1] % 16 == 0), (
372 + "Expected weight_scale.dim(1) to be divisible by 16")
373 + assert (layer.weight_scale.dtype == torch.float8_e4m3fn), (
374 + "Weight Block scale must be represented as FP8-E4M3")
375 + swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)
376 +
377 + layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
378 + requires_grad=False)
379 +
380 + def apply(
381 + self,
382 + layer: torch.nn.Module,
383 + x: torch.Tensor,
384 + bias: Optional[torch.Tensor] = None,
385 + ) -> torch.Tensor:
386 + output_dtype = x.dtype
387 +
388 + # for input only the contracting dimension has a constraint.
389 + x_m, _ = x.shape
390 + w_n, _ = layer.weight.shape
391 + output_shape = [x_m, w_n]
392 +
393 + # quantize BF16 or FP16 to (FP4 and interleaved block scale)
394 + s_quant = 1 / layer.input_scale
395 + x_fp4, x_blockscale = scaled_fp4_quant(x, s_quant)
396 +
397 + # validate dtypes of quantized input, input block scale,
398 + # weight and weight_blockscale
399 + assert (x_fp4.dtype == torch.uint8)
400 + assert (layer.weight.dtype == torch.uint8)
401 + assert (x_blockscale.dtype == torch.float8_e4m3fn)
402 + assert (layer.weight_scale_swizzled.dtype == torch.float8_e4m3fn)
403 + assert (layer.alpha.dtype == torch.float32)
404 +
405 + out = cutlass_scaled_fp4_mm(x_fp4, layer.weight, x_blockscale,
406 + layer.weight_scale_swizzled, layer.alpha,
407 + output_dtype)
408 + if bias is not None:
409 + out = out + bias
410 + return out.view(*output_shape)