# SPDX-License-Identifier: Apache-2.0 """Hugging Face ``transformers`` implementation of raon-1-mini. Architecture (released 28.8B / 4.3B-active checkpoint): 25 decoder layers on a 4-stream Raon_mHC residual, 18 Raon Delta Attention (RDA) layers + 7 NoPE gated Multi-head Latent Attention (MLA) layers, a dense SiTU-GLU MLP in layer 0 and a 128-expert top-6 sigmoid-routed MoE with one shared expert everywhere else, and an untied output head. Every equation follows the training implementation's plain-PyTorch reference: * Raon_mHC selector, router, SiTU nonlinearity, attention softmax, RDA gates and the RDA recurrent state are computed in fp32; GEMM activations live in the model dtype. * RDA prefill/decode use a pure-PyTorch recurrence; no custom kernel package is required. * The parameter names are the release checkpoint's tensor names. * The released checkpoint stores routed experts as MXFP8 (``quant_method: mxfp8``) and loads them as E4M3 weights + E8M0 scales on every device. On sm_100/sm_103 the routed experts run as native grouped MXFP8 GEMMs (activations quantized per call with the production rule, W8A8); elsewhere ``routed_linear`` expands each used expert to the activation dtype per call (W8A16). Passing ``dequantize=True`` in the quantization config instead dequantizes exactly to BF16 at load. """ from __future__ import annotations import contextlib import copy import math from typing import Any, ClassVar import torch import torch.nn.functional as F from torch import nn from transformers import initialization as init from transformers.generation import GenerationMixin from transformers.integrations.finegrained_fp8 import FP8Linear from transformers.utils.quantization_config import FineGrainedFP8Config from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel from transformers.utils import logging from .configuration_raon1_mini import Raon1MiniConfig logger = logging.get_logger(__name__) # --------------------------------------------------------------------------------------------- # primitives # --------------------------------------------------------------------------------------------- class Raon1MiniRMSNorm(nn.Module): """RMSNorm with fp32 statistics, plain (non zero-centered) gamma, no bias.""" def __init__(self, dim: int, eps: float) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, hidden: torch.Tensor, *, keep_fp32: bool = False) -> torch.Tensor: value = hidden.float() normed = ( value * torch.rsqrt(value.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight.float() ) return normed if keep_fp32 else normed.to(hidden.dtype) def extra_repr(self) -> str: return f"{tuple(self.weight.shape)}, eps={self.eps}" def situ_glu( gate: torch.Tensor, value: torch.Tensor, gate_beta: float, value_beta: float ) -> torch.Tensor: """SiTU-GLU in fp32: ``gate_beta*tanh(g/gate_beta)*sigmoid(g) * value_beta*tanh(v/value_beta)``.""" g = gate.float() v = value.float() return (gate_beta * torch.tanh(g / gate_beta) * torch.sigmoid(g)) * ( value_beta * torch.tanh(v / value_beta) ) def helmert_basis(n: int) -> torch.Tensor: """Orthonormal ``Q`` in ``R^{n x (n-1)}`` with ``Q^T Q = I`` and ``Q^T 1 = 0``.""" q = torch.zeros(n, n - 1, dtype=torch.float32) for k in range(1, n): q[:k, k - 1] = 1.0 q[k, k - 1] = -float(k) q[:, k - 1] /= math.sqrt(k * (k + 1.0)) return q _RAON_MHC_CONSTANTS: dict[ tuple[int, torch.device], tuple[torch.Tensor, torch.Tensor, torch.Tensor] ] = {} def _raon_mhc_constants( n: int, device: torch.device ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """``(helmert Q [n, n-1], I [n, n], B = (I + 11^T/n)/2 [n, n])`` in fp32 on ``device``. Built on demand instead of as module buffers: ``from_pretrained`` materialises non-persistent buffers on the meta device without values. """ key = (n, device) if key not in _RAON_MHC_CONSTANTS: eye = torch.eye(n, dtype=torch.float32, device=device) _RAON_MHC_CONSTANTS[key] = ( helmert_basis(n).to(device), eye, 0.5 * (eye + torch.full_like(eye, 1.0 / n)), ) return _RAON_MHC_CONSTANTS[key] # --------------------------------------------------------------------------------------------- # Raon_mHC multi-stream residual # --------------------------------------------------------------------------------------------- class RaonMHCStreamMix(nn.Module): """The ``res.*`` parameters of one Raon_mHC site: the stream-to-stream residual chart.""" def __init__(self, num_streams: int, hidden_size: int) -> None: super().__init__() m = (num_streams - 1) ** 2 self.phi_z = nn.Parameter(torch.zeros(m, num_streams * hidden_size)) self.alpha_z = nn.Parameter(torch.ones(())) self.b_z = nn.Parameter(torch.zeros(m)) self.eta = nn.Parameter(torch.zeros((), dtype=torch.float32)) def forward(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: return self.phi_z, self.alpha_z, self.b_z, self.eta class RaonMHC(nn.Module): """One Raon_mHC site (attention or MLP): selects the branch input from the streams and writes back. ``route`` returns fp32 coefficients ``h_pre [N, n]``, ``h_post [N, n]`` and ``h_res [N, n, n]`` for ``N = batch*seq`` rows (Raon_mHC chart: ``E = Q Z Q^T``, ``rho = max(0, max(-E/B))``, ``Phi = B + E/(1+rho)``, ``H = (1-gamma) I + gamma Phi`` with ``B = (I + 11^T/n)/2`` and ``gamma = sigmoid(eta)``). """ def __init__(self, config: Raon1MiniConfig) -> None: super().__init__() n, c = config.raon_mhc_streams, config.hidden_size self.num_streams = n self.eps = config.raon_mhc_selector_eps self.post_factor = config.raon_mhc_post_factor self.phi_pre = nn.Parameter(torch.zeros(n, n * c)) self.phi_post = nn.Parameter(torch.zeros(n, n * c)) self.alpha_pre = nn.Parameter(torch.ones(())) self.alpha_post = nn.Parameter(torch.ones(())) self.b_pre = nn.Parameter(torch.zeros(n)) self.b_post = nn.Parameter(torch.zeros(n)) self.res = RaonMHCStreamMix(n, c) def route( self, stream: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return self(stream) def forward( self, stream: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: phi_z, alpha_z, b_z, eta = self.res() n = self.num_streams rows = stream.shape[0] * stream.shape[1] flat = stream.reshape(rows, n * stream.shape[-1]) weight = torch.cat([self.phi_pre, self.phi_post, phi_z], dim=0) # RMS scale applied AFTER the (model-dtype) selector GEMM, as in training. inv = torch.rsqrt(flat.float().pow(2).mean(-1, keepdim=True) + self.eps) z = F.linear(flat, weight.to(flat.dtype)).float() * inv t_pre = self.alpha_pre.float() * z[:, :n] + self.b_pre.float() t_post = self.alpha_post.float() * z[:, n : 2 * n] + self.b_post.float() t_z = alpha_z.float() * z[:, 2 * n :] + b_z.float() h_pre = torch.sigmoid(t_pre) h_post = self.post_factor * torch.sigmoid(t_post) nm = n - 1 helmert, eye, chart_b = _raon_mhc_constants(n, z.device) e = torch.einsum("im,bmk,jk->bij", helmert, t_z.reshape(rows, nm, nm), helmert) rho = (-e / chart_b).amax(dim=(-2, -1)).clamp_min(0.0) phi = chart_b + e / (1.0 + rho)[:, None, None] gamma = torch.sigmoid(eta.float()) h_res = (1.0 - gamma) * eye + gamma * phi return h_pre, h_post, h_res @staticmethod def aggregate( stream_f32: torch.Tensor, h_pre: torch.Tensor, dtype: torch.dtype ) -> torch.Tensor: """``[N, n, C]`` fp32 streams -> ``[N, C]`` branch input in the model dtype.""" return torch.einsum("bn,bnc->bc", h_pre, stream_f32).to(dtype) @staticmethod def mix( stream_f32: torch.Tensor, branch_out: torch.Tensor, h_post: torch.Tensor, h_res: torch.Tensor, dtype: torch.dtype, ) -> torch.Tensor: """Residual mixing plus the branch write-back, rounded once to the model dtype.""" mixed = torch.einsum("bij,bjc->bic", h_res, stream_f32) return (mixed + h_post.unsqueeze(-1) * branch_out.float().unsqueeze(1)).to( dtype ) # --------------------------------------------------------------------------------------------- # RDA (Raon Delta Attention) # --------------------------------------------------------------------------------------------- class Raon1MiniShortConv(nn.Module): """Depthwise causal conv (no bias) whose weight is stored squeezed as ``[channels, kernel]``.""" def __init__(self, channels: int, kernel_size: int) -> None: super().__init__() self.weight = nn.Parameter(torch.zeros(channels, kernel_size)) self.kernel_size = kernel_size def forward( self, x: torch.Tensor, state: torch.Tensor | None ) -> tuple[torch.Tensor, torch.Tensor]: """``x [B, T, C]`` raw rows -> silu(conv) ``[B, T, C]`` in x's dtype, plus the new raw window ``[B, C, K-1]``.""" channels = x.shape[-1] k = self.kernel_size signal = x.transpose(1, 2) # [B, C, T] padded = ( F.pad(signal, (k - 1, 0)) if state is None else torch.cat([state.to(signal.dtype), signal], dim=-1) ) # fp32 accumulation and silu, one rounding on the output (causal_conv1d semantics). out = F.conv1d( padded.float(), self.weight.float().unsqueeze(1), groups=channels ) new_state = padded[..., -(k - 1) :] return F.silu(out).to(x.dtype).transpose(1, 2), new_state def torch_recurrent_rda( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, log_alpha: torch.Tensor, beta: torch.Tensor, state: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: """V-first delta rule, token serial, fp32 state ``[B, H, V, K]``. Reference for the fused kernels. ``q/k/v [B, T, H, D]`` (q already l2-normalised and scaled, k l2-normalised), ``log_alpha [B, T, H, K]`` fp32 log decay, ``beta [B, T, H]``. Returns ``out [B, T, H, V]`` in v's dtype and the final state. """ batch, seq, heads, key_dim = k.shape value_dim = v.shape[-1] if state is None: state = torch.zeros( batch, heads, value_dim, key_dim, dtype=torch.float32, device=q.device ) else: state = state.float() out = torch.empty( batch, seq, heads, value_dim, dtype=torch.float32, device=q.device ) qf, kf, vf, bf = q.float(), k.float(), v.float(), beta.float() for t in range(seq): state = state * torch.exp(log_alpha[:, t]).unsqueeze(-2) predicted = torch.einsum("bhvk,bhk->bhv", state, kf[:, t]) error = vf[:, t] - predicted state = state + bf[:, t][..., None, None] * torch.einsum( "bhv,bhk->bhvk", error, kf[:, t] ) out[:, t] = torch.einsum("bhvk,bhk->bhv", state, qf[:, t]) return out.to(v.dtype), state class Raon1MiniRDA(nn.Module): """Raon Delta Attention layer: gated delta-rule linear attention with short convolutions and a write gate.""" def __init__(self, config: Raon1MiniConfig) -> None: super().__init__() hidden = config.hidden_size self.num_heads = config.rda_num_heads self.head_dim = config.rda_head_dim self.proj_size = self.num_heads * self.head_dim self.gate_lower_bound = config.rda_gate_lower_bound self.norm_eps = config.rms_norm_eps kernel = config.rda_conv_kernel_size self.q_proj = nn.Linear(hidden, self.proj_size, bias=False) self.k_proj = nn.Linear(hidden, self.proj_size, bias=False) self.v_proj = nn.Linear(hidden, self.proj_size, bias=False) self.q_conv1d = Raon1MiniShortConv(self.proj_size, kernel) self.k_conv1d = Raon1MiniShortConv(self.proj_size, kernel) self.v_conv1d = Raon1MiniShortConv(self.proj_size, kernel) self.f_a_proj = nn.Linear(hidden, config.rda_decay_rank, bias=False) self.f_b_proj = nn.Linear(config.rda_decay_rank, self.proj_size, bias=False) self.b_proj = nn.Linear(hidden, self.num_heads, bias=False) self.rda_w_proj = nn.Linear(hidden, self.num_heads, bias=False) self.g_proj = nn.Linear(hidden, self.proj_size, bias=False) self.A_log = nn.Parameter(torch.zeros(self.num_heads, dtype=torch.float32)) self.dt_bias = nn.Parameter(torch.zeros(self.proj_size, dtype=torch.float32)) self.o_norm = Raon1MiniRMSNorm(self.head_dim, config.rms_norm_eps) self.o_proj = nn.Linear(self.proj_size, hidden, bias=False) def forward( self, hidden: torch.Tensor, cache: Raon1MiniDynamicCache | None, layer_idx: int, token_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch, seq, _ = hidden.shape dtype = hidden.dtype if token_mask is not None: # Left padding: padded rows must not enter the conv window or the recurrence. hidden = hidden * token_mask[..., None].to(dtype) q = self.q_proj(hidden) k = self.k_proj(hidden) v = self.v_proj(hidden) gate = self.g_proj(hidden) beta_raw = self.b_proj(hidden) rda_delta = self.rda_w_proj(hidden) decay_raw = self.f_b_proj(self.f_a_proj(hidden)) conv_state = cache.conv_states[layer_idx] if cache is not None else None q_hist = k_hist = v_hist = None if conv_state is not None: q_hist, k_hist, v_hist = conv_state.split(self.proj_size, dim=1) q, q_hist = self.q_conv1d(q, q_hist) k, k_hist = self.k_conv1d(k, k_hist) v, v_hist = self.v_conv1d(v, v_hist) if cache is not None: cache.conv_states[layer_idx] = torch.cat([q_hist, k_hist, v_hist], dim=1) q = q.reshape(batch, seq, self.num_heads, self.head_dim) k = k.reshape(batch, seq, self.num_heads, self.head_dim) v = v.reshape(batch, seq, self.num_heads, self.head_dim) # _prepare_after_conv: l2-normalise q/k in fp32 (eps 1e-6), scale q by key_dim**-0.5. q = F.normalize(q.float(), dim=-1, eps=1e-6).to(dtype) * (self.head_dim**-0.5) k = F.normalize(k.float(), dim=-1, eps=1e-6).to(dtype) # Log decay (per-head scale, lower bound): fp32. decay = decay_raw.float().reshape( batch, seq, self.num_heads, self.head_dim ) + self.dt_bias.float().reshape(self.num_heads, self.head_dim) scaled = torch.exp(self.A_log.float())[None, None, :, None] * decay log_alpha = ( -F.softplus(scaled) if self.gate_lower_bound is None else self.gate_lower_bound * torch.sigmoid(scaled) ) beta = torch.sigmoid(beta_raw.float()).to(dtype) # RDA: write gate sigmoid(beta+delta), erase gate sigmoid(beta); realised as a fp32 value scale. ratio = torch.exp( F.logsigmoid(beta_raw.float() + rda_delta.float()) - F.logsigmoid(beta_raw.float()) ) v = (v.float() * ratio.unsqueeze(-1)).to(dtype) recurrent_state = ( cache.recurrent_states[layer_idx] if cache is not None else None ) out, final_state = self._core( q, k, v, log_alpha.contiguous(), beta, recurrent_state ) if cache is not None: cache.recurrent_states[layer_idx] = final_state # Per-head RMSNorm then sigmoid output gate, both in fp32, one rounding. normed = self.o_norm(out.reshape(-1, self.head_dim), keep_fp32=True) gated = (normed * torch.sigmoid(gate.float().reshape(-1, self.head_dim))).to( dtype ) return self.o_proj(gated.reshape(batch, seq, self.proj_size)) def _core(self, q, k, v, log_alpha, beta, state): return torch_recurrent_rda(q, k, v, log_alpha, beta, state) # --------------------------------------------------------------------------------------------- # NoPE gated MLA # --------------------------------------------------------------------------------------------- class Raon1MiniMLA(nn.Module): """Multi-head latent attention without positional encoding and with a sigmoid output gate. The 64 "rope" channels are never rotated: they are extra key/query channels shared across heads (``k_pe`` is expanded over heads), so scores are 192-wide dot products scaled by ``1/sqrt(192)``. """ def __init__(self, config: Raon1MiniConfig) -> None: super().__init__() hidden = config.hidden_size self.num_heads = config.num_attention_heads self.qk_nope = config.qk_nope_head_dim self.qk_pe = config.qk_rope_head_dim self.v_dim = config.v_head_dim self.q_head_dim = self.qk_nope + self.qk_pe self.kv_lora_rank = config.kv_lora_rank self.softmax_scale = self.q_head_dim**-0.5 self.use_output_gate = config.mla_use_output_gate self.q_a_proj = nn.Linear(hidden, config.q_lora_rank, bias=False) self.q_a_layernorm = Raon1MiniRMSNorm( config.q_lora_rank, config.mla_latent_norm_eps ) self.q_b_proj = nn.Linear( config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False ) self.kv_a_proj_with_mqa = nn.Linear( hidden, config.kv_lora_rank + self.qk_pe, bias=False ) self.kv_a_layernorm = Raon1MiniRMSNorm( config.kv_lora_rank, config.mla_latent_norm_eps ) self.kv_b_proj = nn.Linear( config.kv_lora_rank, self.num_heads * (self.qk_nope + self.v_dim), bias=False, ) if self.use_output_gate: self.g_proj = nn.Linear(hidden, self.num_heads * self.v_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.v_dim, hidden, bias=False) def forward( self, hidden: torch.Tensor, cache: Raon1MiniDynamicCache | None, layer_idx: int, attention_mask: torch.Tensor | None, ) -> torch.Tensor: batch, seq, _ = hidden.shape dtype = hidden.dtype q_lat = self.q_a_layernorm(self.q_a_proj(hidden)) q = self.q_b_proj(q_lat).reshape(batch, seq, self.num_heads, self.q_head_dim) kv_full = self.kv_a_proj_with_mqa(hidden) kv_lat, k_pe = torch.split(kv_full, [self.kv_lora_rank, self.qk_pe], dim=-1) kv = self.kv_b_proj(self.kv_a_layernorm(kv_lat)).reshape( batch, seq, self.num_heads, self.qk_nope + self.v_dim ) k_nope, value = torch.split(kv, [self.qk_nope, self.v_dim], dim=-1) key = torch.cat( [k_nope, k_pe.unsqueeze(2).expand(batch, seq, self.num_heads, self.qk_pe)], dim=-1, ) query = q.transpose(1, 2) # [B, H, T, 192] key = key.transpose(1, 2) value = value.transpose(1, 2) if cache is not None: key, value = cache.update(key, value, layer_idx) kv_len = key.shape[2] if attention_mask is None and seq > 1 and kv_len != seq: # Chunked prefill after cached tokens: explicit causal mask with the cache offset. offset = kv_len - seq attention_mask = torch.ones( seq, kv_len, dtype=torch.bool, device=hidden.device ).tril(offset)[None, None] context = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, is_causal=attention_mask is None and seq > 1, scale=self.softmax_scale, ) context = context.transpose(1, 2).reshape( batch, seq, self.num_heads * self.v_dim ) if self.use_output_gate: # The training NoPE path evaluates the gate in the activation dtype. context = context * torch.sigmoid(self.g_proj(hidden)) return self.o_proj(context.to(dtype)) # --------------------------------------------------------------------------------------------- # MLPs # --------------------------------------------------------------------------------------------- class Raon1MiniMLP(nn.Module): """SiTU-GLU MLP: ``down(situ_glu(gate(x), up(x)))``. Used for the dense layer and the shared expert.""" def __init__(self, config: Raon1MiniConfig, intermediate_size: int) -> None: super().__init__() self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False) self.gate_beta = config.activation_situ_beta self.value_beta = config.activation_situ_linear_beta def forward(self, hidden: torch.Tensor) -> torch.Tensor: act = situ_glu( self.gate_proj(hidden), self.up_proj(hidden), self.gate_beta, self.value_beta, ) return self.down_proj(act.to(hidden.dtype)) def dequantize_mxfp8( weight: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype ) -> torch.Tensor: """E4M3 ``[N, K]`` and E8M0 bytes ``[N, K/32]`` -> ``value * 2**(scale - 127)`` in ``dtype``. The fp32 product is exact and every released weight is representable in BF16, so the cast is exact. """ rows, cols = weight.shape blocks = weight.float().view(rows, cols // 32, 32) factor = torch.exp2(_scale_bytes(scale).float() - 127.0).unsqueeze(-1) return (blocks * factor).view(rows, cols).to(dtype) def _scale_bytes(scale: torch.Tensor) -> torch.Tensor: """E8M0 scales as their raw exponent bytes, whether stored as uint8 or ``float8_e8m0fnu``.""" return scale if scale.dtype == torch.uint8 else scale.view(torch.uint8) def quantize_mxfp8(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """``[M, K]`` -> E4M3 ``[M, K]`` and E8M0 bytes ``[M, K/32]`` with the production (TE) rule. The block exponent is ``amax / 448`` rounded up to a power of two on its fp32 bits, and the values are ``x * 2**(127 - e)`` cast round-to-nearest with saturation, as vLLM serving does. """ rows, cols = hidden.shape blocks = hidden.float().view(rows, cols // 32, 32) bits = (blocks.abs().amax(-1) * (1.0 / 448.0)).view(torch.int32) exp = (bits >> 23) & 0xFF man = bits & 0x7FFFFF round_up = (man > 0) & (exp != 0xFE) & ~((exp == 0) & (man <= 0x400000)) e = (exp + round_up.int()).clamp(max=254) # 2**(127 - e) built from its bits; e == 254 needs the subnormal 2**-127. reciprocal = ((254 - e) << 23).view(torch.float32) reciprocal = torch.where(e == 254, torch.full_like(reciprocal, 2.0**-127), reciprocal) values = (blocks * reciprocal.unsqueeze(-1)).clamp(-448.0, 448.0) return values.to(torch.float8_e4m3fn).view(rows, cols), e.to(torch.uint8) def to_blocked_scales(scale: torch.Tensor) -> torch.Tensor: """``[rows, cols]`` E8M0 bytes -> cuBLAS 128x4 swizzled layout (rows padded to 128, cols to 4).""" rows, cols = scale.shape row_tiles, col_tiles = -(-rows // 128), -(-cols // 4) padded = scale.new_zeros(row_tiles * 128, col_tiles * 4) padded[:rows, :cols] = scale tiles = padded.view(row_tiles, 128, col_tiles, 4).permute(0, 2, 1, 3) return tiles.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16).flatten() def _grouped_activation_scales( scale: torch.Tensor, tokens_per_expert: torch.Tensor, cols: int ) -> torch.Tensor: """Expert-sorted activation scales -> each expert's rows padded to 128, then swizzled. ``scaled_grouped_mm`` reads group ``g``'s scales from a 128-row aligned tile range, so the rows are scattered to their padded positions without a host sync. """ groups = tokens_per_expert.shape[0] rows = scale.shape[0] padded = (tokens_per_expert + 127) // 128 * 128 padded_start = torch.cumsum(padded, 0) - padded start = torch.cumsum(tokens_per_expert, 0) - tokens_per_expert expert = torch.repeat_interleave( torch.arange(groups, device=scale.device), tokens_per_expert, output_size=rows ) dest = padded_start[expert] + torch.arange(rows, device=scale.device) - start[expert] buffer = scale.new_zeros(rows + groups * 128, cols) buffer[dest, : scale.shape[1]] = scale return to_blocked_scales(buffer).view(-1, cols) def mxfp8_grouped_mm( hidden: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, tokens_per_expert: torch.Tensor, offsets: torch.Tensor, ) -> torch.Tensor: """Expert-sorted ``[M, K]`` x E4M3 ``[E, N, K]`` (swizzled scales ``[E, *]``) -> ``[M, N]``. Activations are quantized to MXFP8 per call and both operands run through the native block-scaled tensor cores (W8A8, fp32 accumulation, output in the activation dtype). """ values, scale = quantize_mxfp8(hidden) cols = -(-scale.shape[1] // 4) * 4 act_scale = _grouped_activation_scales(scale, tokens_per_expert, cols) return F.scaled_grouped_mm( values, weight.transpose(-2, -1), act_scale.view(torch.float8_e8m0fnu), F.ScalingType.BlockWise1x32, weight_scale.view(torch.float8_e8m0fnu), F.ScalingType.BlockWise1x32, swizzle_a=F.SwizzleType.SWIZZLE_32_4_4, swizzle_b=F.SwizzleType.SWIZZLE_32_4_4, offs=offsets, output_dtype=hidden.dtype, ) def native_mxfp8_available(device: torch.device) -> bool: """Block-scaled grouped GEMM needs datacenter Blackwell (sm_100/sm_103) and torch's MX API.""" return ( device.type == "cuda" and torch.cuda.get_device_capability(device)[0] == 10 and hasattr(F, "scaled_grouped_mm") and hasattr(getattr(F, "ScalingType", None), "BlockWise1x32") and hasattr(getattr(F, "SwizzleType", None), "SWIZZLE_32_4_4") ) def routed_linear(linear: nn.Linear, hidden: torch.Tensor) -> torch.Tensor: """Routed-expert projection in the activation dtype, also when MXFP8 storage was kept.""" weight = linear.weight if weight.dtype == torch.float8_e4m3fn: weight = dequantize_mxfp8(weight, linear.weight_scale_inv, hidden.dtype) return F.linear(hidden, weight) class Raon1MiniRoutedLinear(nn.Linear): def forward(self, hidden: torch.Tensor) -> torch.Tensor: return routed_linear(self, hidden) class Raon1MiniMXFP8Linear(FP8Linear): def forward(self, hidden: torch.Tensor) -> torch.Tensor: return routed_linear(self, hidden) def _routed_projection(config: Raon1MiniConfig, in_features: int, out_features: int) -> nn.Linear: quantization = getattr(config, "quantization_config", None) if isinstance(quantization, dict): quantization = FineGrainedFP8Config.from_dict(quantization) if getattr(quantization, "quant_method", None) == "mxfp8": if not quantization.dequantize: return Raon1MiniMXFP8Linear( in_features, out_features, block_size=quantization.weight_block_size, activation_scheme=quantization.activation_scheme, scale_fmt="ue8m0", has_bias=False, ) return Raon1MiniRoutedLinear(in_features, out_features, bias=False) class Raon1MiniExpert(nn.Module): """One routed SiTU-GLU expert: ``down_proj(situ_glu(gate_proj(x), up_proj(x)))``.""" def __init__(self, config: Raon1MiniConfig) -> None: super().__init__() self.gate_proj = _routed_projection(config, config.hidden_size, config.moe_intermediate_size) self.up_proj = _routed_projection(config, config.hidden_size, config.moe_intermediate_size) self.down_proj = _routed_projection(config, config.moe_intermediate_size, config.hidden_size) self.gate_beta = config.activation_situ_beta self.value_beta = config.activation_situ_linear_beta def forward(self, hidden: torch.Tensor) -> torch.Tensor: act = situ_glu( self.gate_proj(hidden), self.up_proj(hidden), self.gate_beta, self.value_beta, ) return self.down_proj(act.to(hidden.dtype)) class Raon1MiniRouter(nn.Module): """Sigmoid router with the aux-loss-free expert bias (``e_score_correction_bias``) used for selection only.""" def __init__(self, config: Raon1MiniConfig) -> None: super().__init__() self.weight = nn.Parameter(torch.zeros(config.num_experts, config.hidden_size)) self.e_score_correction_bias = nn.Parameter( torch.zeros(config.num_experts, dtype=torch.float32) ) self.top_k = config.num_experts_per_token self.renormalize = config.moe_renormalize self.scaling = config.routed_scaling_factor if config.moe_router_activation_func != "sigmoid": raise NotImplementedError("raon-1-mini uses the sigmoid router.") def forward(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """``[N, C]`` -> (top-k expert indices ``[N, k]``, fp32 routing weights ``[N, k]``).""" logits = F.linear(hidden.float(), self.weight.float()) scores = torch.sigmoid(logits) _, top = torch.topk( scores + self.e_score_correction_bias.float(), self.top_k, dim=-1, sorted=True, ) chosen = torch.gather(scores, 1, top) if self.renormalize: chosen = chosen / (chosen.sum(-1, keepdim=True) + 1e-20) return top, chosen * self.scaling class Raon1MiniMoE(nn.Module): """Routed experts plus the always-on shared expert. With MXFP8 storage on sm_100/sm_103 the routed experts run as two native grouped MXFP8 GEMMs; everywhere else (or after ``set_mxfp8_moe_backend("dequantize")``) each used expert is expanded to the activation dtype per call, so FP8 storage is kept on every device. """ def __init__(self, config: Raon1MiniConfig) -> None: super().__init__() self.num_experts = config.num_experts self.top_k = config.num_experts_per_token self.intermediate_size = config.moe_intermediate_size self.gate_beta = config.activation_situ_beta self.value_beta = config.activation_situ_linear_beta self.gate = Raon1MiniRouter(config) self.experts = nn.ModuleList( [Raon1MiniExpert(config) for _ in range(config.num_experts)] ) shared_size = config.moe_intermediate_size * config.num_shared_experts self.shared_experts = ( Raon1MiniMLP(config, shared_size) if config.num_shared_experts else None ) # Toggled model-wide by ``Raon1MiniPreTrainedModel.set_mxfp8_moe_backend``. self.use_native_mxfp8 = True self._packed: tuple[torch.Tensor, ...] | None = None def _native_weights(self, hidden: torch.Tensor) -> tuple[torch.Tensor, ...] | None: """Stacked E4M3 ``gate_up [E, 2I, H]`` / ``down [E, H, I]`` plus swizzled scales, or None. The per-expert parameters are rebound to views of the stacks on first use, so the layer keeps one copy of its weights and the checkpoint names are unchanged. """ if not self.use_native_mxfp8: return None if self._packed is not None: return self._packed first = self.experts[0].gate_proj.weight if ( first.dtype != torch.float8_e4m3fn or hidden.dtype != torch.bfloat16 or first.device != hidden.device or not native_mxfp8_available(hidden.device) or any(p.device != first.device for p in self.experts.parameters()) ): return None size = self.intermediate_size with torch.no_grad(): gate_up = torch.stack( [torch.cat([e.gate_proj.weight, e.up_proj.weight]) for e in self.experts] ) gate_up_scale = torch.stack( [torch.cat([_scale_bytes(e.gate_proj.weight_scale_inv), _scale_bytes(e.up_proj.weight_scale_inv)]) for e in self.experts] ) down = torch.stack([e.down_proj.weight for e in self.experts]) down_scale = torch.stack([_scale_bytes(e.down_proj.weight_scale_inv) for e in self.experts]) for idx, expert in enumerate(self.experts): expert.gate_proj.weight.data = gate_up[idx, :size] expert.up_proj.weight.data = gate_up[idx, size:] expert.down_proj.weight.data = down[idx] expert.gate_proj.weight_scale_inv.data = gate_up_scale[idx, :size] expert.up_proj.weight_scale_inv.data = gate_up_scale[idx, size:] expert.down_proj.weight_scale_inv.data = down_scale[idx] gate_up_blocked = torch.stack([to_blocked_scales(s) for s in gate_up_scale]) down_blocked = torch.stack([to_blocked_scales(s) for s in down_scale]) self._packed = (gate_up, gate_up_blocked, down, down_blocked) return self._packed def _native_routed( self, flat: torch.Tensor, top: torch.Tensor, probs: torch.Tensor, packed: tuple[torch.Tensor, ...], ) -> torch.Tensor: """Sort the token-expert pairs by expert, run up/gate and down as grouped MXFP8 GEMMs.""" gate_up, gate_up_blocked, down, down_blocked = packed experts = top.flatten() order = torch.argsort(experts, stable=True) token = order // self.top_k tokens_per_expert = torch.bincount(experts, minlength=self.num_experts).int() offsets = torch.cumsum(tokens_per_expert, 0, dtype=torch.int32) hidden = mxfp8_grouped_mm( flat.index_select(0, token), gate_up, gate_up_blocked, tokens_per_expert, offsets ) act = situ_glu( hidden[:, : self.intermediate_size], hidden[:, self.intermediate_size :], self.gate_beta, self.value_beta, ).to(flat.dtype) expert_out = mxfp8_grouped_mm(act, down, down_blocked, tokens_per_expert, offsets) # Back to (token, slot) order and a fixed-order sum over the k slots, so the result is # deterministic (an index_add_ with k duplicate rows per token would race on atomics). weighted = torch.empty(experts.shape[0], flat.shape[1], dtype=torch.float32, device=flat.device) weighted[order] = expert_out.float() * probs.flatten().index_select(0, order).unsqueeze(-1) return weighted.view(flat.shape[0], self.top_k, flat.shape[1]).sum(1) def forward(self, hidden: torch.Tensor) -> torch.Tensor: batch, seq, dim = hidden.shape flat = hidden.reshape(-1, dim) top, probs = self.gate(flat) packed = self._native_weights(flat) if flat.shape[0] else None if packed is not None: out = self._native_routed(flat, top, probs, packed).to(hidden.dtype) if self.shared_experts is not None: out = out + self.shared_experts(flat) return out.reshape(batch, seq, dim) out = torch.zeros(flat.shape[0], dim, dtype=torch.float32, device=flat.device) hit = torch.bincount(top.flatten(), minlength=self.num_experts) for expert_idx in hit.nonzero().flatten().tolist(): rows, slot = torch.where(top == expert_idx) expert_out = self.experts[expert_idx](flat.index_select(0, rows)) out.index_add_( 0, rows, expert_out.float() * probs[rows, slot].unsqueeze(-1) ) out = out.to(hidden.dtype) if self.shared_experts is not None: out = out + self.shared_experts(flat) return out.reshape(batch, seq, dim) # --------------------------------------------------------------------------------------------- # decoder layer # --------------------------------------------------------------------------------------------- class Raon1MiniDecoderLayer(nn.Module): """One Raon_mHC-wired layer: attention site then MLP site, each with its own selector.""" def __init__(self, config: Raon1MiniConfig, layer_idx: int) -> None: super().__init__() self.layer_idx = layer_idx self.is_rda = config.is_rda_layer(layer_idx) self.is_moe = config.is_moe_layer(layer_idx) self.layer_type = "linear_attention" if self.is_rda else "full_attention" self.input_layernorm = Raon1MiniRMSNorm(config.hidden_size, config.rms_norm_eps) self.post_attention_layernorm = Raon1MiniRMSNorm( config.hidden_size, config.rms_norm_eps ) self.self_attn = Raon1MiniRDA(config) if self.is_rda else Raon1MiniMLA(config) self.self_attn_raon_mhc = RaonMHC(config) self.mlp_raon_mhc = RaonMHC(config) if self.is_moe: self.moe = Raon1MiniMoE(config) else: self.mlp = Raon1MiniMLP(config, config.intermediate_size) def forward( self, stream: torch.Tensor, *, attention_mask: torch.Tensor | None, token_mask: torch.Tensor | None, cache: Raon1MiniDynamicCache | None, ) -> torch.Tensor: batch, seq, n, dim = stream.shape dtype = stream.dtype rows = stream.reshape(batch * seq, n, dim) h_pre, h_post, h_res = self.self_attn_raon_mhc.route(stream) sf = rows.float() branch_in = self.input_layernorm( RaonMHC.aggregate(sf, h_pre, dtype) ).reshape(batch, seq, dim) if self.is_rda: branch_out = self.self_attn(branch_in, cache, self.layer_idx, token_mask) else: branch_out = self.self_attn( branch_in, cache, self.layer_idx, attention_mask ) rows = RaonMHC.mix( sf, branch_out.reshape(batch * seq, dim), h_post, h_res, dtype ) h_pre, h_post, h_res = self.mlp_raon_mhc.route(rows.reshape(batch, seq, n, dim)) sf = rows.float() branch_in = self.post_attention_layernorm( RaonMHC.aggregate(sf, h_pre, dtype) ).reshape(batch, seq, dim) branch_out = ( self.moe(branch_in) if self.is_moe else self.mlp(branch_in) ) rows = RaonMHC.mix( sf, branch_out.reshape(batch * seq, dim), h_post, h_res, dtype ) return rows.reshape(batch, seq, n, dim) # --------------------------------------------------------------------------------------------- # cache # --------------------------------------------------------------------------------------------- class Raon1MiniDynamicCache: """Hybrid cache: RDA conv windows + fp32 recurrent states, MLA key/value tensors. Per layer: ``conv_states[i]`` is ``[B, 3*proj, K-1]`` of RAW (pre-conv) q|k|v rows, ``recurrent_states[i]`` is ``[B, H, V, K]`` fp32; ``key_cache[i]`` / ``value_cache[i]`` are ``[B, H, T, 192]`` / ``[B, H, T, 128]``. Implements the hybrid-cache interface ``generate`` uses. """ is_compileable = False def __init__(self, config: Raon1MiniConfig) -> None: self.layer_types = config.layer_types self.transformer_layers = [ i for i, t in enumerate(self.layer_types) if t == "full_attention" ] self.last_linear_layer = ( len(self.layer_types) - 1 - self.layer_types[::-1].index("linear_attention") ) num = config.num_hidden_layers self.conv_states: list[torch.Tensor | None] = [None] * num self.recurrent_states: list[torch.Tensor | None] = [None] * num self.key_cache: list[torch.Tensor | None] = [None] * num self.value_cache: list[torch.Tensor | None] = [None] * num def __len__(self) -> int: return len(self.layer_types) def update( self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: dict | None = None, ): if self.key_cache[layer_idx] is None: self.key_cache[layer_idx] = key_states self.value_cache[layer_idx] = value_states else: self.key_cache[layer_idx] = torch.cat( [self.key_cache[layer_idx], key_states], dim=2 ) self.value_cache[layer_idx] = torch.cat( [self.value_cache[layer_idx], value_states], dim=2 ) return self.key_cache[layer_idx], self.value_cache[layer_idx] def reorder_cache(self, beam_idx: torch.LongTensor) -> None: for i in range(len(self.layer_types)): for store in ( self.key_cache, self.value_cache, self.conv_states, self.recurrent_states, ): if store[i] is not None: store[i] = store[i].index_select(0, beam_idx.to(store[i].device)) def get_seq_length(self, layer_idx: int | None = 0) -> int: layer_idx = ( self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx ) if self.key_cache[layer_idx] is None: return 0 return self.key_cache[layer_idx].shape[-2] def get_mask_sizes( self, query_length: int, layer_idx: int ) -> tuple[int, int]: return query_length + self.get_seq_length(layer_idx), 0 def get_query_offset(self, layer_idx: int = 0) -> int: return self.get_seq_length(layer_idx) @property def has_previous_state(self) -> bool: return self.conv_states[self.last_linear_layer] is not None # --------------------------------------------------------------------------------------------- # models # --------------------------------------------------------------------------------------------- @contextlib.contextmanager def _keep_mxfp8_storage(): """Keep MXFP8 storage on devices where Transformers would force a load-time dequantize. The fine-grained FP8 quantizer switches to ``dequantize=True`` without a GPU or below compute capability 8.9 because its own kernels need FP8 tensor cores. Raon never calls those kernels (``Raon1MiniMoE`` reads the E4M3 weights and scales itself and expands them per call where native MXFP8 is missing), so for a pre-quantized ``mxfp8`` load the override is only active while a Raon model loads; ``dequantize=True`` passed by the caller is still honoured. """ try: from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer except ImportError: yield return original = FineGrainedFP8HfQuantizer.validate_environment def validate_environment(self, *args: Any, **kwargs: Any) -> None: config = self.quantization_config if self.pre_quantized and config.quant_method == "mxfp8" and not config.dequantize: from transformers.utils import is_accelerate_available if not is_accelerate_available(): raise ImportError("Loading an MXFP8 model requires accelerate (`pip install accelerate`)") return original(self, *args, **kwargs) FineGrainedFP8HfQuantizer.validate_environment = validate_environment try: yield finally: FineGrainedFP8HfQuantizer.validate_environment = original class Raon1MiniPreTrainedModel(PreTrainedModel): config: Raon1MiniConfig config_class = Raon1MiniConfig base_model_prefix = "model" supports_gradient_checkpointing = False _no_split_modules: ClassVar[list[str]] = ["Raon1MiniDecoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_sdpa = True _supports_flash_attn = False _is_stateful = True # The checkpoint stores these in fp32 and the math reads them in fp32; never downcast them on load. _keep_in_fp32_modules: ClassVar[list[str]] = [ "A_log", "dt_bias", "e_score_correction_bias", "eta", ] _keep_in_fp32_modules_strict: ClassVar[list[str]] = [ "A_log", "dt_bias", "e_score_correction_bias", "eta", ] def __init__(self, config: Raon1MiniConfig, *inputs: Any, **kwargs: Any) -> None: super().__init__(config, *inputs, **kwargs) quantization = getattr(config, "quantization_config", None) if isinstance(quantization, dict) and quantization.get("quant_method") == "mxfp8": from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING # Without this loader the E4M3 weights would be cast without their scales. if "mxfp8" not in AUTO_QUANTIZER_MAPPING: raise ImportError( "This checkpoint stores routed experts as MXFP8; use a Transformers " "release with the mxfp8 quantizer (validated with 5.16.1)." ) self._mxfp8_storage_checked = False @classmethod def from_pretrained(cls, *args: Any, **kwargs: Any): with _keep_mxfp8_storage(): return super().from_pretrained(*args, **kwargs) def set_mxfp8_moe_backend(self, backend: str) -> None: """``"auto"``: native grouped MXFP8 GEMM where supported; ``"dequantize"``: per-call BF16.""" if backend not in ("auto", "dequantize"): raise ValueError(f"backend must be 'auto' or 'dequantize', got {backend!r}.") for module in self.modules(): if isinstance(module, Raon1MiniMoE): module.use_native_mxfp8 = backend == "auto" def _check_mxfp8_storage(self) -> None: """Reject MXFP8 storage outside the routed experts (a ``dequantize=False`` skip-list error).""" if self._mxfp8_storage_checked: return for name, module in self.named_modules(): weight = getattr(module, "weight", None) if ( isinstance(weight, torch.Tensor) and weight.dtype == torch.float8_e4m3fn and ".moe.experts." not in f".{name}" ): raise ValueError( f"{name} was converted to FP8, but only routed experts are MXFP8 in this " "checkpoint. Pass the repository's quantization_config (its ignored_layers) " "together with dequantize=False." ) self._mxfp8_storage_checked = True @torch.no_grad() def _init_weights(self, module: nn.Module) -> None: # Go through transformers.initialization: those initialisers are no-ops for parameters that # from_pretrained already loaded, so a plain in-place normal_() here would clobber checkpoints. super()._init_weights(module) std = self.config.initializer_range if isinstance(module, Raon1MiniRMSNorm): init.ones_(module.weight) elif isinstance(module, Raon1MiniShortConv): init.normal_(module.weight, mean=0.0, std=std) elif isinstance(module, Raon1MiniRouter): init.normal_(module.weight, mean=0.0, std=std) init.zeros_(module.e_score_correction_bias) elif isinstance(module, Raon1MiniRDA): init.zeros_(module.A_log) init.zeros_(module.dt_bias) elif isinstance(module, RaonMHCStreamMix): init.normal_(module.phi_z, mean=0.0, std=std) init.ones_(module.alpha_z) init.zeros_(module.b_z) init.zeros_(module.eta) elif isinstance(module, RaonMHC): init.normal_(module.phi_pre, mean=0.0, std=std) init.normal_(module.phi_post, mean=0.0, std=std) init.ones_(module.alpha_pre) init.ones_(module.alpha_post) init.zeros_(module.b_pre) init.zeros_(module.b_post) class Raon1MiniModel(Raon1MiniPreTrainedModel): """Embedding -> 4-stream expansion -> decoder layers -> mean reduction -> final norm.""" def __init__(self, config: Raon1MiniConfig) -> None: super().__init__(config) self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, config.pad_token_id ) self.layers = nn.ModuleList( [Raon1MiniDecoderLayer(config, i) for i in range(config.num_hidden_layers)] ) self.norm = Raon1MiniRMSNorm(config.hidden_size, config.rms_norm_eps) self.post_init() def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens def set_input_embeddings(self, value: nn.Embedding) -> None: self.embed_tokens = value def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Raon1MiniDynamicCache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, **kwargs: Any, ) -> BaseModelOutputWithPast | tuple: use_cache = self.config.use_cache if use_cache is None else use_cache output_hidden_states = ( self.config.output_hidden_states if output_hidden_states is None else output_hidden_states ) return_dict = self.config.return_dict if return_dict is None else return_dict self._check_mxfp8_storage() if (input_ids is None) == (inputs_embeds is None): raise ValueError("Pass exactly one of input_ids or inputs_embeds.") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if use_cache and past_key_values is None: past_key_values = Raon1MiniDynamicCache(self.config) causal_mask = create_causal_mask( config=self.config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=position_ids, ) token_mask = None if ( attention_mask is not None and attention_mask.dim() == 2 and not bool(attention_mask.all()) ): token_mask = attention_mask[:, -inputs_embeds.shape[1]:] n = self.config.raon_mhc_streams stream = inputs_embeds.unsqueeze(2).expand(-1, -1, n, -1).contiguous() all_hidden_states = () if output_hidden_states else None for layer in self.layers: if output_hidden_states: # Expose [batch, tokens, hidden] representations by reducing the Raon_mHC streams, # as at the model output. Intermediate states are not final-normalized. all_hidden_states += (stream.float().mean(2).to(inputs_embeds.dtype),) stream = layer( stream, attention_mask=causal_mask, token_mask=token_mask, cache=past_key_values, ) hidden = self.norm(stream.float().mean(2).to(inputs_embeds.dtype)) if output_hidden_states: all_hidden_states += (hidden,) output = BaseModelOutputWithPast( last_hidden_state=hidden, past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, ) return output if return_dict else output.to_tuple() class Raon1MiniPresencePenaltyLogitsProcessor: """Subtract ``penalty`` from the logits of every token generated so far. Matches vLLM's presence penalty: prompt tokens are not penalized and a token is penalized once however often it appeared. ``prompt_length`` is the padded prompt width, the same for every row of a batch. """ def __init__(self, penalty: float, prompt_length: int) -> None: self.penalty = penalty self.prompt_length = prompt_length def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: generated = input_ids[:, self.prompt_length:] if generated.shape[1] == 0: return scores seen = torch.zeros_like(scores, dtype=torch.bool).scatter_(1, generated, True) return scores - self.penalty * seen.to(scores.dtype) class Raon1MiniForCausalLM(Raon1MiniPreTrainedModel, GenerationMixin): """raon-1-mini causal language model with the untied ``lm_head``.""" _tied_weights_keys = None @classmethod def _supports_default_dynamic_cache(cls) -> bool: # RDA conv windows and recurrent state require the Raon hybrid cache. return False def generate(self, *args, presence_penalty: float | None = None, **kwargs): # GenerationConfig rejects keys it does not define, so carry this one in ourselves. if presence_penalty is not None: config = copy.deepcopy(kwargs.get("generation_config") or self.generation_config) config.presence_penalty = presence_penalty kwargs["generation_config"] = config return super().generate(*args, **kwargs) def _get_logits_processor(self, generation_config, input_ids_seq_length=None, *args, **kwargs): # Transformers has no presence penalty; generation_config.json sets one (1.5) # as the serving default, so apply it ahead of temperature/top-k/top-p as vLLM does. processors = super()._get_logits_processor(generation_config, input_ids_seq_length, *args, **kwargs) penalty = getattr(generation_config, "presence_penalty", None) if penalty: processors.insert(0, Raon1MiniPresencePenaltyLogitsProcessor(float(penalty), input_ids_seq_length)) return processors def __init__(self, config: Raon1MiniConfig) -> None: super().__init__(config) self.model = Raon1MiniModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.embed_tokens = value def get_output_embeddings(self) -> nn.Linear: return self.lm_head def set_output_embeddings(self, value: nn.Linear) -> None: self.lm_head = value def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Raon1MiniDynamicCache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, output_hidden_states: bool | None = None, return_dict: bool | None = None, **kwargs: Any, ) -> CausalLMOutputWithPast | tuple: return_dict = self.config.return_dict if return_dict is None else return_dict self._check_mxfp8_storage() outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_hidden_states=output_hidden_states, return_dict=True, **kwargs, ) hidden = outputs.last_hidden_state slice_indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) logits = self.lm_head(hidden[:, slice_indices, :]) loss = None if labels is not None: loss = self.loss_function( logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs, ) output = CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, ) return output if return_dict else output.to_tuple() __all__ = [ "Raon1MiniDecoderLayer", "Raon1MiniDynamicCache", "Raon1MiniForCausalLM", "Raon1MiniRDA", "RaonMHC", "Raon1MiniMLA", "Raon1MiniMoE", "Raon1MiniModel", "Raon1MiniPreTrainedModel", "torch_recurrent_rda", ]