본문으로 건너뛰기

[flashinfer] FlashInfer PrimTS, Sage Attention 도입으로 8비트 양자화 및 성능 최적화 달성

PR 링크: flashinfer-ai/flashinfer#5123 상태: Merged | 변경: +7212 / -2174

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 증가시키면서, 추론 시 메모리 대역폭과 연산 효율성을 극대화하는 것이 중요한 과제가 되었습니다. 특히, 어텐션 메커니즘은 LLM의 핵심 구성 요소이지만 상당한 계산량과 메모리 사용량을 차지합니다. 이러한 문제를 해결하기 위해 FlashInfer는 지속적으로 최적화된 커널을 개발해왔습니다. 이번 PR은 FlashInfer의 PrimTS(Task-Scheduled) 블록 희소(Block Sparse) 커널에 Sage Attention을 도입하여, 8비트 양자화(Quantization)를 지원하고 추론 성능을 크게 향상시키는 것을 목표로 합니다.

이 PR은 특히 TensorRT-LLM의 sageQuant 레이아웃을 활용하여 8비트 양자화된 Q/K/V 데이터를 효율적으로 처리할 수 있도록 합니다. 또한, 기존의 블록 희소 디코드 커널에 조밀 모드(Dense Mode)와 블록 희소 모드 모두에서 Sage Attention을 통합하며, 이를 통해 다양한 시나리오에서의 성능 향상을 기대할 수 있습니다.

코드 분석

이번 PR은 주로 flashinfer/attention/prims_ts/ 디렉토리 내의 여러 파일에 걸쳐 변경 사항을 포함하고 있습니다. 핵심 변경 사항을 파일별로 살펴보겠습니다.

1. flashinfer/attention/prims_ts/README.md 및 __init__.py

  • 변경 사항: 새로운 SageAttentionConfig 및 SageAttentionParams 클래스를 API로 노출하고, README에 Sage Attention에 대한 상세한 설명과 사용 예제를 추가했습니다.

  • 코드 인용 (README.md):

## Sage attention

Sage attention runs the contiguous decode kernel on 8-bit Q, K and V with
dequantization scales, in the dense and block-sparse modes of
`BlockSparseTSWrapper` and `block_sparse_attention`. The plan fixes the recipe
with `sage_config=SageAttentionConfig(...)`; every run supplies the scales as
`sage=SageAttentionParams(...)`:

```text
S[r][c] = sfQ[blkQ(r)] * sfK[blkK(c)] * (Q8 . K8^T)[r][c]
P       = softmax_row(S), quantized to E4M3 as 448 * p
O[r][d] = sfV[d] * (sum_c P8[r][c] * V8[c][d]) / (448 * l[r])  (+ v_mean[d])

FlashInfer does not quantize: the 8-bit tensors and scales come from the caller, for example TensorRT-LLM's sageQuant. v_mean adds a per-channel mean back after normalization, for callers that quantized V - v_mean. Every scale must be finite and positive; the kernel does not check the values.


- **설명**: README 파일은 Sage Attention의 작동 방식, 지원되는 데이터 타입, 파라미터(`SageAttentionConfig`, `SageAttentionParams`) 및 수치적 연산 과정을 명확하게 설명합니다. 특히, Q, K, V의 8비트 양자화와 스케일 팩터(`sfQ`, `sfK`, `sfV`)를 활용하여 연산을 수행함을 보여줍니다. `__init__.py`는 이러한 새로운 클래스들을 외부에서 접근 가능하도록 export 합니다.

### 2. `flashinfer/attention/prims_ts/_block_sparse/common.py`

- **변경 사항**: 조밀 모드(dense mode) 계획 입력에 대한 유효성 검사 함수 `_validate_dense_contiguous_plan_inputs`가 추가되었습니다.

- **코드 인용**: 
```python
def _validate_dense_contiguous_plan_inputs(
    inputs: Iterable[tuple[str, object, object]],
) -> None:
    """Reject block-sparse-only inputs given to a dense contiguous plan.

    Each entry is ``(name, value, unused value)``.
    """

    for name, value, unused in inputs:
        if value != unused:
            raise ValueError(f"{name} is unsupported by a dense contiguous plan")
  • 설명: 이 함수는 use_block_sparse=False로 설정된 조밀 모드에서 블록 희소 관련 파라미터가 전달되는 것을 방지합니다. 이는 API 사용의 일관성을 유지하고 예상치 못한 동작을 방지하는 데 중요합니다.

3. flashinfer/attention/prims_ts/kernels/fmha_decode/fmha_decode_config.py

  • 변경 사항: Sage Attention을 위한 새로운 레시피 규칙이 추가되었고, INT8 Q/K는 SM100에서만 지원하도록 제한하는 검증 로직이 추가되었습니다.

  • 코드 인용 (개념적):

# Sage recipe rules added alongside mixed-precision rules
# Check for INT8 Q/K limited to SM100
if cfg.q_data_type == torch.int8 or cfg.k_data_type == torch.int8:
    if not is_sm100_or_later(cfg.arch):
        raise ValueError("INT8 Q/K Sage attention requires SM100 or later")
  • 설명: Sage Attention은 FP8 또는 INT8 데이터 타입을 지원하지만, INT8 연산은 특정 아키텍처(SM100 이상)에서만 가능합니다. 이 변경은 하드웨어 제약 조건을 반영하여 올바른 GPU에서만 INT8 Sage Attention이 활성화되도록 보장합니다. 또한, SMEM(Shared Memory) 레이아웃 검증 로직이 강화되어, Sage Attention 사용 시 SMEM 예산을 초과하는 구성을 방지합니다. 리뷰어 qsang-nv의 지적에 따라, SMEM 예산 검증이 더 엄격해지고 Sage 스케일 링(scale ring)까지 고려하게 되었습니다.

4. flashinfer/attention/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py 및 관련 파일

  • 변경 사항: Sage Attention 관련 스케일(scale) 텐서 처리 로직이 커널 레벨에서 구현되었습니다. Q, K, V 스케일이 소프트맥스 연산 전후에 적용되며, V 스케일은 에필로그(epilogue)에서 채널별로 적용됩니다. INT8 점수는 FP32로 누적되어 소프트맥스 연산에 사용됩니다.

  • 코드 인용 (개념적):

# Scales enter softmax before row max and exponential
# P is E4M3 with the static 448 scale
# V scales apply per channel in the epilogue

# INT8 scores accumulate onto an FP32 bias, so the softmax reads them without conversion.

  • 설명: 커널 레벨에서의 최적화는 Sage Attention의 핵심입니다. Q와 K 스케일은 어텐션 스코어 계산에 영향을 미치고, V 스케일은 최종 출력 계산 시 적용됩니다. 특히 INT8 연산의 경우, 정밀도 손실을 최소화하기 위해 FP32로 누적하는 방식이 사용됩니다. 이는 8비트 양자화의 이점을 최대한 활용하면서도 정확도를 유지하는 중요한 기법입니다.

  • 리뷰 반영: qsang-nv의 지적에 따라, Sage 스케일 리소스가 태스크의 의존성 그래프에 명시적으로 포함되도록 수정되었습니다. 이는 TaskManager가 리소스 의존성을 올바르게 관리하도록 하여 잠재적인 데이터 경쟁이나 데드락을 방지합니다.

5. flashinfer/attention/prims_ts/sage.py

  • 변경 사항: SageAttentionConfig와 SageAttentionParams 클래스가 정의되었습니다. SageAttentionConfig는 계획(plan) 시점에 Sage Attention의 구성을 정의하고, SageAttentionParams는 실행(run) 시점에 실제 스케일 텐서들을 전달합니다.

  • 코드 인용:

class SageAttentionConfig:
    def __init__(self, q_block_size: int = 1, k_block_size: int = 16, ...):
        # ... configuration parameters ...

class SageAttentionParams:
    def __init__(self, q_scale: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor, v_mean: Optional[torch.Tensor] = None, k_summary_scale: Optional[torch.Tensor] = None):
        # ... scale tensors ...
  • 설명: 이 클래스들은 Sage Attention의 설정을 구조화하고, 실행 시 필요한 동적인 스케일 정보를 전달하는 인터페이스를 제공합니다. q_block_size, k_block_size 등은 TensorRT-LLM의 sageQuant 레이아웃과 호환되도록 설계되었습니다. v_mean과 k_summary_scale은 프록시 경로(proxy routes)와 같은 고급 기능에 사용됩니다.

6. 테스트 추가

  • 변경 사항: Sage Attention 및 조밀 모드에 대한 새로운 테스트 케이스가 추가되었습니다 (tests/attention/test_attention_ts_sage.py, tests/attention/test_attention_ts_block_sparse.py 등).

  • 설명: 새로운 기능의 정확성과 안정성을 보장하기 위해 철저한 테스트가 필수적입니다. 이 테스트들은 Sage Attention의 다양한 시나리오(FP8, INT8, 조밀 모드, 블록 희소 모드)를 검증합니다.

왜 이게 좋은가?

이번 PR은 여러 측면에서 중요한 최적화 및 개선을 이루었습니다.

  1. 메모리 대역폭 절감 및 연산 효율성 증대: 8비트 양자화(FP8, INT8)를 지원함으로써, 기존의 FP16/BF16 연산 대비 메모리 사용량을 절반으로 줄이고, 특정 하드웨어에서는 연산 속도를 향상시킬 수 있습니다. 이는 LLM 추론 시 병목 현상인 메모리 대역폭을 완화하는 데 크게 기여합니다.

  2. 성능 향상: PR 설명에 제시된 성능 측정 결과에 따르면, Sage Attention은 다양한 시나리오에서 BF16 대비 상당한 속도 향상을 보여줍니다. 예를 들어, B200 GPU에서:

    • Dense S=10800 H=40: BF16 대비 FP8 Sage가 0.88배, INT8 Sage가 0.89배의 속도를 보였습니다.
    • Block-sparse exact S=15360 H=40, density 0.125: BF16 대비 FP8 Sage가 0.81배, INT8 Sage가 0.82배의 속도를 보였습니다.
    • Block-sparse exact S=75600 H=40, density 0.10: BF16 대비 FP8 Sage가 0.79배, INT8 Sage가 0.76배의 속도를 보였습니다.

    이는 특히 긴 시퀀스 길이(S)와 높은 헤드 수(H)에서 두드러집니다.

  3. 유연성 증가: 조밀 모드(Dense Mode)와 블록 희소 모드 모두에서 Sage Attention을 지원하며, SageAttentionConfig를 통해 다양한 블록 크기(q_block_size, k_block_size) 및 프록시 경로 옵션을 설정할 수 있습니다. 이는 다양한 모델 아키텍처와 하드웨어 구성에 대한 적용 가능성을 높입니다.

  4. TensorRT-LLM과의 호환성: TensorRT-LLM의 sageQuant 레이아웃을 직접 활용함으로써, 기존 TensorRT-LLM 사용자들은 큰 변경 없이 FlashInfer의 최적화된 커널을 사용할 수 있게 됩니다.

  5. 코드 품질 및 안정성 향상: 리뷰어들의 피드백을 반영하여 SMEM 예산 검증, INT8 지원 하드웨어 제한, 태스크 의존성 그래프 명확화 등 코드의 안정성과 견고성을 높이는 수정이 이루어졌습니다. 이는 복잡한 커널 구현에서 발생할 수 있는 잠재적 오류를 줄이는 데 기여합니다.

일반적인 교훈

이 PR은 LLM 추론 성능 최적화를 위한 몇 가지 중요한 교훈을 제공합니다:

  • 양자화의 중요성: FP8/INT8 양자화는 메모리 대역폭 병목 현상을 완화하고 연산 효율성을 높이는 강력한 도구입니다. 하드웨어 지원이 뒷받침된다면 적극적으로 활용해야 합니다.
  • 커널 레벨 최적화: 어텐션 메커니즘과 같이 계산 집약적인 부분은 커널 레벨에서의 정교한 최적화(스케일 적용, 데이터 타입 관리, SMEM 활용 등)를 통해 큰 성능 향상을 얻을 수 있습니다.
  • API 설계의 유연성: 다양한 하드웨어 및 모델 요구사항을 충족시키기 위해, SageAttentionConfig와 같이 유연한 설정 옵션을 제공하는 API 설계가 중요합니다.
  • 철저한 검증 및 리뷰: 복잡한 커널 최적화 작업에서는 리뷰어의 피드백을 적극적으로 반영하여 SMEM 예산, 하드웨어 제약 조건, 태스크 의존성 등 잠재적인 문제를 해결하는 것이 필수적입니다.

References

  • torch.compile - This PR does not directly use torch.compile, but it is a general reference for PyTorch optimization. (Note: The PR itself focuses on custom CUDA kernels, not PyTorch's JIT compilation.)
  • TensorRT-LLM Sage Attention - The PR leverages concepts and layouts from TensorRT-LLM's Sage Attention.
  • FlashInfer PrimTS Documentation - The PR adds Sage Attention to the PrimTS block sparse kernels.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글