[Chore] propagate changes

configuration_rnd.py

@@ -30,7 +30,7 @@ class RND1Config(Qwen3MoeConfig):
         self,
         moe_backend: str = "hf",
         num_diffusion_steps: int = 256,
-        mask_token_id: int = 151669,  # Default for Qwen-based RND1 models
+        mask_token_id: int = 151669,
         use_cache: bool = False,
         **kwargs,
     ):

generation_config.py

@@ -31,12 +31,11 @@ class RND1GenerationConfig(GenerationConfig):
         self,
         max_length: int = 256,
         num_diffusion_steps: int = 256,
-        mask_token_id: int = 151669,  # Default for Qwen-based RND1 models
+        mask_token_id: int = 151669,
         temperature: float = 1.0,
         top_k: Optional[int] = None,
         top_p: Optional[float] = None,
         greedy: bool = True,
-        seed: Optional[int] = None,  # For reproducible generation
         bos_token_id: Optional[int] = None,
         eos_token_id: Optional[int] = None,
         pad_token_id: Optional[int] = None,
@@ -64,7 +63,6 @@ class RND1GenerationConfig(GenerationConfig):
         self.mask_token_id = mask_token_id
         self.greedy = greedy
         self.temperature = float(temperature)  # Ensure it's a float
-        self.seed = seed
 
     def to_dict(self):
         """Convert configuration to dictionary."""
@@ -72,6 +70,4 @@ class RND1GenerationConfig(GenerationConfig):
         output["num_diffusion_steps"] = self.num_diffusion_steps
         output["mask_token_id"] = self.mask_token_id
         output["greedy"] = self.greedy
-        if self.seed is not None:
-            output["seed"] = self.seed
         return output

sampling.py

@@ -16,41 +16,25 @@ def apply_top_k_filtering(logits: torch.Tensor, k: int) -> torch.Tensor:
     """
     Apply top-k filtering to logits: with non-top-k values set to -inf
     """
-    if k is None or k <= 0:
-        return torch.full_like(logits, float("-inf"))
-    k = min(k, logits.size(-1))
-    top_k_values, top_k_indices = torch.topk(logits, k, dim=-1)
-    filtered = torch.full_like(logits, float("-inf"))
-    filtered.scatter_(-1, top_k_indices, top_k_values)
-    return filtered
+    top_k_values, top_k_indices = torch.topk(logits, min(k, logits.size(-1)), dim=-1)
+    filtered_logits = torch.full_like(logits, float('-inf'))
+    filtered_logits.scatter_(-1, top_k_indices, top_k_values)
+    return filtered_logits
 
 
-def apply_top_p_filtering(logits: torch.Tensor, p: float, min_tokens_to_keep: int = 1) -> torch.Tensor:
+def apply_top_p_filtering(logits: torch.Tensor, p: float) -> torch.Tensor:
     """
     Apply top-p (nucleus) filtering to logits: with tokens beyond threshold set to -inf
     """
-    if p <= 0:
-        p = 1e-8
-    if p >= 1:
-        return logits
-    
     sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
 
-    probs = torch.softmax(sorted_logits, dim=-1)
-    cumulative_probs = torch.cumsum(probs, dim=-1)
-
+    # Remove tokens with cumulative probability above threshold
     sorted_indices_to_remove = cumulative_probs > p
-
-    if min_tokens_to_keep > 0:
-        sorted_indices_to_remove[..., :min_tokens_to_keep] = False
-
+    sorted_indices_to_remove[..., 0] = False  # Keep at least one token
     sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-    sorted_indices_to_remove[..., 0] = False
-
-    indices_to_remove = torch.zeros_like(sorted_indices_to_remove)
-    indices_to_remove.scatter_(-1, sorted_indices, sorted_indices_to_remove)
-
-    return logits.masked_fill(indices_to_remove, float("-inf"))
+    indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
+    return logits.masked_fill(indices_to_remove, float('-inf'))
 
 
 @torch.no_grad()
@@ -171,22 +155,13 @@ def diffusion_sample(
         if bos_token_id is not None:
             x[0, 0] = bos_token_id
         if eos_token_id is not None:
-            # If eos_token_id is a list, use the first one
-            if isinstance(eos_token_id, (list, tuple)):
-                x[0, -1] = eos_token_id[0]
-            else:
-                x[0, -1] = eos_token_id
+            x[0, -1] = eos_token_id
         init_maskable = x.eq(mask_token_id)
 
     if bos_token_id is not None:
         init_maskable[:, 0] = False
     if eos_token_id is not None:
-        # Handle both single token and list of tokens
-        if isinstance(eos_token_id, (list, tuple)):
-            for eos_id in eos_token_id:
-                init_maskable &= x.ne(eos_id)
-        else:
-            init_maskable &= x.ne(eos_token_id)
+        init_maskable &= x.ne(eos_token_id)
     init_maskable &= x.ne(pad_token_id)
 
     maskable = init_maskable.clone()
@@ -204,34 +179,28 @@ def diffusion_sample(
             # Fall back to positional argument
             model_output = model(tokens)
 
-        # Apply temperature scaling (if temperature == 0, treat as 1.0 for greedy)
-        logits = model_output.logits
-        if temperature > 0:
-            logits = logits / temperature
+        # Apply temperature scaling (with safety for near-zero temperature)
+        safe_temperature = max(temperature, 1e-8)  # Prevent division by zero
+        logits = model_output.logits / safe_temperature
 
-        # Apply filtering only when not in greedy mode
-        # Order matches reference: top_p before top_k
-        if not greedy:
-            if top_p is not None and 0 < top_p < 1.0:
-                logits = apply_top_p_filtering(logits, top_p)
 
-            if top_k is not None and top_k > 0:
-                logits = apply_top_k_filtering(logits, top_k)
+        # Apply filtering strategies
+        # Note: When both top_k and top_p are provided, they are applied sequentially:
+        # First top_k filters to k tokens, then top_p filters from those k tokens
+        if top_k is not None and top_k > 0:
+            logits = apply_top_k_filtering(logits, top_k)
 
-        # Compute probabilities for sampling and metrics
-        probs = torch.softmax(logits, dim=-1)
-        logp = torch.log(probs + 1e-10)  # Add epsilon for numerical stability
+        if top_p is not None and 0 < top_p < 1.0:
+            logits = apply_top_p_filtering(logits, top_p)
 
+        # Convert to log probabilities
+        logp = torch.log_softmax(logits, dim=-1)
+
+        # Greedy or stochastic sampling
         if greedy:
             pred_next = logp.argmax(-1)
         else:
-            # Sample from categorical distribution with proper RNG handling
-            if generator is not None:
-                # Use multinomial with generator for reproducible sampling
-                pred_next = torch.multinomial(probs.view(-1, probs.size(-1)), 1, generator=generator).squeeze(-1).view(probs.shape[:-1])
-            else:
-                # Sample from categorical using probabilities
-                pred_next = torch.distributions.Categorical(probs=probs).sample()
+            pred_next = torch.distributions.Categorical(logits=logp).sample(generator=generator)
 
         conf_next = torch.gather(logp, -1, pred_next.unsqueeze(-1)).squeeze(-1)