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 7, 12:45 UTC
[Feat] Add CUDA torch fallbacks for fp8_mqa_logits/fp8_paged_mqa_logits_torch function (#35271)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Feb 28, 2026 at 18:12 UTC
7e08c22b8cb65a1bea6b4bf9c52ed6e71d4acc47
3 files changed
+176
-28
vllm/model_executor/layers/sparse_attn_indexer.py
+50
-26
@@ -9,8 +9,13 @@ from vllm.forward_context import get_forward_context
9
from vllm.logger import init_logger
10
from vllm.model_executor.custom_op import CustomOp
11
from vllm.platforms import current_platform
12
-from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
13
-from vllm.utils.import_utils import has_deep_gemm
12
+from vllm.utils.deep_gemm import (
13
+ fp8_mqa_logits,
14
+ fp8_mqa_logits_torch,
15
+ fp8_paged_mqa_logits,
16
+ fp8_paged_mqa_logits_torch,
17
+ is_deep_gemm_supported,
18
+)
19
from vllm.utils.torch_utils import direct_register_custom_op
20
from vllm.v1.attention.backends.mla.indexer import (
21
DeepseekV32IndexerMetadata,
@@ -102,15 +107,23 @@ def sparse_attn_indexer(
107
chunk.block_table,
108
chunk.cu_seq_lens,
109
)
105
-
106
- logits = fp8_mqa_logits(
107
- q_fp8[chunk.token_start : chunk.token_end],
108
- (k_fp8, k_scale.view(torch.float32).flatten()),
109
- weights[chunk.token_start : chunk.token_end],
110
- chunk.cu_seqlen_ks,
111
- chunk.cu_seqlen_ke,
112
- clean_logits=False,
113
- )
110
+ if is_deep_gemm_supported():
111
+ logits = fp8_mqa_logits(
112
+ q_fp8[chunk.token_start : chunk.token_end],
113
+ (k_fp8, k_scale.view(torch.float32).flatten()),
114
+ weights[chunk.token_start : chunk.token_end],
115
+ chunk.cu_seqlen_ks,
116
+ chunk.cu_seqlen_ke,
117
+ clean_logits=False,
118
+ )
119
+ else:
120
+ logits = fp8_mqa_logits_torch(
121
+ q_fp8[chunk.token_start : chunk.token_end],
122
+ (k_fp8, k_scale.view(torch.float32).flatten()),
123
+ weights[chunk.token_start : chunk.token_end],
124
+ chunk.cu_seqlen_ks,
125
+ chunk.cu_seqlen_ke,
126
+ )
127
num_rows = logits.shape[0]
128
129
topk_indices = topk_indices_buffer[
@@ -159,18 +172,26 @@ def sparse_attn_indexer(
172
next_n = padded_q_fp8_decode_tokens.shape[1]
173
assert batch_size == decode_metadata.seq_lens.shape[0]
174
num_padded_tokens = batch_size * next_n
162
-
163
- logits = fp8_paged_mqa_logits(
164
- padded_q_fp8_decode_tokens,
165
- kv_cache,
166
- weights[:num_padded_tokens],
167
- decode_metadata.seq_lens,
168
- decode_metadata.block_table,
169
- decode_metadata.schedule_metadata,
170
- max_model_len=max_model_len,
171
- clean_logits=False,
172
- )
173
-
175
+ if is_deep_gemm_supported():
176
+ logits = fp8_paged_mqa_logits(
177
+ padded_q_fp8_decode_tokens,
178
+ kv_cache,
179
+ weights[:num_padded_tokens],
180
+ decode_metadata.seq_lens,
181
+ decode_metadata.block_table,
182
+ decode_metadata.schedule_metadata,
183
+ max_model_len=max_model_len,
184
+ clean_logits=False,
185
+ )
186
+ else:
187
+ logits = fp8_paged_mqa_logits_torch(
188
+ padded_q_fp8_decode_tokens,
189
+ kv_cache,
190
+ weights[:num_padded_tokens],
191
+ decode_metadata.seq_lens,
192
+ decode_metadata.block_table,
193
+ max_model_len=max_model_len,
194
+ )
195
num_rows = logits.shape[0]
196
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
197
@@ -278,9 +299,12 @@ class SparseAttnIndexer(CustomOp):
299
self.max_model_len = max_model_len
300
self.max_total_seq_len = max_total_seq_len
301
self.topk_indices_buffer = topk_indices_buffer
281
- if current_platform.is_cuda() and not has_deep_gemm():
282
- raise RuntimeError(
283
- "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
302
+ if current_platform.is_cuda() and not is_deep_gemm_supported():
303
+ logger.warning_once(
304
+ "DeepGEMM is not supported or available. SparseAttnIndexer will use a "
305
+ "less efficient PyTorch implementation. "
306
+ "Please make sure you have the required hardware and software setup "
307
+ "for DeepGEMM to achieve optimal performance."
308
)
309
310
def forward_native(
vllm/utils/deep_gemm.py
+121
@@ -418,6 +418,125 @@ def should_use_deepgemm_for_fp8_linear(
418
)
419
420
421
+def fp8_mqa_logits_torch(
422
+ q: torch.Tensor,
423
+ kv: tuple[torch.Tensor, torch.Tensor],
424
+ weights: torch.Tensor,
425
+ cu_seqlen_ks: torch.Tensor,
426
+ cu_seqlen_ke: torch.Tensor,
427
+) -> torch.Tensor:
428
+ """Compute FP8 MQA logits for a single sequence without KV paging (CUDA fallback).
429
+
430
+ This is a pure PyTorch fallback for CUDA when DeepGEMM is not available.
431
+
432
+ Args:
433
+ q: Query tensor of shape [M, H, D]. Casted to
434
+ `torch.float8_e4m3fn` by caller.
435
+ kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
436
+ dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
437
+ [N, 1]) with dtype `torch.float32`.
438
+ weights: weights of shape [M, H], dtype `torch.float32`.
439
+ cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
440
+ shape [M], dtype int32.
441
+ cu_seqlen_ke: End indices (exclusive) for valid K per query position,
442
+ shape [M], dtype int32.
443
+
444
+ Returns:
445
+ Logits tensor of shape [M, N], dtype `torch.float32`.
446
+ """
447
+ kv_fp8, scale = kv
448
+ seq_len_kv = kv_fp8.shape[0]
449
+ k = kv_fp8.to(torch.bfloat16)
450
+ q = q.to(torch.bfloat16)
451
+
452
+ mask_lo = (
453
+ torch.arange(0, seq_len_kv, device=q.device)[None, :] >= cu_seqlen_ks[:, None]
454
+ )
455
+ mask_hi = (
456
+ torch.arange(0, seq_len_kv, device=q.device)[None, :] < cu_seqlen_ke[:, None]
457
+ )
458
+ mask = mask_lo & mask_hi
459
+
460
+ score = torch.einsum("mhd,nd->hmn", q, k).float() * scale
461
+ logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
462
+ logits = logits.masked_fill(~mask, float("-inf"))
463
+
464
+ return logits
465
+
466
+
467
+def fp8_paged_mqa_logits_torch(
468
+ q: torch.Tensor,
469
+ kv_cache: torch.Tensor,
470
+ weights: torch.Tensor,
471
+ context_lens: torch.Tensor,
472
+ block_tables: torch.Tensor,
473
+ max_model_len: int,
474
+) -> torch.Tensor:
475
+ """Compute FP8 MQA logits using paged KV-cache (CUDA fallback).
476
+
477
+ This is a pure PyTorch fallback for CUDA when DeepGEMM is not available.
478
+ Handles head_dim = 132 (128 + 4 for RoPE).
479
+
480
+ Args:
481
+ q: Query tensor of shape [B, next_n, H, D].
482
+ kv_cache: Paged KV-cache in packed FP8+scale layout with shape
483
+ [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
484
+ 4 bytes per (block,pos) store the `float` dequant scale.
485
+ weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
486
+ context_lens: Tensor of shape [B], dtype int32; effective context length
487
+ for each batch element.
488
+ block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
489
+ block indices to physical blocks in the paged cache.
490
+ max_model_len: Maximum sequence length used to size the logits output.
491
+
492
+ Returns:
493
+ Logits tensor of shape [B * next_n, max_model_len], dtype
494
+ `torch.float32`.
495
+ """
496
+ fp8_dtype = current_platform.fp8_dtype()
497
+ batch_size, next_n, heads, dim = q.size()
498
+ kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
499
+ scale = scale.contiguous().view(torch.float)
500
+ q = q.float()
501
+ kv_cache = kv_cache.view(fp8_dtype).float() * scale
502
+ num_blocks, block_size, _, dim = kv_cache.size()
503
+ logits = torch.full(
504
+ [batch_size * next_n, max_model_len],
505
+ float("-inf"),
506
+ device=q.device,
507
+ dtype=torch.float32,
508
+ )
509
+ for i in range(batch_size):
510
+ context_len = context_lens[i].item()
511
+ q_offsets = torch.arange(context_len - next_n, context_len, device=q.device)
512
+ weight_slice = (
513
+ weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
514
+ )
515
+ for block_idx in range(cdiv(context_len, block_size)):
516
+ block_id = block_tables[i][block_idx]
517
+ qx, kx = q[i], kv_cache[block_id]
518
+ k_offsets = torch.arange(
519
+ block_idx * block_size, (block_idx + 1) * block_size, device=q.device
520
+ )
521
+ mask = (k_offsets[None, :] < context_len) & (
522
+ k_offsets[None, :] <= q_offsets[:, None]
523
+ )
524
+ s = torch.where(
525
+ mask[None, :, :],
526
+ (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
527
+ logits.dtype
528
+ ),
529
+ float("-inf"),
530
+ )
531
+ s = torch.relu(s) * weight_slice[..., None]
532
+ s = s.sum(dim=0)
533
+ logits[
534
+ i * next_n : (i + 1) * next_n,
535
+ block_idx * block_size : (block_idx + 1) * block_size,
536
+ ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf"))
537
+ return logits
538
+
539
+
540
__all__ = [
541
"calc_diff",
542
"DeepGemmQuantScaleFMT",
@@ -425,7 +544,9 @@ __all__ = [
544
"m_grouped_fp8_gemm_nt_contiguous",
545
"fp8_m_grouped_gemm_nt_masked",
546
"fp8_mqa_logits",
547
+ "fp8_mqa_logits_torch",
548
"fp8_paged_mqa_logits",
549
+ "fp8_paged_mqa_logits_torch",
550
"get_paged_mqa_logits_metadata",
551
"per_block_cast_to_fp8",
552
"is_deep_gemm_e8m0_used",
vllm/v1/attention/backends/mla/indexer.py
+5
-2
@@ -8,7 +8,10 @@ import torch
8
from vllm.config import VllmConfig
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
-from vllm.utils.deep_gemm import get_paged_mqa_logits_metadata, has_deep_gemm
11
+from vllm.utils.deep_gemm import (
12
+ get_paged_mqa_logits_metadata,
13
+ is_deep_gemm_supported,
14
+)
15
from vllm.utils.platform_utils import num_compute_units
16
from vllm.v1.attention.backend import (
17
AttentionBackend,
@@ -344,7 +347,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
347
seq_lens = common_attn_metadata.seq_lens[:num_decodes]
348
349
# DeepGEMM is required for the paged MQA logits on CUDA devices
347
- if current_platform.is_cuda() and has_deep_gemm():
350
+ if current_platform.is_cuda() and is_deep_gemm_supported():
351
self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
352
seq_lens, self.kv_cache_spec.block_size, self.num_sms
353
)