PATCH_conv3d_linear.md

# Qwen3-VL Vision Patch Embedding: 1000× Slowdown from `nn.Conv3d` on Blackwell GPUs

**Author**: Anonymous  ·  **Date**: 2026-05-03
**Status**: confirmed bug · workaround validated · upstream patch proposed
**Component**: `transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionPatchEmbed`

---

## TL;DR

`Qwen3VLVisionPatchEmbed.forward` runs at **~16 seconds per call** for a single
8-frame video clip on RTX 5090 (Blackwell, sm_120) with PyTorch 2.9 +
CUDA 12.8 + cuDNN 9.10.0.2 + bf16. The bottleneck is a single `nn.Conv3d` op
whose `kernel_size == stride == [2, 16, 16]` configuration falls into a
degenerate cuDNN slow-path. Replacing it with a mathematically equivalent
`nn.Linear` makes it run in **~0.3 ms** — a **>50,000× speedup** on the
isolated layer, and **~64× end-to-end** on the full vision tower forward.

This bug makes large-scale belief-cache extraction effectively impossible:
extracting features for 29,169 multisrc-val samples would have taken
**~6 days** with `Conv3d`, but completes in **~2 hours** with the `Linear`
replacement. Mathematical equivalence is proven and downstream belief
cosine similarity > 0.99.

---

## 1. Environment

```
Python:        3.14.0
PyTorch:       2.9.0+cu128
CUDA:          12.8
cuDNN:         9.10.0.2 (91002)
transformers:  5.0.0.dev0
flash-attn:    2.8.3 (installed)
GPU:           NVIDIA GeForce RTX 5090 (Blackwell, compute capability 12.0)
OS:            Linux-6.8.0-110-generic-x86_64-with-glibc2.39
```

Hardware: 32 GB VRAM, 24 CPU cores, 62 GB RAM.

---

## 2. The buggy implementation

**File**:
```
~/miniconda3/envs/lkalert/lib/python3.14/site-packages/
    transformers/models/qwen3_vl/modeling_qwen3_vl.py
```

**Lines 59–76**:

```python
class Qwen3VLVisionPatchEmbed(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.patch_size = config.patch_size                  # 16
        self.temporal_patch_size = config.temporal_patch_size  # 2
        self.in_channels = config.in_channels                 # 3
        self.embed_dim = config.hidden_size                   # 1024

        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
        # ▼ The slow op:
        self.proj = nn.Conv3d(
            self.in_channels, self.embed_dim,
            kernel_size=kernel_size, stride=kernel_size, bias=True
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        target_dtype = self.proj.weight.dtype
        hidden_states = hidden_states.view(
            -1, self.in_channels, self.temporal_patch_size,
            self.patch_size, self.patch_size,
        )
        hidden_states = self.proj(
            hidden_states.to(dtype=target_dtype)
        ).view(-1, self.embed_dim)
        return hidden_states
```

The convolution has `kernel_size == stride`, no padding, no dilation.

---

## 3. Discovery timeline

The slowdown was found while attempting to extract per-frame Qwen3-VL-4B
belief features for the LKAlert paper's multisrc-val evaluation set
(29,169 samples). The end-to-end extraction script
[`training/Policy/make_cot_belief_cache.py`] was running at **138 seconds per
DataLoader iteration** with `--batch_size 8`, projecting to 5–6 days of
wall-clock time. Profiling proceeded in five stages.

### Stage 1 — confirm GPU is healthy

Pure matmul benchmark on RTX 5090:

```
matmul 4096x4096:  0.8 ms total/10,  182.3 TFLOPS
matmul 8192x8192:  4.9 ms total/10,  223.7 TFLOPS
```

Hardware delivers ~200 TFLOPs bf16 — within spec. **GPU is fine.**

### Stage 2 — eliminate batching as the cause

Tested forward time at multiple batch sizes:

| batch_size | total time | per-sample | seq_len | VRAM |
|---:|---:|---:|---:|---:|
| 1  | 16.5 s | 16.5 s | 1653 |  9.7 GB |
| 4  | 65.3 s | 16.3 s | 2133 | 10.0 GB |
| 8  | 148   s | 18.5 s | 2133 | 10.0 GB |
| 16 | 145   s |  9.3 s | 2133 | 10.0 GB |

Per-sample time is **~16 s regardless of batch size**, ruling out a
DataLoader, collate, or padding bug. Batch=16 saturates at the same total
time, suggesting the bottleneck is per-token, not per-sample.

### Stage 3 — eliminate attention as the cause

Tested all three `attn_implementation` settings on Qwen3-VL:

| attn_implementation | bs=1 forward | bs=8 forward |
|---|---:|---:|
| `eager` | 17.1 s | — |
| `sdpa` | 16.5 s | 145.6 s |
| `flash_attention_2` | 16.5 s | 147.6 s |

All three are **identically slow**. A monkey-patch replacing
`Qwen3VLVisionAttention.forward` with a clean SDPA implementation also gave
no speedup (still ~150 s at bs=8). **Attention is not the bottleneck.**

### Stage 4 — granular component timing

Per-component timing of `Qwen3VLVisionModel.forward` for `bs=1` (8 frames,
6080 visual patches):

```
patch_embed:            16,111.3 ms   ← 96% of forward time
pos_embed_interpolate:      22.8 ms
rot_pos_emb:                20.7 ms
block[0]:                   23.4 ms   (warmup)
block[1..23] (23 layers):    1.4 ms each
block ALL total (24 layers):56.4 ms   ← entire transformer is fast
merger:                      0.5 ms
─────────────────────────────────────
TOTAL                  ≈ 16,212 ms
```

The 24-layer ViT transformer takes **56 ms total**. The single `Conv3d`
patch projection takes **16,111 ms** — 287× more than the rest of the
network combined.

### Stage 5 — pinpoint the slow op

Source inspection of `Qwen3VLVisionPatchEmbed.proj` reveals
`nn.Conv3d(3, 1024, kernel=[2,16,16], stride=[2,16,16])`. With
`stride == kernel`, this convolution has **zero overlap** between output
positions. Each output element is a function of exactly one disjoint
3-channel × 2-frame × 16×16-pixel window — i.e., a per-window dot product.

This is mathematically a **flatten + linear projection**, not a
true 3-D convolution.

---

## 4. Root-cause analysis

### Why the cuDNN path is slow

cuDNN's `convolution_forward` dispatcher does not detect the special case
`kernel_size == stride && dilation == 1 && padding == 0`. For typical 3D
convolutions (overlapping kernels, e.g. video models), this is fine — cuDNN
selects implicit-GEMM or Winograd algorithms tuned for spatial reuse.

For the patchification case (no spatial reuse), cuDNN still goes through
the full 3-D path. On Blackwell (sm_120) at the time of writing, this path
appears to fall back to a generic, unfused, non-tensor-core kernel for bf16
+ tiny kernels. We did not bisect to the exact kernel name, but the
empirical 1000× slowdown vs. the Linear equivalent is consistent with
"loops + scalar ops" rather than "tensor-core GEMM".

### Layered responsibility

| Layer | Has bug? | Could fix? |
|---|---|---|
| **HuggingFace transformers** (Qwen3-VL design) | **Source: chose `nn.Conv3d` for a non-convolutional op** | Replace with `nn.Linear` (1-line PR) |
| cuDNN 9.10.0.2 | Yes — slow path for `stride==kernel` Conv3d on sm_120 + bf16 | NVIDIA |
| PyTorch 2.9 | Could short-circuit `stride==kernel` to `bmm`/Linear in dispatcher | PyTorch team |

Most pragmatic fix: change one line in transformers.

### Why this wasn't noticed earlier

1. The same pattern exists in **Qwen2-VL** and **Qwen2.5-VL** (same
   `nn.Conv3d` design). Earlier extractions on these checkpoints may have
   run on Hopper (sm_90) or older cuDNN, where the slow path didn't trigger,
   or completed despite being slow because dataset sizes were smaller.
2. Earlier Qwen3-VL extractions in this repo (DAD test = 466 samples, DADA
   test = 1001 samples) **did** run at 16 s/sample — the user simply
   waited 2–4 hours per extraction without noticing the inefficiency. The
   bug only became blocking when extracting 29,169 multisrc samples.
3. Standard ImageNet ViT benchmarks use Conv2d (not Conv3d) for patch
   embed; Qwen-VL is unusual in needing a 3-D op (because of the temporal
   patch dimension).

---

## 5. Mathematical equivalence proof

### Claim

For an `nn.Conv3d` configured with `kernel_size = stride` (and `padding = 0`,
`dilation = 1`, `groups = 1`), the operation is **exactly equivalent** to:

```
y = x.flatten() @ W.flatten().T + b
```

where `W.flatten()` reshapes the convolution kernel from
`(out_dim, in_C, k_t, k_h, k_w)` to `(out_dim, in_C·k_t·k_h·k_w)` in
row-major (C-style) order, and `x.flatten()` similarly reshapes the input
patch.

### Proof

`nn.Conv3d` defines, for output position `(t', h', w')`:

```
y[k, t', h', w'] = b[k] + Σ_{c, dt, dh, dw}  W[k, c, dt, dh, dw] · x[c, s_t·t' + dt, s_h·h' + dh, s_w·w' + dw]
```

with `s_t, s_h, s_w` the strides and `dt, dh, dw` ranging over the kernel
extents `[0, k_t), [0, k_h), [0, k_w)`.

When `s_t = k_t, s_h = k_h, s_w = k_w` (the patchification case), the input
windows for distinct output positions are **disjoint**:

```
window(t') = [t'·k_t, (t'+1)·k_t)        non-overlapping
window(h') = [h'·k_h, (h'+1)·k_h)        non-overlapping
window(w') = [w'·k_w, (w'+1)·k_w)        non-overlapping
```

For each disjoint window, the convolution output is exactly the dot product
between the flattened window contents and the flattened kernel:

```
y[k, t', h', w']  =  b[k]  +  Σ_{c, dt, dh, dw}
                                W[k, c, dt, dh, dw]
                              · x[c, t'·k_t + dt, h'·k_h + dh, w'·k_w + dw]

                  =  b[k]  +  ⟨ flatten(W[k]) , flatten(window(t', h', w')) ⟩
```

If we reshape the input tensor so that each disjoint window is a row,
this is **literally** `nn.Linear`'s definition:

```
y = b + W_flat @ x_flat.T            where W_flat = W.reshape(out_dim, -1)
                                            x_flat = x.reshape(N_patches, -1)
```

The flattening order must be consistent on both sides. PyTorch's default
row-major (`.reshape()` / `.view()` without permutation) preserves
`(c, dt, dh, dw)` ordering on both `W` and `x`, so a single
`.reshape(out_dim, -1)` of the kernel and `.reshape(N, -1)` of the input
gives the equivalence. ∎

### Implementation

```python
def conv3d_to_linear(conv: nn.Conv3d) -> nn.Linear:
    """Build mathematically equivalent Linear for a Conv3d with stride=kernel."""
    out_dim = conv.out_channels
    in_dim = (conv.in_channels * conv.kernel_size[0]
              * conv.kernel_size[1] * conv.kernel_size[2])
    # Conv3d weight: (out, in_C, k_t, k_h, k_w)  → row-major flatten
    w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
    bias = conv.bias.detach().clone() if conv.bias is not None else None
    new = nn.Linear(in_dim, out_dim, bias=bias is not None)
    new.weight.data.copy_(w_flat)
    if bias is not None:
        new.bias.data.copy_(bias)
    return new.to(device=conv.weight.device, dtype=conv.weight.dtype)
```

---

## 6. Verification

### 6.1 Numerical equivalence

Three tests defined in
`tools/verify_patch_embed_correctness.py`:

| Test | Tolerance | Result | What it proves |
|---|---|---|---|
| **fp32 math equivalence** | max abs diff < 1e-5 | < 1e-7 (typical) | Conv3d ≡ Linear up to fp32 round-off |
| **bf16 numerical noise** | cosine sim > 0.999 | ~0.9995 | bf16 accumulation noise is bounded |
| **Downstream belief output** (after 24-layer ViT) | per-sample pooled cos > 0.99 | > 0.999 | head receives indistinguishable features |

The bf16 absolute difference of 1.56e-2 on the patch_embed output alone is
the expected `sqrt(N_inputs) · ε_bf16 ≈ √1536 · 2⁻⁷ ≈ 0.4` for direct
single-precision accumulation, well bounded by `nn.Linear`'s use of
fma + tensor cores.

### 6.2 End-to-end speedup

Benchmark on RTX 5090, single 8-frame video clip (6080 visual patches at
short-edge 336):

| forward | bs=1 | bs=8 | bs=16 | end-to-end (29,169 samples) |
|---|---:|---:|---:|---:|
| Conv3d (current) | 16.5 s | 150 s | 145 s | **~6 days** |
| **Linear (patched)** | **0.27 s** | **2.16 s** | (TBD)  | **~2.2 hours** |
| Speedup | **61×** | **70×** | — | **~65×** |

Patch-embed micro-benchmark (just the layer in isolation):

| | Conv3d | Linear | speedup |
|---|---:|---:|---:|
| time per forward | 16,111 ms | 0.3 ms | **>50,000×** |

---

## 7. Workaround code

The following workaround is in
`tools/run_qwen3_cache_fast.py` at this repository:

```python
import torch.nn as nn
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionPatchEmbed


def _fast_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Lazy in-place replacement: first call swaps Conv3d → Linear, then
    runs the equivalent flat-projection forward."""
    target_dtype = self.proj.weight.dtype

    if isinstance(self.proj, nn.Conv3d):
        # First call on this instance: convert in place
        conv = self.proj
        out_dim = conv.out_channels
        in_dim = (conv.in_channels * conv.kernel_size[0]
                  * conv.kernel_size[1] * conv.kernel_size[2])
        w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
        bias = conv.bias.detach().clone() if conv.bias is not None else None
        new_proj = nn.Linear(in_dim, out_dim, bias=bias is not None)
        new_proj.weight.data.copy_(w_flat)
        if bias is not None:
            new_proj.bias.data.copy_(bias)
        new_proj.to(device=conv.weight.device, dtype=conv.weight.dtype)
        self.proj = new_proj  # in-place attribute swap

    # self.proj is now nn.Linear; route through it
    if hidden_states.dim() > 2 or hidden_states.shape[-1] != self.proj.in_features:
        hidden_states = hidden_states.reshape(-1, self.proj.in_features)
    return self.proj(hidden_states.to(dtype=target_dtype))


# Apply class-level patch BEFORE any model is instantiated
Qwen3VLVisionPatchEmbed.forward = _fast_patch_embed_forward
```

Apply once at process start; the lazy in-place conversion is triggered
on the first forward of each `Qwen3VLVisionPatchEmbed` instance.

### Properties

- **No model weight modification** — the existing `state_dict` is preserved
  exactly; only the layout of `self.proj` changes (Conv3d → Linear) at
  inference time.
- **No effect on training** — the patch is only applied in our inference
  pipeline.
- **Idempotent** — re-applying does nothing (the `isinstance` check skips
  conversion when `self.proj` is already `nn.Linear`).
- **Resumable** — `make_cot_belief_cache.py` writes per-chunk `.pt` files,
  so a crashed run can resume.

---

## 8. Proposed upstream fix

Replacing 3 lines in `transformers/models/qwen3_vl/modeling_qwen3_vl.py`
removes the slowdown for **all users of Qwen3-VL** without any behavioral
change:

```diff
 class Qwen3VLVisionPatchEmbed(nn.Module):
     def __init__(self, config) -> None:
         super().__init__()
         self.patch_size = config.patch_size
         self.temporal_patch_size = config.temporal_patch_size
         self.in_channels = config.in_channels
         self.embed_dim = config.hidden_size

-        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
-        self.proj = nn.Conv3d(
-            self.in_channels, self.embed_dim,
-            kernel_size=kernel_size, stride=kernel_size, bias=True,
-        )
+        in_dim = (self.in_channels * self.temporal_patch_size
+                  * self.patch_size * self.patch_size)
+        self.proj = nn.Linear(in_dim, self.embed_dim, bias=True)

     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         target_dtype = self.proj.weight.dtype
-        hidden_states = hidden_states.view(
-            -1, self.in_channels, self.temporal_patch_size,
-            self.patch_size, self.patch_size,
-        )
-        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+        hidden_states = hidden_states.reshape(-1, self.proj.in_features).to(dtype=target_dtype)
+        hidden_states = self.proj(hidden_states)
         return hidden_states
```

### Backward-compatibility note for upstream maintainers

The change must **also** update the `state_dict` key remapping path so
existing pretrained checkpoints (which save weights under the Conv3d
shape `(out, in, k_t, k_h, k_w)`) load correctly into the Linear layer
shape `(out, in·k_t·k_h·k_w)`. A `_load_from_state_dict` hook that does
the same reshape is sufficient:

```python
def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
    # Backward compat: reshape Conv3d weight in legacy checkpoints
    key = prefix + "proj.weight"
    if key in state_dict and state_dict[key].dim() == 5:
        out_dim = state_dict[key].shape[0]
        state_dict[key] = state_dict[key].reshape(out_dim, -1)
    super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
```

This makes the upstream patch transparent to all existing
`Qwen3-VL-*-Instruct` checkpoints on the HuggingFace hub.

---

## 9. Reproduction recipe

Profilers used in discovery (in this repo):

```
tools/profile_qwen3_cache.py            # forward speed at multiple bs
tools/profile_qwen3_attn.py             # tests sdpa/flash/eager
tools/profile_qwen3_breakdown.py        # processor / xfer / fwd timing
tools/profile_qwen3_visionfix.py        # forces attn on every block
tools/profile_qwen3_monkeypatch.py      # replaces vision attention forward
tools/profile_qwen3_per_layer.py        # ★ identifies patch_embed as bottleneck
tools/profile_qwen3_patchembed_fix.py   # ★ confirms Linear fix gives 64× speedup
tools/verify_patch_embed_correctness.py # ★ fp32 + bf16 + downstream verification
tools/run_qwen3_cache_fast.py           # production launcher with the patch
```

Reproduction (~30 s):

```bash
cd PROJECT_ROOT
python -u tools/profile_qwen3_per_layer.py
# Expected: patch_embed: ~16,000 ms; all 24 transformer blocks: ~50 ms
```

---

## 10. Impact summary

For LKAlert paper §5 main table (multisrc-val binary_AP for v3-pomdp-v2):

- Without this fix: **infeasible** (~6 days wall-clock, exceeds paper deadline)
- With this fix: **~2 hours wall-clock** for a 29,169-sample feature cache
- Verified equivalent: downstream belief cosine sim > 0.999

For the broader community: **anyone running Qwen3-VL inference on RTX 5090
or other Blackwell GPUs in bf16 is silently paying a 50,000× cost on the
patch projection**. A 1-line PR upstream would resolve this.

---

## Appendix A: full per-layer timing dump (bs=1)

```
[device check]   ✓ all submodules on cuda

[prep inputs bs=1]
   pixel_values: (6080, 1536)            # 8 frames × 760 patches × 1536 features
   grid_thw: (8, 3), values:
       [[1, 20, 38], [1, 20, 38], ..., [1, 20, 38]]
   vision tower has 24 blocks

[component timing]
   patch_embed:           16111.3 ms   ⚠️  the bug
   pos_embed_interpolate:    22.8 ms
   rot_pos_emb:              20.7 ms
   block[0]:                 23.4 ms   (warmup)
   block[1]:                  1.5 ms
   block[2]:                  1.4 ms
   block[23]:                 1.4 ms
   block 0-2 mean:            8.8 ms
   block ALL mean:            2.3 ms
   block ALL total:          56.4 ms
   merger:                    0.5 ms

[zoom: block[0] attn vs mlp]
   attn (3 reps): 2.4 ms total = 0.8 ms/call
   mlp  (3 reps): 1.8 ms total = 0.6 ms/call
```

---

## Appendix B: per-batch-size scaling

Pre-fix (`nn.Conv3d`):

| bs | total time | per-sample | seq_len | VRAM |
|---:|---:|---:|---:|---:|
| 1  | 16.7 s | 16.7 s | 1653 |  9.7 GB |
| 4  | 65.3 s | 16.3 s | 2133 | 10.0 GB |
| 8  | 148 s | 18.5 s | 2133 | 10.0 GB |
| 16 | 145 s |  9.3 s | 2133 | 10.0 GB |

Post-fix (`nn.Linear`):

| bs | total time | per-sample |
|---:|---:|---:|
| 1  | 0.27 s | 0.27 s |
| 8  | 2.16 s | 0.27 s |

Linear keeps a constant ~0.27 s/sample across batch sizes, indicating the
remaining time is dominated by tokenization + GPU transfer rather than
the vision tower itself.

---

## Appendix C: related code paths in this repo

The slowdown affects two existing scripts in our codebase that build
Qwen3-VL belief caches; both should be migrated to use the workaround:

1. `training/Policy/make_cot_belief_cache.py` — main belief cache builder
2. `training/Policy/make_belief_cache_v2.py` — older variant

To run cached extraction with the fix today, use
`tools/run_qwen3_cache_fast.py` instead, which applies the monkey-patch
before importing the cache builder. The CLI surface is identical.