From f46445fa1502328fcae8dc326cdaf732b10a71be Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Mon, 20 Apr 2026 20:52:28 +0000 Subject: [PATCH 01/16] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 21 ++++ torchao/prototype/gptq/api.py | 157 ++++++++++++++++++++++++- torchao/prototype/gptq/gptq_example.py | 16 ++- 3 files changed, 188 insertions(+), 6 deletions(-) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index 366fe6d242..0bae6ec4da 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -15,6 +15,9 @@ gptq_quantize, ) from torchao.prototype.gptq.observer import GPTQObserverTensor +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerRow from torchao.utils import _is_mslk_available @@ -435,6 +438,13 @@ def test_gptq_quantize_function(self, base_config): pytest.param( Int8WeightOnlyConfig(granularity=PerRow(), version=2), id="int8" ), + pytest.param( + NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ), + id="nvfp4", + ), ], ) def test_gptq_quantize_better_than_naive(self, base_config): @@ -506,6 +516,13 @@ def test_gptq_quantize_better_than_naive(self, base_config): pytest.param( Int8WeightOnlyConfig(granularity=PerRow(), version=2), id="int8" ), + pytest.param( + NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ), + id="nvfp4", + ), ], ) def test_gptq_sqnr(self, base_config): @@ -556,6 +573,10 @@ def test_gptq_sqnr(self, base_config): assert sqnr_gptq > 25, f"GPTQ SQNR: {sqnr_gptq} is too low" elif isinstance(base_config, Int8WeightOnlyConfig): assert sqnr_gptq > 30, f"GPTQ SQNR: {sqnr_gptq} is too low" + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + assert sqnr_gptq > 15, f"GPTQ SQNR: {sqnr_gptq} is too low" + else: + raise AssertionError("unsupported") assert sqnr_gptq > sqnr_rtn, ( f"GPTQ SQNR: {sqnr_gptq} is not better than RTN SQNR: {sqnr_rtn}" ) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index 0059bd1d51..ac6f7cdc6e 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -19,6 +19,25 @@ pack_int4 = None from torchao.core.config import AOBaseConfig +from torchao.prototype.mx_formats.constants import F4_E2M1_MAX +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) +from torchao.prototype.mx_formats.kernels import ( + f4_unpacked_to_f32, + f32_to_f4_unpacked, + pack_uint4, +) +from torchao.prototype.mx_formats.nvfp4_tensor import ( + NVFP4Tensor, + QuantizeTensorToNVFP4Kwargs, + nvfp4_quantize, + per_tensor_amax_to_scale, +) +from torchao.prototype.mx_formats.utils import ( + hp_data_dims_to_swizzled_scale_dims_nvfp4, + to_blocked, +) from torchao.quantization import Int4Tensor, Int8Tensor from torchao.quantization.granularity import PerRow from torchao.quantization.quant_api import ( @@ -33,6 +52,7 @@ CONFIG_TO_TORCHAO_BASE_TENSOR = { Int4WeightOnlyConfig: Int4Tensor, Int8WeightOnlyConfig: Int8Tensor, + NVFP4DynamicActivationNVFP4WeightConfig: NVFP4Tensor, } @@ -61,7 +81,11 @@ class GPTQConfig(AOBaseConfig): """ step: str = "observe" # "observe" or "convert" - base_config: Union[Int4WeightOnlyConfig, Int8WeightOnlyConfig] = None + base_config: Union[ + Int4WeightOnlyConfig, + Int8WeightOnlyConfig, + NVFP4DynamicActivationNVFP4WeightConfig, + ] = None percdamp: float = 0.01 gptq_quantize_block_size: int = 256 @@ -179,6 +203,59 @@ def _int4_row_dequantize_zp( return torch.cat(dequant_chunks, dim=-1) +def _nvfp4_with_precalculated_scales_qdq( + data_hp: torch.Tensor, + per_tensor_scale: torch.Tensor, + block_scale: torch.Tensor, +) -> torch.Tensor: + """ + Same as torchao.prototype.mx_formats.nvfp4_tensor.nvfp4_quantize, but with + per_tensor_scale and block_scale precalculated and the end result dequantized. + """ + assert per_tensor_scale.dtype is torch.float32 + assert block_scale.dtype is torch.float8_e4m3fn + # this function only works for data_hp.shape == (N, k_slice) + # and block_scale.shape == (N,) + assert len(block_scale.shape) == 1 + + scaled_block_scales_fp32 = block_scale.to(torch.float32) + reciprocal_scale = (1.0 / per_tensor_scale) / scaled_block_scales_fp32 + data_scaled = data_hp * reciprocal_scale.unsqueeze(-1) + data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX) + data_lp = f32_to_f4_unpacked(data_scaled) + data_lp_hp = f4_unpacked_to_f32(data_lp) + data_lp_hp_unscaled = data_lp_hp / reciprocal_scale.unsqueeze(-1) + return data_lp_hp_unscaled + + +def _nvfp4_with_precalculated_scales_q( + data_hp: torch.Tensor, + per_tensor_scale: torch.Tensor, + block_scale: torch.Tensor, +) -> torch.Tensor: + """ + Same as torchao.prototype.mx_formats.nvfp4_tensor.nvfp4_quantize, but with + per_tensor_scale and block_scale precalculated. + """ + assert per_tensor_scale.dtype is torch.float32 + assert block_scale.dtype is torch.float8_e4m3fn + + # TODO(future): figure out what to reuse vs leave copy-pasted vs + # nvfp4_tensor.py + scaled_block_scales_fp32 = block_scale.to(torch.float32) + reciprocal_scale = (1.0 / per_tensor_scale) / scaled_block_scales_fp32 + N, K = data_hp.shape + # reshape to 3d to properly broadcast for scaling + data_hp = data_hp.view(N, K // 16, 16) + data_scaled = data_hp * reciprocal_scale.unsqueeze(-1) + data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX) + data_lp = f32_to_f4_unpacked(data_scaled) + data_lp_packed = pack_uint4(data_lp) + # reshape back to 2d + data_lp_packed = data_lp_packed.view(N, K // 2) + return data_lp_packed + + def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): """ This function implements the GPTQ algorithm described in this paper: https://arxiv.org/abs/2210.17323 (Algorithm 1) @@ -241,6 +318,15 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): block_size = get_block_size(W_t.shape, base_config.granularity) block_size = list(block_size) group_size = block_size[-1] + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + group_size = 16 + block_size = [1, group_size] + # for per-tensor nvfp4, we need to calculate the global scale over the + # entire tensor before we enter the GPTQ loop + tensor_amax = torch.max(torch.abs(W_t)) + nvfp4_global_scale = per_tensor_amax_to_scale(tensor_amax) + else: + raise AssertionError("unsupported") assert group_size > 0 @@ -337,6 +423,27 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): ], base_config.granularity, ) + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + tensor_slice = B_cur[ + :, + G_k_start - B_cur_k_start : G_k_end - B_cur_k_start, + ].contiguous() + # quantize this slice using pre-calculated global scale, to + # get the blockwise dynamic scales, they will be frozen + # after this point + scale, _data_lp = nvfp4_quantize( + tensor_slice, + per_tensor_scale=nvfp4_global_scale, + ) + group_qparams.append(scale) + # TODO(future PR): simpler version of `nvfp4_quantize` which + # just calculates the scale, since we are throwing away the + # quantized packed data here. For now, just call the full + # one. + del _data_lp + + else: + raise AssertionError("unsupported") # Quantize each column and propagate errors to subsequent columns for k in range(G_k_start - B_cur_k_start, G_k_end - B_cur_k_start): @@ -354,6 +461,12 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): scale=quantized_tensor.scale, ) dq = q.dequantize(output_dtype=torch.float) + elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + dq = _nvfp4_with_precalculated_scales_qdq( + w_t, + nvfp4_global_scale, + scale.squeeze(-1), + ) err1 = (w_t - dq) / Hinv_cur[k, k] B_cur[:, k:] -= err1.matmul(Hinv_cur[k, k:].unsqueeze(0)) @@ -385,6 +498,48 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): result = Int8Tensor.from_hp( W_t, granularity=base_config.granularity, scale=quantized_tensor.scale ) + else: + N, K = W_t.shape + # TODO(future PR): clean up the line below. Context: the current nvfp4 + # code follows the int4 code - we save the blockwise scales to an array + # and concat it. This leads to the necessity of t().contiguous().t() to + # get the scales back into the right layout for W_t. This is not + # intuitive, likely better to initialize the scales holder ahead of time + # and write the scales directly to their final place. + combined_scale = ( + torch.cat(group_qparams, dim=0).reshape(K // group_size, N).t().contiguous() + ) + qdata = _nvfp4_with_precalculated_scales_q( + W_t, + nvfp4_global_scale, + combined_scale, + ) + + act_quant_kwargs = QuantizeTensorToNVFP4Kwargs( + use_dynamic_per_tensor_scale=base_config.use_dynamic_per_tensor_scale, + use_triton_kernel=base_config.use_triton_kernel, + is_swizzled_scales=True, + ) + + # swizzle the block scales + combined_scale_swizzled = to_blocked(combined_scale).flatten() + scale_N, scale_K = hp_data_dims_to_swizzled_scale_dims_nvfp4(N, K) + combined_scale_swizzled = combined_scale_swizzled.view(scale_N, scale_K) + + result = NVFP4Tensor( + qdata, + combined_scale_swizzled, + block_size=group_size, + orig_dtype=W_t.dtype, + per_tensor_scale=nvfp4_global_scale, + # TODO(future): get act_per_tensor_scale from calibration data? + # for now, set it to None here to calculate it dynamically at + # runtime + act_per_tensor_scale=None, + is_swizzled_scales=True, + use_triton_kernel=base_config.use_triton_kernel, + act_quant_kwargs=act_quant_kwargs, + ) return result diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 57ab7ed974..a8415828c2 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -230,7 +230,8 @@ def main(): # Generate output directory name from args model_name = args.model_id.split("/")[-1] # Get last part of model ID - output_dir = f"{model_name}_{args.quantization}" + # TODO(before land): make output dir configurable + output_dir = f"/home/dev/tmp/20260420_{model_name}_{args.quantization}" if args.quantization != "none": output_dir += f"_gs{args.group_size}" @@ -358,11 +359,16 @@ def main(): "--model_args", f"pretrained={output_dir}", "--tasks", - "leaderboard_bbh", - "--num_fewshot", - "3", + # "leaderboard_bbh", + # "gsm8k", + "wikitext", + # "--num_fewshot", + # "3", "--batch_size", - "auto", + # "auto", + "1", + # "--limit", + # "20", ] print(f"Running command: {' '.join(lm_eval_cmd)}") From 3c92c1a36ffe165c0e6c39f38808754f01494b54 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Mon, 20 Apr 2026 20:52:31 +0000 Subject: [PATCH 02/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 1 + torchao/prototype/gptq/gptq_example.py | 62 +++++++++++++++++++++++--- 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index ac6f7cdc6e..c2076e1d2d 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -120,6 +120,7 @@ def _gptq_config_transform( ) return module elif config.step == "convert": + print('gptq convert') # Quantization phase: tensor should be an GPTQObserverTensor if not isinstance(tensor, GPTQObserverTensor): raise ValueError( diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index a8415828c2..289ec98ebb 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -17,6 +17,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer from torchao.prototype.gptq import GPTQConfig +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerRow @@ -98,6 +101,8 @@ def prepare_dataset( hf_dataset_id = dataset_map.get(dataset_id, dataset_id) # Load dataset and preprocess + # TODO(before land): clean up below + dataset_split = "train" train_dataset_raw = load_dataset(hf_dataset_id, split=dataset_split, streaming=True) train_dataset_raw = train_dataset_raw.shuffle(seed=seed, buffer_size=1_000) @@ -149,7 +154,8 @@ def parse_args(): parser.add_argument( "--num-calibration-samples", type=int, - default=128, + # default=128, + default=2, help="Number of calibration samples to use", ) parser.add_argument( @@ -177,6 +183,9 @@ def parse_args(): "int8-rtn", "int8-gptq-sequential", "int8-gptq-nonsequential", + "nvfp4-rtn", + "nvfp4-gptq-sequential", + "nvfp4-gptq-nonsequential", ], help="Quantization method to use", ) @@ -241,6 +250,8 @@ def main(): "int4-gptq-nonsequential", "int8-gptq-sequential", "int8-gptq-nonsequential", + "nvfp4-gptq-sequential", + "nvfp4-gptq-nonsequential", ]: output_dir += f"_{args.dataset_id}_n{args.num_calibration_samples}" output_dir += f"_damp{args.percdamp}_bs{args.gptq_block_size}" @@ -250,29 +261,51 @@ def main(): # Handle different quantization methods quantization_start_time = time.time() + def skip_lm_head(module, fqn): + # TODO(before land): remove o_proj from below, it's just for debugging + return isinstance(module, torch.nn.Linear) and "lm_head" not in fqn and "o_proj" in fqn + + if args.quantization == "int4-rtn": print("Applying Int4 RTN (Round-To-Nearest) quantization...") config = Int4WeightOnlyConfig(group_size=args.group_size) - quantize_(model, config, filter_fn=None) + quantize_(model, config, filter_fn=skip_lm_head) elif args.quantization == "int8-rtn": print("Applying Int8 RTN (Round-To-Nearest) quantization...") config = Int8WeightOnlyConfig(version=2, granularity=PerRow()) - quantize_(model, config, filter_fn=None) + quantize_(model, config, filter_fn=skip_lm_head) + + elif args.quantization == "nvfp4-rtn": + print("Applying NVFP4 RTN (Round-To-Nearest) quantization...") + + config = NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ) + quantize_(model, config, filter_fn=skip_lm_head) elif args.quantization in [ "int4-gptq-sequential", "int4-gptq-nonsequential", "int8-gptq-sequential", "int8-gptq-nonsequential", + "nvfp4-gptq-sequential", + "nvfp4-gptq-nonsequential", ]: # Determine base config based on quantization type if "int4" in args.quantization: base_config = Int4WeightOnlyConfig(group_size=args.group_size) quant_type = "Int4" - else: # int8 + elif "int8" in args.quantization: base_config = Int8WeightOnlyConfig(granularity=PerRow(), version=2) quant_type = "Int8" + else: # nvfp4 + base_config = NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ) + quant_type = "NVFP4" # First application: wrap weights with GPTQObserverTensor (observe step) print( @@ -284,7 +317,7 @@ def main(): percdamp=args.percdamp, gptq_quantize_block_size=args.gptq_block_size, ) - quantize_(model, observe_config, filter_fn=None) + quantize_(model, observe_config, filter_fn=skip_lm_head) # Prepare calibration dataset print( @@ -316,7 +349,7 @@ def main(): for seq in tqdm(dataset, desc="Calibrating"): model(seq.to(input_device)) # Apply quantization - quantize_(model, convert_config, filter_fn=None) + quantize_(model, convert_config, filter_fn=skip_lm_head) else: # sequential print(f"Applying {quant_type} GPTQ quantization (sequential)...") sequential_quantize(model, dataset, convert_config) @@ -334,6 +367,23 @@ def main(): # Save model to generated output directory print(f"Saving model to {output_dir}...") tokenizer.save_pretrained(output_dir) + print(model) + if model.config.tie_word_embeddings: + model.config.tie_word_embeddings = False + model._tied_weights_keys = {} + model.lm_head.weight = torch.nn.Parameter( + model.lm_head.weight.clone(), requires_grad=False + ) + + # monkey patch huggingface to skip weight tie checks + # TODO(before land): debug why i'm hitting this error here, + # as nvfp4 works in other hf models just fine + # import transformers.modeling_utils + # transformers.modeling_utils.remove_tied_weights_from_state_dict = lambda state_dict, model: state_dict + + import transformers + assert transformers.__version__ == '4.57.6', "unsupported" + model.save_pretrained(output_dir, safe_serialization=False) print("DONE!") From a669b9e8b5b331dec619d5995e2d87674a637450 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 00:28:42 +0000 Subject: [PATCH 03/16] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 18 +++++++++++++++++- torchao/prototype/gptq/gptq_example.py | 16 +++++----------- 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index 0bae6ec4da..a67804590c 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -20,7 +20,10 @@ ) from torchao.quantization import Int4WeightOnlyConfig, Int8WeightOnlyConfig, quantize_ from torchao.quantization.granularity import PerRow -from torchao.utils import _is_mslk_available +from torchao.utils import ( + _is_mslk_available, + is_sm_at_least_100, +) def _calculate_hessian(inputs, device=None): @@ -449,6 +452,13 @@ def test_gptq_quantize_function(self, base_config): ) def test_gptq_quantize_better_than_naive(self, base_config): """Test that GPTQ produces lower error than naive quantization.""" + + if ( + isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig) + and not is_sm_at_least_100() + ): + pytest.skip("CUDA capability >= 10.0 required for nvfp4") + torch.manual_seed(43) # Create weight and realistic Hessian from actual activations @@ -526,6 +536,12 @@ def test_gptq_quantize_better_than_naive(self, base_config): ], ) def test_gptq_sqnr(self, base_config): + if ( + isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig) + and not is_sm_at_least_100() + ): + pytest.skip("CUDA capability >= 10.0 required for nvfp4") + torch.manual_seed(43) model = ToyLinearModel(m=512, n=2048, k=1024).cuda().to(torch.bfloat16) diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index a8415828c2..57ab7ed974 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -230,8 +230,7 @@ def main(): # Generate output directory name from args model_name = args.model_id.split("/")[-1] # Get last part of model ID - # TODO(before land): make output dir configurable - output_dir = f"/home/dev/tmp/20260420_{model_name}_{args.quantization}" + output_dir = f"{model_name}_{args.quantization}" if args.quantization != "none": output_dir += f"_gs{args.group_size}" @@ -359,16 +358,11 @@ def main(): "--model_args", f"pretrained={output_dir}", "--tasks", - # "leaderboard_bbh", - # "gsm8k", - "wikitext", - # "--num_fewshot", - # "3", + "leaderboard_bbh", + "--num_fewshot", + "3", "--batch_size", - # "auto", - "1", - # "--limit", - # "20", + "auto", ] print(f"Running command: {' '.join(lm_eval_cmd)}") From 53bd8d0214d30d1d2c287cd5660f3011afd71226 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 01:00:51 +0000 Subject: [PATCH 04/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 10 +++++++- torchao/prototype/gptq/gptq_example.py | 33 +++++++++++--------------- 2 files changed, 23 insertions(+), 20 deletions(-) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index b7011e68f6..6d5d0f2d80 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -99,6 +99,11 @@ def __post_init__(self): ) +# simple progress counter for GPTQ convert +# TODO(future): make this cleaner, will require a refactor +gptq_convert_layer_counter = 0 + + @register_quantize_module_handler(GPTQConfig) def _gptq_config_transform( module: torch.nn.Module, config: GPTQConfig, *, parameter_name="weight" @@ -120,7 +125,10 @@ def _gptq_config_transform( ) return module elif config.step == "convert": - print("gptq convert") + global gptq_convert_layer_counter + print(f"gptq convert {gptq_convert_layer_counter}") + gptq_convert_layer_counter += 1 + # Quantization phase: tensor should be an GPTQObserverTensor if not isinstance(tensor, GPTQObserverTensor): raise ValueError( diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 041ef572b7..38473d11f9 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -154,8 +154,7 @@ def parse_args(): parser.add_argument( "--num-calibration-samples", type=int, - # default=128, - default=2, + default=128, help="Number of calibration samples to use", ) parser.add_argument( @@ -265,9 +264,8 @@ def main(): def skip_lm_head(module, fqn): # TODO(before land): remove o_proj from below, it's just for debugging return ( - isinstance(module, torch.nn.Linear) - and "lm_head" not in fqn - and "o_proj" in fqn + isinstance(module, torch.nn.Linear) and "lm_head" not in fqn + # and "o_proj" in fqn ) if args.quantization == "int4-rtn": @@ -352,6 +350,13 @@ def skip_lm_head(module, fqn): # Run calibration for seq in tqdm(dataset, desc="Calibrating"): model(seq.to(input_device)) + # Print total # of GPTQ modules + num_gptq_weights = 0 + for name, param in model.named_parameters(): + # TODO(before land): better check + if "GPTQ" in str(type(param)): + num_gptq_weights += 1 + print(f"Total GPTQ weights to convert: {num_gptq_weights}") # Apply quantization quantize_(model, convert_config, filter_fn=skip_lm_head) else: # sequential @@ -372,22 +377,12 @@ def skip_lm_head(module, fqn): print(f"Saving model to {output_dir}...") tokenizer.save_pretrained(output_dir) print(model) - if model.config.tie_word_embeddings: - model.config.tie_word_embeddings = False - model._tied_weights_keys = {} - model.lm_head.weight = torch.nn.Parameter( - model.lm_head.weight.clone(), requires_grad=False - ) - - # monkey patch huggingface to skip weight tie checks - # TODO(before land): debug why i'm hitting this error here, - # as nvfp4 works in other hf models just fine - # import transformers.modeling_utils - # transformers.modeling_utils.remove_tied_weights_from_state_dict = lambda state_dict, model: state_dict - import transformers + # transformers 5.0.0 have a lot of errors with nvfp4 subclasses + # TODO(before land): debug this further + import transformers - assert transformers.__version__ == "4.57.6", "unsupported" + assert transformers.__version__ == "4.57.6", "unsupported" model.save_pretrained(output_dir, safe_serialization=False) From 4c86363542dc359f3a3c684c2931145e50109d25 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 10:21:32 +0000 Subject: [PATCH 05/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index ac6f7cdc6e..1a83aee684 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -319,6 +319,7 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): block_size = list(block_size) group_size = block_size[-1] elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): + assert base_config.use_dynamic_per_tensor_scale, "unsupported" group_size = 16 block_size = [1, group_size] # for per-tensor nvfp4, we need to calculate the global scale over the From 9b7dc74acae46d3f017792490981f6646affa52a Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 13:55:17 +0000 Subject: [PATCH 06/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/gptq_example.py | 133 ++++++++++++++---- ...vfp4_llama3_2_1b_nonsequential_wikitext.sh | 22 +++ 2 files changed, 125 insertions(+), 30 deletions(-) create mode 100755 torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 38473d11f9..08c71cd1d0 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -16,7 +16,12 @@ from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer +from packaging.version import Version + +import transformers + from torchao.prototype.gptq import GPTQConfig +from torchao.prototype.gptq.observer import GPTQObserverTensor from torchao.prototype.mx_formats.inference_workflow import ( NVFP4DynamicActivationNVFP4WeightConfig, ) @@ -96,14 +101,19 @@ def prepare_dataset( dataset_map = { "hellaswag": "Rowan/hellaswag", "ultrachat200k": "HuggingFaceH4/ultrachat_200k", + "c4": ("allenai/c4", "en"), } - hf_dataset_id = dataset_map.get(dataset_id, dataset_id) + dataset_entry = dataset_map.get(dataset_id, dataset_id) + if isinstance(dataset_entry, tuple): + hf_dataset_id, hf_config_name = dataset_entry + else: + hf_dataset_id, hf_config_name = dataset_entry, None # Load dataset and preprocess - # TODO(before land): clean up below - dataset_split = "train" - train_dataset_raw = load_dataset(hf_dataset_id, split=dataset_split, streaming=True) + train_dataset_raw = load_dataset( + hf_dataset_id, hf_config_name, split=dataset_split, streaming=True + ) train_dataset_raw = train_dataset_raw.shuffle(seed=seed, buffer_size=1_000) def preprocess_hellaswag(example): @@ -167,9 +177,15 @@ def parse_args(): "--dataset-id", type=str, default="ultrachat200k", - choices=["hellaswag", "ultrachat200k"], + choices=["hellaswag", "ultrachat200k", "c4"], help="Dataset for calibration (hellaswag or ultrachat200k)", ) + parser.add_argument( + "--dataset-split", + type=str, + default="train_sft", + help="Dataset split to use for calibration", + ) parser.add_argument( "--quantization", type=str, @@ -206,12 +222,62 @@ def parse_args(): default=1024, help="Block size for GPTQ quantization", ) + parser.add_argument( + "--lm-eval-tasks", + type=str, + default="leaderboard_bbh", + help="Comma-separated tasks for lm_eval", + ) + parser.add_argument( + "--num-fewshot", + type=int, + default=3, + help="Number of few-shot examples for lm_eval (0 to disable)", + ) + parser.add_argument( + "--lm-eval-batch-size", + type=str, + default="auto", + help="Batch size for lm_eval (default: auto)", + ) + parser.add_argument( + "--lm-eval-limit", + type=int, + default=None, + help="Limit number of examples per task for lm_eval (default: no limit)", + ) + parser.add_argument( + "--output-dir-prefix", + type=str, + required=True, + help="Prefix for the output directory (e.g. /home/user/tmp/20260420)", + ) + parser.add_argument( + "--skip-lm-eval", + action="store_true", + default=False, + help="Skip running lm_eval after quantization, useful for quickly iterating on lm_eval arguments", + ) + parser.add_argument( + "--o-proj-only", + action="store_true", + default=False, + help="Only quantize `o_proj` layers, useful for faster GPTQ runs for debugging", + ) return parser.parse_args() def main(): args = parse_args() + # lm_eval batch_size="auto" with nvfp4 gptq causes the error in + # MSLK nvfp4 triton kernel, likely an unsupported shape: + # https://gist.github.com/vkuzo/b71ca46365dee017d1602e9638d91603 + # TODO(future): debug and fix this. For now, the workaround is + # for the user to manually specify lm_eval batch_size. + if "nvfp4" in args.quantization: + assert args.lm_eval_batch_size != "auto", "unsupported" + # Map dtype string to torch dtype dtype_map = { "float32": torch.float32, @@ -238,9 +304,7 @@ def main(): # Generate output directory name from args model_name = args.model_id.split("/")[-1] # Get last part of model ID - output_dir = f"{model_name}_{args.quantization}" - # TODO(before land): make output dir configurable - output_dir = f"/home/dev/tmp/20260420_{model_name}_{args.quantization}" + output_dir = f"{args.output_dir_prefix}_{model_name}_{args.quantization}" if args.quantization != "none": output_dir += f"_gs{args.group_size}" @@ -262,21 +326,28 @@ def main(): quantization_start_time = time.time() def skip_lm_head(module, fqn): - # TODO(before land): remove o_proj from below, it's just for debugging + return isinstance(module, torch.nn.Linear) and "lm_head" not in fqn + + def skip_lm_head_o_proj(module, fqn): return ( - isinstance(module, torch.nn.Linear) and "lm_head" not in fqn - # and "o_proj" in fqn + isinstance(module, torch.nn.Linear) + and "lm_head" not in fqn + and "o_proj" in fqn ) + filter_fn_to_use = skip_lm_head + if args.o_proj_only: + filter_fn_to_use = skip_lm_head_o_proj + if args.quantization == "int4-rtn": print("Applying Int4 RTN (Round-To-Nearest) quantization...") config = Int4WeightOnlyConfig(group_size=args.group_size) - quantize_(model, config, filter_fn=skip_lm_head) + quantize_(model, config, filter_fn=filter_fn_to_use) elif args.quantization == "int8-rtn": print("Applying Int8 RTN (Round-To-Nearest) quantization...") config = Int8WeightOnlyConfig(version=2, granularity=PerRow()) - quantize_(model, config, filter_fn=skip_lm_head) + quantize_(model, config, filter_fn=filter_fn_to_use) elif args.quantization == "nvfp4-rtn": print("Applying NVFP4 RTN (Round-To-Nearest) quantization...") @@ -285,7 +356,7 @@ def skip_lm_head(module, fqn): use_dynamic_per_tensor_scale=True, use_triton_kernel=True, ) - quantize_(model, config, filter_fn=skip_lm_head) + quantize_(model, config, filter_fn=filter_fn_to_use) elif args.quantization in [ "int4-gptq-sequential", @@ -319,7 +390,7 @@ def skip_lm_head(module, fqn): percdamp=args.percdamp, gptq_quantize_block_size=args.gptq_block_size, ) - quantize_(model, observe_config, filter_fn=skip_lm_head) + quantize_(model, observe_config, filter_fn=filter_fn_to_use) # Prepare calibration dataset print( @@ -330,7 +401,7 @@ def skip_lm_head(module, fqn): max_seq_length, args.num_calibration_samples, dataset_id=args.dataset_id, - dataset_split="train_sft", + dataset_split=args.dataset_split, seed=42, ) @@ -353,14 +424,14 @@ def skip_lm_head(module, fqn): # Print total # of GPTQ modules num_gptq_weights = 0 for name, param in model.named_parameters(): - # TODO(before land): better check - if "GPTQ" in str(type(param)): + if isinstance(param, GPTQObserverTensor): num_gptq_weights += 1 print(f"Total GPTQ weights to convert: {num_gptq_weights}") # Apply quantization - quantize_(model, convert_config, filter_fn=skip_lm_head) + quantize_(model, convert_config, filter_fn=filter_fn_to_use) else: # sequential print(f"Applying {quant_type} GPTQ quantization (sequential)...") + assert filter_fn_to_use == skip_lm_head, "unsupported" sequential_quantize(model, dataset, convert_config) quantization_end_time = time.time() @@ -380,9 +451,9 @@ def skip_lm_head(module, fqn): # transformers 5.0.0 have a lot of errors with nvfp4 subclasses # TODO(before land): debug this further - import transformers - - assert transformers.__version__ == "4.57.6", "unsupported" + assert Version(transformers.__version__) < Version("5.0.0"), ( + f"transformers {transformers.__version__} is not supported, need < 5.0.0" + ) model.save_pretrained(output_dir, safe_serialization=False) @@ -409,18 +480,20 @@ def skip_lm_head(module, fqn): "--model_args", f"pretrained={output_dir}", "--tasks", - # "leaderboard_bbh", - "wikitext", - # "--num_fewshot", - # "3", + args.lm_eval_tasks, "--batch_size", - # "auto", - "1", - # "--limit", - # "20", + args.lm_eval_batch_size, ] + if args.num_fewshot > 0: + lm_eval_cmd += ["--num_fewshot", str(args.num_fewshot)] + if args.lm_eval_limit is not None: + lm_eval_cmd += ["--limit", str(args.lm_eval_limit)] + print(f"Running command: {' '.join(lm_eval_cmd)}") + if args.skip_lm_eval: + print("Terminating early due to skip_lm_eval=True") + return try: subprocess.run(lm_eval_cmd, check=True) except subprocess.CalledProcessError as e: diff --git a/torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh b/torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh new file mode 100755 index 0000000000..e60d17ab7a --- /dev/null +++ b/torchao/prototype/gptq/gptq_nvfp4_llama3_2_1b_nonsequential_wikitext.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# +# A quick smoke test for non-sequential GPTQ on `unsloth/Llama-3.2-1B` +# + +COMMON_ARGS="--output-dir-prefix /home/dev/tmp/20260421 --model-id unsloth/Llama-3.2-1B --lm-eval-tasks wikitext --num-fewshot 0 --lm-eval-batch-size 16" + +# baseline (bf16) +echo -e "\n\nbaseline (bf16)\n\n" +python -u torchao/prototype/gptq/gptq_example.py $COMMON_ARGS --quantization none +echo -e "done" + +# nvfp4-rtn +echo -e "\n\nnvfp4-rtn\n\n" +python -u torchao/prototype/gptq/gptq_example.py $COMMON_ARGS --quantization nvfp4-rtn +echo -e "done" + +# nvfp4-gptq-nonsequential +echo -e "\n\nnvfp4-gptq-nonsequential\n\n" +python -u torchao/prototype/gptq/gptq_example.py $COMMON_ARGS --quantization nvfp4-gptq-nonsequential --dataset-id c4 --dataset-split train +echo -e "done" From d69b32a6714efbe917ce849353597861d5d18c8a Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 14:03:33 +0000 Subject: [PATCH 07/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/gptq_example.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index 08c71cd1d0..d24097ce96 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -12,14 +12,12 @@ from typing import Any, List, Optional import torch +import transformers from datasets import load_dataset from tqdm import tqdm from transformers import AutoModelForCausalLM, AutoTokenizer from packaging.version import Version - -import transformers - from torchao.prototype.gptq import GPTQConfig from torchao.prototype.gptq.observer import GPTQObserverTensor from torchao.prototype.mx_formats.inference_workflow import ( From 294c9ccfdf495ed7cb365cff7b97d858eb4903ad Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 18:31:12 +0000 Subject: [PATCH 08/16] Update [ghstack-poisoned] --- benchmarks/benchmark_gptq.py | 76 +++++++++++++++++++++++++++++++++++ torchao/prototype/gptq/api.py | 33 ++++++++++++++- 2 files changed, 107 insertions(+), 2 deletions(-) create mode 100644 benchmarks/benchmark_gptq.py diff --git a/benchmarks/benchmark_gptq.py b/benchmarks/benchmark_gptq.py new file mode 100644 index 0000000000..142626286d --- /dev/null +++ b/benchmarks/benchmark_gptq.py @@ -0,0 +1,76 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD 3-Clause license found in the +# LICENSE file in the root directory of this source tree. + +import time + +import fire +import torch + +from torchao.prototype.gptq import GPTQConfig, gptq_quantize +from torchao.prototype.mx_formats.inference_workflow import ( + NVFP4DynamicActivationNVFP4WeightConfig, +) + + +def run( + K: int = 2048, + N: int = 4096, + profile_fname: str = None, +): + print(f"K={K}, N={N}") + + A = torch.randn(K, K, dtype=torch.float32, device="cuda") + H = A.t() @ A + + W_t = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + + config = GPTQConfig( + step="convert", + base_config=NVFP4DynamicActivationNVFP4WeightConfig( + use_dynamic_per_tensor_scale=True, + use_triton_kernel=True, + ), + ) + + # Warmup + print("Warmup...") + gptq_quantize(H.clone(), W_t.clone(), config) + torch.cuda.synchronize() + + num_runs = 5 + if profile_fname is not None: + print("Profiling run...") + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + with_stack=True, + ) as prof: + torch.cuda.synchronize() + start = time.time() + gptq_quantize(H.clone(), W_t.clone(), config) + torch.cuda.synchronize() + elapsed = time.time() - start + print(f"gptq_quantize time: {elapsed:.3f}s") + prof.export_chrome_trace(profile_fname) + print(f"Saved: {profile_fname}") + else: + print(f"Timed run ({num_runs} iterations)...") + times = [] + for _ in range(num_runs): + torch.cuda.synchronize() + start = time.time() + gptq_quantize(H.clone(), W_t.clone(), config) + torch.cuda.synchronize() + times.append(time.time() - start) + avg = sum(times) / len(times) + print(f"gptq_quantize avg time: {avg:.3f}s") + + +if __name__ == "__main__": + fire.Fire(run) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index 14881d974e..d3d7329c0a 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -4,7 +4,9 @@ # This source code is licensed under the BSD 3-Clause license found in the # LICENSE file in the root directory of this source tree. +import os import types +import warnings from dataclasses import dataclass from functools import partial from typing import Union @@ -265,6 +267,33 @@ def _nvfp4_with_precalculated_scales_q( return data_lp_packed +# Set to True to torch.compile the NVFP4 quantize/dequantize functions +# inside gptq_quantize. Gives ~3x speedup. +_use_torch_compile = True + +if _use_torch_compile: + _nvfp4_qdq_fn = torch.compile(_nvfp4_with_precalculated_scales_qdq) + _nvfp4_q_fn = torch.compile(_nvfp4_with_precalculated_scales_q) + + # Triton's default f32 division uses approximate reciprocal which + # introduces ~1 ULP error per division. In GPTQ's error propagation + # loop this compounds across columns. IEEE-compliant division rounding + # eliminates the drift. + import torch._inductor.config as _inductor_config + + if os.environ.get("TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING") == "0": + warnings.warn( + "TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING=0 may cause numerical " + "drift in GPTQ with torch.compile. " + "Consider unsetting it or setting it to 1." + ) + else: + _inductor_config.eager_numerics.division_rounding = True +else: + _nvfp4_qdq_fn = _nvfp4_with_precalculated_scales_qdq + _nvfp4_q_fn = _nvfp4_with_precalculated_scales_q + + def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): """ This function implements the GPTQ algorithm described in this paper: https://arxiv.org/abs/2210.17323 (Algorithm 1) @@ -472,7 +501,7 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): ) dq = q.dequantize(output_dtype=torch.float) elif isinstance(base_config, NVFP4DynamicActivationNVFP4WeightConfig): - dq = _nvfp4_with_precalculated_scales_qdq( + dq = _nvfp4_qdq_fn( w_t, nvfp4_global_scale, scale.squeeze(-1), @@ -519,7 +548,7 @@ def gptq_quantize(H: torch.Tensor, W_t: torch.Tensor, config: GPTQConfig): combined_scale = ( torch.cat(group_qparams, dim=0).reshape(K // group_size, N).t().contiguous() ) - qdata = _nvfp4_with_precalculated_scales_q( + qdata = _nvfp4_q_fn( W_t, nvfp4_global_scale, combined_scale, From 65fae6264a1accca9a1e44b7baf7d3b019fddc4d Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Tue, 21 Apr 2026 18:52:00 +0000 Subject: [PATCH 09/16] Update [ghstack-poisoned] --- test/prototype/gptq/test_gptqv2.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/prototype/gptq/test_gptqv2.py b/test/prototype/gptq/test_gptqv2.py index a67804590c..78af648e14 100644 --- a/test/prototype/gptq/test_gptqv2.py +++ b/test/prototype/gptq/test_gptqv2.py @@ -10,6 +10,13 @@ import torch import torch.nn.functional as F +from torchao.utils import torch_version_at_least + +pytestmark = pytest.mark.skipif( + not torch_version_at_least("2.11.0"), + reason="GPTQ prototype requires PyTorch 2.11+", +) + from torchao.prototype.gptq import ( GPTQConfig, gptq_quantize, From 5ee2ad2b6029eef0dee6f7317bf2a1eacd3560a7 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 11:21:39 +0000 Subject: [PATCH 10/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/api.py | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/torchao/prototype/gptq/api.py b/torchao/prototype/gptq/api.py index d3d7329c0a..6869353273 100644 --- a/torchao/prototype/gptq/api.py +++ b/torchao/prototype/gptq/api.py @@ -14,6 +14,8 @@ import torch import torch.nn as nn +from torchao.utils import torch_version_at_least + try: from mslk.quantize.shuffle import int4_row_quantize_zp, pack_int4 except: @@ -275,20 +277,27 @@ def _nvfp4_with_precalculated_scales_q( _nvfp4_qdq_fn = torch.compile(_nvfp4_with_precalculated_scales_qdq) _nvfp4_q_fn = torch.compile(_nvfp4_with_precalculated_scales_q) - # Triton's default f32 division uses approximate reciprocal which - # introduces ~1 ULP error per division. In GPTQ's error propagation - # loop this compounds across columns. IEEE-compliant division rounding - # eliminates the drift. - import torch._inductor.config as _inductor_config - - if os.environ.get("TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING") == "0": + if torch_version_at_least("2.11.0"): + # Triton's default f32 division uses approximate reciprocal which + # introduces ~1 ULP error per division. In GPTQ's error propagation + # loop this compounds across columns. IEEE-compliant division rounding + # eliminates the drift. + import torch._inductor.config as _inductor_config + + if os.environ.get("TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING") == "0": + warnings.warn( + "TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING=0 may cause numerical " + "drift in GPTQ with torch.compile. " + "Consider unsetting it or setting it to 1." + ) + else: + _inductor_config.eager_numerics.division_rounding = True + else: warnings.warn( - "TORCHINDUCTOR_EMULATE_DIVISION_ROUNDING=0 may cause numerical " - "drift in GPTQ with torch.compile. " - "Consider unsetting it or setting it to 1." + "PyTorch < 2.11.0 detected. Upgrade to PyTorch 2.11.0+ for " + "better GPTQ numerics with torch.compile (IEEE-compliant " + "division rounding)." ) - else: - _inductor_config.eager_numerics.division_rounding = True else: _nvfp4_qdq_fn = _nvfp4_with_precalculated_scales_qdq _nvfp4_q_fn = _nvfp4_with_precalculated_scales_q From 2adda75ff0a6f84596bf1e88782b060fed8e5821 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 11:21:39 +0000 Subject: [PATCH 11/16] Update [ghstack-poisoned] --- torchao/prototype/gptq/gptq_example.py | 28 ++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/torchao/prototype/gptq/gptq_example.py b/torchao/prototype/gptq/gptq_example.py index d24097ce96..016f0ff174 100644 --- a/torchao/prototype/gptq/gptq_example.py +++ b/torchao/prototype/gptq/gptq_example.py @@ -12,12 +12,11 @@ from typing import Any, List, Optional import torch -import transformers from datasets import load_dataset from tqdm import tqdm -from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig +from transformers.quantizers.quantizer_torchao import TorchAoHfQuantizer -from packaging.version import Version from torchao.prototype.gptq import GPTQConfig from torchao.prototype.gptq.observer import GPTQObserverTensor from torchao.prototype.mx_formats.inference_workflow import ( @@ -447,11 +446,24 @@ def skip_lm_head_o_proj(module, fqn): tokenizer.save_pretrained(output_dir) print(model) - # transformers 5.0.0 have a lot of errors with nvfp4 subclasses - # TODO(before land): debug this further - assert Version(transformers.__version__) < Version("5.0.0"), ( - f"transformers {transformers.__version__} is not supported, need < 5.0.0" - ) + if "nvfp4" in args.quantization: + import inspect + + source = inspect.getsource(TorchAoHfQuantizer.get_weight_conversions) + if "_weight_per_tensor_scale" not in source: + raise RuntimeError( + "Your version of `transformers` does not support NVFP4 serialization. " + "Please install a version that includes " + "https://github.com/huggingface/transformers/pull/45573" + ) + + if args.quantization != "none": + # Attach hf_quantizer so save_pretrained uses the flatten path for tensor + # subclasses (e.g. NVFP4Tensor) that don't have a valid storage pointer. + ao_config = base_config if "gptq" in args.quantization else config + torchao_config = TorchAoConfig(quant_type=ao_config) + model.config.quantization_config = torchao_config + model.hf_quantizer = TorchAoHfQuantizer(torchao_config) model.save_pretrained(output_dir, safe_serialization=False) From 5fe657469f1cff71df891fc9ebbdb4b823d4e722 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 18:01:24 +0000 Subject: [PATCH 12/16] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 69 +++++++++++++++++++ torchao/prototype/mx_formats/mx_tensor.py | 4 +- torchao/prototype/mx_formats/nvfp4_tensor.py | 11 +++ 3 files changed, 82 insertions(+), 2 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index cdc2b31aa5..23bc6b5a8b 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -17,6 +17,7 @@ NVFP4DynamicActivationNVFP4WeightConfig, NVFP4WeightOnlyConfig, ) +from torchao.prototype.mx_formats.nvfp4_tensor import NVFP4Tensor from torchao.quantization import quantize_ from torchao.quantization.quantize_.common import KernelPreference from torchao.quantization.utils import compute_error @@ -426,3 +427,71 @@ def test_nvfp4_static_vs_dynamic_quantization(): assert torch.equal(y_dynamic, y_static), ( "Expect dynamic and static quant result to be equal" ) + + +class GroupedMMModel(nn.Module): + """A toy model whose only op in forward is torch._grouped_mm.""" + + def __init__(self, E, K, N, device="cuda", dtype=torch.bfloat16): + super().__init__() + self.weight = nn.Parameter(torch.randn(E, N, K, device=device, dtype=dtype)) + + def forward(self, x, offs): + return torch._grouped_mm(x, self.weight.transpose(-2, -1), offs=offs) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not torch_version_at_least("2.8.0"), reason="torch >= 2.8.0 required" +) +@pytest.mark.skipif( + not is_sm_at_least_100(), reason="CUDA capability >= 10.0 required for NVFP4" +) +@torch.no_grad() +@skip_if_rocm("ROCm float4 gemm require gfx950") +def test_grouped_mm_nvfp4(): + """Smoke test: torch._grouped_mm forward pass with bfloat16 inputs.""" + E, K, N = 4, 128, 256 + m_per_group = [32, 96, 16, 112] + total_m = sum(m_per_group) + + device = "cuda" + dtype = torch.bfloat16 + + model_ref = GroupedMMModel(E, K, N, device=device, dtype=dtype) + model = copy.deepcopy(model_ref) + + x = torch.randn(total_m, K, device=device, dtype=dtype) + offs = torch.tensor( + [sum(m_per_group[: i + 1]) for i in range(E)], + device=device, + dtype=torch.int32, + ) + + # reference path + y_ref = model_ref(x, offs) + + # quantized path + quantize_( + model, + NVFP4DynamicActivationNVFP4WeightConfig( + use_triton_kernel=False, + ), + filter_fn=lambda mod, *args: isinstance(mod, GroupedMMModel) + and hasattr(mod, "weight"), + ) + assert isinstance(model.weight, NVFP4Tensor), ( + f"Expected NVFP4Tensor weight, got {type(model.weight)}" + ) + w_sqnr = compute_error(model_ref.weight, model.weight.dequantize()) + assert w_sqnr > 18.0 + + ww = model.weight + wwt = ww.transpose(-2, -1) + assert tuple(wwt.shape) == (ww.shape[0], ww.shape[2], ww.shape[1]) + + # For now, this is emulated. In the near future we'll hook up + # a real nvfp4 grouped gemm. + y = model(x, offs) + y_sqnr = compute_error(y_ref, y) + assert y_sqnr > 18.0 diff --git a/torchao/prototype/mx_formats/mx_tensor.py b/torchao/prototype/mx_formats/mx_tensor.py index 0f0f45a5fd..c6ca81d08a 100644 --- a/torchao/prototype/mx_formats/mx_tensor.py +++ b/torchao/prototype/mx_formats/mx_tensor.py @@ -486,7 +486,7 @@ def tensor_size_hp_to_fp4x2(orig_size, is_contiguous): else: assert len(orig_size) == 3, "unsupported" # only supporting dim0, dim1, dim2 and dim0, dim2, dim1 orders - new_size = [new_size[0], new_size[2] // 2, new_size[1]] + new_size = [new_size[0], new_size[1] // 2, new_size[2]] return new_size @@ -500,7 +500,7 @@ def tensor_size_fp4x2_to_hp(orig_size, is_contiguous): else: assert len(orig_size) == 3, "unsupported" # only supporting dim0, dim1, dim2 and dim0, dim2, dim1 orders - new_size = [new_size[0], new_size[2] * 2, new_size[1]] + new_size = [new_size[0], new_size[1] * 2, new_size[2]] return new_size diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index 766692bacf..d7c4a3b671 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -656,6 +656,17 @@ def nvfp4_addmm(func, types, args, kwargs): return _addmm_nvfp4_dispatch(input_tensor, weight_tensor, func, bias=bias) +@implements([aten._grouped_mm.default]) +def nvfp4_grouped_mm(func, types, args, kwargs): + mat_a, mat_b = args[0], args[1] + # temporary: dequantize nvfp4 tensor to call original grouped_mm + # TODO(future PR): enable nvfp4 with per-expert outer scale (current code + # has a single outer scale across all experts). + # TODO(future PR): hook up real nvfp4 grouped_mm + mat_b_dq = mat_b.dequantize() + return func(mat_a, mat_b_dq, *args[2:], **kwargs) + + def per_tensor_amax_to_scale(amax: torch.Tensor) -> torch.Tensor: """Convert per-tensor amax to per-tensor scale for NVFP4 quantization. From 5292f2f1ba3047e3c9a22808456853af57d7b2d3 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 19:59:55 +0000 Subject: [PATCH 13/16] Update [ghstack-poisoned] --- .../prototype/mx_formats/test_nvfp4_tensor.py | 60 +++++++++++++++++++ torchao/prototype/mx_formats/nvfp4_tensor.py | 35 ++++++++++- 2 files changed, 92 insertions(+), 3 deletions(-) diff --git a/test/prototype/mx_formats/test_nvfp4_tensor.py b/test/prototype/mx_formats/test_nvfp4_tensor.py index c90a9cc68a..afcb94f0ad 100644 --- a/test/prototype/mx_formats/test_nvfp4_tensor.py +++ b/test/prototype/mx_formats/test_nvfp4_tensor.py @@ -712,3 +712,63 @@ def test_nvfp4_matmul_optional_per_tensor_scale(shapes, a_has_scale, use_triton_ sqnr = compute_error(C_ref, C_nvfp4) SQNR_THRESHOLD = 16.0 assert sqnr >= SQNR_THRESHOLD, f"SQNR {sqnr:.2f} < {SQNR_THRESHOLD}, {a_has_scale=}" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif( + not torch_version_at_least("2.8.0"), reason="torch.compile requires PyTorch 2.8+" +) +def test_nvfp4_per_expert_scale(): + # per-tensor scale reference + E, K, N = 2, 64, 128 + x0 = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + x1 = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 2 + tensor_amax_x0 = torch.max(torch.abs(x0)) + per_tensor_scale_x0 = per_tensor_amax_to_scale(tensor_amax_x0) + tensor_amax_x1 = torch.max(torch.abs(x1)) + per_tensor_scale_x1 = per_tensor_amax_to_scale(tensor_amax_x1) + x0_nvfp4 = NVFP4Tensor.to_nvfp4( + x0, per_tensor_scale=per_tensor_scale_x0, is_swizzled_scales=False + ) + x1_nvfp4 = NVFP4Tensor.to_nvfp4( + x1, per_tensor_scale=per_tensor_scale_x1, is_swizzled_scales=False + ) + + xc = torch.cat([x0, x1], dim=0).view(E, N, K) + scalec = torch.cat( + [ + per_tensor_scale_x0.view( + 1, + ), + per_tensor_scale_x1.view( + 1, + ), + ], + dim=0, + ).view(E, 1, 1) + + xc_nvfp4 = NVFP4Tensor.to_nvfp4( + xc, per_tensor_scale=scalec, is_swizzled_scales=False + ) + + # internals must match + torch.testing.assert_close( + torch.cat([x0_nvfp4.qdata, x1_nvfp4.qdata], dim=0).view(E, N, K // 2), + xc_nvfp4.qdata, + atol=0, + rtol=0, + ) + + torch.testing.assert_close( + torch.cat([x0_nvfp4.scale, x1_nvfp4.scale], dim=0).view(E, N, K // 16), + xc_nvfp4.scale, + atol=0, + rtol=0, + ) + + x0_dq = x0_nvfp4.dequantize() + x1_dq = x1_nvfp4.dequantize() + xc_dq = xc_nvfp4.dequantize() + + xc_dq_ref = torch.cat([x0_dq, x1_dq], dim=0).view(E, N, K) + torch.testing.assert_close(xc_dq_ref, xc_dq, atol=0, rtol=0) diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index d7c4a3b671..af132d96a9 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -104,6 +104,11 @@ def __new__( requires_grad=False, ) + if per_tensor_scale is not None: + # 0: per-tensor scale, with shape (,) + # 3: per-expert scale, with shape (E, 1, 1) + assert len(per_tensor_scale.shape) in (0, 3), "unsupported" + self.qdata = qdata self.scale = scale self.block_size = block_size @@ -224,7 +229,7 @@ def dequantize(self, output_dtype: Optional[torch.dtype] = None) -> torch.Tensor return result def get_hp_scales(self) -> torch.Tensor: - """Get the scales of the NVFP4Tensor in original dtype. + """Get the scales of the NVFP4Tensor in float32. Returns: torch.Tensor: Scales of the NVFP4Tensor @@ -243,9 +248,9 @@ def get_hp_scales(self) -> torch.Tensor: ) return ( - scale_e4m3.to(self.orig_dtype) + scale_e4m3.to(torch.float) if self.per_tensor_scale is None - else self.per_tensor_scale * scale_e4m3.to(self.orig_dtype) + else self.per_tensor_scale * scale_e4m3.to(torch.float) ) @classmethod @@ -283,6 +288,13 @@ def _same_metadata(cls, self: "NVFP4Tensor", src: "NVFP4Tensor") -> bool: implements = NVFP4Tensor.implements +def _assert_no_per_expert_scale(t): + if t.per_tensor_scale is not None: + assert len(t.per_tensor_scale.shape) == 0, ( + "per-expert scale not supported in this codepath" + ) + + # TODO(future PR): move this to AOBaseTensor (will require debugging/fixing CI) @implements([aten._to_copy.default]) def nvfp4_to_copy(func, types, args, kwargs): @@ -365,6 +377,7 @@ def nvfp4_slice(func, types, args, kwargs): assert len(x.shape) == 2, ( f"only rank 2 is supported for slice, got rank {len(x.shape)}" ) + _assert_no_per_expert_scale(x) sliced_data, sliced_scale = _swizzle_aware_slice(x, dim, start, end, step) @@ -388,6 +401,7 @@ def nvfp4_slice(func, types, args, kwargs): def nvfp4_t(func, types, args, kwargs): # For now, only transpose(input, 0, 1) is supported. old = args[0] + _assert_no_per_expert_scale(old) new = NVFP4Tensor( old.qdata.t(), old.scale.t(), @@ -405,6 +419,7 @@ def nvfp4_t(func, types, args, kwargs): @implements([aten.transpose.int]) def nvfp4_transpose(func, types, args, kwargs): old, dim0, dim1 = args + _assert_no_per_expert_scale(old) assert len(old.shape) == 3, f"unsupported rank {len(old.shape)}" valid_3d_dims = ((1, 2), (2, 1), (-1, -2), (-2, -1)) assert (dim0, dim1) in valid_3d_dims, f"transpose unsupported for {dim0=} {dim1=}" @@ -427,6 +442,7 @@ def nvfp4_transpose(func, types, args, kwargs): @implements([aten.view.default]) def nvfp4_view_op(func, types, args, kwargs): data = args[0].qdata + _assert_no_per_expert_scale(args[0]) new_size = args[1] new_size = tensor_size_hp_to_fp4x2(new_size, data.is_contiguous()) new_data = func(data, new_size, *args[2:], **kwargs) @@ -448,6 +464,7 @@ def nvfp4_select(func, types, args, kwargs): old, dim, index = args assert dim == 0, f"NVFP4Tensor aten.select.int with {dim=} is not yet supported" assert len(old.qdata.shape) == len(old.scale.shape), "unsupported" + _assert_no_per_expert_scale(old) new = old.__class__( old.qdata[index], old.scale[index], @@ -476,6 +493,8 @@ def _addmm_nvfp4_dispatch( assert a.block_size == 16, f"NVFP4 requires block_size=16, got {a.block_size}" assert b.block_size == 16, f"NVFP4 requires block_size=16, got {b.block_size}" assert len(a.shape) == 2 and len(b.shape) == 2 + _assert_no_per_expert_scale(a) + _assert_no_per_expert_scale(b) M, K = a.shape[0], a.shape[1] N = b.shape[1] @@ -564,6 +583,7 @@ def nvfp4_linear(func, types, args, kwargs): if not isinstance(weight_tensor, NVFP4Tensor): raise NotImplementedError("NVFP4Tensor: weight must be NVFP4Tensor") + _assert_no_per_expert_scale(weight_tensor) if weight_tensor.act_quant_kwargs is None: # weight_only quant @@ -597,6 +617,7 @@ def nvfp4_mm(func, types, args, kwargs): if not isinstance(weight_tensor, NVFP4Tensor): raise NotImplementedError("NVFP4Tensor: weight must be NVFP4Tensor") + _assert_no_per_expert_scale(weight_tensor) if weight_tensor.act_quant_kwargs is None: weight_dequant = weight_tensor.dequantize(weight_tensor.orig_dtype) @@ -629,6 +650,7 @@ def nvfp4_addmm(func, types, args, kwargs): if not isinstance(weight_tensor, NVFP4Tensor): raise NotImplementedError("NVFP4Tensor: weight must be NVFP4Tensor") + _assert_no_per_expert_scale(weight_tensor) if weight_tensor.act_quant_kwargs is None: weight_dequant = weight_tensor.dequantize(weight_tensor.orig_dtype) @@ -740,6 +762,13 @@ def nvfp4_quantize( # This will likely be calibrated but # we want the per_tensor_scale ~= amax of the block_scale_fp32 block_scale_fp32 = block_scale.to(torch.float32) + + # if the per_tensor_scale is multidimensional, we are in the special + # case that handles the 3d weight tensor with per-expert outer scales. + # Since the code below is 2d, we convert this scale to 2d. + if len(per_tensor_scale.shape) == 3: + per_tensor_scale = per_tensor_scale.squeeze(-1) + # Quantize the blockwise scales w/ the per_tensor_scale scaled_block_scales = block_scale_fp32 / per_tensor_scale scaled_block_scales_fp8 = torch.clamp( From f6792167d217531b73d836811e904fdabed4e858 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Wed, 22 Apr 2026 21:08:34 +0000 Subject: [PATCH 14/16] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 23 ++++---- .../mx_formats/inference_workflow.py | 15 +++++- torchao/prototype/mx_formats/nvfp4_tensor.py | 53 ++++++++++++++++--- 3 files changed, 68 insertions(+), 23 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index 23bc6b5a8b..9091a60594 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -268,17 +268,6 @@ def test_narrow_similar_to_vllm(self): ) self._test_narrow_similar_to_vllm(config) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - @pytest.mark.skipif( - not torch_version_at_least("2.8.0"), - reason="torch.compile requires PyTorch 2.8+", - ) - def test_nvfp4_quantize_3d_param_similar_to_vllm(self): - config = NVFP4WeightOnlyConfig( - use_dynamic_per_tensor_scale=False, - ) - self._test_quantize_3d_param_similar_to_vllm(config) - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.skipif( @@ -459,9 +448,15 @@ def test_grouped_mm_nvfp4(): dtype = torch.bfloat16 model_ref = GroupedMMModel(E, K, N, device=device, dtype=dtype) + + # make the future nvfp4 weight scales interesting + model_ref.weight[0, :, :] *= 10.0 + model_ref.weight[-1, :, :] *= 1e-3 + model = copy.deepcopy(model_ref) x = torch.randn(total_m, K, device=device, dtype=dtype) + offs = torch.tensor( [sum(m_per_group[: i + 1]) for i in range(E)], device=device, @@ -483,6 +478,8 @@ def test_grouped_mm_nvfp4(): assert isinstance(model.weight, NVFP4Tensor), ( f"Expected NVFP4Tensor weight, got {type(model.weight)}" ) + assert model.weight.per_tensor_scale.shape == (E, 1, 1) + # breakpoint() w_sqnr = compute_error(model_ref.weight, model.weight.dequantize()) assert w_sqnr > 18.0 @@ -490,8 +487,6 @@ def test_grouped_mm_nvfp4(): wwt = ww.transpose(-2, -1) assert tuple(wwt.shape) == (ww.shape[0], ww.shape[2], ww.shape[1]) - # For now, this is emulated. In the near future we'll hook up - # a real nvfp4 grouped gemm. y = model(x, offs) y_sqnr = compute_error(y_ref, y) - assert y_sqnr > 18.0 + assert y_sqnr > 15.0 diff --git a/torchao/prototype/mx_formats/inference_workflow.py b/torchao/prototype/mx_formats/inference_workflow.py index d8d6e27842..2c0db1af16 100644 --- a/torchao/prototype/mx_formats/inference_workflow.py +++ b/torchao/prototype/mx_formats/inference_workflow.py @@ -264,6 +264,7 @@ def _nvfp4_inference_linear_transform( return NVFP4ObservedLinear.from_float(module) elif step == QuantizationStep.CONVERT or step == "convert": + assert len(weight.shape) == 2, "3D weights not yet supported here" if not isinstance(module, NVFP4ObservedLinear): return module @@ -314,8 +315,17 @@ def _nvfp4_inference_linear_transform( per_tensor_scale = None if config.use_dynamic_per_tensor_scale: - tensor_amax = torch.max(torch.abs(weight)) - per_tensor_scale = per_tensor_amax_to_scale(tensor_amax) + if len(weight.shape) == 2: + tensor_amax = torch.max(torch.abs(weight)) + per_tensor_scale = per_tensor_amax_to_scale(tensor_amax) + else: + assert len(weight.shape) == 3, f"unsupported {weight.shape=}" + tensor_amax = torch.amax(torch.abs(weight), dim=(1, 2)) + per_tensor_scale = per_tensor_amax_to_scale(tensor_amax) + # 1D -> 3D + per_tensor_scale = per_tensor_scale.view( + per_tensor_scale.shape[0], 1, 1 + ) act_quant_kwargs = QuantizeTensorToNVFP4Kwargs( use_dynamic_per_tensor_scale=config.use_dynamic_per_tensor_scale, @@ -383,6 +393,7 @@ def _nvfp4_weight_only_linear_transform( """Quantization handler for NVFP4WeightOnlyConfig""" weight = module.weight + assert len(weight.shape) == 2, "3D weights not yet supported in this workflow" if weight.shape[-2] % 16 != 0 or weight.shape[-1] % 16 != 0: raise RuntimeError( f"NVFP4 only supports weight shape with last 2 dims divisible by 16, got {weight.shape}" diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index af132d96a9..8b18b90220 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -246,6 +246,7 @@ def get_hp_scales(self) -> torch.Tensor: scale_e4m3 = from_blocked( scale_e4m3, math.prod(leading_dims) * M, K // self.block_size ) + scale_e4m3 = scale_e4m3.view(*leading_dims, M, K // self.block_size) return ( scale_e4m3.to(torch.float) @@ -419,7 +420,6 @@ def nvfp4_t(func, types, args, kwargs): @implements([aten.transpose.int]) def nvfp4_transpose(func, types, args, kwargs): old, dim0, dim1 = args - _assert_no_per_expert_scale(old) assert len(old.shape) == 3, f"unsupported rank {len(old.shape)}" valid_3d_dims = ((1, 2), (2, 1), (-1, -2), (-2, -1)) assert (dim0, dim1) in valid_3d_dims, f"transpose unsupported for {dim0=} {dim1=}" @@ -680,13 +680,52 @@ def nvfp4_addmm(func, types, args, kwargs): @implements([aten._grouped_mm.default]) def nvfp4_grouped_mm(func, types, args, kwargs): + # TODO(before land): wrap MSLK kernel with a custom op, guard for MSLK + # not installed, etc + from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked + from torch.nn.functional import ScalingType, SwizzleType, scaled_grouped_mm + mat_a, mat_b = args[0], args[1] - # temporary: dequantize nvfp4 tensor to call original grouped_mm - # TODO(future PR): enable nvfp4 with per-expert outer scale (current code - # has a single outer scale across all experts). - # TODO(future PR): hook up real nvfp4 grouped_mm - mat_b_dq = mat_b.dequantize() - return func(mat_a, mat_b_dq, *args[2:], **kwargs) + offs = args[2] if len(args) > 2 else kwargs.get("offs", None) + assert offs is not None, "offs is required for nvfp4 grouped_mm" + + # mat_b is a transposed NVFP4Tensor: the model stores weight as (E, N, K) + # and calls weight.transpose(-2, -1). Undo the transpose to get the + # original NVFP4Tensor so we can extract qdata/scale in the right layout. + assert isinstance(mat_b, NVFP4Tensor) + is_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1) + if is_transposed: + mat_b = mat_b.transpose(-2, -1) + + E = offs.shape[0] + m_sizes = torch.diff(offs, prepend=offs.new_zeros(1)).to(torch.int64) + + # For now, quantize activation with per-expert global scales of 1.0 + # TODO(future PR): compute real per-token-group global scales from activation amax, + # will need some slow eager code or a triton kernel + a_global_scale = torch.ones(E, device=mat_a.device, dtype=torch.float32) + mat_a_qdata, mat_a_scale = nvfp4_quantize_stacked(m_sizes, mat_a, a_global_scale) + + # Weight: extract qdata and already-swizzled scales + mat_b_qdata = mat_b.qdata.view(torch.float4_e2m1fn_x2) + # Flatten 3D scale (E, padded_N, padded_K//16) -> 2D (E, padded_N * padded_K//16) + # as required by _scaled_grouped_mm for 2D-3D grouped GEMM + mat_b_scale = mat_b.scale.view(torch.float8_e4m3fn).flatten(1) + b_global_scale = mat_b.per_tensor_scale.view(E) + + # _scaled_grouped_mm expects wq.transpose(-2, -1) + return scaled_grouped_mm( + mat_a_qdata, + mat_b_qdata.transpose(-2, -1), + scale_a=[mat_a_scale, a_global_scale], + scale_recipe_a=[ScalingType.BlockWise1x16, ScalingType.TensorWise], + scale_b=[mat_b_scale, b_global_scale], + scale_recipe_b=[ScalingType.BlockWise1x16, ScalingType.TensorWise], + swizzle_a=SwizzleType.SWIZZLE_32_4_4, + swizzle_b=SwizzleType.SWIZZLE_32_4_4, + offs=offs, + output_dtype=mat_a.dtype, + ) def per_tensor_amax_to_scale(amax: torch.Tensor) -> torch.Tensor: From be9dc1b3998438643cfc8b0db832813db15933b5 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 09:27:57 +0000 Subject: [PATCH 15/16] Update [ghstack-poisoned] --- torchao/prototype/mx_formats/nvfp4_tensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index af132d96a9..52b5876a64 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -56,8 +56,8 @@ class NVFP4Tensor(TorchAOBaseTensor): Attributes: qdata: Packed FP4 data (2 values per byte) scale: Blockwise scales in float8_e4m3fn format (may be swizzled) - per_tensor_scale: Optional global per-tensor scale in float32 format - act_per_tensor_scale: Optional global per-tensor scale in float32 format, for activation + per_tensor_scale: Optional global per-tensor or per-expert scale in float32 format + act_per_tensor_scale: Optional global per-tensor or per-token-group scale in float32 format, for activation block_size (int): Block size for quantization (fixed at 16) orig_dtype (torch.dtype): Original tensor dtype before quantization is_swizzled_scales (bool): Whether scales are stored in swizzled (blocked) format From 83283cfffdb0840f0d459a55882a8ce7fb78fd26 Mon Sep 17 00:00:00 2001 From: Vasiliy Kuznetsov Date: Thu, 23 Apr 2026 10:55:00 +0000 Subject: [PATCH 16/16] Update [ghstack-poisoned] --- .../mx_formats/test_inference_workflow.py | 3 +- torchao/prototype/mx_formats/kernels.py | 84 +++++++++++++++++++ torchao/prototype/mx_formats/nvfp4_tensor.py | 33 ++++---- 3 files changed, 101 insertions(+), 19 deletions(-) diff --git a/test/prototype/mx_formats/test_inference_workflow.py b/test/prototype/mx_formats/test_inference_workflow.py index 9091a60594..de088e36ea 100644 --- a/test/prototype/mx_formats/test_inference_workflow.py +++ b/test/prototype/mx_formats/test_inference_workflow.py @@ -441,7 +441,7 @@ def forward(self, x, offs): def test_grouped_mm_nvfp4(): """Smoke test: torch._grouped_mm forward pass with bfloat16 inputs.""" E, K, N = 4, 128, 256 - m_per_group = [32, 96, 16, 112] + m_per_group = [1, 3, 4, 16] total_m = sum(m_per_group) device = "cuda" @@ -479,7 +479,6 @@ def test_grouped_mm_nvfp4(): f"Expected NVFP4Tensor weight, got {type(model.weight)}" ) assert model.weight.per_tensor_scale.shape == (E, 1, 1) - # breakpoint() w_sqnr = compute_error(model_ref.weight, model.weight.dequantize()) assert w_sqnr > 18.0 diff --git a/torchao/prototype/mx_formats/kernels.py b/torchao/prototype/mx_formats/kernels.py index 7f891d6e39..ce614d4daf 100644 --- a/torchao/prototype/mx_formats/kernels.py +++ b/torchao/prototype/mx_formats/kernels.py @@ -1249,3 +1249,87 @@ def _(x, global_scale=None): scales = scales.view(*orig_leading_dims, -1, padded_cols) xq = xq.view(*orig_leading_dims, -1, N // 2) return scales, xq + + +def mslk_calculate_group_max(x: torch.Tensor, m_sizes: torch.Tensor) -> torch.Tensor: + """Compute per-expert activation global scale (encoding convention). + + Args: + x: [M, K] concatenated activation tensor (bf16/fp16). + m_sizes: [E] int64 tensor of rows per expert. + + Returns: + Per-expert global scale in encoding convention (448 * FP4_MAX / amax), + shape [E], dtype float32. + """ + return _mslk_calculate_group_max_custom_op(x, m_sizes) + + +@torch.library.custom_op("ao::mslk_calculate_group_max", mutates_args=()) +def _mslk_calculate_group_max_custom_op( + x: torch.Tensor, m_sizes: torch.Tensor +) -> torch.Tensor: + assert _mslk_available, ( + "mslk is required for calculate_group_max. " + "Install from https://github.com/meta-pytorch/MSLK" + ) + from mslk.quantize.triton.fp4_quantize import calculate_group_max + + global_scale, _ = calculate_group_max(x, m_sizes) + return global_scale + + +@_mslk_calculate_group_max_custom_op.register_fake +def _(x, m_sizes): + E = m_sizes.shape[0] + return x.new_empty(E, dtype=torch.float32) + + +def mslk_quantize_nvfp4_stacked( + m_sizes: torch.Tensor, + x: torch.Tensor, + global_scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize concatenated MoE activations to NVFP4 with per-expert global scales. + + Args: + m_sizes: [E] int64 tensor of rows per expert. + x: [M, K] concatenated activation tensor (bf16/fp16). + global_scale: [E] fp32 per-expert global scales in encoding convention. + + Returns: + Tuple of (xq, scale): + xq: [M, K//2] float4_e2m1fn_x2 packed FP4 data. + scale: Padded+swizzled float8_e4m3fn block scales. + """ + return _mslk_quantize_nvfp4_stacked_custom_op(m_sizes, x, global_scale) + + +@torch.library.custom_op("ao::mslk_quantize_nvfp4_stacked", mutates_args=()) +def _mslk_quantize_nvfp4_stacked_custom_op( + m_sizes: torch.Tensor, + x: torch.Tensor, + global_scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + assert _mslk_available, ( + "mslk is required for nvfp4_quantize_stacked. " + "Install from https://github.com/meta-pytorch/MSLK" + ) + from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked + + xq, scale = nvfp4_quantize_stacked(m_sizes, x, global_scale) + return xq, scale + + +@_mslk_quantize_nvfp4_stacked_custom_op.register_fake +def _(m_sizes, x, global_scale): + M, K = x.shape[0], x.shape[1] + num_segments = m_sizes.shape[0] + # Upper-bound on padded total rows (each segment can add at most 127 padding rows) + padded_total_M_ub = M + num_segments * 127 + num_scales_per_row = K // 16 + n_col_blocks = triton.cdiv(num_scales_per_row, 4) + padded_cols = n_col_blocks * 4 + xq = x.new_empty(M, K // 2, dtype=torch.float4_e2m1fn_x2) + scale = x.new_empty(padded_total_M_ub, padded_cols, dtype=torch.float8_e4m3fn) + return xq, scale diff --git a/torchao/prototype/mx_formats/nvfp4_tensor.py b/torchao/prototype/mx_formats/nvfp4_tensor.py index cf8ac3865f..034ea377c0 100644 --- a/torchao/prototype/mx_formats/nvfp4_tensor.py +++ b/torchao/prototype/mx_formats/nvfp4_tensor.py @@ -15,7 +15,9 @@ from torchao.prototype.mx_formats.kernels import ( f4_unpacked_to_f32, f32_to_f4_unpacked, + mslk_calculate_group_max, mslk_quantize_nvfp4, + mslk_quantize_nvfp4_stacked, pack_uint4, unpack_uint4, ) @@ -680,9 +682,6 @@ def nvfp4_addmm(func, types, args, kwargs): @implements([aten._grouped_mm.default]) def nvfp4_grouped_mm(func, types, args, kwargs): - # TODO(before land): wrap MSLK kernel with a custom op, guard for MSLK - # not installed, etc - from mslk.quantize.triton.fp4_quantize import nvfp4_quantize_stacked from torch.nn.functional import ScalingType, SwizzleType, scaled_grouped_mm mat_a, mat_b = args[0], args[1] @@ -694,32 +693,32 @@ def nvfp4_grouped_mm(func, types, args, kwargs): # original NVFP4Tensor so we can extract qdata/scale in the right layout. assert isinstance(mat_b, NVFP4Tensor) is_transposed = mat_b.qdata.stride(-2) < mat_b.qdata.stride(-1) - if is_transposed: - mat_b = mat_b.transpose(-2, -1) + assert is_transposed, "unsupported" E = offs.shape[0] m_sizes = torch.diff(offs, prepend=offs.new_zeros(1)).to(torch.int64) - # For now, quantize activation with per-expert global scales of 1.0 - # TODO(future PR): compute real per-token-group global scales from activation amax, - # will need some slow eager code or a triton kernel - a_global_scale = torch.ones(E, device=mat_a.device, dtype=torch.float32) - mat_a_qdata, mat_a_scale = nvfp4_quantize_stacked(m_sizes, mat_a, a_global_scale) + # mslk_calculate_group_max returns encoding scale (448 * FP4_MAX / amax) + # mslk_quantize_nvfp4_stacked needs encoding scale, scaled_grouped_mm + # needs decoding scale (1 / encoding_scale) + a_global_scale_enc = mslk_calculate_group_max(mat_a, m_sizes) + mat_a_qdata, mat_a_scale = mslk_quantize_nvfp4_stacked( + m_sizes, mat_a, a_global_scale_enc + ) + a_global_scale_dec = 1.0 / a_global_scale_enc - # Weight: extract qdata and already-swizzled scales - mat_b_qdata = mat_b.qdata.view(torch.float4_e2m1fn_x2) # Flatten 3D scale (E, padded_N, padded_K//16) -> 2D (E, padded_N * padded_K//16) # as required by _scaled_grouped_mm for 2D-3D grouped GEMM - mat_b_scale = mat_b.scale.view(torch.float8_e4m3fn).flatten(1) + mat_b_t_scale = mat_b.scale.transpose(-2, -1).flatten(1) + # [E, 1, 1] -> E b_global_scale = mat_b.per_tensor_scale.view(E) - # _scaled_grouped_mm expects wq.transpose(-2, -1) return scaled_grouped_mm( mat_a_qdata, - mat_b_qdata.transpose(-2, -1), - scale_a=[mat_a_scale, a_global_scale], + mat_b.qdata.view(torch.float4_e2m1fn_x2), + scale_a=[mat_a_scale, a_global_scale_dec], scale_recipe_a=[ScalingType.BlockWise1x16, ScalingType.TensorWise], - scale_b=[mat_b_scale, b_global_scale], + scale_b=[mat_b_t_scale, b_global_scale], scale_recipe_b=[ScalingType.BlockWise1x16, ScalingType.TensorWise], swizzle_a=SwizzleType.SWIZZLE_32_4_4, swizzle_b=SwizzleType.SWIZZLE_32_4_4,