Skip to content

Consolidate and simplify the measured Vulkan inference improvements - #200

Open
kvark wants to merge 14 commits into
mainfrom
perf/llama-catchup-review
Open

kvark wants to merge 14 commits into
mainfrom
perf/llama-catchup-review

Conversation

@kvark

@kvark kvark commented Sep 20, 2026

Copy link
Copy Markdown
Owner

Consolidates the remaining inference work onto main at 4a66a27, replacing #193 and #196–199 (#194/#195 were included through #193). No reusable command recordings, new test executables, raw benchmark artifacts, or paper/protocol changes.

Simplification

The original seven commits cover attention sequence search, GEMV/layout improvements, pipeline-key resolution, private-state submission tuning, scalar staging search, queued readback and one benchmark guide.

Five follow-up commits remove 707 net lines from that version:

  • 3829907: remove the duplicate legacy epilogue representation and generator; use the existing pointwise DAG. Make attention eligibility exhaustive and deduplicate physical aliases.
  • 45638bf: replace nine mutually exclusive Dispatch fields with one Kernel enum. Serialization stays because the execution-plan cache uses it; invalidate older formats.
  • e374a18: remove unused scalar-fallback state and unnecessary pipeline compilation.
  • 008e655: prepare and resolve only the selected pipeline through one path.
  • bbe601d: use that same preparation path for kernel tuning and ordinary execution.

The attention installer is still substantial. The enum cleanup does not solve the architectural problem of choosing a graph before tuning its implementations.

Correctness follow-ups

2c46afe fixes cooperative forward/dQ sliding-window bounds: start from the first query row, so earlier rows retain their valid keys. The existing test now checks Q/K/V gradients with a window.

789bcb6 gives duplicated cooperative matrix, convolution and attention tiles a single subgroup writer. Naga emits subgroup-scoped matrices; equal-value non-atomic writes from different subgroups still overlap. Only stores are guarded. Staging, matrix arithmetic and workgroup barriers are unchanged; horizontal wrappers forward the subgroup ID. No fixed subgroup width, vendor gate or new tuning state is required.

Verification

CI for the current head passes all six jobs, including Metal, Windows, MSRV, quality and dependency checks. Rust host coverage is 84.07% lines, 84.29% functions, 83.45% regions. WGSL execution is not instrumented; existing CI backprop/bit-exact skips remain.

The cleanup passes 456 library tests and the focused state/staging, convolution, GEMV, integer-dot, pointwise, block-matmul and prologue suites on both RTX 5070 and Arc B570. Intel's pre-existing pairwise-distance gradient discrepancy remains (0.03788206 != 0.03788227); its tolerance is unchanged.

The store follow-up additionally passes all 31 code-generation checks, all-target/all-feature Clippy, and the 26 skinny-matrix, four convolution and Q/K/V-gradient checks on both GPUs. NVIDIA also passes the experimental full attention oracle and 11 convolution-derivative checks. Intel uses fallback: its two explicitly cooperative-only convolution checks stop at the existing capability assertion; the remaining nine pass. Known Naga Workgroup ArrayStride diagnostics remain.

Cleanup regression check

Exact 9d89ddc versus bbe601d, three fresh processes per arm/device, reversing the middle pair. Same official SmolLM2-135M F16 GGUF, 128-token prefill, 32 decode steps, F32 KV, full logits readback, 30-second kernel and separate two-second submission budgets. Reproduction guide.

Median of process medians, milliseconds:

GPU Before prefill After prefill Before decode After decode
RTX 5070 7.239 7.261 1.398 1.390
Arc B570 14.026 14.036 3.667 3.647

No material median regression is observed; this is not a claimed cleanup speedup or new llama.cpp comparison. All 12 output sets are finite and retain all 33 independent CPU-reference token choices. Maximum per-row relative L2 is 1.33e-5 / 1.77e-4. GPUs/builds run sequentially, clocks unlocked; B570 uses the secondary PCIe x1 link.

Egglog and kernel studies, separate from this PR

The source-only experiment reuses the existing egglog rules and extractor. It keeps equivalent graph families alive until after tuning their implementations. Alternatives are lowered before allocation; whole-program comparisons include split partials, reductions and surrounding work. Private immutable inference sessions share compatible weights and reuse only completed kernel searches, without reusing whole-program validation or timing.

Final SmolVLA results, with corrected stores in both control and candidates:

Workload / GPU Held-out latency reduction
SmolVLA action expert / RTX 5070 18.8–18.9%
SmolVLA action expert / Arc B570 37.6–37.7%
Whisper encoder / RTX 5070 18.5–19.4%

SmolVLA's complete 45-plan search takes 32–33 / 43–48 seconds with warm caches; both orders win 40/40 held-out pairs. Whisper's four attention layouts take 2.7–3.9 seconds and win 39/40 pairs in three repeats. Full independent CPU outputs are checked throughout. Short-budget order sensitivity and timing drift are documented, not hidden.

Nsight attribution and negative trials identify SmolVLA matrix/reduction work and Whisper attention as distinct bottlenecks. Host recording remains substantial. No wall-minus-kernel barrier estimate or PyTorch parity claim is made.

These prototype improvements are not enabled by this PR. They use Inferena's deterministic action-expert/encoder workloads, not full pretrained VLA or transcription pipelines. Production integration still needs the representative-data/qualification contract, training/stateful support, and replacement of the live attention installer. Experimental code and results stay on source branches; the submitted cohort is unchanged.

@kvark
kvark force-pushed the perf/llama-catchup-review branch from c208310 to 9d89ddc Compare September 20, 2026 03:59
@kvark
kvark marked this pull request as ready for review September 20, 2026 04:14
Comment thread src/runtime/tuning/attention.rs Outdated
}

fn plain(d: &Dispatch) -> bool {
!d.use_coop

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this approach is not scale-able. We could add another field to Dispatch and this condition wouldn't even no. At the very least, it should pattern match

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed eligibility to an exhaustive Dispatch pattern (3829907). The follow-up Kernel enum replaces nine independent implementation flags/options with one variant, and the pattern accepts only Kernel::Default. Adding a Dispatch field now requires revisiting the match; adding a kernel kind does not silently opt it into this probe.

Comment thread src/runtime/tuning/attention.rs Outdated
if physical
.iter()
.enumerate()
.any(|(i, p)| physical[..i].contains(p))

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is quadratic complexity? could we instead just sort and then dedup?

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changed to sort_unstable/dedup on a clone of the physical binding list (3829907), preserving the original binding order for placement bookkeeping.

Ok(())
}

fn qualify_reference(

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is all very complex. Can't there be a more general way to say "ok, you have a graph, now this specific parameter needs to have the following tuning range". I don't want to see that custom thousand-line boilerplate specific to each op.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant