CUDA C++ மற்றும் Triton இல் GEMM, torch.compile மற்றும் ROCm பெயர்வுத்திறன் (GEMM in CUDA C++ and Triton, torch.compile and ROCm portability)

மின்சாரப் பரிமாற்றக் கோபுரங்களின் விளக்கப் படம்; GPU வன்பொருள் அல்ல.
மூன்று முன்னணித் துறைகளான CUDA C++ மற்றும் Triton இல் GEMM செயலாக்கம், torch.compile (TorchInductor) மூலமான தானியங்கி குறியீட்டு உருவாக்கம் (code generation), AMD ROCm இன் பல்வேறு உற்பத்தியாளர் சாதனங்களிடையேயான பெயர்வுத்திறன் (cross-vendor portability) ஆகியவை பற்றிய இந்தப் பகுப்பாய்வு ஆழ்ந்த நோக்கு, நவீன AI அமைப்புகளின் பின்னாலுள்ள நடைமுறைப் பொறியியல் சமரசங்களை (engineering trade-offs) வெளிப்படுத்துகிறது.
1. அருகருகான பகுப்பாய்வு: CUDA C++ மற்றும் Triton இல் அணிப் பெருக்கல் (Matrix Multiplication, GEMM)
அணிப் பெருக்கல் (Matrix multiplication) (C = A × B) GPU நிரலாக்கத்துக்கான இறுதி அளவுகோலாகும். ஏனெனில் அது கணிப்புச் சுமை மிகுந்தது (compute-bound); நினைவகத்தில் தடைப்படாமல் Tensor Cores க்கு இடைவிடாமல் தரவை ஊட்ட வேண்டும்.
- நவீன CUDA C++ (CUTLASS / CuTe அணுகுமுறை):
- கைமுறை வன்பொருள் ஒருங்கிணைப்பு (Manual Hardware Orchestration): CUDA C++ இல் உச்ச GEMM வேகத்தை அடைய, வேலையைக் கண்டிப்பான படிநிலையில் பிரிக்க வேண்டும்: Grid → Thread Block (CTA) → Warpgroup (4 warps) → Warp → Tensor Core instruction (wgmma).
- வெளிப்படையான குழாய்வழி அமைப்பும் ஒழுங்கு கலப்பும் (Explicit Pipelining & Swizzling): Tensor Cores ஆனவை Tile N ஐப் பெருக்கும்போது, ஒத்திசைவற்ற வன்பொருள் இயந்திரங்கள் Tile N+1 ஐ ஏற்றும் வகையில், பொறியாளர்கள் இரட்டை அல்லது மூன்று தாங்கல் (double- or triple-buffered) Shared Memory நிலைகளை (cp.async அல்லது TMA) கைமுறையாக ஒதுக்க வேண்டும். Shared Memory bank conflicts ஏற்படாமல் தடுக்க XOR memory swizzling ஐயும் பயன்படுத்த வேண்டும்.
- பகுப்பாய்வுத் தீர்ப்பு (Analytical Verdict): ஒப்பற்ற உச்சச் செயல்திறன் (கோட்பாட்டு Tensor Core FLOPS இன் 90%–98%) கிடைக்கிறது. ஆனால் நூற்றுக்கணக்கான வரிகள் கொண்ட, எளிதில் உடையக்கூடிய, கட்டமைப்புக்கே உரிய C++ template குறியீடு தேவைப்படுகிறது.
- OpenAI Triton அணுகுமுறை:
- அறிவிப்பு முறை tile கணிதம் (Declarative Tile Math): Triton இல் உயர்செயல்திறன் கொண்ட tiled GEMM ஏறத்தாழ 40 வரி Python இல் அமைகிறது. பொறியாளர் உள் பரிமாணம் K மீது ஒரு loop எழுதி, 2D blocks ஐ ஏற்றி (tl.load), பகுதிக் கூட்டுத்தொகைகளைச் சேர்க்க tl.dot(a, b) ஐ அழைக்கிறார்.
- compiler இயக்கும் குழாய்வழி அமைப்பு (Compiler-Driven Pipelining): Triton இன் MLIR backend மென்பொருள் குழாய்வழி நிலைகளைத் (num_stages=3) தானாகவே ஊகித்து, Shared Memory swizzling ஐச் சேர்த்து, Tensor Core instructions ஐ வெளியிடுகிறது. மேலும், @triton.autotune இலக்கு GPU இல் வெவ்வேறு tile அளவுகளைச் (BLOCK_M, BLOCK_N, BLOCK_K) சோதித்து, மிக வேகமான கட்டமைப்பைத் தானாகவே தேர்ந்தெடுக்கிறது.
- பகுப்பாய்வுத் தீர்ப்பு (Analytical Verdict): 10× குறைவான குறியீட்டுடன் CUTLASS செயல்திறனின் 85%–95% ஐ வழங்குகிறது. இதனால் நிலையான நூலகங்கள் போதாத தனிப்பயன் இணைக்கப்பட்ட GEMM களுக்கு (எ.கா., GEMM + Bias + GeLU + Quantization) இது மேலானதாக உள்ளது.
2. torch.compile (TorchInductor) பற்றிய பகுப்பாய்வு: PyTorch எவ்வாறு தானாகவே Triton ஐ எழுதுகிறது
PyTorch 2.0 இன் torch.compile, TorchInductor ஐப் பயன்படுத்தி Triton குறியீட்டைத் தானாகவே உருவாக்குவதன் மூலம், eager Python இன் நெகிழ்வுத்தன்மைக்கும் தொகுக்கப்பட்ட kernel இன் வேகத்துக்கும் இடையிலான இடைவெளியை நிரப்புகிறது.
[Eager Python Code]
│ 1. TorchDynamo (Bytecode Capture)
▼
[FX Graph (ATen Ops)]
│ 2. AOTAutograd (Forward + Backward Decomposition)
▼
[TorchInductor Loop IR]
│ 3. Operator Fusion & Memory Layout Analysis
▼
[Generated Triton Python Kernel] ──► [PTX / SASS Binary]- தானியங்கி இணைப்பு (Automated Fusion) எவ்வாறு செயல்படுகிறது:
- ஒரு மாதிரியை (model) @torch.compile கொண்டு அலங்கரிக்கும்போது, TorchDynamo ஆனது Python bytecode ஐ இடைமறித்து, Python interpreter மேலதிகச் சுமை இல்லாமல் ஒரு நிலையான கணிப்பீட்டு வரைபடத்தைப் (static computation graph) பிடிக்கிறது.
- TorchInductor ஏறத்தாழ 2,000 உயர்மட்ட PyTorch operators ஐ ஏறத்தாழ 250 அடிப்படை primitives ஆகக் குறைத்து, பின்னர் நினைவகச் சார்புகளை ஆய்வு செய்து, அடுத்தடுத்துள்ள pointwise (எ.கா., add, mul, silu) மற்றும் reduction (எ.கா., sum, mean, softmax) செயல்பாடுகளை ஒன்றாகத் தொகுக்கிறது.
- முன்னரே தொகுக்கப்பட்ட தனித்தனி C++ kernels ஐ அழைப்பதற்குப் பதிலாக, TorchInductor ஒரு புத்தம் புதிய Triton Python script ஐ வட்டில் (disk) உருவாக்குகிறது (TORCH_LOGS="output_code" மூலம் பார்க்கலாம்). அதை Triton இன் MLIR backend மூலம் தொகுத்து, விளைந்த GPU binary ஐ cache செய்கிறது.
- பகுப்பாய்வு சமரசங்கள் (பலங்கள் எதிர் எல்லைகள்):
- பலம் (Strength): நிலையான மாதிரித் தடைகளில் (bottlenecks) 80%–90% க்கு kernel ஐக் கைமுறையாக எழுதத் தேவையை நீக்குகிறது. ஒற்றை வரிக் குறியீட்டுடன் பயிற்சியையும் (training) அனுமானத்தையும் (inference) வழமையாக 1.5× முதல் 2.5× வரை விரைவுபடுத்துகிறது.
- எல்லை (Limitation): ஆதரிக்கப்படாத இயங்குநிலை Python கட்டுப்பாட்டு ஓட்டத்தை (dynamic Python control flow) எதிர்கொள்ளும்போது graph breaks ஏற்படுகின்றன. மேலும் சிக்கலான படிமுறை மறுவடிவமைப்புகளை (எ.கா., FlashAttention இன் online Softmax tiling) தானாகக் கண்டறிவதில் இன்னும் சிரமப்படுகிறது. இவற்றுக்குக் கையால் எழுதிய Triton அல்லது CUDA kernels இன்னும் தேவைப்படுகின்றன.
3. பல்வேறு உற்பத்தியாளர் சாதனங்களிடையேயான பெயர்வுத்திறன் பற்றிய பகுப்பாய்வு: NVIDIA CUDA எதிர் AMD ROCm
பல ஆண்டுகளாக NVIDIA இன் தனியுரிம CUDA சூழல்தொகுதி (ecosystem) AI கணிப்பீட்டில் கிட்டத்தட்ட ஏகபோகத்தை உருவாக்கியிருந்தது. உயர்மட்ட AI குறியீட்டை உற்பத்தியாளருக்கே உரிய வன்பொருள் பின்புலங்களிலிருந்து (backends) பிரிப்பதன் மூலம், நவீன compiler தொகுதிகள் இந்தப் பாதுகாப்பு அரணை முறையாகத் தகர்த்து வந்துள்ளன.
| Portability Layer (பெயர்வுத்திறன் அடுக்கு) | NVIDIA Ecosystem (NVIDIA சூழல்தொகுதி) | AMD ROCm Ecosystem (AMD ROCm சூழல்தொகுதி) | Analytical Comparison (பகுப்பாய்வு ஒப்பீடு) |
|---|---|---|---|
| Low-Level C++ Kernels (கீழ்மட்ட C++ kernels) | CUDA C++ (nvcc → PTX → SASS) | HIP (hipcc → AMDGPU LLVM → GCN/CDNA) | High Friction (அதிக உராய்வு). AMD இன் hipify கருவி CUDA தொடரியலை HIP ஆக மாற்றுகிறது. ஆனால் warp அளவுகள் வேறுபடுவதால் (NVIDIA இல் 32 threads, AMD CDNA இல் 64 threads), நிலையாக எழுதப்பட்ட warp-shuffle கணிதம் உடைகிறது. |
| Block-Level Kernels (block மட்ட kernels) | OpenAI Triton (NVPTX Backend) | OpenAI Triton (AMDGPU Backend) | Near-Seamless (கிட்டத்தட்ட இடைவெளியற்றது). Triton ஆனது warp அளவுகளை 2D blocks க்குப் பின்னால் மறைப்பதால், ஒரே Triton kernel NVIDIA H100/B200 மற்றும் AMD Instinct MI300X/MI325X ஆகிய இரண்டிலும் இயல்பாக இயங்குகிறது. |
| Framework Level (கட்டமைப்பு மட்டம்) | PyTorch + torch.compile + vLLM | PyTorch + torch.compile + vLLM | Mature Parity (முதிர்ந்த சமநிலை). முக்கிய அனுமான இயந்திரங்களும் (vLLM, SGLang) PyTorch உம், தனிப்பயன் ops ஐ Triton மற்றும் AMD இன் rocBLAS/CK (Composable Kernel) வழியாகச் செலுத்துவதன் மூலம் ROCm இல் கூடுதல் அமைப்பின்றி இயங்குகின்றன. |
- உத்திசார் முடிவுரை: கையால் எழுதப்பட்ட CUDA C++ NVIDIA சிலிக்கனுடன் மட்டுமே பிணைக்கப்பட்டிருப்பதுடன் AMD HIP க்கு கையால் மாற்ற வேண்டியுள்ளது. ஆனால் OpenAI Triton ஒரு பெயர்வுத்திறன் மிக்க வன்பொருள் சுருக்க அடுக்காகச் (hardware abstraction layer) செயல்படுகிறது. தமது தனிப்பயன் kernel களை Triton இல் தரப்படுத்தும் AI அணிகள், தனித்தனி C++ குறியீட்டுத் தளங்களைப் பராமரிக்காமல், அதே குறியீட்டுத் தளத்தை NVIDIA மற்றும் AMD கொத்தணிகளில் (clusters) பயன்படுத்தலாம்.