From 1ecd91fe6e6172643b65c4912703feaadc04fb92 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 21 Jul 2026 20:53:38 +0800 Subject: [PATCH 1/8] feat: DeepSeek-V4 DSA context-parallel (Model 1, CP == TP) Add prefill-only DSA context parallelism for DeepSeek-V4 on Ascend NPU, aligned with vllm-ascend dsa_cp.py. A single switch (--enable_dsa_cp) turns it on; the CP geometry is derived from TP (the TP group IS the CP token-split group), so there is no separate cp_size dimension to mis-configure. This replaces the earlier orthogonal cp_size approach, where the token split used cp_size while the reassembly gather ran over every TP member (cp_size(2) vs cp_group world_size(4)), corrupting the full-token stream. Core mechanism (DSAttention forward): - Derive cp_size/cp_rank/cp_group from TP when --enable_dsa_cp is set; q is computed on this rank's LOCAL token slice. - q_b and attn_sink are loaded FULL-head (replicated via the single-rank group) so each rank runs attention over ALL heads on its local tokens; decode / non-CP prefill slice both back to the TP head shard, keeping those paths byte-identical to the non-CP baseline. - Replace the reassembly gather with all_to_all_4D (scatter=2, gather=1): [T_local, num_heads, head_dim] -> [T_full, n_local_heads, head_dim]. Head shard r maps to TP rank r, so the standard TP o_proj runs unchanged. o_proj moves after the transpose. Metadata (deepseek_v4.h / deepseek_v4_mtp.h) applies the same CP=TP derivation so the precomputed sparse metadata (local cu_seqlens, right-aligned ori_kv via local_kv_start_loc) matches the local token slice the attention runs on. The MTP draft shares the DSAttention layer and needs the identical derivation. Validated on DeepSeek-V4-Flash-w8a8-mtp (dp=4, tp=4): gpqa_diamond recovers from 0.28 (broken cp=2) to ~0.9, matching / exceeding the cp=1 baseline (0.8), with MTP enabled. --- tests/core/layers/npu_torch/CMakeLists.txt | 30 ++ .../npu_torch/dsa_cp_metadata_tests.cpp | 220 +++++++++++ xllm/core/common/global_flags.h | 4 + .../core/framework/config/parallel_config.cpp | 14 + xllm/core/layers/common/dsa_metadata.h | 48 ++- .../layers/common/dsa_metadata_builder.cpp | 125 +++++++ .../core/layers/common/dsa_metadata_builder.h | 33 ++ .../npu_torch/deepseek_sparse_attention.cpp | 346 +++++++++++++++--- .../npu_torch/deepseek_sparse_attention.h | 17 + .../layers/npu_torch/deepseek_v4_indexer.cpp | 17 +- .../layers/npu_torch/deepseek_v4_indexer.h | 8 +- xllm/core/platform/platform.h | 8 + xllm/core/runtime/llm_worker_impl.cpp | 11 +- xllm/core/runtime/worker_impl.cpp | 10 + xllm/models/llm/deepseek_v4.h | 132 ++++++- xllm/models/llm/deepseek_v4_mtp.h | 100 ++++- 16 files changed, 1030 insertions(+), 93 deletions(-) create mode 100644 tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp diff --git a/tests/core/layers/npu_torch/CMakeLists.txt b/tests/core/layers/npu_torch/CMakeLists.txt index 09162f4f2e..7e653993c4 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/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") diff --git a/tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp b/tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp new file mode 100644 index 0000000000..f2dbfa598d --- /dev/null +++ b/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 +#include + +#include + +#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& v) { + return torch::tensor(v, torch::dtype(torch::kInt32).device(torch::kCPU)); +} + +std::vector to_vec(const torch::Tensor& t) { + auto c = t.to(torch::kCPU).to(torch::kInt64).contiguous(); + return std::vector(c.data_ptr(), + c.data_ptr() + 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{0, 0, 0, 0, 0, 0, 6, 13, 15, 15})); + EXPECT_EQ(to_vec(cp.local_seq_lens), + (std::vector{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{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{0, 1, 3, 6, 10, 15, 15, 15, 15, 15})); + EXPECT_EQ(to_vec(cp.local_seq_lens), + (std::vector{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{0, 0, 0, 0, 0, 0, 0, 0, 6, 15})); + EXPECT_EQ(to_vec(cp.local_seq_lens), + (std::vector{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{0, 3, 3})); + EXPECT_EQ(to_vec(r0.local_seq_lens), (std::vector{3, 0})); + EXPECT_EQ(to_vec(r0.local_kv_start_loc), (std::vector{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{0, 0, 2})); + EXPECT_EQ(to_vec(r1.local_seq_lens), (std::vector{0, 2})); + EXPECT_EQ(to_vec(r1.local_kv_start_loc), (std::vector{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{0, 166})); + EXPECT_EQ(to_vec(r0.local_seq_lens), (std::vector{166})); + EXPECT_EQ(to_vec(r0.local_kv_start_loc), (std::vector{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{0, 166})); + EXPECT_EQ(to_vec(r1.local_seq_lens), (std::vector{332})); + EXPECT_EQ(to_vec(r1.local_kv_start_loc), (std::vector{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(), 15.0f); + EXPECT_FLOAT_EQ(cp.local_cos[14][0].item(), 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{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 diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index ecc6a13d03..ab9967b30c 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -111,6 +111,10 @@ DECLARE_int32(ep_size); DECLARE_int32(cp_size); +DECLARE_bool(enable_dsa_cp); + +DECLARE_int32(dsa_cp_kv_interleave_size); + DECLARE_int64(tp_size); DECLARE_int64(sp_size); diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index d29e1e2786..fd243579a9 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -33,6 +33,20 @@ 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_int32(dsa_cp_kv_interleave_size, + 1, + "KV-cache interleave granularity for DSA-CP, aligned with " + "vllm-ascend cp_kv_cache_interleave_size. Must divide block_size. " + "Default 1 (token interleave). Reserved for M2 decode KV " + "sharding; ignored in the M1 prefill-only path."); + 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."); diff --git a/xllm/core/layers/common/dsa_metadata.h b/xllm/core/layers/common/dsa_metadata.h index 12a362ce83..069455da03 100644 --- a/xllm/core/layers/common/dsa_metadata.h +++ b/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 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; diff --git a/xllm/core/layers/common/dsa_metadata_builder.cpp b/xllm/core/layers/common/dsa_metadata_builder.cpp index 706a5bc2ad..182607006f 100644 --- a/xllm/core/layers/common/dsa_metadata_builder.cpp +++ b/xllm/core/layers/common/dsa_metadata_builder.cpp @@ -776,4 +776,129 @@ void DSAMetadataBuilder::build_positions(const ModelInputParams& params, dsa_metadata.c128_pad_positions = build_compressed_positions(128); } +DSACPMetadata DSAMetadataBuilder::build_cp_local_metadata( + const torch::Tensor& query_start_loc, + const torch::Tensor& seq_lens, + int32_t cp_size, + int32_t cp_rank, + const torch::Tensor& full_cos, + const torch::Tensor& full_sin) { + TORCH_CHECK(cp_size >= 1, "cp_size must be >= 1, got ", cp_size); + TORCH_CHECK(cp_rank >= 0 && cp_rank < cp_size, + "cp_rank must be in [0, cp_size), got cp_rank=", + cp_rank, + ", cp_size=", + cp_size); + TORCH_CHECK(query_start_loc.defined() && query_start_loc.dim() == 1, + "query_start_loc must be a defined 1-D tensor"); + TORCH_CHECK(query_start_loc.numel() >= 1, + "query_start_loc must have at least one element (leading 0)"); + TORCH_CHECK(seq_lens.defined() && seq_lens.dim() == 1, + "seq_lens must be a defined 1-D tensor"); + + // Work in int64 on CPU for exact host-side index math, then cast back to + // the caller's dtype at the end so the result matches DSAMetadata + // conventions. + const auto in_dtype = query_start_loc.scalar_type() == torch::kInt64 + ? torch::kInt64 + : torch::kInt32; + auto qsl = query_start_loc.to(torch::kCPU).to(torch::kInt64).contiguous(); + auto slens = seq_lens.to(torch::kCPU).to(torch::kInt64).contiguous(); + + const int64_t num_reqs = qsl.numel() - 1; + TORCH_CHECK(slens.numel() == num_reqs, + "seq_lens size (", + slens.numel(), + ") must equal num_reqs (", + num_reqs, + ") implied by query_start_loc"); + + const int64_t num_input_tokens = + num_reqs > 0 ? qsl[num_reqs].item() : 0; + + // Split the flattened token stream evenly across CP ranks. Padding keeps + // every rank's local slice the same length, matching vllm-ascend. + const int64_t num_tokens_pad = + ((num_input_tokens + cp_size - 1) / cp_size) * cp_size; + const int64_t tokens_per_rank = num_tokens_pad / cp_size; + const int64_t local_start = static_cast(cp_rank) * tokens_per_rank; + const int64_t local_end = local_start + tokens_per_rank; + + auto int_cpu = + torch::TensorOptions().dtype(torch::kInt64).device(torch::kCPU); + + DSACPMetadata cp; + cp.local_start = local_start; + cp.local_end = local_end; + cp.tokens_per_rank = tokens_per_rank; + cp.num_tokens_pad = num_tokens_pad; + + if (num_reqs <= 0) { + cp.local_query_start_loc = torch::zeros({1}, int_cpu).to(in_dtype); + cp.local_seq_lens = torch::zeros({0}, int_cpu).to(in_dtype); + cp.local_kv_start_loc = torch::zeros({1}, int_cpu).to(in_dtype); + } else { + // Intersect each request's global token interval with this rank's local + // token interval, then build the per-rank query_start_loc from lengths. + auto q_start = qsl.slice(/*dim=*/0, /*start=*/0, /*end=*/num_reqs); + auto q_end = qsl.slice(/*dim=*/0, /*start=*/1, /*end=*/num_reqs + 1); + auto local_q_start = q_start.clamp(local_start, local_end); + auto local_q_end = q_end.clamp(local_start, local_end); + auto local_q_lens = local_q_end - local_q_start; // (num_reqs,) + + auto local_qsl = torch::zeros({num_reqs + 1}, int_cpu); + // NOTE: assigning to a slice temporary does NOT write into local_qsl. + // Use narrow().copy_() to write into the underlying storage. + local_qsl.narrow(/*dim=*/0, /*start=*/1, /*length=*/num_reqs) + .copy_(torch::cumsum(local_q_lens, /*dim=*/0)); + + // For requests that cross the local slice boundary, offset removes the + // tokens that live on later ranks so local_seq_lens matches local + // queries. + auto offset = q_end - local_q_end; // (num_reqs,) + auto has_local = (local_q_lens > 0).to(torch::kInt64); + auto local_slens = has_local * (slens - offset); // (num_reqs,) + + cp.local_query_start_loc = local_qsl.to(in_dtype); + cp.local_seq_lens = local_slens.to(in_dtype); + + // Cumulative form of local_seq_lens (leading 0). This is the ori_kv + // cu-seqlens: unlike local_query_start_loc it accounts for the full + // right-aligned KV extent each local query sees, so the last rank's + // entry sums to the full stream length rather than its local query count. + auto local_kv_qsl = torch::zeros({num_reqs + 1}, int_cpu); + local_kv_qsl.narrow(/*dim=*/0, /*start=*/1, /*length=*/num_reqs) + .copy_(torch::cumsum(local_slens, /*dim=*/0)); + cp.local_kv_start_loc = local_kv_qsl.to(in_dtype); + } + + // Slice RoPE tables to the local token interval when provided. Tables are + // expected to cover the PADDED global token stream (length + // num_tokens_pad). + if (full_cos.defined() && full_cos.numel() > 0) { + TORCH_CHECK(full_cos.size(0) >= local_end, + "full_cos length (", + full_cos.size(0), + ") must cover local_end (", + local_end, + "); pass a table padded to num_tokens_pad"); + cp.local_cos = + full_cos.slice(/*dim=*/0, /*start=*/local_start, /*end=*/local_end) + .contiguous(); + } + if (full_sin.defined() && full_sin.numel() > 0) { + TORCH_CHECK(full_sin.size(0) >= local_end, + "full_sin length (", + full_sin.size(0), + ") must cover local_end (", + local_end, + "); pass a table padded to num_tokens_pad"); + cp.local_sin = + full_sin.slice(/*dim=*/0, /*start=*/local_start, /*end=*/local_end) + .contiguous(); + } + + return cp; +} + } // namespace xllm::layer diff --git a/xllm/core/layers/common/dsa_metadata_builder.h b/xllm/core/layers/common/dsa_metadata_builder.h index 4053cea04d..238293749d 100644 --- a/xllm/core/layers/common/dsa_metadata_builder.h +++ b/xllm/core/layers/common/dsa_metadata_builder.h @@ -28,6 +28,7 @@ namespace layer { struct AttentionMetadata; struct DSAMetadata; +struct DSACPMetadata; // Builder class for DSAMetadata. // Builds a complete AttentionMetadata (with dsa_metadata populated) from @@ -54,6 +55,38 @@ class DSAMetadataBuilder { const torch::Tensor& dsa_c4_cos_sin = torch::Tensor(), const torch::Tensor& dsa_c128_cos_sin = torch::Tensor()); + // Build per-rank DSA-CP local token metadata (M0 of the DSA-CP plan). + // + // Splits the flattened query-token stream evenly across cp_size ranks and + // returns the CURRENT rank's (cp_rank) local view. This mirrors vllm-ascend + // _build_local_token_metadata and is intentionally a pure function over + // plain tensors so it can be unit-tested on CPU without a full + // ModelInputParams / RoPE pipeline. + // + // query_start_loc : (num_reqs+1,) int32/int64, global cumulative query + // starts with a leading 0 (i.e. + // DSAMetadata::actual_seq_lengths_query). + // query_start_loc.back() == num_input_tokens. + // seq_lens : (num_reqs,) int32/int64, per-request KV context + // length (DSAMetadata::actual_seq_lengths_kv). + // cp_size : CP group width (>= 1). + // cp_rank : this rank's index within the CP group ([0, cp_size)). + // full_cos/full_sin : optional RoPE tables over the PADDED global token + // stream (length == num_tokens_pad). When defined, the + // result's local_cos/local_sin are sliced to + // [local_start, local_end). Pass undefined tensors to + // skip (e.g. in pure metadata unit tests). + // + // When cp_size == 1 the returned metadata is the identity view (local == + // global) so callers can use a single code path. + static DSACPMetadata build_cp_local_metadata( + const torch::Tensor& query_start_loc, + const torch::Tensor& seq_lens, + int32_t cp_size, + int32_t cp_rank, + const torch::Tensor& full_cos = torch::Tensor(), + const torch::Tensor& full_sin = torch::Tensor()); + private: // Build DSA-specific fields into dsa_metadata. static void build_dsa_fields( diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp index 17ebd0e194..7b5a32eef9 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp @@ -21,13 +21,17 @@ limitations under the License. #include #include #include +#include #include #include +#include "framework/parallel_state/parallel_state.h" #include "kernels/ops_api.h" +#include "layers/common/dsa_metadata_builder.h" #include "xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h" DECLARE_bool(enable_chunked_prefill); +DECLARE_bool(enable_dsa_cp); namespace xllm { namespace layer { namespace { @@ -346,10 +350,21 @@ Dsv4PreprocessOutputs run_dsv4_preprocess_fallback( double eps, const torch::Tensor& q_rms_gamma, const torch::Tensor& cos, - const torch::Tensor& sin) { + const torch::Tensor& sin, + // DSA-CP: when q_hidden_states is defined, the q path is computed from it + // (the local token slice) with q_cos/q_sin, while the kv path keeps using + // the full hidden_states with cos/sin. Undefined => single-input behavior. + const torch::Tensor& q_hidden_states = torch::Tensor(), + const torch::Tensor& q_cos = torch::Tensor(), + const torch::Tensor& q_sin = torch::Tensor()) { Dsv4PreprocessOutputs outputs; - auto q_down = q_a_proj->forward(hidden_states); + const bool split_qkv = q_hidden_states.defined(); + const torch::Tensor& q_input = split_qkv ? q_hidden_states : hidden_states; + const torch::Tensor& q_rope_cos = split_qkv ? q_cos : cos; + const torch::Tensor& q_rope_sin = split_qkv ? q_sin : sin; + + auto q_down = q_a_proj->forward(q_input); if (q_b_proj->uses_w8a8_dynamic_quant()) { xllm::kernel::RmsNormDynamicQuantParams rms_quant_params; rms_quant_params.input = q_down; @@ -365,7 +380,7 @@ Dsv4PreprocessOutputs run_dsv4_preprocess_fallback( q_b_proj->w8a8_dynamic_weight_scale(), q_pertoken_scale, q_b_proj->bias(), - hidden_states.scalar_type()) + q_input.scalar_type()) .view({-1, n_local_heads, head_dim}); } else { outputs.qr = std::get<0>(q_layernorm->forward(q_down)); @@ -385,7 +400,8 @@ Dsv4PreprocessOutputs run_dsv4_preprocess_fallback( outputs.kv = std::get<0>(kv_layernorm->forward(kv_down)); outputs.kv = outputs.kv.view({-1, 1, qk_head_dim}); - apply_partial_rope(outputs.q, nope_head_dim, rope_head_dim, cos, sin); + apply_partial_rope( + outputs.q, nope_head_dim, rope_head_dim, q_rope_cos, q_rope_sin); apply_partial_rope(outputs.kv, nope_head_dim, rope_head_dim, cos, sin); return outputs; @@ -498,6 +514,30 @@ DSAttentionImpl::DSAttentionImpl(const ModelArgs& args, const int64_t tp_size = parallel_args.tp_group_->world_size(); tp_rank_ = parallel_args.tp_group_->rank(); tp_size_ = tp_size; + + // DSA-CP Model 1 (CP == TP): a single switch (FLAGS_enable_dsa_cp) turns on + // context parallelism, and the CP geometry is DERIVED from TP - the TP group + // IS the CP group (matching vllm-ascend dsa_cp.py, which uses get_tp_group() + // as the CP group). The launcher keeps the config cp_size == 1 (world = + // dp * tp), so cp_rank()/cp_group_ from ParallelArgs stay at their trivial + // defaults; we override them here with the TP geometry when the flag is set. + // This avoids the old cp_size(2) vs cp_group world_size(4) mismatch, where + // the token split used cp_size but the gather ran over every TP member. + cp_enabled_ = FLAGS_enable_dsa_cp && tp_size > 1 && + parallel_args.tp_group_ != nullptr; + if (cp_enabled_) { + cp_size_ = tp_size; + cp_rank_ = tp_rank_; + cp_group_ = parallel_args.tp_group_; + } else { + cp_size_ = 1; + cp_rank_ = 0; + cp_group_ = nullptr; + } + // Full-head layout is required whenever DSA-CP is enabled: prefill runs + // attention over ALL heads on the local token slice, then all_to_all_4D + // re-shards heads to TP. Decode / non-CP prefill slice back to n_local_heads. + cp_full_head_ = cp_enabled_; int64_t hidden_size = args.hidden_size(); int64_t num_heads = args.n_heads(); @@ -508,9 +548,13 @@ DSAttentionImpl::DSAttentionImpl(const ModelArgs& args, n_local_heads_ = num_heads / tp_size; n_local_groups_ = o_groups_ / tp_size; + // attn_sink is indexed by attention head. Under DSA-CP the local attention + // runs over full heads, so the sink must cover all heads; otherwise it is + // sharded to this rank's TP head slice (see load_state_dict). + const int64_t sink_heads = cp_full_head_ ? num_heads : n_local_heads_; attn_sink_ = register_parameter( "attn_sink", - torch::zeros({n_local_heads_}, options.dtype(torch::kFloat32)), + torch::zeros({sink_heads}, options.dtype(torch::kFloat32)), /*requires_grad=*/false); q_a_proj_ = register_module( @@ -521,13 +565,22 @@ DSAttentionImpl::DSAttentionImpl(const ModelArgs& args, q_layernorm_ = register_module("q_a_layernorm", RMSNorm(q_lora_rank_, eps_, options)); + // Under DSA-CP the q projection must produce FULL heads on every rank (the + // TP group is the CP token-split group, not a head-split group here), so load + // q_b replicated over a single-rank group. ColumnParallelLinear with a + // world_size==1 group loads the full (unsharded) weight + W8A8 scale through + // its own load_state_dict, so the quantized matmul path is unchanged. Falls + // back to the standard TP column-parallel weight when DSA-CP is off. + ProcessGroup* q_b_group = cp_full_head_ && parallel_args.single_rank_group_ + ? parallel_args.single_rank_group_ + : parallel_args.tp_group_; q_b_proj_ = register_module("q_b_proj", ColumnParallelLinear(q_lora_rank_, num_heads * head_dim_, false, false, quant_args, - parallel_args.tp_group_, + q_b_group, options)); kv_proj_ = register_module( @@ -622,6 +675,54 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, compress_metadata; auto cos = attn_metadata.cos; auto sin = attn_metadata.sin; + + // ===== DSA-CP (prefill-only) local token setup ===== + // Model-side partitioning: worker keeps the FULL token stream, so + // hidden_states here spans all query tokens. Under CP the q path is computed + // only for this rank's local token slice (with RoPE sliced to that range), + // while kv / compressor keep using the full stream. The local attention + // output is gathered back to full-token order before o_proj (see below). + const bool cp_prefill = is_prefill || is_chunked_prefill; + const bool cp_active = cp_enabled_ && cp_prefill; + DSACPMetadata cp_meta; + torch::Tensor q_hidden_local; + torch::Tensor q_cos_local; + torch::Tensor q_sin_local; + int64_t cp_num_full_tokens = + hidden_states.defined() ? hidden_states.size(0) : 0; + if (cp_active) { + cp_meta = DSAMetadataBuilder::build_cp_local_metadata( + attn_metadata.actual_seq_lengths_query, + attn_metadata.actual_seq_lengths_kv, + static_cast(cp_size_), + static_cast(cp_rank_)); + // Clamp the local interval to the real (unpadded) token count so the last + // rank does not slice past the end of hidden_states. + const int64_t local_start = + std::min(cp_meta.local_start, cp_num_full_tokens); + const int64_t local_end = + std::min(cp_meta.local_end, cp_num_full_tokens); + const int64_t local_len = std::max(local_end - local_start, 0); + // Materialize a contiguous local slice: slicing dim-0 at a nonzero + // start yields a storage-offset view, which some NPU matmul/quant + // kernels mishandle (surfaces later as an async blockDim=0 launch + // failure in a downstream op such as the compressor). + q_hidden_local = + hidden_states.slice(/*dim=*/0, local_start, local_end).contiguous(); + if (cos.defined() && cos.size(0) >= local_end) { + q_cos_local = cos.slice(/*dim=*/0, local_start, local_end).contiguous(); + } + if (sin.defined() && sin.size(0) >= local_end) { + q_sin_local = sin.slice(/*dim=*/0, local_start, local_end).contiguous(); + } + (void)local_len; + } + + // Under DSA-CP q_b is full-head, so q must be viewed with all heads; the + // active CP prefill path keeps full heads for attention, while decode / the + // (rare) non-CP prefill under a full-head build slice back to n_local_heads + // below so the sparse-attn + o_proj math matches the TP baseline. + const int64_t q_view_heads = cp_full_head_ ? num_heads_ : n_local_heads_; Dsv4PreprocessOutputs preprocess_outputs = run_dsv4_preprocess_fallback(q_a_proj_, q_layernorm_, @@ -629,7 +730,7 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, kv_proj_, kv_layernorm_, hidden_states, - n_local_heads_, + q_view_heads, head_dim_, qk_head_dim_, nope_head_dim_, @@ -637,11 +738,48 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, eps_, q_rms_gamma_, cos, - sin); + sin, + q_hidden_local, + q_cos_local, + q_sin_local); auto qr = preprocess_outputs.qr; auto qr_pertoken_scale = preprocess_outputs.qr_pertoken_scale; auto q = preprocess_outputs.q; auto kv = preprocess_outputs.kv; + // Full-head build but CP transpose is NOT active this forward (decode, or a + // non-prefill call): slice q to this rank's TP head shard so attention and + // o_proj run exactly like the non-CP baseline. The sink is sliced at use. + if (cp_full_head_ && !cp_active) { + const int64_t shard_start = tp_rank_ * n_local_heads_; + q = q.slice(/*dim=*/1, + /*start=*/shard_start, + /*end=*/shard_start + n_local_heads_) + .contiguous(); + } + + // DSA-CP local cu-seqlens for the indexer / sparse attention, moved to the + // same device as the full metadata tensors. local_query_start_loc sums to + // this rank's local query-token count; local_seq_lens is the per-request KV + // length visible to the local queries (vllm-ascend semantics). + torch::Tensor cp_local_q_cu; + torch::Tensor cp_local_kv; + torch::Tensor cp_local_ori_kv_cu; + if (cp_active) { + const auto cp_dev = attn_metadata.actual_seq_lengths_query.defined() + ? attn_metadata.actual_seq_lengths_query.device() + : torch::Device(torch::kCPU); + cp_local_q_cu = cp_meta.local_query_start_loc.to(cp_dev); + cp_local_kv = cp_meta.local_seq_lens.to(cp_dev); + // ori_kv cu-seqlens: cumsum of local_seq_lens (right-aligned KV extent), + // NOT the local query cu-seqlens. On the last rank the local queries see + // the full KV stream, so reusing local_query_start_loc here would offset + // the causal window and make the queries read the wrong (earlier) KV. + if (cp_meta.local_kv_start_loc.defined()) { + cp_local_ori_kv_cu = cp_meta.local_kv_start_loc.to(cp_dev); + } else { + cp_local_ori_kv_cu = cp_local_q_cu; + } + } // 4) resolve per-layer cache mapping const int64_t compress_ratio_i = static_cast(compress_ratio_); @@ -715,7 +853,12 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, torch::Tensor ori_block_table_for_attn = ori_block_table; const std::string ori_kv_layout = "PA_ND"; const bool use_prefill_attn = is_prefill || is_chunked_prefill; - const bool use_temporary_prefill_kv = is_prefill && !is_chunked_prefill; + // DSA-CP writes the FULL KV into the persistent SWA cache and reads it back + // through the block table, because the temporary PA_ND build assumes q and + // kv share the same (full) token range - which is false under CP (local q, + // full kv). So the temporary-prefill-KV shortcut is disabled when CP is on. + const bool use_temporary_prefill_kv = + is_prefill && !is_chunked_prefill && !cp_active; if (use_temporary_prefill_kv) { const int64_t block_size = ori_kv.defined() && ori_kv.dim() > 1 ? ori_kv.size(1) : 128; @@ -735,19 +878,34 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, // Token compressed cache is PA_ND for both prefill and decode. torch::Tensor cmp_kv_for_attn; auto cmp_kv = kv_cache.get_k_cache(); + torch::Tensor compress_cos; + torch::Tensor compress_sin; + if (compress_ratio_i == 4) { + compress_cos = attn_metadata.c4_cos; + compress_sin = attn_metadata.c4_sin; + } else if (compress_ratio_i == 128) { + compress_cos = attn_metadata.c128_cos; + compress_sin = attn_metadata.c128_sin; + } + // Skip the compressor when there are no compressed positions to write + // (empty / fake warmup shard, common under cp_size>1 where the worker + // injects a 1-token placeholder without rebuilding the compressed- + // position RoPE). cmp_s == rope_sin.size(0); a 0-row rope makes + // aclnnCompressor launch with blockDim=0 and abort. + // blockDim=0 in aclnnCompressor comes from an empty work grid: either no + // compressed positions (rope_sin rows == 0) OR an empty query batch + // (actual_seq_lengths_query has < 2 entries => batch_size 0). The + // cp_size>1 warmup injects a 1-token placeholder without rebuilding the + // DSA metadata, so hidden_states has 1 row while the cu-seqlens/rope + // stay empty. Guard on BOTH so the placeholder shard skips the kernel. + const auto& cmp_q_cu = attn_metadata.actual_seq_lengths_query; + const bool has_compress_positions = + compress_sin.defined() && compress_sin.size(0) > 0 && + hidden_states.defined() && hidden_states.size(0) > 0 && + cmp_q_cu.defined() && cmp_q_cu.numel() > 1; if (compress_ratio_i > 1 && compressor_ && cmp_kv.defined() && cmp_slot.defined() && compressor_kv_state.defined() && - compressor_score_state.defined()) { - torch::Tensor compress_cos; - torch::Tensor compress_sin; - if (compress_ratio_i == 4) { - compress_cos = attn_metadata.c4_cos; - compress_sin = attn_metadata.c4_sin; - } else if (compress_ratio_i == 128) { - compress_cos = attn_metadata.c128_cos; - compress_sin = attn_metadata.c128_sin; - } - + compressor_score_state.defined() && has_compress_positions) { std::tuple compressor_states{ compressor_kv_state, compressor_score_state}; std::tuple compressor_block_tables{ @@ -766,7 +924,10 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, } torch::Tensor compress_topk_idxs; - if (compress_ratio_i == 4 && cmp_kv.defined()) { + // Same empty/fake-shard guard as the compressor: the indexer selects + // top-k over compressed positions, so with no compressed positions it + // has nothing to do and select_qli would launch a 0-block kernel. + if (compress_ratio_i == 4 && cmp_kv.defined() && has_compress_positions) { auto index_cache = kv_cache.get_index_cache(); std::optional indexer_cache_scale = kv_cache.get_indexer_cache_scale(); @@ -786,23 +947,44 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, CHECK(qli_metadata.defined()) << "DSAttention requires precomputed " "qli_metadata for compress_ratio==4."; auto qli_metadata_opt = std::optional(qli_metadata); - compress_topk_idxs = - indexer_->select_qli(hidden_states, - qr, - qr_pertoken_scale, - index_cache, - &indexer_cache_scale_tensor, - indexer_metadata, - cos, - sin, - attn_metadata.c4_cos, - attn_metadata.c4_sin, - attn_metadata.actual_seq_lengths_query, - attn_metadata.actual_seq_lengths_kv, - qli_metadata_opt, - use_prefill_attn, - &indexer_states, - &indexer_block_tables); + // DSA-CP: the indexer selects top-k for the LOCAL query tokens, so it + // consumes the local q hidden slice, local q RoPE and local cu-seqlens. + // NOTE: index-key cache completeness across CP ranks (each rank needs + // the full index keys when kv_split_size==1) must be validated on + // multi-card NPU; that is the M1 acceptance gate (see design doc 9.1). + const torch::Tensor& idx_hidden = + cp_active ? q_hidden_local : hidden_states; + const torch::Tensor& idx_cos = cp_active ? q_cos_local : cos; + const torch::Tensor& idx_sin = cp_active ? q_sin_local : sin; + const torch::Tensor& idx_q_cu = + cp_active ? cp_local_q_cu : attn_metadata.actual_seq_lengths_query; + const torch::Tensor& idx_kv = + cp_active ? cp_local_kv : attn_metadata.actual_seq_lengths_kv; + compress_topk_idxs = indexer_->select_qli( + idx_hidden, + qr, + qr_pertoken_scale, + index_cache, + &indexer_cache_scale_tensor, + indexer_metadata, + idx_cos, + idx_sin, + attn_metadata.c4_cos, + attn_metadata.c4_sin, + idx_q_cu, + idx_kv, + qli_metadata_opt, + use_prefill_attn, + &indexer_states, + &indexer_block_tables, + // DSA-CP: index-cache update (compress_kv) + // must use FULL hidden + FULL cu-seqlens so + // its full c4 RoPE row count matches; the + // local idx_hidden above only drives the + // query/weights/top-k path. + cp_active ? std::optional(hidden_states) : std::nullopt, + cp_active ? as_optional(attn_metadata.actual_seq_lengths_query) + : std::nullopt); CHECK(compress_topk_idxs.defined()) << "DSAttention indexer returned undefined topk indices for " "compress_ratio==4."; @@ -824,11 +1006,29 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, std::optional cu_seqlens_ori_kv_for_attn = std::nullopt; if (use_prefill_attn) { - // Prefill-style sparse metadata uses query cu-seqlens for ori_kv. + // Prefill-style sparse metadata uses query cu-seqlens for ori_kv in the + // non-CP case (q and ori_kv share the full token range). Under CP the + // ori_kv is the full, replicated stream while q is the local slice, so + // this must be the right-aligned KV cu-seqlens (cumsum of local_seq_lens), + // matching seqused_kv below. Using the local query cu-seqlens here offsets + // the last rank's causal window and reads the wrong KV slice. cu_seqlens_ori_kv_for_attn = - as_optional(attn_metadata.actual_seq_lengths_query); + cp_active ? as_optional(cp_local_ori_kv_cu) + : as_optional(attn_metadata.actual_seq_lengths_query); } + const bool pass_cmp_sparse = (compress_ratio_i == 4); + // Sink is per attention head. It is full-head when CP prefill runs attention + // over all heads; otherwise (decode / non-CP under a full-head build) slice + // it to this rank's TP head shard to match the sliced q above. + torch::Tensor attn_sink_for_attn = attn_sink_; + if (attn_sink_loaded_ && cp_full_head_ && !cp_active) { + const int64_t shard_start = tp_rank_ * n_local_heads_; + attn_sink_for_attn = attn_sink_.slice(/*dim=*/0, + /*start=*/shard_start, + /*end=*/shard_start + n_local_heads_) + .contiguous(); + } auto [attn_output, output_lse] = xllm::kernel::npu::sparse_attn_sharedkv( /*q=*/q, /*ori_kv=*/as_optional(ori_kv_for_attn), @@ -836,16 +1036,21 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, : std::nullopt, /*ori_sparse_indices=*/std::nullopt, /*cmp_sparse_indices=*/ - compress_ratio_i == 4 ? as_optional(compress_topk_idxs) : std::nullopt, + pass_cmp_sparse ? as_optional(compress_topk_idxs) : std::nullopt, /*ori_block_table=*/as_optional(ori_block_table_for_attn), /*cmp_block_table=*/ compress_ratio_i > 1 ? as_optional(cmp_block_table) : std::nullopt, - /*cu_seqlens_q=*/as_optional(attn_metadata.actual_seq_lengths_query), + /*cu_seqlens_q=*/ + cp_active ? as_optional(cp_local_q_cu) + : as_optional(attn_metadata.actual_seq_lengths_query), /*cu_seqlens_ori_kv=*/cu_seqlens_ori_kv_for_attn, /*cu_seqlens_cmp_kv=*/std::nullopt, /*seqused_q=*/std::nullopt, - /*seqused_kv=*/as_optional(attn_metadata.actual_seq_lengths_kv), - /*sinks=*/attn_sink_loaded_ ? as_optional(attn_sink_) : std::nullopt, + /*seqused_kv=*/ + cp_active ? as_optional(cp_local_kv) + : as_optional(attn_metadata.actual_seq_lengths_kv), + /*sinks=*/attn_sink_loaded_ ? as_optional(attn_sink_for_attn) + : std::nullopt, /*metadata=*/sparse_metadata, /*softmax_scale=*/softmax_scale_, /*cmp_ratio=*/compress_ratio_i, @@ -862,16 +1067,51 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, scatter_by_slot(ori_kv, ori_slot, kv); } - // 9) output RoPE + projection - auto o = attn_output.view({-1, n_local_heads_, head_dim_}); + // 9) output RoPE + head/token transpose + projection + // Under CP prefill the attention output holds ALL heads for this rank's LOCAL + // token slice ([T_local, num_heads, head_dim]); the inverse RoPE uses the + // local cos/sin (same token range as q). In the non-CP / decode path it holds + // this rank's TP head shard for the full token stream ([T, n_local_heads, + // head_dim]). + const int64_t attn_heads = cp_active ? num_heads_ : n_local_heads_; + auto o = attn_output.view({-1, attn_heads, head_dim_}); + const torch::Tensor& out_cos = cp_active ? q_cos_local : cos; + const torch::Tensor& out_sin = cp_active ? q_sin_local : sin; apply_partial_rope( - o, nope_head_dim_, rope_head_dim_, cos, sin, /*inverse=*/true); + o, nope_head_dim_, rope_head_dim_, out_cos, out_sin, /*inverse=*/true); + + // DSA-CP head/token transpose (vllm-ascend _restore_tp_head_layout). Turn + // "local tokens, full heads" into "full tokens, TP-sharded heads" via a + // single all_to_all over the CP(=TP) group, so the standard TP o_proj below + // runs unchanged. The transpose needs a uniform per-rank token count, so pad + // the local slice up to tokens_per_rank first, then slice the reassembled + // stream back to the real token count. Head shard r maps to TP rank r, which + // matches the o_a_proj / o_b_proj weight sharding. + if (cp_active && cp_group_ != nullptr && cp_size_ > 1) { + const int64_t tpr = cp_meta.tokens_per_rank; + const int64_t local_len = o.size(0); + if (tpr > local_len) { + o = torch::constant_pad_nd( + o, /*pad=*/{0, 0, 0, 0, 0, tpr - local_len}, /*value=*/0); + } + auto a2a = parallel_state::all_to_all_4D(o.unsqueeze(0), + /*scatter_idx=*/2, + /*gather_idx=*/1, + /*async_ops=*/false, + cp_group_); + o = a2a().squeeze(0); // [tpr*cp_size, n_local_heads, head_dim] + if (o.size(0) > cp_num_full_tokens) { + o = o.slice(/*dim=*/0, /*start=*/0, /*end=*/cp_num_full_tokens); + } + o = o.contiguous(); + } const int64_t num_tokens = o.size(0); auto o_group = o.view({num_tokens, n_local_groups_, -1}); auto wo_a = o_a_proj_->weight().view({n_local_groups_, o_lora_rank_, -1}); auto o_low_rank = torch::einsum("tgd,grd->tgr", {o_group, wo_a}); auto output = o_b_proj_->forward(o_low_rank.reshape({num_tokens, -1})); + std::optional final_lse = std::nullopt; (void)output_lse; @@ -893,8 +1133,12 @@ void DSAttentionImpl::load_state_dict(const StateDict& state_dict) { attn_sink = state_dict.get_tensor("attn_sink.weight"); } if (attn_sink.defined()) { + // Under DSA-CP (cp_full_head_) the local attention runs over full heads, so + // keep the full-head sink; otherwise shard it to this rank's TP head slice. + const int64_t expected_sink_heads = + cp_full_head_ ? num_heads_ : n_local_heads_; if (attn_sink.dim() == 1 && attn_sink.size(0) == num_heads_ && - tp_size_ > 1) { + tp_size_ > 1 && !cp_full_head_) { CHECK_EQ(num_heads_ % tp_size_, 0) << "attn_sink full-head tensor size is not divisible by tp_size."; const int64_t shard_size = num_heads_ / tp_size_; @@ -904,9 +1148,9 @@ void DSAttentionImpl::load_state_dict(const StateDict& state_dict) { /*end=*/shard_start + shard_size); } - CHECK(attn_sink.dim() == 1 && attn_sink.size(0) == n_local_heads_) - << "attn_sink shape mismatch, expected [" << n_local_heads_ << "], got " - << attn_sink.sizes(); + CHECK(attn_sink.dim() == 1 && attn_sink.size(0) == expected_sink_heads) + << "attn_sink shape mismatch, expected [" << expected_sink_heads + << "], got " << attn_sink.sizes(); torch::NoGradGuard no_grad; attn_sink_.copy_(attn_sink.to(attn_sink_.device()).to(attn_sink_.dtype())); diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.h b/xllm/core/layers/npu_torch/deepseek_sparse_attention.h index df3f059378..98ece5ef8c 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.h +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.h @@ -95,6 +95,23 @@ class DSAttentionImpl : public torch::nn::Module { int64_t n_local_groups_; int64_t tp_rank_ = 0; int64_t tp_size_ = 1; + + // ===== DSA-CP (context-parallel) members ===== + // cp_size_ / cp_rank_ describe this rank's position in the CP group; when + // cp_size_ > 1 and cp_enabled_ is set the prefill path splits query tokens + // across cp_group_ (vllm-ascend dsa_cp.py style). cp_enabled_ additionally + // requires FLAGS_enable_dsa_cp at construction time. + int64_t cp_size_ = 1; + int64_t cp_rank_ = 0; + bool cp_enabled_ = false; + ProcessGroup* cp_group_ = nullptr; + // DSA-CP Model 1 (CP == TP): when enabled, q_b and attn_sink are loaded + // FULL-head (replicated over the CP/TP group) so each rank can run attention + // over ALL heads on its local token slice, then an all_to_all_4D transpose + // re-shards heads to the standard TP layout before o_proj. Decode / non-CP + // prefill slice the full-head q and sink back to this rank's TP head shard, + // keeping those paths byte-identical to the non-CP baseline. + bool cp_full_head_ = false; int64_t index_n_heads_ = 0; int64_t index_head_dim_ = 0; int64_t index_topk_ = 0; diff --git a/xllm/core/layers/npu_torch/deepseek_v4_indexer.cpp b/xllm/core/layers/npu_torch/deepseek_v4_indexer.cpp index f1923496a6..95a56a5a3e 100644 --- a/xllm/core/layers/npu_torch/deepseek_v4_indexer.cpp +++ b/xllm/core/layers/npu_torch/deepseek_v4_indexer.cpp @@ -406,7 +406,9 @@ torch::Tensor DeepseekV4IndexerImpl::select_qli( const std::optional& qli_metadata, bool with_prefill, std::tuple* compressor_states, - std::tuple* compressor_block_tables) { + std::tuple* compressor_block_tables, + const std::optional& compress_x, + const std::optional& compress_cu_seqlens) { CHECK(index_cache.defined()) << "DeepseekV4Indexer::select_qli: index_cache is undefined"; @@ -421,11 +423,20 @@ torch::Tensor DeepseekV4IndexerImpl::select_qli( auto hadamard = get_hadamard_matrix(attn_metadata, hadamard_matrix_); q = rotate_activation_with_hadamard(q, hadamard, hadamard_scale_); - auto kv = compress_kv(x, + // DSA-CP: the index-key cache is written for the FULL token stream + // (replicated per CP rank, like vllm-ascend _update_indexer_cache), + // so the compressor sees full hidden + full cu-seqlens; its full c4 + // RoPE row count then matches. query/weights/top-k below stay local. + const torch::Tensor& compress_hidden = + compress_x.has_value() ? compress_x.value() : x; + const std::optional& compress_q_cu = + compress_cu_seqlens.has_value() ? compress_cu_seqlens + : actual_seq_lengths_query; + auto kv = compress_kv(compress_hidden, attn_metadata, compressed_cos, compressed_sin, - actual_seq_lengths_query, + compress_q_cu, compressor_states, compressor_block_tables); if (kv.numel() > 0) { diff --git a/xllm/core/layers/npu_torch/deepseek_v4_indexer.h b/xllm/core/layers/npu_torch/deepseek_v4_indexer.h index 3407f0e275..ad1f066bbc 100644 --- a/xllm/core/layers/npu_torch/deepseek_v4_indexer.h +++ b/xllm/core/layers/npu_torch/deepseek_v4_indexer.h @@ -65,7 +65,13 @@ class DeepseekV4IndexerImpl : public torch::nn::Module { bool with_prefill = false, std::tuple* compressor_states = nullptr, std::tuple* compressor_block_tables = - nullptr); + nullptr, + // DSA-CP: when defined, the index-cache-update path (compress_kv) + // uses these FULL-token tensors while query/weights/top-k keep using + // the LOCAL `x`/`actual_seq_lengths_query`. Mirrors vllm-ascend + // _update_indexer_cache(full) vs _indexer_select_topk(local). + const std::optional& compress_x = std::nullopt, + const std::optional& compress_cu_seqlens = std::nullopt); torch::Tensor select_qli( const torch::Tensor& x, diff --git a/xllm/core/platform/platform.h b/xllm/core/platform/platform.h index c2a29ea7e3..63438a4452 100644 --- a/xllm/core/platform/platform.h +++ b/xllm/core/platform/platform.h @@ -56,6 +56,14 @@ class Platform final { // after MLU moves CP input preparation into WorkerImpl. static constexpr bool uses_model_cp_partition() { return is_mlu(); } + // DSA-CP (DeepSeek-V4) uses model-side token partitioning: the DSAttention + // layer keeps the local token slice for q while AllGathering the full token + // stream for KV/compressor inside the layer (vllm-ascend dsa_cp.py style). + // Whether this path is actually taken at runtime is additionally gated by + // FLAGS_enable_dsa_cp and cp_size > 1 at the call sites; this hook only + // reports platform capability. NPU-only for now. + static constexpr bool supports_dsa_cp_model_partition() { return is_npu(); } + static constexpr bool is_ilu() { #if defined(USE_ILU) return true; diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index 24d3df340d..7ddf219540 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -51,7 +51,16 @@ namespace xllm { namespace { bool uses_worker_cp_partition(const runtime::Options& options) { - return options.cp_size() > 1 && !Platform::uses_model_cp_partition(); + // DSA-CP (DeepSeek-V4) partitions tokens MODEL-side: the DSAttention + // layer keeps full tokens and gathers its local attention output back + // to full-token layout internally, so the worker must NOT take the + // worker-side CP path (LmHead all-gather + 3-arg logits(out_hidden), + // which V4 does not implement). Mirrors the same gate in + // WorkerImpl::prepare_work_before_execute. + const bool dsa_cp_model_side = FLAGS_enable_dsa_cp && options.cp_size() > 1 && + Platform::supports_dsa_cp_model_partition(); + return options.cp_size() > 1 && !Platform::uses_model_cp_partition() && + !dsa_cp_model_side; } void wait_input_ready_events(const ForwardInput& input, const Stream& stream) { diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 6390495833..26a9b60a0c 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -822,8 +822,18 @@ void WorkerImpl::prepare_work_before_execute_on_stream( // already-partitioned device tensors; see ForwardInput::cp_partitioned. // Prefill-side CP (partition + ATB cp tensors) applies to PREFILL, // CHUNKED_PREFILL, and MIXED. `no_decode()` wrongly excludes MIXED. + // DSA-CP (DeepSeek-V4) uses MODEL-side token partitioning: the worker keeps + // the FULL token stream and the DSAttention layer slices its local query + // token range internally (vllm-ascend dsa_cp.py style), so the worker-side + // head-tail CP partition + ATB cp-tensor prep must be skipped here. Gated by + // FLAGS_enable_dsa_cp so the legacy worker-side CP path is unchanged when the + // feature is off. NPU-only (Platform::supports_dsa_cp_model_partition()). + const bool dsa_cp_model_side = FLAGS_enable_dsa_cp && + parallel_args_.cp_size() > 1 && + Platform::supports_dsa_cp_model_partition(); const bool needs_cp_prefill_side = parallel_args_.cp_size() > 1 && !Platform::uses_model_cp_partition() && + !dsa_cp_model_side && !input.input_params.meta.batch_forward_type.is_decode(); const bool needs_cp_partition = needs_cp_prefill_side && !input.cp_partitioned; diff --git a/xllm/models/llm/deepseek_v4.h b/xllm/models/llm/deepseek_v4.h index 1a0304b20f..c1c1ade121 100644 --- a/xllm/models/llm/deepseek_v4.h +++ b/xllm/models/llm/deepseek_v4.h @@ -32,6 +32,7 @@ limitations under the License. #include #include +#include "core/common/global_flags.h" #include "core/framework/config/execution_config.h" #include "core/framework/config/kv_cache_config.h" #include "core/framework/state_dict/utils.h" @@ -431,6 +432,20 @@ class DeepseekV4ModelImpl << ", world_size=" << parallel_args.world_size() << ", dp_size=" << parallel_args.dp_size(); tp_num_heads_ = num_heads_ / dp_local_tp_size_; + // DSA-CP Model 1 (CP == TP): a single switch (FLAGS_enable_dsa_cp) enables + // context parallelism and the CP geometry is DERIVED from TP - the TP group + // is the CP token-split group. The launcher keeps config cp_size == 1, so + // the CP width equals the DP-local TP width and the CP rank is the TP rank. + // This must match DSAttentionImpl's constructor so the precomputed sparse + // metadata (local cu_seqlens) is built for the same local token slice the + // attention runs on. When the flag is off, fall back to config cp_size. + if (FLAGS_enable_dsa_cp && dp_local_tp_size_ > 1) { + cp_size_ = dp_local_tp_size_; + cp_rank_ = parallel_args.rank() % dp_local_tp_size_; + } else { + cp_size_ = std::max(parallel_args.cp_size(), 1); + cp_rank_ = std::max(parallel_args.cp_rank(), 0); + } window_size_ = model_args.window_size(); index_n_heads_ = model_args.index_n_heads(); index_head_dim_ = model_args.index_head_dim(); @@ -820,8 +835,17 @@ class DeepseekV4ModelImpl } #endif } + torch::Tensor pre_hc_head_hidden_states; + if (model_args_.num_speculative_tokens() > 0) { + pre_hc_head_hidden_states = h; + } h = hc_head(h); auto [hidden_states, residual_out] = norm_(h, std::nullopt); + if (pre_hc_head_hidden_states.defined()) { + ModelOutput out(hidden_states, residual_out); + out.aux_hidden_states = pre_hc_head_hidden_states.flatten(1); + return out; + } return ModelOutput(hidden_states, residual_out); } @@ -1367,17 +1391,57 @@ class DeepseekV4ModelImpl is_prefill ? as_optional_tensor(dsa.actual_seq_lengths_query) : empty_int32_opt; + // DSA-CP: every sparse-attn / indexer-topk call downstream runs on + // this rank's LOCAL query slice, so the precomputed tiling metadata + // must be built from LOCAL cu-seqlens too. Using the full geometry + // makes the kernel index q past the local row count -> MTE DDR + // out-of-range. Compressor / indexer-cache keep the full geometry + // (they are not driven by this precomputed metadata). + torch::Tensor md_cu_q = dsa.actual_seq_lengths_query; + torch::Tensor md_seqused_kv = dsa.actual_seq_lengths_kv; + // ori_kv cu-seqlens default to the query cu-seqlens: in full (non-CP) + // prefill q and ori_kv span the same token range, so they coincide. Under + // CP they diverge (local q, full right-aligned kv) and md_cu_ori_kv must + // follow the KV extent, not the local query count. + torch::Tensor md_cu_ori_kv = dsa.actual_seq_lengths_query; + int64_t md_batch = batch_size; + int64_t md_max_q = max_seqlen_q; + const bool dsa_cp_md = FLAGS_enable_dsa_cp && cp_size_ > 1 && is_prefill; + if (dsa_cp_md) { + auto cpm = layer::DSAMetadataBuilder::build_cp_local_metadata( + dsa.actual_seq_lengths_query, + dsa.actual_seq_lengths_kv, + static_cast(cp_size_), + static_cast(cp_rank_)); + if (cpm.local_query_start_loc.defined() && cpm.local_seq_lens.defined()) { + md_cu_q = cpm.local_query_start_loc.to(metadata_device); + md_seqused_kv = cpm.local_seq_lens.to(metadata_device); + md_batch = std::max(cpm.local_seq_lens.size(0), 1); + md_max_q = std::max(cpm.tokens_per_rank, 1); + // Right-aligned ori_kv cu-seqlens (cumsum of local_seq_lens). The last + // rank's local queries see the full KV stream, so this must not reuse + // the local query cu-seqlens (that reads the wrong, earlier KV slice). + if (cpm.local_kv_start_loc.defined()) { + md_cu_ori_kv = cpm.local_kv_start_loc.to(metadata_device); + } + } + } + auto md_cu_q_opt = as_optional_tensor(md_cu_q); + auto md_seqused_kv_opt = as_optional_tensor(md_seqused_kv); + auto md_cu_ori_kv_opt = + is_prefill ? as_optional_tensor(md_cu_ori_kv) : empty_int32_opt; + xllm::kernel::SparseAttnSharedkvMetadataParams c1_params; c1_params.num_heads_q = tp_num_heads_; c1_params.num_heads_kv = 1; c1_params.head_dim = head_dim_; - c1_params.cu_seqlens_q = as_optional_tensor(dsa.actual_seq_lengths_query); - c1_params.cu_seqlens_ori_kv = cu_seqlens_ori_kv_opt; + c1_params.cu_seqlens_q = md_cu_q_opt; + c1_params.cu_seqlens_ori_kv = md_cu_ori_kv_opt; c1_params.cu_seqlens_cmp_kv = empty_int32_opt; c1_params.seqused_q = empty_int32_opt; - c1_params.seqused_kv = as_optional_tensor(dsa.actual_seq_lengths_kv); - c1_params.batch_size = batch_size; - c1_params.max_seqlen_q = max_seqlen_q; + c1_params.seqused_kv = md_seqused_kv_opt; + c1_params.batch_size = md_batch; + c1_params.max_seqlen_q = md_max_q; c1_params.max_seqlen_kv = max_seqlen_kv; c1_params.ori_topk = 0; c1_params.cmp_topk = 0; @@ -1396,13 +1460,13 @@ class DeepseekV4ModelImpl c4_params.num_heads_q = tp_num_heads_; c4_params.num_heads_kv = 1; c4_params.head_dim = head_dim_; - c4_params.cu_seqlens_q = as_optional_tensor(dsa.actual_seq_lengths_query); - c4_params.cu_seqlens_ori_kv = cu_seqlens_ori_kv_opt; + c4_params.cu_seqlens_q = md_cu_q_opt; + c4_params.cu_seqlens_ori_kv = md_cu_ori_kv_opt; c4_params.cu_seqlens_cmp_kv = empty_int32_opt; c4_params.seqused_q = empty_int32_opt; - c4_params.seqused_kv = as_optional_tensor(dsa.actual_seq_lengths_kv); - c4_params.batch_size = batch_size; - c4_params.max_seqlen_q = max_seqlen_q; + c4_params.seqused_kv = md_seqused_kv_opt; + c4_params.batch_size = md_batch; + c4_params.max_seqlen_q = md_max_q; c4_params.max_seqlen_kv = max_seqlen_kv; c4_params.ori_topk = 0; c4_params.cmp_topk = sparse_topk; @@ -1421,13 +1485,13 @@ class DeepseekV4ModelImpl c128_params.num_heads_q = tp_num_heads_; c128_params.num_heads_kv = 1; c128_params.head_dim = head_dim_; - c128_params.cu_seqlens_q = as_optional_tensor(dsa.actual_seq_lengths_query); - c128_params.cu_seqlens_ori_kv = cu_seqlens_ori_kv_opt; + c128_params.cu_seqlens_q = md_cu_q_opt; + c128_params.cu_seqlens_ori_kv = md_cu_ori_kv_opt; c128_params.cu_seqlens_cmp_kv = empty_int32_opt; c128_params.seqused_q = empty_int32_opt; - c128_params.seqused_kv = as_optional_tensor(dsa.actual_seq_lengths_kv); - c128_params.batch_size = batch_size; - c128_params.max_seqlen_q = max_seqlen_q; + c128_params.seqused_kv = md_seqused_kv_opt; + c128_params.batch_size = md_batch; + c128_params.max_seqlen_q = md_max_q; c128_params.max_seqlen_kv = max_seqlen_kv; c128_params.ori_topk = 0; c128_params.cmp_topk = 0; @@ -1469,6 +1533,35 @@ class DeepseekV4ModelImpl return; } + // DSA-CP: the indexer runs top-k selection on this rank's LOCAL query + // slice (see deepseek_sparse_attention.cpp: idx_q_cu/idx_kv are the + // local cu-seqlens). The precomputed qli_metadata must describe the same + // LOCAL query geometry, otherwise the indexer collapses its selection to + // position 0 (topk_max==0) and attention degenerates. KV stays FULL + // (CP splits Q only), so max_seqlen_k keeps the full value. This mirrors + // the c1/c4/c128 md_* localization above. + int64_t qli_local_max_q = 0; + if (dsa_cp_md) { + auto qli_cpm = layer::DSAMetadataBuilder::build_cp_local_metadata( + dsa.actual_seq_lengths_query, + dsa.actual_seq_lengths_kv, + static_cast(cp_size_), + static_cast(cp_rank_)); + if (qli_cpm.local_query_start_loc.defined() && + qli_cpm.local_query_start_loc.size(0) > 1 && + qli_cpm.local_seq_lens.defined()) { + // Drop the leading 0 to match the query_lens convention above. + query_lens = qli_cpm.local_query_start_loc + .slice(/*dim=*/0, + /*start=*/1, + /*end=*/qli_cpm.local_query_start_loc.size(0)) + .clone() + .to(metadata_device); + key_lens = qli_cpm.local_seq_lens.to(metadata_device); + qli_local_max_q = std::max(qli_cpm.tokens_per_rank, 1); + } + } + const int64_t global_index_num_heads = std::max(index_n_heads_ > 0 ? index_n_heads_ : num_heads_, 1); CHECK_EQ(global_index_num_heads % dp_local_tp_size_, 0) @@ -1479,8 +1572,11 @@ class DeepseekV4ModelImpl const int64_t index_head_dim = std::max(index_head_dim_ > 0 ? index_head_dim_ : head_dim_, 1); const int64_t qli_max_seqlen_q = - std::max(params.meta.q_max_seq_len, - vector_max_or_zero(params.attention.host.q_seq_lens)); + qli_local_max_q > 0 + ? qli_local_max_q + : std::max( + params.meta.q_max_seq_len, + vector_max_or_zero(params.attention.host.q_seq_lens)); const int64_t qli_max_seqlen_k = std::max( params.meta.kv_max_seq_len, vector_max_or_zero(params.attention.host.kv_seq_lens)); @@ -1529,6 +1625,8 @@ class DeepseekV4ModelImpl int64_t num_heads_ = 0; int64_t tp_num_heads_ = 0; + int64_t cp_size_ = 1; + int64_t cp_rank_ = 0; int64_t dp_local_tp_size_ = 1; int64_t head_dim_ = 0; int64_t window_size_ = 128; diff --git a/xllm/models/llm/deepseek_v4_mtp.h b/xllm/models/llm/deepseek_v4_mtp.h index ea1c3c8d33..fb98538a9d 100644 --- a/xllm/models/llm/deepseek_v4_mtp.h +++ b/xllm/models/llm/deepseek_v4_mtp.h @@ -59,7 +59,8 @@ class DeepseekV4MultiTokenPredictorLayerImpl layer::AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params, - torch::Tensor tokens) { + torch::Tensor tokens, + torch::Tensor* aux_hidden_states) { ModelInputParams modified_input_params = input_params; modified_input_params.embedding.input_embedding = previous_hidden_states; std::optional residual; @@ -70,7 +71,8 @@ class DeepseekV4MultiTokenPredictorLayerImpl attn_metadata, kv_cache, modified_input_params, - tokens); + tokens, + aux_hidden_states); } }; TORCH_MODULE(DeepseekV4MultiTokenPredictorLayer); @@ -96,6 +98,20 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { std::max(parallel_args.world_size() / std::max(parallel_args.dp_size(), 1), 1); + // DSA-CP Model 1 (CP == TP): derive the CP geometry from TP, matching the + // main model (deepseek_v4.h) and the shared DSAttention layer. The MTP + // draft reuses the same DSAttention, which under FLAGS_enable_dsa_cp slices + // q to the LOCAL token slice and runs full-head attention; its precomputed + // sparse metadata must therefore be built for the SAME local slice. The + // launcher keeps config cp_size == 1, so without this the draft would build + // FULL-geometry metadata while the attention runs local -> corrupt drafts. + if (FLAGS_enable_dsa_cp && dp_local_tp_size_ > 1) { + cp_size_ = dp_local_tp_size_; + cp_rank_ = parallel_args.rank() % dp_local_tp_size_; + } else { + cp_size_ = std::max(parallel_args.cp_size(), 1); + cp_rank_ = std::max(parallel_args.cp_rank(), 0); + } CHECK_EQ(num_heads_ % dp_local_tp_size_, 0) << "[DeepseekV4Mtp] n_heads must be divisible by local tp " "size. n_heads=" @@ -292,6 +308,7 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { static_cast(mtp_layers_.size())) << "deepseek_v4_mtp requires kv_caches size >= mtp layer count"; + torch::Tensor aux_hidden_states; for (size_t i = 0; i < mtp_layers_.size(); ++i) { const int32_t layer_id = static_cast(i); prepare_for_layer(*modified_input_params.attn_metadata, layer_id); @@ -301,7 +318,8 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { *modified_input_params.attn_metadata, kv_caches[i], modified_input_params, - tokens); + tokens, + &aux_hidden_states); #if defined(USE_NPU) if (modified_input_params.parallel.layer_synchronizer != nullptr && !modified_input_params.parallel.layer_synchronizer->record_event( @@ -312,7 +330,7 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { } auto [output, _] = final_norm_(hidden_states, std::nullopt); - return ModelOutput(output, std::nullopt); + return ModelOutput(output, torch::Tensor(), aux_hidden_states.flatten(1)); } bool requires_graph_forward_metadata() { return true; } @@ -745,17 +763,59 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { is_prefill ? as_optional_tensor(dsa.actual_seq_lengths_query) : empty_int32_opt; + // DSA-CP: every sparse-attn / indexer-topk call downstream runs on + // this rank's LOCAL query slice, so the precomputed tiling metadata + // must be built from LOCAL cu-seqlens too. Using the full geometry + // makes the kernel index q past the local row count -> MTE DDR + // out-of-range. Compressor / indexer-cache keep the full geometry + // (they are not driven by this precomputed metadata). This mirrors the + // main model (deepseek_v4.h); the MTP draft shares the same DSAttention + // layer, which slices q LOCAL under CP, so its metadata must match. + torch::Tensor md_cu_q = dsa.actual_seq_lengths_query; + torch::Tensor md_seqused_kv = dsa.actual_seq_lengths_kv; + // ori_kv cu-seqlens default to the query cu-seqlens: in full (non-CP) + // prefill q and ori_kv span the same token range, so they coincide. Under + // CP they diverge (local q, full right-aligned kv) and md_cu_ori_kv must + // follow the KV extent, not the local query count. + torch::Tensor md_cu_ori_kv = dsa.actual_seq_lengths_query; + int64_t md_batch = batch_size; + int64_t md_max_q = max_seqlen_q; + const bool dsa_cp_md = FLAGS_enable_dsa_cp && cp_size_ > 1 && is_prefill; + if (dsa_cp_md) { + auto cpm = layer::DSAMetadataBuilder::build_cp_local_metadata( + dsa.actual_seq_lengths_query, + dsa.actual_seq_lengths_kv, + static_cast(cp_size_), + static_cast(cp_rank_)); + if (cpm.local_query_start_loc.defined() && cpm.local_seq_lens.defined()) { + md_cu_q = cpm.local_query_start_loc.to(metadata_device); + md_seqused_kv = cpm.local_seq_lens.to(metadata_device); + md_batch = std::max(cpm.local_seq_lens.size(0), 1); + md_max_q = std::max(cpm.tokens_per_rank, 1); + // Right-aligned ori_kv cu-seqlens (cumsum of local_seq_lens). The last + // rank's local queries see the full KV stream, so this must not reuse + // the local query cu-seqlens (that reads the wrong, earlier KV slice). + if (cpm.local_kv_start_loc.defined()) { + md_cu_ori_kv = cpm.local_kv_start_loc.to(metadata_device); + } + } + } + auto md_cu_q_opt = as_optional_tensor(md_cu_q); + auto md_seqused_kv_opt = as_optional_tensor(md_seqused_kv); + auto md_cu_ori_kv_opt = + is_prefill ? as_optional_tensor(md_cu_ori_kv) : empty_int32_opt; + xllm::kernel::SparseAttnSharedkvMetadataParams c1_params; c1_params.num_heads_q = tp_num_heads_; c1_params.num_heads_kv = 1; c1_params.head_dim = head_dim_; - c1_params.cu_seqlens_q = as_optional_tensor(dsa.actual_seq_lengths_query); - c1_params.cu_seqlens_ori_kv = cu_seqlens_ori_kv_opt; + c1_params.cu_seqlens_q = md_cu_q_opt; + c1_params.cu_seqlens_ori_kv = md_cu_ori_kv_opt; c1_params.cu_seqlens_cmp_kv = empty_int32_opt; c1_params.seqused_q = empty_int32_opt; - c1_params.seqused_kv = as_optional_tensor(dsa.actual_seq_lengths_kv); - c1_params.batch_size = batch_size; - c1_params.max_seqlen_q = max_seqlen_q; + c1_params.seqused_kv = md_seqused_kv_opt; + c1_params.batch_size = md_batch; + c1_params.max_seqlen_q = md_max_q; c1_params.max_seqlen_kv = max_seqlen_kv; c1_params.ori_topk = 0; c1_params.cmp_topk = 0; @@ -774,13 +834,13 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { c4_params.num_heads_q = tp_num_heads_; c4_params.num_heads_kv = 1; c4_params.head_dim = head_dim_; - c4_params.cu_seqlens_q = as_optional_tensor(dsa.actual_seq_lengths_query); - c4_params.cu_seqlens_ori_kv = cu_seqlens_ori_kv_opt; + c4_params.cu_seqlens_q = md_cu_q_opt; + c4_params.cu_seqlens_ori_kv = md_cu_ori_kv_opt; c4_params.cu_seqlens_cmp_kv = empty_int32_opt; c4_params.seqused_q = empty_int32_opt; - c4_params.seqused_kv = as_optional_tensor(dsa.actual_seq_lengths_kv); - c4_params.batch_size = batch_size; - c4_params.max_seqlen_q = max_seqlen_q; + c4_params.seqused_kv = md_seqused_kv_opt; + c4_params.batch_size = md_batch; + c4_params.max_seqlen_q = md_max_q; c4_params.max_seqlen_kv = max_seqlen_kv; c4_params.ori_topk = 0; c4_params.cmp_topk = sparse_topk; @@ -799,13 +859,13 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { c128_params.num_heads_q = tp_num_heads_; c128_params.num_heads_kv = 1; c128_params.head_dim = head_dim_; - c128_params.cu_seqlens_q = as_optional_tensor(dsa.actual_seq_lengths_query); - c128_params.cu_seqlens_ori_kv = cu_seqlens_ori_kv_opt; + c128_params.cu_seqlens_q = md_cu_q_opt; + c128_params.cu_seqlens_ori_kv = md_cu_ori_kv_opt; c128_params.cu_seqlens_cmp_kv = empty_int32_opt; c128_params.seqused_q = empty_int32_opt; - c128_params.seqused_kv = as_optional_tensor(dsa.actual_seq_lengths_kv); - c128_params.batch_size = batch_size; - c128_params.max_seqlen_q = max_seqlen_q; + c128_params.seqused_kv = md_seqused_kv_opt; + c128_params.batch_size = md_batch; + c128_params.max_seqlen_q = md_max_q; c128_params.max_seqlen_kv = max_seqlen_kv; c128_params.ori_topk = 0; c128_params.cmp_topk = 0; @@ -943,6 +1003,8 @@ class DeepseekV4MtpModelImpl final : public torch::nn::Module { int64_t num_heads_ = 0; int64_t tp_num_heads_ = 0; + int64_t cp_size_ = 1; + int64_t cp_rank_ = 0; int64_t dp_local_tp_size_ = 1; int64_t head_dim_ = 0; int64_t window_size_ = 128; From 7add03f5ec32b0fdce1d4df36cd8d8638c237641 Mon Sep 17 00:00:00 2001 From: Mxyzptlk <140030863+LMxyzptlk@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:32:32 +0800 Subject: [PATCH 2/8] Fix formatting in DsaCpMetadataTest for clarity --- tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp b/tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp index f2dbfa598d..3f9457899d 100644 --- a/tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp +++ b/tests/core/layers/npu_torch/dsa_cp_metadata_tests.cpp @@ -151,8 +151,8 @@ TEST(DsaCpMetadataTest, PaddingAndBoundaryCross) { // 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 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); From 33c61eb68c494f05d817007f0002d66dfb0605c2 Mon Sep 17 00:00:00 2001 From: Mxyzptlk <140030863+LMxyzptlk@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:35:14 +0800 Subject: [PATCH 3/8] Refactor cp_enabled assignment and improve formatting --- .../npu_torch/deepseek_sparse_attention.cpp | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp index 7b5a32eef9..30017039c4 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp @@ -523,8 +523,8 @@ DSAttentionImpl::DSAttentionImpl(const ModelArgs& args, // defaults; we override them here with the TP geometry when the flag is set. // This avoids the old cp_size(2) vs cp_group world_size(4) mismatch, where // the token split used cp_size but the gather ran over every TP member. - cp_enabled_ = FLAGS_enable_dsa_cp && tp_size > 1 && - parallel_args.tp_group_ != nullptr; + cp_enabled_ = + FLAGS_enable_dsa_cp && tp_size > 1 && parallel_args.tp_group_ != nullptr; if (cp_enabled_) { cp_size_ = tp_size; cp_rank_ = tp_rank_; @@ -1024,9 +1024,10 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, torch::Tensor attn_sink_for_attn = attn_sink_; if (attn_sink_loaded_ && cp_full_head_ && !cp_active) { const int64_t shard_start = tp_rank_ * n_local_heads_; - attn_sink_for_attn = attn_sink_.slice(/*dim=*/0, - /*start=*/shard_start, - /*end=*/shard_start + n_local_heads_) + attn_sink_for_attn = attn_sink_ + .slice(/*dim=*/0, + /*start=*/shard_start, + /*end=*/shard_start + n_local_heads_) .contiguous(); } auto [attn_output, output_lse] = xllm::kernel::npu::sparse_attn_sharedkv( @@ -1049,8 +1050,8 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, /*seqused_kv=*/ cp_active ? as_optional(cp_local_kv) : as_optional(attn_metadata.actual_seq_lengths_kv), - /*sinks=*/attn_sink_loaded_ ? as_optional(attn_sink_for_attn) - : std::nullopt, + /*sinks=*/ + attn_sink_loaded_ ? as_optional(attn_sink_for_attn) : std::nullopt, /*metadata=*/sparse_metadata, /*softmax_scale=*/softmax_scale_, /*cmp_ratio=*/compress_ratio_i, From 88e4fc981663bffb7261e8e6cf118ad65877800b Mon Sep 17 00:00:00 2001 From: Mxyzptlk <140030863+LMxyzptlk@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:45:53 +0800 Subject: [PATCH 4/8] Fix formatting in deepseek_sparse_attention.cpp --- xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp index 30017039c4..0f9115714e 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp @@ -1051,7 +1051,7 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, cp_active ? as_optional(cp_local_kv) : as_optional(attn_metadata.actual_seq_lengths_kv), /*sinks=*/ - attn_sink_loaded_ ? as_optional(attn_sink_for_attn) : std::nullopt, + attn_sink_loaded_ ? as_optional(attn_sink_for_attn) : std::nullopt, /*metadata=*/sparse_metadata, /*softmax_scale=*/softmax_scale_, /*cmp_ratio=*/compress_ratio_i, From 1cd583b4e024cf4838c7fe23e65eb00c3eed709a Mon Sep 17 00:00:00 2001 From: unknown Date: Wed, 22 Jul 2026 13:03:56 +0800 Subject: [PATCH 5/8] refactor: add enable_dsa_cp to ParallelConfig --- xllm/core/framework/config/parallel_config.cpp | 4 ++++ xllm/core/framework/config/parallel_config.h | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index fd243579a9..bf41fe3eda 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -90,6 +90,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); @@ -105,6 +106,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); @@ -122,6 +124,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( diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..dee7768f86 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/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; From 1f529797428bd2d2f016f1abd3fb021172e02e5f Mon Sep 17 00:00:00 2001 From: unknown Date: Wed, 22 Jul 2026 13:18:00 +0800 Subject: [PATCH 6/8] refactor: drop unused dsa_cp_kv_interleave_size flag Remove the dsa_cp_kv_interleave_size flag (declaration + definition). --- xllm/core/common/global_flags.h | 2 -- xllm/core/framework/config/parallel_config.cpp | 7 ------- 2 files changed, 9 deletions(-) diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index ab9967b30c..e625282036 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -113,8 +113,6 @@ DECLARE_int32(cp_size); DECLARE_bool(enable_dsa_cp); -DECLARE_int32(dsa_cp_kv_interleave_size); - DECLARE_int64(tp_size); DECLARE_int64(sp_size); diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index bf41fe3eda..6c5df95967 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -40,13 +40,6 @@ DEFINE_bool(enable_dsa_cp, "(DSA). When false, DSA attention runs its original non-CP path " "even if cp_size > 1."); -DEFINE_int32(dsa_cp_kv_interleave_size, - 1, - "KV-cache interleave granularity for DSA-CP, aligned with " - "vllm-ascend cp_kv_cache_interleave_size. Must divide block_size. " - "Default 1 (token interleave). Reserved for M2 decode KV " - "sharding; ignored in the M1 prefill-only path."); - 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."); From 613d839bd36f8371a4f0600c77bba757b4143100 Mon Sep 17 00:00:00 2001 From: unknown Date: Wed, 22 Jul 2026 15:52:07 +0800 Subject: [PATCH 7/8] fix: guard DSA-CP empty local slice for chunked prefill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When chunked prefill creates a tail chunk smaller than cp_size (e.g. 1 token, cp_size=2), some ranks get 0 local query tokens. The quantized matmul in q_a_proj (aclnnQuantMatmulV4) crashes on 0-row input tensors. Three guard points when cp_local_empty (cp_active && local_len==0): 1. run_dsv4_preprocess_fallback: skip q projections, still compute kv for cache writes — every rank must write its own KV replica. 2. indexer select_qli: skip when local is empty (M1 kv_split_size=1: all ranks derive identical index-cache entries from the same full hidden_states, so rank 0 already wrote the correct data). 3. sparse_attn_sharedkv: produce zero-row attn_output directly, avoiding a 0-element kernel launch. Empty rank contributes zeros through all_to_all (padded to tokens_per_rank), which are discarded after the post-all_to_all slice back to the real token count. o_proj runs normally on the gathered full-token output. kv_split_size=1 (replicated KV). Skipping the index-cache write for empty ranks is safe only while every rank computes routing from identical full hidden_states. --- .../npu_torch/deepseek_sparse_attention.cpp | 39 +++++++++++++++++-- 1 file changed, 36 insertions(+), 3 deletions(-) diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp index 0f9115714e..197eb9dc11 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp @@ -364,6 +364,22 @@ Dsv4PreprocessOutputs run_dsv4_preprocess_fallback( const torch::Tensor& q_rope_cos = split_qkv ? q_cos : cos; const torch::Tensor& q_rope_sin = split_qkv ? q_sin : sin; + // DSA-CP empty rank: when this rank has no local query tokens, skip the + // quantized q projections (aclnnQuantMatmulV4 does not support 0-row inputs). + // The kv path runs normally on the full hidden stream so cache writes still + // happen on every rank. M1 only: index-cache write (indexer) is skipped in + // forward() when cp_local_empty since all ranks write identical data to their + // replicated index-cache copies (kv_split_size=1). + if (split_qkv && q_input.size(0) == 0) { + auto kv_down = kv_proj->forward(hidden_states); + outputs.kv = std::get<0>(kv_layernorm->forward(kv_down)); + outputs.kv = outputs.kv.view({-1, 1, qk_head_dim}); + apply_partial_rope(outputs.kv, nope_head_dim, rope_head_dim, cos, sin); + outputs.q = torch::zeros({0, n_local_heads, head_dim}, hidden_states.options()); + outputs.qr = torch::Tensor(); // undefined, unused downstream + return outputs; + } + auto q_down = q_a_proj->forward(q_input); if (q_b_proj->uses_w8a8_dynamic_quant()) { xllm::kernel::RmsNormDynamicQuantParams rms_quant_params; @@ -690,6 +706,7 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, torch::Tensor q_sin_local; int64_t cp_num_full_tokens = hidden_states.defined() ? hidden_states.size(0) : 0; + bool cp_local_empty = false; if (cp_active) { cp_meta = DSAMetadataBuilder::build_cp_local_metadata( attn_metadata.actual_seq_lengths_query, @@ -715,7 +732,7 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, if (sin.defined() && sin.size(0) >= local_end) { q_sin_local = sin.slice(/*dim=*/0, local_start, local_end).contiguous(); } - (void)local_len; + cp_local_empty = (local_len == 0); } // Under DSA-CP q_b is full-head, so q must be viewed with all heads; the @@ -927,7 +944,12 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, // Same empty/fake-shard guard as the compressor: the indexer selects // top-k over compressed positions, so with no compressed positions it // has nothing to do and select_qli would launch a 0-block kernel. - if (compress_ratio_i == 4 && cmp_kv.defined() && has_compress_positions) { + // DSA-CP empty rank: skip indexer when this rank has no local query tokens + // (idx_hidden=empty would crash). M1 only (kv_split_size=1): index-cache + // writes are identical across ranks since hidden_states is full on every rank, + // so skipping rank N's write is safe when rank 0 already wrote the same data. + if (compress_ratio_i == 4 && cmp_kv.defined() && has_compress_positions && + !cp_local_empty) { auto index_cache = kv_cache.get_index_cache(); std::optional indexer_cache_scale = kv_cache.get_indexer_cache_scale(); @@ -1004,6 +1026,16 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, << "DSAttention requires precomputed sparse metadata for compress_ratio=" << compress_ratio_i; + // DSA-CP empty rank: skip sparse attention when this rank has no local query + // tokens (q=empty would launch a 0-batch kernel). Produce a zero attention + // output to participate in the all_to_all transpose with padding. + torch::Tensor attn_output; + std::optional output_lse = std::nullopt; + if (cp_local_empty) { + const int64_t attn_heads = cp_active ? num_heads_ : n_local_heads_; + attn_output = torch::zeros({0, attn_heads, head_dim_}, hidden_states.options()); + } else { + // Normal path: run sparse attention std::optional cu_seqlens_ori_kv_for_attn = std::nullopt; if (use_prefill_attn) { // Prefill-style sparse metadata uses query cu-seqlens for ori_kv in the @@ -1030,7 +1062,7 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, /*end=*/shard_start + n_local_heads_) .contiguous(); } - auto [attn_output, output_lse] = xllm::kernel::npu::sparse_attn_sharedkv( + std::tie(attn_output, output_lse) = xllm::kernel::npu::sparse_attn_sharedkv( /*q=*/q, /*ori_kv=*/as_optional(ori_kv_for_attn), /*cmp_kv=*/compress_ratio_i > 1 ? as_optional(cmp_kv_for_attn) @@ -1062,6 +1094,7 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, /*layout_q=*/"TND", /*layout_kv=*/ori_kv_layout, /*return_softmax_lse=*/false); + } // end if !cp_local_empty // 8) Deferred cache write for full prefill. if (use_temporary_prefill_kv) { From 8bd457dd42b56ba4b6c06febb29ac084ca318362 Mon Sep 17 00:00:00 2001 From: unknown Date: Wed, 22 Jul 2026 16:40:40 +0800 Subject: [PATCH 8/8] style: apply clang-format to DSA-CP empty-slice guard --- .../npu_torch/deepseek_sparse_attention.cpp | 128 +++++++++--------- 1 file changed, 66 insertions(+), 62 deletions(-) diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp index 197eb9dc11..685f2c789a 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp @@ -375,7 +375,8 @@ Dsv4PreprocessOutputs run_dsv4_preprocess_fallback( outputs.kv = std::get<0>(kv_layernorm->forward(kv_down)); outputs.kv = outputs.kv.view({-1, 1, qk_head_dim}); apply_partial_rope(outputs.kv, nope_head_dim, rope_head_dim, cos, sin); - outputs.q = torch::zeros({0, n_local_heads, head_dim}, hidden_states.options()); + outputs.q = + torch::zeros({0, n_local_heads, head_dim}, hidden_states.options()); outputs.qr = torch::Tensor(); // undefined, unused downstream return outputs; } @@ -946,8 +947,9 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, // has nothing to do and select_qli would launch a 0-block kernel. // DSA-CP empty rank: skip indexer when this rank has no local query tokens // (idx_hidden=empty would crash). M1 only (kv_split_size=1): index-cache - // writes are identical across ranks since hidden_states is full on every rank, - // so skipping rank N's write is safe when rank 0 already wrote the same data. + // writes are identical across ranks since hidden_states is full on every + // rank, so skipping rank N's write is safe when rank 0 already wrote the same + // data. if (compress_ratio_i == 4 && cmp_kv.defined() && has_compress_positions && !cp_local_empty) { auto index_cache = kv_cache.get_index_cache(); @@ -1033,67 +1035,69 @@ DSAttentionImpl::forward(const DSAMetadata& attn_metadata, std::optional output_lse = std::nullopt; if (cp_local_empty) { const int64_t attn_heads = cp_active ? num_heads_ : n_local_heads_; - attn_output = torch::zeros({0, attn_heads, head_dim_}, hidden_states.options()); + attn_output = + torch::zeros({0, attn_heads, head_dim_}, hidden_states.options()); } else { - // Normal path: run sparse attention - std::optional cu_seqlens_ori_kv_for_attn = std::nullopt; - if (use_prefill_attn) { - // Prefill-style sparse metadata uses query cu-seqlens for ori_kv in the - // non-CP case (q and ori_kv share the full token range). Under CP the - // ori_kv is the full, replicated stream while q is the local slice, so - // this must be the right-aligned KV cu-seqlens (cumsum of local_seq_lens), - // matching seqused_kv below. Using the local query cu-seqlens here offsets - // the last rank's causal window and reads the wrong KV slice. - cu_seqlens_ori_kv_for_attn = - cp_active ? as_optional(cp_local_ori_kv_cu) - : as_optional(attn_metadata.actual_seq_lengths_query); - } + // Normal path: run sparse attention + std::optional cu_seqlens_ori_kv_for_attn = std::nullopt; + if (use_prefill_attn) { + // Prefill-style sparse metadata uses query cu-seqlens for ori_kv in the + // non-CP case (q and ori_kv share the full token range). Under CP the + // ori_kv is the full, replicated stream while q is the local slice, so + // this must be the right-aligned KV cu-seqlens (cumsum of + // local_seq_lens), matching seqused_kv below. Using the local query + // cu-seqlens here offsets the last rank's causal window and reads the + // wrong KV slice. + cu_seqlens_ori_kv_for_attn = + cp_active ? as_optional(cp_local_ori_kv_cu) + : as_optional(attn_metadata.actual_seq_lengths_query); + } - const bool pass_cmp_sparse = (compress_ratio_i == 4); - // Sink is per attention head. It is full-head when CP prefill runs attention - // over all heads; otherwise (decode / non-CP under a full-head build) slice - // it to this rank's TP head shard to match the sliced q above. - torch::Tensor attn_sink_for_attn = attn_sink_; - if (attn_sink_loaded_ && cp_full_head_ && !cp_active) { - const int64_t shard_start = tp_rank_ * n_local_heads_; - attn_sink_for_attn = attn_sink_ - .slice(/*dim=*/0, - /*start=*/shard_start, - /*end=*/shard_start + n_local_heads_) - .contiguous(); - } - std::tie(attn_output, output_lse) = xllm::kernel::npu::sparse_attn_sharedkv( - /*q=*/q, - /*ori_kv=*/as_optional(ori_kv_for_attn), - /*cmp_kv=*/compress_ratio_i > 1 ? as_optional(cmp_kv_for_attn) - : std::nullopt, - /*ori_sparse_indices=*/std::nullopt, - /*cmp_sparse_indices=*/ - pass_cmp_sparse ? as_optional(compress_topk_idxs) : std::nullopt, - /*ori_block_table=*/as_optional(ori_block_table_for_attn), - /*cmp_block_table=*/ - compress_ratio_i > 1 ? as_optional(cmp_block_table) : std::nullopt, - /*cu_seqlens_q=*/ - cp_active ? as_optional(cp_local_q_cu) - : as_optional(attn_metadata.actual_seq_lengths_query), - /*cu_seqlens_ori_kv=*/cu_seqlens_ori_kv_for_attn, - /*cu_seqlens_cmp_kv=*/std::nullopt, - /*seqused_q=*/std::nullopt, - /*seqused_kv=*/ - cp_active ? as_optional(cp_local_kv) - : as_optional(attn_metadata.actual_seq_lengths_kv), - /*sinks=*/ - attn_sink_loaded_ ? as_optional(attn_sink_for_attn) : std::nullopt, - /*metadata=*/sparse_metadata, - /*softmax_scale=*/softmax_scale_, - /*cmp_ratio=*/compress_ratio_i, - /*ori_mask_mode=*/4, - /*cmp_mask_mode=*/3, - /*ori_win_left=*/std::max(window_size_ - 1, 0), - /*ori_win_right=*/0, - /*layout_q=*/"TND", - /*layout_kv=*/ori_kv_layout, - /*return_softmax_lse=*/false); + const bool pass_cmp_sparse = (compress_ratio_i == 4); + // Sink is per attention head. It is full-head when CP prefill runs + // attention over all heads; otherwise (decode / non-CP under a full-head + // build) slice it to this rank's TP head shard to match the sliced q above. + torch::Tensor attn_sink_for_attn = attn_sink_; + if (attn_sink_loaded_ && cp_full_head_ && !cp_active) { + const int64_t shard_start = tp_rank_ * n_local_heads_; + attn_sink_for_attn = attn_sink_ + .slice(/*dim=*/0, + /*start=*/shard_start, + /*end=*/shard_start + n_local_heads_) + .contiguous(); + } + std::tie(attn_output, output_lse) = xllm::kernel::npu::sparse_attn_sharedkv( + /*q=*/q, + /*ori_kv=*/as_optional(ori_kv_for_attn), + /*cmp_kv=*/compress_ratio_i > 1 ? as_optional(cmp_kv_for_attn) + : std::nullopt, + /*ori_sparse_indices=*/std::nullopt, + /*cmp_sparse_indices=*/ + pass_cmp_sparse ? as_optional(compress_topk_idxs) : std::nullopt, + /*ori_block_table=*/as_optional(ori_block_table_for_attn), + /*cmp_block_table=*/ + compress_ratio_i > 1 ? as_optional(cmp_block_table) : std::nullopt, + /*cu_seqlens_q=*/ + cp_active ? as_optional(cp_local_q_cu) + : as_optional(attn_metadata.actual_seq_lengths_query), + /*cu_seqlens_ori_kv=*/cu_seqlens_ori_kv_for_attn, + /*cu_seqlens_cmp_kv=*/std::nullopt, + /*seqused_q=*/std::nullopt, + /*seqused_kv=*/ + cp_active ? as_optional(cp_local_kv) + : as_optional(attn_metadata.actual_seq_lengths_kv), + /*sinks=*/ + attn_sink_loaded_ ? as_optional(attn_sink_for_attn) : std::nullopt, + /*metadata=*/sparse_metadata, + /*softmax_scale=*/softmax_scale_, + /*cmp_ratio=*/compress_ratio_i, + /*ori_mask_mode=*/4, + /*cmp_mask_mode=*/3, + /*ori_win_left=*/std::max(window_size_ - 1, 0), + /*ori_win_right=*/0, + /*layout_q=*/"TND", + /*layout_kv=*/ori_kv_layout, + /*return_softmax_lse=*/false); } // end if !cp_local_empty // 8) Deferred cache write for full prefill.