Repository navigation
Conversation
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
With NVTE_ACTIVATION_TABLE_EXACT_DSILU=1, the dSiLU table covers |x| >= 2^-25 instead of 2^-9 (32 KB instead of 16 KB). Below 2^-25 the linear term already equals util/math.h, so dSiLU read from the table matches util/math.h for every BF16 input. The variable is read once, when the table is built; kernels tell the two layouts apart by the table's size. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…tables Port the register-tiled MXFP8 cast kernel from NVIDIA#3617 onto the shared BF16 activation tables. The kernel copies the table into shared memory with a bulk copy, evaluates the forward activations two lanes at a time with the branch-free lookup, and evaluates the derivatives from the same tables as the generic kernel, so both kernels produce bit-identical output. The kernel runs only when a table is passed. Pad table sizes to 16 bytes for the bulk copy, and add the packed bf16x2/fp32x2 PTX helpers it needs. Tests compare the register-tiled kernel against the generic kernel with a table, byte for byte. Signed-off-by: Teddy Do <tdophung@nvidia.com>
Instead of having each framework allocate the tables and pass them through new
_v2 entry points, TE common builds each table the first time it is needed and
keeps it, once per device and process, as it does for library handles. The
existing nvte_gelu, nvte_silu, nvte_dgelu, nvte_dsilu and
nvte_quantize_dbias_d{gelu,silu} read the tables for BF16 inputs quantized to
MXFP8, with no API change.
The first use may happen while a CUDA graph is being captured. The build then
runs in relaxed capture mode on a private stream and completes before the
captured kernel is recorded, so the replayed graph reads a complete table.
The tables are on by default. NVTE_ACTIVATION_TABLE=0 turns them off, and
NVTE_MXFP8_REGTILE=0 keeps table requests on the generic kernel. When the
CuTeDSL backend is enabled, it runs first and computes the activations itself.
Remove the _v2 entry points, the activation-table quantization config
attribute, the public table build and size functions, and the PyTorch and JAX
table caches.
Signed-off-by: Teddy Do <tdophung@nvidia.com>
…neric kernel Turn the tables off by default: NVTE_ACTIVATION_TABLE=1 enables them, as a new feature whose results differ from today's at -inf (-0 instead of NaN) and for tiny derivative inputs. The generic MXFP8 kernel no longer reads the table from global memory. A lookup there costs it more than the math, so it computes what the table holds from util/math.h instead: the formula with -inf clamped to give -0, and the linear term below the derivative window. Both are branch-free, so the compiler keeps overlapping the math across elements. The output stays bit-identical to the register-tiled kernel, which keeps reading the table from shared memory. Against TE main with the tables on, shapes and modes the register-tiled kernel does not take (dbias, rows not a multiple of 64) go from 0.75-0.96x to 0.94-1.00x. Tests check the computed form against the table specification for every BF16 input. Signed-off-by: Teddy Do <tdophung@nvidia.com>
… column tails The register-tiled kernel took only rows % 64 == 0 and cols % 256 == 0; every other shape, such as MoE token counts, fell back to the generic kernel. It now takes any number of rows and any cols % 16 == 0, the shapes the generic kernel takes. Out-of-range elements load as 0, and the kernel writes the scales the generic kernel writes for its whole 64x64 tiles, so the output, scale padding included, stays byte-identical. Bands inside the tensor run without bounds checks; only bands at its edges pay for them. The derivatives use this kernel for every shape: with the band body behind the edge branch, the compiler schedules the loads better (GB200, 32768x7168 dGeLU 239 -> 220 us). GB200, tables on, normal data, against TE main: 993x7200 1.40-1.86x (was 0.95-1.00x through the generic kernel), 3517x7168 1.97-2.24x, regular shapes unchanged for the forward and 6-7% faster for the derivatives. The test compares the two kernels on tail shapes with every output byte pre-filled, so both must write the same set of bytes. Signed-off-by: Teddy Do <tdophung@nvidia.com>
With NVTE_ACTIVATION_TABLE=0, the default, GeLU, SiLU and their derivatives now also take the register-tiled kernel: it computes the activation with util/math.h instead of reading a table, so its output is byte-identical to the generic kernel's, and the default path gets the register tiling without any change in results. The derivatives take it at every size. The forward activations are compute-bound without a table, and on small tensors too few CTAs hide the latency, so they take it from 2^23 elements on. GB200, BF16 -> MXFP8 E4M3, against TE main: dGeLU 1.81-1.92x, dSiLU 1.97-2.10x, GeLU 1.62-1.76x, SiLU 1.62-1.67x on 8192x7168 to 16384x16384; 993x7200 dGeLU/dSiLU 1.22x/1.48x. The regtile-vs-generic tests run every case with and without a table. Signed-off-by: Teddy Do <tdophung@nvidia.com>
nvte_quantize_dbias_dgelu and nvte_quantize_dbias_dsilu took the generic kernel even with the activation tables on. The register-tiled kernel now handles them and writes the same dbias workspace, which common::reduce_dbias then reduces as before. The generic kernel sums each column of a 64-row chunk strictly in row order, in FP32, before rounding the products to BF16. To give the same bits, a CTA here takes the chunk's two 32-row bands back to back, stages each band's FP32 products in 32 KB of dynamic shared memory, and one thread per column adds them in row order, carrying the sum across the two bands. The staging buffer limits the kernel to four CTAs per SM, so its launch bounds allow 64 registers. GB200, tables on, against TE main: 1.29-1.75x on normal data and up to 2.47x on wide data (32768x7168 dbias_dgelu: 638 -> 391 us), where the generic kernel ran at 0.94-0.99x. The regtile-vs-generic test adds both dbias modes and compares the FP32 workspace and the BF16 dbias byte for byte. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…P8 kernel The register-tiled kernel refused requests with an amax output, a noop tensor or E5M2 output, which then took the generic kernel. - E5M2: the kernel is templated on the FP8 output type; the scale uses that type's max norm and the packed conversion its E5M2 form. - amax: the tensor amax is the max of the column amaxes the kernel already computes, reduced per CTA and combined with one atomic max, the value the generic kernel reduces. It is a separate instantiation, so the kernels without amax keep their register budget. - noop: the generic kernel ignores the noop tensor for every request that computes an activation, and so does this kernel now. GB200, tables on: E4M3 without amax is unchanged; E5M2 runs at the E4M3 speed (2.3-2.8x TE main at 32768x7168); with amax, the forward costs about 1% and the derivatives about 5%. The register-tiled versus generic test gains the output type and amax as dimensions; amax must match bit for bit. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…U and dbias FP16 inputs of dSiLU, dbias_dGeLU and dbias_dSiLU now take the register-tiled kernel. FP16 has no table (the tables hold BF16 inputs), so the kernel computes the activation with util/math.h, and it scales every value in FP32 with mul_cvt_2x, as the generic kernel does, since the reciprocal scale is often outside FP16's range. The input-type-specific parts of the datapath are in Datapath<IType>; the BF16 path is unchanged. The FP32 scaling costs about twice the instructions of the BF16 path, which the compute-bound forward kernels and plain dGeLU cannot absorb (GB200: GeLU 0.74-0.91x, SiLU 0.86-1.26x, dGeLU 0.95-1.15x of TE main), so those keep the generic kernel. GB200 against TE main: dSiLU 1.26-1.74x, dbias_dGeLU 1.16-1.43x, dbias_dSiLU 1.29-1.51x. A new FP16 regtile-vs-generic suite covers every mode, output type, amax setting and tail shape. Signed-off-by: Teddy Do <tdophung@nvidia.com>
common::reduce_dbias sums the FP32 dbias workspace column by column. It ran 256 threads per CTA with four columns each, so a 7168-column workspace got 7 CTAs, and each thread waited for one load per row. Use one thread per column, 64 per CTA, and copy each column into shared memory with cp.async, 32 rows per stage and four stages deep, adding a stage once it has landed. Each column is still added in row order from 0 and rounded once, so the result is unchanged; a new test checks it bit for bit against that sum on the host for FP32, BF16 and FP16 output, on data spanning 2^-30..2^30 with mixed signs. GB200, reduction of the 32768x7168 dbias_dgelu workspace (512x7168): 108.8 us -> 13.2 us. End to end against TE main: dbias_dgelu / dbias_dsilu / dbias 1.07-1.25x without tables and 1.00-1.31x with them across 8192x7168 to 993x7200. Signed-off-by: Teddy Do <tdophung@nvidia.com>
nvte_swiglu and nvte_dswiglu quantized to MXFP8 from BF16 to E4M3 in both directions now take a register-tiled kernel, like the one for GeLU and SiLU: a CTA owns 32-row bands of 256 output columns, keeps them in registers, and computes both scale directions from there. The backward quantizes the two gradient halves one after the other from the same band. The math, the rounding to BF16 and the scale writes (inside the tensor only) are those of the generic gated kernel, so the output is byte-identical. It takes any number of rows and K % 32 == 0, the generic kernel's requirement. NVTE_MXFP8_REGTILE=0 keeps the generic kernel. GB200, against TE main, normal data: SwiGLU 1.39-1.99x (5.6 TB/s on large shapes), dSwiGLU 1.39-1.59x; wide data up to 2.05x and 2.18x. The new test compares the two kernels byte for byte on regular and tail shapes with every output byte pre-filled. Signed-off-by: Teddy Do <tdophung@nvidia.com>
FP32 inputs of GeLU, SiLU, dGeLU and dSiLU, with or without dbias, now take the register-tiled kernel. A tile column pair is two FP32 registers instead of one 32-bit word, so an FP32 CTA has twice the threads (256 for the activations, 512 for the derivatives) and each thread holds half the rows: per-thread registers and threads per SM stay as for 16-bit inputs. A load group is one row, two 16-byte vectors per lane. FP32 inputs have no table; the activation is computed with util/math.h, and every value is scaled in FP32 and converted with mul_cvt_2x, as the generic kernel does, so the output stays byte-identical. GB200 against TE main, normal data: GeLU 1.12-1.35x, SiLU 1.24-1.35x, dGeLU 1.29-1.48x, dSiLU 1.48-1.92x, dbias_dGeLU 1.16-1.32x, dbias_dSiLU 1.35-1.61x (wide data up to 2.37x); forward tensors under 2^23 elements keep the generic kernel as for BF16. A new FP32 regtile-vs-generic suite covers every mode, output type, amax setting and tail shape. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…ster-tiled kernel The register-tiled dSwiGLU kernel quantized the gradient of the activated half and the gradient of the gate half in two passes over each 32-row band, the second re-reading the activation and the incoming gradient and recomputing the sigmoid. It now reads the activation, gate and incoming gradient once, computes both halves' tiles together, and quantizes them one after the other. The math per element is unchanged, so the output is byte-identical to the generic gated kernel as before. GB200, BF16 -> MXFP8 E4M3, against TE main: dSwiGLU 1.74-2.16x on normal data (was 1.39-1.59x) and 1.87-2.76x on wide data; 32768x7168 goes from 550 to 405 us. The test adds an input whose activated half holds every BF16 value. Signed-off-by: Teddy Do <tdophung@nvidia.com>
The register-tiled SwiGLU / dSwiGLU kernel wrote E4M3 only; requests for E5M2 took the generic gated kernel. The kernel is now templated on the FP8 output type, using the output-type-aware scale and convert helpers of the non-gated register-tiled kernel, and the gated dispatch admits E5M2. The output stays byte-identical to the generic gated kernel. GB200, BF16 -> MXFP8 E5M2, against TE main: SwiGLU 1.40-2.00x on normal data and 1.40-2.05x on wide data, dSwiGLU 1.74-2.13x and 1.89-2.75x, the same as E4M3, which is unchanged. The regtile-vs-generic test runs every gated case with both output types. Signed-off-by: Teddy Do <tdophung@nvidia.com>
With NVTE_BUILD_ACTIVATION_WITH_FAST_MATH=ON the activation kernels compile with --use_fast_math, but activation_table.cu did not, so the tables held the accurate results while the kernels computed the fast-math ones. Compile it with the same flags. Fast-math builds remain approximate for GeLU. On GB200, with every BF16 input: SiLU, dSiLU, dGeLU, dbias_dGeLU, SwiGLU and the FP16/FP32 paths match between the register-tiled and generic kernels byte for byte, but GeLU does not, with or without tables. Without tables the register-tiled output matches the generic kernel's rowwise data and differs from its colwise data in 0.02% of bytes, so the generic kernel already disagrees with itself between its two passes: the compiler builds the fast-math GeLU differently in different code. With tables 0.8% of bytes differ. Default builds are unaffected. Signed-off-by: Teddy Do <tdophung@nvidia.com>
NVTE_MXFP8_REGTILE was a development switch for comparing the register-tiled kernels against the generic ones. Remove it with the tests that used it: - The regtile-vs-generic suites in test_cast_mxfp8_act_exact.cu. The exact test already checks every kernel against the activation computed per element and then quantized; it now also covers E5M2 output, the tensor amax, the forward without tables from 2^23 elements (4127x2064, with row and column tails) and a dbias workspace of 130 rows (8257x544). - test_cast_mxfp8_gated_regtile.cu: CastMXFP8_GatedActTestSuite covers the same requests (BF16/FP16/FP32, E4M3/E5M2, forward and backward, tails). - test_reduce_dbias.cu: the exact test checks dbias bit for bit against a host sum in the kernels' order, now with a workspace past the reduction's 128-row pipeline. Signed-off-by: Teddy Do <tdophung@nvidia.com>
Table<IType, Kind> and ExactDsiluTable<IType>, with BF16 specializations only. The tables hold the activation, not its quantization, so the same tables serve any output format; FP16 tables would be further specializations. test_activation_table.cu describes each table by a per-kind Spec (activation, derivative window and slope), so the lookup and check helpers no longer carry dSiLU-specific parameters. Shorten the CMake comment on building activation_table.cu with fast math. Signed-off-by: Teddy Do <tdophung@nvidia.com>
With NVTE_ACTIVATION_TABLE=1 the CuTeDSL backend now follows the table
specification for BF16 GeLU, SiLU, dGeLU and dSiLU, as the CUDA kernels do:
- The dispatcher passes the table and its layout ("default", or "exact" for
the exact dSiLU table) to the CuTeDSL entry point. The layout is part of the
kernel config; kernel names change only when a table is used. The table
pointer is an always-present argument, null without a table, like the noop
flag.
- The register-resident bidimensional kernel copies the table into shared
memory and reads it with packed BF16 lookups ported from
util/activation_table.cuh. With a table it walks several row tiles per CTA
so the table is copied once per CTA, uses 256-column tiles for the
derivatives, and indexes tiles of a run-time row stride from a 64-bit base
with 32-bit offsets. Without a table it is unchanged.
- The other kernels compute what the table holds: -inf gives -0 in the
forward, and the derivatives are the linear term below the table window.
Every BF16 case of the exact-composition test that reaches the register-
resident kernel matches with tables on and off; the remaining mismatches of
the CuTeDSL backend are in its general kernel and do not depend on the tables.
The test gains a 4096x2048 shape, on which the register-resident kernels walk.
GB200, BF16, geometric mean speedup over TE main with tables: GeLU 2.12,
SiLU 2.38, dGeLU 2.06, dSiLU 2.02 (CUDA backend: 2.09, 2.34, 2.14, 2.10).
Signed-off-by: Teddy Do <tdophung@nvidia.com>
dGeLU and dSiLU with dbias, BF16 input, both directions or colwise only and row counts divisible by 64 now take the register-resident bidimensional kernel instead of the general one, with or without activation tables. dbias follows the CUDA kernels' order: each CTA takes the two 32-row tiles of a 64-row chunk back to back, stages the FP32 products in shared memory, and one thread per column adds them in row order, writing one workspace row per chunk for the C++ reduction. The general kernel adds each chunk as two 32-row partials, which differs from CUDA in the last bit; this path matches it. As with tables, the dbias path walks row tiles, uses 256-column tiles, and indexes tiles of a run-time row stride with 32-bit offsets. GB200, BF16, geometric mean speedup over TE main, tables off / on: dbias+dGeLU 1.69 / 1.87 (general kernel: 1.74 / -), dbias+dSiLU 1.78 / 1.90 (general: 1.42 / -); CUDA backend 1.80 / 2.01 and 1.82 / 2.04. Signed-off-by: Teddy Do <tdophung@nvidia.com>
Fused activations (forward, derivative, with or without dbias) that ask for the tensor amax now stay on the register-resident bidimensional kernel. Each thread takes the max of the BF16 values it quantizes, ignoring NaN like the CUDA kernels' fmaxf; the CTA reduces them and updates the amax with an atomic max, as the general kernel does. Plain casts with amax keep the general kernel. The general kernel's scales differ from CUDA for blocks holding Inf or NaN; on this path the BF16 amax cases of the exact-composition test now match, including those with special values. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…rnel Requests with an activation table, dbias or amax now take the register- resident bidimensional kernel for any number of rows, as the CUDA register- tiled kernel does (the others keep whole 32-row tiles). A tensor whose rows end inside a tile launches a predicated kernel: rows past the tensor load as zeros, which leave the column amaxes and dbias unchanged, and are not stored. Tensors of whole tiles launch the unpredicated kernel, so their speed is unchanged; one kernel holding both variants would double the shared memory of the tile body. With dbias, the last 64-row chunk of an odd number of row tiles has no second tile; skipping it gives the same sum as adding its zero rows. The exact-composition test gains a 976x2048 shape: 31 row tiles, the last of 16 rows. Signed-off-by: Teddy Do <tdophung@nvidia.com>
… kernel A derivative holds the gradient and the activation input in registers. With 512-column tiles (16 columns per thread) the register-resident kernel spills at its 64-register budget; 256-column tiles without clusters, and 32-bit offsets within tiles of a run-time row stride, avoid that, as they already did for derivatives read from a table. Derivatives now also take any row count, like the CUDA register-tiled kernel. Only tiling and addressing change; the arithmetic is the same. GB200, BF16 without tables, geometric mean speedup over TE main: dGeLU 1.69 -> 1.87, dSiLU 1.55 -> 1.87 (CUDA backend: 1.68, 1.83). Signed-off-by: Teddy Do <tdophung@nvidia.com>
Requests with tails (derivatives, tables, amax) now also take any column count with N % 16 == 0, as the CUDA register-tiled kernel does. A thread's 8 or 16 columns then lie all inside or all past the tensor, so one flag per thread covers them: threads past the tensor load zeros and store no data or scales, which also covers a 16-column partial MX block. With dbias, workspace columns past the tensor are not written. The predicated kernel loads with predicated copies instead of branches, so a thread's row loads still issue together. GB200, BF16, geometric mean speedup over TE main on 993x7200, 4127x2064, 16384x2080 and 32768x7184, general kernel -> register kernel: tables off dSiLU 1.21 -> 1.63, dbias+dSiLU 1.41 -> 1.60; tables on GeLU 1.27 -> 1.59, dSiLU 1.19 -> 1.48, dbias+dSiLU 1.38 -> 1.71. The predicated kernel spills at the register budget, so tail shapes stay behind whole-tile ones. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…XFP8 kernel Fused activations with FP16 or FP32 input now take the register-resident bidimensional kernel where it is faster than the general one. As in the CUDA register-tiled kernel, these inputs are scaled in FP32: a block's scale is e8m0(max(amax * max_norm_rcp, 0)) and each value is multiplied in FP32 and converted to FP8. The activation is computed in FP32 and rounded to the input type; dbias takes the FP32 products. FP32 tiles use twice the warps with half the rows each. The BF16 datapath is unchanged. FP16 colwise and both-direction outputs match the CUDA kernels byte for byte (mxh2 parity against TE main, special values included). FP32 SiLU still differs from CUDA, as in the general kernel: CuTeDSL's exp is not expf. GB200, geometric mean speedup over TE main, general kernel -> this kernel: FP16 GeLU 1.46 -> 1.62, SiLU 1.41 -> 1.69, dGeLU 1.56 -> 1.69, dSiLU 1.28 -> 1.76, dbias+dSiLU 1.48 -> 1.74; FP32 SiLU 1.32 -> 1.54, dGeLU 1.38 -> 1.52, dSiLU 1.10 -> 1.85, dbias+dSiLU 1.18 -> 1.53. dbias+dGeLU and FP32 GeLU are faster on the general kernel and stay there. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…block scales - SiLU is x * sigmoid(x), as util/math.h computes it; x / (1 + exp(-x)) rounds differently, so FP32 SiLU differed from CUDA. - The general kernel carries one dbias sum down the column across its stages, as the CUDA kernels do, instead of adding 32-row partial sums. - With an activation fused, a block of only NaNs gets amax 0 and so scale 0, as in CUDA's generic kernel, in the register-resident, 2-lane rowwise and general kernels. Plain casts keep the NaN scale of CUDA's cast kernels. - FP32 dbias+dGeLU takes the register-resident kernel, whose dbias matches CUDA (the general kernel's still differs in the last bit for FP32). mxh2 parity against TE main (CUDA), tables off, 702 cases over BF16/FP16/FP32, all modes, layouts and fills including special values: 678 match (PR NVIDIA#3642 alone: 464). The rest are FP32 dbias with rowwise-only output, where CUDA adds row pairs (r, r + 32) before summing. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…ed one The CuTeDSL backend's register-resident kernel now covers GeLU, SiLU, dGeLU and dSiLU (tables, dbias, amax, tails, BF16/FP16/FP32), so the CUDA copy of it in cast/mxfp8/regtile is removed. The CUDA backend computes these in the generic kernel again; with NVTE_ACTIVATION_TABLE=1 it still follows the table specification (kComputedTable), so both backends give the same bytes. The table helpers only that kernel used (packed lookups, kComputed, kSignShift) go with it. CuTeDSL has no gated path yet, so the register-tiled SwiGLU/dSwiGLU kernel stays, moved with the helpers it needs to cast/mxfp8/specialized/gated_bidim.cuh next to the other register-resident MXFP8 kernels. Its speed is unchanged. Signed-off-by: Teddy Do <tdophung@nvidia.com>
for more information, see https://pre-commit.ci
Collaborator
Author
Collaborator
Author
|
the gated kernels with LUT and register-resident approach is also added in this PR in CUDA,. However I would like to move it to CuteDSL. However, @kainzhong is already working on it, so this draft will temporarily have the CUDA version, but will rebase don kaining's CuteDSL version of the gated kernel to add the LUT once that PR is up. This will make this PR also waiting for that. If not, I will drop support for LUT in gateed kernel entirely in this PR and make that a follow up PR |
…nt one Carries the CuTeDSL port of the gated MXFP8 kernel (gated_mxfp8.cuh) from kainzhong#13 (head 4687cc1), which is not merged yet, and adds a register-resident gated kernel on top of it. Revert this commit to drop both until that port lands. From kainzhong#13, unchanged except for formatting (black) and one argument: - quantize_mxfp8_gated_g2s.py: the gated kernel on the general kernel's TMA pipeline, for all gated activations, forward and backward. - MXFP8GatedQuantizeConfig, gated_mxfp8_cutedsl.cuh, cutedsl_utils.h and the dispatch in cast/dispatch/gated.cuh: CuTeDSL first, the CUDA kernel as the fallback. - The gated kernel passes ACTIVATION_TABLE="none" to the general shared-memory quantizer, which takes the activation table layout on this branch. Register-resident gated kernel (quantize_mxfp8_gated_g2r_bidim.py), a port of cast/mxfp8/specialized/gated_bidim.cuh for every gated activation: - BF16 input with columnwise output (rowwise optional). A CTA covers 256 output columns and walks 32-row bands. Each thread loads its rows of the act and gate halves (and of the gradient, in backward) into registers, computes the gated op there with the elementwise math of the shared-memory port, rounds it to BF16 as the CUDA kernels do, and quantizes the band in both directions from registers. The backward quantizes its two output halves one after the other. - The register-tile quantizer of the bidimensional kernel is factored out (quantize_mxfp8_rmem_tile_bidimensional, with its layout, bounds and row-load helpers) and shared by both kernels; the bidimensional kernel is unchanged. - Only the bands that cross a tensor's edge take the predicated body. - MXFP8GatedQuantizeEntry selects it at run time, and the shared-memory port otherwise (FP16/FP32, rowwise-only output, misaligned columnwise scales). mxh2 parity against TE main (CUDA), 1230 gated cases over SwiGLU, GeGLU, ReGLU, QGeGLU and SReGLU in both directions, all layouts, compact and swizzled scales and special-value fills: every case the register-resident kernel takes matches. The shared-memory port differs only in scale bytes of blocks of only NaNs (255 where the CUDA gated kernel writes 0). CastMXFP8_GatedAct passes with the CuTeDSL backend on and off. GB200, BF16, both directions, geometric mean speedup over TE main on 11 shapes: SwiGLU 1.58, dSwiGLU 1.82, GeGLU 1.26, dGeGLU 1.42, ReGLU 1.35, dReGLU 1.13, QGeGLU 1.59, dQGeGLU 2.27, SReGLU 1.42, dSReGLU 1.07 (CUDA backend: SwiGLU 1.65, dSwiGLU 1.87 with the register-tiled kernel, about 1.0 for the others). Co-authored-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Teddy Do <tdophung@nvidia.com>
This branch has not been deployed
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
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.
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: