docs/benchmark.md

Benchmark System

Overview

벤치마크 시스템은 커스텀 CUDA 커널의 forward/backward 성능을 naive PyTorch 구현 대비 측정한다. Triton의 do_bench를 사용하며, 정확도 검증(correctness check) 후 성능을 측정한다.

Directory Structure

benchmarks/
├── run_cases.py              # CLI 진입점
├── common/
│   ├── bench_framework.py    # 벤치마크 유틸리티 (Triton perf_report 기반)
│   └── diff_engine.py        # 정확도 검증 엔진 (DiffCase ABC)
├── cases/                    # 벤치마크 케이스 구현
│   ├── rms.py                # RMSNorm
│   ├── add_rms.py            # Fused Add + RMSNorm
│   ├── poly.py               # PolyNorm
│   ├── mul_poly.py           # Fused Mul + PolyNorm
│   └── grouped_mul_poly.py   # Grouped MoE Fused Mul + PolyNorm
├── benchmark.yaml            # Kubeflow 벤치마크 job config
├── test.yaml                 # Kubeflow 테스트 job config
├── plots/                    # 생성된 플롯 결과
└── results/                  # 타임스탬프별 벤치마크 결과

Usage

python benchmarks/run_cases.py --case <CASE> [OPTIONS]

Arguments

Argument	Default	Choices	Description
--case	(필수)	rms, add_rms, poly, mul_poly, grouped_mul_poly	벤치마크 케이스
--dtype	bf16	fp16, bf16, fp32, all	데이터 타입
--save-path	./configs/	경로	결과 출력 디렉토리
--plot	false	-	플롯 생성 모드
--profile	false	-	Chrome trace 프로파일링 내보내기

Examples

# bf16 기본 벤치마크
python benchmarks/run_cases.py --case grouped_mul_poly

# 모든 dtype + 프로파일링
python benchmarks/run_cases.py --case mul_poly --dtype all --profile --save-path ./results

# 플롯만 생성
python benchmarks/run_cases.py --case rms --plot --save-path ./plots

Benchmark Cases

각 케이스는 DiffCase ABC를 구현하며, naive(PyTorch 참조)와 CUDA 커널을 비교한다.

Case	Naive	CUDA	Inputs
rms	torch.nn.RMSNorm	activation.layers.RMSNorm	x, weight, eps
add_rms	custom FusedAddRMSNorm	activation.layers.FusedAddRMSNorm	x, residual, weight, eps
poly	custom PolyNorm (x^3, x^2, x 조합)	activation.layers.PolyNorm	x, weight(3), bias(1), eps
mul_poly	custom FusedMulPolyNorm	activation.layers.FusedMulPolyNorm	x, mul, weight(3), bias, eps
grouped_mul_poly	fused_mul_grouped_poly_norm_ref	fused_mul_grouped_poly_norm	x, mul, weight(num_experts, 3), bias, offsets

grouped_mul_poly는 추가로 compiled(torch.compile된 naive)와 compiled_cuda(torch.compile된 CUDA) provider도 측정한다.

Execution Flow

1. 정확도 검증 - 3개 config에 대해 calculate_diff() 실행
   • (bs=2, sl=128, hidden=4096)
   • (bs=8, sl=4096, hidden=1280)
   • (bs=1, sl=32768, hidden=1280)
   • forward/backward 모두 atol=1e-2, rtol=1e-2로 비교
2. 벤치마크 실행 - dtype별로 forward/backward 성능 측정
3. 결과 저장 - CSV 파일 (및 선택적으로 플롯/trace)

Configuration Ranges

Standard cases (rms, add_rms, poly, mul_poly):

• Batch sizes: 1, 2, 4, 8
• Sequence lengths: 1024, 2048, 4096, 8192
• Hidden dims: 2048, 4096

Grouped case (grouped_mul_poly):

• Total tokens: 1024 ~ 65536 (bs x sl)
• Hidden dim: 1280 (고정)
• Experts: 48 per rank

--plot 모드에서는 bs=1로 고정하고 seq_len만 sweep한다.

Output

CSV

{save_path}/{case}/{dtype}/ 디렉토리에 저장:

• {case}-{dtype}-fwd-perf.csv - forward 결과
• {case}-{dtype}-bwd-perf.csv - backward 결과

컬럼: dim, batch_size, seq_len, Naive (us), Compiled (us), Cuda (us), SpeedUp (us)

Chrome Trace (--profile)

{save_path}/{case}/{dtype}/traces/ 디렉토리에 JSON 형식으로 저장. chrome://tracing에서 로드하여 GPU 타임라인을 분석할 수 있다.

파일명 패턴: trace_{fwd|bwd}_{naive|compiled|cuda|compiled_cuda}_N{total_tokens}.json

Plot (--plot)

Speedup 비교 플롯 생성. Geometric mean으로 전체 speedup을 집계한다.

Framework Internals

bench_framework.py

Triton의 perf_report/Benchmark를 사용하는 4개 팩토리 함수:

• make_fwd_benchmark_for_case() - forward 벤치마크 (CSV)
• make_bwd_benchmark_for_case() - backward 벤치마크 (CSV)
• make_fwd_benchmark_plot_for_case() - forward 플롯
• make_bwd_benchmark_plot_for_case() - backward 플롯

타이밍은 triton.testing.do_bench()로 측정하며, ms 단위를 us로 변환한다 (time_unit_scale=1000).

diff_engine.py

DiffCase ABC 인터페이스:

• build_inputs(bs, sl, dim) - 입력 텐서 생성
• make_naive() / make_cuda() - 구현체 생성
• forward(module, inputs) - forward 실행
• grad_inputs(inputs) - gradient 대상 텐서 반환

calculate_diff()가 naive와 CUDA 양쪽의 forward output + backward gradient를 torch.testing.assert_close()로 비교한다.

Kubeflow Integration

benchmark.yaml로 클러스터에서 벤치마크를 실행할 수 있다:

• triton, matplotlib, pandas 설치
• C++ extension 빌드 (setup.py)
• GPU warmup (100 iterations matmul)
• 결과를 benchmarks/results/{YY_MM_DD_HH_MM}/에 저장