[flashinfer] FlashInfer, CuTe DSL W4A16 MoE 최적화 및 비트 단위 정확도 테스트 추가

2026년 10월 27일수정: 2026년 10월 27일

PR 링크: flashinfer-ai/flashinfer#7196 상태: Merged | 변경: +1654 / -670

들어가며

최근 AI 모델, 특히 거대 언어 모델(LLM)의 발전은 Mixture-of-Experts (MoE) 아키텍처의 중요성을 더욱 부각시키고 있습니다. MoE는 모델의 파라미터 수를 크게 늘리면서도 추론 시 활성화되는 파라미터 수를 제한하여 효율성을 높이는 기법입니다. FlashInfer는 이러한 MoE 모델의 추론 성능을 극대화하기 위한 라이브러리로, 이번 PR(#7196)은 특히 W4A16 (4-bit Weight, 16-bit Activation) 양자화된 MoE 모델의 성능을 크게 향상시키고, 정확도를 더욱 엄격하게 검증하는 데 초점을 맞추고 있습니다.

본 블로그 글에서는 이 PR이 해결하고자 하는 문제점, 코드 변경 사항의 상세 분석, 그리고 이러한 최적화가 왜 효과적인지에 대해 실제 코드 diff를 인용하며 자세히 설명하고자 합니다.

코드 변경 분석

이번 PR은 주로 flashinfer/fused_moe/cute_dsl/blackwell/ 디렉토리 내의 파일들을 수정하여 W4A16 MoE의 성능과 정확성을 개선했습니다. 핵심 변경 사항은 다음과 같습니다.

1. `moe_finalize.py`: 연산 정밀도 및 순서 최적화

이 파일에서는 MoE 레이어의 최종 연산(finalize) 부분을 수정하여 정확도를 높이고 성능을 개선했습니다. 특히, 부동 소수점 연산의 정밀도와 순서를 조정했습니다.

변경 전:

@dsl_user_op
def _fma_rn(a, b, c, *, loc=None, ip=None):
    # Match the native finalizer's --use_fast_math FP32 FMA, including FTZ.
    return cutlass.Float32(
        llvm.inline_asm(
            T.f32(),
            [cutlass.Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b, c)],
            "fma.rn.ftz.f32 $0, $1, $2, $3;",
            "=f,f,f,f",
            has_side_effects=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
            loc=loc,
            ip=ip,
        )
    )

변경 후:

@dsl_user_op
def _fma_rn(a, b, c, *, loc=None, ip=None):
    # Match the native finalizer's --use_fast_math FP32 FMA, including FTZ.
    return cutlass.Float32(
        llvm.inline_asm(
            T.f32(),
            [cutlass.Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b, c)],
            "fma.rn.f32 $0, $1, $2, $3;", # FTZ 플래그 제거
            "=f,f,f,f",
            has_side_effects=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
            loc=loc,
            ip=ip,
        )
    )

@dsl_user_op
def _mul_rn(a, b, loc=None, ip=None):
    # A separate first multiply preserves signed zero, as in MegaMoE combine.
    return cutlass.Float32(
        llvm.inline_asm(
            T.f32(),
            [cutlass.Float32(x).ir_value(loc=loc, ip=ip) for x in (a, b)],
            "mul.rn.f32 $0, $1, $2;",
            "=f,f,f",
            has_side_effects=False,
            asm_dialect=llvm.AsmDialect.AD_ATT,
            loc=loc,
            ip=ip,
        )
    )

_fma_rn 함수에서 ftz (flush-to-zero) 플래그가 제거되었습니다. 이는 subnormal 값(매우 작은 값)을 0으로 만들지 않고 유지하여 연산의 정밀도를 높이기 위함입니다. 또한, _mul_rn 함수가 새로 추가되었습니다. 이 함수는 첫 번째 곱셈 연산을 수행하며, MegaMoE의 combine 연산과 유사하게 signed zero (부호 있는 0)를 보존하는 데 사용됩니다. 이후 연산은 _fma_rn을 통해 수행됩니다.

변경 후 moe_unpermute 함수 내부 로직:

                        if cutlass.const_expr(slot == 0):
                            accum[element, vector] = _mul_rn(
                                cutlass.Float32(values[element]), weights[slot]
                            )
                        else:
                            accum[element, vector] = _fma_rn(
                                cutlass.Float32(values[element]),
                                weights[slot],
                                accum[element, vector],
                            )

이 변경은 첫 번째 가중치(slot == 0)에 대해서는 _mul_rn을 사용하여 곱셈을 수행하고, 이후 가중치에 대해서는 _fma_rn을 사용하여 FMA(Fused Multiply-Add) 연산을 수행하도록 합니다. 이는 MegaMoE의 연산 방식과 일치하며, 특히 0으로 수렴하는 경우의 정확도를 높입니다.

2. `moe_w4a16_kernel.py`: 런타임 MMA N 최적화 및 BF16 순서 변경

이 파일에서는 W4A16 MoE 커널의 성능을 향상시키기 위한 여러 최적화가 이루어졌습니다. 특히 런타임 시점에 MMA (Tensor Core)의 N 차원 크기를 동적으로 조절하고, BF16 데이터 처리 방식을 개선했습니다.

주요 변경 사항:

Runtime MMA N: 활성화(activations)가 이미 operand B를 차지하는 상황에서, 실시간으로 라우팅된 토큰 수를 기반으로 MMA N 크기를 8 또는 16의 배수로 조정합니다. 이는 CTA(Thread Block)의 활성화 로드를 런타임 중간 지점으로 이동시키고, 비활성 서브타일을 건너뛰며, 마지막 활성 서브타일 이후에 누산기(accumulator)를 해제하는 방식으로 이루어집니다.
Specialized BF16 permutation: 256개의 스레드를 사용하여 유효한 BF16 행만 복사하고, 영구적인 행 스케줄링을 사용합니다. 256비트 복사는 32바이트 정렬된 행/저장 공간에 사용되며, 16바이트 정렬에는 명시적인 128비트 폴백(fallback)이 사용됩니다.
Arithmetic Alignment: SwiGLU, SiTU 및 W4A4/W4A16 최종 처리기를 MegaMoE와 정렬합니다. Combine 연산은 첫 번째 슬롯에 FP32 곱셈을 수행한 후 순서대로 FP32 FMA를 수행하여 subnormals을 보존합니다.

새로운 함수 및 로직 예시:

@dsl_user_op
def compute_non_leader_cta_load_shift(
    valid_tokens_in_tile,
    mma_tiler_n: int,
    loc: Optional[ir.Location] = None,
    ip: Optional[ir.InsertionPoint] = None,
) -> Int32:
    """Token offset that non-leader CTA's TMA-B read must shift by under 2cta.

    Under dynamic UMMA + 2cta:
      - MMA splits N at align16(valid) / 2
      - TMA static partition splits at mma_tiler_n / 2

    The result ∈ (-mma_tiler_n/2, 0]; apply via
    ``cute.domain_offset((shift, 0, 0), real_b)`` on non-leader CTA only.
    """
    return (_align_n(valid_tokens_in_tile, 16) >> Int32(1)) - Int32(mma_tiler_n // 2)

@dsl_user_op
def issue_dynamic_bf16_mma_tile(
    acc_tensor,
    a_frag_tile,
    b_frag_tile,
    k_tile_idx,
    valid_tokens_in_tile,
    mma_tiler_mnk: tuple,
    n_alignment: int = 16,
    loc=None,
    ip=None,
):
    """Issue BF16 K16 atoms with a runtime instruction N extent.

    The caller owns pipeline waits/commits, TMEM allocation, valid-token
    bounds, and the non-leader TMA-B shift. Only positive valid tiles may
    reach this function, and valid_tokens must not exceed the allocated N.
    One-CTA M128 instructions support N in multiples of 8; two-CTA M256
    instructions require multiples of 16. Each CTA retains its M128 x
    allocated-N accumulator layout: runtime N limits the columns written,
    without moving valid ones.

    ``make_trivial_tiled_mma`` uses a single M/N atom, including N192. Reject
    multi-atom fragments explicitly: their per-atom N tails and two-CTA B
    partitions would need separate descriptors and load shifts.
    """
    # ... (implementation details) ...

compute_non_leader_cta_load_shift 함수는 2개의 CTA가 사용될 때, 리더 CTA가 아닌 CTA의 TMA-B 로드 오프셋을 계산합니다. issue_dynamic_bf16_mma_tile 함수는 런타임 시점에 결정된 N 크기를 사용하여 BF16 MMA 타일을 발행하는 로직을 구현합니다. 이는 GPU 하드웨어의 특성을 최대한 활용하여 효율성을 높입니다.

3. Shared Exact Reference 및 테스트 강화

이 PR은 tests/test_helpers/w4a16_reference.py 파일을 통해 BF16 가중치 디코딩, 스케일 레이아웃 반전, 제어된 FP32 GEMM/알파 스케일링, 활성화 및 순서 결합에 대한 비트 단위 정확도 참조를 공유합니다. 또한, 기존의 근사적인 어설션(assertion)을 제로 톨러런스(zero tolerance)로 교체하여 정확성 검증을 강화했습니다.

변경 전 (개념적):

# Approximate assertion
assert torch.allclose(output, expected_output, atol=1e-3)

변경 후 (개념적):

# Exact assertion with zero tolerance
assert (output.view(torch.int16) == expected_output.view(torch.int16)).all()

이러한 변경은 MoE 및 Dense GEMM 테스트에서 W4A16에 대한 비트 단위 정확도 참조를 공유하고, 기존의 근사적인 검증 대신 엄격한 제로 톨러런스 검증을 적용함으로써 모델의 수치적 일관성을 보장합니다.

왜 이것이 좋은 최적화인가?

성능 향상

PR 설명에 포함된 성능 비교 결과는 이 최적화의 효과를 명확히 보여줍니다.

H7168 full PR (Before / After): 전체적으로 약 1.04배의 성능 향상을 보였습니다. 특히 Decode 단계에서 1.03배, Prefill 단계에서 1.05배의 향상이 있었습니다.
Kimi K3 EP8 full PR (Before / After): 마찬가지로 전체적으로 약 1.04배의 성능 향상을 기록했으며, Prefill 단계에서는 1.12배의 상당한 개선을 보였습니다.
TRTLLM W4A4 vs CuTe W4A16: CuTe W4A16 백엔드가 TRTLLM W4A4보다 약 0.84배 (즉, 1.18배 더 빠름)의 성능을 보였습니다. 이는 W4A16 최적화가 다른 백엔드와 비교해도 경쟁력이 있음을 시사합니다.

이러한 성능 향상은 다음과 같은 요인에 기인합니다:

런타임 MMA N 최적화: GPU 하드웨어의 특성을 실시간으로 파악하여 연산 크기를 동적으로 조절함으로써, 불필요한 연산을 줄이고 데이터 활용도를 높였습니다.
BF16 Permutation 최적화: 데이터 복사 및 로딩 방식을 개선하여 메모리 접근 효율성을 높였습니다.
정밀도 및 연산 순서 개선: subnormal 값 보존 및 signed zero 처리와 같은 미세한 연산 조정은 특히 0에 가까운 값이나 특정 연산에서 수치적 안정성을 높이고, 때로는 성능 향상으로 이어질 수 있습니다.

정확성 및 신뢰성 강화

제로 톨러런스 테스트: 근사적인 비교 대신 비트 단위 정확도를 요구하는 테스트는 모델의 수치적 동작이 예상대로 이루어짐을 보장합니다. 이는 특히 딥러닝 모델의 미묘한 수치적 차이가 결과에 큰 영향을 미칠 수 있는 경우에 중요합니다.
MegaMoE와의 정렬: 기존의 검증된 라이브러리(MegaMoE)와의 연산 방식 정렬은 코드의 신뢰성을 높이고, 잠재적인 버그를 줄이는 데 도움이 됩니다.

일반적인 교훈

하드웨어 특성 활용: GPU 아키텍처(예: Tensor Core의 MMA 연산)의 세부 사항을 이해하고 이를 코드에 반영하는 것이 성능 향상의 핵심입니다.
동적 튜닝의 중요성: 고정된 연산 크기 대신 런타임 정보를 활용하여 연산을 동적으로 조절하는 것은 다양한 입력 크기와 조건에서 최적의 성능을 달성하는 데 필수적입니다.
정확성과 성능의 균형: 때로는 성능 향상을 위해 정밀도를 희생하기도 하지만, 이 PR은 오히려 정밀도(subnormal 보존, signed zero 처리)를 높이면서 성능까지 개선했습니다. 이는 최적의 구현을 통해 두 마리 토끼를 잡을 수 있음을 보여줍니다.
엄격한 테스트의 가치: 비트 단위 정확도 테스트는 코드 변경의 영향을 정확히 파악하고, 예상치 못한 부작용을 방지하는 데 결정적인 역할을 합니다.

리뷰 댓글 분석

리뷰 댓글에서 zianglih는 W4A4와 W4A16 디스패치 로직의 혼동 가능성에 대해 지적했습니다. 특히, tests/moe_ep/test_split_layer_w4a16_cudagraph_multirank.py 테스트가 W4A16 경로를 올바르게 타는지, 그리고 W4A4 관련 로직이 불필요하게 실행되지 않는지에 대한 우려를 표했습니다. 또한, W4A16 런처가 독립적인 moe_unpermute를 사용하며 W4A4 래퍼에 진입하지 않는다는 점을 명확히 했습니다. 이 부분은 코드 수정 후 b049d7c4 커밋에서 해결되었으며, 아키텍처 가드(architecture guard)가 올바르게 작동하도록 수정되었습니다.

flashinfer-bot의 CI 결과는 일부 새로운 테스트 케이스(예: tests.moe.test_cake_w4a16_fused_moe)에서 'illegal instruction' 오류가 발생했음을 보여줍니다. 이는 특정 GPU 아키텍처(B200, H100 등) 및 CUDA 버전 조합에서 발생할 수 있는 문제로, 하드웨어 호환성 또는 드라이버 관련 이슈일 가능성이 있습니다. PR 자체의 로직 오류라기보다는 환경 설정이나 특정 하드웨어/소프트웨어 스택에서의 호환성 문제일 수 있습니다.

결론

이번 PR(#7196)은 FlashInfer 라이브러리의 W4A16 MoE 구현에 있어 중요한 성능 향상과 정확성 검증 강화를 이루어냈습니다. 런타임 MMA N 최적화, BF16 처리 개선, 그리고 엄격한 비트 단위 정확도 테스트 도입은 MoE 모델의 효율적이고 신뢰할 수 있는 추론을 가능하게 합니다. 이러한 최적화는 AI 모델의 성능을 극한까지 끌어올리려는 노력의 좋은 예시이며, GPU 하드웨어의 특성을 깊이 이해하고 이를 소프트웨어 최적화에 적용하는 것이 얼마나 중요한지를 다시 한번 보여줍니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [onnxruntime] [CUDA] HyperConnectionLowRankMix 최적화: Speculative Decoding과 Hopper 아키텍처를 위한 튜닝
현재글 : [flashinfer] FlashInfer, CuTe DSL W4A16 MoE 최적화 및 비트 단위 정확도 테스트 추가
다음글 [flashinfer] FlashInfer의 Cake Ulysses All-to-All 최적화: World Size 2 성능 향상 및 Symmetric Handshake 도입

댓글

관련 포스트

PR Analysis 의 다른글