Skip to content

[AMDGPU] Skip redundant per-launch RuntimeContext HtoD - #876

Open
paveltc wants to merge 1 commit into
Genesis-Embodied-AI:mainfrom
AMD-Ecosystem:feat/amdgpu-skip-redundant-context-h2d
Open

[AMDGPU] Skip redundant per-launch RuntimeContext HtoD#876
paveltc wants to merge 1 commit into
Genesis-Embodied-AI:mainfrom
AMD-Ecosystem:feat/amdgpu-skip-redundant-context-h2d

Conversation

@paveltc

@paveltc paveltc commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

[amdgpu] Skip redundant per-launch RuntimeContext HtoD

Summary

On the AMDGPU launcher's default-stream path, launch_llvm_kernel re-uploads the
whole RuntimeContext struct to the per-handle persistent device buffer on
every launch:

memcpy_host_to_device_async(context_pointer, &ctx.get_context(),
                            sizeof(RuntimeContext), active_stream);

Across repeated launches of the same kernel handle that struct is almost always
byte-identical — the device arg-buffer / runtime / result-buffer pointers are
stable, and only the rare checkpoint kernel mutates the checkpoint_*_ptr
fields. Each of these copies costs ~4 µs of host-side API time (measured), which
is pure overhead on the launch-bound path (Genesis dispatches ~100 kernels per
step).

This PR caches the last-uploaded bytes (plus the device address they were
written to) in the per-handle Context and skips the HtoD when both still
match
. Correctness is preserved unconditionally:

  • The compare runs after prepare_streaming_checkpoint_state, so any
    checkpoint_*_ptr change forces a re-upload.
  • A moved arg-buffer (arg_buffer pointer changes) or a reallocated context
    buffer (context_pointer changes) also forces a re-upload.
  • The ephemeral (explicit-stream) path always uploads — it allocates a fresh
    buffer per launch and may run concurrently on pool streams.

Making the struct actually byte-stable

The skip-cache is useless unless RuntimeContext is stable. Profiling showed one
field flipping nearly every launch: result_buffer (offset 24). The launcher
only set it to the device buffer for result_buffer_size > 0 kernels; for
result-less kernels it left the per-launch-varying host pointer that
LaunchContextBuilder puts there. This PR pins result_buffer to the persistent
device buffer unconditionally:

  • Result-producing kernels are unchanged (they already got the device pointer).
  • Result-less kernels never touch the buffer, so pinning it to a stable, valid
    device address is safe — and is arguably more correct, since a host pointer
    sitting in a device-visible field is a latent hazard on AMDGPU (no UVA
    fallback).

Scope

  • quadrants/runtime/amdgpu/kernel_launcher.{cpp,h} only.
  • AMDGPU-local; no codegen change, no ABI change, no public API.
  • Net +40 / −5.

Benchmarks (CDNA3, gfx942)

12-kernel × 5000-launch loop (launch-bound):

per-launch result checksum
baseline 17.7 µs 10223616.000000
this PR 14.4 µs (~19% faster) 10223616.000000 (identical)

Correctness also verified on a workload including a reduction kernel
(result_buffer_size > 0): identical checksum and output sum vs baseline.

End-to-end (Genesis rigid-body scenes, 4096 envs)

The microbench figure is a launch-bound ceiling. On real scenes the async
context copy overlaps with solver compute, so the end-to-end gain scales
inversely with per-step compute.

The numbers below are not from the repository's tests/benchmarks/test_rigid.py
harness. They come from a small custom driver that reuses that file's scene shapes
(anymal / franka / go2 / box-pyramid, same assets and control patterns) but with
its own measurement: a fixed n_envs = 4096 for every scene and a per-step
timing loop (warm 80 steps, best-of-3 over 400 steps) rather than the harness's
warmup/record runtime_fps at the suite's official env counts. It is a paired A/B
(same binary, the optimization gated behind an env flag and toggled back-to-back),
with physics byte-identical between arms:

scene step time steps/s Δ
franka (collision-free) 1.8 ms +4.3%
franka 3.0 ms +3.1%
anymal (no control) 3.4 ms +2.8%
go2 4.2 ms ~+1%
anymal (per-env control) 11.9 ms +0.1%
box pyramid (stacking) 47 ms −0.4%

So: a consistent ~2–4% throughput gain on launch-bound articulated-body
scenes
, tapering to ~0% (never a regression) as scenes become compute/collision
bound. Official test_rigid.py numbers at the suite's env counts are still worth
collecting as a follow-up.

Test plan

  • Byte-identical results vs baseline on result-less and reduction kernels.
  • Builds and runs on main with the AMDGPU backend.
  • Custom end-to-end rigid-body A/B (above, test_rigid-derived scenes at
    4096 envs): positive on launch-bound scenes, neutral on compute-bound, no
    regressions.
  • Official tests/benchmarks/test_rigid.py A/B at the suite's env counts.
  • Confirm parity on a checkpoint-bearing (qd.checkpoint) kernel — the
    compare should force re-upload when checkpoint_*_ptr changes.
  • Confirm parity under an explicit-stream / graph-capture path (must always
    upload).

Note on scope

This captures the host-side per-launch win without the higher-risk
kernarg-by-value codegen change (passing the struct by value would also remove
the per-instruction context pointer-loads in the kernel body, but touches shared
codegen and arch guards). It is deliberately launcher-local so it can land
independently and be reverted trivially.

The AMDGPU launcher re-uploads the whole RuntimeContext struct to the
per-handle persistent device buffer on every launch. Across repeated launches
of the same kernel handle the struct is almost always identical, so most of
these ~4us async copies are pure overhead on the launch-bound path.

Cache the last-uploaded bytes (plus the device address they were written to)
in the per-handle Context and skip the HtoD when both still match. The compare
runs after prepare_streaming_checkpoint_state, so any checkpoint_*_ptr mutation
forces a re-upload; a moved arg-buffer or reallocated context buffer also forces
one. The ephemeral (explicit-stream) path always uploads.

To make the struct actually byte-stable, also pin RuntimeContext.result_buffer
to the persistent device buffer unconditionally. Previously it was only set for
result_buffer_size > 0 kernels, leaving a per-launch-varying host pointer in the
field for result-less kernels - which defeated the cache and was a latent hazard
on AMDGPU (no UVA fallback for host pointers).

Launcher-local, AMDGPU-only; no codegen or ABI change. Measured on CDNA3
(gfx942), 12-kernel x 5000-launch loop: per-launch 17.7us -> 14.4us (~19%),
with byte-identical results on both result-less and reduction kernels.

Co-authored-by: Cursor <cursoragent@cursor.com>
@paveltc paveltc changed the title [amdgpu] Skip redundant per-launch RuntimeContext HtoD [AMDGPU] Skip redundant per-launch RuntimeContext HtoD Aug 20, 2026
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.

2 participants