Skip to content

feat(vision): add the Inkling HMLP vision tower and 40x40 temporal-pair image ingestion #1327

Closed

Description

@inureyes

Summary

Inkling encodes images without a vision transformer: each 40x40 RGB tile is duplicated into a temporal pair [T=2, 40, 40, 3] and pushed through a hierarchical MLP (HMLP) that repeatedly folds space and time into the channel axis and projects, ending in one soft token of the text hidden size per tile. This issue adds the tower (src/vision/encoders/inkling_hmlp.rs), the tile preprocessor, placeholder expansion (<|unused_200054|> repeated once per tile), and the embedding merge, on top of #1318.

Current behavior

mlxcel has no HMLP encoder. The closest pieces are the LLaVA-style merge in src/vision/merge.rs (merge_llava: replace placeholder positions with projected features in order) and the host-side tile preprocessors under src/vision/processors/. src/loading/mod.rs has no Inkling VLM loader.

Expected behavior

Vision config (vision_config): model_type: "inkling_vision", patch_size: 40, temporal_patch_size: 2, n_channels (HF) or num_channels (MLX) : 3, n_layers: 4 (1 on the 0.6B checkpoint), rms_norm_eps: 1e-6, decoder_dmodel / text_hidden_size = text hidden_size (4096 Inkling-Small, 6144 Inkling, 1024 for 0.6B). use_vision_norm: true, vision_encoder_type: "hmlp" are informational.

Preprocessing (processor_config.json image_processor: image_mean = [0.48145466, 0.4578275, 0.40821073], image_std = [0.26862954, 0.26130258, 0.27577711], rescale_factor = 1/255):

img = RGB uint8 [H, W, 3] as f32 in 0..255, laid out [3, H, W]
rows = ceil(H / 40);  cols = W // 40 + 1                     # the trailing column is intentional; it is empty (all padding) when W % 40 == 0
for i in 0..rows, j in 0..cols (row-major):
    tile = img[:, 40i : 40i+40, 40j : 40j+40]                # may be narrower/shorter than 40 at the edges
    padded = full([3, 40, 40], -1.0);  padded[:, :th, :tw] = tile
pixels = padded * (1/255);  pixels = (pixels - mean[c]) / std[c]
pixel_values[n] = repeat(pixels, T=2) laid out [T, 40, 40, 3]          # [N, 2, 40, 40, 3], N = rows * cols

No resizing by default (rescale_image_frac is unset in the published processor_config.json; when a config sets it, resize the long edge by frac, capped at max(rescale_image_max_upscaled_long_edge = 2048, long_edge), with Lanczos, before tiling).

Placeholder expansion: the chat template renders one <|content_image|><|unused_200054|> per image part; replace the single <|unused_200054|> (id 200054) with N copies, in image order.

Tower plan (computed from the config, not stored): prime factors of patch_size reversed and cumulatively multiplied give spatial scales h = [5, 10, 20, 40] and channel widths h_ch = ceil(h^2 * 3 / 64) * 64 = [128, 320, 1216, 4800]; the temporal scale t = [2] with t_ch = ceil(40^2 * 3 * 2 / 64) * 64 = 9600. Candidate grid rows (t, h, w, c): [1,1,1,3], [1,5,5,128], [1,10,10,320], [1,20,20,1216], [1,40,40,4800], [2,40,40,9600]. Pick n_layers + 1 rows by assigning linspace(0, ln(40*40*2*3), n_layers + 1) to ln(t*h*w) with a minimum-cost injective assignment (brute force over permutations; the grid has at most 6 rows), then force the first pick to row 0 and the last to the final row. For n_layers = 4 the picks are rows [0, 1, 2, 4, 5]; for n_layers = 1 they are [0, 5]. Layer i maps row s to row e:

t_fold = e.t / s.t;  hw_fold = e.h / s.h;  in = s.c * t_fold * hw_fold^2
out = text_hidden if i == n_layers - 1 else e.c;  add_norm = i != n_layers - 1

Inkling-Small / Inkling (n_layers = 4): layer 0 [75 -> 128] hw_fold 5, layer 1 [512 -> 320] hw_fold 2, layer 2 [5120 -> 4800] hw_fold 4, layer 3 [9600 -> text_hidden] t_fold 2, no norm. 0.6B (n_layers = 1): one layer [9600 -> 1024] with t_fold 2, hw_fold 40.

Forward per tile batch x: [N, T, Hp, Wp, C]:

for each layer: if t_fold > 1 or hw_fold > 1: x = fold(x, t_fold, hw_fold)
                x = linear(x)                                # bias-free
                if add_norm: x = gelu(RMSNorm(x; norm_i))    # exact (erf) GELU
x = RMSNorm(x; final_norm)   ->   [N, text_hidden]           # after the last layer x is [N, 1, 1, 1, hidden]
fold(x, t, hw): [B, T, H, W, C] -> reshape [B, T/t, t, H/hw, hw, W/hw, hw, C] -> transpose (0, 1, 3, 5, 2, 4, 6, 7) -> reshape [B, T/t, H/hw, W/hw, t*hw*hw*C]

Merge: h = embed(ids) (with embed_norm), then write features[n] into the n-th position where ids == 200054 (masked scatter in order, merge_llava); the text decoder runs unchanged.

Weights (original and MLX layouts are identical): model.visual.layers.linear_{j}.weight [out_j, in_j] (j in 0..n_layers), model.visual.layers.norm_{j}.weight [out_j] (j in 0..n_layers-1), model.visual.final_norm.weight [text_hidden]; quantized MLX variants add .scales/.biases next to each linear_{j}.weight. Rename to vision_tower.encoder_layers.{j}.projection.*, vision_tower.encoder_layers.{j}.layer_norm.weight, vision_tower.final_norm.weight.

Implementation plan

  1. src/vision/encoders/inkling_hmlp.rs: InklingVisionConfig (serde, #[serde(alias = "n_channels")] on num_channels), plan_out_scales() returning Vec<[i64; 4]>, InklingHmlpLayer { projection: UnifiedLinear, layer_norm: Option<RMSNorm>, t_fold, hw_fold }, InklingHmlpEncoder { layers, final_norm } with forward(pixel_values: [N, 2, 40, 40, 3]) -> [N, hidden] using mlxcel_core::{reshape, transpose_axes, gelu}.
  2. src/vision/processors/inkling.rs: tile splitter and normalizer producing [N, 2, 40, 40, 3] f32 plus N per image; hook it into the host preprocessor (src/multimodal/host_preprocessor.rs) and the CLI --image path.
  3. src/vision/inkling_vl.rs: InklingVLModel { text: InklingModel, vision: InklingHmlpEncoder, image_token_id } with get_input_embeddings(ids, pixel_values) = embed + embed_norm + merge_llava, forward_with_embeddings, and the LanguageModel / VLM traits used by src/loading/vlm_*.rs; src/loading/vlm_inkling.rs loads text + vision weights (the text sanitize already drops model.visual.*; the VLM loader keeps them and applies the rename above).
  4. Prompt expansion in src/multimodal/vlm_prompt.rs: count tiles per image and expand <|unused_200054|>; the server image_url content part goes through the same path as other LLaVA-style families.
  5. Detection: ModelType::InklingVLM chosen by detect_text_or_vlm when vision_config is present and model.visual.* weights exist (the text arm stays for text-only use of the same checkpoint).
  6. docs/supported-models.md: mark image support.

Validation

(a) Unit tests, src/vision/encoders/inkling_hmlp_tests.rs:

  • plan_out_scales_matches_reference_rows: n_layers = 4 yields rows [0,1,2,4,5] with the four (in, out, t_fold, hw_fold, norm) tuples above; n_layers = 1 yields [9600 -> hidden].
  • fold_orders_channels_t_then_row_then_col_then_c: a [1, 2, 2, 2, 1] tensor with distinct values folds to the expected 8-vector order.
  • tiler_counts_rows_and_cols_including_trailing_column: a 100x80 image gives rows = 3, cols = 3 (= 80 // 40 + 1) and the trailing column is all (-1/255 - mean) / std.
  • hmlp_output_shape_and_finite: random weights, 3 tiles -> [3, hidden], finite.
  • placeholder_expansion_one_token_per_tile in src/multimodal/vlm_prompt_tests.rs.

(b) Real checkpoint: models/Inkling-Small-mlx-4bit (mlx-community/Inkling-Small-mlx-4bit, about 153 GB):

./target/release/mlxcel generate -m models/Inkling-Small-mlx-4bit --image tests/fixtures/cat.jpg -p "What animal is in the picture and what is it doing?" -n 64

Acceptance: the answer names the animal and its action; the token count before the prompt equals rows * cols per image; on the 0.6B checkpoint (models/Inkling-0.6B-A0.6B) the vision features for one tile match a reference run of the same checkpoint within 1e-3 relative (the 1-layer HMLP).

Acceptance criteria

  • HMLP plan, fold order and tile preprocessing match the formulas above (unit tests)
  • An image question on Inkling-Small 4-bit is answered correctly through both the CLI and /v1/chat/completions
  • docs/supported-models.md updated
  • detection table in src/models/detection.rs updated with a test
  • cargo test --workspace --profile test-fast --features metal,accelerate passes
  • cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

Out of scope

Dependencies

Activity

  1. added
    status:blockedBlocked by dependencies or other issues
    type:enhancementNew features, capabilities, or significant additions
    area:modelsModel architectures, weights, loading, metadata
    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)
    platform:macosmacOS (Apple Silicon) specific
    area:cliCommand-line interface / CLI flags
    modelsize:xlarge4-bit checkpoint > 100GB; exceeds a 128GB dev box, needs bigger hardware
    arch:moeSparse mixture-of-experts decoder
    platform:linuxLinux (CUDA / packaging) specific
    and removed
    platform:macosmacOS (Apple Silicon) specific
    platform:linuxLinux (CUDA / packaging) specific
    on Aug 22, 2026
  2. added a commit that references this issue on Aug 31, 2026
    696fe47
  3. self-assigned this
    on Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

arch:moeSparse mixture-of-experts decoderarea:cliCommand-line interface / CLI flagsarea:inferenceGeneration, sampling, decoding (incl. speculative, DRY)area:modelsModel architectures, weights, loading, metadatamodelsize:xlarge4-bit checkpoint > 100GB; exceeds a 128GB dev box, needs bigger hardwaremodeltype:vlmVision-language modelpriority:mediumMedium prioritystatus:doneCompletedtype:enhancementNew features, capabilities, or significant additions

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions