b11430: hexagon: matmul 및 flash-atten 확장성 업데이트 (#29974)

• hexagon: row-split 멀티코어 환경에서 head-parallel flash_attn 파티셔닝

row-split 모드에서는 각 코어가 모든 MUL_MAT의 출력 행 조각(output row shard)을 계산하지만, 이전에는 flash_attn이 헤드별(heads)이 아닌 Q 토큰별(flat qrow split)로 파티셔닝했습니다. 이 때문에 모든 코어가 전체 KV 캐시(all n_kv_heads)를 읽어야 했고, 이는 flash_attn에서 멀티코어의 메모리 대역폭 이점(memory bandwidth benefit)을 상쇄시켰습니다.

HMX와 HVX flash_attn 커널 모두 n_kv_heads가 n_cores로 나누어 떨어질 때 KV 헤드별로 파티셔닝하도록 변경되었습니다. 즉, i번째 코어가 오직 자신의 헤드 조각(head shard)인 [i*n_kv_heads/N, (i+1)*n_kv_heads/N)만 처리하고, 해당 헤드 조각의 KV 캐시만 읽습니다. n_kv_heads가 n_cores로 나누어 떨어지지 않는 경우(예: 4개 코어에서 2개의 KV 헤드를 가진 Gemma-4)에는 원래 토큰 블록 분할(token-block split)로 폴백합니다.

이는 GGML_HEXAGON_FA_HEAD_SPLIT 플래그에 의해 제어되며 (기본값 1 = 활성화), 이 플래그는 128바이트 커널_파라미터 블롭(kernel_params blob) 제한 내에서 기존 is_dst_fp32 kparams 바이트의 비트 1에 패킹됩니다.

4c row-split 환경에서의 측정된 성능 향상 (PP t/s, ubatch=1024):
Qwen3-0.6B: 6977 -> 11026 (+58%)
llama-3.2-3B: 3717 -> 5522 (+49%)
Qwen3.5-4B: 2739 -> 2855 (+4%)
Gemma-4 MoE: 변화 없음 (MoE FFN이 지배적이며, 폴백 경로 사용)

TG는 변경되지 않았습니다 (flash_attn은 레이어당 matmul+barrier 비용에 비해 디코드 시간에서 작은 부분을 차지합니다).

• hex-fa: kern_params 및 head-split 선택 정리

• hex-fa: run.py에 -fa-head-split 옵션 추가

• hex-mdev: row-split 시나리오에서 감소된 작업량을 반영하도록 matmul 솔버 업데이트

• hex-mmid: 다중 장치(multi-dev) 시나리오에서 expers를 통한 개선된 작업 분할

• hex-fa: 모델/n-hvx/ctx-len 스윕에 기반하여 HMX 게이팅 업데이트

• hex-fa: 호스트에서 softcap/scale 사전 계산

• hexagon: 다중 시퀀스(multi-sequence)에서 HMX를 사용하기 위해 matmul을 2D로 평탄화(flatten)

• hex-mm: kparams 정리 및 collapse를 사용하여 3/4D -> 2D 매핑 적용

• hex-mm: collapse 폴백(fallback)의 오타 수정

• hex-mm: act 텐서에 대한 일관된 명명 규칙 재정비

• hex-mm: 융합된 matmul에서 차원 축소(colapsing dims) 지원 추가

• hex-build: WoS 빌드 오류 수정

• hex-mm: can_collapse에서 dst 스트라이드를 강제하도록 보장

• hex-fa: head-split 검사를 위한 한 줄짜리 커밋(onliner commit) 추가

• hex-fa: 사용하지 않는 local head_split 변수 제거

• hex-fa: can_split 검사 강화

• hex-mm: 융합되지 않은 경로(unfused paths)가 act를 src1 대신 사용하도록 업데이트

• hex-mm: 분할에 대해 모든 dst를 확인하도록 보장

• hexagon: 배치된 HMX matmul 프롤로그에서 두 번째 가중치 청크 주소 수정

• hexagon: HMX matmul에서의 F16 활성화 및 불규칙(ragged) N 처리

• hex-mm: mdev 케이스에서 불규칙/분할 검사 강화

• hex-mm: F16 활성화를 위한 MM 융합 활성화

• hex-mm: 융합된 경로에 타일 크기(tiled sizes)를 솔버로 전달

• hex-mmid: expert 매핑 루프에서 스칼라 나눗셈 제거

• hex-mmid: mdev 분할에 대한 적절한 캐시라인 안전성 강제

• hex-mm: mdev 분할 시나리오 및 꼬리 처리(tail handling)를 위한 솔버 개선

• hex-mm: 중복 검사 제거

• hex-mm: 융합된 HMX MUL_MAT_NX가 양자화 가중치에 대해 최종 부분 타일(final partial tile)을 누락하는 문제 수정

• hex-mm: 불규칙한 형태(ragged shapes) 처리 개선 (vtcm의 스칼라 memset 제거)

공동 작성자:
ebateni [email protected]
공동 작성자:
Jhen-Jie Hong [email protected]
공동 작성자:
Yiwei Shao [email protected]