[vllm] vLLM chunk_kda 커널의 숨겨진 상태(h) 레이아웃 불일치 버그 수정 및 정확도 개선

2026년 4월 30일수정: 2026년 4월 30일

PR 링크: vllm-project/vllm#40956 상태: Merged | 변경: +None / -None

들어가며

vLLM은 대규모 언어 모델(LLM) 추론을 위한 고성능 서빙 엔진으로, 최적화된 CUDA 커널을 통해 뛰어난 처리량을 제공합니다. 이러한 성능의 핵심에는 chunk_kda와 같은 커널들이 있습니다. chunk_kda는 KDA(Key-Delta Attention) 메커니즘을 청크 단위로 처리하여 효율성을 높이는 역할을 합니다. 하지만 때로는 복잡한 커널 간의 상호작용에서 미묘한 버그가 발생할 수 있습니다. 이 PR은 chunk_kda 커널 내에서 h (hidden state) 행렬의 레이아웃 불일치로 인해 발생하는 심각한 정확도 문제를 해결합니다.

이 버그는 chunk_kda의 프리필(prefill) 과정에서 잘못된 출력을 야기했으며, 이는 모델의 전체적인 정확도에 치명적인 영향을 미쳤습니다. 특히, chunk_gated_delta_rule_fwd_h 커널이 h를 (V, K) 레이아웃으로 저장하는 반면, 출력 커널인 chunk_gla_fwd_kernel_o는 이를 (K, V) 레이아웃으로 로드하여 잘못된 서브 블록을 읽고 qg @ S^T 대신 qg @ S를 계산하는 문제가 있었습니다. 이 PR은 이러한 레이아웃 불일치를 바로잡아 chunk_kda를 사용하는 모든 모델의 프리필 정확도를 복원합니다.

코드 분석: 무엇이 왜 좋은 최적화/개선인가

이 PR의 핵심은 vllm/model_executor/layers/fla/ops/kda.py 파일의 chunk_gla_fwd_kernel_o 함수 내에서 h 행렬을 로드하고 사용하는 방식의 수정입니다. 또한, 이 버그를 검증하기 위한 새로운 테스트 케이스가 추가되었습니다.

1. `tests/kernels/test_kda.py` 파일 추가

가장 먼저 눈에 띄는 변경사항은 tests/kernels/test_kda.py 파일이 새로 추가된 것입니다. 이는 chunk_kda Triton 연산자의 정확도를 검증하기 위한 정밀 테스트입니다. chunk_kda의 출력을 float32로 구현된 naive recurrent reference 구현과 비교하여 RMSE-based relative error를 측정합니다.

이 테스트는 다양한 H (헤드 수), D (차원), cu_seqlens (누적 시퀀스 길이) 조합과 torch.float16, torch.bfloat16 데이터 타입을 사용하여 광범위한 시나리오에서 chunk_kda의 정확성을 검증합니다. 특히, initial_state를 사용하여 이전 청크의 상태를 전달하는 시나리오까지 커버하여 실제 프리필 과정을 시뮬레이션합니다.

왜 좋은가:

회귀 방지: 새로운 테스트 케이스는 향후 유사한 레이아웃 또는 계산 오류가 발생했을 때 이를 즉시 감지할 수 있도록 합니다. 이는 대규모 코드베이스에서 안정성을 유지하는 데 필수적입니다.
정확성 검증: naive_recurrent_kda와 비교하여 chunk_kda 커널의 수치적 정확성을 보장합니다. assert_close 함수를 통해 max abs err와 rmse ratio를 엄격하게 검사하여 미세한 오차도 허용하지 않습니다.
디버깅 용이성: 버그 발생 시, 어떤 특정 파라미터 조합에서 문제가 발생하는지 빠르게 파악할 수 있도록 돕습니다. PR 설명에서 vadiklyutiy의 질문에 ChenxiQ가 테스트 실패 로그를 첨부한 것이 이를 증명합니다.

2. `vllm/model_executor/layers/fla/ops/kda.py` 파일 수정

이 파일의 chunk_gla_fwd_kernel_o 함수는 chunk_kda의 출력 계산을 담당하는 Triton 커널입니다. 여기서 h (hidden state) 행렬을 로드하고 사용하는 방식에 오류가 있었습니다.

Before:

        p_h = tl.make_block_ptr(
            h + (i_tg * H + i_h) * K * V,
            (K, V),
            (V, 1),
            (i_k * BK, i_v * BV),
            (BK, BV),
            (1, 0),
        )

        # [BK, BV]
        b_h = tl.load(p_h, boundary_check=(0, 1))
        # works but dkw, owing to divine benevolence
        # [BT, BV]
        if i_k >= 0:
            b_o += tl.dot(b_qg, b_h.to(b_qg.dtype))

After:

        p_h = tl.make_block_ptr(
            h + (i_tg * H + i_h) * K * V,
            (V, K),
            (K, 1),
            (i_v * BV, i_k * BK),
            (BV, BK),
            (1, 0),
        )

        # [BV, BK]
        b_h = tl.load(p_h, boundary_check=(0, 1))
        # [BT, BV]
        if i_k >= 0:
            b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype))

무엇이 왜 좋은 최적화/개선인가:

p_h 블록 포인터 정의 수정:
- shape 인자가 (K, V)에서 (V, K)로 변경되었습니다. 이는 h 행렬이 실제로는 (V, K) 레이아웃으로 저장되어 있음을 반영합니다. 이전에는 (K, V)로 잘못 가정하여 데이터를 로드할 때 잘못된 메모리 영역을 참조했습니다.
- strides 인자가 (V, 1)에서 (K, 1)로 변경되었습니다. 이는 (V, K) 레이아웃에서 행(V)을 건너뛰는 보폭이 K임을 의미합니다.
- offsets 인자가 (i_k * BK, i_v * BV)에서 (i_v * BV, i_k * BK)로 변경되었습니다. 이는 (V, K) 레이아웃에 맞춰 V 차원에 대한 오프셋이 먼저 오고 K 차원에 대한 오프셋이 뒤에 오도록 수정된 것입니다.
- block_shape 인자가 (BK, BV)에서 (BV, BK)로 변경되었습니다. 이는 로드할 블록의 모양이 (V, K) 레이아웃에 맞춰 BV (V 차원의 블록 크기)가 먼저 오고 BK (K 차원의 블록 크기)가 뒤에 오도록 수정된 것입니다.
이러한 변경은 h 행렬의 실제 메모리 레이아웃((V, K))과 tl.make_block_ptr가 데이터를 해석하는 방식 간의 불일치를 해소합니다. FLA's TRANSPOSE_STATE=True 모드에서 h가 (V, K)로 저장되는데, chunk_gla_fwd_kernel_o는 이를 (K, V)로 읽으려 했기 때문에 발생한 오류입니다. 올바른 레이아웃으로 블록 포인터를 설정함으로써 커널이 정확한 데이터를 읽을 수 있게 됩니다.
tl.dot 연산 전 b_h 전치 (tl.trans(b_h)) 추가:
- 이전 코드에서는 b_qg와 b_h를 직접 tl.dot 연산했습니다. b_qg는 [BT, BK] 형태이고, b_h는 [BK, BV]로 로드될 것으로 예상되었으나, 실제로는 [BV, BK] 형태로 로드되었습니다 (위 블록 포인터 수정 후). 따라서 b_qg와 b_h를 직접 곱하면 차원 불일치 또는 잘못된 계산이 발생합니다.
- tl.trans(b_h)를 추가함으로써, [BV, BK] 형태의 b_h가 [BK, BV] 형태로 전치됩니다. 이제 b_qg ([BT, BK])와 tl.trans(b_h) ([BK, BV])는 올바른 차원 매칭을 통해 행렬 곱셈을 수행할 수 있게 됩니다. 이는 qg @ S 연산을 정확하게 수행하는 데 필수적입니다.

이 두 가지 수정은 h 행렬의 메모리 레이아웃에 대한 커널의 가정을 실제와 일치시키고, 그에 따라 올바른 행렬 곱셈을 수행하도록 합니다. 이는 단순히 코드를 변경하는 것을 넘어, 데이터의 물리적 저장 방식과 논리적 처리 방식 간의 정합성을 맞추는 근본적인 버그 수정입니다.

왜 이게 좋은가

이 PR은 단순한 버그 수정 이상의 의미를 가집니다. 이는 대규모 언어 모델의 추론 정확도에 직접적인 영향을 미치는 핵심적인 개선입니다.

모델 정확도 대폭 개선: PR 설명에 따르면, 이 버그 수정 전에는 Kimi-Linear-48B-A3B-Instruct 모델의 GSM8k 벤치마크 정확도가 17.36%에 불과했습니다. 버그 수정 후에는 정확도가 90.37%로 극적으로 상승했습니다. 이는 이 버그가 모델의 추론 결과에 얼마나 치명적인 영향을 미쳤는지 보여줍니다. 잘못된 h 행렬 처리는 모델의 내부 상태를 오염시켜 올바른 예측을 방해했던 것입니다.
견고한 테스트 인프라 구축: 새로운 test_kda.py 파일은 chunk_kda 커널의 정확성을 체계적으로 검증하는 기반을 마련했습니다. pytest를 활용한 다양한 시나리오 테스트는 향후 코드 변경 시 발생할 수 있는 잠재적인 회귀를 조기에 발견하고 방지하는 데 큰 도움이 됩니다. 이는 고성능 시스템에서 안정성을 유지하는 데 필수적인 요소입니다.
근본적인 원인 해결: ChenxiQ의 리뷰 댓글에서 언급된 것처럼, 이 버그는 PR #33291에서 GDN 상태 레이아웃이 (K, V)에서 (V, K)로 변경되었을 때 chunk_gla_fwd_kernel_o 커널에 동일한 수정이 누락되어 발생했습니다. 이 PR은 이러한 근본적인 원인을 파악하고 해결함으로써, 단순히 증상만 고치는 것이 아니라 시스템의 일관성을 회복시켰습니다.
성능 영향: 이 PR은 직접적인 성능 최적화(예: 연산 속도 향상)보다는 정확도 버그 수정에 초점을 맞추고 있습니다. 그러나 잘못된 계산으로 인해 모델이 올바른 결과를 내지 못한다면, 아무리 빠른 연산도 무의미합니다. 따라서 이 수정은 모델이 의도한 대로 작동하게 함으로써 실질적인 유용성을 확보하는 데 기여합니다. 정확한 모델은 더 적은 재시도와 더 신뢰할 수 있는 결과를 의미하며, 이는 간접적으로 전체 시스템의 효율성을 높입니다.

일반적 교훈:

메모리 레이아웃의 중요성: 저수준 커널 프로그래밍, 특히 GPU 프로그래밍에서는 데이터의 메모리 레이아웃이 성능과 정확도에 결정적인 영향을 미칩니다. 커널 간에 데이터를 주고받을 때 레이아웃 일관성을 유지하는 것이 매우 중요합니다.
철저한 테스트의 필요성: 복잡한 시스템에서는 예상치 못한 상호작용으로 버그가 발생할 수 있습니다. 특히 수치 계산의 정확도를 보장하기 위해서는 naive reference 구현과의 비교와 같은 정밀 테스트가 필수적입니다.
변경 사항 전파의 중요성: 한 부분의 변경(예: 데이터 레이아웃 변경)이 시스템의 다른 부분에 미치는 영향을 면밀히 검토하고, 필요한 모든 관련 코드에 변경 사항을 전파해야 합니다.

이 PR은 vLLM과 같은 고성능 딥러닝 추론 엔진의 안정성과 정확성을 보장하는 데 있어 이러한 원칙들이 얼마나 중요한지를 잘 보여주는 사례입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [onnxruntime] ONNX Runtime의 RISC-V Vector(RVV) 최적화: SGEMM과 Softmax 성능을 3배로 끌어올리기
현재글 : [vllm] vLLM chunk_kda 커널의 숨겨진 상태(h) 레이아웃 불일치 버그 수정 및 정확도 개선
다음글 [vllm] vLLM의 분산 추론 성능 극대화: 양방향 KV 캐시 전송을 통한 Prefill 최적화