From 69aa46b2fa6f099e92c345eae1c0c9a42b5d66b4 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Fri, 11 Sep 2026 15:12:24 +0800 Subject: [PATCH] [ROCm] Raise HiCache JIT block quota for mapped-host throughput (#39036) --- python/sglang/kernels/ops/kvcache/hicache.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/kvcache/hicache.py b/python/sglang/kernels/ops/kvcache/hicache.py index 66c5c0feb..6844b34e0 100644 --- a/python/sglang/kernels/ops/kvcache/hicache.py +++ b/python/sglang/kernels/ops/kvcache/hicache.py @@ -3,14 +3,22 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING -from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args +from sglang.kernels.jit.utils import ( + cache_once, + is_hip_runtime, + load_jit, + make_cpp_args, +) from sglang.kernels.kernel_api_logging import debug_kernel_api if TYPE_CHECKING: import torch from tvm_ffi.module import Module -DEFAULT_BLOCK_QUOTA = 2 +_is_hip = is_hip_runtime() + +# ROCm needs a wider grid to saturate mapped-host transfers; CUDA keeps the legacy quota. +DEFAULT_BLOCK_QUOTA = 32 if _is_hip else 2 @cache_once