diff --git a/tests/core/framework/batch/batch_test.cpp b/tests/core/framework/batch/batch_test.cpp index 99d52e4350..7472a3ec32 100644 --- a/tests/core/framework/batch/batch_test.cpp +++ b/tests/core/framework/batch/batch_test.cpp @@ -3089,6 +3089,7 @@ TEST(BatchTest, OverlapMTPReplacementKeepsCompositeKvBlocks) { .sliding_window_size(window_size) .swa_blocks_per_seq(static_cast( get_swa_blocks_per_seq(window_size, base_block_size))) + .swa_num_blocks(20) .max_tokens_per_batch(1280) .max_seqs_per_batch(max_seqs_per_batch) .manager_types({1, 0, 0}) diff --git a/tests/core/framework/block/composite_block_manager_test.cpp b/tests/core/framework/block/composite_block_manager_test.cpp index 8b43c2e441..cbbd071bab 100644 --- a/tests/core/framework/block/composite_block_manager_test.cpp +++ b/tests/core/framework/block/composite_block_manager_test.cpp @@ -47,11 +47,16 @@ BlockManager::Options MakeCompositeOptions(uint32_t base_num_blocks, uint32_t max_seqs_per_batch) { const uint32_t swa_blocks_per_seq = static_cast(get_swa_blocks_per_seq(window_size, block_size)); + const uint32_t burst_blocks = + (kMaxTokensPerBatch + block_size - 1) / block_size; + const uint32_t swa_num_blocks = swa_blocks_per_seq * max_seqs_per_batch + + burst_blocks + max_seqs_per_batch + 2; BlockManager::Options opts; opts.num_blocks(base_num_blocks) .block_size(block_size) .sliding_window_size(window_size) .swa_blocks_per_seq(swa_blocks_per_seq) + .swa_num_blocks(swa_num_blocks) .max_tokens_per_batch(kMaxTokensPerBatch) .max_seqs_per_batch(max_seqs_per_batch) .manager_types({kManagerTypeSlidingWindowBlockManager, @@ -61,6 +66,20 @@ BlockManager::Options MakeCompositeOptions(uint32_t base_num_blocks, return opts; } +void set_swa_capacity_for_token_budget(BlockManager::Options* options, + uint32_t max_tokens_per_batch) { + ASSERT_NE(options, nullptr); + const uint32_t block_size = options->block_size(); + ASSERT_GT(block_size, 0u); + const uint32_t burst_blocks = + (max_tokens_per_batch + block_size - 1) / block_size; + const uint32_t max_seqs = std::max(options->max_seqs_per_batch(), 1u); + const uint32_t swa_num_blocks = + options->swa_blocks_per_seq() * max_seqs + burst_blocks + max_seqs + 2; + options->max_tokens_per_batch(max_tokens_per_batch) + .swa_num_blocks(swa_num_blocks); +} + constexpr uint32_t kBaseBlockSize = 128; constexpr uint32_t kCompressRatio4 = 4; constexpr uint32_t kCompressRatio128 = 128; @@ -548,7 +567,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheHitOnRepeatedPrefix) { base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); // max_tokens_per_batch has to accommodate the prompt so allocate_sequence // does not exceed the SWA burst budget. - opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128); ASSERT_TRUE(opts.enable_prefix_cache()); CompositeBlockManager manager(build_composite_leaves(opts)); @@ -596,7 +615,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheMissCleanly) { const uint32_t max_seqs_per_batch = 4; BlockManager::Options opts = MakeCompositeOptions( base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); - opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128); CompositeBlockManager manager(build_composite_leaves(opts)); const size_t num_tokens = 2 * kBlockSizeRatio128; @@ -629,7 +648,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheEvictsAtC128Capacity) { const uint32_t window_size = 4 * kBaseBlockSize; BlockManager::Options opts = MakeCompositeOptions( base_num_blocks, kBaseBlockSize, window_size, /*max_seqs_per_batch=*/1); - opts.max_tokens_per_batch(2 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 2 * kBlockSizeRatio128); CompositeBlockManager manager(build_composite_leaves(opts)); const size_t num_tokens = 2 * kBlockSizeRatio128; @@ -663,7 +682,7 @@ TEST(CompositeBlockManagerTest, SlidingWindowSlidOutBlocksEnterPrefixCache) { const uint32_t max_seqs_per_batch = 4; BlockManager::Options opts = MakeCompositeOptions( base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); - opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128); ASSERT_TRUE(opts.enable_prefix_cache()); CompositeBlockManager manager(build_composite_leaves(opts)); @@ -691,6 +710,44 @@ TEST(CompositeBlockManagerTest, SlidingWindowSlidOutBlocksEnterPrefixCache) { manager.deallocate_for_sequence(&seq_hit); } +TEST(CompositeBlockManagerTest, + SlidingWindowReclaimsUncachedBlockBeforeFullCacheUnit) { + const uint32_t window_size = 2 * kBaseBlockSize; + BlockManager::Options opts = MakeCompositeOptions( + /*base_num_blocks=*/4096, + kBaseBlockSize, + window_size, + /*max_seqs_per_batch=*/1); + opts.enable_prefix_cache(true).swa_num_blocks( + /*three live blocks plus padding=*/4); + CompositeBlockManager manager(build_composite_leaves(opts)); + + const size_t first_chunk_tokens = 3 * kBaseBlockSize; + const size_t second_chunk_tokens = 4 * kBaseBlockSize; + Sequence seq = + MakeTestSequence(0, std::vector(second_chunk_tokens, 7)); + + ASSERT_TRUE(manager.allocate_sequence(&seq, first_chunk_tokens)); + BlockManager* swa_leaf = manager.leaf_entries().at(BlockType::SWA).leaf.get(); + ASSERT_NE(swa_leaf, nullptr); + EXPECT_EQ(swa_leaf->num_free_blocks(), 0u); + + seq.kv_state().incr_kv_cache_tokens_num(first_chunk_tokens); + + // The first block is outside the two-block window. No C128 cache unit is + // complete yet, so the next growth releases it directly and reuses its + // physical id without publishing a partial DSV4 prefix. + ASSERT_TRUE(manager.allocate_sequence(&seq, second_chunk_tokens)); + const std::vector swa_blocks = SwaBlocks(seq); + ASSERT_EQ(swa_blocks.size(), 4u); + EXPECT_FALSE(swa_blocks.front().is_valid()); + EXPECT_TRUE(swa_blocks.back().is_valid()); + EXPECT_EQ(swa_leaf->num_free_blocks(), 0u); + EXPECT_EQ(swa_leaf->num_blocks_in_prefix_cache(), 0u); + + manager.deallocate_for_sequence(&seq); +} + // The post-grow hook advances KVCacheState::num_cached_blocks incrementally. // Newly allocated blocks are present by then, but the token cursor limits the // published range to blocks completed by the preceding forward. @@ -700,7 +757,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCachePostGrowCursorAdvances) { const uint32_t max_seqs_per_batch = 4; BlockManager::Options opts = MakeCompositeOptions( base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); - opts.max_tokens_per_batch(4 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 4 * kBlockSizeRatio128); CompositeBlockManager manager(build_composite_leaves(opts)); // Two-chunk prompt (2*C128). Chunk 1 is a single C128 block wide. @@ -725,6 +782,55 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCachePostGrowCursorAdvances) { chunk / kBlockSizeRatio4); EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C128), chunk / kBlockSizeRatio128); + EXPECT_EQ(manager.leaf_entries() + .at(BlockType::SWA) + .leaf->num_blocks_in_prefix_cache(), + 4u); + EXPECT_EQ(manager.leaf_entries() + .at(BlockType::C4) + .leaf->num_blocks_in_prefix_cache(), + 32u); + EXPECT_EQ(manager.leaf_entries() + .at(BlockType::C128) + .leaf->num_blocks_in_prefix_cache(), + 1u); + + manager.deallocate_for_sequence(&seq); +} + +TEST(CompositeBlockManagerTest, Dsv4PrefixCacheSkipsPartialCacheUnitTail) { + const uint32_t base_num_blocks = 4096; + const uint32_t window_size = kBaseBlockSize; + const uint32_t max_seqs_per_batch = 4; + BlockManager::Options opts = MakeCompositeOptions( + base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); + set_swa_capacity_for_token_budget(&opts, 2 * kBlockSizeRatio128); + CompositeBlockManager manager(build_composite_leaves(opts)); + + const size_t completed_tokens = kBlockSizeRatio128 + kBlockSizeRatio4; + Sequence seq = + MakeTestSequence(0, std::vector(completed_tokens, 17)); + ASSERT_TRUE(manager.allocate_sequence(&seq, completed_tokens)); + seq.kv_state().incr_kv_cache_tokens_num(completed_tokens); + manager.cache_for_sequence(&seq); + + EXPECT_EQ(manager.leaf_entries() + .at(BlockType::SWA) + .leaf->num_blocks_in_prefix_cache(), + 1u); + EXPECT_EQ(manager.leaf_entries() + .at(BlockType::C4) + .leaf->num_blocks_in_prefix_cache(), + 32u); + EXPECT_EQ(manager.leaf_entries() + .at(BlockType::C128) + .leaf->num_blocks_in_prefix_cache(), + 1u); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::SWA), + kBlockSizeRatio128 / kBaseBlockSize); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C4), + kBlockSizeRatio128 / kBlockSizeRatio4); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C128), 1u); manager.deallocate_for_sequence(&seq); } @@ -739,7 +845,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheExactRepeatPopsOneC128) { const uint32_t max_seqs_per_batch = 4; BlockManager::Options opts = MakeCompositeOptions( base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); - opts.max_tokens_per_batch(4 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 4 * kBlockSizeRatio128); CompositeBlockManager manager(build_composite_leaves(opts)); const size_t num_tokens = 3 * kBlockSizeRatio128; @@ -770,7 +876,7 @@ TEST(CompositeBlockManagerTest, DecodeRoleSkipsSwaPrefixCache) { const uint32_t max_seqs_per_batch = 4; BlockManager::Options opts = MakeCompositeOptions( base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); - opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128); // First seed the cache under a PREFILL role so all three leaves publish // their blocks -- then swap in a DECODE-role composite that shares the // hash space via the same leaf construction path. @@ -818,6 +924,33 @@ TEST(CompositeBlockManagerTest, DecodeRoleSkipsSwaPrefixCache) { decode_manager.deallocate_for_sequence(&seq_d); } +TEST(CompositeBlockManagerTest, DecodeInitialSwaAllocationKeepsOnlyWindowTail) { + const uint32_t window_size = 2 * kBaseBlockSize; + BlockManager::Options opts = MakeCompositeOptions( + /*base_num_blocks=*/4096, + kBaseBlockSize, + window_size, + /*max_seqs_per_batch=*/1); + opts.instance_is_decode(true).enable_prefix_cache(false).swa_num_blocks( + /*two windows plus padding=*/5); + CompositeBlockManager manager(build_composite_leaves(opts)); + + const size_t logical_blocks = 10; + const size_t num_tokens = logical_blocks * kBaseBlockSize; + Sequence seq = MakeTestSequence(0, std::vector(num_tokens, 7)); + + ASSERT_TRUE(manager.allocate_sequence(&seq, num_tokens)); + const std::vector swa = SwaBlocks(seq); + ASSERT_EQ(swa.size(), logical_blocks); + for (size_t i = 0; i < logical_blocks - 2; ++i) { + EXPECT_FALSE(swa[i].is_valid()); + } + EXPECT_TRUE(swa[logical_blocks - 2].is_valid()); + EXPECT_TRUE(swa[logical_blocks - 1].is_valid()); + + manager.deallocate_for_sequence(&seq); +} + // Qwen3.5 GDN D-side (instance_is_decode=true, LINEAR present): the LINEAR // leaf should stop advertising prefix cache, so build_composite_leaves // classifies FLAT_KV_LINEAR down to FLAT_KV and no restore-source mount diff --git a/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp b/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp index 0db6d87d76..5e0a9487f6 100644 --- a/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp +++ b/tests/core/framework/block/hierarchy_block_manager_pool_test.cpp @@ -163,6 +163,7 @@ BlockManagerPool::Options make_typed_cache_options() { .enable_host_offload(true) .sliding_window_size(kWindow) .swa_blocks_per_seq(swa_blocks_per_seq) + .swa_num_blocks(266) .max_tokens_per_batch(32768) .max_seqs_per_batch(4) // SlidingWindow + BlockManagerImpl (C4) + BlockManagerImpl (C128). @@ -385,9 +386,10 @@ TEST(HierarchyBlockManagerPoolTest, sequence.kv_state().set_kv_cache_tokens_num(kPromptTokens); pool.deallocate(&sequence); - // Only completed blocks are offloaded: 156 SWA blocks, 39 C4 blocks, and - // one C128 checkpoint. The partial SWA tail is not inserted or offloaded. - EXPECT_EQ(HierarchyPoolTestPeer::pending_offload_pair_count(pool), 196u); + // Decode allocates only the active SWA window. Here that window is the + // partial tail block, so no SWA block is offloaded. Only the complete C128 + // cache unit (32 C4 blocks plus one C128 checkpoint) is offloaded. + EXPECT_EQ(HierarchyPoolTestPeer::pending_offload_pair_count(pool), 33u); } TEST(HierarchyBlockManagerPoolTest, AllocateSharedMountsMatchesWithoutH2d) { @@ -1307,7 +1309,7 @@ TEST(HierarchyBlockManagerPoolTest, pool.prefetch_from_storage(request); ASSERT_NE(engine.result(), nullptr); - EXPECT_EQ(engine.transfer_infos().size(), 128u + 32u + 1u); + EXPECT_EQ(engine.transfer_infos().size(), 1u + 32u + 1u); const std::vector hits(engine.transfer_infos().size(), 1); ASSERT_TRUE(engine.result()->set_batch_result( @@ -1321,7 +1323,13 @@ TEST(HierarchyBlockManagerPoolTest, ASSERT_TRUE(pool.update_prefetch_result(request, /*timeout=*/0)); pool.allocate_shared(sequence); - EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::SWA), 128u); + const Slice swa_blocks = + sequence->host_kv_state().blocks(BlockType::SWA); + ASSERT_EQ(swa_blocks.size(), 128u); + for (size_t i = 0; i + 1 < swa_blocks.size(); ++i) { + EXPECT_FALSE(swa_blocks[i].is_valid()); + } + EXPECT_TRUE(swa_blocks.back().is_valid()); EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C4), 32u); EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C128), 1u); EXPECT_EQ(sequence->kv_cache_tokens_num(), 16384u); @@ -1331,7 +1339,8 @@ TEST(HierarchyBlockManagerPoolTest, TEST(HierarchyBlockManagerPoolTest, TypedStoragePrefetchKeepsSwaHitsAfterMiddleMiss) { constexpr size_t kPromptTokens = 32769; - constexpr size_t kMissedSwaOrdinal = 100; + constexpr size_t kMissedSwaOrdinal = 0; + constexpr size_t kMissedSwaBlock = 127; BlockManagerPool::Options options = make_typed_cache_options(); options.enable_kvcache_store(true); FakePrefetchEngine engine(/*worker_count=*/2); @@ -1356,7 +1365,7 @@ TEST(HierarchyBlockManagerPoolTest, } ++swa_ordinal; } - ASSERT_EQ(swa_ordinal, 256u); + ASSERT_EQ(swa_ordinal, 2u); ASSERT_LT(missed_result_index, engine.transfer_infos().size()); const size_t split = engine.transfer_infos().size() / 2; @@ -1388,8 +1397,7 @@ TEST(HierarchyBlockManagerPoolTest, const Slice swa_blocks = sequence->host_kv_state().blocks(BlockType::SWA); ASSERT_EQ(swa_blocks.size(), 256u); - EXPECT_FALSE(swa_blocks[kMissedSwaOrdinal].is_valid()); - EXPECT_TRUE(swa_blocks[kMissedSwaOrdinal + 1].is_valid()); + EXPECT_FALSE(swa_blocks[kMissedSwaBlock].is_valid()); EXPECT_TRUE(swa_blocks.back().is_valid()); EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C4), 64u); EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C128), 2u); diff --git a/tests/core/framework/config/disagg_pd_config_test.cpp b/tests/core/framework/config/disagg_pd_config_test.cpp index abd8980ff1..0b4dc032c5 100644 --- a/tests/core/framework/config/disagg_pd_config_test.cpp +++ b/tests/core/framework/config/disagg_pd_config_test.cpp @@ -12,80 +12,3 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ - -#include "core/framework/config/disagg_pd_config.h" - -#include - -#include -#include - -#include "core/framework/config/kv_cache_config.h" -#include "core/framework/config/scheduler_config.h" - -namespace xllm { -namespace { - -struct PrefixRoleCase { - std::string role; - bool keep_prefix_cache; -}; - -void set_values_requiring_mlu_normalization(DisaggPDConfig& disagg_pd_config, - KVCacheConfig& kv_cache_config, - SchedulerConfig& scheduler_config) { - disagg_pd_config.kv_cache_transfer_mode("PULL").enable_pd_ooc(true); - kv_cache_config.kv_cache_dtype("fp8").enable_prefix_cache(true); - scheduler_config.enable_schedule_overlap(true); -} - -void expect_normalized_values(const DisaggPDConfig& disagg_pd_config, - const KVCacheConfig& kv_cache_config, - const SchedulerConfig& scheduler_config) { - EXPECT_EQ(disagg_pd_config.kv_cache_transfer_mode(), "PULL"); - EXPECT_FALSE(disagg_pd_config.enable_pd_ooc()); - EXPECT_EQ(kv_cache_config.kv_cache_dtype(), "auto"); - EXPECT_FALSE(scheduler_config.enable_schedule_overlap()); -} - -TEST(DisaggPDConfigTest, KeepsMluPrefixCacheForPrefillSideRoles) { - const PrefixRoleCase cases[] = { - {"PREFILL", true}, - {"MIX", true}, - {"DECODE", false}, - {"DEFAULT", false}, - }; - - for (const PrefixRoleCase& test_case : cases) { - DisaggPDConfig disagg_pd_config; - KVCacheConfig kv_cache_config; - SchedulerConfig scheduler_config; - disagg_pd_config.instance_role(test_case.role); - set_values_requiring_mlu_normalization( - disagg_pd_config, kv_cache_config, scheduler_config); - - disagg_pd_config.normalize_mlu(kv_cache_config, scheduler_config); - - SCOPED_TRACE(test_case.role); - expect_normalized_values( - disagg_pd_config, kv_cache_config, scheduler_config); - EXPECT_EQ(kv_cache_config.enable_prefix_cache(), - test_case.keep_prefix_cache); - } -} - -TEST(DisaggPDConfigTest, OmitsRemovedHeterogeneousPullOptions) { - const std::vector& option_names = - DisaggPDConfig::option_category().option_names; - EXPECT_EQ( - std::find( - option_names.begin(), option_names.end(), "enable_heterogeneous_pd"), - option_names.end()); - EXPECT_EQ(std::find(option_names.begin(), - option_names.end(), - "enable_pd_parallel_shard_pull"), - option_names.end()); -} - -} // namespace -} // namespace xllm diff --git a/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp b/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp index 99cec10aa2..a35fa0df7d 100644 --- a/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp +++ b/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include "framework/model/model_args.h" @@ -310,13 +311,14 @@ TEST(KVCacheEstimationTest, EstimatesDeepSeekV4Pools) { KVCacheEstimateOptions options; options.dtype = torch::kFloat32; options.kv_cache_dtype = "auto"; - options.cache_size_in_bytes = 2818048; + options.cache_size_in_bytes = + 2818048 + /*two_additional_swa_blocks=*/2 * 90112; options.block_size = 128; options.max_seqs_per_batch = 4; KVCacheCapacity capacity = estimate_kv_cache_capacity(model_args, options); - EXPECT_EQ(capacity.swa_count(), 19); + EXPECT_EQ(capacity.swa_count(), 21); #if defined(USE_MLU) EXPECT_EQ(capacity.c4_count(), 64); EXPECT_EQ(capacity.c128_count(), 2); @@ -328,6 +330,292 @@ TEST(KVCacheEstimationTest, EstimatesDeepSeekV4Pools) { #endif } +TEST(KVCacheEstimationTest, DeepSeekV4RejectsBudgetWithoutCompressedCacheUnit) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(128) + .compress_ratios({1, 4, 128}); + + constexpr int64_t kSwaCount = 4; + constexpr int64_t kSwaBytesPerBlock = + /*c1=*/128 * 16 * 4 + + /*c4=*/128 * (16 * 4 + 2 * 16 * 4 * 2 + 2 * 8 * 4 * 2) + + /*c128=*/128 * (16 * 4 + 16 * 4 * 2); + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = + kSwaCount * kSwaBytesPerBlock + /*remaining_bytes=*/1; + options.block_size = 128; + options.max_seqs_per_batch = 1; + options.max_tokens_per_chunk_for_prefill = 128; + + EXPECT_DEATH( + estimate_kv_cache_capacity(model_args, options), + "minimum DSV4 SWA cache leaves insufficient memory for one compressed " + "cache unit"); +} + +TEST(KVCacheEstimationTest, DeepSeekV4PdPrefillUsesChunkCapacity) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 16 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.max_tokens_per_chunk_for_prefill = 385; + options.enable_disagg_pd = true; + options.instance_role = InstanceRole::PREFILL; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + // W=3, ceil(385/128)=4: 4 * (3 + 4 + 1) usable rows + padding row. + EXPECT_EQ(capacity.swa_count(), 33); +} + +TEST(KVCacheEstimationTest, DeepSeekV4MixUsesChunkCapacity) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 16 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.max_tokens_per_batch = 16384; + options.max_tokens_per_chunk_for_prefill = 385; + options.instance_role = InstanceRole::MIX; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + EXPECT_EQ(capacity.swa_count(), 33); +} + +TEST(KVCacheEstimationTest, + DeepSeekV4PdPrefillWithoutChunkingUsesBatchCapacity) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 16 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.max_tokens_per_batch = 385; + options.max_tokens_per_chunk_for_prefill = 129; + options.enable_chunked_prefill = false; + options.enable_disagg_pd = true; + options.instance_role = InstanceRole::PREFILL; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + // W=3, ceil(385/128)=4: full-prefill mode uses the batch token bound. + EXPECT_EQ(capacity.swa_count(), 33); +} + +TEST(KVCacheEstimationTest, DeepSeekV4PdDecodeUsesTwoWindows) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 16 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.instance_role = InstanceRole::DECODE; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + // W=3: 4 * 3 * 2 usable rows + padding row. + EXPECT_EQ(capacity.swa_count(), 25); +} + +TEST(KVCacheEstimationTest, + DeepSeekV4PdDecodeAccountsForScheduleOverlapSpeculativeReserve) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(128) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 16 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 1; + options.num_speculative_tokens = 65; + options.enable_disagg_pd = true; + options.instance_role = InstanceRole::DECODE; + + options.enable_schedule_overlap = false; + const KVCacheCapacity non_overlap_capacity = + estimate_kv_cache_capacity(model_args, options); + // ceil((128 + 65) / 128)=2 rows per window, two windows plus padding. + EXPECT_EQ(non_overlap_capacity.swa_count(), 5); + + options.enable_schedule_overlap = true; + const KVCacheCapacity overlap_capacity = + estimate_kv_cache_capacity(model_args, options); + // ceil((128 + 2*65) / 128)=3 rows per window, two windows plus padding. + EXPECT_EQ(overlap_capacity.swa_count(), 7); +} + +TEST(KVCacheEstimationTest, DeepSeekV4PrefixCacheKeepsOperationalSwaPool) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 128 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.max_tokens_per_chunk_for_prefill = 385; + options.enable_disagg_pd = true; + options.instance_role = InstanceRole::PREFILL; + options.enable_prefix_cache = true; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + ASSERT_GT(capacity.c128_count(), 0); + EXPECT_EQ(capacity.c4_count(), 32 * capacity.c128_count()); + EXPECT_EQ(capacity.swa_count(), 33); +} + +TEST(KVCacheEstimationTest, + DeepSeekV4RealisticMixBudgetRetainsCompressedPools) { + std::vector compress_ratios{0, 0}; + compress_ratios.reserve(43); + for (int32_t layer_id = 2; layer_id < 43; ++layer_id) { + compress_ratios.emplace_back(layer_id % 2 == 0 ? 4 : 128); + } + + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(43) + .head_dim(512) + .index_head_dim(128) + .window_size(128) + .max_seq_len(1048576) + .compress_ratios(std::move(compress_ratios)); + + KVCacheEstimateOptions options; + options.dtype = torch::kBFloat16; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = int64_t{16} * 1024 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 10; + options.max_tokens_per_batch = 10240; + options.max_tokens_per_chunk_for_prefill = 2048; + options.enable_prefix_cache = true; + options.instance_role = InstanceRole::MIX; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + EXPECT_EQ(capacity.swa_count(), 181); + EXPECT_GT(capacity.c4_count(), 0); + EXPECT_GT(capacity.c128_count(), 0); + EXPECT_EQ(capacity.c4_count(), 32 * capacity.c128_count()); +} + +TEST(KVCacheEstimationTest, DeepSeekV4DecodeKeepsOperationalSwaPool) { + ModelArgs model_args; + model_args.model_type("deepseek_v4") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 4, 128}); + + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 128 * 1024 * 1024; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.enable_disagg_pd = true; + options.instance_role = InstanceRole::DECODE; + options.enable_prefix_cache = true; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + EXPECT_EQ(capacity.swa_count(), 25); + EXPECT_GT(capacity.c128_count(), 0); +} + +TEST(KVCacheEstimationTest, DeepSeekV4SwaOnlyPrefixKeepsOperationalSwaPool) { + ModelArgs model_args; + model_args.model_type("deepseek_v4_dspark") + .n_layers(3) + .head_dim(16) + .index_head_dim(8) + .window_size(257) + .compress_ratios({1, 1, 1}); + + constexpr int64_t kSwaBytesPerBlock = 3 * 128 * 16 * 4; + KVCacheEstimateOptions options; + options.dtype = torch::kFloat32; + options.kv_cache_dtype = "auto"; + options.cache_size_in_bytes = 43 * kSwaBytesPerBlock; + options.block_size = 128; + options.max_seqs_per_batch = 4; + options.max_tokens_per_chunk_for_prefill = 385; + options.enable_disagg_pd = true; + options.instance_role = InstanceRole::PREFILL; + options.enable_prefix_cache = true; + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + EXPECT_EQ(capacity.swa_count(), 33); + EXPECT_EQ(capacity.c4_count(), 0); + EXPECT_EQ(capacity.c128_count(), 0); +} + TEST(KVCacheEstimationTest, EstimatesDeepSeekV4DSparkSwaPool) { ModelArgs model_args; model_args.model_type("deepseek_v4_dspark") @@ -347,7 +635,7 @@ TEST(KVCacheEstimationTest, EstimatesDeepSeekV4DSparkSwaPool) { const KVCacheCapacity capacity = estimate_kv_cache_capacity(model_args, options); - EXPECT_EQ(capacity.swa_count(), 19); + EXPECT_EQ(capacity.swa_count(), 21); EXPECT_EQ(capacity.c4_count(), 0); EXPECT_EQ(capacity.c128_count(), 0); EXPECT_EQ(capacity.n_blocks(), 1); @@ -373,11 +661,16 @@ TEST(KVCacheEstimationTest, KVCacheEstimateOptions target_options; target_options.dtype = torch::kFloat32; target_options.kv_cache_dtype = "auto"; - target_options.cache_size_in_bytes = 2818048; + target_options.cache_size_in_bytes = + 2818048 + /*target_and_draft_swa_growth=*/229376; target_options.block_size = 128; target_options.max_seqs_per_batch = 4; KVCacheEstimateOptions draft_options = target_options; draft_options.is_draft_engine = true; + + const KVCacheCapacity target_only_capacity = + estimate_kv_cache_capacity(target_args, target_options); + target_options.draft_model_args = &draft_args; target_options.draft_options = &draft_options; @@ -385,13 +678,14 @@ TEST(KVCacheEstimationTest, estimate_kv_cache_capacity(target_args, target_options); constexpr int64_t kDraftSwaBytes = - /*layers=*/3 * /*swa_count=*/19 * /*block_size=*/128 * + /*layers=*/3 * /*swa_count=*/21 * /*block_size=*/128 * /*head_dim=*/16 * /*float32_bytes=*/4; EXPECT_LE(capacity.cache_size_in_bytes() + kDraftSwaBytes, target_options.cache_size_in_bytes); - EXPECT_EQ(capacity.swa_count(), 19); - EXPECT_EQ(capacity.c4_count(), 64); - EXPECT_EQ(capacity.c128_count(), 2); + EXPECT_EQ(capacity.swa_count(), target_only_capacity.swa_count()); + EXPECT_EQ(capacity.c4_count() + 32, target_only_capacity.c4_count()); + EXPECT_EQ(capacity.c128_count() + 1, target_only_capacity.c128_count()); + EXPECT_LT(capacity.n_blocks(), target_only_capacity.n_blocks()); } TEST(KVCacheEstimationTest, diff --git a/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp b/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp index f202cbde06..f6c119a070 100644 --- a/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp +++ b/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp @@ -35,6 +35,7 @@ limitations under the License. #include #include +#include "framework/kv_cache/kv_cache_capacity.h" #include "framework/kv_cache/kv_cache_shape.h" #include "framework/kv_cache_transfer/kv_cache_transfer.h" #include "platform/device.h" @@ -655,6 +656,8 @@ TEST(MooncakeKVCacheTransferDefaultTest, .n_heads(4) .n_kv_heads(4) .head_dim(1); + KVCacheCapacity cache_capacity; + cache_capacity.n_blocks(4).block_size(3); transfer.configure_cache_layout(make_args(/*rank=*/0, /*world_size=*/2, @@ -662,10 +665,12 @@ TEST(MooncakeKVCacheTransferDefaultTest, model_args, /*block_token_capacity=*/3, /*is_spec_draft=*/false); + const KVCacheShape main_shape(cache_capacity, model_args, /*world_size=*/2); std::vector main_caches; main_caches.emplace_back( - KVCacheTensors{torch::zeros({4, 3, 2, 1}), torch::zeros({4, 3, 2, 1})}); - transfer.register_kv_cache(main_caches, KVCacheShape(), torch::kFloat32); + KVCacheTensors{torch::zeros(main_shape.key_cache_shape()), + torch::zeros(main_shape.value_cache_shape())}); + transfer.register_kv_cache(main_caches, main_shape, torch::kFloat32); transfer.configure_cache_layout(make_args(/*rank=*/0, /*world_size=*/1, @@ -673,11 +678,12 @@ TEST(MooncakeKVCacheTransferDefaultTest, model_args, /*block_token_capacity=*/3, /*is_spec_draft=*/true); + const KVCacheShape draft_shape(cache_capacity, model_args, /*world_size=*/1); std::vector draft_caches; draft_caches.emplace_back( - KVCacheTensors{torch::zeros({4, 3, 4, 1}), torch::zeros({4, 3, 4, 1})}); - transfer.register_kv_cache_spec( - draft_caches, KVCacheShape(), torch::kFloat32); + KVCacheTensors{torch::zeros(draft_shape.key_cache_shape()), + torch::zeros(draft_shape.value_cache_shape())}); + transfer.register_kv_cache_spec(draft_caches, draft_shape, torch::kFloat32); EXPECT_EQ(transfer.local_cache_layout_.coordinates.tp_size, 2); ASSERT_EQ(transfer.local_cache_layout_.tensors.size(), 4U); diff --git a/tests/core/framework/sampling/CMakeLists.txt b/tests/core/framework/sampling/CMakeLists.txt index d5221dea37..eaee1f5814 100644 --- a/tests/core/framework/sampling/CMakeLists.txt +++ b/tests/core/framework/sampling/CMakeLists.txt @@ -42,6 +42,7 @@ cc_test( SRCS sampler_filter_mask_test.cpp DEPS + :platform :sampler GTest::gtest_main ) diff --git a/tests/core/framework/sampling/sampler_filter_mask_test.cpp b/tests/core/framework/sampling/sampler_filter_mask_test.cpp index 644f4e4129..f4c3862d31 100644 --- a/tests/core/framework/sampling/sampler_filter_mask_test.cpp +++ b/tests/core/framework/sampling/sampler_filter_mask_test.cpp @@ -16,6 +16,7 @@ limitations under the License. #include #include "core/framework/sampling/sampler.h" +#include "platform/platform.h" namespace xllm { namespace { @@ -32,6 +33,17 @@ SamplingParameters make_greedy_params(int64_t batch_size) { return params; } +bool uses_device_random_sample() { + return Platform::is_mlu() || Platform::is_cuda() || Platform::is_dcu(); +} + +torch::Device get_random_sample_device() { + if (uses_device_random_sample()) { + return torch::Device(Platform::type_torch(), 0); + } + return torch::Device(torch::kCPU); +} + TEST(SamplerFilterMaskTest, GreedySamplingHonorsMixedRows) { SamplingParameters params = make_greedy_params(/*batch_size=*/2); params.filter_mask = @@ -47,6 +59,10 @@ TEST(SamplerFilterMaskTest, GreedySamplingHonorsMixedRows) { } TEST(SamplerFilterMaskTest, RandomSamplingCannotSelectDisallowedToken) { + if (uses_device_random_sample() && Platform::device_count() == 0) { + GTEST_SKIP() << "Random sampling backend device is unavailable"; + } + SamplingParameters params = make_greedy_params(/*batch_size=*/1); params.do_sample = torch::ones({1}, torch::kBool); params.all_greedy_sample = false; @@ -54,6 +70,10 @@ TEST(SamplerFilterMaskTest, RandomSamplingCannotSelectDisallowedToken) { params.filter_mask = torch::tensor({{-1.0e9F, 0.0F, -1.0e9F}}); torch::Tensor logits = torch::tensor({{100.0F, 1.0F, 100.0F}}); + const torch::Device device = get_random_sample_device(); + params = params.to(device, torch::kFloat32); + logits = logits.to(device); + Sampler sampler; SampleOutput output = sampler.forward(logits, params); diff --git a/tests/core/kernels/npu/CMakeLists.txt b/tests/core/kernels/npu/CMakeLists.txt index 29f9e9a85b..f1de4769bb 100644 --- a/tests/core/kernels/npu/CMakeLists.txt +++ b/tests/core/kernels/npu/CMakeLists.txt @@ -16,6 +16,16 @@ cc_test( pybind11::embed ) +# Keep the embedded-Python NPU probes isolated from the parallel CTest pool. +# This limits shared-device contention but is not a substitute for validating +# each operator's supported parameter combinations. +set_tests_properties( + NpuXllmOpsTest.Dsv4QuantLightningIndexerPythonWrapperRunsOnA3 + NpuXllmOpsTest.Dsv4QuantLightningIndexerProductionShapeRunsOnA3 + NpuXllmOpsTest.Dsv4SparseAttentionPythonWrapperRunsOnA3 + PROPERTIES RUN_SERIAL TRUE +) + # Temporarily disabled: Dsv4ScatterCache tests are unstable in the current # NPU test environment. # cc_test( diff --git a/tests/core/kernels/npu/npu_xllm_ops_test.cpp b/tests/core/kernels/npu/npu_xllm_ops_test.cpp index 9bf7acf78b..182ae66d05 100644 --- a/tests/core/kernels/npu/npu_xllm_ops_test.cpp +++ b/tests/core/kernels/npu/npu_xllm_ops_test.cpp @@ -77,6 +77,12 @@ bool is_ascend950_device() { std::string(soc_name).find("Ascend950") != std::string::npos; } +bool is_ascend910_93_device() { + const char* soc_name = aclrtGetSocName(); + return soc_name != nullptr && + std::string(soc_name).find("Ascend910_93") != std::string::npos; +} + torch::Tensor expand_kv_heads_reference(const torch::Tensor& tensor, int64_t num_heads) { const int64_t num_kv_heads = tensor.size(1); @@ -244,6 +250,520 @@ TEST_F(NpuXllmOpsTest, EmbeddedInterpreterSeesOps) { .item(); } +TEST_F(NpuXllmOpsTest, Dsv4OpsUseNpuDispatchKeys) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch + +device_ops = ( + "moe_gating_top_k_hash", + "dequant_swiglu_quant", + "hc_pre", + "hc_post", + "compressor", + "sparse_attn_sharedkv", + "quant_lightning_indexer", +) +for op_name in device_ops: + qualname = f"xllm_ops::{op_name}" + assert torch._C._dispatch_has_kernel_for_dispatch_key( + qualname, "PrivateUse1" + ), qualname + assert not torch._C._dispatch_has_kernel_for_dispatch_key( + qualname, "CompositeExplicitAutograd" + ), qualname + +for op_name in ( + "sparse_attn_sharedkv_metadata", + "quant_lightning_indexer_metadata", +): + assert torch._C._dispatch_has_kernel_for_dispatch_key( + f"xllm_ops::{op_name}", "CompositeExplicitAutograd" + ), op_name +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4GroupGemmMatchesInt32Reference) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.moe import _group_gemm + +device = torch.device("privateuseone:0") +torch.manual_seed(20260814) +tokens, experts, input_dim, output_dim = 8, 2, 128, 256 +x_cpu = torch.randint(-4, 5, (tokens, input_dim), dtype=torch.int8) +w_cpu = torch.randint(-4, 5, (experts, input_dim, output_dim), dtype=torch.int8) +group_list_cpu = torch.tensor([4, 4], dtype=torch.int64) + +x = x_cpu.to(device) +w = w_cpu.to(device) +group_list = group_list_cpu.to(device) +out = _group_gemm( + x=x, + weight=w, + scale=None, + per_token_scale=None, + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=1, + output_dtype=torch.int32, +) +torch.npu.synchronize() + +expected = torch.cat(( + x_cpu[:4].to(torch.int32) @ w_cpu[0].to(torch.int32), + x_cpu[4:].to(torch.int32) @ w_cpu[1].to(torch.int32), +), dim=0) +assert out.shape == (tokens, output_dim) +assert out.dtype == torch.int32 +torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4GroupGemmAcceptsScaleAndPerTokenScale) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.moe import _group_gemm + +device = torch.device("privateuseone:0") +tokens, experts, input_dim, output_dim = 8, 2, 128, 128 +x = torch.randint(-4, 5, (tokens, input_dim), dtype=torch.int8, device=device) +w = torch.randint(-4, 5, (experts, input_dim, output_dim), dtype=torch.int8, device=device) +scale = torch.ones((experts, output_dim), dtype=torch.bfloat16, device=device) +per_token_scale = torch.ones((tokens,), dtype=torch.float32, device=device) +group_list = torch.tensor([4, 4], dtype=torch.int64, device=device) + +out = _group_gemm( + x=x, + weight=w, + scale=scale, + per_token_scale=per_token_scale, + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=1, + output_dtype=torch.bfloat16, +) +torch.npu.synchronize() +assert out.shape == (tokens, output_dim) +assert out.dtype == torch.bfloat16 +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4PartialRotaryPythonWrapperRunsOnNpu) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.rotary_embedding import ( + npu_inplace_partial_rotary_mul, +) + +torch.manual_seed(2026) +x_cpu = torch.randn((8, 2, 128), dtype=torch.float32).to(torch.bfloat16) +cos_cpu = torch.randn((8, 64), dtype=torch.float32).to(torch.bfloat16) +sin_cpu = torch.randn((8, 64), dtype=torch.float32).to(torch.bfloat16) + +expected = x_cpu.float().clone() +segment = x_cpu[..., 64:128].float() +swapped = torch.empty_like(segment) +swapped[..., 0::2] = segment[..., 1::2] +swapped[..., 1::2] = segment[..., 0::2] +sign = torch.ones_like(cos_cpu.float()) +sign[..., 0::2] = -1 +expected[..., 64:128] = ( + segment * cos_cpu.float().unsqueeze(1) + + swapped * sin_cpu.float().unsqueeze(1) * sign.unsqueeze(1) +) +expected = expected.to(torch.bfloat16).float() + +x = x_cpu.to("privateuseone:0") +cos = cos_cpu.to(x.device) +sin = sin_cpu.to(x.device) +result = npu_inplace_partial_rotary_mul(x, cos, sin, 64, 64) +torch.npu.synchronize() + +assert result.data_ptr() == x.data_ptr() +torch.testing.assert_close( + x.cpu().float(), expected, atol=2e-2, rtol=2e-2 +) +)PY"); +} + +TEST_F(NpuXllmOpsTest, Dsv4CompressorPythonWrapperRunsOnNpu) { + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.dsa import compressor + +device = torch.device("privateuseone:0") +torch.manual_seed(2025) +batch, tokens, hidden = 1, 128, 1024 +ratio, head_dim, coff, rope_dim = 128, 512, 1, 64 +compressed_tokens = tokens // ratio + +x_cpu = (torch.randn(batch, tokens, hidden) * 0.1).to(torch.float16) +wkv_cpu = (torch.randn(coff * head_dim, hidden) * 0.05).to(torch.float16) +wgate_cpu = (torch.randn(coff * head_dim, hidden) * 0.05).to(torch.float16) +ape_cpu = (torch.randn(ratio, coff * head_dim) * 0.1).float() +norm_cpu = (torch.randn(head_dim) * 0.1 + 1).to(torch.float16) +rope_cos_cpu = ( + torch.randn(batch, compressed_tokens, rope_dim) * 0.1 +).to(torch.float16) +rope_sin_cpu = ( + torch.randn(batch, compressed_tokens, rope_dim) * 0.1 +).to(torch.float16) + +projected_kv = x_cpu.float()[0] @ wkv_cpu.float().T +scores = x_cpu.float()[0] @ wgate_cpu.float().T + ape_cpu +pooled = (torch.softmax(scores, dim=0) * projected_kv).sum(0, keepdim=True) +variance = pooled.square().mean(-1, keepdim=True) +expected = pooled * torch.rsqrt(variance + 1e-6) * norm_cpu.float() +rope_segment = expected[:, -rope_dim:].clone() +half = rope_dim // 2 +rotated = torch.cat((-rope_segment[:, half:], rope_segment[:, :half]), dim=-1) +expected[:, -rope_dim:] = ( + rope_segment * rope_cos_cpu.float()[0] + + rotated * rope_sin_cpu.float()[0] +) +expected = expected.view(batch, compressed_tokens, head_dim).half().float() + +x = x_cpu.to(device) +wkv = wkv_cpu.to(device) +wgate = wgate_cpu.to(device) +ape = ape_cpu.to(device) +norm_weight = norm_cpu.to(device) +rope_sin = rope_sin_cpu.to(device) +rope_cos = rope_cos_cpu.to(device) +kv_state = torch.zeros((1, 128, head_dim), dtype=torch.float32, device=device) +score_state = torch.zeros_like(kv_state) +kv_block_table = torch.tensor([[0]], dtype=torch.int32, device=device) +score_block_table = torch.tensor([[0]], dtype=torch.int32, device=device) + +out, wkv_proj, softmax_res, norm_x, norm_rstd = compressor( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table, + score_block_table, + None, + None, + None, + rope_dim, + ratio, + coff, + 1e-6, + 1, + False, +) +torch.npu.synchronize() + +assert out.shape == (batch, compressed_tokens, head_dim) +assert out.dtype == torch.float16 +assert wkv_proj.numel() == 0 +assert softmax_res.numel() == 0 +assert norm_x.numel() == 0 +assert norm_rstd.numel() == 0 +torch.testing.assert_close( + out.cpu().float(), expected, atol=2e-2, rtol=2e-2 +) +)PY"); +} + +void run_dsv4_quant_lightning_indexer_probe(bool production_shape) { + py::gil_scoped_acquire gil; + py::dict locals; + locals["production_shape"] = production_shape; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.dsa import ( + quant_lightning_indexer, + quant_lightning_indexer_metadata, +) + +device = torch.device("privateuseone:0") +torch.manual_seed(2026) +heads, head_dim, page_size = 64, 128, 128 +batch = 1 +if production_shape: + q_tokens, kv_tokens = 84, 84 + sparse_count, cmp_ratio = 512, 4 + query_layout = "TND" + query_shape = (q_tokens, heads, head_dim) + weights_shape = (q_tokens, heads) + expected_indices_shape = (q_tokens, 1, sparse_count) +else: + q_tokens, kv_tokens = 4, 128 + sparse_count, cmp_ratio = 8, 1 + query_layout = "BSND" + query_shape = (batch, q_tokens, heads, head_dim) + weights_shape = (batch, q_tokens, heads) + expected_indices_shape = (batch, q_tokens, 1, sparse_count) + +query_cpu = torch.randint(-8, 8, query_shape, dtype=torch.int8) +key_cpu = torch.randint(-8, 8, (1, page_size, 1, head_dim), dtype=torch.int8) +query = query_cpu.to(device) +key = key_cpu.to(device) +weights = torch.ones(weights_shape, dtype=torch.float16, device=device) +query_scale = torch.ones_like(weights) +key_scale = torch.ones((1, page_size, 1), dtype=torch.float16, device=device) +query_lens = torch.tensor([q_tokens], dtype=torch.int32, device=device) +key_lens = torch.tensor([kv_tokens], dtype=torch.int32, device=device) +block_table = torch.tensor([[0]], dtype=torch.int32, device=device) +metadata = quant_lightning_indexer_metadata( + heads, + 1, + head_dim, + 0, + 0, + query_lens, + key_lens, + batch, + q_tokens, + kv_tokens, + query_layout, + "PA_BSND", + sparse_count, + 3, + 2**63 - 1, + 2**63 - 1, + cmp_ratio, + "npu", +) +metadata_again = quant_lightning_indexer_metadata( + heads, + 1, + head_dim, + 0, + 0, + query_lens, + key_lens, + batch, + q_tokens, + kv_tokens, + query_layout, + "PA_BSND", + sparse_count, + 3, + 2**63 - 1, + 2**63 - 1, + cmp_ratio, + "npu", +) +indices, values = quant_lightning_indexer( + query, + key, + weights, + query_scale, + key_scale, + 0, + 0, + query_lens, + key_lens, + block_table, + metadata, + query_layout, + "PA_BSND", + sparse_count, + 3, + 2**63 - 1, + 2**63 - 1, + cmp_ratio, + False, +) +torch.npu.synchronize() + +assert indices.shape == expected_indices_shape +assert indices.dtype == torch.int32 +assert values.numel() == 0 +assert values.dtype == torch.float32 +assert torch.equal(metadata.cpu(), metadata_again.cpu()) + +valid_key_count = kv_tokens // cmp_ratio +indices_cpu = indices.cpu() +assert torch.all( + (indices_cpu == -1) + | ((indices_cpu >= 0) & (indices_cpu < valid_key_count)) +) +keys = key_cpu[0, :valid_key_count, 0].float() +token_idx = q_tokens - 1 +if production_shape: + query_token = query_cpu[token_idx] + actual_indices = indices_cpu[token_idx, 0] +else: + query_token = query_cpu[0, token_idx] + actual_indices = indices_cpu[0, token_idx, 0] +dots = query_token.float() @ keys.T +expected_top8 = set(torch.topk(dots.clamp_min(0).sum(0), 8).indices.tolist()) +actual_top8 = set(actual_indices[:8].tolist()) +assert len(expected_top8 & actual_top8) >= 4, ( + sorted(expected_top8), + sorted(actual_top8), +) +)PY", + py::globals(), + locals); +} + +TEST_F(NpuXllmOpsTest, Dsv4QuantLightningIndexerPythonWrapperRunsOnA3) { + if (!is_ascend910_93_device()) { + GTEST_SKIP() << "Atlas A3 is required for this DSV4 QLI operator probe."; + } + run_dsv4_quant_lightning_indexer_probe(/*production_shape=*/false); +} + +TEST_F(NpuXllmOpsTest, Dsv4QuantLightningIndexerProductionShapeRunsOnA3) { + if (!is_ascend910_93_device()) { + GTEST_SKIP() << "Atlas A3 is required for the production-shape QLI probe."; + } + run_dsv4_quant_lightning_indexer_probe(/*production_shape=*/true); +} + +TEST_F(NpuXllmOpsTest, Dsv4SparseAttentionPythonWrapperRunsOnA3) { + if (!is_ascend910_93_device()) { + GTEST_SKIP() << "Atlas A3 is required for this DSV4 sparse-attention " + "operator probe."; + } + py::gil_scoped_acquire gil; + + py::exec(R"PY( +import torch +from xllm.python.kernels_npu.dsa import ( + sparse_attn_sharedkv, + sparse_attn_sharedkv_metadata, +) + +device = torch.device("privateuseone:0") +torch.manual_seed(1234) +batch, q_tokens, kv_tokens = 1, 4, 16 +heads, head_dim, page_size = 64, 512, 16 +query_cpu = (torch.randn(batch, q_tokens, heads, head_dim) * 0.1).half() +kv_cpu = (torch.randn(batch, kv_tokens, 1, head_dim) * 0.1).half() +sinks_cpu = (torch.randn(heads) * 0.1).float() +query = query_cpu.to(device) +ori_kv = kv_cpu.view(1, page_size, 1, head_dim).to(device) +block_table = torch.tensor([[0]], dtype=torch.int32, device=device) +cu_q = torch.tensor([0, q_tokens], dtype=torch.int32, device=device) +cu_kv = torch.tensor([0, kv_tokens], dtype=torch.int32, device=device) +seq_q = torch.tensor([q_tokens], dtype=torch.int32, device=device) +seq_kv = torch.tensor([kv_tokens], dtype=torch.int32, device=device) +sinks = sinks_cpu.to(device) +metadata = sparse_attn_sharedkv_metadata( + heads, + 1, + head_dim, + cu_q, + cu_kv, + None, + seq_q, + seq_kv, + batch, + q_tokens, + kv_tokens, + 0, + 0, + 4, + 4, + 3, + 127, + 0, + "BSND", + "PA_ND", + True, + False, +) +metadata_again = sparse_attn_sharedkv_metadata( + heads, + 1, + head_dim, + cu_q, + cu_kv, + None, + seq_q, + seq_kv, + batch, + q_tokens, + kv_tokens, + 0, + 0, + 4, + 4, + 3, + 127, + 0, + "BSND", + "PA_ND", + True, + False, +) +out, lse = sparse_attn_sharedkv( + query, + ori_kv, + None, + None, + None, + block_table, + None, + None, + None, + None, + None, + seq_kv, + sinks, + metadata, + head_dim**-0.5, + 4, + 4, + 3, + 127, + 0, + "BSND", + "PA_ND", + False, +) +torch.npu.synchronize() + +assert out.shape == query.shape +assert out.dtype == query.dtype +assert lse.numel() == 0 +assert torch.equal(metadata.cpu(), metadata_again.cpu()) + +expected = torch.zeros_like(query_cpu.float()) +keys = kv_cpu[0, :, 0].float() +scale = head_dim**-0.5 +for q_idx in range(q_tokens): + diagonal = kv_tokens - q_tokens + q_idx + left = max(diagonal - 127, 0) + right = diagonal + selected_keys = keys[left:right + 1] + logits = query_cpu[0, q_idx].float() @ selected_keys.T * scale + sink_logits = sinks_cpu[:, None] + normalizer = torch.logsumexp( + torch.cat((logits, sink_logits), dim=1), dim=1 + ) + probabilities = torch.exp(logits - normalizer[:, None]) + expected[0, q_idx] = probabilities @ selected_keys +expected = expected.half().float() +torch.testing.assert_close( + out.cpu().float(), expected, atol=2e-2, rtol=2e-2 +) +)PY"); +} + TEST_F(NpuXllmOpsTest, Qwen35_27B_TP4_FullAttentionMatchesReference) { py::gil_scoped_acquire gil; if (!is_ascend950_device()) { diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index ff7042cda7..5a408e4d8b 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -49,6 +49,7 @@ limitations under the License. #include "core/runtime/acl_graph_executor_impl.h" #include "core/runtime/acl_graph_persistent_param.h" #include "core/runtime/base_executor_impl.h" +#include "core/runtime/decode_graph_bucket.h" #include "core/runtime/dflash_worker_impl.h" #include "core/runtime/options.h" #include "core/runtime/speculative_worker_impl.h" @@ -1162,9 +1163,15 @@ TEST(AclGraphPersistentParamTest, SpecVerifyMetadataUsesTokenCapacity) { options, /*need_update_attn_mask=*/false, /*is_hybrid_linear_attention=*/false); - EXPECT_EQ(persistent_param.q_seq_lens().size(0), 30); - EXPECT_EQ(persistent_param.kv_seq_lens().size(0), 30); - EXPECT_EQ(persistent_param.persistent_block_tables().size(0), 30); + const int64_t expected_token_capacity = + runtime::get_decode_graph_token_bucket( + static_cast(options.max_seqs_per_batch()) * + options.num_decoding_tokens(), + options.enable_graph_mode_decode_no_padding()); + EXPECT_EQ(persistent_param.q_seq_lens().size(0), expected_token_capacity); + EXPECT_EQ(persistent_param.kv_seq_lens().size(0), expected_token_capacity); + EXPECT_EQ(persistent_param.persistent_block_tables().size(0), + expected_token_capacity); constexpr int64_t kValidateRows = 12; const torch::TensorOptions int_options = @@ -1319,13 +1326,22 @@ TEST(AclGraphPersistentParamTest, AuxHiddenStatesUseGraphTokenCapacity) { ::xllm::npu::GraphPersistentParam target_param(args, device, options); target_param.set_aux_hidden_states(aux_hidden_states); - EXPECT_EQ(target_param.aux_hidden_states().size(0), 30); + const int64_t expected_target_capacity = + runtime::get_decode_graph_token_bucket( + static_cast(options.max_seqs_per_batch()) * + options.num_decoding_tokens(), + options.enable_graph_mode_decode_no_padding()); + EXPECT_EQ(target_param.aux_hidden_states().size(0), expected_target_capacity); options.is_draft_engine(true); ::xllm::npu::GraphPersistentParam draft_param(args, device, options); draft_param.set_aux_hidden_states(aux_hidden_states.slice( /*dim=*/0, /*start=*/0, /*end=*/options.max_seqs_per_batch())); - EXPECT_EQ(draft_param.aux_hidden_states().size(0), 10); + const int64_t expected_draft_capacity = + runtime::get_decode_graph_token_bucket( + options.max_seqs_per_batch(), + options.enable_graph_mode_decode_no_padding()); + EXPECT_EQ(draft_param.aux_hidden_states().size(0), expected_draft_capacity); speculative_config.enable_atb_spec_kernel(original_enable_atb_spec_kernel); } diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index 28aa4c8a55..95b641ac3e 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -632,6 +632,43 @@ TEST(BlockManagerPoolTest, AllocateFailureRollsBackSharedPrefixBlocks) { (void)engine.release(); } +TEST(ContinuousSchedulerTest, + BlockExhaustionDefersRequestAfterCurrentBatchHasWork) { + ScopedConfigValue memory_threshold( + SchedulerConfig::get_instance() + .prefill_scheduling_memory_usage_threshold(), + 2.0); + auto engine = std::make_unique( + /*num_blocks=*/4, /*block_size=*/4); + ContinuousScheduler::Options options = + create_scheduler_options(/*max_tokens_per_batch=*/16, + /*max_seqs_per_batch=*/2, + /*num_speculative_tokens=*/0, + /*max_tokens_per_chunk_for_prefill=*/8, + /*dp_size=*/1); + auto scheduler = std::make_unique(engine.get(), options); + + auto scheduled_request = + generate_request_with_prompt_tokens({1, 2, 3, 4, 5, 6, 7, 8}, 1, 30000); + auto deferred_request = generate_request_with_prompt_tokens( + {9, 10, 11, 12, 13, 14, 15, 16}, 1, 30000); + ASSERT_TRUE(scheduler->add_request(scheduled_request)); + ASSERT_TRUE(scheduler->add_request(deferred_request)); + + std::vector batches = scheduler->prepare_batch_test(); + ASSERT_EQ(batches.size(), 1u); + ASSERT_EQ(batches.front().size(), 1u); + EXPECT_EQ(batches.front()[0], scheduled_request->sequences()[0].get()); + EXPECT_EQ(scheduler->get_running_requests().size(), 1u); + EXPECT_EQ(scheduler->get_waiting_requests_num(), 1u); + EXPECT_FALSE(deferred_request->finished()); + EXPECT_EQ( + deferred_request->sequences()[0]->kv_state().num_blocks(BlockType::KV), + 0u); + + (void)engine.release(); +} + TEST(ContinuousSchedulerTest, PDDecodeBestOfNExpandsAndSharesPromptViaPrefixCache) { // Disagg PD decode instance flow: diff --git a/tests/core/scheduler/scheduler_policy_test.cpp b/tests/core/scheduler/scheduler_policy_test.cpp index 7c914aeab3..6782441739 100644 --- a/tests/core/scheduler/scheduler_policy_test.cpp +++ b/tests/core/scheduler/scheduler_policy_test.cpp @@ -15,9 +15,11 @@ limitations under the License. #include "scheduler/scheduler_policy.h" +#include #include #include +#include #include #include #include @@ -27,6 +29,7 @@ limitations under the License. #include #include "continuous_scheduler.h" +#include "core/framework/config/kv_cache_store_config.h" #include "core/framework/config/scheduler_config.h" #include "distributed_runtime/engine.h" #include "framework/block/block_manager_pool.h" @@ -185,6 +188,154 @@ class PendingReleaseBlockManagerPool final : public BlockManagerPool { int32_t allocate_calls_ = 0; }; +class RestoreWaitingBlockManagerPool final : public BlockManagerPool { + public: + RestoreWaitingBlockManagerPool() + : BlockManagerPool(make_options(), /*dp_size=*/1) {} + + bool allocate(Sequence* /*sequence*/, size_t /*num_tokens*/) override { + ++allocate_calls_; + return !pending_async_release_; + } + + void allocate_shared(Sequence* sequence) override { + ++allocate_shared_calls_; + sequence->kv_state().set_prefix_cache_matched(); + } + + bool has_pending_async_block_release() const override { + return pending_async_release_; + } + + void set_pending_async_release(bool pending) { + pending_async_release_ = pending; + } + + int32_t allocate_calls() const { return allocate_calls_; } + + int32_t allocate_shared_calls() const { return allocate_shared_calls_; } + + private: + static BlockManagerPool::Options make_options() { + BlockManagerPool::Options options; + options.num_blocks_ = 256; + options.block_size_ = 128; + options.max_seqs_per_batch_ = 16; + options.enable_prefix_cache_ = true; + return options; + } + + bool pending_async_release_ = true; + int32_t allocate_calls_ = 0; + int32_t allocate_shared_calls_ = 0; +}; + +class DecodeVictimBlockManagerPool final : public BlockManagerPool { + public: + DecodeVictimBlockManagerPool() + : BlockManagerPool(make_options(), /*dp_size=*/1) {} + + bool allocate(Sequence* /*sequence*/, size_t /*num_tokens*/) override { + return false; + } + + bool supports_host_cache_restore() const override { return true; } + + void deallocate(Request* request) override { + BlockManagerPool::deallocate(request); + pending_async_release_ = true; + } + + bool has_pending_async_block_release() const override { + return pending_async_release_; + } + + void complete_async_release() { pending_async_release_ = false; } + + private: + static BlockManagerPool::Options make_options() { + BlockManagerPool::Options options; + options.num_blocks_ = 256; + options.block_size_ = 128; + options.max_seqs_per_batch_ = 16; + options.enable_prefix_cache_ = true; + return options; + } + + bool pending_async_release_ = false; +}; + +class RestorePriorityBlockManagerPool final : public BlockManagerPool { + public: + RestorePriorityBlockManagerPool() + : BlockManagerPool(make_options(), /*dp_size=*/1) {} + + void set_sequences(Sequence* blocked, Sequence* victim) { + blocked_ = blocked; + victim_ = victim; + } + + bool allocate(Sequence* sequence, size_t /*num_tokens*/) override { + if (sequence == blocked_) { + if (free_blocks_ < kBlockedGrowthBlocks) { + return false; + } + free_blocks_ -= kBlockedGrowthBlocks; + return true; + } + if (sequence == victim_) { + if (free_blocks_ < kVictimRestoreBlocks) { + return false; + } + free_blocks_ -= kVictimRestoreBlocks; + return true; + } + return false; + } + + void allocate_shared(Sequence* sequence) override { + if (sequence != victim_) { + return; + } + CHECK_GT(sequence->num_tokens(), 0u); + sequence->kv_state().set_prefix_cache_matched(); + sequence->kv_state().set_kv_cache_tokens_num(sequence->num_tokens() - 1); + } + + bool supports_host_cache_restore() const override { return true; } + + void deallocate(Request* request) override { + BlockManagerPool::deallocate(request); + pending_async_release_ = true; + } + + bool has_pending_async_block_release() const override { + return pending_async_release_; + } + + void complete_release() { + pending_async_release_ = false; + free_blocks_ = kVictimRestoreBlocks; + } + + private: + static BlockManagerPool::Options make_options() { + BlockManagerPool::Options options; + options.num_blocks_ = 10; + options.block_size_ = 128; + options.max_seqs_per_batch_ = 16; + options.enable_prefix_cache_ = true; + return options; + } + + static constexpr size_t kBlockedGrowthBlocks = 1; + static constexpr size_t kVictimRestoreBlocks = 3; + Sequence* blocked_ = nullptr; + Sequence* victim_ = nullptr; + size_t free_blocks_ = 0; + bool pending_async_release_ = false; +}; + class TestUnifiedPolicy final : public UnifiedPolicy { public: using UnifiedPolicy::UnifiedPolicy; @@ -416,6 +567,7 @@ TEST(SchedulerPolicyTest, KvlessCompositeReprobesAfterPartialAllocation) { DequeQueue chunk_queue; DequeQueue decode_queue; std::list> unified_queue; + std::deque decode_restore_waiting; std::vector> running_requests; std::vector running_sequences; std::vector running_sequence_budgets; @@ -425,6 +577,7 @@ TEST(SchedulerPolicyTest, KvlessCompositeReprobesAfterPartialAllocation) { .chunk_queue = chunk_queue, .decode_queue = decode_queue, .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, .running_requests = running_requests, .running_sequences = running_sequences, .running_sequences_budgets = running_sequence_budgets, @@ -477,6 +630,7 @@ TEST(SchedulerPolicyTest, UnifiedRetryRefreshesHostRestoreBeforeChunkSizing) { DequeQueue chunk_queue; DequeQueue decode_queue; std::list> unified_queue{requests[0], requests[1]}; + std::deque decode_restore_waiting; std::vector> running_requests; std::vector running_sequences; std::vector running_sequence_budgets; @@ -486,6 +640,7 @@ TEST(SchedulerPolicyTest, UnifiedRetryRefreshesHostRestoreBeforeChunkSizing) { .chunk_queue = chunk_queue, .decode_queue = decode_queue, .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, .running_requests = running_requests, .running_sequences = running_sequences, .running_sequences_budgets = running_sequence_budgets, @@ -547,6 +702,7 @@ TEST(SchedulerPolicyTest, DefersWhileAsyncBlockReleaseIsPending) { DequeQueue chunk_queue; DequeQueue decode_queue; std::list> unified_queue{requests.front()}; + std::deque decode_restore_waiting; std::vector> running_requests; std::vector running_sequences; std::vector running_sequence_budgets; @@ -556,6 +712,7 @@ TEST(SchedulerPolicyTest, DefersWhileAsyncBlockReleaseIsPending) { .chunk_queue = chunk_queue, .decode_queue = decode_queue, .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, .running_requests = running_requests, .running_sequences = running_sequences, .running_sequences_budgets = running_sequence_budgets, @@ -585,6 +742,435 @@ TEST(SchedulerPolicyTest, DefersWhileAsyncBlockReleaseIsPending) { EXPECT_TRUE(finished.empty()); } +TEST(SchedulerPolicyTest, + RestoreWaitingDefersUntilAsyncReleaseThenSchedulesPrefill) { + constexpr int32_t kPromptTokens = 128; + constexpr int32_t kHostRestoreTokens = 256; + ContinuousScheduler::Options options = create_scheduler_options( + /*max_tokens_per_batch=*/16384, + /*max_seqs_per_batch=*/16, + /*num_speculative_tokens=*/0, + /*max_tokens_per_chunk_for_prefill=*/16384, + /*dp_size=*/1); + BatchMode mode{ + .enable_mix_batch = false, + .enable_chunked_prefill = true, + .priority_strategy = "fcfs", + }; + PrefillFirstPolicy policy(mode, options); + RestoreWaitingBlockManagerPool block_manager_pool; + BlockManagerPool::Options prefix_cache_options; + prefix_cache_options.num_blocks_ = 4; + prefix_cache_options.block_size_ = 128; + prefix_cache_options.max_seqs_per_batch_ = 16; + BlockManagerPool prefix_cache_pool(prefix_cache_options, /*dp_size=*/1); + + std::vector> requests = + generate_request({kPromptTokens}, + {kPromptTokens + 1}, + std::nullopt, + std::nullopt, + /*max_context_len=*/40000); + std::vector> cached_requests = + generate_request({kHostRestoreTokens}, + {1}, + std::nullopt, + std::nullopt, + /*max_context_len=*/40000); + Sequence* sequence = requests.front()->sequences().front().get(); + Sequence* cached_sequence = + cached_requests.front()->sequences().front().get(); + sequence->kv_state().set_kv_cache_tokens_num(kPromptTokens); + for (int32_t token_id = 0; token_id <= kPromptTokens; ++token_id) { + sequence->append_token(token_id); + } + sequence->kv_state().set_kv_cache_tokens_num(0); + ASSERT_TRUE(prefix_cache_pool.allocate(cached_sequence)); + sequence->host_kv_state().mount_composite_shared( + BlockType::KV, cached_sequence->kv_state().take_blocks(BlockType::KV)); + sequence->host_kv_state().set_kv_cache_tokens_num(kHostRestoreTokens); + sequence->host_kv_state().set_prefix_cache_matched(); + sequence->set_host_cache_match(/*restore_tokens=*/kHostRestoreTokens, + /*copy_units=*/2); + EXPECT_EQ(requests.front()->num_prefix_cache_tokens(), 0u); + DequeQueue prefill_queue; + DequeQueue chunk_queue; + DequeQueue decode_queue; + std::list> unified_queue; + std::deque decode_restore_waiting{ + {requests.front(), absl::Now()}}; + std::vector> running_requests; + std::vector running_sequences; + std::vector running_sequence_budgets; + bool last_step_prefill = false; + SchedulerState state{ + .prefill_queue = prefill_queue, + .chunk_queue = chunk_queue, + .decode_queue = decode_queue, + .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, + .running_requests = running_requests, + .running_sequences = running_sequences, + .running_sequences_budgets = running_sequence_budgets, + .kv_cache_manager = &block_manager_pool, + .profile_manager = nullptr, + .response_processor = nullptr, + .last_step_prefill = last_step_prefill, + .options = options, + .min_speculative_tokens_required = 0, + .enable_prefix_cache = true, + .has_linear_attention_layers = false, + }; + ScheduleBudget budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + std::vector> finished; + + policy.schedule(state, budget, finished); + + EXPECT_EQ(block_manager_pool.allocate_calls(), 0); + EXPECT_EQ(decode_restore_waiting.size(), 1u); + EXPECT_TRUE(running_sequences.empty()); + + block_manager_pool.set_pending_async_release(false); + policy.schedule(state, budget, finished); + + EXPECT_EQ(block_manager_pool.allocate_shared_calls(), 1); + EXPECT_EQ(block_manager_pool.allocate_calls(), 1); + EXPECT_TRUE(decode_restore_waiting.empty()); + ASSERT_EQ(running_sequences.size(), 1u); + EXPECT_EQ(running_sequences.front(), + requests.front()->sequences().front().get()); + EXPECT_EQ(requests.front()->num_prefix_cache_tokens(), 0u); + EXPECT_LE(requests.front()->num_prefix_cache_tokens(), + sequence->num_prompt_tokens()); + EXPECT_TRUE(finished.empty()); +} + +TEST(SchedulerPolicyTest, PendingDecodeReleaseStopsFurtherPreemption) { + ScopedConfigValue host_blocks_factor( + KVCacheStoreConfig::get_instance().host_blocks_factor(), 2.0); + ContinuousScheduler::Options options = create_scheduler_options( + /*max_tokens_per_batch=*/16384, + /*max_seqs_per_batch=*/16, + /*num_speculative_tokens=*/0, + /*max_tokens_per_chunk_for_prefill=*/16384, + /*dp_size=*/1); + options.enable_disagg_pd() = true; + options.enable_pd_ooc() = false; + options.instance_role() = InstanceRole::DECODE; + options.enable_schedule_overlap() = false; + BatchMode mode{ + .enable_mix_batch = false, + .enable_chunked_prefill = true, + .priority_strategy = "fcfs", + }; + PrefillFirstPolicy policy(mode, options); + DecodeVictimBlockManagerPool block_manager_pool; + + std::vector> requests = + generate_request({128, 128, 128}, + {2, 2, 2}, + std::nullopt, + std::nullopt, + /*max_context_len=*/40000); + for (const std::shared_ptr& request : requests) { + Sequence* sequence = request->sequences().front().get(); + sequence->kv_state().set_kv_cache_tokens_num(sequence->num_prompt_tokens()); + sequence->append_token(Token(1)); + } + DequeQueue prefill_queue; + DequeQueue chunk_queue; + DequeQueue decode_queue; + decode_queue.push(requests[0], /*if_back=*/true); + decode_queue.push(requests[1], /*if_back=*/true); + decode_queue.push(requests[2], /*if_back=*/true); + std::list> unified_queue; + std::deque decode_restore_waiting; + std::vector> running_requests; + std::vector running_sequences; + std::vector running_sequence_budgets; + bool last_step_prefill = false; + SchedulerState state{ + .prefill_queue = prefill_queue, + .chunk_queue = chunk_queue, + .decode_queue = decode_queue, + .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, + .running_requests = running_requests, + .running_sequences = running_sequences, + .running_sequences_budgets = running_sequence_budgets, + .kv_cache_manager = &block_manager_pool, + .profile_manager = nullptr, + .response_processor = nullptr, + .last_step_prefill = last_step_prefill, + .options = options, + .min_speculative_tokens_required = 0, + .enable_prefix_cache = true, + .has_linear_attention_layers = false, + }; + ScheduleBudget budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + std::vector> finished; + + policy.schedule(state, budget, finished); + + EXPECT_EQ(decode_restore_waiting.size(), 1u); + ASSERT_FALSE(decode_restore_waiting.empty()); + EXPECT_EQ(decode_restore_waiting.front().request.get(), requests[2].get()); + EXPECT_TRUE(requests[2]->preempted()); + EXPECT_TRUE(requests[2]->sequences().front()->is_prefill_stage()); + EXPECT_FALSE(requests[1]->preempted()); + EXPECT_EQ(decode_queue.size(), 2u); + ASSERT_FALSE(decode_queue.empty()); + EXPECT_EQ(decode_queue.top().get(), requests[0].get()); + EXPECT_EQ(decode_queue.back().get(), requests[1].get()); + EXPECT_EQ(budget.num_preempted_requests, 1u); + EXPECT_TRUE(prefill_queue.empty()); + EXPECT_TRUE(running_sequences.empty()); + EXPECT_TRUE(finished.empty()); + + ScheduleBudget pending_budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + policy.schedule(state, pending_budget, finished); + + EXPECT_EQ(decode_restore_waiting.size(), 1u); + EXPECT_FALSE(requests[1]->preempted()); + EXPECT_EQ(decode_queue.size(), 2u); + EXPECT_EQ(decode_queue.back().get(), requests[1].get()); + EXPECT_EQ(pending_budget.num_preempted_requests, 0u); + + block_manager_pool.complete_async_release(); + ScheduleBudget released_budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + policy.schedule(state, released_budget, finished); + + EXPECT_EQ(decode_restore_waiting.size(), 2u); + EXPECT_TRUE(requests[1]->preempted()); + EXPECT_EQ(decode_queue.size(), 1u); + EXPECT_EQ(released_budget.num_preempted_requests, 1u); +} + +TEST(SchedulerPolicyTest, RetriesBlockedDecodeBeforeRestoringVictim) { + ScopedConfigValue host_blocks_factor( + KVCacheStoreConfig::get_instance().host_blocks_factor(), 2.0); + ContinuousScheduler::Options options = create_scheduler_options( + /*max_tokens_per_batch=*/16384, + /*max_seqs_per_batch=*/16, + /*num_speculative_tokens=*/0, + /*max_tokens_per_chunk_for_prefill=*/16384, + /*dp_size=*/1); + options.enable_disagg_pd() = true; + options.enable_pd_ooc() = false; + options.instance_role() = InstanceRole::DECODE; + options.enable_schedule_overlap() = false; + BatchMode mode{ + .enable_mix_batch = false, + .enable_chunked_prefill = true, + .priority_strategy = "fcfs", + }; + PrefillFirstPolicy policy(mode, options); + RestorePriorityBlockManagerPool block_manager_pool; + + std::vector> requests = + generate_request({128, 128}, + {2, 2}, + std::nullopt, + std::nullopt, + /*max_context_len=*/40000); + Sequence* blocked = requests[0]->sequences().front().get(); + Sequence* victim = requests[1]->sequences().front().get(); + for (const std::shared_ptr& request : requests) { + Sequence* sequence = request->sequences().front().get(); + sequence->kv_state().set_kv_cache_tokens_num(sequence->num_prompt_tokens()); + sequence->append_token(Token(1)); + } + block_manager_pool.set_sequences(blocked, victim); + + DequeQueue prefill_queue; + DequeQueue chunk_queue; + DequeQueue decode_queue; + decode_queue.push(requests[0], /*if_back=*/true); + decode_queue.push(requests[1], /*if_back=*/true); + std::list> unified_queue; + std::deque decode_restore_waiting; + std::vector> running_requests; + std::vector running_sequences; + std::vector running_sequence_budgets; + bool last_step_prefill = false; + SchedulerState state{ + .prefill_queue = prefill_queue, + .chunk_queue = chunk_queue, + .decode_queue = decode_queue, + .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, + .running_requests = running_requests, + .running_sequences = running_sequences, + .running_sequences_budgets = running_sequence_budgets, + .kv_cache_manager = &block_manager_pool, + .profile_manager = nullptr, + .response_processor = nullptr, + .last_step_prefill = last_step_prefill, + .options = options, + .min_speculative_tokens_required = 0, + .enable_prefix_cache = true, + .has_linear_attention_layers = false, + }; + ScheduleBudget initial_budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + std::vector> finished; + + policy.schedule(state, initial_budget, finished); + + ASSERT_EQ(decode_restore_waiting.size(), 1u); + EXPECT_EQ(decode_restore_waiting.front().request.get(), requests[1].get()); + EXPECT_EQ(decode_queue.size(), 1u); + EXPECT_EQ(decode_queue.top().get(), requests[0].get()); + EXPECT_TRUE(running_sequences.empty()); + + block_manager_pool.complete_release(); + ScheduleBudget released_budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + + policy.schedule(state, released_budget, finished); + + ASSERT_EQ(running_sequences.size(), 1u); + EXPECT_EQ(running_sequences.front(), blocked); + ASSERT_EQ(decode_restore_waiting.size(), 1u); + EXPECT_EQ(decode_restore_waiting.front().request.get(), requests[1].get()); + EXPECT_TRUE(decode_queue.empty()); +} + +TEST(SchedulerPolicyTest, + DecodeFirstRetriesBlockedDecodeBeforeRestoringVictim) { + ScopedConfigValue host_blocks_factor( + KVCacheStoreConfig::get_instance().host_blocks_factor(), 2.0); + ContinuousScheduler::Options options = create_scheduler_options( + /*max_tokens_per_batch=*/16384, + /*max_seqs_per_batch=*/16, + /*num_speculative_tokens=*/0, + /*max_tokens_per_chunk_for_prefill=*/16384, + /*dp_size=*/1); + options.enable_disagg_pd() = true; + options.enable_pd_ooc() = false; + options.instance_role() = InstanceRole::DECODE; + options.enable_schedule_overlap() = false; + BatchMode mode{ + .enable_mix_batch = true, + .enable_chunked_prefill = true, + .priority_strategy = "fcfs", + }; + DecodeFirstPolicy policy(mode, options); + RestorePriorityBlockManagerPool block_manager_pool; + + std::vector> requests = + generate_request({128, 128}, + {2, 2}, + std::nullopt, + std::nullopt, + /*max_context_len=*/40000); + Sequence* blocked = requests[0]->sequences().front().get(); + Sequence* victim = requests[1]->sequences().front().get(); + for (const std::shared_ptr& request : requests) { + Sequence* sequence = request->sequences().front().get(); + sequence->kv_state().set_kv_cache_tokens_num(sequence->num_prompt_tokens()); + sequence->append_token(Token(1)); + } + block_manager_pool.set_sequences(blocked, victim); + + DequeQueue prefill_queue; + DequeQueue chunk_queue; + DequeQueue decode_queue; + decode_queue.push(requests[0], /*if_back=*/true); + decode_queue.push(requests[1], /*if_back=*/true); + std::list> unified_queue; + std::deque decode_restore_waiting; + std::vector> running_requests; + std::vector running_sequences; + std::vector running_sequence_budgets; + bool last_step_prefill = false; + SchedulerState state{ + .prefill_queue = prefill_queue, + .chunk_queue = chunk_queue, + .decode_queue = decode_queue, + .unified_queue = unified_queue, + .decode_restore_waiting = decode_restore_waiting, + .running_requests = running_requests, + .running_sequences = running_sequences, + .running_sequences_budgets = running_sequence_budgets, + .kv_cache_manager = &block_manager_pool, + .profile_manager = nullptr, + .response_processor = nullptr, + .last_step_prefill = last_step_prefill, + .options = options, + .min_speculative_tokens_required = 0, + .enable_prefix_cache = true, + .has_linear_attention_layers = false, + }; + ScheduleBudget initial_budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + std::vector> finished; + + policy.schedule(state, initial_budget, finished); + + ASSERT_EQ(decode_restore_waiting.size(), 1u); + EXPECT_EQ(decode_restore_waiting.front().request.get(), requests[1].get()); + EXPECT_EQ(decode_queue.size(), 1u); + EXPECT_EQ(decode_queue.top().get(), requests[0].get()); + EXPECT_TRUE(running_sequences.empty()); + + block_manager_pool.complete_release(); + ScheduleBudget released_budget{ + .remaining_token_budget = 16384, + .remaining_seq_budget = 16, + .latency_budget = std::numeric_limits::max(), + .estimate_latency = 0, + .num_preempted_requests = 0, + }; + + policy.schedule(state, released_budget, finished); + + ASSERT_EQ(running_sequences.size(), 1u); + EXPECT_EQ(running_sequences.front(), blocked); + ASSERT_EQ(decode_restore_waiting.size(), 1u); + EXPECT_EQ(decode_restore_waiting.front().request.get(), requests[1].get()); + EXPECT_TRUE(decode_queue.empty()); +} + // TEST-2: // memory or budget not enough TEST(SchedulerPolicyTest, ResourceNotEnough) { diff --git a/tests/python/test_glm5_2_parallel.py b/tests/python/test_glm5_2_parallel.py index 70de4e59f4..707529591a 100644 --- a/tests/python/test_glm5_2_parallel.py +++ b/tests/python/test_glm5_2_parallel.py @@ -74,7 +74,9 @@ def test_full_world_ep_partitions_glm_experts() -> None: assert moe.local_expert_start == 6 assert moe.local_expert_end == 8 + moe.allocate_experts_w13_for_loading() assert moe.experts_w13.shape == (2, 16, 16) + moe.allocate_experts_w2_for_loading() assert moe.experts_w2.shape == (2, 16, 8) diff --git a/tests/python/test_grouped_moe.py b/tests/python/test_grouped_moe.py new file mode 100644 index 0000000000..08c6a923d7 --- /dev/null +++ b/tests/python/test_grouped_moe.py @@ -0,0 +1,149 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contracts for the NPU pre-selected grouped MoE path.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch + +_REPO_ROOT = Path(__file__).parents[2] + + +def _load_npu_moe_module(): + path = _REPO_ROOT / "xllm/python/kernels_npu/moe.py" + spec = importlib.util.spec_from_file_location("pr5_npu_moe", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_selected_expert_moe_matches_native_call_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from xllm.python import kernels + + moe = _load_npu_moe_module() + + hidden = torch.empty(3, 16, dtype=torch.bfloat16) + topk_weights = torch.ones(3, 2, dtype=torch.bfloat16) + topk_ids = torch.tensor([[4, 0], [5, 9], [7, 6]], dtype=torch.int32) + expanded = torch.empty(6, 16, dtype=torch.bfloat16) + row_ids = torch.arange(6, dtype=torch.int32) + expert_tokens = torch.tensor([1, 3, 5, 6, 7, 8], dtype=torch.int64) + quantized = torch.empty(6, 16, dtype=torch.int8) + input_scale = torch.empty(6, dtype=torch.float32) + gemm1 = torch.empty(6, 32, dtype=torch.int32) + activated = torch.empty(6, 16, dtype=torch.int8) + activation_scale = torch.empty(6, dtype=torch.float32) + gemm2 = torch.empty(6, 16, dtype=torch.bfloat16) + calls: list[tuple[str, object]] = [] + + def init_routing(*args, **kwargs): + calls.append(("routing", kwargs)) + return expanded, row_ids, expert_tokens, torch.empty(0) + + def dynamic_quant(value): + assert value is expanded + calls.append(("dynamic_quant", value)) + return quantized, input_scale + + def dequant_swiglu_quant(**kwargs): + calls.append(("dequant_swiglu_quant", kwargs)) + return activated, activation_scale + + gemm_calls: list[dict[str, object]] = [] + + def group_gemm(**kwargs): + gemm_calls.append(kwargs) + return gemm1 if len(gemm_calls) == 1 else gemm2 + + def token_unpermute(**kwargs): + calls.append(("unpermute", kwargs)) + return hidden + + monkeypatch.setattr(moe, "_group_gemm", group_gemm) + monkeypatch.setattr(moe.torch_npu, "npu_moe_init_routing_v2", init_routing) + monkeypatch.setattr(moe.torch_npu, "npu_moe_token_unpermute", token_unpermute) + monkeypatch.setattr(kernels, "dynamic_quant", dynamic_quant, raising=False) + monkeypatch.setattr(kernels, "dequant_swiglu_quant", dequant_swiglu_quant, raising=False) + + result = moe._grouped_moe_with_selected_experts_impl( + hidden, + topk_weights, + topk_ids, + torch.empty(4, 16, 32, dtype=torch.int8), + torch.empty(4, 16, 16, dtype=torch.int8), + torch.empty(4, 32), + torch.empty(4, 16), + num_total_experts=16, + start_expert_id=4, + num_experts_per_rank=4, + swiglu_limit=7.0, + ) + + assert result is hidden + routing = dict(calls)["routing"] + assert isinstance(routing, dict) + assert routing["active_expert_range"] == [4, 8] + assert routing["expert_num"] == 16 + assert routing["quant_mode"] == -1 + + assert len(gemm_calls) == 2 + assert gemm_calls[0]["scale"] is None + assert gemm_calls[0]["per_token_scale"] is None + assert gemm_calls[0]["output_dtype"] == torch.int32 + assert gemm_calls[1]["scale"].dtype == torch.bfloat16 + assert gemm_calls[1]["per_token_scale"] is activation_scale + assert gemm_calls[1]["output_dtype"] == torch.bfloat16 + assert all(torch.equal(call["group_list"], expert_tokens[:4]) for call in gemm_calls) + assert all(call["group_list"].numel() == 4 for call in gemm_calls) + assert all(call["group_list_type"] == 1 for call in gemm_calls) + + dequant = dict(calls)["dequant_swiglu_quant"] + assert isinstance(dequant, dict) + assert dequant["x"] is gemm1 + assert dequant["activation_scale"] is input_scale + assert torch.equal(dequant["group_index"], expert_tokens[:4]) + assert dequant["clamp_limit"] == 7.0 + + unpermute = dict(calls)["unpermute"] + assert isinstance(unpermute, dict) + torch.testing.assert_close( + unpermute["probs"], + torch.tensor([[1, 0], [1, 0], [1, 1]], dtype=torch.bfloat16), + ) + + +def test_selected_expert_moe_rejects_an_invalid_active_range() -> None: + moe = _load_npu_moe_module() + + with pytest.raises(ValueError, match="active expert range"): + moe._grouped_moe_with_selected_experts_impl( + torch.empty(1, 16, dtype=torch.bfloat16), + torch.ones(1, 1, dtype=torch.bfloat16), + torch.zeros(1, 1, dtype=torch.int32), + torch.empty(4, 16, 32, dtype=torch.int8), + torch.empty(4, 16, 16, dtype=torch.int8), + torch.empty(4, 32), + torch.empty(4, 16), + num_total_experts=16, + start_expert_id=14, + num_experts_per_rank=4, + ) diff --git a/tests/python/test_registry.py b/tests/python/test_registry.py index c61b534a03..999e9569b7 100644 --- a/tests/python/test_registry.py +++ b/tests/python/test_registry.py @@ -21,10 +21,10 @@ def test_unsupported_model_fails_before_import(monkeypatch: pytest.MonkeyPatch) -> None: import_model = Mock() - monkeypatch.setattr(registry.current_platform, "device_type", lambda: "npu") + monkeypatch.setattr(registry.current_platform, "device_type", lambda: "cuda") monkeypatch.setattr(registry, "import_module", import_model) - with pytest.raises(NotImplementedError, match="qwen3_5.*npu"): - registry.get_model_class("qwen3_5") + with pytest.raises(NotImplementedError, match="qwen3_vl.*cuda"): + registry.get_model_class("qwen3_vl") import_model.assert_not_called() diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index aa167eef61..8f689410e5 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -517,11 +517,19 @@ KVCacheCapacity LLMEngine::estimate_kv_cache_capacity() { static_cast(options_.num_speculative_tokens()); estimate_options.max_tokens_per_batch = static_cast(options_.max_tokens_per_batch()); + estimate_options.max_tokens_per_chunk_for_prefill = + static_cast(options_.max_tokens_per_chunk_for_prefill()); estimate_options.max_linear_state_cache_slots = options_.max_linear_state_cache_slots(); estimate_options.is_draft_engine = options_.is_draft_engine(); + estimate_options.enable_chunked_prefill = options_.enable_chunked_prefill(); + estimate_options.enable_schedule_overlap = options_.enable_schedule_overlap(); + const KVCacheConfig& kv_cache_config = KVCacheConfig::get_instance(); estimate_options.enable_prefix_cache = - ::xllm::KVCacheConfig::get_instance().enable_prefix_cache(); + kv_cache_config.enable_prefix_cache() && + !kv_cache_config.enable_xtensor(); + estimate_options.enable_disagg_pd = options_.enable_disagg_pd(); + estimate_options.instance_role = options_.instance_role(); if (options_.enable_mtp_draft_body_tp1() && options_.is_draft_engine()) { estimate_options.world_size = 1; estimate_options.n_local_kv_heads = @@ -690,9 +698,14 @@ bool LLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { : semantic_window; const uint32_t swa_blocks_per_seq = static_cast( get_swa_blocks_per_seq(effective_window, block_size)); + CHECK_LE(kv_cache_cap.swa_count(), + static_cast(std::numeric_limits::max())) + << "DSV4 swa_count exceeds uint32_t range: " + << kv_cache_cap.swa_count(); options.sliding_window_size(static_cast(effective_window)) .swa_blocks_per_seq(swa_blocks_per_seq) + .swa_num_blocks(static_cast(kv_cache_cap.swa_count())) .max_tokens_per_batch(options_.max_tokens_per_batch()) .manager_types(std::move(manager_types)) .compress_ratios(std::move(manager_compress_ratios)); diff --git a/xllm/core/distributed_runtime/vlm_engine.cpp b/xllm/core/distributed_runtime/vlm_engine.cpp index abd0eb25f6..9eb9766c5c 100644 --- a/xllm/core/distributed_runtime/vlm_engine.cpp +++ b/xllm/core/distributed_runtime/vlm_engine.cpp @@ -268,11 +268,21 @@ KVCacheCapacity VLMEngine::estimate_kv_cache_capacity() { estimate_options.n_local_linear_v_heads = n_local_linear_v_heads_; estimate_options.max_seqs_per_batch = static_cast(options_.max_seqs_per_batch()); + estimate_options.max_tokens_per_batch = + static_cast(options_.max_tokens_per_batch()); + estimate_options.max_tokens_per_chunk_for_prefill = + static_cast(options_.max_tokens_per_chunk_for_prefill()); estimate_options.max_linear_state_cache_slots = options_.max_linear_state_cache_slots(); estimate_options.is_draft_engine = options_.is_draft_engine(); + estimate_options.enable_chunked_prefill = options_.enable_chunked_prefill(); + estimate_options.enable_schedule_overlap = options_.enable_schedule_overlap(); + const KVCacheConfig& kv_cache_config = KVCacheConfig::get_instance(); estimate_options.enable_prefix_cache = - ::xllm::KVCacheConfig::get_instance().enable_prefix_cache(); + kv_cache_config.enable_prefix_cache() && + !kv_cache_config.enable_xtensor(); + estimate_options.enable_disagg_pd = options_.enable_disagg_pd(); + estimate_options.instance_role = options_.instance_role(); KVCacheCapacity kv_cache_cap = ::xllm::estimate_kv_cache_capacity(args_, estimate_options); diff --git a/xllm/core/framework/block/block_manager.h b/xllm/core/framework/block/block_manager.h index ef8023709c..cf6933805e 100644 --- a/xllm/core/framework/block/block_manager.h +++ b/xllm/core/framework/block/block_manager.h @@ -51,6 +51,8 @@ class BlockManager { PROPERTY(uint32_t, sliding_window_size) = 0; // Base SWA/cache-state block rows retained per sequence. PROPERTY(uint32_t, swa_blocks_per_seq) = 0; + // Total physical SWA rows computed by the KV cache estimator. + PROPERTY(uint32_t, swa_num_blocks) = 0; // Scheduler token budget used to size the shared SWA burst pool. PROPERTY(uint32_t, max_tokens_per_batch) = 0; // For CompositeBlockManager (passed from upstream). @@ -192,9 +194,9 @@ class BlockManager { size_t num_tokens) = 0; // Sliding-window hook: release leading blocks that have slid out of the - // window. The composite calls this on every leaf AFTER a successful - // allocate_sequence commit; non-SWA leaves keep the empty default (no-op). - // Running post-commit means a failed round never releases existing blocks. + // window. The composite calls this on every leaf after a successful commit; + // the SWA leaf may also call it after an allocation shortage before retrying. + // Non-SWA leaves keep the empty default (no-op). virtual void release_out_of_window(Sequence* /*seq*/) {} virtual void release_out_of_window(Sequence* /*seq*/, KVCacheState& /*kv_state*/) {} diff --git a/xllm/core/framework/block/block_manager_pool.cpp b/xllm/core/framework/block/block_manager_pool.cpp index 17258e8711..750db965e4 100644 --- a/xllm/core/framework/block/block_manager_pool.cpp +++ b/xllm/core/framework/block/block_manager_pool.cpp @@ -45,6 +45,7 @@ BlockManagerPool::BlockManagerPool(const Options& options, int32_t dp_size) .enable_host_offload(options_.enable_host_offload()) .sliding_window_size(options_.sliding_window_size()) .swa_blocks_per_seq(options_.swa_blocks_per_seq()) + .swa_num_blocks(options_.swa_num_blocks()) .max_tokens_per_batch(options_.max_tokens_per_batch()) .manager_types(options_.manager_types()) .compress_ratios(options_.compress_ratios()) diff --git a/xllm/core/framework/block/block_manager_pool.h b/xllm/core/framework/block/block_manager_pool.h index 13d5cdfb23..06e5bc774d 100644 --- a/xllm/core/framework/block/block_manager_pool.h +++ b/xllm/core/framework/block/block_manager_pool.h @@ -55,6 +55,9 @@ class BlockManagerPool : public KVCacheManager { PROPERTY(uint32_t, sliding_window_size) = 0; // Base SWA/cache-state block rows retained per sequence. PROPERTY(uint32_t, swa_blocks_per_seq) = 0; + // Total physical SWA rows. This is computed by the KV cache estimator and + // shared with the device tensor shape to keep both id spaces identical. + PROPERTY(uint32_t, swa_num_blocks) = 0; // Scheduler token budget used to size the shared SWA burst pool. PROPERTY(uint32_t, max_tokens_per_batch) = 0; // For CompositeBlockManager. diff --git a/xllm/core/framework/block/composite_block_manager.cpp b/xllm/core/framework/block/composite_block_manager.cpp index f1be66d840..75038d2e4c 100644 --- a/xllm/core/framework/block/composite_block_manager.cpp +++ b/xllm/core/framework/block/composite_block_manager.cpp @@ -36,11 +36,6 @@ namespace { constexpr uint32_t kManagerTypeBlockManagerImpl = 0; constexpr uint32_t kManagerTypeSlidingWindowBlockManager = 1; -uint32_t ceil_div(uint32_t numerator, uint32_t denominator) { - CHECK_GT(denominator, 0u); - return (numerator + denominator - 1) / denominator; -} - // Whether a leaf of the given BlockType participates in prefix cache under // the current role. On the PREFILL side (instance_is_decode == false) every // cache-bearing leaf participates. On the DECODE side we skip SWA and @@ -228,13 +223,9 @@ CompositeBlockManager::LeafMap build_composite_leaves( CHECK_GT(options.block_size(), 0) << "block_size must be positive"; const uint32_t sliding_window_size = std::max(options.sliding_window_size(), 1u); - const uint32_t max_seqs = std::max(options.max_seqs_per_batch(), 1u); - const uint32_t burst_blocks = - ceil_div(std::max(options.max_tokens_per_batch(), 1u), - static_cast(options.block_size())); - // Slack fits the peak "old blocks not yet released + new tail". - const uint32_t swa_total_blocks = - swa_blocks_per_seq * max_seqs + burst_blocks + max_seqs + 2; + const uint32_t swa_total_blocks = options.swa_num_blocks(); + CHECK_GT(swa_total_blocks, 0u) + << "swa_num_blocks must be provided by the KV cache estimator"; const bool swa_prefix_cache = prefix_cache_on && swa_participates; opts.num_blocks(swa_total_blocks) .swa_blocks_per_seq(swa_blocks_per_seq) @@ -314,6 +305,34 @@ void CompositeBlockManager::cache_full_blocks_for_sequence(Sequence* seq) { return; } KVCacheState& kv = seq->kv_state(); + + size_t cacheable_tokens = + std::min(seq->kv_cache_tokens_num(), seq->tokens().size()); + size_t cache_unit_size = 0; + if (combination_ == LeafCombination::SWA_COMPRESSED) { + BlockManager* c128_leaf = leaf_of(BlockType::C128); + CHECK(c128_leaf != nullptr); + cache_unit_size = c128_leaf->block_size(); + CHECK_GT(cache_unit_size, 0u); + + // A DSV4 prefix is restorable only at a complete C128 boundary. Cap the + // boundary by every participating leaf's allocated logical capacity so no + // partially allocated composite unit can become visible. + for (const auto& [type, entry] : leaves_) { + if (!entry.supports_prefix_cache || type == BlockType::EMBEDDING || + type == BlockType::LINEAR) { + continue; + } + const size_t block_size = entry.leaf->block_size(); + CHECK_GT(block_size, 0u); + CHECK_EQ(cache_unit_size % block_size, 0u) + << "DSV4 cache leaf block size must divide the C128 cache unit"; + cacheable_tokens = + std::min(cacheable_tokens, kv.num_blocks(type) * block_size); + } + cacheable_tokens = (cacheable_tokens / cache_unit_size) * cache_unit_size; + } + for (auto& [type, entry] : leaves_) { // EMBEDDING / LINEAR hold no token cache. KV also participates here so a // Host-restored prefix is published immediately after its HBM destination @@ -333,12 +352,48 @@ void CompositeBlockManager::cache_full_blocks_for_sequence(Sequence* seq) { if (blocks == nullptr || blocks->empty()) { continue; } - const size_t num_full = seq->kv_cache_tokens_num() / block_size; - const size_t cached = kv.num_cached_blocks(type); + const size_t num_full = combination_ == LeafCombination::SWA_COMPRESSED + ? cacheable_tokens / block_size + : seq->kv_cache_tokens_num() / block_size; + size_t cached = kv.num_cached_blocks(type); const size_t end = std::min(num_full, blocks->size()); if (end <= cached) { continue; } + + if (combination_ == LeafCombination::SWA_COMPRESSED && + type == BlockType::SWA) { + const size_t blocks_per_unit = cache_unit_size / block_size; + const size_t blocks_per_window = + static_cast(leaf.options().swa_blocks_per_seq()); + CHECK_GT(blocks_per_unit, 0u); + CHECK_GT(blocks_per_window, 0u); + + // SWA only needs the window immediately preceding each restorable C128 + // checkpoint. Advance the cursor over earlier positions without + // inserting them; LinearStatePrefixCache keeps those positions sparse. + const size_t completed_units = cacheable_tokens / cache_unit_size; + for (size_t unit = cached / blocks_per_unit + 1; unit <= completed_units; + ++unit) { + const size_t unit_end = unit * blocks_per_unit; + const size_t window_begin = + unit_end > blocks_per_window ? unit_end - blocks_per_window : 0; + const size_t publish_begin = std::max(cached, window_begin); + if (unit_end > publish_begin) { + seq->update_block_hashes(static_cast(block_size), + leaf.options().hasher_type()); + leaf.cache(seq->tokens().slice(0, unit_end * block_size), + *blocks, + publish_begin, + seq->mm_data(), + seq->block_hashes()); + } + cached = unit_end; + } + kv.set_num_cached_blocks(type, cached); + continue; + } + // Clamp tokens to `end * block_size`. The leaf's cache() re-derives its own // n_blocks bound from `tokens.size() / block_size_`; without this clamp, // chunked prefill (where seq->tokens() spans the whole prompt but only @@ -368,6 +423,11 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, } KVCacheState& kv_state = seq->kv_state(); + // Publish blocks completed by the previous forward before growing again. + // SlidingWindowBlockManager can then release a cached, slid-out block and + // retry when its first allocation attempt exhausts the SWA pool. + cache_full_blocks_for_sequence(seq); + // Fan out growth. Each leaf returns its newly allocated blocks (or nullopt // on failure). Stage keyed by BlockType; commit only after every leaf // succeeds so a failure rolls back cleanly. diff --git a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp index 2c93267ee0..8f982ad549 100644 --- a/xllm/core/framework/block/hierarchy_block_manager_pool.cpp +++ b/xllm/core/framework/block/hierarchy_block_manager_pool.cpp @@ -743,16 +743,52 @@ void HierarchyBlockManagerPool::prefetch_from_storage( << "Mooncake prefetch admission requires an empty Sequence state."; const int32_t dp_rank = BlockManagerPool::get_dp_rank(sequence); + const auto* composite = static_cast( + block_managers_[dp_rank].get()); auto plan = std::make_shared(); plan->sequence = sequence; plan->host_probes = CompositeBlockManager::probe_prefix_cache( sequence, host_block_managers_[dp_rank]); + const bool is_swa_compressed = + composite->leaf_combination() == + CompositeBlockManager::LeafCombination::SWA_COMPRESSED; + size_t cacheable_tokens = sequence->tokens().size(); + size_t cache_unit_size = 0; + if (is_swa_compressed) { + ProbeResult* c128_probe = find_probe(&plan->host_probes, BlockType::C128); + CHECK(c128_probe != nullptr); + cache_unit_size = c128_probe->block_size; + CHECK_GT(cache_unit_size, 0u); + cacheable_tokens = (cacheable_tokens / cache_unit_size) * cache_unit_size; + + size_t cacheable_units = cacheable_tokens / cache_unit_size; + for (const ProbeResult& probe : plan->host_probes) { + CHECK(probe.leaf != nullptr); + CHECK_GT(probe.block_size, 0u); + CHECK_EQ(cache_unit_size % probe.block_size, 0u); + size_t blocks_per_unit = cache_unit_size / probe.block_size; + if (probe.type == BlockType::SWA) { + const size_t blocks_per_window = + static_cast(probe.leaf->options().swa_blocks_per_seq()); + CHECK_GT(blocks_per_window, 0u); + blocks_per_unit = std::min(blocks_per_unit, blocks_per_window); + } + const size_t available_blocks = + probe.leaf->num_free_blocks() + + probe.leaf->num_blocks_in_prefix_cache(); + cacheable_units = + std::min(cacheable_units, available_blocks / blocks_per_unit); + } + cacheable_tokens = cacheable_units * cache_unit_size; + } + for (size_t probe_index = 0; probe_index < plan->host_probes.size(); ++probe_index) { ProbeResult& probe = plan->host_probes[probe_index]; CHECK(probe.leaf != nullptr); - const size_t full_blocks = sequence->tokens().size() / probe.block_size; + CHECK_GT(probe.block_size, 0u); + const size_t full_blocks = cacheable_tokens / probe.block_size; if (probe.blocks.size() > full_blocks) { trim_blocks_from_back(probe.leaf, &probe.blocks, full_blocks); } @@ -760,6 +796,19 @@ void HierarchyBlockManagerPool::prefetch_from_storage( std::vector missing; for (size_t block_index = 0; block_index < full_blocks; ++block_index) { + if (is_swa_compressed && probe.type == BlockType::SWA) { + CHECK_EQ(cache_unit_size % probe.block_size, 0u); + const size_t blocks_per_unit = cache_unit_size / probe.block_size; + const size_t blocks_per_window = + static_cast(probe.leaf->options().swa_blocks_per_seq()); + CHECK_GT(blocks_per_unit, 0u); + CHECK_GT(blocks_per_window, 0u); + if (blocks_per_window < blocks_per_unit && + block_index % blocks_per_unit < + blocks_per_unit - blocks_per_window) { + continue; + } + } if (!probe.blocks[block_index].is_valid()) { missing.emplace_back(block_index); } diff --git a/xllm/core/framework/block/sliding_window_block_manager.cpp b/xllm/core/framework/block/sliding_window_block_manager.cpp index c5adb81d7a..d6823619c1 100644 --- a/xllm/core/framework/block/sliding_window_block_manager.cpp +++ b/xllm/core/framework/block/sliding_window_block_manager.cpp @@ -16,6 +16,7 @@ limitations under the License. #include "sliding_window_block_manager.h" #include +#include #include "framework/prefix_cache/prefix_cache.h" @@ -34,6 +35,67 @@ SlidingWindowBlockManager::SlidingWindowBlockManager(const Options& options) } } +std::optional> +SlidingWindowBlockManager::allocate_for_sequence(Sequence* seq, + size_t num_tokens) { + if (seq == nullptr) { + return std::nullopt; + } + return allocate_for_sequence(seq, seq->kv_state(), num_tokens); +} + +std::optional> +SlidingWindowBlockManager::allocate_for_sequence(Sequence* seq, + KVCacheState& kv_state, + size_t num_tokens) { + if (seq == nullptr) { + return std::nullopt; + } + if (!options_.instance_is_decode() || kv_state.num_blocks(block_type()) > 0) { + std::optional> blocks = + BlockManagerImpl::allocate_for_sequence(seq, kv_state, num_tokens); + if (blocks.has_value()) { + return blocks; + } + + // A block completed by the previous forward may now be outside the active + // window. Release slid-out sequence references, then retry: uncached blocks + // return directly to the free list, while cached checkpoint-window blocks + // can be evicted and reused. Only mutate the sequence when those + // reclaimable blocks make the retry large enough to succeed; otherwise a + // failed composite round must preserve the existing SWA state. + const size_t block_size = options_.block_size(); + CHECK_GT(block_size, 0u); + const size_t held = kv_state.num_blocks(block_type()); + const size_t num_blocks_needed = (num_tokens + block_size - 1) / block_size; + CHECK_GT(num_blocks_needed, held); + const size_t num_additional = num_blocks_needed - held; + const size_t reclaimable = num_reclaimable_out_of_window_blocks( + kv_state, kv_state.kv_cache_tokens_num()); + if (num_free_blocks() + reclaimable < num_additional) { + return std::nullopt; + } + release_out_of_window(seq, kv_state); + return BlockManagerImpl::allocate_for_sequence(seq, kv_state, num_tokens); + } + + const size_t block_size = options_.block_size(); + CHECK_GT(block_size, 0u); + const size_t logical_blocks = (num_tokens + block_size - 1) / block_size; + const size_t active_blocks = std::min( + logical_blocks, static_cast(options_.swa_blocks_per_seq())); + std::vector live_blocks = allocate(active_blocks); + if (live_blocks.size() != active_blocks) { + return std::nullopt; + } + + std::vector sparse_blocks(logical_blocks - active_blocks); + sparse_blocks.insert(sparse_blocks.end(), + std::make_move_iterator(live_blocks.begin()), + std::make_move_iterator(live_blocks.end())); + return sparse_blocks; +} + void SlidingWindowBlockManager::release_out_of_window(Sequence* seq) { if (seq == nullptr) { return; @@ -53,21 +115,8 @@ void SlidingWindowBlockManager::release_out_of_window(Sequence* seq, return; } std::vector& swa_blocks = *kv_state.mutable_blocks(block_type()); - const size_t block_size = options_.block_size(); - if (block_size == 0 || swa_blocks.empty()) { - return; - } - const size_t num_spec_tokens = - static_cast(options_.num_speculative_tokens()); - const size_t sliding_window_tokens = - std::max(options_.sliding_window_size(), 1); - if (cached_tokens < (sliding_window_tokens + num_spec_tokens)) { - return; - } - const size_t skipped_tokens = - cached_tokens - sliding_window_tokens - num_spec_tokens + 1; - const size_t skipped_blocks = skipped_tokens / block_size; - const size_t release_blocks = std::min(skipped_blocks, swa_blocks.size()); + const size_t release_blocks = + num_out_of_window_blocks(kv_state, cached_tokens); if (release_blocks == 0) { return; } @@ -86,6 +135,45 @@ void SlidingWindowBlockManager::release_out_of_window(Sequence* seq, } } +size_t SlidingWindowBlockManager::num_out_of_window_blocks( + const KVCacheState& kv_state, + size_t cached_tokens) const { + const size_t block_size = options_.block_size(); + const size_t held = kv_state.num_blocks(block_type()); + if (block_size == 0 || held == 0) { + return 0; + } + const size_t num_spec_tokens = + static_cast(options_.num_speculative_tokens()); + const size_t sliding_window_tokens = + std::max(options_.sliding_window_size(), 1); + if (cached_tokens < sliding_window_tokens + num_spec_tokens) { + return 0; + } + const size_t skipped_tokens = + cached_tokens - sliding_window_tokens - num_spec_tokens + 1; + const size_t skipped_blocks = skipped_tokens / block_size; + return std::min(skipped_blocks, held); +} + +size_t SlidingWindowBlockManager::num_reclaimable_out_of_window_blocks( + const KVCacheState& kv_state, + size_t cached_tokens) const { + const Slice swa_blocks = kv_state.blocks(block_type()); + const size_t release_blocks = + num_out_of_window_blocks(kv_state, cached_tokens); + const uint32_t max_reclaimable_ref_count = + options_.enable_prefix_cache() ? 2u : 1u; + size_t reclaimable = 0; + for (size_t i = 0; i < release_blocks; ++i) { + if (swa_blocks[i].is_valid() && + swa_blocks[i].ref_count() <= max_reclaimable_ref_count) { + ++reclaimable; + } + } + return reclaimable; +} + std::vector SlidingWindowBlockManager::allocate_shared( const Slice& token_ids, const Slice& /*existed_shared_blocks*/, diff --git a/xllm/core/framework/block/sliding_window_block_manager.h b/xllm/core/framework/block/sliding_window_block_manager.h index ed9d2c5631..5cafd0b862 100644 --- a/xllm/core/framework/block/sliding_window_block_manager.h +++ b/xllm/core/framework/block/sliding_window_block_manager.h @@ -31,6 +31,14 @@ class SlidingWindowBlockManager : public BlockManagerImpl { explicit SlidingWindowBlockManager(const Options& options); ~SlidingWindowBlockManager() override = default; + std::optional> allocate_for_sequence( + Sequence* seq, + size_t num_tokens) override; + std::optional> allocate_for_sequence( + Sequence* seq, + KVCacheState& kv_state, + size_t num_tokens) override; + // Deallocate leading blocks that have slid out of the window; leaves // invalid placeholders in their slots. Called by the composite after a // successful allocate commit. @@ -49,6 +57,10 @@ class SlidingWindowBlockManager : public BlockManagerImpl { uint32_t swa_blocks_per_seq() const { return options_.swa_blocks_per_seq(); } private: + size_t num_out_of_window_blocks(const KVCacheState& kv_state, + size_t cached_tokens) const; + size_t num_reclaimable_out_of_window_blocks(const KVCacheState& kv_state, + size_t cached_tokens) const; void release_out_of_window(Sequence* seq, KVCacheState& kv_state, size_t cached_tokens); diff --git a/xllm/core/framework/config/disagg_pd_config.cpp b/xllm/core/framework/config/disagg_pd_config.cpp index d2eace042d..5ea7d20ebb 100644 --- a/xllm/core/framework/config/disagg_pd_config.cpp +++ b/xllm/core/framework/config/disagg_pd_config.cpp @@ -51,13 +51,6 @@ DEFINE_bool(kv_push_dst_rotate, "KV-split rank to spread incast across D workers."); namespace xllm { -namespace { - -bool supports_prefix_cache(const std::string& instance_role) { - return instance_role == "PREFILL" || instance_role == "MIX"; -} - -} // namespace void DisaggPDConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_disagg_pd); @@ -122,13 +115,6 @@ void DisaggPDConfig::normalize_mlu(KVCacheConfig& kv_cache_config, << "forcing enable_schedule_overlap=false."; scheduler_config.enable_schedule_overlap(false); } - if (kv_cache_config.enable_prefix_cache() && - !supports_prefix_cache(instance_role())) { - LOG(WARNING) << "MLU disaggregated PD role " << instance_role() - << " does not support prefix cache; " - << "forcing enable_prefix_cache=false."; - kv_cache_config.enable_prefix_cache(false); - } if (enable_pd_ooc()) { LOG(WARNING) << "MLU disaggregated PD does not support pd_ooc; " << "forcing enable_pd_ooc=false."; diff --git a/xllm/core/framework/kv_cache/kv_cache_estimation.cpp b/xllm/core/framework/kv_cache/kv_cache_estimation.cpp index c2b6456fa9..25624944d9 100644 --- a/xllm/core/framework/kv_cache/kv_cache_estimation.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_estimation.cpp @@ -301,12 +301,21 @@ int64_t calculate_linear_state_blocks(int64_t cache_size_in_bytes, return std::min(auto_blocks, max_blocks); } -Dsv4KVCacheEstimateCost estimate_dsv4_kv_cache_cost( +constexpr int64_t kDsv4SwaPaddingBlocks = 1; + +bool is_dsv4_prefill_or_mix_role(InstanceRole instance_role) { + return instance_role == InstanceRole::DEFAULT || + instance_role == InstanceRole::PREFILL || + instance_role == InstanceRole::MIX; +} + +int64_t calculate_dsv4_minimum_swa_count( const ModelArgs& model_args, const KVCacheEstimateOptions& options) { const int64_t max_seqs = std::max(options.max_seqs_per_batch, static_cast(1)); const int64_t block_size = options.block_size; + CHECK_GT(block_size, 0) << "DSV4 block_size must be positive"; const int64_t semantic_window = std::max(model_args.window_size(), 1); const int64_t max_model_len = model_args.max_seq_len(); const int64_t window_size = @@ -314,9 +323,57 @@ Dsv4KVCacheEstimateCost estimate_dsv4_kv_cache_cost( : semantic_window; const int64_t swa_blocks_per_seq = get_swa_blocks_per_seq(window_size, block_size); - const int64_t burst_blocks = util::ceil_div( - std::max(options.max_tokens_per_batch, static_cast(1)), - block_size); + + if (is_dsv4_prefill_or_mix_role(options.instance_role)) { + const bool use_chunk_limit = options.enable_chunked_prefill && + options.max_tokens_per_chunk_for_prefill > 0; + const int64_t chunk_tokens = + use_chunk_limit + ? options.max_tokens_per_chunk_for_prefill + : std::max(options.max_tokens_per_batch, static_cast(1)); + const int64_t chunk_blocks = util::ceil_div(chunk_tokens, block_size); + return max_seqs * (swa_blocks_per_seq + chunk_blocks + 1) + + kDsv4SwaPaddingBlocks; + } + CHECK(options.instance_role == InstanceRole::DECODE) + << "unsupported DSV4 SWA estimation instance_role=" + << static_cast(options.instance_role); + const int64_t speculative_tokens = + std::max(options.num_speculative_tokens, static_cast(0)); + const int64_t speculative_reserve_tokens = + speculative_tokens * (options.enable_schedule_overlap ? 2 : 1); + const int64_t decode_window_blocks = + util::ceil_div(window_size + speculative_reserve_tokens, block_size); + return max_seqs * decode_window_blocks * 2 + kDsv4SwaPaddingBlocks; +} + +int64_t dsv4_common_unit_bytes(const Dsv4KVCacheEstimateCost& cache_cost, + int64_t common_blocks_per_unit) { + CHECK_GT(cache_cost.manager_blocks_per_unit, 0); + CHECK_EQ(common_blocks_per_unit % cache_cost.manager_blocks_per_unit, 0); + const int64_t compressed_units = + common_blocks_per_unit / cache_cost.manager_blocks_per_unit; + return compressed_units * cache_cost.token_unit_bytes; +} + +void set_dsv4_compressed_counts(const Dsv4KVCacheEstimateCost& cache_cost, + int64_t token_unit_count, + KVCacheCapacity* kv_cache_cap) { + CHECK(kv_cache_cap != nullptr); + if (cache_cost.n_c4_layers > 0 && cache_cost.n_c128_layers > 0) { + kv_cache_cap->c128_count(token_unit_count); + kv_cache_cap->c4_count(32 * token_unit_count); + } else if (cache_cost.n_c4_layers > 0) { + kv_cache_cap->c4_count(token_unit_count); + } else if (cache_cost.n_c128_layers > 0) { + kv_cache_cap->c128_count(token_unit_count); + } +} + +Dsv4KVCacheEstimateCost estimate_dsv4_kv_cache_cost( + const ModelArgs& model_args, + const KVCacheEstimateOptions& options) { + const int64_t block_size = options.block_size; const int64_t head_dim = model_args.head_dim(); const int64_t index_head_dim = std::max(model_args.index_head_dim(), 1); @@ -326,8 +383,17 @@ Dsv4KVCacheEstimateCost estimate_dsv4_kv_cache_cost( static_cast(torch::elementSize(options.dtype)); Dsv4KVCacheEstimateCost cache_cost; - cache_cost.swa_count = - swa_blocks_per_seq * max_seqs + burst_blocks + max_seqs + 2; + cache_cost.swa_count = calculate_dsv4_minimum_swa_count(model_args, options); + LOG(INFO) << "DSV4 minimum SWA block request: model_type=" + << model_args.model_type() + << ", is_draft_engine=" << options.is_draft_engine + << ", enable_disagg_pd=" << options.enable_disagg_pd + << ", instance_role=" << static_cast(options.instance_role) + << ", max_seqs_per_batch=" << options.max_seqs_per_batch + << ", max_tokens_per_batch=" << options.max_tokens_per_batch + << ", max_tokens_per_chunk_for_prefill=" + << options.max_tokens_per_chunk_for_prefill + << ", swa_count=" << cache_cost.swa_count; for (int64_t i = 0; i < model_args.n_layers(); ++i) { const int32_t ratio = i < static_cast(compress_ratios.size()) ? compress_ratios[static_cast(i)] @@ -341,21 +407,21 @@ Dsv4KVCacheEstimateCost estimate_dsv4_kv_cache_cost( const int64_t n_c1_layers = model_args.n_layers() - cache_cost.n_c4_layers - cache_cost.n_c128_layers; - const int64_t swa_bytes_per_c1_layer = - cache_cost.swa_count * block_size * head_dim * dtype_size; - const int64_t swa_bytes_per_c4_layer = - cache_cost.swa_count * - (block_size * head_dim * dtype_size + - block_size * (2 * head_dim * float32_size) * 2 + - block_size * (2 * index_head_dim * float32_size) * 2); - const int64_t swa_bytes_per_c128_layer = - cache_cost.swa_count * (block_size * head_dim * dtype_size + - block_size * head_dim * float32_size * 2); - + const int64_t swa_bytes_per_c1_block = block_size * head_dim * dtype_size; + const int64_t swa_bytes_per_c4_block = + block_size * head_dim * dtype_size + + block_size * (2 * head_dim * float32_size) * 2 + + block_size * (2 * index_head_dim * float32_size) * 2; + const int64_t swa_bytes_per_c128_block = + block_size * head_dim * dtype_size + + block_size * head_dim * float32_size * 2; + + cache_cost.swa_bytes_per_block = + n_c1_layers * swa_bytes_per_c1_block + + cache_cost.n_c4_layers * swa_bytes_per_c4_block + + cache_cost.n_c128_layers * swa_bytes_per_c128_block; cache_cost.constant_swa_bytes = - n_c1_layers * swa_bytes_per_c1_layer + - cache_cost.n_c4_layers * swa_bytes_per_c4_layer + - cache_cost.n_c128_layers * swa_bytes_per_c128_layer; + cache_cost.swa_count * cache_cost.swa_bytes_per_block; const DeepSeekV4CachePolicy cache_policy = get_dsv4_cache_policy(options.dtype); @@ -389,10 +455,12 @@ void init_dsv4_counts(const ModelArgs& model_args, CHECK(kv_cache_cap != nullptr); const Dsv4KVCacheEstimateCost cache_cost = estimate_dsv4_kv_cache_cost(model_args, options); - int64_t token_mem = std::max( - static_cast(0), - kv_cache_cap->cache_size_in_bytes() - cache_cost.constant_swa_bytes); - + CHECK_GE(kv_cache_cap->cache_size_in_bytes(), cache_cost.constant_swa_bytes) + << "no memory for the minimum DSV4 SWA cache, required=" + << readable_size(cache_cost.constant_swa_bytes) + << ", available=" << readable_size(kv_cache_cap->cache_size_in_bytes()); + int64_t token_mem = + kv_cache_cap->cache_size_in_bytes() - cache_cost.constant_swa_bytes; if (options.draft_model_args != nullptr) { CHECK(options.draft_options != nullptr) << "DSV4 draft options must be provided with draft model args"; @@ -404,24 +472,35 @@ void init_dsv4_counts(const ModelArgs& model_args, *options.draft_model_args, *options.draft_options); const int64_t constant_bytes = cache_cost.constant_swa_bytes + draft_cost.constant_swa_bytes; - CHECK_GT(kv_cache_cap->cache_size_in_bytes(), constant_bytes) + CHECK_GE(kv_cache_cap->cache_size_in_bytes(), constant_bytes) << "no memory left for speculative target/draft fixed kv cache " "allocation"; - const int64_t token_unit_bytes = - cache_cost.token_unit_bytes + draft_cost.token_unit_bytes; - CHECK_GT(token_unit_bytes, 0) - << "mtp target and draft token unit bytes must be positive"; + const int64_t common_blocks_per_unit = std::max( + cache_cost.manager_blocks_per_unit, draft_cost.manager_blocks_per_unit); + const int64_t target_common_unit_bytes = + dsv4_common_unit_bytes(cache_cost, common_blocks_per_unit); + const int64_t draft_common_unit_bytes = + dsv4_common_unit_bytes(draft_cost, common_blocks_per_unit); + const int64_t combined_unit_bytes = + target_common_unit_bytes + draft_common_unit_bytes; + const int64_t remaining_bytes = + kv_cache_cap->cache_size_in_bytes() - constant_bytes; + if (combined_unit_bytes > 0) { + CHECK_GE(remaining_bytes, combined_unit_bytes) + << "minimum DSV4 target/draft SWA caches leave insufficient memory " + "for one compressed cache unit, swa_required=" + << readable_size(constant_bytes) + << ", compressed_unit_required=" << readable_size(combined_unit_bytes) + << ", available=" + << readable_size(kv_cache_cap->cache_size_in_bytes()); + } const int64_t token_unit_count = - (kv_cache_cap->cache_size_in_bytes() - constant_bytes) / - token_unit_bytes; - CHECK_GT(token_unit_count, 0) - << "no memory left for speculative target/draft kv cache token " - "blocks"; + combined_unit_bytes > 0 ? remaining_bytes / combined_unit_bytes : 0; const int64_t adjusted_cache_size_in_bytes = cache_cost.constant_swa_bytes + - token_unit_count * cache_cost.token_unit_bytes; + token_unit_count * target_common_unit_bytes; CHECK_GT(adjusted_cache_size_in_bytes, 0) << "no memory left for speculative target/draft kv cache allocation"; LOG(INFO) << "speculative kv cache capacity adjusted from " @@ -429,32 +508,34 @@ void init_dsv4_counts(const ModelArgs& model_args, << readable_size(adjusted_cache_size_in_bytes) << ", target_constant_bytes=" << cache_cost.constant_swa_bytes << ", draft_constant_bytes=" << draft_cost.constant_swa_bytes - << ", target_token_unit_bytes=" << cache_cost.token_unit_bytes - << ", draft_token_unit_bytes=" << draft_cost.token_unit_bytes + << ", target_common_unit_bytes=" << target_common_unit_bytes + << ", draft_common_unit_bytes=" << draft_common_unit_bytes + << ", common_blocks_per_unit=" << common_blocks_per_unit << ", token_unit_count=" << token_unit_count; kv_cache_cap->cache_size_in_bytes(adjusted_cache_size_in_bytes); - token_mem = token_unit_count * cache_cost.token_unit_bytes; + token_mem = token_unit_count * target_common_unit_bytes; } else { CHECK(options.draft_options == nullptr) << "DSV4 draft options require draft model args"; + if (cache_cost.token_unit_bytes > 0) { + CHECK_GE(token_mem, cache_cost.token_unit_bytes) + << "minimum DSV4 SWA cache leaves insufficient memory for one " + "compressed cache unit, swa_required=" + << readable_size(cache_cost.constant_swa_bytes) + << ", compressed_unit_required=" + << readable_size(cache_cost.token_unit_bytes) << ", available=" + << readable_size(kv_cache_cap->cache_size_in_bytes()); + } } kv_cache_cap->swa_count(cache_cost.swa_count); kv_cache_cap->c4_count(0); kv_cache_cap->c128_count(0); - if (cache_cost.n_c4_layers > 0 && cache_cost.n_c128_layers > 0) { - if (cache_cost.token_unit_bytes > 0 && token_mem > 0) { - kv_cache_cap->c128_count(token_mem / cache_cost.token_unit_bytes); - kv_cache_cap->c4_count(32 * kv_cache_cap->c128_count()); - } - } else if (cache_cost.n_c4_layers > 0) { - if (cache_cost.token_unit_bytes > 0 && token_mem > 0) { - kv_cache_cap->c4_count(token_mem / cache_cost.token_unit_bytes); - } - } else if (cache_cost.n_c128_layers > 0) { - if (cache_cost.token_unit_bytes > 0 && token_mem > 0) { - kv_cache_cap->c128_count(token_mem / cache_cost.token_unit_bytes); - } + // Keep SWA at the operational minimum calculated above. Prefix-cache entries + // share this pool; any remaining memory is reserved for compressed history. + if (cache_cost.token_unit_bytes > 0 && token_mem > 0) { + const int64_t token_unit_count = token_mem / cache_cost.token_unit_bytes; + set_dsv4_compressed_counts(cache_cost, token_unit_count, kv_cache_cap); } CHECK_GT(kv_cache_cap->swa_count(), 0) << "DSV4 swa_count must be > 0"; diff --git a/xllm/core/framework/kv_cache/kv_cache_estimation.h b/xllm/core/framework/kv_cache/kv_cache_estimation.h index daf00bac77..5cd22df4d1 100644 --- a/xllm/core/framework/kv_cache/kv_cache_estimation.h +++ b/xllm/core/framework/kv_cache/kv_cache_estimation.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include "common/types.h" #include "framework/kv_cache/kv_cache_capacity.h" #include "framework/kv_cache/layerwise_split_layout.h" @@ -41,10 +42,15 @@ struct KVCacheEstimateOptions { int64_t max_seqs_per_batch = 0; int64_t num_speculative_tokens = 0; int64_t max_tokens_per_batch = 0; + int64_t max_tokens_per_chunk_for_prefill = 0; int64_t max_linear_state_cache_slots = 0; bool is_draft_engine = false; bool enable_prefix_cache = false; int32_t layerwise_split_size = 1; + bool enable_chunked_prefill = true; + bool enable_schedule_overlap = true; + bool enable_disagg_pd = false; + InstanceRole instance_role = InstanceRole::DEFAULT; const ModelArgs* draft_model_args = nullptr; const KVCacheEstimateOptions* draft_options = nullptr; }; @@ -53,6 +59,7 @@ struct Dsv4KVCacheEstimateCost { int64_t swa_count = 0; int64_t n_c4_layers = 0; int64_t n_c128_layers = 0; + int64_t swa_bytes_per_block = 0; int64_t constant_swa_bytes = 0; int64_t token_unit_bytes = 0; int64_t manager_blocks_per_unit = 1; diff --git a/xllm/core/framework/request/sequence_kv_state.h b/xllm/core/framework/request/sequence_kv_state.h index 6e7be644a7..22f0547053 100644 --- a/xllm/core/framework/request/sequence_kv_state.h +++ b/xllm/core/framework/request/sequence_kv_state.h @@ -108,8 +108,9 @@ class KVCacheState { // inserted. For sparse SWA this is a logical position and may span invalid // placeholders. Grows monotonically: // - Admission mount: set to that type's retained probe-vector length. - // - Pre-grow hook: after inserting a run [cursor, end), advance cursor to - // `end`. + // - Pre-grow hook: after processing positions through `end`, advance the + // cursor to `end`. Sparse SWA may intentionally insert only the trailing + // window of each cache unit while skipping earlier positions. // - reset(): cleared alongside the rest of the sequence's cache state. size_t num_cached_blocks(BlockType type) const; // Per-type cursor table. Callers that need a stable snapshot must copy it. diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index e757e41f3a..acabe8cdb1 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -446,4 +446,5 @@ std::tuple apply_npu_mega_moe( int64_t dispatch_quant_out_dtype = 0, int64_t topo_type = 0, int64_t rank_num_per_server = 2); + } // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_ops_library.cpp b/xllm/core/kernels/npu/npu_ops_library.cpp index acd5536383..936fb3a36d 100644 --- a/xllm/core/kernels/npu/npu_ops_library.cpp +++ b/xllm/core/kernels/npu/npu_ops_library.cpp @@ -32,6 +32,7 @@ limitations under the License. #include "kernels/npu/xllm_ops/xllm_ops_api.h" #include "npu_ops_api.h" +#include "triton_npu/torch_api/triton_ops_api.h" namespace xllm { @@ -43,6 +44,159 @@ torch::Tensor rms_norm_npu(const torch::Tensor& input, return xllm::kernel::npu::rms_norm(input, weight, eps, "rmsnorm"); } +torch::Tensor rms_norm_gated_npu(const torch::Tensor& input, + const torch::Tensor& gate, + const torch::Tensor& weight, + double eps) { + return xllm::kernel::npu::layer_norm_fwd_aclnn(input, + weight, + /*bias=*/torch::Tensor(), + eps, + /*z=*/gate, + /*group_size=*/input.size(-1), + /*norm_before_gate=*/true, + /*is_rms_norm=*/true); +} + +torch::Tensor l2_norm_npu(torch::Tensor input, double eps) { + return xllm::kernel::npu::npu_l2norm_last_dim(input, eps); +} + +torch::Tensor causal_conv1d_prefill_npu(torch::Tensor x, + torch::Tensor weight, + torch::Tensor conv_state, + torch::Tensor state_indices, + torch::Tensor has_initial_state, + torch::Tensor query_start_loc) { + // Python layer stores weight as [dim, kernel_width]; CANN expects + // [kernel_width, dim]. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + + // Convert device tensors to host vectors for IntArrayRef parameters. + auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous(); + auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous(); + auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous(); + + std::vector qsl_vec(qsl_cpu.data_ptr(), + qsl_cpu.data_ptr() + qsl_cpu.numel()); + std::vector si_vec(si_cpu.data_ptr(), + si_cpu.data_ptr() + si_cpu.numel()); + std::vector ism_vec(ism_cpu.data_ptr(), + ism_cpu.data_ptr() + ism_cpu.numel()); + + constexpr int64_t kActivationSilu = 1; + constexpr int64_t kPadSlotId = -1; + constexpr int64_t kRunModeForward = 0; + + return xllm::kernel::npu::causal_conv1d( + x, + weight, + conv_state, + /*bias_opt=*/std::nullopt, + torch::IntArrayRef(qsl_vec), + torch::IntArrayRef(si_vec), + torch::IntArrayRef(ism_vec), + /*num_accepted_tokens_opt=*/torch::IntArrayRef{}, + kActivationSilu, + kPadSlotId, + kRunModeForward); +} + +std::tuple +causal_conv1d_qkv_prefill_npu(torch::Tensor x, + torch::Tensor weight, + torch::Tensor conv_state, + torch::Tensor state_indices, + torch::Tensor has_initial_state, + torch::Tensor query_start_loc, + int64_t num_qk_heads, + int64_t num_v_heads, + int64_t head_k_dim, + int64_t head_v_dim) { + // Python layer stores weight as [dim, kernel_width]; CANN expects + // [kernel_width, dim]. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + + auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous(); + auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous(); + auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous(); + + std::vector qsl_vec(qsl_cpu.data_ptr(), + qsl_cpu.data_ptr() + qsl_cpu.numel()); + std::vector si_vec(si_cpu.data_ptr(), + si_cpu.data_ptr() + si_cpu.numel()); + std::vector ism_vec(ism_cpu.data_ptr(), + ism_cpu.data_ptr() + ism_cpu.numel()); + + return xllm::kernel::npu::causal_conv1d_qkv(x, + weight, + conv_state, + torch::IntArrayRef(qsl_vec), + torch::IntArrayRef(si_vec), + torch::IntArrayRef(ism_vec), + num_qk_heads, + num_v_heads, + head_k_dim, + head_v_dim); +} + +std::tuple chunk_gated_delta_rule_npu( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor initial_state, + torch::Tensor cu_seqlens) { + return xllm::kernel::npu::npu_mega_chunk_gdn( + q, + k, + v, + g, + beta, + /*scale=*/std::nullopt, + /*initial_state=*/initial_state, + /*output_final_state=*/true, + /*cu_seqlens=*/cu_seqlens, + /*q_seq_lens=*/{}, + /*use_qk_l2norm_in_kernel=*/true); +} + +torch::Tensor fused_sigmoid_gating_delta_rule_decode_npu( + torch::Tensor a_log, + torch::Tensor a, + torch::Tensor dt_bias, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor b, + torch::Tensor ssm_state, + torch::Tensor state_indices, + torch::Tensor cu_seqlens, + double scale) { + auto a_log_f32 = a_log.to(torch::kFloat32); + auto dt_bias_f32 = dt_bias.to(torch::kFloat32); + return xllm::kernel::npu::npu_fused_sigmoid_gating_delta_rule_update( + a_log_f32, + a, + dt_bias_f32, + q, + k, + v, + b, + ssm_state, + state_indices, + cu_seqlens, + /*scale=*/static_cast(scale), + /*use_qk_l2norm_in_kernel=*/true, + /*softplus_beta=*/1.0f, + /*softplus_threshold=*/20.0f); +} + std::tuple fused_add_rms_norm_npu( torch::Tensor& input, torch::Tensor& residual, @@ -293,10 +447,46 @@ void ensure_xllm_ops_registered() { // compiled only under USE_NPU (mutually exclusive with USE_CUDA). TORCH_LIBRARY(xllm_ops, m) { m.def("rms_norm(Tensor input, Tensor weight, float eps) -> Tensor"); + m.def( + "rms_norm_gated(Tensor input, Tensor gate, Tensor weight, float eps) -> " + "Tensor"); + m.def("l2_norm(Tensor input, float eps) -> Tensor"); + m.def( + "chunk_gated_delta_rule(Tensor q, Tensor k, Tensor v, Tensor g, " + "Tensor beta, Tensor initial_state, Tensor cu_seqlens) -> " + "(Tensor, Tensor)"); + m.def( + "causal_conv1d_prefill(Tensor x, Tensor weight, Tensor(a!) conv_state, " + "Tensor state_indices, Tensor has_initial_state, " + "Tensor query_start_loc) -> Tensor"); + m.def( + "causal_conv1d_qkv_prefill(Tensor x, Tensor weight, " + "Tensor(a!) conv_state, Tensor state_indices, " + "Tensor has_initial_state, Tensor query_start_loc, " + "int num_qk_heads, int num_v_heads, " + "int head_k_dim, int head_v_dim) -> (Tensor, Tensor, Tensor)"); + m.def( + "fused_sigmoid_gating_delta_rule_decode(Tensor a_log, Tensor a, " + "Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor b, " + "Tensor(a!) ssm_state, Tensor state_indices, Tensor cu_seqlens, " + "float scale) -> Tensor"); m.def( "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor " "weight, " "float eps) -> (Tensor, Tensor)"); + // Fused RMSNorm + dynamic per-token int8 quant (W8A8 query preprocess). + // Returns (qr_int8, qr_pertoken_scale) matching C++ rms_norm_dynamic_quant + // (npu_ops_api.h:122), used by the DSV4 indexer build_query path. + m.def( + "rms_norm_dynamic_quant(Tensor input, Tensor weight, float eps) -> " + "(Tensor, Tensor)"); + // In-place partial rotary embedding (interleaved). x is 4D [B,N,S,D], r1/r2 + // are cos/sin [B,1,1,rope_head_dim]; partial_slice=[rope_start, + // rope_head_dim]. Mirrors C++ apply_partial_rope + // (deepseek_sparse_attention.cpp:151) used by the DSV4 indexer build_query. + m.def( + "npu_inplace_partial_rotary_mul(Tensor(a!) x, Tensor r1, Tensor r2, " + "str rotary_mode, int[] partial_slice) -> ()"); m.def("silu_and_mul(Tensor input) -> Tensor"); m.def( "inplace_partial_rotary_mul(Tensor(a!) input, Tensor cosine, Tensor " @@ -394,11 +584,75 @@ TORCH_LIBRARY(xllm_ops, m) { "cache_mode, int quant_mode, bool do_rms_norm, int " "wdkv_split_count, bool q_down_out_flag) -> (Tensor, Tensor(a!), " "Tensor, Tensor(b!), Tensor)"); + // ---- DeepSeek-V4 DSA kernels ---- + // MoE hash routing gate (returns routed output, expert_idx, token_unpermute). + m.def( + "moe_gating_top_k_hash(Tensor x, int k, Tensor? bias, Tensor? input_ids, " + "Tensor? tid2eid, int k_group, int group_count, float " + "routed_scaling_factor, " + "float eps, int group_select_mode, int renorm, int norm_type, bool " + "out_flag) -> (Tensor, Tensor, Tensor)"); + // Dequant + SwiGLU + quant (fused, replaces manual dequant loop). + m.def( + "dequant_swiglu_quant(Tensor x, Tensor? weight_scale, Tensor? " + "activation_scale, Tensor? bias, Tensor? quant_scale, Tensor? " + "quant_offset, Tensor? group_index, bool activate_left, int quant_mode, " + "int swiglu_mode, float clamp_limit, float glu_alpha, float glu_bias) " + "-> (Tensor, Tensor)"); + // HyperConnection pre/post (hc_pre returns attn_input, post, comb). + m.def( + "hc_pre(Tensor x, Tensor hc_fn, Tensor hc_scale, Tensor hc_base, " + "int hc_mult, int hc_sinkhorn_iters, float norm_eps, float hc_eps) " + "-> (Tensor, Tensor, Tensor)"); + m.def( + "hc_post(Tensor x, Tensor residual, Tensor post, Tensor comb) -> " + "Tensor"); + // Compressor: NSA-style KV pooling. kv_state/score_state are in-place (Ref). + // Returns (cmp_kv, wkv_proj, softmax_res, norm_x, norm_rstd). + m.def( + "compressor(Tensor x, Tensor wkv, Tensor wgate, Tensor(a!) kv_state, " + "Tensor(b!) score_state, Tensor ape, Tensor norm_weight, Tensor " + "rope_sin, Tensor rope_cos, Tensor? kv_block_table, Tensor? " + "score_block_table, Tensor? cu_seqlens, Tensor? seqused, Tensor? " + "start_pos, int rope_head_dim, int cmp_ratio, int coff, float " + "norm_eps, int rotary_mode, bool enable_grad) -> (Tensor, Tensor, " + "Tensor, Tensor, Tensor)"); + // Two-stage sparse attention over original + compressed KV. + m.def( + "sparse_attn_sharedkv(Tensor q, Tensor? ori_kv, Tensor? cmp_kv, " + "Tensor? ori_sparse_indices, Tensor? cmp_sparse_indices, Tensor? " + "ori_block_table, Tensor? cmp_block_table, Tensor? cu_seqlens_q, " + "Tensor? cu_seqlens_ori_kv, Tensor? cu_seqlens_cmp_kv, Tensor? " + "seqused_q, Tensor? seqused_kv, Tensor? sinks, Tensor? metadata, " + "float softmax_scale, int cmp_ratio, int ori_mask_mode, int " + "cmp_mask_mode, int ori_win_left, int ori_win_right, str layout_q, " + "str layout_kv, bool return_softmax_lse) -> (Tensor, Tensor)"); + // AICPU tiling metadata builder for sparse_attn_sharedkv. + m.def( + "sparse_attn_sharedkv_metadata(int num_heads_q, int num_heads_kv, int " + "head_dim, Tensor? cu_seqlens_q, Tensor? cu_seqlens_ori_kv, Tensor? " + "cu_seqlens_cmp_kv, Tensor? seqused_q, Tensor? seqused_kv, int " + "batch_size, int max_seqlen_q, int max_seqlen_kv, int ori_topk, int " + "cmp_topk, int cmp_ratio, int ori_mask_mode, int cmp_mask_mode, int " + "ori_win_left, int ori_win_right, str layout_q, str layout_kv, bool " + "has_ori_kv, bool has_cmp_kv) -> Tensor"); } TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("rms_norm", TORCH_FN(xllm::rms_norm_npu)); + m.impl("rms_norm_gated", TORCH_FN(xllm::rms_norm_gated_npu)); + m.impl("l2_norm", TORCH_FN(xllm::l2_norm_npu)); + m.impl("chunk_gated_delta_rule", TORCH_FN(xllm::chunk_gated_delta_rule_npu)); + m.impl("causal_conv1d_prefill", TORCH_FN(xllm::causal_conv1d_prefill_npu)); + m.impl("causal_conv1d_qkv_prefill", + TORCH_FN(xllm::causal_conv1d_qkv_prefill_npu)); + m.impl("fused_sigmoid_gating_delta_rule_decode", + TORCH_FN(xllm::fused_sigmoid_gating_delta_rule_decode_npu)); m.impl("fused_add_rms_norm", TORCH_FN(xllm::fused_add_rms_norm_npu)); + m.impl("rms_norm_dynamic_quant", + TORCH_FN(xllm::kernel::npu::rms_norm_dynamic_quant)); + m.impl("npu_inplace_partial_rotary_mul", + TORCH_FN(xllm::kernel::npu::npu_inplace_partial_rotary_mul)); m.impl("silu_and_mul", TORCH_FN(xllm::silu_and_mul_npu)); m.impl("inplace_partial_rotary_mul", TORCH_FN(xllm::inplace_partial_rotary_mul_npu)); @@ -423,6 +677,15 @@ TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("sparse_flash_attention_out", TORCH_FN(xllm::kernel::npu::sparse_flash_attention_out)); m.impl("mla_preprocess_v2", TORCH_FN(xllm::kernel::npu::mla_preprocess_v2)); + m.impl("moe_gating_top_k_hash", + TORCH_FN(xllm::kernel::npu::moe_gating_top_k_hash)); + m.impl("dequant_swiglu_quant", + TORCH_FN(xllm::kernel::npu::dequant_swiglu_quant)); + m.impl("hc_pre", TORCH_FN(xllm::kernel::npu::hc_pre)); + m.impl("hc_post", TORCH_FN(xllm::kernel::npu::hc_post)); + m.impl("compressor", TORCH_FN(xllm::kernel::npu::compressor)); + m.impl("sparse_attn_sharedkv", + TORCH_FN(xllm::kernel::npu::sparse_attn_sharedkv)); } // build_cp_context is pure host index math with no Tensor input, so the @@ -431,4 +694,11 @@ TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { // graph capture), so it needs no fake/meta registration. TORCH_LIBRARY_IMPL(xllm_ops, CompositeExplicitAutograd, m) { m.impl("build_cp_context", TORCH_FN(xllm::build_cp_context_npu)); + // These metadata factories allow every Tensor argument to be omitted, so + // there may be no device key to dispatch on. Their implementations select + // the output NPU device explicitly (or inherit it from an optional Tensor). + m.impl("sparse_attn_sharedkv_metadata", + TORCH_FN(xllm::kernel::npu::sparse_attn_sharedkv_metadata)); + m.impl("quant_lightning_indexer_metadata", + TORCH_FN(xllm::kernel::npu::quant_lightning_indexer_metadata)); } diff --git a/xllm/core/runtime/acl_graph_persistent_param.cpp b/xllm/core/runtime/acl_graph_persistent_param.cpp index 420d8600ec..e3945e8c3f 100644 --- a/xllm/core/runtime/acl_graph_persistent_param.cpp +++ b/xllm/core/runtime/acl_graph_persistent_param.cpp @@ -33,6 +33,7 @@ limitations under the License. #include "core/framework/speculative/mtp_async_state.h" #include "core/kernels/npu/tilelang/tilelang_ops_api.h" #include "core/layers/common/expanded_decode_metadata_builder.h" +#include "core/runtime/decode_graph_bucket.h" #include "core/util/utils.h" // ATB includes @@ -58,13 +59,16 @@ int64_t get_decode_graph_capacity(const runtime::Options& options) { int64_t get_decode_graph_token_capacity(const runtime::Options& options) { CHECK_GT(options.num_decoding_tokens(), 0) << "num_decoding_tokens must be > 0 for graph token capacity"; + int64_t token_capacity = options.max_seqs_per_batch(); if (::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { - return options.max_seqs_per_batch(); + return runtime::get_decode_graph_token_bucket( + token_capacity, options.enable_graph_mode_decode_no_padding()); } if (options.enable_speculative_decode() && !options.is_draft_engine()) { - return options.max_seqs_per_batch() * options.num_decoding_tokens(); + token_capacity *= options.num_decoding_tokens(); } - return options.max_seqs_per_batch(); + return runtime::get_decode_graph_token_bucket( + token_capacity, options.enable_graph_mode_decode_no_padding()); } float get_dp_ep_all2all_buffer_factor(int64_t length) { diff --git a/xllm/core/runtime/py_executor_impl.cpp b/xllm/core/runtime/py_executor_impl.cpp index 8cc46c9c47..93786848da 100644 --- a/xllm/core/runtime/py_executor_impl.cpp +++ b/xllm/core/runtime/py_executor_impl.cpp @@ -17,7 +17,7 @@ limitations under the License. #include #include -#include +#include #include #include @@ -58,9 +58,7 @@ void clear_python_object(py::object& object) { object = py::object(); } -} // namespace - -PYBIND11_EMBEDDED_MODULE(xllm_runtime, m) { +void register_xllm_runtime_module(py::module_& m) { register_attention_metadata_views(m); #if defined(USE_NPU) @@ -75,6 +73,23 @@ PYBIND11_EMBEDDED_MODULE(xllm_runtime, m) { #endif } +void ensure_xllm_runtime_module() { + py::module_ sys = py::module_::import("sys"); + py::dict modules = py::reinterpret_borrow(sys.attr("modules")); + const py::str module_name("xllm_runtime"); + if (modules.contains(module_name)) { + return; + } + + py::object module_object = + py::module_::import("types").attr("ModuleType")(module_name); + py::module_ module = py::reinterpret_borrow(module_object); + register_xllm_runtime_module(module); + modules[module_name] = module; +} + +} // namespace + PyExecutorImpl::PyExecutorImpl(CausalLM* model, const ModelArgs& args, const torch::Device& device, @@ -87,7 +102,7 @@ PyExecutorImpl::PyExecutorImpl(CausalLM* model, CHECK(py_causal_lm_ != nullptr) << "PyExecutorImpl requires PyCausalLM"; py::gil_scoped_acquire gil; - py::module_::import("xllm_runtime"); + ensure_xllm_runtime_module(); py::module_ executor_module = py::module_::import("xllm.python.model_executor.executor"); py_executor_ = executor_module.attr("ModelExecutor")( diff --git a/xllm/core/runtime/speculative_worker_impl.cpp b/xllm/core/runtime/speculative_worker_impl.cpp index 178402bfa2..9e94edabe0 100644 --- a/xllm/core/runtime/speculative_worker_impl.cpp +++ b/xllm/core/runtime/speculative_worker_impl.cpp @@ -83,11 +83,19 @@ KVCacheEstimateOptions make_kv_cache_estimate_options( static_cast(options.num_speculative_tokens()); estimate_options.max_tokens_per_batch = static_cast(options.max_tokens_per_batch()); + estimate_options.max_tokens_per_chunk_for_prefill = + static_cast(options.max_tokens_per_chunk_for_prefill()); estimate_options.max_linear_state_cache_slots = options.max_linear_state_cache_slots(); estimate_options.is_draft_engine = options.is_draft_engine(); + estimate_options.enable_chunked_prefill = options.enable_chunked_prefill(); + estimate_options.enable_schedule_overlap = options.enable_schedule_overlap(); + const KVCacheConfig& kv_cache_config = KVCacheConfig::get_instance(); estimate_options.enable_prefix_cache = - KVCacheConfig::get_instance().enable_prefix_cache(); + kv_cache_config.enable_prefix_cache() && + !kv_cache_config.enable_xtensor(); + estimate_options.enable_disagg_pd = options.enable_disagg_pd(); + estimate_options.instance_role = options.instance_role(); return estimate_options; } diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index 2c5cb9e47d..bcab82dca3 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -49,6 +49,14 @@ limitations under the License. namespace xllm { +namespace { + +constexpr absl::Duration kDecodeRestoreTimeout = absl::Seconds(60); +constexpr char kDecodeRestoreTimeoutMessage[] = + "Decode request could not reacquire device KV cache within 60 seconds"; + +} // namespace + void CancelRequestQueue::submit(std::shared_ptr request) { std::lock_guard lock(mutex_); requests_.emplace_back(std::move(request)); @@ -297,6 +305,35 @@ void ContinuousScheduler::clear_mtp_bootstrap(Request* request) { sequence->clear_mtp_bootstrap_embedding(); } +void ContinuousScheduler::drain_decode_restore_waiting( + std::vector>& finished) { + const absl::Time now = absl::Now(); + for (auto it = decode_restore_waiting_.begin(); + it != decode_restore_waiting_.end();) { + std::shared_ptr& request = it->request; + CHECK(request != nullptr); + request->update_connection_status(); + if (request->finished() || request->cancelled()) { + clear_mtp_bootstrap(request.get()); + kv_cache_manager_->deallocate(request.get()); + finished.emplace_back(request); + it = decode_restore_waiting_.erase(it); + continue; + } + if (now - it->started_at < kDecodeRestoreTimeout) { + ++it; + continue; + } + + clear_mtp_bootstrap(request.get()); + kv_cache_manager_->deallocate(request.get()); + response_processor_->process_failed_request( + request, + {StatusCode::RESOURCE_EXHAUSTED, kDecodeRestoreTimeoutMessage}); + it = decode_restore_waiting_.erase(it); + } +} + std::vector ContinuousScheduler::prepare_batch() { Timer timer; drain_prefetched_requests(); @@ -305,6 +342,7 @@ std::vector ContinuousScheduler::prepare_batch() { // Common phases (strategy-independent) policy_->drain_request_queue(state, request_queue_); auto finished = policy_->collect_finished(state); + drain_decode_restore_waiting(finished); // Initialize budget ScheduleBudget budget; @@ -351,6 +389,7 @@ SchedulerState ContinuousScheduler::make_state() { .chunk_queue = *chunk_queue_, .decode_queue = *decode_queue_, .unified_queue = unified_queue_, + .decode_restore_waiting = decode_restore_waiting_, .running_requests = running_requests_, .running_sequences = running_sequences_, .running_sequences_budgets = running_sequences_budgets_, @@ -962,8 +1001,9 @@ void ContinuousScheduler::preempt_all_running_requests() { } void ContinuousScheduler::abort_all_running_requests() { - const size_t total_to_abort = - running_requests_.size() + decode_queue_->size() + chunk_queue_->size(); + const size_t total_to_abort = running_requests_.size() + + decode_queue_->size() + chunk_queue_->size() + + decode_restore_waiting_.size(); if (total_to_abort == 0) { return; } @@ -1008,6 +1048,12 @@ void ContinuousScheduler::abort_all_running_requests() { decode_queue_->pop_top(); } + // 4. Decode victims that are waiting for D2H publication or HBM capacity. + while (!decode_restore_waiting_.empty()) { + abort_one(decode_restore_waiting_.front().request); + decode_restore_waiting_.pop_front(); + } + // Clear running state. running_requests_.clear(); running_sequences_.clear(); diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h index a5059d073a..eb759ccbb5 100644 --- a/xllm/core/scheduler/continuous_scheduler.h +++ b/xllm/core/scheduler/continuous_scheduler.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -47,6 +48,11 @@ class RequestPriorityQueue; class SchedulerPolicy; struct SchedulerState; +struct DecodeRestoreEntry { + std::shared_ptr request; + absl::Time started_at; +}; + // BatchMode captures the scheduling policy configuration. // The concrete SchedulerPolicy subclass is selected based on these fields: // - enable_mix_batch=false → PrefillFirstPolicy (exclusive batch) @@ -195,7 +201,7 @@ class ContinuousScheduler : public Scheduler { uint32_t get_waiting_requests_num() const override { return prefill_queue_->size() + chunk_queue_->size() + - num_prefetch_pending_requests(); + decode_restore_waiting_.size() + num_prefetch_pending_requests(); } size_t num_prefetch_pending_requests() const; @@ -223,6 +229,10 @@ class ContinuousScheduler : public Scheduler { result.emplace_back(copied_waiting_queue->top()); copied_waiting_queue->pop_top(); } + result.reserve(result.size() + decode_restore_waiting_.size()); + for (const DecodeRestoreEntry& entry : decode_restore_waiting_) { + result.emplace_back(entry.request); + } return result; } @@ -361,6 +371,10 @@ class ContinuousScheduler : public Scheduler { // Decode queue: holds all decode-stage requests. std::unique_ptr decode_queue_; + // Decode victims that wait for D2H publication and device KV capacity before + // re-entering the existing Prefill/H2D restore path. + std::deque decode_restore_waiting_; + // Unified queue: used by UnifiedPolicy only (all requests in one queue). std::list> unified_queue_; @@ -373,7 +387,8 @@ class ContinuousScheduler : public Scheduler { virtual bool if_queue_not_empty() { return !prefill_queue_->empty() || !chunk_queue_->empty() || - !decode_queue_->empty() || !unified_queue_.empty(); + !decode_queue_->empty() || !decode_restore_waiting_.empty() || + !unified_queue_.empty(); } // tokenizer @@ -403,6 +418,9 @@ class ContinuousScheduler : public Scheduler { void apply_cancel_requests(); + void drain_decode_restore_waiting( + std::vector>& finished); + std::vector schedule_request(const absl::Duration& timeout); virtual void update_token_latency_metrics(std::vector& sequences); diff --git a/xllm/core/scheduler/decode_first_policy.cpp b/xllm/core/scheduler/decode_first_policy.cpp index a9b6b0ca4b..0a85abde1e 100644 --- a/xllm/core/scheduler/decode_first_policy.cpp +++ b/xllm/core/scheduler/decode_first_policy.cpp @@ -55,6 +55,7 @@ void DecodeFirstPolicy::schedule( // Step 1: schedule decode requests first (decode-maximal batching). schedule_decode_from_queue(&state.decode_queue, state, budget); + schedule_decode_restore(state, budget); const bool has_decode = !state.running_sequences.empty(); // Step 2: schedule prefill requests (continuations from chunk_queue first, diff --git a/xllm/core/scheduler/prefill_first_policy.cpp b/xllm/core/scheduler/prefill_first_policy.cpp index db7284f1cd..67b61b8f5b 100644 --- a/xllm/core/scheduler/prefill_first_policy.cpp +++ b/xllm/core/scheduler/prefill_first_policy.cpp @@ -88,6 +88,27 @@ void PrefillFirstPolicy::schedule( } reset_batch_state(state); + // A completed D2H release must first be available to the decode request + // that triggered preemption. Do not let the restore waiter reclaim the + // released blocks before that request can retry. + const bool retry_decode_before_restore = + !state.decode_restore_waiting.empty() && + !state.kv_cache_manager->has_pending_async_block_release(); + if (retry_decode_before_restore) { + budget.latency_budget = options_.max_global_tpot_ms(); + adjust_latency_budget_and_reorder(&state.decode_queue, + /*second_queue=*/nullptr, + budget.latency_budget, + /*for_prefill=*/false, + state); + schedule_decode_from_queue(&state.decode_queue, state, budget); + if (!state.running_sequences.empty() || + state.kv_cache_manager->has_pending_async_block_release()) { + return; + } + } + schedule_decode_restore(state, budget); + // === Schedule phase === budget.latency_budget = options_.max_global_ttft_ms(); @@ -112,7 +133,7 @@ void PrefillFirstPolicy::schedule( } // If no prefill sequences were scheduled, try decode. - if (state.running_sequences.empty()) { + if (!retry_decode_before_restore && state.running_sequences.empty()) { budget.latency_budget = options_.max_global_tpot_ms(); adjust_latency_budget_and_reorder(&state.decode_queue, /*second_queue=*/nullptr, diff --git a/xllm/core/scheduler/scheduler_policy.cpp b/xllm/core/scheduler/scheduler_policy.cpp index 78591325d8..8a9f849540 100644 --- a/xllm/core/scheduler/scheduler_policy.cpp +++ b/xllm/core/scheduler/scheduler_policy.cpp @@ -15,6 +15,7 @@ limitations under the License. #include "scheduler/scheduler_policy.h" +#include #include #include @@ -27,6 +28,7 @@ limitations under the License. #include "core/framework/config/kv_cache_store_config.h" #include "core/framework/config/parallel_config.h" #include "core/framework/config/scheduler_config.h" +#include "core/framework/config/speculative_config.h" #include "framework/batch/batch_factory.h" #include "framework/request/priority_comparator.h" #include "util/timer.h" @@ -449,6 +451,43 @@ void SchedulerPolicy::allocate_shared_blocks_for(Sequence* seq, } } +void SchedulerPolicy::schedule_decode_restore(SchedulerState& state, + ScheduleBudget& budget) { + if (state.decode_restore_waiting.empty() || + state.kv_cache_manager->has_pending_async_block_release() || + budget_exhausted(budget)) { + return; + } + + DecodeRestoreEntry& entry = state.decode_restore_waiting.front(); + const std::shared_ptr& request = entry.request; + CHECK(request != nullptr); + CHECK_EQ(request->sequences().size(), 1u); + Sequence* sequence = request->sequences().front().get(); + CHECK(sequence != nullptr); + + const size_t num_tokens = + compute_prefill_tokens(sequence, budget.remaining_token_budget, state); + if (num_tokens == 0 || num_tokens > budget.remaining_token_budget || + budget.remaining_seq_budget == 0) { + return; + } + + size_t actual_tokens = 0; + if (!allocate_for_prefill(sequence, num_tokens, &actual_tokens, state)) { + return; + } + + CHECK_LE(actual_tokens, budget.remaining_token_budget); + state.running_requests.emplace_back(request); + state.running_sequences.emplace_back(sequence); + state.running_sequences_budgets.emplace_back(actual_tokens); + budget.remaining_token_budget -= actual_tokens; + --budget.remaining_seq_budget; + cache_in_batch_prefix({sequence}, {actual_tokens}, state); + state.decode_restore_waiting.pop_front(); +} + // ============================================================================= // Decode scheduling // ============================================================================= @@ -648,8 +687,15 @@ void SchedulerPolicy::schedule_decode_from_queue(RequestPriorityQueue* queue, break; } - // Blocks exhausted: first try preempting from chunk_queue (has KV blocks - // to free), then from the decode queue's lowest priority. + // Blocks exhausted: wait for an in-flight async release before selecting + // another victim. The released blocks remain unavailable until the + // transfer completes. + if (state.kv_cache_manager->has_pending_async_block_release()) { + return; + } + + // First try preempting from chunk_queue (has KV blocks to free), then + // from the decode queue's lowest priority. if (!has_enough_blocks && !state.chunk_queue.empty()) { std::shared_ptr request_to_preempt = state.chunk_queue.back(); ++budget.num_preempted_requests; @@ -657,16 +703,20 @@ void SchedulerPolicy::schedule_decode_from_queue(RequestPriorityQueue* queue, state.chunk_queue.pop_back(); request_to_preempt->set_preempted(); state.prefill_queue.push(request_to_preempt); + if (state.kv_cache_manager->has_pending_async_block_release()) { + return; + } continue; } if (!has_enough_blocks && queue->size() > 1) { std::shared_ptr request_to_preempt = queue->back(); if (request_to_preempt.get() != request.get()) { ++budget.num_preempted_requests; - state.kv_cache_manager->deallocate(request_to_preempt.get()); + enqueue_decode_restore(request_to_preempt, state); queue->pop_back(); - request_to_preempt->set_preempted(); - state.prefill_queue.push(request_to_preempt); + if (state.kv_cache_manager->has_pending_async_block_release()) { + return; + } continue; } } @@ -686,6 +736,42 @@ void SchedulerPolicy::schedule_decode_from_queue(RequestPriorityQueue* queue, } } +bool SchedulerPolicy::should_wait_for_decode_restore( + const std::shared_ptr& request, + const SchedulerState& state) const { + if (!state.options.enable_disagg_pd() || state.options.enable_pd_ooc() || + !state.options.instance_role().has_value() || + state.options.instance_role().value() != InstanceRole::DECODE || + !state.enable_prefix_cache || state.options.enable_schedule_overlap() || + !state.kv_cache_manager->supports_host_cache_restore() || + ::xllm::KVCacheStoreConfig::get_instance().host_blocks_factor() <= 1.0 || + request == nullptr || request->sequences().size() != 1 || + request->check_beam_search()) { + return false; + } + + return state.options.num_speculative_tokens() == 0 || + ::xllm::SpeculativeConfig::get_instance().speculative_algorithm() == + "MTP"; +} + +void SchedulerPolicy::enqueue_decode_restore( + const std::shared_ptr& request, + SchedulerState& state) { + const bool should_wait = should_wait_for_decode_restore(request, state); + if (should_wait) { + clear_mtp_bootstrap(request.get(), state); + } + state.kv_cache_manager->deallocate(request.get()); + request->set_preempted(); + if (should_wait) { + state.decode_restore_waiting.emplace_back( + DecodeRestoreEntry{request, absl::Now()}); + return; + } + state.prefill_queue.push(request); +} + // ============================================================================= // Preemption and error handling // ============================================================================= @@ -775,7 +861,8 @@ void SchedulerPolicy::report_metrics(const SchedulerState& state, double elapsed_seconds, size_t num_preempted_requests) { GAUGE_SET(num_running_requests, state.running_requests.size()); - GAUGE_SET(num_waiting_requests, state.prefill_queue.size()); + GAUGE_SET(num_waiting_requests, + state.prefill_queue.size() + state.decode_restore_waiting.size()); GAUGE_SET(num_preempted_requests, num_preempted_requests); GAUGE_SET(num_running_sequences, state.running_sequences.size()); GAUGE_SET(kv_cache_utilization_perc, diff --git a/xllm/core/scheduler/scheduler_policy.h b/xllm/core/scheduler/scheduler_policy.h index b022d35e29..5dd2959b8f 100644 --- a/xllm/core/scheduler/scheduler_policy.h +++ b/xllm/core/scheduler/scheduler_policy.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -47,6 +48,7 @@ struct SchedulerState { RequestPriorityQueue& chunk_queue; RequestPriorityQueue& decode_queue; std::list>& unified_queue; + std::deque& decode_restore_waiting; // Current batch state (reset each step). std::vector>& running_requests; @@ -159,6 +161,7 @@ class SchedulerPolicy { SchedulerState& state, bool skip_shared = false); void allocate_shared_blocks_for(Sequence* seq, SchedulerState& state); + void schedule_decode_restore(SchedulerState& state, ScheduleBudget& budget); // ===== Decode scheduling ===== void schedule_decode_from_queue(RequestPriorityQueue* queue, @@ -183,6 +186,10 @@ class SchedulerPolicy { size_t allocated_seqs, double allocated_estimate_latency, bool budget_exhausted); + bool should_wait_for_decode_restore(const std::shared_ptr& request, + const SchedulerState& state) const; + void enqueue_decode_restore(const std::shared_ptr& request, + SchedulerState& state); // ===== Helpers ===== void cache_in_batch_prefix(const std::vector& sequences, diff --git a/xllm/models/llm/npu/deepseek_v2.h b/xllm/models/llm/npu/deepseek_v2.h index b3b9a8a6a1..821697fee1 100644 --- a/xllm/models/llm/npu/deepseek_v2.h +++ b/xllm/models/llm/npu/deepseek_v2.h @@ -175,7 +175,8 @@ class DeepseekV2ModelImpl : public torch::nn::Module { if (::xllm::KVCacheConfig::get_instance().enable_prefix_cache() && !input_params.meta.batch_forward_type.is_decode()) { attn_mask = attn_mask_.get_attn_mask(512, dtype_, device_); - } else if (input_params.meta.batch_forward_type.is_prefill()) { + } else if (input_params.meta.batch_forward_type.is_prefill() || + input_params.meta.batch_forward_type.is_chunked_prefill()) { attn_mask = attn_mask_.get_attn_mask(128, dtype_, device_); } else if (num_speculative_tokens_ > 0) { // TODO :the judgement of gen_free_mask need more check diff --git a/xllm/models/llm/py_causal_lm.cpp b/xllm/models/llm/py_causal_lm.cpp index cedfbdc5da..e424c67309 100644 --- a/xllm/models/llm/py_causal_lm.cpp +++ b/xllm/models/llm/py_causal_lm.cpp @@ -214,7 +214,7 @@ py::dict PyCausalLM::build_config_dict( void PyCausalLM::load_model(std::unique_ptr loader) { py::gil_scoped_acquire gil; auto& state_dicts = loader->get_state_dicts(); - py::module_::import("xllm_weight_loader"); + ensure_xllm_weight_loader_module(); py::list py_state_dicts; for (const auto& sd : state_dicts) { diff --git a/xllm/models/py_model_helper.cpp b/xllm/models/py_model_helper.cpp index be35a38320..a8b292f00c 100644 --- a/xllm/models/py_model_helper.cpp +++ b/xllm/models/py_model_helper.cpp @@ -15,7 +15,7 @@ limitations under the License. // Infrastructure for the embedded Python model executor: // - Interpreter lifecycle (ensure_python_interpreter) -// - Weight loading (PyStateDict + PYBIND11_EMBEDDED_MODULE) +// - Weight loading (PyStateDict Python binding) // - Config serialization (dtype_to_string, PyDictVisitor) #include "models/py_model_helper.h" @@ -147,11 +147,22 @@ py::list PyStateDict::keys() const { return result; } -PYBIND11_EMBEDDED_MODULE(xllm_weight_loader, m) { - py::class_(m, "StateDict") +void ensure_xllm_weight_loader_module() { + py::module_ sys = py::module_::import("sys"); + py::dict modules = py::reinterpret_borrow(sys.attr("modules")); + const py::str module_name("xllm_weight_loader"); + if (modules.contains(module_name)) { + return; + } + + py::object module_object = + py::module_::import("types").attr("ModuleType")(module_name); + py::module_ module = py::reinterpret_borrow(module_object); + py::class_(module, "StateDict") .def("get_tensor", &PyStateDict::get_tensor, py::arg("name")) .def("has", &PyStateDict::has, py::arg("name")) .def("keys", &PyStateDict::keys); + modules[module_name] = module; } } // namespace xllm diff --git a/xllm/models/py_model_helper.h b/xllm/models/py_model_helper.h index 42419199fd..9315ebc134 100644 --- a/xllm/models/py_model_helper.h +++ b/xllm/models/py_model_helper.h @@ -30,6 +30,10 @@ namespace xllm { // Initializes the embedded CPython interpreter (idempotent, process-wide). void ensure_python_interpreter(); +// Makes the internal StateDict binding importable as xllm_weight_loader. +// The caller must hold the Python GIL. +void ensure_xllm_weight_loader_module(); + // Convert torch dtype to the string form used by Python model config. std::string dtype_to_string(const torch::TensorOptions& options); diff --git a/xllm/python/kernels_cuda/__init__.py b/xllm/python/kernels_cuda/__init__.py index 36cbeca543..592c57fccb 100644 --- a/xllm/python/kernels_cuda/__init__.py +++ b/xllm/python/kernels_cuda/__init__.py @@ -48,6 +48,7 @@ chunk_gated_delta_rule, fused_gdn_prefill_post_conv, fused_recurrent_gated_delta_rule_packed_decode, + gdn_prefill_prepare, resolve_gdn_prefill_backend, ) from .linear import prepare_row_parallel_weight @@ -115,6 +116,7 @@ "causal_conv1d_prefill", "causal_conv1d_decode", "resolve_gdn_prefill_backend", + "gdn_prefill_prepare", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", diff --git a/xllm/python/kernels_cuda/gated_delta_net.py b/xllm/python/kernels_cuda/gated_delta_net.py index b7409c3bb7..377bc87c7d 100644 --- a/xllm/python/kernels_cuda/gated_delta_net.py +++ b/xllm/python/kernels_cuda/gated_delta_net.py @@ -199,9 +199,58 @@ def _chunk_gated_delta_rule_fake( return torch.empty_like(v), torch.empty_like(initial_state) +def gdn_prefill_prepare( + mixed_qkv: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split + l2norm + gating for prefill. + + Encapsulates the CUDA-optimal fusion strategy: causal_conv1d_prefill + produces packed convolved output, then fused_gdn_prefill_post_conv + splits into Q/K/V with l2norm and computes gating in one kernel. + + Returns: + (q, k, v, g, beta) with shapes [T, H, D] / [T, H]. + """ + from .causal_conv1d import causal_conv1d_prefill + + convolved = causal_conv1d_prefill( + mixed_qkv, + weight, + conv_state, + state_indices, + has_initial_state, + cu_seqlens, + ) + q, k, v, g, beta = fused_gdn_prefill_post_conv( + convolved, + a, + b, + a_log, + dt_bias, + num_key_heads, + key_head_dim, + value_head_dim, + ) + return q, k, v, g, beta + + __all__ = [ "GdnPrefillBackend", "resolve_gdn_prefill_backend", + "gdn_prefill_prepare", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", diff --git a/xllm/python/kernels_npu/__init__.py b/xllm/python/kernels_npu/__init__.py index 5626492d7d..364218602c 100644 --- a/xllm/python/kernels_npu/__init__.py +++ b/xllm/python/kernels_npu/__init__.py @@ -27,18 +27,19 @@ from typing import Any _EXPORTS = { - "activation": ("dequant_swiglu_quant", "silu_and_mul"), + "activation": ("silu_and_mul",), "attention": ( "batch_matmul_transpose", "reshape_paged_cache", "update_decode_graph_metadata", "vision_fusion_attention", ), - "causal_conv1d": ("causal_conv1d_decode", "causal_conv1d_prefill"), + "causal_conv1d": ("causal_conv1d_decode", "causal_conv1d_qkv_prefill"), "gated_delta_net": ( "chunk_gated_delta_rule", - "fused_gdn_prefill_post_conv", + "fused_gdn_gating", "fused_recurrent_gated_delta_rule_packed_decode", + "gdn_prefill_prepare", "resolve_gdn_prefill_backend", ), "linear": ("prepare_quant_weight", "prepare_row_parallel_weight"), @@ -56,6 +57,9 @@ "grouped_moe", "moe_expert_compute", "moe_fused_topk", + "dequant_swiglu_quant", + "grouped_moe_with_selected_experts", + "moe_gating_top_k_hash", "moe_gate_routing", "moe_gmm1", "moe_gmm2_combine", @@ -68,12 +72,14 @@ "fused_add_rms_norm_dynamic_quant", "l2_norm", "rms_norm", + "rms_norm_dynamic_quant", "rms_norm_gated", ), "quantization": ("dynamic_quant", "quant_matmul", "quantize_per_tensor"), "rotary_embedding": ( "fused_qk_norm_rope", "interleaved_rotary_embedding", + "npu_inplace_partial_rotary_mul", "mrope", "vision_rotary_mul", ), @@ -86,12 +92,22 @@ "sparse_flash_attention", "sparse_flash_attention_out", ), + "dsa": ( + "compressor", + "hc_post", + "hc_pre", + "quant_lightning_indexer", + "quant_lightning_indexer_metadata", + "sparse_attn_sharedkv", + "sparse_attn_sharedkv_metadata", + ), } __all__ = [ "rms_norm", "fused_add_rms_norm", "fused_add_rms_norm_dynamic_quant", + "rms_norm_dynamic_quant", "l2_norm", "rms_norm_gated", "silu_and_mul", @@ -102,6 +118,7 @@ "batch_matmul_transpose", "fused_qk_norm_rope", "interleaved_rotary_embedding", + "npu_inplace_partial_rotary_mul", "mrope", "vision_rotary_mul", "moe_fused_topk", @@ -114,6 +131,7 @@ "moe_token_dispatch", "moe_gmm1", "moe_gmm2_combine", + "grouped_moe_with_selected_experts", "prepare_grouped_moe_weights", "supports_cutlass_moe", "prepare_row_parallel_weight", @@ -133,10 +151,18 @@ "scatter_nd_update", "sparse_flash_attention", "sparse_flash_attention_out", - "causal_conv1d_prefill", "causal_conv1d_decode", + "compressor", + "dequant_swiglu_quant", + "hc_pre", + "hc_post", + "moe_gating_top_k_hash", + "quant_lightning_indexer", + "quant_lightning_indexer_metadata", + "sparse_attn_sharedkv", + "sparse_attn_sharedkv_metadata", "resolve_gdn_prefill_backend", - "fused_gdn_prefill_post_conv", + "gdn_prefill_prepare", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", ] diff --git a/xllm/python/kernels_npu/_custom_op.py b/xllm/python/kernels_npu/_custom_op.py index 2957a5cc85..7feedb082f 100644 --- a/xllm/python/kernels_npu/_custom_op.py +++ b/xllm/python/kernels_npu/_custom_op.py @@ -66,6 +66,77 @@ def _rms_norm_fake( return torch.empty_like(input) +def _rms_norm_gated_fake( + input: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + del gate, weight, eps + return torch.empty_like(input) + + +def _l2_norm_fake( + input: torch.Tensor, + eps: float, +) -> torch.Tensor: + del eps + return torch.empty_like(input) + + +def _chunk_gated_delta_rule_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + del q, k, g, beta, cu_seqlens + num_seqs = initial_state.shape[0] + return torch.empty_like(v), torch.empty_like(initial_state) + + +def _causal_conv1d_qkv_prefill_fake( + x: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + query_start_loc: torch.Tensor, + num_qk_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del weight, conv_state, state_indices, has_initial_state, query_start_loc + num_tokens = x.shape[0] + opts = x.options().dtype(torch.bfloat16) + q = torch.empty(1, num_tokens, num_qk_heads, head_k_dim, **opts) + k = torch.empty(1, num_tokens, num_qk_heads, head_k_dim, **opts) + v = torch.empty(1, num_tokens, num_v_heads, head_v_dim, **opts) + return q, k, v + + +def _fused_sigmoid_gating_delta_rule_decode_fake( + a_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + ssm_state: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + scale: float, +) -> torch.Tensor: + del a_log, a, dt_bias, k, b, ssm_state, state_indices, cu_seqlens, scale + # Output shape matches v: [num_tokens, num_value_heads, value_dim] + return torch.empty_like(v) + + def _fused_add_rms_norm_fake( input: torch.Tensor, residual: torch.Tensor, @@ -522,7 +593,383 @@ def _sparse_flash_attention_out_fake( return output +# --------------------------------------------------------------------------- +# DeepSeek-V4 DSA kernel fakes +# --------------------------------------------------------------------------- + +# Matches kDsaMetadataBufferElements in xllm_ops_api.h. +_DSA_METADATA_BUFFER_ELEMENTS = 1024 + + +def _rms_norm_dynamic_quant_fake( + input: torch.Tensor, weight: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: + del weight, eps + return input.new_empty(input.shape, dtype=torch.int8), input.new_empty(input.shape[:-1], dtype=torch.float32) + + +def _npu_inplace_partial_rotary_mul_fake( + x: torch.Tensor, + r1: torch.Tensor, + r2: torch.Tensor, + rotary_mode: str, + partial_slice: list[int], +) -> None: + del x, r1, r2, rotary_mode, partial_slice + + +def _hc_pre_fake( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + hc_sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del hc_fn, hc_scale, hc_base, hc_sinkhorn_iters, norm_eps, hc_eps + if x.dim() == 4: + y_shape = (x.size(0), x.size(1), x.size(3)) + post_shape = (x.size(0), x.size(1), hc_mult) + comb_shape = (x.size(0), x.size(1), hc_mult, hc_mult) + else: + y_shape = (x.size(0), x.size(2)) + post_shape = (x.size(0), hc_mult) + comb_shape = (x.size(0), hc_mult, hc_mult) + attn_input = x.new_empty(y_shape, dtype=x.dtype) + post = x.new_empty(post_shape, dtype=torch.float32) + comb = x.new_empty(comb_shape, dtype=torch.float32) + return attn_input, post, comb + + +def _hc_post_fake( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + del post, comb + # hc_post returns [T, hc_mult, hidden] (the merged residual streams). + return residual.new_empty(residual.shape, dtype=residual.dtype) + + +def _compressor_fake( + x: torch.Tensor, + wkv: torch.Tensor, + wgate: torch.Tensor, + kv_state: torch.Tensor, + score_state: torch.Tensor, + ape: torch.Tensor, + norm_weight: torch.Tensor, + rope_sin: torch.Tensor, + rope_cos: torch.Tensor, + kv_block_table: torch.Tensor | None, + score_block_table: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + seqused: torch.Tensor | None, + start_pos: torch.Tensor | None, + rope_head_dim: int, + cmp_ratio: int, + coff: int, + norm_eps: float, + rotary_mode: int, + enable_grad: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + del ( + wkv, + wgate, + kv_state, + score_state, + ape, + rope_cos, + kv_block_table, + score_block_table, + cu_seqlens, + seqused, + start_pos, + rope_head_dim, + norm_eps, + rotary_mode, + ) + head_dim = norm_weight.size(0) + if x.dim() == 3: + compressed_seq = (x.size(1) + cmp_ratio - 1) // cmp_ratio + cmp_kv_shape = (x.size(0), compressed_seq, head_dim) + grad_shapes = ( + (x.size(0), x.size(1), coff * head_dim), + (x.size(0), compressed_seq, coff * cmp_ratio, head_dim), + (x.size(0), compressed_seq, head_dim), + (x.size(0), compressed_seq), + ) + else: + compressed_seq = rope_sin.size(0) + cmp_kv_shape = (compressed_seq, head_dim) + grad_shapes = ( + (x.size(0), coff * head_dim), + (compressed_seq, coff * cmp_ratio, head_dim), + (compressed_seq, head_dim), + (compressed_seq,), + ) + outputs = [x.new_empty(cmp_kv_shape, dtype=x.dtype)] + if enable_grad: + outputs.extend(x.new_empty(shape, dtype=x.dtype) for shape in grad_shapes) + else: + outputs.extend(x.new_empty((0,), dtype=x.dtype) for _ in grad_shapes) + return tuple(outputs) # type: ignore[return-value] + + +def _sparse_attn_sharedkv_fake( + q: torch.Tensor, + ori_kv: torch.Tensor | None, + cmp_kv: torch.Tensor | None, + ori_sparse_indices: torch.Tensor | None, + cmp_sparse_indices: torch.Tensor | None, + ori_block_table: torch.Tensor | None, + cmp_block_table: torch.Tensor | None, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + sinks: torch.Tensor | None, + metadata: torch.Tensor | None, + softmax_scale: float, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + return_softmax_lse: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + del ( + ori_kv, + cmp_kv, + ori_sparse_indices, + cmp_sparse_indices, + ori_block_table, + cmp_block_table, + sinks, + metadata, + softmax_scale, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + ) + out = q.new_empty(q.shape, dtype=q.dtype) + lse_shape = (*q.shape[:-1], 1) if return_softmax_lse else (0,) + lse = q.new_empty(lse_shape, dtype=torch.float32) + return out, lse + + +def _sparse_attn_sharedkv_metadata_fake( + num_heads_q: int, + num_heads_kv: int, + head_dim: int, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_kv: int, + ori_topk: int, + cmp_topk: int, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + has_ori_kv: bool, + has_cmp_kv: bool, +) -> torch.Tensor: + del ( + num_heads_q, + num_heads_kv, + head_dim, + batch_size, + max_seqlen_q, + max_seqlen_kv, + ori_topk, + cmp_topk, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + has_ori_kv, + has_cmp_kv, + ) + for tensor in ( + cu_seqlens_q, + cu_seqlens_ori_kv, + cu_seqlens_cmp_kv, + seqused_q, + seqused_kv, + ): + if tensor is not None: + return tensor.new_empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32) + return torch.empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32, device="npu") + + +def _quant_lightning_indexer_fake( + query: torch.Tensor, + key: torch.Tensor, + weights: torch.Tensor, + query_dequant_scale: torch.Tensor, + key_dequant_scale: torch.Tensor, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + block_table: torch.Tensor | None, + metadata: torch.Tensor | None, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + return_value: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + del ( + weights, + query_dequant_scale, + key_dequant_scale, + query_quant_mode, + key_quant_mode, + block_table, + metadata, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + ) + key_head_num = key.size(1) if layout_key == "TND" else key.size(2) + if layout_query == "BSND": + out_shape = (query.size(0), query.size(1), key_head_num, sparse_count) + else: + out_shape = (query.size(0), key_head_num, sparse_count) + out = query.new_zeros(out_shape, dtype=torch.int32) + val = ( + query.new_empty(out_shape, dtype=torch.float32) if return_value else query.new_empty((0,), dtype=torch.float32) + ) + return out, val + + +def _quant_lightning_indexer_metadata_fake( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_k: int, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + device: str, +) -> torch.Tensor: + del ( + num_heads_q, + num_heads_k, + head_dim, + query_quant_mode, + key_quant_mode, + batch_size, + max_seqlen_q, + max_seqlen_k, + layout_query, + layout_key, + sparse_count, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + ) + for tensor in (actual_seq_lengths_query, actual_seq_lengths_key): + if tensor is not None: + return tensor.new_empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32) + return torch.empty((_DSA_METADATA_BUFFER_ELEMENTS,), dtype=torch.int32, device=device) + + +def _moe_gating_top_k_hash_fake( + x: torch.Tensor, + k: int, + bias: torch.Tensor | None, + input_ids: torch.Tensor | None, + tid2eid: torch.Tensor | None, + k_group: int, + group_count: int, + routed_scaling_factor: float, + eps: float, + group_select_mode: int, + renorm: int, + norm_type: int, + out_flag: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del bias, input_ids, tid2eid, k_group, group_count, routed_scaling_factor + del eps, group_select_mode, renorm, norm_type, out_flag + y_shape = (*x.shape[:-1], k) + y = x.new_empty(y_shape, dtype=x.dtype) + expert_idx = x.new_empty(y_shape, dtype=torch.int32) + out = x.new_empty(x.shape, dtype=torch.float32) + return y, expert_idx, out + + +def _dequant_swiglu_quant_fake( + x: torch.Tensor, + weight_scale: torch.Tensor | None, + activation_scale: torch.Tensor | None, + bias: torch.Tensor | None, + quant_scale: torch.Tensor | None, + quant_offset: torch.Tensor | None, + group_index: torch.Tensor | None, + activate_left: bool, + quant_mode: int, + swiglu_mode: int, + clamp_limit: float, + glu_alpha: float, + glu_bias: float, +) -> tuple[torch.Tensor, torch.Tensor]: + del weight_scale, activation_scale, bias, quant_scale, quant_offset + del group_index, activate_left, quant_mode, swiglu_mode + del clamp_limit, glu_alpha, glu_bias + # Output is half of input's last dim (SwiGLU splits gate/up). + out_dim = x.size(-1) // 2 + act_quantized = x.new_empty((*x.shape[:-1], out_dim), dtype=torch.int8) + act_scale = x.new_empty(x.shape[:-1], dtype=torch.float32) + return act_quantized, act_scale + + register_fake("xllm_ops::rms_norm", _rms_norm_fake) +register_fake("xllm_ops::rms_norm_gated", _rms_norm_gated_fake) +register_fake("xllm_ops::l2_norm", _l2_norm_fake) +register_fake("xllm_ops::chunk_gated_delta_rule", _chunk_gated_delta_rule_fake) +register_fake("xllm_ops::causal_conv1d_qkv_prefill", _causal_conv1d_qkv_prefill_fake) +register_fake( + "xllm_ops::fused_sigmoid_gating_delta_rule_decode", + _fused_sigmoid_gating_delta_rule_decode_fake, +) register_fake("xllm_ops::fused_add_rms_norm", _fused_add_rms_norm_fake) register_fake("xllm_ops::silu_and_mul", _silu_and_mul_fake) register_fake("xllm_ops::reshape_paged_cache", _reshape_paged_cache_fake) @@ -548,3 +995,15 @@ def _sparse_flash_attention_out_fake( register_fake("xllm_ops::scatter_nd_update", _scatter_nd_update_fake) register_fake("xllm_ops::sparse_flash_attention", _sparse_flash_attention_fake) register_fake("xllm_ops::sparse_flash_attention_out", _sparse_flash_attention_out_fake) +register_fake("xllm_ops::rms_norm_dynamic_quant", _rms_norm_dynamic_quant_fake) +register_fake( + "xllm_ops::npu_inplace_partial_rotary_mul", + _npu_inplace_partial_rotary_mul_fake, +) +register_fake("xllm_ops::compressor", _compressor_fake) +register_fake("xllm_ops::moe_gating_top_k_hash", _moe_gating_top_k_hash_fake) +register_fake("xllm_ops::dequant_swiglu_quant", _dequant_swiglu_quant_fake) +register_fake("xllm_ops::hc_pre", _hc_pre_fake) +register_fake("xllm_ops::hc_post", _hc_post_fake) +register_fake("xllm_ops::sparse_attn_sharedkv", _sparse_attn_sharedkv_fake) +register_fake("xllm_ops::sparse_attn_sharedkv_metadata", _sparse_attn_sharedkv_metadata_fake) diff --git a/xllm/python/kernels_npu/causal_conv1d.py b/xllm/python/kernels_npu/causal_conv1d.py index 5f2fbe3a19..d506260d7d 100644 --- a/xllm/python/kernels_npu/causal_conv1d.py +++ b/xllm/python/kernels_npu/causal_conv1d.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NPU causal-convolution kernels. +"""NPU causal-convolution kernels (PyTorch small-op implementation). -Neither has an NPU kernel yet. The signatures are the contract an NPU -implementation has to meet; see ``kernels_cuda/causal_conv1d.py`` and the -Triton launcher it calls for the reference behaviour. +Implements the same semantics as the CUDA Triton reference in +``kernels_cuda/triton/causal_conv1d.py`` using only standard PyTorch +operations. Performance is not optimized; correctness and precision +alignment are the goals. """ from __future__ import annotations @@ -24,32 +25,35 @@ import torch -def causal_conv1d_prefill( +def causal_conv1d_qkv_prefill( value: torch.Tensor, weight: torch.Tensor, conv_state: torch.Tensor, state_indices: torch.Tensor, has_initial_state: torch.Tensor, query_start_loc: torch.Tensor, -) -> torch.Tensor: - """Convolve a variable-length batch and update the convolution states. - - Args: - value: Packed activations of shape ``[num_tokens, channels]``. - weight: Depthwise kernel of shape ``[channels, kernel_size]``. - conv_state: Per-sequence convolution state, updated in place. - state_indices: State slot of every sequence. - has_initial_state: Whether a sequence continues an earlier state. - query_start_loc: Start offset of every sequence in ``value``. + num_qk_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split into Q/K/V for prefill. Returns: - Convolved activations with the shape and dtype of ``value``. + (q, k, v) with shapes [1, T, num_qk_heads, head_k_dim], + [1, T, num_qk_heads, head_k_dim], [1, T, num_v_heads, head_v_dim]. """ - del value, weight, conv_state, state_indices, has_initial_state - del query_start_loc - raise NotImplementedError( - "causal_conv1d_prefill has no NPU kernel; see " - "kernels_cuda/triton/causal_conv1d.py for the reference implementation" + return torch.ops.xllm_ops.causal_conv1d_qkv_prefill( + value, + weight, + conv_state, + state_indices, + has_initial_state.to(torch.int64), + query_start_loc, + num_qk_heads, + num_v_heads, + head_k_dim, + head_v_dim, ) @@ -70,11 +74,20 @@ def causal_conv1d_decode( Returns: Convolved activations with the shape and dtype of ``value``. """ - del value, weight, conv_state, state_indices - raise NotImplementedError( - "causal_conv1d_decode has no NPU kernel; see " - "kernels_cuda/triton/causal_conv1d.py for the reference implementation" + from .tilelang.causal_conv1d_decode import causal_conv1d_decode as _tl_decode + + # TileLang expects conv_state as [slots, dim, state_len] (PyTorch convention) + # Our cache is [slots, state_len, dim], so transpose before calling + conv_state_pt = conv_state.transpose(1, 2).contiguous() + result = _tl_decode( + x=value, + conv_state=conv_state_pt, + weight=weight, + conv_state_indices=state_indices, ) + # TileLang writes back to conv_state_pt, copy back to original layout + conv_state.copy_(conv_state_pt.transpose(1, 2)) + return result -__all__ = ["causal_conv1d_prefill", "causal_conv1d_decode"] +__all__ = ["causal_conv1d_qkv_prefill", "causal_conv1d_decode"] diff --git a/xllm/python/kernels_npu/dsa.py b/xllm/python/kernels_npu/dsa.py new file mode 100644 index 0000000000..4e977636b5 --- /dev/null +++ b/xllm/python/kernels_npu/dsa.py @@ -0,0 +1,306 @@ +# Copyright 2026 The xLLM Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""NPU DeepSeek-V4 DSA kernels. + +These wrap the AscendC operators registered as ``torch.ops.xllm_ops.*`` by +``core/kernels/npu/npu_ops_library.cpp``. They drive the two-stage sparse +attention (original + compressed KV), the KV compressor, the quantized +lightning indexer, and the HyperConnection pre/post used by DeepSeek-V4's DSA +attention path. +""" + +from __future__ import annotations + +import torch + + +def hc_pre( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + hc_sinkhorn_iters: int, + norm_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """HyperConnection pre: mix hc_mult streams into one sub-block input. + + Returns ``(attn_input, post, comb)`` where post/comb feed ``hc_post``. + """ + return torch.ops.xllm_ops.hc_pre(x, hc_fn, hc_scale, hc_base, hc_mult, hc_sinkhorn_iters, norm_eps, hc_eps) + + +def hc_post( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + """HyperConnection post: combine sub-block output with the residual streams.""" + return torch.ops.xllm_ops.hc_post(x, residual, post, comb) + + +def compressor( + x: torch.Tensor, + wkv: torch.Tensor, + wgate: torch.Tensor, + kv_state: torch.Tensor, + score_state: torch.Tensor, + ape: torch.Tensor, + norm_weight: torch.Tensor, + rope_sin: torch.Tensor, + rope_cos: torch.Tensor, + kv_block_table: torch.Tensor | None, + score_block_table: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + seqused: torch.Tensor | None, + start_pos: torch.Tensor | None, + rope_head_dim: int, + cmp_ratio: int, + coff: int, + norm_eps: float, + rotary_mode: int, + enable_grad: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pool KV along the token axis by ``cmp_ratio`` (NSA-style compressor). + + ``kv_state`` and ``score_state`` are updated in place. + + Returns ``(cmp_kv, wkv_proj, softmax_res, norm_x, norm_rstd)``; only + ``cmp_kv`` is consumed by the DSA path. + """ + # C++ moves DSA metadata to the active device before dispatch. Keep this + # adapter deterministic; experimental clone/noalias paths do not belong in + # the public binding. + kv_block_table = kv_block_table.to(x.device) if kv_block_table is not None else None + score_block_table = score_block_table.to(x.device) if score_block_table is not None else None + cu_seqlens = cu_seqlens.to(x.device) if cu_seqlens is not None else None + seqused = seqused.to(x.device) if seqused is not None else None + start_pos = start_pos.to(x.device) if start_pos is not None else None + return torch.ops.xllm_ops.compressor( + x, + wkv, + wgate, + kv_state, + score_state, + ape, + norm_weight, + rope_sin, + rope_cos, + kv_block_table, + score_block_table, + cu_seqlens, + seqused, + start_pos, + rope_head_dim, + cmp_ratio, + coff, + norm_eps, + rotary_mode, + enable_grad, + ) + + +def sparse_attn_sharedkv( + q: torch.Tensor, + ori_kv: torch.Tensor | None, + cmp_kv: torch.Tensor | None, + ori_sparse_indices: torch.Tensor | None, + cmp_sparse_indices: torch.Tensor | None, + ori_block_table: torch.Tensor | None, + cmp_block_table: torch.Tensor | None, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + sinks: torch.Tensor | None, + metadata: torch.Tensor | None, + softmax_scale: float, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + return_softmax_lse: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Two-stage sparse attention over original and compressed KV.""" + return torch.ops.xllm_ops.sparse_attn_sharedkv( + q, + ori_kv, + cmp_kv, + ori_sparse_indices, + cmp_sparse_indices, + ori_block_table, + cmp_block_table, + cu_seqlens_q, + cu_seqlens_ori_kv, + cu_seqlens_cmp_kv, + seqused_q, + seqused_kv, + sinks, + metadata, + softmax_scale, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + return_softmax_lse, + ) + + +def sparse_attn_sharedkv_metadata( + num_heads_q: int, + num_heads_kv: int, + head_dim: int, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_ori_kv: torch.Tensor | None, + cu_seqlens_cmp_kv: torch.Tensor | None, + seqused_q: torch.Tensor | None, + seqused_kv: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_kv: int, + ori_topk: int, + cmp_topk: int, + cmp_ratio: int, + ori_mask_mode: int, + cmp_mask_mode: int, + ori_win_left: int, + ori_win_right: int, + layout_q: str, + layout_kv: str, + has_ori_kv: bool, + has_cmp_kv: bool, +) -> torch.Tensor: + """Build the AICPU tiling metadata for :func:`sparse_attn_sharedkv`.""" + return torch.ops.xllm_ops.sparse_attn_sharedkv_metadata( + num_heads_q, + num_heads_kv, + head_dim, + cu_seqlens_q, + cu_seqlens_ori_kv, + cu_seqlens_cmp_kv, + seqused_q, + seqused_kv, + batch_size, + max_seqlen_q, + max_seqlen_kv, + ori_topk, + cmp_topk, + cmp_ratio, + ori_mask_mode, + cmp_mask_mode, + ori_win_left, + ori_win_right, + layout_q, + layout_kv, + has_ori_kv, + has_cmp_kv, + ) + + +def quant_lightning_indexer( + query: torch.Tensor, + key: torch.Tensor, + weights: torch.Tensor, + query_dequant_scale: torch.Tensor, + key_dequant_scale: torch.Tensor, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + block_table: torch.Tensor | None, + metadata: torch.Tensor | None, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + return_value: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select the compressed key blocks each query attends to (int8 q/k).""" + return torch.ops.xllm_ops.quant_lightning_indexer( + query, + key, + weights, + query_dequant_scale, + key_dequant_scale, + query_quant_mode, + key_quant_mode, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table, + metadata, + layout_query, + layout_key, + sparse_count, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + return_value, + ) + + +def quant_lightning_indexer_metadata( + num_heads_q: int, + num_heads_k: int, + head_dim: int, + query_quant_mode: int, + key_quant_mode: int, + actual_seq_lengths_query: torch.Tensor | None, + actual_seq_lengths_key: torch.Tensor | None, + batch_size: int, + max_seqlen_q: int, + max_seqlen_k: int, + layout_query: str, + layout_key: str, + sparse_count: int, + sparse_mode: int, + pre_tokens: int, + next_tokens: int, + cmp_ratio: int, + device: str, +) -> torch.Tensor: + """Build the AICPU tiling metadata for :func:`quant_lightning_indexer`.""" + return torch.ops.xllm_ops.quant_lightning_indexer_metadata( + num_heads_q, + num_heads_k, + head_dim, + query_quant_mode, + key_quant_mode, + actual_seq_lengths_query, + actual_seq_lengths_key, + batch_size, + max_seqlen_q, + max_seqlen_k, + layout_query, + layout_key, + sparse_count, + sparse_mode, + pre_tokens, + next_tokens, + cmp_ratio, + device, + ) diff --git a/xllm/python/kernels_npu/gated_delta_net.py b/xllm/python/kernels_npu/gated_delta_net.py index 3fa3577a38..f4f19d74df 100644 --- a/xllm/python/kernels_npu/gated_delta_net.py +++ b/xllm/python/kernels_npu/gated_delta_net.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NPU gated-delta-network kernels. +"""NPU gated-delta-network kernels (PyTorch small-op implementation). -None has an NPU kernel yet. The signatures are the contract an NPU -implementation has to meet; see ``kernels_cuda/gated_delta_net.py`` and the -Triton launchers it calls for the reference behaviour. +Implements the same semantics as the CUDA Triton references in +``kernels_cuda/triton/gdn_prefill.py`` and ``kernels_cuda/triton/gated_delta_net.py`` +using only standard PyTorch operations. Performance is not optimized; +correctness and precision alignment are the goals. """ from __future__ import annotations @@ -25,64 +26,97 @@ import torch -GdnPrefillBackend = Literal["flashinfer", "triton"] +GdnPrefillBackend = Literal["pytorch_naive"] def resolve_gdn_prefill_backend( capability: tuple[int, int] | None = None, ) -> GdnPrefillBackend: - """Select the prefill backend of the active device. + """Select the prefill backend for NPU. Args: - capability: Device capability to resolve for; ``None`` reads it from - the current device. + capability: Ignored on NPU. Returns: - The name to pass as ``backend`` to :func:`chunk_gated_delta_rule`. + The backend name to pass to :func:`chunk_gated_delta_rule`. """ del capability - raise NotImplementedError( - "resolve_gdn_prefill_backend has no NPU implementation; gated delta networks are not supported on NPU yet" + return "pytorch_naive" + + +def fused_gdn_gating( + a_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute decay gate g and beta from raw projections via TileLang kernel.""" + from .tilelang.fused_gdn_gating import fused_gdn_gating_kernel_jit + + num_batches, num_heads = a.shape + kernel = fused_gdn_gating_kernel_jit( + num_batches=num_batches, + compile_max_batch=num_batches, + num_heads=num_heads, ) + g_out = torch.empty(1, num_batches, num_heads, dtype=torch.float32, device=a.device) + beta_out = torch.empty(1, num_batches, num_heads, dtype=a.dtype, device=a.device) + kernel( + a_log.to(torch.float32).contiguous(), + a.contiguous(), + b.contiguous(), + dt_bias.to(torch.float32).contiguous(), + g_out.squeeze(0), + beta_out.squeeze(0), + num_batches, + 1.0, # softplus_beta + 20.0, # softplus_threshold + ) + return g_out, beta_out + -def fused_gdn_prefill_post_conv( +def gdn_prefill_prepare( mixed_qkv: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, a: torch.Tensor, b: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor, num_key_heads: int, + num_value_heads: int, key_head_dim: int, value_head_dim: int, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Split the post-convolution projection and build the recurrence gates. +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split + l2norm + gating for prefill. - Args: - mixed_qkv: Packed projection of shape ``[num_tokens, qkv_size]``. - a: Gate projection of shape ``[num_tokens, num_value_heads]``. - b: Beta projection of shape ``[num_tokens, num_value_heads]``. - a_log: Per-head log decay of shape ``[num_value_heads]``. - dt_bias: Per-head timestep bias of shape ``[num_value_heads]``. - num_key_heads: Key heads on this rank. - key_head_dim: Size of one key head. - value_head_dim: Size of one value head. + Encapsulates the NPU-optimal fusion strategy: causal_conv1d_qkv does + conv + split + l2norm in one kernel, then fused_gdn_gating computes + decay and beta independently. Returns: - Query, key, value, the decay gate and beta. + (q, k, v, g, beta) with shapes [T, H, D] / [T, H]. """ - del mixed_qkv, a, b, a_log, dt_bias - del num_key_heads, key_head_dim, value_head_dim - raise NotImplementedError( - "fused_gdn_prefill_post_conv has no NPU kernel; see " - "kernels_cuda/triton/gdn_prefill.py for the reference implementation" + from .causal_conv1d import causal_conv1d_qkv_prefill + + q, k, v = causal_conv1d_qkv_prefill( + mixed_qkv, + weight, + conv_state, + state_indices, + has_initial_state, + cu_seqlens, + num_key_heads, + num_value_heads, + key_head_dim, + value_head_dim, ) + g, beta = fused_gdn_gating(a_log, a, b, dt_bias) + return q.squeeze(0), k.squeeze(0), v.squeeze(0), g.squeeze(0), beta.squeeze(0) def fused_recurrent_gated_delta_rule_packed_decode( @@ -104,17 +138,52 @@ def fused_recurrent_gated_delta_rule_packed_decode( a_log: Per-head log decay of shape ``[num_value_heads]``. dt_bias: Per-head timestep bias of shape ``[num_value_heads]``. initial_state: Recurrent state pool, updated in place. + Shape ``[num_slots, num_value_heads, value_dim, key_dim]``. state_indices: State slot of every sequence. scale: Query scale. Returns: Output of shape ``[batch_size, 1, num_value_heads, value_head_dim]``. """ - del mixed_qkv, a, b, a_log, dt_bias, initial_state, state_indices, scale - raise NotImplementedError( - "fused_recurrent_gated_delta_rule_packed_decode has no NPU kernel; see " - "kernels_cuda/triton/gated_delta_net.py for the reference implementation" + batch = mixed_qkv.shape[0] + num_value_heads, value_dim, key_dim = initial_state.shape[-3:] + qkv_dim = mixed_qkv.shape[1] + query_key_dim = qkv_dim - num_value_heads * value_dim + query_dim = query_key_dim // 2 + num_key_heads = query_dim // key_dim + + # Split mixed_qkv + q_flat = mixed_qkv[:, :query_dim] + k_flat = mixed_qkv[:, query_dim : 2 * query_dim] + v_flat = mixed_qkv[:, 2 * query_dim :] + + q = q_flat.view(batch, num_key_heads, key_dim) + k = k_flat.view(batch, num_key_heads, key_dim) + v = v_flat.view(batch, num_value_heads, value_dim) + + # Kernel expects [batch, seq_len, heads, dim] — add seq dim for decode + q = q.unsqueeze(1) # [batch, 1, num_key_heads, key_dim] + k = k.unsqueeze(1) # [batch, 1, num_key_heads, key_dim] + v = v.unsqueeze(1) # [batch, 1, num_value_heads, value_dim] + + # Kernel does l2norm, gating, GQA expansion, and recurrence internally. + # cu_seqlens for decode: each seq has 1 token. + cu_seqlens = torch.arange(batch + 1, dtype=torch.int32, device=mixed_qkv.device) + + output = torch.ops.xllm_ops.fused_sigmoid_gating_delta_rule_decode( + a_log, + a.unsqueeze(1), + dt_bias, + q.contiguous(), + k.contiguous(), + v.contiguous(), + b.unsqueeze(1), + initial_state, + state_indices, + cu_seqlens, + scale, ) + return output.unsqueeze(1) def chunk_gated_delta_rule( @@ -136,22 +205,34 @@ def chunk_gated_delta_rule( g: Decay gate of shape ``[num_tokens, num_value_heads]``. beta: Beta with the shape of ``g``. initial_state: Recurrent state each sequence starts from. + Shape ``[batch, num_value_heads, value_dim, key_dim]``. cu_seqlens: Cumulative sequence lengths. - backend: Name returned by :func:`resolve_gdn_prefill_backend`. + backend: Ignored on NPU. Returns: The output with the shape of ``v`` and the final recurrent state. """ - del q, k, v, g, beta, initial_state, cu_seqlens, backend - raise NotImplementedError( - "chunk_gated_delta_rule has no NPU kernel; see kernels_cuda/triton/fla/ for the reference implementation" + del backend + # npu_mega_chunk_gdn expects [B, T, H, D] layout with B=1 for packed input + # Cast g and beta to match C++ layer behavior (bf16 round-trip) + g_input = g.to(v.dtype) + beta_input = beta.to(v.dtype) + output, final_state = torch.ops.xllm_ops.chunk_gated_delta_rule( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + g_input.unsqueeze(0), + beta_input.unsqueeze(0), + initial_state, + cu_seqlens, ) + return output.squeeze(0), final_state __all__ = [ "GdnPrefillBackend", "resolve_gdn_prefill_backend", - "fused_gdn_prefill_post_conv", + "gdn_prefill_prepare", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", ] diff --git a/xllm/python/kernels_npu/moe.py b/xllm/python/kernels_npu/moe.py index b5e0b81f2f..6b57559005 100644 --- a/xllm/python/kernels_npu/moe.py +++ b/xllm/python/kernels_npu/moe.py @@ -44,6 +44,72 @@ def _grouped_matmul_swiglu_quant_v2( ) +def dequant_swiglu_quant( + x: torch.Tensor, + weight_scale: torch.Tensor | None, + activation_scale: torch.Tensor | None, + bias: torch.Tensor | None = None, + quant_scale: torch.Tensor | None = None, + quant_offset: torch.Tensor | None = None, + group_index: torch.Tensor | None = None, + activate_left: bool = True, + quant_mode: int = 1, + swiglu_mode: int = 1, + clamp_limit: float = 0.0, + glu_alpha: float = 1.0, + glu_bias: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply the fused dequantization, SwiGLU, and dynamic quantization.""" + return torch.ops.xllm_ops.dequant_swiglu_quant( + x, + weight_scale, + activation_scale, + bias, + quant_scale, + quant_offset, + group_index, + activate_left, + quant_mode, + swiglu_mode, + clamp_limit, + glu_alpha, + glu_bias, + ) + + +def moe_gating_top_k_hash( + x: torch.Tensor, + k: int, + bias: torch.Tensor | None, + input_ids: torch.Tensor | None, + tid2eid: torch.Tensor | None, + k_group: int, + group_count: int, + routed_scaling_factor: float, + eps: float, + group_select_mode: int, + renorm: int, + norm_type: int, + out_flag: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Select experts with the DeepSeek-V4 hash-routing gate.""" + return torch.ops.xllm_ops.moe_gating_top_k_hash( + x, + k, + bias, + input_ids, + tid2eid, + k_group, + group_count, + routed_scaling_factor, + eps, + group_select_mode, + renorm, + norm_type, + out_flag, + ) + + def supports_cutlass_moe(device: torch.device) -> bool: """Return whether ``device`` has the native expert GEMMs. @@ -180,6 +246,159 @@ def grouped_moe( ) +def _group_gemm( + *, + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor | None, + per_token_scale: torch.Tensor | None, + group_list: torch.Tensor, + split_item: int, + group_type: int, + group_list_type: int, + output_dtype: torch.dtype | None, +) -> torch.Tensor: + outputs = torch.ops.npu.npu_grouped_matmul( + x=[x], + weight=[weight], + scale=None if scale is None else [scale], + per_token_scale=None if per_token_scale is None else [per_token_scale], + group_list=group_list, + split_item=split_item, + group_type=group_type, + group_list_type=group_list_type, + output_dtype=output_dtype, + ) + return outputs[0] + + +def _grouped_moe_with_selected_experts_impl( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + """Run grouped quantized experts with pre-computed routing (no gate). + + The routing and W8A8 grouped-matmul sequence mirrors the native NPU + ``FusedMoEImpl::select_experts`` and ``forward_expert`` paths. + """ + num_tokens = hidden_states.shape[0] + expert_num = num_total_experts if num_total_experts > 0 else w13.shape[0] + local_expert_count = num_experts_per_rank if num_experts_per_rank > 0 else w13.shape[0] + active_range = [start_expert_id, start_expert_id + local_expert_count] + if start_expert_id < 0 or active_range[1] > expert_num: + raise ValueError(f"active expert range {active_range} is outside [0, {expert_num})") + if w13.shape[0] != local_expert_count or w2.shape[0] != local_expert_count: + raise ValueError("local expert count must match the first dimension of w13 and w2") + expanded_hidden, expanded_row_idx, expert_tokens, _ = torch_npu.npu_moe_init_routing_v2( + hidden_states, + topk_ids.to(torch.int32), + scale=None, + active_num=num_tokens * topk_ids.size(-1), + expert_num=expert_num, + expert_tokens_num_type=1, + expert_tokens_num_flag=True, + active_expert_range=active_range, + quant_mode=-1, + ) + from xllm.python import kernels as _kernels + + sorted_hidden_i8, pertoken_scale = _kernels.dynamic_quant(expanded_hidden) + if pertoken_scale is None: + raise RuntimeError("dynamic_quant did not return a per-token scale") + if expert_tokens.numel() < local_expert_count: + raise RuntimeError("npu_moe_init_routing_v2 returned fewer groups than local experts") + group_list = expert_tokens[:local_expert_count].to(torch.int64) + gemm1_out = _group_gemm( + x=sorted_hidden_i8, + weight=w13, + scale=None, + per_token_scale=None, + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=1, + output_dtype=torch.int32, + ) + act_i8, act_pt = _kernels.dequant_swiglu_quant( + x=gemm1_out, + weight_scale=w13_scale, + activation_scale=pertoken_scale, + bias=None, + quant_scale=None, + quant_offset=None, + group_index=group_list, + activate_left=True, + quant_mode=1, + swiglu_mode=1, + clamp_limit=swiglu_limit, + glu_alpha=1.0, + glu_bias=0.0, + ) + del w13_offset, w2_offset + output = _group_gemm( + x=act_i8, + weight=w2, + scale=w2_scale.to(hidden_states.dtype), + per_token_scale=act_pt, + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=1, + output_dtype=hidden_states.dtype, + ) + local_mask = (topk_ids >= active_range[0]) & (topk_ids < active_range[1]) + local_topk_weights = topk_weights * local_mask.to(topk_weights.dtype) + return torch_npu.npu_moe_token_unpermute( + permuted_tokens=output, + sorted_indices=expanded_row_idx.abs(), + probs=local_topk_weights.to(output.dtype), + ) + + +@torch.library.custom_op("xllm_python::grouped_moe_with_selected_experts", mutates_args=()) +def grouped_moe_with_selected_experts( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + return _grouped_moe_with_selected_experts_impl( + hidden_states, + topk_weights, + topk_ids, + w13, + w2, + w13_scale, + w2_scale, + w13_offset, + w2_offset, + num_total_experts, + start_expert_id, + num_experts_per_rank, + swiglu_limit, + ) + + @grouped_moe.register_fake def _grouped_moe_fake( hidden_states: torch.Tensor, @@ -213,6 +432,27 @@ def _grouped_moe_fake( return torch.empty_like(hidden_states) +@grouped_moe_with_selected_experts.register_fake +def _grouped_moe_with_selected_experts_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_offset: torch.Tensor | None = None, + w2_offset: torch.Tensor | None = None, + num_total_experts: int = -1, + start_expert_id: int = 0, + num_experts_per_rank: int = -1, + swiglu_limit: float = 0.0, +) -> torch.Tensor: + del topk_weights, topk_ids, w13, w2, w13_scale, w2_scale, w13_offset, w2_offset + del num_total_experts, start_expert_id, num_experts_per_rank, swiglu_limit + return torch.empty_like(hidden_states) + + def moe_fused_topk( gating_output: torch.Tensor, topk: int, @@ -514,11 +754,14 @@ def _moe_gmm2_combine_fake( __all__ = [ + "dequant_swiglu_quant", + "moe_gating_top_k_hash", "supports_cutlass_moe", "prepare_grouped_moe_weights", "grouped_moe", "moe_gate_routing", "moe_expert_compute", + "grouped_moe_with_selected_experts", "moe_fused_topk", "cutlass_fused_moe", "fused_moe", diff --git a/xllm/python/kernels_npu/normalization.py b/xllm/python/kernels_npu/normalization.py index 19164de26a..4daeb57422 100644 --- a/xllm/python/kernels_npu/normalization.py +++ b/xllm/python/kernels_npu/normalization.py @@ -39,6 +39,7 @@ def fused_add_rms_norm_dynamic_quant( output_mask=[True, False], ) return outputs[0], outputs[3], outputs[2] +rms_norm_dynamic_quant = torch.ops.xllm_ops.rms_norm_dynamic_quant def l2_norm(value: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: @@ -51,10 +52,7 @@ def l2_norm(value: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: Returns: A tensor with the shape and dtype of ``value``. """ - del value, eps - raise NotImplementedError( - "l2_norm has no NPU kernel; see kernels_cuda/triton/l2_norm.py for the reference implementation" - ) + return torch.ops.xllm_ops.l2_norm(value, eps) def rms_norm_gated( @@ -63,27 +61,25 @@ def rms_norm_gated( weight: torch.Tensor, eps: float = 1e-6, ) -> torch.Tensor: - """Apply RMSNorm to ``value`` and gate the result with ``gate``. + """Apply RMSNorm to ``value`` and gate the result with ``silu(gate)``. Args: value: Tensor to normalize. - gate: Gate applied after normalization, same shape as ``value``. + gate: Gate applied after normalization (SiLU is applied internally). weight: RMSNorm weight over the last dimension. eps: RMSNorm epsilon. Returns: A tensor with the shape and dtype of ``value``. """ - del value, gate, weight, eps - raise NotImplementedError( - "rms_norm_gated has no NPU kernel; see kernels_cuda/triton/rms_norm.py for the reference implementation" - ) + return torch.ops.xllm_ops.rms_norm_gated(value, gate, weight, eps) __all__ = [ "rms_norm", "fused_add_rms_norm", "fused_add_rms_norm_dynamic_quant", + "rms_norm_dynamic_quant", "l2_norm", "rms_norm_gated", ] diff --git a/xllm/python/kernels_npu/rotary_embedding.py b/xllm/python/kernels_npu/rotary_embedding.py index 6874a54aa9..4d793fcac1 100644 --- a/xllm/python/kernels_npu/rotary_embedding.py +++ b/xllm/python/kernels_npu/rotary_embedding.py @@ -169,9 +169,41 @@ def vision_rotary_mul( return torch_npu.npu_rotary_mul(value.unsqueeze(0).contiguous(), cos_full, sin_full).squeeze(0) +def npu_inplace_partial_rotary_mul( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rope_start_dim: int, + rope_head_dim: int, + inverse: bool = False, +) -> torch.Tensor: + """In-place partial interleaved RoPE on the ``[rope_start_dim:]`` slice. + + Mirrors C++ ``apply_partial_rope`` (deepseek_sparse_attention.cpp:151-190): + x is 3D ``[M, n_head, head_dim]``; cos/sin are 2D ``[M, rope_head_dim]`` + (per-token, no head dim). Reshaped to 4D for the NPU kernel + (``aclnnInplacePartialRotaryMul``, rotary_mode="interleave", + partial_slice=[rope_start_dim, rope_start_dim+rope_head_dim] -- a half-open + range, NOT [start, length]). Modifies x in place. + """ + x4d = x.unsqueeze(2) # [M, n_head, 1, head_dim] + cos4d = cos.view(cos.size(0), 1, 1, cos.size(1)) + sin_cache = -sin if inverse else sin + sin4d = sin_cache.view(sin.size(0), 1, 1, sin.size(1)) + torch.ops.xllm_ops.npu_inplace_partial_rotary_mul( + x4d, + cos4d, + sin4d, + "interleave", + [int(rope_start_dim), int(rope_start_dim + rope_head_dim)], + ) + return x + + __all__ = [ "fused_qk_norm_rope", "interleaved_rotary_embedding", "mrope", "vision_rotary_mul", + "npu_inplace_partial_rotary_mul", ] diff --git a/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py b/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py index aa3815ef7c..8a43ebae94 100644 --- a/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py +++ b/xllm/python/kernels_npu/tilelang/causal_conv1d_decode.py @@ -17,8 +17,6 @@ tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True, } -_decode_kernel_cache = {} - def build_causal_conv1d_decode_kernel( width: int, @@ -168,31 +166,6 @@ def _build_decode_kernel_jit( ) -def get_decode_kernel( - width: int, - dim: int, - dtype_str: str = "bfloat16", - has_silu: bool = True, -) -> torch.nn.Module: - dim_chunks = (dim + DIM_PER_CORE - 1) // DIM_PER_CORE - cache_key = ( - width, - dim_chunks, - DIM_PER_CORE, - dtype_str, - has_silu, - ) - if cache_key not in _decode_kernel_cache: - _decode_kernel_cache[cache_key] = _build_decode_kernel_jit( - width, - dim_chunks, - DIM_PER_CORE, - dtype_str, - has_silu, - ) - return _decode_kernel_cache[cache_key] - - def causal_conv1d_decode( x: torch.Tensor, conv_state: torch.Tensor, @@ -268,7 +241,8 @@ def causal_conv1d_decode( initial_state_mode = torch.ones(batch, dtype=torch.int32, device=conv_state.device) - kernel = get_decode_kernel(width, dim, "bfloat16", has_silu) + dim_chunks = (dim + DIM_PER_CORE - 1) // DIM_PER_CORE + kernel = _build_decode_kernel_jit(width, dim_chunks, DIM_PER_CORE, "bfloat16", has_silu) output = kernel( x_kernel, weight_t, @@ -288,3 +262,31 @@ def causal_conv1d_decode( output = output.to(torch.float16) return output + + +if __name__ == "__main__": + import torch + import torch_npu + + batch = 1 + dim = 8192 + kernel_width = 4 + state_len = kernel_width - 1 # 3 + slots = 202 + + device = "npu:0" + tilelang.disable_cache() + tilelang.cache.clear_cache() + + x = torch.randn(batch, dim, dtype=torch.bfloat16, device=device) + weight = torch.randn(dim, kernel_width, dtype=torch.bfloat16, device=device) + # conv_state in PyTorch convention: [slots, dim, state_len] + conv_state = torch.zeros(slots, dim, state_len, dtype=torch.bfloat16, device=device) + state_indices = torch.tensor([1], dtype=torch.int32, device=device) + + print(f"x={x.shape}, weight={weight.shape}, conv_state={conv_state.shape}, state_indices={state_indices}") + print("Calling causal_conv1d_decode...") + result = causal_conv1d_decode(x, conv_state, weight, conv_state_indices=state_indices) + print(f"result={result.shape}, dtype={result.dtype}") + print(f"result first 5: {result[0, :5].tolist()}") + print("SUCCESS") diff --git a/xllm/python/layers/gated_delta_net.py b/xllm/python/layers/gated_delta_net.py index cd081b396e..02667f71c2 100644 --- a/xllm/python/layers/gated_delta_net.py +++ b/xllm/python/layers/gated_delta_net.py @@ -125,15 +125,25 @@ def _conv_prefill( state_indices: torch.Tensor, has_initial_state: torch.Tensor, cu_seqlens: torch.Tensor, - ) -> torch.Tensor: + a: torch.Tensor, + b: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: self._conv_state_dim_first(conv_state) - return kernels.causal_conv1d_prefill( + return kernels.gdn_prefill_prepare( mixed_qkv, self.conv1d_weight, conv_state, state_indices, has_initial_state, cu_seqlens, + a, + b, + self.A_log, + self.dt_bias, + self.num_k_heads, + self.num_v_heads, + self.key_head_dim, + self.value_head_dim, ) def _conv_decode( @@ -152,17 +162,16 @@ def _conv_decode( def _gdn_prefill( self, - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, ssm_state: torch.Tensor, state_indices: torch.Tensor, has_initial_state: torch.Tensor, cu_seqlens: torch.Tensor, ) -> torch.Tensor: - # TODO: Fuse cache gather/zero/scatter and null-row output masking into - # the GDN kernels. The current staging is intentionally kept for this - # correctness PR and should be removed in the next performance PR. non_null_state = state_indices > 0 use_initial_state = non_null_state & has_initial_state cache_indices = state_indices.to(torch.long) @@ -173,16 +182,6 @@ def _gdn_prefill( initial_state, torch.zeros_like(initial_state), ) - q, k, v, g, beta = kernels.fused_gdn_prefill_post_conv( - mixed_qkv=mixed_qkv, - a=a, - b=b, - a_log=self.A_log, - dt_bias=self.dt_bias, - num_key_heads=self.num_k_heads, - key_head_dim=self.key_head_dim, - value_head_dim=self.value_head_dim, - ) output, final_state = kernels.chunk_gated_delta_rule( q, k, @@ -193,11 +192,12 @@ def _gdn_prefill( cu_seqlens, self.gdn_prefill_backend, ) + num_tokens = q.shape[0] sequence_lengths = cu_seqlens.diff().to(dtype=torch.long) token_mask = torch.repeat_interleave( non_null_state, sequence_lengths, - output_size=mixed_qkv.shape[0], + output_size=num_tokens, ) output = torch.where(token_mask[:, None, None], output, 0.0) ssm_state.index_copy_( @@ -256,21 +256,25 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor: has_initial_state = has_initial_state.to(device=hidden.device, dtype=torch.bool) if has_initial_state.shape != state_indices.shape: raise ValueError("has_initial_state must match linear_state_indices") - mixed_qkv = self._conv_prefill( + q, k, v, g, beta = self._conv_prefill( mixed_qkv, conv_state, state_indices, has_initial_state, cu_seqlens, + a, + b, ) else: mixed_qkv = self._conv_decode(mixed_qkv, conv_state, state_indices) if is_prefill: output = self._gdn_prefill( - mixed_qkv, - a, - b, + q, + k, + v, + g, + beta, ssm_state, state_indices, has_initial_state, diff --git a/xllm/python/model_platform_support.py b/xllm/python/model_platform_support.py index caae712526..711d648c90 100644 --- a/xllm/python/model_platform_support.py +++ b/xllm/python/model_platform_support.py @@ -16,7 +16,7 @@ MODEL_PLATFORM_SUPPORT: dict[str, dict[str, bool]] = { "qwen3": {"cuda": True, "npu": True}, - "qwen3_5": {"cuda": True, "npu": False}, + "qwen3_5": {"cuda": True, "npu": True}, "qwen3_vl": {"cuda": False, "npu": True}, "deepseek_v32": {"cuda": False, "npu": True}, "glm5_2": {"cuda": False, "npu": True}, diff --git a/xllm/python/models/glm5_2.py b/xllm/python/models/glm5_2.py index dc7c4bddc2..2cddedc6fe 100644 --- a/xllm/python/models/glm5_2.py +++ b/xllm/python/models/glm5_2.py @@ -646,13 +646,11 @@ def load_weights( self.model.layers[i].mlp.process_weights_after_loading() else: se = p + "mlp.experts." + moe_layer = self.model.layers[i].mlp + moe_layer.allocate_experts_w13_for_loading() w13_param = self.get_parameter(p + "mlp.experts_w13") - w2_param = self.get_parameter(p + "mlp.experts_w2") w13_scale = self.get_buffer(p + "mlp.experts_w13_scale") w13_offset = self.get_buffer(p + "mlp.experts_w13_offset") - w2_scale = self.get_buffer(p + "mlp.experts_w2_scale") - w2_offset = self.get_buffer(p + "mlp.experts_w2_offset") - moe_layer = self.model.layers[i].mlp expert_start = moe_layer.local_expert_start expert_end = moe_layer.local_expert_end shard_world = cfg.moe_tp_size if cfg.ep_size > 1 else cfg.tp_size @@ -665,9 +663,6 @@ def load_weights( uw = loader.load_tensor(se + f"{j}.up_proj.weight") us = loader.load_tensor(se + f"{j}.up_proj.weight_scale") uo = loader.load_tensor(se + f"{j}.up_proj.weight_offset") - dw = loader.load_tensor(se + f"{j}.down_proj.weight") - ds = loader.load_tensor(se + f"{j}.down_proj.weight_scale") - do = loader.load_tensor(se + f"{j}.down_proj.weight_offset") w13_param.data[local_idx].copy_( torch.cat( [ @@ -695,6 +690,16 @@ def load_weights( dim=0, ).contiguous() ) + + moe_layer.allocate_experts_w2_for_loading() + w2_param = self.get_parameter(p + "mlp.experts_w2") + w2_scale = self.get_buffer(p + "mlp.experts_w2_scale") + w2_offset = self.get_buffer(p + "mlp.experts_w2_offset") + for j in range(expert_start, expert_end): + local_idx = j - expert_start + dw = loader.load_tensor(se + f"{j}.down_proj.weight") + ds = loader.load_tensor(se + f"{j}.down_proj.weight_scale") + do = loader.load_tensor(se + f"{j}.down_proj.weight_offset") w2_param.data[local_idx].copy_(loader.shard(dw, 1, shard_world, shard_rank).contiguous()) w2_scale.data[local_idx].copy_(ds.contiguous()) w2_offset.data[local_idx].copy_(do.contiguous()) diff --git a/xllm/python/models/qwen3_5.py b/xllm/python/models/qwen3_5.py index 0c2c12d1d5..6fd43452de 100644 --- a/xllm/python/models/qwen3_5.py +++ b/xllm/python/models/qwen3_5.py @@ -418,6 +418,13 @@ def __init__(self, cfg: Qwen3_5Config, dtype: torch.dtype, device: torch.device) self.norm = GemmaRMSNorm(cfg.hidden_size, cfg.rms_norm_eps, dtype=dtype, device=device) def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # TODO: The current tilelang-ascend version has a cache bug that prevents kernels with + # dynamic symbols from being cached, causing the service to crash. This is a temporary + # workaround; we will resubmit once tilelang-ascend fixes the issue. + import tilelang + + tilelang.disable_cache() + tilelang.cache.clear_cache() hidden = self.embed_tokens(input_ids) residual: torch.Tensor | None = None for layer in self.layers: