From 03e2b91db69d507235c2f18787eeffcc3e1afef1 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 21 Oct 2026 01:16:10 -0700 Subject: [PATCH] [executorch][muse] Fix the DFlash final-block draft counts test CudaDFlashNative.FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows has failed in unittest-cuda-muse-glimmer-batching since #24043 added draft forward count checks to it: Expected: (drained.draft_forwards - drain.draft_forwards) > (drained.static_draft_forwards - drain.static_draft_forwards), actual: 1 vs 1 The test expected the decode at position 78 to drain the 79 hidden rows dynamically before the static draft. But run_target drains every owner it touches right after appending hidden rows (#24042), so the 78-row prefill already drains down to the fresh context (32). The decode then holds at most one block of rows and takes a single static draft, and the dynamic forwards fall before the window the test measured. Measure from before the prefill instead: the prefill drains through dynamic drafts only, and the decode runs exactly one static draft. --- .../tests/cuda_dflash_native_test.cpp | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp index 80ed5e774e0..05deebd100a 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp @@ -598,20 +598,23 @@ TEST_F(CudaDFlashNative, FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows) { ASSERT_TRUE(executor_); const auto sid = open(); b::BatchOutput out; + const auto start = executor_->stats(); ASSERT_TRUE(executor_->execute({{input(sid, 0, Tokens(78, 1), false)}}, out)); Peer::check(*executor_, sid, 78); EXPECT_EQ(executor_->stats().hidden_rows_d2h, 78); + // 78 prefill rows exceed the fresh context, so the prefill itself drains + // them through the dynamic draft. + const auto prefilled = executor_->stats(); + EXPECT_GT(prefilled.draft_forwards - start.draft_forwards, 0); + EXPECT_EQ(prefilled.static_draft_forwards, start.static_draft_forwards); fixture_->rejected.fill(0); - const auto drain = executor_->stats(); ASSERT_TRUE(executor_->execute({{input(sid, 78, {1})}}, out)); EXPECT_EQ(out.outputs[0]->tokens, predictions(78, 2)); Peer::check(*executor_, sid, 80); - // 79 rows exceed one static block: drain dynamically, then draft statically. - const auto drained = executor_->stats(); - EXPECT_EQ(drained.static_draft_forwards - drain.static_draft_forwards, 1); - EXPECT_GT( - drained.draft_forwards - drain.draft_forwards, - drained.static_draft_forwards - drain.static_draft_forwards); + // The drained remainder fits one block: the decode drafts statically. + const auto drafted = executor_->stats(); + EXPECT_EQ(drafted.draft_forwards - prefilled.draft_forwards, 1); + EXPECT_EQ(drafted.static_draft_forwards - prefilled.static_draft_forwards, 1); fixture_->rejected.fill(15); const auto before = executor_->stats(); ASSERT_TRUE(executor_->execute({{input(sid, 80, {prediction(79)})}}, out));