@@ -156,6 +156,122 @@ bool copy_host(
156156 cudaStreamSynchronize (stream) == cudaSuccess;
157157}
158158
159+ } // namespace
160+
161+ namespace cuda_dflash_detail {
162+ HiddenSnapshot::~HiddenSnapshot () {
163+ // Queued copies may still read the staging if the fence failed; leak it.
164+ if (fence () && device_) {
165+ (void )cudaFree (device_);
166+ }
167+ }
168+
169+ bool HiddenSnapshot::fence () {
170+ if (stream_set_ &&
171+ cudaStreamSynchronize (static_cast <cudaStream_t>(stream_)) !=
172+ cudaSuccess) {
173+ return false ;
174+ }
175+ stream_set_ = false ;
176+ return true ;
177+ }
178+
179+ bool HiddenSnapshot::begin (std::size_t rows, std::size_t row_bytes) {
180+ if (!fence () || rows == 0 || row_bytes == 0 ||
181+ rows > std::numeric_limits<std::size_t >::max () / row_bytes) {
182+ return false ;
183+ }
184+ rows_ = rows;
185+ row_bytes_ = row_bytes;
186+ captured_ = 0 ;
187+ device_source_.reset ();
188+ stream_ = nullptr ;
189+ stream_set_ = false ;
190+ return true ;
191+ }
192+
193+ bool HiddenSnapshot::capture (
194+ const aten::Tensor& taps,
195+ std::size_t offset,
196+ std::size_t rows) {
197+ const bool device = taps.device_type () == aten::DeviceType::CUDA ;
198+ if (offset != captured_ || offset > rows_ || rows == 0 ||
199+ rows > rows_ - offset || !dense (taps) ||
200+ taps.nbytes () < rows * row_bytes_ ||
201+ (!device && taps.device_type () != aten::DeviceType::CPU ) ||
202+ (device_source_ && *device_source_ != device)) {
203+ return false ;
204+ }
205+ device_source_ = device;
206+ const auto bytes = rows * row_bytes_;
207+ if (device) {
208+ const auto stream =
209+ ::executorch::extension::cuda::getCallerStream ().value_or(
210+ cudaStreamPerThread);
211+ if (stream_set_ && stream_ != stream) {
212+ return false ;
213+ }
214+ stream_ = stream;
215+ stream_set_ = true ;
216+ if (capacity_ < rows_ * row_bytes_) {
217+ if (device_ && cudaFree (device_) != cudaSuccess) {
218+ return false ;
219+ }
220+ device_ = nullptr ;
221+ capacity_ = 0 ;
222+ if (cudaMalloc (&device_, rows_ * row_bytes_) != cudaSuccess) {
223+ return false ;
224+ }
225+ capacity_ = rows_ * row_bytes_;
226+ }
227+ if (cudaMemcpyAsync (
228+ static_cast <std::uint8_t *>(device_) + offset * row_bytes_,
229+ taps.const_data_ptr (),
230+ bytes,
231+ cudaMemcpyDeviceToDevice,
232+ stream) != cudaSuccess) {
233+ return false ;
234+ }
235+ } else {
236+ host_.resize (rows_ * row_bytes_);
237+ std::memcpy (
238+ host_.data () + offset * row_bytes_, taps.const_data_ptr (), bytes);
239+ }
240+ captured_ += rows;
241+ return true ;
242+ }
243+
244+ bool HiddenSnapshot::copy_rows (
245+ std::vector<std::uint8_t >& destination,
246+ std::size_t offset,
247+ std::size_t rows) {
248+ if (offset > captured_ || rows > captured_ - offset || row_bytes_ == 0 ||
249+ destination.size () % row_bytes_ != 0 ||
250+ rows > (destination.max_size () - destination.size ()) / row_bytes_) {
251+ return false ;
252+ }
253+ const auto start = destination.size ();
254+ const auto bytes = rows * row_bytes_;
255+ destination.resize (start + bytes);
256+ if (bytes == 0 ) {
257+ return true ;
258+ }
259+ if (on_device ()) {
260+ stream_set_ = true ;
261+ return cudaMemcpyAsync (
262+ destination.data () + start,
263+ static_cast <std::uint8_t *>(device_) + offset * row_bytes_,
264+ bytes,
265+ cudaMemcpyDeviceToHost,
266+ static_cast <cudaStream_t>(stream_)) == cudaSuccess;
267+ }
268+ std::memcpy (
269+ destination.data () + start, host_.data () + offset * row_bytes_, bytes);
270+ return true ;
271+ }
272+ } // namespace cuda_dflash_detail
273+
274+ namespace {
159275Result<int > positive_int (Module& module , const char * name) {
160276 ET_ASSIGN_OR_RETURN (value, detail::read_required_positive_int (module , name));
161277 ET_CHECK_OR_RETURN_ERROR (
@@ -576,20 +692,48 @@ bool MuseGlimmerCudaDFlashExecutor::run_target(
576692 return false ;
577693 }
578694 predictions.resize (step.logit_indices .size ());
695+ if (!drain_context &&
696+ !verification_hidden_.begin (step.tokens .size (), row_bytes_)) {
697+ return false ;
698+ }
699+ if (drain_context) {
700+ ++stats_.target_preamble_calls ;
701+ } else {
702+ ++stats_.target_verify_calls ;
703+ }
579704 std::size_t predicted = 0 ;
580705 const bool ok =
581706 target_->execute_raw (step, [&](const cb::RawTargetOutput& output) {
582- if (!output.taps ||
707+ if (drain_context) {
708+ ++stats_.target_preamble_forwards ;
709+ } else {
710+ ++stats_.target_verify_forwards ;
711+ }
712+ if (!output.taps || output.taps ->dim () != 3 ||
583713 output.taps ->scalar_type () != aten::ScalarType::BFloat16 ||
584714 output.taps ->size (2 ) != hidden_width_) {
585715 return false ;
586716 }
587- std::vector<std::uint8_t > copied (
588- static_cast <std::size_t >(output.length ) * row_bytes_);
589- // Copy every real tap before invoking the sampler; no borrowed GPU
590- // output survives any subsequent Module::execute call.
591- if (!copy_host (copied.data (), *output.taps , copied.size ())) {
592- return false ;
717+ std::vector<std::uint8_t > copied;
718+ // The next sampler or target slice may invalidate these borrowed taps.
719+ if (drain_context) {
720+ copied.resize (static_cast <std::size_t >(output.length ) * row_bytes_);
721+ if (!copy_host (copied.data (), *output.taps , copied.size ())) {
722+ return false ;
723+ }
724+ if (output.taps ->device_type () == aten::DeviceType::CUDA ) {
725+ stats_.hidden_rows_d2h += output.length ;
726+ stats_.hidden_bytes_d2h += copied.size ();
727+ }
728+ } else {
729+ if (!verification_hidden_.capture (
730+ *output.taps , output.offset , output.length )) {
731+ return false ;
732+ }
733+ if (verification_hidden_.on_device ()) {
734+ stats_.hidden_rows_d2d += output.length ;
735+ stats_.hidden_bytes_d2d += output.length * row_bytes_;
736+ }
593737 }
594738 auto tokens =
595739 target_->argmax_rows (output.logits , output.selected_indices .size ());
@@ -600,6 +744,9 @@ bool MuseGlimmerCudaDFlashExecutor::run_target(
600744 predictions[output.selected_indices [row]] = (*tokens)[row];
601745 ++predicted;
602746 }
747+ if (!drain_context) {
748+ return true ;
749+ }
603750 std::unordered_set<SessionId> touched;
604751 for (int row = 0 ; row < output.length ; ++row) {
605752 const int index = output.offset + row;
@@ -800,29 +947,52 @@ bool MuseGlimmerCudaDFlashExecutor::speculate(
800947 if (!run_target (verification, token_owners, predictions, false )) {
801948 return false ;
802949 }
950+ std::vector<df::GreedyBurst> bursts;
803951 for (std::size_t i = 0 ; i < owners.size (); ++i) {
804952 auto burst = df::verify_greedy (
805953 proposals[i],
806954 std::vector<Token>(
807955 predictions.begin () + i * width,
808956 predictions.begin () + (i + 1 ) * width));
809957 if (!burst ||
810- !rewind (
811- sessions_.at (owners[i]),
812- starts[i] + static_cast <int >(burst->tokens .size ()))) {
958+ !verification_hidden_.copy_rows (
959+ sessions_.at (owners[i]).hidden , i * width, burst->accepted + 1 )) {
960+ // Earlier owners' D2H copies are still queued into their hidden rows.
961+ (void )verification_hidden_.fence ();
813962 return false ;
814963 }
964+ if (verification_hidden_.on_device ()) {
965+ stats_.hidden_rows_d2h += burst->accepted + 1 ;
966+ stats_.hidden_bytes_d2h += (burst->accepted + 1 ) * row_bytes_;
967+ }
968+ bursts.push_back (std::move (*burst));
969+ }
970+ // All host backing remains stable until this coalesced fence, even when
971+ // verification spanned several target callbacks/static padding shapes.
972+ if (!verification_hidden_.fence ()) {
973+ return false ;
974+ }
975+ for (std::size_t i = 0 ; i < owners.size (); ++i) {
976+ auto & state = sessions_.at (owners[i]);
977+ const auto & burst = bursts[i];
978+ const int committed = starts[i] + static_cast <int >(burst.accepted ) + 1 ;
979+ // Rejected KV was written speculatively, but rejected hidden was never
980+ // appended. Do not use the ordinary rewind/hidden trimming path here.
981+ if (!target_->cache_control ().rewind (state.target_id , committed)) {
982+ return false ;
983+ }
984+ state.frontier = committed;
815985 // The regular preamble's prediction is new output, now committed by
816986 // verification. Append only predictions after that anchor. The last
817987 // prediction remains uncommitted: frontier = input.end + out.size() - 1.
818988 auto & output_tokens = out.outputs [indices[i]]->tokens ;
819989 output_tokens.insert (
820- output_tokens.end (), burst-> tokens .begin (), burst-> tokens .end ());
990+ output_tokens.end (), burst. tokens .begin (), burst. tokens .end ());
821991 ++stats_.cycles ;
822992 stats_.proposed += proposals_;
823- stats_.accepted += burst-> accepted ;
824- stats_.bonuses += burst-> bonus ;
825- stats_.rejected += !burst-> bonus ;
993+ stats_.accepted += burst. accepted ;
994+ stats_.bonuses += burst. bonus ;
995+ stats_.rejected += !burst. bonus ;
826996 }
827997 }
828998 return true ;
@@ -832,11 +1002,16 @@ bool MuseGlimmerCudaDFlashExecutor::fail_batch(
8321002 const batching::BatchInput& batch,
8331003 batching::BatchOutput& out) {
8341004 (void )batch;
1005+ // A failed callback/copy may still have queued work touching host backing;
1006+ // if that cannot be fenced, keep the backing alive rather than release it.
1007+ const bool fenced = verification_hidden_.fence ();
8351008 if (!failed_) {
8361009 failed_ = true ;
8371010 for (auto & item : sessions_) {
8381011 item.second .poisoned = true ;
839- item.second .hidden .clear ();
1012+ if (fenced) {
1013+ item.second .hidden .clear ();
1014+ }
8401015 }
8411016 // An abandoned declaration/forward affects the entire shared draft pool.
8421017 // Never allocate another draft SID or remove stale IDs after this clear.
0 commit comments