Initialize KVCache page_tables to zero to prevent NaNs being introd…

…uced
nod-ai · Jan 16, 2025 · f0fb408 · f0fb408
1 parent e34ffec
commit f0fb408
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/shortfin/python/shortfin_apps/llm/components/kvcache/page_pool.py b/shortfin/python/shortfin_apps/llm/components/kvcache/page_pool.py
@@ -104,6 +104,10 @@ def __init__(self, *, devices: Sequence[sf.ScopedDevice], config: PagePoolConfig
             page_table = sf.array.device_array.for_device(
                 device, page_table_shape, self.config.dtype
             )
+            page_table_host = page_table.for_transfer()
+            with page_table_host.map(discard=True) as m:
+                m.fill(0)
+            page_table_host.copy_to(page_table)
             self.page_tables.append(page_table)
 
     def acquire_free_pages(self, count: int) -> list[PageInfo] | None: