[Config] Unify indexer cache dtype under attention_config.indexer_kv_dtype (#52550)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Yongye Zhu committed
Aug 16, 2026 at 19:20 UTC
967e104fad12758e1fc4f16b7f9fdebfc3b6b1bb
6 files changed
+65
-24
tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml
+1
-1
@@ -15,7 +15,7 @@ server_args: >-
15
--block-size 256
16
--gpu-memory-utilization 0.5
17
--kv-cache-dtype fp8
18
- --attention_config.use_fp4_indexer_cache=True
18
+ --attention_config.indexer_kv_dtype=mxfp4
19
--max-num-batched-tokens 16384
20
--max-num-seqs 128
21
--speculative-config '{"method":"dspark",
tests/evals/gsm8k/configs/moe-refactor/DeepSeek-V4-Flash-deepgemm-mega-moe.yaml
+1
-1
@@ -2,4 +2,4 @@ model_name: "deepseek-ai/DeepSeek-V4-Flash"
2
accuracy_threshold: 0.95
3
num_questions: 1319
4
num_fewshot: 5
5
-server_args: "--trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 2 --attention_config.use_fp4_indexer_cache=True --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --speculative_config.method=mtp --speculative_config.num_speculative_tokens=2"
5
+server_args: "--trust-remote-code --kv-cache-dtype fp8 --block-size 256 --enable-expert-parallel --tensor-parallel-size 2 --attention_config.indexer_kv_dtype=mxfp4 --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 --enable-auto-tool-choice --reasoning-parser deepseek_v4 --speculative_config.method=mtp --speculative_config.num_speculative_tokens=2"
vllm/config/attention.py
+34
-7
@@ -7,10 +7,13 @@ from typing import Any, Literal
7
from pydantic import field_validator
8
9
from vllm.config.utils import config
10
+from vllm.logger import init_logger
11
from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
12
from vllm.v1.attention.backends.registry import AttentionBackendEnum
13
13
-IndexerKVDType = Literal["bf16", "fp8", "mxfp4", "nvfp4"]
14
+logger = init_logger(__name__)
15
+
16
+IndexerKVDType = Literal["auto", "bf16", "fp8", "mxfp4", "nvfp4"]
17
MiniMaxM3MSADecodeBackend = Literal["triton", "cutlass"]
18
19
@@ -65,12 +68,15 @@ class AttentionConfig:
68
use_prefill_query_quantization: bool = False
69
"""If set, quantize query for attention in prefill."""
70
68
- use_fp4_indexer_cache: bool = False
69
- """If set, use fp4 indexer cache for dsv32 family model (not support yet)"""
71
+ use_fp4_indexer_cache: bool | None = None
72
+ """Deprecated alias for `indexer_kv_dtype`; use that instead. True maps to
73
+ `mxfp4`, False is a no-op (it selected the model default already)."""
74
71
- indexer_kv_dtype: IndexerKVDType = "bf16"
72
- """Data type for the sparse-attention indexer K cache. Quantized formats
73
- (fp8, mxfp4, nvfp4) require indexer kernel support in the backend."""
75
+ indexer_kv_dtype: IndexerKVDType = "auto"
76
+ """Data type for the sparse-attention indexer K cache. "auto" picks the
77
+ model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse
78
+ indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel
79
+ support in the backend."""
80
81
use_non_causal: bool = False
82
"""Whether to use non-causal (bidirectional) attention."""
@@ -114,6 +120,26 @@ class AttentionConfig:
120
# layers still use the platform's normal automatic backend.
121
self.backend = None
122
123
+ if self.use_fp4_indexer_cache is not None:
124
+ logger.warning(
125
+ "use_fp4_indexer_cache is deprecated and will be removed in "
126
+ "v0.19. Use indexer_kv_dtype instead (True -> 'mxfp4')."
127
+ )
128
+ if self.use_fp4_indexer_cache:
129
+ if self.indexer_kv_dtype not in ("auto", "mxfp4"):
130
+ raise ValueError(
131
+ "use_fp4_indexer_cache=True conflicts with "
132
+ f"indexer_kv_dtype={self.indexer_kv_dtype!r}. Set only "
133
+ "indexer_kv_dtype."
134
+ )
135
+ self.indexer_kv_dtype = "mxfp4"
136
+
137
+ def resolve_indexer_kv_dtype(self, default: IndexerKVDType) -> IndexerKVDType:
138
+ """Resolve `indexer_kv_dtype`, substituting `default` for "auto"."""
139
+ if self.indexer_kv_dtype == "auto":
140
+ return default
141
+ return self.indexer_kv_dtype
142
+
143
def compute_hash(self) -> str:
144
"""
145
Provide a hash that uniquely identifies all the configs
@@ -124,7 +150,8 @@ class AttentionConfig:
150
"""
151
from vllm.config.utils import get_hash_factors, hash_factors
152
127
- ignored_factors: set[str] = set()
153
+ # Folded into indexer_kv_dtype by __post_init__.
154
+ ignored_factors: set[str] = {"use_fp4_indexer_cache"}
155
factors = get_hash_factors(self, ignored_factors)
156
return hash_factors(factors)
157
vllm/models/deepseek_v4/attention.py
+2
-1
@@ -56,6 +56,7 @@ from vllm.utils.multi_stream_utils import (
56
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
57
from vllm.v1.attention.backends.mla.indexer import (
58
DeepseekV4IndexerBackend,
59
+ dsa_indexer_uses_fp4,
60
get_max_prefill_buffer_size,
61
)
62
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
@@ -805,7 +806,7 @@ class DeepseekV4Indexer(nn.Module):
806
self.q_lora_rank = q_lora_rank # 1536
807
self.compress_ratio = compress_ratio
808
self.eager_scratch_pool = eager_scratch_pool
808
- self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache
809
+ self.use_fp4_kv = dsa_indexer_uses_fp4(vllm_config)
810
logger.info_once(
811
"Using %s indexer cache for Lightning Indexer.",
812
"MXFP4" if self.use_fp4_kv else "FP8",
vllm/models/minimax_m3/nvidia/model.py
+3
-1
@@ -498,7 +498,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
498
set_default_quant_scales(self, register_buffer=True)
499
# Indexer side-cache dtype, mirroring --kv-cache-dtype for the main
500
# cache (--attention-config '{"indexer_kv_dtype": ...}').
501
- self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype
501
+ self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
502
+ "bf16"
503
+ )
504
505
# Shared top-k buffer: the indexer writes the selected blocks into it and
506
# the attend impl reads them back (so nothing crosses the eager break as a
vllm/v1/attention/backends/mla/indexer.py
+24
-13
@@ -40,6 +40,28 @@ from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
40
41
logger = init_logger(__name__)
42
43
+# The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout)
44
+# and mxfp4 is the opt-in Blackwell path.
45
+DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4")
46
+
47
+
48
+def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
49
+ """Whether the DeepSeek sparse indexer should use the MXFP4 K cache."""
50
+ kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("fp8")
51
+ if kv_dtype not in DSA_INDEXER_KV_DTYPES:
52
+ raise ValueError(
53
+ f"indexer_kv_dtype={kv_dtype!r} is not supported by the DeepSeek "
54
+ f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})."
55
+ )
56
+ use_fp4 = kv_dtype == "mxfp4"
57
+ if use_fp4 and not current_platform.is_device_capability_family(100):
58
+ raise ValueError(
59
+ "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs "
60
+ "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
61
+ "earlier architectures are not supported."
62
+ )
63
+ return use_fp4
64
+
65
66
@triton.jit
67
def _prepare_uniform_decode_kernel(
@@ -526,18 +548,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
548
if self.vllm_config.speculative_config
549
else 0
550
)
529
- self.use_fp4_indexer_cache = (
530
- self.vllm_config.attention_config.use_fp4_indexer_cache
531
- )
532
-
533
- assert (
534
- current_platform.is_device_capability_family(100)
535
- or not self.use_fp4_indexer_cache
536
- ), (
537
- "use_fp4_indexer_cache requires Blackwell datacenter GPUs "
538
- "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
539
- "earlier architectures are not supported."
540
- )
551
+ self.use_fp4_indexer_cache = dsa_indexer_uses_fp4(self.vllm_config)
552
553
next_n = self.num_speculative_tokens + 1
554
self.decode_threshold = next_n
@@ -546,7 +557,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
557
self.supports_varlen = _supports_varlen_paged_mqa_logits()
558
logger.info_once(
559
"DSA indexer decode path: use_flattening=%s supports_varlen=%s "
549
- "(next_n=%d, use_fp4_indexer_cache=%s)",
560
+ "(next_n=%d, use_fp4_cache=%s)",
561
self.use_flattening,
562
self.supports_varlen,
563
next_n,