feat: DeepSeek-V4 DSA context-parallel on NPU

tests/core/layers/npu_torch/CMakeLists.txt

@@ -57,3 +57,33 @@ target_link_options(npu_linear_w8a8_dynamic_test PRIVATE
   "-Wl,--whole-archive"
   "${CMAKE_BINARY_DIR}/third_party/spdlog/libspdlog.a"
   "-Wl,--no-whole-archive")
+
+# Pure-CPU metadata unit test for DSA-CP (M0). No device ops; validates the
+# host-side token-partition math in DSAMetadataBuilder::build_cp_local_metadata.
+cc_test(
+  NAME
+    npu_dsa_cp_metadata_test
+  SRCS
+    dsa_cp_metadata_tests.cpp
+  DEPS
+    :common_layers
+    :parallel_state
+    :model
+    :state_dict
+    glog::glog
+    torch
+    GTest::gtest_main
+)
+
+target_link_libraries(npu_dsa_cp_metadata_test
+                      PRIVATE
+                      ascendcl
+                      hccl
+                      c_sec
+                      nnopbase
+                      atb)
+
+target_link_options(npu_dsa_cp_metadata_test PRIVATE
+  "-Wl,--whole-archive"
+  "${CMAKE_BINARY_DIR}/third_party/spdlog/libspdlog.a"
+  "-Wl,--no-whole-archive")

tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp

@@ -0,0 +1,220 @@
+/* Copyright 2025-2026 The xLLM Authors.
+
+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://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,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Pure-CPU unit tests for DSAMetadataBuilder::build_cp_local_metadata (M0 of
+// the DeepSeek-V4 DSA-CP plan). These tests do not touch any NPU device; they
+// validate the host-side token-partition math against hand-computed golden
+// vectors, including the worked example from the vllm-ascend
+// `_build_local_token_metadata` docstring.
+
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "core/layers/common/dsa_metadata.h"
+#include "core/layers/common/dsa_metadata_builder.h"
+
+namespace xllm::layer {
+namespace {
+
+torch::Tensor i32(const std::vector<int32_t>& v) {
+  return torch::tensor(v, torch::dtype(torch::kInt32).device(torch::kCPU));
+}
+
+std::vector<int64_t> to_vec(const torch::Tensor& t) {
+  auto c = t.to(torch::kCPU).to(torch::kInt64).contiguous();
+  return std::vector<int64_t>(c.data_ptr<int64_t>(),
+                              c.data_ptr<int64_t>() + c.numel());
+}
+
+// query_start_loc / seq_lens for the docstring example:
+//   9 requests, seq_lens = [1..9], pure prefill so q_lens == seq_lens,
+//   query_start_loc = [0,1,3,6,10,15,21,28,36,45], num_input_tokens = 45.
+torch::Tensor docstring_qsl() {
+  return i32({0, 1, 3, 6, 10, 15, 21, 28, 36, 45});
+}
+torch::Tensor docstring_seq_lens() { return i32({1, 2, 3, 4, 5, 6, 7, 8, 9}); }
+
+}  // namespace
+
+// Rank 1 of the docstring example is the canonical golden vector.
+TEST(DsaCpMetadataTest, DocstringRank1) {
+  auto cp = DSAMetadataBuilder::build_cp_local_metadata(
+      docstring_qsl(), docstring_seq_lens(), /*cp_size=*/3, /*cp_rank=*/1);
+
+  EXPECT_EQ(cp.num_tokens_pad, 45);
+  EXPECT_EQ(cp.tokens_per_rank, 15);
+  EXPECT_EQ(cp.local_start, 15);
+  EXPECT_EQ(cp.local_end, 30);
+  EXPECT_EQ(to_vec(cp.local_query_start_loc),
+            (std::vector<int64_t>{0, 0, 0, 0, 0, 0, 6, 13, 15, 15}));
+  EXPECT_EQ(to_vec(cp.local_seq_lens),
+            (std::vector<int64_t>{0, 0, 0, 0, 0, 6, 7, 2, 0}));
+  // local_kv_start_loc is the leading-0 cumsum of local_seq_lens. The last
+  // rank's queries see the full KV, so this differs from local_query_start_loc.
+  EXPECT_EQ(to_vec(cp.local_kv_start_loc),
+            (std::vector<int64_t>{0, 0, 0, 0, 0, 0, 6, 13, 15, 15}));
+}
+
+TEST(DsaCpMetadataTest, DocstringRank0) {
+  auto cp = DSAMetadataBuilder::build_cp_local_metadata(
+      docstring_qsl(), docstring_seq_lens(), /*cp_size=*/3, /*cp_rank=*/0);
+  EXPECT_EQ(cp.local_start, 0);
+  EXPECT_EQ(cp.local_end, 15);
+  EXPECT_EQ(to_vec(cp.local_query_start_loc),
+            (std::vector<int64_t>{0, 1, 3, 6, 10, 15, 15, 15, 15, 15}));
+  EXPECT_EQ(to_vec(cp.local_seq_lens),
+            (std::vector<int64_t>{1, 2, 3, 4, 5, 0, 0, 0, 0}));
+}
+
+TEST(DsaCpMetadataTest, DocstringRank2) {
+  auto cp = DSAMetadataBuilder::build_cp_local_metadata(
+      docstring_qsl(), docstring_seq_lens(), /*cp_size=*/3, /*cp_rank=*/2);
+  EXPECT_EQ(cp.local_start, 30);
+  EXPECT_EQ(cp.local_end, 45);
+  EXPECT_EQ(to_vec(cp.local_query_start_loc),
+            (std::vector<int64_t>{0, 0, 0, 0, 0, 0, 0, 0, 6, 15}));
+  EXPECT_EQ(to_vec(cp.local_seq_lens),
+            (std::vector<int64_t>{0, 0, 0, 0, 0, 0, 0, 8, 9}));
+}
+
+// The local query lengths across all ranks must partition the global tokens
+// exactly (no token dropped or double-counted).
+TEST(DsaCpMetadataTest, RanksPartitionAllQueryTokens) {
+  const int32_t cp_size = 3;
+  int64_t total = 0;
+  for (int32_t r = 0; r < cp_size; ++r) {
+    auto cp = DSAMetadataBuilder::build_cp_local_metadata(
+        docstring_qsl(), docstring_seq_lens(), cp_size, r);
+    total += to_vec(cp.local_query_start_loc).back();  // local total tokens
+  }
+  EXPECT_EQ(total, 45);
+}
+
+// cp_size == 1 must be the identity view.
+TEST(DsaCpMetadataTest, CpSizeOneIsIdentity) {
+  auto cp = DSAMetadataBuilder::build_cp_local_metadata(
+      docstring_qsl(), docstring_seq_lens(), /*cp_size=*/1, /*cp_rank=*/0);
+  EXPECT_EQ(cp.num_tokens_pad, 45);
+  EXPECT_EQ(cp.tokens_per_rank, 45);
+  EXPECT_EQ(cp.local_start, 0);
+  EXPECT_EQ(cp.local_end, 45);
+  EXPECT_EQ(to_vec(cp.local_query_start_loc), to_vec(docstring_qsl()));
+  EXPECT_EQ(to_vec(cp.local_seq_lens), to_vec(docstring_seq_lens()));
+  // cp_size==1: local_kv_start_loc is the cumsum of the full seq_lens, which
+  // for pure prefill equals the query cu-seqlens.
+  EXPECT_EQ(to_vec(cp.local_kv_start_loc), to_vec(docstring_qsl()));
+}
+
+// Non-divisible token count must pad up to a multiple of cp_size, and a
+// boundary-crossing request must land its tail on the later rank with the full
+// KV length visible there.
+TEST(DsaCpMetadataTest, PaddingAndBoundaryCross) {
+  auto qsl = i32({0, 3, 5});    // 2 reqs, 5 tokens
+  auto seq_lens = i32({3, 2});  // pure prefill
+
+  auto r0 = DSAMetadataBuilder::build_cp_local_metadata(
+      qsl, seq_lens, /*cp_size=*/2, /*cp_rank=*/0);
+  EXPECT_EQ(r0.num_tokens_pad, 6);
+  EXPECT_EQ(r0.tokens_per_rank, 3);
+  EXPECT_EQ(r0.local_start, 0);
+  EXPECT_EQ(r0.local_end, 3);
+  EXPECT_EQ(to_vec(r0.local_query_start_loc), (std::vector<int64_t>{0, 3, 3}));
+  EXPECT_EQ(to_vec(r0.local_seq_lens), (std::vector<int64_t>{3, 0}));
+  EXPECT_EQ(to_vec(r0.local_kv_start_loc), (std::vector<int64_t>{0, 3, 3}));
+
+  auto r1 = DSAMetadataBuilder::build_cp_local_metadata(
+      qsl, seq_lens, /*cp_size=*/2, /*cp_rank=*/1);
+  EXPECT_EQ(r1.local_start, 3);
+  EXPECT_EQ(r1.local_end, 6);
+  EXPECT_EQ(to_vec(r1.local_query_start_loc), (std::vector<int64_t>{0, 0, 2}));
+  EXPECT_EQ(to_vec(r1.local_seq_lens), (std::vector<int64_t>{0, 2}));
+  EXPECT_EQ(to_vec(r1.local_kv_start_loc), (std::vector<int64_t>{0, 0, 2}));
+}
+
+// A single long request split across ranks is the real-world DSA-CP prefill
+// case (the gpqa regression): every rank's local queries causally see the
+// FULL KV stream, so local_kv_start_loc (ori_kv cu-seqlens) must sum to the
+// full length on every rank -- and must DIFFER from local_query_start_loc on
+// the later ranks, which is exactly the bug the ori_kv path had.
+TEST(DsaCpMetadataTest, SingleLongRequestOriKvSeesFullStream) {
+  auto qsl = i32({0, 332});    // 1 request, 332 tokens
+  auto seq_lens = i32({332});  // pure prefill
+
+  auto r0 = DSAMetadataBuilder::build_cp_local_metadata(
+      qsl, seq_lens, /*cp_size=*/2, /*cp_rank=*/0);
+  // rank0 owns queries [0,166): it only sees KV [0,166), so here the ori_kv
+  // cu-seqlens coincides with the local query cu-seqlens.
+  EXPECT_EQ(to_vec(r0.local_query_start_loc), (std::vector<int64_t>{0, 166}));
+  EXPECT_EQ(to_vec(r0.local_seq_lens), (std::vector<int64_t>{166}));
+  EXPECT_EQ(to_vec(r0.local_kv_start_loc), (std::vector<int64_t>{0, 166}));
+
+  auto r1 = DSAMetadataBuilder::build_cp_local_metadata(
+      qsl, seq_lens, /*cp_size=*/2, /*cp_rank=*/1);
+  // rank1 owns queries [166,332): its last query (pos 331) sees the FULL KV
+  // [0,332). local_query_start_loc sums to 166 (local q count) but the ori_kv
+  // cu-seqlens must sum to 332 (full KV) -- they MUST differ.
+  EXPECT_EQ(to_vec(r1.local_query_start_loc), (std::vector<int64_t>{0, 166}));
+  EXPECT_EQ(to_vec(r1.local_seq_lens), (std::vector<int64_t>{332}));
+  EXPECT_EQ(to_vec(r1.local_kv_start_loc), (std::vector<int64_t>{0, 332}));
+  EXPECT_NE(to_vec(r1.local_kv_start_loc), to_vec(r1.local_query_start_loc));
+}
+
+// RoPE tables padded to num_tokens_pad must be sliced to [local_start,
+// local_end).
+TEST(DsaCpMetadataTest, RopeTablesSlicedToLocalInterval) {
+  const int64_t rope_dim = 4;
+  // arange table so the slice is trivially checkable: row i == [i,i,i,i].
+  auto full = torch::arange(45, torch::dtype(torch::kFloat32))
+                  .view({45, 1})
+                  .expand({45, rope_dim})
+                  .contiguous();
+
+  auto cp = DSAMetadataBuilder::build_cp_local_metadata(docstring_qsl(),
+                                                        docstring_seq_lens(),
+                                                        /*cp_size=*/3,
+                                                        /*cp_rank=*/1,
+                                                        full,
+                                                        full);
+
+  ASSERT_TRUE(cp.local_cos.defined());
+  ASSERT_TRUE(cp.local_sin.defined());
+  EXPECT_EQ(cp.local_cos.size(0), 15);
+  EXPECT_EQ(cp.local_cos.size(1), rope_dim);
+  // Rows should be positions 15..29.
+  EXPECT_FLOAT_EQ(cp.local_cos[0][0].item<float>(), 15.0f);
+  EXPECT_FLOAT_EQ(cp.local_cos[14][0].item<float>(), 29.0f);
+}
+
+// Empty batch must not crash and returns a single-element leading-0 cumsum.
+TEST(DsaCpMetadataTest, EmptyBatch) {
+  auto qsl = i32({0});
+  auto seq_lens = torch::zeros({0}, torch::dtype(torch::kInt32));
+  auto cp = DSAMetadataBuilder::build_cp_local_metadata(
+      qsl, seq_lens, /*cp_size=*/2, /*cp_rank=*/0);
+  EXPECT_EQ(cp.num_tokens_pad, 0);
+  EXPECT_EQ(to_vec(cp.local_query_start_loc), (std::vector<int64_t>{0}));
+  EXPECT_EQ(cp.local_seq_lens.numel(), 0);
+}
+
+// Invalid cp_rank must be rejected.
+TEST(DsaCpMetadataTest, RejectsInvalidRank) {
+  EXPECT_THROW(DSAMetadataBuilder::build_cp_local_metadata(
+                   docstring_qsl(), docstring_seq_lens(), 2, 2),
+               c10::Error);
+}
+
+}  // namespace xllm::layer

xllm/core/common/global_flags.h

@@ -105,6 +105,8 @@ DECLARE_int32(ep_size);
 
 DECLARE_int32(cp_size);
 
+DECLARE_bool(enable_dsa_cp);
+
 DECLARE_int64(tp_size);
 
 DECLARE_int64(sp_size);

xllm/core/framework/config/parallel_config.cpp

@@ -33,6 +33,13 @@ DEFINE_int32(kv_split_size,
              "AllGather); other K (K divides cp_size) means KV is sharded "
              "across K ranks while token-CP still uses cp_size.");
 
+DEFINE_bool(enable_dsa_cp,
+            false,
+            "Enable DeepSeek-V4 DSA context parallel (prefill-only, M1). "
+            "Only takes effect when cp_size > 1 and the model is DeepSeek-V4 "
+            "(DSA). When false, DSA attention runs its original non-CP path "
+            "even if cp_size > 1.");
+
 DEFINE_int64(tp_size, 1, "Tensor parallelism size, only used for DiT model.");
 
 DEFINE_int64(sp_size, 1, "Sequence parallelism size, only used for DiT model.");
@@ -76,6 +83,7 @@ void ParallelConfig::from_flags() {
   XLLM_CONFIG_ASSIGN_FROM_FLAG(ep_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(cp_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(kv_split_size);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_dsa_cp);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(tp_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(cfg_size);
@@ -91,6 +99,7 @@ void ParallelConfig::from_json(const JsonReader& json) {
   XLLM_CONFIG_ASSIGN_FROM_JSON(dp_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(ep_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(cp_size);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(enable_dsa_cp);
   XLLM_CONFIG_ASSIGN_FROM_JSON(tp_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size);
   XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size);
@@ -108,6 +117,8 @@ void ParallelConfig::append_config_json(
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, dp_size);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, ep_size);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, cp_size);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, enable_dsa_cp);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, tp_size);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, sp_size);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(

xllm/core/framework/config/parallel_config.h

@@ -44,6 +44,7 @@ class ParallelConfig final {
         {"dp_size",
          "ep_size",
          "cp_size",
+         "enable_dsa_cp",
          "tp_size",
          "sp_size",
          "cfg_size",
@@ -65,6 +66,9 @@ class ParallelConfig final {
   // 0 means follow cp_size (legacy KV-split width).
   PROPERTY(int32_t, kv_split_size) = 1;
 
+  // Enable DeepSeek-V4 DSA context parallel (prefill-only, M1).
+  PROPERTY(bool, enable_dsa_cp) = false;
+
   PROPERTY(int64_t, tp_size) = 1;
 
   PROPERTY(int64_t, sp_size) = 1;

xllm/core/layers/common/dsa_metadata.h

@@ -53,6 +53,42 @@ struct DSACompressedAttentionMetadata {
   int64_t max_context_len = 0;
 };
 
+// Context-parallel (DSA-CP) per-rank metadata. Aligned with vllm-ascend
+// DSACPMetadata (vllm_ascend/attention/context_parallel/dsa_cp.py). Produced
+// by DSAMetadataBuilder::build_cp_local_metadata when cp_size > 1.
+//
+// The query token stream is split evenly across the CP group. Each field
+// below describes the CURRENT rank's local token slice:
+//   - local_query_start_loc: (num_reqs+1,) rebased cumulative query starts
+//     for the local slice, always starting at 0.
+//   - local_seq_lens:        (num_reqs,)   per-request KV length visible to
+//     the local slice (0 for requests with no local query tokens).
+//   - local_kv_start_loc:    (num_reqs+1,) cumulative form of local_seq_lens
+//     with a leading 0. This is the cu-seqlens the sparse-attn / metadata
+//     kernels expect for the (full, replicated) ori_kv path: it encodes the
+//     right-aligned KV extent each local query causally sees. It differs from
+//     local_query_start_loc on any rank whose local queries can see KV tokens
+//     that live on earlier ranks (e.g. the last rank sees the full stream),
+//     so the ori_kv path must use this, not local_query_start_loc.
+//   - local_start / local_end: [start, end) token interval owned by this rank
+//     in the padded global token stream.
+//   - tokens_per_rank:       num_tokens_pad / cp_size.
+//   - num_tokens_pad:        num_input_tokens padded to a multiple of
+//     cp_size.
+//   - local_cos / local_sin: RoPE tables sliced to [local_start, local_end).
+//     Undefined when no full RoPE table is supplied to the builder.
+struct DSACPMetadata {
+  torch::Tensor local_query_start_loc;
+  torch::Tensor local_seq_lens;
+  torch::Tensor local_kv_start_loc;
+  int64_t local_start = 0;
+  int64_t local_end = 0;
+  int64_t tokens_per_rank = 0;
+  int64_t num_tokens_pad = 0;
+  torch::Tensor local_cos;
+  torch::Tensor local_sin;
+};
+
 // DSAMetadata contains DeepSeek V4 sparse attention specific metadata,
 // aligned with Python DSAMetadata(AttentionMetadata) class.
 // It is built once at the beginning of model forward pass and reused by
@@ -80,9 +116,19 @@ struct DSAMetadata {
   // not perform host/device copies in this mode.
   bool is_acl_graph = false;
 
-  // cp_input_dict: context-parallel inputs placeholder (reserved, optional)
+  // cp_input_dict: context-parallel inputs placeholder (reserved, optional).
+  // Kept for backward compatibility with any callers that still read it;
+  // structured CP metadata lives in cp_metadata below.
   std::unordered_map<std::string, torch::Tensor> cp_input_dict;
 
+  // ===== DSA-CP (context-parallel) fields =====
+  // cp_enabled: true when cp_size > 1 and DSA-CP is active for this forward.
+  bool cp_enabled = false;
+  int32_t cp_size = 1;
+  int32_t cp_rank = 0;
+  // Per-rank local token metadata; only meaningful when cp_enabled == true.
+  DSACPMetadata cp_metadata;
+
   // NPU-only DSA metadata fields.
   // RoPE caches selected for the current layer's q/kv/output RoPE.
   torch::Tensor cos;