reference/code/Diffulex-mbd-lms/diffulex/strategy/dmax/engine/request.py

import os

from diffulex.config import Config
from diffulex.engine.request import AutoReq
from diffulex.engine.status import DllmReqStatus
from diffulex.sampling_params import SamplingParams
from diffulex.attention.metadata import is_warming_up
from diffulex.strategy_template.token_merging_multi_block.engine.request import (
    TokenMergingMultiBlockReqTemplate,
)

DMAX_FORCE_PREFILL_ACTIVE = os.environ.get("DIFFULEX_DMAX_FORCE_PREFILL_ACTIVE", "0") == "1"


@AutoReq.register("dmax")
class DMaxReq(TokenMergingMultiBlockReqTemplate):
    """Req for DMax-style block diffusion with token merging."""

    def __init__(
        self,
        token_ids: list[int],
        sampling_params: SamplingParams = SamplingParams(),
        config: Config | None = None,
    ):
        super().__init__(token_ids, sampling_params)
        if config is None:
            raise ValueError("DMaxReq requires config to initialize token-merge state.")
        if is_warming_up():
            # Used for warming up token merging
            self.init_token_merging_multi_block(config)

    def lazy_activate(self):
        if not DMAX_FORCE_PREFILL_ACTIVE:
            return super().lazy_activate()

        self.log_status()
        self.status = self.status_history[-1]
        if self.is_pending or self.is_decoding or self.is_prefilling:
            # Keep active-block iterations on the prefix-cache prefill path.
            # This matches reference generate_spd more closely than switching
            # to the 32-token decode path after the first NFE.
            self.status = DllmReqStatus.PREFILLING