"vscode:/vscode.git/clone" did not exist on "5be3c06485003425e0a6892fac4fc33157d47ab3"
Unverified Commit 899453ac authored by jianan-gu's avatar jianan-gu Committed by GitHub
Browse files

Use explicit uint64 dtype for Tensor data_ptr() to avoid overflow (#11994)

parent ce832d70
...@@ -89,6 +89,7 @@ def write_cache_indices( ...@@ -89,6 +89,7 @@ def write_cache_indices(
prefix_pointers = torch.tensor( prefix_pointers = torch.tensor(
[t.data_ptr() for t in prefix_tensors], [t.data_ptr() for t in prefix_tensors],
device=req_to_token_pool.device, device=req_to_token_pool.device,
dtype=torch.uint64,
) )
# TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start) # TODO: some tensors can be reused for ForwardBatchInfo (e.g., extend_lens, cumsum_start)
write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)]( write_req_to_token_pool_triton[(req_pool_indices_tensor.shape[0],)](
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment