Support DCP with FlashInfer MLA (#43729)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Jun 29, 2026 at 13:24 UTC
a309d4fe60bef7657b88805a1fcc9b014c414314
2 files changed
+11
-5
docs/design/attention_backends.md
+1
-1
@@ -220,7 +220,7 @@ MLA decode backends are selected using the standard
220
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
221
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
222
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
223
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
223
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
224
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
225
| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
226
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
vllm/v1/attention/backends/mla/flashinfer_mla.py
+10
-4
@@ -113,6 +113,8 @@ g_fi_workspace = torch.zeros(
113
114
115
class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
116
+ can_return_lse_for_decode: bool = True
117
+
118
def __init__(
119
self,
120
num_heads: int,
@@ -196,7 +198,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
198
if is_quantized_kv_cache(self.kv_cache_dtype):
199
self.bmm2_scale *= layer._k_scale_float
200
199
- o = trtllm_batch_decode_with_kv_cache_mla(
201
+ return_lse = self.need_to_return_lse_for_decode
202
+ kernel_out = trtllm_batch_decode_with_kv_cache_mla(
203
query=q,
204
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
205
workspace_buffer=self._workspace_buffer,
@@ -208,11 +211,14 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
211
max_seq_len=attn_metadata.max_seq_len,
212
bmm1_scale=self.bmm1_scale,
213
bmm2_scale=self.bmm2_scale,
214
+ return_lse=return_lse,
215
)
216
+ if return_lse:
217
+ o, lse = kernel_out
218
+ else:
219
+ o, lse = kernel_out, None
220
221
# Flatten the output for consistent shape
222
o = o.view(-1, o.shape[-2], o.shape[-1])
223
216
- # TODO: Return LSE pending support from Flashinfer API:
217
- # https://github.com/flashinfer-ai/flashinfer/pull/1566
218
- return o, None
224
+ return o, lse