Patch away the qwen3vl conv3d that only has slow bf16 kernels

extensions_built_in/captioner/Qwen3VLCaptioner.py

@@ -5,6 +5,8 @@ from transformers import (
 )
 from collections import OrderedDict
 
+import torch
+import torch.nn.functional as F
 from optimum.quanto import freeze
 from toolkit.basic import flush
 from toolkit.util.quantize import quantize, get_qtype
@@ -14,6 +16,29 @@ import transformers
 import logging
 import warnings
 
+
+def patch_qwen_vl_patch_embed(model):
+    """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain
+    linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and
+    falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear
+    (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves.
+    Returns the number of patch_embed modules patched."""
+    patched = 0
+    for module in model.modules():
+        proj = getattr(module, "proj", None)
+        if (
+            isinstance(proj, torch.nn.Conv3d)
+            and tuple(proj.kernel_size) == tuple(proj.stride)
+        ):
+            def fast_forward(hidden_states, _proj=proj):
+                w = _proj.weight.reshape(_proj.weight.shape[0], -1)
+                x = hidden_states.view(-1, w.shape[1]).to(w.dtype)
+                return F.linear(x, w, _proj.bias)
+
+            module.forward = fast_forward
+            patched += 1
+    return patched
+
 # transformers.logging.set_verbosity_error()
 warnings.filterwarnings("ignore")
 logging.disable(logging.WARNING)
@@ -35,6 +60,8 @@ class Qwen3VLCaptioner(BaseCaptioner):
             dtype=self.torch_dtype,
             device_map="cpu",
         )
+        # swap the slow bf16 Conv3d patch_embed for an equivalent fast linear
+        patch_qwen_vl_patch_embed(self.model)
         if not self.caption_config.low_vram:
             self.model.to(self.device_torch)
         if self.caption_config.quantize: