vrr/kvcache-ai-ktransformers
2
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2025-09-07 04:59:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

raise exception on device error (#1342)
Some checks failed
Book-CI / test (push) Has been cancelled
Details
Deploy / deploy (macos-latest) (push) Has been cancelled
Details
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Details
Deploy / deploy (windows-latest) (push) Has been cancelled
Details

Browse source 
* display the unavailable torch device on error

* Raise exception on device error

---------

Signed-off-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>
This commit is contained in:
Emmanuel Ferdman 2025-05-29 10:37:41 +03:00 • committed by GitHub
parent ce75fcd7dd
commit d8bc6402b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
ktransformers/util/utils.py
View file

@ -243,7 +243,7 @@ def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cud
elif torch.xpu.is_available(): elif torch.xpu.is_available():
torch.xpu.set_device(torch_device) torch.xpu.set_device(torch_device)
else: else:
RuntimeError("The device: {torch_device} is not available") raise RuntimeError(f"The device: {torch_device} is not available")
inputs_embeds = model.model.embed_tokens(cur_token.to("cpu")).to(torch_device) inputs_embeds = model.model.embed_tokens(cur_token.to("cpu")).to(torch_device)
# with torch.cuda.stream(custom_stream): # with torch.cuda.stream(custom_stream):
logits=model(inputs_embeds=inputs_embeds, logits=model(inputs_embeds=inputs_embeds,
@ -284,7 +284,7 @@ def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cud
elif torch.xpu.is_available(): elif torch.xpu.is_available():
torch.xpu.set_device(torch_device) torch.xpu.set_device(torch_device)
else: else:
RuntimeError("The device: {torch_device} is not available") raise RuntimeError(f"The device: {torch_device} is not available")
with torch.no_grad(): with torch.no_grad():
stream = TextStreamer(tokenizer) stream = TextStreamer(tokenizer)