[fix] gate M2T unmasking on sampled-token confidence p_s

modeling_llada2_moe.py

@@ -1520,14 +1520,14 @@ class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
             neg_inf = torch.full_like(p0, -float("inf"))
 
             # 3. M2T (mask -> token): threshold-gated with a per-step floor.
-            # Note: when temperature > 0 the gate uses greedy confidence ``p0``, not ``p_s``.
-            # ``threshold`` decides if a position is ready to unmask; sampling only picks which
-            # token (``x_s``) is written, so the final token's own probability is not constrained.
+            # The gate uses ``p_s`` -- the confidence of the token that will actually be written
+            # (``x_s``) -- so a position is only unmasked when the sampled token itself is
+            # confident. When temperature == 0, ``p_s == p0``, so this reduces to greedy behavior.
             mt2_index = torch.zeros(block_length, dtype=torch.bool, device=device)
             if mask_index.any():
                 if step_id < len(transfer_schedule):
                     num_need = transfer_schedule[step_id].item() + new_mask_count
-                    mask_conf = torch.where(mask_index, p0, neg_inf)
+                    mask_conf = torch.where(mask_index, p_s, neg_inf)
                     high_conf = (mask_conf > threshold) & mask_index
                     if high_conf.sum().item() >= num_need:
                         mt2_index = high_conf