from typing import Callable, Optional, Union

import torch
from torch import nn
import torch.nn.functional as F
from functools import partial

from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
    SequenceClassifierOutputWithPast,
    TokenClassifierOutput,
)
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import auto_docstring, can_return_tuple, logging
from chengyue_script.configuration_dfm import DFM_QwenConfig
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
from math import ceil
import time

logger = logging.get_logger(__name__)
use_flex_attention = False


@torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs", dynamic=False)
def fused_flex_attention(q, k, v, mask=None):
    # Ensure dtype consistency required by flex_attention
    target_dtype = q.dtype
    if k.dtype != target_dtype:
        k = k.to(target_dtype)
    if v.dtype != target_dtype:
        v = v.to(target_dtype)
    return flex_attention(q, k, v, block_mask=mask, enable_gqa=True)

def block_diff_mask(b, h, q_idx, kv_idx, block_size=None, n=None):
    """
    Constructs the specialized block diffusion attention mask for training
    composed of three masks:
    - **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks
    - **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context
    - **Block Causal Mask (M_BC)**: Attention to update x0

    Args:
        b, h: Batch and head indices (ignored for mask logic).
        q_idx, kv_idx: Query and Key indices.
        seq_len: Total sequence length.
        block_size: Defines the block structure.

    Returns:
        A boolean attention mask.
    """
    # Indicate whether token belongs to xt or x0
    x0_flag_q = (q_idx >= n)
    x0_flag_kv = (kv_idx >= n)

    # Compute block indices
    block_q = torch.where(x0_flag_q == 1,
                        (q_idx - n) // block_size,
                        q_idx // block_size)
    block_kv = torch.where(x0_flag_kv == 1,
                        (kv_idx - n) // block_size,
                        kv_idx // block_size)

    # **1. Block Diagonal Mask (M_BD) **
    block_diagonal = (block_q == block_kv) & (x0_flag_q == x0_flag_kv)

    # **2. Offset Block-Causal Mask (M_OBC) **
    offset_block_causal = (
    (block_q > block_kv)
    & (x0_flag_kv == 1)
    & (x0_flag_q == 0)
    )

    # **3. Block-Causal Mask (M_BC) **
    block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1)

    # **4. Combine Masks **
    return block_diagonal | offset_block_causal | block_causal

def eval_block_diff_mask(b, h, q_idx, kv_idx, block_size=None):
    # Compute block indices
    block_q = q_idx // block_size
    block_kv = kv_idx // block_size

    return block_q == block_kv



def categorical(probs: torch.Tensor) -> torch.Tensor:
    r"""Categorical sampler according to weights in the last dimension of ``probs`` using :func:`torch.multinomial`.

    Args:
        probs (Tensor): probabilities.

    Returns:
        Tensor: Samples.
    """

    return torch.multinomial(probs.flatten(0, -2), 1, replacement=True).view(
        *probs.shape[:-1]
    )



class DFM_QwenMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        position_ids (`torch.Tensor`, *optional*):
            Deprecated and unused.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    """
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
    **kwargs
):
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights


class DFM_QwenAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: DFM_QwenConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=True)
        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True)
        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        # query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
        if self.training:
            #split q into two parts
            q_1 = query_states[:,:,:query_states.shape[2]//2]
            q_2 = query_states[:,:,query_states.shape[2]//2:]
            #split k into two parts
            k_1 = key_states[:,:,:key_states.shape[2]//2]
            k_2 = key_states[:,:,key_states.shape[2]//2:]
            q_1, k_1 = apply_rotary_pos_emb(q_1, k_1, cos, sin)
            q_2, k_2 = apply_rotary_pos_emb(q_2, k_2, cos, sin)
            query_states = torch.cat((q_1, q_2), dim=-2)
            key_states = torch.cat((k_1, k_2), dim=-2)

        if past_key_value is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

        if not use_flex_attention:
            attention_interface: Callable = eager_attention_forward
            if self.config._attn_implementation != "eager":
                attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

            attn_output, attn_weights = attention_interface(
                self,
                query_states,
                key_states,
                value_states,
                attention_mask,
                dropout=0.0 if not self.training else self.attention_dropout,
                scaling=self.scaling,
                sliding_window=self.sliding_window,  # main diff with Llama
                **kwargs,
            )
        else:
            attn_output = fused_flex_attention(query_states, key_states, value_states, mask=attention_mask if self.training else None)
            attn_output = attn_output.transpose(1, 2).contiguous()

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output


@use_kernel_forward_from_hub("RMSNorm")
class DFM_QwenRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        DFM_QwenRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


class DFM_QwenDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: DFM_QwenConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size

        self.self_attn = DFM_QwenAttention(config=config, layer_idx=layer_idx)

        self.mlp = DFM_QwenMLP(config)
        self.input_layernorm = DFM_QwenRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = DFM_QwenRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.attention_type = config.layer_types[layer_idx]

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
        **kwargs
    ) -> tuple[torch.Tensor]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        # Self Attention
        hidden_states = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states
        return hidden_states


@auto_docstring
class DFM_QwenPreTrainedModel(PreTrainedModel):
    config_class = DFM_QwenConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["DFM_QwenDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    _supports_flex_attn = True
    _supports_cache_class = True
    _supports_quantized_cache = True
    _supports_static_cache = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": DFM_QwenDecoderLayer,
        "attentions": DFM_QwenAttention,
    }

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, DFM_QwenRMSNorm):
            module.weight.data.fill_(1.0)


class DFM_QwenRotaryEmbedding(nn.Module):
    def __init__(self, config: DFM_QwenConfig, device=None):
        super().__init__()
        # BC: "rope_type" was originally "type"
        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
        else:
            self.rope_type = "default"
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
        position_ids_expanded = position_ids[:, None, :].float()

        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling

        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


@auto_docstring
class DFM_QwenModel(DFM_QwenPreTrainedModel):
    def __init__(self, config: DFM_QwenConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [DFM_QwenDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = DFM_QwenRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = DFM_QwenRotaryEmbedding(config=config)
        self.gradient_checkpointing = False
        self.has_sliding_layers = "sliding_attention" in self.config.layer_types

        self.a = config.a
        self.c = config.c
        self.bd_size = config.bd_size

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def metric(self, z):
        z_flattened = z.view(-1, z.shape[-1])
        z = F.normalize(z, p=2, dim=-1)
        z_flattened = F.normalize(z_flattened, p=2, dim=-1)
        embedding = F.normalize(self.embed_tokens.weight, p=2, dim=-1)
        d = (torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
            torch.sum(embedding**2, dim=1) - 2 * \
            torch.einsum('bd,dn->bn', z_flattened, torch.einsum('n d -> d n', embedding))) ** 2
        return d

    def get_prob_distribution(self, z, t):
        b, s = z.shape[:2]
        d = self.metric(z)
        d = d.reshape(b, s, -1)
        beta_t = self.c * ((t / (1 - t)) ** self.a)
        if beta_t.shape[0] == b:
            beta_t = beta_t.reshape(b, 1, 1)
        d = d * (-1) * beta_t
        d = torch.softmax(d, dim=-1)
        return d
    
    def gen_mask(self, seqlen, block_size, B, H):
        if use_flex_attention:
            mask = create_block_mask(
                partial(block_diff_mask, block_size=block_size, n=seqlen),
                B=B, H=H, Q_LEN=seqlen*2, KV_LEN=seqlen*2)
        else:
            indices = torch.arange(seqlen*2)
            mask = block_diff_mask(
                b=None, h=None, 
                q_idx=indices[:, None], 
                kv_idx=indices[None, :], 
                block_size=block_size, 
                n=seqlen)
        return mask
    
    def eval_mask(self, seqlen, block_size):
        indices = torch.arange(seqlen)
        mask = eval_block_diff_mask(
            b=None, h=None, 
            q_idx=indices[:, None], 
            kv_idx=indices[None, :], 
            block_size=block_size)
        return mask
    
    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs
    ) -> BaseModelOutputWithPast:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)
        
        if self.training:
        #     # input_ids: [batch_size, seq_len]
        #     # labels: [batch_size, seq_len]
        #     # 将input_ids reshape为[batch_size*seq_len//bd_size, bd_size]

            noisy_input_ids = input_ids.clone()
            emb_x_1 = inputs_embeds.reshape(input_ids.shape[0] * input_ids.shape[1] // self.bd_size, self.bd_size, -1)
            t = torch.rand(emb_x_1.shape[0], device=input_ids.device)
            prob_x_t = self.get_prob_distribution(emb_x_1, t)
            b, s = prob_x_t.shape[:2]
            x_t = torch.multinomial(prob_x_t.reshape(b*s, -1), num_samples=1, replacement=False)
            x_t = x_t.reshape(labels.shape)

            noisy_input_ids[labels != -100] = x_t[labels != -100]
            noisy_inputs_embeds = self.embed_tokens(noisy_input_ids)
            
            inputs_embeds = torch.cat([noisy_inputs_embeds, inputs_embeds], dim=1)
            attention_mask = self.gen_mask(labels.shape[1], self.bd_size, input_ids.shape[0], self.config.num_attention_heads)
        else:
            attention_mask = self.eval_mask(input_ids.shape[1], self.bd_size)

        # It may already have been prepared by e.g. `generate`
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            # Prepare mask arguments
            mask_kwargs = {
                "config": self.config,
                "input_embeds": inputs_embeds,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            # Create the masks
            causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
            }
            # The sliding window alternating layers are not always activated depending on the config
            if self.has_sliding_layers:
                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)

        hidden_states = inputs_embeds

        # create position embeddings to be shared across the decoder layers
        position_embeddings = self.rotary_emb(hidden_states, position_ids)

        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            hidden_states = decoder_layer(
                hidden_states,
                # attention_mask=causal_mask_mapping[decoder_layer.attention_type],
                attention_mask=attention_mask.to(device=hidden_states.device),
                position_ids=position_ids,
                past_key_value=past_key_values,
                use_cache=use_cache,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **kwargs,
            )

        hidden_states = self.norm(hidden_states)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


@auto_docstring
class DFM_QwenForCausalLM(DFM_QwenPreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]
    _tp_plan = {"lm_head": "colwise_rep"}
    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

    def __init__(self, config):
        super().__init__(config)
        self.model = DFM_QwenModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs
    ) -> CausalLMOutputWithPast:

        outputs: BaseModelOutputWithPast = self.model(
            input_ids=input_ids,
            labels=labels,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            cache_position=cache_position,
            **kwargs,
        )

        hidden_states = outputs.last_hidden_state
        if self.training:
            hidden_states = hidden_states[:, :hidden_states.shape[1]//2, :]
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        loss = None
        if labels is not None:
            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @torch.no_grad()
    def sample(
        self,
        input_ids,
        x_init,
        n_steps,
        tokenizer,
        split_idx,
        dtype_categorical = torch.float32,
        time_grid = torch.tensor([0.0, 1.0]),
        output_file='decoded_output.txt'
    ):
        print('register sample')
        # x1 = input_ids[:, split_idx:]
        # t = torch.tensor([0.5], device=x1.device)
        # prob_x_t = self.model.get_prob_distribution(self.model.embed_tokens(x1), t)
        # b, s = prob_x_t.shape[:2]
        # x_t = torch.multinomial(prob_x_t.reshape(b*s, -1), num_samples=1, replacement=False)
        # x_t = x_t.reshape(b, s)
        # x_init = torch.cat([input_ids[:, :split_idx], x_t], dim=1)
        # time_grid = torch.tensor([0.7, 1.0], device=x_init.device)
        x_init = torch.randint(size=(input_ids.shape[0], 256 - input_ids.shape[1]), high=self.vocab_size, device=self.device)
        # x_init = torch.randint(size=(input_ids.shape[0], 8-input_ids.shape[1]%8), high=self.vocab_size, device=self.device)

        # prompt_length = split_idx
        prompt_length = input_ids.shape[1]
        x_init = torch.cat([input_ids, x_init], dim=1)
        # Initialize the current state `x_t` with the initial state `X_0`.
        time_grid = time_grid.to(device=x_init.device)

        # If step_size is float then t discretization is uniform with step size set by step_size.
        t_init = time_grid[0].item()
        t_final = time_grid[-1].item()
        step_size = (t_final - t_init) / n_steps

        t_discretization = torch.tensor(
            [t_init + step_size * i for i in range(n_steps)] + [t_final],
            device=x_init.device,
        )

        x_t = x_init.clone()

        original_x_t = x_t.clone()
        batch_size = original_x_t.shape[0]
        # clear output file
        with open(output_file, 'w') as f:
            f.write('begin sampling\n')
        for i in range(n_steps):
            t = t_discretization[i : i + 1]
            h = t_discretization[i + 1 : i + 2] - t_discretization[i : i + 1]
            # Sample x_1 ~ p_1|t( \cdot |x_t)
            logits = self.forward(input_ids=x_t).logits
            p_1t = torch.softmax(logits, dim=-1)
            p_1t = torch.cat([p_1t[:, :1, :], p_1t[:, :-1, :]], dim=1)
            # x_1 = categorical(p_1t.to(dtype=dtype_categorical)).reshape(batch_size, -1)
            x_1 = p_1t.argmax(dim=-1)
            x_1[:, :prompt_length] = x_t[:, :prompt_length]
            with open(output_file, 'a') as f:
                f.write('-' * 50 + '\n')
                f.write(f'step {i} t={t}\n')
                f.write('decoded x_1:\n')
                f.write(tokenizer.decode(x_1[0]))
                f.write('\n' + '-' * 50 + '\n')

            # Checks if final step
            if i == n_steps - 1:
                x_t = x_1

            else:
                # The text part
                # Compute p_t(x|x_1)
                emb_x_1 = self.model.embed_tokens(x_1)
                prob_x_t = self.model.get_prob_distribution(emb_x_1, t) 
                prob_x_t = prob_x_t.reshape(-1, prob_x_t.shape[-1])
                # Comptute the metric
                emb_x_t = self.model.embed_tokens(x_t)
                emb_x_t_flattened = F.normalize(emb_x_t.view(-1, emb_x_t.shape[-1]), p=2, dim=-1)
                emb_x_1_flattened = F.normalize(emb_x_1.view(-1, emb_x_1.shape[-1]), p=2, dim=-1)
                distance_x_t_2_x_1 = (torch.sum(emb_x_t_flattened ** 2, dim=1, keepdim=True) + torch.sum(emb_x_1_flattened**2, dim=1, keepdim=True) - 2 * torch.einsum('bd,bd->b', emb_x_t_flattened, emb_x_1_flattened).unsqueeze(1)) ** 2
                distance_x_1_2_x = self.model.metric(emb_x_1)
                distance = F.relu(distance_x_t_2_x_1 - distance_x_1_2_x)
                if t ==0 :
                    d_beta_t = 0
                else:
                    d_beta_t = self.model.c * self.model.a * ((t / (1 - t)) ** (self.model.a - 1)) *  1 / ((1 - t) ** 2)
                # get u
                u = prob_x_t * d_beta_t * distance
                u = u.reshape(x_1.shape[0], x_1.shape[1], -1)

                # Set u_t(x_t|x_t,x_1) = 0
                delta_t = F.one_hot(x_t, num_classes=self.vocab_size)
                u = torch.where(
                    delta_t.to(dtype=torch.bool), torch.zeros_like(u), u
                )

                # Sample x_t ~ u_t( \cdot |x_t,x_1)
                intensity = u.sum(dim=-1)  # Assuming u_t(xt|xt,x1) := 0
                with open(output_file, 'a') as f:
                    f.write(f'step {i} t={t} intensity={intensity}\n')
                mask_jump = torch.rand(
                    size=x_t.shape, device=x_t.device
                ) < 1 - torch.exp(-h * intensity)
                if mask_jump.sum() > 0:
                    x_t[mask_jump] = categorical(
                        u[mask_jump].to(dtype=dtype_categorical)
                    )
                original_x_t[:, prompt_length:] = x_t[:, prompt_length:]
                x_t = original_x_t.clone()
            with open(output_file, 'a') as f:
                f.write('decoded x_t:\n')
                f.write(tokenizer.decode(x_t[0]))
                f.write('\n' + '-' * 50 + '\n')

            t = t + h

        return x_t


    @torch.no_grad()
    def bd_sample(
        self,
        input_ids,
        n_steps,
        tokenizer,
        block_size,
        max_tokens, 
        dtype_categorical = torch.float32,
        time_grid = torch.tensor([0.0, 1.0]),
        output_file='decoded_output.txt'
    ):
        self.model.bd_size = block_size
        with open(output_file, 'w') as f:
            f.write('begin bd sampling\n')

        num_blocks = max_tokens // block_size
        for block_idx in range(num_blocks):
            x_init = torch.randint(size=(input_ids.shape[0], block_size-input_ids.shape[1]%block_size), high=self.vocab_size, device=self.device)

            # prompt_length = split_idx
            prompt_length = input_ids.shape[1]
            x_init = torch.cat([input_ids, x_init], dim=1)
            # Initialize the current state `x_t` with the initial state `X_0`.
            time_grid = time_grid.to(device=x_init.device)

            # If step_size is float then t discretization is uniform with step size set by step_size.
            t_init = time_grid[0].item()
            t_final = time_grid[-1].item()
            step_size = (t_final - t_init) / n_steps

            t_discretization = torch.tensor(
                [t_init + step_size * i for i in range(n_steps)] + [t_final],
                device=x_init.device,
            )

            x_t = x_init.clone()

            original_x_t = x_t.clone()
            batch_size = x_t.shape[0]
            # clear output file

            for i in range(n_steps):
                t = t_discretization[i : i + 1]
                h = t_discretization[i + 1 : i + 2] - t_discretization[i : i + 1]
                # Sample x_1 ~ p_1|t( \cdot |x_t)
                logits = self.forward(input_ids=x_t).logits
                p_1t = torch.softmax(logits, dim=-1)
                p_1t = torch.cat([p_1t[:, :1, :], p_1t[:, :-1, :]], dim=1)
                x_1 = categorical(p_1t.to(dtype=dtype_categorical)).reshape(batch_size, -1)
                # x_1 = p_1t.argmax(dim=-1)
                x_1[:, :prompt_length] = x_t[:, :prompt_length]
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write(f'step {i} t={t}\n')
                    f.write('decoded x_1:\n')
                    f.write(tokenizer.decode(x_1[0]))
                    f.write('\n' + '-' * 50 + '\n')

                # Checks if final step
                if i == n_steps - 1:
                    x_t = x_1

                else:
                    # The text part
                    # Compute p_t(x|x_1)

                    emb_x_1 = self.model.embed_tokens(x_1)
                    prob_x_t = self.model.get_prob_distribution(emb_x_1, t) 
                    prob_x_t = prob_x_t.reshape(-1, prob_x_t.shape[-1])
                    # Comptute the metric
                    emb_x_t = self.model.embed_tokens(x_t)
                    emb_x_t_flattened = F.normalize(emb_x_t.view(-1, emb_x_t.shape[-1]), p=2, dim=-1)
                    emb_x_1_flattened = F.normalize(emb_x_1.view(-1, emb_x_1.shape[-1]), p=2, dim=-1)
                    distance_x_t_2_x_1 = (torch.sum(emb_x_t_flattened ** 2, dim=1, keepdim=True) + torch.sum(emb_x_1_flattened**2, dim=1, keepdim=True) - 2 * torch.einsum('bd,bd->b', emb_x_t_flattened, emb_x_1_flattened).unsqueeze(1)) ** 2
                    distance_x_1_2_x = self.model.metric(emb_x_1)
                    distance = F.relu(distance_x_t_2_x_1 - distance_x_1_2_x)
                    if t ==0 :
                        d_beta_t = 0
                    else:
                        d_beta_t = self.model.c * self.model.a * ((t / (1 - t)) ** (self.model.a - 1)) *  1 / ((1 - t) ** 2)
                    # get u
                    u = prob_x_t * d_beta_t * distance
                    u = u.reshape(x_1.shape[0], x_1.shape[1], -1)

                    # Set u_t(x_t|x_t,x_1) = 0
                    delta_t = F.one_hot(x_t, num_classes=self.vocab_size)
                    u = torch.where(
                        delta_t.to(dtype=torch.bool), torch.zeros_like(u), u
                    )

                    # Sample x_t ~ u_t( \cdot |x_t,x_1)
                    intensity = u.sum(dim=-1)  # Assuming u_t(xt|xt,x1) := 0
                    with open(output_file, 'a') as f:
                        f.write(f'step {i} t={t} intensity={intensity}\n')
                    mask_jump = torch.rand(
                        size=x_t.shape, device=x_t.device
                    ) < 1 - torch.exp(-h * intensity)
                    if mask_jump.sum() > 0:
                        x_t[mask_jump] = categorical(
                            u[mask_jump].to(dtype=dtype_categorical)
                        )
                    original_x_t[:, prompt_length:] = x_t[:, prompt_length:]
                    x_t = original_x_t.clone()

                    # x_t = x_1
                with open(output_file, 'a') as f:
                    f.write('decoded x_t:\n')
                    f.write(tokenizer.decode(x_t[0]))
                    f.write('\n' + '-' * 50 + '\n')

                t = t + h
            input_ids = x_t
        return input_ids
    
    @torch.no_grad()
    def mdm_sample(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        mask_id=151665,
        threshold=0.9,
        output_file='decoded_output.txt',
        stop_token=151645,
        top_k=50,
        top_p=0.9,
        repetition_penalty=1.1,
        temperature=1.0
    ):
        self.model.bd_size = block_size
        with open(output_file, 'w') as f:
            f.write('begin mdm sampling with sampling controls\n')

        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]
        for block_idx in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            # 初始化x_init为mask_id
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
            x_t = x_init.clone()
            step = 0
            while True:
                mask_idx = (x_t[:, prompt_length:] == mask_id)
                if mask_idx.sum() == 0:
                    # output = self.forward(input_ids=x_t)
                    # logits = output.logits
                    # # 使用新的采样控制函数
                    # next_token = self.sample_with_sampling_controls(
                    #     logits[:, -1:, :], x_t, top_k, top_p, repetition_penalty, temperature
                    # )
                    # x_t = torch.cat([x_t, next_token], dim=1)
                    # step += 1
                    # with open(output_file, 'a') as f:
                    #     f.write('-' * 50 + '\n')
                    #     f.write(f'step {step}\n')
                    #     f.write('decoded x_t:\n')
                    #     f.write(tokenizer.decode(x_t[0]))
                    #     f.write('\n' + '-' * 50 + '\n')
                    break
                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break
                logits = self.forward(input_ids=x_t).logits

                # 使用采样控制函数获取最可能的token和confidence
                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                # predicted_tokens, confidence = self.sample_with_sampling_controls(
                #     logits[:, prompt_length:, :], x_t[:, prompt_length:], top_k, top_p, repetition_penalty, temperature, return_confidence=True
                # )
                
                # # # 只取prompt_length之后的部分
                # x_1 = predicted_tokens
                # x1_p = confidence

                x_1 = logits.argmax(dim=-1)[:, prompt_length:]
                probs = torch.softmax(logits, dim=-1)[:, prompt_length:]
                x1_p = torch.gather(probs, dim=-1, index=x_1.unsqueeze(-1)).squeeze(-1)
                
                # 选出confidence大于threshold的token
                x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                unmask_idx = (x1_p > threshold)
                if unmask_idx.sum() > 0:
                    x_t[:, prompt_length:][unmask_idx] = x_1[unmask_idx]
                    eos_prob_this_position = None
                    max_eos_prob = None
                else:
                    # 选出confidence最大的那一个token
                    token_position = x1_p.argmax()
                    x_t[:, prompt_length:][0, token_position] = x_1[0, token_position]
                    eos_prob_this_position = probs[0, token_position, stop_token]
                    max_eos_prob = probs[0, :, stop_token].max()
                step += 1
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write(f'step {step}\n')
                    f.write(f'eos_prob_this_position: {eos_prob_this_position}\n')
                    f.write(f'max_eos_prob this block: {max_eos_prob}\n')
                    f.write('decoded x_t:\n')
                    f.write(tokenizer.decode(x_t[0]))
                    f.write('\n' + '-' * 50 + '\n')

            input_ids = x_t
        # 截断stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        return input_ids

    @torch.no_grad()
    def mdm_sample_cache(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        mask_id=151665,
        threshold=0.95,
        output_file='decoded_output.txt',
        stop_token=151645
    ):
        self.model.bd_size = block_size
        with open(output_file, 'w') as f:
            f.write('begin bd sampling\n')

        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]

        if input_ids.shape[1] > block_size:
            output = self.forward(input_ids=input_ids[:, :(input_ids.shape[1] // block_size * block_size)], use_cache=True, update_past_key_values=True)
            logits, past_key_values = output.logits, output.past_key_values
            if input_ids.shape[1] % block_size == 0:
                next_token = logits[:, -1:, :].argmax(dim=-1)
                input_ids = torch.cat([input_ids, next_token], dim=1)
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write('decoded x_t:\n')
                    f.write(tokenizer.decode(input_ids[0]))
                    f.write('\n' + '-' * 50 + '\n')
        else:
            past_key_values = None

        for block_idx in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            # 初始化x_init为mask_id
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
            if block_idx == num_blocks - 1:
                x_init[:, -1] = stop_token
                
            x_init = torch.cat([input_ids, x_init], dim=1)

            x_t = x_init.clone()
            step = 0
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True)
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    x_t = torch.cat([x_t, next_token], dim=1)
                    step += 1
                    with open(output_file, 'a') as f:
                        f.write('-' * 50 + '\n')
                        f.write(f'step {step}\n')
                        f.write(f'past_key_values is None: {past_key_values is None}\n')
                        f.write(f'x_t.shape: {x_t.shape}\n')
                        f.write('decoded x_t:\n')
                        f.write(tokenizer.decode(x_t[0]))
                        f.write('\n' + '-' * 50 + '\n')
                    break
                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break
                logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                p_1t = torch.softmax(logits, dim=-1)
                p_1t = torch.cat([p_1t[:, :1, :], p_1t[:, :-1, :]], dim=1)
                x_1 = p_1t.argmax(dim=-1)
                # 选出p_1t中概率大于threshold的token
                x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                unmask_idx = (x1_p > threshold)

                if unmask_idx.sum() > 0:
                    x_t[:, -block_size:][unmask_idx] = x_1[unmask_idx]
                else:
                    # 选出p_1t中概率最大的那一个token
                    token_position = x1_p.argmax()
                    x_t[:, -block_size:][0, token_position] = x_1[0, token_position]

                step += 1
                with open(output_file, 'a') as f:
                    f.write('-' * 50 + '\n')
                    f.write(f'step {step}\n')
                    f.write(f'past_key_values is None: {past_key_values is None}\n')
                    f.write(f'x_t.shape: {x_t.shape}\n')
                    f.write('decoded x_t:\n')
                    f.write(tokenizer.decode(x_t[0]))
                    f.write('\n' + '-' * 50 + '\n')

            input_ids = x_t
        # 截断stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        return input_ids

    @torch.no_grad()
    def mdm_sample_cache_small_block(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        mask_id=151665,
        threshold=0.95,
        small_block_size=32,
        stop_token=151645,
        temperature=0.7,
        top_p=0.95,
        use_block_cache=False,
        block_cache_refresh_interval=4,
    ):
        self.model.bd_size = block_size
        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]

        if input_ids.shape[1] > block_size:
            output = self.forward(input_ids=input_ids[:, :(input_ids.shape[1] // block_size * block_size)], use_cache=True, update_past_key_values=True)
            logits, past_key_values = output.logits, output.past_key_values
            if input_ids.shape[1] % block_size == 0:
                next_token = logits[:, -1:, :].argmax(dim=-1)
                input_ids = torch.cat([input_ids, next_token], dim=1)
        else:
            past_key_values = None

        num_small_blocks = block_size // small_block_size

        for block_idx in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            # 初始化x_init为mask_id
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
            # if block_idx == num_blocks - 1:
            #     x_init[:, prompt_length:prompt_length+6] = torch.tensor([198, 334, 19357, 21806, 21756, 334], device=self.device)
            #     x_init[:, -3:] = torch.tensor([3417, 13, stop_token], device=self.device)
                
            x_t = x_init.clone()
            block_past_key_values = None
            step = 0
            while True:
                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True)
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    x_t = torch.cat([x_t, next_token], dim=1)
                    break
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, small_block_start_idx:small_block_end_idx].sum() == 0:
                            break
                        if stop_token in x_t[:, prompt_length:]:
                            stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                            if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                                break
                        if use_block_cache:
                            if step % block_cache_refresh_interval == 0 or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
                                logits, block_past_key_values = output.logits, output.block_past_key_values
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                                if block_size == small_block_end_idx:
                                    logits = logits[:, -block_size+small_block_start_idx:]
                                else:
                                    logits = logits[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx]
                            else:
                                if block_size == small_block_end_idx:
                                    logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                else:
                                    logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                        else:
                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                            if block_size == small_block_end_idx:
                                logits = logits[:, -block_size+small_block_start_idx:]
                            else:
                                logits = logits[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx]
                        step += 1
                        # logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                        p_1t = torch.softmax(logits, dim=-1)
                        x_1 = p_1t.argmax(dim=-1)

                        # 选出p_1t中概率大于threshold的token
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx[:, small_block_start_idx:small_block_end_idx], x1_p, -torch.inf)
                        # x1_p[:, small_block_end_idx:] = -torch.inf
                        unmask_idx = (x1_p > threshold)
                        update_start_time = time.time()
                        if unmask_idx.sum() > 0:
                            if block_size == small_block_end_idx:
                                x_t[:, -block_size+small_block_start_idx:][unmask_idx] = x_1[unmask_idx]
                            else:
                                x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][unmask_idx] = x_1[unmask_idx]
                        else:
                            # 选出p_1t中概率最大的那一个token
                            token_position = x1_p.argmax()
                            if block_size == small_block_end_idx:
                                x_t[:, -block_size+small_block_start_idx:][0, token_position] = x_1[0, token_position]
                            else:
                                x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][0, token_position] = x_1[0, token_position]
                        update_end_time = time.time()
                        # print(f"block_idx: {block_idx}, small_block_idx: {small_block_idx}, step: {step}, decoded: {tokenizer.decode(x_t[0])}")

            input_ids = x_t
        # 截断stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        return input_ids


    @torch.no_grad()
    def mdm_sample_cache_batch_small_block(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        small_block_size,
        min_len,
        seq_len,
        mask_id=151665,
        threshold=0.95,
        stop_token=151645,
        use_block_cache=False,
        block_cache_refresh_interval=16,
    ):
        generate_start_time = time.time()
        num_blocks = max_tokens // block_size + seq_len.max().item() // block_size

        if min_len > block_size:
            start_time = time.time()
            output = self.forward(input_ids=input_ids[:, :(min_len // block_size * block_size)], use_cache=True, update_past_key_values=True, block_size=block_size)
            end_time = time.time()
            # print(f"Prefill time: {end_time - start_time} seconds")
            logits, past_key_values = output.logits, output.past_key_values
            if min_len % block_size == 0:
                predict_sample_idx = (seq_len == min_len)
                predict_logits = logits[predict_sample_idx, -1:, :]
                next_token = predict_logits.argmax(dim=-1)
                if input_ids.shape[1] <= min_len:
                    input_ids = torch.cat([input_ids, next_token], dim=1)
                else:
                    input_ids[predict_sample_idx, min_len] = next_token

        else:
            past_key_values = None

        seq_block_idx = seq_len // block_size
        finished_flag = torch.zeros((input_ids.shape[0]), device=self.device, dtype=torch.bool)

        start_block_idx = min_len // block_size
        num_small_blocks = block_size // small_block_size
        for block_idx in range(start_block_idx, num_blocks):
            if finished_flag.all():
                break
            # 如果seq_block_idx都一样，代表所有样本都解码到这个block
            if (seq_block_idx == block_idx).all():
                # 初始化x_init为mask_id
                x_init = mask_id * torch.ones((input_ids.shape[0], block_size-input_ids.shape[1]%block_size), device=self.device, dtype=torch.long)
                x_init = torch.cat([input_ids, x_init], dim=1)
                input_ids = x_init
            else:
                x_init = input_ids[:, :(block_idx + 1)*block_size]

            x_init[finished_flag, -block_size:] = tokenizer.pad_token_id
            x_t = x_init.clone()
            step = 0
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    for sample_idx in range(x_t.shape[0]):
                        if finished_flag[sample_idx] and seq_len[sample_idx] < (block_idx + 1) * block_size:
                            stop_token_idx = (x_t[sample_idx, seq_len[sample_idx]:] == stop_token).nonzero()[0][0]
                            x_t[sample_idx, seq_len[sample_idx]+stop_token_idx+1:] = tokenizer.pad_token_id
                    if finished_flag.all():
                        break
                    start_time = time.time()
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True, block_size=block_size)
                    end_time = time.time()
                    # print(f"Prefill time: {end_time - start_time} seconds")
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    next_token[finished_flag] = tokenizer.pad_token_id
                    x_t = torch.cat([x_t, next_token], dim=1)
                    step += 1

                    break
                
                block_start_time = time.time()
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, small_block_start_idx:small_block_end_idx].sum() == 0:
                            break
                        logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, block_size=block_size).logits
                        p_1t = torch.softmax(logits, dim=-1)
                        p_1t = torch.cat([p_1t[:, :1, :], p_1t[:, :-1, :]], dim=1)
                        x_1 = p_1t.argmax(dim=-1)
                        # 选出p_1t中概率大于threshold的token
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                        x1_p[:, small_block_end_idx:] = -torch.inf
                        
                        # 找到概率大于 threshold 的 mask
                        above_th_mask = x1_p > threshold  # 1 kernel

                        # 找到每行最大位置
                        max_prob_idx = x1_p.argmax(dim=-1)  # 1 kernel
                        max_prob_mask = torch.zeros_like(above_th_mask)
                        max_prob_mask[torch.arange(x_1.shape[0], device=x1_p.device), max_prob_idx] = True  # 1 kernel scatter

                        # 合并条件：大于 threshold 或 最大值位置；并且原来 mask_idx 为 True
                        unmask_idx = (above_th_mask | max_prob_mask) & mask_idx  # 1 kernel

                        # 用 torch.where 避免非连续索引赋值
                        block_view = x_t[:, -block_size:]
                        block_view[:] = torch.where(unmask_idx, x_1, block_view)

                        # 是否到达 stop_token
                        finished_row_flags = ((x_1 == stop_token) & unmask_idx).any(dim=1)  # 1 kernel
                        finished_flag |= finished_row_flags

                        # 在每个迭代开始前，创建事件对象
                        # gen_start = torch.cuda.Event(enable_timing=True)
                        # gen_end = torch.cuda.Event(enable_timing=True)
                        # upd_start = torch.cuda.Event(enable_timing=True)
                        # upd_end = torch.cuda.Event(enable_timing=True)

                        # # 检查 mask_idx
                        # mask_idx = (x_t[:, -block_size:] == mask_id)
                        # if mask_idx[:, small_block_start_idx:small_block_end_idx].sum() == 0:
                        #     break

                        # # ---- Generate 计时 ----
                        # gen_start.record()  # 记录开始
                        # logits = self.forward(
                        #     input_ids=x_t[:, -block_size:],
                        #     use_cache=True,
                        #     past_key_values=past_key_values,
                        #     update_past_key_values=False,
                        #     block_size=block_size
                        # ).logits
                        # gen_end.record()    # 记录结束

                        # p_1t = torch.softmax(logits, dim=-1)
                        # p_1t = torch.cat([p_1t[:, :1, :], p_1t[:, :-1, :]], dim=1)
                        # x_1 = p_1t.argmax(dim=-1)
                        # x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        # x1_p = torch.where(mask_idx, x1_p, -torch.inf)
                        # x1_p[:, small_block_end_idx:] = -torch.inf

                        # # ---- Update 计时 ----
                        # upd_start.record()

                        # above_th_mask = x1_p > threshold
                        # max_prob_idx = x1_p.argmax(dim=-1)
                        # max_prob_mask = torch.zeros_like(above_th_mask)
                        # max_prob_mask[torch.arange(x_1.shape[0], device=x1_p.device), max_prob_idx] = True
                        # unmask_idx = (above_th_mask | max_prob_mask) & mask_idx
                        # block_view = x_t[:, -block_size:]
                        # block_view[:] = torch.where(unmask_idx, x_1, block_view)
                        # finished_row_flags = ((x_1 == stop_token) & unmask_idx).any(dim=1)
                        # finished_flag |= finished_row_flags

                        # upd_end.record()

                        # # ---- 同步后获取时间（毫秒）----
                        # torch.cuda.synchronize()  
                        # gen_time_ms = gen_start.elapsed_time(gen_end)  
                        # upd_time_ms = upd_start.elapsed_time(upd_end)  

                        # print(f"Generate time: {gen_time_ms:.3f} ms")
                        # print(f"Update time: {upd_time_ms:.3f} ms")

                        # for sample_idx in range(x_1.shape[0]):
                        #     if x1_p[sample_idx].max() == -torch.inf:
                        #         continue
                        #     if unmask_idx[sample_idx].sum() > 0:
                        #         x_t[sample_idx, -block_size:][unmask_idx[sample_idx]] = x_1[sample_idx, unmask_idx[sample_idx]]
                        #         finished_flag[sample_idx] = finished_flag[sample_idx] | (x_1[sample_idx, unmask_idx[sample_idx]] == stop_token).any()
                        #     else:
                        #         # 选出p_1t中概率最大的那一个token
                        #         token_position = x1_p[sample_idx].argmax()
                        #         x_t[:, -block_size:][sample_idx, token_position] = x_1[sample_idx, token_position]
                        #         finished_flag[sample_idx] = finished_flag[sample_idx] | (x_1[sample_idx, token_position] == stop_token).any()
                        step += 1
                block_end_time = time.time()
            if finished_flag.all():
                input_ids = x_t
                break
            print(f"block_idx: {block_idx}, block_cnt: {step}, decoding samples: {(seq_block_idx == block_idx).sum()}, finished samples: {finished_flag.sum()}, block time: {block_end_time - block_start_time} seconds")
            input_ids[:, :(block_idx + 1)*block_size] = x_t[:, :-1]
            if (seq_block_idx == block_idx).all():
                input_ids = torch.cat([input_ids, x_t[:, -1:]], dim=1)
            else:
                # try:
                if input_ids.shape[1] <= (block_idx + 1)*block_size:
                    input_ids = x_t
                else:
                    input_ids[seq_block_idx == block_idx, (block_idx + 1)*block_size] = x_t[seq_block_idx == block_idx, (block_idx + 1)*block_size]
                # except:
                #     print(f"input_ids.shape: {input_ids.shape}, x_t.shape: {x_t.shape}")
                #     print(f"seq_block_idx: {seq_block_idx}, block_idx: {block_idx}")
                #     raise ValueError("input_ids and x_t have different shapes")
            seq_block_idx[seq_block_idx == block_idx] = block_idx + 1

        generate_end_time = time.time()
        # print(f"Generate time: {generate_end_time - generate_start_time} seconds")

        return input_ids


    @torch.no_grad()
    def mdm_sample_cache_batch_small_block_with_block_cache(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        small_block_size,
        min_len,
        seq_len,
        mask_id=151665,
        threshold=0.95,
        stop_token=151645,
        use_block_cache=False,
        block_cache_refresh_interval=16,
        dynamic_threshold=False,
        factor=1
    ):
        num_blocks = max_tokens // block_size + seq_len.max().item() // block_size

        if min_len > block_size:
            output = self.forward(input_ids=input_ids[:, :(min_len // block_size * block_size)], use_cache=True, update_past_key_values=True, block_size=block_size)
            logits, past_key_values = output.logits, output.past_key_values
            if min_len % block_size == 0:
                predict_sample_idx = (seq_len == min_len)
                predict_logits = logits[predict_sample_idx, -1:, :]
                next_token = predict_logits.argmax(dim=-1)
                if input_ids.shape[1] <= min_len:
                    input_ids = torch.cat([input_ids, next_token], dim=1)
                else:
                    input_ids[predict_sample_idx, min_len] = next_token.squeeze(dim=-1)

        else:
            past_key_values = None

        seq_block_idx = seq_len // block_size
        finished_flag = torch.zeros((input_ids.shape[0]), device=self.device, dtype=torch.bool)

        start_block_idx = min_len // block_size
        num_small_blocks = block_size // small_block_size
        for block_idx in range(start_block_idx, num_blocks):
            if finished_flag.all():
                break
            # 如果seq_block_idx都一样，代表所有样本都解码到这个block
            if (seq_block_idx == block_idx).all():
                # 初始化x_init为mask_id
                x_init = mask_id * torch.ones((input_ids.shape[0], block_size-input_ids.shape[1]%block_size), device=self.device, dtype=torch.long)
                x_init = torch.cat([input_ids, x_init], dim=1)
                input_ids = x_init
            else:
                x_init = input_ids[:, :(block_idx + 1)*block_size]

            x_init[finished_flag, -block_size:] = tokenizer.pad_token_id
            x_t = x_init.clone()
            step = 0
            block_past_key_values = None
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    for sample_idx in range(x_t.shape[0]):
                        if finished_flag[sample_idx] and seq_len[sample_idx] < (block_idx + 1) * block_size:
                            stop_token_idx = (x_t[sample_idx, seq_len[sample_idx]:] == stop_token).nonzero()[0][0]
                            x_t[sample_idx, seq_len[sample_idx]+stop_token_idx+1:] = tokenizer.pad_token_id
                    if finished_flag.all():
                        break
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True, block_size=block_size)
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    next_token[finished_flag] = tokenizer.pad_token_id
                    x_t = torch.cat([x_t, next_token], dim=1)
                    step += 1

                    break
                
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size

                    start = -block_size + small_block_start_idx
                    end = None if block_size == small_block_end_idx else -block_size + small_block_end_idx
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, start:end].sum() == 0:
                            break

                        if use_block_cache:
                            if step % block_cache_refresh_interval == 0 or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
                                logits, block_past_key_values = output.logits, output.block_past_key_values
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                                logits = logits[:, start:end]
                            else:
                                logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                        else:
                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                            logits = logits[:, start:end]

                        p_1t = torch.softmax(logits, dim=-1)
                        x_1 = p_1t.argmax(dim=-1)
                        # 选出p_1t中概率大于threshold的token
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx[:, small_block_start_idx:small_block_end_idx], x1_p, -torch.inf)
                        
                        if not dynamic_threshold:
                            unmask_idx = (x1_p > threshold)
                            max_prob_idx = x1_p.argmax(dim=-1)
                            unmask_idx[torch.arange(x_1.shape[0]), max_prob_idx] = True
                            unmask_idx = unmask_idx & mask_idx[:, small_block_start_idx:small_block_end_idx]
                            
                        else:
                            unmask_idx = torch.zeros_like(x1_p, dtype=torch.bool, device=x1_p.device)
                            for sample_idx in range(x_1.shape[0]):
                                if x1_p[sample_idx].max() == -torch.inf:
                                    continue
                                
                                num_mask_token = mask_idx[sample_idx, small_block_start_idx:small_block_end_idx].sum().item()
                                ns=list(range(1,num_mask_token+1))
                                es=[factor/(n+1) for n in ns]
                                threshs=[1-e for e in es]
                                threshs[0]=-1

                                # at least one token is transferred
                                sorted_confidence=torch.sort(x1_p[sample_idx],dim=-1,descending=True)[0]
                                for top_i in range(len(threshs)):
                                    if sorted_confidence[top_i]<threshs[top_i]:
                                        break

                                if top_i == len(threshs)-1:
                                    top_i+=1

                                _, select_index = torch.topk(x1_p[sample_idx], k=top_i)
                                unmask_idx[sample_idx, select_index] = True
                                
                        if block_size == small_block_end_idx:
                            x_t[:, -block_size+small_block_start_idx:][unmask_idx] = x_1[unmask_idx]
                        else:
                            x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][unmask_idx] = x_1[unmask_idx]
                        finished_row_flags = ((x_1 == stop_token) & unmask_idx).any(dim=1) # shape: [B]
                        finished_flag = finished_flag | finished_row_flags

                        step += 1

            if finished_flag.all():
                input_ids = x_t
                break
            input_ids[:, :(block_idx + 1)*block_size] = x_t[:, :-1]
            if (seq_block_idx == block_idx).all():
                input_ids = torch.cat([input_ids, x_t[:, -1:]], dim=1)
            else:
                if input_ids.shape[1] <= (block_idx + 1)*block_size:
                    input_ids = x_t
                else:
                    input_ids[seq_block_idx == block_idx, (block_idx + 1)*block_size] = x_t[seq_block_idx == block_idx, (block_idx + 1)*block_size]
            seq_block_idx[seq_block_idx == block_idx] = block_idx + 1

        return input_ids

    @torch.no_grad()
    def mdm_sample_cache_small_block_for_demo(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        mask_id=151665,
        threshold=0.95,
        small_block_size=32,
        stop_token=151645,
        temperature=0.7,
        top_p=0.95,
        use_block_cache=False,
        block_cache_refresh_interval=4,
        output_file='decoded_output.txt',
    ):
        nfe = 0
        self.model.bd_size = block_size
        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]

        with open(output_file, 'w') as f:
            f.write('input_ids per position:\n')
            for i in range(input_ids.shape[1]):
                decoded_token = tokenizer.decode(input_ids[:, i])
                f.write(f'position {i} token: {repr(decoded_token)}')
                f.write('\n')
            f.write('-' * 50 + '\n')

        if input_ids.shape[1] > block_size:
            output = self.forward(input_ids=input_ids[:, :(input_ids.shape[1] // block_size * block_size)], use_cache=True, update_past_key_values=True)
            logits, past_key_values = output.logits, output.past_key_values
            nfe += 1
            if input_ids.shape[1] % block_size == 0:
                next_token = logits[:, -1:, :].argmax(dim=-1)
                input_ids = torch.cat([input_ids, next_token], dim=1)
            with open(output_file, 'a') as f:
                f.write(f'nfe: {nfe}\n')
                f.write('initialization cache\n')
        else:
            past_key_values = None

        num_small_blocks = block_size // small_block_size

        for block_idx in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            # 初始化x_init为mask_id
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
            # if block_idx == num_blocks - 1:
            #     x_init[:, prompt_length:prompt_length+6] = torch.tensor([198, 334, 19357, 21806, 21756, 334], device=self.device)
            #     x_init[:, -3:] = torch.tensor([3417, 13, stop_token], device=self.device)
                
            x_t = x_init.clone()
            block_past_key_values = None
            step = 0
            while True:
                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    nfe += 1
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True)
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    x_t = torch.cat([x_t, next_token], dim=1)
                    with open(output_file, 'a') as f:
                        f.write(f'nfe: {nfe}\n')
                        f.write('calculate cache and produce the first token of the next block\n')
                        f.write('tokens per position:\n')
                        for i in range(x_t.shape[1]):
                            decoded_token = tokenizer.decode(x_t[:, i])
                            f.write(f'position {i} token: {repr(decoded_token)}')
                            f.write('\n')
                        f.write('\n' + '-' * 50 + '\n')
                    break
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, small_block_start_idx:small_block_end_idx].sum() == 0:
                            break
                        if stop_token in x_t[:, prompt_length:]:
                            stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                            if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                                break
                        if use_block_cache:
                            if step % block_cache_refresh_interval == 0 or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
                                nfe += 1
                                with open(output_file, 'a') as f:
                                    f.write(f'nfe: {nfe}\n')
                                    f.write('refresh small block cache\n')
                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
                                logits, block_past_key_values = output.logits, output.block_past_key_values
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                                if block_size == small_block_end_idx:
                                    logits = logits[:, -block_size+small_block_start_idx:]
                                else:
                                    logits = logits[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx]
                            else:
                                nfe += 1
                                with open(output_file, 'a') as f:
                                    f.write(f'nfe: {nfe}\n')
                                    f.write('use small block cache\n')
                                if block_size == small_block_end_idx:
                                    logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                else:
                                    logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                        else:
                            nfe += 1
                            with open(output_file, 'a') as f:
                                f.write(f'nfe: {nfe}\n')
                                f.write('no block cache\n')
                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                            if block_size == small_block_end_idx:
                                logits = logits[:, -block_size+small_block_start_idx:]
                            else:
                                logits = logits[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx]
                        step += 1
                        # logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                        p_1t = torch.softmax(logits, dim=-1)
                        x_1 = p_1t.argmax(dim=-1)

                        # 选出p_1t中概率大于threshold的token
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx[:, small_block_start_idx:small_block_end_idx], x1_p, -torch.inf)
                        # x1_p[:, small_block_end_idx:] = -torch.inf
                        unmask_idx = (x1_p > threshold)
                        update_start_time = time.time()
                        if unmask_idx.sum() > 0:
                            if block_size == small_block_end_idx:
                                x_t[:, -block_size+small_block_start_idx:][unmask_idx] = x_1[unmask_idx]
                            else:
                                x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][unmask_idx] = x_1[unmask_idx]
                        else:
                            # 选出p_1t中概率最大的那一个token
                            token_position = x1_p.argmax()
                            if block_size == small_block_end_idx:
                                x_t[:, -block_size+small_block_start_idx:][0, token_position] = x_1[0, token_position]
                            else:
                                x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][0, token_position] = x_1[0, token_position]
                        update_end_time = time.time()
                        with open(output_file, 'a') as f:
                            f.write('tokens per position:\n')
                            for i in range(x_t.shape[1]):
                                decoded_token = tokenizer.decode(x_t[:, i])
                                f.write(f'position {i} token: {repr(decoded_token)}')
                                f.write('\n')
                            f.write('\n' + '-' * 50 + '\n')
                        # print(f"block_idx: {block_idx}, small_block_idx: {small_block_idx}, step: {step}, decoded: {tokenizer.decode(x_t[0])}")

            input_ids = x_t
        # 截断stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        return input_ids
        
    
    @torch.no_grad()
    def sample_with_top_p(self, logits, top_p=0.95, temperature=1.0):
        """
        Sample tokens from logits using top-p (nucleus) sampling.
        
        Args:
            logits: Tensor of shape [batch_size, seq_len, vocab_size]
            top_p: Float, cumulative probability threshold for nucleus sampling
            temperature: Float, temperature for sampling (0 means greedy)
            
        Returns:
            x_1: Sampled token indices, shape [batch_size, seq_len]
            p_1t: Probability distribution after top-p filtering, shape [batch_size, seq_len, vocab_size]
        """
        # Apply temperature
        if temperature == 0 or temperature < 1e-6:
            # Greedy sampling
            x_1 = logits.argmax(dim=-1)
            p_1t = F.softmax(logits, dim=-1)
            return x_1, p_1t
        
        logits = logits / temperature
        
        # Convert logits to probabilities
        probs = F.softmax(logits, dim=-1)  # [batch_size, seq_len, vocab_size]
        
        # Sort probabilities in descending order
        sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
        
        # Compute cumulative probabilities
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
        
        # Remove tokens with cumulative probability above the threshold
        sorted_indices_to_remove = cumulative_probs > top_p
        
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0
        
        # Scatter sorted tensors back to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(
            dim=-1, index=sorted_indices, src=sorted_indices_to_remove
        )
        
        # Set probabilities of removed indices to 0
        probs = probs.masked_fill(indices_to_remove, 0.0)
        
        # Renormalize probabilities
        probs = probs / probs.sum(dim=-1, keepdim=True)
        
        # Sample from the filtered distribution
        x_1 = torch.multinomial(probs.view(-1, probs.size(-1)), num_samples=1).view(probs.shape[:-1])
        
        return x_1, probs
    
    @torch.no_grad()
    def batch_sample(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        small_block_size,
        min_len,
        seq_len,
        mask_id=151665,
        threshold=0.95,
        stop_token=151645,
        use_block_cache=False,
        top_p=0.95,
        temperature=0
    ):
        num_blocks = max_tokens // block_size + seq_len.max().item() // block_size
        batch_size = input_ids.shape[0]
        nfe = 0
        sample_nfe = {}

        if min_len > block_size:
            output = self.forward(input_ids=input_ids[:, :(min_len // block_size * block_size)], use_cache=True, update_past_key_values=True, block_size=block_size)
            logits, past_key_values = output.logits, output.past_key_values
            if min_len % block_size == 0:
                predict_sample_idx = (seq_len == min_len)
                predict_logits = logits[predict_sample_idx, -1:, :]
                next_token = predict_logits.argmax(dim=-1)
                if input_ids.shape[1] <= min_len:
                    input_ids = torch.cat([input_ids, next_token], dim=1)
                else:
                    input_ids[predict_sample_idx, min_len] = next_token.squeeze(dim=-1)
        else:
            past_key_values = None

        seq_block_idx = seq_len // block_size
        finished_flag = torch.zeros((batch_size), device=self.device, dtype=torch.bool)

        start_block_idx = min_len // block_size
        num_small_blocks = block_size // small_block_size

        sample_indices = torch.arange(batch_size, device=self.device)
        finished_samples = {}
        for block_idx in range(start_block_idx, num_blocks):
            if finished_flag.all():
                break
            # 如果seq_block_idx都一样，代表所有样本都解码到这个block
            if (seq_block_idx == block_idx).all():
                # 初始化x_init为mask_id
                x_init = mask_id * torch.ones((input_ids.shape[0], block_size-input_ids.shape[1]%block_size), device=self.device, dtype=torch.long)
                x_init = torch.cat([input_ids, x_init], dim=1)
                input_ids = x_init
            else:
                x_init = input_ids[:, :(block_idx + 1)*block_size]

            x_init[finished_flag, -block_size:] = tokenizer.pad_token_id
            x_t = x_init.clone()
            step = 0
            block_past_key_values = None
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    for sample_idx in range(x_t.shape[0]):
                        if finished_flag[sample_idx] and seq_len[sample_idx] < (block_idx + 1) * block_size:
                            stop_token_idx = (x_t[sample_idx, seq_len[sample_idx]:] == stop_token).nonzero()[0][0]
                            x_t[sample_idx, seq_len[sample_idx]+stop_token_idx+1:] = tokenizer.pad_token_id
                    if finished_flag.all():
                        break
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True, block_size=block_size)
                    logits, past_key_values = output.logits, output.past_key_values
                    next_token = logits[:, -1:, :].argmax(dim=-1)
                    next_token[finished_flag] = tokenizer.pad_token_id
                    x_t = torch.cat([x_t, next_token], dim=1)
                    nfe += 1
                    step += 1

                    break
                
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size

                    start = -block_size + small_block_start_idx
                    end = None if block_size == small_block_end_idx else -block_size + small_block_end_idx
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, start:end].sum() == 0:
                            break
                        
                        if use_block_cache:
                            if block_past_key_values is None or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
                                logits, block_past_key_values = output.logits, output.block_past_key_values
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                                logits = logits[:, start:end]
                                nfe += 1
                            else:
                                logits = self.forward(input_ids=x_t[:,start:end], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                                nfe += 1
                        else:
                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                            logits = torch.cat([logits[:, :1, :], logits[:, :-1, :]], dim=1)
                            logits = logits[:, start:end]
                            nfe += 1
                        x_1, p_1t = self.sample_with_top_p(logits, top_p=top_p, temperature=temperature)
                        # Select tokens with probability greater than threshold from p_1t
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx[:, start:end], x1_p, -torch.inf)

                        unmask_idx = (x1_p > threshold)
                        max_prob_idx = x1_p.argmax(dim=-1)
                        unmask_idx[torch.arange(x_1.shape[0]), max_prob_idx] = True
                        unmask_idx = unmask_idx & mask_idx[:, start:end]

                        x_t[:, start:end][unmask_idx] = x_1[unmask_idx]

                        finished_row_flags = ((x_1 == stop_token) & unmask_idx).any(dim=1) # shape: [B]
                        finished_flag = finished_flag | finished_row_flags

                        step += 1

            if input_ids.shape[1] ==  x_t.shape[1]:
                input_ids = x_t
            else:
                input_ids[:, :(block_idx + 1)*block_size] = x_t[:, :-1]
                if (seq_block_idx == block_idx).all():
                    input_ids = torch.cat([input_ids, x_t[:, -1:]], dim=1)
                else:
                    if input_ids.shape[1] <= (block_idx + 1)*block_size:
                        input_ids = x_t
                    else:
                        input_ids[seq_block_idx == block_idx, (block_idx + 1)*block_size] = x_t[seq_block_idx == block_idx, (block_idx + 1)*block_size]
            seq_block_idx[seq_block_idx == block_idx] = block_idx + 1
            if finished_flag.any():
                for sample_idx in range(x_t.shape[0]):
                    if finished_flag[sample_idx]:
                        original_idx = sample_indices[sample_idx].item()
                        finished_samples[original_idx] = x_t[sample_idx:sample_idx+1].clone().squeeze(dim=0)
                        # sample_nfe[original_idx] = nfe
                sample_indices = sample_indices[~finished_flag]
                input_ids = input_ids[~finished_flag]
                seq_block_idx = seq_block_idx[~finished_flag]
                seq_len = seq_len[~finished_flag]
                x_t = x_t[~finished_flag]

                for layer_id in range(len(past_key_values)):
                    past_key_values.key_cache[layer_id] = past_key_values.key_cache[layer_id][~finished_flag]
                    past_key_values.value_cache[layer_id] = past_key_values.value_cache[layer_id][~finished_flag]

                finished_flag = finished_flag[~finished_flag]



        # # add not finished samples since max_tokens is reached
        if len(finished_samples) < batch_size:
            for sample_idx in range(x_t.shape[0]):
                original_idx = sample_indices[sample_idx].item()
                finished_samples[original_idx] = x_t[sample_idx:sample_idx+1].clone().squeeze(dim=0)
        
        assert len(finished_samples) == batch_size
        return finished_samples, nfe


    @torch.no_grad()
    def batch_sample_no_token_shift(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        small_block_size,
        min_len,
        seq_len,
        mask_id=151665,
        threshold=0.95,
        stop_token=151645,
        use_block_cache=False,
        top_p=0.95,
        temperature=0
    ):
        num_blocks = max_tokens // block_size + seq_len.max().item() // block_size
        batch_size = input_ids.shape[0]
        nfe = 0
        sample_nfe = {}

        if min_len > block_size:
            output = self.forward(input_ids=input_ids[:, :(min_len // block_size * block_size)], use_cache=True, update_past_key_values=True, block_size=block_size)
            logits, past_key_values = output.logits, output.past_key_values
        else:
            past_key_values = None

        seq_block_idx = seq_len // block_size
        finished_flag = torch.zeros((batch_size), device=self.device, dtype=torch.bool)

        start_block_idx = min_len // block_size
        num_small_blocks = block_size // small_block_size

        sample_indices = torch.arange(batch_size, device=self.device)
        finished_samples = {}
        for block_idx in range(start_block_idx, num_blocks):
            if finished_flag.all():
                break
            # 如果seq_block_idx都一样，代表所有样本都解码到这个block
            if (seq_block_idx == block_idx).all():
                # 初始化x_init为mask_id
                x_init = mask_id * torch.ones((input_ids.shape[0], block_size-input_ids.shape[1]%block_size), device=self.device, dtype=torch.long)
                x_init = torch.cat([input_ids, x_init], dim=1)
                input_ids = x_init
            else:
                x_init = input_ids[:, :(block_idx + 1)*block_size]

            x_init[finished_flag, -block_size:] = tokenizer.pad_token_id
            x_t = x_init.clone()
            step = 0
            block_past_key_values = None
            while True:
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    for sample_idx in range(x_t.shape[0]):
                        if finished_flag[sample_idx] and seq_len[sample_idx] < (block_idx + 1) * block_size:
                            stop_token_idx = (x_t[sample_idx, seq_len[sample_idx]:] == stop_token).nonzero()[0][0]
                            x_t[sample_idx, seq_len[sample_idx]+stop_token_idx+1:] = tokenizer.pad_token_id
                    if finished_flag.all():
                        break
                    output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True, block_size=block_size)
                    past_key_values = output.past_key_values
                    step += 1

                    break
                
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size

                    start = -block_size + small_block_start_idx
                    end = None if block_size == small_block_end_idx else -block_size + small_block_end_idx
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, start:end].sum() == 0:
                            break
                        
                        if use_block_cache:
                            if block_past_key_values is None or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
                                block_past_key_values = output.block_past_key_values
                                logits = logits[:, start:end]
                                nfe += 1
                            else:
                                logits = self.forward(input_ids=x_t[:,start:end], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                nfe += 1
                        else:
                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                            logits = logits[:, start:end]
                            nfe += 1
                        x_1, p_1t = self.sample_with_top_p(logits, top_p=top_p, temperature=temperature)
                        # Select tokens with probability greater than threshold from p_1t
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx[:, start:end], x1_p, -torch.inf)

                        unmask_idx = (x1_p > threshold)
                        max_prob_idx = x1_p.argmax(dim=-1)
                        unmask_idx[torch.arange(x_1.shape[0]), max_prob_idx] = True
                        unmask_idx = unmask_idx & mask_idx[:, start:end]

                        x_t[:, start:end][unmask_idx] = x_1[unmask_idx]

                        finished_row_flags = ((x_1 == stop_token) & unmask_idx).any(dim=1) # shape: [B]
                        finished_flag = finished_flag | finished_row_flags

                        step += 1

            input_ids[:, :(block_idx + 1)*block_size] = x_t
            seq_block_idx[seq_block_idx == block_idx] = block_idx + 1
            if finished_flag.any():
                for sample_idx in range(x_t.shape[0]):
                    if finished_flag[sample_idx]:
                        original_idx = sample_indices[sample_idx].item()
                        finished_samples[original_idx] = x_t[sample_idx:sample_idx+1].clone().squeeze(dim=0)
                        # sample_nfe[original_idx] = nfe
                sample_indices = sample_indices[~finished_flag]
                input_ids = input_ids[~finished_flag]
                seq_block_idx = seq_block_idx[~finished_flag]
                seq_len = seq_len[~finished_flag]
                x_t = x_t[~finished_flag]

                for layer_id in range(len(past_key_values)):
                    past_key_values.key_cache[layer_id] = past_key_values.key_cache[layer_id][~finished_flag]
                    past_key_values.value_cache[layer_id] = past_key_values.value_cache[layer_id][~finished_flag]

                finished_flag = finished_flag[~finished_flag]

        # add not finished samples since max_tokens is reached
        if len(finished_samples) < batch_size:
            for sample_idx in range(x_t.shape[0]):
                original_idx = sample_indices[sample_idx].item()
                finished_samples[original_idx] = x_t[sample_idx:sample_idx+1].clone().squeeze(dim=0)
                sample_nfe[original_idx] = nfe
        
        assert len(finished_samples) == batch_size
        return finished_samples, nfe


    @torch.no_grad()
    def mdm_sample_cache_small_block_no_token_shift(
        self,
        input_ids,
        tokenizer,
        block_size,
        max_tokens, 
        mask_id=151665,
        threshold=0.95,
        small_block_size=32,
        stop_token=151645,
        temperature=0.7,
        top_p=0.95,
        use_block_cache=False,
        block_cache_refresh_interval=4,
    ):
        nfe = 0
        self.model.bd_size = block_size
        num_blocks = max_tokens // block_size
        original_input_length = input_ids.shape[1]

        if input_ids.shape[1] > block_size:
            past_key_values = self.forward(input_ids=input_ids[:, :(input_ids.shape[1] // block_size * block_size)], use_cache=True, update_past_key_values=True).past_key_values
            nfe += 1
        else:
            past_key_values = None

        num_small_blocks = block_size // small_block_size

        for block_idx in range(num_blocks):
            if stop_token in input_ids[:, original_input_length:]:
                break
            prompt_length = input_ids.shape[1]
            x_init = mask_id * torch.ones((input_ids.shape[0], block_size-prompt_length%block_size), device=self.device, dtype=torch.long)
            x_init = torch.cat([input_ids, x_init], dim=1)
                
            x_t = x_init.clone()
            block_past_key_values = None
            step = 0
            while True:
                if stop_token in x_t[:, prompt_length:]:
                    stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                    if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                        break
                mask_idx = (x_t[:, -block_size:] == mask_id)
                # 解码完整的一个block，更新cache，并且生成下一个token
                if mask_idx.sum() == 0:
                    nfe += 1
                    past_key_values = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=True).past_key_values
                    break
                for small_block_idx in range(num_small_blocks):
                    small_block_start_idx = small_block_idx * small_block_size
                    small_block_end_idx = small_block_start_idx + small_block_size
                    while True:
                        mask_idx = (x_t[:, -block_size:] == mask_id)
                        if mask_idx[:, small_block_start_idx:small_block_end_idx].sum() == 0:
                            break
                        if stop_token in x_t[:, prompt_length:]:
                            stop_token_idx = (x_t[:, prompt_length:] == stop_token).nonzero()[0][1]
                            if (x_t[:, prompt_length:prompt_length+stop_token_idx] == mask_id).sum() == 0:
                                break
                        if use_block_cache:
                            if step % block_cache_refresh_interval == 0 or (x_t[:, -block_size+small_block_start_idx] == mask_id).any():
                                nfe += 1
                                output = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True)
                                logits, block_past_key_values = output.logits, output.block_past_key_values
                                if block_size == small_block_end_idx:
                                    logits = logits[:, -block_size+small_block_start_idx:]
                                else:
                                    logits = logits[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx]
                            else:
                                nfe += 1
                                if block_size == small_block_end_idx:
                                    logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                                else:
                                    logits = self.forward(input_ids=x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx], use_cache=True, past_key_values=past_key_values, update_past_key_values=False, use_block_cache=True, block_past_key_values=block_past_key_values, replace_position=small_block_start_idx).logits
                        else:
                            nfe += 1
                            logits = self.forward(input_ids=x_t[:, -block_size:], use_cache=True, past_key_values=past_key_values, update_past_key_values=False).logits
                            if block_size == small_block_end_idx:
                                logits = logits[:, -block_size+small_block_start_idx:]
                            else:
                                logits = logits[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx]
                        step += 1
                        p_1t = torch.softmax(logits, dim=-1)
                        x_1 = p_1t.argmax(dim=-1)

                        # 选出p_1t中概率大于threshold的token
                        x1_p = torch.squeeze(torch.gather(p_1t, dim=-1, index=torch.unsqueeze(x_1, -1)), -1)
                        x1_p = torch.where(mask_idx[:, small_block_start_idx:small_block_end_idx], x1_p, -torch.inf)
                        unmask_idx = (x1_p > threshold)
                        if unmask_idx.sum() > 0:
                            if block_size == small_block_end_idx:
                                x_t[:, -block_size+small_block_start_idx:][unmask_idx] = x_1[unmask_idx]
                            else:
                                x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][unmask_idx] = x_1[unmask_idx]
                        else:
                            # 选出p_1t中概率最大的那一个token
                            token_position = x1_p.argmax()
                            if block_size == small_block_end_idx:
                                x_t[:, -block_size+small_block_start_idx:][0, token_position] = x_1[0, token_position]
                            else:
                                x_t[:, -block_size+small_block_start_idx:-block_size+small_block_end_idx][0, token_position] = x_1[0, token_position]
                        # print(f"block_idx: {block_idx}, small_block_idx: {small_block_idx}, step: {step}, decoded: {tokenizer.decode(x_t[0])}")

            input_ids = x_t
        # 截断stop_token
        if stop_token in input_ids[:, original_input_length:]:
            stop_token_idx = (input_ids[:, original_input_length:] == stop_token).nonzero()[0][1]
            input_ids = input_ids[:, :stop_token_idx+original_input_length+1]
        return input_ids
        
    