[Bugfix][MLA] Read sparse model settings from text config (#56160)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Sep 11, 2026 at 19:02 UTC
dc07f1638f73814b95776832b85df1cc92850416
11 files changed
+32
-22
tests/v1/attention/test_dspark_noncausal_sparse_mla.py
+3
-1
@@ -168,6 +168,7 @@ def _run_sparse_backend_vs_sdpa(
168
)
169
model_config = vllm_config.model_config
170
model_config.hf_text_config = SimpleNamespace(
171
+ index_topk=topk_tokens,
172
q_lora_rank=None,
173
kv_lora_rank=kv_lora_rank,
174
qk_nope_head_dim=qk_nope_head_dim,
@@ -175,6 +176,7 @@ def _run_sparse_backend_vs_sdpa(
176
v_head_dim=v_head_dim,
177
model_type="deepseek_v2",
178
)
179
+ del model_config.hf_config.index_topk # Composite configs only nest this field.
180
model_config.dtype = dtype
181
model_config.get_num_attention_heads = MethodType(
182
lambda self, parallel_config: num_heads, model_config
@@ -287,7 +289,7 @@ def _run_sparse_backend_vs_sdpa(
289
causal_reference = torch.cat(causal_reference_outputs, dim=0)
290
291
vllm_config.cache_config.cache_dtype = kv_cache_dtype
290
- vllm_config.model_config.hf_config.index_topk = topk_tokens
292
+ vllm_config.model_config.hf_text_config.index_topk = topk_tokens
293
294
common_attn_metadata = create_common_attn_metadata(
295
batch_spec, block_size, device, arange_block_indices=True
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+5
-3
@@ -68,7 +68,9 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
68
config = SimpleNamespace(
69
scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
70
model_config=SimpleNamespace(
71
- hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0], index_kpool=32)
71
+ hf_text_config=SimpleNamespace(
72
+ compress_ratios=[0, 0, 4, 128, 0], index_kpool=32
73
+ )
74
),
75
parallel_config=SimpleNamespace(
76
decode_context_parallel_size=1,
@@ -89,7 +91,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
91
def test_compressed_slot_mapping_warmup_includes_index_kpool():
92
config = SimpleNamespace(
93
cache_config=SimpleNamespace(block_size=256),
92
- model_config=SimpleNamespace(hf_config=SimpleNamespace(index_kpool=32)),
94
+ model_config=SimpleNamespace(hf_text_config=SimpleNamespace(index_kpool=32)),
95
)
96
97
keys = CompressedSlotMappingKernel().get_warmup_keys(config)
@@ -101,7 +103,7 @@ def test_index_conversion_warmup_uses_physical_block_stride():
103
cache_config=SimpleNamespace(block_size=64),
104
model_config=SimpleNamespace(
105
max_model_len=1024,
104
- hf_config=SimpleNamespace(index_topk=2048),
106
+ hf_text_config=SimpleNamespace(index_topk=2048),
107
),
108
parallel_config=SimpleNamespace(
109
decode_context_parallel_size=1,
tests/v1/attention/test_sparse_mla_backends.py
+8
-2
@@ -423,6 +423,7 @@ def test_sparse_backend_decode_correctness(
423
)
424
model_config = vllm_config.model_config
425
model_config.hf_text_config = SimpleNamespace(
426
+ index_topk=topk_tokens,
427
q_lora_rank=None,
428
kv_lora_rank=kv_lora_rank,
429
qk_nope_head_dim=qk_nope_head_dim,
@@ -430,6 +431,7 @@ def test_sparse_backend_decode_correctness(
431
v_head_dim=v_head_dim,
432
model_type="deepseek_v2",
433
)
434
+ del model_config.hf_config.index_topk # Composite configs only nest this field.
435
model_config.dtype = dtype
436
model_config.get_num_attention_heads = MethodType(
437
lambda self, parallel_config: num_heads,
@@ -591,7 +593,7 @@ def test_sparse_backend_decode_correctness(
593
sdpa_reference = torch.cat(reference_outputs, dim=0)
594
595
vllm_config.cache_config.cache_dtype = kv_cache_dtype
594
- vllm_config.model_config.hf_config.index_topk = topk_tokens
596
+ vllm_config.model_config.hf_text_config.index_topk = topk_tokens
597
598
common_attn_metadata = create_common_attn_metadata(
599
batch_spec,
@@ -1149,6 +1151,7 @@ def test_sparse_backend_prefill_correctness(
1151
)
1152
model_config = vllm_config.model_config
1153
model_config.hf_text_config = SimpleNamespace(
1154
+ index_topk=topk_tokens,
1155
q_lora_rank=None,
1156
kv_lora_rank=kv_lora_rank,
1157
qk_nope_head_dim=qk_nope_head_dim,
@@ -1156,6 +1159,7 @@ def test_sparse_backend_prefill_correctness(
1159
v_head_dim=v_head_dim,
1160
model_type="deepseek_v2",
1161
)
1162
+ del model_config.hf_config.index_topk # Composite configs only nest this field.
1163
model_config.dtype = dtype
1164
model_config.model_arch_config.total_num_attention_heads = num_heads
1165
model_config.get_num_attention_heads = MethodType(
@@ -1255,7 +1259,7 @@ def test_sparse_backend_prefill_correctness(
1259
ref_output = torch.cat(reference_outputs, dim=0)
1260
1261
vllm_config.cache_config.cache_dtype = kv_cache_dtype
1258
- vllm_config.model_config.hf_config.index_topk = topk_tokens
1262
+ vllm_config.model_config.hf_text_config.index_topk = topk_tokens
1263
1264
common_attn_metadata = create_common_attn_metadata(
1265
batch_spec,
@@ -1747,6 +1751,7 @@ def _build_sparse_dcp_vllm_config(
1751
model_config = vllm_config.model_config
1752
model_config.dtype = torch.bfloat16
1753
model_config.hf_text_config = SimpleNamespace(
1754
+ index_topk=topk_tokens,
1755
q_lora_rank=None,
1756
kv_lora_rank=kv_lora_rank,
1757
qk_nope_head_dim=qk_nope_head_dim,
@@ -1754,6 +1759,7 @@ def _build_sparse_dcp_vllm_config(
1759
v_head_dim=v_head_dim,
1760
model_type="deepseek_v2",
1761
)
1762
+ del model_config.hf_config.index_topk # Composite configs only nest this field.
1763
model_config.get_num_attention_heads = MethodType(
1764
lambda self, parallel_config: local_heads, model_config
1765
)
vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py
+2
-2
@@ -173,7 +173,7 @@ class PackDCPTopkCandidatesKernel(
173
if dcp_world_size <= 1:
174
return []
175
cp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
176
- topk = vllm_config.model_config.hf_config.index_topk
176
+ topk = vllm_config.model_config.hf_text_config.index_topk
177
if topk <= 0:
178
return []
179
@@ -546,7 +546,7 @@ class StableTopKFromGatheredCandidatesKernel(
546
dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
547
if dcp_world_size <= 1:
548
return []
549
- topk = vllm_config.model_config.hf_config.index_topk
549
+ topk = vllm_config.model_config.hf_text_config.index_topk
550
if topk <= 0:
551
return []
552
return self._trace_dispatch(self.dispatch)(
vllm/model_executor/layers/attention/sparse_mla_attention.py
+2
-2
@@ -132,7 +132,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
132
self.device = device
133
self.model_config = vllm_config.model_config
134
self.mla_dims = get_mla_dims(self.model_config)
135
- self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk
135
+ self.topk_tokens: int = vllm_config.model_config.hf_text_config.index_topk
136
self.req_id_per_token_buffer = torch.empty(
137
(vllm_config.scheduler_config.max_num_batched_tokens,),
138
dtype=torch.int32,
@@ -202,7 +202,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
202
scheduler_config = vllm_config.scheduler_config
203
cache_config = vllm_config.cache_config
204
model_config = vllm_config.model_config
205
- topk_tokens = model_config.hf_config.index_topk
205
+ topk_tokens = model_config.hf_text_config.index_topk
206
207
workspace_size = min(
208
max(
vllm/v1/attention/backends/mla/compressor_utils.py
+3
-3
@@ -97,10 +97,10 @@ class CompressedSlotMappingKernel(
97
)
98
99
def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
100
- hf_config = vllm_config.model_config.hf_config
100
+ hf_text_config = vllm_config.model_config.hf_text_config
101
configured_ratios = (
102
- *(getattr(hf_config, "compress_ratios", None) or ()),
103
- getattr(hf_config, "index_kpool", 1) or 1,
102
+ *(getattr(hf_text_config, "compress_ratios", None) or ()),
103
+ getattr(hf_text_config, "index_kpool", 1) or 1,
104
)
105
compress_ratios = tuple(
106
dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+2
-2
@@ -99,8 +99,8 @@ class FlashAttnMLASparseBackend(AttentionBackend):
99
if vllm_config.parallel_config.decode_context_parallel_size > 1:
100
return "FlashAttention MLA Sparse does not support DCP for now"
101
102
- hf_config = vllm_config.model_config.hf_config
103
- if not hasattr(hf_config, "index_topk"):
102
+ hf_text_config = vllm_config.model_config.hf_text_config
103
+ if not hasattr(hf_text_config, "index_topk"):
104
return "FlashAttention MLA Sparse requires model with index_topk"
105
return None
106
vllm/v1/attention/backends/mla/indexer.py
+4
-4
@@ -396,7 +396,7 @@ class BuildPrefillChunkMetadataKernel(
396
397
def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
398
max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
399
- hf_config = vllm_config.model_config.hf_config
399
+ hf_text_config = vllm_config.model_config.hf_text_config
400
parallel_config = vllm_config.parallel_config
401
dcp_world = parallel_config.decode_context_parallel_size
402
dcp_interleave = parallel_config.cp_kv_cache_interleave_size
@@ -405,12 +405,12 @@ class BuildPrefillChunkMetadataKernel(
405
dict.fromkeys(
406
max(1, int(ratio))
407
for ratio in (
408
- *(getattr(hf_config, "compress_ratios", None) or (1,)),
409
- getattr(hf_config, "index_kpool", 1) or 1,
408
+ *(getattr(hf_text_config, "compress_ratios", None) or (1,)),
409
+ getattr(hf_text_config, "index_kpool", 1) or 1,
410
)
411
)
412
)
413
- index_kpool = getattr(hf_config, "index_kpool", None)
413
+ index_kpool = getattr(hf_text_config, "index_kpool", None)
414
if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios:
415
compress_ratios = compress_ratios + (index_kpool,)
416
return self._trace_dispatch(self.dispatch)(
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+1
-1
@@ -435,7 +435,7 @@ class ROCMAiterMLASparseMetadataBuilder(
435
436
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
437
self.mla_dims = get_mla_dims(self.model_config)
438
- self.topk_tokens = vllm_config.model_config.hf_config.index_topk
438
+ self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk
439
attention_context = vllm_config.compilation_config.static_forward_context
440
# Sink decode must use AITER's nonpersistent path. In particular,
441
# gfx942 has no persistent+LSE kernel, and its metadata heuristic
vllm/v1/attention/backends/mla/sparse_utils.py
+1
-1
@@ -232,7 +232,7 @@ class ConvertReqIndexToGlobalIndexKernel(
232
dcp_size = vllm_config.parallel_config.decode_context_parallel_size
233
dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
234
dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
235
- num_topk_tokens = vllm_config.model_config.hf_config.index_topk
235
+ num_topk_tokens = vllm_config.model_config.hf_text_config.index_topk
236
max_num_blocks = cdiv(
237
vllm_config.model_config.max_model_len,
238
block_size * dcp_size,
vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+1
-1
@@ -128,7 +128,7 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]
128
129
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
130
self.mla_dims = get_mla_dims(self.model_config)
131
- self.topk_tokens = vllm_config.model_config.hf_config.index_topk
131
+ self.topk_tokens = vllm_config.model_config.hf_text_config.index_topk
132
self.topk_tokens_tensor = torch.tensor(
133
[self.topk_tokens], device=device, dtype=torch.int32
134
)