NVFP4 fused SwiGLU: take a partial last M tile, so ragged widths stop falling back to linear + silu_mul

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

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<std::uint8_t*>(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<std::uint8_t*>(activation_scales),
-                          CU_TENSOR_MAP_DATA_TYPE_UINT8, Schedule::kBlockM,
-                          (static_cast<std::uint64_t>(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<std::uint8_t*>(activation_scales), CU_TENSOR_MAP_DATA_TYPE_UINT8,
+        Schedule::kBlockM,
+        (static_cast<std::uint64_t>(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<std::uint8_t*>(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<Geometry, M256N128S3>(
         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<Geometry, M256N128S3>
-        <<<grid, M256N128S3::kThreads, kSharedBytes, stream>>>(descriptors, alpha, output);
+        <<<grid, M256N128S3::kThreads, kSharedBytes, stream>>>(descriptors, alpha, output, tokens);
     CUDA_CHECK(cudaGetLastError());
 }
 

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<uint4>(shared_output + token_local * kOutputStride + row_vector * 8);
         store_vec(output + static_cast<std::int64_t>(token) * kIntermediate + pair_begin +

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,

tests/ops/linear_swiglu/test_nvfp4.cpp

@@ -11,11 +11,14 @@ int main() {
 
     try {
         constexpr std::array<std::int32_t, 4> 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<std::int32_t, 17> 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<std::int32_t, 21> 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},