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

Avoid host synchronization in DeepSeek V4.1 CP prefill

Browse Source 
(cherry picked from commit 510596019fd5782d494c9b5227bf784d46661f6f)
This commit is contained in:
abing
2026-09-20 22:07:30 +08:00
committed by minke.yu
parent 12e3b82e52
commit 4f22146e51
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+5 -3
View File
@@ -2914,7 +2914,7 @@ class DeepseekV4AttnBackend(
q_lora[:num_local], q_lora[:num_local],
positions[:num_local].to(torch.int64), positions[:num_local].to(torch.int64),
forward_batch, forward_batch,
torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device), self._move_to_device(q_lens_cpu),
q_lens_cpu, q_lens_cpu,
) )
@@ -3280,7 +3280,9 @@ class DeepseekV4AttnBackend(
continue continue
j = torch.arange(lc, device=device) j = torch.arange(lc, device=device)
slot_chunks.append( slot_chunks.append(
self.req_to_token[req_pool_indices[r], j * ratio].to(torch.int64) self.req_to_token[req_pool_indices[r : r + 1], j * ratio].to(
torch.int64
)
// ratio // ratio
) )
start += lc start += lc
@@ -3304,7 +3306,7 @@ class DeepseekV4AttnBackend(
weights = indexer.head_weights(x).float() weights = indexer.head_weights(x).float()
compress_lens = ((pos + 1) // ratio).to(torch.int32) compress_lens = ((pos + 1) // ratio).to(torch.int32)
ks = torch.repeat_interleave( ks = torch.repeat_interleave(
torch.tensor(starts, dtype=torch.int32, device=device), self._move_to_device(starts),
q_lens.to(torch.int64), q_lens.to(torch.int64),
output_size=num_tokens, output_size=num_tokens,
) )