This commit is contained in:
simon-mo 2024-03-15 22:23:38 -07:00
parent cf6ff18246
commit ad50bf4b25
2 changed files with 3 additions and 3 deletions

View File

@ -31,4 +31,4 @@ jobs:
ruff vllm tests ruff vllm tests
- name: Spelling check with codespell - name: Spelling check with codespell
run: | run: |
codespell --toml pyproject.toml codespell --toml pyproject.toml

View File

@ -36,8 +36,8 @@ def test_contexted_kv_attention(
torch.cuda.manual_seed(0) torch.cuda.manual_seed(0)
torch.set_default_device(device) torch.set_default_device(device)
# Need this, otherwise when we capture the graph the process for GPU 1 would run on both # Need this, otherwise when we capture the graph the process for GPU 1 would
# GPU0 and GPU1 and things would hang # run on both GPU0 and GPU1 and things would hang
# #
# see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523 # see also similar issue: https://github.com/Dao-AILab/flash-attention/issues/523
torch.cuda.set_device(device) torch.cuda.set_device(device)