[flashinfer] MXFP4 MoE 기본값을 expert fill과 FC1 깊이로 고르기: 상속된 튜닝을 버리고 최대 13.85% 지연 단축

PR 링크: flashinfer-ai/flashinfer#9033 상태: Merged | 변경: +877 / -8

들어가며

FlashInfer의 moe_ep(Mixture-of-Experts Expert Parallel) 백엔드에는 SM100(Blackwell)용 MXFP8 activation × MXFP4 weight 조합인 StandardMxfp8Mxfp4Backend가 있습니다. 문제는 이 백엔드가 MXFP8 백엔드의 튜닝 기본값을 그대로 상속하고 있었다는 점입니다. 상속된 기본값은 다음과 같았습니다.

• tile (256, 128, 128)
• grouped scheduling, hint 512
• atomic work claim
• token-input flag batch 4

MXFP4 weight는 MXFP8보다 메모리 footprint와 연산 특성이 다릅니다. 그래서 한 가지 설정으로 모든 batch 크기를 커버하면 작은 batch에서는 tile이 과하게 커서 낭비가 생기고, 큰 batch에서는 스케줄링 방식이 최적이 아닙니다. 이 PR은 expert당 예상 routed row 수와 active cluster당 FC1 tile 깊이라는 두 지표로 기본값을 고르도록 바꿉니다. 사용자가 명시한 override는 그대로 존중합니다.

코드 분석

1. _default_knobs 훅 추출 (standard_cutedsl.py, 공통 베이스)

먼저 서브클래스가 기본값을 바꿀 수 있도록 확장 지점을 만들었습니다.

Before

knobs = default_knobs(
    fleet.max_tokens_per_rank,
    dtype="nvfp4" if kind == "nvfp4" else "mxfp8",
    combine_dtype=wire,
    enable_in_kernel_fc2_reduce=k.enable_in_kernel_fc2_reduce,
)

After

def _default_knobs(self, fleet):
    """Return the inherited block-scaled profile for this buffer capacity."""
    from .....kernel_src.sm100.cutedsl_megamoe import default_knobs
    return default_knobs(
        fleet.max_tokens_per_rank,
        dtype="nvfp4" if self._kind() == "nvfp4" else "mxfp8",
        ...
    )

# _workspace_config 내부
knobs = self._default_knobs(fleet)

전형적인 Template Method 패턴입니다. NVFP4/MXFP8 경로의 동작은 바뀌지 않고, MXFP4 서브클래스만 이 메서드를 override합니다.

2. schedule_mode knob 추가와 검증

이전에는 schedule policy가 "grouped"로 하드코딩되어 있었습니다.

Before

schedule_policy=("grouped", knobs.get("group_hint")),

After

schedule = knobs.get("schedule_mode", "grouped")
if schedule not in ("grouped", "phase_interleave"):
    raise MoEEpConfigError(f"Unsupported schedule_mode {schedule!r}.")
if schedule == "phase_interleave" and balance != "atomic_counter":
    raise MoEEpConfigError(
        "schedule_mode='phase_interleave' requires "
        "load_balance_mode='atomic_counter'."
    )
...
schedule_policy=(schedule, knobs.get("group_hint")),

phase_interleave는 FC1과 FC2 phase의 work를 섞어 배치하는 스케줄입니다. CTA가 동적으로 work를 가져가야 하므로 atomic_counter claim이 필수입니다. 이 제약을 config 단계에서 막아 두면 잘못된 조합이 kernel launch 시점에 hang이나 오동작으로 드러나는 일을 예방할 수 있습니다.

3. MXFP4 전용 휴리스틱: 핵심 로직

_N64_MAX_ROWS = 64
_PHASE_INTERLEAVE_MIN_ROWS = 128
_PHASE_INTERLEAVE_MIN_TILES_PER_CLUSTER = 7

def _default_knobs(self, fleet):
    knobs = super()._default_knobs(fleet)
    routes = self.ep_world_size * fleet.max_tokens_per_rank * self._config.top_k
    rows = routes / fleet.num_experts
    small = rows < self._N64_MAX_ROWS
    tile = (256, 64 if small else 128, 128)
    knobs.update(
        mma_tiler_mnk=tile,
        group_hint=4096,
        flag_batch=1,
        load_balance_mode="atomic_counter" if small else "static",
    )

1단계: expert fill로 N tile을 고릅니다. expert당 예상 row 수는 world_size × capacity × top_k / num_experts입니다. 이 값이 64 미만이면 N=128 tile의 절반 이상이 padding이 됩니다. 이때는 N64로 줄여 낭비를 없애고, 작은 work 단위의 불균형을 흡수하도록 atomic claim을 씁니다. row가 충분하면 N128과 오버헤드가 없는 static(grid-stride) claim을 씁니다.

2단계: FC1 깊이로 phase interleave를 켭니다.

if rows >= self._PHASE_INTERLEAVE_MIN_ROWS:
    cluster = knobs.get("cluster_shape_mnk", (2, 1, 1))
    clusters = self._active_clusters(cluster[0] * cluster[1])
    token_tile_routes = fleet.num_experts * tile[1]
    token_tiles = (routes + token_tile_routes - 1) // token_tile_routes
    weight_tiles = (2 * self._config.intermediate_size + tile[0] - 1) // tile[0]
    fc1_tiles = (
        (fleet.num_experts // self.ep_world_size) * token_tiles * weight_tiles
    )
    if fc1_tiles >= clusters * self._PHASE_INTERLEAVE_MIN_TILES_PER_CLUSTER:
        knobs.update(schedule_mode="phase_interleave", group_hint=None,
                     load_balance_mode="atomic_counter")

phase interleave에는 prologue 비용이 있습니다. 그래서 FC1 work가 얕으면 섞을 것이 별로 없어 손해가 됩니다. 이 코드는 local expert 수 × token tile × weight tile(2 * intermediate_size는 gate+up projection)로 FC1 tile 총량을 계산합니다. 그 값이 active cluster당 7 tile 이상일 때만 interleave를 켭니다. Reviewer Notes에 따르면 이 7이라는 값은 측정된 5.3과 7.9 tiles/cluster 사이에서 고른 경계입니다. 모든 geometry에서의 최적값은 아니라고 솔직하게 밝히고 있습니다.

또 하나 짚을 점은 선택 기준이 실제 token 수가 아니라 buffer capacity(max_tokens_per_rank)라는 것입니다. CUDA graph replay에서는 실제 token이 capacity보다 적을 수 있습니다. 그래도 graph에 capture된 kernel 구성은 고정되어야 하므로 의도된 설계입니다.

4. Override 존중

overrides = self._config.knobs
if isinstance(overrides, dict) and set(overrides) & {
    "mma_tiler_mnk", "cluster_shape_mnk", ..., "schedule_mode", "work_id_mode",
}:
    policy = overrides.get("schedule_policy")
    if overrides.get("schedule_mode") == "phase_interleave" or (...):
        knobs["load_balance_mode"] = "atomic_counter"
        knobs["group_hint"] = None
    return knobs

사용자가 geometry나 scheduling 관련 knob을 하나라도 지정하면 휴리스틱을 건너뜁니다. 측정 범위 밖의 설정이기 때문입니다. 다만 phase-only hint만 줬을 때는 호환되는 claim mode와 hint를 함께 맞춰 줍니다. 그래서 schedule_mode만 지정해도 2번의 검증에 걸리지 않습니다.

5. 분산 일관성: collective 에러 전파

EP 환경에서는 모든 rank가 같은 kernel config를 써야 합니다. 그런데 _active_clusters는 각 rank의 GPU에서 occupancy를 probe하므로 rank마다 결과가 다를 수 있습니다.

def _gather_workspace_proposals(self, error, proposal=None):
    reports = [None] * self.ep_world_size
    dist.all_gather_object(reports, (error, proposal), group=self.ep_comm_group)
    failures = [f"rank {rank}: {error}" for rank, (error, _) in enumerate(reports) if error is not None]
    if failures:
        raise MoEEpConfigError(...)

prepare_workspace는 각 rank의 (config, active_clusters) proposal을 모은 뒤 하나라도 다르면 예외를 던집니다. 핵심은 이 합의가 pool lookup과 symmetric allocation 이전에 일어난다는 점입니다. 한 rank만 예외로 빠져나가고 나머지가 NVSHMEM symmetric allocation 같은 collective에서 기다리면 deadlock이 됩니다. 에러를 all_gather_object로 공유하면 모든 rank가 함께 실패합니다. 리뷰 과정에서도 achartier가 이 경로의 문제를 지적했고, b1a28a181 이후 커밋에서 validation 에러를 로컬이 아닌 collective로 보고하도록 수정되었습니다.

왜 이게 좋은가

GB300, H4096/I6400, 8 experts, top-2 조건에서 upstream 기본값과 비교한 결과입니다(cold-L2 CUPTI, CUDA graph).

EP	Tokens/rank	Upstream (µs)	This PR (µs)	변화
2	32	70.896	66.528	-6.16%
2	256	90.460	81.669	-9.72%
2	1024	159.881	146.005	-8.68%
4	256	90.389	77.868	-13.85%
4	768	137.952	124.540	-9.72%
4	1024	163.582	148.446	-9.25%

17개 측정점 모두에서 지연이 줄었습니다. 두 arm은 같은 weight와 routing에서 bitwise 동일한 출력을 내도록 검증되었습니다.

일반적 교훈

1. 상속된 튜닝은 기술 부채입니다. dtype이 바뀌면 최적 tile과 스케줄도 바뀝니다. 기본값을 hook으로 분리해 두면 서브클래스별 튜닝이 쉬워집니다.
2. 휴리스틱은 물리적 의미가 있는 지표로 만드세요. token 수 같은 단일 축 대신 "expert당 row"(padding 낭비)와 "cluster당 tile"(병렬 깊이)을 쓰면 다른 모델 shape에도 일반화될 가능성이 높습니다.
3. 명시적 override는 휴리스틱보다 우선합니다. 자동 선택은 측정된 범위 안에서만 동작하게 하세요.
4. 분산 환경의 config 결정은 collective로 합의하세요. 로컬 probe 결과에 의존하는 결정은 rank 간 불일치나 deadlock 위험이 있으므로 allocation 전에 반드시 합의 단계를 둬야 합니다.

참고 자료

• torch.distributed.all_gather_object — rank 간 에러/proposal 공유에 사용한 collective API
• torch.cuda.device — occupancy probe 시 EP 디바이스 컨텍스트 지정
• torch.cuda.CUDAGraph — capacity 기반 선택의 근거가 되는 graph replay
• dataclasses.replace — rank 필드를 정규화해 proposal 비교에 사용

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