trainkernel 0.1.0__py3-none-any.whl

trainkernel/runtime.py

@@ -0,0 +1,142 @@
+"""CUDA extension compilation and a checked, static execution boundary."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+import hashlib
+from typing import Any
+
+import torch
+from torch.utils._pytree import tree_flatten
+
+from .codegen import GeneratedKernel, generate
+from .ir import Program
+
+
+@dataclass
+class ExecutionBuffers:
+    """Caller-owned storage. Reuse only when previous work on it is ordered."""
+
+    outputs: tuple[torch.Tensor, ...]
+    workspace: torch.Tensor
+
+
+class CompiledProgram:
+    """One specialized kernel. Calls allocate independent buffers for stream safety.
+
+    The result has no autograd registration: compile a functional training step
+    containing explicit differentiation to compute gradients and parameter updates.
+    """
+
+    def __init__(self, program: Program, generated: GeneratedKernel, module: Any, device: torch.device, inplace_outputs: dict[int, int] | None = None):
+        self.program = program
+        self.generated = generated
+        self.source = generated.source
+        self.module = module
+        self.device = device
+        self.inplace_outputs = dict(inplace_outputs or {})
+
+    def resource_usage(self) -> dict[str, int]:
+        """Compiler resource attributes; these do not require GPU counters."""
+        with torch.cuda.device(self.device):
+            values = self.module.resource_usage()
+        return dict(zip(("registers_per_thread", "local_bytes_per_thread", "shared_bytes_per_block", "resident_blocks_per_sm"), values))
+
+    def _validate_inputs(self, inputs: tuple[torch.Tensor, ...]) -> None:
+        if len(inputs) != len(self.program.inputs):
+            raise ValueError(f"trainkernel: expected {len(self.program.inputs)} inputs, got {len(inputs)}")
+        for i, (tensor, spec) in enumerate(zip(inputs, self.program.inputs)):
+            if not isinstance(tensor, torch.Tensor):
+                raise TypeError(f"trainkernel: input {i} must be a tensor")
+            if tensor.is_neg() or tensor.is_conj():
+                raise ValueError("trainkernel: resolve lazy negative or conjugate input views before execution")
+            if tensor.device != self.device or tuple(tensor.shape) != spec.shape or tuple(tensor.stride()) != spec.strides or tensor.dtype != getattr(torch, spec.dtype):
+                raise ValueError(f"trainkernel: input {i} does not match compiled device, shape, stride, or dtype")
+        if any(inputs[i].requires_grad for i in self.inplace_outputs.values()):
+            raise ValueError("trainkernel: inplace targets must be detached from autograd")
+
+    def allocate(self, *inputs: torch.Tensor) -> ExecutionBuffers:
+        if self.inplace_outputs:
+            self._validate_inputs(inputs)
+        return ExecutionBuffers(
+            tuple(inputs[self.inplace_outputs[i]] if i in self.inplace_outputs else torch.empty(spec.shape, dtype=getattr(torch, spec.dtype), device=self.device) for i, spec in enumerate(self.program.outputs)),
+            torch.empty(self.generated.workspace_elements, dtype=torch.float32, device=self.device),
+        )
+
+    def run_into(self, *inputs: torch.Tensor, buffers: ExecutionBuffers) -> tuple[torch.Tensor, ...]:
+        self._validate_inputs(inputs)
+        if len(buffers.outputs) != len(self.program.outputs):
+            raise ValueError("trainkernel: incorrect output buffer count")
+        if buffers.workspace.device != self.device or buffers.workspace.dtype != torch.float32 or not buffers.workspace.is_contiguous() or buffers.workspace.numel() < self.generated.workspace_elements:
+            raise ValueError("trainkernel: workspace does not match compiled requirements")
+        # Caller buffers must be disjoint from each other and the inputs. A
+        # conservative storage-level check prevents writes clobbering dependencies.
+        output_storages = []
+        allowed_storages = set()
+        for slot, (out, spec) in enumerate(zip(buffers.outputs, self.program.outputs)):
+            if out.device != self.device or out.dtype != getattr(torch, spec.dtype) or tuple(out.shape) != spec.shape or not out.is_contiguous():
+                raise ValueError("trainkernel: output buffer does not match compiled requirements")
+            if out.numel():
+                output_storages.append(out.untyped_storage().data_ptr())
+            if slot in self.inplace_outputs:
+                target_index = self.inplace_outputs[slot]
+                target = inputs[target_index]
+                if out.data_ptr() != target.data_ptr() or tuple(out.stride()) != tuple(target.stride()):
+                    raise ValueError("trainkernel: inplace output must bind its declared input")
+                if target.numel():
+                    storage = target.untyped_storage().data_ptr()
+                    if any(i != target_index and value.numel() and value.untyped_storage().data_ptr() == storage for i, value in enumerate(inputs)):
+                        raise ValueError("trainkernel: inplace target must not alias another input")
+                    allowed_storages.add(storage)
+        if buffers.workspace.numel():
+            output_storages.append(buffers.workspace.untyped_storage().data_ptr())
+        input_storages = {tensor.untyped_storage().data_ptr() for tensor in inputs if tensor.numel()}
+        if len(set(output_storages)) != len(output_storages) or input_storages.intersection(output_storages) - allowed_storages:
+            raise ValueError("trainkernel: execution buffers must not alias inputs or each other")
+        self.module.launch(list(inputs), list(buffers.outputs), buffers.workspace)
+        if self.inplace_outputs:
+            torch.autograd.graph.increment_version([inputs[i] for i in self.inplace_outputs.values()])
+        return buffers.outputs
+
+    def __call__(self, *inputs: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        return self.run_into(*inputs, buffers=self.allocate(*inputs))
+
+
+def compile_program(program: Program, example_inputs: tuple[Any, ...] | list[Any], *, blocks_per_sm: int = 1, max_blocks: int | None = None, fuse_pointwise: bool = True, fuse_epilogues: bool = True, gemm_k_tile: int = 32, gemm_tile_n: int = 32, shared_skew: int = 0, split_k_gemv: bool = False, threads_per_block: int = 128, inline_input_casts: bool = True, inplace_outputs: dict[int, int] | None = None, verbose: bool = False) -> CompiledProgram:
+    """Compile for the current input device, requiring Linux, CUDA and SM80+."""
+    flat, _ = tree_flatten(example_inputs)
+    if not flat or any(not isinstance(x, torch.Tensor) for x in flat):
+        raise TypeError("trainkernel: nonempty tensor inputs required")
+    if len(flat) != len(program.inputs):
+        raise ValueError("trainkernel: input count differs from the program signature")
+    if not torch.cuda.is_available() or any(x.device.type != "cuda" for x in flat):
+        raise RuntimeError("trainkernel: CUDA SM80+ GPU required; CPU IR inspection is available through capture()")
+    device = flat[0].device
+    if any(x.device != device for x in flat):
+        raise ValueError("trainkernel: all inputs must use the same CUDA device")
+    capability = torch.cuda.get_device_capability(device)
+    if capability < (8, 0):
+        raise RuntimeError("trainkernel: NVIDIA compute capability 8.0 or newer required")
+    generated = generate(program, blocks_per_sm=blocks_per_sm, max_blocks=max_blocks, fuse_pointwise=fuse_pointwise, fuse_epilogues=fuse_epilogues, gemm_k_tile=gemm_k_tile, gemm_tile_n=gemm_tile_n, shared_skew=shared_skew, split_k_gemv=split_k_gemv, threads_per_block=threads_per_block, inline_input_casts=inline_input_casts, inplace_outputs=inplace_outputs)
+    if any(flat[i].requires_grad for i in (inplace_outputs or {}).values()):
+        raise ValueError("trainkernel: inplace targets must be detached from autograd")
+    from torch.utils.cpp_extension import CUDA_HOME, load_inline
+    if CUDA_HOME is None:
+        raise RuntimeError("trainkernel: CUDA toolkit with nvcc is required to compile generated kernels")
+    name = "trainkernel_" + hashlib.sha256((generated.source + str(capability) + torch.__version__).encode()).hexdigest()[:20]
+    # Restrict compilation to the actual target rather than relying on a global
+    # TORCH_CUDA_ARCH_LIST, which may target a different GPU or unsupported SM.
+    arch = f"-gencode=arch=compute_{capability[0]}{capability[1]},code=sm_{capability[0]}{capability[1]}"
+    with torch.cuda.device(device):
+        module = load_inline(
+            name=name,
+            cpp_sources="void launch(std::vector<torch::Tensor>, std::vector<torch::Tensor>, torch::Tensor); std::vector<int64_t> resource_usage();",
+            cuda_sources=generated.source,
+            functions=["launch", "resource_usage"],
+            extra_cuda_cflags=["-O3", "-std=c++17", "--fmad=false", arch],
+            extra_cflags=["-O3"],
+            with_cuda=True,
+            verbose=verbose,
+        )
+    result = CompiledProgram(program, generated, module, device, inplace_outputs)
+    result._validate_inputs(tuple(flat))
+    return result

trainkernel/wmma.py

@@ -0,0 +1,130 @@
+"""Optional 32x64 WMMA GEMM source emitter.
+
+This module deliberately has no dependency on the scheduler.  The callbacks
+are the small boundary needed to retain scheduler-specific addressing and
+epilogue fusion while keeping this experimental tile independently testable.
+"""
+
+from dataclasses import dataclass
+from typing import Callable
+
+
+Read = Callable[[str, str], str]
+Store = Callable[[str, str], str]
+Epilogue = Callable[[str, str], str]
+
+
+@dataclass(frozen=True)
+class WMMAEmission:
+    source: str
+    tile_m: int = 32
+    tile_n: int = 64
+    warps: int = 4
+    shared_a_elements: int = 0
+    shared_b_elements: int = 0
+    shared_c_elements: int = 2048
+
+    # A/B are FP16/BF16 (2 bytes each); C is FP32 (4 bytes).  Keeping these
+    # separate avoids treating the mixed-width allocation as one element type.
+    @property
+    def shared_a_bytes(self) -> int:
+        return self.shared_a_elements * 2
+
+    @property
+    def shared_b_bytes(self) -> int:
+        return self.shared_b_elements * 2
+
+    @property
+    def shared_c_bytes(self) -> int:
+        return self.shared_c_elements * 4
+
+    @property
+    def shared_bytes(self) -> int:
+        return self.shared_a_bytes + self.shared_b_bytes + self.shared_c_bytes
+
+    @property
+    def shared_elements(self) -> int:
+        return self.shared_a_elements + self.shared_b_elements + self.shared_c_elements
+
+
+def emit_wmma(
+    m: int,
+    k: int,
+    n: int,
+    read_a: Read,
+    read_b: Read,
+    store: Store,
+    epilogue: Epilogue | None = None,
+    *,
+    k_tile: int = 32,
+    value_type: str = "__nv_bfloat16",
+    cast: str = "__float2bfloat16_rn",
+) -> WMMAEmission:
+    """Emit one 32x64 low-precision GEMM work tile.
+
+    ``read_a(row, k)`` and ``read_b(k, col)`` return device expressions.
+    ``store(index, value)`` must return a complete CUDA statement.  The
+    optional epilogue receives the logical output index and accumulator
+    expression, and may return a replacement expression.
+    """
+    if min(m, k, n) < 0 or k_tile not in (16, 32):
+        raise ValueError("WMMA dimensions must be non-negative and k_tile must be 16 or 32")
+    ep = epilogue or (lambda _index, value: value)
+    a_tile = 32 * k_tile
+    b_tile = k_tile * 64
+    chunks = (k + k_tile - 1) // k_tile
+    # Each warp owns one 16-row A slab and two 16-column output slabs.
+    lines = [
+        "{",
+        f"  auto* wmma_a = reinterpret_cast<{value_type}*>(shared_a);",
+        f"  auto* wmma_b = reinterpret_cast<{value_type}*>(shared_b);",
+        "  const int wmma_warp = threadIdx.x >> 5;",
+        f"  for (int64_t wtile = blockIdx.x; wtile < {((m + 31) // 32) * ((n + 63) // 64)}LL; wtile += gridDim.x) {{",
+        f"    const int64_t wrow = (wtile / {max(1, (n + 63) // 64)}LL) * 32LL;",
+        f"    const int64_t wcol = (wtile % {max(1, (n + 63) // 64)}LL) * 64LL;",
+        "    nvcuda::wmma::fragment<nvcuda::wmma::accumulator,16,16,16,float> acc0, acc1;",
+        "    nvcuda::wmma::fill_fragment(acc0, 0.0f);",
+        "    nvcuda::wmma::fill_fragment(acc1, 0.0f);",
+        f"    for (int64_t wkk = 0; wkk < {k}LL; wkk += {k_tile}) {{",
+        f"      for (int p = threadIdx.x; p < {max(a_tile, b_tile)}; p += blockDim.x) {{",
+        f"        if (p < {a_tile}) {{ const int64_t ar = wrow + p / {k_tile}, ak = wkk + p % {k_tile}; wmma_a[p] = {cast}((ar < {m}LL && ak < {k}LL) ? {read_a('ar', 'ak')} : 0.0f); }}",
+        f"        if (p < {b_tile}) {{ const int64_t bk = wkk + p / 64, bc = wcol + p % 64; wmma_b[p] = {cast}((bk < {k}LL && bc < {n}LL) ? {read_b('bk', 'bc')} : 0.0f); }}",
+        "      }",
+        "      __syncthreads();",
+        "      if (wmma_warp < 4) {",
+        f"        const int arow = (wmma_warp >> 1) * 16;",
+        "        nvcuda::wmma::fragment<nvcuda::wmma::matrix_a,16,16,16," + value_type + ",nvcuda::wmma::row_major> af;",
+        "        nvcuda::wmma::fragment<nvcuda::wmma::matrix_b,16,16,16," + value_type + ",nvcuda::wmma::row_major> bf0, bf1;",
+        f"        nvcuda::wmma::load_matrix_sync(af, wmma_a + arow * {k_tile}, {k_tile});",
+        "        nvcuda::wmma::load_matrix_sync(bf0, wmma_b + (wmma_warp & 1) * 16, 64);",
+        "        nvcuda::wmma::load_matrix_sync(bf1, wmma_b + (wmma_warp & 1) * 16 + 32, 64);",
+        "        nvcuda::wmma::mma_sync(acc0, af, bf0, acc0);",
+        "        nvcuda::wmma::mma_sync(acc1, af, bf1, acc1);",
+    ]
+    if k_tile == 32:
+        lines += [
+            "        nvcuda::wmma::load_matrix_sync(af, wmma_a + arow * 32 + 16, 32);",
+            "        nvcuda::wmma::load_matrix_sync(bf0, wmma_b + 16 * 64 + (wmma_warp & 1) * 16, 64);",
+            "        nvcuda::wmma::load_matrix_sync(bf1, wmma_b + 16 * 64 + (wmma_warp & 1) * 16 + 32, 64);",
+            "        nvcuda::wmma::mma_sync(acc0, af, bf0, acc0);",
+            "        nvcuda::wmma::mma_sync(acc1, af, bf1, acc1);",
+        ]
+    lines += [
+        "      }",
+        "      __syncthreads();",
+        "    }",
+        "    if (wmma_warp < 4) {",
+        "      nvcuda::wmma::store_matrix_sync(shared_c + (wmma_warp >> 1) * 16 * 64 + (wmma_warp & 1) * 16, acc0, 64, nvcuda::wmma::mem_row_major);",
+        "      nvcuda::wmma::store_matrix_sync(shared_c + (wmma_warp >> 1) * 16 * 64 + (wmma_warp & 1) * 16 + 32, acc1, 64, nvcuda::wmma::mem_row_major);",
+        "    }",
+        "    __syncthreads();",
+        "    for (int p = threadIdx.x; p < 2048; p += blockDim.x) {",
+        "      const int64_t rr = wrow + p / 64, cc = wcol + p % 64;",
+        f"      if (rr < {m}LL && cc < {n}LL) {{ const int64_t oi = rr * {n}LL + cc; {store('oi', ep('oi', 'shared_c[p]'))} }}",
+        "    }",
+        "    __syncthreads();",
+        "  }",
+        "}",
+    ]
+    return WMMAEmission("\n".join(lines), shared_a_elements=a_tile, shared_b_elements=b_tile)
+

trainkernel-0.1.0.dist-info/METADATA

@@ -0,0 +1,169 @@
+Metadata-Version: 2.4
+Name: trainkernel
+Version: 0.1.0
+Summary: Compile a complete static PyTorch training step into one CUDA kernel
+Author: Arya Somu
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/arpjw/trainkernel
+Project-URL: Documentation, https://github.com/arpjw/trainkernel#readme
+Project-URL: Issues, https://github.com/arpjw/trainkernel/issues
+Project-URL: Source, https://github.com/arpjw/trainkernel
+Keywords: pytorch,cuda,compiler,training,gpu
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: torch>=2.8
+Requires-Dist: ninja>=1.11
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: twine>=5; extra == "dev"
+Dynamic: license-file
+
+# TrainKernel
+
+TrainKernel compiles a supported static PyTorch training step into one
+cooperative CUDA megakernel. The generated kernel can include the forward pass,
+loss, gradients, and SGD update, so a captured step has no per-operator Python
+launch overhead.
+
+The current backend schedules operators in topological order, uses grid barriers for dependencies, inlines single-consumer pointwise operations, and reuses temporary storage after its last reader. BF16 and FP16 matrix multiplications use WMMA tensor cores. FP32 uses a tiled SIMT implementation for accuracy.
+
+This is a scoped v0.1 research release for static functional training on NVIDIA
+SM80 GPUs. It is not a general PyTorch compiler and does not claim universal
+speedups.
+
+| A100 workload | TrainKernel | Fastest measured Inductor | TrainKernel / Inductor |
+| --- | ---: | ---: | ---: |
+| tiny online MLP | 15.28 μs | 22.85 μs | **0.669, 1.49x faster** |
+| online MLP | 20.32 μs | 24.24 μs | **0.838, 1.19x faster** |
+| small MLP | 66.04 μs | 22.92 μs | 2.894, slower |
+| medium MLP | 255.32 μs | 42.19 μs | 6.052, slower |
+
+These are complete forward, backward, and update steps captured as CUDA Graphs.
+The [held-out report](https://github.com/arpjw/trainkernel/blob/main/results/a100-optimization/README.md) publishes all seeds,
+raw timings, correctness checks, baseline failures, and the two regressions.
+
+An independent release run adds a one-hot least-squares classification head.
+Across three seeds, TrainKernel's in-place path measured 13.07 us versus
+18.63 us for the fastest measured Inductor configuration at the median, or
+1.43x faster. The [release evidence](https://github.com/arpjw/trainkernel/blob/main/results/release-validation/README.md)
+retains all raw samples and correctness failures. Its cold-cache setup took
+53.62 seconds and required an estimated 7.82 million updates to amortize, so
+this is a steady-state small-workload result rather than a claim about
+short-lived jobs.
+
+## Install
+
+Linux, Python 3.10+, PyTorch 2.8+, an NVIDIA SM80+ GPU, a CUDA toolkit with `nvcc`, and a C++ compiler are required for generated kernel execution. CPU capture and IR checks run without CUDA.
+
+```bash
+python -m pip install trainkernel
+```
+
+For development:
+
+```bash
+python -m venv --system-site-packages .venv
+source .venv/bin/activate
+python -m pip install -e '.[dev]'
+python -m pytest -q
+python -m examples.online_training
+```
+
+The backend targets SM80+, including A100 and H100. Release evidence is A100-only; H100 has not been tested and Hopper TMA/WGMMA paths are not implemented.
+
+## Compile a training step
+
+```python
+import torch
+import trainkernel
+
+def train_step(weights, x, target):
+    def loss(w):
+        return ((torch.relu(x @ w) - target) ** 2).mean()
+    gradients, value = torch.func.grad_and_value(loss)(weights)
+    return value, gradients, weights - 1e-3 * gradients
+
+compiled = trainkernel.compile_training(train_step, (weights, x, target))
+loss, gradients, weights = compiled(weights, x, target)
+```
+
+`compile_training` returns a typed `CompiledTrainingStep`. It verifies that
+gradients match state structure and derives the update mapping, including nested
+state trees. The supplied state must be detached from autograd because the
+training function computes its own gradients. `compiled.capture(...)` creates a
+CUDA Graph after warmup, restores the supplied state after warmup and capture,
+and returns a reusable `CapturedProgram`; replay outputs borrow captured buffers
+with fixed shapes and addresses. See the [100-update example](https://github.com/arpjw/trainkernel/blob/main/examples/online_training.py).
+
+Parameters may be nested dictionaries or tuples. Inputs and operations must stay on one device; resolve lazy negative or conjugate input views before compiling. Updated parameters are returned explicitly. The compiler proves direct last-read state commits where safe and stages all other updates transactionally; input mutation and arbitrary aliasing remain unsupported. Different shapes, strides, devices, or dtypes require another compilation.
+
+Inspect the compiler without a GPU:
+
+```python
+program, output_tree = trainkernel.capture(train_step, cpu_examples)
+from trainkernel.codegen import generate
+cuda_source = generate(program).source
+```
+
+## Benchmark
+
+```bash
+bash scripts/run_train_loop_suite.sh results/a100-optimization/heldout
+```
+
+For a clean release-candidate check, including the useful classification-head
+workload and package build:
+
+```bash
+bash scripts/run_release_validation.sh results/release-validation
+```
+
+The [training-loop harness](https://github.com/arpjw/trainkernel/blob/main/benchmarks/train_loop.py) uses FP32 master weights, BF16 compute, FP32 loss, and FP32 SGD updates. It checks every loss, gradient, and updated parameter for 100 steps and separately checks a captured 20-step replay. Timed graphs include state feedback over four resident batches. They compare pure and in-place megakernels, captured eager PyTorch, and default and eager-rounding Inductor variants. CUDA event timings exclude compilation, capture, reset, and warmup, and retain raw samples. Numerical failures remain visible for all baselines. The [methodology](https://github.com/arpjw/trainkernel/blob/main/docs/a100-methodology.md) explains the held-out seeds and fixed tolerance policy; the [result](https://github.com/arpjw/trainkernel/blob/main/results/a100-optimization/README.md) includes both wins and regressions. The older [native-BF16 fixed-step harness](https://github.com/arpjw/trainkernel/blob/main/benchmarks/train_mlp.py) remains available for the historical report.
+
+## Supported subset
+
+- FP32, BF16, and FP16 strided tensors with static shapes.
+- Matrix multiplication and `addmm`, broadcast arithmetic, ReLU and its backward, scalar powers, sums and means, dtype casts, fills, and supported storage views.
+- Pure tensor functions, including explicit differentiation through `torch.func.grad_and_value`.
+
+Unsupported operations produce an error. There is no silent eager fallback. Attention, dropout/RNG, convolutions, dynamic shapes, arbitrary input mutation, distributed collectives, and asynchronous task-queue scheduling are not implemented.
+
+Temporary values use FP32 workspace storage with explicit rounding at low-precision boundaries. Fusion reduces materialization but does not keep an entire model on-chip. The tested larger GEMMs run faster in Inductor; the report keeps those regressions visible.
+
+## Design and evidence
+
+See [compiler design](https://github.com/arpjw/trainkernel/blob/main/docs/design.md), [A100 optimization evidence](https://github.com/arpjw/trainkernel/blob/main/results/a100-optimization/README.md), [independent release validation](https://github.com/arpjw/trainkernel/blob/main/results/release-validation/README.md), and [GPU access record](https://github.com/arpjw/trainkernel/blob/main/docs/infrastructure.md). CPU tests validate IR semantics, views, reductions, diagnostics, and arena lifetimes. The independent release suite passed 116 tests that compile and execute generated CUDA, compare training outputs across precisions, and exercise CUDA Graph replay.
+
+## Related work and roadmap
+
+Most adjacent public megakernel systems currently focus on inference:
+[HazyResearch Megakernels](https://github.com/HazyResearch/Megakernels) targets
+low-latency Llama inference, [MegaBake](https://github.com/megabake/megabake)
+captures PyTorch inference workloads for SM90, [AutoMegaKernel](https://github.com/RightNow-AI/AutoMegaKernel)
+generates Llama-family inference kernels, and [Mirage Persistent Kernel](https://arxiv.org/abs/2512.22219)
+targets multi-GPU model inference. TrainKernel instead explores compiling the
+stateful forward, loss, backward, and optimizer update of a small training step
+on A100 into one kernel. It does not claim priority for megakernel compilation
+in general.
+
+The next engineering priorities are a measured shape-aware dispatcher that can
+delegate GEMM-heavy steps to Inductor, an experimental `torch.compile` backend
+after AOTAutograd, AdamW, and additional real online-learning workloads. These
+are roadmap items, not v0.1 capabilities.
+
+## License
+
+Apache License 2.0. See [LICENSE](https://github.com/arpjw/trainkernel/blob/main/LICENSE) and [NOTICE](https://github.com/arpjw/trainkernel/blob/main/NOTICE).

trainkernel-0.1.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+trainkernel/__init__.py,sha256=q3Rn3ug1CeCDHN3nxOiYwFBxSj38VssMqdH-1XGYCwk,407
+trainkernel/codegen.py,sha256=mzIFB9agOqBaZvN4d6yyECPSs7mGnV5ghcQLC-iEGfs,52881
+trainkernel/frontend.py,sha256=zeC06BmpclBf5eAL7kisVbffkUGGejmmVAUzXp5zhEU,15438
+trainkernel/gemv.py,sha256=T466_07Jj3sH9RGj_XKD8ZYr_AEni1QJAh6UEaign3c,2049
+trainkernel/ir.py,sha256=ghIch3x1XXni4C5xo9fKyhBcw3s-FKblKVGXaxMcLg8,2525
+trainkernel/passes.py,sha256=pw3VtNeI3CSxUAMFpgOL1cZy3llVFha4M8yHvDYi9Ks,3382
+trainkernel/py.typed,sha256=1F5juyAFRQIScl2WbVd2VzUvkxbW-SYhkQnmeXncq_4,60
+trainkernel/reference.py,sha256=8qZTyq6PYx_VaCn6wBZz_V3AEy1rpnd0saAUzyG8a3M,2468
+trainkernel/runtime.py,sha256=CDhAW4ar5BGGsO9v9P96IqE1JjCuNX_gVzBc0O3s6qo,9053
+trainkernel/wmma.py,sha256=TVjFn4OMUBEc_1C_7BcRkXvvjPmmpLFUsmnMs-i_PRY,5796
+trainkernel-0.1.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+trainkernel-0.1.0.dist-info/licenses/NOTICE,sha256=K2yZZQoC774cNzyge_5x1Dr2wE132L3Lf_Yegg8cVxg,174
+trainkernel-0.1.0.dist-info/METADATA,sha256=FEuQCQaLUWGDFwHXwppMYAnOzvjM1T33SZ6P6jOf5VI,9638
+trainkernel-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+trainkernel-0.1.0.dist-info/top_level.txt,sha256=-EBMWciAioOLO1Pbf9uC9AUqPPTJCbPsfjY_zRE9nEs,12
+trainkernel-0.1.0.dist-info/RECORD,,

trainkernel-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+