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

[AMD] Fix Tensor Memory Aliasing (#19928)

Browse Source
This commit is contained in:
Bingxu Chen
2026-03-07 08:06:10 -08:00
committed by GitHub
parent 7da590d4d0
commit 17721b00fd
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+3 -3
View File
@@ -313,8 +313,8 @@ class Indexer(MultiPlatformOp):
q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope) q_rope, k_rope = self.rotary_emb(positions, q_rope, k_rope)
query[..., : self.rope_head_dim] = q_rope query[..., : self.rope_head_dim] = q_rope.clone()
key[..., : self.rope_head_dim] = k_rope key[..., : self.rope_head_dim] = k_rope.clone()
if enable_dual_stream: if enable_dual_stream:
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
@@ -352,7 +352,7 @@ class Indexer(MultiPlatformOp):
) )
_, k_rope = self.rotary_emb(positions, k_rope, k_rope) _, k_rope = self.rotary_emb(positions, k_rope, k_rope)
key[..., : self.rope_head_dim] = k_rope key[..., : self.rope_head_dim] = k_rope.clone()
key = rotate_activation(key) key = rotate_activation(key)
return key return key
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+1 -1
View File
@@ -424,7 +424,7 @@ class DeepseekMHAForwardMixin:
) )
else: else:
latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1) latent_cache[:, :, : self.kv_lora_rank] = kv_a.unsqueeze(1)
latent_cache[:, :, self.kv_lora_rank :] = k_pe latent_cache[:, :, self.kv_lora_rank :] = k_pe.clone()
# Save latent cache # Save latent cache
forward_batch.token_to_kv_pool.set_kv_buffer( forward_batch.token_to_kv_pool.set_kv_buffer(