refactor: 位置编码改用 position_ids [B,S]，简化 attention mask 构建

- RotaryEmbedding/CacheView 接受 position_ids 替代 start_pos

- process_attention_mask 用 position_ids >= arange 做逐位置 causal

- 训练/无 KV cache 时 position_ids=None 内部自动处理

- 移除 executor/benchmark 中冗余的 input_mask 构造

astrai/inference/cache.py

@@ -241,13 +241,14 @@ class PagedCache:
         self,
         layer_id: int,
         page_table: Tensor,
-        start_pos: int,
+        position_ids: Tensor,
         k: Tensor,
         v: Tensor,
     ) -> None:
         seq_len = k.size(1)
         if seq_len == 0:
             return
+        start_pos = position_ids[0, 0].item()
         page_size = self.page_size
         written = 0
         first_page = start_pos // page_size
@@ -288,8 +289,8 @@ class CacheView:
         self._page_table = page_table
         self._total_len = total_len
 
-    def write(self, layer_id: int, start_pos: int, k: Tensor, v: Tensor) -> None:
-        self._cache.write(layer_id, self._page_table, start_pos, k, v)
+    def write(self, layer_id: int, position_ids: Tensor, k: Tensor, v: Tensor) -> None:
+        self._cache.write(layer_id, self._page_table, position_ids, k, v)
 
     def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
         return self._cache.gather(layer_id, self._page_table, self._total_len)

astrai/inference/executor.py

@@ -40,9 +40,6 @@ class Executor:
 
         seq_len = prompt_len - start_pos
         input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
-        input_mask = torch.ones(
-            batch_sz, prompt_len, dtype=torch.bool, device=self.device
-        )
 
         for i, t in enumerate(tasks):
             input_ids[i] = torch.tensor(
@@ -55,8 +52,30 @@ class Executor:
         with torch.inference_mode():
             self.model(
                 input_ids,
-                input_mask=input_mask,
-                start_pos=start_pos,
+                position_ids=torch.arange(
+                    start_pos, prompt_len, dtype=torch.long, device=self.device
+                )
+                .unsqueeze(0)
+                .expand(batch_sz, -1),
+                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
+            )
+
+        for i, t in enumerate(tasks):
+            input_ids[i] = torch.tensor(
+                t.prompt_ids[start_pos:prompt_len], device=self.device
+            )
+
+        task_ids = [t.task_id for t in tasks]
+        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
+
+        with torch.inference_mode():
+            self.model(
+                input_ids,
+                position_ids=torch.arange(
+                    start_pos, prompt_len, dtype=torch.long, device=self.device
+                )
+                .unsqueeze(0)
+                .expand(batch_sz, -1),
                 paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
             )
 
@@ -72,8 +91,6 @@ class Executor:
             device=self.device,
         )
 
-        active_mask = torch.ones((batch_sz, 1), dtype=torch.bool, device=self.device)
-
         task_ids = [t.task_id for t in tasks]
         page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
         total_len = start_pos + 1
@@ -85,9 +102,10 @@ class Executor:
         with torch.inference_mode():
             outputs = self.model(
                 input_ids.unsqueeze(1),
-                input_mask=active_mask,
                 paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
-                start_pos=start_pos,
+                position_ids=torch.full(
+                    (batch_sz, 1), start_pos, dtype=torch.long, device=self.device
+                ),
             )
             logits = outputs["logits"][:, -1, :]