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][deepseek] fix flashmla kernel selection (#25956)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Oct 1, 2025 at 00:30 UTC
a2e6fa7e035ff058fc37fdaaf014707efff2fcf3
1 file changed
+1
-1
vllm/attention/ops/flashmla.py
+1
-1
@@ -136,7 +136,7 @@ def flash_mla_with_kvcache(
136
descale_k is None
137
), "descale_q and descale_k should be both None or both not None"
138
139
- if (descale_q is not None) and (descale_k is not None):
139
+ if indices is None and q.element_size() == 1:
140
out, softmax_lse = torch.ops._flashmla_extension_C.fwd_kvcache_mla_fp8(
141
q, k_cache, head_dim_v, cache_seqlens, block_table, softmax_scale,
142
causal, tile_scheduler_metadata, num_splits, descale_q, descale_k)