Views
No views yet
1tl_methods = [
2 'PropagateNan', 'TRITON_MAX_TENSOR_NUMEL', 'abs', 'advance', 'arange',
3 'argmax', 'argmin', 'associative_scan', 'atomic_add', 'atomic_and',
4 'atomic_cas', 'atomic_max', 'atomic_min', 'atomic_or', 'atomic_xchg',
5 'atomic_xor', 'bfloat16', 'block_type', 'broadcast', 'broadcast_to',
6 'cast', 'cat', 'cdiv', 'ceil', 'clamp', 'const', 'const_pointer_type',
7 'constexpr', 'cos', 'cumprod', 'cumsum', 'debug_barrier', 'device_assert',
8 'device_print', 'div_rn', 'dot', 'dtype', 'erf', 'exp', 'exp2',
9 'expand_dims', 'fdiv', 'flip', 'float16', 'float32', 'float64',
10 'float8e4b15', 'float8e4b8', 'float8e4nv', 'float8e5', 'float8e5b16',
11 'floor', 'fma', 'full', 'function_type', 'histogram',
12 'inline_asm_elementwise', 'int1', 'int16', 'int32', 'int64', 'int8',
13 'interleave', 'join', 'load', 'log', 'log2', 'make_block_ptr', 'max',
14 'max_constancy', 'max_contiguous', 'maximum', 'min', 'minimum',
15 'multiple_of', 'num_programs', 'pair_uniform_to_normal', 'permute',
16 'philox', 'pi32_t', 'pointer_type', 'program_id', 'rand', 'rand4x',
17 'randint', 'randint4x', 'randn', 'randn4x', 'range', 'ravel', 'reduce',
18 'reshape', 'rsqrt', 'sigmoid', 'sin', 'softmax', 'sort', 'split', 'sqrt',
19 'sqrt_rn', 'static_assert', 'static_print', 'static_range', 'store',
20 'str_to_ty', 'sum', 'swizzle2d', 'tensor', 'trans', 'uint16', 'uint32',
21 'uint64', 'uint8', 'uint_to_uniform_float', 'umulhi', 'view', 'void',
22 'where', 'xor_sum', 'zeros', 'zeros_like'
23]
24
25
26def get_user_prompt(name, pytorch_impl):
27 prompt = f"""Convert this PyTorch module implementation into an equivalent Triton kernel:
28
29<torch_code>
30{pytorch_impl}
31</torch_code>
32
33The Triton kernel should:
341. Import torch, triton, and triton.language as tl and other necessary modules
352. Use @triton.jit decorator on the kernel implementation (not the entrypoint function)
363. Have proper grid and block sizes
374. Use a mask in the load/store operations
385. Use typed constants (tl.constexpr)
396. Handle tensor dimensions correctly
407. Return output matching PyTorch's implementation
418. Do not include any test code in your response, only the Triton kernel implementation and entrypoint function
42
43The triton.language (tl) module supports the following methods: {", ".join(tl_methods)}
44
45The entrypoint function must be named: {name}_triton
46The Triton kernel implementation (called by the entrypoint) must be named: {name}_kernel
47
48No computation logic should be done within the entrypoint function. All computation logic should be done within the Triton kernel implementation.
49
50The final generated code in the response must start with <triton_code> and end with </triton_code> tags."""
51
52 return prompt
53
54
55SYSTEM_PROMPT = """You are a helpful assistant that converts PyTorch code into Triton kernels."""
56
57messages = [
58 {"role": "system", "content": SYSTEM_PROMPT},
59 {"role": "user", "content": get_user_prompt(name, code)},
60]
61
62...1import torch
2import torch.nn as nn
3
4class Model(nn.Module):
5 """
6 Simple model that performs a LeakyReLU activation.
7 """
8 def __init__(self, negative_slope: float = 0.01):
9 """
10 Initializes the LeakyReLU module.
11
12 Args:
13 negative_slope (float, optional): The negative slope of the activation function. Defaults to 0.01.
14 """
15 super(Model, self).__init__()
16 self.negative_slope = negative_slope
17
18 def forward(self, x: torch.Tensor) -> torch.Tensor:
19 """
20 Applies LeakyReLU activation to the input tensor.
21
22 Args:
23 x (torch.Tensor): Input tensor of any shape.
24
25 Returns:
26 torch.Tensor: Output tensor with LeakyReLU applied, same shape as input.
27 """
28 return torch.nn.functional.leaky_relu(x, negative_slope=self.negative_slope)
29
30batch_size = 16
31dim = 16384
32
33def get_inputs():
34 x = torch.randn(batch_size, dim)
35 return [x]
36
37def get_init_inputs():
38 return [] # No special initialization inputs needed