# `scripts.add_mtp_weights`

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

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

## 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
    options:
      members:
        - DEFAULT_MLX_MODEL
        - DEFAULT_SOURCE_MODEL
        - MTP_SHARD_NAME
        - MTP_SHARD_URL
        - find_snapshot_dir
        - download_mtp_shard
        - 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.find_snapshot_dir" markdown="1">
<summary><code>scripts.add_mtp_weights.find_snapshot_dir</code> · function</summary>

```python
scripts.add_mtp_weights.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 #L38-L51](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L38-L51).

</details>

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

```python
scripts.add_mtp_weights.download_mtp_shard(dest_path: Path, source_model: str) -> Path
```

Download the MTP shard using curl with resume support.

**Parameters**

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

**Returns**

- Type: `Path`
- Direct return expressions: `shard_path`

**Exceptions and behavior**

Function `download_mtp_shard` calls `shard_path.exists`, `print`, `shard_path.stat`, `subprocess.run`; can raise `RuntimeError`; returns `shard_path`.
Directly raised exceptions: `RuntimeError`.

[View source #L54-L82](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L54-L82).

</details>

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

```python
scripts.add_mtp_weights.extract_and_quantize_mtp_weights(shard_path: Path, snapshot_dir: Path, quantization_bits: int = 6) -> not annotated
```

Extract MTP weights, quantize, and save to MLX model directory.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `shard_path` | `Path` | `yes` | `none` | Required positional or keyword input. |
| `snapshot_dir` | `Path` | `yes` | `none` | Required positional or keyword input. |
| `quantization_bits` | `int` | `no` | `6` | Optional positional or keyword input; defaults to `6`. |

**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`, `print`, `mx.load`, `str`; returns `(mtp_output_file, list(quantized_weights.keys()))`.
No direct `raise` statement appears in this definition.

[View source #L85-L195](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L85-L195).

</details>

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

```python
scripts.add_mtp_weights.extract_and_quantize_mtp_weights._quantize_one(key, weight) -> not annotated
```

Quantize a single weight, apply norm adjustment, return dict entries.

**Parameters**

| Name | Type | Required | Default | Description |
| --- | --- | --- | --- | --- |
| `key` | `not annotated` | `yes` | `none` | Required positional or keyword input. |
| `weight` | `not annotated` | `yes` | `none` | Required positional or keyword input. |

**Returns**

- Type: `not annotated`
- 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 #L137-L160](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L137-L160).

</details>

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

```python
scripts.add_mtp_weights.update_model_index(snapshot_dir: Path, mtp_keys: list) -> 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` | `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 #L198-L221](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L198-L221).

</details>

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

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

Update config.json to enable MTP.

**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`, `json.dump`, `print`.
No direct `raise` statement appears in this definition.

[View source #L224-L236](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L224-L236).

</details>

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

```python
scripts.add_mtp_weights.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 #L239-L337](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L239-L337).

</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.find_snapshot_dir) | function | `find_snapshot_dir(model_path: str) -> Path` | Find the latest snapshot directory in HF cache structure. | [#L38-L51](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L38-L51) |
| [`download_mtp_shard`](#contract-scripts.add_mtp_weights.download_mtp_shard) | function | `download_mtp_shard(dest_path: Path, source_model: str) -> Path` | Download the MTP shard using curl with resume support. | [#L54-L82](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L54-L82) |
| [`extract_and_quantize_mtp_weights`](#contract-scripts.add_mtp_weights.extract_and_quantize_mtp_weights) | function | `extract_and_quantize_mtp_weights(shard_path: Path, snapshot_dir: Path, quantization_bits: int = 6) -> not annotated` | Extract MTP weights, quantize, and save to MLX model directory. | [#L85-L195](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L85-L195) |
| [`extract_and_quantize_mtp_weights._quantize_one`](#contract-scripts.add_mtp_weights.extract_and_quantize_mtp_weights._quantize_one) | nested function | `extract_and_quantize_mtp_weights._quantize_one(key, weight) -> not annotated` | Quantize a single weight, apply norm adjustment, return dict entries. | [#L137-L160](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L137-L160) |
| [`update_model_index`](#contract-scripts.add_mtp_weights.update_model_index) | function | `update_model_index(snapshot_dir: Path, mtp_keys: list) -> not annotated` | Update model.safetensors.index.json to include MTP weight keys. | [#L198-L221](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L198-L221) |
| [`update_config`](#contract-scripts.add_mtp_weights.update_config) | function | `update_config(snapshot_dir: Path) -> not annotated` | Update config.json to enable MTP. | [#L224-L236](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L224-L236) |
| [`main`](#contract-scripts.add_mtp_weights.main) | function | `main() -> not annotated` | Function `main` calls `argparse.ArgumentParser`, `parser.add_argument`, `parser.parse_args`, `print`. | [#L239-L337](https://github.com/waybarrios/vllm-mlx/blob/a69d47912bcb21d8fe04d48f75fa896b620ffcfa/scripts/add_mtp_weights.py#L239-L337) |