From 2f745ff14669af7ecff6f5ea1c0df42f3eac0e66 Mon Sep 17 00:00:00 2001 From: "q.yao" Date: Sun, 18 Jan 2026 21:00:56 +0800 Subject: [PATCH] fix dllm mask on set_step --- lmdeploy/pytorch/strategies/dllm/sequence.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lmdeploy/pytorch/strategies/dllm/sequence.py b/lmdeploy/pytorch/strategies/dllm/sequence.py index 0348e8bae5..03ad19e75d 100644 --- a/lmdeploy/pytorch/strategies/dllm/sequence.py +++ b/lmdeploy/pytorch/strategies/dllm/sequence.py @@ -200,7 +200,8 @@ def set_step(self, step: int): start = min(step, self.num_history_ids) end = self.num_history_ids if end > start: - self.history_dllm_mask[start:end] = DLLM_MASKED + to_change_mask = self.history_dllm_mask[start:] + to_change_mask[to_change_mask == DLLM_CACHED] = DLLM_UNMASKED super().set_step(step)