From 4a7d270b58ce3e80761748821ab5fb2775c61a60 Mon Sep 17 00:00:00 2001 From: Rasmus Larsen Date: Sun, 26 Jul 2026 18:54:36 +0300 Subject: [PATCH] Add 2048-token pages to FMHA batch prefill Co-authored-by: OpenAI Codex Signed-off-by: Rasmus Larsen --- example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py index b687cb9761..2340bf8c90 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py @@ -56,7 +56,7 @@ K0_MAX_SUBMAX_MAP = {32: 32, 64: 64, 96: 128, 128: 128, 256: 256} -SUPPORTED_PAGE_SIZE = [1, 16, 1024] +SUPPORTED_PAGE_SIZE = [1, 16, 1024, 2048] SUPPORTED_KV_MEMORY_LAYOUT = ["vectorized", "linear"] SUPPORTED_KV_LOOKUP_TABLE = ["vllm", "sglang"] KV_MEMORY_LAYOUT_ENUM_MAP = { @@ -858,7 +858,7 @@ def get_fwd_blobs( ): continue - # Generate kernels for both page_size=16 and page_size=1024 + # Generate kernels for every supported paged-KV cache size. for page_size in SUPPORTED_PAGE_SIZE: if page_size == 1 and pipeline.F_kv_memory_layout != "linear": continue