Merge pull request #2987 from hx89:hx/nvfp4_no_rht

Google-ML-Automation · Google-ML-Automation · commit 9fe6cdfc5106 · 2026-02-05T14:38:02.000-08:00
PiperOrigin-RevId: 866117394
diff --git a/src/MaxText/configs/types.py b/src/MaxText/configs/types.py
@@ -84,6 +84,7 @@ class QuantizationType(str, Enum):
   TE_FP8_CS = "te_fp8_currentscaling"
   TE_MXFP8 = "te_mxfp8"
   TE_NVFP4 = "te_nvfp4"
+  TE_NVFP4_NO_RHT = "te_nvfp4_no_rht"
 
 
 class KvQuantAxis(str, Enum):
diff --git a/src/MaxText/layers/quantizations.py b/src/MaxText/layers/quantizations.py
@@ -749,6 +749,7 @@ def _get_recipe(recipe_name: str):
         "te_fp8_currentscaling": recipe.Float8CurrentScaling,
         "te_mxfp8": recipe.MXFP8BlockScaling,
         "te_nvfp4": recipe.NVFP4BlockScaling,  # pytype: disable=module-attr
+        "te_nvfp4_no_rht": functools.partial(recipe.NVFP4BlockScaling, disable_rht=True),  # pytype: disable=module-attr
     }
     if recipe_name not in RECIPES:
       raise ValueError(f"Invalid TransformerEngine recipe: {recipe_name}")

Original file line number	Diff line number	Diff line change
`@@ -749,6 +749,7 @@ def _get_recipe(recipe_name: str):`
`749`	`749`	`"te_fp8_currentscaling": recipe.Float8CurrentScaling,`
`750`	`750`	`"te_mxfp8": recipe.MXFP8BlockScaling,`
`751`	`751`	`"te_nvfp4": recipe.NVFP4BlockScaling, # pytype: disable=module-attr`
	`752`	`+ "te_nvfp4_no_rht": functools.partial(recipe.NVFP4BlockScaling, disable_rht=True), # pytype: disable=module-attr`
`752`	`753`	`}`
`753`	`754`	`if recipe_name not in RECIPES:`
`754`	`755`	`raise ValueError(f"Invalid TransformerEngine recipe: {recipe_name}")`