diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2870ec9a15..0bc7ca7aa4 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -231,8 +231,6 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): reorder_batch_threshold: int = 1 - natively_supported_next_n_fp4: list[int] = [1, 2] - # TODO (matt): integrate kernel with next_n = 4 support @classmethod def get_cudagraph_support( @@ -267,15 +265,21 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): next_n = self.num_speculative_tokens + 1 self.reorder_batch_threshold += self.num_speculative_tokens - # NOTE(zyongye) fp4 indexer cache only natively supports next_n in - # natively_supported_next_n_fp4; for other next_n values we fall back - # to the flattening path. Outside the SM100 datacenter family the FP8 - # paged MQA logits kernel has the same [1, 2] constraint (deepgemm - # smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there too. - self.use_flattening = ( - self.use_fp4_indexer_cache - or not current_platform.is_device_capability_family(100) - ) and next_n not in self.natively_supported_next_n_fp4 + # NOTE: SM100 datacenter GPUs support any next_n natively via the + # multi-atom paged MQA logits kernels (FP8 and FP4 indexer + # caches). Outside the SM100 family the FP8 + # paged MQA logits kernel only supports next_n in (1, 2) + # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there. + self.use_flattening = not current_platform.is_device_capability_family( + 100 + ) and next_n not in (1, 2) + logger.info_once( + "DSA indexer decode path: use_flattening=%s " + "(next_n=%d, use_fp4_indexer_cache=%s)", + self.use_flattening, + next_n, + self.use_fp4_indexer_cache, + ) sm_count = num_compute_units(self.device.index) self.num_sms = sm_count