# `scripts.add_mtp_weights_qwen35`

Add MTP (Multi-Token Prediction) weights to an existing MLX Qwen3.5 model.

[View the complete module source at #L1-L470](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L1-L470).

## API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

::: scripts.add_mtp_weights_qwen35
    options:
      members:
        - MODEL_CONFIGS
        - find_snapshot_dir
        - fetch_shard_index
        - identify_mtp_shards
        - download_shards
        - extract_and_quantize_mtp_weights
        - update_model_index
        - update_config
        - main
      filters: []
      show_if_no_docstring: true

## Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.find_snapshot_dir" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.find_snapshot_dir</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.find_snapshot_dir(model_path: str) -> Path
```

Find the latest snapshot directory in HF cache structure.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `model_path` | `str` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `Path`
- Direct return expressions: `Path(model_path)`; `snapshots[-1]`

**Exceptions and behavior**

Function `find_snapshot_dir` calls `Path`, `snapshots_dir.exists`, `(Path(model_path) / 'config.json').exists`, `FileNotFoundError`; can raise `FileNotFoundError`; has 2 explicit return paths.
Directly raised exceptions: `FileNotFoundError`.

[View source #L53-L63](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L53-L63).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.fetch_shard_index" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.fetch_shard_index</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.fetch_shard_index(source_model: str, download_dir: Path) -> dict
```

Fetch model.safetensors.index.json from HuggingFace.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `source_model` | `str` | `yes` | `none` | Required positional or keyword input. |
| `download_dir` | `Path` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `dict`
- Direct return expressions: `json.load(f)`

**Exceptions and behavior**

Function `fetch_shard_index` calls `print`, `subprocess.run`, `str`, `RuntimeError`; can raise `RuntimeError`; returns `json.load(f)`.
Directly raised exceptions: `RuntimeError`.

[View source #L66-L80](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L66-L80).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.identify_mtp_shards" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.identify_mtp_shards</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.identify_mtp_shards(index: dict) -> tuple[dict[str, str], set[str]]
```

Identify which shards contain MTP weights.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `index` | `dict` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `tuple[dict[str, str], set[str]]`
- Direct return expressions: `(mtp_keys, shards_needed)`

**Exceptions and behavior**

Function `identify_mtp_shards` calls `index.get`, `set`, `weight_map.items`, `key.startswith`; returns `(mtp_keys, shards_needed)`.
No direct `raise` statement appears in this definition.

[View source #L83-L98](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L83-L98).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.download_shards" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.download_shards</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.download_shards(shards: set[str], source_model: str, download_dir: Path) -> dict[str, Path]
```

Download required shards using curl with resume support.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `shards` | `set[str]` | `yes` | `none` | Required positional or keyword input. |
| `source_model` | `str` | `yes` | `none` | Required positional or keyword input. |
| `download_dir` | `Path` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `dict[str, Path]`
- Direct return expressions: `shard_paths`

**Exceptions and behavior**

Function `download_shards` calls `sorted`, `shard_path.exists`, `shard_path.stat`, `print`; can raise `RuntimeError`; returns `shard_paths`.
Directly raised exceptions: `RuntimeError`.

[View source #L101-L130](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L101-L130).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights(mtp_keys: dict[str, str], shard_paths: dict[str, Path], snapshot_dir: Path, is_moe: bool, num_experts: int, no_quantize: bool = False) -> not annotated
```

Extract MTP weights from BF16 shards, optionally quantize, and save.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `mtp_keys` | `dict[str, str]` | `yes` | `none` | Required positional or keyword input. |
| `shard_paths` | `dict[str, Path]` | `yes` | `none` | Required positional or keyword input. |
| `snapshot_dir` | `Path` | `yes` | `none` | Required positional or keyword input. |
| `is_moe` | `bool` | `yes` | `none` | Required positional or keyword input. |
| `num_experts` | `int` | `yes` | `none` | Required positional or keyword input. |
| `no_quantize` | `bool` | `no` | `False` | Optional positional or keyword input; defaults to `False`. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `(mtp_output_file, list(quantized_weights.keys()))`

**Exceptions and behavior**

Function `extract_and_quantize_mtp_weights` calls `mx.set_default_device`, `open`, `json.load`, `config.get`; returns `(mtp_output_file, list(quantized_weights.keys()))`.
No direct `raise` statement appears in this definition.

[View source #L133-L272](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L133-L272).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights._quantize_one" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights._quantize_one</code> · nested function</summary>

```python
scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights._quantize_one(key: str, weight: mx.array) -> dict[str, mx.array]
```

Quantize a single weight, apply norm adjustment.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `key` | `str` | `yes` | `none` | Required positional or keyword input. |
| `weight` | `mx.array` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `dict[str, mx.array]`
- Direct return expressions: `{key: weight}`; `{key: q_w, key.replace('.weight', '.scales'): q_s, key.replace('.weight', '.biases'): q_b}`

**Exceptions and behavior**

Nested Function `extract_and_quantize_mtp_weights._quantize_one` calls `any`, `key.endswith`, `mx.eval`, `print`; has 2 explicit return paths.
No direct `raise` statement appears in this definition.

[View source #L204-L229](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L204-L229).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.update_model_index" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.update_model_index</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.update_model_index(snapshot_dir: Path, mtp_keys: list[str]) -> not annotated
```

Update model.safetensors.index.json to include MTP weight keys.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `snapshot_dir` | `Path` | `yes` | `none` | Required positional or keyword input. |
| `mtp_keys` | `list[str]` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- Direct return expressions: `None`

**Exceptions and behavior**

Function `update_model_index` calls `index_path.exists`, `print`, `open`, `json.load`; returns `None`.
No direct `raise` statement appears in this definition.

[View source #L275-L294](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L275-L294).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.update_config" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.update_config</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.update_config(snapshot_dir: Path) -> not annotated
```

Update config.json to signal MTP availability.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `snapshot_dir` | `Path` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Function `update_config` calls `open`, `json.load`, `config.get`, `text_config.get`.
No direct `raise` statement appears in this definition.

[View source #L297-L322](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L297-L322).

</details>

<details class="api-contract" id="contract-scripts.add_mtp_weights_qwen35.main" markdown="1">
<summary><code>scripts.add_mtp_weights_qwen35.main</code> · function</summary>

```python
scripts.add_mtp_weights_qwen35.main() -> not annotated
```

Function `main` calls `argparse.ArgumentParser`, `parser.add_argument`, `parser.parse_args`, `print`.

**Parameters**

This callable has no explicit inputs.

**Returns**

- Type: `not annotated`

**Exceptions and behavior**

Function `main` calls `argparse.ArgumentParser`, `parser.add_argument`, `parser.parse_args`, `print`.
No direct `raise` statement appears in this definition.

[View source #L325-L466](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L325-L466).

</details>

## Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

| Symbol | Kind | Signature and inputs | What it does | Source |
| --- | --- | --- | --- | --- |
| [`find_snapshot_dir`](#contract-scripts.add_mtp_weights_qwen35.find_snapshot_dir) | function | `find_snapshot_dir(model_path: str) -> Path` | Find the latest snapshot directory in HF cache structure. | [#L53-L63](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L53-L63) |
| [`fetch_shard_index`](#contract-scripts.add_mtp_weights_qwen35.fetch_shard_index) | function | `fetch_shard_index(source_model: str, download_dir: Path) -> dict` | Fetch model.safetensors.index.json from HuggingFace. | [#L66-L80](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L66-L80) |
| [`identify_mtp_shards`](#contract-scripts.add_mtp_weights_qwen35.identify_mtp_shards) | function | `identify_mtp_shards(index: dict) -> tuple[dict[str, str], set[str]]` | Identify which shards contain MTP weights. | [#L83-L98](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L83-L98) |
| [`download_shards`](#contract-scripts.add_mtp_weights_qwen35.download_shards) | function | `download_shards(shards: set[str], source_model: str, download_dir: Path) -> dict[str, Path]` | Download required shards using curl with resume support. | [#L101-L130](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L101-L130) |
| [`extract_and_quantize_mtp_weights`](#contract-scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights) | function | `extract_and_quantize_mtp_weights(mtp_keys: dict[str, str], shard_paths: dict[str, Path], snapshot_dir: Path, is_moe: bool, num_experts: int, no_quantize: bool = False) -> not annotated` | Extract MTP weights from BF16 shards, optionally quantize, and save. | [#L133-L272](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L133-L272) |
| [`extract_and_quantize_mtp_weights._quantize_one`](#contract-scripts.add_mtp_weights_qwen35.extract_and_quantize_mtp_weights._quantize_one) | nested function | `extract_and_quantize_mtp_weights._quantize_one(key: str, weight: mx.array) -> dict[str, mx.array]` | Quantize a single weight, apply norm adjustment. | [#L204-L229](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L204-L229) |
| [`update_model_index`](#contract-scripts.add_mtp_weights_qwen35.update_model_index) | function | `update_model_index(snapshot_dir: Path, mtp_keys: list[str]) -> not annotated` | Update model.safetensors.index.json to include MTP weight keys. | [#L275-L294](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L275-L294) |
| [`update_config`](#contract-scripts.add_mtp_weights_qwen35.update_config) | function | `update_config(snapshot_dir: Path) -> not annotated` | Update config.json to signal MTP availability. | [#L297-L322](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L297-L322) |
| [`main`](#contract-scripts.add_mtp_weights_qwen35.main) | function | `main() -> not annotated` | Function `main` calls `argparse.ArgumentParser`, `parser.add_argument`, `parser.parse_args`, `print`. | [#L325-L466](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights_qwen35.py#L325-L466) |