# Code study — modular walkthrough

Source of truth: [`reference/transformers-inkling/modular_inkling.py`](../reference/transformers-inkling/modular_inkling.py)  

Line numbers from [`derived/class_index.md`](../derived/class_index.md) (regenerate after refresh).

**Rule:** upstream edits go in `modular_inkling.py` only; `modeling_*.py` / `configuration_*.py` are generated.

---

## 1. Config surface

| Class | Responsibility |
|---|---|
| `InklingTextConfig` | LM dims, MoE, SWA/global, SConv, μP, MTP fields |
| `InklingVisionConfig` | hMLP patch/temporal |
| `InklingAudioConfig` | dMel dims |
| `InklingConfig` | top-level MM; wires `mtp_config` into text; special token ids |

Important remaps (checkpoint → code):

| Hub / ckpt | Code |
|---|---|
| `sconv_kernel_size` | `conv_kernel_size` |
| `model_max_length` | `max_position_embeddings` |
| `dense_intermediate_size` | dense `intermediate_size` |
| `mtp_config.num_nextn_predict_layers` | `num_mtp_layers` |

`__post_init__` builds:

- `layer_types[i] ∈ {hybrid_sliding, hybrid}` from `local_layer_ids`  
- `mlp_layer_types[i] ∈ {dense, sparse}` from `dense_mlp_idx`  
- `number_of_conv_states = 4`

---

## 2. Relative logits → attention

### `InklingRelativeLogits`

Learned bank `[d_rel, rel_extent]`. Forward:

1. `rel_logits = relative_states @ proj` → `[B,H,Q,extent]`  
2. `distance = q_pos - k_pos`  
3. gather; zero if distance &lt; 0 or ≥ extent  

Past-only bias; mask still enforces causality.

### `InklingAttention`

1. Project Q, K, V, R  
2. **SConv on K and V** (`conv_idx` 0,1)  
3. Reshape; **RMSNorm Q and K**  
4. Cache update if `past_key_values`  
5. Build `position_bias` from R + positions  
6. Optional **log length scaling** (global, long n) on Q and bias  
7. Attention interface with `scaling = 1/head_dim`  
8. `o_proj`  

SWA vs global: head/KV counts, window, `rel_extent` only.

`eager_attention_forward` shows the bias add path explicitly — other backends must accept `position_bias`.

---

## 3. SConv module

`InklingShortConvolution`:

- depthwise `Conv1d`, kernel = config, padding = k−1  
- residual: `out = conv(x) + x`  
- fp32 compute  
- decode: `causal_conv1d_update` with cached state  
- prefill: `causal_conv1d_fn`  

Used 4× per `InklingDecoderLayer`.

---

## 4. MoE path

### `InklingTopkRouter`

- Weight over `n_routed + n_shared`  
- `sigmoid` scores  
- top-k on **routed + e_score_correction_bias**  
- joint `logsigmoid` → logsumexp normalize selected routed **and** shared  
- `× route_scale × global_scale`  
- returns routed weights, indices, `shared_gammas`

### `InklingExperts` / `InklingSharedExperts` / `InklingMoE`

```
y = routed_experts(h, idx, w) + shared_experts(h, gammas)
```

Dense path: `InklingMLP` = Gemma3 SwiGLU + `global_scale`.

---

## 5. Decoder layer residual

```
h = h + attn_sconv(attn(LN(h)))
h = h + mlp_sconv(mlp(LN(h)))
```

---

## 6. Text model masks

`InklingTextModel` builds a **dict**:

| Key | Consumer |
|---|---|
| `full_attention` | global layers |
| `sliding_attention` | SWA layers |
| `linear_attention` | SConv padding (`conv_mask`) |

Embed path: `embed_norm(embed_tokens(ids))`.

---

## 7. Multimodal model

`InklingModel`:

1. text embeds  
2. vision tower → scatter @ `image_token_id`  
3. audio tower → scatter @ `audio_token_id`  
4. run language model  

Placeholder count must equal feature count (`torch_compilable_check`).

`InklingForConditionalGeneration`:

- `lm_head` over full vocab; slice to unpadded in forward  
- μP divide by `logits_mup_width_multiplier`  
- generation: pass `pixel_values` / `audio_input_ids` only on first step when caching  

MTP: `_keys_to_ignore_on_load_unexpected = [r"model\.mtp\..*"]`.

---

## 8. Vision / audio modules

| Module | Notes |
|---|---|
| `plan_out_scales` | prime-factor schedule for hMLP dims |
| `InklingVisionEncoderLayer` | fold space-time → Linear → (RMSNorm+GELU) |
| `InklingVisionModel` | stack → final norm → per-patch vectors |
| `InklingAudioModelEmbeddings` | HiggsAudioV2-style multi-codebook embeds |
| `InklingAudioModel` | embed + norm |

---

## 9. Inheritance map

| Piece | From |
|---|---|
| RMSNorm | Llama |
| Dense MLP base | Gemma3 |
| Experts | Mixtral |
| Conv pad helper | Qwen3-Next |
| Audio embeddings | Higgs Audio v2 |

---

## 10. Read checklist

Open modular and, without scrolling randomly, find:

- [ ] Where `scaling = 1.0 / self.head_dim` is set  
- [ ] Where shared logits are concatenated for joint normalize  
- [ ] All four `conv_idx` constructions  
- [ ] Where `masked_scatter` merges image features  
- [ ] Where logits are divided by μP multiplier  

Next: [02-chat-protocol.md](02-chat-protocol.md)