[sglang] EAGLE 트리 검증 어텐션 최적화: Triton 안정 압축과 FA4 고정 길이 쿼리 레이아웃

PR 링크: sgl-project/sglang#49477 상태: Merged | 변경: +536 / -140

들어가며

EAGLE 계열 speculative decoding에서 topk > 1이면 draft 토큰은 체인이 아니라 트리 모양이 됩니다. target 모델은 이 트리 전체(예: W=128 토큰)를 한 번의 forward로 검증하는데, 각 query 토큰이 실제로 attend해야 하는 draft 토큰은 자기 조상 체인뿐입니다. 트리는 넓지만(W=64, 128) 깊이는 얕습니다(steps=8).

SGLang의 FA3/FA4 백엔드는 이를 cascade attention으로 처리합니다.

1. prefix attention: 모든 query가 공유하는 이전 KV 캐시에 대한 attention
2. suffix attention: 각 query가 자기 조상 draft 토큰에만 attend
3. 두 결과를 LSE(log-sum-exp)로 merge

기존 구현에는 세 가지 비효율이 있었습니다.

• suffix용 page table을 만들 때 mask 추출 → torch.where → torch.sort → gather로 이어지는 여러 PyTorch 커널을 실행했습니다.
• suffix page table의 폭이 트리 전체 폭(W)이라, 깊이가 8인 트리도 128칸짜리 table을 봤습니다.
• FA4에서 query 1개짜리 suffix 시퀀스 수백 개를 varlen scheduler로 처리해 스케줄링 오버헤드가 컸습니다.

이 PR은 세 문제를 모두 다룹니다.

코드 분석

1. metadata.py: Triton 기반 안정 압축(stable compaction)

Before — flashattention_backend.py에서 PyTorch op을 여러 개 이어 붙였습니다.

mask = forward_batch.spec_info.custom_mask[
    mask_extraction_indices
].view(-1, self.speculative_num_draft_tokens)
...
keys = torch.where(
    mask, col_indices, col_indices + self.speculative_num_draft_tokens
)
_, sort_order = torch.sort(keys, dim=1)
non_masked_page_table = (
    self.req_to_token_pool.req_to_token[forward_batch.req_pool_indices, :]
    .gather(1, cols)
    .repeat_interleave(self.speculative_num_draft_tokens, dim=0)
)
metadata_expand.page_table = non_masked_page_table.gather(1, sort_order)
metadata_expand.cache_seqlens_int32 = mask.sum(dim=1).to(torch.int32)

이 코드는 visible한 슬롯을 앞으로 모으려고 정렬을 씁니다. 그런데 키가 "visible이면 원래 인덱스, 아니면 +W"라서 결과는 사실상 안정 파티션(stable partition)입니다. O(W log W) 정렬까지 갈 필요가 없는 문제입니다. 여기에 arange, cumsum, repeat_interleave, 인덱싱, gather 2회, sum까지 더해 커널 런치가 10개를 넘고, 중간 텐서도 매번 할당됩니다.

After — query row당 program 하나를 띄우는 Triton 커널 하나로 바꿨습니다.

valid = tl.load(custom_mask + mask_offset, cols < WIDTH, 0).to(tl.int32)
rank = tl.cumsum(valid)
count = tl.sum(valid, 0)
# Stable partition: visible slots first, preserving order in both groups.
destination = tl.where(valid != 0, rank - 1, count + cols - rank)
slots = tl.load(req_to_token + req * req_stride + seq_len + cols, cols < WIDTH, 0)
tl.store(page_table + row * table_stride + destination, slots, cols < WIDTH)
tl.store(cache_seqlens + row, count)

동작 방식은 이렇습니다.

• prefix sum(tl.cumsum)으로 각 visible 원소의 목적지는 rank - 1입니다.
• invisible 원소의 목적지는 count + (cols - rank)이고, 여기서 cols - rank는 자기보다 앞선 invisible 원소의 개수입니다.
• 정렬 없이 O(W) scatter로 두 그룹의 상대 순서를 모두 유지합니다. 고전적인 scan 기반 stream compaction 패턴입니다.

CUDA Graph를 고려한 처리도 들어 있습니다. 그래프 padding row는 입력 mask에 대응하는 행이 없으므로 tl.where(row < num_actual_tokens, mask_offset, 0)로 offset을 0으로 고정해 out-of-bounds 읽기를 막습니다. 또 PR 설명에 따르면 이렇게 만든 메타데이터는 레이어마다 다시 만들지 않고 모든 레이어에서 재사용합니다.

2. flashattention_backend.py: suffix table을 조상 깊이로 자르기

suffix_page_table = suffix_metadata.page_table
if forward_batch.spec_algorithm.is_eagle():
    # Each query sees its ancestors, not the whole tree.
    suffix_page_table = suffix_page_table[
        :, : forward_batch.spec_info.max_tree_depth
    ]

앞 단계의 압축 덕분에 visible 슬롯은 항상 row 앞쪽에 모여 있습니다. 따라서 max_tree_depth 이후 열은 볼 필요가 없습니다. W=128이어도 깊이가 8이면 table view는 [N, 8]이 됩니다. 이 슬라이싱은 view라 복사 비용이 없습니다. 압축이 슬라이싱을 가능하게 만드는 구조입니다.

3. 기존 fused suffix-attention/merge 커널을 EAGLE에도 적용

Before: fused 경로는 UNO에서만 쓰였습니다.

use_cascade_attn
and forward_batch.spec_algorithm.is_uno()
and can_use_fused_suffix_attention_merge(...)

After: suffix 폭이 작으면 EAGLE도 같은 경로를 탑니다.

and (
    forward_batch.spec_algorithm.is_uno()
    or (
        forward_batch.spec_algorithm.is_eagle()
        and suffix_page_table.shape[1] <= 64
        and layer.head_dim <= 128
    )
)

깊이 ≤ 64로 잘린 table이라면 suffix attention과 LSE merge를 커널 하나에서 끝낼 수 있습니다. 새 커널을 만들지 않고 기존 커널의 적용 조건만 넓힌 점이 깔끔합니다. 이를 위해 q = q.contiguous().view(...)처럼 contiguous Q를 변수에 다시 대입하는 리팩토링도 함께 들어갔습니다. prefix와 suffix 단계가 같은 연속 메모리 Q를 공유하게 하려는 것입니다.

4. FA4 고정 길이 쿼리 레이아웃

fixed_length = ver == 4 and not kwargs
result = self.flash_attn_with_kvcache(
    q=q.unsqueeze(1) if fixed_length else q,
    ...
    cu_seqlens_q=None if fixed_length else metadata.cu_seqlens_q,
    cu_seqlens_k_new=None if fixed_length else metadata.cu_seqlens_k,
    max_seqlen_q=1,
    ...
)
if fixed_length:
    output, lse, *rest = result
    return output.squeeze(1), lse.squeeze(-1).T, *rest

suffix 시퀀스는 모두 query가 정확히 1개입니다. 그런데 cu_seqlens_q를 넘기면 FA4는 varlen scheduler를 씁니다. 대신 Q를 [N, 1, H, D] batch 레이아웃으로 바꾸면 SM90에서 직접 tile scheduling이 선택됩니다. KV 길이는 cache_seqlens로 여전히 가변입니다. 끝에서 squeeze와 .T로 output과 LSE를 varlen 레이아웃([H, N])으로 되돌려 merge 코드와 호환되게 합니다. modifier나 scale tensor가 kwargs로 들어오면 기존 varlen 레이아웃을 유지하는 보수적인 fallback도 있습니다.

리뷰 논의

리뷰어 BBuf는 다음 문제를 지적했습니다. "기존 FA4 verify 테스트는 깊이가 65라서 새 fused 경로(≤64)를 우회합니다. fused 케이스는 모두 FA3이고 Blackwell에서는 skip됩니다." 즉 새 FA4 경로는 실제로 테스트되지 않고 있었습니다. 작성자는 얕은 분기 트리로 fused dispatch를 확인하는 test_fa4_verify_fused를 추가했습니다. 이 테스트는 BF16 출력을 FP32 reference와 비교하며, eager 모드와 입력을 갱신한 CUDA Graph replay 두 경우 모두에서 B300(SM103)으로 검증했습니다. dispatch 조건이 까다로운 최적화일수록 "경로가 실제로 실행되는가"를 테스트로 고정하는 것이 중요하다는 좋은 사례입니다.

왜 이게 좋은가

H800, Qwen3-8B + EAGLE3로 측정한 54개 트리 구성 결과입니다.

Backend	TPS Δ 중앙값 (범위)	E2E p50 Δ	E2E p95 Δ
FA3	+6.25% (+1.44% ~ +11.30%)	-5.79%	-7.30%
FA4	+16.96% (+2.61% ~ +342.54%)	-14.24%	-15.54%

• FA4, W128/S8/K16, input 128, C32 구성에서는 TPS가 164.52 → 728.29(+342%) 로 올랐습니다. 넓은 트리와 높은 동시성에서 varlen 스케줄링과 메타데이터 오버헤드가 지배적이었다는 뜻입니다.
• standalone draft attention 기준 FA4는 중앙값 1.152×, 최대 2.333×였고, FA3는 1.000×로 회귀가 없었습니다.
• B300에서도 FA4가 최대 +156.80%였습니다.
• spec 비활성화나 chain spec(topk=1)에서는 ±0.5% 안팎이라 영향이 없었습니다. 유일한 회귀는 KV 압박이 있던 FA3/page64 한 구성의 p95 +2.61%입니다.

일반적 교훈

1. sort로 partition을 흉내 내지 말 것: 키가 2개 그룹뿐이라면 scan 기반 stable compaction(O(N))이 정렬보다 빠르고 단순합니다.
2. 작은 텐서 연산 체인은 커널 하나로 합칠 것: decode 단계에서는 연산량보다 커널 런치 수와 중간 할당이 병목입니다.
3. 데이터 레이아웃 정리가 다음 최적화를 연다: 압축 덕분에 깊이 슬라이싱이 가능해졌고, 슬라이싱 덕분에 fused 커널 조건(≤64)을 만족하게 됐습니다.
4. 커널 API가 고르는 스케줄러를 의식할 것: 같은 연산이라도 [N,1,H,D]처럼 shape만 바꿔 varlen보다 효율적인 fixed-length 경로를 탈 수 있습니다.

참고 자료

• triton.language.cumsum — 안정 압축 커널에서 rank 계산에 사용한 prefix sum
• triton.language.where — 목적지 인덱스 선택에 사용
• triton.next_power_of_2 — BLOCK 크기 계산
• torch.sort — 기존 구현에서 제거된 정렬 연산
• torch.gather — 기존 구현에서 제거된 gather 연산
• FlashAttention GitHub — flash_attn_with_kvcache / varlen API 원본 구현

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