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},