fix: stabilize D2F LoRA decoding under preempt-rebuild scenarios#40
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (6)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR adds optional mask-token-id plumbing from server args into engine startup, changes request and sampler logic to handle resumed prefill and terminal-context blocks, and updates LoRA loading to resolve target modules through packed-module metadata. ChangesMask token override and resumed prefill handling
Packed-module LoRA mapping
Sequence Diagram(s)sequenceDiagram
participant parse_args
participant ServerArgs
participant DiffulexEngine
participant maybe_override_mask_token_id
participant tokenizer
parse_args->>ServerArgs: populate mask_token_id from ns.mask_token_id
DiffulexEngine->>maybe_override_mask_token_id: call with mask_token_id_explicit
maybe_override_mask_token_id->>tokenizer: read tokenizer.mask_token_id
maybe_override_mask_token_id->>DiffulexEngine: update config.mask_token_id when still default
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Osc-7
marked this pull request as ready for review
April 22, 2026 02:32
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
diffulex/strategy_template/multi_block/engine/request.py (1)
11-25:⚠️ Potential issue | 🟡 MinorInitialize
_resume_prefill_untilin_restore_req_runtime_stateto preventAttributeErroron unpickling.
_resume_prefill_untilis initialized only ininit_multi_block(line 31) but not restored in_restore_req_runtime_state. When a multi-block request is unpickled via__setstate__, if the attribute is missing from the pickled state (edge case: old pickled requests, failed migrations, etc.), subsequent accesses at lines 232, 415, 507, 550 will raiseAttributeError. The sampler already uses defensivegetattr(req, "_resume_prefill_until", 0)at line 99 ofno_shift.py, establishing the pattern.Initialize it in
_restore_req_runtime_stateto match the mixin pattern shown in the test suite:Suggested fix
def _restore_req_runtime_state(self): super()._restore_req_runtime_state() if not self.is_multi_block: return + self._resume_prefill_until = getattr(self, "_resume_prefill_until", 0) + dllm_blocks = self.dllm_blocks🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex/strategy_template/multi_block/engine/request.py` around lines 11 - 25, The _restore_req_runtime_state method must ensure the _resume_prefill_until attribute exists after unpickling to avoid AttributeError; update the _restore_req_runtime_state implementation to initialize self._resume_prefill_until = getattr(self, "_resume_prefill_until", 0) (or assign 0 if missing) before any sampler or block logic runs—do this alongside the existing multi-block handling in _restore_req_runtime_state (the same place that accesses dllm_blocks and dllm_block_buffer) so it matches initialization in init_multi_block and the defensive pattern used in no_shift.py.
🧹 Nitpick comments (2)
diffulex_bench/configs/llada_instruct_gsm8k.yml (1)
19-21: Prefill concurrency is effectively capped to 1 with these values.With
max_num_batched_tokens == max_model_len(Line 20 and Line 19), warmup prefill concurrency becomes1even thoughmax_num_reqsis24(Line 21). If this profile is meant for stability-only that’s fine, but if benchmark throughput matters, consider settingmax_num_batched_tokensabovemax_model_len(or add a comment explaining the intentional tradeoff).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex_bench/configs/llada_instruct_gsm8k.yml` around lines 19 - 21, The config sets max_num_batched_tokens equal to max_model_len which forces warmup prefill concurrency to 1 despite max_num_reqs being 24; either raise max_num_batched_tokens above max_model_len (e.g., to a multiple of max_model_len) to allow higher prefill concurrency or add a comment nearby documenting that this equality is intentional to cap concurrency for stability; update the values or add the explanatory comment referencing the keys max_model_len, max_num_batched_tokens, and max_num_reqs.diffulex/strategy_template/multi_block/engine/request.py (1)
397-401: Redundantcontinuein the preempt demotion loop.If
block.is_to_cacheis True,block.is_in_cacheis False (the statuses are mutually exclusive), so theif block.is_in_cache:branch is already a no-op for those blocks. Thecontinueat line 399 therefore skips nothing.Either drop the condition, or if the intent is to document that TO_CACHE blocks within
rebuild_untilare deliberately preserved, convert it to a clarifying comment. As written, future readers may infer that TO_CACHE would otherwise be demoted here, which is not the case.♻️ Suggested simplification
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 - # Scheduler preemption frees this req's page_table immediately after moving it - # back to WAITING. Any block still marked IN_CACHE would then point to KV pages - # that no longer belong to the req. Demote those blocks back to TO_CACHE so a - # resumed req rebuilds its cached prefix instead of reusing stale block state. - 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 + # Scheduler preemption frees this req's page_table immediately after moving it + # back to WAITING. Any block still marked IN_CACHE would then point to KV pages + # that no longer belong to the req. Demote IN_CACHE back to TO_CACHE so a + # resumed req rebuilds its cached prefix instead of reusing stale block state. + # TO_CACHE blocks are left untouched. + for block in self.dllm_blocks: + if block.is_in_cache: + block.status = DllmBlockStatus.TO_CACHE🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex/strategy_template/multi_block/engine/request.py` around lines 397 - 401, The preempt demotion loop contains a redundant continue: in the loop over self.dllm_blocks, the check "if rebuild_until > 0 and block.end <= rebuild_until and block.is_to_cache: continue" is unnecessary because block.is_to_cache and block.is_in_cache are mutually exclusive; remove the continue (or replace the condition with a short clarifying comment) so the code simply checks "if block.is_in_cache: block.status = DllmBlockStatus.TO_CACHE" and, if you keep a note, add a comment near the loop that TO_CACHE blocks within rebuild_until are intentionally preserved.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@diffulex/sampler/base/no_shift.py`:
- Around line 18-31: The method _maybe_log_prefill_alignment currently only
tracks req_ids in self._debug_prefill_alignment_logged and never emits logs,
wasting CPU because prefill_block_summaries is always built; fix by either
removing this dead scaffolding or implementing real gated logging: in the hot
path where prefill_block_summaries is constructed (the code that builds
per-block summaries and calls _maybe_log_prefill_alignment), wrap that
construction and the call with a logger.isEnabledFor(logging.DEBUG) check, and
update _maybe_log_prefill_alignment to call logger.debug(...) with a concise
summary (using req_id and the passed block_summaries) while still using
self._debug_prefill_alignment_logged to de-duplicate; reference
_maybe_log_prefill_alignment, prefill_block_summaries (the summary construction
site), and self._debug_prefill_alignment_logged when making the changes.
In `@diffulex/utils/loader.py`:
- Around line 54-57: Normalize target_modules to a sequence before the
comparison: if target_modules is a str (or other non-iterable intended as a
single module), wrap it into a list so the any(...) check in the loader logic
works on module names rather than characters. Update the code surrounding
target_modules (used with name, leaf, rev_mapping and should_apply) to coerce
strings into [target_modules] (or otherwise ensure it's a list/tuple) before
computing effective and should_apply.
---
Outside diff comments:
In `@diffulex/strategy_template/multi_block/engine/request.py`:
- Around line 11-25: The _restore_req_runtime_state method must ensure the
_resume_prefill_until attribute exists after unpickling to avoid AttributeError;
update the _restore_req_runtime_state implementation to initialize
self._resume_prefill_until = getattr(self, "_resume_prefill_until", 0) (or
assign 0 if missing) before any sampler or block logic runs—do this alongside
the existing multi-block handling in _restore_req_runtime_state (the same place
that accesses dllm_blocks and dllm_block_buffer) so it matches initialization in
init_multi_block and the defensive pattern used in no_shift.py.
---
Nitpick comments:
In `@diffulex_bench/configs/llada_instruct_gsm8k.yml`:
- Around line 19-21: The config sets max_num_batched_tokens equal to
max_model_len which forces warmup prefill concurrency to 1 despite max_num_reqs
being 24; either raise max_num_batched_tokens above max_model_len (e.g., to a
multiple of max_model_len) to allow higher prefill concurrency or add a comment
nearby documenting that this equality is intentional to cap concurrency for
stability; update the values or add the explanatory comment referencing the keys
max_model_len, max_num_batched_tokens, and max_num_reqs.
In `@diffulex/strategy_template/multi_block/engine/request.py`:
- Around line 397-401: The preempt demotion loop contains a redundant continue:
in the loop over self.dllm_blocks, the check "if rebuild_until > 0 and block.end
<= rebuild_until and block.is_to_cache: continue" is unnecessary because
block.is_to_cache and block.is_in_cache are mutually exclusive; remove the
continue (or replace the condition with a short clarifying comment) so the code
simply checks "if block.is_in_cache: block.status = DllmBlockStatus.TO_CACHE"
and, if you keep a note, add a comment near the loop that TO_CACHE blocks within
rebuild_until are intentionally preserved.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 7bbaa9ce-b02f-4e4c-a2d5-c1076ad99b9d
📒 Files selected for processing (9)
.gitignorediffulex/sampler/base/no_shift.pydiffulex/sampler/base/shift.pydiffulex/strategy_template/multi_block/engine/request.pydiffulex/utils/loader.pydiffulex_bench/configs/dream_base_gsm8k.ymldiffulex_bench/configs/fast_dllm_v2_gsm8k.ymldiffulex_bench/configs/llada_instruct_gsm8k.ymldiffulex_bench/configs/sdar_chat_gsm8k.yml
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
diffulex/strategy_template/multi_block/engine/request.py (1)
586-598:⚠️ Potential issue | 🔴 CriticalSet a terminal context for every truncation path.
Line 586 only sets
terminal_context_blockfor EOS/max-token/model-len stops, but Line 592 also gates deactivation onmax_nfe_reachedandmax_repetition_run_reached. For those two stop reasons,last_block_finishedstays false becauseterminal_context_blockis stillNone, so the request can livelock instead of completing.🐛 Proposed fix
- if self.eos_token_generated or self.max_new_tokens_reached or self.max_model_len_reached: + if self.is_truncated: self.set_terminal_context_block(self.dllm_block_buffer.last_valid_block)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex/strategy_template/multi_block/engine/request.py` around lines 586 - 598, The code only calls set_terminal_context_block(self.dllm_block_buffer.last_valid_block) when eos_token_generated/max_new_tokens_reached/max_model_len_reached are true, but later the completion check also includes max_nfe_reached and max_repetition_run_reached; ensure terminal_context_block is set for those two stop reasons too. Update the logic around set_terminal_context_block so that if any of eos_token_generated, max_new_tokens_reached, max_model_len_reached, max_nfe_reached, or max_repetition_run_reached is true you call set_terminal_context_block(self.dllm_block_buffer.last_valid_block) before evaluating the final conditional that checks last_block_finished; reference the flags named above and the methods/attributes set_terminal_context_block, terminal_context_block, last_block_finished, and dllm_block_buffer.last_valid_block when making the change.diffulex/engine/engine.py (1)
75-84:⚠️ Potential issue | 🟠 MajorMove tokenizer loading and mask-token normalization before spawning worker processes.
With the
spawncontext,configis serialized whenprocess.start()executes. Worker ranks are spawned on lines 29–31 before the tokenizer is loaded andconfig.mask_token_idis normalized on lines 32–35, causing workers to retain the stale/default mask token ID while rank 0 receives the corrected value.Proposed fix
self.ps = [] self.events = [] + self.tokenizer = AutoTokenizer.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) + ctx = mp.get_context("spawn") for i in range(1, self.model_parallel_world_size): event = ctx.Event() process = ctx.Process(target=AutoModelRunner.from_config, args=(config, i, event)) process.start() self.ps.append(process) self.events.append(event) - self.tokenizer = AutoTokenizer.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)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex/engine/engine.py` around lines 75 - 84, Move the tokenizer initialization and mask-token normalization so they run before spawning worker processes: load AutoTokenizer.from_pretrained and set config.tokenizer_vocab_size, config.eos and call maybe_override_mask_token_id(config, tokenizer) prior to the loop that creates ctx.Event and ctx.Process and starts AutoModelRunner.from_config; this ensures the serialized config passed to each child (via the ctx.Process start loop that appends to self.ps/self.events and targets AutoModelRunner.from_config) contains the corrected mask_token_id and tokenizer-derived fields.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@diffulex/engine/engine.py`:
- Around line 32-40: maybe_override_mask_token_id is unconditionally overwriting
user-provided mask_token_id; detect whether the user explicitly passed
--mask-token-id by checking config_kwargs in DiffulexEngine.__init__ (set a
boolean mask_token_id_explicit) and pass that flag into
maybe_override_mask_token_id; inside maybe_override_mask_token_id, skip the
tokenizer-based override when mask_token_id_explicit is True so CLI/config
values are preserved, and ensure you set/pass the explicit flag before worker
processes are spawned so non-rank-0 runners receive the correct config.
---
Outside diff comments:
In `@diffulex/engine/engine.py`:
- Around line 75-84: Move the tokenizer initialization and mask-token
normalization so they run before spawning worker processes: load
AutoTokenizer.from_pretrained and set config.tokenizer_vocab_size, config.eos
and call maybe_override_mask_token_id(config, tokenizer) prior to the loop that
creates ctx.Event and ctx.Process and starts AutoModelRunner.from_config; this
ensures the serialized config passed to each child (via the ctx.Process start
loop that appends to self.ps/self.events and targets
AutoModelRunner.from_config) contains the corrected mask_token_id and
tokenizer-derived fields.
In `@diffulex/strategy_template/multi_block/engine/request.py`:
- Around line 586-598: The code only calls
set_terminal_context_block(self.dllm_block_buffer.last_valid_block) when
eos_token_generated/max_new_tokens_reached/max_model_len_reached are true, but
later the completion check also includes max_nfe_reached and
max_repetition_run_reached; ensure terminal_context_block is set for those two
stop reasons too. Update the logic around set_terminal_context_block so that if
any of eos_token_generated, max_new_tokens_reached, max_model_len_reached,
max_nfe_reached, or max_repetition_run_reached is true you call
set_terminal_context_block(self.dllm_block_buffer.last_valid_block) before
evaluating the final conditional that checks last_block_finished; reference the
flags named above and the methods/attributes set_terminal_context_block,
terminal_context_block, last_block_finished, and
dllm_block_buffer.last_valid_block when making the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 17fad98f-c049-44e2-a280-cc0a73608e66
📒 Files selected for processing (6)
diffulex/engine/engine.pydiffulex/sampler/base/no_shift.pydiffulex/server/args.pydiffulex/strategy_template/multi_block/engine/request.pydiffulex/utils/loader.pyexamples/streamlit_block_append_chat.py
✅ Files skipped from review due to trivial changes (1)
- examples/streamlit_block_append_chat.py
🚧 Files skipped from review as they are similar to previous changes (1)
- diffulex/utils/loader.py
There was a problem hiding this comment.
♻️ Duplicate comments (1)
diffulex/engine/engine.py (1)
24-55:⚠️ Potential issue | 🟠 MajorPreserve explicit
mask_token_idoverrides from CLI/config.
maybe_override_mask_token_idstill unconditionally replacesconfig.mask_token_idwhen tokenizer/hf_config has a different value, so explicit--mask-token-idcan still be ignored (seediffulex/server/args.pyLine 92-93 passing it through). This was previously flagged and remains unresolved.🐛 Proposed fix
-def maybe_override_mask_token_id(config: Config, tokenizer) -> None: +def maybe_override_mask_token_id( + config: Config, + tokenizer, + *, + mask_token_id_explicit: bool = False, +) -> None: @@ + if mask_token_id_explicit: + return + tokenizer_mask_token_id = getattr(tokenizer, "mask_token_id", None) @@ - maybe_override_mask_token_id(config, self.tokenizer) + maybe_override_mask_token_id( + config, + self.tokenizer, + mask_token_id_explicit="mask_token_id" in config_kwargs, + )Also applies to: 77-77
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex/engine/engine.py` around lines 24 - 55, The function maybe_override_mask_token_id currently unconditionally replaces config.mask_token_id from the tokenizer branch, which lets explicit CLI/config overrides be lost; change both override sites in maybe_override_mask_token_id (the tokenizer branch that sets config.mask_token_id from tokenizer.mask_token_id and the hf_config branch) to only assign when the current config.mask_token_id equals the class default sentinel (Config.mask_token_id). In other words, obtain default_mask_token_id = Config.mask_token_id and wrap the tokenizer override in a guard like "and config.mask_token_id == default_mask_token_id" (the hf_config branch already has this check—ensure the same logic is applied to the tokenizer branch and any duplicate occurrence around the later block).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@diffulex/engine/engine.py`:
- Around line 24-55: The function maybe_override_mask_token_id currently
unconditionally replaces config.mask_token_id from the tokenizer branch, which
lets explicit CLI/config overrides be lost; change both override sites in
maybe_override_mask_token_id (the tokenizer branch that sets
config.mask_token_id from tokenizer.mask_token_id and the hf_config branch) to
only assign when the current config.mask_token_id equals the class default
sentinel (Config.mask_token_id). In other words, obtain default_mask_token_id =
Config.mask_token_id and wrap the tokenizer override in a guard like "and
config.mask_token_id == default_mask_token_id" (the hf_config branch already has
this check—ensure the same logic is applied to the tokenizer branch and any
duplicate occurrence around the later block).
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 2ea6e083-500e-40c3-9eb8-32544fe5a09c
📒 Files selected for processing (11)
diffulex/engine/engine.pydiffulex/strategy_template/multi_block/engine/request.pydiffulex_bench/configs/dream_base_gsm8k.ymldiffulex_bench/configs/llada_instruct_gsm8k.ymldocs/cookbook/benchmark.mddocs/cookbook/index.mddocs/cookbook/models/dream_base.mddocs/cookbook/models/fast_dllm_v2.mddocs/cookbook/models/llada_instruct.mddocs/cookbook/models/sdar_chat.mddocs/cookbook/server.md
✅ Files skipped from review due to trivial changes (7)
- docs/cookbook/models/fast_dllm_v2.md
- docs/cookbook/server.md
- docs/cookbook/benchmark.md
- docs/cookbook/models/llada_instruct.md
- docs/cookbook/models/dream_base.md
- docs/cookbook/index.md
- docs/cookbook/models/sdar_chat.md
🚧 Files skipped from review as they are similar to previous changes (2)
- diffulex_bench/configs/llada_instruct_gsm8k.yml
- diffulex_bench/configs/dream_base_gsm8k.yml
There was a problem hiding this comment.
♻️ Duplicate comments (1)
diffulex/engine/engine.py (1)
24-33:⚠️ Potential issue | 🟡 MinorExplicit
mask_token_idis still ambiguous when user passes the default value.The helper infers “explicitness” by value (
config.mask_token_id == Config.mask_token_id), so an explicitly provided default can still be overridden. Please pass an explicit boolean fromconfig_kwargsto preserve user intent in all cases.🐛 Minimal fix
-def maybe_override_mask_token_id(config: Config, tokenizer) -> None: +def maybe_override_mask_token_id( + config: Config, + tokenizer, + *, + mask_token_id_explicit: bool = False, +) -> None: @@ + if mask_token_id_explicit: + return + default_mask_token_id = Config.mask_token_id @@ - maybe_override_mask_token_id(config, self.tokenizer) + maybe_override_mask_token_id( + config, + self.tokenizer, + mask_token_id_explicit="mask_token_id" in config_kwargs, + )Also applies to: 35-42, 85-85
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@diffulex/engine/engine.py` around lines 24 - 33, The current maybe_override_mask_token_id logic infers whether the user explicitly set mask_token_id by comparing values, which misclassifies an explicitly-provided default; change the function to accept an explicit boolean flag (e.g., mask_token_id_explicit) passed from config_kwargs and use that flag instead of value equality to decide whether to skip tokenizer/HF overrides; update the call sites that construct Config (where config_kwargs are built) to set this boolean when the user provided mask_token_id, and apply the same change to the other similar checks referenced (the duplicate logic around lines 35-42 and the check at 85) to use the explicit boolean parameter so user intent is preserved.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@diffulex/engine/engine.py`:
- Around line 24-33: The current maybe_override_mask_token_id logic infers
whether the user explicitly set mask_token_id by comparing values, which
misclassifies an explicitly-provided default; change the function to accept an
explicit boolean flag (e.g., mask_token_id_explicit) passed from config_kwargs
and use that flag instead of value equality to decide whether to skip
tokenizer/HF overrides; update the call sites that construct Config (where
config_kwargs are built) to set this boolean when the user provided
mask_token_id, and apply the same change to the other similar checks referenced
(the duplicate logic around lines 35-42 and the check at 85) to use the explicit
boolean parameter so user intent is preserved.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a1f832dc-d836-4e03-a73e-4c0f992d3fd4
📒 Files selected for processing (6)
diffulex/engine/engine.pydocs/cookbook/models/dream_base.mddocs/cookbook/models/fast_dllm_v2.mddocs/cookbook/models/llada_instruct.mddocs/cookbook/models/sdar_chat.mddocs/cookbook/server.md
✅ Files skipped from review due to trivial changes (3)
- docs/cookbook/models/llada_instruct.md
- docs/cookbook/models/fast_dllm_v2.md
- docs/cookbook/models/dream_base.md
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/cookbook/models/sdar_chat.md
- docs/cookbook/server.md
…uest 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.
Osc-7
force-pushed
the
fix-llada-lora-bug
branch
from
June 24, 2026 19:24
4550c0b to
92cc4e6
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@diffulex/engine/engine.py`:
- Around line 376-389: Validate the `sampling_params` input in `engine.py`
before the `zip(prompts, sampling_params)` loop inside the request scheduling
logic. The issue is that a list whose length differs from `prompts` will be
silently truncated by `zip`, so add an explicit length check after the
`sampling_params` normalization and before building `req_id_to_prompt_id` in the
scheduling path. If the lengths do not match, raise a clear error from this
request-adding flow so every prompt has exactly one corresponding
`sampling_params` entry.
- Around line 91-98: Refresh the derived runtime config after
tokenizer-dependent mutations in Engine initialization. In `Engine.__init__`,
after `auto_tokenizer_from_pretrained`, `maybe_override_mask_token_id`, and the
updates to `config.tokenizer_vocab_size` and `config.eos`, also propagate the
final values into `config.runtime.model` so `tokenizer_vocab_size`, `eos`, and
`mask_token_id` stay in sync for consumers of `config.runtime`.
In `@diffulex/sampler/base/no_shift.py`:
- Line 72: The zip loop in the sampling path is vulnerable to silently
truncating mismatched inputs; update the loop in the no_shift sampler’s request
iteration to use zip with strict checking so temperature, request, and logits
lengths must match. Make this change in the enumerate(zip(...)) loop inside the
relevant sampling function in no_shift.py, ensuring any mismatch fails fast
instead of producing partial output.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 95882a22-7886-4aaf-bac5-3125a01082e1
📒 Files selected for processing (6)
diffulex/engine/engine.pydiffulex/engine/request.pydiffulex/sampler/base/no_shift.pydiffulex/sampler/base/shift.pydiffulex/server/args.pydiffulex/utils/loader.py
🚧 Files skipped from review as they are similar to previous changes (3)
- diffulex/server/args.py
- diffulex/utils/loader.py
- diffulex/sampler/base/shift.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@diffulex/engine/engine.py`:
- Around line 376-389: Validate the `sampling_params` input in `engine.py`
before the `zip(prompts, sampling_params)` loop inside the request scheduling
logic. The issue is that a list whose length differs from `prompts` will be
silently truncated by `zip`, so add an explicit length check after the
`sampling_params` normalization and before building `req_id_to_prompt_id` in the
scheduling path. If the lengths do not match, raise a clear error from this
request-adding flow so every prompt has exactly one corresponding
`sampling_params` entry.
- Around line 91-98: Refresh the derived runtime config after
tokenizer-dependent mutations in Engine initialization. In `Engine.__init__`,
after `auto_tokenizer_from_pretrained`, `maybe_override_mask_token_id`, and the
updates to `config.tokenizer_vocab_size` and `config.eos`, also propagate the
final values into `config.runtime.model` so `tokenizer_vocab_size`, `eos`, and
`mask_token_id` stay in sync for consumers of `config.runtime`.
In `@diffulex/sampler/base/no_shift.py`:
- Line 72: The zip loop in the sampling path is vulnerable to silently
truncating mismatched inputs; update the loop in the no_shift sampler’s request
iteration to use zip with strict checking so temperature, request, and logits
lengths must match. Make this change in the enumerate(zip(...)) loop inside the
relevant sampling function in no_shift.py, ensuring any mismatch fails fast
instead of producing partial output.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 95882a22-7886-4aaf-bac5-3125a01082e1
📒 Files selected for processing (6)
diffulex/engine/engine.pydiffulex/engine/request.pydiffulex/sampler/base/no_shift.pydiffulex/sampler/base/shift.pydiffulex/server/args.pydiffulex/utils/loader.py
🚧 Files skipped from review as they are similar to previous changes (3)
- diffulex/server/args.py
- diffulex/utils/loader.py
- diffulex/sampler/base/shift.py
🛑 Comments failed to post (3)
diffulex/engine/engine.py (2)
91-98: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Refresh derived config after tokenizer overrides.
Config.__post_init__has already builtconfig.runtimebefore these tokenizer-derived mutations. Workers receive the mutatedconfig, but consumers readingconfig.runtime.model.*can still see staletokenizer_vocab_size,eos, or defaultmask_token_id.Proposed fix
maybe_override_mask_token_id( config, self.tokenizer, mask_token_id_explicit="mask_token_id" in config_kwargs, ) + hf_config = getattr(config, "hf_config", None) + for target in (hf_config, getattr(hf_config, "text_config", None)): + if target is not None: + setattr(target, "mask_token_id", config.mask_token_id) + config._build_runtime_config() ctx = mp.get_context("spawn")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.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, ) hf_config = getattr(config, "hf_config", None) for target in (hf_config, getattr(hf_config, "text_config", None)): if target is not None: setattr(target, "mask_token_id", config.mask_token_id) config._build_runtime_config()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@diffulex/engine/engine.py` around lines 91 - 98, Refresh the derived runtime config after tokenizer-dependent mutations in Engine initialization. In `Engine.__init__`, after `auto_tokenizer_from_pretrained`, `maybe_override_mask_token_id`, and the updates to `config.tokenizer_vocab_size` and `config.eos`, also propagate the final values into `config.runtime.model` so `tokenizer_vocab_size`, `eos`, and `mask_token_id` stay in sync for consumers of `config.runtime`.
376-389: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate
sampling_paramslength before zipping.When a list length differs from
prompts,zip()silently drops extras, leaving missing outputs or ignored params.Proposed fix
if not isinstance(sampling_params, list): sampling_params = [sampling_params] * len(prompts) + elif len(sampling_params) != len(prompts): + raise ValueError( + "sampling_params length must match prompts length, " + f"got {len(sampling_params)} and {len(prompts)}." + ) total_expected_tokens = sum(max(getattr(sp, "max_tokens", 0) or 0, 0) for sp in sampling_params)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.if not isinstance(sampling_params, list): sampling_params = [sampling_params] * len(prompts) elif len(sampling_params) != len(prompts): raise ValueError( "sampling_params length must match prompts length, " f"got {len(sampling_params)} and {len(prompts)}." ) total_expected_tokens = sum(max(getattr(sp, "max_tokens", 0) or 0, 0) for sp in sampling_params) req_id_to_prompt_id = {} for prompt_id, (prompt, sp) in tqdm( enumerate(zip(prompts, sampling_params)), total=len(prompts), desc="Adding Requests to Scheduler", dynamic_ncols=True, disable=not use_tqdm, leave=False, mininterval=1.0, ):🧰 Tools
🪛 Ruff (0.15.18)
[warning] 382-382:
zip()without an explicitstrict=parameterAdd explicit value for parameter
strict=(B905)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@diffulex/engine/engine.py` around lines 376 - 389, Validate the `sampling_params` input in `engine.py` before the `zip(prompts, sampling_params)` loop inside the request scheduling logic. The issue is that a list whose length differs from `prompts` will be silently truncated by `zip`, so add an explicit length check after the `sampling_params` normalization and before building `req_id_to_prompt_id` in the scheduling path. If the lengths do not match, raise a clear error from this request-adding flow so every prompt has exactly one corresponding `sampling_params` entry.Source: Linters/SAST tools
diffulex/sampler/base/no_shift.py (1)
72-72: 🎯 Functional Correctness | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Verify Python/tooling target supports zip(..., strict=...) fd -HI "pyproject.toml|setup.cfg|tox.ini|.python-version|runtime.txt|.tool-versions" . \ | xargs -I{} sh -c 'echo "### {}"; sed -n "1,220p" "{}"'Repository: SJTU-DENG-Lab/Diffulex
Length of output: 2197
🏁 Script executed:
#!/bin/bash set -euo pipefail # Inspect the relevant file and nearby call sites for length alignment. file="diffulex/sampler/base/no_shift.py" wc -l "$file" sed -n '1,180p' "$file" printf '\n--- call sites / related helpers ---\n' rg -n "split_logits|temperatures|reqs" diffulex/sampler -SRepository: SJTU-DENG-Lab/Diffulex
Length of output: 15849
🏁 Script executed:
#!/bin/bash set -euo pipefail # Check whether Python 3.11+ is the project target and whether strict zip is available in that target. python3 - <<'PY' import sys print(sys.version) print("zip strict supported:", sys.version_info >= (3, 10)) PYRepository: SJTU-DENG-Lab/Diffulex
Length of output: 237
🏁 Script executed:
#!/bin/bash set -euo pipefail sed -n '1,180p' diffulex/sampler/base/no_shift.py printf '\n---\n' rg -n "split_logits|temperatures|reqs" diffulex/sampler -SRepository: SJTU-DENG-Lab/Diffulex
Length of output: 15778
🏁 Script executed:
#!/bin/bash set -euo pipefail # Inspect the sampler call chain and how temperatures are shaped/passed in. rg -n "temperatures\s*=" diffulex -S printf '\n--- base/shift.py ---\n' sed -n '1,140p' diffulex/sampler/base/shift.py printf '\n--- llada2.py relevant sections ---\n' sed -n '250,360p' diffulex/sampler/llada2.py sed -n '520,590p' diffulex/sampler/llada2.pyRepository: SJTU-DENG-Lab/Diffulex
Length of output: 14590
Use
strict=Truein this zip loop.
zip(temperatures, reqs, split_logits)can silently drop requests if one sequence is shorter, which would produce partial sampling output. This codebase targets Python 3.11+, sostrict=Trueis available here.Suggested patch
- for idx, (temperature, req, req_logits) in enumerate(zip(temperatures, reqs, split_logits)): + for idx, (temperature, req, req_logits) in enumerate( + zip(temperatures, reqs, split_logits, strict=True) + ):📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.for idx, (temperature, req, req_logits) in enumerate( zip(temperatures, reqs, split_logits, strict=True) ):🧰 Tools
🪛 Ruff (0.15.18)
[warning] 72-72:
zip()without an explicitstrict=parameterAdd explicit value for parameter
strict=(B905)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@diffulex/sampler/base/no_shift.py` at line 72, The zip loop in the sampling path is vulnerable to silently truncating mismatched inputs; update the loop in the no_shift sampler’s request iteration to use zip with strict checking so temperature, request, and logits lengths must match. Make this change in the enumerate(zip(...)) loop inside the relevant sampling function in no_shift.py, ensuring any mismatch fails fast instead of producing partial output.Source: Linters/SAST tools
Contributor
|
lgtm |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR fixes multiple instability issues in D2F LoRA decoding, especially under preempt + rebuild scenarios.
The main problems were caused by inconsistent indexing and request state handling during prefill and decode phases, which could lead to incorrect logits alignment and unstable generation.
Key Changes
LoRA
packed_modules_mappingD2F / Decode flow
Sampler
Request / KV cache
Results
Notes
Checklist
Summary by CodeRabbit
New Features
Bug Fixes