perf: 消除 PagedCache.write 中的 position_ids GPU 同步,解码提速 15%
- CacheView.write 用 total_len - k.size(1) 推导 start_pos,替代 position_ids[0,0].item() - 移除 GQA/MLA/DecoderBlock 中不再使用的 position_ids 参数 - PagedCache.write 参数 position_ids:Tensor → start_pos:int
This commit is contained in:
@@ -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, torch.zeros(1, 2, dtype=torch.long), k, v)
|
||||
cache.write(0, page_table, 0, 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, torch.zeros(1, 8, dtype=torch.long), k, v)
|
||||
cache.write(0, page_table, 0, 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, torch.zeros(1, 6, dtype=torch.long), k, v)
|
||||
cache.write(0, page_table, 0, k, v)
|
||||
|
||||
gk, gv = cache.gather(0, page_table, 5)
|
||||
assert gk.shape == (1, 5, 2, 8)
|
||||
|
||||
Reference in New Issue
Block a user