Generate the complete kernel file content for extraction.
(
op_type: str,
rank: int,
pct_total: float,
model_shape: Dict[str, int],
model_name: str,
gpu_time_ms: float,
starter_code: str,
backend: str = "triton",
)
| 288 | |
| 289 | |
| 290 | def generate_kernel_file( |
| 291 | op_type: str, |
| 292 | rank: int, |
| 293 | pct_total: float, |
| 294 | model_shape: Dict[str, int], |
| 295 | model_name: str, |
| 296 | gpu_time_ms: float, |
| 297 | starter_code: str, |
| 298 | backend: str = "triton", |
| 299 | ) -> str: |
| 300 | """Generate the complete kernel file content for extraction.""" |
| 301 | |
| 302 | half_shape = scale_shape(model_shape, 0.5) |
| 303 | double_shape = scale_shape(model_shape, 2.0) |
| 304 | |
| 305 | shape_display = shape_to_display(model_shape) |
| 306 | half_display = shape_to_display(half_shape) |
| 307 | double_display = shape_to_display(double_shape) |
| 308 | |
| 309 | tolerances = TOLERANCES_MAP.get(op_type, { |
| 310 | "float16": {"atol": 1e-2, "rtol": 1e-2}, |
| 311 | "bfloat16": {"atol": 2e-2, "rtol": 2e-2}, |
| 312 | "float32": {"atol": 1e-4, "rtol": 1e-4}, |
| 313 | }) |
| 314 | |
| 315 | flops_fn_body = FLOPS_FN_SRC.get(op_type, 'return 0') |
| 316 | bytes_fn_body = BYTES_FN_SRC.get(op_type, 'return 0') |
| 317 | |
| 318 | # Extract the kernel code body (imports + jit functions + kernel_fn) |
| 319 | kernel_body = extract_kernel_body(starter_code) |
| 320 | |
| 321 | # Build the file |
| 322 | lines = [] |
| 323 | |
| 324 | # Header docstring |
| 325 | lines.append('"""') |
| 326 | lines.append(f"AutoKernel -- Extracted kernel from model profiling.") |
| 327 | lines.append(f"Op type: {op_type}") |
| 328 | lines.append(f"Rank: {rank} ({pct_total}% of GPU time)") |
| 329 | lines.append(f"Model shape: {shape_display}") |
| 330 | lines.append(f"") |
| 331 | lines.append(f"This kernel was extracted from profiling {model_name}.") |
| 332 | lines.append(f"The agent optimizes this to maximize throughput at the model-specific shapes.") |
| 333 | lines.append('"""') |
| 334 | lines.append("") |
| 335 | |
| 336 | # KERNEL_TYPE and BACKEND |
| 337 | lines.append(f'KERNEL_TYPE = "{op_type}"') |
| 338 | if backend == "cuda": |
| 339 | lines.append(f'BACKEND = "cuda"') |
| 340 | lines.append("") |
| 341 | |
| 342 | # Model-specific shapes |
| 343 | lines.append("# Model-specific shapes (the shapes that matter for THIS model)") |
| 344 | lines.append(f"MODEL_SHAPES = {repr(model_shape)}") |
| 345 | lines.append("") |
| 346 | |
| 347 | # Benchmark config |
no test coverage detected