[Performance]: Optimize native Rust/CUDA Laya inference #66

Proposal to improve performance

Track focused performance improvements to the native Rust/CUDA Laya implementation as a child of #14, starting from the eager encoder in #49. Follow the approach used in #55: profile a frozen baseline, change one bottleneck at a time, preserve precision/rounding, and validate each change with controlled comparisons.

Use #39 for the shared benchmark methodology and eventual comparison against the original Laya runtime. This issue owns implementation optimizations and their focused validation.

Starting evidence

A local experiment switches both kernel loading and dispatch in src/models/laya/src/encoder.rs from the existing rope_original entry point to rope, using the same compiled bundle and arithmetic.

Warm Encoder::run results on the same reserved GPU:

One feasibility sequence was excluded. Each configuration reused one worker, with two measured passes in AB/BA order; each block had ten excluded warmups and 50 measured requests. The two-pass mean ranges were disjoint for all five workloads. These are repeated frozen fixtures, with limited run-to-run coverage.

The timing includes host packing, input uploads, 28 encoder layers, two decision transformer layers, and completion synchronization. Tokenization, scoring/decoding, loading, compilation, warmup, hidden readback, and HTTP are outside this measurement.

Native hidden outputs were bitwise identical in 26/26 comparisons across 13 fixtures. Direct RoPE comparisons passed at 160 supported shapes with both global/local tables (320 comparisons), preserving V bytes.

Validation prerequisite: the official Laya 0.3.20 oracle fails in both unchanged #49 and the candidate at the same score/encoder27 comparison: nRMS 0.00725476 exceeds the existing 0.005 limit. Resolve the baseline/reference discrepancy before claiming full upstream parity or merge readiness; retain the existing tolerances.

Implementation priorities

• Reproduce and resolve the baseline/reference parity discrepancy with pinned dependencies and kernel provenance.
• Integrate the selected RoPE dispatch with supported-shape checks and published performance evidence.
• Tune small-shape GEMM geometry, starting with down/output projections. After the RoPE change, the four main encoder GEMMs account for approximately 72–74% of traced GPU kernel time on the short/long single-question cases. Test smaller tiles while preserving accumulation and BF16 rounding; Nsight Compute counters were unavailable, so occupancy/bandwidth causes remain unverified.
• Add shape-specific CUDA Graph capture/replay for the encoder and decision transformer, following [RFC]: Native Laya inference with Rust and CUDA #14's cache and lifecycle design. Measure capture/first-shape costs separately from replay latency.
• Reduce repeated upload synchronization and per-launch host work where profiling supports it: group uploads with safe host-buffer lifetimes and reuse resolved kernel/weight handles.

Keep changes in focused PRs consistent with #14's delivery plan. Place model dispatch/execution under src/models/laya/ and reusable CUDA operations under src/backends/cuda/. Further optimization priorities remain hypotheses until measured.

Validation and acceptance

• For each comparison, freeze source/binary/checkpoint/input hashes; declare the hypothesis, independent variable, success criterion, and run limit before execution. Keep exact GPU, affinity, precision, inputs, readiness, and cache conditions fixed.
• Run a CPU preflight, excluded feasibility run, then two measured runs per configuration by default. Publish commands, raw samples, traces, warmup/capture costs, and variability; report inconclusive results when appropriate.
• Preserve supported question types, output semantics, and established numerical tolerances. Cover changing inputs, padding, reused workspaces, and shape boundaries; validate eager/Graph agreement when Graphs are added.
• Report encoder, complete engine, and HTTP timings with explicit boundaries. Use [Performance]: Reproducible Laya benchmarks against the original runtime #39 for original-runtime quality/parity and end-to-end comparisons once the native serving path is runnable.

Recorded baseline and environment

• Encoder: Run Laya encoder and decision layers on CUDA #49 at ff4d8fa13b2c8d52027aa2c565a0b97940d1c0ca; main base 354704b6dafeb2dfc3cacd2fc2700da63a9b0b6f.
• Kernel builder: linear3735/system1-omni@5ff41a5; both configurations reused the same bundle.
• Checkpoint: convaiinnovations/laya@55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851, English 421M, maximum 512 tokens.
• GPU 2, UUID GPU-cbf66259-f4ab-0ede-1811-82037dde5924, NUMA 0 / CPUs 0–15. Reserved CUDA probing reports SM90, 132 SMs, and about 140 GB; NVML's reported L20X/8.9 label conflicts with that probe. Record both when reproducing.
• Torch 2.13.0+cu129, Transformers 5.10.2, TileLang 0.1.12. The cause of the oracle discrepancy is unresolved.
• Raw results and profiling artifacts are retained locally and should accompany the follow-up PR; the initial RoPE candidate is committed locally and has not been pushed.

Before submitting

• Searched existing/open and past issues and reviewed relevant implementation documentation; related work is linked above.