""" Prefill and Decode wrapper models for CoreML conversion. Key design: all dynamic indexing is eliminated. Positions are encoded into cos/sin/mask inputs by the caller, not computed inside the model. This makes the models fully traceable by torch.jit.trace and convertible by coremltools. """ import torch import torch.nn as nn import torch.nn.functional as F from attention import ( LlamaDecoderLayerDecode, LlamaRMSNorm, precompute_rope_frequencies, ) # Model constants NUM_LAYERS = 30 HIDDEN_SIZE = 576 NUM_HEADS = 9 NUM_KV_HEADS = 3 HEAD_DIM = 64 INTERMEDIATE_SIZE = 1536 VOCAB_SIZE = 20802 RMS_NORM_EPS = 1e-5 ROPE_THETA = 100000.0 MAX_CONTEXT = 512 PREFILL_SEQ_LEN = 512 SPEAKER_DIM = 128 class PlaprePico(nn.Module): """Generates one token at a time using the KV cache. Also used for token-by-token prefill. Speaker conditioning is handled internally: at position 0, pass is_speaker_step=1.0 and the raw speaker_embedding. The model projects it and replaces the token embedding. Inputs: input_ids: (1, 1) int32 causal_mask: (1, 1, 1, 2048) float16 — 0 or -inf cos: (1, 1, 1, 64) float16 — RoPE cos for current position sin: (1, 1, 1, 64) float16 — RoPE sin for current position update_mask: (1, 1, 2048, 1) float16 — one-hot at current position speaker_embedding: (1, 128) float16 — raw speaker embedding (used at position 0) is_speaker_step: (1,) float16 — 1.0 at position 0, 0.0 otherwise State buffers: k_cache_0..29, v_cache_0..29: (1, 3, 2048, 64) float16 Output: logits: (1, 1, 20802) float16 """ def __init__(self): super().__init__() self.embed_tokens = nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE) self.speaker_proj = nn.Linear(SPEAKER_DIM, HIDDEN_SIZE, bias=True) self.layers = nn.ModuleList( [ LlamaDecoderLayerDecode( hidden_size=HIDDEN_SIZE, num_heads=NUM_HEADS, num_kv_heads=NUM_KV_HEADS, head_dim=HEAD_DIM, intermediate_size=INTERMEDIATE_SIZE, rms_norm_eps=RMS_NORM_EPS, max_context=MAX_CONTEXT, ) for _ in range(NUM_LAYERS) ] ) self.norm = LlamaRMSNorm(HIDDEN_SIZE, eps=RMS_NORM_EPS) # KV cache state buffers for i in range(NUM_LAYERS): self.register_buffer( f"k_cache_{i}", torch.zeros(1, NUM_KV_HEADS, MAX_CONTEXT, HEAD_DIM, dtype=torch.float16), ) self.register_buffer( f"v_cache_{i}", torch.zeros(1, NUM_KV_HEADS, MAX_CONTEXT, HEAD_DIM, dtype=torch.float16), ) def forward( self, input_ids: torch.Tensor, causal_mask: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, update_mask: torch.Tensor, speaker_embedding: torch.Tensor, is_speaker_step: torch.Tensor, ) -> torch.Tensor: hidden = self.embed_tokens(input_ids) # (1, 1, 576) # Speaker conditioning: at position 0, replace token embedding with # speaker_proj(embedding). Use is_speaker_step as a blend factor: # hidden = (1 - flag) * embed + flag * speaker_proj(emb) # When flag=0 this is a no-op; when flag=1 it fully replaces. spk_proj = self.speaker_proj(speaker_embedding).unsqueeze(1) # (1, 1, 576) flag = is_speaker_step.view(1, 1, 1) # (1, 1, 1) for broadcast hidden = hidden * (1.0 - flag) + spk_proj * flag # Unrolled layer loop — explicit self.k_cache_N access required for # torch.jit.trace to generate prim::SetAttr, which coremltools converts # to coreml_update_state. getattr() in a loop doesn't work. hidden, new_k_0, new_v_0 = self.layers[0](hidden, cos, sin, causal_mask, self.k_cache_0, self.v_cache_0, update_mask) self.k_cache_0 += (new_k_0 - self.k_cache_0) self.v_cache_0 += (new_v_0 - self.v_cache_0) hidden, new_k_1, new_v_1 = self.layers[1](hidden, cos, sin, causal_mask, self.k_cache_1, self.v_cache_1, update_mask) self.k_cache_1 += (new_k_1 - self.k_cache_1) self.v_cache_1 += (new_v_1 - self.v_cache_1) hidden, new_k_2, new_v_2 = self.layers[2](hidden, cos, sin, causal_mask, self.k_cache_2, self.v_cache_2, update_mask) self.k_cache_2 += (new_k_2 - self.k_cache_2) self.v_cache_2 += (new_v_2 - self.v_cache_2) hidden, new_k_3, new_v_3 = self.layers[3](hidden, cos, sin, causal_mask, self.k_cache_3, self.v_cache_3, update_mask) self.k_cache_3 += (new_k_3 - self.k_cache_3) self.v_cache_3 += (new_v_3 - self.v_cache_3) hidden, new_k_4, new_v_4 = self.layers[4](hidden, cos, sin, causal_mask, self.k_cache_4, self.v_cache_4, update_mask) self.k_cache_4 += (new_k_4 - self.k_cache_4) self.v_cache_4 += (new_v_4 - self.v_cache_4) hidden, new_k_5, new_v_5 = self.layers[5](hidden, cos, sin, causal_mask, self.k_cache_5, self.v_cache_5, update_mask) self.k_cache_5 += (new_k_5 - self.k_cache_5) self.v_cache_5 += (new_v_5 - self.v_cache_5) hidden, new_k_6, new_v_6 = self.layers[6](hidden, cos, sin, causal_mask, self.k_cache_6, self.v_cache_6, update_mask) self.k_cache_6 += (new_k_6 - self.k_cache_6) self.v_cache_6 += (new_v_6 - self.v_cache_6) hidden, new_k_7, new_v_7 = self.layers[7](hidden, cos, sin, causal_mask, self.k_cache_7, self.v_cache_7, update_mask) self.k_cache_7 += (new_k_7 - self.k_cache_7) self.v_cache_7 += (new_v_7 - self.v_cache_7) hidden, new_k_8, new_v_8 = self.layers[8](hidden, cos, sin, causal_mask, self.k_cache_8, self.v_cache_8, update_mask) self.k_cache_8 += (new_k_8 - self.k_cache_8) self.v_cache_8 += (new_v_8 - self.v_cache_8) hidden, new_k_9, new_v_9 = self.layers[9](hidden, cos, sin, causal_mask, self.k_cache_9, self.v_cache_9, update_mask) self.k_cache_9 += (new_k_9 - self.k_cache_9) self.v_cache_9 += (new_v_9 - self.v_cache_9) hidden, new_k_10, new_v_10 = self.layers[10](hidden, cos, sin, causal_mask, self.k_cache_10, self.v_cache_10, update_mask) self.k_cache_10 += (new_k_10 - self.k_cache_10) self.v_cache_10 += (new_v_10 - self.v_cache_10) hidden, new_k_11, new_v_11 = self.layers[11](hidden, cos, sin, causal_mask, self.k_cache_11, self.v_cache_11, update_mask) self.k_cache_11 += (new_k_11 - self.k_cache_11) self.v_cache_11 += (new_v_11 - self.v_cache_11) hidden, new_k_12, new_v_12 = self.layers[12](hidden, cos, sin, causal_mask, self.k_cache_12, self.v_cache_12, update_mask) self.k_cache_12 += (new_k_12 - self.k_cache_12) self.v_cache_12 += (new_v_12 - self.v_cache_12) hidden, new_k_13, new_v_13 = self.layers[13](hidden, cos, sin, causal_mask, self.k_cache_13, self.v_cache_13, update_mask) self.k_cache_13 += (new_k_13 - self.k_cache_13) self.v_cache_13 += (new_v_13 - self.v_cache_13) hidden, new_k_14, new_v_14 = self.layers[14](hidden, cos, sin, causal_mask, self.k_cache_14, self.v_cache_14, update_mask) self.k_cache_14 += (new_k_14 - self.k_cache_14) self.v_cache_14 += (new_v_14 - self.v_cache_14) hidden, new_k_15, new_v_15 = self.layers[15](hidden, cos, sin, causal_mask, self.k_cache_15, self.v_cache_15, update_mask) self.k_cache_15 += (new_k_15 - self.k_cache_15) self.v_cache_15 += (new_v_15 - self.v_cache_15) hidden, new_k_16, new_v_16 = self.layers[16](hidden, cos, sin, causal_mask, self.k_cache_16, self.v_cache_16, update_mask) self.k_cache_16 += (new_k_16 - self.k_cache_16) self.v_cache_16 += (new_v_16 - self.v_cache_16) hidden, new_k_17, new_v_17 = self.layers[17](hidden, cos, sin, causal_mask, self.k_cache_17, self.v_cache_17, update_mask) self.k_cache_17 += (new_k_17 - self.k_cache_17) self.v_cache_17 += (new_v_17 - self.v_cache_17) hidden, new_k_18, new_v_18 = self.layers[18](hidden, cos, sin, causal_mask, self.k_cache_18, self.v_cache_18, update_mask) self.k_cache_18 += (new_k_18 - self.k_cache_18) self.v_cache_18 += (new_v_18 - self.v_cache_18) hidden, new_k_19, new_v_19 = self.layers[19](hidden, cos, sin, causal_mask, self.k_cache_19, self.v_cache_19, update_mask) self.k_cache_19 += (new_k_19 - self.k_cache_19) self.v_cache_19 += (new_v_19 - self.v_cache_19) hidden, new_k_20, new_v_20 = self.layers[20](hidden, cos, sin, causal_mask, self.k_cache_20, self.v_cache_20, update_mask) self.k_cache_20 += (new_k_20 - self.k_cache_20) self.v_cache_20 += (new_v_20 - self.v_cache_20) hidden, new_k_21, new_v_21 = self.layers[21](hidden, cos, sin, causal_mask, self.k_cache_21, self.v_cache_21, update_mask) self.k_cache_21 += (new_k_21 - self.k_cache_21) self.v_cache_21 += (new_v_21 - self.v_cache_21) hidden, new_k_22, new_v_22 = self.layers[22](hidden, cos, sin, causal_mask, self.k_cache_22, self.v_cache_22, update_mask) self.k_cache_22 += (new_k_22 - self.k_cache_22) self.v_cache_22 += (new_v_22 - self.v_cache_22) hidden, new_k_23, new_v_23 = self.layers[23](hidden, cos, sin, causal_mask, self.k_cache_23, self.v_cache_23, update_mask) self.k_cache_23 += (new_k_23 - self.k_cache_23) self.v_cache_23 += (new_v_23 - self.v_cache_23) hidden, new_k_24, new_v_24 = self.layers[24](hidden, cos, sin, causal_mask, self.k_cache_24, self.v_cache_24, update_mask) self.k_cache_24 += (new_k_24 - self.k_cache_24) self.v_cache_24 += (new_v_24 - self.v_cache_24) hidden, new_k_25, new_v_25 = self.layers[25](hidden, cos, sin, causal_mask, self.k_cache_25, self.v_cache_25, update_mask) self.k_cache_25 += (new_k_25 - self.k_cache_25) self.v_cache_25 += (new_v_25 - self.v_cache_25) hidden, new_k_26, new_v_26 = self.layers[26](hidden, cos, sin, causal_mask, self.k_cache_26, self.v_cache_26, update_mask) self.k_cache_26 += (new_k_26 - self.k_cache_26) self.v_cache_26 += (new_v_26 - self.v_cache_26) hidden, new_k_27, new_v_27 = self.layers[27](hidden, cos, sin, causal_mask, self.k_cache_27, self.v_cache_27, update_mask) self.k_cache_27 += (new_k_27 - self.k_cache_27) self.v_cache_27 += (new_v_27 - self.v_cache_27) hidden, new_k_28, new_v_28 = self.layers[28](hidden, cos, sin, causal_mask, self.k_cache_28, self.v_cache_28, update_mask) self.k_cache_28 += (new_k_28 - self.k_cache_28) self.v_cache_28 += (new_v_28 - self.v_cache_28) hidden, new_k_29, new_v_29 = self.layers[29](hidden, cos, sin, causal_mask, self.k_cache_29, self.v_cache_29, update_mask) self.k_cache_29 += (new_k_29 - self.k_cache_29) self.v_cache_29 += (new_v_29 - self.v_cache_29) hidden = self.norm(hidden) logits = F.linear(hidden, self.embed_tokens.weight) return logits