[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
  • For roughly 17–30 KiB of shared memory, use the full-width kernel for at most 16 row tiles or when the total row tiles exceed the current GPU’s SM count. The SM-count cutoff replaces the previous hard-coded 128-tile threshold. Otherwise use the generic rowwise kernel.
• Optimized row scaling kernel is also extended to be use for NVFP4.
  • Apt unit tests are also added
• 16 byte Alignment requirement for Swizzling kernels
  • Previously existing kernels had the minimum 4 byte alignment requirement. But even that was never checked
  • Now we have kernels for most of the use-cases doing a int4 read/write and hence 16 byte alignment is a must for both input and output scale pointers.

Blackwell Perf results

MXFP8 rowwise Shape Before (µs) After (µs) Speedup 4096×128 2.111 1.183 1.784× 4096×256 2.818 1.289 2.186× 4096×512 4.413 1.247 3.538× 4096×768 6.014 1.537 3.912× 4096×1024 8.204 1.664 4.932× 4096×1280 10.278 1.793 5.731× 4096×1536 13.576 1.729 7.852× 4096×1664 22.635 1.811 12.502× 4096×1792 25.265 1.806 13.990× 4096×1920 2.083 1.826 1.141× 4096×2048 2.049 1.505 1.362× 4096×2560 2.303 1.913 1.204× 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×

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

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

[High impact] The PR appears safe to merge; no new actionable issue was found.

Summary

This PR makes the optimized swizzle kernels the default and removes the obsolete narrow kernels.

• Scale swizzling now picks optimized rowwise and columnwise kernels automatically.
• Swizzle calls now reject scale buffers that are not 16-byte aligned.

The earlier threads are unnumbered, so they have no entries in previousFindings. Their reported problems are addressed by the current size checks, padded NVFP4 tests, selected-buffer dimensions, and explicit alignment requirement.

_{Reviews (13) · 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

[vthumbe1503]

/te-ci L0

[Oleg-Goncharov]

LGTM overall, just a few minor fixes

[vthumbe1503]

/te-ci L0

[Oleg-Goncharov]

LGTM

[vthumbe1503]

/te-ci L0

[Oleg-Goncharov]

LGTM

[vthumbe1503]

Pipeline: 77344495