update sparse_attn_indexer

vllm/model_executor/models/deepseek_v2.py

@@ -872,38 +872,21 @@ class Indexer(nn.Module):
                 -1) * q_scale * self.softmax_scale * self.n_head**-0.5
             weights = weights.squeeze(-1)
 
-        if not current_platform.is_rocm() or torch.cuda.get_device_properties("cuda").gcnArchName.split(':')[0] == "gfx938":
-            return torch.ops.vllm.sparse_attn_indexer(
-                hidden_states,
-                self.k_cache.prefix,
-                self.k_cache.kv_cache[0],
-                q_fp8,
-                k,
-                weights,
-                self.quant_block_size,
-                self.scale_fmt,
-                self.topk_tokens,
-                self.head_dim,
-                self.max_model_len,
-                self.max_total_seq_len,
-                self.topk_indices_buffer,
-            )
-        else:
-            return torch.ops.vllm.sparse_attn_indexer(
-                hidden_states,
-                self.k_cache.prefix,
-                self.k_cache.kv_cache[0],
-                q,
-                k,
-                weights,
-                self.quant_block_size,
-                self.scale_fmt,
-                self.topk_tokens,
-                self.head_dim,
-                self.max_model_len,
-                self.max_total_seq_len,
-                self.topk_indices_buffer,
-            )
+        return torch.ops.vllm.sparse_attn_indexer(
+            hidden_states,
+            self.k_cache.prefix,
+            self.k_cache.kv_cache[0],
+            q_fp8 if not current_platform.is_rocm() or torch.cuda.get_device_properties("cuda").gcnArchName.split(':')[0] == "gfx938" else q,
+            k,
+            weights,
+            self.quant_block_size,
+            self.scale_fmt,
+            self.topk_tokens,
+            self.head_dim,
+            self.max_model_len,
+            self.max_total_seq_len,
+            self.topk_indices_buffer,
+        )
 
 
 class DeepseekV2MLAAttention(nn.Module):