[vllm] vLLM, DeepSeek V3.2 커널의 대규모 토큰 처리 안정성 강화

PR 링크: vllm-project/vllm#52381 상태: Merged | 변경: +79 / -6

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부십니다. 더 빠르고 효율적인 추론을 위한 기술 개발 또한 활발히 이루어지고 있으며, vLLM은 이러한 흐름의 선두에 서 있는 프로젝트 중 하나입니다. vLLM은 KV 캐싱 최적화와 같은 혁신적인 기법을 통해 LLM 추론 속도를 크게 향상시켰습니다.

하지만 모델의 복잡성이 증가함에 따라, 특정 모델 아키텍처나 연산에서 예상치 못한 문제가 발생하기도 합니다. 이번 PR은 vLLM에서 DeepSeek V3.2 모델을 지원하는 과정에서 발견된, 특히 fused_norm_rope 및 fused_q 커널에서 대규모 단일 반복(single-iteration) 토큰 처리 시 발생하는 안정성 문제를 해결하는 데 중점을 둡니다.

기존 코드에서는 CUDA 그리드 차원 중 하나에 토큰 수를 매핑했는데, 이 경우 65,536개의 토큰을 처리할 때 CUDA 그리드 블록의 최대 한계인 65,535개를 초과하여 Triton Error [CUDA]: invalid argument 오류가 발생했습니다. 이 글에서는 해당 문제를 어떻게 해결했는지, 그리고 그 원리가 무엇인지 심층적으로 분석해보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 두 가지 커널, fused_norm_rope와 fused_q에서 CUDA 커널 실행 시 그리드 차원의 매핑 방식을 변경하는 것입니다. 특히, 토큰 수가 CUDA 그리드 차원의 제한을 초과하지 않도록 조정하는 데 초점을 맞추고 있습니다.

1. vllm/models/deepseek_v32/common/kernels.py 변경 사항

이 파일은 DeepSeek V3.2 모델의 핵심 연산을 Triton 커널로 구현한 부분입니다. 두 가지 주요 함수인 _fused_norm_rope_kernel과 _fused_q_kernel 내부의 그리드 차원 매핑 로직이 수정되었습니다.

_fused_norm_rope_kernel

기존 코드에서는 tok_idx (토큰 인덱스)를 tl.program_id(1)로, pid (프로그램 ID)를 tl.program_id(0)으로 사용했습니다. 이는 (4, num_tokens) 형태의 그리드 차원을 의미합니다. 즉, 두 번째 차원(grid-y)에 num_tokens가 매핑됩니다.

Before:

    pid = tl.program_id(0)
    tok_idx = tl.program_id(1)

수정된 코드에서는 이 순서를 바꾸어 tok_idx를 tl.program_id(0)으로, pid를 tl.program_id(1)으로 변경했습니다. 또한, tok_idx를 tl.int64 타입으로 변환하여 안전한 주소 연산을 보장합니다. 이는 (num_tokens, 4) 형태의 그리드 차원을 의미하며, 첫 번째 차원(grid-x)에 num_tokens가 매핑됩니다.

After:

    tok_idx = tl.program_id(0).to(tl.int64)
    pid = tl.program_id(1)

그리고 이 커널을 호출하는 fused_norm_rope 함수에서도 그리드 차원 설정이 변경되었습니다.

Before:

    _fused_norm_rope_kernel[(4, num_tokens)](

After:

    _fused_norm_rope_kernel[(num_tokens, 4)](

_fused_q_kernel

_fused_q_kernel 역시 유사한 변경이 이루어졌습니다. 기존에는 pid가 tl.program_id(0), tok_idx가 tl.program_id(1)이었습니다. 이는 (3, num_tokens, grid_heads) 형태의 그리드 차원을 의미하며, 두 번째 차원(grid-y)에 num_tokens가 매핑됩니다.

Before:

    pid = tl.program_id(0)
    tok_idx = tl.program_id(1)

수정 후에는 tok_idx가 tl.program_id(0)으로, pid가 tl.program_id(1)으로 변경되었고, tok_idx는 tl.int64로 변환됩니다. 이는 (num_tokens, 3, grid_heads) 형태의 그리드 차원을 의미하며, 첫 번째 차원(grid-x)에 num_tokens가 매핑됩니다.

After:

    tok_idx = tl.program_id(0).to(tl.int64)
    pid = tl.program_id(1)

마찬가지로 fused_q 함수에서도 그리드 차원 설정이 변경되었습니다.

Before:

    _fused_q_kernel[(3, num_tokens, grid_heads)](

After:

    _fused_q_kernel[(num_tokens, 3, grid_heads)](

2. tests/kernels/test_fused_deepseek_v32_norm_rope.py 변경 사항

이번 PR은 단순히 코드를 수정하는 것을 넘어, 변경된 로직이 올바르게 동작하는지 검증하기 위한 새로운 테스트 케이스를 추가했습니다.

test_fused_norm_rope_supports_large_token_count

이 테스트 함수는 fused_norm_rope 커널이 65,536개의 토큰을 처리할 때 발생하는 문제를 재현하고, 수정 후에는 정상적으로 동작하는지 확인합니다. num_tokens를 65536으로 설정하고 fused_norm_rope 함수를 호출하여 예외가 발생하지 않는지, 그리고 결과가 올바른지 검증합니다.

test_fused_q_triton_supports_large_token_count

fused_q 커널의 Triton fallback 경로에 대해서도 유사한 테스트가 추가되었습니다. 이 테스트 역시 num_tokens를 65536으로 설정하고 fused_q 함수를 호출하여 대규모 토큰 처리 시 발생하는 오류를 방지하고 정확성을 검증합니다.

이러한 테스트 케이스 추가는 코드 변경의 안정성을 높이고 향후 유사한 문제가 재발하는 것을 방지하는 데 중요한 역할을 합니다.

왜 이게 좋은가?

이번 PR의 핵심은 CUDA 그리드 차원 매핑 방식을 변경하여 CUDA의 제한 사항을 우회하고 커널의 안정성을 높인 것입니다. 구체적으로 다음과 같은 장점들이 있습니다.

1. CUDA 그리드 차원 제한 극복: CUDA는 grid-y 차원의 최대 블록 수를 65,535개로 제한합니다. 기존 코드에서는 이 grid-y 차원에 토큰 수를 직접 매핑했기 때문에, 65,536개 이상의 토큰을 처리할 때 invalid argument 오류가 발생했습니다. 변경된 코드에서는 토큰 수를 grid-x 차원에 매핑함으로써 이 제한을 효과적으로 우회했습니다. grid-x 차원은 훨씬 더 큰 값을 지원하므로, 대규모 토큰 처리가 가능해집니다.
2. 안정성 향상: 이 변경을 통해 DeepSeek V3.2 모델과 같이 많은 토큰을 한 번에 처리해야 하는 경우에도 커널이 안정적으로 동작하게 됩니다. 이는 LLM 추론 시 발생할 수 있는 치명적인 오류를 방지하고, 서비스의 신뢰도를 높이는 데 기여합니다.
3. 성능 유지: PR 설명에 따르면, 이 변경은 커널의 수학적 연산이나 CuTeDSL fused_q 경로에는 영향을 미치지 않습니다. 즉, 성능 저하 없이 안정성만 향상시킨 것입니다. 오히려 불필요한 오류 처리나 재시도 로직이 줄어들면서 미미한 성능 향상이 있을 수도 있습니다.
4. 테스트 커버리지 확대: 새로운 테스트 케이스 추가는 해당 변경 사항의 유효성을 검증하고, 엣지 케이스(edge case)에 대한 코드의 견고함을 높입니다. 이는 향후 유지보수 및 기능 확장에 긍정적인 영향을 미칩니다.

일반적 교훈:

• 하드웨어/플랫폼 제약사항 인지: GPU 커널 개발 시, 사용하는 하드웨어(CUDA, ROCm 등) 및 라이브러리(Triton 등)의 명시적인 제약사항(예: 그리드 차원 제한)을 반드시 인지하고 설계해야 합니다.
• 엣지 케이스 테스트의 중요성: 특히 경계값(boundary value)에서 발생하는 문제는 재현하기 어렵지만, 서비스 안정성에 치명적일 수 있습니다. 따라서 엣지 케이스에 대한 철저한 테스트 케이스 작성이 필수적입니다.
• 유연한 그리드 차원 매핑: 연산의 특성(예: 토큰 수, 헤드 수 등)에 따라 그리드 차원을 유연하게 매핑하고, 각 차원의 제약사항을 고려하여 최적의 구성을 찾아야 합니다.

리뷰어 피드백

이번 PR에는 WoosukKwon님의 /ci run 명령이 두 차례 포함되어 있습니다. 이는 CI(Continuous Integration) 파이프라인을 수동으로 트리거하여 변경 사항이 CI 환경에서 정상적으로 빌드되고 테스트를 통과하는지 확인하려는 의도로 보입니다. 이는 코드 변경의 안정성을 확보하기 위한 좋은 관행입니다.

References

• Triton Documentation
• CUDA C++ Programming Guide
• DeepSeek V3.2 Model Architecture (참고용, 직접적인 문서 링크는 아님)

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