From 59d1bccb4979efac73ed25a083af7345661d6224 Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Sat, 5 Sep 2026 07:17:57 +0000 Subject: [PATCH 1/2] perf(runtime): keep the Linear Attention state in L2 across the decode graph Every decode round reads the whole Linear Attention continuation state and writes it back at a fixed address, and recurrent_fold_kernel streams all of it in a single launch. That is the one decode consumer with both a fixed footprint and enough reuse to be worth an L2 reservation, so the captured decode graph now carries an access policy window over the state pool with hitProp = Persisting. The window is installed before cudaStreamBeginCapture, because capture bakes the stream window into every kernel node and does not re-read it at replay, and it is taken off the stream right after, so eager phases keep the default access policy window. The device-wide reserve is a separate setting and is not restored: it stays in force outside the graph too, which is where the prefill cost below comes from. The reserve is clamped to persistingL2CacheMaxSize (60.00 MiB on this card), and when the pool is larger than the reserve the hit ratio is lowered to the fraction that physically fits, which is what CUDA requires: a larger window at hitRatio 1 makes the lines evict each other. Nothing at all is installed - no window, and no reserve either - when the state pool is wider than accessPolicyMaxWindowSize (128.00 MiB here, read off the device). Past that limit a window is not a shorter window over the pool, it is a window over the pool's prefix, while the device-wide reserve is still taken in full and paid for in prefill. The threshold is the window limit and not the reserve limit: on Qwen3.6-35B-A3B the window is the whole 122.81 MiB pool against a 60.00 MiB reserve, and that asymmetry is the configuration every number below was taken at. Qwen3.6-35B-A3B on RTX 5090, pp2048+tg384, prefill chunk 8192, MTP draft tokens 3, eight alternating passes with rotated arm order and the first dropped as warm-up, delta taken as pairs inside a pass, median over the remaining seven: decode throughput +0.789% (sem 0.043%), and +0.539% (sem 0.072%) on a second card prefill throughput -0.709% (sem 0.054%) total wall time +0.581% (sem 0.035%) Run-to-run spread is wider than the within-run sem, so read the decode figure as a +0.5..0.8% range rather than as a point. A second ladder on the second card reads +0.660% (sem 0.145%), 0.121pp from the first, so the absolute level is not claimed to replicate across cards. Two later ladders on that same card, taken to check the guard, put the same arm at +0.752% and +0.674% and the guarded arm at +0.778% and +0.777%, which stays inside that range and widens the same-card spread to 0.213pp. One contrast does replicate, to 0.009 and 0.060pp: the same window with the prefill cache reset put back. That is a fact about those two arms and not a property of pairing, since every figure here is a pair taken inside a pass including the headline, and it is not offered as the headline, because both of its arms carry the change while master carries neither. Prefill pays because the reservation is also in force while a chunk streams through L2; at this prompt and generation length the decode gain more than covers it, and a prefill-dominated mix would not be a good trade. Qwen3.8-27B NVFP4 is the target the guard exists for, and it is a second target of this same runtime path rather than a hypothetical one. Its state pool is 307 888 128 B, 2.29x the window limit, so the window used to cover 43.6% of the pool while the reserve was paid for all of it. Same protocol, same card, two clean ladders per column: without the guard with the guard decode -0.039% / -0.041% -0.013% / -0.008% prefill -0.550% / -0.521% -0.001% / +0.005% total wall time -0.089% / -0.108% -0.013% / -0.018% Half a percent of prefill for nothing in decode, and now nothing either way: every figure in the right column is inside the base spread of its own ladder, 0.027% and 0.062%. Walking the captured graph shows the mechanism off rather than weakened: 0 of 811 kernel nodes carry a window, against 811 of 811 without the guard, and the persisting limit stays at the 18 874 368 B driver default instead of being raised. On Qwen3.6-35B-A3B the guard does not fire and changes nothing: the audit prints the same geometry either way - 520 of 520 kernel nodes windowed, window 128 778 240 B, reserve 62 914 560 B, hitRatio 0.4885 - and two clean ladders put the guarded arm against the unguarded one at +0.030pp and -0.072pp of decode, 0.080pp and 0.052pp of prefill and 0.002pp and 0.016pp of total wall time, on base spreads of 0.053% and 0.066%. Nothing reports the truncation on its own - cudaDeviceSetLimit and cudaStreamSetAttribute both return cudaSuccess whether the window covers the pool or a fraction of it - so the guard says so once per process on stderr rather than disabling itself silently, which is the failure it exists to prevent. recurrent_fold_kernel goes 54944 -> 41216 ns median (-25.0%), recurrent_record_kernel 5344 -> 4737 ns (-11.4%) and recurrent_bf16_direct_kernel 2656 -> 2096 ns (-21.1%). The one kernel that gets materially slower is the prefill reduce, 38193 -> 48288.5 ns (+26.4%); across all 84 kernels, which are the same set in both runs, no other median moves by more than 8.3% in either direction, and total kernel time over the run drops 1434.0 -> 1430.3 ms. At 41216 ns recurrent_fold_kernel moves its 61.41 MiB of state in and out at 3124.5 GB/s, which is 2.08x the measured DRAM copy ceiling of 1503.4 GB/s on this card, so a large share of that traffic cannot be coming from DRAM. Output is bit-identical: 21 of 21 comparisons over three prompts under greedy decoding on the earlier revision, and re-run with the guard in place, 12 of 12 on Qwen3.6-35B-A3B and 12 of 12 on Qwen3.8-27B - the target where the guard fires - with the reflexive base-against-base comparison taken first every time and exactly three distinct digests per model. ctest is 104/104 on both arms from one build directory on ad0f3d38, the base this commit sits on, whose tests/CMakeLists.txt registers exactly those 104 targets; 487f8977 registers 114 and neither arm was re-run there. The run reads 464.9 s on the base and 460.5 s on the branch, and the three opt-in real-artifact tests reachable here were run explicitly and pass, including the one on a 307 888 128 B state pool, which is the guarded path. Co-Authored-By: Claude Opus 5 Base and provenance. This change now targets upstream 487f8977; it was written and measured on ad0f3d38, and nothing in this message has been re-measured on the new base. Two instrument changes came with the rebase. The bench flag --mtp-draft-tokens no longer exists: the spelling is --spec mtp --draft-tokens N, and --draft-tokens 0 is rejected for MTP (src/product/speculative_options.h:41), so the zero-draft arm is spelled by passing neither flag. And ninfer_bench now disables the engine context cache itself (bench/targets/qwen3_6_27b/ninfer_bench.cpp:157), where on ad0f3d38 that line did not exist and the cache stayed at its enabled default. Per-request prefix reuse was already off on both bases (same file, line 65), so no request here could take a prefix hit; how far a constructed-but-unused cache reaches into prefill and TTFT has not been measured, and these figures describe a configuration the current bench no longer builds. tests/CMakeLists.txt registers 114 targets on 487f8977 rather than 104, with 20 named targets explicitly artifact-gated rather than 19; no ctest run quoted here was repeated on the new base, and no round "114 of 114" is claimed for it: ninfer_attn_input_proj_test is red on the bare base (upstream defect, issue #196), and ctest and a direct run of the same binary have been seen to disagree on this host. --- src/CMakeLists.txt | 1 + src/core/l2_persist.cpp | 70 +++++++++++++++++++ src/core/l2_persist.h | 29 ++++++++ src/targets/qwen3_6/impl/runtime/graph_impl.h | 22 ++++++ 4 files changed, 122 insertions(+) create mode 100644 src/core/l2_persist.cpp create mode 100644 src/core/l2_persist.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 3e99d8c7e4..0902b4860f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -27,6 +27,7 @@ add_library(ninfer_core STATIC core/layout.cpp core/tensor.cpp core/decode_graph.cpp + core/l2_persist.cpp runtime/engine/admission_policy.cpp runtime/engine/context_cost.cpp runtime/engine/context_cost_defaults.cpp diff --git a/src/core/l2_persist.cpp b/src/core/l2_persist.cpp new file mode 100644 index 0000000000..f7b9aa35ca --- /dev/null +++ b/src/core/l2_persist.cpp @@ -0,0 +1,70 @@ +#include "core/l2_persist.h" + +#include +#include + +namespace ninfer::l2p { + +bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes) { + if (base == nullptr || bytes == 0) { return false; } + + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return false; } + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, device) != cudaSuccess) { return false; } + if (prop.persistingL2CacheMaxSize <= 0 || prop.accessPolicyMaxWindowSize <= 0) { return false; } + + // An access policy window can never be longer than accessPolicyMaxWindowSize. A pool wider + // than that does not get a shorter window over the whole pool, it gets a window over the + // pool's prefix, while the device-wide reserve is taken in full and paid for everywhere, + // prefill included. On a target where that happens the reservation is a pure loss, so do + // nothing at all: no window, and no reserve either. + if (bytes > static_cast(prop.accessPolicyMaxWindowSize)) { + // Nothing else reports this: the driver returns cudaSuccess whether the window covers + // the pool or a fraction of it, so a silently truncated window is indistinguishable + // from a working one. Say it once per process instead. + static std::atomic reported{false}; + if (!reported.exchange(true)) { + std::fprintf(stderr, + "ninfer: L2 persistence disabled: the %zu byte Linear Attention state " + "pool is wider than the %zu byte access policy window of this device\n", + bytes, static_cast(prop.accessPolicyMaxWindowSize)); + } + return false; + } + + const std::size_t window = bytes; + std::size_t reserve = window; + if (reserve > static_cast(prop.persistingL2CacheMaxSize)) { + reserve = static_cast(prop.persistingL2CacheMaxSize); + } + if (cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, reserve) != cudaSuccess) { + return false; + } + + float hit_ratio = 1.0F; + if (window > reserve) { + hit_ratio = static_cast(static_cast(reserve) / static_cast(window)); + } + + cudaStreamAttrValue value{}; + value.accessPolicyWindow.base_ptr = const_cast(base); + value.accessPolicyWindow.num_bytes = window; + value.accessPolicyWindow.hitRatio = hit_ratio; + value.accessPolicyWindow.hitProp = cudaAccessPropertyPersisting; + value.accessPolicyWindow.missProp = cudaAccessPropertyNormal; + (void)cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &value); + return true; +} + +void unpin(cudaStream_t stream) { + cudaStreamAttrValue value{}; + value.accessPolicyWindow.base_ptr = nullptr; + value.accessPolicyWindow.num_bytes = 0; + value.accessPolicyWindow.hitRatio = 0.0F; + value.accessPolicyWindow.hitProp = cudaAccessPropertyNormal; + value.accessPolicyWindow.missProp = cudaAccessPropertyNormal; + (void)cudaStreamSetAttribute(stream, cudaStreamAttributeAccessPolicyWindow, &value); +} + +} // namespace ninfer::l2p diff --git a/src/core/l2_persist.h b/src/core/l2_persist.h new file mode 100644 index 0000000000..14911f3b81 --- /dev/null +++ b/src/core/l2_persist.h @@ -0,0 +1,29 @@ +#pragma once + +#include + +#include + +namespace ninfer::l2p { + +/** + * Pins a device range in L2 for the graph captured next on this stream. + * + * Must be called BEFORE cudaStreamBeginCapture: capture bakes the stream access policy window + * into every kernel node of the graph and does not re-read the stream attribute at replay. + * The reserve is clamped to cudaDeviceProp::persistingL2CacheMaxSize, and when the range is + * larger than the reserve the hit ratio is lowered to the fraction that physically fits, which + * is what CUDA requires: a window larger than the reserve at hitRatio 1 makes the lines evict + * each other. + * + * Does nothing - no window and no reserve - when the range is wider than + * cudaDeviceProp::accessPolicyMaxWindowSize, because past that limit the window covers only a + * prefix of the range while the reserve is still taken in full. Returns whether a window was + * installed, so the caller knows whether it has anything to take off the stream. + */ +bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes); + +/** Removes the window from the stream, so eager phases run under the default policy. */ +void unpin(cudaStream_t stream); + +} // namespace ninfer::l2p diff --git a/src/targets/qwen3_6/impl/runtime/graph_impl.h b/src/targets/qwen3_6/impl/runtime/graph_impl.h index e0687bdd1e..2d866a5f23 100644 --- a/src/targets/qwen3_6/impl/runtime/graph_impl.h +++ b/src/targets/qwen3_6/impl/runtime/graph_impl.h @@ -1,8 +1,11 @@ #include "targets/qwen3_6/impl/runtime/instance.h" #include "targets/qwen3_6/impl/runtime/schedule.h" +#include "core/l2_persist.h" #include "core/nvtx.h" +#include +#include #include namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule { @@ -23,7 +26,26 @@ void run_prepared(Context& state, DecodeGraphExecutable* executable, Body&& body template void capture_graph(Context& state, DecodeGraphDefinition& definition, Body&& body) { state.execution.work.reset(); + // Keep the Linear Attention continuation state resident in L2 for the captured decode + // graph. Every decode round reads the whole pool and writes it back at a fixed address, + // and recurrent_fold_kernel streams all of it in a single launch, so it is the one decode + // consumer with both a fixed footprint and enough reuse to be worth an L2 reservation. + // The window must be installed before capture: the graph bakes it into its kernel nodes + // and does not re-read the stream attribute at replay. + const LinearAttentionStateAllLayersView linear = + state.execution.linear_attention.all_layers_view(); + const auto low = reinterpret_cast(linear.conv_layer0.data); + const auto span = + static_cast(static_cast(linear.recurrent_layer_stride_bytes) * + (linear.spec.layers - 1U) + + reinterpret_cast(linear.recurrent_layer0.data) - + low + linear.recurrent_layer0.bytes()); + // A target whose state pool does not fit an access policy window gets nothing installed, + // and then there is nothing to take off the stream either. + const bool pinned = + l2p::pin_range(state.execution.device.stream, reinterpret_cast(low), span); definition.capture(state.execution.device.stream, body); + if (pinned) { l2p::unpin(state.execution.device.stream); } } } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule From 0deee4d8d155a6b7c46a768580c5bd7718e56b40 Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Mon, 7 Sep 2026 15:58:12 +0200 Subject: [PATCH 2/2] perf(runtime): give the persisting-L2 set-aside an owner that gives it back Codex flagged that `cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, ...)` in `pin_range` is context state that nothing restores: after an Engine with graphs is destroyed, an Engine that installs no access policy window - graphs disabled, or a state pool wider than `accessPolicyMaxWindowSize` - inherits the enlarged set-aside and pays for it while using none of it. The finding is correct; this commit adds the owner. `l2p::Reservation` is a counted handle living in `DeviceContext`. The limit is per device, not per process, so requests are reconciled per device ordinal: the driver limit is held at the largest live request for that device and returns to the limit that was in force before the first request once the last one is released. `pin_range` takes the set-aside through the handle instead of writing the limit itself. `DeviceContext` outlives every graph captured under the window - Program and EngineCore hold it by reference and it is a member of `Engine::Impl` declared before `core` - so the owner is wider than the graphs it serves rather than narrower. Behaviour on the measured configuration is unchanged while the reserve is live: the 35B state pool is 128 778 240 B, the reserve clamps to 62 914 560 B, and that is above this driver's 18 874 368 B default in both the old code and the new. What differs is teardown. Three properties the owner has to have, each of which cost a defect before it was written down: * Its fields are under the same lock as the registry. With them outside, two threads on one handle publish two requests and withdraw one, and the device is stranded at a raised limit with no window - the exact state this commit exists to prevent. Reproduced, then fixed. * A failed device bind is reported, not swallowed. `cudaDeviceSetLimit` writes the current device, so a release whose `cudaSetDevice` failed would otherwise write one device's bookkeeping to another. * The bookkeeping is surrendered only once the restoring write lands, and a request that records no baseline may not take one back on failure. Otherwise a refused request - one above `persistingL2CacheMaxSize` returns `cudaErrorInvalidValue` on healthy hardware - discards the baseline a failed restore was still holding, and the process forgets what it owed. `tests/test_l2_reservation.cpp` covers nesting, equal-sized twins, 32 cycles without baseline drift, `request(0)` as a non-release, moves, replacement on a live handle, a refused request, and concurrent use of one handle. Its discrimination is measured rather than assumed: six mutations of the source were built and run against it, and the cases above were written until the survivors died. Two mutants survive by construction - both need a limit write to fail, which one healthy card cannot stage - and are recorded as uncovered rather than passed. `ctest` 114 of 115. The one failure is `ninfer_attn_input_proj_test`, which fails the same way on untouched master and is the subject of issue #196; it reaches neither `pin_range` nor `capture_graph`. Co-Authored-By: Claude Opus 5 --- src/core/device.cu | 3 +- src/core/device.h | 6 + src/core/l2_persist.cpp | 180 ++++++++++++- src/core/l2_persist.h | 56 +++- src/targets/qwen3_6/impl/runtime/graph_impl.h | 3 +- tests/CMakeLists.txt | 2 + tests/test_l2_reservation.cpp | 242 ++++++++++++++++++ 7 files changed, 483 insertions(+), 9 deletions(-) create mode 100644 tests/test_l2_reservation.cpp diff --git a/src/core/device.cu b/src/core/device.cu index 54db2d31bd..08636a0b6a 100644 --- a/src/core/device.cu +++ b/src/core/device.cu @@ -86,7 +86,7 @@ DeviceContext::~DeviceContext() { DeviceContext::DeviceContext(DeviceContext&& other) noexcept : device(other.device), stream(other.stream), transfer_stream(other.transfer_stream), - props(other.props) { + props(other.props), persisting_l2(std::move(other.persisting_l2)) { other.stream = nullptr; other.transfer_stream = nullptr; } @@ -102,6 +102,7 @@ DeviceContext& DeviceContext::operator=(DeviceContext&& other) noexcept { props = other.props; stream = other.stream; transfer_stream = other.transfer_stream; + persisting_l2 = std::move(other.persisting_l2); other.stream = nullptr; other.transfer_stream = nullptr; diff --git a/src/core/device.h b/src/core/device.h index 4bc06d96b3..b121d14883 100644 --- a/src/core/device.h +++ b/src/core/device.h @@ -1,5 +1,7 @@ #pragma once +#include "core/l2_persist.h" + #include #include @@ -23,6 +25,10 @@ struct DeviceContext { cudaStream_t stream = nullptr; cudaStream_t transfer_stream = nullptr; cudaDeviceProp props{}; + // Owns the device-wide persisting-L2 set-aside for as long as this context exists. The limit + // is context state and outlives both the stream attribute and the graph the window is baked + // into, so it is given back here rather than left standing for whatever runs next. + l2p::Reservation persisting_l2; explicit DeviceContext(int device_id = 0); ~DeviceContext(); diff --git a/src/core/l2_persist.cpp b/src/core/l2_persist.cpp index f7b9aa35ca..efb13f79a7 100644 --- a/src/core/l2_persist.cpp +++ b/src/core/l2_persist.cpp @@ -2,10 +2,184 @@ #include #include +#include +#include +#include namespace ninfer::l2p { +namespace { -bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes) { +// The limit is one number per device, so its owners have to be reconciled in one place. Requests +// are counted rather than replaced: withdrawing one owner must not shrink the set-aside another +// owner is still replaying a graph under. +struct Registry { + std::mutex mutex; + std::map> live; + // The limit in force before the first request on that device, restored when the last one goes + // away. Recorded under the same lock that publishes the first request, so a second owner + // cannot record an already raised limit as the one to return to, and kept until a restore + // actually lands: a failed write must not lose the value it was going to write. + std::map baseline; +}; + +// Deliberately never destroyed. A handle owned by an object with static storage duration would +// otherwise reach a destroyed registry during exit, and the cost of the leak is one map. +Registry& registry() { + static Registry* const instance = new Registry(); + return *instance; +} + +// cudaDeviceSetLimit writes the calling thread's current device, so a release has to name the +// device its request was made against rather than whichever one happens to be bound. A failed +// bind is reported rather than swallowed: writing the limit anyway would apply one device's +// bookkeeping to another. +class BoundDevice { +public: + explicit BoundDevice(int device) { + int current = 0; + if (cudaGetDevice(¤t) != cudaSuccess) { return; } + if (current == device) { + bound_ = true; + return; + } + if (cudaSetDevice(device) != cudaSuccess) { return; } + previous_ = current; + bound_ = true; + } + + ~BoundDevice() { + if (previous_ >= 0) { (void)cudaSetDevice(previous_); } + } + + [[nodiscard]] bool bound() const noexcept { return bound_; } + + BoundDevice(const BoundDevice&) = delete; + BoundDevice& operator=(const BoundDevice&) = delete; + BoundDevice(BoundDevice&&) = delete; + BoundDevice& operator=(BoundDevice&&) = delete; + +private: + int previous_ = -1; + bool bound_ = false; +}; + +bool apply_limit(int device, std::size_t bytes) { + const BoundDevice bound(device); + if (!bound.bound()) { return false; } + return cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, bytes) == cudaSuccess; +} + +// Drops one request and writes back what that device should hold now. Called with the lock held. +// When that was the last request, the baseline is given up only once the restoring write lands, so +// a device whose limit could not be written keeps the value it still owes a restore to. A device +// that still has owners needs no such care: the next withdrawal recomputes the target from what +// remains, so a write that did not land is corrected rather than remembered. +void withdraw(Registry& reg, int device, std::size_t bytes) { + const auto slot = reg.live.find(device); + if (slot == reg.live.end()) { return; } + auto& live = slot->second; + const auto it = live.find(bytes); + if (it == live.end()) { return; } + live.erase(it); + + const auto base = reg.baseline.find(device); + if (!live.empty()) { + (void)apply_limit(device, *live.rbegin()); + return; + } + if (base == reg.baseline.end()) { + // No recorded baseline means nothing is known to restore to; leaving the limit alone is + // the only honest option. Unreachable while a request was live, kept as a guard. + reg.live.erase(slot); + return; + } + if (apply_limit(device, base->second)) { + reg.live.erase(slot); + reg.baseline.erase(base); + } +} + +} // namespace + +Reservation::Reservation(Reservation&& other) noexcept { + const std::lock_guard guard(registry().mutex); + bytes_ = other.bytes_; + device_ = other.device_; + held_ = other.held_; + other.bytes_ = 0; + other.held_ = false; +} + +Reservation& Reservation::operator=(Reservation&& other) noexcept { + if (this == &other) { return *this; } + Registry& reg = registry(); + const std::lock_guard guard(reg.mutex); + if (held_) { withdraw(reg, device_, bytes_); } + bytes_ = other.bytes_; + device_ = other.device_; + held_ = other.held_; + other.bytes_ = 0; + other.held_ = false; + return *this; +} + +bool Reservation::request(std::size_t bytes) { + // Zero is not a release: a handle that asks for nothing still owes its graph the set-aside it + // already holds. + if (bytes == 0) { return false; } + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return false; } + + Registry& reg = registry(); + const std::lock_guard guard(reg.mutex); + if (held_ && device_ == device && bytes_ == bytes) { return true; } + + auto& live = reg.live[device]; + // A withdrawal whose write did not land leaves an empty request set behind but keeps the + // baseline it still owes a restore to, so an empty set is not proof that this call is the + // first owner. Only a call that records the baseline may take it back on failure. + const bool recorded_baseline = reg.baseline.find(device) == reg.baseline.end(); + if (recorded_baseline) { + std::size_t current = 0; + const BoundDevice bound(device); + if (!bound.bound() || + cudaDeviceGetLimit(¤t, cudaLimitPersistingL2CacheSize) != cudaSuccess) { + if (live.empty()) { reg.live.erase(device); } + return false; + } + reg.baseline[device] = current; + } + + // Put the new request in force before withdrawing the old one. A refusal must not leave a + // graph that is already captured running under a set-aside that has been handed back. + live.insert(bytes); + if (!apply_limit(device, *live.rbegin())) { + const auto it = live.find(bytes); + if (it != live.end()) { live.erase(it); } + if (live.empty()) { + reg.live.erase(device); + if (recorded_baseline) { reg.baseline.erase(device); } + } + return false; + } + + if (held_) { withdraw(reg, device_, bytes_); } + bytes_ = bytes; + device_ = device; + held_ = true; + return true; +} + +void Reservation::release() noexcept { + Registry& reg = registry(); + const std::lock_guard guard(reg.mutex); + if (!held_) { return; } + held_ = false; + withdraw(reg, device_, bytes_); + bytes_ = 0; +} + +bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes, Reservation& reservation) { if (base == nullptr || bytes == 0) { return false; } int device = 0; @@ -38,9 +212,7 @@ bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes) { if (reserve > static_cast(prop.persistingL2CacheMaxSize)) { reserve = static_cast(prop.persistingL2CacheMaxSize); } - if (cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, reserve) != cudaSuccess) { - return false; - } + if (!reservation.request(reserve)) { return false; } float hit_ratio = 1.0F; if (window > reserve) { diff --git a/src/core/l2_persist.h b/src/core/l2_persist.h index 14911f3b81..53eb8927e9 100644 --- a/src/core/l2_persist.h +++ b/src/core/l2_persist.h @@ -6,6 +6,52 @@ namespace ninfer::l2p { +/** + * Ownership handle for a device's persisting-L2 set-aside. + * + * cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize) is context state, not stream or graph state: + * it outlives both the access policy window that motivated it and the graph the window was baked + * into, so it needs an owner that gives it back. A handle registers one request against one + * device; the limit is held at the largest live request for that device, and returns to the limit + * that was in force before the first request once the last one is released. Without that, work + * started after the requesting owner is gone carries a set-aside while installing no window, which + * costs what an unused reserve costs and buys nothing. + * + * The limit is per device, not per process, so requests are reconciled per device ordinal and a + * release binds its own device before writing the limit back; a bind that fails is reported rather + * than written to whichever device happens to be current. Handles are safe to use from several + * threads. + */ +class Reservation { +public: + Reservation() noexcept = default; + + ~Reservation() { release(); } + + Reservation(const Reservation&) = delete; + Reservation& operator=(const Reservation&) = delete; + Reservation(Reservation&& other) noexcept; + Reservation& operator=(Reservation&& other) noexcept; + + /** + * Requests `bytes` of set-aside on the calling thread's current device, replacing whatever + * this handle asked for before. The limit becomes the largest request live on that device. + * + * Returns whether the request is in force. A refused request changes nothing: the previous + * request of this handle stays live, because a graph may already be captured under it. A + * request of zero bytes is refused for the same reason, rather than treated as a release. + */ + bool request(std::size_t bytes); + + /** Withdraws this handle's request, lowering or restoring its device's limit accordingly. */ + void release() noexcept; + +private: + std::size_t bytes_ = 0; + int device_ = 0; + bool held_ = false; +}; + /** * Pins a device range in L2 for the graph captured next on this stream. * @@ -16,12 +62,16 @@ namespace ninfer::l2p { * is what CUDA requires: a window larger than the reserve at hitRatio 1 makes the lines evict * each other. * + * The set-aside is taken through `reservation`, which must outlive every graph captured under this + * window: the window is replayed from the graph long after this call returns, and the set-aside + * has to still be in force for it to mean anything. + * * Does nothing - no window and no reserve - when the range is wider than * cudaDeviceProp::accessPolicyMaxWindowSize, because past that limit the window covers only a - * prefix of the range while the reserve is still taken in full. Returns whether a window was - * installed, so the caller knows whether it has anything to take off the stream. + * prefix of the range while the device-wide reserve is still taken in full. Returns whether a + * window was installed, so the caller knows whether it has anything to take off the stream. */ -bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes); +bool pin_range(cudaStream_t stream, const void* base, std::size_t bytes, Reservation& reservation); /** Removes the window from the stream, so eager phases run under the default policy. */ void unpin(cudaStream_t stream); diff --git a/src/targets/qwen3_6/impl/runtime/graph_impl.h b/src/targets/qwen3_6/impl/runtime/graph_impl.h index 2d866a5f23..d820b3d0b9 100644 --- a/src/targets/qwen3_6/impl/runtime/graph_impl.h +++ b/src/targets/qwen3_6/impl/runtime/graph_impl.h @@ -43,7 +43,8 @@ void capture_graph(Context& state, DecodeGraphDefinition& definition, Body&& bod // A target whose state pool does not fit an access policy window gets nothing installed, // and then there is nothing to take off the stream either. const bool pinned = - l2p::pin_range(state.execution.device.stream, reinterpret_cast(low), span); + l2p::pin_range(state.execution.device.stream, reinterpret_cast(low), span, + state.execution.device.persisting_l2); definition.capture(state.execution.device.stream, body); if (pinned) { l2p::unpin(state.execution.device.stream); } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 34760fdf64..a7076a047c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -69,6 +69,8 @@ ninfer_add_test(ninfer_pretty_logging_test LIBRARIES ninfer_product_logging) ninfer_add_test(ninfer_device_test SOURCES test_device.cpp) ninfer_add_test(ninfer_decode_graph_test SOURCES test_decode_graph.cpp) +ninfer_add_test(ninfer_l2_reservation_test SOURCES test_l2_reservation.cpp) +set_tests_properties(ninfer_l2_reservation_test PROPERTIES SKIP_RETURN_CODE 77) ninfer_add_test(ninfer_tensor_test SOURCES test_tensor.cpp) ninfer_add_test(ninfer_arena_test SOURCES test_arena.cpp) ninfer_add_test(ninfer_admission_policy_test SOURCES test_admission_policy.cpp) diff --git a/tests/test_l2_reservation.cpp b/tests/test_l2_reservation.cpp new file mode 100644 index 0000000000..e3e9270b90 --- /dev/null +++ b/tests/test_l2_reservation.cpp @@ -0,0 +1,242 @@ +// The persisting-L2 set-aside is device-context state that outlives the access policy window and +// the graph the window is baked into. These cases pin the property that makes it safe to touch at +// all: whatever the handles do, the device is left with the limit it had before the first request. +#include "core/l2_persist.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +int failures = 0; + +std::size_t limit() { + std::size_t value = 0; + if (cudaDeviceGetLimit(&value, cudaLimitPersistingL2CacheSize) != cudaSuccess) { + return static_cast(-1); + } + return value; +} + +void expect_at_least(const char* what, std::size_t got, std::size_t want) { + if (got < want) { + std::cout << what << ": limit " << got << " does not cover " << want << '\n'; + ++failures; + } +} + +// The set-aside has to come down as well as go up, and on a card whose driver default already +// exceeds a small request an "at least" check cannot see that. This is the direction that fails +// when a withdrawal stops lowering the limit to the request that survived it. +void expect_below(const char* what, std::size_t got, std::size_t want) { + if (got >= want) { + std::cout << what << ": limit " << got << " still carries " << want << '\n'; + ++failures; + } +} + +void expect_equal(const char* what, std::size_t got, std::size_t want) { + if (got != want) { + std::cout << what << ": limit " << got << ", expected " << want << '\n'; + ++failures; + } +} + +bool cuda_unavailable() { + int count = 0; + return cudaGetDeviceCount(&count) != cudaSuccess || count <= 0; +} + +} // namespace + +int main() { + if (cuda_unavailable()) { + std::cout << "SKIP: no usable CUDA device\n"; + return 77; + } + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, 0) != cudaSuccess || prop.persistingL2CacheMaxSize <= 0) { + std::cout << "SKIP: device reserves no L2 for persisting accesses\n"; + return 77; + } + + // The driver rounds a set-aside up, so coverage is checked with >= and only the restored value + // is checked for equality - it is the one the driver itself produced. + const std::size_t ceiling = static_cast(prop.persistingL2CacheMaxSize); + const std::size_t small = ceiling / 8; + const std::size_t large = ceiling / 2; + if (small == 0) { + std::cout << "SKIP: persisting L2 reserve too small to subdivide\n"; + return 77; + } + const std::size_t baseline = limit(); + + { + ninfer::l2p::Reservation outer; + if (!outer.request(small)) { + std::cout << "a request the device advertises room for was refused\n"; + return 1; + } + expect_at_least("one request covers what it asked for", limit(), small); + { + ninfer::l2p::Reservation inner; + if (!inner.request(large)) { + std::cout << "a nested request was refused\n"; + return 1; + } + expect_at_least("the larger of two live requests is in force", limit(), large); + } + expect_at_least("the survivor of two requests keeps its own", limit(), small); + expect_below("the departed request is no longer in force", limit(), large); + } + expect_equal("the last release restores the limit it found", limit(), baseline); + + // Two handles asking for the same size are two requests, not one: withdrawing either must + // leave the other covered. This is bookkeeping a set would get wrong and a multiset gets right. + { + ninfer::l2p::Reservation first; + ninfer::l2p::Reservation second; + if (!first.request(large) || !second.request(large)) { + std::cout << "two handles of equal size were not both admitted\n"; + return 1; + } + first.release(); + expect_at_least("equal-sized twin still covered after its peer leaves", limit(), large); + } + expect_equal("equal-sized twins restore the limit once both leave", limit(), baseline); + + // Re-asking the same size must not ratchet the recorded baseline: it is read when the first + // request goes in, and a cycle of release and request reads it again. + { + ninfer::l2p::Reservation handle; + for (int cycle = 0; cycle < 32; ++cycle) { + if (!handle.request(large)) { + std::cout << "a repeated request was refused on cycle " << cycle << '\n'; + return 1; + } + expect_at_least("a repeated request stays in force", limit(), large); + handle.release(); + } + } + expect_equal("32 request/release cycles do not drift the baseline", limit(), baseline); + + // Zero is not a release. A handle that is asked for nothing keeps what its graph is running + // under, and says it did not take the request. + { + ninfer::l2p::Reservation handle; + if (!handle.request(large)) { + std::cout << "a request before the zero case was refused\n"; + return 1; + } + if (handle.request(0)) { + std::cout << "a zero-byte request was reported as taken\n"; + ++failures; + } + expect_at_least("a zero-byte request leaves the live one alone", limit(), large); + } + expect_equal("the zero case restores the limit on release", limit(), baseline); + + { + ninfer::l2p::Reservation source; + if (!source.request(large)) { + std::cout << "a request before the move case was refused\n"; + return 1; + } + ninfer::l2p::Reservation sink(std::move(source)); + expect_at_least("a moved handle carries its request", limit(), large); + source.release(); // The moved-from handle owns nothing and must not disturb the limit. + expect_at_least("releasing a moved-from handle changes nothing", limit(), large); + } + expect_equal("the moved-to handle restores the limit when it dies", limit(), baseline); + + // A live handle asked for a different size replaces its own request rather than adding one. + // Nothing in the product exercises this today, but the header advertises it, and it is the + // subtlest path in the file: the new request goes in before the old one comes out. + { + ninfer::l2p::Reservation handle; + if (!handle.request(large) || !handle.request(small)) { + std::cout << "a handle could not replace its own request\n"; + return 1; + } + expect_at_least("a replacement covers what it asked for", limit(), small); + expect_below("the replaced request is not left standing", limit(), large); + } + expect_equal("the replacement case restores the limit on release", limit(), baseline); + + // A request the device cannot honour must be refused without disturbing a live one. This is + // the only case that reaches the rollback in request(), because nothing else here makes the + // write fail, and a rollback that gives back another owner's baseline is exactly the leak + // this file exists to prevent. + { + ninfer::l2p::Reservation live_one; + if (!live_one.request(large)) { + std::cout << "a request before the refusal case was refused\n"; + return 1; + } + ninfer::l2p::Reservation refused; + if (refused.request(ceiling * 2)) { + std::cout << "a request beyond the device ceiling was reported as taken\n"; + ++failures; + } + expect_at_least("a refused request leaves the live one covered", limit(), large); + } + expect_equal("a refused request leaves nothing behind", limit(), baseline); + + // A refused request must not leave a baseline of its own behind. Nothing already here can see + // that it did: the registry looks the same either way and the limit never moved. It becomes + // visible only when the process already holds a set-aside of its own by the time the next + // request arrives - then a baseline stashed by the refusal is handed back instead of that one. + { + ninfer::l2p::Reservation refused_cold; + if (refused_cold.request(ceiling * 2)) { + std::cout << "a cold request beyond the device ceiling was reported as taken\n"; + ++failures; + } + } + if (cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, large) == cudaSuccess) { + // Whatever the driver rounded that to is this process's own set-aside now, and it is what + // a later release owes back - not the limit that was in force before the refusal. + const std::size_t theirs = limit(); + { + ninfer::l2p::Reservation ours; + if (!ours.request(small)) { + std::cout << "a request after the cold refusal was refused\n"; + return 1; + } + } + expect_equal("a refused request does not stale the baseline", limit(), theirs); + (void)cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, baseline); + } + expect_equal("the cold refusal case leaves the limit as it found it", limit(), baseline); + + // The handle's own fields are part of what the registry lock protects. Before that was true, + // two threads on one handle could publish two requests and withdraw one, stranding the device + // at a raised limit with no window installed - the exact state this file exists to prevent. + { + ninfer::l2p::Reservation shared; + std::vector threads; + threads.reserve(4); + for (int t = 0; t < 4; ++t) { + threads.emplace_back([&shared, large]() { + for (int i = 0; i < 2000; ++i) { + (void)shared.request(large); + shared.release(); + } + }); + } + for (std::thread& thread : threads) { thread.join(); } + shared.release(); + } + expect_equal("concurrent use of one handle leaves no request behind", limit(), baseline); + + if (failures != 0) { + std::cout << failures << " L2 reservation checks failed\n"; + return 1; + } + std::cout << "ok\n"; + return 0; +}