[Triton] matmul 커널 시그니처에 input microblock size 추가

2026년 3월 25일

들어가며

Triton의 matmul 커널에서 microscaled 연산(MXFP, NVFP4)을 사용할 때, microblock size(스케일 블록 크기)를 텐서 shape에서 추론하던 방식에서 명시적으로 PrecisionConfig에 지정하는 방식으로 변경하는 PR이다. NVFP4(block size 16)와 MXFP(block size 32)를 구분하기 위해 필요하다.

핵심 코드 분석

Before

a_microblock_size = None if a_scale is None else a.shape[-1] // a_scale.shape[-1]
b_microblock_size = None if b_scale is None else b.shape[-2] // b_scale.shape[-2]

microblock size를 텐서 shape 비율로 추론했다. 이 방식은 MXFP와 NVFP4를 구분할 수 없었다.

After

@dataclass
class PrecisionConfig:
    a_mx_scale: torch.Tensor | Tensor | None = None
    a_microblock_size: int | None = None  # 새로 추가
    b_mx_scale: torch.Tensor | Tensor | None = None
    b_microblock_size: int | None = None  # 새로 추가

# 사용 시 명시적 지정
pc = PrecisionConfig(
    b_mx_scale=b_scale,
    b_microblock_size=MXFP_BLOCK_SIZE.value,  # 32
)

검증 로직도 추가되었다:

assert b_scale is None or b_microblock_size is not None, (
    "precision_config.b_microblock_size is required when precision_config.b_mx_scale is set"
)