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 6, 00:45 UTC
[Revert] Remove CUDA torch fallbacks for fp8_mqa_logits/fp8_paged_mqa_logits_torch function (#37968)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Mar 25, 2026 at 14:19 UTC
09c3dc91862eb3770c2456f570e95ebc81161b7d
3 files changed
+24
-172
vllm/model_executor/layers/sparse_attn_indexer.py
+22
-50
@@ -9,13 +9,7 @@ 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 (
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
-)
12
+from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits, has_deep_gemm
13
from vllm.utils.torch_utils import direct_register_custom_op
14
from vllm.v1.attention.backends.mla.indexer import (
15
DeepseekV32IndexerMetadata,
@@ -114,23 +108,14 @@ def sparse_attn_indexer(
108
chunk.block_table,
109
chunk.cu_seq_lens,
110
)
117
- if is_deep_gemm_supported():
118
- logits = fp8_mqa_logits(
119
- q_fp8[chunk.token_start : chunk.token_end],
120
- (k_fp8, k_scale.view(torch.float32).flatten()),
121
- weights[chunk.token_start : chunk.token_end],
122
- chunk.cu_seqlen_ks,
123
- chunk.cu_seqlen_ke,
124
- clean_logits=False,
125
- )
126
- else:
127
- logits = fp8_mqa_logits_torch(
128
- q_fp8[chunk.token_start : chunk.token_end],
129
- (k_fp8, k_scale.view(torch.float32).flatten()),
130
- weights[chunk.token_start : chunk.token_end],
131
- chunk.cu_seqlen_ks,
132
- chunk.cu_seqlen_ke,
133
- )
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
num_rows = logits.shape[0]
120
121
topk_indices = topk_indices_buffer[
@@ -194,26 +179,16 @@ def sparse_attn_indexer(
179
next_n = padded_q_fp8_decode_tokens.shape[1]
180
assert batch_size == decode_metadata.seq_lens.shape[0]
181
num_padded_tokens = batch_size * next_n
197
- if is_deep_gemm_supported():
198
- logits = fp8_paged_mqa_logits(
199
- padded_q_fp8_decode_tokens,
200
- kv_cache,
201
- weights[:num_padded_tokens],
202
- decode_metadata.seq_lens,
203
- decode_metadata.block_table,
204
- decode_metadata.schedule_metadata,
205
- max_model_len=max_model_len,
206
- clean_logits=False,
207
- )
208
- else:
209
- logits = fp8_paged_mqa_logits_torch(
210
- padded_q_fp8_decode_tokens,
211
- kv_cache,
212
- weights[:num_padded_tokens],
213
- decode_metadata.seq_lens,
214
- decode_metadata.block_table,
215
- max_model_len=max_model_len,
216
- )
182
+ logits = fp8_paged_mqa_logits(
183
+ padded_q_fp8_decode_tokens,
184
+ kv_cache,
185
+ weights[:num_padded_tokens],
186
+ decode_metadata.seq_lens,
187
+ decode_metadata.block_table,
188
+ decode_metadata.schedule_metadata,
189
+ max_model_len=max_model_len,
190
+ clean_logits=False,
191
+ )
192
num_rows = logits.shape[0]
193
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
194
@@ -333,12 +308,9 @@ class SparseAttnIndexer(CustomOp):
308
self.max_model_len = max_model_len
309
self.max_total_seq_len = max_total_seq_len
310
self.topk_indices_buffer = topk_indices_buffer
336
- if current_platform.is_cuda() and not is_deep_gemm_supported():
337
- logger.warning_once(
338
- "DeepGEMM is not supported or available. SparseAttnIndexer will use a "
339
- "less efficient PyTorch implementation. "
340
- "Please make sure you have the required hardware and software setup "
341
- "for DeepGEMM to achieve optimal performance."
311
+ if current_platform.is_cuda() and not has_deep_gemm():
312
+ raise RuntimeError(
313
+ "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed."
314
)
315
316
def forward_native(
vllm/utils/deep_gemm.py
-121
@@ -415,125 +415,6 @@ def should_use_deepgemm_for_fp8_linear(
415
)
416
417
418
-def fp8_mqa_logits_torch(
419
- q: torch.Tensor,
420
- kv: tuple[torch.Tensor, torch.Tensor],
421
- weights: torch.Tensor,
422
- cu_seqlen_ks: torch.Tensor,
423
- cu_seqlen_ke: torch.Tensor,
424
-) -> torch.Tensor:
425
- """Compute FP8 MQA logits for a single sequence without KV paging (CUDA fallback).
426
-
427
- This is a pure PyTorch fallback for CUDA when DeepGEMM is not available.
428
-
429
- Args:
430
- q: Query tensor of shape [M, H, D]. Casted to
431
- `torch.float8_e4m3fn` by caller.
432
- kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
433
- dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
434
- [N, 1]) with dtype `torch.float32`.
435
- weights: weights of shape [M, H], dtype `torch.float32`.
436
- cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
437
- shape [M], dtype int32.
438
- cu_seqlen_ke: End indices (exclusive) for valid K per query position,
439
- shape [M], dtype int32.
440
-
441
- Returns:
442
- Logits tensor of shape [M, N], dtype `torch.float32`.
443
- """
444
- kv_fp8, scale = kv
445
- seq_len_kv = kv_fp8.shape[0]
446
- k = kv_fp8.to(torch.bfloat16)
447
- q = q.to(torch.bfloat16)
448
-
449
- mask_lo = (
450
- torch.arange(0, seq_len_kv, device=q.device)[None, :] >= cu_seqlen_ks[:, None]
451
- )
452
- mask_hi = (
453
- torch.arange(0, seq_len_kv, device=q.device)[None, :] < cu_seqlen_ke[:, None]
454
- )
455
- mask = mask_lo & mask_hi
456
-
457
- score = torch.einsum("mhd,nd->hmn", q, k).float() * scale
458
- logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0)
459
- logits = logits.masked_fill(~mask, float("-inf"))
460
-
461
- return logits
462
-
463
-
464
-def fp8_paged_mqa_logits_torch(
465
- q: torch.Tensor,
466
- kv_cache: torch.Tensor,
467
- weights: torch.Tensor,
468
- context_lens: torch.Tensor,
469
- block_tables: torch.Tensor,
470
- max_model_len: int,
471
-) -> torch.Tensor:
472
- """Compute FP8 MQA logits using paged KV-cache (CUDA fallback).
473
-
474
- This is a pure PyTorch fallback for CUDA when DeepGEMM is not available.
475
- Handles head_dim = 132 (128 + 4 for RoPE).
476
-
477
- Args:
478
- q: Query tensor of shape [B, next_n, H, D].
479
- kv_cache: Paged KV-cache in packed FP8+scale layout with shape
480
- [num_blocks, block_size, 1, D+4], dtype `torch.uint8`. The last
481
- 4 bytes per (block,pos) store the `float` dequant scale.
482
- weights: Tensor of shape [B * next_n, H], dtype `torch.float32`.
483
- context_lens: Tensor of shape [B], dtype int32; effective context length
484
- for each batch element.
485
- block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical
486
- block indices to physical blocks in the paged cache.
487
- max_model_len: Maximum sequence length used to size the logits output.
488
-
489
- Returns:
490
- Logits tensor of shape [B * next_n, max_model_len], dtype
491
- `torch.float32`.
492
- """
493
- fp8_dtype = current_platform.fp8_dtype()
494
- batch_size, next_n, heads, dim = q.size()
495
- kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
496
- scale = scale.contiguous().view(torch.float)
497
- q = q.float()
498
- kv_cache = kv_cache.view(fp8_dtype).float() * scale
499
- num_blocks, block_size, _, dim = kv_cache.size()
500
- logits = torch.full(
501
- [batch_size * next_n, max_model_len],
502
- float("-inf"),
503
- device=q.device,
504
- dtype=torch.float32,
505
- )
506
- for i in range(batch_size):
507
- context_len = context_lens[i].item()
508
- q_offsets = torch.arange(context_len - next_n, context_len, device=q.device)
509
- weight_slice = (
510
- weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous()
511
- )
512
- for block_idx in range(cdiv(context_len, block_size)):
513
- block_id = block_tables[i][block_idx]
514
- qx, kx = q[i], kv_cache[block_id]
515
- k_offsets = torch.arange(
516
- block_idx * block_size, (block_idx + 1) * block_size, device=q.device
517
- )
518
- mask = (k_offsets[None, :] < context_len) & (
519
- k_offsets[None, :] <= q_offsets[:, None]
520
- )
521
- s = torch.where(
522
- mask[None, :, :],
523
- (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(
524
- logits.dtype
525
- ),
526
- float("-inf"),
527
- )
528
- s = torch.relu(s) * weight_slice[..., None]
529
- s = s.sum(dim=0)
530
- logits[
531
- i * next_n : (i + 1) * next_n,
532
- block_idx * block_size : (block_idx + 1) * block_size,
533
- ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float("-inf"))
534
- return logits
535
-
536
-
418
__all__ = [
419
"calc_diff",
420
"DeepGemmQuantScaleFMT",
@@ -541,9 +422,7 @@ __all__ = [
422
"m_grouped_fp8_gemm_nt_contiguous",
423
"fp8_m_grouped_gemm_nt_masked",
424
"fp8_mqa_logits",
544
- "fp8_mqa_logits_torch",
425
"fp8_paged_mqa_logits",
546
- "fp8_paged_mqa_logits_torch",
426
"get_paged_mqa_logits_metadata",
427
"per_block_cast_to_fp8",
428
"is_deep_gemm_e8m0_used",
vllm/v1/attention/backends/mla/indexer.py
+2
-1
@@ -9,6 +9,7 @@ from vllm.logger import init_logger
9
from vllm.platforms import current_platform
10
from vllm.utils.deep_gemm import (
11
get_paged_mqa_logits_metadata,
12
+ has_deep_gemm,
13
is_deep_gemm_supported,
14
)
15
from vllm.utils.math_utils import cdiv
@@ -449,7 +450,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
450
batch_size = num_decodes
451
452
# DeepGEMM is required for the paged MQA logits on CUDA devices
452
- if current_platform.is_cuda() and is_deep_gemm_supported():
453
+ if current_platform.is_cuda() and has_deep_gemm():
454
self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
455
seq_lens,
456
self.kv_cache_spec.block_size,