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 14, 05:30 UTC
[Bugfix] fixes the decoding metadata of dense mla's fp8 kvcache. (#27144)
Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Tao He committed
Oct 22, 2025 at 02:27 UTC
250fb1b8ea836ab4cdd2890581815a9c19b8b8ed
3 files changed
+10
-1
cmake/external_projects/flashmla.cmake
+2
-1
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22
- GIT_TAG 5f65b85703c7ed75fda01e06495077caad207c3f
22
+ GIT_TAG 28417e516fcbf6257a422ba117ef5b6f44da5682
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
@@ -66,6 +66,7 @@ if(FLASH_MLA_ARCHS)
66
${flashmla_SOURCE_DIR}/csrc/extension/torch_api.cpp
67
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/pybind.cpp
68
${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_fp8_sm90.cu
69
+ ${flashmla_SOURCE_DIR}/csrc/extension/sm90/dense_fp8/flash_fwd_mla_metadata.cu
70
)
71
72
set(FlashMLA_INCLUDES
vllm/attention/ops/flashmla.py
+6
@@ -102,6 +102,12 @@ def get_mla_metadata(
102
(num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
103
- num_splits: (batch_size + 1), dtype torch.int32.
104
"""
105
+ if is_fp8_kvcache and topk is None:
106
+ return torch.ops._flashmla_extension_C.get_mla_decoding_metadata_dense_fp8(
107
+ cache_seqlens,
108
+ num_q_tokens_per_head_k,
109
+ num_heads_k,
110
+ )
111
return torch.ops._flashmla_C.get_mla_decoding_metadata(
112
cache_seqlens,
113
num_q_tokens_per_head_k,
vllm/v1/attention/backends/mla/flashmla.py
+2
@@ -91,6 +91,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
91
92
self.cg_buf_tile_scheduler_metadata = None
93
self.cg_buf_num_splits = None
94
+ self.is_fp8_kvcache = vllm_config.cache_config.cache_dtype.startswith("fp8")
95
96
device_properties = torch.cuda.get_device_properties(self.device)
97
num_sms = device_properties.multi_processor_count
@@ -123,6 +124,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
124
seq_lens_device,
125
self.num_q_heads,
126
1, # MQA for the decode path
127
+ is_fp8_kvcache=self.is_fp8_kvcache,
128
)
129
130
# TODO: we can disambiguate between decode and mixed-prefill decode here