@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b10f41c894

Fork 0

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

[SM100] Enable fp8 compute for prefill MLA (#30746)

Signed-off-by: Pavani Majety <pmajety@nvidia.com>

Pavani Majety committed Dec 22, 2025 at 11:15 UTC b10f41c894bc55c8dd8234a40edd1971be762805

3 files changed +117 -18

tests/v1/attention/test_mla_backends.py

+4 -3

@@ -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 FullAttentionSpec

30 +from vllm.v1.kv_cache_interface import MLAAttentionSpec

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: FullAttentionSpec,

292 + kv_cache_spec: MLAAttentionSpec,

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 = FullAttentionSpec(

743 + backend_kv_cache_spec = MLAAttentionSpec(

744 block_size=block_size,

745 num_kv_heads=vllm_config.model_config.get_num_kv_heads(

746 vllm_config.parallel_config

@@ -748,6 +748,7 @@ 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,

752 )

753

754 backend_output = run_attention_backend(

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

-1

@@ -325,7 +325,6 @@ 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,

328 routing_method_type=routing_method_type,

329 do_finalize=True,

330 )[0]

vllm/v1/attention/backends/mla/common.py

+113 -14

@@ -355,6 +355,7 @@ class MLACommonPrefillMetadata:

355 max_query_len: int

356 chunked_context: ChunkedContextMetadata | None = None

357 query_seq_lens: torch.Tensor | None = None

358 + q_data_type: torch.dtype | None = None

359

360

361 @dataclass

@@ -539,6 +540,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

540 metadata_cls if metadata_cls is not None else MLACommonMetadata

541 )

542 self.kv_cache_spec = kv_cache_spec

543 + self.q_data_type = (

544 + current_platform.fp8_dtype()

545 + if (kv_cache_spec is not None and "fp8" in kv_cache_spec.cache_dtype_str)

546 + else vllm_config.model_config.dtype

547 + )

548 scheduler_config = vllm_config.scheduler_config

549 self.model_config = vllm_config.model_config

550 parallel_config = vllm_config.parallel_config

@@ -681,7 +687,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

687

688 # For main run, qo_indptr == kv_indptr

689 kv_indptr = qo_indptr.clone()

684 -

690 # Prepare main prefill

691 self._fi_prefill_main.plan(

692 qo_indptr=qo_indptr,

@@ -694,7 +699,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

699 sm_scale=self._global_hyperparameters.sm_scale,

700 window_left=self._global_hyperparameters.window_left,

701 logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

697 - q_data_type=self.model_config.dtype,

702 + q_data_type=self.q_data_type,

703 )

704

705 # Prepare context prefills

@@ -713,7 +718,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

718 sm_scale=self._global_hyperparameters.sm_scale,

719 window_left=self._global_hyperparameters.window_left,

720 logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

716 - q_data_type=self.model_config.dtype,

721 + q_data_type=self.q_data_type,

722 )

723

724 prefill.prefill_main = self._fi_prefill_main

@@ -970,6 +975,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

975 query_start_loc=prefill_query_start_loc,

976 max_query_len=max_query_len,

977 chunked_context=chunked_context_metadata,

978 + q_data_type=self.q_data_type,

979 )

980

981 if self._use_cudnn_prefill:

@@ -1370,8 +1376,15 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1376 return attn_out

1377

1378 def _run_prefill_new_tokens_fa(

1373 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

1379 + self,

1380 + prefill: MLACommonPrefillMetadata,

1381 + q,

1382 + k,

1383 + v,

1384 + return_softmax_lse,

1385 + fp8_attention: bool,

1386 ):

1387 + logger.debug_once("Running FlashAttention prefill new tokens", scope="local")

1388 return self._flash_attn_varlen_diff_headdims(

1389 q=q,

1390 k=k,

@@ -1386,11 +1399,23 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1399 )

1400

1401 def _run_prefill_new_tokens_fi(

1389 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

1402 + self,

1403 + prefill: MLACommonPrefillMetadata,

1404 + q,

1405 + k,

1406 + v,

1407 + return_softmax_lse,

1408 + fp8_attention: bool,

1409 ):

1410 + logger.debug_once("Running FlashInfer prefill new tokens", scope="local")

1411 assert isinstance(prefill, FlashInferPrefillMetadata)

1412 assert prefill.prefill_main is not None

1393 -

1413 + if fp8_attention:

1414 + logger.debug_once("Running Flashinfer prefill in FP8")

1415 + fp8_dtype = current_platform.fp8_dtype()

1416 + q = q.to(fp8_dtype)

1417 + k = k.to(fp8_dtype)

1418 + v = v.to(fp8_dtype)

1419 ret = prefill.prefill_main.run(

1420 q=q,

1421 k=k,

@@ -1403,10 +1428,18 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1428 return ret

1429

1430 def _run_prefill_new_tokens_cudnn(

1406 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

1431 + self,

1432 + prefill: MLACommonPrefillMetadata,

1433 + q,

1434 + k,

1435 + v,

1436 + return_softmax_lse,

1437 + fp8_attention: bool,

1438 ):

1439 + logger.debug_once("Running Cudnn prefill new tokens", scope="local")

1440 assert isinstance(prefill, CudnnPrefillMetadata)

1441 assert prefill.query_seq_lens is not None

1442 + assert fp8_attention is False, "Cudnn prefill does not support fp8 attention"

1443 output, lse = cudnn_batch_prefill_with_kv_cache(

1444 q=q,

1445 k_cache=k,

@@ -1428,9 +1461,19 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1461 return output

1462

1463 def _run_prefill_context_chunk_fa(

1431 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

1464 + self,

1465 + prefill: MLACommonPrefillMetadata,

1466 + chunk_idx: int,

1467 + q,

1468 + k,

1469 + v,

1470 + fp8_attention: bool,

1471 ):

1472 + logger.debug_once("Running FlashAttention prefill context chunk", scope="local")

1473 assert prefill.chunked_context is not None

1474 + assert fp8_attention is False, (

1475 + "FlashAttention prefill does not support fp8 attention"

1476 + )

1477 return self._flash_attn_varlen_diff_headdims(

1478 q=q,

1479 k=k,

@@ -1445,10 +1488,22 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1488 )

1489

1490 def _run_prefill_context_chunk_fi(

1448 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

1491 + self,

1492 + prefill: MLACommonPrefillMetadata,

1493 + chunk_idx: int,

1494 + q,

1495 + k,

1496 + v,

1497 + fp8_attention: bool,

1498 ):

1499 + logger.debug_once("Running FlashInfer prefill context chunk", scope="local")

1500 assert isinstance(prefill, FlashInferPrefillMetadata)

1451 -

1501 + if fp8_attention:

1502 + logger.debug_once("Running FlashInfer prefill in FP8")

1503 + fp8_dtype = current_platform.fp8_dtype()

1504 + q = q.to(fp8_dtype)

1505 + k = k.to(fp8_dtype)

1506 + v = v.to(fp8_dtype)

1507 attn_out, lse = prefill.prefill_chunks[chunk_idx].run(

1508 q=q,

1509 k=k,

@@ -1460,12 +1515,20 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1515 return attn_out, lse.transpose(0, 1).contiguous()

1516

1517 def _run_prefill_context_chunk_cudnn(

1463 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

1518 + self,

1519 + prefill: MLACommonPrefillMetadata,

1520 + chunk_idx: int,

1521 + q,

1522 + k,

1523 + v,

1524 + fp8_attention: bool,

1525 ):

1526 + logger.debug_once("Running Cudnn prefill context chunk", scope="local")

1527 assert isinstance(prefill, CudnnPrefillMetadata)

1528 assert prefill.chunked_context is not None

1529 assert prefill.chunked_context.seq_lens[chunk_idx] is not None

1530 assert prefill.query_seq_lens is not None

1531 + assert fp8_attention is False, "Cudnn prefill does not support fp8 attention"

1532 return cudnn_batch_prefill_with_kv_cache(

1533 q=q,

1534 k_cache=k,

@@ -1485,13 +1548,27 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1548 )

1549

1550 def _run_prefill_new_tokens_trtllm_ragged(

1488 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

1551 + self,

1552 + prefill: MLACommonPrefillMetadata,

1553 + q,

1554 + k,

1555 + v,

1556 + return_softmax_lse,

1557 + fp8_attention: bool,

1558 ):

1559 + logger.debug_once("Running TRT-LLM ragged prefill new tokens", scope="local")

1560 """TRT-LLM ragged attention for new tokens (causal)."""

1561 from flashinfer.prefill import trtllm_ragged_attention_deepseek

1562

1563 assert prefill.query_seq_lens is not None

1564

1565 + if fp8_attention:

1566 + logger.debug_once("Running TRT-LLM ragged prefill in FP8")

1567 + fp8_dtype = current_platform.fp8_dtype()

1568 + q = q.to(fp8_dtype)

1569 + k = k.to(fp8_dtype)

1570 + v = v.to(fp8_dtype)

1571 +

1572 ret = trtllm_ragged_attention_deepseek(

1573 query=q,

1574 key=k,

@@ -1518,8 +1595,15 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1595 return ret

1596

1597 def _run_prefill_context_chunk_trtllm_ragged(

1521 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

1598 + self,

1599 + prefill: MLACommonPrefillMetadata,

1600 + chunk_idx: int,

1601 + q,

1602 + k,

1603 + v,

1604 + fp8_attention: bool,

1605 ):

1606 + logger.debug_once("Running TRT-LLM ragged prefill context chunk", scope="local")

1607 """TRT-LLM ragged attention for context chunks (non-causal)."""

1608 from flashinfer.prefill import trtllm_ragged_attention_deepseek

1609

@@ -1535,6 +1619,13 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1619 )

1620 self._workspace_buffer.fill_(0)

1621

1622 + if fp8_attention:

1623 + logger.debug_once("Running TRT-LLM ragged prefill context chunk in FP8")

1624 + fp8_dtype = current_platform.fp8_dtype()

1625 + q = q.to(fp8_dtype)

1626 + k = k.to(fp8_dtype)

1627 + v = v.to(fp8_dtype)

1628 +

1629 attn_out, lse = trtllm_ragged_attention_deepseek(

1630 query=q,

1631 key=k,

@@ -1687,6 +1778,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1778 kv_c_and_k_pe_cache: torch.Tensor,

1779 attn_metadata: MLACommonMetadata,

1780 k_scale: torch.Tensor,

1781 + fp8_attention: bool,

1782 ):

1783 assert attn_metadata.prefill is not None

1784 prefill_metadata = attn_metadata.prefill

@@ -1725,6 +1817,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1817 q=q,

1818 k=k,

1819 v=v,

1820 + fp8_attention=fp8_attention,

1821 )

1822

1823 if output is None:

@@ -1753,6 +1846,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1846 attn_metadata: MLACommonMetadata,

1847 k_scale: torch.Tensor,

1848 dcp_world_size: int,

1849 + fp8_attention: bool,

1850 ):

1851 assert k_scale is None, "DCP not support scaled kvcache now."

1852 assert attn_metadata.prefill is not None

@@ -1829,6 +1923,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1923 q=q,

1924 k=k,

1925 v=v,

1926 + fp8_attention=fp8_attention,

1927 )

1928

1929 if output is None:

@@ -1859,6 +1954,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1954 attn_metadata: MLACommonMetadata,

1955 k_scale: torch.Tensor,

1956 output: torch.Tensor,

1957 + fp8_attention: bool = False,

1958 ) -> None:

1959 # TODO (zyongye): Prefill function here

1960 assert attn_metadata.prefill is not None

@@ -1878,6 +1974,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1974 k=k,

1975 v=v,

1976 return_softmax_lse=has_context,

1977 + fp8_attention=fp8_attention,

1978 )

1979

1980 if has_context:

@@ -1890,11 +1987,12 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1987 attn_metadata,

1988 k_scale=None,

1989 dcp_world_size=self.dcp_world_size,

1990 + fp8_attention=fp8_attention,

1991 )

1992 )

1993 else:

1994 context_output, context_lse = self._compute_prefill_context(

1897 - q, kv_c_and_k_pe_cache, attn_metadata, k_scale

1995 + q, kv_c_and_k_pe_cache, attn_metadata, k_scale, fp8_attention

1996 )

1997

1998 # unpad if necessary

@@ -2015,6 +2113,7 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

2113 attn_metadata,

2114 layer._k_scale,

2115 output=output[num_decode_tokens:],

2116 + fp8_attention=fp8_attention,

2117 )

2118

2119 if has_decode: