tilelang-developer

Name: tilelang-developer
Author: yzlnew

15installs10forks

How to install tilelang-developer?

yzlnew/infra-skills/tilelang-developer

$npx -y skills add yzlnew/infra-skills --skill tilelang-developer

Installs into the current project.

Summary

Write, optimize, and debug high-performance AI compute kernels using TileLang (a Python DSL for GPU programming). Use when the user requests: (1) Writing custom GPU kernels for AI workloads (GEMM, Attention, MLA, etc.), (2) Optimizing existing TileLang code for NVIDIA, AMD, or As

Files · 1

View on GitHub

SKILL.md

1	# TileLang Developer
2
3	Write high-performance AI compute kernels using TileLang - a tile-based programming model that bridges the gap between CUDA's low-level control and high-level abstractions.
4
5	## When to Use This Skill
6
7	Use this skill when the user needs to:
8	- Implement custom GPU kernels for AI operations (matrix multiplication, attention mechanisms, etc.)
9	- Optimize performance-critical operators for modern GPUs (NVIDIA Ampere/Hopper, AMD MI300X, Ascend NPU)
10	- Debug TileLang code or resolve performance issues
11	- Port kernels across different hardware platforms
12	- Understand or modify existing TileLang implementations
13
14	## Kernel Development Workflow
15
16	Follow these steps when writing a TileLang kernel:
17
18	### Step 1: Analyze Requirements
19
20	Gather essential information:
21
22	Input/Output Specifications:
23	- Tensor shapes (M, N, K dimensions)
24	- Data types (float16, float32, bfloat16, int8)
25	- Memory layout (row-major, column-major)
26
27	Hardware Target:
28	- NVIDIA GPU (Ampere A100, Hopper H100, etc.)
29	- AMD GPU (MI300X, etc.)
30	- Huawei Ascend NPU
31
32	Performance Goals:
33	- Target throughput or latency
34	- Memory bandwidth constraints
35	- Comparison baseline (cuBLAS, vendor libraries)
36
37	Ask clarifying questions if details are missing.
38
39	### Step 2: Set Up Kernel Structure
40
41	Create the basic kernel scaffold:
42
43	```python
44	import tilelang
45	import tilelang.language as T
46
47	@tilelang.jit(target="cuda", out_idx=[2]) # Specify output indices
48	def kernel_name(M, N, K, block_M, block_N, block_K):
49	@T.prim_func
50	def main(
51	A: T.Buffer((M, K), "float16"),
52	B: T.Buffer((K, N), "float16"),
53	C: T.Buffer((M, N), "float16")
54	):
55	# Kernel logic will go here
56	pass
57
58	return main
59	```
60
61	Key decisions:
62	- `target`: "cuda" (NVIDIA), "hip" (AMD), or "cpu"
63	- `out_idx`: List indices of output parameters
64	- Block dimensions: Typical values are 64, 128, or 256
65
66	### Step 3: Define Grid and Memory Hierarchy
67
68	Set up computation grid and allocate memory:
69
70	```python
71	# Define grid dimensions
72	with T.Kernel(
73	T.ceildiv(N, block_N), # Grid X
74	T.ceildiv(M, block_M), # Grid Y
75	threads=128
76	) as (bx, by):
77
78	# Allocate shared memory (L1 cache)
79	A_shared = T.alloc_shared((block_M, block_K), "float16")
80	B_shared = T.alloc_shared((block_K, block_N), "float16")
81
82	# Allocate register fragments (accumulators)
83	C_local = T.alloc_fragment((block_M, block_N), "float32")
84
85	# CRITICAL: Apply swizzle layout to avoid bank conflicts
86	T.annotate_layout({
87	A_shared: T.make_swizzled_layout(A_shared),
88	B_shared: T.make_swizzled_layout(B_shared)
89	})
90	```
91
92	Memory hierarchy:
93	- Global Memory (HBM): Input/output tensors, slowest
94	- Shared Memory (L1): Explicitly managed cache, ~164KB on A100
95	- Registers: Fastest, used for accumulators and temporaries
96
97	Critical optimization: Always apply `T.make_swizzled_layout` to shared memory to eliminate bank conflicts.
98
99	### Step 4: Implement Computation Logic
100
101	Use TileLang primitives for data movement and computation:
102
103	```python
104	# Initialize accumulator
105	T.clear(C_local)
106
107	# Main computation loop with software pipelining
108	for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
109	# Load tiles from global to shared memory
110	T.copy(A[by * block_M, k * block_K], A_shared)
111	T.copy(B[k * block_K, bx * block_N], B_shared)
112
113	# Compute using Tensor Cores
114	T.gemm(A_shared, B_shared, C_local, transpose_B=False)
115
116	# Write results back
117	T.copy(C_local, C[by * block_M, bx * block_N])
118	```
119
120	Key primitives:
121	- `T.copy`: Intelligent data transfer (auto-selects cp.async, TMA, etc.)
122	- `T.gemm`: Matrix multiplication using Tensor Cores
123	- `T.Pipelined`: Software pipelining to overlap compute and memory transfer
124	- `T.Parallel`: Element-wise parallel operations
125
126	Pipeline stages:
127	- `num_stages=2`: Double buffering
128	- `num_stages=3`: Triple buffering (recommended for most workloads)
129	- `num_stages=4+`: Diminishing returns, increases shared memory usage
130
131	### Step 5: Validate and Test
132
133	Generate test code to verify correctness:
134
135	```python
136	# Example instantiation
137	func = kernel_name(
138	M=1024, N=1024, K=1024,
139	block_M=128, block_N=128, block_K=32
140	)
141
142	# Test against reference implementation
143	import torch
144	A = torch.randn(1024, 1024, dtype=torch.float16, device='cuda')
145	B = torch.randn(1024, 1024, dtype=torch.float16, device='cuda')
146	C_tilelang = torch.empty(1024, 1024, dtype=torch.float16, device='cuda')
147	C_refer

People who liked this also viewed

6 picks

Type

tilelang-developer

15installs10forks

View on GitHub

How to install tilelang-developer?

yzlnew/infra-skills/tilelang-developer

$npx -y skills add yzlnew/infra-skills --skill tilelang-developer

Installs into the current project.

Summary

Files · 1

View on GitHub

SKILL.md

1	# TileLang Developer
2
3	Write high-performance AI compute kernels using TileLang - a tile-based programming model that bridges the gap between CUDA's low-level control and high-level abstractions.
4
5	## When to Use This Skill
6
7	Use this skill when the user needs to:
8	- Implement custom GPU kernels for AI operations (matrix multiplication, attention mechanisms, etc.)
9	- Optimize performance-critical operators for modern GPUs (NVIDIA Ampere/Hopper, AMD MI300X, Ascend NPU)
10	- Debug TileLang code or resolve performance issues
11	- Port kernels across different hardware platforms
12	- Understand or modify existing TileLang implementations
13
14	## Kernel Development Workflow
15
16	Follow these steps when writing a TileLang kernel:
17
18	### Step 1: Analyze Requirements
19
20	Gather essential information:
21
22	Input/Output Specifications:
23	- Tensor shapes (M, N, K dimensions)
24	- Data types (float16, float32, bfloat16, int8)
25	- Memory layout (row-major, column-major)
26
27	Hardware Target:
28	- NVIDIA GPU (Ampere A100, Hopper H100, etc.)
29	- AMD GPU (MI300X, etc.)
30	- Huawei Ascend NPU
31
32	Performance Goals:
33	- Target throughput or latency
34	- Memory bandwidth constraints
35	- Comparison baseline (cuBLAS, vendor libraries)
36
37	Ask clarifying questions if details are missing.
38
39	### Step 2: Set Up Kernel Structure
40
41	Create the basic kernel scaffold:
42
43	```python
44	import tilelang
45	import tilelang.language as T
46
47	@tilelang.jit(target="cuda", out_idx=[2]) # Specify output indices
48	def kernel_name(M, N, K, block_M, block_N, block_K):
49	@T.prim_func
50	def main(
51	A: T.Buffer((M, K), "float16"),
52	B: T.Buffer((K, N), "float16"),
53	C: T.Buffer((M, N), "float16")
54	):
55	# Kernel logic will go here
56	pass
57
58	return main
59	```
60
61	Key decisions:
62	- `target`: "cuda" (NVIDIA), "hip" (AMD), or "cpu"
63	- `out_idx`: List indices of output parameters
64	- Block dimensions: Typical values are 64, 128, or 256
65
66	### Step 3: Define Grid and Memory Hierarchy
67
68	Set up computation grid and allocate memory:
69
70	```python
71	# Define grid dimensions
72	with T.Kernel(
73	T.ceildiv(N, block_N), # Grid X
74	T.ceildiv(M, block_M), # Grid Y
75	threads=128
76	) as (bx, by):
77
78	# Allocate shared memory (L1 cache)
79	A_shared = T.alloc_shared((block_M, block_K), "float16")
80	B_shared = T.alloc_shared((block_K, block_N), "float16")
81
82	# Allocate register fragments (accumulators)
83	C_local = T.alloc_fragment((block_M, block_N), "float32")
84
85	# CRITICAL: Apply swizzle layout to avoid bank conflicts
86	T.annotate_layout({
87	A_shared: T.make_swizzled_layout(A_shared),
88	B_shared: T.make_swizzled_layout(B_shared)
89	})
90	```
91
92	Memory hierarchy:
93	- Global Memory (HBM): Input/output tensors, slowest
94	- Shared Memory (L1): Explicitly managed cache, ~164KB on A100
95	- Registers: Fastest, used for accumulators and temporaries
96
97	Critical optimization: Always apply `T.make_swizzled_layout` to shared memory to eliminate bank conflicts.
98
99	### Step 4: Implement Computation Logic
100
101	Use TileLang primitives for data movement and computation:
102
103	```python
104	# Initialize accumulator
105	T.clear(C_local)
106
107	# Main computation loop with software pipelining
108	for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
109	# Load tiles from global to shared memory
110	T.copy(A[by * block_M, k * block_K], A_shared)
111	T.copy(B[k * block_K, bx * block_N], B_shared)
112
113	# Compute using Tensor Cores
114	T.gemm(A_shared, B_shared, C_local, transpose_B=False)
115
116	# Write results back
117	T.copy(C_local, C[by * block_M, bx * block_N])
118	```
119
120	Key primitives:
121	- `T.copy`: Intelligent data transfer (auto-selects cp.async, TMA, etc.)
122	- `T.gemm`: Matrix multiplication using Tensor Cores
123	- `T.Pipelined`: Software pipelining to overlap compute and memory transfer
124	- `T.Parallel`: Element-wise parallel operations
125
126	Pipeline stages:
127	- `num_stages=2`: Double buffering
128	- `num_stages=3`: Triple buffering (recommended for most workloads)
129	- `num_stages=4+`: Diminishing returns, increases shared memory usage
130
131	### Step 5: Validate and Test
132
133	Generate test code to verify correctness:
134
135	```python
136	# Example instantiation
137	func = kernel_name(
138	M=1024, N=1024, K=1024,
139	block_M=128, block_N=128, block_K=32
140	)
141
142	# Test against reference implementation
143	import torch
144	A = torch.randn(1024, 1024, dtype=torch.float16, device='cuda')
145	B = torch.randn(1024, 1024, dtype=torch.float16, device='cuda')
146	C_tilelang = torch.empty(1024, 1024, dtype=torch.float16, device='cuda')
147	C_refer

People who liked this also viewed

6 picks

Type