where is the script for replacing torch graphs with triton kernels

[JINO-ROHIT]

Honestly just trying to read and learn from this repo, I was wondering if someone can direct/explain me the piece where instead of torch graphs being compiled, the triton kernels kick in.

Thanks very much

[173787247]

There isn’t one standalone “replace torch graphs with triton” script. Unsloth mostly monkey-patches Transformers forward methods to call hand-written Triton kernels.

Good reading order:

1. unsloth/kernels/ — the Triton kernels (@triton.jit): RMSNorm (rms_layernorm.py), RoPE (rope_embedding.py), SwiGLU (swiglu.py), CE loss (cross_entropy_loss.py), etc. Start with kernels/__init__.py for exports.

2. unsloth/models/llama.py — canonical example of the patch + fast-forward wiring:

   • class-level patch: LlamaAttention.forward = LlamaAttention_fast_forward, same for decoder/model/causal LM
   • inside those fast forwards you’ll see direct calls like fast_rope_embedding, fast_rms_layernorm, fast_swiglu_inference

3. Load path — FastLanguageModel.from_pretrained() → patch_model_and_tokenizer() applies the above when you load a model.

4. Other architectures — same pattern in unsloth/models/qwen2.py, qwen3.py, mistral.py, etc. (often reusing the Llama fast forwards).

5. torch.compile shows up in a few specialized paths (e.g. flex attention / some encoder models), but the main finetuning speedups are from the explicit Triton kernels + forward patching, not a global torch-graph → triton rewrite.

If you also have unsloth_zoo installed, additional temporary patches may live there; the main repo’s kernels/__init__.py checks for it.

Happy to point at a specific op (attention vs MLP vs loss) if you say which model you’re reading.