fix: stabilize D2F LoRA decoding under preempt-rebuild scenarios

diffulex/engine/engine.py

@@ -25,6 +25,32 @@
 logger = get_logger(__name__)
 
 
+def maybe_override_mask_token_id(
+    config: Config,
+    tokenizer,
+    *,
+    mask_token_id_explicit: bool = False,
+) -> None:
+    """Resolve mask token id from tokenizer artifacts when using the default."""
+
+    if mask_token_id_explicit:
+        return
+
+    default_mask_token_id = Config.__dataclass_fields__["mask_token_id"].default
+    tokenizer_mask_token_id = getattr(tokenizer, "mask_token_id", None)
+    if (
+        tokenizer_mask_token_id is not None
+        and config.mask_token_id == default_mask_token_id
+        and int(tokenizer_mask_token_id) != default_mask_token_id
+    ):
+        logger.warning(
+            "Overriding default mask_token_id from %s to tokenizer mask_token_id %s.",
+            config.mask_token_id,
+            tokenizer_mask_token_id,
+        )
+        config.mask_token_id = int(tokenizer_mask_token_id)
+
+
 def _set_parent_death_signal(sig: int = signal.SIGTERM) -> None:
     if os.name != "posix":
         return
@@ -62,6 +88,15 @@ def __init__(self, model, **kwargs):
             )
         self.ps = []
         self.events = []
+        self.tokenizer = auto_tokenizer_from_pretrained(config.model, use_fast=True, trust_remote_code=True)
+        config.tokenizer_vocab_size = len(self.tokenizer)
+        config.eos = self.tokenizer.eos_token_id
+        maybe_override_mask_token_id(
+            config,
+            self.tokenizer,
+            mask_token_id_explicit="mask_token_id" in config_kwargs,
+        )
+
         ctx = mp.get_context("spawn")
         for i in range(1, self.model_parallel_world_size):
             event = ctx.Event()
@@ -75,21 +110,6 @@ def __init__(self, model, **kwargs):
         self._install_signal_handlers()
 
         try:
-            self.tokenizer = auto_tokenizer_from_pretrained(config.model, use_fast=True, trust_remote_code=True)
-            config.tokenizer_vocab_size = len(self.tokenizer)
-            config.eos = self.tokenizer.eos_token_id
-
-            if (
-                getattr(self.tokenizer, "mask_token_id", None) is not None
-                and config.mask_token_id != self.tokenizer.mask_token_id
-            ):
-                logger.warning(
-                    "Overriding mask_token_id from %s to tokenizer mask_token_id %s.",
-                    config.mask_token_id,
-                    self.tokenizer.mask_token_id,
-                )
-                config.mask_token_id = self.tokenizer.mask_token_id
-
             self.model_runner = AutoModelRunner.from_config(config, 0, self.events)
             self.scheduler: SchedulerBase | DataParallelScheduler = AutoScheduler.from_config(config)
         except BaseException:

diffulex/engine/request.py

@@ -120,6 +120,8 @@ def init_multi_block(self, config: Config):
         self.is_multi_block = True
         self.status_history = [self.status]
         self.completion_reason = None
+        self._resume_prefill_until = 0
+        self._terminal_context_block_id: int | None = None
 
         self.block_size = config.block_size
         self.buffer_size = config.buffer_size
@@ -341,6 +343,8 @@ def chunk_size(self) -> int:
     @property
     def running_len(self) -> int:
         if self.is_prefilling:
+            if self._resume_prefill_until > 0:
+                return self._resume_prefill_until
             return (
                 (self.padded_prefix_len - self.block_size) + self.dllm_block_buffer.num_valid_blocks * self.block_size
                 if self.is_padded
@@ -419,7 +423,7 @@ def to_cache_seq_end(self) -> tuple[int, int]:
     @property
     def has_to_cache_blocks(self) -> bool:
         if self.is_prefilling:
-            return True
+            return self._prefill_visible_to_cache_last_global_id() is not None
         if self.is_decoding:
             return len(self.dllm_block_buffer.to_cache_blocks) > 0
         return False
@@ -431,14 +435,67 @@ def has_to_cache_block(self) -> bool:
     @property
     def to_cache_last_token_id(self) -> int:
         if self.is_prefilling:
-            return self.to_cache_len - 1 if self.to_cache_len > 0 else 0
+            window_start = int(self.contiguous_in_cache_prefix_len)
+            last_global = self._prefill_visible_to_cache_last_global_id()
+            if last_global is None:
+                return 0
+            return int(last_global - window_start)
         n = len(self.dllm_block_buffer.to_cache_blocks) * self.block_size
         return n - 1 if n > 0 else 0
 
+    def _prefill_visible_to_cache_last_global_id(self) -> int | None:
+        if not self.is_prefilling:
+            return None
+
+        window_start = int(self.contiguous_in_cache_prefix_len)
+        window_end = int(self.running_len)
+        if window_end <= window_start:
+            return None
+
+        last_global = None
+        for block in self.dllm_blocks:
+            if block.end <= window_start:
+                continue
+            if block.start >= window_end:
+                break
+            if not block.is_to_cache:
+                continue
+            candidate = min(block.end, window_end) - 1
+            if candidate < window_start:
+                continue
+            last_global = candidate if last_global is None else max(last_global, candidate)
+        return last_global
+
     @property
     def last_block_finished(self) -> bool:
-        inspected_block = self.dllm_block_buffer.first_running_block.prev_block
-        return inspected_block is not None and inspected_block.is_complete and inspected_block.is_last_in_context
+        terminal_block = self.terminal_context_block
+        return terminal_block is not None and terminal_block.is_complete
+
+    @property
+    def terminal_context_block(self) -> DllmBlock | None:
+        block_id = getattr(self, "_terminal_context_block_id", None)
+        if block_id is None or not (0 <= int(block_id) < len(self.dllm_blocks)):
+            return None
+        return self.dllm_blocks[int(block_id)]
+
+    def set_terminal_context_block(self, block: DllmBlock | None) -> None:
+        while block is not None and block.is_dummy:
+            block = block.prev_block
+        if block is None:
+            self._terminal_context_block_id = None
+            return
+
+        terminal_block_id = int(block.block_id)
+        self._terminal_context_block_id = terminal_block_id
+
+        for dllm_block in self.dllm_blocks:
+            if dllm_block.is_dummy or dllm_block.block_id > terminal_block_id:
+                dllm_block.make_out_of_context()
+            elif dllm_block.block_id < terminal_block_id:
+                dllm_block.make_in_context()
+            else:
+                dllm_block.make_in_context()
+                dllm_block.make_last_in_context()
 
     @property
     def pure_prefill_without_mask_token(self) -> bool:
@@ -463,6 +520,16 @@ def make_pending(self):
     def preempt(self):
         self.lazy_activate()
         self.log_status()
+        if self.is_multi_block:
+            rebuild_until = 0
+            if self.is_decoding:
+                rebuild_until = int(self.running_seq_start)
+                self._resume_prefill_until = rebuild_until
+            for block in self.dllm_blocks:
+                if rebuild_until > 0 and block.end <= rebuild_until and block.is_to_cache:
+                    continue
+                if block.is_in_cache:
+                    block.status = DllmBlockStatus.TO_CACHE
         self.status = DllmReqStatus.WAITING
 
     @property
@@ -476,6 +543,9 @@ def lazy_activate(self):
         self.log_status()
 
         self.status = self.status_history[-1]
+        if self._resume_prefill_until > 0:
+            self.status = DllmReqStatus.PREFILLING
+            return
         if self.is_pending:
             self.status = DllmReqStatus.PREFILLING
         elif self.is_prefilling:
@@ -506,6 +576,15 @@ def deactivate(self, reason: str | None = None):
     def step(self):
         self.lazy_activate()
 
+        if (
+            self.is_decoding
+            and not self.dllm_block_buffer.active_blocks
+            and not self.dllm_block_buffer.to_cache_blocks
+        ):
+            head_block = self.dllm_block_buffer.first_running_block
+            if head_block.is_in_cache and not head_block.is_last_in_context:
+                head_block.status = DllmBlockStatus.TO_CACHE
+
         for block in self.dllm_block_buffer.active_blocks:
             block.total_steps += 1
 
@@ -534,9 +613,13 @@ def push_back_dummy_block(self):
         )
 
         dllm_block.post_init_dllm_block(self, self.dllm_block_buffer)
-        if (self.max_new_tokens_reached or self.max_model_len_reached) and dllm_block.prev_block.is_in_context:
-            dllm_block.make_last_in_context()
-        elif dllm_block.prev_block.is_out_of_context or dllm_block.prev_block.is_last_in_context:
+        if (
+            self.max_new_tokens_reached
+            or self.max_model_len_reached
+            or self.terminal_context_block is not None
+            or dllm_block.prev_block.is_out_of_context
+            or dllm_block.prev_block.is_last_in_context
+        ):
             dllm_block.make_out_of_context()
 
         self.dllm_blocks.append(dllm_block)
@@ -549,6 +632,13 @@ def maybe_postprocess_prefix_blocks(self):
         for block_id in range(self.num_prefix_blocks):
             self.dllm_blocks[block_id].in_cache()
 
+        if self._resume_prefill_until > 0:
+            for block in self.dllm_blocks:
+                if block.end > self._resume_prefill_until:
+                    break
+                if block.is_to_cache and block.is_complete:
+                    block.in_cache()
+
     def apply_cached_prefix_pages(self):
         if not self.is_multi_block:
             return
@@ -585,8 +675,13 @@ def postprocess(self):
             elif block.is_dummy or block.is_active or block.is_in_cache:
                 block_id += 1
 
+        if self._resume_prefill_until > 0 and self.contiguous_in_cache_prefix_len >= self._resume_prefill_until:
+            self._resume_prefill_until = 0
+
+        if self.is_truncated:
+            self.set_terminal_context_block(self.dllm_block_buffer.last_valid_block)
+
         if self.eos_token_generated:
-            self.dllm_block_buffer.last_valid_block.make_last_in_context()
             self.meet_eos = True
             self.dllm_block_buffer.maybe_fix_context_management()
 

diffulex/sampler/base/no_shift.py

@@ -38,15 +38,9 @@ def _prefill_mask_token_local_ids(req: DllmReq, block, req_logits: torch.Tensor)
             return local_ids
 
         if min(local_ids) < 0 or max(local_ids) >= req_logits.shape[0]:
-            raise IndexError(
-                "Prefill mask-token logits index out of bounds: "
-                f"req_id={getattr(req, 'req_id', '?')}, "
-                f"block_id={getattr(block, 'block_id', '?')}, "
-                f"in_cache_len={prefix_offset}, "
-                f"global_ids={block.mask_token_global_ids}, "
-                f"local_ids={local_ids}, "
-                f"req_logits_len={req_logits.shape[0]}"
-            )
+            # Mixed prefill batches can yield partial q_len for a req in one step.
+            # Skip this block this step and retry when its logits slice is present.
+            return []
         return local_ids
 
     def forward(
@@ -96,11 +90,15 @@ def forward(
 
                     with record_function("diffulex.sampler.no_shift.mask_logits"):
                         if attn_metadata.is_prefill[idx]:
+                            if getattr(req, "_resume_prefill_until", 0) > 0 and getattr(block, "start", 0) >= req.running_len:
+                                continue
                             # Prefix-cache prefill can produce q_len=0 for some requests in mixed batches.
                             # In that case there are no logits to sample for this req in this step.
                             if req_logits.shape[0] == 0:
                                 continue
                             local_ids = self._prefill_mask_token_local_ids(req, block, req_logits)
+                            if not local_ids:
+                                continue
                             mask_token_logits = req_logits[local_ids, ...]
                         else:
                             buf_offset = block.start - req.dllm_block_buffer.first_running_block.start

diffulex/sampler/base/shift.py

@@ -31,7 +31,9 @@ def evict_req_states(self, req_ids: list[int] | list[str]) -> None:
     def _fetch_last_logits(self, logits: torch.Tensor, req: DllmReq) -> torch.Tensor:
         req_id_str = str(req.req_id)
         if req.has_to_cache_block:
-            return self._cache_last_logits(req_id_str, logits[req.to_cache_last_token_id])
+            idx = int(req.to_cache_last_token_id)
+            if 0 <= idx < logits.shape[0]:
+                return self._cache_last_logits(req_id_str, logits[idx])
 
         if req_id_str in self.req_last_logits_map:
             return self.req_last_logits_map[req_id_str]
@@ -114,6 +116,8 @@ def forward(
                     if shifted_logits.shape[0] == 0:
                         continue
                     local_ids = DllmSamplerNoShiftBase._prefill_mask_token_local_ids(req, block, shifted_logits)
+                    if not local_ids:
+                        continue
                     mask_token_logits = shifted_logits[local_ids, ...]
                 else:
                     buf_offset = block.start - req.dllm_block_buffer.first_running_block.start

diffulex/server/args.py

@@ -23,6 +23,7 @@ class ServerArgs:
     model_name: str = "dream"
     decoding_strategy: str = "d2f"
     sampling_mode: str = "naive"
+    mask_token_id: int | None = None
     tensor_parallel_size: int = 1
     data_parallel_size: int = 1
     master_addr: str = "localhost"
@@ -119,6 +120,8 @@ def engine_kwargs(self) -> dict:
             "lora_path": self.lora_path,
             "pre_merge_lora": self.pre_merge_lora,
         }
+        if self.mask_token_id is not None:
+            kwargs["mask_token_id"] = self.mask_token_id
         if self.profiler is not None:
             kwargs["profiler_config"] = {
                 "profiler": self.profiler,
@@ -144,6 +147,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument("--model-name", default="dream")
     parser.add_argument("--decoding-strategy", default="d2f")
     parser.add_argument("--sampling-mode", default="naive", choices=["naive", "edit"])
+    parser.add_argument("--mask-token-id", type=int, default=None)
     parser.add_argument("--tensor-parallel-size", type=int, default=1)
     parser.add_argument("--data-parallel-size", type=int, default=1)
     parser.add_argument("--master-addr", default="localhost")
@@ -205,6 +209,7 @@ def parse_args(argv: Sequence[str] | None = None) -> ServerArgs:
         model_name=ns.model_name,
         decoding_strategy=ns.decoding_strategy,
         sampling_mode=ns.sampling_mode,
+        mask_token_id=ns.mask_token_id,
         tensor_parallel_size=ns.tensor_parallel_size,
         data_parallel_size=ns.data_parallel_size,
         master_addr=ns.master_addr,