fix: prefill 时 attention mask 长度不足导致 expand 崩溃

- astrai/inference/scheduler.py: prefill input_mask 由 [batch, seq_len] 改为 [batch, prompt_len]，覆盖全部 KV 位置 - astrai/model/transformer.py: seq_mask is None 分支补全 start_pos + seq_len 列，避免 expand 非 singleton 维度不匹配
2026-05-10 19:56:41 +08:00
parent 3da428e0e4
commit c95ace41aa
2 changed files with 2 additions and 2 deletions
@@ -29,7 +29,7 @@ def process_attention_mask(

    if seq_mask is None:
        if start_pos != 0:
-            seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
+            seq_mask = torch.ones((1, start_pos + seq_len), dtype=torch.bool, device=device)
        else:
            return None