# Multi Token Prediction (MTP)

## Why We Need MTP

MTP boosts inference performance by parallelizing the prediction of multiple tokens, shifting from single-token to multi-token generation. This approach significantly increases generation throughput and achieves multiplicative acceleration in inference speed—all without compromising output quality.

## How to Use MTP

To enable MTP for DeepSeek-V3 models, add the following parameter when starting the service:

--speculative_config ' {"method": "mtp", "num_speculative_tokens": 1, "disable_padded_drafter_batch": False} '

- `num_speculative_tokens`: The number of speculative tokens that enables the model to predict multiple tokens at once, if provided. It will default to the number in the draft model config if present, otherwise, it is required.
- `disable_padded_drafter_batch`: Disable input padding for speculative decoding. If set to True, speculative input batches can contain sequences of different lengths, which may only be supported by certain attention backends. This currently only affects the MTP method of speculation, default is False.

## How It Works

### Module Architecture

```shell
vllm_ascend
├── sample
│   ├── rejection_sample.py
├── spec_decode
│   ├── mtp_proposer.py
└───────────
```

**1. sample**

- *rejection_sample.py*: During decoding, the main model processes the previous round’s output token and the predicted token together (computing 1+k tokens simultaneously). The first token is always correct, while the second token—referred to as the **bonus token**—is uncertain since it is derived from speculative prediction, thus we employ **Greedy Strategy** and **Rejection Sampling Strategy** to determine whether the bonus token should be accepted. The module structure consists of an `AscendRejectionSampler` class with a forward method that implements the specific sampling logic.

```shell
rejection_sample.py
├── AscendRejectionSampler
│   ├── forward
```

**2. spec_decode**

This section encompasses the model preprocessing for spec-decode, primarily structured as follows: it includes loading the model, executing a dummy run, and generating token IDs. These steps collectively form the model data construction and forward invocation for a single spec-decode operation.

- *mtp_proposer.py*: Configure vLLM-Ascend to use speculative decoding where proposals are generated by DeepSeek MTP layer.

```shell
mtp_proposer.py
├── Proposer
│   ├── load_model
│   ├── dummy_run
│   ├── generate_token_ids
│   ├── _prepare_inputs
│   ├── _propose
```

### Algorithm

**1. Rejection Sampling**

- *Greedy Strategy*

Verify whether the token generated by the main model matches the speculative token predicted by MTP in the previous round. If they match exactly, accept the bonus token; otherwise, reject it and any subsequent tokens derived from that speculation.

- *Rejection Sampling Strategy*

This method introduces stochasticity in rejection sampling.

For each draft token, acceptance is determined by verifying whether the inequality `P_target / P_draft ≥ U` holds, where `P_target` represents the probability assigned to the current draft token by the target model, `P_draft` denotes the probability assigned by the draft model, and `U` is a random number sampled uniformly from the interval [0, 1).

The decision logic for each draft token is as follows: if the inequality `P_target / P_draft ≥ U` holds, the draft token is accepted as output; conversely, if `P_target / P_draft < U`, the draft token is rejected.

When a draft token is rejected, a recovery sampling process is triggered where a "recovered token" is resampled from the adjusted probability distribution defined as `Q = max(P_target - P_draft, 0)`. In the current MTP implementation, since `P_draft` is not provided and defaults to 1, the formulas simplify such that token acceptance occurs when `P_target ≥ U` and the recovery distribution becomes `Q = max(P_target - 1, 0)`.

**2. Performance**

If the bonus token is accepted, the MTP model performs inference for (num_speculative + 1) tokens, including original main model output token and bonus token. If rejected, inference is performed for fewer tokens, depending on how many tokens are accepted.

## DFX

### Method Validation

- Currently, the spec_decode scenario only supports methods such as n-gram, EAGLE, EAGLE3, and MTP. If an incorrect parameter is passed for the method, the code will raise an error to alert the user that an incorrect method was provided.

```python
def get_spec_decode_method(method,
                           vllm_config,
                           device,
                           runner):
    if method == "ngram":
        return NgramProposer(vllm_config, device, runner)
    elif method in ["eagle", "eagle3"]:
        return EagleProposer(vllm_config, device, runner)
    elif method == 'mtp':
        return MtpProposer(vllm_config, device, runner)
    else:
        raise ValueError("Unknown speculative decoding method: "
                         f"{method}")
```

### Integer Validation

- The current npu_fused_infer_attention_score operator only supports integers less than 16 per decode round. Therefore, the maximum supported value for MTP is 15. If a value greater than 15 is provided, the code will raise an error and alert the user.

```python
if self.speculative_config:
    spec_token_num = self.speculative_config.num_speculative_tokens
    self.decode_threshold += spec_token_num
    assert self.decode_threshold <= 16, f"decode_threshold exceeded \
        npu_fused_infer_attention_score TND layout's limit of 16, \
        got {self.decode_threshold}"
```

## Limitations

- Due to the fact that only a single layer of weights is exposed in DeepSeek's MTP, the accuracy and performance are not effectively guaranteed in scenarios where MTP > 1 (especially MTP ≥ 3). Moreover, due to current operator limitations, MTP supports a maximum of 15.
- In the fullgraph mode with MTP > 1, the capture size of each ACLGraph must be an integer multiple of (num_speculative_tokens + 1).
