A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 14:45 UTC
[Bugfix][ROCm] Fix Accuracy Drop in Sparse Indexer on gfx950 (#43781)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>
kliuae committed
May 28, 2026 at 18:37 UTC
a583c84e2be9b99376dacee0e455ad65ed8fafb7
2 files changed
+14
-4
vllm/model_executor/models/deepseek_v2.py
+9
-1
@@ -612,6 +612,7 @@ class Indexer(nn.Module):
612
cache_config: CacheConfig | None,
613
topk_indices_buffer: torch.Tensor | None,
614
prefix: str = "",
615
+ is_inplace_rope: bool = False,
616
):
617
super().__init__()
618
self.vllm_config = vllm_config
@@ -673,15 +674,21 @@ class Indexer(nn.Module):
674
self.topk_indices_buffer,
675
)
676
677
+ self.is_inplace_rope = is_inplace_rope
678
+
679
def forward(
680
self, hidden_states: torch.Tensor, qr: torch.Tensor, positions, rotary_emb
681
) -> torch.Tensor:
682
q, _ = self.wq_b(qr)
683
q = q.view(-1, self.n_head, self.head_dim)
684
682
- if current_platform.is_rocm():
685
+ if current_platform.is_rocm() and self.is_inplace_rope:
686
# This path should works on all platform, will remove extra
687
# branches in the future
688
+ # This fast path relies on rotary_emb mutating q and k inplace.
689
+ # On ROCm, this is only valid for kernels used as custom ops.
690
+ # In pytorch-native rope for inductor fusion, rotated q/k tensors
691
+ # are not mutated inplace but returned as new tensors.
692
# Fused wk + weights_proj: one GEMM, then split
693
kw, _ = self.wk_weights_proj(hidden_states)
694
k = kw[:, : self.head_dim]
@@ -1002,6 +1009,7 @@ class DeepseekV2MLAAttention(nn.Module):
1009
cache_config,
1010
topk_indices_buffer,
1011
f"{prefix}.indexer",
1012
+ is_inplace_rope=self.indexer_rope_emb.enabled(),
1013
)
1014
1015
# Enable IndexCache for DeepSeek models to reduce redundant top-k
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+5
-3
@@ -107,6 +107,7 @@ def indexer_k_quant_and_cache_triton(
107
kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype)
108
kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32)
109
head_tile_size = head_tile_size // kv_cache.element_size()
110
+ layout = "NORMAL" if block_size == 1 else "SHUFFLE"
111
grid = (num_tokens,)
112
_indexer_k_quant_and_cache_kernel[grid](
113
k,
@@ -118,7 +119,7 @@ def indexer_k_quant_and_cache_triton(
119
block_size,
120
num_tokens,
121
head_dim,
121
- "SHUFFLE",
122
+ layout,
123
block_tile_size,
124
head_tile_size,
125
IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
@@ -229,6 +230,7 @@ def cp_gather_indexer_k_quant_cache_triton(
230
k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
231
grid = (num_tokens,)
232
k_fp8_scale = k_fp8_scale.view(torch.float32)
233
+ layout = "NORMAL" if block_size == 1 else "SHUFFLE"
234
_cp_gather_indexer_quant_cache_kernel[grid](
235
k_cache_value,
236
k_cache_scale,
@@ -241,7 +243,7 @@ def cp_gather_indexer_k_quant_cache_triton(
243
block_table_stride,
244
k_cache_value.stride(0),
245
k_cache_scale.stride(0),
244
- "SHUFFLE",
246
+ layout,
247
head_dim,
248
block_tile_size,
249
head_tile_size,
@@ -433,7 +435,7 @@ def rocm_fp8_paged_mqa_logits(
435
block_tables,
436
max_model_len,
437
ChunkK=256,
436
- Preshuffle=block_size == 64,
438
+ Preshuffle=block_size > 1,
439
KVBlockSize=block_size,
440
WavePerEU=2,
441
)