Function _export_triton_kernel

export_hf.py:587–667

Export a Triton kernel to HF Kernels format. Triton kernels are already Python, so the export is simpler: package the Triton code as a Python module. No CUDA compilation needed -- the Triton JIT compiler handles everything at runtime.

(
    source: str,
    name: str,
    output_dir: str,
    repo_id: str,
)

# ---------------------------------------------------------------------------

def _export_triton_kernel(
    source: str,
    name: str,
    output_dir: str,
    repo_id: str,
) -> None:
    """
    Export a Triton kernel to HF Kernels format.

    Triton kernels are already Python, so the export is simpler: package
    the Triton code as a Python module. No CUDA compilation needed -- the
    Triton JIT compiler handles everything at runtime.
    """
    # Create directory structure
    project_dir = os.path.join(output_dir, name)
    module_dir = os.path.join(project_dir, name)

    os.makedirs(module_dir, exist_ok=True)

    # 1. Write the Triton kernel as kernel.py inside the module
    triton_code = extract_triton_code(source)
    kernel_py_path = os.path.join(module_dir, "kernel.py")
    with open(kernel_py_path, "w", encoding="utf-8") as f:
        f.write(triton_code.strip())
        f.write("\n")
    print(f"  Created {os.path.relpath(kernel_py_path, output_dir)}")

    # 2. Write __init__.py
    init_py_path = os.path.join(module_dir, "__init__.py")
    # For Triton, functions are the Python entry points (kernel_fn)
    functions = [{"name": "kernel_fn"}]
    init_py_content = generate_init_py(name, functions, repo_id, backend="triton")
    with open(init_py_path, "w", encoding="utf-8") as f:
        f.write(init_py_content)
    print(f"  Created {os.path.relpath(init_py_path, output_dir)}")

    # 3. Write a minimal pyproject.toml for the Triton package
    pyproject_path = os.path.join(project_dir, "pyproject.toml")
    pyproject_content = textwrap.dedent(f"""\
        [project]
        name = "{name}"
        version = "0.1.0"
        description = "Optimized Triton GPU kernel exported from AutoKernel"
        requires-python = ">=3.10"
        dependencies = [
            "torch>=2.4.0",
            "triton>=3.3.0",
        ]
    """)
    with open(pyproject_path, "w", encoding="utf-8") as f:
        f.write(pyproject_content)
    print(f"  Created {os.path.relpath(pyproject_path, output_dir)}")

    # 4. Write a README for the Hub repo
    readme_path = os.path.join(project_dir, "README.md")
    readme_content = textwrap.dedent(f"""\
        # {name}