# How DeepGEMM JIT Compilation Works: Environment Variables and Implementation

> Explore how DeepGEMM JIT compilation creates GPU kernels from Python at runtime. Discover environment variables like DG_JIT_ to control this process.

- Repository: [DeepSeek/DeepGEMM](https://github.com/deepseek-ai/DeepGEMM)
- Tags: internals
- Published: 2026-04-19

---

**DeepGEMM generates GPU kernels just-in-time (JIT) by compiling Python-defined Triton or TileLang kernels to CUDA binaries at runtime, caching results in `$HOME/.deep_gemm` and exposing control via environment variables prefixed with `DG_JIT_`.**

DeepGEMM is an open-source library for high-performance FP8 matrix multiplication. Understanding its **JIT compilation** pipeline is essential for debugging performance issues and optimizing kernel launch latency. The system transparently converts Python kernel definitions into executable GPU code while providing extensive environment controls for cache management, compiler selection, and diagnostic output.

## The DeepGEMM JIT Compilation Pipeline

The JIT process transforms high-level Python kernels into optimized CUDA binaries through several distinct stages.

### Kernel Definition and Registration

Kernels originate in Python using framework-specific decorators. Legacy components use **`@triton.jit`** (defined in [`deep_gemm/legacy/m_grouped_gemm.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/legacy/m_grouped_gemm.py)【^11-L12】), while newer implementations use **`@tilelang.jit`** (found in `third_party/tilelang_ops/`). These decorators register the kernel's abstract syntax tree and metadata for later code generation.

### Compiler Selection: NVCC vs NVRTC

When a kernel is first invoked, DeepGEMM's C++ infrastructure (in [`csrc/jit/compiler.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/jit/compiler.hpp)) determines whether to use **NVCC** (static compilation) or **NVRTC** (runtime compilation). This decision hinges on the **`DG_JIT_USE_NVRTC`** environment variable (default `0`). The compiler singleton initializes this selection during construction【^54-L60】.

### Filesystem Caching Mechanism

To avoid redundant compilation, DeepGEMM maintains a persistent cache:

- **Default location**: `$HOME/.deep_gemm`
- **Override**: Set **`DG_JIT_CACHE_DIR`** to a custom path【^48-L52】
- **Cache keys**: Generated from kernel signatures, compilation flags, and source code hashes

On subsequent calls with identical parameters, the system loads the pre-compiled `.cubin` file directly, bypassing the compiler entirely.

## Environment Variables Controlling DeepGEMM JIT Compilation

DeepGEMM exposes fine-grained control over the JIT process through numerous environment variables.

### Compilation Backend and Paths

| Variable | Default | Description |
|----------|---------|-------------|
| **`DG_JIT_USE_NVRTC`** | `0` | Use NVRTC instead of NVCC for faster compilation |
| **`DG_JIT_NVCC_COMPILER`** | system NVCC | Override path to NVCC executable【^95-L99】 |
| **`DG_JIT_CACHE_DIR`** | `$HOME/.deep_gemm` | Filesystem cache location【^48-L52】 |
| **`DG_JIT_CPP_STANDARD`** | `20` | C++ standard version (`-std=c++20`)【^55-L58】 |

### Debugging and Diagnostic Output

Enable verbose compilation and intermediate dumps:

- **`DG_JIT_DEBUG`** – Print compilation commands and debug info
- **`DG_JIT_PRINT_COMPILER_COMMAND`** – Echo exact compiler invocations
- **`DG_JIT_PRINT_LOAD_TIME`** – Log kernel loading latency
- **`DG_JIT_PTXAS_VERBOSE`** – Enable verbose PTX assembler output
- **`DG_JIT_PTXAS_CHECK`** – Validate local memory usage via regex inspection【^44-L47】

### Intermediate Representation Dumps

Dump compiler artifacts for inspection:

- **`DG_JIT_DUMP_PTX`** – Save PTX assembly to cache directory
- **`DG_JIT_DUMP_SASS`** – Save SASS (binary microcode) to cache directory
- **`DG_JIT_DUMP_ASM`** – Save assembly output

### Runtime API Selection

DeepGEMM defaults to the CUDA Driver API for kernel loading. To use the CUDA Runtime API (requires CUDA ≥ 12.8):

- Set **`DG_JIT_USE_RUNTIME_API=1`** – This sets a compile-time macro processed in [`setup.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/setup.py)【^25-L31】

## Practical Code Examples

### Example 1: Triton JIT with Custom Cache and Debugging

```python
import os
import torch
import deep_gemm
from deep_gemm.legacy.m_grouped_gemm import m_grouped_bf16_gemm_nt_contiguous_tl

# Configure JIT environment before any kernel invocation

os.environ["DG_JIT_CACHE_DIR"] = "/tmp/deep_gemm_cache"
os.environ["DG_JIT_DEBUG"] = "1"
os.environ["DG_JIT_USE_NVRTC"] = "1"  # Use NVRTC for faster compilation

# Prepare dummy BF16 inputs

A = torch.randn(1024, 256, dtype=torch.bfloat16, device="cuda")
B = torch.randn(1, 256, 1024, dtype=torch.bfloat16, device="cuda")
D = torch.empty(1024, 1024, dtype=torch.bfloat16, device="cuda")
m_idx = torch.arange(1024, dtype=torch.int32, device="cuda")

# First call triggers JIT compilation and caching

m_grouped_bf16_gemm_nt_contiguous_tl(A, B, D, m_idx)

# Subsequent calls load from cache instantly

m_grouped_bf16_gemm_nt_contiguous_tl(A, B, D, m_idx)

```

### Example 2: TileLang JIT with PTX Dump

```python
import os
import torch
import deep_gemm
from third_party.tilelang_ops.swiglu_apply_weight_to_fp8 import apply_weight_to_fp8

# Enable PTX dumping and command printing

os.environ["DG_JIT_DUMP_PTX"] = "1"
os.environ["DG_JIT_PRINT_COMPILER_COMMAND"] = "1"

# Prepare FP8 dummy tensors

weight = torch.randn(128, 64, dtype=torch.float16, device="cuda")
output = torch.empty(128, 64, dtype=torch.float8_e4m3fn, device="cuda")

# Compilation triggers PTX output to DG_JIT_CACHE_DIR

apply_weight_to_fp8(weight, output)

```

## Summary

- **DeepGEMM JIT compilation** converts `@triton.jit` and `@tilelang.jit` Python kernels into CUDA binaries at runtime through a C++ compiler infrastructure in [`csrc/jit/compiler.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/jit/compiler.hpp).
- **Compiler selection** defaults to NVCC but switches to NVRTC when `DG_JIT_USE_NVRTC=1` is set, affecting compilation speed and compatibility.
- **Persistent caching** stores compiled kernels in `$HOME/.deep_gemm` (override with `DG_JIT_CACHE_DIR`) using content-addressable hashes to eliminate redundant compilation.
- **Debugging capabilities** include dumping PTX/SASS (`DG_JIT_DUMP_PTX`, `DG_JIT_DUMP_SASS`), verbose output (`DG_JIT_DEBUG`), and local memory validation (`DG_JIT_PTXAS_CHECK`).
- **Runtime API support** can be enabled via `DG_JIT_USE_RUNTIME_API=1` for CUDA 12.8+ environments, switching from Driver API to Runtime API execution.

## Frequently Asked Questions

### What triggers JIT compilation in DeepGEMM?

JIT compilation triggers the first time a specific kernel configuration (matrix shape, data type, or grouped layout) is encountered. The Python wrappers in [`deep_gemm/legacy/m_grouped_gemm.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/legacy/m_grouped_gemm.py) invoke the underlying C++ compiler infrastructure, which checks the filesystem cache in `DG_JIT_CACHE_DIR` before deciding whether to compile or load an existing binary.

### How do I switch between NVCC and NVRTC for JIT compilation?

Set the environment variable `DG_JIT_USE_NVRTC` to `1` to use NVRTC (NVIDIA Runtime Compilation), or leave it as `0` (default) to use NVCC. NVRTC typically offers faster compilation times but may have different optimization characteristics compared to the static NVCC compiler. This logic is handled in [`csrc/jit/compiler.hpp`](https://github.com/deepseek-ai/DeepGEMM/blob/main/csrc/jit/compiler.hpp) during the compiler singleton initialization.

### Where does DeepGEMM store compiled kernels and how can I clear the cache?

By default, DeepGEMM stores compiled kernels in `$HOME/.deep_gemm`. You can override this location by setting `DG_JIT_CACHE_DIR` to a custom path. To clear the cache, simply delete the directory contents; the next kernel invocation will trigger fresh compilation. The cache uses cryptographic hashes of source code and compilation flags as filenames to ensure correctness.

### What is the difference between the Driver API and Runtime API execution modes?

By default, DeepGEMM uses the **CUDA Driver API** to load and launch kernels. Setting `DG_JIT_USE_RUNTIME_API=1` switches to the **CUDA Runtime API**, which requires CUDA version 12.8 or higher. The Runtime API mode is configured during package installation via [`setup.py`](https://github.com/deepseek-ai/DeepGEMM/blob/main/setup.py) and affects how compiled cubins are loaded into the process.