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] Fix IMA FlashMLA full cuda-graph and DP + Update FlashMLA (#21691)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Lucas Wilkinson committed
Aug 8, 2025 at 19:09 UTC
cd9b9de1fb009cf607403ba08961f2a3f869931d
3 files changed
+42
-27
cmake/external_projects/flashmla.cmake
+4
-4
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA.git
22
- GIT_TAG 575f7724b9762f265bbee5889df9c7d630801845
22
+ GIT_TAG 0e43e774597682284358ff2c54530757b654b8d1
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
@@ -37,9 +37,9 @@ cuda_archs_loose_intersection(FLASH_MLA_ARCHS "9.0a" "${CUDA_ARCHS}")
37
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.3 AND FLASH_MLA_ARCHS)
38
set(FlashMLA_SOURCES
39
${flashmla_SOURCE_DIR}/csrc/flash_api.cpp
40
- ${flashmla_SOURCE_DIR}/csrc/flash_fwd_mla_bf16_sm90.cu
41
- ${flashmla_SOURCE_DIR}/csrc/flash_fwd_mla_fp16_sm90.cu
42
- ${flashmla_SOURCE_DIR}/csrc/flash_fwd_mla_metadata.cu)
40
+ ${flashmla_SOURCE_DIR}/csrc/kernels/splitkv_mla.cu
41
+ ${flashmla_SOURCE_DIR}/csrc/kernels/mla_combine.cu
42
+ ${flashmla_SOURCE_DIR}/csrc/kernels/get_mla_metadata.cu)
43
44
set(FlashMLA_INCLUDES
45
${flashmla_SOURCE_DIR}/csrc/cutlass/include
vllm/attention/ops/flashmla.py
-1
@@ -91,7 +91,6 @@ def flash_mla_with_kvcache(
91
out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
92
q,
93
k_cache,
94
- None,
94
head_dim_v,
95
cache_seqlens,
96
block_table,
vllm/v1/attention/backends/mla/flashmla.py
+38
-22
@@ -70,6 +70,22 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
70
self.cg_buf_tile_scheduler_metadata = None
71
self.cg_buf_num_splits = None
72
73
+ device_properties = torch.cuda.get_device_properties(self.device)
74
+ num_sms = device_properties.multi_processor_count
75
+
76
+ if self.compilation_config.full_cuda_graph:
77
+ self.cg_buf_tile_scheduler_metadata = torch.zeros(
78
+ # Upper bound on size (<= #SMs, TileSchedulerMetaDataSize)
79
+ # TileSchedulerMetaDataSize = 8
80
+ (num_sms, 8),
81
+ device=self.device,
82
+ dtype=torch.int32,
83
+ )
84
+ self.cg_buf_num_splits = torch.empty(
85
+ (vllm_config.scheduler_config.max_num_seqs + 1),
86
+ device=self.device,
87
+ dtype=torch.int32)
88
+
89
def _build_decode(self, block_table_tensor: torch.Tensor,
90
seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:
91
tile_scheduler_metadata, num_splits = \
@@ -80,28 +96,28 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
96
)
97
98
if self.compilation_config.full_cuda_graph:
83
- # First time around (CUDAGraph capture), allocate the static buffer
84
- if self.cg_buf_tile_scheduler_metadata is None:
85
- self.cg_buf_tile_scheduler_metadata = tile_scheduler_metadata
86
- self.cg_buf_num_splits = num_splits
87
- else:
88
- assert self.cg_buf_num_splits is not None
89
-
90
- # Metadata per-SM, fixed size (#SMs, TileMetadataSize)
91
- assert (self.cg_buf_tile_scheduler_metadata.size() ==
92
- tile_scheduler_metadata.size())
93
- self.cg_buf_tile_scheduler_metadata.\
94
- copy_(tile_scheduler_metadata)
95
- tile_scheduler_metadata = self.cg_buf_tile_scheduler_metadata
96
-
97
- # Num splits is per-batch, varying size (batch_size,)
98
- n = num_splits.size(0)
99
- # make sure static buffer is large enough
100
- assert n <= self.cg_buf_num_splits.size(0)
101
- num_splits_view = self.cg_buf_num_splits[:n]
102
- num_splits_view.copy_(num_splits)
103
- self.cg_buf_num_splits[n:].fill_(0) # fill the rest with 0s
104
- num_splits = num_splits_view
99
+ assert self.cg_buf_tile_scheduler_metadata is not None
100
+ assert self.cg_buf_num_splits is not None
101
+
102
+ sm_parts = tile_scheduler_metadata.size(0)
103
+ # Metadata per-SM, upper bound on size (<= #SMs, TileMetadataSize)
104
+ assert sm_parts <= self.cg_buf_tile_scheduler_metadata.size(0)
105
+ tile_scheduler_metadata_view = \
106
+ self.cg_buf_tile_scheduler_metadata[:sm_parts]
107
+ tile_scheduler_metadata_view.copy_(tile_scheduler_metadata)
108
+ tile_scheduler_metadata = tile_scheduler_metadata_view
109
+
110
+ # Num splits is per-batch, varying size (batch_size,)
111
+ n = num_splits.size(0)
112
+ # make sure static buffer is large enough
113
+ assert n <= self.cg_buf_num_splits.size(0)
114
+ num_splits_view = self.cg_buf_num_splits[:n]
115
+ num_splits_view.copy_(num_splits)
116
+ # Num splits needs to monotonically increasing
117
+ # (with: https://github.com/vllm-project/FlashMLA/pull/3, otherwise
118
+ # it needs to monotonically increasing by 1)
119
+ self.cg_buf_num_splits[n:].fill_(num_splits[-1])
120
+ num_splits = num_splits_view
121
122
return FlashMLADecodeMetadata(
123
block_table=block_table_tensor,