mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[bugfix] fix bfloat16 for get_tensor (#1216) 

* [bugfix] fix bfloat16 for get_tensor

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>

* update

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>

* update

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>

---------

Signed-off-by: Xuchun Shang <xuchun.shang@gmail.com>
This commit is contained in:
Xuchun Shang 2025-12-16 15:26:16 +08:00 committed by GitHub
parent 80431fb489
commit 284a3edfec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mooncake-integration/store/store_py.cpp
View File

@ -378,6 +378,11 @@ class MooncakeStorePyWrapper {
}
pybind11::object tensor =
torch_module().attr("from_numpy")(np_array);
if (dtype_enum == TensorDtype::BFLOAT16) {
tensor = tensor.attr("view")(torch_module().attr("bfloat16"));
} else if (dtype_enum == TensorDtype::FLOAT16) {
tensor = tensor.attr("view")(torch_module().attr("float16"));
}
return tensor;
} catch (const pybind11::error_already_set &e) {
@ -494,6 +499,13 @@ class MooncakeStorePyWrapper {
np_array = np_array.attr("reshape")(shape_tuple);
}
pybind11::object tensor = torch.attr("from_numpy")(np_array);
if (dtype_enum == TensorDtype::BFLOAT16) {
tensor =
tensor.attr("view")(torch_module().attr("bfloat16"));
} else if (dtype_enum == TensorDtype::FLOAT16) {
tensor =
tensor.attr("view")(torch_module().attr("float16"));
}
results_list.append(tensor);
}
} catch (const pybind11::error_already_set &e) {