디코드 컨텍스트 병렬화

디코드 컨텍스트 병렬화 (Decode Context Parallelism)

Decode context parallelism (DCP)은 요청의 MLA KV 캐시를 기존 텐서 병렬 그룹 내의 랭크들에 스트라이프(스트라이프)로 분산하는 기법입니다. DCP 크기가 c일 때 랭크 r은 p mod c = r인 위치 p를 소유하므로, 각 랭크는 캐시의 약 1/c만 저장하고 읽어요.

출처: 문서

본문

DCP 크기가 c일 때, 랭크 r은 p mod c = r인 위치 p를 소유하므로 각 랭크는 캐시의 약 1/c를 저장하고 읽습니다. 그러면 로컬 어텐션은 컨텍스트의 일부만 보게 됩니다. MLA 경로는 부분 출력과 log-sum-exp (LSE)를 반환하고, 둘을 하나의 packed all-to-all로 교환한 뒤 정확히 병합합니다. DCP는 TP와 attention data parallelism (DPA)을 보완하지 대체하지는 않아요. 주요 경로는 absorbed MLA decode와 static target verification입니다. dcp_size=1은 기존의 비-DCP 동작을 유지합니다.

Kimi K3에서 DCP는 MLA 레이어에만 적용됩니다. 요청 인덱스 기반 KDA 상태는 변경되지 않으므로, DCP는 긴 컨텍스트 KV 용량을 늘리지만 KDA 동시성 상한을 올리지는 않습니다.

DCP 활성화 (Enable DCP)

Setting	Behavior
`--dcp-size N`	DCP를 활성화하고 가상 용량과 페이지 크기를 `N`배 확장. 별칭: `--decode-context-parallel-size`.
`--dcp-comm-backend`	MLA 부분 출력 병합 방식을 선택: `ag_rs`, `a2a`, or `fi_a2a`. 미설정 시 Blackwell에서 DCP 그룹이 MNNVL 도메인을 공유하면 `fi_a2a`, 아니면 CUDA/ROCm에서 `a2a`.
`--dcp-replicate-q-proj`	지원되는 MLA 가중치에 대해 query all-gather를 제거. 모델별 기본값을 끄려면 `--no-dcp-replicate-q-proj` 사용.
`--enable-dp-attention`	DCP 그룹이 attention TP 안에 중첩될 때만 합성.

DCP 그룹은 하나의 attention-TP 그룹과 하나의 attention-DP 복제본 안에 들어가야 합니다:

attn_tp_size = tp_size / attn_dp_size
attn_tp_size % dcp_size == 0

예를 들어 TP=64, DP=4, DCP=16은 4개의 유효한 16-랭크 attention 복제본을 만듭니다. DCP=32는 복제본 경계를 가로지르므로 유효하지 않습니다.

현재 시작 시 `tp_size % dcp_size == 0`만 확인하며 위의 더 강한 조건은 확인하지 않습니다. 그 검사가 고쳐질 때까지, DPA+DCP 배포는 토폴로지를 고를 때 containment(포함)를 직접 강제해야 합니다.

TP 8 안에서 DCP 8로 DeepSeek-V3.1의 MLA decode:

sglang serve \
  --model-path deepseek-ai/DeepSeek-V3.1 \
  --trust-remote-code \
  --tp-size 8 \
  --dcp-size 8 \
  --host 0.0.0.0 \
  --port 30000

32 GPU에서 DPA와 DCP 8로 Kimi K3. 모델 오버라이드가 기본적으로 replicated Q를 켭니다. 통신 백엔드는 MNNVL 시스템에서 일반적으로 fi_a2a, 그 외에는 a2a로 결정됩니다:

sglang serve \
  --trust-remote-code \
  --model-path moonshotai/Kimi-K3 \
  --tp-size 32 \
  --ep-size 32 \
  --enable-dp-attention \
  --dp-size 4 \
  --dcp-size 8 \
  --host 0.0.0.0 \
  --port 30000

DCP를 DPA·expert parallelism과 결합한 대규모 프리셋은 Kimi K3 cookbook에서 확인하세요.

왜 디코드 컨텍스트 병렬화인가 (Why decode context parallelism?)

Absorbed MLA는 모든 쿼리 헤드에 하나의 잠재 LV KV 표현을 공유합니다. TP는 헤드와 가중치를 분할할 수 있지만 그 캐시는 분할할 수 없어서, 모든 attention-TP 랭크가 동일한 컨텍스트를 저장하고 읽습니다.

flowchart LR
    subgraph TP["TP only: every rank stores the full MLA context"]
        direction TB
        T0["Rank 0<br/>tokens 0 1 2 3 4 5 6 7"]
        T1["Rank 1<br/>tokens 0 1 2 3 4 5 6 7"]
        T2["Rank 2<br/>tokens 0 1 2 3 4 5 6 7"]
        T3["Rank 3<br/>tokens 0 1 2 3 4 5 6 7"]
    end

    TP -- "stripe by position" --> DCP

    subgraph DCP["DCP4: one interleaved shard per rank"]
        direction TB
        D0["Rank 0<br/>tokens 0 4"]
        D1["Rank 1<br/>tokens 1 5"]
        D2["Rank 2<br/>tokens 2 6"]
        D3["Rank 3<br/>tokens 3 7"]
    end

DCP는 그 캐시를 논리적 토큰 위치에 따라 라운드로빈으로 스트라이프하고, 모든 랭크에 공유 가상 인덱스를 유지하며, 소유자 로컬 샤드에 대해 어텐션을 실행하고, LSE로 부분 출력을 병합해 일반적인 TP-로컬 헤드 레이아웃을 만듭니다.

아키텍처 (Architecture)

가상 KV 레이아웃 (Virtual KV layout)

각 랭크가 물리 토큰 용량 C와 물리 페이지 크기 P를 가진다고 하겠습니다. DCP는 공유 가상 할당자를 노출합니다:

virtual capacity  = C * c
virtual page size = P * c
owner(v)          = v mod c
physical(v)       = floor(v / c)

모든 랭크는 동일한 요청-대-토큰 맵을 봅니다. 쓰기 시에는 자신의 위치를 유지하고 로컬 물리 풀로 압축합니다. 하나의 넓힌 가상 페이지는 각 랭크에서 하나의 물리 페이지로 매핑되므로, 시퀀스가 커져도 레이아웃은 토큰 하나 이내로 균형을 유지합니다.

MLA decode 데이터플로우 (MLA decode dataflow)

flowchart LR
    H["Hidden states"] --> Q["Full DCP-group query heads"]

    Q --> A0["Rank 0 local attention<br/>partial O0 + LSE0"]
    Q --> A1["Rank 1 local attention<br/>partial O1 + LSE1"]
    Q --> AN["Rank c-1 local attention<br/>partial O + LSE"]

    A0 --> A2A["Packed all-to-all"]
    A1 --> A2A
    AN --> A2A
    A2A --> MERGE["Exact LSE merge"]
    MERGE --> OUT["TP-local head shard"]

컨텍스트 파티션 r에 대해 커널은 로컬 정규화 출력 o_r과 lse_r을 반환합니다. 전역 결과는 lse = logsumexp_r(lse_r)이고 o = sum_r(exp(lse_r - lse) * o_r)입니다. 이는 부동 소수점 감소 순서를 제외하면 정확합니다. 병합은 어텐션 백엔드의 LSE 밑(base-e 또는 base-2)을 사용합니다.

쿼리-프로젝션 복제 (Query-projection replication)

보통 각 레이어는 absorbed 및 rotary 쿼리 컴포넌트를 all-gather합니다. --dcp-replicate-q-proj는 쿼리-프로젝션과 w_kc 가중치를 시작 시 한 번 gather하고 전체 DCP-그룹 쿼리를 로컬로 계산합니다. 레이어당 collective 하나를 줄이는 대신 가중치 메모리와 GEMM 작업이 더 필요해요. 이 경로는 unquantized BF16/FP16 가중치만 사용하며, 다른 레이어는 쿼리 all-gather로 폴백합니다.

통신 백엔드 (Communication backends)

Backend	DCP collectives per MLA layer	Notes
`ag_rs`	Query AG + LSE AG + FP32 output RS	a2a 경로가 없는 플랫폼용 폴백; `--dcp-comm-backend ag_rs`로 선택
`a2a`, gathered Q	Query AG + packed NCCL A2A	Collective 2개
`a2a`, replicated Q	One packed NCCL A2A	표준 저지연 경로
`fi_a2a`, replicated Q	One FlashInfer MNNVL A2A	Blackwell 필요, 하나의 MNNVL 도메인 내 DCP 그룹 필요 (GB200/GB300 패브릭 또는 단일 노드)

--dcp-comm-backend가 미설정이면 모든 DCP 시작이 동일하게 해석합니다: Blackwell에서 DCP 그룹이 하나의 MNNVL 도메인을 공유하면 fi_a2a, 아니면 CUDA/ROCm에서 a2a, a2a 경로가 없는 플랫폼에서만 ag_rs. 쿼리 복제는 모델에 따라 결정됩니다. Kimi K3는 a2a 계열에서 기본으로 활성화하고, decode 백엔드는 cutedsl_mla입니다.

DCP는 컨텍스트 비의존 decode collective를 추가하는 대신, 컨텍스트 의존 KV 저장·읽기를 약 c만큼 줄입니다. Extend는 그 decode 비용 모델 밖에 있어요. 캐시된 프리픽스 샤드를 gather하고, 토큰 순서를 복원하고, 새 토큰을 추가할 수 있으며, 그 작업은 프리픽스 길이에 따라 커집니다.

합성 (Compositions)

DCP × 추측 디코딩 (DCP × speculative decoding)

드래프트 KV 캐시는 복제됩니다. 모든 DCP 랭크가 전체 드래프트 컨텍스트를 저장합니다. DCP는 target MLA KV만 스트라이프하므로 target 캐시는 절약하지만 드래프트 캐시는 절약하지 않습니다.

Target verify·decode는 DCP 인식 MLA 커널 cutedsl_mla를 사용합니다. 이것은:

순환 토큰 맵(owner(p) = p mod c)에서 랭크 로컬 페이지 테이블을 만듭니다.
cp_world, cp_rank, 전역 KV 길이를 전달해 인과성(causality)이 전체 시퀀스를 계속 볼 수 있게 하면서 그 샤드에 대해 어텐션을 실행합니다.
랭크 로컬 (output, LSE)를 반환하고, 일반적인 DCP all-to-all 병합이 결합합니다.

드래프트 단계는 복제된 드래프트 캐시에 대해 같은 백엔드를 사용하므로 DCP 메타데이터를 건너뜁니다. Kimi K3 DCP는 cutedsl_mla를 자동으로 선택합니다.

DCP × PD 분리 (DCP × PD disaggregation)

밀집 prefill 페이지는 striped decode 스토리지에 직접 복사할 수 없습니다. PD 전송자는 decode 랭크의 DCP 메타데이터를 사용해 올바른 행을 선택·패킹합니다. 전송 엔진은 PD disaggregation을 참고하세요.

동일한 DCP 크기는 일치하는 DCP 랭크가 필요하며 기존 페이지 경로를 사용합니다.
DCP1 prefill → DCP decode는 MLA·hybrid-MLA 풀에 대해 토큰 재배치(token relayout)를 사용합니다.
다른 모든 DCP 크기 전환은 거부됩니다.

플랜은 정확한 토큰 수를 포함하므로 부분 최종 페이지에서 stale 행이 새지 않습니다. Mooncake와 NIXL은 같은 플랜을 사용합니다. KDA 상태는 attention-TP 매핑을 유지하고 DCP 필터링을 우회합니다.

이 합성에는 Mooncake 또는 NIXL, 일치하는 물리 페이지 크기와 KV dtype, prefill attention CP 1, decode chunk cache가 필요합니다. 여기서 decode radix cache와 HiCache는 지원되지 않습니다.

DCP × HiCache L2

HiCache는 L1·L2에 걸쳐 하나의 넓힌 논리 페이지 공간을 유지합니다. 컨트롤러는 P * c 슬롯을 보지만 각 GPU·호스트 버퍼는 자신의 P 로컬 행만 저장합니다. 캐시 계층은 HiCache를 참고하세요.

H2D·D2H 전에 두 인덱스 리스트 모두 이 랭크가 소유한 위치를 유지하고 i를 floor(i / c)로 매핑합니다. 전송은 전체 넓힌 페이지를 다루며 각 랭크가 자신의 샤드를 독립적으로 옮기고 DCP collective는 없습니다. Kimi K3의 경우 MLA KV에 적용되며, KDA/Mamba 상태는 자신의 요청 인덱스 호스트 풀을 사용합니다.

현재 범위는 MLA L1/L2입니다. L3, LMCache, HiSparse, non-MLA KV 호스트 풀, 추측 디코딩, PD decode는 이 조합에서 지원되지 않습니다.

디코드 컨텍스트 병렬화