[vllm] [vLLM] GLM-5.3-Flash: 롱 컨텍스트 Prefill 성능을 32% 끌어올린 TP 쿼리 로우 샤딩 분석

2026년 10월 11일수정: 2026년 10월 11일

PR 링크: vllm-project/vllm#54951 상태: Merged | 변경: +436 / -38

들어가며: 중복 계산의 늪

대규모 언어 모델(LLM)의 컨텍스트 길이가 1M(백만) 토큰을 넘어서는 시대가 되었습니다. 하지만 컨텍스트가 길어질수록 Prefill 단계에서의 연산 비용은 기하급수적으로 증가합니다. 특히 GLM-5.3-Flash와 같은 모델에서 사용하는 Sparse Indexer는 MQA(Multi-Query Attention) 스코어링과 Top-k 선택 과정을 거치는데, 기존 vLLM의 Tensor Parallel(TP) 구현에서는 이 과정이 모든 Rank에서 중복으로 계산(Redundant Computation)되고 있었습니다.

즉, TP4 환경이라면 4개의 GPU가 똑같은 쿼리 로우(Query Rows)에 대해 똑같은 스코어링을 수행하고 있었던 것입니다. 이번 PR(#54951)은 이러한 낭비를 제거하기 위해 쿼리 로우를 TP Rank별로 샤딩(Sharding)하고, 결과인 int32 인덱스만 all_gatherv로 모으는 최적화를 도입했습니다.

핵심 아이디어: 쿼리 로우 샤딩 (Query-Row Sharding)

이 최적화의 핵심은 "쿼리 로우는 서로 독립적이다"라는 점에 착안한 시퀀스 병렬화(Sequence Parallelism)의 일종입니다.

Before: 모든 TP Rank가 전체 쿼리 로우 $N$개에 대해 전체 Key 이력을 대상으로 스코어링을 수행.
After: 각 Rank는 $N/TP$개의 쿼리 로우만 담당하여 스코어링을 수행. 단, 각 로우는 여전히 전체(Full) Key 이력을 참조함 (Attention 근사가 아님).
Reassembly: 계산된 int32 타입의 Top-k 인덱스들만 all_gatherv를 통해 다시 합침.

이 방식이 효과적인 이유는 연산량($O(Rows \times Context)$)은 컨텍스트 길이에 비례해 늘어나지만, 통신량($O(Rows \times Topk)$)은 컨텍스트 길이와 무관하게 고정되어 있기 때문입니다. 따라서 컨텍스트가 길어질수록 이득이 극대화됩니다.

코드 분석: 무엇이 바뀌었나?

1. `PrefillRowShard` 클래스 도입 (`vllm/model_executor/layers/sparse_attn_indexer.py`)

기존에는 모든 로우를 한꺼번에 처리했지만, 이제는 각 Rank가 담당할 범위를 계산하고 해당 윈도우만 처리하도록 변경되었습니다.

# After: 각 Rank의 담당 범위를 관리하는 클래스
class PrefillRowShard:
    def __init__(self, num_decode_tokens: int, sizes: list[int], rank: int):
        self.num_decode_tokens = num_decode_tokens
        self.sizes = sizes
        self.start = num_decode_tokens + sum(sizes[:rank])
        self.stop = self.start + sizes[rank]

    def narrow(self, chunk) -> "tuple[int, int, torch.Tensor, torch.Tensor] | None":
        """현재 Rank가 담당하는 쿼리 로우 범위로 chunk를 좁힘"""
        start = max(chunk.token_start, self.start)
        end = min(chunk.token_end, self.stop)
        if start >= end:
            return None
        lo, hi = start - chunk.token_start, end - chunk.token_start
        return start, end, chunk.cu_seqlen_ks[lo:hi], chunk.cu_seqlen_ke[lo:hi]

2. 비용 기반 밸런싱 샤딩 (`vllm/v1/attention/backends/mla/indexer.py`)

단순히 로우 개수를 $N/TP$로 나누는 것이 아니라, 각 로우가 가진 인과적(Causal) 히스토리 길이를 고려하여 연산 비용이 균등하도록 샤딩합니다.

# After: 비용 균형을 맞춘 샤딩 로직
def balanced_prefill_row_shard(
    seq_lens: torch.Tensor, query_lens: torch.Tensor, ratio: int, tp_size: int
) -> list[int] | None:
    # ... (생략) ...
    # 각 로우의 스코어링 비용(쿼리 로우 x 참조할 Key 길이)을 계산하여 누적합 기반으로 분할
    costs = (seq_lens - query_lens // 2) * query_lens
    total_cost = costs.sum().item()
    target_cost_per_rank = total_cost / tp_size
    # ... (생략) ...
    return shard_sizes

3. `all_gatherv`를 통한 인덱스 재조립

스코어링이 끝난 후, 각 Rank가 가진 부분적인 Top-k 인덱스를 하나로 합칩니다.

# After: 통신을 통한 결과 취합
def exchange_topk(self, topk_buffer: torch.Tensor, width: int) -> None:
    local_slice = topk_buffer[self.start : self.stop]
    # PyNCCL을 사용하여 각 Rank의 결과를 취합
    gathered = get_tp_group().all_gatherv(local_slice, dim=0, sizes=self.sizes)
    topk_buffer[self.num_decode_tokens : self.num_decode_tokens + sum(self.sizes)] = gathered

리뷰어와의 논의: 안정성과 호환성

리뷰 과정에서 몇 가지 중요한 기술적 논의가 있었습니다.

Speculative Decoding & Model Runner V2: 리뷰어 @ZJY0516의 요청으로 최신 V2 러너와 MTP(Multi-Token Prediction) 환경에서 재검증이 이루어졌습니다. 결과적으로 1M 컨텍스트에서 26%의 TTFT 감소를 유지하며 호환성을 입증했습니다.
메모리 프로파일링: @yewentao256은 all_gatherv가 사용하는 추가 GPU 메모리가 프로파일링(Dummy run)에 반영되지 않는 점을 지적했습니다. 이를 해결하기 위해 sparse_attn_indexer의 프로파일링 단계에서 최악의 경우(Worst-case)의 통신 버퍼 크기를 미리 예약하도록 수정되었습니다.
활성화 임계값 (Threshold): MIN_TP_SHARD_ROWS_PER_RANK를 16,384로 설정했습니다. 이는 통신 오버헤드가 연산 절감량보다 커지는 짧은 컨텍스트에서의 성능 저하를 막기 위한 보수적인 수치입니다.

왜 이게 좋은 최적화인가?

성능 수치 (GLM-5.3-Flash, 4× H200, TP4)

Input Context	TTFT Reduction	Throughput Increase
128K	3.69%	2.79%
512K	14.68%	15.78%
1M (1,000,000)	25.42%	32.59%

일반적인 교훈

Compute vs Communication: 분산 시스템 최적화의 핵심은 연산과 통신의 트레이드오프입니다. 이번 PR은 $O(N^2)$에 가까운 연산을 포기하고 $O(N)$의 통신을 선택함으로써 롱 컨텍스트에서 압도적인 승리를 거두었습니다.
Cost-Balanced Partitioning: 단순 균등 분할이 아닌 실제 연산 부하(Causal history)를 고려한 샤딩은 분산 환경에서 스트래글러(Straggler) 문제를 방지하는 좋은 사례입니다.
Exact Equivalence: 이 최적화는 Attention 메커니즘 자체를 근사화하지 않습니다. 수학적으로 동일한 결과를 내면서 성능만 개선했다는 점에서 매우 안전하고 강력합니다.

마치며

이번 최적화는 vLLM이 롱 컨텍스트 서빙의 강자임을 다시 한번 증명했습니다. 특히 AMD ROCm 플랫폼으로도 즉시 포팅되어 MI325X 등에서도 성능 향상이 확인되었다는 점은 이 기법의 범용성을 잘 보여줍니다. 1M 이상의 초거대 컨텍스트를 다루는 서비스라면 이번 업데이트는 필수적이라고 할 수 있습니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [sglang] DeepSeek-V4.1 성능 최적화: num_stages=1 적용으로 _hc_mix_stats_partial_kernel 속도 향상
현재글 : [vllm] [vLLM] GLM-5.3-Flash: 롱 컨텍스트 Prefill 성능을 32% 끌어올린 TP 쿼리 로우 샤딩 분석
다음글 [uvicorn] Uvicorn의 리로드 패턴 최적화: O(N²)에서 O(N log N)으로의 성능 개선