mstar.model.components.decoder_layer

Pre-norm transformer decoder layers.

DecoderLayer: the standard pre-norm block (norm → attn → residual, norm → mlp → residual). Composes any nn.Module for the attention, MLP, and norms — so models pick the variants they want and pass them in.

GatedDecoderLayer: variant for adaRMS conditioning (pi05 action expert). The norms are AdaRMSNorm and return (normed, gate); the residual is x + gate * y instead of x + y.

Classes

DecoderLayer(self_attn, mlp, ...)	Standard pre-norm transformer decoder layer.
GatedDecoderLayer(self_attn, mlp, ...)	Pre-norm decoder layer with adaRMS gated residuals.

class mstar.model.components.decoder_layer.DecoderLayer(self_attn, mlp, input_layernorm, post_attention_layernorm)

Bases: Module

Standard pre-norm transformer decoder layer.

Computes:

residual = x x = input_layernorm(x); x = self_attn(x); x = residual + x residual = x x = post_attention_layernorm(x); x = mlp(x); x = residual + x

Parameters:

• self_attn (nn.Module) – attention module taking (hidden_states, label=, layer_idx=) and returning a tensor.

• mlp (nn.Module) – feedforward module taking and returning a tensor.

• input_layernorm (nn.Module) – pre-attention norm.

• post_attention_layernorm (nn.Module) – pre-FFN norm.

forward(hidden_states)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.decoder_layer.GatedDecoderLayer(self_attn, mlp, input_layernorm, post_attention_layernorm)

Bases: Module

Pre-norm decoder layer with adaRMS gated residuals.

The norms must be AdaRMSNorm-shaped: forward(x, cond) returns (normed, gate). The residual becomes x + gate * y.

Used by pi05’s action expert; adarms_cond is the shared condition vector consumed by both norms.

Parameters:

• self_attn (nn.Module)

• mlp (nn.Module)

• input_layernorm (nn.Module)

• post_attention_layernorm (nn.Module)

forward(hidden_states, adarms_cond)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

• hidden_states (Tensor)

• adarms_cond (Tensor)

Return type:

Tensor