Minor code cleanup / improvement for PREBUILT_EXTEND mode (#12948)

This commit is contained in:
Liangsheng Yin
2025-11-10 11:59:54 +08:00
committed by GitHub
parent 83f89cc615
commit 61c7fe7aed
7 changed files with 50 additions and 55 deletions

View File

@@ -79,7 +79,7 @@ def compute_split_seq_index(
elif forward_mode.is_target_verify() or forward_mode.is_decode():
assert token_num_per_seq is not None
return (num_tokens // token_num_per_seq) // 2
elif forward_mode.is_idle() or forward_mode.is_prebuilt_extend():
elif forward_mode.is_idle() or forward_mode.is_prebuilt():
assert num_tokens == 0
return 0
else:
@@ -381,7 +381,7 @@ class TboDPAttentionPreparer:
or local_batch.forward_mode.is_decode()
):
num_tokens = local_batch.batch_size() * token_num_per_seq
elif local_batch.forward_mode.is_prebuilt_extend():
elif local_batch.forward_mode.is_prebuilt():
num_tokens = 0
else:
num_tokens = local_batch.extend_num_tokens