update sparse_attn_indexer

vllm/model_executor/models/deepseek_v2.py

@@ -872,28 +872,11 @@ 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,
+            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,