A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf] Use native DSA indexer decode path for next_n > 2 on SM100 (#45322)
Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
qizixi committed
Jun 12, 2026 at 12:54 UTC
efe7adb5e145de0de2a691cc86756f088f4f01d0
1 file changed
+15
-11
vllm/v1/attention/backends/mla/indexer.py
+15
-11
@@ -231,8 +231,6 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
231
232
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
233
reorder_batch_threshold: int = 1
234
- natively_supported_next_n_fp4: list[int] = [1, 2]
235
- # TODO (matt): integrate kernel with next_n = 4 support
234
235
@classmethod
236
def get_cudagraph_support(
@@ -267,15 +265,21 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
265
266
next_n = self.num_speculative_tokens + 1
267
self.reorder_batch_threshold += self.num_speculative_tokens
270
- # NOTE(zyongye) fp4 indexer cache only natively supports next_n in
271
- # natively_supported_next_n_fp4; for other next_n values we fall back
272
- # to the flattening path. Outside the SM100 datacenter family the FP8
273
- # paged MQA logits kernel has the same [1, 2] constraint (deepgemm
274
- # smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there too.
275
- self.use_flattening = (
276
- self.use_fp4_indexer_cache
277
- or not current_platform.is_device_capability_family(100)
278
- ) and next_n not in self.natively_supported_next_n_fp4
268
+ # NOTE: SM100 datacenter GPUs support any next_n natively via the
269
+ # multi-atom paged MQA logits kernels (FP8 and FP4 indexer
270
+ # caches). Outside the SM100 family the FP8
271
+ # paged MQA logits kernel only supports next_n in (1, 2)
272
+ # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there.
273
+ self.use_flattening = not current_platform.is_device_capability_family(
274
+ 100
275
+ ) and next_n not in (1, 2)
276
+ logger.info_once(
277
+ "DSA indexer decode path: use_flattening=%s "
278
+ "(next_n=%d, use_fp4_indexer_cache=%s)",
279
+ self.use_flattening,
280
+ next_n,
281
+ self.use_fp4_indexer_cache,
282
+ )
283
284
sm_count = num_compute_units(self.device.index)
285
self.num_sms = sm_count