Skip to content

Commit

Permalink
Avoid cuda-dependent code for CPU-only inference (#499)
Browse files Browse the repository at this point in the history
* Avoid cuda-dependent code for CPU-only inference

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
dur-randir and pre-commit-ci[bot] committed Aug 25, 2024
1 parent e029f38 commit 9e2f5e6
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions tools/llama/generate.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import queue
import threading
import time
from contextlib import nullcontext
from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Optional, Tuple, Union
Expand Down Expand Up @@ -181,8 +182,12 @@ def decode_n_tokens(
else:
window = previous_tokens[:, i - win_size : i]

with torch.backends.cuda.sdp_kernel(
enable_flash=False, enable_mem_efficient=False, enable_math=True
with (
torch.backends.cuda.sdp_kernel(
enable_flash=False, enable_mem_efficient=False, enable_math=True
)
if torch.cuda.is_available()
else nullcontext()
): # Actually better for Inductor to codegen attention here
next_token = decode_one_token(
model=model,
Expand Down

0 comments on commit 9e2f5e6

Please sign in to comment.