modeling_ember.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint
from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
from transformers.models.llama.modeling_llama import LlamaRMSNorm
from transformers.models.llama.modeling_llama import LlamaMLP
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
from transformers.cache_utils import DynamicCache

try:
    from .configuration_ember import EmberConfig
except ImportError:
    from configuration_ember import EmberConfig

try:
    from flash_attn import flash_attn_varlen_func
    FLASH_ATTN_AVAILABLE = True
except ImportError:
    FLASH_ATTN_AVAILABLE = False

@torch._dynamo.disable()
def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
    return flash_attn_varlen_func(
        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
        dropout_p=dropout_p, causal=True,
    )

class ClampedLlamaMLP(LlamaMLP):
    def forward(self, x):
        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
        up = self.up_proj(x)
        return self.down_proj(gate * up)

class XSAAttention(nn.Module):
    def __init__(self, config, layer_idx=None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.recurrent_cache_idx = None
        self._use_recurrent_slot = False
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
        self.attention_bias = getattr(config, "attention_bias", False)

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)

        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)

    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None,
                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None,
                expected_batch_size=None, cu_seqlens=None, max_seqlen=None, **kwargs):
        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)

        if hidden_states.ndim == 2:
            if expected_batch_size is None:
                raise RuntimeError(
                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
                    f"without an expected_batch_size to safely restore the batch dim."
                )
            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)

        bsz, q_len, _ = hidden_states.size()

        if expected_batch_size is not None and bsz != expected_batch_size:
            raise RuntimeError(
                f"XSAAttention: hidden_states batch size {bsz} does not match "
                f"expected_batch_size {expected_batch_size}."
            )

        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)

        query_states = self.q_norm(query_states)
        key_states = self.k_norm(key_states)

        cos, sin = position_embeddings

        use_flash = (
            cu_seqlens is not None
            and past_kv is None
            and getattr(self.config, "use_flash_attn", False)
            and FLASH_ATTN_AVAILABLE
        )

        if use_flash:
            total = bsz * q_len
            q = query_states.reshape(total, self.num_heads, self.head_dim)
            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim)
            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim)

            # FA2 FIX: Strictly cast to bf16 to prevent fp32 leaks from RoPE/RMSNorm
            q = q.to(torch.bfloat16)
            k = k.to(torch.bfloat16)
            v = v.to(torch.bfloat16)

            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)

            attn_output = _flash_varlen(
                q, k, v, cu_seqlens, max_seqlen,
                self.config.attention_dropout if self.training else 0.0,
            )

            if getattr(self.config, 'xsa_projection', True):
                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
                v_grouped = v.unsqueeze(2)
                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
                scale = (dot_yv / dot_vv).to(y.dtype)
                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)

            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
            return (attn_output, None)

        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)

        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        current_v = value_states

        target_idx = self.layer_idx
        if self._use_recurrent_slot and self.recurrent_cache_idx is not None:
            target_idx = self.recurrent_cache_idx

        if past_kv is not None:
            while len(past_kv) <= target_idx:
                past_kv.update(
                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
                    len(past_kv)
                )
            key_states, value_states = past_kv.update(key_states, value_states, target_idx)

        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)

        kv_len = key_states.shape[-2]

        if attention_mask is not None:
            if attention_mask.ndim == 2:
                if attention_mask.shape[-1] < kv_len:
                    attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
                elif attention_mask.shape[-1] > kv_len:
                    attention_mask = attention_mask[:, -kv_len:]

                pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min

                if q_len > 1:
                    if cache_position is None:
                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
                    kv_positions = torch.arange(kv_len, device=query_states.device)

                    neg_inf = torch.finfo(query_states.dtype).min
                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
                    attn_mask = causal_mask[None, None, :, :] + pad_mask

                    diag_idx = torch.arange(q_len, device=attn_mask.device)
                    start_idx = attn_mask.shape[-1] - q_len
                    attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
                else:
                    attn_mask = pad_mask
            else:
                if attention_mask.shape[0] != bsz:
                    raise RuntimeError(
                        f"attention_mask batch size {attention_mask.shape[0]} does not "
                        f"match hidden_states batch size {bsz}."
                    )
                attn_mask = attention_mask.to(dtype=query_states.dtype)
            is_causal = False
        else:
            is_causal = True
            attn_mask = None

        attn_output = F.scaled_dot_product_attention(
            query_states, key_states, value_states, attn_mask=attn_mask,
            dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
        )

        if getattr(self.config, 'xsa_projection', True):
            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
            v_grouped = current_v.unsqueeze(2)
            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
            scale = (dot_yv / dot_vv).to(y.dtype)
            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)

        return (attn_output, None)

@torch._dynamo.disable()
def _checkpointed_layer_forward(layer, hidden_states, attention_mask, position_ids,
                                 cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen):
    out = layer(
        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
        past_key_value=None, use_cache=False,
        cache_position=cache_position, position_embeddings=(cos, sin),
        expected_batch_size=expected_batch_size,
        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
    )
    hs_out = out[0] if isinstance(out, tuple) else out
    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
        raise RuntimeError(
            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
            f"batch size {expected_batch_size}."
        )
    return hs_out

class EmberModel(LlamaModel):
    def __init__(self, config):
        super().__init__(config)

        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"

        if getattr(config, "use_flash_attn", False) and not FLASH_ATTN_AVAILABLE:
            raise ImportError(
                "config.use_flash_attn=True but flash_attn is not importable. "
                "Install the FA2 wheel or set use_flash_attn=False."
            )

        p1 = config.prelude_layers
        r1 = p1 + config.recurrent_layers

        for i, layer in enumerate(self.layers):
            layer.self_attn = XSAAttention(config, layer_idx=i)
            layer.mlp = ClampedLlamaMLP(config)

        for i, layer in enumerate(self.layers[p1:r1]):
            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i

        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", True)

    def gradient_checkpointing_enable(self):
        self.gradient_checkpointing = True

    def gradient_checkpointing_disable(self):
        self.gradient_checkpointing = False

    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
                past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False,
                cache_position=None, return_dict=True, cu_seqlens=None, max_seqlen=None, **kwargs):
        if use_cache is None:
            use_cache = False

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]

        if cache_position is None:
            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(hidden_states, position_ids)
        cos, sin = position_embeddings

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        p1 = self.config.prelude_layers
        r1 = p1 + self.config.recurrent_layers
        c1 = r1 + self.config.coda_layers

        prelude = self.layers[:p1]
        recurrent = self.layers[p1:r1]
        coda = self.layers[r1:c1]

        use_ckpt = self.training and self.gradient_checkpointing and not use_cache

        def run_layer(layer, hs):
            if cu_seqlens is not None:
                torch._dynamo.mark_dynamic(cu_seqlens, 0)

            out = layer(
                hs, attention_mask=attention_mask, position_ids=position_ids,
                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
                cache_position=cache_position, position_embeddings=position_embeddings,
                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
            )
            hs_out = out[0] if isinstance(out, tuple) else out
            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
                raise RuntimeError(
                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
                    f"batch size {bsz}."
                )
            return hs_out

        def run_layer_maybe_ckpt(layer, hs):
            if use_ckpt:
                return torch.utils.checkpoint.checkpoint(
                    _checkpointed_layer_forward,
                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
                    cu_seqlens, max_seqlen,
                    use_reentrant=False,
                )
            return run_layer(layer, hs)

        for layer in prelude:
            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)

        if self.training:
            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02

        for layer in recurrent:
            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)

        if self.training:
            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02

        for layer in recurrent:
            layer.self_attn._use_recurrent_slot = True
            try:
                hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
            finally:
                layer.self_attn._use_recurrent_slot = False

        for layer in coda:
            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)

class EmberForCausalLM(LlamaForCausalLM):
    config_class = EmberConfig
    def __init__(self, config):
        super(LlamaForCausalLM, self).__init__(config)
        self.model = EmberModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def gradient_checkpointing_enable(self, **kwargs):
        self.model.gradient_checkpointing_enable()

    def gradient_checkpointing_disable(self):
        self.model.gradient_checkpointing_disable()

    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None,
                use_cache=None, num_logits_to_keep=0, position_ids=None, past_key_values=None,
                cache_position=None, cu_seqlens=None, max_seqlen=None, **kwargs):
        if use_cache is None:
            use_cache = False if (self.training or labels is not None) else True

        if num_logits_to_keep == 0 and "logits_to_keep" in kwargs:
            num_logits_to_keep = kwargs["logits_to_keep"]

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            cu_seqlens=cu_seqlens,
            max_seqlen=max_seqlen,
        )
        hidden_states = outputs[0]

        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
            raise RuntimeError(
                f"EmberModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
                f"expected batch size {expected_bsz}."
            )

        loss = None
        logits = None

        if labels is not None:
            shift_hidden = hidden_states[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()

            num_chunks = 8
            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
            l_chunks = shift_labels.chunk(num_chunks, dim=0)

            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
            total_tokens = 0
            for h_c, l_c in zip(h_chunks, l_chunks):
                logits_c = self.lm_head(h_c)
                chunk_loss = F.cross_entropy(
                    logits_c.view(-1, logits_c.size(-1)).float(),
                    l_c.view(-1),
                    reduction="sum",
                )
                total_loss = total_loss + chunk_loss
                total_tokens += l_c.numel()
            loss = (total_loss / total_tokens).to(hidden_states.dtype)
        else:
            slice_hidden = hidden_states if num_logits_to_keep == 0 else hidden_states[:, -num_logits_to_keep:, :]
            logits = self.lm_head(slice_hidden)

        return CausalLMOutputWithPast(
            loss=loss, logits=logits, past_key_values=outputs.past_key_values
        )