tilekernels-gpu-kernels

# TileKernels GPU Kernel Library > Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection. TileKernels is a high-performance GPU kernel library for LLM operations (MoE routing, FP8/FP4 quantization, transpose, engram gating, Manifold HyperConnection) written in TileLang — a Python DSL for expressing GPU kernels with automatic optimization. Kernels target NVIDIA SM90/SM100 (Hopper/Blackwell) architectures and approach hardware performance limits. ## Requirements - Python 3.10+ - PyTorch 2.10+ - TileLang 0.1.9+ - NVIDIA SM90 or SM100 GPU (H100/H200/B100/B200) - CUDA Toolkit 13.1+ ## Installation ```bash # Development install (recommended for extending/modifying kernels) pip install -e ".[dev]" # Release install pip install tile-kernels ``` ## Project Structure ``` tile_kernels/ ├── moe/ # MoE routing: top-k selection, token-to-expert mapping, weight normalization ├── quant/ # FP8/FP4/E5M6 quantization with fused SwiGLU ops ├── transpose/ # Batched matrix transpose ├── engram/ # Engram gating: fused RMSNorm, forward/backward, weight gradient reduction ├── mhc/ # Manifold HyperConnection: Sinkhorn normalization, mix split/apply ├── modeling/ # High-level torch.autograd.Function wrappers ├── torch/ # PyTorch reference implementations for validation └── testing/ # Test and benchmark utilities ``` ## Key Modules and Usage ### MoE Routing Kernels ```python import torch from tile_kernels.moe import ( topk_gating, # Top-k expert selection and scoring token_to_expert_map, # Token-to-expert mapping fused_expand_reduce, # Fused expansion/reduction weight_normalize, # Weight normalization ) # Top-k gating: select top-k experts per token # logits: [num_tokens, num_experts] logits = torch.randn(1024, 256, device="cuda", dtype=torch.float32) topk_weights, topk_indices = topk_gating(logits, top_k=8) # topk_weights: [num_tokens, top_k], topk_indices: [num_tokens, top_k] ``` ### Quantization Kernels ```python import torch from tile_kernels.quant import ( per_token_cast_fp8, # Per-token FP8 quantization per_block_cast_fp8, # Per-block FP8 quantization per_channel_cast_fp8, # Per-channel FP8 quantization fused_swiglu_quant_fp8, # Fused SwiGLU + FP8 quantization ) # Per-token FP8 quantization x = torch.randn(1024, 4096, device="cuda", dtype=torch.bfloat16) x_fp8, scale = per_token_cast_fp8(x) # x_fp8: [1024, 4096] in torch.float8_e4m3fn # scale: [1024, 1] per-token scales # Per-block FP8 quantization (common for weight quantization) w = torch.randn(8192, 4096, device="cuda", dtype=torch.bfloat16) w_fp8, scale = per_block_cast_fp8(w, block_size=128) # Fused SwiGLU + FP8 cast (saves memory bandwidth) gate = torch.randn(1024, 8192, device="cuda", dtype=torch.bfloat16) up = torch.randn(1024, 8192, device="cuda", dtype=torch.bfloat16) out_fp8, scale = fused_swiglu_quant_fp8(gate, up) ``` ### Transpose Kernels ```python import torch from tile_kernels.transpose import batched_transpose # Batched transpose for MoE weight manipulation # x: [batch, M, N] x = torch.randn(32, 1024, 4096, device="cuda", dtype=torch.bfloat16) x_T = batched_transpose(x) # x_T: [batch, N, M] = [32, 4096, 1024] ``` ### Engram Gating Kernels ```python import torch from tile_kernels.engram import ( engram_gate_forward, # Forward pass with fused RMSNorm engram_gate_backward, # Backward pass engram_weight_grad, # Weight gradient reduction ) # Forward pass hidden = torch.randn(1024, 2048, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 2048, device="cuda", dtype=torch.bfloat16) output, norm_hidden = engram_gate_forward(hidden, weight) ``` ### Manifold HyperConnection (mHC) Kernels ```python import torch from tile_kernels.mhc import ( sinkhorn_normalize, # Sinkhorn normalization mhc_mix_split, # Mix splitting mhc_mix_apply, # Mix application ) # Sinkhorn normalization for connection weights conn_weights = torch.randn(8, 64, device="cuda", dtype=torch.float32) normalized = sinkhorn_normalize(conn_weights, num_iters=20) ``` ### High-Level Modeling Layers ```python import torch from tile_kernels.modeling import EngramGateLayer, MHCPipeline # EngramGate as a trainable nn.Module-compatible layer # Uses torch.autograd.Function internally gate_layer = EngramGateLayer( hidden_size=2048, num_experts=256, ).cuda() hidden_states = torch.randn(1024, 2048, device="cuda", dtype=torch.bfloat16) gate_output = gate_layer(hidden_states) # Manifold HyperConnection pipeline mhc = MHCPipeline( num_connections=8, hidden_size=2048, ).cuda() ``` ## Testing ```bash # Test a single module (correctness only, 4 parallel workers) pytest tests/transpose/test_transpose.py -n 4 # Test with benchmarking pytest tests/transpose/test_transpose.py --run-benchmark # Test MoE kernels pytest tests/moe/ -n 4 # Test quantization kernels pytest tests/quant/ -n 4 # Test engram kernels pytest tests/engram/ -n 4 # Full pressure test (all tests, 2 repetitions, 4 workers) TK_FULL_TEST=1 pytest -n 4 --count 2 # Test specific quantization variant pytest tests/quant/test_fp8_cast.py -n 4 --run-benchmark ``` ## Common Patterns ### Pattern: Fused MoE Forward Pass ```python import torch from tile_kernels.moe import topk_gating, token_to_expert_map from tile_kernels.quant import per_token_cast_fp8 def moe_dispatch(hidden_states, gate_weight, top_k=8): """Full MoE dispatch using TileKernels.""" # 1. Compute gating logits logits = torch.mm(hidden_states, gate_weight.T) # [T, E] # 2. Top-k expert selection topk_weights, topk_indices = topk_gating(logits, top_k=top_k) # 3. Build token-to-expert routing map routing_map = token_to_expert_map(topk_indices, num_experts=gate_weight.shape[0]) # 4. Quantize activations before expert computation hidden_fp8, scale = per_token_cast_fp8(hidden_states) return hidden_fp8, scale, topk_weights, routing_map ``` ### Pattern: Using PyTorch Reference Implementations for Validation ```python import torch from tile_kernels.quant import per_token_cast_fp8 from tile_kernels.torch import per_token_cast_fp8 as per_token_cast_fp8_ref # Compare kernel output vs PyTorch reference x = torch.randn(512, 4096, device="cuda", dtype=torch.bfloat16) out_kernel, scale_kernel = per_token_cast_fp8(x) out_ref, scale_ref = per_token_cast_fp8_ref(x) # Validate torch.testing.assert_close( out_kernel.float(), out_ref.float(), atol=1e-2, rtol=1e-2 ) print("Kernel matches reference ✓") ``` ### Pattern: Benchmarking a Kernel ```python import torch from tile_kernels.testing import benchmark_kernel from tile_kernels.transpose import batched_transpose x = torch.randn(64, 4096, 4096, device="cuda", dtype=torch.bfloat16) # Using the testing utility result = benchmark_kernel( fn=batched_transpose, args=(x,), warmup=25, rep=100, ) print(f"Latency: {result.mean:.3f} ms, Bandwidth: {result.gbps:.1f} GB/s") ``` ### Pattern: Custom TileLang Kernel (extending the library) ```python # tile_kernels follow TileLang DSL patterns import tilelang import tilelang.language as T def make_elementwise_scale_kernel(M, N, dtype="float16"): @T.prim_func def scale_kernel( A: T.Buffer((M, N), dtype), scale: T.Buffer((M,), "float32"), B: T.Buffer((M, N), dtype), ): # TileLang kernel body for i, j in T.grid(M, N): B[i, j] = T.cast( T.cast(A[i, j], "float32") * scale[i], dtype ) return scale_kernel # Compile and use kernel = tilelang.compile(make_elementwise_scale_kernel(1024, 4096)) ``` ## Architecture-Specific Notes - **SM90 (Hopper: H100/H200)**: Full support, primary target - **SM100 (Blackwell: B100/B200)**: Full support - Kernels use hardware-specific features (tensor memory accelerator, async copy, warp-specialized pipelines) — do NOT run on older GPUs (Ampere/Ada) ## Troubleshooting ### CUDA Architecture Mismatch ``` RuntimeError: CUDA error: no kernel image is available for execution on the device ``` → You need SM90 or SM100. Check with: `python -c "import torch; print(torch.cuda.get_device_capability())"` ### TileLang Version Mismatch ``` ImportError: cannot import name 'xyz' from 'tilelang' ``` → Ensure TileLang >= 0.1.9: `pip install tilelang>=0.1.9` ### CUDA Toolkit Version ``` error: identifier "__nv_fp8_e4m3" is undefined ``` → Requires CUDA 13.1+. Check: `nvcc --version` ### Out of Shared Memory → Kernels are tuned for specific tile sizes. If you hit shared memory limits, reduce batch size or sequence length, or file an issue. ### Running Tests Without Benchmark Flag ```bash # Benchmarks are opt-in to avoid slow CI pytest tests/ -n 4 # Fast correctness only pytest tests/ -n 4 --run-benchmark # Include performance numbers ``` ## Citation ```bibtex @misc{tilekernels, title={TileKernels}, author={Xiangwen Wang, Chenhao Xu, Huanqi Cao, Rui Tian, Weilin Zhao, Kuai Yu and Chenggang Zhao}, year={2026}, publisher = {GitHub}, howpublished = {\url{https://github.com/deepseek-ai/TileKernels}}, } ```