Skip to content

Commit f4d695c

Browse files
[executorch][muse] Copy only committed DFlash hidden rows to host (#24040)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #23982 by @Gasoonjia ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/81/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/81/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/80/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/81/orig Differential Revision: [D125705714](https://our.internmc.facebook.com/intern/diff/D125705714/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>
1 parent 8d984e4 commit f4d695c

3 files changed

Lines changed: 333 additions & 17 deletions

File tree

‎examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp‎

Lines changed: 190 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -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 {
159275
Result<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.

‎examples/models/muse-glimmer/runtime/cuda_dflash_executor.h‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,36 @@ bool trim_hidden(
4040
int draft_frontier,
4141
int position);
4242

43+
// Reusable owned storage for borrowed target taps. Capture precedes every
44+
// sampler/next slice; host destinations stay alive until fence().
45+
class HiddenSnapshot {
46+
public:
47+
~HiddenSnapshot();
48+
HiddenSnapshot() = default;
49+
HiddenSnapshot(const HiddenSnapshot&) = delete;
50+
HiddenSnapshot& operator=(const HiddenSnapshot&) = delete;
51+
HiddenSnapshot(HiddenSnapshot&&) = delete;
52+
HiddenSnapshot& operator=(HiddenSnapshot&&) = delete;
53+
bool begin(std::size_t rows, std::size_t row_bytes);
54+
bool capture(const aten::Tensor& taps, std::size_t offset, std::size_t rows);
55+
bool copy_rows(
56+
std::vector<std::uint8_t>& destination,
57+
std::size_t offset,
58+
std::size_t rows);
59+
bool fence();
60+
bool on_device() const {
61+
return device_source_.value_or(false);
62+
}
63+
64+
private:
65+
void* device_ = nullptr;
66+
void* stream_ = nullptr;
67+
std::size_t capacity_ = 0, rows_ = 0, row_bytes_ = 0, captured_ = 0;
68+
std::vector<std::uint8_t> host_;
69+
std::optional<bool> device_source_;
70+
bool stream_set_ = false;
71+
};
72+
4373
struct DraftContext {
4474
std::int32_t id;
4575
int position;
@@ -66,6 +96,14 @@ struct CudaDFlashStats {
6696
std::uint64_t executions = 0;
6797
std::uint64_t draft_forwards = 0;
6898
std::uint64_t warmed_context_tokens = 0;
99+
// Calls are packed execute_raw attempts; forwards count slice callbacks.
100+
std::uint64_t target_preamble_calls = 0;
101+
std::uint64_t target_verify_calls = 0;
102+
std::uint64_t target_preamble_forwards = 0;
103+
std::uint64_t target_verify_forwards = 0;
104+
// Actual CUDA transfers, excluding static padding and CPU-resident taps.
105+
std::uint64_t hidden_rows_d2h = 0, hidden_bytes_d2h = 0;
106+
std::uint64_t hidden_rows_d2d = 0, hidden_bytes_d2d = 0;
69107
std::uint64_t cycles = 0; // per-owner verification cycles
70108
std::uint64_t proposed = 0;
71109
std::uint64_t accepted = 0;
@@ -115,8 +153,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
115153
std::int32_t target_id;
116154
std::int32_t draft_id;
117155
int frontier = 0;
118-
// Host-owned BF16 rows [draft frontier, target frontier). Copied/fenced
119-
// before a sampler or subsequent forward can overwrite borrowed GPU taps.
156+
// Host-owned committed BF16 rows [draft frontier, target frontier).
157+
// Verification appends only the consumed anchor and accepted proposals.
120158
std::vector<std::uint8_t> hidden;
121159
std::optional<batching::SamplingParams> sampling;
122160
std::optional<std::uint64_t> seed;
@@ -163,6 +201,7 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
163201
batching::SessionId next_session_ = 1;
164202
std::unordered_map<batching::SessionId, State> sessions_;
165203
CudaDFlashStats stats_;
204+
cuda_dflash_detail::HiddenSnapshot verification_hidden_;
166205
bool initialized_ = false;
167206
// A failed co-batch invalidates every draft owner. Only model restart
168207
// recovers.

0 commit comments

Comments
 (0)