diff --git "a/modeling_sdar.py" "b/modeling_sdar.py"
new file mode 100644--- /dev/null
+++ "b/modeling_sdar.py"
@@ -0,0 +1,4362 @@
+# 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",
+]