From 8881952fdbad553314ea95c4e8f449fbe245134d Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:31:18 +0000 Subject: [PATCH] perf(nvfp4): make the TMA route read activation scales tile-contiguous The activation scale descriptor is a 16-byte box over the scale plane, so one stage of the TMA pipeline asks for BlockM separate 16-byte transfers to deliver 4 KiB. The weight codes beside it travel as one 64-byte box per row and cost nothing like that. A standalone harness that runs the producer alone - the same four descriptors, the same 29,696-byte stage, the same three-deep mbarrier ring, no math - shows where the stage time goes: the codes alone move at 6977 GB/s; adding the weight scales costs 14.8% more time for 4.2% more bytes; adding the activation scales costs 73.0% more time for 16.7% more bytes. Requesting the same bytes as one [BlockM tokens, 16 groups] tile restores the codes-only rate, 7009 GB/s against 4732 for the same payload. The byte order inside a tile is unchanged, so the shared image the consumer reads is identical and the consumer is untouched: only the address the quantizer writes to, the descriptor shape and the request coordinate move. The quantizer writes the plane blocked when the GEMM that follows will take the TMA route and keeps it row-major otherwise, because the small-T and GEMV routes read the same buffer. The predicate is the route's own; the fused SwiGLU route admits every multiple of 256 from 256 up, which is wider than the shared one. RTX 5090, sm_120a, 510 W cap, 3105 MHz boost, CUDA 13.1.115, Release. Both builds come from one tree, interleaved ABAB, three rounds, one process per run. Operator, T=4096, NVFP4 A4: shape before us after us ratio GB/s TFLOP/s linear_add [5120, 17408] 759.808 612.352 0.806 364 -> 452 961 -> 1192 linear_add [5120, 6144] 303.104 262.144 0.865 501 -> 580 850 -> 983 linear_swiglu [34816, 5120] 1456.128 1214.464 0.834 196 -> 234 1003 -> 1202 Against the hardware limits: the GDDR7 spec peak is 1792 GB/s, so these shapes sit between 11% and 28% of it before and between 13% and 32% after - the route is nowhere near bandwidth-bound at either end, and what the change buys is transaction width rather than bandwidth. For the tensor core, the block-scaled NVFP4 MMA measures 2117 TFLOP/s in isolation on this part, and 1897 when fed in this kernel's own ratio of 4 ldmatrix.x4 and 8 ldmatrix.x2 per 32 MMA; against that 1897 the three shapes move from 51%, 45% and 53% to 63%, 52% and 63%. Those two ceilings were taken at full boost and these runs did not sustain it, so the tensor-core fractions are a lower bound; the bandwidth fractions are not affected, memory speed being independent of the SM clock. Controls in the same runs, below the 256-token floor where the route is not taken: linear_add at T=64 and T=128, fused SwiGLU at T=64. All three are 1.000 to the printed digit, which is what shows the two builds are otherwise identical. End-to-end confirmation, Qwen3.6-27B NVFP4, 32768-token prompt, prefill chunk 4096, KV int8: prefill 9426.4 -> 9811.8 tok/s (+4.09%), 3.4762 -> 3.3397 s, with round-to-round spreads of 101.4 and 98.6 tok/s. Decode is the untouched control at 87.945 -> 87.958 tok/s (+0.01%): its token counts never reach the tile. Resource deltas. Workspace capacity and peak, weight capacity and KV payload are byte-identical between the builds as the engine reports them: 639,926,272, 17,206,931,200 and 1,107,296,256. No kernel changes registers or shared memory, and none of the 37 functions in the touched translation units spills. The binary gains three kernel instantiations: the quantizer template takes a flat/blocked flag, so each of the three activation geometries is emitted twice, at 40 registers and no shared memory. Numerical validation. Values do not change - only the address a scale byte is written to - and greedy generation matches the previous binary byte for byte on four prompts of 8K, 64K and two of 55K tokens. ctest is 102 of 102, with six cases skipped for want of real model artifacts. clang-format 23.1.0 reports eight violations across these files. All of them are on lines this change does not touch, so they are left as they are. --- .../nvfp4/nvfp4_attn_input_w4a4.cu | 2 +- .../nvfp4/nvfp4_gdn_input_w4a4.cu | 2 +- src/ops/linear/nvfp4/nvfp4_w4a4.cu | 26 +++++++++++------ src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh | 26 +++++++++++++++-- src/ops/linear/nvfp4/nvfp4_w4a4_plan.h | 11 +++++++- src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh | 28 +++++++++++++------ .../linear_add/nvfp4/nvfp4_linear_add_w4a4.cu | 2 +- .../nvfp4/nvfp4_linear_swiglu_plan.cpp | 4 ++- .../nvfp4/nvfp4_linear_swiglu_w4a4.cu | 3 +- .../nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu | 20 ++++++++----- .../nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh | 10 +++++-- 11 files changed, 100 insertions(+), 34 deletions(-) diff --git a/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_w4a4.cu b/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_w4a4.cu index 56835368cf..66cceea17b 100644 --- a/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_w4a4.cu +++ b/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_w4a4.cu @@ -83,7 +83,7 @@ void launch_gemm(const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tenso void nvfp4_attn_input_w4a4_launch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate, Tensor& k, Tensor& v, Nvfp4W4a4Workspace workspace, cudaStream_t stream) { - launch_nvfp4_w4a4_quantize(x, weight, workspace, stream); + launch_nvfp4_w4a4_quantize(x, weight, workspace, nvfp4_w4a4_tma_route(x.ne[1]), stream); const std::int32_t tokens = x.ne[1]; if (tokens >= 1024 && (tokens % kTmaBlockM) == 0) { const float alpha = 1.0F / (weight.input_scale_divisor * weight.weight_scale_divisor); diff --git a/src/ops/gdn_input_proj/nvfp4/nvfp4_gdn_input_w4a4.cu b/src/ops/gdn_input_proj/nvfp4/nvfp4_gdn_input_w4a4.cu index f56a4cd0fe..f33fca8e96 100644 --- a/src/ops/gdn_input_proj/nvfp4/nvfp4_gdn_input_w4a4.cu +++ b/src/ops/gdn_input_proj/nvfp4/nvfp4_gdn_input_w4a4.cu @@ -37,7 +37,7 @@ void launch_gemm(const Weight& weight, Tensor& qkv, Tensor& z, Nvfp4W4a4Workspac void nvfp4_gdn_input_w4a4_launch(const Tensor& x, const Weight& weight, Tensor& qkv, Tensor& z, Nvfp4W4a4Workspace workspace, cudaStream_t stream) { - launch_nvfp4_w4a4_quantize(x, weight, workspace, stream); + launch_nvfp4_w4a4_quantize(x, weight, workspace, nvfp4_w4a4_tma_route(x.ne[1]), stream); const std::int32_t tokens = x.ne[1]; if (tokens >= 1024 && (tokens % kTmaBlockM) == 0) { const float alpha = 1.0F / (weight.input_scale_divisor * weight.weight_scale_divisor); diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4.cu b/src/ops/linear/nvfp4/nvfp4_w4a4.cu index 156871fbbe..1fb451acb2 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4.cu +++ b/src/ops/linear/nvfp4/nvfp4_w4a4.cu @@ -38,14 +38,21 @@ void launch_gemm(const Weight& weight, Tensor& out, Nvfp4W4a4Workspace workspace template void launch_quantize_exact(const Tensor& x, const Weight& weight, Nvfp4W4a4Workspace workspace, - cudaStream_t stream) { + bool blocked_scales, cudaStream_t stream) { const std::int32_t tokens = x.ne[1]; constexpr int kThreads = 256; const std::int32_t tasks = tokens * ActivationGeometry::kGroupsPerRow; - nvfp4_w4a4_quantize_kernel - <<<(tasks + kThreads - 1) / kThreads, kThreads, 0, stream>>>( + const int blocks = (tasks + kThreads - 1) / kThreads; + if (blocked_scales) { + nvfp4_w4a4_quantize_kernel + <<>>(static_cast(x.data), + workspace.codes, workspace.scales, tokens, + weight.input_scale_divisor); + } else { + nvfp4_w4a4_quantize_kernel<<>>( static_cast(x.data), workspace.codes, workspace.scales, tokens, weight.input_scale_divisor); + } CUDA_CHECK(cudaGetLastError()); } @@ -89,19 +96,22 @@ void launch_problem(const Weight& weight, Tensor& out, Nvfp4W4a4Workspace worksp } // namespace void launch_nvfp4_w4a4_quantize(const Tensor& x, const Weight& weight, Nvfp4W4a4Workspace workspace, - cudaStream_t stream) { + bool blocked_scales, cudaStream_t stream) { if (workspace.codes == nullptr || workspace.scales == nullptr) { throw std::invalid_argument("nvfp4 W4A4 requires caller workspace"); } switch (weight.k) { case Nvfp4Activation5120Geometry::kInputRows: - launch_quantize_exact(x, weight, workspace, stream); + launch_quantize_exact(x, weight, workspace, blocked_scales, + stream); return; case Nvfp4Activation6144Geometry::kInputRows: - launch_quantize_exact(x, weight, workspace, stream); + launch_quantize_exact(x, weight, workspace, blocked_scales, + stream); return; case Nvfp4Activation17408Geometry::kInputRows: - launch_quantize_exact(x, weight, workspace, stream); + launch_quantize_exact(x, weight, workspace, blocked_scales, + stream); return; default: throw std::invalid_argument("nvfp4 W4A4 quantize: unsupported K"); @@ -110,7 +120,7 @@ void launch_nvfp4_w4a4_quantize(const Tensor& x, const Weight& weight, Nvfp4W4a4 void launch_nvfp4_w4a4(const Tensor& x, const Weight& weight, Tensor& out, Nvfp4W4a4Workspace workspace, cudaStream_t stream) { - launch_nvfp4_w4a4_quantize(x, weight, workspace, stream); + launch_nvfp4_w4a4_quantize(x, weight, workspace, nvfp4_w4a4_tma_route(x.ne[1]), stream); const std::int32_t tokens = x.ne[1]; switch (resolve_nvfp4_problem(weight.n, weight.k)) { case Nvfp4Problem::AttnInput: diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh b/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh index 675df5b992..47cb9d718d 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_mma.cuh @@ -384,11 +384,29 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 } } -template +// The TMA route reads a whole [BlockM tokens, 16 groups] tile of activation scales per request. +// Row-major by token that tile is 16 bytes per row with the scale plane's row stride, so one stage +// costs BlockM tiny requests. Writing the same bytes tile-contiguous makes the request wide. The +// byte order inside a tile is unchanged, so the shared image the consumer reads is identical. +template +__device__ __forceinline__ std::int64_t nvfp4_blocked_scale_offset(int token, int group) { + constexpr int kGroupsPerRow = Geometry::kInputRows / 16; + constexpr int kGroupsPerTile = 16; + constexpr int kTilesPerPlane = kGroupsPerRow / kGroupsPerTile; + const int token_tile = token / BlockM; + const int group_tile = group / kGroupsPerTile; + const int tile = token_tile * kTilesPerPlane + group_tile; + return static_cast(tile) * BlockM * kGroupsPerTile + + static_cast(token - token_tile * BlockM) * kGroupsPerTile + + (group - group_tile * kGroupsPerTile); +} + +template __global__ __launch_bounds__(Threads, 512 / Threads) void nvfp4_w4a4_quantize_kernel( const __nv_bfloat16* __restrict__ input, std::uint8_t* __restrict__ codes, std::uint8_t* __restrict__ scales, std::int32_t tokens, float input_scale_divisor) { static_assert(Threads == 128 || Threads == 256 || Threads == 512); + static_assert(!BlockedScales || (Geometry::kInputRows / 16) % 16 == 0); constexpr int kGroupsPerRow = Geometry::kInputRows / 16; const int task = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); @@ -403,7 +421,11 @@ __global__ __launch_bounds__(Threads, 512 / Threads) void nvfp4_w4a4_quantize_ke auto* code_destination = codes + static_cast(token) * Geometry::kCodeBytesPerRow + group * 8; store_vec(code_destination, make_uint2(quantized.codes_lo, quantized.codes_hi)); - scales[static_cast(token) * kGroupsPerRow + group] = quantized.scale; + if constexpr (BlockedScales) { + scales[nvfp4_blocked_scale_offset(token, group)] = quantized.scale; + } else { + scales[static_cast(token) * kGroupsPerRow + group] = quantized.scale; + } } } // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_plan.h b/src/ops/linear/nvfp4/nvfp4_w4a4_plan.h index 9695374c2c..179405f591 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_plan.h +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_plan.h @@ -50,8 +50,17 @@ inline std::size_t nvfp4_w4a4_workspace_capacity_bytes(std::int32_t tokens, return layout.peak_bytes(1); } +// The TMA route reads activation scales one [256 tokens, 16 groups] tile per request, and wants +// that tile contiguous. Every route that shares make_nvfp4_w4a4_tma_descriptors asks the same +// question, so the predicate lives here. +inline constexpr std::int32_t kNvfp4TmaBlockM = 256; + +[[nodiscard]] inline bool nvfp4_w4a4_tma_route(std::int32_t tokens) { + return tokens >= 1024 && (tokens % kNvfp4TmaBlockM) == 0; +} + void launch_nvfp4_w4a4_quantize(const Tensor& x, const Weight& weight, Nvfp4W4a4Workspace workspace, - cudaStream_t stream); + bool blocked_scales, cudaStream_t stream); void launch_nvfp4_w4a4(const Tensor& x, const Weight& weight, Tensor& out, Nvfp4W4a4Workspace workspace, cudaStream_t stream); diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh b/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh index 6c0d1e3e28..8b5a475976 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh @@ -56,10 +56,13 @@ Nvfp4W4a4TmaDescriptors make_nvfp4_w4a4_tma_descriptors(const std::uint8_t* acti std::int32_t tokens) { static_assert(BlockM == 128 || BlockM == 256); constexpr std::uint32_t kCodeColumns = 64; - // TMA's innermost box is at least one 16-byte transaction. A K128 tile consumes the - // first eight bytes of each row; the second half is harmless look-ahead. - constexpr std::uint32_t kScaleColumns = 16; - constexpr std::uint32_t kBlockN = 128; + // Activation scales arrive tile-contiguous: one [BlockM tokens, 16 groups] tile is BlockM + // bytes wide and 16 rows tall, so the request is wide instead of BlockM separate 16-byte ones. + // A K128 tile consumes the first eight of the sixteen group bytes; the rest is look-ahead. + constexpr std::uint32_t kScaleTileGroups = 16; + constexpr std::uint64_t kScaleTilesPerPlane = + static_cast(Geometry::kGroupsPerRow) / kScaleTileGroups; + constexpr std::uint32_t kBlockN = 128; constexpr std::uint64_t kWeightScaleBytes = static_cast(Geometry::kOutputRows) * Geometry::kInputRows / 16; @@ -73,9 +76,10 @@ Nvfp4W4a4TmaDescriptors make_nvfp4_w4a4_tma_descriptors(const std::uint8_t* acti Geometry::kCodeBytesPerRow, Geometry::kOutputRows, Geometry::kCodeBytesPerRow, kCodeColumns, kBlockN, CU_TENSOR_MAP_SWIZZLE_64B, "encode weight codes TMA"); descriptors.a_scales = nvfp4_make_tma_2d( - const_cast(activation_scales), CU_TENSOR_MAP_DATA_TYPE_UINT8, - Geometry::kGroupsPerRow, tokens, Geometry::kGroupsPerRow, kScaleColumns, BlockM, - CU_TENSOR_MAP_SWIZZLE_NONE, "encode activation scales TMA"); + const_cast(activation_scales), CU_TENSOR_MAP_DATA_TYPE_UINT8, BlockM, + (static_cast(tokens) / BlockM) * kScaleTilesPerPlane * kScaleTileGroups, + BlockM, BlockM, kScaleTileGroups, CU_TENSOR_MAP_SWIZZLE_NONE, + "encode 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, "encode weight scales TMA"); @@ -220,8 +224,14 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 nvfp4_tma_load_2d(tensors.b_codes[stage], &descriptors.b_codes, k_tile * Schedule::kCodeRowBytes, row_begin, &shared.full[stage]); if (load_scales) { - nvfp4_tma_load_2d(tensors.a_scale4[(k_tile / 2) & 1], &descriptors.a_scales, - (k_tile / 2) * 16, token_begin, &shared.full[stage]); + // The box is tile-contiguous, so its address is a tile index rather than a + // (byte column, token row) pair; the two-slot buffer and the even-tile guard + // are unchanged. + constexpr int kScaleTilesPerPlane = Geometry::kGroupsPerRow / 16; + const int scale_tile = + (token_begin / Schedule::kBlockM) * kScaleTilesPerPlane + k_tile / 2; + nvfp4_tma_load_2d(tensors.a_scale4[(k_tile / 2) & 1], &descriptors.a_scales, 0, + scale_tile * 16, &shared.full[stage]); } const int b_scale_row = ((row_begin / 128) * Geometry::kScaleTilesPerRow + k_tile * Schedule::kK64PerStage) * diff --git a/src/ops/linear_add/nvfp4/nvfp4_linear_add_w4a4.cu b/src/ops/linear_add/nvfp4/nvfp4_linear_add_w4a4.cu index 4dc44ba376..0fdf7448a7 100644 --- a/src/ops/linear_add/nvfp4/nvfp4_linear_add_w4a4.cu +++ b/src/ops/linear_add/nvfp4/nvfp4_linear_add_w4a4.cu @@ -56,7 +56,7 @@ void launch_problem(const Weight& weight, Tensor& residual, Nvfp4W4a4Workspace w void nvfp4_linear_add_w4a4_launch(const Tensor& x, const Weight& weight, Tensor& residual, Nvfp4W4a4Workspace workspace, cudaStream_t stream) { - launch_nvfp4_w4a4_quantize(x, weight, workspace, stream); + launch_nvfp4_w4a4_quantize(x, weight, workspace, nvfp4_w4a4_tma_route(x.ne[1]), stream); const std::int32_t tokens = x.ne[1]; const Nvfp4Problem problem = resolve_nvfp4_problem(weight.n, weight.k); if (tokens >= 1024 && (tokens % kTmaBlockM) == 0) { 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 5912b5d660..11ea435e6f 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp @@ -126,7 +126,9 @@ void nvfp4_linear_swiglu_dispatch(const Tensor& x, const Weight& weight, Tensor& case Nvfp4LinearSwiGluRoute::TmaFusedW4A4: { auto scope = workspace.scope(); const Nvfp4W4a4Workspace scratch = allocate_fused_workspace(workspace, x.ne[1]); - launch_nvfp4_w4a4_quantize(x, weight, scratch, stream); + // Inside the TMA case, so this route always reads tile-contiguous scales. Note its own + // predicate admits every multiple of 256 from 256 up, which is wider than the shared one. + launch_nvfp4_w4a4_quantize(x, weight, scratch, true, stream); const float alpha = 1.0F / (weight.input_scale_divisor * weight.weight_scale_divisor); launch_nvfp4_linear_swiglu_w4a4_tma( scratch.codes, scratch.scales, static_cast(weight.qdata), diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4.cu b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4.cu index 13e1b60528..f6755bd693 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4.cu +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4.cu @@ -85,7 +85,8 @@ void launch(const Tensor& x, const Weight& weight, Tensor& out, WorkspaceArena& auto scope = workspace.scope(); const Nvfp4W4a4Workspace scratch = allocate_nvfp4_w4a4_workspace(workspace, x.ne[1], Geometry::kInputRows); - launch_nvfp4_w4a4_quantize(x, weight, scratch, stream); + // This route runs the MMA kernel, never TMA, so the scale plane stays row-major. + launch_nvfp4_w4a4_quantize(x, weight, scratch, false, stream); launch_gemm(weight, out, scratch, x.ne[1], stream); } 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 0127a8d1d5..1e2a990b4b 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 @@ -19,9 +19,12 @@ Nvfp4W4a4TmaDescriptors make_descriptors(const std::uint8_t* activation_codes, const std::uint8_t* activation_scales, const std::uint8_t* weight_codes, const std::uint8_t* weight_scales, std::int32_t tokens) { - constexpr std::uint32_t kCodeColumns = 64; - constexpr std::uint32_t kScaleColumns = 16; - constexpr std::uint32_t kPairN = Schedule::kBlockN / 2; + constexpr std::uint32_t kCodeColumns = 64; + // Activation scales arrive tile-contiguous, one [BlockM tokens, 16 groups] tile per request. + constexpr std::uint32_t kScaleTileGroups = 16; + constexpr std::uint64_t kScaleTilesPerPlane = + static_cast(Geometry::kGroupsPerRow) / kScaleTileGroups; + constexpr std::uint32_t kPairN = Schedule::kBlockN / 2; constexpr std::uint64_t kWeightScaleBytes = static_cast(Geometry::kOutputRows) * Geometry::kInputRows / 16; @@ -34,10 +37,13 @@ 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, - Geometry::kGroupsPerRow, tokens, Geometry::kGroupsPerRow, kScaleColumns, Schedule::kBlockM, - CU_TENSOR_MAP_SWIZZLE_NONE, "encode LinearSwiGLU activation scales 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"); 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, 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 7f85a36c85..3b922afbf1 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 @@ -119,8 +119,14 @@ __global__ __launch_bounds__( &descriptors.b_codes, k_tile * Schedule::kCodeRowBytes, pair_begin + kIntermediate, &shared.full[stage]); if (load_scales) { - nvfp4_tma_load_2d(tensors.a_scale4[(k_tile / 2) & 1], &descriptors.a_scales, - (k_tile / 2) * 16, token_begin, &shared.full[stage]); + // The box is tile-contiguous, so its address is a tile index rather than a + // (byte column, token row) pair; the two-slot buffer and the even-tile guard + // are unchanged. + constexpr int kScaleTilesPerPlane = Geometry::kGroupsPerRow / 16; + const int scale_tile = + (token_begin / Schedule::kBlockM) * kScaleTilesPerPlane + k_tile / 2; + nvfp4_tma_load_2d(tensors.a_scale4[(k_tile / 2) & 1], &descriptors.a_scales, 0, + scale_tile * 16, &shared.full[stage]); } const int gate_scale_row = ((pair_begin / 128) * Geometry::kScaleTilesPerRow +