PATCHES.md

# PATCHES.md — MiniMax-H3 MLX 8-bit pipeline modifications (complete port)

All modifications relative to upstream `PipeNetwork/minimax-h3-mlx` base commit `b2f7e4d2b7861cefe68b75e4b59ab81cc4e7c318` (2026-08-10). The first committed patch is `video-lab-8bit-fixes` `7210b93e6df86bf9c7206091c9542b6983c10c30` (2026-08-17, two files); fixes S8/S9/S10 extend the same branch in the working tree (dit.py two-phase, pipeline release + TeaCache hook, text_encoder processor+scatter, teacache.py, generate.py flags, processor/ + vision shards). Everything else is byte-identical to the base commit.

| | |
|---|---|
| Base commit | `b2f7e4d2b7861cefe68b75e4b59ab81cc4e7c318` |
| Branch | `video-lab-8bit-fixes` (+ working-tree S8/S9/S10 on `video-lab`) |
| Patch commit (committed) | `7210b93e6df86bf9c7206091c9542b6983c10c30` |
| Date (committed) | 2026-08-17 |
| Working-tree S8/S9/S10 | 2026-08-18–2026-08-19 (pipeline.py, dit.py extension, teacache.py, text_encoder.py B1/B2, generate.py, processor/ + vision shards) |
| Validated with | Real generations incl. Turbo LoRA merges + I2V encode checks (see README Validation report + tmp/S9_teacache_spike_RESULT.md) |
| License | Port: Apache-2.0 (PipeNetwork). Weights: MiniMax H3 Community License (`LICENSE`). Turbo LoRA: Apache-2.0 (larryvrh). |

---

## (a) `minimax_h3_mlx/dit.py` — DiT attention QKV blocked layout + TeaCache two-phase hooks

**File:** `minimax_h3_mlx/dit.py` (~20 lines → ~230 lines added vs base)

**QKV blocked layout (S7, committed 7210b93, lines ~12–17 docstring + lines 153–165 in `Attention.__call__`):**
- Raw-checkpoint `attn.qkv_proj` rows are **blocked** `[q | k | v]` — three contiguous thirds of width `heads * head_dim` — not per-head interleaved `[h0: q,k,v][h1: q,k,v]...`. The old code reshaped to `(B, S, heads, 3, head_dim)` and indexed `qkv[...,0]` which scrambled heads. Fix splits and reshapes each third independently (same layout `mlx-serve` consumes via `splitEqual(qkv,3)`; video VAE `to_qkv` remains per-head interleaved — different codebase, different convention):
  ```python
  # minimax_h3_mlx/dit.py:160-164 (Attention)
  qkv = self.qkv_proj(x)
  q, k, v = mx.split(qkv, 3, axis=-1)
  q = q.reshape(B, S, self.heads, self.head_dim)
  k = k.reshape(B, S, self.heads, self.head_dim)
  v = v.reshape(B, S, self.heads, self.head_dim)
  ```
- Module docstring (lines 12–17, 157–158) now documents the blocked layout and the VAE contrast.

**TeaCache two-phase forward (S9, working-tree, lines ~335–510):**
- New helpers `TYPE_CHECKING` import of `ModulationCache`, `_resolve_cache_layer` (lines 347–354), `_prepare_stages` (355–391, shared input projections + token_refiner + packed buffer + temb + adaln indices so the prefix graph is built exactly once), `_run_blocks` (393–409, iterates `blocks[i]` with optional `ModulationCache.get(i)`), `_finish_heads` (411–423, `final_layer.norm_out` + `video_out`/`audio_out` heads).
- Public two-phase entry points: `forward_to_cache_layer` (425–464, partial forward up to and including `cache_layer` — default last block `len(blocks)-1` — returning `(hidden, ctx)` where `hidden` is pre-final-norm residual after that block and `ctx` carries rotary/temb/adaln/layer bookkeeping) and `finish_from_cache_layer` (466–487, completes from cached `hidden`/`ctx` without recomputing prefix, returns `(video_velocity, audio_velocity)`). `__call__` (489–510) now delegates to `_prepare_stages`/`_run_blocks`/`_finish_heads` and gains `return_hidden` for probing. Saving per skip = suffix `blocks[layer+1..49]` + `final_layer.norm_out` + both heads — with default last block only heads+norm are saved (~<5% wall); moving to block 40 saves ~20% blocks per skip (see S9 spike result: 0/7 skips at last block, no wall saving).

---

## (b) `minimax_h3_mlx/text_encoder.py` — 8-bit quantized loader + processor fallback (B1) + positional scatter (B2)

**File:** `minimax_h3_mlx/text_encoder.py` (+84 lines vs base, ~6 → ~17 imports)

**8-bit quantized loader (S7, committed 7210b93, lines ~64–110 + ~143–180):**
- Detects MLX affine 8-bit pack by presence of `quant_config.json` (line 71: `self.quantized = (model_dir / "quant_config.json").exists()`). Before loading, quantizes the module tree so packed U32/scales/biases key 1:1 — same recipe `load.py` replays from `quant_config.json`:
  ```python
  # lines 90-92
  nn.quantize(self.language, group_size=64, bits=8,
              class_predicate=lambda _path, m: isinstance(m, nn.Linear))
  ```
- Vision tower: if `load_vision=True` and `vision_quant_config.json` exists, replays its recorded `quantized` paths (89 tensors) with same group/bits (lines 98–110):
  ```python
  with open(model_dir / "vision_quant_config.json") as fh: vq = json.load(fh)
  paths = set(vq.get("quantized", ()))
  nn.quantize(self.vision, group_size=vq.get("group_size",64), bits=vq.get("bits",8),
              class_predicate=lambda _p, m: isinstance(m, nn.Linear) and _p in paths)
  ```
- Quantized shards loaded raw (no `astype(dtype)`, lines 166–170: `buckets[bucket][path] = tensor` when `self.quantized` else `tensor.astype(dtype)`), preserving packed integers. Only dense path casts.
- Serve pack omits `model.norm` (never evaluated — H3 conditions pre-norm); loader fabricates zeros like dense converter (lines 176–180):
  ```python
  buckets["language"]["norm.weight"] = mx.zeros(tuple(module.norm.weight.shape), dtype=mx.bfloat16)
  ```

**Processor property with torch-free fallback (B1, S10, lines 207–233):**
- `@property processor` (lines 207–233) tries `AutoProcessor.from_pretrained(processor_dir)` (full Qwen3VLProcessor with torch video/image sub-processors). On `ImportError` (no torch/torchvision; `Qwen3VLVideoProcessor` hard-requires them, and transformers 5.15 `AutoImageProcessor` is itself gated), falls back to PIL-only `Qwen2VLImageProcessorPil`:
  ```python
  try: self._processor = AutoProcessor.from_pretrained(processor_dir)
  except Exception:
      try: image_processor = AutoImageProcessor.from_pretrained(processor_dir)
      except Exception:
          from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import Qwen2VLImageProcessorPil
          image_processor = Qwen2VLImageProcessorPil.from_pretrained(processor_dir)
      self._processor = SimpleNamespace(image_processor=image_processor)
  ```
  The facade only ever reads `.image_processor` (`build_request` line 251), tokenizer comes from separate `tokenizer` property (lines 197–205), so image-only processor is sufficient for I2V.

**Encode positional scatter fix (B2, S10, lines 324–344):**
- `encode()` (lines 306–344) previously did `mx.where(image_mask[...,None], hidden[None], inputs_embeds)` which mis-broadcasts `(1,3096,1)` vs `(1,3072,5120)` because vision hidden rows are compact while image pads sit at arbitrary positions in the longer token row — a **positional scatter**, not a broadcast.
- Fix validates `hidden.shape[0] == num_image` (lines 330–334), then does positional REPLACE scatter via `scatter_add` (mlx 0.32 has no `nonzero`/`index_put`, only `at[].add`):
  ```python
  pos_np = np.nonzero(np.array(image_mask))[0]   # line 339 (mlx has no nonzero, use numpy)
  pos = mx.array(pos_np)                         # line 340
  target = inputs_embeds[0]                      # line 341
  inputs_embeds = target.at[pos].add(hidden.astype(target.dtype) - target[pos])  # line 342
  inputs_embeds = inputs_embeds[None]            # line 343
  ```
  Validated with `tmp/s10_validate_encode.py` (1,3096,5120) finite, 3074 video-tagged rows (start+3072 pads+end). `build_request` (lines 237–277) tags whole vision block as `TAG_VIDEO` (not text) — the DiT AdaLN key.

---

## (c) `minimax_h3_mlx/pipeline.py` — `release_text_encoder` headroom (S8/S8b) + TeaCache hook (S9) + checkpoint resume

**File:** `minimax_h3_mlx/pipeline.py` (+~205 lines vs base)

**Release text encoder (S8/S8b, lines 59–73 `__init__`, 77–118 `from_pretrained`, 228–331 `__call__`, 470–532 `_release_text_encoder_now`):**
- `MiniMaxH3Pipeline.__init__` gains `release_text_encoder: bool` (line 66, stored as `self._release_text_encoder` line 73). `from_pretrained` threads it (lines 84, 118).
- In `__call__` (lines 328–331), right after `prompt_embeds` built (line 1, text conditioning + vision rows already encoded, step-4 noise already drawn — trajectory fixed), opt-in `release_text_encoder or self._release_text_encoder` drops the encoder:
  ```python
  if (release_text_encoder or self._release_text_encoder) and getattr(self, "text_encoder", None) is not None:
      self._release_text_encoder_now()
  ```
- `_release_text_encoder_now` (lines 470–532) drops `self.text_encoder` (`del` + `gc.collect()` pattern), sizes correctly via `mlx.utils.tree_flatten` (lines 484,496,510 — previous `module.parameters().values()` summed dicts not arrays → 0.0 GB), calls `mx.clear_cache()` to return unified Metal memory (line 528), logs `released text encoder after conditioning (freeing ~N GB)` (line 532). Second call is no-op (line 471 guard). Safe: inner estimator errors just leave freed-size as "(sizing indeterminado)" while delete still happens. Measured: frees ~27.5 GB / ~22 GB resident before denoise loop; pipeline then only touches DiT+VAEs.
- CLI: `scripts/generate.py` `--release-encoder` (line 56) → `pipeline(release_text_encoder=args.release_encoder)` (line 96).
- Checkpoint resume (bonus, lines 371–389): `resume_from`/`checkpoint_dir` floats inside the loop (not a pipeline-level patch for HF, but present in `pipeline.py` working tree).

**TeaCache hook (S9, lines 43, 228–230, 342–410):**
- Imports `TeaCacheConfig/TeaCacheController` (line 43), `__call__` gains `teacache`, `teacache_config`, `teacache_layer` (lines 228–230 docstring 244–252). Controller instantiated when `teacache=True` (lines 342–350, resolves `tc_layer = num_layers-1 if None else int`, validates range). Per-step two-phase: `hidden, ctx = dit.forward_to_cache_layer(..., cache_layer=tc_layer)` → `mx.eval(hidden)` → `skip = controller.decide(i, hidden)` → on skip reuse `cached_preds = (video_pred, audio_pred)` joint (lines 391–399), else `finish_from_cache_layer` and cache preds (lines 406–410). Log: `step N/T teacache skip (rel_l1=..., thresh=...)` (lines 398–401). Default last-block hook measured **0/7 skips** at 768x448 turbo8 (see teacache.py below → OFF recommended).

**Other:** `load.py` one-line import fix for quantized path (line 1 added), `scripts/generate.py` adds `--teacache*`/`--release-encoder` flags (lines 45–56, 85–96).

---

## (d) `minimax_h3_mlx/teacache.py` — two-phase TeaCache controller (S9, new file)

**File:** `minimax_h3_mlx/teacache.py` (11 KB, new — 223 lines)

- **Intent:** feature-caching (ByteDance/ali-vilab line) — each step runs partial forward to a probe hidden, compares against previous step's feature, reuses previous `(video, audio)` velocity if similar. Joint reuse (both modalities from one forward must be reused together). Noise/keyframe sampling untouched → seed still fixes trajectory.
- **`TeaCacheConfig` (lines 40–71):** `rel_l1_thresh=0.2` (default, ref 0.15 for 10s 544x960), `metric='rel_l1'|'cosine'`, `start_at_step=3`, `compute_last_step=True` (protects first 3 and last step), `cache_type='avg'|'last'`. `__post_init__` validates.
- **Pure helpers (lines 87–134):** `relative_l1_distance` (mean|curr-prev|/mean|prev|, lines 87–100), `cosine_similarity` (lines 103–113), `teacache_gate` (1-min(dist,1) for rel_l1 so higher=more similar, lines 115–124), `teacache_should_skip` (lines 126–134: `rel_l1 <= thresh` or `cosine >= thresh`).
- **`TeaCacheController` (lines 148–242):** stateful per-run (`total_steps`, `cached_feature`, `computed/skipped`, `by_step`). `decide(step_index, feature)` (lines 171–210): `must_compute` if `step==0` or `step<start` or `cache is None` or `last step` → prime cache; else compute `metric_value` (cosine or rel_l1 distance), `skip = metric >=/<= thresh`, smoothing `avg` folds current into cache as `(prev+curr)/2` on skip, else replaces. Stats: `last_gate` (similarity-domain), `last_metric` (raw distance/similarity for honest log), `metric_name`, `to_stats()`.
- **Pipeline integration:** see (c) above; `dit.py` hooks provide the probe feature. **Measurement (M4 Max, 768x448 turbo8, 7 forwards):** default last-block `thresh 0.2` → 0 skips (+9.7% overhead), `0.35` → 1/7 skip (rel_l1 0.297) still no wall saving (partial forward already 49/50 blocks, only `final_layer.norm_out`+heads saved). `tmp/S9_teacache_spike_RESULT.md` — **verdict: TeaCache OFF by default**; aggressive needs earlier hook `layer 40` + `thresh 0.25–0.35` for double-digit % (trades quality).

---

## (e) `processor/` + vision shards (S10 I2V)

**Files added to checkpoint (not code patches but data required for I2V):**
- `processor/` (7 files, ~11.6 MB): `preprocessor_config.json` (390 B), `video_preprocessor_config.json` (385 B), `chat_template.json` (5.5 KB), `tokenizer_config.json` (11 KB), `tokenizer.json` (6.7 MB), `vocab.json` (2.6 MB), `merges.txt` (1.6 MB) — copied from `models/minimax-h3-ckpt/processor/` (`Qwen3VLProcessor`, `Qwen2VLImageProcessorFast` + `Qwen3VLVideoProcessor` configs; runtime uses PIL fallback per (b) B1).
- `text_encoder/model-00005-of-00008-vision.safetensors` (291,980,081 B), `model-00006-of-00008-vision.safetensors` (341,371,441 B), `model-00007-of-00008-vision.safetensors` (129,267,789 B) — 529 `model.visual.*` tensors (vision tower: depth 27, hidden 1152, patch 16, merge 2, out 5120). Built by `scripts/rebuild_h3_text_encoder_vision_8bit.py` which transposes `patch_embed.proj.weight` from torch `(O,C,kD,kH,kW)` to MLX channels-last and quantizes 89 paths.
- `text_encoder/vision_quant_config.json` (2,461 B, 89 quantized paths: `blocks.*.attn.qkv/proj`, `blocks.*.mlp.linear_fc1`, `merger.*`, `deepstack_merger_list.*`).

**Why this matters:** without `processor/` the pipeline cannot build `pixel_values`/`image_grid_thw` (I2V falls back to T2V). Without vision shards `text_encoder.encode(prompt, images)` raises `load_vision=False` or missing-tensor error. With both, `encode()` (b-B2) scatters 3072 vision patches into `inputs_embeds` and hits the DiT as `TAG_VIDEO` rows, then `pipeline._encode_keyframes` patches conditioning latents — FL2VA keyframe path (anchors `first`/`last`) validated end-to-end (see README `## Image-to-Video`).

---

## Other minor port fixes

- `minimax_h3_mlx/load.py`: quantization structure replay for DiT (already in committed 7210b93) — one-line import guard retained.
- `scripts/generate.py`: flags `--teacache`, `--teacache-thresh` (default 0.2), `--teacache-metric` (rel_l1/cosine), `--teacache-layer`, `--teacache-start` (default 3), `--release-encoder` (all S8/S9).
- `tests/test_release_text_encoder.py`, `tests/test_teacache.py`: unit/mocked tests for (c)/(d) (not shipped in HF package, kept in repo).

## Reproducing the diff

```bash
git -C models/minimax-h3-mlx diff b2f7e4d2 -- minimax_h3_mlx/dit.py minimax_h3_mlx/text_encoder.py minimax_h3_mlx/pipeline.py minimax_h3_mlx/teacache.py scripts/generate.py
# committed two-file patch:
git -C models/minimax-h3-mlx show 7210b93 --stat
```

## License

- **Weights**: MiniMax H3 Community License (included as `LICENSE`) — not open source; territorial exclusions apply (EU/UK/KR/US excluded per mlx-serve pack note).
- **Port code**: Apache-2.0 (PipeNetwork/minimax-h3-mlx).
- **Turbo LoRA**: Apache-2.0 (larryvrh/MiniMax-H3-Turbo-Lora).