Add triton to kernel bench

README.md

@@ -86,6 +86,9 @@ It is easier to get started with a single problem. This will fetch the problem,
 
 python3 scripts/generate_and_eval_single_sample.py dataset_src="huggingface" level=2 problem_id=40
 
+# to generate and evaluate triton kernels, use the following
+python3 scripts/generate_and_eval_single_sample.py dataset_src="huggingface" level=2 problem_id=40
+
 # dataset_src could be "local" or "huggingface"
 # add .verbose_logging for more visbility
 ```
@@ -111,8 +114,13 @@ python3 scripts/benchmark_eval_analysis.py run_name=test_hf_level_1 level=1 hard
 If you are using a different hardware, you can generate the baseline time with `scripts/generate_baseline_time.py` script.
 We provide some reference baseline times a variety of NVIDIA GPUs across generations in `results/timing`, but we recommend you to generate your own baseline time for more accurate results (cluster power, software version, all affects timing result). See `results/timing/README.md` for more details.
 
+### Framework Support
+KernelBench is created to target PyTorch -> PyTorch + custom inline-compiled CUDA kernels.
+Now we are also extending support for Triton: PyTorch -> PyTorch + custom jit-compiled Triton kernels; please set `framework="triton"` (default "cuda") in your command.
+
+Future stable support coming for [ThunderKittens](https://github.com/HazyResearch/ThunderKittens).
+
 ## 🛣️ Upcoming Roadmap
-- [ ] Triton Variant (Ongoing)
 - [ ] Easy to use CoLab Notebook Example
 - [ ] Push button flow on Modal / Cloud Provider 
 - [ ] Integrate with more frameworks, such as [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)

scripts/generate_and_eval_single_sample.py

@@ -8,7 +8,7 @@
 
 from src.dataset import construct_kernelbench_dataset
 from src.eval import eval_kernel_against_ref
-from src.prompt_constructor import prompt_generate_custom_cuda_from_prompt_template
+from src.prompt_constructor import prompt_generate_custom_kernel_from_prompt_template
 from src.utils import extract_first_code, query_server, set_gpu_arch, read_file, create_inference_server_from_presets
 
 """
@@ -43,7 +43,12 @@ def __init__(self):
 
         # Inference config
         self.server_type = "deepseek"
-        self.model_name = "deepseek-coder"
+
+        # for local
+        self.server_port = 30000
+        self.server_address = "localhost"
+
+        self.model_name = "deepseek-chat"
         self.max_tokens = 4096
         self.temperature = 0.0
 
@@ -56,6 +61,9 @@ def __init__(self):
         self.log_generated_kernel = False
         self.log_eval_result = False
 
+        # todo: make this an enum
+        self.framework = "cuda" # cuda or triton
+
     def verbose_logging(self):
         self.log = True
         self.log_prompt = True
@@ -93,6 +101,7 @@ def main(config: EvalConfig):
     print(f"Start Generation + Evaluation for Level {config.level} Problem {config.problem_id}")
 
     assert config.problem_id <= num_problems, f"Problem ID {config.problem_id} out of range for Level {config.level}"
+    assert config.framework in ["cuda", "triton"], "Framework must be either cuda or triton"
 
 
     # 1. Fetch Problem
@@ -123,31 +132,36 @@ def main(config: EvalConfig):
                                                         temperature=config.temperature,
                                                         max_tokens=config.max_tokens,
                                                         verbose=config.verbose,
-                                                        time_generation=True)
+                                                        time_generation=True,
+                                                        server_port=config.server_port,
+                                                        server_address=config.server_address)
 
 
 
-    custom_cuda_prompt = prompt_generate_custom_cuda_from_prompt_template(ref_arch_src)
+    custom_kernel_prompt = prompt_generate_custom_kernel_from_prompt_template(ref_arch_src, framework=config.framework)
     if config.log_prompt:
         with open(os.path.join(config.logdir, f"prompt_level_{config.level}_problem_{config.problem_id}.txt"), "w") as f:
-            f.write(custom_cuda_prompt)
+            f.write(custom_kernel_prompt)
+
+    # DEBUG
+    # return 
 
     # Query server with constructed prompt
-    custom_cuda = inference_server(custom_cuda_prompt)
-    custom_cuda = extract_first_code(custom_cuda, ["python", "cpp"])
+    custom_kernel = inference_server(custom_kernel_prompt)
+    custom_kernel = extract_first_code(custom_kernel, ["python", "cpp"])
     # check LLM is able to generate custom CUDA code
-    assert custom_cuda is not None, "Custom CUDA code generation failed"
+    assert custom_kernel is not None, "Custom CUDA code generation failed"
 
     # this should be optional
     if config.log:
         with open(os.path.join(config.logdir, f"generated_kernel_level_{config.level}_problem_{config.problem_id}.py"), "w") as f:
-            f.write(custom_cuda)
+            f.write(custom_kernel)
 
     # 3. Evaluate Kernel
     # NOTE: no need to wrap around process here as only a single sample
     # see batch eval for examples of process isolation
     kernel_exec_result = eval_kernel_against_ref(
-        ref_arch_src, custom_cuda, verbose=config.verbose, measure_performance=True, num_correct_trials=5, num_perf_trials=100
+        ref_arch_src, custom_kernel, verbose=config.verbose, measure_performance=True, num_correct_trials=5, num_perf_trials=100
     )
 
     print(f"Evaluation result for level {config.level} problem {config.problem_id}:\n{kernel_exec_result}")

scripts/generate_and_eval_single_sample_modal.py

@@ -8,9 +8,8 @@
 from datasets import load_dataset
 
 #from src.dataset import construct_kernelbench_dataset
-from src.eval import eval_kernel_against_ref
-from src.prompt_constructor import prompt_generate_custom_cuda_from_prompt_template
-from src.utils import extract_first_code, query_server, set_gpu_arch, read_file, create_inference_server_from_presets
+from src.prompt_constructor import prompt_generate_custom_kernel_from_prompt_template
+from src.utils import extract_first_code, read_file, create_inference_server_from_presets
 
 app = modal.App("eval_single_sample")
 
@@ -62,6 +61,9 @@ def __init__(self):
         self.log_prompt = False
         self.log_generated_kernel = False
         self.log_eval_result = False
+
+        # todo: make this an enum
+        self.framework = "cuda" # cuda or triton
 
     def verbose_logging(self):
         self.log = True
@@ -106,15 +108,15 @@ def __repr__(self):
 class EvalFunc:
 
     @modal.method()
-    def eval_single_sample_modal(self, ref_arch_src, custom_cuda, verbose, gpu_arch):
+    def eval_single_sample_modal(self, ref_arch_src, custom_kernel, verbose, gpu_arch):
         # 3. Evaluate Kernel
         # NOTE: no need to wrap around process here as only a single sample
         # see batch eval for examples of process isolation
         from src.eval import eval_kernel_against_ref
         from src.utils import set_gpu_arch
         set_gpu_arch(gpu_arch)
         return eval_kernel_against_ref(
-            ref_arch_src, custom_cuda, verbose=verbose, measure_performance=True, num_correct_trials=5, num_perf_trials=100
+            ref_arch_src, custom_kernel, verbose=verbose, measure_performance=True, num_correct_trials=5, num_perf_trials=100
         )
 
 @pydra.main(base=EvalConfig)
@@ -140,6 +142,7 @@ def main(config: EvalConfig):
     print(f"Start Generation + Evaluation for Level {config.level} Problem {config.problem_id}")
 
     assert config.problem_id <= num_problems, f"Problem ID {config.problem_id} out of range for Level {config.level}"
+    assert config.framework in ["cuda", "triton"], "Framework must be either cuda or triton"
 
 
     # 1. Fetch Problem
@@ -172,26 +175,24 @@ def main(config: EvalConfig):
                                                         verbose=config.verbose, 
                                                         time_generation=True)
 
-
-
-    custom_cuda_prompt = prompt_generate_custom_cuda_from_prompt_template(ref_arch_src)
+    custom_kernel_prompt = prompt_generate_custom_kernel_from_prompt_template(ref_arch_src, framework=config.framework)
     if config.log_prompt:
         with open(os.path.join(config.logdir, f"prompt_level_{config.level}_problem_{config.problem_id}.txt"), "w") as f:
-            f.write(custom_cuda_prompt)
+            f.write(custom_kernel_prompt)
 
     # Query server with constructed prompt
-    custom_cuda = inference_server(custom_cuda_prompt)
-    custom_cuda = extract_first_code(custom_cuda, ["python", "cpp"])
+    custom_kernel = inference_server(custom_kernel_prompt)
+    custom_kernel = extract_first_code(custom_kernel, ["python", "cpp"])
     # check LLM is able to generate custom CUDA code
-    assert custom_cuda is not None, "Custom CUDA code generation failed"
+    assert custom_kernel is not None, "Custom CUDA code generation failed"
 
     # this should be optional
     if config.log:
         with open(os.path.join(config.logdir, f"generated_kernel_level_{config.level}_problem_{config.problem_id}.py"), "w") as f:
-            f.write(custom_cuda)
+            f.write(custom_kernel)
 
     with app.run():
-        kernel_exec_result = EvalFunc.with_options(gpu=config.gpu)().eval_single_sample_modal.remote(ref_arch_src, custom_cuda, config.verbose, gpu_arch_mapping[config.gpu])
+        kernel_exec_result = EvalFunc.with_options(gpu=config.gpu)().eval_single_sample_modal.remote(ref_arch_src, custom_kernel, config.verbose, gpu_arch_mapping[config.gpu])
 
         print(f"Evaluation result for level {config.level} problem {config.problem_id}:\n{kernel_exec_result}")
 

scripts/generate_samples.py

@@ -10,8 +10,8 @@
 
 from src.dataset import construct_kernelbench_dataset
 from src.eval import eval_kernel_against_ref
-from src.prompt_constructor import prompt_generate_custom_cuda_from_prompt_template
-from src.utils import extract_first_code, set_gpu_arch, read_file, create_inference_server_from_presets, maybe_multithread
+from src.prompt_constructor import prompt_generate_custom_kernel_from_prompt_template
+from src.utils import extract_first_code, extract_last_code, set_gpu_arch, read_file, create_inference_server_from_presets, maybe_multithread
 
 """
 Batch Generate Samples for Particular Level
@@ -46,8 +46,16 @@ def __init__(self):
         # Inference config
         self.server_type = "deepseek"
         self.model_name = "deepseek-coder"
+        self.server_port = 30000
+        self.server_address = "localhost"
+
+        # ex. non-reasoning model 
+        self.is_reasoning_model = False
         self.max_tokens = 4096
         self.temperature = 0.0
+
+        # ex. reasoning model
+        # self.is_reasoning_model = True
 
         # Logging
         # Top Directory to Store Runs
@@ -62,6 +70,9 @@ def __init__(self):
 
         self.log_prompt = False
 
+        # todo: make this an enum
+        self.framework = "cuda" # cuda or triton
+
     def greedy(self):
         # For greedy decoding, epsecially baseline eval
         self.greedy_sample = True
@@ -93,29 +104,31 @@ def generate_sample_single(work: WorkArgs, config: GenerationConfig, dataset, in
     # Extract problem number from problem name (e.g. "1" from "1_Square_matrix_multiplication_.py")
     problem_number = int(problem_name.split("_")[0])
     assert problem_number == work.problem_id, f"Problem number in filename ({problem_number}) does not match config problem_id ({config.problem_id})"
+    assert config.framework in ["cuda", "triton"], "Framework must be either cuda or triton"
 
 
 
     # Construct Prompt   
-    custom_cuda_prompt = prompt_generate_custom_cuda_from_prompt_template(ref_arch_src)
+    custom_kernel_prompt = prompt_generate_custom_kernel_from_prompt_template(ref_arch_src, framework=config.framework)
     if config.log_prompt:
         prompt_path = os.path.join(run_dir, f"level_{config.level}_problem_{work.problem_id}_sample_{work.sample_id}_prompt.txt")
         with open(prompt_path, "w") as f:
-            f.write(custom_cuda_prompt)
+            f.write(custom_kernel_prompt)
 
     # Query server with constructed prompt
-    custom_cuda = inference_server(custom_cuda_prompt)
-    custom_cuda = extract_first_code(custom_cuda, ["python", "cpp"])
+    custom_kernel = inference_server(custom_kernel_prompt)
+    # custom_kernel = extract_first_code(custom_kernel, ["python", "cpp"])
+    custom_kernel = extract_last_code(custom_kernel, ["python", "cpp"])
     # check LLM is able to generate custom CUDA code
-    assert custom_cuda is not None, "Custom CUDA code generation failed"
+    assert custom_kernel is not None, "Custom CUDA code generation failed"
 
     if config.verbose:
         print(f"Generated sample {work.sample_id} for problem {problem_number}: {problem_name}")
 
     # Store to local file
     kernel_path = os.path.join(run_dir, f"level_{config.level}_problem_{work.problem_id}_sample_{work.sample_id}_kernel.py")
     with open(kernel_path, "w") as f:
-        f.write(custom_cuda)
+        f.write(custom_kernel)
 
     return True
 
@@ -188,7 +201,10 @@ def main(config: GenerationConfig):
                                                         model_name=config.model_name,
                                                         temperature=config.temperature,
                                                         max_tokens=config.max_tokens,
-                                                        verbose=config.verbose)
+                                                        verbose=config.verbose,
+                                                        is_reasoning_model=config.is_reasoning_model,
+                                                        server_port=config.server_port,
+                                                        server_address=config.server_address)
 
     # Launch workers
     generation_results = maybe_multithread(generate_sample_launcher, 

scripts/greedy_analysis.py

@@ -0,0 +1,46 @@
+import json, os
+
+"""
+Analyze the greedy eval results for a run of a particular level
+"""
+from src.dataset import construct_kernelbench_dataset
+
+run_name = "test_hf_level_1"  # Replace this with your run name
+level = 1  # change if needed
+
+dataset = construct_kernelbench_dataset(level)
+
+
+# load json
+eval_file_path = f"runs/{run_name}/eval_results.json"
+assert os.path.exists(eval_file_path), f"Eval file does not exist at {eval_file_path}"
+
+
+with open(eval_file_path, "r") as f:
+    eval_results = json.load(f)
+
+
+# Initialize counters
+total_count = len(dataset)
+total_eval = len(eval_results)
+compiled_count = 0
+correct_count = 0
+
+# Count results
+for entry in eval_results.values():
+    if entry["compiled"] == True:
+        compiled_count += 1
+    if entry["correctness"] == True:
+        correct_count += 1
+
+# Print results
+print("-" * 128)
+print(f"Eval Summary for {run_name}")
+print("-" * 128)
+print(f"Total test cases with Eval Results: {total_eval} out of {total_count}")
+print(f"Successfully compiled: {compiled_count}")
+print(f"Functionally correct: {correct_count}")
+
+print(f"\nSuccess rates:")
+print(f"Compilation rate: {compiled_count/total_count*100:.1f}%")
+print(f"Correctness rate: {correct_count/total_count*100:.1f}%")