PERF MSCFormer: temporal and spatial convolutions of the patch embedding as one convolution (train step 0.63x Gaudi eager, 0.66x lazy, 0.63x CPU at 21 x 6000) - #1381

Draft: one gate cell needs a decision (see "Open cell" below).

Each of MSCFormer's three patch-embedding branches runs a temporal convolution (1 -> 16 filters, kernel 85 / 65 / 45, bias) and then a depthwise spatial convolution (16 -> 16, all channels, bias) with nothing in between. Both are linear, so this PR runs them as one convolution over the channels with the composed weight spatial.weight.squeeze(1) * temporal.weight.squeeze(1) (filters, n_chans, kernel) and the composed bias (the algebra of the EEGConformer #1373 and MIRepNet #1377 stems). The rest of each branch (batch norm, ELU, pooling, dropout) is unchanged. Same parameters, names, state dict and init; gradients reach both original weights through the composition. The (batch, 16, n_chans, n_times) intermediate goes away (516 MB per branch at 21 x 6000, batch 64).

A spatial-first form (project the channels, then a depthwise temporal convolution) was faster still at 21 x 6000 (Gaudi eager 0.58x, CPU 0.38x) but 1.18x slower on CPU at the registry input (CPU depthwise long convolutions at batch 2); the composed form is not.

Tier: numerical only (Tier C numerical-only cells disclosed below, accepted by Bruno on 2026-10-11 at 16:28). Same initial state dict (0.0), same dropout draws, argmax identical, TorchScript scripted == eager (0.0).

Speed (train step = forward + cross-entropy + backward + AdamW, paired in one process, master and PR alternated, median of per-round ratios PR / master)

Gaudi small-input cells: 60 rounds; at these inputs the Gaudi steps are host-bound (host 31-55 ms of 34-72 ms) and vary by +-10-15 % between rounds (master -> master controls in earlier runs: 0.987-1.053).

Open cell

At the registry input (batch 2, 3 channels) in Gaudi lazy mode the train step is 1.05-1.10x master in four independent runs (1.052, 1.088, 1.086, 1.102), and Gaudi eager evaluation at batch 2-32 is 1.06-1.11x: the weight and bias composition adds about five small ops per branch per call, which costs host time when the step is host-bound. Everywhere else the PR is faster or within noise. Per the speed rule this needs a decision before merging.

Numerics (max abs vs master, same weights and input; yardstick = master fp32 vs master fp64)

• CPU, registry / 22 x 1000 / 21 x 6000, seeds 0 and 1 (composed form): eval out up to 1.9x master's own yardstick (1.9e-7 vs 1.0e-7 at the registry input; 1.1-1.5x elsewhere), train out <= 3.3e-7 (yardsticks 0.08-0.41), grads global <= 6.1e-5 (yardsticks 0.9-7.4), per tensor up to 2.8x (the convolution biases in front of batch norm and attention key biases, whose exact gradient is 0), batch-norm buffers up to 3.4x (<= 6.0e-8 absolute). bfloat16/float16 forwards run.
• Gaudi, 21 x 6000, batch 16 (PR - master vs master - CPU fp64): eager eval 2.4e-7 vs 3.6e-5, train 3.3e-7 vs 0.13, grads 1.1e-5 vs 1.74 (per tensor up to 2.7x, same exact-zero-gradient tensors), buffers 6.0e-8 vs 1.8e-7; lazy eval 3.0e-7, train 6.3e-7, grads 3.1e-5 (per tensor 2.2x), buffers 1.9x.

No released checkpoints for MSCFormer.

Tests (Voyager, CPU)

pytest test/unit_tests/models -k MSCFormer: 31 passed, 2 skipped (new test_mscformer.py: the patch embedding equals its branches run in order, float64 values and grads). pytest test/unit_tests/models/test_return_features.py: 137 passed, 5 skipped.

Files: braindecode/models/mscformer.py, test/unit_tests/models/test_mscformer.py, docs/whats_new.rst.