Make fused_ff_qkv contiguous

[Michal-Novomestsky]

Currently, walrus is painfully slow to run (at least on my cluster - L40S is taking ~6-7s per forward pass for a 4x32x32x6 2D model with 6 timesteps and batch size 4).

After some sleuthing, it appears to stem from a single line of code:

fused_ff_qkv=rearrange(x, "b c h w d -> b h w d c")

This breaks the contiguity of fused_ff_qkv which causes ATen to dispatch to $BHWDC$[1, fused_dims] seperate GEMVs instead of a single [B*H*W*D*C, fused_dims] GEMM. In my case, this amounts to 4*32*32*1*6 = 24576 individual GEMV calls. Changing this gives me speedups of ~6x.

That said, idk what kind of workloads you guys use, especially at train time, so I'm unsure if the extra tensor copy will kill your VRAM, but that's nothing a little if not self.training can't fix.

[Michal-Novomestsky]

For more info, I quickly vibecoded some logging to compare the difference. Note that one aten::bmm is killing us. That's ff, q, k, v = self.fused_ff_qkv(fused_ff_qkv) if I'm not mistaken.

Before:

device: NVIDIA L40S torch 2.5.1+cu124
peaks: 181 TFLOP/s bf16, 92 TFLOP/s fp32, 864 GB/s
cudnn.benchmark=False resolution=128 n_steps=6
torch._C._cuda_set_sync_debug_mode(debug_mode)
latent grid 32x32x1 = 1024 spatial tokens x 6 timesteps
processor cost 15.07 TFLOP per sample
--- 1x1x1 Conv3d vs equivalent Linear (temporal qkv projection) ---
float32 conv3d 1.01 ms ( 72.5 TFLOP/s) linear 2.03 ms ( 36.0 TFLOP/s) conv is 0.5x slower
bfloat16 conv3d 0.50 ms ( 144.8 TFLOP/s) linear 0.29 ms ( 249.0 TFLOP/s) conv is 1.7x slower
--- float32 (no autocast) ---
B cold ms warm ms ms/sample TFLOP/s % peak
1 3254.6 3296.9 3296.9 4.6 5.0%
2 6754.7 6755.3 3377.6 4.5 4.9%
4 13617.2 13661.7 3415.4 4.4 4.8%
8 27407.7 27492.5 3436.6 4.4 4.8%
stage breakdown at B=8 (ms per forward):
blocks 27452.46 (100%)
decoder[2] 17.56 (0%)
encoder[2] 17.48 (0%)
jitter 0.77 (0%)
(unaccounted) 4.24 (0%) <- host-side prep, hook copies, stack/permute
--- bfloat16 (autocast) ---
B cold ms warm ms ms/sample TFLOP/s % peak
1 2412.2 2178.9 2178.9 6.9 3.8%
2 4499.1 4489.7 2244.9 6.7 3.7%
4 9169.8 9169.0 2292.2 6.6 3.6%
8 18479.1 18471.3 2308.9 6.5 3.6%
stage breakdown at B=8 (ms per forward):
blocks 18434.80 (100%)
encoder[2] 18.27 (0%)
decoder[2] 14.85 (0%)
jitter 0.51 (0%)
(unaccounted) 2.84 (0%) <- host-side prep, hook copies, stack/permute
--- materialised memory traffic (one forward, views excluded) ---
303.1 GB written by 9765 aten calls at B=4, bfloat16
>= 702 ms of pure memory time at peak bandwidth
mul 72.84 GB
_to_copy 58.99 GB
add 41.90 GB
bmm 25.29 GB
native_layer_norm 22.65 GB
clone 19.62 GB
pow 11.51 GB
convolution 11.29 GB
cat 11.09 GB
stack 8.33 GB
silu 5.97 GB
neg 4.15 GB
_scaled_dot_product_flash_attention 2.83 GB
mm 2.77 GB
addmm 2.77 GB
--- implicit device syncs: 5 per forward ---
5x called a synchronizing CUDA operation (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:150.)
--- torch.profiler, top ops by self device time ---
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg # of Calls 
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ aten::bmm 0.05% 4.375ms 0.52% 47.955ms 399.625us 7.915s 85.74% 7.915s 65.957ms 120 void cutlass::Kernel2<cutlass_80_tensorop_bf16_s1681... 0.00% 0.000us 0.00% 0.000us 0.000us 7.868s 85.23% 7.868s 196.711ms 40 aten::copy_ 0.08% 7.131ms 4.72% 434.240ms 263.655us 330.319ms 3.58% 330.319ms 200.558us 1647 aten::mul 0.07% 6.093ms 69.39% 6.380s 8.307ms 226.568ms 2.45% 226.568ms 295.010us 768 void at::native::elementwise_kernel<128, 2, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 217.705ms 2.36% 217.705ms 429.398us 507 aten::cat 0.04% 3.594ms 0.46% 42.309ms 208.417us 149.913ms 1.62% 149.913ms 738.488us 203 aten::cudnn_convolution 0.05% 4.185ms 0.66% 60.282ms 735.145us 132.465ms 1.43% 132.465ms 1.615ms 82 sm86_xmma_fprop_implicit_gemm_bf16bf16_bf16f32_f32_n... 0.00% 0.000us 0.00% 0.000us 0.000us 111.300ms 1.21% 111.300ms 1.391ms 80 aten::native_layer_norm 0.02% 2.122ms 0.29% 26.591ms 166.196us 102.329ms 1.11% 169.701ms 1.061ms 160 void at::native::(anonymous namespace)::vectorized_l... 0.00% 0.000us 0.00% 0.000us 0.000us 102.329ms 1.11% 102.329ms 639.554us 160 void at::native::unrolled_elementwise_kernel<at::nat... 0.00% 0.000us 0.00% 0.000us 0.000us 94.957ms 1.03% 94.957ms 212.909us 446 aten::add 0.02% 2.057ms 0.61% 56.490ms 198.909us 94.150ms 1.02% 94.150ms 331.515us 284 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 89.550ms 0.97% 89.550ms 362.552us 247 aten::add_ 0.03% 2.695ms 0.33% 30.630ms 125.532us 86.481ms 0.94% 86.481ms 354.429us 244 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 85.761ms 0.93% 85.761ms 708.765us 121 void at::native::vectorized_elementwise_kernel<4, at... 0.00% 0.000us 0.00% 0.000us 0.000us 81.610ms 0.88% 81.610ms 291.465us 280 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 67.547ms 0.73% 67.547ms 337.734us 200 void cutlass::Kernel2<cutlass_80_tensorop_bf16_s1681... 0.00% 0.000us 0.00% 0.000us 0.000us 59.451ms 0.64% 59.451ms 743.138us 80 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 44.693ms 0.48% 44.693ms 372.440us 120 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 41.748ms 0.45% 41.748ms 171.802us 243 aten::addmm 0.01% 1.171ms 0.11% 10.530ms 263.238us 38.640ms 0.42% 38.640ms 966.008us 40 void at::native::vectorized_elementwise_kernel<4, at... 0.00% 0.000us 0.00% 0.000us 0.000us 38.150ms 0.41% 38.150ms 476.874us 80 aten::pow 0.01% 1.359ms 0.05% 4.193ms 48.758us 35.091ms 0.38% 37.374ms 434.584us 86 void at::native::vectorized_elementwise_kernel<4, at... 0.00% 0.000us 0.00% 0.000us 0.000us 35.091ms 0.38% 35.091ms 422.780us 83 aten::_flash_attention_forward 0.01% 647.225us 0.12% 11.309ms 282.735us 33.271ms 0.36% 33.271ms 831.784us 40 ------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ Self CPU time total: 9.194s
Self CUDA time total: 9.232s

After:

device: NVIDIA L40S torch 2.5.1+cu124
peaks: 181 TFLOP/s bf16, 92 TFLOP/s fp32, 864 GB/s
cudnn.benchmark=False resolution=128 n_steps=6
torch._C._cuda_set_sync_debug_mode(debug_mode)
latent grid 32x32x1 = 1024 spatial tokens x 6 timesteps
processor cost 15.07 TFLOP per sample
--- 1x1x1 Conv3d vs equivalent Linear (temporal qkv projection) ---
float32 conv3d 1.00 ms ( 72.8 TFLOP/s) linear 1.93 ms ( 37.9 TFLOP/s) conv is 0.5x slower
bfloat16 conv3d 0.50 ms ( 145.5 TFLOP/s) linear 0.31 ms ( 232.4 TFLOP/s) conv is 1.6x slower
--- float32 (no autocast) ---
B cold ms warm ms ms/sample TFLOP/s % peak
1 495.2 501.3 501.3 30.1 32.8%
2 1143.2 1133.9 566.9 26.6 29.0%
4 2423.3 2419.3 604.8 24.9 27.2%
8 4890.9 4931.2 616.4 24.4 26.7%
stage breakdown at B=8 (ms per forward):
blocks 4883.80 (99%)
decoder[2] 17.61 (0%)
encoder[2] 17.01 (0%)
jitter 0.48 (0%)
(unaccounted) 12.27 (0%) <- host-side prep, hook copies, stack/permute
--- bfloat16 (autocast) ---
B cold ms warm ms ms/sample TFLOP/s % peak
1 270.8 225.8 225.8 66.7 36.9%
2 614.0 609.3 304.7 49.5 27.3%
4 1398.5 1397.2 349.3 43.1 23.8%
8 2815.4 2816.4 352.0 42.8 23.6%
stage breakdown at B=8 (ms per forward):
blocks 2781.14 (99%)
encoder[2] 18.27 (1%)
decoder[2] 14.73 (1%)
jitter 0.48 (0%)
(unaccounted) 1.77 (0%) <- host-side prep, hook copies, stack/permute
--- materialised memory traffic (one forward, views excluded) ---
286.5 GB written by 9605 aten calls at B=4, bfloat16
>= 663 ms of pure memory time at peak bandwidth
mul 72.84 GB
_to_copy 58.99 GB
native_layer_norm 22.65 GB
add 22.52 GB
clone 22.38 GB
addmm 22.15 GB
pow 11.51 GB
convolution 11.29 GB
cat 11.09 GB
stack 8.33 GB
silu 5.97 GB
bmm 5.91 GB
neg 4.15 GB
_scaled_dot_product_flash_attention 2.83 GB
mm 2.77 GB
--- implicit device syncs: 5 per forward ---
5x called a synchronizing CUDA operation (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:150.)
--- torch.profiler, top ops by self device time ---
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg # of Calls 
------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ aten::copy_ 0.63% 8.608ms 22.11% 302.405ms 183.609us 328.939ms 23.51% 328.939ms 199.720us 1647 void at::native::elementwise_kernel<128, 2, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 235.457ms 16.83% 235.457ms 430.452us 547 aten::mul 0.42% 5.744ms 7.75% 105.997ms 138.017us 225.991ms 16.15% 225.991ms 294.260us 768 void cutlass::Kernel2<cutlass_80_tensorop_bf16_s1681... 0.00% 0.000us 0.00% 0.000us 0.000us 176.992ms 12.65% 176.992ms 1.475ms 120 aten::addmm 0.18% 2.502ms 0.66% 9.004ms 112.555us 157.428ms 11.25% 157.428ms 1.968ms 80 aten::cat 0.27% 3.633ms 2.32% 31.686ms 156.089us 148.539ms 10.61% 148.539ms 731.721us 203 aten::cudnn_convolution 0.23% 3.208ms 3.64% 49.733ms 606.501us 128.237ms 9.16% 128.237ms 1.564ms 82 sm86_xmma_fprop_implicit_gemm_bf16bf16_bf16f32_f32_n... 0.00% 0.000us 0.00% 0.000us 0.000us 107.065ms 7.65% 107.065ms 1.338ms 80 void at::native::unrolled_elementwise_kernel<at::nat... 0.00% 0.000us 0.00% 0.000us 0.000us 94.297ms 6.74% 94.297ms 211.429us 446 aten::add 0.15% 2.018ms 2.96% 40.440ms 142.395us 94.103ms 6.72% 94.103ms 331.350us 284 aten::native_layer_norm 0.16% 2.140ms 14.84% 203.031ms 1.269ms 89.507ms 6.40% 156.795ms 979.969us 160 void at::native::(anonymous namespace)::vectorized_l... 0.00% 0.000us 0.00% 0.000us 0.000us 89.507ms 6.40% 89.507ms 559.421us 160 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 89.182ms 6.37% 89.182ms 361.062us 247 void at::native::vectorized_elementwise_kernel<4, at... 0.00% 0.000us 0.00% 0.000us 0.000us 81.564ms 5.83% 81.564ms 291.300us 280 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 66.241ms 4.73% 66.241ms 331.207us 200 aten::bmm 0.16% 2.136ms 3.84% 52.532ms 656.656us 44.800ms 3.20% 44.800ms 560.001us 80 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 44.669ms 3.19% 44.669ms 372.241us 120 void at::native::vectorized_elementwise_kernel<4, at... 0.00% 0.000us 0.00% 0.000us 0.000us 38.128ms 2.72% 38.128ms 476.601us 80 aten::pow 0.09% 1.174ms 2.06% 28.172ms 327.585us 35.076ms 2.51% 37.359ms 434.409us 86 void at::native::vectorized_elementwise_kernel<4, at... 0.00% 0.000us 0.00% 0.000us 0.000us 35.076ms 2.51% 35.076ms 422.598us 83 aten::_flash_attention_forward 0.05% 627.150us 0.10% 1.396ms 34.892us 31.928ms 2.28% 31.928ms 798.204us 40 void pytorch_flash::flash_fwd_kernel<pytorch_flash::... 0.00% 0.000us 0.00% 0.000us 0.000us 31.928ms 2.28% 31.928ms 798.204us 40 Memcpy DtoD (Device -> Device) 0.00% 0.000us 0.00% 0.000us 0.000us 31.830ms 2.27% 31.830ms 306.056us 104 aten::add_ 0.19% 2.541ms 1.40% 19.120ms 93.724us 29.015ms 2.07% 29.015ms 142.233us 204 void at::native::elementwise_kernel<128, 4, at::nati... 0.00% 0.000us 0.00% 0.000us 0.000us 28.338ms 2.03% 28.338ms 349.850us 81 ------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ Self CPU time total: 1.368s
Self CUDA time total: 1.399s

[Michal-Novomestsky]

Please do test this yourselves before merging! Perhaps this is some local quirk of my cluster.