Add non-TMA MXFP8 cast-only kernels for specialized rowwise-only and row+colwise

[tdophung]

Description

The specialized cast-only path currently stages its tile: the rowwise kernel through shared memory, the bidimensional one through a TMA pipeline. For a cast with no bias or activation there is nothing to stage for, the tile is read once and consumed immediately, so both can be done entirely in registers instead.

Adds two kernels under cast/mxfp8/specialized/ (NOTE: both of these has been adapted from the winning kernel in Kernel Factory to work in TE, after some suggestions to KF on where I think the bottleneck is)

cast_rowwise.cu

Two lanes cooperate on each 32-element scaling block, each lane's half being exactly one 256-bit load. No shared memory and no barrier at all; the tensor is a flat sequence of independent scaling blocks whenever the scale array is packed.

cast_bidim.cu

A CTA owns a 32-row band, which is exactly the colwise block height, so the colwise reduction closes inside the CTA and the tile drives both passes from registers. Shared memory is used only for the cross-warp column fold.

Both pick their launch configuration from a documented size-tier table and use L2 eviction policies plus a software prefetch one resident CTA-wave ahead.

New PTX wrappers in util/ptx.cuh:

• L2 cache-policy creation
• 256/128-bit non-coherent loads
• 128/64/8-bit stores carrying a policy
• L2 prefetch
• packed BF16 reciprocal-scale helper
• mul_cvt_4x overload taking two BF16 pairs with independent scales, which the colwise pass needs since every column pair has its own scale.

Dispatch routes to these kernels for BF16 input with a NON-SWIZZLE scale layout only, and falls through to the existing kernels otherwise. FP16 input and GEMM-swizzled scale layouts are not yet covered.

Measured on GR10x (CC 10.7, CUDA 13.4), medians of 20 with the clocks warmed:

rowwise, geomean over 12 small shapes (up to M_size = 4k) 10.949 us -> 9.913 us
rowwise, geomean over 6 standard shapes (M_size = 16k - 64k) 19.411 us -> 17.934 us
bidimensional, geomean over 6 shapes (M_size 8k - 64k) 29.595 us -> 25.730 us

The margin narrows as the shapes grow and the kernels become DRAM-bound; at 65536x16384 rowwise reaches 89% of peak DRAM throughput based on NCU.

Registers drop from 48 to 24-30 (rowwise) and from 70 to 46-64 (bidimensional), with occupancy rising correspondingly.

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

See in description

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

[sylvesterkaczmarek]

I think this needs correctness regression coverage before the new kernels become the default dispatch for these shapes.

Could we compare the new path against the existing MXFP8 implementation for at least rowwise and bidimensional scaling, including boundary/alignment cases and a padded rowwise scale stride?

Since the change also introduces new scale arithmetic and PTX helpers, coverage for zero, Inf/NaN and both supported FP8 output formats would be particularly useful. Performance measurements look good, but currently a numerical regression would have no dedicated test catching it.