# Project export: KERneL This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small. ## Project metadata - Hackathon: TreeHacks 2025 - Tagline: Efficient CUDA kernels unlock NVIDIA GPUs' full potential to accelerate AI. But only highly skilled developers can code them. That's why we created a CUDA dev tool for all, with SOTA reasoning LLMs. - Devpost: https://devpost.com/software/kernel-ls5pqa - GitHub: https://github.com/radi-cho/KERneL - Video: https://www.youtube.com/embed/-GqtCtDWJmc?enablejsapi=1&hl=en_US&rel=0&start=&version=3&wmode=transparent - Team: 4 GitHub contributor(s) — Radi Cho (46 commits), Kyle Kun Hyung Roh (21 commits), Luke (18 commits), lalala-e (4 commits) ## Devpost submission (written by the team) ### Inspiration Optimizing code for GPUs is notoriously difficult despite its great potential in AI training and high-performance computing. Writing efficient CUDA kernels requires deep expertise, making GPU acceleration inaccessible to many developers. We were inspired by the idea of democratizing GPU programming—leveraging LLMs to generate optimized GPU kernels from high-level descriptions. Our goal? Lower the barrier of entry, so that anyone can harness the power of NVIDIA GPUs without needing years of CUDA experience. ### What it does KERneL is an AI-powered kernel generation tool that takes high-level code descriptions and transforms them into optimized CUDA kernels. By leveraging large language models (LLMs), KERneL automates the process of writing efficient GPU code, allowing developers to: ✅ Generate CUDA kernels from simple prompts ✅ Optimize performance with AI-assisted tuning ✅ Reduce the learning curve for GPU acceleration With KERneL, more developers can tap into GPU computing for AI, graphics, and scientific simulations effortlessly. ### How we built it Backend: We utilized NVIDIA Build Cloud to integrate LLMs like Qwen 2.5 7B and DeepSeek R1 for kernel preprocessing and generation, as well as OpenAI’s state-of-the-art tools for additional kernel optimization. Frontend: A user-friendly interface built with Streamlit allows seamless interaction for generating, testing, and refining CUDA kernels. Cloud Computing: Hosted on Brev.dev Cloud with H100 GPUs for high-performance compute tasks, integrated through an ngrok tunnel for secure access. ### Challenges we ran into Backend-frontend integration: Coordinating data flow between the Flask API, compute instances, and the NVIDIA Build Cloud was complex. Kernel validation: Ensuring the generated CUDA kernels met performance expectations required extensive testing and refinement. LLM prompt engineering: Guiding the language models to understand PyTorch timing details and compiler logs was time-intensive. ### Accomplishments we're proud of Achieves lower latency compared to a variety of PyTorch Dynamo compiled models Provide great access to beginners interested in fully leveraging FLOP utilization Successful generation and compilation of CUDA kernel code We fixed the challenges! Above all, we are proud that we cooperatively finished the project and helped each others out to augment our skillsets. ### What's next Next updates will include the following features: Ensure live interaction with user in computational graph, allowing them to visualize which parts will be fused. Copilot mode with assistive agent. ## README (from the GitHub repository) # CUDA Kernel Generation with PyTorch Welcome to our project! This repository provides examples and tools to help you explore and generate efficient CUDA kernels using PyTorch. Whether you're new to CUDA programming or a seasoned developer, you'll find examples and test cases to deepen your understanding of CUDA kernel generation. ## 🚀 Introduction CUDA (Compute Unified Device Architecture) is the backbone of modern GPU programming, enabling highly efficient computation for deep learning, scientific simulations, and more. Writing efficient CUDA kernels, however, is often considered complex and time-consuming. This project provides: - Easy-to-follow examples of custom CUDA kernel generation. - Test cases for advanced CUDA-powered operations, like attention mechanisms and relative position embeddings. - Tools for benchmarking and optimizing PyTorch code with CUDA. ## 🌟 Features - **Custom CUDA Kernels**: Examples of transforming PyTorch operations into CUDA for faster performance. - **Advanced Attention Mechanisms**: Relative position embeddings, sliding window attention, PrefixLM, ALiBi, and more. - **Customizable Block Masks**: Define unique attention patterns with CUDA. - **Ease of Use**: All examples come with ready-to-use PyTorch models and initialization inputs. --- ## 📘 Getting Started Follow these instructions to get started with the repository. ### Prerequisites - Python 3.8+ - PyTorch with CUDA support - A CUDA-enabled GPU ### Test Cases #### Extremely basic example ```python import torch import torch.nn as nn import time #Define a simple PyTorch module class SequentialOperations(nn.Module): def __init__(self): super(SequentialOperations, self).__init__() self.relu = nn.ReLU() def forward(self, x): # Sequential operations x = x.cos() x = x.square() x = x.sin() x = self.relu(x) return x ``` #### Relative Position Embeddings ```python import torch import torch.nn as nn import torch.nn.functional as F # Placeholder for the custom CUDA function def relative_attention(query, key, value): """ Placeholder function for attention with relative position encoding. This function will be replaced by a CUDA kernel. Args: query (torch.Tensor): Query tensor of shape (B, H, S, D). key (torch.Tensor): Key tensor of shape (B, H, S, D). value (torch.Tensor): Value tensor of shape (B, H, S, D). Returns: torch.Tensor: Output tensor of shape (B, H, S, D). """ B, H, S, D = query.shape scores = torch.einsum("bhqd,bhkd->bhqk", query, key) # Compute QK^T for q_idx in range(S): for kv_idx in range(S): scores[:, :, q_idx, kv_idx] += q_idx - kv_idx # Apply relative position bias attention_weights = F.softmax(scores, dim=-1) # Compute weighted sum return torch.einsum("bhqk,bhvd->bhqd", attention_weights, value) class Model(nn.Module): """ Model that performs scaled dot-product attention with relative position encoding. """ def __init__(self): super(Model, self).__init__() def forward(self, query, key, value): """ Compute attention with relative position encoding. Args: query (torch.Tensor): Query tensor of shape (B, H, S, D). key (torch.Tensor): Key tensor of shape (B, H, S, D). value (torch.Tensor): Value tensor of shape (B, H, S, D). Returns: torch.Tensor: Output tensor of shape (B, H, S, D). """ return relative_attention(query, key, value) #Define the batch size, number of heads, sequence length, and embedding dimension B = 8 # Batch size H = 16 # Number of attention heads S = 2048 # Sequence length D = 128 # Embedding dimension per head def get_inputs(): """ Generate random input tensors for query, key, and value. """ query = torch.randn(B, H, S, D).cuda() key = torch.randn(B, H, S, D).cuda() value = torch.randn(B, H, S, D).cuda() return [query, key, value] def get_init_inputs(): """ No special initialization inputs needed for this model. """ return [] ``` #### Custom Sliding Window Attention ```python import torch import torch.nn as nn import torch.nn.functional as F SLIDING_WINDOW = 2048 #Placeholder for the custom CUDA function def sliding_window_attention(query, key, value): """ Placeholder function for sliding window causal attention. This function will be replaced by a CUDA kernel. Args: query (torch.Tensor): Query tensor of shape (B, H, S, D). key (torch.Tensor): Key tensor of shape (B, H, S, D). value (torch.Tensor): Value tensor of shape (B, H, S, D). Returns: torch.Tensor: Output tensor of shape (B, H, S, D). """ B, H, S, D = query.shape scores = torch.einsum("bhqd,bhkd->bhqk", query, key) # Compute QK^T for q_idx in range(S): for kv_idx in range(S): causal_mask = q_idx >= kv_idx window_mask = q_idx - kv_idx <= SLIDING_WINDOW scores[:, :, q_idx, kv_idx] *= (causal_mask & window_mask) attention_weights = F.softmax(scores, dim=-1) return torch.einsum("bhqk,bhvd->bhqd", attention_weights, value) # Compute weighted sum class Model(nn.Module): """ Model that performs sliding window causal attention. """ def __init__(self): super(Model, self).__init__() def forward(self, query, key, value): """ Compute sliding window causal attention. Args: query (torch.Tensor): Query tensor of shape (B, H, S, D). key (torch.Tensor): Key tensor of shape (B, H, S, D). value (torch.Tensor): Value tensor of shape (B, H, S, D). Returns: torch.Tensor: Output tensor of shape (B, H, S, D). """ return sliding_window_attention(query, key, value) #Define the batch size, number of heads, sequence length, and embedding dimension B = 8 # Batch size H = 16 # Number of attention heads S = 2048 # Sequence length D = 128 # Embedding dimension per head def get_inputs(): """ Generate random input tensors for query, key, and value. """ query = torch.randn(B, H, S, D).cuda() key = torch.randn(B, H, S, D).cuda() value = torch.randn(B, H, S, D).cuda() return [query, key, value] def get_init_inputs(): """ No special initialization inputs needed for this model. """ return [] ``` #### PrefixLM Attention ```python import torch import torch.nn as nn import torch.nn.functional as F #Placeholder for the custom CUDA function def prefix_attention(query, key, value, prefix_length): """ Placeholder function for PrefixLM attention with dynamic prefix-based and causal masking. This function will be replaced by a CUDA kernel. Args: query (torch.Tensor): Query tensor of shape (B, H, S, D). key (torch.Tensor): Key tensor of shape (B, H, S, D). value (torch.Tensor): Value tensor of shape (B, H, S, D). prefix_length (torch.Tensor): Tensor of shape (B,) indicating the prefix length for each sequence. Returns: torch.Tensor: Output tensor of shape (B, H, S, D). """ B, H, S, D = query.shape scores = torch.einsum("bhqd,bhkd->bhqk", query, key) # Compute QK^T for b in range(B): for q_idx in range(S): for kv_idx in range(S): causal_mask = q_idx >= kv_idx prefix_mask = kv_idx < prefix_length[b] scores[b, :, q_idx, kv_idx] *= (causal_mask or prefix_mask) attention_weights = F.softmax(scores, dim=-1) return torch.einsum("bhqk,bhvd->bhqd", attention_weights, value) # Compute weighted sum class Model(nn.Module): """ Model that performs PrefixLM attention with dynamic prefix-based and causal masking. """ def __init__(self): super(Model, self).__init__() [README truncated for size] ## Detected evidence (automated analysis) Indexed codebase: 20 recognized source files, 87 KB. - C++ (language) — detected in the code - Python (language) — detected in the code - Streamlit (technology) — claimed on Devpost, not found in the code ## Codebase structure (from repository index) ### Files (24 of 24) ``` adding_red.py alternate_prompt.txt api_query.py kernel_api.py llm_query_sample.py prompt_construction.py prompt_postfix.txt prompt_prefix.txt prompts/model_cot_fuse_gelu.py prompts/model_cot_mnist2.py prompts/model_cot_tiled_matmul.py prompts/model_ex_fuse_gelu.py prompts/model_ex_mnist2.py prompts/model_ex_tiled_matmul.py prompts/model_new_ex_fuse_gelu.py prompts/model_new_ex_mnist2.py prompts/model_new_ex_tiled_matmul.py README.md sample/kernel.cpp sample/kernel.cu sample/kernel.py u3.py ui.py utils.py ``` ### Dependencies No dependency index available. ### Recent commits (newest first) - Update utils.py - Update api_query.py - Update README.md - Update num_trials. - Safety check on UI. - Updates. - Merge branch 'master' of https://github.com/radi-cho/KERneL - UI updates. - updated api for reasoning api query - Merge branch 'master' of https://github.com/radi-cho/KERneL - Reorder items in UI. - updated api query - updated readme - Merge branch 'master' of https://github.com/radi-cho/KERneL - Update UI outputs. - updated README - Update printing options. - Merge branch 'master' of https://github.com/radi-cho/KERneL - Update output tracking. - readme ## Key source files (fetched from GitHub, selected and truncated for size) ### llm_query_sample.py ```python def llm_query(pytorch_function, additional_context): function_name = "diag_matmul_cuda" cuda = """ #include #include __global__ void diag_matmul_kernel( const float* diag, const float* mat, float* out, const int N, const int M) { const int row = blockIdx.y * blockDim.y + threadIdx.y; const int col = blockIdx.x * blockDim.x + threadIdx.x; if (row < N && col < M) { out[row * M + col] = diag[row] * mat[row * M + col]; } } torch::Tensor diag_matmul_cuda(torch::Tensor diag, torch::Tensor mat) { const int N = diag.size(0); const int M = mat.size(1); auto out = torch::zeros({N, M}, mat.options()); const dim3 threads(16, 16); const dim3 blocks((M + threads.x - 1) / threads.x, (N + threads.y - 1) / threads.y); diag_matmul_kernel<<>>( diag.data_ptr(), mat.data_ptr(), out.data_ptr(), N, M); return out; } """ cpp = "torch::Tensor diag_matmul_cuda(torch::Tensor diag, torch::Tensor mat);" return function_name, cuda, cpp ``` ### utils.py ```python import re from typing import Union, Tuple, List import os import time from datetime import datetime from openai import OpenAI NVIDIA_API_KEY = #redacted NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1" def initialize_client(api_key = None, base_url = None): if api_key is None: api_key = os.getenv("API_KEY") if not api_key: raise ValueError("API key not found. Please set the API_KEY environment variable.") return OpenAI( api_key = api_key ) else: base_url = NVIDIA_BASE_URL return OpenAI( base_url = base_url, api_key = api_key ) def extract_method_name(cpp_signature: str) -> str: """ Extracts the method name from a C++ function signature. Args: cpp_signature (str): The C++ function signature as a string. Returns: str: The method name, or an empty string if not found. """ # Use a regular expression to find the method name match = re.search(r'[\w:]+::(\w+)\s*\(', cpp_signature) if match: return match.group(1) # Group 1 is the method name else: # If the namespace (::) is absent, handle it separately match = re.search(r'(\w+)\s*\(', cpp_signature) if match: return match.group(1) # Group 1 is the method name return "" # Return an empty string if no match is found def extract_from_text(full_response: str, flags: List[Tuple[str, str]]) -> Tuple[str, str, str]: try: results = [] for idx, (flag_start, flag_stop) in enumerate(flags): idx_start = full_response.find(flag_start) + len(flag_start) idx_end = full_response.find(flag_stop) results.append(full_response[idx_start:idx_end].strip() if idx_start != -1 and idx_end != -1 else "") return tuple(results) except Exception as e: print(f"Error parsing response: {e}") return tuple([None for _ in flags]) def save_reasoning(results: List[str], filename: str = "results.txt"): """ Saves the generated results to a specified file. Args: results (List[str]): The list of string results to save. filename (str): The filename where results will be saved. """ # Ensure the output directory exists directory = os.path.dirname(filename) if directory and not os.path.exists(directory): os.makedirs(directory) with open(filename, "w") as file: for idx, result in enumerate(results, start=1): file.write(f"Result {idx}:\n") file.write(f"{result}\n") file.write("-" * 50 + "\n") print(f"Results saved to {filename}") def save_kernels(kernels: List[Tuple[str, str, str]], directory="sample"): # Get the current date and time for the filename timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") directory = f"{directory}/{timestamp}" # Ensure the output directory exists if not os.path.exists(directory): os.makedirs(directory) for idx, (method_name, cuda_kernel, cpp_kernel_signature) in enumerate(kernels): cpp_filename = os.path.join(directory, f"/kernel_{idx}.cpp") cu_filename = os.path.join(directory, f"kernel_{idx}.cu") with open(cpp_filename, "w") as cpp_file: cpp_file.write(cpp_kernel_signature) print(f"Saved C++ kernel to {cpp_filename}") # Save the cuda_kernel to a file with open(cu_filename, "w") as cu_file: cu_file.write(cuda_kernel) print(f"Saved CUDA kernel to {cu_filename}") ``` ### kernel_api.py ```python import sys import torch import torch.nn as nn import importlib.util from uuid import uuid4 from flask import Flask, request, jsonify from torch.utils.cpp_extension import load_inline from api_query import generate_multiple_kernels, get_init_and_input_function from llm_query_sample import llm_query torch.set_printoptions(edgeitems=1, threshold=10, linewidth=100) TASKS = {} def initialize_python_module(source, module_name="dynamic_kernel_module"): try: spec = importlib.util.spec_from_loader(module_name, loader=None) dynamic_module = importlib.util.module_from_spec(spec) exec(source, dynamic_module.__dict__) sys.modules[module_name] = dynamic_module return True, dynamic_module except Exception as e: return False, str(e) def initialize_kernel_module(cuda_sources, cpp_sources, function_name, kernel_name="dynamic_cuda_kernel"): try: kernel_module = load_inline( name=kernel_name, cuda_sources=cuda_sources, cpp_sources=cpp_sources, functions=[function_name], verbose=True ) return True, kernel_module except Exception as e: return False, str(e) def time_execution_with_cuda_event(model, inputs, num_trials): device = torch.cuda.current_device() if isinstance(model, torch.nn.Module): model.to(device=device) inputs = [inp.to(device=device) for inp in inputs] elapsed_times = 0 for trial in range(num_trials): start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) start_event.record() if trial == num_trials - 1: output = model(*inputs) else: model(*inputs) end_event.record() torch.cuda.synchronize(device=device) elapsed_time_ms = start_event.elapsed_time(end_event) elapsed_times += elapsed_time_ms return elapsed_times / num_trials, output app = Flask(__name__) @app.route('/initialize_task', methods=['POST']) def initialize_task(): try: data = request.get_json() python_source = data.get("python_source", "") num_trials = data.get("num_trials", 1) model_init_code, get_input_function_code = get_init_and_input_function(python_source) imports = """ import torch import torch.nn as nn import numpy as np """ python_source = imports + "\n\n" + python_source + "\n\n" + model_init_code + "\n" + get_input_function_code initialized, result = initialize_python_module(python_source) if initialized: model, get_inputs = getattr(result, "model"), getattr(result, "get_inputs") inputs = get_inputs() task_id = str(uuid4()) average_time, output = time_execution_with_cuda_event(model, inputs, num_trials) TASKS[task_id] = [python_source, model, inputs, [["", average_time, output]]] response = { "status": "Task initialized successfully", "torch_time": average_time, "task_id": task_id, "output": str(output) } return jsonify(response), 200 except Exception as e: response = {"error": str(e)} return jsonify(response), 400 @app.route('/get_kernel', methods=['POST']) def get_kernel(): try: data = request.get_json() num_trials = data.get("num_trials", 1) task_id = data.get("task_id", "") source, _, inputs, history = TASKS[task_id] additional_context = "" if len(history) > 1: additional_context = f"Your previous attempt with source: {history[-1][0]}\n\n runs in time {history[-1][1]} compared to a native PyTorch compilation which runs in time {history[0][1]}." function_name, cuda_sources, cpp_sources = generate_multiple_kernels(pytorch_function=source, additional_context=additional_context) initialized, result = initialize_kernel_module(cuda_sources, cpp_sources, function_name) if initialized: model = getattr(result, function_name) average_time, output = time_execution_with_cuda_event(model, inputs, num_trials) response = { "task_id": task_id, "status": "Kernel compiled successfully", "kernel_code": cuda_sources, "kernel_time": average_time, "output": str(output) } TASKS[task_id][3].append([cuda_sources, average_time, output]) return jsonify(response), 200 else: response = { "task_id": task_id, "status": "Failed during kernel compilation", "error": result } return jsonify(response), 200 except Exception as e: return jsonify({ "task_id": task_id, "error": str(e) }), 400 if __name__ == '__main__': app.run(debug=False) ``` ### u3.py ```python import streamlit as st import requests import random import types from streamlit_ace import st_ace import time import torch from torch import nn import torchvision # Streamlit UI Setup - Remove Top Blank Space st.set_page_config(page_title="Python to CUDA Kernel Optimization", layout="wide") # 🔹 Two-Column Layout (Editors at the Top) col1, col2 = st.columns(2) # 🔹 Left Side: Python Code Input (Editable, Starts at the Top) with col1: st.markdown("✏️ **Enter Python Code**", unsafe_allow_html=True) python_code = st_ace( language="python", theme="monokai", placeholder="Write or paste your Python code here...", height=400, key="python_code_editor" ) # **Tensor Input Dimension Section** st.markdown("📏 **Specify Tensor Input Dimensions**", unsafe_allow_html=True) tensor_dim = st.text_input("Enter tensor dimensions (e.g., 1, 3, 224, 224)", "1, 3, 224, 224") # 🔹 Hardware Selection & Optimization Duration st.markdown("⚙️ **Optimization Settings**", unsafe_allow_html=True) hardware = st.selectbox("💻 Select Hardware", ["NVIDIA H100", "NVIDIA A100"]) num_trials = st.slider("🔁 Number of Trials", 10, 200, 100) # 🔹 Right Side: CUDA Kernel Code Output (Read-Only, Starts at the Top) with col2: st.markdown("⚡ **Generated CUDA Kernel Code**", unsafe_allow_html=True) # Placeholder for CUDA Kernel Code cuda_code_container = st.empty() # 🔹 Performance Metrics (Below CUDA Code) st.markdown("📊 **Performance Metrics**", unsafe_allow_html=True) torch_time_text = st.empty() kernel_time_text = st.empty() from torchview import draw_graph # ensure torchview is installed st.markdown("🖼 **Model Computational Graph**", unsafe_allow_html=True) if st.button("🖥 Generate Computational Graph"): # Show a loading message for exactly 3 seconds. with st.spinner("Loading computational graph..."): time.sleep(2) # 1. Parse tensor dimensions and create dummy input. try: input_dims = tuple(map(int, tensor_dim.split(','))) dummy_input = torch.rand(input_dims) except Exception as e: st.error(f"Error parsing tensor dimensions: {e}") st.stop() # 2. Execute the user code in a fresh module namespace. user_module = types.ModuleType("user_module") try: exec(compile(python_code, "", "exec"), user_module.__dict__) except Exception as e: st.error(f"Error executing user code: {e}") st.stop() # 3. Extract the model (the first nn.Module instance found) model = next((v for k, v in user_module.__dict__.items() if isinstance(v, nn.Module)), None) if model is None: st.error("No valid nn.Module instance found in the provided code.") st.write("Module keys:", list(user_module.__dict__.keys())) st.stop() try: # 4. Use TorchView to generate the computational graph. # draw_graph returns a path to an image file by default, # so we can then display that image using st.image. graph_path = draw_graph(model, input_size=input_dims, roll=True) st.write(graph_path.visual_graph) except Exception as e: st.error(f"TorchView error: {e}") st.exception(e) # 🔹 Button to Send Python Code st.markdown("⚙️ **Transform Python to CUDA Kernel**", unsafe_allow_html=True) if st.button("🚀 Generate kernel"): if python_code.strip(): st.info("📡 Initializing task on server...") try: # 🔹 Step 1: Send Python Code to Initialize Task payload = {"python_source": python_code, "num_trials": num_trials} response = requests.post( "https://d4fa-209-20-157-139.ngrok-free.app/initialize_task", json=payload ) if response.status_code == 200: data = response.json() task_id = data.get("task_id") torch_time = data.get("torch_time", "N/A") torch_time_text.markdown(f"🔥 **Torch Execution Time:** `{torch_time} ms`") if task_id: st.info("📡 Initializing CUDA kernel...") # 🔹 Step 2: Request CUDA Kernel Initialization kernel_payload = {"task_id": task_id, "num_trials": num_trials} kernel_response = requests.post( "https://d4fa-209-20-157-139.ngrok-free.app/get_kernel", json=kernel_payload ) if kernel_response.status_code == 200: kernel_data = kernel_response.json() kernel_time = kernel_data.get("kernel_time", "N/A") kernel_time_text.markdown(f"⚡ **CUDA Execution Time:** `{kernel_time} ms`") if "kernel_code" in kernel_data: # ✅ Render CUDA Kernel Code dynamically cuda_code_container.markdown("```cpp\n" + kernel_data["kernel_code"] + "\n```") st.success("✅ CUDA kernel compiled successfully!") else: st.warning("⚠️ CUDA kernel not received.") else: st.error(f"❌ Kernel error: {kernel_response.text}") else: st.error("❌ Task ID missing from response.") else: st.error(f"❌ Error from server: {response.status_code} - {response.text}") except requests.exceptions.RequestException as e: st.error(f"❌ Failed to connect to server: {e}") else: st.warning("⚠️ Please enter Python code before applying.") ``` ### prompt_construction.py ```python import os PROBLEM_STATEMENT_CLEANED = """You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups.\n\nYou have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.\n """ PROBLEM_INSTRUCTION_CLEANED = """ Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! \n """ PROBLEM_STATEMENT = """You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. \n You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.\n """ PROBLEM_INSTRUCTION = """ Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! \n """ def read_file(file_path) -> str: if not os.path.exists(file_path): print(f"File {file_path} does not exist") return "" try: with open(file_path, "r") as file: return file.read() except Exception as e: print(f"Error reading file {file_path}: {e}") return "" def prompt_generate_ex_with_CoT_template(ref_arch_src: str, cot_example: str) -> str: """ Generate a prompt with a CoT example following a template Avaliable CoT examples: - ex_fuse_gelu: fused gelu - ex_mnist2: fused convolutions and relus - ex_tiled_matmul: tiled matrix multiplication """ # I updated this to allow CoT. Also explicilty state think step by step. PROBLEM_INSTRUCTION_COT = """ Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Do not output testing code. In the end, make sure the final code block contains code for output architecture ModelNew with cuda code.\n Let's think step by step.\n """ prompt = PROBLEM_STATEMENT_CLEANED assert cot_example in ["ex_fuse_gelu", "ex_mnist2", "ex_tiled_matmul"] REPO_TOP_PATH = "" # k = 2 example_fuse_gelu = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_ex_fuse_gelu.py") ) example_fuse_gelu_cot = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_cot_fuse_gelu.py") ) example_fuse_gelu_new = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_new_ex_fuse_gelu.py") ) example_fuse_gelu_desc = "This given architecture is for a fused gelu: " # k = 3 example_mnist2 = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_ex_mnist2.py") ) example_mnist2_cot = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_cot_mnist2.py") ) example_mnist2_new = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_new_ex_mnist2.py") ) exmaple_mnist2_desc = "This given architecture is for a model with fused convolutions and relus: " # k = 4 example_tiled_matmul = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_ex_tiled_matmul.py") ) example_tiled_matmul_cot = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_cot_tiled_matmul.py") ) example_tiled_matmul_new = read_file( os.path.join(REPO_TOP_PATH, "prompts/model_new_ex_tiled_matmul.py") ) example_tiled_matmul_desc = "This given architecture is for a model with tiled matrix multiplication: " match cot_example: case "ex_fuse_gelu": base = example_fuse_gelu cot = example_fuse_gelu_cot kernel = example_fuse_gelu_new desc = example_fuse_gelu_desc case "ex_mnist2": base = example_mnist2 cot = example_mnist2_cot kernel = example_mnist2_new desc = exmaple_mnist2_desc case "ex_tiled_matmul": base = example_tiled_matmul cot = example_tiled_matmul_cot kernel = example_tiled_matmul_new desc = example_tiled_matmul_desc case _: raise ValueError(f"Invalid CoT example: {cot_example} not found in CoT examples") # construct example with # NOTE: we only do one example with CoT for now # 1. ref_src problem -> 2. Instruction -> 3. CoT -> 4. Solution prompt += f""" Here is an example architecture:\n\n ``` {base} ```\n {PROBLEM_INSTRUCTION_COT} \n {cot} \n ``` {kernel} ```\n\n """ # show task to solve prompt += f""" Task:\n\n Here is an example architecture:\n\n ``` {ref_arch_src} ```\n """ prompt += PROBLEM_INSTRUCTION_COT return prompt def prompt_fix_correctness(ref_arch_src, custom_cuda, metadat [truncated — 580 more characters] ``` ### adding_red.py ```python import streamlit as st import requests import random import torch import torch.nn as nn import torchviz from streamlit_ace import st_ace from io import BytesIO import tempfile import os # Ensure Graphviz is installed and available in PATH os.environ["PATH"] += os.pathsep + "/usr/local/bin:/usr/bin:/bin:/opt/homebrew/bin:/usr/sbin:/sbin" # Streamlit UI Setup - Remove Top Blank Space st.set_page_config(page_title="Python to CUDA Kernel Optimization", layout="wide") # **🔹 Two-Column Layout (Editors at the Top)** col1, col2 = st.columns(2) # **🔹 Left Side: Python Code Input (Editable, Starts at the Top)** with col1: st.markdown("✏️ **Enter Python Code**", unsafe_allow_html=True) python_code = st_ace( language="python", # Python Syntax Highlighting theme="monokai", # Dark theme placeholder="Write or paste your Python code here...", height=400, # **Slightly reduced height** key="python_code_editor" ) # **Tensor Input Dimension Section** st.markdown("📏 **Specify Tensor Input Dimensions**", unsafe_allow_html=True) tensor_dim = st.text_input("Enter tensor dimensions (e.g., 1, 3, 224, 224)", "1, 3, 224, 224") # **Hardware Selection & Optimization Duration (Always Visible)** st.markdown("⚙️ **Optimization Settings**", unsafe_allow_html=True) hardware = st.selectbox("💻 Select Hardware", ["NVIDIA H100", "NVIDIA A100"]) optimization_time = st.slider("⏳ Optimization Duration (mins)", 1, 15, 5) # **🔹 Right Side: CUDA Kernel Code Output (Read-Only, Starts at the Top)** with col2: st.markdown("⚡ **Generated CUDA Kernel Code**", unsafe_allow_html=True) # **Placeholder CUDA Kernel Code** cuda_kernel_pseudo = """ __global__ void kernel_function(float *input, float *output, int N) { int idx = threadIdx.x + blockIdx.x * blockDim.x; if (idx < N) { output[idx] = input[idx] * input[idx]; // Example computation } } """ # ✅ Syntax-Highlighted **Read-Only** CUDA Output st_ace( value=cuda_kernel_pseudo, language="c_cpp", # Use C++ mode since CUDA mode isn't available theme="monokai", readonly=True, # **Ensures it's not editable** height=400, # **Slightly reduced height** key="cuda_code_output" ) # **Live Updating Performance Graph (Always Visible)** st.markdown("📊 **Live Performance Graph**", unsafe_allow_html=True) # Generate random initial data random_values = [round(random.uniform(1.0, 2.5), 2) for _ in range(7)] constant_values = [2.0] * 7 # Constant red line at value 2.0 # Chart.js Script chart_html = f""" """ st.components.v1.html(chart_html, height=250) # **🔹 Visualization Panel** st.markdown("🖼 **Model Computational Graph**", unsafe_allow_html=True) if st.button("🖥 Generate Computational Graph"): try: input_dims = tuple(map(int, tensor_dim.split(','))) dummy_input = torch.randn(input_dims) local_scope = {} exec(compile(python_code, "", "exec"), local_scope) model = next((v for v in local_scope.values() if isinstance(v, nn.Module)), None) if model is None: raise ValueError("No valid PyTorch model found in the input code.") dot = torchviz.make_dot(model(dummy_input), params=dict(model.named_parameters())) with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile: dot.render(tmpfile.name, format='png') image_path = tmpfile.name + ".png" os.rename(tmpfile.name, image_path) st.image(image_path, caption="Computational Graph") os.remove(image_path) # Clean up the f [truncated — 581 more characters] ``` ### ui.py ```python import streamlit as st import requests import random import types import itertools from streamlit_ace import st_ace import time import torch from torch import nn # Streamlit UI Setup - Remove Top Blank Space st.set_page_config(page_title="Python to CUDA Kernel Optimization", layout="wide") # 🔹 Two-Column Layout (Editors at the Top) col1, col2 = st.columns(2) # 🔹 Left Side: Python Code Input (Editable, Starts at the Top) with col1: st.markdown("✏️ **Enter Python Code**", unsafe_allow_html=True) python_code = st_ace( language="python", theme="monokai", placeholder="Write or paste your Python code here...", height=400, key="python_code_editor" ) # **Tensor Input Dimension Section** # st.markdown("📏 **Specify Tensor Input Dimensions**", unsafe_allow_html=True) # tensor_dim = st.text_input("Enter tensor dimensions (e.g., 1, 3, 224, 224)", "1, 3, 224, 224") # 🔹 Hardware Selection & Optimization Duration st.markdown("⚙️ **Optimization Settings**", unsafe_allow_html=True) hardware = st.selectbox("💻 Select Hardware", ["NVIDIA H100"]) # num_trials = st.slider("🔁 Number of Trials", 10, 200, 100) st.markdown("⚙️ **Translate Python to CUDA Kernel**", unsafe_allow_html=True) button_clicked = st.button("🚀 Generate kernel") torch_output = st.empty() cuda_output = st.empty() # 🔹 Right Side: CUDA Kernel Code Output (Read-Only, Starts at the Top) with col2: status_text = st.empty() # 🔹 Performance Metrics (Below CUDA Code) st.markdown("📊 **Performance Metrics**", unsafe_allow_html=True) torch_time_text = st.empty() kernel_time_text = st.empty() # Placeholder for CUDA Kernel Code st.markdown("⚡ **Generated CUDA Kernel Code**", unsafe_allow_html=True) cuda_code_container = st.empty() # from torchview import draw_graph # ensure torchview is installed # st.markdown("🖼 **Model Computational Graph**", unsafe_allow_html=True) # if st.button("🖥 Generate Computational Graph"): # # Show a loading message for exactly 3 seconds. # with st.spinner("Loading computational graph..."): # time.sleep(2) # # 1. Parse tensor dimensions and create dummy input. # try: # input_dims = tuple(map(int, tensor_dim.split(','))) # dummy_input = torch.rand(input_dims) # except Exception as e: # st.error(f"Error parsing tensor dimensions: {e}") # st.stop() # # 2. Execute the user code in a fresh module namespace. # user_module = types.ModuleType("user_module") # try: # exec(compile(python_code, "", "exec"), user_module.__dict__) # except Exception as e: # st.error(f"Error executing user code: {e}") # st.stop() # # 3. Extract the model (the first nn.Module instance found) # model = next((v for k, v in user_module.__dict__.items() if isinstance(v, nn.Module)), None) # if model is None: # st.error("No valid nn.Module instance found in the provided code.") # st.write("Module keys:", list(user_module.__dict__.keys())) # st.stop() # try: # # 4. Use TorchView to generate the computational graph. # # draw_graph returns a path to an image file by default, # # so we can then display that image using st.image. # graph_path = draw_graph(model, input_size=input_dims, roll=True) # st.write(graph_path.visual_graph) # except Exception as e: # st.error(f"TorchView error: {e}") # st.exception(e) if button_clicked: if python_code.strip(): status_text.info("📡 Initializing task on server...") try: # 🔹 Step 1: Send Python Code to Initialize Task payload = {"python_source": python_code, "num_trials": 1} response = requests.post( "https://d4fa-209-20-157-139.ngrok-free.app/initialize_task", json=payload ) if response.status_code == 200: data = response.json() task_id = data.get("task_id") torch_time = data.get("torch_time", "N/A") torch_time_text.markdown(f"🔥 **Torch Execution Time:** `{torch_time:.3f} ms`") torch_output.markdown("🔥 **Torch output sample:** \n```\n" + data["output"] + "\n```") if task_id: status_text.info("📡 Generating and compiling CUDA kernel...") # 🔹 Step 2: Request CUDA Kernel Initialization kernel_payload = {"task_id": task_id, "num_trials": 1} kernel_response = requests.post( "https://d4fa-209-20-157-139.ngrok-free.app/get_kernel", json=kernel_payload ) if kernel_response.status_code == 200: kernel_data = kernel_response.json() kernel_time = float(kernel_data.get("kernel_time", "0")) kernel_time_text.markdown(f"⚡ **CUDA Execution Time:** `{kernel_time:.3f} ms`") if "kernel_code" in kernel_data: # ✅ Render CUDA Kernel Code dynamically cuda_code_container.markdown("```cpp\n" + kernel_data["kernel_code"] + "\n```") cuda_output.markdown("**CUDA output sample:** \n```\n" + kernel_data["output"] + "\n```") algorithms = ["Torch", "CUDA"] runtimes = [torch_time, kernel_time] # Replace with actual runtimes # Convert data to JavaScript-friendly format labels_js = str(algorithms) data_js = str(runtimes) # Chart.js HTML [truncated — 2392 more characters] ``` ### sample/kernel.cpp ```c++ torch::Tensor diag_matmul_cuda(torch::Tensor diag, torch::Tensor mat); ``` ### prompts/model_ex_fuse_gelu.py ```python import torch import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() def forward(self, x): return F.gelu(x, approximate='tanh') def get_inputs(): # randomly generate input tensors based on the model architecture x = torch.randn(1024, 1024).cuda() return [x] def get_init_inputs(): # randomly generate tensors required for initialization based on the model architecture return [] ``` ### prompts/model_ex_tiled_matmul.py ```python import torch import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() def forward(self, a, b): return a@b def get_inputs(): # randomly generate input tensors based on the model architecture a = torch.randn(1024, 1024).cuda() b = torch.randn(1024, 1024).cuda() return [a, b] def get_init_inputs(): # randomly generate tensors required for initialization based on the model architecture return [] ``` [9 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]