Template export autograd codegen

Author: drisspg

examples/export_autograd_triton/rms_norm.py

@@ -11,13 +11,13 @@
     export_autograd_triton,
     load_exported_module,
 )
-from transformer_nuggets.utils.benchmark import benchmark_cuda_function_in_microseconds
+from transformer_nuggets.utils.benchmark import benchmark_cuda_function_in_microseconds_triton
 
 
 HIDDEN_SIZE = 4096
 SAMPLE_TOKENS = 128
 MAX_TOKENS = 2048
-BENCHMARK_TOKENS = 512
+TOKEN_COUNTS = (1, 17, SAMPLE_TOKENS, 512)
 DTYPE = torch.bfloat16
 EPS = 1e-5
 
@@ -99,7 +99,7 @@ def _export_rms_norm(
 
 
 def _validate(compiled_fn: Callable, weight: torch.Tensor) -> None:
-    for tokens in (1, 17, SAMPLE_TOKENS, BENCHMARK_TOKENS):
+    for tokens in TOKEN_COUNTS:
         runtime_x = torch.randn(
             tokens,
             HIDDEN_SIZE,
@@ -122,26 +122,26 @@ def _validate(compiled_fn: Callable, weight: torch.Tensor) -> None:
 
 
 def _benchmark_memory_bandwidth(label: str, compiled_fn: Callable, weight: torch.Tensor) -> None:
-    benchmark_x = torch.randn(
-        BENCHMARK_TOKENS,
-        HIDDEN_SIZE,
-        device="cuda",
-        dtype=DTYPE,
-    )
+    print("forward bandwidth (assumes 2 x reads + 1 weight read + 1 output write):")
+    for tokens in TOKEN_COUNTS:
+        benchmark_x = torch.randn(
+            tokens,
+            HIDDEN_SIZE,
+            device="cuda",
+            dtype=DTYPE,
+        )
 
-    def run_forward():
-        with torch.no_grad():
-            return compiled_fn(benchmark_x, weight, eps=EPS)
+        def run_forward():
+            with torch.no_grad():
+                return compiled_fn(benchmark_x, weight, eps=EPS)
 
-    time_us = benchmark_cuda_function_in_microseconds(run_forward, NUM_ITERS=100)
-    bandwidth_gb_s = _forward_memory_bytes(benchmark_x) / (time_us * 1e-6) / 1e9
-    print(
-        f"{label} forward: {time_us:.2f} us, ~{bandwidth_gb_s:.1f} GB/s effective memory bandwidth"
-    )
+        time_us = benchmark_cuda_function_in_microseconds_triton(run_forward)
+        bandwidth_gb_s = _forward_memory_bytes(benchmark_x) / (time_us * 1e-6) / 1e9
+        print(f"  tokens={tokens}: {time_us:.2f} us, ~{bandwidth_gb_s:.1f} GB/s")
 
 
 def _forward_memory_bytes(x: torch.Tensor) -> int:
-    return 3 * x.numel() * x.element_size()
+    return 4 * x.numel() * x.element_size()
 
 
 def _print_artifact_summary(output_path: Path) -> None:

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
     "pandas",
     "seaborn",
     "typer",
+    "jinja2",
 ]
 
 [project.optional-dependencies]

transformer_nuggets/export_autograd_triton/codegen.py

@@ -7,227 +7,33 @@
 import re
 from textwrap import indent
 
+from jinja2 import Environment, FileSystemLoader, StrictUndefined
+
 from transformer_nuggets.export_autograd_triton.clean_triton import clean_triton_module
 from transformer_nuggets.export_autograd_triton.specs import CapturedSpecialization
 
 
+_TEMPLATE_ENV = Environment(
+    loader=FileSystemLoader(Path(__file__).with_name("templates")),
+    undefined=StrictUndefined,
+    keep_trailing_newline=True,
+)
+
+
 def generate_autograd_source(
     fn: Callable[..., object],
     exported_name: str,
     specializations: list[CapturedSpecialization],
 ) -> str:
     signature = inspect.signature(fn)
-    wrapper = _generate_public_wrapper(exported_name, signature)
-    specs_literal = pprint.pformat(
-        [_spec_to_metadata(index, spec) for index, spec in enumerate(specializations)],
-        width=100,
-        sort_dicts=False,
+    return _TEMPLATE_ENV.get_template("autograd_module.py.j2").render(
+        specs_literal=pprint.pformat(
+            [_spec_to_metadata(index, spec) for index, spec in enumerate(specializations)],
+            width=100,
+            sort_dicts=False,
+        ),
+        wrapper=_generate_public_wrapper(exported_name, signature).rstrip(),
     )
-    return f"""from __future__ import annotations
-
-import importlib.util
-from pathlib import Path
-
-import torch
-
-_ARTIFACTS_DIR = Path(__file__).with_name(f"{{Path(__file__).stem}}_artifacts")
-
-_SPECS = {specs_literal}
-
-
-def _load_runner(module_filename):
-    if module_filename is None:
-        return None
-    module_path = _ARTIFACTS_DIR / module_filename
-    spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
-    module = importlib.util.module_from_spec(spec)
-    if spec.loader is None:
-        raise RuntimeError(f"Could not load compiled source module {{module_path}}")
-    spec.loader.exec_module(module)
-    return module.call
-
-
-_FORWARD_RUNNERS = [_load_runner(spec["forward_module"]) for spec in _SPECS]
-_BACKWARD_RUNNERS = [_load_runner(spec["backward_module"]) for spec in _SPECS]
-
-
-class _CompiledAutogradFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, spec_id, *runtime_tensors):
-        spec = _SPECS[spec_id]
-        outputs = _FORWARD_RUNNERS[spec_id](list(runtime_tensors))
-        if not isinstance(outputs, tuple):
-            outputs = tuple(outputs)
-        user_outputs = outputs[: spec["num_user_outputs"]]
-        residuals = outputs[spec["num_user_outputs"] :]
-        named_residuals = tuple(zip(spec["forward_residual_names"], residuals, strict=True))
-        ctx.spec_id = spec_id
-        ctx.runtime_tensor_count = len(runtime_tensors)
-        ctx.saved_tensor_residual_names = tuple(
-            name for name, residual in named_residuals if isinstance(residual, torch.Tensor)
-        )
-        ctx.saved_non_tensor_residuals = {{
-            name: residual for name, residual in named_residuals if not isinstance(residual, torch.Tensor)
-        }}
-        ctx.save_for_backward(
-            *(residual for _, residual in named_residuals if isinstance(residual, torch.Tensor))
-        )
-        non_differentiable_outputs = tuple(
-            output
-            for output, is_differentiable in zip(
-                user_outputs,
-                spec["differentiable_output_mask"],
-                strict=True,
-            )
-            if isinstance(output, torch.Tensor) and not is_differentiable
-        )
-        if non_differentiable_outputs:
-            ctx.mark_non_differentiable(*non_differentiable_outputs)
-        if spec["output_kind"] == "single":
-            return user_outputs[0]
-        return tuple(user_outputs)
-
-    @staticmethod
-    def backward(ctx, *grad_outputs):
-        spec = _SPECS[ctx.spec_id]
-        grad_outputs = tuple(
-            grad_output.contiguous() if isinstance(grad_output, torch.Tensor) else grad_output
-            for index, grad_output in enumerate(grad_outputs)
-            if spec["differentiable_output_mask"][index]
-        )
-        backward_runner = _BACKWARD_RUNNERS[ctx.spec_id]
-        if backward_runner is None:
-            raise RuntimeError("Selected export_autograd_triton specialization has no backward graph")
-        saved_residuals = dict(ctx.saved_non_tensor_residuals)
-        saved_residuals.update(zip(ctx.saved_tensor_residual_names, ctx.saved_tensors, strict=True))
-        backward_saved_inputs = tuple(
-            saved_residuals[name] for name in spec["backward_saved_input_names"]
-        )
-        grads = backward_runner(list(backward_saved_inputs + grad_outputs))
-        if not isinstance(grads, tuple):
-            grads = tuple(grads)
-        if len(grads) < ctx.runtime_tensor_count:
-            grads = grads + (None,) * (ctx.runtime_tensor_count - len(grads))
-        if len(grads) > ctx.runtime_tensor_count:
-            grads = grads[: ctx.runtime_tensor_count]
-        return (None, *grads)
-
-
-def _run_with_bound_args(bound_args):
-    spec_id = _select_spec(bound_args)
-    spec = _SPECS[spec_id]
-    runtime_tensors = tuple(bound_args[name] for name in spec["runtime_tensor_names"])
-    if spec["needs_autograd"]:
-        result = _CompiledAutogradFunction.apply(spec_id, *runtime_tensors)
-        if spec["output_kind"] == "list":
-            return list(result)
-        return result
-    return _run_forward_only(spec_id, runtime_tensors)
-
-
-def _run_forward_only(spec_id, runtime_tensors):
-    spec = _SPECS[spec_id]
-    outputs = _FORWARD_RUNNERS[spec_id](list(runtime_tensors))
-    if not isinstance(outputs, tuple):
-        outputs = tuple(outputs)
-    user_outputs = outputs[: spec["num_user_outputs"]]
-    if spec["output_kind"] == "single":
-        return user_outputs[0]
-    if spec["output_kind"] == "list":
-        return list(user_outputs)
-    return tuple(user_outputs)
-
-
-def _select_spec(bound_args):
-    matches = []
-    failures = []
-    for spec_id, spec in enumerate(_SPECS):
-        reasons = _mismatch_reasons(spec, bound_args)
-        if reasons:
-            failures.append((spec, reasons))
-        else:
-            matches.append(spec_id)
-    if len(matches) == 1:
-        return matches[0]
-    if len(matches) > 1:
-        names = ", ".join(_SPECS[spec_id]["name"] for spec_id in matches)
-        raise RuntimeError(f"Ambiguous export_autograd_triton specializations matched: {{names}}")
-    raise RuntimeError(_format_no_match(bound_args, failures))
-
-
-def _mismatch_reasons(spec, bound_args):
-    reasons = []
-    symbols = {{}}
-    for name, expected in spec["static_args"].items():
-        actual = bound_args.get(name)
-        if actual != expected:
-            reasons.append(f"static {{name}}={{actual!r}} != {{expected!r}}")
-    for guard in spec["tensor_guards"]:
-        tensor = bound_args.get(guard["name"])
-        reason = _tensor_mismatch_reason(tensor, guard, symbols)
-        if reason is not None:
-            reasons.append(reason)
-    return reasons
-
-
-def _tensor_mismatch_reason(tensor, guard, symbols):
-    name = guard["name"]
-    if not isinstance(tensor, torch.Tensor):
-        return f"{{name}} is {{type(tensor).__name__}}, expected Tensor"
-    expected_dtype = getattr(torch, guard["dtype"])
-    if tensor.dtype is not expected_dtype:
-        return f"{{name}} dtype {{tensor.dtype}} != torch.{{guard['dtype']}}"
-    if tensor.device.type != guard["device_type"]:
-        return f"{{name}} device type {{tensor.device.type}} != {{guard['device_type']}}"
-    if tensor.device.index != guard["device_index"]:
-        return f"{{name}} device index {{tensor.device.index}} != {{guard['device_index']}}"
-    if tensor.dim() != guard["rank"]:
-        return f"{{name}} rank {{tensor.dim()}} != {{guard['rank']}}"
-    shape_reason = _shape_mismatch_reason(name, tensor, guard, symbols)
-    if shape_reason is not None:
-        return shape_reason
-    if tuple(tensor.stride()) != tuple(guard["stride"]):
-        return f"{{name}} stride {{tuple(tensor.stride())}} != {{tuple(guard['stride'])}}"
-    return None
-
-
-def _shape_mismatch_reason(name, tensor, guard, symbols):
-    for dim, expected in enumerate(guard["shape"]):
-        actual = int(tensor.shape[dim])
-        if isinstance(expected, dict):
-            symbol = expected["symbol"]
-            if expected["min"] is not None and actual < expected["min"]:
-                return f"{{name}} shape dim {{dim}}={{actual}} is less than {{symbol}} min {{expected['min']}}"
-            if expected["max"] is not None and actual > expected["max"]:
-                return f"{{name}} shape dim {{dim}}={{actual}} is greater than {{symbol}} max {{expected['max']}}"
-            if symbol in symbols and symbols[symbol] != actual:
-                return f"{{name}} shape dim {{dim}}={{actual}} does not equal {{symbol}}={{symbols[symbol]}}"
-            symbols[symbol] = actual
-        elif actual != expected:
-            return f"{{name}} shape dim {{dim}}={{actual}} != {{expected}}"
-    return None
-
-
-def _format_no_match(bound_args, failures):
-    lines = ["No export_autograd_triton specialization matched the runtime inputs.", "Received:"]
-    for name, value in bound_args.items():
-        if isinstance(value, torch.Tensor):
-            lines.append(
-                f"  {{name}}: dtype={{value.dtype}}, device={{value.device}}, "
-                f"shape={{tuple(value.shape)}}, stride={{tuple(value.stride())}}"
-            )
-        else:
-            lines.append(f"  {{name}}: {{value!r}}")
-    lines.append("Expected specializations:")
-    for spec, reasons in failures:
-        lines.append(f"  - {{spec['name']}}:")
-        for reason in reasons:
-            lines.append(f"      {{reason}}")
-    return "\\n".join(lines)
-
-
-{wrapper}
-"""
 
 
 def write_autograd_source(