Class MTPDecodingConfig

tensorrt_llm/llmapi/llm_args.py:1110–1167

class MTPDecodingConfig(DecodingBaseConfig):
    num_nextn_predict_layers: int = 1
    use_relaxed_acceptance_for_thinking: bool = False
    relaxed_topk: int = 1
    relaxed_delta: float = 0.
    use_mtp_vanilla: bool = False
    mtp_eagle_one_model: bool = True

    # TODO: remove this after distinguishing `max_draft_len` and `num_nextn_predict_layers`
    # Now we need a flag when MTPDecodingConfig is updated by PyTorchModelEngine.
    num_nextn_predict_layers_from_model_config: int = 1

    # When encounter <think>, start thinking phase.
    # When encounter </think>, end thinking phase.
    # <think> [thinking phase] </think> [real output]
    begin_thinking_phase_token: int = 128798
    end_thinking_phase_token: int = 128799

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        if 'num_nextn_predict_layers' in kwargs:
            self.max_draft_len = kwargs['num_nextn_predict_layers']
            self.max_total_draft_tokens = kwargs[
                'num_nextn_predict_layers']  # Current MTP only support linear tree

        if not self.mtp_eagle_one_model:
            logger.warning(
                "2-model style MTP is deprecated. The mtp_eagle_one_model flag will do nothing "
                "in release 1.3. After that, the flag will be removed entirely."
            )

    @classmethod
    def from_dict(cls, data: dict):
        out = cls(**data)
        out.max_draft_len = out.num_nextn_predict_layers
        out.max_total_draft_tokens = out.num_nextn_predict_layers  # Current MTP only support linear tree
        return out

    decoding_type: ClassVar[str] = "MTP"

    def supports_backend(self, backend: str) -> bool:
        return backend == "pytorch"

    @functools.cached_property
    def num_capture_layers(self) -> int:
        if not self.use_mtp_vanilla and not self.mtp_eagle_one_model:
            return 1
        return 0

    @functools.cached_property
    def spec_dec_mode(self):
        from tensorrt_llm._torch.speculative.interface import \
            SpeculativeDecodingMode as TorchSpeculativeDecodingMode
        if self.num_nextn_predict_layers_from_model_config == 1 and not self.use_mtp_vanilla and self.mtp_eagle_one_model:
            return TorchSpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL
        elif self.num_nextn_predict_layers_from_model_config == 1 and not self.use_mtp_vanilla and not self.mtp_eagle_one_model:
            return TorchSpeculativeDecodingMode.MTP_EAGLE
        return TorchSpeculativeDecodingMode.MTP