aprender-serve 0.64.0

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447

/// PARITY-082b: Wave quantization problem
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_082b_wave_quantization() {
    println!("PARITY-082b: Wave Quantization Problem");
    println!("=======================================");
    println!();

    // The problem Stream-K solves
    let sms = 128u32; // RTX 4090

    // Example: 200 output tiles
    let total_tiles = 200u32;
    let full_waves = total_tiles / sms;
    let remainder = total_tiles % sms;

    println!("  Traditional Tile-Centric (200 tiles, 128 SMs):");
    println!("  -----------------------------------------------");
    println!("    Full waves: {}", full_waves);
    println!("    Remainder tiles: {}", remainder);
    println!(
        "    Last wave utilization: {:.1}%",
        remainder as f32 / sms as f32 * 100.0
    );
    println!();

    // Efficiency loss
    let total_sm_slots = (full_waves + 1) * sms;
    let utilization = total_tiles as f32 / total_sm_slots as f32 * 100.0;
    let waste = total_sm_slots - total_tiles;

    println!("  Efficiency Analysis:");
    println!(
        "    Total SM slots used: {} ({} waves × {} SMs)",
        total_sm_slots,
        full_waves + 1,
        sms
    );
    println!("    Actual tiles: {}", total_tiles);
    println!("    Wasted slots: {}", waste);
    println!("    Overall utilization: {:.1}%", utilization);
    println!();

    // Stream-K solution
    println!("  Stream-K Solution:");
    println!("  -------------------");
    println!("    • Divide K dimension into segments");
    println!("    • Each SM processes multiple segments");
    println!("    • Final reduction combines partial results");
    println!("    • Achieves ~100% utilization");

    assert!(
        utilization < 95.0,
        "PARITY-082b: Traditional has poor utilization on irregular sizes"
    );
    assert!(true, "PARITY-082b: Wave quantization documented");
}

/// PARITY-082c: Work-stealing implementation
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_082c_work_stealing() {
    println!("PARITY-082c: Work-Stealing Implementation");
    println!("==========================================");
    println!();

    println!("  Work Queue Structure:");
    println!("  ----------------------");
    println!("    __device__ int global_tile_idx;  // Atomic counter");
    println!();
    println!("    __global__ void streamk_gemm(...) {{");
    println!("        while (true) {{");
    println!("            int tile = atomicAdd(&global_tile_idx, 1);");
    println!("            if (tile >= total_tiles) break;");
    println!("            ");
    println!("            // Compute tile coordinates");
    println!("            int tile_m = tile / tiles_n;");
    println!("            int tile_n = tile % tiles_n;");
    println!("            ");
    println!("            // Process all K iterations for this tile");
    println!("            compute_tile(tile_m, tile_n);");
    println!("        }}");
    println!("    }}");
    println!();

    // Atomic overhead analysis
    let tiles_per_sm = 10u32; // Average tiles per SM
    let atomic_latency_cycles = 100u32; // Approximate
    let compute_cycles_per_tile = 50000u32; // Approximate

    let overhead_pct = (atomic_latency_cycles * tiles_per_sm) as f32
        / (compute_cycles_per_tile * tiles_per_sm) as f32
        * 100.0;

    println!("  Atomic Overhead Analysis:");
    println!("  --------------------------");
    println!("    Tiles per SM: {}", tiles_per_sm);
    println!("    Atomic latency: ~{} cycles", atomic_latency_cycles);
    println!("    Compute per tile: ~{} cycles", compute_cycles_per_tile);
    println!("    Overhead: {:.2}%", overhead_pct);
    println!();

    println!("  Optimization: Tile Batching");
    println!("  ---------------------------");
    println!("    • Each SM claims batch of tiles (e.g., 4)");
    println!("    • Reduces atomic contention 4x");
    println!("    • Still maintains load balance");

    assert!(
        overhead_pct < 1.0,
        "PARITY-082c: Atomic overhead should be <1%"
    );
    assert!(true, "PARITY-082c: Work-stealing documented");
}

/// PARITY-082d: Partial result accumulation
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_082d_partial_accumulation() {
    println!("PARITY-082d: Partial Result Accumulation");
    println!("=========================================");
    println!();

    // Stream-K splits K dimension across CTAs
    // Need to combine partial results

    println!("  K-Splitting Strategy:");
    println!("  ----------------------");
    println!("    • Split K dimension into segments");
    println!("    • Each CTA computes C_partial = A_segment × B_segment");
    println!("    • Final: C = Σ C_partial");
    println!();

    let k = 4096u32;
    let tile_k = 32u32;
    let k_splits = 4u32;
    let k_per_split = k / k_splits;
    let iters_per_split = k_per_split / tile_k;

    println!("  Example (K={}, {} splits):", k, k_splits);
    println!("    K per split: {}", k_per_split);
    println!("    Tile-K iterations per split: {}", iters_per_split);
    println!();

    // Reduction strategies
    println!("  Reduction Strategies:");
    println!("  ----------------------");
    println!("  1. Global Memory Atomics:");
    println!("     atomicAdd(&C[i][j], partial);");
    println!("     Pro: Simple");
    println!("     Con: High contention for small tiles");
    println!();
    println!("  2. Two-Phase Reduction:");
    println!("     Phase 1: Write partials to scratch");
    println!("     Phase 2: Dedicated reduction kernel");
    println!("     Pro: No atomics");
    println!("     Con: Extra memory, kernel launch");
    println!();
    println!("  3. Cooperative Groups:");
    println!("     grid.sync() between compute and reduce");
    println!("     Pro: Single kernel");
    println!("     Con: Requires cooperative launch");

    // Memory for partials
    let m = 1024u32;
    let n = 768u32;
    let partial_mem = m * n * k_splits * 4; // F32

    println!();
    println!("  Partial Storage ({}×{}, {} splits):", m, n, k_splits);
    println!("    Memory: {} MB", partial_mem / 1024 / 1024);

    assert!(
        k_splits >= 2,
        "PARITY-082d: K-splitting requires at least 2 splits"
    );
    assert!(true, "PARITY-082d: Partial accumulation documented");
}

/// PARITY-082e: Tile rasterization order
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_082e_tile_rasterization() {
    println!("PARITY-082e: Tile Rasterization Order");
    println!("======================================");
    println!();

    // Tile ordering affects cache efficiency
    println!("  Rasterization Orders:");
    println!("  -----------------------");
    println!();

    // Row-major
    println!("  1. Row-Major (default):");
    println!("     ┌───┬───┬───┬───┐");
    println!("     │ 0 │ 1 │ 2 │ 3 │");
    println!("     ├───┼───┼───┼───┤");
    println!("     │ 4 │ 5 │ 6 │ 7 │");
    println!("     └───┴───┴───┴───┘");
    println!("     Con: Poor B-matrix locality");
    println!();

    // Morton/Z-order
    println!("  2. Morton Order (Z-curve):");
    println!("     ┌───┬───┬───┬───┐");
    println!("     │ 0 │ 1 │ 4 │ 5 │");
    println!("     ├───┼───┼───┼───┤");
    println!("     │ 2 │ 3 │ 6 │ 7 │");
    println!("     └───┴───┴───┴───┘");
    println!("     Pro: Better 2D locality");
    println!();

    // Swizzled
    println!("  3. Swizzled (Stream-K default):");
    println!("     Tiles assigned based on SM topology");
    println!("     Consecutive SMs get spatially close tiles");
    println!("     Maximizes L2 cache hits");
    println!();

    // Cache analysis
    let l2_cache = 72 * 1024 * 1024u64; // RTX 4090: 72 MB
    let tile_m = 128u32;
    let tile_n = 128u32;
    let tile_a_size = tile_m as u64 * 4096 * 4; // A tile row
    let tile_b_size = 4096u64 * tile_n as u64 * 4; // B tile column

    println!("  L2 Cache Analysis:");
    println!("    RTX 4090 L2: {} MB", l2_cache / 1024 / 1024);
    println!("    A tile ({}×K): {} KB", tile_m, tile_a_size / 1024);
    println!("    B tile (K×{}): {} KB", tile_n, tile_b_size / 1024);
    println!(
        "    Tiles fitting in L2: ~{}",
        l2_cache / (tile_a_size + tile_b_size)
    );

    assert!(true, "PARITY-082e: Tile rasterization documented");
}

/// PARITY-082f: Stream-K summary
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_082f_streamk_summary() {
    println!("PARITY-082f: Stream-K Summary");
    println!("==============================");
    println!();

    println!("  ╔═══════════════════════════════════════════════════════════════╗");
    println!("  ║          PARITY-082: Stream-K Decomposition Complete          ║");
    println!("  ╠═══════════════════════════════════════════════════════════════╣");
    println!("  ║                                                               ║");
    println!("  ║  Key Concepts:                                                ║");
    println!("  ║  ─────────────                                                ║");
    println!("  ║  1. Work-centric vs tile-centric decomposition               ║");
    println!("  ║  2. Global work queue with atomic tile claiming              ║");
    println!("  ║  3. K-splitting for partial result accumulation              ║");
    println!("  ║  4. Swizzled rasterization for cache efficiency              ║");
    println!("  ║                                                               ║");
    println!("  ║  Performance Benefits:                                        ║");
    println!("  ║  ────────────────────                                         ║");
    println!("  ║  • >95% SM utilization (vs ~75% traditional)                 ║");
    println!("  ║  • 1.2x speedup on irregular matrices                        ║");
    println!("  ║  • Eliminates wave quantization waste                        ║");
    println!("  ║                                                               ║");
    println!("  ╚═══════════════════════════════════════════════════════════════╝");
    println!();

    println!("  NEXT: PARITY-083 - Irregular matrix handling");

    assert!(true, "PARITY-082f: Summary complete");
}

// ==================== PARITY-083: Irregular Matrix Handling ====================
// Efficient handling of non-power-of-2 dimensions

/// PARITY-083a: LLM matrix shapes
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_083a_llm_matrix_shapes() {
    println!("PARITY-083a: LLM Matrix Shapes");
    println!("===============================");
    println!();

    // Real LLM dimensions from various models
    let models = [
        ("Phi-2 2.7B", 2560, 10240, 2560),
        ("Llama-2 7B", 4096, 11008, 4096),
        ("Mistral 7B", 4096, 14336, 4096),
        ("Llama-2 13B", 5120, 13824, 5120),
        ("Llama-2 70B", 8192, 28672, 8192),
    ];

    println!("  Common LLM Hidden/FFN Dimensions:");
    println!("  -----------------------------------");
    println!(
        "  {:15} {:>8} {:>8} {:>8}",
        "Model", "Hidden", "FFN", "Output"
    );
    println!(
        "  {:15} {:>8} {:>8} {:>8}",
        "─────", "──────", "───", "──────"
    );

    for (name, hidden, ffn, output) in models {
        println!("  {:15} {:>8} {:>8} {:>8}", name, hidden, ffn, output);
    }
    println!();

    // Check which are powers of 2
    println!("  Power-of-2 Analysis:");
    for (name, hidden, ffn, _) in models {
        let hidden_pow2 = hidden & (hidden - 1) == 0;
        let ffn_pow2 = ffn & (ffn - 1) == 0;
        println!(
            "    {}: hidden={} ({}), FFN={} ({})",
            name,
            hidden,
            if hidden_pow2 { "" } else { "" },
            ffn,
            if ffn_pow2 { "" } else { "" }
        );
    }
    println!();

    println!("  Key Insight:");
    println!("  • Most FFN dimensions are NOT powers of 2");
    println!("  • Traditional GEMM loses efficiency on these shapes");
    println!("  • Stream-K handles irregular shapes efficiently");

    assert!(true, "PARITY-083a: LLM shapes documented");
}

/// PARITY-083b: Padding overhead analysis
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_083b_padding_overhead() {
    println!("PARITY-083b: Padding Overhead Analysis");
    println!("=======================================");
    println!();

    // Traditional approach: pad to tile boundary
    let tile_size = 128u32;

    let dimensions = [
        (4096u32, 11008u32), // Llama-2 7B
        (4096, 14336),       // Mistral 7B
        (5120, 13824),       // Llama-2 13B
    ];

    println!("  Padding Overhead (tile_size={}):", tile_size);
    println!("  ─────────────────────────────────");
    println!(
        "  {:>8} {:>8} {:>8} {:>8} {:>8}",
        "M", "N", "Padded_M", "Padded_N", "Overhead"
    );

    for (m, n) in dimensions {
        let padded_m = m.div_ceil(tile_size) * tile_size;
        let padded_n = n.div_ceil(tile_size) * tile_size;
        let original = m as u64 * n as u64;
        let padded = padded_m as u64 * padded_n as u64;
        let overhead = (padded as f64 / original as f64 - 1.0) * 100.0;

        println!(
            "  {:>8} {:>8} {:>8} {:>8} {:>7.1}%",
            m, n, padded_m, padded_n, overhead
        );
    }
    println!();

    // Compute waste
    println!("  Wasted Computation:");
    for (m, n) in dimensions {
        let padded_m = m.div_ceil(tile_size) * tile_size;
        let padded_n = n.div_ceil(tile_size) * tile_size;
        let k = 4096u64; // Example
        let wasted_flops = 2 * (padded_m as u64 * padded_n as u64 - m as u64 * n as u64) * k;

        println!(
            "    {}×{}: {:.2} GFLOP wasted per forward",
            m,
            n,
            wasted_flops as f64 / 1e9
        );
    }

    assert!(true, "PARITY-083b: Padding overhead documented");
}

/// PARITY-083c: Predicated execution
#[test]
#[cfg(feature = "cuda")]
#[serial_test::serial]
fn test_parity_083c_predicated_execution() {
    println!("PARITY-083c: Predicated Execution");
    println!("===================================");
    println!();

    println!("  Predicated vs Padded Execution:");
    println!("  ─────────────────────────────────");
    println!();

    println!("  Padded Approach:");
    println!("    • Pad input matrices to tile boundary");
    println!("    • All threads compute (including padded region)");
    println!("    • Discard padded outputs");
    println!("    • Simple but wasteful");
    println!();

    println!("  Predicated Approach:");
    println!("    • No padding required");
    println!("    • Threads check bounds before load/store");
    println!("    • Out-of-bounds threads contribute zero");
    println!("    • More efficient for irregular shapes");
    println!();

    // PTX predicated load example
    println!("  PTX Predicated Load:");
    println!("  ─────────────────────");
    println!("    setp.lt.u32 %p1, %tid_m, %M;  // p1 = (row < M)");
    println!("    setp.lt.u32 %p2, %tid_n, %N;  // p2 = (col < N)");
    println!("    and.pred %p3, %p1, %p2;       // p3 = in_bounds");
    println!("    @%p3 ld.global.f32 %val, [%addr];");
    println!("    @!%p3 mov.f32 %val, 0.0;       // zero if OOB");
    println!();

    // Overhead analysis
    println!("  Predicate Overhead:");
    println!("    • 2-3 extra instructions per boundary check");
    println!("    • ~5% overhead for small tiles");
    println!("    • <1% overhead for large tiles (amortized)");
    println!();

    println!("  Stream-K + Predication:");
    println!("    • Combine work-stealing with bounds checking");
    println!("    • No wasted computation on irregular shapes");
    println!("    • Best of both worlds");

    assert!(true, "PARITY-083c: Predicated execution documented");
}