Add config option to set attrs in TE quantization recipe

[timmoon10]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

This PR adds a config option so that users can set arbitrary attrs in TE quantization recipes, reducing the need to create new config options for every experimental recipe variant that may come along in the future.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[timmoon10]

This code path uses TE support for UE5M3 scales, which is pending in NVIDIA/TransformerEngine#3325. If you set this envvar, then you're almost certainly an expert user who's explicitly enabling UE5M3 support.

Also, all of TEFusedMLPWithGroupedLinear should be considered as a temporary hack. The right design is to use TEFusedMLP once TE supports dense MLP fusions. CC @sraman-rgb

[fanshiqing]

/ok to test fb906dc

[fanshiqing]

/claude review

[claude]

Light review — the recipe_attrs plumbing itself looks correct and consistent across all four call sites (get_fp8_recipe, get_fp4_recipe, _get_fp8_model_init_for_quant_recipe, _get_fp8_autocast_for_quant_recipe), each applying the setattr loop after every branch has assigned a recipe. Three things worth addressing:

1. shared_experts.py: recipe validation for the fused grouped-SwiGLU path was dropped

_get_fused_grouped_swiglu_recipe previously required fp4_recipe=nvfp4 or fp8_recipe=mxfp8 and raised otherwise. After this change any FP8/FP4 recipe is accepted:

if self.config.fp8:
    self._fused_grouped_swiglu_recipe = get_fp8_recipe(self.config)
elif self.config.fp4:
    self._fused_grouped_swiglu_recipe = get_fp4_recipe(self.config)

With e.g. fp8_recipe=delayed or tensorwise this now returns a TEDelayedScaling/Float8CurrentScaling recipe, the ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 fusion never fires, and the user silently gets the slow path instead of the previous explicit error. Routing through get_fp*_recipe to pick up recipe_attrs is the right call — but consider keeping a recipe check, ideally in _validate_fused_grouped_swiglu so it fails at construction time rather than on first forward:

fp4_recipe = getattr(self.config.fp4_recipe, "value", self.config.fp4_recipe)
fp8_recipe = getattr(self.config.fp8_recipe, "value", self.config.fp8_recipe)
if not (
    (self.config.fp8 and fp8_recipe in ("mxfp8", "custom"))
    or (self.config.fp4 and fp4_recipe in ("nvfp4", "custom"))
):
    raise ValueError(
        f"{self.__class__.__name__} requires fp8_recipe=mxfp8 or fp4_recipe=nvfp4 "
        f"(or custom), but got fp8={self.config.fp8}, fp8_recipe={self.config.fp8_recipe}, "
        f"fp4={self.config.fp4}, fp4_recipe={self.config.fp4_recipe}."
    )

2. No test coverage for the new config options

fp8_recipe_attrs / fp4_recipe_attrs / TEQuantizationRecipe.recipe_attrs are a new user-facing feature with no tests, and the new FP4_RECIPE=nvfp4_ue5m3 branch in TEFusedMLPWithGroupedLinear.forward is untested as well. tests/unit_tests/transformer/test_te_fused_mlp_with_grouped_linear_spec.py already parametrizes over FP4_RECIPE and monkeypatches the recipe constructors, so adding a case there is cheap. A small unit test asserting get_fp8_recipe(config) reflects config.fp8_recipe_attrs (and the FP4 equivalent) would also guard against the setattr loop being lost in a future refactor of these branchy functions.

3. Pre-existing typo adjacent to your change (megatron/core/extensions/transformer_engine.py:366)

The FP4 branch of _get_fp8_autocast_for_quant_recipe interpolates the FP8 field in its error message:

raise ValueError(f"Unhandled fp4 recipe: {qrecipe.fp8_quantization_recipe}")

Should be qrecipe.fp4_quantization_recipe (as it correctly is at line 297). Not introduced by this PR, but you are adding lines five below it — trivial drive-by fix if you are willing.