Skip to content

GitLab

Commit 0ec3aa4b authored by zhuwenwen's avatar zhuwenwen
Browse files

Merge branch 'v0.6.2-dev' of ssh://10.6.10.68:10022/dcutoolkit/deeplearing/vllm into v0.6.2-dev

parents 305b5a09 dab3026c
Show whitespace changes
Inline Side-by-side
Showing with 5 additions and 4 deletions
vllm/worker/model_runner.py
View file @ 0ec3aa4b
...@@ -855,11 +855,12 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]): ...@@ -855,11 +855,12 @@ class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):
tree_attention_masks_list = [] tree_attention_masks_list = []
for inter_data in self.inter_data_list: for inter_data in self.inter_data_list:
for i in range(len(inter_data.seq_lens)): for i in range(len(inter_data.seq_lens)):
if inter_data.tree_attn_masks:
tree_attn_masks = inter_data.tree_attn_masks[i] tree_attn_masks = inter_data.tree_attn_masks[i]
if tree_attn_masks is not None: if tree_attn_masks is not None:
tree_attention_masks_list.append(tree_attn_masks) tree_attention_masks_list.append(tree_attn_masks)
tree_attention_masks_tensor = None tree_attention_masks_tensor = None
if len(tree_attention_masks_list) > 0: if tree_attention_masks_list:
tree_attention_masks_tensor = torch.stack(tree_attention_masks_list, dim=0) tree_attention_masks_tensor = torch.stack(tree_attention_masks_list, dim=0)
tree_attention_masks_tensor = tree_attention_masks_tensor.contiguous() tree_attention_masks_tensor = tree_attention_masks_tensor.contiguous()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment