From 61be95e18f9e5d8e4e9950b90f74b7bb7c6e077d Mon Sep 17 00:00:00 2001 From: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com> Date: Wed, 30 Sep 2026 07:21:43 -0700 Subject: [PATCH 01/26] [None][feat] Add DeepSeek V4.1 support Signed-off-by: ziyixiong-nv <219238287+ziyixiong-nv@users.noreply.github.com> --- .gitignore | 1 + 3rdparty/fetch_content.json | 5 +- ...eepgemm_configurable_barrier_timeout.patch | 82 - .../kv_cache_manager_v2/blockRadixTree.cpp | 7 +- .../kv_cache_manager_v2/config.h | 16 + .../kv_cache_manager_v2/kvCache.cpp | 306 +- .../kv_cache_manager_v2/kvCache.h | 41 +- .../kv_cache_manager_v2/kvCacheManager.cpp | 10 +- .../kv_cache_manager_v2/lifeCycleRegistry.cpp | 3 +- .../kv_cache_manager_v2/lifeCycleRegistry.h | 43 +- .../kv_cache_manager_v2/page.cpp | 23 +- .../batch_manager/kv_cache_manager_v2/page.h | 2 + .../kv_cache_manager_v2/storage/config.cpp | 9 +- .../kv_cache_manager_v2/storageManager.cpp | 32 +- cpp/tensorrt_llm/deep_gemm/CMakeLists.txt | 30 +- .../deep_jit_shim/elfutils/libdwfl.h | 72 + .../compressorKernels/compressorKernels.cu | 45 +- .../kernels/deepseekV4QNormKernel.cu | 40 +- .../kernels/deepseekV4QNormKernel.h | 4 +- .../mhcKernels/fused_tf32_pmap_gemm.cuh | 20 +- .../kernels/mhcKernels/mhcFusedHcKernel.cu | 144 +- .../kernels/mhcKernels/mhcKernels.cu | 77 +- .../kernels/mhcKernels/mhcKernels.h | 13 +- .../kernels/mhcKernels/mhc_fused_fma.cuh | 16 +- .../moe/communication/moeAlltoAllKernels.cu | 7 +- .../moe/routing/customMoeRoutingKernels.cu | 37 +- .../moe/routing/customMoeRoutingKernels.h | 3 +- .../trtllmGenKernels/fmha/fmhaReduction.cu | 14 +- .../batch_manager/kvCacheManagerV2.cpp | 82 +- cpp/tensorrt_llm/thop/CMakeLists.txt | 9 + cpp/tensorrt_llm/thop/deepseekV4QNormOp.cpp | 10 +- cpp/tensorrt_llm/thop/mhcOp.cpp | 76 +- .../thop/moe/routing/moeGateOp.cpp | 23 +- .../kvCacheManagerV2ColdPageTest.cpp | 245 + ...CacheManagerV2DefaultColdPageCodecTest.cpp | 48 +- .../kvCacheManagerV2StatsTest.cpp | 434 ++ docker/Dockerfile.multi | 6 + docker/Makefile | 8 +- docker/common/install_base.sh | 11 + docs/source/features/disagg-serving.md | 31 + docs/source/models/supported-models.md | 2 + pyproject.toml | 6 + .../_torch/attention/backends/interface.py | 9 + .../attention/backends/sparse/csa2/README.md | 422 ++ .../backends/sparse/csa2/__init__.py | 3 + .../attention/backends/sparse/csa2/backend.py | 845 +++ .../backends/sparse/csa2/cache_manager.py | 1233 ++++ .../attention/backends/sparse/csa2/ced.py | 47 + .../backends/sparse/csa2/compressor.py | 106 + .../backends/sparse/csa2/decoder_replay.py | 394 ++ .../attention/backends/sparse/csa2/indexer.py | 1508 +++++ .../attention/backends/sparse/csa2/kernel.py | 5041 +++++++++++++++++ .../backends/sparse/csa2/metadata.py | 2659 +++++++++ .../attention/backends/sparse/csa2/module.py | 652 +++ .../attention/backends/sparse/csa2/params.py | 266 + .../backends/sparse/csa2/quantization.py | 388 ++ .../attention/backends/sparse/csa2/weights.py | 192 + .../attention/backends/sparse/dsa/indexer.py | 305 +- .../attention/backends/sparse/dsa/metadata.py | 7 +- .../attention/backends/sparse/registry.py | 8 + .../_torch/attention/backends/trtllm.py | 6 + tensorrt_llm/_torch/configs/__init__.py | 16 + tensorrt_llm/_torch/configs/deepseek_v41.py | 1314 +++++ .../_torch/custom_ops/cpp_custom_ops.py | 2 + .../_torch/disaggregation/base/backend.py | 6 +- .../_torch/disaggregation/native/auxiliary.py | 8 +- .../_torch/disaggregation/native/transfer.py | 2 + .../disaggregation/resource/cache_reuse.py | 26 +- .../disaggregation/resource/kv_extractor.py | 23 +- .../_torch/disaggregation/transceiver.py | 23 +- tensorrt_llm/_torch/host_staging.py | 61 + tensorrt_llm/_torch/model_config.py | 101 +- tensorrt_llm/_torch/models/__init__.py | 2 + tensorrt_llm/_torch/models/_arch_index.py | 5 + .../checkpoints/hf/checkpoint_loader.py | 30 +- .../_torch/models/deepseek_v41_processor.py | 387 ++ .../_torch/models/deepseek_v41_vision.py | 187 + .../_torch/models/modeling_deepseekv4.py | 536 +- .../_torch/models/modeling_deepseekv41.py | 2568 +++++++++ .../models/modeling_deepseekv41_multimodal.py | 247 + tensorrt_llm/_torch/models/modeling_dspark.py | 104 +- .../_torch/models/modeling_dspark_v41.py | 318 ++ .../_torch/models/modeling_speculative.py | 11 + .../_torch/modules/engram/__init__.py | 18 +- .../_torch/modules/engram/_decode_hash.py | 85 + tensorrt_llm/_torch/modules/engram/engram.py | 1020 +++- .../_torch/modules/engram/functional.py | 190 + .../_torch/modules/engram/projection.py | 95 + tensorrt_llm/_torch/modules/linear.py | 11 +- .../_torch/modules/mhc/hyper_connection.py | 265 +- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 99 +- .../_torch/moe/fused_moe/configurable_moe.py | 7 + .../_torch/moe/fused_moe/interface.py | 12 + .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 8 +- .../fused_moe/mega_moe/mega_moe_deepgemm.py | 39 +- .../_torch/moe/fused_moe/moe_scheduler.py | 66 +- tensorrt_llm/_torch/moe/fused_moe/routing.py | 98 +- tensorrt_llm/_torch/pyexecutor/_util.py | 14 +- tensorrt_llm/_torch/pyexecutor/ced_replay.py | 58 + .../_torch/pyexecutor/config_utils.py | 24 + .../_torch/pyexecutor/cuda_graph_runner.py | 60 +- .../pyexecutor/engine/runners/common.py | 10 +- .../engine/runners/decoder/runner.py | 442 +- .../kv_cache/kv_cache_manager_v2.py | 70 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 57 +- .../_torch/pyexecutor/model_loader.py | 2 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 300 +- .../_torch/pyexecutor/resource_manager.py | 6 +- .../_torch/pyexecutor/scheduler/adp_router.py | 18 +- .../_torch/pyexecutor/scheduler/scheduler.py | 15 +- .../pyexecutor/scheduler/scheduler_v2.py | 123 +- tensorrt_llm/_torch/speculative/dspark.py | 55 +- tensorrt_llm/_torch/speculative/utils.py | 5 + tensorrt_llm/commands/serve.py | 7 + tensorrt_llm/disaggregated_params.py | 20 + tensorrt_llm/executor/result.py | 32 +- tensorrt_llm/inputs/utils.py | 20 +- tensorrt_llm/llmapi/disagg_utils.py | 9 +- tensorrt_llm/llmapi/llm_args.py | 79 +- tensorrt_llm/llmapi/reasoning_parser.py | 12 + .../runtime/kv_cache_manager_v2/__init__.py | 4 + .../runtime/kv_cache_manager_v2/__init__.pyi | 46 +- .../kv_cache_manager_v2/_introspection.py | 5 + tensorrt_llm/serve/openai_client.py | 122 +- tensorrt_llm/serve/openai_disagg_service.py | 50 +- tensorrt_llm/serve/openai_protocol.py | 29 + .../serve/tool_parser/deepseekv32_parser.py | 16 +- .../serve/tool_parser/deepseekv41_parser.py | 27 + .../serve/tool_parser/tool_parser_factory.py | 4 + .../tokenizer/deepseek_v41/__init__.py | 18 + .../tokenizer/deepseek_v41/tokenizer.py | 638 +++ tensorrt_llm/tokenizer/tokenizer.py | 2 + tensorrt_llm/usage/architecture_allowlist.py | 2 + .../usage/llm_args_golden_manifest.json | 6 + .../defs/accuracy/references/gsm8k.yaml | 9 + .../defs/accuracy/test_llm_api_pytorch.py | 164 + tests/integration/defs/cpp/conftest.py | 8 + tests/integration/defs/deepseek_v41_gsm8k.py | 198 + .../integration/defs/deepseek_v41_serving.py | 129 + .../defs/disaggregated/test_disaggregated.py | 250 +- tests/integration/defs/perf/_model_paths.py | 4 + .../integration/defs/test_deepseek_v41_ced.py | 235 + tests/integration/defs/test_unittests.py | 56 + .../test_lists/qa/llm_function_core.txt | 3 + tests/integration/test_lists/waives.txt | 1 - .../deepseek_v41_flash_grace_blackwell.yaml | 125 + .../kernels/test_deepseek_v4_q_norm.py | 98 +- .../_torch/attention/sparse/csa2/__init__.py | 2 + .../_torch/attention/sparse/csa2/_utils.py | 430 ++ .../attention/sparse/csa2/test_backend.py | 1516 +++++ .../sparse/csa2/test_cache_manager.py | 574 ++ .../sparse/csa2/test_context_only_cache.py | 236 + .../sparse/csa2/test_context_only_metadata.py | 108 + .../attention/sparse/csa2/test_indexer.py | 2552 +++++++++ .../sparse/csa2/test_indexer_descriptors.py | 292 + .../attention/sparse/csa2/test_kernel.py | 1884 ++++++ .../attention/sparse/csa2/test_metadata.py | 2327 ++++++++ .../sparse/csa2/test_metadata_uploads.py | 151 + .../attention/sparse/csa2/test_module.py | 1949 +++++++ .../deepseek_v4/test_compressor_kernel.py | 16 + .../test_deepseek_v4_variant_guard.py | 54 + .../deepseek_v4/test_indexer_slot_mapping.py | 210 + .../attention/sparse/dsa/test_dsa_indexer.py | 7 +- .../sparse/dsa/test_req_idx_per_token.py | 6 + .../test_disagg_coordinator_receive.py | 1 + .../_torch/executor/engine/test_common.py | 15 +- .../kv_cache/test_dual_pool_kv_cache.py | 1 + .../kv_cache/test_kv_cache_estimation.py | 51 + .../kv_cache/test_kv_cache_v2_scheduler.py | 162 + .../_torch/executor/test_adp_router.py | 2 +- .../executor/test_context_only_warmup.py | 143 + .../executor/test_mla_workspace_reserve.py | 2 + .../test_model_loader_checkpoint_io.py | 106 +- .../executor/test_multimodal_scheduler.py | 1 + .../_torch/executor/test_py_executor.py | 437 +- .../executor/test_pytorch_model_engine.py | 258 +- .../test_pytorch_model_engine_warmup.py | 8 +- .../_torch/executor/test_recompute_pause.py | 2 + .../executor/test_remote_tail_adp_phase.py | 678 +++ .../executor/test_token_budget_fallback.py | 16 + .../unittest/_torch/misc/test_host_staging.py | 64 + .../modeling/_deepseek_v41_test_utils.py | 444 ++ .../_torch/modeling/test_deepseek_v41_ced.py | 453 ++ .../test_deepseek_v41_decoder_global.py | 123 + .../test_deepseek_v41_encoder_replay.py | 260 + .../modeling/test_deepseek_v41_multimodal.py | 302 + .../test_deepseek_v41_multimodal_backbone.py | 282 + .../modeling/test_deepseek_v41_snapshots.py | 428 ++ .../modeling/test_modeling_deepseekv41.py | 3548 ++++++++++++ tests/unittest/_torch/modules/test_engram.py | 840 +-- .../_torch/modules/test_engram_decode_hash.py | 515 ++ .../modules/test_engram_history_uploads.py | 185 + .../_torch/modules/test_engram_replay.py | 148 + tests/unittest/_torch/modules/test_mhc.py | 481 +- .../_torch/modules/test_mxfp8_linear.py | 4 + .../moe/test_deepseek_v41_parallel_moe.py | 271 + .../moe/test_deepseek_v41_shared_moe.py | 301 + .../_torch/moe/test_kimi_k3_situ_moe.py | 21 +- tests/unittest/_torch/moe/test_moe_comm.py | 16 + tests/unittest/_torch/moe/test_moe_impl.py | 4 +- tests/unittest/_torch/moe/test_moe_routing.py | 78 + .../test_deepseek_v41_attention_dp.py | 947 ++++ .../_torch/multi_gpu/test_engram_sharding.py | 157 + .../hw_agnostic/test_dspark_v41.py | 533 ++ .../hw_agnostic/test_dspark_worker.py | 63 + .../_torch/test_custom_config_registration.py | 37 +- tests/unittest/_torch/test_model_config.py | 68 + .../disaggregated/test_cache_reuse_adapter.py | 133 +- .../test_context_only_cache_transfer.py | 276 + .../disaggregated/test_csa2_kv_transfer.py | 268 + .../test_deepseek_v4_kv_transfer.py | 7 + .../test_disagg_openai_client.py | 380 ++ .../disaggregated/test_disagg_utils.py | 12 + .../test_disaggregated_params.py | 24 + .../test_openai_disagg_service.py | 183 + .../test_verified_reuse_admission.py | 20 +- .../test_kv_cache_concurrency.py | 51 + .../test_optional_checkpoint.py | 63 + tests/unittest/llmapi/_run_mpi_comm_task.py | 15 + .../unittest/llmapi/apps/test_tool_parsers.py | 72 + .../llmapi/test_deepseek_v41_tokenizer.py | 576 ++ tests/unittest/llmapi/test_executor.py | 71 + tests/unittest/llmapi/test_llm_args.py | 57 + .../unittest/llmapi/test_reasoning_parser.py | 40 +- .../others/test_deepseek_v41_gsm8k.py | 233 + 225 files changed, 53649 insertions(+), 1463 deletions(-) delete mode 100644 3rdparty/patches/deepgemm_configurable_barrier_timeout.patch create mode 100644 cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/README.md create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/__init__.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/backend.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/cache_manager.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/ced.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/compressor.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/decoder_replay.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/indexer.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/kernel.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/metadata.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/module.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/params.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/quantization.py create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/csa2/weights.py create mode 100644 tensorrt_llm/_torch/configs/deepseek_v41.py create mode 100644 tensorrt_llm/_torch/host_staging.py create mode 100644 tensorrt_llm/_torch/models/deepseek_v41_processor.py create mode 100644 tensorrt_llm/_torch/models/deepseek_v41_vision.py create mode 100644 tensorrt_llm/_torch/models/modeling_deepseekv41.py create mode 100644 tensorrt_llm/_torch/models/modeling_deepseekv41_multimodal.py create mode 100644 tensorrt_llm/_torch/models/modeling_dspark_v41.py create mode 100644 tensorrt_llm/_torch/modules/engram/_decode_hash.py create mode 100644 tensorrt_llm/_torch/modules/engram/functional.py create mode 100644 tensorrt_llm/_torch/modules/engram/projection.py create mode 100644 tensorrt_llm/_torch/pyexecutor/ced_replay.py create mode 100644 tensorrt_llm/serve/tool_parser/deepseekv41_parser.py create mode 100644 tensorrt_llm/tokenizer/deepseek_v41/__init__.py create mode 100644 tensorrt_llm/tokenizer/deepseek_v41/tokenizer.py create mode 100644 tests/integration/defs/deepseek_v41_gsm8k.py create mode 100644 tests/integration/defs/deepseek_v41_serving.py create mode 100644 tests/integration/defs/test_deepseek_v41_ced.py create mode 100644 tests/scripts/perf-sanity/aggregated/deepseek_v41_flash_grace_blackwell.yaml create mode 100644 tests/unittest/_torch/attention/sparse/csa2/__init__.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/_utils.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_backend.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_context_only_cache.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_context_only_metadata.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_indexer.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_indexer_descriptors.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_kernel.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_metadata.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_metadata_uploads.py create mode 100644 tests/unittest/_torch/attention/sparse/csa2/test_module.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_variant_guard.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_indexer_slot_mapping.py create mode 100644 tests/unittest/_torch/executor/test_context_only_warmup.py create mode 100644 tests/unittest/_torch/executor/test_remote_tail_adp_phase.py create mode 100644 tests/unittest/_torch/misc/test_host_staging.py create mode 100644 tests/unittest/_torch/modeling/_deepseek_v41_test_utils.py create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_ced.py create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_decoder_global.py create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_encoder_replay.py create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_multimodal_backbone.py create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_snapshots.py create mode 100644 tests/unittest/_torch/modeling/test_modeling_deepseekv41.py create mode 100644 tests/unittest/_torch/modules/test_engram_decode_hash.py create mode 100644 tests/unittest/_torch/modules/test_engram_history_uploads.py create mode 100644 tests/unittest/_torch/modules/test_engram_replay.py create mode 100644 tests/unittest/_torch/moe/test_deepseek_v41_parallel_moe.py create mode 100644 tests/unittest/_torch/moe/test_deepseek_v41_shared_moe.py create mode 100644 tests/unittest/_torch/multi_gpu/test_deepseek_v41_attention_dp.py create mode 100644 tests/unittest/_torch/multi_gpu/test_engram_sharding.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_v41.py create mode 100644 tests/unittest/disaggregated/test_context_only_cache_transfer.py create mode 100644 tests/unittest/disaggregated/test_csa2_kv_transfer.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_optional_checkpoint.py create mode 100644 tests/unittest/llmapi/test_deepseek_v41_tokenizer.py create mode 100644 tests/unittest/others/test_deepseek_v41_gsm8k.py diff --git a/.gitignore b/.gitignore index 6c115ad2a40c..8754313d88fb 100644 --- a/.gitignore +++ b/.gitignore @@ -81,6 +81,7 @@ ad-test-workspace/ */tllm_debug/** *.patch !cpp/tensorrt_llm/deep_ep/*.patch +!3rdparty/patches/*.patch !3rdparty/vendor_patches/*.patch examples/disaggregated/slurm/benchmark/logs/ scripts/attribution/data/checksum_to_paths.yml diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 25f64fbda4f4..89bef45b8482 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -39,10 +39,9 @@ { "name": "deepgemm", "git_repository": "https://github.com/deepseek-ai/DeepGEMM", - "git_tag": "8b1392b978f5a03c828dd1711090d7fb50958b8a", + "git_tag": "b64107f2b9599ca76445b7f62eedf66bae1d095b", "git_submodules_recurse": true, - "source_subdir": "dont-add-this-project-with-add-subdirectory", - "patch_file": "patches/deepgemm_configurable_barrier_timeout.patch" + "source_subdir": "dont-add-this-project-with-add-subdirectory" }, { "name": "eigen", diff --git a/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch b/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch deleted file mode 100644 index b9d89bcf6d4c..000000000000 --- a/3rdparty/patches/deepgemm_configurable_barrier_timeout.patch +++ /dev/null @@ -1,82 +0,0 @@ ---- a/csrc/jit/compiler.hpp -+++ b/csrc/jit/compiler.hpp -@@ -56,12 +56,18 @@ public: - "--ptxas-options=--register-usage-level=10", - get_env("DG_JIT_CPP_STANDARD", 20)); - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0) or get_env("DG_JIT_PTXAS_CHECK", 0)) - flags += " --ptxas-options=--verbose,--warn-on-local-memory-usage"; - if (get_env("DG_JIT_WITH_LINEINFO", 0)) - flags += " -Xcompiler -rdynamic -lineinfo"; -+ -+ // Allow overriding the in-kernel barrier timeout, see `comm/barrier.cuh` -+ // NOTES: `flags` is part of the kernel signature, so changing this value -+ // naturally invalidates the JIT cache -+ if (const auto barrier_timeout_seconds = get_env("DG_JIT_BARRIER_TIMEOUT_SECONDS", 0); barrier_timeout_seconds > 0) -+ flags += fmt::format(" -DDG_BARRIER_TIMEOUT_SECONDS={}", barrier_timeout_seconds); - } - - virtual ~Compiler() = default; - - std::filesystem::path make_tmp_dir() const { - return make_dirs(cache_dir_path / "tmp"); ---- a/deep_gemm/include/deep_gemm/comm/barrier.cuh -+++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh -@@ -5,14 +5,22 @@ - #include - #include - #include - - namespace deep_gemm::comm { - --// 60s timeout, at 2 GHz --constexpr int64_t kNumTimeoutCycles = 60ll * 2000000000ll; -+// Barrier timeout in seconds, overridable at JIT compile time. The host-side compiler -+// turns the `DG_JIT_BARRIER_TIMEOUT_SECONDS` environment variable into this definition, -+// so slow-but-healthy rendezvous (e.g. large expert-parallel groups whose first -+// collective follows a long warm-up) can be distinguished from a genuine hang. -+#ifndef DG_BARRIER_TIMEOUT_SECONDS -+#define DG_BARRIER_TIMEOUT_SECONDS 60 -+#endif -+ -+// Timeout in cycles, at 2 GHz -+constexpr int64_t kNumTimeoutCycles = static_cast(DG_BARRIER_TIMEOUT_SECONDS) * 2000000000ll; - - CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { - // Perform cluster_sync with `barrier.cluster.arrive.relaxed` - // This is slightly faster than `cute::cluster_sync` but has weaker memory ordering guarantee - cute::cluster_arrive_relaxed(); - cute::cluster_wait(); -@@ -31,14 +39,14 @@ CUTLASS_DEVICE void grid_sync(const layout::Workspace& workspace, - count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); - uint32_t new_value; - const auto start_clock = clock64(); - do { - new_value = ptx::ld_acq(count_ptr); - if (clock64() - start_clock >= kNumTimeoutCycles) { -- printf("DeepGEMM grid sync timeout: sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", -- sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); -+ printf("DeepGEMM grid sync timeout (%ds): sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", -+ DG_BARRIER_TIMEOUT_SECONDS, sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); - DG_DEVICE_ASSERT(false and "Grid sync timeout"); - } - } while (((new_value ^ old_value) & kFinishSumTag) == 0); - } - sync_scope(); - } -@@ -72,14 +80,14 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, - if (thread_idx == 0) { - ptx::red_add(counter_ptr, 1); - const int target = signal_sign ? 0 : static_cast(kNumRanks); - const auto start_clock = clock64(); - while (ptx::ld_acq_sys(signal_ptr) != target) { - if (clock64() - start_clock >= kNumTimeoutCycles) { -- printf("DeepGEMM NVLink barrier timeout: rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", -- sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); -+ printf("DeepGEMM NVLink barrier timeout (%ds): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", -+ DG_BARRIER_TIMEOUT_SECONDS, sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); - DG_DEVICE_ASSERT(false and "NVLink barrier timeout"); - } - } - } - } - diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index d38b7ec2d7ef..1b015fc48af6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -482,7 +482,8 @@ std::vector> Block::clearStaleBlocksAfterPageUnlink( // If this is a full-attention block or a sink block: evict subtree. // Mirrors Python: pages = remove_subtree(self) - if (alc && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks})) + if (alc && alc->requiresPrefixCoverage() + && (!alc->windowSize.has_value() || block.ordinal() < BlockOrdinal{alc->numSinkBlocks})) { pruneStart = block.prev; detachedBlocks.push_back(removeSubtree(block)); @@ -887,6 +888,10 @@ std::vector BlockRadixTree::pruneMatch( bool shortened = false; for (auto [lcId, attn] : attnLcs) { + if (!attn->requiresPrefixCoverage()) + { + continue; + } auto const staleRange = attn->getStaleRange(numTok, mTokensPerBlock); int const staleBeg = staleRange.beg.value(); int const staleEnd = staleRange.end.value(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index bb40369610e3..a6cf8530f21b 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -137,6 +137,19 @@ inline void validateNoDuplicateBufferRoles(std::vector const& buff } // namespace detail +//! Controls whether an attention lifecycle participates in prefix matching and page reuse. +enum class AttentionReusePolicy +{ + //! Commit reusable pages and require live-window coverage when matching a prefix. + REQUIRED = 0, + //! Keep request-local working pages; never publish or reuse them through the radix tree. + PRIVATE = 1, + //! Commit reusable pages without constraining the REQUIRED prefix match. The caller selects + //! available groups at resume; missing or unselected groups use fresh pages for recomputation. + //! These pages may be evicted independently without removing the REQUIRED prefix. + OPTIONAL = 2, +}; + struct AttentionLayerConfig { static constexpr LayerType type = LayerType::ATTENTION; @@ -150,6 +163,9 @@ struct AttentionLayerConfig // nullopt or 0 = no sink tokens. std::optional numSinkTokens; + // Private working pages never constrain prefix matching or enter the tree. + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + [[nodiscard]] std::optional windowSize() const noexcept { return slidingWindowSize; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index 3a75bbaf2a9e..5a0aaba48fa1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -102,6 +102,7 @@ KvCache::KvCache(KvCacheManager& manager, ReuseScope reuseScope, std::optional(mHistoryLength)); @@ -150,6 +151,128 @@ StorageManager* KvCache::storageManager() const return &mManager->storage(); } +void KvCache::_collectReuseCandidates() +{ + auto const end = numCommittedTokens(); + auto const& lifeCycles = mManager->lifeCycles(); + mOptionalCandidates.resize(lifeCycles.size()); + for (LifeCycleId lc{0}; lc < lifeCycles.size(); ++lc) + { + auto const* attention = std::get_if(&lifeCycles[lc]); + auto const policy = attention ? attention->reusePolicy : AttentionReusePolicy::REQUIRED; + ReuseGroupStatus status{lc, policy, end, end > 0 && policy == AttentionReusePolicy::REQUIRED, {}}; + if (status.complete) + { + if (!attention) + { + status.coverage.emplace_back(end, end); // Exact recurrent-state endpoint. + } + else + { + auto stale = getStaleRange(lifeCycles[lc], end, mTokensPerBlock); + if (stale.beg > BlockOrdinal{0}) + { + status.coverage.emplace_back(0, std::min(end, stale.beg.value() * mTokensPerBlock)); + } + if (stale.end.value() * mTokensPerBlock < end) + { + status.coverage.emplace_back(stale.end.value() * mTokensPerBlock, end); + } + } + } + // OPTIONAL checks coverage at the REQUIRED endpoint without shortening it. + // Use the ordinary SWA live ranges, including sinks and multi-page windows. + if (attention && attention->allowsCheckpoint() && end > 0 && end % mTokensPerBlock == 0 + && attention->windowSize) + { + auto const stale = getStaleRange(lifeCycles[lc], end, mTokensPerBlock); + auto& candidates = mOptionalCandidates[lc]; + status.complete = true; + auto collectRange = [&](BlockOrdinal beg, BlockOrdinal limit) + { + for (auto ordinal = beg; ordinal < limit; ++ordinal) + { + auto const& block = mBlocks[ordinal].treeBlock; + auto* page = block && !block->isOrphan() ? block->getPage(lc) : nullptr; + if (!page || page->numTokensInBlock < mTokensPerBlock) + { + status.complete = false; + return; + } + candidates.emplace_back(ordinal, page->hold()); + } + if (beg < limit) + { + status.coverage.emplace_back(beg.value() * mTokensPerBlock, limit.value() * mTokensPerBlock); + } + }; + collectRange(BlockOrdinal{0}, stale.beg); + if (status.complete) + { + collectRange(stale.end, BlockOrdinal{end / mTokensPerBlock}); + } + if (!status.complete) + { + candidates.clear(); + status.coverage.clear(); + } + } + mReuseStatus.push_back(std::move(status)); + } +} + +bool KvCache::_needsFreshPrefixPages(LifeCycleId lc) const +{ + return !requiresPrefixCoverage(mManager->lifeCycles()[lc]) + && std::find(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end(), lc) + == mSelectedOptionalGroups.end(); +} + +void KvCache::_selectOptionalGroups(std::vector const& groups) +{ + // Validate everything before touching the previous attempt's mappings. + for (auto lc : groups) + { + if (lc < LifeCycleId{0} || lc >= mOptionalCandidates.size() + || mReuseStatus[lc.value()].policy != AttentionReusePolicy::OPTIONAL || !mReuseStatus[lc.value()].complete) + { + throw std::invalid_argument("Selected OPTIONAL group has no reusable state at the claimed prefix"); + } + } + for (auto lc : mSelectedOptionalGroups) + { + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + { + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = std::monostate{}; + } + } + for (auto lc : groups) + { + auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]); + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + { + if (!stale.contains(ordinal)) + { + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] = holder; + } + } + } + mSelectedOptionalGroups = groups; + std::sort(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end()); + mSelectedOptionalGroups.erase( + std::unique(mSelectedOptionalGroups.begin(), mSelectedOptionalGroups.end()), mSelectedOptionalGroups.end()); + for (LifeCycleId lc{0}; lc < mOptionalCandidates.size(); ++lc) + { + if (std::find(groups.begin(), groups.end(), lc) == groups.end() + && mReuseStatus[lc.value()].policy == AttentionReusePolicy::OPTIONAL) + { + mOptionalCandidates[lc].clear(); + mReuseStatus[lc.value()].complete = false; + mReuseStatus[lc.value()].coverage.clear(); + } + } +} + std::vector KvCache::_activePages() const { std::vector result; @@ -172,6 +295,10 @@ std::vector KvCache::_activePages() const // Attention lifecycle: yield non-stale blocks (sink + window). LifeCycle const& lc = lcs.getLifeCycle(lcId); + if (mNeverResumed && _needsFreshPrefixPages(lcId)) + { + continue; // Allocated by first resume; private pages are never claimed from the tree. + } auto staleRange = _getStaleRange(mHistoryLength, lc); BlockOrdinal staleBeg = staleRange.beg; BlockOrdinal staleEnd = staleRange.end; @@ -378,7 +505,8 @@ void KvCache::activate() } } -bool KvCache::resume(std::optional stream, std::optional isDecoding) +bool KvCache::resume(std::optional stream, std::optional isDecoding, + std::optional> optionalReuseGroups) { KVCM2_API_GUARD(); TLLM_CHECK(mStatus == Status::SUSPENDED); @@ -397,6 +525,14 @@ bool KvCache::resume(std::optional stream, std::optional isDecod throw std::invalid_argument("Cannot return a decoding cache to prefill"); auto restorePhase = FuncGuard([&]() { mIsDecoding = oldIsDecoding; }); mIsDecoding = isDecoding.value_or(mIsDecoding); + if (mNeverResumed) + { + _selectOptionalGroups(optionalReuseGroups.value_or(std::vector{})); + } + else if (optionalReuseGroups.has_value()) + { + throw std::invalid_argument("OPTIONAL reuse selection is only valid before the first resume"); + } // Check utilization against threshold. auto const utilizations = mManager->storage().getUtilization(kHotLevel); @@ -410,6 +546,26 @@ bool KvCache::resume(std::optional stream, std::optional isDecod auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); LifeCycleId numLc = storageMgr.numLifeCycles(); + // Snapshot source levels before allocation or activation can migrate the selected pages. + ReusedBlocksByLevelByLifeCycle optionalReusedByLevel; + if (mNeverResumed && _shouldRecordStats()) + { + for (auto lc : mSelectedOptionalGroups) + { + auto& byLevel = optionalReusedByLevel[lc]; + byLevel.full = CountsByLevel(storageMgr.numCacheLevels(), 0); + byLevel.partial = CountsByLevel(storageMgr.numCacheLevels(), 0); + auto const stale = _getStaleRange(mHistoryLength, mManager->lifeCycles()[lc]); + for (auto const& [ordinal, holder] : mOptionalCandidates[lc]) + { + if (!stale.contains(ordinal)) + { + ++byLevel.full.at(holder->page->cacheLevel); + } + } + } + } + // Pre-allocate GPU slots for deferred copies (partial blocks + SSM) and scratch slots // before locking, so we never end up in a state where pages are locked but we can't allocate. TypedVec> deferredSlots(numLc); @@ -432,6 +588,22 @@ bool KvCache::resume(std::optional stream, std::optional isDecod } } }); + TypedVec> privateOrdinals(numLc); + TypedVec> privateSlots(numLc); + auto releasePrivateSlots = FuncGuard( + [&]() + { + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + for (auto& slot : privateSlots[lc]) + { + if (slot.hasValidSlot()) + { + storageMgr.releaseSlot(lc, kHotLevel, std::move(slot)); + } + } + } + }); // Compute scratch slot deltas UNCONDITIONALLY (mirrors Python: _take_excess_scratch_slots // is called outside _never_resumed). @@ -448,8 +620,24 @@ bool KvCache::resume(std::optional stream, std::optional isDecod for (LifeCycleId lc{0}; lc < numLc; ++lc) { bool isSsm = ssmLcId.has_value() && lc == *ssmLcId; - if (isSsm || hasPartial) + auto const& lifeCycle = mManager->lifeCycles()[lc]; + if (_needsFreshPrefixPages(lc)) + { + auto const stale = _getStaleRange(mHistoryLength, lifeCycle); + for (BlockOrdinal ord{0}; ord < stale.beg; ++ord) + { + privateOrdinals[lc].push_back(ord); + } + for (BlockOrdinal ord = stale.end; ord < mBlocks.size(); ++ord) + { + privateOrdinals[lc].push_back(ord); + } + numSlotsNeeded[lc] += static_cast(privateOrdinals[lc].size()); + } + else if (isSsm || hasPartial) + { numSlotsNeeded[lc] += 1; + } } } @@ -486,7 +674,14 @@ bool KvCache::resume(std::optional stream, std::optional isDecod if (!tmpSlots[lc].empty()) { // Mirrors Python: `if self._never_resumed and (... SsmLifeCycle or has_partial):` - bool needsDeferred = mNeverResumed && ((ssmLcId.has_value() && lc == *ssmLcId) || hasPartial); + bool const isPrivate = _needsFreshPrefixPages(lc); + bool needsDeferred + = mNeverResumed && !isPrivate && ((ssmLcId.has_value() && lc == *ssmLcId) || hasPartial); + for (size_t i = 0; i < privateOrdinals[lc].size(); ++i) + { + privateSlots[lc].push_back(std::move(tmpSlots[lc].back())); + tmpSlots[lc].pop_back(); + } if (needsDeferred) { // Python uses pop() here: reserve one slot for deferred copy, then treat the rest as scratch. @@ -524,6 +719,35 @@ bool KvCache::resume(std::optional stream, std::optional isDecod mStatus = Status::ACTIVE; auto rollbackActivation = FuncGuard([&]() { _deactivate(); }); + // Private recovery pages have no reusable source. Their contents are + // initialized by the caller before inference, after all allocations succeed. + std::vector privateReadyEvents; + for (auto const& slots : privateSlots) + { + for (auto const& slot : slots) + { + privateReadyEvents.push_back(slot.readyEvent.handle()); + } + } + streamWaitEvents(reinterpret_cast(cudaStream()), std::move(privateReadyEvents)); + for (LifeCycleId lc{0}; lc < numLc; ++lc) + { + for (size_t i = 0; i < privateSlots[lc].size(); ++i) + { + auto const ordinal = privateOrdinals[lc][i]; + auto page = makeShared(*this, ordinal, lc, kHotLevel, kDefaultBeamIndex); + page->setSlot(privateSlots[lc][i]); + mBlocks[ordinal].pages[kDefaultBeamIndex][lc] + = page->lock(*this, kDefaultBeamIndex, ordinal, lc, /*skipWait=*/true); + if (mPendingStats.recordAllocationRange(lc, ordinal, ordinal + 1, + /*beamWidth=*/1, /*countAsMissed=*/true, /*countAsGeneration=*/false, _shouldRecordManagerStats(), + _shouldRecordRequestStats())) + { + mManager->markStatsDirty(id); + } + } + } + // Deferred copy: for partial blocks and SSM, copy from now-locked source pages // to pre-allocated GPU slots, then unlock sources and replace with new pages. if (mNeverResumed) @@ -637,6 +861,19 @@ bool KvCache::resume(std::optional stream, std::optional isDecod mBlocks[lastOrdinal].treeBlock = nullptr; } + if (mNeverResumed && _shouldRecordStats()) + { + for (auto const& [lc, byLevel] : optionalReusedByLevel) + { + int const reused = static_cast(countsByLevelTotal(byLevel.full)); + if (mPendingStats.recordReuse( + lc, reused, 0, byLevel, _shouldRecordManagerStats(), _shouldRecordRequestStats())) + { + mManager->markStatsDirty(id); + } + } + } + mOptionalCandidates.clear(); // Deferred copies survive a failed decode admission and must not be repeated on retry. mNeverResumed = false; if (mIsDecoding) @@ -719,6 +956,7 @@ void KvCache::suspend() KVCM2_API_GUARD(); auto const apiLock = mManager->lockExclusive(); TLLM_CHECK_DEBUG(mStatus == Status::ACTIVE); + mOptionalCandidates.clear(); TLLM_CHECK_DEBUG(_checkSanity()); TLLM_CHECK_DEBUG(!mFinishEvent.has_value()); @@ -776,6 +1014,7 @@ void KvCache::close() if (mStatus == Status::CLOSED) return; + mOptionalCandidates.clear(); discardPendingStats(); stopCommitting(); TLLM_CHECK_DEBUG(_checkSanity()); @@ -1221,7 +1460,10 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) std::vector attachedLcs; for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) { - (void) attn; + if (!attn->requiresPrefixCoverage()) + { + continue; + } auto& bp = beamBlock[lcIdx]; if (blockPageIsNull(bp) || treeBlock->pageCoverage(lcIdx) >= numTokens) { @@ -1874,8 +2116,13 @@ std::vector KvCache::_unlockStaleBlocks(int newHistoryLeng { auto& sb = mBlocks[ord]; bool isCommitted = sb.isCommitted(); - bool holdForCommit - = !mManager->commitMinSnapshot() && !isCommitted && (mCommitState == CommitState::ALLOWED); + // A partial tail can expire pages still needed at the preceding + // aligned checkpoint. Keep its entire live window until publication. + int const alignedEnd = newHistoryLength / mTokensPerBlock * mTokensPerBlock; + bool const keepAlignedCheckpoint = alc.allowsCheckpoint() + && ord < BlockOrdinal{alignedEnd / mTokensPerBlock} && !_getStaleRange(alignedEnd, lc).contains(ord); + bool holdForCommit = !alc.isPrivate() && (!mManager->commitMinSnapshot() || keepAlignedCheckpoint) + && !isCommitted && (mCommitState == CommitState::ALLOWED); for (BeamIndex bi{0}; bi < sb.pages.size(); ++bi) { @@ -1923,8 +2170,10 @@ TypedVec KvCache::_takeUncommittedPage( TypedVec result(numLc, TakenPage{nullptr, false}); for (LifeCycleId lc{0}; lc < numLc; ++lc) { - if (skipLc.has_value() && lc == *skipLc) + if ((skipLc.has_value() && lc == *skipLc) || isPrivateLifeCycle(mManager->lifeCycles()[lc])) + { continue; + } auto& bp = sb.pages[beamIdx][lc]; if (auto* lock = std::get_if(&bp)) { @@ -2042,6 +2291,10 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) { if (ssmLcId.has_value() && lc == *ssmLcId) continue; + if (isPrivateLifeCycle(mManager->lifeCycles()[lc])) + { + continue; + } auto& bp = sb.pages[kDefaultBeamIndex][lc]; if (blockPageIsNull(bp)) continue; @@ -2580,6 +2833,10 @@ void KvCache::_setupForReuse(BlockRadixTree::ReuseMatch const& match) if (ssmLcId.has_value() && lcId == *ssmLcId) continue; + if (!requiresPrefixCoverage(allLc[lcId])) + { + continue; + } auto staleRange = getStaleRange(allLc[lcId], numTokens, mTokensPerBlock); BlockOrdinal staleStart = staleRange.beg; BlockOrdinal staleEnd = staleRange.end; @@ -2746,8 +3003,19 @@ SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const auto page = blockPageGetPage(beamBlock[lcId]); // committed->block may differ from `ret`: a page can be moved to a longer // sibling block, or replaced there by one with larger token coverage. - auto committed = dynamicPointerCast(page); - TLLM_CHECK(committed); + auto const* attention = std::get_if(&mManager->lifeCycles()[lcId]); + if (attention && attention->allowsCheckpoint()) + { + TLLM_CHECK(dynamicPointerCast(page) || dynamicPointerCast(page)); + } + else if (!requiresPrefixCoverage(mManager->lifeCycles()[lcId])) + { + TLLM_CHECK(dynamicPointerCast(page)); + } + else + { + TLLM_CHECK(dynamicPointerCast(page)); + } } } } @@ -2804,6 +3072,12 @@ bool KvCache::_checkSanity() const auto const& staleRange = staleRanges[lc]; auto const& scratchRange = scratchRangesVec[lc]; + bool const isPrivate = !requiresPrefixCoverage(lcs[lc]); + if (_needsFreshPrefixPages(lc) && mNeverResumed) + { + TLLM_CHECK_DEBUG(blockPageIsNull(bp)); + continue; + } if (scratchRange.contains(ordinal)) { @@ -2812,7 +3086,7 @@ bool KvCache::_checkSanity() const } else if (staleRange.beg <= ordinal && ordinal < staleRange.end) { - if (isCommitted || mCommitState != CommitState::ALLOWED) + if (isPrivateLifeCycle(lcs[lc]) || isCommitted || mCommitState != CommitState::ALLOWED) { TLLM_CHECK_DEBUG(blockPageIsNull(bp)); } @@ -2841,7 +3115,17 @@ bool KvCache::_checkSanity() const if (!blockPageIsNull(bp)) { auto page = blockPageGetPage(bp); - TLLM_CHECK_DEBUG(isCommitted == (dynamicPointerCast(page) != nullptr)); + auto const* attention = std::get_if(&lcs[lc]); + auto committed = dynamicPointerCast(page); + if (attention && attention->allowsCheckpoint() && committed) + { + TLLM_CHECK_DEBUG(isCommitted && committed->numTokensInBlock > 0 + && committed->numTokensInBlock <= mTokensPerBlock); + } + else + { + TLLM_CHECK_DEBUG((isCommitted && !isPrivate) == (committed != nullptr)); + } } } } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index 90fe45a5fc49..8444d3049f72 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -88,8 +88,8 @@ struct SeqBlock bool ret = treeBlock != nullptr; if (TLLM_UNLIKELY(gDebug)) { - // When committed: must have one canonical beam, and all non-null - // pages must be CommittedPage. + // Committed blocks have one canonical beam; publication concerns reusable pages. + // Request-private state stays uncommitted. if (ret) { TLLM_CHECK(pages.size() == BeamIndex{1}); @@ -98,7 +98,7 @@ struct SeqBlock if (!blockPageIsNull(bp)) { auto pg = blockPageGetPage(bp); - TLLM_CHECK(!pg || dynamicPointerCast(pg)); + TLLM_CHECK(!pg || dynamicPointerCast(pg) || pg->isPrivate()); } } else @@ -204,6 +204,25 @@ class PageStorageSnapshot std::vector mReadyEvents; }; +// Availability at the claimed prefix, indexed by LayerGroupId. OPTIONAL +// candidates are held until the first successful resume or close. +// Describes the initial prefix claim, not subsequent request progress. +struct ReuseGroupStatus +{ + // Lifecycle group returned by getLayerGroupId(layerId), not a physical pool index. + LayerGroupId groupId; + // REQUIRED constrains the prefix match; OPTIONAL is selected at resume; PRIVATE is never reused. + AttentionReusePolicy policy; + // Exclusive token endpoint of the prefix claimed by REQUIRED groups, shared by all groups. + int endpoint; + // All state needed by this group at endpoint is available; does not imply GPU residency or selection. + // False for an empty prefix, PRIVATE groups, or missing/discarded OPTIONAL candidates. + bool complete; + // Retained token ranges [begin, end), including block-rounded SWA windows and sinks. + // SSM uses (endpoint, endpoint) for its exact state; empty when complete is false. + std::vector> coverage; +}; + // --------------------------------------------------------------------------- // KvCache — manages the per-sequence KV cache state. // Mirrors Python's _KVCache. @@ -243,7 +262,8 @@ class KvCache : public std::enable_shared_from_this // isDecoding defaults to the current phase. A cache starts in prefill and cannot return to it // after decode admission. Set true when admitting a suspended request directly to decode. // Returns false if utilization too high or out of memory. - bool resume(std::optional stream = std::nullopt, std::optional isDecoding = std::nullopt); + bool resume(std::optional stream = std::nullopt, std::optional isDecoding = std::nullopt, + std::optional> optionalReuseGroups = std::nullopt); // Enter decode only after prefill has submitted its final KV accesses. Offloads complete sparse // history, deferring pages still needed on GPU by another owner. Retries deferred pages even @@ -405,6 +425,11 @@ class KvCache : public std::enable_shared_from_this //! survive, so only the tail stops counting as a local hit. void dropPartialBlockCachedTokenAttribution(); + std::vector const& reuseStatus() const noexcept + { + return mReuseStatus; + } + // Internal diagnostic: prefix supported by the attention pages alone, // before recurrent-state (SSM) snapshot pruning shortened the reuse. int numReusableTokensBeforeHybridPruning() const noexcept @@ -782,6 +807,14 @@ class KvCache : public std::enable_shared_from_this PendingStats mPendingStats; + std::vector mReuseStatus; + TypedVec>>> mOptionalCandidates; + std::vector mSelectedOptionalGroups; + void _collectReuseCandidates(); + void _selectOptionalGroups(std::vector const& groups); + // PRIVATE and unselected OPTIONAL groups need fresh pages for the claimed prefix. + bool _needsFreshPrefixPages(LifeCycleId lc) const; + // SWA scratch slot support. bool mEnableSwaScratchReuse = false; TypedVec> mScratchSlots; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 74e5c481c40d..07608ed8362d 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -26,6 +26,7 @@ #include "tensorrt_llm/common/logger.h" #include #include +#include #include #include #include @@ -163,11 +164,16 @@ void KvCacheManager::shutdown() // Post-condition: after clearing all reusable blocks and with no active KvCaches, // no evictable pages should remain. - for (auto const& lvl : mStorage->mLevels) + for (CacheLevel level{0}; level < mStorage->numCacheLevels(); ++level) { + auto const& lvl = mStorage->mLevels[level]; for (PoolGroupIndex pgIdx{0}; pgIdx < lvl.numPoolGroups(); ++pgIdx) { - TLLM_CHECK_DEBUG(lvl.controller.numEvictablePages(pgIdx) == 0); + auto const stats = mStorage->getStatistics(level, pgIdx); + TLLM_CHECK_WITH_INFO(stats.free == stats.total && stats.evictable == 0, + "KV shutdown retained slots at level %d, pool %d: total=%" PRId64 ", free=%" PRId64 + ", evictable=%" PRId64, + level.value(), pgIdx.value(), stats.total, stats.free, stats.evictable); } } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp index 9f4a4f0eec29..a4c57cacea07 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp @@ -41,7 +41,8 @@ LifeCycle makeLifeCycle(LayerConfig const& layer, int tokensPerBlock) else { bool const isSparse = !cfg.buffers.empty() && cfg.buffers.front().isSparse; - return AttnLifeCycle::make(cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse); + return AttnLifeCycle::make( + cfg.slidingWindowSize, cfg.numSinkTokens, tokensPerBlock, isSparse, cfg.reusePolicy); } }, layer); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h index c62d9b40d2c2..bdf550e6fcc1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h @@ -41,6 +41,22 @@ struct AttnLifeCycle std::optional windowSize; // nullopt = no sliding window int numSinkBlocks = 0; // divUp(numSinkTokens, tokensPerBlock) bool isSparse = false; + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; + + bool isPrivate() const noexcept + { + return reusePolicy == AttentionReusePolicy::PRIVATE; + } + + bool requiresPrefixCoverage() const noexcept + { + return reusePolicy == AttentionReusePolicy::REQUIRED; + } + + bool allowsCheckpoint() const noexcept + { + return reusePolicy == AttentionReusePolicy::OPTIONAL; + } HalfOpenRange getStaleRange(int historyLength, int tokensPerBlock) const { @@ -58,7 +74,8 @@ struct AttnLifeCycle bool operator==(AttnLifeCycle const& o) const noexcept { - return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse; + return windowSize == o.windowSize && numSinkBlocks == o.numSinkBlocks && isSparse == o.isSparse + && reusePolicy == o.reusePolicy; } bool operator<(AttnLifeCycle const& o) const noexcept @@ -71,18 +88,22 @@ struct AttnLifeCycle { return numSinkBlocks < o.numSinkBlocks; } - return isSparse < o.isSparse; + if (isSparse != o.isSparse) + { + return isSparse < o.isSparse; + } + return reusePolicy < o.reusePolicy; } - static AttnLifeCycle make( - std::optional ws, std::optional numSinkTokens, int tokensPerBlock, bool isSparse = false) + static AttnLifeCycle make(std::optional ws, std::optional numSinkTokens, int tokensPerBlock, + bool isSparse = false, AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED) { TLLM_CHECK_DEBUG(tokensPerBlock > 0); TLLM_CHECK_DEBUG(!ws.has_value() || *ws > 0); TLLM_CHECK_DEBUG(!numSinkTokens.has_value() || *numSinkTokens >= 0); TLLM_CHECK_DEBUG((!numSinkTokens.has_value() || *numSinkTokens == 0) || ws.has_value()); int sinkBlocks = divUp(numSinkTokens.value_or(0), tokensPerBlock); - return AttnLifeCycle{ws, sinkBlocks, isSparse}; + return AttnLifeCycle{ws, sinkBlocks, isSparse, reusePolicy}; } }; @@ -114,6 +135,18 @@ struct SsmLifeCycle // --------------------------------------------------------------------------- using LifeCycle = std::variant; +inline bool isPrivateLifeCycle(LifeCycle const& lc) +{ + auto const* attn = std::get_if(&lc); + return attn && attn->isPrivate(); +} + +inline bool requiresPrefixCoverage(LifeCycle const& lc) +{ + auto const* attn = std::get_if(&lc); + return !attn || attn->requiresPrefixCoverage(); +} + // Free function: compute stale range via std::visit. inline HalfOpenRange getStaleRange(LifeCycle const& lc, int historyLength, int tokensPerBlock) { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index 7ada3a2e97e7..2fb89b143ea6 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -42,6 +42,11 @@ Page::Page(StorageManager* mgr, LifeCycleId lc, CacheLevel level, Priority prio) { } +bool Page::isPrivate() const +{ + return !isCommitted() && !requiresPrefixCoverage(manager->lifeCycles()[lifeCycle]); +} + Page::~Page() { KVCM2_POISON_ON_EXCEPT( @@ -259,8 +264,11 @@ PageHolder::~PageHolder() // page with a larger recorded token count) is unreachable for reuse, so keeping it // in the eviction LRU would just pin a slot until memory pressure hits. auto* cp = static_cast(page.get()); - if (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp)) + if (page->scheduledForEviction() + && (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp))) + { manager->excludeFromEviction(*page); + } } else { @@ -531,6 +539,15 @@ void SharedPageLock::acquirePageIndex() void SharedPageLock::releasePageIndex() { + if (mUser.ordinal == kBadBlockOrdinal) + { + // SSM state has one live page per beam and no attention page-index + // entry. updateBasePageIndex deliberately returns BAD for this ordinal. + auto const& lc = page()->manager->lifeCycles().getLifeCycle(mUser.lifeCycle); + auto const* attention = std::get_if(&lc); + TLLM_CHECK_DEBUG(std::holds_alternative(lc) || (attention && attention->allowsCheckpoint())); + return; + } int oldBaseIndex = mUser.kvCache->updateBasePageIndex(mUser.beamIndex, mUser.ordinal, mUser.lifeCycle, kBadPageIndex.value()); // SSM pages use kBadBlockOrdinal, for which updateBasePageIndex returns kBadPageIndex. @@ -717,7 +734,9 @@ void ScratchSlotLock::unlock() { TLLM_CHECK_DEBUG(mSlot.hasValidSlot()); mSlot.readyEvent = mOwner->finishEvent(); - mOwner->storageManager()->releaseSlot(mLifeCycle, kHotLevel, std::move(mSlot)); + // Slot's implicit move leaves its optional ID engaged. Transfer ownership + // explicitly so destruction cannot release this physical slot a second time. + mOwner->storageManager()->releaseSlot(mLifeCycle, kHotLevel, detachSlot()); TLLM_CHECK_DEBUG(!mSlot.hasValidSlot()); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index c60e8eafc411..9ecd538fe079 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -62,6 +62,8 @@ class Page : public Slot, public EnableSharedFromThis virtual bool isCommitted() const = 0; + bool isPrivate() const; + PageStatus status() const noexcept; bool scheduledForEviction() const noexcept diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index bdb429ba9074..c9e3718b14b9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -206,15 +206,17 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) slotGroups.push_back(std::move(var)); } - // Keep sparse and dense lifecycles in separate GPU pools even when slot sizes match. + // Distinct sparsity and reuse policies need separate GPU pools and eviction queues. struct PoolGroupKey { + AttentionReusePolicy reusePolicy = AttentionReusePolicy::REQUIRED; std::vector slotSizes; bool isSparse = false; bool operator<(PoolGroupKey const& other) const { - return std::tie(slotSizes, isSparse) < std::tie(other.slotSizes, other.isSparse); + return std::tie(reusePolicy, slotSizes, isSparse) + < std::tie(other.reusePolicy, other.slotSizes, other.isSparse); } }; @@ -224,7 +226,8 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) auto const sizes = sg.slotSizeList(); auto const* attn = std::get_if(®istry.getLifeCycle(sg.lifeCycleId)); bool const isSparse = attn != nullptr && attn->isSparse; - poolGroups[{.slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); + auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; + poolGroups[{.reusePolicy = policy, .slotSizes = sizes.raw(), .isSparse = isSparse}].push_back(std::move(sg)); } StorageConfig out; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index 94ca946bba55..dc2a7519ffaa 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -17,6 +17,7 @@ #include "kv_cache_manager_v2/storageManager.h" #include "kv_cache_manager_v2/batchedPageCopy.h" +#include "kv_cache_manager_v2/blockRadixTree.h" #include "kv_cache_manager_v2/common.h" #include "kv_cache_manager_v2/copyEngine.h" #include "kv_cache_manager_v2/exceptions.h" @@ -397,12 +398,14 @@ StorageManager::StorageManager(LifeCycleRegistry const& lifeCycles, StorageConfi TypedVec coldGrouping(numLifeCycles()); TypedVec coldSlotDescList; - std::map coldGroupByPageBytes; + std::map, PoolGroupIndex> coldGroupByPageBytes; for (LifeCycleId lifeCycle{0}; lifeCycle < numLifeCycles(); ++lifeCycle) { size_t const coldPageBytes = coldPageBytesByLifeCycle[lifeCycle]; + auto const* attn = std::get_if(&mLifeCycles[lifeCycle]); + auto const policy = attn ? attn->reusePolicy : AttentionReusePolicy::REQUIRED; auto [it, inserted] = coldGroupByPageBytes.emplace( - coldPageBytes, PoolGroupIndex{static_cast(coldSlotDescList.size().value())}); + std::make_pair(policy, coldPageBytes), PoolGroupIndex{static_cast(coldSlotDescList.size().value())}); PoolGroupIndex const coldPgIdx = it->second; if (inserted) { @@ -744,7 +747,14 @@ bool StorageManager::isEvictable(Page const& page, std::optional lev { PageStatus s = page.status(); CacheLevel lvl = level.value_or(page.cacheLevel); - return (s == PageStatus::DROPPABLE && page.isCommitted()) || (s == PageStatus::HELD && lvl < numCacheLevels() - 1); + if (s == PageStatus::DROPPABLE && page.isCommitted()) + { + // A migration batch can outlive checkpoint invalidation. Do not put + // detached bundle members back in an LRU that would retain their slots. + auto const& committed = static_cast(page); + return committed.block && !committed.block->isOrphan() && committed.block->holdsPage(committed); + } + return s == PageStatus::HELD && lvl < numCacheLevels() - 1; } // --------------------------------------------------------------------------- @@ -1079,10 +1089,7 @@ void StorageManager::_prepareFreeSlots(TypedVecstatus() != PageStatus::HELD) - { - lvl.controller.scheduleForEviction(*page); - } + scheduleForEviction(*page); } } } @@ -1943,7 +1950,16 @@ TypedVec StorageManager::projectPoolGroupRatio( TypedVec StorageManager::ratioFromBatch(BatchDesc const& batch, int tokensPerBlock, std::optional const& swaScratchReuse, size_t granularity) const { - auto const numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); + auto numSlots = computeSlotsForBatch(batch, tokensPerBlock, swaScratchReuse); + // A soft preference, not a feasibility floor: keep one optional endpoint + // per typical request in addition to its private working pages. + for (auto const& [lc, attention] : mLifeCycles.attentionLifeCycles()) + { + if (attention->allowsCheckpoint()) + { + numSlots[lc] += slotCountValueFromSize(batch.kvCaches.size()); + } + } return normalizeToRatio(slotsToBytes(numSlots, granularity)); } diff --git a/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt b/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt index bedc7cad6f62..1660f43fa1cb 100644 --- a/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt +++ b/cpp/tensorrt_llm/deep_gemm/CMakeLists.txt @@ -12,10 +12,11 @@ set(DEEP_GEMM_SOURCE_DIR ${CMAKE_BINARY_DIR}/_deps/deepgemm-src) get_filename_component(DEEP_GEMM_SOURCE_DIR ${DEEP_GEMM_SOURCE_DIR} ABSOLUTE) # Check if submodules are initialized -if(NOT EXISTS ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include) +if(NOT EXISTS ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include + OR NOT EXISTS ${DEEP_GEMM_SOURCE_DIR}/third-party/deep_jit/include) message( FATAL_ERROR - "DeepGEMM submodules not initialized. Something went wrong with FetchContent.\n" + "DeepGEMM submodules (cutlass, deep_jit) not initialized. Something went wrong with FetchContent.\n" ) endif() @@ -83,7 +84,8 @@ foreach(SOURCE_FILE ${DEEP_GEMM_ALL_FILES}) PROPERTY CMAKE_CONFIGURE_DEPENDS ${SOURCE_FILE}) endforeach() -# Copy third-party includes (cutlass and fmt) to the include directory +# Copy third-party includes (cutlass) to the include directory. DeepJIT is +# header-only and host-side only, so it is not needed at kernel JIT time. set(DEEP_GEMM_INCLUDE_DEST ${DEEP_GEMM_PYTHON_DEST}/include) file(MAKE_DIRECTORY ${DEEP_GEMM_INCLUDE_DEST}) file(COPY ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include/cute @@ -119,9 +121,25 @@ target_compile_definitions(deep_gemm_cpp_tllm # Include directories matching deep_gemm setup.py target_include_directories( deep_gemm_cpp_tllm - PRIVATE ${CUDA_INCLUDE_DIRS} ${DEEP_GEMM_SOURCE_DIR}/deep_gemm/include - ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include - ${DEEP_GEMM_SOURCE_DIR}/third-party/fmt/include) + PRIVATE ${CUDA_INCLUDE_DIRS} + ${CUDA_TOOLKIT_ROOT_DIR}/include/cccl + ${DEEP_GEMM_SOURCE_DIR}/deep_gemm/include + ${DEEP_GEMM_SOURCE_DIR}/third-party/deep_jit/include + ${DEEP_GEMM_SOURCE_DIR}/third-party/cutlass/include) + +# DeepJIT resolves libdwfl through dlopen at runtime but includes the elfutils +# header for its types. Hosts without the elfutils development headers use the +# minimal declaration shim under deep_jit_shim/ instead. +include(CheckIncludeFileCXX) +check_include_file_cxx("elfutils/libdwfl.h" DEEP_GEMM_HAVE_LIBDWFL_HEADER) +if(NOT DEEP_GEMM_HAVE_LIBDWFL_HEADER) + message( + STATUS + "elfutils/libdwfl.h not found; DeepJIT builds against the libdwfl declaration shim" + ) + target_include_directories(deep_gemm_cpp_tllm + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/deep_jit_shim) +endif() # Link libraries (matching deep_gemm setup.py: cuda, cudart + torch) target_link_libraries( diff --git a/cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h b/cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h new file mode 100644 index 000000000000..5058a21b29d7 --- /dev/null +++ b/cpp/tensorrt_llm/deep_gemm/deep_jit_shim/elfutils/libdwfl.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Minimal declarations of the libdwfl (elfutils) entry points that DeepJIT's + * exception backtraces resolve through dlopen("libdw.so.1"). DeepJIT only + * includes for these types and prototypes and never links + * libdw, so build hosts without the elfutils development headers use this + * shim instead; when the real header is installed it takes precedence. + */ + +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" +{ +#endif + + typedef uint64_t Dwarf_Addr; + typedef uint64_t Dwarf_Word; + typedef uint64_t GElf_Addr; + typedef uint32_t GElf_Word; + + typedef struct Elf Elf; + typedef struct Dwfl Dwfl; + typedef struct Dwfl_Module Dwfl_Module; + typedef struct Dwfl_Line Dwfl_Line; + typedef struct Dwfl_Shim_Elf64_Shdr GElf_Shdr; + + typedef struct + { + int (*find_elf)( + Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, char** file_name, Elf** elfp); + int (*find_debuginfo)(Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, + char const* file_name, char const* debuglink_file, GElf_Word debuglink_crc, char** debuginfo_file_name); + int (*section_address)(Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, + char const* secname, GElf_Word shndx, GElf_Shdr const* shdr, Dwarf_Addr* addr); + char** debuginfo_path; + } Dwfl_Callbacks; + + Dwfl* dwfl_begin(Dwfl_Callbacks const* callbacks); + int dwfl_linux_proc_report(Dwfl* dwfl, pid_t pid); + int dwfl_report_end(Dwfl* dwfl, + int (*removed)(Dwfl_Module* mod, void* userdata, char const* name, Dwarf_Addr low_addr, void* arg), void* arg); + Dwfl_Module* dwfl_addrmodule(Dwfl* dwfl, Dwarf_Addr address); + char const* dwfl_module_addrname(Dwfl_Module* mod, GElf_Addr address); + Dwfl_Line* dwfl_module_getsrc(Dwfl_Module* mod, Dwarf_Addr address); + char const* dwfl_lineinfo( + Dwfl_Line* line, Dwarf_Addr* addr, int* linep, int* colp, Dwarf_Word* mtime, Dwarf_Word* length); + int dwfl_linux_proc_find_elf( + Dwfl_Module* mod, void** userdata, char const* modname, Dwarf_Addr base, char** file_name, Elf** elfp); + +#ifdef __cplusplus +} +#endif diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index 9379dbad7310..001772ed74eb 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -205,9 +205,10 @@ struct VecType<16> // Cache scale / pre-store quantization type for postProcessScatterKernel. // The store dtype is implied: kNone keeps the input dtype (elem_bytes -// controls bf16 vs fp32), kFP8* writes one byte per element, and -// kMXFP4Blockwise writes packed FP4 (two values per byte) plus per-32 -// UE8M0 scale bytes. +// controls bf16 vs fp32), kFP8* writes one byte per element, and the FP4 +// types write packed FP4 (two values per byte) plus one scale byte per +// block -- per-32 UE8M0 for kMXFP4Blockwise, per-16 E4M3 for +// kNVFP4Blockwise. enum class CacheScaleType { kNone = 0, @@ -308,8 +309,13 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo { using KvScoreElemT = typename std::conditional::type; using StateElemT = typename std::conditional::type; - // DeepSeek-V4 model configures compress_ratio to be 4 or 128. - static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + // DeepSeek-V4 configures compress_ratio to 4 or 128; DeepSeek-V4.1 uses 2. + // (V4.1's other ratio, 1, needs no reduction at all and never reaches here.) + static_assert(COMPRESS_RATIO == 2 || COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + // Overlapping pooling windows are a V4 ratio-4 property, not a property of + // small ratios: V4.1's ratio-2 compressor keeps one window in flight and + // sizes its state (max_batch, ratio, head_dim) with no overlap half + // (model.py:451-456), so CR=2 takes the non-overlap path like CR=128. constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); constexpr int ELEM_BYTES_FOR_VEC = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; @@ -663,9 +669,9 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // HD — HEAD_DIM in {128, 512} // KV_EB — kv_score element bytes in {2 (bf16), 4 (fp32)} // STATE_EB — paged state element bytes in {2 (bf16), 4 (fp32)} -// CR — COMPRESS_RATIO in {4, 128} +// CR — COMPRESS_RATIO in {2 (V4.1), 4, 128} // NRW — NUM_RED_WARPS — 4 when CR=128 (multi-warp Phase 3 reduction -// hides DRAM latency for the heavier R=128 chunk); 1 when CR=4. +// hides DRAM latency for the heavier R=128 chunk); 1 when CR=2 or 4. // // Multi-warp SMEM budget (per block): 3 * NRW * ELEM_PER_BLOCK * sizeof(float). // HD=128: ELEM_PER_BLOCK=128 → 6 KB @@ -680,6 +686,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // Master list. Order does not matter; the dispatcher walks linearly. // clang-format off #define FOREACH_DECODE_CONFIG(F) \ + /* CR=2 (DeepSeek-V4.1): single-warp, same reason as CR=4 and half the loop. */ \ + FOREACH_DECODE_DTYPE(F, 128, 2, 1) FOREACH_DECODE_DTYPE(F, 512, 2, 1) \ /* CR=4: single-warp only (small reduction; multi-warp would over-subscribe). */ \ FOREACH_DECODE_DTYPE(F, 128, 4, 1) FOREACH_DECODE_DTYPE(F, 512, 4, 1) \ /* CR=128: multi-warp Phase 3 reduction for every supported next_n. */ \ @@ -722,8 +730,8 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, int32_t const* new_tokens_per_seq, cudaStream_t stream) { - TLLM_CHECK_WITH_INFO( - compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128"); + TLLM_CHECK_WITH_INFO(compress_ratio == 2 || compress_ratio == 4 || compress_ratio == 128, + "pagedKvCompressLaunch only supports compress_ratio 2 (V4.1), 4 or 128, got %d", compress_ratio); TLLM_CHECK_WITH_INFO( (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), "pagedKvCompressLaunch only supports bf16/fp32 kv_score and paged state"); @@ -873,7 +881,8 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl { using KvScoreElemT = typename std::conditional::type; using StateElemT = typename std::conditional::type; - static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + // See pagedKvCompressKernel: V4.1 adds CR=2, which is non-overlapping. + static_assert(COMPRESS_RATIO == 2 || COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); constexpr int ELEM_BYTES_FOR_VEC @@ -1228,8 +1237,10 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl INST_PREFILL(HD, 2, 2, CR, NRW) \ INST_PREFILL(HD, 2, 4, CR, NRW) INST_PREFILL(HD, 4, 2, CR, NRW) INST_PREFILL(HD, 4, 4, CR, NRW) +INST_PREFILL_DTYPES(128, 2, 1) INST_PREFILL_DTYPES(128, 4, 1) INST_PREFILL_DTYPES(128, 128, 4) +INST_PREFILL_DTYPES(512, 2, 1) INST_PREFILL_DTYPES(512, 4, 1) INST_PREFILL_DTYPES(512, 128, 4) #undef INST_PREFILL_DTYPES @@ -1256,8 +1267,8 @@ void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_ int out_elem_bytes, cudaStream_t stream) { bool const overlap = (compress_ratio == 4); - TLLM_CHECK_WITH_INFO( - compress_ratio == 4 || compress_ratio == 128, "prefillReductionLaunch only supports compress_ratio 4 or 128"); + TLLM_CHECK_WITH_INFO(compress_ratio == 2 || compress_ratio == 4 || compress_ratio == 128, + "prefillReductionLaunch only supports compress_ratio 2 (V4.1), 4 or 128, got %d", compress_ratio); TLLM_CHECK_WITH_INFO( (kv_score_elem_bytes == 2 || kv_score_elem_bytes == 4) && (state_elem_bytes == 2 || state_elem_bytes == 4), "prefillReductionLaunch only supports bf16/fp32 kv_score and paged state"); @@ -1304,14 +1315,18 @@ void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_ if (head_dim == 512) { - if (compress_ratio == 4) + if (compress_ratio == 2) + DISPATCH_PREFILL_DTYPE(512, 2, 1); + else if (compress_ratio == 4) DISPATCH_PREFILL_DTYPE(512, 4, 1); else DISPATCH_PREFILL_DTYPE(512, 128, 4); } else { - if (compress_ratio == 4) + if (compress_ratio == 2) + DISPATCH_PREFILL_DTYPE(128, 2, 1); + else if (compress_ratio == 4) DISPATCH_PREFILL_DTYPE(128, 4, 1); else DISPATCH_PREFILL_DTYPE(128, 128, 4); @@ -1340,6 +1355,8 @@ void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_ // - kFP8Blockwise: one fp8 byte per value + one fp32 scale per 128 values. // - kMXFP4Blockwise: packed fp4 (two values per byte) + one ue8m0 byte // per 32 values. +// - kNVFP4Blockwise: packed fp4 (two values per byte) + one e4m3 byte +// per 16 values (V4.1 tech report §2.4.4). // // Pipeline (10 steps, all in fp32 registers): // 1. Vectorized load compressed token from kv_comp diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu index 47d7640e7e62..e5a307f25074 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.cu @@ -168,6 +168,13 @@ void dispatchDeepseekV4QNorm( // (kPairsPerLane-1) iterations cover the nope range and the last iteration // covers the rope range exactly. +// `applyNorm`: false makes the RMS scale an identity, for models that have no +// per-head query norm at all (V4.1 deleted V4's). The norm here carries no learned +// weight -- it is one per-row scalar -- so skipping it is exactly `normScale = 1`, +// and the warp reduction that produced it goes with it. Runtime rather than a fifth +// template parameter: the branch is warp-uniform, and the instantiation set is +// already 8-way (2 dtypes x kFuseRope x kWideVec). +// // `kFuseRope`: rotate the rope segment here too and store it FP8 in the same row. // Register-local -- `kRopePairs == kWarpSize`, so lane `l` owns rope pair `l` and // needs only one float2 of cos/sin. @@ -180,7 +187,7 @@ void dispatchDeepseekV4QNorm( template __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8_e4m3* __restrict__ quant_q_nope, T* __restrict__ q_pe_out, float const* __restrict__ quant_scale_qkv_ptr, int totalRows, - int quantQNopeRowStrideBytes, float eps, float2 const* __restrict__ cos_sin_cache = nullptr, + int quantQNopeRowStrideBytes, float eps, bool applyNorm, float2 const* __restrict__ cos_sin_cache = nullptr, int const* __restrict__ cache_seq_lens = nullptr, int num_heads = 0, int seq_len = 0, int const* __restrict__ cu_q_seqlens = nullptr, int num_seqs = 0, int num_heads_shift = -1, int seq_len_shift = -1) { @@ -272,8 +279,15 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 } } - sumSquares = warpReduceSum(sumSquares); - float const normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); + // The per-lane squares above stay unconditional: they ride along with loads this + // path needs anyway, and branching inside the unrolled loop would cost the norming + // model. The cross-lane reduction is what is worth skipping. + float normScale = 1.0F; + if (applyNorm) + { + sumSquares = warpReduceSum(sumSquares); + normScale = rsqrtf(sumSquares / static_cast(kHeadDim) + eps); + } float const fp8Scale = normScale * quantScale; // Position depends on the token, which every lane of the warp shares, so this is @@ -410,8 +424,8 @@ __global__ void deepseekV4QNormFusedKernel(T const* __restrict__ input, __nv_fp8 template void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q_pe_out, void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, float eps, - void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, int const* cu_q_seqlens, - int num_seqs, cudaStream_t stream) + bool applyNorm, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, + int const* cu_q_seqlens, int num_seqs, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(headDim == 512, "deepseekV4QNormFused only supports head_dim=512, got %d", headDim); TLLM_CHECK_WITH_INFO(nopeDim == 448, "deepseekV4QNormFused only supports nope_dim=448, got %d", nopeDim); @@ -457,7 +471,7 @@ void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q tensorrt_llm::common::launchWithPdlWhenEnabled("deepseekV4QNormFused", deepseekV4QNormFusedKernel, grid, block, 0, stream, static_cast(input), static_cast<__nv_fp8_e4m3*>(quant_q_nope), static_cast(q_pe_out), - static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, + static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, applyNorm, static_cast(cos_sin_cache), cache_seq_lens, num_heads, seq_len, cu_q_seqlens, num_seqs, num_heads_shift, seq_len_shift); } @@ -466,7 +480,7 @@ void dispatchDeepseekV4QNormFused(void const* input, void* quant_q_nope, void* q tensorrt_llm::common::launchWithPdlWhenEnabled("deepseekV4QNormFused", deepseekV4QNormFusedKernel, grid, block, 0, stream, static_cast(input), static_cast<__nv_fp8_e4m3*>(quant_q_nope), static_cast(q_pe_out), - static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, + static_cast(quant_scale_qkv_ptr), totalRows, quantQNopeRowStrideBytes, eps, applyNorm, static_cast(nullptr), static_cast(nullptr), 0, 0, static_cast(nullptr), 0, -1, -1); } @@ -504,8 +518,8 @@ void invokeDeepseekV4QNorm( void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* q_pe_out, void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, - bool isBfloat16, float eps, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, - int const* cu_q_seqlens, int num_seqs, cudaStream_t stream) + bool isBfloat16, float eps, bool applyNorm, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, + int seq_len, int const* cu_q_seqlens, int num_seqs, cudaStream_t stream) { if (totalRows == 0) { @@ -515,14 +529,14 @@ void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* if (isBfloat16) { dispatchDeepseekV4QNormFused<__nv_bfloat16>(input, quant_q_nope, q_pe_out, quant_scale_qkv_ptr, totalRows, - headDim, nopeDim, quantQNopeRowStrideBytes, eps, cos_sin_cache, cache_seq_lens, num_heads, seq_len, - cu_q_seqlens, num_seqs, stream); + headDim, nopeDim, quantQNopeRowStrideBytes, eps, applyNorm, cos_sin_cache, cache_seq_lens, num_heads, + seq_len, cu_q_seqlens, num_seqs, stream); } else { dispatchDeepseekV4QNormFused(input, quant_q_nope, q_pe_out, quant_scale_qkv_ptr, totalRows, headDim, - nopeDim, quantQNopeRowStrideBytes, eps, cos_sin_cache, cache_seq_lens, num_heads, seq_len, cu_q_seqlens, - num_seqs, stream); + nopeDim, quantQNopeRowStrideBytes, eps, applyNorm, cos_sin_cache, cache_seq_lens, num_heads, seq_len, + cu_q_seqlens, num_seqs, stream); } } diff --git a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h index a09b68f6f6cc..82fa43dc30c5 100644 --- a/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h +++ b/cpp/tensorrt_llm/kernels/deepseekV4QNormKernel.h @@ -50,8 +50,8 @@ void invokeDeepseekV4QNorm( // the plain no-RoPE behaviour. void invokeDeepseekV4QNormFusedFp8(void const* input, void* quant_q_nope, void* q_pe_out, void const* quant_scale_qkv_ptr, int totalRows, int headDim, int nopeDim, int quantQNopeRowStrideBytes, - bool isBfloat16, float eps, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, int seq_len, - int const* cu_q_seqlens, int num_seqs, cudaStream_t stream); + bool isBfloat16, float eps, bool applyNorm, void const* cos_sin_cache, int const* cache_seq_lens, int num_heads, + int seq_len, int const* cu_q_seqlens, int num_seqs, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index b997621a8fcc..28fc74cfe6a1 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -802,7 +802,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 template + uint32_t kNumPmapThreads, uint32_t kNumSplits = 1, bool kFuseNorm = false, bool kLaggedPre = false> __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_allinone_tf32_pmap_gemm_atomic_impl(const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, // residual_prev, bf16 @@ -825,7 +825,11 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) // norm_weight must be bf16 [HIDDEN]; norm_eps is the RMSNorm epsilon. // When kFuseNorm is false, norm_weight/norm_eps are ignored. __nv_bfloat16 const* __restrict__ norm_weight, float norm_eps, float rms_eps, float hc_pre_eps, - float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat) + float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat, + // kLaggedPre (DeepSeek-V4.1 wiring): layer_input is collapsed with the previous sublayer's + // gate pre_mix_ext [M, HC_MULT] and this token's own gate is written to pre_mix_out [M, HC_MULT]. + // Both are ignored when kLaggedPre is false (V4). + float const* __restrict__ pre_mix_ext = nullptr, float* __restrict__ pre_mix_out = nullptr) { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; @@ -1417,6 +1421,11 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float v = y_row[lane_bf] * rstd * s0 + hc_base[lane_bf]; pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; + if constexpr (kLaggedPre) + { + if (warp_in_team == 0) + pre_mix_out[static_cast(tok) * HC_MULT + lane_bf] = pre_mix_local; + } if (warp_in_team == 0) { @@ -1477,7 +1486,12 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float pm[HC_MULT]; #pragma unroll for (uint32_t j = 0; j < HC_MULT; ++j) - pm[j] = __shfl_sync(0xffffffff, pre_mix_local, j); + { + if constexpr (kLaggedPre) + pm[j] = pre_mix_ext[static_cast(tok) * HC_MULT + j]; + else + pm[j] = __shfl_sync(0xffffffff, pre_mix_local, j); + } // Layer_input[tok, h] = sum_j pm[j] * residual_cur[tok, j, h]. // When WARPS_PER_TOK>1, warp_in_team 0..WARPS_PER_TOK-1 together cover diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index c098a4203b0e..012ecd2b01e4 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -104,6 +104,7 @@ static constexpr uint32_t FHC_BLOCK_K = 64; #ifdef TRTLLM_MHC_ENABLE_FUSED_HC static constexpr uint32_t FHC_HIDDEN_FLASH = 4096; static constexpr uint32_t FHC_HIDDEN_PRO = 7168; +static constexpr uint32_t FHC_HIDDEN_V41 = 5120; // DeepSeek-V4.1-Flash static constexpr uint32_t FHC_BLOCK_M = 64; static constexpr uint32_t FHC_BLOCK_N = 32; static constexpr uint32_t FHC_SWIZZLE_CD = 128; @@ -120,12 +121,13 @@ static constexpr uint32_t FHC_NUM_PMAP_TH = 128; template