class DecoderModel(nn.Module):
"""
A Transformer-based decoder model.
This model consists of an embedding layer, a stack of Transformer blocks,
an optional final layer normalization (for Pre-LN architectures), and a
language modeling head to predict token logits.
"""
def __init__(
self,
vocab_size: int,
hidden_size: int,
num_layers: int,
num_heads: int,
max_seq_len: int = 512,
intermediate_size: int | None = None,
pos_encoding_learned: bool = False,
embedding_dropout_p: float = 0.1,
attn_dropout_p: float = 0.1,
mlp_dropout_p: float = 0.1,
mlp_activation: str | nn.Module = "gelu",
norm_eps: float = 1e-5,
norm_first: bool = True,
is_causal: bool = True, # Default to True for a decoder model
padding_idx: int | None = None,
qkv_bias: bool = True, # Bias for QKV in MHA within TransformerBlock
mlp_bias: bool = True, # Bias for MLP in TransformerBlock
lm_head_bias: bool = True, # Bias for the final LM head
num_experts: int = 0,
top_k: int = 0,
num_kv_heads: int | None = None, # For GQA support
use_glu: bool = False,
norm_impl: str = "layer_norm",
device: torch.device | str | None = None,
dtype: torch.dtype | None = None,
attn_impl: str = "mha",
mlp_impl: str = "mlp",
gradient_checkpointing: bool = False,
window_size: int | None = None,
use_rope: bool = False,
rope_theta: float = 10000.0,
use_alibi: bool = False, # ALiBi linear-bias PE (BLOOM-style, mha backend)
attn_sparse: dict | None = None, # sparse/streaming scheme config (dict or None)
):
"""
Initializes the DecoderModel.
"""
super().__init__()
# ALiBi is a second positional-encoding system; it cannot coexist with
# RoPE, and this milestone wires it only into the mha attention
# backend (flash_attn/MLA/paged reject loudly rather than silently
# dropping the PE — RIL — ALiBi milestone).
if use_alibi and use_rope:
raise ValueError("use_alibi=True and use_rope=True are mutually exclusive position encodings")
if use_alibi and attn_impl != "mha":
raise NotImplementedError(
f"use_alibi=True is only supported with attn_impl='mha', got attn_impl={attn_impl!r}; "
"flash_attn/MLA have no additive-bias channel in this milestone"
)
# Sparse/streaming scheme (RIL TASK-248): the pattern mask is consumed only
# by backends that route through the sdpa wrapper (mha and mla). The
# flash_attn backend explicitly ignores ``attn_mask``, so combining a
# sparse scheme with it would silently run dense attention — refuse loudly
# instead of advertising a scheme that has no effect.
if attn_sparse is not None and attn_impl == "flash_attn":
raise NotImplementedError(
"attn_sparse is not supported with attn_impl='flash_attn'; flash_attn "
"ignores the boolean sparse mask and would silently run dense "
"attention. Use attn_impl='mha' or 'mla' (both route through the "
"SDPA wrapper that consumes the mask)"
)
factory_kwargs = make_factory_kwargs(device, dtype)
resolved_norm_factory = _resolve_norm_factory(norm_impl)
self.hidden_size = hidden_size
self.num_heads = num_heads
self.max_seq_len = max_seq_len
self.norm_first = norm_first # Store for final norm logic
# ``norm_impl`` is a model-defining flag (LayerNorm vs RMSNorm); store
# it so hf_publisher can persist it and save->load roundtrips keep the
# same normalization function (RIL ISS-062).
self.norm_impl = norm_impl
# RoPE is model-defining too (real Llama/Mistral inject position in
# attention, not additive embeddings); store so hf_publisher persists
# and the loader honors it (RIL ISS-062).
self.use_rope = use_rope
self.rope_theta = rope_theta
self.use_alibi = use_alibi
# Sliding-window attention size is model-defining (Mistral uses a
# 4096-token window): store it so hf_publisher persists it and the
# loader honors an external checkpoint's ``sliding_window`` instead of
# running full-context attention past the window (RIL ISS-242).
self.window_size = window_size
# Sparse/streaming attention scheme is model-defining too: when set, the
# forward builds the dispatched mask from the current sequence length so
# training/inference can select a scheme by name (RIL TASK-243). Stored
# as an immutable mapping snapshot.
self.attn_sparse = dict(attn_sparse) if attn_sparse else None
# Bias flags are model-defining too: real Llama/Mistral are bias-free
# (qkv/mlp/lm_head), while our grown-from-scratch default is biased.
# Store them so hf_publisher persists the actual values and the loader
# can honor external bias-free checkpoints (RIL ISS-062).
self.qkv_bias = bool(qkv_bias)
self.mlp_bias = bool(mlp_bias)
self.lm_head_bias = bool(lm_head_bias)
self._gradient_checkpointing = gradient_checkpointing
self.embedding_layer = EmbeddingLayer(
vocab_size=vocab_size,
hidden_size=hidden_size,
max_seq_len=max_seq_len,
pos_encoding_learned=pos_encoding_learned,
dropout_p=embedding_dropout_p,
padding_idx=padding_idx,
use_rope=use_rope,
**factory_kwargs,
)
if intermediate_size is None:
intermediate_size = 4 * hidden_size
# ALiBi linear-bias PE: one shared bias module (its ``num_heads`` x
# ``max_seq_len`` score cache) threaded into every attention block.
self.alibi = (
ALiBiPositionBias(num_heads=num_heads, max_seq_len=max_seq_len, **factory_kwargs) if use_alibi else None
)
self.transformer_blocks = nn.ModuleList(
[
TransformerBlock(
hidden_size=hidden_size,
num_heads=num_heads,
intermediate_size=intermediate_size,
attn_dropout_p=attn_dropout_p,
mlp_dropout_p=mlp_dropout_p,
mlp_activation=mlp_activation,
norm_eps=norm_eps,
norm_first=norm_first,
is_causal=is_causal, # Pass overall model's causality default
qkv_bias=qkv_bias,
mlp_bias=mlp_bias,
num_experts=num_experts,
top_k=top_k,
num_kv_heads=num_kv_heads,
use_glu=use_glu, # Pass use_glu
norm_type=resolved_norm_factory,
window_size=window_size, # Pass window_size
attn_impl=attn_impl, # Pass attn_impl
mlp_impl=mlp_impl, # Pass mlp_impl
max_seq_len=max_seq_len, # RoPE context for the block
use_rope=use_rope,
rope_theta=rope_theta,
alibi=self.alibi,
**factory_kwargs,
)
for _ in range(num_layers)
]
)
self.final_norm = None
if self.norm_first:
self.final_norm = resolved_norm_factory(hidden_size, eps=norm_eps, **factory_kwargs)
self.lm_head = nn.Linear(hidden_size, vocab_size, bias=lm_head_bias, **factory_kwargs)
self.max_seq_len = max_seq_len
def forward(
self,
input_ids: torch.Tensor,
attn_mask: torch.Tensor | None = None,
kv_caches: list[KVCache] | None = None,
use_cache: bool = False,
position_ids: torch.Tensor | None = None,
batch_indices: torch.Tensor | None = None,
paged_kv_cache: object | None = None,
) -> torch.Tensor | tuple[torch.Tensor, list[KVCache] | None]:
"""
Forward pass of the DecoderModel.
Args:
input_ids: Input token IDs of shape [B, S].
attn_mask: Optional attention mask broadcastable to SDPA.
kv_caches: Pre-allocated KV caches, one per transformer layer.
use_cache: When True, update ``kv_caches`` in place and return them.
position_ids: Explicit position IDs of shape [B, S].
batch_indices: Cache slot indices for continuous batching.
paged_kv_cache: Block-allocator KV cache; when set, ``kv_caches``
is unused and the model routes K/V through the block
allocator + ``paged_attention_forward``.
Returns:
Logits tensor, or ``(logits, kv_caches)`` when ``use_cache=True``;
in the paged path the second element is ``None``.
"""
if self._gradient_checkpointing and use_cache:
raise ValueError("Gradient checkpointing is incompatible with use_cache=True. ")
if use_cache and kv_caches is None and paged_kv_cache is None:
raise ValueError("use_cache=True requires either kv_caches or paged_kv_cache to be set.")
if kv_caches is not None and paged_kv_cache is not None:
raise ValueError("Pass either kv_caches or paged_kv_cache, not both.")
start_pos = 0
if kv_caches is not None and kv_caches[0].seq_len > 0:
start_pos = kv_caches[0].seq_len
# Sparse/streaming scheme selected by name in the model config: unless the
# caller supplied an explicit mask (which always wins), build the
# pattern-only mask over the *key history*. A KV-cache decode step has
# ``Sq``=current-token rows but ``Sk``=accumulated-key columns, so the
# mask cannot be square here — otherwise sink/window would never
# constrain the cached past keys (RIL TASK-245).
if attn_mask is None and self.attn_sparse is not None:
key_len = start_pos + input_ids.shape[1]
attn_mask = _build_config_attention_mask(self, input_ids.shape[1], key_len=key_len)
hidden_states = self.embedding_layer(input_ids, start_pos=start_pos, position_ids=position_ids)
for i, block in enumerate(self.transformer_blocks):
kv_cache = kv_caches[i] if kv_caches is not None else None
if self._gradient_checkpointing and self.training:
hidden_states = checkpoint(
block,
hidden_states,
attn_mask,
None,
None,
False,
use_reentrant=False,
)
else:
block_outputs = block(
hidden_states,
attn_mask=attn_mask,
is_causal=None,
kv_cache=kv_cache,
use_cache=use_cache,
batch_indices=batch_indices,
start_pos=position_ids if (batch_indices is not None and position_ids is not None) else start_pos,
paged_kv_cache=paged_kv_cache,
layer_idx=i,
)
if paged_kv_cache is not None:
# Paged path returns the output directly.
hidden_states = block_outputs
elif use_cache:
hidden_states, _current_kv = block_outputs
else:
hidden_states = block_outputs
if self.final_norm is not None:
hidden_states = self.final_norm(hidden_states)
logits = self.lm_head(hidden_states)
if use_cache:
# In the paged path (``paged_kv_cache`` set) ``kv_caches`` is
# None and the second element is None; the engine unpacks it
# as ``logits, _``.
return logits, kv_caches
return logits
@property
def gradient_checkpointing(self) -> bool:
"""Whether gradient checkpointing is enabled."""
return self._gradient_checkpointing
def enable_gradient_checkpointing(self) -> None:
"""Enable gradient checkpointing to reduce memory usage during training."""
self._gradient_checkpointing = True
def disable_gradient_checkpointing(self) -> None:
"""Disable gradient checkpointing."""
self._gradient_checkpointing = False