已合并
Fix potential integer overflow issues in npu_paged_cache_load #3413
wang-guangbin创建于 2025年10月31日
Fix potential integer overflow issues in npu_paged_cache_load #3413
已合并
共 2 个文件变更+5-3
| @@ -46,7 +46,11 @@ std::tuple<at::Tensor, at::Tensor> npu_paged_cache_load( | |||
| 46 | context_lens.numel() > 0 ? context_lens[-1].item<int32_t>() : 0; | 46 | context_lens.numel() > 0 ? context_lens[-1].item<int32_t>() : 0; |
| 47 | } else { | 47 | } else { |
| 48 | for (int i = 0; i < context_lens.numel(); i++) { | 48 | for (int i = 0; i < context_lens.numel(); i++) { |
王 | |||
| 49 | - num_tokens += context_lens[i].item<int32_t>(); | 49 | + int32_t context_val = context_lens[i].item<int32_t>(); |
| 50 | + TORCH_CHECK(context_val >= 0, "Invalid context_lens: negative value encountered"); | ||
| 51 | + TORCH_CHECK(num_tokens <= INT32_MAX - context_val, | ||
| 52 | + "Integer overflow in accumulation: sum exceeds int32_t max in npu_paged_cache_load"); | ||
| 53 | + num_tokens += context_val; | ||
| 50 | } | 54 | } |
| 51 | } | 55 | } |
| 52 | at::Tensor key = | 56 | at::Tensor key = |
| @@ -185,7 +185,6 @@ class TestPagedCacheLoadSeqStarts(TestCase): | |||
| 185 | ) | 185 | ) |
| 186 | 186 | ||
| 187 | 187 | ||
| 188 | - | ||
| 189 | def test_atb_paged_cache_load_out(self): | 188 | def test_atb_paged_cache_load_out(self): |
| 190 | kv_lora_rank = 512 | 189 | kv_lora_rank = 512 |
| 191 | qk_rope_head_dim = 64 | 190 | qk_rope_head_dim = 64 |
| @@ -243,7 +242,6 @@ class TestPagedCacheLoadSeqStarts(TestCase): | |||
| 243 | self.assertRtolEqual(expected, torch_npu_result) | 242 | self.assertRtolEqual(expected, torch_npu_result) |
| 244 | 243 | ||
| 245 | 244 | ||
| 246 | - | ||
| 247 | def test_atb_paged_cache_load(self): | 245 | def test_atb_paged_cache_load(self): |
| 248 | kv_lora_rank = 512 | 246 | kv_lora_rank = 512 |
| 249 | qk_rope_head_dim = 64 | 247 | qk_rope_head_dim = 64 |
这里for循环内部的代码还是都保持下退4格吧，保持编码风格对齐