Add Kernel Optimization Template to PromptManager

[kaiming-cheng]

This PR adds a new Jinja2 template for bottleneck-guided kernel optimization.

Each optimization round targets exactly one root cause with one recommended fix

kernel_optimization.j2

• ROOFLINE ANALYSIS section with SOL%, efficiency, headroom
• BOTTLENECK ANALYSIS expects singular bottleneck with root_cause and recommended_fix
• PERFORMANCE TARGET targets 10% improvement over current best or Eager baseline
• TARGET GPU includes all fields from gpu_specs_database

prompt_manager.py

• Added kernel_optimization template
• New render_kernel_optimization_prompt() with explicit params: category, summary, reasoning, root_cause, recommended_fix

Example Usage

1. Generate bottleneck analysis from NCU data

  from kernel_perf_agent.kernel_opt.diagnose_prompt.judger_prompt import (                                                                                                             
      build_bottleneck_prompt,                                                                                                                                                         
      parse_bottleneck_response,                                                                                                                                                       
  )                                                                                                                                                                                    
  from kernel_perf_agent.kernel_opt.roofline.ncu_roofline import RooflineAnalyzer                                                                                                      
                                                                                                                                                                                       
  # Analyze roofline                                                                                                                                                                   
  analyzer = RooflineAnalyzer()                                                                                                                                                        
  roofline = analyzer.analyze(ncu_metrics)                                                                                                                                             
                                                                                                                                                                                       
  # Build prompt for LLM                                                                                                                                                               
  prompt = build_bottleneck_prompt(                                                                                                                                                    
      kernel_code=kernel_src,                                                                                                                                                          
      ncu_metrics=ncu_metrics,                                                                                                                                                         
      roofline=roofline,                                                                                                                                                               
      gpu_specs=gpu_specs,                                                                                                                                                             
      num_bottlenecks=2,                                                                                                                                                               
      num_causes=2,                                                                                                                                                                    
      num_fixes=1,                                                                                                                                                                     
  )                                                                                                                                                                                    
                                                                                                                                                                                       
  # Call LLM and parse response                                                                                                                                                        
  llm_response = call_llm(prompt)                                                                                                                                                      
  bottlenecks = parse_bottleneck_response(llm_response)                                                                                                                                
  # Returns: [BottleneckResult, BottleneckResult]                                                                                                                                      

2. Render optimization prompt (one cause, one fix)

  from triton_kernel_agent.prompt_manager import PromptManager                                                                                                                         
                                                                                                                                                                      
  prompt_manager = PromptManager()                                                                                                                                                     
                                                                                                                                                                                       
  # Pick first bottleneck, first cause/fix pair                                                                                                                                        
  b = bottlenecks[0]                                                                                                                                                                   
                                                                                                                                                                                       
  prompt = prompt_manager.render_kernel_optimization_prompt(                                                                                                                           
      problem_description="Fused attention kernel",                                                                                                                                    
      kernel_code=current_kernel,                                                                                                                                                      
      gpu_specs=gpu_specs,                                                                                                                                                             
      roofline={                                                                                                                                                                       
          "bottleneck": roofline.bottleneck,                                                                                                                                           
          "compute_sol_pct": roofline.compute_sol_pct,                                                                                                                                 
          "memory_sol_pct": roofline.memory_sol_pct,                                                                                                                                   
          "efficiency_pct": roofline.efficiency_pct,                                                                                                                                   
          "headroom_pct": roofline.headroom_pct,                                                                                                                                       
          "at_roofline": roofline.at_roofline,                                                                                                                                         
          "uses_tensor_cores": roofline.uses_tensor_cores,                                                                                                                             
          "warnings": roofline.warnings,                                                                                                                                               
      },                                                                                                                                                                               
      category=b.category,                                                                                                                                                             
      summary=b.summary,                                                                                                                                                               
      reasoning=b.reasoning,                                                                                                                                                           
      root_cause=b.root_causes[0],                                                                                                                                                     
      recommended_fix=b.recommended_fixes[0],                                                                                                                                          
      pytorch_baseline_ms=1.234,                                                                                                                                                       
      current_best_ms=0.987,  # from previous iteration                                                                                                                                
  )                                                                                                                                                                                    
                                                                                                                                                                                       
  # Call LLM to generate optimized kernel                                                                                                                                              
  optimized_kernel = call_llm(prompt)        

[Jack-Khuu]

Not in this PR so you don't need to change it, but once everything lands we should consider creating arg dataclasses, so that we aren't passing around 10+ arg to functions

[Jack-Khuu]

Is this just to prevent folks from editing it by accident? If so consider GPU_SPECS_DATABASE = MappingProxyType(GPU_SPECS_DATABASE)

This makes it so that you can read from the DB, but not write. Removes the need for this file

[Jack-Khuu]

If we aren't customizing to_dict, have callers use asdict