diff --git a/pyproject.toml b/pyproject.toml
index ae6a649bef..fd88925c49 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,12 @@ include-package-data = false
 exclude = ["images*"]
 
 [project.optional-dependencies]
+triton = [
+    "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp39-cp39-win_amd64.whl ; python_version=='3.9' and platform_system == 'Windows'",
+    "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp310-cp310-win_amd64.whl ; python_version=='3.10' and platform_system == 'Windows'",
+    "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp311-cp311-win_amd64.whl ; python_version=='3.11' and platform_system == 'Windows'",
+    "triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
+]
 huggingface = [
     "unsloth_zoo>=2024.12.3",
     "packaging",
@@ -49,6 +55,7 @@ huggingface = [
     "protobuf<4.0.0",
     "huggingface_hub",
     "hf_transfer",
+    "unsloth[triton]",
 ]
 cu118only = [
     "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.22.post7%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9' and platform_system == 'Linux'",
diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 1d9a0c1334..c654343282 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -77,7 +77,7 @@ pass
 def GemmaDecoderLayer_fast_forward(
     self,
     hidden_states:        torch.Tensor,
-    causal_mask:          Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None,
+    causal_mask:          Optional[BlockDiagonalCausalMask] = None,
     attention_mask:       Optional[torch.Tensor] = None,
     position_ids:         Optional[torch.LongTensor] = None,
     past_key_value:       Optional[Tuple[torch.Tensor]] = None,
diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py
index e47a7434f2..4a6f3e97cc 100644
--- a/unsloth/models/gemma2.py
+++ b/unsloth/models/gemma2.py
@@ -75,7 +75,7 @@ pass
 def Gemma2Attention_fast_forward(
     self,
     hidden_states:        torch.Tensor,
-    causal_mask:          Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None,
+    causal_mask:          Optional[BlockDiagonalCausalMask] = None,
     attention_mask:       Optional[torch.Tensor] = None,
     position_ids:         Optional[torch.LongTensor] = None,
     past_key_value:       Optional[Tuple[torch.Tensor]] = None,