refactor: address DeepSeek V4 DSpark review feedback.

Author: chenchuw886

docs/src/content/docs/en/cookbook/autoregressive_models/deepseek/deepseek_v4.md

@@ -218,8 +218,8 @@ block geometry and is out of the trained distribution; use it only after an
 acceptance/performance evaluation. Context parallelism (`cp_size > 1`) is not
 supported on this path yet.
 
-xLLM supports two SAS modes. The default compatibility mode works with CANN
-9.0 and needs no extra option. If the installed SAS operator accepts a
+On NPU, xLLM supports two SAS modes. The default compatibility mode works with
+CANN 9.0 and needs no extra option. If the installed SAS operator accepts a
 non-empty `ori_sparse_indices`, set `--enable_dspark_native_sas=true` to use
 the complete DSpark SWA window. Older operators terminate during tiling, so
 xLLM cannot safely detect this capability automatically.

docs/src/content/docs/zh/cookbook/autoregressive_models/deepseek/deepseek_v4.md

@@ -215,8 +215,8 @@ DeepSeek-V4-Flash-0731 权重目录即可：
 `dspark_block_size=5` 训练。改用其他 gamma 会改变扩散块几何，超出训练分布，
 需要重新验证接受率和性能。当前路径暂不支持 `cp_size > 1`。
 
-xLLM 支持两种 SAS 模式。默认兼容模式适配 CANN 9.0，无需增加参数；若当前
-SAS 算子支持非空 `ori_sparse_indices`，可设置
+在 NPU 上，xLLM 支持两种 SAS 模式。默认兼容模式适配 CANN 9.0，无需增加参数；
+若当前 SAS 算子支持非空 `ori_sparse_indices`，可设置
 `--enable_dspark_native_sas=true`，使用完整的 DSpark SWA 窗口。旧版算子会在
 tiling 阶段直接终止进程，因此无法安全地自动探测该能力。
 

tests/core/framework/config/config_json_test.cpp

@@ -24,6 +24,7 @@ limitations under the License.
 #include "core/common/global_flags.h"
 #include "core/framework/config/config_utils.h"
 #include "core/framework/config/execution_config.h"
+#include "core/framework/config/kernel_config.h"
 #include "core/framework/config/kv_cache_config.h"
 #include "core/framework/config/model_config.h"
 #include "core/framework/config/parallel_config.h"
@@ -51,6 +52,16 @@ inline constexpr std::string_view kMalformedConfig = R"json({
   "block_size":
 })json";
 
+#if !defined(USE_NPU)
+TEST(KernelConfigTest, RejectsNpuOnlyDsparkNativeSas) {
+  JsonReader json_config =
+      config::parse_json_string(R"json({"enable_dspark_native_sas":true})json");
+  KernelConfig kernel_config;
+  EXPECT_DEATH(kernel_config.from_json(json_config),
+               "enable_dspark_native_sas is only supported on NPU");
+}
+#endif
+
 class ConfigJsonFileFlagGuard final {
  public:
   explicit ConfigJsonFileFlagGuard(const std::string& config_json_file)

tests/core/layers/npu_torch/deepseek_v4_indexer_tests.cpp

@@ -203,6 +203,31 @@ TEST_F(DeepseekV4IndexerTest, DSparkSparseTilingUsesSupportedWindow) {
       deepseek_v4_uses_prefill_sparse_metadata(params.meta.batch_forward_type));
 }
 
+TEST_F(DeepseekV4IndexerTest, DSparkNativeSwaIndicesAreSharedByQueryRows) {
+  const torch::Tensor block_table =
+      torch::tensor({{10, 11, 12}}, torch::kInt32);
+  const torch::Tensor query_cu_seq_lens = torch::tensor({0, 3}, torch::kInt32);
+  const torch::Tensor seq_lens = torch::tensor({6}, torch::kInt32);
+
+  const torch::Tensor indices =
+      build_dspark_swa_indices(block_table,
+                               query_cu_seq_lens,
+                               seq_lens,
+                               /*window_size=*/4,
+                               /*num_speculative_tokens=*/3,
+                               /*cache_block_size=*/4);
+
+  ASSERT_EQ(indices.dim(), 3);
+  ASSERT_EQ(indices.size(0), 3);
+  ASSERT_EQ(indices.size(1), 1);
+  ASSERT_EQ(indices.size(2), 128);
+  const torch::Tensor expected_prefix =
+      torch::tensor({40, 41, 42, 43, 44, 45, -1}, torch::kInt32);
+  for (int64_t row = 0; row < indices.size(0); ++row) {
+    EXPECT_TRUE(torch::equal(indices[row][0].slice(0, 0, 7), expected_prefix));
+  }
+}
+
 TEST_F(DeepseekV4IndexerTest, DsaDummyAttentionUsesPositionDevice) {
   ModelInputParams params;
   params.meta.batch_forward_type = BatchForwardType::DECODE;

tests/core/runtime/acl_graph_executor_test.cpp

@@ -20,6 +20,8 @@ limitations under the License.
 #include <torch_npu/torch_npu.h>
 
 #include <cstdlib>
+#include <filesystem>
+#include <fstream>
 #include <memory>
 #include <optional>
 #include <vector>
@@ -30,6 +32,7 @@ limitations under the License.
 #include "core/framework/block/block.h"
 #include "core/framework/block/block_manager_impl.h"
 #include "core/framework/config/execution_config.h"
+#include "core/framework/config/kernel_config.h"
 #include "core/framework/config/speculative_config.h"
 #include "core/framework/kv_cache/kv_cache.h"
 #include "core/framework/kv_cache/kv_cache_utils.h"
@@ -47,6 +50,7 @@ limitations under the License.
 #include "core/runtime/acl_graph_executor_impl.h"
 #include "core/runtime/acl_graph_persistent_param.h"
 #include "core/runtime/base_executor_impl.h"
+#include "core/runtime/block_diffusion_model_config.h"
 #include "core/runtime/dflash_worker_impl.h"
 #include "core/runtime/mtp_async_state.h"
 #include "core/runtime/options.h"
@@ -1211,6 +1215,77 @@ TEST(DSparkWorkerOptionsTest, PreservesDraftBlockSize) {
   EXPECT_EQ(dflash_detail::draft_model_num_speculative_tokens(options), 0);
 }
 
+TEST(BlockDiffusionConfigTest, MapsCheckpointLayersToNpuCapturePoints) {
+  EXPECT_EQ(block_diffusion::map_target_layer_ids_to_capture_points(
+                std::vector<int32_t>{0, 40, 42}),
+            (std::vector<int32_t>{1, 41, 43}));
+}
+
+TEST(BlockDiffusionConfigTest, PreservesDeepseekV4NpuDraftArguments) {
+  const std::filesystem::path config_dir =
+      std::filesystem::path(::testing::TempDir()) /
+      "xllm_block_diffusion_config_test";
+  std::filesystem::create_directories(config_dir);
+  {
+    std::ofstream config_file(config_dir / "config.json");
+    config_file << R"json({"dspark_target_layer_ids":[40,41,42]})json";
+  }
+
+  runtime::Options target_options;
+  target_options.speculative_algorithm("DSpark")
+      .draft_model_path(config_dir.string())
+      .num_speculative_tokens(5)
+      .is_draft_engine(false);
+  ModelArgs target_args;
+  target_args.model_type("deepseek_v4")
+      .n_layers(43)
+      .dspark_num_layers(3)
+      .dspark_block_size(0)
+      .compress_ratios({1, 1, 4});
+  block_diffusion::configure_model_args(
+      target_args, target_options, /*model_weights_path=*/"unused");
+  EXPECT_EQ(target_args.model_type(), "deepseek_v4");
+  EXPECT_EQ(target_args.n_layers(), 43);
+  EXPECT_EQ(target_args.dspark_block_size(), 0);
+  EXPECT_EQ(target_args.layers_to_capture(),
+            (std::vector<int32_t>{41, 42, 43}));
+  EXPECT_EQ(target_args.compress_ratios(), (std::vector<int32_t>{1, 1, 4}));
+
+  runtime::Options draft_options;
+  draft_options.speculative_algorithm("DSpark")
+      .num_speculative_tokens(5)
+      .is_draft_engine(true);
+  ModelArgs draft_args;
+  draft_args.model_type("deepseek_v4")
+      .n_layers(43)
+      .n_hash_layers(2)
+      .dspark_num_layers(3)
+      .dspark_block_size(0)
+      .compress_ratios({1, 1, 4});
+
+  KernelConfig& kernel_config = KernelConfig::get_instance();
+  const bool original_native_sas = kernel_config.enable_dspark_native_sas();
+  kernel_config.enable_dspark_native_sas(false);
+  block_diffusion::configure_model_args(
+      draft_args, draft_options, config_dir.string());
+  kernel_config.enable_dspark_native_sas(original_native_sas);
+
+  EXPECT_EQ(draft_args.model_type(), "deepseek_v4_dspark");
+  EXPECT_EQ(draft_args.n_layers(), 3);
+  EXPECT_EQ(draft_args.n_hash_layers(), 0);
+  EXPECT_EQ(draft_args.dspark_block_size(), 5);
+  EXPECT_FALSE(draft_args.dspark_use_native_sas());
+  EXPECT_EQ(draft_args.layers_to_capture(), (std::vector<int32_t>{41, 42, 43}));
+  EXPECT_EQ(draft_args.compress_ratios(), (std::vector<int32_t>{1, 1, 1}));
+
+  std::filesystem::remove_all(config_dir);
+}
+
+TEST(DSparkNativeSasConfigTest, DefaultsToCompatibilityMode) {
+  KernelConfig config;
+  EXPECT_FALSE(config.enable_dspark_native_sas());
+}
+
 TEST(DSparkWorkerInputTest, InvalidatesTargetAttentionMetadataOnly) {
   ModelInputParams params;
   params.attn_metadata = std::make_shared<layer::AttentionMetadata>();
@@ -1234,16 +1309,17 @@ TEST(DSparkWorkerWeightsTest, PreservesDeepseekDraftHeadAndEmbedding) {
 TEST(DSparkSasFallbackTest, ChoosesCompatibleRowsUnlessNativeIsEnabled) {
   ModelArgs draft_args;
   draft_args.model_type("deepseek_v4_dspark");
-  EXPECT_TRUE(dflash_detail::uses_dsa_block_parallel_query_rows(
-      draft_args, /*sample_from_anchor=*/true));
-  EXPECT_FALSE(dflash_detail::uses_native_dspark_sas(
-      draft_args, /*sample_from_anchor=*/true));
+  EXPECT_EQ(dflash_detail::classify_dspark_sas_mode(
+                draft_args, /*sample_from_anchor=*/true),
+            dflash_detail::DSparkSasMode::COMPATIBILITY);
 
   draft_args.dspark_use_native_sas(true);
-  EXPECT_FALSE(dflash_detail::uses_dsa_block_parallel_query_rows(
-      draft_args, /*sample_from_anchor=*/true));
-  EXPECT_TRUE(dflash_detail::uses_native_dspark_sas(
-      draft_args, /*sample_from_anchor=*/true));
+  EXPECT_EQ(dflash_detail::classify_dspark_sas_mode(
+                draft_args, /*sample_from_anchor=*/true),
+            dflash_detail::DSparkSasMode::NATIVE);
+  EXPECT_EQ(dflash_detail::classify_dspark_sas_mode(
+                draft_args, /*sample_from_anchor=*/false),
+            dflash_detail::DSparkSasMode::NOT_DSPARK);
 }
 
 TEST(DSparkWorkerWeightsTest, DedicatedVocabularyOverridesFallbackInAnyOrder) {

tests/models/model_registry_test.cpp

@@ -4,7 +4,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    https://www.apache.org/licenses/LICENSE-2.0
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,

xllm/core/common/global_flags.h

@@ -82,8 +82,6 @@ DECLARE_bool(enable_opt_validate_probs);
 
 DECLARE_bool(enable_mtp_draft_body_tp1);
 
-DECLARE_bool(enable_dspark_native_sas);
-
 DECLARE_int32(speculative_suffix_cache_max_depth);
 
 DECLARE_double(speculative_suffix_max_spec_factor);
@@ -403,6 +401,8 @@ DECLARE_bool(enable_aclnn_matmul);
 
 DECLARE_bool(enable_aclnn_swiglu);
 
+DECLARE_bool(enable_dspark_native_sas);
+
 DECLARE_bool(enable_flashcomm1);
 
 DECLARE_int32(flashcomm1_min_prefill_tokens);

xllm/core/distributed_runtime/speculative_output_metrics.h

@@ -4,7 +4,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    https://www.apache.org/licenses/LICENSE-2.0
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -43,13 +43,11 @@ SpeculativeOutputStats calculate_speculative_output_stats_typed(
     for (int64_t column = 0; column < token_width; ++column) {
       if (data[row * token_width + column] >= static_cast<T>(0)) {
         ++stats.committed_tokens;
-      }
-    }
-    for (int64_t position = 0; position < num_speculative_tokens; ++position) {
-      // Column 0 is always the first committed token. Draft position p was
-      // accepted exactly when output column p+1 is non-negative.
-      if (data[row * token_width + position + 1] >= static_cast<T>(0)) {
-        ++stats.accepted_per_position[static_cast<size_t>(position)];
+        // Column 0 is always the first committed token. Draft position p was
+        // accepted exactly when output column p+1 is non-negative.
+        if (column > 0 && column <= num_speculative_tokens) {
+          ++stats.accepted_per_position[static_cast<size_t>(column - 1)];
+        }
       }
     }
   }

xllm/core/distributed_runtime/worker_service.cpp

@@ -61,8 +61,10 @@ int32_t get_num_decode_seqs_for_schedule_overlap(const ForwardInput& input) {
       unpacked_input.sampling_params.sample_idxes.size(0));
 }
 
-void record_speculative_metrics_from_output(const torch::Tensor& next_tokens,
-                                            const runtime::Options& options) {
+void record_speculative_metrics_from_output(
+    const torch::Tensor& next_tokens,
+    const runtime::Options& options,
+    const std::vector<std::string>& position_labels) {
   if (!options.enable_speculative_decode() || !next_tokens.defined() ||
       next_tokens.dim() != 2 || next_tokens.numel() == 0) {
     return;
@@ -75,6 +77,7 @@ void record_speculative_metrics_from_output(const torch::Tensor& next_tokens,
       token_width != num_speculative_tokens + 1) {
     return;
   }
+  CHECK_EQ(position_labels.size(), static_cast<size_t>(num_speculative_tokens));
 
   torch::Tensor tokens = next_tokens.contiguous();
   worker_service_detail::SpeculativeOutputStats stats =
@@ -93,7 +96,7 @@ void record_speculative_metrics_from_output(const torch::Tensor& next_tokens,
         stats.accepted_per_position[static_cast<size_t>(position)];
     num_accepted_tokens += accepted;
     MULTI_COUNTER_ADD(speculative_num_accepted_tokens_per_pos,
-                      std::to_string(position),
+                      position_labels[static_cast<size_t>(position)],
                       accepted);
   }
   COUNTER_ADD(speculative_num_drafts_total, batch_size);
@@ -109,6 +112,21 @@ void record_speculative_metrics_from_output(const torch::Tensor& next_tokens,
   }
 }
 
+std::vector<std::string> build_speculative_position_labels(
+    const runtime::Options& options) {
+  const int32_t num_speculative_tokens = options.num_speculative_tokens();
+  if (num_speculative_tokens <= 0) {
+    return {};
+  }
+
+  std::vector<std::string> labels;
+  labels.reserve(static_cast<size_t>(num_speculative_tokens));
+  for (int32_t position = 0; position < num_speculative_tokens; ++position) {
+    labels.emplace_back(std::to_string(position));
+  }
+  return labels;
+}
+
 torch::Tensor clone_cpu_tensor_view(const torch::Tensor& tensor) {
   if (!tensor.defined()) {
     return tensor;
@@ -128,7 +146,10 @@ void stabilize_schedule_overlap_host_views(ForwardInput& input) {
 
 WorkerService::WorkerService(runtime::Options options,
                              const torch::Device& device)
-    : options_(options), device_(device), initialized_(false) {
+    : options_(options),
+      speculative_position_labels_(build_speculative_position_labels(options)),
+      initialized_(false),
+      device_(device) {
   device_.set_device();
   device_.init_device_context();
   stream_ = device_.get_stream_from_pool();
@@ -143,9 +164,10 @@ WorkerService::WorkerService(runtime::Options options,
                              const torch::Device& device,
                              std::unique_ptr<Worker> worker)
     : options_(options),
+      speculative_position_labels_(build_speculative_position_labels(options)),
+      initialized_(true),
       device_(device),
-      worker_(std::move(worker)),
-      initialized_(true) {
+      worker_(std::move(worker)) {
   device_.set_device();
   device_.init_device_context();
   stream_ = device_.get_stream_from_pool();
@@ -281,7 +303,8 @@ void WorkerService::step(ForwardInput& fwd_input,
         } else {
           stream_->synchronize();
         }
-        record_speculative_metrics_from_output(next_tokens, options_);
+        record_speculative_metrics_from_output(
+            next_tokens, options_, speculative_position_labels_);
       }
     }
   } else {
@@ -911,7 +934,8 @@ void WorkerService::GetLastStepResult(
                 device_.index());
 #endif
           }
-          record_speculative_metrics_from_output(next_tokens, options_);
+          record_speculative_metrics_from_output(
+              next_tokens, options_, speculative_position_labels_);
 
           if (next_tokens.defined() || !dit_images.empty() ||
               !dit_text_output.empty() ||

xllm/core/distributed_runtime/worker_service.h

@@ -16,6 +16,7 @@ limitations under the License.
 #pragma once
 
 #include <string>
+#include <vector>
 
 #include "runtime/forward_shared_memory_manager.h"
 #include "runtime/worker.h"
@@ -170,6 +171,7 @@ class WorkerService : public proto::DistributeWorker {
  private:
   // runtime options
   runtime::Options options_;
+  std::vector<std::string> speculative_position_labels_;
 
   bool initialized_;