kv: thread compute stream through activate() so membership publish is stream-ordered

[ranxianglei]

Fixes the race described in #210.

commit_activation() defaulted to the legacy default stream (stream 0) on the bind_sequence_kv() -> activate() paths, so the paged-cache membership publish memcpy was unordered vs the compute stream. Every conversational turn that re-activates a retained KV catalog re-opened the race; real agent traffic (multi-turn + concurrent second lane) crashed with a device-side assert at device.cu:132 after ~54 requests, locking up the whole GPU.

Fix (3 lines): thread cudaStream_t through activate() and pass device.stream at both call sites — matching the existing correct usage at program_impl.h:9714.

Verification: replay of the exact real-world ticket that crashed at request #54 → 66 requests through the original crash point, server alive, full 102-line diagnostic output, quality unchanged. Six synthetic reproducers never triggered the crash before or after (they lack the retained-endpoint + catalog re-activation + concurrent-lane combination, which is why the bug hid from them).

Caveats and remaining suspects (page-release fencing, staged tail COW) are documented in #210 — not claimed fixed here.

[Xtravaganz]

save and crash test case, can't test it @ the moment, llm is in use.

tests\test_kv_cache.cpp
line 420-581

// Demonstrates that release_page() returns a physical page to the free list without
// zeroing or protecting its content.  Stale user data persists across release + re-allocate
// on the same physical index — this is a structural precondition for the dangerous race.
int exercise_page_release_fence(ninfer::DeviceContext& context) {
    int failures = 0;
    ninfer::KVPageGeometry geometry{
        .planes = {{ninfer::DType::I8, 8, 2, 256}},
    };
    PlannedCache plan = plan_cache(2, 2, 1, geometry);
    ninfer::DeviceArena arena(plan.bytes);
    ninfer::DeviceKVPagePool pool({arena.base(), arena.capacity()}, plan.pages);

    const ninfer::HostKVPageLayout host_layout =
        ninfer::plan_host_kv_page_layout(pool.geometry());
    const ninfer::HostKVPageLayout layouts[] = {host_layout};
    ninfer::HostKVArena host_arena(host_layout.page_stride * 6, layouts);

    // 1. Allocate page 0 and write known data.
    std::vector<ninfer::DeviceKVPageLease> pages = materialize(pool, 1);
    std::optional<ninfer::HostKVAllocation> write =
        host_arena.allocate(host_layout, 1);
    ninfer::HostKVAllocationView write_view = host_arena.writable_view(*write);
    std::memset(write_view.data(), 0xAB, host_layout.page_stride);
    pool.copy_from_host(host_arena.view(*write),
                        std::span<const ninfer::DeviceKVPageHandle>(&pages[0].handle(), 1),
                        context.stream);
    context.synchronize();

    // 2. Release and re-allocate — gets the same physical index.
    pages[0].release();
    pages.clear();
    pages = materialize(pool, 1);

    // 3. Read back without writing new data.
    std::optional<ninfer::HostKVAllocation> readback =
        host_arena.allocate(host_layout, 1);
    ninfer::HostKVAllocationView readback_view = host_arena.writable_view(*readback);
    std::memset(readback_view.data(), 0, host_layout.page_stride);
    pool.copy_to_host(std::span<const ninfer::DeviceKVPageHandle>(&pages[0].handle(), 1),
                      readback_view, context.stream);
    context.synchronize();

    // 4. The re-allocated page still holds the old (0xAB) data — release_page()
    //    does not clear or fence the page.
    const ninfer::HostKVAllocationConstView readback_contents = host_arena.view(*readback);
    failures += expect(page_payload_equal(readback_contents, 0, host_arena.view(*write), 0),
                       "re-allocated page does not contain stale data from prior lease; "
                       "pool zeroed or invalidated the page on release");

    (void)write_view;
    (void)readback_view;
    return failures;
}

// Dangerous-only test: writes to the same physical page from two independent
// non-blocking streams without ordering, reproducing the Xid-79 crash pattern.
// The transfer-stream write is still in-flight when release_page() returns the
// page to the free list.  After re-allocation, the compute stream writes new
// data to the same GPU address.  This concurrent-write is undefined behaviour;
// on Blackwell it can produce cudaErrorLaunchFailure → Xid-79 → GPU lockup.
//
// Run with --dangerous on hardware that can tolerate a GPU reset.
int exercise_page_release_race(ninfer::DeviceContext& context) {
    int failures = 0;
    ninfer::KVPageGeometry geometry{
        .planes = {{ninfer::DType::I8, 8, 2, 256}},
    };
    PlannedCache plan = plan_cache(2, 2, 1, geometry);
    ninfer::DeviceArena arena(plan.bytes);
    ninfer::DeviceKVPagePool pool({arena.base(), arena.capacity()}, plan.pages);

    const ninfer::HostKVPageLayout host_layout =
        ninfer::plan_host_kv_page_layout(pool.geometry());
    const ninfer::HostKVPageLayout layouts[] = {host_layout};
    ninfer::HostKVArena host_arena(host_layout.page_stride * 6, layouts);

    // 1. Allocate page, write 0xAB on transfer_stream — do NOT synchronize.
    std::vector<ninfer::DeviceKVPageLease> pages = materialize(pool, 1);
    std::optional<ninfer::HostKVAllocation> old =
        host_arena.allocate(host_layout, 1);
    ninfer::HostKVAllocationView old_view = host_arena.writable_view(*old);
    std::memset(old_view.data(), 0xAB, host_layout.page_stride);
    pool.copy_from_host(host_arena.view(*old),
                        std::span<const ninfer::DeviceKVPageHandle>(&pages[0].handle(), 1),
                        context.transfer_stream);

    // 2. Release — no fence, page returns to free list while transfer_stream
    //    write is still in-flight.
    pages[0].release();
    pages.clear();

    // 3. Re-allocate same physical index, write 0xCD on context.stream.
    //    Both streams may issue concurrent writes to the same GPU address.
    pages = materialize(pool, 1);
    std::optional<ninfer::HostKVAllocation> fresh =
        host_arena.allocate(host_layout, 1);
    ninfer::HostKVAllocationView fresh_view = host_arena.writable_view(*fresh);
    std::memset(fresh_view.data(), 0xCD, host_layout.page_stride);
    pool.copy_from_host(host_arena.view(*fresh),
                        std::span<const ninfer::DeviceKVPageHandle>(&pages[0].handle(), 1),
                        context.stream);
    context.synchronize();

    // 4. Read back (may crash before reaching here).
    std::optional<ninfer::HostKVAllocation> readback =
        host_arena.allocate(host_layout, 1);
    ninfer::HostKVAllocationView readback_view = host_arena.writable_view(*readback);
    std::memset(readback_view.data(), 0, host_layout.page_stride);
    pool.copy_to_host(std::span<const ninfer::DeviceKVPageHandle>(&pages[0].handle(), 1),
                      readback_view, context.stream);
    context.synchronize();

    // 5. If we survive, verify the payload matches the new write.  A stale 0xAB
    //    would prove the transfer-stream write raced past release + re-allocate.
    const ninfer::HostKVAllocationConstView rb = host_arena.view(*readback);
    const ninfer::HostKVAllocationConstView fv = host_arena.view(*fresh);
    failures += expect(page_payload_equal(rb, 0, fv, 0),
                       "page payload after release+reallocate showed stale transfer-stream "
                       "data; release_page() is missing a CUDA ordering fence");

    (void)old_view;
    (void)fresh_view;
    (void)readback_view;
    return failures;
}

} // namespace

int main(int argc, char* argv[]) {
    int device_count              = 0;
    const cudaError_t count_error = cudaGetDeviceCount(&device_count);
    if (cuda_unavailable(count_error) || (count_error == cudaSuccess && device_count == 0)) {
        std::cout << "SKIP: no usable CUDA device\n";
        return 77;
    }
    if (count_error != cudaSuccess) {
        std::cerr << "cudaGetDeviceCount failed: " << cudaGetErrorString(count_error) << '\n';
        return 1;
    }

    bool dangerous = false;
    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--dangerous") == 0) { dangerous = true; }
    }

    try {
        ninfer::DeviceContext context(0);
        int failures = 0;

        std::cout << "  page_release_fence (safe) ... ";
        failures += exercise_page_release_fence(context);
        std::cout << (failures == 0 ? "PASS" : "FAIL") << '\n';

        std::cout << "  page_release_race (dangerous) ... ";
        if (dangerous) {
            failures += exercise_page_release_race(context);
            std::cout << (failures == 0 ? "PASS" : "FAIL") << '\n';
        } else {
            std::cout << "SKIP  (use --dangerous to enable)\n";
        }

        failures += exercise_reservation_and_mapping(context);

Here is the final shape:

exercise_page_release_fence() (line 423) safe, deterministic, runs always. Writes 0xAB to a page, syncs, releases, re-allocates, reads back, asserts the stale 0xAB data is still there. This proves release_page() does not zero or fence the page content, a precondition for the dangerous race.

exercise_page_release_race() (line 482) gated behind --dangerous . Writes 0xAB on transfer_stream , releases the page without any sync, re-allocates, writes 0xCD on context.stream . Two cudaStreamNonBlocking streams issuing concurrent writes to the same GPU address. This is the actual Xid-79 reproducer, it can crash the GPU on Blackwell (sm_120a).

main() (line 548): parses --dangerous from argv . Safe test runs unconditionally; dangerous test prints SKIP unless --dangerous is passed.

[ranxianglei]

Ran your safe test on our hardware (RTX Pro 6000 Blackwell, sm_120a, CUDA 13.0, gcc):

page_release_fence (safe) ... PASS
page_release_race (dangerous) ... SKIP  (use --dangerous to enable)
Paged KV physical-container checks passed

Fence test PASS confirms the structural precondition on real hardware: the re-allocated page still carries the stale 0xAB payload — release_page() neither zeroes nor fences. Exactly as you described.

Two notes:

1. Portability: your snippet doesn't compile on gcc — &pages[0].handle() is address-of-rvalue at 6 call sites (MSVC is more permissive). I added a one_page_span() helper (single DeviceKVPageHandle storage + span), no behavior change. Committed at 93a6088 alongside your tests verbatim otherwise.

2. Dangerous race: that one is the actual Xid-79 reproducer and we have two inference services live on the only spare GPU right now. We'll run --dangerous in a free-GPU window (expected: crash or stale-0xAB payload → either confirms the missing ordering fence at the pool layer, complementing the commit_activation stream fix already in this PR which addresses the publish path). Will report back.

[ranxianglei]

Free-GPU window secured — ran the dangerous race test on the same hardware (RTX Pro 6000 Blackwell, sm_120a, CUDA 13.0):

page_release_fence (safe) ... PASS
page_release_race (dangerous) ... PASS   × 3 runs (15,000 interleavings total)
Paged KV physical-container checks passed

Result: did not reproduce — no Xid, no GPU lockup, no stale-0xAB contamination detected by the readback check, across 3 independent runs at 5,000 repeats each. GPU healthy afterwards, services restored.

Reading of the negative result (offered with the usual caveat that absence of reproduction ≠ absence of the race):

• The safe test proves the structural precondition on this hardware: released pages keep stale payload. So the precondition half of your scenario is confirmed.
• The interleaving half (delayed 0xAB write landing after the fresh 0xCD write on the same physical address) never materialized in 15k trials. Plausible causes: (a) same-address write-write ordering is enforced at the L2/memory pipeline level on Blackwell, so the hazard may only bite on cross-page or read-modify-write patterns rather than same-address blind overwrites; (b) allocator reuse timing — the just-released page goes straight back to the same slot, so the two streams' windows rarely interleave as drawn; (c) genuine race with probability < 1/15k under this driver stack.

If you want to chase it further, a stronger variant would be one of: escalating page churn (fill the whole pool between release and realloc so the physical slot changes), adding a concurrent reader of the old logical slot, or write via a different aperture (host-mapped vs device). Happy to run any of those in the next window — they're all drop-in modifications of your harness, which by the way is excellent: the --dangerous gate is exactly the right shape for fleet hardware.

[Xtravaganz]

Ran your safe test on our hardware (RTX Pro 6000 Blackwell, sm_120a, CUDA 13.0, gcc):

page_release_fence (safe) ... PASS
page_release_race (dangerous) ... SKIP  (use --dangerous to enable)
Paged KV physical-container checks passed

Fence test PASS confirms the structural precondition on real hardware: the re-allocated page still carries the stale 0xAB payload — release_page() neither zeroes nor fences. Exactly as you described.

Two notes:

1. Portability: your snippet doesn't compile on gcc — &pages[0].handle() is address-of-rvalue at 6 call sites (MSVC is more permissive). I added a one_page_span() helper (single DeviceKVPageHandle storage + span), no behavior change. Committed at 93a6088 alongside your tests verbatim otherwise.
2. Dangerous race: that one is the actual Xid-79 reproducer and we have two inference services live on the only spare GPU right now. We'll run --dangerous in a free-GPU window (expected: crash or stale-0xAB payload → either confirms the missing ordering fence at the pool layer, complementing the commit_activation stream fix already in this PR which addresses the publish path). Will report back.

Will you include the tests in this MR, or should I create a separate test MR for them? Of course, it should include the points mentioned above. Thank you!

[ranxianglei]

Split the tests into their own branch as you suggested: ranxianglei:kv-release-tests (commit 93a6088 — your test logic verbatim + the one_page_span() portability helper).

Odd GitHub hiccup: opening a new PR from my fork to this repo is being rejected right now (GraphQL permission error / REST 404 — token scopes and the branch are fine; PR #213 was created from the same fork this morning). Rather than fight it, feel free to open the MR from that branch yourself, or cherry-pick it — either works for me.

Suggested MR description (use freely):

Tests from #211 review by @Xtravaganz. Safe fence test: PASS on Blackwell sm_120a. Dangerous race test: 3 runs × 5,000 interleavings did not reproduce (no Xid, no contamination) — readback negative-result analysis in the #211 thread. gcc portability fix included (address-of-rvalue at 6 sites).

[ranxianglei]

Apologies for the earlier detour — your tests are now part of this PR directly (no cherry-picking needed on your side):

• 0687a66 adds the tests verbatim from your review, with the one_page_span() helper so they compile on gcc (avoids the &pages[0].handle() address-of-rvalue issue you flagged — value captured into a named local, address taken from that).

Also cleaned up the branch history while I was here: this PR's head previously carried an unrelated W8 commit that belongs to #213; both branches are now rebased onto current master and contain only their own changes (#211 = the stream-ordering fix + these tests, #213 = the W8 variant). Diff should be easier to review now (3 files).

Safe-mode verification on RTX Pro 6000 (sm_120a, CUDA 13.0, gcc): page_release_fence (safe) ... PASS, race variant left behind --dangerous as intended.