[vthumbe1503]

Description

This PR introduced rowwise/colwise swizzling kernels targeted at narrow scale width use-cases. The scope of this PR is to

• Colwise Swizzling: Enable the colwise swizzling optimized kernel(swizzle_columnwise_chunk_kernel) for both narrow scale-width and wide scale-width use-cases. And completely removes the previously existing kernel for narrow scale width kernel. The previously existing kernel for wide scale-width is retained since that is still used for grouped swizzling. Using the kernel in grouped swizzling will be a follow up PR.
• Rowwise swizzling: Enable rowwise swizzling optimized kernel(swizzle_rowwise_full_width_kernel) for narrow scale-width use-cases. For wide scale-width, we fall back to our previosuly existing rowwise scaling kernel which operates on 128128N tile where N is vectorization's bit width. Fall back is implemented through shared memory heuristics
  • If scale is sufficiently narrow (smem ~= 128*padded_k <= 17KB), new kernel is better
  • If smem is between (17,30)KB then using the new kernel turns out to be better if the number of row tiles is between 16 and 256.
• Optimized row scaling kernel is also extended to be use for NVFP4.
  • Apt unit tests are also added

Blackwell Perf results

MXFP8 rowwise Shape Before (µs) After (µs) Speedup 4096×1536 13.576 1.729 7.852× 4096×512 4.413 1.247 3.538× 4096×128 2.111 1.183 1.784× 4096×2048 2.049 1.505 1.362× 4096×4096 2.551 2.130 1.198× 2048×7168 3.033 2.886 1.051× 4096×7168 3.519 3.517 1.001×

MXFP8 columnwise Shape Before (µs) After (µs) Speedup 4096×1536 13.736 1.601 8.580× 4096×512 5.548 1.600 3.467× 4096×2048 5.114 1.631 3.135× 4224×2048 4.990 1.631 3.059× 4096×7168 5.370 2.427 2.213× 7168×4096 5.288 2.431 2.175× 2048×7168 3.422 1.665 2.055×

MXFP8 multi-tensor rowwise Shape N Before (µs) After (µs) Speedup 7168×4096 8 8.388 4.428 1.894× 4096×2048 8 4.846 2.589 1.872× 4096×128 4 2.497 1.471 1.697× 16384×128 4 2.465 1.473 1.673× 256×256 4 2.079 1.377 1.510× 2048×7168 15 7.213 6.162 1.171× 4096×7168 8 7.107 6.085 1.168×

MXFP8 multi-tensor columnwise Shape N Before (µs) After (µs) Speedup 4096×2048 8 5.989 3.159 1.896× 4224×2048 20 7.721 4.839 1.596× 1024×7168 32 10.777 7.419 1.453× 4224×2048 13 6.167 4.520 1.364× 4096×7168 8 6.547 5.190 1.262× 7168×4096 8 6.561 5.804 1.130×

Rubin perf results

MXFP8 rowwise Shape Before (µs) After (µs) Speedup 4096×1536 11.433 1.615 7.077× 4096×512 3.340 1.181 2.827× 4096×128 1.605 1.156 1.388× 4096×2048 1.884 1.369 1.376× 4096×4096 2.310 2.037 1.134× 2048×7168 2.681 2.482 1.080× 4096×7168 3.148 3.148 1.000×

MXFP8 columnwise Shape Before (µs) After (µs) Speedup 4096×1536 10.314 1.495 6.900× 4096×2048 4.443 1.516 2.931× 4096×512 4.246 1.482 2.864× 4224×2048 4.355 1.516 2.873× 4096×7168 4.818 2.239 2.152× 7168×4096 4.700 2.233 2.105× 2048×7168 3.046 1.581 1.926×

MXFP8 multi-tensor rowwise Shape N Before (µs) After (µs) Speedup 7168×4096 8 7.063 3.773 1.872× 4096×2048 8 4.302 2.411 1.784× 4096×128 4 2.170 1.371 1.583× 16384×128 4 2.144 1.434 1.495× 256×256 4 1.769 1.360 1.300× 4096×7168 8 6.174 5.529 1.117× 2048×7168 15 6.201 5.559 1.115×

MXFP8 multi-tensor columnwise Shape N Before (µs) After (µs) Speedup 4096×2048 8 5.525 2.731 2.023× 4224×2048 20 7.835 4.609 1.700× 4224×2048 13 5.349 3.572 1.498× 4096×7168 8 5.704 4.294 1.328× 1024×7168 32 8.724 6.689 1.304× 7168×4096 8 5.665 4.837 1.171×

One note on the perf of the rowwise scaling kernel 4096 * 1536 has much better speed up compared to 4096*2048: This is because narrow width kernel is chosen up until about 1536 columns. After that limit the generic kernel is used. Generic kernel parallelizes decently well but still not better than the newly introduced kernel at 2048 width. Hence 2048 has a lesser speedup relative to baseline Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[vthumbe1503]

/te-ci

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no outstanding actionable findings remain.

Summary

The PR enables optimized scale-swizzling kernels and removes obsolete narrow-width kernels.

• Scale swizzling now chooses optimized kernels from the workload.
• Batched swizzling now uses optimized paths for MXFP8 and NVFP4.
• Swizzle calls now require 16-byte-aligned scale buffers.

No new actionable issues were found. CUDA tests were not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Check tensors and alignment] --> B{Scale layout}
  B -->|Columnwise MXFP8| C[Columnwise chunk kernel]
  B -->|Rowwise or NVFP4| D{Width and tile count}
  D -->|Fits heuristic| E[Rowwise full-width kernel]
  D -->|Otherwise| F[Generic rowwise kernel]

Loading

_{Reviews (6) · Last reviewed commit: "Optimize scale swizzle kernels" · Reviewed by Greptile}

[ptrendx]

@vthumbe1503 Please sign every commit in this PR. See https://github.com/NVIDIA/TransformerEngine/blob/main/CONTRIBUTING.rst#sign-your-work

[vthumbe1503]

/te-ci L0