문서

빠른 백엔드

TileLang 이식성

laya/tl_kernels.py의 커널 다섯 개는 compile_cpu를 통해 TileLang의 CPU(c) 타깃에서 실행할 수 있습니다. 이는 명시적인 스칼라 fp32 특수화이며, Agent.accelerate의 CPU 가속 백엔드가 아닙니다. TileLang과 로컬 C++ 컴파일러가 필요합니다. 런타임 서비스 의존성은 도입하지 않습니다.

import torch
from laya import tl_kernels as K

A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)

compile_cpu(factory, *args, **kwargs)는 다섯 개 GPU 팩토리와 같은 형상 및 연산 옵션을 받습니다. cpu=True, dtype="float32", 타깃 c를 선택하고 벡터화를 비활성화합니다. 모든 부동소수점 입력과 출력은 연속된 CPU float32 텐서여야 합니다. 어텐션 길이는 int32로 유지됩니다. 호출하기 전에 16비트 활성화를 .cpu().float().contiguous()로 변환하십시오. LayerNorm은 여전히 잔차 텐서를 제자리에서 갱신하고, RoPE는 여전히 Q/K 열을 제자리에서 갱신합니다. 컴파일된 커널을 보관하여 동적 차원을 재사용하십시오.

CPU 특수화는 fragment와 shared 할당을 로컬 버퍼로 바꾸고, 직렬 T.grid 루프를 사용하며, 어텐션의 GPU swizzle 주석을 생략합니다. GEMM은 TileLang의 스칼라 CPU 구현을 사용하고 리덕션은 로컬 버퍼를 사용합니다. 원래의 타일 알고리즘, 패딩 조건자, 어텐션 마스크, 온라인 소프트맥스는 GPU 구현과 공유된 채로 남습니다. CPU 중간 결과는 어텐션 확률을 포함해 fp32로 유지됩니다. GPU 중간 결과는 원래 dtype을 유지합니다. CPU 가속이나 전체 모델 CPU 백엔드는 주장하지 않습니다.

TileLang 0.1.14에서 재탐침

Linux, Python 3.12.13, torch 2.11.0+cu130, TileLang 0.1.14, RTX 4070 Ti SUPER에서 관측했습니다. 이전의 이식성 우려는 수정되지 않은 GPU 특수화를 컴파일하는 것에 관한 것이며, CPU로의 lower 가능성에 관한 것이 아닙니다. 베이스 커밋 fa9a2a7에서는 이 문서 페이지가 checkout에 없었습니다.

테스트는 factory.get_tir(...)를 target="c"와 GPU의 FAST 패스 구성으로 컴파일합니다. 다음은 정확한 진단 텍스트입니다(소스 위치와 스택 트레이스는 생략). 모든 커널에 대해 bf16과 fp16을 모두 탐침했습니다.

커널과 탐침 차원 첫 bf16 실패 첫 fp16 실패
gemm_kernel(128, 64) Check failed: layout_map.count(buffer) != 0 (0 vs. 0) : The layout for fragment C_l can not be inferred correctly. 동일
gemm_geglu_kernel(64, 64) CPU fill only supports local and global buffers, but got dst scope `local.fragment`. 동일
add_ln_kernel(128) CPU reduce only supports local src and local/local.var dst buffers, got src scope `local.fragment` and dst scope `local.fragment`. 동일
rope_kernel(2, 64) Cannot convert type bfloat16 to C type C++ 컴파일러: error: no matching function for call to ‘vec_type<float, 4>::vec_type(half4&)’
attn_kernel(1, 64, 2, 64) Check failed: layout_map.count(buffer) != 0 (0 vs. 0) : The layout for fragment s_c can not be inferred correctly. 동일

처음 네 개의 fragment/리덕션 실패는 tvm.error.InternalError입니다. bf16 코드 생성 실패도 InternalError입니다. FP16 RoPE는 RuntimeError: Compilation Failed!를 발생시키고, 이어서 컴파일러 호출과 소스가 나옵니다. 위 진단은 stderr로 출력됩니다. 생성된 벡터 캐스트도 float 벡터를 다시 half 벡터로 변환하는 데 실패합니다.

dtype 지원을 fragment 지원과 분리하기 위해, 테스트는 각 커널을 cpu=True(로컬 버퍼와 직렬 루프)로 다시 컴파일하고, bf16을 유지하며 벡터화를 비활성화합니다. 그러면 다섯 개 모두 정확히 다음으로 실패합니다:

Cannot convert type bfloat16 to C type

따라서 fragment는 GEMM, GEGLU, 어텐션의 첫 번째 장애물이고, fragment 리덕션은 LayerNorm의 첫 번째 장애물이며, bf16은 독립적으로 다섯 개 모두의 장애물입니다. RoPE에는 fragment 할당도 리덕션도 없습니다. GEMM과 GEGLU에는 명시적 T.reduce_* 연산이 없습니다. 어텐션의 리덕션은 처음에는 레이아웃 실패에 가려집니다. fp32 전용의 별도 fragment 탐침은 GEMM이나 bf16 없이 reduce_sum과 reduce_max 모두에서 필 오류와 리덕션 오류를 재현합니다. scope만 바꾸는 것도 GEMM 버퍼에서 이 의미 오류를 일으켰습니다(영향받은 다른 버퍼는 Ci, x, s):

[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.

직렬 fp32 특수화는 이러한 장애물을 제거합니다. 진단 단언은 버전 0.1.14에 고정되어 있고 다른 버전에서는 건너뜁니다. 그 버전에서는 오류를 다시 탐침해야 합니다. 수치 CPU 테스트는 다른 버전에서도 계속 실행됩니다.

수치 검사

python -m pytest tests/test_fast_cpu.py -q -s를 실행하십시오. 위 환경에서 58개 통과했습니다. CPU 전용 검사는 CUDA가 필요하지 않습니다. CUDA가 없을 때 건너뛰는 것은 GPU 비교뿐입니다. 커버리지에는 불균일한 GEMM M/N/K 타일, 모든 GEMM 에필로그, GEGLU, 모든 잔차/바이어스 조합, 큰 잔차 값, RoPE 위치 랩어라운드와 변경되지 않은 V 열, 정적/동적 어텐션 형상, 슬라이딩 윈도, 부분 타일, 불균일한 길이, 빈 시퀀스, 유한한 패딩 출력이 포함됩니다.

시드는 1234입니다. GPU 비교는 동일한 입력 값을 bf16 또는 fp16으로 반올림한 뒤 CPU 실행을 위해 fp32로 승격합니다. 이는 유계 픽스처에 대한 절대 허용 오차이며, 임의의 크기나 모델 깊이에 대한 보장이 아닙니다. 어텐션 비교는 tests/test_fast.py에서처럼 유효한 query 행을 사용합니다.

커널 CPU 대 fp32 참조 최대값 CPU 대 GPU 최대값(두 dtype) CPU/GPU 허용 오차
GEMM 2.38419e-7 0.00770831 0.05
GEGLU 2.98023e-8 0.000208303 0.05
LayerNorm 1.07288e-6 0.0156183 0.05
RoPE 0 0.0130053 0.05
Attention 5.96046e-7 0.00377572 0.02

CPU/참조 허용 오차는 2e-5(RoPE는 2e-6)입니다. 잔차 스트림 갱신은 잔차가 없는 경우를 포함해 정확히 같습니다.

GPU 보존 증거

cpu 키워드의 기본값은 false입니다. 기존 bf16/fp16 기본값, GPU 할당 scope, 병렬 루프, swizzle, FAST 옵션은 변경되지 않았습니다. 기본 fp32 GPU 호출은 여전히 ValueError를 발생시킵니다.

전후 실행은 각각 git show fa9a2a7:laya/tl_kernels.py로 추출한 원본 모듈과 변경된 모듈을 사용했습니다. 베이스라인 실행에서는 원본을 importlib를 통해 laya.tl_kernels로 로드했습니다. 나머지 Laya 코드와 Python 환경은 그대로였습니다.

  • python -m pytest tests/test_fast.py -q, full-forward 테스트는 아래에 식별된 캐시된 영어 체크포인트를 로드하도록 구성했습니다: 변경 전 13개 통과, 변경 후 13개 통과. 처음에 체크포인트 없이는 11개 통과 / 2개 건너뜀이었습니다. 두 전체 실행 모두 체크포인트의 기존 온도 클램핑 경고(choice:11+=0.10058280825614929 -> 0.5)를 출력했습니다.
  • 기존 커널 테스트의 시드 캡처: 잔차 갱신을 포함해 24개 출력 텐서가 비트 단위로 동일합니다. 전후 최대 차이는 0입니다.
  • 위 다섯 개 탐침 형상에 대해 두 dtype으로 생성한 CUDA 소스: 10/10이 바이트 단위로 동일합니다. 이는 원래 fast 스위트에 없는 RoPE도 포함합니다.
  • python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ... 및 동등한 fp16 명령: dtype당 60개 상태에 걸친 288개 질문. 전후 모든 JSON 레코드(fp32, stock, fast 확률)는 정확히 동일하게 비교되고, 전후 최대 확률 차이는 0입니다.

MODEL은 캐시된 convaiinnovations/laya 영어 스냅샷 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851이었습니다. 실행에는 HF_HUB_OFFLINE=1을 사용했습니다. 다국어나 타입 지정 의사결정 체크포인트 비교는 여기서 주장하지 않습니다.

dtype type n fast와 stock 최대 차, 전 = 후 fast/stock argmax 일치, 전 = 후
bf16 choice 48 0.0310 47/48
bf16 noul 180 0.0756 180/180
bf16 score 60 0.0152 60/60
fp16 choice 48 0.0069 48/48
fp16 noul 180 0.0092 180/180
fp16 score 60 0.0040 60/60

저장소 검사

모든 명령은 /home/ckl/projects/S/laya/.venv/bin/python을 사용했습니다. ruff와 zensical은 그 가상 환경의 bin 디렉터리에서 가져왔습니다.

명령 결과
ruff check laya/ --select=E9,F63,F7,F82,F401,F811 --line-length=120 All checks passed!
python -m compileall -q laya/ tests/ 종료 코드 0, 출력 없음
python tests/test_router.py 703개 통과, 0개 실패
python tests/test_criteria.py 198개 통과, 0개 실패
python tests/test_hooks.py 240개 통과, 0개 실패
python tests/test_hooks_api.py 415개 통과, 0개 실패
python tests/test_packaging.py 131개 통과, 0개 실패
uv pip install --python /home/ckl/projects/S/laya/.venv/bin/python -r requirements-docs.txt 패키지 3개 확인(이미 설치됨). 가상 환경에 pip가 없습니다
zensical build --strict --clean No issues found. griffe: 줄은 없습니다

새 스위트는 패키징 테스트의 기존 예외 목록에, 선택적 TileLang extra와 C++ 컴파일러가 필요하다고 등록되어 있습니다. API 계약 테스트는 베이스 CI 환경에 TileLang을 임포트하지 않고, 추가된 키워드 전용 CPU 선택자, 변경되지 않은 GPU dtype 기본값, compile_cpu 시그니처를 고정합니다.

AOTInductor

DecisionModel은 내보내고, .pt2 패키지로 컴파일하고, torch._inductor.aoti_load_package로 로드할 수 있습니다. 패키지는 의사결정 로짓과 행동 로짓을 모두 반환합니다. 토큰화, 패딩, 온도 캘리브레이션, 답변 포매팅은 호출자의 몫으로 남고, Agent 백엔드를 추가하거나 기본 실행 경로를 바꾸지 않습니다.

닫힌 PR #472는 이전의 dtype 장애물을 기록했습니다. 행동 head의 풀링된 상태와 신뢰도 특징은 모델 가중치가 명시적으로 bf16이더라도 fp32로 계산됩니다. autocast가 없을 때 첫 행동 linear는 따라서 fp32 입력과 bf16 가중치를 받았습니다. 이제 forward는 연결된 입력을 head의 가중치 dtype으로, autocast 밖에서만 캐스팅합니다. Softmax, 엔트로피, 반환되는 의사결정 로짓은 fp32 계산을 유지합니다. 기존 fp32 파라미터 eager 추론은 fp16/bf16 AMP를 포함해 수치를 유지합니다. 가중치/autocast의 혼합 dtype도 autocast의 원래 변환을 유지합니다.

bf16으로 변환한 별도의 평가용 복사본을 autocast 밖에서 내보내십시오. token 축을 16 * Dim("tokens16", ...)로 선언해 어텐션 정렬 가드를 만족시키십시오. 행과 마커 개수는 독립적으로 동적일 수 있습니다. autocast 아래에서 캡처한 내보내기를 그 컨텍스트 밖에서 패키징할 수 있다고 가정하지 마십시오. 이 레시피에는 명시적 가중치 dtype을 사용하십시오.

오프라인 재현

단언 기반 검사는 기본적으로 네트워크 접근 없이 무작위로 초기화된 작은 ModernBERT를 사용합니다. 실제 모델 측정에는 로컬 체크포인트 디렉터리를 전달하십시오. CUDA, AOTInductor API 또는 C++ 컴파일러가 없으면 명시적 SKIP이 됩니다. 지원되는 설치에서 컴파일과 패리티 실패는 검사를 실패시킵니다.

python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
  python scripts/check_aoti.py --model /path/to/local/multilingual \
  --output-dir /tmp/laya-aoti

출력 디렉터리에는 decision.pt2, results.json, inputs.pt, eager.pt가 들어 있습니다. JSON은 완전한 nvidia-smi 스냅샷, 내보내기/패키징/로드 시간, 패키지 바이트 수, 지연 시간, 두 head의 최대 절대 로짓/확률 델타를 포함합니다. 바이너리 산출물과 컴파일러 캐시는 의도적으로 저장소 밖에 둡니다.

RTX 4070 Ti SUPER에서 측정

2026-10-03, Linux x86-64, Python 3.12.13, torch 2.11.0+cu130, transformers 5.17.0, NVIDIA 드라이버 615.71.09, 16,376 MiB VRAM. 체크포인트: convaiinnovations/laya, 리비전 1c5edc17a7acd8701df6fc341c0d179f1c62c982의 multilingual 하위 디렉터리. 베이스라인은 main fa9a2a7입니다. 완전한 머신 스냅샷과 반올림하지 않은 측정은 aoti_multilingual_rtx4070.json에 있습니다.

측정 전 후
내보내기 5.00 s 3.89 s
패키징 16.59 s 후 실패 60.02 s
패키지 크기 산출물 없음 645,729,621 바이트(615.82 MiB)
이미 초기화된 프로세스에서 로드 사용 불가 0.447 s
빈 캐시의 새 프로세스에서 로드 사용 불가 4.867 s

재현된 실패는 내보내기 성공 후 패키징에서 발생합니다: mat1 and mat2 must have the same dtype, but got Float and BFloat16. 각 실행은 별도의 초기 상태가 빈 Inductor 캐시를 사용했습니다. AMP 컴파일이 패키징 전에 실행되었으므로 패키지 시간은 새 인터프리터의 콜드 스타트 측정이 아닙니다. 새 프로세스 검사는 패키지와 저장된 입력만 로드했고(체크포인트 없음), torch.compile과 패키징 진입점은 차단했습니다. 두 head의 답변을 재현했습니다. PyTorch는 그래도 작은 CPU AVX 기능 탐침을 초기 상태가 빈 캐시에 컴파일했습니다. 모델 그래프나 CUDA 커널은 컴파일되지 않았습니다. 따라서 “로드 시 컴파일러 활동이 없다”는 일반화된 주장은 이 런타임에서 부정확합니다.

같은 저장된 입력에는 세 배치에 걸친 일곱 개 의사결정 행이 있고, 실제로 토큰화된 청구/환불 텍스트, choice/score/noul 질문, 패딩, 불균일한 유효 마커 개수를 담고 있습니다. 각 지연 시간은 열 번 워밍업 후 30 forward를 다섯 그룹으로 나눈 중앙값이며, 동기화된 벽시계 계측을 사용합니다. torch.compile(dynamic=True)는 기본 모드를 사용합니다. 이러한 지연 시간 수치에는 명시적 CUDA 그래프, 토큰화, 내보내기, 컴파일이 포함되지 않습니다.

행 × token × 마커 슬롯 AMP eager 전 → 후(ms) AMP compiled 전 → 후(ms) 명시적 bf16 eager(ms) 명시적 bf16 compiled(ms) AOTI bf16(ms)
2 × 128 × 3 15.1173 → 14.1200 6.7117 → 6.9076 13.9635 4.8835 2.2674
1 × 256 × 3 15.4687 → 14.4156 6.7274 → 5.4772 14.7994 4.7250 2.4154
4 × 160 × 5 14.8452 → 14.5384 7.3872 → 5.5934 14.7115 5.1393 3.6018

여기서 AMP는 fp32 파라미터와 bf16 autocast를 뜻합니다. 명시적 bf16은 가중치와 잔차 스트림이 bf16이고 autocast가 없음을 뜻합니다. 패키지에 가장 가까운 실행 비교에는 그 열을 사용하십시오. 명시적 bf16 eager/compiled와 AOTI는 이 수정 전에는 사용할 수 없었습니다.

비교, 일곱 행 전체의 최대값 의사결정 로짓 의사결정 확률 행동 로짓 행동 확률
기존 eager 전 대 후, fp32와 fp16/bf16 AMP 0 0 0 0
AOTI 대 명시적 bf16 eager 0.5625 0.00659859 0 0
AOTI 대 기존 bf16 AMP eager 0.4375 0.00769910 8.0 0

두 AOTI 비교 모두에서 의사결정과 행동의 argmax가 7/7 행에서 일치합니다. 확률은 보정되지 않은 raw softmax 출력입니다. 행동 분포는 이러한 입력에서 포화되어 있으므로, 확률 델타가 0이라고 해서 AMP에 대한 기저 로짓이 동일하다는 뜻은 아닙니다. 명시적 bf16 compiled forward도 명시적 bf16 eager와 다릅니다(의사결정 로짓 최대 델타 0.5, 확률 델타 0.00659859). 저정밀도 컴파일은 비트 단위로 정확하지 않습니다.

GPU는 데스크톱 애플리케이션과 다른 Python 작업과 공유되었습니다. CPU 컴파일도 공유되었고, 클럭은 고정되지 않았습니다. 이는 관측된 시간이며, 격리된 가속 주장이 아닙니다. 실행을 감싸는 nvidia-smi 스냅샷:

실행 / 스냅샷 GPU 사용률 사용 VRAM 전력 온도 / 상태
전 / 시작 28% 2,817 MiB 12 W 33°C / P8
전 / 종료 10% 8,560 MiB 23 W 36°C / P3
후 / 시작 0% 2,634 MiB 12 W 33°C / P8
후 / 종료 73% 6,476 MiB 160 W 42°C / P2

남은 제약

  • 산출물은 이 체크포인트, 정밀도, PyTorch/런타임 스택, GPU 타깃에 특정합니다. 다른 하드웨어나 PyTorch 버전으로의 이식성은 테스트하지 않았습니다.
  • 이 검사는 행 1–8, 16의 배수인 토큰 32–512, 마커 슬롯 2–8을 내보냅니다. 이 범위의 모든 지점이 아니라, 내보내기 예제와 다른 형상을 포함한 세 가지 형상을 시험합니다. 유효한 옵션 하나는 패딩된 마커 슬롯을 쓸 수 있습니다. 진짜 1슬롯 텐서는 별도의 단일 옵션 분기를 타며 별도 내보내기가 필요합니다. 임의의 토큰 길이와 프로덕션 버킷팅/디스패치 계층은 이 변경 범위 밖입니다.
  • 일곱 행은 패키징 회귀를 검증하는 것이지, 광범위한 체크포인트 정확도나 보정된 신뢰도 패리티가 아닙니다. 내보내기 복사본을 bf16으로 캐스팅하는 것은 autocast 아래에서 fp32 가중치를 유지하는 것과 다릅니다. 기존 eager 경로 자체는 변경되지 않았습니다.
  • 스크립트는 산출물을 만들기 위해 CUDA 지원 PyTorch 빌드와 로컬 컴파일러 툴체인이 필요합니다. .pt2는 모델과 CUDA 커널을 내장합니다. 호스팅 서비스는 관여하지 않습니다.