From 7c3b676a012b1e3a5a7e4f018e09412e2dfed0e7 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 9 Jul 2026 15:07:24 +0000
Subject: [PATCH] DeepSeek-V4: eager attention and trainable FP8 grouped
 experts

deepseek_v4 ships a custom attention that is not compatible with the sdpa
and flash paths, so add it to _EAGER_ONLY_PREFIXES to load with eager.

Its fused experts load as FP8GroupedLinear, whose forward calls a grouped
matmul kernel with no autograd formula, so loss.backward() fails during
finetuning. Patch the forward to dequantize the frozen fp8 weight and run a
differentiable grouped matmul while training, keeping the fused fp8 kernel
for inference.
---
 unsloth/kernels/fp8.py   | 27 +++++++++++++++++++++++++++
 unsloth/models/_utils.py |  3 ++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py
index 4efc4bd5d..17b0cd49a 100644
--- a/unsloth/kernels/fp8.py
+++ b/unsloth/kernels/fp8.py
@@ -42,6 +42,11 @@ except:
         "Unsloth: FP8 models need importing FP8Linear from `transformers.integrations.finegrained_fp8` but we don't see it."
     )
 
+try:
+    from transformers.integrations.finegrained_fp8 import FP8GroupedLinear
+except:
+    FP8GroupedLinear = None
+
 try:
     from transformers.integrations.fbgemm_fp8 import FbgemmFp8Linear
 except:
@@ -688,3 +693,25 @@ if FbgemmFp8Linear is not None:
     FbgemmFp8Linear.forward = module_forward_patch(fbgemm_fp8_linear, "weight_scale")
 if FP8Linear is not None:
     FP8Linear.forward = module_forward_patch(fp8_block_quant_linear, "weight_scale_inv")
+
+# FP8GroupedLinear.forward calls a grouped matmul kernel with no autograd
+# formula, so backward fails while training. For training, dequantize the frozen
+# fp8 weight and run a differentiable grouped matmul; inference keeps the kernel.
+# torch._grouped_mm matches but has no backward here; _scaled_grouped_mm needs
+# fp8 activations, so bmm on the dequantized weight is exact and universal.
+if FP8GroupedLinear is not None:
+    _fp8_grouped_forward_orig = FP8GroupedLinear.forward
+
+    def _fp8_grouped_forward(self, x):
+        if self.weight.element_size() > 1 or not torch.is_grad_enabled():
+            return _fp8_grouped_forward_orig(self, x)
+        hidden_dim = x.shape[-1]
+        W = weight_dequant(self.weight, self.weight_scale_inv.float()).to(x.dtype)
+        w = W.view(self.n_groups, -1, hidden_dim).transpose(1, 2)
+        xg = x.reshape(-1, self.n_groups, hidden_dim).transpose(0, 1)
+        y = torch.bmm(xg, w).transpose(0, 1).reshape(*x.shape[:-2], self.n_groups, -1)
+        if self.has_bias:
+            y = y + self.bias.view(self.n_groups, -1)
+        return y
+
+    FP8GroupedLinear.forward = _fp8_grouped_forward
diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py
index fa0e0b1c4..878d6edcc 100644
--- a/unsloth/models/_utils.py
+++ b/unsloth/models/_utils.py
@@ -438,7 +438,8 @@ DISABLE_SDPA_MODEL_NAMES = [
     "gpt_oss",
 ]
 _FLASH_EXCLUDED_MODELS = ("gpt_oss",)
-_EAGER_ONLY_PREFIXES = ("gemma3n",)
+# deepseek_v4 ships sdpa/flash-incompatible custom attention; force eager.
+_EAGER_ONLY_PREFIXES = ("gemma3n", "deepseek_v4")
 _FLASH_ATTENTION_MAX_HEAD_DIM = 256
 _FLASH_ATTENTION_DISABLED_WARNED = set()