[JAX] Clamped Swiglu Integration

[vthumbe1503]

Description

Jax integration for clamped swiglu

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[jberchtold-nvidia]

/te-ci L2 jax

[jberchtold-nvidia]

Overall looks good pending CI, left some comments and questions, mostly minor things. Thanks for updating the Flax modules as well, great work on this PR!

[jberchtold-nvidia]

/te-ci L2 jax

[phu0ngng]

Hi,
I left some minor comments. LGTM otherwise. Thanks!

[phu0ngng]

Hi, let's move this class to cpp_extensions.py and expose it as a part of the tex.

[jberchtold-nvidia]

He had it there originally and I told him to move it here. Since this param will be user facing, they'll need to create it themselves and import from cpp_extensions. My understanding is that we wanted to encapsulate and hide those internals from users of the higher-level VJPs or Flax layers, so I suggested moving it here.

But do you think it's more important to prevent the dependency from cpp_extensions on this type in the higher-level VJP even tho users will need to use cpp_extensions types directly? Lmk what you think, thanks!

[phu0ngng]

The reasons I suggested moving it to the tex are that:

• Users don't really use this activation VJP, so they will still need to import from this file even if they don't use other stuff.
• We have other Params/Configs that are placed in the tex (i.e., AmaxScope is in cpp_extensions/quantization.py, CollectiveOp is in cpp_extensions/gemm.py).

At the same time, I agree that exposing it here may be a cleaner approach. If we decide to keep it here, perhaps we should start a follow-up work on which functions/symbols we want to export from each VJP file so that users can do import transformer_engine.jax as te_jax.

[jberchtold-nvidia]

Sounds good, I'm okay with reverting back to having it in cpp_extensions (tex) for now then to be consistent with other structures until we decide on a way to organize these structures. Apologies for the churn, Varun

[vthumbe1503]

No issues @jberchtold-nvidia

[phu0ngng]

Hi, should we make a default ActParam instead of None here?

[vthumbe1503]

I am actually following generally safer practice of not providing default values other than None for user defined data type in the function definition. Although ActivationParams Object today is immutable, it might be dangerous in future in case it isnt anymore. Since the default values created are just instantiated and shared during import time. And accidental change of the param in one function can impact the other which can be harder to debug. Although this can be solved with default factory and such a things to happen is very unrealistic, so let me know what you think.

[phu0ngng]

Hi,
Since we will lower this struct for all activations (even if we don't use it), should we have a more general name for it instead?
Something like:

struct ActParams {
  float swiglu_limit;
  float swiglu_alpha;
 };

[jberchtold-nvidia]

Varun made a different version here with a generic outer struct and a nested inner struct for ClampedSwiglu, which I believe is what you have in mind. I haven't reviewed this other version yet, planning to do it today
vthumbe1503#2

[vthumbe1503]

Have moved the generic approach to this PR itself

[Copilot]

Pull Request Overview

This PR integrates Clamped SwiGLU activation into the JAX backend, adding support for the GPT-OSS variant of SwiGLU that includes parameter-based clamping and scaling functionality.

• Adds ClampedSwiGLU activation implementation across PyTorch, JAX, and C++ backends
• Implements parameter passing infrastructure for activation functions requiring additional parameters
• Extends existing activation framework to support parameterized activations

Reviewed Changes

Copilot reviewed 23 out of 23 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
transformer_engine/pytorch/ops/basic/activation.py	Adds ClampedSwiGLU class with limit and alpha parameters
transformer_engine/pytorch/ops/basic/init.py	Exports ClampedSwiGLU activation
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Exposes clamped_swiglu functions to Python
transformer_engine/pytorch/csrc/extensions/activation.cpp	Refactors activation helpers to support parameterized functions
transformer_engine/pytorch/csrc/extensions.h	Declares clamped activation function signatures
transformer_engine/jax/layernorm_mlp.py	Adds activation_params parameter to layernorm_mlp
transformer_engine/jax/flax/transformer.py	Adds mlp_activation_params to TransformerLayer
transformer_engine/jax/flax/module.py	Extends LayerNormMLP with activation_params support
transformer_engine/jax/csrc/extensions/pybind.cpp	Adds CLAMPED_SWIGLU enum value
transformer_engine/jax/csrc/extensions/activation.cpp	Implements clamped activation FFI handlers
transformer_engine/jax/csrc/extensions.h	Defines ActivationConfig structs
transformer_engine/jax/cpp_extensions/activation.py	Implements ActivationParams framework and clamped functions
transformer_engine/jax/activation.py	Updates activation functions to accept parameters
transformer_engine/common/util/vectorized_pointwise.h	Updates kernels for parameterized activations
transformer_engine/common/util/math.h	Adds clamped SiLU mathematical functions
transformer_engine/common/util/cast_gated_kernels.cuh	Updates gated kernels for parameter support
transformer_engine/common/include/transformer_engine/activation.h	Adds clamped activation API declarations
transformer_engine/common/activation/swiglu.cu	Implements nvte_clamped_swiglu functions
transformer_engine/common/activation/relu.cu	Updates function calls with parameters
transformer_engine/common/activation/gelu.cu	Updates function calls with parameters
transformer_engine/common/activation/activation_template.h	Updates templates for parameter passing
tests/pytorch/test_fusible_ops.py	Adds ClampedSwiGLU test cases
tests/jax/test_custom_call_compute.py	Updates tests for parameterized activations

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[vthumbe1503]

/te-ci L2 jax

[vthumbe1503]

/te-ci L2 jax