源码导读：Triton kernel 文件结构

本章给一张"哪个文件做什么"的源码地图，覆盖 Unsloth 的 kernel 层、patcher 层、loader 层。 目标是让你给 Unsloth 提 PR 时知道改哪个文件，或者照搬一段 Triton 代码到自己项目里。

13.1顶层目录地图

unsloth/
├── unsloth/                            ★ 主代码
│   ├── __init__.py                     入口，含 MLX 检测
│   ├── _gpu_init.py                    GPU 检测 + 早期环境变量
│   ├── _auto_install.py                自动安装 Triton / flash-attn
│   ├── device_type.py                  设备分类
│   ├── chat_templates.py               ★ 所有模型的 chat template + train_on_responses_only
│   ├── tokenizer_utils.py              tokenizer 修复
│   ├── trainer.py                      UnslothTrainer (continue pre-training 用)
│   ├── save.py                         ★ 所有 save_pretrained_* 实现
│   │
│   ├── kernels/                        ★ 所有 Triton kernel
│   │   ├── __init__.py
│   │   ├── cross_entropy_loss.py       chunked CE
│   │   ├── fast_lora.py                ★ fused LoRA 前后向
│   │   ├── flex_attention.py           PyTorch FlexAttention 集成
│   │   ├── fp8.py                      H100 fp8 训练
│   │   ├── geglu.py                    GeGLU (Gemma 用)
│   │   ├── layernorm.py                LayerNorm
│   │   ├── moe/                         MoE 专家 kernel
│   │   ├── rms_layernorm.py            ★ RMSNorm fused
│   │   ├── rope_embedding.py           ★ RoPE fused
│   │   ├── swiglu.py                   SwiGLU fused
│   │   └── utils.py
│   │
│   ├── models/                          ★ 各架构 patch + Fast*Model
│   │   ├── __init__.py
│   │   ├── _utils.py                   通用工具 + gradient ckpt
│   │   ├── loader.py                   ★ FastLanguageModel.from_pretrained 入口
│   │   ├── loader_utils.py
│   │   ├── mapper.py                   HF 模型名 ↔ Unsloth 预量化版映射
│   │   ├── llama.py                    Llama / Llama-3
│   │   ├── llama4.py                   Llama-4 MoE
│   │   ├── mistral.py                  Mistral / Mixtral
│   │   ├── qwen2.py / qwen3.py / qwen3_moe.py
│   │   ├── gemma.py / gemma2.py
│   │   ├── cohere.py                   Command-R
│   │   ├── granite.py / falcon_h1.py / glm4_moe.py / ...
│   │   ├── vision.py                   FastVisionModel
│   │   ├── sentence_transformer.py     SentenceTransformerModel
│   │   ├── dpo.py                      PatchDPOTrainer
│   │   ├── rl.py / rl_replacements.py  PatchFastRL (GRPO 等)
│   │
│   ├── dataprep/                        数据预处理工具
│   ├── optimizers/                      自家 optimizer（如果有）
│   ├── ollama_template_mappers.py       Modelfile 模板生成
│   ├── registry/                        模型注册表
│   └── utils/                            通用工具
│
├── studio/                              Unsloth Studio（Tauri GUI）
├── unsloth_cli/                         CLI 入口
├── unsloth-cli.py                       CLI 主脚本
├── cli.py
├── install.sh / install.ps1            ★ 安装脚本
├── tests/                                测试 + Colab notebook
│   ├── notebooks/                       50+ 个官方 Colab
│   ├── python/                          单元测试
│   ├── qlora / saving / security / ...
│   └── run_all.sh
├── scripts/                              辅助脚本
├── pyproject.toml
└── README.md

13.2五条主线代码

13.3调用栈：从用户代码到 Triton kernel

13.4unsloth/models/loader.py：入口

这是最重要的入口文件。FastLanguageModel.from_pretrained() 大致流程：

class FastLanguageModel:
    @staticmethod
    def from_pretrained(model_name, max_seq_length, dtype, load_in_4bit, ...):
        # 1. 探测模型架构
        config = AutoConfig.from_pretrained(model_name)
        model_type = config.model_type   # "llama" / "qwen2" / "gemma" / ...

        # 2. 选 patcher 类
        if model_type == "llama":
            from .llama import FastLlamaModel as PatcherClass
        elif model_type == "qwen3":
            from .qwen3 import FastQwen3Model as PatcherClass
        elif model_type == "gemma":
            from .gemma import FastGemmaModel as PatcherClass
        # ... 各种 if elif

        # 3. 调用 patcher 的 pre_patch（在加载模型前 patch 类）
        PatcherClass.pre_patch()

        # 4. 加载实际权重（这一步会触发 patch 后的 forward）
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            quantization_config=bnb_config if load_in_4bit else None,
            torch_dtype=dtype,
            ...
        )

        # 5. tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        tokenizer = fix_tokenizer(tokenizer)         # padding_side / eos / pad

        # 6. post-patch
        PatcherClass.post_patch(model, tokenizer)

        return model, tokenizer

13.5unsloth/models/llama.py：单个架构 patcher

每种架构一个 patcher 类。以 llama.py 为例，大致结构：

class FastLlamaModel:
    @staticmethod
    def pre_patch():
        # 1. 替换 LlamaAttention.forward
        from transformers.models.llama.modeling_llama import LlamaAttention
        LlamaAttention.forward = unsloth_attention_forward

        # 2. 替换 RMSNorm
        from transformers.models.llama.modeling_llama import LlamaRMSNorm
        LlamaRMSNorm.forward = unsloth_rms_forward

        # 3. 替换 RoPE
        from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
        LlamaRotaryEmbedding.forward = unsloth_rotary_forward

        # 4. 替换 MLP（SwiGLU）
        from transformers.models.llama.modeling_llama import LlamaMLP
        LlamaMLP.forward = unsloth_mlp_forward

        # 5. 替换 cross_entropy_loss
        import transformers
        transformers.modeling_utils.cross_entropy_loss = unsloth_cross_entropy

    @staticmethod
    def post_patch(model, tokenizer):
        # 1. 设 padding_side
        # 2. 处理 EOS / PAD token
        # 3. 关键：注入 use_cache=False 等训练设置

    @staticmethod
    def get_peft_model(model, r, ...):
        # 1. 调 peft.get_peft_model
        # 2. 把 peft.LoraLayer 替换为 fast_lora.Linear
        # 3. 处理 use_gradient_checkpointing="unsloth"

13.6unsloth/kernels/fast_lora.py

fused LoRA 实现。核心是几个 Triton kernel：

读这个文件能学到的 Triton 风格：

• autotune 配置：声明几组 BLOCK_M / BLOCK_N / num_warps，让 Triton 选最佳；
• guarded loads：用 mask 处理边界，避免越界；
• 反向 kernel 复用前向 layout，省一次内存读写。

13.7unsloth/kernels/cross_entropy_loss.py

chunked CE 是 Unsloth 显存优势的核心。关键设计：

@triton.jit
def _cross_entropy_forward(
    logits_ptr, labels_ptr, loss_ptr, ...,
    VOCAB_SIZE, CHUNK_SIZE,
):
    # 处理一个 (chunk_size, vocab_size) 的 logits 块
    # 1. 算 max(logits)（数值稳定）
    # 2. 算 logsumexp
    # 3. 算 -log_softmax(label)
    # 4. 累加 loss


def chunked_ce_forward(hidden, lm_head, labels, chunk_size=1024):
    # 不实例化完整 logits 矩阵
    losses = []
    for chunk_start in range(0, hidden.shape[0], chunk_size):
        h_chunk = hidden[chunk_start:chunk_start+chunk_size]
        l_chunk = labels[chunk_start:chunk_start+chunk_size]
        # 当场算 logits，当场算 loss，当场释放 logits
        logits = h_chunk @ lm_head.T
        loss = triton_ce(logits, l_chunk)
        losses.append(loss)
    return sum(losses) / total_tokens

反向也分块，整个训练过程都不会出现完整 logits 矩阵。

13.8unsloth/kernels/rms_layernorm.py

RMSNorm fused 实现。前向：

@triton.jit
def _rms_layernorm_forward(
    Y_ptr, X_ptr, W_ptr, R_ptr,
    n_cols, eps,
    BLOCK_SIZE: tl.constexpr,
):
    row_idx = tl.program_id(0)
    X_row = X_ptr + row_idx * n_cols
    Y_row = Y_ptr + row_idx * n_cols

    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(X_row + cols, mask=mask, other=0.0).to(tl.float32)

    # 算 1/RMS（用 rsqrt）
    sum_sq = tl.sum(x * x)
    rms = tl.math.rsqrt(sum_sq / n_cols + eps)
    tl.store(R_ptr + row_idx, rms)             # 反向需要

    w = tl.load(W_ptr + cols, mask=mask, other=0.0)
    y = (x * rms).to(Y_ptr.dtype.element_ty) * w
    tl.store(Y_row + cols, y, mask=mask)

整层 norm 用一个 kernel，对 hidden state 只读写一次（vs. 原 PyTorch 5 次）。

13.9unsloth/chat_templates.py

这个文件维护所有支持模型的 chat template + tokenizer 修复，是用户最常碰的工具文件。

支持的模板："llama-3", "chatml", "gemma", "mistral", "phi-3", "qwen-2.5", "qwen-3", "unsloth"（自家通用）。

13.10unsloth/save.py

所有 save_pretrained_* 实现都在这里：

读这里能理解 GGUF 量化的整个流程（自动 clone llama.cpp + 编译 + 调 convert + quantize）。

13.11阅读路线总结

1. 跑通 ch03 hello world：建立感觉；
2. 读 unsloth/__init__.py：看 MLX 检测和 import 入口；
3. 读 unsloth/models/loader.py：理解 FastLanguageModel 的实现；
4. 读 unsloth/models/llama.py：理解 patcher 风格；
5. 读 unsloth/kernels/rms_layernorm.py：学一个简单 Triton kernel；
6. 读 unsloth/kernels/fast_lora.py：硬核 Triton 实战；
7. 读 unsloth/kernels/cross_entropy_loss.py：理解 chunked 思想；
8. 读 unsloth/save.py：理解 GGUF 导出全流程；
9. 看一份 tests/notebooks/Llama3_(8B)-Alpaca.ipynb：作为参考实现。

13.12想给 Unsloth 提 PR？从这里入手

13.13真实源码树（v2026）

unsloth/
├── _gpu_init.py:92            ★ enforce unsloth_zoo >= 2026.5.2
├── import_fixes.py            HF transformers / PEFT / TRL 兼容修补
├── device_type.py              CUDA / ROCm / Intel XPU / MLX
├── trainer.py                  UnslothTrainer + Q-GaLore（line 219-339）
├── chat_templates.py（127KB） 500+ chat 模板
├── tokenizer_utils.py
├── save.py:1339                save_to_gguf / push_to_hub_merged
│
├── models/                    ★ 按模型一个文件
│   ├── loader.py:241          FastLanguageModel / FastModel
│   ├── llama.py（149KB）       Llama 全家 + get_peft_model:2833
│   ├── mistral.py              含 patch_mistral_nemo_attention:416
│   ├── qwen2.py / qwen3.py / qwen3_moe.py
│   ├── gemma.py / gemma2.py
│   ├── granite.py:135          含 Granite-VL
│   ├── cohere.py / falcon_h1.py
│   ├── vision.py（73KB）       FastVisionModel + get_peft_model:1340
│   ├── rl.py（98KB）           GRPO/DPO/ORPO/KTO patch
│   ├── sentence_transformer.py（98KB）  embedding 模型
│   ├── _utils.py（128KB）       patch_fast_lora:2544、patch_llama_rope_scaling:2060
│   └── mapper.py / dpo.py
│
├── kernels/                   ★ Triton 实现
│   ├── fast_lora.py             apply_lora_mlp_swiglu:235 / apply_lora_qkv:535
│   ├── rope_embedding.py        Fast_RoPE_Embedding:185
│   ├── rms_layernorm.py         Fast_RMS_Layernorm:162
│   ├── cross_entropy_loss.py    chunked CE
│   ├── swiglu.py / geglu.py / layernorm.py
│   ├── flex_attention.py        PyTorch 2.6+
│   ├── fp8.py                    ★ 2026 新
│   └── moe/                      MoE 专用
│
├── dataprep/                  RawTextDataLoader / TextPreprocessor
├── registry/                  model 元数据
├── optimizers/                Q-GaLore
└── utils/                      packing / 注意力 dispatch

13.14unsloth_zoo 依赖（外置库）

不能少：_gpu_init.py:92 硬卡 unsloth_zoo >= 2026.5.2。

13.15延伸阅读