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 11, 01:15 UTC
Revert "[SM100] Enable fp8 compute for prefill MLA (#30746)" (#31197)
Signed-off-by: Pavani Majety <pmajety@nvidia.com>
Pavani Majety committed
Dec 22, 2025 at 18:15 UTC
3e10262356657f3ee7673335832599f910c28d89
3 files changed
+18
-116
tests/v1/attention/test_mla_backends.py
+3
-4
@@ -27,7 +27,7 @@ from vllm.utils.math_utils import cdiv
27
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
28
from vllm.v1.attention.backends.mla.common import QueryLenSupport
29
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
30
-from vllm.v1.kv_cache_interface import MLAAttentionSpec
30
+from vllm.v1.kv_cache_interface import FullAttentionSpec
31
32
BACKENDS_TO_TEST = [
33
AttentionBackendEnum.CUTLASS_MLA,
@@ -289,7 +289,7 @@ class MockMLAAttentionLayer(AttentionLayerBase):
289
290
def run_attention_backend(
291
backend: AttentionBackendEnum,
292
- kv_cache_spec: MLAAttentionSpec,
292
+ kv_cache_spec: FullAttentionSpec,
293
layer_names: list[str],
294
vllm_config,
295
device: torch.device,
@@ -740,7 +740,7 @@ def test_backend_correctness(
740
kv_cache = kv_cache_per_block_size[block_size]
741
742
# Create kv_cache_spec with the correct block_size for this backend
743
- backend_kv_cache_spec = MLAAttentionSpec(
743
+ backend_kv_cache_spec = FullAttentionSpec(
744
block_size=block_size,
745
num_kv_heads=vllm_config.model_config.get_num_kv_heads(
746
vllm_config.parallel_config
@@ -748,7 +748,6 @@ def test_backend_correctness(
748
head_size=vllm_config.model_config.get_head_size(),
749
dtype=vllm_config.model_config.dtype,
750
sliding_window=vllm_config.model_config.get_sliding_window(),
751
- cache_dtype_str=vllm_config.cache_config.cache_dtype,
751
)
752
753
backend_output = run_attention_backend(
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+1
@@ -325,6 +325,7 @@ def flashinfer_trtllm_fp4_moe(
325
local_expert_offset=layer.ep_rank * layer.local_num_experts,
326
local_num_experts=layer.local_num_experts,
327
routed_scaling_factor=None,
328
+ tile_tokens_dim=None,
329
routing_method_type=routing_method_type,
330
do_finalize=True,
331
)[0]
vllm/v1/attention/backends/mla/common.py
+14
-112
@@ -541,11 +541,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
541
metadata_cls if metadata_cls is not None else MLACommonMetadata
542
)
543
self.kv_cache_spec = kv_cache_spec
544
- self.q_data_type = (
545
- current_platform.fp8_dtype()
546
- if (kv_cache_spec is not None and "fp8" in kv_cache_spec.cache_dtype_str)
547
- else vllm_config.model_config.dtype
548
- )
544
scheduler_config = vllm_config.scheduler_config
545
self.model_config = vllm_config.model_config
546
parallel_config = vllm_config.parallel_config
@@ -689,6 +684,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
684
685
# For main run, qo_indptr == kv_indptr
686
kv_indptr = qo_indptr.clone()
687
+
688
# Prepare main prefill
689
self._fi_prefill_main.plan(
690
qo_indptr=qo_indptr,
@@ -701,7 +697,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
697
sm_scale=self._global_hyperparameters.sm_scale,
698
window_left=self._global_hyperparameters.window_left,
699
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
704
- q_data_type=self.q_data_type,
700
+ q_data_type=self.model_config.dtype,
701
)
702
703
# Prepare context prefills
@@ -720,7 +716,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
716
sm_scale=self._global_hyperparameters.sm_scale,
717
window_left=self._global_hyperparameters.window_left,
718
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
723
- q_data_type=self.q_data_type,
719
+ q_data_type=self.model_config.dtype,
720
)
721
722
prefill.prefill_main = self._fi_prefill_main
@@ -973,7 +969,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
969
query_start_loc=prefill_query_start_loc,
970
max_query_len=max_query_len,
971
chunked_context=chunked_context_metadata,
976
- q_data_type=self.q_data_type,
972
)
973
974
if self._use_cudnn_prefill:
@@ -1384,15 +1379,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1379
return attn_out
1380
1381
def _run_prefill_new_tokens_fa(
1387
- self,
1388
- prefill: MLACommonPrefillMetadata,
1389
- q,
1390
- k,
1391
- v,
1392
- return_softmax_lse,
1393
- fp8_attention: bool,
1382
+ self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
1383
):
1395
- logger.debug_once("Running FlashAttention prefill new tokens", scope="local")
1384
return self._flash_attn_varlen_diff_headdims(
1385
q=q,
1386
k=k,
@@ -1407,23 +1395,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1395
)
1396
1397
def _run_prefill_new_tokens_fi(
1410
- self,
1411
- prefill: MLACommonPrefillMetadata,
1412
- q,
1413
- k,
1414
- v,
1415
- return_softmax_lse,
1416
- fp8_attention: bool,
1398
+ self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
1399
):
1418
- logger.debug_once("Running FlashInfer prefill new tokens", scope="local")
1400
assert isinstance(prefill, FlashInferPrefillMetadata)
1401
assert prefill.prefill_main is not None
1421
- if fp8_attention:
1422
- logger.debug_once("Running Flashinfer prefill in FP8")
1423
- fp8_dtype = current_platform.fp8_dtype()
1424
- q = q.to(fp8_dtype)
1425
- k = k.to(fp8_dtype)
1426
- v = v.to(fp8_dtype)
1402
+
1403
ret = prefill.prefill_main.run(
1404
q=q,
1405
k=k,
@@ -1436,18 +1412,10 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1412
return ret
1413
1414
def _run_prefill_new_tokens_cudnn(
1439
- self,
1440
- prefill: MLACommonPrefillMetadata,
1441
- q,
1442
- k,
1443
- v,
1444
- return_softmax_lse,
1445
- fp8_attention: bool,
1415
+ self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
1416
):
1447
- logger.debug_once("Running Cudnn prefill new tokens", scope="local")
1417
assert isinstance(prefill, CudnnPrefillMetadata)
1418
assert prefill.query_seq_lens is not None
1450
- assert fp8_attention is False, "Cudnn prefill does not support fp8 attention"
1419
output, lse = cudnn_batch_prefill_with_kv_cache(
1420
q=q,
1421
k_cache=k,
@@ -1469,19 +1437,9 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1437
return output
1438
1439
def _run_prefill_context_chunk_fa(
1472
- self,
1473
- prefill: MLACommonPrefillMetadata,
1474
- chunk_idx: int,
1475
- q,
1476
- k,
1477
- v,
1478
- fp8_attention: bool,
1440
+ self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
1441
):
1480
- logger.debug_once("Running FlashAttention prefill context chunk", scope="local")
1442
assert prefill.chunked_context is not None
1482
- assert fp8_attention is False, (
1483
- "FlashAttention prefill does not support fp8 attention"
1484
- )
1443
return self._flash_attn_varlen_diff_headdims(
1444
q=q,
1445
k=k,
@@ -1496,22 +1454,10 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1454
)
1455
1456
def _run_prefill_context_chunk_fi(
1499
- self,
1500
- prefill: MLACommonPrefillMetadata,
1501
- chunk_idx: int,
1502
- q,
1503
- k,
1504
- v,
1505
- fp8_attention: bool,
1457
+ self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
1458
):
1507
- logger.debug_once("Running FlashInfer prefill context chunk", scope="local")
1459
assert isinstance(prefill, FlashInferPrefillMetadata)
1509
- if fp8_attention:
1510
- logger.debug_once("Running FlashInfer prefill in FP8")
1511
- fp8_dtype = current_platform.fp8_dtype()
1512
- q = q.to(fp8_dtype)
1513
- k = k.to(fp8_dtype)
1514
- v = v.to(fp8_dtype)
1460
+
1461
attn_out, lse = prefill.prefill_chunks[chunk_idx].run(
1462
q=q,
1463
k=k,
@@ -1523,20 +1469,12 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1469
return attn_out, lse.transpose(0, 1).contiguous()
1470
1471
def _run_prefill_context_chunk_cudnn(
1526
- self,
1527
- prefill: MLACommonPrefillMetadata,
1528
- chunk_idx: int,
1529
- q,
1530
- k,
1531
- v,
1532
- fp8_attention: bool,
1472
+ self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
1473
):
1534
- logger.debug_once("Running Cudnn prefill context chunk", scope="local")
1474
assert isinstance(prefill, CudnnPrefillMetadata)
1475
assert prefill.chunked_context is not None
1476
assert prefill.chunked_context.seq_lens[chunk_idx] is not None
1477
assert prefill.query_seq_lens is not None
1539
- assert fp8_attention is False, "Cudnn prefill does not support fp8 attention"
1478
return cudnn_batch_prefill_with_kv_cache(
1479
q=q,
1480
k_cache=k,
@@ -1556,28 +1494,14 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1494
)
1495
1496
def _run_prefill_new_tokens_trtllm_ragged(
1559
- self,
1560
- prefill: MLACommonPrefillMetadata,
1561
- q,
1562
- k,
1563
- v,
1564
- return_softmax_lse,
1565
- fp8_attention: bool,
1497
+ self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
1498
):
1567
- logger.debug_once("Running TRT-LLM ragged prefill new tokens", scope="local")
1499
"""TRT-LLM ragged attention for new tokens (causal)."""
1500
from flashinfer.prefill import trtllm_ragged_attention_deepseek
1501
1502
assert prefill.query_seq_lens is not None
1503
assert prefill.workspace_buffer is not None
1504
1574
- if fp8_attention:
1575
- logger.debug_once("Running TRT-LLM ragged prefill in FP8")
1576
- fp8_dtype = current_platform.fp8_dtype()
1577
- q = q.to(fp8_dtype)
1578
- k = k.to(fp8_dtype)
1579
- v = v.to(fp8_dtype)
1580
-
1505
ret = trtllm_ragged_attention_deepseek(
1506
query=q,
1507
key=k,
@@ -1604,15 +1528,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1528
return ret
1529
1530
def _run_prefill_context_chunk_trtllm_ragged(
1607
- self,
1608
- prefill: MLACommonPrefillMetadata,
1609
- chunk_idx: int,
1610
- q,
1611
- k,
1612
- v,
1613
- fp8_attention: bool,
1531
+ self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v
1532
):
1615
- logger.debug_once("Running TRT-LLM ragged prefill context chunk", scope="local")
1533
"""TRT-LLM ragged attention for context chunks (non-causal)."""
1534
from flashinfer.prefill import trtllm_ragged_attention_deepseek
1535
@@ -1629,13 +1546,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1546
)
1547
prefill.workspace_buffer.fill_(0)
1548
1632
- if fp8_attention:
1633
- logger.debug_once("Running TRT-LLM ragged prefill context chunk in FP8")
1634
- fp8_dtype = current_platform.fp8_dtype()
1635
- q = q.to(fp8_dtype)
1636
- k = k.to(fp8_dtype)
1637
- v = v.to(fp8_dtype)
1638
-
1549
attn_out, lse = trtllm_ragged_attention_deepseek(
1550
query=q,
1551
key=k,
@@ -1788,7 +1698,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1698
kv_c_and_k_pe_cache: torch.Tensor,
1699
attn_metadata: MLACommonMetadata,
1700
k_scale: torch.Tensor,
1791
- fp8_attention: bool,
1701
):
1702
assert attn_metadata.prefill is not None
1703
prefill_metadata = attn_metadata.prefill
@@ -1827,7 +1736,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1736
q=q,
1737
k=k,
1738
v=v,
1830
- fp8_attention=fp8_attention,
1739
)
1740
1741
if output is None:
@@ -1856,7 +1764,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1764
attn_metadata: MLACommonMetadata,
1765
k_scale: torch.Tensor,
1766
dcp_world_size: int,
1859
- fp8_attention: bool,
1767
):
1768
assert k_scale is None, "DCP not support scaled kvcache now."
1769
assert attn_metadata.prefill is not None
@@ -1933,7 +1840,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1840
q=q,
1841
k=k,
1842
v=v,
1936
- fp8_attention=fp8_attention,
1843
)
1844
1845
if output is None:
@@ -1964,7 +1870,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1870
attn_metadata: MLACommonMetadata,
1871
k_scale: torch.Tensor,
1872
output: torch.Tensor,
1967
- fp8_attention: bool = False,
1873
) -> None:
1874
# TODO (zyongye): Prefill function here
1875
assert attn_metadata.prefill is not None
@@ -1984,7 +1889,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1889
k=k,
1890
v=v,
1891
return_softmax_lse=has_context,
1987
- fp8_attention=fp8_attention,
1892
)
1893
1894
if has_context:
@@ -1997,12 +1901,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1901
attn_metadata,
1902
k_scale=None,
1903
dcp_world_size=self.dcp_world_size,
2000
- fp8_attention=fp8_attention,
1904
)
1905
)
1906
else:
1907
context_output, context_lse = self._compute_prefill_context(
2005
- q, kv_c_and_k_pe_cache, attn_metadata, k_scale, fp8_attention
1908
+ q, kv_c_and_k_pe_cache, attn_metadata, k_scale
1909
)
1910
1911
# unpad if necessary
@@ -2123,7 +2026,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
2026
attn_metadata,
2027
layer._k_scale,
2028
output=output[num_decode_tokens:],
2126
- fp8_attention=fp8_attention,
2029
)
2030
2031
if has_decode: