compile=True와 TileLang 빠른 경로: 엔지니어링 노트

이 노트는 compile=True와 fast=True가 README가 말하는 것 너머에서 어떻게 동작하는지 다룹니다. #472, #576, #718 작업 중 RTX 4070 Ti SUPER, torch 2.11, tilelang 0.1.14에서 측정한 것입니다. 다음 사람이 다시 측정하지 않도록 여기 남깁니다.

백엔드 선택

Agent(..., backend="auto")와 laya.load(..., backend="auto")는 백엔드 클래스 계층을 선택합니다. 기본값은 여전히 eager입니다. 명시적 backend=는 compile과 fast보다 우선하며, 생략하면 두 플래그의 기존 동작이 유지됩니다.

• eager: 지원되는 모든 장치에서의 표준 PyTorch forward.
• compile: CUDA 전용 torch.compile로, 동적 형상과 reduce-overhead 모드, 버킷 패딩, 영속 inductor 캐시, 설치 시 warmup을 사용합니다. compile=True와 같은 독립 차원 범위를 재사용하며, 기존 기본 모드와 CPU 지원을 유지합니다. LAYA_COMPILE_WARMUP=0으로 백엔드 warmup을 미루고 LAYA_INDUCTOR_CACHE_DIR로 캐시 디렉터리(기본값 ~/.cache/laya/inductor)를 고를 수 있습니다.
• tilelang: 현재 빠른 경로를 감싼 어댑터로, 에이전트의 bf16 또는 fp16 dtype을 사용합니다.
• auto: TileLang이 설치되어 있고 지원되는 ModernBERT 인코더와 dtype이 있으면 CUDA에서 TileLang, 그렇지 않으면 CUDA에서 compile, 다른 장치에서는 eager.
• onnx: laya.load(..., backend="onnx", onnx_path="model.onnx")는 기존 ONNXAgent를 반환합니다. onnx_path가 없으면 laya.onnx를 사용합니다.

사용할 수 없는 백엔드는 해석된 백엔드를 밝히는 RuntimeWarning을 내고 eager로 폴백합니다. 백엔드를 강제하려면 agent.set_backend("tilelang", strict=True)를 사용하십시오. 전환은 활성 추론을 기다립니다. agent.backend는 활성 이름을 보고하고 agent.backend_object는 설치된 객체를 노출합니다. agent.set_backend("compile", warmup=False)는 추론까지 컴파일을 미루므로 컴파일 오류가 그때 요청에서 드러납니다. agent.warmup()은 계속 사용할 수 있습니다. agent.deaccelerate()는 클래스 계층을 통해 설치된 백엔드를 제거합니다.

Router는 명시적 선택을 Router(agent_kwargs={"backend": "auto"})로 전달합니다. 기본적으로 백엔드 인자를 넘기지 않아 기존 Agent 유사 생성자와의 호환성을 유지합니다. 범위가 지정된 CPU OOM 재시도는 백엔드를 분리하고 모델이 원래 장치로 돌아올 때 복원합니다.

compile=True는 어텐션 마스크를 실체화합니다

Eager SDPA는 ModernBERT의 (rows, 1, L, L) 어텐션 마스크를 브로드캐스트 뷰로 취합니다. compile=True가 쓰는 동적 형상 아래에서는 inductor가 마지막 차원이 정렬되었음을 증명할 수 없습니다. 그래서 마스크를 모든 head로 확장하고 rows x heads x L x L의 실제 버퍼로 패딩합니다. head가 12개인 bf16에서 그것은 32행 x 1024토큰일 때 0.8 GB입니다.

• 여유 있는 GPU. 버퍼는 대역폭을 소모하며, 긴 배치마다 수십 ms가 듭니다.
• 거의 꽉 찬 GPU. 캐싱 할당자가 요동치며, 같은 호출이 수십 초가 걸릴 수 있습니다.

바쁜 GPU에서 긴 배치로 컴파일한다면 배치 크기를 제한하거나 (predict_batch(..., batch_size=)) or use fast=True. TileLang 어텐션은 패킹된 QKV 버퍼를 읽고 시퀀스 길이로 마스킹하므로 그런 버퍼가 없습니다.

콜드 스타트

• 첫 컴파일. 그래프마다 수십 초가 걸립니다. compile=True는 두 개의 그래프가 필요합니다. 배치용 하나와 torch가 특수화하는 단일 행용 하나입니다. compile=True는 이제 로드하는 동안 agent.warmup()을 호출합니다. compile_warmup=False는 지연 컴파일을 복원하고, agent.warmup(shapes=...)은 계속 수동으로 사용할 수 있습니다. eager와 TileLang 로드는 자동으로 워밍하지 않습니다. 이 형상들은 흔한 요청을 커버할 뿐, 가능한 모든 형상 가드를 커버하지는 않습니다.
• warmup 실패. 자동 warmup은 최선 노력입니다. 실패하면 오류(기저 컴파일러 오류 포함)를 밝히는 RuntimeWarning을 내고, 로드는 torch.compile 래퍼와 컴파일 설정을 그대로 둔 채 돌아옵니다. 예를 들어 MSVC가 없는 Windows는 warmup이 실패해도 compile=True로 로드할 수 있습니다. 이후 요청은 컴파일된 모델을 그대로 사용하며 컴파일 실패를 드러냅니다. Laya는 그것들을 eager 실행으로 전환하지 않습니다. 명시적 agent.warmup() 호출도 실패를 전파하며, 실패한 자동 warmup 이후에도 마찬가지입니다. 따라서 성공한 로드가 컴파일된 추론의 준비를 보장하지는 않습니다.
• Laya 캐시 옵트인. laya.load(..., compile=True, compile_cache=True)는 프로세스 전역 TORCHINDUCTOR_CACHE_DIR이 없을 때만 그것을 설정하며, $XDG_CACHE_HOME/laya/torchinductor, XDG가 설정되지 않았거나 절대 경로가 아니면 ~/.cache/laya/torchinductor로 설정합니다. 기존 설정(이전 PyTorch 컴파일이 설정한 것 포함)이 우선합니다. 디렉터리는 로드 시 생성되며 파일 시스템 오류는 전파됩니다. compile_cache=False(기본값), eager, TileLang 로드는 환경을 건드리지 않습니다. 이것은 오래된 캐시를 옮기거나 삭제하지 않습니다. 컨테이너는 여전히 영속 home/볼륨이 필요합니다. 캐시 호환성과 무효화는 PyTorch가 관리합니다. GPU, torch, 컴파일러, 모델 또는 입력 가드 변경은 재컴파일을 요구할 수 있습니다.
• 재시작을 넘어서. Inductor의 FX 그래프 캐시는 컴파일된 그래프를 TORCHINDUCTOR_CACHE_DIR 아래에 보관합니다. 기본값은 /tmp 아래인데, 재부팅이나 컨테이너 재시작을 견디지 못합니다. 영속 디렉터리나 컨테이너의 볼륨으로 설정하면, 두 번째 프로세스가 그래프를 컴파일하는 대신 로드합니다. #472의 측정에서 그것은 워밍업을 약 120초에서 약 50초로 줄였습니다.

선택적 CUDA 그래프

agent = laya.load("convaiinnovations/laya", compile=True,
                  compile_cache=True, compile_mode="reduce-overhead")

compile_mode의 기본값은 "default"입니다. 활성 컴파일 경로에서는 "default"와 "reduce-overhead"만 허용됩니다. eager와 TileLang 로드는 컴파일 옵션을 무시합니다. CPU 컴파일은 여전히 동작하지만 CUDA 그래프 기록은 CUDA에서만 적용됩니다. CUDA 모드는 PyTorch의 torch.compiler.cudagraph_mark_step_begin API를 필요로 합니다. 그것이 없는 더 오래된 빌드는 명시적 오류를 발생시킵니다.

동적 Dynamo 그래프가 형상 독립적 CUDA 그래프를 의미하지는 않습니다. 새로운 구체 형상은 새로운 Dynamo 그래프 없이 다시 warmup과 기록을 요구할 수 있습니다. 두 가지 기본 합성 warmup 형상이 모든 요청 형상을 미리 기록하지는 않습니다. 반복되는 형상은 이득을 볼 수 있지만, 변하는 형상은 추가 지연을 치르고 그래프 풀을 유지할 수 있습니다. PyTorch는 지원되지 않는 연산이나 구성에 대해 CUDA 그래프를 건너뛸 수 있습니다. 이 모드를 설정하는 것은 캡처를 보장하지 않습니다.

Laya는 컴파일된 각 CUDA forward를 새 단계로 표시하고, 그 forward들을 에이전트들 사이에서 직렬화하며, 잠금을 해제하기 전에 두 출력 텐서를 컴파일된 그래프 밖에서 복제합니다. 이는 유지된 출력이 재생을 넘어 유효하게 유지되도록 하며, 대가는 두 번의 복사와 직렬화된 forward 실행입니다. 이 잠금은 애플리케이션이 소유한 무관한 컴파일 모델을 조정하지 않습니다. CUDA 그래프 반복을 공유하거나 사용자 정의 스트림을 사용하는 호출자는 자체 조정을 관리해야 합니다. 디스크 캐시는 프로세스 간에 컴파일된 코드를 재사용하며, 살아 있는 CUDA 그래프 기록이나 그 장치 메모리는 아닙니다.

benchmarks/bench_compile_defaults.py로 콜드/재시작 타이밍, 메모리, 캐시 카운터를 재현하고 기록된 측정을 참고하십시오.

AOTInductor: 아직

체크포인트와 GPU 아키텍처마다 미리 컴파일된 산출물을 배포하면 (torch._inductor.aoti_compile_and_package) 컴파일을 완전히 없앨 수 있습니다. torch 2.11에서는 패키징 단계에서 멈춥니다:

• 익스포트는 됩니다. DecisionModel의 torch.export는 동적 행, 마커, 토큰과 함께 약 5초에 성공합니다. 토큰은 16의 배수로 선언해야 하며(16 * Dim(...)), 평범한 범위는 익스포터 자체의 L % 8 정렬 가드를 통과하지 못합니다. 이는 위와 같은 마스크 정렬입니다.
• 패키징은 실패합니다. 어떻게 실패하는지는 프로그램이 어떻게 익스포트되었는지에 달려 있습니다:
  • autocast 아래에서 프로그램은 dtype 단언을 지니는데, AOTI가 autocast 밖에서 그것에 걸립니다: Tensor dtype mismatch! Expected: torch.bfloat16, Got: torch.float32.
  • autocast 없이 bf16 사본에서 트레이싱이 forward 안에서 실패합니다: mat1 and mat2 must have the same dtype. DecisionModel.forward는 풀링된 상태와 신뢰도 특징을 action head 이전에 fp32로 업캐스트하며, autocast가 보통 그것을 조정합니다.

따라서 산출물 경로에는 dtype을 명시한 action head가 필요합니다. 입력을 head의 dtype으로 캐스팅하거나 head를 fp32로 실행하십시오.

TileLang 이식성: 커널은 CUDA 전용

tilelang은 CUDA, HIP, Metal, WebGPU와 C 백엔드에 대한 타깃을 등록합니다. AMD나 Apple 하드웨어 없이 답할 수 있는 질문은 laya/tl_kernels.py가 CPU용으로 낮춰지기는 하는지였습니다. Linux x86-64에서 tilelang.compile(kernel.prim_func, target=...)으로 탐침했습니다:

모든 Laya 커널은 세 가지 이유 중 하나로 "c"에서 실패합니다:

C 백엔드는 다음을 받습니다:

• fp32 요소별 루프(T.Parallel)
• 평범한 루프로 낮춰지는 T.Pipelined
• 스칼라 삼중 루프로 낮춰지는, T.alloc_local 누산기를 가진 T.gemm

따라서 CPU 버전은 타깃 플래그가 아니라 두 번째 커널 집합이 될 것입니다. 그 GEMM은 블로킹 없는 스칼라 루프가 될 것이고, 기본 forward가 CPU에서 이미 쓰는 MKL/oneDNN 경로와 경쟁하지 못할 것입니다. HIP과 Metal에서 먼저 확인해야 할 것도 같은 세 가지 구문입니다. 프래그먼트, GemmWarpPolicy를 가진 T.gemm, bf16/fp16 지원입니다.

두 번째 표의 첫 행을 재현하려면: tilelang.compile(K.gemm_kernel(768, 768).prim_func, target="c").