[BUGFIX] 解决mqa_logits oom

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -30,6 +30,7 @@ elif current_platform.is_xpu():
     from vllm._ipex_ops import ipex_ops as ops
 
 logger = init_logger(__name__)
+_GLOBAL_LOGITS_BUFFERS = {}
 
 @maybe_transfer_kv_layer
 def sparse_attn_indexer(
@@ -50,7 +51,21 @@ def sparse_attn_indexer(
     # careful! this will be None in dummy run
     attn_metadata = get_forward_context().attn_metadata
     fp8_dtype = current_platform.fp8_dtype()
-
+    if q_fp8.dtype == fp8_dtype:
+        MAX_ELEMENTS = 65536 * 65536
+    elif q_fp8.dtype in (torch.bfloat16, torch.float16):
+        MAX_ELEMENTS = 16384 * 32768
+    else:
+        MAX_ELEMENTS = 16384 * 32768 
+
+    device = q_fp8.device
+    if device not in _GLOBAL_LOGITS_BUFFERS or _GLOBAL_LOGITS_BUFFERS[device].numel() < MAX_ELEMENTS:
+        _GLOBAL_LOGITS_BUFFERS[device] = torch.empty(
+            MAX_ELEMENTS, 
+            dtype=torch.float32, 
+            device=device
+        )
+    logits_buffer = _GLOBAL_LOGITS_BUFFERS[device]
     # assert isinstance(attn_metadata, dict)
     if not isinstance(attn_metadata, dict):
         # Reserve workspace for indexer during profiling run
@@ -144,14 +159,17 @@ def sparse_attn_indexer(
             num_q = q_all.shape[0]
             num_k = k_fp8.shape[0]
 
-            MAX_ELEMENTS = 1024 * 1024 * 1024 # 4GB
-            if (num_q <= 65536 and num_k <= 65536): # if num_q <= 65536 and num_k <= 65536 and (num_q * num_k <= MAX_ELEMENTS):
-                MAX_Q_CHUNK = max(1, num_q)
-            else:
-                MAX_Q_CHUNK = max(1024, MAX_ELEMENTS // max(1, num_k))
-                MAX_Q_CHUNK = min(MAX_Q_CHUNK, max(1, num_q))
+            is_q_fp16_bf16 = q_all.dtype in (torch.float16, torch.bfloat16)
+            align_size = 128 if is_q_fp16_bf16 else 1
+            
+            kv_seq_len_aligned = (num_k + align_size - 1) // align_size * align_size
+
+            current_capacity = logits_buffer.numel()
+            MAX_Q_CHUNK = current_capacity // max(1, kv_seq_len_aligned)
+            if align_size > 1:
+                MAX_Q_CHUNK = (MAX_Q_CHUNK // align_size) * align_size
+            MAX_Q_CHUNK = max(1, MAX_Q_CHUNK)
 
-            #存储q的起始和终止地址
             slices = []
 
             for start_idx in range(0, num_q, MAX_Q_CHUNK):
@@ -168,6 +186,12 @@ def sparse_attn_indexer(
                 ks_slice = ks_all[q_start:q_end]
                 ke_slice = ke_all[q_start:q_end]
 
+                q_len = q_end - q_start
+                q_seq_len_aligned = (q_len + align_size - 1) // align_size * align_size
+
+                required_size = q_seq_len_aligned * kv_seq_len_aligned
+                logits_slice_view = logits_buffer[:required_size].view(q_seq_len_aligned, kv_seq_len_aligned)
+
                 if not current_platform.is_rocm():
                     logits_slice = fp8_mqa_logits(
                         q_slice,
@@ -177,7 +201,7 @@ def sparse_attn_indexer(
                         ke_slice,
                     )
                 elif get_gcn_arch_name() == "gfx938":
-                    logits_slice = op.mqa_logits(
+                    op.mqa_logits(
                         q_slice,  
                         k_fp8, 
                         weights_slice, 
@@ -188,10 +212,12 @@ def sparse_attn_indexer(
                         q_slice.shape[1],
                         q_slice.shape[2],
                         k_scale.view(torch.float32).flatten(),
-                        True
+                        True,
+                        logits_slice_view 
                     )
+                    logits_slice = logits_slice_view[:q_len, :num_k]
                 else:
-                    logits_slice = op.mqa_logits(
+                    op.mqa_logits(
                         q_slice,  
                         k_fp8, 
                         weights_slice.to(torch.float32), 
@@ -202,8 +228,10 @@ def sparse_attn_indexer(
                         q_slice.shape[1],
                         q_slice.shape[2],
                         None,
-                        True
+                        True,
+                        logits_slice_view 
                     )
+                    logits_slice = logits_slice_view[:q_len, :num_k]
 
                 num_rows_slice = logits_slice.shape[0]
                 
@@ -461,5 +489,3 @@ class SparseAttnIndexer(CustomOp):
                 self.max_total_seq_len,
                 self.topk_indices_buffer,
             )
\ No newline at end of file
-        
-