Add non-TMA MXFP8 cast-only kernels for specialized rowwise-only and row+colwise - #3459
Add non-TMA MXFP8 cast-only kernels for specialized rowwise-only and row+colwise #3459tdophung wants to merge 6 commits into
Conversation
sylvesterkaczmarek
left a comment
sylvesterkaczmarek
left a comment
There was a problem hiding this comment.
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.
| e8m0_t *scale_out = reinterpret_cast<e8m0_t *>(scales); | ||
|
|
||
| // A padded scale array breaks the flat block view the fast path relies on. | ||
| if (scale_stride != blocks_per_row) { |
There was a problem hiding this comment.
I think currently there is no normal way from TE to reach this place because in transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh, line 731-741:
const bool is_full_rowwise_chunk = (cols % 128 == 0);
...
const bool scaling_type_has_specialized_support =
(scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk &&
rowwise_specialized_grid_fits) ||
...
So we only use specialized kernels when cols % 128 == 0.
Here this if says scale_stride != cols / 32, meaning there is padding, which is never the case for cols % 128 == 0?
| } | ||
|
|
||
| const int64_t blocks_done = grid * blocks_per_cta; | ||
| if (blocks_done < num_blocks) { |
There was a problem hiding this comment.
Another option is to handle the remainder directly in quantize_contiguous_kernel by masking to skip MX blocks whose indices are beyond num_blocks.
You can template the continuous kernel with CHECK_BOUNDS and gate boundary check under this constexpr, and decide whether there is a remainder to handle in your host code so we don't pay the masking cost when we don't have to.
The tradeoff is then:
- Checked contiguous kernel: one launch, but additional bounds checks and predication.
- Separate remainder kernel: check-free contiguous kernel, but an additional kernel launch.
I can ask my agent to run a benchmark later if you haven't tried this
There was a problem hiding this comment.
I vide coded the masked version: https://github.com/kainzhong/TransformerEngine/tree/mxfp8-register-cast
The benchmark script shows:
┌───────────────────────┬──────────┬──────────────────────┬──────────┬──────────┐
│ shape │ leftover │ baseline (2 kernels) │ current │ Δ │
├───────────────────────┼──────────┼──────────────────────┼──────────┼──────────┤
│ 8160×1920 │ 128 │ 10.85 µs │ 9.06 µs │ −16.5% │
├───────────────────────┼──────────┼──────────────────────┼──────────┼──────────┤
│ 4064×4480 │ 128 │ 11.89 µs │ 10.11 µs │ −15.0% │
├───────────────────────┼──────────┼──────────────────────┼──────────┼──────────┤
│ 65504×1920 │ 128 │ 57.44 µs │ 55.36 µs │ −3.6% │
├───────────────────────┼──────────┼──────────────────────┼──────────┼──────────┤
│ 32736×4480 │ 128 │ 66.06 µs │ 63.99 µs │ −3.1% │
├───────────────────────┼──────────┼──────────────────────┼──────────┼──────────┤
│ 8 shapes, no leftover │ 0 │ — │ — │ ≤0.02 µs │
└───────────────────────┴──────────┴──────────────────────┴──────────┴──────────┘
Looks like the masking version is a bit faster?
There was a problem hiding this comment.
Those numbers make the masked single-kernel path look preferable: roughly 3-16% better on leftover shapes with effectively no regression on the no-leftover cases. My remaining concern is correctness coverage around boundary/alignment, padded stride, and both FP8 formats; once that variant lands here I am happy to recheck.
| constexpr int32_t kNarrowMinBlocksPerSm = 6; | ||
|
|
||
| // Thread-block cluster widths. Clustering lets neighbouring CTAs, which read | ||
| // adjacent columns of the same rows, share L2 traffic. |
There was a problem hiding this comment.
Hum but L2 cache is already shared by all CTAs? How does clustering help L2 traffic (I'm not very familiar with CTA cluster so I don't really know how this works)?
…XFP8 The narrow tile was selected once the grid got deep, but measurement shows that switch was a straight loss: it halves the work per thread while the per-thread fixed cost stays, costing ~10% on large shapes. The wide tile now applies whenever the column count divides it; narrow remains the fallback for counts it does not divide. The software L2 prefetch was compensating for the narrow tile. With the tile fixed it never helps at any shape measured up to 65536x8192, and costs ~3.4% from 8192x8192 upward, so it is removed along with its plumbing. Bidim vs upstream main, ncu on cold memory, 12 shapes: geomean 0.957 -> 0.864, and the +5.8% regressions at K=7168 are gone.
Shapes whose MX block count does not divide evenly among CTAs needed a second launch to mop up the leftovers. Instead the grid is rounded up and the kernel predicates its accesses under a CHECK_BOUNDS template parameter, so shapes that divide evenly instantiate it false and generate exactly the code they did before. Measured rowwise, wall clock, rotating shapes, median of 25: 8160x1920 10.66 -> 8.96 us (-15.9%) 4064x4480 11.10 -> 9.47 us (-14.7%) 65504x1920 39.94 -> 38.40 us (-3.9%) 32736x4480 45.60 -> 44.13 us (-3.2%) four shapes that divide evenly: within 0.1-0.7%, i.e. unchanged Output and scales are bit-identical to both the two-kernel path and to the upstream kernel on every shape checked. Approach and implementation by Kaining Zhong (github.com/kainzhong), from review of NVIDIA#3459.
The kernel can now emit scales in the GEMM-swizzled layout, bit-exactly against the staged kernel across shapes covering power-of-two, add-correction and small divisors. Recovering a row index from the flat block index needs a division, and a 64-bit one more than doubled the instruction count (30.9M -> 64.0M), so the host passes a magic reciprocal instead (46.7M). It is not yet profitable, so dispatch still sends the swizzled path to the staged kernel and the per-path gate records why. The swizzled layout packs a 512-byte tile as 128 rows x 4 scale columns; this kernel walks a flat row-major block sequence, so a warp's scales land four bytes to a sector. Measured 7-12% slower, with L2 read sectors identical and write sectors 8.7% higher -- the whole gap. Packing the four contiguous columns via shuffle does not recover it (the next contiguous run comes from rows 32 apart), so closing this needs a two-dimensional traversal rather than a tweak. Also correct the thread-block cluster comment: clustering does not reduce L2 traffic. Ablating it leaves DRAM bytes and L2 hit rate unchanged; it is worth 3.5-4.5% because co-scheduling on one GPC lifts DRAM row-buffer locality.
Shapes whose MX block count does not divide evenly among CTAs needed a second launch to mop up the leftovers. Instead the grid is rounded up and the kernel predicates its accesses under a CHECK_BOUNDS template parameter, so shapes that divide evenly instantiate it false and generate exactly the code they did before. Measured rowwise, wall clock, rotating shapes, median of 25: 8160x1920 10.66 -> 8.96 us (-15.9%) 4064x4480 11.10 -> 9.47 us (-14.7%) 65504x1920 39.94 -> 38.40 us (-3.9%) 32736x4480 45.60 -> 44.13 us (-3.2%) four shapes that divide evenly: within 0.1-0.7%, i.e. unchanged Output and scales are bit-identical to both the two-kernel path and to the upstream kernel on every shape checked. Approach and implementation by Kaining Zhong (github.com/kainzhong), from review of NVIDIA#3459. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
f19950f to
98d18c7
Compare
The packed 32-bit scale store assumes a 4-aligned group of MX blocks never straddles a row, which holds only when blocks_per_row is a multiple of 4. Dispatch guarantees that via cols % 128 == 0, but the kernel documented it in a comment and enforced nothing, so a caller reaching the launcher directly with, say, cols = 96 would have written a byte belonging to the next row and corrupted scales silently. Also scope the 32-bit block-count limit to the swizzled path, which is the only one that needs it; it was constraining the packed path as well.
Same hook and pinned version the repository's pre-commit uses (mirrors-clang-format v18.1.6), so this subsumes the auto-fix commit pre-commit.ci pushed on the previous tip.
a905df8 to
fe7bfde
Compare
3 participants
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: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
Changes
See in description
Checklist: