From 4b8dd8fdd2990c3279e2c1c0d05185120f13fced Mon Sep 17 00:00:00 2001 From: Vikash Loomba Date: Wed, 12 Aug 2026 09:20:48 -0700 Subject: [PATCH 1/2] test(rocm): ReshapeAndCache->PagedAttention composition at real dims (#41) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The in-tree paged-attention case hand-builds a contiguous KV cache; the real model path writes it with ReshapeAndCache and reads it back. This case is that composition at real model dims (Dh=256, Hq=8, Hkv=2, block_size 16) with a shuffled block table and non-sequential slot mapping — the layout a stride/scatter bug would live in and the contiguous case cannot see. Surfaced by the #41 Qwen3.5-0.8B divergence investigation: every compositional piece of the ROCm attention path now validates in isolation, which is what localizes the residual divergence to bf16-softmax accumulation rather than a kernel defect. Evidence (4x gfx1100, ROCm 7.14, Release): the new case passes 7/7 vs the CPU oracle; full cross-device suite green. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: pi:kimi-k3 [pi] --- tests/vt/test_backend_cross_device.cpp | 80 ++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index cd3820c63..3071d2992 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1980,6 +1980,86 @@ TEST_CASE("MoeRouterTopK matches the CPU oracle (f32 and bf16 logits)") { } } +TEST_CASE("ReshapeAndCache->PagedAttention composition matches CPU (real dims, shuffled blocks)") { + // The "paged attention" case above hand-builds a contiguous KV cache; the + // real model path writes it with ReshapeAndCache and reads it back. This + // case is that composition, at real model dims (Dh=256, Hq=8, Hkv=2, + // block_size 16), a shuffled block table, and a non-sequential slot mapping + // — the layout a stride/scatter bug would live in and the contiguous case + // cannot see. + constexpr int64_t T = 20, Hq = 8, Hkv = 2, Dh = 256, BS = 16; + constexpr int64_t kBlocks = 4; // 4 blocks x 16 slots = 64 >= 20 + constexpr int64_t qstride = Hq * Dh + 64; // padded row (fused-view shape) + const size_t qn = static_cast(T) * qstride; + const size_t kvn = static_cast(T) * Hkv * Dh; + const size_t cachen = static_cast(kBlocks) * BS * Hkv * Dh; + const std::vector q = RandomVec(qn, 711); + const std::vector k = RandomVec(kvn, 712); + const std::vector v = RandomVec(kvn, 713); + // Non-sequential slot mapping (reverse-ish) to exercise the scatter. + std::vector slots(T); + for (int64_t i = 0; i < T; ++i) slots[i] = (i * 7 + 3) % (kBlocks * BS); + std::vector block_table = {3, 1, 2, 0}; // shuffled physical blocks + std::vector seq_lens = {T}; + std::vector qsl = {0, T}; + vt::PagedAttentionArgs pa; + pa.scale = 1.0f / std::sqrt(static_cast(Dh)); + pa.causal = true; + + std::vector ref_out(static_cast(T) * Hq * Dh, 0.0f); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector q_host = q, ck = k, cv = v; + std::vector ckc(cachen, 0.0f), cvc(cachen, 0.0f); + std::vector cslots = slots; + std::vector cbt = block_table, csl = seq_lens, cqsl = qsl; + Tensor tq = Tensor::Contiguous(q_host.data(), DType::kF32, cd, {T, Hq, Dh}); // contiguous (op contract) + Tensor tk = Tensor::Contiguous(ck.data(), DType::kF32, cd, {T, Hkv, Dh}); + Tensor tv = Tensor::Contiguous(cv.data(), DType::kF32, cd, {T, Hkv, Dh}); + Tensor tkc = Tensor::Contiguous(ckc.data(), DType::kF32, cd, {kBlocks, BS, Hkv, Dh}); + Tensor tvc = Tensor::Contiguous(cvc.data(), DType::kF32, cd, {kBlocks, BS, Hkv, Dh}); + Tensor tsm = Tensor::Contiguous(cslots.data(), DType::kI64, cd, {T}); + vt::ReshapeAndCache(cq, tk, tv, tkc, tvc, tsm); + Tensor tbt = Tensor::Contiguous(cbt.data(), DType::kI32, cd, {1, kBlocks}); + Tensor tsl = Tensor::Contiguous(csl.data(), DType::kI32, cd, {1}); + Tensor tqsl = Tensor::Contiguous(cqsl.data(), DType::kI32, cd, {2}); + Tensor to = Tensor::Contiguous(ref_out.data(), DType::kF32, cd, {T, Hq, Dh}); + vt::PagedAttention(cq, to, tq, tkc, tvc, tbt, tsl, tqsl, pa); + cpu.DestroyQueue(cq); + } + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kPagedAttention, dt) || !OpAvailable(vt::OpId::kReshapeAndCache, dt)) continue; + CAPTURE(DeviceName(dt)); + vt::Backend& dev = vt::GetBackend(dt); + Queue q_ = dev.CreateQueue(); + const Device d{dt, 0}; + DevBuf dq(dev, q_, qn), dk(dev, q_, kvn), dv(dev, q_, kvn), + dkc(dev, q_, cachen), dvc(dev, q_, cachen), dout(dev, q_, static_cast(T) * Hq * Dh); + DevBufBytes dsm(dev, q_, T * 8), dbt(dev, q_, kBlocks * 4), dsl_(dev, q_, 4), dqsl(dev, q_, 8); + dq.Upload(q); dk.Upload(k); dv.Upload(v); + dkc.Upload(std::vector(cachen, 0.0f)); dvc.Upload(std::vector(cachen, 0.0f)); + dsm.Upload(slots.data()); dbt.Upload(block_table.data()); + dsl_.Upload(seq_lens.data()); dqsl.Upload(qsl.data()); + Tensor tq = Tensor::Contiguous(dq.ptr(), DType::kF32, d, {T, Hq, Dh}); // contiguous (op contract) + Tensor tk = Tensor::Contiguous(dk.ptr(), DType::kF32, d, {T, Hkv, Dh}); + Tensor tv = Tensor::Contiguous(dv.ptr(), DType::kF32, d, {T, Hkv, Dh}); + Tensor tkc = Tensor::Contiguous(dkc.ptr(), DType::kF32, d, {kBlocks, BS, Hkv, Dh}); + Tensor tvc = Tensor::Contiguous(dvc.ptr(), DType::kF32, d, {kBlocks, BS, Hkv, Dh}); + Tensor tsm = Tensor::Contiguous(dsm.ptr(), DType::kI64, d, {T}); + vt::ReshapeAndCache(q_, tk, tv, tkc, tvc, tsm); + Tensor tbt = Tensor::Contiguous(dbt.ptr(), DType::kI32, d, {1, kBlocks}); + Tensor tsl = Tensor::Contiguous(dsl_.ptr(), DType::kI32, d, {1}); + Tensor tqsl = Tensor::Contiguous(dqsl.ptr(), DType::kI32, d, {2}); + Tensor to = Tensor::Contiguous(dout.ptr(), DType::kF32, d, {T, Hq, Dh}); + vt::PagedAttention(q_, to, tq, tkc, tvc, tbt, tsl, tqsl, pa); + CHECK(Nmse(ref_out, dout.Download()) <= kNmseTol); + dev.DestroyQueue(q_); + } +} + + TEST_CASE("reference tier: an op with no native kernel matches the CPU oracle (unified only)") { constexpr int64_t kRows = 7, kCols = 48; constexpr size_t kN = kRows * kCols; From ce40575b7548e57a071fe8357e9c6bbc29b8c268 Mon Sep 17 00:00:00 2001 From: Vikash Loomba Date: Fri, 14 Aug 2026 01:45:39 -0700 Subject: [PATCH 2/2] fix(ROCM): the composition test's slot mapping derives from the block table + an anti-vacuity guard -- the #497 review rework MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CHANGES_REQUESTED review (localai-org-maint-bot, 2026-08-14), both findings accepted: 1. The first version scattered token i to slot (i*7+3)%64 while the attention read walked logical positions through block table {3,1,2,0} — the write set and the read set were disjoint by construction, so most positions compared unwritten zeros on both backends and a scatter/layout defect could not have failed the case. The slot mapping now derives from the logical position through the shuffled table (slots[i] = block_table[i/BS]*BS + i%BS), which is the engine's real mapping. 2. Anti-vacuity guard: a corrupted block table (the first two LOGICAL blocks swapped — both hold real tokens) must change the attention output; if the composition compared only unwritten slots it would not. (The first draft of the guard swapped two blocks OUTSIDE the logical range and was itself vacuous — the guard proved the guard; the committed version swaps the in-range mapping.) Also removed the unused padded qstride — the query tensor is declared contiguous, so the padding exercised nothing. Gates (gfx1100, flock): test_backend_cross_device 20/20 with the corrected composition. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: pi:kimi-k3 [pi] --- tests/vt/test_backend_cross_device.cpp | 39 ++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index 3071d2992..153319a1b 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1989,17 +1989,22 @@ TEST_CASE("ReshapeAndCache->PagedAttention composition matches CPU (real dims, s // cannot see. constexpr int64_t T = 20, Hq = 8, Hkv = 2, Dh = 256, BS = 16; constexpr int64_t kBlocks = 4; // 4 blocks x 16 slots = 64 >= 20 - constexpr int64_t qstride = Hq * Dh + 64; // padded row (fused-view shape) - const size_t qn = static_cast(T) * qstride; + const size_t qn = static_cast(T) * Hq * Dh; const size_t kvn = static_cast(T) * Hkv * Dh; const size_t cachen = static_cast(kBlocks) * BS * Hkv * Dh; const std::vector q = RandomVec(qn, 711); const std::vector k = RandomVec(kvn, 712); const std::vector v = RandomVec(kvn, 713); - // Non-sequential slot mapping (reverse-ish) to exercise the scatter. - std::vector slots(T); - for (int64_t i = 0; i < T; ++i) slots[i] = (i * 7 + 3) % (kBlocks * BS); + // The slot mapping must DERIVE from the logical position through the + // (shuffled) block table — exactly what the engine produces — otherwise the + // attention read of logical position p lands on a slot nothing wrote and + // both backends compare zeros (review on #497: the first version's + // (i*7+3)%64 scatter was disjoint from the block table, so the composition + // exercised mostly-unwritten cache). std::vector block_table = {3, 1, 2, 0}; // shuffled physical blocks + std::vector slots(T); + for (int64_t i = 0; i < T; ++i) + slots[i] = static_cast(block_table[static_cast(i / BS)]) * BS + (i % BS); std::vector seq_lens = {T}; std::vector qsl = {0, T}; vt::PagedAttentionArgs pa; @@ -2055,6 +2060,30 @@ TEST_CASE("ReshapeAndCache->PagedAttention composition matches CPU (real dims, s Tensor to = Tensor::Contiguous(dout.ptr(), DType::kF32, d, {T, Hq, Dh}); vt::PagedAttention(q_, to, tq, tkc, tvc, tbt, tsl, tqsl, pa); CHECK(Nmse(ref_out, dout.Download()) <= kNmseTol); + + // Anti-vacuity guard (review on #497): a WRONG physical mapping must NOT + // reproduce the reference — if the composition were vacuous (reads never + // hitting writes), a corrupted table would compare equal. Blocks 0 and 2 + // both carry real tokens under the true table, so swapping them must + // change the output. + // Swap the mapping of the first two LOGICAL blocks — both hold real + // tokens (0-15 and 16-19), so the read path changes. (The first version + // of this guard swapped two blocks OUTSIDE the logical range and was + // itself vacuous — the guard proved the guard.) + std::vector bad_table = {1, 3, 2, 0}; + DevBufBytes dbt_bad(dev, q_, kBlocks * 4); + dbt_bad.Upload(bad_table.data()); + DevBuf dout2(dev, q_, static_cast(T) * Hq * Dh); + Tensor tbt2 = Tensor::Contiguous(dbt_bad.ptr(), DType::kI32, d, {1, kBlocks}); + Tensor to2 = Tensor::Contiguous(dout2.ptr(), DType::kF32, d, {T, Hq, Dh}); + vt::PagedAttention(q_, to2, tq, tkc, tvc, tbt2, tsl, tqsl, pa); + const std::vector bad_out = dout2.Download(); + bool any_diff = false; + for (size_t i = 0; i < ref_out.size(); ++i) + if (std::fabs(bad_out[i] - ref_out[i]) > 1e-3f) { any_diff = true; break; } + CHECK_MESSAGE(any_diff, + "a corrupted block table must change the attention output — " + "otherwise the composition test is vacuous"); dev.DestroyQueue(q_); } }