[NPU] Support Kimi-K3 DCP with shared compact communication by qybnb · Pull Request #40825 · sgl-project/sglang

qybnb · 2026-09-23T02:34:01Z

Motivation

Enable decode context parallelism (DCP) for Kimi-K3 on Ascend NPU, including DSPARK speculative decoding. Shard target-model MLA KV cache across DCP ranks to reduce per-rank KV storage for a given logical context capacity.

Modifications

Add sharded MLA KV allocation, reads, and writes while preserving the scheduler's logical token address space. Reuse the shared DCP-aware Triton KV writer, prefix-index builder, and page-table kernel.
Adapt Ascend FIA to rank-local KV pages, sequence lengths, and causal masks for decode and DSPARK target verification. Gather sharded prefix KV when needed for prefill/extend.
Reuse Q AllGather and compact DCP AllToAll with BF16/FP16 partial outputs and FP32 LSE, followed by the shared Triton LSE merge. Handle empty KV shards safely.
Integrate local metadata and sequence-length updates with NPUGraph capture/replay, and add unit and NPU kernel tests.

Supported scope: Kimi-K3 MLA with FP16/BF16 KV cache and non-NZ FIA V1. Speculative decoding supports DSPARK with static verification; mixed-chunk execution and replicated Q projection are not supported.

Run

#!/usr/bin/env bash
set -eo pipefail

CODE_PATH=XXX
MODEL_PATH=XXX
DRAFT_MODEL_PATH=XXX

cd "$CODE_PATH"
mkdir -p logs/manual_dcp

unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY all_proxy
export NO_PROXY=XXX
export no_proxy="$NO_PROXY"

unset PYTHONPATH
unset ASCEND_CUSTOM_OPP_PATH
source /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.sh

export PYTHONUNBUFFERED=1

export SGLANG_SET_CPU_AFFINITY=1
export SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS=1
export SGLANG_NPU_USE_TRITON_PREFIX_KV_CACHE_STORE=1
export SGLANG_K3_SHARED_EXPERTS_ATTN_TP=1
export SGLANG_K3_DENSE_MLP_ATTN_TP=1
export SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1

export SGLANG_RAGGED_VERIFY_MODE=static
export SGLANG_DSPARK_FOLDED_PROPOSAL=0
export SGLANG_DSPARK_FOLDED_SAMPLING=0
export SGLANG_DSPARK_STACKED_CTX_KV=0
export SGLANG_DSPARK_EMBED_IN_GRAPH=0

export ASCEND_USE_FIA=1
export SGLANG_USE_FIA_NZ=0
export SGLANG_NPU_USE_FIAS_V2_BSND=0

export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
export HCCL_SOCKET_IFNAME=enp196s0f0
export GLOO_SOCKET_IFNAME=enp196s0f0
export HCCL_IF_BASE_PORT=18000
export HCCL_NPU_SOCKET_PORT_RANGE=18000-18100
export HCCL_BUFFSIZE=200
export DEEPEP_HCCL_BUFFSIZE=1800
export HCCL_OP_EXPANSION_MODE=AIV
export STREAMS_PER_DEVICE=32

export DEEP_NORMAL_MODE_USE_INT8_QUANT=1
export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=128
export DEEPEP_NORMAL_LONG_SEQ_ROUND=64
export DEEPEP_NORMAL_LONG_SEQ_PER_ROUND_TOKENS=512

unset DEEP_USE_MODE DEEPEP_MC2_COMM_ALG DCP_HCCL_BUFFSIZE
unset MOE_ENABLE_TOPK_NEG_ONE SGLANG_NPU_FUSED_MOE_MODE
unset ENABLE_PROFILING SGLANG_NPU_PROFILING
unset SGLANG_K3_TRACE_STATE_FILE SGLANG_K3_TRACE_HIDDEN_FILE
unset ASCEND_RT_VISIBLE_DEVICES ASCEND_LAUNCH_BLOCKING
unset SGLANG_ENABLE_CP_V2 SGLANG_ENABLE_SPEC_V2
unset SGLANG_MAMBA_SSM_DTYPE SGLANG_MAMBA_CONV_DTYPE

NODE_IPS=(XXX)
LOCAL_IPS=" $(hostname -I) "
NODE_RANK=-1

for i in "${!NODE_IPS[@]}"; do
if [[ "$LOCAL_IPS" == " ${NODE_IPS[$i]} " ]]; then
NODE_RANK="$i"
break
fi
done

if [[ "$NODE_RANK" == "-1" ]]; then
echo "ERROR: 无法识别节点，hostname -I: $LOCAL_IPS" >&2
exit 1
fi

LOG_FILE="$PWD/logs/manual_dcp/rank${NODE_RANK}_$(date -u +%Y%m%dT%H%M%SZ).log"
echo "node_rank=$NODE_RANK log=$LOG_FILE"

python3 -m sglang.launch_server
--model-loader-extra-config '{"enable_multithread_load": true}'
--dist-init-addr XXX
--nnodes 4
--node-rank "$NODE_RANK"
--model-path "$MODEL_PATH"
--tokenizer-path "$MODEL_PATH"
--trust-remote-code
--attention-backend ascend
--device npu
--quantization modelslim
--dtype bfloat16
--tp-size 64
--enable-dp-attention
--dp-size 4
--enable-dp-lm-head
--enable-shared-experts-attn-tp
--enable-dense-mlp-attn-tp
--dcp-size 2
--dcp-comm-backend a2a
--mem-fraction-static 0.72
--chunked-prefill-size 8192
--max-running-requests 64
--cuda-graph-bs-decode 1 4 16
--disable-radix-cache
--disable-custom-all-reduce
--enable-metrics
--reasoning-parser kimi_k3
--moe-a2a-backend deepep
--deepep-mode auto
--expert-balancedness-report-mode off
--speculative-algorithm DSPARK
--speculative-draft-model-path "$DRAFT_MODEL_PATH"
--speculative-draft-model-quantization unquant
--speculative-dspark-block-size 7
--speculative-attention-mode decode
--speculative-draft-attention-backend ascend
--linear-attn-verify-backend triton
--speculative-eagle-topk 1
--watchdog-timeout 9000
--host 0.0.0.0
--port 15010
2>&1 | tee "$LOG_FILE"

Accuracy Tests

gpqa accuracy 0.9445

Speed Tests and Profiling

test with bs=4 128k in 1k out no radix cache

no dcp

dcp=2

Kimi-K3 KV Cache Capacity

Configuration	Maximum logical KV capacity (tokens per DP)	Capacity relative to DCP=1
DCP=1 (no KV sharding)	379,520	1.00×
DCP=2	687,872	1.8125×

DCP=2 provides 1.81× the logical KV capacity of DCP=1, an increase of 81.25%.

CI States

Latest PR Test (Base): ✅ Run #37773169564
Latest PR Test (Extra): ❌ Run #37773169050
Latest PR Test (AMD ROCm 10): ❌ Run #37773169638

qybnb · 2026-09-30T03:44:00Z

+                (self.k_buffer[layer_id - self.start_layer], cache_k),
+                (self.v_buffer[layer_id - self.start_layer], cache_v),
+            ):
+                src = src.reshape(-1, src.shape[-1])


因为存储布局不同：

GPU：一个合并 buffer
[latent | RoPE]
→ 一次 kernel 写两部分

当前 NPU：两个独立 buffer
k_buffer：latent
v_buffer：RoPE
→ 两次 kernel，各写一部分

qybnb · 2026-09-30T03:46:59Z

+                if src.stride(-1) != 1:
+                    src = src.contiguous()
+                dst = dst.view(-1, src.shape[-1])
+                set_mla_kv_buffer_kernel[(loc.numel(), 1)](


开启 DCP 后，它对每个 token 做三件事：

判断是否属于当前 rank：loc % DCP_SIZE == DCP_RANK。

转换成本地地址：local_loc = loc // DCP_SIZE。

属于自己就写入，不属于自己就跳过；保留地址 loc=0 也跳过。

QSA's SimpleNamespace mock in test_qsa.py does not carry the new attribute, so the hard read in HybridLinearAttnBackend.__init__ crashed the base-b-kernel-unit-test-4-gpu-b200 suite (1 failed: test_qsa_glue_fetches_indexer_metadata_without_model_unwrap). Fall back to the AttentionBackend default (True), matching the neighboring getattr style.

qybnb · 2026-09-30T09:07:59Z

/rerun-failed-ci

Preserve token-interleaved K3 MLA alongside upstream page-interleaved DSA DCP. Keep K3 virtual write locations, isolate DSA metadata and graph buffers, and retain native draft pool page geometry. Add regression coverage for the merged paths without rewriting existing history.

qybnb · 2026-10-06T07:01:40Z

@@ -34,8 +34,11 @@ def _jit_mla_output_gate_module() -> Module:


 def covered(x: torch.Tensor, gate: torch.Tensor) -> bool:
+    # transfer_to_npu aliases Tensor.is_cuda to Tensor.is_npu; this JIT is CUDA-only.
    return (
        not _is_npu


防止 NPU 误调用 CUDA 专用的 MLA output gate 融合算子。

原来通过以下条件判断能否使用 CUDA 算子：

x.is_cuda and gate.is_cuda

但当前 NPU 环境加载的 transfer_to_npu 兼容层会把 Tensor.is_cuda 映射成 Tensor.is_npu。所以 NPU 张量也可能通过这个判断，然后尝试调用 nvcc 编译 CUDA，导致启动报错。
参考pr36406

qybnb · 2026-10-06T07:22:28Z

/rerun-failed-ci

qybnb · 2026-10-06T07:59:06Z

@@ -136,6 +136,7 @@ def __init__(
        parallel = get_parallel()
        self.mla_dcp_graph = (
            model_runner.model_config.attention_arch == AttentionArch.MLA
+            and getattr(model_runner.model_config, "index_head_dim", None) is None


用于排除DSA，DSA和MLA的分片方式（token/page）、算子不同

qybnb · 2026-10-06T08:30:52Z

+            indices,
+            self.dcp_size,
+            self.dcp_rank,
+            self.page_size if self.index_head_dim is not None else 1,


interleave_size,将全局逻辑地址转换为本地地址，DSA采用page，MLA采用1 token

qybnb · 2026-10-06T09:19:17Z

+            _is_npu
+            and get_parallel().dcp_enabled
+            and not model_runner.is_draft_worker
+            and getattr(model_runner.model_config, "index_head_dim", None) is not None


DSA走当前转换路径，K3走和GPU相同的公共triton

qybnb · 2026-10-06T09:36:54Z

@@ -45,7 +196,6 @@ def build_mla_dcp_local_block_tables(
        None,
        req_to_token.stride(0),
        block_tables.stride(0),


主线kernel方法删除了这个入参

qybnb · 2026-10-08T11:06:18Z

/rerun-failed-ci

qybnb · 2026-10-08T13:00:15Z



 class AscendAttnBackend(AttentionBackend):
+    dcp_use_packed_kv = False


不同主要因为下游读取 KV 的接口不同：

使用公共 packed 路径的 GPU backend：从一个打包的 MLA KV 临时 buffer 读取，所以需要 indptr 标出每个请求的索引段、
indices 指出读哪些位置。

当前 NPU backend：KV latent 和 K-RoPE 分开存放，按本地 prefix 索引分别读取、聚合，再交给 Prefill Attention。它不读取
那个公共 packed buffer，因此也用不到对应的 indptr/indices。

所以这个标记的作用是：公共代码生成前缀读取索引后，判断是否还需要继续构造 packed buffer 及其 metadata，避免 NPU 做无用
的准备。
NPU DSA：由 Ascend backend 准备本地页表、长度和稀疏索引，给稀疏 Attention 使用；跳过整个公共 planner。

NPU K3：Prefill 需要读出并聚合完整 prefix。我们复用了公共 planner 的“前缀读取索引构造”，但不需要后面的 packed
buffer，所以用 dcp_use_packed_kv=False 提前返回

qybnb · 2026-10-08T15:33:01Z

/rerun-failed-ci

qybnb · 2026-10-08T15:33:14Z

/rerun-failed-ci

qybnb added 7 commits September 16, 2026 08:33

[NPU] Support Kimi-K3 decode context parallelism

ea41a99

[NPU] Add a group-local HCCL buffer for DCP

16f8c1c

fix(ci): address NPU DCP lint failures

49d0647

refactor(npu): isolate Kimi-K3 DCP metadata

0fcf4d7

refactor(npu): use standard HCCL buffer for DCP

4cc040a

perf(npu): compact Kimi-K3 DCP state exchange

5051451

refactor(npu): reuse shared DCP metadata and graph buffers

72de60d

qybnb commented Sep 30, 2026

View reviewed changes

qybnb and others added 2 commits September 30, 2026 04:01

refactor(npu): keep MLA KV reads in native cache dtype

70fc82f

qybnb added 5 commits October 6, 2026 03:17

refactor(npu): isolate dense MLA DCP metadata from DSA

b49fc9c

fix(npu): allow K3 DCP startup and exclude CUDA output gate JIT

cbb7cfb

chore: remove NPU DCP regression test changes from PR

0224a7b

refactor(npu): keep upstream DSA metadata logic inline

ee4d5bf

qybnb commented Oct 6, 2026

View reviewed changes

qybnb and others added 9 commits October 6, 2026 09:47

refactor(npu): simplify dense MLA DCP predicates

840a35f

Merge branch 'main' into feature/kimi-k3-npu-dcp-reuse

cbc00ae

test(npu): cover page and token DCP retraction layouts

1670159

refactor(npu): clarify DCP routing and explicit MLA cache writes

3e6b38e

refactor(npu): restore existing MLA cache-write ownership

73031a7

Merge branch 'main' into feature/kimi-k3-npu-dcp-reuse

9531c10

test(npu): sync graph runner mocks with current runtime APIs

677eabe

Merge main and reconcile NPU graph test mocks

45e5b7e

Merge main and preserve NPU DCP write-index ownership

6f40e5f

refactor(dcp): restore upstream LSE merge behavior

096434d

qybnb commented Oct 8, 2026

View reviewed changes



		class AscendAttnBackend(AttentionBackend):
		dcp_use_packed_kv = False

Conversation

qybnb commented Sep 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Run

Accuracy Tests

Speed Tests and Profiling

Kimi-K3 KV Cache Capacity

CI States

Uh oh!

qybnb Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

qybnb Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

qybnb commented Sep 30, 2026

Uh oh!

qybnb Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

qybnb commented Oct 6, 2026

Uh oh!

qybnb Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

qybnb Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

qybnb Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

qybnb Oct 6, 2026

Choose a reason for hiding this comment

Uh oh!

qybnb commented Oct 8, 2026

Uh oh!

qybnb Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

qybnb commented Oct 8, 2026

Uh oh!

qybnb commented Oct 8, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

qybnb commented Sep 23, 2026 •

edited by github-actions Bot

Loading

qybnb Oct 6, 2026 •

edited

Loading

qybnb Oct 8, 2026 •

edited

Loading