From 07e9ec2b66056d7e73854ae78bbe3783fdd66604 Mon Sep 17 00:00:00 2001
From: Mohit Sharma <mohit21sharma.ms@gmail.com>
Date: Tue, 10 Dec 2024 10:54:52 +0000
Subject: [PATCH] update partition size

---
 server/text_generation_server/layers/attention/rocm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/text_generation_server/layers/attention/rocm.py b/server/text_generation_server/layers/attention/rocm.py
index d65054a1..a401b589 100644
--- a/server/text_generation_server/layers/attention/rocm.py
+++ b/server/text_generation_server/layers/attention/rocm.py
@@ -179,7 +179,7 @@ def paged_attention(
                 1.0,
                 1.0,
                 None,
-                512,
+                _PARTITION_SIZE,
             )
 
     return out