Backbone¶

def backbone_forward(
    self,
    embeds: torch.Tensor,
    token_type: torch.Tensor,
    cache: dict[str, Any] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, dict[str, Any] | None]:
    """Run the Llama backbone over the token sequence.

    Args:
        embeds: ``[B, T_total, D]`` embedding tensor from ``StepEmbedder``.
        token_type: ``[B, T_total]`` int64 ``TokenType`` ids; ``PAD`` positions
            are masked out from attention.
        cache: KV-cache dict from a previous call, or ``None`` for full prefill.
            Reads and writes the ``"backbone"`` key.
        use_cache: If ``True``, return an updated KV-cache dict.
        cache_position: Unused; present for interface compatibility.
        **kwargs: Forwarded to the underlying ``LlamaModel``.

    Returns:
        Tuple of ``(hidden_states [B, T_total, D], cache_dict | None)``.
    """
    cache = cache or {}

    has_padding = bool((token_type == TokenType.PAD).any())
    attention_mask = (token_type != TokenType.PAD).long() if has_padding else None

    out = self.backbone(
        inputs_embeds=embeds,
        past_key_values=cache.get("backbone", None),
        use_cache=use_cache,
        position_ids=None,
        attention_mask=attention_mask,
        **kwargs,
    )

    new_cache = {"backbone": out.past_key_values} if use_cache else None
    return out.last_hidden_state, new_cache

def backbone_forward(
    self,
    embeds: torch.Tensor,
    token_type: torch.Tensor,
    cache: dict[str, Any] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, dict[str, Any] | None]:
    """Run the Qwen3 backbone over the token sequence.

    Args:
        embeds: ``[B, T_total, D]`` embedding tensor from ``StepEmbedder``.
        token_type: ``[B, T_total]`` int64 ``TokenType`` ids; ``PAD`` positions
            are masked out from attention.
        cache: KV-cache dict from a previous call, or ``None`` for full prefill.
            Reads and writes the ``"backbone"`` key.
        use_cache: If ``True``, return an updated KV-cache dict.
        cache_position: Unused; present for interface compatibility.
        **kwargs: Forwarded to the underlying ``Qwen3Model``.

    Returns:
        Tuple of ``(hidden_states [B, T_total, D], cache_dict | None)``.
    """
    cache = cache or {}

    has_padding = bool((token_type == TokenType.PAD).any())
    attention_mask = (token_type != TokenType.PAD).long() if has_padding else None

    out = self.backbone(
        inputs_embeds=embeds,
        past_key_values=cache.get("backbone", None),
        use_cache=use_cache,
        position_ids=None,
        attention_mask=attention_mask,
        **kwargs,
    )

    new_cache = {"backbone": out.past_key_values} if use_cache else None
    return out.last_hidden_state, new_cache

def backbone_forward(
    self,
    embeds: torch.Tensor,
    token_type: torch.Tensor,
    cache: dict[str, Any] | None = None,
    use_cache: bool = False,
    cache_position: torch.Tensor | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, dict[str, Any] | None]:
    """Pass embeddings through unchanged; always returns ``None`` for cache.

    Args:
        embeds: ``[B, T_total, D]`` embedding tensor.
        token_type: Ignored.
        cache: Ignored.
        use_cache: Ignored.
        cache_position: Ignored.

    Returns:
        Tuple of ``(embeds unchanged, None)``.
    """
    return embeds, None

def build(self, hidden_dim: int) -> LlamaModel:
    """Instantiate a ``LlamaModel`` with this config.

    Args:
        hidden_dim: Model hidden dimension ``D``; must be divisible by ``num_heads``.

    Returns:
        ``LlamaModel`` with the final norm replaced by ``nn.Identity``.
    """
    _disable_cudnn_sdp()
    if hidden_dim % self.num_heads != 0:
        raise ValueError(
            f"hidden_dim ({hidden_dim}) must be divisible by num_heads ({self.num_heads})."
        )
    ffn_size = self.intermediate_size if self.intermediate_size is not None else hidden_dim * self.expand
    config_kwargs: dict = dict(
        vocab_size=1,
        hidden_size=hidden_dim,
        num_attention_heads=self.num_heads,
        num_key_value_heads=self.num_key_value_heads,
        intermediate_size=ffn_size,
        max_position_embeddings=self.max_position_embeddings,
        attention_dropout=0.0,
        attention_bias=self.attention_bias,
        rms_norm_eps=self.rms_norm_eps,
        num_hidden_layers=self.num_layers,
    )
    if self.rope_parameters is not None:
        config_kwargs["rope_parameters"] = self.rope_parameters
    config = LlamaConfig(**config_kwargs)
    config._attn_implementation = "sdpa"
    model = LlamaModel(config)
    model.norm = nn.Identity()  # type: ignore[assignment]
    return model

def build(self, hidden_dim: int) -> Qwen3Model:
    """Instantiate a ``Qwen3Model`` with this config.

    Args:
        hidden_dim: Model hidden dimension ``D``. When ``head_dim`` is ``None``,
            must be divisible by ``num_heads``.

    Returns:
        ``Qwen3Model`` with the final norm replaced by ``nn.Identity``.
    """
    _disable_cudnn_sdp()
    if self.head_dim is None:
        if hidden_dim % self.num_heads != 0:
            raise ValueError(
                f"hidden_dim ({hidden_dim}) must be divisible by num_heads ({self.num_heads})."
            )
        resolved_head_dim = hidden_dim // self.num_heads
    else:
        resolved_head_dim = int(self.head_dim)
    ffn_size = self.intermediate_size if self.intermediate_size is not None else hidden_dim * self.expand
    config_kwargs: dict = dict(
        vocab_size=1,
        hidden_size=hidden_dim,
        num_attention_heads=self.num_heads,
        num_key_value_heads=self.num_key_value_heads,
        head_dim=resolved_head_dim,
        intermediate_size=ffn_size,
        max_position_embeddings=self.max_position_embeddings,
        attention_dropout=0.0,
        attention_bias=self.attention_bias,
        rms_norm_eps=self.rms_norm_eps,
        num_hidden_layers=self.num_layers,
        use_sliding_window=self.use_sliding_window,
    )
    if self.rope_parameters is not None:
        config_kwargs["rope_parameters"] = self.rope_parameters
    config = Qwen3Config(**config_kwargs)
    config._attn_implementation = "sdpa"
    model = Qwen3Model(config)
    model.norm = nn.Identity()  # type: ignore[assignment]
    return model