| { |
| "architecture": "causal-encoder-decoder", |
| "num_encoder_layers": 4, |
| "num_decoder_layers": 4, |
| "hidden_size": 384, |
| "num_attention_heads": 6, |
| "head_dim": 64, |
| "vocab_size": 8192, |
| "local_window": 64, |
| "max_position_embeddings": 512, |
| "position_encoding": "rope", |
| "normalization": "rmsnorm", |
| "activation": "swiglu", |
| "input": { |
| "name": "input_ids", |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch=1", |
| "sequence" |
| ] |
| }, |
| "mask_semantics": "int64: 1=allowed, 0=blocked; first T attention columns are global, next T are local", |
| "attention_probs_semantics": "softmax probabilities actually multiplied by concatenated values; first T keys global, next T local", |
| "outputs": { |
| "token_embeddings": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Learned token embedding before encoder blocks" |
| }, |
| "encoder_0_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "encoder_0_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "encoder_0_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "encoder_0_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "encoder_0_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: the block's norm1" |
| }, |
| "encoder_0_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: the block's norm1" |
| }, |
| "encoder_0_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: the block's norm1" |
| }, |
| "encoder_0_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "encoder_0_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "encoder_0_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "encoder_0_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "encoder_0_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "encoder_0_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "encoder_0_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "encoder_0_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "encoder_0_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "encoder_0_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "encoder_0_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "encoder_0_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "encoder_0_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "encoder_0_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "encoder_0_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "encoder_0_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "encoder_0_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "encoder_0_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "encoder_0_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "encoder_1_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "encoder_1_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "encoder_1_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "encoder_1_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "encoder_1_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: the block's norm1" |
| }, |
| "encoder_1_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: the block's norm1" |
| }, |
| "encoder_1_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: the block's norm1" |
| }, |
| "encoder_1_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "encoder_1_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "encoder_1_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "encoder_1_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "encoder_1_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "encoder_1_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "encoder_1_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "encoder_1_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "encoder_1_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "encoder_1_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "encoder_1_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "encoder_1_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "encoder_1_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "encoder_1_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "encoder_1_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "encoder_1_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "encoder_1_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "encoder_1_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "encoder_1_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "encoder_2_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "encoder_2_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "encoder_2_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "encoder_2_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "encoder_2_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: the block's norm1" |
| }, |
| "encoder_2_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: the block's norm1" |
| }, |
| "encoder_2_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: the block's norm1" |
| }, |
| "encoder_2_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "encoder_2_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "encoder_2_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "encoder_2_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "encoder_2_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "encoder_2_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "encoder_2_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "encoder_2_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "encoder_2_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "encoder_2_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "encoder_2_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "encoder_2_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "encoder_2_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "encoder_2_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "encoder_2_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "encoder_2_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "encoder_2_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "encoder_2_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "encoder_2_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "encoder_3_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "encoder_3_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "encoder_3_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "encoder_3_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "encoder_3_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: the block's norm1" |
| }, |
| "encoder_3_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: the block's norm1" |
| }, |
| "encoder_3_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: the block's norm1" |
| }, |
| "encoder_3_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "encoder_3_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "encoder_3_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "encoder_3_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "encoder_3_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "encoder_3_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "encoder_3_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "encoder_3_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "encoder_3_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "encoder_3_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "encoder_3_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "encoder_3_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "encoder_3_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "encoder_3_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "encoder_3_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "encoder_3_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "encoder_3_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "encoder_3_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "encoder_3_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "encoder_pre_final_norm": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Encoder hidden state before encoder final RMSNorm" |
| }, |
| "encoder_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Normalized encoder representation used as decoder global K/V source" |
| }, |
| "decoder_0_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "decoder_0_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "decoder_0_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "decoder_0_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "decoder_0_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: encoder_output" |
| }, |
| "decoder_0_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: encoder_output" |
| }, |
| "decoder_0_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: encoder_output" |
| }, |
| "decoder_0_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "decoder_0_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "decoder_0_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "decoder_0_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "decoder_0_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "decoder_0_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "decoder_0_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "decoder_0_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "decoder_0_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "decoder_0_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "decoder_0_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "decoder_0_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "decoder_0_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "decoder_0_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "decoder_0_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "decoder_0_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "decoder_0_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "decoder_0_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "decoder_0_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "decoder_1_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "decoder_1_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "decoder_1_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "decoder_1_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "decoder_1_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: encoder_output" |
| }, |
| "decoder_1_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: encoder_output" |
| }, |
| "decoder_1_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: encoder_output" |
| }, |
| "decoder_1_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "decoder_1_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "decoder_1_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "decoder_1_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "decoder_1_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "decoder_1_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "decoder_1_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "decoder_1_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "decoder_1_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "decoder_1_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "decoder_1_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "decoder_1_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "decoder_1_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "decoder_1_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "decoder_1_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "decoder_1_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "decoder_1_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "decoder_1_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "decoder_1_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "decoder_2_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "decoder_2_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "decoder_2_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "decoder_2_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "decoder_2_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: encoder_output" |
| }, |
| "decoder_2_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: encoder_output" |
| }, |
| "decoder_2_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: encoder_output" |
| }, |
| "decoder_2_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "decoder_2_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "decoder_2_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "decoder_2_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "decoder_2_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "decoder_2_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "decoder_2_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "decoder_2_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "decoder_2_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "decoder_2_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "decoder_2_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "decoder_2_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "decoder_2_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "decoder_2_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "decoder_2_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "decoder_2_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "decoder_2_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "decoder_2_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "decoder_2_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "decoder_3_input": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input before attention RMSNorm" |
| }, |
| "decoder_3_norm1": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block input after attention RMSNorm; local K/V source" |
| }, |
| "decoder_3_q_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query linear projection before RoPE" |
| }, |
| "decoder_3_q": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Query after RoPE, used in attention scores" |
| }, |
| "decoder_3_global_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key projection before RoPE; source: encoder_output" |
| }, |
| "decoder_3_global_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global key after RoPE, used in attention scores; source: encoder_output" |
| }, |
| "decoder_3_global_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Global value projection used in attention output; source: encoder_output" |
| }, |
| "decoder_3_local_k_pre_rope": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key projection before RoPE" |
| }, |
| "decoder_3_local_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local key after RoPE, used in attention scores" |
| }, |
| "decoder_3_local_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "sequence", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 64 |
| ], |
| "meaning": "Local value projection used in attention output" |
| }, |
| "decoder_3_combined_k": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global keys followed by local keys on the sequence axis" |
| }, |
| "decoder_3_combined_v": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "global+local-key", |
| "head_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 8, |
| 64 |
| ], |
| "meaning": "Global values followed by local values on the sequence axis" |
| }, |
| "decoder_3_global_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed global key positions; causal" |
| }, |
| "decoder_3_local_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 4 |
| ], |
| "meaning": "Allowed local key positions; causal and within local window" |
| }, |
| "decoder_3_attention_mask": { |
| "dtype": "int64", |
| "semantic_shape": [ |
| "batch", |
| "broadcast-head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 1, |
| 4, |
| 8 |
| ], |
| "meaning": "Global mask followed by local mask; 1 allowed, 0 blocked" |
| }, |
| "decoder_3_attn_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Scaled Q times transposed combined K, before mask" |
| }, |
| "decoder_3_attn_masked_scores": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Attention scores with blocked positions set to negative infinity" |
| }, |
| "decoder_3_attn_probs": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "head", |
| "query", |
| "global+local-key" |
| ], |
| "example_shape": [ |
| 1, |
| 6, |
| 4, |
| 8 |
| ], |
| "meaning": "Softmax probabilities actually multiplied by combined V" |
| }, |
| "decoder_3_attn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Attention result after output projection, before residual add" |
| }, |
| "decoder_3_after_attn_residual": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after attention residual add" |
| }, |
| "decoder_3_norm2": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Hidden state after FFN RMSNorm" |
| }, |
| "decoder_3_ffn_gate": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU gate linear projection" |
| }, |
| "decoder_3_ffn_up": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU up linear projection" |
| }, |
| "decoder_3_ffn_activation": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "ffn_dim" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 1024 |
| ], |
| "meaning": "SwiGLU product SiLU(gate) times up" |
| }, |
| "decoder_3_ffn_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "FFN down projection before residual add" |
| }, |
| "decoder_3_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Block output after FFN residual add" |
| }, |
| "decoder_output": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Last decoder block output before final RMSNorm" |
| }, |
| "final_hidden": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "feature" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 384 |
| ], |
| "meaning": "Decoder hidden state after final RMSNorm" |
| }, |
| "logits": { |
| "dtype": "float32", |
| "semantic_shape": [ |
| "batch", |
| "sequence", |
| "vocab" |
| ], |
| "example_shape": [ |
| 1, |
| 4, |
| 8192 |
| ], |
| "meaning": "Full-sequence vocabulary logits from tied LM head" |
| } |
| } |
| } |
|
|
