Make fused_ff_qkv contiguous - #12
Open
Michal-Novomestsky wants to merge 1 commit into
Open
Conversation
Author
|
For more info, I quickly vibecoded some logging to compare the difference. Note that one 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.232sAfter: 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
marked this pull request as ready for review
August 22, 2026 07:41
Author
|
Please do test this yourselves before merging! Perhaps this is some local quirk of my cluster. |
Author
|
@payelmuk150 Perhaps this could be of interest to you? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Currently, walrus is painfully slow to run (at least on my cluster - L40S is taking ~6-7s per forward pass for a 4x32x32x6 2D field with 6 timesteps and batch size 4).
After some sleuthing, it appears to stem from a single line of code:
This breaks the contiguity of$BHWDC$
fused_ff_qkvwhich causes ATen to dispatch to[1, fused_dims]seperate GEMVs instead of a single[B*H*W*D*C, fused_dims]GEMM. In my case, this amounts to4*32*32*1*6 = 24576individual GEMV calls. Changing this gives me speedups of ~6.5x.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.trainingcan't fix.