From c43f6bc986330d0684af12c2ee0c893eaafce429 Mon Sep 17 00:00:00 2001 From: Osc-7 <1449028496@qq.com> Date: Mon, 20 Apr 2026 10:35:47 +0000 Subject: [PATCH 1/5] fix: align LoRA target matching with packed_modules_mapping --- diffulex/utils/loader.py | 30 +++++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/diffulex/utils/loader.py b/diffulex/utils/loader.py index e54c6dd1..88b3a478 100755 --- a/diffulex/utils/loader.py +++ b/diffulex/utils/loader.py @@ -23,19 +23,38 @@ def load_lora_config(lora_path: str) -> dict: return {} -def enable_lora_for_model(model: nn.Module, lora_config: dict): - """Enable LoRA for existing linear layers in the model.""" +def enable_lora_for_model( + model: nn.Module, + lora_config: dict, + packed_modules_mapping: dict | None = None, +): + """Enable LoRA for existing linear layers in the model. + + `target_modules` from PEFT adapter_config refers to the *checkpoint* leaf + names (e.g. `attn_out`). When the local model class re-names a layer (e.g. + LLaDA's `attn_out` is implemented as `self_attn.o_proj`), the mapping is + declared in `packed_modules_mapping` as `{ckpt_leaf: (local_dotted_name, _)}`. + We must consult that mapping here, otherwise renamed targets silently miss + `__init_lora__` and the loaded LoRA tensors get dropped at apply time. + """ r = lora_config.get("r", 16) lora_alpha = lora_config.get("lora_alpha", 32.0) lora_dropout = lora_config.get("lora_dropout", 0.0) target_modules = lora_config.get("target_modules", []) + rev_mapping = {} + if packed_modules_mapping: + for ckpt_leaf, (local_dotted, _) in packed_modules_mapping.items(): + local_leaf = local_dotted.split(".")[-1] + rev_mapping[local_leaf] = ckpt_leaf + for name, module in model.named_modules(): if hasattr(module, "__init_lora__"): should_apply = True if target_modules: leaf = name.split(".")[-1] if name else name - should_apply = any(target == leaf for target in target_modules) + effective = rev_mapping.get(leaf, leaf) + should_apply = any(target == effective for target in target_modules) if should_apply: module.__init_lora__(r, lora_alpha, lora_dropout) return model @@ -177,13 +196,14 @@ def load_model(model: nn.Module, config: Config): # Enable LoRA for linear layers if LoRA is enabled if config.use_lora and config.lora_path: lora_config = load_lora_config(config.lora_path) + packed_modules_mapping_for_lora = getattr(model, "packed_modules_mapping", None) if lora_config: logger.info(f"LoRA Config Loaded: {lora_config}") - model = enable_lora_for_model(model, lora_config) + model = enable_lora_for_model(model, lora_config, packed_modules_mapping_for_lora) else: logger.info("No adapter_config.json found, using default LoRA parameters") default_config = {"r": 16, "lora_alpha": 32.0, "lora_dropout": 0.0} - model = enable_lora_for_model(model, default_config) + model = enable_lora_for_model(model, default_config, packed_modules_mapping_for_lora) # Load base model weights packed_modules_mapping = getattr(model, "packed_modules_mapping", {}) From 5653c6992b2b5ade52a468dcfe4ced11db59c3ad Mon Sep 17 00:00:00 2001 From: Osc-7 <1449028496@qq.com> Date: Tue, 21 Apr 2026 14:47:08 +0000 Subject: [PATCH 2/5] fix(d2f): stabilize llada decode after preempt with rebuild-aware request flow - add preempt-rebuild request state handling in multi_block request template - prevent decode livelock by recycling cached head block when frontier is lost - align no_shift prefill sampling with resume_prefill_until window to avoid out-of-window block sampling This keeps the fix scoped to request/no_shift without experimental guards or benchmark config changes. --- diffulex/engine/request.py | 111 +++++++++++++++++++++++++++--- diffulex/sampler/base/no_shift.py | 16 ++--- 2 files changed, 110 insertions(+), 17 deletions(-) diff --git a/diffulex/engine/request.py b/diffulex/engine/request.py index 4d658070..9b097d31 100755 --- a/diffulex/engine/request.py +++ b/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() diff --git a/diffulex/sampler/base/no_shift.py b/diffulex/sampler/base/no_shift.py index 77809fa4..c1ea5259 100644 --- a/diffulex/sampler/base/no_shift.py +++ b/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 From 5128f6604a0bc9db95a768883aafec981058b69b Mon Sep 17 00:00:00 2001 From: Osc-7 <1449028496@qq.com> Date: Tue, 21 Apr 2026 16:44:17 +0000 Subject: [PATCH 3/5] fix(sampler): align prefill indices for preempt-rebuild stability --- diffulex/sampler/base/shift.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/diffulex/sampler/base/shift.py b/diffulex/sampler/base/shift.py index bd94b3c8..ae60e19d 100644 --- a/diffulex/sampler/base/shift.py +++ b/diffulex/sampler/base/shift.py @@ -31,7 +31,20 @@ 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]) + logger.warning( + "Invalid to_cache_last_token_id for req %s: idx=%s, logits_len=%s; fallback to last row.", + req_id_str, + idx, + logits.shape[0], + ) + if logits.shape[0] > 0: + return self._cache_last_logits(req_id_str, logits[-1]) + raise ValueError( + f"Cannot fetch last logits for req {req.req_id}: empty logits tensor with invalid index {idx}" + ) if req_id_str in self.req_last_logits_map: return self.req_last_logits_map[req_id_str] @@ -114,6 +127,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 From cc376f23c736b147639c62a0cfa5f43d3c00ac8c Mon Sep 17 00:00:00 2001 From: Osc-7 <1449028496@qq.com> Date: Tue, 21 Apr 2026 18:38:54 +0000 Subject: [PATCH 4/5] fix(request): bound prefill to-cache index to visible logits window --- diffulex/sampler/base/shift.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/diffulex/sampler/base/shift.py b/diffulex/sampler/base/shift.py index ae60e19d..8ef87732 100644 --- a/diffulex/sampler/base/shift.py +++ b/diffulex/sampler/base/shift.py @@ -34,17 +34,6 @@ def _fetch_last_logits(self, logits: torch.Tensor, req: DllmReq) -> torch.Tensor idx = int(req.to_cache_last_token_id) if 0 <= idx < logits.shape[0]: return self._cache_last_logits(req_id_str, logits[idx]) - logger.warning( - "Invalid to_cache_last_token_id for req %s: idx=%s, logits_len=%s; fallback to last row.", - req_id_str, - idx, - logits.shape[0], - ) - if logits.shape[0] > 0: - return self._cache_last_logits(req_id_str, logits[-1]) - raise ValueError( - f"Cannot fetch last logits for req {req.req_id}: empty logits tensor with invalid index {idx}" - ) if req_id_str in self.req_last_logits_map: return self.req_last_logits_map[req_id_str] From 92cc4e68fd0a2b13210f93e8a119b4b1f7ee5796 Mon Sep 17 00:00:00 2001 From: Osc-7 <1449028496@qq.com> Date: Wed, 24 Jun 2026 18:54:49 +0000 Subject: [PATCH 5/5] fix: preserve mask token and lora target handling --- diffulex/engine/engine.py | 50 ++++++++++++++------ diffulex/server/args.py | 5 ++ diffulex/utils/loader.py | 2 + test/python/engine/test_engine_mask_token.py | 32 +++++++++++++ test/python/model/test_lora_loader.py | 45 ++++++++++++++++++ test/python/server/test_args.py | 6 +++ 6 files changed, 125 insertions(+), 15 deletions(-) create mode 100644 test/python/engine/test_engine_mask_token.py create mode 100644 test/python/model/test_lora_loader.py diff --git a/diffulex/engine/engine.py b/diffulex/engine/engine.py index 4e0a7ccc..83c9bf28 100644 --- a/diffulex/engine/engine.py +++ b/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: diff --git a/diffulex/server/args.py b/diffulex/server/args.py index 6e699c6a..2687f9e0 100644 --- a/diffulex/server/args.py +++ b/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, diff --git a/diffulex/utils/loader.py b/diffulex/utils/loader.py index 88b3a478..89257bbb 100755 --- a/diffulex/utils/loader.py +++ b/diffulex/utils/loader.py @@ -41,6 +41,8 @@ def enable_lora_for_model( lora_alpha = lora_config.get("lora_alpha", 32.0) lora_dropout = lora_config.get("lora_dropout", 0.0) target_modules = lora_config.get("target_modules", []) + if isinstance(target_modules, str): + target_modules = [target_modules] rev_mapping = {} if packed_modules_mapping: diff --git a/test/python/engine/test_engine_mask_token.py b/test/python/engine/test_engine_mask_token.py new file mode 100644 index 00000000..1c21d5f6 --- /dev/null +++ b/test/python/engine/test_engine_mask_token.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from diffulex.engine.engine import maybe_override_mask_token_id + + +def test_mask_token_override_uses_tokenizer_when_config_is_default(): + config = SimpleNamespace(mask_token_id=151666) + tokenizer = SimpleNamespace(mask_token_id=126336) + + maybe_override_mask_token_id(config, tokenizer) + + assert config.mask_token_id == 126336 + + +def test_mask_token_override_preserves_explicit_value(): + config = SimpleNamespace(mask_token_id=151666) + tokenizer = SimpleNamespace(mask_token_id=126336) + + maybe_override_mask_token_id(config, tokenizer, mask_token_id_explicit=True) + + assert config.mask_token_id == 151666 + + +def test_mask_token_override_preserves_non_default_value(): + config = SimpleNamespace(mask_token_id=151665) + tokenizer = SimpleNamespace(mask_token_id=126336) + + maybe_override_mask_token_id(config, tokenizer) + + assert config.mask_token_id == 151665 diff --git a/test/python/model/test_lora_loader.py b/test/python/model/test_lora_loader.py new file mode 100644 index 00000000..fc0d9650 --- /dev/null +++ b/test/python/model/test_lora_loader.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import torch.nn as nn + +from diffulex.utils.loader import enable_lora_for_model + + +class DummyLoraModule(nn.Module): + def __init__(self): + super().__init__() + self.initialized = False + + def __init_lora__(self, r, lora_alpha, lora_dropout): + self.initialized = True + self.lora_args = (r, lora_alpha, lora_dropout) + + +class DummyModel(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = DummyLoraModule() + self.self_attn = nn.Module() + self.self_attn.o_proj = DummyLoraModule() + + +def test_enable_lora_accepts_string_target_modules(): + model = DummyModel() + + enable_lora_for_model(model, {"target_modules": "q_proj"}) + + assert model.q_proj.initialized is True + assert model.self_attn.o_proj.initialized is False + + +def test_enable_lora_matches_packed_module_checkpoint_names(): + model = DummyModel() + + enable_lora_for_model( + model, + {"target_modules": ["attn_out"]}, + packed_modules_mapping={"attn_out": ("self_attn.o_proj", None)}, + ) + + assert model.q_proj.initialized is False + assert model.self_attn.o_proj.initialized is True diff --git a/test/python/server/test_args.py b/test/python/server/test_args.py index 6a2dfc55..1ad23b54 100644 --- a/test/python/server/test_args.py +++ b/test/python/server/test_args.py @@ -53,3 +53,9 @@ def test_server_args_forward_recent_engine_config_fields(): "delay_iterations": 0, "max_iterations": 3, } + + +def test_server_args_forward_explicit_mask_token_id(): + args = parse_args(["--model", "/tmp/model", "--mask-token-id", "126336"]) + + assert args.engine_kwargs()["mask_token_id"] == 126336