mstar.model.components

Reusable transformer building blocks.

Components in this package are model-agnostic and meant to be shared across model implementations. Anything model-specific (vision encoders, audio codecs, multimodal preprocessing, etc.) stays in the per-model components/ directory.

The TP-aware versions of these blocks (parallel linears, etc.) will land here too in a follow-up; the current shapes intentionally leave room for that.

class mstar.model.components.AdaRMSNorm(hidden_size, cond_dim, eps=1e-6)

Bases: Module

RMSNorm with adaRMS conditioning.

A per-norm nn.Linear(cond_dim, hidden_size*3) maps a shared condition vector to (scale, shift, gate). The normalization is rmsnorm(x) * (1 + scale) + shift and the gate is returned for the enclosing decoder layer to apply at the residual.

The dense.weight and dense.bias are zero-initialized so the norm starts as the identity (matches HF Gemma / lerobot openpi).

Parameters:

• hidden_size (int)

• cond_dim (int)

• eps (float)

forward(x, 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:

• x (Tensor)

• cond (Tensor)

Return type:

tuple[Tensor, Tensor]

class mstar.model.components.Attention(*, hidden_size, num_heads, num_kv_heads, head_dim, qkv_bias=False, o_bias=False, qk_norm=False, rms_norm_eps=1e-6, rope_theta=10_000.0, rope_scale=1.0, rope_low_freq_factor=1.0, rope_high_freq_factor=1.0, rope_old_context_len=8192, input_hidden_size=None, attn_key='attn', kv_key='kv', pos_key='rope')

Bases: Module

Parameters:

• hidden_size (int)

• num_heads (int)

• num_kv_heads (int)

• head_dim (int)

• qkv_bias (bool)

• o_bias (bool)

• qk_norm (bool)

• rms_norm_eps (float)

• rope_theta (float)

• rope_scale (float)

• rope_low_freq_factor (float)

• rope_high_freq_factor (float)

• rope_old_context_len (int)

• input_hidden_size (int | None)

• attn_key (str)

• kv_key (str)

• pos_key (str | None)

bind_resources(resources)

Resolve the resources this layer calls. See NodeSubmodule.bind_node_resources.

.get: a layer may be bound on a node that owns only some of them.

Parameters:

resources (dict)

Return type:

None

consolidate_qkv_weight()

Fuse q_proj/k_proj/v_proj weights into a single qkv_proj_weight buffer and null out the originals. Idempotent.

Return type:

None

forward(hidden_states)

The label and layer index are cursors on the resources, set by the caller running the layer stack (attend.bind_step once, then attend.set_layer_idx per layer) rather than passed in per call.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

class mstar.model.components.CrossAttention(*, hidden_size, num_heads, head_dim, q_bias=True, k_bias=False, v_bias=True, o_bias=True, source='default', cross_key=None, context_kv_key=None)

Bases: Module

Multi-head cross-attention over an engine-managed encoder-context KV.

For encoder-decoder models (Whisper, etc.): the decoder attends to a fixed encoder context whose K/V are computed once at prefill and written into a KV resource of their own, under the context_label the model’s CrossAttentionConfig names (see issue #160). Q is projected per step; K/V projections are exposed via compute_kv so the submodule can write them into that cache at encode time.

Q/K/V/O are separate nn.Linear matching the HF layout. Subclasses override projection details (bias flags, a compute_kv that reshapes for a model-specific pool layout) as needed; the default matches Whisper (q/v/o biased, k unbiased).

TODO(#160): the projections are plain nn.Linear — this module is not yet TP/SP-compatible (no column/row-parallel splits over heads). A tensor-parallel cross-attention variant is needed to serve the decoder under TP alongside the self-attention path.

Parameters:

• hidden_size (int)

• num_heads (int)

• head_dim (int)

• q_bias (bool)

• k_bias (bool)

• v_bias (bool)

• o_bias (bool)

• source (str)

• cross_key (str | None)

• context_kv_key (str | None)

bind_resources(resources)

Resolve this source’s cross-attention resource and the cache holding its context. See NodeSubmodule.bind_node_resources.

Parameters:

resources (dict)

Return type:

None

bind_step(label)

Parameters:

label (str)

Return type:

None

compute_kv(encoder_states)

Project the encoder context to K/V for the cross-attention pool.

(enc_len, hidden) -> (k, v), each (enc_len, num_heads, head_dim). Override to reshape for a model-specific pool layout.

Parameters:

encoder_states (Tensor)

Return type:

tuple[Tensor, Tensor]

forward(hidden_states)

Label and layer index come off the resources’ cursors; see Attention.forward.

Parameters:

hidden_states (Tensor)

Return type:

Tensor

set_layer_idx(layer_idx)

Parameters:

layer_idx (int)

Return type:

None

class mstar.model.components.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.FusedColumnLinear(input_size, shard_sizes, bias=False, dtype=None)

Bases: Module

Linear whose output is the concatenation of several shards along dim 0, fused into a single weight (and optional bias).

Parameters:

• input_size (int) – input feature dim.

• shard_sizes (dict[str | int, int]) – maps a shard id to its output size. Shards are laid out along dim 0 in iteration order, so the order of this dict defines the layout (and the split used at forward time).

• bias (bool) – whether to include a (fused) bias.

• dtype (torch.dtype | None)

The fused weight / bias carry a weight_loader(param, tensor, shard_id) method that copies one checkpoint shard into its slice of the fused parameter, dispatched by shard_id (a key of shard_sizes). Wire the per-shard checkpoint keys to it with the loader’s stacked-param rules.

forward(x)

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:

x (Tensor)

Return type:

Tensor

weight_loader(param, loaded_weight, loaded_shard_id=None)

Parameters:

• param (Parameter)

• loaded_weight (Tensor)

• loaded_shard_id (str | int | None)

Return type:

None

class mstar.model.components.FusedGatedMLP(hidden_size, intermediate_size, activation='silu', bias=False)

Bases: Module

SwiGLU-style gated MLP with the gate + up projections fused into a single FusedColumnLinear (one GEMM instead of two): down(act(gate) * up). Unlike GatedMLP (separate Linears + post-load consolidate_gate_up_weight), this fuses from construction and loads the separate gate_proj / up_proj checkpoint tensors straight into the fused parameter via the loader’s stacked-param rules (gate is shard 0, up is shard 1). Use for models that fuse but don’t need TP. :param hidden_size: input/output feature dim. :param intermediate_size: gate/up output and down input feature dim. :param activation: activation applied to the gate path. Either an HF

string (silu / gelu / gelu_tanh) or a callable.

Parameters:

• bias (bool) – whether the linears have a bias term.

• hidden_size (int)

• intermediate_size (int)

• activation (str | Callable)

forward(x)

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:

x (Tensor)

Return type:

Tensor

class mstar.model.components.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

class mstar.model.components.GatedMLP(hidden_size, intermediate_size, activation='silu', bias=False)

Bases: Module

SwiGLU-style gated MLP: down(act(gate(x)) * up(x)).

Parameters:

• hidden_size (int) – input/output feature dim.

• intermediate_size (int) – gate/up output and down input feature dim.

• activation (str | Callable) – activation applied to the gate path. Either an HF string (silu / gelu / gelu_tanh) or a callable.

• bias (bool) – whether the linears have a bias term.

consolidate_gate_up_weight()

Fuse gate_proj and up_proj weights into a single gate_up_proj_weight buffer and null out the originals. Idempotent; safe to call multiple times.

Return type:

None

forward(x)

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:

x (Tensor)

Return type:

Tensor

class mstar.model.components.MLP(input_size, intermediate_size, output_size=None, activation='silu', bias=True)

Bases: Module

Plain two-layer MLP: out(act(in(x))).

Used for small projection MLPs that aren’t gated (Talker resize projection, bagel timestep embedder MLP, etc.). For SwiGLU-style transformer FFNs, use GatedMLP.

Parameters:

• input_size (int) – input feature dim.

• intermediate_size (int) – hidden feature dim.

• output_size (int | None) – output feature dim. Defaults to input_size if None.

• activation (str | Callable) – activation between the two linears.

• bias (bool) – whether the linears have a bias term.

forward(x)

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:

x (Tensor)

Return type:

Tensor

class mstar.model.components.ParallelSparseMoeBlock(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, norm_topk_prob=True, router=None, comm_group=None)

Bases: Module

TP-aware Top-K sparse MoE.

When tp_size == 1, the forward is identical to SparseMoeBlock (full fused kernel, no communication). When tp_size > 1, expert weights are sharded along the intermediate dimension and an all-reduce is inserted between the down-projection GEMM and the top-k sum-reduce.

Parameters:

• hidden_size (int)

• num_experts (int)

• num_experts_per_tok (int)

• moe_intermediate_size (int)

• norm_topk_prob (bool)

• router (nn.Module | None)

• comm_group (CommGroup | None)

forward(hidden_states, router_states=None, *, return_router_states=False)

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)

• router_states (Tensor | None)

• return_router_states (bool)

class mstar.model.components.ParallelSparseMoeBlockWithSharedExpert(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, shared_expert, norm_topk_prob=False, router=None, comm_group=None)

Bases: Module

TP-aware Top-K sparse MoE with a shared expert + sigmoid gating.

The shared expert should be a ParallelGatedMLP constructed with the same comm_group so its all-reduce is handled internally.

Parameters:

• hidden_size (int)

• num_experts (int)

• num_experts_per_tok (int)

• moe_intermediate_size (int)

• shared_expert (nn.Module)

• norm_topk_prob (bool)

• router (nn.Module | None)

• comm_group (CommGroup | None)

forward(hidden_states, router_states=None, *, return_router_states=False)

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)

• router_states (Tensor | None)

• return_router_states (bool)

class mstar.model.components.RMSNorm(hidden_size, eps=1e-6, gemma_mode=False)

Bases: Module

RMSNorm with optional Gemma-style (1 + weight) scaling.

Parameters:

• hidden_size (int) – feature dimension to normalize over.

• eps (float) – variance epsilon.

• gemma_mode (bool) – if True, use (1 + weight) and a fp32 manual implementation (matches HF Gemma exactly; the loaded checkpoint weight is centered around zero, not one). If False, use weight and dispatch to FlashInfer’s fused kernel.

extra_repr()

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

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.SparseMoeBlock(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, norm_topk_prob=True, router=None)

Bases: Module

Top-K sparse MoE with fused expert weights, no shared expert.

Expert weights match the HF fused checkpoint layout:

• experts.gate_up_proj: (num_experts, 2 * moe_intermediate_size, hidden_size)

• experts.down_proj: (num_experts, hidden_size, moe_intermediate_size)

Parameters:

• hidden_size (int)

• num_experts (int)

• num_experts_per_tok (int)

• moe_intermediate_size (int)

• norm_topk_prob (bool)

• router (nn.Module | None)

forward(hidden_states, router_states=None, *, return_router_states=False)

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)

• router_states (Tensor | None)

• return_router_states (bool)

class mstar.model.components.SparseMoeBlockWithSharedExpert(hidden_size, num_experts, num_experts_per_tok, moe_intermediate_size, shared_expert, norm_topk_prob=False, router=None)

Bases: Module

Top-K sparse MoE with a shared expert + sigmoid gating.

Final output is:

out = routed(x) + sigmoid(shared_gate(x)) * shared_expert(x)

The shared expert is supplied by the caller (any nn.Module matching the hidden_size → hidden_size interface). The routed path uses the same fused checkpoint layout as SparseMoeBlock.

Parameters:

• hidden_size (int)

• num_experts (int)

• num_experts_per_tok (int)

• moe_intermediate_size (int)

• shared_expert (nn.Module)

• norm_topk_prob (bool)

• router (nn.Module | None)

forward(hidden_states, router_states=None, *, return_router_states=False)

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)

• router_states (Tensor | None)

• return_router_states (bool)

class mstar.model.components.TopKRouter(hidden_size, num_experts, num_experts_per_tok, norm_topk_prob=True)

Bases: Module

Softmax top-k router. This is the default router for all MoE blocks.

See the module docstring for the router contract. This router is stateless. It ignores router_states. It returns None for router_states_next.

Parameters:

• hidden_size (int) – input hidden dimension.

• num_experts (int) – total number of routed experts.

• num_experts_per_tok (int) – number of experts each token is dispatched to (top-k).

• norm_topk_prob (bool) – if True, renormalize the top-k probabilities so they sum to 1.

forward(hidden_states, router_states=None)

Parameters:

• hidden_states (Tensor) – (tokens, hidden_size) router input.

• router_states (Tensor | None) – unused (stateless router); accepted for interface compatibility with stateful routers.

Returns:

(tokens, top_k) top-k probabilities

(optionally renormalized).

selected_experts: (tokens, top_k) int64 indices. router_states_next: always None (stateless).

Return type:

routing_weights

mstar.model.components.dispatch_experts_fused(hidden_states, gate_up_proj, down_proj, num_experts, selected_experts, routing_weights)

Naive per-expert dispatch using the fused HF checkpoint layout.

Used as a fallback when the Triton fused-MoE kernel isn’t available. Loops over the experts that received any tokens and runs SwiGLU per expert.

Parameters:

• hidden_states (Tensor) – (tokens, hidden_size).

• gate_up_proj (Tensor) – (num_experts, 2 * moe_intermediate_size, hidden_size).

• down_proj (Tensor) – (num_experts, hidden_size, moe_intermediate_size).

• selected_experts (Tensor) – (tokens, top_k) int64.

• routing_weights (Tensor) – (tokens, top_k) float.

• num_experts (int)

Return type:

Tensor

Modules

attention	Multi-head attention with GQA, optional QK-norm, and pluggable RoPE.
decoder_layer	Pre-norm transformer decoder layers.
distributed	Tensor-parallel building blocks.
linear	Non-distributed fused linear projections.
mlp	MLP and SwiGLU-style GatedMLP.
moe	Mixture-of-Experts blocks.
norm	RMSNorm and AdaRMSNorm.