From 4f080a2cb8894046edbde7075d66ce1fc71f737f Mon Sep 17 00:00:00 2001 From: surprisely <46776020+surprisely@users.noreply.github.com> Date: Sat, 11 Jul 2026 19:04:04 +0000 Subject: [PATCH 1/3] perf(graph): bound paged captures by batch capacity Propagate max_batch_size through the Python and C++ cache configuration so Paged CUDA Graph captures only scheduler-reachable batch shapes. --- csrc/cache/kv_cache.cpp | 15 ++++++++++++++- csrc/cache/kv_cache.hpp | 6 ++++++ csrc/engine/compiler/paged_compiler.cpp | 16 ++++++++++++++++ csrc/pybind11/cache/cache.hpp | 8 ++++++-- python/infinilm/cache/cache.py | 8 ++++++-- python/infinilm/llm/model_runner/model_runner.py | 4 +++- 6 files changed, 51 insertions(+), 6 deletions(-) diff --git a/csrc/cache/kv_cache.cpp b/csrc/cache/kv_cache.cpp index 88d7367e9..1d30b72de 100644 --- a/csrc/cache/kv_cache.cpp +++ b/csrc/cache/kv_cache.cpp @@ -77,11 +77,19 @@ infinicore::Tensor create_layer_kv_cache( // ========================== // PagedKVCacheConfig // ========================== +// Preserve the legacy Graph capture ceiling for two-argument C++ callers. PagedKVCacheConfig::PagedKVCacheConfig( size_t num_blocks, size_t block_size) + : PagedKVCacheConfig(num_blocks, block_size, 512) {} + +PagedKVCacheConfig::PagedKVCacheConfig( + size_t num_blocks, + size_t block_size, + size_t max_batch_size) : num_blocks_(num_blocks), - block_size_(block_size) { + block_size_(block_size), + max_batch_size_(max_batch_size) { } std::unique_ptr @@ -99,6 +107,11 @@ PagedKVCacheConfig::block_size() const { return block_size_; } +size_t +PagedKVCacheConfig::max_batch_size() const { + return max_batch_size_; +} + namespace PagedKVCache { // ========================== // PagedKVCache diff --git a/csrc/cache/kv_cache.hpp b/csrc/cache/kv_cache.hpp index 4d0a9a704..7ca94a586 100644 --- a/csrc/cache/kv_cache.hpp +++ b/csrc/cache/kv_cache.hpp @@ -40,14 +40,20 @@ class PagedKVCacheConfig final : public CacheConfig { PagedKVCacheConfig( size_t num_blocks, size_t block_size = 256); + PagedKVCacheConfig( + size_t num_blocks, + size_t block_size, + size_t max_batch_size); std::unique_ptr unique_copy() const override; size_t num_blocks() const; size_t block_size() const; + size_t max_batch_size() const; private: size_t num_blocks_; size_t block_size_; + size_t max_batch_size_; }; namespace PagedKVCache { diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index df3fd1cb4..d711c942d 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -38,6 +38,22 @@ PagedCompiler::PagedCompiler(const std::shared_ptr &model, RankBa for (size_t b = 256; b <= 512; b += 64) { decode_batch_sizes_.push_back(b); } + + const auto *config = dynamic_cast(model_->get_cache_config()); + if (config == nullptr) { + return; + } + const size_t max_batch_size = config->max_batch_size(); + if (max_batch_size == 0) { + throw std::invalid_argument("Paged Graph max_batch_size must be greater than zero"); + } + decode_batch_sizes_.erase( + std::remove_if(decode_batch_sizes_.begin(), decode_batch_sizes_.end(), + [max_batch_size](size_t b) { return b > max_batch_size; }), + decode_batch_sizes_.end()); + if (decode_batch_sizes_.empty() || decode_batch_sizes_.back() != max_batch_size) { + decode_batch_sizes_.push_back(max_batch_size); + } } void PagedCompiler::compile() { diff --git a/csrc/pybind11/cache/cache.hpp b/csrc/pybind11/cache/cache.hpp index 492f6c302..439c9da2c 100644 --- a/csrc/pybind11/cache/cache.hpp +++ b/csrc/pybind11/cache/cache.hpp @@ -34,15 +34,19 @@ inline void bind_cache(py::module &m) { infinilm::cache::CacheConfig, std::shared_ptr>(m, "PagedKVCacheConfig") .def( - py::init(), + py::init(), py::arg("num_blocks"), - py::arg("block_size") = 256) + py::arg("block_size") = 256, + py::arg("max_batch_size") = 512) .def( "num_blocks", &infinilm::cache::PagedKVCacheConfig::num_blocks) .def( "block_size", &infinilm::cache::PagedKVCacheConfig::block_size) + .def( + "max_batch_size", + &infinilm::cache::PagedKVCacheConfig::max_batch_size) .def("__repr__", [](const infinilm::cache::PagedKVCacheConfig &) { return ""; }); diff --git a/python/infinilm/cache/cache.py b/python/infinilm/cache/cache.py index 4a9bcc446..b8b8e41c4 100644 --- a/python/infinilm/cache/cache.py +++ b/python/infinilm/cache/cache.py @@ -18,5 +18,9 @@ def __init__( class PagedKVCacheConfig(CacheConfig, _infinilm.PagedKVCacheConfig): - def __init__(self, num_blocks: int, block_size: int = 256): - _infinilm.PagedKVCacheConfig.__init__(self, num_blocks, block_size) + def __init__( + self, num_blocks: int, block_size: int = 256, max_batch_size: int = 512 + ): + _infinilm.PagedKVCacheConfig.__init__( + self, num_blocks, block_size, max_batch_size + ) diff --git a/python/infinilm/llm/model_runner/model_runner.py b/python/infinilm/llm/model_runner/model_runner.py index bbf3ccd37..ef6d9e13e 100644 --- a/python/infinilm/llm/model_runner/model_runner.py +++ b/python/infinilm/llm/model_runner/model_runner.py @@ -59,7 +59,9 @@ def __init__(self, config: EngineConfig): ) elif config.cache_type == "paged": cache_config = PagedKVCacheConfig( - num_blocks=config.num_blocks, block_size=config.block_size + num_blocks=config.num_blocks, + block_size=config.block_size, + max_batch_size=config.max_batch_size, ) logger.info(f"Using Paged KV Cache with num_blocks={config.num_blocks}") else: From 0b595a31841953a5940e20673c3d67b4239a765a Mon Sep 17 00:00:00 2001 From: surprisely <46776020+surprisely@users.noreply.github.com> Date: Sat, 11 Jul 2026 19:04:05 +0000 Subject: [PATCH 2/3] perf(graph): skip empty non-quant runtime resets Skip recursive runtime-state reset only when the model is explicitly non-quantized. Quantized Marlin lock and workspace resets remain unchanged. --- csrc/models/infinilm_model.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index c2cc68762..a16951f85 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -94,6 +94,10 @@ void InfinilmModel::process_weights_after_loading() { } void InfinilmModel::reset_runtime_state() const { + // Only quantized kernels currently keep resettable runtime state. + if (model_config_ && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) { + return; + } for (const auto &[_, sub] : children()) { reset_runtime_state_recursive_(sub.get()); } From ecee461c1e471e960627ba256f36b674017f0ceb Mon Sep 17 00:00:00 2001 From: surprisely <46776020+surprisely@users.noreply.github.com> Date: Sun, 12 Jul 2026 17:56:48 +0800 Subject: [PATCH 3/3] fix(graph): preserve runtime resets for other models --- csrc/models/infinilm_model.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index a16951f85..2a77be12a 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -94,8 +94,10 @@ void InfinilmModel::process_weights_after_loading() { } void InfinilmModel::reset_runtime_state() const { - // Only quantized kernels currently keep resettable runtime state. - if (model_config_ && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) { + // Non-quantized Qwen3-MoE currently has no resettable module state. + if (model_config_ + && model_config_->get_or("model_type", "") == "qwen3_moe" + && model_config_->get_quant_scheme() == quantization::QuantScheme::NONE) { return; } for (const auto &[_, sub] : children()) {