# This file is modified based on https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen3/modeling_qwen3.py.
#
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/qwen3/modular_qwen3.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_qwen3.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# coding=utf-8
# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Callable, Optional, Tuple, Union, List, TypedDict
import math
import time
import sys
import re

import torch
from torch import nn
from einops import rearrange

from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
from transformers.generation import GenerationMixin
from transformers.integrations import use_kernel_forward_from_hub
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
    QuestionAnsweringModelOutput,
    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
try:
    from transformers.utils import LossKwargs
except ImportError:
    class LossKwargs(TypedDict, total=False):
        pass

from transformers.utils import auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging
# Pre-register submodules to handle hyphenated directory names
import importlib.util as _ilu, os as _os, sys as _sys
_dir = _os.path.dirname(_os.path.abspath(__file__))
for _n in ["configuration_sdar", "fused_linear_diffusion_cross_entropy", "gap_sdar_training"]:
    _fqn = f"{__name__.rsplit(chr(46), 1)[0]}.{_n}" if chr(46) in (__name__ or "") else _n
    if _fqn not in _sys.modules:
        _sp = _ilu.spec_from_file_location(_fqn, _os.path.join(_dir, f"{_n}.py"))
        _md = _ilu.module_from_spec(_sp)
        _sys.modules[_fqn] = _md
        _sp.loader.exec_module(_md)
from .configuration_sdar import SDARConfig
from .fused_linear_diffusion_cross_entropy import FusedLinearDiffusionCrossEntropyLoss
from .gap_sdar_training import (
    apply_gap_remask,
    build_rollout_scope_mask,
    build_rollout_p_mask,
    get_num_transfer_tokens,
    select_policy_transfer_tokens,
    select_teacher_forced_rollout_tokens,
)

try:
    from flash_attn.ops.triton.layer_norm import rms_norm_fn as flash_rms_norm
except ImportError:
    flash_rms_norm = None

import torch.nn.functional as F
try:
    from flash_attn import flash_attn_func, flash_attn_varlen_func
    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
except:
    pass

try:
    from liger_kernel.ops.swiglu import LigerSiLUMulFunction  # noqa: F401
    liger_kernel_is_available = True
except ImportError:
    liger_kernel_is_available = False


if is_torch_flex_attn_available():
    from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
    from transformers.integrations.flex_attention import make_flex_block_causal_mask


logger = logging.get_logger(__name__)


def _gap_env_int(name: str, default: int = 0) -> int:
    raw = _os.getenv(name)
    if raw is None:
        return default
    try:
        return int(raw.strip())
    except Exception:
        return default


def _gap_env_flag(name: str, default: bool = False) -> bool:
    raw = _os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def _gap_is_rank0() -> bool:
    if not torch.distributed.is_available() or not torch.distributed.is_initialized():
        return True
    try:
        return torch.distributed.get_rank() == 0
    except Exception:
        return True


def _gap_debug_enabled() -> bool:
    return _gap_env_flag("SDAR_GAP_COARSE_MARKERS", False) and _gap_is_rank0()


def _gap_stderr(message: str) -> None:
    try:
        _os.write(2, (message + "\n").encode("utf-8", errors="replace"))
    except Exception:
        pass


def modify_padded_position_ids_2d(position_ids: torch.LongTensor) -> torch.LongTensor:
    """
    使用完全向量化的 PyTorch 操作修改一个 batch 的 packed position_ids。
    这个函数假设输入是一个 2D Tensor，形状为 (batch_size, sequence_length)。
    它会独立地处理 batch 中的每一行。

    Args:
        position_ids: 二维 PyTorch Tensor, shape (batch_size, sequence_length).

    Returns:
        修改后的 position_ids Tensor, shape (batch_size, sequence_length).
    """
    if position_ids.dim() != 2:
        raise ValueError(f"Input tensor must be 2D, but got {position_ids.dim()} dimensions.")
        
    batch_size, seq_len = position_ids.shape
    device = position_ids.device

    col_indices = torch.arange(seq_len, device=device, dtype=position_ids.dtype).expand(batch_size, -1)
    mask = (position_ids != 0)

    masked_indices = col_indices * mask
    last_nonzero_idx = torch.max(masked_indices, dim=1).values
    has_nonzero = torch.any(mask, dim=1)
    pad_start_idx = torch.where(has_nonzero, last_nonzero_idx + 1, torch.tensor(0, device=device, dtype=position_ids.dtype))

    padding_mask = col_indices >= pad_start_idx.unsqueeze(1)
    new_pad_values = col_indices - pad_start_idx.unsqueeze(1)
    position_ids = torch.where(padding_mask, new_pad_values, position_ids)

    return position_ids


def calculate_token_nums(position_ids: torch.Tensor):
    """
    使用 PyTorch 高效计算一个批次中每个打包序列的长度。

    Args:
        position_ids (torch.Tensor): 一个 2D Tensor，形状为 (batch_size, sequence_length)。
                                     例如：tensor([[0,1,2,3,4,0,1,2,3,4,5,0,1,2,3,0,0,0]])
    Returns:
        list[list[int]]: 一个嵌套列表，包含每个批次项中各个序列的长度。
                         例如：[[5, 6, 4, 1, 1, 1]]
    """
    # 检查输入是否为 2D Tensor
    if position_ids.dim() != 2:
        raise ValueError(f"输入必须是 2D Tensor，但得到了 {position_ids.dim()}D")

    all_lengths = []
    
    # 我们按批次逐行处理。因为每行的序列长度数量不同（ragged），
    # 所以 Python 循环在批次维度上是最高效且最清晰的写法。
    # 循环内部的操作是完全向量化的。
    for pids_row in position_ids:
        # 获取当前行的总长度
        seq_len = pids_row.shape[0]
        
        # 1. 找到所有值为 0 的元素的索引
        # pids_row == 0 会返回一个布尔 Tensor: [True, False, ..., True, ...]
        # torch.nonzero 会返回这些 True 值的索引
        # .flatten() 将其从 (N, 1) 形状的 Tensor 变为 (N,) 形状
        zero_indices = torch.nonzero(pids_row == 0).flatten()
        
        # 2. 将序列的总长度作为一个额外的切分点添加到末尾
        # 这对于计算最后一个序列的长度至关重要
        # 注意：要确保新创建的 tensor 和原始 tensor 在同一个设备上 (cpu/cuda)
        split_points = torch.cat([
            zero_indices,
            torch.tensor([seq_len], device=pids_row.device, dtype=zero_indices.dtype)
        ])
        
        # 3. 计算相邻切分点之间的差值，这就是我们想要的长度
        # torch.diff([a, b, c, d]) 会返回 [b-a, c-b, d-c]
        lengths = torch.diff(split_points)

        all_lengths.append(lengths)

    return all_lengths


def forward_add_noise_packed(
    inputs_ids: torch.Tensor,
    num_tokens_list: List[torch.Tensor],
    prompt_mask: torch.Tensor,
    mask_id: int,
    eps: float = 1e-3,
    max_tries: int = 10,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    为一批打包（packed）序列的 token ID 添加噪声。

    此函数保留了为每个逻辑样本（在每个批次项内拼接）生成独立随机噪声率的逻辑。
    它会随机将一部分 token 的 ID 替换为 mask_id。
    这个过程会避开被 prompt_mask 标记的位置。

    Args:
        inputs_ids (torch.Tensor): 
            输入的 token ID 张量，形状为 (bsz, total_tokens)。
        num_tokens_list (List[torch.Tensor]): 
            一个张量列表，长度为 bsz。列表中的每个张量记录了对应批次项中
            每个逻辑样本的长度。例如: [tensor([len1, len2]), tensor([len3, len4, len5])].
        prompt_mask (torch.Tensor): 
            布尔型张量，形状为 (bsz, total_tokens)，值为 True 的位置表示是 prompt，
            不应添加噪声。
        mask_id (int): 
            用于替换的 mask token 的 ID。
        eps (float): 
            微小值，用于防止噪声率 t 恰好为 0，确保 p_mask > 0。
        max_tries (int): 
            为确保至少一个非 prompt token 被 mask，对每个批次项尝试的最大次数。

    Returns:
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        - noisy_input_ids (torch.Tensor): 
            添加噪声后的 token ID 张量，形状为 (bsz, total_tokens)。
        - final_masked_indices (torch.Tensor): 
            布尔型张量，标记了哪些位置被实际 mask 了，形状为 (bsz, total_tokens)。
        - p_masks (torch.Tensor): 
            一个一维张量，包含了被 mask 的 token 对应的实际噪声率。
    """
    # 1. 验证和获取形状
    bsz, total_tokens = inputs_ids.shape
    device = inputs_ids.device

    # 检查输入的一致性
    assert len(num_tokens_list) == bsz, f"num_tokens_list 的长度 ({len(num_tokens_list)}) 必须等于 bsz ({bsz})"
    assert prompt_mask.shape == (bsz, total_tokens), f"prompt_mask 形状不匹配, 期望 {(bsz, total_tokens)}, 得到 {prompt_mask.shape}"

    # 准备结果容器
    noisy_ids_list = []
    final_masked_indices_list = []
    p_masks_per_token_list = []

    # 2. 在批次维度上迭代
    # 这是处理不同打包结构最直接有效的方法
    for i in range(bsz):
        # 提取当前批次项的数据
        current_ids = inputs_ids[i:i+1] # shape: (1, total_tokens)
        current_num_tokens = num_tokens_list[i]
        current_prompt_mask = prompt_mask[i:i+1] # shape: (1, total_tokens)
        
        num_samples_in_item = len(current_num_tokens)
        # 验证当前批次项的 token 总数是否匹配
        assert total_tokens == torch.sum(current_num_tokens), \
            f"批次项 {i} 的 num_tokens 之和 ({torch.sum(current_num_tokens)}) 与 total_tokens ({total_tokens}) 不匹配"

        eligible_for_masking = ~current_prompt_mask

        # 如果没有任何 token 可以被 mask，直接使用原始输入，并设置 p_mask 为 eps
        if not eligible_for_masking.any():
            noisy_ids_list.append(current_ids)
            final_masked_indices_list.append(torch.zeros_like(current_prompt_mask, dtype=torch.bool))
            # p_mask_per_token 的形状应为 (1, total_tokens) 以便后续拼接
            p_masks_per_token_list.append(torch.full((1, total_tokens), eps, device=device, dtype=torch.float))
            continue

        # --- 尝试生成 mask，确保至少 mask 一个 token ---
        final_masked_indices_item = torch.zeros_like(current_prompt_mask, dtype=torch.bool)
        p_mask_per_token = None
        
        for _ in range(max_tries):
            # 为每个逻辑样本生成一个独立的噪声率 t
            t = torch.rand(num_samples_in_item, device=device)
            p_mask_per_sample = (1 - eps) * t + eps

            # 将每个样本的噪声率扩展到其所有 token 上
            p_mask_per_token_1d = torch.repeat_interleave(p_mask_per_sample, current_num_tokens)
            p_mask_per_token = p_mask_per_token_1d.unsqueeze(0) # shape: (1, total_tokens)

            # 根据噪声率生成随机 mask
            masked_indices = torch.rand_like(p_mask_per_token) < p_mask_per_token
            # 应用 prompt mask，确保 prompt 不被 mask
            final_masked_indices_item = masked_indices & eligible_for_masking

            # 如果成功 mask 了至少一个 token，则跳出尝试循环
            if final_masked_indices_item.any():
                break
        
        # 如果 max_tries 之后仍然没有 mask 任何 token (极小概率)，就强制 mask 一个可 mask 的 token
        if not final_masked_indices_item.any():
            eligible_indices = torch.nonzero(eligible_for_masking.squeeze(0), as_tuple=True)[0]
            if len(eligible_indices) > 0:
                # 随机选择一个可 mask 的位置
                random_choice = torch.randint(0, len(eligible_indices), (1,)).item()
                force_mask_idx = eligible_indices[random_choice]
                final_masked_indices_item[0, force_mask_idx] = True


        # --- 根据最终的 mask 生成带噪声的 IDs ---
        noisy_ids_item = torch.where(
            final_masked_indices_item,
            mask_id,
            current_ids
        )
        
        # 保存这个批次项的结果
        noisy_ids_list.append(noisy_ids_item)
        final_masked_indices_list.append(final_masked_indices_item)
        p_masks_per_token_list.append(p_mask_per_token)

    # 3. 将列表中的结果堆叠成最终的批处理张量
    noisy_input_ids = torch.cat(noisy_ids_list, dim=0)
    final_masked_indices = torch.cat(final_masked_indices_list, dim=0)
    p_mask_full = torch.cat(p_masks_per_token_list, dim=0)
    
    # 4. 提取被 mask 位置对应的噪声率
    p_masks = p_mask_full[final_masked_indices]
    
    return noisy_input_ids, final_masked_indices, p_masks


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 block_attn_mask(num_tokens, block_size, device):
    masks = []
    for i in range(len(num_tokens)):
        cur_masks = []
        for num in num_tokens[i]:
            # 全部返回 n*n 而非 2n*2n
            single_mask = block_diff_mask(
                b=None,
                h=None,
                q_idx=torch.arange(num * 2, device=device)[:, None],
                kv_idx=torch.arange(num * 2, device=device)[None, :],
                block_size=block_size,
                n=num,
            )
            cur_masks.append(single_mask)
        masks.append(torch.block_diag(*cur_masks))
    masks = torch.stack(masks, dim=0)
    return masks


# @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs")  # Commented out to prevent Dynamo compile errors with Tensor masks
def fused_flex_attention(query, key, value, attention_mask, **kwargs):
    return flex_attention(query, key, value, block_mask=attention_mask, **kwargs)


@use_kernel_forward_from_hub("RMSNorm")
class SDARRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        SDARRMSNorm 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)
        hidden_states = hidden_states.to(input_dtype)
        if flash_rms_norm is not None:
            return flash_rms_norm(
                hidden_states, weight=self.weight, bias=None, eps=self.variance_epsilon
            )
        return self.weight * hidden_states

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


class SDARMLP(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):
        if liger_kernel_is_available:
            return self.down_proj(LigerSiLUMulFunction.apply(self.gate_proj(x), self.up_proj(x)))
        else:
            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 SDARAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: SDARConfig, 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.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads

        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        self.k_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.v_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
        )
        # unlike olmo, only on the head dim!
        self.q_norm = SDARRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        # thus post q_norm does not need reshape
        self.k_norm = SDARRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.sliding_window = config.sliding_window
        if not (
            self.config.use_sliding_window
            and getattr(self.config, "sliding_window", None) is not None
            and self.layer_idx >= self.config.max_window_layers
        ):
            self.sliding_window = 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]
        bsz, q_len = input_shape
        hidden_shape = (*input_shape, -1, self.head_dim)

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

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(
            query_states, key_states, cos, sin)
        query_states = query_states.to(value_states.dtype)
        key_states = key_states.to(value_states.dtype)

        if past_key_value is not None and kwargs.get("store_kv", False):
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            key_states, value_states = past_key_value.update(
                key_states, value_states, self.layer_idx)
        elif past_key_value is not None and not kwargs.get("store_kv", False) and len(past_key_value) > self.layer_idx:
            # only retrive, do not store kv
            past_key_states, past_value_states = past_key_value[self.layer_idx]
            key_states = torch.cat(
                [past_key_states, key_states], dim=-2)
            value_states = torch.cat(
                [past_value_states, value_states], dim=-2)

        if self.training:
            if isinstance(attention_mask, torch.Tensor) or attention_mask is None:
                attn_output = torch.nn.functional.scaled_dot_product_attention(
                    query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=0.0, is_causal=(attention_mask is None)
                )
                attn_weights = None
            else:
                attn_output, attn_weights = fused_flex_attention(
                        query=query_states,
                        key=key_states,
                        value=value_states,
                        attention_mask=attention_mask,
                        enable_gqa=True,
                        scale=self.scaling,
                        return_lse=True
                    )
                attn_weights = attn_weights.to(
                    value_states.dtype) if attn_weights is not None else None
                attn_output = rearrange(attn_output, 'b h l d -> b l (h d)')
        else:
            attention_mask = attention_mask.bool() if attention_mask is not None else None
            attn_weights = None
            if torch.all(attention_mask):  # decoding
                query_states = query_states.transpose(1, 2)
                key_states = key_states.transpose(1, 2)
                value_states = value_states.transpose(1, 2)
                attn_output = flash_attn_func(
                    query_states,
                    key_states,
                    value_states,
                    causal=False,
                    softmax_scale=self.scaling
                )
                attn_output = rearrange(attn_output, 'b l h d -> b l (h d)')
            else:  # prefilling
                attn_output = F.scaled_dot_product_attention(
                    query=query_states,
                    key=key_states,
                    value=value_states,
                    attn_mask=attention_mask,
                    is_causal=False,
                    scale=self.scaling,
                    enable_gqa=True
                )
                attn_output = rearrange(attn_output, 'b h l d -> b l (h d)')
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights  # , attn_weights


class SDARDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: SDARConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = SDARAttention(config=config, layer_idx=layer_idx)
        self.mlp = SDARMLP(config)
        self.input_layernorm = SDARRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = SDARRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps)
        if (
            config.sliding_window and config._attn_implementation != "flash_attention_2"
        ):  # diff with Llama is this warning
            logger.warning_once(
                f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
                "unexpected results may be encountered."
            )

    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,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
        store_kv: Optional[bool] = False,
        cache_position: Optional[torch.LongTensor] = None,
        # necessary, but kept here for BC
        position_embeddings: Optional[Tuple[torch.Tensor,
                                            torch.Tensor]] = None,
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        hidden_states, self_attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            store_kv=store_kv,
            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

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (self_attn_weights,)

        return outputs


@auto_docstring
class SDARPreTrainedModel(PreTrainedModel):
    config_class = SDARConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["SDARDecoderLayer"]
    _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

    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, SDARRMSNorm):
            module.weight.data.fill_(1.0)


class SDARRotaryEmbedding(nn.Module):
    def __init__(self, config: SDARConfig, device=None):
        super().__init__()
        # BC: "rope_type" was originally "type"
        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
            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()
    # power user: used with advanced RoPE types (e.g. dynamic rope)
    @dynamic_rope_update
    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 SDARModel(SDARPreTrainedModel):
    def __init__(self, config: SDARConfig):
        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(
            [SDARDecoderLayer(config, layer_idx)
             for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = SDARRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = SDARRotaryEmbedding(config=config)
        self.gradient_checkpointing = False

        # 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

    @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,
        use_cache: Optional[bool] = None,
        store_kv: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **flash_attn_kwargs: Unpack[FlashAttentionKwargs],
    ) -> BaseModelOutputWithPast:
        r"""
        store_kv (`bool`, *optional*):
            Whether to keep KV states in the custom SDAR cache path during generation/inference.
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError(
                "You must specify exactly one of input_ids or inputs_embeds")

        if self.gradient_checkpointing and self.training and use_cache:
            logger.warning_once(
                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
            )
            use_cache = False

        # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
        if not isinstance(past_key_values, (type(None), Cache)):
            raise ValueError(
                "The `past_key_values` should be either a `Cache` object or `None`.")

        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)

        attention_mask = self._update_causal_mask(
            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
        )

        hidden_states = inputs_embeds

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

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None

        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                store_kv=store_kv,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **flash_attn_kwargs,
            )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

    def _update_causal_mask(
        self,
        attention_mask: Union[torch.Tensor, "BlockMask"],
        input_tensor: torch.Tensor,
        cache_position: torch.Tensor,
        past_key_values: Cache,
        output_attentions: bool = False,
    ):
        # Training can pass a precomputed flex-attention BlockMask even when the
        # Transformers-selected backend is not "flex_attention". In that case the
        # mask should bypass tensor-only causal-mask preparation entirely.
        if attention_mask is not None and not isinstance(attention_mask, torch.Tensor):
            assert isinstance(attention_mask, BlockMask)
            return attention_mask

        if self.config._attn_implementation == "flash_attention_2":
            if attention_mask is not None and past_key_values is not None:
                is_padding_right = attention_mask[:, -
                                                  1].sum().item() != input_tensor.size()[0]
                if is_padding_right:
                    raise ValueError(
                        "You are attempting to perform batched generation with padding_side='right'"
                        " this may lead to unexpected behaviour for Flash Attention version of Qwen3. Make sure to "
                        " call `tokenizer.padding_side  = 'left'` before tokenizing the input. "
                    )
            if attention_mask is not None and 0.0 in attention_mask:
                return attention_mask
            return None
        if self.config._attn_implementation == "flex_attention":
            if isinstance(attention_mask, torch.Tensor):
                seq_len_q, seq_len_kv = attention_mask.shape
                assert seq_len_q == seq_len_kv, f"got {attention_mask.shape=}"
                attention_mask = create_block_mask(
                    # 2d bool tensor, shape: [2*seqlen, 2*seqlen]
                    lambda b, h, q_idx, kv_idx: attention_mask[q_idx, kv_idx],
                    B=None, H=None, Q_LEN=seq_len_q, KV_LEN=seq_len_kv,
                )
            else:
                # Here we pass in flex mask computed externally
                assert isinstance(attention_mask, BlockMask)
            return attention_mask

        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
        # to infer the attention mask.
        past_seen_tokens = past_key_values.get_seq_length(
        ) if past_key_values is not None else 0
        using_static_cache = isinstance(past_key_values, StaticCache)
        using_sliding_window_cache = isinstance(
            past_key_values, SlidingWindowCache)

        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
        if (
            self.config._attn_implementation == "sdpa"
            and not (using_static_cache or using_sliding_window_cache)
            and not output_attentions
        ):
            if AttentionMaskConverter._ignore_causal_mask_sdpa(
                attention_mask,
                inputs_embeds=input_tensor,
                past_key_values_length=past_seen_tokens,
                sliding_window=self.config.sliding_window,
                is_training=self.training,
            ):
                return None

        dtype = input_tensor.dtype
        min_dtype = torch.finfo(dtype).min
        sequence_length = input_tensor.shape[1]
        # SlidingWindowCache or StaticCache
        if using_sliding_window_cache or using_static_cache:
            target_length = past_key_values.get_max_cache_shape()
        # DynamicCache or no cache
        else:
            target_length = (
                attention_mask.shape[-1]
                if isinstance(attention_mask, torch.Tensor)
                else past_seen_tokens + sequence_length + 1
            )

        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
            attention_mask,
            sequence_length=sequence_length,
            target_length=target_length,
            dtype=dtype,
            cache_position=cache_position,
            batch_size=input_tensor.shape[0],
            config=self.config,
            past_key_values=past_key_values,
        )

        if (
            self.config._attn_implementation == "sdpa"
            and attention_mask is not None
            and attention_mask.device.type in ["cuda", "xpu", "npu"]
            and not output_attentions
        ):
            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
            # Details: https://github.com/pytorch/pytorch/issues/110213
            causal_mask = AttentionMaskConverter._unmask_unattended(
                causal_mask, min_dtype)

        return causal_mask

    @staticmethod
    def _prepare_4d_causal_attention_mask_with_cache_position(
        attention_mask: torch.Tensor,
        sequence_length: int,
        target_length: int,
        dtype: torch.dtype,
        cache_position: torch.Tensor,
        batch_size: int,
        config: SDARConfig,
        past_key_values: Cache,
    ):
        """
        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

        Args:
            attention_mask (`torch.Tensor`):
                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
            sequence_length (`int`):
                The sequence length being processed.
            target_length (`int`):
                The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
            dtype (`torch.dtype`):
                The dtype to use for the 4D attention mask.
            cache_position (`torch.Tensor`):
                Indices depicting the position of the input sequence tokens in the sequence.
            batch_size (`torch.Tensor`):
                Batch size.
            config (`SDARConfig`):
                The model's configuration class
            past_key_values (`Cache`):
                The cache class that is being used currently to generate
        """
        if attention_mask is not None and attention_mask.dim() == 4:
            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
            causal_mask = attention_mask
        else:
            min_dtype = torch.finfo(dtype).min
            causal_mask = torch.full(
                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device
            )
            diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape(
                -1, 1
            )
            text_config = config.get_text_config()
            if getattr(text_config, "use_sliding_window", True) and text_config.sliding_window is not None:
                # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
                # the check is needed to verify is current checkpoint was trained with sliding window or not
                if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:
                    sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= (
                        cache_position.reshape(-1, 1) -
                        text_config.sliding_window
                    )
                    diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
            causal_mask *= diagonal_attend_mask
            causal_mask = causal_mask[None, None,
                                      :, :].expand(batch_size, 1, -1, -1)
            if attention_mask is not None:
                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
                if attention_mask.shape[-1] > target_length:
                    attention_mask = attention_mask[:, :target_length]
                mask_length = attention_mask.shape[-1]
                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
                    causal_mask.device
                )
                padding_mask = padding_mask == 0
                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
                    padding_mask, min_dtype
                )
        return causal_mask


class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs):
    ...


class GapRemaskHead(nn.Module):
    def __init__(self, hidden_size: int, confidence_feature_dim: int = 4):
        super().__init__()
        self.up_proj = nn.Linear(hidden_size, hidden_size, bias=True)
        self.conf_proj = nn.Linear(confidence_feature_dim, hidden_size, bias=True)
        self.down_proj = nn.Linear(hidden_size, 1, bias=True)
        self.base_scale = nn.Parameter(torch.tensor(1.0))
        self.residual_scale = nn.Parameter(torch.tensor(1.0))

    @staticmethod
    def _score_to_logit(score: torch.Tensor) -> torch.Tensor:
        score = score.clamp(min=1e-4, max=1.0 - 1e-4)
        return torch.logit(score)

    @staticmethod
    def build_confidence_features(lm_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        probs = lm_logits.softmax(dim=-1)
        top2 = torch.topk(probs, k=2, dim=-1).values
        top1_prob = top2[..., 0]
        top2_prob = top2[..., 1]
        low_conf = 1.0 - top1_prob
        margin = top1_prob - top2_prob
        entropy = -(probs * probs.clamp_min(1e-8).log()).sum(dim=-1)
        entropy = entropy / math.log(lm_logits.shape[-1])
        features = torch.stack((top1_prob, low_conf, margin, entropy), dim=-1)
        return features, low_conf

    def forward(self, hidden_states: torch.Tensor, lm_logits: Optional[torch.Tensor] = None) -> torch.Tensor:
        residual_hidden = self.up_proj(hidden_states)
        if lm_logits is not None:
            conf_features, low_conf = self.build_confidence_features(lm_logits.float())
            residual_hidden = residual_hidden + self.conf_proj(conf_features.to(hidden_states.dtype))
            base_logit = self._score_to_logit(low_conf)
        else:
            base_logit = hidden_states.new_zeros(hidden_states.shape[:-1])
        residual_logit = self.down_proj(F.gelu(residual_hidden)).squeeze(-1)
        return self.base_scale * base_logit + self.residual_scale * residual_logit

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        legacy_key_map = {
            f"{prefix}0.weight": f"{prefix}up_proj.weight",
            f"{prefix}0.bias": f"{prefix}up_proj.bias",
            f"{prefix}2.weight": f"{prefix}down_proj.weight",
            f"{prefix}2.bias": f"{prefix}down_proj.bias",
        }
        for legacy_key, new_key in legacy_key_map.items():
            if legacy_key in state_dict and new_key not in state_dict:
                state_dict[new_key] = state_dict.pop(legacy_key)
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )


@auto_docstring
class SDARForCausalLM(SDARPreTrainedModel, 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 = SDARModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(
            config.hidden_size, config.vocab_size, bias=False)
        self.gap_remask_head = GapRemaskHead(config.hidden_size)
        self.gap_value_head = nn.Linear(config.hidden_size, 1, bias=False)
        self._puma_streaming_state = None
        self._puma_streaming_context = {"slot_offset": 0, "buffer_size": None}
        self._gap_collect_branch_debug = False
        self._last_branch_debug = None

        # Initialize weights and apply final processing
        self.post_init()
        with torch.no_grad():
            self.gap_value_head.weight.zero_()

    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

    def _get_gap_reward_tokenizer(self):
        tokenizer = getattr(self, "_gap_reward_tokenizer", None)
        if tokenizer is None:
            tokenizer = getattr(self, "_grpo_tokenizer", None)
        return tokenizer

    def _get_gap_reference_model(self):
        kl_coef = float(getattr(self.config, "gap_grpo_kl_coef", 0.0) or 0.0)
        ref_model_path = getattr(self.config, "gap_grpo_ref_model_path", None)
        if kl_coef <= 0.0:
            return None
        if not ref_model_path:
            raise ValueError("gap_grpo_kl_coef is positive, but gap_grpo_ref_model_path is not set.")

        ref_model = self.__dict__.get("_gap_reference_model", None)
        if ref_model is None:
            dtype = self.lm_head.weight.dtype
            ref_model = self.__class__.from_pretrained(
                ref_model_path,
                torch_dtype=dtype,
                low_cpu_mem_usage=True,
            )
            ref_model.eval()
            ref_model.requires_grad_(False)
            self.__dict__["_gap_reference_model"] = ref_model

        device = next(self.parameters()).device
        ref_device = next(ref_model.parameters()).device
        if ref_device != device:
            ref_model.to(device)
        ref_model.eval()
        return ref_model

    def _decode_gap_response_tokens(self, token_ids: torch.Tensor) -> str:
        tokenizer = self._get_gap_reward_tokenizer()
        if tokenizer is None:
            return ""
        if torch.is_tensor(token_ids):
            token_ids = token_ids.detach().to("cpu").tolist()
        token_ids = [int(token_id) for token_id in token_ids if int(token_id) >= 0]
        try:
            return tokenizer.decode(token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False).strip()
        except TypeError:
            tokens = tokenizer.convert_ids_to_tokens(token_ids, skip_special_tokens=True)
            tokens = [token for token in tokens if isinstance(token, str)]
            if not tokens:
                return ""
            return tokenizer.convert_tokens_to_string(tokens).strip()

    def _get_gap_eval_stop_words(self) -> list[str]:
        tokenizer = self._get_gap_reward_tokenizer()
        stop_words: list[str] = []
        raw_stop_words = _os.getenv("SDAR_GAP_GRPO_STOP_WORDS", "")
        if raw_stop_words.strip():
            stop_words.extend([item for item in raw_stop_words.split("|||") if item])
        eos_token_id = getattr(self.config, "eos_token_id", None)
        if tokenizer is not None and eos_token_id is not None:
            eos_ids = [eos_token_id] if isinstance(eos_token_id, int) else list(eos_token_id)
            for token_id in eos_ids:
                try:
                    stop_words.append(tokenizer.decode(int(token_id)))
                except Exception:
                    pass
        if tokenizer is not None and getattr(tokenizer, "eos_token", None):
            stop_words.append(tokenizer.eos_token)
        seen = set()
        return [word for word in stop_words if word and not (word in seen or seen.add(word))]

    def _get_gap_eval_stop_sequences(self) -> list[list[int]]:
        cached = getattr(self, "_gap_eval_stop_sequences", None)
        if cached is not None:
            return cached
        tokenizer = self._get_gap_reward_tokenizer()
        sequences: list[list[int]] = []
        if tokenizer is not None:
            for stop in self._get_gap_eval_stop_words():
                try:
                    sequence = tokenizer.encode(stop, add_special_tokens=False)
                except Exception:
                    sequence = []
                if sequence:
                    sequences.append([int(token_id) for token_id in sequence])
        self._gap_eval_stop_sequences = sequences
        return sequences

    @staticmethod
    def _match_gap_eval_stop_sequences(token_ids: torch.LongTensor, stop_sequences: list[list[int]]) -> torch.BoolTensor:
        if token_ids.dim() == 1:
            token_ids = token_ids.unsqueeze(0)
        matches = torch.zeros(token_ids.shape[0], dtype=torch.bool, device=token_ids.device)
        if token_ids.numel() == 0 or not stop_sequences:
            return matches
        for sequence in stop_sequences:
            seq_len = len(sequence)
            if seq_len <= 0:
                continue
            if seq_len == 1:
                matches |= token_ids.eq(int(sequence[0])).any(dim=1)
                continue
            if token_ids.shape[1] < seq_len:
                continue
            sequence_tensor = token_ids.new_tensor(sequence)
            matches |= token_ids.unfold(1, seq_len, 1).eq(sequence_tensor).all(dim=-1).any(dim=1)
        return matches

    def _strip_gap_eval_stop_text(self, text: str) -> str:
        text = text or ""
        for stop in self._get_gap_eval_stop_words():
            text = text.split(stop)[0]
        return text

    @staticmethod
    def _truncate_gap_debug_text(text: str, max_chars: int) -> str:
        text = (text or "").replace("\n", "\\n").strip()
        if max_chars > 0 and len(text) > max_chars:
            return text[: max_chars - 3] + "..."
        return text

    def _apply_gap_grpo_remask_guards(
        self,
        candidate_mask: torch.BoolTensor,
        target_scope_mask: torch.BoolTensor,
        masked_indices: Optional[torch.BoolTensor] = None,
    ) -> torch.BoolTensor:
        prefix_guard_tokens = _gap_env_int("SDAR_GAP_GRPO_REMASK_PREFIX_GUARD_TOKENS", 0)
        tail_guard_blocks = _gap_env_int("SDAR_GAP_GRPO_REMASK_TAIL_GUARD_BLOCKS", 0)
        exclude_frontier_blocks = _gap_env_int("SDAR_GAP_GRPO_EXCLUDE_FRONTIER_BLOCKS", 0)
        middle_window_blocks = _gap_env_int("SDAR_GAP_GRPO_MIDDLE_WINDOW_BLOCKS", 0)
        if prefix_guard_tokens <= 0 and tail_guard_blocks <= 0 and exclude_frontier_blocks <= 0 and middle_window_blocks <= 0:
            return candidate_mask

        guarded_mask = candidate_mask.clone()
        block_size = max(1, int(getattr(self.config, "block_size", 1)))
        for row_idx in range(guarded_mask.shape[0]):
            target_positions = torch.nonzero(target_scope_mask[row_idx], as_tuple=False).flatten()
            if target_positions.numel() == 0:
                guarded_mask[row_idx] = False
                continue

            allowed_start = int(target_positions[0].item()) + max(0, prefix_guard_tokens)
            allowed_end = int(target_positions[-1].item()) + 1 - max(0, tail_guard_blocks) * block_size
            if masked_indices is not None:
                active_target_mask = masked_indices[row_idx] & target_scope_mask[row_idx]
                active_positions = torch.nonzero(active_target_mask, as_tuple=False).flatten()
                if active_positions.numel() == 0:
                    guarded_mask[row_idx] = False
                    continue
                frontier_block_start = (int(active_positions[0].item()) // block_size) * block_size
                rolling_end = frontier_block_start + block_size - max(0, exclude_frontier_blocks) * block_size
                allowed_end = min(allowed_end, rolling_end)
                if middle_window_blocks > 0:
                    allowed_start = max(allowed_start, allowed_end - middle_window_blocks * block_size)
            elif middle_window_blocks > 0:
                allowed_end = min(allowed_end, allowed_start + middle_window_blocks * block_size)
            if allowed_end <= allowed_start:
                guarded_mask[row_idx] = False
                continue

            positions = torch.arange(guarded_mask.shape[1], device=guarded_mask.device)
            middle_mask = positions.ge(allowed_start) & positions.lt(allowed_end)
            guarded_mask[row_idx] &= middle_mask
        return guarded_mask

    def _maybe_capture_gap_branch_debug(
        self,
        clean_input_ids: torch.LongTensor,
        target_scope_mask: torch.BoolTensor,
        shared_state_input_ids: torch.LongTensor,
        baseline_terminal: Optional[torch.LongTensor],
        baseline_terminal_reward: Optional[torch.FloatTensor],
        baseline_reward: Optional[torch.FloatTensor],
        sampled_terminal: torch.LongTensor,
        target_scope_flat: torch.BoolTensor,
        sampled_terminal_reward: torch.FloatTensor,
        rewards: torch.FloatTensor,
        reward_gain: torch.FloatTensor,
        remask_rate: torch.FloatTensor,
        baseline_remaining_mask_rate: Optional[torch.FloatTensor],
        sampled_remaining_mask_rate: torch.FloatTensor,
        sampled_full: torch.BoolTensor,
        full_candidate_mask: torch.BoolTensor,
        baseline_debug_mask: Optional[torch.BoolTensor] = None,
        sampled_debug_mask: Optional[torch.BoolTensor] = None,
        baseline_stop_hit: Optional[torch.BoolTensor] = None,
        sampled_stop_hit: Optional[torch.BoolTensor] = None,
        baseline_length_cap_hit: Optional[torch.BoolTensor] = None,
        sampled_length_cap_hit: Optional[torch.BoolTensor] = None,
    ) -> None:
        if not getattr(self, "_gap_collect_branch_debug", False):
            return

        batch_size = clean_input_ids.shape[0]
        num_samples = rewards.shape[0]
        max_examples = int(getattr(self.config, "gap_grpo_branch_debug_max_examples", 1))
        max_branches = int(getattr(self.config, "gap_grpo_branch_debug_max_branches", 3))
        max_chars = int(getattr(self.config, "gap_grpo_branch_debug_max_chars", 160))
        example_count = max(1, min(batch_size, max_examples))
        branch_count = num_samples if max_branches <= 0 else max(1, min(num_samples, max_branches))
        has_baseline = (
            baseline_terminal is not None
            and baseline_terminal_reward is not None
            and baseline_reward is not None
            and baseline_remaining_mask_rate is not None
        )
        baseline_text_mask = baseline_debug_mask if baseline_debug_mask is not None else target_scope_mask
        sampled_text_mask = sampled_debug_mask if sampled_debug_mask is not None else target_scope_flat

        snapshot = {
            "candidate_count_mean": float(full_candidate_mask.to(torch.float32).sum(dim=-1).mean().item()) if full_candidate_mask.numel() > 0 else 0.0,
            "baseline_remaining_mask_rate": float(baseline_remaining_mask_rate.mean().item()) if baseline_remaining_mask_rate is not None else None,
            "sampled_remaining_mask_rate": float(sampled_remaining_mask_rate.mean().item()),
            "sampled_remaining_mask_rate_max": float(sampled_remaining_mask_rate.max().item()) if sampled_remaining_mask_rate.numel() > 0 else 0.0,
            "has_baseline": has_baseline,
            "examples": [],
        }

        for row_idx in range(example_count):
            answer_positions = torch.nonzero(target_scope_mask[row_idx], as_tuple=False).flatten()
            prompt_length = int(answer_positions[0].item()) if answer_positions.numel() > 0 else clean_input_ids.shape[1]
            visible_answer_mask = target_scope_mask[row_idx] & shared_state_input_ids[row_idx].ne(self.config.mask_token_id)
            visible_answer_positions = torch.nonzero(visible_answer_mask, as_tuple=False).flatten()
            if visible_answer_positions.numel() > 0:
                shared_prefix_end = int(visible_answer_positions[-1].item()) + 1
                shared_prefix_text = self._truncate_gap_debug_text(
                    self._decode_gap_response_tokens(shared_state_input_ids[row_idx, prompt_length:shared_prefix_end]),
                    max_chars,
                )
            else:
                shared_prefix_end = prompt_length
                shared_prefix_text = ""
            block_size = max(1, int(getattr(self.config, "block_size", 1)))
            shared_visible_answer_tokens = int(visible_answer_positions.numel())
            shared_visible_full_blocks = shared_visible_answer_tokens // block_size
            frontier_answer_block = shared_visible_answer_tokens // block_size
            shared_answer_tokens = int(answer_positions.numel())
            gold_text = self._truncate_gap_debug_text(
                self._decode_gap_response_tokens(clean_input_ids[row_idx][target_scope_mask[row_idx]]),
                max_chars,
            )
            if has_baseline:
                assert baseline_terminal is not None
                baseline_text = self._truncate_gap_debug_text(
                    self._decode_gap_response_tokens(baseline_terminal[row_idx][baseline_text_mask[row_idx]]),
                    max_chars,
                )
            else:
                baseline_text = ""
            branch_entries = []
            branch_texts = []
            for branch_idx in range(branch_count):
                flat_idx = branch_idx * batch_size + row_idx
                branch_text = self._truncate_gap_debug_text(
                    self._decode_gap_response_tokens(sampled_terminal[flat_idx][sampled_text_mask[flat_idx]]),
                    max_chars,
                )
                branch_stop = bool(sampled_stop_hit[branch_idx, row_idx].item()) if sampled_stop_hit is not None else None
                branch_cap = bool(sampled_length_cap_hit[branch_idx, row_idx].item()) if sampled_length_cap_hit is not None else None
                branch_texts.append(branch_text)
                branch_entries.append(
                    {
                        "branch_idx": int(branch_idx),
                        "terminal_reward": float(sampled_terminal_reward[branch_idx, row_idx].item()),
                        "reward": float(rewards[branch_idx, row_idx].item()),
                        "gain": float(reward_gain[branch_idx, row_idx].item()),
                        "remask_rate": float(remask_rate[branch_idx, row_idx].item()),
                        "remaining_mask_rate": float(sampled_remaining_mask_rate[branch_idx, row_idx].item()),
                        "remask_tokens": int(sampled_full[branch_idx, row_idx].sum().item()),
                        "remask_block_span": "",
                        "remask_answer_offset_span": "",
                        "stop_hit": branch_stop,
                        "length_cap_hit": branch_cap,
                        "text": branch_text,
                    }
                )
                branch_answer_positions = torch.nonzero(sampled_text_mask[flat_idx], as_tuple=False).flatten()
                remask_positions = torch.nonzero(
                    sampled_full[branch_idx, row_idx] & sampled_text_mask[flat_idx],
                    as_tuple=False,
                ).flatten()
                if remask_positions.numel() > 0 and branch_answer_positions.numel() > 0:
                    answer_offsets = torch.searchsorted(branch_answer_positions, remask_positions)
                    block_ids = torch.div(answer_offsets, block_size, rounding_mode="floor")
                    branch_entries[-1]["remask_block_span"] = f"{int(block_ids.min().item())}-{int(block_ids.max().item())}"
                    branch_entries[-1]["remask_answer_offset_span"] = f"{int(answer_offsets.min().item())}-{int(answer_offsets.max().item())}"

            base_stop = bool(baseline_stop_hit[row_idx].item()) if (has_baseline and baseline_stop_hit is not None) else None
            base_cap = bool(baseline_length_cap_hit[row_idx].item()) if (has_baseline and baseline_length_cap_hit is not None) else None
            snapshot["examples"].append(
                {
                    "example_idx": int(row_idx),
                    "prompt_length": prompt_length,
                    "answer_token_count": shared_answer_tokens,
                    "shared_visible_answer_tokens": shared_visible_answer_tokens,
                    "shared_visible_full_blocks": shared_visible_full_blocks,
                    "frontier_answer_block": frontier_answer_block,
                    "shared_prefix_text": shared_prefix_text,
                    "gold": gold_text,
                    "baseline_text": baseline_text,
                    "baseline_terminal_reward": float(baseline_terminal_reward[row_idx].item()) if has_baseline else None,
                    "baseline_reward": float(baseline_reward[row_idx].item()) if has_baseline else None,
                    "baseline_remaining_mask_rate": float(baseline_remaining_mask_rate[row_idx].item()) if has_baseline else None,
                    "baseline_stop_hit": base_stop,
                    "baseline_length_cap_hit": base_cap,
                    "unique_branch_texts": int(len(set(branch_texts))),
                    "branches": branch_entries,
                }
            )

        self._last_branch_debug = snapshot
        self._gap_collect_branch_debug = False

    @staticmethod
    def _extract_gap_boxed_answer(text: str) -> str:
        marker = "\\boxed"
        start = text.rfind(marker)
        if start == -1:
            return ""
        brace_start = text.find("{", start)
        if brace_start == -1:
            return ""
        depth = 0
        chars = []
        for ch in text[brace_start + 1:]:
            if ch == "{":
                depth += 1
                chars.append(ch)
            elif ch == "}":
                if depth == 0:
                    return "".join(chars).strip()
                depth -= 1
                chars.append(ch)
            else:
                chars.append(ch)
        return ""

    @staticmethod
    def _normalize_gap_answer_text(text: str) -> str:
        text = (text or "").strip()
        boxed = SDARForCausalLM._extract_gap_boxed_answer(text)
        if boxed:
            text = boxed
        text = text.strip().lower()
        if text.endswith("."):
            text = text[:-1]
        text = text.replace("\\left", "").replace("\\right", "")
        text = text.replace("$", "").replace(" ", "")
        text = text.replace("\\,", "").replace(",", "")
        return text

    @staticmethod
    def _normalize_gap_latex_answer_text(text: str) -> str:
        text = (text or "").strip().lower()
        text = text.replace("\\left", "").replace("\\right", "")
        text = text.replace("\\dfrac", "\\frac").replace("\\tfrac", "\\frac")
        text = text.replace("\\cdot", "*").replace("\\times", "*")
        text = text.replace("\\,", "").replace("\\!", "").replace("\\;", "").replace("\\:", "")
        text = text.replace("$", "").replace(",", "")
        text = re.sub(r"\\(?:mathrm|text)\{([^{}]*)\}", r"\1", text)
        text = re.sub(r"\s+", "", text)
        if text.startswith("{") and text.endswith("}"):
            text = text[1:-1]
        return text

    @staticmethod
    def _gap_latex_answer_has_symbolic_token(text: str) -> bool:
        text = SDARForCausalLM._normalize_gap_latex_answer_text(text)
        symbolic_patterns = (
            "\\pi",
            "\\sqrt",
            "\\sin",
            "\\cos",
            "\\tan",
            "\\log",
            "\\ln",
            "\\theta",
            "\\alpha",
            "\\beta",
            "\\gamma",
            "\\infty",
        )
        if any(pattern in text for pattern in symbolic_patterns):
            return True
        return bool(re.search(r"[a-z]", text.replace("\\frac", "")))

    @staticmethod
    def _gap_strict_boxed_answer_compatible(pred_text: str, gold_text: str) -> bool:
        pred_boxed = SDARForCausalLM._extract_gap_boxed_answer(pred_text)
        gold_boxed = SDARForCausalLM._extract_gap_boxed_answer(gold_text)
        if not pred_boxed or not gold_boxed:
            return False

        pred_norm = SDARForCausalLM._normalize_gap_latex_answer_text(pred_boxed)
        gold_norm = SDARForCausalLM._normalize_gap_latex_answer_text(gold_boxed)
        if not pred_norm or not gold_norm:
            return False
        if pred_norm == gold_norm:
            return True

        # OpenCompass' postprocessor can reduce symbolic boxed answers such as
        # \pi/12 to the trailing number 12. Do not let that fallback mark a
        # symbolic answer correct unless both sides retain the same symbols.
        pred_symbolic = SDARForCausalLM._gap_latex_answer_has_symbolic_token(pred_boxed)
        gold_symbolic = SDARForCausalLM._gap_latex_answer_has_symbolic_token(gold_boxed)
        if pred_symbolic != gold_symbolic:
            return False

        if "\\pi" in pred_norm or "\\pi" in gold_norm:
            return "\\pi" in pred_norm and "\\pi" in gold_norm
        return True

    def _get_gap_opencompass_math_tools(self):
        cached = getattr(self, "_gap_opencompass_math_tools", None)
        if cached is not None:
            return cached
        repo_root = _os.getenv("SDAR_REPO_ROOT", "/work/leotsia0416/projects/SDAR")
        oc_root = _os.getenv("SDAR_OPENCOMPASS_ROOT", _os.path.join(repo_root, "evaluation", "opencompass"))
        if oc_root and oc_root not in sys.path:
            sys.path.insert(0, oc_root)
        try:
            from opencompass.datasets.math import MATHEvaluator, math_postprocess_sdar

            evaluator = MATHEvaluator(version="v2")
            cached = (math_postprocess_sdar, evaluator)
        except Exception:
            cached = None
        self._gap_opencompass_math_tools = cached
        return cached

    def _score_gap_answer_with_opencompass(self, pred_text: str, gold_text: str) -> bool:
        tools = self._get_gap_opencompass_math_tools()
        pred_text = self._strip_gap_eval_stop_text(pred_text)
        gold_text = self._strip_gap_eval_stop_text(gold_text)
        if not self._gap_strict_boxed_answer_compatible(pred_text, gold_text):
            return False
        if tools is not None:
            postprocess, evaluator = tools
            try:
                pred = postprocess(pred_text)
                gold = postprocess(gold_text)
                return bool(pred and gold and evaluator.is_equiv(pred, gold))
            except Exception:
                pass
        gold_norm = self._normalize_gap_answer_text(gold_text)
        pred_norm = self._normalize_gap_answer_text(pred_text)
        return bool(gold_norm and gold_norm == pred_norm)

    @staticmethod
    def _compute_gap_answer_block_ids(answer_mask: torch.BoolTensor, block_size: int) -> tuple[torch.LongTensor, int]:
        seq_len = answer_mask.shape[0]
        block_ids = torch.full((seq_len,), -1, dtype=torch.long, device=answer_mask.device)
        answer_positions = torch.nonzero(answer_mask, as_tuple=False).flatten()
        if answer_positions.numel() == 0:
            return block_ids, 0
        answer_order = torch.arange(answer_positions.numel(), device=answer_mask.device, dtype=torch.long)
        block_ids[answer_positions] = torch.div(answer_order, block_size, rounding_mode="floor")
        total_blocks = int(block_ids[answer_positions[-1]].item()) + 1
        return block_ids, total_blocks

    @staticmethod
    def _compute_gap_prefix_progress_limits(
        labels: torch.LongTensor,
        num_tokens,
        block_size: int,
        rollout_steps: int,
    ) -> torch.LongTensor:
        target_mask = labels.ne(-100)
        limits = torch.ones(labels.shape[0], dtype=torch.long, device=labels.device)
        for batch_idx, packed_lengths in enumerate(num_tokens):
            cursor = 0
            max_blocks = 0
            for sample_len_tensor in packed_lengths:
                sample_len = int(sample_len_tensor.item())
                sample_end = cursor + sample_len
                answer_tokens = int(target_mask[batch_idx, cursor:sample_end].sum().item())
                if answer_tokens > 0:
                    max_blocks = max(max_blocks, int(math.ceil(answer_tokens / max(1, block_size))))
                cursor = sample_end
            limits[batch_idx] = max(1, max_blocks * max(1, rollout_steps))
        return limits

    def _build_gap_prefix_teacher_forced_state(
        self,
        clean_input_ids: torch.LongTensor,
        labels: torch.LongTensor,
        num_tokens,
        progress_units: torch.LongTensor,
    ) -> torch.LongTensor:
        rollout_steps = max(1, int(getattr(self.config, "gap_rollout_steps", self.config.block_size)))
        block_size = int(self.config.block_size)
        transfer_schedule = get_num_transfer_tokens(block_size, rollout_steps).to(clean_input_ids.device)
        cumulative_transfers = torch.cat(
            (
                torch.zeros(1, dtype=torch.long, device=clean_input_ids.device),
                transfer_schedule.cumsum(dim=0),
            ),
            dim=0,
        )
        target_mask = labels.ne(-100)
        noisy_input_ids = torch.where(
            target_mask,
            torch.full_like(clean_input_ids, self.config.mask_token_id),
            clean_input_ids,
        )

        for batch_idx, packed_lengths in enumerate(num_tokens):
            progress = max(0, int(progress_units[batch_idx].item()))
            full_blocks = progress // rollout_steps
            frontier_stage = progress % rollout_steps
            frontier_visible_tokens = int(cumulative_transfers[min(frontier_stage, rollout_steps)].item())
            visible_answer_tokens = full_blocks * block_size + frontier_visible_tokens
            cursor = 0
            for sample_len_tensor in packed_lengths:
                sample_len = int(sample_len_tensor.item())
                sample_end = cursor + sample_len
                sample_target_mask = target_mask[batch_idx, cursor:sample_end]
                answer_positions = torch.nonzero(sample_target_mask, as_tuple=False).flatten()
                if answer_positions.numel() > 0 and visible_answer_tokens > 0:
                    reveal_count = min(int(answer_positions.numel()), visible_answer_tokens)
                    reveal_positions = answer_positions[:reveal_count]
                    local_noisy = noisy_input_ids[batch_idx, cursor:sample_end]
                    local_clean = clean_input_ids[batch_idx, cursor:sample_end]
                    local_noisy[reveal_positions] = local_clean[reveal_positions]
                cursor = sample_end

        return noisy_input_ids

    def _should_use_gap_prefix_frontier_state(self) -> bool:
        return (
            getattr(self.config, "gap_rollout_strategy", "low_confidence_dynamic") == "sequential"
            and getattr(self.config, "gap_rollout_scope", "all") == "frontier_block"
        )

    @staticmethod
    def _sample_from_logits(
        logits: torch.Tensor,
        temperature: float,
        top_k: int,
        top_p: float,
    ) -> torch.LongTensor:
        original_shape = logits.shape[:-1]
        vocab_size = logits.shape[-1]
        logits = logits.reshape(-1, vocab_size)

        if temperature <= 0.0:
            return logits.argmax(dim=-1).reshape(*original_shape)

        sample_logits = logits / max(temperature, 1e-5)
        candidate_indices = None
        if top_k > 0 and top_k < sample_logits.shape[-1]:
            sample_logits, candidate_indices = torch.topk(sample_logits, k=top_k, dim=-1)

        if 0.0 < top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(sample_logits, descending=True, dim=-1)
            sorted_probs = sorted_logits.softmax(dim=-1)
            cumulative_probs = sorted_probs.cumsum(dim=-1)
            sorted_remove = cumulative_probs > top_p
            sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
            sorted_remove[..., 0] = False
            remove_mask = torch.zeros_like(sorted_remove, dtype=torch.bool)
            remove_mask.scatter_(dim=-1, index=sorted_indices, src=sorted_remove)
            sample_logits = sample_logits.masked_fill(remove_mask, float("-inf"))

        probs = sample_logits.softmax(dim=-1)
        sampled_local = torch.multinomial(probs, num_samples=1).squeeze(-1)
        if candidate_indices is None:
            return sampled_local.reshape(*original_shape)

        sampled = candidate_indices.gather(dim=-1, index=sampled_local.unsqueeze(-1)).squeeze(-1)
        return sampled.reshape(*original_shape)

    @staticmethod
    def _sample_from_logits_with_eval_scores(
        logits: torch.Tensor,
        temperature: float,
        top_k: int,
        top_p: float,
    ) -> tuple[torch.LongTensor, torch.FloatTensor]:
        original_shape = logits.shape[:-1]
        vocab_size = logits.shape[-1]
        logits = logits.reshape(-1, vocab_size)

        if temperature <= 0.0:
            probs = logits.softmax(dim=-1)
            token = logits.argmax(dim=-1)
            token_prob = probs.gather(-1, token.unsqueeze(-1)).squeeze(-1)
            return token.reshape(*original_shape), token_prob.reshape(*original_shape)

        sample_logits = logits / max(temperature, 1e-5)
        original_probs = sample_logits.softmax(dim=-1)
        if top_k > 0 and top_k < sample_logits.shape[-1]:
            values, _ = torch.topk(sample_logits, k=top_k, dim=-1)
            min_values = values[..., -1, None]
            sample_logits = torch.where(
                sample_logits < min_values,
                torch.full_like(sample_logits, float("-inf")),
                sample_logits,
            )
        if 0.0 < top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(sample_logits, descending=True, dim=-1)
            cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
            sorted_remove = cumulative_probs > top_p
            sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
            sorted_remove[..., 0] = False
            remove_mask = torch.zeros_like(sorted_remove, dtype=torch.bool)
            remove_mask.scatter_(dim=-1, index=sorted_indices, src=sorted_remove)
            sample_logits = sample_logits.masked_fill(remove_mask, float("-inf"))

        probs = sample_logits.softmax(dim=-1)
        token = torch.multinomial(probs, num_samples=1).squeeze(-1)
        token_prob = original_probs.gather(-1, token.unsqueeze(-1)).squeeze(-1)
        return token.reshape(*original_shape), token_prob.reshape(*original_shape)

    @staticmethod
    def _resolve_gap_valid_length(num_tokens_for_row) -> int:
        total = 0
        for sample_len_tensor in num_tokens_for_row:
            total += int(sample_len_tensor.item())
        return total

    @staticmethod
    def _to_gap_block_mask(attention_mask):
        if attention_mask is None or not isinstance(attention_mask, torch.Tensor):
            return attention_mask

        if attention_mask.dim() == 4:
            base_mask = attention_mask[:, 0].to(dtype=torch.bool)
        elif attention_mask.dim() == 3:
            base_mask = attention_mask.to(dtype=torch.bool)
        elif attention_mask.dim() == 2:
            base_mask = attention_mask.to(dtype=torch.bool).unsqueeze(0)
        else:
            raise ValueError(f"Unsupported GAP rollout attention mask rank: {attention_mask.dim()}")

        return create_block_mask(
            lambda b, h, q_idx, kv_idx: base_mask[b, q_idx, kv_idx],
            B=base_mask.size(0),
            H=None,
            Q_LEN=base_mask.size(1),
            KV_LEN=base_mask.size(2),
        )

    @staticmethod
    def _select_gap_eval_window_transfer_tokens(
        masked_indices: torch.BoolTensor,
        proposal_scores_full: torch.FloatTensor,
        block_size: int,
        num_transfer_tokens: int,
        strategy: str,
        confidence_threshold: float,
    ) -> torch.BoolTensor:
        reveal_mask = torch.zeros_like(masked_indices)
        if num_transfer_tokens <= 0:
            return reveal_mask

        batch_size, seq_len = masked_indices.shape
        for batch_idx in range(batch_size):
            for block_start in range(0, seq_len, block_size):
                block_end = min(block_start + block_size, seq_len)
                block_mask = masked_indices[batch_idx, block_start:block_end]
                if not block_mask.any():
                    continue

                masked_local_indices = torch.nonzero(block_mask, as_tuple=False).flatten()
                block_scores = proposal_scores_full[batch_idx, block_start:block_end][masked_local_indices]
                chosen = select_teacher_forced_rollout_tokens.__globals__["_select_block_positions"](
                    block_scores=block_scores,
                    masked_local_indices=masked_local_indices,
                    num_transfer_tokens=num_transfer_tokens,
                    strategy=strategy,
                    confidence_threshold=confidence_threshold,
                )
                reveal_mask[batch_idx, block_start:block_end][chosen] = True
        return reveal_mask

    def _build_gap_eval_rollout_prefix_cache(
        self,
        x: torch.LongTensor,
        attention_mask: torch.Tensor,
        position_ids: torch.LongTensor,
        prefix_token_end: int,
    ) -> DynamicCache:
        cache = DynamicCache()
        if prefix_token_end <= 0:
            return cache

        cur_x = x[:, :prefix_token_end]
        cur_attn_mask = attention_mask[:, :, :prefix_token_end, :prefix_token_end]
        cur_position_ids = position_ids[:, :prefix_token_end]
        if self.model.training and not _gap_env_flag("SDAR_GAP_GRPO_STRICT_EVAL_DECODE", False):
            cur_attn_mask = self._to_gap_block_mask(cur_attn_mask)
        self.model(
            cur_x,
            attention_mask=cur_attn_mask,
            position_ids=cur_position_ids,
            past_key_values=cache,
            use_cache=True,
            store_kv=True,
        )
        return cache

    def _commit_gap_eval_rollout_block(
        self,
        cache: DynamicCache,
        x: torch.LongTensor,
        attention_mask: torch.Tensor,
        position_ids: torch.LongTensor,
        block_start: int,
        block_end: int,
    ) -> None:
        if block_end <= block_start:
            return
        cur_x = x[:, block_start:block_end]
        cur_attn_mask = attention_mask[:, :, block_start:block_end, :block_end]
        cur_position_ids = position_ids[:, block_start:block_end]
        if self.model.training and not _gap_env_flag("SDAR_GAP_GRPO_STRICT_EVAL_DECODE", False):
            cur_attn_mask = self._to_gap_block_mask(cur_attn_mask)
        self.model(
            cur_x,
            attention_mask=cur_attn_mask,
            position_ids=cur_position_ids,
            past_key_values=cache,
            use_cache=True,
            store_kv=True,
        )

    def _decode_gap_eval_window(
        self,
        window_inputs: torch.LongTensor,
        window_attention_mask: torch.Tensor,
        window_position_ids: torch.LongTensor,
        prefix_cache: DynamicCache,
        mask_id: int,
        block_length: int,
        denoising_steps: int,
        temperature: float,
        top_k: int,
        top_p: float,
        rollout_strategy: str,
        confidence_threshold: float,
    ) -> torch.LongTensor:
        rollout_steps = max(1, int(denoising_steps))
        transfer_schedule = get_num_transfer_tokens(block_length, rollout_steps).to(window_inputs.device)
        current_inputs = window_inputs.clone()
        current_stage = 0
        strict_eval_decode = _gap_env_flag("SDAR_GAP_GRPO_STRICT_EVAL_DECODE", False)
        block_attention_mask = self._to_gap_block_mask(window_attention_mask) if self.model.training and not strict_eval_decode else window_attention_mask

        while True:
            masked_indices = current_inputs.eq(mask_id)
            if not masked_indices.any():
                break

            outputs = self.model(
                current_inputs,
                attention_mask=block_attention_mask,
                position_ids=window_position_ids,
                past_key_values=prefix_cache,
                use_cache=True,
                store_kv=False,
                output_attentions=False,
                output_hidden_states=False,
                return_dict=True,
            )
            logits = self.lm_head(outputs.last_hidden_state).float()
            logits[..., mask_id] = float("-inf")
            proposal_ids, proposal_scores = self._sample_from_logits_with_eval_scores(
                logits=logits,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
            )
            proposal_ids = torch.where(masked_indices, proposal_ids, current_inputs)
            proposal_scores_full = torch.where(
                masked_indices,
                proposal_scores,
                torch.full_like(proposal_scores, float("-inf")),
            )

            if current_stage >= rollout_steps:
                current_inputs[masked_indices] = proposal_ids[masked_indices]
                break

            reveal_mask = self._select_gap_eval_window_transfer_tokens(
                masked_indices=masked_indices,
                proposal_scores_full=proposal_scores_full,
                block_size=block_length,
                num_transfer_tokens=int(transfer_schedule[current_stage].item()),
                strategy=rollout_strategy,
                confidence_threshold=confidence_threshold,
            )
            fill_mask = reveal_mask if reveal_mask.any() else masked_indices
            current_inputs[fill_mask] = proposal_ids[fill_mask]
            current_stage += 1

        return current_inputs

    @staticmethod
    def _compute_gap_answer_start(labels_row: torch.LongTensor, valid_length: int) -> int:
        valid_length = max(0, int(valid_length))
        if valid_length <= 0:
            return 0
        answer_mask = labels_row[:valid_length].ne(-100)
        if not answer_mask.any():
            return valid_length
        return int(torch.nonzero(answer_mask, as_tuple=False)[0].item())

    def _should_stop_gap_eval_rollout(
        self,
        generated_prefix: torch.LongTensor,
    ) -> bool:
        stop_sequences = self._get_gap_eval_stop_sequences()
        if not stop_sequences:
            return False
        return bool(self._match_gap_eval_stop_sequences(generated_prefix, stop_sequences).any().item())

    @torch.no_grad()
    def _rollout_gap_row_to_terminal_eval_style(
        self,
        clean_input_ids_row: torch.LongTensor,
        noisy_input_ids_row: torch.LongTensor,
        labels_row: torch.LongTensor,
        position_ids_row: torch.LongTensor,
        valid_length: int,
        rollout_strategy: str,
        rollout_confidence_threshold: float,
        sample_temperature: float,
        sample_top_k: int,
        sample_top_p: float,
    ) -> torch.LongTensor:
        block_length = int(self.config.block_size)
        valid_length = max(0, int(valid_length))
        if valid_length <= 0:
            return noisy_input_ids_row

        answer_mask = labels_row[:valid_length].ne(-100)
        if not answer_mask.any():
            return noisy_input_ids_row

        prompt_length = int(torch.nonzero(answer_mask, as_tuple=False)[0].item())
        total_length = int(math.ceil(valid_length / max(1, block_length)) * block_length)
        x = torch.full((1, total_length), self.config.mask_token_id, dtype=noisy_input_ids_row.dtype, device=noisy_input_ids_row.device)
        x[:, :valid_length] = noisy_input_ids_row[:valid_length].unsqueeze(0)

        block_mask = torch.tril(torch.ones(total_length // block_length, total_length // block_length, device=x.device), diagonal=0)
        attention_mask = (
            block_mask.repeat_interleave(block_length, dim=0)
            .repeat_interleave(block_length, dim=1)
            .unsqueeze(0)
            .unsqueeze(1)
        )
        valid_positions = torch.zeros((1, total_length), dtype=attention_mask.dtype, device=x.device)
        valid_positions[:, :valid_length] = 1
        attention_mask = attention_mask * valid_positions[:, None, None, :]
        attention_mask = attention_mask * valid_positions[:, None, :, None]

        rollout_position_ids = torch.zeros((1, total_length), dtype=position_ids_row.dtype, device=x.device)
        rollout_position_ids[:, :valid_length] = position_ids_row[:valid_length].unsqueeze(0)
        if total_length > valid_length:
            start_position = int(position_ids_row[valid_length - 1].item()) + 1 if valid_length > 0 else 0
            rollout_position_ids[:, valid_length:] = torch.arange(
                start_position,
                start_position + (total_length - valid_length),
                dtype=position_ids_row.dtype,
                device=x.device,
            ).unsqueeze(0)

        num_blocks = total_length // block_length
        prefill_blocks = prompt_length // block_length
        prefill_length = prefill_blocks * block_length
        prefix_cache = self._build_gap_eval_rollout_prefix_cache(
            x=x,
            attention_mask=attention_mask,
            position_ids=rollout_position_ids,
            prefix_token_end=prefill_length,
        )
        window_start_block = prefill_blocks
        finished = False

        for block_idx in range(prefill_blocks, num_blocks):
            if finished:
                break

            while block_idx - window_start_block >= 1:
                commit_start = window_start_block * block_length
                commit_end = commit_start + block_length
                self._commit_gap_eval_rollout_block(
                    cache=prefix_cache,
                    x=x,
                    attention_mask=attention_mask,
                    position_ids=rollout_position_ids,
                    block_start=commit_start,
                    block_end=commit_end,
                )
                window_start_block += 1

            window_token_start = window_start_block * block_length
            window_token_end = (block_idx + 1) * block_length
            window_slice = slice(window_token_start, window_token_end)
            window_inputs = x[:, window_slice].clone()
            window_attention_mask = attention_mask[:, :, window_token_start:window_token_end, :window_token_end]
            window_position_ids = rollout_position_ids[:, window_slice]

            window_inputs = self._decode_gap_eval_window(
                window_inputs=window_inputs,
                window_attention_mask=window_attention_mask,
                window_position_ids=window_position_ids,
                prefix_cache=prefix_cache,
                mask_id=self.config.mask_token_id,
                block_length=block_length,
                denoising_steps=max(1, int(getattr(self.config, "gap_rollout_steps", block_length))),
                temperature=sample_temperature,
                top_k=sample_top_k,
                top_p=sample_top_p,
                rollout_strategy=rollout_strategy,
                confidence_threshold=rollout_confidence_threshold,
            )
            x[:, window_slice] = window_inputs

            generated_prefix = x[:, prompt_length: min(valid_length, window_token_end)]
            if self._should_stop_gap_eval_rollout(generated_prefix[0]):
                eos_token_id = getattr(self.config, "eos_token_id", None)
                if isinstance(eos_token_id, (list, tuple)):
                    eos_fill = int(eos_token_id[0])
                elif eos_token_id is None:
                    eos_fill = self.config.mask_token_id
                else:
                    eos_fill = int(eos_token_id)
                if window_token_end < valid_length:
                    x[:, window_token_end:valid_length] = eos_fill
                finished = True

        terminal_row = noisy_input_ids_row.clone()
        terminal_row[:valid_length] = x[0, :valid_length]
        return terminal_row

    @torch.no_grad()
    def _rollout_gap_group_to_terminal_eval_style(
        self,
        clean_input_ids: torch.LongTensor,
        noisy_input_ids: torch.LongTensor,
        labels: torch.LongTensor,
        position_ids: torch.LongTensor,
        valid_lengths: list[int],
        prompt_lengths: list[int],
        rollout_strategy: str,
        rollout_confidence_threshold: float,
        sample_temperature: float,
        sample_top_k: int,
        sample_top_p: float,
    ) -> torch.LongTensor:
        batch_size = noisy_input_ids.shape[0]
        block_length = int(self.config.block_size)
        heartbeat_interval = _gap_env_int("SDAR_GAP_ROLLOUT_HEARTBEAT_INTERVAL_BLOCKS", 0)
        heartbeat_enabled = heartbeat_interval > 0 and _gap_is_rank0()
        debug_step = getattr(self, "_gap_debug_global_step", -1)
        valid_lengths = [max(0, int(v)) for v in valid_lengths]
        prompt_lengths = [max(0, int(p)) for p in prompt_lengths]
        total_lengths = [
            int(math.ceil(valid_length / max(1, block_length)) * block_length)
            for valid_length in valid_lengths
        ]
        max_total_length = max(total_lengths, default=0)
        if max_total_length <= 0:
            return noisy_input_ids

        device = noisy_input_ids.device
        dtype = noisy_input_ids.dtype
        x = torch.full(
            (batch_size, max_total_length),
            self.config.mask_token_id,
            dtype=dtype,
            device=device,
        )
        for row_idx, valid_length in enumerate(valid_lengths):
            if valid_length > 0:
                x[row_idx, :valid_length] = noisy_input_ids[row_idx, :valid_length]

        num_blocks = max_total_length // block_length
        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=device), diagonal=0)
        attention_mask = (
            block_mask.repeat_interleave(block_length, dim=0)
            .repeat_interleave(block_length, dim=1)
            .unsqueeze(0)
            .unsqueeze(1)
            .expand(batch_size, -1, -1, -1)
            .clone()
        )

        valid_positions = torch.zeros((batch_size, max_total_length), dtype=attention_mask.dtype, device=device)
        for row_idx, valid_length in enumerate(valid_lengths):
            if valid_length > 0:
                valid_positions[row_idx, :valid_length] = 1
        attention_mask = attention_mask * valid_positions[:, None, None, :]
        attention_mask = attention_mask * valid_positions[:, None, :, None]

        rollout_position_ids = torch.zeros((batch_size, max_total_length), dtype=position_ids.dtype, device=device)
        for row_idx, valid_length in enumerate(valid_lengths):
            if valid_length <= 0:
                continue
            rollout_position_ids[row_idx, :valid_length] = position_ids[row_idx, :valid_length]
            if max_total_length > valid_length:
                start_position = int(position_ids[row_idx, valid_length - 1].item()) + 1
                rollout_position_ids[row_idx, valid_length:max_total_length] = torch.arange(
                    start_position,
                    start_position + (max_total_length - valid_length),
                    dtype=position_ids.dtype,
                    device=device,
                )

        prefill_blocks = min(prompt_length // block_length for prompt_length in prompt_lengths)
        prefill_length = prefill_blocks * block_length
        prefix_cache = self._build_gap_eval_rollout_prefix_cache(
            x=x,
            attention_mask=attention_mask,
            position_ids=rollout_position_ids,
            prefix_token_end=prefill_length,
        )
        window_start_block = prefill_blocks
        row_num_blocks = torch.tensor(
            [max(0, total_length // block_length) for total_length in total_lengths],
            dtype=torch.long,
            device=device,
        )
        finished = row_num_blocks.le(prefill_blocks)

        for block_idx in range(prefill_blocks, num_blocks):
            if bool(finished.all().item()):
                break

            active_rows = (~finished) & row_num_blocks.gt(block_idx)
            if not bool(active_rows.any().item()):
                break

            if heartbeat_enabled and ((block_idx - prefill_blocks) % heartbeat_interval == 0):
                logger.info(
                    "[GAP rollout heartbeat] step=%s block=%s/%s active_rows=%s finished=%s",
                    debug_step,
                    int(block_idx - prefill_blocks),
                    int(max(num_blocks - prefill_blocks, 0)),
                    int(active_rows.sum().item()),
                    int(finished.sum().item()),
                )

            while block_idx - window_start_block >= 1:
                commit_start = window_start_block * block_length
                commit_end = commit_start + block_length
                self._commit_gap_eval_rollout_block(
                    cache=prefix_cache,
                    x=x,
                    attention_mask=attention_mask,
                    position_ids=rollout_position_ids,
                    block_start=commit_start,
                    block_end=commit_end,
                )
                window_start_block += 1

            window_token_start = window_start_block * block_length
            window_token_end = (block_idx + 1) * block_length
            window_slice = slice(window_token_start, window_token_end)
            frozen_window_inputs = x[:, window_slice].clone()
            window_inputs = x[:, window_slice].clone()
            window_attention_mask = attention_mask[:, :, window_token_start:window_token_end, :window_token_end]
            window_position_ids = rollout_position_ids[:, window_slice]

            window_inputs = self._decode_gap_eval_window(
                window_inputs=window_inputs,
                window_attention_mask=window_attention_mask,
                window_position_ids=window_position_ids,
                prefix_cache=prefix_cache,
                mask_id=self.config.mask_token_id,
                block_length=block_length,
                denoising_steps=max(1, int(getattr(self.config, "gap_rollout_steps", block_length))),
                temperature=sample_temperature,
                top_k=sample_top_k,
                top_p=sample_top_p,
                rollout_strategy=rollout_strategy,
                confidence_threshold=rollout_confidence_threshold,
            )
            window_inputs[~active_rows] = frozen_window_inputs[~active_rows]
            x[:, window_slice] = window_inputs

            for row_idx in torch.nonzero(active_rows, as_tuple=False).flatten().tolist():
                row_prompt_length = prompt_lengths[row_idx]
                row_valid_length = valid_lengths[row_idx]
                generated_prefix = x[row_idx, row_prompt_length:min(row_valid_length, window_token_end)]
                if self._should_stop_gap_eval_rollout(generated_prefix):
                    eos_token_id = getattr(self.config, "eos_token_id", None)
                    if isinstance(eos_token_id, (list, tuple)):
                        eos_fill = int(eos_token_id[0])
                    elif eos_token_id is None:
                        eos_fill = self.config.mask_token_id
                    else:
                        eos_fill = int(eos_token_id)
                    if window_token_end < row_valid_length:
                        x[row_idx, window_token_end:row_valid_length] = eos_fill
                    finished[row_idx] = True

            finished |= row_num_blocks.le(block_idx + 1)

        terminal_input_ids = noisy_input_ids.clone()
        for row_idx, valid_length in enumerate(valid_lengths):
            if valid_length > 0:
                terminal_input_ids[row_idx, :valid_length] = x[row_idx, :valid_length]
        return terminal_input_ids

    @torch.no_grad()
    def _rollout_gap_group_eval_with_sampled_remask(
        self,
        noisy_input_ids: torch.LongTensor,
        labels: torch.LongTensor,
        position_ids: torch.LongTensor,
        valid_lengths: list[int],
        prompt_lengths: list[int],
        num_samples: int,
        rollout_strategy: str,
        rollout_confidence_threshold: float,
        sample_temperature: float,
        sample_top_k: int,
        sample_top_p: float,
        enable_remask_actions: bool = True,
    ) -> tuple[torch.LongTensor, torch.BoolTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.BoolTensor, torch.BoolTensor, torch.BoolTensor]:
        base_batch_size = noisy_input_ids.shape[0]
        block_length = int(self.config.block_size)
        num_samples = max(1, int(num_samples))
        batch_size = base_batch_size * num_samples
        source_valid_lengths = [max(0, int(v)) for _ in range(num_samples) for v in valid_lengths]
        prompt_lengths = [max(0, int(p)) for _ in range(num_samples) for p in prompt_lengths]
        eval_gen_length = _gap_env_int("SDAR_GAP_GRPO_EVAL_GEN_LENGTH", 1024)
        eval_gen_length = max(1, int(eval_gen_length))
        valid_lengths = [prompt_length + eval_gen_length for prompt_length in prompt_lengths]
        total_lengths = [
            int(math.ceil(valid_length / max(1, block_length)) * block_length)
            for valid_length in valid_lengths
        ]
        max_total_length = max(total_lengths, default=0)
        if max_total_length <= 0:
            empty = noisy_input_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(batch_size, -1)
            zeros = torch.zeros((num_samples, base_batch_size), dtype=torch.float32, device=noisy_input_ids.device)
            empty_mask = torch.zeros_like(empty, dtype=torch.bool)
            empty_flags = torch.zeros((num_samples, base_batch_size), dtype=torch.bool, device=noisy_input_ids.device)
            return empty, empty_mask, zeros, zeros.mean(dim=0), zeros, empty_flags, empty_flags, torch.zeros((num_samples, base_batch_size, empty.shape[1]), dtype=torch.bool, device=noisy_input_ids.device)

        strict_eval_decode = _gap_env_flag("SDAR_GAP_GRPO_STRICT_EVAL_DECODE", False)
        restore_model_training = bool(strict_eval_decode and self.model.training)
        if restore_model_training:
            self.model.eval()

        device = noisy_input_ids.device
        dtype = noisy_input_ids.dtype
        x = torch.full((batch_size, max_total_length), self.config.mask_token_id, dtype=dtype, device=device)
        repeated_noisy = noisy_input_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(batch_size, -1)
        repeated_positions = position_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(batch_size, -1)
        for row_idx, prompt_length in enumerate(prompt_lengths):
            prompt_copy_len = min(prompt_length, repeated_noisy.shape[1])
            if prompt_copy_len > 0:
                x[row_idx, :prompt_copy_len] = repeated_noisy[row_idx, :prompt_copy_len]

        num_blocks = max_total_length // block_length
        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=device), diagonal=0)
        attention_mask = (
            block_mask.repeat_interleave(block_length, dim=0)
            .repeat_interleave(block_length, dim=1)
            .unsqueeze(0)
            .unsqueeze(1)
            .expand(batch_size, -1, -1, -1)
            .clone()
        )
        valid_positions = torch.zeros((batch_size, max_total_length), dtype=attention_mask.dtype, device=device)
        for row_idx, valid_length in enumerate(valid_lengths):
            if valid_length > 0:
                valid_positions[row_idx, :valid_length] = 1
        attention_mask = attention_mask * valid_positions[:, None, None, :]
        attention_mask = attention_mask * valid_positions[:, None, :, None]

        rollout_position_ids = torch.zeros((batch_size, max_total_length), dtype=position_ids.dtype, device=device)
        for row_idx, valid_length in enumerate(valid_lengths):
            if valid_length <= 0:
                continue
            source_valid_length = min(source_valid_lengths[row_idx], repeated_positions.shape[1])
            copy_len = min(source_valid_length, valid_length)
            if copy_len > 0:
                rollout_position_ids[row_idx, :copy_len] = repeated_positions[row_idx, :copy_len]
                start_position = int(repeated_positions[row_idx, copy_len - 1].item()) + 1
            else:
                start_position = 0
            if valid_length > copy_len:
                rollout_position_ids[row_idx, copy_len:valid_length] = torch.arange(
                    start_position,
                    start_position + (valid_length - copy_len),
                    dtype=position_ids.dtype,
                    device=device,
                )

        prefill_blocks = min(prompt_length // block_length for prompt_length in prompt_lengths)
        prefill_length = prefill_blocks * block_length
        prefix_cache = self._build_gap_eval_rollout_prefix_cache(
            x=x,
            attention_mask=attention_mask,
            position_ids=rollout_position_ids,
            prefix_token_end=prefill_length,
        )
        window_start_block = prefill_blocks
        row_num_blocks = torch.tensor([max(0, total_length // block_length) for total_length in total_lengths], dtype=torch.long, device=device)
        finished = row_num_blocks.le(prefill_blocks)

        logprob_sums = torch.zeros(batch_size, dtype=torch.float32, device=device)
        entropy_sums = torch.zeros(batch_size, dtype=torch.float32, device=device)
        candidate_counts = torch.zeros(batch_size, dtype=torch.float32, device=device)
        remask_counts = torch.zeros(batch_size, dtype=torch.float32, device=device)
        action_full_flat = torch.zeros((batch_size, max_total_length), dtype=torch.bool, device=device)
        stop_hits = torch.zeros(batch_size, dtype=torch.bool, device=device)
        prefix_guard_tokens = _gap_env_int("SDAR_GAP_GRPO_REMASK_PREFIX_GUARD_TOKENS", 0)
        tail_guard_blocks = _gap_env_int("SDAR_GAP_GRPO_REMASK_TAIL_GUARD_BLOCKS", 1)
        window_blocks = max(1, _gap_env_int("SDAR_GAP_GRPO_MIDDLE_WINDOW_BLOCKS", int(getattr(self.config, "gap_grpo_candidate_window_blocks", 4) or 4)))
        exclude_frontier_blocks = max(0, _gap_env_int("SDAR_GAP_GRPO_EXCLUDE_FRONTIER_BLOCKS", 0))
        remask_interval_blocks = max(1, _gap_env_int("SDAR_GAP_GRPO_REMASK_INTERVAL_BLOCKS", 1))
        max_remask_candidates = max(1, int(getattr(self.config, "gap_remask_adv_max_candidates", 1) or 1))
        prob_eps = float(getattr(self.config, "gap_grpo_sample_prob_eps", 1e-4) or 1e-4)

        for block_idx in range(prefill_blocks, num_blocks):
            if bool(finished.all().item()):
                break
            active_rows = (~finished) & row_num_blocks.gt(block_idx)
            if not bool(active_rows.any().item()):
                break

            while block_idx - window_start_block >= window_blocks:
                commit_start = window_start_block * block_length
                commit_end = commit_start + block_length
                self._commit_gap_eval_rollout_block(
                    cache=prefix_cache,
                    x=x,
                    attention_mask=attention_mask,
                    position_ids=rollout_position_ids,
                    block_start=commit_start,
                    block_end=commit_end,
                )
                window_start_block += 1

            window_token_start = window_start_block * block_length
            window_token_end = (block_idx + 1) * block_length
            window_slice = slice(window_token_start, window_token_end)
            frozen_window_inputs = x[:, window_slice].clone()
            window_inputs = x[:, window_slice].clone()
            window_attention_mask = attention_mask[:, :, window_token_start:window_token_end, :window_token_end]
            window_position_ids = rollout_position_ids[:, window_slice]

            window_inputs = self._decode_gap_eval_window(
                window_inputs=window_inputs,
                window_attention_mask=window_attention_mask,
                window_position_ids=window_position_ids,
                prefix_cache=prefix_cache,
                mask_id=self.config.mask_token_id,
                block_length=block_length,
                denoising_steps=max(1, int(getattr(self.config, "gap_rollout_steps", block_length))),
                temperature=sample_temperature,
                top_k=sample_top_k,
                top_p=sample_top_p,
                rollout_strategy=rollout_strategy,
                confidence_threshold=rollout_confidence_threshold,
            )

            global_positions = torch.arange(window_token_start, window_token_end, device=device).unsqueeze(0)
            rolling_end = (block_idx + 1 - exclude_frontier_blocks) * block_length
            rolling_start = max(prefill_blocks * block_length, rolling_end - window_blocks * block_length)
            candidate_mask = (
                global_positions.ge(rolling_start)
                & global_positions.lt(rolling_end)
                & window_inputs.ne(self.config.mask_token_id)
                & active_rows.unsqueeze(1)
            )
            for row_idx, prompt_length in enumerate(prompt_lengths):
                candidate_mask[row_idx] &= global_positions[0].ge(prompt_length + max(0, prefix_guard_tokens))
            if tail_guard_blocks > 0:
                tail_guard_start = window_token_end - tail_guard_blocks * block_length
                candidate_mask &= global_positions.lt(tail_guard_start)

            generated_blocks = (block_idx - prefill_blocks) + 1
            remask_active = (generated_blocks - 1) % remask_interval_blocks == 0

            if enable_remask_actions and remask_active and candidate_mask.any():
                score_attention_mask = self._to_gap_block_mask(window_attention_mask) if self.model.training and not strict_eval_decode else window_attention_mask
                outputs = self.model(
                    window_inputs,
                    attention_mask=score_attention_mask,
                    position_ids=window_position_ids,
                    past_key_values=prefix_cache,
                    use_cache=True,
                    store_kv=False,
                    output_hidden_states=True,
                    return_dict=True,
                )
                score_logits = self.lm_head(outputs.last_hidden_state).float()
                action_probs = torch.sigmoid(self.gap_remask_head(outputs.last_hidden_state, score_logits)).float().clamp(prob_eps, 1.0 - prob_eps)
                sampled_actions = torch.bernoulli(action_probs).to(torch.bool) & candidate_mask
                for row_idx in range(batch_size):
                    row_candidates = torch.nonzero(candidate_mask[row_idx], as_tuple=False).flatten()
                    if row_candidates.numel() == 0:
                        continue
                    row_actions = torch.nonzero(sampled_actions[row_idx], as_tuple=False).flatten()
                    if row_actions.numel() > max_remask_candidates:
                        keep = torch.topk(action_probs[row_idx, row_actions], k=max_remask_candidates, sorted=False).indices
                        new_row = torch.zeros_like(sampled_actions[row_idx])
                        new_row[row_actions[keep]] = True
                        sampled_actions[row_idx] = new_row
                    p = action_probs[row_idx, row_candidates]
                    a = sampled_actions[row_idx, row_candidates].to(p.dtype)
                    logprob_sums[row_idx] += (a * p.log() + (1.0 - a) * (1.0 - p).log()).sum()
                    entropy_sums[row_idx] += (-(p * p.log() + (1.0 - p) * (1.0 - p).log())).sum()
                    candidate_counts[row_idx] += float(row_candidates.numel())
                if sampled_actions.any():
                    window_inputs[sampled_actions] = self.config.mask_token_id
                    global_action_positions = global_positions.expand(batch_size, -1)[sampled_actions]
                    global_action_rows = torch.nonzero(sampled_actions, as_tuple=False)[:, 0]
                    valid_action_mask = global_action_positions.lt(action_full_flat.shape[1])
                    if valid_action_mask.any():
                        action_full_flat[global_action_rows[valid_action_mask], global_action_positions[valid_action_mask]] = True
                    remask_counts += sampled_actions.to(torch.float32).sum(dim=1)
                    window_inputs = self._decode_gap_eval_window(
                        window_inputs=window_inputs,
                        window_attention_mask=window_attention_mask,
                        window_position_ids=window_position_ids,
                        prefix_cache=prefix_cache,
                        mask_id=self.config.mask_token_id,
                        block_length=block_length,
                        denoising_steps=max(1, int(getattr(self.config, "gap_rollout_steps", block_length))),
                        temperature=sample_temperature,
                        top_k=sample_top_k,
                        top_p=sample_top_p,
                        rollout_strategy=rollout_strategy,
                        confidence_threshold=rollout_confidence_threshold,
                    )

            window_inputs[~active_rows] = frozen_window_inputs[~active_rows]
            x[:, window_slice] = window_inputs
            for row_idx in torch.nonzero(active_rows, as_tuple=False).flatten().tolist():
                generated_prefix = x[row_idx, prompt_lengths[row_idx]:min(valid_lengths[row_idx], window_token_end)]
                if self._should_stop_gap_eval_rollout(generated_prefix):
                    eos_token_id = getattr(self.config, "eos_token_id", None)
                    if isinstance(eos_token_id, (list, tuple)):
                        eos_fill = int(eos_token_id[0])
                    elif eos_token_id is None:
                        eos_fill = self.config.mask_token_id
                    else:
                        eos_fill = int(eos_token_id)
                    if window_token_end < valid_lengths[row_idx]:
                        x[row_idx, window_token_end:valid_lengths[row_idx]] = eos_fill
                    stop_hits[row_idx] = True
                    finished[row_idx] = True
            finished |= row_num_blocks.le(block_idx + 1)

        terminal = x.clone()
        eval_target_mask = torch.zeros_like(terminal, dtype=torch.bool)
        for row_idx, valid_length in enumerate(valid_lengths):
            prompt_length = min(prompt_lengths[row_idx], valid_length)
            if valid_length > prompt_length:
                eval_target_mask[row_idx, prompt_length:valid_length] = True
        length_cap_hits = ~stop_hits
        logprob = (logprob_sums / candidate_counts.clamp_min(1.0)).view(num_samples, base_batch_size)
        entropy = (entropy_sums / candidate_counts.clamp_min(1.0)).view(num_samples, base_batch_size).mean(dim=0)
        remask_rate = (remask_counts / candidate_counts.clamp_min(1.0)).view(num_samples, base_batch_size)
        action_full = action_full_flat.view(num_samples, base_batch_size, -1)
        if restore_model_training:
            self.model.train()
        return (
            terminal,
            eval_target_mask,
            logprob,
            entropy,
            remask_rate,
            stop_hits.view(num_samples, base_batch_size),
            length_cap_hits.view(num_samples, base_batch_size),
            action_full,
        )

    @torch.no_grad()
    def _compute_gap_masked_proposals(
        self,
        clean_input_ids: torch.LongTensor,
        noisy_input_ids: torch.LongTensor,
        position_ids: torch.LongTensor,
        masked_indices: torch.BoolTensor,
        num_tokens,
        sample_temperature: float = 0.0,
        sample_top_k: int = 0,
        sample_top_p: float = 1.0,
    ) -> tuple[torch.LongTensor, torch.FloatTensor]:
        concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, logits_to_keep_half, logits_to_keep, _ = self.build_bd_training_inputs(
            inputs_ids=clean_input_ids,
            noisy_inputs_ids=noisy_input_ids,
            position_ids=position_ids,
            logits_to_keep_half=masked_indices,
            num_tokens=num_tokens,
        )
        outputs = self.model(
            input_ids=concat_inputs_ids,
            attention_mask=flex_attention_mask_3d,
            position_ids=concat_position_ids,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=True,
        )
        hidden_states = outputs.last_hidden_state[logits_to_keep].contiguous()
        proposal_logits = self.lm_head(hidden_states).float()
        proposal_ids = self._sample_from_logits(
            logits=proposal_logits,
            temperature=sample_temperature,
            top_k=sample_top_k,
            top_p=sample_top_p,
        )
        proposal_probs = proposal_logits.softmax(dim=-1).gather(-1, proposal_ids.unsqueeze(-1)).squeeze(-1)
        proposal_ids_full = torch.full_like(noisy_input_ids, self.config.mask_token_id)
        proposal_scores_full = torch.full(
            noisy_input_ids.shape,
            float("-inf"),
            dtype=proposal_probs.dtype,
            device=proposal_probs.device,
        )
        proposal_ids_full[masked_indices] = proposal_ids
        proposal_scores_full[masked_indices] = proposal_probs
        return proposal_ids_full, proposal_scores_full

    @torch.no_grad()
    def _rollout_gap_state_to_terminal(
        self,
        clean_input_ids: torch.LongTensor,
        noisy_input_ids: torch.LongTensor,
        labels: torch.LongTensor,
        position_ids: torch.LongTensor,
        num_tokens,
        start_stage: int = 0,
        rollout_strategy: Optional[str] = None,
        rollout_confidence_threshold: Optional[float] = None,
        rollout_scope: Optional[str] = None,
        sample_temperature: float = 0.0,
        sample_top_k: int = 0,
        sample_top_p: float = 1.0,
    ) -> torch.LongTensor:
        rollout_strategy = rollout_strategy or getattr(self.config, "gap_rollout_strategy", "low_confidence_dynamic")
        rollout_confidence_threshold = float(
            rollout_confidence_threshold
            if rollout_confidence_threshold is not None
            else getattr(self.config, "gap_rollout_confidence_threshold", 0.95)
        )
        _ = rollout_scope or getattr(self.config, "gap_grpo_terminal_rollout_scope", None) or getattr(self.config, "gap_rollout_scope", "all")
        terminal_input_ids = noisy_input_ids.clone()
        rollout_groups: dict[tuple[int, int], list[int]] = {}
        valid_lengths: list[int] = []
        prompt_lengths: list[int] = []
        block_length = int(self.config.block_size)

        for row_idx in range(terminal_input_ids.shape[0]):
            valid_length = self._resolve_gap_valid_length(num_tokens[row_idx])
            prompt_length = self._compute_gap_answer_start(labels[row_idx], valid_length)
            valid_lengths.append(valid_length)
            prompt_lengths.append(prompt_length)
            total_length = int(math.ceil(valid_length / max(1, block_length)) * block_length) if valid_length > 0 else 0
            group_key = (prompt_length // max(1, block_length), total_length)
            rollout_groups.setdefault(group_key, []).append(row_idx)

        was_base_training = self.model.training
        if was_base_training:
            self.model.eval()
        try:
            for (_, _), row_indices in rollout_groups.items():
                group_clean = clean_input_ids[row_indices]
                group_noisy = terminal_input_ids[row_indices]
                group_labels = labels[row_indices]
                group_positions = position_ids[row_indices]
                group_valid_lengths = [valid_lengths[idx] for idx in row_indices]
                group_prompt_lengths = [prompt_lengths[idx] for idx in row_indices]
                group_terminal = self._rollout_gap_group_to_terminal_eval_style(
                    clean_input_ids=group_clean,
                    noisy_input_ids=group_noisy,
                    labels=group_labels,
                    position_ids=group_positions,
                    valid_lengths=group_valid_lengths,
                    prompt_lengths=group_prompt_lengths,
                    rollout_strategy=rollout_strategy,
                    rollout_confidence_threshold=rollout_confidence_threshold,
                    sample_temperature=sample_temperature,
                    sample_top_k=sample_top_k,
                    sample_top_p=sample_top_p,
                )
                terminal_input_ids[row_indices] = group_terminal
        finally:
            if was_base_training:
                self.model.train()

        return terminal_input_ids

    @torch.no_grad()
    def _compute_gap_terminal_answer_rewards(
        self,
        clean_input_ids: torch.LongTensor,
        terminal_input_ids: torch.LongTensor,
        target_scope_mask: torch.BoolTensor,
        pred_scope_mask: Optional[torch.BoolTensor] = None,
    ) -> torch.FloatTensor:
        if target_scope_mask.dim() == 1:
            target_scope_mask = target_scope_mask.unsqueeze(0)
        if pred_scope_mask is None:
            pred_scope_mask = target_scope_mask
        elif pred_scope_mask.dim() == 1:
            pred_scope_mask = pred_scope_mask.unsqueeze(0)
        if clean_input_ids.dim() == 1:
            clean_input_ids = clean_input_ids.unsqueeze(0)
        if terminal_input_ids.dim() == 1:
            terminal_input_ids = terminal_input_ids.unsqueeze(0)

        if clean_input_ids.shape[0] == 1 and terminal_input_ids.shape[0] > 1:
            clean_input_ids = clean_input_ids.expand(terminal_input_ids.shape[0], -1)
        if target_scope_mask.shape[0] == 1 and terminal_input_ids.shape[0] > 1:
            target_scope_mask = target_scope_mask.expand(terminal_input_ids.shape[0], -1)
        if pred_scope_mask.shape[0] == 1 and terminal_input_ids.shape[0] > 1:
            pred_scope_mask = pred_scope_mask.expand(terminal_input_ids.shape[0], -1)

        rewards = torch.zeros(terminal_input_ids.shape[0], dtype=torch.float32, device=terminal_input_ids.device)
        for row_idx in range(terminal_input_ids.shape[0]):
            gold_text = self._decode_gap_response_tokens(clean_input_ids[row_idx][target_scope_mask[row_idx]])
            pred_text = self._decode_gap_response_tokens(terminal_input_ids[row_idx][pred_scope_mask[row_idx]])
            rewards[row_idx] = 1.0 if self._score_gap_answer_with_opencompass(pred_text, gold_text) else 0.0
        return rewards

    def _compute_gap_sequence_logprob_means(
        self,
        model,
        input_ids: torch.LongTensor,
        position_ids: torch.LongTensor,
        target_scope_mask: torch.BoolTensor,
        require_grad: bool = False,
        chunk_size: int = 0,
    ) -> torch.FloatTensor:
        if target_scope_mask.dim() == 1:
            target_scope_mask = target_scope_mask.unsqueeze(0)
        if input_ids.dim() == 1:
            input_ids = input_ids.unsqueeze(0)
        if position_ids.dim() == 1:
            position_ids = position_ids.unsqueeze(0)
        position_ids = modify_padded_position_ids_2d(position_ids)
        chunk_size = int(chunk_size or input_ids.shape[0] or 1)

        def _score_chunk(chunk_input_ids, chunk_position_ids, chunk_target_scope_mask):
            num_tokens = calculate_token_nums(chunk_position_ids)
            masked_indices = chunk_target_scope_mask.to(dtype=torch.bool)
            concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, _, logits_to_keep, _ = model.build_bd_training_inputs(
                inputs_ids=chunk_input_ids,
                noisy_inputs_ids=chunk_input_ids,
                position_ids=chunk_position_ids,
                logits_to_keep_half=masked_indices,
                num_tokens=num_tokens,
            )
            was_training = model.training
            model.train()
            try:
                outputs = model.model(
                    input_ids=concat_inputs_ids,
                    attention_mask=flex_attention_mask_3d,
                    position_ids=concat_position_ids,
                    output_attentions=False,
                    output_hidden_states=False,
                    return_dict=True,
                )
            finally:
                if not was_training:
                    model.eval()
            hidden_states = outputs.last_hidden_state[logits_to_keep].contiguous()
            logits = model.lm_head(hidden_states).float()
            target_ids = chunk_input_ids[masked_indices]
            target_batch = torch.nonzero(masked_indices, as_tuple=False)[:, 0]
            token_logprobs = logits.log_softmax(dim=-1).gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)
            batch_sums = torch.zeros(chunk_input_ids.shape[0], dtype=token_logprobs.dtype, device=token_logprobs.device)
            batch_counts = torch.zeros_like(batch_sums)
            batch_sums.scatter_add_(0, target_batch, token_logprobs)
            batch_counts.scatter_add_(0, target_batch, torch.ones_like(token_logprobs))
            return batch_sums / batch_counts.clamp_min(1.0)

        def _forward():
            if input_ids.shape[0] <= chunk_size:
                return _score_chunk(input_ids, position_ids, target_scope_mask)

            outputs = []
            for start in range(0, input_ids.shape[0], chunk_size):
                end = min(start + chunk_size, input_ids.shape[0])
                outputs.append(
                    _score_chunk(
                        input_ids[start:end],
                        position_ids[start:end],
                        target_scope_mask[start:end],
                    )
                )
            return torch.cat(outputs, dim=0)

        if require_grad:
            return _forward()
        with torch.no_grad():
            return _forward()

    @staticmethod
    def _extend_gap_position_ids_to_length(
        position_ids: torch.LongTensor,
        target_length: int,
    ) -> torch.LongTensor:
        if position_ids.dim() == 1:
            position_ids = position_ids.unsqueeze(0)
        target_length = int(target_length)
        if position_ids.shape[1] == target_length:
            return position_ids

        batch_size, source_length = position_ids.shape
        output = torch.zeros(
            (batch_size, target_length),
            dtype=position_ids.dtype,
            device=position_ids.device,
        )
        copy_length = min(source_length, target_length)
        if copy_length > 0:
            output[:, :copy_length] = position_ids[:, :copy_length]
        if target_length > copy_length:
            if copy_length > 0:
                start_positions = output[:, copy_length - 1] + 1
            else:
                start_positions = torch.zeros(batch_size, dtype=position_ids.dtype, device=position_ids.device)
            offsets = torch.arange(
                target_length - copy_length,
                dtype=position_ids.dtype,
                device=position_ids.device,
            ).unsqueeze(0)
            output[:, copy_length:] = start_positions.unsqueeze(1) + offsets
        return output

    @staticmethod
    def _compute_gap_group_advantages(
        rewards: torch.FloatTensor,
        baseline_reward: Optional[torch.FloatTensor],
        advantage_eps: float,
    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
        if baseline_reward is not None:
            centered_rewards = rewards - baseline_reward.unsqueeze(0)
        else:
            centered_rewards = rewards
        reward_mean = centered_rewards.mean(dim=0, keepdim=True)
        reward_std = centered_rewards.std(dim=0, keepdim=True)
        advantages = (centered_rewards - reward_mean) / reward_std.clamp_min(advantage_eps)
        advantages = torch.where(reward_std.gt(advantage_eps), advantages, torch.zeros_like(advantages))
        return advantages, reward_std

    def _compute_gap_current_remask_action_logprob(
        self,
        remask_logits: torch.FloatTensor,
        masked_indices: torch.BoolTensor,
        full_candidate_mask: torch.BoolTensor,
        sampled_full: torch.BoolTensor,
        sample_prob_eps: float,
    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
        candidate_mask_flat = full_candidate_mask[masked_indices]
        batch_size = full_candidate_mask.shape[0]
        candidate_count = int(full_candidate_mask.sum().item())
        if candidate_count <= 0:
            num_samples = sampled_full.shape[0]
            zero = remask_logits.sum() * 0.0
            return zero.expand(num_samples, batch_size), zero.expand(batch_size)

        flat_probs = torch.sigmoid(remask_logits[candidate_mask_flat]).clamp(
            min=sample_prob_eps,
            max=1.0 - sample_prob_eps,
        )
        batch_ids = torch.nonzero(full_candidate_mask, as_tuple=False)[:, 0]
        candidate_counts_per_batch = torch.zeros(batch_size, dtype=flat_probs.dtype, device=flat_probs.device)
        candidate_counts_per_batch.scatter_add_(
            0,
            batch_ids,
            torch.ones_like(batch_ids, dtype=flat_probs.dtype),
        )
        candidate_counts_per_batch = candidate_counts_per_batch.clamp_min(1.0)

        num_samples = sampled_full.shape[0]
        aligned_actions = torch.zeros(
            (num_samples, batch_size, full_candidate_mask.shape[1]),
            dtype=torch.bool,
            device=full_candidate_mask.device,
        )
        copy_len = min(full_candidate_mask.shape[1], sampled_full.shape[-1])
        if copy_len > 0:
            aligned_actions[:, :, :copy_len] = sampled_full[:, :, :copy_len]
        sampled_flat = aligned_actions[:, full_candidate_mask]

        sample_logprob = (
            sampled_flat.to(flat_probs.dtype) * flat_probs.log().unsqueeze(0)
            + (~sampled_flat).to(flat_probs.dtype) * (1.0 - flat_probs).log().unsqueeze(0)
        )
        sample_entropy = -(
            flat_probs * flat_probs.log() + (1.0 - flat_probs) * (1.0 - flat_probs).log()
        )
        logprob_per_batch = torch.zeros((num_samples, batch_size), dtype=sample_logprob.dtype, device=sample_logprob.device)
        entropy_per_batch = torch.zeros(batch_size, dtype=sample_entropy.dtype, device=sample_entropy.device)
        logprob_per_batch.scatter_add_(1, batch_ids.unsqueeze(0).expand(num_samples, -1), sample_logprob)
        entropy_per_batch.scatter_add_(0, batch_ids, sample_entropy)
        logprob_per_batch = logprob_per_batch / candidate_counts_per_batch.unsqueeze(0)
        entropy_per_batch = entropy_per_batch / candidate_counts_per_batch
        return logprob_per_batch, entropy_per_batch

    def _compute_gap_grpo_loss(
        self,
        clean_input_ids: torch.LongTensor,
        labels: torch.LongTensor,
        position_ids: torch.LongTensor,
        num_tokens,
        remask_logits: torch.FloatTensor,
        grpo_hidden_states: torch.FloatTensor,
        gap_outputs,
        masked_indices: torch.BoolTensor,
        rollout_strategy: str,
        rollout_confidence_threshold: float,
        target_scope_mask: torch.BoolTensor,
    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
        grpo_weight = float(getattr(self.config, "gap_grpo_loss_weight", 0.0) or 0.0)
        num_samples = int(getattr(self.config, "gap_grpo_num_samples", 0) or 0)
        if grpo_weight <= 0.0 or num_samples <= 0:
            zero = remask_logits.sum() * 0.0
            return zero, {}

        use_eval_rollout_actions = _gap_env_flag("SDAR_GAP_GRPO_USE_EVAL_ROLLOUT_ACTIONS", False)
        if use_eval_rollout_actions:
            batch_size = clean_input_ids.shape[0]
            entropy_coef = float(getattr(self.config, "gap_grpo_entropy_coef", 0.0) or 0.0)
            use_baseline_branch = bool(getattr(self.config, "gap_grpo_use_baseline_branch", True))
            terminal_weight = float(getattr(self.config, "gap_grpo_terminal_reward_weight", 1.0) or 1.0)
            format_weight = float(getattr(self.config, "gap_grpo_format_reward_weight", 0.0) or 0.0)
            remask_penalty = float(getattr(self.config, "gap_grpo_remask_penalty", 0.0) or 0.0)
            advantage_eps = float(getattr(self.config, "gap_grpo_advantage_eps", 1e-4) or 1e-4)
            rollout_temperature = float(getattr(self.config, "gap_grpo_rollout_temperature", 0.0) or 0.0)
            rollout_top_k = int(getattr(self.config, "gap_grpo_rollout_top_k", 0) or 0)
            rollout_top_p = float(getattr(self.config, "gap_grpo_rollout_top_p", 1.0) or 1.0)
            terminal_rollout_strategy = getattr(self.config, "gap_grpo_terminal_rollout_strategy", None) or rollout_strategy
            valid_lengths = [self._resolve_gap_valid_length(num_tokens[row_idx]) for row_idx in range(batch_size)]
            prompt_lengths = [
                self._compute_gap_answer_start(labels[row_idx], valid_lengths[row_idx])
                for row_idx in range(batch_size)
            ]

            with torch.no_grad():
                if use_baseline_branch:
                    (
                        baseline_terminal,
                        baseline_eval_target_mask,
                        _baseline_logprob,
                        _baseline_entropy,
                        _baseline_remask_rate,
                        baseline_stop_hit,
                        baseline_length_cap_hit,
                        _baseline_sampled_full,
                    ) = self._rollout_gap_group_eval_with_sampled_remask(
                        noisy_input_ids=gap_outputs.z_accept,
                        labels=labels,
                        position_ids=position_ids,
                        valid_lengths=valid_lengths,
                        prompt_lengths=prompt_lengths,
                        num_samples=1,
                        rollout_strategy=terminal_rollout_strategy,
                        rollout_confidence_threshold=rollout_confidence_threshold,
                        sample_temperature=rollout_temperature,
                        sample_top_k=rollout_top_k,
                        sample_top_p=rollout_top_p,
                        enable_remask_actions=False,
                    )
                    baseline_terminal_reward = self._compute_gap_terminal_answer_rewards(
                        clean_input_ids=clean_input_ids,
                        terminal_input_ids=baseline_terminal,
                        target_scope_mask=target_scope_mask,
                        pred_scope_mask=baseline_eval_target_mask,
                    )
                    baseline_format_reward = torch.zeros_like(baseline_terminal_reward)
                    if format_weight > 0.0:
                        for row_idx in range(batch_size):
                            pred_text = self._decode_gap_response_tokens(baseline_terminal[row_idx][baseline_eval_target_mask[row_idx]])
                            baseline_format_reward[row_idx] = 1.0 if self._extract_gap_boxed_answer(pred_text) else 0.0
                else:
                    baseline_terminal = None
                    baseline_eval_target_mask = None
                    baseline_terminal_reward = None
                    baseline_format_reward = None
                    baseline_stop_hit = None
                    baseline_length_cap_hit = None

                (
                    sampled_terminal,
                    sampled_eval_target_mask,
                    logprob_per_batch,
                    entropy_per_batch,
                    remask_rate,
                    sampled_stop_hit,
                    sampled_length_cap_hit,
                    sampled_full,
                ) = (
                    self._rollout_gap_group_eval_with_sampled_remask(
                        noisy_input_ids=gap_outputs.z_accept,
                        labels=labels,
                        position_ids=position_ids,
                        valid_lengths=valid_lengths,
                        prompt_lengths=prompt_lengths,
                        num_samples=num_samples,
                        rollout_strategy=terminal_rollout_strategy,
                        rollout_confidence_threshold=rollout_confidence_threshold,
                        sample_temperature=rollout_temperature,
                        sample_top_k=rollout_top_k,
                        sample_top_p=rollout_top_p,
                        enable_remask_actions=True,
                    )
                )
                clean_flat = clean_input_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
                target_scope_flat = target_scope_mask.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
                sampled_terminal_reward = self._compute_gap_terminal_answer_rewards(
                    clean_input_ids=clean_flat,
                    terminal_input_ids=sampled_terminal,
                    target_scope_mask=target_scope_flat,
                    pred_scope_mask=sampled_eval_target_mask,
                ).view(num_samples, batch_size)
                sampled_format_reward = torch.zeros_like(sampled_terminal_reward)
                if format_weight > 0.0:
                    for branch_idx in range(num_samples):
                        for row_idx in range(batch_size):
                            flat_idx = branch_idx * batch_size + row_idx
                            pred_text = self._decode_gap_response_tokens(sampled_terminal[flat_idx][sampled_eval_target_mask[flat_idx]])
                            sampled_format_reward[branch_idx, row_idx] = 1.0 if self._extract_gap_boxed_answer(pred_text) else 0.0

            rewards = terminal_weight * sampled_terminal_reward + format_weight * sampled_format_reward - remask_penalty * remask_rate
            if use_baseline_branch:
                assert baseline_terminal_reward is not None and baseline_format_reward is not None
                baseline_reward = terminal_weight * baseline_terminal_reward + format_weight * baseline_format_reward
            else:
                baseline_reward = None
            kl_coef = float(getattr(self.config, "gap_grpo_kl_coef", 0.0) or 0.0)
            reference_model = self._get_gap_reference_model() if kl_coef > 0.0 else None
            if reference_model is not None:
                position_flat = position_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
                position_flat = self._extend_gap_position_ids_to_length(position_flat, sampled_terminal.shape[1])
                kl_chunk_size = 1
                actor_logprob = self._compute_gap_sequence_logprob_means(
                    model=self,
                    input_ids=sampled_terminal,
                    position_ids=position_flat,
                    target_scope_mask=sampled_eval_target_mask,
                    require_grad=False,
                    chunk_size=kl_chunk_size,
                ).view(num_samples, batch_size)
                reference_logprob = self._compute_gap_sequence_logprob_means(
                    model=reference_model,
                    input_ids=sampled_terminal,
                    position_ids=position_flat,
                    target_scope_mask=sampled_eval_target_mask,
                    require_grad=False,
                    chunk_size=kl_chunk_size,
                ).view(num_samples, batch_size)
                sampled_kl = (actor_logprob - reference_logprob).clamp_min(0.0)
                rewards = rewards - kl_coef * sampled_kl.detach()
            else:
                actor_logprob = None
                reference_logprob = None
                sampled_kl = None
            reward_gain = rewards - baseline_reward.unsqueeze(0) if baseline_reward is not None else rewards
            advantages, reward_std = self._compute_gap_group_advantages(
                rewards=rewards,
                baseline_reward=baseline_reward,
                advantage_eps=advantage_eps,
            )
            current_logprob_per_batch, current_entropy_per_batch = self._compute_gap_current_remask_action_logprob(
                remask_logits=remask_logits,
                masked_indices=masked_indices,
                full_candidate_mask=gap_outputs.full_candidate_mask,
                sampled_full=sampled_full,
                sample_prob_eps=float(getattr(self.config, "gap_grpo_sample_prob_eps", 1e-4) or 1e-4),
            )
            valid_advantage_mask = advantages.ne(0.0)
            if valid_advantage_mask.any():
                policy_terms = -(advantages.detach() * current_logprob_per_batch)
                policy_loss = policy_terms[valid_advantage_mask].mean()
                entropy_loss = current_entropy_per_batch.mean()
            else:
                policy_loss = current_logprob_per_batch.sum() * 0.0
                entropy_loss = current_entropy_per_batch.sum() * 0.0
            entropy_bonus = current_entropy_per_batch.mean()
            total_loss = grpo_weight * (policy_loss - entropy_coef * entropy_loss)
            total_loss = total_loss + remask_logits.sum() * 0.0 + self.gap_value_head(grpo_hidden_states[:1]).sum() * 0.0
            if use_baseline_branch:
                assert baseline_eval_target_mask is not None and baseline_terminal is not None
                baseline_target_counts = baseline_eval_target_mask.to(torch.float32).sum(dim=-1).clamp_min(1.0)
                baseline_remaining_mask_rate = (
                    baseline_terminal.eq(self.config.mask_token_id) & baseline_eval_target_mask
                ).to(torch.float32).sum(dim=-1) / baseline_target_counts
            else:
                baseline_remaining_mask_rate = None
            sampled_target_counts = sampled_eval_target_mask.to(torch.float32).sum(dim=-1).clamp_min(1.0)
            sampled_remaining_mask_rate = (
                sampled_terminal.eq(self.config.mask_token_id) & sampled_eval_target_mask
            ).to(torch.float32).sum(dim=-1).view(num_samples, batch_size) / sampled_target_counts.view(num_samples, batch_size)

            action_candidate_mask = sampled_full.any(dim=0)
            self._maybe_capture_gap_branch_debug(
                clean_input_ids=clean_input_ids,
                target_scope_mask=target_scope_mask,
                shared_state_input_ids=gap_outputs.z_accept,
                baseline_terminal=baseline_terminal,
                baseline_terminal_reward=baseline_terminal_reward,
                baseline_reward=baseline_reward,
                sampled_terminal=sampled_terminal,
                target_scope_flat=target_scope_flat,
                sampled_terminal_reward=sampled_terminal_reward,
                rewards=rewards,
                reward_gain=reward_gain,
                remask_rate=remask_rate,
                baseline_remaining_mask_rate=baseline_remaining_mask_rate,
                sampled_remaining_mask_rate=sampled_remaining_mask_rate,
                sampled_full=sampled_full,
                full_candidate_mask=action_candidate_mask,
                baseline_debug_mask=baseline_eval_target_mask,
                sampled_debug_mask=sampled_eval_target_mask,
                baseline_stop_hit=baseline_stop_hit.view(batch_size) if baseline_stop_hit is not None else None,
                sampled_stop_hit=sampled_stop_hit,
                baseline_length_cap_hit=baseline_length_cap_hit.view(batch_size) if baseline_length_cap_hit is not None else None,
                sampled_length_cap_hit=sampled_length_cap_hit,
            )

            gain_positive_rate = (reward_gain > 0).to(torch.float32).mean()
            gain_negative_rate = (reward_gain < 0).to(torch.float32).mean()
            branch_positive_rate = (sampled_terminal_reward > 0).to(torch.float32).mean()
            zero_metric = rewards.new_tensor(0.0)
            metrics = {
                "grpo_reward": rewards.mean().detach(),
                "grpo_terminal_reward": sampled_terminal_reward.mean().detach(),
                "grpo_reward_gain": reward_gain.mean().detach(),
                "grpo_reward_std": reward_std.mean().detach(),
                "grpo_group_advantage_abs": advantages.abs().mean().detach(),
                "grpo_value": grpo_hidden_states.new_tensor(0.0),
                "grpo_value_advantage_abs": reward_gain.abs().mean().detach(),
                "grpo_value_loss": reward_gain.pow(2).mean().detach(),
                "grpo_entropy": entropy_bonus.detach(),
                "grpo_policy_active_rate": valid_advantage_mask.to(torch.float32).mean().detach(),
                "grpo_rollout_logprob": logprob_per_batch.mean().detach(),
                "grpo_current_logprob": current_logprob_per_batch.mean().detach(),
                "grpo_branch_correct_rate": sampled_terminal_reward.mean().detach(),
                "grpo_branch_positive_rate": branch_positive_rate.detach(),
                "grpo_branch_negative_rate": (1.0 - branch_positive_rate).detach(),
                "grpo_baseline_correct_rate": ((baseline_terminal_reward > 0).to(torch.float32).mean().detach() if baseline_terminal_reward is not None else zero_metric),
                "grpo_baseline_reward": (baseline_reward.mean().detach() if baseline_reward is not None else zero_metric),
                "grpo_baseline_remaining_mask_rate": (baseline_remaining_mask_rate.mean().detach() if baseline_remaining_mask_rate is not None else zero_metric),
                "grpo_sampled_remaining_mask_rate": sampled_remaining_mask_rate.mean().detach(),
                "grpo_sampled_remaining_mask_rate_max": sampled_remaining_mask_rate.max().detach(),
                "grpo_gain_positive_rate": gain_positive_rate.detach(),
                "grpo_gain_negative_rate": gain_negative_rate.detach(),
                "grpo_gain_tie_rate": (1.0 - gain_positive_rate - gain_negative_rate).detach(),
                "grpo_any_win_rate": (reward_gain > 0).any(dim=0).to(torch.float32).mean().detach(),
                "grpo_any_lose_rate": (reward_gain < 0).any(dim=0).to(torch.float32).mean().detach(),
                "grpo_loss": total_loss.detach(),
            }
            if sampled_kl is not None:
                metrics["grpo_kl"] = sampled_kl.mean().detach()
                metrics["grpo_actor_logp"] = actor_logprob.mean().detach()
                metrics["grpo_ref_logp"] = reference_logprob.mean().detach()
            return total_loss, metrics

        full_candidate_mask = gap_outputs.full_candidate_mask
        candidate_mask_flat = full_candidate_mask[masked_indices]
        candidate_count = int(full_candidate_mask.sum().item())
        if candidate_count <= 0:
            zero = remask_logits.sum() * 0.0 + self.gap_value_head(grpo_hidden_states[:1]).sum() * 0.0
            return zero, {}

        timing_interval = _gap_env_int("SDAR_GAP_GRPO_TIMING_INTERVAL", 0)
        debug_step = int(getattr(self, "_gap_debug_global_step", -1))
        should_log_timing = timing_interval > 0 and _gap_is_rank0() and debug_step >= 0 and (debug_step % timing_interval == 0)
        timing_marks = {}

        def _mark(name: str) -> None:
            if should_log_timing:
                timing_marks[name] = time.perf_counter()

        sample_prob_eps = float(getattr(self.config, "gap_grpo_sample_prob_eps", 1e-4) or 1e-4)
        entropy_coef = float(getattr(self.config, "gap_grpo_entropy_coef", 0.0) or 0.0)
        terminal_weight = float(getattr(self.config, "gap_grpo_terminal_reward_weight", 1.0) or 1.0)
        format_weight = float(getattr(self.config, "gap_grpo_format_reward_weight", 0.0) or 0.0)
        remask_penalty = float(getattr(self.config, "gap_grpo_remask_penalty", 0.0) or 0.0)
        advantage_eps = float(getattr(self.config, "gap_grpo_advantage_eps", 1e-4) or 1e-4)
        value_loss_weight = float(getattr(self.config, "gap_grpo_value_loss_weight", 0.0) or 0.0)
        value_baseline_weight = float(getattr(self.config, "gap_grpo_value_baseline_weight", 0.0) or 0.0)
        use_baseline_branch = bool(getattr(self.config, "gap_grpo_use_baseline_branch", True))
        rollout_temperature = float(getattr(self.config, "gap_grpo_rollout_temperature", 0.0) or 0.0)
        rollout_top_k = int(getattr(self.config, "gap_grpo_rollout_top_k", 0) or 0)
        rollout_top_p = float(getattr(self.config, "gap_grpo_rollout_top_p", 1.0) or 1.0)
        kl_coef = float(getattr(self.config, "gap_grpo_kl_coef", 0.0) or 0.0)
        reference_model = self._get_gap_reference_model() if kl_coef > 0.0 else None
        terminal_rollout_strategy = (
            getattr(self.config, "gap_grpo_terminal_rollout_strategy", None)
            or rollout_strategy
        )
        terminal_rollout_scope = (
            getattr(self.config, "gap_grpo_terminal_rollout_scope", None)
            or getattr(self.config, "gap_rollout_scope", "all")
        )

        flat_probs = torch.sigmoid(remask_logits[candidate_mask_flat]).clamp(min=sample_prob_eps, max=1.0 - sample_prob_eps)
        batch_ids = torch.nonzero(full_candidate_mask, as_tuple=False)[:, 0]
        batch_size = full_candidate_mask.shape[0]
        candidate_counts_per_batch = torch.zeros(batch_size, dtype=torch.float32, device=flat_probs.device)
        candidate_counts_per_batch.scatter_add_(
            0,
            batch_ids,
            torch.ones_like(batch_ids, dtype=torch.float32),
        )
        candidate_counts_per_batch = candidate_counts_per_batch.clamp_min(1.0)
        flat_values = self.gap_value_head(grpo_hidden_states[candidate_mask_flat]).squeeze(-1).float()
        value_per_batch = torch.zeros(batch_size, dtype=flat_values.dtype, device=flat_values.device)
        value_per_batch.scatter_add_(0, batch_ids, flat_values)
        value_per_batch = value_per_batch / candidate_counts_per_batch

        if use_baseline_branch:
            _mark("baseline_start")
            with torch.no_grad():
                baseline_terminal = self._rollout_gap_state_to_terminal(
                    clean_input_ids=clean_input_ids,
                    noisy_input_ids=gap_outputs.z_accept,
                    labels=labels,
                    position_ids=position_ids,
                    num_tokens=num_tokens,
                    start_stage=0,
                    rollout_strategy=terminal_rollout_strategy,
                    rollout_confidence_threshold=rollout_confidence_threshold,
                    rollout_scope=terminal_rollout_scope,
                    sample_temperature=rollout_temperature,
                    sample_top_k=rollout_top_k,
                    sample_top_p=rollout_top_p,
                )
                baseline_terminal_reward = self._compute_gap_terminal_answer_rewards(
                    clean_input_ids=clean_input_ids,
                    terminal_input_ids=baseline_terminal,
                    target_scope_mask=target_scope_mask,
                )
                if format_weight > 0.0:
                    baseline_format_reward = torch.zeros_like(baseline_terminal_reward)
                    for row_idx in range(batch_size):
                        pred_text = self._decode_gap_response_tokens(baseline_terminal[row_idx][target_scope_mask[row_idx]])
                        baseline_format_reward[row_idx] = 1.0 if self._extract_gap_boxed_answer(pred_text) else 0.0
                else:
                    baseline_format_reward = torch.zeros_like(baseline_terminal_reward)
            _mark("baseline_end")
        else:
            baseline_terminal = None
            baseline_terminal_reward = None
            baseline_format_reward = None

        sampled_flat = torch.bernoulli(flat_probs.unsqueeze(0).expand(num_samples, -1)).to(dtype=torch.bool)
        sampled_full = torch.zeros(
            (num_samples, batch_size, full_candidate_mask.shape[1]),
            dtype=torch.bool,
            device=full_candidate_mask.device,
        )
        sampled_full[:, full_candidate_mask] = sampled_flat

        sample_logprob = (
            sampled_flat.to(flat_probs.dtype) * flat_probs.log().unsqueeze(0)
            + (~sampled_flat).to(flat_probs.dtype) * (1.0 - flat_probs).log().unsqueeze(0)
        )
        sample_entropy = -(
            flat_probs * flat_probs.log() + (1.0 - flat_probs) * (1.0 - flat_probs).log()
        )
        logprob_per_batch = torch.zeros((num_samples, batch_size), dtype=sample_logprob.dtype, device=sample_logprob.device)
        entropy_per_batch = torch.zeros(batch_size, dtype=sample_entropy.dtype, device=sample_entropy.device)
        logprob_per_batch.scatter_add_(1, batch_ids.unsqueeze(0).expand(num_samples, -1), sample_logprob)
        entropy_per_batch.scatter_add_(0, batch_ids, sample_entropy)
        logprob_per_batch = logprob_per_batch / candidate_counts_per_batch.unsqueeze(0)
        entropy_per_batch = entropy_per_batch / candidate_counts_per_batch

        _mark("sampled_start")
        with torch.no_grad():
            sampled_states = gap_outputs.z_accept.unsqueeze(0).repeat(num_samples, 1, 1)
            sampled_states[sampled_full] = self.config.mask_token_id
            sampled_states_flat = sampled_states.view(num_samples * batch_size, -1)
            clean_flat = clean_input_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
            labels_flat = labels.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
            position_flat = position_ids.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
            target_scope_flat = target_scope_mask.unsqueeze(0).repeat(num_samples, 1, 1).view(num_samples * batch_size, -1)
            num_tokens_flat = [num_tokens[row_idx % batch_size] for row_idx in range(num_samples * batch_size)]

            sampled_terminal = self._rollout_gap_state_to_terminal(
                clean_input_ids=clean_flat,
                noisy_input_ids=sampled_states_flat,
                labels=labels_flat,
                position_ids=position_flat,
                num_tokens=num_tokens_flat,
                start_stage=0,
                rollout_strategy=terminal_rollout_strategy,
                rollout_confidence_threshold=rollout_confidence_threshold,
                rollout_scope=terminal_rollout_scope,
                sample_temperature=rollout_temperature,
                sample_top_k=rollout_top_k,
                sample_top_p=rollout_top_p,
            )
            sampled_terminal_reward = self._compute_gap_terminal_answer_rewards(
                clean_input_ids=clean_flat,
                terminal_input_ids=sampled_terminal,
                target_scope_mask=target_scope_flat,
            ).view(num_samples, batch_size)

            if format_weight > 0.0:
                sampled_format_reward = torch.zeros_like(sampled_terminal_reward)
                for branch_idx in range(num_samples):
                    for row_idx in range(batch_size):
                        flat_idx = branch_idx * batch_size + row_idx
                        pred_text = self._decode_gap_response_tokens(sampled_terminal[flat_idx][target_scope_flat[flat_idx]])
                        sampled_format_reward[branch_idx, row_idx] = 1.0 if self._extract_gap_boxed_answer(pred_text) else 0.0
            else:
                sampled_format_reward = torch.zeros_like(sampled_terminal_reward)
        _mark("sampled_end")

        if use_baseline_branch:
            assert baseline_terminal is not None
            baseline_target_counts = target_scope_mask.to(torch.float32).sum(dim=-1).clamp_min(1.0)
            baseline_remaining_mask_rate = (
                baseline_terminal.eq(self.config.mask_token_id) & target_scope_mask
            ).to(torch.float32).sum(dim=-1) / baseline_target_counts
        else:
            baseline_remaining_mask_rate = None
        sampled_target_counts = target_scope_flat.to(torch.float32).sum(dim=-1).clamp_min(1.0)
        sampled_remaining_mask_rate = (
            sampled_terminal.eq(self.config.mask_token_id) & target_scope_flat
        ).to(torch.float32).sum(dim=-1).view(num_samples, batch_size) / sampled_target_counts.view(num_samples, batch_size)

        _mark("reward_start")
        remask_rate = sampled_full.to(torch.float32).sum(dim=-1) / candidate_counts_per_batch.unsqueeze(0)
        rewards = (
            terminal_weight * sampled_terminal_reward
            + format_weight * sampled_format_reward
            - remask_penalty * remask_rate
        )
        if use_baseline_branch:
            assert baseline_terminal_reward is not None and baseline_format_reward is not None
            baseline_reward = (
                terminal_weight * baseline_terminal_reward
                + format_weight * baseline_format_reward
            )
        else:
            baseline_reward = None

        if reference_model is not None:
            kl_chunk_size = 1
            actor_logprob = self._compute_gap_sequence_logprob_means(
                model=self,
                input_ids=sampled_terminal,
                position_ids=position_flat,
                target_scope_mask=target_scope_flat,
                require_grad=False,
                chunk_size=kl_chunk_size,
            ).view(num_samples, batch_size)
            reference_logprob = self._compute_gap_sequence_logprob_means(
                model=reference_model,
                input_ids=sampled_terminal,
                position_ids=position_flat,
                target_scope_mask=target_scope_flat,
                require_grad=False,
                chunk_size=kl_chunk_size,
            ).view(num_samples, batch_size)
            sampled_kl = (actor_logprob - reference_logprob).clamp_min(0.0)
            rewards = rewards - kl_coef * sampled_kl.detach()
        else:
            actor_logprob = None
            reference_logprob = None
            sampled_kl = None

        reward_gain = rewards - baseline_reward.unsqueeze(0) if baseline_reward is not None else rewards
        _mark("reward_end")
        group_advantages, reward_std = self._compute_gap_group_advantages(
            rewards=rewards,
            baseline_reward=baseline_reward,
            advantage_eps=advantage_eps,
        )
        critic_advantages = reward_gain - value_per_batch.detach().unsqueeze(0)
        if value_baseline_weight > 0.0:
            blend_weight = max(0.0, min(1.0, value_baseline_weight))
            advantages = (1.0 - blend_weight) * group_advantages + blend_weight * critic_advantages
        else:
            advantages = group_advantages
        if num_samples == 1:
            advantages = critic_advantages if value_baseline_weight > 0.0 else reward_gain

        self._maybe_capture_gap_branch_debug(
            clean_input_ids=clean_input_ids,
            target_scope_mask=target_scope_mask,
            shared_state_input_ids=gap_outputs.z_accept,
            baseline_terminal=baseline_terminal,
            baseline_terminal_reward=baseline_terminal_reward,
            baseline_reward=baseline_reward,
            sampled_terminal=sampled_terminal,
            target_scope_flat=target_scope_flat,
            sampled_terminal_reward=sampled_terminal_reward,
            rewards=rewards,
            reward_gain=reward_gain,
            remask_rate=remask_rate,
            baseline_remaining_mask_rate=baseline_remaining_mask_rate,
            sampled_remaining_mask_rate=sampled_remaining_mask_rate,
            sampled_full=sampled_full,
            full_candidate_mask=full_candidate_mask,
        )

        gain_positive_rate = (reward_gain > 0).to(torch.float32).mean()
        gain_negative_rate = (reward_gain < 0).to(torch.float32).mean()
        gain_tie_rate = 1.0 - gain_positive_rate - gain_negative_rate
        branch_positive_rate = (sampled_terminal_reward > 0).to(torch.float32).mean()
        branch_negative_rate = 1.0 - branch_positive_rate
        baseline_correct_rate = (baseline_terminal_reward > 0).to(torch.float32).mean() if baseline_terminal_reward is not None else rewards.new_tensor(0.0)
        any_win_rate = (reward_gain > 0).any(dim=0).to(torch.float32).mean()
        any_lose_rate = (reward_gain < 0).any(dim=0).to(torch.float32).mean()

        valid_advantage_mask = advantages.ne(0.0)
        if valid_advantage_mask.any():
            policy_terms = -(advantages.detach() * logprob_per_batch)
            policy_loss = policy_terms[valid_advantage_mask].mean()
            entropy_loss = entropy_per_batch.mean()
        else:
            policy_loss = logprob_per_batch.sum() * 0.0
            entropy_loss = entropy_per_batch.sum() * 0.0
        entropy_bonus = entropy_per_batch.mean()
        value_loss = ((value_per_batch.unsqueeze(0) - reward_gain.detach()) ** 2).mean()
        total_loss = grpo_weight * (policy_loss - entropy_coef * entropy_loss)
        if value_loss_weight > 0.0:
            total_loss = total_loss + value_loss_weight * value_loss
        else:
            total_loss = total_loss + value_loss * 0.0
        metrics = {
            "grpo_reward": rewards.mean().detach(),
            "grpo_terminal_reward": sampled_terminal_reward.mean().detach(),
            "grpo_reward_gain": reward_gain.mean().detach(),
            "grpo_reward_std": reward_std.mean().detach(),
            "grpo_group_advantage_abs": group_advantages.abs().mean().detach(),
            "grpo_value": value_per_batch.mean().detach(),
            "grpo_value_advantage_abs": critic_advantages.abs().mean().detach(),
            "grpo_value_loss": value_loss.detach(),
            "grpo_entropy": entropy_bonus.detach(),
            "grpo_policy_active_rate": valid_advantage_mask.to(torch.float32).mean().detach(),
            "grpo_branch_correct_rate": sampled_terminal_reward.mean().detach(),
            "grpo_branch_positive_rate": branch_positive_rate.detach(),
            "grpo_branch_negative_rate": branch_negative_rate.detach(),
            "grpo_baseline_correct_rate": baseline_correct_rate.detach(),
            "grpo_baseline_reward": baseline_reward.mean().detach() if baseline_reward is not None else rewards.new_tensor(0.0),
            "grpo_baseline_remaining_mask_rate": baseline_remaining_mask_rate.mean().detach() if baseline_remaining_mask_rate is not None else rewards.new_tensor(0.0),
            "grpo_sampled_remaining_mask_rate": sampled_remaining_mask_rate.mean().detach(),
            "grpo_sampled_remaining_mask_rate_max": sampled_remaining_mask_rate.max().detach(),
            "grpo_gain_positive_rate": gain_positive_rate.detach(),
            "grpo_gain_negative_rate": gain_negative_rate.detach(),
            "grpo_gain_tie_rate": gain_tie_rate.detach(),
            "grpo_any_win_rate": any_win_rate.detach(),
            "grpo_any_lose_rate": any_lose_rate.detach(),
            "grpo_loss": total_loss.detach(),
        }
        if sampled_kl is not None:
            metrics["grpo_kl"] = sampled_kl.mean().detach()
            metrics["grpo_actor_logp"] = actor_logprob.mean().detach()
            metrics["grpo_ref_logp"] = reference_logprob.mean().detach()
        if should_log_timing:
            baseline_sec = max(0.0, timing_marks.get("baseline_end", 0.0) - timing_marks.get("baseline_start", 0.0))
            sampled_sec = max(0.0, timing_marks.get("sampled_end", 0.0) - timing_marks.get("sampled_start", 0.0))
            reward_sec = max(0.0, timing_marks.get("reward_end", 0.0) - timing_marks.get("reward_start", 0.0))
            logger.info(
                "[GAP grpo timing] step=%s batch=%s num_samples=%s candidate_tokens=%s baseline_sec=%.2f sampled_sec=%.2f reward_sec=%.2f",
                debug_step,
                batch_size,
                num_samples,
                candidate_count,
                baseline_sec,
                sampled_sec,
                reward_sec,
            )
        return total_loss, metrics

    def _compute_gap_sft_ce_anchor(
        self,
        clean_input_ids: torch.LongTensor,
        clean_labels: torch.LongTensor,
        clean_position_ids: torch.LongTensor,
    ) -> torch.Tensor:
        clean_position_ids = modify_padded_position_ids_2d(clean_position_ids)
        target_mask = clean_labels.ne(-100)
        num_tokens = calculate_token_nums(clean_position_ids)
        concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, _, logits_to_keep, _ = self.build_bd_training_inputs(
            inputs_ids=clean_input_ids,
            noisy_inputs_ids=clean_input_ids,
            position_ids=clean_position_ids,
            logits_to_keep_half=target_mask,
            num_tokens=num_tokens,
        )
        outputs = self.model(
            input_ids=concat_inputs_ids,
            attention_mask=flex_attention_mask_3d,
            position_ids=concat_position_ids,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=True,
        )
        hidden_states = outputs.last_hidden_state[logits_to_keep].contiguous()
        logits = self.lm_head(hidden_states).float()
        target_ids = clean_input_ids[target_mask]
        return nn.functional.cross_entropy(logits, target_ids, reduction="mean")

    def set_puma_streaming_context(self, slot_offset: int, buffer_size: int) -> None:
        self._puma_streaming_context = {
            "slot_offset": max(int(slot_offset), 0),
            "buffer_size": max(int(buffer_size), 1),
        }

    def reset_puma_streaming_state(self) -> None:
        self._puma_streaming_state = None
        self._puma_streaming_context = {"slot_offset": 0, "buffer_size": None}

    def _ensure_puma_streaming_state(self, buffer_size, inputs_ids, labels, position_ids):
        expected_shape = (buffer_size, inputs_ids.shape[1])
        state = self._puma_streaming_state
        if (
            state is None
            or tuple(state["clean_input_ids"].shape) != expected_shape
            or state["clean_input_ids"].device != inputs_ids.device
            or state["clean_input_ids"].dtype != inputs_ids.dtype
        ):
            state = {
                "clean_input_ids": torch.zeros(expected_shape, dtype=inputs_ids.dtype, device=inputs_ids.device),
                "labels": torch.full(expected_shape, -100, dtype=labels.dtype, device=labels.device),
                "position_ids": torch.zeros(expected_shape, dtype=position_ids.dtype, device=position_ids.device),
                "noisy_inputs_ids": torch.zeros(expected_shape, dtype=inputs_ids.dtype, device=inputs_ids.device),
                "stages": torch.full((buffer_size,), self.config.block_size, dtype=torch.long, device=inputs_ids.device),
                "max_progress": torch.full((buffer_size,), self.config.block_size, dtype=torch.long, device=inputs_ids.device),
                "active": torch.zeros(buffer_size, dtype=torch.bool, device=inputs_ids.device),
            }
            self._puma_streaming_state = state

        return state

    def build_bd_training_inputs(self, inputs_ids, noisy_inputs_ids, position_ids, logits_to_keep_half, num_tokens=None):
        bsz, seq_len = inputs_ids.shape
        if num_tokens is None:
            num_tokens = calculate_token_nums(position_ids)
        router_noisy_part_list = []
        for i in range(bsz):
            cur_router_noisy_part = (torch.arange(num_tokens[i].shape[0] *2) % 2 == 0).to(inputs_ids.device)
            cur_router_noisy_part = cur_router_noisy_part.repeat_interleave(num_tokens[i].repeat_interleave(2))
            router_noisy_part_list.append(cur_router_noisy_part)
        router_noisy_part = torch.stack(router_noisy_part_list, dim=0)

        # concated inputs_ids: (bzs, seq_len x 2)
        concat_inputs_ids = inputs_ids.repeat(1, 2)
        # concated logits_to_keep: (bsz, seq_len x 2)
        logits_to_keep = torch.zeros(
                    bsz, 2 * seq_len, dtype=torch.bool, device=inputs_ids.device)
        # concated position_ids: (bsz, seq_len x 2)
        concat_position_ids = torch.zeros(
                    bsz, 2 * seq_len, dtype=position_ids.dtype, device=position_ids.device)
        for i in range(bsz):
            concat_inputs_ids[i][router_noisy_part[i]] = noisy_inputs_ids[i]
            concat_inputs_ids[i][~router_noisy_part[i]] = inputs_ids[i]

            logits_to_keep[i][router_noisy_part[i]] = logits_to_keep_half[i]

            concat_position_ids[i][router_noisy_part[i]] = position_ids[i]
            concat_position_ids[i][~router_noisy_part[i]] = position_ids[i]

        # create flex_attention mask
        attention_mask = block_attn_mask(num_tokens, self.config.block_size, inputs_ids.device)
        flex_attention_mask_3d = create_block_mask(
                            lambda b, h, q_idx, kv_idx: attention_mask[b, q_idx, kv_idx],
                            B=attention_mask.size(0), H=None,
                            Q_LEN=attention_mask.size(1), KV_LEN=attention_mask.size(2),
        )

        return concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, logits_to_keep_half, logits_to_keep, num_tokens

    def prepare_for_bd_training(self, inputs_ids, position_ids, prompt_mask):
        num_tokens = calculate_token_nums(position_ids) # List[torch.Tensor]
        noisy_inputs_ids, logits_to_keep_half, p_mask = forward_add_noise_packed(
            inputs_ids=inputs_ids,
            num_tokens_list=num_tokens,
            prompt_mask=prompt_mask,
            mask_id=self.config.mask_token_id,
        )
        concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, logits_to_keep_half, logits_to_keep, _ = self.build_bd_training_inputs(
            inputs_ids=inputs_ids,
            noisy_inputs_ids=noisy_inputs_ids,
            position_ids=position_ids,
            logits_to_keep_half=logits_to_keep_half,
            num_tokens=num_tokens,
        )

        return {
            "concat_inputs_ids": concat_inputs_ids,
            "concat_position_ids": concat_position_ids,
            "flex_attention_mask_3d": flex_attention_mask_3d,
            "logits_to_keep_half": logits_to_keep_half,
            "logits_to_keep": logits_to_keep,
            "p_mask": p_mask,
            "noisy_inputs_ids": noisy_inputs_ids,
            "num_tokens": num_tokens,
        }

    def prepare_for_puma_streaming_training(self, inputs_ids, labels, position_ids, use_remask_aux: bool = False):
        batch_size = inputs_ids.shape[0]
        rollout_steps = int(getattr(self.config, "gap_rollout_steps", self.config.block_size))
        rollout_steps = max(1, rollout_steps)
        transfer_schedule = get_num_transfer_tokens(self.config.block_size, rollout_steps).to(inputs_ids.device)
        context = getattr(self, "_puma_streaming_context", {}) or {}
        buffer_size = int(context.get("buffer_size") or batch_size)
        slot_offset = int(context.get("slot_offset", 0))
        slot_indices = (torch.arange(batch_size, device=inputs_ids.device) + slot_offset) % buffer_size
        state = self._ensure_puma_streaming_state(buffer_size, inputs_ids, labels, position_ids)
        if self._should_use_gap_prefix_frontier_state():
            need_refill = (~state["active"][slot_indices]) | (state["stages"][slot_indices] >= state["max_progress"][slot_indices])
        else:
            need_refill = (~state["active"][slot_indices]) | (state["stages"][slot_indices] >= rollout_steps)
        if need_refill.any():
            refill_slots = slot_indices[need_refill]
            refill_inputs = inputs_ids[need_refill].detach().clone()
            refill_labels = labels[need_refill].detach().clone()
            refill_position_ids = position_ids[need_refill].detach().clone()
            refill_num_tokens = calculate_token_nums(refill_position_ids)
            refill_max_progress = self._compute_gap_prefix_progress_limits(
                labels=refill_labels,
                num_tokens=refill_num_tokens,
                block_size=self.config.block_size,
                rollout_steps=rollout_steps,
            )
            state["clean_input_ids"][refill_slots] = refill_inputs
            state["labels"][refill_slots] = refill_labels
            state["position_ids"][refill_slots] = refill_position_ids
            refill_progress = torch.zeros(refill_inputs.shape[0], dtype=torch.long, device=inputs_ids.device)
            refill_noisy = self._build_gap_prefix_teacher_forced_state(
                clean_input_ids=refill_inputs,
                labels=refill_labels,
                num_tokens=refill_num_tokens,
                progress_units=refill_progress,
            )
            state["noisy_inputs_ids"][refill_slots] = refill_noisy
            state["stages"][refill_slots] = 0
            state["max_progress"][refill_slots] = refill_max_progress
            state["active"][refill_slots] = True

        clean_input_ids = state["clean_input_ids"][slot_indices].clone()
        clean_labels = state["labels"][slot_indices].clone()
        clean_position_ids = state["position_ids"][slot_indices].clone()
        current_stages = state["stages"][slot_indices].clone()
        num_tokens = calculate_token_nums(clean_position_ids)
        if self._should_use_gap_prefix_frontier_state():
            noisy_inputs_ids = self._build_gap_prefix_teacher_forced_state(
                clean_input_ids=clean_input_ids,
                labels=clean_labels,
                num_tokens=num_tokens,
                progress_units=current_stages,
            )
        else:
            noisy_inputs_ids = state["noisy_inputs_ids"][slot_indices].clone()
        target_scope_mask = clean_labels.ne(-100)
        logits_to_keep_half = noisy_inputs_ids.eq(self.config.mask_token_id) & target_scope_mask
        if not logits_to_keep_half.any() and target_scope_mask.any():
            fallback_index = torch.nonzero(target_scope_mask, as_tuple=False)[0]
            noisy_inputs_ids[fallback_index[0], fallback_index[1]] = self.config.mask_token_id
            logits_to_keep_half[fallback_index[0], fallback_index[1]] = True

        p_mask = build_rollout_p_mask(
            masked_indices=logits_to_keep_half,
            labels=clean_labels,
            num_tokens=num_tokens,
            target_scope_mask=target_scope_mask,
            per_block=True,
            block_size=self.config.block_size,
        )
        concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, logits_to_keep_half, logits_to_keep, _ = self.build_bd_training_inputs(
            inputs_ids=clean_input_ids,
            noisy_inputs_ids=noisy_inputs_ids,
            position_ids=clean_position_ids,
            logits_to_keep_half=logits_to_keep_half,
            num_tokens=num_tokens,
        )

        return {
            "concat_inputs_ids": concat_inputs_ids,
            "concat_position_ids": concat_position_ids,
            "flex_attention_mask_3d": flex_attention_mask_3d,
            "logits_to_keep_half": logits_to_keep_half,
            "logits_to_keep": logits_to_keep,
            "p_mask": p_mask,
            "noisy_inputs_ids": noisy_inputs_ids,
            "num_tokens": num_tokens,
            "target_scope_mask": target_scope_mask,
            "loss_target_count": target_scope_mask.sum().clamp_min(1),
            "gap_training_mode": "puma",
            "use_remask_aux": use_remask_aux,
            "rollout_depth": current_stages.float() / rollout_steps,
            "rollout_progress_units": current_stages,
            "next_transfer_tokens": transfer_schedule[current_stages.clamp_max(rollout_steps - 1)],
            "rollout_strategy": getattr(self.config, "gap_rollout_strategy", "low_confidence_dynamic"),
            "rollout_confidence_threshold": float(getattr(self.config, "gap_rollout_confidence_threshold", 0.95)),
            "terminal_rollout_strategy": getattr(self.config, "gap_grpo_terminal_rollout_strategy", None)
            or getattr(self.config, "gap_rollout_strategy", "low_confidence_dynamic"),
            "terminal_rollout_scope": getattr(self.config, "gap_grpo_terminal_rollout_scope", None)
            or getattr(self.config, "gap_rollout_scope", "all"),
            "clean_input_ids": clean_input_ids,
            "clean_labels": clean_labels,
            "clean_position_ids": clean_position_ids,
            "streaming_slot_indices": slot_indices,
            "streaming_refills": need_refill.sum(),
            "max_rollout_progress": self._compute_gap_prefix_progress_limits(
                labels=clean_labels,
                num_tokens=num_tokens,
                block_size=self.config.block_size,
                rollout_steps=rollout_steps,
            ) if self._should_use_gap_prefix_frontier_state() else state["max_progress"][slot_indices].clone(),
        }

    def advance_puma_streaming_state(self, training_batch, proposal_scores_full):
        state = self._puma_streaming_state
        if state is None:
            return

        slot_indices = training_batch.get("streaming_slot_indices")
        if slot_indices is None:
            return

        rollout_steps = int(getattr(self.config, "gap_rollout_steps", self.config.block_size))
        rollout_steps = max(1, rollout_steps)
        if self._should_use_gap_prefix_frontier_state():
            next_stages = training_batch["rollout_progress_units"].to(torch.long) + 1
            max_progress = training_batch.get("max_rollout_progress")
            if max_progress is None:
                max_progress = torch.full_like(next_stages, rollout_steps)
            finished_mask = next_stages >= max_progress
            next_stages = torch.where(finished_mask, max_progress, next_stages)
            state["stages"][slot_indices] = next_stages
            state["active"][slot_indices] = True
            return
        reveal_mask = select_teacher_forced_rollout_tokens(
            masked_indices=training_batch["logits_to_keep_half"],
            proposal_scores_full=proposal_scores_full,
            num_tokens=training_batch["num_tokens"],
            block_size=self.config.block_size,
            num_transfer_tokens=training_batch["next_transfer_tokens"],
            strategy=training_batch["rollout_strategy"],
            confidence_threshold=training_batch["rollout_confidence_threshold"],
            scope=getattr(self.config, "gap_rollout_scope", "all"),
        )
        next_noisy_inputs_ids = training_batch["noisy_inputs_ids"].clone()
        next_noisy_inputs_ids[reveal_mask] = training_batch["clean_input_ids"][reveal_mask]
        next_stages = training_batch["rollout_depth"].to(torch.long) + 1
        finished_mask = next_stages >= rollout_steps
        finished_mask |= ~(next_noisy_inputs_ids.eq(self.config.mask_token_id) & training_batch["clean_labels"].ne(-100)).any(dim=1)
        next_stages = torch.where(
            finished_mask,
            torch.full_like(next_stages, rollout_steps),
            next_stages,
        )

        state["noisy_inputs_ids"][slot_indices] = next_noisy_inputs_ids.detach()
        state["stages"][slot_indices] = next_stages
        state["active"][slot_indices] = True

    def prepare_for_teacher_forced_rollout_training(
        self,
        inputs_ids,
        labels,
        position_ids,
        prompt_mask,
        output_attentions,
        output_hidden_states,
        cache_position,
        **kwargs,
    ):
        num_tokens = calculate_token_nums(position_ids)
        answer_mask = ~prompt_mask
        gap_training_mode = getattr(self.config, "gap_training_mode", "remask")
        if gap_training_mode not in {"puma", "remask"}:
            raise ValueError(f"Unsupported GAP training mode: {gap_training_mode}")

        use_remask_aux = gap_training_mode == "remask"
        if getattr(self.config, "gap_puma_streaming", True):
            return self.prepare_for_puma_streaming_training(
                inputs_ids,
                labels,
                position_ids,
                use_remask_aux=use_remask_aux,
            )

        noisy_inputs_ids = torch.where(
            answer_mask,
            torch.full_like(inputs_ids, self.config.mask_token_id),
            inputs_ids,
        )

        rollout_steps = int(getattr(self.config, "gap_rollout_steps", self.config.block_size))
        rollout_steps = max(1, rollout_steps)
        transfer_schedule = get_num_transfer_tokens(self.config.block_size, rollout_steps).tolist()
        rollout_depth = int(torch.randint(0, rollout_steps, (1,), device=inputs_ids.device).item())
        rollout_strategy = getattr(self.config, "gap_rollout_strategy", "low_confidence_dynamic")
        rollout_confidence_threshold = float(getattr(self.config, "gap_rollout_confidence_threshold", 0.95))
        default_rollout_scope = "all"
        rollout_scope = getattr(self.config, "gap_rollout_scope", default_rollout_scope)
        completed_steps = 0

        if self._should_use_gap_prefix_frontier_state():
            max_progress = self._compute_gap_prefix_progress_limits(
                labels=labels,
                num_tokens=num_tokens,
                block_size=self.config.block_size,
                rollout_steps=rollout_steps,
            )
            rollout_progress = torch.floor(
                torch.rand(inputs_ids.shape[0], device=inputs_ids.device) * max_progress.to(dtype=torch.float32)
            ).to(dtype=torch.long)
            noisy_inputs_ids = self._build_gap_prefix_teacher_forced_state(
                clean_input_ids=inputs_ids,
                labels=labels,
                num_tokens=num_tokens,
                progress_units=rollout_progress,
            )
            completed_steps = int(rollout_progress.float().mean().item())
        else:
            rollout_depth = int(torch.randint(0, rollout_steps, (1,), device=inputs_ids.device).item())

            for step_idx in range(rollout_depth):
                masked_indices = noisy_inputs_ids.eq(self.config.mask_token_id) & answer_mask
                if not masked_indices.any():
                    break

                concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, logits_to_keep_half, logits_to_keep, _ = self.build_bd_training_inputs(
                    inputs_ids=inputs_ids,
                    noisy_inputs_ids=noisy_inputs_ids,
                    position_ids=position_ids,
                    logits_to_keep_half=masked_indices,
                    num_tokens=num_tokens,
                )
                outputs = self.model(
                    input_ids=concat_inputs_ids,
                    attention_mask=flex_attention_mask_3d,
                    position_ids=concat_position_ids,
                    output_attentions=output_attentions,
                    output_hidden_states=output_hidden_states,
                    return_dict=True,
                    cache_position=cache_position,
                    **kwargs,
                )
                hidden_states = outputs.last_hidden_state[logits_to_keep].contiguous()
                proposal_logits = self.lm_head(hidden_states).float()
                proposal_ids = proposal_logits.argmax(dim=-1)
                proposal_scores = (
                    proposal_logits.gather(-1, proposal_ids.unsqueeze(-1)).squeeze(-1)
                    - torch.logsumexp(proposal_logits, dim=-1)
                ).exp()
                proposal_scores_full = torch.full(
                    noisy_inputs_ids.shape,
                    float("-inf"),
                    dtype=proposal_scores.dtype,
                    device=proposal_scores.device,
                )
                proposal_scores_full[masked_indices] = proposal_scores
                reveal_mask = select_teacher_forced_rollout_tokens(
                    masked_indices=masked_indices,
                    proposal_scores_full=proposal_scores_full,
                    num_tokens=num_tokens,
                    block_size=self.config.block_size,
                    num_transfer_tokens=int(transfer_schedule[step_idx]),
                    strategy=rollout_strategy,
                    confidence_threshold=rollout_confidence_threshold,
                    scope=rollout_scope,
                )
                if not reveal_mask.any():
                    break
                noisy_inputs_ids[reveal_mask] = inputs_ids[reveal_mask]
                completed_steps += 1

        remaining_masked = noisy_inputs_ids.eq(self.config.mask_token_id) & answer_mask
        target_scope_mask = labels.ne(-100)
        logits_to_keep_half = remaining_masked
        if not logits_to_keep_half.any() and answer_mask.any():
            fallback_mask = target_scope_mask if target_scope_mask.any() else labels.ne(-100)
            if not fallback_mask.any():
                fallback_mask = answer_mask
            fallback_index = torch.nonzero(fallback_mask, as_tuple=False)[0]
            noisy_inputs_ids[fallback_index[0], fallback_index[1]] = self.config.mask_token_id
            logits_to_keep_half[fallback_index[0], fallback_index[1]] = True
            target_scope_mask[fallback_index[0], fallback_index[1]] = labels[fallback_index[0], fallback_index[1]].ne(-100)
        p_mask = build_rollout_p_mask(
            masked_indices=logits_to_keep_half,
            labels=labels,
            num_tokens=num_tokens,
            target_scope_mask=target_scope_mask,
            per_block=True,
            block_size=self.config.block_size,
        )
        concat_inputs_ids, concat_position_ids, flex_attention_mask_3d, logits_to_keep_half, logits_to_keep, _ = self.build_bd_training_inputs(
            inputs_ids=inputs_ids,
            noisy_inputs_ids=noisy_inputs_ids,
            position_ids=position_ids,
            logits_to_keep_half=logits_to_keep_half,
            num_tokens=num_tokens,
        )

        return {
            "concat_inputs_ids": concat_inputs_ids,
            "concat_position_ids": concat_position_ids,
            "flex_attention_mask_3d": flex_attention_mask_3d,
            "logits_to_keep_half": logits_to_keep_half,
            "logits_to_keep": logits_to_keep,
            "p_mask": p_mask,
            "noisy_inputs_ids": noisy_inputs_ids,
            "num_tokens": num_tokens,
            "target_scope_mask": target_scope_mask,
            "loss_target_count": target_scope_mask.sum().clamp_min(1),
            "gap_training_mode": "puma",
            "use_remask_aux": use_remask_aux,
            "rollout_depth": (
                rollout_progress.float() / rollout_steps if self._should_use_gap_prefix_frontier_state() else completed_steps
            ),
            "rollout_progress_units": rollout_progress if self._should_use_gap_prefix_frontier_state() else torch.full(
                (inputs_ids.shape[0],),
                int(completed_steps),
                dtype=torch.long,
                device=inputs_ids.device,
            ),
            "next_transfer_tokens": int(transfer_schedule[min(completed_steps, rollout_steps - 1)]),
            "rollout_strategy": rollout_strategy,
            "rollout_confidence_threshold": rollout_confidence_threshold,
            "terminal_rollout_strategy": getattr(self.config, "gap_grpo_terminal_rollout_strategy", None)
            or rollout_strategy,
            "terminal_rollout_scope": getattr(self.config, "gap_grpo_terminal_rollout_scope", None)
            or rollout_scope,
            "max_rollout_progress": self._compute_gap_prefix_progress_limits(
                labels=labels,
                num_tokens=num_tokens,
                block_size=self.config.block_size,
                rollout_steps=rollout_steps,
            ) if self._should_use_gap_prefix_frontier_state() else torch.full(
                (inputs_ids.shape[0],),
                rollout_steps,
                dtype=torch.long,
                device=inputs_ids.device,
            ),
        }

    @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,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs: Unpack[KwargsForCausalLM],
    ) -> CausalLMOutputWithPast:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Example:

        ```python
        >>> from transformers import AutoTokenizer, SDARForCausalLM

        >>> model = SDARForCausalLM.from_pretrained("DiffuOpen/SDAR-1.7B-Chat")
        >>> tokenizer = AutoTokenizer.from_pretrained("DiffuOpen/SDAR-1.7B-Chat")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        if self.training:
            assert inputs_embeds is None, "only support input_ids during training"
            prompt_mask = (labels == -100) if labels is not None else None
            coarse_debug = _gap_debug_enabled()
            coarse_step = int(getattr(self, "_gap_debug_global_step", -1))
            if position_ids is None:
                position_ids = torch.arange(
                    input_ids.shape[1], device=input_ids.device, dtype=torch.long
                ).unsqueeze(0).expand(input_ids.shape[0], -1)
            position_ids = modify_padded_position_ids_2d(position_ids)
            if coarse_debug:
                logger.info(
                    "[GAP coarse] step=%s entering_prepare_rollout batch=%s seq=%s gap_enable=%s",
                    coarse_step,
                    int(input_ids.shape[0]),
                    int(input_ids.shape[1]),
                    bool(getattr(self.config, "gap_enable", False)),
                )
            if getattr(self.config, "gap_enable", False):
                training_batch = self.prepare_for_teacher_forced_rollout_training(
                    input_ids,
                    labels,
                    position_ids,
                    prompt_mask,
                    output_attentions,
                    output_hidden_states,
                    cache_position,
                    **kwargs,
                )
            else:
                training_batch = self.prepare_for_bd_training(input_ids, position_ids, prompt_mask)
            if coarse_debug:
                logger.info(
                    "[GAP coarse] step=%s finished_prepare_rollout mode=%s logits_to_keep=%s target_scope=%s",
                    coarse_step,
                    str(training_batch.get("gap_training_mode", "bd")),
                    int(training_batch["logits_to_keep"].sum().item()) if "logits_to_keep" in training_batch else -1,
                    int(training_batch["target_scope_mask"].sum().item()) if "target_scope_mask" in training_batch else -1,
                )
            train_input_ids = training_batch.get("clean_input_ids", input_ids)
            train_labels = training_batch.get("clean_labels", labels)
            train_position_ids = training_batch.get("clean_position_ids", position_ids)
            pre_timing_interval = _gap_env_int("SDAR_GAP_GRPO_TIMING_INTERVAL", 0)
            pre_debug_step = int(getattr(self, "_gap_debug_global_step", -1))
            should_log_pre_timing = (
                pre_timing_interval > 0
                and _gap_is_rank0()
                and pre_debug_step >= 0
                and (pre_debug_step % pre_timing_interval == 0)
            )
            pre_forward_start = time.perf_counter() if should_log_pre_timing else 0.0
            if coarse_debug:
                logger.info("[GAP coarse] step=%s entering_model_forward", coarse_step)
            outputs = self.model(
                input_ids=training_batch["concat_inputs_ids"],
                attention_mask=training_batch["flex_attention_mask_3d"],
                position_ids=training_batch["concat_position_ids"],
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=True,
                cache_position=cache_position,
                **kwargs,
            )
            pre_forward_sec = (time.perf_counter() - pre_forward_start) if should_log_pre_timing else 0.0
            if coarse_debug:
                logger.info(
                    "[GAP coarse] step=%s finished_model_forward hidden_shape=%s",
                    coarse_step,
                    tuple(outputs.last_hidden_state.shape),
                )
            hidden_states = outputs.last_hidden_state
            assert train_labels is not None, "Labels must be provided for training."
            answer_len = (train_labels != -100).sum()
            hidden_states = hidden_states[training_batch["logits_to_keep"]].contiguous()
            p_mask = training_batch["p_mask"]
            diffusion_loss_weight = float(getattr(self.config, "gap_diffusion_loss_weight", 1.0) or 0.0)
            rollout_depth_metric = training_batch.get("rollout_depth", 0)
            if torch.is_tensor(rollout_depth_metric):
                rollout_depth_metric = rollout_depth_metric.detach().to(torch.float32).mean()
            else:
                rollout_depth_metric = hidden_states.new_tensor(float(rollout_depth_metric))
            if getattr(self.config, "gap_enable", False) and training_batch.get("gap_training_mode") == "puma":
                loss_fct = FusedLinearDiffusionCrossEntropyLoss(reduction='sum')
                diffusion_obj = loss_fct(
                    x=hidden_states,
                    target=train_labels[training_batch["logits_to_keep_half"]].contiguous(),
                    weight=self.lm_head.weight,
                    bias=self.lm_head.bias,
                    p_mask=training_batch["p_mask"],
                )
                diffusion_loss = diffusion_obj / answer_len.to(diffusion_obj.dtype)
                use_remask_aux = bool(training_batch.get("use_remask_aux", False))
                proposal_scores_full = None
                proposal_ids = None
                if "streaming_slot_indices" in training_batch or use_remask_aux:
                    if coarse_debug:
                        logger.info("[GAP coarse] step=%s entering_proposal_prep", coarse_step)
                    proposal_start = time.perf_counter() if should_log_pre_timing else 0.0
                    proposal_logits = self.lm_head(hidden_states).float()
                    proposal_ids = proposal_logits.argmax(dim=-1)
                    proposal_scores = (
                        proposal_logits.gather(-1, proposal_ids.unsqueeze(-1)).squeeze(-1)
                        - torch.logsumexp(proposal_logits, dim=-1)
                    ).exp()
                    proposal_scores_full = torch.full(
                        training_batch["noisy_inputs_ids"].shape,
                        float("-inf"),
                        dtype=proposal_scores.dtype,
                        device=proposal_scores.device,
                    )
                    proposal_scores_full[training_batch["logits_to_keep_half"]] = proposal_scores
                    proposal_sec = (time.perf_counter() - proposal_start) if should_log_pre_timing else 0.0
                    if coarse_debug:
                        logger.info(
                            "[GAP coarse] step=%s finished_proposal_prep proposal_shape=%s",
                            coarse_step,
                            tuple(proposal_logits.shape),
                        )
                else:
                    proposal_sec = 0.0
                if "streaming_slot_indices" in training_batch:
                    self.advance_puma_streaming_state(training_batch, proposal_scores_full)
                weighted_diffusion_loss = diffusion_loss * diffusion_loss_weight
                loss = weighted_diffusion_loss
                loss_metrics = {"rollout_depth": rollout_depth_metric}
                streaming_refills_metric = training_batch.get("streaming_refills")
                if streaming_refills_metric is not None:
                    if torch.is_tensor(streaming_refills_metric):
                        streaming_refills_metric = streaming_refills_metric.detach().to(torch.float32)
                    else:
                        streaming_refills_metric = hidden_states.new_tensor(float(streaming_refills_metric))
                    loss_metrics["streaming_refills"] = streaming_refills_metric
                if use_remask_aux:
                    if coarse_debug:
                        logger.info("[GAP coarse] step=%s entering_remask_prep", coarse_step)
                    remask_prep_start = time.perf_counter() if should_log_pre_timing else 0.0
                    remask_scope = getattr(self.config, "gap_remask_scope", "frontier_block")
                    if (
                        _gap_env_int("SDAR_GAP_GRPO_REMASK_PREFIX_GUARD_TOKENS", 0) > 0
                        or _gap_env_int("SDAR_GAP_GRPO_REMASK_TAIL_GUARD_BLOCKS", 0) > 0
                    ):
                        remask_scope = "all"
                    remask_candidate_mask = select_policy_transfer_tokens(
                        masked_indices=training_batch["logits_to_keep_half"],
                        proposal_scores_full=proposal_scores_full,
                        num_tokens=training_batch["num_tokens"],
                        block_size=self.config.block_size,
                        num_transfer_tokens=training_batch["next_transfer_tokens"],
                        strategy=training_batch["rollout_strategy"],
                        confidence_threshold=training_batch["rollout_confidence_threshold"],
                        scope=remask_scope,
                    )
                    remask_candidate_mask = self._apply_gap_grpo_remask_guards(
                        remask_candidate_mask,
                        training_batch["target_scope_mask"],
                        training_batch["logits_to_keep_half"],
                    )
                    remask_logits = self.gap_remask_head(hidden_states, proposal_logits)
                    gap_outputs = apply_gap_remask(
                        noisy_input_ids=training_batch["noisy_inputs_ids"],
                        clean_input_ids=train_input_ids,
                        labels=train_labels,
                        masked_indices=training_batch["logits_to_keep_half"],
                        p_mask=training_batch["p_mask"],
                        proposal_ids=proposal_ids,
                        remask_logits=remask_logits,
                        candidate_mask_full=remask_candidate_mask,
                        mask_token_id=self.config.mask_token_id,
                        remask_threshold=getattr(self.config, "gap_remask_threshold", 0.5),
                        remask_loss_weight=getattr(self.config, "gap_remask_loss_weight", 1.0),
                        remask_default_p_mask=getattr(self.config, "gap_remask_default_p_mask", 1.0),
                        block_size=self.config.block_size,
                        supervision=getattr(self.config, "gap_remask_supervision", "adv_bce"),
                        target_scope_mask=training_batch["target_scope_mask"],
                    )
                    remask_prep_sec = (time.perf_counter() - remask_prep_start) if should_log_pre_timing else 0.0
                    if coarse_debug:
                        logger.info(
                            "[GAP coarse] step=%s finished_remask_prep candidate_tokens=%s",
                            coarse_step,
                            int(gap_outputs.full_candidate_mask.sum().item()) if getattr(gap_outputs, "full_candidate_mask", None) is not None else -1,
                        )
                    remask_loss = gap_outputs.remask_loss
                    loss = loss + remask_loss
                    loss_metrics["diffusion_loss"] = diffusion_loss.detach()
                    loss_metrics["weighted_diffusion_loss"] = weighted_diffusion_loss.detach()
                    loss_metrics["remask_loss"] = remask_loss.detach()
                    for metric_name, metric_value in gap_outputs.metrics.items():
                        if metric_name == "remask_loss":
                            continue
                        loss_metrics[metric_name] = hidden_states.new_tensor(float(metric_value))
                    if should_log_pre_timing:
                        logger.info(
                            "[GAP pre-grpo timing] step=%s model_forward_sec=%.2f proposal_sec=%.2f remask_prep_sec=%.2f candidate_tokens=%s",
                            pre_debug_step,
                            pre_forward_sec,
                            proposal_sec,
                            remask_prep_sec,
                            int(gap_outputs.full_candidate_mask.sum().item()) if getattr(gap_outputs, "full_candidate_mask", None) is not None else -1,
                        )
                        logger.info("[GAP pre-grpo timing] step=%s entering_grpo_loss=1", pre_debug_step)
                    max_grpo_target_tokens = _gap_env_int("SDAR_GAP_GRPO_MAX_TARGET_TOKENS", 0)
                    max_grpo_valid_tokens = _gap_env_int("SDAR_GAP_GRPO_MAX_VALID_TOKENS", 0)
                    target_token_counts = training_batch["target_scope_mask"].sum(dim=1)
                    max_target_tokens = int(target_token_counts.max().item()) if target_token_counts.numel() > 0 else 0
                    max_valid_tokens = int(train_input_ids.shape[1])
                    skip_long_target = max_grpo_target_tokens > 0 and max_target_tokens > max_grpo_target_tokens
                    skip_long_valid = max_grpo_valid_tokens > 0 and max_valid_tokens > max_grpo_valid_tokens
                    if skip_long_target or skip_long_valid:
                        grpo_loss = remask_logits.sum() * 0.0 + self.gap_value_head(hidden_states[:1]).sum() * 0.0
                        grpo_metrics = {
                            "grpo_skipped_long_batch": hidden_states.new_tensor(1.0),
                            "grpo_max_target_tokens": hidden_states.new_tensor(float(max_target_tokens)),
                            "grpo_max_valid_tokens": hidden_states.new_tensor(float(max_valid_tokens)),
                        }
                        if coarse_debug:
                            _gap_stderr(
                                "[GAP grpo skip] "
                                f"step={coarse_step} max_target_tokens={max_target_tokens} "
                                f"target_limit={max_grpo_target_tokens} max_valid_tokens={max_valid_tokens} "
                                f"valid_limit={max_grpo_valid_tokens}"
                            )
                    else:
                        if coarse_debug:
                            _gap_stderr(f"[GAP coarse] step={coarse_step} entering_grpo_loss")
                        grpo_loss, grpo_metrics = self._compute_gap_grpo_loss(
                            clean_input_ids=train_input_ids,
                            labels=train_labels,
                            position_ids=train_position_ids,
                            num_tokens=training_batch["num_tokens"],
                            remask_logits=remask_logits,
                            grpo_hidden_states=hidden_states,
                            gap_outputs=gap_outputs,
                            masked_indices=training_batch["logits_to_keep_half"],
                            rollout_strategy=training_batch.get("terminal_rollout_strategy", training_batch["rollout_strategy"]),
                            rollout_confidence_threshold=training_batch["rollout_confidence_threshold"],
                            target_scope_mask=training_batch["target_scope_mask"],
                        )
                    if coarse_debug:
                        _gap_stderr(f"[GAP coarse] step={coarse_step} finished_grpo_loss")
                    loss = loss + grpo_loss
                    for metric_name, metric_value in grpo_metrics.items():
                        loss_metrics[metric_name] = metric_value.detach()
                    sft_ce_weight = float(getattr(self.config, "gap_grpo_sft_ce_weight", 0.0) or 0.0)
                    if sft_ce_weight > 0.0:
                        if coarse_debug:
                            logger.info("[GAP coarse] step=%s entering_sft_ce", coarse_step)
                        ce_timing_interval = _gap_env_int("SDAR_GAP_GRPO_TIMING_INTERVAL", 0)
                        ce_debug_step = int(getattr(self, "_gap_debug_global_step", -1))
                        should_log_ce_timing = (
                            ce_timing_interval > 0
                            and _gap_is_rank0()
                            and ce_debug_step >= 0
                            and (ce_debug_step % ce_timing_interval == 0)
                        )
                        ce_start = time.perf_counter() if should_log_ce_timing else 0.0
                        sft_ce_loss = self._compute_gap_sft_ce_anchor(
                            clean_input_ids=train_input_ids,
                            clean_labels=train_labels,
                            clean_position_ids=train_position_ids,
                        )
                        ce_sec = (time.perf_counter() - ce_start) if should_log_ce_timing else 0.0
                        weighted_sft_ce_loss = sft_ce_loss * sft_ce_weight
                        loss = loss + weighted_sft_ce_loss
                        loss_metrics["grpo_sft_ce_loss"] = sft_ce_loss.detach()
                        loss_metrics["grpo_sft_ce_weighted"] = weighted_sft_ce_loss.detach()
                        if coarse_debug:
                            logger.info("[GAP coarse] step=%s finished_sft_ce", coarse_step)
                        if should_log_ce_timing:
                            logger.info(
                                "[GAP ce timing] step=%s ce_sec=%.2f sft_ce_weight=%.3f",
                                ce_debug_step,
                                ce_sec,
                                sft_ce_weight,
                            )
                else:
                    loss_metrics["diffusion_loss"] = diffusion_loss.detach()
                    loss_metrics["weighted_diffusion_loss"] = weighted_diffusion_loss.detach()
                self._last_loss_metrics = loss_metrics
            elif getattr(self.config, "gap_enable", False):
                raise ValueError(f'Unsupported GAP training batch mode: {training_batch.get("gap_training_mode")}')
            else:
                loss_fct = FusedLinearDiffusionCrossEntropyLoss(reduction='sum')
                loss = loss_fct(  # it will return (sum_loss, unreduced_loss)
                        # conduct `view(-1, V)` inside the function
                        x=hidden_states,
                        target=train_labels[training_batch["logits_to_keep_half"]].contiguous(),
                        weight=self.lm_head.weight,
                        bias=self.lm_head.bias,
                        p_mask=training_batch["p_mask"],
                    )
                diffusion_loss = loss / answer_len
                loss = diffusion_loss
                self._last_loss_metrics = None
            logits = None
        else:
            # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
            self._last_loss_metrics = None
            outputs: BaseModelOutputWithPast = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                cache_position=cache_position,
                **kwargs,
            )

            hidden_states = outputs.last_hidden_state
            # 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
            hidden_states = hidden_states[:, slice_indices, :].contiguous()
            fuse_linear_and_cross_entropy = self.config.fuse_cross_entropy and self.training
            if fuse_linear_and_cross_entropy:
                # When using fused_linear_ce_loss, we do not compute the whole logits on HBM
                logits = None
            else:
                logits = self.lm_head(hidden_states)

            loss = None
            if labels is not None:
                # FusedLinearCrossEntropyLoss will be implemented by monkey patch when training
                # We don't use it when inferencing
                loss_fct = nn.CrossEntropyLoss()  # nn.CE
                loss = loss_fct(
                    logits.view(-1, self.config.vocab_size), labels.view(-1))

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


__all__ = [
    "SDARForCausalLM",
    "SDARModel",
    "SDARPreTrainedModel",
]