chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Minor code style fix for dllm (#14836)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-12-11 10:35:42 +09:00
committed by GitHub
parent 312df1d6c0
commit bd7824b24d
3 changed files with 43 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/sglang/srt/managers/schedule_batch.py
View File

@@ -710,7 +710,6 @@ class Req:
self.dimensions = dimensions
# For diffusion LLM
self.dllm_ids = []
self.dllm_block_offset = 0
self.dllm_config = dllm_config
@@ -786,22 +785,19 @@ class Req:
def is_dllm(self):
return self.dllm_config is not None
def _init_fill_ids_for_dllm(self):
if not self.fill_ids:
self.fill_ids = (
self.origin_input_ids
+ [self.dllm_config.mask_id] * self.dllm_config.block_size
)
else:
self.dllm_block_offset += self.dllm_config.block_size
self.fill_ids += [self.dllm_config.mask_id] * self.dllm_config.block_size
def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None):
if self.is_dllm():
if not self.fill_ids:
self.dllm_ids = (
self.origin_input_ids
+ [
self.dllm_config.mask_id,
]
* self.dllm_config.block_size
)
else:
self.dllm_block_offset += self.dllm_config.block_size
self.dllm_ids += [
self.dllm_config.mask_id
] * self.dllm_config.block_size
self.fill_ids = self.dllm_ids
self._init_fill_ids_for_dllm()
else:
self.fill_ids = self.origin_input_ids + self.output_ids
@@ -1322,9 +1318,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
), f"Expected {len(self.out_cache_loc)}, got {self.extend_num_tokens}"
def prepare_for_extend(self):
self.forward_mode = (
ForwardMode.DLLM_EXTEND if self.is_dllm() else ForwardMode.EXTEND
)
self.forward_mode = ForwardMode.EXTEND
if self.is_dllm():
# For DLLM, we use a separate forward mode
self.forward_mode = ForwardMode.DLLM_EXTEND
# Init tensors
reqs = self.reqs