From 4f22146e510020ac61ec50b82704d97416cb8d8a Mon Sep 17 00:00:00 2001 From: abing Date: Fri, 18 Sep 2026 02:28:12 -0700 Subject: [PATCH] Avoid host synchronization in DeepSeek V4.1 CP prefill (cherry picked from commit 510596019fd5782d494c9b5227bf784d46661f6f) --- python/sglang/srt/layers/attention/deepseek_v4_backend.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 563a829f0..980a7303f 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -2914,7 +2914,7 @@ class DeepseekV4AttnBackend( q_lora[:num_local], positions[:num_local].to(torch.int64), forward_batch, - torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device), + self._move_to_device(q_lens_cpu), q_lens_cpu, ) @@ -3280,7 +3280,9 @@ class DeepseekV4AttnBackend( continue j = torch.arange(lc, device=device) 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 ) start += lc @@ -3304,7 +3306,7 @@ class DeepseekV4AttnBackend( weights = indexer.head_weights(x).float() compress_lens = ((pos + 1) // ratio).to(torch.int32) ks = torch.repeat_interleave( - torch.tensor(starts, dtype=torch.int32, device=device), + self._move_to_device(starts), q_lens.to(torch.int64), output_size=num_tokens, )