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 9, 20:45 UTC
[FlashMLA] Update FlashMLA to expose new arguments (#32810)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jan 21, 2026 at 22:02 UTC
889722f3bfdacf54cf2016b94631690b327b9076
8 files changed
+133
-217
.gitignore
+3
@@ -7,6 +7,9 @@ vllm/vllm_flash_attn/*
7
# OpenAI triton kernels copied from source
8
vllm/third_party/triton_kernels/*
9
10
+# FlashMLA interface copied from source
11
+vllm/third_party/flashmla/flash_mla_interface.py
12
+
13
# triton jit
14
.triton
15
cmake/external_projects/flashmla.cmake
+19
-2
@@ -19,7 +19,7 @@ else()
19
FetchContent_Declare(
20
flashmla
21
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
22
- GIT_TAG 526781394b33d9888e4c41952e692266267dd8bf
22
+ GIT_TAG c2afa9cb93e674d5a9120a170a6da57b89267208
23
GIT_PROGRESS TRUE
24
CONFIGURE_COMMAND ""
25
BUILD_COMMAND ""
@@ -30,6 +30,24 @@ endif()
30
FetchContent_MakeAvailable(flashmla)
31
message(STATUS "FlashMLA is available at ${flashmla_SOURCE_DIR}")
32
33
+# Vendor FlashMLA interface into vLLM with torch-ops shim.
34
+set(FLASHMLA_VENDOR_DIR "${CMAKE_SOURCE_DIR}/vllm/third_party/flashmla")
35
+file(MAKE_DIRECTORY "${FLASHMLA_VENDOR_DIR}")
36
+file(READ "${flashmla_SOURCE_DIR}/flash_mla/flash_mla_interface.py"
37
+ FLASHMLA_INTERFACE_CONTENT)
38
+string(REPLACE "import flash_mla.cuda as flash_mla_cuda"
39
+ "import vllm._flashmla_C\nflash_mla_cuda = torch.ops._flashmla_C"
40
+ FLASHMLA_INTERFACE_CONTENT
41
+ "${FLASHMLA_INTERFACE_CONTENT}")
42
+file(WRITE "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
43
+ "${FLASHMLA_INTERFACE_CONTENT}")
44
+
45
+# Install the generated flash_mla_interface.py to the wheel
46
+# Use COMPONENT _flashmla_C to ensure it's installed with the C extension
47
+install(FILES "${FLASHMLA_VENDOR_DIR}/flash_mla_interface.py"
48
+ DESTINATION vllm/third_party/flashmla/
49
+ COMPONENT _flashmla_C)
50
+
51
# The FlashMLA kernels only work on hopper and require CUDA 12.3 or later.
52
# Only build FlashMLA kernels if we are building for something compatible with
53
# sm90a
@@ -79,7 +97,6 @@ if(FLASH_MLA_ARCHS)
97
98
# sm100 dense prefill & backward
99
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_fwd_sm100.cu
82
- ${flashmla_SOURCE_DIR}/csrc/sm100/prefill/dense/fmha_cutlass_bwd_sm100.cu
100
101
# sm100 sparse prefill
102
${flashmla_SOURCE_DIR}/csrc/sm100/prefill/sparse/fwd/head64/instantiations/phase1_k512.cu
setup.py
+10
@@ -646,6 +646,9 @@ class precompiled_wheel_utils:
646
triton_kernels_regex = re.compile(
647
r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
648
)
649
+ flashmla_regex = re.compile(
650
+ r"vllm/third_party/flashmla/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
651
+ )
652
file_members = list(
653
filter(lambda x: x.filename in files_to_copy, wheel.filelist)
654
)
@@ -657,6 +660,9 @@ class precompiled_wheel_utils:
660
lambda x: triton_kernels_regex.match(x.filename), wheel.filelist
661
)
662
)
663
+ file_members += list(
664
+ filter(lambda x: flashmla_regex.match(x.filename), wheel.filelist)
665
+ )
666
667
for file in file_members:
668
print(f"[extract] {file.filename}")
@@ -925,6 +931,10 @@ if _is_cuda():
931
):
932
# FA3 requires CUDA 12.3 or later
933
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
934
+ if envs.VLLM_USE_PRECOMPILED or (
935
+ CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
936
+ ):
937
+ # FlashMLA requires CUDA 12.9 or later
938
# Optional since this doesn't get built (produce an .so file) when
939
# not targeting a hopper system
940
ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
tests/v1/attention/test_sparse_mla_backends.py
-1
@@ -53,7 +53,6 @@ SPARSE_BACKEND_BATCH_SPECS["large_q_pure_prefill"] = BatchSpec(
53
54
def _float_to_e8m0_truncate(f: float) -> float:
55
"""Simulate SM100's float -> e8m0 -> bf16 scale conversion.
56
-
56
e8m0 format only stores the exponent (power of 2).
57
cudaRoundZero truncates toward zero, meaning we round down to the
58
nearest power of 2.
vllm/third_party/flashmla/__init__.py
new
+1
@@ -0,0 +1 @@
1
+# Sources copied from FlashMLA
vllm/v1/attention/backends/mla/flashmla.py
+45
-50
@@ -32,8 +32,11 @@ from vllm.v1.attention.backends.utils import (
32
reshape_query_for_spec_decode,
33
)
34
from vllm.v1.attention.ops.flashmla import (
35
+ FlashMLASchedMeta,
36
flash_mla_with_kvcache,
37
+ flash_mla_with_kvcache_fp8,
38
get_mla_metadata,
39
+ get_mla_metadata_dense_fp8,
40
is_flashmla_dense_supported,
41
)
42
from vllm.v1.kv_cache_interface import AttentionSpec
@@ -93,8 +96,7 @@ class FlashMLABackend(MLACommonBackend):
96
97
@dataclass
98
class FlashMLADecodeMetadata(MLACommonDecodeMetadata):
96
- tile_scheduler_metadata: torch.Tensor
97
- num_splits: torch.Tensor
99
+ scheduler_metadata: FlashMLASchedMeta
100
101
102
@dataclass
@@ -158,46 +160,25 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
160
# we use the max but all should be the same due to uniform length requirement
161
max_query_len = query_lens_cpu.max().item()
162
num_q_tokens_per_head_k = max_query_len * self.num_q_heads // 1
161
- tile_scheduler_metadata, num_splits = get_mla_metadata(
163
+ scheduler_metadata, _ = get_mla_metadata(
164
seq_lens_device,
165
num_q_tokens_per_head_k,
166
1, # MQA for the decode path
167
is_fp8_kvcache=self.is_fp8_kvcache,
168
)
167
-
168
- # TODO: we can disambiguate between decode and mixed-prefill decode here
169
- # so we can only use the persistent buffer if a cudagraph is actually
170
- # being used.
171
- if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
172
- assert self.cg_buf_tile_scheduler_metadata is not None
173
- assert self.cg_buf_num_splits is not None
174
-
175
- sm_parts = tile_scheduler_metadata.size(0)
176
- # Metadata per-SM, upper bound on size (<= #SMs, TileMetadataSize)
177
- assert sm_parts <= self.cg_buf_tile_scheduler_metadata.size(0)
178
- tile_scheduler_metadata_view = self.cg_buf_tile_scheduler_metadata[
179
- :sm_parts
180
- ]
181
- tile_scheduler_metadata_view.copy_(tile_scheduler_metadata)
182
- tile_scheduler_metadata = tile_scheduler_metadata_view
183
-
184
- # Num splits is per-batch, varying size (batch_size,)
185
- n = num_splits.size(0)
186
- # make sure static buffer is large enough
187
- assert n <= self.cg_buf_num_splits.size(0)
188
- num_splits_view = self.cg_buf_num_splits[:n]
189
- num_splits_view.copy_(num_splits)
190
- # Num splits needs to monotonically increasing
191
- # (with: https://github.com/vllm-project/FlashMLA/pull/3, otherwise
192
- # it needs to monotonically increasing by 1)
193
- self.cg_buf_num_splits[n:].fill_(num_splits[-1])
194
- num_splits = num_splits_view
169
+ if self.is_fp8_kvcache:
170
+ tile_scheduler_metadata, num_splits = get_mla_metadata_dense_fp8(
171
+ seq_lens_device,
172
+ num_q_tokens_per_head_k,
173
+ 1, # MQA for the decode path
174
+ )
175
+ scheduler_metadata.tile_scheduler_metadata = tile_scheduler_metadata
176
+ scheduler_metadata.num_splits = num_splits
177
178
return FlashMLADecodeMetadata(
179
block_table=block_table_tensor,
180
seq_lens=seq_lens_device,
199
- tile_scheduler_metadata=tile_scheduler_metadata,
200
- num_splits=num_splits,
181
+ scheduler_metadata=scheduler_metadata,
182
dcp_tot_seq_lens=dcp_tot_seq_lens_device,
183
)
184
@@ -272,9 +253,8 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
253
num_decodes = attn_metadata.num_decodes
254
q = reshape_query_for_spec_decode(q, num_decodes)
255
275
- tile_scheduler_metadata = attn_metadata.decode.tile_scheduler_metadata
276
- num_splits = attn_metadata.decode.num_splits
277
- if vllm_is_batch_invariant():
256
+ scheduler_metadata = attn_metadata.decode.scheduler_metadata
257
+ if vllm_is_batch_invariant() and not self.kv_cache_dtype.startswith("fp8"):
258
device = q.device
259
dtype = torch.int32
260
@@ -301,20 +281,35 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
281
# Non-split path ignores num_splits, but the API requires it:
282
# zeros of length B+1
283
num_splits = torch.zeros((B + 1,), dtype=dtype, device=device)
304
-
305
- o, lse = flash_mla_with_kvcache(
306
- q=q,
307
- k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1
308
- block_table=attn_metadata.decode.block_table,
309
- cache_seqlens=attn_metadata.decode.seq_lens,
310
- head_dim_v=self.kv_lora_rank,
311
- tile_scheduler_metadata=tile_scheduler_metadata,
312
- num_splits=num_splits,
313
- softmax_scale=self.scale,
314
- causal=True,
315
- descale_q=layer._q_scale.reshape(1),
316
- descale_k=layer._k_scale.reshape(1),
317
- )
284
+ scheduler_metadata.tile_scheduler_metadata = tile_scheduler_metadata
285
+ scheduler_metadata.num_splits = num_splits
286
+
287
+ if self.kv_cache_dtype.startswith("fp8"):
288
+ o, lse = flash_mla_with_kvcache_fp8(
289
+ q=q,
290
+ k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1
291
+ block_table=attn_metadata.decode.block_table,
292
+ cache_seqlens=attn_metadata.decode.seq_lens,
293
+ head_dim_v=self.kv_lora_rank,
294
+ tile_scheduler_metadata=scheduler_metadata.tile_scheduler_metadata,
295
+ num_splits=scheduler_metadata.num_splits,
296
+ softmax_scale=self.scale,
297
+ causal=True,
298
+ descale_q=layer._q_scale.reshape(1),
299
+ descale_k=layer._k_scale.reshape(1),
300
+ )
301
+ else:
302
+ o, lse = flash_mla_with_kvcache(
303
+ q=q,
304
+ k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1
305
+ block_table=attn_metadata.decode.block_table,
306
+ cache_seqlens=attn_metadata.decode.seq_lens,
307
+ head_dim_v=self.kv_lora_rank,
308
+ tile_scheduler_metadata=scheduler_metadata,
309
+ softmax_scale=self.scale,
310
+ causal=True,
311
+ is_fp8_kvcache=False,
312
+ )
313
314
o = reshape_attn_output_for_spec_decode(o)
315
vllm/v1/attention/backends/mla/flashmla_sparse.py
+8
-29
@@ -33,7 +33,8 @@ from vllm.v1.attention.backends.utils import (
33
split_prefill_chunks,
34
)
35
from vllm.v1.attention.ops.flashmla import (
36
- flash_mla_sparse_prefill,
36
+ FlashMLASchedMeta,
37
+ flash_mla_sparse_fwd,
38
flash_mla_with_kvcache,
39
get_mla_metadata,
40
)
@@ -142,8 +143,7 @@ class FlashMLASparseMetadata(AttentionMetadata):
143
144
@dataclass
145
class FP8KernelMetadata:
145
- scheduler_metadata: torch.Tensor | None
146
- num_splits: torch.Tensor
146
+ scheduler_metadata: FlashMLASchedMeta
147
dummy_block_table: torch.Tensor
148
cache_lens: torch.Tensor
149
@@ -468,7 +468,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
468
padded_heads = self.fp8_decode_padded_heads
469
470
# Build metadata for all tokens as a single batch
471
- tile_scheduler_metadata, num_splits = get_mla_metadata(
471
+ scheduler_metadata, _ = get_mla_metadata(
472
cache_seqlens=self.topk_tokens_tensor[:1], # Single batch
473
num_q_tokens_per_head_k=num_tokens * padded_heads,
474
topk=self.topk_tokens,
@@ -477,17 +477,8 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
477
is_fp8_kvcache=True,
478
)
479
480
- num_sm_parts = tile_scheduler_metadata.size(0)
481
- tile_scheduler_metadata_buffer = self.tile_scheduler_metadata_buffer[
482
- :num_sm_parts
483
- ]
484
- tile_scheduler_metadata_buffer.copy_(tile_scheduler_metadata)
485
- num_splits_view = self.num_splits_buffer[:2]
486
- num_splits_view.copy_(num_splits)
487
-
480
fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata(
489
- scheduler_metadata=tile_scheduler_metadata_buffer,
490
- num_splits=num_splits_view,
481
+ scheduler_metadata=scheduler_metadata,
482
cache_lens=self.max_model_len_tensor[:1],
483
dummy_block_table=self.dummy_block_table[:1],
484
)
@@ -620,7 +611,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
611
612
# Use padded head count since that's what the kernel will see
613
padded_heads = self.fp8_decode_padded_heads
623
- tile_scheduler_metadata, num_splits = get_mla_metadata(
614
+ scheduler_metadata, _ = get_mla_metadata(
615
cache_seqlens=self.topk_tokens_tensor[:num_decodes],
616
num_q_tokens_per_head_k=decode_query_len * padded_heads,
617
topk=self.topk_tokens,
@@ -629,19 +620,8 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
620
is_fp8_kvcache=True,
621
)
622
632
- num_sm_parts = tile_scheduler_metadata.size(0)
633
- # Copy to persistent buffer for full-CG support
634
- tile_scheduler_metadata_buffer = self.tile_scheduler_metadata_buffer[
635
- :num_sm_parts
636
- ]
637
- tile_scheduler_metadata_buffer.copy_(tile_scheduler_metadata)
638
- # num_splits has size [num_decodes + 1]
639
- num_splits_view = self.num_splits_buffer[: num_decodes + 1]
640
- num_splits_view.copy_(num_splits)
641
-
623
kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata(
643
- scheduler_metadata=tile_scheduler_metadata_buffer,
644
- num_splits=num_splits_view,
624
+ scheduler_metadata=scheduler_metadata,
625
dummy_block_table=self.dummy_block_table[:num_decodes],
626
cache_lens=self.max_model_len_tensor[:num_decodes],
627
)
@@ -949,7 +929,6 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
929
head_dim_v=512,
930
cache_seqlens=kernel_metadata.cache_lens,
931
tile_scheduler_metadata=kernel_metadata.scheduler_metadata,
952
- num_splits=kernel_metadata.num_splits,
932
is_fp8_kvcache=True,
933
indices=topk_indices,
934
softmax_scale=self.softmax_scale,
@@ -985,7 +964,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):
964
q = q_padded
965
966
topk_indices = topk_indices.view(num_tokens, 1, -1)
988
- output = flash_mla_sparse_prefill(
967
+ output = flash_mla_sparse_fwd(
968
q, kv_c_and_k_pe_cache, topk_indices, self.softmax_scale
969
)[0]
970
output = output[:, : self.num_heads, :]
vllm/v1/attention/ops/flashmla.py
+47
-135
@@ -78,50 +78,49 @@ def is_flashmla_sparse_supported() -> tuple[bool, str | None]:
78
return True, None
79
80
81
-def get_mla_metadata(
81
+def _raise_flashmla_unavailable(*_args, **_kwargs):
82
+ _, reason = _is_flashmla_available()
83
+ raise RuntimeError(reason or "FlashMLA is not available")
84
+
85
+
86
+if _is_flashmla_available()[0]:
87
+ from vllm.third_party.flashmla.flash_mla_interface import ( # noqa: F401
88
+ FlashMLASchedMeta,
89
+ flash_attn_varlen_func,
90
+ flash_attn_varlen_kvpacked_func,
91
+ flash_attn_varlen_qkvpacked_func,
92
+ flash_mla_sparse_fwd,
93
+ flash_mla_with_kvcache,
94
+ get_mla_metadata,
95
+ )
96
+else:
97
+
98
+ class FlashMLASchedMeta: # type: ignore[no-redef]
99
+ pass
100
+
101
+ flash_attn_varlen_func = _raise_flashmla_unavailable # type: ignore[assignment]
102
+ flash_attn_varlen_kvpacked_func = _raise_flashmla_unavailable # type: ignore[assignment]
103
+ flash_attn_varlen_qkvpacked_func = _raise_flashmla_unavailable # type: ignore[assignment]
104
+ flash_mla_sparse_fwd = _raise_flashmla_unavailable # type: ignore[assignment]
105
+ flash_mla_with_kvcache = _raise_flashmla_unavailable # type: ignore[assignment]
106
+ get_mla_metadata = _raise_flashmla_unavailable # type: ignore[assignment]
107
+
108
+
109
+def get_mla_metadata_dense_fp8(
110
cache_seqlens: torch.Tensor,
111
num_q_tokens_per_head_k: int,
112
num_heads_k: int,
85
- num_heads_q: int | None = None,
86
- is_fp8_kvcache: bool = False,
87
- topk: int | None = None,
113
) -> tuple[torch.Tensor, torch.Tensor]:
89
- """
90
- Arguments:
91
- - cache_seqlens: (batch_size), dtype torch.int32.
92
- - num_q_tokens_per_head_k:
93
- Equals to num_q_tokens_per_q_seq * num_heads_q // num_heads_k.
94
- - num_heads_k: The number of k heads.
95
- - num_heads_q:
96
- The number of q heads.
97
- This argument is optional when sparse attention is not enabled
98
- - is_fp8_kvcache: Whether the k_cache and v_cache are in fp8 format.
99
- - topk: If not None, sparse attention will be enabled,
100
- and only tokens in the `indices` array
101
- passed to `flash_mla_with_kvcache_sm90` will be attended to.
102
-
103
- Returns:
104
- - tile_scheduler_metadata:
105
- (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
106
- - num_splits: (batch_size + 1), dtype torch.int32.
107
- """
108
- if is_fp8_kvcache and topk is None:
109
- return torch.ops._flashmla_extension_C.get_mla_decoding_metadata_dense_fp8(
110
- cache_seqlens,
111
- num_q_tokens_per_head_k,
112
- num_heads_k,
113
- )
114
- return torch.ops._flashmla_C.get_mla_decoding_metadata(
114
+ if not _is_flashmla_available()[0]:
115
+ _raise_flashmla_unavailable()
116
+ return torch.ops._flashmla_extension_C.get_mla_decoding_metadata_dense_fp8(
117
cache_seqlens,
118
num_q_tokens_per_head_k,
119
num_heads_k,
118
- num_heads_q,
119
- is_fp8_kvcache,
120
- topk,
120
)
121
122
124
-def flash_mla_with_kvcache(
123
+def flash_mla_with_kvcache_fp8(
124
q: torch.Tensor,
125
k_cache: torch.Tensor,
126
block_table: torch.Tensor,
@@ -133,114 +132,27 @@ def flash_mla_with_kvcache(
132
causal: bool = False,
133
descale_q: torch.Tensor | None = None,
134
descale_k: torch.Tensor | None = None,
136
- is_fp8_kvcache: bool = False,
137
- indices: torch.Tensor | None = None,
135
) -> tuple[torch.Tensor, torch.Tensor]:
139
- """
140
- Arguments:
141
- - q: (batch_size, seq_len_q, num_heads_q, head_dim).
142
- - k_cache: (num_blocks, page_block_size, num_heads_k, head_dim).
143
- - block_table: (batch_size, max_num_blocks_per_seq), torch.int32.
144
- - cache_seqlens: (batch_size), torch.int32.
145
- - head_dim_v: Head dimension of v.
146
- - tile_scheduler_metadata:
147
- (num_sm_parts, TileSchedulerMetaDataSize), torch.int32,
148
- returned by get_mla_metadata.
149
- - num_splits:
150
- (batch_size + 1), torch.int32, returned by get_mla_metadata.
151
- - softmax_scale: float.
152
- The scale of QK^T before applying softmax.
153
- Default to 1 / sqrt(head_dim).
154
- - causal: bool. Whether to apply causal attention mask.
155
- - descale_q: (batch_size),
156
- torch.float32. Descaling factors for Q, used for fp8 quantization.
157
- - descale_k: (batch_size),
158
- torch.float32. Descaling factors for K, used for fp8 quantization.
159
- - is_fp8_kvcache: bool.
160
- Whether the k_cache and v_cache are in fp8 format.
161
- For the format of FP8 KV cache, please refer to README.md
162
- - indices: (batch_size, seq_len_q, topk), torch.int32.
163
- If not None, sparse attention will be enabled,
164
- and only tokens in the `indices` array will be attended to.
165
- Invalid indices should be set to -1 or numbers >= total_seq_len_kv.
166
- For details about how to set up `indices`, please refer to README.md.
167
-
168
- Returns:
169
- - out: (batch_size, seq_len_q, num_heads_q, head_dim_v).
170
- - softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32.
171
- """
136
+ if not _is_flashmla_available()[0]:
137
+ _raise_flashmla_unavailable()
138
if softmax_scale is None:
139
softmax_scale = q.shape[-1] ** (-0.5)
174
- if indices is not None:
175
- # NOTE (zyongye): sparse attention is also causal
176
- # since it only attend to the tokens before
177
- # but here `causal` should not be specified
178
- assert not causal, "causal must be `false` if sparse attention is enabled."
179
- assert (descale_q is None) == (descale_k is None), (
180
- "descale_q and descale_k should be both None or both not None"
140
+ out, softmax_lse = torch.ops._flashmla_extension_C.fwd_kvcache_mla_fp8(
141
+ q,
142
+ k_cache,
143
+ head_dim_v,
144
+ cache_seqlens,
145
+ block_table,
146
+ softmax_scale,
147
+ causal,
148
+ tile_scheduler_metadata,
149
+ num_splits,
150
+ descale_q,
151
+ descale_k,
152
)
182
-
183
- if indices is None and q.element_size() == 1:
184
- out, softmax_lse = torch.ops._flashmla_extension_C.fwd_kvcache_mla_fp8(
185
- q,
186
- k_cache,
187
- head_dim_v,
188
- cache_seqlens,
189
- block_table,
190
- softmax_scale,
191
- causal,
192
- tile_scheduler_metadata,
193
- num_splits,
194
- descale_q,
195
- descale_k,
196
- )
197
- else:
198
- out, softmax_lse = torch.ops._flashmla_C.fwd_kvcache_mla(
199
- q,
200
- k_cache,
201
- head_dim_v,
202
- cache_seqlens,
203
- block_table,
204
- softmax_scale,
205
- causal,
206
- tile_scheduler_metadata,
207
- num_splits,
208
- is_fp8_kvcache,
209
- indices,
210
- )
153
return out, softmax_lse
154
155
214
-def flash_mla_sparse_prefill(
215
- q: torch.Tensor,
216
- kv: torch.Tensor,
217
- indices: torch.Tensor,
218
- sm_scale: float,
219
- d_v: int = 512,
220
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
221
- """
222
- Sparse attention prefill kernel
223
-
224
- Args:
225
- - q: [s_q, h_q, d_qk], bfloat16
226
- - kv: [s_kv, h_kv, d_qk], bfloat16
227
- - indices: [s_q, h_kv, topk], int32.
228
- Invalid indices should be set to -1 or numbers >= s_kv
229
- - sm_scale: float
230
- - d_v: The dimension of value vectors. Can only be 512
231
-
232
- Returns:
233
- - (output, max_logits, lse)
234
- About the definition of output,
235
- max_logits and lse, please refer to README.md
236
- - output: [s_q, h_q, d_v], bfloat16
237
- - max_logits: [s_q, h_q], float
238
- - lse: [s_q, h_q], float, 2-based log-sum-exp
239
- """
240
- results = torch.ops._flashmla_C.sparse_prefill_fwd(q, kv, indices, sm_scale, d_v)
241
- return results
242
-
243
-
156
#
157
# TODO: Add fake functions
158
#