지난 외전 4에서는, 실제 문서를 입력할 수 있도록 개조한 llama-bench를 사용하여, 실제 문서 6종류와 원래 무작위 토큰 7종류의 입력을 통해 PP 속도에 차이가 있음을 확인했습니다.
그때 Vulkan 버전의 PP 속도 차이가 ROCm 버전에 비해 큰 것이 신경 쓰였습니다.
입력 로딩 시에 차이가 발생한다는 것은 원인이 Flash Attention과 관련이 아닐까? 하고 Chappy님께 상담했더니, Chappy님은 멈춰주기는커녕 조사 방법을 여러 가지 제안해 주셔서 또다시 깊은 수렁에 빠졌습니다.
이후부터 Prompt Processing(입력을 로드하는 처리)를 PP, Token Generation(출력 생성)을 TG로 표기합니다. 둘 다 단위는 tok/s이며, 높을수록 빠릅니다.
- Vulkan 버전에서 일본어 문학과 C/C++ 코드의 PP 속도가 다른 주된 요인은,
Flash Attention (FA)의 GPU 처리 시간 차이에 대응하고 있었습니다. 컨텍스트 길이 64k의 GPU 연산 시간을 합산하면, FA만으로 약 32.9초의 차이가 있었습니다. - 동일한 FA 셰이더, 동일한 레이어, 동일한 KV 길이를 RGP(Radeon GPU Profiler)로 비교해도, FA 본체는 문학에서 약
311~314ms, 코드에서 약 76ms였습니다. 약 4.1배의 차이를 두 쌍의 측정으로 확인했습니다. 다만, 이는 64k 전체가 아닌 KV=49,152에서의 FA 1 노드 비교입니다. - Vulkan이 사용하는 마스크(mask)의 분류 결과를 조사해 보니, 16×64 타일을 통째로 건너뛸 수 있는 비율은 문학에서 35.68%, 코드에서 **75.03%**였습니다. 입력에 따라 실제로 처리하는 타일이 크게 다르다는 것이 FA 시간 차이의 유력한 설명이 될 것 같습니다. - 마스크 타일 수의 차이만으로 4.1배의 차이를 완전히 설명할 수는 없었고, ROCm과의 차이점 전체를 모두 해명할 수도 없었습니다.
소스, Windows 빌드, 검증 스크립트 등은 여기서 공개하고 있습니다.
이번 조사는 investigation/r5-vulkan-off-fa-dispatch-20261010
브랜치에서 실시했습니다.
속도와 관련된 부분만 먼저 말씀드립니다.
| 항목 | 내용 |
|---|
| PC | GMKtec NucBox EVO-X2 |
| ... | |
| 이번에는 지난 7종류 중, 컨텍스트 128k에서 가장 빨랐던 C/C++ 코드(llama.cpp 소스코드 일부)와 가장 느렸던 일본어 문학(『吾輩は猫である』)의 두 가지를 검증 대상으로 삼았습니다. 이후 GPU 내부 조사는 주로 컨텍스트 64k에서 진행했습니다. | |
지난번 64k PP 결과에서 이 두 종류만 추출하여 보여드립니다.
Vulkan은 코드가 명확히 빠른데, ROCm에서는 오히려 문학이 조금 더 빠릅니다.
| 64k, 61,789 tokens | Vulkan QSA union OFF | ROCm |
|---|---|
| 일본어 문학 | 285.479 tok/s | 404.027 tok/s |
| C/C++ 코드 | 329.465 tok/s | 394.669 tok/s |
- Attention과 query, KV: Attention은 현재 처리하는 토큰(query)이 과거의 토큰(Key/Value, 약어: KV) 중 어떤 것을 참고할지 계산하는 처리입니다. 입력이 길수록 과거의 KV가 늘어나기 때문에 무거운 처리가 됩니다. 'KV=49,152'는 과거 49,152 토분량의 KV가 쌓인 시점입니다. - Flash Attention (FA)와 타일: Attention을 GPU 상에서 효율적으로 계산하는 구현 방식입니다. 큰 계산을 작은 블록(타일)으로 나누어 메모리 전송 등을 줄입니다. - Sparse Attention / QSA: 모든 과거 토큰을 똑같이 참조하는 것이 아니라, 필요한 위치를 골라 Attention을 계산하는 방식입니다. QSA는 Qwen Sparse Attention이며, 이번 모델에서는 선택한 과거 위치에 대응하는 마스크(mask)를 만듭니다. - QSA의 mask: 참조해도 되는 위치는 기본적으로
0,
참조하지 못하게 하는 위치는 -inf
등 등으로 표현한다. QSA가 '어떤 것을 참조할지'를 결정하고, FA는 그 마스크(mask)를 보고 Attention을 계산한다. QSA와 FA는 같은 것이 아니다. -
Vulkan과 ROCm: 둘 다 AMD GPU 위에서 계산에 사용할 수 있는 백엔드이지만, GPU로 전달하는 프로그램의 구현은 다르다. GPU 위에서 동작하는 함수를 Vulkan에서는 '셰이더(shader)'(.comp 파일)라고 하고, ROCm/HIP에서는 **'커널(kernel)'**이라고 부르는 경우가 많다. dispatch는 그 GPU 함수를 시작(launch)하는 것을 의미한다.
Vulkan의 FA 구현에는 조건에 따라 Sparse FA 경로를 선택하는 메커니즘이 있다. 이는 선택된 부분만 모아서 처리하여 계산량을 줄이기 위한 경로로, 고속화 메커니즘이지만 현재 사용 중인 GPU, GPU 드라이버, 모델 등과의 조합 문제로 오히려 느려지고 있을 가능성도 있다.
이는 소스를 변경하지 않고 GGML_VK_FA_SPARSE_DISABLE라는 환경 변수로 비활성화할 수 있기 때문에, 가장 먼저 시도해 보았다.
| 64k, Vulkan QSA union OFF (b11535) | 일반 (환경변수 없음) | Sparse FA 비활성화 (=1) |
|---|
| 일본어 문학 | 286.158 tok/s | 285.796 tok/s |
| C/C++ 코드 | 330.443 tok/s | 330.384 tok/s |
64k로 PP(처리량)를 측정해 보니, Sparse FA 환경 변수 유무에 따라 속도는 거의 차이가 없었다. 지난번의 64k Vulkan QSA union OFF 결과와도 비슷한 값이어서, 결과적으로는 이전 결과를 재확인하는 정도였다.
주의할 점은, 이 환경 변수는 0을 지정하더라도 변수가 존재하는 것만으로 Sparse FA가 비활성화되는 사양이라는 것이다. 일반 측정을 할 때는 환경 변지 자체를 삭제했다.
이 시점에서는 'Sparse FA가 원인이 아닌 것 같다'고 생각했다. 나중에 진단 로그를 넣어 확인해 보니, 애초에 이번 64k PP에서는 Sparse FA가 선택되지 않았다.
다음으로는 무엇에 시간이 걸리고 있는지 알아보기 위해, FA의 경로 선택/호출 횟수/KV 구간 등을 CSV로 출력하는 기능을 추가했다. 주로 ggml/src/ggml-vulkan/ggml-vulkan.cpp에 집계 및 CSV 출력 처리를 추가했고, 카운터 정의는 ggml-vulkan-types.h에도 추가했다.
환경 변수 GGML_VK_FA_DISPATCH_DIAG=1로 활성화하여, GPU에서 실행된 시간이 아니라 CPU 측에서 어떤 FA 경로를 선택했는지를 기록하도록 했다.
진단 기능을 넣은 Vulkan b11551로, 64k의 C/C++ 코드와 일본어 문학을 읽히고 집계 CSV를 비교해 본 결과, 두 개의 CSV는 바이트 단위로 완벽하게 일치했다.
| 진단 항목 | 두 입력에서의 결과 | 의미 |
|---|
| FA calls | 732회 | 12층 × 61 마이크로 배치 (60×1,024 + 끝 349 토큰) |
| 선택된 FA 경로 | FA_COOPMAT1 | Cooperative Matrix를 사용하는 Vulkan 측의 FA 경로 |
sparse | 0 | Sparse FA는 선택되지 않음 |
mask_opt | 1 | 마스크 타일 최적화가 유효함 |
Br / Bc | 16 / 64 | 마스크나 FA 처리에 사용되는 query 측/KV 측의 블록 크기 |
split_k | 1 | 이번 호출에서는 KV 분할을 추가하지 않는 설정 |
KV 구간・n_kv_max 집계 | 일치 | KV 길이 분포나, 마스크와 관련된 참조 범위 상한도 동일 |
이로써 코드의 경우에만 다른 FA 셰이더로 분기했다는 가능성은 거의 사라졌다.
다만, 여기서 일치한 것은 FA 호출 조건일 뿐, 마스크의 **내용물(중심)**은 아니다. 같은 프로그램을 호출했더라도, GPU 내부에서 실제로 몇 번 계산할지는 다를 수 있다.
진단 CSV에서는 차이가 보이지 않았기 때문에, 다음으로는 제10회에서도 사용했던 Vulkan GPU performance logger로 조사해 볼 생각이다. 환경 변수로 GGML_VK_PERF_LOGGER=1로 설정하면, GPU 연산별 시간을 로그 출력해 주는 기능이다. 이는 직접 만든 것이 아니라 원래 llama.cpp에 있는 기능이다.
로그 출력 자체에 오버헤드가 있기 때문에, 이 상태에서의 PP 속도는 일반적인 속도 지표로 사용하지 않습니다. 이번에는 컨텍스트 길이 64k, 61 마이크로 배치 전체 로그를 합산하여 비교했습니다.
| GPU 연산 종류 (61 블록 합산) | 일본어 문학 | C/C++ 코드 | 코드 - 문학 |
|---|
FLASH_ATTN_EXT (FA) | 102.460초 | 69.565초 | −32.895초 |
MoE MUL_MAT_ID* | 40.790초 | 45.240초 | +4.450초 |
LIGHTNING_INDEXER | 1.168초 | 1.170초 | +0.002초 |
| 기타 GPU 연산 | 72.819초 | 71.726초 | −1.093초 |
| GPU 연산 합계 | 217.237초 | 187.701초 | −29.537초 |
**MoE (Mixture of Experts)**는 입력에 따라 사용할 전문가(Expert)를 선택하는 메커니즘입니다. MUL_MAT_ID*는 해당 Expert의 행렬 연산과 관련이 있습니다. LIGHTNING_INDEXER는 QSA에서 참조 위치를 선택하기 위한 처리와 관련됩니다.
C/C++ 코드는 FA 부분에서는 32.90초 더 빠르고, MoE 부분에서는 4.45초 더 느립니다. GPU 연산 합계에서는 29.54초 더 빠릅니다.
다음으로, FA 시간을 KV 구간별로 나누어 보았습니다.
| KV 구간 | 블록 수 | 일본어 문학의 FA | C/C++ 코드의 FA | 문학 - 코드 |
|---|
| 1~16k | 16 | 4.594초 | 4.388초 | +0.207초 |
| 17~32k | 16 | 16.636초 | 16.624초 | +0.013초 |
| 33~48k | 16 | 36.779초 | 26.929초 | +9.849초 |
| 49~60k | 12 | 43.465초 | 21.051초 | +22.414초 |
| 마지막 349 토큰 | 1 | 0.986초 | 0.573초 | +0.412초 |
※ 이 초수는 구간 내 여러 마이크로 배치를 합산한 값이며, 한 번의 FA 호출 시간이 아닙니다.
KV 17~32k에서는 거의 동일하지만, 33k를 넘어서면서 차이가 벌어지고 49~60k에서 커집니다. 이번 QSA union(제8회에 추가된 고속화)은 계속 OFF였기 때문에, 32k를 기준으로 union이 ON으로 전환된 것은 아닙니다.
여기서 멈추면 좋았을 텐데, Flash Attention 자체는 Vulkan에서도 ROCm에서도 작동해야 하는데, Vulkan에서만 차이가 큰 이유가 궁금해서 Flash Attention의 구체적인 처리를 조사하게 되었습니다.
AMD의 Radeon Developer Panel / Radeon GPU Profiler를 사용하여 GPU 연산 내부와 시간을 더 자세히 조사해 보았습니다. 이 역시 제10회에 한 번 사용한 적이 있습니다.
llama-cli나 llama-bench의 전체 GPU 이벤트 1회분을 가져오려고 하면 엄청난 데이터 양이 되기 때문에, 실행 중 일부를 지정하여 데이터를 가져올 필요가 있었는데, 이 위치 맞추기가 의외로 어려웠습니다. 위치 맞추기에 대한 이야기는 본론에서 벗어나므로 부록에 적겠습니다.
여러 번 시도한 결과, Vulkan의 FA에 EVOX2_FA_MAIN나 EVOX2_FA_MASK_OPT라는 레이블을 붙이고 RGP의 Auto Capture에서 Dispatch range로 위치를 지정하는 방식으로, 동일한 레이어・동일한 KV 길이・동일한 셰이더를 비교할 수 있었습니다.
조건은 Vulkan QSA union OFF, KV=49152, N=1024, Qcur-19 (view) (permuted), Br=16 / Bc=64, n_kv_max=2051이었습니다. RGP에서 확인된 FA 본체는 flash_attn_f32_f16_aligned_cm1이고, 마스크 전처리는 fa_mask_opt였습니다. 동일한 조건으로 각 입력을 2회 수집했습니다.
FA 본체는 문학/코드 기준으로 4.141배, 4.092배입니다. 두 경우 모두 거의 4.1배의 차이가 있었습니다.
반면, mask 전처리는 문학과 코드에서 거의 동일합니다(코드가 약간 느림). 즉, mask를 준비하는 처리가 아니라, 준비된 mask를 사용하여 Attention을 계산하는 FA 본체에서 큰 차이가 발생하고 있습니다.
※스크린샷 오른쪽의 시간 그래프 끝부분은, 문학이 약 540,000,000μs, 코드가 약 260,000,000μs로 축척이 다르기 때문에, 실제는 그래프를 나란히 놓은 모습보다 차이가 더 큽니다.
앞서 언급된 Vulkan GPU performance logger에서의 초 단위 숫자는 64k 입력의 전체 61 마이크로 배치, 여러 레이어를 합산한 것입니다. 여기의 ms 단위 숫자는 KV=49152 시점의 하나의 FA 노드입니다. 측정 범위가 다르기 때문에 단순히 더하거나 같은 배율을 전체에 적용할 수는 없습니다.
셰이더도 KV 길이도 같은데, FA 본체의 속도가 이렇게 다릅니다. 이 차이가 어디서 오는지 Chappy님께 소스 코드를 통해 조사해 달라고 했습니다.
소스 코드를 조사한 결과, 핵심이 될 만한 두 가지가 있었습니다.
flash_attn_mask_opt.comp:
FA에 전달되는 QSA mask를 16 query × 64 KV의 직사각형 단위로 조사하여 상태를 작은 분류 정보(2비트)로 요약합니다.
flash_attn_cm1.comp:
FA 본체입니다. QSA mask의 분류 정보를 읽고, 타일 전체가 참조 불가능하면 해당 타일에 대해 continue 하여 K/V 읽기나 연산을 건너뜁니다.
mask의 분류는 크게 다음과 같이 나뉩니다.
| 분류 | mask 내 상태 | FA 본체에서의 처리 |
|---|
ALL_NEG_INF | 16×64 전체가 참조 불가능(-inf) | 타일 전체를 건너뛸 수 있음 |
ALL_ZERO | 16×64 전체가 mask 0로 구성됨 | mask에 의한 제한이 없는 타일 |
| Mixed/Other | 유효/무효가 혼재하는 등 | 타일 전체를 일괄 건너뛸 수 없고, 내부에서 추가 판정이 필요할 수 있음 |
여기서 중요한 것은, FA의 셰이더가 같더라도 입력에 따라 생성되는 mask는 같지 않을 수 있다는 점입니다.
RGP로 본 fa_mask_opt의 시간은 작았지만, 이는 mask를 분류하는 전처리 시간입니다. 건너뛸지 여부를 판단하여 실제로 K/V를 처리하는 것은 후단의 FA 본체이므로, FA 본체의 실행 시간에 큰 차이가 나도 이상하지 않습니다.
소스 코드 조사 결과에 따르면, Attention을 계산할 때 타일을 통째로 건너뛸 수 있는지, 1개 타일 내에 무효/유효가 혼재하는지에 따라 처리 시간이 달라질 것 같아, GPU 측에서 생성된 2비트 분류 결과를 읽어와 CSV에 기록하는 처리를 추가해 보았습니다.
이번에 구현한 진단 코드는 GGML_VK_FA_MASK_DIAG=1로 활성화되며, GPU가 만든 QSA mask 분류 정보를 지정된 FA 호출에서 단 한 번만 복사하여 읽고 출력하는 것입니다. Vulkan b11571을 사용하고, RGP와 동일한 KV=49152 / Qcur-19를 대상으로 측정했습니다.
N=1024와 Br=16에서 query 측은 64행이 되고, KV=49152와 Bc=64에서 KV 측은 768열이 됩니다. 따라서 분류된 mask 타일은 64×768=49,152개가 됩니다.
C/C++ 코드가 통째로 건너뛸 수 있는 타일이 19,343개나 더 많았다. 게다가 일부 행만 극단적으로 다른 것이 아니라, 64행 전체에서 코드 쪽이 건너뛰기 가능한 타일 수가 더 많았다. 행별 중앙값은 문학 279개, 코드 577개(각 행 768타일)였다.
여기서 말하는 75.03%는 '입력 토큰의 75%를 버리고 있다'라는 의미가 아니다. 16×64 마스크 타일 중에서, 타일 단위로 모아서 계산을 건너뛸 수 있는 부분이 75.03%였다는 의미이다.
그렇다면 이것만으로 FA 본체의 약 4.1배의 차이를 설명할 수 있냐고 묻는다면, 그렇게 단순하지 않다. 건너뛰어지지 않는 타일 수의 비율은 31,614 / 12,271 ≒ 2.58배로,
RGP FA 본체 시간 비 약 4.1배와 일치하지 않는다. 각 타일의 배치, GPU의 workgroup이나 wave 실행 효율, 혼재된 타일 내부에서의 판별, 메모리 접근 등도 영향을 미칠 것 같지만, 이번에는 그 정도까지 분해할 수는 없었다.
그렇다면, C/C++ 코드가 일본어 문학보다 건너뛸 수 있는 타일이 많은 이유는 무엇일까?
그 이유를 직접 알아내기는 어려울 것 같아, Qwen3.8-Flash-Next의 tokenizer를 사용해 이번에 사용한 실제 입력 텍스트인 일본어 문학 '吾輩は猫である'는 본문 시작부터, C/C++ 코드(llama.cpp)는 첫 함수 정의의 시작부터 각각 16 토큰과 64 토큰 범위로 잘라보았다.
| 입력 | 16/64 토큰 범위 |
|---|
| 일본어 문학 16토큰 | |
| C/C++ 코드 16토큰 | |
| 일본어 문학 64토큰 | |
| C/C++ 코드 64토큰 | |
이미지로 비유하자면 일본어는 짧은 한 문장 정도가 16토큰이고, 일반적인 문장은 중간까지, 소스 코드도 짧은 한 문장 정도이다. 64토큰은 일본어의 경우 여러 문장, 소스 코드의 경우 매우 짧은 함수 하나 정도에 해당한다.
※'吾輩は猫である'에서 16토큰, 64토큰이 한 문장의 잘리는 지점에 걸쳐 있는 것은 아마 우연일 것이다.
※llama.cpp 코드의 64토큰 끝에 붙어있는 gg는 첫 번째 함수 정의가 63토큰으로 끝나고, 64토큰째에 두 번째 함수 정의인 ggml_tensor * llama_adapter_cvec::apply_to(ggml_context * ctx, ggml_tensor * cur, int il)의 시작 부분이 들어간 것이다.
두 문서의 토큰 경계를 표시해 보았다. 아래 표는 각 행 16토큰×12행=총 192토큰을 나타낸다.
│ = 토큰 경계, · = 반각 공백, □ = 전각 공백, ↵ = 개행
※아래 예시는 입력 문서의 본문, 첫 함수 정의의 시작을 지정하여 tokenize한 것이다. 앞뒤 문자열이 바뀌면 토큰 경계가 바뀔 가능성이 있으므로, 실제 측정 시점의 실제 토큰 경계를 보여준 것은 아니다.
※'吾輩は猫である'에 나오는 │ ⟨E7⟩ │ ⟨8D B0⟩ │는 한자 '獰'가 문자 코드로 분해되어 2개의 토큰으로 나뉜 것이다. 특이한 문자는 여러 개의 토큰으로 분할될 수 있는 것 같다.
일본어 문학 '나의 할아버지' (吾輩は猫である)
001–016 │ □ │ 輩 │ 《 │ わ │ が │ はい │ 》 │ は │ 猫 │ である │ 。 │ 名前 │ はまだ │ 無い │ 。
017–032 │ ↵ │ □ │ どこ │ で │ 生 │ れた │ かと │ んと │ 見 │ 当 │ 《 │ けん │ とう │ 》 │ が │ つか
033–048 │ ぬ │ 。 │ 何でも │ 薄 │ 暗 │ い │ じめ │ じめ │ した │ 所で │ ニ │ ャ │ ー │ ニ │ ャ │ ー
...
C/C++ 코드 (llama.cpp)
001–016 │ gg │ ml │ _tensor │ ·* │ ·llama │ _adapter │ _c │ vec │ :: │ tensor │ _for │ (int │ ·il │ ) │ ·const │ ·{
017–032 │ ↵ │ ··· │ ·if │ ·( │ il │ ·< │ · │ 0 │ ·|| │ ·il │ ·< │ ·layer │ _start │ ·|| │ ·il │ ·>
033–048 │ ·layer │ _end │ ·|| │ ·( │ size │ _t │ ) │ ·il │ ·>= │ ·tensors │ .size │ ()) │ ·{ │ ↵ │ ······· │ ·return
...
'나의 할아버지'를 기반으로 16 토큰 query × 64 토큰 KV의 마스크를 생각해 보면, 예를 들어
- 129번째 토큰의 'しい'을 query로, 1-64 토큰을 KV로 하여, query의 'しい'과 KV의 각 토큰 '□', '吾'... 각각에 대해 QSA 조건으로 참조 가능(유효)/참조 불가(무효)를 계산한다. 이것이 1번째 줄의 64개 데이터이다.
- 다음으로 130번째 토큰의 'とも'를 query로, 1-64 토큰을 KV로 하여 같은 처리를 한다. 이것이 2번째 줄의 64개 데이터이다.
- 같은 것을 query=144 토큰까지 반복하면, 16행 × 64개의 유효/무효 데이터가 만들어진다. 이것이 1타일이다.
- 이 16×64 데이터가 모두 무효가 되는 타일은 처리를 건너뛸 수 있다.
- 마찬가지로, query=129-144 토큰 × KV=65-128 토큰, query=145-160 토큰 × KV=1-64 토큰 등의 타일도 만들 수 있다.
하지만, 토큰으로 분해해 봐도 역시 C/C++ 코드가 일본어 문학보다 건너뛸 수 있는 타일이 많은 이유는 알 수 없다.
대략적으로 말하면 무효한 query와 KV의 조합이 C/C++ 코드에서는 64토큰 정도 단위로 모여 있고, 일본어 문학에서는 흩어져 있다는 것이겠지만, 왜 그런지는 모르겠다.
FA 처리 전에 위치 인코딩 등의 처리가 들어가기 때문에, 단순히 토큰의 의미만으로 참조 가능(유효)/참조 불가(무효)가 결정되는 것은 아닌 것 같아서, 제대로 조사하려면 실제 마스크 패턴과 토큰의 관계를 조사할 필요가 있고, 꽤 어려워 보인다.
그래서 이번에는 깊이 파고드는 것을 그만두고, 구현 조사로 돌아가자.
Vulkan의 C/C++ 코드와 일본어 문학의 차이는 많이 알게 되었지만, 지난 번 외전 4에서 측정한 결과에서는 ROCm의 경우 차이가 작았고, 64k에서는 일본어 문학이 C/C++ 코드보다 약간 빨랐다.
궁금해서 ROCm의 Flash Attention 처리가 어떻게 되어 있는지 소스 코드로 조사해 보았다.
그 결과, ROCm/HIP의 FA는 ggml-hip
가 ggml-cuda
계열의 구현을 공유하고 있어, Vulkan의 flash_attn_cm1.comp와는 다른 구현이었다.
- Vulkan: 마스크를 16×64 타일로 분류하여, 과거 KV 중간에 흩어져 있는
ALL_NEG_INF
타일을 FA 본체에서 건너뛸 수 있다. -
- ROCm/HIP: 이번에 조사한 MMA 계열 경로에서는, CUDA용 compact sparse gather 경로는 HIP에서 유효하지 않다. 하지만 마스크 최적화가 전혀 없는 것은 아니다.
flash_attn_mask_to_KV_max
라는 처리가, 끝 부분 쪽에 연속되는 불필요한 KV 블록을 조사하여 계산할 KV 범위의 종단을 짧게 할 수 있는 장치가 있다.
즉, Vulkan은 중간에 있는 건너뛸 수 있는 타일을 보고 있고, ROCm은 끝 쪽의 연속된 불필요 범위를 조사하는 경로가 있다. 메커니즘이 다르기 때문에, 입력에 따른 시간 차이가 다르게 나타날 수도 있다.
다만, 소스에 적힌 것만으로는 실제로 어떤 커널이 작동하는지 알 수 없기 때문에, 이쪽도 Radeon Developer Panel / RGP로 조사해 보았습니다.
64k, ROCm b11543으로 측정된 결과는 다음과 같습니다. 두 입력 모두 동일한 flash_attn_ext_f16<256,256,16,4,false,false,false>라는 AMD WMMA 계열의 FA 특수화와, flash_attn_mask_to_KV_max<16>
| ROCm/HIP RGP 이벤트 | 일본어 문학 | C/C++ 코드 |
|---|---|
| flash_attn_mask_to_KV_max<16> | 46.914µs | 68.492µs |
| flash_attn_ext_f16<256,256,16,4,false,false,false> | 104.911ms | 111.198ms |
| Total wavefronts | 3,072 | 3,072 |
| VGPR / SGPR | 256 / 102 | 256 / 102 |
| LDS | 36,864 bytes | 36,864 bytes |
| Occupancy | 4 / 16 | 4 / 16 |
FA 본체의 시간은 코드가 약 6.0% 더 길었습니다. Vulkan의 동일 노드에서 봤던 약 4.1배라는 큰 차이와는 상당히 다릅니다. 끝 쪽 마스크를 순회하는 커널이 실제로 작동했음도 확인했습니다. 순회 처리 자체는 FA 본체보다 훨씬 짧지만, 끝 쪽에서 실제로 KV를 얼마나 단축할 수 있었는지는 측정하지 않았습니다.
다만, 여기는 Vulkan만큼 엄격하게 위치 맞춤을 할 수는 없습니다. Vulkan처럼 KV=49152 / Qcur-19와 같은 라벨을 붙여주지 않았고, ROCm으로 측정한 FA의 정확한 KV 길이/레이어/웜업 여부 본 측정인지 확정할 수 없습니다.
따라서 이 104.911ms 대 111.198ms는, 측정된 동일 종류의 커널에서는 극단적인 차이가 보이지 않았다는 범위의 이야기입니다. Vulkan에서 엄격하게 위치를 맞춘 FA 1 노드와 같은 조건으로 수치 비교한 것은 아닙니다.
시간과 체력이 있었다면 ROCm에도 라벨을 붙여서 위치 맞춤까지 하고 싶었지만, 이번에는 시간이 다 되어 종료했습니다.
이번 조사는 처음에 Sparse FA를 의심했던 것부터 시작하여, 호출 조건 비교, GPU 연산별 시간 측정, RGP로 동일한 FA 노드를 포착하고, 마지막으로 마스크의 타일 분류 개수를 세는 단계까지 진행되었습니다.