Split-KV가 확산 드래프터의 긴 문맥 학습 통신을 줄인다

무엇인가

이 논문은 확산 드래프터(diffusion drafter)를 긴 문맥에서 학습할 때 생기는 분산 학습 병목을 다룬다. 확산 드래프터는 추론에서 여러 드래프트 토큰을 병렬로 노이즈 제거해 speculative decoding을 가속하는데, 학습에서는 앵커 샘플링으로 소수의 질의를 만들고 훨씬 긴 문맥에 블록 희소 마스크를 적용한 교차 어텐션을 수행한다. 저자들은 DFlash 드래프터로 Qwen3.5-9B와 Qwen3.5-35B-A3B를 SFT하며 GSM8K, MATH-500, HumanEval, MBPP, MT-Bench 평균 acceptance length를 측정했고, 질의 예산을 고정한 채 문맥을 늘리면 평균이 최대 0.82%만 변했으며 질의 예산을 늘리면 마지막 두 배 증가가 각각 0.98%, 1.57% 개선에 그쳤다고 보고한다. 따라서 질의 길이가 문맥보다 훨씬 짧은 Q≪S가 실용적 영역이라는 것이 논문의 출발점이다.

어떻게 동작하나

Split-KV는 이 Q≪S를 이용하는 정확한 context parallelism 방법이다. 질의 원본 X와 키·값 원본 U를 P개 장치에 나누되, 각 장치의 키·값 투영 k_p, v_p는 로컬에 남기고 투영된 질의 q만 all-gather한다. 각 장치는 로컬 키·값에 대해 블록 희소 마스크를 적용한 어텐션을 계산해 로컬 log-sum-exp ℓ_p와 출력 o_p를 얻고, ℓ=log Σ exp(ℓ_p), o=Σ exp(ℓ_p−ℓ)o_p로 전체 결과를 정확히 재구성한다. 출력 투영이 선형이라는 점을 이용해 합산과 reduce-scatter를 함께 수행하며, 역전파에서는 질의 그래디언트만 통신하고 dk_p, dv_p는 로컬에 둔다. 이 설계로 레이어당 통신량은 (P−1)/P×4QH 원소가 되어 문맥 길이 S와 무관해지고, 순전파 피크 메모리는 2(S+Q)H/P+2QH가 된다.

무엇과 다른가

앵커 샘플링은 부하 분산도 바꾼다. 인과적 자기어텐션의 위치별 부하는 선형이지만, 앵커 기반 어텐션에서는 위치 x의 부하가 B Σ 1[s_i≤x<a_i]로 표현되는 비선형·샘플 의존 프로파일이 된다. Split-KV는 모든 장치가 전체 질의를 보는 점을 이용해 문맥 청크 j의 비용을 c_j = B Σ (min(a_i, λ_j+m) − max(s_i, λ_j))^+로 계산하고, 가장 비용이 큰 청크를 아직 J개 미만을 받은 가장 부하가 낮은 랭크에 배정하는 LPT 전략을 쓴다. 이 계획은 통신 없이 만들어지고, 학습 스텝마다 한 번의 all-to-all로 문맥 원본 행을 재분배하며 드래프트 레이어 전체에 걸쳐 비용이 분산된다. 논문의 예에서 Zigzag는 랭크 부하를 [32,32,38,42]로 만들었지만 LPT는 [36,36,36,36]으로 맞춰 불균형 비율 ρ를 1.167에서 1.000으로 낮춘다. 구현은 PyTorch 기반 VeOmni에서 FSDP2, 시퀀스 병렬, 전문가 병렬을 조합했고 약 4.3K줄의 Python으로 구성되며, 가변 로컬 키·값 길이를 4096 배수로 올려 정적 형상 버킷을 만들어 FlexAttention 재컴파일을 줄인다.

어떻게 쓰나

실험은 최대 8개 노드, 노드당 8개 가속기에서 진행됐고 노드 내 600GB/s 양방향, 노드 간 400Gb/s RDMA 환경을 쓴다. Qwen3.5-35B-A3B와 Qwen3.5-27B용 DFlash 드래프터를 오프라인으로 학습해 타깃 모델 특징 추출 시간을 제외한 드래프터 학습 스텝 시간을 측정했고, 마이크로배치당 프롬프트 대 응답 비율은 1:3, 블록 크기 B=16, DP 크기는 1로 두고 CP 크기 P를 8, 16, 32, 64로 늘렸다. 비교 대상은 Ulysses, Ring Attention, RingX, LV-XAttn이다. 약한 스케일링에서 P를 8에서 64로 늘리며 샘플당 문맥을 128k에서 1M로 키우고 질의는 8k로 고정했을 때, Split-KV는 여덟 개의 모델·스케일 구성 중 여섯에서 가장 빨랐고 하나에서는 Ulysses와 동률이었으며, 더 빠른 경우 1.01~2.81배 속도를 냈다. 8배 확장에서 스텝 시간은 35B-A3B가 3.67~4.02초, 27B가 3.33~3.94초에 머물렀고 P=64에서 약한 스케일링 효율 91%와 85%, 가장 빠른 베이스라인 대비 2.81배와 2.19배를 기록했다.

전제와 한계

강한 스케일링에서는 b=4, 샘플당 문맥 256k, 샘플당 질의 8k로 전체 작업량을 고정한 채 P를 8에서 64로 늘렸다. Split-KV는 여덟 구성 중 일곱에서 가장 빨랐고 1.04~3.67배 속도를 냈으며, 35B-A3B의 스텝 시간을 3.46초에서 0.78초로, 27B를 2.98초에서 0.87초로 줄여 각각 4.44배와 3.43배 개선을 보였다. 다만 P=64에서 강한 스케일링 효율은 55%와 43%로, 로컬 어텐션 작업은 줄어드는데 전역 질의 페이로드와 집합 통신 지연·동기화는 고정되거나 커지는 비용이 남는다. 질의 대 문맥 비율 실험에서는 P=32, 샘플당 문맥 128k, 패킹된 질의 64k를 고정하고 Q:S를 1:4에서 1:32까지 낮췄을 때, Split-KV의 전체 스텝 시간은 35B-A3B에서 1.11배, 27B에서 1.44배만 증가했다. 같은 조건에서 Ulysses는 2.18배와 2.05배, Ring Attention은 5.63배와 6.32배, RingX는 4.35배와 4.09배 증가했고, LV-XAttn은 증가폭이 1.17배와 1.22배로 작았지만 순차 홉 스케줄 때문에 Split-KV보다 1.98~3.48배 느렸다.

부하 분산 품질과 계획 비용은 동일한 Split-KV 전송 위에서 청크 배정 방식만 바꿔 분리 평가했다. 샘플당 질의 8k, b=4에서 문맥 길이를 16k에서 256k로 늘리는 실험과 CP 크기를 8에서 64로 늘리는 실험 모두에서, 불균형을 100(ρ−1)로 측정했을 때 No balance는 평균 46.49%와 55.26%, Zigzag는 7.94%와 8.01%, DCP는 0.66%와 0.80%, LPT는 0.38%와 0.21%를 기록했다. 계획 지연은 35B-A3B 기준 Zigzag가 0.01~0.05ms, LPT가 0.38~1.24ms였고, DCP는 문맥 길이에 따라 16.92~148.70ms, P에 따라 148.38~1164.64ms로 커졌다. 즉 LPT는 DCP보다 거의 세 자릿수 빠르게 계획하면서 더 낮은 불균형을 달성했고, No balance 대비 종단 학습 스텝 시간을 1.15~1.42배 줄였다.

개발자 관점에서 이 논문은 긴 문맥 확산 드래프터 학습에서 통신량을 문맥 길이가 아니라 질의 길이에 묶는 설계와, 샘플마다 달라지는 어텐션 비용을 싸게 분산 배치하는 방법을 보여준다. 다만 이 이득은 Q≪S와 앵커 기반 블록 희소 교차 어텐션 패턴에 의존하므로, 질의 예산이 커지거나 일반 인과적 자기어텐션을 학습하는 경우에는 전제가 달라진다. 논문은 별도 한계 절을 두지 않지만 P가 커질수록 집합 통신 지연과 동기화로 약한·강한 스케일링 효율이 떨어지는 점, 전체 질의 all-gather로 추가 질의 메모리가 드는 점, 4096 배수 정적 버킷을 위한 패딩이 필요할 수 있다는 점을 본문에서 밝힌다.