Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

temporarily disable zero on device memory during kv_cache alloc #1269

Merged
merged 4 commits into from
Feb 21, 2025
Merged

temporarily disable zero on device memory during kv_cache alloc #1269

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
15b251d
temporarily disable zero on device memory during kv_cache alloc
guschmue Feb 20, 2025
f08f30a
fix comment
guschmue Feb 20, 2025
ec1c6da
make the comment easier to understand
guschmue Feb 20, 2025
5da2f9e
Merge branch 'main' into gs/fix-zero
guschmue Feb 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/models/kv_cache.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,10 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state)
sb_kv_caches_.empty() ? OrtValue::CreateTensor(Allocator(), shape_, type_)
: sb_kv_caches_[i]->CreateTensorOnStaticBuffer(shape_, type_));
// Zero the memory so we don't leak any data from the previous run
ByteWrapTensor(Device(), *presents_.back()).Zero();
// WebGPU device has no Zero() implementation yet. Since this zeroing is optional we disable it for WebGPU for now
if (Device().GetType() != DeviceType::WEBGPU) {
ByteWrapTensor(Device(), *presents_.back()).Zero();
}
}
} catch (const Ort::Exception&) {
std::ostringstream oss;
Expand Down
Loading