chore(train): engine-10 continued: llama-bench's generation range, the rotations' weights before the PDL wait

TORAD.md

@@ -97,6 +97,8 @@ which pins a commit of this branch as a submodule.
 | `818965d72` | the mma flash attention compiled its tile code once per stream-k role (a block that ends inside a tile, one that ends a tile it did not start, one that does a whole tile) as template parameters, though only the epilogue's destination depends on the role: Ternary Bonsai 2 27B's decode kernel (head 256, q4_0 K/V read raw) was 184 KB of SASS in three copies. A per-block timeline on an RTX 5080 (252 blocks, 63 a KV head) had the 4 blocks that end a tile ending last, 8-9 us after the median block, their prologues 7.8-10.5 us against 3.7-4.2: only they ran the tile-ending copy, cold in every cache after 75-283 MB of K/V through the 64 MB L2. The roles are arguments of one copy (71 KB, 255 registers, 16 bytes of stack against 24; libggml-cuda 79.1 -> 73.2 MB). test-backend-ops perf, the served layout, 4 rounds against the parent's library: a token -2.8 / -4.2 / -2.2 / -0.9 % and a 4-row verify +0.4 / -3.1 / -1.0 / -1.1 % at 16,384 / 65,536 / 131,072 / 245,760 cells; llama-bench at 16,384 inside its noise. FLASH_ATTN_EXT 3,220/3,220; the served decode's texts identical | none: the same arithmetic and stores, only the code's layout changed |
 | `6aa14679c` | the live-tile split (serving with more than one slot and `--kv-unified`) ran a memset and a scan of the mask before every flash attention, 3.2-4.5 us a call in the graph on an RTX 5080, though the live steps depend on the mask and the split alone and every attention layer of a graph reads the same mask (Ternary Bonsai 2 27B: 16 layers). `ggml_cuda_fattn_kv_live_context` keeps a graph evaluation's first scan, keyed by the mask tensor and its data, shape and strides, the stream and the split; the other layers read it and any other key scans again; it resets where the graph evaluator resets its other per-graph records, so a CUDA graph captures one scan. test-backend-ops perf (its graph repeats the op: only the first copy scans), 4 rounds: a token -13.2 / -3.6 / -2.0 / -1.1 % and a 4-row verify -13.5 / -3.7 / -1.2 / -0.7 % at 16,384 / 65,536 / 131,072 / 245,760 cells (3.3-4.2 us a call). FLASH_ATTN_EXT 3,220/3,220; the served decode at 15,616 tokens, 4 slots: the same tokens in every leg, plain and drafted, its rates inside a loaded host's noise | `GGML_CUDA_FATTN_LIVE_SCAN_EACH_LEGACY=1` |
 | `32e695ecf` | the live-tile split launched every block the SMs hold, where the plain stream-k split rounds down to a multiple of the output tiles. 01f4f0fda's 3 blocks an SM made a decode token's split 510 blocks on an RTX 5090 (127.5 per KV head) and 210 on an RTX 5070 Ti (52.5): each head's blocks started half a block's steps from its neighbour's. A cell's K and V rows hold the heads side by side (q4_0, head 256: 144 bytes each) and the L2 fetches 64-byte units, so neighbouring heads share one, fetched once only when they read the row together: 12 units a row where 9 hold it. ncu on the 5070 Ti at 131,072 cells: DRAM read 202-207 MB against 151.0 MB of K and V; rounded down to a multiple of the Q tile's output tiles (208 blocks), 152-153 MB. test-backend-ops perf, a token, the same library with and without the switch: 5090 -15.7 / -7.8 / -17.0 / -20.9 %, 5070 Ti -2.1 / -18.3 / -21.4 / -23.1 % at 16,384 / 65,536 / 131,072 / 245,760 cells; on the 5090 131,072 and 245,760 cells take 115.72 and 196.44 us, where 4104c47 (340 blocks) took 123.69 and 204.66 and 4b61c54 145.50 and 255.48. A 4-row verify (2 blocks an SM: 340 and 140 blocks) does not move; the 5080's 252 is a multiple already. The served decode on the 5090 at 245K tokens (the gate's pool, 4 x 294,912 with `--kv-unified`, its four questions, 3 legs each): plain 98.75 -> 107.92 tok/s (+9.3 %), drafted 205.43 -> 208.90 (+1.7 %), the texts first differing 35-132 tokens in: the split sums in another order. Held to engine-10's G1 bars at that shape with top-5 log-probabilities, 384 tokens: plain 4 first differences, each at a tie, |dlogprob| p99 0.095 and max 0.128; drafted 4, each at a tie, p99 0.070, max 0.078. The gate's one-slot legs prefill through the live-tile path (the mask-prefix hint covers at most 16 rows), so their bits move too; their decode takes the mask-prefix path, whose split was rounded already, and its rate does not move (101.5 against 100.7 tok/s at 245K). FLASH_ATTN_EXT 3,220/3,220 on the 5070 Ti and the 5080 | `GGML_CUDA_FATTN_LIVE_BLOCKS_ALIGN_LEGACY=1` |
+| `e75e318e5` | `llama-bench` puts an NVTX range `gen` (domain `llama-bench`, a registered string) around each repetition's generation, so `nsys profile -c nvtx -p gen@llama-bench --capture-range-end=stop` records the tokens and none of the depth's prefill before them; a build without the NVTX headers compiles it out. rig's roofline probe traced the whole run, and at a depth of 245,760 on an RTX 5080 (Ternary Bonsai 2 27B, graphs off, `GGML_CUDA_NVTX=1`) the prefill's ~1.19M kernels kept nsys past its 30-minute limit; started at the range the capture takes 179 s, a 2.3 MB report of the 16 tokens' 18,816 kernels. At 16,384 it reads what the full capture reads: kernels 9,553.6 against 9,559.8 us a token, the same 23 groups, the largest group apart by 6.3 us | (a profiler that does not ask for the range ignores it) |
+| `fb03d294b` | `rms_norm_fwht_cuda` and `fwht_cuda_block` request their weights (the norm's `w` and the signs, which no kernel writes) before the PDL dependency wait, as the PQ2_0 matmul does its own. The matmul after a rotation starts under it (PDL), and its ring fill and the L2 prefetch of its head (`16ad036c4`) take DRAM for the microseconds the rotation runs; read after the wait, the rotation's weights (DRAM misses each token) queued behind them. RTX 5080, Ternary Bonsai 2 27B, one build both ways: nsys over 15 tokens (graphs off, the capture started at `e75e318e5`'s range), the FFN-input norm after an out-projection ends 6.1-6.2 -> 4.7-4.8 us past it at 16,384 and 5.5-6.1 -> 4.5-4.8 at 245,760, the layer-input norm 2.1 -> 1.7 us; llama-bench with graphs on, legs N-L-L-N: tg128 at 16,384 +0.81 % (105.48 / 104.21 / 104.40 / 104.82 tok/s), tg64 at 245,760 +1.63 % (69.35 / 67.93 / 67.48 / 68.27), all four pairs faster; `test-backend-ops MUL_MAT_HADAMARD` 42/42 both ways; llama-server at one slot and at 4 slots with `--kv-unified`, 128 greedy tokens with top-5 log-probabilities: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_FWHT_PREWAIT_LEGACY=1` |
 
 Every switch in the last column is read once per process and parses as an integer: a `*_LEGACY` switch set to `0`
 is the same as unset (the change stays on), and `=0` turns off `GGML_CUDA_LORA_RANK1_FUSE` and

ggml/src/ggml-cuda/fwht.cu

@@ -215,10 +215,12 @@ struct fwht_src_view {
     int64_t s0, s1, s2, s3;
 };
 
+// prewait: the signs, which no kernel writes, are requested before the dependency wait (fwht_prewait)
 template <int N, int NT, typename T, bool has_signs, bool has_view = false>
 __launch_bounds__(NT, 1)
 __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows, const float scale,
-                                const float * signs, const int n_blk, const bool pdl_trigger, const fwht_src_view view) {
+                                const float * signs, const int n_blk, const bool pdl_trigger, const bool prewait,
+                                const fwht_src_view view) {
     if (pdl_trigger) {
         ggml_cuda_pdl_lc();
     }
@@ -238,8 +240,21 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows
 
     const int tid = threadIdx.x;
 
-    ggml_cuda_pdl_sync();
     const float * signs_row = has_signs ? signs + (r % n_blk) * N : nullptr;
+    float sg[NE];
+    const auto load_signs = [&]() {
+#pragma unroll
+        for (int i = 0; i < NE; ++i) {
+            sg[i] = signs_row[i * NT + tid];
+        }
+    };
+    if (has_signs && prewait) {
+        load_signs();
+    }
+    ggml_cuda_pdl_sync();
+    if (has_signs && !prewait) {
+        load_signs();
+    }
 
     float reg[NE];
 #pragma unroll
@@ -255,7 +270,7 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows
             reg[i] = fwht_load(src[i * NT + tid]) * scale;
         }
         if (has_signs) {
-            reg[i] *= signs_row[i * NT + tid];
+            reg[i] *= sg[i];
         }
     }
 
@@ -265,10 +280,12 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows
 
 // rms_norm with its weight multiply (as rms_norm_f32<1024, true>), the sign flip and the transform in one launch.
 // Block (c, t) reduces token t's whole row as rms_norm does, then transforms its chunk c. normed: the multiply's result or nullptr.
+// prewait: w and the signs, which no kernel writes, are requested before the dependency wait (fwht_prewait).
 template <int N>
 __launch_bounds__(1024, 1)
 __global__ void rms_norm_fwht_cuda(const float * x, const float * w, const float * signs, float * normed, float * dst,
-                                   const int ncols, const float eps, const float scale, const bool pdl_trigger) {
+                                   const int ncols, const float eps, const float scale, const bool pdl_trigger,
+                                   const bool prewait) {
     if (pdl_trigger) {
         ggml_cuda_pdl_lc();
     }
@@ -284,7 +301,23 @@ __global__ void rms_norm_fwht_cuda(const float * x, const float * w, const float
 
     float tmp = 0.0f;
 
+    float wr[NE];
+    float sg[NE];
+    const auto load_weights = [&]() {
+#pragma unroll
+        for (int i = 0; i < NE; ++i) {
+            const int col = blockIdx.x * N + i * NT + tid;
+            wr[i] = w[col];
+            sg[i] = signs[col];
+        }
+    };
+    if (prewait) {
+        load_weights();
+    }
     ggml_cuda_pdl_sync();
+    if (!prewait) {
+        load_weights();
+    }
     for (int col = tid; col < ncols; col += NT) {
         const float xi = x[col];
         tmp += xi * xi;
@@ -298,12 +331,12 @@ __global__ void rms_norm_fwht_cuda(const float * x, const float * w, const float
 #pragma unroll
     for (int i = 0; i < NE; ++i) {
         const int   col = blockIdx.x * N + i * NT + tid;
-        const float v   = rms_scale * x[col] * w[col];
+        const float v   = rms_scale * x[col] * wr[i];
         if (normed != nullptr) {
             normed[e0 + i * NT + tid] = v;
         }
         reg[i] = v * scale;
-        reg[i] *= signs[col];
+        reg[i] *= sg[i];
     }
 
     fwht_block_transform<N, NT>(reg, s);
@@ -324,6 +357,15 @@ static bool fwht_pdl_trigger() {
     return pdl_trigger;
 }
 
+// rms_norm_fwht_cuda and fwht_cuda_block request their weights (the norm's w and the signs, which no kernel writes)
+// before the dependency wait, so they arrive while the kernel before them ends. The matmul after a rotation starts under
+// it (PDL), and its ring fill and the L2 prefetch of its head (ggml_cuda_pq2_prefetch) take DRAM for the microseconds the
+// rotation runs: read after the wait, the weights queued behind them. GGML_CUDA_FWHT_PREWAIT_LEGACY=1 reads them after.
+static bool fwht_prewait() {
+    static const bool prewait = !ggml_env_switch("GGML_CUDA_FWHT_PREWAIT_LEGACY");
+    return prewait;
+}
+
 template <typename T>
 static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float * dst_d,
                         const int n, const int64_t rows, const float scale,
@@ -340,6 +382,7 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float
         ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream);
 
     const bool pdl_trigger = fwht_pdl_trigger();
+    const bool prewait     = fwht_prewait();
 
     switch (n) {
 #define FWHT_CASE(NN) \
@@ -376,14 +419,14 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float
             const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(g, b, 0, stream); \
             if constexpr (std::is_same_v<T, float>) { \
                 if (view) { \
-                    ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, true, true>, lp, src_d, dst_d, rows, scale, signs, n_blk, pdl_trigger, *view); \
+                    ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, true, true>, lp, src_d, dst_d, rows, scale, signs, n_blk, pdl_trigger, prewait, *view); \
                     return true; \
                 } \
             } \
             if (signs) { \
-                ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, true>,  lp, src_d, dst_d, rows, scale, signs, n_blk, pdl_trigger, fwht_src_view{}); \
+                ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, true>,  lp, src_d, dst_d, rows, scale, signs, n_blk, pdl_trigger, prewait, fwht_src_view{}); \
             } else { \
-                ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, false>, lp, src_d, dst_d, rows, scale, nullptr, 1, pdl_trigger, fwht_src_view{}); \
+                ggml_cuda_kernel_launch(fwht_cuda_block<NN, FWHT_BLOCK_THREADS, T, false>, lp, src_d, dst_d, rows, scale, nullptr, 1, pdl_trigger, prewait, fwht_src_view{}); \
             } \
             return true; \
         }
@@ -508,18 +551,19 @@ void ggml_cuda_op_rms_norm_fwht(ggml_backend_cuda_context & ctx, const ggml_tens
     const float   scale    = 1 / sqrtf(n);
 
     const bool pdl_trigger = fwht_pdl_trigger();
+    const bool prewait     = fwht_prewait();
 
     const dim3 grid(ncols / n, ntok, 1), block(1024, 1, 1);
     const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(grid, block, 0, ctx.stream());
     switch (n) {
         case 1024:
-            ggml_cuda_kernel_launch(rms_norm_fwht_cuda<1024>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger);
+            ggml_cuda_kernel_launch(rms_norm_fwht_cuda<1024>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait);
             break;
         case 2048:
-            ggml_cuda_kernel_launch(rms_norm_fwht_cuda<2048>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger);
+            ggml_cuda_kernel_launch(rms_norm_fwht_cuda<2048>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait);
             break;
         default:
-            ggml_cuda_kernel_launch(rms_norm_fwht_cuda<4096>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger);
+            ggml_cuda_kernel_launch(rms_norm_fwht_cuda<4096>, lp, x_d, w_d, signs_d, normed_d, dst_d, ncols, eps, scale, pdl_trigger, prewait);
             break;
     }
 }

tools/llama-bench/CMakeLists.txt

@@ -8,6 +8,14 @@ set_target_properties(${TARGET} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)
 target_include_directories(${TARGET} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 target_link_libraries(${TARGET} PUBLIC llama-common llama ${CMAKE_THREAD_LIBS_INIT})
 
+# the generation's NVTX range (llama-bench.cpp): NVTX is header-only, and a CUDA build has it in the toolkit
+if (GGML_CUDA)
+    find_package(CUDAToolkit)
+    if (TARGET CUDA::nvtx3)
+        target_link_libraries(${TARGET} PRIVATE CUDA::nvtx3)
+    endif()
+endif()
+
 if(LLAMA_TOOLS_INSTALL)
     install(TARGETS ${TARGET} LIBRARY)
 endif()

tools/llama-bench/llama-bench.cpp

@@ -36,12 +36,44 @@
 #    include <windows.h>
 #endif
 
+#if __has_include(<nvtx3/nvToolsExt.h>)
+#    include <nvtx3/nvToolsExt.h>
+#    define LLAMA_BENCH_NVTX
+#endif
+
 // utils
 static uint64_t get_time_ns() {
     using clock = std::chrono::high_resolution_clock;
     return std::chrono::nanoseconds(clock::now().time_since_epoch()).count();
 }
 
+// An NVTX range "gen" in the domain "llama-bench" around each repetition's generation, so a profiler can capture the
+// tokens and none of the depth prefill before them: nsys profile -c nvtx -p gen@llama-bench --capture-range-end=stop
+// records the first repetition's generation (Ternary Bonsai 2 27B at a depth of 245,760: 16 tokens launch 18,816 kernels,
+// the prefill before them 1.19M). The message is a registered string, which nsys matches without NSYS_NVTX_PROFILER_REGISTER_ONLY=0. Without
+// the NVTX headers (a build without the CUDA toolkit) the range compiles out.
+struct bench_phase_range {
+#ifdef LLAMA_BENCH_NVTX
+    static nvtxDomainHandle_t domain() {
+        static const nvtxDomainHandle_t d = nvtxDomainCreateA("llama-bench");
+        return d;
+    }
+
+    explicit bench_phase_range(const char * name) {
+        nvtxEventAttributes_t attr = {};
+        attr.version            = NVTX_VERSION;
+        attr.size               = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
+        attr.messageType        = NVTX_MESSAGE_TYPE_REGISTERED;
+        attr.message.registered = nvtxDomainRegisterStringA(domain(), name);
+        nvtxDomainRangePushEx(domain(), &attr);
+    }
+
+    ~bench_phase_range() { nvtxDomainRangePop(domain()); }
+#else
+    explicit bench_phase_range(const char *) {}
+#endif
+};
+
 static bool tensor_buft_override_equal(const llama_model_tensor_buft_override& a, const llama_model_tensor_buft_override& b) {
     if (a.pattern != b.pattern) {
         // cString comparison that may be null
@@ -2508,6 +2540,7 @@ int llama_bench(int argc, char ** argv) {
                     fprintf(stderr, "llama-bench: benchmark %d/%zu: generation run %d/%d\n", params_idx, params_count,
                             i + 1, params.reps);
                 }
+                const bench_phase_range gen_range("gen");
                 bool res = test_gen(ctx, t.n_gen, t.n_threads);
                 if (!res) {
                     fprintf(stderr, "%s: error: failed to run gen\n", __func__);