Greptile Overview

Greptile Summary

This PR adds a new triton_extensions module to enable JAX primitives to use Triton kernels. The implementation provides:

• triton_call_lowering: A utility function that handles MLIR lowering for Triton kernels, including support for autotuned kernels with multiple configurations
• Kernel compilation with caching: Compiled kernels are cached by signature, constants, and compile options to avoid redundant compilation
• JAX-Triton bridge: Uses JAX's internal gpu_triton APIs to create kernel calls and lower them through the FFI

Key implementation details:

• Autotuned kernels compile all configurations upfront for runtime selection via TritonAutotunedKernelCall
• Grid dimensions are normalized to 3D tuples for the GPU backend
• The test demonstrates proper grid calculation using the minimum BLOCK_SIZE from autotuner configs to ensure all elements are processed regardless of which config is selected at runtime

Confidence Score: 4/5

• This PR is safe to merge - it adds a new module without modifying existing code paths, and includes test coverage for the core functionality.
• Score of 4 reflects: clean implementation using JAX's official gpu_triton APIs, test coverage demonstrating the feature works, and proper handling of autotuned kernels. Minor concerns include the test primitive not implementing all BasePrimitive abstract methods (limiting it to single-device JIT), and limited test coverage (single shape/dtype).
• The test primitive in tests/jax/test_triton_custom_calls.py is intentionally limited scope - consider expanding test coverage for production use.

Important Files Changed

File Analysis

Sequence Diagram

sequenceDiagram
    participant User as JAX Primitive
    participant TCL as triton_call_lowering
    participant CT as compile_triton
    participant GPU as gpu_triton (JAX)
    participant FFI as JAX FFI

    User->>TCL: Call with kernel_fn, ctx, arrays, grid
    TCL->>TCL: Get compute capability
    TCL->>TCL: Build signature from avals
    TCL->>TCL: Normalize grid to 3D

    alt Autotuned Kernel
        loop For each config
            TCL->>CT: Compile with config params
            CT->>CT: Check cache
            CT->>CT: Build ASTSource
            CT->>CT: Compile to PTX
            CT->>GPU: Create TritonKernel
            GPU-->>CT: Return kernel
            CT-->>TCL: Return cached/compiled kernel
            TCL->>GPU: Create TritonKernelCall
        end
        TCL->>GPU: Create TritonAutotunedKernelCall
    else Regular Kernel
        TCL->>CT: Compile kernel
        CT-->>TCL: Return kernel
        TCL->>GPU: Create TritonKernelCall
    end

    TCL->>GPU: Serialize to protobuf
    TCL->>FFI: ffi_lowering with compressed proto
    FFI-->>User: Return MLIR result

/te-ci JAX L0

/te-ci JAX L0

/te-ci JAX L0