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 构造
This commit is contained in:
2026-05-14 13:26:31 +08:00
parent df0845e916
commit c0effc9f5b
6 changed files with 104 additions and 76 deletions
+3 -3
View File
@@ -244,7 +244,7 @@ def test_paged_cache_write_gather_single_page():
k = torch.randn(1, 2, 2, 8)
v = torch.randn(1, 2, 2, 8)
cache.write(0, page_table, 0, k, v)
cache.write(0, page_table, torch.zeros(1, 2, dtype=torch.long), k, v)
gk, gv = cache.gather(0, page_table, 2)
assert torch.allclose(gk, k)
@@ -263,7 +263,7 @@ def test_paged_cache_write_cross_page():
k = torch.randn(1, 8, 2, 8)
v = torch.randn(1, 8, 2, 8)
cache.write(0, page_table, 0, k, v)
cache.write(0, page_table, torch.zeros(1, 8, dtype=torch.long), k, v)
gk, gv = cache.gather(0, page_table, 8)
assert torch.allclose(gk, k)
@@ -281,7 +281,7 @@ def test_paged_cache_gather_truncates_to_total_len():
page_table = torch.tensor([[0, 1]], dtype=torch.long)
k = torch.randn(1, 6, 2, 8)
v = torch.randn(1, 6, 2, 8)
cache.write(0, page_table, 0, k, v)
cache.write(0, page_table, torch.zeros(1, 6, dtype=torch.long), k, v)
gk, gv = cache.gather(0, page_table, 5)
assert gk.shape == (1, 5, 2, 8)