reference/code/Diffulex-mbd-lms/diffulex/sampler/llada2.py

from __future__ import annotations

import os
import torch
import torch.nn.functional as F

from diffulex.mixin import EditSamplerMixin, TokenMergeSamplerMixin
from diffulex.sampler.auto_sampler import AutoSampler
from diffulex.sampler.base import DllmSamplerNoShiftBase


class LLaDA2AcceptedIdsMixin:
    def _compute_accepted_ids(
        self,
        block,
        confidence: torch.Tensor,
        initial_confidence: torch.Tensor,
        sampled_tokens: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        accept_threshold = block.thresholds.accept_threshold
        pre_block_complete = block.prev_block.is_semi_complete if block.prev_block else True

        high_conf_indices = torch.where(initial_confidence > accept_threshold)[0]
        if pre_block_complete:
            if len(high_conf_indices) == 0:
                _, transfer_index = torch.topk(confidence, 1)
                return transfer_index
            transfer_index = torch.tensor([], device=sampled_tokens.device, dtype=torch.long)
            return torch.unique(torch.cat([transfer_index, high_conf_indices]))
        return high_conf_indices


class LLaDA2Sampler(LLaDA2AcceptedIdsMixin, DllmSamplerNoShiftBase):
    def __init__(self, config=None):
        del config
        super().__init__()


class LLaDA2dot1Sampler(EditSamplerMixin, LLaDA2Sampler):
    def __init__(self, config=None):
        super().__init__(config=config)
        self.edit_threshold = float(getattr(config, "edit_threshold", 0.0))
        self.max_post_edit_steps = int(getattr(config, "max_post_edit_steps", 16))
        self.penalty_lambda = float(getattr(config, "penalty_lambda", 0.0))
        self._last_block_state_map: dict[str, dict[str, dict]] = {}

    def _reset_block_state_map(self) -> None:
        self._last_block_state_map = {}

    @staticmethod
    def _sample_argmax(logits: torch.Tensor, temperature: float) -> torch.Tensor:
        if temperature > 0:
            gumbel = -torch.log(-torch.log(torch.rand_like(logits, dtype=torch.float32).clamp_(1e-6, 1 - 1e-6)))
            logits = logits.to(torch.float32) + gumbel * temperature
        return torch.argmax(logits, dim=-1)

    @staticmethod
    def _extract_block_logits(req, req_logits: torch.Tensor, block, is_prefill: bool) -> torch.Tensor | None:
        if req_logits.shape[0] == 0:
            return None
        if is_prefill:
            prefix_offset = int(req.contiguous_in_cache_prefix_len)
            local_start = int(block.start - prefix_offset)
            local_end = int(block.end - prefix_offset)
            if local_start < 0 or local_end > req_logits.shape[0]:
                return None
            return req_logits[local_start:local_end, ...]

        buf_offset = int(block.start - req.dllm_block_buffer.first_running_block.start)
        local_start = buf_offset
        local_end = buf_offset + int(block.block_size)
        if local_start < 0 or local_end > req_logits.shape[0]:
            return None
        return req_logits[local_start:local_end, ...]

    def forward(
        self,
        reqs,
        logits: torch.Tensor,
        temperatures: torch.Tensor,
        top_p=None,
        top_k=None,
        margin_confidence=False,
        neg_entropy=False,
        **kwargs,
    ):
        del top_p, top_k, margin_confidence, neg_entropy
        attn_metadata = self.fetch_attn_metadata()
        split_logits = self._split_logits_per_req(attn_metadata, reqs, logits)

        empty_per_req = {str(req.req_id): {} for req in reqs}
        sample_output = self.output_cls(
            true_local_ids_map=dict(empty_per_req),
            accepted_ids_map=dict(empty_per_req),
            sampled_tokens_map=dict(empty_per_req),
            mask_token_rel_ids_map=dict(empty_per_req),
            confidence_map=dict(empty_per_req),
            initial_confidence_map=dict(empty_per_req),
        )
        return self._postprocess_sample_output(
            reqs=reqs,
            split_logits=split_logits,
            temperatures=temperatures,
            sample_output=sample_output,
            attn_metadata=attn_metadata,
            **kwargs,
        )

    def _build_edit_writes_map(
        self,
        reqs,
        split_logits,
        temperatures: torch.Tensor,
        sample_output,
        attn_metadata,
        **kwargs,
    ) -> dict[str, dict[str, dict[int, int]]]:
        del sample_output, kwargs
        edit_writes_map: dict[str, dict[str, dict[int, int]]] = {}
        self._reset_block_state_map()

        for req_idx, (req, req_logits) in enumerate(zip(reqs, split_logits)):
            req_id_str = str(req.req_id)
            req_edit_writes: dict[str, dict[int, int]] = {}
            req_block_states: dict[str, dict] = {}

            for block_id, block in enumerate(req.dllm_blocks):
                if not block.is_active:
                    continue

                post_edit_steps = int(getattr(block, "post_edit_steps", 0))
                total_steps = int(getattr(block, "total_steps", 0))
                editable_start = int(getattr(block, "editable_start", 0) or 0)
                accept_threshold = float(block.thresholds.accept_threshold)
                mask_id = int(block.mask_token_id)

                block_logits = self._extract_block_logits(
                    req, req_logits, block, attn_metadata.is_prefill[req_idx]
                )
                if block_logits is None or block_logits.shape[0] != int(block.block_size):
                    continue

                block_tokens = torch.tensor(block.token_ids, dtype=torch.long, device=block_logits.device)
                block_size = int(block.block_size)

                # Argmax + confidence over full block
                temperature = float(temperatures[req_idx].item())
                x = self._sample_argmax(block_logits, temperature)
                logits_fp32 = block_logits.to(torch.float32)

                # Penalty lambda: penalize predicting the adjacent previous token
                if self.penalty_lambda > 0 and block_size > 1:
                    prev_ids = block_tokens[:-1]
                    logits_fp32[1:, :].scatter_(
                        1, prev_ids.unsqueeze(-1), -self.penalty_lambda, reduce="add"
                    )
                    x = torch.argmax(logits_fp32, dim=-1)

                top1_logits = logits_fp32.gather(dim=-1, index=x.unsqueeze(-1)).squeeze(-1)
                logsumexp = torch.logsumexp(logits_fp32, dim=-1)
                p = torch.exp(top1_logits - logsumexp)

                mask_index = block_tokens.eq(mask_id)
                has_mask = bool(mask_index.any().item())

                # M2T: mask-to-token transfers
                mask_transfer_index = torch.zeros(block_size, dtype=torch.bool, device=block_tokens.device)
                if has_mask:
                    eligible_mask = mask_index & (
                        torch.arange(block_size, device=block_tokens.device) >= editable_start
                    )
                    confidence_at_mask = torch.where(
                        eligible_mask, p,
                        torch.tensor(-float("inf"), device=p.device, dtype=p.dtype),
                    )
                    mask_transfer_index = confidence_at_mask > accept_threshold
                    if not mask_transfer_index.any():
                        _, select_index = torch.topk(confidence_at_mask, k=1)
                        mask_transfer_index = torch.zeros(block_size, dtype=torch.bool, device=block_tokens.device)
                        mask_transfer_index[select_index] = True
                else:
                    post_edit_steps += 1

                # T2T: token-to-token edits on editable, non-mask positions
                editable_positions = torch.arange(block_size, device=block_tokens.device) >= editable_start
                edit_positions = ~mask_index & editable_positions
                edit_transfer_index = (
                    (p > self.edit_threshold) & (block_tokens != x) & edit_positions
                )

                transfer_index = mask_transfer_index | edit_transfer_index

                # Hard upper bound: block_size + max_post_edit_steps
                max_steps = block_size + self.max_post_edit_steps
                timed_out = total_steps >= max_steps

                # Determine finished state
                finished = False
                if timed_out:
                    finished = True
                elif not transfer_index.any():
                    finished = True
                elif not has_mask and post_edit_steps > self.max_post_edit_steps:
                    finished = True

                # Build block_writes
                block_writes: dict[int, int] = {}
                if timed_out:
                    # Force-fill remaining editable mask positions with argmax
                    eligible_mask = mask_index & (
                        torch.arange(block_size, device=block_tokens.device) >= editable_start
                    )
                    for rel_idx in torch.nonzero(eligible_mask, as_tuple=False).flatten().tolist():
                        block_writes[int(rel_idx)] = int(x[int(rel_idx)].item())
                    # Also include any T2T edits that would have happened
                    for rel_idx in torch.nonzero(edit_transfer_index, as_tuple=False).flatten().tolist():
                        rel_idx_int = int(rel_idx)
                        if rel_idx_int >= editable_start:
                            block_writes[rel_idx_int] = int(x[rel_idx_int].item())
                elif not finished:
                    for rel_idx in torch.nonzero(transfer_index, as_tuple=False).flatten().tolist():
                        rel_idx_int = int(rel_idx)
                        if rel_idx_int >= editable_start:
                            block_writes[rel_idx_int] = int(x[rel_idx_int].item())

                # Persist counters on block
                block.post_edit_steps = post_edit_steps

                # Block state for scheduler
                same_as_previous = not bool(transfer_index.any().item()) and not timed_out
                comparable = ~mask_index & editable_positions
                if comparable.any() and not timed_out:
                    same_token_ratio = float(
                        x[comparable].eq(block_tokens[comparable]).to(torch.float32).mean().item()
                    )
                else:
                    same_token_ratio = 1.0
                all_confident = bool((p >= accept_threshold).all().item()) if p.numel() > 0 else True

                req_block_states[str(block.block_id)] = {
                    "committable": finished or (same_as_previous and not has_mask),
                    "same_as_previous": same_as_previous,
                    "same_token_ratio": same_token_ratio,
                    "all_confident": all_confident,
                }

                if block_writes:
                    req_edit_writes[str(block.block_id)] = block_writes

            edit_writes_map[req_id_str] = req_edit_writes
            self._last_block_state_map[req_id_str] = req_block_states

        return edit_writes_map

    def _postprocess_sample_output(
        self,
        reqs,
        split_logits,
        temperatures: torch.Tensor,
        sample_output,
        attn_metadata,
        **kwargs,
    ):
        sample_output = super()._postprocess_sample_output(
            reqs=reqs,
            split_logits=split_logits,
            temperatures=temperatures,
            sample_output=sample_output,
            attn_metadata=attn_metadata,
            **kwargs,
        )
        sample_output.block_state_map = self._last_block_state_map
        return sample_output


class LLaDA2DMaxSampler(TokenMergeSamplerMixin, LLaDA2dot1Sampler):
    def __init__(self, config=None):
        super().__init__(config=config)
        self._token_merge_mode = str(getattr(config, "token_merge_mode", "dmax_topk"))
        self._enable_token_merge = bool(
            self._token_merge_mode in {"dmax_topk", "iter_smooth_topk"}
            and float(getattr(config, "token_merge_weight", 1.0)) > 0.0
        )
        self._last_block_state_map: dict[str, dict[str, dict]] = {}
        self._fast_prob_path = os.getenv("DIFFULEX_DMAX_SAMPLER_FAST", "1") != "0"
        del config

    def _compute_accepted_ids(
        self,
        block,
        confidence: torch.Tensor,
        initial_confidence: torch.Tensor,
        sampled_tokens: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        del block, confidence, initial_confidence, kwargs
        return torch.empty(0, dtype=torch.long, device=sampled_tokens.device)

    def forward(
        self,
        reqs,
        logits: torch.Tensor,
        temperatures: torch.Tensor,
        top_p=None,
        top_k=None,
        margin_confidence=False,
        neg_entropy=False,
        **kwargs,
    ):
        # DMax derives writes/token-merge state from full block logits in
        # _build_edit_writes_map. Running the generic mask-token sampler first
        # duplicates argmax/softmax work and its accepted-id output is unused.
        del top_p, top_k, margin_confidence, neg_entropy
        attn_metadata = self.fetch_attn_metadata()
        split_logits = self._split_logits_per_req(attn_metadata, reqs, logits)

        empty_per_req = {str(req.req_id): {} for req in reqs}
        sample_output = self.output_cls(
            true_local_ids_map=dict(empty_per_req),
            accepted_ids_map=dict(empty_per_req),
            sampled_tokens_map=dict(empty_per_req),
            mask_token_rel_ids_map=dict(empty_per_req),
            confidence_map=dict(empty_per_req),
            initial_confidence_map=dict(empty_per_req),
        )
        return self._postprocess_sample_output(
            reqs=reqs,
            split_logits=split_logits,
            temperatures=temperatures,
            sample_output=sample_output,
            attn_metadata=attn_metadata,
            **kwargs,
        )

    @staticmethod
    def _sample_argmax(logits: torch.Tensor, temperature: float) -> torch.Tensor:
        if temperature > 0:
            gumbel = -torch.log(-torch.log(torch.rand_like(logits, dtype=torch.float32).clamp_(1e-6, 1 - 1e-6)))
            logits = logits.to(torch.float32) + gumbel * temperature
        return torch.argmax(logits, dim=-1)

    def _extract_block_logits(self, req, req_logits: torch.Tensor, block, is_prefill: bool) -> torch.Tensor | None:
        if req_logits.shape[0] == 0:
            return None
        if is_prefill:
            prefix_offset = int(req.contiguous_in_cache_prefix_len)
            local_start = int(block.start - prefix_offset)
            local_end = int(block.end - prefix_offset)
            if local_start < 0 or local_end > req_logits.shape[0]:
                return None
            return req_logits[local_start:local_end, ...]

        buf_offset = int(block.start - req.dllm_block_buffer.first_running_block.start)
        local_start = buf_offset
        local_end = buf_offset + int(block.block_size)
        if local_start < 0 or local_end > req_logits.shape[0]:
            return None
        return req_logits[local_start:local_end, ...]

    def _build_dmax_block_outputs(
        self,
        block,
        block_tokens: torch.Tensor,
        block_logits: torch.Tensor,
        temperature: float,
    ) -> tuple[dict[int, int], dict[int, dict | None], dict]:
        editable_start = int(getattr(block, "editable_start", 0) or 0)
        if editable_start >= int(block.block_size):
            return {}, {}, {
                "committable": True,
                "same_as_previous": True,
                "same_token_ratio": 1.0,
                "all_confident": True,
            }

        mask_id = int(block.mask_token_id)
        accept_threshold = float(block.thresholds.accept_threshold)
        full_block_before = block_tokens.clone()
        top1_tokens = self._sample_argmax(block_logits, temperature)
        mask_index = full_block_before.eq(mask_id)
        mask_positions = torch.nonzero(mask_index, as_tuple=False).flatten()
        if self._fast_prob_path:
            # Fast path: only compute exact confidence for mask positions.
            top1_confidence = torch.ones(top1_tokens.shape, dtype=torch.float32, device=top1_tokens.device)
            if mask_positions.numel() > 0:
                mask_logits_fp32 = block_logits.index_select(0, mask_positions).to(torch.float32)
                mask_top1 = top1_tokens.index_select(0, mask_positions).unsqueeze(-1)
                mask_top1_logits = mask_logits_fp32.gather(dim=-1, index=mask_top1).squeeze(-1)
                mask_lse = torch.logsumexp(mask_logits_fp32, dim=-1)
                top1_confidence[mask_positions] = torch.exp(mask_top1_logits - mask_lse)
        else:
            logits_fp32 = block_logits.to(torch.float32)
            top1_logits = logits_fp32.gather(dim=-1, index=top1_tokens.unsqueeze(-1)).squeeze(-1)
            logsumexp = torch.logsumexp(logits_fp32, dim=-1)
            top1_confidence = torch.exp(top1_logits - logsumexp)

        target_tokens = full_block_before.clone()
        token_index = full_block_before.ne(mask_id)
        if bool(token_index.any().item()):
            target_tokens[token_index] = top1_tokens[token_index]

        decode_positions = torch.empty(0, dtype=torch.long, device=full_block_before.device)
        below_threshold_positions = torch.empty(0, dtype=torch.long, device=full_block_before.device)
        if bool(mask_index.any().item()):
            mask_confidence = top1_confidence[mask_positions]
            below_threshold = torch.nonzero(mask_confidence < accept_threshold, as_tuple=False).flatten()
            if below_threshold.numel() > 0:
                below_threshold_positions = mask_positions[below_threshold]
            if below_threshold.numel() == 0:
                decode_upto = int(mask_positions.numel())
            elif int(below_threshold[0].item()) == 0:
                decode_upto = 1
            else:
                decode_upto = int(below_threshold[0].item())
            decode_positions = mask_positions[:decode_upto]
            if decode_positions.numel() > 0:
                target_tokens[decode_positions] = top1_tokens[decode_positions]

        block_writes: dict[int, int] = {}
        token_merge_entries: dict[int, dict | None] = {}
        changed_positions = torch.nonzero(target_tokens.ne(full_block_before), as_tuple=False).flatten()
        same_as_previous = not bool(changed_positions.numel())
        comparable_positions = torch.nonzero(
            full_block_before.ne(mask_id) & target_tokens.ne(mask_id),
            as_tuple=False,
        ).flatten()
        if comparable_positions.numel() > 0:
            same_token_ratio = float(
                target_tokens.index_select(0, comparable_positions)
                .eq(full_block_before.index_select(0, comparable_positions))
                .to(torch.float32)
                .mean()
                .item()
            )
        else:
            same_token_ratio = 1.0
        all_confident = bool((top1_confidence >= 0.9).all().item()) if top1_confidence.numel() > 0 else True
        for rel_idx in changed_positions.tolist():
            if int(rel_idx) < editable_start:
                continue
            token = int(target_tokens[rel_idx].item())
            if token != int(full_block_before[rel_idx].item()):
                block_writes[int(rel_idx)] = token

        if self._enable_token_merge:
            non_mask_positions = torch.nonzero(target_tokens.ne(mask_id), as_tuple=False).flatten()
            for rel_idx in non_mask_positions.tolist():
                if int(rel_idx) < editable_start:
                    continue
                token = int(target_tokens[rel_idx].item())
                if self._token_merge_mode == "dmax_topk":
                    descriptor = self._build_manual_token_merge_descriptor(
                        token=token,
                        confidence=float(top1_confidence[rel_idx].item()),
                        mask_id=mask_id,
                    )
                else:
                    row_probs = F.softmax(block_logits[rel_idx].to(torch.float32), dim=-1)
                    descriptor = self._build_token_merge_descriptor(
                        probs=row_probs,
                        token=token,
                        mask_id=mask_id,
                    )
                token_merge_entries[int(rel_idx)] = descriptor

        return block_writes, token_merge_entries, {
            "committable": bool(same_as_previous or all_confident),
            "same_as_previous": bool(same_as_previous),
            "same_token_ratio": same_token_ratio,
            "all_confident": bool(all_confident),
        }

    def _reset_block_state_map(self) -> None:
        self._last_block_state_map = {}

    def _build_edit_writes_map(
        self,
        reqs,
        split_logits,
        temperatures: torch.Tensor,
        sample_output,
        attn_metadata,
        **kwargs,
    ) -> dict[str, dict[str, dict[int, int]]]:
        del sample_output, kwargs
        edit_writes_map: dict[str, dict[str, dict[int, int]]] = {}
        self._reset_token_merge_map()
        self._reset_block_state_map()
        for req_idx, (req, req_logits) in enumerate(zip(reqs, split_logits)):
            req_id_str = str(req.req_id)
            req_edit_writes: dict[str, dict[int, int]] = {}
            req_token_merge: dict[int, dict | None] = {}
            req_block_states: dict[str, dict] = {}
            for block_id, block in enumerate(req.dllm_blocks):
                if not block.is_active:
                    continue
                req_block_states[str(block_id)] = {
                    "committable": False,
                    "same_as_previous": False,
                    "same_token_ratio": 0.0,
                    "all_confident": False,
                }
                block_logits = self._extract_block_logits(req, req_logits, block, attn_metadata.is_prefill[req_idx])
                if block_logits is None or block_logits.shape[0] != int(block.block_size):
                    continue
                for rel_idx in range(int(block.block_size)):
                    req_token_merge.setdefault(int(block.start + rel_idx), None)
                block_tokens = torch.tensor(block.token_ids, dtype=torch.long, device=block_logits.device)
                block_writes, token_merge_entries, block_state = self._build_dmax_block_outputs(
                    block=block,
                    block_tokens=block_tokens,
                    block_logits=block_logits,
                    temperature=float(temperatures[req_idx].item()),
                )
                req_block_states[str(block_id)] = block_state
                if block_writes:
                    req_edit_writes[str(block_id)] = block_writes
                for rel_idx, descriptor in token_merge_entries.items():
                    req_token_merge[int(block.start + rel_idx)] = descriptor
            edit_writes_map[req_id_str] = req_edit_writes
            self._set_token_merge_entries(req_id_str, req_token_merge)
            self._last_block_state_map[req_id_str] = req_block_states
        return edit_writes_map

    def _postprocess_sample_output(
        self,
        reqs,
        split_logits,
        temperatures: torch.Tensor,
        sample_output,
        attn_metadata,
        **kwargs,
    ):
        sample_output = super()._postprocess_sample_output(
            reqs=reqs,
            split_logits=split_logits,
            temperatures=temperatures,
            sample_output=sample_output,
            attn_metadata=attn_metadata,
            **kwargs,
        )
        sample_output.block_state_map = self._last_block_state_map
        return sample_output


def build_llada2_sampler(config=None):
    sampling_mode = str(getattr(config, "sampling_mode", "naive"))
    if getattr(config, "decoding_strategy", None) == "dmax":
        return LLaDA2DMaxSampler(config)
    if sampling_mode == "edit" and getattr(config, "model_name", None) in [
        "llada2",
        "llada2_moe",
        "llada2_mini",
        "llada2dot1_mini",
    ]:
        return LLaDA2dot1Sampler(config)
    if sampling_mode == "naive" and getattr(config, "model_name", None) in ["llada2", "llada2_moe", "llada2_mini"]:
        return LLaDA2Sampler(config)


AutoSampler.register("llada2", build_llada2_sampler, use_full_config=True)
AutoSampler.register("llada2_moe", build_llada2_sampler, use_full_config=True)
AutoSampler.register("llada2_mini", build_llada2_sampler, use_full_config=True)
AutoSampler.register("llada2dot1_mini", build_llada2_sampler, use_full_config=True)
AutoSampler.register("llada2_mini_dmax", build_llada2_sampler, use_full_config=True)