[ROCm][DSv4] Enable DSpark adaptive verification (#52362)
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com> Signed-off-by: larryli2-amd <larryli2@amd.com> Co-authored-by: larryli2-amd <larryli2@amd.com>
Tuukka Sarvi committed
Sep 21, 2026 at 18:03 UTC
e34685dfc0c9f46ae58650fe7b08cc89b615a5fe
6 files changed
+725
-7
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+259
-1
@@ -188,6 +188,7 @@ def _ref_sparse_decode_ragged(
188
extra_rows: list[list[int]] | None = None,
189
main_use_fnuz: bool = False,
190
extra_use_fnuz: bool = False,
191
+ extra_block_size: int | None = None,
192
) -> torch.Tensor:
193
q_f32 = q.float()
194
out = torch.empty_like(q_f32)
@@ -209,7 +210,10 @@ def _ref_sparse_decode_ragged(
210
)
211
row_kv.append(
212
_read_fp8_ds_mla_cache_rows(
212
- extra_cache, extra_slots, block_size, extra_use_fnuz
213
+ extra_cache,
214
+ extra_slots,
215
+ extra_block_size or block_size,
216
+ extra_use_fnuz,
217
)
218
)
219
@@ -244,6 +248,12 @@ def _ragged_from_rows(
248
)
249
250
251
+def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]:
252
+ ends = indptr.cpu().tolist()
253
+ values = indices[: ends[-1]].cpu().tolist()
254
+ return [values[start:end] for start, end in zip(ends, ends[1:])]
255
+
256
+
257
def _launch_sparse_decode_reduce(
258
part_m: torch.Tensor,
259
part_l: torch.Tensor,
@@ -1192,6 +1202,254 @@ def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None:
1202
torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2)
1203
1204
1205
+@requires_gfx950
1206
+@torch.inference_mode()
1207
+def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
1208
+ """Replay the production ROCm decode path after device-only reallocation."""
1209
+ from tests.v1.attention.utils import create_vllm_config
1210
+ from vllm.models.deepseek_v4.amd.rocm import (
1211
+ DeepseekV4ROCMAiterMLASparseMetadataBuilder,
1212
+ DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
1213
+ )
1214
+ from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
1215
+ from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
1216
+ from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec
1217
+
1218
+ device = torch.device("cuda")
1219
+ num_reqs = 3
1220
+ upper_query_len = 8
1221
+ graph_tokens = num_reqs * upper_query_len
1222
+ block_size = 256
1223
+ compressed_block_size = block_size // 128
1224
+ window_size = 32
1225
+
1226
+ vllm_config = create_vllm_config(
1227
+ model_name="facebook/opt-125m",
1228
+ max_model_len=1024,
1229
+ block_size=block_size,
1230
+ max_num_seqs=num_reqs,
1231
+ max_num_batched_tokens=graph_tokens,
1232
+ hf_config_override={
1233
+ "compress_ratios": [128],
1234
+ "index_topk": 2048,
1235
+ "sliding_window": window_size,
1236
+ },
1237
+ )
1238
+ vllm_config.speculative_config = SimpleNamespace(
1239
+ num_speculative_tokens=upper_query_len - 1,
1240
+ parallel_drafting=False,
1241
+ enable_adaptive_verification=True,
1242
+ use_dspark=lambda: True,
1243
+ )
1244
+ mla_spec = MLAAttentionSpec(
1245
+ block_size=block_size,
1246
+ num_kv_heads=1,
1247
+ head_size=HEAD_DIM,
1248
+ dtype=torch.bfloat16,
1249
+ tokens_per_state=128,
1250
+ cache_dtype_str="fp8_ds_mla",
1251
+ model_version="deepseek_v4",
1252
+ )
1253
+ swa_spec = SlidingWindowMLASpec(
1254
+ block_size=block_size,
1255
+ num_kv_heads=1,
1256
+ head_size=HEAD_DIM,
1257
+ dtype=torch.bfloat16,
1258
+ sliding_window=window_size,
1259
+ cache_dtype_str="fp8_ds_mla",
1260
+ model_version="deepseek_v4",
1261
+ )
1262
+ mla_builder = DeepseekV4ROCMAiterMLASparseMetadataBuilder(
1263
+ mla_spec, ["c128a"], vllm_config, device
1264
+ )
1265
+ swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
1266
+ swa_spec, ["c128a"], vllm_config, device
1267
+ )
1268
+ assert (
1269
+ mla_builder.get_cudagraph_support(vllm_config, mla_spec)
1270
+ == AttentionCGSupport.ALWAYS
1271
+ )
1272
+ assert (
1273
+ swa_builder.get_cudagraph_support(vllm_config, swa_spec)
1274
+ == AttentionCGSupport.ALWAYS
1275
+ )
1276
+
1277
+ seq_lens = torch.tensor([520, 528, 536], dtype=torch.int32, device=device)
1278
+ seq_lens_cpu = seq_lens.cpu()
1279
+ cpu_query_start_loc = torch.arange(
1280
+ 0,
1281
+ graph_tokens + 1,
1282
+ upper_query_len,
1283
+ dtype=torch.int32,
1284
+ )
1285
+ max_blocks = (int(seq_lens_cpu.max()) + block_size - 1) // block_size
1286
+ block_table = torch.arange(
1287
+ num_reqs * max_blocks, dtype=torch.int32, device=device
1288
+ ).view(num_reqs, max_blocks)
1289
+
1290
+ def build_metadata(query_lens: list[int]):
1291
+ query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device)
1292
+ query_start_loc = torch.cat(
1293
+ [
1294
+ torch.zeros(1, dtype=torch.int32, device=device),
1295
+ query_lens_tensor.cumsum(0),
1296
+ ]
1297
+ )
1298
+ active_tokens = sum(query_lens)
1299
+ positions = torch.zeros(graph_tokens, dtype=torch.int64, device=device)
1300
+ position_rows = [
1301
+ torch.arange(
1302
+ seq_len - query_len,
1303
+ seq_len,
1304
+ dtype=torch.int64,
1305
+ device=device,
1306
+ )
1307
+ for seq_len, query_len in zip(seq_lens_cpu.tolist(), query_lens)
1308
+ ]
1309
+ positions[:active_tokens] = torch.cat(position_rows)
1310
+ slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device)
1311
+ slot_mapping[:active_tokens] = torch.arange(
1312
+ active_tokens, dtype=torch.int64, device=device
1313
+ )
1314
+ common = CommonAttentionMetadata(
1315
+ query_start_loc=query_start_loc,
1316
+ query_start_loc_cpu=cpu_query_start_loc,
1317
+ seq_lens=seq_lens,
1318
+ seq_lens_cpu_upper_bound=seq_lens_cpu,
1319
+ num_reqs=num_reqs,
1320
+ num_actual_tokens=graph_tokens,
1321
+ max_query_len=upper_query_len,
1322
+ max_seq_len=int(seq_lens_cpu.max()),
1323
+ block_table_tensor=block_table,
1324
+ slot_mapping=slot_mapping,
1325
+ positions=positions,
1326
+ causal=True,
1327
+ )
1328
+ return (
1329
+ mla_builder.build_for_cudagraph_capture(common),
1330
+ swa_builder.build_for_cudagraph_capture(common),
1331
+ )
1332
+
1333
+ mla_metadata, swa_metadata = build_metadata([8, 8, 8])
1334
+ assert mla_metadata.for_cudagraph_capture
1335
+ metadata_ptrs = (
1336
+ mla_metadata.c128a_decode_topk_ragged_indices.data_ptr(),
1337
+ mla_metadata.c128a_decode_topk_ragged_indptr.data_ptr(),
1338
+ swa_metadata.decode_swa_ragged_indices.data_ptr(),
1339
+ swa_metadata.decode_swa_ragged_indptr.data_ptr(),
1340
+ )
1341
+
1342
+ torch.manual_seed(19)
1343
+ num_heads = 16
1344
+ q = (
1345
+ torch.randn(
1346
+ graph_tokens,
1347
+ num_heads,
1348
+ HEAD_DIM,
1349
+ dtype=torch.bfloat16,
1350
+ device=device,
1351
+ )
1352
+ * 0.125
1353
+ )
1354
+ num_cache_blocks = num_reqs * max_blocks
1355
+ swa_cache = _pack_fp8_ds_mla_cache(
1356
+ torch.randn(
1357
+ num_cache_blocks * block_size,
1358
+ HEAD_DIM,
1359
+ dtype=torch.bfloat16,
1360
+ device=device,
1361
+ )
1362
+ * 0.125,
1363
+ block_size,
1364
+ use_fnuz=False,
1365
+ )
1366
+ compressed_cache = _pack_fp8_ds_mla_cache(
1367
+ torch.randn(
1368
+ num_cache_blocks * compressed_block_size,
1369
+ HEAD_DIM,
1370
+ dtype=torch.bfloat16,
1371
+ device=device,
1372
+ )
1373
+ * 0.125,
1374
+ compressed_block_size,
1375
+ use_fnuz=False,
1376
+ )
1377
+ attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
1378
+ out = torch.empty_like(q)
1379
+ monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
1380
+
1381
+ def run_decode(mla_md, swa_md) -> None:
1382
+ mod.rocm_sparse_attn_decode(
1383
+ q=q,
1384
+ kv_cache=compressed_cache,
1385
+ swa_k_cache=swa_cache,
1386
+ swa_only=False,
1387
+ topk_indices=mla_md.c128a_global_decode_topk_indices,
1388
+ topk_lens=mla_md.c128a_decode_topk_lens,
1389
+ swa_indices=swa_md.decode_swa_indices,
1390
+ swa_lens=swa_md.decode_swa_lens,
1391
+ swa_ragged_indices=swa_md.decode_swa_ragged_indices,
1392
+ swa_ragged_indptr=swa_md.decode_swa_ragged_indptr,
1393
+ topk_ragged_indices=mla_md.c128a_decode_topk_ragged_indices,
1394
+ topk_ragged_indptr=mla_md.c128a_decode_topk_ragged_indptr,
1395
+ attn_sink=attn_sink,
1396
+ scale=HEAD_DIM**-0.5,
1397
+ head_dim=HEAD_DIM,
1398
+ nope_head_dim=NOPE_HEAD_DIM,
1399
+ rope_head_dim=ROPE_HEAD_DIM,
1400
+ output=out,
1401
+ extra_cache_nan_free=True,
1402
+ adaptive_splits=True,
1403
+ )
1404
+
1405
+ def reference(mla_md, swa_md) -> torch.Tensor:
1406
+ return _ref_sparse_decode_ragged(
1407
+ q=q,
1408
+ main_cache=swa_cache,
1409
+ main_rows=_rows_from_ragged(
1410
+ swa_md.decode_swa_ragged_indices,
1411
+ swa_md.decode_swa_ragged_indptr,
1412
+ ),
1413
+ scale=HEAD_DIM**-0.5,
1414
+ attn_sink=attn_sink,
1415
+ block_size=block_size,
1416
+ extra_cache=compressed_cache,
1417
+ extra_rows=_rows_from_ragged(
1418
+ mla_md.c128a_decode_topk_ragged_indices,
1419
+ mla_md.c128a_decode_topk_ragged_indptr,
1420
+ ),
1421
+ extra_block_size=compressed_block_size,
1422
+ )
1423
+
1424
+ run_decode(mla_metadata, swa_metadata)
1425
+ torch.accelerator.synchronize()
1426
+ expected_full = reference(mla_metadata, swa_metadata)
1427
+ torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2)
1428
+
1429
+ graph = torch.cuda.CUDAGraph()
1430
+ with torch.cuda.graph(graph):
1431
+ run_decode(mla_metadata, swa_metadata)
1432
+ torch.accelerator.synchronize()
1433
+ captured_full = out.clone()
1434
+
1435
+ reallocated_mla, reallocated_swa = build_metadata([3, 8, 5])
1436
+ assert metadata_ptrs == (
1437
+ reallocated_mla.c128a_decode_topk_ragged_indices.data_ptr(),
1438
+ reallocated_mla.c128a_decode_topk_ragged_indptr.data_ptr(),
1439
+ reallocated_swa.decode_swa_ragged_indices.data_ptr(),
1440
+ reallocated_swa.decode_swa_ragged_indptr.data_ptr(),
1441
+ )
1442
+ expected_reallocated = reference(reallocated_mla, reallocated_swa)
1443
+ run_decode(reallocated_mla, reallocated_swa)
1444
+ torch.accelerator.synchronize()
1445
+ torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
1446
+ assert not torch.equal(captured_full, expected_reallocated)
1447
+
1448
+ graph.replay()
1449
+ torch.accelerator.synchronize()
1450
+ torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
1451
+
1452
+
1453
# ---------------------------------------------------------------------------
1454
# o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
1455
# ---------------------------------------------------------------------------
tests/models/test_deepseek_v4_dspark_rocm.py
new
+130
@@ -0,0 +1,130 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.models.deepseek_v4.amd import dspark as dspark_module
10
+from vllm.models.deepseek_v4.amd.dspark import DSparkDeepseekV4ForCausalLM
11
+
12
+
13
+def _make_uninitialized_model(confidence_head):
14
+ model = DSparkDeepseekV4ForCausalLM.__new__(DSparkDeepseekV4ForCausalLM)
15
+ object.__setattr__(
16
+ model,
17
+ "model",
18
+ SimpleNamespace(confidence_head=confidence_head),
19
+ )
20
+ return model
21
+
22
+
23
+def _prepare_loader_model(model, named_parameters):
24
+ object.__setattr__(
25
+ model,
26
+ "config",
27
+ SimpleNamespace(
28
+ n_routed_experts=1,
29
+ expert_dtype="fp4",
30
+ num_attention_heads=1,
31
+ ),
32
+ )
33
+ object.__setattr__(model, "named_parameters", lambda: named_parameters)
34
+
35
+
36
+def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch):
37
+ monkeypatch.setattr(
38
+ dspark_module,
39
+ "fused_moe_make_expert_params_mapping",
40
+ lambda *args, **kwargs: [],
41
+ )
42
+ monkeypatch.setattr(
43
+ dspark_module, "get_tensor_model_parallel_world_size", lambda: 1
44
+ )
45
+ monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0)
46
+
47
+
48
+@pytest.mark.cpu_test
49
+def test_deepseek_v4_rocm_dspark_maps_enabled_confidence_head():
50
+ model = _make_uninitialized_model(object())
51
+
52
+ assert (
53
+ model._remap_dspark_name("mtp.2.confidence_head.proj.weight")
54
+ == "model.confidence_head.proj.weight"
55
+ )
56
+
57
+
58
+@pytest.mark.cpu_test
59
+def test_deepseek_v4_rocm_dspark_skips_disabled_confidence_head():
60
+ model = _make_uninitialized_model(None)
61
+
62
+ assert model._remap_dspark_name("mtp.2.confidence_head.proj.weight") is None
63
+
64
+
65
+@pytest.mark.cpu_test
66
+def test_deepseek_v4_rocm_dspark_disables_unloaded_confidence_head(
67
+ monkeypatch: pytest.MonkeyPatch,
68
+):
69
+ model = _make_uninitialized_model(object())
70
+ _prepare_loader_model(model, [])
71
+ _disable_distributed_loader_paths(monkeypatch)
72
+
73
+ assert model.load_weights([]) == set()
74
+ assert model.model.confidence_head is None
75
+
76
+
77
+@pytest.mark.cpu_test
78
+def test_deepseek_v4_rocm_dspark_loads_complete_confidence_head(
79
+ monkeypatch: pytest.MonkeyPatch,
80
+):
81
+ class FakeParameter:
82
+ loaded_weight = None
83
+
84
+ def weight_loader(self, param, loaded_weight):
85
+ assert param is self
86
+ self.loaded_weight = loaded_weight
87
+
88
+ confidence_head = object()
89
+ parameter = FakeParameter()
90
+ model = _make_uninitialized_model(confidence_head)
91
+ _prepare_loader_model(
92
+ model,
93
+ [("model.confidence_head.proj.weight", parameter)],
94
+ )
95
+ _disable_distributed_loader_paths(monkeypatch)
96
+ loaded_weight = torch.tensor([[1.0, 2.0]])
97
+
98
+ loaded = model.load_weights([("mtp.2.confidence_head.proj.weight", loaded_weight)])
99
+
100
+ assert loaded == {"model.confidence_head.proj.weight"}
101
+ assert parameter.loaded_weight is loaded_weight
102
+ assert model.model.confidence_head is confidence_head
103
+
104
+
105
+@pytest.mark.cpu_test
106
+def test_deepseek_v4_rocm_dspark_confidence_is_probability():
107
+ class ConfidenceHead:
108
+ def __call__(self, head_hidden, markov_embed):
109
+ return (head_hidden[:, 0] + markov_embed[:, 0]).float()
110
+
111
+ model = _make_uninitialized_model(ConfidenceHead())
112
+ head_hidden = torch.tensor([[0.0], [1.0]], dtype=torch.bfloat16)
113
+ markov_embed = torch.tensor([[0.0], [-2.0]], dtype=torch.bfloat16)
114
+
115
+ confidence = model.compute_confidence(head_hidden, markov_embed)
116
+
117
+ torch.testing.assert_close(
118
+ confidence,
119
+ torch.sigmoid(torch.tensor([0.0, -1.0])),
120
+ )
121
+ assert torch.all((confidence >= 0) & (confidence <= 1))
122
+
123
+
124
+@pytest.mark.cpu_test
125
+def test_deepseek_v4_rocm_dspark_confidence_requires_a_head():
126
+ model = _make_uninitialized_model(None)
127
+ empty = torch.zeros((1, 1))
128
+
129
+ with pytest.raises(RuntimeError, match="confidence_head"):
130
+ model.compute_confidence(empty, empty)
tests/v1/attention/test_deepseek_v4_rocm_adaptive.py
new
+249
@@ -0,0 +1,249 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.models.deepseek_v4.amd.rocm import (
10
+ DeepseekV4ROCMAiterMLASparseMetadataBuilder,
11
+ DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
12
+)
13
+from vllm.v1.attention.backend import AttentionCGSupport
14
+from vllm.v1.attention.backends.mla import indexer
15
+from vllm.v1.attention.backends.mla.indexer import (
16
+ DeepseekV4IndexerBackend,
17
+ DeepseekV32IndexerMetadataBuilder,
18
+ DeepseekV41IndexerBackend,
19
+)
20
+
21
+
22
+def _make_indexer_config(*, architecture: str, adaptive: bool):
23
+ return SimpleNamespace(
24
+ model_config=SimpleNamespace(architectures=[architecture]),
25
+ num_speculative_tokens=1,
26
+ speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive),
27
+ )
28
+
29
+
30
+def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None:
31
+ # Platform predicates are mutually exclusive in production. Override both
32
+ # so these ROCm policy tests do not inherit CUDA capabilities from the CI
33
+ # host running them.
34
+ monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: False)
35
+ monkeypatch.setattr(indexer.current_platform, "is_rocm", lambda: True)
36
+
37
+
38
+def _make_indexer_builder(*, adaptive: bool, capacity: int = 12):
39
+ builder = DeepseekV32IndexerMetadataBuilder.__new__(
40
+ DeepseekV32IndexerMetadataBuilder
41
+ )
42
+ builder.vllm_config = SimpleNamespace(
43
+ speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive)
44
+ )
45
+ builder.supports_varlen = False
46
+ builder.decode_seq_lens_buffer = torch.zeros(capacity, dtype=torch.int32)
47
+ builder.expanded_block_table_buffer = torch.zeros((capacity, 2), dtype=torch.int32)
48
+ builder.decode_lens_buffer = torch.zeros(capacity, dtype=torch.int32)
49
+ builder.arange_buffer = torch.arange(capacity, dtype=torch.int32)
50
+ return builder
51
+
52
+
53
+@pytest.mark.cpu_test
54
+def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs():
55
+ adaptive_config = SimpleNamespace(
56
+ speculative_config=SimpleNamespace(enable_adaptive_verification=True)
57
+ )
58
+ fixed_config = SimpleNamespace(
59
+ speculative_config=SimpleNamespace(enable_adaptive_verification=False)
60
+ )
61
+
62
+ for builder_cls in (
63
+ DeepseekV4ROCMAiterMLASparseMetadataBuilder,
64
+ DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
65
+ ):
66
+ assert (
67
+ builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace())
68
+ == AttentionCGSupport.ALWAYS
69
+ )
70
+ assert (
71
+ builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace())
72
+ == AttentionCGSupport.UNIFORM_BATCH
73
+ )
74
+
75
+
76
+@pytest.mark.cpu_test
77
+def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch):
78
+ _mock_rocm_platform(monkeypatch)
79
+ adaptive_config = _make_indexer_config(
80
+ architecture="DeepseekV4ForCausalLM", adaptive=True
81
+ )
82
+ fixed_config = _make_indexer_config(
83
+ architecture="DeepseekV4ForCausalLM", adaptive=False
84
+ )
85
+
86
+ assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch()
87
+ assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
88
+ assert (
89
+ DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder
90
+ )
91
+ assert indexer._use_flattening(adaptive_config)
92
+ assert (
93
+ DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
94
+ adaptive_config, SimpleNamespace()
95
+ )
96
+ == AttentionCGSupport.ALWAYS
97
+ )
98
+ assert not indexer._use_flattening(fixed_config)
99
+ assert (
100
+ DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
101
+ fixed_config, SimpleNamespace()
102
+ )
103
+ == AttentionCGSupport.UNIFORM_BATCH
104
+ )
105
+
106
+
107
+@pytest.mark.cpu_test
108
+@pytest.mark.parametrize(
109
+ "architecture",
110
+ ["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"],
111
+)
112
+def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4(
113
+ monkeypatch: pytest.MonkeyPatch,
114
+ architecture: str,
115
+):
116
+ _mock_rocm_platform(monkeypatch)
117
+ adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True)
118
+
119
+ assert not indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
120
+ assert not indexer._use_flattening(adaptive_config)
121
+ assert (
122
+ DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
123
+ adaptive_config, SimpleNamespace()
124
+ )
125
+ == AttentionCGSupport.UNIFORM_BATCH
126
+ )
127
+
128
+ if architecture == "DeepseekV41ForCausalLM":
129
+ assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch()
130
+
131
+
132
+@pytest.mark.cpu_test
133
+def test_rocm_adaptive_indexer_preserves_single_request_uniform_path(
134
+ monkeypatch: pytest.MonkeyPatch,
135
+):
136
+ builder = _make_indexer_builder(adaptive=True, capacity=8)
137
+
138
+ class FakeUniformKernel:
139
+ called = False
140
+
141
+ def __call__(
142
+ self,
143
+ seq_lens,
144
+ decode_seq_lens,
145
+ block_table,
146
+ expanded_block_table,
147
+ decode_lens,
148
+ num_decode_tokens,
149
+ max_decode_len,
150
+ ):
151
+ self.called = True
152
+ assert num_decode_tokens == 4
153
+ assert block_table.stride(0) == expanded_block_table.stride(0) == 2
154
+ decode_seq_lens[:max_decode_len] = torch.arange(
155
+ seq_lens[0] - max_decode_len + 1,
156
+ seq_lens[0] + 1,
157
+ dtype=torch.int32,
158
+ )
159
+ expanded_block_table[:max_decode_len] = block_table[0]
160
+ decode_lens[:max_decode_len] = 1
161
+
162
+ fake_kernel = FakeUniformKernel()
163
+ monkeypatch.setattr(indexer, "_PREPARE_UNIFORM_DECODE_KERNEL", fake_kernel)
164
+
165
+ seq_lens, block_table, decode_lens, batch_size, requires_padding = (
166
+ builder._prepare_decode_tensors(
167
+ seq_lens=torch.tensor([10], dtype=torch.int32),
168
+ block_table=torch.tensor([[1, 2]], dtype=torch.int32),
169
+ decode_lens=torch.tensor([4], dtype=torch.int32),
170
+ decode_lens_cpu=torch.tensor([4], dtype=torch.int32),
171
+ query_start_loc=torch.tensor([0], dtype=torch.int32),
172
+ num_decodes=1,
173
+ num_decode_tokens=4,
174
+ use_native=False,
175
+ next_n=8,
176
+ max_decode_len=4,
177
+ )
178
+ )
179
+
180
+ assert fake_kernel.called
181
+ torch.testing.assert_close(seq_lens, torch.tensor([7, 8, 9, 10], dtype=torch.int32))
182
+ torch.testing.assert_close(
183
+ block_table,
184
+ torch.tensor([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=torch.int32),
185
+ )
186
+ torch.testing.assert_close(decode_lens, torch.ones(4, dtype=torch.int32))
187
+ assert batch_size == 4
188
+ assert not requires_padding
189
+
190
+
191
+@pytest.mark.cpu_test
192
+def test_rocm_adaptive_indexer_replays_changed_allocations_with_stable_buffers():
193
+ builder = _make_indexer_builder(adaptive=True)
194
+ buffer_ptrs = (
195
+ builder.decode_seq_lens_buffer.data_ptr(),
196
+ builder.expanded_block_table_buffer.data_ptr(),
197
+ builder.decode_lens_buffer.data_ptr(),
198
+ )
199
+ block_table = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
200
+ decode_lens_cpu = torch.tensor([4, 4, 0], dtype=torch.int32)
201
+
202
+ first = builder._prepare_decode_tensors(
203
+ seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32),
204
+ block_table=block_table,
205
+ decode_lens=torch.tensor([7, 1, 0], dtype=torch.int32),
206
+ decode_lens_cpu=decode_lens_cpu,
207
+ query_start_loc=torch.tensor([0, 7, 8], dtype=torch.int32),
208
+ num_decodes=3,
209
+ num_decode_tokens=10,
210
+ use_native=False,
211
+ next_n=8,
212
+ max_decode_len=4,
213
+ )
214
+ torch.testing.assert_close(
215
+ first[0],
216
+ torch.tensor([14, 15, 16, 17, 18, 19, 20, 20, 0, 0], dtype=torch.int32),
217
+ )
218
+ torch.testing.assert_close(
219
+ first[1][:, 0],
220
+ torch.tensor([1, 1, 1, 1, 1, 1, 1, 3, 0, 0], dtype=torch.int32),
221
+ )
222
+
223
+ second = builder._prepare_decode_tensors(
224
+ seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32),
225
+ block_table=block_table,
226
+ decode_lens=torch.tensor([1, 7, 0], dtype=torch.int32),
227
+ decode_lens_cpu=decode_lens_cpu,
228
+ query_start_loc=torch.tensor([0, 1, 8], dtype=torch.int32),
229
+ num_decodes=3,
230
+ num_decode_tokens=10,
231
+ use_native=False,
232
+ next_n=8,
233
+ max_decode_len=4,
234
+ )
235
+ torch.testing.assert_close(
236
+ second[0],
237
+ torch.tensor([20, 14, 15, 16, 17, 18, 19, 20, 0, 0], dtype=torch.int32),
238
+ )
239
+ torch.testing.assert_close(
240
+ second[1][:, 0],
241
+ torch.tensor([1, 3, 3, 3, 3, 3, 3, 3, 0, 0], dtype=torch.int32),
242
+ )
243
+ torch.testing.assert_close(second[2], torch.ones(10, dtype=torch.int32))
244
+ assert second[3:] == (10, False)
245
+ assert buffer_ptrs == (
246
+ builder.decode_seq_lens_buffer.data_ptr(),
247
+ builder.expanded_block_table_buffer.data_ptr(),
248
+ builder.decode_lens_buffer.data_ptr(),
249
+ )
vllm/models/deepseek_v4/amd/dspark.py
+28
-5
@@ -44,6 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
44
)
45
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
46
from vllm.model_executor.models.qwen3_dspark import (
47
+ DSparkConfidenceHead,
48
DSparkMarkovHead,
49
)
50
from vllm.model_executor.models.utils import maybe_prefix
@@ -102,7 +103,7 @@ class DSparkDeepseekV4Model(nn.Module):
103
]
104
)
105
105
- # Heads: final norm + hc_head, and the Markov head
106
+ # Heads: final norm + hc_head, and the Markov + confidence heads
107
# Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
108
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
109
hc_dim = self.hc_mult * config.hidden_size
@@ -125,6 +126,12 @@ class DSparkDeepseekV4Model(nn.Module):
126
config.dspark_markov_rank,
127
prefix=maybe_prefix(prefix, "markov_head"),
128
)
129
+ self.confidence_head: DSparkConfidenceHead | None = None
130
+ if getattr(config, "enable_confidence_head", True):
131
+ self.confidence_head = DSparkConfidenceHead(
132
+ config.hidden_size + config.dspark_markov_rank,
133
+ prefix=maybe_prefix(prefix, "confidence_head"),
134
+ )
135
136
# MHC head CustomOp dispatcher (aiter / tilelang / triton / torch),
137
# replacing the direct nvidia tilelang kernel call.
@@ -358,6 +365,17 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
365
def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
366
return self.model.markov_head.bias(markov_embed, self.logits_processor)
367
368
+ def compute_confidence(
369
+ self, head_hidden: torch.Tensor, markov_embed: torch.Tensor
370
+ ) -> torch.Tensor:
371
+ """Per-position acceptance probability for each drafted token."""
372
+ if self.model.confidence_head is None:
373
+ raise RuntimeError(
374
+ "compute_confidence() requires a confidence head, but the "
375
+ "checkpoint did not provide confidence_head weights."
376
+ )
377
+ return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))
378
+
379
# --- Weight loading ----------------------------------------------------
380
381
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -391,6 +409,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
409
410
params_dict = dict(self.named_parameters())
411
loaded_params: set[str] = set()
412
+ loaded_confidence_head = False
413
414
tp_size = get_tensor_model_parallel_world_size()
415
tp_rank = get_tensor_model_parallel_rank()
@@ -403,6 +422,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
422
if mapped is None:
423
continue
424
name = mapped
425
+ if "confidence_head." in name:
426
+ loaded_confidence_head = True
427
428
# ``.scale`` -> per-method scale suffix.
429
if name.endswith(".scale"):
@@ -439,8 +460,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
460
continue
461
462
# Stacked rules only apply to decoder-layer weights. Head-stack params
442
- # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
443
- # "markov_w1" would collide with the "w1" shard rule.
463
+ # (main_proj/norm/hc_head/markov_head/confidence_head) load directly —
464
+ # otherwise e.g. "markov_w1" would collide with the "w1" shard rule.
465
is_layer_param = name.startswith("model.layers.")
466
for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
467
if not is_layer_param or weight_name not in name:
@@ -469,6 +490,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
490
weight_loader(param, loaded_weight)
491
loaded_params.add(name)
492
493
+ if self.model.confidence_head is not None and not loaded_confidence_head:
494
+ self.model.confidence_head = None
495
logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
496
return loaded_params
497
@@ -482,8 +505,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
505
return None
506
stage = int(m.group(1))
507
rest = m.group(2)
485
- # The confidence head is not wired into inference yet; drop its weights.
486
- if rest.startswith("confidence_head."):
508
+ if rest.startswith("confidence_head.") and self.model.confidence_head is None:
509
return None
510
# Head-stack params live at model level (mtp.last), context combiner at
511
# model level (mtp.0); everything else is a per-layer decoder block.
@@ -493,6 +515,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
515
"hc_head_base",
516
"hc_head_scale",
517
"markov_head.",
518
+ "confidence_head.",
519
)
520
if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
521
head_prefixes
vllm/models/deepseek_v4/amd/rocm.py
+31
@@ -8,6 +8,7 @@ from typing import cast
8
import torch
9
10
from vllm import envs
11
+from vllm.config import VllmConfig
12
from vllm.distributed import (
13
get_tensor_model_parallel_world_size,
14
tensor_model_parallel_all_reduce,
@@ -26,6 +27,7 @@ from vllm.platforms.rocm import _ON_GFX950
27
from vllm.triton_utils import tl, triton
28
from vllm.utils.multi_stream_utils import execute_in_parallel
29
from vllm.v1.attention.backend import (
30
+ AttentionCGSupport,
31
CommonAttentionMetadata,
32
)
33
from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -38,6 +40,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
40
rocm_sparse_attn_decode,
41
rocm_sparse_attn_prefill,
42
)
43
+from vllm.v1.kv_cache_interface import KVCacheSpec
44
from vllm.v1.worker.workspace import current_workspace_manager
45
46
logger = init_logger(__name__)
@@ -445,6 +448,20 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
448
449
450
class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
451
+ @classmethod
452
+ def get_cudagraph_support(
453
+ cls,
454
+ vllm_config: VllmConfig,
455
+ kv_cache_spec: KVCacheSpec,
456
+ ) -> AttentionCGSupport:
457
+ spec_config = vllm_config.speculative_config
458
+ if spec_config is not None and spec_config.enable_adaptive_verification:
459
+ # All per-token metadata is built from device query boundaries into
460
+ # persistent buffers, so adaptive verification can replay varlen
461
+ # FULL decode graphs after reallocating drafts across requests.
462
+ return AttentionCGSupport.ALWAYS
463
+ return super().get_cudagraph_support(vllm_config, kv_cache_spec)
464
+
465
def __init__(self, *args, **kwargs):
466
super().__init__(*args, **kwargs)
467
self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None
@@ -512,6 +529,20 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBui
529
530
531
class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
532
+ @classmethod
533
+ def get_cudagraph_support(
534
+ cls,
535
+ vllm_config: VllmConfig,
536
+ kv_cache_spec: KVCacheSpec,
537
+ ) -> AttentionCGSupport:
538
+ spec_config = vllm_config.speculative_config
539
+ if spec_config is not None and spec_config.enable_adaptive_verification:
540
+ # SWA indices, lengths, and token-to-request mappings are built from
541
+ # device boundaries into persistent buffers, so adaptive verification
542
+ # can replay varlen FULL decode graphs safely.
543
+ return AttentionCGSupport.ALWAYS
544
+ return super().get_cudagraph_support(vllm_config, kv_cache_spec)
545
+
546
# Keep fused multi-step decode disabled until update_draft_decode_metadata()
547
# also refreshes the ROCm-specific ragged SWA indices and indptrs.
548
supports_draft_decode_metadata_update = False
vllm/v1/attention/backends/mla/indexer.py
+28
-1
@@ -237,6 +237,15 @@ class KpoolTailBackend(DeepseekV32IndexerBackend):
237
238
239
class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
240
+ @classmethod
241
+ def supports_device_cpu_query_lens_mismatch(cls) -> bool:
242
+ # ROCm runs adaptive verification through the per-token flattened
243
+ # indexer path, which derives row ownership from device decode lengths.
244
+ return (
245
+ current_platform.is_rocm()
246
+ or super().supports_device_cpu_query_lens_mismatch()
247
+ )
248
+
249
@staticmethod
250
def get_name() -> str:
251
return "DEEPSEEK_V4_INDEXER"
@@ -254,6 +263,12 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
263
264
265
class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
266
+ @classmethod
267
+ def supports_device_cpu_query_lens_mismatch(cls) -> bool:
268
+ # The ROCm flattened-query support above is validated for the
269
+ # DeepSeek-V4 adaptive DSpark path only.
270
+ return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch()
271
+
272
@staticmethod
273
def get_name() -> str:
274
return "DEEPSEEK_V41_INDEXER"
@@ -802,6 +817,15 @@ def _supports_flattened_device_query_lens() -> bool:
817
)
818
819
820
+def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool:
821
+ model_config = vllm_config.model_config
822
+ return (
823
+ current_platform.is_rocm()
824
+ and model_config is not None
825
+ and "DeepseekV4ForCausalLM" in model_config.architectures
826
+ )
827
+
828
+
829
def _supports_native_decode(next_n: int) -> bool:
830
"""Whether decode can pass `next_n` Q rows per request to the kernel
831
instead of flattening to one single-token row per query, which re-reads
@@ -822,7 +846,10 @@ def _use_flattening(vllm_config: VllmConfig) -> bool:
846
return not _supports_native_decode(next_n) or (
847
speculative_config is not None
848
and speculative_config.enable_adaptive_verification
825
- and _supports_flattened_device_query_lens()
849
+ and (
850
+ _supports_flattened_device_query_lens()
851
+ or _rocm_supports_flattened_device_query_lens(vllm_config)
852
+ )
853
)
854
855