[vllm] MLA DCP 디코드에서 A2A 전에 W_UV 업프로젝션을 먼저 적용해 통신량 4배 줄이기

PR 링크: vllm-project/vllm#64994 상태: Merged | 변경: +241 / -14

들어가며

vLLM의 Decode Context Parallel(DCP)은 긴 KV cache를 여러 rank에 나눠 저장하고, 각 rank가 자기 몫의 KV에 대해 attention을 계산한 뒤 LSE(log-sum-exp) 기반으로 partial output을 merge합니다. --dcp-comm-backend a2a를 쓰면 이 merge는 NCCL all-to-all로 head들을 재분배한 뒤 수행됩니다.

문제는 MLA(Multi-head Latent Attention)의 구조입니다. MLA 디코드에서 attention output은 아직 latent 공간(Kimi-K3 기준 kv_lora_rank L=512)에 있고, 이후 W_UV로 V 차원(V=128)으로 up-projection됩니다. 기존 순서는 다음과 같았습니다.

1. 각 rank가 latent partial output(L=512) + LSE를 all-to-all
2. LSE merge
3. W_UV 적용 (512 → 128)

즉 통신에는 head당 512 + LSE 2개 = 514개의 bf16 값이 실립니다. 이 PR은 순서를 바꿔 먼저 W_UV를 모든 DCP-gathered head에 적용한 뒤 all-to-all을 수행합니다. 그러면 head당 128 + 2 = 130개만 전송하면 됩니다. 핵심 근거는 간단합니다: LSE merge는 head별로 partial output들의 가중합(선형 연산)이므로, 선형 projection W_UV와 순서를 바꿔도 수학적으로 동일하고 rounding만 달라집니다.

$$\sum_r w_r (o_r W_{UV}) = \left(\sum_r w_r o_r\right) W_{UV}$$

코드 분석

1. mla_attention.py — decode 경로 분기

먼저 레이어에 gathered W_UV를 담을 필드가 추가됩니다.

self.W_UK_T_dcp_qrep: torch.Tensor | None = None
self.W_UV_dcp: torch.Tensor | None = None

process_weights_after_loading에서, DCP manager가 v_up_proj_combine을 제공할 때만 DCP group 전체의 W_UV를 all-gather해 둡니다.

if (
    self.dcp_manager is not None
    and self.dcp_manager.v_up_proj_combine is not None
):
    replace_parameter(
        self,
        "W_UV_dcp",
        get_dcp_group().all_gather(self.W_UV.contiguous(), dim=0),
    )

왜 gather가 필요할까요? all-to-all 이전 시점에 각 rank는 모든 head(H = N × dcp_world_size)의 partial latent output을 갖고 있습니다. 이를 V 공간으로 올리려면 자기 local head뿐 아니라 전체 head의 W_UV가 필요합니다.

forward 경로는 다음과 같이 바뀝니다.

Before

if num_merge_rows == attn_out.shape[0]:
    attn_out = self.dcp_manager.combine(
        attn_out,
        lse,
        ...
# v_up projection
self._v_up_proj(attn_out, out=mqa_output_slice)

After

if self.W_UV_dcp is not None:
    assert self.dcp_manager.v_up_proj_combine is not None
    self.dcp_manager.v_up_proj_combine(
        attn_out,
        lse,
        self.W_UV_dcp,
        mqa_output_slice,
        seq_lens=seq_lens,
        query_start_loc=query_start_loc,
    )
elif num_merge_rows == attn_out.shape[0]:
    attn_out = self.dcp_manager.combine(
    ...
# v_up projection, unless the DCP merge already applied it
if self.W_UV_dcp is None:
    self._v_up_proj(attn_out, out=mqa_output_slice)

v_up_proj_combine은 projection과 merge를 한 번에 수행하고 최종 결과를 mqa_output_slice에 바로 씁니다. 그래서 이후의 _v_up_proj는 건너뜁니다. 기존 경로는 elif로 그대로 남아 있어, 조건이 맞지 않으면 동작이 완전히 동일합니다.

2. vllm/v1/attention/ops/dcp.py — 새 combine 함수와 선택 로직

새 함수 dcp_a2a_v_up_proj_lse_reduce가 추가되었고, MLADCPManager는 NCCL all-to-all이면서 PCP를 쓰지 않을 때만 v_up_proj_combine을 설정합니다. 테스트가 이 선택 조건을 명시적으로 고정합니다.

# fallback(NCCL a2a) 백엔드
assert isinstance(manager.v_up_proj_combine, functools.partial)
assert manager.v_up_proj_combine.func is dcp_manager.dcp_a2a_v_up_proj_lse_reduce
assert manager.v_up_proj_combine.keywords == {
    "cp_group": group,
    "is_lse_base_on_e": True,
}

# symmetric-memory direct 백엔드, PCP 조합
assert manager.v_up_proj_combine is None

symmetric-memory direct A2A, ag_rs, PCP 경로는 손대지 않았다는 것을 테스트로 보장하는 구조입니다.

3. 작은 배치를 위한 threshold

리뷰어 shen-shanshan이 핵심적인 지적을 했습니다. project-first에서는 BMM이 local head가 아닌 전체 H개 head에 대해 돌고, gathered W_UV_dcp 전체를 매 step 읽어야 합니다. 이 비용은 배치 크기와 무관하게 고정적인 반면 A2A 절감분은 토큰 수에 비례하므로, 저동시성에서는 오히려 손해일 수 있다는 것입니다.

작성자는 이를 반영해 _V_UP_PROJ_FIRST_MIN_TOKENS(=8) 미만에서는 merge-first로 되돌아가고, 이때는 W_UV_dcp에서 자기 rank 몫을 슬라이스(복사 없는 view)해 projection합니다. 테스트도 threshold 경계 양쪽을 모두 검증합니다.

@pytest.mark.parametrize(
    "num_tokens",
    [_V_UP_PROJ_FIRST_MIN_TOKENS - 1, _V_UP_PROJ_FIRST_MIN_TOKENS],
)
def test_distributed_v_up_proj_a2a_matches_combine_then_v_up_proj(...)

4. 정합성 테스트

새 2-rank 테스트는 기존 경로(dcp_a2a_lse_reduce 후 local W_UV로 BMM)를 reference로 삼아 비교합니다. 특히 rank 0이 첫 요청의 KV를 갖지 않아 출력이 NaN인 경우를 일부러 넣었습니다.

# Rank 0 holds no KV for the first request; its output there is undefined.
seq_lens = torch.tensor([int(rank != 0), 3], device=device, dtype=torch.int32)
if rank == 0:
    cp_attn_out[:2] = float("nan")
...
local_w_uv = w_uv[rank * h_per_rank : (rank + 1) * h_per_rank]
expected = torch.bmm(latent.transpose(0, 1).float(), local_w_uv.float())

project-first에서는 NaN이 projection을 거쳐 그대로 전파되므로, merge 단계의 마스킹이 projection 이후 값에도 올바르게 적용되는지가 중요합니다. 작성자에 따르면 head 순서를 틀리게 하면 이 테스트가 실패해, gathered weight의 head 정렬까지 검증됩니다.

왜 이게 좋은가

수치

• A2A payload: 160 토큰 기준 rank당 레이어당 15.8 MB → 4.0 MB (약 4배 감소)
• 레이어 microbenchmark (8x MI355X, 96 heads, cold weights):

Tokens	merge first (µs)	project first (µs)
1	25.1	26.2
8	28.8	27.3
16	34.5	25.1
160	84.8	53.9
360	153.7	80.1

14 토큰에서는 +1.1 µs 손해, 16 토큰 이상부터 1248% 이득입니다. 그래서 threshold 8이 합리적인 경계입니다.

• End-to-end (Kimi-K3 1P1D, TP8/DCP8, MI355X): C16/C32에서 throughput +0.6%/+1.1%, P90 interactivity +1.4%/+2.8%, decode step -1.1%/-1.7%. GSM8K는 두 순서 모두 0.958로 동일합니다.
• 수치 정확도: fp64 reference 대비 Frobenius 상대 오차가 두 순서 모두 2.30e-3으로 동일합니다.

비용

레이어당 gathered W_UV_dcp가 추가되며 Kimi-K3 DCP8에서 12 MiB, GPU당 약 0.32 GiB입니다. 리뷰어가 local W_UV를 view로 만들어 메모리를 아끼자고 했지만, 작성자는 local W_UV와 W_UK_T가 이미 dequantized kv_b_proj의 view라 추가 storage가 없음을 설명했습니다.

일반적 교훈

1. 선형 연산은 통신 경계를 넘어 이동시킬 수 있다. reduce/merge가 선형이면, 차원을 줄이는 projection을 통신 앞으로 당겨 payload를 줄일 수 있습니다. MLA처럼 latent 차원(512)이 출력 차원(128)보다 큰 구조에서 특히 효과적입니다.
2. 연산 재배치는 compute를 늘린다. 통신을 줄이는 대신 BMM이 N배 head에 대해 돌고, 더 큰 weight를 읽습니다. 작은 배치에서는 latency-bound인 통신 절감이 미미하므로 token threshold 기반 fallback이 필수입니다.
3. 적용 범위를 좁게, 테스트로 고정하라. quantized W_UV, symmetric memory, PCP, ag_rs 경로는 그대로 두고, 각 선택 조건을 assert로 고정해 회귀를 막았습니다.
4. rounding 차이를 정량화하라. "수학적으로 같다"에서 그치지 않고 fp64 reference 대비 오차와 GSM8K로 실제 품질 영향을 확인한 점이 설득력을 높입니다.

참고 자료

• torch.bmm — up-projection에 사용되는 batched matmul
• torch.distributed.all_to_all_single — DCP a2a 통신의 기반 collective
• torch.distributed.all_gather — W_UV_dcp 생성에 쓰인 all-gather
• functools.partial — v_up_proj_combine 바인딩에 사용

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.