nemo_automodel.components.models.qwen3_5.layers

Fused elementwise kernels for the dense Qwen3.5 decoder.

The reference Qwen3.5 decoder runs its norms, residual adds, SwiGLU activation and GatedDeltaNet decay gate as chains of eager elementwise kernels, each a full pass over the activation: about seven forward launches per zero-centered RMSNorm and nine per gated RMSNorm, with more in backward. On CUDA these modules compile the same math into fused Triton kernels for forward and backward. CPU and DTensor inputs keep the eager chain, so results match the reference modules.

Module Contents

Classes

Name	Description
`Qwen3_5FusedMLP`	Dense SwiGLU MLP whose `silu(gate) * up` runs as one fused kernel on CUDA.
`Qwen3_5FusedRMSNorm`	Zero-centered Qwen3.5 RMSNorm computed by one fused kernel on CUDA.
`Qwen3_5FusedRMSNormGated`	GatedDeltaNet output norm with SiLU gate computed by one fused kernel on CUDA.
`Qwen3_5FusedSSMGate`	fp32 GatedDeltaNet decay gate computed by one fused kernel on CUDA.

Functions

Name	Description
`_add_zero_centered_rms_norm`	Add `residual` to `x` and normalize the sum, as an add followed by `Qwen3NextRMSNorm`.
`_fusable`	-
`_rms_norm_gated_silu`	Compute `weight * norm(hidden_states) * silu(gate)` as `Qwen3_5MoeRMSNormGated`.
`_silu_mul`	Compute the SwiGLU product `silu(gate) * up` as the eager `MLP` does.
`_ssm_gate`	Compute the GatedDeltaNet log-decay `-exp(A_log) * softplus(a + dt_bias)` as `SSMGate`.
`_zero_centered_rms_norm`	Compute `norm(x) * (1 + weight)` in fp32, exactly as `Qwen3NextRMSNorm`.

Data

_add_zero_centered_rms_norm_compiled

_rms_norm_gated_silu_compiled

_silu_mul_compiled

_ssm_gate_compiled

_zero_centered_rms_norm_compiled

API

class nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedMLP()

Bases: MLP

Dense SwiGLU MLP whose silu(gate) * up runs as one fused kernel on CUDA.

nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedMLP.forward(
    x: torch.Tensor
) -> torch.Tensor

Project, apply the fused SwiGLU product, and project back.

Parameters:

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

Returns: torch.Tensor

Tensor of shape […, hidden].

class nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedRMSNorm()

Bases: Qwen3NextRMSNorm

Zero-centered Qwen3.5 RMSNorm computed by one fused kernel on CUDA.

nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedRMSNorm.forward(
    x: torch.Tensor,
    residual: torch.Tensor | None = None
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Normalize the last axis in fp32 and scale by 1 + weight.

Parameters:

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

residual

torch.Tensor | NoneDefaults to None

Optional tensor of x’s shape and dtype. When given, the residual add is fused into the norm: x + residual is normalized and returned alongside its normalization.

Returns: torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Tensor of shape […, hidden] in x’s dtype, or, with residual, the

class nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedRMSNormGated()

Bases: Qwen3_5MoeRMSNormGated

GatedDeltaNet output norm with SiLU gate computed by one fused kernel on CUDA.

nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedRMSNormGated.forward(
    hidden_states: torch.Tensor,
    gate: torch.Tensor
) -> torch.Tensor

Normalize the last axis, apply the gain, then multiply by silu(gate).

Parameters:

hidden_states

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

gate

torch.Tensor

Tensor of shape […, hidden] holding the output-gate logits.

Returns: torch.Tensor

Tensor of shape […, hidden] in hidden_states’ dtype.

class nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedSSMGate()

Bases: SSMGate

fp32 GatedDeltaNet decay gate computed by one fused kernel on CUDA.

nemo_automodel.components.models.qwen3_5.layers.Qwen3_5FusedSSMGate.forward(
    a: torch.Tensor
) -> torch.Tensor

Compute -exp(A_log) * softplus(a + dt_bias) in fp32.

Parameters:

torch.Tensor

Tensor of shape […, value_heads] from the decay projection.

Returns: torch.Tensor

fp32 tensor of shape […, value_heads].

nemo_automodel.components.models.qwen3_5.layers._add_zero_centered_rms_norm(
    x: torch.Tensor,
    residual: torch.Tensor,
    weight: torch.Tensor,
    eps: float
) -> tuple[torch.Tensor, torch.Tensor]

Add residual to x and normalize the sum, as an add followed by Qwen3NextRMSNorm.

Parameters:

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

residual

torch.Tensor

Tensor of the same shape and dtype as x.

weight

torch.Tensor

Tensor of shape [hidden] holding the zero-centered gain.

eps

float

Epsilon added to the mean square.

Returns: torch.Tensor

The sum x + residual and its normalization, both of shape […, hidden]

nemo_automodel.components.models.qwen3_5.layers._fusable(
    tensors: torch.Tensor = ()
) -> bool

nemo_automodel.components.models.qwen3_5.layers._rms_norm_gated_silu(
    hidden_states: torch.Tensor,
    gate: torch.Tensor,
    weight: torch.Tensor,
    eps: float
) -> torch.Tensor

Compute weight * norm(hidden_states) * silu(gate) as Qwen3_5MoeRMSNormGated.

The normalized activations are rounded to the input dtype before the weight multiply, matching the reference module.

Parameters:

hidden_states

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

gate

torch.Tensor

Tensor of shape […, hidden] holding the output-gate logits.

weight

torch.Tensor

Tensor of shape [hidden] holding the multiplicative gain.

eps

float

Epsilon added to the mean square.

Returns: torch.Tensor

Tensor of shape […, hidden] in hidden_states’ dtype.

nemo_automodel.components.models.qwen3_5.layers._silu_mul(
    gate: torch.Tensor,
    up: torch.Tensor
) -> torch.Tensor

Compute the SwiGLU product silu(gate) * up as the eager MLP does.

Parameters:

gate

torch.Tensor

Tensor of shape […, intermediate] from the gate projection.

torch.Tensor

Tensor of the same shape and dtype from the up projection.

Returns: torch.Tensor

Tensor of shape […, intermediate] in gate’s dtype.

nemo_automodel.components.models.qwen3_5.layers._ssm_gate(
    a: torch.Tensor,
    a_log: torch.Tensor,
    dt_bias: torch.Tensor
) -> torch.Tensor

Compute the GatedDeltaNet log-decay -exp(A_log) * softplus(a + dt_bias) as SSMGate.

Parameters:

torch.Tensor

Tensor of shape […, value_heads] from the decay projection.

a_log

torch.Tensor

Tensor of shape [value_heads] holding A_log.

dt_bias

torch.Tensor

Tensor of shape [value_heads].

Returns: torch.Tensor

fp32 tensor of shape […, value_heads].

nemo_automodel.components.models.qwen3_5.layers._zero_centered_rms_norm(
    x: torch.Tensor,
    weight: torch.Tensor,
    eps: float
) -> torch.Tensor

Compute norm(x) * (1 + weight) in fp32, exactly as Qwen3NextRMSNorm.

Parameters:

torch.Tensor

Tensor of shape […, hidden], with arbitrary leading dimensions.

weight

torch.Tensor

Tensor of shape [hidden] holding the zero-centered gain.

eps

float

Epsilon added to the mean square.

Returns: torch.Tensor

Tensor of shape […, hidden] in x’s dtype.

nemo_automodel.components.models.qwen3_5.layers._add_zero_centered_rms_norm_compiled = torch.compile(_add_zero_centered_rms_norm, dynamic=True, options={'emulate_preci...

nemo_automodel.components.models.qwen3_5.layers._rms_norm_gated_silu_compiled = torch.compile(_rms_norm_gated_silu, dynamic=True, options={'emulate_precision_ca...

nemo_automodel.components.models.qwen3_5.layers._silu_mul_compiled = torch.compile(_silu_mul, dynamic=True, options={'emulate_precision_casts': True}...

nemo_automodel.components.models.qwen3_5.layers._ssm_gate_compiled = torch.compile(_ssm_gate, dynamic=True)

nemo_automodel.components.models.qwen3_5.layers._zero_centered_rms_norm_compiled = torch.compile(_zero_centered_rms_norm, dynamic=True)