Update rocm.py

aa6b01c6 · xuxzh1 · c5eae25b · aa6b01c6
Commit aa6b01c6 authored Feb 26, 2025 by xuxzh1 🎱
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

server/text_generation_server/layers/attention/rocm.py server/text_generation_server/layers/attention/rocm.py +2 -0

No files found.
--- a/server/text_generation_server/layers/attention/rocm.py
+++ b/server/text_generation_server/layers/attention/rocm.py
@@ -114,6 +114,7 @@ def paged_attention(
            None,
            "auto",
            1.0,
+            1.0
        )
    else:
        # Run PagedAttention V2.
@@ -149,6 +150,7 @@ def paged_attention(
            None,
            "auto",
            1.0,
+            1.0
        )
        # else:
        #     paged_attention_custom(