minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CP] Register KV cache allgather buffer with symmetric memory (#24040)

Browse Source 
Signed-off-by: wangfakang <fakangwang@gmail.com>
This commit is contained in:
sky
2026-05-06 23:24:36 +08:00
committed by GitHub
parent c4c5541618
commit bfc1aeae13
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/utils/cp_utils.py
+9 -6
View File
@@ -179,12 +179,15 @@ def cp_all_gather_reorganized_into_tensor_kv_cache(
input_tensor = F.pad(input_tensor, padding, mode="constant", value=0) input_tensor = F.pad(input_tensor, padding, mode="constant", value=0)
# Create output tensor with proper shape for all dimensions # Create output tensor with proper shape for all dimensions
input_tensor_full = torch.empty( with use_symmetric_memory(
max_len * cp_size, get_attention_cp_group(), disabled=not is_allocation_symmetric()
*input_tensor.shape[1:], ):
device=input_tensor.device, input_tensor_full = torch.empty(
dtype=input_tensor.dtype, max_len * cp_size,
) *input_tensor.shape[1:],
device=input_tensor.device,
dtype=input_tensor.dtype,
)
get_attention_cp_group().cp_all_gather_into_tensor_async( get_attention_cp_group().cp_all_gather_into_tensor_async(
input_tensor_full, input_tensor, stream input_tensor_full, input_tensor, stream