# monkeypatch.lora_kernels – Axolotl

# monkeypatch.lora_kernels

`monkeypatch.lora_kernels`

Module for patching custom LoRA Triton kernels and `torch.autograd` functions.

## Classes

Name | Description
--- | ---
FakeMLP | placeholder MLP for triton patching

### FakeMLP

    monkeypatch.lora_kernels.FakeMLP(gate_proj, up_proj, down_proj)

placeholder MLP for triton patching

## Functions

Name | Description
--- | ---
apply_lora_kernel_patches | Applies optimized Triton kernel patches to a PEFT model.
get_attention_cls_from_config | Get the appropriate attention class by inspecting the model config.
get_layers | Get the layers of the model. Handles text-only and multimodal models.
original_apply_o | Original implementation of output projection without optimizations.
original_apply_qkv | Original implementation of QKV projection without optimizations.
patch_self_attn_lora | Given an `axolotl` config, this method patches the inferred attention class forward

### apply_lora_kernel_patches

    monkeypatch.lora_kernels.apply_lora_kernel_patches(model, cfg)

Applies optimized Triton kernel patches to a PEFT model.

Patches a PEFT model with optimized implementations for MLP and attention computations. The optimizations include custom Triton kernels for activation functions and specialized autograd functions for LoRA computations.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
model | PeftModelForCausalLM | A PEFT model to be patched with optimized kernels. | required
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required

#### Returns

Name | Type | Description
--- | --- | ---
PeftModelForCausalLM | PeftModelForCausalLM | The patched model with optimized kernels.

#### Raises

Name | Type | Description
--- | --- | ---
| TypeError | If the provided model is not a `PeftModelForCausalLM`.
| NotImplementedError | If the model type is not supported.
| AssertionError | If multiple adapters are active (currently unsupported).

#### Note

The optimizations require LoRA adapters with no dropout and no bias terms. The function will skip patching if these conditions aren’t met.

### get_attention_cls_from_config

    monkeypatch.lora_kernels.get_attention_cls_from_config(cfg)

Get the appropriate attention class by inspecting the model config. Uses dynamic import to support any model architecture that follows the standard transformers naming convention.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required

#### Returns

Name | Type | Description
--- | --- | ---
| Type[nn.Module] | The appropriate attention class for the model.

#### Raises

Name | Type | Description
--- | --- | ---
| ValueError | If `base_model` not specified or attention class cannot be imported
| ImportError | If the model module or attention class doesn’t exist

### get_layers

    monkeypatch.lora_kernels.get_layers(model)

Get the layers of the model. Handles text-only and multimodal models.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
model | PeftModelForCausalLM | A PEFT model. | required

#### Returns

Name | Type | Description
--- | --- | ---
| list[nn.Module] | A list of layers.

### original_apply_o

    monkeypatch.lora_kernels.original_apply_o(self, hidden_states)

Original implementation of output projection without optimizations.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
self | nn.Module | The attention module instance. | required
hidden_states | torch.Tensor | Input tensor of shape `[`batch_size, seq_len, hidden_dim]`. | required

#### Returns

Name | Type | Description
--- | --- | ---
| torch.Tensor | The output projection result.

### original_apply_qkv

    monkeypatch.lora_kernels.original_apply_qkv(self, hidden_states)

Original implementation of QKV projection without optimizations.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
self | nn.Module | The attention module instance. | required
hidden_states | torch.Tensor | Input tensor of shape [batch_size, seq_len, hidden_dim]. | required

#### Returns

Name | Type | Description
--- | --- | ---
| tuple[torch.Tensor, torch.Tensor, torch.Tensor] | A tuple `(query_states, key_states, value_states)` containing the projected states for query, key, and value.

### patch_self_attn_lora

    monkeypatch.lora_kernels.patch_self_attn_lora(cfg)

Given an `axolotl` config, this method patches the inferred attention class forward pass with optimized LoRA implementations.

It modifies the attention class to use optimized QKV and output projections. The original implementation is preserved and can be restored if needed.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required

#### Raises

Name | Type | Description
--- | --- | ---
| AssertionError | If the required code blocks are not found in the attention implementation.