From e52e1be3c750db748f829a048aa6b887799b8b25 Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Sun, 6 Sep 2026 20:24:47 +0000 Subject: [PATCH 1/2] perf(sparse-moe): keep two Q4 group quads in flight in the routed gate/up dot product dot_two_rows walks the routed gate/up K dimension in quads of four Q4 groups: one quad is one 128-byte warp transaction per row, and each lane owns eight consecutive K values so a single mantissa decode feeds eight FP32 FMAs. The loop issues one quad's code words and scales and then decodes them immediately, so for the whole of that decode and its sixteen FMAs the load unit has nothing outstanding. This commit issues both quads of a pair before either is decoded. The groups are still visited in ascending order, so the order of additions into acc0 and acc1 is unchanged. 50 insertions, 23 deletions, one file. The static_assert next to the constant pins the divisibility of the full K, not of the span the loop actually walks ([k_begin/kGroupK, k_end/kGroupK)); the two coincide only because every call site passes the full range, and a comment above the assertion says so rather than leaving it to be found. The change rests on a measurement of what limits this path. Holding the kernel, the token count, the grid, the warp count and the instruction count fixed and varying only the number of distinct experts the routing touches, the mandatory weight traffic moves 3.30x while the kernel moves 1.62x: on this artifact the routed gate/up path demands 30.8% of the measured 1689.4 GB/s DRAM read ceiling on a concentrated route, 55.1% on a trace-like one and 62.9% on an independent one, and extrapolated to zero experts it still costs 16.8 us. It is not a tensor-core path either - cuobjdump -sass over sparse_moe_decode_kernels.cu.o reports zero HMMA, IMMA or OMMA. What is left is memory latency, and one quad in flight does not cover it. Claim level: operator. The change is confined to one function body; it adds no template parameter, no instantiation and no launch-shape change, and the same body is also reached by the T=1 decode path (sparse_moe_d3_nine_warp_kernel), which is covered by the same gate. Measured on one RTX 5090 (sm_120a, CUDA 13.1.115, driver 580.159.03, 525 W cap, Release, -DCMAKE_CUDA_COMPILER=/usr/local/cuda-13.1/bin/nvcc), artifact Qwen3.6-35B-A3B (groupwise-int, codecs q4-q5), base ad0f3d38. Baseline and candidate are two binaries built from the same tree and differing only by this diff; no environment variable selects behaviour. Every measurement is a separate GPU-exclusive capture with its own card-witness window, and any window with a foreign process in it was discarded and re-taken: 13 attempts on 3 cells were, twelve of them in the research series that is not submitted, and none at all in the submitted series. The timestamp that opens the window is written before the run is handed to the lock wrapper, so the window also covers the wait for the card lock; the bias is one-sided and can only reject a good cell. Kernel medians come from nsys -t cuda --cuda-graph-trace=node, not from wall clock: the bench's wall clock for this operation lands on a 2.048 us lattice and the operation's own wall median is 55.296 us at trace-like, so one division is 3.7% of it - larger than the effect being measured. Operator, ninfer_sparse_moe_bench --codec q4-q5 --tokens 4 --distribution --execution graph --cache cold --warmup 3 --repeat 20 under nsys, two passes per arm with the arm order rotated between passes, both passes shown: distribution same (8 unique experts): 21472/21376 -> 20448/20352 ns -4.78% distribution trace-like (22 unique experts): 28768/28736 -> 26464/26784 ns -7.40% distribution independent(31 unique experts): 34624/34592 -> 32576/32448 ns -6.06% The depth itself was chosen by a dose rather than assumed. On a separate research build in which the depth is an environment-selected template argument, trace-like gives depth 1 (master) 28928/28832 ns, depth 2 26464/26559 ns, depth 4 27103/27039 ns: the optimum is interior. Two was then rebuilt as a constant and re-measured as the two binaries above. Two knobs that did not work are reported because they bound the claim. Shrinking the grid and running a grid-stride loop is not available here: the launch already issues exactly one CTA per unit of work (6144 = 4 tokens x 3 path blocks x 512 columns), and every narrower grid measured is worse - the same work split rewritten as a strided loop costs +2.22%, half the columns +19.72%, a quarter +65.59%. The two knobs together are worse than either alone (+53.63%). These four figures come from the same research build as the depth dose above, in which both knobs are environment-selected template arguments; none of that scaffolding is part of this commit. The same feed-depth change on the routed down kernel (sparse_moe_d4_token_kernel) is -0.15%, i.e. nothing, because its k loop has two iterations, not eight; it is not included here. Round level, same binary pair, nsys over ninfer_bench --weights qwen3_6_35b_a3b.ninfer -pg 1000,128 --prefill-chunk 8192 --spec mtp --draft-tokens 3 -r 1 --warmup 1, two arms x two passes: the routed gate/up kernel median 24416/24416 -> 23744/23744 ns (-2.75%) over 5240 launches, kernel total 130512610/130622538 -> 126452934/126528751 ns (-3.11/-3.13%). In the same traces the routed down kernel and the W8 draft instantiation of the same kernel are the controls and do not move. The operator fixture overstates this change by 2.69x: its best case is -7.40% on trace-like, the same kernel inside the real round moves -2.75%. The number that belongs in a decision is the round one. One regression, measured and reported rather than left out. dot_two_rows is shared with sparse_moe_d3_nine_warp_kernel, the T=1 decode path, and in the same two traces it gets slower: median 11136 -> 11872 ns (+6.61%) and 11200 -> 11888 ns (+6.14%), minimum 10656 -> 11392 (+6.91%) and 10560 -> 11456 (+8.48%), total over 40 launches 452610 -> 480832 (+6.24%) and 451263 -> 482238 (+6.86%). The minima move with the medians in both passes, so it is real. It costs 28222 and 30975 ns, 0.004% of the trace's kernel time and 0.7% of what the change saves on the routed path. Separating the two callers of the helper would remove it at no cost to the routed path. End to end, same binaries without a profiler, ninfer_bench -pg 1000,128 --prefill-chunk 8192 --spec mtp --draft-tokens 3 -r 4 --warmup 1, four passes with the arm order rotated: decode 414.25/413.36/413.73/413.11 -> 416.40/415.80/415.93/415.82 tok/s, mean 413.61 -> 415.99, +0.57%, with every candidate cell above every baseline cell; the spread inside an arm is 0.28% on the baseline and 0.14% on the candidate, so the baseline's own spread is half the size of the effect and it is the separation of the cells, not the spread, that is claimed. The prefill column of the same runs, which this change cannot reach because prefill dispatches different kernels, moves -0.17%, i.e. the stand's systematic works against the claim rather than for it. Speculation acceptance (0.3567567568) and round counts (248/0) are identical across all eight cells, so the two arms did the same work. An end-to-end number does not isolate an operator change; it is reported as corroboration, not as the claim. Numerics: the output is bit-identical. ninfer_sparse_moe_test with NINFER_OP_REPORT_STATS=1 over all sixteen registered cells (q4+q5 T=1,2,46,47,768,4097; q4+q6 T=1,2,46,47,768; w8+w8 T=1,2,19,20,768) produces byte-identical statistics on both binaries, and both report OK. That the oracle can move was shown with a sabotage build that rounds each decoded Q4 weight to BF16 before the FMA: it moves rel_l2 by 5.04-5.34x on exactly the four cells this path serves (q4+q5 and q4+q6 at T=2 and T=46) and leaves the other twelve byte-identical. That sabotage moves the statistic without crossing the 1.2e-2 limit, so what is demonstrated is the sensitivity of the number, not a failing verdict; this is sufficient here only because the claim is bit-identity, which is strictly stronger than passing the tolerance. Rebuilding the same honest source into a different binary reproduces all sixteen cells byte for byte, so the comparison is not a false alarm. Where that strength control does not reach: the four cells the sabotage moves are T=2 and T=46; both T=1 cells are unmoved by it. T=1 is exactly the sparse_moe_d3_nine_warp_kernel path this commit also alters. Bit-identity there is asserted by the same instrument, whose sensitivity has been demonstrated on the small-T routed path and not on the T=1 path. ctest was run in full on a clean ad0f3d38 and on this commit from the same build tree: 104 of 104 pass on both, with the same six opt-in real-artifact tests skipped on both. clang-format reports zero violations on the base file, ten on the first version of this change (all on lines it adds), and zero after reformatting; both 19.1.7 and 23.1.0 agree. The timings and the ctest run above were taken before the reformatting and carry over because the reformatted source compiles to an object whose SASS and -res-usage are identical once the translation-unit hash is normalised away. Resources, cuobjdump --dump-resource-usage over sparse_moe_decode_kernels.cu.o with the translation-unit hash normalised away and the bodies matched by demangled name with template arguments: 5 of 33 bodies change register count, all five Q4 instantiations of the paths this commit touches, and shared memory, LOCAL and STACK change on none of them. The instantiation the artifact selects at mtp3, sparse_moe_d3_path_tiled_kernel, is 40 registers before and after. The other four go 39 -> 40, 39 -> 40, 39 -> 40 and 37 -> 39, and sparse_moe_d3_nine_warp_kernel goes 37 -> 40. At 96 threads per block, 40 registers allow 17 resident blocks per SM against the 16 the thread limit allows, so occupancy stays limited by threads and shared memory, not by registers. Every W8 instantiation and every sparse_moe_d4 instantiation is unchanged. Workspace is unchanged. Checks not run, and the limitation each leaves: the gate's sensitivity on the T=1 path is not demonstrated, only its byte-for-byte agreement; the +6% regression on sparse_moe_d3_nine_warp_kernel is reported and not fixed; no accuracy evaluation was run, which is acceptable only because the output is bit-identical; the change was measured on one artifact and one codec pair (q4-q5, with q4-q6 covered by the gate but not by the operator bench), so the effect on q4-q6 and on w8-w8 is asserted from the code path taken rather than measured; and no profiler counters were available on this machine, so every attribution above rests on ablations and on byte budgets computed from geometry, not on measured DRAM traffic - the percentages of the DRAM ceiling are demand, not a counter reading. Base and provenance. This change now targets upstream 487f8977; it was written and measured on ad0f3d38, and nothing in this message has been re-measured on the new base. Two instrument changes came with the rebase. The bench flag --mtp-draft-tokens no longer exists: the spelling is --spec mtp --draft-tokens N, and --draft-tokens 0 is rejected for MTP (src/product/speculative_options.h:41), so the zero-draft arm is spelled by passing neither flag. And ninfer_bench now disables the engine context cache itself (bench/targets/qwen3_6_27b/ninfer_bench.cpp:157), where on ad0f3d38 that line did not exist and the cache stayed at its enabled default. Per-request prefix reuse was already off on both bases (same file, line 65), so no request here could take a prefix hit; how far a constructed-but-unused cache reaches into prefill and TTFT has not been measured, and these figures describe a configuration the current bench no longer builds. tests/CMakeLists.txt registers 114 targets on 487f8977 rather than 104, with 20 named targets explicitly artifact-gated rather than 19; no ctest run quoted here was repeated on the new base, and no round "114 of 114" is claimed for it: ninfer_attn_input_proj_test is red on the bare base (upstream defect, issue #196), and ctest and a direct run of the same binary have been seen to disagree on this host. --- .../decode/sparse_moe_decode_kernels.cu | 73 +++++++++++++------ 1 file changed, 50 insertions(+), 23 deletions(-) diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 12a68e0c18..831e626f82 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -168,32 +168,59 @@ __device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const st // Four adjacent Q4 groups form one 128-byte warp transaction. Each lane owns eight // consecutive K values, so one mantissa decode feeds eight FP32 FMAs instead of issuing // four scalar code-pair/decode iterations. + // + // Two such quads are in flight at a time. Both quads' code words and scales are issued + // before either is decoded, so the decode and the FMA chain of the first quad cover the + // memory latency of the second. One quad in flight leaves the load unit idle for the + // whole decode, and this path is latency-bound rather than bandwidth-bound: at the + // operator bench's trace-like point it demands 55% of the measured DRAM read ceiling. + // + // The groups are still visited in ascending order, so the order of additions into acc0 + // and acc1 is unchanged and the output is bit-identical. + constexpr int kQuadsInFlight = 2; + // What the pairing needs is that the span this loop walks, (k_end - k_begin) / + // Codec::kGroupK, is a multiple of 4 * kQuadsInFlight. That is a property of the + // arguments and cannot be asserted here; every caller passes k_begin = 0 and + // k_end = K, which reduces it to the compile-time condition below. + static_assert(kGroups % (4 * kQuadsInFlight) == 0); const int lane_group = lane >> 3; const int lane_in_group = lane & 7; - for (int group_base = first_group; group_base < last_group; group_base += 4) { - const int group = group_base + lane_group; - const std::int64_t index0 = static_cast(row0) * kGroups + group; - const std::int64_t index1 = static_cast(row1) * kGroups + group; - const std::uint32_t packed0 = - *reinterpret_cast(codes + index0 * 32 + lane_in_group * 4); - const std::uint32_t packed1 = - *reinterpret_cast(codes + index1 * 32 + lane_in_group * 4); - const auto scale0 = *reinterpret_cast(scales + index0 * 2); - const auto scale1 = *reinterpret_cast(scales + index1 * 2); - float weights0[8]; - float weights1[8]; - Q4SimtDecodeAtom::decode_eight(packed0, scale0, weights0); - Q4SimtDecodeAtom::decode_eight(packed1, scale1, weights1); - const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); - const float2 x0 = bf16x2_bits_to_float2(input.x); - const float2 x1 = bf16x2_bits_to_float2(input.y); - const float2 x2 = bf16x2_bits_to_float2(input.z); - const float2 x3 = bf16x2_bits_to_float2(input.w); - const float values[8] = {x0.x, x0.y, x1.x, x1.y, x2.x, x2.y, x3.x, x3.y}; + for (int group_base = first_group; group_base < last_group; + group_base += 4 * kQuadsInFlight) { + std::uint32_t packed0[kQuadsInFlight]; + std::uint32_t packed1[kQuadsInFlight]; + std::uint16_t scale0[kQuadsInFlight]; + std::uint16_t scale1[kQuadsInFlight]; +#pragma unroll + for (int quad = 0; quad < kQuadsInFlight; ++quad) { + const int group = group_base + quad * 4 + lane_group; + const std::int64_t index0 = static_cast(row0) * kGroups + group; + const std::int64_t index1 = static_cast(row1) * kGroups + group; + packed0[quad] = *reinterpret_cast(codes + index0 * 32 + + lane_in_group * 4); + packed1[quad] = *reinterpret_cast(codes + index1 * 32 + + lane_in_group * 4); + scale0[quad] = *reinterpret_cast(scales + index0 * 2); + scale1[quad] = *reinterpret_cast(scales + index1 * 2); + } +#pragma unroll + for (int quad = 0; quad < kQuadsInFlight; ++quad) { + const int group = group_base + quad * 4 + lane_group; + float weights0[8]; + float weights1[8]; + Q4SimtDecodeAtom::decode_eight(packed0[quad], scale0[quad], weights0); + Q4SimtDecodeAtom::decode_eight(packed1[quad], scale1[quad], weights1); + const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); + const float2 x0 = bf16x2_bits_to_float2(input.x); + const float2 x1 = bf16x2_bits_to_float2(input.y); + const float2 x2 = bf16x2_bits_to_float2(input.z); + const float2 x3 = bf16x2_bits_to_float2(input.w); + const float values[8] = {x0.x, x0.y, x1.x, x1.y, x2.x, x2.y, x3.x, x3.y}; #pragma unroll - for (int item = 0; item < 8; ++item) { - acc0 = fmaf(weights0[item], values[item], acc0); - acc1 = fmaf(weights1[item], values[item], acc1); + for (int item = 0; item < 8; ++item) { + acc0 = fmaf(weights0[item], values[item], acc0); + acc1 = fmaf(weights1[item], values[item], acc1); + } } } } else if constexpr (Codec::kD3SingleValuePerLane) { From a04b5fe40175f34c5266f2b9171d01ede32c383c Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:28:58 +0200 Subject: [PATCH 2/2] perf(sparse_moe): keep one quad in flight on the T = 1 decode path Pairing the Q4 group quads wins on the routed path-tiled kernel and loses 6-11% on sparse_moe_d3_nine_warp_kernel, which serves every token of non-speculative decode. Give the two callers separate depths: dot_two_rows takes QuadsInFlight, the nine-warp caller passes 1. Its SASS is then byte-identical to master's. Co-Authored-By: Claude Opus 5 --- .../decode/sparse_moe_decode_kernels.cu | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 831e626f82..6933be7989 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -153,7 +153,7 @@ struct W8Codec { } }; -template +template __device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const std::uint8_t* high, const std::uint8_t* scales, int row0, int row1, const __nv_bfloat16* x, int k_begin, int k_end, @@ -169,15 +169,20 @@ __device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const st // consecutive K values, so one mantissa decode feeds eight FP32 FMAs instead of issuing // four scalar code-pair/decode iterations. // - // Two such quads are in flight at a time. Both quads' code words and scales are issued - // before either is decoded, so the decode and the FMA chain of the first quad cover the - // memory latency of the second. One quad in flight leaves the load unit idle for the - // whole decode, and this path is latency-bound rather than bandwidth-bound: at the - // operator bench's trace-like point it demands 55% of the measured DRAM read ceiling. + // QuadsInFlight quads are in flight at a time. All of their code words and scales are + // issued before any is decoded, so the decode and the FMA chain of one quad cover the + // memory latency of the next. One quad in flight leaves the load unit idle for the whole + // decode, which pays off where the path is latency-bound rather than bandwidth-bound: at + // the operator bench's trace-like point it demands 55% of the measured DRAM read ceiling. // - // The groups are still visited in ascending order, so the order of additions into acc0 - // and acc1 is unchanged and the output is bit-identical. - constexpr int kQuadsInFlight = 2; + // The two callers want different depths and get them. The T >= 2 path-tiled kernel takes + // the default 2; sparse_moe_d3_nine_warp_kernel, the T = 1 decode path, keeps 1, because + // at one token per launch it has a fifth of the occupancy and pairing costs it 6% in the + // round and 10% on the operator bench. + // + // The groups are still visited in ascending order at either depth, so the order of + // additions into acc0 and acc1 is unchanged and the output is bit-identical. + constexpr int kQuadsInFlight = QuadsInFlight; // What the pairing needs is that the span this loop walks, (k_end - k_begin) / // Codec::kGroupK, is a multiple of 4 * kQuadsInFlight. That is a property of the // arguments and cannot be asserted here; every caller passes k_begin = 0 and @@ -273,9 +278,9 @@ __global__ void sparse_moe_d3_nine_warp_kernel( pdl::wait_for_dependencies(); const int expert = ids[warp]; const int row_base = expert * 1024; - dot_two_rows(routed_codes, routed_high, routed_scales, row_base + j, - row_base + kIntermediate + j, x_shared, 0, kHidden, gate, - up); + dot_two_rows(routed_codes, routed_high, routed_scales, + row_base + j, row_base + kIntermediate + j, x_shared, + 0, kHidden, gate, up); } else { dot_two_rows(shared_codes, nullptr, shared_scales, j, kIntermediate + j, x_shared, 0, kHidden, gate, up);