diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 256b50d49a..5b68732f7a 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -188,6 +188,7 @@ def _ref_sparse_decode_ragged( extra_rows: list[list[int]] | None = None, main_use_fnuz: bool = False, extra_use_fnuz: bool = False, + extra_block_size: int | None = None, ) -> torch.Tensor: q_f32 = q.float() out = torch.empty_like(q_f32) @@ -209,7 +210,10 @@ def _ref_sparse_decode_ragged( ) row_kv.append( _read_fp8_ds_mla_cache_rows( - extra_cache, extra_slots, block_size, extra_use_fnuz + extra_cache, + extra_slots, + extra_block_size or block_size, + extra_use_fnuz, ) ) @@ -244,6 +248,12 @@ def _ragged_from_rows( ) +def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]: + ends = indptr.cpu().tolist() + values = indices[: ends[-1]].cpu().tolist() + return [values[start:end] for start, end in zip(ends, ends[1:])] + + def _launch_sparse_decode_reduce( part_m: torch.Tensor, part_l: torch.Tensor, @@ -1192,6 +1202,254 @@ def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None: torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2) +@requires_gfx950 +@torch.inference_mode() +def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None: + """Replay the production ROCm decode path after device-only reallocation.""" + from tests.v1.attention.utils import create_vllm_config + from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLASparseMetadataBuilder, + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + ) + from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata + from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod + from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec + + device = torch.device("cuda") + num_reqs = 3 + upper_query_len = 8 + graph_tokens = num_reqs * upper_query_len + block_size = 256 + compressed_block_size = block_size // 128 + window_size = 32 + + vllm_config = create_vllm_config( + model_name="facebook/opt-125m", + max_model_len=1024, + block_size=block_size, + max_num_seqs=num_reqs, + max_num_batched_tokens=graph_tokens, + hf_config_override={ + "compress_ratios": [128], + "index_topk": 2048, + "sliding_window": window_size, + }, + ) + vllm_config.speculative_config = SimpleNamespace( + num_speculative_tokens=upper_query_len - 1, + parallel_drafting=False, + enable_adaptive_verification=True, + use_dspark=lambda: True, + ) + mla_spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=HEAD_DIM, + dtype=torch.bfloat16, + tokens_per_state=128, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + ) + swa_spec = SlidingWindowMLASpec( + block_size=block_size, + num_kv_heads=1, + head_size=HEAD_DIM, + dtype=torch.bfloat16, + sliding_window=window_size, + cache_dtype_str="fp8_ds_mla", + model_version="deepseek_v4", + ) + mla_builder = DeepseekV4ROCMAiterMLASparseMetadataBuilder( + mla_spec, ["c128a"], vllm_config, device + ) + swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder( + swa_spec, ["c128a"], vllm_config, device + ) + assert ( + mla_builder.get_cudagraph_support(vllm_config, mla_spec) + == AttentionCGSupport.ALWAYS + ) + assert ( + swa_builder.get_cudagraph_support(vllm_config, swa_spec) + == AttentionCGSupport.ALWAYS + ) + + seq_lens = torch.tensor([520, 528, 536], dtype=torch.int32, device=device) + seq_lens_cpu = seq_lens.cpu() + cpu_query_start_loc = torch.arange( + 0, + graph_tokens + 1, + upper_query_len, + dtype=torch.int32, + ) + max_blocks = (int(seq_lens_cpu.max()) + block_size - 1) // block_size + block_table = torch.arange( + num_reqs * max_blocks, dtype=torch.int32, device=device + ).view(num_reqs, max_blocks) + + def build_metadata(query_lens: list[int]): + query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device) + query_start_loc = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + query_lens_tensor.cumsum(0), + ] + ) + active_tokens = sum(query_lens) + positions = torch.zeros(graph_tokens, dtype=torch.int64, device=device) + position_rows = [ + torch.arange( + seq_len - query_len, + seq_len, + dtype=torch.int64, + device=device, + ) + for seq_len, query_len in zip(seq_lens_cpu.tolist(), query_lens) + ] + positions[:active_tokens] = torch.cat(position_rows) + slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device) + slot_mapping[:active_tokens] = torch.arange( + active_tokens, dtype=torch.int64, device=device + ) + common = CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=cpu_query_start_loc, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=seq_lens_cpu, + num_reqs=num_reqs, + num_actual_tokens=graph_tokens, + max_query_len=upper_query_len, + max_seq_len=int(seq_lens_cpu.max()), + block_table_tensor=block_table, + slot_mapping=slot_mapping, + positions=positions, + causal=True, + ) + return ( + mla_builder.build_for_cudagraph_capture(common), + swa_builder.build_for_cudagraph_capture(common), + ) + + mla_metadata, swa_metadata = build_metadata([8, 8, 8]) + assert mla_metadata.for_cudagraph_capture + metadata_ptrs = ( + mla_metadata.c128a_decode_topk_ragged_indices.data_ptr(), + mla_metadata.c128a_decode_topk_ragged_indptr.data_ptr(), + swa_metadata.decode_swa_ragged_indices.data_ptr(), + swa_metadata.decode_swa_ragged_indptr.data_ptr(), + ) + + torch.manual_seed(19) + num_heads = 16 + q = ( + torch.randn( + graph_tokens, + num_heads, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + num_cache_blocks = num_reqs * max_blocks + swa_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_cache_blocks * block_size, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + block_size, + use_fnuz=False, + ) + compressed_cache = _pack_fp8_ds_mla_cache( + torch.randn( + num_cache_blocks * compressed_block_size, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + * 0.125, + compressed_block_size, + use_fnuz=False, + ) + attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device) + out = torch.empty_like(q) + monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1) + + def run_decode(mla_md, swa_md) -> None: + mod.rocm_sparse_attn_decode( + q=q, + kv_cache=compressed_cache, + swa_k_cache=swa_cache, + swa_only=False, + topk_indices=mla_md.c128a_global_decode_topk_indices, + topk_lens=mla_md.c128a_decode_topk_lens, + swa_indices=swa_md.decode_swa_indices, + swa_lens=swa_md.decode_swa_lens, + swa_ragged_indices=swa_md.decode_swa_ragged_indices, + swa_ragged_indptr=swa_md.decode_swa_ragged_indptr, + topk_ragged_indices=mla_md.c128a_decode_topk_ragged_indices, + topk_ragged_indptr=mla_md.c128a_decode_topk_ragged_indptr, + attn_sink=attn_sink, + scale=HEAD_DIM**-0.5, + head_dim=HEAD_DIM, + nope_head_dim=NOPE_HEAD_DIM, + rope_head_dim=ROPE_HEAD_DIM, + output=out, + extra_cache_nan_free=True, + adaptive_splits=True, + ) + + def reference(mla_md, swa_md) -> torch.Tensor: + return _ref_sparse_decode_ragged( + q=q, + main_cache=swa_cache, + main_rows=_rows_from_ragged( + swa_md.decode_swa_ragged_indices, + swa_md.decode_swa_ragged_indptr, + ), + scale=HEAD_DIM**-0.5, + attn_sink=attn_sink, + block_size=block_size, + extra_cache=compressed_cache, + extra_rows=_rows_from_ragged( + mla_md.c128a_decode_topk_ragged_indices, + mla_md.c128a_decode_topk_ragged_indptr, + ), + extra_block_size=compressed_block_size, + ) + + run_decode(mla_metadata, swa_metadata) + torch.accelerator.synchronize() + expected_full = reference(mla_metadata, swa_metadata) + torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_decode(mla_metadata, swa_metadata) + torch.accelerator.synchronize() + captured_full = out.clone() + + reallocated_mla, reallocated_swa = build_metadata([3, 8, 5]) + assert metadata_ptrs == ( + reallocated_mla.c128a_decode_topk_ragged_indices.data_ptr(), + reallocated_mla.c128a_decode_topk_ragged_indptr.data_ptr(), + reallocated_swa.decode_swa_ragged_indices.data_ptr(), + reallocated_swa.decode_swa_ragged_indptr.data_ptr(), + ) + expected_reallocated = reference(reallocated_mla, reallocated_swa) + run_decode(reallocated_mla, reallocated_swa) + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2) + assert not torch.equal(captured_full, expected_reallocated) + + graph.replay() + torch.accelerator.synchronize() + torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2) + + # --------------------------------------------------------------------------- # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum) # --------------------------------------------------------------------------- diff --git a/tests/models/test_deepseek_v4_dspark_rocm.py b/tests/models/test_deepseek_v4_dspark_rocm.py new file mode 100644 index 0000000000..efa97e10d7 --- /dev/null +++ b/tests/models/test_deepseek_v4_dspark_rocm.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.deepseek_v4.amd import dspark as dspark_module +from vllm.models.deepseek_v4.amd.dspark import DSparkDeepseekV4ForCausalLM + + +def _make_uninitialized_model(confidence_head): + model = DSparkDeepseekV4ForCausalLM.__new__(DSparkDeepseekV4ForCausalLM) + object.__setattr__( + model, + "model", + SimpleNamespace(confidence_head=confidence_head), + ) + return model + + +def _prepare_loader_model(model, named_parameters): + object.__setattr__( + model, + "config", + SimpleNamespace( + n_routed_experts=1, + expert_dtype="fp4", + num_attention_heads=1, + ), + ) + object.__setattr__(model, "named_parameters", lambda: named_parameters) + + +def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + dspark_module, + "fused_moe_make_expert_params_mapping", + lambda *args, **kwargs: [], + ) + monkeypatch.setattr( + dspark_module, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_maps_enabled_confidence_head(): + model = _make_uninitialized_model(object()) + + assert ( + model._remap_dspark_name("mtp.2.confidence_head.proj.weight") + == "model.confidence_head.proj.weight" + ) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_skips_disabled_confidence_head(): + model = _make_uninitialized_model(None) + + assert model._remap_dspark_name("mtp.2.confidence_head.proj.weight") is None + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_disables_unloaded_confidence_head( + monkeypatch: pytest.MonkeyPatch, +): + model = _make_uninitialized_model(object()) + _prepare_loader_model(model, []) + _disable_distributed_loader_paths(monkeypatch) + + assert model.load_weights([]) == set() + assert model.model.confidence_head is None + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_loads_complete_confidence_head( + monkeypatch: pytest.MonkeyPatch, +): + class FakeParameter: + loaded_weight = None + + def weight_loader(self, param, loaded_weight): + assert param is self + self.loaded_weight = loaded_weight + + confidence_head = object() + parameter = FakeParameter() + model = _make_uninitialized_model(confidence_head) + _prepare_loader_model( + model, + [("model.confidence_head.proj.weight", parameter)], + ) + _disable_distributed_loader_paths(monkeypatch) + loaded_weight = torch.tensor([[1.0, 2.0]]) + + loaded = model.load_weights([("mtp.2.confidence_head.proj.weight", loaded_weight)]) + + assert loaded == {"model.confidence_head.proj.weight"} + assert parameter.loaded_weight is loaded_weight + assert model.model.confidence_head is confidence_head + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_confidence_is_probability(): + class ConfidenceHead: + def __call__(self, head_hidden, markov_embed): + return (head_hidden[:, 0] + markov_embed[:, 0]).float() + + model = _make_uninitialized_model(ConfidenceHead()) + head_hidden = torch.tensor([[0.0], [1.0]], dtype=torch.bfloat16) + markov_embed = torch.tensor([[0.0], [-2.0]], dtype=torch.bfloat16) + + confidence = model.compute_confidence(head_hidden, markov_embed) + + torch.testing.assert_close( + confidence, + torch.sigmoid(torch.tensor([0.0, -1.0])), + ) + assert torch.all((confidence >= 0) & (confidence <= 1)) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_dspark_confidence_requires_a_head(): + model = _make_uninitialized_model(None) + empty = torch.zeros((1, 1)) + + with pytest.raises(RuntimeError, match="confidence_head"): + model.compute_confidence(empty, empty) diff --git a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py new file mode 100644 index 0000000000..914cd249fe --- /dev/null +++ b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.deepseek_v4.amd.rocm import ( + DeepseekV4ROCMAiterMLASparseMetadataBuilder, + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, +) +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.mla import indexer +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, +) + + +def _make_indexer_config(*, architecture: str, adaptive: bool): + return SimpleNamespace( + model_config=SimpleNamespace(architectures=[architecture]), + num_speculative_tokens=1, + speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive), + ) + + +def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None: + # Platform predicates are mutually exclusive in production. Override both + # so these ROCm policy tests do not inherit CUDA capabilities from the CI + # host running them. + monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: False) + monkeypatch.setattr(indexer.current_platform, "is_rocm", lambda: True) + + +def _make_indexer_builder(*, adaptive: bool, capacity: int = 12): + builder = DeepseekV32IndexerMetadataBuilder.__new__( + DeepseekV32IndexerMetadataBuilder + ) + builder.vllm_config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive) + ) + builder.supports_varlen = False + builder.decode_seq_lens_buffer = torch.zeros(capacity, dtype=torch.int32) + builder.expanded_block_table_buffer = torch.zeros((capacity, 2), dtype=torch.int32) + builder.decode_lens_buffer = torch.zeros(capacity, dtype=torch.int32) + builder.arange_buffer = torch.arange(capacity, dtype=torch.int32) + return builder + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs(): + adaptive_config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=True) + ) + fixed_config = SimpleNamespace( + speculative_config=SimpleNamespace(enable_adaptive_verification=False) + ) + + for builder_cls in ( + DeepseekV4ROCMAiterMLASparseMetadataBuilder, + DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + ): + assert ( + builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace()) + == AttentionCGSupport.ALWAYS + ) + assert ( + builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace()) + == AttentionCGSupport.UNIFORM_BATCH + ) + + +@pytest.mark.cpu_test +def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch): + _mock_rocm_platform(monkeypatch) + adaptive_config = _make_indexer_config( + architecture="DeepseekV4ForCausalLM", adaptive=True + ) + fixed_config = _make_indexer_config( + architecture="DeepseekV4ForCausalLM", adaptive=False + ) + + assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch() + assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config) + assert ( + DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder + ) + assert indexer._use_flattening(adaptive_config) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + adaptive_config, SimpleNamespace() + ) + == AttentionCGSupport.ALWAYS + ) + assert not indexer._use_flattening(fixed_config) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + fixed_config, SimpleNamespace() + ) + == AttentionCGSupport.UNIFORM_BATCH + ) + + +@pytest.mark.cpu_test +@pytest.mark.parametrize( + "architecture", + ["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"], +) +def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4( + monkeypatch: pytest.MonkeyPatch, + architecture: str, +): + _mock_rocm_platform(monkeypatch) + adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True) + + assert not indexer._rocm_supports_flattened_device_query_lens(adaptive_config) + assert not indexer._use_flattening(adaptive_config) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + adaptive_config, SimpleNamespace() + ) + == AttentionCGSupport.UNIFORM_BATCH + ) + + if architecture == "DeepseekV41ForCausalLM": + assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch() + + +@pytest.mark.cpu_test +def test_rocm_adaptive_indexer_preserves_single_request_uniform_path( + monkeypatch: pytest.MonkeyPatch, +): + builder = _make_indexer_builder(adaptive=True, capacity=8) + + class FakeUniformKernel: + called = False + + def __call__( + self, + seq_lens, + decode_seq_lens, + block_table, + expanded_block_table, + decode_lens, + num_decode_tokens, + max_decode_len, + ): + self.called = True + assert num_decode_tokens == 4 + assert block_table.stride(0) == expanded_block_table.stride(0) == 2 + decode_seq_lens[:max_decode_len] = torch.arange( + seq_lens[0] - max_decode_len + 1, + seq_lens[0] + 1, + dtype=torch.int32, + ) + expanded_block_table[:max_decode_len] = block_table[0] + decode_lens[:max_decode_len] = 1 + + fake_kernel = FakeUniformKernel() + monkeypatch.setattr(indexer, "_PREPARE_UNIFORM_DECODE_KERNEL", fake_kernel) + + seq_lens, block_table, decode_lens, batch_size, requires_padding = ( + builder._prepare_decode_tensors( + seq_lens=torch.tensor([10], dtype=torch.int32), + block_table=torch.tensor([[1, 2]], dtype=torch.int32), + decode_lens=torch.tensor([4], dtype=torch.int32), + decode_lens_cpu=torch.tensor([4], dtype=torch.int32), + query_start_loc=torch.tensor([0], dtype=torch.int32), + num_decodes=1, + num_decode_tokens=4, + use_native=False, + next_n=8, + max_decode_len=4, + ) + ) + + assert fake_kernel.called + torch.testing.assert_close(seq_lens, torch.tensor([7, 8, 9, 10], dtype=torch.int32)) + torch.testing.assert_close( + block_table, + torch.tensor([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=torch.int32), + ) + torch.testing.assert_close(decode_lens, torch.ones(4, dtype=torch.int32)) + assert batch_size == 4 + assert not requires_padding + + +@pytest.mark.cpu_test +def test_rocm_adaptive_indexer_replays_changed_allocations_with_stable_buffers(): + builder = _make_indexer_builder(adaptive=True) + buffer_ptrs = ( + builder.decode_seq_lens_buffer.data_ptr(), + builder.expanded_block_table_buffer.data_ptr(), + builder.decode_lens_buffer.data_ptr(), + ) + block_table = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) + decode_lens_cpu = torch.tensor([4, 4, 0], dtype=torch.int32) + + first = builder._prepare_decode_tensors( + seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32), + block_table=block_table, + decode_lens=torch.tensor([7, 1, 0], dtype=torch.int32), + decode_lens_cpu=decode_lens_cpu, + query_start_loc=torch.tensor([0, 7, 8], dtype=torch.int32), + num_decodes=3, + num_decode_tokens=10, + use_native=False, + next_n=8, + max_decode_len=4, + ) + torch.testing.assert_close( + first[0], + torch.tensor([14, 15, 16, 17, 18, 19, 20, 20, 0, 0], dtype=torch.int32), + ) + torch.testing.assert_close( + first[1][:, 0], + torch.tensor([1, 1, 1, 1, 1, 1, 1, 3, 0, 0], dtype=torch.int32), + ) + + second = builder._prepare_decode_tensors( + seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32), + block_table=block_table, + decode_lens=torch.tensor([1, 7, 0], dtype=torch.int32), + decode_lens_cpu=decode_lens_cpu, + query_start_loc=torch.tensor([0, 1, 8], dtype=torch.int32), + num_decodes=3, + num_decode_tokens=10, + use_native=False, + next_n=8, + max_decode_len=4, + ) + torch.testing.assert_close( + second[0], + torch.tensor([20, 14, 15, 16, 17, 18, 19, 20, 0, 0], dtype=torch.int32), + ) + torch.testing.assert_close( + second[1][:, 0], + torch.tensor([1, 3, 3, 3, 3, 3, 3, 3, 0, 0], dtype=torch.int32), + ) + torch.testing.assert_close(second[2], torch.ones(10, dtype=torch.int32)) + assert second[3:] == (10, False) + assert buffer_ptrs == ( + builder.decode_seq_lens_buffer.data_ptr(), + builder.expanded_block_table_buffer.data_ptr(), + builder.decode_lens_buffer.data_ptr(), + ) diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py index edca59032a..c3f249fa80 100644 --- a/vllm/models/deepseek_v4/amd/dspark.py +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -44,6 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.qwen3_dspark import ( + DSparkConfidenceHead, DSparkMarkovHead, ) from vllm.model_executor.models.utils import maybe_prefix @@ -102,7 +103,7 @@ class DSparkDeepseekV4Model(nn.Module): ] ) - # Heads: final norm + hc_head, and the Markov head + # Heads: final norm + hc_head, and the Markov + confidence heads # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) hc_dim = self.hc_mult * config.hidden_size @@ -125,6 +126,12 @@ class DSparkDeepseekV4Model(nn.Module): config.dspark_markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) + self.confidence_head: DSparkConfidenceHead | None = None + if getattr(config, "enable_confidence_head", True): + self.confidence_head = DSparkConfidenceHead( + config.hidden_size + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "confidence_head"), + ) # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch), # replacing the direct nvidia tilelang kernel call. @@ -358,6 +365,17 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed, self.logits_processor) + def compute_confidence( + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor + ) -> torch.Tensor: + """Per-position acceptance probability for each drafted token.""" + if self.model.confidence_head is None: + raise RuntimeError( + "compute_confidence() requires a confidence head, but the " + "checkpoint did not provide confidence_head weights." + ) + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed)) + # --- Weight loading ---------------------------------------------------- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -391,6 +409,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + loaded_confidence_head = False tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() @@ -403,6 +422,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): if mapped is None: continue name = mapped + if "confidence_head." in name: + loaded_confidence_head = True # ``.scale`` -> per-method scale suffix. if name.endswith(".scale"): @@ -439,8 +460,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): continue # Stacked rules only apply to decoder-layer weights. Head-stack params - # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g. - # "markov_w1" would collide with the "w1" shard rule. + # (main_proj/norm/hc_head/markov_head/confidence_head) load directly — + # otherwise e.g. "markov_w1" would collide with the "w1" shard rule. is_layer_param = name.startswith("model.layers.") for param_name, weight_name, stacked_shard_id in stacked_params_mapping: if not is_layer_param or weight_name not in name: @@ -469,6 +490,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): weight_loader(param, loaded_weight) loaded_params.add(name) + if self.model.confidence_head is not None and not loaded_confidence_head: + self.model.confidence_head = None logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -482,8 +505,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): return None stage = int(m.group(1)) rest = m.group(2) - # The confidence head is not wired into inference yet; drop its weights. - if rest.startswith("confidence_head."): + if rest.startswith("confidence_head.") and self.model.confidence_head is None: return None # Head-stack params live at model level (mtp.last), context combiner at # model level (mtp.0); everything else is a per-layer decoder block. @@ -493,6 +515,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module): "hc_head_base", "hc_head_scale", "markov_head.", + "confidence_head.", ) if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( head_prefixes diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index a70a23e3ab..b9478b508f 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -8,6 +8,7 @@ from typing import cast import torch from vllm import envs +from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, @@ -26,6 +27,7 @@ from vllm.platforms.rocm import _ON_GFX950 from vllm.triton_utils import tl, triton from vllm.utils.multi_stream_utils import execute_in_parallel from vllm.v1.attention.backend import ( + AttentionCGSupport, CommonAttentionMetadata, ) from vllm.v1.attention.backends.mla.sparse_swa import ( @@ -38,6 +40,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( rocm_sparse_attn_decode, rocm_sparse_attn_prefill, ) +from vllm.v1.kv_cache_interface import KVCacheSpec from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) @@ -445,6 +448,20 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.enable_adaptive_verification: + # All per-token metadata is built from device query boundaries into + # persistent buffers, so adaptive verification can replay varlen + # FULL decode graphs after reallocating drafts across requests. + return AttentionCGSupport.ALWAYS + return super().get_cudagraph_support(vllm_config, kv_cache_spec) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None @@ -512,6 +529,20 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBui class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): + @classmethod + def get_cudagraph_support( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> AttentionCGSupport: + spec_config = vllm_config.speculative_config + if spec_config is not None and spec_config.enable_adaptive_verification: + # SWA indices, lengths, and token-to-request mappings are built from + # device boundaries into persistent buffers, so adaptive verification + # can replay varlen FULL decode graphs safely. + return AttentionCGSupport.ALWAYS + return super().get_cudagraph_support(vllm_config, kv_cache_spec) + # Keep fused multi-step decode disabled until update_draft_decode_metadata() # also refreshes the ROCm-specific ragged SWA indices and indptrs. supports_draft_decode_metadata_update = False diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 00d14ceeb7..1ede0442bc 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -237,6 +237,15 @@ class KpoolTailBackend(DeepseekV32IndexerBackend): class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + # ROCm runs adaptive verification through the per-token flattened + # indexer path, which derives row ownership from device decode lengths. + return ( + current_platform.is_rocm() + or super().supports_device_cpu_query_lens_mismatch() + ) + @staticmethod def get_name() -> str: return "DEEPSEEK_V4_INDEXER" @@ -254,6 +263,12 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend): + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + # The ROCm flattened-query support above is validated for the + # DeepSeek-V4 adaptive DSpark path only. + return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch() + @staticmethod def get_name() -> str: return "DEEPSEEK_V41_INDEXER" @@ -802,6 +817,15 @@ def _supports_flattened_device_query_lens() -> bool: ) +def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool: + model_config = vllm_config.model_config + return ( + current_platform.is_rocm() + and model_config is not None + and "DeepseekV4ForCausalLM" in model_config.architectures + ) + + def _supports_native_decode(next_n: int) -> bool: """Whether decode can pass `next_n` Q rows per request to the kernel instead of flattening to one single-token row per query, which re-reads @@ -822,7 +846,10 @@ def _use_flattening(vllm_config: VllmConfig) -> bool: return not _supports_native_decode(next_n) or ( speculative_config is not None and speculative_config.enable_adaptive_verification - and _supports_flattened_device_query_lens() + and ( + _supports_flattened_device_query_lens() + or _rocm_supports_flattened_device_query_lens(vllm_config) + ) )