[flashinfer] SM120 VSA 커널의 마지막 Partial Wave 패킹: 호스트 플래너만으로 평균 4.3% 가속

PR 링크: flashinfer-ai/flashinfer#8829 상태: Merged | 변경: +359 / -57

들어가며

FlashInfer의 BlockSparseAttentionWrapper(backend="vsa_sm120_blk64")는 SM120/SM121(RTX PRO 6000 Blackwell 등)용 Cake block-sparse attention 커널입니다. 두 개의 persistent kernel이 같은 host plan을 공유합니다.

• vsa_sm120_blk64: SM당 128-thread CTA 2개
• vsa_sm120_blk64_pair: SM당 4-warp group 2개로 구성된 CTA 1개

어느 쪽이든 SM마다 가상 CTA(virtual CTA)가 2개이며, 작업 단위는 (head, query block) 쌍인 "row"입니다. 기존 플래너는 모든 row를 동일한 part 크기로 나누는 uniform plan을 만들고 static grid-stride 방식으로 가상 CTA에 배분했습니다.

문제는 마지막 partial wave입니다. 188 SM × 2 = 376개 가상 CTA에 item이 예컨대 500개라면, 두 번째 wave에는 124개만 남아 나머지 252개 가상 CTA가 놀게 됩니다. row가 가상 CTA당 3개 미만일 정도로 적으면 이 tail이 전체 시간에서 차지하는 비중이 큽니다. 이 PR은 커널·API·numerics는 그대로 두고 호스트 플래너만 수정해 이 tail을 메웁니다.

코드 분석

모든 변경은 flashinfer/cake_vsa_sm120.py에 있습니다.

1. 패킹 상수: 언제 패킹할 것인가

_PACK_COST = {"vsa_sm120_blk64": (3, 8), "vsa_sm120_blk64_pair": (5, 9)}
_PACK_SPLIT_COST = 1
_PACK_ROWS = 3
_PACK_KV_BLOCKS = (1024, 4096)
_MIN_PIECE = 2

• _PACK_ROWS = 3: 가상 CTA당 row가 3개 미만일 때만 패킹을 고려합니다. wave가 많으면 tail 비중이 작아 이득이 없기 때문입니다.
• _PACK_KV_BLOCKS = (1024, 4096): 64×64 블록, head dim 128, 16-bit 기준으로 4096 블록 = 128MB(RTX PRO 6000 L2 전체), 1024 블록 = 32MB(L2의 1/4)입니다. 모든 KV head의 K/V가 이를 넘으면 패킹하지 않습니다. row 개수가 가상 CTA보다 적으면 더 엄격한 32MB 한도를 씁니다.

왜 L2가 기준일까요? 패킹은 row를 여러 가상 CTA에 걸쳐 자르므로, 동시에 실행되는 row 조합이 uniform plan과 달라지고 KV 접근의 지역성이 깨집니다. working set이 L2에 들어가면 이 비용이 거의 없지만, 넘어가면 DRAM 트래픽이 늘어 오히려 느려진다는 것이 실측 결과입니다. 주석에 따르면 이 상수들은 118개 selection에 fit되었고, 선택된 plan은 uniform 대비 6.7% 빠르며 네 후보(uniform/packed × 두 커널) 중 최선과 비교해 0.24%만 느립니다.

2. build_work 리팩터링: _table 분리

Before에는 build_work 안에서 split·정렬·int32 변환까지 한 번에 처리했습니다.

# Before (build_work 내부)
item_order = torch.sort(-(2 * size + is_split.long()), stable=True).indices
work = work[item_order].to(torch.int32).contiguous()
groups = int(split.sum())
slots = int(split_parts.sum())

After에는 uniform table 생성을 _table로 분리해, 그 결과를 패킹의 입력으로 재사용합니다.

# After
def _table(row_count, row_first, row_head, row_block, part):
    """The uniform work table of the rows split at ``part`` ..."""
    ...
    return work[item_order], int(split.sum()), int(split_parts.sum())

이 분리 덕분에 "uniform plan을 만든 뒤 → 패킹 후보를 만들고 → 비용 모델로 비교"하는 흐름이 자연스러워졌습니다.

3. _pack: full wave는 고정, tail만 재배치

핵심 알고리즘입니다. 먼저 full wave의 item은 grid-stride 위치(p % resident) 그대로 둡니다.

held = len(work) // resident * resident
lists = [[] for _ in range(resident)]
for p in range(held):
    w = work[p]
    lists[p % resident].append([w[0], w[1], w[2], w[3], w[4] >= 0])
tail = [[w[0], w[1], w[2], w[3], w[4] >= 0] for w in work[held:]]

full wave를 건드리지 않는 이유는 docstring에 명시되어 있습니다: "so that the rows in flight together stay those of the uniform plan". 동시에 실행되는 row 집합을 유지해 L2 지역성을 보존하려는 의도입니다.

그다음 tail item을 가상 CTA 0, 1, 2…에 순서대로 채우는 fill(capacity)를 정의하고, 이를 이진 탐색으로 감쌉니다.

piece = capacity - used[v] - item_cost - _SPLIT_COST
if piece >= _MIN_PIECE and rest - piece >= 1:
    out[v].append([head, block, begin, begin + piece, True])
    begin += piece
    cut = True
v += 1

total = sum(base) + sum(cost(e) for e in tail)
low, high = max(max(base), -(-total // resident)), max(base) + total
...
while low < high:
    mid = (low + high) // 2
    got = fill(mid)
    if got is None:
        low = mid + 1
    else:
        high, best = mid, got

가상 CTA가 가득 차는 지점에서 row를 임의 크기로 자릅니다. 기존 uniform plan이 _PART_DIVISORS = (1, 2, 3, 4, 6, ...)처럼 정해진 분할만 허용한 것과 대조적입니다. 하한은 "이미 가장 무거운 CTA의 부하"와 "전체 평균" 중 큰 값이므로, 이진 탐색은 tail 전체를 담을 수 있는 최소 공통 부하(makespan)를 찾습니다. 자를 때는 _SPLIT_COST를 비용에 포함하고, _MIN_PIECE = 2 블록 미만의 조각은 만들지 않아 merge 오버헤드가 이득을 잠식하지 않도록 합니다.

마지막으로 grid-stride가 요구하는 "리스트 길이가 L+1 또는 L" 조건을 맞추는 단계가 있습니다. 두 개 이상 짧은 리스트는 가장 큰 조각을 반으로 쪼개거나, 가장 긴 리스트의 마지막 item을 가져옵니다. 이후 split row의 조각들을 block index 순서로 정렬해 group/slot 번호를 부여합니다.

for ps in sorted(pieces.values(), key=lambda ps: min(e[2] for e in ps)):
    ps.sort(key=lambda e: e[2])

이 덕분에 merge가 항상 part 순서로 이루어지고, 출력이 run-to-run deterministic하게 유지됩니다.

4. _pack_time: SM 공유를 반영한 비용 모델

sms = resident // _VIRTUAL_PER_SM
span = max(
    10 * min(a, b) + lone_rate * abs(a - b)
    for a, b in zip(load[:sms], load[sms:], strict=True)
)

한 SM의 두 가상 CTA가 함께 돌 때는 rate 10(tenths), 하나가 먼저 끝나 혼자 남으면 lone_rate(8 또는 9)로 계산합니다. 혼자 남은 CTA는 SM 자원을 독점해 더 빨리 진행된다는 하드웨어 특성을 반영한 것입니다. 각 커널의 packed plan은 이 모델에서 uniform보다 빠를 때만 채택되며, 일부 shape에서는 같은 work table로 다른 커널을 고르는 결과가 나오기도 합니다(std_s2048_d50, 1.076).

왜 이게 좋은가

성능

RTX PRO 6000 (188 SMs), cold L2, CUPTI 측정 기준입니다.

범위	Geomean speedup
#8799의 22개 shape	1.028
그 중 packed 7개	1.078
추가 26개 shape	1.059
전체 100개 shape	1.043 (최저 0.970)

shared_h8_s3072_k18은 1.259, prefix_h6_m6144_p4_k32는 1.228까지 빨라졌습니다. 반면 wan14b_720p처럼 wave가 많은 대형 shape은 plan이 unchanged로 남아 회귀가 없습니다. 이게 바로 _PACK_ROWS와 _PACK_KV_BLOCKS 가드의 효과입니다.

리뷰에서 보강된 부분

리뷰 과정에서 테스트가 강화되었습니다. split row의 partial slot이 서로 겹치지 않고 다른 row와도 disjoint한지 검증하고, cross-kernel 테스트에서 pack=True와 pack=False table을 비교하며, 각 커널의 packed table을 FP32 reference와 BF16/FP16으로 대조하는 테스트가 추가되었습니다. 임의 크기 분할은 workspace slot 충돌 같은 버그가 숨기 쉬운 영역이라 이런 불변식 테스트가 특히 중요합니다.

일반적 교훈

1. Persistent kernel의 tail effect는 호스트에서 해결할 수 있다. 커널을 고치지 않아도 work table만 바꿔 wave quantization 손실을 줄일 수 있습니다.
2. 바꿀 것만 바꾼다. full wave는 그대로 두고 tail만 재배치해 캐시 지역성과 기존 동작을 보존했습니다.
3. 분할 비용을 모델에 넣고, 하드웨어 한계로 가드한다. split/merge 비용, SM 공유 rate, L2 크기를 반영하지 않으면 load balancing이 오히려 손해가 됩니다.
4. 결정성을 지켜라. 조각을 index 순서로 merge해 numerics contract를 깨지 않았습니다.

참고 자료

• FlashInfer Documentation — BlockSparseAttentionWrapper를 포함한 FlashInfer 공식 문서
• torch.sort — uniform work table의 stable 정렬에 사용된 함수
• torch.cumsum — split group 번호 부여에 사용된 함수

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