From fdd1dd1b975c019c94d7eba9d4e9b509407cd21f Mon Sep 17 00:00:00 2001 From: MichaelDementii Date: Wed, 16 Sep 2026 17:52:30 +0200 Subject: [PATCH] perf(nvfp4): take a partial last M tile in the fused SwiGLU TMA route 5f5fccab gave the shared W4A4 TMA kernel a partial last M tile, and the five Linear shapes plus linear_add, gdn_input_proj and attn_input_proj route ragged widths to it now. The fused SwiGLU route was left out of that change: it dispatches its own kernel, and which implementation should serve its ragged widths depended on whether 05507ab0 was in the tree. Both are on master, so this finishes it. The route required tokens % 256 == 0. A prefill chunk only has to be a multiple of 128 and a prompt's last chunk is whatever is left, so on most real widths it was not taken and linear + silu_mul ran instead. That cost twice over. The composition materialises the 34816xT projection in BF16 and reads it back where the fused kernel writes 17408xT once; and because a ragged width could reach the composition, the capacity the Op reports had to cover that projection at the interval's widest ragged width. linear_swiglu_workspace_capacity_bytes for [1, 8192] returned 566.43 MiB, reserved whether or not a ragged width was ever asked for. It now returns 35.34 MiB. The kernel takes the real token count, its grid is a ceiling division, and the last M tile may be partial - the same three parts the shared kernel has. The activation-scale plane is described over the padded token count, since dim1 of that descriptor counts whole tiles; the plane is already allocated and zero-filled over nvfp4_w4a4_padded_tokens and its extent already checked by launch_nvfp4_w4a4_quantize, both from 5f5fccab, so the producer side needs nothing. The launcher's whole-tile precondition becomes a positive-count one, which is what the shared descriptor builder settled on in the same commit, and the header now states the precondition that replaces it: the scale plane must span the padded token count. The epilogue bounds itself - a padded row reads only memory this route owns, so it computes without reaching past anything, and it owns no output, so its store is dropped. The ragged floor is two whole M tiles. While the grid is two tiles the partial one is half of it, and there the two implementations are close enough that the sign of the difference is not reproducible: measured at every integer, the composition is ahead through T=263 by up to 0.69 % and the fused route from T=264, but at T=264 itself four sessions read -0.69, -0.08, +0.31 and +0.59 %, and inside the four-arm session alone the six passes span -0.76 to +0.61 %. A constant placed there is calibrated on noise and on one machine's clock, so that band keeps the composition it already had. The whole-tile floor at 256 does not move. What this gives up is measured and stated below. Four arms in one session: master, this change, the same change with the floor at one tile, and the master binary again under a second label as an identical-baseline null. Six passes, the arm order rotated through six permutations, 291 widths sampled at every integer for 48 past each tile boundary, --warmup 20 --repeat 300, a 256 MiB L2 flush before each sample. A pass is dropped for a cell when the two identical binaries disagree by more than 1 %; 35 of 1746 were. band widths median best worst rerouted, 513..3000 208 -9.03% -16.09% -5.93% whole tiles, same route on both 7 +0.05% -0.42% +0.32% 257..511, same route on both 63 +0.03% -0.14% +0.15% T < 256, same route on both 13 0.00% -0.79% 0.00% No measured rerouted width is slower; 208 of the 2488 ragged widths in that range were sampled. The last three rows are controls: both arms run the same implementation there, so they measure the instrument plus whatever the mechanism costs. The worst is T=1024 at +0.32 % against a null of 0.117 %, and the widest swing is T=192 at -0.79 % against a null of 0.064, on a width neither arm reroutes. A separate three-arm session reproduces the rerouted band at -8.82 % median and -5.88 % worst. The floor at one tile, measured in the same session, would take the 257..511 band to a median of -0.77 % and -6.75 % at its top, at the cost of seven widths - 257 through 263 - regressing by 0.14 to 0.69 %. Above 512 the two floors are indistinguishable: -9.00 % against -9.03 %. Absolute performance of the final implementation, median of 300 cold samples, with the benchmark's own bandwidth references: T us GB/s % of 1792.0 % of 1674.5 sustained read 1 72.928 1375.5 76.8 82.1 4 81.040 1239.5 69.2 74.0 8 72.960 1379.3 77.0 82.4 128 83.168 1275.0 71.1 76.1 512 169.360 728.3 40.6 43.5 1024 274.624 533.1 29.8 31.8 513 201.968 610.9 34.1 36.5 1025 320.736 456.6 25.5 27.3 1500 390.384 430.0 24.0 25.7 The bench defines no dense NVFP4 tensor peak, so no Tensor Core utilisation is stated. End to end on qwen3.8-27b NVFP4, ninfer_bench -p 8192, four passes with the arm order rotated: at --prefill-chunk 1408, a legal product width that is ragged against the M tile and above the floor, prefill goes from 9208.5 to 9421.4 tok/s, +2.37 % [+2.23 .. +2.47] against a null of +0.02 % [-0.05 .. +0.27]. At --prefill-chunk 1536, six whole tiles and the same route on both arms, -0.05 % [-0.13 .. +0.10]. Output is unchanged everywhere master already measured: NINFER_OP_REPORT_STATS=1 at %.17g over the whole suite, not only the Ops on this route, gives 11900 shared cases and 13928 lines, all byte-identical. No width in master's suite changes route, so nothing there should differ. Giving master's tree the new case list and reading both arms at identical widths, three of 29 cases differ - 513, 767 and 1025, exactly the rerouted ones - and at each the fused route sits closer to the FP64 oracle than the composition it replaces, 0.3567 to 0.3535, 0.3376 to 0.3358 and 0.3311 to 0.3293 as a fraction of the criterion. The composition rounds the projection to BF16 and applies silu to the rounded value; the fused kernel applies it to the fp32 accumulator and rounds once. Tests gain 513, 767 and 1025, and 511 which straddles the floor from below. Each was checked against a build with one part of the mechanism broken. A grid left as a floor division fails at 513, 767 and 1025 with a non-finite value, the unwritten tail read back. A scale plane described over the real token count instead of the padded one fails at 767 alone, with actual=0 reference=5.48743 - at 513 and 1025 the last tile carries one token, and one zeroed column in five hundred stays inside the A4 profile's relative-L2 allowance of 1.6e-1. An unbounded epilogue store fails at 513 with cudaErrorIllegalAddress, overrunning the output by 255 tokens. No numeric case can catch a wrong value of the floor: both routes are correct, and the harness compares numerics and the capacity the same constant computes, so both sides move together. The routing is pinned by a strength control instead. Scaling the fused epilogue by two fails at exactly 256, 512, 513, 767, 1024 and 1025 and passes at every other case in the list, including 257 and 511. ctest 120/120 on both arms. The SASS census of the fused kernel reports 1402 instructions on each side with nine opcode counts differing: the added bound and its parameter load. --- .../nvfp4/nvfp4_linear_swiglu_plan.cpp | 36 ++++++++++++++----- .../nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu | 32 ++++++++++------- .../nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh | 8 ++++- .../nvfp4_linear_swiglu_w4a4_tma_launch.h | 10 ++++++ tests/ops/linear_swiglu/test_nvfp4.cpp | 13 ++++--- 5 files changed, 73 insertions(+), 26 deletions(-) diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp index ec56e64ca5..31e103a1b1 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp @@ -24,6 +24,19 @@ enum class Nvfp4LinearSwiGluRoute { constexpr std::int32_t kFusedMaxTokens = 128; +// A ragged width is worth the fused route once its partial M tile is not half of the grid. While +// the grid is two tiles the fused kernel and the public linear + silu_mul composition are close +// enough on an RTX 5090 that the sign of the difference does not survive a second session, so +// that band keeps the composition it already had. +constexpr std::int32_t kRaggedTmaFloor = 2 * kNvfp4TmaBlockM; +// The capacity function walks one step down from this floor to find the widest width the +// composition still serves. That walk is only correct while the floor is at least one whole tile: +// below that the composition owns widths above the floor as well, and one step does not reach +// them. +static_assert(kRaggedTmaFloor >= kNvfp4TmaBlockM, + "a ragged floor inside the first M tile would leave the composition band above it " + "out of the workspace capacity"); + Nvfp4LinearSwiGluRoute resolve_route(LinearPolicy policy, std::int32_t tokens) { if (tokens <= 0) { throw std::invalid_argument("nvfp4 linear_swiglu: T must be positive"); } if (!valid_linear_policy(policy)) { @@ -38,8 +51,10 @@ Nvfp4LinearSwiGluRoute resolve_route(LinearPolicy policy, std::int32_t tokens) { if (tokens <= 4) { return Nvfp4LinearSwiGluRoute::SmallTFusedA16; } if (tokens <= kFusedMaxTokens) { return Nvfp4LinearSwiGluRoute::FusedW4A4; } // This route dispatches its own fused kernel rather than a Linear shape's, so it carries its - // own condition; the call site below forces the matching scale layout. - if (tokens >= kNvfp4TmaBlockM && (tokens % kNvfp4TmaBlockM) == 0) { + // own condition - a whole tile, or a ragged width with enough behind it; the call site below + // forces the matching scale layout either way. + if (tokens >= kNvfp4TmaBlockM && + ((tokens % kNvfp4TmaBlockM) == 0 || tokens >= kRaggedTmaFloor)) { return Nvfp4LinearSwiGluRoute::TmaFusedW4A4; } return Nvfp4LinearSwiGluRoute::LinearW4A4Post; @@ -96,14 +111,19 @@ std::size_t nvfp4_linear_swiglu_workspace_capacity_bytes(LinearPolicy policy, if (min_tokens <= kFusedMaxTokens && max_tokens >= 5) { maximum = fused_workspace_bytes(std::min(max_tokens, kFusedMaxTokens)); } - if (max_tokens >= kNvfp4TmaBlockM) { - const std::int32_t largest_fused = max_tokens - (max_tokens % kNvfp4TmaBlockM); - if (largest_fused >= std::max(min_tokens, kNvfp4TmaBlockM)) { - maximum = std::max(maximum, fused_workspace_bytes(largest_fused)); - } + // From the ragged floor upward the TMA route owns every width, so the widest it can be asked + // for is the interval's own maximum; below the floor it owns only whole tiles, and the guard + // then rejects an interval that contains none. + const std::int32_t largest_fused = + max_tokens >= kRaggedTmaFloor ? max_tokens : max_tokens - (max_tokens % kNvfp4TmaBlockM); + if (largest_fused >= std::max(min_tokens, kNvfp4TmaBlockM)) { + maximum = std::max(maximum, fused_workspace_bytes(largest_fused)); } - std::int32_t last_baseline = max_tokens; + // The composition keeps the gap below the floor, so it no longer has to be sized for the + // interval's maximum. Capping at the floor puts the search where the only TMA widths are whole + // tiles, and no two consecutive widths are both whole tiles, so one step is always enough. + std::int32_t last_baseline = std::min(max_tokens, kRaggedTmaFloor - 1); if (resolve_route(policy, last_baseline) == Nvfp4LinearSwiGluRoute::TmaFusedW4A4) { --last_baseline; } diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu index 1a3f87b288..7ddf3e2e08 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu @@ -13,6 +13,9 @@ namespace ninfer::ops::detail { namespace { using M256N128S3 = Nvfp4W4a4TmaSchedule<256, 3, 1>; +// The descriptor below divides the padded token count by this tile to get the plane's tile count, +// and nvfp4_w4a4_padded_tokens rounds up to kNvfp4TmaBlockM, so the two have to be the same number +// or that division stops being exact and the last tile addresses past the plane. static_assert(M256N128S3::kBlockM == kNvfp4TmaBlockM, "the fused descriptor below addresses the quantizer's scale tiles"); @@ -40,13 +43,17 @@ Nvfp4W4a4TmaDescriptors make_descriptors(const std::uint8_t* activation_codes, const_cast(weight_codes), CU_TENSOR_MAP_DATA_TYPE_UINT8, Geometry::kCodeBytesPerRow, Geometry::kOutputRows, Geometry::kCodeBytesPerRow, kCodeColumns, kPairN, CU_TENSOR_MAP_SWIZZLE_64B, "encode LinearSwiGLU weight codes TMA"); - descriptors.a_scales = - nvfp4_make_tma_2d(const_cast(activation_scales), - CU_TENSOR_MAP_DATA_TYPE_UINT8, Schedule::kBlockM, - (static_cast(tokens) / Schedule::kBlockM) * - kScaleTilesPerPlane * kScaleTileGroups, - Schedule::kBlockM, Schedule::kBlockM, kScaleTileGroups, - CU_TENSOR_MAP_SWIZZLE_NONE, "encode LinearSwiGLU activation scales TMA"); + // dim1 counts whole token tiles, so the plane is described over the padded token count the + // quantizer writes. Describing it over the real count instead would leave the last tile's + // scales outside the descriptor, where TMA zero-fills them, and every token in that tile would + // be written out as zero. + descriptors.a_scales = nvfp4_make_tma_2d( + const_cast(activation_scales), CU_TENSOR_MAP_DATA_TYPE_UINT8, + Schedule::kBlockM, + (static_cast(nvfp4_w4a4_padded_tokens(tokens)) / Schedule::kBlockM) * + kScaleTilesPerPlane * kScaleTileGroups, + Schedule::kBlockM, Schedule::kBlockM, kScaleTileGroups, CU_TENSOR_MAP_SWIZZLE_NONE, + "encode LinearSwiGLU activation scales TMA"); descriptors.b_scales = nvfp4_make_tma_2d(const_cast(weight_scales), CU_TENSOR_MAP_DATA_TYPE_UINT8, 16, kWeightScaleBytes / 16, 16, 16, 64, CU_TENSOR_MAP_SWIZZLE_NONE, @@ -61,9 +68,8 @@ void launch_nvfp4_linear_swiglu_w4a4_tma(const std::uint8_t* activation_codes, const std::uint8_t* weight_codes, const std::uint8_t* weight_scales, __nv_bfloat16* output, std::int32_t tokens, float alpha, cudaStream_t stream) { - if (tokens < M256N128S3::kBlockM || (tokens % M256N128S3::kBlockM) != 0) { - throw std::invalid_argument( - "nvfp4 LinearSwiGLU TMA requires a positive M256 full-tile token count"); + if (tokens <= 0) { + throw std::invalid_argument("nvfp4 LinearSwiGLU TMA needs a positive token count"); } using Geometry = Nvfp4N34816K5120; @@ -79,9 +85,11 @@ void launch_nvfp4_linear_swiglu_w4a4_tma(const std::uint8_t* activation_codes, const Nvfp4W4a4TmaDescriptors descriptors = make_descriptors( activation_codes, activation_scales, weight_codes, weight_scales, tokens); constexpr int kPairN = M256N128S3::kBlockN / 2; - const dim3 grid((Geometry::kOutputRows / 2) / kPairN, tokens / M256N128S3::kBlockM); + // The last M tile may be partial; the kernel bounds its stores by the real token count. + const dim3 grid((Geometry::kOutputRows / 2) / kPairN, + (tokens + M256N128S3::kBlockM - 1) / M256N128S3::kBlockM); nvfp4_linear_swiglu_w4a4_tma_kernel - <<>>(descriptors, alpha, output); + <<>>(descriptors, alpha, output, tokens); CUDA_CHECK(cudaGetLastError()); } diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh index 2ad6f3de5f..ec2eae0f1c 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh @@ -50,7 +50,8 @@ __global__ __launch_bounds__( Nvfp4W4a4TmaDescriptors descriptors, float alpha, - __nv_bfloat16* __restrict__ output) { + __nv_bfloat16* __restrict__ output, + int token_count) { static_assert(Geometry::kOutputRows == 34816); static_assert(Geometry::kInputRows == 5120); static_assert((Geometry::kInputRows % Schedule::kBlockK) == 0); @@ -272,6 +273,11 @@ __global__ __launch_bounds__( const int token_local = task / kVectorsPerRow; const int row_vector = task - token_local * kVectorsPerRow; const int token = token_begin + token_local; + // The last M tile may be partial. A padded row reads only memory this route owns - its + // codes are TMA zero-fill, because the code descriptor's row extent is the real token + // count, and its scales are the zeroes in the padded plane - so it computes without + // reaching past anything. It owns no output, so its store is dropped here. + if (token >= token_count) { continue; } const uint4 values = load_vec(shared_output + token_local * kOutputStride + row_vector * 8); store_vec(output + static_cast(token) * kIntermediate + pair_begin + diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma_launch.h b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma_launch.h index 4cd5d90368..15678d9f33 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma_launch.h +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma_launch.h @@ -7,6 +7,16 @@ namespace ninfer::ops::detail { +/** + * Fused NVFP4 LinearSwiGLU over the W4A4 TMA route, for the 34816x5120 geometry. tokens may be any + * positive count; the last M tile may be partial. + * + * activation_scales must be the tiled layout written over whole M tiles, that is a plane of + * nvfp4_w4a4_padded_tokens(tokens) rows with the padding zeroed - what + * allocate_nvfp4_w4a4_workspace sizes and launch_nvfp4_w4a4_quantize fills and checks. The + * descriptor this builds covers that whole plane, so a caller that sizes the plane at tokens + * instead has it read past the end. + */ void launch_nvfp4_linear_swiglu_w4a4_tma(const std::uint8_t* activation_codes, const std::uint8_t* activation_scales, const std::uint8_t* weight_codes, diff --git a/tests/ops/linear_swiglu/test_nvfp4.cpp b/tests/ops/linear_swiglu/test_nvfp4.cpp index 53cb824459..90fa065785 100644 --- a/tests/ops/linear_swiglu/test_nvfp4.cpp +++ b/tests/ops/linear_swiglu/test_nvfp4.cpp @@ -11,11 +11,14 @@ int main() { try { constexpr std::array kA16Cases{1, 4, 8, 16}; - // 255, 256 and 257 straddle the fused route floor. This change does not move that - // route, but 257 reaches it through the baseline composition, whose own linear now - // runs the ragged TMA path. - constexpr std::array kA4Cases{2, 4, 5, 16, 56, 64, 65, 96, 97, - 112, 128, 129, 255, 256, 257, 512, 1024}; + // 511 and 513 straddle the ragged floor: the first still reaches the composition, the + // second is the narrowest ragged width the fused route admits and leaves one real token + // in a third M tile. 767 leaves that tile all but full, and 1025 leaves one after four + // whole tiles. 255/256 straddle the whole-tile floor, which does not move; 257 stays as + // the width just past it that the composition keeps. + constexpr std::array kA4Cases{2, 4, 5, 16, 56, 64, 65, + 96, 97, 112, 128, 129, 255, 256, + 257, 511, 512, 513, 767, 1024, 1025}; int failures = 0; failures += run_profile("LinearSwiGLU NVFP4_A16", {QType::NVFP4, 34816, 5120, 17408, 1801U, ActivationCompute::A16},