From e43db1f5e2eb7ca16fc284a810c548a0d3d44b14 Mon Sep 17 00:00:00 2001 From: PatrykSaffer Date: Tue, 6 Oct 2026 22:03:05 +0200 Subject: [PATCH] Pad to 64 not 128 Q heads in flashMLA sparse (#60029) Signed-off-by: Patryk Saffer Co-authored-by: Misha Goin --- vllm/models/hy_v4/nvidia/flashmla_sparse.py | 6 +++--- vllm/v1/attention/backends/mla/flashmla_sparse.py | 12 +++++------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/vllm/models/hy_v4/nvidia/flashmla_sparse.py b/vllm/models/hy_v4/nvidia/flashmla_sparse.py index 1d62f2f556..52656b06a6 100644 --- a/vllm/models/hy_v4/nvidia/flashmla_sparse.py +++ b/vllm/models/hy_v4/nvidia/flashmla_sparse.py @@ -228,9 +228,9 @@ class HYV4FlashMLASparseImpl(FlashMLASparseImpl): -1, 1, kv_c_and_k_pe_cache.shape[-1] ) - # NOTE(Chen): kernel requires num_local_head to be a multiple of - # 64 on hopper and 128 on blackwell. Pad from q's head count, not - # self.num_heads: under DCP the heads are all-gathered before this. + # NOTE(Chen): kernel requires num_local_head to be a multiple of 64. + # Pad from q's head count, not self.num_heads: under DCP the heads are + # all-gathered before this. if actual_num_heads is None: actual_num_heads = q.shape[1] padded_num_heads = ( diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index dea9b81116..d441934c02 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -730,10 +730,8 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): **mla_args, ) self.softmax_scale = scale - # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell - self.prefill_padding = ( - 128 if current_platform.is_device_capability_family(100) else 64 - ) + # Prefill BF16 kernel requires 64 heads on Hopper and Blackwell. + self.prefill_padding = 64 self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads) # The kernels read queries as wide as the cache rows. A NoPE model on @@ -1323,9 +1321,9 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): -1, 1, kv_c_and_k_pe_cache.shape[-1] ) - # NOTE(Chen): kernel requires num_local_head to be a multiple of - # 64 on hopper and 128 on blackwell. Pad from q's head count, not - # self.num_heads: under DCP the heads are all-gathered before this. + # NOTE(Chen): kernel requires num_local_head to be a multiple of 64. + # Pad from q's head count, not self.num_heads: under DCP the heads are + # all-gathered before this. if actual_num_heads is None: actual_num_heads = q.shape[1] padded_num_heads = (