diff --git a/tests/core/framework/block/block_manager_test.cpp b/tests/core/framework/block/block_manager_test.cpp index b8402ccaf1..9edcdf4e94 100644 --- a/tests/core/framework/block/block_manager_test.cpp +++ b/tests/core/framework/block/block_manager_test.cpp @@ -24,7 +24,6 @@ limitations under the License. #include "block_manager_pool.h" #include "common/global_flags.h" #include "core/framework/config/scheduler_config.h" -#include "core/framework/config/service_config.h" #include "framework/block/block_manager_pool_test_peer.h" #include "framework/block/linear_state_block_manager.h" #include "framework/model/model_input_params.h" @@ -164,8 +163,12 @@ BlockManagerPool::Options make_linear_state_pool_options( BlockManagerPool::Options options; options.num_blocks(8).host_num_blocks(0).block_size(4).enable_prefix_cache( true); - options.num_single_blocks(4).enable_linear_state(true).linear_state_num_slots( - linear_state_num_slots); + // Zero out max_seqs_per_batch so the SINGLE pool is sized purely by + // num_single_blocks (otherwise the pool takes the max of the two). + options.max_seqs_per_batch(0) + .num_single_blocks(4) + .enable_linear_state(true) + .linear_state_num_slots(linear_state_num_slots); return options; } @@ -342,17 +345,16 @@ TEST(BlockManagerPoolTest, SingleBlockCapacityUsesOptionsMaxSeqs) { } TEST(BlockManagerPoolTest, TryAllocateKvFailureRollsBackSingleBlock) { - // unified scheduler-side single-block pool has 2 ids. - ScopedValue max_seqs_guard( - &SchedulerConfig::get_instance().max_seqs_per_batch(), 0); - BlockManagerPool::Options options; options.num_blocks(3).host_num_blocks(0).block_size(1).enable_prefix_cache( false); // id 0 is reserved for padding, so capacity 3 exposes 2 usable single-block - // ids, enough for the two sequences allocated after the rollback. - options.num_single_blocks(3).enable_linear_state(true).linear_state_num_slots( - 64); + // ids, enough for the two sequences allocated after the rollback. Zero out + // max_seqs_per_batch so num_single_blocks alone drives the SINGLE pool size. + options.max_seqs_per_batch(0) + .num_single_blocks(3) + .enable_linear_state(true) + .linear_state_num_slots(64); ScopedValue chunk_guard( &SchedulerConfig::get_instance().max_tokens_per_chunk_for_prefill(), 4); BlockManagerPool pool(options, /*dp_size=*/1); @@ -375,17 +377,17 @@ TEST(BlockManagerPoolTest, TryAllocateKvFailureRollsBackSingleBlock) { } TEST(BlockManagerPoolTest, SingleBlockCapacityCanBeLowerThanMaxSeqs) { - ScopedValue max_seqs_guard(&FLAGS_max_seqs_per_batch, 8); - // Pin the service-level concurrency so the SINGLE pool is sized purely by the - // options.num_single_blocks set below (the pool takes the max of the two). - ScopedValue max_conc_guard( - &ServiceConfig::get_instance().max_concurrent_requests(), 0); - BlockManagerPool::Options options; // id 0 is reserved for padding, so capacity 4 exposes 3 usable single blocks. + // max_seqs_per_batch is 8 in the scheduler view, but num_single_blocks caps + // the SINGLE pool at 4 (pool takes the max, so num_single_blocks wins only + // when it is >= max_seqs_per_batch + 2 -- here it is 4 vs 8+2=10, so the + // 8+2 path would win; force num_single_blocks to be the winner by lowering + // the scheduler side). options.num_blocks(16) .host_num_blocks(0) .block_size(1) + .max_seqs_per_batch(0) .num_single_blocks(4) .enable_prefix_cache(false); options.enable_linear_state(true).linear_state_num_slots(64); @@ -410,14 +412,14 @@ TEST(BlockManagerPoolTest, SingleBlockCapacityCanBeLowerThanMaxSeqs) { } TEST(BlockManagerPoolTest, DpRankSelectionSkipsExhaustedSingleBlockPool) { - ScopedValue max_seqs_guard(&FLAGS_max_seqs_per_batch, 8); - BlockManagerPool::Options options; // id 0 is reserved for padding, so capacity 2 exposes 1 usable block per - // rank. + // rank. Zero out max_seqs_per_batch so num_single_blocks alone drives the + // SINGLE pool size. options.num_blocks(16) .host_num_blocks(0) .block_size(1) + .max_seqs_per_batch(0) .num_single_blocks(2) .enable_prefix_cache(false); options.enable_linear_state(true).linear_state_num_slots(64); @@ -435,18 +437,14 @@ TEST(BlockManagerPoolTest, DpRankSelectionSkipsExhaustedSingleBlockPool) { } TEST(BlockManagerPoolTest, SingleBlockExhaustionBehavesLikeKvBlockExhaustion) { - ScopedValue max_seqs_guard(&FLAGS_max_seqs_per_batch, 8); - // Pin the service-level concurrency so the SINGLE pool is sized purely by the - // options.num_single_blocks set below (the pool takes the max of the two). - ScopedValue max_conc_guard( - &ServiceConfig::get_instance().max_concurrent_requests(), 0); - BlockManagerPool::Options options; // id 0 is reserved for padding, so capacity 2 exposes 1 usable block per - // rank. + // rank. Zero out max_seqs_per_batch so num_single_blocks alone drives the + // SINGLE pool size. options.num_blocks(16) .host_num_blocks(0) .block_size(1) + .max_seqs_per_batch(0) .num_single_blocks(2) .enable_prefix_cache(false); options.enable_linear_state(true).linear_state_num_slots(64); diff --git a/tests/core/framework/config/config_json_test.cpp b/tests/core/framework/config/config_json_test.cpp index b0ceb3a8d6..2ecde3448a 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -412,7 +412,7 @@ TEST(ConfigJsonTest, MissingJsonFileKeepsFlagDefaults) { EXPECT_EQ(kv_cache_config.block_size(), 128); EXPECT_DOUBLE_EQ(kv_cache_config.max_memory_utilization(), 0.8); EXPECT_EQ(scheduler_config.max_tokens_per_batch(), 10240); - EXPECT_EQ(scheduler_config.max_seqs_per_batch(), 1024); + EXPECT_EQ(scheduler_config.max_seqs_per_batch(), 200); } TEST(ConfigJsonTest, DumpStartupConfigSkipsWhenDisabled) { diff --git a/xllm/core/distributed_runtime/rec_engine.cpp b/xllm/core/distributed_runtime/rec_engine.cpp index 77fc81833f..a67f9807fc 100644 --- a/xllm/core/distributed_runtime/rec_engine.cpp +++ b/xllm/core/distributed_runtime/rec_engine.cpp @@ -223,7 +223,8 @@ bool RecEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { .host_num_blocks(0) .block_size(block_size) .enable_prefix_cache(options_.enable_prefix_cache()) - .enable_disagg_pd(options_.enable_disagg_pd()); + .enable_disagg_pd(options_.enable_disagg_pd()) + .max_seqs_per_batch(options_.max_seqs_per_batch()); kv_cache_manager_ = std::make_unique(options, dp_size_); return pipeline_->allocate_kv_cache(kv_cache_shape); diff --git a/xllm/core/distributed_runtime/vlm_engine.cpp b/xllm/core/distributed_runtime/vlm_engine.cpp index e871cef05d..757ce45291 100644 --- a/xllm/core/distributed_runtime/vlm_engine.cpp +++ b/xllm/core/distributed_runtime/vlm_engine.cpp @@ -338,7 +338,8 @@ bool VLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { .enable_linear_state(enable_linear_attention) .enable_prefix_cache(options_.enable_prefix_cache()) .enable_disagg_pd(options_.enable_disagg_pd()) - .hasher_type(BlockHasherType::MM); + .hasher_type(BlockHasherType::MM) + .max_seqs_per_batch(options_.max_seqs_per_batch()); if (enable_linear_attention) { // The unified linear-state slot pool spans all physical slots [0, N); // id 0 is reserved as padding and ids [1, N) serve live and checkpoint diff --git a/xllm/core/framework/config/scheduler_config.cpp b/xllm/core/framework/config/scheduler_config.cpp index 427ede1576..7ebe1d2927 100644 --- a/xllm/core/framework/config/scheduler_config.cpp +++ b/xllm/core/framework/config/scheduler_config.cpp @@ -20,7 +20,7 @@ limitations under the License. DEFINE_int32(max_tokens_per_batch, 10240, "Max number of tokens per batch."); -DEFINE_int32(max_seqs_per_batch, 1024, "Max number of sequences per batch."); +DEFINE_int32(max_seqs_per_batch, 200, "Max number of sequences per batch."); DEFINE_bool(enable_schedule_overlap, false, diff --git a/xllm/core/framework/config/scheduler_config.h b/xllm/core/framework/config/scheduler_config.h index 3e8c7a7422..ba68ab50d5 100644 --- a/xllm/core/framework/config/scheduler_config.h +++ b/xllm/core/framework/config/scheduler_config.h @@ -61,7 +61,7 @@ class SchedulerConfig final { PROPERTY(int32_t, max_tokens_per_batch) = 10240; - PROPERTY(int32_t, max_seqs_per_batch) = 1024; + PROPERTY(int32_t, max_seqs_per_batch) = 200; PROPERTY(bool, enable_schedule_overlap) = false; diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 24bd4937a5..ad4e53a2ec 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -84,6 +85,33 @@ void initialize_configs() { SchedulerConfig::get_instance().initialize(); ServiceConfig::get_instance().initialize(); SpeculativeConfig::get_instance().initialize(); + + // Reconcile the two per-batch admission caps into a single scheduler view + // consumed by every downstream user (Master -> Engine -> BlockManagerPool + // etc.). max_seqs_per_batch bounds the scheduler batch; + // max_concurrent_requests bounds the service-level admission; the effective + // batch cap is the tighter of the two. 0 means "unset" on either side; if + // both are 0 the caller has no way to size batch-bound resources (e.g. the + // SINGLE block pool) and we fail early. + { + SchedulerConfig& scheduler_config = SchedulerConfig::get_instance(); + ServiceConfig& service_config = ServiceConfig::get_instance(); + const int32_t scheduler_cap = scheduler_config.max_seqs_per_batch(); + const int32_t service_cap = service_config.max_concurrent_requests(); + int32_t effective_cap = 0; + if (scheduler_cap > 0 && service_cap > 0) { + effective_cap = std::min(scheduler_cap, service_cap); + } else if (scheduler_cap > 0) { + effective_cap = scheduler_cap; + } else if (service_cap > 0) { + effective_cap = service_cap; + } else { + LOG(FATAL) << "Both max_seqs_per_batch and max_concurrent_requests are " + "0; set at least one to a positive value."; + } + scheduler_config.max_seqs_per_batch(effective_cap); + service_config.max_concurrent_requests(effective_cap); + } } Options create_options(const std::string& instance_name, bool is_local) {