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 = (