| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Standalone HF inference for Clover's causal-prefix/block-bidirectional MoE. |
| |
| The architecture fixes softmax-blocksize GDN with independent right convolution, |
| a causal prefix, bidirectional block preview, and same-position logits. |
| Checkpoint names remain unchanged. Preview never mutates a clean cache; completed |
| blocks are replayed causally before their state is persisted. The implementation |
| uses standalone PyTorch operations defined in this file. |
| """ |
|
|
| from copy import deepcopy |
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel, GenerationMixin, GenerationConfig |
| from transformers.modeling_outputs import ( |
| BaseModelOutputWithPast, |
| CausalLMOutputWithPast, |
| ) |
| from transformers.generation.utils import GenerateDecoderOnlyOutput |
|
|
| from .configuration_clover_moe import CloverMoeConfig |
| from .sampling_clover_moe import ( |
| PENALTY_DEFAULTS, |
| CloverPenaltyState, |
| predict_with_confidence, |
| validate_penalties, |
| ) |
|
|
|
|
| |
| def bd_gdn_conv1d( |
| mixed_qkv: torch.Tensor, |
| *, |
| clean_conv_states: torch.Tensor, |
| cache_indices: torch.Tensor, |
| conv_weights: torch.Tensor, |
| block_size: int, |
| active_valid_mask: torch.Tensor, |
| conv_right_weight: torch.Tensor, |
| ) -> torch.Tensor: |
| """Apply Clover independent right convolution and SiLU without updating cache. |
| |
| ``clean_conv_states`` contains the last ``kernel_size - 1`` committed QKV |
| projections. Preview tokens are appended only in a temporary tensor. |
| """ |
| num_tokens, channels = mixed_qkv.shape |
| batch_size = cache_indices.numel() |
| if num_tokens != batch_size * block_size: |
| raise ValueError( |
| "BD-GDN preview requires one fixed-size block per request: " |
| f"{num_tokens=} != {batch_size=} * {block_size=}" |
| ) |
| if active_valid_mask.numel() != num_tokens: |
| raise ValueError(f"{active_valid_mask.numel()=} does not match {num_tokens=}") |
|
|
| safe_cache_indices = cache_indices.to(torch.long).clamp_min(0) |
| history = clean_conv_states.index_select(0, safe_cache_indices) |
| if history.shape[1] != channels: |
| raise ValueError( |
| f"BD-GDN conv cache has {history.shape[1]} channels, expected {channels}" |
| ) |
|
|
| valid = ( |
| active_valid_mask.reshape(batch_size, block_size) |
| & (cache_indices >= 0)[:, None] |
| ) |
| active_requests = valid.any(dim=1) |
| history = history * active_requests[:, None, None].to(history.dtype) |
| current = mixed_qkv.reshape(batch_size, block_size, channels) |
| current = current * valid[:, :, None].to(current.dtype) |
| convolution_input = torch.cat( |
| (history, current.transpose(1, 2)), |
| dim=-1, |
| ) |
|
|
| weight = conv_weights.reshape(channels, 1, -1) |
| output = F.conv1d( |
| convolution_input.float(), |
| weight.float(), |
| groups=channels, |
| ) |
| width = weight.shape[-1] |
| expected_shape = (channels, 1, width - 1) |
| if tuple(conv_right_weight.shape) != expected_shape: |
| raise ValueError(f"Clover right weight shape must be {expected_shape}") |
| for distance in range(1, min(width, block_size)): |
| output[:, :, :-distance] += ( |
| current[:, distance:, :].transpose(1, 2).float() |
| * conv_right_weight[:, 0, distance - 1].float()[None, :, None] |
| ) |
| output = F.silu(output) |
|
|
| output = output.transpose(1, 2) |
| output = output * valid[:, :, None].to(output.dtype) |
| return output.reshape(num_tokens, channels).to(mixed_qkv.dtype) |
|
|
|
|
| def bd_gdn_local_memory( |
| query: torch.Tensor, |
| key: torch.Tensor, |
| value: torch.Tensor, |
| g: torch.Tensor, |
| *, |
| clean_ssm_states: torch.Tensor, |
| cache_indices: torch.Tensor, |
| active_valid_mask: torch.Tensor, |
| scale: float, |
| softmax_logit_scale: torch.Tensor, |
| softmax_history_bias: torch.Tensor, |
| softmax_history_c: torch.Tensor, |
| ) -> torch.Tensor: |
| """Clover bidirectional softmax preview with a read-only clean-state bucket. |
| |
| Normalized q/k determine block attention logits without beta. The history |
| bucket logit is bias + history_c * log(valid_block_size / 32), and its |
| value is the gate-decayed recurrent readout from clean [B, H, V, K] state. |
| """ |
| if query.ndim != 4 or query.shape[0] != 1: |
| raise ValueError(f"Expected query shape [1, T, H, K], got {query.shape}") |
| num_tokens = query.shape[1] |
| batch_size = cache_indices.numel() |
| if num_tokens % batch_size != 0: |
| raise ValueError( |
| f"BD-GDN token count {num_tokens} is not divisible by batch {batch_size}" |
| ) |
| block_size = num_tokens // batch_size |
| if active_valid_mask.numel() != num_tokens: |
| raise ValueError(f"{active_valid_mask.numel()=} does not match {num_tokens=}") |
|
|
| num_value_heads = value.shape[2] |
| if num_value_heads % query.shape[2] or num_value_heads % key.shape[2]: |
| raise ValueError("BD-GDN value heads must be divisible by query and key heads") |
| query = query.repeat_interleave(num_value_heads // query.shape[2], dim=2) |
| key = key.repeat_interleave(num_value_heads // key.shape[2], dim=2) |
|
|
| valid = active_valid_mask.reshape(batch_size, block_size) |
| valid_float = valid[:, :, None].float() |
| query_float = query.reshape( |
| batch_size, block_size, num_value_heads, query.shape[-1] |
| ).float() |
| key_float = key.reshape( |
| batch_size, block_size, num_value_heads, key.shape[-1] |
| ).float() |
| value_float = value.reshape( |
| batch_size, block_size, num_value_heads, value.shape[-1] |
| ).float() |
| query_float = query_float * torch.rsqrt( |
| query_float.square().sum(dim=-1, keepdim=True) + 1e-6 |
| ) |
| key_float = key_float * torch.rsqrt( |
| key_float.square().sum(dim=-1, keepdim=True) + 1e-6 |
| ) |
|
|
| safe_cache_indices = cache_indices.to(torch.long).clamp_min(0) |
| clean_state = clean_ssm_states.index_select(0, safe_cache_indices).float() |
| active_requests = valid.any(dim=1) |
| clean_state = clean_state * active_requests[:, None, None, None].float() |
|
|
| scaled_query = query_float * scale |
| clean_output = torch.einsum( |
| "blhk,bhvk->blhv", |
| scaled_query, |
| clean_state, |
| ) |
| decay = g.reshape(batch_size, block_size, num_value_heads).float().exp() |
| history = clean_output * decay[..., None] |
|
|
| for name, scalar in ( |
| ("logit_scale", softmax_logit_scale), |
| ("history_bias", softmax_history_bias), |
| ("history_c", softmax_history_c), |
| ): |
| if scalar.numel() != num_value_heads: |
| raise ValueError(f"Clover {name} must have one entry per value head") |
| logit_scale = softmax_logit_scale.reshape(-1).float().to(query_float.device) |
| history_bias = softmax_history_bias.reshape(-1).float().to(query_float.device) |
| history_c = softmax_history_c.reshape(-1).float().to(query_float.device) |
| |
| logits = torch.einsum("blhk,bshk->bhls", query_float, key_float) |
| logits = logits * logit_scale.view(1, -1, 1, 1) |
| logits = logits.masked_fill(~valid[:, None, None, :], float("-inf")) |
| history_logit = history_bias.view(1, -1, 1, 1).expand(batch_size, -1, block_size, 1) |
| valid_block_size = valid.sum(dim=-1).clamp_min(1).float() |
| block_size_log_ratio = torch.log(valid_block_size / 32.0) |
| history_logit = history_logit + ( |
| history_c.view(1, -1, 1, 1) * block_size_log_ratio.view(batch_size, 1, 1, 1) |
| ) |
| weights = torch.softmax(torch.cat((logits, history_logit), dim=-1), dim=-1) |
| local_output = torch.einsum( |
| "bhls,bshv->blhv", |
| weights[..., :-1], |
| value_float * valid_float[..., None], |
| ) |
| history_weight = weights[..., -1].permute(0, 2, 1) |
| output = history_weight[..., None] * history + local_output |
| output = output * valid_float[..., None] |
| return output.to(query.dtype).reshape( |
| 1, num_tokens, num_value_heads, value.shape[-1] |
| ) |
|
|
|
|
| @dataclass |
| class CloverCache: |
| layers: list |
| attention_mask: torch.Tensor |
|
|
| def get_seq_length(self, layer_idx=0): |
| return self.attention_mask.shape[-1] |
|
|
|
|
| class CloverRMSNorm(nn.Module): |
| def __init__(self, size, eps, zero_centered=True): |
| super().__init__() |
| self.weight = nn.Parameter( |
| torch.zeros(size) if zero_centered else torch.ones(size) |
| ) |
| self.eps, self.zero_centered = eps, zero_centered |
|
|
| def forward(self, x, gate=None): |
| y = x.float() * torch.rsqrt( |
| x.float().square().mean(-1, keepdim=True) + self.eps |
| ) |
| weight = 1 + self.weight.float() if self.zero_centered else self.weight.float() |
| y = y * weight |
| if gate is not None: |
| y = y * F.silu(gate.float()) |
| return y.to(x.dtype) |
|
|
|
|
| def recurrent_commit_torch(q, k, v, g, beta, state, valid): |
| h = v.shape[2] |
| q = q.float() * torch.rsqrt(q.float().square().sum(-1, keepdim=True) + 1e-6) |
| k = k.float() * torch.rsqrt(k.float().square().sum(-1, keepdim=True) + 1e-6) |
| q = q.repeat_interleave(h // q.shape[2], 2) / q.shape[-1] ** 0.5 |
| k = k.repeat_interleave(h // k.shape[2], 2) |
| state = state.float() |
| out = [] |
| for t in range(v.shape[1]): |
| decayed = state * g[:, t].float().exp()[..., None, None] |
| remembered = torch.einsum("bhvk,bhk->bhv", decayed, k[:, t]) |
| delta = (v[:, t].float() - remembered) * beta[:, t].float()[..., None] |
| updated = decayed + delta[..., None] * k[:, t, :, None, :] |
| state = torch.where(valid[:, t, None, None, None], updated, state) |
| y = torch.einsum("bhvk,bhk->bhv", state, q[:, t]) |
| out.append(y * valid[:, t, None, None]) |
| return torch.stack(out, dim=1).to(v.dtype), state |
|
|
|
|
| class CloverGatedDeltaNet(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.config = c |
| self.kh, self.vh = c.linear_num_key_heads, c.linear_num_value_heads |
| self.kd, self.vd = c.linear_key_head_dim, c.linear_value_head_dim |
| self.channels = 2 * self.kh * self.kd + self.vh * self.vd |
| self.conv1d = nn.Conv1d( |
| self.channels, |
| self.channels, |
| c.linear_conv_kernel_dim, |
| groups=self.channels, |
| bias=False, |
| ) |
| self.conv_right_weight = nn.Parameter( |
| torch.zeros(self.channels, 1, c.linear_conv_kernel_dim - 1) |
| ) |
| self.in_proj_qkv = nn.Linear(c.hidden_size, self.channels, bias=False) |
| self.in_proj_z = nn.Linear(c.hidden_size, self.vh * self.vd, bias=False) |
| self.in_proj_a = nn.Linear(c.hidden_size, self.vh, bias=False) |
| self.in_proj_b = nn.Linear(c.hidden_size, self.vh, bias=False) |
| self.A_log = nn.Parameter(torch.zeros(self.vh)) |
| self.dt_bias = nn.Parameter(torch.ones(self.vh)) |
| self.bd_softmax_logit_scale = nn.Parameter(torch.full((self.vh,), 10.0)) |
| self.bd_softmax_history_bias = nn.Parameter(torch.full((self.vh,), 5.0)) |
| self.bd_softmax_history_c = nn.Parameter(torch.full((self.vh,), 1.0)) |
| self.norm = CloverRMSNorm(self.vd, c.rms_norm_eps, zero_centered=False) |
| self.out_proj = nn.Linear(self.vh * self.vd, c.hidden_size, bias=False) |
|
|
| def forward(self, x, cache, valid, preview): |
| b, t, _ = x.shape |
| x = x * valid[..., None] |
| raw = self.in_proj_qkv(x) |
| width = self.config.linear_conv_kernel_dim |
| if cache is None: |
| conv_state = raw.new_zeros((b, self.channels, width - 1)) |
| ssm_state = torch.zeros( |
| (b, self.vh, self.vd, self.kd), device=x.device, dtype=torch.float32 |
| ) |
| else: |
| conv_state, ssm_state = cache["conv"], cache["ssm"] |
| if preview: |
| flat = raw.reshape(-1, self.channels).contiguous() |
| kw = dict( |
| clean_conv_states=conv_state, |
| cache_indices=torch.arange(b, device=x.device), |
| conv_weights=self.conv1d.weight, |
| block_size=t, |
| active_valid_mask=valid.reshape(-1), |
| conv_right_weight=self.conv_right_weight, |
| ) |
| mixed = bd_gdn_conv1d(flat, **kw).view(b, t, self.channels) |
| next_conv = conv_state |
| else: |
| if bool(valid.all()): |
| joined = torch.cat((conv_state, raw.transpose(1, 2)), -1) |
| mixed = causal_conv_silu(joined, self.conv1d.weight).transpose(1, 2) |
| next_conv = ( |
| joined[..., -(width - 1) :] if width > 1 else joined[..., :0] |
| ).clone() |
| else: |
| history, pieces = conv_state, [] |
| for i in range(t): |
| window = torch.cat((history, raw[:, i, :, None]), -1) |
| y = F.silu((window * self.conv1d.weight[:, 0]).float().sum(-1)).to( |
| raw.dtype |
| ) |
| pieces.append(y * valid[:, i, None]) |
| history = torch.where( |
| valid[:, i, None, None], window[..., 1:], history |
| ) |
| mixed, next_conv = torch.stack(pieces, 1), history.clone() |
| q, k, v = mixed.split( |
| (self.kh * self.kd, self.kh * self.kd, self.vh * self.vd), -1 |
| ) |
| q, k, v = ( |
| q.reshape(b, t, self.kh, self.kd), |
| k.reshape(b, t, self.kh, self.kd), |
| v.reshape(b, t, self.vh, self.vd), |
| ) |
| g = -self.A_log.float().exp() * F.softplus( |
| self.in_proj_a(x).float() + self.dt_bias.float() |
| ) |
| if preview: |
| qf = q.reshape(1, b * t, self.kh, self.kd) |
| kf = k.reshape(1, b * t, self.kh, self.kd) |
| vf = v.reshape(1, b * t, self.vh, self.vd) |
| kw = dict( |
| clean_ssm_states=ssm_state, |
| cache_indices=torch.arange(b, device=x.device), |
| active_valid_mask=valid.reshape(-1), |
| scale=self.kd**-0.5, |
| softmax_logit_scale=self.bd_softmax_logit_scale, |
| softmax_history_bias=self.bd_softmax_history_bias, |
| softmax_history_c=self.bd_softmax_history_c, |
| ) |
| y = bd_gdn_local_memory( |
| qf, |
| kf, |
| vf, |
| g.reshape(1, b * t, self.vh), |
| **kw, |
| ).view(b, t, self.vh, self.vd) |
| next_ssm = ssm_state |
| else: |
| beta = self.in_proj_b(x).sigmoid() |
| y, next_ssm = recurrent_commit_torch(q, k, v, g, beta, ssm_state, valid) |
| |
| z = self.in_proj_z(x).reshape(b, t, self.vh, self.vd) |
| y = self.norm(y, z) |
| y = self.out_proj(y.reshape(b, t, -1)) |
| return y, {"conv": next_conv, "ssm": next_ssm} |
|
|
|
|
| class CloverAttention(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.config = c |
| h, d, kv = c.num_attention_heads, c.head_dim, c.num_key_value_heads |
| self.q_proj = nn.Linear(c.hidden_size, h * d * 2, bias=c.attention_bias) |
| self.k_proj = nn.Linear(c.hidden_size, kv * d, bias=c.attention_bias) |
| self.v_proj = nn.Linear(c.hidden_size, kv * d, bias=c.attention_bias) |
| self.o_proj = nn.Linear(h * d, c.hidden_size, bias=c.attention_bias) |
| self.q_norm = CloverRMSNorm(d, c.rms_norm_eps) |
| self.k_norm = CloverRMSNorm(d, c.rms_norm_eps) |
|
|
| def rotary(self, x, positions): |
| rd = int(self.config.head_dim * self.config.partial_rotary_factor) |
| inv = 1.0 / self.config.rope_parameters["rope_theta"] ** ( |
| torch.arange(0, rd, 2, device=x.device, dtype=torch.float32) / rd |
| ) |
| phase = positions.float()[..., None] * inv |
| cos, sin = phase.cos().to(x.dtype)[:, None], phase.sin().to(x.dtype)[:, None] |
| first, second = x[..., :rd].chunk(2, -1) |
| rotated = torch.cat( |
| (first * cos - second * sin, second * cos + first * sin), -1 |
| ) |
| return torch.cat((rotated, x[..., rd:]), -1) |
|
|
| def forward(self, x, cache, valid, prefix_valid, positions, preview): |
| b, t, _ = x.shape |
| c = self.config |
| q, gate = ( |
| self.q_proj(x) |
| .view(b, t, c.num_attention_heads, 2 * c.head_dim) |
| .chunk(2, -1) |
| ) |
| q = self.rotary(self.q_norm(q).transpose(1, 2), positions) |
| k = self.rotary( |
| self.k_norm( |
| self.k_proj(x).view(b, t, c.num_key_value_heads, c.head_dim) |
| ).transpose(1, 2), |
| positions, |
| ) |
| v = self.v_proj(x).view(b, t, c.num_key_value_heads, c.head_dim).transpose(1, 2) |
| if cache is not None: |
| k = torch.cat((cache["key"], k), 2) |
| v = torch.cat((cache["value"], v), 2) |
| next_state = {"key": k, "value": v} |
| prefix = prefix_valid.shape[-1] |
| key_valid = torch.cat((prefix_valid, valid), 1) |
| allow = key_valid[:, None, None, :].expand(b, 1, t, prefix + t) |
| if not preview: |
| causal = ( |
| torch.arange(prefix + t, device=x.device)[None, :] |
| <= prefix + torch.arange(t, device=x.device)[:, None] |
| ) |
| allow = allow & causal[None, None] |
| allow = allow & valid[:, None, :, None] |
| k = k.repeat_interleave(c.num_attention_heads // c.num_key_value_heads, 1) |
| v = v.repeat_interleave(c.num_attention_heads // c.num_key_value_heads, 1) |
| if c.attention_backend == "sdpa": |
| y = F.scaled_dot_product_attention(q, k, v, attn_mask=allow, dropout_p=0.0) |
| else: |
| scores = (q.float() @ k.float().transpose(-1, -2)) * c.head_dim**-0.5 |
| scores = scores.masked_fill(~allow, -float("inf")) |
| probs = torch.nan_to_num(scores.softmax(-1)).to(v.dtype) |
| y = probs @ v |
| y = y.transpose(1, 2).reshape(b, t, -1) * gate.reshape(b, t, -1).sigmoid() |
| return self.o_proj(y), next_state |
|
|
|
|
| def coreset_torch(logits, valid, top_k, top_p): |
| b, t, e = logits.shape |
| ids = logits.float().argsort(dim=-1, descending=True, stable=True)[..., :top_k] |
| scores = logits.float().softmax(-1).gather(-1, ids) * valid[..., None] |
| totals = logits.new_zeros((b, e), dtype=torch.float32) |
| totals.scatter_add_(1, ids.reshape(b, -1), scores.reshape(b, -1)) |
| probs = totals / totals.sum(-1, keepdim=True).clamp_min(1e-9) |
| order = probs.argsort(dim=-1, descending=True, stable=True) |
| sorted_probs = probs.gather(-1, order) |
| keep = (sorted_probs.cumsum(-1) - sorted_probs) < top_p |
| keep[:, 0] = True |
| allowed = torch.zeros_like(keep).scatter(1, order, keep) |
| return logits.masked_fill(valid[..., None] & ~allowed[:, None, :], -float("inf")) |
|
|
|
|
| def linear_fp32_output(x, weight): |
| """Keep the down-projection accumulator in FP32 until route weighting.""" |
| if ( |
| x.is_cuda |
| and x.dtype in (torch.float16, torch.bfloat16) |
| and not torch.is_grad_enabled() |
| ): |
| try: |
| return torch.mm(x, weight.T, out_dtype=torch.float32) |
| except TypeError: |
| |
| pass |
| with torch.autocast(device_type=x.device.type, enabled=False): |
| return F.linear(x.float(), weight.float()) |
|
|
|
|
| def causal_conv_silu(joined, weight): |
| """Match serving's input-dtype products, FP32 sum and fused activation.""" |
| width = weight.shape[-1] |
| length = joined.shape[-1] - width + 1 |
| total = torch.zeros( |
| (*joined.shape[:-1], length), device=joined.device, dtype=torch.float32 |
| ) |
| for tap in range(width): |
| total.add_( |
| (joined[..., tap : tap + length] * weight[:, 0, tap][None, :, None]).float() |
| ) |
| return F.silu(total).to(joined.dtype) |
|
|
|
|
| class CloverMLP(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.gate_proj = nn.Linear( |
| c.hidden_size, c.shared_expert_intermediate_size, bias=False |
| ) |
| self.up_proj = nn.Linear( |
| c.hidden_size, c.shared_expert_intermediate_size, bias=False |
| ) |
| self.down_proj = nn.Linear( |
| c.shared_expert_intermediate_size, c.hidden_size, bias=False |
| ) |
|
|
| def forward(self, x): |
| gate, up = self.gate_proj(x), self.up_proj(x) |
| return self.down_proj((F.silu(gate.float()) * up.float()).to(x.dtype)) |
|
|
|
|
| class CloverExperts(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.gate_up_proj = nn.Parameter( |
| torch.empty(c.num_experts, 2 * c.moe_intermediate_size, c.hidden_size) |
| ) |
| self.down_proj = nn.Parameter( |
| torch.empty(c.num_experts, c.hidden_size, c.moe_intermediate_size) |
| ) |
|
|
| def forward(self, x, ids, weights): |
| out = torch.zeros_like(x, dtype=torch.float32) |
| |
| for expert in ids.unique().tolist(): |
| token, slot = torch.where(ids == expert) |
| gate, up = F.linear(x[token], self.gate_up_proj[expert]).chunk(2, -1) |
| hidden = (F.silu(gate.float()) * up.float()).to(x.dtype) |
| y = linear_fp32_output(hidden, self.down_proj[expert]) |
| routed = (y * weights[token, slot, None].float()).to(x.dtype) |
| out.index_add_(0, token, routed.float()) |
| return out.to(x.dtype) |
|
|
|
|
| class CloverMoeBlock(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.config = c |
| self.gate = nn.Linear(c.hidden_size, c.num_experts, bias=False) |
| self.experts = CloverExperts(c) |
| if c.num_shared_experts: |
| self.shared_expert = CloverMLP(c) |
| self.shared_expert_gate = nn.Linear(c.hidden_size, 1, bias=False) |
|
|
| def forward(self, x, valid, preview): |
| c = self.config |
| b, t, d = x.shape |
| logits = self.gate(x) |
| if preview and c.moe_block_coreset_enabled: |
| logits = coreset_torch( |
| logits, valid, c.num_experts_per_tok, c.moe_block_coreset_top_p |
| ) |
| probs = logits.float().softmax(-1) |
| ids = logits.float().argsort(dim=-1, descending=True, stable=True)[ |
| ..., : c.num_experts_per_tok |
| ] |
| weights = probs.gather(-1, ids) |
| weights = weights / weights.sum(-1, keepdim=True).clamp_min(1e-12) |
| y = self.experts( |
| x.reshape(-1, d), |
| ids.reshape(-1, c.num_experts_per_tok), |
| weights.reshape(-1, c.num_experts_per_tok), |
| ).view(b, t, d) |
| if c.num_shared_experts: |
| y = y + self.shared_expert(x) * self.shared_expert_gate(x).sigmoid() |
| return y |
|
|
|
|
| class CloverDecoderLayer(nn.Module): |
| def __init__(self, c, i): |
| super().__init__() |
| self.preview_type = c.layer_types[i] |
| if self.preview_type == "linear_attention": |
| self.linear_attn = CloverGatedDeltaNet(c) |
| else: |
| self.self_attn = CloverAttention(c) |
| self.input_layernorm = CloverRMSNorm(c.hidden_size, c.rms_norm_eps) |
| self.post_attention_layernorm = CloverRMSNorm(c.hidden_size, c.rms_norm_eps) |
| self.mlp = CloverMoeBlock(c) |
|
|
| def forward(self, x, cache, valid, prefix_valid, positions, preview): |
| hidden = self.input_layernorm(x) |
| if self.preview_type == "linear_attention": |
| hidden, next_cache = self.linear_attn(hidden, cache, valid, preview) |
| else: |
| hidden, next_cache = self.self_attn( |
| hidden, cache, valid, prefix_valid, positions, preview |
| ) |
| x = x + hidden |
| return ( |
| x + self.mlp(self.post_attention_layernorm(x), valid, preview), |
| next_cache, |
| ) |
|
|
|
|
| class CloverMoePreTrainedModel(PreTrainedModel): |
| config_class = CloverMoeConfig |
| base_model_prefix = "model" |
| _no_split_modules = ["CloverDecoderLayer"] |
|
|
| def _init_weights(self, module): |
| if isinstance(module, (nn.Linear, nn.Conv1d, nn.Embedding)): |
| nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) |
| if getattr(module, "bias", None) is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, CloverExperts): |
| nn.init.normal_(module.gate_up_proj, std=self.config.initializer_range) |
| nn.init.normal_(module.down_proj, std=self.config.initializer_range) |
| elif isinstance(module, CloverRMSNorm): |
| nn.init.constant_(module.weight, 0.0 if module.zero_centered else 1.0) |
| elif isinstance(module, CloverGatedDeltaNet): |
| nn.init.zeros_(module.A_log) |
| nn.init.ones_(module.dt_bias) |
| if hasattr(module, "conv_right_weight"): |
| nn.init.zeros_(module.conv_right_weight) |
| nn.init.constant_(module.bd_softmax_logit_scale, 10.0) |
| nn.init.constant_(module.bd_softmax_history_bias, 5.0) |
| nn.init.constant_(module.bd_softmax_history_c, 1.0) |
|
|
|
|
| class CloverMoeModel(CloverMoePreTrainedModel): |
| def __init__(self, config): |
| super().__init__(config) |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList( |
| [CloverDecoderLayer(config, i) for i in range(config.num_hidden_layers)] |
| ) |
| self.norm = CloverRMSNorm(config.hidden_size, config.rms_norm_eps) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.embed_tokens = value |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| position_ids=None, |
| past_key_values=None, |
| inputs_embeds=None, |
| use_cache=None, |
| return_dict=True, |
| preview=False, |
| **kwargs, |
| ): |
| if kwargs: |
| raise TypeError(f"Unsupported Clover forward arguments: {sorted(kwargs)}") |
| if (input_ids is None) == (inputs_embeds is None): |
| raise ValueError("Specify exactly one of input_ids or inputs_embeds") |
| x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds |
| b, t, _ = x.shape |
| if t == 0: |
| raise ValueError("Empty input is not supported") |
| if past_key_values is not None and not isinstance(past_key_values, CloverCache): |
| raise TypeError("past_key_values must be CloverCache") |
| prefix_valid = ( |
| past_key_values.attention_mask |
| if past_key_values is not None |
| else torch.zeros((b, 0), device=x.device, dtype=torch.bool) |
| ) |
| if prefix_valid.shape[0] != b: |
| raise ValueError("Cache batch size mismatch") |
| valid = torch.ones((b, t), device=x.device, dtype=torch.bool) |
| if attention_mask is not None: |
| if attention_mask.shape not in ((b, t), (b, prefix_valid.shape[-1] + t)): |
| raise ValueError( |
| "attention_mask must describe current tokens or the complete cached sequence" |
| ) |
| if attention_mask.shape[-1] != t and not torch.equal( |
| attention_mask[:, :-t].bool(), prefix_valid |
| ): |
| raise ValueError("attention_mask changes the committed prefix") |
| valid = attention_mask[:, -t:].to(device=x.device, dtype=torch.bool) |
| if preview and (t > self.config.block_size or self.config.block_size % t): |
| raise ValueError("Preview block size must divide configured block_size") |
| if position_ids is None: |
| positions = prefix_valid.sum(-1, keepdim=True) + valid.long().cumsum(-1) - 1 |
| positions = positions.masked_fill(~valid, 0).clamp_min(0) |
| else: |
| if position_ids.shape != (b, t): |
| raise ValueError("position_ids must have shape [batch, tokens]") |
| positions = position_ids |
| if torch.any((positions >= self.config.max_position_embeddings) & valid): |
| raise ValueError("Sequence exceeds max_position_embeddings") |
| next_layers = [] |
| for i, layer in enumerate(self.layers): |
| state = past_key_values.layers[i] if past_key_values is not None else None |
| x, state = layer(x, state, valid, prefix_valid, positions, preview) |
| next_layers.append(state) |
| use_cache = ( |
| getattr(self.config, "use_cache", True) if use_cache is None else use_cache |
| ) |
| cache = None |
| if use_cache: |
| cache = ( |
| past_key_values |
| if preview |
| else CloverCache(next_layers, torch.cat((prefix_valid, valid), 1)) |
| ) |
| output = BaseModelOutputWithPast( |
| last_hidden_state=self.norm(x), past_key_values=cache |
| ) |
| return output if return_dict else output.to_tuple() |
|
|
|
|
| class CloverMoeModelLM(CloverMoePreTrainedModel, GenerationMixin): |
| _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.model = CloverMoeModel(config) |
| self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) |
| self.post_init() |
|
|
| def get_input_embeddings(self): |
| return self.model.embed_tokens |
|
|
| def set_input_embeddings(self, value): |
| self.model.embed_tokens = value |
|
|
| def get_output_embeddings(self): |
| return self.lm_head |
|
|
| def set_output_embeddings(self, value): |
| self.lm_head = value |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| position_ids=None, |
| past_key_values=None, |
| inputs_embeds=None, |
| use_cache=None, |
| return_dict=True, |
| preview=False, |
| labels=None, |
| **kwargs, |
| ): |
| if labels is not None: |
| raise NotImplementedError( |
| "This release implements inference, not a training loss" |
| ) |
| out = 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, |
| return_dict=True, |
| preview=preview, |
| **kwargs, |
| ) |
| output = CausalLMOutputWithPast( |
| logits=self.lm_head(out.last_hidden_state), |
| past_key_values=out.past_key_values, |
| ) |
| return output if return_dict else output.to_tuple() |
|
|
| @torch.inference_mode() |
| def generate( |
| self, inputs=None, generation_config=None, attention_mask=None, **kwargs |
| ): |
| """Block LowConfidence decoding. Returns prompt+completion, like HF generate. |
| |
| Supported: max_new_tokens/max_length, greedy or temperature/top-k/top-p |
| sampling, EOS/padding, return_dict_in_generate, output_scores, streamer, |
| and stopping_criteria. Response penalties and length/thinking controls follow |
| SGLang LowConfidence semantics. Custom overrides: block_size, threshold. |
| Beam search and speculative decoding are not exposed. |
| Batched requests run sequentially to keep independent clean-state caches. |
| """ |
| input_ids = kwargs.pop("input_ids", inputs) |
| if input_ids is None or input_ids.ndim != 2: |
| raise ValueError("input_ids must be [batch, tokens]") |
| gen = deepcopy( |
| generation_config or self.generation_config or GenerationConfig() |
| ) |
| custom = deepcopy(self.config.denoising_config) |
| for key in ("block_size", "threshold"): |
| if key in kwargs: |
| custom[key] = kwargs.pop(key) |
| block_size = custom.get("block_size", self.config.block_size) |
| threshold = custom["threshold"] |
| if block_size <= 0 or self.config.block_size % block_size: |
| raise ValueError( |
| "block_size must be positive and divide the configured maximum block size" |
| ) |
| if not 0 <= threshold <= 1: |
| raise ValueError("threshold must be in [0, 1]") |
| streamer = kwargs.pop("streamer", None) |
| stopping = kwargs.pop("stopping_criteria", None) |
| allowed = { |
| "max_new_tokens", |
| "max_length", |
| "do_sample", |
| "temperature", |
| "top_k", |
| "top_p", |
| "eos_token_id", |
| "pad_token_id", |
| "return_dict_in_generate", |
| "output_scores", |
| "num_beams", |
| "num_return_sequences", |
| "use_cache", |
| } |
| allowed.update(PENALTY_DEFAULTS) |
| for key in list(kwargs): |
| if key not in allowed: |
| raise TypeError(f"Unsupported generation argument: {key}") |
| setattr(gen, key, kwargs.pop(key)) |
| defaults = GenerationConfig().to_dict() |
| for key, value in gen.to_dict().items(): |
| if ( |
| key not in allowed |
| and not key.startswith("_") |
| and key not in ("transformers_version", "bos_token_id") |
| ): |
| if value is None or ( |
| value is False and defaults.get(key) in (None, False) |
| ): |
| continue |
| if value != defaults.get(key): |
| raise NotImplementedError( |
| f"Unsupported nondefault generation setting: {key}" |
| ) |
| validate_penalties(gen, self.config.vocab_size) |
| if (gen.num_beams or 1) != 1 or (gen.num_return_sequences or 1) != 1: |
| raise NotImplementedError( |
| "Clover block generation supports one sequence per prompt, without beams" |
| ) |
| if gen.temperature is None: |
| gen.temperature = 1.0 |
| if gen.use_cache is False: |
| raise NotImplementedError("Clover generation requires clean-state caching") |
| if gen.do_sample and gen.temperature <= 0: |
| raise ValueError("Sampling requires temperature > 0") |
| if gen.top_p is not None and not 0 < gen.top_p <= 1: |
| raise ValueError("top_p must be in (0,1]") |
| if input_ids.shape[0] > 1 and (streamer is not None or gen.output_scores): |
| raise ValueError( |
| "Streaming and output_scores currently require batch size 1" |
| ) |
| if attention_mask is None: |
| attention_mask = torch.ones_like(input_ids) |
| if attention_mask.shape != input_ids.shape: |
| raise ValueError("attention_mask shape mismatch") |
| eos = ( |
| gen.eos_token_id |
| if gen.eos_token_id is not None |
| else getattr(self.config, "eos_token_id", None) |
| ) |
| eos = [] if eos is None else ([eos] if isinstance(eos, int) else list(eos)) |
| eos = sorted(set(eos) | set(gen.stop_token_ids)) |
| pad = gen.pad_token_id |
| if pad is None: |
| pad = getattr(self.config, "pad_token_id", None) |
| if pad is None: |
| pad = eos[0] if eos else 0 |
| if gen.output_scores and gen.do_sample: |
| raise ValueError("output_scores is available for greedy generation only") |
| results = [] |
| all_scores = [] |
| for row in range(input_ids.shape[0]): |
| prompt = input_ids[row][attention_mask[row].bool()].unsqueeze(0) |
| if prompt.shape[1] == 0: |
| raise ValueError("Each prompt must have at least one token") |
| |
| |
| maximum = gen.max_new_tokens |
| if maximum is None: |
| maximum = max(0, (gen.max_length or 20) - prompt.shape[1]) |
| if maximum < 0: |
| raise ValueError("max_new_tokens must be nonnegative") |
| |
| |
| preview_length = ((maximum + block_size - 1) // block_size) * block_size |
| if prompt.shape[1] + preview_length > self.config.max_position_embeddings: |
| raise ValueError( |
| "Prompt plus full preview blocks exceeds context length: " |
| f"{prompt.shape[1]} + {preview_length} > " |
| f"{self.config.max_position_embeddings} tokens" |
| ) |
| completed = prompt |
| if streamer is not None: |
| streamer.put(prompt.cpu()) |
| cache = None |
| if maximum: |
| |
| cache = self.model(prompt, use_cache=True).past_key_values |
| response = [] |
| penalties = CloverPenaltyState( |
| gen, self.config.vocab_size, prompt.device, eos + [pad] |
| ) |
| stopped = False |
| while len(response) < maximum and not stopped: |
| remaining = min(block_size, maximum - len(response)) |
| block = torch.full( |
| (1, block_size), |
| self.config.mask_token_id, |
| device=prompt.device, |
| dtype=prompt.dtype, |
| ) |
| valid = torch.ones_like(block, dtype=torch.bool) |
| pending = valid.clone() |
| token_scores = [None] * block_size |
| for _ in range(block_size): |
| out = self( |
| block, |
| attention_mask=valid, |
| past_key_values=cache, |
| use_cache=True, |
| preview=True, |
| ) |
| logits = out.logits[0].float() |
| logits[:, self.config.mask_token_id] = -float("inf") |
| penalties.apply(logits, block, self.config.mask_token_id) |
| prediction, confidence = predict_with_confidence(logits, gen) |
| confidence = confidence.masked_fill(~pending[0], -float("inf")) |
| transfer = (confidence > threshold) & pending[0] |
| if not bool(transfer.any()): |
| transfer[confidence.argmax()] = True |
| if gen.output_scores: |
| for pos in transfer.nonzero().flatten().tolist(): |
| token_scores[pos] = logits[pos : pos + 1].detach().clone() |
| block[0, transfer] = prediction[transfer] |
| pending[0, transfer] = False |
| if not bool(pending.any()): |
| break |
| if bool(pending.any()): |
| raise RuntimeError( |
| "Denoising step limit reached with unresolved masks" |
| ) |
| |
| |
| new = block[0, :remaining] |
| |
| |
| for i, token in enumerate(new.tolist()): |
| response.append(token) |
| completed = torch.cat((completed, new[i : i + 1].view(1, 1)), 1) |
| if gen.output_scores: |
| all_scores.append(token_scores[i]) |
| if streamer is not None: |
| streamer.put(new[i : i + 1].cpu()) |
| if token in eos or ( |
| stopping is not None |
| and bool( |
| torch.as_tensor( |
| stopping( |
| completed, |
| tuple(all_scores) if gen.output_scores else None, |
| ) |
| ).all() |
| ) |
| ): |
| stopped = True |
| break |
| penalties.commit(new[: i + 1]) |
| if len(response) < maximum and not stopped: |
| |
| |
| cache = self.model( |
| new.view(1, -1), past_key_values=cache, use_cache=True |
| ).past_key_values |
| if streamer is not None: |
| streamer.end() |
| results.append(torch.cat((input_ids[row], input_ids.new_tensor(response)))) |
| length = max(r.numel() for r in results) |
| sequences = input_ids.new_full((len(results), length), pad) |
| for i, result in enumerate(results): |
| sequences[i, : result.numel()] = result |
| if gen.return_dict_in_generate: |
| return GenerateDecoderOnlyOutput( |
| sequences=sequences, |
| scores=tuple(all_scores) if gen.output_scores else None, |
| ) |
| return sequences |
|
|
| def diffusion_generate(self, *args, **kwargs): |
| return self.generate(*args, **kwargs) |
|
|
