Gemma4UnifiedAssistantForCausalLM

Model type	Gemma4UnifiedAssistantForCausalLM
Class	Gemma4AssistantCausalLMModel
Task	gemma4-assistant
Source	models/gemma4_assistant.py

Description

Gemma4-Assistant speculative-decoding draft model.

See module docstring for the architecture overview and limitations.

Weight module layout (mirrors HF Gemma4AssistantForCausalLM):

• self.pre_projection Linear[2*backbone → hidden] (no bias)

• self.model.layers.{i}.{...} one Gemma4AssistantDecoderLayer per layer

• self.model.norm final RMSNorm on hidden

• self.lm_head Linear[hidden → vocab] (no bias; tied to embed_tokens upstream)

• self.post_projection Linear[hidden → backbone] (no bias)

Usage

mobius build --model <MODEL_ID> output_dir/

from mobius import build

model = build("<MODEL_ID>")