feat: support python dspark acl graphs with data parallelism.

tests/core/framework/speculative/spec_input_builder_test.cpp

@@ -20,6 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "core/framework/model/mtp_topk_state.h"
+#include "core/util/tensor_helper.h"
 #include "framework/model/model_input_params.h"
 #include "models/llm/mlu/mtp_topk_state.h"
 #include "runtime/forward_params.h"
@@ -110,6 +111,51 @@ TEST(SpecDecodeInputBuilderTest, DraftInputsSingleRowPerSeq) {
   EXPECT_EQ(buf.out_kv_seq_lens, to_layout_seq_lens({6, 10}));
 }
 
+TEST(SpecDecodeInputBuilderTest, EmptyBlockDraftMaterializesWholeQuery) {
+  LlmForwardInput input;
+  input.input_params.parallel.dp_global_token_nums = {0, 21};
+  input.input_params.parallel.raw_dp_global_token_nums = {0, 21};
+  input.input_params.parallel.dp_global_sequence_nums = {0, 3};
+  input.runtime.input_host_buffer_has_layout = true;
+  materialize_empty_block_draft_rows(input, 7, 42);
+
+  const auto& params = input.input_params;
+  EXPECT_EQ(params.meta.num_sequences, 1);
+  EXPECT_EQ(params.meta.actual_num_sequences, 0);
+  EXPECT_EQ(params.meta.q_max_seq_len, 7);
+  EXPECT_EQ(params.meta.kv_max_seq_len, 1);
+  EXPECT_EQ(tensor_to_vector<int32_t>(input.token_ids_host),
+            std::vector<int32_t>(7, 42));
+  EXPECT_EQ(tensor_to_vector<int32_t>(input.positions_host),
+            std::vector<int32_t>({0, 1, 2, 3, 4, 5, 6}));
+  EXPECT_EQ(params.attention.host.q_seq_lens, to_layout_seq_lens({7}));
+  EXPECT_EQ(params.attention.host.q_cu_seq_lens, std::vector<int32_t>({0, 7}));
+  EXPECT_EQ(params.attention.host.kv_seq_lens, to_layout_seq_lens({1}));
+  EXPECT_EQ(params.attention.host.new_cache_slots, std::vector<int32_t>(7, -1));
+  EXPECT_EQ(params.parallel.dp_global_token_nums,
+            std::vector<int32_t>({7, 21}));
+  EXPECT_EQ(params.parallel.raw_dp_global_token_nums,
+            std::vector<int32_t>({0, 21}));
+  EXPECT_EQ(params.parallel.dp_global_sequence_nums,
+            std::vector<int32_t>({0, 3}));
+  EXPECT_FALSE(input.runtime.device_tensors_ready);
+  EXPECT_FALSE(input.runtime.input_host_buffer_has_layout);
+}
+
+TEST(SpecDecodeInputBuilderTest, ActiveBlockDraftOnlyFillsIdlePeerCount) {
+  LlmForwardInput input;
+  input.input_params.meta.num_sequences = 3;
+  input.input_params.parallel.dp_global_token_nums = {21, 0};
+  input.token_ids = torch::tensor({11, 12}, torch::kInt32);
+  materialize_empty_block_draft_rows(input, 7, 42);
+
+  EXPECT_EQ(input.input_params.meta.num_sequences, 3);
+  EXPECT_EQ(input.input_params.parallel.dp_global_token_nums,
+            std::vector<int32_t>({21, 7}));
+  EXPECT_EQ(tensor_to_vector<int32_t>(input.token_ids),
+            std::vector<int32_t>({11, 12}));
+}
+
 TEST(SpecDecodeInputBuilderTest, VlmVerifyRowsPreserveVisionState) {
   VlmForwardInput input;
   input.input_params.meta.num_sequences = 2;

tests/python/test_decode_acl_graph.py

@@ -83,40 +83,182 @@ def prepare_metadata(metadata: SimpleNamespace, *, device_kv_lengths: bool) -> S
         )
 
 
-def _block_draft_runner() -> BlockDraftAclGraphRunner:
+def _block_draft_runner(max_batch: int = 1, **kwargs: object) -> BlockDraftAclGraphRunner:
     attention_backend = _FakeBlockDraftAttentionBackend()
     return BlockDraftAclGraphRunner(
         nn.Identity(),
         attention_backend,
         torch.device("cpu"),
-        max_batch=1,
+        max_batch=max_batch,
         max_model_len=8,
+        **kwargs,
     )
 
 
 def _block_draft_metadata(block_table: torch.Tensor) -> SimpleNamespace:
     sequence_count, block_cols = block_table.shape
     query_width = 2
     token_count = sequence_count * query_width
+    query_ends = [query_width * (index + 1) for index in range(sequence_count)]
     return SimpleNamespace(
         is_prefill=False,
         is_chunked_prefill=True,
         is_spec_verify=False,
         q_seq_lens=torch.full((sequence_count,), query_width, dtype=torch.int32),
-        q_cu_seq_lens=torch.tensor([0, token_count], dtype=torch.int32),
-        q_cu_seq_lens_host_values=[token_count],
+        q_cu_seq_lens=torch.tensor([0, *query_ends], dtype=torch.int32),
+        q_cu_seq_lens_host_values=query_ends,
         q_seq_lens_host=None,
         block_table=block_table,
         slot_mapping=torch.arange(token_count, dtype=torch.int32),
         kv_seq_lens=torch.full((sequence_count,), block_cols * 4, dtype=torch.int32),
         kv_cu_seq_lens=None,
-        kv_seq_lens_host_values=[block_cols * 4],
-        paged_kv_indptr=torch.tensor([0, block_cols], dtype=torch.int32),
-        paged_kv_indices=torch.arange(block_cols, dtype=torch.int32),
-        paged_kv_last_page_len=torch.tensor([4], dtype=torch.int32),
+        kv_seq_lens_host_values=[block_cols * 4] * sequence_count,
+        paged_kv_indptr=torch.arange(sequence_count + 1, dtype=torch.int32) * block_cols,
+        paged_kv_indices=torch.arange(sequence_count * block_cols, dtype=torch.int32),
+        paged_kv_last_page_len=torch.full((sequence_count,), 4, dtype=torch.int32),
     )
 
 
+@pytest.mark.parametrize("dp_rank,sequence_count", [(0, 1), (1, 3)])
+def test_dspark_dp_graph_uses_group_wide_complete_blocks(dp_rank: int, sequence_count: int) -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2, dp_rank=dp_rank)
+    metadata = _block_draft_metadata(torch.zeros(sequence_count, 3, dtype=torch.int32))
+    metadata.dp_execution_token_counts = (2, 6)
+    metadata.dp_is_decode = (1, 1)
+    input_ids = torch.arange(sequence_count * 2, dtype=torch.int32)
+
+    assert runner.can_execute(input_ids, metadata)
+    assert runner._padded_sequence_count(metadata, sequence_count, 2) == 4
+    entry = runner._allocate_entry(input_ids, input_ids, metadata, sequence_count=4)
+    runner._fill_entry(entry, input_ids, input_ids, metadata)
+    static = entry.static_metadata
+    assert entry.static_input_ids.numel() == 8
+    assert static.q_cu_seq_lens.tolist() == [0, 2, 4, 6, 8]
+    assert static.q_seq_lens.tolist() == [2] * 4
+    assert static.dp_execution_token_counts == (8, 8)
+    assert static.dp_is_decode == (1, 1)
+    assert static.slot_mapping[sequence_count * 2 :].tolist() == [-1] * (8 - sequence_count * 2)
+    assert static.kv_seq_lens.tolist() == [12] * sequence_count + [1] * (4 - sequence_count)
+    assert static.prepared_attention_state.actual_seq_q == [2, 4, 6, 8]
+    assert static.prepared_attention_state.actual_seq_kv == [16] * 4
+
+
+def test_dspark_dp_graph_admits_an_idle_rank_full_block() -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2, dp_rank=0)
+    metadata = _block_draft_metadata(torch.zeros(1, 1, dtype=torch.int32))
+    metadata.dp_execution_token_counts = (2, 6)
+    metadata.dp_global_sequence_nums = (0, 3)
+    metadata.dp_is_decode = (1, 1)
+    metadata.slot_mapping.fill_(-1)
+    metadata.kv_seq_lens.fill_(1)
+    metadata.kv_seq_lens_host_values = [1]
+    metadata.paged_kv_last_page_len.fill_(1)
+    input_ids = torch.zeros(2, dtype=torch.int32)
+
+    assert runner.can_execute(input_ids, metadata)
+    entry = runner._allocate_entry(input_ids, input_ids, metadata, sequence_count=4)
+    runner._fill_entry(entry, input_ids, input_ids, metadata)
+    assert entry.static_metadata.slot_mapping.tolist() == [-1] * 8
+    assert entry.static_metadata.kv_seq_lens.tolist() == [1] * 4
+
+
+@pytest.mark.parametrize(
+    "counts,phases,error",
+    [
+        ((2,), (1, 1), "every DP rank"),
+        ((2, 3), (1, 1), "complete query blocks"),
+        ((4, 6), (1, 1), "local execution count"),
+        ((2, 0), (1, 1), "complete query blocks"),
+    ],
+)
+def test_dspark_dp_graph_rejects_invalid_counts(counts: tuple[int, ...], phases: tuple[int, ...], error: str) -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2)
+    metadata = _block_draft_metadata(torch.zeros(1, 1, dtype=torch.int32))
+    metadata.dp_execution_token_counts = counts
+    metadata.dp_is_decode = phases
+    with pytest.raises(ValueError, match=error):
+        runner.can_execute(torch.zeros(2, dtype=torch.int32), metadata)
+
+
+@pytest.mark.parametrize("counts,phases", [((2, 10), (1, 1)), ((2, 6), (1, 0))])
+def test_dspark_dp_graph_falls_back_group_wide(counts: tuple[int, ...], phases: tuple[int, ...]) -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2)
+    metadata = _block_draft_metadata(torch.zeros(1, 1, dtype=torch.int32))
+    metadata.dp_execution_token_counts = counts
+    metadata.dp_is_decode = phases
+    assert not runner.can_execute(torch.zeros(2, dtype=torch.int32), metadata)
+
+
+def test_dspark_dp_replay_clears_rows_from_a_larger_local_batch() -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2)
+    metadata = _block_draft_metadata(torch.full((3, 2), 10, dtype=torch.int32))
+    input_ids = torch.arange(6, dtype=torch.int32) + 10
+    metadata.kv_cu_seq_lens = torch.tensor([0, 8, 16, 24], dtype=torch.int32)
+    entry = runner._allocate_entry(input_ids, input_ids, metadata, sequence_count=4)
+    runner._fill_entry(entry, input_ids, input_ids, metadata)
+    smaller = _block_draft_metadata(torch.tensor([[20]], dtype=torch.int32))
+    runner._fill_entry(entry, input_ids[:2], input_ids[:2], smaller)
+
+    assert entry.static_input_ids.tolist() == [10, 11, 0, 0, 0, 0, 0, 0]
+    assert entry.static_positions.tolist() == [10, 11, 0, 0, 0, 0, 0, 0]
+    static = entry.static_metadata
+    assert static.slot_mapping.tolist() == [0, 1, -1, -1, -1, -1, -1, -1]
+    assert static.kv_seq_lens.tolist() == [4, 1, 1, 1]
+    assert static.kv_cu_seq_lens.tolist() == [0, 4, 5, 6, 7]
+    assert static.paged_kv_indptr.tolist() == [0, 1, 1, 1, 1]
+    assert static.paged_kv_last_page_len.tolist() == [4, 1, 1, 1]
+    assert static.block_table[1:].count_nonzero().item() == 0
+
+
+def test_dspark_dp_graph_key_tracks_peer_page_capacity() -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2)
+    input_ids = torch.arange(2, dtype=torch.int32)
+    local_key = runner._graph_key(input_ids, _block_draft_metadata(torch.zeros(1, 1, dtype=torch.int32)), 4, 2)
+    with (
+        patch("xllm.python.model_executor.runners.decode_acl_graph.dist.is_initialized", return_value=True),
+        patch(
+            "xllm.python.distributed.all_gather",
+            side_effect=[torch.tensor([1, 1]), torch.tensor([1, 2]), torch.tensor([1, 1])],
+        ),
+    ):
+        first = runner._synchronize_dp_graph_key(local_key, input_ids)
+        changed_peer = runner._synchronize_dp_graph_key(local_key, input_ids)
+        restored = runner._synchronize_dp_graph_key(local_key, input_ids)
+    assert first[:-1] == local_key[:-1]
+    assert first != changed_peer
+    assert first == restored
+
+
+def test_dspark_dp_warmup_recaptures_when_any_rank_is_missing() -> None:
+    runner = _block_draft_runner(max_batch=8, dp_size=2)
+    metadata = _block_draft_metadata(torch.zeros(1, 1, dtype=torch.int32))
+    metadata.dp_execution_token_counts = (2, 6)
+    metadata.dp_is_decode = (1, 1)
+    input_ids = torch.arange(2, dtype=torch.int32)
+    local_key = runner._graph_key(input_ids, metadata, 4, 2)
+    graph_key = (*local_key[:-1], (1, 1))
+    runner._graphs[graph_key] = object()
+    expected_eplb = Mock()
+
+    def presence(count: torch.Tensor, *, group_name: str) -> None:
+        assert group_name == "dp"
+        count.fill_(1)
+
+    def prepare(*args: object, graph_key: tuple[object, ...], eplb: object | None = None) -> None:
+        assert eplb is expected_eplb
+        assert graph_key not in runner._graphs
+        runner._graphs[graph_key] = object()
+
+    with (
+        patch("xllm.python.model_executor.runners.decode_acl_graph.dist.is_initialized", return_value=True),
+        patch("xllm.python.distributed.all_gather", return_value=torch.tensor([1, 1])),
+        patch("xllm.python.distributed.all_reduce_", side_effect=presence),
+        patch.object(runner, "_prepare_graph_entry", side_effect=prepare) as capture,
+    ):
+        assert runner.warmup(input_ids, input_ids, metadata, eplb=expected_eplb) == graph_key
+    capture.assert_called_once()
+
+
 def test_dspark_graph_key_buckets_live_page_table_width() -> None:
     runner = _block_draft_runner()
     input_ids = torch.zeros(2, dtype=torch.int32)

tests/python/test_model_executor.py

@@ -52,6 +52,9 @@
     get_forward_context,
     record_layer_event,
 )
+from xllm.python.model_executor.runners.block_draft_acl_graph import (  # noqa: E402
+    BlockDraftAclGraphRunner,
+)
 from xllm.python.model_executor.runners.decode_acl_graph import (  # noqa: E402
     DecodeAclGraphRunner,
 )
@@ -2322,6 +2325,26 @@ def test_executor_selects_acl_graph_input_owner(model_type: str, dp_size: int, p
     assert len(runner.layer_caches) == 2
 
 
+@pytest.mark.parametrize("dp_size", [1, 2, 4])
+def test_executor_selects_dspark_acl_graph_with_dp(dp_size: int) -> None:
+    config = {
+        "model_type": "DSparkDraftModel",
+        "dp_size": dp_size,
+        "dp_rank": dp_size - 1,
+        "python_graph_backend": "aclgraph",
+        "max_position_embeddings": 128,
+    }
+    with patch(
+        "xllm.python.model_executor.executor._create_attention_backend", return_value=_PreparedStubAttentionBackend()
+    ):
+        executor = ModelExecutor(_FakeModel(), config, max_seqs_per_batch=5, acl_graph_decode_batch_size_limit=2)
+    runner = executor.decode_graph_runner
+    assert isinstance(runner, BlockDraftAclGraphRunner)
+    assert runner.dp_size == dp_size
+    assert runner.dp_rank == dp_size - 1
+    assert runner.max_batch == min((5 + dp_size - 1) // dp_size, 2)
+
+
 def test_eager_runner_skips_cp_for_empty_rank() -> None:
     runner = _make_eager_runner()
     metadata = SimpleNamespace(

xllm/core/framework/speculative/spec_input_builder.cpp

@@ -489,6 +489,44 @@ void update_input_params(Params& input_params,
 torch::Tensor make_cpu_int_tensor(const std::vector<int32_t>& values) {
   return make_pinned_cpu_tensor(values);
 }
+
+void materialize_empty_block_draft_rows(LlmForwardInput& input,
+                                        int32_t query_width,
+                                        int32_t mask_token_id) {
+  CHECK_GT(query_width, 0);
+  auto& params = input.input_params;
+  for (int32_t& count : params.parallel.dp_global_token_nums) {
+    if (count == 0) {
+      count = query_width;
+    }
+  }
+  if (params.meta.num_sequences != 0) {
+    return;
+  }
+  const torch::TensorOptions options =
+      torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU);
+  input.token_ids_host = torch::full({query_width}, mask_token_id, options);
+  input.positions_host = torch::arange(query_width, options);
+  input.token_ids = input.token_ids_host;
+  input.positions = input.positions_host;
+  input.runtime.device_tensors_ready = false;
+  input.runtime.input_host_buffer_has_layout = false;
+  params.meta.num_sequences = 1;
+  params.meta.actual_num_sequences = 0;
+  params.meta.q_max_seq_len = query_width;
+  params.meta.kv_max_seq_len = 1;
+  params.attention.host.q_seq_lens.clear();
+  append_seq_len_by_layout(params.attention.host.q_seq_lens, query_width);
+  params.attention.host.q_cu_seq_lens = {0, query_width};
+  params.attention.host.kv_seq_lens.clear();
+  append_seq_len_by_layout(params.attention.host.kv_seq_lens, 1);
+  params.attention.host.new_cache_slots.assign(query_width, -1);
+  params.attention.host.block_tables = torch::zeros({1, 1}, options);
+  params.attention.device.paged_kv_indptr = torch::tensor({0, 1}, options);
+  params.attention.device.paged_kv_indices = torch::tensor({0}, options);
+  params.attention.device.paged_kv_last_page_len = torch::tensor({1}, options);
+}
+
 template <typename Input>
 void set_token_position_tensors(Input& input,
                                 const std::vector<int32_t>& token_ids,

xllm/core/framework/speculative/spec_input_builder.h

@@ -180,6 +180,13 @@ void update_input_params(Params& input_params,
 // Packs a host int32 vector into a pinned CPU tensor for async H2D staging.
 torch::Tensor make_cpu_int_tensor(const std::vector<int32_t>& values);
 
+// Empty DP ranks participate in block-draft execution with a complete query
+// block. Fill their execution counts on every rank and materialize an idle
+// rank's host input without writing any real cache slots.
+void materialize_empty_block_draft_rows(LlmForwardInput& input,
+                                        int32_t query_width,
+                                        int32_t mask_token_id);
+
 // Stages token_ids/positions into both the host and device tensors of `input`
 // with async H2D copies, toggling device_tensors_ready around the write.
 template <typename Input>

xllm/core/runtime/dflash_worker_impl.cpp

@@ -28,6 +28,7 @@ limitations under the License.
 #include "common/metrics.h"
 #include "core/framework/config/execution_config.h"
 #include "core/framework/config/kernel_config.h"
+#include "core/framework/config/model_config.h"
 #include "core/framework/config/scheduler_config.h"
 #include "core/framework/config/speculative_config.h"
 #include "core/framework/speculative/adaptive_pruning_helpers.h"
@@ -593,6 +594,13 @@ std::optional<ForwardOutput> DFlashWorkerImpl::step_empty(
   }
   scale_speculative_parallel_token_counts(query_input.input_params,
                                           draft_width);
+  if (sample_from_anchor() && !use_block_parallel_rows &&
+      ModelConfig::is_python_model_impl(context_.get_model_impl())) {
+    c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard();
+    specBuilder::materialize_empty_block_draft_rows(
+        query_input, draft_width, mask_token_id_);
+    query_input.input_params.attention.rebuild_device_buffer(device_.unwrap());
+  }
   // Warmup only: prime the draft; its output is unused. Keep it alive until the
   // sync below so the no-sync draft input is not freed while the target forward
   // launched next can reuse the buffer.
@@ -1259,6 +1267,11 @@ void DFlashWorkerImpl::prepare_query_inputs(const LlmForwardInput& input,
                                    /*update_block_tables=*/
                                    use_block_parallel_rows);
   scale_speculative_parallel_token_counts(input_params, query_width);
+  if (sample_from_anchor() && !use_block_parallel_rows &&
+      ModelConfig::is_python_model_impl(context_.get_model_impl())) {
+    specBuilder::materialize_empty_block_draft_rows(
+        query_input, query_width, mask_token_id_);
+  }
   input_params.attention.rebuild_device_buffer(device_);
 
   query_input.sampling_params.selected_token_idxes =

xllm/python/model_executor/executor.py

@@ -415,14 +415,14 @@ def __init__(
                 )
             else:
                 if config.get("model_type") == "DSparkDraftModel":
-                    if dp_size != 1:
-                        raise NotImplementedError("DSpark ACL graphs currently require dp_size == 1")
                     self.decode_graph_runner = BlockDraftAclGraphRunner(
                         execution_model,
                         self.attention_backend,
                         device,
-                        local_graph_sequence_capacity,
+                        local_graph_sequence_capacity * dp_size,
                         int(config["max_position_embeddings"]),
+                        dp_size,
+                        dp_rank,
                     )
                 else:
                     self.decode_graph_runner = DecodeAclGraphRunner(