diff --git a/.gitignore b/.gitignore index 1bd6c0306045..9b5a47c5deeb 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,8 @@ __pycache__/ *.cache *.nsys-rep *.npy +*.so +*.whl .VSCodeCounter cpp/build* cpp/Release diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp index 35b08198fd74..ee8d645bceb2 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp @@ -1,5 +1,6 @@ #include "tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h" #include "tensorrt_llm/common/logger.h" +#include "tensorrt_llm/common/memoryUtils.h" #include #include #include @@ -7,6 +8,9 @@ #include #include +namespace tc = tensorrt_llm::common; +using namespace tensorrt_llm::runtime; + namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -142,4 +146,30 @@ CUresult copyHostToHost( return cuLaunchHostFunc(stream, hostFnHostToHostCopy, data); } +void copyBatchBlockOffsets(ITensor& output, SizeType32 batchSize, std::vector const& batchBlockIndices, + SizeType32 numPools, SizeType32 offset) noexcept +{ + auto* dstPtr = bufferCast(output); + auto const& dstShape = output.getShape(); + + for (SizeType32 poolIdx = 0; poolIdx < numPools; poolIdx++) + { + for (SizeType32 batchIdx = 0; batchIdx < batchSize; batchIdx++) + { + auto const& blockIndices = batchBlockIndices[poolIdx * batchSize + batchIdx]; + auto const addr = reinterpret_cast(blockIndices.addr); + auto const length = blockIndices.length; + auto const dstIndexK = tc::flat_index(dstShape.d, poolIdx, batchIdx + offset, 0, 0); + auto const dstIndexV = tc::flat_index(dstShape.d, poolIdx, batchIdx + offset, 1, 0); + auto addrK = dstPtr + dstIndexK; + auto addrV = dstPtr + dstIndexV; + for (SizeType32 i = 0; i < length; i++) + { + addrK[i] = addr[i]; + addrV[i] = addr[i] + 1; + } + } + } +} + } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h index 4a886fdc9a1c..f26ea0beae39 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h @@ -17,10 +17,17 @@ #pragma once +#include "tensorrt_llm/kernels/kvCacheIndex.h" +#include "tensorrt_llm/runtime/iBuffer.h" +#include "tensorrt_llm/runtime/iTensor.h" #include #include #include +namespace tk = tensorrt_llm::kernels; +using SizeType32 = tensorrt_llm::runtime::SizeType32; +using ITensor = tensorrt_llm::runtime::ITensor; + namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { struct DiskAddress @@ -38,6 +45,12 @@ struct Task SrcAddr src; }; +struct BlockIndices +{ + MemAddress addr; + SizeType32 length; +}; + CUresult copyDiskToDisk( std::vector> const& tasks, ssize_t numBytes, CUstream stream) noexcept; CUresult copyDiskToHost( @@ -52,4 +65,8 @@ CUresult copyDeviceToHost( std::vector> const& tasks, ssize_t numBytes, CUstream stream) noexcept; CUresult copyDeviceToDevice( std::vector> const& tasks, ssize_t numBytes, CUstream stream) noexcept; + +void copyBatchBlockOffsets(ITensor& output, SizeType32 batchSize, std::vector const& batchBlockIndices, + SizeType32 numPools, SizeType32 offset) noexcept; + } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp index c2dc730b8061..4751a12f6bdf 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp @@ -17,14 +17,31 @@ #include "kvCacheManagerV2Utils.h" #include "tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h" +#include "tensorrt_llm/runtime/iTensor.h" +#include "tensorrt_llm/runtime/torchView.h" +#include #include +#include #include +#include +namespace tr = tensorrt_llm::runtime; namespace nb = nanobind; +using SizeType32 = tensorrt_llm::runtime::SizeType32; + namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { +std::optional from_torch(std::optional torchPtr) +{ + if (torchPtr) + { + return tr::TorchView::of(torchPtr.value()); + } + return std::nullopt; +} + void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) { // Bind DiskAddress struct @@ -54,6 +71,11 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) .def_rw("dst", &Task::dst) .def_rw("src", &Task::src); + nb::class_(module, "BlockIndices") + .def(nb::init(), nb::arg("addr"), nb::arg("length")) + .def_rw("addr", &BlockIndices::addr) + .def_rw("length", &BlockIndices::length); + // Bind copy functions module.def( "copy_disk_to_disk", @@ -103,6 +125,18 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) { return copyDeviceToDevice(tasks, numBytes, reinterpret_cast(stream)); }, nb::arg("tasks"), nb::arg("num_bytes"), nb::arg("stream"), nb::call_guard(), "Copy data from device to device using CUDA kernels"); + + module.def( + "copy_batch_block_offsets", + [](at::Tensor output, SizeType32 batchSize, std::vector const& batchBlockIndices, + SizeType32 numPools, SizeType32 offset) + { + auto _output = from_torch(output); + TLLM_CHECK_WITH_INFO(_output.has_value(), "Invalid output tensor."); + copyBatchBlockOffsets(*(_output.value()), batchSize, batchBlockIndices, numPools, offset); + }, + nb::arg("output"), nb::arg("batch_size"), nb::arg("batch_block_indices"), nb::arg("num_pools"), + nb::arg("offset"), nb::call_guard(), "Copy batch block indices to output tensor"); } } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.h b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.h index 7da561a98b5b..e73166491b8b 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.h +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.h @@ -17,6 +17,7 @@ #pragma once +#include "tensorrt_llm/nanobind/common/customCasters.h" #include namespace nb = nanobind; diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp index ae4936a4df8e..6377f3a3c5fd 100644 --- a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp @@ -63,15 +63,15 @@ void initBindings(nb::module_& m) new (&self) tle::DecodingMode(nb::cast(state[0])); }; nb::class_(m, "DecodingMode") - .def("Auto", &tle::DecodingMode::Auto) - .def("TopK", &tle::DecodingMode::TopK) - .def("TopP", &tle::DecodingMode::TopP) - .def("TopKTopP", &tle::DecodingMode::TopKTopP) - .def("BeamSearch", &tle::DecodingMode::BeamSearch) - .def("Medusa", &tle::DecodingMode::Medusa) - .def("Lookahead", &tle::DecodingMode::Lookahead) - .def("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens) - .def("Eagle", &tle::DecodingMode::Eagle) + .def_static("Auto", &tle::DecodingMode::Auto) + .def_static("TopK", &tle::DecodingMode::TopK) + .def_static("TopP", &tle::DecodingMode::TopP) + .def_static("TopKTopP", &tle::DecodingMode::TopKTopP) + .def_static("BeamSearch", &tle::DecodingMode::BeamSearch) + .def_static("Medusa", &tle::DecodingMode::Medusa) + .def_static("Lookahead", &tle::DecodingMode::Lookahead) + .def_static("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens) + .def_static("Eagle", &tle::DecodingMode::Eagle) .def("isAuto", &tle::DecodingMode::isAuto) .def("isTopK", &tle::DecodingMode::isTopK) .def("isTopP", &tle::DecodingMode::isTopP) diff --git a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp index 700133a46b2a..88e3a582e4f1 100644 --- a/cpp/tensorrt_llm/nanobind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/thop/bindings.cpp @@ -37,38 +37,32 @@ void initBindings(nb::module_& m) m.def("attention", &torch_ext::attention, // Parameters with default values using std::nullopt for optional arguments - nb::arg("q"), nb::arg("k") = std::nullopt, nb::arg("v") = std::nullopt, nb::arg("output"), - nb::arg("output_sf") = std::nullopt, nb::arg("out_dtype") = std::nullopt, nb::arg("workspace_") = std::nullopt, - nb::arg("sequence_length"), nb::arg("host_past_key_value_lengths"), nb::arg("host_total_kv_lens"), - nb::arg("context_lengths"), nb::arg("host_context_lengths"), nb::arg("host_request_types"), - nb::arg("kv_cache_block_offsets") = std::nullopt, nb::arg("host_kv_cache_block_offsets") = std::nullopt, - nb::arg("host_kv_cache_pool_pointers") = std::nullopt, nb::arg("host_kv_cache_pool_mapping") = std::nullopt, - nb::arg("cache_indirection") = std::nullopt, nb::arg("kv_scale_orig_quant") = std::nullopt, - nb::arg("kv_scale_quant_orig") = std::nullopt, nb::arg("out_scale") = std::nullopt, - nb::arg("rotary_inv_freq") = std::nullopt, nb::arg("rotary_cos_sin") = std::nullopt, - nb::arg("latent_cache") = std::nullopt, nb::arg("q_pe") = std::nullopt, - nb::arg("block_ids_per_seq") = std::nullopt, nb::arg("attention_sinks") = std::nullopt, nb::arg("is_fused_qkv"), - nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"), nb::arg("layer_idx"), nb::arg("num_heads"), - nb::arg("num_kv_heads"), nb::arg("head_size"), nb::arg("tokens_per_block") = std::nullopt, + nb::arg("q"), nb::arg("k").none(), nb::arg("v").none(), nb::arg("output"), nb::arg("output_sf").none(), + nb::arg("out_dtype").none(), nb::arg("workspace_").none(), nb::arg("sequence_length"), + nb::arg("host_past_key_value_lengths"), nb::arg("host_total_kv_lens"), nb::arg("context_lengths"), + nb::arg("host_context_lengths"), nb::arg("host_request_types"), nb::arg("kv_cache_block_offsets").none(), + nb::arg("host_kv_cache_block_offsets").none(), nb::arg("host_kv_cache_pool_pointers").none(), + nb::arg("host_kv_cache_pool_mapping").none(), nb::arg("cache_indirection").none(), + nb::arg("kv_scale_orig_quant").none(), nb::arg("kv_scale_quant_orig").none(), nb::arg("out_scale").none(), + nb::arg("rotary_inv_freq").none(), nb::arg("rotary_cos_sin").none(), nb::arg("latent_cache").none(), + nb::arg("q_pe").none(), nb::arg("block_ids_per_seq").none(), nb::arg("attention_sinks").none(), + nb::arg("is_fused_qkv"), nb::arg("update_kv_cache"), nb::arg("predicted_tokens_per_seq"), nb::arg("layer_idx"), + nb::arg("num_heads"), nb::arg("num_kv_heads"), nb::arg("head_size"), nb::arg("tokens_per_block").none(), nb::arg("max_num_requests"), nb::arg("max_context_length"), nb::arg("attention_window_size"), nb::arg("sink_token_length"), nb::arg("beam_width"), nb::arg("mask_type"), nb::arg("quant_mode"), nb::arg("q_scaling"), nb::arg("position_embedding_type"), nb::arg("rotary_embedding_dim"), nb::arg("rotary_embedding_base"), nb::arg("rotary_embedding_scale_type"), nb::arg("rotary_embedding_scales"), nb::arg("rotary_embedding_max_position_info"), nb::arg("use_paged_context_fmha"), - nb::arg("attention_input_type") = std::nullopt, nb::arg("is_mla_enable"), - nb::arg("chunked_prefill_buffer_batch_size") = std::nullopt, nb::arg("q_lora_rank") = std::nullopt, - nb::arg("kv_lora_rank") = std::nullopt, nb::arg("qk_nope_head_dim") = std::nullopt, - nb::arg("qk_rope_head_dim") = std::nullopt, nb::arg("v_head_dim") = std::nullopt, - nb::arg("mrope_rotary_cos_sin") = std::nullopt, nb::arg("mrope_position_deltas") = std::nullopt, - nb::arg("mla_tensor_params"), nb::arg("attention_chunk_size") = std::nullopt, - nb::arg("softmax_stats_tensor") = std::nullopt, nb::arg("spec_decoding_bool_params"), - nb::arg("spec_decoding_tensor_params"), nb::arg("sparse_kv_indices") = std::nullopt, - nb::arg("sparse_kv_offsets") = std::nullopt, nb::arg("sparse_attn_indices") = std::nullopt, - nb::arg("sparse_attn_offsets") = std::nullopt, nb::arg("sparse_attn_indices_block_size"), - nb::arg("sparse_mla_topk") = std::nullopt, nb::arg("cu_q_seqlens") = std::nullopt, - nb::arg("cu_kv_seqlens") = std::nullopt, nb::arg("fmha_scheduler_counter") = std::nullopt, - nb::arg("mla_bmm1_scale") = std::nullopt, nb::arg("mla_bmm2_scale") = std::nullopt, - nb::arg("quant_q_buffer") = std::nullopt, "Multi-head attention operation", - nb::call_guard()); + nb::arg("attention_input_type").none(), nb::arg("is_mla_enable"), + nb::arg("chunked_prefill_buffer_batch_size").none(), nb::arg("q_lora_rank").none(), + nb::arg("kv_lora_rank").none(), nb::arg("qk_nope_head_dim").none(), nb::arg("qk_rope_head_dim").none(), + nb::arg("v_head_dim").none(), nb::arg("mrope_rotary_cos_sin").none(), nb::arg("mrope_position_deltas").none(), + nb::arg("mla_tensor_params"), nb::arg("attention_chunk_size").none(), nb::arg("softmax_stats_tensor").none(), + nb::arg("spec_decoding_bool_params"), nb::arg("spec_decoding_tensor_params"), + nb::arg("sparse_kv_indices").none(), nb::arg("sparse_kv_offsets").none(), nb::arg("sparse_attn_indices").none(), + nb::arg("sparse_attn_offsets").none(), nb::arg("sparse_attn_indices_block_size"), + nb::arg("sparse_mla_topk").none(), nb::arg("cu_q_seqlens").none(), nb::arg("cu_kv_seqlens").none(), + nb::arg("fmha_scheduler_counter").none(), nb::arg("mla_bmm1_scale").none(), nb::arg("mla_bmm2_scale").none(), + nb::arg("quant_q_buffer").none(), "Multi-head attention operation", nb::call_guard()); } } // namespace tensorrt_llm::nanobind::thop diff --git a/cpp/tensorrt_llm/pybind/executor/bindings.cpp b/cpp/tensorrt_llm/pybind/executor/bindings.cpp index bbb843bedba2..9a043f1e8492 100644 --- a/cpp/tensorrt_llm/pybind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/executor/bindings.cpp @@ -64,15 +64,15 @@ void initBindings(pybind11::module_& m) return tle::DecodingMode(state[0].cast()); }; py::class_(m, "DecodingMode") - .def("Auto", &tle::DecodingMode::Auto) - .def("TopK", &tle::DecodingMode::TopK) - .def("TopP", &tle::DecodingMode::TopP) - .def("TopKTopP", &tle::DecodingMode::TopKTopP) - .def("BeamSearch", &tle::DecodingMode::BeamSearch) - .def("Medusa", &tle::DecodingMode::Medusa) - .def("Lookahead", &tle::DecodingMode::Lookahead) - .def("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens) - .def("Eagle", &tle::DecodingMode::Eagle) + .def_static("Auto", &tle::DecodingMode::Auto) + .def_static("TopK", &tle::DecodingMode::TopK) + .def_static("TopP", &tle::DecodingMode::TopP) + .def_static("TopKTopP", &tle::DecodingMode::TopKTopP) + .def_static("BeamSearch", &tle::DecodingMode::BeamSearch) + .def_static("Medusa", &tle::DecodingMode::Medusa) + .def_static("Lookahead", &tle::DecodingMode::Lookahead) + .def_static("ExplicitDraftTokens", &tle::DecodingMode::ExplicitDraftTokens) + .def_static("Eagle", &tle::DecodingMode::Eagle) .def("isAuto", &tle::DecodingMode::isAuto) .def("isTopK", &tle::DecodingMode::isTopK) .def("isTopP", &tle::DecodingMode::isTopP) diff --git a/cpp/tensorrt_llm/pybind/thop/bindings.cpp b/cpp/tensorrt_llm/pybind/thop/bindings.cpp index 0f8bffefea9b..bb4ab514d415 100644 --- a/cpp/tensorrt_llm/pybind/thop/bindings.cpp +++ b/cpp/tensorrt_llm/pybind/thop/bindings.cpp @@ -37,38 +37,36 @@ void initBindings(pybind11::module_& m) m.def("attention", &torch_ext::attention, // Parameters with default values using std::nullopt for optional arguments - py::arg("q"), py::arg("k") = std::nullopt, py::arg("v") = std::nullopt, py::arg("output"), - py::arg("output_sf") = std::nullopt, py::arg("out_dtype") = std::nullopt, py::arg("workspace_") = std::nullopt, + py::arg("q"), py::arg("k").none(true), py::arg("v").none(true), py::arg("output"), + py::arg("output_sf").none(true), py::arg("out_dtype").none(true), py::arg("workspace_").none(true), py::arg("sequence_length"), py::arg("host_past_key_value_lengths"), py::arg("host_total_kv_lens"), py::arg("context_lengths"), py::arg("host_context_lengths"), py::arg("host_request_types"), - py::arg("kv_cache_block_offsets") = std::nullopt, py::arg("host_kv_cache_block_offsets") = std::nullopt, - py::arg("host_kv_cache_pool_pointers") = std::nullopt, py::arg("host_kv_cache_pool_mapping") = std::nullopt, - py::arg("cache_indirection") = std::nullopt, py::arg("kv_scale_orig_quant") = std::nullopt, - py::arg("kv_scale_quant_orig") = std::nullopt, py::arg("out_scale") = std::nullopt, - py::arg("rotary_inv_freq") = std::nullopt, py::arg("rotary_cos_sin") = std::nullopt, - py::arg("latent_cache") = std::nullopt, py::arg("q_pe") = std::nullopt, - py::arg("block_ids_per_seq") = std::nullopt, py::arg("attention_sinks") = std::nullopt, py::arg("is_fused_qkv"), - py::arg("update_kv_cache"), py::arg("predicted_tokens_per_seq"), py::arg("layer_idx"), py::arg("num_heads"), - py::arg("num_kv_heads"), py::arg("head_size"), py::arg("tokens_per_block") = std::nullopt, + py::arg("kv_cache_block_offsets").none(true), py::arg("host_kv_cache_block_offsets").none(true), + py::arg("host_kv_cache_pool_pointers").none(true), py::arg("host_kv_cache_pool_mapping").none(true), + py::arg("cache_indirection").none(true), py::arg("kv_scale_orig_quant").none(true), + py::arg("kv_scale_quant_orig").none(true), py::arg("out_scale").none(true), + py::arg("rotary_inv_freq").none(true), py::arg("rotary_cos_sin").none(true), py::arg("latent_cache").none(true), + py::arg("q_pe").none(true), py::arg("block_ids_per_seq").none(true), py::arg("attention_sinks").none(true), + py::arg("is_fused_qkv"), py::arg("update_kv_cache"), py::arg("predicted_tokens_per_seq"), py::arg("layer_idx"), + py::arg("num_heads"), py::arg("num_kv_heads"), py::arg("head_size"), py::arg("tokens_per_block").none(true), py::arg("max_num_requests"), py::arg("max_context_length"), py::arg("attention_window_size"), py::arg("sink_token_length"), py::arg("beam_width"), py::arg("mask_type"), py::arg("quant_mode"), py::arg("q_scaling"), py::arg("position_embedding_type"), py::arg("rotary_embedding_dim"), py::arg("rotary_embedding_base"), py::arg("rotary_embedding_scale_type"), py::arg("rotary_embedding_scales"), py::arg("rotary_embedding_max_position_info"), py::arg("use_paged_context_fmha"), - py::arg("attention_input_type") = std::nullopt, py::arg("is_mla_enable"), - py::arg("chunked_prefill_buffer_batch_size") = std::nullopt, py::arg("q_lora_rank") = std::nullopt, - py::arg("kv_lora_rank") = std::nullopt, py::arg("qk_nope_head_dim") = std::nullopt, - py::arg("qk_rope_head_dim") = std::nullopt, py::arg("v_head_dim") = std::nullopt, - py::arg("mrope_rotary_cos_sin") = std::nullopt, py::arg("mrope_position_deltas") = std::nullopt, - py::arg("mla_tensor_params"), py::arg("attention_chunk_size") = std::nullopt, - py::arg("softmax_stats_tensor") = std::nullopt, py::arg("spec_decoding_bool_params"), - py::arg("spec_decoding_tensor_params"), py::arg("sparse_kv_indices") = std::nullopt, - py::arg("sparse_kv_offsets") = std::nullopt, py::arg("sparse_attn_indices") = std::nullopt, - py::arg("sparse_attn_offsets") = std::nullopt, py::arg("sparse_attn_indices_block_size"), - py::arg("sparse_mla_topk") = std::nullopt, py::arg("cu_q_seqlens") = std::nullopt, - py::arg("cu_kv_seqlens") = std::nullopt, py::arg("fmha_scheduler_counter") = std::nullopt, - py::arg("mla_bmm1_scale") = std::nullopt, py::arg("mla_bmm2_scale") = std::nullopt, - py::arg("quant_q_buffer") = std::nullopt, "Multi-head attention operation", + py::arg("attention_input_type").none(true), py::arg("is_mla_enable"), + py::arg("chunked_prefill_buffer_batch_size").none(true), py::arg("q_lora_rank").none(true), + py::arg("kv_lora_rank").none(true), py::arg("qk_nope_head_dim").none(true), + py::arg("qk_rope_head_dim").none(true), py::arg("v_head_dim").none(true), + py::arg("mrope_rotary_cos_sin").none(true), py::arg("mrope_position_deltas").none(true), + py::arg("mla_tensor_params"), py::arg("attention_chunk_size").none(true), + py::arg("softmax_stats_tensor").none(true), py::arg("spec_decoding_bool_params"), + py::arg("spec_decoding_tensor_params"), py::arg("sparse_kv_indices").none(true), + py::arg("sparse_kv_offsets").none(true), py::arg("sparse_attn_indices").none(true), + py::arg("sparse_attn_offsets").none(true), py::arg("sparse_attn_indices_block_size"), + py::arg("sparse_mla_topk").none(true), py::arg("cu_q_seqlens").none(true), py::arg("cu_kv_seqlens").none(true), + py::arg("fmha_scheduler_counter").none(true), py::arg("mla_bmm1_scale").none(true), + py::arg("mla_bmm2_scale").none(true), py::arg("quant_q_buffer").none(true), "Multi-head attention operation", py::call_guard()); } } // namespace tensorrt_llm::pybind::thop diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 9b37f8c7b296..550c9c912f66 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -8,9 +8,12 @@ TorchCompileConfig) example_prompts = [ - "Hello, my name is", + "Hello, my name is John and I am a software engineer. I have experience in developing web applications using JavaScript and React. I am also familiar with Python and Django. I am a quick learner and enjoy working in a team environment.", + "Hello, my name is John and I am a software engineer. I have experience in developing web applications using JavaScript and React. I am also familiar with Python and Django. I am a quick learner and enjoy working in a team environment.", "The capital of France is", "The future of AI is", + "Hi, how are you?", + "What is the capital of the United States?", ] @@ -85,6 +88,9 @@ def add_llm_args(parser): default=False, action='store_true') parser.add_argument("--tokens_per_block", type=int, default=32) + parser.add_argument('--use_kv_cache_manager_v2', + default=False, + action='store_true') # Runtime parser.add_argument('--disable_overlap_scheduler', @@ -182,6 +188,7 @@ def setup_llm(args, **kwargs): free_gpu_memory_fraction=args.kv_cache_fraction, dtype=args.kv_cache_dtype, tokens_per_block=args.tokens_per_block, + use_kv_cache_manager_v2=args.use_kv_cache_manager_v2, ) spec_decode_algo = args.spec_decode_algo.upper( diff --git a/mypyclib/mypyclib/__init__.py b/mypyclib/mypyclib/__init__.py new file mode 100644 index 000000000000..0746fd3bffe9 --- /dev/null +++ b/mypyclib/mypyclib/__init__.py @@ -0,0 +1,17 @@ +from .sampler_utils import ( + _apply_embedding_bias_impl, + _group_requests_by_strategy_key_impl, + _request_get_sampling_params_impl, + _request_strategy_impl, + _speculation_could_use_rejection_sampling_impl, + resolve_sampling_strategy_impl, +) + +__all__ = [ + "resolve_sampling_strategy_impl", + "_request_get_sampling_params_impl", + "_group_requests_by_strategy_key_impl", + "_request_strategy_impl", + "_speculation_could_use_rejection_sampling_impl", + "_apply_embedding_bias_impl", +] diff --git a/mypyclib/mypyclib/sampler_utils.py b/mypyclib/mypyclib/sampler_utils.py new file mode 100644 index 000000000000..69b90ac77ea0 --- /dev/null +++ b/mypyclib/mypyclib/sampler_utils.py @@ -0,0 +1,213 @@ +import sys +from collections import defaultdict +from collections.abc import Iterable +from itertools import repeat +from typing import Callable, Optional, TypeVar, cast + +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.sampling_utils import ( + GREEDY, + GenericStrategyKeyType, + Strategy, + UtilsSamplingParams, +) +from tensorrt_llm.sampling_params import SamplingParams + +if sys.version_info[:2] >= (3, 12): + pass +else: + pass + +T = TypeVar("T") + + +def resolve_sampling_strategy_impl(params: UtilsSamplingParams, *, vocab_size: int) -> Strategy: + # The semantics are specified in the doc-string of SamplingParams + + temperature = params.temperature + top_p = params.top_p + top_k = params.top_k + + if SamplingParams.params_imply_greedy_decoding( + temperature=temperature, + top_p=top_p, + top_k=top_k, + ): + return GREEDY + + # --- resolving default values + # NB: not greedy, hence temperature != 0 if specified + temperature = temperature or 1.0 + + # NB: not greedy, hence top_p != 0 if specified + top_p = top_p or 1.0 + # NB: not greedy, hence top_k != 1 if specified + # (0 and vocab_size are equivalent) + top_k = top_k or vocab_size + + assert top_k > 1, "non-greedy sampling requires valid top_k" + need_top_k = top_k < vocab_size + assert top_p > 0, "non-greedy sampling requires valid top_p" + need_top_p = top_p < 1 + + if need_top_p: + if need_top_k: + return ("top_k_top_p", top_k, top_p, temperature) + return ("top_p", top_p, temperature) + if need_top_k: + return ("top_k", top_k, temperature) + return ("temperature", temperature) + + +# Due to tensorrt_llm::runtime::SamplingConfig using vectors, params +# in LlmRequest.sampling_params are either None or single-element lists. +# This helper method simplifies code using such params. +def _unwrap_singleton(p: Optional[list[T]]) -> Optional[T]: + if p is None: + return None + (t,) = p + return t + + +def _request_get_sampling_params_impl(request: LlmRequest) -> UtilsSamplingParams: + sampling_config = request.sampling_config + temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) + top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) + top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) + + return UtilsSamplingParams( + temperature=temperature, + top_p=top_p, + top_k=top_k, + ) + + +def _request_strategy_impl(request: LlmRequest, *, vocab_size: int) -> Strategy: + params = _request_get_sampling_params_impl(request) + return resolve_sampling_strategy_impl(params, vocab_size=vocab_size) + + +def _speculation_could_use_rejection_sampling_impl( + request: LlmRequest, strategy: Optional[Strategy] = None +) -> bool: + if strategy is None: + strategy = _request_strategy_impl( + request, + vocab_size=2**31, # vocab_size does not affect greediness + ) + return get_draft_token_length(request) > 0 and strategy != GREEDY + + +def _group_requests_by_strategy_key_impl( + requests: Iterable[LlmRequest], + *, + strategy_to_key: Callable[[Strategy, bool], GenericStrategyKeyType], + pin_memory: bool = False, + vocab_size: int, +) -> dict[tuple[GenericStrategyKeyType, bool], tuple[torch.Tensor, list[Strategy]]]: + # NB: Client code relies on request indices in returned torch.Tensor being sorted. + group_dict: dict[tuple[GenericStrategyKeyType, bool], tuple[list[int], list[Strategy]]] = ( + defaultdict(lambda: ([], [])) + ) + + for req_index, req in enumerate(requests): + strategy = _request_strategy_impl(req, vocab_size=vocab_size) + speculation_needs_probs = ( + # NB: This criterion needs to be consistent with the gating of rejection sampling in + # process_draft_tokens. + _speculation_could_use_rejection_sampling_impl(req, strategy) + ) + strategy_key = strategy_to_key(strategy, speculation_needs_probs) + group_dict_entry = group_dict[(strategy_key, speculation_needs_probs)] + group_dict_entry[0].append(req_index) + group_dict_entry[1].append(strategy) + return { + group_key: ( + torch.tensor(indices, pin_memory=pin_memory, dtype=torch.int32), + strategies, + ) + for group_key, (indices, strategies) in group_dict.items() + } + + +def _apply_embedding_bias_impl( + logits: torch.Tensor, + requests: list[LlmRequest], + request_steps: torch.Tensor, +) -> None: + """Apply embedding bias (aka logit bias) to logits. + + Arguments: + request_steps: Number of steps/tokens for each request. + + Modifies logits in-place. + """ + # NB: Unfortunately, Torch provides no combination of torch.index_select (similar to + # torch.Tensor.gather -- allows one-to-many mapping) and addition, analogous to how + # torch.Tensor.scatter_add_ (and it's variant torch.Tensor.index_add_ -- allows + # many-to-one mapping) combine addition with torch.Tensor.scatter_. + # + # Notwithstanding the previous point, there are two options: + # (i) materialize a permuted bias tensor with repeated consecutive rows via + # torch.repeat_interleave and then use torch.Tensor.index_add_ (poor write + # locality / risk of false sharing) + # (ii) materialize the correctly ordered bias tensor via torch.index_select and then + # perform a masked addition (poor read locality for request batches randomly + # mixing uniform and heterogeneous bias tensors, i.e., mixing slices with high + # and low reuse). + # Since read-caching is expected to help in typical cases, option (ii) is implemented here. + + # Track which logits require logit bias application + logits_bias_mask = torch.zeros((logits.size(0),), dtype=torch.bool, pin_memory=True) + + _next_bias_index = 0 + + def provision_bias_index() -> int: + nonlocal _next_bias_index + bias_index = _next_bias_index + _next_bias_index += 1 + return bias_index + + # Indices of unique bias tensors + # + # NB: hash(torch.Tensor) is equivalent to id(torch.Tensor), and does not + # depend on tensor contents, cf. https://github.com/pytorch/pytorch/issues/2569 + bias_to_index: dict[torch.Tensor, int] = defaultdict(provision_bias_index) + + # Source indices for bias application + bias_gather_indices: list[int] = [] + + # Collect bias information + req_bias = None + for i, (req, steps) in enumerate(zip(requests, request_steps)): + steps = int(steps.item()) + req_bias = req._py_embedding_bias_1d + if req_bias is not None: + logits_bias_mask[i : (i + steps)] = True + req_bias_index = bias_to_index[req_bias] + bias_gather_indices.extend(repeat(req_bias_index, steps)) + + if not bias_to_index: + return + assert req_bias is not None # otherwise bias_to_index is empty + + bias_gather_indices_cuda = torch.tensor( + bias_gather_indices, pin_memory=True, dtype=torch.int32 + ).to(logits.device, non_blocking=True) + logits_bias_mask_cuda = logits_bias_mask.to(logits.device, non_blocking=True) + biases_tensor = torch.empty((len(bias_to_index), *req_bias.shape), pin_memory=True) + biases_tensor = torch.stack( + tuple(bias_to_index.keys()), + out=biases_tensor, + ) + biases_tensor_cuda = biases_tensor.to(logits.device, non_blocking=True) + + biases_tensor_cuda = torch.index_select(biases_tensor_cuda, 0, bias_gather_indices_cuda) + # NB: Avoiding logits[bias_scatter_indices] += biases_tensor (and torch.Tensor.scatter_add_), because it + # is unclear if this allows for repeated indices, cf. + # https://docs.pytorch.org/docs/2.8/generated/torch.Tensor.index_put_.html#torch-tensor-index-put + # and thus introduces read-after-write dependencies (including possible false + # sharing). + logits[logits_bias_mask_cuda] += biases_tensor_cuda diff --git a/mypyclib/setup.py b/mypyclib/setup.py new file mode 100644 index 000000000000..0119cca99ab3 --- /dev/null +++ b/mypyclib/setup.py @@ -0,0 +1,12 @@ +from setuptools import setup +from mypyc.build import mypycify +import os + +os.environ["TRTLLM_BUILD_MYPYCLIB"] = "1" + +setup( + name="mypyclib", + version="0.1.0", + packages=["mypyclib"], + ext_modules=mypycify(["mypyclib/__init__.py", "mypyclib/sampler_utils.py"]), +) diff --git a/tensorrt_llm/_torch/attention_backend/interface.py b/tensorrt_llm/_torch/attention_backend/interface.py index 2e0dc20f2985..132469aa2073 100644 --- a/tensorrt_llm/_torch/attention_backend/interface.py +++ b/tensorrt_llm/_torch/attention_backend/interface.py @@ -19,7 +19,7 @@ from ..memory_buffer_utils import Buffers from ..metadata import KVCacheParams -from ..pyexecutor.resource_manager import KVCacheManager +from ..pyexecutor.resource_manager import KVCacheManager, KVCacheManagerV2 from ..utils import get_model_extra_attrs @@ -53,7 +53,7 @@ class AttentionMetadata: # The max number of sequences in a single batch. max_num_sequences: Optional[int] = None # The KV cache manager. - kv_cache_manager: KVCacheManager + kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2] mapping: Optional[Mapping] = None enable_flash_mla: bool = False diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 464396ab785f..c54c9985481c 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -860,11 +860,11 @@ def prepare(self) -> None: assert self.request_ids is not None if self.kv_cache_manager is not None: # Copy blocks for all context requests - self.kv_cache_manager.impl.copy_batch_block_offsets( + self.kv_cache_manager.copy_batch_block_offsets( self.host_kv_cache_block_offsets, self.request_ids[:self.num_contexts], 1, 0) # Copy blocks for all generation requests - self.kv_cache_manager.impl.copy_batch_block_offsets( + self.kv_cache_manager.copy_batch_block_offsets( self.host_kv_cache_block_offsets, self.request_ids[self.num_contexts:], self.beam_width, self.num_contexts) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 8da982aba2bd..16197a6c3746 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -10,13 +10,14 @@ MODEL_CLASS_VISION_ENCODER_MAPPING from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str from tensorrt_llm.bindings.executor import DecodingMode -from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, - EagleDecodingConfig, KvCacheConfig, - MTPDecodingConfig, PeftCacheConfig, - SamplerType, SchedulerConfig, - SparseAttentionConfig, - SpeculativeConfig, TorchLlmArgs) + +# isort: off +from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, + KvCacheConfig, MTPDecodingConfig, PeftCacheConfig, SamplerType, + SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs) from tensorrt_llm.logger import logger +# isort: on from tensorrt_llm.lora_helper import (LoraConfig, get_default_trtllm_modules_to_hf_modules) from tensorrt_llm.lora_manager import load_torch_lora @@ -33,12 +34,13 @@ from .mamba_cache_manager import MambaHybridCacheManager from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (KVCacheManager, PeftCacheManager, - ResourceManager, ResourceManagerType) +from .resource_manager import (KVCacheManager, KVCacheManagerV2, + PeftCacheManager, ResourceManager, + ResourceManagerType) from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, TRTLLMSampler) -from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, - SimpleScheduler) +from .scheduler import (BindMicroBatchScheduler, GuaranteedNoEvictScheduler, + MaxUtilizationScheduler, SimpleScheduler) from .seq_slot_manager import SeqSlotManager GB = 1 << 30 @@ -76,6 +78,7 @@ def __init__( speculative_config: SpeculativeConfig, sparse_attention_config: SparseAttentionConfig, profiling_stage_data: Optional[dict], + scheduler_config: Optional[SchedulerConfig] = None, ): self._model_engine = model_engine self._draft_model_engine = draft_model_engine @@ -94,8 +97,14 @@ def __init__( self._net_max_seq_len = net_max_seq_len self._dummy_reqs = None self._profiling_stage_data = profiling_stage_data + self._scheduler_config = scheduler_config + self._capacity_scheduler_policy = ( + scheduler_config.capacity_scheduler_policy if scheduler_config + is not None else CapacitySchedulerPolicy.GUARANTEED_NO_EVICT) self._kv_cache_manager_cls = get_kv_cache_manager_cls( model_engine.model.model_config) + if self._kv_cache_manager_cls == KVCacheManager and kv_cache_config.use_kv_cache_manager_v2: + self._kv_cache_manager_cls = KVCacheManagerV2 def _get_kv_size_per_token(self): model_config = self._model_engine.model.model_config @@ -467,11 +476,13 @@ def _create_kv_cache_manager( mapping=mapping, dtype=kv_cache_dtype, spec_config=spec_config, + vocab_size=config.vocab_size, max_beam_width=self._max_beam_width, is_draft=model_engine.is_draft_model, kv_connector_manager=self._kv_connector_manager if not estimating_kv_cache else None, sparse_attn_config=sparse_attn_config, + capacity_scheduler_policy=self._capacity_scheduler_policy, ) elif is_nemotron_hybrid(config): if self._max_beam_width > 1: @@ -588,6 +599,7 @@ def _create_kv_cache_manager( mapping=mapping, dtype=kv_cache_dtype, spec_config=spec_config, + vocab_size=config.vocab_size, max_num_tokens=self._max_num_tokens, model_config=binding_model_config, max_beam_width=self._max_beam_width, @@ -595,6 +607,7 @@ def _create_kv_cache_manager( kv_connector_manager=self._kv_connector_manager if not estimating_kv_cache else None, sparse_attn_config=sparse_attn_config, + capacity_scheduler_policy=self._capacity_scheduler_policy, ) # KVCacheManager (Non-draft) modifies the max_seq_len field, update it to self if model_engine.kv_cache_manager_key == ResourceManagerType.KV_CACHE_MANAGER: @@ -779,12 +792,18 @@ def create_py_executor_instance( if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 - capacity_scheduler = BindCapacityScheduler( - scheduler_capacity, - kv_cache_manager.impl if kv_cache_manager is not None else None, - peft_cache_manager.impl if peft_cache_manager is not None else None, - scheduler_config.capacity_scheduler_policy, - two_step_lookahead=mapping.has_pp()) + # Select the capacity scheduler based on the scheduler_config policy + capacity_scheduler_policy = (scheduler_config.capacity_scheduler_policy + if scheduler_config is not None else + CapacitySchedulerPolicy.GUARANTEED_NO_EVICT) + if capacity_scheduler_policy == CapacitySchedulerPolicy.MAX_UTILIZATION: + capacity_scheduler = MaxUtilizationScheduler( + scheduler_capacity, + kv_cache_manager if kv_cache_manager is not None else None) + else: + capacity_scheduler = GuaranteedNoEvictScheduler( + scheduler_capacity, + kv_cache_manager if kv_cache_manager is not None else None) mb_scheduler = BindMicroBatchScheduler(max_batch_size, max_num_tokens, ctx_chunk_config) scheduler = SimpleScheduler(capacity_scheduler, mb_scheduler) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index d924f2ea457c..7a0b0b8840e7 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -15,8 +15,8 @@ from ..modules.multi_stream_utils import with_multi_stream from ..speculative.eagle3 import Eagle3ResourceManager from ..utils import make_weak_ref, piecewise_cuda_graph -from .resource_manager import (BaseResourceManager, ResourceManager, - ResourceManagerType) +from .resource_manager import (BaseResourceManager, KVCacheManager, + ResourceManager, ResourceManagerType) from .scheduler import ScheduledRequests # A large prime number used for dummy request IDs to avoid collisions @@ -372,17 +372,22 @@ def _get_padded_batch(self, batch: ScheduledRequests, # This is not strictly required, but we should probably # respect the requirement just in case that changes in the future. if self.padding_dummy_request is None: - available_blocks = kv_cache_manager.get_num_free_blocks() - # No padding if not enough KV cache space - if available_blocks < 1: - return 0 + if isinstance(kv_cache_manager, KVCacheManager): + available_blocks = kv_cache_manager.get_num_free_blocks() + # No padding if not enough KV cache space + if available_blocks < 1: + return 0 self.padding_dummy_request = kv_cache_manager.add_dummy_requests( [CUDA_GRAPH_DUMMY_REQUEST_ID], is_gen=True, max_num_draft_tokens=runtime_draft_len, use_mrope=self.config.use_mrope, - max_beam_width=self.config.max_beam_width)[0] + max_beam_width=self.config.max_beam_width) + if self.padding_dummy_request is None: + return 0 + else: + self.padding_dummy_request = self.padding_dummy_request[0] self.padding_dummy_request.is_cuda_graph_dummy = True spec_res_mgr = resource_manager.get_resource_manager( ResourceManagerType.SPEC_RESOURCE_MANAGER) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9b679c4ae116..5aaec7ec9bbd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -8,7 +8,7 @@ import weakref from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch._dynamo.config @@ -61,7 +61,8 @@ from .llm_request import LlmRequest, get_draft_token_length from .model_loader import ModelLoader, _construct_checkpoint_loader from .resource_manager import (BaseResourceManager, KVCacheManager, - ResourceManager, ResourceManagerType) + KVCacheManagerV2, ResourceManager, + ResourceManagerType) from .sampler import SampleStateTensors from .scheduler import ScheduledRequests @@ -140,6 +141,7 @@ def __init__( torch.nn.Module]] = None, model: Optional[torch.nn.Module] = None, ): + gc.disable() self.forward_pass_callable = None self.ub_buffers = None ( @@ -586,8 +588,8 @@ def _run_torch_compile_warmup(self, resource_manager: ResourceManager): self.kv_cache_manager_key) curr_max_num_tokens = min( kv_cache_manager.get_num_available_tokens( - self.original_max_draft_len), self.max_num_tokens, - self.batch_size * (self.max_seq_len - 1)) + max_num_draft_tokens=self.original_max_draft_len), + self.max_num_tokens, self.batch_size * (self.max_seq_len - 1)) warmup_requests_configs = { (1, 1), # Specialize for 1 token. @@ -624,8 +626,8 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): self.kv_cache_manager_key) curr_max_num_tokens = min( kv_cache_manager.get_num_available_tokens( - self.original_max_draft_len), self.max_num_tokens, - self.batch_size * (self.max_seq_len - 1)) + max_num_draft_tokens=self.original_max_draft_len), + self.max_num_tokens, self.batch_size * (self.max_seq_len - 1)) cache_path = os.environ.get("TLLM_AUTOTUNER_CACHE_PATH", None) with self.no_cuda_graph(), autotune(cache_path=cache_path, @@ -797,7 +799,7 @@ def _create_warmup_request( ResourceManagerType.SPEC_RESOURCE_MANAGER) available_tokens = kv_cache_manager.get_num_available_tokens( - self.runtime_draft_len) + max_num_draft_tokens=self.runtime_draft_len) available_blocks = kv_cache_manager.get_num_free_blocks() if num_tokens > self.max_num_tokens or num_tokens > available_tokens: return None @@ -829,7 +831,8 @@ def _create_warmup_request( num_left_over_tokens / kv_cache_manager.tokens_per_block) + num_gen_tokens - if blocks_to_use > available_blocks: + if blocks_to_use > available_blocks and isinstance( + kv_cache_manager, KVCacheManager): return None if num_ctx_tokens > 0: @@ -844,6 +847,9 @@ def _create_warmup_request( max_num_draft_tokens=self.runtime_draft_len, use_mrope=self.use_mrope) + if ctx_requests is None: + return None + if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests( request_ids=list(range(num_ctx_requests))) @@ -856,6 +862,12 @@ def _create_warmup_request( is_gen=True, max_num_draft_tokens=self.max_total_draft_tokens, use_mrope=self.use_mrope) + + if gen_requests is None: + for r in ctx_requests: + kv_cache_manager.free_resources(r) + return None + if spec_resource_manager is not None: spec_resource_manager.add_dummy_requests(request_ids=list( range(num_ctx_requests, num_ctx_requests + num_gen_tokens))) @@ -892,7 +904,11 @@ def _create_cuda_graph_warmup_request( max_beam_width=self.max_beam_width, num_extra_decoding_steps=num_extra_decoding_steps) - available_tokens = kv_cache_manager.get_num_available_tokens(draft_len) + if requests is None: + return None + + available_tokens = kv_cache_manager.get_num_available_tokens( + batch_size=batch_size, max_num_draft_tokens=draft_len) # Add one dummy request with the maximum possible sequence length. token_num = max(1, min(available_tokens, self.max_seq_len - 1)) @@ -914,7 +930,14 @@ def _create_cuda_graph_warmup_request( max_num_draft_tokens=draft_len, use_mrope=self.use_mrope, max_beam_width=self.max_beam_width, - num_extra_decoding_steps=num_extra_decoding_steps)[0] + num_extra_decoding_steps=num_extra_decoding_steps) + + if max_seq_len_request is None: + for r in requests: + kv_cache_manager.free_resources(r) + return None + else: + max_seq_len_request = max_seq_len_request[0] # Insert the longest request first to simulate padding for the CUDA graph. requests.insert(0, max_seq_len_request) @@ -958,7 +981,8 @@ def _update_draft_inference_state_for_warmup( req.py_is_first_draft = True req.py_draft_tokens = [] - def _set_up_attn_metadata(self, kv_cache_manager: KVCacheManager): + def _set_up_attn_metadata(self, kv_cache_manager: Union[KVCacheManager, + KVCacheManagerV2]): enable_context_mla_with_cached_kv = is_mla( self.model.model_config.pretrained_config) and ( self.attn_runtime_features.cache_reuse @@ -1290,7 +1314,7 @@ def _prepare_multimodal_indices(self, input_ids: list[int]): def _prepare_tp_inputs( self, scheduled_requests: ScheduledRequests, - kv_cache_manager: KVCacheManager, + kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], attn_metadata: AttentionMetadata, spec_metadata: Optional[SpecMetadata] = None, new_tensors_device: Optional[SampleStateTensors] = None, @@ -2510,7 +2534,7 @@ def _get_lora_params_from_requests(self, def _prepare_inputs( self, scheduled_requests: ScheduledRequests, - kv_cache_manager: KVCacheManager, + kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], attn_metadata: AttentionMetadata, spec_metadata: Optional[SpecMetadata] = None, new_tensors_device: Optional[SampleStateTensors] = None, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index b82d36fe8e83..90f2c5ecee1b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -596,6 +596,7 @@ def drafting_loop_wrapper(model): speculative_config=spec_config, profiling_stage_data=profiling_stage_data, sparse_attention_config=sparse_attention_config, + scheduler_config=scheduler_config, ) estimating_kv_cache = kv_cache_creator.try_prepare_estimation() with allocation_scope( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index cba535e108a4..586e5d57ff05 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -6,17 +6,36 @@ from typing import (TYPE_CHECKING, Dict, Iterable, List, Optional, Set, Tuple, Union) +import numpy as np import torch import tensorrt_llm import tensorrt_llm.bindings -from tensorrt_llm._utils import mpi_disabled +from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, + get_size_in_bytes, mpi_disabled) from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE -from tensorrt_llm.llmapi.llm_args import (KvCacheConfig, PeftCacheConfig, +from tensorrt_llm.llmapi.llm_args import (CapacitySchedulerPolicy, + KvCacheConfig, PeftCacheConfig, PybindMirror) from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.lora_manager import LoraManager, LoraModelConfig from tensorrt_llm.runtime import ModelConfig as ModelConfigPython +from tensorrt_llm.runtime.kv_cache_manager_v2 import (AttentionLayerConfig, + BufferConfig, + DiskCacheTierConfig, + GpuCacheTierConfig, + HostCacheTierConfig) +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManager as KVCacheManagerPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import LayerId, _KVCache +from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole +from tensorrt_llm.runtime.kv_cache_manager_v2._copy_engine import \ + copy_batch_block_offsets as copy_batch_block_offsets_nanobind +from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError +from tensorrt_llm.runtime.kv_cache_manager_v2._utils import (exact_div, + typed_range) from tensorrt_llm.sampling_params import SamplingParams from ..._utils import (binding_to_str_dtype, get_size_in_bytes, mpi_rank, @@ -61,6 +80,14 @@ class ResourceManagerType(enum.Enum): SPEC_RESOURCE_MANAGER = "SPEC_RESOURCE_MANAGER" +class Role: + KEY = DataRole("key") + VALUE = DataRole("value") + KEY_BLOCK_QUANT = DataRole("key_block_quant") + VALUE_BLOCK_QUANT = DataRole("value_block_quant") + ALL = DataRole("all") + + def compute_page_count(token_count: int, tokens_per_page: int) -> int: return (token_count + tokens_per_page) // tokens_per_page @@ -301,6 +328,9 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], (self.blocks_in_primary_pool, self.blocks_in_secondary_pool) } + print(f"blocks_per_window: {blocks_per_window}") + print(f"self.max_attention_window_vec: {self.max_attention_window_vec}") + # Validate and adjust attention windows against their upper bounds if needed blocks_per_window, self.max_seq_len, self.max_attention_window_vec = self._validate_and_adjust_attention_windows( max_attention_window_vec=self.max_attention_window_vec, @@ -311,6 +341,9 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], max_beam_width=max_beam_width, ) + print(f"blocks_per_window: {blocks_per_window}") + print(f"self.max_attention_window_vec: {self.max_attention_window_vec}") + if kv_cache_type != CacheTypeCpp.SELF: assert len( blocks_per_window @@ -353,6 +386,8 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], 'indexer_k_cache_index_head_dim': indexer_k_cache_index_head_dim } + print(f"kv_connector_manager: {self.kv_connector_manager}") + if self.event_buffer_max_size > 0: if mapping.enable_attention_dp: kwargs['event_manager'] = KVCacheEventManagerCpp( @@ -366,6 +401,9 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], kwargs['event_manager'] = KVCacheEventManagerCpp( max_kv_event_entries=self.event_buffer_max_size) + if tensorrt_llm.mpi_rank() == 0: + print(f"kwargs: {kwargs}") + self.impl = KVCacheManagerCpp(**kwargs) self.impl.allocate_pools(False) @@ -379,6 +417,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], dim=-1) self.kv_cache_pool_mapping = self.impl.get_layer_to_pool_mapping() + print(f"kv_cache_pool_mapping: {self.kv_cache_pool_mapping}") self.num_pools = self.impl.num_pools self.max_blocks_per_seq = self.impl.max_blocks_per_seq self.enable_block_reuse = kv_cache_config.enable_block_reuse @@ -404,6 +443,7 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: remaining_tokens / self.tokens_per_block) return need_blocks + @nvtx_range("prepare_resources_kv_cache_manager_v1") def prepare_resources(self, scheduled_batch: ScheduledRequests): with request_context(self.is_draft, scheduled_batch): context_batch = scheduled_batch.context_requests @@ -411,6 +451,9 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # allocate KV Cache for req in context_batch: req_beam_width = req.sampling_config.beam_width + # print( + # f"prepare_resources: {req.py_request_id}, req.is_first_context_chunk: {req.is_first_context_chunk}, req.context_current_position: {req.context_current_position}, req.context_chunk_size: {req.context_chunk_size}, {req.cached_tokens}" + # ) if 'cp_type' in self.mapping.cp_config and CpType.STAR == self.mapping.cp_config[ 'cp_type']: if req.ctx_iters == 0: @@ -427,6 +470,9 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self.impl.add_sequence(req.py_request_id, req.prompt_len, req_beam_width, req) + # print( + # f"after add sequence: {req.py_request_id}, {req.prompt_len}, {req.context_current_position}, {req.context_chunk_size}, {req.cached_tokens}" + # ) for _ in range(self.num_extra_kv_tokens): self.impl.add_token(req.py_request_id) for _ in range(get_draft_token_length(req)): @@ -809,7 +855,9 @@ def get_num_free_blocks(self) -> int: def get_num_kv_blocks(self, num_tokens: int) -> int: return (num_tokens + self.tokens_per_block - 1) // self.tokens_per_block - def get_num_available_tokens(self, max_num_draft_tokens: int = 0) -> int: + def get_num_available_tokens(self, + max_num_draft_tokens: int = 0, + **kwargs) -> int: return (self.get_num_free_blocks() * self.tokens_per_block - self.num_extra_kv_tokens - max_num_draft_tokens) @@ -1191,11 +1239,911 @@ def _set_temp_attention_window_inputs( else: return None + @nvtx_range("copy_batch_block_offsets") + def copy_batch_block_offsets(self, dst_tensor: torch.Tensor, + request_ids: List[int], beam_width: int, + offset: int): + self.impl.copy_batch_block_offsets(dst_tensor, request_ids, beam_width, + offset) + def reset_reuse_state(self): """Reset the reuse state of the KV cache manager.""" self.impl.reset_reuse_state() +class KVCacheManagerV2(BaseResourceManager): + + def __init__( + self, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config=None, + layer_mask: Optional[List[bool]] = None, + vocab_size: int = None, + max_num_tokens: int = 8192, + model_config: Optional[ModelConfigCpp] = None, + max_beam_width: int = 1, + is_draft: bool = False, + kv_connector_manager: Optional[KvCacheConnectorManager] = None, + capacity_scheduler_policy: + CapacitySchedulerPolicy = CapacitySchedulerPolicy.GUARANTEED_NO_EVICT, + **kwargs, + ) -> None: + self.mapping = mapping + self.dtype = dtype + self.capacity_scheduler_policy = capacity_scheduler_policy + + assert self.dtype != DataType.NVFP4, "NVFP4 is not supported for KVCacheManagerV2" + assert kv_connector_manager is None, "kv_connector_manager is not supported for KVCacheManagerV2" + assert max_beam_width == 1, "max_beam_width must be 1 for KVCacheManagerV2" + + self.kv_cache_type = kv_cache_type + self.pp_layers, self.num_layers = get_pp_layers( + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + self.is_draft = is_draft + self.num_local_layers = len(self.pp_layers) + self.layer_offsets = { + idx: offset + for offset, idx in enumerate(self.pp_layers) + } + + tp_size = mapping.tp_size + if mapping.enable_attention_dp: + tp_size = 1 + + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.tokens_per_block = tokens_per_block + self.max_seq_len = max_seq_len + self.max_batch_size = max_batch_size + self.kv_factor = 1 if kv_cache_type == CacheTypeCpp.SELFKONLY else 2 + from ..speculative import get_num_extra_kv_tokens + self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + + self.event_buffer_max_size = kv_cache_config.event_buffer_max_size + + assert self.event_buffer_max_size == 0, "event_buffer_max_size must be 0" + + # Determine max_attention_window_vec + if kv_cache_config.max_attention_window is not None: + + self.max_attention_window_vec = kv_cache_config.max_attention_window.copy( + ) # Make a copy to avoid modifying original + # Clamp all window sizes to max_seq_len before calculating the + # number of KV cache blocks. This prevents the KV cache pool from + # being skewed by the largest window values. + self.max_attention_window_vec = [ + min(max_seq_len, w) for w in self.max_attention_window_vec + ] + + self.max_attention_window_vec = [ + None if w == max_seq_len else w + for w in self.max_attention_window_vec + ] + + else: + self.max_attention_window_vec = [None] + + if isinstance(num_kv_heads, int): + self.num_kv_heads_per_layer = [ + (num_kv_heads + tp_size - 1) // tp_size + for _ in range(self.num_local_layers) + ] + self.total_num_kv_heads_per_layer = [ + (num_kv_heads + tp_size - 1) // tp_size + for _ in range(self.num_layers) + ] + else: + assert len(num_kv_heads) == self.num_layers + + def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], + kv_head: Optional[int]): + if kv_head is not None: + num_kv_heads_per_layer.append( + (kv_head + tp_size - 1) // tp_size) + else: + num_kv_heads_per_layer.append(0) + + self.num_kv_heads_per_layer = [] + if self.num_local_layers > 0: + for i in self.pp_layers: + kv_head = num_kv_heads[i] + append_to_kv_heads_per_layer(self.num_kv_heads_per_layer, + kv_head) + + self.total_num_kv_heads_per_layer = [] + for i in range(self.num_layers): + kv_head = num_kv_heads[i] + append_to_kv_heads_per_layer(self.total_num_kv_heads_per_layer, + kv_head) + + free_mem_fraction = (kv_cache_config.free_gpu_memory_fraction + if kv_cache_config.free_gpu_memory_fraction + is not None else 0.9) + + self.is_vswa = len(set(self.max_attention_window_vec)) > 1 + + assert free_mem_fraction < 1.0, f"Invalid freeMemFraction, freeMemFraction {free_mem_fraction} must be smaller than 1.0" + + free_mem, total_mem = torch.cuda.mem_get_info() + max_tokens = free_mem_fraction * free_mem / self.get_cache_bytes_per_token( + ) + + self.kv_connector_manager = kv_connector_manager + + if kv_cache_config.max_tokens is not None: + config_max_tokens = int( + math.ceil(kv_cache_config.max_tokens / + kv_cache_config.max_util_for_resume) * 1.2) + if kv_cache_config.free_gpu_memory_fraction is not None: + max_tokens = min(config_max_tokens, max_tokens) + logger.warning( + f'Both free_gpu_memory_fraction and max_tokens are set (to {free_mem_fraction} and {max_tokens} with free memory {free_mem / (1 << 32)} of total memory {total_mem / (1<<32)}, respectively). The smaller value will be used.' + ) + else: + max_tokens = config_max_tokens + logger.info( + f"max_tokens is set by kv_cache_config.max_tokens: {max_tokens}" + ) + + # Rough estimate of the quota needed for the KV cache + # TODO: Consider vswa case + self.quota = GpuCacheTierConfig( + quota=int(max_tokens * self.get_cache_bytes_per_token())) + + logger.info( + f"Allocated {self.quota.quota / (1 << 30)} GiB in paged KV cache.") + + buffer_type = [Role.KEY] + if kv_cache_type != CacheTypeCpp.SELFKONLY: + buffer_type.append(Role.VALUE) + + config = KVCacheManagerConfigPy( + tokens_per_block=tokens_per_block, + vocab_size=vocab_size, + cache_tiers=[ + self.quota, + # Magic Number for now + HostCacheTierConfig(quota=8000 << 20), + DiskCacheTierConfig(quota=1 << 30, path="/workspace/"), + ], + max_util_for_resume=kv_cache_config.max_util_for_resume, + layers=[ + AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig( + role=role, + size=self.get_cache_bytes_per_token( + local_layer_idx=layer_id, data_role=role) * + tokens_per_block, + ) for role in buffer_type + ], + sliding_window_size=self.max_attention_window_vec[ + layer_id % len(self.max_attention_window_vec)], + num_sink_tokens=None, + ) for layer_id in typed_range(LayerId(self.num_local_layers)) + ], + ) + + self.kv_cache_manager_py_config = config + + self.impl = KVCacheManagerPy(config) + + self.num_pools = len(self.impl.layer_grouping) + + self.layer_to_pool_mapping_dict: dict[int, int] = { + layer_id: self.impl.get_layer_group_id(layer_id) + for layer_id in typed_range(LayerId(self.num_local_layers)) + } + self.pool_to_layer_mapping_dict: dict[int, list[int]] = {} + + self.kv_cache_pool_pointers = torch.tensor([[ + self.impl.get_mem_pool_base_address( + self.impl.layer_grouping[pool_id][0], Role.KEY), 0 + ] for pool_id in range(self.num_pools)], + dtype=torch.int64, + device="cpu", + pin_memory=True) + + kv_cache_pool_mapping_list = [] + for layer_id in typed_range(LayerId(self.num_local_layers)): + layer_group_id = self.impl.get_layer_group_id(layer_id) + offset = exact_div( + self.impl.get_mem_pool_base_address(layer_id, Role.KEY) - + int(self.kv_cache_pool_pointers[layer_group_id][0]), + self.get_cache_bytes_per_token(layer_id, Role.KEY) * + self.kv_factor * self.tokens_per_block) + kv_cache_pool_mapping_list.append([layer_group_id, offset]) + + self.kv_cache_pool_mapping = torch.tensor(kv_cache_pool_mapping_list, + dtype=torch.int32, + device="cpu", + pin_memory=True) + self.max_blocks_per_seq = (max_seq_len + tokens_per_block - + 1) // tokens_per_block + + self.kv_cache_map: dict[int, _KVCache] = {} + + max_num_tokens = self.get_num_available_tokens() + + if max_seq_len > max_num_tokens: + logger.warning( + f"max_seq_len {max_seq_len} is greater than max_num_tokens {max_num_tokens} that can be allocated in kv cache manager, setting max_seq_len to {max_num_tokens}" + ) + self.max_seq_len = max_num_tokens + + self.enable_block_reuse = kv_cache_config.enable_block_reuse + + @property + def blocks_in_primary_pool(self) -> int: + """ + Get the number of blocks in the primary pool. + """ + return self.impl.get_page_index_upper_bound(0, Role.KEY) + + def get_buffers(self, layer_idx: int) -> Optional[torch.Tensor]: + layer_offset = self.layer_offsets[layer_idx] + addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY) + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + addr_value = self.impl.get_mem_pool_base_address( + layer_offset, Role.VALUE) + page_size_key = self.impl.get_page_stride(layer_offset, Role.KEY) + page_size_value = self.impl.get_page_stride(layer_offset, + Role.VALUE) + + assert addr_key + page_size_value == addr_value and page_size_key == page_size_value + + return convert_to_torch_tensor( + TensorWrapper( + addr_key, + self.dtype, + ( + self.impl.get_page_index_upper_bound( + layer_offset, Role.KEY), + self.kv_factor, + self.tokens_per_block, + self.num_kv_heads_per_layer[layer_offset], + self.head_dim, + ), + )) + + def get_num_available_tokens(self, + *, + batch_size: int = 1, + max_num_draft_tokens: int = 0) -> int: + if max_num_draft_tokens > 0: + raise ValueError( + "max_num_draft_tokens is not supported for KVCacheManagerV2") + # Multiplied by 0.95 to make `max_util_for_resume` happy + return int( + self.impl.clamp_max_seq_len_for_mem(batch_size) * + self.kv_cache_manager_py_config.max_util_for_resume * + 0.95) - self.num_extra_kv_tokens - max_num_draft_tokens + + def get_num_free_blocks(self) -> int: + # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks. + # However, since we only use this function when the kv cache manager is empty, so it is safe to do so. + assert len( + self.kv_cache_map + ) == 0, "get_num_free_blocks is only used when the kv cache manager is empty" + max_num_pages = max([ + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) + for layer_id in typed_range(LayerId(self.num_local_layers)) + ]) + return max_num_pages // self.kv_factor + + @nvtx_range("prepare_resources_kv_cache_manager_v2") + def prepare_resources(self, scheduled_batch: ScheduledRequests): + """Unified prepare_resources that dispatches based on capacity_scheduler_policy.""" + if self.capacity_scheduler_policy == CapacitySchedulerPolicy.MAX_UTILIZATION: + self._prepare_resources_max_utilization(scheduled_batch) + else: + # Default to GUARANTEED_NO_EVICT + self._prepare_resources_guaranteed_no_evict(scheduled_batch) + + @nvtx_range("prepare_resources_kv_cache_manager_v2_guaranteed_no_evict") + def _prepare_resources_guaranteed_no_evict( + self, scheduled_batch: ScheduledRequests): + with request_context(self.is_draft, scheduled_batch): + context_batch = scheduled_batch.context_requests + generation_batch = scheduled_batch.generation_requests + # allocate KV Cache + for req in context_batch: + beam_width = req.sampling_config.beam_width + if 'cp_type' in self.mapping.cp_config and CpType.STAR == self.mapping.cp_config[ + 'cp_type']: + raise RuntimeError( + "Star attention is not supported for kv cache manager v2" + ) + else: + if req.is_first_context_chunk and self._kv_connector_should_add_sequence( + req): + # Last token cannot be recovered, so we don't include it in the input tokens to look up for the block that can be reused. + kv_cache = self.impl.create_kv_cache( + req.lora_task_id, + req.get_tokens(0)[:-1] + if self.enable_block_reuse else None) + assert beam_width == 1, "Currently, KVCacheManagerV2 only supports beam width 1" + assert req.py_request_id not in self.kv_cache_map, f"req.py_request_id {req.py_request_id} already in kv_cache_map" + self.kv_cache_map[req.py_request_id] = kv_cache + if not self.enable_block_reuse: + assert kv_cache.num_committed_tokens == 0 + kv_cache.stop_committing() + else: + req.context_current_position = kv_cache.num_committed_tokens + chunk_size = req.context_chunk_size + if req.context_current_position + req.context_chunk_size < req.prompt_len: + floored_end_position = ( + req.context_current_position + + req.context_chunk_size + ) // self.tokens_per_block * self.tokens_per_block + chunk_size = floored_end_position - req.context_current_position + + req.context_chunk_size = min( + chunk_size, + req.prompt_len - req.context_current_position) + + success = kv_cache.resume( + torch.cuda.current_stream().cuda_stream) + assert success + + kv_cache.capacity = req.prompt_len + + if self.kv_connector_manager is not None: + block_ids = self.get_cache_indices(req) + self.kv_connector_manager.update_state_after_alloc( + req, block_ids) + + for req in generation_batch: + kv_cache = self.kv_cache_map[req.py_request_id] + kv_cache.capacity += 1 + + if self.kv_connector_manager is not None: + self.kv_connector_manager.build_scheduler_output( + scheduled_batch, self) + + @nvtx_range("prepare_resources_kv_cache_manager_v2_max_utilization") + def _prepare_resources_max_utilization(self, + scheduled_batch: ScheduledRequests): + evicted_requests = [] + no_scheduled = [] + with request_context(self.is_draft, scheduled_batch): + context_batch = scheduled_batch.context_requests + generation_batch = scheduled_batch.generation_requests + + new_generation_batch: RequestList = [] + + for req in generation_batch: + if req in evicted_requests: + continue + kv_cache = self.kv_cache_map[req.py_request_id] + + if not kv_cache.is_active: + result = kv_cache.resume( + torch.cuda.current_stream().cuda_stream) + if not result: + no_scheduled.append(req) + continue + # Max Utilization Scheduler: Try to increase capacity for generation + # Recursively try to evict requests until we have enough capacity + max_eviction_attempts = len(generation_batch) - len( + evicted_requests) + capacity_increased = False + + for _ in range(max_eviction_attempts): + try: + kv_cache.capacity += 1 + new_generation_batch.append(req) + capacity_increased = True + break + except OutOfPagesError: + evicted = self._try_evict_requests_for_capacity( + scheduled_batch, req, kv_cache.capacity + 1, + kv_cache, new_generation_batch) + if evicted is None: + # No more requests to evict + break + if evicted in new_generation_batch: + new_generation_batch.remove(evicted) + evicted_requests.append(evicted) + + if not capacity_increased: + # Could not increase capacity even after evicting all possible requests + no_scheduled.append(req) + continue + + scheduled_batch.generation_requests = new_generation_batch + + # allocate KV Cache + + new_context_batch: RequestList = [] + # print(f"len(context_batch): {len(context_batch)}") + for req in context_batch: + beam_width = req.sampling_config.beam_width + if 'cp_type' in self.mapping.cp_config and CpType.STAR == self.mapping.cp_config[ + 'cp_type']: + raise RuntimeError( + "Star attention is not supported for kv cache manager v2" + ) + else: + kv_cache = None + # print( + # f"req.py_request_id: {req.py_request_id if hasattr(req, 'py_request_id') else req.request_id}, req.is_first_context_chunk: {req.is_first_context_chunk}, self._kv_connector_should_add_sequence(req): {self._kv_connector_should_add_sequence(req)}" + # ) + if req.is_first_context_chunk and self._kv_connector_should_add_sequence( + req): + if req.py_request_id in self.kv_cache_map: + kv_cache = self.kv_cache_map[req.py_request_id] + else: + # Last token cannot be recovered, so we don't include it in the input tokens to look up for the block that can be reused. + kv_cache = self.impl.create_kv_cache( + req.lora_task_id, + req.get_tokens(0)[:-1] + if self.enable_block_reuse else None) + assert beam_width == 1, "Currently, KVCacheManagerV2 only supports beam width 1" + assert req.py_request_id not in self.kv_cache_map, f"req.py_request_id {req.py_request_id} already in kv_cache_map" + self.kv_cache_map[req.py_request_id] = kv_cache + if not self.enable_block_reuse: + assert kv_cache.num_committed_tokens == 0 + kv_cache.stop_committing() + else: + req.context_current_position = kv_cache.num_committed_tokens + chunk_size = req.context_chunk_size + if req.context_current_position + req.context_chunk_size < req.prompt_len: + floored_end_position = ( + req.context_current_position + + req.context_chunk_size + ) // self.tokens_per_block * self.tokens_per_block + chunk_size = floored_end_position - req.context_current_position + + req.context_chunk_size = min( + chunk_size, + req.prompt_len - req.context_current_position) + + success = kv_cache.resume( + torch.cuda.current_stream().cuda_stream) + if not success: + # print( + # f"resume no success failed to resume kv_cache for request {req.py_request_id if hasattr(req, 'py_request_id') else req.request_id}" + # ) + no_scheduled.append(req) + continue + try: + kv_cache.capacity = req.prompt_len + new_context_batch.append(req) + except OutOfPagesError: + # print( + # f"OutOfPagesError failed to resume kv_cache for request {req.py_request_id if hasattr(req, 'py_request_id') else req.request_id}" + # ) + no_scheduled.append(req) + kv_cache.suspend() + continue + + if self.kv_connector_manager is not None: + block_ids = self.get_cache_indices(req) + self.kv_connector_manager.update_state_after_alloc( + req, block_ids) + else: + assert req.py_request_id in self.kv_cache_map, f"req.py_request_id {req.py_request_id} not in kv_cache_map" + kv_cache = self.kv_cache_map[req.py_request_id] + assert kv_cache.status is _KVCache.Status.ACTIVE, f"kv_cache {req.py_request_id} is not active" + new_context_batch.append(req) + + scheduled_batch.context_requests = new_context_batch + + if self.kv_connector_manager is not None: + self.kv_connector_manager.build_scheduler_output( + scheduled_batch, self) + + def _try_evict_requests_for_capacity(self, + scheduled_batch: ScheduledRequests, + current_req: LlmRequest, + needed_capacity: int, current_kv_cache, + new_generation_batch) -> LlmRequest: + """ + Try to evict requests to make room for capacity allocation. + + Based on TestBatching pattern (lines 387-393): + - Find requests that can be evicted (from the end - LIFO) + - Suspend their kv_caches + - Move them from scheduled to paused + + Args: + scheduled_batch: Current scheduled batch + current_req: Request that needs capacity + needed_capacity: Required capacity + current_kv_cache: KV cache of current request + + Returns: + List of evicted requests + """ + evicted_request = None + all_scheduled_requests = scheduled_batch.context_requests + scheduled_batch.generation_requests + + # Try to evict from the end (LIFO - Last In First Out) + # Don't evict the current request itself + # for req in reversed(scheduled_batch.generation_requests): + for req in reversed(new_generation_batch): + # for req in scheduled_batch.generation_requests: + if req == current_req: + continue + + req_id = req.py_request_id if hasattr( + req, 'py_request_id') else req.request_id + + # Check if this request has a kv_cache + if req_id not in self.kv_cache_map: + continue + + kv_cache = self.kv_cache_map[req_id] + + if kv_cache.status is not _KVCache.Status.ACTIVE: + continue + + # Only evict requests that are in generation or have started context processing + # (similar to TestBatching line 223-227) + can_evict = False + if req.state == LlmRequestState.GENERATION_IN_PROGRESS: + can_evict = True + elif req.state == LlmRequestState.CONTEXT_INIT and \ + hasattr(req, 'context_current_position') and \ + req.context_current_position > 0: + can_evict = True + + if not can_evict: + continue + + # Suspend the kv_cache + kv_cache.suspend() + evicted_request = req + break + if evicted_request is None: + for req in all_scheduled_requests: + if req == current_req: + continue + req_id = req.py_request_id if hasattr( + req, 'py_request_id') else req.request_id + + # Check if this request has a kv_cache + if req_id not in self.kv_cache_map: + continue + + kv_cache = self.kv_cache_map[req_id] + if kv_cache.status is not _KVCache.Status.ACTIVE: + continue + can_evict = False + if req.state == LlmRequestState.GENERATION_IN_PROGRESS: + can_evict = True + elif req.state == LlmRequestState.CONTEXT_INIT and \ + hasattr(req, 'context_current_position') and \ + req.context_current_position > 0: + can_evict = True + + if not can_evict: + continue + kv_cache.suspend() + evicted_request = req + break + return evicted_request + + def _kv_connector_should_add_sequence(self, request: LlmRequest) -> bool: + return self.kv_connector_manager is None or self.kv_connector_manager.should_add_sequence( + request) + + def get_kv_cache_stats(self): + + class KVCacheStatus: + + def __init__(self, allocated_bytes: int): + self.allocated_bytes = 0 + + return KVCacheStatus(allocated_bytes=self.quota.quota) + + def add_dummy_requests( + self, + request_ids: List[int], + # Note that token_nums should be past_kv_len + input_len (without + # spec decoding). The draft tokens will be added in this function, + # so we don't need to take care of it in the caller. When preparing + # token_nums, we should not take the draft tokens into account, so + # don't use the kv_cache_manager.max_seq_len, which includes both + # extra tokens and draft tokens. + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + max_num_draft_tokens: int = 0, + use_mrope: bool = False, + max_beam_width: int = 1, + num_extra_decoding_steps: + int = 0, # TODO: support num_extra_decoding_steps + ): + + beam_width = max_beam_width + requests = [] + for i, req_id in enumerate(request_ids): + # exact choice of n can be ignored for dummy requests + sampling_params = SamplingParams(n=beam_width, + best_of=beam_width, + use_beam_search=beam_width > 1) + # Here 1+max_num_draft_tokens is used to extend the prompt length to + # a non-zero number to skip illegal memory access issue in MLA kernel + # during warmup. + token_num = token_nums[ + i] if token_nums is not None else 1 + max_num_draft_tokens + # TODO: support cross attention + encoder_input_tokens = None + # Using 1 instead of 0 prevents NaN during warmup in e.g. Deepseek + input_tokens = [1 for _ in range(token_num)] + req = LlmRequest(request_id=req_id, + max_new_tokens=1, + input_tokens=input_tokens, + sampling_config=SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + encoder_input_tokens=encoder_input_tokens) + req.is_dummy_request = True + req.paged_kv_block_ids = [] + if prepare_resource: + kv_cache = self.impl.create_kv_cache(req.lora_task_id, + input_tokens) + assert kv_cache.num_committed_tokens == 0 + success = kv_cache.resume( + torch.cuda.current_stream().cuda_stream) + if not success: + for r in requests: + self.free_resources(r) + self.free_resources(req) + return None + kv_cache.stop_committing() + kv_cache.capacity = token_num + self.kv_cache_map[req_id] = kv_cache + + if is_gen: + req.state = LlmRequestState.GENERATION_IN_PROGRESS + req.prompt_len = token_num - 1 + req.py_prompt_len = req.prompt_len + + # TODO: Planning to get dummy_data from each model. Before that, we need to add dummy mrop_config to the request here. + if use_mrope: + dummy_mrope_position_ids = torch.arange( + 0, token_num, dtype=torch.int32).expand(3, 1, -1).clone() + req.py_multimodal_data = { + "mrope_config": { + "mrope_position_ids": dummy_mrope_position_ids + } + } + if is_gen: + dummy_mrope_position_deltas = torch.zeros( + 1, dtype=torch.int32).unsqueeze(0) + req.py_multimodal_data["mrope_config"][ + "mrope_position_deltas"] = dummy_mrope_position_deltas + requests.append(req) + + return requests + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + if request.py_request_id in self.kv_cache_map: + kv_cache = self.kv_cache_map.pop(request.py_request_id) + kv_cache.close() + + def get_batch_cache_indices( + self, + request_ids: List[int], + *, + layer_id: int = 0, + is_kv_aggregate: bool = True) -> List[List[int]]: + + return self.get_batch_cache_indices_by_pool_id( + request_ids, + pool_id=self.layer_to_pool_mapping_dict[layer_id], + is_kv_aggregate=is_kv_aggregate) + + def get_batch_cache_indices_by_pool_id( + self, + request_ids: List[int], + *, + pool_id: int = 0, + is_kv_aggregate: bool = True) -> List[List[int]]: + + if is_kv_aggregate: + # Div by kv_factor to index kv cache with size [num_blocks, kv_factor, tokens_per_block, num_kv_heads, head_dim] + div_factor = self.kv_factor + else: + div_factor = 1 + + return [(np.asarray(self.kv_cache_map[req_id].get_page_indices(pool_id), + copy=False) // div_factor).tolist() + for req_id in request_ids] + + def get_cache_bytes_per_token( + self, + local_layer_idx: Optional[int] = None, + data_role: Role = Role.ALL): # None means all layers/data_roles + if self.dtype not in ( + DataType.FP8, + DataType.HALF, + DataType.BF16, + DataType.FLOAT, + DataType.NVFP4, + ): + raise ValueError(f"Cannot support {self.dtype} KV cache.") + + if data_role == Role.ALL: + kv_factor = self.kv_factor + elif data_role in [ + Role.KEY, Role.VALUE, Role.KEY_BLOCK_QUANT, + Role.VALUE_BLOCK_QUANT + ]: + if data_role in [Role.KEY_BLOCK_QUANT, Role.VALUE_BLOCK_QUANT]: + assert self.dtype == DataType.NVFP4, "NVFP4 is the only supported dtype for block quant data roles" + if data_role == Role.VALUE: + assert self.kv_cache_type != CacheTypeCpp.SELFKONLY, "SELFKONLY is the only supported cache type for value data role" + kv_factor = 1 + else: + raise ValueError(f"Invalid data role: {data_role}") + + if local_layer_idx is None: + cache_size_per_token = (kv_factor * + sum(self.num_kv_heads_per_layer) * + self.head_dim) + else: + cache_size_per_token = ( + kv_factor * self.num_kv_heads_per_layer[local_layer_idx] * + self.head_dim) + + cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, + self.dtype) + + if data_role in [Role.KEY, Role.VALUE]: + return cache_size_bytes_per_token + + quant_size_per_token = 0 + + if self.dtype == DataType.NVFP4: + quant_size_per_token = self.calculate_scaling_factor_size_bytes( + cache_size_per_token, + quant_vector_size=16, + scaling_factor_dtype=DataType.FP8, + ) + + if data_role in [Role.KEY_BLOCK_QUANT, Role.VALUE_BLOCK_QUANT]: + return quant_size_per_token + + return cache_size_bytes_per_token + quant_size_per_token + + @staticmethod + def calculate_scaling_factor_size_bytes( + cache_size: int, quant_vector_size: int, + scaling_factor_dtype: DataType) -> int: + assert cache_size % quant_vector_size == 0, "NVFP4 cache size must be divisible by quant vector size" + return get_size_in_bytes(cache_size // quant_vector_size, + scaling_factor_dtype) + + def shutdown(self): + for kv_cache in self.kv_cache_map.values(): + kv_cache.close() + self.kv_cache_map.clear() + + def get_max_resource_count(self) -> int: + # TODO: implement this + return 1 + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + # TODO: implement this + # context_token_count = request.orig_prompt_len + # num_context_blocks = context_token_count // self.tokens_per_block + # remaining_tokens = context_token_count + request.max_new_tokens - num_context_blocks * self.tokens_per_block + # need_blocks = num_context_blocks + math.ceil( + # remaining_tokens / self.tokens_per_block) + # return need_blocks + return 0 + + # TODO: refactor get_cache_size_per_token and get_cache_bytes_per_token to use the same logic + @staticmethod + def get_cache_size_per_token(model_config: ModelConfigPython, + mapping: Mapping, **kwargs): + # get kv cache dtype bytes + mem_per_token = 2 + quant_config = model_config.quant_config + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( + ): + mem_per_token = 1 + + # get num key value heads + config = model_config.pretrained_config + num_key_value_heads = getattr(config, 'num_key_value_heads', + config.num_attention_heads) + if isinstance(num_key_value_heads, Iterable): + num_key_value_heads = sum(num_key_value_heads) / len( + num_key_value_heads) + + # get head dim + mla = hasattr(config, "kv_lora_rank") + if mla: + head_dim = config.kv_lora_rank + config.qk_rope_head_dim + kv_factor = 1 + else: + tp_size = 1 if mapping.enable_attention_dp else mapping.tp_size + head_dim = getattr(config, "head_dim", None) + if not isinstance(head_dim, int): + head_dim = config.hidden_size // config.num_attention_heads + head_dim = head_dim * num_key_value_heads // tp_size + kv_factor = 2 + + # provide at least 1 layer to prevent division by zero cache size + num_attention_layers = max( + len(mapping.pp_layers(model_config.get_num_attention_layers())), 1) + mem_per_token *= num_attention_layers * head_dim + + # K and V + mem_per_token *= kv_factor + return mem_per_token + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None): + for req in scheduled_batch.context_requests: + if req.py_request_id not in self.kv_cache_map: + continue + kv_cache = self.kv_cache_map[req.py_request_id] + if self.enable_block_reuse and not req.is_dummy_request: + if req.context_current_position > kv_cache.num_committed_tokens: + try: + kv_cache.commit( + req.get_tokens(0)[kv_cache.num_committed_tokens:req. + context_current_position]) + except Exception as e: + logger.warning( + f"Error committing tokens for request {req.py_request_id if hasattr(req, 'py_request_id') else req.request_id}: {e}" + ) + continue + kv_cache.stop_committing() + else: + kv_cache.history_length = req.context_current_position + + for req in scheduled_batch.generation_requests: + if req.py_request_id not in self.kv_cache_map: + continue + kv_cache = self.kv_cache_map[req.py_request_id] + kv_cache.history_length = req.max_beam_num_tokens - 1 + + @nvtx_range("copy_batch_block_offsets") + def copy_batch_block_offsets(self, dst_tensor: torch.Tensor, + request_ids: List[int], beam_width: int, + offset: int): + assert beam_width == 1, "beam_width must be 1 for KVCacheManager" + + batch_cache_indices = [] + + for pool_idx in range(self.num_pools): + for req_id in request_ids: + batch_cache_indices.append( + self.kv_cache_map[req_id].get_page_indices( + pool_idx).buffer_info()) + + if len(batch_cache_indices) > 0: + copy_batch_block_offsets_nanobind(dst_tensor, len(request_ids), + batch_cache_indices, + self.num_pools, offset) + + class SlotManager: def __init__(self, max_num_requests: int): diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 5757f8efbc70..2cfbdda4305d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -57,6 +57,7 @@ from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE from ..speculative.spec_tree_manager import SpecTreeManager +from ..utils import MYPYCLIB_ENABLED from .finish_reason import FinishedState from .llm_request import LlmRequest, LlmRequestState, get_draft_token_length from .resource_manager import ResourceManager, ResourceManagerType @@ -250,6 +251,10 @@ def _unwrap_singleton(p: Optional[list[T]]) -> Optional[T]: def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: + if MYPYCLIB_ENABLED: + import mypyclib + + return mypyclib._request_get_sampling_params_impl(request) sampling_config = request.sampling_config temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) @@ -263,6 +268,10 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: def _request_strategy(request: LlmRequest, *, vocab_size: int) -> Strategy: + if MYPYCLIB_ENABLED: + import mypyclib + + return mypyclib._request_strategy_impl(request, vocab_size=vocab_size) params = _request_get_sampling_params(request) return resolve_sampling_strategy(params, vocab_size=vocab_size) @@ -274,6 +283,13 @@ def _group_requests_by_strategy_key( pin_memory: bool = False, vocab_size: int, ) -> dict[tuple[GenericStrategyKeyType, bool], tuple[torch.Tensor, List[Strategy]]]: + if MYPYCLIB_ENABLED: + import mypyclib + + return mypyclib._group_requests_by_strategy_key_impl( + requests, strategy_to_key=strategy_to_key, pin_memory=pin_memory, vocab_size=vocab_size + ) + # NB: Client code relies on request indices in returned torch.Tensor being sorted. group_dict: dict[tuple[GenericStrategyKeyType, bool], tuple[list[int], list[Strategy]]] = ( defaultdict(lambda: ([], [])) @@ -1033,6 +1049,10 @@ def _process_draft_tokens_rejection_sampling( def _speculation_could_use_rejection_sampling( request: LlmRequest, strategy: Optional[Strategy] = None ) -> bool: + if MYPYCLIB_ENABLED: + import mypyclib + + return mypyclib._speculation_could_use_rejection_sampling_impl(request, strategy) if strategy is None: strategy = _request_strategy( request, @@ -1205,6 +1225,11 @@ def _apply_embedding_bias( Modifies logits in-place. """ + if MYPYCLIB_ENABLED: + import mypyclib + + return mypyclib._apply_embedding_bias_impl(logits, requests, request_steps) + # NB: Unfortunately, Torch provides no combination of torch.index_select (similar to # torch.Tensor.gather -- allows one-to-many mapping) and addition, analogous to how # torch.Tensor.scatter_add_ (and it's variant torch.Tensor.index_add_ -- allows diff --git a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py index 35e64afe4c2a..070a8cacfd99 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampling_utils.py @@ -32,6 +32,7 @@ else: from typing_extensions import override +from ..utils import MYPYCLIB_ENABLED TemperatureOnly: TypeAlias = tuple[Literal["temperature"], float] TopK: TypeAlias = tuple[Literal["top_k"], int, float] @@ -54,6 +55,11 @@ class UtilsSamplingParams: def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) -> Strategy: # The semantics are specified in the doc-string of SamplingParams + if MYPYCLIB_ENABLED: + import mypyclib + + return mypyclib.resolve_sampling_strategy_impl(params, vocab_size=vocab_size) + temperature = params.temperature top_p = params.top_p top_k = params.top_k diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler.py index c71c4596ed7f..e243d786020f 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler.py @@ -99,6 +99,40 @@ def schedule_request( self.peft_cache_manager) +class MaxUtilizationScheduler(CapacityScheduler): + + # only schedule requests has no_schedule_until_state <= state < no_schedule_after_state + no_schedule_until_state = LlmRequestState.CONTEXT_INIT + no_schedule_after_state = LlmRequestState.GENERATION_COMPLETE + + def __init__(self, max_num_requests: int, kv_cache_manager): + """ + Args: + max_num_requests: Maximum number of concurrent requests + kv_cache_manager: KV cache manager instance (KVCacheManagerV2) + """ + super(MaxUtilizationScheduler, self).__init__() + self.max_num_requests = max_num_requests + self.kv_cache_manager = kv_cache_manager + + def schedule_request( + self, active_requests: RequestList + ) -> tuple[list[LlmRequest], list[LlmRequest], list[LlmRequest]]: + + scheduled_requests = [] + + for request in active_requests: + req_state = request.state + # if request cannot be scheduled yet or request should no longer be scheduled, skip + if len(scheduled_requests) >= self.max_num_requests: + break + elif req_state.value < self.no_schedule_until_state.value or req_state.value >= self.no_schedule_after_state.value: + continue + scheduled_requests.append(request) + + return scheduled_requests, [], [] + + class GuaranteedNoEvictScheduler(CapacityScheduler): # only schedule requests has no_schedule_until_state <= state < no_schedule_after_state no_schedule_until_state = LlmRequestState.CONTEXT_INIT @@ -150,7 +184,7 @@ def schedule_request( assert len(scheduled_requests) > 0, ( "no pending request can get enough resource to complete, " "please increase KV cache pool size.") - return scheduled_requests, [] + return scheduled_requests, [], [] class MicroBatchScheduler(ABC): diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 9301578b4a2d..ea1d522b31e4 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -8,6 +8,7 @@ import torch from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.math_utils import ceil_div, pad_up from tensorrt_llm.quantization.utils import fp4_utils @@ -31,6 +32,13 @@ start=0, ) +MYPYCLIB_ENABLED = os.getenv("TRTLLM_ENABLE_MYPYCLIB", "0") == "1" + +if MYPYCLIB_ENABLED: + logger.info("mypyclib is enabled") +else: + logger.info("mypyclib is disabled") + # IMPORTANT: Keep the same order of activation functions in this enum and the enum in # cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 189b96d8d66d..8fc9b93f5361 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -208,6 +208,12 @@ def binding_to_str_dtype(binding_dtype) -> str: return ret +def binding_to_torch_dtype(binding_dtype) -> str: + ret = _binding_to_str_dtype.get(binding_dtype) + assert ret is not None, f'Unsupported binding dtype: {binding_dtype}' + return str_dtype_to_torch(ret) + + def binding_dtype_size(dtype: DataType): return _binding_dtype_size[dtype] @@ -968,7 +974,7 @@ class TensorWrapper: def __init__( self, data_ptr: int, - dtype: Union[torch.dtype, str, np.dtype, trt.DataType], + dtype: Union[torch.dtype, str, np.dtype, trt.DataType, DataType], shape: Sequence[int], strides: Optional[Sequence[int]] = None, ): @@ -990,7 +996,8 @@ def shape(self): return getattr(self, "_shape", None) @dtype.setter - def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType]): + def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType, + DataType]): if isinstance(dtype, torch.dtype): self._dtype = dtype elif isinstance(dtype, str): @@ -999,6 +1006,8 @@ def dtype(self, dtype: Union[torch.dtype, str, np.dtype, trt.DataType]): self._dtype = np_dtype_to_torch(dtype) elif isinstance(dtype, trt.DataType): self._dtype = trt_dtype_to_torch(dtype) + elif isinstance(dtype, DataType): + self._dtype = binding_to_torch_dtype(dtype) else: raise TypeError(f"Unsupported dtype: {dtype}") diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index ede008e12d40..8f8f88fed879 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1509,6 +1509,15 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): tokens_per_block: int = Field(default=32, description="The number of tokens per block.") + use_kv_cache_manager_v2: bool = Field( + default=False, description="Whether to use the v2 KV cache.") + + max_util_for_resume: float = Field( + default=0.9, + description= + "The maximum utilization of the KV cache for resume. Default is 90%. Only used when using KV cache manager v2." + ) + def _to_pybind(self): return _KvCacheConfig( enable_block_reuse=self.enable_block_reuse, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py index c897dd9f0c2e..fb9046378e4f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py @@ -1,9 +1,11 @@ import atexit +import os import sys import threading from _thread import LockType from collections.abc import Callable, Iterator from dataclasses import dataclass +import torch # avoid importing the whole tensorrt_llm module, which takes time during debugging. from importlib.util import find_spec @@ -18,6 +20,7 @@ if "tensorrt_llm" in sys.modules: from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( # noqa # type: ignore + BlockIndices, DiskAddress, DiskToDiskTask, DiskToHostTask, @@ -31,12 +34,16 @@ copy_host_to_disk, copy_host_to_host, ) + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets as nb_copy_batch_block_offsets, +) else: # fast path for dev, avoids importing the whole tensorrt_llm module spec = find_spec("kv_cache_manager_v2") assert spec is not None and spec.origin is not None with DynamicPathManager(str(Path(spec.origin).parent.parent.parent)): from bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( # noqa + BlockIndices, DiskAddress, DiskToDiskTask, DiskToHostTask, @@ -50,6 +57,9 @@ copy_host_to_disk, copy_host_to_host, ) + from bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets as nb_copy_batch_block_offsets, +) class CopyTask(NamedTuple): @@ -179,6 +189,22 @@ def get_copier(dst: CacheTier, src: CacheTier) -> Copier | HomoTuple[Copier]: return copiers[dst][src] +def copy_batch_block_offsets( + dst_tensor: torch.Tensor, + batch_size: int, + batch_cache_indices: list[tuple[int, int]], + num_pools: int, + offset: int, +) -> None: + nb_copy_batch_block_offsets( + dst_tensor, + batch_size, + [BlockIndices(addr, length) for addr, length in batch_cache_indices], + num_pools, + offset, + ) + + @dataclass(slots=True) class GrainMetadata: mutex: LockType diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 5cd51296e91c..fc2a2252613c 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -528,6 +528,7 @@ def resume(self, cuda_stream: CudaStream | None = None) -> bool: try: locks = batched_lock_to_gpu(self, tasks) except OutOfPagesError: + print("OutOfPagesError") return False for (ordinal, beam_idx, lc_idx), lock in zip(self._active_pages(), locks): beam_block = self._block(ordinal, beam_idx) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index 5bd6405fdfc7..25b8eb8d1756 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py @@ -149,7 +149,7 @@ def layer_grouping(self) -> HomoTuple[HomoTuple[LayerId]]: return tuple(tuple(grouping[i]) for i in typed_range(num_life_cycles)) # @TODO: need updating when dynamic resizing is supported. - def clamp_max_seq_len_for_mem(self, batch_size: int, model_max_seq_len: int) -> int: + def clamp_max_seq_len_for_mem(self, batch_size: int) -> int: "Get the max possible sequence length limited by the GPU memory pools." assert batch_size > 0 tokens_per_block = self.tokens_per_block @@ -184,13 +184,14 @@ def is_enough(num_blocks: int) -> bool: assert is_enough(1) lb = 1 - ub = div_up(model_max_seq_len, tokens_per_block) - if is_enough(ub): - return model_max_seq_len - while lb < ub: + ub = lb + while is_enough(ub): + lb = ub + ub *= 2 + while lb < ub - 1: mid = (lb + ub) // 2 if is_enough(mid): lb = mid else: - ub = mid - 1 - return min(lb * tokens_per_block, model_max_seq_len) + ub = mid + return lb * tokens_per_block diff --git a/tests/aa/__init__.py b/tests/aa/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/aa/fake_engine.py b/tests/aa/fake_engine.py new file mode 100644 index 000000000000..cbdc42c8a214 --- /dev/null +++ b/tests/aa/fake_engine.py @@ -0,0 +1,190 @@ +import itertools +from collections.abc import Sequence +from functools import cached_property +from importlib.util import find_spec +from typing import NamedTuple + +if find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2 import ( + BeamIndex, + CudaStream, + KVCacheManagerConfig, + LayerId, + TokenIdExt, + _KVCache, + ) + from kv_cache_manager_v2._common import BAD_PAGE_INDEX, NDEBUG, MemAddress + from kv_cache_manager_v2._config import AttentionLayerConfig, DataRole + from kv_cache_manager_v2._utils import ( + div_up, + exact_div, + get_uniform_attribute, + overlap, + typed_range, + value_or, + ) +else: + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BeamIndex, + CudaStream, + KVCacheManagerConfig, + LayerId, + TokenIdExt, + _KVCache, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX, NDEBUG, MemAddress + from tensorrt_llm.runtime.kv_cache_manager_v2._config import AttentionLayerConfig, DataRole + from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + div_up, + exact_div, + get_uniform_attribute, + overlap, + typed_range, + value_or, + ) + +import os + +from dynamic_path_manager import DynamicPathManager + +with DynamicPathManager(os.path.dirname(os.path.abspath(__file__))): + from kernels import check_values, fill_values + +class Step(NamedTuple): + kv_cache: _KVCache + input: list[TokenIdExt] # when empty, just check history + history: list[TokenIdExt] + + +class Role: + """Constants for data roles in KV cache management.""" + + KEY = DataRole("key") + VALUE = DataRole("value") + KEY_BLOCK_QUANT = DataRole("key_block_quant") + VALUE_BLOCK_QUANT = DataRole("value_block_quant") + + +roles = (Role.KEY, Role.VALUE, Role.KEY_BLOCK_QUANT, Role.VALUE_BLOCK_QUANT) + + +class FakeEngine: + cfg: KVCacheManagerConfig + + def __init__(self, config: KVCacheManagerConfig) -> None: + super().__init__() + self.cfg = config + + @property + def tokens_per_block(self) -> int: + return self.cfg.tokens_per_block + + @cached_property + def layers(self) -> dict[LayerId, AttentionLayerConfig]: + return { + layer.layer_id: layer + for layer in sorted(self.cfg.layers, key=lambda layer: layer.layer_id) + } + + def execute(self, batch: Sequence[Step], stream: CudaStream) -> None: + assert batch + manager = get_uniform_attribute(batch, lambda step: step.kv_cache.manager) + for kv_cache, input, history in batch: + for layer_id, layer_cfg in self.layers.items(): + for buf_id, buf in enumerate(layer_cfg.buffers): + role = buf.role + assert NDEBUG or buf.size == manager.get_page_stride(layer_id, role) + for beam in typed_range(kv_cache.beam_width): + # check history + self._check_pages(kv_cache, layer_id, buf_id, beam, history, stream) + # write new token + if input: + self._write_new_tokens( + kv_cache, len(history), layer_id, buf_id, beam, input, stream + ) + + def _check_pages( + self, + kv_cache: _KVCache, + layer_id: LayerId, + buf_id: int, + beam: BeamIndex, + history: Sequence[TokenIdExt], + stream: CudaStream, + ): + manager = kv_cache.manager + tokens_per_block = self.tokens_per_block + layer_cfg = self.layers[layer_id] + buf = layer_cfg.buffers[buf_id] + role = buf.role + token_bytes = exact_div(buf.size, tokens_per_block) + pool = manager.get_mem_pool_base_address(layer_id, role) + print(f"pool: {pool}, layer_id: {layer_id}, role: {role}") + stride = manager.get_page_stride(layer_id, role) + lc_id = manager._storage._layer_to_life_cycle_ids[layer_id] + pages = kv_cache.get_page_indices(lc_id, beam) + capacity = kv_cache.capacity + history_len = len(history) + assert len(history) == history_len + window = ( + (0, capacity) + if layer_cfg.window_size is None + else (max(0, history_len + 1 - layer_cfg.window_size), capacity) + ) + sink = value_or(layer_cfg.num_sink_tokens, 0) + # check history + for ordinal, page in enumerate(pages): + if page == BAD_PAGE_INDEX: + continue + page_range = (tokens_per_block * ordinal, tokens_per_block * (ordinal + 1)) + need_page = overlap(page_range, (0, sink)) or overlap(page_range, window) + if need_page: + assert page != BAD_PAGE_INDEX + else: + assert kv_cache.history_length != history_len or page == BAD_PAGE_INDEX + addr = MemAddress(pool + stride * page) + tokens = history[tokens_per_block * ordinal : tokens_per_block * (ordinal + 1)] + check_values(addr, token_bytes, layer_id, buf_id, beam, tokens, stream) + + def _write_new_tokens( + self, + kv_cache: _KVCache, + history_len: int, + layer_id: LayerId, + buf_id: int, + beam: BeamIndex, + input: Sequence[TokenIdExt], + stream: CudaStream, + ): + manager = kv_cache.manager + tokens_per_block = self.tokens_per_block + layer_cfg = self.layers[layer_id] + buf = layer_cfg.buffers[buf_id] + role = buf.role + token_bytes = exact_div(buf.size, self.tokens_per_block) + pool = manager.get_mem_pool_base_address(layer_id, role) + stride = manager.get_page_stride(layer_id, role) + lc_id = manager._storage._layer_to_life_cycle_ids[layer_id] + pages = kv_cache.get_page_indices(lc_id, beam)[ + : div_up(history_len + len(input), tokens_per_block) + ] + capacity = kv_cache.capacity + input_range = (history_len, history_len + len(input)) + assert input_range[1] <= capacity + ordinal_beg = input_range[0] // tokens_per_block + pages = itertools.islice(pages, ordinal_beg, None) + ordinal = None + for i, page in enumerate(pages): + ordinal = ordinal_beg + i + assert page != BAD_PAGE_INDEX + page_range = (tokens_per_block * ordinal, tokens_per_block * (ordinal + 1)) + batch_range = tuple(i for i in overlap(input_range, page_range)) + assert batch_range + tokens = input[(batch_range[0] - history_len) : (batch_range[1] - history_len)] + addr = MemAddress( + pool + stride * page + token_bytes * (batch_range[0] % tokens_per_block) + ) + # print('layer_id={}, buf_id={}, beam={}, i={}, addr={}, tokens={}'.format( + # layer_id, buf_id, beam, i, addr, tokens)) + fill_values(addr, token_bytes, layer_id, buf_id, beam, tokens, stream) + assert ordinal is None or ordinal + 1 == div_up(input_range[1], tokens_per_block) diff --git a/tests/aa/kernels.py b/tests/aa/kernels.py new file mode 100644 index 000000000000..2f37b6097877 --- /dev/null +++ b/tests/aa/kernels.py @@ -0,0 +1,318 @@ +import contextlib +import ctypes +from collections.abc import Sequence +from functools import lru_cache +from importlib.util import find_spec +from typing import Final, Iterator + +import cuda.bindings.driver as drv +from cuda.core.experimental import Kernel, Program, ProgramOptions +from cuda.core.experimental._module import ObjectCode + +if find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2._common import CudaStream, LayerId, MemAddress, TokenIdExt + from kv_cache_manager_v2._utils import _unwrap, div_up, exact_div +else: + from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( + CudaStream, + LayerId, + MemAddress, + TokenIdExt, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2._utils import _unwrap, div_up, exact_div + +_SLEEP_TIME_NS: Final[int] = 0 + + +@contextlib.contextmanager +def enable_kernel_delay() -> Iterator[None]: + global _SLEEP_TIME_NS + _SLEEP_TIME_NS = 30_000 + yield + _SLEEP_TIME_NS = 0 + + +@lru_cache(maxsize=None) +def get_program(debug: bool, max_tokens: int, sleep_time: int) -> ObjectCode: + assert max_tokens > 0 and (max_tokens & (max_tokens - 1)) == 0, ( + "max_tokens must be a power of 2" + ) + code = r""" +#if !defined(__CUDACC_RTC__) +#include +#include +#endif + +#ifdef NDEBUG +__device__ inline void check(bool condition) { + if (!condition) { + asm volatile("trap;" ::: "memory"); + } +} +#else +#define check assert +#endif + +using uint32_t = unsigned int; +using uint16_t = unsigned short; + +constexpr uint32_t sleepTime = SLEEP_TIME_NS; + +struct alignas(16) Value { + uint32_t token; + uint32_t layer; + uint32_t role; + uint32_t beam; + + __device__ inline bool operator==(Value const& other) const { + return token == other.token && layer == other.layer && role == other.role && beam == other.beam; + } + __device__ inline bool operator!=(Value const& other) const { + return !(*this == other); + } +}; + +constexpr uint32_t kMAX_TOKENS = MAX_TOKENS; + +struct Tokens { + uint32_t tokens[kMAX_TOKENS]; +}; + +extern "C" __global__ void fillValues(Value* data, uint32_t valuesPerToken, uint32_t layer, + uint32_t buf_id, uint32_t beam, __grid_constant__ const Tokens tokens, uint32_t numTokens) { + if (sleepTime > 0) { + __nanosleep(sleepTime); + } + check(numTokens <= kMAX_TOKENS); + auto const tid = (static_cast(blockIdx.x) * blockDim.x) + threadIdx.x; + auto const idxToken = tid / valuesPerToken; + if (idxToken >= numTokens) { + return; + } + auto const token = tokens.tokens[idxToken]; + auto const value = Value{token, layer, buf_id, beam}; + data[tid] = value; +} + +__device__ inline void assertEq(Value const& a, Value const& b) { +#ifndef NDEBUG + if (a != b) { + printf("(%d, %d, %d, %d) != (%d, %d, %d, %d)\n", + a.token, a.layer, a.role, a.beam, + b.token, b.layer, b.role, b.beam); + } +#endif + check(a == b); +} + +extern "C" __global__ void checkValues(Value const* data, uint32_t valuesPerToken, uint32_t layer, + uint32_t buf_id, uint32_t beam, __grid_constant__ const Tokens tokens, uint32_t numTokens) { + if (sleepTime > 0) { + __nanosleep(sleepTime); + } + check(numTokens <= kMAX_TOKENS); + auto const tid = (static_cast(blockIdx.x) * blockDim.x) + threadIdx.x; + auto const idxToken = tid / valuesPerToken; + if (idxToken >= numTokens) { + return; + } + auto const token = tokens.tokens[idxToken]; + auto const value = Value{token, layer, buf_id, beam}; + assertEq(data[tid], value); +} + """ + macros = [("MAX_TOKENS", str(max_tokens)), ("SLEEP_TIME_NS", str(sleep_time))] + program_options = ProgramOptions(std="c++17", lineinfo=True, debug=debug, define_macro=macros) # type: ignore[arg-type] + if not debug: + program_options.use_fast_math = True + prog = Program(code, code_type="c++", options=program_options) + mod = prog.compile("cubin", name_expressions=("fillValues", "checkValues")) + return mod + + +def get_kernel(name: str, num_tokens: int, sleep_time: int) -> tuple[Kernel, int]: + assert num_tokens > 0 + + @lru_cache(maxsize=None) + def impl(name: str, max_tokens: int, sleep_time: int) -> Kernel: + assert name in ("fillValues", "checkValues") + assert max_tokens != 0 and (max_tokens & (max_tokens - 1)) == 0, ( + "max_tokens must be a power of 2" + ) + debug = False + # debug = not NDEBUG + return get_program(debug, max_tokens, sleep_time).get_kernel(name) + + # Round up to the next power of two + max_tokens = 2 ** ((num_tokens - 1).bit_length()) + return impl(name, max_tokens, sleep_time), max_tokens + + +class Value(ctypes.Structure): + _fields_ = [ + ("token", ctypes.c_uint32), + ("layer", ctypes.c_uint32), + ("buf_id", ctypes.c_uint32), + ("beam", ctypes.c_uint32), + ] + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Value): + return NotImplemented + return ( + self.token == other.token + and self.layer == other.layer + and self.buf_id == other.buf_id + and self.beam == other.beam + ) + + def __str__(self) -> str: + return ( + f"Value(token={self.token}, layer={self.layer}, buf_id={self.buf_id}, beam={self.beam})" + ) + + +@lru_cache(maxsize=None) +def _get_ctypes_struct(max_tokens: int) -> type[ctypes.Structure]: + class Tokens(ctypes.Structure): + _fields_ = [ + ("tokens", ctypes.c_uint32 * max_tokens), + ] + + Tokens.__name__ = f"Tokens_{max_tokens}" + return Tokens + + +def _make_tokens(tokens: Sequence[TokenIdExt], max_tokens: int) -> ctypes.Structure: + assert len(tokens) <= max_tokens + padded = list(tokens) + [0] * (max_tokens - len(tokens)) + Tokens = _get_ctypes_struct(max_tokens) + return Tokens( + tokens=(ctypes.c_uint32 * max_tokens)( + *[ + t if isinstance(t, int) else int.from_bytes(t[:4], "little", signed=False) + for t in padded + ] + ) + ) + + +def fill_values( + address: MemAddress, + bytes_per_token: int, + layer: LayerId, + buf_id: int, + beam: int, + tokens: Sequence[TokenIdExt], + stream: CudaStream, +): + values_per_token = exact_div(bytes_per_token, ctypes.sizeof(Value)) + num_tokens = len(tokens) + if num_tokens == 0: + return + kernel, max_tokens = get_kernel("fillValues", len(tokens), _SLEEP_TIME_NS) + args = ( + address, + values_per_token, + layer, + buf_id, + beam, + _make_tokens(tokens, max_tokens), + num_tokens, + ) + arg_types = ( + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + None, + ctypes.c_uint32, + ) + num_threads = values_per_token * num_tokens + cta_size = 256 + _unwrap( + drv.cuLaunchKernel( + kernel._handle, + div_up(num_threads, cta_size), + 1, + 1, + cta_size, + 1, + 1, + 0, + stream, + (args, arg_types), + 0, + ) + ) + + +def check_values( + address: MemAddress, + bytes_per_token: int, + layer: LayerId, + buf_id: int, + beam: int, + tokens: Sequence[TokenIdExt], + stream: CudaStream, +): + values_per_token = exact_div(bytes_per_token, ctypes.sizeof(Value)) + num_tokens = len(tokens) + if num_tokens == 0: + return + kernel, max_tokens = get_kernel("checkValues", len(tokens), _SLEEP_TIME_NS) + args = ( + address, + values_per_token, + layer, + buf_id, + beam, + _make_tokens(tokens, max_tokens), + num_tokens, + ) + arg_types = ( + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + ctypes.c_uint32, + None, + ctypes.c_uint32, + ) + num_threads = values_per_token * num_tokens + cta_size = 256 + _unwrap( + drv.cuLaunchKernel( + kernel._handle, + div_up(num_threads, cta_size), + 1, + 1, + cta_size, + 1, + 1, + 0, + stream, + (args, arg_types), + 0, + ) + ) + + +def debug_dump_tokens( + addr: MemAddress, token_bytes: int, num_tokens: int, stream: CudaStream +) -> Iterator[Value]: + if num_tokens == 0: + return + val_size = ctypes.sizeof(Value) + values_per_token = exact_div(token_bytes, val_size) + host_buf = (Value * values_per_token * num_tokens)() + ptr = ctypes.addressof(host_buf) + _unwrap(drv.cuMemcpyDtoHAsync(ptr, addr, num_tokens * token_bytes, stream)) + _unwrap(drv.cuStreamSynchronize(stream)) + for i in range(num_tokens): + token = host_buf[i] + value = Value.from_buffer_copy(token[0]) + for j in range(1, values_per_token): + assert token[j] == token[0] + yield value diff --git a/tests/aa/test_attention.py b/tests/aa/test_attention.py new file mode 100644 index 000000000000..f4b427ff93d7 --- /dev/null +++ b/tests/aa/test_attention.py @@ -0,0 +1,788 @@ +import math +import unittest +from dataclasses import dataclass +from typing import Callable, Optional, Sequence + +import flashinfer +import pytest +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend import (AttentionBackend, + FlashInferAttention, + VanillaAttention) +from tensorrt_llm._torch.attention_backend.interface import \ + PredefinedAttentionMask +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import \ + KVCacheManagerV2 as KVCacheManager +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +try: + from cuda.bindings import driver as cuda + from cuda.bindings import nvrtc +except ImportError: + from cuda import cuda, nvrtc + + +def ASSERT_DRV(err): + if isinstance(err, cuda.CUresult): + if err != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError("Cuda Error: {}".format(err)) + elif isinstance(err, nvrtc.nvrtcResult): + if err != nvrtc.nvrtcResult.NVRTC_SUCCESS: + raise RuntimeError("Nvrtc Error: {}".format(err)) + else: + raise RuntimeError("Unknown error type: {}".format(err)) + + +def getSMVersion(): + # Init + (err, ) = cuda.cuInit(0) + ASSERT_DRV(err) + + # Device + err, cuDevice = cuda.cuDeviceGet(0) + ASSERT_DRV(err) + + # Get target architecture + err, sm_major = cuda.cuDeviceGetAttribute( + cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, + cuDevice) + ASSERT_DRV(err) + err, sm_minor = cuda.cuDeviceGetAttribute( + cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, + cuDevice) + ASSERT_DRV(err) + + return sm_major * 10 + sm_minor + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, + None, :, :].expand(batch, num_key_value_heads, + n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, + head_dim) + + +atol = 1e-2 +rtol = 1e-3 +fp8_atol = 5e-2 + + +@dataclass(kw_only=True, frozen=True) +class Scenario: + dtype: torch.dtype = torch.float16 + kvcache_dtype: torch.dtype = torch.float16 + num_layers: int + num_heads: int = 64 + num_kv_heads: int = 16 + head_dim: int = 128 + page_size: int = 256 + """flash-attention requires `page_size` to be a multiple of 256""" + num_pages: int = 4 + qo_len: int = 32 + """setting kv_len to non-zero to test cross attention""" + kv_len: int = 0 + causal: bool = True + batch_size: int = 7 + + @property + def cross(self) -> bool: + return self.kv_len != 0 + + @property + def kv_len_resolved(self) -> int: + return self.kv_len or self.qo_len + + @property + def num_kv_groups(self) -> int: + return self.num_heads // self.num_kv_heads + + @property + def kv_cache_len(self) -> int: + return self.page_size * self.num_pages + + @property + def past_kv_len(self) -> int: + return self.kv_cache_len - self.kv_len_resolved + + @property + def max_num_pages(self) -> int: + return self.batch_size * self.num_pages + + @property + def nnz_qo(self): + return self.batch_size * self.qo_len + + @property + def nnz_kv(self): + return self.batch_size * self.kv_len_resolved + + def __post_init__(self) -> None: + assert self.kv_len <= self.kv_cache_len, "KV len larger than cache len" + assert (self.kv_len != 0 or self.qo_len + <= self.kv_cache_len), "Seq len larger than cache len" + assert not (self.cross + and self.causal), "Cross attention cannot be causal" + + +@dataclass(kw_only=True, frozen=True) +class PagedScenario(Scenario): + num_generations: int + + @property + def num_contexts(self) -> int: + return self.batch_size - self.num_generations + + @property + def num_ctx_q_tokens(self) -> int: + return self.num_contexts * self.qo_len + + @property + def num_ctx_kv_tokens(self) -> int: + return self.num_contexts * self.kv_len_resolved + + @property + def nnz_qo(self) -> int: + return self.num_ctx_q_tokens + self.num_generations + + @property + def nnz_kv(self) -> int: + n = self.num_ctx_kv_tokens + if not self.cross: + n += self.num_generations + return n + + +paged_backends = { + VanillaAttention: False, + FlashInferAttention: True, +} + + +def kv_cache_manager_from(Attention: type[AttentionBackend], + s: Scenario) -> KVCacheManager: + paged = paged_backends[Attention] + + num_blocks = s.max_num_pages if paged else s.batch_size + tokens_per_block = s.page_size if paged else s.kv_cache_len + num_layers = s.num_layers + num_kv_heads = s.num_kv_heads + head_dim = s.head_dim + max_seq_len = num_blocks * tokens_per_block + batch_size = s.batch_size + + if s.kvcache_dtype == torch.float16: + kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF + elif s.kvcache_dtype == torch.bfloat16: + kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 + elif s.kvcache_dtype == torch.float8_e4m3fn: + kv_cache_dtype = tensorrt_llm.bindings.DataType.FP8 + else: + raise ValueError("Invalid dtype for unit test") + + kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block) + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + + cache_type = (tensorrt_llm.bindings.internal.batch_manager.CacheType.CROSS + if s.cross else + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF) + + result = KVCacheManager( + kv_cache_config, + cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + ) + + return result + + +def produce_outputs( + Attention: type[AttentionBackend], + q_at_layer: torch.Tensor, + kv: Optional[torch.Tensor], + s: Scenario, + *, + kv_cache: torch.Tensor, + num_cached_tokens: Callable[[int], int] | int, + num_contexts: int | None = None, + seq_lens: torch.Tensor, + seq_lens_kv: Optional[torch.Tensor] = None, + quant_config: Optional[QuantConfig] = None, +) -> list[torch.Tensor]: + num_cached_tokens_per_seq = [(num_cached_tokens if isinstance( + num_cached_tokens, int) else num_cached_tokens(i)) + for i in range(s.batch_size)] + + kv_cache_params = KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=num_cached_tokens_per_seq) + kv_cache_manager = kv_cache_manager_from(Attention, s) + request_ids = list(range(s.batch_size)) + seq_lens_append = seq_lens_kv if seq_lens_kv is not None else seq_lens + token_nums = (torch.tensor(num_cached_tokens_per_seq) + + seq_lens_append).tolist() + kv_cache_manager.add_dummy_requests(request_ids, token_nums) + for layer_idx in range(s.num_layers): + buffer = kv_cache_manager.get_buffers(layer_idx) + batch_cache_indices_flatten = [ + i for j in kv_cache_manager.get_batch_cache_indices(request_ids, 0) + for i in j + ] + print(batch_cache_indices_flatten) + kv_cache_layer = kv_cache[layer_idx] + print(buffer.shape, buffer.dtype, kv_cache_layer.shape, + kv_cache_layer.dtype) + print(buffer.data_ptr(), buffer[1].data_ptr(), kv_cache.data_ptr()) + for ii, idx in enumerate(batch_cache_indices_flatten): + buffer[idx].copy_(kv_cache_layer[ii]) + + metadata = Attention.Metadata( + num_contexts=num_contexts if num_contexts is not None else s.batch_size, + kv_cache_params=kv_cache_params, + seq_lens=seq_lens, + seq_lens_kv=seq_lens_kv, + max_num_requests=s.batch_size, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=token_nums, + ) + metadata.prepare() + mask = PredefinedAttentionMask.CAUSAL if s.causal else PredefinedAttentionMask.FULL + outputs = [] + for i in range(s.num_layers): + q = q_at_layer[i] + if kv is not None: + k, v = kv[i][0], kv[i][1] + else: + k, v = None, None + attention = Attention( + layer_idx=i, + num_heads=s.num_heads, + num_kv_heads=s.num_kv_heads, + head_dim=s.head_dim, + quant_config=quant_config, + ) + o = attention.forward(q, k, v, metadata, attention_mask=mask) + assert list(o.shape) == [s.nnz_qo, s.num_heads * s.head_dim] + outputs.append(o) + kv_cache_manager.shutdown() + return outputs + + +def allclose( + ref: Sequence[torch.Tensor], + impls: dict[str, Sequence[torch.Tensor]], + *, + layer=0, + atol=atol, + rtol=rtol, +): + for name, outputs in impls.items(): + print(f"{name} output: ", float(outputs[layer].abs().mean())) + print("ref outputs: ", float(ref[layer].abs().mean())) + for name, outputs in impls.items(): + print(f"{name} & ref diff: ", + float((ref[layer] - outputs[layer]).abs().mean())) + for name, outputs in impls.items(): + torch.testing.assert_close( + outputs[layer], + ref[layer], + atol=atol, + rtol=rtol, + ), + + +def test_flashinfer_prefill(): + s = Scenario(num_layers=1) + dtype = s.dtype + num_layers = s.num_layers + num_qo_heads = s.num_heads + num_kv_heads = s.num_kv_heads + num_kv_groups = s.num_kv_groups + head_dim = s.head_dim + page_size = s.page_size + num_pages = s.num_pages + kv_cache_len = s.kv_cache_len + qo_len = s.qo_len + past_kv_len = s.past_kv_len + batch_size = s.batch_size + nnz_qo = s.nnz_qo + max_num_pages = s.max_num_pages + + # allocate 128MB workspace buffer + workspace_buffer = torch.empty(128 * 1024 * 1024, + dtype=torch.uint8, + device="cuda") + paged_wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, "NHD", backend="fa2") + ragged_wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, "NHD", backend="fa2") + paged_kv_indices = torch.arange(max_num_pages).int().cuda() + paged_kv_indptr = torch.arange(0, batch_size + 1).int().cuda() * num_pages + # 1 <= paged_kv_last_page_len <= page_size + paged_kv_last_page_len = torch.full((batch_size, ), page_size).int().cuda() + qo_indptr = torch.arange(0, batch_size + 1).int().to("cuda") * qo_len + kv_indptr = torch.arange(0, batch_size + 1).int().to("cuda") * kv_cache_len + + q_at_layer = (torch.randn(num_layers, + nnz_qo, + num_qo_heads, + head_dim, + device="cuda").to(dtype).cuda()) + kv_cache_at_layer = torch.randn(num_layers, + max_num_pages, + 2, + page_size, + num_kv_heads, + head_dim, + device="cuda").to(s.kvcache_dtype) + kv_data = (kv_cache_at_layer.transpose(1, 2).contiguous().view( + num_layers, 2, batch_size, kv_cache_len, num_kv_heads, head_dim)) + + causal_mask = torch.full( + (qo_len, kv_cache_len), + fill_value=torch.finfo(dtype).min, + dtype=dtype, + device="cuda", + ) + cache_position = torch.arange(past_kv_len, kv_cache_len).cuda() + bool_causal_mask = torch.arange( + kv_cache_len).cuda() <= cache_position.reshape(-1, 1) + causal_mask *= ~bool_causal_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + + # create auxiliary data structures for batch prefill attention + paged_wrapper.plan( + qo_indptr, + paged_kv_indptr, + paged_kv_indices, + paged_kv_last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + causal=True, + ) + ragged_wrapper.plan( + qo_indptr, + kv_indptr, + num_qo_heads, + num_kv_heads, + head_dim, + causal=True, + ) + flashinfer_outputs = [] + for i in range(num_layers): + q = q_at_layer[i] + kv_cache = kv_cache_at_layer[i] + o = paged_wrapper.run(q, kv_cache) + k = kv_data[i][0] + v = kv_data[i][1] + k = k.view(-1, num_kv_heads, head_dim) + v = v.view(-1, num_kv_heads, head_dim) + ragged_o = ragged_wrapper.run(q, k, v) + assert list(o.shape) == [nnz_qo, num_qo_heads, head_dim] + print("paged output: ", float(o.abs().mean())) + print("ragged output: ", float(ragged_o.abs().mean())) + print("paged & ragged diff: ", float((ragged_o - o).abs().mean())) + assert torch.allclose(o, ragged_o, atol=atol, rtol=rtol) + flashinfer_outputs.append(o) + + sdpa_outputs = [] + for i in range(num_layers): + q = q_at_layer[i] + k = kv_data[i][0] + v = kv_data[i][1] + q = q.view(batch_size, qo_len, num_qo_heads, head_dim) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + k = repeat_kv(k, num_kv_groups) + v = repeat_kv(v, num_kv_groups) + o = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=causal_mask, + ) + o = o.transpose(1, 2).contiguous().view(nnz_qo, num_qo_heads, head_dim) + sdpa_outputs.append(o) + + ref_outputs = [] + for i in range(num_layers): + q = q_at_layer[i] + k = kv_data[i][0] + v = kv_data[i][1] + q = q.view(batch_size, qo_len, num_qo_heads, head_dim) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + k = repeat_kv(k, num_kv_groups) + v = repeat_kv(v, num_kv_groups) + + attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_dim) + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = torch.nn.functional.softmax(attn_weights, + dim=-1, + dtype=torch.float32).to( + q.dtype) + o = torch.matmul(attn_weights, v) + + o = o.transpose(1, 2).contiguous().view(nnz_qo, num_qo_heads, head_dim) + ref_outputs.append(o) + + allclose( + ref_outputs, + { + "flashinfer": flashinfer_outputs, + "sdpa": sdpa_outputs, + }, + ) + + +@pytest.mark.parametrize( + "s", + [ + Scenario(num_layers=1), + Scenario(num_layers=1, causal=False), + Scenario(num_layers=1, qo_len=32, kv_len=32, causal=False), + Scenario(num_layers=1, qo_len=32, kv_len=64, causal=False), + ], + ids=["typical", "non-causal", "cross", "cross-diff-kv-len"], +) +def test_attention_backend(s: Scenario): + dtype = s.dtype + num_layers = s.num_layers + num_heads = s.num_heads + num_kv_heads = s.num_kv_heads + num_kv_groups = s.num_kv_groups + head_dim = s.head_dim + page_size = s.page_size + kv_cache_len = s.kv_cache_len + qo_len = s.qo_len + kv_len = s.kv_len_resolved + past_kv_len = s.past_kv_len + batch_size = s.batch_size + nnz_qo = s.nnz_qo + nnz_kv = s.nnz_kv + causal = s.causal + + q_at_layer = torch.randn(num_layers, + nnz_qo, + num_heads * head_dim, + device="cuda").to(dtype) + flashinfer_kv_cache = torch.randn(num_layers, + s.max_num_pages, + 2, + page_size, + num_kv_heads, + head_dim, + device="cuda").to(s.kvcache_dtype) + ref_kv_cache = (flashinfer_kv_cache.transpose(1, 2).contiguous().view( + num_layers, 2, batch_size, kv_cache_len, num_kv_heads, head_dim)) + kv = torch.randn(num_layers, + 2, + nnz_kv, + num_kv_heads * head_dim, + device="cuda").to(dtype) + + def produce(Attention: type[AttentionBackend], kv_cache: torch.Tensor): + return produce_outputs( + Attention, + q_at_layer, + kv, + s, + kv_cache=kv_cache, + num_cached_tokens=past_kv_len, + seq_lens=torch.full((batch_size, ), qo_len).int(), + seq_lens_kv=torch.full( + (batch_size, ), kv_len).int() if s.cross else None, + ) + + flashinfer_outputs = produce(FlashInferAttention, flashinfer_kv_cache) + sdpa_outputs = produce(VanillaAttention, ref_kv_cache.transpose(1, 2)) + + # Test reference attention + if causal: + causal_mask = torch.full( + (qo_len, kv_cache_len), + fill_value=torch.finfo(dtype).min, + dtype=dtype, + device="cuda", + ) + cache_position = torch.arange(past_kv_len, kv_cache_len).cuda() + bool_causal_mask = torch.arange( + kv_cache_len).cuda() <= cache_position.reshape(-1, 1) + causal_mask *= ~bool_causal_mask + causal_mask = causal_mask[None, + None, :, :].expand(batch_size, 1, -1, -1) + else: + causal_mask = 0 + + ref_outputs = [] + for i in range(num_layers): + q = q_at_layer[i] + ref_kv_cache[i][:, :, past_kv_len:kv_cache_len] = kv[i].view( + 2, batch_size, kv_len, num_kv_heads, head_dim) + k = ref_kv_cache[i][0] + v = ref_kv_cache[i][1] + q = q.view(batch_size, qo_len, num_heads, head_dim) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + k = repeat_kv(k, num_kv_groups) + v = repeat_kv(v, num_kv_groups) + + attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_dim) + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = torch.nn.functional.softmax(attn_weights, + dim=-1, + dtype=torch.float32).to( + q.dtype) + o = torch.matmul(attn_weights, v) + + o = o.transpose(1, 2).contiguous().view(nnz_qo, num_heads * head_dim) + ref_outputs.append(o) + + allclose( + ref_outputs, + { + "flashinfer": flashinfer_outputs, + "sdpa": sdpa_outputs, + }, + ) + + del flashinfer_kv_cache + del ref_kv_cache + torch.cuda.empty_cache() + + +def generate_causal_mask(seq_lens, qo_lens, batch_size, dtype): + causal_masks = [] + max_seq_len = int(seq_lens.max()) + max_qo_len = int(qo_lens.max()) + for i in range(batch_size): + kv_len = seq_lens[i] + qo_len = qo_lens[i] + past_kv_len = kv_len - qo_len + causal_mask = torch.full( + (qo_len, kv_len), + fill_value=torch.finfo(dtype).min, + dtype=dtype, + device="cuda", + ) + cache_position = torch.arange(past_kv_len, kv_len).cuda() + causal_mask *= torch.arange(kv_len).cuda() > cache_position.reshape( + -1, 1) + causal_mask = torch.nn.functional.pad( + causal_mask, + (0, max_seq_len - kv_len, 0, max_qo_len - qo_len), + "constant", + torch.finfo(dtype).min, + ) + causal_masks.append(causal_mask) + causal_mask = torch.stack(causal_masks).view(batch_size, 1, max_qo_len, + max_seq_len) + return causal_mask + + +@pytest.mark.parametrize( + "s", + [ + PagedScenario(num_layers=4, num_generations=5), + PagedScenario(num_layers=4, num_generations=5, kv_len=64, causal=False), + PagedScenario( + num_layers=4, num_generations=5, kvcache_dtype=torch.float8_e4m3fn), + PagedScenario( + num_layers=4, + num_generations=5, + kv_len=64, + causal=False, + kvcache_dtype=torch.float8_e4m3fn, + ), + ], + ids=["fp16", "fp16-cross", "fp8", "fp8-cross"], +) +def test_attention_backend_ifb(s: PagedScenario): + dtype = s.dtype + is_fp8 = s.kvcache_dtype == torch.float8_e4m3fn + if is_fp8 and getSMVersion() < 89: + pytest.skip("This test is not supported in pre-Ada architecture.") + torch.manual_seed(0) + num_layers = s.num_layers + num_heads = s.num_heads + num_kv_heads = s.num_kv_heads + num_kv_groups = s.num_kv_groups + head_dim = s.head_dim + page_size = s.page_size + kv_cache_len = s.kv_cache_len + past_kv_len = s.past_kv_len + qo_len = s.qo_len + kv_len = s.kv_len_resolved + batch_size = s.batch_size + num_generations = s.num_generations + num_contexts = s.num_contexts + num_ctx_q_tokens = s.num_ctx_q_tokens + num_ctx_kv_tokens = s.num_ctx_kv_tokens + nnz_qo = s.nnz_qo + nnz_kv = s.nnz_kv + cross = s.cross + + q_at_layer = torch.randn(num_layers, nnz_qo, + num_heads * head_dim).half().cuda() + flashinfer_kv_cache = torch.randn(num_layers, + s.max_num_pages, + 2, + page_size, + num_kv_heads, + head_dim, + device="cuda").to(s.kvcache_dtype) + ref_kv_cache = (flashinfer_kv_cache.transpose(1, 2).contiguous().view( + num_layers, 2, batch_size, kv_cache_len, num_kv_heads, head_dim)) + vanilla_kv_cache = ref_kv_cache.transpose(1, 2).contiguous() + kv = torch.randn(num_layers, + 2, + nnz_kv, + num_kv_heads * head_dim, + device="cuda").to(dtype) + + # Test flashinfer attention + + context_lens = torch.full((num_contexts, ), qo_len).int() + qo_lens = torch.concat([context_lens, torch.ones(num_generations).int()]) + + if cross: + context_lens_kv = torch.full((num_contexts, ), kv_len).int() + seq_lens_kv = torch.concat( + [context_lens_kv, + torch.zeros(num_generations).int()]) + else: + seq_lens_kv = None + + num_cached_tokens_prefill = past_kv_len + num_cached_tokens_decode = kv_cache_len - (0 if cross else 1) + + def produce(Attention: type[AttentionBackend], kv_cache: torch.Tensor): + return produce_outputs( + Attention, + q_at_layer, + kv, + s, + kv_cache=kv_cache, + num_cached_tokens=lambda i: + (num_cached_tokens_prefill + if i < num_contexts else num_cached_tokens_decode), + seq_lens=qo_lens, + seq_lens_kv=seq_lens_kv, + num_contexts=num_contexts, + quant_config=(QuantConfig( + quant_algo=QuantAlgo.FP8, + kv_cache_quant_algo=QuantAlgo.FP8, + ) if is_fp8 else None), + ) + + flashinfer_outputs = produce(FlashInferAttention, flashinfer_kv_cache) + vanilla_outputs = produce(VanillaAttention, vanilla_kv_cache) + + # Test reference attention + + kv_lens = torch.full((batch_size, ), kv_cache_len).int().cuda() + causal_mask = (generate_causal_mask(kv_lens, qo_lens, batch_size, dtype) + if s.causal else 0) + + ref_outputs = [] + for i in range(num_layers): + q = q_at_layer[i] + ref_kv_cache[i][:, :num_contexts, past_kv_len:kv_cache_len] = kv[ + i][:, :num_ctx_kv_tokens].view(2, num_contexts, kv_len, + num_kv_heads, head_dim) + if not cross: + ref_kv_cache[i][:, num_contexts:, + -1:] = kv[i][:, num_ctx_kv_tokens:].view( + 2, num_generations, 1, num_kv_heads, head_dim) + k = ref_kv_cache[i][0] + v = ref_kv_cache[i][1] + ctx_q, gen_q = q.split([num_ctx_q_tokens, num_generations]) + gen_q = torch.nn.functional.pad(gen_q.unsqueeze(1), + (0, 0, 0, qo_len - 1), "constant", + 0).view(num_generations * qo_len, + num_heads * head_dim) + q = torch.cat([ctx_q, gen_q], dim=0) + q = q.view(batch_size, qo_len, num_heads, head_dim) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous().to(dtype) + v = v.transpose(1, 2).contiguous().to(dtype) + k = repeat_kv(k, num_kv_groups) + v = repeat_kv(v, num_kv_groups) + + attn_weights = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(head_dim) + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = torch.nn.functional.softmax(attn_weights, + dim=-1, + dtype=torch.float32).to( + q.dtype) + o = torch.matmul(attn_weights, v) + + o = o.transpose(1, 2).contiguous() + o = o.view(batch_size, qo_len, -1) + ctx_o, gen_o = o.split([num_contexts, num_generations], dim=0) + gen_o = gen_o[:, 0].view(num_generations, num_heads * head_dim) + ctx_o = ctx_o.view(num_ctx_q_tokens, num_heads * head_dim) + o = torch.cat([ctx_o, gen_o], dim=0) + assert list(o.shape) == [nnz_qo, num_heads * head_dim] + ref_outputs.append(o) + + for i in range(num_layers): + print(f"validate accuracy for layer {i}") + + allclose( + ref_outputs, + { + "flashinfer": flashinfer_outputs, + "vanilla": vanilla_outputs + }, + layer=i, + atol=fp8_atol if is_fp8 else atol, + ) + assert torch.allclose( + flashinfer_outputs[i], + vanilla_outputs[i], + atol=fp8_atol if is_fp8 else atol, + rtol=rtol, + ) + + del flashinfer_kv_cache + del ref_kv_cache + del vanilla_kv_cache + torch.cuda.empty_cache() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/aa/test_attention_mla.py b/tests/aa/test_attention_mla.py new file mode 100644 index 000000000000..36f85f55add8 --- /dev/null +++ b/tests/aa/test_attention_mla.py @@ -0,0 +1,784 @@ +import math +import random +from dataclasses import dataclass +from typing import List + +import pytest +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionInputType, MLAParams, PositionalEmbeddingParams, RopeParams) +from tensorrt_llm._torch.attention_backend.utils import get_attention_backend +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, + LlmRequestState, + SamplingConfig) +from tensorrt_llm._torch.pyexecutor.resource_manager import (KVCacheManager, + KVCacheManagerV2) +from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm.sampling_params import SamplingParams + + +# Copied from transformers.models.llama.modeling_llama.rotate_half +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + +def calculate_ref_result_ctx(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + compressed_kv: torch.Tensor, k_pe: torch.Tensor, + rope_cos_sin: torch.Tensor, num_heads: int, + num_kv_heads: int, qk_nope_head_dim: int, + qk_rope_head_dim: int, v_head_dim: int, + sequence_lengths: List[int], q_scaling: float): + """ + use standard attention to calculate the reference result by iterating over each request + q shape: (total_tokens, num_heads * qk_head_dim) + k shape: (total_tokens, num_kv_heads * qk_head_dim) + v shape: (total_tokens, num_kv_heads * v_head_dim) + compressed_kv shape: (total_tokens, kv_lora_rank) + k_pe shape: (total_tokens, qk_rope_head_dim) + rope_cos_sin shape: (max_position_embeddings, 2, qk_rope_head_dim) + """ + num_requests = len(sequence_lengths) + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + # Reshape inputs for reference calculation + q_reshaped = [] + k_reshaped = [] + v_reshaped = [] + k_pe_ref_list = [] + total_tokens = 0 + for i in range(num_requests): + q_seq = q[total_tokens:total_tokens + sequence_lengths[i]].unflatten( + -1, [num_heads, qk_head_dim]) + k_seq = k[total_tokens:total_tokens + sequence_lengths[i]].unflatten( + -1, [num_kv_heads, qk_head_dim]) + v_seq = v[total_tokens:total_tokens + sequence_lengths[i]].unflatten( + -1, [num_kv_heads, v_head_dim]) + + q_pe_seq = q_seq[..., -qk_rope_head_dim:] + k_pe_seq = k_pe[total_tokens:total_tokens + + sequence_lengths[i]].unsqueeze(-2) + cos, sin = rope_cos_sin[:sequence_lengths[i]].chunk(2, dim=-2) + q_pe_seq = q_pe_seq.unflatten(-1, [-1, 2]).transpose( + -2, -1).flatten(start_dim=-2) + k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose( + -2, -1).flatten(start_dim=-2) + q_pe_seq = ((q_pe_seq * cos) + + (rotate_half(q_pe_seq) * sin)).to(dtype=q_seq.dtype) + k_pe_seq = ((k_pe_seq * cos) + + (rotate_half(k_pe_seq) * sin)).to(dtype=k_seq.dtype) + q_seq[..., -qk_rope_head_dim:] = q_pe_seq + k_seq[..., -qk_rope_head_dim:] = k_pe_seq + k_pe_ref_list.append(k_pe_seq) + + q_reshaped.append(q_seq.transpose( + 0, 1)) # (num_heads, seq_len, qk_head_dim) + k_reshaped.append(k_seq.transpose( + 0, 1)) # (num_kv_heads, seq_len, qk_head_dim) + v_reshaped.append(v_seq.transpose( + 0, 1)) # (num_kv_heads, seq_len, v_head_dim) + + total_tokens += sequence_lengths[i] + + # Calculate reference result batch by batch + ref_results = [] + for i in range(num_requests): + q = q_reshaped[i] # (num_heads, seq_len, qk_head_dim) + k = k_reshaped[i] # (num_kv_heads, seq_len, qk_head_dim) + v = v_reshaped[i] # (num_kv_heads, seq_len, v_head_dim) + + # Handle grouped-query attention if num_heads > num_kv_heads + if num_heads > num_kv_heads: + assert num_heads % num_kv_heads == 0 + num_kv_groups = num_heads // num_kv_heads + k = repeat_kv(k.unsqueeze(0), num_kv_groups).squeeze(0) + v = repeat_kv(v.unsqueeze(0), num_kv_groups).squeeze(0) + + # Compute attention scores + attn_weights = torch.matmul(q, k.transpose( + 1, 2)) / (q_scaling * math.sqrt(qk_head_dim)) + + # For causal mask, we block future tokens (upper triangular above the diagonal) + seq_len = q.shape[1] + causal_mask = torch.triu(torch.ones(seq_len, + seq_len, + device=q.device, + dtype=torch.bool), + diagonal=1) + attn_weights = attn_weights.masked_fill(causal_mask, float('-inf')) + + # Apply softmax to get attention probabilities + attn_weights = torch.nn.functional.softmax(attn_weights, + dim=-1, + dtype=torch.float32).to( + q.dtype) + + # Apply attention weights to values + attn_output = torch.matmul(attn_weights, + v) # (num_heads, seq_len, v_head_dim) + + # Reshape back to (seq_len, num_heads*v_head_dim) + attn_output = attn_output.transpose(0, 1).contiguous().view( + sequence_lengths[i], num_heads * v_head_dim) + ref_results.append(attn_output) + + ref_result = torch.cat(ref_results) + k_pe_ref = torch.cat(k_pe_ref_list).squeeze(-2) + latent_cache = torch.cat([compressed_kv, k_pe_ref], dim=-1) + return ref_result, latent_cache + + +def calculate_ref_result_gen(fused_q: torch.Tensor, q_pe: torch.Tensor, + compressed_kv: torch.Tensor, k_pe: torch.Tensor, + latent_cache: torch.Tensor, + rope_cos_sin: torch.Tensor, num_heads: int, + kv_lora_rank: int, qk_nope_head_dim: int, + qk_rope_head_dim: int, sequence_lengths: List[int], + q_scaling: float): + """ + use standard attention to calculate the reference result by iterating over each request + fused_q shape: (num_tokens, num_heads * (kv_lora_rank + qk_rope_head_dim)) + q_pe shape: (num_tokens, num_heads, qk_rope_head_dim) + compressed_kv shape: (num_requests, kv_lora_rank) + k_pe shape: (num_requests, qk_rope_head_dim) + latent_cache shape: (total_tokens, kv_lora_rank + qk_rope_head_dim) + rope_cos_sin shape: (max_position_embeddings, 2, qk_rope_head_dim) + """ + num_requests = len(sequence_lengths) + seq_len_q = fused_q.shape[0] // num_requests + + # Reshape inputs for reference calculation + q_reshaped = [] + k_reshaped = [] + v_reshaped = [] + latent_cache_list = [] + total_tokens = 0 + for i in range(num_requests): + fused_q_seq = fused_q[i * seq_len_q:(i + 1) * seq_len_q].unflatten( + -1, [num_heads, kv_lora_rank + qk_rope_head_dim]) + q_pe_seq = q_pe[i * seq_len_q:(i + 1) * seq_len_q] + compressed_kv_seq = compressed_kv[i * seq_len_q:(i + 1) * + seq_len_q].unsqueeze(-2) + k_pe_seq = k_pe[i * seq_len_q:(i + 1) * seq_len_q].unsqueeze(-2) + latent_cache_seq = latent_cache[total_tokens:total_tokens + + sequence_lengths[i]].unsqueeze(-2) + + cos, sin = rope_cos_sin[sequence_lengths[i]:sequence_lengths[i] + + seq_len_q].chunk(2, dim=-2) + q_pe_seq = q_pe_seq.unflatten(-1, [-1, 2]).transpose( + -2, -1).flatten(start_dim=-2) + k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose( + -2, -1).flatten(start_dim=-2) + q_pe_seq = ((q_pe_seq * cos) + + (rotate_half(q_pe_seq) * sin)).to(dtype=q_pe_seq.dtype) + k_pe_seq = ((k_pe_seq * cos) + + (rotate_half(k_pe_seq) * sin)).to(dtype=k_pe_seq.dtype) + fused_q_seq[..., -qk_rope_head_dim:] = q_pe_seq + latent_cache_seq = torch.cat([ + latent_cache_seq, + torch.cat([compressed_kv_seq, k_pe_seq], dim=-1) + ], + dim=0) + latent_cache_list.append(latent_cache_seq) + + q_reshaped.append(fused_q_seq.transpose( + 0, 1)) # (num_heads, seq_len_q, kv_lora_rank + qk_rope_head_dim) + k_reshaped.append(latent_cache_seq.transpose( + 0, 1)) # (1, seq_len_kv, kv_lora_rank + qk_rope_head_dim) + v_reshaped.append(latent_cache_seq[..., :kv_lora_rank].transpose( + 0, 1)) # (1, seq_len_kv, kv_lora_rank) + + total_tokens += sequence_lengths[i] + + # Calculate reference result batch by batch + ref_results = [] + for i in range(num_requests): + q = q_reshaped[ + i] # (num_heads, seq_len_q, kv_lora_rank + qk_rope_head_dim) + k = k_reshaped[i] # (1, seq_len_kv, kv_lora_rank + qk_rope_head_dim) + v = v_reshaped[i] # (1, seq_len_kv, kv_lora_rank) + + # Handle grouped-query attention + k = repeat_kv(k.unsqueeze(0), num_heads).squeeze(0) + v = repeat_kv(v.unsqueeze(0), num_heads).squeeze(0) + + # Compute attention scores + attn_weights = torch.matmul(q, k.transpose(1, 2)) / ( + q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) + + # Use MTP mask by default if seqlen_q > 1. + seq_len_q = q.shape[1] + seq_len_kv = k.shape[1] + mask = torch.zeros(seq_len_q, + seq_len_kv, + device=q.device, + dtype=torch.bool) + for qi in range(seq_len_q): + for ki in range(seq_len_kv - seq_len_q + 1 + qi, seq_len_kv): + mask[qi, ki] = 1 + attn_weights = attn_weights.masked_fill(mask, float('-inf')) + # Apply softmax to get attention probabilities + attn_weights = torch.nn.functional.softmax(attn_weights, + dim=-1, + dtype=torch.float32).to( + q.dtype) + + # Apply attention weights to values + attn_output = torch.matmul(attn_weights, + v) # (num_heads, 1, kv_lora_rank) + + # Reshape back to (seq_len_q, num_heads*kv_lora_rank) + attn_output = attn_output.transpose(0, 1).contiguous().view( + -1, num_heads * kv_lora_rank) + ref_results.append(attn_output) + + ref_result = torch.cat(ref_results) + latent_cache = torch.cat(latent_cache_list).squeeze(-2) + return ref_result, latent_cache + + +@dataclass(kw_only=True, frozen=True) +class Scenario: + dtype: torch.dtype = torch.bfloat16 + kv_cache_dtype: torch.dtype = torch.bfloat16 + num_layers: int = 1 + num_heads: int = 128 + num_kv_heads: int = 128 + q_lora_rank: int = 1536 + kv_lora_rank: int = 512 + qk_nope_head_dim: int = 128 + qk_rope_head_dim: int = 64 + v_head_dim: int = 128 + hidden_size: int = 7168 + max_position_embeddings: int = 163840 + rope_theta: float = 10000.0 + rope_beta_fast: int = 32 + rope_beta_slow: int = 1 + rope_factor: float = 40.0 + rope_mscale: float = 1.0 + rope_mscale_all_dim: float = 1.0 + rope_original_max_position_embeddings: int = 4096 + rope_type: str = "yarn" + model_type: str = "deepseek_v3" + kv_cache_tokens_per_block: int = 64 + + +@dataclass(kw_only=True, frozen=True) +class RopeConfig: + hidden_size: int = 7168 + num_attention_heads: int = 128 + rope_scaling: dict = { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40.0, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + max_position_embeddings: int = 163840 + rope_theta: float = 10000.0 + qk_rope_head_dim: int = 64 + model_type: str = "deepseek_v3" + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, + None, :, :].expand(batch, num_key_value_heads, + n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, + head_dim) + + +# Set seed for reproducibility. +random.seed(0) +min_context_sequence_length = 1 +max_context_sequence_length = 1000 +min_num_contexts = 1 +max_num_contexts = 10 +random_context_sequence_lengths = [ + random.randint(min_context_sequence_length, max_context_sequence_length) + for _ in range(random.randint(min_num_contexts, max_num_contexts)) +] + +# Define test data +context_sequence_lengths = [ + [10, 12, 5], + [100, 300, 20, 10], + [253, 253, 253, 253], + [100, 1110, 1000, 1000], + random_context_sequence_lengths, +] +# Use MTP by default if seqlen_q > 1. +generation_seq_len_q = [1, 4] +num_generation_steps = [10] + +kv_cache_dtype_list = [torch.bfloat16] +if torch.cuda.get_device_capability() in [(8, 9), (9, 0), (10, 0), (12, 0)]: + kv_cache_dtype_list.append(torch.float8_e4m3fn) +scenarios = [ + Scenario(kv_cache_dtype=kv_cache_dtype, num_layers=num_layers) + for kv_cache_dtype in kv_cache_dtype_list for num_layers in [1, 2] +] + +accuracy_dict = { + torch.bfloat16: (3e-2, 3e-3), + torch.float8_e4m3fn: (4e-1, 4e-2), +} + + +# Convert parameterized tests to pytest parametrize +@pytest.mark.parametrize("scenario", scenarios, ids=lambda x: f"scenario: {x}") +@pytest.mark.parametrize("context_sequence_lengths", + context_sequence_lengths, + ids=lambda x: f"context_sequence_lengths: {x}") +@pytest.mark.parametrize("generation_seq_len_q", + generation_seq_len_q, + ids=lambda x: f"generation_seq_len_q: {x}") +@pytest.mark.parametrize("num_generation_steps", + num_generation_steps, + ids=lambda x: f"num_generation_steps: {x}") +def test_attention_mla(scenario: Scenario, + context_sequence_lengths: List[int], + generation_seq_len_q: int, + num_generation_steps: List[int], + v2_kv_cache: bool = False): + """Test MLA computation for both context and generation phases""" + num_heads = scenario.num_heads + num_kv_heads = scenario.num_kv_heads + q_lora_rank = scenario.q_lora_rank + kv_lora_rank = scenario.kv_lora_rank + qk_nope_head_dim = scenario.qk_nope_head_dim + qk_rope_head_dim = scenario.qk_rope_head_dim + v_head_dim = scenario.v_head_dim + rope_config = RopeConfig( + hidden_size=scenario.hidden_size, + num_attention_heads=scenario.num_heads, + rope_scaling={ + "beta_fast": scenario.rope_beta_fast, + "beta_slow": scenario.rope_beta_slow, + "factor": scenario.rope_factor, + "mscale": scenario.rope_mscale, + "mscale_all_dim": scenario.rope_mscale_all_dim, + "original_max_position_embeddings": + scenario.rope_original_max_position_embeddings, + "type": scenario.rope_type, + }, + max_position_embeddings=scenario.max_position_embeddings, + rope_theta=scenario.rope_theta, + qk_rope_head_dim=scenario.qk_rope_head_dim, + model_type=scenario.model_type, + ) + kv_cache_tokens_per_block = scenario.kv_cache_tokens_per_block + num_layers = scenario.num_layers + device = torch.device('cuda') + dtype = scenario.dtype + kv_cache_dtype = scenario.kv_cache_dtype + + print( + f"--------------------------------Test for scenario: {scenario} start--------------------------------" + ) + + _run_test_for_backend("TRTLLM", num_heads, num_kv_heads, num_layers, + q_lora_rank, kv_lora_rank, qk_nope_head_dim, + qk_rope_head_dim, v_head_dim, rope_config, + kv_cache_tokens_per_block, device, dtype, + kv_cache_dtype, context_sequence_lengths, + generation_seq_len_q, num_generation_steps, + v2_kv_cache) + + +def _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers, + q_lora_rank, kv_lora_rank, qk_nope_head_dim, + qk_rope_head_dim, v_head_dim, rope_config, + kv_cache_tokens_per_block, device, dtype, + kv_cache_dtype, context_sequence_lengths, + generation_seq_len_q, num_generation_steps, + v2_kv_cache): + AttentionCls = get_attention_backend(backend_name) + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + # Set seed for reproducibility. + torch.manual_seed(123) + + # Create inputs + inputs_per_layer = [] + for layer_idx in range(num_layers): + ctx_compressed_kv = torch.cat([ + torch.empty( + [ctx_len, kv_lora_rank], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for ctx_len in context_sequence_lengths + ]) + ctx_k_pe = torch.cat([ + torch.empty( + [ctx_len, qk_rope_head_dim], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for ctx_len in context_sequence_lengths + ]) + ctx_q = torch.cat([ + torch.empty( + [ctx_len, num_heads * qk_head_dim], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for ctx_len in context_sequence_lengths + ]) + ctx_kv = torch.cat([ + torch.empty( + [ctx_len, num_kv_heads * (qk_nope_head_dim + v_head_dim)], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for ctx_len in context_sequence_lengths + ]) + # ctx_v.stride(0) == num_kv_heads * (qk_nope_head_dim + v_head_dim) + ctx_k_nope, ctx_v = ctx_kv.split( + [num_kv_heads * qk_nope_head_dim, num_kv_heads * v_head_dim], + dim=-1) + ctx_k_nope = ctx_k_nope.view(-1, num_kv_heads, qk_nope_head_dim) + ctx_k = torch.cat([ + ctx_k_nope, + ctx_k_pe.view(-1, 1, qk_rope_head_dim).expand(-1, num_kv_heads, -1) + ], + dim=-1) + ctx_k = ctx_k.view(-1, num_kv_heads * qk_head_dim) + + gen_compressed_kv_list = [ + torch.cat([ + torch.empty( + [generation_seq_len_q, kv_lora_rank], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for _ in context_sequence_lengths + ]) for _ in range(num_generation_steps) + ] + gen_k_pe_list = [ + torch.cat([ + torch.empty( + [generation_seq_len_q, qk_rope_head_dim], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for _ in context_sequence_lengths + ]) for _ in range(num_generation_steps) + ] + gen_fused_q_list = [ + torch.cat([ + torch.empty( + [ + generation_seq_len_q, num_heads * + (kv_lora_rank + qk_rope_head_dim) + ], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for _ in context_sequence_lengths + ]) for _ in range(num_generation_steps) + ] + gen_q_pe_list = [ + torch.cat([ + torch.empty( + [generation_seq_len_q, num_heads, qk_rope_head_dim], + dtype=dtype, + device=device, + ).uniform_(-1, 1) for _ in context_sequence_lengths + ]) for _ in range(num_generation_steps) + ] + + inputs = { + "ctx_compressed_kv": ctx_compressed_kv, + "ctx_k_pe": ctx_k_pe, + "ctx_q": ctx_q, + "ctx_k": ctx_k, + "ctx_v": ctx_v, + "gen_compressed_kv_list": gen_compressed_kv_list, + "gen_k_pe_list": gen_k_pe_list, + "gen_fused_q_list": gen_fused_q_list, + "gen_q_pe_list": gen_q_pe_list, + } + inputs_per_layer.append(inputs) + print(f"context sequence lengths: {context_sequence_lengths}") + for key, val in inputs.items(): + if key.endswith("_list"): + print(f"{key}: [{val[0].shape}] * {len(val)}") + else: + print(f"{key}: {val.shape}") + + rope_cos_sin = torch.tensor( + RopeEmbeddingUtils.create_sinusoidal_positions_yarn( + rope_config.max_position_embeddings, + rope_config.qk_rope_head_dim, + rope_config.rope_theta, + rope_config.rope_scaling['factor'], + rope_config.rope_scaling['original_max_position_embeddings'], + rope_config.rope_scaling['beta_fast'], + rope_config.rope_scaling['beta_slow'], + rope_config.rope_scaling['mscale'], + rope_config.rope_scaling['mscale_all_dim'], + )[1], + dtype=torch.float32, + device=device, + ).reshape(rope_config.max_position_embeddings, -1, 2).transpose(-2, -1) + + # Setup attention module and metadata + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, + rope=RopeParams.from_config(rope_config), + is_neox=False, + ) + mla_params = MLAParams( + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + v_head_dim=v_head_dim, + predicted_tokens_per_seq=1, + ) + + def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + mscale_all_dim = pos_embd_params.rope.mscale_all_dim + scaling_factor = pos_embd_params.rope.scale + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + q_scaling = 1.0 / (mscale * mscale) + + quant_config = None + if kv_cache_dtype == torch.float8_e4m3fn: + quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8.value) + ctx_layers = [ + AttentionCls( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=qk_head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + ) for layer_idx in range(num_layers) + ] + gen_layers = [ + AttentionCls( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=kv_lora_rank + qk_rope_head_dim, + num_kv_heads=1, + quant_config=quant_config, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + ) for layer_idx in range(num_layers) + ] + + # NOTE: set up metadata, refer to tensorrt_llm/_torch/pyexecutor/model_engine.py + # all layers share the same metadata + mapping = Mapping(world_size=1, tp_size=1, rank=0) + max_tokens = ( + max_context_sequence_length + + (num_generation_steps + 1) * generation_seq_len_q + + kv_cache_tokens_per_block - 1 + ) // kv_cache_tokens_per_block * kv_cache_tokens_per_block * max_num_contexts + kv_cache_cls = KVCacheManagerV2 if v2_kv_cache else KVCacheManager + kv_cache_manager = kv_cache_cls( + KvCacheConfig( + max_tokens=max_tokens, + enable_block_reuse=False, + ), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=kv_lora_rank + qk_rope_head_dim, + tokens_per_block=kv_cache_tokens_per_block, + max_seq_len=max(context_sequence_lengths) + + (num_generation_steps + 1) * generation_seq_len_q, + max_batch_size=len(context_sequence_lengths), + mapping=mapping, + dtype=str_dtype_to_binding(torch_dtype_to_str(kv_cache_dtype)), + ) + request_list = [] + for req_id, ctx_len in enumerate(context_sequence_lengths): + req = LlmRequest( + request_id=req_id, + max_new_tokens=num_generation_steps + 1, + input_tokens=[1] * ctx_len, + sampling_config=SamplingConfig( + SamplingParams()._get_sampling_config()), + is_streaming=False, + ) + req.paged_kv_block_ids = [] + beam_width = 1 + if v2_kv_cache: + kv_cache = kv_cache_manager.impl.create_kv_cache(None, None) + success = kv_cache.resume(torch.cuda.current_stream().cuda_stream) + assert success + kv_cache.capacity = ctx_len + kv_cache_manager.kv_cache_map[req_id] = (kv_cache, ctx_len) + else: + kv_cache_manager.impl.add_sequence(req_id, ctx_len, beam_width, req) + request_list.append(req) + attn_metadata = AttentionCls.Metadata( + seq_lens=torch.tensor(context_sequence_lengths, dtype=torch.int), + request_ids=list(range(len(context_sequence_lengths))), + max_num_requests=len(context_sequence_lengths), + num_contexts=len(context_sequence_lengths), + prompt_lens=context_sequence_lengths, + max_num_tokens=max(context_sequence_lengths), + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0 for _ in context_sequence_lengths], + ), + mapping=mapping, + ) + attn_metadata.prepare() + + # run forward for each step and each layer + latent_cache_ref_all_list = [None for _ in range(num_layers)] + for step in range(num_generation_steps + 1): + if step > 0: + for req_id in range(len(context_sequence_lengths)): + for _ in range(generation_seq_len_q): + kv_cache_manager.impl.add_token(req_id) + attn_metadata = AttentionCls.Metadata( + seq_lens=torch.tensor([generation_seq_len_q] * + len(context_sequence_lengths), + dtype=torch.int), + request_ids=list(range(len(context_sequence_lengths))), + max_num_requests=len(context_sequence_lengths), + num_contexts=0, + prompt_lens=context_sequence_lengths, + max_num_tokens=max(context_sequence_lengths), + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[ + ctx_len + (step - 1) * generation_seq_len_q + for ctx_len in context_sequence_lengths + ], + ), + mapping=mapping, + enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), + ) + attn_metadata.prepare() + for layer_idx in range(num_layers): + print( + f"--------------------------------step {step} layer {layer_idx} start--------------------------------" + ) + if step == 0: + q = inputs_per_layer[layer_idx]["ctx_q"] + k = inputs_per_layer[layer_idx]["ctx_k"] + v = inputs_per_layer[layer_idx]["ctx_v"] + compressed_kv = inputs_per_layer[layer_idx]["ctx_compressed_kv"] + k_pe = inputs_per_layer[layer_idx]["ctx_k_pe"] + latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + # q/k will be modified in the forward pass, so we need to clone them + # we should not clone v because we need to keep the stride of v + result = ctx_layers[layer_idx].forward( + q.clone(), + k.clone(), + v, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=latent_cache, + ) + ref_result, latent_cache_ref = calculate_ref_result_ctx( + q, + k, + v, + compressed_kv, + k_pe, + rope_cos_sin, + num_heads, + num_kv_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + context_sequence_lengths, + q_scaling, + ) + latent_cache_ref_all_list[layer_idx] = latent_cache_ref + for req in request_list: + req.state = LlmRequestState.GENERATION_IN_PROGRESS + else: + fused_q = inputs_per_layer[layer_idx]["gen_fused_q_list"][step - + 1] + q_pe = inputs_per_layer[layer_idx]["gen_q_pe_list"][step - 1] + compressed_kv = inputs_per_layer[layer_idx][ + "gen_compressed_kv_list"][step - 1] + k_pe = inputs_per_layer[layer_idx]["gen_k_pe_list"][step - 1] + latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + result = gen_layers[layer_idx].forward( + fused_q, + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=latent_cache, + q_pe=q_pe, + ) + ref_result, latent_cache_ref = calculate_ref_result_gen( + fused_q, + q_pe, + compressed_kv, + k_pe, + latent_cache_ref_all_list[layer_idx], + rope_cos_sin, + num_heads, + kv_lora_rank, + qk_nope_head_dim, + qk_rope_head_dim, + [ + ctx_len + (step - 1) * generation_seq_len_q + for ctx_len in context_sequence_lengths + ], + q_scaling, + ) + latent_cache_ref_all_list[layer_idx] = latent_cache_ref + # Compare results + print( + f"{backend_name} output mean: {result.abs().mean().item()}, max: {result.abs().max().item()}" + ) + print( + f"Reference output mean: {ref_result.abs().mean().item()}, max: {ref_result.abs().max().item()}" + ) + print( + f"Difference mean: {(result - ref_result).abs().mean().item()}, max: {(result - ref_result).abs().max().item()}" + ) + + # Assert results are close + atol, rtol = accuracy_dict[kv_cache_dtype] + assert torch.allclose(result, ref_result, atol=atol, rtol=rtol), \ + f"Results for MLA in {backend_name} backend don't match reference implementation at layer {layer_idx} in step {step}" + + print( + f"Test for MLA in {backend_name} backend passed at layer {layer_idx} in step {step}" + ) + print( + f"--------------------------------step {step} layer {layer_idx} end--------------------------------" + ) + + print(f"Test for MLA in {backend_name} backend passed") + kv_cache_manager.shutdown() + + +if __name__ == "__main__": + test_attention_mla(Scenario(num_layers=1), [1024, 12, 5], 1, 10, True) diff --git a/tests/aa/test_naive.py b/tests/aa/test_naive.py new file mode 100755 index 000000000000..2d5c16ec5af0 --- /dev/null +++ b/tests/aa/test_naive.py @@ -0,0 +1,653 @@ +import functools +import gc +import itertools +import os +import random +import time +import unittest +from contextlib import contextmanager +from importlib.util import find_spec +from random import randbytes +from statistics import median +from typing import Iterator, NamedTuple, cast + +if find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + DiskCacheTierConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheManager, + KVCacheManagerConfig, + LayerId, + TokenId, + TokenIdExt, + _KVCache, + ) + from kv_cache_manager_v2._block_radix_tree import traverse_post_order + from kv_cache_manager_v2._common import CudaStream, PageStatus, SlidingWindowSize + from kv_cache_manager_v2._exceptions import OutOfPagesError + from kv_cache_manager_v2._utils import ( + TemporaryCudaStream, + init_cuda_once, + remove_if, + round_up, + typed_range, + unwrap_weakref, + ) +else: + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + DiskCacheTierConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheManager, + KVCacheManagerConfig, + LayerId, + TokenId, + TokenIdExt, + _KVCache, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import traverse_post_order + from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( + CudaStream, + PageStatus, + SlidingWindowSize, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError + from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + TemporaryCudaStream, + init_cuda_once, + remove_if, + round_up, + typed_range, + unwrap_weakref, + ) + +from dynamic_path_manager import DynamicPathManager +from parameterized import parameterized + +with DynamicPathManager(os.path.dirname(os.path.abspath(__file__))): + from fake_engine import FakeEngine, Role, Step + from kernels import enable_kernel_delay + +from fake_engine import FakeEngine, Role, Step +from kernels import enable_kernel_delay + +seed = int.from_bytes(os.urandom(8), "little") +print(f"seed: {seed}") +random.seed(seed) +DBG_PRINT = int(os.environ.get("DBG_PRINT", "0")) != 0 +PRINT_TIME = int(os.environ.get("PRINT_TIME", "0")) != 0 + + +@contextmanager +def ref_cycle_check_context(): + """Context manager for reference cycle check.""" + import gc + + gc.collect() + gc.garbage.clear() + gc.set_debug(gc.DEBUG_SAVEALL | gc.DEBUG_COLLECTABLE) + + def on_gc_event(phase, info): + # phase is "start" or "stop" + # info contains keys like: "generation", "collected", "uncollectable", "duration" + if phase == "stop": + collected = info.get("collected", 0) + uncollectable = info.get("uncollectable", 0) + if collected != 0 or uncollectable != 0: + import pdb + + pdb.set_trace() + assert collected == 0 and uncollectable == 0 + + gc.callbacks.append(on_gc_event) + try: + yield + finally: + gc.collect() + gc.callbacks.pop() + gc.set_debug(0) + + +def assert_no_ref_cycle(func): + """Decorator to wrap test methods with GC debugging context.""" + + @functools.wraps(func) + def wrapper(self, *args, **kwargs): + with ref_cycle_check_context(): + result = func(self, *args, **kwargs) + return result + + return wrapper + + +class TestKVCacheManagerV2(unittest.TestCase): + engine: FakeEngine + cfg: KVCacheManagerConfig + manager: KVCacheManager + _token_id_gen: Iterator[int] + + def setUp(self) -> None: + init_cuda_once() + self._token_id_gen = itertools.count() + gc.collect() + gc.disable() + + def tearDown(self) -> None: + gc.enable() + if hasattr(self, "manager"): + del self.manager + + def next_token(self) -> TokenIdExt: + token_id = next(self._token_id_gen) + if token_id % 100 == 99: + return randbytes(32) + else: + return TokenId(token_id) + + def prepare( + self, + gpu_quota: int, + host_quota: int, + disk_quota: int, + num_layers: int, + window_size: SlidingWindowSize, + sink_tokens: int, + tokens_per_block: int = 32, + kv_buf_size: int = 8192, + ): + self._init_cfg( + tokens_per_block, + gpu_quota, + host_quota, + disk_quota, + num_layers, + window_size, + sink_tokens, + kv_buf_size, + ) + self.engine = FakeEngine(self.cfg) + self.manager = KVCacheManager(self.cfg) + + def _init_cfg( + self, + tokens_per_block: int, + gpu_quota: int, + host_quota: int, + disk_quota: int, + num_layers: int, + window_size: SlidingWindowSize, + sink_tokens: int, + kv_buf_size: int = 8192, + block_quant_buf_size: int | None = None, + ): + layer_buffers = [ + BufferConfig(role=Role.KEY, size=kv_buf_size), + BufferConfig(role=Role.VALUE, size=kv_buf_size), + ] + if block_quant_buf_size is not None: + layer_buffers.extend( + [ + BufferConfig(role=Role.KEY_BLOCK_QUANT, size=block_quant_buf_size), + BufferConfig(role=Role.VALUE_BLOCK_QUANT, size=block_quant_buf_size), + ] + ) + cache_tiers = [ + GpuCacheTierConfig(quota=gpu_quota), + HostCacheTierConfig(quota=host_quota), + DiskCacheTierConfig(quota=disk_quota, path="/workspace/"), + ] + self.cfg = KVCacheManagerConfig( + tokens_per_block=tokens_per_block, + vocab_size=4096, + cache_tiers=[t for t in cache_tiers if t.quota > 0], + layers=[ + AttentionLayerConfig( + layer_id=layer_id, + buffers=layer_buffers, + sliding_window_size=window_size if layer_id % 2 == 0 else None, + num_sink_tokens=sink_tokens if layer_id % 2 == 0 else None, + ) + for layer_id in typed_range(LayerId(num_layers)) + ], + ) + + +class TestNoBatching(TestKVCacheManagerV2): + class Request(NamedTuple): + id: int + kv_cache: _KVCache + prompt: list[TokenIdExt] + decode_len: int + + def new_request( + self, req_id: int, lora_task_id: int | None, prompt_len: int, decode_len: int + ) -> Request: + prompt = [self.next_token() for _ in range(prompt_len)] + return self.Request( + req_id, self.manager.create_kv_cache(lora_task_id, prompt), prompt, decode_len + ) + + def run_request(self, req: Request, interval: int, refcheck: bool) -> float: + req_id, kv_cache, prompt, decode_len = req + assert kv_cache.status == _KVCache.Status.ACTIVE + stream = kv_cache.cuda_stream + tic = time.perf_counter() + # prefill + num_reused = kv_cache.num_committed_tokens + # workaround a mypyc bug: exception in property setter is not propagated + # kv_cache.capacity = round_up(len(prompt), interval) + if not kv_cache.resize(round_up(len(prompt), interval)): + raise OutOfPagesError("Not enough pages in GPU memory") + capacity = kv_cache.capacity + history = prompt[:num_reused] + input = prompt[num_reused:] + if refcheck: + self.engine.execute([Step(kv_cache, input, history)], stream) + if input: + kv_cache.commit(input) + history.extend(input) + # decode + for _ in range(decode_len): + required_capacity = len(history) + 1 + if required_capacity > capacity: + kv_cache.commit(history[kv_cache.history_length :]) + # workaround a mypyc bug: exception in property setter is not propagated + # kv_cache.capacity = round_up(required_capacity, interval) + if not kv_cache.resize(round_up(required_capacity, interval)): + raise OutOfPagesError("Not enough pages in GPU memory") + capacity = kv_cache.capacity + input_token = self.next_token() + if refcheck: + self.engine.execute([Step(kv_cache, [input_token], history)], stream) + history.append(input_token) + kv_cache.commit(history[kv_cache.history_length :]) + # last check + if refcheck: + self.engine.execute([Step(kv_cache, [], history)], stream) + toc = time.perf_counter() + time_taken = toc - tic + # print(f"Time taken: {time_taken} seconds") + return time_taken + + def run_naive(self, seq_len: int, interval: int = 1, refcheck: bool = True) -> float: + prompt_len = 1 + decode_len = seq_len - prompt_len + + req_id = 0 + lora_task_id = None + req0 = self.new_request(req_id, lora_task_id, prompt_len, decode_len) + with TemporaryCudaStream([]) as s: + stream = s.handle + kv_cache = req0.kv_cache + success = kv_cache.resume(stream) + assert success + time_taken = self.run_request(req0, interval, refcheck) + + s.take_finish_event().synchronize() + kv_cache.close() + self.manager.clear_reusable_blocks() + return time_taken + + def test_shrink_capacity(self) -> None: + self.prepare(32 << 20, 32 << 20, 1 << 30, 36, 128, 1, kv_buf_size=32768) + seq_len = 32 * 10 + req0 = self.new_request(0, None, 32, seq_len - 32) + with TemporaryCudaStream([]) as s: + stream = s.handle + kv_cache = req0.kv_cache + success = kv_cache.resume(stream) + assert success + success = kv_cache.resize(seq_len) + assert success + for capacity in range(seq_len, len(req0.prompt), -1): + success = kv_cache.resize(capacity) + assert success + s.take_finish_event() + kv_cache.close() + + # @assert_no_ref_cycle + def test_sol_mem_utilization(self) -> None: + self.prepare(32 << 20, 32 << 20, 1 << 30, 36, 128, 1, kv_buf_size=32768) + # if we have n blocks, we need 8192*2*18*(1+5+n) bytes of memory. For the (1+5+n), 1 is for sink + # blocks, 5 is for SWA (window=128), n is for full attention. + max_seq_len = 32 * 22 # 23 blocks will require more than 32MB memory + seq_len = max_seq_len + + # create a request and suspend it. It shall not consume any GPU memory after suspend. + req0 = self.new_request(0, None, 256, seq_len - 256) + with TemporaryCudaStream([]) as s: + stream = s.handle + success = req0.kv_cache.resume(stream) + assert success + self.run_request(req0, 32, False) + s.take_finish_event() + req0.kv_cache.suspend() + + # run another request that will take all the GPU memory + req1 = self.new_request(0, None, 256, seq_len - 256) + with TemporaryCudaStream([]) as s: + stream = s.handle + success = req1.kv_cache.resume(stream) + assert success + self.run_request(req1, 1, True) + s.take_finish_event() + + req1.kv_cache.close() + req0.kv_cache.close() + + # run another longer request and expect OutOfPagesError + # This also tests eviction to disk. + self.assertRaises(OutOfPagesError, lambda: self.run_naive(seq_len + 1, 1, False)) + + @parameterized.expand([(1,), (2,), (4,)]) + # @assert_no_ref_cycle + def test_cache_reuse(self, num_reusable_requests: int) -> None: + self.prepare(32 << 20, 32 << 20, 1 << 30, 36, 128, 1, kv_buf_size=32768) + # if we have n blocks, we need 8192*2*18*(1+5+n) bytes of memory. For the (1+5+n), 1 is for sink + # blocks, 5 is for SWA (window=128), n is for full attention. + max_seq_len = 32 * 22 # 23 blocks will require more than 32MB memory + seq_len = max_seq_len + + req_id_gen = itertools.count() + reusable_requests = [] + with TemporaryCudaStream([]) as s: + stream = s.handle + for _ in range(num_reusable_requests): + req = self.new_request(next(req_id_gen), None, 256, seq_len - 256) + reusable_requests.append(req) + success = req.kv_cache.resume(stream) + assert success + self.run_request(req, 32, True) + req.kv_cache.close() + s.take_finish_event() + + for root_block in self.manager._radix_tree.next.values(): + for block0 in root_block.next.values(): + for block in traverse_post_order(block0): + for page in block.storage: + if page is not None: + assert unwrap_weakref(page).status == PageStatus.DROPPABLE + + req0 = reusable_requests[0] + prompt1 = req0.kv_cache._committed_tokens[: (seq_len // 2 - 7)] + # request id must be same as req0 because we wrote it into the kv cache. + req1 = self.Request( + next(req_id_gen), + self.manager.create_kv_cache(None, prompt1), + prompt1, + seq_len - len(prompt1), + ) + assert req1.kv_cache.num_committed_tokens == len(prompt1) + with TemporaryCudaStream([]) as s: + stream = s.handle + success = req1.kv_cache.resume(stream) + assert success + self.run_request(req1, 32, True) + s.take_finish_event() + req1.kv_cache.close() + + self.manager.clear_reusable_blocks() + + # @assert_no_ref_cycle + def test_naive(self) -> None: + self.prepare(256 << 20, 256 << 20, 1 << 30, 36, 128, 48) + self.run_naive(512, 1, True) + + @parameterized.expand([(2**i, False) for i in range(12)]) + # @parameterized.expand([(32, True)]) + # @assert_no_ref_cycle + def test_naive_perf(self, interval, profile: bool) -> None: + if not PRINT_TIME: + self.skipTest("Skipping perf test") + self.prepare(256 << 20, 256 << 20, 1 << 30, 36, 128, 48) + seq_len = 10240 + self.run_naive(seq_len, interval, False) # warm up for numba jit + profiler = None + if profile: + import cProfile + + profiler = cProfile.Profile() + profiler.enable() + time_taken = [ + self.run_naive(seq_len, interval, False) for _ in range(11 if profiler is None else 1) + ] + median_time_taken = median(time_taken) + if PRINT_TIME: + print( + f"Throughput: {round(seq_len / median_time_taken)} tokens/sec for interval {interval}" + ) + if profiler is not None: + profiler.disable() + profiler.print_stats(sort="cumtime") + profiler.dump_stats("profiler.prof") + + +class TestBatching(TestKVCacheManagerV2): + num_requests: int + avg_length: int + past_sequences: list[list[TokenIdExt]] + seq_len_dict: dict[_KVCache, int] + batch: list[Step] + suspended: list[Step] + num_created: int + num_finished: int + req_id_gen: Iterator[int] + acc_num_prompt_tokens: int + acc_num_decode_tokens: int + interval: int + enable_reuse: bool + + def setUp(self) -> None: + super().setUp() + self.past_sequences = list[list[TokenIdExt]]() + self.seq_len_dict = dict[_KVCache, int]() + self.batch = list[Step]() + self.suspended = list[Step]() + self.num_finished = 0 + self.num_created = 0 + self.req_id_gen = itertools.count() + self.acc_num_prompt_tokens = 0 + self.acc_num_decode_tokens = 0 + self.enable_reuse = False + + def gen_request(self) -> Step: + if self.num_created >= self.num_requests: + raise ValueError("Too many requests created") + + token_id_gen = cast(Iterator[TokenId], self._token_id_gen) + + def gen_length() -> int: + return random.randint(int(self.avg_length * 0.6), int(self.avg_length * 1.4)) + + if self.enable_reuse: + if len(self.past_sequences) >= 32 and random.random() < 0.2: + # continued multi-round dialog + prompt = random.choice(self.past_sequences) + [ + next(token_id_gen) for _ in range(gen_length()) + ] + else: + # new dialog + if len(self.past_sequences) < 32 or random.random() < 0.5: + # completely new prompt + prompt = [next(token_id_gen) for _ in range(gen_length())] + else: + # with reused tokens + reused = random.choice(self.past_sequences) + prompt = reused[: random.randint(0, min(gen_length(), len(reused)))] + [ + next(token_id_gen) for _ in range(gen_length()) + ] + else: + prompt = [next(token_id_gen) for _ in range(gen_length())] + decode_len = gen_length() + lora_task_id = None + kv_cache = self.manager.create_kv_cache( + lora_task_id, prompt[:-1] if self.enable_reuse else None, id=next(self.req_id_gen) + ) + DBG_PRINT and print( # type: ignore[arg-type] + f"created {kv_cache.id} with {kv_cache.num_committed_tokens} tokens reused" + ) + history = prompt[: kv_cache.num_committed_tokens] + input = prompt[kv_cache.num_committed_tokens :] + seq_len = len(prompt) + decode_len + self.seq_len_dict[kv_cache] = seq_len + self.num_created += 1 + assert input + self.acc_num_prompt_tokens += len(prompt) + self.acc_num_decode_tokens += decode_len + return Step(kv_cache, input, history) + + def update_batch(self, stream: CudaStream) -> None: + for s in self.batch: + assert s.input + if self.enable_reuse: + s.kv_cache.commit(s.input) + else: + s.kv_cache.history_length += len(s.input) + s.history.extend(s.input) + s.input.clear() + # remove finished requests first + removed = remove_if( + self.batch, + lambda step: len(step.history) >= self.seq_len_dict[step.kv_cache], + ) + for kv_cache, _, _ in removed: + seq_len = self.seq_len_dict[kv_cache] + if seq_len < self.avg_length * 3: + self.past_sequences.append(kv_cache._committed_tokens[:seq_len]) + kv_cache.close() + self.seq_len_dict.pop(kv_cache) + self.num_finished += 1 + # fill input for remaining requests and increase capacity for them + token_id_gen = cast(Iterator[TokenId], self._token_id_gen) + for s in self.batch: + assert not s.input + length = min(self.interval, self.seq_len_dict[s.kv_cache] - len(s.history)) + s.input.extend(next(token_id_gen) for _ in range(length)) + for i in itertools.count(): + if i >= len(self.batch): + break + s = self.batch[i] + while i < len(self.batch) and not s.kv_cache.resize( + len(s.history) + len(s.input), None + ): + last = self.batch.pop() + DBG_PRINT and print(f"suspending {last.kv_cache.id}") # type: ignore[arg-type] + last.kv_cache.suspend() + self.suspended.append(last) + + # try to add new requests + suspended = self.suspended + while suspended or self.num_created < self.num_requests: + if not suspended: + assert self.num_created < self.num_requests + suspended.append(self.gen_request()) + if suspended: + step = suspended[-1] + kv_cache = step.kv_cache + ok = kv_cache.resume(stream) + if ( + ok + and not self.enable_reuse + and kv_cache._commit_state == _KVCache.CommitState.ALLOWED + ): + kv_cache.stop_committing() + ok = ok and kv_cache.resize(len(step.history) + len(step.input), None) + if ok: + DBG_PRINT and print(f"activating {step.kv_cache.id}") # type: ignore[arg-type] + self.batch.append(suspended.pop()) + else: + if kv_cache.status == _KVCache.Status.ACTIVE: + kv_cache.suspend() + break + + DBG_PRINT and print( # type: ignore[arg-type] + f"update_batch: found {len(removed)} finished requests, now with {len(self.batch)} requests" + ) + + @parameterized.expand( + [ + (1000, 1000, 1024, True, 32, 32), + (100, 100, 128, False, 1, 128), + (100, 100, 128, False, 4, 64), + ] + ) + # @assert_no_ref_cycle + def test_inflight_batching( + self, + num_requests: int, + avg_length: int, + gpu_quota_mb: int, + skip_execution: bool, + interval: int, + tokens_per_block: int, + ): + self.prepare( + gpu_quota_mb << 20, 4 << 30, 0 << 30, 36, 128, 0, tokens_per_block=tokens_per_block + ) + self.num_requests = num_requests + self.avg_length = avg_length + self.interval = interval + profile = False + profiler = None + if profile: + import cProfile + + profiler = cProfile.Profile() + profiler.enable() + tic = time.perf_counter() + with TemporaryCudaStream([]) as stream, enable_kernel_delay(): + i = itertools.count() + self.update_batch(stream.handle) + while self.num_finished < self.num_requests: + DBG_PRINT and print( # type: ignore[arg-type] + f"Executing batch {next(i)} with size {len(self.batch)}" + ) + assert self.batch + if not skip_execution: + self.engine.execute(self.batch, stream.handle) + self.update_batch(stream.handle) + toc = time.perf_counter() + if profiler is not None: + profiler.disable() + profiler.print_stats(sort="cumtime") + profiler.dump_stats("profiler.prof") + if DBG_PRINT or PRINT_TIME: + print( + f"Time taken: {toc - tic} seconds (num_prompt_tokens: {self.acc_num_prompt_tokens}, " + f"num_decode_tokens: {self.acc_num_decode_tokens})" + ) + stream.take_finish_event().synchronize() + + +class TestDisagg(TestKVCacheManagerV2): + @parameterized.expand([512]) + # @assert_no_ref_cycle + def test_disagg(self, prompt_len: int) -> None: + self.prepare(128 << 20, 128 << 20, 1 << 30, 36, 128, 0) + lora_task_id = None + prompt = [self.next_token() for _ in range(prompt_len)] + kv_cache = self.manager.create_kv_cache(lora_task_id, prompt) + assert kv_cache.num_committed_tokens == 0 + with TemporaryCudaStream([]) as stream: + success = kv_cache.resume(stream.handle) + assert success + success = kv_cache.resize(prompt_len, prompt_len) + assert success + + def transfer() -> None: + return None + + transfer() + kv_cache.commit(prompt) + kv_cache.close() + stream.take_finish_event().synchronize() + + +if __name__ == "__main__": + a = TestNoBatching() + a.test_sol_mem_utilization() diff --git a/tests/integration/defs/accuracy/references/gpqa_diamond.yaml b/tests/integration/defs/accuracy/references/gpqa_diamond.yaml index 404dcbb6b705..6c16a8bad679 100644 --- a/tests/integration/defs/accuracy/references/gpqa_diamond.yaml +++ b/tests/integration/defs/accuracy/references/gpqa_diamond.yaml @@ -79,3 +79,5 @@ GPT-OSS/120B-MXFP4: - quant_algo: W4A16_MXFP4 kv_cache_quant_algo: FP8 accuracy: 65.0 + - kv_cache_quant_algo: FP8 + accuracy: 65.0 diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index e2bca37c5110..78fbb23f86b6 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -249,6 +249,8 @@ GPT-OSS/120B-MXFP4: - quant_algo: W4A16_MXFP4 kv_cache_quant_algo: FP8 accuracy: 90.3 + - kv_cache_quant_algo: FP8 + accuracy: 90.3 GPT-OSS/20B-MXFP4: - accuracy: 85.0 - quant_algo: W4A8_MXFP4_MXFP8 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index af38a021c743..dbabe9fa0b2e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -22,10 +22,11 @@ from tensorrt_llm._torch.model_config import MoeLoadBalancerConfig from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton import \ IS_TRITON_KERNELS_AVAILABLE -from tensorrt_llm.llmapi import (AutoDecodingConfig, CudaGraphConfig, - EagleDecodingConfig, KvCacheConfig, MoeConfig, - MTPDecodingConfig, NGramDecodingConfig, - SamplingParams, TorchCompileConfig) +from tensorrt_llm.llmapi import (AutoDecodingConfig, CapacitySchedulerPolicy, + CudaGraphConfig, EagleDecodingConfig, + KvCacheConfig, MoeConfig, MTPDecodingConfig, + NGramDecodingConfig, SamplingParams, + SchedulerConfig, TorchCompileConfig) from tensorrt_llm.quantization import QuantAlgo from ..conftest import (get_device_count, get_device_memory, llm_models_root, @@ -3624,7 +3625,7 @@ class TestGPTOSS(LlmapiAccuracyTestHarness): "apply_chat_template": True, } - MODEL_PATH = f"{llm_models_root()}/gpt_oss/gpt-oss-120b" + MODEL_PATH = f"openai/gpt-oss-120b" @pytest.mark.parametrize( "kv_cache_dtype", @@ -3639,7 +3640,10 @@ class TestGPTOSS(LlmapiAccuracyTestHarness): ]) def test_w4_1gpu(self, kv_cache_dtype, moe_backend, cuda_graph, overlap_scheduler, mocker): - mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 8192) + model = f"{llm_models_root()}/gpt_oss/gpt-oss-120b" + mocker.patch.object(GSM8K, "NUM_SAMPLES", 10) + mocker.patch.object(GSM8K, "MAX_OUTPUT_LEN", 32768) + mocker.patch.object(GSM8K, "MAX_INPUT_LEN", 20) mocker.patch.dict(GSM8K.EVALUATE_KWARGS, {"scores_filter": "exact_match,flexible-extract"}) if moe_backend == "TRITON" and not IS_TRITON_KERNELS_AVAILABLE: @@ -3650,15 +3654,19 @@ def test_w4_1gpu(self, kv_cache_dtype, moe_backend, cuda_graph, cuda_graph_config=CudaGraphConfig() if cuda_graph else None) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5, - dtype=kv_cache_dtype) + dtype=kv_cache_dtype, + use_kv_cache_manager_v2=True) - llm = LLM(self.MODEL_PATH, + llm = LLM(model, tensor_parallel_size=1, pipeline_parallel_size=1, moe_expert_parallel_size=1, kv_cache_config=kv_cache_config, **pytorch_config, - moe_config=MoeConfig(backend=moe_backend)) + moe_config=MoeConfig(backend=moe_backend), + scheduler_config=SchedulerConfig( + capacity_scheduler_policy=CapacitySchedulerPolicy. + MAX_UTILIZATION)) with llm: model_name = "GPT-OSS/20B-MXFP4" @@ -3714,7 +3722,8 @@ def test_w4_4gpus(self, kv_cache_dtype, moe_backend, tp_size, pp_size, cuda_graph_config=CudaGraphConfig() if cuda_graph else None) kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, - dtype=kv_cache_dtype) + dtype=kv_cache_dtype, + use_kv_cache_manager_v2=True) max_seq_len = MAX_INPUT_LEN + MAX_OUTPUT_LEN llm = LLM(self.MODEL_PATH, diff --git a/tests/unittest/kv_cache_manager_v2/fake_engine.py b/tests/unittest/kv_cache_manager_v2/fake_engine.py index 68bd0a83282c..ef288f61102c 100644 --- a/tests/unittest/kv_cache_manager_v2/fake_engine.py +++ b/tests/unittest/kv_cache_manager_v2/fake_engine.py @@ -52,6 +52,27 @@ with DynamicPathManager(os.path.dirname(os.path.abspath(__file__))): from kernels import check_values, fill_values +from kernels import check_values, fill_values + +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + BeamIndex, + CudaStream, + KVCacheManagerConfig, + LayerId, + TokenIdExt, + _KVCache, +) +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX, NDEBUG, MemAddress +from tensorrt_llm.runtime.kv_cache_manager_v2._config import AttentionLayerConfig, DataRole +from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + div_up, + exact_div, + get_uniform_attribute, + overlap, + typed_range, + value_or, +) + class Step(NamedTuple): kv_cache: _KVCache @@ -122,6 +143,7 @@ def _check_pages( role = buf.role token_bytes = exact_div(buf.size, tokens_per_block) pool = manager.get_mem_pool_base_address(layer_id, role) + print(f"pool: {pool}, layer_id: {layer_id}, role: {role}") stride = manager.get_page_stride(layer_id, role) lc_id = manager._storage._layer_to_life_cycle_ids[layer_id] pages = kv_cache.get_page_indices(lc_id, beam) diff --git a/tests/unittest/kv_cache_manager_v2/kernels.py b/tests/unittest/kv_cache_manager_v2/kernels.py index 9b81bbf5e9be..673d2b9a0d3c 100644 --- a/tests/unittest/kv_cache_manager_v2/kernels.py +++ b/tests/unittest/kv_cache_manager_v2/kernels.py @@ -9,6 +9,14 @@ from cuda.core.experimental import Kernel, Program, ProgramOptions from cuda.core.experimental._module import ObjectCode +from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( + CudaStream, + LayerId, + MemAddress, + TokenIdExt, +) +from tensorrt_llm.runtime.kv_cache_manager_v2._utils import _unwrap, div_up, exact_div + if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: from kv_cache_manager_v2._common import CudaStream, LayerId, MemAddress, TokenIdExt from kv_cache_manager_v2._utils import _unwrap, div_up, exact_div diff --git a/tests/unittest/kv_cache_manager_v2/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2/test_kv_cache_manager_v2.py index 6defee901bfd..c902b5480052 100755 --- a/tests/unittest/kv_cache_manager_v2/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2/test_kv_cache_manager_v2.py @@ -691,4 +691,5 @@ def transfer() -> None: if __name__ == "__main__": - unittest.main() + a = TestNoBatching() + a.test_sol_mem_utilization()