From 9b8c41171aa7582c0cbd9fe550109d9c87d77fd5 Mon Sep 17 00:00:00 2001 From: YC Yen-Ching Tseng Date: Wed, 17 Jun 2026 18:01:43 +0800 Subject: [PATCH] [AMD] Fall back to layer_first layout for kernel write-back on ROCm (#28473) --- python/sglang/srt/server_args.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index cfe1237f2..4a0eb4421 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4025,6 +4025,20 @@ class ServerArgs: "Page first layout is not supported with direct IO backend, switching to page first direct layout" ) + # The page_first kernel write-back relies on the CUDA-only JIT staged + # kernel. On ROCm it falls back to a kernel that requires CUDA index + # tensors and crashes on host write-back, so use layer_first there. + if ( + self.hicache_mem_layout == "page_first" + and self.hicache_io_backend == "kernel" + and is_hip() + ): + self.hicache_mem_layout = "layer_first" + logger.warning( + "page_first kernel write-back requires the CUDA JIT kernel; " + "falling back to layer_first layout on ROCm." + ) + def _resolve_storage_layout_compatibility(self): if ( self.hicache_storage_backend != "mooncake"