From 58597602697ea5e748fcfffdbaa1bbef30cc7648 Mon Sep 17 00:00:00 2001 From: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Date: Thu, 30 Apr 2026 05:51:41 +0000 Subject: [PATCH 01/58] [None][feat] Add DeepSeek-V4 model support Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit 344b9e93090aa6f4359a88ef82d81799525c110c) Signed-off-by: Fanrong Li --- 3rdparty/CMakeLists.txt | 9 +- cpp/tensorrt_llm/CMakeLists.txt | 2 + cpp/tensorrt_llm/kernels/CMakeLists.txt | 6 + cpp/tensorrt_llm/kernels/IndexerTopK.h | 65 +- .../kernels/compressorKernels/CMakeLists.txt | 25 + .../compressorKernels/compressorKernels.cu | 1792 ++++++++++++ .../compressorKernels/compressorKernels.h | 99 + .../kernels/customMoeRoutingKernels.cu | 172 ++ .../kernels/customMoeRoutingKernels.h | 12 +- cpp/tensorrt_llm/kernels/indexerTopK.cu | 848 ++---- .../kernels/mhcKernels/CMakeLists.txt | 37 + .../mhcKernels/fused_tf32_pmap_gemm.cuh | 1331 +++++++++ .../kernels/mhcKernels/mhcFusedHcKernel.cu | 537 ++++ .../kernels/mhcKernels/mhcKernels.cu | 772 +++++ .../kernels/mhcKernels/mhcKernels.h | 139 + .../kernels/mhcKernels/mhc_fused_fma.cuh | 898 ++++++ .../routing/RoutingKernelTopK.cuh | 93 + .../trtllmGenKernels/blockScaleMoe/runner.h | 5 +- cpp/tensorrt_llm/thop/CMakeLists.txt | 6 +- cpp/tensorrt_llm/thop/IndexerTopKOp.cpp | 94 +- cpp/tensorrt_llm/thop/compressorOp.cpp | 166 ++ cpp/tensorrt_llm/thop/mhcOp.cpp | 208 ++ cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp | 92 + cpp/tensorrt_llm/thop/moeGateOp.cpp | 76 + docs/source/models/supported-models.md | 3 + examples/llm-api/quickstart_advanced.py | 2 + examples/models/core/deepseek_v4/README.md | 488 ++++ jenkins/L0_MergeRequest.groovy | 18 +- pyproject.toml | 2 +- .../sparse/deepseek_v4/__init__.py | 4 + .../sparse/deepseek_v4/cache_manager.py | 830 ++++++ .../sparse/deepseek_v4/compressor.py | 334 +++ .../sparse/deepseek_v4/deepseek_v4.py | 1209 ++++++++ .../_torch/attention_backend/sparse/dsa.py | 1694 ++++++----- .../_torch/attention_backend/sparse/kernel.py | 263 ++ .../_torch/attention_backend/sparse/utils.py | 5 + .../_torch/attention_backend/trtllm.py | 28 +- tensorrt_llm/_torch/configs/__init__.py | 4 +- tensorrt_llm/_torch/configs/deepseekv4.py | 177 ++ .../_torch/custom_ops/cpp_custom_ops.py | 18 +- .../_torch/custom_ops/torch_custom_ops.py | 6 +- tensorrt_llm/_torch/memory_buffer_utils.py | 21 + tensorrt_llm/_torch/model_config.py | 213 +- tensorrt_llm/_torch/models/__init__.py | 2 + .../_torch/models/modeling_deepseekv4.py | 2521 ++++++++++++++++ .../_torch/models/modeling_speculative.py | 9 + tensorrt_llm/_torch/models/modeling_utils.py | 20 +- tensorrt_llm/_torch/modules/attention.py | 745 +++-- .../_torch/modules/engram/__init__.py | 18 + tensorrt_llm/_torch/modules/engram/engram.py | 846 ++++++ .../_torch/modules/fused_moe/__init__.py | 5 +- .../modules/fused_moe/configurable_moe.py | 2 + .../modules/fused_moe/fused_moe_cute_dsl.py | 3 +- .../modules/fused_moe/fused_moe_cutlass.py | 25 +- .../modules/fused_moe/fused_moe_deepgemm.py | 24 +- .../modules/fused_moe/fused_moe_triton.py | 25 +- .../modules/fused_moe/fused_moe_trtllm_gen.py | 32 +- .../modules/fused_moe/fused_moe_vanilla.py | 3 +- .../modules/fused_moe/fused_moe_wide_ep.py | 17 +- .../_torch/modules/fused_moe/routing.py | 184 +- tensorrt_llm/_torch/modules/gated_mlp.py | 9 +- tensorrt_llm/_torch/modules/linear.py | 1 - tensorrt_llm/_torch/modules/mhc/__init__.py | 0 .../_torch/modules/mhc/hyper_connection.py | 280 ++ tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 1085 +++++++ .../_torch/modules/rotary_embedding.py | 17 +- tensorrt_llm/_torch/modules/swiglu.py | 19 +- tensorrt_llm/_torch/pyexecutor/_util.py | 14 +- .../_torch/pyexecutor/config_utils.py | 1 + .../_torch/pyexecutor/model_engine.py | 10 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 39 +- tensorrt_llm/_torch/speculative/utils.py | 14 +- tensorrt_llm/_torch/utils.py | 1 + tensorrt_llm/_utils.py | 3 + tensorrt_llm/bench/build/dataclasses.py | 5 + .../bench/dataclasses/configuration.py | 3 +- tensorrt_llm/inputs/utils.py | 5 +- tensorrt_llm/llmapi/__init__.py | 5 +- tensorrt_llm/llmapi/llm_args.py | 42 +- tensorrt_llm/llmapi/tokenizer.py | 2 + .../tokenizer/deepseek_v4/__init__.py | 18 + .../tokenizer/deepseek_v4/tokenizer.py | 97 + tensorrt_llm/tokenizer/tokenizer.py | 1 + .../tools/layer_wise_benchmarks/runner.py | 4 +- .../test_lists/test-db/l0_b200.yml | 3 + .../test_lists/test-db/l0_b300.yml | 2 + .../_torch/attention/sparse/__init__.py | 0 .../attention/sparse/deepseek_v4/__init__.py | 0 .../deepseek_v4/test_compressor_kernel.py | 2549 +++++++++++++++++ .../deepseek_v4/test_compressor_module.py | 2457 ++++++++++++++++ .../deepseek_v4/test_compressor_tf32.py | 108 + .../test_deepseek_v4_cache_manager.py | 891 ++++++ .../test_deepseek_v4_indices_transform.py | 540 ++++ .../deepseek_v4/test_deepseek_v4_o_proj.py | 293 ++ .../test_deepseek_v4_sparse_mla.py | 1464 ++++++++++ .../_torch/attention/sparse/dsa/__init__.py | 0 .../sparse/{ => dsa}/test_dsa_indexer.py | 339 ++- .../sparse/dsa/test_dsa_sparse_mla.py | 1009 +++++++ .../sparse/{ => dsa}/test_short_seq_mha.py | 2 +- .../attention/sparse/kernel/__init__.py | 0 .../sparse/{ => kernel}/test_flash_mla.py | 0 .../test_triton.py} | 225 ++ .../attention/sparse/rocketkv/__init__.py | 0 .../sparse/{ => rocketkv}/test_rocketkv.py | 4 +- .../attention/sparse/test_dsa_fp4_indexer.py | 2 +- .../sparse/test_sparse_mla_forward.py | 1602 +++++++++-- .../attention/sparse/test_triton_topk.py | 228 -- .../modeling/test_modeling_deepseekv4.py | 765 +++++ tests/unittest/_torch/modules/test_engram.py | 588 ++++ tests/unittest/_torch/modules/test_mhc.py | 1029 +++++++ .../_torch/modules/test_rotary_embedding.py | 258 +- .../_torch/thop/parallel/test_indexer_topk.py | 23 +- .../_torch/thop/serial/test_moe_gate.py | 350 +++ .../api_stability/references/llm.yaml | 2 +- .../llmapi/test_deepseek_v4_tokenizer.py | 97 + 115 files changed, 31449 insertions(+), 2340 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt create mode 100644 cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu create mode 100644 cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh create mode 100644 cpp/tensorrt_llm/thop/compressorOp.cpp create mode 100644 cpp/tensorrt_llm/thop/mhcOp.cpp create mode 100644 cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/moeGateOp.cpp create mode 100644 examples/models/core/deepseek_v4/README.md create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py create mode 100644 tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py create mode 100644 tensorrt_llm/_torch/configs/deepseekv4.py create mode 100644 tensorrt_llm/_torch/models/modeling_deepseekv4.py create mode 100644 tensorrt_llm/_torch/modules/engram/__init__.py create mode 100644 tensorrt_llm/_torch/modules/engram/engram.py create mode 100644 tensorrt_llm/_torch/modules/mhc/__init__.py create mode 100644 tensorrt_llm/_torch/modules/mhc/hyper_connection.py create mode 100644 tensorrt_llm/_torch/modules/mhc/mhc_cuda.py create mode 100644 tensorrt_llm/tokenizer/deepseek_v4/__init__.py create mode 100644 tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py create mode 100644 tests/unittest/_torch/attention/sparse/__init__.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/__init__.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py create mode 100644 tests/unittest/_torch/attention/sparse/dsa/__init__.py rename tests/unittest/_torch/attention/sparse/{ => dsa}/test_dsa_indexer.py (91%) create mode 100644 tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py rename tests/unittest/_torch/attention/sparse/{ => dsa}/test_short_seq_mha.py (99%) create mode 100644 tests/unittest/_torch/attention/sparse/kernel/__init__.py rename tests/unittest/_torch/attention/sparse/{ => kernel}/test_flash_mla.py (100%) rename tests/unittest/_torch/attention/sparse/{test_triton_bmm.py => kernel/test_triton.py} (66%) create mode 100644 tests/unittest/_torch/attention/sparse/rocketkv/__init__.py rename tests/unittest/_torch/attention/sparse/{ => rocketkv}/test_rocketkv.py (99%) delete mode 100644 tests/unittest/_torch/attention/sparse/test_triton_topk.py create mode 100644 tests/unittest/_torch/modeling/test_modeling_deepseekv4.py create mode 100644 tests/unittest/_torch/modules/test_engram.py create mode 100644 tests/unittest/_torch/modules/test_mhc.py create mode 100644 tests/unittest/_torch/thop/serial/test_moe_gate.py create mode 100644 tests/unittest/llmapi/test_deepseek_v4_tokenizer.py diff --git a/3rdparty/CMakeLists.txt b/3rdparty/CMakeLists.txt index b9982a166b52..27033824920b 100644 --- a/3rdparty/CMakeLists.txt +++ b/3rdparty/CMakeLists.txt @@ -129,13 +129,20 @@ foreach(DEP_IDX RANGE ${DEP_COUNT_MINUS_ONE}) if(DEP_PATCH_FILE AND NOT DEP_PATCH_FILE STREQUAL "") set(_patch_file "${CMAKE_CURRENT_SOURCE_DIR}/${DEP_PATCH_FILE}") + # Use a two-step patch command so incremental builds succeed when the patch + # is already applied: first try to apply; if that fails, verify the patch + # was previously applied (reverse dry-run) — only fail if neither succeeds. list( APPEND FETCH_ARGS PATCH_COMMAND + ${CMAKE_COMMAND} + -E + env + -- bash -c - "patch -p1 --forward --batch --dry-run -i '${_patch_file}' && patch -p1 --forward --batch -i '${_patch_file}' || echo 'Patch already applied, skipping.'" + "patch -p1 --forward --batch -i '${_patch_file}' || patch -p1 --reverse --dry-run --batch -i '${_patch_file}'" ) endif() diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 06aefb3886ba..91beff7cf5bd 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -200,6 +200,8 @@ set(TRTLLM_LINK_LIBS layers_src runtime_src testing_src + mhcKernels_src + compressorKernels_src userbuffers_src ${DECODER_SHARED_TARGET_0} ${DECODER_SHARED_TARGET_1}) diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index e0e498fa89eb..0c6f4f234dec 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -30,6 +30,8 @@ add_subdirectory(dsv3MinLatencyKernels) add_subdirectory(causalConv1d) add_subdirectory(fusedGatedRMSNormQuant) add_subdirectory(mamba2MTPSSMCache) +add_subdirectory(mhcKernels) +add_subdirectory(compressorKernels) file(GLOB_RECURSE SRC_CPP *.cpp) file(GLOB_RECURSE SRC_CU *.cu) @@ -55,6 +57,10 @@ list(FILTER SRC_CU EXCLUDE REGEX "userbuffers/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedLayernormKernels/.*") list(FILTER SRC_CU EXCLUDE REGEX "fusedGatedRMSNormQuant/.*") list(FILTER SRC_CU EXCLUDE REGEX "mamba2MTPSSMCache/.*") +list(FILTER SRC_CPP EXCLUDE REGEX "mhcKernels/.*") +list(FILTER SRC_CU EXCLUDE REGEX "mhcKernels/.*") +list(FILTER SRC_CPP EXCLUDE REGEX "compressorKernels/.*") +list(FILTER SRC_CU EXCLUDE REGEX "compressorKernels/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 324e597dc779..30fa8e7a5ee4 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -27,47 +27,56 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -/// Indexer TopK decode. Three tiers: -/// - GVR Heuristic (preIdx provided, K in {512,1024,2048}, numColumns in -/// [kSeqSmall, splitWorkThreshold), numRows below the -/// architecture-derived wave/L2 bound). -/// - Single-block (numColumns < split-work threshold) -/// - Multi-pass radix (numColumns >= split-work threshold; requires -/// `scratch` sized via indexerTopKDecodeScratchBytes, -/// zero-init on first call and may be reused). -/// -/// `is_prefill = true` forces single-block (split-work suppressed). -void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, - int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - float* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, - bool is_prefill = false); - -/// Size of the multi-pass radix `scratch` buffer for these shapes. -size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int topK); +/// fp32 indexer TopK decode — L2-aware BS-threshold dispatcher with four +/// fallback tiers: +/// - GVR Heuristic (preIdx provided, kSeqSmall ≤ N < splitWork, BS < kBsLarge, K ∈ {512,1024,2048}) +/// - Insertion sort (N < kSortingAlgorithmThreshold) +/// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) +/// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) +void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, + int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, + int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, + int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1, + cudaStream_t const stream = 0); -/// bf16 overload; same contract. +/// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except +/// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and +/// the split-work tier is unsupported (the bf16/fp16 entry does not expose +/// the float aux buffers required for split-work). Insertion + radix tiers +/// share topKPerRowDecode with fp32 — histogram and sort run on float keys +/// after a static_cast(InputT) at HBM-read sites. +/// +/// Aborts with TLLM_CHECK if numColumns ≥ splitWorkThreshold; callers in +/// that regime must use the fp32 entry. void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, - int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, cudaStream_t const stream = 0, - void* scratch = nullptr, size_t scratchBytes = 0, bool is_prefill = false); + int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, + cudaStream_t const stream = 0); -/// fp16 overload; same contract. +/// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, cudaStream_t const stream = 0, void* scratch = nullptr, size_t scratchBytes = 0, - bool is_prefill = false); + __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048, cudaStream_t const stream = 0); -/// True iff invokeIndexerTopKDecode would pick the GVR tier for this shape: -/// K in {512,1024,2048}, numColumns in [kSeqSmall, splitWorkThreshold), and -/// numRows below the architecture-derived wave/L2 bound. Lets callers -/// provision preIdx / heuristicScratch only when needed. +/// Returns true iff invokeIndexerTopKDecode would route to the GVR Heuristic +/// kernel for this (numRows, numColumns, topK) triple, assuming valid preIdx +/// is provided and stride1 == 1. Useful for callers that need to provision a +/// preIdx tensor or heuristicScratch buffer only when GVR will be selected. +/// +/// Mirrors the gating logic of the dispatcher: K ∈ {512, 1024, 2048}, +/// numColumns ∈ [kSeqSmall, splitWorkThreshold), numRows < kBsLarge, where +/// kBsLarge = min(kBsWave, kBsL2) and kBsL2 scales with bytesPerElem. +/// +/// @param numRows logits rows (batch · next_n) +/// @param numColumns logits columns (max sequence length) +/// @param topK requested output size +/// @param bytesPerElem element size of logits (4 for fp32, 2 for bf16/fp16) bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytesPerElem = 4); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt new file mode 100644 index 000000000000..160dbced4b19 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/CMakeLists.txt @@ -0,0 +1,25 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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. +# + +set(SRC_CU compressorKernels.cu) + +add_library(compressorKernels_src OBJECT ${SRC_CU}) +set_property(TARGET compressorKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET compressorKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) +target_compile_options(compressorKernels_src + PRIVATE $<$:--use_fast_math>) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu new file mode 100644 index 000000000000..7fbe9034d50e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -0,0 +1,1792 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// ============================================================================ +// Compressor Kernels — DeepSeek-V4 KV Cache Compression +// ============================================================================ +// +// This file implements CUDA kernels for KV cache compression in the DeepSeek-V4 +// sparse attention system. The compressor reduces sequences of input tokens +// into fewer compressed tokens via learned weighted averaging (online softmax), +// then post-processes and scatters results into a paged KV cache. +// +// Three kernels are provided: +// +// 1. pagedKvCompressKernel — Decode path (single/few new tokens per batch). +// Loads prior compressor state from paged memory, performs online softmax +// with the new token(s), writes updated state back, and emits a compressed +// output token when compress_ratio tokens have been accumulated. +// +// 2. prefillReductionKernel — Prefill path (many tokens per batch). +// Processes full chunks of compress_ratio tokens in one shot via online +// softmax reduction over the input sequence. Also saves compressor state +// for any remainder tokens that don't form a complete chunk. +// +// 3. postProcessScatterKernel — Fused post-processing + paged cache write. +// Takes compressed output tokens and applies: RMSNorm → RoPE → Hadamard +// transform → optional V4-Pro QDQ → scatter to paged KV cache. Supports +// default, FP8, and V4-Pro MXFP8/MXFP4 QDQ cache modes. +// Keeps all intermediate values in float32 registers to avoid extra DRAM +// round-trips. +// +// Vectorization strategy: +// All kernels use 128-bit vectorized loads/stores (float4 / 8×bf16). +// VEC = number of elements per thread, chosen so that NTHRD = HEAD_DIM/VEC >= 32. +// For HEAD_DIM=128, bf16: VEC=4, NTHRD=32. For HEAD_DIM=512, bf16: VEC=8, NTHRD=64. +// +// Overlap mode (compress_ratio=4): +// When enabled, state_dim = 2*head_dim and the compressor uses overlapping +// windows: each compressed output is derived from both the previous and current +// chunk of compress_ratio tokens (previous chunk → first head_dim features, +// current chunk → second head_dim features). This doubles the state stored +// per position but improves compression quality. +// +// Template parameters: +// HEAD_DIM — Head dimension (128 or 512) +// ELEM_BYTES — Element size in bytes (2=bf16, 4=fp32) +// SCALE_TYPE — Output cache scale/dtype for postProcessScatterKernel +// ============================================================================ + +#include "tensorrt_llm/kernels/compressorKernels/compressorKernels.h" + +#include "tensorrt_llm/common/assert.h" +#include +#include +#include +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::compressor +{ + +// ============================================================================ +// Helper functions +// ============================================================================ + +// Full-warp butterfly reductions via __shfl_xor_sync (all 32 lanes participate). +__device__ inline float warpReduceSum(float val) +{ + for (int mask = 16; mask > 0; mask >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, mask); + return val; +} + +__device__ inline float warpReduceMax(float val) +{ + for (int mask = 16; mask > 0; mask >>= 1) + val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, mask)); + return val; +} + +// Runtime-dispatched element load (bf16 or fp32 → float). Used in the decode +// kernel where elem_bytes is a runtime parameter from paged state buffers. +__device__ inline float loadAsFloat(void const* base, int64_t offset, int elem_bytes) +{ + if (elem_bytes == 2) + return __bfloat162float(reinterpret_cast<__nv_bfloat16 const*>(base)[offset]); + else + return reinterpret_cast(base)[offset]; +} + +__device__ inline void storeFromFloat(void* base, int64_t offset, float val, int elem_bytes) +{ + if (elem_bytes == 2) + reinterpret_cast<__nv_bfloat16*>(base)[offset] = __float2bfloat16_rn(val); + else + reinterpret_cast(base)[offset] = val; +} + +// Bit-hack ceil(log2(x)) for x>0: equivalent to V4 reference fast_log2_ceil. +__device__ inline int fastLog2Ceil(float x) +{ + uint32_t const bits = __float_as_uint(x); + int const exp_part = static_cast((bits >> 23) & 0xFFu) - 127; + uint32_t const man_bits = bits & 0x007FFFFFu; + return exp_part + (man_bits != 0u ? 1 : 0); +} + +// Bit-hack 2^n for integer n: equivalent to V4 reference fast_pow2. +__device__ inline float fastPow2(int n) +{ + uint32_t const bits = static_cast(n + 127) << 23; + return __uint_as_float(bits); +} + +// V4-Pro fast_round_scale: 2^ceil(log2(amax * max_value_inv)). Operand order +// (multiplication vs `amax / max_value`) matches V4 byte-for-byte; using fp32 +// bit hacks avoids log2f/exp2f rounding so the resulting power-of-2 is exact. +__device__ inline float roundedPow2Scale(float amax, float max_value_inv, float min_amax) +{ + float const clamped_amax = fmaxf(amax, min_amax); + return fastPow2(fastLog2Ceil(clamped_amax * max_value_inv)); +} + +// Hardware FP4 (e2m1) round-trip via __nv_fp4_e2m1 (cuda_fp4.h). The +// constructor uses the SM100 PTX `cvt.rn.satfinite.e2m1x2.f32` cast +// (round-to-nearest-even with finite saturation), matching V4-Pro's +// Cast(FP4) byte-for-byte. Verified against the V4-Pro reference LUT +// (e2m1 levels {0, 0.5, 1, 1.5, 2, 3, 4, 6} with midpoint thresholds +// {0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0}) over 1.1M sweep inputs covering +// every exact level, every midpoint ±2 ULPs, out-of-range, subnormals, and +// dense random fp32 -- zero mismatches. The Python test reference still +// uses the explicit LUT so the kernel is checked against an independent +// software model rather than the HW cast itself. +__device__ inline uint8_t toUe8m0(float val) +{ + __nv_fp8_e8m0 out; + out.__x = __nv_cvt_float_to_e8m0(val, __NV_SATFINITE, cudaRoundPosInf); + return out.__x; +} + +__device__ inline uint8_t packE2M1x2(float lo, float hi) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "mov.b32 %0, {byte0, byte0, byte0, byte0};\n" + "}" + : "=r"(val) + : "f"(lo), "f"(hi)); + return static_cast(val); +#else + return 0; +#endif +} + +// Vectorized load/store types: maps byte-width to CUDA vector type. +template +struct VecType; + +template <> +struct VecType<4> +{ + using type = unsigned int; +}; // 32-bit: 2 bf16 or 1 fp32 + +template <> +struct VecType<8> +{ + using type = uint2; +}; // 64-bit: 4 bf16 or 2 fp32 + +template <> +struct VecType<16> +{ + using type = uint4; +}; // 128-bit: 8 bf16 or 4 fp32 + +// 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. +enum class CacheScaleType +{ + kNone = 0, + kFP8PerTensor = 1, // FP8 E4M3 with static scale=1.0 + kFP8Blockwise = 2, // FP8 E4M3 with per-128-element fp32 scales + kMXFP4Blockwise = 3 // packed FP4 E2M1 with per-32 UE8M0 scales +}; + +// ============================================================================ +// Decode Kernel: pagedKvCompressKernel +// +// Template: +// NEXT_N: number of new tokens per sequence in this decode step (1-4) +// +// Grid: (batch_size) — one block per batch element +// Block: (NTHRD) where NTHRD = HEAD_DIM / VEC (>= 32 threads) +// +// Algorithm per batch element: +// For each new token in the decode step: +// 1. Load existing compressor state (partial kv/score) from paged cache +// 2. Perform online softmax: accumulate new token's contribution using +// the numerically stable running max + weighted sum formulation +// 3. Write updated state back to paged cache +// 4. If compress_ratio tokens accumulated → emit compressed output, +// reset state for next compression window +// +// Each thread handles VEC contiguous elements of head_dim. In overlap mode +// (state_dim = 2*head_dim), Phase 1 iterates over 2 column halves. +// +// Memory layout: +// kv_score: [total_tokens, 2 * state_dim] — interleaved KV and score projections +// paged_kv: paged cache for compressor KV state +// paged_score: paged cache for compressor score state (with APE bias) +// output: [total_comp_tokens, head_dim] — compressed output tokens +// ============================================================================ + +// Helper: vectorized online softmax step reading from paged KV/score state. +// Loads one position's KV and score from paged memory and updates the running +// online softmax accumulators (rmax, rsum, rwsum) per element. +// APE is already baked into paged_score (added during Phase 1), so no APE +// addition is performed here. +template +__device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_kv_raw, + void const* __restrict__ paged_score_raw, + int64_t page_sd, // page_size * state_dim (in elements) + int state_dim, + int phys_kv, // physical page index for kv + int phys_sc, // physical page index for score + int blk_off, // offset within page + int kv_col_off, // column offset (0 or HEAD_DIM) + int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) +{ + using IoElemT = typename std::conditional::type; + constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + using IoVecT = typename VecType::type; + + auto const* kv = reinterpret_cast(paged_kv_raw); + auto const* sc = reinterpret_cast(paged_score_raw); + + int64_t base_kv = static_cast(phys_kv) * page_sd + blk_off * state_dim + kv_col_off; + int64_t base_sc = static_cast(phys_sc) * page_sd + blk_off * state_dim + kv_col_off; + + IoVecT k_raw = reinterpret_cast(&kv[base_kv])[tid]; + IoVecT s_raw = reinterpret_cast(&sc[base_sc])[tid]; + IoElemT const* ke = reinterpret_cast(&k_raw); + IoElemT const* se = reinterpret_cast(&s_raw); + +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float kf[4] = {static_cast(ke[i]), static_cast(ke[i + 1]), static_cast(ke[i + 2]), + static_cast(ke[i + 3])}; + // score already includes APE (added during Phase 1 store) + float sf[4] = {static_cast(se[i]), static_cast(se[i + 1]), static_cast(se[i + 2]), + static_cast(se[i + 3])}; + // Online softmax: maintain running (max, sum_exp, weighted_sum) per element. + // nm = new max, sc_f = rescale factor for old accumulators, tm = exp(score - new_max). + // Final output: rwsum / rsum = weighted average of KV values. +#pragma unroll + for (int j = 0; j < 4; j++) + { + float nm = fmaxf(rmax[i + j], sf[j]); + float sc_f = expf(rmax[i + j] - nm); + float tm = expf(sf[j] - nm); + rsum[i + j] = rsum[i + j] * sc_f + tm; + rwsum[i + j] = rwsum[i + j] * sc_f + kf[j] * tm; + rmax[i + j] = nm; + } + } +} + +template +__global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, + void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, + int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, + int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ cu_seq_lens, + int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, int out_elem_bytes) +{ + using IoElemT = 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"); + constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); + constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + using IoVecT = typename VecType::type; + static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); + + // HEAD_BLOCKS: split head_dim across blockIdx.y for better SM utilisation. + // For HD=512 bf16: NTHRD_BASE=64 → HEAD_BLOCKS=2, NTHRD_INNER=32. + // For HD=128 bf16: NTHRD_BASE=32 → HEAD_BLOCKS=1, NTHRD_INNER=32. + // For HD=512 fp32: NTHRD_BASE=128 → HEAD_BLOCKS=4, NTHRD_INNER=32. + constexpr int NTHRD_BASE = HEAD_DIM / VEC; + constexpr int HEAD_BLOCKS = (NTHRD_BASE > 32) ? (NTHRD_BASE / 32) : 1; + constexpr int NTHRD_INNER = NTHRD_BASE / HEAD_BLOCKS; // always <= 32 + // ELEM_PER_BLOCK: head_dim elements handled by one blockIdx.y block. + // = NTHRD_INNER * VEC = HEAD_DIM / HEAD_BLOCKS. + // Used for the multi-warp shared-memory merge layout so that each block + // only allocates storage for its own head slice, not the full HEAD_DIM. + constexpr int ELEM_PER_BLOCK = NTHRD_INNER * VEC; + + // state_dim is fully determined by template parameters. + constexpr int STATE_DIM = IS_OVERLAP ? 2 * HEAD_DIM : HEAD_DIM; + constexpr int64_t TWO_SD = 2 * STATE_DIM; + constexpr int COFF = IS_OVERLAP ? 2 : 1; + + int const tid = threadIdx.x % NTHRD_INNER; + int const warp_id = threadIdx.x / NTHRD_INNER; // 0..NUM_RED_WARPS-1 + int const batch_idx = blockIdx.x; + int const head_blk = blockIdx.y; + int const eff_tid = head_blk * NTHRD_INNER + tid; + + int const sp = start_pos_arr[batch_idx]; + int const kv_len = kv_lens[batch_idx]; + int const in_off = cu_seq_lens[batch_idx]; + int const out_off = cu_kv_comp[batch_idx]; + int64_t const page_sd = static_cast(page_size) * STATE_DIM; + + auto const* kv_score = reinterpret_cast(kv_score_raw); + auto* paged_kv = reinterpret_cast(paged_kv_raw); + auto* paged_score = reinterpret_cast(paged_score_raw); + + // ================================================================ + // Phase 1: Write NEXT_N new tokens' KV and score state to paged cache. + // + // Only warp 0 participates (all warps share the same eff_tid mapping). + // When NUM_RED_WARPS == 1, the guard compiles away. + // ================================================================ + if (warp_id == 0) + { +#pragma unroll + for (int t = 0; t < NEXT_N; t++) + { + int token_idx = sp + t; + if (token_idx < kv_len) + { + int ape_idx = token_idx % COMPRESS_RATIO; + int log_blk = token_idx / page_size; + int blk_off = token_idx % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + + for (int col_idx = 0; col_idx < COFF; col_idx++) + { + int const col = col_idx * HEAD_DIM; + int64_t const src = static_cast(in_off + t) * TWO_SD + col; + int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * STATE_DIM + col; + int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * STATE_DIM + col; + + IoVecT kv_raw = reinterpret_cast(&kv_score[src])[eff_tid]; + IoVecT sc_raw = reinterpret_cast(&kv_score[src + STATE_DIM])[eff_tid]; + + reinterpret_cast(&paged_kv[dkv])[eff_tid] = kv_raw; + + IoElemT const* sc_e = reinterpret_cast(&sc_raw); + IoVecT sc_out; + IoElemT* sc_o = reinterpret_cast(&sc_out); +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float4 av + = *reinterpret_cast(&ape[ape_idx * STATE_DIM + col + eff_tid * VEC + i]); + sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); + sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); + sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); + sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); + } + reinterpret_cast(&paged_score[dsc])[eff_tid] = sc_out; + } + } + } + } + + if constexpr (NUM_RED_WARPS > 1) + { + __syncthreads(); + } + + // ================================================================ + // Phase 2: Count how many complete compression windows finished. + // ================================================================ + int last_token_idx = sp + NEXT_N - 1; + int num_compressions = (last_token_idx + 1) / COMPRESS_RATIO - sp / COMPRESS_RATIO; + + // ================================================================ + // Phase 3: Online softmax reduction over each complete chunk. + // + // When NUM_RED_WARPS > 1, the compress_ratio positions are split + // across warps. Each warp reduces its partition independently, then + // partial (rmax, rsum, rwsum) accumulators are merged via shared + // memory using the log-sum-exp identity. + // ================================================================ + for (int c = 0; c < NEXT_N; c++) + { + if (c >= num_compressions) + break; + + int compress_idx = sp / COMPRESS_RATIO + c; + int curr_chunk_start = compress_idx * COMPRESS_RATIO; + + float rmax[VEC], rsum[VEC], rwsum[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + rmax[i] = -INFINITY; + rsum[i] = 0.0f; + rwsum[i] = 0.0f; + } + + constexpr int positions_per_warp = COMPRESS_RATIO / NUM_RED_WARPS; + int const my_r_start = warp_id * positions_per_warp; + int const my_r_end = (warp_id == NUM_RED_WARPS - 1) ? COMPRESS_RATIO : (my_r_start + positions_per_warp); + + if constexpr (IS_OVERLAP) + { + int prev_start = curr_chunk_start - COMPRESS_RATIO; + if (prev_start >= 0) + { + if (page_size >= COMPRESS_RATIO) + { + int log_blk_prev = prev_start / page_size; + int phys_kv_prev = block_table_kv[batch_idx * max_blocks + log_blk_prev]; + int phys_sc_prev = block_table_score[batch_idx * max_blocks + log_blk_prev]; + int chunk_off_prev = prev_start % page_size; + for (int r = my_r_start; r < my_r_end; r++) + { + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv_prev, phys_sc_prev, chunk_off_prev + r, 0, eff_tid, rmax, rsum, rwsum); + } + } + else + { + for (int r = my_r_start; r < my_r_end; r++) + { + int pos = prev_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); + } + } + } + + if (page_size >= COMPRESS_RATIO) + { + int log_blk_cur = curr_chunk_start / page_size; + int phys_kv_cur = block_table_kv[batch_idx * max_blocks + log_blk_cur]; + int phys_sc_cur = block_table_score[batch_idx * max_blocks + log_blk_cur]; + int chunk_off_cur = curr_chunk_start % page_size; + for (int r = my_r_start; r < my_r_end; r++) + { + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv_cur, phys_sc_cur, chunk_off_cur + r, HEAD_DIM, eff_tid, rmax, rsum, rwsum); + } + } + else + { + for (int r = my_r_start; r < my_r_end; r++) + { + int pos = curr_chunk_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, blk_off, HEAD_DIM, eff_tid, rmax, rsum, rwsum); + } + } + } + else + { + if (page_size >= COMPRESS_RATIO) + { + int log_blk = curr_chunk_start / page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + int chunk_off = curr_chunk_start % page_size; + for (int r = my_r_start; r < my_r_end; r++) + { + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, chunk_off + r, 0, eff_tid, rmax, rsum, rwsum); + } + } + else + { + for (int r = my_r_start; r < my_r_end; r++) + { + int pos = curr_chunk_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); + } + } + } + + // Multi-warp merge epilogue (compiled away when NUM_RED_WARPS == 1). + if constexpr (NUM_RED_WARPS > 1) + { + // Shared-memory layout: [NUM_RED_WARPS * ELEM_PER_BLOCK] per array. + // ELEM_PER_BLOCK = NTHRD_INNER * VEC = HEAD_DIM / HEAD_BLOCKS, i.e. + // the number of head_dim elements covered by this block (blockIdx.y). + // Using ELEM_PER_BLOCK (not HEAD_DIM) avoids 4x over-allocation when + // HEAD_BLOCKS > 1 (e.g. HD=512 fp32 has HEAD_BLOCKS=4). + extern __shared__ float smem[]; + float* s_rmax = smem; + float* s_rsum = s_rmax + NUM_RED_WARPS * ELEM_PER_BLOCK; + float* s_rwsum = s_rsum + NUM_RED_WARPS * ELEM_PER_BLOCK; + + // local_elem: index within this block's head slice [0, ELEM_PER_BLOCK). + // = tid * VEC + i (tid = threadIdx.x % NTHRD_INNER, same as eff_tid - head_blk*NTHRD_INNER) +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const local_elem = tid * VEC + i; + s_rmax[warp_id * ELEM_PER_BLOCK + local_elem] = rmax[i]; + s_rsum[warp_id * ELEM_PER_BLOCK + local_elem] = rsum[i]; + s_rwsum[warp_id * ELEM_PER_BLOCK + local_elem] = rwsum[i]; + } + __syncthreads(); + + if (warp_id == 0) + { + for (int w = 1; w < NUM_RED_WARPS; w++) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const local_elem = tid * VEC + i; + float const m2 = s_rmax[w * ELEM_PER_BLOCK + local_elem]; + float const s2 = s_rsum[w * ELEM_PER_BLOCK + local_elem]; + float const ws2 = s_rwsum[w * ELEM_PER_BLOCK + local_elem]; + + float const nm = fmaxf(rmax[i], m2); + float const sc1 = expf(rmax[i] - nm); + float const sc2 = expf(m2 - nm); + rsum[i] = rsum[i] * sc1 + s2 * sc2; + rwsum[i] = rwsum[i] * sc1 + ws2 * sc2; + rmax[i] = nm; + } + } + } + __syncthreads(); + } + + bool const should_write = (NUM_RED_WARPS == 1) || (warp_id == 0); + if (should_write) + { + int64_t const out_base = static_cast(out_off + c) * HEAD_DIM + eff_tid * VEC; + if (out_elem_bytes == 2) + { + __nv_bfloat16 packed[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + packed[i] = __float2bfloat16_rn(rwsum[i] / rsum[i]); + using OutVecT = typename VecType::type; + *reinterpret_cast(&reinterpret_cast<__nv_bfloat16*>(output_raw)[out_base]) + = *reinterpret_cast(packed); + } + else + { + float result[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + result[i] = rwsum[i] / rsum[i]; +#pragma unroll + for (int i = 0; i < VEC; i += 4) + *reinterpret_cast(&reinterpret_cast(output_raw)[out_base + i]) + = *reinterpret_cast(&result[i]); + } + } + } +} + +// Explicit instantiations for decode kernel (NUM_RED_WARPS defaults to 1) +#define INST_DECODE(HD, EB, CR, NN, NRW) \ + template __global__ void pagedKvCompressKernel(void const*, float const*, void*, void*, \ + int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, int, \ + int, int); + +#define INST_DECODE_NN(HD, EB, CR) \ + INST_DECODE(HD, EB, CR, 1, 1) \ + INST_DECODE(HD, EB, CR, 2, 1) INST_DECODE(HD, EB, CR, 3, 1) INST_DECODE(HD, EB, CR, 4, 1) + +INST_DECODE_NN(128, 2, 4) +INST_DECODE_NN(128, 2, 128) +INST_DECODE_NN(128, 4, 4) +INST_DECODE_NN(128, 4, 128) +INST_DECODE_NN(512, 2, 4) +INST_DECODE_NN(512, 2, 128) +INST_DECODE_NN(512, 4, 4) +INST_DECODE_NN(512, 4, 128) + +// 4-warp parallel reduction variants for large compress_ratio. +// HD=128: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. +// HD=512 bf16: ELEM_PER_BLOCK=256, smem = 3*4*256*4 = 12 KB. +// HD=512 fp32: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. +// NEXT_N=1 (single-token decode) and NEXT_N=2 (MTP speculative decode). +INST_DECODE(128, 2, 128, 1, 4) +INST_DECODE(128, 4, 128, 1, 4) +INST_DECODE(128, 2, 128, 2, 4) +INST_DECODE(128, 4, 128, 2, 4) +INST_DECODE(512, 2, 128, 1, 4) +INST_DECODE(512, 4, 128, 1, 4) +INST_DECODE(512, 2, 128, 2, 4) +INST_DECODE(512, 4, 128, 2, 4) + +#undef INST_DECODE_NN +#undef INST_DECODE + +// ============================================================================ +// Decode Launch Wrapper +// +// Dispatches to the correct template instantiation based on head_dim, elem_bytes, +// and next_n (number of new tokens per decode step, capped at 4). +// Grid is 2D: (batch_size, head_blocks) where head_blocks = NTHRD_BASE / 32. +// For HD=512 bf16: head_blocks=2; for HD=128 bf16: head_blocks=1. +// ============================================================================ + +// Forward declaration (defined in prefill section below). +static inline int prefillNthreads(int head_dim, int io_elem_bytes); + +void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, + int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, + int32_t const* start_pos, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, + int max_blocks, int head_dim, int compress_ratio, int next_n, int io_elem_bytes, int out_elem_bytes, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO( + compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128"); + + // Compute HEAD_BLOCKS: mirrors the compile-time constant in the kernel. + // VEC = max_vec if HEAD_DIM/max_vec >= 32, else HEAD_DIM/32. + // NTHRD_BASE = HEAD_DIM / VEC; HEAD_BLOCKS = NTHRD_BASE / 32 (or 1 if <= 32). + // This spreads the head_dim across multiple blocks for better SM utilisation. + int const max_vec_elem = 16 / io_elem_bytes; + int const vec = (head_dim / max_vec_elem >= 32) ? max_vec_elem : (head_dim / 32); + int const nthrd_base = head_dim / vec; // mirrors kernel's NTHRD_BASE = HEAD_DIM / VEC + int const head_blocks = (nthrd_base > 32) ? (nthrd_base / 32) : 1; + int const nthreads_inner = nthrd_base / head_blocks; // = min(32, nthrd_base) = always 32 + + // For large compress_ratio, use 4-warp parallel reduction to cut the serial + // softmax loop from COMPRESS_RATIO iterations to COMPRESS_RATIO/4 per warp. + // Supported configs: CR=128, (HD=128 or HD=512), NEXT_N=1 or NEXT_N=2. + // + // smem per block = 3 * MULTI_WARP * ELEM_PER_BLOCK * sizeof(float) + // where ELEM_PER_BLOCK = nthreads_inner * vec = HEAD_DIM / HEAD_BLOCKS. + // HD=128: ELEM_PER_BLOCK=128 → 6 KB. + // HD=512 bf16 (vec=8, HEAD_BLOCKS=2): ELEM_PER_BLOCK=256 → 12 KB. + // HD=512 fp32 (vec=4, HEAD_BLOCKS=4): ELEM_PER_BLOCK=128 → 6 KB. + constexpr int MULTI_WARP = 4; + bool const use_multi_warp = (compress_ratio == 128 && next_n <= 2); + int const num_red_warps = use_multi_warp ? MULTI_WARP : 1; + int const nthreads = nthreads_inner * num_red_warps; + int const elem_per_block = nthreads_inner * vec; // = HEAD_DIM / HEAD_BLOCKS + int const smem_bytes = use_multi_warp ? (3 * MULTI_WARP * elem_per_block * static_cast(sizeof(float))) : 0; + + dim3 grid(batch_size, head_blocks); + +#define LAUNCH_DECODE(HD, EB, CR, NN) \ + pagedKvCompressKernel<<>>(kv_score, ape, paged_kv, \ + paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ + page_size, max_blocks, out_elem_bytes) + +#define LAUNCH_DECODE_MW(HD, EB, CR, NN) \ + pagedKvCompressKernel<<>>(kv_score, ape, paged_kv, \ + paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ + page_size, max_blocks, out_elem_bytes) + +#define DISPATCH_NN_MW(HD, EB, CR) \ + switch (next_n) \ + { \ + case 2: LAUNCH_DECODE_MW(HD, EB, CR, 2); break; \ + default: LAUNCH_DECODE_MW(HD, EB, CR, 1); break; \ + } + +#define DISPATCH_NN(HD, EB, CR) \ + switch (next_n) \ + { \ + case 1: LAUNCH_DECODE(HD, EB, CR, 1); break; \ + case 2: LAUNCH_DECODE(HD, EB, CR, 2); break; \ + case 3: LAUNCH_DECODE(HD, EB, CR, 3); break; \ + default: LAUNCH_DECODE(HD, EB, CR, 4); break; \ + } + + if (use_multi_warp) + { + // Multi-warp path: HD=128 or HD=512, NEXT_N=1 or NEXT_N=2, CR=128. + if (io_elem_bytes == 4) + { + if (head_dim == 512) + { + DISPATCH_NN_MW(512, 4, 128); + } + else + { + DISPATCH_NN_MW(128, 4, 128); + } + } + else + { + if (head_dim == 512) + { + DISPATCH_NN_MW(512, 2, 128); + } + else + { + DISPATCH_NN_MW(128, 2, 128); + } + } + } + else if (compress_ratio == 4) + { + if (io_elem_bytes == 4) + { + if (head_dim == 512) + { + DISPATCH_NN(512, 4, 4); + } + else + { + DISPATCH_NN(128, 4, 4); + } + } + else + { + if (head_dim == 512) + { + DISPATCH_NN(512, 2, 4); + } + else + { + DISPATCH_NN(128, 2, 4); + } + } + } + else if (io_elem_bytes == 4) + { + if (head_dim == 512) + { + DISPATCH_NN(512, 4, 128); + } + else + { + DISPATCH_NN(128, 4, 128); + } + } + else + { + if (head_dim == 512) + { + DISPATCH_NN(512, 2, 128); + } + else + { + DISPATCH_NN(128, 2, 128); + } + } + +#undef DISPATCH_NN +#undef DISPATCH_NN_MW +#undef LAUNCH_DECODE_MW +#undef LAUNCH_DECODE +} + +// ============================================================================ +// Prefill Kernel: prefillReductionKernel +// +// Template: +// +// Grid: (batch_size, max_outputs_per_batch) — one block per compressed output +// Block: (NTHRD) where NTHRD = HEAD_DIM / VEC (>= 32 threads) +// +// Unlike the decode kernel (which operates token-by-token from paged state), +// the prefill kernel processes the full input sequence at once. Each block +// reads compress_ratio consecutive input rows from kv_score and reduces them +// via online softmax to produce one compressed output. +// +// The last block (local_output_idx == num_outputs - 1) also handles saving +// compressor state for any remainder tokens that don't form a full chunk. +// This state is written to paged kv/score caches for use in future decode steps. +// +// Memory layout: +// kv_score: [total_tokens, 2*state_dim] — interleaved KV and score from linear projection +// paged_kv: paged cache for compressor state (remainder) +// paged_score: paged cache for compressor score state (remainder, with APE) +// output: [total_comp_tokens, head_dim] — compressed output tokens +// ============================================================================ + +// Per-element online softmax step on VEC elements via 128-bit vectorized loads. +// Reads directly from the kv_score input buffer (not paged state) since prefill +// has the full sequence available. +template +__device__ __forceinline__ void prefillSoftmaxVec(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, + int64_t row_elem, // (input_offset + row_idx) * two_sd + int kv_col_off, // column offset into kv_score row (0 or HEAD_DIM) + int ape_base, // r * state_dim + ape_col_off + int state_dim, int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) +{ + using IoElemT = typename std::conditional::type; + constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + using IoVecT = typename VecType::type; + + auto const* kv = reinterpret_cast(kv_score_raw); + + IoVecT k_raw = reinterpret_cast(&kv[row_elem + kv_col_off])[tid]; + IoVecT s_raw = reinterpret_cast(&kv[row_elem + state_dim + kv_col_off])[tid]; + IoElemT const* ke = reinterpret_cast(&k_raw); + IoElemT const* se = reinterpret_cast(&s_raw); + +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float4 av = *reinterpret_cast(&ape[ape_base + tid * VEC + i]); + float kf[4] = {static_cast(ke[i]), static_cast(ke[i + 1]), static_cast(ke[i + 2]), + static_cast(ke[i + 3])}; + float sf[4] = {static_cast(se[i]) + av.x, static_cast(se[i + 1]) + av.y, + static_cast(se[i + 2]) + av.z, static_cast(se[i + 3]) + av.w}; +#pragma unroll + for (int j = 0; j < 4; j++) + { + float nm = fmaxf(rmax[i + j], sf[j]); + float sc = expf(rmax[i + j] - nm); + float tm = expf(sf[j] - nm); + rsum[i + j] = rsum[i + j] * sc + tm; + rwsum[i + j] = rwsum[i + j] * sc + kf[j] * tm; + rmax[i + j] = nm; + } + } +} + +template +__global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, + void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, + int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, + int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ cu_seq_lens, + int32_t const* __restrict__ cu_kv_comp, int page_size, int state_dim, int max_blocks, int out_elem_bytes) +{ + using IoElemT = typename std::conditional::type; + static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); + constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); + + constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + using IoVecT = typename VecType::type; + static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); + + int const tid = threadIdx.x; + int const batch_idx = blockIdx.x; + int const local_output_idx = blockIdx.y; + + int const sp = start_pos_arr[batch_idx]; + int const kv_len = kv_lens[batch_idx]; + int const input_offset = cu_seq_lens[batch_idx]; + int const output_offset = cu_kv_comp[batch_idx]; + + // Absolute compression index range. Window k covers [k*R, (k+1)*R). + // We emit one output per complete window that gains at least one new token. + int const first_abs_idx = sp / COMPRESS_RATIO; + int const last_abs_idx = kv_len / COMPRESS_RATIO; // exclusive + int const actual_num_outputs = last_abs_idx - first_abs_idx; + // Keep at least one block so the last CTA can still persist remainder state + // even when this chunk has no full compression window. + int const num_outputs = max(actual_num_outputs, 1); + + if (local_output_idx >= num_outputs) + return; + + constexpr int coff = IS_OVERLAP ? 2 : 1; + bool const should_compress = (local_output_idx < actual_num_outputs); + + // Absolute window for this output block. + int const abs_idx = first_abs_idx + local_output_idx; + int const win_start = abs_idx * COMPRESS_RATIO; + + auto const* kv_score = reinterpret_cast(kv_score_raw); + auto* paged_kv = reinterpret_cast(paged_kv_raw); + auto* paged_score = reinterpret_cast(paged_score_raw); + + int64_t const two_sd = 2 * state_dim; + int64_t const page_sd = static_cast(page_size) * state_dim; + + // ================================================================ + // Phase 1: State Update (all output blocks, for block reuse support) + // + // 1a runs on every block that has a full chunk (should_compress), writing + // each block's own window to paged cache so any prefix slice is valid for + // block reuse. 1b (remainder) still runs only on the last block. + // ================================================================ + + // Helper: write a contiguous block of positions to paged KV/score cache. + // Only new positions (>= sp) are written; positions already persisted from + // prior calls are skipped by starting the loop at write_r_start. + // APE index is r (position within the compression window), matching the + // index used during the original decode/prefill that first established + // the window alignment. + auto write_to_paged = [&](int range_start, int range_end, int write_r_start) + { + for (int r = write_r_start; r < range_end - range_start; r++) + { + int const pos = range_start + r; + int const input_row = pos - sp; // >= 0 since pos >= sp (loop starts at write_r_start) + int const log_blk = pos / page_size; + int const blk_off = pos % page_size; + int const phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int const phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + + for (int col_idx = 0; col_idx < coff; col_idx++) + { + int const col = col_idx * HEAD_DIM; + int64_t const src = static_cast(input_offset + input_row) * two_sd + col; + int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * state_dim + col; + int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * state_dim + col; + + IoVecT kv_raw = reinterpret_cast(&kv_score[src])[tid]; + IoVecT sc_raw = reinterpret_cast(&kv_score[src + state_dim])[tid]; + reinterpret_cast(&paged_kv[dkv])[tid] = kv_raw; + + IoElemT const* sc_e = reinterpret_cast(&sc_raw); + IoVecT sc_out; + IoElemT* sc_o = reinterpret_cast(&sc_out); +#pragma unroll + for (int i = 0; i < VEC; i += 4) + { + float4 av = *reinterpret_cast(&ape[r * state_dim + col + tid * VEC + i]); + sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); + sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); + sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); + sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); + } + reinterpret_cast(&paged_score[dsc])[tid] = sc_out; + } + } + }; + + // 1a. Full chunk for this output block. + // Positions [win_start, sp) are already in paged cache from a prior call; + // start the loop at write_r_start to skip them without a per-iteration branch. + if (should_compress) + { + int const write_r_start = (win_start < sp) ? (sp - win_start) : 0; + write_to_paged(win_start, win_start + COMPRESS_RATIO, write_r_start); + } + + // 1b. Remainder tokens (last block only). + // Tokens past the last complete window are persisted so a later call can + // continue the same compression window. + // rem_start_pos < sp when the chunk has no full window (actual_num_outputs == 0); + // rem_write_start skips those already-paged positions without a per-iteration branch. + if (local_output_idx == num_outputs - 1) + { + int const rem_start_pos = last_abs_idx * COMPRESS_RATIO; + int const rem_count = kv_len - rem_start_pos; + int const rem_write_start = (rem_start_pos < sp) ? (sp - rem_start_pos) : 0; + write_to_paged(rem_start_pos, rem_start_pos + rem_count, rem_write_start); + } + + // ================================================================ + // Phase 2: Online softmax reduction (vectorized) + // + // Each block reduces compress_ratio rows into one output via per-element + // online softmax: output[d] = sum_r(kv[r,d] * softmax(score[r,d] + ape[r,d])) + // In overlap mode, combines previous chunk's first-half and current chunk's + // second-half features (same as decode kernel). + // ================================================================ + if (!should_compress) + return; + + float rmax[VEC], rsum[VEC], rwsum[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + rmax[i] = -INFINITY; + rsum[i] = 0.0f; + rwsum[i] = 0.0f; + } + + // Helper: online-softmax reduction over a contiguous window of COMPRESS_RATIO positions. + // Positions [range_start, range_start + new_start) come from paged cache (APE already + // fused during the call that wrote them); positions [range_start + new_start, range_start + // + COMPRESS_RATIO) are new tokens read from kv_score with live APE addition. + // Two branch-free loops avoid a per-iteration if/else on pos < sp. + // kv_col_off — column offset into kv_score / paged state (0 or HEAD_DIM for overlap) + // ape_col_off — column offset into APE table for the score column (same as kv_col_off) + auto reduce_window = [&](int range_start, int kv_col_off, int ape_col_off) + { + // Precompute split once for the whole window. + int const new_start = (range_start < sp) ? min(sp - range_start, COMPRESS_RATIO) : 0; + + // Paged portion — APE already fused when tokens were stored. + for (int r = 0; r < new_start; r++) + { + int const pos = range_start + r; + int log_blk = pos / page_size; + int blk_off = pos % page_size; + int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; + int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, state_dim, phys_kv, + phys_sc, blk_off, kv_col_off, tid, rmax, rsum, rwsum); + } + + // New-token portion — read from kv_score and add APE live. + for (int r = new_start; r < COMPRESS_RATIO; r++) + { + int const pos = range_start + r; + int const input_row = pos - sp; + int64_t const row = static_cast(input_offset + input_row) * two_sd; + prefillSoftmaxVec( + kv_score_raw, ape, row, kv_col_off, r * state_dim + ape_col_off, state_dim, tid, rmax, rsum, rwsum); + } + }; + + if constexpr (IS_OVERLAP) + { + // Overlap mode: each output combines + // prev-segment (first head_dim, kv_col=0) from window (abs_idx-1) + // curr-segment (second head_dim, kv_col=HD) from window abs_idx + if (abs_idx > 0) + reduce_window((abs_idx - 1) * COMPRESS_RATIO, 0, 0); // prev window, first half + reduce_window(win_start, HEAD_DIM, HEAD_DIM); // curr window, second half + } + else + { + // Non-overlap mode: reduce one ratio-sized window. + reduce_window(win_start, 0, 0); + } + + // ================================================================ + // Store output (vectorized) + // ================================================================ + int64_t const out_base = static_cast(output_offset + local_output_idx) * HEAD_DIM + tid * VEC; + + if (out_elem_bytes == 2) + { + __nv_bfloat16 packed[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + packed[i] = __float2bfloat16_rn(rwsum[i] / rsum[i]); + // VEC * 2 bytes: VEC=4 → 8B (uint2), VEC=8 → 16B (uint4) + using OutVecT = typename VecType::type; + *reinterpret_cast(&reinterpret_cast<__nv_bfloat16*>(output_raw)[out_base]) + = *reinterpret_cast(packed); + } + else + { + float result[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + result[i] = rwsum[i] / rsum[i]; + // VEC * 4 bytes: VEC=4 → 16B (uint4/float4), VEC=8 → 32B (2×float4) +#pragma unroll + for (int i = 0; i < VEC; i += 4) + *reinterpret_cast(&reinterpret_cast(output_raw)[out_base + i]) + = *reinterpret_cast(&result[i]); + } +} + +// Explicit instantiations +#define INST_PREFILL(HD, EB, CR) \ + template __global__ void prefillReductionKernel(void const*, float const*, void*, void*, \ + int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, int, \ + int, int, int); + +INST_PREFILL(128, 2, 4) +INST_PREFILL(128, 2, 128) +INST_PREFILL(128, 4, 4) +INST_PREFILL(128, 4, 128) +INST_PREFILL(512, 2, 4) +INST_PREFILL(512, 2, 128) +INST_PREFILL(512, 4, 4) +INST_PREFILL(512, 4, 128) +#undef INST_PREFILL + +// ============================================================================ +// Prefill Launch Wrapper +// +// Grid is (batch_size, max(max_outputs, 1)). Blocks for local_output_idx >= num_outputs +// early-exit inside the kernel. +// ============================================================================ + +// Compute threads per block: mirrors compile-time NTHRD = HEAD_DIM / VEC. +static inline int prefillNthreads(int head_dim, int io_elem_bytes) +{ + int max_vec = 16 / io_elem_bytes; + int vec = (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); + return head_dim / vec; +} + +void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, + int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, + int32_t const* start_pos, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, + int max_blocks, int head_dim, int compress_ratio, int max_outputs, int io_elem_bytes, 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"); + int const nthreads = prefillNthreads(head_dim, io_elem_bytes); + int const coff = overlap ? 2 : 1; + int const state_dim = coff * head_dim; + dim3 grid(batch_size, max(max_outputs, 1)); + +#define LAUNCH_PREFILL(HD, EB, CR) \ + prefillReductionKernel<<>>(kv_score, ape, paged_kv, paged_score, \ + block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, page_size, state_dim, \ + max_blocks, out_elem_bytes) + + if (io_elem_bytes == 4) + { + if (head_dim == 512) + { + if (compress_ratio == 4) + LAUNCH_PREFILL(512, 4, 4); + else + LAUNCH_PREFILL(512, 4, 128); + } + else + { + if (compress_ratio == 4) + LAUNCH_PREFILL(128, 4, 4); + else + LAUNCH_PREFILL(128, 4, 128); + } + } + else + { + if (head_dim == 512) + { + if (compress_ratio == 4) + LAUNCH_PREFILL(512, 2, 4); + else + LAUNCH_PREFILL(512, 2, 128); + } + else + { + if (compress_ratio == 4) + LAUNCH_PREFILL(128, 2, 4); + else + LAUNCH_PREFILL(128, 2, 128); + } + } + +#undef LAUNCH_PREFILL +} + +// ============================================================================ +// Postprocess + Scatter Kernel: postProcessScatterKernel +// +// Template: +// +// Grid: (total_tokens) — one block per compressed token +// Block: (NTHRD = HEAD_DIM / VEC) threads, always >= 32 +// Smem: HEAD_DIM * sizeof(float) — used for cross-warp Hadamard butterfly +// +// This kernel fuses all post-compression processing with the paged cache write +// into a single kernel launch, keeping data in float32 registers throughout. +// This eliminates the DRAM round-trip that a split postprocess+scatter would need. +// +// SCALE_TYPE alone determines the output cache layout: +// - kNone: ELEM_BYTES per value (bf16 / fp32) into kv_cache. +// - kFP8PerTensor: one fp8 byte per value, implicit scale=1.0. +// - 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. +// +// Pipeline (10 steps, all in fp32 registers): +// 1. Vectorized load compressed token from kv_comp +// 2. RMSNorm: compute sum-of-squares → cross-warp reduce → rsqrt → scale +// 3. Apply RMSNorm weights +// 4. RoPE: interleaved even/odd rotation on rope_dim elements (skip nope_dim) +// 5. Hadamard butterfly transform (3 phases: local → warp shuffle → shared mem) +// 6. Scale by 1/sqrt(HEAD_DIM) (Hadamard normalization) +// 7. Optionally write postprocessed result to kv_out (for callers that need it) +// 8. Binary search cu_kv_comp to find batch_idx for this token +// 9. Compute paged cache destination (logical→physical block via block table) +// 10. Store to cache (layout by SCALE_TYPE). +// ============================================================================ + +template +__global__ void postProcessScatterKernel(void const* __restrict__ kv_comp, // [total_tokens, head_dim] input + void* __restrict__ kv_out, // [total_tokens, head_dim] postprocessed output (may be nullptr) + void const* __restrict__ rms_weight, // [head_dim] + float rms_eps, + float const* __restrict__ cos_sin_table, // [max_pos, 2, rope_dim/2] + int32_t const* __restrict__ position_ids, // [total_tokens] + int nope_dim, int rope_dim, + // scatter params + void* __restrict__ kv_cache, // paged cache buffer + int32_t const* __restrict__ num_outputs_arr, int32_t const* __restrict__ cu_kv_comp, + int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ block_offsets, + bool const* __restrict__ compressed_mask, int batch_size, int tokens_per_block, int max_blocks, + int cache_stride_blk_bytes, int total_tokens, int num_scale_blocks, void* __restrict__ quant_output, + void* __restrict__ scale_output) +{ + using ElemT = typename std::conditional::type; + constexpr int MAX_VEC = 16 / ELEM_BYTES; + constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); + constexpr int NTHRD = HEAD_DIM / VEC; + constexpr int VEC_BYTES = VEC * ELEM_BYTES; + using VecT = typename VecType::type; + + int const token_idx = blockIdx.x; + if (token_idx >= total_tokens) + return; + + // Per-token mask: precomputed on host, skips padded generation slots + // and batches that produced no compressed tokens. + if (!compressed_mask[token_idx]) + return; + + // ================================================================ + // Step 0: Find owning batch via binary search on cu_kv_comp. + // ================================================================ + int batch_idx, local_output_idx; + if (batch_size <= 1) + { + batch_idx = 0; + local_output_idx = token_idx; + } + else + { + int lo = 0, hi = batch_size; + while (lo < hi) + { + int mid = (lo + hi) >> 1; + if (cu_kv_comp[mid + 1] <= token_idx) + lo = mid + 1; + else + hi = mid; + } + batch_idx = lo; + if (batch_idx >= batch_size) + return; + local_output_idx = token_idx - cu_kv_comp[batch_idx]; + } + + if (local_output_idx >= num_outputs_arr[batch_idx]) + return; + + int const tid = threadIdx.x; + extern __shared__ float smem[]; + + // ================================================================ + // Step 1: Vectorized load from kv_comp + // ================================================================ + auto const* src = reinterpret_cast( + reinterpret_cast(kv_comp) + static_cast(token_idx) * HEAD_DIM); + VecT raw_in = src[tid]; + ElemT const* in_elems = reinterpret_cast(&raw_in); + + float v[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + v[i] = static_cast(in_elems[i]); + + // ================================================================ + // Step 2: RMSNorm + // ================================================================ + float local_sq = 0.f; +#pragma unroll + for (int i = 0; i < VEC; i++) + local_sq += v[i] * v[i]; + + float warp_sum = warpReduceSum(local_sq); + + constexpr int NUM_WARPS = (NTHRD + 31) / 32; + int const warp_id = tid / 32; + int const lane_id = tid % 32; + + if (lane_id == 0) + smem[warp_id] = warp_sum; + __syncthreads(); + + if (warp_id == 0) + { + float s = (lane_id < NUM_WARPS) ? smem[lane_id] : 0.0f; + for (int offset = 16; offset > 0; offset >>= 1) + s += __shfl_xor_sync(0xFFFFFFFF, s, offset); + if (lane_id == 0) + smem[0] = s; + } + __syncthreads(); + float const rms_scale = rsqrtf(smem[0] / static_cast(HEAD_DIM) + rms_eps); + + // ================================================================ + // Step 3: Load weight, apply RMSNorm + // ================================================================ + auto const* wt_src = reinterpret_cast(reinterpret_cast(rms_weight)); + VecT raw_w = wt_src[tid]; + ElemT const* w_elems = reinterpret_cast(&raw_w); + +#pragma unroll + for (int i = 0; i < VEC; i++) + v[i] = v[i] * rms_scale * static_cast(w_elems[i]); + + // ================================================================ + // Step 4: RoPE (Rotary Positional Embedding) + // + // Applied only to elements in [nope_dim, nope_dim+rope_dim). + // Uses interleaved even/odd pairs: (x_even, x_odd) → rotated by (cos, sin). + // cos_sin_table layout: [max_pos, rope_dim] where first half is cos, second is sin. + // ================================================================ + int const half_rope = rope_dim / 2; + int const pos_id = position_ids[token_idx]; + +#pragma unroll + for (int i = 0; i < VEC; i += 2) + { + int const elem_idx = tid * VEC + i; + if (elem_idx >= nope_dim) + { + int const rope_idx = elem_idx - nope_dim; + int const d = rope_idx >> 1; + float const cos_v = cos_sin_table[pos_id * rope_dim + d]; + float const sin_v = cos_sin_table[pos_id * rope_dim + half_rope + d]; + float const x_even = v[i]; + float const x_odd = v[i + 1]; + v[i] = x_even * cos_v - x_odd * sin_v; + v[i + 1] = x_odd * cos_v + x_even * sin_v; + } + } + + // ================================================================ + // Step 5: Hadamard butterfly transform (rotate activation) + // + // Implements the Walsh-Hadamard transform H_n * v via butterfly network. + // H_n has the recursive structure: H_n = [[H_{n/2}, H_{n/2}], [H_{n/2}, -H_{n/2}]] + // which decomposes into log2(HEAD_DIM) butterfly stages. + // + // Three phases handle increasing stride lengths: + // A) Local: strides < VEC — within each thread's register file + // B) Warp shuffle: strides VEC..32*VEC-1 — via __shfl_xor_sync + // C) Shared memory: strides >= 32*VEC — via XOR-swizzled smem + // + // The XOR swizzle pattern `idx ^ ((idx >> 3) & 0x1F)` ensures bank-conflict- + // free access to shared memory across all butterfly stride patterns. + // + // Skipped entirely when ROTATE_ACTIVATION=false. + // ================================================================ + if constexpr (ROTATE_ACTIVATION) + { + + // Phase A: local butterfly (strides 1..VEC-1, within each thread's VEC registers) +#pragma unroll + for (int stride = 1; stride < VEC; stride <<= 1) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + if ((i & stride) == 0) + { + float a = v[i], b = v[i ^ stride]; + v[i] = a + b; + v[i ^ stride] = a - b; + } + } + } + + // Phase B: warp shuffle butterfly (strides VEC..32*VEC-1, within a single warp) + if constexpr (NTHRD > 1) + { + constexpr int SHFL_END = (NTHRD <= 32) ? NTHRD : 32; +#pragma unroll + for (int ts = 1; ts < SHFL_END; ts <<= 1) + { + int const stride = ts * VEC; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + float partner = __shfl_xor_sync(0xFFFFFFFF, v[i], ts); + int const elem_idx = tid * VEC + i; + v[i] = (elem_idx & stride) ? (partner - v[i]) : (v[i] + partner); + } + } + } + + // Phase C: cross-warp butterfly via XOR-swizzled shared memory (strides >= 32*VEC) + // Only needed when NTHRD > 32 (i.e., multiple warps, e.g., HEAD_DIM=512, VEC=8 → 64 threads) + if constexpr (NTHRD > 32) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const idx = tid * VEC + i; + smem[idx ^ ((idx >> 3) & 0x1F)] = v[i]; + } + __syncthreads(); + + for (int stride = 32 * VEC; stride < HEAD_DIM; stride <<= 1) + { +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const idx = tid * VEC + i; + int const partner_idx = idx ^ stride; + float const a = smem[idx ^ ((idx >> 3) & 0x1F)]; + float const b = smem[partner_idx ^ ((partner_idx >> 3) & 0x1F)]; + v[i] = (idx & stride) ? (b - a) : (a + b); + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < VEC; i++) + { + int const idx = tid * VEC + i; + smem[idx ^ ((idx >> 3) & 0x1F)] = v[i]; + } + __syncthreads(); + } + } + + // ================================================================ + // Step 6: Scale by Hadamard factor + // ================================================================ + float const had_scale = rsqrtf(static_cast(HEAD_DIM)); + +#pragma unroll + for (int i = 0; i < VEC; i++) + v[i] *= had_scale; + + } // ROTATE_ACTIVATION + + // ================================================================ + // Step 7: Write postprocessed output to kv_out (if requested) + // ================================================================ + if (kv_out != nullptr) + { + VecT raw_out; + ElemT* out_elems = reinterpret_cast(&raw_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + out_elems[i] = static_cast(v[i]); + + auto* dst + = reinterpret_cast(reinterpret_cast(kv_out) + static_cast(token_idx) * HEAD_DIM); + dst[tid] = raw_out; + } + + // ================================================================ + // Step 9: Compute paged cache destination address. + // Map (batch_idx, local_output_idx) → logical block → physical block + // via the block table. block_base points to the start of the physical + // page; token_offset is the slot within that page. + // ================================================================ + int const start_pos = start_pos_arr[batch_idx]; + int const cache_pos = start_pos + local_output_idx; + int const logical_block = cache_pos / tokens_per_block; + int const token_offset = cache_pos % tokens_per_block; + int const phys_block = block_offsets[batch_idx * max_blocks + logical_block]; + + uint8_t* block_base + = reinterpret_cast(kv_cache) + static_cast(phys_block) * cache_stride_blk_bytes; + + // ================================================================ + // Step 11: Store to cache (compile-time dispatch on cache dtype/scale type) + // + // Cache addressing is byte-based: block_base points to the start of + // the physical block, cache_stride_blk_bytes is the total block size. + // ================================================================ + if constexpr (SCALE_TYPE == CacheScaleType::kNone) + { + // Default mode: float→bf16/fp32 pack + vectorized store. + // Cache layout per block: [tokens_per_block * HEAD_DIM] elements of ElemT. + VecT raw_out; + ElemT* out_elems = reinterpret_cast(&raw_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + out_elems[i] = static_cast(v[i]); + + ElemT* row_base = reinterpret_cast(block_base) + token_offset * HEAD_DIM; + reinterpret_cast(row_base)[tid] = raw_out; + } + else if constexpr (SCALE_TYPE == CacheScaleType::kFP8PerTensor) + { + // FP8 per-tensor: direct float→fp8_e4m3fn cast (implicit scale=1.0). + // Cache layout per block: [tokens_per_block * HEAD_DIM] bytes of fp8. + uint8_t fp8_bytes[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + __nv_fp8_e4m3 fp8_val(v[i]); + fp8_bytes[i] = *reinterpret_cast(&fp8_val); + } + + using Fp8VecT = typename VecType::type; + uint8_t* fp8_dst = block_base + token_offset * HEAD_DIM + tid * VEC; + *reinterpret_cast(fp8_dst) = *reinterpret_cast(fp8_bytes); + } + else if constexpr (SCALE_TYPE == CacheScaleType::kFP8Blockwise) + { + // FP8 blockwise: per-128-element quantization with explicit scales. + // GROUP_SIZE = number of threads that share one scale factor. + // For HD=512, VEC=8: GROUP_SIZE=16 threads → 128 elements per scale block. + // + // Cache layout per block: + // [fp8_data: tokens_per_block * HEAD_DIM bytes] + // [scales: tokens_per_block * (HEAD_DIM/128) * 4 bytes] + constexpr int GROUP_SIZE = 128 / VEC; + + // Step 11a: Compute per-group amax via warp shuffle reduction. + // GROUP_SIZE <= 16 (< warp), so shuffle is sufficient (no smem needed). + float local_amax = 0.f; +#pragma unroll + for (int i = 0; i < VEC; i++) + local_amax = fmaxf(local_amax, fabsf(v[i])); + +#pragma unroll + for (int offset = GROUP_SIZE / 2; offset > 0; offset >>= 1) + local_amax = fmaxf(local_amax, __shfl_xor_sync(0xFFFFFFFF, local_amax, offset)); + + // Step 11b: Compute scale and inverse scale for quantization. + // 448.0 is the max representable value for fp8_e4m3fn. + float const scale = local_amax / 448.0f; + float const inv_scale = (local_amax > 0.f) ? (448.0f / local_amax) : 1.0f; + + // Step 11c: Quantize to FP8 and store data. + uint8_t fp8_bytes[VEC]; +#pragma unroll + for (int i = 0; i < VEC; i++) + { + __nv_fp8_e4m3 fp8_val(v[i] * inv_scale); + fp8_bytes[i] = *reinterpret_cast(&fp8_val); + } + + using Fp8VecT = typename VecType::type; + uint8_t* fp8_dst = block_base + token_offset * HEAD_DIM + tid * VEC; + *reinterpret_cast(fp8_dst) = *reinterpret_cast(fp8_bytes); + + // Step 11d: Store scale factor (one thread per 128-element group writes it). + if (tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + float* scale_dst = reinterpret_cast(block_base + tokens_per_block * HEAD_DIM + + (token_offset * num_scale_blocks + scale_idx) * sizeof(float)); + *scale_dst = scale; + } + + // Step 11e: Optionally write FP8 data and scales to output buffers. + // Used by the indexer compressor which returns (fp8_data, scales) to Python + // for downstream sparse attention indexing. + if (quant_output != nullptr) + { + uint8_t* fp8_out_dst + = reinterpret_cast(quant_output) + static_cast(token_idx) * HEAD_DIM + tid * VEC; + *reinterpret_cast(fp8_out_dst) = *reinterpret_cast(fp8_bytes); + } + if (scale_output != nullptr && tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + reinterpret_cast(scale_output)[static_cast(token_idx) * num_scale_blocks + scale_idx] + = scale; + } + } + else if constexpr (SCALE_TYPE == CacheScaleType::kMXFP4Blockwise) + { + constexpr int GROUP_SIZE = 32 / VEC; + constexpr int PACKED_VEC_BYTES = VEC / 2; + constexpr float kFp4Max = 6.0f; + constexpr float kFp4MaxInv = 1.0f / kFp4Max; + constexpr float kFp4MinAmax = kFp4Max * 1.1754943508222875e-38f; + + float local_amax = 0.f; +#pragma unroll + for (int i = 0; i < VEC; i++) + local_amax = fmaxf(local_amax, fabsf(v[i])); + +#pragma unroll + for (int offset = GROUP_SIZE / 2; offset > 0; offset >>= 1) + local_amax = fmaxf(local_amax, __shfl_xor_sync(0xFFFFFFFF, local_amax, offset)); + + float const scale = roundedPow2Scale(local_amax, kFp4MaxInv, kFp4MinAmax); + + uint8_t fp4_bytes[PACKED_VEC_BYTES]; +#pragma unroll + for (int i = 0; i < VEC; i += 2) + { + fp4_bytes[i / 2] = packE2M1x2(v[i] / scale, v[i + 1] / scale); + } + + int const packed_head_dim = HEAD_DIM / 2; + uint8_t* fp4_dst = block_base + token_offset * packed_head_dim + tid * PACKED_VEC_BYTES; +#pragma unroll + for (int i = 0; i < PACKED_VEC_BYTES; ++i) + fp4_dst[i] = fp4_bytes[i]; + + if (tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + uint8_t* scale_dst + = block_base + tokens_per_block * packed_head_dim + token_offset * num_scale_blocks + scale_idx; + *scale_dst = toUe8m0(scale); + } + + if (quant_output != nullptr) + { + uint8_t* fp4_out_dst = reinterpret_cast(quant_output) + + static_cast(token_idx) * packed_head_dim + tid * PACKED_VEC_BYTES; +#pragma unroll + for (int i = 0; i < PACKED_VEC_BYTES; ++i) + fp4_out_dst[i] = fp4_bytes[i]; + } + if (scale_output != nullptr && tid % GROUP_SIZE == 0) + { + int const scale_idx = tid / GROUP_SIZE; + reinterpret_cast(scale_output)[static_cast(token_idx) * num_scale_blocks + scale_idx] + = toUe8m0(scale); + } + } +} + +// Explicit instantiations — fused postprocess+scatter. +// kNone supports bf16 (EB=2) and fp32 (EB=4) input types; the quantized +// scale types only support bf16 input since the compressor output is bf16. +// Each combination is instantiated with ROTATE_ACTIVATION=true and false. +#define INST_PPS(HD, EB, CST, AR) \ + template __global__ void postProcessScatterKernel(void const*, void*, void const*, float, \ + float const*, int32_t const*, int, int, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, \ + bool const*, int, int, int, int, int, int, void*, void*); + +#define INST_PPS_AR(HD, EB, CST) \ + INST_PPS(HD, EB, CST, true) \ + INST_PPS(HD, EB, CST, false) + +INST_PPS_AR(128, 2, CacheScaleType::kNone) +INST_PPS_AR(128, 4, CacheScaleType::kNone) +INST_PPS_AR(512, 2, CacheScaleType::kNone) +INST_PPS_AR(512, 4, CacheScaleType::kNone) +INST_PPS_AR(128, 2, CacheScaleType::kFP8PerTensor) +INST_PPS_AR(512, 2, CacheScaleType::kFP8PerTensor) +INST_PPS_AR(128, 2, CacheScaleType::kFP8Blockwise) +INST_PPS_AR(512, 2, CacheScaleType::kFP8Blockwise) +INST_PPS_AR(128, 2, CacheScaleType::kMXFP4Blockwise) +INST_PPS_AR(512, 2, CacheScaleType::kMXFP4Blockwise) +#undef INST_PPS_AR +#undef INST_PPS + +// ============================================================================ +// Postprocess + Scatter Launch Wrapper +// +// Derives cache layout parameters (cache_stride_blk_bytes, num_scale_blocks) +// from cache dtype / scale type, then dispatches to the appropriate template instantiation. +// ============================================================================ + +// Compute number of threads per block, mirroring the compile-time VEC/NTHRD logic. +// Ensures NTHRD >= 32 by reducing VEC when HEAD_DIM is small. +static inline int compressorNthreads(int head_dim, int elem_bytes) +{ + int max_vec = 16 / elem_bytes; + int vec = (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); + return head_dim / vec; +} + +void postProcessScatterLaunch(void const* kv_comp, void* kv_out, void const* rms_weight, float rms_eps, + float const* cos_sin_table, int32_t const* position_ids, int nope_dim, int rope_dim, void* kv_cache, + int32_t const* num_outputs, int32_t const* cu_kv_comp, int32_t const* start_pos, int32_t const* block_offsets, + bool const* compressed_mask, int batch_size, int tokens_per_block, int head_dim, int max_blocks_per_seq, + int elem_bytes, int total_tokens, int cache_scale_type, bool rotate_activation, void* quant_output, + void* scale_output, cudaStream_t stream) +{ + if (total_tokens == 0) + { + return; + } + + TLLM_CHECK_WITH_INFO( + cache_scale_type >= 0 && cache_scale_type <= 3, "Invalid cache_scale_type: %d", cache_scale_type); + auto const cst = static_cast(cache_scale_type); + + bool const is_quantized_store = (cst != CacheScaleType::kNone); + int const nthreads = compressorNthreads(head_dim, elem_bytes); + int const smem_bytes = head_dim * sizeof(float); + TLLM_CHECK_WITH_INFO(cst != CacheScaleType::kMXFP4Blockwise || head_dim % 32 == 0, + "MXFP4 cache requires head_dim divisible by 32, got %d", head_dim); + TLLM_CHECK_WITH_INFO(cst != CacheScaleType::kMXFP4Blockwise || head_dim % 2 == 0, + "FP4 packed cache requires even head_dim, got %d", head_dim); + TLLM_CHECK_WITH_INFO(!is_quantized_store || elem_bytes == 2, + "Quantized cache modes require bf16 compressor output, got elem_bytes=%d", elem_bytes); + + // Derive cache block stride (in bytes) and scale block count from the + // scale type. Each physical cache block stores tokens_per_block tokens: + // none: tpb * HD * elem_bytes + // fp8 pertensor: tpb * HD + // fp8 blockwise: tpb * HD + tpb * (HD/128)*4 + // mxfp4: tpb * (HD/2) + tpb * (HD/32) + int num_scale_blocks = 0; + int cache_stride_blk_bytes = 0; + switch (cst) + { + case CacheScaleType::kFP8PerTensor: cache_stride_blk_bytes = tokens_per_block * head_dim; break; + case CacheScaleType::kFP8Blockwise: + num_scale_blocks = head_dim / 128; + cache_stride_blk_bytes = tokens_per_block * head_dim + tokens_per_block * num_scale_blocks * 4; + break; + case CacheScaleType::kMXFP4Blockwise: + num_scale_blocks = head_dim / 32; + cache_stride_blk_bytes = tokens_per_block * (head_dim / 2) + tokens_per_block * num_scale_blocks; + break; + default: cache_stride_blk_bytes = tokens_per_block * head_dim * elem_bytes; break; + } + +#define LAUNCH_PPS(HD, EB, CST, AR) \ + postProcessScatterKernel<<>>(kv_comp, kv_out, \ + rms_weight, rms_eps, cos_sin_table, position_ids, nope_dim, rope_dim, kv_cache, num_outputs, cu_kv_comp, \ + start_pos, block_offsets, compressed_mask, batch_size, tokens_per_block, max_blocks_per_seq, \ + cache_stride_blk_bytes, total_tokens, num_scale_blocks, quant_output, scale_output) + +#define DISPATCH_ROTATE(HD, EB, CST) \ + if (rotate_activation) \ + { \ + LAUNCH_PPS(HD, EB, CST, true); \ + } \ + else \ + { \ + LAUNCH_PPS(HD, EB, CST, false); \ + } +#define DISPATCH_HD_EB(CST) \ + if (elem_bytes == 4) \ + { \ + switch (head_dim) \ + { \ + case 128: DISPATCH_ROTATE(128, 4, CST); break; \ + default: DISPATCH_ROTATE(512, 4, CST); break; \ + } \ + } \ + else \ + { \ + switch (head_dim) \ + { \ + case 128: DISPATCH_ROTATE(128, 2, CST); break; \ + default: DISPATCH_ROTATE(512, 2, CST); break; \ + } \ + } +#define DISPATCH_HD_BF16(CST) \ + switch (head_dim) \ + { \ + case 128: DISPATCH_ROTATE(128, 2, CST); break; \ + default: DISPATCH_ROTATE(512, 2, CST); break; \ + } + + if (cst == CacheScaleType::kFP8PerTensor) + { + DISPATCH_HD_BF16(CacheScaleType::kFP8PerTensor); + } + else if (cst == CacheScaleType::kFP8Blockwise) + { + DISPATCH_HD_BF16(CacheScaleType::kFP8Blockwise); + } + else if (cst == CacheScaleType::kMXFP4Blockwise) + { + DISPATCH_HD_BF16(CacheScaleType::kMXFP4Blockwise); + } + else + { + DISPATCH_HD_EB(CacheScaleType::kNone); + } + +#undef DISPATCH_HD_BF16 +#undef DISPATCH_HD_EB +#undef DISPATCH_ROTATE +#undef LAUNCH_PPS +} + +} // namespace kernels::compressor + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h new file mode 100644 index 000000000000..2f5ea155447e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h @@ -0,0 +1,99 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ +#pragma once + +#include +#include +#include +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::compressor +{ + +// Decode kernel: write NEXT_N tokens to paged cache + conditional compression +// via online softmax. Overlap is derived from compress_ratio (ratio=4). +// +// Grid: (batch_size, cdiv(state_dim, block_size)) +// One thread per state_dim element across all phases. +// state_dim is a constexpr derived from COMPRESS_RATIO and HEAD_DIM inside the kernel. +void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or fp32) + float const* ape, // [compress_ratio, state_dim] + void* paged_kv, // [num_blocks, page_size, state_dim] + void* paged_score, // [num_blocks, page_size, state_dim] + int32_t const* block_table_kv, // [bsz, max_blocks] + int32_t const* block_table_score, // [bsz, max_blocks] + void* output, // [total_outputs, head_dim] + int32_t const* kv_lens, // [bsz] + int32_t const* start_pos, // [bsz] + int32_t const* cu_seq_lens, // [bsz+1] + int32_t const* cu_kv_comp, // [bsz+1] + int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int next_n, + int io_elem_bytes, // bytes per element for kv_score/paged (2=bf16, 4=fp32) + int out_elem_bytes, // bytes per element for output + cudaStream_t stream); + +// Prefill kernel: bulk compression with per-token gather/scatter + state update. +// Writes remainder tokens to paged cache, then performs online softmax reduction. +// +// Grid: (batch_size, max_outputs_per_batch, num_head_chunks) +// Each block computes one compressed output for one head_dim chunk. +void prefillReductionLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or fp32) + float const* ape, // [compress_ratio, state_dim] + void* paged_kv, // [num_blocks, page_size, state_dim] + void* paged_score, // [num_blocks, page_size, state_dim] + int32_t const* block_table_kv, // [bsz, max_blocks] + int32_t const* block_table_score, // [bsz, max_blocks] + void* output, // [total_outputs, head_dim] + int32_t const* kv_lens, // [bsz] + int32_t const* start_pos, // [bsz] + int32_t const* cu_seq_lens, // [bsz+1] + int32_t const* cu_kv_comp, // [bsz+1] + int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int max_outputs, int io_elem_bytes, + int out_elem_bytes, cudaStream_t stream); + +// RMSNorm + RoPE + Hadamard + paged scatter in a single kernel launch. +// Optionally writes postprocessed result to kv_out (nullptr to skip). +// +// Grid: (total_tokens) -- one block per compressed token +// Block: (head_dim / VEC), always >= 32 +void postProcessScatterLaunch(void const* kv_comp, // [total_tokens, head_dim] input + void* kv_out, // [total_tokens, head_dim] postprocessed output (nullptr to skip) + void const* rms_weight, // [head_dim] + float rms_eps, + float const* cos_sin_table, // [max_pos, 2, rope_dim/2] + int32_t const* position_ids, // [total_tokens] + int nope_dim, int rope_dim, + void* kv_cache, // paged cache buffer + int32_t const* num_outputs, // [bsz] + int32_t const* cu_kv_comp, // [bsz+1] + int32_t const* start_pos, // [bsz] + int32_t const* block_offsets, // [bsz, max_blocks] + bool const* compressed_mask, // [total_tokens] — per-token mask, false ⇒ skip + int batch_size, int tokens_per_block, int head_dim, int max_blocks_per_seq, int elem_bytes, int total_tokens, + int cache_scale_type, // 0=none (bf16/fp32 by elem_bytes), 1=fp8_pertensor, + // 2=fp8_blockwise, 3=mxfp4 (packed FP4) + bool rotate_activation, // whether to apply Hadamard transform (false to skip) + void* quant_output, // optional fp8/fp4 packed output (nullptr if unused) + void* scale_output, // optional scale output (float* for fp8, uint8_t* for fp4) + cudaStream_t stream); + +} // namespace kernels::compressor + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu index 58ccb4ac8ea8..6c2c32cf4be4 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu @@ -15,6 +15,7 @@ */ #include "moeTopKFuncs.cuh" +#include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaTypeUtils.cuh" #include "tensorrt_llm/common/envUtils.h" @@ -309,6 +310,177 @@ INSTANTIATE_RENORM_MOE_ROUTING(half, __nv_bfloat16, int32_t, true); INSTANTIATE_RENORM_MOE_ROUTING(__nv_bfloat16, __nv_bfloat16, int32_t, true); #endif +static constexpr int kTOPK = 6; + +// CUDA kernel for gate forward +// Input: pre-computed scores from linear(x, weight) done outside kernel +// Template parameters: +// nExperts: number of experts +// topK: number of top experts to select +// hash: true for hash mode, false for topk mode +// One warp per row (batch element) +template +__global__ void gate_forward_kernel( + float const* __restrict__ scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + float const* __restrict__ bias, // [nExperts] (only used when hash=false) + int const* __restrict__ input_ids, // [batch_size] (only used when hash=true) + int const* __restrict__ tid2eid, // [vocab_size, topK] (only used when hash=true) + float* __restrict__ out_weights, // [batch_size, topK] + int* __restrict__ out_indices, // [batch_size, topK] + int batch_size, float route_scale) +{ + // Compile-time constants + constexpr int kExpertsPerThread = nExperts / WARP_SIZE; + constexpr int kWarpsPerBlock = 4; // Adjust based on occupancy needs + + // Shared memory for original scores (one array per warp in the block) + __shared__ float smem_scores[kWarpsPerBlock][nExperts]; + + // One warp per batch element + int const global_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE; + int const local_warp_id = (threadIdx.x / WARP_SIZE) % kWarpsPerBlock; + int const lane_id = threadIdx.x % WARP_SIZE; + + if (global_warp_id >= batch_size) + return; + + auto warp = cg::tiled_partition(cg::this_thread_block()); + + // Pointer to this warp's shared memory and input scores + float* my_smem = smem_scores[local_warp_id]; + float const* scores_row = scores_in + global_warp_id * nExperts; + +// Load scores, apply score function (softplus + sqrt), and store to shared memory +#pragma unroll + for (int e = 0; e < kExpertsPerThread; ++e) + { + int expert_id = lane_id + e * WARP_SIZE; + float s = scores_row[expert_id]; + float sp = log1pf(expf(s)); + float score = sqrtf(sp); + my_smem[expert_id] = score; // Store original score to shared memory + } + __syncwarp(); // Ensure all scores are written before reading + + // Output: each of first K lanes holds one value + float my_topk_value = 0.0f; + int my_topk_index = 0; + + if constexpr (hash) + { + // Hash mode: directly read from shared memory + int token_id = input_ids[global_warp_id]; + int const* expert_ids = tid2eid + token_id * topK; + + if (lane_id < topK) + { + int expert_id = expert_ids[lane_id]; + my_topk_index = expert_id; + my_topk_value = my_smem[expert_id]; // Direct lookup from shared memory + } + } + else + { + // Topk mode: load from shared memory, add bias to registers for topk + float scores[kExpertsPerThread]; + int indices[kExpertsPerThread]; + +#pragma unroll + for (int e = 0; e < kExpertsPerThread; ++e) + { + int expert_id = lane_id + e * WARP_SIZE; + indices[e] = expert_id; + scores[e] = my_smem[expert_id] + bias[expert_id]; // Add bias for topk selection + } + + // Use reduceTopK to find top-k experts + float topk_values[topK]; + int32_t topk_indices[topK]; + constexpr float minValue = -1e30f; + reduce_topk::reduceTopK( + warp, topk_values, topk_indices, scores, indices, minValue, topK); + + // Gather original weights (without bias) from shared memory + if (lane_id < topK) + { + int expert_id = topk_indices[lane_id]; + my_topk_index = expert_id; + my_topk_value = my_smem[expert_id]; // Read original score (no bias) + } + } + + // Reduce to get sum (first K lanes have values, others have 0) + float weight_sum = cg::reduce(warp, my_topk_value, cg::plus{}); + + // Normalize weights and write output (first K lanes) + if (lane_id < topK) + { + out_weights[global_warp_id * topK + lane_id] = (my_topk_value / weight_sum) * route_scale; + out_indices[global_warp_id * topK + lane_id] = my_topk_index; + } +} + +// C++ wrapper function (output tensors passed as parameters) +// All tensors are float32 +template +void launch_gate_forward_kernel(float* scores_in, float* bias, int* input_ids, int* tid2eid, float* out_weights, + int* out_indices, int batch_size, float route_scale, cudaStream_t stream) +{ + constexpr int warps_per_block = 4; + constexpr int threads_per_block = warps_per_block * WARP_SIZE; + int const blocks = (batch_size + warps_per_block - 1) / warps_per_block; + + gate_forward_kernel<<>>( + scores_in, bias, input_ids, tid2eid, out_weights, out_indices, batch_size, route_scale); +} + +void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + void* bias, // nullptr if hash mode + void* input_ids, // nullptr if non-hash mode + void* tid2eid, // nullptr if non-hash mode + void* out_weights, // [batch_size, topK] - pre-allocated + void* out_indices, // [batch_size, topK] - pre-allocated + int batch_size, int n_experts, float route_scale, bool is_hash, cudaStream_t stream) +{ + auto* scores = static_cast(scores_in); + auto* bias_ptr = static_cast(bias); + auto* input_ids_ptr = static_cast(input_ids); + auto* tid2eid_ptr = static_cast(tid2eid); + auto* weights = static_cast(out_weights); + auto* indices = static_cast(out_indices); + + switch (n_experts) + { + case 256: + if (is_hash) + { + launch_gate_forward_kernel<256, true>( + scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, route_scale, stream); + } + else + { + launch_gate_forward_kernel<256, false>( + scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); + } + break; + case 384: + if (is_hash) + { + launch_gate_forward_kernel<384, true>( + scores, nullptr, input_ids_ptr, tid2eid_ptr, weights, indices, batch_size, route_scale, stream); + } + else + { + launch_gate_forward_kernel<384, false>( + scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); + } + break; + default: + TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384"); + } + sync_check_cuda_error(stream); +} + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h index f8240b436393..367676fcc334 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -29,6 +29,16 @@ namespace kernels template void invokeCustomMoeRouting(InputT* routerLogits, OutputT* topkValues, IdxT* topkIndices, int64_t const numTokens, int64_t const numExperts, int64_t const topK, cudaStream_t const stream); + +// Gate forward function for custom MoE routing +// All tensors are expected to be float32 for scores/weights, int32 for indices +void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + void* bias, // nullptr if hash mode + void* input_ids, // nullptr if non-hash mode + void* tid2eid, // nullptr if non-hash mode + void* out_weights, // [batch_size, topK] - pre-allocated + void* out_indices, // [batch_size, topK] - pre-allocated + int batch_size, int n_experts, float route_scale, bool is_hash, cudaStream_t stream); } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index ad1e728b3298..5ecea6ce5537 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -271,18 +271,13 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i // The threshold bin. thresholdBinIdx = smemThresholdBinIdx[0]; - // Skip auto-promote at step 0 when we'll continue: half-precision bins - // don't align with step 2's full-precision bit-pattern filter, so a - // step-0 promote would be double-counted at step 2. - bool const step0WillContinue = (step == 0) && (smemFinalBinSize[0] > kNumFinalItems); - auto processBins = [&](InputT logitIn, int idx) { float const logit = static_cast(logitIn); if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); - if (binIdx < thresholdBinIdx && !step0WillContinue) + if (binIdx < thresholdBinIdx) { // The element is part of the top-k selection int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1); @@ -368,21 +363,28 @@ __device__ bool processHistogramStep(int const* indices, InputT const* logits, i return smemFinalBinSize[0] > kNumFinalItems; } -// Smem holder for topKPerRowJob's final-sort. Always reserves -// BlockRadixSort::TempStorage so the sort algorithm can be picked at runtime -// (the union's size is dominated by FinalItems = 16 KB at our shapes). -template -struct TopKSmem +// Follows half - 11 - 11 - 10 bit iterations +template +static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, + int* outIndices, float* outLogits, int stride1, int topK) { + // The number of slots for the final pass. static constexpr int kNumFinalItems = 2048; + // The number of elements per thread for the final sort. static constexpr int kNumFinalItemsPerThread = kNumFinalItems / kNumThreadsPerBlock; + // The class to sort the elements during the final pass. using FinalSort = cub::BlockRadixSort; - using FinalSortTempStorage = typename FinalSort::TempStorage; + using FinalSortTempStorage = std::conditional_t; + // The class to compute the inclusive prefix-sum over the histogram. using Scan = cub::BlockScan; + // The structure to store the final items (for the final pass). struct FinalItems { + // Shared memory to store the indices for the final pass. int indices[kNumFinalItems]; + // Shared memory to store the logits for the final pass. float logits[kNumFinalItems]; }; @@ -392,47 +394,34 @@ struct TopKSmem int data[kNumBins]; }; - union Final + // Shared memory to compute the block sort. + __shared__ union { FinalItems items; FinalSortTempStorage finalSort; Histogram histo; - }; - - Final smemFinal; - int smemThresholdBinIdx[1]; - int smemFinalDstIdx[1]; - int smemFinalBinSize[1]; - int smemFoundTopKValues[1]; -}; - -// Follows half - 11 - 11 - 10 bit iterations -template -static __device__ void topKPerRowJob(int const* indices, InputT const* logits, int rowStart, int rowEnd, - int* outIndices, float* outLogits, int stride1, int topK, TopKSmem& smem) -{ - static constexpr int kNumFinalItems = TopKSmem::kNumFinalItems; - static constexpr int kNumFinalItemsPerThread = TopKSmem::kNumFinalItemsPerThread; - using FinalSort = typename TopKSmem::FinalSort; - - auto& smemFinal = smem.smemFinal; - int* smemThresholdBinIdx = smem.smemThresholdBinIdx; - int* smemFinalDstIdx = smem.smemFinalDstIdx; - int* smemFinalBinSize = smem.smemFinalBinSize; - int* smemFoundTopKValues = smem.smemFoundTopKValues; + } smemFinal; // Shared memory to store the selected indices. // If we are processing using multiple blocks, we need to store the logits and // indices. extern __shared__ int32_t smemOutput[]; + // Shared memory to store the threshold bin. + __shared__ int smemThresholdBinIdx[1]; + // Shared memory counter to register the candidates for the final phase. + __shared__ int smemFinalDstIdx[1]; + // Shared memory to determine if the threshold bin fits in the final items. + __shared__ int smemFinalBinSize[1]; + // Shared memory to keep track of the top-k values found so far by the + // previous iterations + __shared__ int smemFoundTopKValues[1]; + // The length of the row. int rowLen = rowEnd - rowStart; // Shortcut if the length of the row is smaller than Top-K. Indices are not - // sorted by their corresponding logit. Unreachable when mergeBlocks=true: - // both merge callers pass rowLen = numBlocksPerRow * topK > topK. + // sorted by their corresponding logit. if (rowLen <= topK) { for (int rowIt = threadIdx.x; rowIt < rowLen; rowIt += kNumThreadsPerBlock) @@ -502,12 +491,10 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i if (!continueToNextStep) { - // Sort the threshold-bin candidates. Insertion sort wins below ~512 - // items (O(n^2/T) with no fixed cost); BlockRadixSort wins above - // (constant cost padded to kNumFinalItems = 2048). - constexpr int kInsertionSortBranchThreshold = 512; - int const finalCount = smemFinalDstIdx[0]; - if (finalCount > kInsertionSortBranchThreshold) + // The histogram did not proceed to the final 10 bits, therefore we need to + // sort the final items The logits of the elements to be sorted in the final + // pass. + if constexpr (useRadixSort) { // Sorting with radix sort float finalLogits[kNumFinalItemsPerThread]; @@ -525,7 +512,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { int srcIdx = ii * kNumThreadsPerBlock + threadIdx.x; - if (srcIdx < finalCount) + if (srcIdx < smemFinalDstIdx[0]) { finalLogits[ii] = smemFinal.items.logits[srcIdx]; finalIndices[ii] = smemFinal.items.indices[srcIdx]; @@ -558,13 +545,13 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } else { - // Sorting with insertion sort. + // Sorting with insertion sort auto baseIdx = smemFoundTopKValues[0]; - for (int i = threadIdx.x; i < finalCount; i += kNumThreadsPerBlock) + for (int i = threadIdx.x; i < smemFinalDstIdx[0]; i += kNumThreadsPerBlock) { int outIndex = 0; auto logit = smemFinal.items.logits[i]; - for (int j = 0; j < finalCount; j++) + for (int j = 0; j < smemFinalDstIdx[0]; j++) { auto otherLogit = smemFinal.items.logits[j]; if (logit < otherLogit || (logit == otherLogit && i < j)) @@ -572,6 +559,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndex++; } } + // Store if outIndex is in bounds if (outIndex + baseIdx < topK) { smemOutput[outIndex + baseIdx] = smemFinal.items.indices[i]; @@ -593,10 +581,6 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i outIndices[i] = smemOutput[i]; outLogits[i] = reinterpret_cast(smemOutput + topK)[i]; } - else if constexpr (mergeBlocks) - { - outIndices[i] = smemOutput[i]; - } else { if (stride1 == 1) @@ -613,7 +597,7 @@ static __device__ void topKPerRowJob(int const* indices, InputT const* logits, i } } // namespace -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill(InputT const* logits, int const* rowStarts, int const* rowEnds, int* outIndices, int stride0, int stride1, int const topK, int const offsetIndex) @@ -635,17 +619,18 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( outIndices += static_cast(rowIdx) * topK; logits += static_cast(rowIdx) * stride0; - __shared__ TopKSmem smem; - topKPerRowJob( - nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK, smem); + topKPerRowJob( + nullptr, logits, rowStart, rowEnd, outIndices, nullptr, stride1, topK); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif } -template +template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(InputT const* logits, int const* seqLens, - int* outIndices, int stride0, int stride1, int const topK, int next_n, float* outLogits = nullptr, + int* outIndices, int stride0, int stride1, int const topK, int next_n, int compressRatio, + float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) @@ -660,7 +645,8 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I // The range of logits within the row. int rowStart = 0; int seq_len = seqLens[rowIdx / next_n]; - int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1; + int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + int rowEnd = actual_kv_len / compressRatio; // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) @@ -683,9 +669,8 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I } logits += static_cast(rowIdx) * stride0; - __shared__ TopKSmem smem; - topKPerRowJob( - indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK, smem); + topKPerRowJob( + indices, logits, rowStart, rowEnd, outIndices, outLogits, stride1, topK); #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaTriggerProgrammaticLaunchCompletion(); #endif @@ -694,415 +679,10 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I namespace { -// Multi-pass radix: 4 launches (half-precision top-11 bits, then float bits -// [21..32), [10..21), [0..10)). Per-row state lives in DRAM scratch; candidate -// data uses two ping-pong buffers. Pass 3's last block emits the final top-K. -static constexpr int kRadixBins = 2048; - -// Per-row scratch state. The last block of each pass (picked via -// `finishedBlocks`) scans the global histogram and writes the next pass's -// threshold. -struct alignas(64) RadixState -{ - int candCount; // candidates entering current pass - int outIdx; // running outIndices write position - int kRemaining; // topK - outIdx - int filterCnt; // running candidate-buf write position - int thresholdBin; // prior pass's threshold; overwritten this pass's last block - int finishedBlocks; // last-block atomic, reset between passes - int thresholdLess; // count of bins below thresholdBin; pass-3 emit routes - // ties (bin == threshold) into a disjoint slot range. - int padding[1]; -}; - -// Common last-block trailer: prefix-scan the merged global histogram, locate -// the bin where the running count crosses kRemaining, stash to state for the -// next pass. step < 3 also resets the per-row global histogram. In step 1 -// the trailer writes st.candCount = full row length for pass 2 to consume -// (pass 1 reads the length inline from seqLens, so no init kernel is needed). -template -__device__ __forceinline__ void radixLastBlockTrailer(int* gHist, RadixState& st, int topK, int rowFullLen) -{ - using Scan = cub::BlockScan; - __shared__ typename Scan::TempStorage scanStorage; - __shared__ int s_thresholdBin; - __shared__ int s_runningBefore; - __shared__ int s_thresholdCount; - - if (threadIdx.x == 0) - { - s_thresholdBin = -1; - s_runningBefore = 0; - s_thresholdCount = 0; - } - __syncthreads(); - - // kRemaining for THIS pass: - // step 1: topK (no auto-promotes have happened yet). - // step 2/3: topK - outIdx (where outIdx was atomic-incremented during - // the filter loop). The histogram for the next pass's - // threshold pick must target this fresh value, not the - // stale state.kRemaining left over from the previous pass. - int const kRem = (step == 1) ? topK : (topK - st.outIdx); - if (kRem <= 0) - { - // All top-k slots already filled by auto-promotes in this pass — - // no need for a next pass to emit anything. - if (threadIdx.x == 0) - { - st.thresholdBin = -1; - st.candCount = 0; - st.kRemaining = 0; - st.finishedBlocks = 0; - st.filterCnt = 0; - (void) rowFullLen; - } - if constexpr (step < 3) - { - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - } - return; - } - constexpr int kRoundsPerScan = kRadixBins / kThreads; - int running = 0; - for (int r = 0; r < kRoundsPerScan; ++r) - { - int bin = r * kThreads + threadIdx.x; - int c = gHist[bin]; - int prefix, total; - Scan(scanStorage).ExclusiveSum(c, prefix, total); - prefix += running; - int next = prefix + c; - if (prefix < kRem && next >= kRem && s_thresholdBin == -1) - { - atomicCAS(&s_thresholdBin, -1, bin); - if (s_thresholdBin == bin) - { - s_runningBefore = prefix; - s_thresholdCount = c; - } - } - running += total; - __syncthreads(); - if (s_thresholdBin != -1) - break; - } - __syncthreads(); - - if (threadIdx.x == 0) - { - // If the cumsum over the whole histogram never reached kRem the row - // has fewer items than we still need (e.g. decode rows shorter than - // topK). Set thresholdBin to a sentinel above any valid bin so the - // next pass's `bin < thresholdBin` test accepts every surviving - // candidate as auto-promote. - st.thresholdBin = (s_thresholdBin == -1) ? kRadixBins : s_thresholdBin; - // s_runningBefore is the count of histogram items in bins < threshold. - // In the sentinel case it stays at its init 0; that's fine because - // pass-3's inline final-emit only uses thresholdLess to position the - // ties (bin == threshold) write base, and the sentinel branch has - // no items in the threshold bin (the sentinel is above all valid bins). - st.thresholdLess = s_runningBefore; - st.finishedBlocks = 0; - if constexpr (step == 1) - { - // Pass 2 also scans the full row from `logits`, so it reads - // st.candCount = full row length. Pass 1 did not write to it - // and the state struct started at zero (cudaMemsetAsync). - st.candCount = rowFullLen; - (void) s_thresholdCount; - (void) kRem; - } - else - { - st.candCount = st.filterCnt; - int newKRem = topK - st.outIdx; - if (newKRem < 0) - newKRem = 0; - st.kRemaining = newKRem; - st.filterCnt = 0; - } - } - if constexpr (step < 3) - { - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - } -} - -// Fused histogram + filter pass kernel. -// -// step == 1: pass 1, no filter, just histogram top-11 bits of the row. -// step == 2: pass 2 reads `logits` (pass 1 didn't write a candidate buffer), -// for each item: -// bin1 < thresholdBin1 → write to outIndices (auto-promote) -// bin1 == thresholdBin1 → append to candBufOut, count its bin2 -// in the histogram for pass 2 -// else → drop -// step == 3: same as step 2 but reads `candBufIn` (pass 2's output) and -// uses extractBinIdx<2> for the prior-bits check, extractBinIdx<3> -// for the histogram. -// -// Last block of every pass runs `radixLastBlockTrailer` to compute the -// next pass's threshold and reset cross-pass state. -template -static __global__ __launch_bounds__(kThreads) void radixPassKernel(InputT const* logits, int const* seqLens, - int* outIndices, int const* candBufIn, int* candBufOut, int* histograms, RadixState* state, int stride0, int next_n, - int topK) -{ - int rowIdx = blockIdx.y; - int blockInRow = blockIdx.x; - int blocksPerRow = gridDim.x; - - RadixState& st = state[rowIdx]; - int* gHist = histograms + static_cast(rowIdx) * kRadixBins; - - __shared__ int sHist[kRadixBins]; - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - sHist[i] = 0; - __syncthreads(); - - if constexpr (step == 1) - { - // Read seqLen inline so pass 1 does not depend on st.candCount being - // pre-initialised by a separate init kernel; the cudaMemsetAsync that - // zeroes state+histograms together is enough. Pass-1 trailer below - // writes st.candCount = seqLens[rowIdx] for pass 2 to consume. - int const rowEnd = seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1; - InputT const* in = logits + static_cast(rowIdx) * stride0; - size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; - size_t numThreads = static_cast(blocksPerRow) * kThreads; - auto f = [&](InputT vIn, size_t /*idx*/) - { - float const v = static_cast(vIn); - uint32_t bin = extractBinIdx(v); - atomicAdd(&sHist[bin], 1); - }; - vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); - } - else if constexpr (step == 2) - { - int const rowEnd = st.candCount; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int* candArr = candBufOut + static_cast(rowIdx) * stride0; - int const prevThresh = st.thresholdBin; - size_t threadRank = static_cast(blockInRow) * kThreads + threadIdx.x; - size_t numThreads = static_cast(blocksPerRow) * kThreads; - auto f = [&](InputT vIn, size_t i) - { - float const v = static_cast(vIn); - int bin1 = static_cast(extractBinIdx<1>(v)); - if (bin1 < prevThresh) - { - int pos = atomicAdd(&st.outIdx, 1); - if (pos < topK) - outIdxArr[pos] = static_cast(i); - } - else if (bin1 == prevThresh) - { - int pos = atomicAdd(&st.filterCnt, 1); - candArr[pos] = static_cast(i); - uint32_t bin2 = extractBinIdx(v); - atomicAdd(&sHist[bin2], 1); - } - }; - vectorized_process(threadRank, numThreads, in, static_cast(rowEnd), f); - } - else // step == 3 - { - int const candCnt = st.candCount; - int const* candArrIn = candBufIn + static_cast(rowIdx) * stride0; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int* candArrOut = candBufOut + static_cast(rowIdx) * stride0; - int const prevThresh = st.thresholdBin; - for (int i = blockInRow * kThreads + threadIdx.x; i < candCnt; i += blocksPerRow * kThreads) - { - int srcIdx = candArrIn[i]; - float v = static_cast(in[srcIdx]); - int bin2 = static_cast(extractBinIdx<2>(v)); - if (bin2 < prevThresh) - { - int pos = atomicAdd(&st.outIdx, 1); - if (pos < topK) - outIdxArr[pos] = srcIdx; - } - else if (bin2 == prevThresh) - { - int pos = atomicAdd(&st.filterCnt, 1); - candArrOut[pos] = srcIdx; - uint32_t bin3 = extractBinIdx(v); - atomicAdd(&sHist[bin3], 1); - } - } - } - __syncthreads(); - - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - { - int c = sHist[i]; - if (c) - atomicAdd(&gHist[i], c); - } - - __threadfence(); - __shared__ int isLast; - if (threadIdx.x == 0) - { - int prev = atomicAdd(&st.finishedBlocks, 1); - isLast = (prev == blocksPerRow - 1) ? 1 : 0; - } - __syncthreads(); - if (!isLast) - return; - - int const rowFullLen = (step == 1) ? (seqLens[rowIdx / next_n] - next_n + (rowIdx % next_n) + 1) : 0; - radixLastBlockTrailer(gHist, st, topK, rowFullLen); - - if constexpr (step == 3) - { - // Final emit, folded into the last block of pass 3. Scan candBufOut - // (top 22 bits == thresholdBin2) and route items into outIndices by - // bin3 vs thresholdBin3. Two-counter scheme so threshold-bin ties - // don't race definite top-k items on the same atomic: - // bin3 < thresh3 → slots [ltBase, ltBase + prefix3) - // bin3 == thresh3 → slots [ltBase + prefix3, topK) - __syncthreads(); - int const filterCnt = st.candCount; - int const thresh3 = st.thresholdBin; - int const prefix3 = st.thresholdLess; - int const* candArr = candBufOut + static_cast(rowIdx) * stride0; - InputT const* in = logits + static_cast(rowIdx) * stride0; - int* outIdxArr = outIndices + static_cast(rowIdx) * topK; - int const ltBase = st.outIdx; // already at outBase here - int const eqBase = ltBase + prefix3; - int const eqCap = topK - eqBase; // ≥ 0 by trailer invariant - __shared__ int sEqEmitted; - if (threadIdx.x == 0) - sEqEmitted = 0; - __syncthreads(); - for (int i = threadIdx.x; i < filterCnt; i += kThreads) - { - int srcIdx = candArr[i]; - float v = static_cast(in[srcIdx]); - int bin3 = static_cast(extractBinIdx<3>(v)); - if (bin3 < thresh3) - { - // atomicAdd on st.outIdx is safe: by construction exactly - // prefix3 items fall in this branch, so pos stays in - // [ltBase, eqBase) which is strictly inside [0, topK). - int pos = atomicAdd(&st.outIdx, 1); - outIdxArr[pos] = srcIdx; - } - else if (bin3 == thresh3) - { - int pos = atomicAdd(&sEqEmitted, 1); - if (pos < eqCap) - outIdxArr[eqBase + pos] = srcIdx; - } - } - __syncthreads(); - if (threadIdx.x == 0) - { - int eq = sEqEmitted < eqCap ? sEqEmitted : eqCap; - int filled = eqBase + eq; - if (filled > topK) - filled = topK; - for (int i = filled; i < topK; ++i) - outIdxArr[i] = -1; - } - // Reset the per-row global histogram and st.outIdx so the next call - // sees a clean state without a per-call cudaMemsetAsync. Caller must - // zero-initialize the scratch buffer before the first call. - __syncthreads(); - for (int i = threadIdx.x; i < kRadixBins; i += kThreads) - gHist[i] = 0; - if (threadIdx.x == 0) - st.outIdx = 0; - } -} - -// Scratch layout (uint8 buffer, 64-byte aligned regions): -// RadixState[numRows] -// int histograms[numRows * kRadixBins] (zeroed on first call by -// torch::zeros allocator; pass-3 -// trailer zeroes for subsequent -// calls) -// int candBuf1[numRows * stride0] (pass 2 → pass 3 input) -// int candBuf2[numRows * stride0] (pass 3 → fused final filter) -static size_t radixScratchBytes(int numRows, int numColumns) -{ - auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; - size_t s = 0; - s += roundUp(sizeof(RadixState) * numRows); - s += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); - s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - s += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - return s; -} - -template -static void launchMultiPassRadix(void* scratch, InputT const* logits, int const* seqLens, int* outIndices, int numRows, - int numColumns, int topK, int stride0, int next_n, cudaLaunchAttribute const* attrs, cudaStream_t stream) -{ - auto roundUp = [](size_t x) { return (x + 63) & ~size_t(63); }; - char* base = static_cast(scratch); - RadixState* state = reinterpret_cast(base); - base += roundUp(sizeof(RadixState) * numRows); - int* histograms = reinterpret_cast(base); - base += roundUp(sizeof(int) * static_cast(numRows) * kRadixBins); - int* candBuf1 = reinterpret_cast(base); - base += roundUp(sizeof(int) * static_cast(numRows) * numColumns); - int* candBuf2 = reinterpret_cast(base); - - int sm_cnt = 132; - { - int dev = 0; - cudaGetDevice(&dev); - cudaDeviceGetAttribute(&sm_cnt, cudaDevAttrMultiProcessorCount, dev); - } - // Block fan-out heuristic: target ~4 active blocks/SM (one wave at the - // achievable occupancy of radixPassKernel<512, 1>), with a per-block work - // floor of 2048 items (4 items/thread at 512-wide). - int targetTotalBlocks = sm_cnt * 4; - int numBlocksPerRow = (targetTotalBlocks + numRows - 1) / numRows; - int maxByCols = numColumns / 2048; - if (numBlocksPerRow > maxByCols) - numBlocksPerRow = maxByCols; - if (numBlocksPerRow < 1) - numBlocksPerRow = 1; - - constexpr int kPassThreads = 512; - - auto launchPass = [&](void const* kernel, int const* candIn, int* candOut) - { - cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(numBlocksPerRow, numRows); - cfg.blockDim = kPassThreads; - cfg.dynamicSmemBytes = 0; - cfg.stream = stream; - cfg.numAttrs = 1; - cfg.attrs = const_cast(attrs); - void* args[] = {(void*) &logits, (void*) &seqLens, (void*) &outIndices, (void*) &candIn, (void*) &candOut, - (void*) &histograms, (void*) &state, (void*) &stride0, (void*) &next_n, (void*) &topK}; - cudaLaunchKernelExC(&cfg, kernel, args); - }; - - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) nullptr, (int*) nullptr); - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) nullptr, candBuf1); - // Pass 3 emits the final top-K inline in its last-block trailer (see - // radixPassKernel) instead of requiring a separate filter launch. - launchPass( - reinterpret_cast(&radixPassKernel), (int const*) candBuf1, candBuf2); -} - -// Architecture-derived GVR eligibility bounds (cached per-process). +// Scheme X bound calculator — shared between fp32 and bf16/fp16 dispatchers. +// Caches hardware attrs (SM count, L2 capacity) and the small-N threshold +// once per process via std::call_once. Per-call cost is just two reads +// from cached static variables plus a small arithmetic block, no syscalls. struct SchemeXBounds { int smCount; @@ -1151,33 +731,89 @@ inline SchemeXBounds getSchemeXBounds(int numColumns, int bytesPerElem) return b; } -// Unified dispatcher (fp32 / bf16 / fp16). Each tier's kernel is templated on -// InputT and casts to float at HBM-read sites; the scratch buffer is uint8. -template -void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +} // anonymous namespace + +void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, + int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, + int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, + int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) { - // Split-work cutoff: matches main's 200k default. is_prefill forces - // single-block via a 1<<30 threshold no shape can reach: prefill chunks are - // bounded by max_num_tokens, well below the multi-pass radix crossover at - // any practical setting. - int const adaptiveSplitWorkThreshold = is_prefill ? (1 << 30) : 200 * 1000; - int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : adaptiveSplitWorkThreshold; + + // INVARIANT: kSortingAlgorithmThreshold is the ORIGINAL TRT-LLM Radix-path + // internal boundary (Insertion vs Radix-radix). v1.2.X dispatcher leaves it + // at 12288 — the GVR Heuristic axis (kSeqSmall, see below) is INDEPENDENT, + // so when canUseHeuristic is false (e.g. preIdx missing, BS too large, or + // numColumns < kSeqSmall), this function falls back to BYTE-IDENTICAL + // original radix dispatcher behavior. Do not touch this constant. + constexpr int kSortingAlgorithmThreshold = 12288; + constexpr int kDefaultSplitWorkThreshold = 200 * 1000; constexpr int kNumThreadsPerBlock = 512; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // ======================================================================== + // Small-N dispatch axis. + // + // GVR Heuristic Top-K has a *fixed* per-launch overhead from Phase-1 + // (preIdx stats reduction over M=2048) and Phase-4 (2048-bin histogram + // snap), totaling ~11 µs regardless of N. For small N (≤16K), this + // fixed cost dominates and the kernel loses to the existing + // insertion-sort/radix path. Empirically (random data, B200 BS=1): + // N=8192 : Heuristic 16.5 µs vs Radix 11.2 µs (radix 1.47× faster) + // N=16384 : Heuristic 21.9 µs vs Radix 22.0 µs (parity) + // N=32768 : Heuristic 26.1 µs vs Radix 32.9 µs (heuristic 1.26× faster) + // N=131072 : Heuristic 43.4 µs vs Radix 76.1 µs (heuristic 1.75× faster) + // + // Route N < kSeqSmall to the existing Radix/Insertion path (which itself + // splits at kSortingAlgorithmThreshold=12288). kSeqSmall is set at the + // empirical crossover point. + // + // ======================================================================== + // Architecture-derived BS-threshold dispatch — jointly bounded by + // occupancy AND L2 cache capacity. + // + // Two physical constraints bound when the per-row heuristic kernel + // remains faster than a radix streaming kernel: + // + // (A) Occupancy bound — 3·SM − SM/8 (wave geometry + setup margin) + // Each CTA uses ~58 KB SMEM (fixed, independent of N), so B200's + // 228 KB dynamic SMEM allows max 3 CTA/SM. Above 3·SM rows per + // launch, tail-wave imbalance causes stragglers. The -SM/8 margin + // (~1/8 wave) covers CTA setup + L2 ingestion overhead. + // On B200(148 SM): 3×148 − 18 = 426. + // + // (B) L2 cache bound — 0.9·L2 / (4·N) per-CTA logits fit + // Each CTA streams its row (N×4B) through L2 per Phase-2 iter. + // With num_concurrent_CTAs × N × 4B > L2, eviction dominates. + // On B200(126 MB L2) with N=70K: 0.9·126MB/(4·70690) ≈ 440, + // which is ~ equal to (A)=426 — the two constraints cross over + // near the SWE-Bench data point. + // For N > 73K the L2 bound tightens below (A) and must take + // over; e.g. N=128K → kBsL2=238, N=196K → kBsL2=155. + // + // Dispatch threshold = min(kBsWave, kBsL2), still data-agnostic (only + // queries hardware attrs). At N≈70K both bounds produce ~426, so the + // L2 axis is a no-op there; for larger N it auto-tightens the threshold. + // + // Small-N lower bound `kSeqSmall` (default 12288) lets the Heuristic + // axis take over wherever the original Radix-radix branch would have + // triggered. Random-data benchmarks suggest the crossover is 16384, + // but workloads with strongly preIdx-correlated logits make P1 stats + // accurate and P2 converge in 1-2 iterations, shifting the real + // crossover into the [12288, 16384] band. Configurable via + // TRTLLM_HEURISTIC_NMIN env (>=1024). + // ======================================================================== + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/4); + int const kBsWave = bounds.kBsWave; + int const kBsL2 = bounds.kBsL2; + int const kBsLarge = bounds.kBsLarge; + int const kSeqSmall = bounds.kSeqSmall; - // GVR eligibility (matches main's rule): supported K, stride1 contiguous, - // preIdx + scratch provided, numColumns in [kSeqSmall, splitWorkThreshold), - // and numRows below the architecture-derived wave/L2 bound. is_prefill - // suppresses GVR through effectiveSplitWorkThreshold being huge. - auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); - bool const canUseHeuristic = preIdx != nullptr && stride1 == 1 && isSupportedTopK && preIdxCount == topK - && preIdxStride >= preIdxCount && heuristicScratch != nullptr && numColumns >= bounds.kSeqSmall - && numColumns < effectiveSplitWorkThreshold && numRows < bounds.kBsLarge; + bool const canUseHeuristic = compressRatio == 1 && preIdx != nullptr && stride1 == 1 && isSupportedTopK + && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold + && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; - // Env-gated dispatch trace (TRTLLM_SCHEMEX_DEBUG=1). + // Optional env-gated dispatch trace (set TRTLLM_SCHEMEX_DEBUG=1 to enable) { static std::once_flag sDebugOnceFlag; static bool sDebug = false; @@ -1189,8 +825,12 @@ void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* }); if (sDebug) { - fprintf(stderr, "[Scheme X] numRows=%d numColumns=%d kBsLarge=%d kSeqSmall=%d -> %s path\n", numRows, - numColumns, bounds.kBsLarge, bounds.kSeqSmall, canUseHeuristic ? "Heuristic" : "Radix"); + fprintf(stderr, + "[Scheme X] numRows=%d numColumns=%d kBsWave=%d kBsL2=%d kBsLarge=%d kSeqSmall=%d smCount=%d " + "L2=%dMB -> %s path%s\n", + numRows, numColumns, kBsWave, kBsL2, kBsLarge, kSeqSmall, bounds.smCount, + bounds.l2Bytes / (1024 * 1024), canUseHeuristic ? "Heuristic" : "Radix", + (numColumns < kSeqSmall) ? " (small-N route)" : ""); } } @@ -1199,11 +839,30 @@ void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, preIdxStride, preIdxCount, numRows, stream); } + else if (numColumns < kSortingAlgorithmThreshold) + { + // Use insertion sort + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } else if (numColumns < effectiveSplitWorkThreshold) { - // Single-block tier: one CTA per row. - auto* kernel_instance = &topKPerRowDecode; + // From this threshold, use radix sort instead + auto* kernel_instance = &topKPerRowDecode; + cudaLaunchConfig_t config; config.gridDim = numRows; config.blockDim = kNumThreadsPerBlock; @@ -1214,74 +873,192 @@ void invokeIndexerTopKDecodeImpl(InputT const* logits, int const* seqLens, int* attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - /*outLogits=*/nullptr, /*numBlocksToMerge=*/0, /*indices=*/nullptr); + compressRatio, nullptr, 0, nullptr); } else { - // Multi-pass radix. radixPassKernel reads logits contiguously, so - // strided inputs would rank the wrong values — gate on stride1 == 1. - // (The single-block tier handles stride1 != 1 via topKPerRowJob's - // strided fallback.) - TLLM_CHECK_WITH_INFO(stride1 == 1, "indexer top-k split-work tier (multi-pass radix) requires stride1 == 1."); - TLLM_CHECK_WITH_INFO(scratch != nullptr && scratchBytes >= radixScratchBytes(numRows, numColumns), - "indexer top-k split-work tier: scratch buffer missing or too small."); - cudaLaunchAttribute radixAttrs[1]; - radixAttrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - radixAttrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); - launchMultiPassRadix( - scratch, logits, seqLens, indices, numRows, numColumns, topK, stride0, next_n, radixAttrs, stream); + // Long sequences are run in two steps + constexpr auto multipleBlocksPerRowConfig = 10; + auto* kernel_instance_part1 = &topKPerRowDecode; + cudaLaunchConfig_t config_part1; + config_part1.gridDim = dim3(numRows, multipleBlocksPerRowConfig); + config_part1.blockDim = kNumThreadsPerBlock; + config_part1.dynamicSmemBytes = 2 * topK * sizeof(int32_t); + config_part1.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config_part1.numAttrs = 1; + config_part1.attrs = attrs; + + cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, + next_n, compressRatio, outLogitsAux, 0, nullptr); + + constexpr int kNumThreadsPerBlockMerge = 1024; + auto* kernel_instance_part2 = &topKPerRowDecode; + cudaLaunchConfig_t config_part2; + config_part2.gridDim = numRows; + config_part2.blockDim = kNumThreadsPerBlockMerge; + config_part2.dynamicSmemBytes = topK * sizeof(int32_t); + config_part2.stream = stream; + // Reuse attrs array since part1 kernel has already been launched + config_part2.numAttrs = 1; + config_part2.attrs = attrs; + + cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, + multipleBlocksPerRowConfig * topK, 1, topK, next_n, 1, nullptr, multipleBlocksPerRowConfig, outIndicesAux); } sync_check_cuda_error(stream); } -} // anonymous namespace +// ============================================================================ +// bf16 / fp16 dispatcher overloads +// ============================================================================ +// Reuses the BS-threshold + small-N dispatch axes (kBsLarge, kSeqSmall) from +// the fp32 dispatcher, except kBsL2 uses sizeof(InputT) bytes/element instead +// of 4 — L2 footprint is half, so bf16/fp16 path remains valid for larger BS +// than fp32 at the same N. +// +// Fallback chain when GVR-Heuristic preconditions are not met (preIdx +// missing, BS too large, or numColumns < kSeqSmall): +// numColumns < kSortingAlgorithmThreshold (12288) → insertion sort +// kSortingAlgorithmThreshold ≤ numColumns < splitWorkThreshold → radix sort +// numColumns ≥ splitWorkThreshold (200K default) → unsupported +// +// Insertion + radix tiers use the same topKPerRowDecode kernel as fp32 with +// InputT propagated through; the histogram and sort steps operate on float +// keys after a static_cast(InputT) at HBM-read sites, so accuracy is +// identical to casting input to fp32 before the kernel. +// +// The split-work tier requires float aux buffers (outLogitsAux / +// outIndicesAux) that the bf16/fp16 entry does not expose; callers in that +// regime must use the fp32 entry. -void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, - int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, - int const* preIdx, int const preIdxStride, int const preIdxCount, float* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) +namespace { - invokeIndexerTopKDecodeImpl(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, - is_prefill); -} -size_t indexerTopKDecodeScratchBytes(int numRows, int numColumns, int /*topK*/) +template +void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, + int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, + int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, + int const compressRatio, cudaStream_t const stream) { - return radixScratchBytes(numRows, numColumns); + static_assert(std::is_same_v || std::is_same_v, + "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); + + constexpr int kSortingAlgorithmThreshold = 12288; + constexpr int kDefaultSplitWorkThreshold = 200 * 1000; + constexpr int kNumThreadsPerBlock = 512; + int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; + + // bf16/fp16: bytes_per_element = sizeof(InputT) = 2 → kBsL2 doubles vs fp32. + auto const bounds = getSchemeXBounds(numColumns, /*bytesPerElem=*/static_cast(sizeof(InputT))); + int const kBsLarge = bounds.kBsLarge; + int const kSeqSmall = bounds.kSeqSmall; + + bool const isSupportedTopK = (topK == 512 || topK == 1024 || topK == 2048); + bool const canUseHeuristic = compressRatio == 1 && preIdx != nullptr && stride1 == 1 && isSupportedTopK + && preIdxCount == topK && preIdxStride >= preIdxCount && numColumns < effectiveSplitWorkThreshold + && numColumns >= kSeqSmall && heuristicScratch != nullptr && numRows < kBsLarge; + + if (canUseHeuristic) + { + launchHeuristicTopKDecode(logits, seqLens, preIdx, indices, heuristicScratch, stride0, next_n, topK, + preIdxStride, preIdxCount, numRows, stream); + } + else if (numColumns < kSortingAlgorithmThreshold) + { + // Insertion sort path — InputT propagated; histogram/sort run on float keys. + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } + else if (numColumns < effectiveSplitWorkThreshold) + { + // Radix sort path — InputT propagated; histogram/sort run on float keys. + auto* kernel_instance = &topKPerRowDecode; + + cudaLaunchConfig_t config; + config.gridDim = numRows; + config.blockDim = kNumThreadsPerBlock; + config.dynamicSmemBytes = topK * sizeof(int32_t); + config.stream = stream; + cudaLaunchAttribute attrs[1]; + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); + config.numAttrs = 1; + config.attrs = attrs; + + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, + compressRatio, nullptr, 0, nullptr); + } + else + { + TLLM_CHECK_WITH_INFO(false, + "indexer_topk_decode bf16/fp16 path does not support numColumns >= splitWorkThreshold " + "(split-work path requires float aux buffers not exposed in the bf16/fp16 entry). " + "Got numColumns=%d splitWorkThreshold=%d. Use the fp32 entry for this regime.", + numColumns, effectiveSplitWorkThreshold); + } + + sync_check_cuda_error(stream); } +} // anonymous namespace + void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, - __nv_bfloat16* heuristicScratch, cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) + __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream) { - invokeIndexerTopKDecodeImpl<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, - scratchBytes, is_prefill); + invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, + stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); } void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, - cudaStream_t const stream, void* scratch, size_t scratchBytes, bool is_prefill) + int const compressRatio, cudaStream_t const stream) { - invokeIndexerTopKDecodeImpl<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, stream, scratch, scratchBytes, - is_prefill); + invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK, cudaStream_t const stream) { + constexpr int kSortingAlgorithmThreshold = 12288; constexpr int kNumThreadsPerBlock = 512; - // One launch over all rows; the per-row sort algorithm is picked at - // runtime inside topKPerRowJob. - topKPerRowPrefill<<>>( - logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); + int numInsertionBlocks = std::min(numRows, kSortingAlgorithmThreshold); + topKPerRowPrefill + <<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, 0); + + if (numRows > kSortingAlgorithmThreshold) + { + int numRadixBlocks = numRows - kSortingAlgorithmThreshold; + topKPerRowPrefill + <<>>( + logits, rowStarts, rowEnds, indices, stride0, stride1, topK, kSortingAlgorithmThreshold); + } sync_check_cuda_error(stream); } @@ -1293,7 +1070,6 @@ bool canIndexerTopKDecodeUseGvr(int numRows, int numColumns, int topK, int bytes { return false; } - // Mirrors the dispatcher's effectiveSplitWorkThreshold default. constexpr int kDefaultSplitWorkThreshold = 200 * 1000; auto const bounds = getSchemeXBounds(numColumns, bytesPerElem); return numColumns >= bounds.kSeqSmall && numColumns < kDefaultSplitWorkThreshold && numRows < bounds.kBsLarge; diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt new file mode 100644 index 000000000000..ae7ed2d085c8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt @@ -0,0 +1,37 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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. +# + +# Library sources are listed explicitly. Profiling/benchmark/test harnesses +# in this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are +# standalone binaries with their own main() and must NOT be compiled into the +# trtllm shared library. +set(SRC_CU + mhcKernels.cu + mhcFusedHcKernel.cu +) +add_library(mhcKernels_src OBJECT ${SRC_CU}) + +set_property(TARGET mhcKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET mhcKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) +target_compile_options(mhcKernels_src PRIVATE $<$:--use_fast_math>) + +# Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh, reduction.cuh) +# Used by the tcgen05.mma-based fused post-mapping + GEMM kernel on SM100. +target_include_directories(mhcKernels_src + PRIVATE + ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include +) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh new file mode 100644 index 000000000000..24e9d3d656a5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -0,0 +1,1331 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Fused post-mapping + TF32 HC prenorm GEMM on B200 (SM100). +// +// Mathematical formula: +// new_r[i, hc, h] = post_mix[i, hc] * x[i, h] +// + sum_j comb_mix[i, j, hc] * residual[i, j, h] +// D[i, n] = sum_{hc, h} new_r[i, hc, h] * W[n, hc, h] +// sqr[i] = sum_{hc, h} new_r[i, hc, h]^2 (bf16-rounded) +// +// Shape: +// residual [M, HC_MULT, hidden] bf16 +// x [M, hidden] bf16 +// post_mix [M, HC_MULT] fp32 +// comb_mix [M, HC_MULT, HC_MULT] fp32 +// W [N, HC_MULT*hidden] fp32 (TF32) +// D [M, N] fp32 +// sqr [M] fp32 +// +// Kernel architecture (derived from DeepGEMM sm100_tf32_hc_prenorm_gemm): +// - BLOCK_M x BLOCK_N x BLOCK_K = 64 x 32 x 64, TF32 MMA on tcgen05 +// - 256 threads per CTA (warps 0..3 = MMA group, warps 4..7 = pmap group) +// - new_r is NEVER materialized in GMEM - it is computed on-chip directly +// into TMEM (fp32) where UMMA consumes it. +// - Iteration order: outer h_tile (slow), inner hc_idx (fast). residual+x +// SMEM is reused for HC_MULT=4 consecutive MMA stages. +// +// Barriers: +// full_B[N_B_STAGES] TMA -> MMA (B arrived in SMEM) +// empty_B[N_B_STAGES] MMA -> TMA (B slot empty) +// full_input[N_INPUT] TMA -> pmap (residual+x arrived) +// empty_input[N_INPUT] pmap -> TMA (input slot empty) +// full_cast[2] pmap -> MMA (A ready in TMEM) +// empty_cast[2] MMA -> pmap (TMEM slot empty) +// tmem_full[1] MMA -> epilogue + +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include + +#include + +namespace deep_gemm +{ +using math::swap; +} + +#include +#include +#include +#include +#include +#include +#include + +namespace fused_mhc +{ + +// Reuse DeepGEMM's swizzle helper +using deep_gemm::sm100::make_umma_desc; +using deep_gemm::sm100::get_num_aligned_tmem_cols; +using deep_gemm::sm100::tcgen05_before_thread_sync; +using deep_gemm::sm100::tcgen05_after_thread_sync; +using deep_gemm::sm100::advance_umma_desc_lo; +using deep_gemm::tma_copy; +using deep_gemm::math::constexpr_ceil_div; +using deep_gemm::ptx::get_lane_idx; +using deep_gemm::utils::PatternVisitor; + +template +__device__ __forceinline__ uint32_t get_swizzled_smem_offset(uint32_t const& offset, uint32_t const& lane_idx) +{ + auto const& bank_group_idx = offset + lane_idx * (kSwizzleMode / kSwizzleBase); + constexpr uint32_t kNumBankGroups = 128 / kSwizzleBase; + constexpr bool kHasShortcut = (kSwizzleMode / kSwizzleBase) == kNumBankGroups; + auto row = kHasShortcut ? (offset / kNumBankGroups + lane_idx) : (bank_group_idx / kNumBankGroups); + auto col = kHasShortcut ? (offset) : (bank_group_idx % kNumBankGroups); + col ^= row % (kSwizzleMode / kSwizzleBase); + return row * 128 + col * kSwizzleBase; +} + +__device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uint32_t b, uint32_t c, uint32_t d) +{ + asm volatile( + "stmatrix.sync.aligned.x4.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n" ::"l"(__cvta_generic_to_shared(smem_dst)), + "r"(a), "r"(b), "r"(c), "r"(d)); +} + +template +__global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf32_pmap_gemm_rout_atomic_impl( + const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual, + const __grid_constant__ cute::TmaDescriptor tensor_map_x, const __grid_constant__ cute::TmaDescriptor tensor_map_b, + const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, + float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0) + float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum) +{ // [M] (caller memsets to 0) +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; + constexpr uint32_t H_TILES_PER_HC = HIDDEN / BLOCK_K; + static_assert(H_TILES_PER_HC % kNumSplits == 0, "H_TILES_PER_HC must be divisible by kNumSplits"); + constexpr uint32_t H_TILES_PER_SPLIT = H_TILES_PER_HC / kNumSplits; + constexpr uint32_t kNumCastStages = 4; + constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); + constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); + constexpr uint32_t kSwizzleXMode = kSwizzleAMode; + constexpr uint32_t kSwizzleResMode = kSwizzleAMode; + constexpr uint32_t kSwizzleRoutMode = kSwizzleAMode; + constexpr auto kMajorA = cute::UMMA::Major::K; + constexpr auto kMajorB = cute::UMMA::Major::K; + static_assert(HIDDEN % BLOCK_K == 0, "HIDDEN must be multiple of BLOCK_K"); + static_assert(N_B_STAGES >= HC_MULT, "N_B_STAGES must be >= HC_MULT"); + static_assert(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); + static_assert(kNumMMAThreads == 128, "Invalid MMA threads"); + static_assert(kNumPmapThreads == 128, "Invalid pmap threads"); + static_assert(BLOCK_M == 64, "Invalid block M"); + static_assert(HC_MULT == 4, "Only HC_MULT=4 supported"); + static_assert(kSwizzleCDMode == 128, "Atomic variant expects kSwizzleCDMode=128"); + static_assert(SHAPE_N <= BLOCK_N, "SHAPE_N must fit within BLOCK_N"); + + auto const warp_idx = cutlass::canonical_warp_idx_sync(); + auto const lane_idx = get_lane_idx(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_B_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); + constexpr uint32_t SMEM_RES_PER_ISTG = BLOCK_M * HC_MULT * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_X_PER_ISTG = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_RC_PER_HC = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); // 8 KB + constexpr uint32_t SMEM_RC_SIZE = HC_MULT * SMEM_RC_PER_HC; // 32 KB + + constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); + + // Prefetch TMA descriptors + if (warp_idx == 0 and cute::elect_one_sync()) + { + cute::prefetch_tma_descriptor(&tensor_map_residual); + cute::prefetch_tma_descriptor(&tensor_map_x); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_residual_out); + } + + // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] + auto smem_cd = reinterpret_cast(smem_buffer); + uint8_t* cursor = smem_buffer + SMEM_CD_SIZE; + auto smem_b = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_B_PER_STAGE); }); + cursor += N_B_STAGES * SMEM_B_PER_STAGE; + auto smem_res = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; + auto smem_x_stg = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; + auto smem_post = reinterpret_cast(cursor); + cursor += SMEM_POST_SIZE; + auto smem_comb = reinterpret_cast(cursor); + cursor += SMEM_COMB_SIZE; + auto smem_rc = reinterpret_cast(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16, single-buffered + cursor += SMEM_RC_SIZE; + + cursor = reinterpret_cast((reinterpret_cast(cursor) + 7) & ~uintptr_t(7)); + auto barrier_start_ptr = reinterpret_cast(cursor); + auto full_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + i; }); + auto empty_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + N_B_STAGES + i; }); + auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; }); + auto empty_input + = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; }); + auto full_cast = PatternVisitor( + [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); + auto empty_cast = PatternVisitor([=](uint32_t const& i) + { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); + auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + auto tmem_ptr_in_smem = reinterpret_cast(cursor); + + if (warp_idx == 1 and cute::elect_one_sync()) + { +#pragma unroll + for (uint32_t i = 0; i < N_B_STAGES; ++i) + { + full_B[i]->init(1); + empty_B[i]->init(1); + } +#pragma unroll + for (uint32_t i = 0; i < N_INPUT_STAGES; ++i) + { + full_input[i]->init(1); + empty_input[i]->init(kNumPmapThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumCastStages; ++i) + { + full_cast[i]->init(kNumPmapThreads); + empty_cast[i]->init(1); + } + tmem_full_barrier->init(1); + cutlass::arch::fence_barrier_init(); + } + else if (warp_idx == 2) + { + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + __syncthreads(); + + const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); + const uint32_t m_block_idx = block_idx / kNumSplits; + const uint32_t k_split_idx = block_idx % kNumSplits; + const uint32_t m_offset = m_block_idx * BLOCK_M; + const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT; + constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT; + + // Prologue: pmap warp group loads post_mix, comb_mix into SMEM + if (warp_idx >= kNumMMAThreads / 32) + { + const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; +#pragma unroll + for (uint32_t t = 0; t < 2; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT) + { + uint32_t m = idx / HC_MULT; + uint32_t hc = idx % HC_MULT; + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? post_mix[gmem_m * HC_MULT + hc] : 0.f; + smem_post[idx] = v; + } + } +#pragma unroll + for (uint32_t t = 0; t < 8; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT * HC_MULT) + { + uint32_t m = idx / (HC_MULT * HC_MULT); + uint32_t jk = idx % (HC_MULT * HC_MULT); + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? comb_mix[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; + smem_comb[idx] = v; + } + } + } + __syncthreads(); + + if (warp_idx < kNumMMAThreads / 32) + { + // ----- TMA warp (warp 0) ----- + if (warp_idx == 0 and cute::elect_one_sync()) + { + uint32_t b_stage = 0; + uint32_t i_stage = 0; + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t h_tile = h_tile_start + ht; + empty_input[i_stage]->wait(((ht / N_INPUT_STAGES) & 1) ^ 1); + uint32_t m_idx = m_block_idx * BLOCK_M; + uint32_t h_idx = h_tile * BLOCK_K; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + tma_copy(&tensor_map_residual, full_input[i_stage], + smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); + } + tma_copy( + &tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx); + constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG; + full_input[i_stage]->arrive_and_expect_tx(kInputBytes); + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + empty_B[b_stage]->wait(((s / N_B_STAGES) & 1) ^ 1); + uint32_t k_idx = hc * HIDDEN + h_idx; + tma_copy( + &tensor_map_b, full_B[b_stage], smem_b[b_stage], k_idx, 0); + full_B[b_stage]->arrive_and_expect_tx(SMEM_B_PER_STAGE); + b_stage = (b_stage + 1) % N_B_STAGES; + ++s; + } + i_stage = (i_stage + 1) % N_INPUT_STAGES; + } + } + + // ----- MMA issue warp (warp 1) ----- + if (warp_idx == 1) + { + constexpr uint32_t UMMA_M = BLOCK_M; + constexpr uint32_t UMMA_N = BLOCK_N; + constexpr uint32_t UMMA_K = 32 / sizeof(float); + constexpr uint32_t BLOCK_SWIZZLED_BK = kSwizzleBMode / sizeof(float); + using umma_t = cute::SM100_MMA_TF32_TS; + auto instr_desc = cute::UMMA::make_instr_desc(); + auto const& runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + static_assert(N_B_STAGES <= 32, "Too many B stages"); + auto b_desc = make_umma_desc(smem_b[0], 0, 0); + uint32_t const& b_desc_lo = lane_idx < N_B_STAGES ? b_desc.lo + lane_idx * SMEM_B_PER_STAGE / 16 : 0u; + + for (uint32_t s = 0; s < num_total_stages; ++s) + { + const uint32_t b_stage = s % N_B_STAGES; + const uint32_t cast_stage_idx = s % kNumCastStages; + full_cast[cast_stage_idx]->wait((s / kNumCastStages) & 1); + full_B[b_stage]->wait((s / N_B_STAGES) & 1); + tcgen05_after_thread_sync(); + auto const& b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(b_stage)); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) + { + uint32_t const& atom_idx = (k * UMMA_K) / BLOCK_SWIZZLED_BK; + uint32_t const& in_atom_idx = (k * UMMA_K) % BLOCK_SWIZZLED_BK; + uint32_t const& offset = atom_idx * BLOCK_N * BLOCK_SWIZZLED_BK; + b_desc.lo = advance_umma_desc_lo( + b_desc_base_lo, offset, in_atom_idx); + umma_t::fma(BLOCK_K * cast_stage_idx + k * UMMA_K, b_desc, BLOCK_K * kNumCastStages, s > 0 or k > 0, + runtime_instr_desc); + } + cutlass::arch::umma_arrive(reinterpret_cast(empty_cast[cast_stage_idx])); + cutlass::arch::umma_arrive(reinterpret_cast(empty_B[b_stage])); + } + cutlass::arch::umma_arrive(reinterpret_cast(tmem_full_barrier)); + } + + // ----- Epilogue (warps 0..3, 128 threads) ----- + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(float); + static_assert(BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); + + tmem_full_barrier->wait(0); + tcgen05_after_thread_sync(); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_N / kNumElemsPerBankGroup; ++i) + { + uint32_t tmem_addr = BLOCK_K * kNumCastStages + i * kNumElemsPerBankGroup; + auto smem_ptr = reinterpret_cast(smem_cd) + warp_idx * BLOCK_M / 4 * kSwizzleCDMode + + get_swizzled_smem_offset(i, lane_idx); + uint32_t values[kNumElemsPerBankGroup]; + static_assert(kNumElemsPerBankGroup == 4, "Invalid type"); + cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); + cutlass::arch::fence_view_async_tmem_load(); + if (BLOCK_M == 128 or (BLOCK_M == 64 and lane_idx < 16)) + deep_gemm::ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); + if constexpr (BLOCK_M == 64) + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kNumMMAThreads, 0); + + constexpr uint32_t kTotalOut = BLOCK_M * SHAPE_N; + const uint32_t tid = threadIdx.x; +#pragma unroll + for (uint32_t k = tid; k < kTotalOut; k += kNumMMAThreads) + { + uint32_t m = k / SHAPE_N; + uint32_t n = k - m * SHAPE_N; + uint32_t gm = m_block_idx * BLOCK_M + m; + if (gm < shape_m) + { + uint32_t col_group = n >> 2; + uint32_t in_group = n & 3; + uint32_t phys_col_grp = col_group ^ (m & 7); + uint32_t byte = m * kSwizzleCDMode + phys_col_grp * kNumBankGroupBytes + in_group * sizeof(float); + float val = *reinterpret_cast(reinterpret_cast(smem_cd) + byte); + atomicAdd(&D[gm * SHAPE_N + n], val); + } + } + + if (warp_idx == 1) + cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } + else + { + // ----- Pmap warp group (warps 4..7, 128 threads) ----- + const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; + const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; + const uint32_t lower_row = upper_row + 8; + const uint32_t col_lane = lane_idx % 4; + + float pm_u[HC_MULT], pm_l[HC_MULT]; + float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT]; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + pm_u[hc] = smem_post[upper_row * HC_MULT + hc]; + pm_l[hc] = smem_post[lower_row * HC_MULT + hc]; + } +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + } + } + + float sqr_u = 0.f, sqr_l = 0.f; + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); + constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; + constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; + static_assert(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "BLOCK_K must match swizzle A mode"); + static_assert(kNumLoads % 2 == 0, "kNumLoads must be even for LDSM.x4"); + + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t i_stage = ht % N_INPUT_STAGES; + full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1); + + uint32_t x_vals[2][kNumLoads]; + { + uint8_t const* x_base + = reinterpret_cast(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr = x_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1], + x_vals[1][i + 1], const_cast(smem_ptr)); + } + } + + uint32_t r_vals[HC_MULT][2][kNumLoads]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + uint8_t const* r_base = reinterpret_cast(smem_res[i_stage]) + + j * BLOCK_M * BLOCK_K * sizeof(nv_bfloat16) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleResMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = r_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(r_vals[j][0][i + 0], r_vals[j][1][i + 0], + r_vals[j][0][i + 1], r_vals[j][1][i + 1], const_cast(smem_ptr)); + } + } + + float2 xf[2][kNumLoads]; +#pragma unroll + for (uint32_t u = 0; u < 2; ++u) + { +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + xf[u][i] = __bfloat1622float2(*reinterpret_cast(&x_vals[u][i])); + } + } + + if constexpr (kEarlyRelease) + { + empty_input[i_stage]->arrive(); + } + + // Wait for previous ht's residual_out TMA_STOREs to drain before we + // overwrite single-buffered smem_rc with new hc values. + if (ht > 0) + { + cute::tma_store_wait<0>(); + } + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + const uint32_t cast_stage_idx = s % kNumCastStages; + empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1); + + uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads]; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y}; + float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + float2 ruj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][0][i])); + float2 rlj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][1][i])); + nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x); + nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y); + nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x); + nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y); + } + nv_bfloat162 b_up = __float22bfloat162_rn(nu); + nv_bfloat162 b_lo = __float22bfloat162_rn(nl); + uint32_t b_up_bits = *reinterpret_cast(&b_up); + uint32_t b_lo_bits = *reinterpret_cast(&b_lo); + rc_u_buf[i] = b_up_bits; + rc_l_buf[i] = b_lo_bits; + float2 ru = __bfloat1622float2(b_up); + float2 rl = __bfloat1622float2(b_lo); + sqr_u = fmaf(ru.x, ru.x, sqr_u); + sqr_u = fmaf(ru.y, ru.y, sqr_u); + sqr_l = fmaf(rl.x, rl.x, sqr_l); + sqr_l = fmaf(rl.y, rl.y, sqr_l); + cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast(&ru.x), + *reinterpret_cast(&ru.y), *reinterpret_cast(&rl.x), + *reinterpret_cast(&rl.y), cast_stage_idx * BLOCK_K + i * 8); + } + cutlass::arch::fence_view_async_tmem_store(); + tcgen05_before_thread_sync(); + full_cast[cast_stage_idx]->arrive(); + ++s; + + // STSM bf16 new_r values into smem_rc[hc] sub-region for this warp. + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = rc_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]); + } + } + if constexpr (!kEarlyRelease) + { + empty_input[i_stage]->arrive(); + } + + // Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows. + cute::tma_store_fence(); + if (cute::elect_one_sync()) + { + const uint32_t h_idx = (h_tile_start + ht) * BLOCK_K; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; + cute::SM90_TMA_STORE_2D::copy(&tensor_map_residual_out, rc_base, hc * HIDDEN + h_idx, + m_offset + sub_warp_idx * BLOCK_M_PER_WARP); + cute::tma_store_arrive(); + } + } + } + + // Drain any in-flight residual_out TMA stores before exit. + cute::tma_store_wait<0>(); + + // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 2); + if (col_lane == 0) + { + uint32_t gm_u = m_block_idx * BLOCK_M + upper_row; + uint32_t gm_l = m_block_idx * BLOCK_M + lower_row; + if (gm_u < shape_m) + atomicAdd(&sqr_sum[gm_u], sqr_u); + if (gm_l < shape_m) + atomicAdd(&sqr_sum[gm_l], sqr_l); + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_100a"); +#endif +} + +// ============================================================================ +// ALL-IN-ONE variant (Path D, tf32 tcgen05 MMA analogue of Path F). +// +// Single-kernel fusion of: +// 1) post_mapping : residual_cur[j,h] = post_mix_prev[j]*x_prev[h] +// + sum_k comb_mix_prev[k,j]*residual_prev[k,h] +// 2) pre_GEMM : D[i,n] = sum_{hc,h} residual_cur[i,hc,h] * W_T[n, hc*HIDDEN+h] +// sqr[i] = sum_{hc,h} residual_cur[i,hc,h]^2 +// 3) bigFuse : rmsnorm+sigmoid+sinkhorn on D/sqr -> post_mix_out, +// comb_mix_out, pre_mix; layer_input = pre_mix @ residual_cur. +// +// Semantics match Path F (fused_pmap_gemm_fma_allinone) exactly. The only +// algorithmic difference vs Path F is that we use tf32 tcgen05 MMA for the +// residual_cur @ W_T GEMM (instead of CUDA-core FMA). +// +// Pipelining, warp layout, and SMEM layout inherit from Path B +// (fused_tf32_pmap_gemm_rout_atomic_impl): the pmap warp group computes +// residual_cur, TMA-stores it to GMEM, and STSM-casts it into TMEM; the MMA +// warp group consumes TMEM for the GEMM; the epilogue atomicAdd's into +// y_acc / sqr_sum. Phase 3 elects the last-home CTA (per m_block) via +// atomicAdd on done_counter; Phase 4 runs bigfuse inline on that CTA only, +// reloading residual_cur from GMEM and writing layer_input + post_mix_out + +// comb_mix_out. +// +// Caller MUST zero D (y_acc), sqr_sum (r_acc), done_counter before launch. +// ============================================================================ + +template +__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 + const __grid_constant__ cute::TmaDescriptor tensor_map_x, // x_prev, bf16 + const __grid_constant__ cute::TmaDescriptor tensor_map_b, // W_T, tf32 + const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out, // residual_cur, bf16 (TMA store) + __nv_bfloat16 const* __restrict__ residual_cur_ptr, // same buffer as TMA target + __nv_bfloat16* __restrict__ layer_input_out, // [M, HIDDEN] bf16 + float* __restrict__ D, // [M, SHAPE_N] fp32 (y_acc, caller zeros) + float* __restrict__ sqr_sum, // [M] fp32 (r_acc, caller zeros) + int* __restrict__ done_counter, // [ceil(M/BLOCK_M)] int (caller zeros) + float const* __restrict__ post_mix_prev, // [M, HC_MULT] + float const* __restrict__ comb_mix_prev, // [M, HC_MULT, HC_MULT] + float const* __restrict__ hc_scale, // [3] + float const* __restrict__ hc_base, // [HC_MULT*(2+HC_MULT)] + float* __restrict__ post_mix_out, // [M, HC_MULT] + float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] + float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat) +{ +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t HC_MULT2 = HC_MULT * HC_MULT; + constexpr uint32_t HC_MULT3 = HC_MULT * (2 + HC_MULT); + constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; + constexpr uint32_t H_TILES_PER_HC = HIDDEN / BLOCK_K; + static_assert(H_TILES_PER_HC % kNumSplits == 0, "H_TILES_PER_HC must be divisible by kNumSplits"); + constexpr uint32_t H_TILES_PER_SPLIT = H_TILES_PER_HC / kNumSplits; + constexpr uint32_t kNumCastStages = 4; + constexpr uint32_t kSwizzleAMode = cute::min(BLOCK_K * sizeof(nv_bfloat16), 128); + constexpr uint32_t kSwizzleBMode = cute::min(BLOCK_K * sizeof(float), 128); + constexpr uint32_t kSwizzleXMode = kSwizzleAMode; + constexpr uint32_t kSwizzleResMode = kSwizzleAMode; + constexpr uint32_t kSwizzleRoutMode = kSwizzleAMode; + constexpr auto kMajorA = cute::UMMA::Major::K; + constexpr auto kMajorB = cute::UMMA::Major::K; + static_assert(HIDDEN % BLOCK_K == 0, "HIDDEN must be multiple of BLOCK_K"); + static_assert(N_B_STAGES >= HC_MULT, "N_B_STAGES must be >= HC_MULT"); + static_assert(kSwizzleCDMode / sizeof(float) == BLOCK_N, "Invalid block N"); + static_assert(kNumMMAThreads == 128, "Invalid MMA threads"); + static_assert(kNumPmapThreads == 128, "Invalid pmap threads"); + static_assert(BLOCK_M == 64, "Invalid block M"); + static_assert(HC_MULT == 4, "Only HC_MULT=4 supported"); + static_assert(kSwizzleCDMode == 128, "Atomic variant expects kSwizzleCDMode=128"); + static_assert(SHAPE_N <= BLOCK_N, "SHAPE_N must fit within BLOCK_N"); + static_assert(SHAPE_N == HC_MULT3, "Path D expects SHAPE_N == HC_MULT*(2+HC_MULT)=24"); + + auto const warp_idx = cutlass::canonical_warp_idx_sync(); + auto const lane_idx = get_lane_idx(); + + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_CD_SIZE = BLOCK_M * kSwizzleCDMode; + constexpr uint32_t SMEM_B_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(float); + constexpr uint32_t SMEM_RES_PER_ISTG = BLOCK_M * HC_MULT * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_X_PER_ISTG = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_POST_SIZE = BLOCK_M * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_COMB_SIZE = BLOCK_M * HC_MULT * HC_MULT * sizeof(float); + constexpr uint32_t SMEM_RC_PER_HC = BLOCK_M * BLOCK_K * sizeof(nv_bfloat16); // 8 KB + constexpr uint32_t SMEM_RC_SIZE = HC_MULT * SMEM_RC_PER_HC; // 32 KB + + constexpr uint32_t kNumTmemCols = get_num_aligned_tmem_cols(); + + // Prefetch TMA descriptors + if (warp_idx == 0 and cute::elect_one_sync()) + { + cute::prefetch_tma_descriptor(&tensor_map_residual); + cute::prefetch_tma_descriptor(&tensor_map_x); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_residual_out); + } + + // SMEM layout: [cd, B stages, res stages, x stages, post, comb, rc (HC_MULT slices)] + auto smem_cd = reinterpret_cast(smem_buffer); + uint8_t* cursor = smem_buffer + SMEM_CD_SIZE; + auto smem_b = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_B_PER_STAGE); }); + cursor += N_B_STAGES * SMEM_B_PER_STAGE; + auto smem_res = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; + auto smem_x_stg = PatternVisitor( + [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); + cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; + auto smem_post = reinterpret_cast(cursor); + cursor += SMEM_POST_SIZE; + auto smem_comb = reinterpret_cast(cursor); + cursor += SMEM_COMB_SIZE; + auto smem_rc = reinterpret_cast(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16 + cursor += SMEM_RC_SIZE; + + cursor = reinterpret_cast((reinterpret_cast(cursor) + 7) & ~uintptr_t(7)); + auto barrier_start_ptr = reinterpret_cast(cursor); + auto full_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + i; }); + auto empty_B = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + N_B_STAGES + i; }); + auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; }); + auto empty_input + = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; }); + auto full_cast = PatternVisitor( + [=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; }); + auto empty_cast = PatternVisitor([=](uint32_t const& i) + { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; }); + auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages; + + cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier); + auto tmem_ptr_in_smem = reinterpret_cast(cursor); + + if (warp_idx == 1 and cute::elect_one_sync()) + { +#pragma unroll + for (uint32_t i = 0; i < N_B_STAGES; ++i) + { + full_B[i]->init(1); + empty_B[i]->init(1); + } +#pragma unroll + for (uint32_t i = 0; i < N_INPUT_STAGES; ++i) + { + full_input[i]->init(1); + empty_input[i]->init(kNumPmapThreads); + } +#pragma unroll + for (uint32_t i = 0; i < kNumCastStages; ++i) + { + full_cast[i]->init(kNumPmapThreads); + empty_cast[i]->init(1); + } + tmem_full_barrier->init(1); + cutlass::arch::fence_barrier_init(); + } + else if (warp_idx == 2) + { + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + __syncthreads(); + + const uint32_t block_idx = __shfl_sync(0xffffffff, blockIdx.x, 0); + const uint32_t m_block_idx = block_idx / kNumSplits; + const uint32_t k_split_idx = block_idx % kNumSplits; + const uint32_t m_offset = m_block_idx * BLOCK_M; + const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT; + constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT; + + // Prologue: pmap warp group loads post_mix_prev, comb_mix_prev into SMEM + if (warp_idx >= kNumMMAThreads / 32) + { + const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads; +#pragma unroll + for (uint32_t t = 0; t < 2; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT) + { + uint32_t m = idx / HC_MULT; + uint32_t hc = idx % HC_MULT; + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? post_mix_prev[gmem_m * HC_MULT + hc] : 0.f; + smem_post[idx] = v; + } + } +#pragma unroll + for (uint32_t t = 0; t < 8; ++t) + { + uint32_t idx = pmap_tid + t * kNumPmapThreads; + if (idx < BLOCK_M * HC_MULT * HC_MULT) + { + uint32_t m = idx / (HC_MULT * HC_MULT); + uint32_t jk = idx % (HC_MULT * HC_MULT); + uint32_t gmem_m = m_offset + m; + float v = (gmem_m < shape_m) ? comb_mix_prev[gmem_m * HC_MULT * HC_MULT + jk] : 0.f; + smem_comb[idx] = v; + } + } + } + __syncthreads(); + + if (warp_idx < kNumMMAThreads / 32) + { + // ----- TMA warp (warp 0) ----- + if (warp_idx == 0 and cute::elect_one_sync()) + { + uint32_t b_stage = 0; + uint32_t i_stage = 0; + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t h_tile = h_tile_start + ht; + empty_input[i_stage]->wait(((ht / N_INPUT_STAGES) & 1) ^ 1); + uint32_t m_idx = m_block_idx * BLOCK_M; + uint32_t h_idx = h_tile * BLOCK_K; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + tma_copy(&tensor_map_residual, full_input[i_stage], + smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); + } + tma_copy( + &tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx); + constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG; + full_input[i_stage]->arrive_and_expect_tx(kInputBytes); + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + empty_B[b_stage]->wait(((s / N_B_STAGES) & 1) ^ 1); + uint32_t k_idx = hc * HIDDEN + h_idx; + tma_copy( + &tensor_map_b, full_B[b_stage], smem_b[b_stage], k_idx, 0); + full_B[b_stage]->arrive_and_expect_tx(SMEM_B_PER_STAGE); + b_stage = (b_stage + 1) % N_B_STAGES; + ++s; + } + i_stage = (i_stage + 1) % N_INPUT_STAGES; + } + } + + // ----- MMA issue warp (warp 1) ----- + if (warp_idx == 1) + { + constexpr uint32_t UMMA_M = BLOCK_M; + constexpr uint32_t UMMA_N = BLOCK_N; + constexpr uint32_t UMMA_K = 32 / sizeof(float); + constexpr uint32_t BLOCK_SWIZZLED_BK = kSwizzleBMode / sizeof(float); + using umma_t = cute::SM100_MMA_TF32_TS; + auto instr_desc = cute::UMMA::make_instr_desc(); + auto const& runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + static_assert(N_B_STAGES <= 32, "Too many B stages"); + auto b_desc = make_umma_desc(smem_b[0], 0, 0); + uint32_t const& b_desc_lo = lane_idx < N_B_STAGES ? b_desc.lo + lane_idx * SMEM_B_PER_STAGE / 16 : 0u; + + for (uint32_t s = 0; s < num_total_stages; ++s) + { + const uint32_t b_stage = s % N_B_STAGES; + const uint32_t cast_stage_idx = s % kNumCastStages; + full_cast[cast_stage_idx]->wait((s / kNumCastStages) & 1); + full_B[b_stage]->wait((s / N_B_STAGES) & 1); + tcgen05_after_thread_sync(); + auto const& b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(b_stage)); +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) + { + uint32_t const& atom_idx = (k * UMMA_K) / BLOCK_SWIZZLED_BK; + uint32_t const& in_atom_idx = (k * UMMA_K) % BLOCK_SWIZZLED_BK; + uint32_t const& offset = atom_idx * BLOCK_N * BLOCK_SWIZZLED_BK; + b_desc.lo = advance_umma_desc_lo( + b_desc_base_lo, offset, in_atom_idx); + umma_t::fma(BLOCK_K * cast_stage_idx + k * UMMA_K, b_desc, BLOCK_K * kNumCastStages, s > 0 or k > 0, + runtime_instr_desc); + } + cutlass::arch::umma_arrive(reinterpret_cast(empty_cast[cast_stage_idx])); + cutlass::arch::umma_arrive(reinterpret_cast(empty_B[b_stage])); + } + cutlass::arch::umma_arrive(reinterpret_cast(tmem_full_barrier)); + } + + // ----- Epilogue (warps 0..3, 128 threads) ----- + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(float); + static_assert(BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); + + tmem_full_barrier->wait(0); + tcgen05_after_thread_sync(); + +#pragma unroll + for (uint32_t i = 0; i < BLOCK_N / kNumElemsPerBankGroup; ++i) + { + uint32_t tmem_addr = BLOCK_K * kNumCastStages + i * kNumElemsPerBankGroup; + auto smem_ptr = reinterpret_cast(smem_cd) + warp_idx * BLOCK_M / 4 * kSwizzleCDMode + + get_swizzled_smem_offset(i, lane_idx); + uint32_t values[kNumElemsPerBankGroup]; + static_assert(kNumElemsPerBankGroup == 4, "Invalid type"); + cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, values[0], values[1], values[2], values[3]); + cutlass::arch::fence_view_async_tmem_load(); + if (BLOCK_M == 128 or (BLOCK_M == 64 and lane_idx < 16)) + deep_gemm::ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); + if constexpr (BLOCK_M == 64) + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kNumMMAThreads, 0); + + constexpr uint32_t kTotalOut = BLOCK_M * SHAPE_N; + const uint32_t tid = threadIdx.x; +#pragma unroll + for (uint32_t k = tid; k < kTotalOut; k += kNumMMAThreads) + { + uint32_t m = k / SHAPE_N; + uint32_t n = k - m * SHAPE_N; + uint32_t gm = m_block_idx * BLOCK_M + m; + if (gm < shape_m) + { + uint32_t col_group = n >> 2; + uint32_t in_group = n & 3; + uint32_t phys_col_grp = col_group ^ (m & 7); + uint32_t byte = m * kSwizzleCDMode + phys_col_grp * kNumBankGroupBytes + in_group * sizeof(float); + float val = *reinterpret_cast(reinterpret_cast(smem_cd) + byte); + atomicAdd(&D[gm * SHAPE_N + n], val); + } + } + + if (warp_idx == 1) + cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } + else + { + // ----- Pmap warp group (warps 4..7, 128 threads) ----- + const uint32_t sub_warp_idx = warp_idx - kNumMMAThreads / 32; + const uint32_t upper_row = sub_warp_idx * 16 + lane_idx / 4; + const uint32_t lower_row = upper_row + 8; + const uint32_t col_lane = lane_idx % 4; + + float pm_u[HC_MULT], pm_l[HC_MULT]; + float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT]; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + pm_u[hc] = smem_post[upper_row * HC_MULT + hc]; + pm_l[hc] = smem_post[lower_row * HC_MULT + hc]; + } +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc]; + } + } + + float sqr_u = 0.f, sqr_l = 0.f; + constexpr uint32_t kNumBankGroupBytes = 16; + constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16); + constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup; + constexpr uint32_t BLOCK_M_PER_WARP = BLOCK_M / 4; + static_assert(BLOCK_K * sizeof(nv_bfloat16) == kSwizzleAMode, "BLOCK_K must match swizzle A mode"); + static_assert(kNumLoads % 2 == 0, "kNumLoads must be even for LDSM.x4"); + + uint32_t s = 0; + for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht) + { + const uint32_t i_stage = ht % N_INPUT_STAGES; + full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1); + + uint32_t x_vals[2][kNumLoads]; + { + uint8_t const* x_base + = reinterpret_cast(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr = x_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1], + x_vals[1][i + 1], const_cast(smem_ptr)); + } + } + + uint32_t r_vals[HC_MULT][2][kNumLoads]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + uint8_t const* r_base = reinterpret_cast(smem_res[i_stage]) + + j * BLOCK_M * BLOCK_K * sizeof(nv_bfloat16) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleResMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = r_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(r_vals[j][0][i + 0], r_vals[j][1][i + 0], + r_vals[j][0][i + 1], r_vals[j][1][i + 1], const_cast(smem_ptr)); + } + } + + float2 xf[2][kNumLoads]; +#pragma unroll + for (uint32_t u = 0; u < 2; ++u) + { +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + xf[u][i] = __bfloat1622float2(*reinterpret_cast(&x_vals[u][i])); + } + } + + // Wait for previous ht's residual_out TMA_STOREs to drain before we + // overwrite single-buffered smem_rc with new hc values. + if (ht > 0) + { + cute::tma_store_wait<0>(); + } + +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + const uint32_t cast_stage_idx = s % kNumCastStages; + empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1); + + uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads]; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; ++i) + { + float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y}; + float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + float2 ruj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][0][i])); + float2 rlj = __bfloat1622float2(*reinterpret_cast(&r_vals[j][1][i])); + nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x); + nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y); + nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x); + nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y); + } + nv_bfloat162 b_up = __float22bfloat162_rn(nu); + nv_bfloat162 b_lo = __float22bfloat162_rn(nl); + uint32_t b_up_bits = *reinterpret_cast(&b_up); + uint32_t b_lo_bits = *reinterpret_cast(&b_lo); + rc_u_buf[i] = b_up_bits; + rc_l_buf[i] = b_lo_bits; + float2 ru = __bfloat1622float2(b_up); + float2 rl = __bfloat1622float2(b_lo); + sqr_u = fmaf(ru.x, ru.x, sqr_u); + sqr_u = fmaf(ru.y, ru.y, sqr_u); + sqr_l = fmaf(rl.x, rl.x, sqr_l); + sqr_l = fmaf(rl.y, rl.y, sqr_l); + cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast(&ru.x), + *reinterpret_cast(&ru.y), *reinterpret_cast(&rl.x), + *reinterpret_cast(&rl.y), cast_stage_idx * BLOCK_K + i * 8); + } + cutlass::arch::fence_view_async_tmem_store(); + tcgen05_before_thread_sync(); + full_cast[cast_stage_idx]->arrive(); + ++s; + + // STSM bf16 new_r values into smem_rc[hc] sub-region for this warp. + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; +#pragma unroll + for (uint32_t i = 0; i < kNumLoads; i += 2) + { + auto smem_ptr + = rc_base + get_swizzled_smem_offset(i + lane_idx / 16, lane_idx % 16); + stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]); + } + } + empty_input[i_stage]->arrive(); + + // Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows. + cute::tma_store_fence(); + if (cute::elect_one_sync()) + { + const uint32_t h_idx = (h_tile_start + ht) * BLOCK_K; +#pragma unroll + for (uint32_t hc = 0; hc < HC_MULT; ++hc) + { + uint8_t* rc_base = reinterpret_cast(smem_rc) + hc * SMEM_RC_PER_HC + + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode; + cute::SM90_TMA_STORE_2D::copy(&tensor_map_residual_out, rc_base, hc * HIDDEN + h_idx, + m_offset + sub_warp_idx * BLOCK_M_PER_WARP); + cute::tma_store_arrive(); + } + } + } + + // Drain any in-flight residual_out TMA stores before exit. + cute::tma_store_wait<0>(); + + // Warp-reduce sqr across 4 col_lanes then atomicAdd to global. + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1); + sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1); + sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 2); + if (col_lane == 0) + { + uint32_t gm_u = m_block_idx * BLOCK_M + upper_row; + uint32_t gm_l = m_block_idx * BLOCK_M + lower_row; + if (gm_u < shape_m) + atomicAdd(&sqr_sum[gm_u], sqr_u); + if (gm_l < shape_m) + atomicAdd(&sqr_sum[gm_l], sqr_l); + } + } + + // ======================================================================== + // Phase 3: cross-split barrier. + // For kNumSplits == 1 we only need a block-scope fence + __syncthreads. + // For kNumSplits > 1 ALL splits participate in Phase 4: each CTA + // increments done_counter, then spin-waits until the counter reaches + // kNumSplits. Phase 4 work is then partitioned across CTAs by + // k_split_idx — CTA i processes tokens + // [i * TOKS_PER_CTA, (i+1) * TOKS_PER_CTA) + // where TOKS_PER_CTA = BLOCK_M / kNumSplits. This replaces the old + // "last-home CTA does all Phase 4 work serially" design that bottlenecked + // Path D at BLOCK_M=64 (1 CTA's Phase 4 = ~28 µs regardless of M). + // ======================================================================== + if constexpr (kNumSplits == 1) + { + __threadfence_block(); + __syncthreads(); + } + else + { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + atomicAdd(&done_counter[m_block_idx], 1); + // Spin-wait until all kNumSplits CTAs finish Phase 2. The + // atomicAdd(..., 0) is a zero-increment load with full device + // coherence — cheap on B200 and avoids an extra flag allocation. + while (atomicAdd(&done_counter[m_block_idx], 0) < static_cast(kNumSplits)) + { + /* spin */ + } + } + __syncthreads(); + } + + // ======================================================================== + // Phase 4: inline bigFuse for this CTA's subset of BLOCK_M tokens. + // y_acc[tok, 0..HC_MULT) -> pre_mix (sigmoid) + // y_acc[tok, HC_MULT..2*HC_MULT) -> post_mix_out + // y_acc[tok, 2*HC_MULT..HC_MULT3) -> comb_mix_out (sinkhorn) + // layer_input[tok, h] = sum_j pre_mix[tok, j] * residual_cur[tok, j, h] + // + // Token subset: CTA k_split_idx handles tokens + // [k_split_idx * TOKS_PER_CTA, (k_split_idx + 1) * TOKS_PER_CTA), + // where TOKS_PER_CTA = ceil(BLOCK_M / kNumSplits). At kNumSplits == 1 + // this is the whole m_block (64 tokens); at kNumSplits == 8 it's 8 tokens + // spread over 8 CTAs running Phase 4 concurrently — ~8× the Phase-4 + // throughput of the old single-last-home-CTA design. + // + // Within a CTA, tokens are parallelized across warps: each warp handles + // max(1, TOKS_PER_CTA / NUM_WARPS_BF) tokens serially. Within a warp, + // lanes 0..HC_MULT-1 compute rmsnorm/sigmoid/sinkhorn; pre_mix is + // broadcast via __shfl_sync so all 32 lanes run the layer_input dot + // product across HIDDEN=4096. + // ======================================================================== + constexpr uint32_t BLOCK_SIZE_BF = kNumMMAThreads + kNumPmapThreads; // 256 + constexpr uint32_t WARP_SIZE_BF = 32; + constexpr uint32_t NUM_WARPS_BF = BLOCK_SIZE_BF / WARP_SIZE_BF; // 8 + constexpr uint32_t TOKS_PER_CTA = (BLOCK_M + kNumSplits - 1) / kNumSplits; + // When TOKS_PER_CTA < NUM_WARPS_BF (large kNumSplits / small BLOCK_M case), + // have WARPS_PER_TOK warps cooperate on the HIDDEN-stride layer_input loop + // so all 8 warps stay active. Example at BLOCK_M=64, kNumSplits=16: + // TOKS_PER_CTA=4, WARPS_PER_TOK=2, TOKS_PER_PASS=4, TOKEN_PASSES=1 — + // 2 warps per token each sweep HIDDEN/2, zero idle warps. + constexpr uint32_t WARPS_PER_TOK = (NUM_WARPS_BF > TOKS_PER_CTA) ? (NUM_WARPS_BF / TOKS_PER_CTA) : 1u; + constexpr uint32_t TOKS_PER_PASS = NUM_WARPS_BF / WARPS_PER_TOK; + constexpr uint32_t TOKEN_PASSES = (TOKS_PER_CTA + TOKS_PER_PASS - 1) / TOKS_PER_PASS; + constexpr uint32_t BF16_VEC_LI = 8; + static_assert(HIDDEN % (WARPS_PER_TOK * WARP_SIZE_BF * BF16_VEC_LI) == 0, + "HIDDEN must be a multiple of WARPS_PER_TOK * WARP_SIZE * BF16_VEC_LI"); + const uint32_t tid_bf = threadIdx.x; + const uint32_t lane_bf = tid_bf % WARP_SIZE_BF; + const uint32_t warp_bf = tid_bf / WARP_SIZE_BF; + const uint32_t warp_tok_pos = warp_bf / WARPS_PER_TOK; // which token in a pass + const uint32_t warp_in_team = warp_bf % WARPS_PER_TOK; // which warp inside team + const uint32_t cta_tok_base = k_split_idx * TOKS_PER_CTA; + +#pragma unroll 1 + for (uint32_t pass = 0; pass < TOKEN_PASSES; ++pass) + { + const uint32_t t_in_cta = pass * TOKS_PER_PASS + warp_tok_pos; + if (t_in_cta >= TOKS_PER_CTA) + continue; + const uint32_t t = cta_tok_base + t_in_cta; + if (t >= BLOCK_M) + continue; + const uint32_t tok = m_offset + t; + if (tok >= shape_m) + continue; + + // Lanes 0..HC_MULT-1 compute rmsnorm / sigmoid / sinkhorn; pre_mix is + // held in `pre_mix_local` on lanes 0..HC_MULT-1 and later broadcast to + // all 32 lanes via __shfl_sync. All warps in a team redundantly run + // these ~tens of FLOPs (cheap) to avoid a cross-warp SMEM sync; only + // warp_in_team==0 writes comb_mix_out / post_mix_out to GMEM. + float pre_mix_local = 0.f; + if (lane_bf < HC_MULT) + { + float const r_val = sqr_sum[tok]; + float y_local[HC_MULT3]; + float const* y_row = D + static_cast(tok) * SHAPE_N; +#pragma unroll + for (uint32_t c = 0; c < HC_MULT3; ++c) + y_local[c] = y_row[c]; + + // Reference bigfuse divides by HIDDEN (per-head count), matching Path A. + float const rstd = rsqrtf(r_val / static_cast(HIDDEN) + rms_eps); + float const s0 = hc_scale[0]; + float const s1 = hc_scale[1]; + float const s2 = hc_scale[2]; + + float v = y_local[lane_bf] * rstd * s0 + hc_base[lane_bf]; + pre_mix_local = 1.0f / (1.0f + __expf(-v)) + hc_pre_eps; + + v = y_local[HC_MULT + lane_bf] * rstd * s1 + hc_base[HC_MULT + lane_bf]; + float post_val = 1.0f / (1.0f + __expf(-v)) * hc_post_mult_value; + if (warp_in_team == 0) + { + post_mix_out[tok * HC_MULT + lane_bf] = post_val; + } + + float cm_vals[HC_MULT]; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = y_local[2 * HC_MULT + lane_bf * HC_MULT + k] * rstd * s2 + + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; + + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; + float const rowMax + = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = __expf(cm_vals[k] - rowMax); + float rs = cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] = cm_vals[k] / rs + hc_sinkhorn_eps; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] /= (cs + hc_sinkhorn_eps); + } + for (uint32_t it = 1; it < sinkhorn_repeat; ++it) + { + rs = cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_vals[k] /= rs; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] /= (cs + hc_sinkhorn_eps); + } + } + if (warp_in_team == 0) + { + float* cm_out_ptr = comb_mix_out + tok * HC_MULT2; +#pragma unroll + for (uint32_t k = 0; k < HC_MULT; ++k) + cm_out_ptr[lane_bf * HC_MULT + k] = cm_vals[k]; + } + } + + // Broadcast pre_mix[j] from lane j to all 32 lanes (intra-warp shfl). + float pm[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + 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 + // HIDDEN in strides of WARPS_PER_TOK * 32 * 8. When WARPS_PER_TOK==1, + // each warp sweeps HIDDEN alone (same as the original single-warp case). + __nv_bfloat16 const* rbase = residual_cur_ptr + static_cast(tok) * HC_MULT * HIDDEN; + __nv_bfloat16* obase = layer_input_out + static_cast(tok) * HIDDEN; + + constexpr uint32_t H_STRIDE = WARPS_PER_TOK * WARP_SIZE_BF * BF16_VEC_LI; + const uint32_t h_start = warp_in_team * WARP_SIZE_BF * BF16_VEC_LI + lane_bf * BF16_VEC_LI; + +#pragma unroll + for (uint32_t h = h_start; h < HIDDEN; h += H_STRIDE) + { + // Issue all HC_MULT=4 residual_cur reads first so their L2 latency + // is hidden by the bf16→fp32 arithmetic that follows. The compiler + // schedules the 4 independent __ldg's in parallel. + uint4 raws[HC_MULT]; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + raws[j] = __ldg(reinterpret_cast(&rbase[j * HIDDEN + h])); + } + float acc_li[BF16_VEC_LI] = {}; +#pragma unroll + for (uint32_t j = 0; j < HC_MULT; ++j) + { + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raws[j]); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + { + float2 f = __bfloat1622float2(pairs[v]); + acc_li[2 * v + 0] += pm[j] * f.x; + acc_li[2 * v + 1] += pm[j] * f.y; + } + } + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (uint32_t v = 0; v < BF16_VEC_LI / 2; ++v) + opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_100a"); +#endif +} + +} // namespace fused_mhc + +#pragma clang diagnostic pop diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu new file mode 100644 index 000000000000..e85c868509fa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -0,0 +1,537 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Single-launch fused hyper-connection boundary op. +// +// This implementation wraps the SM100 tcgen05-based residual-out variant +// (fused_tf32_pmap_gemm_rout_atomic_impl) to produce residual_cur, D, and +// sqr_sum in a single kernel launch; the big-fuse postlogue kernel then +// consumes (D, sqr_sum, residual_cur) to emit (post_mix_cur, comb_mix_cur, +// layer_input_cur). Semantically identical to +// +// residual_cur = prev_mHC.post_mapping(x_prev, residual_prev, ...) +// post_mix_cur, comb_mix_cur, layer_input_cur = self.pre_mapping(residual_cur) +// +// but exposed as a single entry point. + +#include "fused_tf32_pmap_gemm.cuh" +#include "mhcKernels.h" +#include "mhc_fused_fma.cuh" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc +{ + +// ---- Single-launch workspace zero kernel ----------------------------------- +// +// Replaces 2-3 separate cudaMemsetAsync calls for the atomic accumulator +// workspaces (y_acc, r_acc, optional done_counter). Avoids per-memset launch +// latency that is visible at small M / high-frequency inference. +namespace +{ + +__global__ void fhcZeroWorkspacesKernel( + float* __restrict__ y_acc, uint32_t y_elems, float* __restrict__ r_acc, uint32_t r_elems, + int* __restrict__ done_counter, uint32_t done_elems) +{ + uint32_t const tid = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t const stride = gridDim.x * blockDim.x; + for (uint32_t i = tid; i < y_elems; i += stride) + { + y_acc[i] = 0.0f; + } + for (uint32_t i = tid; i < r_elems; i += stride) + { + r_acc[i] = 0.0f; + } + if (done_counter != nullptr) + { + for (uint32_t i = tid; i < done_elems; i += stride) + { + done_counter[i] = 0; + } + } +} + +inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint32_t r_elems, + int* done_counter, uint32_t done_elems, cudaStream_t stream) +{ + uint32_t const total = y_elems + r_elems + (done_counter != nullptr ? done_elems : 0u); + if (total == 0u) + { + return; + } + constexpr uint32_t kBlock = 256; + // Cap grid so we don't over-launch for small workspaces; 148 SMs on B200. + uint32_t const num_blocks = min(static_cast((total + kBlock - 1) / kBlock), 148u * 8u); + fhcZeroWorkspacesKernel<<>>( + y_acc, y_elems, r_acc, r_elems, done_counter, done_elems); +} + +} // namespace + + +// ---- mHC fused kernel shape constants (mirrors the Python module) ---- +static constexpr uint32_t FHC_SHAPE_N = 24; // HC_MULT * (2 + HC_MULT) = 4 * 6 = 24 +static constexpr uint32_t FHC_HIDDEN = 4096; // only this hidden size is currently wired up +static constexpr uint32_t FHC_HC_MULT = 4; +static constexpr uint32_t FHC_BLOCK_M = 64; +static constexpr uint32_t FHC_BLOCK_N = 32; +static constexpr uint32_t FHC_BLOCK_K = 64; +static constexpr uint32_t FHC_SWIZZLE_CD = 128; +static constexpr uint32_t FHC_N_B_STAGES = 12; +static constexpr uint32_t FHC_N_INPUT_STG = 2; +static constexpr uint32_t FHC_NUM_MMA_TH = 128; +static constexpr uint32_t FHC_NUM_PMAP_TH = 128; + +static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter, + uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, uint32_t swizzleBytes, uint32_t elemBytes) +{ + CUtensorMap tm; + if (swizzleBytes != 0) + { + smemInner = swizzleBytes / elemBytes; + } + uint64_t gmemDims[2] = {gmemInner, gmemOuter}; + uint32_t smemDims[2] = {smemInner, smemOuter}; + uint64_t gmemStrides[1] = {gmemOuterStrideBytes}; + uint32_t elemStrides[2] = {1, 1}; + CUtensorMapSwizzle swizzle = (swizzleBytes == 128) ? CU_TENSOR_MAP_SWIZZLE_128B + : (swizzleBytes == 64) ? CU_TENSOR_MAP_SWIZZLE_64B + : (swizzleBytes == 32) ? CU_TENSOR_MAP_SWIZZLE_32B + : CU_TENSOR_MAP_SWIZZLE_NONE; + CUresult rc = cuTensorMapEncodeTiled(&tm, dtype, 2, base, gmemDims, gmemStrides, smemDims, elemStrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TLLM_CHECK_WITH_INFO(rc == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed (%d)", static_cast(rc)); + return tm; +} + +static constexpr uint32_t fhcSmemSize() +{ + constexpr uint32_t SMEM_CD = FHC_BLOCK_M * FHC_SWIZZLE_CD; + constexpr uint32_t SMEM_B = FHC_BLOCK_N * FHC_BLOCK_K * sizeof(float); + constexpr uint32_t SMEM_RES_ISTG = FHC_BLOCK_M * FHC_HC_MULT * FHC_BLOCK_K * sizeof(__nv_bfloat16); + constexpr uint32_t SMEM_X_ISTG = FHC_BLOCK_M * FHC_BLOCK_K * sizeof(__nv_bfloat16); + constexpr uint32_t SMEM_POST = FHC_BLOCK_M * FHC_HC_MULT * sizeof(float); + constexpr uint32_t SMEM_COMB = FHC_BLOCK_M * FHC_HC_MULT * FHC_HC_MULT * sizeof(float); + constexpr uint32_t SMEM_RC = FHC_HC_MULT * FHC_BLOCK_M * FHC_BLOCK_K * sizeof(__nv_bfloat16); + constexpr uint32_t kNumCast = 4; + constexpr uint32_t barriers = 2 * FHC_N_B_STAGES + 2 * FHC_N_INPUT_STG + 2 * kNumCast + 1; + // 4 bytes for the tmem ptr word + 32 bytes padding for alignment headroom. + return SMEM_CD + FHC_N_B_STAGES * SMEM_B + FHC_N_INPUT_STG * (SMEM_RES_ISTG + SMEM_X_ISTG) + SMEM_POST + SMEM_COMB + + SMEM_RC + barriers * 8 + 4 + 32; +} + +using FusedRoutFn = void (*)( + uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, float*, float const*, float const*, float*); + +template +static FusedRoutFn fhcInstance() +{ + return &fused_mhc::fused_tf32_pmap_gemm_rout_atomic_impl; +} + +static FusedRoutFn pickFhc(uint32_t ks) +{ + switch (ks) + { + case 1: return fhcInstance<1>(); + case 2: return fhcInstance<2>(); + case 4: return fhcInstance<4>(); + case 8: return fhcInstance<8>(); + case 16: return fhcInstance<16>(); + case 32: return fhcInstance<32>(); + case 64: return fhcInstance<64>(); + default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported kNumSplits=%u", ks); return nullptr; + } +} + +// Heuristic split-K selection driven by M. Keeps large-M waves compute-bound +// (low splits) and low-M waves fully occupied (high splits). +// +// NOTE: split > 1 introduces atomic accumulation into D and sqr_sum. Atomic +// ordering is non-deterministic across runs, which breaks CUDA-graph bit-exact +// replay equality. Consumers that require bit-exact determinism should keep +// M small enough to land in the splits=1 bucket, or supply an explicit +// deterministic variant if needed. +static uint32_t pickKSplits(int M) +{ + if (M <= 512) + return 16; + if (M <= 1024) + return 8; + if (M <= 2048) + return 4; + if (M <= 4096) + return 2; + return 1; +} + +static int selectBigFuseBS(int M) +{ + if (M >= 4096) + return 128; + if (M >= 256) + return 256; + return 512; +} + +void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev, + float const* comb_mix_prev, float const* w_t, float const* hc_scale, float const* hc_base, + __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, __nv_bfloat16* layer_input_cur, + float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int num_k_splits, + int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, + int sinkhorn_repeat, cudaStream_t stream) +{ + if (M <= 0) + return; + + TLLM_CHECK_WITH_INFO(hidden_size == static_cast(FHC_HIDDEN), + "mhcFusedHcLaunch: hidden_size=%d not supported (only %u)", hidden_size, FHC_HIDDEN); + TLLM_CHECK_WITH_INFO(hc_mult == static_cast(FHC_HC_MULT), + "mhcFusedHcLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT); + + constexpr uint32_t SHAPE_K = FHC_HC_MULT * FHC_HIDDEN; + + uint32_t const m_u = static_cast(M); + uint32_t const ks = (num_k_splits > 0) ? static_cast(num_k_splits) : pickKSplits(M); + int const bs = (bigfuse_block_size > 0) ? bigfuse_block_size : selectBigFuseBS(M); + + // ---- Zero workspace buffers (atomic accumulators) ---- + fhcZeroWorkspaces(y_acc_workspace, static_cast(M) * FHC_SHAPE_N, r_acc_workspace, + static_cast(M), /*done_counter=*/nullptr, /*done_elems=*/0, stream); + + // ---- Build TMA descriptors ---- + CUtensorMap desc_res = makeTma2D(const_cast<__nv_bfloat16*>(residual_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), + /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + + CUtensorMap desc_x = makeTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, FHC_HIDDEN, + m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(FHC_HIDDEN) * sizeof(__nv_bfloat16), + /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + + CUtensorMap desc_b = makeTma2D(const_cast(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N, + FHC_BLOCK_K, FHC_BLOCK_N, static_cast(SHAPE_K) * sizeof(float), + /*swizzleBytes=*/128, sizeof(float)); + + CUtensorMap desc_res_out = makeTma2D(residual_cur, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, + /*smemOuter=*/16, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), + /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + + // ---- Step 1: fused post-mapping + TF32 GEMM + sqrsum + residual_out ---- + constexpr uint32_t fused_smem = fhcSmemSize(); + FusedRoutFn fa = pickFhc(ks); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + reinterpret_cast(fa), cudaFuncAttributeMaxDynamicSharedMemorySize, fused_smem)); + + uint32_t const m_tiles = (m_u + FHC_BLOCK_M - 1) / FHC_BLOCK_M; + dim3 const grid(m_tiles * ks); + dim3 const block(FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH); + fa<<>>( + m_u, desc_res, desc_x, desc_b, desc_res_out, y_acc_workspace, post_mix_prev, comb_mix_prev, r_acc_workspace); + + // ---- Step 2: big-fuse postlogue (RMS + sigmoid + Sinkhorn + pre-apply) ---- + // Delegate to mhcBigFuseLaunch (defined in mhcKernels.cu) to avoid + // instantiating the mhcBigFuseKernel template in this TU. + mhcBigFuseLaunch(y_acc_workspace, r_acc_workspace, residual_cur, hc_scale, hc_base, post_mix_cur, comb_mix_cur, + layer_input_cur, M, /*K=*/hidden_size, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, + sinkhorn_repeat, /*num_splits=*/1, /*block_size=*/bs, stream); +} + +// =================================================================== +// FMA-path fused hyper-connection boundary launcher. +// +// Uses `fused_fma_kernels::fused_pmap_gemm_fma_ksplit` +// which computes pmap inline in registers, emits residual_cur to HBM, and +// writes split GEMM partials to Yp[ks, M, N] / Rp[ks, M]. The same bigfuse +// kernel used by mhc_pre_mapping consumes those split buffers via num_splits. +// =================================================================== + +using FmaKsplitFn = void (*)(__nv_bfloat16 const*, __nv_bfloat16 const*, float const*, float const*, float const*, + float*, float*, int, int, int, __nv_bfloat16*); + +template +static FmaKsplitFn fhcFmaInstance() +{ + return &fused_fma_kernels::fused_pmap_gemm_fma_ksplit; +} + +// Valid (tile_n, num_k_splits) combinations the fused_hc FMA path supports. +// Keep this limited to the small/mid-M sweet spots from profile_fair_report v4. +static FmaKsplitFn pickFhcFma(int tile_n, int ks) +{ +#define FHCFMA_CASE(TN, KS) \ + if (tile_n == (TN) && ks == (KS)) \ + return fhcFmaInstance() + + FHCFMA_CASE(1, 1); + FHCFMA_CASE(1, 2); + FHCFMA_CASE(1, 4); + FHCFMA_CASE(1, 8); + FHCFMA_CASE(2, 1); + FHCFMA_CASE(2, 2); + FHCFMA_CASE(2, 4); + FHCFMA_CASE(2, 8); + FHCFMA_CASE(3, 1); + FHCFMA_CASE(3, 2); + FHCFMA_CASE(3, 4); + FHCFMA_CASE(4, 1); + FHCFMA_CASE(4, 2); + FHCFMA_CASE(6, 1); + FHCFMA_CASE(8, 1); + FHCFMA_CASE(12, 1); + FHCFMA_CASE(24, 1); +#undef FHCFMA_CASE + TLLM_CHECK_WITH_INFO(false, "mhcFusedHcFmaLaunch: unsupported (tile_n=%d, ks=%d)", tile_n, ks); + return nullptr; +} + +void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev, + float const* comb_mix_prev, float const* w_t, float const* hc_scale, float const* hc_base, + __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, __nv_bfloat16* layer_input_cur, + float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int tile_n, int num_k_splits, + int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, + int sinkhorn_repeat, cudaStream_t stream) +{ + if (M <= 0) + return; + + TLLM_CHECK_WITH_INFO(hc_mult == static_cast(FHC_HC_MULT), + "mhcFusedHcFmaLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT); + TLLM_CHECK_WITH_INFO( + FHC_SHAPE_N % tile_n == 0, "mhcFusedHcFmaLaunch: SHAPE_N=%u not divisible by tile_n=%d", FHC_SHAPE_N, tile_n); + + int const K = hc_mult * hidden_size; + int const N = static_cast(FHC_SHAPE_N); + + // ---- Step 1: fused pmap + GEMM + sqrsum + residual_cur (FMA ksplit) ---- + FmaKsplitFn fn = pickFhcFma(tile_n, num_k_splits); + dim3 const grid(static_cast(M), static_cast(N / tile_n), static_cast(num_k_splits)); + dim3 const block(256); + fn<<>>(residual_prev, x_prev, post_mix_prev, comb_mix_prev, w_t, y_acc_workspace, + r_acc_workspace, hidden_size, N, K, residual_cur); + + // ---- Step 2: big-fuse postlogue (reduces ks splits internally) ---- + mhcBigFuseLaunch(y_acc_workspace, r_acc_workspace, residual_cur, hc_scale, hc_base, post_mix_cur, comb_mix_cur, + layer_input_cur, M, /*K=*/hidden_size, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, + sinkhorn_repeat, /*num_splits=*/num_k_splits, bigfuse_block_size, stream); +} + +// =================================================================== +// All-in-one single-kernel fused hyper-connection launcher (TF32 tcgen05). +// +// Wraps fused_allinone_tf32_pmap_gemm_atomic_impl: pmap + GEMM + bigfuse +// all fused into one kernel. Phase 3 elects last-home CTA per m-block via +// atomic done_counter; Phase 4 runs bigfuse inline on the elected CTA. +// =================================================================== + +// Path D (all-in-one) reuses Path B's SMEM layout: CD + B stages + res/x +// stages + post + comb + rc (HC_MULT slices) + barriers. The inline bigfuse +// tail uses static __shared__ scratch (s_pre_mix, s_is_last) which is NOT +// counted in the dynamic SMEM budget, so this size matches fhcSmemSize(). +static constexpr uint32_t fhcAllInOneSmemSize() +{ + return fhcSmemSize(); +} + +using FusedAllInOneFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, __nv_bfloat16 const*, + __nv_bfloat16*, float*, float*, int*, float const*, float const*, float const*, float const*, float*, float*, float, + float, float, float, uint32_t); + +template +static FusedAllInOneFn fhcAllInOneInstance() +{ + return &fused_mhc::fused_allinone_tf32_pmap_gemm_atomic_impl; +} + +static FusedAllInOneFn pickFhcAllInOne(uint32_t ks) +{ + switch (ks) + { + case 1: return fhcAllInOneInstance<1>(); + case 2: return fhcAllInOneInstance<2>(); + case 4: return fhcAllInOneInstance<4>(); + case 8: return fhcAllInOneInstance<8>(); + case 16: return fhcAllInOneInstance<16>(); + case 32: return fhcAllInOneInstance<32>(); + case 64: return fhcAllInOneInstance<64>(); + default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcAllInOneLaunch: unsupported kNumSplits=%u", ks); return nullptr; + } +} + +void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, + float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale, + float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, + __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int* done_counter_workspace, int M, + int hidden_size, int hc_mult, int num_k_splits, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, + float hc_post_mult_value, int sinkhorn_repeat, cudaStream_t stream) +{ + if (M <= 0) + return; + + TLLM_CHECK_WITH_INFO(hidden_size == static_cast(FHC_HIDDEN), + "mhcFusedHcAllInOneLaunch: hidden_size=%d not supported (only %u)", hidden_size, FHC_HIDDEN); + TLLM_CHECK_WITH_INFO(hc_mult == static_cast(FHC_HC_MULT), + "mhcFusedHcAllInOneLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT); + + constexpr uint32_t SHAPE_K = FHC_HC_MULT * FHC_HIDDEN; + + uint32_t const m_u = static_cast(M); + uint32_t const ks = (num_k_splits > 0) ? static_cast(num_k_splits) : 1u; + uint32_t const m_tiles = (m_u + FHC_BLOCK_M - 1) / FHC_BLOCK_M; + + // ---- Zero workspace buffers (atomic accumulators + done counter) ---- + fhcZeroWorkspaces(y_acc_workspace, static_cast(M) * FHC_SHAPE_N, r_acc_workspace, + static_cast(M), done_counter_workspace, m_tiles, stream); + + // ---- Build TMA descriptors ---- + CUtensorMap desc_res = makeTma2D(const_cast<__nv_bfloat16*>(residual_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), + /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + + CUtensorMap desc_x = makeTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, FHC_HIDDEN, + m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast(FHC_HIDDEN) * sizeof(__nv_bfloat16), + /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + + CUtensorMap desc_b = makeTma2D(const_cast(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N, + FHC_BLOCK_K, FHC_BLOCK_N, static_cast(SHAPE_K) * sizeof(float), + /*swizzleBytes=*/128, sizeof(float)); + + CUtensorMap desc_res_out = makeTma2D(residual_cur, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, SHAPE_K, m_u, FHC_BLOCK_K, + /*smemOuter=*/16, static_cast(SHAPE_K) * sizeof(__nv_bfloat16), + /*swizzleBytes=*/128, sizeof(__nv_bfloat16)); + + // ---- Launch the single all-in-one kernel ---- + constexpr uint32_t fused_smem = fhcAllInOneSmemSize(); + FusedAllInOneFn fa = pickFhcAllInOne(ks); + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + reinterpret_cast(fa), cudaFuncAttributeMaxDynamicSharedMemorySize, fused_smem)); + + dim3 const grid(m_tiles * ks); + dim3 const block(FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH); + fa<<>>(m_u, desc_res, desc_x, desc_b, desc_res_out, residual_cur, layer_input_cur, + y_acc_workspace, r_acc_workspace, done_counter_workspace, post_mix_prev, comb_mix_prev, hc_scale, hc_base, + post_mix_cur, comb_mix_cur, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, + static_cast(sinkhorn_repeat)); +} + +// =================================================================== +// All-in-one single-kernel fused hyper-connection launcher (FMA). +// +// Wraps fused_pmap_gemm_fma_allinone: pmap + FMA GEMM + bigfuse +// all fused into one kernel. Same last-home CTA election pattern as the +// TF32 all-in-one path. Preferred for small-M (M <= 32). +// =================================================================== + +using FmaAllInOneFn = void (*)(__nv_bfloat16 const*, __nv_bfloat16 const*, float const*, float const*, float const*, + float const*, float const*, __nv_bfloat16*, float*, float*, __nv_bfloat16*, float*, float*, int*, int, int, int, + float, float, float, float, int); + +template +static FmaAllInOneFn fhcFmaAllInOneInstance() +{ + return &fused_fma_kernels::fused_pmap_gemm_fma_allinone; +} + +static FmaAllInOneFn pickFhcFmaAllInOne(int tile_n, int ks, int tile_m) +{ +#define FHC_FMA_AIO_CASE(TN, KS, TM) \ + if (tile_n == (TN) && ks == (KS) && tile_m == (TM)) \ + return fhcFmaAllInOneInstance() + + FHC_FMA_AIO_CASE(1, 1, 1); + FHC_FMA_AIO_CASE(1, 2, 1); + FHC_FMA_AIO_CASE(2, 1, 1); + FHC_FMA_AIO_CASE(2, 2, 1); + FHC_FMA_AIO_CASE(3, 1, 1); + FHC_FMA_AIO_CASE(4, 1, 1); + FHC_FMA_AIO_CASE(6, 1, 1); + FHC_FMA_AIO_CASE(8, 1, 1); + FHC_FMA_AIO_CASE(12, 1, 1); + FHC_FMA_AIO_CASE(24, 1, 1); + FHC_FMA_AIO_CASE(1, 1, 2); + FHC_FMA_AIO_CASE(1, 2, 2); + FHC_FMA_AIO_CASE(2, 1, 2); + FHC_FMA_AIO_CASE(2, 2, 2); + FHC_FMA_AIO_CASE(3, 1, 2); + FHC_FMA_AIO_CASE(4, 1, 2); + FHC_FMA_AIO_CASE(6, 1, 2); + FHC_FMA_AIO_CASE(8, 1, 2); + FHC_FMA_AIO_CASE(12, 1, 2); + FHC_FMA_AIO_CASE(24, 1, 2); + FHC_FMA_AIO_CASE(1, 1, 4); + FHC_FMA_AIO_CASE(1, 2, 4); + FHC_FMA_AIO_CASE(2, 1, 4); + FHC_FMA_AIO_CASE(2, 2, 4); + FHC_FMA_AIO_CASE(3, 1, 4); + FHC_FMA_AIO_CASE(4, 1, 4); + FHC_FMA_AIO_CASE(6, 1, 4); + FHC_FMA_AIO_CASE(8, 1, 4); + FHC_FMA_AIO_CASE(12, 1, 4); + FHC_FMA_AIO_CASE(24, 1, 4); +#undef FHC_FMA_AIO_CASE + TLLM_CHECK_WITH_INFO( + false, "mhcFusedHcFmaAllInOneLaunch: unsupported (tile_n=%d, ks=%d, tile_m=%d)", tile_n, ks, tile_m); + return nullptr; +} + +void mhcFusedHcFmaAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, + float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale, + float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, + __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int* done_counter_workspace, int M, + int hidden_size, int hc_mult, int tile_n, int num_k_splits, int tile_m, float rms_eps, float hc_pre_eps, + float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat, cudaStream_t stream) +{ + if (M <= 0) + return; + + TLLM_CHECK_WITH_INFO(hc_mult == static_cast(FHC_HC_MULT), + "mhcFusedHcFmaAllInOneLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT); + TLLM_CHECK_WITH_INFO(FHC_SHAPE_N % tile_n == 0, + "mhcFusedHcFmaAllInOneLaunch: SHAPE_N=%u not divisible by tile_n=%d", FHC_SHAPE_N, tile_n); + + int const K = hc_mult * hidden_size; + int const N = static_cast(FHC_SHAPE_N); + int const m_batches = (M + tile_m - 1) / tile_m; + + // ---- Zero workspace buffers (atomic accumulators + done counter) ---- + fhcZeroWorkspaces(y_acc_workspace, static_cast(M) * static_cast(N), r_acc_workspace, + static_cast(M), done_counter_workspace, static_cast(m_batches), stream); + + FmaAllInOneFn fn = pickFhcFmaAllInOne(tile_n, num_k_splits, tile_m); + dim3 const grid( + static_cast(m_batches), static_cast(N / tile_n), static_cast(num_k_splits)); + dim3 const block(256); + fn<<>>(residual_prev, x_prev, post_mix_prev, comb_mix_prev, w_t, hc_scale, hc_base, + residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace, r_acc_workspace, + done_counter_workspace, M, K, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, + sinkhorn_repeat); +} + +} // namespace kernels::mhc + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu new file mode 100644 index 000000000000..a579e2ec5c89 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu @@ -0,0 +1,772 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "mhcKernels.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc +{ + +// =================================================================== +// Kernel 1: big_fuse — one CTA per token +// +// Template parameters: +// NUM_SPLITS: split-K reduction (1 = direct, >1 = sum across splits) +// BLOCK_SIZE: threads per CTA (multiple of 32, >= 64) +// +// Phase 1a (warp 0, lanes 0-3): RMS norm + sigmoid → s_pre_mix, post_mix +// ── __syncthreads ── +// Phase 1b (warp 0, lanes 0-3) ‖ Phase 2 (remaining warps) — overlapped +// 1b: parallel Sinkhorn (4 lanes, __shfl_xor col normalize) → comb_mix +// 2: stream residual × pre_mix → layer_input +// =================================================================== + +template +__launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __restrict__ y_acc, + float const* __restrict__ r_acc, __nv_bfloat16 const* __restrict__ residual, float const* __restrict__ hc_scale, + float const* __restrict__ hc_base, float* __restrict__ post_mix, float* __restrict__ comb_mix, + __nv_bfloat16* __restrict__ layer_input, int M, int K, int hidden_size, float rms_eps, float hc_pre_eps, + float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat) +{ + constexpr int HC_MULT = 4; + constexpr int HC_MULT2 = HC_MULT * HC_MULT; // 16 comb_mix entries + constexpr int HC_MULT3 = HC_MULT * (2 + HC_MULT); // 24 = pre(4) + post(4) + comb(16) + constexpr int WARP_SIZE = 32; + constexpr int BF16_VEC = 8; // bf16 elements per uint4 load + + int const token = blockIdx.x; + int const tid = threadIdx.x; + int const warp_id = tid / WARP_SIZE; + int const lane = tid % WARP_SIZE; + + __shared__ float s_pre_mix[HC_MULT]; + + float cm[HC_MULT]; + + // ---- Phase 1a (warp 0, lanes 0..3): split-K reduce → RMS norm → sigmoid ---- + // Each lane handles one of the 4 hc_mult slots. + // Produces: s_pre_mix[4] (shared), post_mix[4] (global), cm[4×4] (registers). + if (warp_id == 0 && lane < HC_MULT) + { + float r_val; + float y_local[HC_MULT3]; + + if constexpr (NUM_SPLITS == 1) + { + r_val = r_acc[token]; + float const* y_row = y_acc + token * HC_MULT3; +#pragma unroll + for (int c = 0; c < HC_MULT3; c++) + y_local[c] = y_row[c]; + } + else + { + // Reduce across split-K partials + r_val = 0.0f; +#pragma unroll + for (int c = 0; c < HC_MULT3; c++) + y_local[c] = 0.0f; + for (int s = 0; s < NUM_SPLITS; s++) + { + r_val += r_acc[s * M + token]; + float const* y_row = y_acc + (static_cast(s) * M + token) * HC_MULT3; +#pragma unroll + for (int c = 0; c < HC_MULT3; c++) + y_local[c] += y_row[c]; + } + } + + // RMS norm: rstd = 1/sqrt(mean(x²) + eps) + float const rstd = rsqrtf(r_val / static_cast(K) + rms_eps); + float const s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; + + // y_local layout: [pre_mix(4) | post_mix(4) | comb_mix(4×4)] + // pre_mix: sigmoid(norm * scale0 + base) + eps → shared for Phase 2 + float v = y_local[lane] * rstd * s0 + hc_base[lane]; + s_pre_mix[lane] = 1.0f / (1.0f + expf(-v)) + hc_pre_eps; + + // post_mix: sigmoid(norm * scale1 + base) * mult → global + v = y_local[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; + post_mix[token * HC_MULT + lane] = 1.0f / (1.0f + expf(-v)) * hc_post_mult_value; + + // comb_mix init: norm * scale2 + base → cm[4] per lane (one row of 4×4 matrix) +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm[k] = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; + } + + __syncthreads(); + + // ---- Phase 1b (warp 0, lanes 0..3): Sinkhorn normalization ---- + // Each lane holds one row of the 4×4 comb matrix. + // Row normalize via local sum, column normalize via __shfl_xor across 4 lanes. + if (warp_id == 0 && lane < HC_MULT) + { + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; // 0xf for HC_MULT=4 + + // Softmax rows: subtract the row max to avoid inf / inf when comb logits + // are large. + float const rowMax = fmaxf(fmaxf(cm[0], cm[1]), fmaxf(cm[2], cm[3])); +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm[k] = expf(cm[k] - rowMax); + float rs = cm[0] + cm[1] + cm[2] + cm[3]; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm[k] = cm[k] / rs + hc_sinkhorn_eps; + + // Column normalize: sum across lanes (rows) via butterfly shuffle +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + float cs = cm[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm[k] /= (cs + hc_sinkhorn_eps); + } + + // Remaining Sinkhorn iterations: alternate row / column normalize + for (int it = 1; it < sinkhorn_repeat; it++) + { + rs = cm[0] + cm[1] + cm[2] + cm[3] + hc_sinkhorn_eps; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm[k] /= rs; + +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + float cs = cm[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm[k] /= (cs + hc_sinkhorn_eps); + } + } + + // Write 4×4 comb_mix to global (lane = row index, k = col index) + float* cm_out = comb_mix + token * HC_MULT2; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm_out[lane * HC_MULT + k] = cm[k]; + } + + // ---- Phase 2 (warps 1..N, overlapped with Phase 1b): weighted residual sum ---- + // layer_input[h] = sum_j pre_mix[j] * residual[j][h] + // Vectorized: 8 bf16 per thread per iteration via uint4 (LDG.128). + if (warp_id > 0) + { + float pm[HC_MULT]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + pm[j] = s_pre_mix[j]; + + __nv_bfloat16 const* rbase = residual + static_cast(token) * HC_MULT * hidden_size; + __nv_bfloat16* obase = layer_input + static_cast(token) * hidden_size; + + int const p2_tid = tid - WARP_SIZE; + constexpr int p2_threads = BLOCK_SIZE - WARP_SIZE; + + for (int h = p2_tid * BF16_VEC; h < hidden_size; h += p2_threads * BF16_VEC) + { + float acc[BF16_VEC] = {}; + +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + uint4 raw = *reinterpret_cast(&rbase[j * hidden_size + h]); + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raw); +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + { + float2 f = __bfloat1622float2(pairs[v]); + acc[2 * v + 0] += pm[j] * f.x; + acc[2 * v + 1] += pm[j] * f.y; + } + } + + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + opairs[v] = __float22bfloat162_rn(make_float2(acc[2 * v], acc[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; + } + } +} + +#define INST_BIGFUSE(NS, BS) \ + template __global__ void mhcBigFuseKernel(float const*, float const*, __nv_bfloat16 const*, float const*, \ + float const*, float*, float*, __nv_bfloat16*, int, int, int, float, float, float, float, int); + +INST_BIGFUSE(1, 128) +INST_BIGFUSE(1, 256) +INST_BIGFUSE(1, 512) +INST_BIGFUSE(2, 128) +INST_BIGFUSE(2, 256) +INST_BIGFUSE(2, 512) +INST_BIGFUSE(4, 128) +INST_BIGFUSE(4, 256) +INST_BIGFUSE(4, 512) +INST_BIGFUSE(8, 128) +INST_BIGFUSE(8, 256) +INST_BIGFUSE(8, 512) +INST_BIGFUSE(16, 128) +INST_BIGFUSE(16, 256) +INST_BIGFUSE(16, 512) +#undef INST_BIGFUSE + +// =================================================================== +// Kernel 3: gemm_sqrsum_fma — split-N FP32 FMA GEMM with fused sqrsum +// +// Y[m, n] = dot(X[m, :], W_T[n, :]) for each (m, n) +// R[m] = sum_k X[m, k]^2 (only computed once per row) +// +// Grid: (M, ceil(N / N_PER_BLOCK)) +// Block: 256 threads = 8 warps × 32 lanes +// +// Each block handles 1 row × N_PER_BLOCK output columns × full K. +// The sqrsum R[row] is fused into the first N-tile (n_start==0) to +// avoid a separate reduction pass. +// +// Data layout: X is row-major bf16, W_T is row-major fp32 (transposed). +// Uses PTX prmt.b32 for bf16→fp32 zero-extension instead of cvt, +// and explicit cache hints (ld.global.cs for X, L1::evict_last for W_T) +// since X is reused across N-tiles but W_T is streamed once. +// +// Warp-level reduction via __shfl_xor → shared memory → cross-warp sum. +// N_PER_BLOCK columns are processed in groups of NUM_WARPS (8), mapping +// each warp to one output column for the final shared-memory reduce. +// =================================================================== + +template +__launch_bounds__(256) __global__ void mhcGemmSqrsumFmaKernel(__nv_bfloat16 const* __restrict__ X, + float const* __restrict__ W_T, float* __restrict__ Y, float* __restrict__ R, int M, int N, int K) +{ + constexpr int BLOCK_SIZE = 256; + constexpr int WARP_SIZE = 32; + constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; + constexpr int K_VEC = 4; // bf16 elements loaded per thread per iteration + constexpr int K_STEP = BLOCK_SIZE * K_VEC; // K elements consumed per block per iteration + + int const tid = threadIdx.x; + int const warp_id = tid / WARP_SIZE; + int const lane = tid % WARP_SIZE; + int const row = blockIdx.x; + int const n_start = blockIdx.y * N_PER_BLOCK; + + if (row >= M || n_start >= N) + return; + __nv_bfloat16 const* x_row = X + static_cast(row) * K; + + float acc[N_PER_BLOCK]; +#pragma unroll + for (int n = 0; n < N_PER_BLOCK; n++) + acc[n] = 0.0f; + float sqr = 0.0f; + bool const do_sqr = (n_start == 0); + + // Main loop: process K in chunks of K_STEP (256 threads × 4 = 1024 elements) + for (int k_base = 0; k_base + K_STEP <= K; k_base += K_STEP) + { + int const my_k = k_base + tid * K_VEC; + + // Load 4 bf16 from X as packed u32 pair (cache-streaming hint: X reused across N-tiles) + unsigned xp0, xp1; + asm volatile("ld.global.cs.v2.b32 {%0, %1}, [%2];" : "=r"(xp0), "=r"(xp1) : "l"(x_row + my_k)); + + // bf16→fp32 via prmt.b32: zero-extend each 16-bit half into a 32-bit float + // prmt selectors: 0x5410 extracts low halfword, 0x7610 extracts high halfword + float xv0, xv1, xv2, xv3; + asm volatile( + "{ .reg .b32 t0, t1, t2, t3;\n\t" + " prmt.b32 t0, %4, %5, %6;\n\t" + " prmt.b32 t1, %4, %5, %7;\n\t" + " prmt.b32 t2, %4, %8, %6;\n\t" + " prmt.b32 t3, %4, %8, %7;\n\t" + " mov.b32 %0, t0;\n\t mov.b32 %1, t1;\n\t" + " mov.b32 %2, t2;\n\t mov.b32 %3, t3; }" + : "=f"(xv0), "=f"(xv1), "=f"(xv2), "=f"(xv3) + : "r"(0), "r"(xp0), "r"(0x5410), "r"(0x7610), "r"(xp1)); + + if (do_sqr) + { + sqr = fmaf(xv0, xv0, sqr); + sqr = fmaf(xv1, xv1, sqr); + sqr = fmaf(xv2, xv2, sqr); + sqr = fmaf(xv3, xv3, sqr); + } + + // Dot product: acc[n] += X[k:k+4] · W_T[n][k:k+4] + // W_T uses evict_last hint since each weight row is accessed only once +#pragma unroll + for (int n = 0; n < N_PER_BLOCK; n++) + { + float w0, w1, w2, w3; + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w0), "=f"(w1), "=f"(w2), "=f"(w3) + : "l"(W_T + static_cast(n_start + n) * K + my_k)); + acc[n] = fmaf(xv0, w0, acc[n]); + acc[n] = fmaf(xv1, w1, acc[n]); + acc[n] = fmaf(xv2, w2, acc[n]); + acc[n] = fmaf(xv3, w3, acc[n]); + } + } + + // Tail: handle remaining K elements (K % K_STEP), scalar loads + { + int const tail_start = K - (K % K_STEP); + for (int kk = tail_start + tid; kk < K; kk += BLOCK_SIZE) + { + float xv; + asm volatile( + "{ .reg .b16 tmp;\n\t" + " ld.global.cs.b16 tmp, [%1];\n\t" + " cvt.f32.bf16 %0, tmp; }" + : "=f"(xv) + : "l"(x_row + kk)); + if (do_sqr) + sqr = fmaf(xv, xv, sqr); +#pragma unroll + for (int n = 0; n < N_PER_BLOCK; n++) + { + float wv = W_T[static_cast(n_start + n) * K + kk]; + acc[n] = fmaf(xv, wv, acc[n]); + } + } + } + + // ---- Cross-warp reduction ---- + // N_PER_BLOCK columns are grouped into chunks of NUM_WARPS (8). + // Within each group: intra-warp butterfly → smem → one warp reads the column sums. + // s_warp[warp_id][col] holds per-warp partial; extra column SQRSUM_SLOT for sqrsum. + constexpr int COLS_PER_GROUP = NUM_WARPS; + constexpr int SQRSUM_SLOT = COLS_PER_GROUP; + constexpr int N_GROUPS = (N_PER_BLOCK + COLS_PER_GROUP - 1) / COLS_PER_GROUP; + constexpr unsigned FULL_WARP_MASK = 0xffffffff; + + __shared__ float s_warp[NUM_WARPS][COLS_PER_GROUP + 1]; + +#pragma unroll + for (int g = 0; g < N_GROUPS; g++) + { + constexpr int LAST_COLS = N_PER_BLOCK - (N_GROUPS - 1) * COLS_PER_GROUP; + int const n_cols = (g == N_GROUPS - 1) ? LAST_COLS : COLS_PER_GROUP; + + // Intra-warp butterfly reduction for each column in this group +#pragma unroll + for (int n = 0; n < COLS_PER_GROUP; n++) + { + if (n < n_cols) + { +#pragma unroll + for (int off = WARP_SIZE / 2; off > 0; off >>= 1) + acc[g * COLS_PER_GROUP + n] += __shfl_xor_sync(FULL_WARP_MASK, acc[g * COLS_PER_GROUP + n], off); + } + } + if (g == 0 && do_sqr) + { +#pragma unroll + for (int off = WARP_SIZE / 2; off > 0; off >>= 1) + sqr += __shfl_xor_sync(FULL_WARP_MASK, sqr, off); + } + + // Lane 0 of each warp writes its partial sum to shared memory + if (lane == 0) + { +#pragma unroll + for (int n = 0; n < COLS_PER_GROUP; n++) + if (n < n_cols) + s_warp[warp_id][n] = acc[g * COLS_PER_GROUP + n]; + if (g == 0 && do_sqr) + s_warp[warp_id][SQRSUM_SLOT] = sqr; + } + __syncthreads(); + + // Final cross-warp sum: warp_id maps to output column + if (lane == 0 && warp_id < n_cols) + { + float val = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; w++) + val += s_warp[w][warp_id]; + Y[static_cast(row) * N + n_start + g * COLS_PER_GROUP + warp_id] = val; + } + if (g == 0 && do_sqr && lane == 0 && warp_id == 0) + { + float sq = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; w++) + sq += s_warp[w][SQRSUM_SLOT]; + R[row] = sq; + } + __syncthreads(); + } +} + +#define INST_FMA(NPB) \ + template __global__ void mhcGemmSqrsumFmaKernel( \ + __nv_bfloat16 const*, float const*, float*, float*, int, int, int); + +INST_FMA(1) +INST_FMA(2) +INST_FMA(3) +INST_FMA(4) +INST_FMA(6) +INST_FMA(8) +INST_FMA(12) +INST_FMA(24) +#undef INST_FMA + +// =================================================================== +// Kernel 4: post_mapping — one CTA per token, 256 threads +// +// Computes: out[token][j][h] = post[j] * x[h] + sum_k comb[k][j] * residual[k][h] +// +// post_mix: [B, HC_MULT] — per-token gating weights +// comb_mix: [B, HC_MULT, HC_MULT] — per-token combination matrix (from Sinkhorn) +// residual: [B, HC_MULT, hidden] — bf16 multi-head residual +// x: [B, hidden] — bf16 layer output +// out: [B, HC_MULT, hidden] — bf16 result +// +// Vectorized: 8 bf16 per thread per iteration via uint4 (LDG.128). +// =================================================================== + +__launch_bounds__(256) __global__ void mhcPostMappingKernel(__nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ x, float const* __restrict__ post_mix, float const* __restrict__ comb_mix, + __nv_bfloat16* __restrict__ out, int hidden_size) +{ + constexpr int HC_MULT = 4; + constexpr int BLOCK_SIZE = 256; + constexpr int BF16_VEC = 8; + + int const token = blockIdx.x; + int const tid = threadIdx.x; + + // Load post_mix and comb_mix into shared memory (only first 20 threads active) + __shared__ float s_post[HC_MULT]; + __shared__ float s_comb[HC_MULT][HC_MULT]; + + if (tid < HC_MULT) + s_post[tid] = post_mix[token * HC_MULT + tid]; + if (tid < HC_MULT * HC_MULT) + { + int const r = tid / HC_MULT; + int const c = tid % HC_MULT; + s_comb[r][c] = comb_mix[token * HC_MULT * HC_MULT + r * HC_MULT + c]; + } + __syncthreads(); + + // Cache post_mix and comb_mix in registers + float pm[HC_MULT]; + float comb[HC_MULT][HC_MULT]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + pm[j] = s_post[j]; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + comb[k][j] = s_comb[k][j]; + + long long const tok_res = static_cast(token) * HC_MULT * hidden_size; + long long const tok_x = static_cast(token) * hidden_size; + + for (int h = tid * BF16_VEC; h < hidden_size; h += BLOCK_SIZE * BF16_VEC) + { + // Load x[h:h+8] as bf16 → float + uint4 x_raw = *reinterpret_cast(&x[tok_x + h]); + __nv_bfloat162 const* xp = reinterpret_cast<__nv_bfloat162 const*>(&x_raw); + float xf[BF16_VEC]; +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + { + float2 f = __bfloat1622float2(xp[v]); + xf[2 * v + 0] = f.x; + xf[2 * v + 1] = f.y; + } + + // acc[j][v] = post[j] * x[v] (initialize with the x contribution) + float acc[HC_MULT][BF16_VEC]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + acc[j][v] = pm[j] * xf[v]; + + // acc[j][v] += sum_k comb[k][j] * residual[k][v] +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + uint4 r_raw = *reinterpret_cast(&residual[tok_res + k * hidden_size + h]); + __nv_bfloat162 const* rp = reinterpret_cast<__nv_bfloat162 const*>(&r_raw); + float rf[BF16_VEC]; +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + { + float2 f = __bfloat1622float2(rp[v]); + rf[2 * v + 0] = f.x; + rf[2 * v + 1] = f.y; + } + +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + acc[j][v] = fmaf(comb[k][j], rf[v], acc[j][v]); + } + + // Store acc → out[j][h:h+8] as bf16 +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + uint4 o_raw; + __nv_bfloat162* op = reinterpret_cast<__nv_bfloat162*>(&o_raw); +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + op[v] = __float22bfloat162_rn(make_float2(acc[j][2 * v], acc[j][2 * v + 1])); + *reinterpret_cast(&out[tok_res + j * hidden_size + h]) = o_raw; + } + } +} + +// =================================================================== +// Kernel 5: hc_head_apply — RMS norm → sigmoid → weighted sum +// +// Used for the HC head's final reduction: collapse `mult` input streams +// into a single hidden-size output per token. +// +// Per token: +// 1. Compute mix weights: sigmoid(GEMM_output * rstd * scale + base) + eps +// 2. Weighted sum: out[h] = sum_j mix[j] * x[j][h] +// +// mult is typically 4 (hc_mult) but passed as a runtime parameter for +// flexibility with different HC configurations. +// =================================================================== + +__launch_bounds__(256) __global__ + void mhcHcHeadApplyKernel(float const* __restrict__ mixes, float const* __restrict__ sqrsum, + __nv_bfloat16 const* __restrict__ x, __nv_bfloat16* __restrict__ out, float const* __restrict__ scale, + float const* __restrict__ base, int mult, int hidden_size, int K, float norm_eps, float eps) +{ + constexpr int BLOCK_SIZE = 256; + constexpr int BF16_VEC = 8; + constexpr int MAX_MULT = 8; + + int const token = blockIdx.x; + int const tid = threadIdx.x; + + // Phase 1: first `mult` threads compute sigmoid gating weights + __shared__ float s_pre[MAX_MULT]; + + if (tid < mult) + { + float sq = sqrsum[token]; + float rstd = rsqrtf(sq / static_cast(K) + norm_eps); + float m = mixes[token * mult + tid]; + float val = m * rstd * scale[0] + base[tid]; + float sig = 1.0f / (1.0f + expf(-val)); + s_pre[tid] = sig + eps; + } + __syncthreads(); + + // Phase 2: all threads compute weighted sum across `mult` input streams + __nv_bfloat16 const* xbase = x + static_cast(token) * mult * hidden_size; + __nv_bfloat16* obase = out + static_cast(token) * hidden_size; + + for (int h = tid * BF16_VEC; h < hidden_size; h += BLOCK_SIZE * BF16_VEC) + { + float acc[BF16_VEC] = {}; + + for (int j = 0; j < mult; j++) + { + float pm = s_pre[j]; + uint4 raw = *reinterpret_cast(&xbase[j * hidden_size + h]); + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raw); +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + { + float2 f = __bfloat1622float2(pairs[v]); + acc[2 * v + 0] += pm * f.x; + acc[2 * v + 1] += pm * f.y; + } + } + + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (int v = 0; v < BF16_VEC / 2; v++) + opairs[v] = __float22bfloat162_rn(make_float2(acc[2 * v], acc[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; + } +} + +// =================================================================== +// Offline-tuned config selection +// +// Thresholds determined by bench_mhc_gridsearch.cu on B200 (148 SMs). +// Fixed params: K=16384, hidden_size=4096, N∈{4,24}. +// Run bench_mhc_gridsearch and update thresholds if GPU changes. +// =================================================================== + +// FMA: select tileN given M and N. (B200, 148 SMs, K=16384) +// N=4 (HCHead): valid tileN ∈ {1, 2, 4} +// N=24 (pre_mapping): valid tileN ∈ {1, 2, 3, 4, 6, 8, 12, 24} +static int selectFmaTileN(int M, int N) +{ + if (N <= 4) + { + if (M <= 128) + return 1; + return 2; + } + if (M <= 32) + return 1; + return 8; +} + +// BigFuse: select BLOCK_SIZE given M. Thresholds match the production autotuner +// fallback; we use 128 threads for tiny-M waves (too few tokens to hide global +// scoreboard latency at 256) and 256 otherwise. 512 is not picked here because +// it only wins when M is very large, which path B/D do not route through this +// helper. +static int selectBigFuseBlockSize(int M) +{ + if (M <= 16) + return 128; + return 256; +} + +// =================================================================== +// Launch wrappers +// +// Two-level dispatch for BigFuse (NUM_SPLITS × BLOCK_SIZE) and a +// switch-dispatch for the FMA GEMM tile size. These are the public +// C++ entry points called by the PyTorch custom-op layer. +// =================================================================== + +template +static void mhcBigFuseDispatch(float const* y_acc, float const* r_acc, __nv_bfloat16 const* residual, + float const* hc_scale, float const* hc_base, float* post_mix, float* comb_mix, __nv_bfloat16* layer_input, int M, + int K, int hidden_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, + int sinkhorn_repeat, int block_size, cudaStream_t stream) +{ + dim3 grid(static_cast(M)); + +#define LAUNCH_BF(BS) \ + mhcBigFuseKernel<<>>(y_acc, r_acc, residual, hc_scale, hc_base, post_mix, \ + comb_mix, layer_input, M, K, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, \ + sinkhorn_repeat) + + if (block_size >= 512) + { + LAUNCH_BF(512); + } + else if (block_size >= 256) + { + LAUNCH_BF(256); + } + else + { + LAUNCH_BF(128); + } +#undef LAUNCH_BF +} + +void mhcBigFuseLaunch(float const* y_acc, float const* r_acc, __nv_bfloat16 const* residual, float const* hc_scale, + float const* hc_base, float* post_mix, float* comb_mix, __nv_bfloat16* layer_input, int M, int K, int hidden_size, + float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat, + int num_splits, int block_size, cudaStream_t stream) +{ + if (M <= 0) + return; + + TLLM_CHECK_WITH_INFO(num_splits == 1 || num_splits == 2 || num_splits == 4 || num_splits == 8 || num_splits == 16, + "mhcBigFuseLaunch: only num_splits ∈ {1,2,4,8,16} supported, got %d", num_splits); + + int const bs = (block_size > 0) ? block_size : selectBigFuseBlockSize(M); + +#define DISPATCH_BF(NS) \ + mhcBigFuseDispatch(y_acc, r_acc, residual, hc_scale, hc_base, post_mix, comb_mix, layer_input, M, K, \ + hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, bs, stream) + + switch (num_splits) + { + case 1: DISPATCH_BF(1); break; + case 2: DISPATCH_BF(2); break; + case 4: DISPATCH_BF(4); break; + case 8: DISPATCH_BF(8); break; + case 16: DISPATCH_BF(16); break; + } +#undef DISPATCH_BF +} + +void mhcGemmSqrsumFmaLaunch(__nv_bfloat16 const* x, float const* w_t, float* y, float* r, int M, int N, int K, + int tile_n, int tile_m, cudaStream_t stream) +{ + if (M <= 0) + return; + + (void) tile_m; // reserved for future multi-row kernel variant + + int const tileN = (tile_n > 0) ? tile_n : selectFmaTileN(M, N); + + TLLM_CHECK_WITH_INFO(N % tileN == 0, "mhcGemmSqrsumFmaLaunch: N=%d not divisible by tile_n=%d", N, tileN); + +#define LAUNCH_FMA(TN) mhcGemmSqrsumFmaKernel<<>>(x, w_t, y, r, M, N, K) + + switch (tileN) + { + case 1: LAUNCH_FMA(1); break; + case 2: LAUNCH_FMA(2); break; + case 3: LAUNCH_FMA(3); break; + case 4: LAUNCH_FMA(4); break; + case 6: LAUNCH_FMA(6); break; + case 8: LAUNCH_FMA(8); break; + case 12: LAUNCH_FMA(12); break; + default: LAUNCH_FMA(24); break; + } +#undef LAUNCH_FMA +} + +void mhcHcHeadApplyLaunch(float const* mixes, float const* sqrsum, __nv_bfloat16 const* x, __nv_bfloat16* out, + float const* scale, float const* base, int M, int mult, int hidden_size, int K, float norm_eps, float eps, + cudaStream_t stream) +{ + dim3 grid(static_cast(M)); + dim3 block(256); + + mhcHcHeadApplyKernel<<>>( + mixes, sqrsum, x, out, scale, base, mult, hidden_size, K, norm_eps, eps); +} + +void mhcPostMappingLaunch(__nv_bfloat16 const* residual, __nv_bfloat16 const* x, float const* post_mix, + float const* comb_mix, __nv_bfloat16* out, int B, int hidden_size, cudaStream_t stream) +{ + dim3 grid(static_cast(B)); + dim3 block(256); + + mhcPostMappingKernel<<>>(residual, x, post_mix, comb_mix, out, hidden_size); +} + +} // namespace kernels::mhc + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h new file mode 100644 index 000000000000..c64e494b50a1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h @@ -0,0 +1,139 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ +#pragma once + +#include +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc +{ + +void mhcBigFuseLaunch(float const* y_acc, float const* r_acc, __nv_bfloat16 const* residual, float const* hc_scale, + float const* hc_base, float* post_mix, float* comb_mix, __nv_bfloat16* layer_input, int M, int K, int hidden_size, + float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat, + int num_splits, int block_size, cudaStream_t stream); + +void mhcGemmSqrsumFmaLaunch(__nv_bfloat16 const* x, float const* w_t, float* y, float* r, int M, int N, int K, + int tile_n, int tile_m, cudaStream_t stream); + +void mhcHcHeadApplyLaunch(float const* mixes, float const* sqrsum, __nv_bfloat16 const* x, __nv_bfloat16* out, + float const* scale, float const* base, int M, int mult, int hidden_size, int K, float norm_eps, float eps, + cudaStream_t stream); + +void mhcPostMappingLaunch(__nv_bfloat16 const* residual, __nv_bfloat16 const* x, float const* post_mix, + float const* comb_mix, __nv_bfloat16* out, int B, int hidden_size, cudaStream_t stream); + +// Single-launch fused hyper-connection boundary op (SM100 only). +// +// Produces (residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur) in two +// kernel launches: (1) tcgen05 TF32 GEMM fused with post-mapping and sqr-sum +// that emits D, sqr_sum, and residual_cur, and (2) big-fuse postlogue that +// consumes those to emit post_mix_cur, comb_mix_cur, layer_input_cur. +// +// Workspace tensors are caller-allocated and zeroed by this launcher: +// y_acc_workspace: fp32 [M, 24] +// r_acc_workspace: fp32 [M] +// +// Shape constraints (B200 / sm_100a): hidden_size == 4096, hc_mult == 4. +// Passing num_k_splits == 0 or bigfuse_block_size == 0 falls back to the +// internal heuristics. +void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev, + float const* comb_mix_prev, float const* w_t, float const* hc_scale, float const* hc_base, + __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, __nv_bfloat16* layer_input_cur, + float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int num_k_splits, + int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, + int sinkhorn_repeat, cudaStream_t stream); + +// FMA-path fused hyper-connection boundary launcher. +// +// Composes fused_pmap_gemm_fma_ksplit (writes Yp[ks, M, 24], Rp[ks, M], and +// residual_cur[M, HC_MULT, HIDDEN]) with mhcBigFuseKernel which +// reduces across the split axis internally. Used as the small-M backend in +// the mhcFusedHc autotuner. +// +// Supported tactics (hit by the autotuner); see pickFhcFma for the exact +// (tile_n, num_k_splits) table: +// tile_n ∈ {1, 2, 3, 4, 6, 8, 12, 24} +// num_k_splits ∈ {1, 2, 4, 8} (ks=1 ⇔ seq FMA, no HIDDEN split). Larger ks +// requires smaller tile_n: {1,2}→{1,2,4,8}, {3}→{1,2,4}, {4}→{1,2}, +// {6,8,12,24}→{1}. +// bigfuse_block_size ∈ {128, 256, 512} +// +// Workspace shapes (caller-allocated): +// y_acc_workspace: fp32 [num_k_splits, M, 24] +// r_acc_workspace: fp32 [num_k_splits, M] +// No pre-zeroing required (ksplit kernel writes = not +=). +void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev, + float const* comb_mix_prev, float const* w_t, float const* hc_scale, float const* hc_base, + __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, __nv_bfloat16* layer_input_cur, + float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int tile_n, int num_k_splits, + int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, + int sinkhorn_repeat, cudaStream_t stream); + +// Single-kernel all-in-one fused hyper-connection boundary launcher (TF32 tcgen05 path). +// +// Wraps fused_allinone_tf32_pmap_gemm_atomic_impl: pmap + TF32 GEMM + bigfuse +// are all fused into ONE kernel launch. Phase 3 uses an atomic done-counter +// to elect the last-home CTA per m-block, which then executes the bigfuse +// epilogue inline (no second kernel, no launch-latency stall). Preferred for +// mid/large M where the extra kernel launch overhead of the 2-kernel path is +// visible (see design_doc_fused_hc.md). +// +// Workspace tensors are caller-allocated and zeroed by this launcher: +// y_acc_workspace: fp32 [M, 24] +// r_acc_workspace: fp32 [M] +// done_counter_workspace: int32 [ceil(M / 64)] +// +// Shape constraints (B200 / sm_100a): hidden_size == 4096, hc_mult == 4. +// Passing num_k_splits == 0 falls back to internal heuristics. +void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, + float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale, + float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, + __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int* done_counter_workspace, int M, + int hidden_size, int hc_mult, int num_k_splits, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, + float hc_post_mult_value, int sinkhorn_repeat, cudaStream_t stream); + +// Single-kernel all-in-one fused hyper-connection boundary launcher (FMA path). +// +// Wraps fused_pmap_gemm_fma_allinone: pmap + fp32 FMA GEMM + +// bigfuse are all fused into ONE kernel launch. Uses the same atomic +// last-home election pattern as the TF32 all-in-one path. Preferred for +// small-M (M <= 32) where the TF32 tcgen05 path cannot saturate the MMA pipe. +// +// Supported tactics; see pickFhcFmaAllInOne for the exact table: +// tile_n ∈ {1, 2, 3, 4, 6, 8, 12, 24} +// num_k_splits ∈ {1, 2}. ks=2 is only valid with tile_n ∈ {1, 2}; all +// other tile_n require ks=1. +// tile_m ∈ {1, 2, 4} (tokens per CTA) +// +// Workspace shapes (caller-allocated, zeroed by this launcher): +// y_acc_workspace: fp32 [M, 24] +// r_acc_workspace: fp32 [M] +// done_counter_workspace: int32 [ceil(M / tile_m)] +void mhcFusedHcFmaAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, + float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale, + float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, + __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int* done_counter_workspace, int M, + int hidden_size, int hc_mult, int tile_n, int num_k_splits, int tile_m, float rms_eps, float hc_pre_eps, + float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat, cudaStream_t stream); + +} // namespace kernels::mhc + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh new file mode 100644 index 000000000000..e3b01fab6b7c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh @@ -0,0 +1,898 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// CUDA-core FMA fused post_mapping + pre_mapping kernels for the mhc +// fused_hc pipeline. Provides the small-M variants: +// - fused_pmap_gemm_fma_ksplit (Path E, 2-kernel: this + mhcBigFuseKernel) +// - fused_pmap_gemm_fma_allinone (Path F, 1-kernel: pmap + GEMM + bigFuse) +#pragma once + +#include +#include + +namespace fused_fma_kernels +{ + +// =================================================================== +// Fused K-split pmap+GEMM FMA: new_r is NEVER materialized to HBM. Each CTA +// scans one HIDDEN-slice inline: for each h in its slice it loads residual[j,h] +// for ALL HC j-indices plus x[h], computes FULL new_r[j,h] = pm[j]*x[h] + Σ_k +// comb[k,j]*r[k,h] in registers, then multiplies by W_T and accumulates partial +// acc+sqr. Writes Yp[k_split, token, n] and Rp[k_split, token]. A reducer +// kernel sums over splits. Eliminates the 32-KB/token new_r round-trip at tiny M. +// +// CORRECTNESS: We split the HIDDEN dimension (not HC). For each fixed h, +// new_r[*, h] depends only on x[h] and residual[*, h]; the HC summation is +// fully local to h, so a HIDDEN-split is algebraically exact. The final GEMM +// acc[n] = Σ_{j,h} new_r[j,h] * W_T[n, j*HIDDEN+h] is then partitioned into +// disjoint h-ranges across splits and reduced by mhcBigFuseKernel. +// +// Grid: (M, N/N_PER_BLOCK, NUM_K_SPLITS). Block: 256 threads. +// NUM_K_SPLITS must evenly divide HIDDEN (typically ∈ {2, 4, 8}). +// BF16_VEC_OVERRIDE lets the caller force a smaller per-thread vector (e.g. 4) +// so that BLOCK_SIZE*BF16_VEC == HIDDEN_PER_SPLIT and every thread is active. +// Pass 0 to auto-pick: 8 for ks=2, 4 for ks=4, 2 for ks=8. +// =================================================================== +// Optional WRITE_RESIDUAL path: when true, the CTA with (blockIdx.y == 0) emits +// the post-mapped residual (bf16 [M, HC_MULT, hidden]) into residual_out, using +// the same disjoint h-slice it already owns. When false the residual_out arg +// may be a nullptr; the parameter stays in the signature for simplicity. +template +__launch_bounds__(256) __global__ void fused_pmap_gemm_fma_ksplit(__nv_bfloat16 const* __restrict__ residual_in, + __nv_bfloat16 const* __restrict__ x_in, float const* __restrict__ post_mix, float const* __restrict__ comb_mix, + float const* __restrict__ W_T, float* __restrict__ Yp, float* __restrict__ Rp, int hidden_size, int N, int K, + __nv_bfloat16* __restrict__ residual_out = nullptr) +{ + static_assert(NUM_K_SPLITS == 1 || NUM_K_SPLITS == 2 || NUM_K_SPLITS == 4 || NUM_K_SPLITS == 8, + "ksplit ∈ {1, 2, 4, 8} supported"); + constexpr int HC_MULT = 4; + constexpr int BLOCK_SIZE = 256; + constexpr int WARP_SIZE = 32; + constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; + // Default BF16_VEC by ks so that BLOCK_SIZE*BF16_VEC = HIDDEN/ks at HIDDEN=4096. + // ks=1 takes the full HIDDEN in one slice; reuses the BF16_VEC=8 (uint4) load path. + constexpr int BF16_VEC = (BF16_VEC_OVERRIDE != 0) ? BF16_VEC_OVERRIDE + : ((NUM_K_SPLITS == 1 || NUM_K_SPLITS == 2) ? 8 + : (NUM_K_SPLITS == 4) ? 4 + : 2); + static_assert(BF16_VEC == 2 || BF16_VEC == 4 || BF16_VEC == 8, "BF16_VEC must be 2, 4, or 8"); + + int const token = blockIdx.x; + int const n_start = blockIdx.y * N_PER_BLOCK; + int const k_split = blockIdx.z; + int const tid = threadIdx.x; + int const warp_id = tid / WARP_SIZE; + int const lane = tid % WARP_SIZE; + bool const do_sqr = (n_start == 0); + + int const hidden_per_split = hidden_size / NUM_K_SPLITS; + int const h_start = k_split * hidden_per_split; + int const h_end = h_start + hidden_per_split; + + __shared__ float s_post[HC_MULT]; + __shared__ float s_comb[HC_MULT][HC_MULT]; + + if (tid < HC_MULT) + s_post[tid] = post_mix[token * HC_MULT + tid]; + if (tid < HC_MULT * HC_MULT) + { + int const r = tid / HC_MULT; + int const c = tid % HC_MULT; + s_comb[r][c] = comb_mix[token * HC_MULT * HC_MULT + r * HC_MULT + c]; + } + __syncthreads(); + + float pm[HC_MULT]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + pm[j] = s_post[j]; + + float acc[N_PER_BLOCK]; +#pragma unroll + for (int n = 0; n < N_PER_BLOCK; n++) + acc[n] = 0.0f; + float sqr = 0.0f; + + long long const tok_res = static_cast(token) * HC_MULT * hidden_size; + long long const tok_x = static_cast(token) * hidden_size; + + // Iterate h over this CTA's HIDDEN slice. Each thread owns BF16_VEC + // contiguous positions per step; stride BLOCK_SIZE*BF16_VEC. + for (int h = h_start + tid * BF16_VEC; h < h_end; h += BLOCK_SIZE * BF16_VEC) + { + // Load x[h..h+BF16_VEC). Vector width depends on BF16_VEC. + float xf[BF16_VEC]; + if constexpr (BF16_VEC == 8) + { + uint4 x_raw = *reinterpret_cast(&x_in[tok_x + h]); + __nv_bfloat162 const* xp = reinterpret_cast<__nv_bfloat162 const*>(&x_raw); +#pragma unroll + for (int v = 0; v < 4; v++) + { + float2 f = __bfloat1622float2(xp[v]); + xf[2 * v + 0] = f.x; + xf[2 * v + 1] = f.y; + } + } + else if constexpr (BF16_VEC == 4) + { + uint2 x_raw = *reinterpret_cast(&x_in[tok_x + h]); + __nv_bfloat162 const* xp = reinterpret_cast<__nv_bfloat162 const*>(&x_raw); +#pragma unroll + for (int v = 0; v < 2; v++) + { + float2 f = __bfloat1622float2(xp[v]); + xf[2 * v + 0] = f.x; + xf[2 * v + 1] = f.y; + } + } + else + { // BF16_VEC == 2 + unsigned x_raw = *reinterpret_cast(&x_in[tok_x + h]); + __nv_bfloat162 xp = *reinterpret_cast<__nv_bfloat162 const*>(&x_raw); + float2 f = __bfloat1622float2(xp); + xf[0] = f.x; + xf[1] = f.y; + } + + // Seed new_r[j, v] = pm[j] * x[v] for all j + float new_r[HC_MULT][BF16_VEC]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + new_r[j][v] = pm[j] * xf[v]; + +// Add Σ_k comb[k,j] * r[k, v] — FULL HC sum (all k ∈ [0, HC_MULT)) +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + float rf_k[BF16_VEC]; + __nv_bfloat16 const* r_ptr = &residual_in[tok_res + k * hidden_size + h]; + if constexpr (BF16_VEC == 8) + { + uint4 r_raw = *reinterpret_cast(r_ptr); + __nv_bfloat162 const* rp = reinterpret_cast<__nv_bfloat162 const*>(&r_raw); +#pragma unroll + for (int v = 0; v < 4; v++) + { + float2 f = __bfloat1622float2(rp[v]); + rf_k[2 * v + 0] = f.x; + rf_k[2 * v + 1] = f.y; + } + } + else if constexpr (BF16_VEC == 4) + { + uint2 r_raw = *reinterpret_cast(r_ptr); + __nv_bfloat162 const* rp = reinterpret_cast<__nv_bfloat162 const*>(&r_raw); +#pragma unroll + for (int v = 0; v < 2; v++) + { + float2 f = __bfloat1622float2(rp[v]); + rf_k[2 * v + 0] = f.x; + rf_k[2 * v + 1] = f.y; + } + } + else + { // BF16_VEC == 2 + unsigned r_raw = *reinterpret_cast(r_ptr); + __nv_bfloat162 rp = *reinterpret_cast<__nv_bfloat162 const*>(&r_raw); + float2 f = __bfloat1622float2(rp); + rf_k[0] = f.x; + rf_k[1] = f.y; + } +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + new_r[j][v] = fmaf(s_comb[k][j], rf_k[v], new_r[j][v]); + } + + if (do_sqr) + { +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + sqr = fmaf(new_r[j][v], new_r[j][v], sqr); + } + + // Optional: emit the post-mapped residual to HBM (bf16) for downstream + // bigfuse. Guard on WRITE_RESIDUAL (compile-time) AND do_sqr so only + // one n-tile per (token, k_split) writes — other n-tiles compute the + // same new_r but we don't want duplicate writes. + if constexpr (WRITE_RESIDUAL) + { + if (do_sqr) + { +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + __nv_bfloat16* o_ptr = residual_out + tok_res + j * hidden_size + h; + if constexpr (BF16_VEC == 8) + { + uint4 packed; + __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&packed); +#pragma unroll + for (int v = 0; v < 4; v++) + pairs[v] = __float22bfloat162_rn(make_float2(new_r[j][2 * v], new_r[j][2 * v + 1])); + *reinterpret_cast(o_ptr) = packed; + } + else if constexpr (BF16_VEC == 4) + { + uint2 packed; + __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&packed); +#pragma unroll + for (int v = 0; v < 2; v++) + pairs[v] = __float22bfloat162_rn(make_float2(new_r[j][2 * v], new_r[j][2 * v + 1])); + *reinterpret_cast(o_ptr) = packed; + } + else + { // BF16_VEC == 2 + __nv_bfloat162 packed = __float22bfloat162_rn(make_float2(new_r[j][0], new_r[j][1])); + *reinterpret_cast(o_ptr) = *reinterpret_cast(&packed); + } + } + } + } + +// GEMM partial: for each n, contract over ALL HC j and this h-slice. +// W_T layout: W_T[n, K=HC_MULT*hidden_size], inner dim packed as (j, h). +#pragma unroll + for (int n = 0; n < N_PER_BLOCK; n++) + { +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + long long const w_base = static_cast(n_start + n) * K + j * hidden_size + h; + float w[BF16_VEC]; + if constexpr (BF16_VEC == 8) + { + float w0, w1, w2, w3, w4, w5, w6, w7; + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w0), "=f"(w1), "=f"(w2), "=f"(w3) + : "l"(W_T + w_base)); + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w4), "=f"(w5), "=f"(w6), "=f"(w7) + : "l"(W_T + w_base + 4)); + w[0] = w0; + w[1] = w1; + w[2] = w2; + w[3] = w3; + w[4] = w4; + w[5] = w5; + w[6] = w6; + w[7] = w7; + } + else if constexpr (BF16_VEC == 4) + { + float w0, w1, w2, w3; + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w0), "=f"(w1), "=f"(w2), "=f"(w3) + : "l"(W_T + w_base)); + w[0] = w0; + w[1] = w1; + w[2] = w2; + w[3] = w3; + } + else + { // BF16_VEC == 2 + float w0, w1; + asm volatile("ld.global.L1::evict_last.v2.f32 {%0, %1}, [%2];" + : "=f"(w0), "=f"(w1) + : "l"(W_T + w_base)); + w[0] = w0; + w[1] = w1; + } +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + acc[n] = fmaf(new_r[j][v], w[v], acc[n]); + } + } + } + + constexpr int COLS_PER_GROUP = NUM_WARPS; + constexpr int SQRSUM_SLOT = COLS_PER_GROUP; + constexpr int N_GROUPS = (N_PER_BLOCK + COLS_PER_GROUP - 1) / COLS_PER_GROUP; + constexpr unsigned FULL_WARP_MASK = 0xffffffff; + + __shared__ float s_warp[NUM_WARPS][COLS_PER_GROUP + 1]; + +#pragma unroll + for (int g = 0; g < N_GROUPS; g++) + { + constexpr int LAST_COLS = N_PER_BLOCK - (N_GROUPS - 1) * COLS_PER_GROUP; + int const n_cols = (g == N_GROUPS - 1) ? LAST_COLS : COLS_PER_GROUP; +#pragma unroll + for (int n = 0; n < COLS_PER_GROUP; n++) + { + if (n < n_cols) + { +#pragma unroll + for (int off = WARP_SIZE / 2; off > 0; off >>= 1) + acc[g * COLS_PER_GROUP + n] += __shfl_xor_sync(FULL_WARP_MASK, acc[g * COLS_PER_GROUP + n], off); + } + } + if (g == 0 && do_sqr) + { +#pragma unroll + for (int off = WARP_SIZE / 2; off > 0; off >>= 1) + sqr += __shfl_xor_sync(FULL_WARP_MASK, sqr, off); + } + if (lane == 0) + { +#pragma unroll + for (int n = 0; n < COLS_PER_GROUP; n++) + if (n < n_cols) + s_warp[warp_id][n] = acc[g * COLS_PER_GROUP + n]; + if (g == 0 && do_sqr) + s_warp[warp_id][SQRSUM_SLOT] = sqr; + } + __syncthreads(); + if (lane == 0 && warp_id < n_cols) + { + float val = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; w++) + val += s_warp[w][warp_id]; + // Yp layout: [NUM_K_SPLITS, M, N] (M = gridDim.x) + int const n_global = n_start + g * COLS_PER_GROUP + warp_id; + Yp[(static_cast(k_split) * static_cast(gridDim.x) + token) * N + n_global] = val; + } + if (g == 0 && do_sqr && lane == 0 && warp_id == 0) + { + float sq = 0.0f; +#pragma unroll + for (int w = 0; w < NUM_WARPS; w++) + sq += s_warp[w][SQRSUM_SLOT]; + // Rp layout: [NUM_K_SPLITS, M] + Rp[static_cast(k_split) * static_cast(gridDim.x) + token] = sq; + } + __syncthreads(); + } +} + +// =================================================================== +// fused_pmap_gemm_fma_allinone +// ------------------------------------------------------------------- +// 1-kernel FMA path: pmap + GEMM + bigFuse, all in a single launch. +// Parallelism: TM tokens per CTA × TN N-columns × KS K-slices. +// Grid: (ceil(M / TM), FULL_N / TN, KS) +// Block: 256 threads. +// +// Uses atomicAdd into y_acc[M, FULL_N] and r_acc[M] for cross-CTA reduction +// (KS > 1 and/or TN < FULL_N). Uses per-token-batch atomic counter to elect +// the last-home CTA, which __threadfences, reloads the now-complete y_acc / +// r_acc / residual_cur rows, and runs bigFuse inline to produce +// post_mix_out, comb_mix_out, layer_input_out. +// +// Caller MUST zero y_acc[M, FULL_N], r_acc[M], done_counter[ceil(M / TM)] +// before launch. FULL_N must equal HC_MULT * (2 + HC_MULT) = 24. +// =================================================================== +template +__launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat16 const* __restrict__ residual_in, + __nv_bfloat16 const* __restrict__ x_in, float const* __restrict__ post_mix_prev, + float const* __restrict__ comb_mix_prev, float const* __restrict__ W_T, float const* __restrict__ hc_scale, + float const* __restrict__ hc_base, __nv_bfloat16* __restrict__ residual_out, float* __restrict__ post_mix_out, + float* __restrict__ comb_mix_out, __nv_bfloat16* __restrict__ layer_input_out, float* __restrict__ y_acc, + float* __restrict__ r_acc, int* __restrict__ done_counter, int M, int K, int hidden_size, float rms_eps, + float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat) +{ + constexpr int HC_MULT = 4; + constexpr int HC_MULT2 = HC_MULT * HC_MULT; + constexpr int HC_MULT3 = HC_MULT * (2 + HC_MULT); + static_assert(TM >= 1, "TM>=1"); + static_assert(FULL_N % TN == 0, "TN must divide FULL_N"); + static_assert(HC_MULT3 == FULL_N, "FULL_N must be HC_MULT*(2+HC_MULT)=24"); + static_assert(KS == 1 || KS == 2 || KS == 4 || KS == 8, "KS in {1,2,4,8}"); + constexpr int BLOCK_SIZE = 256; + constexpr int WARP_SIZE = 32; + constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE; + constexpr int BF16_VEC = (BF16_VEC_OVERRIDE != 0) ? BF16_VEC_OVERRIDE + : ((KS == 1 || KS == 2) ? 8 + : (KS == 4) ? 4 + : 2); + static_assert(BF16_VEC == 2 || BF16_VEC == 4 || BF16_VEC == 8, "BF16_VEC in {2,4,8}"); + constexpr int CTAS_PER_BATCH = (FULL_N / TN) * KS; + + int const base_tok = blockIdx.x * TM; + int const n_start = blockIdx.y * TN; + int const k_split = blockIdx.z; + int const tid = threadIdx.x; + int const warp_id = tid / WARP_SIZE; + int const lane = tid % WARP_SIZE; + bool const is_n0 = (n_start == 0); + + if (base_tok >= M) + return; + + int const hidden_per_split = hidden_size / KS; + int const h_lo = k_split * hidden_per_split; + int const h_hi = h_lo + hidden_per_split; + + bool valid[TM]; +#pragma unroll + for (int t = 0; t < TM; t++) + valid[t] = (base_tok + t) < M; + + __shared__ float s_pm[TM][HC_MULT]; + __shared__ float s_cm[TM][HC_MULT][HC_MULT]; + + { + constexpr int N_PM = TM * HC_MULT; + for (int i = tid; i < N_PM; i += BLOCK_SIZE) + { + int const t = i / HC_MULT; + int const j = i % HC_MULT; + float v = 0.0f; + if (valid[t]) + v = post_mix_prev[(base_tok + t) * HC_MULT + j]; + s_pm[t][j] = v; + } + } + { + constexpr int N_CM = TM * HC_MULT * HC_MULT; + for (int i = tid; i < N_CM; i += BLOCK_SIZE) + { + int const t = i / HC_MULT2; + int const r = (i / HC_MULT) % HC_MULT; + int const c = i % HC_MULT; + float v = 0.0f; + if (valid[t]) + v = comb_mix_prev[(base_tok + t) * HC_MULT2 + r * HC_MULT + c]; + s_cm[t][r][c] = v; + } + } + __syncthreads(); + + float acc[TM][TN]; + float sqr[TM]; +#pragma unroll + for (int t = 0; t < TM; t++) + { + sqr[t] = 0.0f; +#pragma unroll + for (int n = 0; n < TN; n++) + acc[t][n] = 0.0f; + } + + // Phase 1: pmap + GEMM + residual_out write + sqrsum accumulate. + // Iterate h over this CTA's HIDDEN slice; each thread owns BF16_VEC. + for (int h = h_lo + tid * BF16_VEC; h < h_hi; h += BLOCK_SIZE * BF16_VEC) + { +#pragma unroll + for (int t = 0; t < TM; t++) + { + if (!valid[t]) + continue; + long long const tok_x = static_cast(base_tok + t) * hidden_size; + long long const tok_r = static_cast(base_tok + t) * HC_MULT * hidden_size; + + float xf[BF16_VEC]; + if constexpr (BF16_VEC == 8) + { + uint4 x_raw = *reinterpret_cast(&x_in[tok_x + h]); + __nv_bfloat162 const* xp = reinterpret_cast<__nv_bfloat162 const*>(&x_raw); +#pragma unroll + for (int v = 0; v < 4; v++) + { + float2 f = __bfloat1622float2(xp[v]); + xf[2 * v + 0] = f.x; + xf[2 * v + 1] = f.y; + } + } + else if constexpr (BF16_VEC == 4) + { + uint2 x_raw = *reinterpret_cast(&x_in[tok_x + h]); + __nv_bfloat162 const* xp = reinterpret_cast<__nv_bfloat162 const*>(&x_raw); +#pragma unroll + for (int v = 0; v < 2; v++) + { + float2 f = __bfloat1622float2(xp[v]); + xf[2 * v + 0] = f.x; + xf[2 * v + 1] = f.y; + } + } + else + { + unsigned x_raw = *reinterpret_cast(&x_in[tok_x + h]); + __nv_bfloat162 xp = *reinterpret_cast<__nv_bfloat162 const*>(&x_raw); + float2 f = __bfloat1622float2(xp); + xf[0] = f.x; + xf[1] = f.y; + } + + float rf[HC_MULT][BF16_VEC]; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + __nv_bfloat16 const* r_ptr = &residual_in[tok_r + k * hidden_size + h]; + if constexpr (BF16_VEC == 8) + { + uint4 r_raw = *reinterpret_cast(r_ptr); + __nv_bfloat162 const* rp = reinterpret_cast<__nv_bfloat162 const*>(&r_raw); +#pragma unroll + for (int v = 0; v < 4; v++) + { + float2 f = __bfloat1622float2(rp[v]); + rf[k][2 * v + 0] = f.x; + rf[k][2 * v + 1] = f.y; + } + } + else if constexpr (BF16_VEC == 4) + { + uint2 r_raw = *reinterpret_cast(r_ptr); + __nv_bfloat162 const* rp = reinterpret_cast<__nv_bfloat162 const*>(&r_raw); +#pragma unroll + for (int v = 0; v < 2; v++) + { + float2 f = __bfloat1622float2(rp[v]); + rf[k][2 * v + 0] = f.x; + rf[k][2 * v + 1] = f.y; + } + } + else + { + unsigned r_raw = *reinterpret_cast(r_ptr); + __nv_bfloat162 rp = *reinterpret_cast<__nv_bfloat162 const*>(&r_raw); + float2 f = __bfloat1622float2(rp); + rf[k][0] = f.x; + rf[k][1] = f.y; + } + } + + float pm[HC_MULT]; + float cm[HC_MULT][HC_MULT]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + pm[j] = s_pm[t][j]; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + cm[k][j] = s_cm[t][k][j]; + + float new_r[HC_MULT][BF16_VEC]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + new_r[j][v] = pm[j] * xf[v]; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + new_r[j][v] = fmaf(cm[k][j], rf[k][v], new_r[j][v]); + + if (is_n0) + { +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + __nv_bfloat16* o_ptr = residual_out + tok_r + j * hidden_size + h; + if constexpr (BF16_VEC == 8) + { + uint4 packed; + __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&packed); +#pragma unroll + for (int v = 0; v < 4; v++) + pairs[v] = __float22bfloat162_rn(make_float2(new_r[j][2 * v], new_r[j][2 * v + 1])); + *reinterpret_cast(o_ptr) = packed; + } + else if constexpr (BF16_VEC == 4) + { + uint2 packed; + __nv_bfloat162* pairs = reinterpret_cast<__nv_bfloat162*>(&packed); +#pragma unroll + for (int v = 0; v < 2; v++) + pairs[v] = __float22bfloat162_rn(make_float2(new_r[j][2 * v], new_r[j][2 * v + 1])); + *reinterpret_cast(o_ptr) = packed; + } + else + { + __nv_bfloat162 packed = __float22bfloat162_rn(make_float2(new_r[j][0], new_r[j][1])); + *reinterpret_cast(o_ptr) = *reinterpret_cast(&packed); + } + } + +#pragma unroll + for (int j = 0; j < HC_MULT; j++) +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + sqr[t] = fmaf(new_r[j][v], new_r[j][v], sqr[t]); + } + +#pragma unroll + for (int n = 0; n < TN; n++) + { +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + long long const w_base = static_cast(n_start + n) * K + j * hidden_size + h; + float w[BF16_VEC]; + if constexpr (BF16_VEC == 8) + { + float w0, w1, w2, w3, w4, w5, w6, w7; + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w0), "=f"(w1), "=f"(w2), "=f"(w3) + : "l"(W_T + w_base)); + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w4), "=f"(w5), "=f"(w6), "=f"(w7) + : "l"(W_T + w_base + 4)); + w[0] = w0; + w[1] = w1; + w[2] = w2; + w[3] = w3; + w[4] = w4; + w[5] = w5; + w[6] = w6; + w[7] = w7; + } + else if constexpr (BF16_VEC == 4) + { + float w0, w1, w2, w3; + asm volatile("ld.global.L1::evict_last.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(w0), "=f"(w1), "=f"(w2), "=f"(w3) + : "l"(W_T + w_base)); + w[0] = w0; + w[1] = w1; + w[2] = w2; + w[3] = w3; + } + else + { + float w0, w1; + asm volatile("ld.global.L1::evict_last.v2.f32 {%0, %1}, [%2];" + : "=f"(w0), "=f"(w1) + : "l"(W_T + w_base)); + w[0] = w0; + w[1] = w1; + } +#pragma unroll + for (int v = 0; v < BF16_VEC; v++) + acc[t][n] = fmaf(new_r[j][v], w[v], acc[t][n]); + } + } + } + } + + // Phase 2: warp-reduce + SMEM cross-warp + atomicAdd into y_acc / r_acc. + constexpr int COLS_PER_GROUP = NUM_WARPS; + constexpr int SQRSUM_SLOT = COLS_PER_GROUP; + constexpr int N_GROUPS = (TN + COLS_PER_GROUP - 1) / COLS_PER_GROUP; + constexpr unsigned FULL_WARP_MASK = 0xffffffff; + + __shared__ float s_warp[NUM_WARPS][COLS_PER_GROUP + 1]; + +#pragma unroll + for (int t = 0; t < TM; t++) + { + if (!valid[t]) + continue; +#pragma unroll + for (int g = 0; g < N_GROUPS; g++) + { + constexpr int LAST_COLS = TN - (N_GROUPS - 1) * COLS_PER_GROUP; + int const n_cols = (g == N_GROUPS - 1) ? LAST_COLS : COLS_PER_GROUP; +#pragma unroll + for (int n = 0; n < COLS_PER_GROUP; n++) + { + if (n < n_cols) + { +#pragma unroll + for (int off = WARP_SIZE / 2; off > 0; off >>= 1) + acc[t][g * COLS_PER_GROUP + n] + += __shfl_xor_sync(FULL_WARP_MASK, acc[t][g * COLS_PER_GROUP + n], off); + } + } + if (g == 0 && is_n0) + { +#pragma unroll + for (int off = WARP_SIZE / 2; off > 0; off >>= 1) + sqr[t] += __shfl_xor_sync(FULL_WARP_MASK, sqr[t], off); + } + if (lane == 0) + { +#pragma unroll + for (int n = 0; n < COLS_PER_GROUP; n++) + if (n < n_cols) + s_warp[warp_id][n] = acc[t][g * COLS_PER_GROUP + n]; + if (g == 0 && is_n0) + s_warp[warp_id][SQRSUM_SLOT] = sqr[t]; + } + __syncthreads(); + if (lane == 0 && warp_id < n_cols) + { + float val = 0.0f; +#pragma unroll + for (int ww = 0; ww < NUM_WARPS; ww++) + val += s_warp[ww][warp_id]; + long long y_idx + = static_cast(base_tok + t) * FULL_N + n_start + g * COLS_PER_GROUP + warp_id; + if constexpr (KS == 1) + { + y_acc[y_idx] = val; + } + else + { + atomicAdd(y_acc + y_idx, val); + } + } + if (g == 0 && is_n0 && lane == 0 && warp_id == 0) + { + float sq = 0.0f; +#pragma unroll + for (int ww = 0; ww < NUM_WARPS; ww++) + sq += s_warp[ww][SQRSUM_SLOT]; + if constexpr (KS == 1) + { + r_acc[base_tok + t] = sq; + } + else + { + atomicAdd(r_acc + (base_tok + t), sq); + } + } + __syncthreads(); + } + } + + // Phase 3: elect last-home CTA per (ceil(M/TM)) token-batch via counter. + __shared__ int s_is_last; + if constexpr (CTAS_PER_BATCH == 1) + { + s_is_last = 1; + } + else + { + __threadfence(); + __syncthreads(); + if (tid == 0) + { + int const prev = atomicAdd(&done_counter[blockIdx.x], 1); + s_is_last = (prev == CTAS_PER_BATCH - 1) ? 1 : 0; + } + __syncthreads(); + } + if (!s_is_last) + return; + + // Phase 4: inline bigFuse for the TM tokens in this batch. + // Layout: FULL_N = HC_MULT*(2+HC_MULT) = 24 + // y_local[0..HC_MULT) → s_pre_mix (sigmoid gate) + // y_local[HC_MULT..2*HC_MULT) → post_mix_out (sigmoid*hc_post_mult) + // y_local[2*HC_MULT..FULL_N) → comb_mix_out (Sinkhorn) + __shared__ float s_pre_mix[TM][HC_MULT]; + +#pragma unroll + for (int t = 0; t < TM; t++) + { + if (!valid[t]) + { + __syncthreads(); + continue; + } + int const tok = base_tok + t; + float cm_vals[HC_MULT]; + + if (warp_id == 0 && lane < HC_MULT) + { + float const r_val = r_acc[tok]; + float y_local[HC_MULT3]; + float const* y_row = y_acc + static_cast(tok) * FULL_N; +#pragma unroll + for (int c = 0; c < HC_MULT3; c++) + y_local[c] = y_row[c]; + + float const rstd = rsqrtf(r_val / static_cast(K) + rms_eps); + float const s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; + + float v = y_local[lane] * rstd * s0 + hc_base[lane]; + s_pre_mix[t][lane] = 1.0f / (1.0f + expf(-v)) + hc_pre_eps; + + v = y_local[HC_MULT + lane] * rstd * s1 + hc_base[HC_MULT + lane]; + post_mix_out[tok * HC_MULT + lane] = 1.0f / (1.0f + expf(-v)) * hc_post_mult_value; + +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm_vals[k] + = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; + + constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; + float const rowMax + = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm_vals[k] = expf(cm_vals[k] - rowMax); + float rs = cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3]; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm_vals[k] = cm_vals[k] / rs + hc_sinkhorn_eps; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] /= (cs + hc_sinkhorn_eps); + } + for (int it = 1; it < sinkhorn_repeat; it++) + { + rs = cm_vals[0] + cm_vals[1] + cm_vals[2] + cm_vals[3] + hc_sinkhorn_eps; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm_vals[k] /= rs; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + { + float cs = cm_vals[k]; + cs += __shfl_xor_sync(LANE_MASK, cs, 1); + cs += __shfl_xor_sync(LANE_MASK, cs, 2); + cm_vals[k] /= (cs + hc_sinkhorn_eps); + } + } + float* cm_out_ptr = comb_mix_out + tok * HC_MULT2; +#pragma unroll + for (int k = 0; k < HC_MULT; k++) + cm_out_ptr[lane * HC_MULT + k] = cm_vals[k]; + } + __syncthreads(); + + // layer_input: warp>0 threads process hidden in parallel. + if (warp_id > 0) + { + float pm[HC_MULT]; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + pm[j] = s_pre_mix[t][j]; + + __nv_bfloat16 const* rbase = residual_out + static_cast(tok) * HC_MULT * hidden_size; + __nv_bfloat16* obase = layer_input_out + static_cast(tok) * hidden_size; + int const p2_tid = tid - WARP_SIZE; + constexpr int p2_threads = BLOCK_SIZE - WARP_SIZE; + constexpr int BF16_VEC_LI = 8; + + for (int h = p2_tid * BF16_VEC_LI; h < hidden_size; h += p2_threads * BF16_VEC_LI) + { + float acc_li[BF16_VEC_LI] = {}; +#pragma unroll + for (int j = 0; j < HC_MULT; j++) + { + uint4 raw = *reinterpret_cast(&rbase[j * hidden_size + h]); + __nv_bfloat162 const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raw); +#pragma unroll + for (int v = 0; v < BF16_VEC_LI / 2; v++) + { + float2 f = __bfloat1622float2(pairs[v]); + acc_li[2 * v + 0] += pm[j] * f.x; + acc_li[2 * v + 1] += pm[j] * f.y; + } + } + uint4 out_raw; + __nv_bfloat162* opairs = reinterpret_cast<__nv_bfloat162*>(&out_raw); +#pragma unroll + for (int v = 0; v < BF16_VEC_LI / 2; v++) + opairs[v] = __float22bfloat162_rn(make_float2(acc_li[2 * v], acc_li[2 * v + 1])); + *reinterpret_cast(&obase[h]) = out_raw; + } + } + __syncthreads(); + } +} + +} // namespace fused_fma_kernels diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh index 90cb76c507e1..520eaf68b8cf 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/routing/RoutingKernelTopK.cuh @@ -237,6 +237,99 @@ struct Sort<4, RedType> } }; +template +struct Sort<5, RedType> +{ + static __device__ void run(RedType* topK) + { + TOPK_SWAP(0, 1); + TOPK_SWAP(3, 4); + TOPK_SWAP(2, 4); + TOPK_SWAP(2, 3); + TOPK_SWAP(1, 4); + TOPK_SWAP(0, 3); + TOPK_SWAP(0, 2); + TOPK_SWAP(1, 3); + TOPK_SWAP(1, 2); + } +}; + +template +struct Sort<6, RedType> +{ + static __device__ void run(RedType* topK) + { + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 2); + TOPK_SWAP(0, 1); + TOPK_SWAP(4, 5); + TOPK_SWAP(3, 5); + TOPK_SWAP(3, 4); + TOPK_SWAP(0, 3); + TOPK_SWAP(1, 4); + TOPK_SWAP(2, 5); + TOPK_SWAP(2, 4); + TOPK_SWAP(1, 3); + TOPK_SWAP(2, 3); + } +}; + +template +struct Sort<7, RedType> +{ + static __device__ void run(RedType* topK) + { + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 2); + TOPK_SWAP(0, 1); + TOPK_SWAP(3, 4); + TOPK_SWAP(5, 6); + TOPK_SWAP(3, 5); + TOPK_SWAP(4, 6); + TOPK_SWAP(4, 5); + TOPK_SWAP(0, 4); + TOPK_SWAP(0, 3); + TOPK_SWAP(1, 5); + TOPK_SWAP(2, 6); + TOPK_SWAP(2, 5); + TOPK_SWAP(1, 3); + TOPK_SWAP(2, 4); + TOPK_SWAP(2, 3); + } +}; + +template +struct Sort<8, RedType> +{ + static __device__ void run(RedType* topK) + { + TOPK_SWAP(0, 1); + TOPK_SWAP(2, 3); + TOPK_SWAP(4, 5); + TOPK_SWAP(6, 7); + + TOPK_SWAP(0, 2); + TOPK_SWAP(1, 3); + TOPK_SWAP(4, 6); + TOPK_SWAP(5, 7); + + TOPK_SWAP(1, 2); + TOPK_SWAP(5, 6); + + TOPK_SWAP(0, 4); + TOPK_SWAP(1, 5); + TOPK_SWAP(2, 6); + TOPK_SWAP(3, 7); + + TOPK_SWAP(2, 4); + TOPK_SWAP(3, 5); + + TOPK_SWAP(1, 2); + TOPK_SWAP(3, 4); + TOPK_SWAP(5, 6); + } +}; + //////////////////////////////////////////////////////////////////////////////////////////////////// template diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h index 883ad64a6554..30554747d820 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h @@ -85,8 +85,10 @@ enum class RoutingMethodType : int64_t MiniMax2 = 5, // SigmoidRenorm: Sigmoid -> TopK -> Renormalize SigmoidRenorm = 6, + // DeepSeek-V4 + DeepSeekV4 = 7, // Unspecified - Unspecified = 7, + Unspecified = 8, }; inline int32_t maybeGetMinTokenCount(int32_t numPaddedTokens, int32_t hiddenSize, int32_t dtypeSizeBits) @@ -107,6 +109,7 @@ inline std::string serializeMoeRoutingMethodType(RoutingMethodType routingMethod case RoutingMethodType::RenormalizeNaive: return "RenormalizeNaive"; case RoutingMethodType::MiniMax2: return "MiniMax2"; case RoutingMethodType::SigmoidRenorm: return "SigmoidRenorm"; + case RoutingMethodType::DeepSeekV4: return "DeepSeekV4"; default: TLLM_CHECK_WITH_INFO(false, "Invalid routing method"); return ""; }; } diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 1942a00556b4..ce6217f0fbd1 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -84,6 +84,7 @@ add_library( mambaConv1dOp.cpp moeOp.cpp moeUtilOp.cpp + moeGateOp.cpp moeCommOp.cpp moeAlltoAllOp.cpp moeLoadBalanceOp.cpp @@ -99,6 +100,7 @@ add_library( IndexerKCacheGatherOp.cpp IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp + mlaRopeInplaceOp.cpp ncclCommunicatorOp.cpp allocateOutput.cpp parallelDecodeKVCacheUpdateOp.cpp @@ -123,7 +125,9 @@ add_library( fusedGemmAllreduceOp.cpp convertReqIndexToGlobalOp.cpp trtllmGenQKVProcessOp.cpp - inplaceSliceCopyOp.cpp) + inplaceSliceCopyOp.cpp + mhcOp.cpp + compressorOp.cpp) set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_libraries( th_common PRIVATE ${TORCH_LIBRARIES} th_utils ${Python3_LIBRARIES} diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index 696b8d471043..34edb1af57c0 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -38,9 +38,9 @@ namespace torch_ext void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, th::Tensor const& indices, int64_t next_n, int64_t index_topk, std::optional const& pre_idx, - std::optional const& heuristic_scratch, std::optional const& done_counter_scratch, - std::optional const& scratch, bool is_prefill) + std::optional const& heuristic_scratch, int64_t compress_ratio) { + TORCH_CHECK(compress_ratio > 0, "compress_ratio must be greater than 0"); TORCH_CHECK(logits.is_cuda() && seq_lens.is_cuda() && indices.is_cuda(), "logits, seq_lens, and indices must be CUDA tensors"); @@ -110,70 +110,43 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t heuristicScratchPtr = scratchTensor.data_ptr(); } - // Multi-pass radix scratch. Pre-allocate it for CUDA-graph stable - // addresses, or let the wrapper allocate internally below. - void* multiPassScratchPtr = nullptr; - size_t multiPassScratchBytes = 0; - if (scratch.has_value()) - { - auto const& t = scratch.value(); - TORCH_CHECK(t.is_cuda(), "scratch must be a CUDA tensor"); - TORCH_CHECK(t.device() == logits.device(), "scratch must be on the same device as logits"); - TORCH_CHECK(t.is_contiguous(), "scratch must be contiguous"); - TORCH_CHECK(t.scalar_type() == at::ScalarType::Byte, "scratch must be uint8"); - multiPassScratchPtr = t.data_ptr(); - multiPassScratchBytes = static_cast(t.numel()); - } - - // 0 lets invokeIndexerTopKDecode pick its numRows-aware threshold. - int32_t const splitWorkThreshold = 0; + int32_t splitWorkThreshold = 200 * 1000; auto stream = at::cuda::getCurrentCUDAStream(logits.get_device()); - // Mirror the kernel's threshold so we only allocate when split-work fires. - int32_t const adaptiveSplitWorkThreshold = 200 * 1000; - th::Tensor scratch_internal; - if (num_columns >= adaptiveSplitWorkThreshold && multiPassScratchPtr == nullptr) - { - size_t const bytes = tk::indexerTopKDecodeScratchBytes(num_rows, num_columns, static_cast(index_topk)); - // Zero-init: the radix kernel reads the per-row state on first call. - scratch_internal - = th::zeros({static_cast(bytes)}, th::TensorOptions().dtype(th::kByte).device(logits.device())); - multiPassScratchPtr = scratch_internal.data_ptr(); - multiPassScratchBytes = bytes; - } - - // `done_counter_scratch` is kept on the op signature for source compat - // but unused since the fused split-work tier was removed. Warn once so - // callers still passing it know to drop the argument. - if (done_counter_scratch.has_value()) - { - TORCH_WARN_ONCE( - "indexer_topk_decode: `done_counter_scratch` is deprecated and ignored since the fused split-work tier " - "was removed; drop this argument. The multi-pass radix path allocates its scratch internally, or you may " - "pass a pre-allocated buffer via `scratch` for CUDA-graph capture."); - } - if (logits_dtype == at::ScalarType::Float) { + // fp32 path — full Scheme X v1.2 dispatcher (GVR / Insertion / Radix / + // Radix-split-work). aux_logits/aux_indices needed only by split-work. + th::Tensor aux_indices = th::empty({0}, th::TensorOptions().dtype(th::kInt32).device(logits.device())); + th::Tensor aux_logits = th::empty({0}, th::TensorOptions().dtype(th::kFloat32).device(logits.device())); + constexpr auto multipleBlocksPerRowConfig = 10; + if (num_columns >= splitWorkThreshold) + { + aux_indices = th::empty({num_rows, multipleBlocksPerRowConfig, index_topk}, + th::TensorOptions().dtype(th::kInt32).device(logits.device())); + aux_logits = th::empty({num_rows, multipleBlocksPerRowConfig, index_topk}, + th::TensorOptions().dtype(th::kFloat32).device(logits.device())); + } tk::invokeIndexerTopKDecode(logits.data_ptr(), seq_lens.data_ptr(), indices.data_ptr(), - splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), - static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast(heuristicScratchPtr), stream, multiPassScratchPtr, multiPassScratchBytes, is_prefill); + aux_logits.data_ptr(), aux_indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, + logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, + preIdxStride, preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), + stream); } else if (logits_dtype == at::ScalarType::BFloat16) { tk::invokeIndexerTopKDecode(reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, - preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), stream, multiPassScratchPtr, - multiPassScratchBytes, is_prefill); + preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), + static_cast(compress_ratio), stream); } else // Half { tk::invokeIndexerTopKDecode(reinterpret_cast<__half const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast<__half*>(heuristicScratchPtr), stream, multiPassScratchPtr, multiPassScratchBytes, is_prefill); + static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), stream); } } @@ -188,6 +161,8 @@ void indexer_topk_prefill(th::Tensor const& logits, th::Tensor const& row_starts TORCH_CHECK(indices.dim() == 2, "indices must be a 2D Tensor"); TORCH_CHECK(logits.dim() == 2, "logits must be a 2D Tensor"); + TORCH_CHECK(indices.size(1) >= index_topk, "indices.size(1) must be >= index_topk, got ", indices.size(1), " < ", + index_topk); auto const inputSize = logits.sizes(); auto const numRows64 = inputSize[0]; @@ -213,16 +188,6 @@ void indexer_topk_prefill(th::Tensor const& logits, th::Tensor const& row_starts static_cast(logits_stride_1), static_cast(index_topk), stream); } -// Returns the size in bytes of the `scratch` buffer required by -// indexer_topk_decode's multi-pass radix path for the given shape. Callers -// can allocate `torch.empty(size, dtype=torch.uint8, device='cuda')` and -// pass the result as the `scratch` argument. -int64_t indexer_topk_decode_scratch_bytes(int64_t num_rows, int64_t num_columns, int64_t index_topk) -{ - return static_cast(tk::indexerTopKDecodeScratchBytes( - static_cast(num_rows), static_cast(num_columns), static_cast(index_topk))); -} - } // end namespace torch_ext TRTLLM_NAMESPACE_END @@ -231,9 +196,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, " - "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, Tensor? done_counter_scratch=None, " - "Tensor? scratch=None, bool is_prefill=False) -> ()"); - m.def("indexer_topk_decode_scratch_bytes(int num_rows, int num_columns, int index_topk) -> int"); + "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, int compress_ratio=1) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) @@ -241,11 +204,6 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("indexer_topk_decode", &tensorrt_llm::torch_ext::indexer_topk_decode); } -TORCH_LIBRARY_IMPL(trtllm, CompositeExplicitAutograd, m) -{ - m.impl("indexer_topk_decode_scratch_bytes", &tensorrt_llm::torch_ext::indexer_topk_decode_scratch_bytes); -} - TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( diff --git a/cpp/tensorrt_llm/thop/compressorOp.cpp b/cpp/tensorrt_llm/thop/compressorOp.cpp new file mode 100644 index 000000000000..3de3a9821e95 --- /dev/null +++ b/cpp/tensorrt_llm/thop/compressorOp.cpp @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "tensorrt_llm/kernels/compressorKernels/compressorKernels.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +namespace tk = tensorrt_llm::kernels::compressor; + +namespace +{ + +// Decode kernel: write tokens to paged cache + conditional compression +void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf16 + torch::Tensor ape, // [compress_ratio, state_dim] fp32 + torch::Tensor paged_kv, // [num_blocks, page_size, state_dim] bf16 + torch::Tensor paged_score, // [num_blocks, page_size, state_dim] bf16 + torch::Tensor block_table_kv, // [bsz, max_blocks] int32 + torch::Tensor block_table_score, // [bsz, max_blocks] int32 + torch::Tensor output, // [total_outputs, head_dim] bf16 + torch::Tensor kv_lens, // [bsz] int32 + torch::Tensor start_pos, // [bsz] int32 + torch::Tensor cu_seq_lens, // [bsz+1] int32 + torch::Tensor cu_kv_comp, // [bsz+1] int32 + int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t next_n) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + int io_eb = static_cast(kv_score.element_size()); + int out_eb = static_cast(output.element_size()); + + tk::pagedKvCompressLaunch(kv_score.data_ptr(), ape.data_ptr(), paged_kv.data_ptr(), paged_score.data_ptr(), + block_table_kv.data_ptr(), block_table_score.data_ptr(), output.data_ptr(), + kv_lens.data_ptr(), start_pos.data_ptr(), cu_seq_lens.data_ptr(), + cu_kv_comp.data_ptr(), + static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), + static_cast(head_dim), static_cast(compress_ratio), static_cast(next_n), io_eb, out_eb, stream); +} + +// Prefill kernel: bulk compression with state update +void compressorPrefillReductionOp(torch::Tensor kv_score, torch::Tensor ape, torch::Tensor paged_kv, + torch::Tensor paged_score, torch::Tensor block_table_kv, torch::Tensor block_table_score, torch::Tensor output, + torch::Tensor kv_lens, torch::Tensor start_pos, torch::Tensor cu_seq_lens, torch::Tensor cu_kv_comp, + int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t max_outputs) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + int io_eb = static_cast(kv_score.element_size()); + int out_eb = static_cast(output.element_size()); + + tk::prefillReductionLaunch(kv_score.data_ptr(), ape.data_ptr(), paged_kv.data_ptr(), paged_score.data_ptr(), + block_table_kv.data_ptr(), block_table_score.data_ptr(), output.data_ptr(), + kv_lens.data_ptr(), start_pos.data_ptr(), cu_seq_lens.data_ptr(), + cu_kv_comp.data_ptr(), + static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), + static_cast(head_dim), static_cast(compress_ratio), static_cast(max_outputs), io_eb, out_eb, + stream); +} + +// Fused postprocess + scatter: RMSNorm + RoPE + Hadamard + paged scatter in one kernel +void compressorPostProcessScatterOp(torch::Tensor kv_comp, // [total_tokens, head_dim] input + std::optional kv_out, // [total_tokens, head_dim] output (optional) + torch::Tensor rms_weight, // [head_dim] + double rms_eps, + torch::Tensor cos_sin_table, // [max_pos, 2, rope_dim/2] + torch::Tensor position_ids, // [total_tokens] + int64_t nope_dim, int64_t rope_dim, + torch::Tensor kv_cache, // paged cache buffer + torch::Tensor num_outputs, // [bsz] int32 + torch::Tensor cu_kv_comp, // [bsz+1] int32 + torch::Tensor start_pos, // [bsz] int32 + torch::Tensor block_offsets, // [bsz, max_blocks] int32 + torch::Tensor compressed_mask, // [total_tokens] bool — per-token mask + int64_t tokens_per_block, int64_t cache_scale_type, bool rotate_activation, + std::optional quant_output, std::optional scale_output) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + + TORCH_CHECK(cos_sin_table.scalar_type() == at::kFloat, + "cos_sin_table must be float32, got ", cos_sin_table.scalar_type()); + TORCH_CHECK(cos_sin_table.is_contiguous(), "cos_sin_table must be contiguous"); + TORCH_CHECK(position_ids.scalar_type() == at::kInt, + "position_ids must be int32, got ", position_ids.scalar_type()); + TORCH_CHECK(position_ids.is_contiguous(), "position_ids must be contiguous"); + TORCH_CHECK(compressed_mask.scalar_type() == at::kBool, + "compressed_mask must be bool, got ", compressed_mask.scalar_type()); + + tk::postProcessScatterLaunch(kv_comp.data_ptr(), kv_out.has_value() ? kv_out->data_ptr() : nullptr, + rms_weight.data_ptr(), static_cast(rms_eps), cos_sin_table.data_ptr(), + position_ids.data_ptr(), static_cast(nope_dim), static_cast(rope_dim), + kv_cache.data_ptr(), num_outputs.data_ptr(), cu_kv_comp.data_ptr(), + start_pos.data_ptr(), block_offsets.data_ptr(), + reinterpret_cast(compressed_mask.data_ptr()), + static_cast(num_outputs.size(0)), // batch_size + static_cast(tokens_per_block), + static_cast(kv_comp.size(1)), // head_dim + static_cast(block_offsets.size(1)), // max_blocks + static_cast(kv_comp.element_size()), // elem_bytes + static_cast(kv_comp.size(0)), // total_tokens + static_cast(cache_scale_type), rotate_activation, + quant_output.has_value() ? quant_output->data_ptr() : nullptr, + scale_output.has_value() ? scale_output->data_ptr() : nullptr, stream); +} + +} // anonymous namespace + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "compressor_paged_kv_compress(" + "Tensor kv_score, Tensor ape, " + "Tensor(a!) paged_kv, Tensor(b!) paged_score, " + "Tensor block_table_kv, Tensor block_table_score, " + "Tensor(c!) output, " + "Tensor kv_lens, Tensor start_pos, " + "Tensor cu_seq_lens, Tensor cu_kv_comp, " + "int batch_size, int page_size, " + "int head_dim, int compress_ratio, " + "int next_n) -> ()"); + + m.def( + "compressor_prefill_reduction(" + "Tensor kv_score, Tensor ape, " + "Tensor(a!) paged_kv, Tensor(b!) paged_score, " + "Tensor block_table_kv, Tensor block_table_score, " + "Tensor(c!) output, " + "Tensor kv_lens, Tensor start_pos, " + "Tensor cu_seq_lens, Tensor cu_kv_comp, " + "int batch_size, int page_size, " + "int head_dim, int compress_ratio, " + "int max_outputs) -> ()"); + + m.def( + "compressor_postprocess_scatter(" + "Tensor kv_comp, Tensor(a!)? kv_out, " + "Tensor rms_weight, float rms_eps, " + "Tensor cos_sin_table, Tensor position_ids, " + "int nope_dim, int rope_dim, " + "Tensor(b!) kv_cache, " + "Tensor num_outputs, Tensor cu_kv_comp, " + "Tensor start_pos, Tensor block_offsets, " + "Tensor compressed_mask, " + "int tokens_per_block, int cache_scale_type, " + "bool rotate_activation, " + "Tensor(c!)? quant_output, Tensor(d!)? scale_output) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("compressor_paged_kv_compress", &compressorPagedKvCompressOp); + m.impl("compressor_prefill_reduction", &compressorPrefillReductionOp); + m.impl("compressor_postprocess_scatter", &compressorPostProcessScatterOp); +} diff --git a/cpp/tensorrt_llm/thop/mhcOp.cpp b/cpp/tensorrt_llm/thop/mhcOp.cpp new file mode 100644 index 000000000000..cb87fac0651f --- /dev/null +++ b/cpp/tensorrt_llm/thop/mhcOp.cpp @@ -0,0 +1,208 @@ +/* + * Copyright (c) 2022-2025, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "tensorrt_llm/kernels/mhcKernels/mhcKernels.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +namespace tk = tensorrt_llm::kernels::mhc; + +namespace +{ + +void mhcBigFuseOp(torch::Tensor y_acc, torch::Tensor r_acc, torch::Tensor residual, torch::Tensor hc_scale, + torch::Tensor hc_base, torch::Tensor post_mix, torch::Tensor comb_mix, torch::Tensor layer_input, int64_t M, + int64_t K, int64_t hidden_size, double rms_eps, double hc_pre_eps, double hc_sinkhorn_eps, + double hc_post_mult_value, int64_t sinkhorn_repeat, int64_t num_splits, int64_t block_size) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + + tk::mhcBigFuseLaunch(y_acc.data_ptr(), r_acc.data_ptr(), + reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()), hc_scale.data_ptr(), + hc_base.data_ptr(), post_mix.data_ptr(), comb_mix.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(layer_input.data_ptr()), static_cast(M), + static_cast(K), static_cast(hidden_size), static_cast(rms_eps), static_cast(hc_pre_eps), + static_cast(hc_sinkhorn_eps), static_cast(hc_post_mult_value), static_cast(sinkhorn_repeat), + static_cast(num_splits), static_cast(block_size), stream); +} + +void mhcGemmSqrsumFmaOp(torch::Tensor x, torch::Tensor w, torch::Tensor y, torch::Tensor r, int64_t M, int64_t N, + int64_t K, int64_t tile_n, int64_t tile_m) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + + tk::mhcGemmSqrsumFmaLaunch(reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), w.data_ptr(), + y.data_ptr(), r.data_ptr(), static_cast(M), static_cast(N), static_cast(K), + static_cast(tile_n), static_cast(tile_m), stream); +} + +void mhcHcHeadApplyOp(torch::Tensor mixes, torch::Tensor sqrsum, torch::Tensor x, torch::Tensor out, + torch::Tensor scale, torch::Tensor base_t, int64_t M, int64_t mult, int64_t hidden_size, int64_t K, double norm_eps, + double eps) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + + tk::mhcHcHeadApplyLaunch(mixes.data_ptr(), sqrsum.data_ptr(), + reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), scale.data_ptr(), + base_t.data_ptr(), static_cast(M), static_cast(mult), static_cast(hidden_size), + static_cast(K), static_cast(norm_eps), static_cast(eps), stream); +} + +void mhcPostMappingOp(torch::Tensor residual, torch::Tensor x, torch::Tensor post_mix, torch::Tensor comb_mix, + torch::Tensor out, int64_t B, int64_t hidden_size) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + + tk::mhcPostMappingLaunch(reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()), post_mix.data_ptr(), + comb_mix.data_ptr(), reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(B), + static_cast(hidden_size), stream); +} + +void mhcFusedHcOp(torch::Tensor x_prev, torch::Tensor residual_prev, torch::Tensor post_mix_prev, + torch::Tensor comb_mix_prev, torch::Tensor w, torch::Tensor hc_scale, torch::Tensor hc_base, + torch::Tensor residual_cur, torch::Tensor post_mix_cur, torch::Tensor comb_mix_cur, torch::Tensor layer_input_cur, + torch::Tensor y_acc_workspace, torch::Tensor r_acc_workspace, torch::Tensor done_counter_workspace, int64_t M, + int64_t hidden_size, int64_t hc_mult, double rms_eps, double hc_pre_eps, double hc_sinkhorn_eps, + double hc_post_mult_value, int64_t sinkhorn_repeat, int64_t backend, int64_t tile_n, int64_t num_k_splits, + int64_t bigfuse_block_size, int64_t tile_m) +{ + auto stream = at::cuda::getCurrentCUDAStream(); + + // backend codes: + // 0 = fused_half_mma (2-kernel: fused_tf32_pmap_gemm_rout + mhcBigFuseKernel) + // 1 = fused_half_fma (2-kernel: fused_pmap_gemm_fma_ksplit + mhcBigFuseKernel) + // 2 = fused_all_mma (1-kernel: fused_allinone_tf32_pmap_gemm_atomic_impl, Path D) + // 3 = fused_all_fma (1-kernel: fused_pmap_gemm_fma_allinone, Path F) + if (backend == 3) + { + tk::mhcFusedHcFmaAllInOneLaunch(reinterpret_cast<__nv_bfloat16 const*>(x_prev.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(residual_prev.data_ptr()), + post_mix_prev.data_ptr(), comb_mix_prev.data_ptr(), w.data_ptr(), + hc_scale.data_ptr(), hc_base.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(residual_cur.data_ptr()), post_mix_cur.data_ptr(), + comb_mix_cur.data_ptr(), reinterpret_cast<__nv_bfloat16*>(layer_input_cur.data_ptr()), + y_acc_workspace.data_ptr(), r_acc_workspace.data_ptr(), + done_counter_workspace.data_ptr(), static_cast(M), static_cast(hidden_size), + static_cast(hc_mult), static_cast(tile_n), static_cast(num_k_splits), + static_cast(tile_m), static_cast(rms_eps), static_cast(hc_pre_eps), + static_cast(hc_sinkhorn_eps), static_cast(hc_post_mult_value), + static_cast(sinkhorn_repeat), stream); + return; + } + if (backend == 2) + { + tk::mhcFusedHcAllInOneLaunch(reinterpret_cast<__nv_bfloat16 const*>(x_prev.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(residual_prev.data_ptr()), + post_mix_prev.data_ptr(), comb_mix_prev.data_ptr(), w.data_ptr(), + hc_scale.data_ptr(), hc_base.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(residual_cur.data_ptr()), post_mix_cur.data_ptr(), + comb_mix_cur.data_ptr(), reinterpret_cast<__nv_bfloat16*>(layer_input_cur.data_ptr()), + y_acc_workspace.data_ptr(), r_acc_workspace.data_ptr(), + done_counter_workspace.data_ptr(), static_cast(M), static_cast(hidden_size), + static_cast(hc_mult), static_cast(num_k_splits), static_cast(rms_eps), + static_cast(hc_pre_eps), static_cast(hc_sinkhorn_eps), static_cast(hc_post_mult_value), + static_cast(sinkhorn_repeat), stream); + return; + } + if (backend == 1) + { + tk::mhcFusedHcFmaLaunch(reinterpret_cast<__nv_bfloat16 const*>(x_prev.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(residual_prev.data_ptr()), + post_mix_prev.data_ptr(), comb_mix_prev.data_ptr(), w.data_ptr(), + hc_scale.data_ptr(), hc_base.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(residual_cur.data_ptr()), post_mix_cur.data_ptr(), + comb_mix_cur.data_ptr(), reinterpret_cast<__nv_bfloat16*>(layer_input_cur.data_ptr()), + y_acc_workspace.data_ptr(), r_acc_workspace.data_ptr(), static_cast(M), + static_cast(hidden_size), static_cast(hc_mult), static_cast(tile_n), + static_cast(num_k_splits), static_cast(bigfuse_block_size), static_cast(rms_eps), + static_cast(hc_pre_eps), static_cast(hc_sinkhorn_eps), static_cast(hc_post_mult_value), + static_cast(sinkhorn_repeat), stream); + return; + } + + tk::mhcFusedHcLaunch(reinterpret_cast<__nv_bfloat16 const*>(x_prev.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(residual_prev.data_ptr()), post_mix_prev.data_ptr(), + comb_mix_prev.data_ptr(), w.data_ptr(), hc_scale.data_ptr(), hc_base.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(residual_cur.data_ptr()), post_mix_cur.data_ptr(), + comb_mix_cur.data_ptr(), reinterpret_cast<__nv_bfloat16*>(layer_input_cur.data_ptr()), + y_acc_workspace.data_ptr(), r_acc_workspace.data_ptr(), static_cast(M), + static_cast(hidden_size), static_cast(hc_mult), static_cast(num_k_splits), + static_cast(bigfuse_block_size), static_cast(rms_eps), static_cast(hc_pre_eps), + static_cast(hc_sinkhorn_eps), static_cast(hc_post_mult_value), static_cast(sinkhorn_repeat), + stream); +} + +} // anonymous namespace + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "mhc_big_fuse(" + "Tensor y_acc, Tensor r_acc, Tensor residual, " + "Tensor hc_scale, Tensor hc_base, " + "Tensor(a!) post_mix, Tensor(b!) comb_mix, Tensor(c!) layer_input, " + "int M, int K, int hidden_size, " + "float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, " + "float hc_post_mult_value, int sinkhorn_repeat, int num_splits, " + "int block_size=0) -> ()"); + + m.def( + "mhc_gemm_sqrsum_fma(" + "Tensor x, Tensor w, Tensor(a!) y, Tensor(b!) r, " + "int M, int N, int K, " + "int tile_n=0, int tile_m=0) -> ()"); + + m.def( + "mhc_hc_head_apply(" + "Tensor mixes, Tensor sqrsum, Tensor x, Tensor(a!) out, " + "Tensor scale, Tensor base_t, " + "int M, int mult, int hidden_size, int K, " + "float norm_eps, float eps) -> ()"); + + m.def( + "mhc_post_mapping(" + "Tensor residual, Tensor x, " + "Tensor post_mix, Tensor comb_mix, Tensor(a!) out, " + "int B, int hidden_size) -> ()"); + + m.def( + "mhc_fused_hc(" + "Tensor x_prev, Tensor residual_prev, " + "Tensor post_mix_prev, Tensor comb_mix_prev, " + "Tensor w, Tensor hc_scale, Tensor hc_base, " + "Tensor(a!) residual_cur, Tensor(b!) post_mix_cur, " + "Tensor(c!) comb_mix_cur, Tensor(d!) layer_input_cur, " + "Tensor(e!) y_acc_workspace, Tensor(f!) r_acc_workspace, " + "Tensor(g!) done_counter_workspace, " + "int M, int hidden_size, int hc_mult, " + "float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, " + "float hc_post_mult_value, int sinkhorn_repeat, " + "int backend=0, int tile_n=0, int num_k_splits=0, int bigfuse_block_size=0, " + "int tile_m=1) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("mhc_big_fuse", &mhcBigFuseOp); + m.impl("mhc_gemm_sqrsum_fma", &mhcGemmSqrsumFmaOp); + m.impl("mhc_hc_head_apply", &mhcHcHeadApplyOp); + m.impl("mhc_post_mapping", &mhcPostMappingOp); + m.impl("mhc_fused_hc", &mhcFusedHcOp); +} diff --git a/cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp b/cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp new file mode 100644 index 000000000000..073abdf5b871 --- /dev/null +++ b/cpp/tensorrt_llm/thop/mlaRopeInplaceOp.cpp @@ -0,0 +1,92 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 1993-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. + */ + +#include "tensorrt_llm/kernels/mlaKernels.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include +#include + +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +// Apply RoPE in-place to only the last rope_dim elements of each head. +// data: [num_tokens, num_heads, nope_dim + rope_dim] +// position_ids: [num_tokens] +// cos_sin_cache: [max_positions, 2, rope_dim/2] float +// is_neox: true = neox style (first/second half), false = interleaved style (even/odd) +void mla_rope_inplace(torch::Tensor data, torch::Tensor position_ids, torch::Tensor cos_sin_cache, int64_t num_heads, + int64_t nope_dim, int64_t rope_dim, bool inverse, bool is_neox) +{ + auto stream = at::cuda::getCurrentCUDAStream(data.get_device()); + int const num_tokens = data.size(0); + auto const dtype = data.scalar_type(); + + TORCH_CHECK(data.dim() == 3, "data must be 3D [num_tokens, num_heads, head_dim]"); + TORCH_CHECK(data.size(1) == num_heads, "data.size(1) must equal num_heads"); + TORCH_CHECK(data.size(2) == nope_dim + rope_dim, "data.size(2) must equal nope_dim + rope_dim"); + TORCH_CHECK(data.is_contiguous(), "data must be contiguous"); + TORCH_CHECK(position_ids.dim() == 1 && position_ids.size(0) == num_tokens, "position_ids shape mismatch"); + TORCH_CHECK(position_ids.scalar_type() == torch::kInt32, "position_ids must be int32"); + TORCH_CHECK(cos_sin_cache.scalar_type() == torch::kFloat32, "cos_sin_cache must be float32"); + + if (dtype == torch::kBFloat16) + { + tk::invokeMLARoPEInplace(static_cast<__nv_bfloat16*>(data.data_ptr()), position_ids.data_ptr(), + cos_sin_cache.data_ptr(), num_tokens, static_cast(num_heads), static_cast(nope_dim), + static_cast(rope_dim), inverse, is_neox, stream); + } + else if (dtype == torch::kFloat16) + { + tk::invokeMLARoPEInplace(static_cast(data.data_ptr()), position_ids.data_ptr(), + cos_sin_cache.data_ptr(), num_tokens, static_cast(num_heads), static_cast(nope_dim), + static_cast(rope_dim), inverse, is_neox, stream); + } + else + { + TORCH_CHECK(false, "mla_rope_inplace: unsupported dtype, expected bf16 or fp16"); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "mla_rope_inplace(" + "Tensor(a!) data" + ", Tensor position_ids" + ", Tensor cos_sin_cache" + ", int num_heads" + ", int nope_dim" + ", int rope_dim" + ", bool inverse" + ", bool is_neox" + ") -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("mla_rope_inplace", &tensorrt_llm::torch_ext::mla_rope_inplace); +} diff --git a/cpp/tensorrt_llm/thop/moeGateOp.cpp b/cpp/tensorrt_llm/thop/moeGateOp.cpp new file mode 100644 index 000000000000..944f886c1403 --- /dev/null +++ b/cpp/tensorrt_llm/thop/moeGateOp.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) 2022-2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/kernels/customMoeRoutingKernels.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +namespace th = torch; +namespace tl = tensorrt_llm; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +static constexpr int kTOPK = 6; + +// PyTorch wrapper function that calls the kernel implementation +void gate_forward(th::Tensor scores_in, // [batch_size, nExperts] - pre-computed from linear(x, weight) + th::Tensor bias, // empty tensor if hash mode + th::Tensor input_ids, // empty tensor if non-hash mode + th::Tensor tid2eid, // empty tensor if non-hash mode + th::Tensor out_weights, // [batch_size, topK] - pre-allocated + th::Tensor out_indices, // [batch_size, topK] - pre-allocated + int64_t topk, double route_scale, bool is_hash) +{ + TORCH_CHECK(topk == kTOPK, "topk must be ", kTOPK); + auto const n_experts = scores_in.size(1); + TORCH_CHECK(n_experts == 256 || n_experts == 384, "n_experts must be 256 or 384"); + TORCH_CHECK(scores_in.scalar_type() == torch::kFloat32, "scores_in must be float32"); + TORCH_CHECK(out_weights.scalar_type() == torch::kFloat32, "out_weights must be float32"); + TORCH_CHECK(out_indices.scalar_type() == torch::kInt32, "out_indices must be int32"); + TORCH_CHECK( + scores_in.is_cuda() && out_weights.is_cuda() && out_indices.is_cuda(), "All tensors must be CUDA tensors"); + TORCH_CHECK( + scores_in.get_device() == out_weights.get_device() && scores_in.get_device() == out_indices.get_device(), + "All tensors must be on the same device"); + + auto const batch_size = scores_in.size(0); + auto stream = at::cuda::getCurrentCUDAStream(scores_in.get_device()); + + // Call the kernel implementation from kernels namespace + kernels::gate_forward(scores_in.data_ptr(), is_hash ? nullptr : bias.data_ptr(), + is_hash ? input_ids.data_ptr() : nullptr, is_hash ? tid2eid.data_ptr() : nullptr, + out_weights.data_ptr(), out_indices.data_ptr(), batch_size, static_cast(n_experts), + static_cast(route_scale), is_hash, stream); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "gate_forward(Tensor scores_in, Tensor bias, Tensor input_ids, Tensor tid2eid, Tensor out_weights, " + "Tensor out_indices, int topk, float route_scale, bool is_hash) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("gate_forward", &tensorrt_llm::torch_ext::gate_forward); +} diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 3bfbece3a549..31bc1ed6112b 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -12,6 +12,7 @@ The following is a table of supported models for the PyTorch backend: | `DeepSeekV2ForCausalLM` [^5] | DeepSeek V2 | `deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct` | | `DeepseekV3ForCausalLM` | DeepSeek-V3, Kimi-K2 | `deepseek-ai/DeepSeek-V3` | | `DeepseekV32ForCausalLM` | DeepSeek-V3.2 | `deepseek-ai/DeepSeek-V3.2` | +| `DeepseekV4ForCausalLM` [^11] | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Pro` | | `ExaoneForCausalLM` [^5] | EXAONE 3.5 | `LGAI-EXAONE/EXAONE-3.5-32B-Instruct` | | `Exaone4ForCausalLM` | EXAONE 4.0 | `LGAI-EXAONE/EXAONE-4.0-32B` | | `ExaoneMoEForCausalLM` | K-EXAONE | `LGAI-EXAONE/K-EXAONE-236B-A23B` | @@ -62,6 +63,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | -------------------------------- | ----------------- | ---------- | -------------------------- | --------------------- | --------------- | --- | ---------------- | ----------------- | ------ | ------------- | ---------------- | -------------- | ------------------------ | --------------------- | --------------- | | `DeepseekV3ForCausalLM` | Yes | Yes | Yes | Yes | Yes [^1] | Yes | No | No | No | Yes | Yes | Yes [^2] | N/A | Yes | Yes | | `DeepseekV32ForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Yes | Yes | +| `DeepseekV4ForCausalLM` [^11] | Yes | Yes | Yes | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | Yes | Untested | Untested | | `Glm4MoeForCausalLM` | Yes | Yes | Yes | Untested | Yes | Yes | No | No | No | Yes | Yes | Untested | N/A | Yes | Yes | | `Qwen3MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | N/A | Yes | Yes | | `Qwen3NextForCausalLM` [^3] | Yes | Yes | Yes | Untested | Yes | No | No | No | No | Yes | Yes | No | No | Untested | Untested | @@ -82,6 +84,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^8]: Supports text and image inputs. The vision tower runs in BF16 even when the text decoder is quantized (FP8 block-scale or NVFP4). The text decoder is also usable standalone (text-only) via the `Step3p5ForCausalLM` architecture. [^9]: Audio modality only supported on E2B/E4B variants. [^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path. +[^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism. # Multimodal Feature Support Matrix (PyTorch Backend) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 10901d87c520..ab76328840c0 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -224,6 +224,7 @@ def add_llm_args(parser): type=str, default=None, nargs='+') + parser.add_argument('--tokenizer_dir', type=str, default=None) return parser @@ -363,6 +364,7 @@ def setup_llm(args, **kwargs): gather_generation_logits=args.return_generation_logits, max_beam_width=args.max_beam_width, orchestrator_type=args.orchestrator_type, + tokenizer=args.tokenizer_dir, **kwargs) use_beam_search = args.max_beam_width > 1 diff --git a/examples/models/core/deepseek_v4/README.md b/examples/models/core/deepseek_v4/README.md new file mode 100644 index 000000000000..fe7d4976c114 --- /dev/null +++ b/examples/models/core/deepseek_v4/README.md @@ -0,0 +1,488 @@ +# DeepSeek-V4 + +This guide walks you through the examples to run DeepSeek-V4 models using NVIDIA TensorRT LLM with +the PyTorch backend. + +DeepSeek-V4 uses the `DeepseekV4ForCausalLM` architecture in TensorRT LLM. Compared with +DeepSeek-V3/R1/V3.2, it has a separate model implementation and sparse attention path. Use the +commands in this guide as starting points and tune the parallelism and memory settings for your +checkpoint and workload. + +Please refer to [this guide](https://nvidia.github.io/TensorRT-LLM/installation/build-from-source-linux.html) +for how to build TensorRT LLM from source and start a TRT-LLM Docker container. + +> [!NOTE] +> This guide assumes that you replace placeholder values such as `` with the +> appropriate paths. Commands in this guide target the PyTorch backend. + + +## Table of Contents + +- [DeepSeek-V4](#deepseek-v4) + - [Table of Contents](#table-of-contents) + - [Hardware Requirements](#hardware-requirements) + - [Downloading the Model Weights](#downloading-the-model-weights) + - [Quick Start](#quick-start) + - [Run a single inference](#run-a-single-inference) + - [Run chat-style prompts](#run-chat-style-prompts) + - [Multi-Token Prediction (MTP)](#multi-token-prediction-mtp) + - [Benchmarking](#benchmarking) + - [Evaluation](#evaluation) + - [Serving](#serving) + - [trtllm-serve](#trtllm-serve) + - [OpenAI-compatible request](#openai-compatible-request) + - [Advanced Configuration](#advanced-configuration) + - [Parallelism](#parallelism) + - [Sparse attention](#sparse-attention) + - [KV cache](#kv-cache) + - [Quantized checkpoints](#quantized-checkpoints) + - [Notes and Troubleshooting](#notes-and-troubleshooting) + + +## Hardware Requirements + +DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`) in the current PyTorch backend +implementation. Pre-Blackwell GPUs are not supported for this model path. + +DeepSeek-V4 has two model scales, and each scale provides Base and Instruct checkpoints. The table +below follows the model list published on the +[DeepSeek-V4 Hugging Face model card](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro): + +| Checkpoint | Total Params | Activated Params | Context Length | Precision | +| --- | --- | --- | --- | --- | +| DeepSeek-V4-Flash-Base | 284B | 13B | 1M | FP8 Mixed | +| DeepSeek-V4-Flash | 284B | 13B | 1M | FP4 + FP8 Mixed | +| DeepSeek-V4-Pro-Base | 1.6T | 49B | 1M | FP8 Mixed | +| DeepSeek-V4-Pro | 1.6T | 49B | 1M | FP4 + FP8 Mixed | + +The minimum number of GPUs depends on the model scale, checkpoint precision, KV cache budget, +maximum sequence length, and runtime batch size. For initial bring-up, an 8xB200 node is enough for +Flash checkpoints and the FP4 + FP8 mixed DeepSeek-V4-Pro checkpoint. DeepSeek-V4-Pro-Base is larger +because it uses FP8 mixed precision; if you want to keep the deployment on a single node, use an +8xB300 node. Multi-node Blackwell deployments are still recommended for larger KV cache budgets, +longer context windows, or higher throughput targets. Tune `--tp_size`, `--ep_size`, +`--max_num_tokens`, and the KV cache memory fraction for your deployment target. + +DeepSeek-V4 requires KV cache block sizes of 128 or 256 tokens. TensorRT LLM defaults DeepSeek-V4 to +`tokens_per_block=128`, but scripts that set their own KV cache config should pass this explicitly. + + +## Downloading the Model Weights + +Choose one of the DeepSeek-V4 checkpoint IDs: + +| Checkpoint | Hugging Face model ID | Prompt format | +| --- | --- | --- | +| DeepSeek-V4-Flash-Base | `deepseek-ai/DeepSeek-V4-Flash-Base` | Raw completion | +| DeepSeek-V4-Flash | `deepseek-ai/DeepSeek-V4-Flash` | Chat/Instruct | +| DeepSeek-V4-Pro-Base | `deepseek-ai/DeepSeek-V4-Pro-Base` | Raw completion | +| DeepSeek-V4-Pro | `deepseek-ai/DeepSeek-V4-Pro` | Chat/Instruct | + +Then download the weights: + +```bash +git lfs install +MODEL_ID=deepseek-ai/DeepSeek-V4-Flash +git clone https://huggingface.co/${MODEL_ID} +``` + +At minimum, the checkpoint config should identify the architecture as DeepSeek-V4: + +```json +{ + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4" +} +``` + +Do not replace the full checkpoint config with this minimal snippet. TensorRT LLM also reads +DeepSeek-V4-specific sparse attention fields such as `compress_ratios`, `window_size` or +`sliding_window`, and indexer settings from the checkpoint config unless you provide a complete +override through `sparse_attention_config`. + + +## Quick Start + +### Run a single inference + +To quickly run DeepSeek-V4, use [examples/llm-api/quickstart_advanced.py](../../../llm-api/quickstart_advanced.py): + +```bash +cd examples/llm-api +python quickstart_advanced.py \ + --model_dir \ + --tp_size 8 \ + --moe_ep_size 8 \ + --tokens_per_block 128 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --kv_cache_fraction 0.5 +``` + +The command above assumes one 8-GPU node. If you use a different number of GPUs, adjust `--tp_size` +and `--moe_ep_size` so that the requested parallelism matches your available world size. DeepSeek-V4 +checkpoints advertise a 1M-token context window; for bring-up, set `--max_seq_len` and the KV cache +memory fraction explicitly, then increase them according to your memory budget. + +### Run chat-style prompts + +DeepSeek-V4 Instruct checkpoints (`DeepSeek-V4-Flash` and `DeepSeek-V4-Pro`) use the checkpoint +reference chat/message format. TensorRT LLM provides a `deepseek_v4` tokenizer wrapper for this +format. Use `custom_tokenizer="deepseek_v4"` only with Instruct checkpoints and chat-style prompts. + +Base checkpoints (`DeepSeek-V4-Flash-Base` and `DeepSeek-V4-Pro-Base`) are completion models. For +Base checkpoints, do not apply a chat template and do not pass `custom_tokenizer="deepseek_v4"`; +send raw text prompts instead. + +```python +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig + +def main(): + llm = LLM( + model="", + backend="pytorch", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + custom_tokenizer="deepseek_v4", + kv_cache_config=KvCacheConfig( + tokens_per_block=128, + free_gpu_memory_fraction=0.5, + ), + max_seq_len=4096, + max_num_tokens=8192, + ) + + messages = [{"role": "user", "content": "Explain TensorRT LLM in one paragraph."}] + prompt = llm.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + outputs = llm.generate([prompt], SamplingParams(max_tokens=128)) + print(outputs[0].outputs[0].text) + + +if __name__ == "__main__": + main() +``` + +### Multi-Token Prediction (MTP) + +If the checkpoint contains MTP layers, run MTP speculative decoding with the one-model flow: + +```bash +cd examples/llm-api +python quickstart_advanced.py \ + --model_dir \ + --tp_size 8 \ + --moe_ep_size 8 \ + --tokens_per_block 128 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --kv_cache_fraction 0.5 \ + --spec_decode_algo MTP \ + --spec_decode_max_draft_len N \ + --use_one_model +``` + +`N` is the number of draft tokens to predict. Start with `N=1` for bring-up, then increase it after +validating accuracy and latency for your workload. + + +## Benchmarking + +The following example prepares a synthetic dataset and runs `trtllm-bench` throughput on one 8-GPU +Blackwell node: + +```bash +trtllm-bench --model \ + --model_path \ + prepare-dataset \ + --output /tmp/deepseek_v4_1k1k.txt \ + token-norm-dist \ + --input-mean 1024 \ + --output-mean 1024 \ + --input-stdev 0 \ + --output-stdev 0 \ + --num-requests 256 + +cat > /tmp/deepseek_v4_config.yml < \ + --model_path \ + throughput \ + --tp 8 \ + --ep 8 \ + --dataset /tmp/deepseek_v4_1k1k.txt \ + --max_batch_size 256 \ + --max_num_tokens 8192 \ + --concurrency 2048 \ + --num_requests 6144 \ + --kv_cache_free_gpu_mem_fraction 0.9 \ + --config /tmp/deepseek_v4_config.yml +``` + +The example enables attention DP because it is typically beneficial for high-throughput, large-batch +workloads. It also uses FP8 KV cache (`kv_cache_config.dtype: fp8`), which is the recommended +starting point for benchmarking DeepSeek-V4 throughput. For checkpoints with MTP layers, enable MTP +for benchmarking as well: use `num_nextn_predict_layers: 1` for throughput-oriented runs, and use +`num_nextn_predict_layers: 3` for low-latency runs. When `enable_attention_dp` is enabled, +`--max_batch_size` is the maximum batch size per local rank; use `--concurrency` high enough to +saturate all ranks. Tune `--max_batch_size`, `--max_num_tokens`, `--concurrency`, MTP depth, and the +KV cache memory fraction for the target ISL/OSL distribution. + + +## Evaluation + +Evaluate model accuracy using `trtllm-eval`. The following commands are for Instruct checkpoints and +apply the DeepSeek-V4 chat template through `--custom_tokenizer deepseek_v4` and +`--apply_chat_template`. For Base checkpoints, remove both flags because Base models expect raw +completion prompts. `--custom_tokenizer` is a top-level `trtllm-eval` option, so keep it before the +dataset subcommand such as `mmlu`, `gsm8k`, or `gpqa_diamond`. + +1. Prepare a configuration file: + +```bash +cat > ./deepseek_v4_config.yml < \ + --tp_size 8 \ + --ep_size 8 \ + --max_batch_size 16 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_config.yml \ + mmlu \ + --apply_chat_template +``` + +3. Evaluate GSM8K with an Instruct checkpoint: + +```bash +trtllm-eval --model \ + --tp_size 8 \ + --ep_size 8 \ + --max_batch_size 16 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_config.yml \ + gsm8k \ + --apply_chat_template \ + --system_prompt "Solve the problem carefully. End your response with a final line exactly in the form #### , using the simplest numeric form without units or trailing zeros." +``` + +The `--system_prompt` constrains the answer format so that the lm-eval `strict-match` +regex (which expects a final `#### ` line) can pick up the model's answer. +Without it, DeepSeek-V4 Instruct checkpoints often return the correct value in a +free-form sentence, which `flexible-extract` recovers but `strict-match` does not. + +4. Evaluate GPQA Diamond with an Instruct checkpoint: + +```bash +trtllm-eval --model \ + --tp_size 8 \ + --ep_size 8 \ + --max_batch_size 16 \ + --max_num_tokens 8192 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_config.yml \ + gpqa_diamond \ + --apply_chat_template +``` + + +## Serving + +### trtllm-serve + +Create a serving config: + +```bash +cat > ./deepseek_v4_serve.yml < \ + --backend pytorch \ + --host 0.0.0.0 \ + --port 8000 \ + --tp_size 8 \ + --ep_size 8 \ + --max_seq_len 4096 \ + --custom_tokenizer deepseek_v4 \ + --config ./deepseek_v4_serve.yml +``` + +The `/v1/chat/completions` API applies chat formatting on the server side, so clients should send +OpenAI-style `messages` rather than preformatted prompt strings. For Base checkpoints, use the same +command but remove `--custom_tokenizer deepseek_v4`. Increase `max_seq_len`, `max_batch_size`, and +the KV cache memory fraction after validating the memory budget for your target deployment. + +### OpenAI-compatible request + +For Instruct checkpoints, send a chat-completions request: + +```bash +curl http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "", + "messages": [ + { + "role": "user", + "content": "Write a short summary of TensorRT LLM." + } + ], + "stream": true, + "max_tokens": 128 + }' +``` + +For Base checkpoints, use the text completions API with a raw prompt: + +```bash +curl http://localhost:8000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "", + "prompt": "TensorRT LLM is", + "stream": true, + "max_tokens": 128 + }' +``` + + +## Advanced Configuration + +### Parallelism + +DeepSeek-V4 supports the same main PyTorch backend parallelism knobs used by other large MoE models: + +- Tensor parallelism (`--tp_size` or `tensor_parallel_size`) shards attention and dense weights. +- Pipeline parallelism (`--pp_size` or `pipeline_parallel_size`) distributes model layers across + pipeline stages, which can help fit larger checkpoints or larger KV cache budgets across more + GPUs. +- Expert parallelism (`--ep_size` or `moe_expert_parallel_size`) distributes routed experts. +- Attention DP (`enable_attention_dp: true`) keeps attention data-parallel across ranks and is + commonly used for high-throughput, large-batch serving. + +For latency-oriented tests, start without attention DP. For throughput-oriented tests, enable +attention DP in YAML: + +```yaml +enable_attention_dp: true +attention_dp_config: + batching_wait_iters: 0 + enable_balance: true + timeout_iters: 60 +``` + +When attention DP is enabled, remember that `max_batch_size` is local-rank batch size. Increase +`concurrency` and `num_requests` accordingly when benchmarking. + +### Sparse attention + +If `sparse_attention_config` is not provided, TensorRT LLM configures DeepSeek-V4 sparse attention +from the model config. It reads fields such as `compress_ratios`, `window_size` or `sliding_window`, +and indexer settings, then constructs the corresponding `DeepSeekV4SparseAttentionConfig`. + +If `sparse_attention_config` is provided, user values override the corresponding sparse attention +settings, subject to the current implementation constraints: `window_size` must be `128`, and +`compress_ratios` must use supported ratios (`1`, `4`, or `128`). If checkpoint `compress_ratios` +are present and longer than the user-provided list, TensorRT LLM keeps the checkpoint list to avoid +silently changing the sparse attention layout. + +Example YAML override: + +```yaml +sparse_attention_config: + algorithm: deepseek_v4 + window_size: 128 + index_topk: 512 +``` + +### KV cache + +DeepSeek-V4 uses `DeepseekV4CacheManager`, a `KvCacheManagerV2` subclass. This manager can describe +different cache layer types per model layer, so DeepSeek-V4 can map sliding-window, compressed, +indexer, and compressor-state caches according to the sparse attention layout from the model config +or user-provided `sparse_attention_config`. + +DeepSeek-V4 KV cache requires: + +- `tokens_per_block` set to `128` or `256`. +- `max_beam_width=1`. +- Blackwell GPUs for the current implementation. + +Use a lower `free_gpu_memory_fraction`, `max_batch_size`, or `max_num_tokens` if the workload runs +out of memory during initialization or prefill. + +### Quantized checkpoints + +TensorRT LLM detects supported quantization metadata from the checkpoint directory, including +`hf_quant_config.json`, `quantization_config`, or `dtypes.json`. For DeepSeek-V4 checkpoints with +MXFP4 routed MoE expert weights, TensorRT LLM automatically applies the routed-expert quantization +configuration. + + +## Notes and Troubleshooting + +- `DeepseekV4CacheManager requires tokens_per_block in [128, 256]`: pass + `--tokens_per_block 128` in `quickstart_advanced.py` or set + `kv_cache_config.tokens_per_block: 128` in YAML. +- `DeepSeek-V4 is not supported on pre-blackwell GPUs`: run on Blackwell GPUs (`SM100+`). +- Out-of-memory during initialization or prefill: reduce `max_batch_size`, `max_num_tokens`, or + `kv_cache_config.free_gpu_memory_fraction`. For bring-up on 8xB200, set `max_seq_len` explicitly + instead of using the checkpoint's 1M-token context length. +- Chat formatting issues with `trtllm-serve` or `trtllm-eval` on Instruct checkpoints: pass + `--custom_tokenizer deepseek_v4`. Do not use this tokenizer wrapper for Base checkpoints. +- Tool-call chat formatting is not supported by the DeepSeek-V4 tokenizer wrapper yet. diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 97f28dca435f..d7ccaccb7092 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1428,15 +1428,15 @@ def launchJob(pipeline, jobName, reuseBuild, enableFailFast, globalVars, platfor def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) { stages = [ - "Release-Check": { - script { - if (GEN_POST_MERGE_BUILDS_ONLY) { - echo "Skipping Release-Check (GenPostMergeBuilds mode: builds only)" - return - } - launchReleaseCheck(this, globalVars) - } - }, + // "Release-Check": { + // script { + // if (GEN_POST_MERGE_BUILDS_ONLY) { + // echo "Skipping Release-Check (GenPostMergeBuilds mode: builds only)" + // return + // } + // launchReleaseCheck(this, globalVars) + // } + // }, "x86_64-Linux": { script { // CBTS deliberately does NOT short-circuit at the arch / Build diff --git a/pyproject.toml b/pyproject.toml index fd5f14508851..245d9cbe14c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ column_limit = 80 [tool.codespell] skip = ".git,3rdparty,triton_kernels,tests/integration/test_input_files**,**.jsonl,**.json" exclude-file = "examples/models/core/whisper/tokenizer.py" -ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut,indext,asend,medias" +ignore-words-list = "rouge,inout,atleast,strat,nd,subtile,thrid,improbe,NotIn,te,iteract,anythin,tru,Tracin,vEw,dOut,indext,asend,medias,ElemT" [tool.autoflake] in-place = true diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py new file mode 100644 index 000000000000..1a6b89f2992a --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py @@ -0,0 +1,4 @@ +from .cache_manager import DeepseekV4CacheManager +from .deepseek_v4 import DeepseekV4AttentionType, DeepseekV4TrtllmAttention + +__all__ = ["DeepseekV4CacheManager", "DeepseekV4AttentionType", "DeepseekV4TrtllmAttention"] diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py new file mode 100644 index 000000000000..c10dd440af25 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -0,0 +1,830 @@ +from collections import defaultdict +from typing import Dict, List, Optional, Tuple + +import torch + +from tensorrt_llm._torch.pyexecutor import llm_request +from tensorrt_llm._torch.pyexecutor.resource_manager import GPU_LEVEL, KVCacheManagerV2, Role +from tensorrt_llm._utils import ( + TensorWrapper, + convert_to_torch_tensor, + get_size_in_bytes, + prefer_pinned, +) +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime import ModelConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BatchDesc, + BufferConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheDesc, + LayerId, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy + +from .compressor import KVCacheDtype +from .deepseek_v4 import ( + DEEPSEEK_V4_SPARSE_RATIO, + DeepseekV4AttentionType, + compress_ratio_has_attention, + get_attn_dim, + get_token_bytes, + is_compress_layer, + is_overlap_compressor, + is_sparse_layer, +) + + +def _estimate_bytes_per_token( + head_dim: int, + index_head_dim: int, + compress_ratios: List[int], + has_fp8_kv_cache, + attn_types: set[DeepseekV4AttentionType] | None = None, +) -> int: + total_bytes = 0 + for ratio in compress_ratios: + for attn in DeepseekV4AttentionType: + if attn_types is not None and attn not in attn_types: + continue + if compress_ratio_has_attention(ratio, attn): + total_bytes += _get_attn_bytes_per_token( + head_dim, + index_head_dim, + ratio, + attn, + has_fp8_kv_cache, + ) + return total_bytes + + +def _get_attn_bytes_per_token( + head_dim: int, + index_head_dim: int, + compress_ratio: int, + attn_type: DeepseekV4AttentionType, + has_fp8_kv_cache: bool, +) -> int: + token_bytes = get_token_bytes( + head_dim, + index_head_dim, + compress_ratio, + attn_type, + has_fp8_kv_cache, + ) + if attn_type in [DeepseekV4AttentionType.COMPRESS, DeepseekV4AttentionType.INDEXER_COMPRESS]: + token_bytes //= compress_ratio + return token_bytes + + +class DeepseekV4CacheManager(KVCacheManagerV2): + fixed_size_attention = { + DeepseekV4AttentionType.SWA, + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + } + # This tensor is for compatibility with AttentionOp, it only contains swa attention. + # kv_cache_pool_pointers contains pool pointers swa pool, shape: [1, 2] + # It assume the KVCacheManagerPy has only one pool for swa attention. + # The second column is always 0. + kv_cache_pool_pointers: torch.Tensor + # This tensor is for compatibility with AttentionOp, it only contains swa attention. + # kv_cache_pool_mapping contains pool id and layer offset for each layer's swa attention, + # shape: [num_local_layers, 2] + kv_cache_pool_mapping: torch.Tensor + # The block size of the (indexer) compressed cache. + # For other attention types, block size is tokens_per_block. + compressed_block_sizes: List[int] + + def __init__( + self, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: int = 1, + max_batch_size: int, + max_beam_width: int = 1, + tokens_per_block: int, + max_seq_len: int, + vocab_size: int, + mapping: Mapping, + dtype: DataType = DataType.BF16, + compressor_dtype: DataType = DataType.FLOAT, + sparse_attn_config: DeepSeekV4SparseAttentionConfig, + max_input_len: Optional[int] = None, + max_num_tokens: Optional[int] = None, + **kwargs, + ) -> None: + # DeepSeek-V4 specific attributes initialization + assert kv_cache_type == CacheTypeCpp.SELFKONLY, "DeepSeek-V4 only supports SELFKONLY" + assert num_kv_heads == 1, "DeepSeek-V4 only supports num_kv_heads == 1" + assert len(sparse_attn_config.compress_ratios) >= num_layers, ( + "The length of compress ratios must be >= the number of layers" + ) + assert dtype in [DataType.BF16, DataType.FP8], ( + f"Unsupported dtype: {dtype}, only support BF16 and FP8" + ) + assert compressor_dtype == DataType.FLOAT, ( + f"Unsupported compressor dtype: {compressor_dtype}, only support FP32/TF32" + ) + + assert tokens_per_block in [128, 256], ( + f"DeepseekV4CacheManager requires tokens_per_block in [128, 256], got {tokens_per_block}. " + f"Set kv_cache_config.tokens_per_block to 128 or 256." + ) + + self.index_head_dim = sparse_attn_config.index_head_dim + self._compress_ratios = sparse_attn_config.compress_ratios + # When MTP is enabled, enlarge the sliding window sizes by + # max_draft_len so that rewinding rejected draft tokens can still + # reach the KV entries that would otherwise have been evicted by the + # sliding-window policy. + spec_config = kwargs.get("spec_config", None) + self._max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 + self._swa_window_size = sparse_attn_config.window_size + self._compressor_dtype = compressor_dtype + # If MTP is enabled, append compress ratios for MTP virtual layers. + # MTP adds (max_draft_len - 1) extra layers that mirror the last real + # layer's attention pattern. Only NEW entries are appended; existing + # per-layer ratios are never modified, so a real layer with ratio==1 + # stays SWA-only. + if self._max_draft_len > 0: + self._compress_ratios = self._compress_ratios + [self._compress_ratios[-1]] * ( + self._max_draft_len - 1 + ) + self.compressed_block_sizes = [tokens_per_block // ratio for ratio in self._compress_ratios] + + # Indexer compressor cache: FP8 blockwise (1 byte/value + per-128 fp32 + # scale). The bf16 / fp8_pertensor presets are reserved for the + # main-attention compressor and are not legal indexer dtypes. + self._indexer_cache_dtype = KVCacheDtype.FP8_BLOCKWISE + self._indexer_dtype = DataType.FP8 + self.use_fp4 = False + self.quant_block_size = 128 + self._indexer_data_size = self.index_head_dim + self._indexer_scale_size = get_size_in_bytes( + self.index_head_dim // self.quant_block_size, DataType.FLOAT + ) + assert self.index_head_dim % self.quant_block_size == 0, ( + f"indexer_head_dim {self.index_head_dim} must be divisible by {self.quant_block_size}" + ) + + # _build_cache_config() needs them to build constraints + self._max_input_len = max_input_len + self._max_num_tokens = max_num_tokens + + # General initialization + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + max_batch_size=max_batch_size, + max_beam_width=max_beam_width, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + vocab_size=vocab_size, + mapping=mapping, + dtype=dtype, + **kwargs, + ) + self.is_vswa = True # DeepSeek-V4 must has VSWA + + # DeepSeek-V4 expects cache of all layers with the same attention type and compress ratio + # to be in the same pool and have the same scale. + self._assert_layer_pool_scale() + + # For DeepSeek-V4 Attention, the base pointer for SWA pool + # Use first PP layer instead of hardcoded 0 for pipeline parallelism. + first_pp_layer = self.pp_layers[0] + self.swa_pool_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_pp_layer, DeepseekV4AttentionType.SWA], Role.KEY + ) + + self.compress_pool_ptrs = {} + # Find first PP layer with each compress ratio for pool pointer lookup. + pp_compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + if 4 in pp_compress_ratios: # indexer compressor + first_layer_with_4 = self.pp_layers[pp_compress_ratios.index(4)] + self.compress_pool_ptrs[4] = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_layer_with_4, DeepseekV4AttentionType.COMPRESS], + Role.KEY, + ) + if 128 in pp_compress_ratios: # compressor + first_layer_with_128 = self.pp_layers[pp_compress_ratios.index(128)] + self.compress_pool_ptrs[128] = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[ + first_layer_with_128, DeepseekV4AttentionType.COMPRESS + ], + Role.KEY, + ) + # Use pinned staging buffer to avoid pageable H2D memcpy + max_num_sequences = max_batch_size * mapping.pp_size + self._host_block_offsets_staging = torch.empty( + (max_num_sequences + 1) * max_beam_width, + 2, # key and value + self.max_blocks_per_seq, + dtype=torch.int32, + pin_memory=prefer_pinned(), + device="cpu", + ) + + def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> torch.Tensor: + """ + Get the buffers for a specific layer and attention type. + + Args: + layer_idx: The layer index + attn_type: The attention type + + Returns: + The buffer tensor (shape: [num_blocks, tokens_per_block, attn_dim]) + For blockwise FP8 layers, shape is [num_blocks, tokens_per_block, attn_dim + scale_size] + """ + layer_id = self._layer_attn_to_layer_id[(layer_idx, attn_type)] + addr = self.impl.get_mem_pool_base_address(layer_id, Role.KEY) + + block_size = self.tokens_per_block + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + block_size = self.compressed_block_sizes[layer_idx] + + attn_dim = get_attn_dim( + self.head_dim, self.index_head_dim, self._compress_ratios[layer_idx], attn_type + ) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + # Indexer always pack data + per-block scales into the same row. + dim_per_token = self._indexer_data_size + self._indexer_scale_size + else: + dim_per_token = attn_dim + + shape = ( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY), + block_size, + dim_per_token, + ) + + dtype = self.dtype + # (indexer) compressor state and score use compressor_dtype + if attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + dtype = self._compressor_dtype + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + dtype = self._indexer_dtype + + return convert_to_torch_tensor(TensorWrapper(addr, dtype, shape)) + + def _get_window_size(self, compress_ratio: int, attn_type: DeepseekV4AttentionType) -> int: + if attn_type == DeepseekV4AttentionType.SWA: + base_window_size = self._swa_window_size + elif attn_type in self.fixed_size_attention: + state_factor = 2 if is_overlap_compressor(compress_ratio) else 1 + base_window_size = state_factor * compress_ratio + else: + raise ValueError(f"Unsupported fixed-size attention type: {attn_type}") + return base_window_size + self._max_draft_len + + def _build_pool_mapping_tensors(self) -> Tuple[torch.Tensor, torch.Tensor]: + first_pp_layer = self.pp_layers[0] + swa_bytes_per_block = self._get_attn_bytes_per_block( + DeepseekV4AttentionType.SWA, first_pp_layer + ) + swa_pool_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[first_pp_layer, DeepseekV4AttentionType.SWA], Role.KEY + ) + + def _get_layer_offset(pp_layer: int) -> int: + buffer_ptr = self.impl.get_mem_pool_base_address( + self._layer_attn_to_layer_id[pp_layer, DeepseekV4AttentionType.SWA], Role.KEY + ) + return (buffer_ptr - swa_pool_ptr) // swa_bytes_per_block + + # Tensors for compatibility with AttentionOp, only contains swa attention. + # Assume the SWA of all layers share the same pool. + # shape: [1, 2] + kv_cache_pool_pointers = torch.tensor( + [[swa_pool_ptr, 0]], + dtype=torch.int64, + device="cpu", + pin_memory=prefer_pinned(), + ) + # shape: [num_local_layers, 2] + kv_cache_pool_mapping = torch.tensor( + [[0, _get_layer_offset(pp_layer)] for pp_layer in self.pp_layers], + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + return kv_cache_pool_pointers, kv_cache_pool_mapping + + def get_cache_indices( + self, + request_id: int, + layer_idx: int, + attn_type: DeepseekV4AttentionType, + ) -> List[int]: + """ + Get the cache block indices for a batch of requests at a specific layer and attention type. + + Args: + request_id: The request id + layer_idx: The layer index + attn_type: The attention type + + Returns: + The cache block indices, shape (max_blocks_per_seq,) + """ + layer_id = self._layer_attn_to_layer_id[(layer_idx, attn_type)] + pool_id = self.layer_to_pool_mapping_dict[layer_id] + base_indices = self.kv_cache_map[request_id].get_base_page_indices(pool_id).tolist() + converter = self.impl.get_page_index_converter(layer_id, Role.KEY) + return converter(base_indices) + + def _get_cache_quota(self, max_tokens: int) -> int: + quota = int(max_tokens * self.get_cache_bytes_per_token()) + # Add extra quota to ensure sufficient space for small max_tokens cases. + quota += len(DeepseekV4AttentionType) * (2 << 20) + return quota + + def _build_cache_config( + self, + kv_cache_config: KvCacheConfig, + *, + tokens_per_block: int, + vocab_size: int | None, + cache_tiers: List[GpuCacheTierConfig | HostCacheTierConfig], + ) -> KVCacheManagerConfigPy: + """ + Create the cache manager config for DeepSeek-V4. + """ + layers: List[AttentionLayerConfig] = [] + layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} + + def _add_layer( + layer_idx: int, attn_type: DeepseekV4AttentionType, sliding_window_size: int | None + ): + nonlocal layers, layer_attn_to_layer_id + layer_id = LayerId(len(layers)) + # update the mapping from layer index and attention type to layer id + layer_attn_to_layer_id[layer_idx, attn_type] = layer_id + # add the layer to the layers list + layer_config = AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig( + role=Role.KEY, size=self._get_attn_bytes_per_block(attn_type, layer_idx) + ) + ], + sliding_window_size=sliding_window_size, + num_sink_tokens=None, + ) + layers.append(layer_config) + + # create the layer config for DeepSeek-V4 + for layer in self.pp_layers: + compress_ratio = self._compress_ratios[layer] + is_compress = is_compress_layer(compress_ratio) + is_sparse = is_sparse_layer(compress_ratio) + + # sliding window attention pool + _add_layer( + layer, + DeepseekV4AttentionType.SWA, + self._get_window_size(compress_ratio, DeepseekV4AttentionType.SWA), + ) + + if is_compress: + # compressed attention pool + _add_layer(layer, DeepseekV4AttentionType.COMPRESS, None) + # compressor state, managed as a sliding window attention cache, + # including compressor kv states and compressor score states. + # Add max_draft_len so rewind after rejected draft tokens can + # still reach past states. + compressor_window = self._get_window_size( + compress_ratio, DeepseekV4AttentionType.COMPRESSOR_STATE + ) + _add_layer(layer, DeepseekV4AttentionType.COMPRESSOR_STATE, compressor_window) + _add_layer(layer, DeepseekV4AttentionType.COMPRESSOR_SCORE, compressor_window) + + # sparse attention layer has indexer + if is_sparse: + # indexer kv cache pool, dim is indexer_head_dim + _add_layer(layer, DeepseekV4AttentionType.INDEXER_COMPRESS, None) + # indexer has its own compressor, so a separate compressor state + # similarly, indexer compressor state is managed as a sliding window attention cache + indexer_compressor_window = self._get_window_size( + compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE + ) + _add_layer( + layer, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + indexer_compressor_window, + ) + _add_layer( + layer, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + indexer_compressor_window, + ) + # the mapping from layer index and attention type to layer id + self._layer_attn_to_layer_id = layer_attn_to_layer_id + # number of layers in the KVCacheManagerPy + self._num_manager_layers = len(layers) + + # Build constraints and typical_step for better pool ratio. + max_batch_size = self.max_batch_size + max_seq_len = self.max_seq_len + max_input_len = self._max_input_len + max_num_tokens = self._max_num_tokens + max_draft_len = self._max_draft_len + + # For aggregated serving in large batch size: + # Use 1 context request + (max_batch_size - 1) generation requests as + # the typical step. An all-generation typical_step over-provisions the + # compressed-cache pool at the expense of the SWA pool, starving the + # SWA pool and artificially capping the achievable batch size. + typical_step = BatchDesc( + kv_caches=[ + KVCacheDesc(capacity=max_seq_len, history_length=0), + ] + + [KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - max_draft_len - 1)] + * (max_batch_size - 1), + ) + + constraints = [] + # Constraint 1: one context request at max_seq_len, this case is used in warmup stage. + # max_draft_len is already included in max_seq_len, so no need to add it again. + constraints.append(BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=0)])) + + # Constraint 2: when user given max_input_len and max_num_tokens, we can build constraints for context requests. + if max_input_len is not None and max_num_tokens is not None: + # There are at most max_batch_size generation requests, and each generation request consumes + # (max_draft_len + 1) tokens, so the remaining tokens are for context requests. + max_context_tokens = max_num_tokens - max_batch_size * (max_draft_len + 1) + if max_context_tokens > 0: + max_num_context = min(max_context_tokens // max_input_len, max_batch_size) + constraints.append( + BatchDesc( + [KVCacheDesc(capacity=max_input_len + max_draft_len + 1, history_length=0)] + * max_num_context + ) + ) + + return KVCacheManagerConfigPy( + tokens_per_block=tokens_per_block, + vocab_size=vocab_size, + cache_tiers=cache_tiers, + max_util_for_resume=kv_cache_config.max_util_for_resume, + layers=layers, + typical_step=typical_step, + constraints=constraints, + ) + + def _assert_layer_pool_scale(self) -> None: + attn_ratio_to_pool_id = defaultdict[DeepseekV4AttentionType, dict[int, int]](lambda: {}) + attn_ratio_to_scale = defaultdict[DeepseekV4AttentionType, dict[int, int]](lambda: {}) + + comb = [ + (attn_type, layer_idx) + for attn_type in DeepseekV4AttentionType + for layer_idx in self.pp_layers + if compress_ratio_has_attention(self._compress_ratios[layer_idx], attn_type) + ] + for attn_type, layer_idx in comb: + compress_ratio = self._compress_ratios[layer_idx] + layer_id = self._layer_attn_to_layer_id[layer_idx, attn_type] + pool_id = self.layer_to_pool_mapping_dict[layer_id] + converter = self.impl.get_page_index_converter(layer_id, Role.KEY) + scale = converter.scale + + # check if the pool id is consistent + if compress_ratio in attn_ratio_to_pool_id[attn_type]: + other_pool_id = attn_ratio_to_pool_id[attn_type][compress_ratio] + assert other_pool_id == pool_id, ( + f"Layer {layer_idx} with compress ratio {compress_ratio}, " + f"its attention type {attn_type.name} has pool id {pool_id}, " + f"but another layer with the same compress ratio and attention type has pool id {other_pool_id}." + "DeepSeek-V4 expects they share the same pool." + ) + else: + attn_ratio_to_pool_id[attn_type][compress_ratio] = pool_id + + # check if the scale is consistent + if compress_ratio in attn_ratio_to_scale[attn_type]: + other_scale = attn_ratio_to_scale[attn_type][compress_ratio] + assert other_scale == scale, ( + f"Layer {layer_idx} with compress ratio {compress_ratio}, " + f"its attention type {attn_type.name} has scale {scale}, " + f"but another layer with the same compress ratio and attention type has scale {other_scale}." + "DeepSeek-V4 expects they share the same scale." + ) + else: + attn_ratio_to_scale[attn_type][compress_ratio] = scale + + # check if all swa attentions are in the same pool and have the same scale + swa_pool_ids = set(attn_ratio_to_pool_id[DeepseekV4AttentionType.SWA].values()) + swa_scales = set(attn_ratio_to_scale[DeepseekV4AttentionType.SWA].values()) + assert len(swa_pool_ids) == 1, "All swa attentions must be in the same pool" + assert len(swa_scales) == 1, "All swa attentions must have the same scale" + + # Ensure all compress ratios have SWA entries, not just PP-local ones. + # The attention metadata uses compress_ratio=1 as a hardcoded SWA key, + # but with pipeline parallelism a PP stage may not have any layers with + # ratio 1. Since all SWA layers share the same pool, we populate entries + # for every compress ratio in the model. + swa_pool_id = next(iter(swa_pool_ids)) + swa_scale = next(iter(swa_scales)) + for ratio in set(self._compress_ratios): + if ratio not in attn_ratio_to_pool_id[DeepseekV4AttentionType.SWA]: + attn_ratio_to_pool_id[DeepseekV4AttentionType.SWA][ratio] = swa_pool_id + attn_ratio_to_scale[DeepseekV4AttentionType.SWA][ratio] = swa_scale + + self._attn_ratio_to_pool_id = attn_ratio_to_pool_id + self._attn_ratio_to_scale = attn_ratio_to_scale + + def _get_attn_bytes_per_block( + self, + attn_type: DeepseekV4AttentionType, + layer_idx: int, + ) -> int: + """ + Get the cache bytes per token for a specific attention type and layer. + """ + has_fp8_kv_cache = self.dtype == DataType.FP8 + token_bytes = get_token_bytes( + self.head_dim, + self.index_head_dim, + self._compress_ratios[layer_idx], + attn_type, + has_fp8_kv_cache, + ) + + block_size = self.tokens_per_block + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + block_size = self.compressed_block_sizes[layer_idx] + + return token_bytes * block_size + + def get_cache_bytes_per_token(self) -> int: + """Get the average cache bytes per token for DeepSeek-V4.""" + has_fp8_kv_cache = self.dtype == DataType.FP8 + compress_ratios = [self._compress_ratios[layer] for layer in self.pp_layers] + return _estimate_bytes_per_token( + self.head_dim, + self.index_head_dim, + compress_ratios, + has_fp8_kv_cache, + ) + + def get_max_resource_count(self) -> int: + # Keep scheduler capacity tied to physical GPU KV quota in bytes. + return int(self.impl.get_quota(GPU_LEVEL)) + + def _is_context_request(self, request: llm_request.LlmRequest) -> bool: + if getattr(request, "is_context_init_state", False): + return True + return getattr(request, "state", None) == llm_request.LlmRequestState.CONTEXT_INIT + + def _is_generation_request(self, request: llm_request.LlmRequest) -> bool: + if ( + getattr(request, "is_generation_in_progress_state", False) + or getattr(request, "is_generation_to_complete_state", False) + or getattr(request, "is_disagg_generation_init_state", False) + ): + return True + return getattr(request, "state", None) in ( + llm_request.LlmRequestState.GENERATION_IN_PROGRESS, + llm_request.LlmRequestState.GENERATION_TO_COMPLETE, + ) + + def _get_context_bytes(self, request: llm_request.LlmRequest) -> int: + prompt_len = max(0, getattr(request, "prompt_len", request.orig_prompt_len)) + total_tokens = prompt_len + self.num_extra_kv_tokens + return total_tokens * self.get_cache_bytes_per_token() + + def _get_generation_bytes(self, request: llm_request.LlmRequest) -> int: + prompt_len = max(0, getattr(request, "prompt_len", request.orig_prompt_len)) + max_new_tokens = max(0, request.max_new_tokens) + total_tokens = prompt_len + max_new_tokens + self.num_extra_kv_tokens + has_fp8_kv_cache = self.dtype == DataType.FP8 + total_bytes = 0 + for layer in self.pp_layers: + compress_ratio = self._compress_ratios[layer] + for attn_type in DeepseekV4AttentionType: + if not compress_ratio_has_attention(compress_ratio, attn_type): + continue + token_bytes = _get_attn_bytes_per_token( + self.head_dim, + self.index_head_dim, + compress_ratio, + attn_type, + has_fp8_kv_cache, + ) + attn_tokens = total_tokens + if attn_type in self.fixed_size_attention: + attn_tokens = self._get_window_size(compress_ratio, attn_type) + total_bytes += attn_tokens * token_bytes + return total_bytes + + def get_needed_resource_to_completion(self, request: llm_request.LlmRequest) -> int: + if self._is_generation_request(request): + return self._get_generation_bytes(request) + if self._is_context_request(request): + return self._get_context_bytes(request) + raise ValueError(f"Unsupported request state: {request.state}") + + def get_layer_bytes_per_token( + self, + local_layer_idx: int, + data_role: Role, + ) -> int: + raise NotImplementedError( + "DeepSeek-V4 doesn't support get_layer_bytes_per_token, use _get_attn_bytes_per_block" + ) + + def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor: + """ + Get the buffers for the indexer k cache for a specific layer. + """ + buffer = self.get_buffers(layer_idx, DeepseekV4AttentionType.INDEXER_COMPRESS).unsqueeze(2) + return buffer.view(torch.uint8) + + def get_batch_indexer_k_cache_indices(self, request_ids: List[int]) -> List[List[int]]: + """ + Get the indices for the indexer k cache for a batch of requests. + """ + return self.get_batch_attn_offset( + request_ids, + # use beam_width=1 and num_contexts=0 since we don't support beam search + 1, + 0, + len(request_ids), + DeepseekV4AttentionType.INDEXER_COMPRESS, + DEEPSEEK_V4_SPARSE_RATIO, + ).tolist() + + def copy_batch_block_offsets( + self, + dst_tensor: torch.Tensor, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + ) -> None: + """For compatibility with AttentionOp, copy only the SWA block offsets.""" + offsets = self.get_batch_attn_offset( + request_ids, + beam_width, + num_contexts, + num_seqs, + DeepseekV4AttentionType.SWA, + # all compress ratios have SWA attention and they are in the same pool + self._compress_ratios[self.pp_layers[0]], + ) + self._host_block_offsets_staging[:num_seqs, :, :] = offsets[:, None, :] + dst_tensor[0, :num_seqs, :, :].copy_( + self._host_block_offsets_staging[:num_seqs, :, :], non_blocking=True + ) + + def get_batch_attn_offset( + self, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + attn_type: DeepseekV4AttentionType, + compress_ratio: int, + ) -> torch.Tensor: + """ + Get the block offsets for a specific attention type for a batch of requests. + + Args: + request_ids: The request ids + beam_width: The beam width + num_contexts: The number of context requests + num_seqs: The number of sequence requests + attn_type: The attention type + compress_ratio: The compress ratio. Used for non-SWA attention types. + + Returns: + The block offsets, shape (num_seqs, max_blocks_per_seq) + """ + assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" + assert attn_type == DeepseekV4AttentionType.SWA or compress_ratio is not None, ( + "compress_ratio must be provided for non-SWA attention types" + ) + + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + assert copy_idx.shape[0] == num_seqs + + pool_id = self._attn_ratio_to_pool_id[attn_type][compress_ratio] + scale = self._attn_ratio_to_scale[attn_type][compress_ratio] + offsets = self.host_kv_cache_block_offsets[pool_id, copy_idx, 0] * scale + offsets[offsets == -scale] = -1 + return offsets + + def get_batch_block_offsets( + self, + request_ids: List[int], + num_contexts: int, + attention_type_set: set, + ) -> Dict[Tuple[int, "DeepseekV4AttentionType"], torch.Tensor]: + """Get block offsets for all attention types in a single call. + + Calls get_copy_index once and deduplicates offset computation by + (pool_id, scale) to avoid redundant work. + + Args: + request_ids: The request ids. + num_contexts: The number of context requests. + attention_type_set: Set of (compress_ratio, attention_type) tuples. + + Returns: + Dict mapping (compress_ratio, attention_type) -> offset tensor. + """ + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, 1) + + offset_cache = {} # (pool_id, scale) -> offsets tensor + result = {} + for compress_ratio, attention_type in attention_type_set: + pool_id = self._attn_ratio_to_pool_id[attention_type][compress_ratio] + scale = self._attn_ratio_to_scale[attention_type][compress_ratio] + cache_key = (pool_id, scale) + if cache_key not in offset_cache: + offsets = self.host_kv_cache_block_offsets[pool_id, copy_idx, 0] * scale + offsets[offsets == -scale] = -1 + offset_cache[cache_key] = offsets + result[(compress_ratio, attention_type)] = offset_cache[cache_key] + return result + + @staticmethod + def get_cache_size_per_token(model_config: ModelConfig, mapping: Mapping, **kwargs): + config = model_config.pretrained_config + head_dim = config.kv_lora_rank + config.qk_rope_head_dim + index_head_dim = model_config.sparse_attention_config.index_head_dim + is_disagg = kwargs.get("is_disagg", False) + pp_layers = mapping.pp_layers(model_config.get_num_attention_layers(is_disagg=is_disagg)) + compress_ratios = [ + model_config.sparse_attention_config.compress_ratios[layer] for layer in pp_layers + ] + quant_config = model_config.quant_config + if quant_config is not None: + has_fp8_kv_cache = quant_config.quant_mode.has_fp8_kv_cache() + else: + has_fp8_kv_cache = False + return _estimate_bytes_per_token( + head_dim, + index_head_dim, + compress_ratios, + has_fp8_kv_cache, + ) + + def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: + some_checks_unavailable = False + has_invalid_values = torch.tensor( + [False], dtype=torch.bool, device=torch.cuda.current_device() + ) + pool_handled = set() + + # Handle each layer from start to end to traverse the whole KV cache. + for (layer, attn), layer_id in self._layer_attn_to_layer_id.items(): + pool_id = self.layer_to_pool_mapping_dict[layer_id] + if pool_id in pool_handled: + continue + buffer = self.get_buffers(layer, attn) + # process in chunks of 256 pages to avoid OoM + for i in range(0, buffer.shape[0], 256): + buffer_slice = buffer[i : i + 256] + try: + has_invalid_values.logical_or_(torch.isnan(buffer_slice).any()) + has_invalid_values.logical_or_(torch.isinf(buffer_slice).any()) + except NotImplementedError: + some_checks_unavailable = True + if fill_with_zero: + buffer.zero_() + pool_handled.add(pool_id) + torch.cuda.synchronize() + + if some_checks_unavailable: + logger.warning( + "`torch.isnan` or `torch.isinf` is not implemented for current kv cache dtype, " + "related checks are skipped" + ) + return bool(has_invalid_values) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py new file mode 100644 index 000000000000..d8c3d2620413 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -0,0 +1,334 @@ +import os +from contextlib import contextmanager +from enum import IntEnum +from typing import TYPE_CHECKING, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from tensorrt_llm._torch.attention_backend.interface import MLAParams, PositionalEmbeddingParams +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._torch.utils import maybe_compile + +if TYPE_CHECKING: + from .deepseek_v4 import DeepseekV4TrtllmAttentionMetadata + +# When set to "1", forces wkv_gate to use full FP32 computation (via nn.Linear) +# instead of the default TF32 path (via F.linear + allow_tf32). +_USE_FP32_COMPRESSOR = os.environ.get("DEEPSEEK_V4_COMPRESSOR_FP32", "0") == "1" + + +@maybe_compile(dynamic=True) +def _to_float(x: torch.Tensor) -> torch.Tensor: + """Cast to float32 for TF32 GEMM (following DSA pattern).""" + return x.float() + + +@contextmanager +def _tf32_matmul_enabled(): + """Temporarily enable TF32 for FP32 matmul in this scope. + + Forces PyTorch/cuBLASLt to use CUBLAS_COMPUTE_32F_FAST_TF32 which + guarantees TF32 tensor cores. Plain CUBLAS_COMPUTE_32F (used by + torch.ops.trtllm.cublas_mm) falls back to SIMT SGEMM based on + cuBLASLt heuristics for small M. + """ + prev = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = prev + + +class KVCacheDtype(IntEnum): + """KV cache dtype/layout preset (values match C++ cache_scale_type parameter). + + The store dtype and scale layout are implied by this value: + - NONE: keeps the input dtype (bf16/fp32, decided by the + caller's tensor element size). + - FP8_PERTENSOR: 1 byte per value (FP8 E4M3) with implicit scale=1. + - FP8_BLOCKWISE: 1 byte per value + 1 fp32 scale per 128 values. + - MXFP4_BLOCKWISE: packed FP4 (½ byte per value) + 1 UE8M0 byte per + 32 values. + + Storage size in bytes per logical element is therefore:: + + size_per_value = { + NONE: elem_bytes, # caller-side + FP8_PERTENSOR: 1, + FP8_BLOCKWISE: 1 + 4 / 128, # data + fp32 scale + MXFP4_BLOCKWISE: 0.5 + 1 / 32, # nibble + ue8m0 byte + }[kv_cache_dtype] + """ + + NONE = 0 + FP8_PERTENSOR = 1 # FP8 E4M3 with implicit scale=1 + FP8_BLOCKWISE = 2 # FP8 E4M3 with per-128 fp32 scales + MXFP4_BLOCKWISE = 3 # packed FP4 E2M1 with per-32 UE8M0 scales + + +_KV_CACHE_DTYPE_MAP = { + "default": KVCacheDtype.NONE, + "bf16": KVCacheDtype.NONE, + "fp8_pertensor": KVCacheDtype.FP8_PERTENSOR, + "fp8_blockwise": KVCacheDtype.FP8_BLOCKWISE, + "mxfp4": KVCacheDtype.MXFP4_BLOCKWISE, +} + + +def resolve_kv_cache_dtype(kv_cache_dtype: Union[str, KVCacheDtype]) -> KVCacheDtype: + if isinstance(kv_cache_dtype, str): + return _KV_CACHE_DTYPE_MAP[kv_cache_dtype] + return kv_cache_dtype + + +class Compressor(nn.Module): + """KV compressor using Triton kernels with paged memory management. + + Args: + mla_params: MLA parameters containing hidden_size and head dimensions + layer_idx: Layer index for cache management + compress_ratio: Compression ratio (e.g., 4 compresses 4 tokens into 1) + norm_eps: RMSNorm epsilon + skip_create_weights_in_init: Whether to skip weight initialization + pos_embd_params: Positional embedding parameters for RoPE + dtype: Data type for computation + kv_cache_dtype: Cache preset string or KVCacheDtype. + rotate_activation: Whether to apply Hadamard transform in postprocessing (False to skip) + """ + + def __init__( + self, + mla_params: MLAParams, + layer_idx: int, + compress_ratio: int, + norm_eps: float, + skip_create_weights_in_init: bool, + pos_embd_params: PositionalEmbeddingParams, + dtype: Optional[torch.dtype] = torch.bfloat16, + kv_cache_dtype: Union[str, KVCacheDtype] = KVCacheDtype.NONE, + is_indexer: bool = False, + rotate_activation: bool = False, + ): + super().__init__() + # Dimensions + self.dim = mla_params.hidden_size + self.head_dim = mla_params.qk_rope_head_dim + mla_params.qk_nope_head_dim + self.rope_head_dim = mla_params.qk_rope_head_dim + self.nope_head_dim = mla_params.qk_nope_head_dim + + # Compression config + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.state_dim = 2 * self.head_dim if self.overlap else self.head_dim + + # Cache config + self.layer_idx = layer_idx + self.kv_cache_dtype: KVCacheDtype = resolve_kv_cache_dtype(kv_cache_dtype) + self.is_indexer = is_indexer + self.rotate_activation = rotate_activation + + # Modules + self.wkv_gate = Linear( + self.dim, + self.state_dim * 2, + bias=False, + dtype=torch.float32, + quant_config=None, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=True, + ) + self.norm = RMSNorm(hidden_size=self.head_dim, eps=norm_eps, dtype=dtype) + self.rotary_emb = RotaryEmbedding( + pos_embd_params.rope, + head_dim=self.rope_head_dim, + is_neox=pos_embd_params.is_neox, + ) + + # Learnable absolute positional encoding for compression + self.ape = nn.Parameter(torch.empty(compress_ratio, self.state_dim, dtype=torch.float32)) + + def forward( + self, + x: torch.Tensor, + metadata: "DeepseekV4TrtllmAttentionMetadata", + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Forward pass for paged KV compression. + + Args: + x: Input tensor [num_tokens, dim] + metadata: Attention metadata with cache info + + Returns: + (kv_data, scale) tuple: + - default / fp8_pertensor main compressor: (kv_comp, None) + - default indexer: (kv_out, None) bf16 + - fp8_blockwise indexer: (fp8_output, fp32 scale) + - mxfp4 indexer: (fp4_output, ue8m0 scale) + - no compressed tokens: (None, None) + """ + # Import at runtime to avoid circular dependency + from .deepseek_v4 import DeepseekV4AttentionType + + # Extract metadata + num_contexts = metadata.num_contexts + num_generations = metadata.num_generations + num_ctx_tokens = metadata.num_ctx_tokens + bsz = num_contexts + num_generations + + # Determine attention types based on whether this is an indexer compressor + if self.is_indexer: + compress_type = DeepseekV4AttentionType.INDEXER_COMPRESS + state_type = DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE + score_type = DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE + else: + compress_type = DeepseekV4AttentionType.COMPRESS + state_type = DeepseekV4AttentionType.COMPRESSOR_STATE + score_type = DeepseekV4AttentionType.COMPRESSOR_SCORE + + # Get cache buffers + kv_cache = metadata.kv_cache_manager.get_buffers(self.layer_idx, compress_type) + paged_kv_state = metadata.kv_cache_manager.get_buffers(self.layer_idx, state_type) + paged_score_state = metadata.kv_cache_manager.get_buffers(self.layer_idx, score_type) + + # Get block tables + block_table = metadata.block_tables[(self.compress_ratio, compress_type)] + block_table_kv_state = metadata.block_tables[(self.compress_ratio, state_type)] + block_table_score_state = metadata.block_tables[(self.compress_ratio, score_type)] + + # Get tokens_per_block from cache manager + # state_tokens_per_block: for state/score caches (used in compress kernels) + # compress_tokens_per_block: for compressed KV cache (used in scatter) + state_tokens_per_block = metadata.kv_cache_manager.tokens_per_block + compress_tokens_per_block = metadata.kv_cache_manager.compressed_block_sizes[self.layer_idx] + + # Get compression metadata + cu_new_comp_kv = metadata.cu_new_comp_kv_cuda[self.compress_ratio] + kv_lens = metadata.kv_lens_cuda_runtime + total_num_comp_tokens = metadata.num_total_compressed_tokens[self.compress_ratio] + num_comp_tokens = metadata.new_comp_kv_lens_cuda[self.compress_ratio][:bsz] + max_ctx_comp_kv_lens = metadata.max_ctx_compressed_tokens[self.compress_ratio] + + # Project input to KV and score. + # Default: TF32 via F.linear under allow_tf32 context (explicit + # CUBLAS_COMPUTE_32F_FAST_TF32 -> TF32 tensor cores on Ampere+). + # Fallback: strict FP32 via nn.Linear when DEEPSEEK_V4_COMPRESSOR_FP32=1. + if _USE_FP32_COMPRESSOR: + kv_score = self.wkv_gate(_to_float(x)) + else: + with _tf32_matmul_enabled(): + kv_score = F.linear(_to_float(x), self.wkv_gate.weight) + + # Allocate output buffer + kv_comp = torch.empty(total_num_comp_tokens, self.head_dim, device=x.device, dtype=x.dtype) + + # Run compression kernels + if num_contexts > 0: + torch.ops.trtllm.compressor_prefill_reduction( + kv_score[:num_ctx_tokens], + self.ape, + paged_kv_state, + paged_score_state, + block_table_kv_state[:num_contexts], + block_table_score_state[:num_contexts], + kv_comp, + kv_lens[:num_contexts], + metadata.cached_token_lens_cuda[:num_contexts], + metadata.cu_seq_lens_cuda, + cu_new_comp_kv[: num_contexts + 1], + num_contexts, + state_tokens_per_block, + self.head_dim, + self.compress_ratio, + max_ctx_comp_kv_lens, + ) + + if num_generations > 0: + gen_kv_lens = kv_lens[num_contexts:] + next_n = metadata.num_gen_tokens_per_seq + # Pass full kv_score (not sliced) with the generation portion of + # cu_seq_lens so the kernel reads at the correct absolute offsets. + torch.ops.trtllm.compressor_paged_kv_compress( + kv_score, + self.ape, + paged_kv_state, + paged_score_state, + block_table_kv_state[num_contexts:], + block_table_score_state[num_contexts:], + kv_comp, + gen_kv_lens, + gen_kv_lens - next_n, + metadata.cu_seq_lens_cuda[num_contexts:], + cu_new_comp_kv[num_contexts:], + num_generations, + state_tokens_per_block, + self.head_dim, + self.compress_ratio, + next_n, + ) + + # Scatter to cache with appropriate quantization (all modes fused) + start_pos = metadata.past_kv_lens_cuda[self.compress_ratio][:bsz] + total_tokens = kv_comp.shape[0] + + # Allocate optional returned postprocess buffers for indexer paths. + kv_out = None + quant_output = None + scale_output = None + if self.is_indexer: + if self.kv_cache_dtype == KVCacheDtype.NONE: + kv_out = torch.empty_like(kv_comp) + elif self.kv_cache_dtype == KVCacheDtype.FP8_BLOCKWISE: + num_scale_blocks = self.head_dim // 128 + quant_output = torch.empty( + total_tokens, self.head_dim, dtype=torch.uint8, device=kv_comp.device + ) + scale_output = torch.empty( + total_tokens, num_scale_blocks, dtype=torch.float32, device=kv_comp.device + ) + elif self.kv_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: + num_scale_blocks = self.head_dim // 32 + quant_output = torch.empty( + total_tokens, self.head_dim // 2, dtype=torch.uint8, device=kv_comp.device + ) + scale_output = torch.empty( + total_tokens, num_scale_blocks, dtype=torch.uint8, device=kv_comp.device + ) + + position_ids = metadata.compressed_position_ids_cuda[self.compress_ratio][:total_tokens] + compressed_mask = metadata.compressed_mask_cuda[self.compress_ratio][:total_tokens] + + # Fused postprocess + scatter: RMSNorm + RoPE + Hadamard + paged cache write + torch.ops.trtllm.compressor_postprocess_scatter( + kv_comp, + kv_out, + self.norm.weight, + self.norm.variance_epsilon, + self.rotary_emb.rotary_cos_sin, + position_ids, + self.nope_head_dim, + self.rope_head_dim, + kv_cache, + num_comp_tokens, + cu_new_comp_kv, + start_pos, + block_table, + compressed_mask, + compress_tokens_per_block, + int(self.kv_cache_dtype), + self.rotate_activation, + quant_output, + scale_output, + ) + + if quant_output is not None: + if self.kv_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: + return quant_output.view(torch.float4_e2m1fn_x2), scale_output + return quant_output.view(torch.float8_e4m3fn), scale_output + if kv_out is not None: + return kv_out, None + return kv_comp, None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py new file mode 100644 index 000000000000..bf187391dcb0 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -0,0 +1,1209 @@ +import math +from enum import Enum +from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple + +import torch + +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionForwardArgs, + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.modules.linear import Linear # noqa: E402 (avoid cycle) +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding +from tensorrt_llm._torch.utils import maybe_compile +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.utils import fp8_utils + +from ..dsa import DSAtrtllmAttentionMetadata, Indexer, rotate_activation +from ..kernel import deepseek_v4_local_to_global_indices +from .compressor import Compressor, KVCacheDtype, resolve_kv_cache_dtype + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig + +DEEPSEEK_V4_SPARSE_RATIO = 4 +DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO = 4 + + +class DeepseekV4AttentionType(Enum): + SWA = 0 + COMPRESS = 1 + COMPRESSOR_STATE = 2 + COMPRESSOR_SCORE = 3 + INDEXER_COMPRESS = 4 + INDEXER_COMPRESSOR_STATE = 5 + INDEXER_COMPRESSOR_SCORE = 6 + + +def is_overlap_compressor(compress_ratio: int) -> bool: + """ + Check if the compressor of the given layer is working in the overlap mode. + """ + return compress_ratio == DEEPSEEK_V4_OVERLAP_COMPRESSOR_RATIO + + +def is_sparse_layer(compress_ratio: int) -> bool: + """ + Check if the given layer is a sparse layer. + """ + return compress_ratio == DEEPSEEK_V4_SPARSE_RATIO + + +def is_compress_layer(compress_ratio: int) -> bool: + """ + Check if the given layer is a compress layer. + """ + return compress_ratio > 1 + + +def compress_ratio_has_attention(compress_ratio: int, attn_type: DeepseekV4AttentionType) -> bool: + """ + Check if the given compress ratio has the given attention type. + """ + is_sparse = is_sparse_layer(compress_ratio) + is_compress = is_compress_layer(compress_ratio) + + if attn_type == DeepseekV4AttentionType.SWA: + return True + elif attn_type == DeepseekV4AttentionType.COMPRESS: + return is_compress + elif attn_type == DeepseekV4AttentionType.COMPRESSOR_STATE: + return is_compress + elif attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: + return is_compress + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + return is_sparse + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE: + return is_sparse + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: + return is_sparse + + +def get_attn_dim( + head_dim: int, index_head_dim: int, compress_ratio: int, attn_type: DeepseekV4AttentionType +) -> int: + """ + Get the dimension of the attention type for a specific layer. + """ + state_factor = 2 if is_overlap_compressor(compress_ratio) else 1 + if attn_type == DeepseekV4AttentionType.SWA: + return head_dim + elif attn_type == DeepseekV4AttentionType.COMPRESS: + return head_dim + elif attn_type == DeepseekV4AttentionType.COMPRESSOR_STATE: + return state_factor * head_dim + elif attn_type == DeepseekV4AttentionType.COMPRESSOR_SCORE: + return state_factor * head_dim + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + return index_head_dim + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE: + return state_factor * index_head_dim + elif attn_type == DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE: + return state_factor * index_head_dim + + +def get_token_bytes( + head_dim: int, + index_head_dim: int, + compress_ratio: int, + attn_type: DeepseekV4AttentionType, + has_fp8_kv_cache: bool, + indexer_k_cache_dtype: str = "fp8_blockwise", +) -> int: + """ + Get the token bytes for a specific layer and attention type. + + Args: + head_dim: The head dimension + index_head_dim: The index head dimension + compress_ratio: The compress ratio + attn_type: The attention type + has_fp8_kv_cache: Whether the KV cache uses FP8 quantization + indexer_k_cache_dtype: Indexer compressor cache preset string + (controls INDEXER_COMPRESS dtype + scale layout independently + from has_fp8_kv_cache, which controls the main attention path). + + Returns: + The number of bytes per token, including scaling factor + """ + if not compress_ratio_has_attention(compress_ratio, attn_type): + raise ValueError( + f"Layer with compress ratio {compress_ratio} does not have attention type {attn_type}" + ) + + attn_dim = get_attn_dim(head_dim, index_head_dim, compress_ratio, attn_type) + + # Default dtype is bfloat16 (2 bytes), or fp8 (1 byte) when FP8 kv cache is enabled + dtype_bytes = 1 if has_fp8_kv_cache else 2 + # (indexer) compressor state and score always use float32 + if attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + dtype_bytes = 4 # (indexer) compressor state and score use float32 + # Indexer cache always packs data + per-block scales into one row. Only + # FP8 blockwise and MXFP4 are valid here -- bf16 / fp8_pertensor + # are reserved for the main-attention compressor. + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + indexer_cache_dtype = resolve_kv_cache_dtype(indexer_k_cache_dtype) + if indexer_cache_dtype == KVCacheDtype.FP8_BLOCKWISE: + # 1 byte per value + 1 fp32 scale per 128 values. + return attn_dim + index_head_dim // 128 * 4 + if indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE: + # ½ byte per value + 1 ue8m0 byte per 32 values. + return index_head_dim // 2 + index_head_dim // 32 + raise ValueError( + f"Unsupported indexer_k_cache_dtype {indexer_k_cache_dtype!r}; " + "expected 'fp8_blockwise' or 'mxfp4'." + ) + + return attn_dim * dtype_bytes + + +class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): + # The set of compress ratios for the layers + compress_ratio_set: Set[int] + # The set of (compress ratio, attention type) for the layers + attention_type_set: Set[Tuple[int, DeepseekV4AttentionType]] + # The number of total compressed tokens for each compress ratio + num_total_compressed_tokens: Dict[int, int] = {} + # The max number of context compressed tokens for each compress ratio + max_ctx_compressed_tokens: Dict[int, int] = {} + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def __post_init__(self): + super().__post_init__() + assert self.sparse_attention_config.window_size == 128, ( + f"Dual-pool sparse MLA requires window_size == 128, which equals to the" + f"TileSizeKV of the FMHA kernel. (got {self.sparse_attention_config.window_size})." + ) + capture_graph = self.is_cuda_graph + self.compress_ratio_set = set(self.compress_ratios) + # Cache a sorted list for deterministic iteration order across + # torch.compile traces. Avoids repeated set→list conversion and + # guarantees stable specialization keys. + self._compress_ratios_sorted = sorted(self.compress_ratio_set) + _supported_ratios = {1, 4, 128} + _unsupported = self.compress_ratio_set - _supported_ratios + assert not _unsupported, ( + f"Unsupported compress ratios {_unsupported}. Only {_supported_ratios} are supported." + ) + + attention_types = [] + for compress_ratio in self.compress_ratio_set: + if compress_ratio == 1: + attention_types.append((1, DeepseekV4AttentionType.SWA)) + elif compress_ratio == 4: + attention_types.append((1, DeepseekV4AttentionType.SWA)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESS)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_STATE)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_SCORE)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESS)) + attention_types.append( + (compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE) + ) + attention_types.append( + (compress_ratio, DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE) + ) + else: + attention_types.append((1, DeepseekV4AttentionType.SWA)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESS)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_STATE)) + attention_types.append((compress_ratio, DeepseekV4AttentionType.COMPRESSOR_SCORE)) + self.attention_type_set = set(attention_types) + + # Create buffers for the compressor + # cu_seq_lens_cuda is the cumulative sequence lengths for the requests + self.cu_seq_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences + 1,), + dtype=torch.int, + cache_name="cu_seq_lens_cuda", + capture_graph=capture_graph, + ) + self.cu_seq_lens = torch.empty_like( + self.cu_seq_lens_cuda, device="cpu", pin_memory=prefer_pinned() + ) + self.cu_seq_lens[0] = 0 + + # new_comp_kv_lens_cuda is the number of new compressed tokens for the requests + self.new_comp_kv_lens_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + dtype=torch.int, + cache_name=f"new_comp_kv_lens_cuda_{compress_ratio}", + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + + # cu_new_comp_kv_cuda is the cumulative number of new compressed tokens for the requests + self.cu_new_comp_kv_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences + 1,), + dtype=torch.int, + cache_name=f"cu_new_comp_kv_cuda_{compress_ratio}", + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + + # compressed_kv_lens_cuda is the number of compressed tokens for the requests + self.compressed_kv_lens_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + dtype=torch.int, + cache_name=f"compressed_kv_lens_cuda_{compress_ratio}", + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + + # past_kv_lens_cuda is the number of past compressed tokens for the requests + self.past_kv_lens_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + dtype=torch.int, + cache_name=f"past_kv_lens_cuda_{compress_ratio}", + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + + # compressed_position_ids_cuda is the compressed position ids for the requests + self.compressed_position_ids_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + dtype=torch.int, + cache_name=f"compressed_position_ids_cuda_{compress_ratio}", + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + + # compressed_mask_cuda: per-token bool mask for postprocess scatter. + # Precomputed from new_comp_kv_lens to skip padded generation slots. + self.compressed_mask_cuda = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + dtype=torch.bool, + cache_name=f"compressed_mask_cuda_{compress_ratio}", + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + self.compressed_mask = { + compress_ratio: torch.empty_like( + self.compressed_mask_cuda[compress_ratio], device="cpu", pin_memory=prefer_pinned() + ) + for compress_ratio in self.compress_ratio_set + } + + # empty topk indices buffer with all -1s in the tensor + self.empty_topk_indices_buffer = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens, self.sparse_mla_topk), + cache_name="empty_topk_indices_buffer", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.empty_topk_indices_buffer.fill_(-1) + + # SWA local indices + self.swa_local_indices_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens, self.sparse_attention_config.window_size), + cache_name="swa_local_indices_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + # Compute max_compressed_indices for CUDA graph compatibility. + # For ratio=4, the indexer selects index_topk compressed tokens. + # For ratio=128, use max_seq_len / 128 rounded up to next power of 2 + raw_128 = math.ceil(self.max_seq_len / 128) + po2_128 = 1 << (raw_128 - 1).bit_length() if raw_128 > 0 else 1 + self.max_compressed_indices = { + 1: 0, # No compressed indices + 4: self.sparse_mla_topk, # index_topk from indexer + 128: po2_128, # All compressed tokens, rounded to power of 2 + } + + # Compressed local indices for compress_ratio=128 + # Note: ratio=4 uses dynamic topk_indices from indexer, so we only pre-allocate for ratio=128 + self.compressed_local_indices_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens, self.max_compressed_indices[128]), + cache_name="compressed_local_indices_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + self.block_tables = { + attention_type: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq), + cache_name=f"block_tables_{attention_type}", + dtype=torch.int32, + capture_graph=capture_graph, + ) + for attention_type in self.attention_type_set + } + self.host_block_tables = { + attention_type: torch.empty_like( + self.block_tables[attention_type], device="cpu", pin_memory=prefer_pinned() + ) + for attention_type in self.attention_type_set + } + + # sparse_mla_topk_lens: actual token count per token for each compress_ratio (SWA + compressed) + # Shape: [max_num_tokens] per compress_ratio + self.sparse_mla_topk_lens = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"sparse_mla_topk_lens_{compress_ratio}", + dtype=torch.int32, + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + } + + # cached_token_lens_cuda: number of tokens already cached per request + self.cached_token_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + cache_name="cached_token_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.cached_token_lens_cpu = torch.empty_like( + self.cached_token_lens_cuda, device="cpu", pin_memory=prefer_pinned() + ) + + # Cache buffer data pointers are constant after KV cache allocation, + # so compute them once during initialization instead of every prepare(). + self._init_cache_buffer_data_pointers() + + def prepare_for_indexer_k_cache(self): + """Optimized: bulk tensor copy instead of per-row Python loop. + + Note: must use num_contexts=0 for get_batch_attn_offset because the + indexer k cache always uses generation-style copy indices regardless + of request type. + """ + num_seqs = self.num_seqs + offsets = self.kv_cache_manager.get_batch_attn_offset( + self.request_ids, + 1, # beam_width + 0, # num_contexts=0 (indexer always uses gen-style copy index) + num_seqs, + DeepseekV4AttentionType.INDEXER_COMPRESS, + DEEPSEEK_V4_SPARSE_RATIO, + ) + num_cols = offsets.shape[1] + self.host_indexer_k_cache_block_offsets[:num_seqs, :num_cols] = offsets[:num_seqs] + self.indexer_k_cache_block_offsets[:num_seqs].copy_( + self.host_indexer_k_cache_block_offsets[:num_seqs], + non_blocking=True, + ) + + def prepare_for_block_tables(self): + """Prepare block tables for all attention types. + + Delegates offset computation to DeepseekV4CacheManager.get_batch_block_offsets + (single get_copy_index call, deduplicated by pool), then copies to device. + """ + num_seqs = self.num_seqs + offsets_map = self.kv_cache_manager.get_batch_block_offsets( + self.request_ids, self.num_contexts, self.attention_type_set + ) + + # Phase 1: write host buffers. + for key, offsets in offsets_map.items(): + self.host_block_tables[key][:num_seqs] = offsets[:num_seqs] + + # Phase 2: batch H2D copies. + for key in self.attention_type_set: + self.block_tables[key][:num_seqs].copy_( + self.host_block_tables[key][:num_seqs], non_blocking=True + ) + + def prepare_for_deepseek_v4_indices(self, token_positions=None): + """Prepare SWA/compressed local indices and sparse_mla_topk_lens.""" + window_size = self.sparse_attention_config.window_size + device = self.swa_local_indices_cuda.device + + if token_positions is None: + # Initial prepare() path — build token_positions from CPU data. + num_tokens = self.num_tokens + num_requests = self.num_seqs + + # cu_seq_lens_cuda must already be populated before this call + cu_seq_lens = self.cu_seq_lens_cuda[: num_requests + 1] + cached_tokens = self.cached_token_lens_cuda[:num_requests] + + token_idx = torch.arange(num_tokens, dtype=torch.int32, device=device) + req_idx = torch.searchsorted(cu_seq_lens[1:].to(torch.int32), token_idx, right=True) + offsets = token_idx - cu_seq_lens[req_idx].to(torch.int32) + token_positions = cached_tokens[req_idx].to(torch.int32) + offsets + + self._prepare_deepseek_v4_indices_compiled( + token_positions, + window_size, + self.max_compressed_indices[128], + self.sparse_mla_topk, + self.swa_local_indices_cuda, + self.compressed_local_indices_cuda, + self.sparse_mla_topk_lens, + self._compress_ratios_sorted, + ) + + @staticmethod + @maybe_compile(dynamic=True, options={"max-autotune": True}) + def _prepare_deepseek_v4_indices_compiled( + token_positions: torch.Tensor, + window_size: int, + max_compressed_indices_128: int, + sparse_mla_topk: int, + swa_local_indices_buf: torch.Tensor, + compressed_local_indices_buf: torch.Tensor, + sparse_mla_topk_lens_bufs: Dict[int, torch.Tensor], + compress_ratios: list, + ): + """Build SWA indices, compressed indices, and topk_lens in one fused graph.""" + device = token_positions.device + num_tokens = token_positions.shape[0] + + # ── SWA local indices ── + positions = token_positions.unsqueeze(1) # [num_tokens, 1] + swa_offsets = torch.arange(window_size, dtype=torch.int32, device=device) + swa_start = (positions - window_size + 1).clamp(min=0) + swa_indices = swa_start + swa_offsets + swa_indices = torch.where(swa_indices > positions, -1, swa_indices).to(torch.int32) + swa_local_indices_buf[:num_tokens] = swa_indices + + # ── Compressed local indices (ratio=128) ── + # Hardcoded 128: compressed_local_indices only applies to ratio=128 + # layers. This is a design constraint — see compress_ratio_set + # validation in __post_init__ which restricts ratios to {1, 4, 128}. + num_valid = (token_positions + 1) // 128 + comp_col = torch.arange(max_compressed_indices_128, dtype=torch.int32, device=device) + valid_mask = comp_col.unsqueeze(0) < num_valid.unsqueeze(1) + comp_indices = torch.where( + valid_mask, + comp_col.unsqueeze(0).expand(num_tokens, -1), + torch.full( + (num_tokens, max_compressed_indices_128), -1, dtype=torch.int32, device=device + ), + ) + compressed_local_indices_buf[:num_tokens] = comp_indices + + # ── sparse_mla_topk_lens per compress_ratio ── + # Dual-pool layout: + # - compress_ratio=1: min(kv_len, window_size) — actual valid SWA count + # - compress_ratio=4: window_size + min(kv_len // 4, sparse_mla_topk) + # (fixed 128 SWA slots + compressed valid count) + # - compress_ratio=128: window_size + kv_len // 128 + # (fixed 128 SWA slots + compressed valid count) + # For ratio>1, the SWA region always occupies exactly window_size (128) + # slots. Invalid SWA positions are padded with -1 in the index buffer. + kv_lens = token_positions + 1 + for compress_ratio in compress_ratios: + if compress_ratio == 1: + total_count = kv_lens.clamp(max=window_size) + elif compress_ratio == 4: + compressed_count = (kv_lens // compress_ratio).clamp(max=sparse_mla_topk) + total_count = window_size + compressed_count + elif compress_ratio == 128: + compressed_count = kv_lens // compress_ratio + total_count = window_size + compressed_count + else: + raise ValueError(f"Unsupported compress_ratio: {compress_ratio}") + sparse_mla_topk_lens_bufs[compress_ratio][:num_tokens] = total_count.to(torch.int32) + + def _init_cache_buffer_data_pointers(self): + # If MTP is enabled, enlarge the compress ratios by max_draft_tokens - 1 + extend_compress_ratios = self.compress_ratios + [self.compress_ratios[-1]] * ( + self.max_draft_tokens - 1 + ) + self.swa_buffer_ptrs = { + layer_idx: self.kv_cache_manager.get_buffers( + layer_idx, DeepseekV4AttentionType.SWA + ).data_ptr() + for layer_idx in self.kv_cache_manager.pp_layers + } + self.compressed_buffer_ptrs = { + layer_idx: self.kv_cache_manager.get_buffers( + layer_idx, DeepseekV4AttentionType.COMPRESS + ).data_ptr() + for layer_idx in self.kv_cache_manager.pp_layers + if is_compress_layer(extend_compress_ratios[layer_idx]) + } + + # Per-ratio base pointer for sparse MLA global-index conversion. + # Ratio 1 is the SWA pool; ratios greater than 1 are compressed pools. + self.sparse_mla_base_ptrs = { + 1: self.kv_cache_manager.swa_pool_ptr, + } + for ratio, compress_pool_ptr in self.kv_cache_manager.compress_pool_ptrs.items(): + self.sparse_mla_base_ptrs[ratio] = compress_pool_ptr + + def prepare(self): + TrtllmAttentionMetadata.prepare(self) + + num_requests = self.num_contexts + self.num_generations + num_gen_tokens = self.num_tokens - self.num_ctx_tokens + + self.cached_token_lens_cpu[:num_requests] = torch.tensor( + self.kv_cache_params.num_cached_tokens_per_seq[:num_requests], + dtype=torch.int32, + ) + cached_token_lens = self.cached_token_lens_cpu + kv_lens = cached_token_lens[:num_requests] + self.seq_lens_kv[:num_requests] + + self.cached_token_lens_cuda[:num_requests].copy_( + cached_token_lens[:num_requests], non_blocking=True + ) + + # Prepare cu_seq_lens early — needed by prepare_for_deepseek_v4_indices + self.cu_seq_lens[1 : num_requests + 1] = self.seq_lens.cumsum(0) + self.cu_seq_lens_cuda[: num_requests + 1].copy_( + self.cu_seq_lens[: num_requests + 1], non_blocking=True + ) + + # For indices conversion + self.prepare_for_indices_conversion() + + has_sparse_layers = DEEPSEEK_V4_SPARSE_RATIO in self.compress_ratio_set + + # For indexer k cache (only needed when sparse layers exist) + if has_sparse_layers: + self.prepare_for_indexer_k_cache() + + # For spec decode + self.prepare_for_spec_decode(kv_lens) + + # For block offsets + self.prepare_for_block_tables() + + # For DeepSeek-V4 indices + self.prepare_for_deepseek_v4_indices() + + # Prepare metadata for indexer (only needed when sparse layers exist) + if has_sparse_layers: + DeepseekV4Indexer.prepare(metadata=self) + + # --- Per-ratio metadata --- + # 1) CPU-side: compute scalar metadata (num_total_compressed_tokens, etc.) + # 2) CUDA-side: fill *_cuda buffers via prepare_compressed_kv_metadata() + num_gen_tokens_per_seq = ( + num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 + ) + self.num_gen_tokens_per_seq = num_gen_tokens_per_seq + num_contexts = self.num_contexts + num_generations = self.num_generations + kv_lens_slice = kv_lens[:num_requests] + cached_slice = cached_token_lens[:num_requests] + + if num_contexts > 0: + # Prefill path: need per-request tensor ops for ctx scalar metadata. + for compress_ratio in self.compress_ratio_set: + new_comp_kv_lens = kv_lens_slice // compress_ratio - cached_slice // compress_ratio + cu_new = new_comp_kv_lens.cumsum(0) + num_ctx_compressed_tokens = cu_new[num_contexts - 1].item() + num_gen_compressed_tokens = num_generations * ( + (num_gen_tokens_per_seq + compress_ratio - 1) // compress_ratio + ) + self.num_total_compressed_tokens[compress_ratio] = ( + num_ctx_compressed_tokens + num_gen_compressed_tokens + ) + self.max_ctx_compressed_tokens[compress_ratio] = ( + new_comp_kv_lens[:num_contexts].max().item() + ) + else: + # Decode-only: scalars depend only on num_generations and + # num_gen_tokens_per_seq, no per-request tensor ops needed. + for compress_ratio in self.compress_ratio_set: + self.num_total_compressed_tokens[compress_ratio] = num_generations * ( + (num_gen_tokens_per_seq + compress_ratio - 1) // compress_ratio + ) + self.max_ctx_compressed_tokens[compress_ratio] = 0 + + # 2) CUDA-side: fill *_cuda buffers on device. + kv_lens_cuda = ( + self.cached_token_lens_cuda[:num_requests] + self._seq_lens_cuda[:num_requests] + ) + cached_tokens_cuda = self.cached_token_lens_cuda[:num_requests] + self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda) + + self._compute_compressed_mask( + self.new_comp_kv_lens_cuda, + self.cu_new_comp_kv_cuda, + self.compressed_mask_cuda, + num_requests, + self.num_total_compressed_tokens, + self._compress_ratios_sorted, + ) + + def prepare_compressed_kv_metadata( + self, + kv_lens: torch.Tensor, + cached_tokens: torch.Tensor, + ): + """Compute per-ratio compressed KV lens and position IDs on device. + + Shared by prepare() and on_update_kv_lens() to avoid duplicated logic. + + Args: + kv_lens: Total KV lengths per request (device tensor, [batch_size]). + cached_tokens: Cached token counts per request (device tensor, [batch_size]). + """ + batch_size = kv_lens.shape[0] + num_contexts = self.num_contexts + num_generations = self.num_generations + + self._compute_per_ratio_kv_lens( + kv_lens, + cached_tokens, + batch_size, + self.compressed_kv_lens_cuda, + self.past_kv_lens_cuda, + self.new_comp_kv_lens_cuda, + self.cu_new_comp_kv_cuda, + self._compress_ratios_sorted, + ) + + if num_contexts > 0: + self._compute_ctx_compressed_position_ids( + self.past_kv_lens_cuda, + self.cu_new_comp_kv_cuda, + self.compressed_position_ids_cuda, + num_contexts, + self._compress_ratios_sorted, + ) + + if self.num_gen_tokens_per_seq > 0 and num_generations > 0: + # Extract output_offset as Python int per ratio to avoid + # tensor-scalar slice inside compiled function. + # For decode-only batches (num_contexts == 0), offset is 0. + gen_output_offsets = { + r: self.cu_new_comp_kv_cuda[r][num_contexts].item() if num_contexts > 0 else 0 + for r in self._compress_ratios_sorted + } + self._compute_gen_compressed_position_ids( + self.past_kv_lens_cuda, + self.compressed_position_ids_cuda, + num_contexts, + num_generations, + self.num_gen_tokens_per_seq, + self._compress_ratios_sorted, + gen_output_offsets, + ) + + def on_update_kv_lens(self): + """Recompute kv-lens-dependent DeepSeek-V4 metadata on device.""" + super().on_update_kv_lens() + + batch_size = self.num_seqs + num_tokens = self.num_tokens + kv_lens = self.kv_lens_cuda[:batch_size] + seq_lens = self._seq_lens_cuda[:batch_size] + cached_tokens = kv_lens - seq_lens + + num_gen_tokens = num_tokens - self.num_ctx_tokens + self.num_gen_tokens_per_seq = ( + num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 + ) + + self.prepare_compressed_kv_metadata(kv_lens, cached_tokens) + + self._compute_compressed_mask( + self.new_comp_kv_lens_cuda, + self.cu_new_comp_kv_cuda, + self.compressed_mask_cuda, + batch_size, + self.num_total_compressed_tokens, + self._compress_ratios_sorted, + ) + + token_positions = self._compute_token_positions( + seq_lens, + cached_tokens, + batch_size, + num_tokens, + self.cu_seq_lens_cuda, + self.req_idx_per_token, + ) + + self.prepare_for_deepseek_v4_indices(token_positions) + + @staticmethod + @maybe_compile(dynamic=True, options={"max-autotune": True}) + def _compute_per_ratio_kv_lens( + kv_lens: torch.Tensor, + cached_tokens: torch.Tensor, + batch_size: int, + compressed_kv_lens_bufs: Dict[int, torch.Tensor], + past_kv_lens_bufs: Dict[int, torch.Tensor], + new_comp_kv_lens_bufs: Dict[int, torch.Tensor], + cu_new_comp_kv_bufs: Dict[int, torch.Tensor], + compress_ratios: list, + ): + """Compute per-ratio compressed/past/new kv lens and cu_new_comp.""" + for compress_ratio in compress_ratios: + compressed_kv = (kv_lens // compress_ratio).to(torch.int) + compressed_kv_lens_bufs[compress_ratio][:batch_size] = compressed_kv + + past_kv = (cached_tokens // compress_ratio).to(torch.int) + past_kv_lens_bufs[compress_ratio][:batch_size] = past_kv + + new_comp = compressed_kv - past_kv + new_comp_kv_lens_bufs[compress_ratio][:batch_size] = new_comp + + cu_new_comp = cu_new_comp_kv_bufs[compress_ratio] + cu_new_comp[: batch_size + 1] = torch.nn.functional.pad( + torch.cumsum(new_comp, dim=0), (1, 0) + ) + + @staticmethod + def _compute_token_positions( + seq_lens: torch.Tensor, + cached_tokens: torch.Tensor, + batch_size: int, + num_tokens: int, + cu_seq_lens_buf: torch.Tensor, + req_idx_per_token_buf: torch.Tensor, + ) -> torch.Tensor: + """Compute cu_seq_lens, req_idx_per_token, and token_positions (eager).""" + device = seq_lens.device + + # cu_seq_lens + cu_seq_lens_buf[: batch_size + 1] = torch.nn.functional.pad( + torch.cumsum(seq_lens.to(torch.int), dim=0), (1, 0) + ) + + # req_idx_per_token via searchsorted + token_idx = torch.arange(num_tokens, dtype=torch.int32, device=device) + req_idx = torch.searchsorted( + cu_seq_lens_buf[1 : batch_size + 1].to(torch.int32), token_idx, right=True + ) + req_idx_per_token_buf[:num_tokens] = req_idx + + # token positions + base_pos = cached_tokens[req_idx].to(torch.int32) + offsets = token_idx - cu_seq_lens_buf[req_idx].to(torch.int32) + return base_pos + offsets + + @staticmethod + @maybe_compile(dynamic=True, options={"max-autotune": True}) + def _compute_gen_compressed_position_ids( + past_kv_lens_bufs: Dict[int, torch.Tensor], + compressed_position_ids_bufs: Dict[int, torch.Tensor], + num_contexts: int, + num_generations: int, + num_gen_tokens_per_seq: int, + compress_ratios: list, + gen_output_offsets: Dict[int, int], + ): + """Generation compressed position IDs. + + gen_output_offsets: dict mapping compress_ratio -> Python int offset, + pre-extracted by the caller to avoid tensor-scalar .item() inside + compiled code. 0 for decode-only batches.""" + device = past_kv_lens_bufs[compress_ratios[0]].device + batch_size = num_contexts + num_generations + for compress_ratio in compress_ratios: + gen_past = past_kv_lens_bufs[compress_ratio][num_contexts:batch_size] + new_gen_comp = (num_gen_tokens_per_seq + compress_ratio - 1) // compress_ratio + gen_offsets = torch.arange(new_gen_comp, dtype=torch.int32, device=device) + gen_pos = gen_past.unsqueeze(1) + gen_offsets.unsqueeze(0) + gen_comp = num_generations * new_gen_comp + result = (gen_pos.reshape(-1) * compress_ratio).to(torch.int) + output_offset = gen_output_offsets[compress_ratio] + compressed_position_ids_bufs[compress_ratio][ + output_offset : output_offset + gen_comp + ] = result + + @staticmethod + @maybe_compile(dynamic=True, options={"max-autotune": True}) + def _compute_compressed_mask( + new_comp_kv_lens_bufs: Dict[int, torch.Tensor], + cu_new_comp_kv_bufs: Dict[int, torch.Tensor], + compressed_mask_bufs: Dict[int, torch.Tensor], + batch_size: int, + num_total_compressed_tokens: Dict[int, int], + compress_ratios: list, + ): + """Compute per-token compressed_mask on device (graph-safe, no .item()). + + For each token in [0, total_tokens), determine which sequence it + belongs to via searchsorted on cu_new_comp_kv, then compare its + within-sequence offset against the actual new_comp_kv_len. + Context tokens (offset < actual) are always True. + Generation tokens whose offset >= actual new_comp are padding → False. + """ + device = new_comp_kv_lens_bufs[compress_ratios[0]].device + for compress_ratio in compress_ratios: + total_tokens = num_total_compressed_tokens[compress_ratio] + new_comp = new_comp_kv_lens_bufs[compress_ratio][:batch_size] + cu = cu_new_comp_kv_bufs[compress_ratio][: batch_size + 1] + + token_idx = torch.arange(total_tokens, dtype=torch.int32, device=device) + seq_idx = torch.searchsorted(cu[1:], token_idx, right=True) + seq_idx = seq_idx.clamp_(max=batch_size - 1) + offset_in_seq = token_idx - cu[seq_idx] + compressed_mask_bufs[compress_ratio][:total_tokens] = offset_in_seq < new_comp[seq_idx] + + @staticmethod + def _compute_ctx_compressed_position_ids( + past_kv_lens_bufs: Dict[int, torch.Tensor], + cu_new_comp_kv_bufs: Dict[int, torch.Tensor], + compressed_position_ids_bufs: Dict[int, torch.Tensor], + num_contexts: int, + compress_ratios: list, + ): + """Context-only compressed position IDs (eager, data-dependent shapes).""" + device = past_kv_lens_bufs[compress_ratios[0]].device + for compress_ratio in compress_ratios: + past_kv = past_kv_lens_bufs[compress_ratio] + cu_new_comp = cu_new_comp_kv_bufs[compress_ratio] + + total_ctx_comp = cu_new_comp[num_contexts] + ctx_idx = torch.arange(total_ctx_comp, dtype=torch.int32, device=device) + ctx_cu = cu_new_comp[: num_contexts + 1].to(torch.int32) + ctx_req = torch.searchsorted(ctx_cu[1:], ctx_idx, right=True) + ctx_offset = ctx_idx - ctx_cu[ctx_req] + compressed_position_ids_bufs[compress_ratio][:total_ctx_comp] = ( + (past_kv[:num_contexts][ctx_req] + ctx_offset) * compress_ratio + ).to(torch.int) + + +class DeepseekV4Indexer(Indexer): + def __init__( + self, + quant_config: Optional[QuantConfig], + pos_embd_params: Optional[PositionalEmbeddingParams], + mla_params: Optional[MLAParams], + skip_create_weights_in_init: bool, + sparse_attention_config: "SparseAttentionConfig", + dtype: Optional[torch.dtype], + compress_ratio: int = 1, + layer_idx: int = 0, + aux_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__( + quant_config, + pos_embd_params, + mla_params, + skip_create_weights_in_init, + sparse_attention_config, + dtype, + compress_ratio, + layer_idx, + aux_stream, + ) + # Override base Indexer.weights_proj to bf16 (matches V4 checkpoint). + self.weights_proj = Linear( + self.hidden_size, + self.n_heads, + bias=False, + dtype=dtype, + quant_config=None, + skip_create_weights_in_init=skip_create_weights_in_init, + use_custom_cublas_mm=True, + ) + self.rotary_emb = RotaryEmbedding( + pos_embd_params.rope, + head_dim=self.rope_dim, + is_neox=False, + ) + rms_norm_eps = 1e-6 + index_head_dim = sparse_attention_config.index_head_dim + indexer_mla_params = MLAParams( + hidden_size=mla_params.hidden_size, + qk_rope_head_dim=mla_params.qk_rope_head_dim, + qk_nope_head_dim=index_head_dim - mla_params.qk_rope_head_dim, + ) + self.indexer_k_cache_dtype = getattr( + sparse_attention_config, "indexer_k_cache_dtype", "fp8_blockwise" + ) + self.indexer_cache_dtype = resolve_kv_cache_dtype(self.indexer_k_cache_dtype) + self.compressor = Compressor( + indexer_mla_params, + layer_idx, + compress_ratio, + rms_norm_eps, + skip_create_weights_in_init, + pos_embd_params, + dtype=dtype, + kv_cache_dtype=self.indexer_k_cache_dtype, + is_indexer=True, + rotate_activation=True, + ) + + def post_load_weights(self): + # V4 does not use the V3 fused fp32 wk+weights_proj GEMM, and the + # base concat would now hit an fp32/bf16 dtype mismatch. + return + + def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): + """Project Q and apply RoPE. + + Returns q with layout [num_tokens, n_heads, head_dim] where + head_dim = nope_dim + rope_dim, RoPE already applied in-place. + """ + q = self.wq_b(qr) + q = q.view(-1, self.n_heads, self.head_dim) + # Fused in-place RoPE on the rope portion of each head + nope_dim = self.head_dim - self.rope_dim + torch.ops.trtllm.mla_rope_inplace( + q, + position_ids.view(-1), + self.rotary_emb.rotary_cos_sin, + self.n_heads, + nope_dim, + self.rope_dim, + False, + self.rotary_emb.is_neox, + ) + return q + + def _update_k_cache(self, k_fp8, k_scale, metadata): + """Overwrite the fused compressor's FP8 cache with reference-style FP4-QAT values.""" + super()._update_k_cache(k_fp8, k_scale, metadata) + + def forward( + self, + qr: torch.Tensor, + hidden_states: torch.Tensor, + metadata: DeepseekV4TrtllmAttentionMetadata, + position_ids: torch.Tensor, + ): + if self.indexer_cache_dtype not in ( + KVCacheDtype.FP8_PERTENSOR, + KVCacheDtype.FP8_BLOCKWISE, + ): + raise NotImplementedError( + "DeepSeek-V4 indexer BMM currently consumes FP8 K. " + f"Indexer cache preset {self.indexer_k_cache_dtype!r} is supported by " + "the compressor/cache scatter path, but needs the matching BF16/FP4 " + "indexer BMM path before end-to-end indexer execution can use it." + ) + # compress k + k_fp8, k_scale = self.compressor(hidden_states, metadata) + + # multi-stream q proj/rope and weights proj + q, weights = maybe_execute_in_parallel( + lambda: self._qk_projection_and_rope(qr, position_ids), + lambda: self.weights_proj(hidden_states), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + + # Rotate + quantize (layout matches compressor K: [nope|pe]) + q = rotate_activation(q) + q = q.view(-1, self.head_dim) + q_fp8, q_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( + q, use_ue8m0=self.scale_fmt == "ue8m0" + ) + q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim) + q_scale = q_scale.view(-1, self.n_heads, 1) + + # weights scale + weights = self._weight_scale(weights, q_scale) + + # If there are no compressed tokens, return an topk indices buffer with all -1s in the tensor. + if k_fp8 is None: + topk_indices = metadata.empty_topk_indices_buffer[: hidden_states.shape[0]] + else: + topk_indices = self.sparse_attn_indexer( + metadata, hidden_states, q_fp8, k_fp8, k_scale, weights + ) + return topk_indices + + +class DeepseekV4TrtllmAttention(TrtllmAttention): + Metadata = DeepseekV4TrtllmAttentionMetadata + + def __init__( + self, + layer_idx: int, + num_heads: int, + head_dim: int, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + q_scaling: Optional[float] = None, + pos_embd_params: Optional[PositionalEmbeddingParams] = None, + mla_params: Optional[MLAParams] = None, + skip_create_weights_in_init: bool = False, + attention_chunk_size: Optional[int] = None, + sparse_attention_config: Optional["SparseAttentionConfig"] = None, + dtype: Optional[torch.dtype] = None, + aux_stream: Optional[torch.cuda.Stream] = None, + **kwargs, + ): + assert sparse_attention_config is not None, ( + "sparse_attention_config is required for DeepseekV4TrtllmAttention and cannot be None" + ) + TrtllmAttention.__init__( + self, + layer_idx, + num_heads, + head_dim, + sparse_attention_config=sparse_attention_config, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + skip_create_weights_in_init=skip_create_weights_in_init, + attention_chunk_size=attention_chunk_size, + **kwargs, + ) + + self.compress_ratio = sparse_attention_config.compress_ratios[layer_idx] + + if self.compress_ratio == 4: + self.indexer = DeepseekV4Indexer( + quant_config, + pos_embd_params, + mla_params, + skip_create_weights_in_init, + sparse_attention_config, + dtype, + self.compress_ratio, + layer_idx, + aux_stream, + ) + + if self.compress_ratio > 1: + rms_norm_eps = 1e-6 + has_fp8_kv_cache = False + if quant_config is not None: + has_fp8_kv_cache = quant_config.layer_quant_mode.has_fp8_kv_cache() + kv_cache_dtype = "fp8_pertensor" if has_fp8_kv_cache else "default" + self.compressor = Compressor( + mla_params, + layer_idx, + self.compress_ratio, + rms_norm_eps, + skip_create_weights_in_init, + pos_embd_params, + kv_cache_dtype=kv_cache_dtype, + dtype=dtype, + rotate_activation=False, + ) + + def forward(self, *args, **kwargs): + attn_sink = getattr(self, "attn_sink", None) + if attn_sink is not None and kwargs.get("attention_sinks") is None: + kwargs["attention_sinks"] = attn_sink.data + return super().forward(*args, **kwargs) + + def sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DeepseekV4TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + """Convert local indices (SWA + compressed) to global pool indices.""" + layer_idx = self.layer_idx + kv_cache_manager = metadata.kv_cache_manager + attention_input_type = forward_args.attention_input_type + + swa_pool_base_ptr = metadata.sparse_mla_base_ptrs[1] + + # Get cached buffer pointers + swa_buffer_ptr = metadata.swa_buffer_ptrs[layer_idx] + + # Token stride + index_head_dim = self.sparse_attention_config.index_head_dim + has_fp8_kv_cache = False + if self.quant_config is not None: + has_fp8_kv_cache = self.quant_config.layer_quant_mode.has_fp8_kv_cache() + token_stride = get_token_bytes( + self.head_dim, + index_head_dim, + self.compress_ratio, + DeepseekV4AttentionType.SWA, + has_fp8_kv_cache, + ) + + # Select token range based on phase + if attention_input_type == AttentionInputType.context_only: + start_idx = 0 + end_idx = metadata.num_ctx_tokens + elif attention_input_type == AttentionInputType.generation_only: + start_idx = metadata.num_ctx_tokens + end_idx = metadata.num_tokens + else: + start_idx = 0 + end_idx = metadata.num_tokens + + # Use global req_id directly + req_id = metadata.req_idx_per_token[start_idx:end_idx] + swa_local_indices = metadata.swa_local_indices_cuda[start_idx:end_idx] + block_table_swa = metadata.block_tables[(1, DeepseekV4AttentionType.SWA)] + + if self.compress_ratio > 1: + compressed_buffer_ptr = metadata.compressed_buffer_ptrs[layer_idx] + compress_pool_base_ptr = metadata.sparse_mla_base_ptrs[self.compress_ratio] + block_table_compressed = metadata.block_tables[ + (self.compress_ratio, DeepseekV4AttentionType.COMPRESS) + ] + if self.compress_ratio == 4: + topk_indices = forward_args.topk_indices + assert topk_indices is not None, "topk_indices is required when compress_ratio=4" + compressed_local_indices = topk_indices + else: + compressed_local_indices = metadata.compressed_local_indices_cuda[start_idx:end_idx] + else: + compressed_buffer_ptr = 0 + compress_pool_base_ptr = 0 + block_table_compressed = None + compressed_local_indices = None + + global_indices = deepseek_v4_local_to_global_indices( + req_id=req_id, + block_table_swa=block_table_swa, + swa_local_indices=swa_local_indices, + swa_pool_base_ptr=swa_pool_base_ptr, + swa_buffer_ptr=swa_buffer_ptr, + tokens_per_block=kv_cache_manager.tokens_per_block, + token_stride=token_stride, + block_table_compressed=block_table_compressed, + compressed_local_indices=compressed_local_indices, + compress_pool_base_ptr=compress_pool_base_ptr, + compressed_buffer_ptr=compressed_buffer_ptr, + compress_ratio=self.compress_ratio, + num_compressed_indices=metadata.max_compressed_indices[self.compress_ratio], + ) + + return global_indices, None + + def sparse_kv_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + metadata: DeepseekV4TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + return None, None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index ab16aac4d716..22619142303b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 """Dense Sparse Attention (DSA) backend for TRT-LLM with indexer-based TopK selection.""" import math import threading @@ -12,28 +14,30 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.interface import ( - AttentionForwardArgs, AttentionInputType, MLAParams, - PositionalEmbeddingParams) -from tensorrt_llm._torch.attention_backend.trtllm import ( - TrtllmAttention, TrtllmAttentionMetadata) + AttentionForwardArgs, + MLAParams, + PositionalEmbeddingParams, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.distributed.ops import allgather from tensorrt_llm._torch.modules.layer_norm import LayerNorm from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.modules.multi_stream_utils import \ - maybe_execute_in_parallel +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.utils import maybe_compile -from tensorrt_llm._utils import get_size_in_bytes, get_sm_version, prefer_pinned +from tensorrt_llm._utils import get_size_in_bytes, get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.bindings.internal.batch_manager import \ - CacheType as CacheTypeCpp -from tensorrt_llm.deep_gemm import (fp8_fp4_mqa_logits, - fp8_fp4_paged_mqa_logits, fp8_mqa_logits, - fp8_paged_mqa_logits, - get_paged_mqa_logits_metadata) +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.deep_gemm import ( + fp8_fp4_mqa_logits, + fp8_fp4_paged_mqa_logits, + fp8_mqa_logits, + fp8_paged_mqa_logits, + get_paged_mqa_logits_metadata, +) from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -42,6 +46,9 @@ ModelConfig = tensorrt_llm.bindings.ModelConfig if TYPE_CHECKING: + from tensorrt_llm._torch.speculative.interface import SpecMetadata + from tensorrt_llm._torch.speculative.spec_tree_manager import \ + SpecTreeManager from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig # Optional import: fast-hadamard-transform causes CI build issues (requires wheel+torch pre-installed) @@ -94,13 +101,27 @@ def warmup_heuristic_topk_decode(top_k: int = 2048, indices = torch.empty((1, top_k), dtype=torch.int32, device=device) pre_idx = torch.zeros((1, hint_size), dtype=torch.int32, device=device) scratch = torch.empty((top_k, ), dtype=torch.float32, device=device) + # The default warmup geometry (num_cols=4096) falls below kSeqSmall=12288 + # and routes to the Radix path with blocks_per_row=2 (num_rows=1 sweeps + # bp ∈ [2, maxByCols=2]). The cpp op rejects blocks_per_row > 1 without + # caller-owned radix aux scratch, so supply worst-case (kMaxBlocksPerRowDecode=10) + # buffers here. Cost is negligible (~80 KB) and the warmup is a one-shot. + _radix_max_bp = 10 + radix_aux_indices = torch.empty((1, _radix_max_bp, top_k), + dtype=torch.int32, + device=device) + radix_aux_logits = torch.empty((1, _radix_max_bp, top_k), + dtype=torch.float32, + device=device) torch.ops.trtllm.indexer_topk_decode(logits, seq_lens, indices, 1, top_k, pre_idx=pre_idx, - heuristic_scratch=scratch) + heuristic_scratch=scratch, + radix_aux_indices=radix_aux_indices, + radix_aux_logits=radix_aux_logits) torch.cuda.synchronize() @@ -187,7 +208,7 @@ def _compute_slot_mappings( quant_block_size: int, data_bytes_per_token: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute flat byte indices for indexer K data and scales from global token positions. + """Compute flat byte indices for FP8/FP4 data and scales from global token positions. Shared by Indexer.prepare() (CPU) and on_update_kv_lens() (GPU) to avoid duplicating the slot mapping arithmetic. @@ -196,14 +217,14 @@ def _compute_slot_mappings( global_positions: Per-token absolute position in the KV sequence. block_offsets: [num_seqs, max_blocks_per_seq] block offset table. req_indices: Per-token request index. - head_dim: Indexer head dimension (used for the scale-size formula). + head_dim: Indexer head dimension (logical element count). tokens_per_block: Tokens stored per cache block. quant_block_size: Quantization block size. data_bytes_per_token: Bytes of quantized data per token in the cache - pool. FP8 stores one byte per element (= head_dim). FP4 packs two - E2M1 codes per byte (= head_dim // 2). Defaults to ``head_dim`` - when unset, preserving the FP8 layout for callers that haven't - threaded the FP4 dtype through. + (default: head_dim, i.e. one byte per FP8 value). For MXFP4 this + should be ``head_dim // 2`` since two E2M1 codes pack into one byte. + The scale layout is identical at head_dim=128: 4 bytes per token + (one float32 for FP8, four packed UE8M0 exponents for FP4). Returns: (fp8_indices, scale_indices): Flat byte offsets into the cache pool. @@ -234,6 +255,24 @@ def _compute_slot_mappings( return fp8_indices, scale_indices +def _unravel_indices(flat_indices: torch.Tensor, + shape: Tuple[int, ...]) -> Tuple[torch.Tensor, ...]: + """ + Unravel indices into multiple dimensions. + """ + d3 = shape[3] + i3 = flat_indices % d3 + flat_indices = flat_indices // d3 + d2 = shape[2] + i2 = flat_indices % d2 + flat_indices = flat_indices // d2 + d1 = shape[1] + i1 = flat_indices % d1 + flat_indices = flat_indices // d1 + i0 = flat_indices + return i0, i1, i2, i3 + + def rotate_activation(x: torch.Tensor) -> torch.Tensor: """Apply Hadamard rotation to activation tensor for DSA sparse attention.""" assert x.dtype == torch.bfloat16 @@ -241,7 +280,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: if not HAS_FAST_HADAMARD: # Fallback: skip transformation (acceptable for test/dev) logger.warning_once( - "fast-hadamard-transform not available. DSA sparse attention will skip " + "fast-hadamard-transform not available. Sparse MLA will skip " "hadamard transformation. Install with: " "pip install git+https://github.com/Dao-AILab/fast-hadamard-transform.git", key="fast_hadamard_import_missing") @@ -365,11 +404,25 @@ def split_prefill_chunks( return chunk_groups +def _select_indexer_compress_ratio(compress_ratios: List[int]) -> int: + if 4 in compress_ratios: + return 4 + if 1 in compress_ratios: + return 1 + return 0 + + +def _effective_compress_ratio_divisor(compress_ratio: int) -> int: + return compress_ratio if compress_ratio > 1 else 1 + + def compute_cu_seqlen_kv_bounds_with_cache( seq_lens: torch.Tensor, num_contexts: int, num_ctx_tokens: int, cached_token_lens: Optional[torch.Tensor] = None, + kv_lens: Optional[torch.Tensor] = None, + compress_ratio: int = 1, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute attention window bounds for batched sequences with causal attention, @@ -380,6 +433,8 @@ def compute_cu_seqlen_kv_bounds_with_cache( num_contexts: Number of sequences in the batch num_ctx_tokens: Total number of context tokens across all sequences in current batch cached_token_lens: Cached KV token lengths [num_contexts], dtype=torch.int32 (optional) + kv_lens: KV token lengths [num_contexts], dtype=torch.int32 (optional) + compress_ratio: Compression ratio for KV tokens Returns: cu_seqlen_ks: Start index in KV for each Q token [num_ctx_tokens] @@ -387,7 +442,9 @@ def compute_cu_seqlen_kv_bounds_with_cache( """ device = seq_lens.device # Total KV lengths per request - kv_lens = seq_lens if cached_token_lens is None else cached_token_lens + seq_lens # [num_contexts] + if kv_lens is None: + kv_lens = seq_lens if cached_token_lens is None else cached_token_lens + seq_lens # [num_contexts] + kv_lens = kv_lens // compress_ratio # Cumulative KV offsets: where each request's KV sequence starts in global KV space cu_kv_offsets = torch.cat([ @@ -418,9 +475,11 @@ def compute_cu_seqlen_kv_bounds_with_cache( if cached_token_lens is not None: cached_per_token = torch.repeat_interleave(cached_token_lens, seq_lens) # [num_ctx_tokens] - cu_seqlen_ke = cu_seqlen_ks + cached_per_token + local_q_positions + 1 # [num_ctx_tokens] + cu_seqlen_ke = cu_seqlen_ks + (cached_per_token + local_q_positions + + 1) // compress_ratio # [num_ctx_tokens] else: - cu_seqlen_ke = cu_seqlen_ks + local_q_positions + 1 # [num_ctx_tokens] + cu_seqlen_ke = cu_seqlen_ks + (local_q_positions + + 1) // compress_ratio # [num_ctx_tokens] return cu_seqlen_ks, cu_seqlen_ke @@ -447,10 +506,16 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # 1. Request-level: Pack multiple small requests into one chunk (up to indexer_max_chunk_size) # 2. Intra-request: Split large requests into Q-blocks when seq_len > max_chunk_size indexer_max_chunk_size: int - # Topk for sparse MLA + # TopK for static token sparse attention num_sparse_topk: int + # TopK for dynamic sparse MLA + sparse_mla_topk: int # max number of draft tokens max_draft_tokens: int = 0 + # Indexer head dimension + indexer_head_dim: int = 128 + # Indexer quant block size + indexer_quant_block_size: int = 128 # Enable indexer skip for short sequences enable_indexer_skip: bool = False # Whether skip the indexer for context requests @@ -472,6 +537,11 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # (otherwise the populated buffers would mismatch the kernel reshape). dsl_expand_factor: int = 1 dsl_atom: int = 1 + # Compression ratio for KV tokens + compress_ratios: List[int] = [1] + # Number of compressed KV tokens for context requests + num_ctx_kv_tokens: int = 0 + gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None def __init__(self, *args, **kwargs): """Initialize DSA metadata with SM count and indexer chunk size.""" @@ -499,26 +569,210 @@ def __init__(self, *args, **kwargs): def __post_init__(self): """Allocate indexer K-cache buffers and heuristic TopK metadata.""" super().__post_init__() - assert isinstance(self.kv_cache_manager, DSACacheManager), \ - f"DSAtrtllmAttentionMetadata requires DSACacheManager, got {type(self.kv_cache_manager)}" + if not isinstance(self.kv_cache_manager, DSACacheManager): + has_deepseek_v4_cache_interface = all( + hasattr(self.kv_cache_manager, attr) + for attr in ("compressed_block_sizes", "get_cache_indices")) + assert has_deepseek_v4_cache_interface, ( + "DSAtrtllmAttentionMetadata requires DSACacheManager-compatible " + f"cache manager, got {type(self.kv_cache_manager)}") self.num_sparse_topk = self.sparse_attention_config.index_topk + self.sparse_mla_topk = self.num_sparse_topk + self.indexer_head_dim = self.sparse_attention_config.index_head_dim + self.indexer_quant_block_size = 128 self.enable_indexer_skip = self.sparse_attention_config.skip_indexer_for_short_seqs capture_graph = self.is_cuda_graph + # Get compression ratio from sparse attention config + if hasattr(self.sparse_attention_config, 'compress_ratios'): + self.compress_ratios = self.sparse_attention_config.compress_ratios + + # Effective tokens-per-block for the indexer k-cache slot mapping. + # DeepSeek-V4's indexer cache uses layer-dependent compressed block sizes + # (tokens_per_block // compress_ratio), so slot mappings must be built + # against that stride — not kv_cache_manager.tokens_per_block directly. + tpb = self.kv_cache_manager.tokens_per_block + self._indexer_compress_ratio = _select_indexer_compress_ratio( + self.compress_ratios) + if hasattr(self.kv_cache_manager, 'compressed_block_sizes'): + tpb = tpb // _effective_compress_ratio_divisor( + self._indexer_compress_ratio) + self._tokens_per_block = tpb + + self.create_buffers_for_mla_rope_append(capture_graph=capture_graph) + self.create_buffers_for_indexer(capture_graph=capture_graph) - self.indexer_k_cache_block_offsets = self.get_empty( - self.cuda_graph_buffers, - [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], - cache_name="indexer_k_cache_block_offsets", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.host_indexer_k_cache_block_offsets = torch.zeros_like( - self.indexer_k_cache_block_offsets, + def prepare(self): + super().prepare() + self._invalidate_pool_view_cache() + + # Get kv lengths + assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" + cached_token_lens = torch.tensor( + self.kv_cache_params.num_cached_tokens_per_seq, + dtype=torch.int, device='cpu', - pin_memory=prefer_pinned(), ) + if self.enable_helix: + # For Helix CP, inactive ranks only attend to previously cached + # tokens (no new token appended), while active ranks add new tokens. + # This mirrors the kv_lens logic in TrtllmAttentionMetadata.prepare(). + active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] + kv_lens = cached_token_lens.clone() + kv_lens[active_rank] += self.seq_lens_kv[active_rank] + else: + kv_lens = cached_token_lens + self.seq_lens_kv + + # For mla_rope_append_paged_kv_assign_q + self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) + + # Prepare to support skip indexer + self.prepare_for_skip_indexer(kv_lens) + + # For indices conversion + self.prepare_for_indices_conversion() + # For indexer k cache + self.prepare_for_indexer_k_cache() + + # For spec decode + self.prepare_for_spec_decode(kv_lens) + + # Prepare metadata for indexer + Indexer.prepare(metadata=self) + + def get_indexer_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: + if self._indexer_compress_ratio <= 1: + return kv_lens + return kv_lens // self._indexer_compress_ratio + + def get_indexer_max_seq_len(self) -> int: + if self._indexer_compress_ratio <= 1: + return self.kv_cache_manager.max_seq_len + return max( + 1, + self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) + + def on_update_kv_lens(self): + # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. + # Especially for the changes in the _preprocess_inputs() of model_engine.py. + # + # NOTE: + # In overlap scheduler + speculative decoding, kv_lens_cuda can be corrected at runtime + # (inside _preprocess_inputs) to account for variable accepted tokens. The indexer + # slot_mapping_* buffers also depend on these effective cached lengths. If we do not + # refresh slot mappings here, indexer K-cache updates can be written with stale offsets. + + # _preprocess_inputs() also uses this as a general hook to "invalidate per-forward-pass + # caches so they are recomputed (and captured) on every _forward_step". Invalidate the + # pool_view cache here so it is recomputed on the next + # transform_local_topk_and_prepare_pool_view() call. + self._invalidate_pool_view_cache() + + if self.kv_cache_manager is not None and self.num_tokens > 0: + seq_lens = self.seq_lens_cuda[:self.num_seqs] + # Runtime cached lengths after overlap/spec-dec correction. + start_positions = self.kv_lens_cuda[:self.num_seqs] - seq_lens + + # Reuse request-per-token mapping prepared in metadata.prepare(). + # This avoids repeat_interleave in graph-capture mode. + req_indices = self.req_idx_per_token[:self.num_tokens].to( + dtype=torch.int64) + seq_starts = torch.cumsum( + seq_lens, dim=0, dtype=torch.int64) - seq_lens.to(torch.int64) + token_offsets = torch.arange( + self.num_tokens, device=seq_lens.device, + dtype=torch.int64) - seq_starts[req_indices] + + global_positions = start_positions[req_indices] + token_offsets + # Honor MXFP4 indexer K cache layout (½ byte per value vs FP8's + # 1 byte) when the cache manager exposes a use_fp4 flag. + index_head_dim = self.kv_cache_manager.index_head_dim + use_fp4 = getattr(self.kv_cache_manager, 'use_fp4', False) + data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim + fp8_indices, scale_indices = _compute_slot_mappings( + global_positions, + self.indexer_k_cache_block_offsets, + req_indices, + index_head_dim, + self._tokens_per_block, + self.kv_cache_manager.quant_block_size, + data_bytes_per_token=data_bytes_per_token, + ) + self.slot_mapping_fp8[:self.num_tokens] = fp8_indices + self.slot_mapping_scale[:self.num_tokens] = scale_indices + + if self.num_generations > 0: + torch.cumsum( + self.kv_lens_cuda[self.num_contexts:self. + num_seqs], # num_contexts should be 0 + dim=0, + dtype=torch.int64, + out=self.gen_kv_indptr[1:self.num_generations + 1]) + torch.cumsum( + (self.kv_lens_cuda[self.num_contexts:self.num_seqs] - + self.seq_lens_cuda[self.num_contexts:self.num_seqs]), + dim=0, + dtype=torch.int64, + out=self.gen_cached_token_indptr[1:self.num_generations + 1]) + gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] + gen_indexer_kv_lens = self.get_indexer_kv_lens(gen_kv_lens) + self.gen_indexer_kv_lens_cuda_runtime = gen_indexer_kv_lens + next_n_cap = self.kv_lens_cuda_2d.shape[1] + self.kv_lens_cuda_2d[:self.num_generations, :next_n_cap].copy_( + gen_indexer_kv_lens.unsqueeze(-1).expand(-1, next_n_cap)) + scheduler_metadata_buffer = get_paged_mqa_logits_metadata( + gen_indexer_kv_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, + self.num_sms) + self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, + non_blocking=True) + if (self.max_draft_tokens > 0 + and not self.use_expanded_buffers_for_mtp): + scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( + self.kv_lens_cuda_2d[:self.num_generations, :next_n_cap], + _DG_SCHEDULE_BLOCK_KV, self.num_sms) + self.scheduler_metadata_buffer_full_next_n.copy_( + scheduler_metadata_buffer_full_next_n, non_blocking=True) + if self.use_expanded_buffers_for_mtp: + num_draft_tokens = 1 + self.max_draft_tokens + num_tokens = self.num_generations * num_draft_tokens + kv_lens_expanded = torch.stack([gen_indexer_kv_lens] * + num_draft_tokens, + dim=0) + self.kv_lens_expanded_cuda[:num_tokens] = \ + kv_lens_expanded.transpose(0, 1).contiguous().flatten() + scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( + self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + _DG_SCHEDULE_BLOCK_KV, self.num_sms) + self.scheduler_metadata_buffer_expanded.copy_( + scheduler_metadata_buffer_expanded, non_blocking=True) + if self.expand_for_dsl and self.dsl_expand_factor > 1: + expand_factor = self.dsl_expand_factor + num_tokens = self.num_generations * expand_factor + gen_kv_lens_expanded = gen_indexer_kv_lens.repeat_interleave( + expand_factor) + self.kv_lens_expanded_cuda[:num_tokens].copy_( + gen_kv_lens_expanded) + scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( + self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + _DG_SCHEDULE_BLOCK_KV, self.num_sms) + self.scheduler_metadata_buffer_expanded.copy_( + scheduler_metadata_buffer_expanded, non_blocking=True) + self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) + + def update_for_spec_dec(self): + super().update_for_spec_dec() + # host + self.max_ctx_kv_len = 0 + self.num_ctx_cached_tokens = 0 + self.max_gen_seq_len = 1 + + # device + self.on_update_kv_lens() + + # Create buffers for mla_rope_append_paged_kv_assign_q + def create_buffers_for_mla_rope_append(self, capture_graph=False): + # New context buffers for dsa if not self.enable_context_mla_with_cached_kv: self.ctx_cached_token_indptr = self.get_empty( self.cuda_graph_buffers, @@ -545,12 +799,6 @@ def __post_init__(self): pin_memory=prefer_pinned(), ) - # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. - # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory. - if self.enable_context_mla_with_cached_kv: - self.slot_mapping_fp8_fullkv = None - self.slot_mapping_scale_fullkv = None - # New generation buffers for dsa self.gen_cached_token_indptr = self.get_empty( self.cuda_graph_buffers, @@ -576,6 +824,21 @@ def __post_init__(self): device='cpu', pin_memory=prefer_pinned(), ) + + def create_buffers_for_indexer(self, capture_graph=False): + self.indexer_k_cache_block_offsets = self.get_empty( + self.cuda_graph_buffers, + [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], + cache_name="indexer_k_cache_block_offsets", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.host_indexer_k_cache_block_offsets = torch.zeros_like( + self.indexer_k_cache_block_offsets, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Indexer metadata # Separate slot mappings for non-interleaved layout (flat byte indices) self.slot_mapping_fp8 = self.get_empty( @@ -602,6 +865,11 @@ def __post_init__(self): device='cpu', pin_memory=prefer_pinned(), ) + # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. + # These buffers will be allocated dynamically in Indexer.prepare() based on actual total_kv_len to save memory. + if self.enable_context_mla_with_cached_kv: + self.slot_mapping_fp8_fullkv = None + self.slot_mapping_scale_fullkv = None # Per-token request index buffer for topk_indices conversion self.req_idx_per_token = self.get_empty( self.cuda_graph_buffers, @@ -610,10 +878,12 @@ def __post_init__(self): dtype=torch.int32, capture_graph=capture_graph, ) + self.host_req_idx_per_token = torch.empty_like( + self.req_idx_per_token, device='cpu', pin_memory=prefer_pinned()) # Block table for topk_indices conversion (shared for context and generation) self.block_table = self.get_empty( self.cuda_graph_buffers, - (self.max_num_requests, self.kv_cache_manager.max_blocks_per_seq), + [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], cache_name="block_table", dtype=torch.int32, capture_graph=capture_graph, @@ -710,6 +980,34 @@ def __post_init__(self): capture_graph=capture_graph, ) + # Persistent scratch for Radix-split-work indexer path (blocks_per_row > 1). + # Mirrors the fix the Heuristic path applied above: per-call th::empty + # inside indexer_topk_decode produces stale pointers under CUDA Graph + # replay when the caching allocator is perturbed by chunked prefill at + # high CONC. Sized to the worst case kMaxBlocksPerRowDecode=10 from + # cpp/tensorrt_llm/kernels/indexerTopK.cu. Allocated unconditionally: + # even with enable_heuristic_topk=True, the dispatcher can fall back + # to Radix when canUseHeuristic returns False (small numColumns, etc.). + _radix_max_blocks_per_row = 10 + _radix_max_gen_tokens = self.max_num_sequences * (1 + + self.max_draft_tokens) + self.radix_aux_indices = self.get_empty( + self.cuda_graph_buffers, + (_radix_max_gen_tokens, _radix_max_blocks_per_row, + self.num_sparse_topk), + cache_name="radix_aux_indices", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.radix_aux_logits = self.get_empty( + self.cuda_graph_buffers, + (_radix_max_gen_tokens, _radix_max_blocks_per_row, + self.num_sparse_topk), + cache_name="radix_aux_logits", + dtype=torch.float32, + capture_graph=capture_graph, + ) + # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) @@ -811,6 +1109,34 @@ def update_spec_dec_param( capture_graph=capture_graph, ) + def _get_pool_block_indices(self) -> torch.Tensor: + """Extract memory pool block indices from host_kv_cache_block_offsets. + + The C++ setOffsets() encodes offsets as: + encoded = memPoolBlockIndex * numLayers * kvFactor + For SELFKONLY (MLA/DSA), kvFactor=1, so: + memPoolBlockIndex = encoded // num_local_layers + + Returns a (num_seqs, max_blocks_per_seq) int32 CPU tensor with valid + pool indices clamped to [0, blocks_in_primary_pool - 1]. + """ + num_local_layers = self.kv_cache_manager.num_local_layers + max_pool_idx = self.kv_cache_manager.blocks_in_primary_pool - 1 + # DSA uses SELFKONLY mode where only key cache is stored (kv_factor=1). + # host_kv_cache_block_offsets shape: (num_pools, max_batch*beam, 2, max_blocks_per_seq) + # Note: dim=2 is always 2 in the tensor layout (K and V slots), but for + # SELFKONLY only the K slot (index 0) contains valid data. + assert self.kv_cache_manager.kv_factor == 1, \ + f"DSA requires SELFKONLY mode (kv_factor=1), got kv_factor={self.kv_cache_manager.kv_factor}" + # Pool 0, first num_seqs entries, field 0 (key offsets) + encoded = self.kv_cache_manager.host_kv_cache_block_offsets[ + 0, :self.num_seqs, 0, :] + pool_indices = encoded // num_local_layers + # Clamp for safety: handles garbage padding from torch.empty in uninitialized slots + pool_indices = pool_indices.clamp(min=0, + max=max_pool_idx).to(torch.int32) + return pool_indices + def _invalidate_pool_view_cache(self): """Invalidate the cached pool view and related step-invariant values. @@ -914,188 +1240,23 @@ def prepare_dense_topk_indices(self, self.host_topk_indices_buffer[gen_range, :], non_blocking=True) - def _get_pool_block_indices(self) -> torch.Tensor: - """Extract memory pool block indices from host_kv_cache_block_offsets. - - The C++ setOffsets() encodes offsets as: - encoded = memPoolBlockIndex * numLayers * kvFactor - For SELFKONLY (MLA/DSA), kvFactor=1, so: - memPoolBlockIndex = encoded // num_local_layers - - Returns a (num_seqs, max_blocks_per_seq) int32 CPU tensor with valid - pool indices clamped to [0, blocks_in_primary_pool - 1]. - """ - num_local_layers = self.kv_cache_manager.num_local_layers - max_pool_idx = self.kv_cache_manager.blocks_in_primary_pool - 1 - # DSA uses SELFKONLY mode where only key cache is stored (kv_factor=1). - # host_kv_cache_block_offsets shape: (num_pools, max_batch*beam, 2, max_blocks_per_seq) - # Note: dim=2 is always 2 in the tensor layout (K and V slots), but for - # SELFKONLY only the K slot (index 0) contains valid data. - assert self.kv_cache_manager.kv_factor == 1, \ - f"DSA requires SELFKONLY mode (kv_factor=1), got kv_factor={self.kv_cache_manager.kv_factor}" - # Pool 0, first num_seqs entries, field 0 (key offsets) - encoded = self.kv_cache_manager.host_kv_cache_block_offsets[ - 0, :self.num_seqs, 0, :] - pool_indices = encoded // num_local_layers - # Clamp for safety: handles garbage padding from torch.empty in uninitialized slots - pool_indices = pool_indices.clamp(min=0, - max=max_pool_idx).to(torch.int32) - return pool_indices - - def prepare(self): - """Prepare DSA metadata: compute slot mappings, block tables, and prefill chunks.""" - super().prepare() - self._invalidate_pool_view_cache() - - # Get kv lengths - assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" - cached_token_lens = torch.tensor( - self.kv_cache_params.num_cached_tokens_per_seq, - dtype=torch.int, - device='cpu', - ) - if self.enable_helix: - # For Helix CP, inactive ranks only attend to previously cached - # tokens (no new token appended), while active ranks add new tokens. - # This mirrors the kv_lens logic in TrtllmAttentionMetadata.prepare(). - active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] - kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.seq_lens_kv[active_rank] - else: - kv_lens = cached_token_lens + self.seq_lens_kv - - # Prepare to support skip indexer - num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens - if self.num_contexts > 0 and self.enable_indexer_skip: - # Minus the number of extra KV tokens because when using one-model MTP, the - # draft layers needs more KV tokens for the next draft forwards. - self.skip_indexer_for_ctx_reqs = kv_lens[:self.num_contexts].max( - ).item() <= self.num_sparse_topk - num_extra_kv_tokens - else: - self.skip_indexer_for_ctx_reqs = False - - if self.num_generations > 0 and self.enable_indexer_skip: - # Minus the number of extra KV tokens because when using one-model MTP, the - # draft layers needs more KV tokens for the next draft forwards. - self.skip_indexer_for_gen_reqs = kv_lens[ - self.num_contexts:self.num_seqs].max().item( - ) <= self.num_sparse_topk - num_extra_kv_tokens - else: - self.skip_indexer_for_gen_reqs = False - self.prepare_dense_topk_indices(kv_lens) - - # Build indexer_k_cache_block_offsets using pool block indices derived - # from host_kv_cache_block_offsets (populated by super().prepare()). - # This correctly resolves block IDs to memory pool indices, which is - # required when host cache offload is enabled (block IDs != pool indices - # for onboarded secondary blocks). - if self.kv_cache_manager is not None: - pool_indices = self._get_pool_block_indices() - self.host_indexer_k_cache_block_offsets[:self.num_seqs].copy_( - pool_indices) - self.indexer_k_cache_block_offsets[:self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[:self.num_seqs], - non_blocking=True) - # Safety clamp: prevent OOB from CUDA graph padding entries which - # may contain stale negative or out-of-range values after block - # eviction/onboarding with host cache offload. - self.indexer_k_cache_block_offsets.clamp_(min=0) - - # Build req_idx_per_token for topk_indices conversion - host_req_idx_per_token = torch.repeat_interleave(torch.arange( - self.num_seqs, dtype=torch.int32), - self.seq_lens, - dim=0) - self.req_idx_per_token[:self.num_tokens].copy_(host_req_idx_per_token, - non_blocking=True) - - # Build block_table for topk_indices conversion (actual block allocation) - if self.kv_cache_manager is not None: - tokens_per_block = self.kv_cache_manager.tokens_per_block - num_blocks_per_seq = (kv_lens[:self.num_seqs] + tokens_per_block - - 1) // tokens_per_block - max_blocks_used = num_blocks_per_seq.max().item( - ) if self.num_seqs > 0 else 1 - # pool_indices already has correct values; set padding to -1 - host_block_table = pool_indices[:, :max_blocks_used].clone() - for i in range(self.num_seqs): - if num_blocks_per_seq[i] < max_blocks_used: - host_block_table[i, num_blocks_per_seq[i]:] = -1 - # Copy to GPU - self.block_table[:self.num_seqs, :max_blocks_used].copy_( - host_block_table, non_blocking=True) - - # For mla_rope_append_paged_kv_assign_q - if self.num_contexts > 0: - self.num_ctx_cached_tokens = cached_token_lens[:self. - num_contexts].sum( - ).item() - self.max_ctx_kv_len = kv_lens[:self.num_contexts].max().item() - self.max_ctx_seq_len = self.seq_lens[:self.num_contexts].max().item( - ) - # context cached token indptr - torch.cumsum( - cached_token_lens[:self.num_contexts], - dim=0, - dtype=torch.int64, - out=self.host_ctx_cached_token_indptr[1:self.num_contexts + 1]) - self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_( - self.host_ctx_cached_token_indptr[:self.num_contexts + 1], - non_blocking=True) - # context kv indptr - torch.cumsum(kv_lens[:self.num_contexts], - dim=0, - dtype=torch.int64, - out=self.host_ctx_kv_indptr[1:self.num_contexts + 1]) - self.ctx_kv_indptr[:self.num_contexts + 1].copy_( - self.host_ctx_kv_indptr[:self.num_contexts + 1], - non_blocking=True) - else: - self.num_ctx_cached_tokens = 0 - self.max_ctx_kv_len = 0 - self.max_ctx_seq_len = 0 - - if self.num_generations > 0: - self.max_gen_seq_len = self.seq_lens[self.num_contexts:self. - num_seqs].max().item() - # generation cached token indptr - torch.cumsum( - cached_token_lens[self.num_contexts:self.num_seqs], - dim=0, - dtype=torch.int64, - out=self.host_gen_cached_token_indptr[1:self.num_generations + - 1]) - self.gen_cached_token_indptr[:self.num_generations + 1].copy_( - self.host_gen_cached_token_indptr[:self.num_generations + 1], - non_blocking=True) - # generation kv indptr - torch.cumsum(kv_lens[self.num_contexts:self.num_seqs], - dim=0, - dtype=torch.int64, - out=self.host_gen_kv_indptr[1:self.num_generations + - 1]) - self.gen_kv_indptr[:self.num_generations + 1].copy_( - self.host_gen_kv_indptr[:self.num_generations + 1], - non_blocking=True) - else: - self.max_gen_seq_len = 0 - + def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # Because the fp8_paged_mqa_logits only supports seq_len == 1/2/4 (i.e., max_draft_tokens == 0/1/3) on sm100, and # seq_len == 1/2 (i.e., max_draft_tokens == 0/1) on sm90, for other cases, we need to flatten the q tensor and # expand the kv_lens and block_table for MTP support. - # The CuTe DSL kernel supports arbitrary next_n natively, so it never needs expansion. # TODO: # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90. # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. - _use_dsl = self.sparse_attention_config.use_cute_dsl_paged_mqa_logits - self.use_expanded_buffers_for_mtp = (not _use_dsl and ( - (self.max_draft_tokens > 1 and get_sm_version() == 90) or - ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) - and get_sm_version() >= 100))) + use_dsl = self.sparse_attention_config.use_cute_dsl_paged_mqa_logits + self.use_expanded_buffers_for_mtp = (not use_dsl and ( + (self.max_draft_tokens > 1 and get_sm_version() == 90) + or ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) + and get_sm_version() >= 100))) if self.use_expanded_buffers_for_mtp: # Expand kv_lens_cuda (only generation) num_tokens = self.num_generations * (1 + self.max_draft_tokens) - gen_kv_lens = kv_lens[self.num_contexts:self.num_seqs] + gen_kv_lens = self.get_indexer_kv_lens( + kv_lens[self.num_contexts:self.num_seqs]) gen_kv_lens_expanded = torch.stack([gen_kv_lens] * (1 + self.max_draft_tokens), dim=0) @@ -1107,7 +1268,7 @@ def prepare(self): # Expand indexer_k_cache_block_offsets (only generation) # host_indexer_k_cache_block_offsets already contains correct pool - # indices from _get_pool_block_indices() above. + # indices from _get_pool_block_indices() in prepare_for_indexer_k_cache(). if self.kv_cache_manager is not None and self.num_generations > 0: max_len = self.host_indexer_k_cache_block_offsets.shape[1] gen_block_tensor = self.host_indexer_k_cache_block_offsets[ @@ -1121,32 +1282,15 @@ def prepare(self): non_blocking=True) self.block_table_expanded.clamp_(min=0) - # CuTe DSL FP4 paged MQA logits kernel natively supports - # next_n ∈ {1, 2, 3} only. For next_n ≥ 4 atom-split is mandatory. - # For next_n ∈ {2, 3} atom-split is also beneficial when the wave-aware - # picker decides more SM utilization outweighs the (expand_factor)x - # HBM cost (e.g., low batch with idle SMs). The expanded buffers - # (`kv_lens_expanded_cuda` / `block_table_expanded` / - # `scheduler_metadata_buffer_expanded`, all sized for the worst-case - # `1+max_draft_tokens` factor) are reused: under `_use_dsl=True` the - # existing FP8 DG expand path never writes to them (gated on - # `not _use_dsl`), so there's no conflict. - # Trigger relaxed to `max_draft_tokens >= 1` (i.e., next_n >= 2) so the - # picker can choose to expand when waves vs HBM trade-off favors it. - # Trigger atom-split for both FP4 and FP8 DSL paths. FP4 kernel - # supports atom ∈ {1, 2, 3}; FP8 supports {1, 2, 3, 4}. Picker is - # given the appropriate kernel_atoms set so it only enumerates - # decompositions the kernel can handle. - self.expand_for_dsl = (_use_dsl and self.kv_cache_manager is not None + self.expand_for_dsl = (use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1) if self.expand_for_dsl and self.num_generations > 0: next_n = 1 + self.max_draft_tokens kernel_atoms = (1, 2, 3) if self.kv_cache_manager.use_fp4 else (1, 2, 3, 4) - # Wave-aware picker. max_ctx ≈ longest gen kv_len (decode iter - # upper-bound observed at this prepare). num_sms is hardware. - gen_kv_lens = kv_lens[self.num_contexts:self.num_seqs] + gen_kv_lens = self.get_indexer_kv_lens( + kv_lens[self.num_contexts:self.num_seqs]) max_ctx = int( gen_kv_lens.max().item()) if gen_kv_lens.numel() else 0 expand_factor, atom = _pick_dsl_expand( @@ -1158,8 +1302,6 @@ def prepare(self): ) self.dsl_expand_factor = expand_factor self.dsl_atom = atom - # Only populate when picker chose to actually split (factor > 1); - # factor=1 means kernel-native, no expansion needed. if expand_factor > 1: num_tokens = self.num_generations * expand_factor gen_kv_lens_expanded = gen_kv_lens.repeat_interleave( @@ -1168,172 +1310,155 @@ def prepare(self): gen_kv_lens_expanded) self.kv_lens_expanded_cuda[:num_tokens].copy_( self.kv_lens_expanded_host[:num_tokens], non_blocking=True) - if self.kv_cache_manager is not None: - max_len = self.host_indexer_k_cache_block_offsets.shape[1] - gen_block_tensor = self.host_indexer_k_cache_block_offsets[ - self.num_contexts:self.num_seqs, :max_len] - expanded_blocks = gen_block_tensor.repeat_interleave( - expand_factor, dim=0) - self.host_block_table_expanded[:num_tokens, :max_len].copy_( - expanded_blocks, non_blocking=True) - self.block_table_expanded[:num_tokens].copy_( - self.host_block_table_expanded[:num_tokens], - non_blocking=True) - self.block_table_expanded.clamp_(min=0) + max_len = self.host_indexer_k_cache_block_offsets.shape[1] + gen_block_tensor = self.host_indexer_k_cache_block_offsets[ + self.num_contexts:self.num_seqs, :max_len] + expanded_blocks = gen_block_tensor.repeat_interleave( + expand_factor, dim=0) + self.host_block_table_expanded[:num_tokens, :max_len].copy_( + expanded_blocks, non_blocking=True) + self.block_table_expanded[:num_tokens].copy_( + self.host_block_table_expanded[:num_tokens], + non_blocking=True) + self.block_table_expanded.clamp_(min=0) else: - # Reset cache; forward path uses kernel-native next_n. self.dsl_expand_factor = 1 self.dsl_atom = 1 + self.max_draft_tokens - # Prepare metadata for indexer - Indexer.prepare(metadata=self) - - def on_update_kv_lens(self): - """Refresh indexer slot mappings after KV lengths change at runtime.""" - # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. - # Especially for the changes in the _preprocess_inputs() of model_engine.py. - # - # NOTE: - # In overlap scheduler + speculative decoding, kv_lens_cuda can be corrected at runtime - # (inside _preprocess_inputs) to account for variable accepted tokens. The indexer - # slot_mapping_* buffers also depend on these effective cached lengths. If we do not - # refresh slot mappings here, indexer K-cache updates can be written with stale offsets. - - # _preprocess_inputs() also uses this as a general hook to "invalidate per-forward-pass - # caches so they are recomputed (and captured) on every _forward_step". Invalidate the - # pool_view cache here so it is recomputed on the next - # transform_local_topk_and_prepare_pool_view() call. - self._invalidate_pool_view_cache() - - if self.kv_cache_manager is not None and self.num_tokens > 0: - seq_lens = self.seq_lens_cuda[:self.num_seqs] - # Runtime cached lengths after overlap/spec-dec correction. - start_positions = self.kv_lens_cuda[:self.num_seqs] - seq_lens - - # Reuse request-per-token mapping prepared in metadata.prepare(). - # This avoids repeat_interleave in graph-capture mode. - req_indices = self.req_idx_per_token[:self.num_tokens].to( - dtype=torch.int64) - seq_starts = torch.cumsum( - seq_lens, dim=0, dtype=torch.int64) - seq_lens.to(torch.int64) - token_offsets = torch.arange( - self.num_tokens, device=seq_lens.device, - dtype=torch.int64) - seq_starts[req_indices] + def prepare_for_indexer_k_cache(self): + # Build indexer_k_cache_block_offsets using pool block indices derived + # from host_kv_cache_block_offsets (populated by super().prepare()). + # This correctly resolves block IDs to memory pool indices, which is + # required when host cache offload is enabled (block IDs != pool indices + # for onboarded secondary blocks). + if self.kv_cache_manager is None: + return + pool_indices = self._get_pool_block_indices() + self.host_indexer_k_cache_block_offsets[:self.num_seqs].copy_( + pool_indices) + self.indexer_k_cache_block_offsets[:self.num_seqs].copy_( + self.host_indexer_k_cache_block_offsets[:self.num_seqs], + non_blocking=True) + # Safety clamp: prevent OOB from CUDA graph padding entries which + # may contain stale negative or out-of-range values after block + # eviction/onboarding with host cache offload. + self.indexer_k_cache_block_offsets.clamp_(min=0) - global_positions = start_positions[req_indices] + token_offsets - # Under FP4 the indexer cache stores two E2M1 codes per byte, so - # the per-token data footprint is head_dim // 2; otherwise it is - # head_dim (one FP8 byte per element). Feed the real byte count - # into _compute_slot_mappings so scatter/gather see offsets that - # match the pool layout produced by createIndexerKCachePools. - use_fp4 = self.kv_cache_manager.use_fp4 - index_head_dim = self.kv_cache_manager.index_head_dim - data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim - fp8_indices, scale_indices = _compute_slot_mappings( - global_positions, - self.indexer_k_cache_block_offsets, - req_indices, - index_head_dim, - self.kv_cache_manager.tokens_per_block, - self.kv_cache_manager.quant_block_size, - data_bytes_per_token=data_bytes_per_token, - ) - self.slot_mapping_fp8[:self.num_tokens] = fp8_indices - self.slot_mapping_scale[:self.num_tokens] = scale_indices + # Build block_table for topk_indices conversion (actual block allocation) + cached_token_lens = torch.tensor( + self.kv_cache_params.num_cached_tokens_per_seq, + dtype=torch.int, + device='cpu', + ) + if self.enable_helix: + active_rank = ~self.helix_is_inactive_rank_cpu[:self.num_seqs] + kv_lens = cached_token_lens.clone() + kv_lens[active_rank] += self.seq_lens_kv[active_rank] + else: + kv_lens = cached_token_lens + self.seq_lens_kv + tokens_per_block = self.kv_cache_manager.tokens_per_block + num_blocks_per_seq = (kv_lens[:self.num_seqs] + tokens_per_block - + 1) // tokens_per_block + max_blocks_used = num_blocks_per_seq.max().item( + ) if self.num_seqs > 0 else 1 + # pool_indices already has correct values; set padding to -1 + host_block_table = pool_indices[:, :max_blocks_used].clone() + for i in range(self.num_seqs): + if num_blocks_per_seq[i] < max_blocks_used: + host_block_table[i, num_blocks_per_seq[i]:] = -1 + # Copy to GPU + self.block_table[:self.num_seqs, :max_blocks_used].copy_( + host_block_table, non_blocking=True) + + def prepare_for_skip_indexer(self, kv_lens: torch.Tensor): + num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens + if self.num_contexts > 0 and self.enable_indexer_skip: + # Minus the number of extra KV tokens because when using one-model MTP, the + # draft layers needs more KV tokens for the next draft forwards. + self.skip_indexer_for_ctx_reqs = kv_lens[:self.num_contexts].max( + ).item() <= self.num_sparse_topk - num_extra_kv_tokens + else: + self.skip_indexer_for_ctx_reqs = False - if self.num_generations > 0: + if self.num_generations > 0 and self.enable_indexer_skip: + # Minus the number of extra KV tokens because when using one-model MTP, the + # draft layers needs more KV tokens for the next draft forwards. + self.skip_indexer_for_gen_reqs = kv_lens[ + self.num_contexts:self.num_seqs].max().item( + ) <= self.num_sparse_topk - num_extra_kv_tokens + else: + self.skip_indexer_for_gen_reqs = False + self.prepare_dense_topk_indices(kv_lens) + + def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, + kv_lens: torch.Tensor): + if self.num_contexts > 0: + self.num_ctx_cached_tokens = cached_token_lens[:self. + num_contexts].sum( + ).item() + self.max_ctx_kv_len = kv_lens[:self.num_contexts].max().item() + self.max_ctx_seq_len = self.seq_lens[:self.num_contexts].max().item( + ) + # context cached token indptr torch.cumsum( - self.kv_lens_cuda[self.num_contexts:self. - num_seqs], # num_contexts should be 0 + cached_token_lens[:self.num_contexts], dim=0, dtype=torch.int64, - out=self.gen_kv_indptr[1:self.num_generations + 1]) + out=self.host_ctx_cached_token_indptr[1:self.num_contexts + 1]) + self.ctx_cached_token_indptr[:self.num_contexts + 1].copy_( + self.host_ctx_cached_token_indptr[:self.num_contexts + 1], + non_blocking=True) + # context kv indptr + torch.cumsum(kv_lens[:self.num_contexts], + dim=0, + dtype=torch.int64, + out=self.host_ctx_kv_indptr[1:self.num_contexts + 1]) + self.ctx_kv_indptr[:self.num_contexts + 1].copy_( + self.host_ctx_kv_indptr[:self.num_contexts + 1], + non_blocking=True) + else: + self.num_ctx_cached_tokens = 0 + self.max_ctx_kv_len = 0 + self.max_ctx_seq_len = 0 + + if self.num_generations > 0: + self.max_gen_seq_len = self.seq_lens[self.num_contexts:self. + num_seqs].max().item() + # generation cached token indptr torch.cumsum( - (self.kv_lens_cuda[self.num_contexts:self.num_seqs] - - self.seq_lens_cuda[self.num_contexts:self.num_seqs]), + cached_token_lens[self.num_contexts:self.num_seqs], dim=0, dtype=torch.int64, - out=self.gen_cached_token_indptr[1:self.num_generations + 1]) - # Write 2D kv_lens in-place (broadcast same kv_len across next_n - # positions). .expand() returns a view and .copy_() writes into the - # pre-allocated destination, so this is CUDA-graph-friendly. - gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] - next_n_cap = self.kv_lens_cuda_2d.shape[1] - self.kv_lens_cuda_2d[:self.num_generations, :next_n_cap].copy_( - gen_kv_lens.unsqueeze(-1).expand(-1, next_n_cap)) - # Build the next_n=1 schedule (used by MTP draft layers and any - # non-MTP forward). Reshape the contiguous gen slice of - # kv_lens_cuda to (num_gen, 1) — slicing kv_lens_cuda_2d's first - # column would be non-contiguous and would fail the metadata - # kernel's is_contiguous assertion. - context_lens_next_n1 = gen_kv_lens.view(-1, 1) - # `_DG_SCHEDULE_BLOCK_KV` (= 64) instead of cache `tokens_per_block`: - # see module-level constant comment for the SPLIT_KV=256 alignment. - scheduler_metadata_buffer = get_paged_mqa_logits_metadata( - context_lens_next_n1, _DG_SCHEDULE_BLOCK_KV, self.num_sms) - self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, - non_blocking=True) - # When MTP is on without the expanded-tokens path, also populate - # the full-next_n schedule for the main forward call. The metadata - # kernel reads next_n from context_lens.size(1), so we must pass - # the wider slice here. - if (self.max_draft_tokens > 0 - and not self.use_expanded_buffers_for_mtp): - context_lens_full_next_n = self.kv_lens_cuda_2d[:self. - num_generations, : - next_n_cap] - scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( - context_lens_full_next_n, _DG_SCHEDULE_BLOCK_KV, - self.num_sms) - self.scheduler_metadata_buffer_full_next_n.copy_( - scheduler_metadata_buffer_full_next_n, non_blocking=True) - if self.use_expanded_buffers_for_mtp: - num_draft_tokens = 1 + self.max_draft_tokens - num_tokens = self.num_generations * num_draft_tokens - kv_lens_expanded = torch.stack([gen_kv_lens] * num_draft_tokens, - dim=0) - self.kv_lens_expanded_cuda[:num_tokens] = \ - kv_lens_expanded.transpose(0, 1).contiguous().flatten() - # New API requires 2D; each expanded token becomes a (1,) row. - kv_lens_expanded_2d = self.kv_lens_expanded_cuda[: - num_tokens].view( - -1, 1) - scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - kv_lens_expanded_2d, _DG_SCHEDULE_BLOCK_KV, self.num_sms) - self.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - # DSL atom-split path: mirror the prepare()-time build so that - # overlap-scheduler / spec-dec runtime corrections to kv_lens_cuda - # propagate into kv_lens_expanded_cuda and the matching schedule. - # Reuse the cached (dsl_expand_factor, dsl_atom) — re-running the - # picker here would let the split decision drift between prepare - # and forward, breaking CUDA graph capture. - if self.expand_for_dsl and self.dsl_expand_factor > 1: - expand_factor = self.dsl_expand_factor - num_tokens = self.num_generations * expand_factor - gen_kv_lens_expanded = gen_kv_lens.repeat_interleave( - expand_factor) - self.kv_lens_expanded_cuda[:num_tokens].copy_( - gen_kv_lens_expanded) - kv_lens_expanded_2d = self.kv_lens_expanded_cuda[: - num_tokens].view( - -1, 1) - scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - kv_lens_expanded_2d, _DG_SCHEDULE_BLOCK_KV, self.num_sms) - self.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) - - def update_for_spec_dec(self): - """Reset context/generation counters and refresh slot mappings for speculative decoding.""" - super().update_for_spec_dec() - # host - self.max_ctx_kv_len = 0 - self.num_ctx_cached_tokens = 0 - self.max_gen_seq_len = 1 + out=self.host_gen_cached_token_indptr[1:self.num_generations + + 1]) + self.gen_cached_token_indptr[:self.num_generations + 1].copy_( + self.host_gen_cached_token_indptr[:self.num_generations + 1], + non_blocking=True) + # generation kv indptr + torch.cumsum(kv_lens[self.num_contexts:self.num_seqs], + dim=0, + dtype=torch.int64, + out=self.host_gen_kv_indptr[1:self.num_generations + + 1]) + self.gen_kv_indptr[:self.num_generations + 1].copy_( + self.host_gen_kv_indptr[:self.num_generations + 1], + non_blocking=True) + else: + self.max_gen_seq_len = 0 - # device - self.on_update_kv_lens() + def prepare_for_indices_conversion(self): + # Build req_idx_per_token for topk_indices conversion + # Use pinned staging buffer to avoid pageable H2D memcpy + self.host_req_idx_per_token[:self.num_tokens] = ( + torch.repeat_interleave( + torch.arange(self.num_seqs, dtype=torch.int32), + self.seq_lens, + dim=0, + )) + self.req_idx_per_token[:self.num_tokens].copy_( + self.host_req_idx_per_token[:self.num_tokens], + non_blocking=True, + ) @maybe_compile(dynamic=True) @@ -1366,6 +1491,75 @@ def _tf32_matmul_enabled(): torch.backends.cuda.matmul.allow_tf32 = prev +@dataclass +class IndexerParams: + """ + Parameters for indexer. + """ + num_contexts: int + num_generations: int + num_ctx_tokens: int + head_dim: int + quant_block_size: int + tokens_per_block: int + compress_ratio: int + request_ids: List[int] + num_past_tokens: List[int] + seq_lens: torch.Tensor + # Bytes of quantized data per token in the indexer K cache; defaults to + # head_dim (one FP8 byte per element). For MXFP4 use head_dim // 2. + data_bytes_per_token: Optional[int] = None + + def __post_init__(self): + # Pre-compute frequently used tensors once instead of on every property access + num_past_tokens_tensor = torch.tensor(self.num_past_tokens, + dtype=torch.int32) + self._num_past_tokens_tensor = num_past_tokens_tensor + compress_ratio = self.compress_ratio + self._cached_kv_tokens = num_past_tokens_tensor // compress_ratio + self._all_kv_tokens = (self.seq_lens + + num_past_tokens_tensor) // compress_ratio + self._new_kv_tokens = self._all_kv_tokens - self._cached_kv_tokens + self._kv_lens = self._all_kv_tokens + if self.data_bytes_per_token is None: + self.data_bytes_per_token = self.head_dim + self._scale_size = self.head_dim // self.quant_block_size * 4 + self._block_stride = self.tokens_per_block * ( + self.data_bytes_per_token + self._scale_size) + + @property + def batch_size(self): + return len(self.request_ids) + + @property + def kv_lens(self): + return self._kv_lens + + @property + def total_tokens(self): + return self.seq_lens.sum().item() + + @property + def cached_kv_tokens(self): + return self._cached_kv_tokens + + @property + def all_kv_tokens(self): + return self._all_kv_tokens + + @property + def new_kv_tokens(self): + return self._new_kv_tokens + + @property + def scale_size(self): + return self._scale_size + + @property + def block_stride(self): + return self._block_stride + + class Indexer(nn.Module): """DSA sparse attention indexer that selects top-K KV cache entries per token.""" @@ -1376,6 +1570,7 @@ def __init__(self, skip_create_weights_in_init: bool, sparse_attention_config: "SparseAttentionConfig", dtype: Optional[torch.dtype], + compress_ratio: int = 1, layer_idx: int = 0, aux_stream: Optional[torch.cuda.Stream] = None): """Initialize indexer with projection weights, norms, and TopK configuration.""" @@ -1387,6 +1582,7 @@ def __init__(self, self.head_dim = sparse_attention_config.index_head_dim # 128 self.index_topk = sparse_attention_config.index_topk # 2048 self.layer_idx = layer_idx + self.compress_ratio = compress_ratio self.wq_b = Linear( self.q_lora_rank, @@ -1429,13 +1625,14 @@ def __init__(self, self.softmax_scale = self.head_dim**-0.5 # TODO: make it configurable from hf config self.scale_fmt = "ue8m0" - # indexer_k_dtype controls both Q and K precision. DeepGEMM's - # fp8_fp4_mqa_logits / fp8_fp4_paged_mqa_logits kernels only dispatch - # to FP4xFP4 or FP8xFP8 (no mixed-precision variant). The DeepGEMM - # kernel asserts SM100 + head_dim=128 at launch time under FP4. - self.use_fp4 = sparse_attention_config.indexer_k_dtype == "fp4" self.aux_stream = aux_stream self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + # `indexer_k_dtype` is the single user-facing FP4 knob; the V4 + # config inherits this field from the DSA base config. DeepGEMM's + # fp8_fp4_mqa_logits / fp8_fp4_paged_mqa_logits kernels only + # dispatch to FP4xFP4 or FP8xFP8 (no mixed-precision variant) and + # assert SM100 + head_dim=128 at launch time under FP4. + self.use_fp4 = sparse_attention_config.indexer_k_dtype == "fp4" self.use_cute_dsl_topk = (sparse_attention_config.use_cute_dsl_topk and IS_CUTLASS_DSL_AVAILABLE) self.use_cute_dsl_paged_mqa_logits = ( @@ -1463,6 +1660,10 @@ def __init__(self, # attribute queries do not end up frozen into a captured graph. warmup_heuristic_topk_decode(top_k=self.index_topk) + # Fused wk + weights_proj weight for single FP32 cuBLAS GEMM + # (populated in post_load_weights; maps to TF32 tensor cores on Ampere+) + self._fused_wk_wp_weight: Optional[torch.Tensor] = None + def post_load_weights(self): """Fuse wk + weights_proj into single FP32 weight for F.linear GEMM under allow_tf32 (TF32 tensor cores on Ampere+).""" # wk: [head_dim, hidden_size] + weights_proj: [n_heads, hidden_size] @@ -1489,6 +1690,8 @@ def prepare_one_prefill_chunk( Note: Cached token counts are derived from metadata.host_ctx_cached_token_indptr """ device = metadata.cu_seqlen_ks.device + compress_ratio = _effective_compress_ratio_divisor( + _select_indexer_compress_ratio(metadata.compress_ratios)) if len(chunk_specs) == 1: # Single request or intra-request Q-block req_idx, token_start_in_req, token_end_in_req, req_cum_start = chunk_specs[ @@ -1500,24 +1703,32 @@ def prepare_one_prefill_chunk( metadata.host_ctx_cached_token_indptr[req_idx + 1] - metadata.host_ctx_cached_token_indptr[req_idx]).item() + # Total compressed KV tokens for this request + req_kv_len = (num_cached + token_end_in_req) // compress_ratio + # For intra-request chunks: Q block attends to all previous K in the request # Q tokens [token_start_in_req:token_end_in_req] within the request's current tokens - # K tokens [0:num_cached + token_end_in_req] within the request (causal attention) + # K tokens [0:req_kv_len] in compressed KV space cu_seqlen_ks = torch.zeros(num_q_tokens, dtype=torch.int32, device='cpu') - cu_seqlen_ke = torch.arange(token_start_in_req + 1, - token_end_in_req + 1, - dtype=torch.int32, - device='cpu') + num_cached + cu_seqlen_ke = (torch.arange(token_start_in_req + 1, + token_end_in_req + 1, + dtype=torch.int32, + device='cpu') + + num_cached) // compress_ratio # Q token range in batch (indices into context tokens in the current batch) token_start = req_cum_start + token_start_in_req token_end = req_cum_start + token_end_in_req - # K token range: index into full KV slot mapping (cached + current batch context tokens) - kv_offset_in_extended = metadata.host_ctx_kv_indptr[req_idx].item() - total_kv_for_req = num_cached + token_end_in_req + # K token range: index into full KV slot mapping in compressed KV space + # For req_idx=0 the offset is 0; for other indices compute compressed cumulative offset + kv_offset_in_extended = sum( + (metadata.host_ctx_kv_indptr[j + 1] - + metadata.host_ctx_kv_indptr[j]).item() // compress_ratio + for j in range(req_idx)) + total_kv_for_req = req_kv_len k_token_start = kv_offset_in_extended k_token_end = kv_offset_in_extended + total_kv_for_req @@ -1525,40 +1736,48 @@ def prepare_one_prefill_chunk( # Multi-request chunk: batch multiple full requests together # Extract sequence lengths for these requests req_seq_lens = [] - req_cached_lens = [] + req_num_past_tokens = [] + req_kv_lens = [] first_req_idx = chunk_specs[0][0] for spec in chunk_specs: req_idx, token_start_in_req, token_end_in_req, _ = spec req_seq_lens.append(token_end_in_req - token_start_in_req) # Get cached token count from metadata - num_cached = ( + num_past_tokens = ( metadata.host_ctx_cached_token_indptr[req_idx + 1] - metadata.host_ctx_cached_token_indptr[req_idx]).item() - req_cached_lens.append(num_cached) + req_num_past_tokens.append(num_past_tokens) + req_kv_lens.append( + (num_past_tokens + req_seq_lens[-1]) // compress_ratio) req_seq_lens_tensor = torch.tensor(req_seq_lens, dtype=torch.int32, device='cpu') - req_cached_lens_tensor = torch.tensor(req_cached_lens, - dtype=torch.int32, - device='cpu') + req_num_past_tokens_tensor = torch.tensor(req_num_past_tokens, + dtype=torch.int32, + device='cpu') + req_kv_lens_tensor = torch.tensor(req_kv_lens, + dtype=torch.int32, + device='cpu') num_q_tokens = sum(req_seq_lens) # Compute causal attention bounds for batched requests cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( req_seq_lens_tensor, len(chunk_specs), num_q_tokens, - req_cached_lens_tensor) + req_num_past_tokens_tensor, req_kv_lens_tensor, compress_ratio) # Global Q token ranges (indices into ctx tokens in the current batch) token_start = chunk_specs[0][3] # req_cum_start of first request token_end = token_start + num_q_tokens # K token range: index into full kv slot mapping (cached + current ctx tokens within the batch) - kv_offset_in_extended = metadata.host_ctx_kv_indptr[ - first_req_idx].item() - total_kv_len = sum(req_seq_lens_tensor + - req_cached_lens_tensor).item() + # Must use compressed offsets + kv_offset_in_extended = sum( + (metadata.host_ctx_kv_indptr[j + 1] - + metadata.host_ctx_kv_indptr[j]).item() // compress_ratio + for j in range(first_req_idx)) + total_kv_len = sum(req_kv_lens) k_token_start = kv_offset_in_extended k_token_end = kv_offset_in_extended + total_kv_len @@ -1568,8 +1787,10 @@ def prepare_one_prefill_chunk( f"Indexer.prepare_one_prefill_chunk - cu_seqlen_ke length mismatch: {cu_seqlen_ke.shape[0]} != {num_q_tokens}" return IndexerPrefillChunkMetadata( - cu_seqlen_ks=cu_seqlen_ks.to(device, non_blocking=True), - cu_seqlen_ke=cu_seqlen_ke.to(device, non_blocking=True), + cu_seqlen_ks=maybe_pin_memory(cu_seqlen_ks).to(device, + non_blocking=True), + cu_seqlen_ke=maybe_pin_memory(cu_seqlen_ke).to(device, + non_blocking=True), token_start=token_start, token_end=token_end, k_token_start=k_token_start, @@ -1577,41 +1798,40 @@ def prepare_one_prefill_chunk( ) @staticmethod - def recompute_slot_mappings(metadata: DSAtrtllmAttentionMetadata): - """Recompute only slot_mapping_fp8/scale from the current block offsets. - - This is the subset of prepare() that maps each token to its flat cache - position. It is safe to call in isolation (e.g. during draft KV-cache - replay) because it only touches slot-mapping buffers and reads - block-offset / sequence metadata that the caller has already set up. + def prepare_for_update_k_cache(metadata: DSAtrtllmAttentionMetadata, + indexer_params: IndexerParams): """ - kv_cache_manager = metadata.kv_cache_manager - if kv_cache_manager is None or not hasattr(kv_cache_manager, - 'index_head_dim'): - return - - seq_lens = metadata.seq_lens - head_dim = kv_cache_manager.index_head_dim - tokens_per_block = kv_cache_manager.tokens_per_block - quant_block_size = kv_cache_manager.quant_block_size - use_fp4 = kv_cache_manager.use_fp4 - # FP4 packs two E2M1 codes per byte; FP8 stores one byte per element. - data_bytes_per_token = head_dim // 2 if use_fp4 else head_dim - cached_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq - total_tokens = seq_lens.sum().item() - - start_positions = torch.tensor(cached_tokens, dtype=torch.int32) - batch_size = len(metadata.request_ids) + Prepare indexer for the update_k_cache stage. + Compute slot_mapping for all requests (both context and generation) + This maps each token to its flat cache position for vectorized KV cache updates + """ + batch_size = indexer_params.batch_size + tokens_per_block = indexer_params.tokens_per_block + head_dim = indexer_params.head_dim + new_kv_tokens = indexer_params.new_kv_tokens + total_new_kv_tokens = new_kv_tokens.sum().item() + data_bytes_per_token = head_dim // 2 if metadata.kv_cache_manager.use_fp4 else head_dim + + # Compute global positions for all kv tokens in the batch (fully vectorized) req_indices = torch.repeat_interleave( - torch.arange(batch_size, dtype=torch.int64, device='cpu'), seq_lens) - - token_offsets = torch.cat([ - torch.arange(seq_lens[i].item(), dtype=torch.int64, device='cpu') - for i in range(batch_size) - ]) - - global_positions = start_positions[req_indices] + token_offsets + torch.arange(batch_size, dtype=torch.int64, device='cpu'), + new_kv_tokens) + # Vectorized token_offsets: arange(total) - cumulative start per request + cu_new_kv = torch.zeros(batch_size + 1, dtype=torch.int64, device='cpu') + cu_new_kv[1:] = new_kv_tokens.to(torch.int64).cumsum(0) + token_offsets = ( + torch.arange(total_new_kv_tokens, dtype=torch.int64, device='cpu') - + cu_new_kv[:-1].repeat_interleave(new_kv_tokens)) + global_positions = indexer_params.cached_kv_tokens[ + req_indices] + token_offsets + + # Block indices/pos for all kv tokens in the batch + block_indices_in_seq = global_positions // tokens_per_block + + max_blocks = metadata.host_indexer_k_cache_block_offsets.shape[1] + assert (block_indices_in_seq < max_blocks).all(), \ + f"Block index out of bounds: max={max_blocks}, got indices up to {block_indices_in_seq.max().item()}" fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( global_positions, @@ -1619,203 +1839,99 @@ def recompute_slot_mappings(metadata: DSAtrtllmAttentionMetadata): req_indices, head_dim, tokens_per_block, - quant_block_size, + indexer_params.quant_block_size, data_bytes_per_token=data_bytes_per_token, ) - metadata.host_slot_mapping_fp8[:total_tokens] = fp8_flat_indices - metadata.host_slot_mapping_scale[:total_tokens] = scale_flat_indices + metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices + metadata.host_slot_mapping_scale[: + total_new_kv_tokens] = scale_flat_indices - metadata.slot_mapping_fp8[:total_tokens].copy_( - metadata.host_slot_mapping_fp8[:total_tokens], non_blocking=True) - metadata.slot_mapping_scale[:total_tokens].copy_( - metadata.host_slot_mapping_scale[:total_tokens], non_blocking=True) + metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_( + metadata.host_slot_mapping_fp8[:total_new_kv_tokens], + non_blocking=True) + metadata.slot_mapping_scale[:total_new_kv_tokens].copy_( + metadata.host_slot_mapping_scale[:total_new_kv_tokens], + non_blocking=True) @staticmethod - def prepare(metadata: DSAtrtllmAttentionMetadata): + def prepare_for_chunked_prefill(metadata: DSAtrtllmAttentionMetadata, + indexer_params: IndexerParams): """ - Prepare indexer for the forward pass. - This should be called during metadata.prepare() stage. - - - Computes slot_mapping for KV cache updates - - Prepares schedule_metadata for fp8_paged_mqa_logits - - Stores generation request IDs for decode phase + Prepare indexer for the chunked prefill. """ - kv_cache_manager = metadata.kv_cache_manager - num_contexts = metadata.num_contexts - num_generations = metadata.num_generations - num_ctx_tokens = metadata.num_ctx_tokens - seq_lens = metadata.seq_lens - tokens_per_block = kv_cache_manager.tokens_per_block - - # Prepare for prefill phase if there are context requests - if num_contexts > 0: - # Compute attention window bounds for each query token in batched sequences - # cu_seqlen_ks[i]: start index in global KV for query token i - # cu_seqlen_ke[i]: end index (exclusive) in global KV for query token i - host_seq_lens = seq_lens[:num_contexts] - cached_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq - host_cached_tokens = torch.tensor(cached_tokens[:num_contexts], - dtype=torch.int32, - device='cpu') + num_contexts = indexer_params.num_contexts + seq_lens = indexer_params.seq_lens + tokens_per_block = indexer_params.tokens_per_block + head_dim = indexer_params.head_dim + + # When MLA chunked prefill is active, it already handles chunking + # Indexer should just process the current MLA chunk as a single chunk + has_mla_chunked_prefill = (metadata.enable_context_mla_with_cached_kv + and + metadata.runtime_features.chunked_prefill) + if has_mla_chunked_prefill: + chunk_specs = [(i, 0, seq_lens[i].item(), + seq_lens[:i].sum().item() if i > 0 else 0) + for i in range(num_contexts)] + metadata.indexer_prefill_chunks = [ + Indexer.prepare_one_prefill_chunk( + metadata, + chunk_specs, + ) + ] + else: + # Use indexer's own chunking logic to prevent L^2 complexity of indexer MQA logits computation for long sequences. + # This is only used when MLA chunked prefill is not enabled. + chunk_groups = split_prefill_chunks( + seq_lens[:num_contexts], + metadata.indexer_max_chunk_size, + start_idx=0, + ) - # When MLA chunked prefill is active, it already handles chunking - # Indexer should just process the current MLA chunk as a single chunk - has_mla_chunked_prefill = ( - metadata.enable_context_mla_with_cached_kv - and metadata.runtime_features.chunked_prefill) - - if has_mla_chunked_prefill: - # MLA chunked prefill is active - use single-chunk pattern for - # indexer prefill chunks. - chunk_specs = [(i, 0, host_seq_lens[i].item(), - host_seq_lens[:i].sum().item() if i > 0 else 0) - for i in range(num_contexts)] + if len(chunk_groups + ) > 1 or metadata.enable_context_mla_with_cached_kv: metadata.indexer_prefill_chunks = [ Indexer.prepare_one_prefill_chunk( metadata, chunk_specs, - ) + ) for chunk_specs in chunk_groups ] else: - # Use indexer's own chunking logic to prevent L^2 complexity of indexer MQA logits computation for long sequences. - # This is only used when MLA chunked prefill is not enabled. - chunk_groups = split_prefill_chunks( - host_seq_lens, - metadata.indexer_max_chunk_size, - start_idx=0, - ) - - if len(chunk_groups - ) > 1 or metadata.enable_context_mla_with_cached_kv: - metadata.indexer_prefill_chunks = [ - Indexer.prepare_one_prefill_chunk( - metadata, - chunk_specs, - ) for chunk_specs in chunk_groups - ] - else: - metadata.indexer_prefill_chunks = None - - host_cu_seqlen_ks, host_cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - host_seq_lens, num_contexts, num_ctx_tokens, host_cached_tokens) - - metadata.cu_seqlen_ks[:num_ctx_tokens].copy_(host_cu_seqlen_ks, - non_blocking=True) - metadata.cu_seqlen_ke[:num_ctx_tokens].copy_(host_cu_seqlen_ke, - non_blocking=True) - - # Prepare for decode phase if there are generation requests - if num_generations > 0: - # Prepare schedule metadata for fp8_paged_mqa_logits - # This is a preprocessing step that computes scheduling information for the kernel - if not metadata.use_expanded_buffers_for_mtp: - # Write 2D kv_lens (broadcast same kv_len across next_n positions). - gen_seq_lens = metadata.kv_lens_cuda_runtime[ - num_contexts:num_contexts + num_generations] - next_n_cap = metadata.kv_lens_cuda_2d.shape[1] - metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( - gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap)) - # Build the next_n=1 schedule (used by MTP draft layers). - # Use the contiguous 1D gen slice reshaped to (num_gen, 1); - # slicing kv_lens_cuda_2d's first column would be a strided - # view that fails the metadata kernel's contiguous assertion. - context_lens_next_n1 = gen_seq_lens.view(-1, 1) - # `_DG_SCHEDULE_BLOCK_KV` (= 64) instead of cache `tokens_per_block`: - # see module-level constant comment for the SPLIT_KV=256 alignment. - scheduler_metadata_buffer = get_paged_mqa_logits_metadata( - context_lens_next_n1, _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms) - metadata.scheduler_metadata_buffer.copy_( - scheduler_metadata_buffer, non_blocking=True) - # MTP main forward uses next_n = 1 + max_draft_tokens; build - # a separate schedule because the metadata kernel reads next_n - # from context_lens.size(1). - if metadata.max_draft_tokens > 0: - context_lens_full_next_n = metadata.kv_lens_cuda_2d[: - num_generations, : - next_n_cap] - scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( - context_lens_full_next_n, _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms) - metadata.scheduler_metadata_buffer_full_next_n.copy_( - scheduler_metadata_buffer_full_next_n, - non_blocking=True) - else: - # Expand schedule metadata buffer (only generation). The new - # DeepGEMM API requires 2D; each expanded token becomes a (1,) - # row. - num_tokens = metadata.num_generations * ( - 1 + metadata.max_draft_tokens) - kv_lens_expanded_2d = metadata.kv_lens_expanded_cuda[: - num_tokens].view( - -1, 1) - scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - kv_lens_expanded_2d, _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms) - metadata.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - - # DSL atom-split schedule. Picker decision was cached on - # `metadata.dsl_{expand_factor, atom}` at metadata prepare - # time; only build the expanded schedule when picker chose to - # split (factor > 1). Runtime mutually exclusive with the `else` - # branch above (latter requires `use_expanded_buffers_for_mtp` - # which is False under DSL). - if metadata.expand_for_dsl and metadata.num_generations > 0 \ - and metadata.dsl_expand_factor > 1: - expand_factor = metadata.dsl_expand_factor - num_tokens = metadata.num_generations * expand_factor - kv_lens_expanded_2d = metadata.kv_lens_expanded_cuda[: - num_tokens].view( - -1, 1) - scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - kv_lens_expanded_2d, _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms) - metadata.scheduler_metadata_buffer_expanded.copy_( - scheduler_metadata_buffer_expanded, non_blocking=True) - - # Compute slot_mapping for all requests (both context and generation) - Indexer.recompute_slot_mappings(metadata) + metadata.indexer_prefill_chunks = None # When chunked prefill or KVCache reuse is enabled, we need to gather the full KV for indexer's logit computation. # Indexer's own chunking does not need full KV gathering, instead it gathers only the current chunk with loop-based gathering. - _need_full_kv_gathering = num_contexts > 0 and metadata.enable_context_mla_with_cached_kv - if _need_full_kv_gathering: - head_dim = kv_cache_manager.index_head_dim - quant_block_size = kv_cache_manager.quant_block_size - use_fp4 = kv_cache_manager.use_fp4 - data_bytes_per_token = head_dim // 2 if use_fp4 else head_dim - cached_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq - start_positions = torch.tensor(cached_tokens, dtype=torch.int32) - - total_kv_len = metadata.host_ctx_kv_indptr[num_contexts].item() - total_kv_per_request = seq_lens[: - num_contexts] + start_positions[: - num_contexts] + if metadata.enable_context_mla_with_cached_kv: + # Use kv_lens which correctly computes (raw_past + seq_lens) // compress_ratio. + total_kv_per_request = indexer_params.kv_lens[:num_contexts] + total_kv_len = total_kv_per_request.sum().item() host_slot_mapping_fp8_fullkv = torch.empty( total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()) host_slot_mapping_scale_fullkv = torch.empty( total_kv_len, dtype=torch.int64, pin_memory=prefer_pinned()) - fullkv_req_indices = torch.repeat_interleave( + req_indices = torch.repeat_interleave( torch.arange(num_contexts, dtype=torch.int64, device='cpu'), total_kv_per_request) - kv_positions = torch.cat([ - torch.arange(total_kv_per_request[i].item(), - dtype=torch.int64, - device='cpu') for i in range(num_contexts) - ]) + cu_kv = torch.zeros(num_contexts + 1, + dtype=torch.int64, + device='cpu') + cu_kv[1:] = total_kv_per_request.to(torch.int64).cumsum(0) + kv_positions = ( + torch.arange(total_kv_len, dtype=torch.int64, device='cpu') - + cu_kv[:-1].repeat_interleave(total_kv_per_request)) fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( kv_positions, metadata.host_indexer_k_cache_block_offsets, - fullkv_req_indices, + req_indices, head_dim, tokens_per_block, - quant_block_size, - data_bytes_per_token=data_bytes_per_token, + indexer_params.quant_block_size, + data_bytes_per_token=head_dim // + 2 if metadata.kv_cache_manager.use_fp4 else head_dim, ) host_slot_mapping_fp8_fullkv[:total_kv_len] = fp8_flat_indices @@ -1833,6 +1949,124 @@ def prepare(metadata: DSAtrtllmAttentionMetadata): metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + @staticmethod + def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): + """ + Prepare scheduler metadata for the DeepGEMM decode MQA kernel. + """ + num_contexts = metadata.num_contexts + num_generations = metadata.num_generations + if not metadata.use_expanded_buffers_for_mtp: + gen_seq_lens = metadata.get_indexer_kv_lens( + metadata.kv_lens_cuda_runtime[num_contexts:num_contexts + + num_generations]) + metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens + next_n_cap = metadata.kv_lens_cuda_2d.shape[1] + metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( + gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap)) + scheduler_metadata_buffer = get_paged_mqa_logits_metadata( + gen_seq_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, + metadata.num_sms) + metadata.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, + non_blocking=True) + if metadata.max_draft_tokens > 0: + scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( + metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap], + _DG_SCHEDULE_BLOCK_KV, metadata.num_sms) + metadata.scheduler_metadata_buffer_full_next_n.copy_( + scheduler_metadata_buffer_full_next_n, non_blocking=True) + else: + # Expand schedule metadata buffer (only generation). The DeepGEMM + # API requires 2D; each expanded token becomes a (1,) row. + num_tokens = metadata.num_generations * (1 + + metadata.max_draft_tokens) + scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( + metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + _DG_SCHEDULE_BLOCK_KV, metadata.num_sms) + metadata.scheduler_metadata_buffer_expanded.copy_( + scheduler_metadata_buffer_expanded, non_blocking=True) + + @staticmethod + def prepare(metadata: DSAtrtllmAttentionMetadata): + """ + Prepare indexer for the forward pass. + This should be called during metadata.prepare() stage. + + - Computes slot_mapping for KV cache updates + - Prepares schedule_metadata for fp8_paged_mqa_logits + - Stores generation request IDs for decode phase + """ + + # Skip indexer preparation if the kv_cache_manager doesn't have index_head_dim. + # This can happen when the metadata is being used with a draft KV cache manager + # during MTP speculative decoding, which uses a regular KVCacheManager instead + # of DSACacheManager. + kv_cache_manager = metadata.kv_cache_manager + if kv_cache_manager is None or not hasattr(kv_cache_manager, + 'index_head_dim'): + return + + num_contexts = metadata.num_contexts + num_generations = metadata.num_generations + num_ctx_tokens = metadata.num_ctx_tokens + request_ids = metadata.request_ids + seq_lens = metadata.seq_lens + head_dim = metadata.indexer_head_dim + quant_block_size = metadata.indexer_quant_block_size + num_past_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq + compress_ratio = _effective_compress_ratio_divisor( + _select_indexer_compress_ratio(metadata.compress_ratios)) + # MXFP4 indexer K cache packs two E2M1 codes per byte so the data + # footprint is half head_dim; FP8 keeps one byte per element. + use_fp4 = getattr(kv_cache_manager, 'use_fp4', False) + data_bytes_per_token = head_dim // 2 if use_fp4 else head_dim + + indexer_params = IndexerParams( + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + head_dim=head_dim, + quant_block_size=quant_block_size, + tokens_per_block=metadata._tokens_per_block, + compress_ratio=compress_ratio, + request_ids=request_ids, + num_past_tokens=num_past_tokens, + seq_lens=seq_lens, + data_bytes_per_token=data_bytes_per_token, + ) + # Store compressed KV token count for context requests + metadata.num_ctx_kv_tokens = indexer_params.new_kv_tokens[: + num_contexts].sum( + ).item() + + # Prepare for update_k_cache + Indexer.prepare_for_update_k_cache(metadata, indexer_params) + + # Prepare for prefill phase if there are context requests + if num_contexts > 0: + # Compute attention window bounds for each query token in batched sequences + # cu_seqlen_ks[i]: start index in global KV for query token i + # cu_seqlen_ke[i]: end index (exclusive) in global KV for query token i + host_seq_lens = seq_lens[:num_contexts] + host_num_past_tokens = indexer_params._num_past_tokens_tensor[: + num_contexts] + host_kv_lens = indexer_params.kv_lens[:num_contexts] + host_cu_seqlen_ks, host_cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( + host_seq_lens, num_contexts, num_ctx_tokens, + host_num_past_tokens, host_kv_lens, compress_ratio) + + metadata.cu_seqlen_ks[:num_ctx_tokens].copy_( + maybe_pin_memory(host_cu_seqlen_ks), non_blocking=True) + metadata.cu_seqlen_ke[:num_ctx_tokens].copy_( + maybe_pin_memory(host_cu_seqlen_ke), non_blocking=True) + Indexer.prepare_for_chunked_prefill(metadata, indexer_params) + + # Prepare for decode phase if there are generation requests + if num_generations > 0: + # Prepare schedule metadata for fp8_paged_mqa_logits + # This is a preprocessing step that computes scheduling information for the kernel + Indexer.prepare_scheduler_metadata(metadata) + def _update_k_cache(self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata) -> None: """ @@ -1854,29 +2088,96 @@ def _update_k_cache(self, k_fp8: torch.Tensor, k_scale: torch.Tensor, # The C++ op reinterprets k_fp8 (FP8) and k_scale (float32) as raw # bytes internally and only reads the first num_tokens entries from # the slot mapping buffers, avoiding Python-side view/slice overhead. + if k_scale.element_size() == 1: + # The op expects 4 byte elements. + k_scale = k_scale.view(torch.int32) torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache, metadata.slot_mapping_fp8, metadata.slot_mapping_scale, num_tokens) + def _gather_k_cache_for_chunk( + self, + metadata: DSAtrtllmAttentionMetadata, + chunk: IndexerPrefillChunkMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Gather K values from indexer cache for a specific chunk. + + Uses pre-computed extended slot mappings that cover cached + current batch context tokens. + chunk.k_token_start/k_token_end directly index into the extended slot mapping. + + Args: + metadata: Attention metadata + chunk: Chunk metadata with k_token_start/end as indices into extended slot mapping + + Returns: + k_fp8: FP8 quantized k tensor, shape [num_k_tokens, head_dim] + k_scale: Scaling factors, shape [num_k_tokens, 1] + """ + assert metadata.slot_mapping_fp8_fullkv is not None, \ + "_gather_k_cache_for_chunk requires extended slot mappings (only available with cached tokens)" + + k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers( + self.layer_idx) + + head_dim = self.head_dim + scale_size = 4 # float32 = 4 bytes + + # Extract slot mappings using chunk's k_token_start/end + # These indices point directly into the extended slot mapping array + k_token_start = chunk.k_token_start + k_token_end = chunk.k_token_end + num_k_tokens = k_token_end - k_token_start + + slot_mapping_fp8_chunk = metadata.slot_mapping_fp8_fullkv[ + k_token_start:k_token_end] + slot_mapping_scale_chunk = metadata.slot_mapping_scale_fullkv[ + k_token_start:k_token_end] + + # Vectorized gather using pre-computed slot mappings + # Gather FP8 data + byte_offsets_fp8 = torch.arange( + head_dim, device=k_cache.device).unsqueeze(0) # [1, head_dim] + gather_indices_fp8 = slot_mapping_fp8_chunk.unsqueeze( + 1) + byte_offsets_fp8 # [num_k_tokens, head_dim] + assert (gather_indices_fp8 + >= k_cache.numel()).sum() == 0, "Out-of-bounds access detected" + gather_indices_fp8 = _unravel_indices(gather_indices_fp8, k_cache.shape) + k_fp8_bytes = k_cache[gather_indices_fp8] + k_fp8 = k_fp8_bytes.view(torch.float8_e4m3fn).view( + num_k_tokens, head_dim) + + # Gather scale data + byte_offsets_scale = torch.arange( + scale_size, device=k_cache.device).unsqueeze(0) # [1, 4] + gather_indices_scale = slot_mapping_scale_chunk.unsqueeze( + 1) + byte_offsets_scale # [num_k_tokens, 4] + assert (gather_indices_scale + >= k_cache.numel()).sum() == 0, "Out-of-bounds access detected" + gather_indices_scale = _unravel_indices(gather_indices_scale, + k_cache.shape) + k_scale_bytes = k_cache[gather_indices_scale] + k_scale = k_scale_bytes.view(torch.float32).view(num_k_tokens, 1) + + return k_fp8, k_scale + def _call_mqa_logits(self, q_fp8: torch.Tensor, k_fp8: torch.Tensor, k_scale: torch.Tensor, weights: torch.Tensor, cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, q_scale: Optional[torch.Tensor]) -> torch.Tensor: - """Dispatch to fp8_mqa_logits or fp8_fp4_mqa_logits based on use_fp4. + """Dispatch fp8_mqa_logits vs fp8_fp4_mqa_logits based on use_fp4. - For FP4 the gather output is typed as FP8 for historical reasons; - reinterpret the bytes as the expected int8 / int32 layouts. DeepGEMM - asserts kv_sf is 1D in both modes, so flatten the scale here. + For FP4 the gather output keeps the legacy float8_e4m3fn dtype for + API compatibility; reinterpret the bytes as the int8 / int32 layout + the DeepGEMM kernel expects. The scale tensor is collapsed to 1D for + the kv side and 2D for the q side per the kernel's asserts. """ if self.use_fp4: k_fp4_bytes = k_fp8.view(torch.int8) k_scale_int32 = k_scale.view(torch.int32).reshape(-1) - # q_scale arrives here as (chunk_tokens, n_heads, 1) — fused_cat_fp4 - # emits one int32 per (token, head) carrying four UE8M0 exponents, - # pre_indexer_proj reshapes back to (N, n_heads, 1), and - # sparse_attn_indexer chunk-slices on axis 0. The DeepGEMM FP4 - # kernel asserts q_sf is 2D, so collapse the trailing unit axis. + # q_scale arrives as (chunk_tokens, n_heads, 1); the FP4 kernel + # asserts q_sf is 2D so collapse the trailing unit axis. q_scale_2d = q_scale.reshape(-1, self.n_heads) return fp8_fp4_mqa_logits( (q_fp8, q_scale_2d), @@ -1896,7 +2197,7 @@ def _call_paged_mqa_logits(self, q_decode: torch.Tensor, scheduler_metadata_buffer: torch.Tensor, max_seq_len: int, q_scale: Optional[torch.Tensor]) -> torch.Tensor: - """Dispatch to fp8_paged_mqa_logits or fp8_fp4_paged_mqa_logits.""" + """Dispatch fp8_paged_mqa_logits vs fp8_fp4_paged_mqa_logits.""" if self.use_fp4: return fp8_fp4_paged_mqa_logits( (q_decode, q_scale), k_cache, weights_decode, context_lens, @@ -1915,20 +2216,22 @@ def sparse_attn_indexer( weights: torch.Tensor, use_custom_topk: bool = True, q_scale: Optional[torch.Tensor] = None, + update_k_cache: bool = True, ) -> torch.Tensor: """Run the indexer TopK kernel for both prefill and decode phases. q_scale is only consumed by the FP4 dispatch; FP8 path ignores it. """ - # DSACacheManager hardcodes quant_block_size=128 (see its __init__); - # FP4 uses per-block-32 UE8M0 scales but packs four of them into one - # int32 so the cache-layout contribution is the same 4 bytes/token as - # FP8 per-block-128 at head_dim=128. + # DSACacheManager / DeepseekV4CacheManager force quant_block_size to + # 128 (FP8 path) or 32 (MXFP4 path); both round-trip to the same + # 4-byte scale word per token at index_head_dim=128, so the slot + # mapping arithmetic is unchanged in either mode. assert metadata.kv_cache_manager is None or \ - metadata.kv_cache_manager.quant_block_size == 128, \ + metadata.kv_cache_manager.quant_block_size in (32, 128), \ f"Unexpected quant_block_size {metadata.kv_cache_manager.quant_block_size if metadata.kv_cache_manager else 'N/A'}" # Update the indexer k cache before prefill chunks gather from it. - self._update_k_cache(k_fp8, k_scale, metadata) + if update_k_cache: + self._update_k_cache(k_fp8, k_scale, metadata) num_contexts = metadata.num_contexts num_generations = metadata.num_generations @@ -1960,9 +2263,19 @@ def sparse_attn_indexer( k_cache_4d = metadata.kv_cache_manager.get_indexer_k_cache_buffers( self.layer_idx) - + # FP4 packs two codes per byte so the gathered row holds half + # as many bytes as in the FP8 path. The scale (4 bytes) is the + # same in both modes because FP4 packs four UE8M0 exponents + # into one int32 to match FP8's float32 scale width. gather_head_dim = self.head_dim // 2 if self.use_fp4 else self.head_dim + for chunk in metadata.indexer_prefill_chunks: + # Skip chunks with no compressed KV tokens (e.g., warmup + # sequences shorter than compress_ratio produce zero KV). + if chunk.k_token_start >= chunk.k_token_end: + topk_indices_buffer[ + chunk.token_start:chunk.token_end, :].fill_(-1) + continue num_k_tokens = chunk.k_token_end - chunk.k_token_start chunk_k_fp8, chunk_k_scale = torch.ops.trtllm.indexer_k_cache_gather_op( k_cache_4d, metadata.slot_mapping_fp8_fullkv, @@ -1997,7 +2310,8 @@ def sparse_attn_indexer( logits, chunk.cu_seqlen_ks[chunk_q_start:chunk_q_end], chunk.cu_seqlen_ke[chunk_q_start:chunk_q_end], - topk_indices_buffer[global_q_start:global_q_end, :]) + topk_indices_buffer[global_q_start:global_q_end, :], + self.index_topk) else: topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), @@ -2030,8 +2344,12 @@ def sparse_attn_indexer( metadata.mapping, dim=0, sizes=q_sizes) + elif metadata.num_ctx_kv_tokens == 0: + # No compressed KV tokens — fill with -1 (no valid indices) + topk_indices_buffer[:num_ctx_tokens, :].fill_(-1) else: # Fallback: single-pass indexer prefill (TODO: remove this once chunked prefill is fully tested) + num_ctx_kv_tokens = metadata.num_ctx_kv_tokens cu_seqlen_ks = metadata.cu_seqlen_ks[:num_ctx_tokens] cu_seqlen_ke = metadata.cu_seqlen_ke[:num_ctx_tokens] @@ -2039,8 +2357,8 @@ def sparse_attn_indexer( ...] if self.use_fp4 else None logits = self._call_mqa_logits( q_fp8[:num_ctx_tokens, ...], - k_fp8[:num_ctx_tokens, ...], - k_scale[:num_ctx_tokens, ...], + k_fp8[:num_ctx_kv_tokens, ...], + k_scale[:num_ctx_kv_tokens, ...], weights[:num_ctx_tokens, ...], cu_seqlen_ks, cu_seqlen_ke, @@ -2049,7 +2367,8 @@ def sparse_attn_indexer( if use_custom_topk: torch.ops.trtllm.indexer_topk_prefill( logits, cu_seqlen_ks, cu_seqlen_ke, - topk_indices_buffer[:num_ctx_tokens, :]) + topk_indices_buffer[:num_ctx_tokens, :], + self.index_topk) else: topk_indices = logits.topk(min(self.index_topk, logits.shape[-1]), @@ -2070,8 +2389,32 @@ def sparse_attn_indexer( topk_indices_buffer[:num_ctx_tokens, :] = \ metadata.topk_indices_buffer[:num_ctx_tokens, :] + # Prefill→decode GVR handoff: seed each finishing-prefill sequence's + # heuristic_prev_topk slot with its own last-context-token top-K, so + # the FIRST decode step of that sequence gets a warm-started preIdx + # (~60-75% set-overlap with the eventual decode top-K on this + # workload) instead of the all-zero / all-(-1) cold start that the + # default `heuristic_prev_topk.zero_()` initialization leaves behind. + # Without this, GVR P2 secant on decode step 0 runs from a benign + # but uninformative seed (kernel +1 offset on zeros → all indices + # point at compressed-token position 1), wasting iterations. + # Slot convention (mirrors the existing decode write-back at the + # bottom of the decode block): new gens from finishing prefill + # append after currently-active gens, i.e., slots + # [num_generations : num_generations + num_contexts]. + if (self._enable_heuristic_topk and has_prefill + and not metadata.skip_indexer_for_ctx_reqs): + local_layer = metadata.kv_cache_manager.layer_offsets[ + self.layer_idx] + ctx_seq_lens = metadata.seq_lens[:num_contexts] + # Per-sequence last context-token offset (exclusive cumsum minus 1). + last_ctx_idx = (torch.cumsum(ctx_seq_lens, dim=0) - + 1).to(dtype=torch.long) + metadata.heuristic_prev_topk[ + local_layer, num_generations:num_generations + + num_contexts].copy_(topk_indices_buffer[last_ctx_idx, :]) + if has_decode and not metadata.skip_indexer_for_gen_reqs: - max_seq_len = metadata.kv_cache_manager.max_seq_len # Get decode lengths per request (from seq_lens) for validation gen_seq_lens = metadata.seq_lens[num_contexts:num_contexts + num_generations] @@ -2120,10 +2463,12 @@ def sparse_attn_indexer( weights_decode = weights[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, ...] - # Get k cache and call fp8_paged_mqa_logits with prepared decode metadata + # Get k cache and call fp8_paged_mqa_logits / fp8_fp4_paged_mqa_logits + # with prepared decode metadata. # [num_blocks, tokens_per_block, 1, head_dim + scale_size] k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers( self.layer_idx) + indexer_max_seq_len = metadata.get_indexer_max_seq_len() if self.use_cute_dsl_paged_mqa_logits: # DSL kernel design: 1 atom per q (atom = real next_n positions), @@ -2136,8 +2481,8 @@ def sparse_attn_indexer( # the same KV length on this path (kv_lens_cuda_2d broadcasts), # so passing the 1D contiguous kv_lens slice for context_lens # avoids materializing a 2D contiguous tensor per call. - dsl_context_lens = metadata.kv_lens_cuda_runtime[ - num_contexts:num_contexts + num_generations] + dsl_context_lens = metadata.gen_indexer_kv_lens_cuda_runtime + assert dsl_context_lens is not None # Wave-aware atom-split: the picker in `_pick_dsl_expand` caches # (factor, atom) on metadata with invariant # `factor * atom == 1 + max_draft_tokens` (the target/verify-time @@ -2156,7 +2501,7 @@ def sparse_attn_indexer( # separate args and requires q.dtype == uint8 (q_decode # came in via the FP8 plumbing as int8; reinterpret with # no copy). sf_q is the q_scale slice reshaped to - # (B, next_n, H) int32 — mirrors the non-DSL FP4 branch. + # (B, next_n, H) int32 -- mirrors the non-DSL FP4 branch. decode_q_scale = q_scale[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, ...] decode_q_scale = decode_q_scale.view( @@ -2181,12 +2526,13 @@ def sparse_attn_indexer( logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits( dsl_q, decode_q_scale, k_cache, weights_decode, dsl_context_lens, dsl_block_table, dsl_schedule_meta, - max_seq_len) + indexer_max_seq_len) else: # FP8 DSL kernel natively supports next_n ∈ {1, 2, 3, 4}. # Atom-split benefits small-batch / low-ntask configs by - # raising SM utilization at the cost of factor× KV HBM - # re-reads; guard logic shared via dsl_atom_split above. + # raising SM utilization at the cost of factorx KV HBM + # re-reads. Context lengths are indexer KV lengths, which + # may be compressed for DeepSeek-V4. dsl_q = q_decode fp8_ctx_lens = dsl_context_lens fp8_block_table = block_table @@ -2203,7 +2549,8 @@ def sparse_attn_indexer( metadata.scheduler_metadata_buffer_expanded) logits_decode = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( dsl_q, k_cache, weights_decode, fp8_ctx_lens, - fp8_block_table, fp8_schedule_meta, max_seq_len) + fp8_block_table, fp8_schedule_meta, + indexer_max_seq_len) else: decode_q_scale = q_scale[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, @@ -2216,7 +2563,8 @@ def sparse_attn_indexer( q_decode.shape[0], q_decode.shape[1], self.n_heads) logits_decode = self._call_paged_mqa_logits( q_decode, k_cache, weights_decode, context_lens, - block_table, scheduler_metadata_buffer, max_seq_len, + block_table, scheduler_metadata_buffer, + indexer_max_seq_len, decode_q_scale) if use_custom_topk: @@ -2243,9 +2591,11 @@ def sparse_attn_indexer( # so we cap it at 256 for now and fall back to the CUDA C++ # indexer_topk_decode. This limit can be removed if GPU memory # is not a bottleneck. - if self.use_cute_dsl_topk and num_gen_tokens <= 256: + if (self.use_cute_dsl_topk and num_gen_tokens <= 256 + and (self.compress_ratio == 1 or next_n == 1)): torch.ops.trtllm.cute_dsl_indexer_topk_decode( - logits_decode, gen_kv_lens_cuda, + logits_decode, context_lens + if self.compress_ratio > 1 else gen_kv_lens_cuda, topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, :], self.index_topk, next_n) @@ -2258,19 +2608,22 @@ def sparse_attn_indexer( next_n, self.index_topk, pre_idx=pre_idx, - heuristic_scratch=heuristic_scratch) + heuristic_scratch=heuristic_scratch, + compress_ratio=self.compress_ratio, + radix_aux_indices=metadata.radix_aux_indices, + radix_aux_logits=metadata.radix_aux_logits) else: # padded positions = torch.arange( - max_seq_len, device=q_decode.device).unsqueeze(0).expand( + logits_decode.shape[-1], + device=q_decode.device).unsqueeze(0).expand( num_gen_tokens, -1) row_indices = torch.arange(num_gen_tokens, device=q_decode.device) // next_n next_n_offset = torch.arange(num_gen_tokens, device=q_decode.device) % next_n - index_end_pos = ( - metadata.kv_lens_cuda_runtime[num_contexts + row_indices] - - next_n + next_n_offset).unsqueeze(1) + index_end_pos = (context_lens[row_indices] - next_n + + next_n_offset).unsqueeze(1) # index_end_pos: [B * N, 1] mask = positions <= index_end_pos # mask: [B * N, L] @@ -2348,7 +2701,7 @@ def pre_indexer_proj( torch.Tensor]: """Pure token-wise projections (CUDA-graph-capturable). - Runs cublas_mm, qk_projection_and_rope, FP8 quantize, and weight + Runs cublas_mm, qk_projection_and_rope, FP8/FP4 quantize, and weight scaling. Does NOT touch the k cache or any batch-specific metadata, so this can safely run inside a captured CUDA graph partition. @@ -2392,7 +2745,14 @@ def pre_indexer_proj( # internally, so weights carry only softmax_scale * n_heads^-0.5. q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim // 2) q_scale = q_scale.view(-1, self.n_heads, 1) - weights = weights * self.weight_scale_factor + # DeepGEMM's fp8_fp4_(paged_)mqa_logits asserts + # `weights.scalar_type() == kFloat`. Unlike the FP8 branch (whose + # `_weight_scale` multiplies by the fp32 `q_scale` tensor and so + # implicitly upcasts), this branch's only multiplier is a Python + # float — `bf16 * float` stays bf16 and trips the kernel assert. + # Cast explicitly so this path doesn't silently rely on an + # upstream `_to_float`. + weights = weights.float() * self.weight_scale_factor else: q_fp8 = q_fp8.view(-1, self.n_heads, self.head_dim) q_scale = q_scale.view(-1, self.n_heads, 1) @@ -2458,10 +2818,14 @@ def __init__( attention_chunk_size=attention_chunk_size, **kwargs) - self.indexer = Indexer(quant_config, pos_embd_params, mla_params, + self.indexer = Indexer(quant_config, + pos_embd_params, + mla_params, skip_create_weights_in_init, - sparse_attention_config, dtype, layer_idx, - aux_stream) + sparse_attention_config, + dtype, + layer_idx=layer_idx, + aux_stream=aux_stream) def sparse_attn_predict( self, @@ -2472,11 +2836,9 @@ def sparse_attn_predict( ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Transform local TopK indices to global paged KV cache indices.""" # Transform the local topk indices to global topk indices in paged kv cache - is_generation = (forward_args.attention_input_type == - AttentionInputType.generation_only) topk_indices_global, _ = transform_local_topk_and_prepare_pool_view( forward_args.topk_indices, metadata, - self.get_local_layer_idx(metadata), is_generation) + self.get_local_layer_idx(metadata), forward_args.is_generation) # TODO: Use sparse_attn_indexer to predict the indices for DSA attention # return self.indexer(q, k, metadata, hidden_states, qr, position_ids) @@ -2619,6 +2981,14 @@ def get_indexer_k_cache_buffers(self, layer_idx: int): return self.indexer_k_cache_pool_per_layer[layer_offset].view( self.num_blocks, block_size, 1, per_token_size) + def get_batch_indexer_k_cache_indices( + self, request_ids: List[int]) -> List[List[int]]: + """ + Get the indices for the indexer k cache for a specific batch of requests. + """ + # All of layers share the same cache indices, so we use layer index 0. + return self.get_batch_cache_indices(request_ids, 0) + def shutdown(self): """Release indexer K-cache pool references before C++ buffer cleanup.""" # Clear Python references BEFORE C++ frees the underlying CUDA buffers diff --git a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py index 65e31f209182..bce4472b4942 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/kernel.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/kernel.py @@ -2079,3 +2079,266 @@ def triton_gather_k_cache( k_fp8 = out_fp8.view(torch.float8_e4m3fn) k_scale = out_scale.view(torch.float32).view(num_k_tokens, 1) return k_fp8, k_scale + + +######################################################## +# DeepSeek-V4 local to global index transformation kernel +######################################################## + + +@triton.jit +def _deepseek_v4_local_to_global_kernel( + req_id_ptr, + block_table_swa_ptr, + block_table_compressed_ptr, + swa_local_indices_ptr, + compressed_local_indices_ptr, + out_ptr, + swa_buffer_offset_in_tokens, + compressed_buffer_offset_in_tokens, + tokens_per_block_swa: tl.constexpr, + tokens_per_block_compressed: tl.constexpr, + max_blocks_swa, + max_blocks_compressed, + num_swa_indices: tl.constexpr, + num_compressed_indices: tl.constexpr, + total_output_indices, + has_compressed: tl.constexpr, + bt_swa_stride0, + bt_swa_stride1, + bt_compressed_stride0, + bt_compressed_stride1, + swa_indices_stride0, + swa_indices_stride1, + compressed_indices_stride0, + compressed_indices_stride1, + out_stride0, + out_stride1, +): + """ + Triton kernel for converting local indices to global KV cache pool indices. + + Dual-pool output layout: + - SWA region [0, num_swa_indices): indices relative to swa_pool_base_ptr + - Compress region [num_swa_indices, num_swa_indices + num_compressed_indices): + indices relative to compress_pool_base_ptr + - Invalid positions padded with -1 at their fixed positions. + + This enables the FMHA kernel to determine which TMA descriptor to use based + solely on tile index (tile 0 = SWA via tmaKSecondary_, rest = compress via tmaK_). + """ + token_id = tl.program_id(0) + + # Load request ID for this token + req = tl.load(req_id_ptr + token_id) + + # Load all SWA local indices for this token + swa_ids = tl.arange(0, num_swa_indices) + swa_ptr = swa_local_indices_ptr + token_id * swa_indices_stride0 + swa_ids * swa_indices_stride1 + swa_local_idx = tl.load(swa_ptr) + + # Compute global indices for all SWA positions + swa_valid_mask = swa_local_idx >= 0 + swa_block_ordinal = swa_local_idx // tokens_per_block_swa + swa_token_in_block = swa_local_idx % tokens_per_block_swa + swa_valid_block = swa_block_ordinal < max_blocks_swa + swa_full_mask = swa_valid_mask & swa_valid_block + + swa_bt_ptr = block_table_swa_ptr + req * bt_swa_stride0 + swa_block_ordinal * bt_swa_stride1 + swa_page_index = tl.load(swa_bt_ptr, mask=swa_full_mask, other=0) + + swa_global_index = swa_buffer_offset_in_tokens + swa_page_index * tokens_per_block_swa + swa_token_in_block + swa_global_index = tl.where(swa_full_mask, swa_global_index, -1) + + # Store SWA results at fixed positions [0, num_swa_indices) + swa_out_ptr = out_ptr + token_id * out_stride0 + swa_ids * out_stride1 + tl.store(swa_out_ptr, swa_global_index) + + if has_compressed: + # Load all compressed local indices for this token + compressed_ids = tl.arange(0, num_compressed_indices) + compressed_ptr = (compressed_local_indices_ptr + + token_id * compressed_indices_stride0 + + compressed_ids * compressed_indices_stride1) + compressed_local_idx = tl.load(compressed_ptr) + + # Compute global indices for all compressed positions + compressed_valid_mask = compressed_local_idx >= 0 + compressed_block_ordinal = compressed_local_idx // tokens_per_block_compressed + compressed_token_in_block = compressed_local_idx % tokens_per_block_compressed + compressed_valid_block = compressed_block_ordinal < max_blocks_compressed + compressed_full_mask = compressed_valid_mask & compressed_valid_block + + compressed_bt_ptr = (block_table_compressed_ptr + + req * bt_compressed_stride0 + + compressed_block_ordinal * bt_compressed_stride1) + compressed_page_index = tl.load(compressed_bt_ptr, + mask=compressed_full_mask, + other=0) + + compressed_global_index = ( + compressed_buffer_offset_in_tokens + + compressed_page_index * tokens_per_block_compressed + + compressed_token_in_block) + compressed_global_index = tl.where(compressed_full_mask, + compressed_global_index, -1) + + # Store compressed results at fixed positions [num_swa_indices, total) + compressed_write_pos = num_swa_indices + compressed_ids + compressed_out_ptr = out_ptr + token_id * out_stride0 + compressed_write_pos * out_stride1 + tl.store(compressed_out_ptr, compressed_global_index) + + +def deepseek_v4_local_to_global_indices( + req_id: torch.Tensor, # int32 [num_tokens] + block_table_swa: torch.Tensor, # int32 [num_requests, max_blocks_swa] + swa_local_indices: torch.Tensor, # int32 [num_tokens, num_swa_indices] + swa_pool_base_ptr: int, # int64: base address of SWA pool + swa_buffer_ptr: int, # int64: base address of SWA buffer + tokens_per_block: int, # tokens per block for SWA + token_stride: int, # bytes per token (use SWA token stride) + # Optional compressed arguments (for compress_ratio > 1) + block_table_compressed: torch.Tensor | None = None, + compressed_local_indices: torch.Tensor | None = None, + compress_pool_base_ptr: int = 0, # int64: base address of compress pool + compressed_buffer_ptr: int = 0, + compress_ratio: int = 1, + num_compressed_indices: int = 0, # max number of compressed indices +) -> torch.Tensor: + """ + Convert local token indices to global KV cache pool indices. + + Dual-pool output layout: + - SWA region [0, window_size): indices relative to swa_pool_base_ptr + - Compress region [window_size, window_size + num_compressed_indices): + indices relative to compress_pool_base_ptr + - Invalid positions padded with -1 at their fixed positions. + + For compress_ratio=1: Only SWA region, all indices relative to swa_pool_base_ptr. + For compress_ratio>1: SWA region + compress region with separate base pointers. + + Args: + req_id: Request ID per token [num_tokens], int32 + block_table_swa: SWA block table [num_requests, max_blocks_swa], int32 + swa_local_indices: Local indices for SWA cache [num_tokens, num_swa_indices], int32 + Use -1 for invalid/padding indices. + swa_pool_base_ptr: Base address of SWA pool + swa_buffer_ptr: Base address of SWA buffer (base_pool_ptr + buffer_offset_in_slot) + tokens_per_block: Number of tokens per block for SWA cache + token_stride: Bytes per token (use SWA token stride) + block_table_compressed: Compressed block table [num_requests, max_blocks_compressed], int32 (optional) + compressed_local_indices: Local indices for compressed cache [num_tokens, num_compressed_indices], int32 (optional) + Use -1 for invalid/padding indices. + compress_pool_base_ptr: Base address of compress pool + compressed_buffer_ptr: Base address of compressed buffer (optional) + compress_ratio: Compression ratio (1: no compression, >1: with compression) + num_compressed_indices: Max number of compressed indices for CUDA graph compatibility + Output width = num_swa_indices + num_compressed_indices. + + Returns: + global_indices: int32 [num_tokens, num_swa_indices + num_compressed_indices] + """ + assert req_id.dtype == torch.int32, f"req_id must be int32, got {req_id.dtype}" + assert block_table_swa.dtype == torch.int32, f"block_table_swa must be int32, got {block_table_swa.dtype}" + assert swa_local_indices.dtype == torch.int32, f"swa_local_indices must be int32, got {swa_local_indices.dtype}" + + num_tokens = req_id.shape[0] + num_swa_indices = swa_local_indices.shape[1] + + assert swa_local_indices.shape[0] == num_tokens + + has_compressed = compress_ratio > 1 + + # Compute SWA buffer offset relative to swa_pool_base_ptr in tokens + swa_buffer_offset_in_tokens = (swa_buffer_ptr - + swa_pool_base_ptr) // token_stride + + if has_compressed: + assert block_table_compressed is not None, "block_table_compressed required when compress_ratio > 1" + assert compressed_local_indices is not None, "compressed_local_indices required when compress_ratio > 1" + assert ( + block_table_compressed.dtype == torch.int32 + ), f"block_table_compressed must be int32, got {block_table_compressed.dtype}" + assert ( + compressed_local_indices.dtype == torch.int32 + ), f"compressed_local_indices must be int32, got {compressed_local_indices.dtype}" + assert compressed_local_indices.shape[0] == num_tokens + + tokens_per_block_compressed = tokens_per_block // compress_ratio + # Compute compressed buffer offset relative to compress_pool_base_ptr in tokens + assert ( + compressed_buffer_ptr - compress_pool_base_ptr + ) % token_stride == 0, "compressed_buffer_ptr must be aligned to token_stride" + compressed_buffer_offset_in_tokens = ( + compressed_buffer_ptr - compress_pool_base_ptr) // token_stride + _, max_blocks_compressed = block_table_compressed.shape + block_table_compressed_c = block_table_compressed.contiguous() + compressed_local_indices_c = compressed_local_indices.contiguous() + else: + # Dummy values + tokens_per_block_compressed = tokens_per_block + compressed_buffer_offset_in_tokens = 0 + max_blocks_compressed = 1 + block_table_compressed_c = torch.zeros((1, 1), + dtype=torch.int32, + device=req_id.device) + compressed_local_indices_c = torch.zeros((1, 1), + dtype=torch.int32, + device=req_id.device) + + total_output_indices = num_swa_indices + num_compressed_indices + _, max_blocks_swa = block_table_swa.shape + + # Ensure contiguous tensors + req_id_c = req_id.contiguous() + block_table_swa_c = block_table_swa.contiguous() + swa_local_indices_c = swa_local_indices.contiguous() + + # Create output tensor + out = torch.empty((num_tokens, total_output_indices), + dtype=torch.int32, + device=req_id.device) + + # Grid: one program per token + grid = (num_tokens, ) + + # Get strides + bt_swa_stride0, bt_swa_stride1 = block_table_swa_c.stride() + bt_compressed_stride0, bt_compressed_stride1 = block_table_compressed_c.stride( + ) + swa_indices_stride0, swa_indices_stride1 = swa_local_indices_c.stride() + compressed_indices_stride0, compressed_indices_stride1 = compressed_local_indices_c.stride( + ) + out_stride0, out_stride1 = out.stride() + + # Launch kernel + _deepseek_v4_local_to_global_kernel[grid]( + req_id_c, + block_table_swa_c, + block_table_compressed_c, + swa_local_indices_c, + compressed_local_indices_c, + out, + swa_buffer_offset_in_tokens, + compressed_buffer_offset_in_tokens, + tokens_per_block, + tokens_per_block_compressed, + max_blocks_swa, + max_blocks_compressed, + num_swa_indices, + num_compressed_indices, + total_output_indices, + has_compressed, + bt_swa_stride0, + bt_swa_stride1, + bt_compressed_stride0, + bt_compressed_stride1, + swa_indices_stride0, + swa_indices_stride1, + compressed_indices_stride0, + compressed_indices_stride1, + out_stride0, + out_stride1, + ) + + return out diff --git a/tensorrt_llm/_torch/attention_backend/sparse/utils.py b/tensorrt_llm/_torch/attention_backend/sparse/utils.py index ba37cecc9c05..56db0d9856f7 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/utils.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/utils.py @@ -1,6 +1,7 @@ from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from .deepseek_v4 import DeepseekV4CacheManager, DeepseekV4TrtllmAttention from .dsa import DSACacheManager, DSATrtllmAttention from .rocket import (RocketKVCacheManager, RocketTrtllmAttention, RocketVanillaAttention) @@ -12,6 +13,8 @@ def get_sparse_attn_kv_cache_manager( return RocketKVCacheManager elif sparse_attn_config.algorithm == "dsa": return DSACacheManager + elif sparse_attn_config.algorithm == "deepseek_v4": + return DeepseekV4CacheManager elif sparse_attn_config.algorithm == "skip_softmax": return KVCacheManager else: @@ -36,6 +39,8 @@ def get_trtllm_sparse_attn_attention_backend( return RocketTrtllmAttention elif sparse_attn_config.algorithm == "dsa": return DSATrtllmAttention + elif sparse_attn_config.algorithm == "deepseek_v4": + return DeepseekV4TrtllmAttention elif sparse_attn_config.algorithm == "skip_softmax": return TrtllmAttention else: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 9fe9f070301f..4b56dbfb7db0 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -15,7 +15,7 @@ from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings.internal import thop from tensorrt_llm.functional import AttentionMaskType -from tensorrt_llm.llmapi import SkipSoftmaxAttentionConfig +from tensorrt_llm.llmapi import DeepSeekV4SparseAttentionConfig, SkipSoftmaxAttentionConfig from tensorrt_llm.models.modeling_utils import QuantConfig from ..utils import (compute_swizzled_sf_shape, get_global_attrs, @@ -1814,6 +1814,32 @@ def forward( sparse_args.sparse_attn_offsets = at_off sparse_args.sparse_attn_indices_block_size = ( self.sparse_attention_config.get_indices_block_size()) + if isinstance(self.sparse_attention_config, + DeepSeekV4SparseAttentionConfig): + # TODO: refactor this part of code. + ratio = self.compress_ratio + if hasattr(metadata, 'sparse_mla_topk_lens'): + num_ctx_tokens = metadata.num_ctx_tokens + num_tokens = metadata.num_tokens + topk_lens = metadata.sparse_mla_topk_lens[ratio] + attention_input_type = forward_args.attention_input_type + if attention_input_type == AttentionInputType.context_only: + sparse_args.sparse_mla_topk_lens = topk_lens[: + num_ctx_tokens] + elif attention_input_type == AttentionInputType.generation_only: + sparse_args.sparse_mla_topk_lens = topk_lens[ + num_ctx_tokens:num_tokens] + else: + sparse_args.sparse_mla_topk_lens = topk_lens[: + num_tokens] + window_size = self.sparse_attention_config.window_size + compressed_len = metadata.max_compressed_indices[ratio] + metadata.num_sparse_topk = compressed_len + window_size + if ratio > 1: + # host_kv_cache_pool_pointers carries the SWA pool. The optional extra + # pointer is the compressed pool used for non-SWA sparse MLA tokens. + sparse_args.compressed_kv_cache_pool_ptr = ( + metadata.sparse_mla_base_ptrs[ratio]) # Compute FlashMLA tile-scheduler metadata once per forward pass. # The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index c4d9fb25ecfd..ab5cf9e81650 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -1,4 +1,5 @@ from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config +from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config from tensorrt_llm._torch.configs.laguna import LagunaConfig @@ -19,6 +20,7 @@ def _register_custom_configs_with_transformers() -> None: custom_configs = { "deepseek_v32": DeepseekV3Config, "kimi_k2": DeepseekV3Config, + "deepseek_v4": DeepseekV4Config, "laguna": LagunaConfig, } for model_type, config_class in custom_configs.items(): @@ -30,4 +32,4 @@ def _register_custom_configs_with_transformers() -> None: _register_custom_configs_with_transformers() del _register_custom_configs_with_transformers -__all__ = ["DeepseekV3Config", "LagunaConfig"] +__all__ = ["DeepseekV3Config", "DeepseekV4Config", "LagunaConfig"] diff --git a/tensorrt_llm/_torch/configs/deepseekv4.py b/tensorrt_llm/_torch/configs/deepseekv4.py new file mode 100644 index 000000000000..68dc481b4021 --- /dev/null +++ b/tensorrt_llm/_torch/configs/deepseekv4.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +# Default Engram embedding vocabulary size per n-gram level. +# Derived from the DeepSeek-V3 tokenizer vocab size (129280) scaled by 5. +_DEFAULT_ENGRAM_VOCAB_SIZE = 129280 * 5 + + +# This is a temporary workaround for DeepSeek-V4 model as HF does not support it yet +# TODO: Remove this once HF supports DeepSeek-V4 +class DeepseekV4Config(PretrainedConfig): + model_type = "deepseek_v4" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=129280, + hidden_size=4096, + intermediate_size=14336, + moe_intermediate_size=2048, + num_hidden_layers=43, + num_nextn_predict_layers=0, + num_attention_heads=64, + num_key_value_heads=64, + n_shared_experts=1, + n_routed_experts=256, + ep_size=1, + routed_scaling_factor=1.5, + kv_lora_rank=448, + q_lora_rank=1024, + qk_rope_head_dim=64, + v_head_dim=512, + qk_nope_head_dim=448, + topk_method="noaux_tc", + n_group=8, + topk_group=4, + num_experts_per_tok=6, + moe_layer_freq=1, + first_k_dense_replace=3, + norm_topk_prob=True, + scoring_func="sqrtsoftplus", + hidden_act="silu", + max_position_embeddings=65536, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=0, + eos_token_id=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + index_head_dim=128, + index_n_heads=64, + index_topk=512, + o_groups=8, + o_lora_rank=1024, + n_hash_layers=3, + hc_mult=4, + hc_sinkhorn_iters=20, + hc_eps=1e-6, + window_size=128, + compress_rope_theta=160000, + compress_ratios=None, + swiglu_limit=10.0, + has_engram=False, + engram_vocab_size=None, + engram_max_ngram_size=3, + engram_n_embed_per_ngram=512, + engram_n_head_per_ngram=8, + engram_kernel_size=4, + engram_pad_id=2, + engram_layer_ids=None, + engram_seed=0, + **kwargs, + ): + # DeepSeek-V4 HF config uses `sliding_window`, `num_hash_layers`, + # and `head_dim` for these internal fields. + # Accept them as aliases for the internal names the rest of the code uses. + if "sliding_window" in kwargs: + window_size = kwargs.pop("sliding_window") + if "num_hash_layers" in kwargs: + n_hash_layers = kwargs.pop("num_hash_layers") + if "head_dim" in kwargs: + v_head_dim = kwargs.pop("head_dim") + if "score_func" in kwargs: + scoring_func = kwargs.pop("score_func") + + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_nextn_predict_layers = num_nextn_predict_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.window_size = window_size + + # indexer + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.index_topk = index_topk + + # output projection + self.o_groups = o_groups + self.o_lora_rank = o_lora_rank + + # hash layer and mhc + self.n_hash_layers = n_hash_layers + self.hc_mult = hc_mult + self.hc_sinkhorn_iters = hc_sinkhorn_iters + self.hc_eps = hc_eps + + # kv compression + self.compress_rope_theta = compress_rope_theta + self.compress_ratios = compress_ratios + self.swiglu_limit = swiglu_limit + + # Engram configuration + self.has_engram = has_engram + self.engram_vocab_size = ( + engram_vocab_size + if engram_vocab_size is not None + else [_DEFAULT_ENGRAM_VOCAB_SIZE, _DEFAULT_ENGRAM_VOCAB_SIZE] + ) + self.engram_max_ngram_size = engram_max_ngram_size + self.engram_n_embed_per_ngram = engram_n_embed_per_ngram + self.engram_n_head_per_ngram = engram_n_head_per_ngram + self.engram_kernel_size = engram_kernel_size + self.engram_pad_id = engram_pad_id + self.engram_layer_ids = engram_layer_ids if engram_layer_ids is not None else [1, 15] + self.engram_seed = engram_seed + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index fa5cbd696d8d..efbb06238225 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -249,7 +249,8 @@ def _(logits, next_n, index_topk, pre_idx=None, - heuristic_scratch=None): + heuristic_scratch=None, + compress_ratio=1): # In-place operation, no return value (void function) pass @@ -263,6 +264,21 @@ def _(a, b, a_scale, b_scale): n = b.shape[0] return a.new_empty((m, n), dtype=torch.bfloat16) + @torch.library.register_fake("trtllm::gate_forward") + def _( + scores_in: torch.Tensor, + bias: torch.Tensor, + input_ids: torch.Tensor, + tid2eid: torch.Tensor, + out_weights: torch.Tensor, + out_indices: torch.Tensor, + topk: int, + route_scale: float, + is_hash: bool, + ) -> None: + # In-place operation, no return value (void function) + pass + @torch.library.register_fake("tensorrt_llm::quantize_e4m3_per_tensor") def _(input: torch.Tensor): scale_shape = [1] * input.dim() diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 7aaa7fb9d2ed..f5338c433625 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1914,7 +1914,8 @@ def _(a, b, a_scale, b_scale, tune_max_num_tokens=4096): @torch.library.custom_op("trtllm::silu_and_mul", mutates_args=()) def silu_and_mul(x: torch.Tensor, scale: Optional[torch.Tensor] = None, - dtype: Optional[torch.dtype] = None) -> torch.Tensor: + dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None) -> torch.Tensor: b, n = x.shape assert n % 2 == 0 @@ -1933,8 +1934,10 @@ def grid(meta: Mapping[str, int]) -> tuple[int, int]: x_ptr=x, x_stride=x.stride(0), d=d, + swiglu_limit=swiglu_limit or 0.0, BLOCK_SIZE=1024, HAS_O_SCALE=scale is not None, + HAS_SWIGLU_LIMIT=swiglu_limit is not None and swiglu_limit > 0.0, ) return o @@ -1945,6 +1948,7 @@ def _( x: torch.Tensor, scale: Optional[torch.Tensor] = None, dtype: Optional[torch.dtype] = None, + swiglu_limit: Optional[float] = None, ) -> torch.Tensor: b, n = x.shape diff --git a/tensorrt_llm/_torch/memory_buffer_utils.py b/tensorrt_llm/_torch/memory_buffer_utils.py index 3be29e30a745..9dfcca890278 100644 --- a/tensorrt_llm/_torch/memory_buffer_utils.py +++ b/tensorrt_llm/_torch/memory_buffer_utils.py @@ -130,6 +130,27 @@ def get_memory_buffers(): return _buffer +def clear_memory_buffers(): + """Clear all non-reserved buffers from the global buffer pool. + + This should be called after warmup phases (e.g., autotuner warmup) that + may allocate large workspace buffers sized for the maximum token count. + Clearing these prevents them from inflating the memory baseline used by + the KV cache profiler, which would otherwise reduce the memory available + for activations during inference. + """ + global _buffer + for buffer_name, blocks in list(_buffer.buffers.items()): + remaining = [b for b in blocks if b.is_reserved] + freed = [b for b in blocks if not b.is_reserved] + for block in freed: + del block.buffer + if remaining: + _buffer.buffers[buffer_name] = remaining + else: + del _buffer.buffers[buffer_name] + + _shared_pool = None diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index d1a2fa5503b8..e7830d8f9fb0 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -2,6 +2,7 @@ import errno import json import os +import struct import tempfile from dataclasses import dataclass, field from pathlib import Path @@ -13,14 +14,21 @@ from transformers.utils import HF_MODULES_CACHE from tensorrt_llm._torch.pyexecutor.config_utils import ( - get_qwen3_hybrid_num_attention_layers, is_nemotron_hybrid, is_qwen3_hybrid, - load_pretrained_config) -from tensorrt_llm._utils import (get_sm_version, is_sm_100f, - torch_dtype_to_binding) + get_qwen3_hybrid_num_attention_layers, + is_hybrid_linear, + is_nemotron_hybrid, + is_qwen3_hybrid, + load_pretrained_config, +) +from tensorrt_llm._utils import get_sm_version, is_sm_100f, torch_dtype_to_binding from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.functional import AllReduceStrategy -from tensorrt_llm.llmapi.llm_args import (DeepSeekSparseAttentionConfig, - KvCacheConfig, MoeLoadBalancerConfig) +from tensorrt_llm.llmapi.llm_args import ( + DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, + KvCacheConfig, + MoeLoadBalancerConfig, +) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -33,12 +41,18 @@ if TYPE_CHECKING: from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp - from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, LoraConfig, - SparseAttentionConfig, - SpeculativeConfig) + from tensorrt_llm.llmapi.llm_args import ( + DecodingBaseConfig, + LoraConfig, + SparseAttentionConfig, + SpeculativeConfig, + ) TConfig = TypeVar("TConfig", bound=transformers.PretrainedConfig) +_DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"} +_DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight" + def _unified_kv_pool_includes_mamba( is_disagg: bool, spec_config: Optional['SpeculativeConfig']) -> bool: @@ -322,6 +336,11 @@ def resolve_moe_backend(moe_backend: str, if moe_backend.upper() != "AUTO": return moe_backend + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + sm_version = get_sm_version() + if 100 <= sm_version < 120: + return "TRTLLM" + if architecture == "GptOssForCausalLM": sm_version = get_sm_version() # Select the best performing backend based on SM version @@ -524,6 +543,91 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): quant_config.exclude_modules = default_exclude return quant_config, layer_quant_config + @staticmethod + def _read_safetensors_header(path: Path) -> Dict[str, Any]: + with open(path, "rb") as f: + header_size = struct.unpack(" Optional[Dict]: + checkpoint_path = Path(checkpoint_dir) + candidates = [] + index_path = checkpoint_path / "model.safetensors.index.json" + if index_path.exists(): + with open(index_path) as f: + index = json.load(f) + shard_name = index.get("weight_map", {}).get(tensor_name) + if shard_name is not None: + candidates.append(checkpoint_path / shard_name) + + candidates.extend(sorted(checkpoint_path.glob("*.safetensors"))) + seen = set() + for candidate in candidates: + if candidate in seen or not candidate.exists(): + continue + seen.add(candidate) + header = ModelConfig._read_safetensors_header(candidate) + if tensor_name in header: + return header[tensor_name] + return None + + @staticmethod + def _detect_deepseek_v4_routed_moe_layout( + checkpoint_dir: str) -> Optional[str]: + tensor_info = ModelConfig._get_safetensors_header_for_tensor( + checkpoint_dir, _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT) + if tensor_info is None: + return None + + dtype = tensor_info.get("dtype") + shape = tensor_info.get("shape", []) + if dtype == "I8" and len(shape) == 2: + return "mxfp4" + if dtype is not None and dtype.startswith("F8_"): + return "fp8_block_scales" + return None + + @staticmethod + def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, + checkpoint_dir: str, + moe_backend: str, + layer_quant_config, + spec_config=None): + layout = ModelConfig._detect_deepseek_v4_routed_moe_layout( + checkpoint_dir) + if layout != "mxfp4": + return layer_quant_config + + experts_quant_config = QuantConfig() + experts_quant_config.quant_algo = ModelConfig.get_mxfp4_quant_algo( + moe_backend) + experts_quant_config.group_size = 32 + experts_quant_config.exclude_modules = [ + 'block.*.attn.out', 'block.*.mlp.gate', 'block.*.attn.qkv', + 'embedding', 'unembedding' + ] + + if layer_quant_config is None: + layer_quant_config = {} + else: + layer_quant_config = dict(layer_quant_config) + + num_moe_layers = pretrained_config.num_hidden_layers + if (spec_config is not None + and spec_config.spec_dec_mode.is_mtp_one_model()): + num_moe_layers += spec_config.num_nextn_predict_layers + + for layer_idx in range(num_moe_layers): + layer_quant_config[ + f"model.layers.{layer_idx}.mlp.experts"] = experts_quant_config + + logger.info( + "Detected DeepSeek-V4 routed MoE MXFP4 checkpoint layout; using " + "%s for routed experts.", experts_quant_config.quant_algo) + return layer_quant_config + @staticmethod def load_quant_config_from_dtypes_json(dtypes_json_file, moe_backend: str): quant_config = QuantConfig() @@ -581,6 +685,30 @@ def from_pretrained(cls, checkpoint_dir: str, trust_remote_code=False, **kwargs): + + def update_sparse_attention_indexer_config(pretrained_config, kwargs): + sparse_attention_config = kwargs.get('sparse_attention_config') + if sparse_attention_config: + index_n_heads = sparse_attention_config.index_n_heads or pretrained_config.index_n_heads + index_head_dim = sparse_attention_config.index_head_dim or pretrained_config.index_head_dim + index_topk = sparse_attention_config.index_topk or pretrained_config.index_topk + indexer_max_chunk_size = sparse_attention_config.indexer_max_chunk_size + skip_indexer_for_short_seqs = sparse_attention_config.skip_indexer_for_short_seqs + else: + index_n_heads = pretrained_config.index_n_heads + index_head_dim = pretrained_config.index_head_dim + index_topk = pretrained_config.index_topk + indexer_max_chunk_size = None + skip_indexer_for_short_seqs = True + indexer_config = {} + indexer_config['index_n_heads'] = index_n_heads + indexer_config['index_head_dim'] = index_head_dim + indexer_config['index_topk'] = index_topk + indexer_config['indexer_max_chunk_size'] = indexer_max_chunk_size + indexer_config[ + 'skip_indexer_for_short_seqs'] = skip_indexer_for_short_seqs + return indexer_config + # Use file lock to prevent race conditions when multiple processes # try to import/cache the same remote model config file with config_file_lock(): @@ -636,6 +764,65 @@ def from_pretrained(cls, indexer_rope_interleave=indexer_rope_interleave, enable_heuristic_topk=enable_heuristic_topk, indexer_k_dtype=indexer_k_dtype) + elif pretrained_config.architectures[ + 0] == "DeepseekV4ForCausalLM": + indexer_config = update_sparse_attention_indexer_config( + pretrained_config, kwargs) + checkpoint_compress_ratios = getattr( + pretrained_config, 'compress_ratios', None) + num_base_layers = pretrained_config.num_hidden_layers + spec_config = kwargs.get('spec_config', None) + mtp_enabled = ( + spec_config is not None + and spec_config.spec_dec_mode.is_mtp_one_model()) + sparse_attention_config = kwargs.get( + 'sparse_attention_config') + checkpoint_window_size = getattr( + pretrained_config, 'window_size', None) + if checkpoint_window_size is None: + checkpoint_window_size = getattr( + pretrained_config, 'sliding_window', None) + if sparse_attention_config: + compress_ratios = sparse_attention_config.compress_ratios + window_size = sparse_attention_config.window_size + if 'window_size' not in sparse_attention_config.model_fields_set: + window_size = checkpoint_window_size + else: + compress_ratios = checkpoint_compress_ratios + window_size = checkpoint_window_size + + if (checkpoint_compress_ratios is not None + and (compress_ratios is None + or len(checkpoint_compress_ratios) > + len(compress_ratios))): + compress_ratios = checkpoint_compress_ratios + + if window_size is None: + window_size = checkpoint_window_size + + # Normalize checkpoint-facing ratio 0 (SWA-only/uncompressed) + # to 1 internally so cache allocation math works. The + # external config keeps the original semantics. + compress_ratios = [ + ratio if ratio > 0 else 1 for ratio in compress_ratios + ] + + # Only synthesize ratios for extra MTP layers. The base + # model ratios must come from the checkpoint or an + # explicit user override; padding a short default list for + # non-MTP changes sparse attention semantics. + if mtp_enabled: + mtp_num_layers = spec_config.num_nextn_predict_layers + total_layers = num_base_layers + mtp_num_layers + if len(compress_ratios) < total_layers: + compress_ratios = list(compress_ratios) + [1] * ( + total_layers - len(compress_ratios)) + + kwargs[ + 'sparse_attention_config'] = DeepSeekV4SparseAttentionConfig( + compress_ratios=compress_ratios, + window_size=window_size, + **indexer_config) else: raise ValueError( "checkpoint_dir is None. Cannot load model config without a valid checkpoint directory." @@ -753,6 +940,11 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config=quant_config, ) + if architecture in _DEEPSEEK_V4_ARCHITECTURES: + layer_quant_config = cls._set_deepseek_v4_routed_moe_quant_config( + pretrained_config, checkpoint_dir, moe_backend, + layer_quant_config, kwargs.get('spec_config', None)) + model_config = cls(pretrained_config=pretrained_config, quant_config=quant_config, quant_config_dict=layer_quant_config, @@ -907,8 +1099,7 @@ def _infer_nemotron_ffn_mult(self): for x in self.pretrained_config.block_configs ]) - from tensorrt_llm._torch.models.modeling_nemotron_nas import \ - _ffn_mult_to_intermediate_size + from tensorrt_llm._torch.models.modeling_nemotron_nas import _ffn_mult_to_intermediate_size mlp_hidden_size = _ffn_mult_to_intermediate_size( biggest_ffn_mult, self.pretrained_config.hidden_size) diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 20e65ab1c119..463ec8ff344e 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -27,6 +27,7 @@ from .modeling_laguna import LagunaForCausalLM from .modeling_llama import LlamaForCausalLM from .modeling_llava_next import LlavaNextModel +from .modeling_deepseekv4 import DeepseekV4ForCausalLM from .modeling_minimaxm2 import MiniMaxM2ForCausalLM from .modeling_mistral import Mistral3VLM, MistralForCausalLM from .modeling_mixtral import MixtralForCausalLM @@ -78,6 +79,7 @@ "Mistral3VLM", "MistralForCausalLM", "MixtralForCausalLM", + "DeepseekV4ForCausalLM", "NemotronH_Nano_VL_V2", "NemotronForCausalLM", "NemotronHForCausalLM", diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py new file mode 100644 index 000000000000..350d5956c99e --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -0,0 +1,2521 @@ +# -------------------------------------------------- +# Portions of this code were derived from DeepSeek‑V3: +# https://github.com/deepseek-ai/DeepSeek-V3 +# +# MIT License + +# Copyright (c) 2023 DeepSeek + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# -------------------------------------------------- + +import copy +import math +import os +from typing import TYPE_CHECKING, Dict, List, Optional + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + +import torch +import triton +import triton.language as tl +from torch import nn +from tqdm import tqdm +from transformers import PretrainedConfig + +import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils +from tensorrt_llm._ipc_utils import can_access_peer +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams +from ..attention_backend.sparse.deepseek_v4.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata +from ..distributed import ( + AllReduce, + AllReduceFusionOp, + AllReduceParams, + MoEAllReduce, + MoEAllReduceParams, + allgather, +) +from ..model_config import ModelConfig +from ..modules.attention import MLA +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding +from ..modules.engram import Engram, EngramConfig, EngramHashProvider +from ..modules.fused_moe import DeepSeekV4MoeRoutingMethod, MoE, MoEWeightLoadingMode, create_moe +from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE +from ..modules.linear import Linear +from ..modules.mhc.hyper_connection import HCHead, HCState, mHC + +# isort: off +from ..modules.fused_moe.routing import ( + get_cached_perfect_router_logits, + precompute_common_perfect_router_logits, +) + +# isort: on +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import TensorParallelMode, WeightsLoadingConfig +from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.rms_norm import RMSNorm +from ..peft.lora.layer import LoraLayer +from ..speculative import SpecMetadata +from ..utils import AuxStreamType, EventType, Fp4QuantizedTensor, create_lm_head_tp_mapping +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, EagerFusionConfig, filter_weights, register_auto_model + + +@triton.jit +def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): + """ + Dequantizes weights using the provided scaling factors and stores the result. + + Args: + x_ptr (tl.pointer): Pointer to the quantized weights. + s_ptr (tl.pointer): Pointer to the scaling factors. + y_ptr (tl.pointer): Pointer to the output buffer for dequantized weights. + M (int): Number of rows in the weight matrix. + N (int): Number of columns in the weight matrix. + BLOCK_SIZE (tl.constexpr): Size of the block for tiling. + + Returns: + None + """ + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + n = tl.cdiv(N, BLOCK_SIZE) + offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offs = offs_m[:, None] * N + offs_n[None, :] + mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + x = tl.load(x_ptr + offs, mask=mask).to(tl.float32) + s = tl.load(s_ptr + pid_m * n + pid_n) + y = x * s + tl.store(y_ptr + offs, y, mask=mask) + + +def weight_dequant(x: torch.Tensor, s: torch.Tensor, block_size: int = 128) -> torch.Tensor: + """ + Dequantizes the given weight tensor using the provided scale tensor. + + Args: + x (torch.Tensor): The quantized weight tensor of shape (M, N). + s (torch.Tensor): The scale tensor of shape (M, N). + block_size (int, optional): The block size to use for dequantization. Defaults to 128. + + Returns: + torch.Tensor: The dequantized weight tensor of the same shape as `x`. + + Raises: + AssertionError: If `x` or `s` are not contiguous or if their dimensions are not 2. + """ + assert x.is_contiguous() and s.is_contiguous(), "Input tensors must be contiguous" + assert x.dim() == 2 and s.dim() == 2, "Input tensors must have 2 dimensions" + if s.dtype != torch.float32: + s = s.float() + M, N = x.size() + y = torch.empty_like(x, dtype=torch.get_default_dtype()) + grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) # noqa: E731 + weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size) + return y + + +def _deepseek_v4_pos_embd_params( + config: PretrainedConfig, model_config: ModelConfig, layer_idx: Optional[int] +) -> PositionalEmbeddingParams: + rope_params = RopeParams.from_config(config) + + compress_ratios = None + if model_config.sparse_attention_config is not None: + compress_ratios = getattr(model_config.sparse_attention_config, "compress_ratios", None) + if not compress_ratios: + compress_ratios = getattr(config, "compress_ratios", None) + + compress_ratio = 0 + if compress_ratios and layer_idx is not None: + compress_ratio = compress_ratios[min(layer_idx, len(compress_ratios) - 1)] + + if compress_ratio > 1: + rope_params.theta = getattr(config, "compress_rope_theta", rope_params.theta) + rope_params.scale_type = RotaryScalingType.yarn + # DeepSeek-V4 reference applies YaRN frequency interpolation but does + # not scale the cos/sin amplitudes. + rope_params.mscale = 0.0 + rope_params.mscale_all_dim = 0.0 + pos_type = PositionEmbeddingType.yarn + else: + # DeepSeek-V4 reference uses base RoPE, without YaRN, for SWA-only + # layers. Internal configs normalize checkpoint ratio 0 to 1, so both + # values are treated as uncompressed here. + rope_params.theta = getattr(config, "rope_theta", rope_params.theta) + rope_params.scale_type = RotaryScalingType.none + rope_params.scale = 1.0 + pos_type = PositionEmbeddingType.rope_gptj + + return PositionalEmbeddingParams( + type=pos_type, + rope=rope_params, + is_neox=False, + ) + + +@torch.compile(dynamic=True) +def moe_reduce_add_shared_output(routed_output, shared_output): + routed_output = torch.sum(routed_output, dim=1, keepdim=False) + return shared_output + routed_output + + +# Per-attention parameter renames inside `attn.` / `mtp.0.attn.`. +# Maps checkpoint name component → model name component. +_ATTN_PARAM_RENAME = { + "wq_a": "q_a_proj", + "wq_b": "q_b_proj", + "wkv": "kv_a_proj_with_mqa", + "wo_b": "o_b_proj", + "q_norm": "q_a_layernorm", + "kv_norm": "kv_a_layernorm", +} + +# Shared expert leaf rename: checkpoint w1/w3/w2 → model gate/up/down. The +# loader's `params_map` then fuses `gate_proj` + `up_proj` into `gate_up_proj`. +_SHARED_EXPERT_RENAME = { + "w1": "gate_proj", + "w3": "up_proj", + "w2": "down_proj", +} + + +def _remap_deepseek_v4_checkpoint_keys( + weights: Dict, num_hidden_layers: int, kv_lora_rank: int = 448 +) -> Dict: + """Convert DeepSeek-V4 checkpoint keys to model named-parameter keys. + + Why: the upstream DS-V4 release uses keys like ``layers.X.attn.wkv.weight``, + ``mtp.0.ffn.experts.0.w1.scale``, ``embed.weight``, ``head.weight``. The + TRT-LLM model exposes them as ``model.layers.X.self_attn.kv_a_proj_with_mqa.weight``, + ``model.layers.{N}.mlp.experts.0.w1.weight_scale_inv``, ``model.embed_tokens.weight``, + ``lm_head.weight``. This pass renames keys, fuses split projections that the + model represents as one tensor, and synthesizes default values for params + the model has but the checkpoint does not (so the existing loader's strict + key lookup does not raise). + + Caveats — limitations carried as TODOs (kept here so future readers can + audit them in one place rather than chasing comments): + * ``self_attn.indexer.wk.weight`` and ``self_attn.indexer.k_norm.{weight,bias}`` + are zero-filled — V4 indexer's k path is served by the compressor, so + the base ``Indexer.wk`` / ``k_norm`` are unused at forward time. + * Routed experts can use either FP8 block scales or the packed MXFP4 + layout. The first routed expert weight determines the scale suffix used + for all routed expert tensors in the shard. + * ``self_attn.o_a_proj`` is loaded by the DeepSeek-V4 loader because it + is a direct MLA parameter, not a child Linear module. + * ``mtp.0.head.weight`` is dropped — DeepSeekV4MTP reuses the main + ``lm_head`` via ``shared_head``. Flash omits this key entirely; Flash-Base + carries it but matches the main head, so we let the main head win. + """ + mtp_layer_prefix = f"model.layers.{num_hidden_layers}" + routed_moe_scale_name = "weight_scale_inv" + for key, value in weights.items(): + if ( + key.startswith("layers.") + and ".ffn.experts." in key + and key.endswith(".weight") + and getattr(value, "ndim", 0) == 2 + and value.dtype in (torch.int8, torch.uint8) + ): + routed_moe_scale_name = "weight_scale" + break + + def _rename_attn_subkey(rest: str) -> Optional[str]: + # rest examples: "wq_a.weight", "wq_a.scale", "wo_a.weight", + # "attn_sink", "compressor.wkv.weight", "indexer.wq_b.scale", + # "kv_norm.weight" + # ``attn_sink`` is loaded by the ``mqa`` branch in the per-module + # loader, which reads it under the parent ``self_attn.attn_sink`` + # key. Pass through unchanged. + if rest == "attn_sink": + return "attn_sink" + # `wo_a` is an nn.Parameter on the model side (not a Linear), so + # `wo_a.weight` carries the value directly into `o_a_proj` without + # a trailing ``.weight``. Retain `.scale` so the loader can dequantize + # FP8 block-scaled checkpoints before assigning the bf16 parameter. + if rest == "wo_a.weight": + return "o_a_proj" + if rest == "wo_a.scale": + return "o_a_proj.weight_scale_inv" + # Compressor / indexer paths — pass through with .scale rename, plus + # wkv+wgate fusion handled separately below. + if rest.startswith("compressor.") or rest.startswith("indexer."): + return rest.replace(".scale", ".weight_scale_inv") + head, sep, tail = rest.partition(".") + new_head = _ATTN_PARAM_RENAME.get(head, head) + if tail == "scale": + tail = "weight_scale_inv" + return f"{new_head}.{tail}" if sep else new_head + + def _rename_ffn_subkey(rest: str) -> str: + # Examples: + # gate.weight / gate.tid2eid → gate.weight / gate.tid2eid + # gate.bias → gate.e_score_correction_bias + # experts... → experts... + # shared_experts.. → shared_experts._proj. + if rest == "gate.bias": + return "gate.e_score_correction_bias" + if rest.startswith("experts.") and rest.endswith(".scale"): + return f"{rest[: -len('.scale')]}.{routed_moe_scale_name}" + rest = rest.replace(".scale", ".weight_scale_inv") + # Non-hashed layers carry the routing logit bias as `gate.bias`; the + # model wires it through `DeepseekV4Gate.e_score_correction_bias`. + if rest == "gate.bias": + return "gate.e_score_correction_bias" + if rest.startswith("shared_experts."): + parts = rest.split(".") + if len(parts) >= 2 and parts[1] in _SHARED_EXPERT_RENAME: + parts[1] = _SHARED_EXPERT_RENAME[parts[1]] + rest = ".".join(parts) + return rest + + def _rename_layer_subkey(rest: str) -> Optional[str]: + # rest examples: "attn_norm.weight", "ffn_norm.weight", + # "hc_attn_fn", "attn.wkv.weight", "ffn.experts.0.w1.weight" + if rest == "attn_norm.weight": + return "input_layernorm.weight" + if rest == "ffn_norm.weight": + return "post_attention_layernorm.weight" + # ``hc_attn_*`` and ``hc_ffn_*`` are loaded by ``load_flat_hc_weights`` + # which builds the lookup key as ``f"{stem}_{a}"`` with the parent + # module path as the stem (e.g. ``model.layers.0.hc_attn_fn``). Pass + # the flat-underscore form through unchanged so that lookup succeeds. + if rest.startswith("hc_attn_") or rest.startswith("hc_ffn_"): + return rest + if rest.startswith("attn."): + new_sub = _rename_attn_subkey(rest[len("attn.") :]) + return None if new_sub is None else f"self_attn.{new_sub}" + if rest.startswith("ffn."): + return f"mlp.{_rename_ffn_subkey(rest[len('ffn.') :])}" + return rest + + out: Dict[str, torch.Tensor] = {} + # Pending fusions: collected first, materialized at the end so we don't + # depend on iteration order. + compressor_split: Dict[str, Dict[str, torch.Tensor]] = {} + + def _record_compressor_part(model_key: str, part: str, tensor: torch.Tensor): + # model_key looks like "...self_attn.compressor..weight" or with + # ".indexer.compressor.". Strip trailing "..weight" → "." + # is wgate or wkv; we want a stable "fusion bucket" key that ends at + # the parent compressor. + bucket = model_key.rsplit(f".{part}.", 1)[0] + compressor_split.setdefault(bucket, {})[part] = tensor + + def _emit_or_collect(model_key: str, tensor: torch.Tensor): + """Route a (model_key, tensor) pair and collect compressor fusions.""" + if ".compressor." in model_key and ( + model_key.endswith(".wkv.weight") or model_key.endswith(".wgate.weight") + ): + part = "wkv" if model_key.endswith(".wkv.weight") else "wgate" + _record_compressor_part(model_key, part, tensor) + return + if ( + routed_moe_scale_name == "weight_scale" + and ".mlp.experts." in model_key + and (model_key.endswith(".weight") or model_key.endswith(".weight_scale")) + and tensor.dtype != torch.uint8 + ): + tensor = tensor.view(torch.uint8) + out[model_key] = tensor + + for k, v in weights.items(): + # Top-level keys that don't go through the layer/mtp branches. + if k == "embed.weight": + out["model.embed_tokens.weight"] = v + continue + if k == "head.weight": + out["lm_head.weight"] = v + continue + if k == "norm.weight": + out["model.norm.weight"] = v + continue + if k.startswith("hc_head_"): + # ``load_flat_hc_weights`` looks up ``f"{stem}_{a}"`` with stem set + # to the module's full name path (``model.hc_head``). Emit the flat + # checkpoint key under the parent prefix so the lookup matches. + out[f"model.{k}"] = v + continue + + # mtp.0.head.weight is intentionally dropped (Flash-Base only); see + # docstring caveats. + if k == "mtp.0.head.weight": + continue + + # mtp.0.* — route to model.layers.{num_hidden_layers}.* + if k.startswith("mtp.0."): + rest = k[len("mtp.0.") :] + # MTP-only keys: enorm, hnorm map directly; norm maps to + # shared_head.norm; hc_head_* maps under shared_head; e_proj / + # h_proj are loaded as two separate Linear modules (no fused + # eh_proj on the MTP layer). + if rest in ("enorm.weight", "hnorm.weight"): + out[f"{mtp_layer_prefix}.{rest}"] = v + continue + if rest == "norm.weight": + out[f"{mtp_layer_prefix}.shared_head.norm.weight"] = v + continue + if rest.startswith("hc_head_"): + # The MTP HCHead is wired at ``...shared_head.hc_head``, so + # ``load_flat_hc_weights`` matches ``names[-1] == "hc_head"`` + # first and computes ``stem = ".".join(names)`` with the full + # module path. Emit at that full path with the flat suffix + # so the lookup matches. + out[f"{mtp_layer_prefix}.shared_head.{rest}"] = v + continue + for proj in ("e_proj", "h_proj"): + if rest.startswith(f"{proj}."): + suffix = rest[len(f"{proj}.") :] + if suffix == "scale": + suffix = "weight_scale_inv" + out[f"{mtp_layer_prefix}.{proj}.{suffix}"] = v + break + else: + # General per-layer transform reused for the MTP layer. + new_rest = _rename_layer_subkey(rest) + if new_rest is None: + continue + _emit_or_collect(f"{mtp_layer_prefix}.{new_rest}", v) + continue + + # layers..* — route to model.layers..* + if k.startswith("layers."): + parts = k.split(".", 2) + if len(parts) < 3: + continue + layer_idx, rest = parts[1], parts[2] + new_rest = _rename_layer_subkey(rest) + if new_rest is None: + continue + _emit_or_collect(f"model.layers.{layer_idx}.{new_rest}", v) + continue + + # Anything else: pass through. This catches future top-level keys + # added by the upstream release without silently dropping them. + out[k] = v + + # Materialize compressor wkv_gate fusion: model has a single Linear with + # out_features = state_dim * 2; checkpoint splits it as `wkv` (kv half) + # and `wgate` (gate half). Concatenating wkv first matches the order the + # compressor kernels expect (kv_score = [kv | gate]). + for bucket, parts in compressor_split.items(): + if "wkv" not in parts or "wgate" not in parts: + # Partial — emit what we have so the loader fails loudly with a + # specific missing key rather than a silent shape mismatch. + for name, tensor in parts.items(): + out[f"{bucket}.{name}.weight"] = tensor + continue + out[f"{bucket}.wkv_gate.weight"] = torch.cat([parts["wkv"], parts["wgate"]], dim=0) + + return out + + +class DeepseekV4WeightLoader: + def __init__(self, model, is_draft_model: bool = False): + self.model = model + self.config = model.config + self.model_config = model.model_config + self.is_draft_model = is_draft_model + + def load_weights(self, weights: Dict, skip_modules: List[str] = []): + # If the checkpoint uses raw DS-V4 keys (layers.X.attn.wkv.weight, + # mtp.0.*, embed.weight, head.weight), rewrite them to the model's + # named-parameter keys before iterating modules. The detection is by + # presence of any top-level "layers." key; HF-style checkpoints use + # "model.layers." and skip this branch. + if any(k == "embed.weight" or k.startswith("layers.") for k in weights): + weights = _remap_deepseek_v4_checkpoint_keys( + weights, + num_hidden_layers=self.config.num_hidden_layers, + kv_lora_rank=self.config.kv_lora_rank, + ) + # Synthesize defaults (with correct shape pulled from the model) + # for parameters the model has but the V4 checkpoint omits. We do + # this in one place vs scattering zero-fills through the per- + # module branches so the missing-key contract is auditable here. + # indexer.k_norm.weight → ones + # indexer.k_norm.bias → zeros + # indexer.wk.weight → zeros (V4 indexer's k path is + # served by compressor; wk unused) + _ones_suffixes = ("self_attn.indexer.k_norm.weight",) + _zeros_suffixes = ( + "self_attn.indexer.k_norm.bias", + "self_attn.indexer.wk.weight", + ) + model_params = dict(self.model.named_parameters()) + for pname, p in model_params.items(): + if pname in weights: + continue + if any(pname.endswith(s) for s in _ones_suffixes): + weights[pname] = torch.ones_like(p, device="cpu") + elif any(pname.endswith(s) for s in _zeros_suffixes): + weights[pname] = torch.zeros_like(p, device="cpu") + + def requantize_weight_with_new_scale( + weight, weight_scale, old_scale_2, new_scale_2, device + ): + """ + Dequantize FP4 weights and requantize with a new scale. + + Args: + weight: FP4 quantized weight tensor 2D [,] + weight_scale: FP8 per-block scaling factors + old_scale_2: original global scale (amax/(448*6)) + new_scale_2: new global scale (amax/(448*6)) + device: target device for computation + + Returns: + (requantized_weight, new_weight_scale) + """ + # Remember original dtype of weight_scale + original_scale_dtype = weight_scale.dtype + original_scale_shape = weight_scale.shape + + # Dequantize + dequant_shape = (weight.shape[0], weight.shape[1] * 2) + weight_dequant = ( + torch.ops.tensorrt_llm.e2m1_and_ufp8sf_scale_to_float_v2( + weight.contiguous(), + weight_scale.flatten().view(fp4_utils.float4_sf_dtype).contiguous(), + old_scale_2, + 16, + 1, + True, + ) + .to(dtype=torch.bfloat16) + .reshape(dequant_shape) + ) + + # Requantize using the new_scale_2 + weight_requant, weight_scale_requant = torch.ops.trtllm.fp4_quantize( + weight_dequant.to(device), + 1.0 / new_scale_2.to(device), + 16, # scaling_vector_size + False, + ) + + # Ensure the returned scale has the same dtype as the input scale + return weight_requant.cpu(), weight_scale_requant.reshape(original_scale_shape).view( + original_scale_dtype + ).cpu() + + def rename_moe_weight(weights: Dict, rename_rules: Dict): + result = {} + for key, value in weights.items(): + new_key = key + for old, new in rename_rules.items(): + new_key = new_key.replace(old, new) + result[new_key] = value + return result + + ## Prepare weights for TP + def split(v, tp_size, idx, dim=0): + if tp_size == 1: + return v + if len(v.shape) == 1: + return torch.chunk(v, tp_size)[idx].contiguous() + return torch.chunk(v, tp_size, dim=dim)[idx].contiguous() + + def split_matrix_tp(v, tensor_parallel, rank, dim): + return split(v, tensor_parallel, rank, dim=dim) + + def load_kv_b_proj_and_k_b_proj_trans(module_name: str, is_scale: bool) -> torch.Tensor: + weight_name = "weight" if not is_scale else "weight_scale_inv" + local_qk_nope_head_dim = qk_nope_head_dim if not is_scale else qk_nope_head_dim // 128 + local_v_head_dim = v_head_dim if not is_scale else v_head_dim // 128 + local_kv_lora_rank = kv_lora_rank if not is_scale else kv_lora_rank // 128 + + kv_b_proj = weights[f"{module_name}.{weight_name}"][:].unflatten( + 0, + [ + num_heads, + local_qk_nope_head_dim + local_v_head_dim, + ], + ) + + if not self.model_config.mapping.enable_attention_dp: + kv_b_proj = split_matrix_tp(kv_b_proj, tp_size, tp_rank, 0) + k_nope_weight, v_weight = kv_b_proj.split( + [local_qk_nope_head_dim, local_v_head_dim], + dim=1, + ) + weight_divisor = 1 if self.model_config.mapping.enable_attention_dp else tp_size + local_num_heads = num_heads // weight_divisor + + k_nope_weight_trans = k_nope_weight.transpose(2, 1).contiguous() + + kv_b_proj = torch.concat( + [ + k_nope_weight.reshape( + local_num_heads * local_qk_nope_head_dim, local_kv_lora_rank + ), + v_weight.reshape(local_num_heads * local_v_head_dim, local_kv_lora_rank), + ], + dim=0, + ) + + return kv_b_proj, k_nope_weight_trans + + def load_kv_b_proj_and_k_b_proj_trans_dequant(module_name: str) -> torch.Tensor: + weight_name = "weight" + local_qk_nope_head_dim = qk_nope_head_dim + local_v_head_dim = v_head_dim + local_kv_lora_rank = kv_lora_rank + + kv_b_proj = weights[f"{module_name}.{weight_name}"][:].cuda() + + weight_name = "weight_scale_inv" + kv_b_proj_scale = weights[f"{module_name}.{weight_name}"][:].cuda() + + kv_b_proj = weight_dequant(kv_b_proj, kv_b_proj_scale) + kv_b_proj = kv_b_proj.unflatten( + 0, + [ + num_heads, + local_qk_nope_head_dim + local_v_head_dim, + ], + ) + if not self.model_config.mapping.enable_attention_dp: + kv_b_proj = split_matrix_tp(kv_b_proj, tp_size, tp_rank, 0) + k_nope_weight, v_weight = kv_b_proj.split( + [local_qk_nope_head_dim, local_v_head_dim], + dim=1, + ) + weight_divisor = 1 if self.model_config.mapping.enable_attention_dp else tp_size + local_num_heads = num_heads // weight_divisor + + k_nope_weight_trans = k_nope_weight.transpose(2, 1).contiguous() + + kv_b_proj = torch.concat( + [ + k_nope_weight.reshape( + local_num_heads * local_qk_nope_head_dim, local_kv_lora_rank + ), + v_weight.reshape(local_num_heads * local_v_head_dim, local_kv_lora_rank), + ], + dim=0, + ) + + return kv_b_proj, k_nope_weight_trans + + def load_o_a_proj(module_name: str, module) -> None: + weight_name = f"{module_name}.o_a_proj" + o_a_proj = weights[weight_name][:].unflatten( + 0, + [ + module.num_groups, + module.o_lora_rank, + ], + ) + + scale_name = f"{weight_name}.weight_scale_inv" + o_a_proj_scale = weights.get(scale_name) + if o_a_proj_scale is not None: + o_a_proj_scale = o_a_proj_scale[:].unflatten( + 0, + [ + module.num_groups, + module.o_lora_rank // 128, + ], + ) + elif o_a_proj.dtype == torch.float8_e4m3fn: + raise KeyError(f"Missing FP8 block scale for {weight_name}") + + if not self.model_config.mapping.enable_attention_dp: + o_a_proj = split_matrix_tp(o_a_proj, tp_size, tp_rank, 0) + if o_a_proj_scale is not None: + o_a_proj_scale = split_matrix_tp(o_a_proj_scale, tp_size, tp_rank, 0) + + if o_a_proj_scale is not None: + o_a_proj = weight_dequant( + o_a_proj.reshape(-1, o_a_proj.shape[-1]).contiguous().cuda(), + o_a_proj_scale.reshape(-1, o_a_proj_scale.shape[-1]).contiguous().cuda(), + ).view(o_a_proj.shape) + + module.o_a_proj.data.copy_( + o_a_proj.reshape(module.o_a_proj.shape).to(module.o_a_proj.dtype) + ) + + if getattr(module, "o_a_proj_scale", None) is not None and o_a_proj_scale is not None: + module.o_a_proj_scale.data.copy_( + o_a_proj_scale.reshape(module.o_a_proj_scale.shape) + ) + + if getattr(module, "o_a_proj_dequant", None) is not None and o_a_proj_scale is not None: + module.o_a_proj_dequant.data.copy_( + o_a_proj.reshape(module.o_a_proj_dequant.shape).to( + module.o_a_proj_dequant.dtype + ) + ) + + def split_kv_b_proj(kv_b_proj: torch.Tensor, is_scale: bool) -> torch.Tensor: + local_qk_nope_head_dim = qk_nope_head_dim if not is_scale else qk_nope_head_dim // 128 + local_v_head_dim = v_head_dim if not is_scale else v_head_dim // 128 + + weight_divisor = 1 if self.model_config.mapping.enable_attention_dp else tp_size + local_num_heads = num_heads // weight_divisor + + k_b_proj, v_b_proj = kv_b_proj.split( + [local_num_heads * local_qk_nope_head_dim, local_num_heads * local_v_head_dim], + dim=0, + ) + k_b_proj = k_b_proj.view([local_num_heads, local_qk_nope_head_dim, -1]) + v_b_proj = v_b_proj.view([local_num_heads, local_v_head_dim, -1]) + + if cp_size > 1: + local_cp_heads = local_num_heads // cp_size + k_b_proj = k_b_proj[cp_rank * local_cp_heads : (cp_rank + 1) * local_cp_heads] + v_b_proj = v_b_proj[cp_rank * local_cp_heads : (cp_rank + 1) * local_cp_heads] + + return k_b_proj, v_b_proj + + is_lite = self.config.q_lora_rank is None + num_heads = self.config.num_attention_heads + qk_nope_head_dim = getattr(self.config, "qk_nope_head_dim", None) + v_head_dim = getattr(self.config, "v_head_dim", getattr(self.config, "head_dim", None)) + kv_lora_rank = getattr(self.config, "kv_lora_rank", None) + + tp_rank = self.model_config.mapping.tp_rank + tp_size = self.model_config.mapping.tp_size + cp_rank = self.model_config.mapping.cp_rank + cp_size = self.model_config.mapping.cp_size + + params_map = {"gate_up_proj": ["gate_proj", "up_proj"]} + all_named_modules = dict(self.model.named_modules()) + + def load_flat_hc_weights(module, names: List[str]) -> bool: + """Load mHC / HCHead from flat ckpt keys: ``_{fn,base,scale}``. + + V4 / DeepSeek-V4 checkpoints store these as flat names (e.g. + ``hc_attn_fn``) rather than structured (``hc_attn.fn``). + ``shared_head.hc_head`` (MTP) is rewritten to flat ``hc_head_*`` + under the same parent. + """ + if names[-1] in ("hc_attn", "hc_ffn", "hc_head"): + stem = ".".join(names) + elif names[-2:] == ["shared_head", "hc_head"]: + stem = ".".join(names[:-2] + ["hc_head"]) + else: + return False + keys = {a: f"{stem}_{a}" for a in ("fn", "base", "scale")} + if not all(k in weights for k in keys.values()): + return False + for attr, key in keys.items(): + getattr(module, attr).data.copy_(weights[key][:]) + return True + + for name, module in tqdm(all_named_modules.items(), desc="Loading weights"): + if name.startswith("draft_model"): + continue + names = name.split(".") + parent_module_name = ".".join(names[:-1]) + + if names[-1] == "mqa": + parent_attn_name = ".".join(names[:-1]) + attn_sink_key = f"{parent_attn_name}.attn_sink" + if attn_sink_key in weights: + sink_full = weights[attn_sink_key][:] + if not self.model_config.mapping.enable_attention_dp: + sink_full = split(sink_full, tp_size, tp_rank) + sink_full = sink_full.to(torch.float32).cuda() + module.attn_sink = nn.Parameter(sink_full, requires_grad=False) + continue + + if len(module._parameters) <= 0: + continue + elif any(skip_module in name for skip_module in skip_modules): + continue + else: + if "model.layers" in name and int(names[2]) >= self.config.num_hidden_layers: + mtp_layer_idx = int(names[2]) - self.config.num_hidden_layers + names[2] = str( + mtp_layer_idx % self.config.num_nextn_predict_layers + + self.config.num_hidden_layers + ) + name = ".".join(names) + if names[-1] == "kv_b_proj": + # TODO: remove weight_dequant after enabling fp8_bmm + dequant_kv_b_proj = ( + self.model_config.quant_config.is_module_excluded_from_quantization( + names[-1] + ) + ) + if dequant_kv_b_proj: + kv_b_proj, k_b_proj_trans = load_kv_b_proj_and_k_b_proj_trans_dequant(name) + else: + kv_b_proj, k_b_proj_trans = load_kv_b_proj_and_k_b_proj_trans( + name, is_scale=False + ) + module.weight.data.copy_(kv_b_proj.reshape(module.weight.shape)) + + attn_module = all_named_modules[parent_module_name] + _, v_b_proj = split_kv_b_proj(module.weight.data, is_scale=False) + attn_module.v_b_proj = nn.Parameter(v_b_proj, requires_grad=False) + + attn_module.k_b_proj_trans.data.copy_( + k_b_proj_trans.reshape(attn_module.k_b_proj_trans.shape) + ) + + if getattr(module, "weight_scale", None) is not None and not dequant_kv_b_proj: + kv_b_proj_scale, k_b_proj_trans_scale = load_kv_b_proj_and_k_b_proj_trans( + name, is_scale=True + ) + module.weight_scale.copy_( + kv_b_proj_scale.reshape(module.weight_scale.shape) + ) + attn_module.k_b_proj_trans_scale.copy_( + k_b_proj_trans_scale.reshape(attn_module.k_b_proj_trans_scale.shape) + ) + + _, v_b_proj_scale = split_kv_b_proj(module.weight_scale.data, is_scale=True) + attn_module.v_b_proj_scale = nn.Parameter( + v_b_proj_scale, requires_grad=False + ) + + if attn_module.k_b_proj_trans_dequant is not None: + attn_module.k_b_proj_trans_dequant.data.copy_( + weight_dequant( + k_b_proj_trans.view(-1, k_b_proj_trans.shape[-1]).cuda(), + k_b_proj_trans_scale.view( + -1, k_b_proj_trans_scale.shape[-1] + ).cuda(), + ) + .view(*attn_module.k_b_proj_trans_dequant.shape) + .to(attn_module.k_b_proj_trans_dequant.dtype) + ) + if attn_module.v_b_proj_dequant is not None: + attn_module.v_b_proj_dequant.data.copy_( + weight_dequant( + v_b_proj.view(-1, v_b_proj.shape[-1]).cuda(), + v_b_proj_scale.view(-1, v_b_proj_scale.shape[-1]).cuda(), + ) + .view(*attn_module.v_b_proj_dequant.shape) + .to(attn_module.v_b_proj_dequant.dtype) + ) + elif names[-1] == "kv_a_proj_with_mqa": + nvfp4_fused_a = ( + self.model_config.get_quant_config().layer_quant_mode.has_nvfp4() + and weights[f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight"].dtype + == fp4_utils.float4_e2m1x2 + and weights[f"{'.'.join(names[:-1])}.q_a_proj.weight"].dtype + == fp4_utils.float4_e2m1x2 + ) + if nvfp4_fused_a: + ########### input_scale + kv_a_proj_with_mqa_input_scale = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.input_scale" + ] + if not is_lite: + q_a_proj_input_scale = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.input_scale" + ] + assert kv_a_proj_with_mqa_input_scale == q_a_proj_input_scale, ( + "kv_a_proj_with_mqa.input_scale and q_a_proj.input_scale should be the same" + ) + # modelopt ckpt stores amax/(448*6), convert to (448*6)/amax + shared_input_scale = kv_a_proj_with_mqa_input_scale + module.input_scale.data.copy_(1.0 / shared_input_scale) + E2M1_MAX = 6.0 + module.inv_input_scale.data.copy_(module.input_scale / E2M1_MAX) + ########### weight_scale_2 + need_requant_kv_a_proj_with_mqa = False + need_requant_q_a_proj = False + kv_a_proj_with_mqa_scale_2 = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale_2" + ] + shared_weight_scale_2 = kv_a_proj_with_mqa_scale_2 + if not is_lite: + q_a_proj_scale_2 = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.weight_scale_2" + ] + if kv_a_proj_with_mqa_scale_2 < q_a_proj_scale_2: + shared_weight_scale_2 = q_a_proj_scale_2 + need_requant_kv_a_proj_with_mqa = True + elif q_a_proj_scale_2 < kv_a_proj_with_mqa_scale_2: + need_requant_q_a_proj = True + + ########### alpha + alpha = shared_input_scale.float() * shared_weight_scale_2.float() + module.alpha.data.copy_(alpha) + module.scalar_alpha = alpha.item() + + ########### weights + kv_a_proj_with_mqa = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight" + ][:] + + if not is_lite: + q_a_proj = weights[f"{'.'.join(names[:-1])}.q_a_proj.weight"][:] + + ########### weight_scale + kv_a_proj_with_mqa_scale = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale" + ][:] + kv_a_proj_with_mqa_scale = torch.ops.trtllm.block_scale_interleave( + kv_a_proj_with_mqa_scale.view(fp4_utils.float4_sf_dtype) + ) + if not is_lite: + q_a_proj_scale = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.weight_scale" + ][:] + q_a_proj_scale = torch.ops.trtllm.block_scale_interleave( + q_a_proj_scale.view(fp4_utils.float4_sf_dtype) + ) + + ########### requantize + if need_requant_kv_a_proj_with_mqa: + # requant kv_a_proj_with_mqa + kv_a_proj_with_mqa, kv_a_proj_with_mqa_scale = ( + requantize_weight_with_new_scale( + kv_a_proj_with_mqa, + kv_a_proj_with_mqa_scale, + kv_a_proj_with_mqa_scale_2, + shared_weight_scale_2, + device=module.weight.device, + ) + ) + if need_requant_q_a_proj: + # requant q_a_proj + q_a_proj, q_a_proj_scale = requantize_weight_with_new_scale( + q_a_proj, + q_a_proj_scale, + q_a_proj_scale_2, + shared_weight_scale_2, + device=module.weight.device, + ) + + ########### fuse and load weights + if not is_lite: + fused_a = torch.cat([q_a_proj, kv_a_proj_with_mqa], dim=0) + else: + fused_a = kv_a_proj_with_mqa + + # For DeepseekV32: kv_a_proj_with_mqa is oversized + # to include indexer k weights, which is filled in post_load_weights. + module.weight.data[0 : fused_a.shape[0]].copy_(fused_a) + + ########### fuse weight_scale + if not is_lite: + fused_a_scale = torch.cat( + [q_a_proj_scale, kv_a_proj_with_mqa_scale], dim=0 + ) + else: + fused_a_scale = kv_a_proj_with_mqa_scale + # For DeepseekV32: kv_a_proj_with_mqa is oversized + # to include indexer k weights, which is filled in post_load_weights. + module.weight_scale.data[0 : fused_a_scale.shape[0]].copy_(fused_a_scale) + else: + fused_a = weights[f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight"][:] + if not is_lite: + q_a_proj = weights[f"{'.'.join(names[:-1])}.q_a_proj.weight"][:] + fused_a = torch.cat([q_a_proj, fused_a], dim=0) + + if f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale_inv" in weights: + fused_a_scale = weights[ + f"{'.'.join(names[:-1])}.kv_a_proj_with_mqa.weight_scale_inv" + ] + if not is_lite: + q_a_proj_scale = weights[ + f"{'.'.join(names[:-1])}.q_a_proj.weight_scale_inv" + ][:] + fused_a_scale = torch.cat([q_a_proj_scale, fused_a_scale], dim=0) + + module.weight_scale.data[0 : fused_a_scale.shape[0]].copy_( + fused_a_scale + ) + # For DeepseekV32: kv_a_proj_with_mqa is oversized + # to include indexer k weights, which is filled in post_load_weights. + module.weight.data[0 : fused_a.shape[0]].copy_(fused_a) + elif names[-1] in params_map: + module_weights = [] + for new_name in params_map[names[-1]]: + module_weights.append( + filter_weights(".".join(names[:-1] + [new_name]), weights) + ) + module.load_weights(weights=module_weights) + elif names[-1] == "experts": + module_weights = filter_weights(name, weights) + module_weights = rename_moe_weight( + module_weights, + { + "down_proj": "w2", + "up_proj": "w3", + "gate_proj": "w1", + }, + ) + module.load_weights(weights=[module_weights]) + elif names[-1] == "backend" and isinstance(module, MoE): + # Special case: ConfigurableMoE.backend (TRTLLMGenFusedMoE) + # Currently saved MoE weights don't include 'backend' in their names. + # After MoE refactoring, ConfigurableMoE now has a backend submodule, + # and weights loading is done in the backend, so module name includes '.backend'. + # We need to use parent module name (without .backend) to match saved weight names. + # After MoE refactoring is fully complete, all paths will follow this branch. + parent_name = ".".join(names[:-1]) + module_weights = filter_weights(parent_name, weights) + module_weights = rename_moe_weight( + module_weights, + { + "down_proj": "w2", + "up_proj": "w3", + "gate_proj": "w1", + }, + ) + module.load_weights(weights=[module_weights]) + elif names[-1] == "self_attn": + if f"{name}.o_a_proj" in weights: + load_o_a_proj(name, module) + attn_sink_key = f"{name}.attn_sink" + if attn_sink_key in weights: + sink_full = weights[attn_sink_key][:] + if not self.model_config.mapping.enable_attention_dp: + sink_full = split(sink_full, tp_size, tp_rank) + sink_full = sink_full.to(torch.float32).cuda() + module.mqa.attn_sink = nn.Parameter(sink_full, requires_grad=False) + continue + elif names[-1] == "mqa": + # DeepseekV4TrtllmAttention owns the optional attn_sink + # (per-head fp32, already TP-sharded). The checkpoint key + # uses the parent attention module name, not the .mqa + # suffix. When the key is absent we leave module.attn_sink + # as None so DeepseekV4TrtllmAttention.forward does not pass + # attention_sinks to the kernel. + parent_attn_name = ".".join(names[:-1]) + attn_sink_key = f"{parent_attn_name}.attn_sink" + if attn_sink_key in weights: + sink_full = weights[attn_sink_key][:] + if not self.model_config.mapping.enable_attention_dp: + sink_full = split(sink_full, tp_size, tp_rank) + sink_full = sink_full.to(torch.float32).cuda() + module.attn_sink = nn.Parameter(sink_full, requires_grad=False) + continue + elif names[-1] == "next_layer_layernorm": + continue + elif isinstance(module, (mHC, HCHead)) and load_flat_hc_weights(module, names): + continue + elif names[-1] in ("engram",): + # Engram is a container module with no direct parameters; + # its leaf sub-modules (multi_head_embedding, kv_proj, + # short_conv) and direct parameters (key_norm_weight, + # query_norm_weight) are loaded via the generic path. + continue + else: + module_weights = filter_weights(name, weights) + if hasattr(module, "load_weights"): + module.load_weights(weights=[module_weights]) + else: + for n, p in module.named_parameters(): + p.data.copy_(module_weights[n][:]) + + +@torch.compile(options={"max-autotune": True}) +def _get_last_token_states(hidden_states, attn_metadata): + last_tokens = ( + torch.cumsum( + attn_metadata.seq_lens_cuda, + dim=0, + dtype=torch.long, + ) + - 1 + ) + return hidden_states[last_tokens] + + +class DeepseekV4MTPHead(nn.Module): + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.hc_head = HCHead(config.hc_mult, config.hidden_size) + self.hc_mult = config.hc_mult + self.hidden_dim = config.hidden_size + + self.mapping_lm_head_tp = None + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: Linear, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + if not return_context_logits: + if attn_metadata is not None: + hidden_states = _get_last_token_states(hidden_states, attn_metadata) + else: + hidden_states = hidden_states[-1].unsqueeze(0) + + hidden_states = hidden_states.reshape(-1, self.hc_mult, self.hidden_dim) + hidden_states = self.hc_head(hidden_states) + hidden_states = self.norm(hidden_states) + + enable_attention_dp = self.model_config.mapping.enable_attention_dp + enable_lm_head_tp_in_adp = ( + enable_attention_dp and self.model_config.mapping.enable_lm_head_tp_in_adp + ) + + # Add pre-lm gather logic + if enable_lm_head_tp_in_adp: + # ADP + LM TP mode: perform All-Gather before LM_head + self.mapping_lm_head_tp = create_lm_head_tp_mapping( + self.model_config.mapping, hidden_states.shape[0] + ) + hidden_states = allgather(hidden_states, self.mapping_lm_head_tp, dim=0) + + # Temporarily disable gather_output when not in ADP mode or (in ADP mode and LM TP is enabled) + if not enable_attention_dp or enable_lm_head_tp_in_adp: + lm_head.gather_output = False + logits = lm_head( + hidden_states, mapping_lm_head_tp=self.mapping_lm_head_tp, is_spec_decoding_head=True + ) + if not enable_attention_dp or enable_lm_head_tp_in_adp: + lm_head.gather_output = True + return logits + + +class DeepseekV4LogitsProcessor(nn.Module): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + hc_head: HCHead, + norm: RMSNorm, + ): + super().__init__() + config = model_config.pretrained_config + self.model_config = model_config + self.hc_mult = config.hc_mult + self.hidden_dim = config.hidden_size + # Keep HCHead and final norm owned by DeepseekV4Model. This processor only + # borrows them, so checkpoint loading and PP weight removal still happen + # through the model's normal module tree. + object.__setattr__(self, "_hc_head", hc_head) + object.__setattr__(self, "_norm", norm) + + def forward( + self, + hidden_states: torch.Tensor, + lm_head: Linear, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + return_context_logits: bool = False, + ) -> torch.Tensor: + if not self.model_config.mapping.is_last_pp_rank(): + return lm_head(hidden_states).float() + + if not return_context_logits: + if attn_metadata is not None: + hidden_states = _get_last_token_states(hidden_states, attn_metadata) + else: + hidden_states = hidden_states[-1] + + hidden_states = hidden_states.reshape(-1, self.hc_mult, self.hidden_dim) + hidden_states = self._hc_head(hidden_states) + hidden_states = self._norm(hidden_states) + return lm_head(hidden_states).float() + + +class DeepseekV4Linear(Linear): + """ + A wrapper around Linear because we may optionally use min-latency kernels depending on input shapes. + """ + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool = True, + dtype: torch.dtype = None, + mapping: Optional[Mapping] = None, + tensor_parallel_mode: Optional[TensorParallelMode] = None, + gather_output: bool = False, # COLUMN parallel only + quant_config: Optional[QuantConfig] = None, + weights_loading_config: Optional[WeightsLoadingConfig] = None, + reduce_output: bool = True, # ROW parallel only + skip_create_weights_in_init: bool = False, + use_custom_cublas_mm: bool = False, + lora: Optional[LoraLayer] = None, + ): + super().__init__( + in_features, + out_features, + bias, + dtype, + mapping, + tensor_parallel_mode, + gather_output, + quant_config, + weights_loading_config, + reduce_output, + skip_create_weights_in_init, + use_custom_cublas_mm, + lora, + ) + + def apply_linear( + self, + input, + bias, + lora_params: Optional[dict] | None = None, + layer_idx: Optional[int] | None = None, + ): + num_tokens = input.shape[0] + if not self.has_any_quant and 1 <= num_tokens <= 16 and get_sm_version() not in [120, 121]: + output = torch.ops.trtllm.dsv3_fused_a_gemm_op(input, self.weight.t(), bias, None) + else: + output = super().apply_linear(input, bias, lora_params, layer_idx) + return output + + +class DeepseekV4Attention(MLA): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: Optional[int] = None, + aux_stream: Optional[torch.cuda.Stream] = None, + mapping_with_cp: Optional[Mapping] = None, + reduce_output: bool = True, + ): + config = model_config.pretrained_config + assert config.qk_rope_head_dim == 64, ( + "DeepseekV4Attention only supports qk_rope_head_dim=64" + ) + assert config.kv_lora_rank == 448, "DeepseekV4Attention only supports kv_lora_rank=448" + predicted_tokens_per_seq = ( + model_config.spec_config.tokens_per_gen_step + if model_config.spec_config is not None + else 1 + ) + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + predicted_tokens_per_seq=predicted_tokens_per_seq, + max_position_embeddings=config.max_position_embeddings, + bias=False, + pos_embd_params=_deepseek_v4_pos_embd_params(config, model_config, layer_idx), + layer_idx=layer_idx, + dtype=config.torch_dtype, + config=model_config, + aux_stream=aux_stream, + num_groups=config.o_groups, + o_lora_rank=config.o_lora_rank, + mapping_with_cp=mapping_with_cp, + reduce_output=reduce_output, + ) + + self.indexer = getattr(self.mqa, "indexer", None) + self.compressor = getattr(self.mqa, "compressor", None) + + +class DeepseekV4Gate(nn.Module): + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + is_hashed: bool, + vocab_size: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + fuse_routing_kernel: bool = True, + apply_routing: bool = False, + moe_backend: str = "CUTLASS", + ): + super().__init__() + self.weight = nn.Parameter( + torch.empty((num_experts, hidden_size), dtype=dtype), requires_grad=False + ) + self.moe_backend = moe_backend + bias_dtype = torch.float32 + + self.is_hashed = is_hashed + self.top_k = top_k + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + + if self.is_hashed: + # self.tid2eid = nn.Parameter( + # torch.empty(vocab_size, top_k, dtype=torch.int32), requires_grad=False + # ) + # WAR to avoid illegal expert indexes in hashed gating + self.tid2eid = nn.Parameter( + torch.stack([torch.randperm(num_experts)[:top_k] for _ in range(vocab_size)]) + .to(torch.int32) + .contiguous(), + requires_grad=False, + ) + self.e_score_correction_bias = None + else: + self.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=bias_dtype), requires_grad=False + ) + + assert not apply_routing, "DeepseekV4Gate routing is called inside MoE" + + def fetch_e_score_correction_bias(): + if not self.is_hashed: + return self.e_score_correction_bias.to(torch.float32) + else: + return None + + self._routing_method = DeepSeekV4MoeRoutingMethod( + top_k=self.top_k, + n_group=self.n_group, + topk_group=self.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + # Pass a callable to fetch the tensor from DeepseekV4Gate at runtime, ensuring it is on the correct device + callable_e_score_correction_bias=fetch_e_score_correction_bias, + callable_tid2eid=lambda: self.tid2eid, + is_hashed=self.is_hashed, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.linear(hidden_states.float(), self.weight.float()) + + def load_weights(self, weights: List[Dict]): + assert len(weights) == 1 + + self.weight.copy_(weights[0]["weight"][:]) + + if self.is_hashed: + self.tid2eid.copy_(weights[0]["tid2eid"][:].to(self.tid2eid.dtype)) + else: + self.e_score_correction_bias.copy_( + weights[0]["e_score_correction_bias"][:].to(self.e_score_correction_bias.dtype) + ) + + @property + def routing_method(self) -> DeepSeekV4MoeRoutingMethod: + return self._routing_method + + def apply( + self, logits: torch.Tensor, input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + # topk routing + return self.routing_method.apply(logits, input_ids) + + def get_experts_per_token(self): + return self.routing_method.top_k + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + *, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + shared_expert_intermediate_size: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + dtype: Optional[torch.dtype] = None, + model_config: ModelConfig = ModelConfig(), + override_quant_config: Optional[QuantConfig] = None, + layer_idx: Optional[int] = None, + ): + from ..distributed import AllReduce + + super().__init__() + config = model_config.pretrained_config + self.top_k = top_k + self.use_dp = model_config.mapping.enable_attention_dp + self.gate = DeepseekV4Gate( + hidden_size, + num_experts, + top_k=top_k, + n_group=config.n_group, + topk_group=config.topk_group, + routed_scaling_factor=config.routed_scaling_factor, + is_hashed=layer_idx < config.n_hash_layers, + vocab_size=config.vocab_size, + dtype=dtype, + fuse_routing_kernel=True, + apply_routing=False, + moe_backend=model_config.moe_backend, + ) + experts_quant_config = self._get_experts_quant_config(model_config, layer_idx) + if override_quant_config is not None and experts_quant_config is model_config.quant_config: + experts_quant_config = override_quant_config + + swiglu_limit = getattr(config, "swiglu_limit", None) + moe_swiglu_limit = None + if swiglu_limit is not None: + supports_swiglu_limit = True + if (model_config.moe_backend or "").upper() == "TRTLLM": + supports_swiglu_limit = False + if experts_quant_config is not None: + mode = experts_quant_config.layer_quant_mode + supports_swiglu_limit = ( + mode.has_nvfp4() + or mode.has_w4a16_mxfp4() + or mode.has_w4a8_mxfp4_fp8() + or mode.has_w4a8_mxfp4_mxfp8() + ) + if supports_swiglu_limit: + moe_load_balancer_config = getattr(model_config, "moe_load_balancer", None) + num_slots = ( + moe_load_balancer_config.num_slots + if moe_load_balancer_config and moe_load_balancer_config.num_slots + else num_experts + ) + local_num_slots = num_slots // model_config.mapping.moe_ep_size + device = "cuda" if torch.cuda.is_available() else "cpu" + moe_swiglu_limit = torch.full( + (local_num_slots,), float(swiglu_limit), dtype=torch.float32, device=device + ) + + self.experts = create_moe( + num_experts=num_experts, + routing_method=self.gate.routing_method, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results=False, # In both low‑latency and attention‑DP modes, FusedMoE skips the in‑op all‑reduce. + model_config=model_config, + override_quant_config=experts_quant_config, + aux_stream_dict=aux_stream_dict, + layer_idx=layer_idx, + # DS-R1 W4A8 is only supported through custom quantization script from + # examples/quantization/quantize_mixed_precision_moe.py + weight_loading_mode=( + MoEWeightLoadingMode.W4A8_CUSTOM + if experts_quant_config.layer_quant_mode.is_int4_weight_only_per_group() + else MoEWeightLoadingMode.VANILLA + ), + swiglu_limit=moe_swiglu_limit, + ) + + self.mapping = model_config.mapping + + # FIXME: incompatible with mixed quantization mode (including excluding modules from quantization) + block_size = 1 + if model_config.quant_config and model_config.quant_config.group_size is not None: + block_size = model_config.quant_config.group_size + + shared_tp_size, self.shared_output_scale = self._compute_shared_expert_tp_size( + shared_expert_intermediate_size, block_size + ) + + self.shared_experts = GatedMLP( + hidden_size=hidden_size, + intermediate_size=shared_expert_intermediate_size, + bias=False, + dtype=dtype, + config=model_config, + overridden_tp_size=shared_tp_size, + reduce_output=False, + swiglu_limit=swiglu_limit, + ) + + self.allreduce = None + if not self.use_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, strategy=model_config.allreduce_strategy + ) + self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.event_dict = {key: torch.cuda.Event() for key in [EventType.Main, EventType.MoeShared]} + + # Store config values for perfect routing. + self.model_config = model_config + self.dtype = dtype + + # Perfect router caching - precompute common logits if enabled. + if os.environ.get("ENABLE_PERFECT_ROUTER", "0") == "1": + precompute_common_perfect_router_logits( + num_experts=num_experts, + experts_per_token=top_k, + moe_ep_size=model_config.mapping.moe_ep_size, + dtype=dtype, + ) + + def _compute_shared_expert_tp_size( + self, intermediate_size: int, block_size: int + ) -> tuple[int, float | None]: + """ + In the case of Deepseek-R1, the TP size of MLP is capped by intermediate_size // block_size. + For example, when the intermediate_size is 2048 and block scaling size is 128, + TP sizes are limited to {1, 2, 4, 8, 16} because of 2048/128 = 16. + + Args: + intermediate_size (int): MLP intermediate size. + block_size (int): The quantization block scale size. In the case of Deepseek FP8 recipe, + it's 128. For NVFP4, it's 16. + + Returns: + tuple[int, float | None]: A tuple containing (shared_tp_size, shared_output_scale). + - shared_tp_size: The computed TP size. + - shared_output_scale: The output scale factor, or None if not needed. + """ + + assert intermediate_size % block_size == 0, ( + "intermediate_size must be divisible by block_size." + ) + + shared_output_scale = None + # The block scale size is 128, which requires shared_expert_intermediate_size to be divisible by 128. + if self.use_dp: + # If using attention DP, the shared experts also use DP instead of TP. + shared_tp_size = 1 + else: + # Due to the restriction of block scale size (i.e., 128), the supported + # TP sizes only include 1, 2, 4, 8, and 16. The math.gcd operation ensures that + # shared_tp_size falls in the supported TP sizes. + shared_tp_size = math.gcd( + intermediate_size // block_size, + self.mapping.tp_size, + ) + # If shared_tp_size has been overridden, the output of shared experts needs + # to be scaled down accordingly before all-reduce. + if shared_tp_size != self.mapping.tp_size: + shared_output_scale = shared_tp_size / self.mapping.tp_size + + return shared_tp_size, shared_output_scale + + @staticmethod + def _get_experts_quant_config(model_config, layer_idx: int) -> QuantConfig: + if getattr(model_config, "quant_config_dict", None) is None: + return model_config.quant_config + return model_config.quant_config_dict.get( + f"model.layers.{layer_idx}.mlp.experts", model_config.quant_config + ) + + def _create_ideal_expert_load_balanced_logits( + self, num_tokens: int, num_experts: int, device: torch.device + ) -> torch.Tensor: + """ + Create ideal logits that produce GPU-aware load balanced expert assignment. + This method uses the global cache to access precomputed logits to optimize performance. + """ + # Use global cached logits. + return get_cached_perfect_router_logits( + num_tokens=num_tokens, + num_experts=num_experts, + experts_per_token=self.top_k, + moe_ep_size=self.model_config.mapping.moe_ep_size, + device=device, + dtype=self.dtype, + ) + + def compute_routed_output( + self, hidden_states, hidden_states_fp4, input_ids, all_rank_num_tokens, do_finalize + ): + # max-throughput + use_dp_padding = False + # Add DP padding on SM120 for context comm performance + # TODO: Move this model-agonostic part to MoE + if self.use_dp and self.mapping.tp_size > 1 and get_sm_version() == 120: + use_dp_padding = True + hidden_states = torch.nn.functional.pad( + hidden_states, (0, 0, 0, max(all_rank_num_tokens) - hidden_states.shape[0]) + ) + input_ids = torch.nn.functional.pad( + input_ids, + (0, max(all_rank_num_tokens) - input_ids.shape[0]), + value=self.model_config.pretrained_config.pad_token_id, + ) + + router_logits = self.gate(hidden_states) + + # Use ideal load balanced logits if enabled, otherwise use gate output. + if os.environ.get("ENABLE_PERFECT_ROUTER", "0") == "1": + # WARNING: This discards the learned gate output and uses ideal logits for perfect load balancing. + # Only use this for testing load balancing strategies, not for actual inference. + # The gate is still computed to maintain realistic performance measurement. + num_tokens, num_experts = router_logits.shape + router_logits = self._create_ideal_expert_load_balanced_logits( + num_tokens=num_tokens, num_experts=num_experts, device=hidden_states.device + ) + + routed_output = self.experts( + hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states, + router_logits, + input_ids=input_ids, + do_finalize=do_finalize, + output_dtype=hidden_states.dtype, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=use_dp_padding, + **({"alltoall_result_do_sum": False} if isinstance(self.experts, WideEPMoE) else {}), + ) + + return routed_output + + def forward( + self, + hidden_states: torch.Tensor, + hidden_states_fp4: Optional[Fp4QuantizedTensor] = None, + input_ids: Optional[torch.Tensor] = None, + all_rank_num_tokens: Optional[list[int]] = None, + final_all_reduce_params: Optional[AllReduceParams] = None, + do_finalize: Optional[bool] = True, + ) -> torch.Tensor: + if not do_finalize: + assert not self.use_dp + + def _compute_shared_output(): + shared_output = self.shared_experts( + hidden_states_fp4 if hidden_states_fp4 is not None else hidden_states + ) + if self.shared_output_scale is not None: + shared_output *= self.shared_output_scale + return shared_output + + def _compute_routed_output(): + routed_output = self.compute_routed_output( + hidden_states, hidden_states_fp4, input_ids, all_rank_num_tokens, do_finalize + ) + return routed_output + + # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames + + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, + _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + ) + + if not do_finalize: + return [shared_output, *routed_output] + else: + if routed_output.dim() == 3: + assert shared_output.numel() * self.top_k == routed_output.numel(), ( + "unmatched tensor shape" + ) + final_hidden_states = moe_reduce_add_shared_output(routed_output, shared_output) + else: + assert shared_output.size() == routed_output.size(), "unmatched tensor shape" + final_hidden_states = shared_output + routed_output + + if not self.use_dp and self.mapping.tp_size > 1: + final_hidden_states = self.allreduce( + final_hidden_states, all_reduce_params=final_all_reduce_params + ) + + return final_hidden_states + + +class DeepseekV4DecoderLayer(DecoderLayer): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + is_separate_draft_engine: bool = False, + mapping_with_cp: Optional[Mapping] = None, + disable_post_moe_fusion: bool = False, + ): + super().__init__() + self.model_config = model_config + self.config = model_config.pretrained_config + config = self.config + + self.hidden_size = config.hidden_size + self.moe_intermediate_size = config.moe_intermediate_size + self.num_experts = config.n_routed_experts + self.num_shared_experts = config.n_shared_experts + self.top_k = config.num_experts_per_tok + + self.mapping = model_config.mapping + mapping = self.mapping + self.enable_attention_dp = mapping.enable_attention_dp + self.mlp_tp_size = mapping.tp_size + self.is_p2p_supported = can_access_peer(mapping) + + self.hc_attn = mHC( + config.hc_mult, + config.hidden_size, + config.hc_sinkhorn_iters, + dtype=torch.float32, + post_mult_value=2.0, + ) + + layer_idx_for_attention = layer_idx + if is_separate_draft_engine: + # KVCacheManager only support 1 layer for separate draft engine + layer_idx_for_attention = layer_idx - model_config.pretrained_config.num_hidden_layers + + self.self_attn = DeepseekV4Attention( + model_config, + layer_idx=layer_idx_for_attention, + aux_stream=aux_stream_dict[AuxStreamType.Attention], + reduce_output=not self.enable_attention_dp and self.mapping.tp_size > 1, + ) + + self.fusion_config = EagerFusionConfig() + self.enable_fusion = os.environ.get("TRTLLM_DEEPSEEK_EAGER_FUSION_DISABLED", "0") == "0" + self.enable_fusion &= not self.enable_attention_dp + + self.hc_ffn = mHC( + config.hc_mult, + config.hidden_size, + config.hc_sinkhorn_iters, + dtype=torch.float32, + post_mult_value=2.0, + ) + + # FIXME: incompatible with mixed quantization mode + quant_config = self._get_decoder_layer_quant_config(model_config, layer_idx) + self.is_nvfp4 = quant_config.layer_quant_mode.has_nvfp4() + assert quant_config.quant_algo is not QuantAlgo.MIXED_PRECISION, ( + "MIXED_PRECISION is ambiguous" + ) + + self.allreduce = None + self.moe_allreduce = None + if not self.enable_attention_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + dtype=config.torch_dtype, + ) + self.moe_allreduce = MoEAllReduce(self.mapping) + + has_tp = mapping.has_tp() + self.fusion_config.PRE_MOE_FUSION = self.enable_fusion and has_tp + # DeepSeek-V4 applies the next RMSNorm after mHC post_mapping and the + # next layer's mHC pre_mapping. Fusing the post-MoE all-reduce with the + # next RMSNorm would normalize the raw MoE output before mHC post_mapping, + # which is not equivalent. + self.fusion_config.POST_MOE_FUSION = False + + self.mlp = DeepseekV4MoE( + num_experts=self.num_experts, + top_k=self.top_k, + hidden_size=self.hidden_size, + intermediate_size=self.moe_intermediate_size, + shared_expert_intermediate_size=self.moe_intermediate_size * self.num_shared_experts, + dtype=config.torch_dtype, + model_config=model_config, + override_quant_config=quant_config, + aux_stream_dict=aux_stream_dict, + layer_idx=layer_idx, + ) + + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + # When enable_attention_dp is True, we normally skip attention all-reduce since each + # DP rank works on different batch elements. However, with CP > 1, attention is split + # across CP ranks for the SAME batch element, so all-reduce is still needed. + has_cp = mapping_with_cp is not None and mapping_with_cp.cp_size > 1 + can_skip_for_attention_dp = self.enable_attention_dp and not has_cp + self.disable_attn_allreduce = ( + self.fusion_config.PRE_MOE_FUSION + or self.fusion_config.PRE_MLP_FUSION + or self.mapping.tp_size == 1 + or can_skip_for_attention_dp + ) + + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.layer_idx = layer_idx + # is_first_layer is baked in at __init__ time so the Python-side branch in + # forward() resolves at CUDA-graph capture time. + self.is_first_layer = layer_idx == 0 + # fused_hc knob: pretrained-config attr `enable_fused_hc` controls whether + # the MHC boundary fusion (`mHC.fused_hc`) is used. When False, fall back + # to the unfused `post_mapping → pre_mapping` chain (same path engram + # layers already take). Env var TRTLLM_MHC_ENABLE_FUSED_HC overrides the + # config attr (set to "1" to force-enable while validating fused_hc). + # Default is disabled because fused_hc must be bit-equivalent to the + # explicit post_mapping -> pre_mapping chain before it can safely replace it. + _env = os.environ.get("TRTLLM_MHC_ENABLE_FUSED_HC") + if _env is not None: + self.enable_fused_hc = _env not in ("0", "false", "False") + else: + self.enable_fused_hc = bool(getattr(config, "enable_fused_hc", False)) + self.next_layer_layernorm: RMSNorm = None + # Finalized in DeepseekV4ForCausalLM.post_load_weights once the full layer + # list is visible: a layer may defer its hc_ffn.post_mapping only if + # the next layer is able to absorb it via fused_hc (i.e. the next + # layer has fused_hc enabled and no engram at its entry). Last layer + # never defers — hc_head consumes the residual directly. + self.defer_post_mapping: bool = False + + # Engram module (optional, for n-gram context augmentation) + self.engram: Optional[Engram] = None + _engram_config = getattr(config, "engram_config", None) + _engram_vocab_sizes_by_layer = getattr(config, "engram_vocab_sizes_by_layer", {}) + if _engram_config is not None and layer_idx in _engram_vocab_sizes_by_layer: + self.engram = Engram( + layer_id=layer_idx, + config=_engram_config, + vocab_sizes_flat=_engram_vocab_sizes_by_layer[layer_idx], + stream=aux_stream_dict[AuxStreamType.EngramPrecompute], + ) + + def _get_decoder_layer_quant_config( + self, model_config: ModelConfig[PretrainedConfig], layer_idx: int + ): + """ + The MTP layer in the nvfp4 checkpoint is unquantized. Because the TRTLLM + moe_backend only supports fp8/fp4 quantization, we need to override + the quant_config for the MTP layer. + """ + quant_config = model_config.quant_config + + layer_name = f"model.layers.{layer_idx}" + if quant_config.is_module_excluded_from_quantization(layer_name): + return QuantConfig( + quant_algo=None, + kv_cache_quant_algo=quant_config.kv_cache_quant_algo, + ) + else: + return model_config.quant_config + + def _compute_mlp_tp_size(self, intermediate_size: int, block_size: int) -> int: + """ + For DeepSeek‑R1, MLP TP size is limited by intermediate_size // block_size + and must also be multiples of gpus_per_node to avoid expensive inter‑node allreduce. + + Args: + intermediate_size (int): MLP intermediate size. + block_size (int): The quantization block scale size. In the case of Deepseek FP8 recipe, + it's 128. For NVFP4, it's 16. + + Returns: + int: The computed tp_size. + """ + + assert intermediate_size % block_size == 0, ( + "intermediate_size must be divisible by block_size." + ) + if self.enable_attention_dp: + # If using attention DP, the MLP also uses DP instead of TP. + mlp_tp_size = 1 + else: + # The two math.gcd operations ensure that mlp_tp_size falls in the candidate TP sizes. + tp = math.gcd( + intermediate_size // block_size, + self.mapping.tp_size, + ) + + if tp > self.mapping.gpus_per_node: + mlp_tp_size = math.gcd( + tp, + self.mapping.gpus_per_node, + ) # Avoid costly inter-node TP + else: + mlp_tp_size = tp + return mlp_tp_size + + def forward( + self, + position_ids: torch.IntTensor, + hc_state, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + spec_metadata: Optional[SpecMetadata] = None, + input_ids: Optional[torch.IntTensor] = None, + engram_embeddings=None, + **kwargs, + ): + """mHC-aware decoder layer with boundary fusion. + + ``hc_state`` carries the mHC pipeline state across layers: + + - ``is_first_layer=True``: ``hc_state`` is the initial residual tensor + ``[B, HC_MULT, hidden]`` (bf16). This layer bootstraps the stream + with ``hc_attn.pre_mapping`` (aka the standalone ``hc_pre``). + - ``is_first_layer=False``: ``hc_state`` is an ``HCState``. If + ``is_deferred`` (fused path), the 4 tensors feed the current + ``hc_attn.fused_hc``. Otherwise ``residual`` is already post-mapped + by the prior layer and this layer just runs ``pre_mapping``. + + Returns an ``HCState``. Fused mode returns a deferred state carrying + this layer's ``hc_ffn`` inputs, so the next layer absorbs + ``hc_ffn.post_mapping`` via its own ``fused_hc``. Engram / unfused + mode resolves the post_mapping in-layer and returns a resolved state. + + Engram (when enabled on this layer) injects a residual modification + between the previous block's post_mapping and this block's pre_mapping. + That mutation breaks the algebraic assumptions of ``fused_hc``, so + engram-enabled layers fall back to the unfused + ``post_mapping → +engram → pre_mapping`` chain at the entry boundary. + The mid-layer ``attn → MoE`` boundary is always safe to fuse. + """ + has_engram = self.engram is not None and engram_embeddings is not None + + # ------------------------------------------------------------------- + # Entry boundary: hc_pre (layer 0) or fused_hc / unfused chain. + # ------------------------------------------------------------------- + residual, post_mix, comb_mix, layer_input = self._entry_boundary( + hc_state, engram_embeddings, has_engram + ) + + # ------------------------------------------------------------------- + # Attention block + # ------------------------------------------------------------------- + x_attn = self.input_layernorm(layer_input) + x_attn = self.self_attn( + position_ids=position_ids, + hidden_states=x_attn, + attn_metadata=attn_metadata, + all_reduce_params=AllReduceParams(enable_allreduce=not (self.disable_attn_allreduce)), + **kwargs, + ) + + # ------------------------------------------------------------------- + # Mid-layer boundary: fuse hc_attn.post_mapping + hc_ffn.pre_mapping. + # No engram concern here because engram only fires at layer entry. + # When enable_fused_hc=False, fall back to the unfused chain. + # ------------------------------------------------------------------- + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + self.fusion_config.POST_MOE_FUSION = False + if self.enable_fused_hc: + residual, post_mix, comb_mix, layer_input = self.hc_ffn.fused_hc( + x_prev=x_attn, + residual_prev=residual, + post_mix_prev=post_mix, + comb_mix_prev=comb_mix, + ) + else: + # Break fused_hc into post_mapping and pre_mapping as separate ops. + residual = self.hc_attn.post_mapping( + x=x_attn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + post_mix, comb_mix, layer_input = self.hc_ffn.pre_mapping(residual) + + # ------------------------------------------------------------------- + # MoE block — returns x_ffn (post-MoE, already normed by + # next_layer_layernorm when POST_MOE_FUSION is on). + # ------------------------------------------------------------------- + x_ffn = self.forward_MoE( + hidden_states=layer_input, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + input_ids=input_ids, + ) + + # Defer this layer's hc_ffn.post_mapping only when the NEXT layer can + # absorb it via fused_hc (see post_load_weights). Otherwise resolve it + # here and hand the next layer a fully post-mapped residual. + if self.defer_post_mapping: + return HCState.deferred( + residual=residual, post_mix=post_mix, comb_mix=comb_mix, x_prev=x_ffn + ) + resolved_residual = self.hc_ffn.post_mapping( + x=x_ffn, + residual=residual, + post_layer_mix=post_mix, + comb_res_mix=comb_mix, + ) + return HCState.resolved(resolved_residual) + + def _entry_boundary(self, hc_state, engram_embeddings, has_engram): + """Resolve the per-layer entry into (residual, post_mix, comb_mix, layer_input). + + Two code paths: + 1. fused: previous layer deferred its post_mapping; fold it into this + layer's pre_mapping via hc_attn.fused_hc. The prev layer only + defers when post_load_weights has proved we can absorb it here + (no engram, fused_hc enabled), so a deferred state at entry is + guaranteed fusable. + 2. unfused: previous layer already resolved its post_mapping (or + this is layer 0); run pre_mapping on the residual directly. + """ + # Fused entry: prev layer deferred its post_mapping; fold it into + # hc_attn.fused_hc. By construction (post_load_weights), a deferred + # state at entry is guaranteed fusable. Layer 0 receives the raw + # residual tensor and has no HCState, so short-circuit on it. + if not self.is_first_layer and hc_state.is_deferred: + return self.hc_attn.fused_hc( + x_prev=hc_state.x_prev, + residual_prev=hc_state.residual, + post_mix_prev=hc_state.post_mix, + comb_mix_prev=hc_state.comb_mix, + ) + + # Unfused entry: layer 0 hands us the initial residual tensor + # [B, HC_MULT, hidden] directly; post-layer-0 hands us a resolved + # HCState. Both collapse to "apply engram delta (if any) then run + # pre_mapping". + residual = hc_state if self.is_first_layer else hc_state.residual + if has_engram: + residual = residual + self.engram(residual, engram_embeddings) + post_mix, comb_mix, layer_input = self.hc_attn.pre_mapping(residual) + return residual, post_mix, comb_mix, layer_input + + def forward_MoE( + self, + hidden_states: torch.Tensor, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + spec_metadata: Optional[SpecMetadata] = None, + input_ids: Optional[torch.IntTensor] = None, + ) -> torch.Tensor: + def _run_MoE(hidden_states, hidden_states_fp4, do_finalize, input_ids): + return self.mlp( + hidden_states, + hidden_states_fp4, + all_rank_num_tokens=attn_metadata.all_rank_num_tokens, + final_all_reduce_params=AllReduceParams( + enable_allreduce=not ( + self.fusion_config.POST_MOE_FUSION or self.mapping.tp_size == 1 + ) + ), + do_finalize=do_finalize, + input_ids=input_ids, + ) + + if self.fusion_config.PRE_MOE_FUSION: + # In DeepSeek-V4 the external residual connection is handled by mHC + # (hc_ffn.post_mapping), so there is no residual to add here. + # Use fused allreduce + RMSNorm (no residual addition). + hidden_states = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + else: + # No fusion: just normalize. + hidden_states = self.post_attention_layernorm(hidden_states) + + # Note: this fusion pattern is only supported for single-node TRTLLM-nvfp4 backend now + do_finalize = self.mapping.is_multi_node() or ( + not ( + self.fusion_config.POST_MOE_FUSION + and hidden_states.shape[0] <= self.moe_allreduce.max_token + and self.model_config.moe_backend == "TRTLLM" + and self.mlp.experts.has_nvfp4 + and self.is_p2p_supported + ) + ) + + hidden_states = _run_MoE( + hidden_states, hidden_states_fp4=None, input_ids=input_ids, do_finalize=do_finalize + ) + + if self.fusion_config.POST_MOE_FUSION: + if do_finalize: + # Fused allreduce + RMSNorm, no residual needed (mHC handles it). + hidden_states = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RMS_NORM, + norm_weight=self.next_layer_layernorm.weight, + eps=self.next_layer_layernorm.variance_epsilon, + trigger_completion_at_end=False, + ), + ) + else: + assert len(hidden_states) == 4, "hidden_states must have 4 elements" + + shared_output = hidden_states[0] + fc2_output = hidden_states[1] + expert_scale_factor = hidden_states[2] + expanded_idx_to_permuted_idx = hidden_states[3] + + moe_all_reduce_params = MoEAllReduceParams( + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + expert_scale_factor=expert_scale_factor, + shared_expert_output=shared_output, + residual=None, + norm_weight=self.next_layer_layernorm.weight, + eps=self.next_layer_layernorm.variance_epsilon, + is_cutlass_min_latency=False, + ) + (hidden_states,) = self.moe_allreduce( + fc2_output, all_reduce_params=moe_all_reduce_params + ) + else: + if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx): + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, None) + + return hidden_states + + +class DeepseekV4MTP(DeepseekV4DecoderLayer): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + is_separate_draft_engine: bool = False, + ): + super().__init__( + model_config, + layer_idx, + aux_stream_dict, + is_separate_draft_engine, + disable_post_moe_fusion=True, + ) + config = model_config.pretrained_config + self.hidden_dim = config.hidden_size + self.moe_intermediate_size = config.moe_intermediate_size + self.num_experts = config.n_routed_experts + self.num_shared_experts = config.n_shared_experts + self.top_k = config.num_experts_per_tok + + self.aux_stream = aux_stream_dict[AuxStreamType.MoeShared] + self.event_dict = {key: torch.cuda.Event() for key in [EventType.Main, EventType.MoeShared]} + + self.enorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + self.hnorm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + self.hc_mult = config.hc_mult + if model_config.mapping.enable_attention_dp: + self.e_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.h_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + else: + self.e_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + tensor_parallel_mode=TensorParallelMode.ROW, + mapping=model_config.mapping, + reduce_output=True, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + self.h_proj = Linear( + config.hidden_size, + config.hidden_size, + bias=False, + dtype=config.torch_dtype, + tensor_parallel_mode=TensorParallelMode.ROW, + mapping=model_config.mapping, + reduce_output=True, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + ) + + self.shared_head = DeepseekV4MTPHead(model_config) + + def forward( + self, + input_ids: torch.IntTensor, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + embed_tokens: Embedding, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + all_rank_num_tokens: Optional[List[int]] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + """Run an MTP layer. + + ``embed_tokens`` is injected by the one-model draft path and shared + with the target model. ``hidden_states`` is the flattened mHC residual + from the target or previous MTP layer: [num_tokens, hc_mult * hidden]. + """ + + def norm_embeds(): + return self.enorm(embed_tokens(input_ids)) # emdedding + + def norm_hidden(): + return self.hnorm(hidden_states.reshape(-1, self.hc_mult, self.hidden_dim)) + + inputs_embeds, hidden_states = maybe_execute_in_parallel( + norm_embeds, + norm_hidden, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], + self.aux_stream, + ) + + # Split hidden_states columnwise based on TP + tp_size = self.model_config.mapping.tp_size + tp_rank = self.model_config.mapping.tp_rank + if tp_size > 1 and not (self.model_config.mapping.enable_attention_dp): + inputs_embeds = torch.chunk(inputs_embeds, tp_size, dim=-1)[tp_rank].contiguous() + hidden_states = torch.chunk(hidden_states, tp_size, dim=-1)[tp_rank].contiguous() + + inputs_embeds = self.e_proj(inputs_embeds).unsqueeze(1) + hidden_states = self.h_proj(hidden_states) + hidden_states = inputs_embeds + hidden_states + + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + hc_state = super().forward( + position_ids=position_ids, + hc_state=HCState.resolved(hidden_states), + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + input_ids=input_ids, + **kwargs, + ) + finally: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + hidden_states = hc_state.residual.flatten(1) + + return hidden_states + + +class DeepseekV4Model(DecoderModel): + def __init__( + self, model_config: ModelConfig[PretrainedConfig], mapping_with_cp: Optional[Mapping] = None + ): + super().__init__(model_config) + config = model_config.pretrained_config + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hc_mult = config.hc_mult + aux_stream_list = [torch.cuda.Stream() for _ in range(5)] + self.aux_stream_dict = { + AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: aux_stream_list[1], + AuxStreamType.MoeBalancer: aux_stream_list[2], + AuxStreamType.MoeOutputMemset: aux_stream_list[3], + AuxStreamType.EngramPrecompute: aux_stream_list[4], + } + + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + ) + + self.hc_head = HCHead(config.hc_mult, config.hidden_size) + + # Engram hash provider (optional, for n-gram context augmentation) + # Must be created before layers so vocab sizes can be stored on config. + self.engram_hash_provider: Optional[EngramHashProvider] = None + self.use_engram = getattr(config, "has_engram", False) + if self.use_engram: + engram_config = EngramConfig( + tokenizer_name_or_path=getattr( + config, "tokenizer_name_or_path", "deepseek-ai/DeepSeek-V3" + ), + engram_vocab_size=config.engram_vocab_size, + max_ngram_size=config.engram_max_ngram_size, + n_embed_per_ngram=config.engram_n_embed_per_ngram, + n_head_per_ngram=config.engram_n_head_per_ngram, + layer_ids=config.engram_layer_ids, + pad_id=config.engram_pad_id, + seed=config.engram_seed, + kernel_size=config.engram_kernel_size, + hidden_size=config.hidden_size, + hc_mult=config.hc_mult, + norm_eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.engram_hash_provider = EngramHashProvider(engram_config) + self.engram_layer_ids = engram_config.layer_ids + # Store engram config and per-layer vocab sizes on the pretrained config so + # DeepseekV4DecoderLayer can read them directly from model_config without extra params. + config.engram_config = engram_config + config.engram_vocab_sizes_by_layer = { + layer_id: [ + x + for y in self.engram_hash_provider.vocab_size_across_layers[layer_id] + for x in y + ] + for layer_id in engram_config.layer_ids + } + + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + model_config, + layer_idx, + self.aux_stream_dict, + mapping_with_cp=mapping_with_cp, + ) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm( + hidden_size=config.hidden_size, eps=config.rms_norm_eps, dtype=config.torch_dtype + ) + + def __pp_init__(self): + self.epilogue.append(self.hc_head) + super().__pp_init__() + + def forward( + self, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # ----------------------------------------------------------------- + # Engram pre-computation (overlapped with main-stream layer forward) + # + # Hash provider internally applies CompressedTokenizer to input_ids, + # so no separate engram tokenizer input is needed. + # + # All precompute() calls are dispatched onto a dedicated engram CUDA + # stream so they run concurrently with the main stream processing the + # earlier transformer layers. A per-layer Event is recorded after + # each precompute; the main stream waits on the event just before it + # needs the result for that specific layer. + # ----------------------------------------------------------------- + engram_embeddings_cache: Optional[Dict] = None + engram_events: Dict[int, torch.cuda.Event] = {} + if self.use_engram and self.engram_hash_provider is not None and input_ids is not None: + hash_cache = self.engram_hash_provider.compute_hashes( + input_ids.view(-1), + seq_lens=attn_metadata.seq_lens_cuda, + ) + engram_embeddings_cache = {} + for layer_id in self.engram_layer_ids: + engram_mod = self.layers[layer_id].engram + if engram_mod is not None: + # precompute() dispatches onto the engram stream internally and + # records sync_event; the main stream will wait on it before use. + engram_embeddings_cache[layer_id] = engram_mod.precompute( + hash_cache[layer_id], dtype=inputs_embeds.dtype + ) + engram_events[layer_id] = engram_mod.sync_event + + mapping = self.model_config.mapping + if mapping.has_pp() and not mapping.is_first_pp_rank(): + hidden_states = torch.empty( + inputs_embeds.shape[0], + self.hc_mult, + inputs_embeds.shape[1], + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + else: + hidden_states = inputs_embeds + hidden_states = hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1) + + # ``hc_state`` carries the mHC pipeline state across layers. Layer 0 + # receives the initial residual tensor and bootstraps via + # hc_attn.pre_mapping ("hc_pre"). Every later layer receives an + # ``HCState``. In fused mode the state is "deferred" (prior layer's + # hc_ffn.post_mapping is folded into this layer's hc_attn.fused_hc); + # in unfused mode the state is "resolved" (residual already + # post-mapped). After the last layer, a deferred state is closed with + # a standalone hc_post; a resolved state feeds hc_head directly. + hc_state = hidden_states + + for idx, decoder_layer in enumerate(self.layers[: self.num_hidden_layers]): + engram_embeddings = None + if engram_embeddings_cache is not None and idx in engram_embeddings_cache: + # Sync: ensure the engram stream has finished precompute for this layer + # before the main stream reads the result. + engram_events[idx].wait(torch.cuda.current_stream()) + engram_embeddings = engram_embeddings_cache[idx] + + hc_state = decoder_layer( + position_ids=position_ids, + hc_state=hc_state, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + input_ids=input_ids, + engram_embeddings=engram_embeddings, + ) + + hidden_states = hc_state.residual.flatten(1) + + return hidden_states + + +@register_auto_model("DeepseekV4ForCausalLM") +class DeepseekV4ForCausalLM(SpecDecOneEngineForCausalLM[DeepseekV4Model, PretrainedConfig]): + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + return {"kv_cache_config": {"tokens_per_block": 128}} + + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + self.mapping_with_cp = None + # Note: Currently the usage of mapping is all over the place making its usage brittle + # in this file. As a temporary WAR, we hold on to an original copy of mapping when CP + # is in action. This shall be passed on to attention which is the only layer that's + # affected by CP. For other layers, CP ranks are repurposed to TP. This shall be undone + # at the end of __init__. + if model_config.mapping.has_cp_helix(): + print( + "[DeepseekV4ForCausalLM::__init__] Repurposing KVP ranks to TP while keeping other details the same." + ) + self.mapping_with_cp = copy.deepcopy(model_config.mapping) + # Repurpose KVP ranks to TP while keeping other details the same. + model_config._frozen = False + model_config.mapping = model_config.mapping.repurpose_helix_cp_to_tp() + model_config._frozen = True + + # Rename some keys of quant_config_dict to support legacy checkpoints + if model_config.quant_config_dict is not None: + model_config = copy.deepcopy(model_config) + quant_config_dict = {} + for key, val in model_config.quant_config_dict.items(): + key_split = key.split(".") + if key_split[-1] == "fused_a": + key = ".".join(key_split[:-1] + ["kv_a_proj_with_mqa"]) + quant_config_dict[key] = val + model_config._frozen = False + model_config.quant_config_dict = quant_config_dict + model_config._frozen = True + + super().__init__( + model=DeepseekV4Model(model_config, mapping_with_cp=self.mapping_with_cp), + model_config=model_config, + ) + self.logits_processor = DeepseekV4LogitsProcessor( + model_config, self.model.hc_head, self.model.norm + ) + + # Exclude Engram weights from quantization. Engram embedding tables + # and small linear projections are not suited for NVFP4/FP8 quant. + if getattr(self.config, "has_engram", False): + if model_config.quant_config.exclude_modules is None: + model_config.quant_config.exclude_modules = [] + model_config.quant_config.exclude_modules.append("*engram*") + + self.model_nextn = 0 + if ( + model_config.spec_config is not None + and model_config.spec_config.spec_dec_mode.is_mtp_one_model() + ): + self.model_nextn = model_config.spec_config.num_nextn_predict_layers + ckpt_nextn = self.config.num_nextn_predict_layers + self.num_hidden_layers = self.config.num_hidden_layers + assert ckpt_nextn > 0, "There is not MTP modules in the checkpoint." + if ckpt_nextn == 1 and not model_config.spec_config.use_mtp_vanilla: + pass + else: + # modify the QuantConfig to support duplicated mtp layers + if model_config.quant_config.exclude_modules is not None: + extend_exclude_modules = [] + for model_mtp_idx in range( + self.num_hidden_layers, self.num_hidden_layers + self.model_nextn + ): + ckpt_mtp_idx = ( + model_mtp_idx - self.num_hidden_layers + ) % ckpt_nextn + self.num_hidden_layers + model_prefix = f"model.layers.{model_mtp_idx}" + ckpt_prefix = f"model.layers.{ckpt_mtp_idx}" + for exclude_module in model_config.quant_config.exclude_modules: + if ckpt_prefix in exclude_module and model_prefix not in exclude_module: + extend_exclude_modules.append( + exclude_module.replace(ckpt_prefix, model_prefix) + ) + self.model_config.quant_config.exclude_modules.extend(extend_exclude_modules) + self.model.layers.extend(self.draft_model.mtp_layers) + + # Undo any manipulations done to mapping. + if self.mapping_with_cp is not None: + print("[DeepseekV4ForCausalLM::__init__] Restoring original mapping.") + model_config._frozen = False + model_config.mapping = self.mapping_with_cp + model_config._frozen = True + + def forward( + self, + attn_metadata: DeepseekV4TrtllmAttentionMetadata, + input_ids: torch.IntTensor = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + return_context_logits: bool = False, + **kwargs, + ) -> torch.Tensor: + return super().forward( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + return_context_logits=return_context_logits, + **kwargs, + ) + + def load_weights(self, weights: Dict): + weight_loader = DeepseekV4WeightLoader(self) + weight_loader.load_weights(weights) + + def post_load_weights(self): + layers = self.model.layers[: self.config.num_hidden_layers] + last_idx = self.config.num_hidden_layers - 1 + for idx, layer in enumerate(layers): + if idx == last_idx: + # The V4 logits path is HCHead -> model.norm -> lm_head, so + # the final decoder layer must not fold model.norm into MoE. + layer.next_layer_layernorm = None + layer.fusion_config.POST_MOE_FUSION = False + layer.defer_post_mapping = False + else: + next_layer = layers[idx + 1] + layer.next_layer_layernorm = next_layer.input_layernorm + # Defer this layer's hc_ffn.post_mapping into the next layer's + # hc_attn.fused_hc only if that next layer can actually absorb + # it: fused_hc enabled on both sides, and no engram at the + # next layer's entry (engram needs the materialized residual + # before pre_mapping runs). + layer.defer_post_mapping = ( + layer.enable_fused_hc + and next_layer.enable_fused_hc + and next_layer.engram is None + ) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 62683b3f62f2..5324416f4eb6 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1437,6 +1437,9 @@ def __init__( case "step3p7" | "step3p5": from .modeling_step3p7 import Step3p7MTP mtp_layer = Step3p7MTP + case "deepseek_v4": + from .modeling_deepseekv4 import DeepseekV4MTP + mtp_layer = DeepseekV4MTP case _: raise ValueError( f"Model type {model_type} not supported for MTP") @@ -1496,6 +1499,12 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], elif model_type == "qwen3_next": from .modeling_qwen3_next import Qwen3NextMTP mtp_layer = Qwen3NextMTP(model_config, layer_idx, aux_stream_dict) + elif model_type == "deepseek_v4": + from .modeling_deepseekv4 import DeepseekV4MTP + mtp_layer = DeepseekV4MTP(model_config, + layer_idx, + aux_stream_dict, + is_separate_draft_engine=True) else: raise ValueError( f"MTPDraftModel does not support model_type: {model_type}") diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 03dbf36d0bc0..6f7a8083c4a6 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -300,11 +300,25 @@ def __pp_init__(self): skip_forward(module) num_hidden_layers = self.model_config.pretrained_config.num_hidden_layers - assert num_hidden_layers >= mapping.pp_size, f"{num_hidden_layers} layers are not enough for PP{mapping.pp_size}" + assert num_hidden_layers >= mapping.pp_size, ( + f"{num_hidden_layers} layers are not enough for PP{mapping.pp_size}" + ) + pp_layer_list = mapping.pp_layers(num_hidden_layers) + total_num_layers = num_hidden_layers + spec_config = getattr(self.model_config, "spec_config", None) + if spec_config is not None: + from ..speculative.utils import get_num_spec_layers + + num_spec_layers = get_num_spec_layers(spec_config) + total_num_layers += num_spec_layers + if num_spec_layers > 0 and mapping.is_last_pp_rank(): + pp_layer_list.extend( + range(total_num_layers - num_spec_layers, total_num_layers)) + if len(pp_layer_list) == 0: + pp_layer_list.append(0) has_pp_layer = len(pp_layer_list) > 0 - for layer_idx in range(num_hidden_layers): - layer = self.layers[layer_idx] + for layer_idx, layer in enumerate(self.layers[:num_hidden_layers]): is_last_layer = (layer_idx == num_hidden_layers - 1) if layer_idx not in pp_layer_list: # keep next layer's input_layernorm's weights for fusion diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index b50273386373..e71eb48c3beb 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -8,29 +8,47 @@ from torch import nn import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils -from tensorrt_llm._utils import (get_sm_version, is_sm_100f, nvtx_range, - nvtx_range_debug) +from tensorrt_llm._utils import get_sm_version, is_sm_100f, nvtx_range, nvtx_range_debug from tensorrt_llm.llmapi.llm_args import SkipSoftmaxAttentionConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from ..attention_backend import (AttentionForwardArgs, AttentionInputType, - AttentionMetadata, FlashInferAttentionMetadata, - TrtllmAttention, TrtllmAttentionMetadata) -from ..attention_backend.interface import (AttentionBackend, AttentionMask, - CustomAttentionMask, - PositionalEmbeddingParams, - PredefinedAttentionMask) +from ..attention_backend import ( + AttentionForwardArgs, + AttentionInputType, + AttentionMetadata, + FlashInferAttentionMetadata, + TrtllmAttention, + TrtllmAttentionMetadata, +) +from ..attention_backend.interface import ( + AttentionBackend, + AttentionMask, + CustomAttentionMask, + PositionalEmbeddingParams, + PredefinedAttentionMask, +) from ..attention_backend.sparse.dsa import ( - DSAtrtllmAttentionMetadata, transform_local_topk_and_prepare_pool_view) + DSAtrtllmAttentionMetadata, + transform_local_topk_and_prepare_pool_view, +) from ..attention_backend.utils import create_attention, get_attention_backend -from ..distributed import (AllReduceParams, HelixAllToAllNative, alltoall_helix, - cp_allgather, reducescatter) +from ..distributed import ( + AllReduceParams, + HelixAllToAllNative, + alltoall_helix, + cp_allgather, + reducescatter, +) from ..model_config import ModelConfig from ..peft.lora.layer import LoraLayer, LoraModuleType -from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, - is_torch_compiling, maybe_compiled_cat, - maybe_compiled_copy_) +from ..utils import ( + Fp4QuantizedTensor, + get_model_extra_attrs, + is_torch_compiling, + maybe_compiled_cat, + maybe_compiled_copy_, +) from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from .multi_stream_utils import maybe_execute_in_parallel from .rms_norm import RMSNorm @@ -1072,11 +1090,17 @@ def mla_custom_op_inplace( latent_cache_gen: Optional[torch.Tensor], ) -> None: metadata, mla_layer = extract_extra_attrs(layer_idx, "mla") - mla_layer.forward_impl(position_ids, - hidden_states, - metadata, - output=output, - latent_cache_gen=latent_cache_gen) + if mla_layer.is_deepseek_v4: + mla_layer.forward_impl_with_deepseek_v4(position_ids, + hidden_states, + metadata, + output=output) + else: + mla_layer.forward_impl(position_ids, + hidden_states, + metadata, + output=output, + latent_cache_gen=latent_cache_gen) @torch.library.custom_op("trtllm::mla_dsa_proj", mutates_args=()) @@ -1246,6 +1270,8 @@ def __init__( config: Optional[ModelConfig] = None, mapping_with_cp: Optional[Mapping] = None, reduce_output: bool = True, + num_groups: int = 1, + o_lora_rank: int = 1024, ): """ Initialize the MLA module. @@ -1268,6 +1294,8 @@ def __init__( dtype (torch.dtype): The data type. dense_bias (bool): Whether to use bias in the output projection layer. config (ModelConfig): The model configuration. + num_groups (int): The number of groups. + o_lora_rank (int): The dimension of the compressed output. """ super().__init__() self.layer_idx = layer_idx @@ -1288,6 +1316,8 @@ def __init__( self.max_position_embeddings = max_position_embeddings self.pos_embd_params = pos_embd_params self.dense_bias = dense_bias + self.num_groups = num_groups + self.o_lora_rank = o_lora_rank if dense_bias is None: self.dense_bias = bias @@ -1307,11 +1337,17 @@ def __init__( self) self.register_to_config = True - # Currently only DSA sparse attention is supported. - if config is not None and config.sparse_attention_config is not None and config.sparse_attention_config.algorithm == "dsa": - self.is_dsa = True - else: - self.is_dsa = False + # Currently only DSA/DeepSeek-V4 sparse attention are supported. + self.is_dsa, self.is_deepseek_v4 = False, False + if config is not None and config.sparse_attention_config is not None: + if config.sparse_attention_config.algorithm == "dsa": + self.is_dsa = True + elif config.sparse_attention_config.algorithm == "deepseek_v4": + self.is_deepseek_v4 = True + else: + raise ValueError( + f"Invalid sparse attention algorithm: {config.sparse_attention_config.algorithm}" + ) # tensor parallel config = config or ModelConfig() @@ -1351,6 +1387,7 @@ def __init__( self.num_heads_tp_cp = self.num_heads_tp // cp_size self.num_key_value_heads_tp = (self.num_key_value_heads + tp_size - 1) // tp_size + self.n_local_groups = self.num_groups // tp_size rms_norm_eps = getattr(config.pretrained_config, "rms_norm_eps", 1e-6) quant_config = config.get_quant_config() @@ -1391,6 +1428,12 @@ def __init__( force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm) + if self.is_deepseek_v4: + # V4 unweighted per-head RMS on Q post-wq_b; q.view(-1, head_dim) at call site. + self.q_b_layernorm = RMSNorm(hidden_size=self.qk_head_dim, + eps=rms_norm_eps, + dtype=dtype, + has_weights=False) else: self.kv_a_proj_with_mqa = Linear( hidden_size, @@ -1419,33 +1462,36 @@ def __init__( use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm) self.q_b_proj = self.q_proj - self.kv_a_layernorm = RMSNorm(hidden_size=kv_lora_rank, + kv_a_layernorm_hidden_size = (self.kv_lora_rank + self.qk_rope_head_dim + if self.is_deepseek_v4 else kv_lora_rank) + self.kv_a_layernorm = RMSNorm(hidden_size=kv_a_layernorm_hidden_size, dtype=dtype, eps=rms_norm_eps) - self.kv_b_proj = Linear( - self.kv_lora_rank, - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), - bias=bias, - dtype=dtype, - mapping=mapping, - tensor_parallel_mode=TensorParallelMode.COLUMN, - quant_config=quant_config, - skip_create_weights_in_init=config.skip_create_weights_in_init, - allreduce_strategy=config.allreduce_strategy, - force_dynamic_quantization=config.force_dynamic_quantization, - use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, - use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm) - # This parameter will view into self.kv_b_proj.weight after loading weights. - # For dummy weight initialization, this parameter is initialized with empty tensor. - # Used in forward_absorption only - self.v_b_proj = nn.Parameter( - torch.empty( - (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank), + if not self.is_deepseek_v4: + self.kv_b_proj = Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=bias, dtype=dtype, - ), - requires_grad=False, - ) + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + allreduce_strategy=config.allreduce_strategy, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm, + use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm) + # This parameter will view into self.kv_b_proj.weight after loading weights. + # For dummy weight initialization, this parameter is initialized with empty tensor. + # Used in forward_absorption only + self.v_b_proj = nn.Parameter( + torch.empty( + (self.num_heads_tp_cp, self.v_head_dim, self.kv_lora_rank), + dtype=dtype, + ), + requires_grad=False, + ) mapping_o = Mapping( world_size=pp_size * dp_size * tp_size * cp_size, @@ -1457,19 +1503,42 @@ def __init__( enable_attention_dp=self.mapping.enable_attention_dp, ) self.mapping_o = mapping_o - self.o_proj = Linear( - self.num_key_value_heads * self.v_head_dim, - self.hidden_size, - bias=self.dense_bias, - dtype=dtype, - mapping=mapping_o, - tensor_parallel_mode=TensorParallelMode.ROW, - quant_config=quant_config, - skip_create_weights_in_init=config.skip_create_weights_in_init, - reduce_output=reduce_output, - allreduce_strategy=config.allreduce_strategy, - force_dynamic_quantization=config.force_dynamic_quantization, - use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm) + if self.is_deepseek_v4: + self.o_a_proj = nn.Parameter( + torch.empty( + (self.n_local_groups, self.o_lora_rank, + self.num_heads * self.qk_head_dim // self.num_groups), + dtype=dtype, + ), + requires_grad=False, + ) + self.o_b_proj = Linear( + self.num_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + dtype=dtype, + mapping=mapping_o, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + reduce_output=reduce_output, + allreduce_strategy=config.allreduce_strategy, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm) + else: + self.o_proj = Linear( + self.num_key_value_heads * self.v_head_dim, + self.hidden_size, + bias=self.dense_bias, + dtype=dtype, + mapping=mapping_o, + tensor_parallel_mode=TensorParallelMode.ROW, + quant_config=quant_config, + skip_create_weights_in_init=config.skip_create_weights_in_init, + reduce_output=reduce_output, + allreduce_strategy=config.allreduce_strategy, + force_dynamic_quantization=config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm) def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: @@ -1495,13 +1564,14 @@ def yarn_get_mscale(scale=1, mscale=1): kv_lora_rank=self.kv_lora_rank, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.kv_lora_rank, + v_head_dim=self.v_head_dim if self.is_deepseek_v4 else self.kv_lora_rank, hidden_size=self.hidden_size, predicted_tokens_per_seq=self.predicted_tokens_per_seq, skip_create_weights_in_init=config.skip_create_weights_in_init, sparse_attention_config=config.sparse_attention_config, dtype=dtype, aux_stream=aux_stream, + rope_append=not self.is_deepseek_v4, ) self.softmax_scale = 1.0 / (math.sqrt(self.qk_head_dim) * q_scaling) @@ -1519,6 +1589,14 @@ def yarn_get_mscale(scale=1, mscale=1): is_neox=pos_embd_params.is_neox, ) + if self.is_deepseek_v4: + self.inverse_rotary_emb = RotaryEmbedding( + pos_embd_params.rope, + head_dim=self.qk_rope_head_dim, + is_neox=pos_embd_params.is_neox, + inverse=True, + ) + # Short-sequence MHA optimization for DSA models: # For short prefill sequences, use MHA (kv_b_proj expansion + standard # attention) instead of the absorption path, which has overhead from @@ -1537,7 +1615,7 @@ def yarn_get_mscale(scale=1, mscale=1): # by DSA for the short-seq path (dense attention, no sparse config). _short_seq_mha = (self.is_dsa and self.short_seq_mha_threshold > 0 and not self.apply_rotary_emb) - if not self.is_dsa or _short_seq_mha: + if (not self.is_dsa or _short_seq_mha) and not self.is_deepseek_v4: self.mha = create_attention( config.attn_backend, self.layer_idx, @@ -1586,22 +1664,29 @@ def create_weights(self): # k_b_proj_trans's dtype must be consistent with self.kv_b_proj, # which can be modified after __init__ - has_fp8_block_scales = ( - self.kv_b_proj.quant_config - and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales()) - - mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype - self.k_b_proj_trans = nn.Parameter( - torch.empty( - (self.num_heads_tp, self.kv_lora_rank, self.qk_nope_head_dim), - dtype=mla_weight_dtype, - ), - requires_grad=False, - ) + if self.is_deepseek_v4: + has_fp8_block_scales = ( + self.o_b_proj.quant_config and + self.o_b_proj.quant_config.quant_mode.has_fp8_block_scales()) + else: + has_fp8_block_scales = ( + self.kv_b_proj.quant_config and + self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales()) + + mla_weight_dtype = torch.float8_e4m3fn if has_fp8_block_scales else self.dtype + self.k_b_proj_trans = nn.Parameter( + torch.empty( + (self.num_heads_tp, self.kv_lora_rank, + self.qk_nope_head_dim), + dtype=mla_weight_dtype, + ), + requires_grad=False, + ) self.k_b_proj_trans_dequant = None self.v_b_proj_dequant = None - if has_fp8_block_scales: + self.o_a_proj_dequant = None + if has_fp8_block_scales and not self.is_deepseek_v4: self.k_b_proj_trans_scale = nn.Parameter( torch.empty( ( @@ -1644,9 +1729,29 @@ def create_weights(self): ), requires_grad=False, ) + elif has_fp8_block_scales: + self.o_a_proj_scale = nn.Parameter( + torch.empty( + (self.n_local_groups, self.o_lora_rank // 128, + self.num_heads * self.qk_head_dim // self.num_groups // + 128), + dtype=torch.float32, + ), + requires_grad=False, + ) + if is_sm_100f(): + self.o_a_proj_dequant = nn.Parameter( + torch.empty( + (self.n_local_groups, self.o_lora_rank, + self.num_heads * self.qk_head_dim // self.num_groups), + dtype=self.dtype, + ), + requires_grad=False, + ) else: self.k_b_proj_trans_scale = None self.v_b_proj_scale = None + self.o_a_proj_scale = None def apply_rope( self, @@ -1698,7 +1803,10 @@ def _attn_forward_gen(self, attn_backend: AttentionBackend, q: torch.Tensor, def create_output(self, hidden_states: torch.Tensor, num_contexts: int): num_tokens = hidden_states.shape[0] - hidden_size = self.o_proj.in_features + if self.is_deepseek_v4: + hidden_size = self.num_heads_tp_cp * self.v_head_dim + else: + hidden_size = self.o_proj.in_features return hidden_states.new_empty([num_tokens, hidden_size], dtype=hidden_states.dtype) @@ -1714,14 +1822,54 @@ def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor: q = (q * attn_scale).to(q.dtype) return q - def forward_impl( - self, - position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - output: torch.Tensor, - latent_cache_gen: Optional[torch.Tensor] = None, - ) -> None: + def _deepseek_v4_o_proj(self, attn_out_latent: torch.Tensor, + position_ids: torch.Tensor) -> torch.Tensor: + num_tokens = attn_out_latent.shape[0] + attn_out_latent = attn_out_latent.view(num_tokens, self.num_heads_tp, + -1) + + # Fused in-place inverse RoPE on the rope portion of each head + torch.ops.trtllm.mla_rope_inplace( + attn_out_latent, position_ids.view(-1), + self.inverse_rotary_emb.rotary_cos_sin, self.num_heads_tp, + self.qk_nope_head_dim, self.qk_rope_head_dim, True, + self.inverse_rotary_emb.is_neox) + + # Output projections + o_lora = torch.empty( + [num_tokens, self.n_local_groups, self.o_lora_rank], + device=attn_out_latent.device, + dtype=attn_out_latent.dtype) + if self.o_a_proj.dtype == torch.bfloat16: + # dim = head_dim * num_head // num_group + # [num_groups, num_tokens, dim] x [num_groups, dim, o_lora_rank] + # -> [num_groups, num_tokens, o_lora_rank] + torch.ops.trtllm.bmm_out( + attn_out_latent.view(num_tokens, self.n_local_groups, + -1).transpose(0, 1), + self.o_a_proj.transpose(1, 2), o_lora.transpose(0, 1)) + elif self.o_a_proj.dtype == torch.float8_e4m3fn: + fp8_block_scaling_bmm_out( + attn_out_latent.view(num_tokens, self.n_local_groups, -1), + self.o_a_proj, + self.o_a_proj_scale, + o_lora.transpose(0, 1), + self.o_a_proj_dequant, + self.use_cute_dsl_blockscaling_bmm, + ) + else: + raise NotImplementedError( + f"Missing bmm impl for dtype: {self.o_a_proj.dtype}.") + o_lora = o_lora.flatten(1) + output = self.o_b_proj(o_lora) + return output + + def forward_impl(self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor, + latent_cache_gen: Optional[torch.Tensor] = None) -> None: """ Forward pass for the MLA module. Writes result into output tensor in-place. @@ -1885,7 +2033,8 @@ def forward_dsa_proj( assert self.mqa is not None, "DSA is only supported in MQA mode" q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1) + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], + -1) q, compressed_kv = maybe_execute_in_parallel( lambda: self.q_a_layernorm(q), @@ -1990,7 +2139,7 @@ def forward_dsa_attn( assert position_ids is not None k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, position_ids) - self.forward_context_dsa( + self.forward_context_sparse_mla( q_ctx, compressed_kv_ctx, k_pe_ctx, @@ -2011,16 +2160,125 @@ def forward_dsa_attn( assert position_ids is not None k_pe_gen = self.apply_rope(q_gen, k_pe_gen, position_ids) - self.forward_generation_dsa( + self.forward_generation_sparse_mla( q_gen, compressed_kv_gen, k_pe_gen, attn_metadata, output[num_ctx_tokens:num_tokens, :], - latent_cache_gen, + latent_cache=latent_cache_gen, topk_indices=topk_indices[num_ctx_tokens:num_tokens, :], ) + def forward_impl_with_deepseek_v4(self, position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor) -> None: + """ + Forward pass for the MLA module with DeepSeek-V4 (always in MQA mode). + + Args: + position_ids (Optional[torch.IntTensor]): The position IDs. + hidden_states (torch.Tensor): The hidden states. + attn_metadata (AttentionMetadata): The attention metadata. + + Returns: + torch.Tensor: The output tensor. + """ + assert self.mha is None and self.mqa is not None, "DeepSeek-V4 is only supported in MQA mode" + # split q, k, v into context and gen batches + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + + hidden_states = hidden_states[:num_tokens, ...] + if position_ids is not None: + position_ids = position_ids[..., :num_tokens] + + q, kv = self.kv_a_proj_with_mqa(hidden_states).split([ + self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim + ], -1) + + q, kv = maybe_execute_in_parallel( + lambda: self.q_a_layernorm(q), + lambda: self.kv_a_layernorm(kv), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + compressed_kv, k_pe = kv.split( + [self.kv_lora_rank, self.qk_rope_head_dim], -1) + qr = q + latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) + + q = self.q_b_proj(q) + # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm reduces per-head. + q = self.q_b_layernorm(q.view(-1, self.qk_head_dim)).view_as(q) + # Indexer + topk_indices = None + if self.indexer is not None: + topk_indices = self.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + ) + + # Compressor + if self.compressor is not None: + self.compressor(hidden_states, attn_metadata) + + assert q.shape[ + 0] == num_tokens, f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}" + + assert output is not None, "output must be provided" + + if num_contexts > 0: + q_ctx = q[:num_ctx_tokens, ...] + topk_indices_ctx = topk_indices[: + num_ctx_tokens, :] if topk_indices is not None else None + compressed_kv_ctx = compressed_kv[:num_ctx_tokens, ...] + k_pe_ctx = k_pe[:num_ctx_tokens, ...] + latent_cache_ctx = latent_cache[:num_ctx_tokens, ...] + if self.apply_rotary_emb: + assert position_ids is not None + k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, position_ids) + + self.forward_context_sparse_mla( + q_ctx, + compressed_kv_ctx, + k_pe_ctx, + attn_metadata, + output[:num_ctx_tokens, :], + position_ids=position_ids[:num_ctx_tokens], + latent_cache=latent_cache_ctx, + topk_indices=topk_indices_ctx, + ) + + if num_generations > 0: + q_gen = q[num_ctx_tokens:, ...] + topk_indices_gen = topk_indices[ + num_ctx_tokens: + num_tokens, :] if topk_indices is not None else None + compressed_kv_gen = compressed_kv[num_ctx_tokens:, ...] + k_pe_gen = k_pe[num_ctx_tokens:, ...] + latent_cache_gen = latent_cache[num_ctx_tokens:, ...] + if self.apply_rotary_emb: + assert position_ids is not None + k_pe_gen = self.apply_rope(q_gen, k_pe_gen, position_ids) + + self.forward_generation_sparse_mla( + q_gen, + compressed_kv_gen, + k_pe_gen, + attn_metadata, + output[num_ctx_tokens:num_tokens, :], + position_ids=position_ids[num_ctx_tokens:num_tokens], + latent_cache=latent_cache_gen, + topk_indices=topk_indices_gen, + ) + def forward_context_default( self, q: torch.Tensor, @@ -2094,7 +2352,7 @@ def _should_use_short_mha(self, attn_metadata: AttentionMetadata, attn_metadata.num_ctx_tokens) return effective_len <= self.short_seq_mha_threshold - def forward_context_dsa( + def forward_context_sparse_mla( self, q: torch.Tensor, compressed_kv: torch.Tensor, @@ -2141,9 +2399,11 @@ def forward_context_dsa( k_pe, attn_metadata, output, + position_ids=position_ids, latent_cache=latent_cache, topk_indices=topk_indices) else: + assert not self.is_deepseek_v4, "DeepSeek-V4 is not supported on pre-blackwell GPUs." return self.forward_sparse_mla_kvcache_bf16(q, latent_cache, attn_metadata, @@ -2151,13 +2411,14 @@ def forward_context_dsa( topk_indices, is_generation=False) - def forward_generation_dsa( + def forward_generation_sparse_mla( self, q: torch.Tensor, compressed_kv: torch.Tensor, k_pe: torch.Tensor, attn_metadata: AttentionMetadata, output: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, latent_cache: Optional[torch.Tensor] = None, topk_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: @@ -2167,9 +2428,11 @@ def forward_generation_dsa( k_pe, attn_metadata, output, + position_ids=position_ids, latent_cache=latent_cache, topk_indices=topk_indices) else: + assert not self.is_deepseek_v4, "DeepSeek-V4 is not supported on pre-blackwell GPUs." return self.forward_sparse_mla_kvcache_bf16(q, latent_cache, attn_metadata, @@ -2569,83 +2832,98 @@ def forward_absorption_generation( dtype=torch.uint8, device=q.device) - fused_q = torch.empty( - [ - num_tokens, self.num_heads_tp, - (self.kv_lora_rank + self.qk_rope_head_dim) - ], - dtype=q.dtype, - device=q.device, - ) - - rope_stream = self.aux_stream if not has_fp8_kv_cache else None - if self.k_b_proj_trans.dtype == torch.bfloat16: - # [num_heads, num_tokens, self.qk_nope_head_dim] - q_nope_t = q_nope.transpose(0, 1) - # [num_heads, num_tokens, self.kv_lora_rank] - q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) - - # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] - # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] - # The output of bmm is written directly into fused_q - maybe_execute_in_parallel( - lambda: self._bmm_bf16_out(q_nope_t, self.k_b_proj_trans, - self.k_b_proj_trans.transpose(1, 2), - q_nope_out), - lambda: self.mqa.mla_rope_generation( - fused_q, - q_pe, - latent_cache, - attn_metadata, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - mla_bmm1_scale, - mla_bmm2_scale, - quant_q_buffer, - ), - self.ln_events[0], - self.ln_events[1], - rope_stream, + if self.is_deepseek_v4: + fused_q = q + self.mqa.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ) + else: + fused_q = torch.empty( + [ + num_tokens, self.num_heads_tp, + (self.kv_lora_rank + self.qk_rope_head_dim) + ], + dtype=q.dtype, + device=q.device, ) - elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: - # [num_heads, num_tokens, self.kv_lora_rank] - q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) + rope_stream = self.aux_stream if not has_fp8_kv_cache else None + if self.k_b_proj_trans.dtype == torch.bfloat16: + # [num_heads, num_tokens, self.qk_nope_head_dim] + q_nope_t = q_nope.transpose(0, 1) + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) + + # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] + # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] + # The output of bmm is written directly into fused_q + maybe_execute_in_parallel( + lambda: self._bmm_bf16_out( + q_nope_t, self.k_b_proj_trans, + self.k_b_proj_trans.transpose(1, 2), q_nope_out), + lambda: self.mqa.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ), + self.ln_events[0], + self.ln_events[1], + rope_stream, + ) - maybe_execute_in_parallel( - lambda: fp8_block_scaling_bmm_out( - q_nope, - self.k_b_proj_trans, - self.k_b_proj_trans_scale, - q_nope_out, - self.k_b_proj_trans_dequant, - self.use_cute_dsl_blockscaling_bmm, - ), - lambda: self.mqa.mla_rope_generation( - fused_q, - q_pe, - latent_cache, - attn_metadata, - cu_q_seqlens, - cu_kv_seqlens, - fmha_scheduler_counter, - mla_bmm1_scale, - mla_bmm2_scale, - quant_q_buffer, - ), - self.ln_events[0], - self.ln_events[1], - rope_stream, - ) - else: - raise NotImplementedError( - f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") + elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) + + maybe_execute_in_parallel( + lambda: fp8_block_scaling_bmm_out( + q_nope, + self.k_b_proj_trans, + self.k_b_proj_trans_scale, + q_nope_out, + self.k_b_proj_trans_dequant, + self.use_cute_dsl_blockscaling_bmm, + ), + lambda: self.mqa.mla_rope_generation( + fused_q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ), + self.ln_events[0], + self.ln_events[1], + rope_stream, + ) + else: + raise NotImplementedError( + f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") - fused_q = fused_q.view([ - num_tokens, - self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim) - ]) + fused_q = fused_q.view([ + num_tokens, + self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim) + ]) # Use generation_only for generation phase and context_only for context phase in DSA attention attention_input_type = AttentionInputType.generation_only @@ -2659,6 +2937,7 @@ def forward_absorption_generation( attn_metadata, attention_input_type=attention_input_type, out_scale=self.out_scale, + output=output if self.is_deepseek_v4 else None, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by `invokeMLARopeGeneration` topk_indices=topk_indices, # used by DSA attention @@ -2672,6 +2951,16 @@ def forward_absorption_generation( ) fused_q = None + if self.is_deepseek_v4: + if self.mapping.has_cp_helix(): + raise RuntimeError( + "DeepSeek-V4 + CP Helix is not supported: " + "_helix_post_process returns a different tensor, " + "bypassing the pre-allocated output buffer.") + assert attn_out_latent.data_ptr() == output.data_ptr(), \ + "Attention backend did not write into the provided output buffer." + return output + # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert (attn_out_latent.shape[0] == q.shape[0] and attn_out_latent.shape[1] @@ -2717,53 +3006,58 @@ def forward_absorption_context( topk_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_tokens = q.shape[0] + q_nope, q_pe = q.view([-1, self.num_heads_tp, self.qk_head_dim]).split( [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] - # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp - fused_q = torch.empty( - [ - num_tokens, self.num_heads_tp, - (self.kv_lora_rank + self.qk_rope_head_dim) - ], - dtype=q.dtype, - device=q.device, - ) - - if self.k_b_proj_trans.dtype == torch.bfloat16: - # [num_heads, num_tokens, self.qk_nope_head_dim] - q_nope_t = q_nope.transpose(0, 1) - # [num_heads, num_tokens, self.kv_lora_rank] - q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) - - # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] - # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] - # The output of bmm is written directly into fused_q - self._bmm_bf16_out(q_nope_t, self.k_b_proj_trans, - self.k_b_proj_trans.transpose(1, 2), q_nope_out) - elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: - # [num_heads, num_tokens, self.kv_lora_rank] - q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) - - fp8_block_scaling_bmm_out( - q_nope, - self.k_b_proj_trans, - self.k_b_proj_trans_scale, - q_nope_out, - self.k_b_proj_trans_dequant, - self.use_cute_dsl_blockscaling_bmm, - ) + if self.is_deepseek_v4: + fused_q = q else: - raise NotImplementedError( - f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") + # fused_q contains 1) the result of the following bmm with shape [num_tokens, num_heads, kv_lora_rank] + # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp + fused_q = torch.empty( + [ + num_tokens, self.num_heads_tp, + (self.kv_lora_rank + self.qk_rope_head_dim) + ], + dtype=q.dtype, + device=q.device, + ) - if self.apply_rotary_emb: - fused_q[..., self.kv_lora_rank:] = q_pe - fused_q = fused_q.view([ - num_tokens, - self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim) - ]) + if self.k_b_proj_trans.dtype == torch.bfloat16: + # [num_heads, num_tokens, self.qk_nope_head_dim] + q_nope_t = q_nope.transpose(0, 1) + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) + + # [num_heads, num_tokens, self.qk_nope_head_dim] x [num_heads, kv_lora_rank, qk_nope_head_dim] + # -> [num_heads, num_tokens, kv_lora_rank] -> [num_tokens, num_heads, kv_lora_rank] + # The output of bmm is written directly into fused_q + self._bmm_bf16_out(q_nope_t, self.k_b_proj_trans, + self.k_b_proj_trans.transpose(1, 2), + q_nope_out) + elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn: + # [num_heads, num_tokens, self.kv_lora_rank] + q_nope_out = fused_q[..., :self.kv_lora_rank].transpose(0, 1) + + fp8_block_scaling_bmm_out( + q_nope, + self.k_b_proj_trans, + self.k_b_proj_trans_scale, + q_nope_out, + self.k_b_proj_trans_dequant, + self.use_cute_dsl_blockscaling_bmm, + ) + else: + raise NotImplementedError( + f"Missing bmm impl for dtype: {self.k_b_proj_trans.dtype}.") + + if self.apply_rotary_emb: + fused_q[..., self.kv_lora_rank:] = q_pe + fused_q = fused_q.view([ + num_tokens, + self.num_heads_tp * (self.kv_lora_rank + self.qk_rope_head_dim) + ]) # Use generation_only for generation phase and context_only for context phase in DSA attention attention_input_type = AttentionInputType.context_only @@ -2776,12 +3070,23 @@ def forward_absorption_context( attn_metadata, attention_input_type=attention_input_type, out_scale=self.out_scale, + output=output if self.is_deepseek_v4 else None, latent_cache=latent_cache, # kvcache and k_pe q_pe=q_pe, # used by `invokeMLARopeGeneration` topk_indices=topk_indices, # used by DSA attention ) fused_q = None + if self.is_deepseek_v4: + if self.mapping.has_cp_helix(): + raise RuntimeError( + "DeepSeek-V4 + CP Helix is not supported: " + "_helix_post_process returns a different tensor, " + "bypassing the pre-allocated output buffer.") + assert attn_out_latent.data_ptr() == output.data_ptr(), \ + "Attention backend did not write into the provided output buffer." + return output + # note: if we do not have CP, then num_heads_tp_cp == num_heads_tp assert (attn_out_latent.shape[0] == q.shape[0] and attn_out_latent.shape[1] @@ -2982,6 +3287,8 @@ def forward( q, compressed_kv, k_pe, latent_cache, indexer_intermediates, position_ids, self.layer_idx_str, attn_output) else: + # DeepSeek-V4 and vanilla MLA both use the single custom op. + # DeepSeek-V4 dispatches to forward_impl_with_deepseek_v4 inside. torch.ops.trtllm.mla_custom_op_inplace(hidden_states, position_ids, self.layer_idx_str, @@ -2992,6 +3299,11 @@ def forward( hidden_states, attn_metadata, output=attn_output) + elif self.is_deepseek_v4: + self.forward_impl_with_deepseek_v4(position_ids, + hidden_states, + attn_metadata, + output=attn_output) else: self.forward_impl(position_ids, hidden_states, @@ -2999,11 +3311,12 @@ def forward( output=attn_output, latent_cache_gen=latent_cache_gen) - attn_output = _helix_cp_output_projection(self.o_proj, attn_output, - attn_metadata, - all_reduce_params, - self.mapping, self.mapping_o, - self.layer_idx) + if self.is_deepseek_v4: + attn_output = self._deepseek_v4_o_proj(attn_output, position_ids) + else: + attn_output = _helix_cp_output_projection( + self.o_proj, attn_output, attn_metadata, all_reduce_params, + self.mapping, self.mapping_o, self.layer_idx) return attn_output def resmooth_parameters(self, @@ -3027,9 +3340,13 @@ def resmooth_parameters(self, return weight_param, scale_param def post_load_weights(self): - has_fp8_block_scales = ( - self.kv_b_proj.quant_config - and self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales()) + # In DeepSeek-V4 mode, kv_b_proj doesn't exist + if self.is_deepseek_v4: + has_fp8_block_scales = False + else: + has_fp8_block_scales = ( + self.kv_b_proj.quant_config and + self.kv_b_proj.quant_config.quant_mode.has_fp8_block_scales()) is_sm120 = get_sm_version() == 120 if is_sm120 and has_fp8_block_scales: self.k_b_proj_trans, self.k_b_proj_trans_scale = self.resmooth_parameters( diff --git a/tensorrt_llm/_torch/modules/engram/__init__.py b/tensorrt_llm/_torch/modules/engram/__init__.py new file mode 100644 index 000000000000..ff1d218a8e4a --- /dev/null +++ b/tensorrt_llm/_torch/modules/engram/__init__.py @@ -0,0 +1,18 @@ +# 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. + +from tensorrt_llm._torch.modules.engram.engram import Engram, EngramConfig, EngramHashProvider + +__all__ = ["Engram", "EngramConfig", "EngramHashProvider"] diff --git a/tensorrt_llm/_torch/modules/engram/engram.py b/tensorrt_llm/_torch/modules/engram/engram.py new file mode 100644 index 000000000000..41ddc5b4b6ab --- /dev/null +++ b/tensorrt_llm/_torch/modules/engram/engram.py @@ -0,0 +1,846 @@ +# 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. +""" +Engram module implementation for TensorRT-LLM PyTorch backend. + +The Engram module provides n-gram based hash embeddings that augment +transformer hidden states with local context information. + +All operations run on device: + - Hash IDs (GPU) → embedding → flatten → linear projections (GEMMs) + → normed keys + projected value + - GPU main stream forward: SDP gating + short conv +""" + +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from tokenizers import Regex, normalizers +from torch import nn +from transformers import AutoTokenizer + +# Default Engram embedding vocabulary size per n-gram level. +# Derived from the DeepSeek-V3 tokenizer vocab size (129280) scaled by 5. +_DEFAULT_ENGRAM_VOCAB_SIZE = 129280 * 5 + + +@dataclass +class EngramConfig: + """Configuration for the Engram module. + + Attributes: + tokenizer_name_or_path: Path or name of the HuggingFace tokenizer. + engram_vocab_size: List of vocabulary sizes for each n-gram level (2-gram, 3-gram, etc.). + max_ngram_size: Maximum n-gram size to compute hashes for. + n_embed_per_ngram: Embedding dimension for each n-gram. + n_head_per_ngram: Number of attention heads per n-gram. + layer_ids: List of layer indices where Engram modules are applied. + pad_id: Token ID used for padding. + seed: Random seed for hash multiplier generation. + kernel_size: Kernel size for the short convolution. + hidden_size: Hidden dimension of the backbone model. + hc_mult: Hyper-connection multiplier (number of residual streams). + norm_eps: Epsilon for RMSNorm. + dtype: Data type for model parameters (e.g., torch.float32, torch.bfloat16). + """ + + tokenizer_name_or_path: str = "deepseek-ai/DeepSeek-V3" + engram_vocab_size: List[int] = field( + default_factory=lambda: [_DEFAULT_ENGRAM_VOCAB_SIZE, _DEFAULT_ENGRAM_VOCAB_SIZE] + ) + max_ngram_size: int = 3 + n_embed_per_ngram: int = 512 + n_head_per_ngram: int = 8 + layer_ids: List[int] = field(default_factory=lambda: [1, 15]) + pad_id: int = 2 + seed: int = 0 + kernel_size: int = 4 + hidden_size: int = 1024 + hc_mult: int = 4 + norm_eps: float = 1e-5 + dtype: Optional[torch.dtype] = None + + +class CompressedTokenizer: + """Tokenizer wrapper that normalizes and compresses tokens. + + This class builds a lookup table mapping original token IDs to + normalized/compressed token IDs, reducing vocabulary size for + more efficient n-gram hashing. + + If a pre-built ``lookup_table`` (torch.Tensor of shape [vocab_size]) + and ``num_new_token`` count are provided, the tokenizer download is + skipped entirely — useful to avoid cold-start network fetches. + """ + + def __init__( + self, + tokenizer_name_or_path: str, + lookup_table: Optional[torch.Tensor] = None, + num_new_token: Optional[int] = None, + ): + if lookup_table is not None and num_new_token is not None: + self.lookup_table = lookup_table + self.num_new_token = num_new_token + return + + self.tokenizer = AutoTokenizer.from_pretrained( + tokenizer_name_or_path, trust_remote_code=True + ) + + SENTINEL = "\ue000" + self.normalizer = normalizers.Sequence( + [ + normalizers.NFKC(), + normalizers.NFD(), + normalizers.StripAccents(), + normalizers.Lowercase(), + normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), + normalizers.Replace(Regex(r"^ $"), SENTINEL), + normalizers.Strip(), + normalizers.Replace(SENTINEL, " "), + ] + ) + + self.lookup_table, self.num_new_token = self._build_lookup_table() + + def __len__(self) -> int: + return self.num_new_token + + def _build_lookup_table(self): + old2new = {} + key2new = {} + new_tokens = [] + + vocab_size = len(self.tokenizer) + for tid in range(vocab_size): + text = self.tokenizer.decode([tid], skip_special_tokens=False) + + if "\ufffd" in text: + key = self.tokenizer.convert_ids_to_tokens(tid) + else: + norm = self.normalizer.normalize_str(text) + key = norm if norm else text + + nid = key2new.get(key) + if nid is None: + nid = len(new_tokens) + key2new[key] = nid + new_tokens.append(key) + old2new[tid] = nid + + lookup = torch.tensor([old2new[tid] for tid in range(vocab_size)], dtype=torch.long) + + return lookup, len(new_tokens) + + def _compress(self, input_ids: torch.Tensor) -> torch.Tensor: + if not isinstance(input_ids, torch.Tensor): + input_ids = torch.tensor(input_ids, dtype=torch.long) + ids = input_ids.long() + # Preserve negative padding IDs while compressing valid token IDs. + vocab_size = len(self.lookup_table) + compressed = self.lookup_table[ids.clamp(0, vocab_size - 1)] + return torch.where(ids < 0, ids, compressed) + + def __call__(self, input_ids): + return self._compress(input_ids) + + +def _is_prime(n: int) -> bool: + """Deterministic Miller-Rabin primality test for n < 3.3e24.""" + if n < 2: + return False + if n < 4: + return True + if n % 2 == 0 or n % 3 == 0: + return False + # Write n-1 as 2^r * d + d, r = n - 1, 0 + while d % 2 == 0: + d //= 2 + r += 1 + # Witnesses sufficient for n < 3.3e24 + for a in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): + if a >= n: + continue + x = pow(a, d, n) + if x == 1 or x == n - 1: + continue + for _ in range(r - 1): + x = pow(x, 2, n) + if x == n - 1: + break + else: + return False + return True + + +def _find_next_prime(start: int, seen_primes: set) -> int: + """Find the next prime number greater than start that is not in seen_primes.""" + candidate = start + 1 + while True: + if _is_prime(candidate) and candidate not in seen_primes: + return candidate + candidate += 1 + + +class NgramHashMapping: + """Computes n-gram hash indices for embedding lookup. + + This class generates per-layer, per-head hash mappings for n-grams, + using prime moduli to reduce hash collisions across heads. + """ + + def __init__( + self, + engram_vocab_size: List[int], + max_ngram_size: int, + n_embed_per_ngram: int, + n_head_per_ngram: int, + layer_ids: List[int], + tokenizer_name_or_path: str, + pad_id: int, + seed: int, + ): + self.vocab_size_per_ngram = engram_vocab_size + self.max_ngram_size = max_ngram_size + self.n_embed_per_ngram = n_embed_per_ngram + self.n_head_per_ngram = n_head_per_ngram + self.pad_id = pad_id + self.layer_ids = layer_ids + + self.compressed_tokenizer = CompressedTokenizer( + tokenizer_name_or_path=tokenizer_name_or_path + ) + self.tokenizer_vocab_size = len(self.compressed_tokenizer) + if self.pad_id is not None: + self.pad_id = int(self.compressed_tokenizer.lookup_table[self.pad_id].item()) + + max_long = torch.iinfo(torch.long).max + M_max = int(max_long // self.tokenizer_vocab_size) + half_bound = max(1, M_max // 2) + PRIME_1 = 10007 + + self.layer_multipliers: Dict[int, torch.Tensor] = {} + + for layer_id in self.layer_ids: + base_seed = int(seed + PRIME_1 * int(layer_id)) + # Use numpy RNG for deterministic reproducibility with existing + # checkpoints — torch.Generator uses a different algorithm and + # would produce incompatible multipliers for the same seed. + g = np.random.default_rng(base_seed) + r = g.integers(low=0, high=half_bound, size=(self.max_ngram_size,), dtype=np.int64) + multipliers = torch.tensor(r * 2 + 1, dtype=torch.long) + self.layer_multipliers[layer_id] = multipliers + + self.vocab_size_across_layers = self._calculate_vocab_size_across_layers() + + def _calculate_vocab_size_across_layers(self) -> Dict[int, List[List[int]]]: + seen_primes = set() + vocab_size_across_layers = {} + + for layer_id in self.layer_ids: + all_ngram_vocab_sizes = [] + for ngram in range(2, self.max_ngram_size + 1): + current_ngram_heads_sizes = [] + + vocab_size = self.vocab_size_per_ngram[ngram - 2] + num_head = self.n_head_per_ngram + current_prime_search_start = vocab_size - 1 + + for _ in range(num_head): + found_prime = _find_next_prime(current_prime_search_start, seen_primes) + seen_primes.add(found_prime) + current_ngram_heads_sizes.append(found_prime) + current_prime_search_start = found_prime + + all_ngram_vocab_sizes.append(current_ngram_heads_sizes) + vocab_size_across_layers[layer_id] = all_ngram_vocab_sizes + + return vocab_size_across_layers + + def _get_ngram_hashes( + self, + input_ids: torch.Tensor, + layer_id: int, + ) -> torch.Tensor: + x = input_ids.long() + (T,) = x.shape + + multipliers = self.layer_multipliers[layer_id] + + def shift_k(k: int) -> torch.Tensor: + if k == 0: + return x + return torch.nn.functional.pad(x, (k, 0), value=self.pad_id)[:T] + + base_shifts = [shift_k(k) for k in range(self.max_ngram_size)] + + all_hashes: List[torch.Tensor] = [] + + for n in range(2, self.max_ngram_size + 1): + n_gram_index = n - 2 + tokens = base_shifts[:n] + mix = tokens[0] * multipliers[0] + for k in range(1, n): + mix = torch.bitwise_xor(mix, tokens[k] * multipliers[k]) + head_vocab_sizes = self.vocab_size_across_layers[layer_id][n_gram_index] + + for j in range(self.n_head_per_ngram): + mod = int(head_vocab_sizes[j]) + all_hashes.append(mix % mod) + + return torch.stack(all_hashes, dim=1) + + def hash(self, input_ids) -> Dict[int, torch.Tensor]: + """Compute hash indices for all configured layers. + + Args: + input_ids: Token IDs of shape ``[T]``. + + Returns: + Dictionary mapping layer_id to hash indices of shape ``[T, num_heads]``. + """ + input_ids = self.compressed_tokenizer(input_ids) + hash_ids_for_all_layers = {} + for layer_id in self.layer_ids: + hash_ids_for_all_layers[layer_id] = self._get_ngram_hashes(input_ids, layer_id=layer_id) + return hash_ids_for_all_layers + + def hash_single_layer(self, input_ids, layer_id: int) -> torch.Tensor: + """Compute hash indices for a single layer. + + Args: + input_ids: Token IDs of shape ``[T]`` (already compressed or raw). + layer_id: The layer to compute hashes for. + + Returns: + Hash indices of shape ``[T, num_heads]``. + """ + if not isinstance(input_ids, torch.Tensor): + input_ids = torch.tensor(input_ids, dtype=torch.long) + return self._get_ngram_hashes(input_ids, layer_id=layer_id) + + +class EngramHashProvider: + """Computes and caches n-gram hash indices for all Engram layers. + + All hash computation runs on GPU using PyTorch ops. CPU-side + ``NgramHashMapping`` is used only at init to derive constants + (lookup table, multipliers, moduli) and then discarded. + + Usage: + # At model initialization + hash_provider = EngramHashProvider(config) + + # At each forward pass + hash_cache = hash_provider.compute_hashes(input_ids) + + # In each Engram layer + precomputed = engram_layer.precompute(hash_cache[layer_id]) + output = engram_layer(hidden_states, precomputed=precomputed) + """ + + def __init__(self, config: EngramConfig): + self.config = config + + # Use NgramHashMapping only to derive constants, then discard it. + hash_mapping = NgramHashMapping( + engram_vocab_size=config.engram_vocab_size, + max_ngram_size=config.max_ngram_size, + n_embed_per_ngram=config.n_embed_per_ngram, + n_head_per_ngram=config.n_head_per_ngram, + layer_ids=config.layer_ids, + tokenizer_name_or_path=config.tokenizer_name_or_path, + pad_id=config.pad_id, + seed=config.seed, + ) + + # Store vocab sizes directly (needed by Engram layer init). + self._vocab_size_across_layers = hash_mapping.vocab_size_across_layers + + # GPU tensors for on-device hash computation. + # Created on CPU and lazily moved to GPU on first use. + self._lookup_table = hash_mapping.compressed_tokenizer.lookup_table.clone() + self._pad_id = hash_mapping.pad_id + + self._multipliers: Dict[int, torch.Tensor] = {} + self._modules: Dict[int, List[torch.Tensor]] = {} + for layer_id in config.layer_ids: + self._multipliers[layer_id] = hash_mapping.layer_multipliers[layer_id].clone() + self._modules[layer_id] = [ + torch.tensor(ngram_head_sizes, dtype=torch.long) + for ngram_head_sizes in hash_mapping.vocab_size_across_layers[layer_id] + ] + + self._device: Optional[torch.device] = None + + # Cache for CUDA graph capture - stores last computed hashes. + # _cached_hashes points to the most-recently used entry. + # _cached_hashes_store keeps ALL per-shape entries alive so that + # CUDA graph tensor addresses remain valid after subsequent compute calls + # change the active shape (e.g. warmup batch_size=4 → 3 → 2 → 1). + self._cached_hashes: Optional[Dict[int, torch.Tensor]] = None + self._cached_hashes_store: Dict[tuple, Dict[int, torch.Tensor]] = {} + + def _ensure_on_device(self, device: torch.device): + """Move hash tensors to the specified device (lazy, once).""" + if self._device == device: + return + self._lookup_table = self._lookup_table.to(device) + for layer_id in list(self._multipliers.keys()): + self._multipliers[layer_id] = self._multipliers[layer_id].to(device) + self._modules[layer_id] = [m.to(device) for m in self._modules[layer_id]] + self._device = device + + def compute_hashes( + self, + input_ids: torch.Tensor, + seq_lens: Optional[torch.Tensor] = None, + ) -> Dict[int, torch.Tensor]: + """Compute hash indices on GPU using PyTorch ops. + + Runs the compressed-tokenizer lookup, n-gram mixing, and modular + hashing entirely on the device where ``input_ids`` resides. + + During CUDA graph capture, returns cached hashes from the warmup pass. + + Args: + input_ids: Flattened token IDs of shape ``[T]`` on GPU. + When the tensor contains multiple packed sequences, pass + ``seq_lens`` so that n-gram hashing does not cross + sequence boundaries. + seq_lens: Optional 1-D tensor of per-sequence lengths. When + provided, shifted n-gram windows that would cross a + sequence boundary are replaced with ``pad_id``. + + Returns: + Dictionary mapping layer_id to hash indices as torch.Tensor + of shape ``[T, num_heads]`` on the same device as input_ids. + """ + if torch.cuda.is_current_stream_capturing(): + if self._cached_hashes is not None: + return self._cached_hashes + raise RuntimeError( + "EngramHashProvider.compute_hashes() called during CUDA " + "graph capture but no cached hashes available. " + "Ensure warmup runs before capture." + ) + + device = input_ids.device + self._ensure_on_device(device) + + # 1. Compress token ids via lookup table (1-D) + ids = input_ids.long().view(-1) + vocab_size = self._lookup_table.shape[0] + compressed = self._lookup_table[ids.clamp(0, vocab_size - 1)] + + (T,) = compressed.shape + + # 2. Pre-compute shifted versions (left-padded with pad_id) + shifts = [compressed] + for k in range(1, self.config.max_ngram_size): + padded = torch.nn.functional.pad(compressed, (k, 0), value=self._pad_id) + shifts.append(padded[:T]) + + # 2b. When input_ids is a packed/flattened tensor containing multiple + # sequences, mask out shifted positions that cross sequence boundaries + # so n-gram hashes stay within each sequence. + if seq_lens is not None and seq_lens.numel() > 1: + cum_lens = torch.cumsum(seq_lens, dim=0) + seq_starts = torch.cat( + [ + torch.zeros(1, device=device, dtype=cum_lens.dtype), + cum_lens[:-1], + ] + ) + positions = torch.arange(T, device=device) + # Map each position to its owning sequence + seq_idx = torch.searchsorted(cum_lens, positions, right=True) + seq_idx = seq_idx.clamp(max=seq_lens.numel() - 1) + pos_in_seq = positions - seq_starts[seq_idx] + + for k in range(1, self.config.max_ngram_size): + boundary_mask = pos_in_seq < k # [T] + shifts[k][boundary_mask] = self._pad_id + + # 3. Compute hashes per layer + result: Dict[int, torch.Tensor] = {} + for layer_id in self.config.layer_ids: + multipliers = self._multipliers[layer_id] + all_hashes: List[torch.Tensor] = [] + + for n in range(2, self.config.max_ngram_size + 1): + n_gram_index = n - 2 + mix = shifts[0] * multipliers[0] + for k in range(1, n): + mix = torch.bitwise_xor(mix, shifts[k] * multipliers[k]) + + moduli = self._modules[layer_id][n_gram_index] + for j in range(self.config.n_head_per_ngram): + head_hash = mix % int(moduli[j].item()) + all_hashes.append(head_hash) + + result[layer_id] = torch.stack(all_hashes, dim=1) + + # Cache for CUDA graph capture. + # Use a per-shape store so tensors captured by CUDA graphs are never + # garbage-collected when a subsequent call uses a different batch size. + # When the same shape is seen again, update the existing tensors + # in-place so CUDA graphs can read the freshly computed values from + # the same memory addresses used at capture time. + shape_key = tuple(tuple(v.shape) for v in result.values()) + if shape_key in self._cached_hashes_store: + cached = self._cached_hashes_store[shape_key] + for layer_id, hashes in result.items(): + cached[layer_id].copy_(hashes) + self._cached_hashes = cached + else: + self._cached_hashes_store[shape_key] = result + self._cached_hashes = result + return self._cached_hashes + + @property + def layer_ids(self) -> List[int]: + """Return the list of layer IDs that have Engram modules.""" + return self.config.layer_ids + + @property + def vocab_size_across_layers(self) -> Dict[int, List[List[int]]]: + """Return the vocabulary sizes for each layer and head.""" + return self._vocab_size_across_layers + + +class MultiHeadEmbedding(nn.Module): + """Multi-head embedding layer with per-head offset handling. + + Each head has its own embedding space, indexed by applying offsets + to the input indices before a single embedding lookup. + """ + + def __init__(self, list_of_N: List[int], D: int, dtype: Optional[torch.dtype] = None): + super().__init__() + self.num_heads = len(list_of_N) + self.embedding_dim = D + + offsets = [0] + for n in list_of_N[:-1]: + offsets.append(offsets[-1] + n) + + self.register_buffer("offsets", torch.tensor(offsets, dtype=torch.long)) + + total_N = sum(list_of_N) + self.embedding = nn.Embedding(num_embeddings=total_N, embedding_dim=D, dtype=dtype) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + """Look up embeddings for multi-head input indices. + + Args: + input_ids: Indices of shape ``[T, num_heads]``. + + Returns: + Embeddings of shape ``[T, num_heads, D]``. + """ + shifted_input_ids = input_ids + self.offsets + # Clamp to valid embedding range to prevent CUDA kernel OOB assertion. + # head_hash values should be in [0, prime-1] by construction, but this + # acts as a defensive guard in case of unexpected inputs. + shifted_input_ids = shifted_input_ids.clamp(0, self.embedding.num_embeddings - 1) + output = self.embedding(shifted_input_ids) + return output + + +class ShortConv(nn.Module): + """Short depthwise convolution with RMSNorm and SiLU activation. + + Applies a causal depthwise convolution across the sequence dimension + with per-hyper-connection-stream normalization. + + Note: During token-by-token autoregressive generation (seq_len=1), + the Conv1d sees only the current token plus zero padding — no ring + buffer or conv state carries over from previous steps. This means + the short-range context capture is limited to prefill. A proper + conv state cache for generation is left as a future enhancement. + """ + + def __init__( + self, + hidden_size: int, + kernel_size: int = 4, + dilation: int = 1, + norm_eps: float = 1e-5, + hc_mult: int = 4, + activation: bool = True, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + self.hc_mult = hc_mult + self.activation = activation + + self.hidden_size = hidden_size + self.norm_eps = norm_eps + + total_channels = hidden_size * hc_mult + self.conv = nn.Conv1d( + in_channels=total_channels, + out_channels=total_channels, + kernel_size=kernel_size, + groups=total_channels, + bias=False, + padding=(kernel_size - 1) * dilation, + dilation=dilation, + dtype=dtype, + ) + + # Stacked RMSNorm weight [hc_mult, D] — single vectorised norm + # instead of hc_mult separate kernel launches. + self.norm_weight = nn.Parameter(torch.ones(hc_mult, hidden_size, dtype=dtype)) + + if self.activation: + self.act_fn = nn.SiLU() + + def load_weights(self, weights): + """Load weights, handling legacy ``norms.{i}.weight`` checkpoint keys. + + Legacy checkpoints store per-stream RMSNorm weights as + ``norms.0.weight``, ``norms.1.weight``, … This method stacks + them into the single ``norm_weight`` parameter. + """ + w = weights[0] if isinstance(weights, list) else weights + # Check for legacy per-stream norm keys + legacy_keys = [f"norms.{i}.weight" for i in range(self.hc_mult)] + if legacy_keys[0] in w: + stacked = torch.stack([w[k][:] for k in legacy_keys], dim=0) + self.norm_weight.data.copy_(stacked) + # Load remaining keys (e.g. conv.weight) via the generic path + for n, p in self.named_parameters(): + if n == "norm_weight": + continue + if n in w: + p.data.copy_(w[n][:]) + else: + for n, p in self.named_parameters(): + if n in w: + p.data.copy_(w[n][:]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply short convolution with normalization. + + Args: + x: Input tensor of shape ``[T, HC_MULT, D]``. + + Returns: + Output tensor of shape ``[T, HC_MULT, D]``. + """ + T, G, C = x.shape + + assert G == self.hc_mult, f"Input groups {G} != hc_mult {self.hc_mult}" + + rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.norm_eps).rsqrt() + x_normed = x * rms * self.norm_weight # [hc_mult, D] broadcasts over [T, G, D] + + x_norm = x_normed.reshape(T, G * C) + # Conv1d expects [N, C, L]; use N=1 so that L=T is the sequence dim. + x_nct = x_norm.unsqueeze(0).transpose(1, 2) # [1, G*C, T] + y_nct = self.conv(x_nct) + # Truncate to maintain causal masking + y_nct = y_nct[..., :T] + + if self.activation: + y_nct = self.act_fn(y_nct) + y = y_nct.transpose(1, 2).squeeze(0).view(T, G, C).contiguous() + + return y + + +class Engram(nn.Module): + """Engram module for n-gram based context augmentation. + + The Engram module computes n-gram hash embeddings from input tokens + and uses gated attention to augment the hidden states of the model. + + All operations run on device. ``precompute()`` is called on a separate + CUDA stream to overlap with other layers; the main stream performs + only SDP gating + short conv. + + Known limitations: + - **No TP sharding**: The embedding table is replicated on every rank. + For large-scale deployments with TP > 1 this means redundant memory. + - **No conv state for generation**: See ``ShortConv`` docstring. + """ + + def __init__( + self, + layer_id: int, + config: EngramConfig, + vocab_sizes_flat: Optional[List[int]] = None, + stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + self.layer_id = layer_id + self.config = config + self.stream = stream + self.sync_event: Optional[torch.cuda.Event] = ( + torch.cuda.Event() if stream is not None else None + ) + + # If vocab_sizes not provided, compute them (for standalone usage) + if vocab_sizes_flat is None: + hash_mapping = NgramHashMapping( + engram_vocab_size=config.engram_vocab_size, + max_ngram_size=config.max_ngram_size, + n_embed_per_ngram=config.n_embed_per_ngram, + n_head_per_ngram=config.n_head_per_ngram, + layer_ids=config.layer_ids, + tokenizer_name_or_path=config.tokenizer_name_or_path, + pad_id=config.pad_id, + seed=config.seed, + ) + vocab_sizes_flat = [ + x for y in hash_mapping.vocab_size_across_layers[layer_id] for x in y + ] + self.hash_mapping = hash_mapping + else: + self.hash_mapping = None + + embed_dim_per_head = config.n_embed_per_ngram // config.n_head_per_ngram + + dtype = config.dtype + + self.multi_head_embedding = MultiHeadEmbedding( + list_of_N=vocab_sizes_flat, + D=embed_dim_per_head, + dtype=dtype, + ) + + self.short_conv = ShortConv( + hidden_size=config.hidden_size, + kernel_size=config.kernel_size, + dilation=config.max_ngram_size, + hc_mult=config.hc_mult, + norm_eps=config.norm_eps, + dtype=dtype, + ) + + engram_hidden_size = (config.max_ngram_size - 1) * config.n_embed_per_ngram + hc = config.hc_mult + D = config.hidden_size + + # Fused projection: one GEMM produces value (1 head) + all keys (hc_mult heads). + # Output layout: [value (D) | key_0 (D) | key_1 (D) | ... | key_{hc-1} (D)] + self.kv_proj = nn.Linear(engram_hidden_size, (1 + hc) * D, bias=False, dtype=dtype) + + # Per-HC-stream RMSNorm weights for keys and queries, stored as + # stacked parameters so we can apply a single vectorised norm + # instead of hc_mult separate kernel launches. + self.key_norm_weight = nn.Parameter(torch.ones(hc, D, dtype=dtype)) + self.query_norm_weight = nn.Parameter(torch.ones(hc, D, dtype=dtype)) + self.norm_eps = config.norm_eps + + # CUDA graph capture cache + self._cached_embeddings: Optional[torch.Tensor] = None + + def precompute( + self, + hash_indices: torch.Tensor, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + """Pre-compute embeddings from hash indices. + + When a stream was supplied at construction time, the work is + dispatched onto that stream (overlapping with the main stream) and + ``sync_event`` is recorded so the caller can synchronize before + consuming the result. When no stream is set the computation runs + on the current stream synchronously and ``sync_event`` is None. + + Args: + hash_indices: Hash indices of shape ``[T, num_heads]`` on GPU. + dtype: Cast embeddings to this dtype after lookup. + If None, uses the embedding's native dtype. + + Returns: + Embedding tensor of shape ``[T, num_heads * embed_dim_per_head]``. + """ + if self.stream is not None: + # Fork: let the engram stream wait for any pending main-stream work + # (e.g. hash computation) before we start the embedding lookup. + self.stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.stream): + embeddings = self.multi_head_embedding(hash_indices) + embeddings = embeddings.flatten(start_dim=-2) + if dtype is not None: + embeddings = embeddings.to(dtype) + self.sync_event.record() + else: + embeddings = self.multi_head_embedding(hash_indices) + embeddings = embeddings.flatten(start_dim=-2) + if dtype is not None: + embeddings = embeddings.to(dtype) + + return embeddings + + def forward( + self, + hidden_states: torch.Tensor, + embeddings: torch.Tensor, + conv_state: Optional[torch.Tensor] = None, + use_cache: bool = False, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """Forward pass of the Engram module. + + Args: + hidden_states: Hidden states of shape ``[T, HC_MULT, D]``. + embeddings: Pre-computed embeddings from ``precompute()``. + conv_state: Optional conv state from a previous decode step + (shape ``[1, C_total, conv_state_size]``). + use_cache: When ``True``, return ``(output, new_conv_state)``. + + Returns: + Output tensor of shape ``[T, HC_MULT, D]`` to be added as + residual, and optionally the updated conv state. + """ + # Fused key/value projection: single GEMM replaces hc_mult + 1 separate GEMMs. + # kv_proj output: [T, (1 + HC) * D] + D = self.config.hidden_size + HC = self.config.hc_mult + kv = self.kv_proj(embeddings) + value_raw, keys = kv.split([D, HC * D], dim=-1) + keys = keys.view(*keys.shape[:-1], HC, D) # [T, HC, D] + + # Vectorised RMSNorm for keys and queries (no per-HC kernel launches). + # rms_norm(x, w) = x / rms(x) * w where rms(x) = sqrt(mean(x^2) + eps) + key_rms = keys.pow(2).mean(dim=-1, keepdim=True).add(self.norm_eps).rsqrt() + normed_keys = keys * key_rms * self.key_norm_weight # [HC, D] broadcasts + + queries = hidden_states # [T, HC, D] + query_rms = queries.pow(2).mean(dim=-1, keepdim=True).add(self.norm_eps).rsqrt() + normed_queries = queries * query_rms * self.query_norm_weight + + # Gating: per-HC dot product between normed keys and queries. + gates = (normed_keys * normed_queries).sum(dim=-1) / math.sqrt(D) # [T, HC] + gates = gates.abs().clamp_min(1e-6).sqrt() * gates.sign() + gates = gates.sigmoid().unsqueeze(-1) # [T, HC, 1] + + value = gates * value_raw.unsqueeze(-2) # [T, HC, D] + + if use_cache: + conv_out, new_conv_state = self.short_conv(value, conv_state=conv_state, use_cache=True) + return value + conv_out, new_conv_state + + output = value + self.short_conv(value) + return output diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/__init__.py index 5cf48dc4fe2c..293e695c452b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/__init__.py @@ -14,8 +14,8 @@ from .routing import (BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, - LoadBalancedMoeRoutingMethod, MiniMaxM2MoeRoutingMethod, - RenormalizeMoeRoutingMethod, + LoadBalancedMoeRoutingMethod, DeepSeekV4MoeRoutingMethod, + MiniMaxM2MoeRoutingMethod, RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, RoutingMethodType, SigmoidRenormMoeRoutingMethod, SparseMixerMoeRoutingMethod, StaticMoeRoutingMethod, @@ -41,6 +41,7 @@ "MoeLoadBalancer", "MoEWeightLoadingMode", "MiniMaxM2MoeRoutingMethod", + "DeepSeekV4MoeRoutingMethod", "RenormalizeMoeRoutingMethod", "SigmoidRenormMoeRoutingMethod", "RenormalizeNaiveMoeRoutingMethod", diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index bafe8d5be4d8..0b28df3156ea 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -174,6 +174,8 @@ def __init__( layer_idx=layer_idx, # ConfigurableMoE needs correct layer_idx for EPLB initialization **kwargs, ) + if override_quant_config is not None: + self.quant_config = override_quant_config # Store model_config and aux_stream_dict for later use (e.g., backend setter) self.model_config = model_config diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py index 8abe42efa267..bb19ee29e1e9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py @@ -852,6 +852,7 @@ def forward_chunk( self, x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor] = None, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -861,7 +862,7 @@ def forward_chunk( # This forward_chunk method is a reference implementation of the legacy path. # Apply routing token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits) + router_logits, input_ids) assert token_selected_experts.shape[ 1] == self.routing_method.experts_per_token assert token_selected_experts.shape == token_final_scales.shape diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 85fdac23361d..93eb1b931955 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -1139,6 +1139,7 @@ def forward_chunk( self, x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor], output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -1156,7 +1157,7 @@ def forward_chunk( # apply routing token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits) + router_logits, input_ids) assert token_selected_experts.shape[ 1] == self.routing_method.experts_per_token assert token_selected_experts.shape == token_final_scales.shape @@ -1394,6 +1395,7 @@ def forward_impl( x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, *, + input_ids: Optional[torch.IntTensor] = None, do_finalize: bool = True, # used by other MoE backends output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, @@ -1445,7 +1447,8 @@ def forward_impl( outputs = self.forward_chunk( x, router_logits, - output_dtype, + input_ids=input_ids, + output_dtype=output_dtype, all_rank_num_tokens=all_rank_num_tokens_padded, use_dp_padding=use_dp_padding, repeating_info=(is_first_call, is_last_call), @@ -1470,17 +1473,21 @@ def forward_impl( x_list = x.split(chunk_size_list) router_logits_list = router_logits.split(chunk_size_list) + input_ids_list = input_ids.split( + chunk_size_list) if input_ids is not None else [None + ] * num_chunks self.event_dict[EventType.Main].record() with torch.cuda.stream(self.aux_stream): self.event_dict[EventType.Main].wait() - def _forward_chunk(x_, router_logits_, idx): + def _forward_chunk(x_, router_logits_, input_ids_, idx): is_first_call = idx == 0 and self.repeat_idx == 0 is_last_call = idx == num_chunks - 1 and self.repeat_idx == self.repeat_count - 1 return self.forward_chunk( x_, router_logits_, + input_ids=input_ids_, all_rank_num_tokens=all_rank_num_tokens_list[idx] if self.use_dp else None, use_dp_padding=use_dp_padding, @@ -1495,8 +1502,8 @@ def _reducescatter_or_allreduce(x_, idx): outputs_list = [] # Postpone reduce-scatter/all-reduce to the next iteration to achieve better overlap - for idx_chunk, (x, router_logits) in enumerate( - zip(x_list, router_logits_list)): + for idx_chunk, (x, router_logits, input_ids) in enumerate( + zip(x_list, router_logits_list, input_ids_list)): if not (self.alltoall_method_type == AlltoallMethodType.NVLinkOneSided or self.alltoall_method_type @@ -1504,17 +1511,19 @@ def _reducescatter_or_allreduce(x_, idx): if idx_chunk % 2 == 0: with torch.cuda.stream(self.aux_stream): outputs = _forward_chunk(x, router_logits, - idx_chunk) + input_ids, idx_chunk) if idx_chunk > 0: outputs_list[-1] = _reducescatter_or_allreduce( outputs_list[-1], idx_chunk - 1) else: - outputs = _forward_chunk(x, router_logits, idx_chunk) + outputs = _forward_chunk(x, router_logits, input_ids, + idx_chunk) with torch.cuda.stream(self.aux_stream): outputs_list[-1] = _reducescatter_or_allreduce( outputs_list[-1], idx_chunk - 1) else: - outputs = _forward_chunk(x, router_logits, idx_chunk) + outputs = _forward_chunk(x, router_logits, input_ids, + idx_chunk) outputs_list.append(outputs) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index 0b775ecc4edc..6e39e22a8c16 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -1118,6 +1118,7 @@ def forward_chunk( self, x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor] = None, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -1130,7 +1131,7 @@ def forward_chunk( # apply routing token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits) + router_logits, input_ids) assert token_selected_experts.shape[ 1] == self.routing_method.experts_per_token assert token_selected_experts.shape == token_final_scales.shape @@ -1179,6 +1180,7 @@ def forward_impl( x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, *, + input_ids: Optional[torch.IntTensor] = None, do_finalize: bool = True, # used by other MoE backends output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, @@ -1215,7 +1217,8 @@ def forward_impl( outputs = self.forward_chunk( x, router_logits, - output_dtype, + input_ids=input_ids, + output_dtype=output_dtype, all_rank_num_tokens=all_rank_num_tokens_padded, use_dp_padding=use_dp_padding, workspace=workspaces[0]) @@ -1248,15 +1251,19 @@ def forward_impl( x_list = x.split(chunk_size_list) router_logits_list = router_logits.split(chunk_size_list) + input_ids_list = input_ids.split( + chunk_size_list) if input_ids is not None else [None + ] * num_chunks self.event_dict[EventType.Main].record() with torch.cuda.stream(self.aux_stream): self.event_dict[EventType.Main].wait() - def _forward_chunk(x_, router_logits_, idx, workspace): + def _forward_chunk(x_, router_logits_, input_ids_, idx, workspace): return self.forward_chunk( x_, router_logits_, + input_ids=input_ids_, all_rank_num_tokens=all_rank_num_tokens_list[idx] if self.use_dp else None, use_dp_padding=use_dp_padding, @@ -1270,19 +1277,20 @@ def _reducescatter_or_allreduce(x_, idx): outputs_list = [] # Postpone reduce-scatter/all-reduce to the next iteration to achieve better overlap - for idx_chunk, (x, router_logits) in enumerate( - zip(x_list, router_logits_list)): + for idx_chunk, (x, router_logits, input_ids_chunk) in enumerate( + zip(x_list, router_logits_list, input_ids_list)): if idx_chunk % 2 == 0: with torch.cuda.stream(self.aux_stream): - outputs = _forward_chunk(x, router_logits, idx_chunk, + outputs = _forward_chunk(x, router_logits, + input_ids_chunk, idx_chunk, workspace_0) if idx_chunk > 0: outputs_list[-1] = _reducescatter_or_allreduce( outputs_list[-1], idx_chunk - 1) else: - outputs = _forward_chunk(x, router_logits, idx_chunk, - workspace_1) + outputs = _forward_chunk(x, router_logits, input_ids_chunk, + idx_chunk, workspace_1) with torch.cuda.stream(self.aux_stream): outputs_list[-1] = _reducescatter_or_allreduce( outputs_list[-1], idx_chunk - 1) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py index 4c46490b619d..625b01f31f27 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py @@ -366,8 +366,11 @@ def load_expert_weights_to_dst( module.w3_w1_weight.data) module.w2_weight.data = update_weight_stride(module.w2_weight.data) - def apply(self, module: torch.nn.Module, x: torch.Tensor, - router_logits: torch.Tensor) -> torch.Tensor: + def apply(self, + module: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor] = None) -> torch.Tensor: # Fetch all the data needed for the Triton kernel hidden_states = x expert_logits = router_logits @@ -592,8 +595,11 @@ def load_expert_weights_to_dst( module.w3_w1_weight.data) module.w2_weight.data = update_weight_stride(module.w2_weight.data) - def apply(self, module: torch.nn.Module, x: torch.Tensor, - router_logits: torch.Tensor) -> torch.Tensor: + def apply(self, + module: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor] = None) -> torch.Tensor: # Fetch all the data needed for the Triton kernel hidden_states, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor( x, module.fc31_input_dequant) @@ -1242,8 +1248,11 @@ def load_quant_scales(self, module: torch.nn.Module, weights: Dict): module.fc2_input_dequant.data.copy_(max_fc2_input_scale, non_blocking=True) - def apply(self, module: torch.nn.Module, x: torch.Tensor, - router_logits: torch.Tensor) -> torch.Tensor: + def apply(self, + module: torch.nn.Module, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor] = None) -> torch.Tensor: # Fetch all the data needed for the Triton kernel if self.activation_dtype == torch.float8_e4m3fn: if module.fc31_input_dequant is None: @@ -1557,6 +1566,7 @@ def forward_impl( x: torch.Tensor, router_logits: torch.Tensor, *, + input_ids: Optional[torch.IntTensor] = None, do_finalize: bool = True, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -1566,7 +1576,8 @@ def forward_impl( assert use_dp_padding is None or not use_dp_padding, \ "TritonFusedMoE does not support use_dp_padding=True" - hidden_states = self.quant_method.apply(self, x, router_logits) + hidden_states = self.quant_method.apply(self, x, router_logits, + input_ids) final_hidden_states = self.reducescatter_or_allreduce( hidden_states, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index e9dac388d34b..83e0342d90d2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -28,8 +28,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantAlgo -from ...custom_ops.trtllm_gen_custom_ops import \ - fp4_block_scale_fake_output_without_finalize +from ...custom_ops.trtllm_gen_custom_ops import fp4_block_scale_fake_output_without_finalize from ...distributed import allgather from ...expert_statistic import ExpertStatistic from ...model_config import ModelConfig @@ -44,8 +43,13 @@ W4A8MXFP4FP8TRTLLMGenFusedMoEMethod, W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod, W4A8NVFP4FP8TRTLLMGenFusedMoEMethod, W4A16MXFP4TRTLLMGenFusedMoEMethod) # isort: on -from .routing import (BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, - DefaultMoeRoutingMethod, MiniMaxM2MoeRoutingMethod) +from .routing import ( + BaseMoeRoutingMethod, + DeepSeekV3MoeRoutingMethod, + DeepSeekV4MoeRoutingMethod, + DefaultMoeRoutingMethod, + MiniMaxM2MoeRoutingMethod, +) @dataclass @@ -580,7 +584,7 @@ def quantize_input(self, x, post_quant_comm: bool = True): elif self.has_w4a8_mxfp4_mxfp8: x, x_sf = self.op_backend.mxfp8_quantize( x, False, alignment=self.quant_method.input_hidden_alignment) - x_row, x_col = x.shape[0], x.shape[1] + x_row = x.shape[0] elif self.has_deepseek_fp8_block_scales: # For SM100+, fp8_quantize_1x128 returns x_sf with shape (blocked_n, num_tokens), # but moe_a2a_dispatch requires all payloads to have first dim = num_tokens. @@ -625,6 +629,14 @@ def _extract_routing_params(self) -> RoutingParams: topk_group=None, routed_scaling_factor=None, ) + elif isinstance(self.routing_method, DeepSeekV4MoeRoutingMethod): + return RoutingParams( + top_k=self.routing_method.top_k, + routing_bias=self.routing_method.e_score_correction_bias, + n_group=self.routing_method.n_group, + topk_group=self.routing_method.topk_group, + routed_scaling_factor=self.routing_method.routed_scaling_factor, + ) else: return RoutingParams( top_k=self.routing_method.top_k, @@ -909,6 +921,7 @@ def forward_impl( x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, *, + input_ids: Optional[torch.IntTensor] = None, do_finalize: bool = True, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -921,6 +934,7 @@ def forward_impl( run_post_quant_allgather = (self.use_dp and self.parallel_size > 1 and not self.enable_alltoall) post_quant_comm = run_post_quant_allgather or self.enable_alltoall + requires_separated_routing = self.routing_method.requires_separated_routing x_sf = None token_selected_experts = None @@ -931,11 +945,11 @@ def forward_impl( is_first_call = self.repeat_idx == 0 is_last_call = self.repeat_idx == self.repeat_count - 1 - if post_quant_comm: + if post_quant_comm or requires_separated_routing: self._load_balancer_start_wait_gpu_stage(is_first_call) token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits) + router_logits, input_ids) token_selected_experts = token_selected_experts.to(torch.int32) if token_final_scales is not None: token_final_scales = token_final_scales.to(torch.bfloat16) @@ -963,6 +977,7 @@ def forward_impl( # Use routed slots for subsequent processing token_selected_experts = token_selected_slots + if post_quant_comm: x, x_sf = self.quantize_input(x) if self.enable_alltoall: @@ -1103,7 +1118,8 @@ def forward_impl( # Call the extracted run_moe interface # Determine router_logits based on post_quant_comm - router_logits_arg = None if post_quant_comm else router_logits + router_logits_arg = None if ( + post_quant_comm or requires_separated_routing) else router_logits final_hidden_states = self.run_moe( x=x, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py index db0afbac1b3f..e93d67d20398 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py @@ -535,6 +535,7 @@ def forward( self, x: torch.Tensor, router_logits: torch.Tensor, + input_ids: Optional[torch.IntTensor] = None, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, **kwargs, @@ -543,7 +544,7 @@ def forward( x = x.view(-1, self.hidden_size) token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits) + router_logits, input_ids) if self.use_dp and self.parallel_size > 1: x, token_selected_experts, token_final_scales = allgather( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py index 1edbf54c2914..0f207ceff0e7 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py @@ -385,6 +385,7 @@ def forward_chunk( x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, use_all_to_all: bool, + input_ids: Optional[torch.IntTensor] = None, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, @@ -407,7 +408,7 @@ def forward_chunk( weight_dtype = self.w3_w1_weight.dtype token_selected_experts, token_final_scales = self.routing_method.apply( - router_logits) + router_logits, input_ids) assert token_selected_experts.shape[ 1] == self.routing_method.experts_per_token @@ -699,6 +700,7 @@ def forward_impl( x: Union[torch.Tensor, Fp4QuantizedTensor], router_logits: torch.Tensor, *, + input_ids: Optional[torch.IntTensor] = None, do_finalize: bool = True, output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, @@ -728,7 +730,8 @@ def forward_impl( x, router_logits, use_all_to_all, - output_dtype, + input_ids=input_ids, + output_dtype=output_dtype, all_rank_num_tokens=all_rank_num_tokens_padded, use_dp_padding=use_dp_padding, repeating_info=(is_first_call, is_last_call), @@ -768,6 +771,9 @@ def split_chunk(split_token_num: int, split_num_chunks: int): x_list = x.split(chunk_size_list) router_logits_list = router_logits.split(chunk_size_list) + input_ids_list = input_ids.split( + chunk_size_list) if input_ids is not None else [None + ] * num_chunks if not use_all_to_all: self.event_dict[EventType.Main].record() @@ -776,8 +782,8 @@ def split_chunk(split_token_num: int, split_num_chunks: int): outputs_list = [] # Postpone reduce-scatter/all-reduce to the next iteration to achieve better overlap - for idx_chunk, (x, router_logits) in enumerate( - zip(x_list, router_logits_list)): + for idx_chunk, (x, router_logits, input_ids_chunk) in enumerate( + zip(x_list, router_logits_list, input_ids_list)): is_first_call = idx_chunk == 0 and self.repeat_idx == 0 is_last_call = idx_chunk == num_chunks - 1 and self.repeat_idx == self.repeat_count - 1 if not use_all_to_all: @@ -787,6 +793,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): x, router_logits, use_all_to_all, + input_ids=input_ids_chunk, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk], use_dp_padding=use_dp_padding, @@ -804,6 +811,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): x, router_logits, use_all_to_all, + input_ids=input_ids_chunk, all_rank_num_tokens=all_rank_num_tokens_list[ idx_chunk], use_dp_padding=use_dp_padding, @@ -821,6 +829,7 @@ def split_chunk(split_token_num: int, split_num_chunks: int): x, router_logits, use_all_to_all, + input_ids=input_ids_chunk, all_rank_num_tokens=all_rank_num_tokens_list[idx_chunk], repeating_info=(is_first_call, is_last_call), alltoall_result_do_sum=alltoall_result_do_sum) diff --git a/tensorrt_llm/_torch/modules/fused_moe/routing.py b/tensorrt_llm/_torch/modules/fused_moe/routing.py index 56724d44346f..eea2b2c49e39 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/routing.py +++ b/tensorrt_llm/_torch/modules/fused_moe/routing.py @@ -248,13 +248,19 @@ class RoutingMethodType(IntEnum): MiniMax2 = 5, # SigmoidRenorm: Sigmoid -> TopK -> Renormalize SigmoidRenorm = 6, + # DeepSeek-V4 + DeepSeekV4 = 7, # Unspecified - Unspecified = 7, + Unspecified = 8, class BaseMoeRoutingMethod(nn.Module): - def apply(self, _router_logits) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + _router_logits, + _input_ids, + ) -> tuple[torch.Tensor, torch.Tensor]: """ Applies the routing method to the router logits. Router logits are usually the output of the router Linear layer, but can be any type for more complex routing methods. @@ -273,6 +279,16 @@ def get_experts_per_token(self) -> int: def experts_per_token(self) -> int: return self.get_experts_per_token() + @property + def requires_separated_routing(self) -> bool: + """Whether routing must be computed externally (in Python) rather than + fused inside the C++ MoE kernel. + + Override to ``True`` when the routing algorithm (e.g. sqrtsoftplus + scoring) is not natively supported by any C++ kernel backend. + """ + return False + @property def routing_method_type(self) -> RoutingMethodType: return RoutingMethodType.Unspecified @@ -290,16 +306,21 @@ def __init__(self, self.output_dtype = output_dtype def apply_pytorch( - self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: topk_values, topk_indices = torch.topk(torch.nn.functional.softmax( router_logits.to(self.output_dtype), dim=-1), k=self.top_k, dim=-1) return topk_indices.to(torch.int32), topk_values - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: num_experts = router_logits.shape[-1] if self.force_enable_pytorch_op or num_experts > 512 or self.top_k > 16: return self.apply_pytorch(router_logits) @@ -460,7 +481,11 @@ def __init__( assert callable(callable_e_score_correction_bias) self.callable_e_score_correction_bias = callable_e_score_correction_bias - def apply(self, logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: return self.routing_impl.apply(logits, self.e_score_correction_bias) @property @@ -476,6 +501,87 @@ def routing_method_type(self): return RoutingMethodType.DeepSeekV3 +class DeepSeekV4MoeRoutingMethod(BaseMoeRoutingMethod): + + def __init__( + self, + top_k: int, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + callable_e_score_correction_bias: Callable[[], torch.Tensor], + callable_tid2eid: Callable[[], torch.Tensor], + is_hashed: bool = True, + ): + super().__init__() + self._top_k = top_k + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.is_hashed = is_hashed + + # Pass a callable to fetch the tensor from DeepseekV4Gate at runtime, ensuring it is on the correct device + assert callable(callable_e_score_correction_bias) + assert callable(callable_tid2eid) + self.callable_e_score_correction_bias = callable_e_score_correction_bias + self.callable_tid2eid = callable_tid2eid + + def apply( + self, + logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: + # gate_forward kernel requires float32 input scores + logits = logits.to(torch.float32) + m = logits.shape[0] + out_weights = torch.empty(m, + self._top_k, + dtype=torch.float32, + device=logits.device) + out_indices = torch.empty(m, + self._top_k, + dtype=torch.int32, + device=logits.device) + if self.is_hashed: + assert input_ids is not None + torch.ops.trtllm.gate_forward( + logits, self.callable_e_score_correction_bias(), input_ids, + self.callable_tid2eid(), out_weights, out_indices, self._top_k, + self.routed_scaling_factor, True) + else: + input_ids_tensor = torch.empty(0, + dtype=torch.int32, + device=logits.device, + requires_grad=False) + tid2eid_tensor = torch.empty(0, + dtype=torch.int32, + device=logits.device, + requires_grad=False) + torch.ops.trtllm.gate_forward( + logits, self.callable_e_score_correction_bias(), + input_ids_tensor, tid2eid_tensor, out_weights, out_indices, + self._top_k, self.routed_scaling_factor, False) + return out_indices, out_weights + + @property + def e_score_correction_bias(self) -> torch.Tensor: + return self.callable_e_score_correction_bias() + + @property + def top_k(self): + return self._top_k + + @property + def requires_separated_routing(self) -> bool: + # C++ MoE kernels don't support DeepSeek-V4's sqrtsoftplus scoring natively. + return True + + @property + def routing_method_type(self): + # Return DeepSeekV3 because C++ MoE kernels don't recognize DeepSeek-V4 routing type + return RoutingMethodType.DeepSeekV3 + + class MiniMaxM2MoeRoutingMethod(BaseMoeRoutingMethod): def __init__( @@ -506,8 +612,11 @@ def get_scores(logits, e_score_correction_bias): return scores, scores_with_bias - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: scores, scores_with_bias = self.get_scores(router_logits, self.e_score_correction_bias) _, topk_idx = torch.topk(scores_with_bias, @@ -542,8 +651,11 @@ def __init__( self.renormalize = renormalize self.output_dtype = output_dtype - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: scores = torch.sigmoid(router_logits) topk_weights, topk_idx = torch.topk(scores, k=self.top_k, @@ -573,16 +685,21 @@ def __init__( self.output_dtype = output_dtype def apply_pytorch( - self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: topk_values, topk_indices = torch.topk(router_logits, k=self.top_k, dim=-1) return topk_indices.to(torch.int32), torch.nn.functional.softmax( topk_values.to(self.output_dtype), dim=-1) - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: num_experts = router_logits.shape[-1] if self.force_enable_pytorch_op or num_experts > 512 or self.top_k > 16: return self.apply_pytorch(router_logits) @@ -602,8 +719,11 @@ def __init__(self, top_k: int, output_dtype: torch.dtype = torch.float32): self.top_k = top_k self.output_dtype = output_dtype - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: topk_values, topk_indices = torch.topk(router_logits, k=self.top_k, dim=-1) @@ -643,8 +763,11 @@ def __init__(self, top_k: int, eps: float): self.top_k = top_k self.eps = eps - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: router_logits = router_logits.float() topk_values = torch.empty(router_logits.shape[0], self.top_k, @@ -696,8 +819,11 @@ def __init__(self, self.routing_tensor = routing_tensor self.routing_scales = routing_scales - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: return self.routing_tensor, self.routing_scales def get_experts_per_token(self): @@ -710,8 +836,11 @@ def __init__(self, top_k: int): super().__init__() self.top_k = top_k - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: balanced_values = torch.ones(router_logits.shape[0], self.top_k, device=router_logits.device, @@ -742,8 +871,11 @@ def __init__(self, top_k: int, output_dtype: torch.dtype = torch.float32): self.top_k = top_k self.output_dtype = output_dtype - def apply(self, - router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def apply( + self, + router_logits: torch.Tensor, + input_ids: Optional[torch.Tensor] = None + ) -> tuple[torch.Tensor, torch.Tensor]: #x = topk(softmax()); x /= x.sum() is mathematically equivalent to softmax(topk) topk_indices, topk_values = self.apply_pytorch(router_logits) return topk_indices.to(torch.int32), topk_values.to(self.output_dtype) @@ -771,6 +903,8 @@ def routing_method_type(self) -> RoutingMethodType: MiniMaxM2MoeRoutingMethod, RoutingMethodType.SigmoidRenorm: SigmoidRenormMoeRoutingMethod, + RoutingMethodType.DeepSeekV4: + DeepSeekV4MoeRoutingMethod, } diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 01293a83faf0..32965bcf3263 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -34,6 +34,7 @@ def __init__( disable_deep_gemm: bool = False, use_custom_cublas_mm: bool = False, is_shared_expert: bool = False, + swiglu_limit: Optional[float] = None, ): super().__init__() @@ -42,6 +43,7 @@ def __init__( self.intermediate_size = intermediate_size self.activation = activation self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm + self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None config = config or ModelConfig() self.mapping = config.mapping @@ -139,13 +141,14 @@ def _apply_activation(self, x, *, has_lora: bool = False): logger.warning( f"GatedMLP._apply_activation: LoRA path active; forcing non-FP8 activation dtype bf16/fp16, layer_idx={self.layer_idx}" ) - return swiglu(x) + return swiglu(x, swiglu_limit=self.swiglu_limit) else: return swiglu(x, quant_scale=self.down_proj.input_scale, - quant_type=torch.float8_e4m3fn) + quant_type=torch.float8_e4m3fn, + swiglu_limit=self.swiglu_limit) else: - return swiglu(x) + return swiglu(x, swiglu_limit=self.swiglu_limit) elif callable(self.activation): return self.activation(x) elif self.activation is None: diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index aae3f3a65ab2..97bde1b70a63 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1116,7 +1116,6 @@ def apply(self, module: Linear, input: torch.Tensor, input, module.weight, module.weight_scale, - disable_ue8m0_cast=True, ) elif get_sm_version() == 120: act_input_fp8, act_input_sf = per_token_quant_and_transform(input) diff --git a/tensorrt_llm/_torch/modules/mhc/__init__.py b/tensorrt_llm/_torch/modules/mhc/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/modules/mhc/hyper_connection.py b/tensorrt_llm/_torch/modules/mhc/hyper_connection.py new file mode 100644 index 000000000000..051eeaa95a48 --- /dev/null +++ b/tensorrt_llm/_torch/modules/mhc/hyper_connection.py @@ -0,0 +1,280 @@ +# Multi-Head Hyper-Connection (mHC) module +# Based on: "Hyper-Connections" (https://arxiv.org/abs/2409.19606) +from dataclasses import dataclass +from typing import Optional + +import torch +from torch import nn + + +@dataclass +class HCState: + """Inter-layer mHC pipeline state. + + Two modes, distinguished by whether ``post_mix`` is populated: + + - **resolved** (``post_mix is None``): ``residual`` is the fully post-mapped + activation for the next layer's ``pre_mapping``. This is what the prior + layer returns when fused_hc is disabled (or engram mutated the residual). + The next layer just runs ``pre_mapping(residual)``. + - **deferred** (``post_mix is not None``): the prior layer deferred its + ``post_mapping``; the 4 tensors carry the inputs needed for the next + layer to absorb it via ``fused_hc``. + + Only ``modeling_deepseekv4.py`` depends on this shape — the kernel-level + ``mHC.fused_hc`` still returns a 4-tuple so low-level callers (tests, + benchmarks) stay unchanged. + """ + + residual: torch.Tensor + post_mix: Optional[torch.Tensor] = None + comb_mix: Optional[torch.Tensor] = None + x_prev: Optional[torch.Tensor] = None + + @property + def is_deferred(self) -> bool: + return self.post_mix is not None + + @classmethod + def resolved(cls, residual: torch.Tensor) -> "HCState": + return cls(residual=residual) + + @classmethod + def deferred( + cls, + residual: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + x_prev: torch.Tensor, + ) -> "HCState": + return cls(residual=residual, post_mix=post_mix, comb_mix=comb_mix, x_prev=x_prev) + +try: + from tensorrt_llm._torch.modules.mhc.mhc_cuda import mhc_hc_head_cuda, mhc_post_mapping_cuda + from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( + mhc_fused_hc as mhc_fused_hc_cuda, + ) + from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( + mhc_pre_mapping_fused as mhc_pre_mapping_fused_cuda, + ) + + _cuda_available = True +except Exception as _e: + _cuda_available = False + mhc_hc_head_cuda = None + mhc_post_mapping_cuda = None + mhc_pre_mapping_fused_cuda = None + mhc_fused_hc_cuda = None + + +class mHC(nn.Module): + def __init__( + self, + mult: int, + hidden_size: int, + sinkhorn_iters: int, + dtype: Optional[torch.dtype] = None, + eps: float = 1e-6, + norm_eps: float = 1e-6, + sinkhorn_eps: float = 1e-6, + post_mult_value: float = 1.0, + n_splits: int = 1, + ): + super().__init__() + self.mult = mult + self.hidden_size = hidden_size + self.sinkhorn_iters = sinkhorn_iters + self.dtype = dtype + self.eps = eps + self.norm_eps = norm_eps + self.sinkhorn_eps = sinkhorn_eps + self.post_mult_value = post_mult_value + self.n_splits = n_splits + self.mix_hc = (2 + self.mult) * self.mult + self.hc_dim = self.mult * self.hidden_size + + # Parameters + self.fn = nn.Parameter( + torch.empty((self.mix_hc, self.hc_dim), dtype=torch.float32), requires_grad=False + ) + self.base = nn.Parameter( + torch.empty((self.mix_hc,), dtype=torch.float32), requires_grad=False + ) + self.scale = nn.Parameter(torch.empty((3,), dtype=torch.float32), requires_grad=False) + + def pre_mapping(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # x: [b,s,hc,d], hc_fn: [mix_hc,hc*d], hc_scale: [3], hc_base: [mix_hc], y: [b,s,hc,d] + if not _cuda_available: + raise RuntimeError( + "Raw CUDA backend is unavailable. " + "Ensure torch.utils.cpp_extension and CUDA toolkit are installed." + ) + assert x.dtype == torch.bfloat16 + assert self.mult == x.shape[-2] + assert self.hidden_size == x.shape[-1] + outer_shape = x.shape[:-2] + residual_flat = x.view(-1, self.mult, self.hidden_size) + num_tokens = residual_flat.shape[0] + + post_mix, comb_mix, layer_input = mhc_pre_mapping_fused_cuda( + residual_flat.view(num_tokens, self.hc_dim), + self.fn.contiguous(), + residual_flat, + self.mult, + self.scale, + self.base, + self.hidden_size, + self.norm_eps, + self.eps, + self.sinkhorn_eps, + self.post_mult_value, + self.sinkhorn_iters, + ) + + post_mix = post_mix.view(*outer_shape, self.mult, 1) + comb_mix = comb_mix.view(*outer_shape, self.mult, self.mult) + layer_input = layer_input.view(*outer_shape, self.hidden_size) + return post_mix, comb_mix, layer_input + + def fused_hc( + self, + x_prev: torch.Tensor, + residual_prev: torch.Tensor, + post_mix_prev: torch.Tensor, + comb_mix_prev: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused post_mapping(from previous mHC) + pre_mapping(from self). + + This is the boundary op between two mHC-wrapped blocks. It consumes the + output of the previous block (``x_prev``) plus the previous block's + residual and mix matrices, then runs the current block's pre_mapping + using ``self``'s parameters. Semantically identical to + + residual_cur = prev_mHC.post_mapping(x_prev, residual_prev, ...) + post_mix, comb_mix, layer_input = self.pre_mapping(residual_cur) + + but exposed as one call so the model forward can say + ``state = mHC_next.fused_hc(...)`` at every layer boundary. + + Args: + x_prev: [..., hidden] bf16 (attn / MoE output of prev block) + residual_prev: [..., mult, hidden] bf16 + post_mix_prev: [..., mult] or [..., mult, 1] fp32 + comb_mix_prev: [..., mult, mult] fp32 + + Returns: + residual_cur: [..., mult, hidden] bf16 (new residual for next post_mapping) + post_mix_cur: [..., mult, 1] fp32 + comb_mix_cur: [..., mult, mult] fp32 + layer_input_cur: [..., hidden] bf16 + """ + if not _cuda_available: + raise RuntimeError( + "Raw CUDA backend is unavailable. " + "Ensure torch.utils.cpp_extension and CUDA toolkit are installed." + ) + assert x_prev.dtype == torch.bfloat16 + assert residual_prev.dtype == torch.bfloat16 + n = self.mult + hidden = self.hidden_size + outer_shape = residual_prev.shape[:-2] + + residual_prev_flat = residual_prev.reshape(-1, n, hidden).contiguous() + B = residual_prev_flat.shape[0] + x_prev_flat = x_prev.reshape(B, hidden).contiguous() + post_mix_prev_flat = post_mix_prev.reshape(B, n) + comb_mix_prev_flat = comb_mix_prev.reshape(B, n, n) + + residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = mhc_fused_hc_cuda( + x_prev_flat, + residual_prev_flat, + post_mix_prev_flat, + comb_mix_prev_flat, + self.fn.contiguous(), + self.scale, + self.base, + n, + hidden, + self.norm_eps, + self.eps, + self.sinkhorn_eps, + self.post_mult_value, + self.sinkhorn_iters, + ) + + residual_cur = residual_cur.view(*outer_shape, n, hidden) + post_mix_cur = post_mix_cur.view(*outer_shape, n, 1) + comb_mix_cur = comb_mix_cur.view(*outer_shape, n, n) + layer_input_cur = layer_input_cur.view(*outer_shape, hidden) + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + + def post_mapping( + self, + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + ) -> torch.Tensor: + if not _cuda_available: + raise RuntimeError( + "Raw CUDA backend is unavailable. " + "Ensure torch.utils.cpp_extension and CUDA toolkit are installed." + ) + outer_shape = residual.shape[:-2] + n = self.mult + hidden = residual.shape[-1] + residual_flat = residual.view(-1, n, hidden) + B = residual_flat.shape[0] + + out = mhc_post_mapping_cuda( + residual_flat, + x.reshape(B, hidden), + post_layer_mix.view(B, n), + comb_res_mix.view(B, n, n), + n, + ) + return out.view(*outer_shape, n, hidden) + + +class HCHead(nn.Module): + def __init__( + self, + mult: int, + hidden_size: int, + eps: float = 1e-6, + norm_eps: float = 1e-6, + ): + super().__init__() + self.mult = mult + self.hidden_size = hidden_size + self.eps = eps + self.norm_eps = norm_eps + self.fn = nn.Parameter( + torch.empty((self.mult, self.mult * self.hidden_size), dtype=torch.float32), + requires_grad=False, + ) + self.base = nn.Parameter( + torch.empty((self.mult,), dtype=torch.float32), requires_grad=False + ) + self.scale = nn.Parameter(torch.empty((1,), dtype=torch.float32), requires_grad=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if not _cuda_available: + raise RuntimeError("CUDA MHC kernels not available") + dtype = x.dtype + x_bf16 = x.to(torch.bfloat16).contiguous() + y = mhc_hc_head_cuda( + x_bf16, + self.fn, + self.scale, + self.base, + self.mult, + self.hidden_size, + norm_eps=self.norm_eps, + eps=self.eps, + ) + return y.to(dtype) + + def skip_forward(self, x: torch.Tensor) -> torch.Tensor: + """Skip HCHead computation for pipeline parallelism on non-last ranks.""" + return x diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py new file mode 100644 index 000000000000..125c5d6adf4d --- /dev/null +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -0,0 +1,1085 @@ +"""MHC CUDA kernels — torch.ops.trtllm interface. + +Kernels (cpp/tensorrt_llm/kernels/mhcKernels/): + - mhc_big_fuse: fused RMS + sigmoid + Sinkhorn + pre_apply_mix + (NUM_SPLITS=1: normal, =16: with split-K reduction) + - mhc_gemm_sqrsum_fma: FP32 FMA GEMM + sqrsum (fused, inline PTX) + - mhc_hc_head_apply: RMS norm + sigmoid + weighted sum + - mhc_post_mapping: out = post * x + comb.T @ residual +DeepGEMM wrapper: + - gemm_rms_dg: TF32 GEMM + sqrsum via DeepGEMM (optional split-K) + +Backend selection for pre_mapping is handled by the autotuner, which profiles +all available backends (FMA, DeepGEMM split-K, DeepGEMM no-split) at warmup. +Falls back to FMA when DeepGEMM is unavailable or autotuner cache misses. +""" + +from functools import lru_cache +from typing import Any, List + +import torch + +from tensorrt_llm._torch.autotuner import ( + AutoTuner, + ConstraintSpec, + DynamicTensorSpec, + OptimizationProfile, + TunableRunner, + TuningConfig, +) +from tensorrt_llm._utils import get_sm_version + + +@lru_cache(maxsize=1) +def _fused_hc_mma_supported() -> bool: + """tcgen05 TF32 MMA paths (Path B / Path D, "fused_*_mma") require SM100+. + + On pre-SM100 GPUs only the FMA paths (Path E / Path F, "fused_*_fma") are + safe to run. Called lazily so the module can still be imported on a host + without a CUDA device (e.g. CPU-only lint / typecheck). + """ + try: + return get_sm_version() >= 100 + except Exception: + return False + +_DG_NUM_SPLITS = 16 + + +# --------------------------------------------------------------------------- +# Python API — low-level (kernel-level interfaces) +# --------------------------------------------------------------------------- + + +def mhc_big_fuse_cuda( + y_acc: torch.Tensor, + r_acc: torch.Tensor, + residual: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + layer_input: torch.Tensor, + M: int, + K: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + num_splits: int = 1, + block_size: int = 0, +): + torch.ops.trtllm.mhc_big_fuse( + y_acc, + r_acc, + residual, + hc_scale, + hc_base, + post_mix, + comb_mix, + layer_input, + M, + K, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + num_splits, + block_size, + ) + + +_dg_fn_cache = None +_dg_available = None + + +def _get_dg_fn(): + """Import tf32_hc_prenorm_gemm (standalone deep_gemm or bundled). + + Returns the function, or None if DeepGEMM is unavailable. + """ + global _dg_fn_cache, _dg_available + if _dg_available is not None: + return _dg_fn_cache + try: + from deep_gemm import tf32_hc_prenorm_gemm + + _dg_fn_cache = tf32_hc_prenorm_gemm + except ImportError: + try: + from tensorrt_llm.deep_gemm import tf32_hc_prenorm_gemm + + _dg_fn_cache = tf32_hc_prenorm_gemm + except ImportError: + _dg_fn_cache = None + _dg_available = _dg_fn_cache is not None + return _dg_fn_cache + + +def mhc_gemm_rms_dg_cuda( + x: torch.Tensor, + w_nk: torch.Tensor, + M: int, + N: int, + K: int, + num_splits: int = 16, +): + """DeepGEMM TF32 GEMM + sqrsum with optional split-K. + + Args: + x: [M, K] bfloat16 input + w_nk: [N, K] float32 weight (K-major, as required by DeepGEMM) + num_splits: 1 for no split-K, >1 for split-K + + Returns (y_acc, r_acc, num_splits): + num_splits == 1: y_acc [M, N] fp32, r_acc [M] fp32 + num_splits > 1: y_acc [num_splits, M, N] fp32, r_acc [num_splits, M] fp32 + """ + dg_fn = _get_dg_fn() + assert dg_fn is not None, "DeepGEMM is not available" + x = x.contiguous() + w_nk = w_nk.contiguous() + + if num_splits <= 1: + y_acc = torch.empty((M, N), dtype=torch.float32, device=x.device) + r_acc = torch.empty((M,), dtype=torch.float32, device=x.device) + dg_fn(x, w_nk, y_acc, r_acc) + return y_acc, r_acc, 1 + else: + y_acc = torch.empty((num_splits, M, N), dtype=torch.float32, device=x.device) + r_acc = torch.empty((num_splits, M), dtype=torch.float32, device=x.device) + dg_fn(x, w_nk, y_acc, r_acc, num_splits=num_splits) + return y_acc, r_acc, num_splits + + +def mhc_gemm_rms_fma_cuda( + x: torch.Tensor, + w: torch.Tensor | None, + M: int, + N: int, + K: int, + w_t: torch.Tensor | None = None, + tile_n: int = 0, + tile_m: int = 0, +): + """Split-N FP32 FMA fused GEMM + sqrsum on CUDA cores (no tensor cores). + + Args: + tile_n: FMA N-tile size. 0 = auto-select. + Valid: {1,2,3,4,6,8,12,24}. + tile_m: FMA M-tile size. 0 = auto-select. + Valid: {1, 2} (2 only with tile_n=8). + + Returns (y_acc [M, N] fp32, r_acc [M] fp32). + """ + x = x.contiguous() + if w_t is None: + w_t = w.t().contiguous() + + y_acc = torch.empty((M, N), dtype=torch.float32, device=x.device) + r_acc = torch.empty((M,), dtype=torch.float32, device=x.device) + + torch.ops.trtllm.mhc_gemm_sqrsum_fma( + x, + w_t, + y_acc, + r_acc, + M, + N, + K, + tile_n, + tile_m, + ) + + return y_acc, r_acc + + +# --------------------------------------------------------------------------- +# Autotuner runner for pre_mapping backend selection +# --------------------------------------------------------------------------- + + +def _mhc_gen_tuning_buckets(x: int): + """Generate M-dimension tuning buckets for MHC pre_mapping. + + Buckets: 1, 2, 4, 8, 16, 32, 64, 128, 256, 384, 512, 768, 1024, ... + Small M uses powers-of-2 for fine granularity; large M uses 128 steps. + """ + buckets = (1, 2, 4, 8, 16, 32, 64, 128) + if x >= 128: + x = min(x, 8192) + x = max(x, 1024) + buckets += tuple(range(256, x + 1, 128)) + return buckets + + +def _mhc_map_to_tuning_bucket(x: int) -> int: + """Map an inference-time M to the nearest tuning bucket (round up).""" + if x <= 128: + v = 1 + while v < x: + v *= 2 + return min(v, 128) + return ((x + 127) // 128) * 128 + + +_FMA_TILE_N_OPTIONS = (1, 2, 3, 4, 6, 8, 12, 24) +_BIGFUSE_BLOCK_SIZE_OPTIONS = (128, 256, 512) + +# Tactic is a tuple: (backend, tile_n, tile_m, bigfuse_bs) +# backend: "fma", "dg_splitk", "dg_nosplit" +# tile_n: FMA N-tile (only used for "fma" backend, 0 otherwise) +# tile_m: FMA M-tile (only used for "fma" backend, 0 otherwise) +# bigfuse_bs: BigFuse block size {128, 256, 512} +_FALLBACK_TACTIC = ("fma", 0, 0, 0) + + +class MhcPreMappingRunner(TunableRunner): + """Profiles the full MHC pre_mapping pipeline (GEMM+sqrsum + big_fuse). + + Tactic format: (backend, tile_n, tile_m, bigfuse_bs) + backend: "fma" | "dg_splitk" | "dg_nosplit" + tile_n: FMA N-tile size (fma backend only) + tile_m: FMA M-tile size (fma backend only; 2 requires tile_n=8) + bigfuse_bs: BigFuse BLOCK_SIZE {128, 256, 512} + + Fallback tactic (-1): uses auto-select for all parameters. + """ + + tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=0, + dim_idx=0, + gen_tuning_buckets=_mhc_gen_tuning_buckets, + map_to_tuning_buckets=_mhc_map_to_tuning_bucket, + ), + ), + # residual (input[2]) dim 0 = M, same as x (input[0]) dim 0 + constraint_specs=( + ConstraintSpec( + input_idx=2, + dim_idx=0, + infer_shape=lambda shapes: shapes[0][0], + ), + ), + ) + + def __init__( + self, + n: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + ): + self.n = n + self.hidden_size = hidden_size + self.rms_eps = rms_eps + self.hc_pre_eps = hc_pre_eps + self.hc_sinkhorn_eps = hc_sinkhorn_eps + self.hc_post_mult_value = hc_post_mult_value + self.sinkhorn_repeat = sinkhorn_repeat + + def unique_id(self): + return (self.n, self.hidden_size) + + def get_valid_tactics( + self, inputs: List[torch.Tensor], profile: OptimizationProfile, **kwargs + ) -> List[Any]: + N = inputs[1].shape[0] # w_t is [N, K] + valid_tile_n = tuple(tn for tn in _FMA_TILE_N_OPTIONS if N % tn == 0) + + tactics = [] + for tn in valid_tile_n: + for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: + tactics.append(("fma", tn, 1, bs)) + if 8 in valid_tile_n: + for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: + tactics.append(("fma", 8, 2, bs)) + + if _get_dg_fn() is not None: + for bs in _BIGFUSE_BLOCK_SIZE_OPTIONS: + tactics.append(("dg_splitk", 0, 0, bs)) + tactics.append(("dg_nosplit", 0, 0, bs)) + + return tactics + + def forward(self, inputs: List[torch.Tensor], *, tactic: Any = -1, **kwargs) -> Any: + x, w_t, residual, hc_scale, hc_base = inputs + + residual = residual.contiguous() + hc_scale = hc_scale.to(torch.float32).contiguous() + hc_base = hc_base.to(torch.float32).contiguous() + + M, K = x.shape + N = w_t.shape[0] + n = self.n + n2 = n * n + + if tactic == -1: + tactic = _FALLBACK_TACTIC + backend, tile_n, tile_m, bigfuse_bs = tactic + + num_splits = 1 + if backend == "dg_splitk": + y_acc, r_acc, num_splits = mhc_gemm_rms_dg_cuda( + x, w_t, M, N, K, num_splits=_DG_NUM_SPLITS + ) + elif backend == "dg_nosplit": + y_acc, r_acc, num_splits = mhc_gemm_rms_dg_cuda(x, w_t, M, N, K, num_splits=1) + else: + y_acc, r_acc = mhc_gemm_rms_fma_cuda( + x, None, M, N, K, w_t=w_t, tile_n=tile_n, tile_m=tile_m + ) + + residual_3d = residual.view(M, n, self.hidden_size) + + post_mix = torch.empty((M, n), dtype=torch.float32, device=x.device) + comb_mix = torch.empty((M, n2), dtype=torch.float32, device=x.device) + layer_input = torch.empty((M, self.hidden_size), dtype=torch.bfloat16, device=x.device) + + mhc_big_fuse_cuda( + y_acc.contiguous(), + r_acc.contiguous(), + residual_3d.contiguous(), + hc_scale, + hc_base, + post_mix, + comb_mix, + layer_input, + M, + K, + self.hidden_size, + self.rms_eps, + self.hc_pre_eps, + self.hc_sinkhorn_eps, + self.hc_post_mult_value, + self.sinkhorn_repeat, + num_splits=num_splits, + block_size=bigfuse_bs, + ) + + return post_mix, comb_mix, layer_input + + +# --------------------------------------------------------------------------- +# Python API — high-level (drop-in for mhc_pre_mapping_fused) +# --------------------------------------------------------------------------- + + +# Process-wide pre_mapping runner cache keyed on mHC config. +_pre_mapping_runner_cache: dict = {} + + +def _get_pre_mapping_runner( + n: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +) -> "MhcPreMappingRunner": + key = ( + n, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + runner = _pre_mapping_runner_cache.get(key) + if runner is None: + runner = MhcPreMappingRunner( + n=n, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + ) + _pre_mapping_runner_cache[key] = runner + return runner + + +def mhc_pre_mapping_fused( + x: torch.Tensor, + w_t: torch.Tensor, + residual: torch.Tensor, + n: int, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +): + """Full pre-mapping pipeline: GEMM+sqrsum -> big_fuse. + + Backend selection is handled by the autotuner at warmup. + Falls back to FMA when DeepGEMM is unavailable or cache misses. + + Args: + w_t: [N, K] float32 weight (row-major, pre-transposed). + """ + runner = _get_pre_mapping_runner( + n=n, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + ) + + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( + "trtllm::mhc_pre_mapping", + [runner], + MhcPreMappingRunner.tuning_config, + [x, w_t, residual, hc_scale, hc_base], + ) + + return runner( + inputs=[x, w_t, residual, hc_scale, hc_base], + tactic=best_tactic, + ) + + +# --------------------------------------------------------------------------- +# Python API — post_mapping +# --------------------------------------------------------------------------- + + +def mhc_post_mapping_cuda( + residual: torch.Tensor, + x: torch.Tensor, + post_mix: torch.Tensor, + comb_mix: torch.Tensor, + n: int, +) -> torch.Tensor: + """Post-mapping: out = post * x + comb.T @ residual. + + Args: + residual: [B, n, hidden_size] bf16 + x: [B, hidden_size] bf16 + post_mix: [B, n] fp32 + comb_mix: [B, n, n] fp32 + n: number of hyper-connection heads + + Returns: [B, n, hidden_size] bf16. + """ + residual = residual.contiguous() + x = x.contiguous() + post_mix = post_mix.to(torch.float32).contiguous() + comb_mix = comb_mix.to(torch.float32).contiguous() + + B = residual.shape[0] + hidden_size = residual.shape[2] + + out = torch.empty((B, n, hidden_size), dtype=torch.bfloat16, device=x.device) + + torch.ops.trtllm.mhc_post_mapping( + residual, + x, + post_mix, + comb_mix, + out, + B, + hidden_size, + ) + + return out + + +# --------------------------------------------------------------------------- +# Python API — fused_hc (post_mapping(prev) + pre_mapping(cur)) +# --------------------------------------------------------------------------- + + +_FUSED_HC_BACKEND_CODE = { + "fused_half_mma": 0, # 2-kernel tf32 tcgen05 (fused pmap_gemm + bigfuse) + "fused_half_fma": 1, # 2-kernel fp32 FMA ksplit (fused pmap_gemm + bigfuse) + "fused_all_mma": 2, # 1-kernel tf32 tcgen05 all-in-one (Path D) + "fused_all_fma": 3, # 1-kernel fp32 FMA all-in-one (Path F) +} + +# Tactics supported by the half-fused FMA path in `mhcFusedHcFmaLaunch` +# (must stay in sync with the C++ pickFhcFma() table). +_FUSED_HC_HALF_FMA_TN_KS = ( + (1, 1), + (1, 2), + (1, 4), + (1, 8), + (2, 1), + (2, 2), + (2, 4), + (2, 8), + (3, 1), + (3, 2), + (3, 4), + (4, 1), + (4, 2), + (6, 1), + (8, 1), + (12, 1), + (24, 1), +) +# Tactics for the half-fused MMA path: (num_k_splits,). Matches Path D +# (pickFhcAllInOne) so the autotuner can compare half-fused vs all-in-one at +# the same ks across the full range. +_FUSED_HC_HALF_MMA_KS = (1, 2, 4, 8, 16, 32, 64) +# Tactics for Path D (all-in-one MMA): (num_k_splits,). No bigfuse_bs — the +# bigfuse runs inline inside the single kernel and uses fixed parameters. +_FUSED_HC_ALL_MMA_KS = (1, 2, 4, 8, 16, 32, 64) +# Tactics for Path F (all-in-one FMA): (tile_n, num_k_splits, tile_m). +# Must stay in sync with the C++ pickFhcFmaAllInOne() table. +_FUSED_HC_ALL_FMA_TN_KS_TM = tuple( + (tn, ks, tm) + for tm in (1, 2, 4) + for tn, ks in ( + (1, 1), + (1, 2), + (2, 1), + (2, 2), + (3, 1), + (4, 1), + (6, 1), + (8, 1), + (12, 1), + (24, 1), + ) +) +# Shared BigFuse block-size options (same as pre_mapping autotuner). +_FUSED_HC_BIGFUSE_BS = _BIGFUSE_BLOCK_SIZE_OPTIONS + + +def _fused_hc_call( + backend_code: int, + tile_n: int, + num_k_splits: int, + bigfuse_bs: int, + tile_m: int, + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + y_acc_ws, + r_acc_ws, + done_counter_ws, + B: int, + hidden_size: int, + n: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +): + torch.ops.trtllm.mhc_fused_hc( + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + y_acc_ws, + r_acc_ws, + done_counter_ws, + B, + hidden_size, + n, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + backend_code, + tile_n, + num_k_splits, + bigfuse_bs, + tile_m, + ) + + +class _FusedHcWorkspaceCache: + """Size-keyed bounded LRU for the 4 outputs + 3 workspaces of mhc_fused_hc. + + The 4 outputs are consumed by the caller (so they can't alias across + calls at different B), but repeatedly calling ``torch.empty`` for each + call inside a CUDA-graph-captured inference loop is wasteful. Keyed on + ``(B, ws_ks, tile_m, device)``: same-shape calls reuse the previously + allocated buffers; tensors are stable across graph captures (same ptr + under the torch allocator's retained block). + + Two bounds keep this from ballooning: + + 1. ``_CACHE_MAX_B``: a per-call threshold. Only cache when the request's + residual_cur footprint is modest (B * n * hidden * 2 bytes; e.g. at + n=4 hidden=4096 that's 32 KB/token, so a 256-token cap = 8 MB/entry). + Prefill shapes (B in the tens of thousands) flow straight through to + ``torch.empty``, which the torch caching allocator already keeps + cheap on repeated calls. + 2. ``_maxsize``: LRU cap on the number of cached entries. Covers the + discrete CUDA-graph decode batch sizes plus a few stragglers; decode + is the only regime that actually benefits from the cache. + + Without these bounds every distinct prefill B leaks ~B * n * hidden * 2 + bytes (≈1.3 GB at B=32768, n=4, hidden=4096). Under a prefill ramp-up + admitting one new ctx request per iter, the leak reaches tens of GB per + rank within a dozen iters and blows past HBM. + """ + + __slots__ = ("n", "hidden_size", "_cache", "_maxsize") + + # Up to 48 distinct entries — covers the 35 CUDA-graph decode buckets + # plus headroom. Each entry at B<=256 is under ~10 MB. + DEFAULT_MAXSIZE = 48 + # Skip the cache above this B; prefill rides the torch allocator. + _CACHE_MAX_B = 256 + + def __init__(self, n: int, hidden_size: int, maxsize: int = DEFAULT_MAXSIZE): + self.n = n + self.hidden_size = hidden_size + self._maxsize = maxsize + from collections import OrderedDict + + self._cache: "OrderedDict" = OrderedDict() + + def get(self, B: int, num_k_splits: int, tile_m: int, device): + n = self.n + hidden_size = self.hidden_size + ws_ks = max(1, num_k_splits) + tm = max(1, tile_m) + m_batches = (B + tm - 1) // tm + n2 = n * n + shape_n = n * (2 + n) + + def _alloc(): + residual_cur = torch.empty((B, n, hidden_size), dtype=torch.bfloat16, device=device) + post_mix_cur = torch.empty((B, n), dtype=torch.float32, device=device) + comb_mix_cur = torch.empty((B, n2), dtype=torch.float32, device=device) + layer_input_cur = torch.empty((B, hidden_size), dtype=torch.bfloat16, device=device) + if ws_ks == 1: + y_acc_ws = torch.empty((B, shape_n), dtype=torch.float32, device=device) + r_acc_ws = torch.empty((B,), dtype=torch.float32, device=device) + else: + y_acc_ws = torch.empty((ws_ks, B, shape_n), dtype=torch.float32, device=device) + r_acc_ws = torch.empty((ws_ks, B), dtype=torch.float32, device=device) + done_counter_ws = torch.empty((m_batches,), dtype=torch.int32, device=device) + return ( + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + y_acc_ws, + r_acc_ws, + done_counter_ws, + ) + + if B > self._CACHE_MAX_B: + return _alloc() + + key = (B, ws_ks, m_batches, device) + hit = self._cache.get(key) + if hit is not None: + self._cache.move_to_end(key) + return hit + entry = _alloc() + self._cache[key] = entry + if len(self._cache) > self._maxsize: + self._cache.popitem(last=False) + return entry + + +def _alloc_fused_hc_outputs( + B: int, n: int, hidden_size: int, num_k_splits: int, tile_m: int, device +): + """Uncached fallback (kept for API compatibility).""" + return _FusedHcWorkspaceCache(n=n, hidden_size=hidden_size).get( + B, num_k_splits, tile_m, device + ) + + +# Fallback tactic: backend, tile_n, num_k_splits, bigfuse_bs, tile_m. +# +# The MMA (tcgen05) default delegates ks/bs to the C++ heuristic (pickKSplits +# + selectBigFuseBS). That backend requires SM100+, so on pre-SM100 GPUs we +# fall back to the half-fused FMA path with an explicit ks=1 tactic that is +# valid everywhere. +_FUSED_HC_FALLBACK_TACTIC_MMA = ("fused_half_mma", 0, 0, 0, 1) +_FUSED_HC_FALLBACK_TACTIC_FMA = ("fused_half_fma", 2, 1, 256, 1) + + +def _get_fused_hc_fallback_tactic(): + return ( + _FUSED_HC_FALLBACK_TACTIC_MMA + if _fused_hc_mma_supported() + else _FUSED_HC_FALLBACK_TACTIC_FMA + ) + + +class MhcFusedHcRunner(TunableRunner): + """Profiles the full mhc_fused_hc pipeline (pmap+GEMM+bigfuse + residual out). + + Tactic format: (backend, tile_n, num_k_splits, bigfuse_bs, tile_m) + backend: "fused_half_mma" | "fused_half_fma" | + "fused_all_mma" | "fused_all_fma" + tile_n: FMA N-tile size (*_fma only; 0 for *_mma) + num_k_splits: HIDDEN-axis split (all backends) + bigfuse_bs: BigFuse CTA BLOCK_SIZE {128, 256, 512} (fused_half_* only; + fused_all_* runs bigfuse inline so this is 0) + tile_m: M tokens per CTA (fused_all_fma only; 1 otherwise) + + Backend selection strategy (M-bucketed): + M <= 32: prefer "fused_all_fma" (Path F) — single-kernel FMA wins + when MMA can't saturate its pipe. + M >= 64: prefer "fused_all_mma" (Path D) — single-kernel TF32 MMA + wins once M fills at least one BLOCK_M=64 tile. + In-between: autotuner chooses between all four backends. + + Fallback (-1): delegates to the C++ heuristic (fused_half_mma + auto ks + + auto bs). + """ + + tuning_config = TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( + input_idx=0, # x_prev [B, hidden] + dim_idx=0, + gen_tuning_buckets=_mhc_gen_tuning_buckets, + map_to_tuning_buckets=_mhc_map_to_tuning_bucket, + ), + ), + constraint_specs=( + # residual_prev (input[1]) dim 0 = M + ConstraintSpec(input_idx=1, dim_idx=0, infer_shape=lambda shapes: shapes[0][0]), + # post_mix_prev (input[2]) dim 0 = M + ConstraintSpec(input_idx=2, dim_idx=0, infer_shape=lambda shapes: shapes[0][0]), + # comb_mix_prev (input[3]) dim 0 = M + ConstraintSpec(input_idx=3, dim_idx=0, infer_shape=lambda shapes: shapes[0][0]), + ), + ) + + def __init__( + self, + n: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + ): + self.n = n + self.hidden_size = hidden_size + self.rms_eps = rms_eps + self.hc_pre_eps = hc_pre_eps + self.hc_sinkhorn_eps = hc_sinkhorn_eps + self.hc_post_mult_value = hc_post_mult_value + self.sinkhorn_repeat = sinkhorn_repeat + self._ws_cache = _FusedHcWorkspaceCache(n=n, hidden_size=hidden_size) + + def unique_id(self): + return (self.n, self.hidden_size) + + def get_valid_tactics(self, inputs, profile: OptimizationProfile, **kwargs): + M = inputs[0].shape[0] + tactics = [] + # The MMA (tcgen05) paths require SM100+. On older archs only the FMA + # paths are compilable/runnable — we simply never emit MMA tactics. + mma_ok = _fused_hc_mma_supported() + # Path F (fused_all_fma, 1-kernel FMA) — preferred at small M (<=32) + # where MMA can't fill BLOCK_M=64. Include for M <= 64 as the + # crossover is measured per-M. + if M <= 64: + for tn, ks, tm in _FUSED_HC_ALL_FMA_TN_KS_TM: + # Skip grids that wildly oversubscribe SMs. + m_batches = (M + tm - 1) // tm + if m_batches * (self.n * (2 + self.n) // tn) * ks > 148 * 4: + continue + # fused_all_fma runs bigfuse inline — no bigfuse_bs tactic axis. + tactics.append(("fused_all_fma", tn, ks, 0, tm)) + # Path D (fused_all_mma, 1-kernel TF32 MMA) — preferred at mid/large + # M (>=64). Include when M >= 48 to overlap with Path F at the + # crossover boundary. + if mma_ok and M >= 48: + for ks in _FUSED_HC_ALL_MMA_KS: + m_tiles = (M + 63) // 64 + if m_tiles * ks > 148 * 4: + continue + tactics.append(("fused_all_mma", 0, ks, 0, 1)) + # Half-fused FMA path (2-kernel) — useful at smallish M; kept as a + # fallback option for the autotuner at small M. + if M <= 512: + for tn, ks in _FUSED_HC_HALF_FMA_TN_KS: + if ks > 1 and M * (self.n * (2 + self.n) // tn) >= 148 * 2: + continue + for bs in _FUSED_HC_BIGFUSE_BS: + tactics.append(("fused_half_fma", tn, ks, bs, 1)) + # Half-fused MMA path (2-kernel) — always an option when tcgen05 is + # available. + if mma_ok: + for ks in _FUSED_HC_HALF_MMA_KS: + m_tiles = (M + 63) // 64 + if m_tiles * ks > 148 * 4: + continue + for bs in _FUSED_HC_BIGFUSE_BS: + tactics.append(("fused_half_mma", 0, ks, bs, 1)) + return tactics + + def forward(self, inputs, *, tactic=-1, **kwargs): + ( + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + ) = inputs + + x_prev = x_prev.contiguous() + residual_prev = residual_prev.contiguous() + post_mix_prev = post_mix_prev.to(torch.float32).contiguous() + comb_mix_prev = comb_mix_prev.to(torch.float32).contiguous() + w_t_cur = w_t_cur.to(torch.float32).contiguous() + hc_scale_cur = hc_scale_cur.to(torch.float32).contiguous() + hc_base_cur = hc_base_cur.to(torch.float32).contiguous() + + if tactic == -1: + tactic = _get_fused_hc_fallback_tactic() + backend, tile_n, num_k_splits, bigfuse_bs, tile_m = tactic + backend_code = _FUSED_HC_BACKEND_CODE[backend] + + B = residual_prev.shape[0] + ( + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + y_acc_ws, + r_acc_ws, + done_counter_ws, + ) = self._ws_cache.get(B, num_k_splits, tile_m, x_prev.device) + + _fused_hc_call( + backend_code, + tile_n, + num_k_splits, + bigfuse_bs, + tile_m, + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + y_acc_ws, + r_acc_ws, + done_counter_ws, + B, + self.hidden_size, + self.n, + self.rms_eps, + self.hc_pre_eps, + self.hc_sinkhorn_eps, + self.hc_post_mult_value, + self.sinkhorn_repeat, + ) + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + + +# Process-wide runner cache keyed on the mHC configuration. Avoids recreating +# a MhcFusedHcRunner (and its workspace cache) on every call, which would also +# defeat the workspace cache inside the runner. +_fused_hc_runner_cache: dict = {} + + +def _get_fused_hc_runner( + n: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +) -> "MhcFusedHcRunner": + key = ( + n, + hidden_size, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + runner = _fused_hc_runner_cache.get(key) + if runner is None: + runner = MhcFusedHcRunner( + n=n, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + ) + _fused_hc_runner_cache[key] = runner + return runner + + +def mhc_fused_hc( + x_prev: torch.Tensor, + residual_prev: torch.Tensor, + post_mix_prev: torch.Tensor, + comb_mix_prev: torch.Tensor, + w_t_cur: torch.Tensor, + hc_scale_cur: torch.Tensor, + hc_base_cur: torch.Tensor, + n: int, + hidden_size: int, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, +): + """Fuse the previous block's post_mapping with the current block's pre_mapping. + + The autotuner chooses between four backends: + * "fused_half_mma" — 2-kernel tcgen05 TF32 pmap+GEMM atomic + bigfuse. + * "fused_half_fma" — 2-kernel pmap inline + FMA GEMM + sqrsum + bigfuse. + * "fused_all_mma" — 1-kernel TF32 tcgen05 all-in-one (Path D). + * "fused_all_fma" — 1-kernel FMA all-in-one (Path F). + + Returns: + residual_cur: [B, n, hidden] bf16 (new residual, input to the next post_mapping) + post_mix_cur: [B, n] fp32 + comb_mix_cur: [B, n*n] fp32 + layer_input_cur: [B, hidden] bf16 (input to this block's attn/MoE) + """ + runner = _get_fused_hc_runner( + n=n, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_pre_eps=hc_pre_eps, + hc_sinkhorn_eps=hc_sinkhorn_eps, + hc_post_mult_value=hc_post_mult_value, + sinkhorn_repeat=sinkhorn_repeat, + ) + + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( + "trtllm::mhc_fused_hc", + [runner], + MhcFusedHcRunner.tuning_config, + [x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t_cur, hc_scale_cur, hc_base_cur], + ) + + return runner( + inputs=[ + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + ], + tactic=best_tactic, + ) + + +# --------------------------------------------------------------------------- +# Python API — HCHead +# --------------------------------------------------------------------------- + + +def mhc_hc_head_cuda( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + mult: int, + hidden_size: int, + norm_eps: float = 1e-5, + eps: float = 1e-5, +) -> torch.Tensor: + """HCHead forward: RMS-normed GEMM -> sigmoid -> weighted sum. + + Args: + x: [M, mult, hidden_size] bf16 input + fn: [mult, K] fp32 weight (K = mult * hidden_size) + scale: [1] fp32 + base: [mult] fp32 + mult: number of hyper-connection heads (typically 4) + hidden_size: per-head dimension + norm_eps: RMS norm epsilon + eps: sigmoid offset epsilon + + Returns: [M, hidden_size] bf16 + """ + M = x.shape[0] + K = mult * hidden_size + + x_flat = x.reshape(M, K).contiguous() + fn_t = fn.to(torch.float32).contiguous() + scale = scale.to(torch.float32).contiguous() + base = base.to(torch.float32).contiguous() + + mixes, sqrsum = mhc_gemm_rms_fma_cuda( + x_flat, + None, + M, + mult, + K, + w_t=fn_t, + ) + + out = torch.empty((M, hidden_size), dtype=torch.bfloat16, device=x.device) + + torch.ops.trtllm.mhc_hc_head_apply( + mixes, + sqrsum, + x.reshape(M, mult, hidden_size).contiguous(), + out, + scale, + base, + M, + mult, + hidden_size, + K, + float(norm_eps), + float(eps), + ) + + return out diff --git a/tensorrt_llm/_torch/modules/rotary_embedding.py b/tensorrt_llm/_torch/modules/rotary_embedding.py index 150541e4d048..db9dac7a51ec 100644 --- a/tensorrt_llm/_torch/modules/rotary_embedding.py +++ b/tensorrt_llm/_torch/modules/rotary_embedding.py @@ -15,11 +15,13 @@ def __init__( *, head_dim: int, is_neox: bool = True, + inverse: bool = False, ): super().__init__() self.rope_params = rope_params self.head_dim = head_dim self.is_neox = is_neox + self.inverse = inverse self.max_positions = rope_params.max_positions self.rotary_cos_sin = rope_params.create_rope_const_params( interleave=False)[1].reshape(rope_params.max_positions, 2, -1) @@ -70,7 +72,8 @@ def rope_target(target): target = RotaryEmbedding.apply_rotary_pos_emb(target, cos, sin, - is_neox=self.is_neox) + is_neox=self.is_neox, + inverse=self.inverse) target = target.transpose(1, 2).contiguous() if remove_input_padding: target = target.view(seq_len, -1) @@ -85,7 +88,8 @@ def apply_rotary_pos_emb(q_or_k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1, - is_neox: bool = True) -> torch.Tensor: + is_neox: bool = True, + inverse: bool = False) -> torch.Tensor: """Applies Rotary Position Embedding to the query and key tensors. Args: @@ -100,6 +104,7 @@ def apply_rotary_pos_emb(q_or_k: torch.Tensor, cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. is_neox (bool): Whether to use Neox style RoPE, True by default. + inverse (bool): Whether to apply inverse RoPE (for de-rotating), False by default. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ @@ -117,8 +122,12 @@ def apply_rotary_pos_emb(q_or_k: torch.Tensor, x1 = q_or_k[..., ::2] x2 = q_or_k[..., 1::2] - o1 = x1 * cos - x2 * sin - o2 = x2 * cos + x1 * sin + if inverse: + o1 = x1 * cos + x2 * sin + o2 = x2 * cos - x1 * sin + else: + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin if is_neox: return torch.cat((o1, o2, q_or_k_pass), dim=-1) diff --git a/tensorrt_llm/_torch/modules/swiglu.py b/tensorrt_llm/_torch/modules/swiglu.py index 34058786ca12..328bb4e9fdcb 100644 --- a/tensorrt_llm/_torch/modules/swiglu.py +++ b/tensorrt_llm/_torch/modules/swiglu.py @@ -1,3 +1,5 @@ +from typing import Optional + import torch import triton # type: ignore[import] import triton.language as tl # type: ignore[import] @@ -25,8 +27,9 @@ def scale_and_clamp(x, scale, dtype): @triton.jit def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, - BLOCK_SIZE: tl.constexpr, - HAS_O_SCALE: tl.constexpr) -> None: + swiglu_limit: tl.constexpr, BLOCK_SIZE: tl.constexpr, + HAS_O_SCALE: tl.constexpr, + HAS_SWIGLU_LIMIT: tl.constexpr) -> None: i = tl.program_id(axis=0).to(tl.int64) j = tl.program_id(axis=1) @@ -39,6 +42,10 @@ def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, a = tl.load(x_row_ptr + offsets, mask=mask).to(tl.float32) b = tl.load(x_row_ptr + offsets + d, mask=mask).to(tl.float32) + if HAS_SWIGLU_LIMIT: + a = tl.minimum(a, swiglu_limit) + b = tl.clamp(b, -swiglu_limit, swiglu_limit) + result = tl.sigmoid(a) * a * b if HAS_O_SCALE: @@ -48,13 +55,17 @@ def silu_and_mul_kernel(o_ptr, o_stride, o_scale_ptr, x_ptr, x_stride, d, tl.store(o_row_ptr + offsets, result, mask=mask) -def swiglu(x, quant_scale: torch.Tensor = None, quant_type=None): +def swiglu(x, + quant_scale: Optional[torch.Tensor] = None, + quant_type=None, + swiglu_limit: Optional[float] = None): if quant_scale is not None: assert quant_type is not None return torch.ops.trtllm.silu_and_mul( x, scale=quant_scale, dtype=quant_type, + swiglu_limit=swiglu_limit, ) - return torch.ops.trtllm.silu_and_mul(x) + return torch.ops.trtllm.silu_and_mul(x, swiglu_limit=swiglu_limit) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 1fef2c239d37..4f6ec92bd6cd 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -798,6 +798,7 @@ def _create_kv_cache_manager( sparse_attn_config=self._sparse_attention_config, max_num_tokens=self._max_num_tokens, max_beam_width=self._max_beam_width, + max_input_len=self._llm_args.max_input_len, kv_connector_manager=self._kv_connector_manager, estimating_kv_cache=estimating_kv_cache, execution_stream=self._execution_stream, @@ -844,6 +845,11 @@ def _should_create_separate_draft_kv_cache(self) -> bool: "Attention DP is enabled, separate draft KV cache is not supported." ) return False + if self._sparse_attention_config is not None: + logger.info( + "Sparse attention is enabled, separate draft KV cache is not supported." + ) + return False return should_use_separate_draft_kv_cache(self._speculative_config) def _get_effective_draft_config(self) -> ModelConfig: @@ -1476,6 +1482,8 @@ def _create_kv_cache_manager( max_num_tokens: int, max_beam_width: int, kv_connector_manager: Optional[KvCacheConnectorManager], + is_disagg: bool = False, + max_input_len: Optional[int] = None, estimating_kv_cache: bool = False, execution_stream: Optional[torch.cuda.Stream] = None, # Optional overrides for one-model draft case (when model_engine is None) @@ -1486,8 +1494,7 @@ def _create_kv_cache_manager( num_layers: Optional[int] = None, num_kv_heads: Optional[Union[int, List[int]]] = None, head_dim: Optional[int] = None, - kv_cache_type=None, - is_disagg: bool = False) -> KVCacheManager: + kv_cache_type=None) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config @@ -1619,6 +1626,7 @@ def _create_kv_cache_manager( vocab_size=config.vocab_size, max_beam_width=max_beam_width, is_draft=is_draft, + max_input_len=max_input_len, kv_connector_manager=kv_connector_manager if not estimating_kv_cache else None, sparse_attn_config=sparse_attn_config, @@ -1626,6 +1634,7 @@ def _create_kv_cache_manager( execution_stream=execution_stream, layer_mask=layer_mask, is_disagg=is_disagg, + max_num_tokens=max_num_tokens, ) elif is_nemotron_hybrid(config): if max_beam_width > 1: @@ -1804,6 +1813,7 @@ def _create_kv_cache_manager( max_num_tokens=max_num_tokens, model_config=binding_model_config, max_beam_width=max_beam_width, + max_input_len=max_input_len, is_draft=is_draft, kv_connector_manager=kv_connector_manager if not estimating_kv_cache else None, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 2ab710bf14ac..ca2abedc0f15 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -445,6 +445,7 @@ def __getitem__(self, key): kimi_k2="DeepseekV3Config", glm_moe_dsa="DeepseekV3Config", laguna="LagunaConfig", + deepseek_v4="DeepseekV4Config", ) # NOTE: HF config.json uses deepseek_v32 as model_type but with same DSV3 config class diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 32f456d76407..0e3c03c7b15e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -44,7 +44,7 @@ from ..distributed import Distributed from ..distributed.communicator import init_pp_comm from ..expert_statistic import ExpertStatistic -from ..memory_buffer_utils import with_shared_pool +from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool from ..metadata import KVCacheParams from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.modeling_multimodal_utils import filter_mm_token_from_input_ids @@ -1223,6 +1223,14 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): ) AutoTuner.get().print_profiling_cache() + # Clear workspace buffers allocated during the autotuner forward pass. + # The autotuner runs a context-only forward with max_num_tokens, which + # causes the global Buffers pool to cache large MoE/GEMM workspaces. + # If not cleared, these inflate the memory baseline seen by the KV cache + # profiler, reducing memory available for activations during inference. + clear_memory_buffers() + torch.cuda.empty_cache() + def _compute_dynamic_draft_len_mapping(self) -> Optional[Dict[int, int]]: """Compute graph_bs → draft_len mapping for dynamic draft length feature. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2bb4ec572f2e..73ef1246bc17 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2464,6 +2464,9 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): kv_cache_dtype_byte_size = getattr(self.model_engine, 'kv_cache_dtype_byte_size', None) + if self._scheduler_manages_kv_suspend: + self.kv_cache_manager.update_context_resources( + sample_state_scheduled_requests) self.resource_manager.update_resources( sample_state_scheduled_requests, attn_metadata, kv_cache_dtype_byte_size) @@ -2948,7 +2951,8 @@ def _executor_loop(self): if scheduled_batch.encoder_requests: self._run_encoder_step(scheduled_batch.encoder_requests) - can_queue, _ = self._can_queue(scheduled_batch) + can_queue, can_queue_this_rank = self._can_queue( + scheduled_batch) if can_queue: if self.kv_cache_transceiver: @@ -2971,7 +2975,15 @@ def _executor_loop(self): # if using a kv connector, we need to call can_queue again since scheduled_batch might have changed if self.kv_connector_manager: - can_queue, _ = self._can_queue(scheduled_batch) + can_queue, can_queue_this_rank = self._can_queue( + scheduled_batch) + + # Revert spurious KV cache capacity growth (see overlap + # loop for the full comment). + if not can_queue and can_queue_this_rank \ + and self._scheduler_manages_kv_suspend: + for req in scheduled_batch.generation_requests: + self.kv_cache_manager.revert_allocate_generation(req) if not can_queue: self._revert_gen_alloc(scheduled_batch) @@ -3079,6 +3091,9 @@ def _executor_loop(self): None) kv_cache_dtype_byte_size = getattr( self.model_engine, 'kv_cache_dtype_byte_size', None) + if self._scheduler_manages_kv_suspend: + self.kv_cache_manager.update_context_resources( + scheduled_batch) self.resource_manager.update_resources( scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) @@ -3373,6 +3388,18 @@ def _executor_loop_overlap(self): # we need to delay the update of the previous batch's sample state, # and let the later iteration to update it. should_process_previous_batch = can_queue or not can_queue_this_rank + + # With attention DP, can_queue=False means another rank has + # an empty batch so no forward pass will run. The V2 + # scheduler already grew each generation request's KV cache + # capacity during scheduling; revert that growth so it does + # not accumulate across skipped iterations and overflow the + # host page-index buffer. + if not can_queue and can_queue_this_rank \ + and self._scheduler_manages_kv_suspend: + for req in scheduled_batch.generation_requests: + self.kv_cache_manager.revert_allocate_generation(req) + if can_queue: # The generation requests that do not have batch_idx @@ -3493,6 +3520,14 @@ def _executor_loop_overlap(self): scheduled_batch, guided_decoder_failed_requests) self._update_request_states(scheduled_batch) + # Update context requests' KV cache so that sliding-window + # blocks freed by this chunk are visible to the next + # iteration's scheduler. + # Only applies to KV cache manager V2 + scheduler V2. + if self._scheduler_manages_kv_suspend: + self.kv_cache_manager.update_context_resources( + scheduled_batch) + if self.previous_batch is not None and should_process_previous_batch: self._process_previous_batch() self.perf_manager.compute_batch_gpu_times( diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 9c4284878b06..9392c0a5847f 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -189,6 +189,16 @@ def get_spec_metadata(spec_config, return None +def get_mtp_hidden_size(model_config) -> int: + pretrained_config = getattr(model_config, "pretrained_config", model_config) + hidden_size = getattr(pretrained_config, "hidden_size", None) + if hidden_size is None: + hidden_size = getattr(model_config, "hidden_size") + if getattr(pretrained_config, "model_type", None) == "deepseek_v4": + return hidden_size * getattr(pretrained_config, "hc_mult", 1) + return hidden_size + + def get_spec_resource_manager(model_engine, draft_model_engine=None): spec_config = model_engine.spec_config if spec_config is None: @@ -211,7 +221,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): return Eagle3ResourceManager( spec_config, model_config.torch_dtype, - model_config.hidden_size, + get_mtp_hidden_size(model_config), max_num_requests, max_seq_len, max_num_tokens, @@ -228,7 +238,7 @@ def get_spec_resource_manager(model_engine, draft_model_engine=None): return MTPHiddenStatesManager( spec_config, model_config.torch_dtype, - model_config.hidden_size, + get_mtp_hidden_size(model_config), max_num_requests, sa_manager=sa_manager, ) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 4e9c92c9ba76..62a69890e93b 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -25,6 +25,7 @@ 'MoeBalancer', 'MoeOutputMemset', 'MoeFc2Alpha', + 'EngramPrecompute', ] AuxStreamType = Enum( 'AuxStreamType', diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index cff3383c8579..520579efd2e9 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -178,6 +178,7 @@ def torch_version(): int64=np.int64, int32=np.int32, int8=np.int8, + uint8=np.uint8, bool=np.bool_, bfloat16=np_bfloat16, fp8=np_float8, @@ -197,6 +198,7 @@ def str_dtype_to_np(dtype): int64=torch.int64, int32=torch.int32, int8=torch.int8, + uint8=torch.uint8, bool=torch.bool, fp8=torch.float8_e4m3fn, ) @@ -215,6 +217,7 @@ def str_dtype_to_torch(dtype): int64=DataType.INT64, int32=DataType.INT32, int8=DataType.INT8, + uint8=DataType.UINT8, bool=DataType.BOOL, fp8=DataType.FP8, ) diff --git a/tensorrt_llm/bench/build/dataclasses.py b/tensorrt_llm/bench/build/dataclasses.py index 0b27ca4cbf9b..37d63062f23a 100755 --- a/tensorrt_llm/bench/build/dataclasses.py +++ b/tensorrt_llm/bench/build/dataclasses.py @@ -207,6 +207,11 @@ def get_param_count_and_checkpoint_size(cls, model_hf_name, hf_model_path): if model_hf_name == "EleutherAI/gpt-j-6b": # GPT-J repo doesn't use safetensor format. param_count = 6053381344 checkpoint_size_in_gb = param_count * 2 / (1024**3) + # TODO: Remove this once HF supports DeepSeek-V4 + elif model_hf_name == "deepseek-ai/DeepSeek-V4": + param_count = 284347051735 + # DeepSeek-V4 ships as an FP8 checkpoint (~1 byte/param). + checkpoint_size_in_gb = param_count / (1024**3) else: model_name_or_path = hf_model_path or model_hf_name metadata = get_safetensors_metadata(model_name_or_path) diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index d88bd2b722fe..7025e372c986 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -82,7 +82,8 @@ def get_llm_args(self) -> Dict: if self.backend in backend_config_map: llm_args.update(backend_config_map[self.backend]()) - kv_cache_config = self.settings_config.get_kvcache_config().__dict__ + kv_cache_config = self.settings_config.get_kvcache_config().model_dump( + exclude_unset=True) backend_cache_config = llm_args.pop("kv_cache_config", {}) llm_args["kv_cache_config"] = backend_cache_config | kv_cache_config diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index eaa98cad44ef..fcc53dbc73f1 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -35,6 +35,7 @@ MultimodalPlaceholderPlacement) from tensorrt_llm.llmapi.llm_utils import ModelLoader from tensorrt_llm.tokenizer import TokenizerBase, TransformersTokenizer +from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer logger = logging.get_logger(__name__) @@ -652,8 +653,8 @@ def apply_chat_template( - STRING: keeps flattened text with pre-inserted placeholders """ - # Handle DeepSeek V32 tokenizer with custom chat template - if isinstance(tokenizer, DeepseekV32Tokenizer): + # Handle DeepSeek tokenizers with custom chat templates. + if isinstance(tokenizer, (DeepseekV32Tokenizer, DeepseekV4Tokenizer)): prompt = tokenizer.apply_chat_template( messages=conversation, tools=tools, diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 7495be04f9bc..cf765bc536bd 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -5,12 +5,14 @@ from ..scheduling_params import SchedulingParams from .build_cache import BuildCacheConfig from .llm import LLM, RequestOutput + # yapf: disable from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType, CacheTransceiverConfig, CalibConfig, CapacitySchedulerPolicy, ContextChunkingPolicy, CudaGraphConfig, DecodeCudaGraphConfig, - DeepSeekSparseAttentionConfig, DFlashDecodingConfig, + DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, DFlashDecodingConfig, DraftTargetDecodingConfig, DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, @@ -88,4 +90,5 @@ 'PrometheusMetricsConfig', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', + 'DeepSeekV4SparseAttentionConfig', ] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 19704e1a85ed..0f8c9b3cc477 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -622,6 +622,37 @@ def needs_separate_short_long_cuda_graphs(self) -> bool: return self.skip_indexer_for_short_seqs +class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): + """ + Configuration for DeepSeek-V4 Sparse Attention. + """ + algorithm: Literal["deepseek_v4"] = "deepseek_v4" + skip_indexer_for_short_seqs: bool = Field( + default=False, + description= + "Whether to skip the MQA and Top-K in the indexer for short sequences.") + + compress_ratios: List[int] = Field( + default_factory=lambda: [1, 1, 4, 128, 4, 128, 4], + description="The compress ratios of each layer. DeepSeek-V4 uses 0 " + "for uncompressed/SWA-only layers; internal allocation code normalizes " + "0 to 1, while checkpoint-facing semantics remain unchanged.") + + window_size: int = Field( + default=128, + description="The window size for slicing window attention part.") + + index_topk: Optional[int] = Field(default=512, + description="The topk for the indexer.") + + def supports_backend(self, backend: str) -> bool: + return backend == "pytorch" + + def needs_separate_short_long_cuda_graphs(self) -> bool: + # DeepSeek-V4 does not support short/long CUDA graph separation. + return False + + class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): """Configuration for skip softmax attention.""" algorithm: Literal["skip_softmax"] = Field(default="skip_softmax") @@ -2702,6 +2733,7 @@ def supports_backend(self, backend: str) -> bool: Union[ RocketSparseAttentionConfig, DeepSeekSparseAttentionConfig, + DeepSeekV4SparseAttentionConfig, SkipSoftmaxAttentionConfig, ], Field(discriminator="algorithm"), @@ -3544,9 +3576,13 @@ def validate_and_init_tokenizer(self): if self.skip_tokenizer_init: self.tokenizer = None elif self.custom_tokenizer: - # If tokenizer is already a tokenizer object, custom_tokenizer is not compatible - if isinstance(self.tokenizer, - (TokenizerBase, PreTrainedTokenizerBase)): + # IPC workers receive the tokenizer object that was already loaded + # in the parent LLM process. Reuse TRT-LLM tokenizer wrappers as-is. + if isinstance(self.tokenizer, TokenizerBase): + return self + # A raw HF tokenizer object would bypass the requested custom + # wrapper, so keep rejecting that combination. + if isinstance(self.tokenizer, PreTrainedTokenizerBase): raise ValueError( "Cannot use custom_tokenizer when tokenizer is already a tokenizer object. " "Please specify a tokenizer path or leave it as None to load from model path." diff --git a/tensorrt_llm/llmapi/tokenizer.py b/tensorrt_llm/llmapi/tokenizer.py index 214ca86f9780..5f593d5d6000 100644 --- a/tensorrt_llm/llmapi/tokenizer.py +++ b/tensorrt_llm/llmapi/tokenizer.py @@ -7,6 +7,7 @@ _xgrammar_tokenizer_info, load_hf_tokenizer, tokenizer_factory) from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer +from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer __all__ = [ "TLLM_INCREMENTAL_DETOKENIZATION_BACKEND", @@ -14,6 +15,7 @@ "TokenizerBase", "TransformersTokenizer", "DeepseekV32Tokenizer", + "DeepseekV4Tokenizer", "tokenizer_factory", "_xgrammar_tokenizer_info", "_llguidance_tokenizer_info", diff --git a/tensorrt_llm/tokenizer/deepseek_v4/__init__.py b/tensorrt_llm/tokenizer/deepseek_v4/__init__.py new file mode 100644 index 000000000000..30c51f77d64d --- /dev/null +++ b/tensorrt_llm/tokenizer/deepseek_v4/__init__.py @@ -0,0 +1,18 @@ +# 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. + +from .tokenizer import DeepseekV4Tokenizer + +__all__ = ["DeepseekV4Tokenizer"] diff --git a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py new file mode 100644 index 000000000000..7e022ea27158 --- /dev/null +++ b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py @@ -0,0 +1,97 @@ +# 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. + +from pathlib import Path +from typing import Any + +from transformers import AutoTokenizer + +from ..tokenizer import TransformersTokenizer + +BOS_TOKEN = "<｜begin▁of▁sentence｜>" +EOS_TOKEN = "<｜end▁of▁sentence｜>" +USER_TOKEN = "<｜User｜>" +ASSISTANT_TOKEN = "<｜Assistant｜>" +THINKING_END_TOKEN = "" + + +def _message_content_to_text(content: Any) -> str: + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + else: + parts.append(str(block)) + return "\n\n".join(parts) + return str(content) + + +class DeepseekV4Tokenizer(TransformersTokenizer): + """DeepSeek-V4 tokenizer with the checkpoint reference chat format.""" + + @classmethod + def from_pretrained( + cls, + path_or_repo_id: str | Path, + *args, + trust_remote_code: bool = False, + revision: str | None = None, + **kwargs, + ) -> "DeepseekV4Tokenizer": + tokenizer = AutoTokenizer.from_pretrained( + path_or_repo_id, + *args, + trust_remote_code=trust_remote_code, + revision=revision, + **kwargs, + ) + return cls(tokenizer) + + def apply_chat_template(self, messages, tools=None, **kwargs): + if tools: + raise NotImplementedError( + "DeepSeek-V4 tool-call chat formatting is not supported yet.") + + add_generation_prompt = kwargs.get("add_generation_prompt", True) + tokenize = kwargs.get("tokenize", False) + + rendered = BOS_TOKEN + for idx, message in enumerate(messages): + role = message.get("role") + content = _message_content_to_text(message.get("content")) + next_role = (messages[idx + 1].get("role") + if idx + 1 < len(messages) else None) + + if role == "system": + rendered += content + elif role in ("user", "developer"): + rendered += USER_TOKEN + content + if next_role == "assistant" or (next_role is None + and add_generation_prompt): + rendered += ASSISTANT_TOKEN + THINKING_END_TOKEN + elif role == "assistant": + rendered += content + EOS_TOKEN + else: + raise NotImplementedError( + f"Unsupported DeepSeek-V4 message role: {role}") + + if tokenize: + return self.encode(rendered, add_special_tokens=False) + return rendered diff --git a/tensorrt_llm/tokenizer/tokenizer.py b/tensorrt_llm/tokenizer/tokenizer.py index 2e86bc0fd218..7c6e1bd69dd0 100644 --- a/tensorrt_llm/tokenizer/tokenizer.py +++ b/tensorrt_llm/tokenizer/tokenizer.py @@ -26,6 +26,7 @@ # Aliases for built-in custom tokenizers. TOKENIZER_ALIASES = { "deepseek_v32": "tensorrt_llm.tokenizer.deepseek_v32.DeepseekV32Tokenizer", + "deepseek_v4": "tensorrt_llm.tokenizer.deepseek_v4.DeepseekV4Tokenizer", } TLLM_INCREMENTAL_DETOKENIZATION_BACKEND = os.environ.get( diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index f2c4beb1fdc7..b411a79e9a02 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -237,8 +237,8 @@ def make_balanced_routing_method( dp_rank, ep_size, ): - def balanced_routing_method(router_logits): - token_selected_experts, token_final_scales = apply_method_orig(router_logits) + def balanced_routing_method(router_logits, input_ids=None): + token_selected_experts, token_final_scales = apply_method_orig(router_logits, input_ids) assert moe_module._routing_results_replaced_at in [None, "make_balanced_routing_method"] if balance_method == BalanceMethod.Balanced: token_selected_experts = get_balanced_selection( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 43f513bbe6aa..0dea81f4fbc1 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -132,6 +132,8 @@ l0_b200: - unittest/_torch/modules/test_triton_linear.py - unittest/_torch/modules/test_group_rmn_norm.py - unittest/_torch/modules/test_rotary_embedding.py + - unittest/_torch/modules/test_engram.py + - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules # ------------- MoE components tests --------------- @@ -179,6 +181,7 @@ l0_b200: - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_e4b_dummy - unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_multimodal_26b_dummy + - unittest/_torch/modeling/test_modeling_deepseekv4.py - unittest/tools/test_layer_wise_benchmarks.py::test_deepseek_r1_ctx_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_nemotron_gen_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_qwen3_next_gen_tep[1] diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index b8368937244c..a5601e39b74b 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -27,6 +27,8 @@ l0_b300: - unittest/_torch/modules/test_triton_linear.py - unittest/_torch/modules/test_group_rmn_norm.py - unittest/_torch/modules/test_rotary_embedding.py + - unittest/_torch/modules/test_engram.py + - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules # ------------- MoE components tests --------------- diff --git a/tests/unittest/_torch/attention/sparse/__init__.py b/tests/unittest/_torch/attention/sparse/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/__init__.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py new file mode 100644 index 000000000000..02cba87f3f5e --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py @@ -0,0 +1,2549 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Test suite for KV compressor kernels. + +Tests cover: prefill/decode corner cases, state updates, varlen, MTP support. +Run: pytest -s tests/unittest/_torch/attention/sparse/test_compressor_kernel.py +""" + +import pytest +import torch +import triton +from utils.util import skip_pre_blackwell + +_CUDA_SUPPORTED_HEAD_DIMS = (128, 512) + + +def prefill_kernel( + kv_score: torch.Tensor, + ape: torch.Tensor, + kv_lens: torch.Tensor, + start_pos: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_new_comp_kv: torch.Tensor, + kv_comp: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + max_outputs: int, + block_table_score: torch.Tensor = None, + compress_ratio: int = None, + head_dim: int = None, + page_size: int = 32, +): + """CUDA prefill kernel wrapper with head_dim skip guard.""" + if head_dim not in _CUDA_SUPPORTED_HEAD_DIMS: + pytest.skip( + f"CUDA prefill kernel only supports head_dim in {_CUDA_SUPPORTED_HEAD_DIMS}, got {head_dim}" + ) + if block_table_score is None: + block_table_score = block_table_kv + + paged_kv_2d = paged_kv.flatten(1) if paged_kv.dim() == 3 else paged_kv + paged_score_2d = paged_score.flatten(1) if paged_score.dim() == 3 else paged_score + + torch.ops.trtllm.compressor_prefill_reduction( + kv_score, + ape, + paged_kv_2d, + paged_score_2d, + block_table_kv, + block_table_score, + kv_comp, + kv_lens, + start_pos, + cu_seq_lens, + cu_new_comp_kv, + kv_lens.shape[0], + page_size, + head_dim, + compress_ratio, + max_outputs, + ) + + +def prepare_compress_output( + cu_new_comp_kv: torch.Tensor, + batch_size: int, + head_dim: int, + device: torch.device, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Pre-allocate output tensor for compression kernels. + + Args: + cu_new_comp_kv: [bsz+1] cumulative output offsets + batch_size: Number of batches + head_dim: Dimension of KV head + device: Target device + dtype: Output dtype (default bfloat16) + + Returns: + kv_comp: [total_outputs, head_dim] output buffer + """ + total_outputs = cu_new_comp_kv[-1].item() + kv_comp = torch.empty(total_outputs, head_dim, device=device, dtype=dtype) + return kv_comp + + +def decode_kernel( + kv_score: torch.Tensor, + ape: torch.Tensor, + kv_lens: torch.Tensor, + start_pos: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_new_comp_kv: torch.Tensor, + kv_comp: torch.Tensor, + paged_kv: torch.Tensor, + paged_score: torch.Tensor, + block_table_kv: torch.Tensor, + block_table_score: torch.Tensor = None, + compress_ratio: int = None, + head_dim: int = None, + page_size: int = 32, + next_n: int = 1, +): + """CUDA decode kernel wrapper with head_dim skip guard.""" + if head_dim not in _CUDA_SUPPORTED_HEAD_DIMS: + pytest.skip( + f"CUDA decode kernel only supports head_dim in {_CUDA_SUPPORTED_HEAD_DIMS}, got {head_dim}" + ) + if block_table_score is None: + block_table_score = block_table_kv + if start_pos is None: + start_pos = kv_lens - next_n + + torch.ops.trtllm.compressor_paged_kv_compress( + kv_score, + ape, + paged_kv, + paged_score, + block_table_kv, + block_table_score, + kv_comp, + kv_lens, + start_pos, + cu_seq_lens, + cu_new_comp_kv, + kv_lens.shape[0], + page_size, + head_dim, + compress_ratio, + next_n, + ) + + +# ============================================================================ +# PyTorch References (mirror model.py logic) +# ============================================================================ + + +def run_pytorch_reference( + new_kv, new_score, ape, kv_state, score_state, token_idx, compress_ratio, head_dim, overlap +): + """Decode reference: single token update + conditional compression.""" + bsz = kv_state.shape[0] + should_compress = (token_idx + 1) % compress_ratio == 0 + ape_ratio = token_idx % compress_ratio + + score_val = new_score + ape[ape_ratio] + kv_val = new_kv + output = None + + if overlap: + update_idx = compress_ratio + ape_ratio + kv_state[:bsz, update_idx] = kv_val.squeeze(1) + score_state[:bsz, update_idx] = score_val.squeeze(1) + if should_compress: + d = head_dim + kv_cat = torch.cat( + [kv_state[:bsz, :compress_ratio, :d], kv_state[:bsz, compress_ratio:, d:]], dim=1 + ) + score_cat = torch.cat( + [score_state[:bsz, :compress_ratio, :d], score_state[:bsz, compress_ratio:, d:]], + dim=1, + ) + output = (kv_cat * score_cat.softmax(dim=1)).sum(dim=1, keepdim=True) + kv_state[:bsz, :compress_ratio] = kv_state[:bsz, compress_ratio:].clone() + score_state[:bsz, :compress_ratio] = score_state[:bsz, compress_ratio:].clone() + else: + kv_state[:bsz, ape_ratio] = kv_val.squeeze(1) + score_state[:bsz, ape_ratio] = score_val.squeeze(1) + if should_compress: + output = (kv_state[:bsz] * score_state[:bsz].softmax(dim=1)).sum(dim=1, keepdim=True) + + return output + + +def run_pytorch_prefill_reference( + kv, score, ape, kv_state, score_state, compress_ratio, head_dim, overlap +): + """Prefill reference: bulk compression + state update.""" + bsz, seqlen, _ = kv.size() + ratio = compress_ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + offset = ratio if overlap else 0 + + # State update + if overlap and cutoff >= ratio: + kv_state[:bsz, :ratio] = kv[:, cutoff - ratio : cutoff] + score_state[:bsz, :ratio] = score[:, cutoff - ratio : cutoff] + ape + + if remainder > 0: + kv, kv_state[:bsz, offset : offset + remainder] = kv.split([cutoff, remainder], dim=1) + score_state[:bsz, offset : offset + remainder] = score[:, cutoff:] + ape[:remainder] + score = score[:, :cutoff] + + if cutoff == 0: + return None + + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + ape + + if overlap: + b, s, r, _ = kv.size() + d = head_dim + + kv_transformed = torch.zeros(b, s, 2 * ratio, d, device=kv.device, dtype=kv.dtype) + score_transformed = torch.full( + (b, s, 2 * ratio, d), float("-inf"), device=score.device, dtype=score.dtype + ) + + kv_transformed[:, :, ratio:] = kv[:, :, :, d:] + score_transformed[:, :, ratio:] = score[:, :, :, d:] + kv_transformed[:, 1:, :ratio] = kv[:, :-1, :, :d] + score_transformed[:, 1:, :ratio] = score[:, :-1, :, :d] + + kv, score = kv_transformed, score_transformed + + output = (kv * score.softmax(dim=2)).sum(dim=2) + return output if seqlen >= ratio else None + + +def run_pytorch_prefill_reference_varlen( + kv_score, ape, kv_lens, start_pos, compress_ratio, head_dim, overlap, kv_state, score_state +): + """Varlen prefill reference: process each batch independently.""" + bsz = kv_lens.shape[0] + coff = 2 if overlap else 1 + state_dim = coff * head_dim + ratio = compress_ratio + offset = ratio if overlap else 0 + + # Compute sequence lengths and cumulative offsets + seq_lens = kv_lens - start_pos + cu_seq_lens = torch.zeros(bsz + 1, device=kv_score.device, dtype=torch.int32) + cu_seq_lens[1:] = torch.cumsum(seq_lens, dim=0) + + outputs = [] + + for b in range(bsz): + seqlen = seq_lens[b].item() + input_start = cu_seq_lens[b].item() + + kv_score_b = kv_score[input_start : input_start + seqlen] + kv = kv_score_b[:, :state_dim].unsqueeze(0) + score = kv_score_b[:, state_dim:].unsqueeze(0) + + remainder = seqlen % ratio + cutoff = seqlen - remainder + + # State update + if overlap and cutoff >= ratio: + kv_state[b : b + 1, :ratio] = kv[:, cutoff - ratio : cutoff] + score_state[b : b + 1, :ratio] = score[:, cutoff - ratio : cutoff] + ape + + if remainder > 0: + kv_state[b : b + 1, offset : offset + remainder] = kv[:, cutoff:] + score_state[b : b + 1, offset : offset + remainder] = ( + score[:, cutoff:] + ape[:remainder] + ) + + if cutoff == 0: + continue + + kv_comp = kv[:, :cutoff].unflatten(1, (-1, ratio)) + score_comp = score[:, :cutoff].unflatten(1, (-1, ratio)) + ape + + if overlap: + _, s, r, _ = kv_comp.size() + d = head_dim + + kv_transformed = torch.zeros(1, s, 2 * ratio, d, device=kv.device, dtype=kv.dtype) + score_transformed = torch.full( + (1, s, 2 * ratio, d), float("-inf"), device=score.device, dtype=score.dtype + ) + kv_transformed[:, :, ratio:] = kv_comp[:, :, :, d:] + score_transformed[:, :, ratio:] = score_comp[:, :, :, d:] + kv_transformed[:, 1:, :ratio] = kv_comp[:, :-1, :, :d] + score_transformed[:, 1:, :ratio] = score_comp[:, :-1, :, :d] + kv_comp, score_comp = kv_transformed, score_transformed + + output = (kv_comp * score_comp.softmax(dim=2)).sum(dim=2) + outputs.append(output.squeeze(0)) + + if outputs: + return torch.cat(outputs, dim=0) # [total_outputs, head_dim] + return torch.empty(0, head_dim, device=kv_score.device, dtype=torch.float32) + + +# ============================================================================ +# Test Utilities +# ============================================================================ + + +def fuse_kv_score(kv, score): + """Fuse kv and score into [*, 2*dim].""" + return torch.cat([kv, score], dim=-1) + + +def prepare_decode_metadata( + batch_size: int, + compress_ratio: int, + head_dim: int, + device: torch.device, + next_n: int = 1, +): + """Prepare decode metadata for tests.""" + max_compressions = (next_n + compress_ratio - 1) // compress_ratio + cu_seq_lens = torch.arange( + 0, (batch_size + 1) * next_n, next_n, device=device, dtype=torch.int32 + ) + cu_outputs = torch.arange( + 0, (batch_size + 1) * max_compressions, max_compressions, device=device, dtype=torch.int32 + ) + return cu_seq_lens, cu_outputs + + +def prepare_prefill_metadata( + kv_lens: torch.Tensor, + start_pos: torch.Tensor, + compress_ratio: int, + head_dim: int, + device: torch.device = None, +): + """Prepare prefill metadata for tests.""" + batch_size = kv_lens.shape[0] + if device is None: + device = kv_lens.device + + if start_pos is None: + start_pos = torch.zeros(batch_size, device=device, dtype=torch.int32) + + seq_lens = kv_lens - start_pos + num_outputs_per_batch = torch.clamp( + kv_lens // compress_ratio - start_pos // compress_ratio, min=1 + ) + + cu_seq_lens = torch.zeros(batch_size + 1, device=device, dtype=torch.int32) + cu_seq_lens[1:] = torch.cumsum(seq_lens, dim=0) + + cu_outputs = torch.zeros(batch_size + 1, device=device, dtype=torch.int32) + cu_outputs[1:] = torch.cumsum(num_outputs_per_batch, dim=0) + + return cu_seq_lens, cu_outputs + + +def create_paged_cache(batch_size, seqlen, compress_ratio, head_dim, overlap, page_size=4): + """Create paged cache tensors.""" + coff = 2 if overlap else 1 + state_dim = coff * head_dim + total_positions = seqlen + (compress_ratio if overlap else 0) + max_blocks = (total_positions + page_size - 1) // page_size + num_blocks = batch_size * max_blocks + + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + return paged_kv, paged_score, block_table, page_size, max_blocks + + +def pack_prefill_inputs(kv_list, score_list): + """Pack variable-length inputs into [m, 2*state_dim] format.""" + seq_lens = torch.tensor([kv.shape[0] for kv in kv_list], device="cuda", dtype=torch.int32) + kv_score = fuse_kv_score(torch.cat(kv_list, dim=0), torch.cat(score_list, dim=0)) + return kv_score, seq_lens + + +# ============================================================================ +# Correctness Tests +# ============================================================================ + +PREFILL_CONFIGS = [ + # Overlap mode (ratio=4) + pytest.param(4, 32, 4, 128, True, id="overlap_large_head_dim"), + pytest.param(1, 20, 4, 512, True, id="overlap_hd512_5chunks"), + # Basic mode (ratio=128) + pytest.param(1, 128, 128, 128, False, id="basic_hd128_eq_ratio"), + pytest.param(1, 256, 128, 512, False, id="basic_hd512_2chunks"), +] + + +@pytest.mark.parametrize("batch_size,seqlen,compress_ratio,head_dim,overlap", PREFILL_CONFIGS) +def test_prefill_corner_cases(batch_size, seqlen, compress_ratio, head_dim, overlap): + """Test prefill kernel corner cases.""" + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + + kv = torch.randn(batch_size, seqlen, state_dim, device="cuda") + score = torch.randn(batch_size, seqlen, state_dim, device="cuda") + ape = torch.randn(compress_ratio, state_dim, device="cuda") + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + paged_kv, paged_score, block_table, page_size, _ = create_paged_cache( + batch_size, seqlen, compress_ratio, head_dim, overlap + ) + + out_py = run_pytorch_prefill_reference( + kv.clone(), + score.clone(), + ape, + kv_state_py, + score_state_py, + compress_ratio, + head_dim, + overlap, + ) + + kv_score = fuse_kv_score(kv.view(-1, state_dim), score.view(-1, state_dim)) + kv_lens = torch.full((batch_size,), seqlen, device="cuda", dtype=torch.int32) + start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens, start_pos, compress_ratio, head_dim, kv_score.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score.device, torch.bfloat16 + ) + seq_lens = kv_lens - start_pos + num_outputs_per_batch = torch.clamp(seq_lens // compress_ratio, min=1) + max_outputs = num_outputs_per_batch.max().item() + + prefill_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + num_chunks = seqlen // compress_ratio + if out_py is None or num_chunks == 0: + pass # No compression expected + elif kv_comp.numel() == 0: + pytest.fail("cuTile returned empty output but PyTorch returned valid output") + else: + out_reshaped = kv_comp.view(batch_size, num_chunks, head_dim) + assert torch.allclose(out_py.to(kv_comp.dtype), out_reshaped, rtol=2e-3, atol=5e-3), ( + f"Output mismatch: max diff = {(out_py.to(kv_comp.dtype) - out_reshaped).abs().max():.6f}" + ) + + +DECODE_CONFIGS = [ + pytest.param(1, 4, 128, True, 12, id="overlap_large_head_dim"), + pytest.param(1, 4, 512, True, 8, id="overlap_hd512"), + pytest.param(1, 128, 128, False, 256, id="basic_hd128_2compressions"), + pytest.param(1, 128, 512, False, 256, id="basic_hd512_2compressions"), +] + + +@pytest.mark.parametrize("batch_size,compress_ratio,head_dim,overlap,num_steps", DECODE_CONFIGS) +def test_decode_corner_cases(batch_size, compress_ratio, head_dim, overlap, num_steps): + """Test decode kernel corner cases.""" + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + page_size = 8 + max_blocks = (num_steps + compress_ratio + page_size - 1) // page_size + + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + + # Pre-compute decode metadata + cu_seq_lens, cu_outputs = prepare_decode_metadata( + batch_size, compress_ratio, head_dim, torch.device("cuda"), next_n=1 + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, torch.device("cuda"), torch.bfloat16 + ) + + # Pre-fill for overlap mode (all batches) + if overlap: + init_kv = torch.randn(compress_ratio, state_dim, device="cuda") + init_score = torch.randn(compress_ratio, state_dim, device="cuda") + kv_state_py[:, :compress_ratio] = init_kv + score_state_py[:, :compress_ratio] = init_score + for b in range(batch_size): + for r in range(compress_ratio): + log_block, offset = r // page_size, r % page_size + phys_block = block_table[b, log_block].item() + paged_kv[phys_block, offset] = init_kv[r] + paged_score[phys_block, offset] = init_score[r] + + for step in range(num_steps): + new_kv = torch.randn(batch_size, state_dim, device="cuda") + new_score = torch.randn(batch_size, state_dim, device="cuda") + + # For overlap mode, account for initial compress_ratio tokens in cache + # token_idx is the absolute position: compress_ratio + step for overlap, step for basic + token_idx = (compress_ratio + step) if overlap else step + total_tokens = token_idx + 1 + + # PyTorch reference + out_py = run_pytorch_reference( + new_kv.unsqueeze(1), + new_score.unsqueeze(1), + ape, + kv_state_py, + score_state_py, + token_idx, + compress_ratio, + head_dim, + overlap, + ) + + # CUDA decode kernel + kv_score = fuse_kv_score(new_kv, new_score) # [bsz, 2*state_dim] + kv_lens = torch.full((batch_size,), total_tokens, device="cuda", dtype=torch.int32) + start_pos = torch.full((batch_size,), token_idx, device="cuda", dtype=torch.int32) + + decode_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + block_table, + compress_ratio, + head_dim, + page_size, + next_n=1, + ) + + should_compress = (step + 1) % compress_ratio == 0 + if should_compress: + if out_py is not None: + for b in range(batch_size): + out_idx = cu_outputs[b].item() + diff = out_py[b, 0, :head_dim].to(kv_comp.dtype) - kv_comp[out_idx, :] + assert torch.allclose( + out_py[b, 0, :head_dim].to(kv_comp.dtype), + kv_comp[out_idx, :], + rtol=1e-2, + atol=1e-3, + ), f"Step {step}, Batch {b}: mismatch diff={(diff).abs().max():.6f}" + + +STATE_UPDATE_CONFIGS = [ + pytest.param(1, 12, 4, 128, True, id="overlap_hd128_3chunks"), + pytest.param(1, 13, 4, 512, True, id="overlap_hd512_3chunks_1remainder"), + pytest.param(1, 256, 128, 128, False, id="basic_hd128_2chunks"), + pytest.param(1, 260, 128, 512, False, id="basic_hd512_2chunks_4remainder"), +] + + +@pytest.mark.parametrize("batch_size,seqlen,compress_ratio,head_dim,overlap", STATE_UPDATE_CONFIGS) +def test_prefill_state_update(batch_size, seqlen, compress_ratio, head_dim, overlap): + """Verify prefill state updates match reference — all chunks, not just the last one. + + Checks two things: + 1. All full-chunk positions [0, cutoff) are written to paged cache with + the correct kv and score+APE values (block reuse requirement: this + section would FAIL against the old kernel that only wrote the last chunk). + 2. The last-chunk + remainder state matches the PyTorch reference state + (continuity requirement for decode following prefill). + """ + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + + kv = torch.arange(batch_size * seqlen * state_dim, dtype=torch.float32, device="cuda").reshape( + batch_size, seqlen, state_dim + ) + score = torch.arange( + batch_size * seqlen * state_dim, dtype=torch.float32, device="cuda" + ).reshape(batch_size, seqlen, state_dim) + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + + paged_kv, paged_score, block_table, page_size, max_blocks = create_paged_cache( + batch_size, seqlen, compress_ratio, head_dim, overlap + ) + + # Run PyTorch reference (sets last-chunk + remainder in kv_state_py/score_state_py) + _ = run_pytorch_prefill_reference( + kv.clone(), + score.clone(), + ape, + kv_state_py, + score_state_py, + compress_ratio, + head_dim, + overlap, + ) + + # Run cuTile kernel + kv_packed = kv.view(-1, state_dim) + score_packed = score.view(-1, state_dim) + kv_score = fuse_kv_score(kv_packed, score_packed) + kv_lens = torch.full((batch_size,), seqlen, device="cuda", dtype=torch.int32) + start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens, start_pos, compress_ratio, head_dim, kv_score.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score.device, torch.bfloat16 + ) + seq_lens = kv_lens - start_pos + num_outputs_per_batch = torch.clamp(seq_lens // compress_ratio, min=1) + max_outputs = num_outputs_per_batch.max().item() + + prefill_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + remainder = seqlen % compress_ratio + cutoff = seqlen - remainder + offset = compress_ratio if overlap else 0 + + # --- Check 1: ALL full-chunk positions [0, cutoff) --- + # Expected: paged_kv[phys, blk_off] == kv[b, p] + # paged_score[phys, blk_off] == score[b, p] + ape[p % compress_ratio] + # This check FAILS against the old kernel (which only wrote the last chunk), + # proving these positions are now correctly populated for block reuse. + for b in range(batch_size): + for p in range(cutoff): + log_block = p // page_size + block_offset = p % page_size + phys_block = block_table[b, log_block].item() + r = p % compress_ratio + + got_kv = paged_kv[phys_block, block_offset].float() + exp_kv = kv[b, p].float() + assert torch.allclose(got_kv, exp_kv, atol=1e-5), ( + f"batch={b} pos={p}: paged_kv mismatch max_diff={(got_kv - exp_kv).abs().max():.6f}" + ) + + got_sc = paged_score[phys_block, block_offset].float() + exp_sc = (score[b, p] + ape[r]).float() + assert torch.allclose(got_sc, exp_sc, atol=1e-5), ( + f"batch={b} pos={p}: paged_score mismatch " + f"max_diff={(got_sc - exp_sc).abs().max():.6f}" + ) + + # --- Check 2: last-chunk + remainder match the PyTorch reference state --- + # The PyTorch reference stores the "virtual" state needed for decode continuation. + # Verify the kernel's paged cache agrees for those positions. + kv_state_kernel = torch.zeros_like(kv_state_py) + score_state_kernel = torch.full_like(score_state_py, float("-inf")) + + for b in range(batch_size): + if overlap and cutoff >= compress_ratio: + for r in range(compress_ratio): + abs_pos = cutoff - compress_ratio + r + log_block = abs_pos // page_size + block_offset = abs_pos % page_size + phys_block = block_table[b, log_block].item() + kv_state_kernel[b, r] = paged_kv[phys_block, block_offset] + score_state_kernel[b, r] = paged_score[phys_block, block_offset] + + if remainder > 0: + for r in range(remainder): + abs_pos = cutoff + r + log_block = abs_pos // page_size + block_offset = abs_pos % page_size + phys_block = block_table[b, log_block].item() + kv_state_kernel[b, offset + r] = paged_kv[phys_block, block_offset] + score_state_kernel[b, offset + r] = paged_score[phys_block, block_offset] + + assert torch.allclose(kv_state_py, kv_state_kernel, atol=1e-5), ( + f"KV state mismatch: {(kv_state_py - kv_state_kernel).abs().max():.6f}" + ) + assert torch.allclose(score_state_py, score_state_kernel, atol=1e-5), ( + f"Score state mismatch: {(score_state_py - score_state_kernel).abs().max():.6f}" + ) + + +PREFILL_ALL_BLOCKS_CONFIGS = [ + pytest.param(1, 12, 4, 128, True, id="overlap_hd128_3chunks"), + pytest.param(1, 20, 4, 512, True, id="overlap_hd512_5chunks"), + pytest.param(1, 256, 128, 128, False, id="basic_hd128_2chunks"), + pytest.param(1, 260, 128, 512, False, id="basic_hd512_2chunks_4remainder"), +] + + +@pytest.mark.parametrize( + "batch_size,seqlen,compress_ratio,head_dim,overlap", PREFILL_ALL_BLOCKS_CONFIGS +) +def test_prefill_paged_state_all_blocks(batch_size, seqlen, compress_ratio, head_dim, overlap): + """Verify prefill writes ALL token positions to paged cache (block reuse requirement). + + Before the fix, only the last full chunk + remainder were written to paged cache. + Every chunk block must now be written so reused requests can read any prefix slice. + + Checks: for every full-chunk position p in [0, cutoff): + paged_kv[phys, blk_off] == kv[b, p] + paged_score[phys, blk_off] == score[b, p] + ape[p % compress_ratio] + """ + coff = 2 if overlap else 1 + state_dim = coff * head_dim + + kv = torch.randn(batch_size, seqlen, state_dim, device="cuda") + score = torch.randn(batch_size, seqlen, state_dim, device="cuda") + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + paged_kv, paged_score, block_table, page_size, _ = create_paged_cache( + batch_size, seqlen, compress_ratio, head_dim, overlap + ) + + kv_score = fuse_kv_score(kv.view(-1, state_dim), score.view(-1, state_dim)) + kv_lens = torch.full((batch_size,), seqlen, device="cuda", dtype=torch.int32) + start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens, start_pos, compress_ratio, head_dim, kv_score.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score.device, torch.bfloat16 + ) + max_outputs = torch.clamp((kv_lens - start_pos) // compress_ratio, min=1).max().item() + + prefill_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + # Verify every full-chunk position is correctly written + cutoff = (seqlen // compress_ratio) * compress_ratio + for b in range(batch_size): + for p in range(cutoff): + log_blk = p // page_size + blk_off = p % page_size + phys = block_table[b, log_blk].item() + r = p % compress_ratio # APE index: position within window + + got_kv = paged_kv[phys, blk_off].float() + exp_kv = kv[b, p].float() + assert torch.allclose(got_kv, exp_kv, atol=1e-5), ( + f"batch={b} pos={p}: paged_kv mismatch max_diff={(got_kv - exp_kv).abs().max():.6f}" + ) + + got_sc = paged_score[phys, blk_off].float() + exp_sc = (score[b, p] + ape[r]).float() + assert torch.allclose(got_sc, exp_sc, atol=1e-5), ( + f"batch={b} pos={p}: paged_score mismatch " + f"max_diff={(got_sc - exp_sc).abs().max():.6f}" + ) + + +BLOCK_REUSE_CONFIGS = [ + # (batch_size, prefix_len, compress_ratio, head_dim, overlap, reuse_at) + # reuse_at must be on a chunk boundary (reuse_at % compress_ratio == 0) + # so the decode uses an EARLIER chunk as overlap window, not the last one. + pytest.param(1, 12, 4, 128, True, 4, id="overlap_hd128_reuse_after_chunk0"), + pytest.param(1, 12, 4, 512, True, 8, id="overlap_hd512_reuse_after_chunk1"), + pytest.param(1, 512, 128, 128, False, 128, id="basic_hd128_reuse_after_chunk0"), + pytest.param(1, 512, 128, 512, False, 256, id="basic_hd512_reuse_after_chunk1"), +] + + +@pytest.mark.parametrize( + "batch_size,prefix_len,compress_ratio,head_dim,overlap,reuse_at", + BLOCK_REUSE_CONFIGS, +) +def test_prefill_block_reuse_decode( + batch_size, prefix_len, compress_ratio, head_dim, overlap, reuse_at +): + """End-to-end block reuse: decode from mid-prefix using paged cache written during prefill. + + Simulates the block reuse scenario: + 1. Request A prefills the full prefix [0 .. prefix_len-1], populating paged cache. + 2. Request B reuses A's paged blocks: decodes starting at position reuse_at, + reading the overlap window from earlier blocks (positions < reuse_at). + + With the old code only the last chunk was in paged cache, so decode triggered at + reuse_at + compress_ratio - 1 would read garbage for the overlap window. + With the fix, all blocks are present and decode outputs must match the reference. + """ + assert reuse_at % compress_ratio == 0, "reuse_at must be on a chunk boundary" + assert reuse_at < prefix_len, "reuse_at must be within the prefix" + + coff = 2 if overlap else 1 + state_len = coff * compress_ratio + state_dim = coff * head_dim + page_size = 8 + decode_steps = compress_ratio * 2 # enough to trigger two compressions + total_len = prefix_len + decode_steps + max_blocks = (total_len + (compress_ratio if overlap else 0) + page_size - 1) // page_size + + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + # --- Step 1: Prefill the full prefix --- + kv_prefill = torch.randn(batch_size, prefix_len, state_dim, device="cuda") + score_prefill = torch.randn(batch_size, prefix_len, state_dim, device="cuda") + + kv_score_prefill = fuse_kv_score( + kv_prefill.view(-1, state_dim), score_prefill.view(-1, state_dim) + ) + kv_lens_pre = torch.full((batch_size,), prefix_len, device="cuda", dtype=torch.int32) + start_pos_pre = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens_pre, start_pos_pre, compress_ratio, head_dim, "cuda" + ) + kv_comp_pre = prepare_compress_output(cu_outputs, batch_size, head_dim, "cuda", torch.bfloat16) + max_outputs = torch.clamp((kv_lens_pre - start_pos_pre) // compress_ratio, min=1).max().item() + + prefill_kernel( + kv_score_prefill, + ape, + kv_lens_pre, + start_pos_pre, + cu_seq_lens, + cu_outputs, + kv_comp_pre, + paged_kv, + paged_score, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + # --- Step 2: Build PyTorch reference state at reuse_at --- + # The state at reuse_at is: the overlap window = tokens [reuse_at-CR .. reuse_at-1], + # and any remainder tokens within [reuse_at-CR .. reuse_at). + # Since reuse_at is on a chunk boundary, remainder = 0. + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + + if overlap and reuse_at >= compress_ratio: + # Overlap window = last full chunk before reuse_at: [reuse_at-CR .. reuse_at-1] + kv_state_py[:, :compress_ratio] = kv_prefill[:, reuse_at - compress_ratio : reuse_at] + score_state_py[:, :compress_ratio] = score_prefill[ + :, reuse_at - compress_ratio : reuse_at + ] + ape.unsqueeze(0) + # (non-overlap, or overlap with reuse_at < compress_ratio: state stays zero/-inf) + + # --- Step 3: Decode from reuse_at and verify outputs match reference --- + cu_seq_lens_dec, cu_outputs_dec = prepare_decode_metadata( + batch_size, compress_ratio, head_dim, torch.device("cuda"), next_n=1 + ) + kv_comp_dec = prepare_compress_output( + cu_outputs_dec, batch_size, head_dim, torch.device("cuda"), torch.bfloat16 + ) + + for i in range(decode_steps): + token_idx = reuse_at + i + total_tokens = token_idx + 1 + + new_kv = torch.randn(batch_size, state_dim, device="cuda") + new_score = torch.randn(batch_size, state_dim, device="cuda") + + # PyTorch reference (uses ground-truth kv_state_py) + out_py = run_pytorch_reference( + new_kv.unsqueeze(1), + new_score.unsqueeze(1), + ape, + kv_state_py, + score_state_py, + token_idx, + compress_ratio, + head_dim, + overlap, + ) + + # CUDA decode (reads overlap window from paged cache written during prefill) + kv_score = fuse_kv_score(new_kv, new_score) + kv_lens = torch.full((batch_size,), total_tokens, device="cuda", dtype=torch.int32) + start_pos = torch.full((batch_size,), token_idx, device="cuda", dtype=torch.int32) + + decode_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens_dec, + cu_outputs_dec, + kv_comp_dec, + paged_kv, + paged_score, + block_table, + block_table, + compress_ratio, + head_dim, + page_size, + next_n=1, + ) + + should_compress = (token_idx + 1) % compress_ratio == 0 + if should_compress and out_py is not None: + for b in range(batch_size): + out_idx = cu_outputs_dec[b].item() + diff = out_py[b, 0, :head_dim].to(kv_comp_dec.dtype) - kv_comp_dec[out_idx, :] + assert torch.allclose( + out_py[b, 0, :head_dim].to(kv_comp_dec.dtype), + kv_comp_dec[out_idx, :], + rtol=1e-2, + atol=1e-3, + ), ( + f"Block reuse decode step {i} (token_idx={token_idx}), " + f"batch {b}: diff={(diff).abs().max():.6f}" + ) + + +PREFILL_VARLEN_CONFIGS = [ + pytest.param([8, 12, 4, 16], 4, 128, True, id="varlen_hd128_overlap"), + pytest.param([128, 256, 64, 192], 128, 128, False, id="varlen_hd128_basic"), + pytest.param([8, 12, 4, 16], 4, 512, True, id="varlen_hd512_overlap"), + pytest.param([128, 256, 64, 192], 128, 512, False, id="varlen_hd512_basic"), +] + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.parametrize("seq_lens_list,compress_ratio,head_dim,overlap", PREFILL_VARLEN_CONFIGS) +def test_prefill_varlen(seq_lens_list, compress_ratio, head_dim, overlap): + """Test prefill with variable-length sequences.""" + batch_size = len(seq_lens_list) + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + + # Create variable-length packed input + total_tokens = sum(seq_lens_list) + kv_packed = torch.randn(total_tokens, state_dim, device="cuda") + score_packed = torch.randn(total_tokens, state_dim, device="cuda") + kv_score = fuse_kv_score(kv_packed, score_packed) + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + # Per-batch metadata + kv_lens = torch.tensor(seq_lens_list, device="cuda", dtype=torch.int32) + start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + + # PyTorch reference state + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + + # Paged cache setup + page_size = 4 + max_seqlen = max(seq_lens_list) + total_positions = max_seqlen + (compress_ratio if overlap else 0) + max_blocks = (total_positions + page_size - 1) // page_size + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + # PyTorch reference (varlen) + out_py = run_pytorch_prefill_reference_varlen( + kv_score.clone(), + ape, + kv_lens, + start_pos, + compress_ratio, + head_dim, + overlap, + kv_state_py, + score_state_py, + ) + + # cuTile kernel + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens, start_pos, compress_ratio, head_dim, kv_score.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score.device, torch.bfloat16 + ) + seq_lens = kv_lens - start_pos + num_outputs_per_batch = torch.clamp(seq_lens // compress_ratio, min=1) + max_outputs = num_outputs_per_batch.max().item() + + prefill_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + # Check output - extract only valid outputs from packed output + # cu_outputs uses min=1 for CUDA graph, but kernel only writes where seqlen >= ratio + actual_outputs_per_batch = [s // compress_ratio for s in seq_lens_list] + total_actual_outputs = sum(actual_outputs_per_batch) + + if total_actual_outputs == 0: + # No compression expected + assert out_py.numel() == 0, "Expected empty output when no compression" + elif out_py.numel() == 0: + pytest.fail("PyTorch returned empty output but expected valid output") + else: + # Extract valid outputs from kernel's packed output + # cu_outputs[b] gives offset for batch b, but includes min=1 padding + valid_outputs = [] + offset = 0 + for b, actual_count in enumerate(actual_outputs_per_batch): + # cu_outputs uses clamped count, so we need to compute actual offset + clamped_count = max(seq_lens_list[b] // compress_ratio, 1) + if actual_count > 0: + valid_outputs.append(kv_comp[offset : offset + actual_count]) + offset += clamped_count + + if valid_outputs: + out_kernel_valid = torch.cat(valid_outputs, dim=0) + assert torch.allclose( + out_py.to(out_kernel_valid.dtype), out_kernel_valid, rtol=2e-3, atol=5e-3 + ), ( + f"Output mismatch: max diff = {(out_py.to(out_kernel_valid.dtype) - out_kernel_valid).abs().max():.6f}" + ) + else: + assert out_py.numel() == 0, "Expected empty output" + + +PREFILL_DECODE_CONFIGS = [ + pytest.param(1, 20, 4, 128, True, 12, id="overlap_hd128_prefill20_decode12"), + pytest.param(1, 256, 128, 128, False, 128, id="basic_hd128_prefill256_decode128"), + pytest.param(1, 20, 4, 512, True, 12, id="overlap_hd512_prefill20_decode12"), + pytest.param(1, 256, 128, 512, False, 128, id="basic_hd512_prefill256_decode128"), +] + + +@pytest.mark.parametrize( + "batch_size,prefill_len,compress_ratio,head_dim,overlap,decode_steps", PREFILL_DECODE_CONFIGS +) +def test_prefill_then_decode( + batch_size, prefill_len, compress_ratio, head_dim, overlap, decode_steps +): + """Test prefill followed by decode (simulates inference).""" + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + page_size = 8 + total_len = prefill_len + decode_steps + max_blocks = (total_len + (compress_ratio if overlap else 0) + page_size - 1) // page_size + + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + # PyTorch reference state + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + + # Paged cache + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + # 1. Prefill + kv_prefill = torch.randn(batch_size, prefill_len, state_dim, device="cuda") + score_prefill = torch.randn(batch_size, prefill_len, state_dim, device="cuda") + + # PyTorch prefill + out_py_prefill = run_pytorch_prefill_reference( + kv_prefill.clone(), + score_prefill.clone(), + ape, + kv_state_py, + score_state_py, + compress_ratio, + head_dim, + overlap, + ) + + # cuTile prefill + kv_packed = kv_prefill.view(-1, state_dim) + score_packed = score_prefill.view(-1, state_dim) + kv_score_prefill = fuse_kv_score(kv_packed, score_packed) + kv_lens_prefill = torch.full((batch_size,), prefill_len, device="cuda", dtype=torch.int32) + start_pos_prefill = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + + # Pre-compute prefill metadata + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens_prefill, start_pos_prefill, compress_ratio, head_dim, kv_score_prefill.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score_prefill.device, torch.bfloat16 + ) + seq_lens = kv_lens_prefill - start_pos_prefill + num_outputs_per_batch = torch.clamp(seq_lens // compress_ratio, min=1) + max_outputs = num_outputs_per_batch.max().item() + + prefill_kernel( + kv_score_prefill, + ape, + kv_lens_prefill, + start_pos_prefill, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + if out_py_prefill is not None and kv_comp.numel() > 0: + num_chunks = prefill_len // compress_ratio + out_reshaped = kv_comp.view(batch_size, num_chunks, head_dim) + assert torch.allclose( + out_py_prefill.to(kv_comp.dtype), out_reshaped, rtol=2e-3, atol=5e-3 + ), ( + f"Prefill output mismatch: {(out_py_prefill.to(kv_comp.dtype) - out_reshaped).abs().max():.6f}" + ) + + # 2. Decode (continue from where prefill left off) + decode_start = (prefill_len // compress_ratio) * compress_ratio + remainder = prefill_len % compress_ratio + + # Pre-compute decode metadata + cu_seq_lens_decode, cu_outputs_decode = prepare_decode_metadata( + batch_size, compress_ratio, head_dim, torch.device("cuda"), next_n=1 + ) + kv_comp_decode = prepare_compress_output( + cu_outputs_decode, batch_size, head_dim, torch.device("cuda"), torch.bfloat16 + ) + + for i in range(decode_steps): + step = decode_start + remainder + i # Continue from where prefill left off + + new_kv = torch.randn(batch_size, state_dim, device="cuda") + new_score = torch.randn(batch_size, state_dim, device="cuda") + + # Both prefill and decode use absolute token positions in the state cache. + # The prefill kernel writes at [cutoff-ratio:cutoff] and [cutoff:cutoff+remainder], + # so the decode kernel continues from position step = prefill_len + i. + token_idx = step + total_tokens = token_idx + 1 + + # PyTorch reference + out_py = run_pytorch_reference( + new_kv.unsqueeze(1), + new_score.unsqueeze(1), + ape, + kv_state_py, + score_state_py, + token_idx, + compress_ratio, + head_dim, + overlap, + ) + + # CUDA decode + kv_score = fuse_kv_score(new_kv, new_score) + kv_lens = torch.full((batch_size,), total_tokens, device="cuda", dtype=torch.int32) + start_pos = torch.full((batch_size,), token_idx, device="cuda", dtype=torch.int32) + + decode_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens_decode, + cu_outputs_decode, + kv_comp_decode, + paged_kv, + paged_score, + block_table, + block_table, + compress_ratio, + head_dim, + page_size, + next_n=1, + ) + + should_compress = (token_idx + 1) % compress_ratio == 0 + if should_compress and out_py is not None: + for b in range(batch_size): + out_idx = cu_outputs_decode[b].item() + diff = out_py[b, 0, :head_dim].to(kv_comp_decode.dtype) - kv_comp_decode[out_idx, :] + assert torch.allclose( + out_py[b, 0, :head_dim].to(kv_comp_decode.dtype), + kv_comp_decode[out_idx, :], + rtol=2e-3, + atol=5e-3, + ), ( + f"Decode step {i} (token_idx={token_idx}), Batch {b}: mismatch diff={(diff).abs().max():.6f}" + ) + + +# MTP (Multi-Token Prediction) Tests +MTP_CONFIGS = [ + pytest.param(1, 4, 128, True, 4, id="overlap_hd128_next4"), + pytest.param(1, 4, 512, True, 3, id="overlap_hd512_next3"), + pytest.param(2, 128, 128, False, 4, id="basic_hd128_multi_batch_next4"), + pytest.param(1, 128, 512, False, 4, id="basic_hd512_next4"), +] + + +@pytest.mark.parametrize("batch_size,compress_ratio,head_dim,overlap,next_n", MTP_CONFIGS) +def test_decode_mtp(batch_size, compress_ratio, head_dim, overlap, next_n): + """Test decode with multiple tokens per request (MTP).""" + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + page_size = 8 + num_steps = compress_ratio * 4 # Enough to trigger multiple compressions + max_blocks = (num_steps + compress_ratio + page_size - 1) // page_size + + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + + # Pre-fill for overlap mode + if overlap: + init_kv = torch.randn(compress_ratio, state_dim, device="cuda") + init_score = torch.randn(compress_ratio, state_dim, device="cuda") + kv_state_py[:, :compress_ratio] = init_kv + score_state_py[:, :compress_ratio] = init_score + for b in range(batch_size): + for r in range(compress_ratio): + log_block, offset = r // page_size, r % page_size + phys_block = block_table[b, log_block].item() + paged_kv[phys_block, offset] = init_kv[r] + paged_score[phys_block, offset] = init_score[r] + + # Process multiple tokens at once + # For overlap mode, account for initial compress_ratio tokens in cache + base_token_idx = compress_ratio if overlap else 0 + + step = 0 + while step < num_steps: + actual_n = min(next_n, num_steps - step) + + # Generate next_n tokens + new_kv = torch.randn(batch_size * actual_n, state_dim, device="cuda") + new_score = torch.randn(batch_size * actual_n, state_dim, device="cuda") + + # PyTorch reference: process one token at a time + py_outputs = [] + for t in range(actual_n): + token_idx = base_token_idx + step + t + kv_t = new_kv[t::actual_n].unsqueeze(1) # Get token t for all batches + score_t = new_score[t::actual_n].unsqueeze(1) + out_py = run_pytorch_reference( + kv_t, + score_t, + ape, + kv_state_py, + score_state_py, + token_idx, + compress_ratio, + head_dim, + overlap, + ) + if out_py is not None: + py_outputs.append((token_idx, out_py)) + + # CUDA decode: process all tokens at once + kv_score = fuse_kv_score(new_kv, new_score) + abs_start = base_token_idx + step + kv_lens = torch.full((batch_size,), abs_start + actual_n, device="cuda", dtype=torch.int32) + start_pos = torch.full((batch_size,), abs_start, device="cuda", dtype=torch.int32) + + # Compute decode metadata for actual_n (may differ from next_n at end of loop) + cu_seq_lens, cu_outputs = prepare_decode_metadata( + batch_size, compress_ratio, head_dim, torch.device("cuda"), next_n=actual_n + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, torch.device("cuda"), torch.bfloat16 + ) + + decode_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + block_table, + compress_ratio, + head_dim, + page_size, + next_n=actual_n, + ) + + # Verify outputs match (packed [total_outputs, head_dim] format) + if len(py_outputs) > 0: + for i, (token_idx, out_py) in enumerate(py_outputs): + for b in range(batch_size): + out_idx = cu_outputs[b].item() + i + diff = out_py[b, 0, :head_dim].to(kv_comp.dtype) - kv_comp[out_idx, :] + assert torch.allclose( + out_py[b, 0, :head_dim].to(kv_comp.dtype), + kv_comp[out_idx, :], + rtol=2e-3, + atol=2e-3, + ), f"Token {token_idx}, Batch {b}: mismatch diff={(diff).abs().max():.6f}" + + step += actual_n + + +CHUNKED_PREFILL_CONFIGS = [ + # (compress_ratio, head_dim, overlap, batch_size, start_pos, new_seqlen) + # overlap=True, aligned start_pos + pytest.param(4, 512, True, 1, 4, 4, id="overlap_sp4_seq4"), + pytest.param(4, 512, True, 1, 4, 8, id="overlap_sp4_seq8"), + pytest.param(4, 512, True, 1, 8, 8, id="overlap_sp8_seq8"), + pytest.param(4, 512, True, 2, 4, 4, id="overlap_batch2_sp4_seq4"), + pytest.param(4, 512, True, 1, 4, 5, id="overlap_sp4_seq5_remainder"), + # overlap=True, unaligned start_pos (sp % ratio != 0) + pytest.param(4, 512, True, 1, 6, 4, id="overlap_sp6_seq4_unaligned"), + pytest.param(4, 512, True, 1, 6, 5, id="overlap_sp6_seq5_unaligned"), + pytest.param(4, 512, True, 1, 5, 7, id="overlap_sp5_seq7_unaligned"), + pytest.param(4, 512, True, 1, 3, 9, id="overlap_sp3_seq9_unaligned"), + # overlap=False, aligned start_pos + pytest.param(128, 128, False, 1, 128, 128, id="nonoverlap_sp128_seq128"), + # overlap=False, unaligned start_pos + pytest.param(128, 128, False, 1, 50, 206, id="nonoverlap_sp50_seq206_unaligned"), + pytest.param(128, 128, False, 1, 5, 251, id="nonoverlap_sp5_seq251_2windows"), +] + + +@pytest.mark.parametrize( + "compress_ratio,head_dim,overlap,batch_size,start_pos_val,new_seqlen", + CHUNKED_PREFILL_CONFIGS, +) +def test_chunked_prefill(compress_ratio, head_dim, overlap, batch_size, start_pos_val, new_seqlen): + """End-to-end chunked prefill: reference (one-shot full prefill) vs + kernel (initial prefill + chunked prefill on same paged cache). + + Validates that kernel Phase 1 correctly persists state and Phase 2 + correctly reads it back in a subsequent call. + """ + device = torch.device("cuda") + coff = 2 if overlap else 1 + state_dim = coff * head_dim + ratio = compress_ratio + page_size = 32 + total_seqlen = start_pos_val + new_seqlen + + total_ref_outputs = total_seqlen // ratio + if total_ref_outputs == 0: + return + + # Shared paged cache (zeros — no prior state) + total_positions = total_seqlen + (ratio if overlap else 0) + max_blocks = (total_positions + page_size - 1) // page_size + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device=device) + paged_score = torch.zeros(num_blocks, page_size, state_dim, device=device) + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32).view( + batch_size, max_blocks + ) + ape = torch.randn(ratio, state_dim, device=device) + + # Generate full sequence data + kv_full = torch.randn(batch_size, total_seqlen, state_dim, device=device) + score_full = torch.randn(batch_size, total_seqlen, state_dim, device=device) + + # ---- Reference: one-shot full prefill (sp=0) ---- + state_len = coff * ratio + kv_state_ref = torch.zeros(batch_size, state_len, state_dim, device=device) + score_state_ref = torch.full((batch_size, state_len, state_dim), float("-inf"), device=device) + ref_out = run_pytorch_prefill_reference( + kv_full.clone(), + score_full.clone(), + ape, + kv_state_ref, + score_state_ref, + ratio, + head_dim, + overlap, + ) + assert ref_out is not None, "Reference should produce output" + + # ---- Kernel call 1: initial prefill (sp=0, kv_len=start_pos_val) ---- + actual_num_out_1 = start_pos_val // ratio + kv_comp_1 = torch.empty(0, head_dim, device=device, dtype=torch.bfloat16) + + if start_pos_val > 0: + kv_score_1 = fuse_kv_score( + kv_full[:, :start_pos_val].reshape(-1, state_dim), + score_full[:, :start_pos_val].reshape(-1, state_dim), + ) + kv_lens_1 = torch.full((batch_size,), start_pos_val, device=device, dtype=torch.int32) + start_pos_1 = torch.zeros(batch_size, device=device, dtype=torch.int32) + cu_seq_lens_1, cu_outputs_1 = prepare_prefill_metadata( + kv_lens_1, + start_pos_1, + ratio, + head_dim, + device, + ) + kv_comp_1 = prepare_compress_output( + cu_outputs_1, + batch_size, + head_dim, + device, + torch.bfloat16, + ) + max_outputs_1 = max(actual_num_out_1, 1) + prefill_kernel( + kv_score_1, + ape, + kv_lens_1, + start_pos_1, + cu_seq_lens_1, + cu_outputs_1, + kv_comp_1, + paged_kv, + paged_score, + block_table, + max_outputs_1, + block_table, + ratio, + head_dim, + page_size, + ) + + # ---- Kernel call 2: chunked prefill (sp=start_pos_val, kv_len=total_seqlen) ---- + actual_num_out_2 = total_seqlen // ratio - start_pos_val // ratio + kv_score_2 = fuse_kv_score( + kv_full[:, start_pos_val:].reshape(-1, state_dim), + score_full[:, start_pos_val:].reshape(-1, state_dim), + ) + kv_lens_2 = torch.full((batch_size,), total_seqlen, device=device, dtype=torch.int32) + start_pos_2 = torch.full((batch_size,), start_pos_val, device=device, dtype=torch.int32) + cu_seq_lens_2, cu_outputs_2 = prepare_prefill_metadata( + kv_lens_2, + start_pos_2, + ratio, + head_dim, + device, + ) + kv_comp_2 = prepare_compress_output( + cu_outputs_2, + batch_size, + head_dim, + device, + torch.bfloat16, + ) + max_outputs_2 = max(actual_num_out_2, 1) + prefill_kernel( + kv_score_2, + ape, + kv_lens_2, + start_pos_2, + cu_seq_lens_2, + cu_outputs_2, + kv_comp_2, + paged_kv, + paged_score, + block_table, + max_outputs_2, + block_table, + ratio, + head_dim, + page_size, + ) + + # ---- Compare per batch: concat(call1_out, call2_out) vs reference ---- + for b in range(batch_size): + parts = [] + if start_pos_val > 0 and actual_num_out_1 > 0: + off_1 = cu_outputs_1[b].item() + parts.append(kv_comp_1[off_1 : off_1 + actual_num_out_1]) + if actual_num_out_2 > 0: + off_2 = cu_outputs_2[b].item() + parts.append(kv_comp_2[off_2 : off_2 + actual_num_out_2]) + + if not parts: + continue + kernel_out_b = torch.cat(parts, dim=0) + ref_out_b = ref_out[b, :total_ref_outputs, :head_dim].to(torch.bfloat16) + + assert torch.allclose(ref_out_b, kernel_out_b, rtol=2e-3, atol=5e-3), ( + f"batch={b} sp={start_pos_val} seqlen={new_seqlen} " + f"max_diff={(ref_out_b - kernel_out_b).abs().max():.6f}" + ) + + +# ============================================================================ +# Postprocess + Scatter Kernel Tests +# ============================================================================ + +try: + _HAS_POSTPROCESS_SCATTER = hasattr(torch.ops.trtllm, "compressor_postprocess_scatter") +except Exception: + _HAS_POSTPROCESS_SCATTER = False + + +def _scatter_reference_to_paged_cache( + kv: torch.Tensor, + num_comp_tokens: torch.Tensor, + cu_kv_comp: torch.Tensor, + kv_cache: torch.Tensor, + block_offsets: torch.Tensor, + tokens_per_block: int, + head_dim: int, +): + """Reference paged scatter without calling any CuTile kernel.""" + batch_size = num_comp_tokens.shape[0] + for b in range(batch_size): + num_tokens_b = num_comp_tokens[b].item() + src_start = cu_kv_comp[b].item() + for t in range(num_tokens_b): + blk = t // tokens_per_block + off = t % tokens_per_block + phys = block_offsets[b, blk].item() + dst = off * head_dim + kv_cache[phys, dst : dst + head_dim] = kv[src_start + t] + + +@pytest.mark.parametrize("rotate_activation", [True, False], ids=["rotate", "no_rotate"]) +@pytest.mark.parametrize( + "batch_size,num_tokens,head_dim,nope_dim,rope_dim,tokens_per_block", + [ + pytest.param(1, 16, 128, 64, 64, 32, id="b1_t16_hd128"), + pytest.param(1, 64, 128, 64, 64, 32, id="b1_t64_hd128"), + pytest.param(2, 32, 128, 64, 64, 32, id="b2_t32_hd128"), + pytest.param(1, 128, 512, 256, 256, 128, id="b1_t128_hd512"), + pytest.param(2, 64, 512, 256, 256, 128, id="b2_t64_hd512"), + pytest.param(4, 32, 512, 256, 256, 128, id="b4_t32_hd512"), + ], +) +@pytest.mark.skipif( + not _HAS_POSTPROCESS_SCATTER, reason="Postprocess/scatter CUDA ops not available" +) +def test_fused_postprocess_scatter( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block, rotate_activation +): + """Compare fused RMSNorm+RoPE+(optional Hadamard)+Scatter vs sequential unfused path. + + Ensures the fused CUDA kernel produces identical paged KV cache output as + the reference pipeline: RMSNorm -> RoPE -> (Hadamard) -> Scatter. + """ + torch.manual_seed(42) + device = "cuda" + + tokens_per_batch = num_tokens + total_tokens = batch_size * tokens_per_batch + + kv_comp = torch.randn(total_tokens, head_dim, device=device, dtype=torch.bfloat16) + rms_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) * 0.1 + 1.0 + rms_eps = 1e-5 + + max_pos = total_tokens + 64 + cos_sin_table = torch.randn(max_pos, 2, rope_dim // 2, device=device, dtype=torch.float32) + position_ids = torch.arange(total_tokens, device=device, dtype=torch.int32) + + max_comp_len = tokens_per_batch + 4 + max_blocks = (max_comp_len + tokens_per_block - 1) // tokens_per_block + num_blocks = batch_size * max_blocks + kv_cache_fused = torch.zeros( + num_blocks, tokens_per_block * head_dim, device=device, dtype=torch.bfloat16 + ) + kv_cache_ref = torch.zeros_like(kv_cache_fused) + + block_offsets = torch.zeros(batch_size, max_blocks, device=device, dtype=torch.int32) + for b in range(batch_size): + block_offsets[b] = torch.arange(b * max_blocks, (b + 1) * max_blocks, dtype=torch.int32) + + num_comp_tokens = torch.full((batch_size,), tokens_per_batch, device=device, dtype=torch.int32) + cu_kv_comp = torch.zeros(batch_size + 1, device=device, dtype=torch.int32) + cu_kv_comp[1:] = num_comp_tokens.cumsum(0) + start_pos = torch.zeros(batch_size, device=device, dtype=torch.int32) + + # --- Reference: unfused pipeline --- + x = _build_postprocess_reference( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + head_dim, + rotate_activation, + ) + + _scatter_reference_to_paged_cache( + x, + num_comp_tokens, + cu_kv_comp, + kv_cache_ref, + block_offsets, + tokens_per_block, + head_dim, + ) + + # --- Fused postprocess + scatter kernel --- + compressed_mask = torch.ones(total_tokens, dtype=torch.bool, device=device) + torch.ops.trtllm.compressor_postprocess_scatter( + kv_comp, + None, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + kv_cache_fused, + num_comp_tokens, + cu_kv_comp, + start_pos, + block_offsets, + compressed_mask, + tokens_per_block, + 0, + rotate_activation, + None, + None, + ) + torch.cuda.synchronize() + + max_diff = (kv_cache_fused.float() - kv_cache_ref.float()).abs().max().item() + assert max_diff < 0.05, f"Postprocess+scatter vs ref max diff = {max_diff}" + + +@pytest.mark.parametrize( + "batch_size, tokens_per_batch_list, head_dim, nope_dim, rope_dim, tokens_per_block, mask_pattern", + [ + # Variable token counts + non-contiguous mask: skip batches 1 and 3 + (4, [8, 4, 6, 2], 128, 64, 64, 4, [True, False, True, False]), + # All masked except one + (4, [4, 4, 4, 4], 128, 64, 64, 4, [False, False, True, False]), + # Alternating mask with head_dim=512 + (4, [4, 8, 4, 8], 512, 384, 128, 4, [True, False, True, False]), + # Single batch masked (edge case) + (1, [8], 128, 64, 64, 4, [False]), + # All unmasked (baseline sanity) + (3, [6, 4, 8], 128, 64, 64, 4, [True, True, True]), + ], +) +@pytest.mark.skipif( + not _HAS_POSTPROCESS_SCATTER, reason="Postprocess/scatter CUDA ops not available" +) +def test_fused_postprocess_scatter_masked_batches( + batch_size, tokens_per_batch_list, head_dim, nope_dim, rope_dim, tokens_per_block, mask_pattern +): + """Verify compressed_mask skips scatter for masked batches with variable token counts. + + Tests: + - Masked batches have zero cache + - Unmasked batches have correct postprocessed values matching reference + - Non-contiguous and all-masked/all-unmasked patterns + """ + torch.manual_seed(42) + device = "cuda" + + num_comp_tokens_list = tokens_per_batch_list + num_comp_tokens = torch.tensor(num_comp_tokens_list, device=device, dtype=torch.int32) + total_tokens = int(num_comp_tokens.sum().item()) + + kv_comp = torch.randn(total_tokens, head_dim, device=device, dtype=torch.bfloat16) * 0.1 + rms_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) * 0.1 + 1.0 + rms_eps = 1e-5 + + max_pos = total_tokens + 64 + cos_sin_table = torch.randn(max_pos, 2, rope_dim // 2, device=device, dtype=torch.float32) + position_ids = torch.arange(total_tokens, device=device, dtype=torch.int32) + + max_tokens = max(num_comp_tokens_list) + max_comp_len = max_tokens + 4 + max_blocks = (max_comp_len + tokens_per_block - 1) // tokens_per_block + num_blocks = batch_size * max_blocks + kv_cache_fused = torch.zeros( + num_blocks, tokens_per_block * head_dim, device=device, dtype=torch.bfloat16 + ) + + block_offsets = torch.zeros(batch_size, max_blocks, device=device, dtype=torch.int32) + for b in range(batch_size): + block_offsets[b] = torch.arange(b * max_blocks, (b + 1) * max_blocks, dtype=torch.int32) + + cu_kv_comp = torch.zeros(batch_size + 1, device=device, dtype=torch.int32) + cu_kv_comp[1:] = num_comp_tokens.cumsum(0) + start_pos = torch.zeros(batch_size, device=device, dtype=torch.int32) + + # Build per-token compressed_mask from per-batch mask_pattern + compressed_mask = torch.zeros(total_tokens, dtype=torch.bool, device=device) + for b in range(batch_size): + if mask_pattern[b]: + start = cu_kv_comp[b].item() + end = cu_kv_comp[b + 1].item() + compressed_mask[start:end] = True + + # Build reference for unmasked batches + ref = _build_postprocess_reference( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + head_dim, + rotate_activation=True, + ) + + torch.ops.trtllm.compressor_postprocess_scatter( + kv_comp, + None, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + kv_cache_fused, + num_comp_tokens, + cu_kv_comp, + start_pos, + block_offsets, + compressed_mask, + tokens_per_block, + 0, + True, + None, + None, + ) + torch.cuda.synchronize() + + for b in range(batch_size): + n_tokens = num_comp_tokens_list[b] + if not mask_pattern[b]: + # Masked: all cache blocks must be zero + for blk_idx in range(max_blocks): + phys = block_offsets[b, blk_idx].item() + assert kv_cache_fused[phys].abs().max().item() == 0.0, ( + f"Masked batch {b} block {blk_idx} should be zero" + ) + else: + # Unmasked: cache values should match reference + token_offset = cu_kv_comp[b].item() + for t in range(n_tokens): + blk = t // tokens_per_block + off = t % tokens_per_block + phys = block_offsets[b, blk].item() + cache_row = kv_cache_fused[phys, off * head_dim : (off + 1) * head_dim] + ref_row = ref[token_offset + t] + max_diff = (cache_row.float() - ref_row.float()).abs().max().item() + assert max_diff < 0.05, f"Unmasked batch {b} token {t}: max_diff={max_diff}" + + +def _build_postprocess_reference( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + head_dim, + rotate_activation=True, + keep_fp32=False, +): + """Reference: RMSNorm -> RoPE -> optional Hadamard on kv_comp. + + Stays fp32 between steps and only truncates to ``kv_comp.dtype`` at the + end -- matches the kernel which keeps activations in fp32 registers + throughout postprocess (no V4-Pro fake-quant bf16 round-trips). + """ + x = kv_comp.clone().float() + var = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(var + rms_eps) + x = rms_weight.float() * x + xn = x[:, :nope_dim] + xp = x[:, nope_dim:] + half_rope = rope_dim // 2 + cos_v = cos_sin_table[position_ids.long(), 0, :] + sin_v = cos_sin_table[position_ids.long(), 1, :] + xp = xp.view(-1, half_rope, 2) + x_even, x_odd = xp[..., 0], xp[..., 1] + xp = torch.stack([x_even * cos_v - x_odd * sin_v, x_odd * cos_v + x_even * sin_v], dim=-1) + xp = xp.view(-1, rope_dim) + x = torch.cat([xn, xp], dim=-1) + if rotate_activation: + try: + from fast_hadamard_transform import hadamard_transform + except ImportError: + pytest.skip("fast_hadamard_transform not installed") + x = hadamard_transform(x, scale=head_dim**-0.5) + if keep_fp32: + return x + return x.to(kv_comp.dtype) + + +def _ceil_pow2_scale(amax: torch.Tensor, max_value_inv: float, min_amax: float) -> torch.Tensor: + """V4-Pro fast_round_scale: 2^ceil(log2(amax * max_value_inv)). + + Uses fp32 bit manipulation to match V4's fast_log2_ceil + fast_pow2 + byte-for-byte (avoids log2/exp2 fp32 rounding, so the resulting + power-of-2 is exact even at 2^k boundaries). + """ + scaled = torch.clamp(amax.float(), min=min_amax) * max_value_inv + bits = scaled.contiguous().view(torch.int32) + exp_part = ((bits >> 23) & 0xFF) - 127 + has_mantissa = (bits & 0x7FFFFF).ne(0).to(torch.int32) + log2_ceil = exp_part + has_mantissa + pow2_bits = (log2_ceil + 127) << 23 + return pow2_bits.view(torch.float32) + + +def _e2m1_nibbles_reference(x: torch.Tensor) -> torch.Tensor: + thresholds = torch.tensor([0.0, 0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], device=x.device) + abs_x = x.abs().clamp(max=6.0) + idx = torch.zeros_like(abs_x, dtype=torch.uint8) + for level_idx in range(1, 8): + take_upper = (abs_x > thresholds[level_idx]) | ( + (abs_x == thresholds[level_idx]) & (level_idx % 2 == 0) + ) + idx = torch.where(take_upper, torch.full_like(idx, level_idx), idx) + sign = torch.signbit(x).to(torch.uint8) << 3 + return idx | sign + + +def _mxfp4_reference(x: torch.Tensor, block_size: int = 32) -> tuple[torch.Tensor, torch.Tensor]: + assert x.shape[-1] % block_size == 0 + packed = torch.empty(*x.shape[:-1], x.shape[-1] // 2, device=x.device, dtype=torch.uint8) + scales = torch.empty( + *x.shape[:-1], x.shape[-1] // block_size, device=x.device, dtype=torch.uint8 + ) + fp4_min_amax = 6.0 * torch.finfo(torch.float32).tiny + for start in range(0, x.shape[-1], block_size): + end = start + block_size + chunk = x[:, start:end].float() + scale = _ceil_pow2_scale(chunk.abs().amax(dim=1, keepdim=True), 1.0 / 6.0, fp4_min_amax) + q_nibbles = _e2m1_nibbles_reference(chunk / scale) + packed[:, start // 2 : end // 2] = q_nibbles[:, 0::2] | (q_nibbles[:, 1::2] << 4) + scale_bits = scale.contiguous().view(torch.int32) + scales[:, start // block_size] = ((scale_bits[:, 0] >> 23) & 0xFF).to(torch.uint8) + return packed, scales + + +def _setup_fused_test_inputs( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block +): + """Common setup for fused postprocess scatter tests.""" + torch.manual_seed(42) + device = "cuda" + total_tokens = batch_size * num_tokens + + kv_comp = torch.randn(total_tokens, head_dim, device=device, dtype=torch.bfloat16) * 0.1 + rms_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) * 0.1 + 1.0 + rms_eps = 1e-5 + max_pos = total_tokens + 64 + cos_sin_table = torch.randn(max_pos, 2, rope_dim // 2, device=device, dtype=torch.float32) + position_ids = torch.arange(total_tokens, device=device, dtype=torch.int32) + + max_comp_len = num_tokens + 4 + max_blocks = (max_comp_len + tokens_per_block - 1) // tokens_per_block + num_blocks = batch_size * max_blocks + + block_offsets = torch.zeros(batch_size, max_blocks, device=device, dtype=torch.int32) + for b in range(batch_size): + block_offsets[b] = torch.arange(b * max_blocks, (b + 1) * max_blocks, dtype=torch.int32) + + num_comp_tokens = torch.full((batch_size,), num_tokens, device=device, dtype=torch.int32) + cu_kv_comp = torch.zeros(batch_size + 1, device=device, dtype=torch.int32) + cu_kv_comp[1:] = num_comp_tokens.cumsum(0) + start_pos = torch.zeros(batch_size, device=device, dtype=torch.int32) + + return ( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + num_blocks, + max_blocks, + block_offsets, + num_comp_tokens, + cu_kv_comp, + start_pos, + total_tokens, + ) + + +@pytest.mark.parametrize("rotate_activation", [True, False], ids=["rotate", "no_rotate"]) +@pytest.mark.parametrize( + "batch_size,num_tokens,head_dim,nope_dim,rope_dim,tokens_per_block", + [ + pytest.param(1, 16, 128, 64, 64, 32, id="b1_t16_hd128"), + pytest.param(2, 32, 128, 64, 64, 32, id="b2_t32_hd128"), + pytest.param(1, 128, 512, 256, 256, 128, id="b1_t128_hd512"), + pytest.param(2, 64, 512, 256, 256, 128, id="b2_t64_hd512"), + ], +) +@pytest.mark.skipif( + not _HAS_POSTPROCESS_SCATTER, reason="Postprocess/scatter CUDA ops not available" +) +def test_fused_postprocess_scatter_fp8_pertensor( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block, rotate_activation +): + """Test fused RMSNorm+RoPE+(optional Hadamard)+FP8PerTensor scatter vs reference. + + FP8 per-tensor uses scale=1.0 (direct float->fp8_e4m3fn cast). + Validates fp8 bytes in cache match reference pipeline. + """ + ( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + num_blocks, + max_blocks, + block_offsets, + num_comp_tokens, + cu_kv_comp, + start_pos, + total_tokens, + ) = _setup_fused_test_inputs( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block + ) + + # FP8 per-tensor: cache_stride_blk_bytes = tpb * hd (1 byte per element) + cache_stride_blk_bytes = tokens_per_block * head_dim + kv_cache_fused = torch.zeros( + num_blocks * cache_stride_blk_bytes, device="cuda", dtype=torch.uint8 + ) + + # Reference: postprocess (returns bf16), then cast to fp8 + ref = _build_postprocess_reference( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + head_dim, + rotate_activation, + ) + ref_fp8 = ref.float().to(torch.float8_e4m3fn) + + # Fused postprocess + scatter (cache_dtype=fp8, scale_type=fp8_pertensor) + compressed_mask = torch.ones(total_tokens, dtype=torch.bool, device="cuda") + torch.ops.trtllm.compressor_postprocess_scatter( + kv_comp, + None, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + kv_cache_fused, + num_comp_tokens, + cu_kv_comp, + start_pos, + block_offsets, + compressed_mask, + tokens_per_block, + 1, + rotate_activation, + None, + None, + ) + torch.cuda.synchronize() + + # Read back and compare fp8 bytes + for b in range(batch_size): + for t in range(num_tokens): + blk = t // tokens_per_block + off = t % tokens_per_block + phys = block_offsets[b, blk].item() + base = phys * cache_stride_blk_bytes + off * head_dim + fused_bytes = kv_cache_fused[base : base + head_dim] + ref_bytes = ref_fp8[b * num_tokens + t].view(torch.uint8) + max_diff = (fused_bytes.int() - ref_bytes.int()).abs().max().item() + assert max_diff <= 1, ( + f"FP8 pertensor mismatch at b={b}, t={t}: max byte diff={max_diff}" + ) + + +@pytest.mark.parametrize("rotate_activation", [True, False], ids=["rotate", "no_rotate"]) +@pytest.mark.parametrize( + "batch_size,num_tokens,head_dim,nope_dim,rope_dim,tokens_per_block", + [ + pytest.param(1, 16, 128, 64, 64, 32, id="b1_t16_hd128"), + pytest.param(2, 32, 128, 64, 64, 32, id="b2_t32_hd128"), + pytest.param(1, 128, 512, 256, 256, 128, id="b1_t128_hd512"), + pytest.param(2, 64, 512, 256, 256, 128, id="b2_t64_hd512"), + ], +) +@pytest.mark.skipif( + not _HAS_POSTPROCESS_SCATTER, reason="Postprocess/scatter CUDA ops not available" +) +def test_fused_postprocess_scatter_fp8_blockwise( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block, rotate_activation +): + """Test fused RMSNorm+RoPE+(optional Hadamard)+FP8Blockwise scatter vs reference. + + FP8 blockwise uses per-128-element scales. Validates: + 1. FP8 data in cache matches reference quantization + 2. Scales in cache are correct + 3. Optional fp8_output and scale_output buffers are populated + """ + ( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + num_blocks, + max_blocks, + block_offsets, + num_comp_tokens, + cu_kv_comp, + start_pos, + total_tokens, + ) = _setup_fused_test_inputs( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block + ) + + num_scale_blocks = head_dim // 128 + cache_stride_blk_bytes = tokens_per_block * head_dim + tokens_per_block * num_scale_blocks * 4 + kv_cache_fused = torch.zeros( + num_blocks * cache_stride_blk_bytes, device="cuda", dtype=torch.uint8 + ) + + # Optional output buffers + fp8_output = torch.zeros(total_tokens, head_dim, device="cuda", dtype=torch.uint8) + scale_output = torch.zeros(total_tokens, num_scale_blocks, device="cuda", dtype=torch.float32) + + # Reference: postprocess (returns bf16) then blockwise quantize + ref = _build_postprocess_reference( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + head_dim, + rotate_activation, + ) + + # Per-128-element quantization reference (from bf16 values) + ref_fp8_list = [] + ref_scale_list = [] + for token_idx in range(total_tokens): + row = ref[token_idx] + fp8_row = torch.zeros(head_dim, dtype=torch.uint8, device="cuda") + scales = torch.zeros(num_scale_blocks, dtype=torch.float32, device="cuda") + for s in range(num_scale_blocks): + chunk = row[s * 128 : (s + 1) * 128].float() + amax = chunk.abs().max() + scale = amax / 448.0 + inv_scale = (448.0 / amax) if amax > 0 else 1.0 + quantized = (chunk * inv_scale).to(torch.float8_e4m3fn) + fp8_row[s * 128 : (s + 1) * 128] = quantized.view(torch.uint8) + scales[s] = scale + ref_fp8_list.append(fp8_row) + ref_scale_list.append(scales) + ref_fp8_all = torch.stack(ref_fp8_list) + ref_scale_all = torch.stack(ref_scale_list) + + # Fused postprocess + scatter (cache_dtype=fp8, scale_type=fp8_blockwise) + compressed_mask = torch.ones(total_tokens, dtype=torch.bool, device="cuda") + torch.ops.trtllm.compressor_postprocess_scatter( + kv_comp, + None, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + kv_cache_fused, + num_comp_tokens, + cu_kv_comp, + start_pos, + block_offsets, + compressed_mask, + tokens_per_block, + 2, + rotate_activation, + fp8_output, + scale_output, + ) + torch.cuda.synchronize() + + # Validate optional output buffers + max_byte_diff = (fp8_output.int() - ref_fp8_all.int()).abs().max().item() + assert max_byte_diff <= 1, f"fp8_output buffer max byte diff = {max_byte_diff}" + + max_scale_diff = (scale_output - ref_scale_all).abs().max().item() + assert max_scale_diff < 5e-3, f"scale_output buffer max scale diff = {max_scale_diff}" + + # Validate cache contents + for b in range(batch_size): + for t in range(num_tokens): + blk = t // tokens_per_block + off = t % tokens_per_block + phys = block_offsets[b, blk].item() + block_base = phys * cache_stride_blk_bytes + + # Check FP8 data + fp8_base = block_base + off * head_dim + fused_bytes = kv_cache_fused[fp8_base : fp8_base + head_dim] + ref_bytes = ref_fp8_all[b * num_tokens + t] + byte_diff = (fused_bytes.int() - ref_bytes.int()).abs().max().item() + assert byte_diff <= 1, ( + f"FP8 blockwise data mismatch at b={b}, t={t}: max byte diff={byte_diff}" + ) + + # Check scales + scale_base = block_base + tokens_per_block * head_dim + for s in range(num_scale_blocks): + scale_off = scale_base + (off * num_scale_blocks + s) * 4 + fused_scale_bytes = kv_cache_fused[scale_off : scale_off + 4] + fused_scale = fused_scale_bytes.view(torch.float32).item() + ref_scale = ref_scale_all[b * num_tokens + t, s].item() + assert abs(fused_scale - ref_scale) < 5e-3, ( + f"Scale mismatch at b={b}, t={t}, s={s}: " + f"fused={fused_scale:.6f} ref={ref_scale:.6f}" + ) + + +@pytest.mark.parametrize("rotate_activation", [True, False], ids=["rotate", "no_rotate"]) +@pytest.mark.parametrize( + "batch_size,num_tokens,head_dim,nope_dim,rope_dim,tokens_per_block", + [ + pytest.param(1, 16, 128, 64, 64, 32, id="b1_t16_indexer"), + pytest.param(2, 32, 128, 64, 64, 32, id="b2_t32_indexer"), + pytest.param(1, 64, 512, 256, 256, 64, id="b1_t64_hd512"), + ], +) +@pytest.mark.skipif( + not _HAS_POSTPROCESS_SCATTER, reason="Postprocess/scatter CUDA ops not available" +) +def test_fused_postprocess_scatter_mxfp4( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block, rotate_activation +): + """Optimized indexer mode: packed FP4 cache data plus per-32 UE8M0 scale bytes.""" + ( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + num_blocks, + _, + block_offsets, + num_comp_tokens, + cu_kv_comp, + start_pos, + total_tokens, + ) = _setup_fused_test_inputs( + batch_size, num_tokens, head_dim, nope_dim, rope_dim, tokens_per_block + ) + + packed_head_dim = head_dim // 2 + num_scale_blocks = head_dim // 32 + cache_stride_blk_bytes = ( + tokens_per_block * packed_head_dim + tokens_per_block * num_scale_blocks + ) + kv_cache_fused = torch.zeros( + num_blocks * cache_stride_blk_bytes, device="cuda", dtype=torch.uint8 + ) + fp4_output = torch.zeros(total_tokens, packed_head_dim, device="cuda", dtype=torch.uint8) + scale_output = torch.zeros(total_tokens, num_scale_blocks, device="cuda", dtype=torch.uint8) + + ref = _build_postprocess_reference( + kv_comp, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + head_dim, + rotate_activation, + keep_fp32=True, + ) + ref_packed, ref_scales = _mxfp4_reference(ref) + + compressed_mask = torch.ones(total_tokens, dtype=torch.bool, device="cuda") + torch.ops.trtllm.compressor_postprocess_scatter( + kv_comp, + None, + rms_weight, + rms_eps, + cos_sin_table, + position_ids, + nope_dim, + rope_dim, + kv_cache_fused, + num_comp_tokens, + cu_kv_comp, + start_pos, + block_offsets, + compressed_mask, + tokens_per_block, + 3, + rotate_activation, + fp4_output, + scale_output, + ) + torch.cuda.synchronize() + + assert torch.equal(fp4_output, ref_packed), "packed FP4 output bytes must match reference" + assert torch.equal(scale_output, ref_scales), "UE8M0 scale bytes must match reference" + + for b in range(batch_size): + for t in range(num_tokens): + blk = t // tokens_per_block + off = t % tokens_per_block + phys = block_offsets[b, blk].item() + block_base = phys * cache_stride_blk_bytes + fp4_base = block_base + off * packed_head_dim + scale_base = block_base + tokens_per_block * packed_head_dim + off * num_scale_blocks + token_idx = b * num_tokens + t + assert torch.equal( + kv_cache_fused[fp4_base : fp4_base + packed_head_dim], + ref_packed[token_idx], + ) + assert torch.equal( + kv_cache_fused[scale_base : scale_base + num_scale_blocks], + ref_scales[token_idx], + ) + + +# ============================================================================ +# Benchmarks: cuTile Kernels vs PyTorch Reference +# ============================================================================ + + +def benchmark_compress_kernel(): + """Benchmark CUDA vs PyTorch compress (decode) kernels using triton.testing.do_bench.""" + + print("\n" + "=" * 70) + print("Compress Kernel Benchmark: CUDA vs PyTorch (decode)") + print("=" * 70) + + configs = [ + # (batch_size, compress_ratio, head_dim, overlap, page_size, name) + (1, 4, 512, True, 32, "b1_r4_d512_overlap"), + (8, 4, 512, True, 32, "b8_r4_d512_overlap"), + (32, 4, 512, True, 32, "b32_r4_d512_overlap"), + (1, 128, 512, False, 32, "b1_r128_d512"), + (8, 128, 512, False, 32, "b8_r128_d512"), + (32, 128, 512, False, 32, "b32_r128_d512"), + (1, 4, 128, True, 8, "b1_r4_d128_overlap"), + (8, 4, 128, True, 8, "b8_r4_d128_overlap"), + (32, 4, 128, True, 8, "b32_r4_d128_overlap"), + ] + + results = [] + for batch_size, compress_ratio, head_dim, overlap, page_size, name in configs: + coff = 2 if overlap else 1 + state_dim = coff * head_dim + max_blocks = (compress_ratio * 2 + page_size - 1) // page_size + + # Prepare inputs + ape = torch.randn(compress_ratio, state_dim, device="cuda") + new_kv = torch.randn(batch_size, state_dim, device="cuda") + new_score = torch.randn(batch_size, state_dim, device="cuda") + kv_score = fuse_kv_score(new_kv, new_score) + + # PyTorch state + state_len = coff * compress_ratio + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full( + (batch_size, state_len, state_dim), float("-inf"), device="cuda" + ) + + # Paged cache + num_blocks = batch_size * max_blocks + paged_kv = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros(num_blocks, page_size, state_dim, device="cuda") + block_table = torch.arange(num_blocks, device="cuda", dtype=torch.int32).view( + batch_size, max_blocks + ) + + # Use step = compress_ratio - 1 to trigger compression. + step = compress_ratio - 1 + token_idx = (compress_ratio + step) if overlap else step + kv_lens = torch.full((batch_size,), token_idx + 1, device="cuda", dtype=torch.int32) + start_pos = torch.full((batch_size,), token_idx, device="cuda", dtype=torch.int32) + cu_seq_lens, cu_outputs = prepare_decode_metadata( + batch_size, compress_ratio, head_dim, torch.device("cuda"), next_n=1 + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, torch.device("cuda"), torch.bfloat16 + ) + + paged_kv_cuda = paged_kv.clone() + paged_score_cuda = paged_score.clone() + kv_comp_cuda = kv_comp.clone() + + def cuda_fn(): + decode_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp_cuda, + paged_kv_cuda, + paged_score_cuda, + block_table, + block_table, + compress_ratio, + head_dim, + page_size, + next_n=1, + ) + + def pytorch_fn(): + run_pytorch_reference( + new_kv.unsqueeze(1), + new_score.unsqueeze(1), + ape, + kv_state_py.clone(), + score_state_py.clone(), + step, + compress_ratio, + head_dim, + overlap, + ) + + cuda_ms = triton.testing.do_bench(cuda_fn, warmup=25, rep=100) + pytorch_ms = triton.testing.do_bench(pytorch_fn, warmup=25, rep=100) + + cuda_us = cuda_ms * 1000 + pytorch_us = pytorch_ms * 1000 + + results.append( + { + "name": name, + "cuda_us": cuda_us, + "pytorch_us": pytorch_us, + "speedup": pytorch_us / cuda_us if cuda_us > 0 else float("inf"), + } + ) + + # Print results + print(f"\n{'Config':<30} {'CUDA (us)':>12} {'PyTorch (us)':>12} {'Speedup':>10}") + print("-" * 70) + for r in results: + print( + f"{r['name']:<30} {r['cuda_us']:>12.2f} {r['pytorch_us']:>12.2f} {r['speedup']:>9.2f}x" + ) + print("=" * 70) + + return results + + +def benchmark_compress_prefill_kernel(): + """Benchmark CUDA vs PyTorch compress (prefill) kernels using triton.testing.do_bench.""" + + print("\n" + "=" * 70) + print("Compress Prefill Kernel Benchmark: CUDA vs PyTorch") + print("=" * 70) + + configs = [ + # (batch_size, seqlen, compress_ratio, head_dim, overlap, page_size, name) + (1, 128, 4, 512, True, 32, "prefill_b1_s128_r4_d512"), + (4, 256, 4, 512, True, 32, "prefill_b4_s256_r4_d512"), + (8, 512, 4, 512, True, 32, "prefill_b8_s512_r4_d512"), + (1, 512, 128, 512, False, 32, "prefill_b1_s512_r128_d512"), + (4, 512, 128, 512, False, 32, "prefill_b4_s512_r128_d512"), + (8, 1024, 128, 512, False, 32, "prefill_b8_s1024_r128_d512"), + (1, 128, 4, 128, True, 8, "prefill_b1_s128_r4_d128"), + (8, 512, 4, 128, True, 8, "prefill_b8_s512_r4_d128"), + (32, 512, 4, 128, True, 8, "prefill_b32_s512_r4_d128"), + ] + + results = [] + for batch_size, seqlen, compress_ratio, head_dim, overlap, page_size, name in configs: + coff = 2 if overlap else 1 + state_dim = coff * head_dim + + # Prepare inputs + kv = torch.randn(batch_size, seqlen, state_dim, device="cuda") + score = torch.randn(batch_size, seqlen, state_dim, device="cuda") + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + # PyTorch state + state_len = coff * compress_ratio + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full( + (batch_size, state_len, state_dim), float("-inf"), device="cuda" + ) + + # Shared inputs + kv_score = fuse_kv_score(kv.view(-1, state_dim), score.view(-1, state_dim)) + kv_lens = torch.full((batch_size,), seqlen, device="cuda", dtype=torch.int32) + start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens, start_pos, compress_ratio, head_dim, kv_score.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score.device, torch.bfloat16 + ) + paged_kv, paged_score, block_table, _, _ = create_paged_cache( + batch_size, seqlen, compress_ratio, head_dim, overlap, page_size + ) + + paged_kv_kernel = paged_kv.clone() + paged_score_kernel = paged_score.clone() + kv_comp_kernel = kv_comp.clone() + + def kernel_fn(): + seq_lens = kv_lens - start_pos + num_outputs_per_batch = torch.clamp(seq_lens // compress_ratio, min=1) + max_outputs = num_outputs_per_batch.max().item() + prefill_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp_kernel, + paged_kv_kernel, + paged_score_kernel, + block_table, + max_outputs, + block_table, + compress_ratio, + head_dim, + page_size, + ) + + def pytorch_fn(): + run_pytorch_prefill_reference( + kv.clone(), + score.clone(), + ape, + kv_state_py.clone(), + score_state_py.clone(), + compress_ratio, + head_dim, + overlap, + ) + + kernel_ms = triton.testing.do_bench(kernel_fn, warmup=25, rep=100) + pytorch_ms = triton.testing.do_bench(pytorch_fn, warmup=25, rep=100) + + kernel_us = kernel_ms * 1000 + pytorch_us = pytorch_ms * 1000 + + results.append( + { + "name": name, + "kernel_us": kernel_us, + "pytorch_us": pytorch_us, + "speedup": pytorch_us / kernel_us if kernel_us > 0 else float("inf"), + } + ) + + # Print results + print(f"\n{'Config':<30} {'CUDA (us)':>12} {'PyTorch (us)':>12} {'Speedup':>10}") + print("-" * 70) + for r in results: + print( + f"{r['name']:<30} {r['kernel_us']:>12.2f} " + f"{r['pytorch_us']:>12.2f} {r['speedup']:>9.2f}x" + ) + print("=" * 70) + + return results + + +def run_all_benchmarks(): + """Run all kernel benchmarks.""" + print("\n" + "=" * 80) + print("Compressor Kernel Benchmarks") + print("=" * 80) + + benchmark_compress_kernel() + benchmark_compress_prefill_kernel() + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py new file mode 100644 index 000000000000..fa0113a1465d --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -0,0 +1,2457 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests comparing Compressor with RefCompressor.""" + +import math +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from tensorrt_llm._torch.attention_backend.interface import ( + MLAParams, + PositionalEmbeddingParams, + PositionEmbeddingType, + RotaryScalingType, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import Compressor +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4AttentionType, +) +from tensorrt_llm._torch.modules.rotary_embedding import RopeParams +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping + + +def _hadamard_transform(x: torch.Tensor, scale: float) -> torch.Tensor: + """Pure-Python Walsh-Hadamard transform (no external dependency). + + Matches the butterfly implementation in the CUDA postProcessScatterKernel. + """ + n = x.shape[-1] + assert n & (n - 1) == 0, "Last dim must be a power of 2" + y = x.float().clone() + stride = 1 + while stride < n: + idx = torch.arange(n, device=x.device) + lo_mask = (idx & stride) == 0 + hi_idx = idx ^ stride + a = y[..., lo_mask].clone() + b = y[..., hi_idx[lo_mask]].clone() + y[..., lo_mask] = a + b + y[..., hi_idx[lo_mask]] = a - b + stride <<= 1 + return (y * scale).to(x.dtype) + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Hadamard rotation matching the CUDA kernel (no fast_hadamard_transform needed).""" + return _hadamard_transform(x, scale=x.size(-1) ** -0.5) + + +# ============================================================================ +# Dummy Metadata for Testing +# ============================================================================ + + +class DummyAttentionMetadata: + """Dummy attention metadata for testing Compressor.forward.""" + + def __init__( + self, + num_contexts: int, + num_generations: int, + num_ctx_tokens: int, + num_tokens: int, + kv_cache_manager: DeepseekV4CacheManager, + block_tables: dict, + cu_seq_lens: dict, + cu_new_comp_kv: dict, + compressed_position_ids: dict, + compressed_kv_lens: dict, + past_kv_lens: dict, + new_comp_kv_lens_cuda: dict, + num_total_compressed_tokens: dict, + max_ctx_compressed_tokens: dict, + slot_mapping_fp8: torch.Tensor = None, + slot_mapping_scale: torch.Tensor = None, + compressed_mask_cuda: dict = None, + ): + self.num_contexts = num_contexts + self.num_generations = num_generations + self.num_ctx_tokens = num_ctx_tokens + self.num_tokens = num_tokens + self.kv_cache_manager = kv_cache_manager + self.block_tables = block_tables + self.cu_seq_lens_cuda = cu_seq_lens + self.cu_new_comp_kv_cuda = cu_new_comp_kv + self.compressed_position_ids_cuda = compressed_position_ids + self.compressed_kv_lens_cuda = compressed_kv_lens + self.past_kv_lens_cuda = past_kv_lens + self.slot_mapping_fp8 = slot_mapping_fp8 + self.slot_mapping_scale = slot_mapping_scale + self.new_comp_kv_lens_cuda = new_comp_kv_lens_cuda + self.num_total_compressed_tokens = num_total_compressed_tokens + self.max_ctx_compressed_tokens = max_ctx_compressed_tokens + self.compressed_mask_cuda = compressed_mask_cuda + self.num_gen_tokens_per_seq = 0 # Set by caller + self.kv_lens_cuda_runtime = None # Set by caller + self.cached_token_lens_cuda = None # Set by caller + + +# ============================================================================ +# Reference Implementation (DO NOT MODIFY) +# ============================================================================ + + +@dataclass +class ModelArgs: + """Model arguments for Compressor.""" + + max_batch_size: int = 16 + max_seq_len: int = 4096 + dim: int = 4096 + head_dim: int = 512 + rope_head_dim: int = 64 + norm_eps: float = 1e-6 + compress_ratios: Tuple[int, ...] = (1, 1, 4, 128, 4, 128, 4) + + +class RMSNorm(nn.Module): + """Root Mean Square Layer Normalization.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor): + dtype = x.dtype + x = x.float() + var = x.square().mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.eps) + return (self.weight * x).to(dtype) + + +class Linear(nn.Module): + """Simple linear layer (fp32 weights for reference).""" + + def __init__(self, in_features: int, out_features: int, dtype=None): + super().__init__() + self.weight = nn.Parameter( + torch.empty(out_features, in_features, dtype=dtype or torch.float32) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, self.weight) + + +def apply_rotary_emb( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + """Apply rotary positional embeddings.""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +class RefCompressor(nn.Module): + """Reference Compressor implementation for testing.""" + + def __init__( + self, args: ModelArgs, compress_ratio: int = 4, head_dim: int = 512, rotate: bool = False + ): + super().__init__() + self.dim = args.dim + self.head_dim = head_dim + self.rope_head_dim = args.rope_head_dim + self.nope_head_dim = head_dim - args.rope_head_dim + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.rotate = rotate + coff = 1 + self.overlap + + self.ape = nn.Parameter( + torch.empty(compress_ratio, coff * self.head_dim, dtype=torch.float32) + ) + self.wkv = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.wgate = Linear(self.dim, coff * self.head_dim, dtype=torch.float32) + self.norm = RMSNorm(self.head_dim, args.norm_eps) + self.kv_cache = None + self.register_buffer( + "kv_state", + torch.zeros( + args.max_batch_size, + coff * compress_ratio, + coff * self.head_dim, + dtype=torch.float32, + ), + persistent=False, + ) + self.register_buffer( + "score_state", + torch.full( + (args.max_batch_size, coff * compress_ratio, coff * self.head_dim), + float("-inf"), + dtype=torch.float32, + ), + persistent=False, + ) + + def overlap_transform(self, tensor: torch.Tensor, value=0): + b, s, _, _ = tensor.size() + ratio, d = self.compress_ratio, self.head_dim + new_tensor = tensor.new_full((b, s, 2 * ratio, d), value) + new_tensor[:, :, ratio:] = tensor[:, :, :, d:] + new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor): + assert self.kv_cache is not None + bsz, seqlen, _ = x.size() + ratio, overlap, d = self.compress_ratio, self.overlap, self.head_dim + dtype = x.dtype + x = x.float() + kv = self.wkv(x) + score = self.wgate(x) + if start_pos == 0: + should_compress = seqlen >= ratio + remainder = seqlen % ratio + cutoff = seqlen - remainder + freqs_cis = freqs_cis[:cutoff:ratio] + offset = ratio if overlap else 0 + if overlap and cutoff >= ratio: + self.kv_state[:bsz, :ratio] = kv[:, cutoff - ratio : cutoff] + self.score_state[:bsz, :ratio] = score[:, cutoff - ratio : cutoff] + self.ape + if remainder > 0: + kv, self.kv_state[:bsz, offset : offset + remainder] = kv.split( + [cutoff, remainder], dim=1 + ) + self.score_state[:bsz, offset : offset + remainder] = ( + score[:, cutoff:] + self.ape[:remainder] + ) + score = score[:, :cutoff] + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + self.ape + if overlap: + kv = self.overlap_transform(kv, 0) + score = self.overlap_transform(score, float("-inf")) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + else: + # Handles any seqlen >= 1 with start_pos > 0 (decode or chunked prefill). + # freqs_cis must be the FULL precomputed array so output freqs can be + # indexed at absolute first-token-of-window positions (win_first = pos+1-ratio). + outputs = [] + output_freqs = [] + for t in range(seqlen): + pos = start_pos + t + kv_t = kv[:, t] + sc_t = score[:, t] + self.ape[pos % ratio] + if overlap: + self.kv_state[:bsz, ratio + pos % ratio] = kv_t + self.score_state[:bsz, ratio + pos % ratio] = sc_t + if (pos + 1) % ratio == 0: + kv_state = torch.cat( + [self.kv_state[:bsz, :ratio, :d], self.kv_state[:bsz, ratio:, d:]], + dim=1, + ) + score_state = torch.cat( + [ + self.score_state[:bsz, :ratio, :d], + self.score_state[:bsz, ratio:, d:], + ], + dim=1, + ) + comp = (kv_state * score_state.softmax(dim=1)).sum(dim=1, keepdim=True) + self.kv_state[:bsz, :ratio] = self.kv_state[:bsz, ratio:] + self.score_state[:bsz, :ratio] = self.score_state[:bsz, ratio:] + outputs.append(comp) + output_freqs.append(freqs_cis[pos + 1 - ratio : pos + 2 - ratio]) + else: + self.kv_state[:bsz, pos % ratio] = kv_t + self.score_state[:bsz, pos % ratio] = sc_t + if (pos + 1) % ratio == 0: + comp = (self.kv_state[:bsz] * self.score_state[:bsz].softmax(dim=1)).sum( + dim=1, keepdim=True + ) + outputs.append(comp) + output_freqs.append(freqs_cis[pos + 1 - ratio : pos + 2 - ratio]) + + if not outputs: + return None + should_compress = True + kv = torch.cat(outputs, dim=1) + freqs_cis = torch.cat(output_freqs, dim=0) + + if not should_compress: + return + # Match the kernel's hand-off: the kernel's kv_comp buffer is bf16, + # so we bf16-truncate the compression result here too. Then upcast + # back to fp32 for the postprocess so RMSNorm/RoPE/Hadamard run at + # full fp32 precision -- matching the kernel which keeps activations + # in fp32 registers throughout postprocess (no V4-Pro fake-quant bf16 + # round-trips between norm / rope / hadamard). Returns fp32 so the + # downstream QDQ reference sees the same full-precision values the + # kernel feeds into its QDQ on registers; only the cache write + # truncates to bf16. + kv = kv.to(dtype).float() + kv = self.norm(kv) + apply_rotary_emb(kv[..., -self.rope_head_dim :], freqs_cis) + if self.rotate: + kv = rotate_activation(kv) + if start_pos == 0: + self.kv_cache[:bsz, : seqlen // ratio] = kv.to(dtype) + else: + first_abs = start_pos // ratio + n_out = kv.size(1) + self.kv_cache[:bsz, first_abs : first_abs + n_out] = kv.to(dtype) + return kv + + +# ============================================================================ +# Test Configuration & Helpers +# ============================================================================ + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +DIM, HEAD_DIM, ROPE_DIM = 4096, 512, 64 +INDEX_HEAD_DIM = 128 # Fixed head_dim for indexer (INDEXER_COMPRESS) +MAX_BATCH, MAX_SEQ, PAGE_SIZE = 16, 4096, 128 +ORI_SEQ_LEN = 65536 +ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW = 40000.0, 4, 32, 1 + + +def precompute_freqs_cis( + dim, seqlen, original_seq_len, base, factor, beta_fast, beta_slow +) -> torch.Tensor: + """Precompute rotary embeddings.""" + + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if seqlen > original_seq_len: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def assert_similar(t1: Optional[torch.Tensor], t2: Optional[torch.Tensor], name: str = "Output"): + """Assert tensors are similar (cosine sim >= 0.999).""" + if t1 is None and t2 is None: + return + assert t1 is not None and t2 is not None, f"{name}: One is None" + assert t1.shape == t2.shape, f"{name}: Shape mismatch {t1.shape} vs {t2.shape}" + t1, t2 = t1.float().flatten(), t2.float().flatten() + cos_sim = F.cosine_similarity(t1.unsqueeze(0), t2.unsqueeze(0)).item() + # Also check magnitude to avoid scaled-but-equal-direction false positives + max_diff = (t1 - t2).abs().max().item() + scale = max(t1.abs().max().item(), t2.abs().max().item(), 1e-3) + rel_err = max_diff / scale + assert cos_sim >= 0.999, f"{name}: cos_sim={cos_sim:.6f}" + assert rel_err <= 5e-2, f"{name}: rel_err={rel_err:.6f}, max_diff={max_diff:.6f}" + + +def dequantize_blockwise_fp8( + kv_fp8: torch.Tensor, kv_scale: torch.Tensor, block_size: int = 128 +) -> torch.Tensor: + """Dequantize blockwise FP8 data. + + Args: + kv_fp8: [num_tokens, head_dim] FP8 tensor + kv_scale: [num_tokens, num_scale_blocks] scale factors (one per 128 elements) + + Returns: + Dequantized float tensor + """ + num_tokens, head_dim = kv_fp8.shape + num_blocks = (head_dim + block_size - 1) // block_size + + kv_float = kv_fp8.float() + kv_dequant = torch.zeros_like(kv_float) + + for b in range(num_blocks): + start = b * block_size + end = min(start + block_size, head_dim) + kv_dequant[:, start:end] = kv_float[:, start:end] * kv_scale[:, b : b + 1] + + return kv_dequant + + +def assert_fp8_similar( + fp8_result: tuple, + ref_result: torch.Tensor, + kv_cache_dtype: str, + name: str = "FP8 Output", +): + """Assert FP8 result matches reference after dequantization. + + Uses FP8-appropriate tolerances: cos_sim >= 0.99, rel_err <= 10% + """ + kv_fp8, scale = fp8_result + + if kv_cache_dtype in ("fp8_blockwise"): + kv_dequant = dequantize_blockwise_fp8(kv_fp8, scale) + else: # fp8_pertensor + kv_dequant = kv_fp8.float() * scale.item() + + ref_float = ref_result.float().flatten() + dequant_flat = kv_dequant.flatten() + + # Cosine similarity (>= 0.99 for FP8) + cos_sim = F.cosine_similarity(dequant_flat.unsqueeze(0), ref_float.unsqueeze(0)).item() + assert cos_sim >= 0.99, f"{name}: cos_sim={cos_sim:.4f} < 0.99" + + # Relative error (<= 10% for FP8) + max_diff = (dequant_flat - ref_float).abs().max().item() + scale_val = max(ref_float.abs().max().item(), 1e-3) + rel_err = max_diff / scale_val + assert rel_err <= 0.1, f"{name}: rel_err={rel_err:.4f} > 0.1" + + +def read_paged_cache_tokens( + kv_cache: torch.Tensor, + block_offsets: torch.Tensor, + batch_idx: int, + num_tokens: int, + tokens_per_block: int, +) -> torch.Tensor: + """Materialize paged compressed cache for a batch into a contiguous view. + + Args: + kv_cache: Cache buffer with shape [num_blocks, tokens_per_block, head_dim] + block_offsets: Block offset table with shape [num_seqs, max_blocks] + batch_idx: Index of the batch to read + num_tokens: Number of tokens to read + tokens_per_block: Tokens per cache block + + Returns: + Tensor containing the read values with shape [num_tokens, head_dim] + """ + blocks_needed = (num_tokens + tokens_per_block - 1) // tokens_per_block + # Extract block indices for this batch + block_indices = block_offsets[batch_idx, :blocks_needed].tolist() + # Read all blocks at once and reshape + _, _, dim_per_token = kv_cache.shape + return kv_cache[block_indices].reshape(-1, dim_per_token)[:num_tokens] + + +def build_fp8_golden_cache( + kv_fp8: torch.Tensor, + kv_scale: torch.Tensor, + cache_shape: tuple, + block_offsets: torch.Tensor, + batch: int, + num_compressed: int, + tokens_per_block: int, + kv_cache_dtype: str, + head_dim: int = HEAD_DIM, +) -> torch.Tensor: + """Build Python golden reference cache from compressor's FP8 output. + + This validates the cache scatter operation by manually placing the compressor's + FP8 output into the expected cache positions. Quantization correctness is + validated separately via assert_fp8_similar. + + Args: + kv_fp8: Compressor's FP8 output [total_tokens, head_dim] + kv_scale: Compressor's scale output + cache_shape: Shape of the cache tensor + block_offsets: Block offset table [num_seqs, max_blocks] + batch: Batch size + num_compressed: Compressed tokens per batch + tokens_per_block: Tokens per cache block + kv_cache_dtype: "fp8_blockwise" or "fp8_pertensor" + head_dim: Head dimension (default HEAD_DIM=512, use INDEX_HEAD_DIM=128 for indexer) + + Returns: + golden_cache: Expected cache tensor for comparison + """ + total_comp_tokens = batch * num_compressed + golden_cache = torch.zeros(cache_shape, device=kv_fp8.device, dtype=torch.uint8) + + # Convert FP8 to bytes + kv_fp8_bytes = kv_fp8.contiguous().view(torch.uint8).view(total_comp_tokens, head_dim) + + if kv_cache_dtype in ("fp8_blockwise"): + # Blockwise: non-interleaved layout per block: + # [k0, k1, ..., kN, scale0, scale1, ..., scaleN] + num_scale_blocks = (head_dim + 127) // 128 + scale_size = num_scale_blocks * 4 + kv_scale_bytes = ( + kv_scale.flatten().contiguous().view(torch.uint8).view(total_comp_tokens, scale_size) + ) + + # Flatten to 2D [num_blocks, block_stride] for flat byte indexing + num_blocks = cache_shape[0] + golden_flat = golden_cache.view(num_blocks, -1) + + # Scatter into cache using non-interleaved offsets + global_token = 0 + for b in range(batch): + for i in range(num_compressed): + block_idx = i // tokens_per_block + pos_in_block = i % tokens_per_block + block_id = int(block_offsets[b, block_idx].item()) + + # FP8 data in first section of block + fp8_start = pos_in_block * head_dim + golden_flat[block_id, fp8_start : fp8_start + head_dim] = kv_fp8_bytes[global_token] + # Scale data in second section of block + scale_start = tokens_per_block * head_dim + pos_in_block * scale_size + golden_flat[block_id, scale_start : scale_start + scale_size] = kv_scale_bytes[ + global_token + ] + global_token += 1 + + # Reshape back to original shape + golden_cache = golden_flat.view(cache_shape) + else: + # Per-tensor: cache layout is [num_blocks, tokens_per_block, head_dim] + # Same as blockwise but without scale bytes + global_token = 0 + for b in range(batch): + for i in range(num_compressed): + block_idx = i // tokens_per_block + pos_in_block = i % tokens_per_block + block_id = int(block_offsets[b, block_idx].item()) + + golden_cache[block_id, pos_in_block, :head_dim] = kv_fp8_bytes[global_token] + global_token += 1 + + return golden_cache + + +def assert_fp8_cache_match( + kernel_cache: torch.Tensor, + golden_cache: torch.Tensor, + kv_cache_dtype: str, + name: str = "FP8 Cache", + head_dim: int = HEAD_DIM, +): + """Assert kernel cache matches Python golden reference.""" + # Convert both to uint8 bytes for comparison (kernel_cache may be Float8_e4m3fn) + kernel_bytes = kernel_cache.view(torch.uint8) + golden_bytes = golden_cache.view(torch.uint8) + + if torch.equal(kernel_bytes, golden_bytes): + return # Perfect match + + diff_mask = kernel_bytes != golden_bytes + num_diffs = diff_mask.sum().item() + total_bytes = kernel_bytes.numel() + + # Build detailed error message + msg = ( + f"{name}: {num_diffs}/{total_bytes} byte differences ({100 * num_diffs / total_bytes:.4f}%)" + ) + + if kv_cache_dtype in ("fp8_blockwise"): + # Reshape to original layout for detailed analysis + kernel_reshaped = kernel_bytes.view(kernel_cache.shape) + golden_reshaped = golden_bytes.view(golden_cache.shape) + fp8_diffs = ( + (kernel_reshaped[:, :, :head_dim] != golden_reshaped[:, :, :head_dim]).sum().item() + ) + scale_diffs = ( + (kernel_reshaped[:, :, head_dim:] != golden_reshaped[:, :, head_dim:]).sum().item() + ) + + msg += f" [FP8: {fp8_diffs}, Scale: {scale_diffs}]" + + # Show first few differences + diff_indices = torch.nonzero(diff_mask.view(-1))[:5] + for idx in diff_indices: + flat_idx = idx.item() + msg += f"\n Byte {flat_idx}: kernel={kernel_cache.view(-1)[flat_idx].item()}, " + msg += f"golden={golden_cache.view(-1)[flat_idx].item()}" + + raise AssertionError(msg) + + +def run_ref_segmented_forward( + ref: RefCompressor, + tokens: torch.Tensor, + freqs_cis: torch.Tensor, + segments: list[tuple[int, int]], +) -> Optional[torch.Tensor]: + """Run ref forward per segment, concatenating non-None outputs. + + For start_pos == 0 a freq slice is passed (forward indexes from 0). + For start_pos > 0 the FULL freqs_cis array is passed so forward() can + index at absolute win_first positions. + """ + outputs = [] + cursor = 0 + for start_pos, seg_len in segments: + seg_tokens = tokens[:, cursor : cursor + seg_len] + if start_pos > 0: + out = ref(seg_tokens, start_pos, freqs_cis) + else: + out = ref(seg_tokens, start_pos, freqs_cis[start_pos : start_pos + seg_len]) + if out is not None: + outputs.append(out) + cursor += seg_len + assert cursor <= tokens.size(1), "Segment lengths exceed provided tokens" + if not outputs: + return None + return torch.cat(outputs, dim=1) + + +class CompressorWrapper: + """Wrapper around Compressor to manage caches and provide a simpler test interface.""" + + # Class-level constants for DeepseekV4CacheManager + WINDOW_SIZE = 128 + VOCAB_SIZE = 129280 + + def __init__( + self, + compress_ratio: int = 4, + rotate: bool = False, + layer_idx: int = 0, + kv_cache_dtype: str = "default", + is_indexer: bool = False, + ): + self.compress_ratio = compress_ratio + self.overlap = compress_ratio == 4 + self.layer_idx = layer_idx + self.kv_cache_dtype = kv_cache_dtype + self.is_indexer = is_indexer + + # Create MLAParams + # For indexer mode, use INDEX_HEAD_DIM instead of HEAD_DIM + target_head_dim = INDEX_HEAD_DIM if is_indexer else HEAD_DIM + mla_params = MLAParams( + hidden_size=DIM, + qk_rope_head_dim=ROPE_DIM, + qk_nope_head_dim=target_head_dim - ROPE_DIM, + ) + + # Create RoPE - use no scaling to match precompute_freqs_cis behavior + # (precompute_freqs_cis only applies yarn scaling when seqlen > ORI_SEQ_LEN) + rope_params = RopeParams( + dim=ROPE_DIM, + theta=ROPE_THETA, + max_positions=4096, + beta_fast=BETA_FAST, + beta_slow=BETA_SLOW, + scale=1.0, # No scaling + mscale=1.0, + mscale_all_dim=1.0, + original_max_positions=ORI_SEQ_LEN, + scale_type=RotaryScalingType.none, # Match precompute_freqs_cis (no yarn for pos < ORI_SEQ_LEN) + ) + + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, # Basic RoPE without yarn + rope=rope_params, + is_neox=False, + ) + + # Create the Compressor + self.compressor = Compressor( + mla_params=mla_params, + layer_idx=layer_idx, + compress_ratio=compress_ratio, + norm_eps=1e-6, + skip_create_weights_in_init=False, + pos_embd_params=pos_embd_params, + dtype=DTYPE, + kv_cache_dtype=kv_cache_dtype, + is_indexer=is_indexer, + rotate_activation=rotate, + ).to(DEVICE) + + # Create DeepseekV4CacheManager + self.cache_manager = self._create_deepseek_v4_cache_manager(compress_ratio) + # COMPRESS cache has tokens_per_block from cache manager's compressed_block_sizes + self.tokens_per_block = self.cache_manager.compressed_block_sizes[self.layer_idx] + + # Track active requests for the cache manager + self.active_requests: Dict[int, LlmRequest] = {} + self.next_request_id = 0 + + # Store reference to kv_cache for test compatibility (read from cache manager) + self._update_kv_cache_reference() + + # Create block_offsets for test compatibility (will be updated per forward) + max_compressed = MAX_SEQ // compress_ratio + max_comp_blocks = (max_compressed + self.tokens_per_block - 1) // self.tokens_per_block + self.block_offsets = torch.zeros( + MAX_BATCH, max_comp_blocks, device=DEVICE, dtype=torch.int32 + ) + + def cleanup(self): + """Free all active requests and shut down the cache manager. + + Must be called before the wrapper is discarded, otherwise the cache + manager's destructor will fail with ResourceBusyError and GPU memory + will leak across tests. + """ + if not hasattr(self, "cache_manager"): + return + for req in self.active_requests.values(): + self.cache_manager.free_resources(req) + self.active_requests.clear() + self.cache_manager.shutdown() + + def __del__(self): + try: + self.cleanup() + except Exception: + pass + + def _create_deepseek_v4_cache_manager(self, compress_ratio: int) -> DeepseekV4CacheManager: + """Create a DeepseekV4CacheManager for testing.""" + # Single layer with the given compress ratio + compress_ratios = [compress_ratio] + + # Create sparse attention config + sparse_attn_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=INDEX_HEAD_DIM, + window_size=self.WINDOW_SIZE, + compress_ratios=compress_ratios, + ) + + # Create KV cache config + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_tokens=MAX_SEQ * MAX_BATCH, + event_buffer_max_size=0, + ) + + # Create mapping (single GPU, no parallelism) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + if self.kv_cache_dtype in ["fp8_pertensor", "fp8_blockwise"]: + cache_dtype = DataType.FP8 + else: + cache_dtype = DataType.BF16 + + # Create cache manager + cache_manager = DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(compress_ratios), + num_kv_heads=1, + head_dim=HEAD_DIM, + tokens_per_block=PAGE_SIZE, + max_seq_len=MAX_SEQ, + max_batch_size=MAX_BATCH, + max_input_len=MAX_SEQ, + mapping=mapping, + dtype=cache_dtype, + compressor_dtype=DataType.FLOAT, # State caches always use FP32 + vocab_size=self.VOCAB_SIZE, + max_num_tokens=MAX_SEQ + MAX_BATCH, + sparse_attn_config=sparse_attn_config, + ) + + return cache_manager + + def _update_kv_cache_reference(self): + """Update the kv_cache reference from cache manager for test compatibility.""" + compress_type = ( + DeepseekV4AttentionType.INDEXER_COMPRESS + if self.is_indexer + else DeepseekV4AttentionType.COMPRESS + ) + self.kv_cache = self.cache_manager.get_buffers(self.layer_idx, compress_type) + + def _create_request(self, request_id: int, prompt_len: int) -> LlmRequest: + """Helper to create a test LlmRequest (following test_deepseek_v4_cache_manager pattern). + + Args: + request_id: Unique request identifier + prompt_len: Prompt length (number of tokens) + + Returns: + LlmRequest instance + """ + input_tokens = list(range(prompt_len)) + request = LlmRequest( + request_id=request_id, + max_new_tokens=1024, + input_tokens=input_tokens, + sampling_config=SamplingConfig(), + is_streaming=False, + ) + return request + + def _prepare_requests_for_batch( + self, + bsz: int, + seq_lens: torch.Tensor, + start_pos: torch.Tensor, + is_prefill: torch.Tensor, + batch_indices: torch.Tensor = None, + ) -> Tuple[List[LlmRequest], ScheduledRequests]: + """Prepare requests for a batch using the KVCacheV2 scheduler allocation flow. + + Args: + batch_indices: Optional tensor mapping batch positions to external batch indices. + If provided, requests are stored/retrieved using these indices. + If None, sequential indices (0, 1, 2...) are used. + + Returns: + Tuple of (requests list, scheduled_batch) for later update_resources call + """ + requests = [] + context_requests = [] + generation_requests = [] + + # Use batch_indices if provided, otherwise use sequential indices + if batch_indices is not None: + ext_indices = [int(batch_indices[b].item()) for b in range(bsz)] + else: + ext_indices = list(range(bsz)) + + # Separate prefill and generation indices. + prefill_indices = [] + gen_indices = [] + for b in range(bsz): + if is_prefill[b]: + prefill_indices.append(b) + else: + gen_indices.append(b) + + # Handle prefill requests, including chunked prefill with start_pos > 0. + for b in prefill_indices: + ext_idx = ext_indices[b] + req = self.active_requests.get(ext_idx) + total_prompt_len = int((start_pos[b] + seq_lens[b]).item()) + chunk_size = int(seq_lens[b].item()) + if req is None: + assert int(start_pos[b].item()) == 0, ( + "Chunked prefill requests must reuse an existing request created by " + "the initial context chunk" + ) + req = self._create_request(self.next_request_id, total_prompt_len) + self.active_requests[ext_idx] = req + self.next_request_id += 1 + req.state = LlmRequestState.CONTEXT_INIT + req.context_current_position = int(start_pos[b].item()) + req.prompt_len = total_prompt_len + req.py_prompt_len = total_prompt_len + req.context_chunk_size = chunk_size + req.py_draft_tokens = [] + context_requests.append(req) + + # Handle generation requests - reuse existing requests or create new ones + for b in gen_indices: + ext_idx = ext_indices[b] + req = self.active_requests.get(ext_idx) + req_seq_len = int(seq_lens[b].item()) + if req is None: + # Need to create a new request for generation (prefill was done in previous call) + pos = int(start_pos[b].item()) + req = self._create_request(self.next_request_id, pos) + req.state = LlmRequestState.GENERATION_IN_PROGRESS + req.context_current_position = pos + req.add_new_token(pos, 0) # Simulate having processed tokens + self.active_requests[ext_idx] = req + self.next_request_id += 1 + else: + # Existing request from previous prefill - mark as generation + req.state = LlmRequestState.GENERATION_IN_PROGRESS + # Scheduler v2 allocates 1 sampled token plus any draft tokens. + # Mirror that contract so multi-token generation reserves enough KV slots. + req.py_draft_tokens = [0] * max(req_seq_len - 1, 0) + generation_requests.append(req) + + # Build final request list in batch order + for b in range(bsz): + ext_idx = ext_indices[b] + requests.append(self.active_requests[ext_idx]) + + # Build scheduled batch and allocate buffers like KVCacheV2Scheduler. + scheduled_batch = ScheduledRequests() + for req in context_requests: + scheduled_batch.append_context_request(req) + scheduled_batch.generation_requests = generation_requests + for req in context_requests: + assert self.cache_manager.prepare_context(req), ( + f"Failed to prepare context for request {req.py_request_id}" + ) + assert self.cache_manager.resize_context(req, req.context_chunk_size), ( + f"Failed to resize context for request {req.py_request_id}" + ) + for req in generation_requests: + assert self.cache_manager.try_allocate_generation(req), ( + f"Failed to allocate generation KV cache for request {req.py_request_id}" + ) + + return requests, scheduled_batch + + def _get_block_table_for_request( + self, + req: LlmRequest, + attn_type: DeepseekV4AttentionType, + ) -> torch.Tensor: + """Get block table for a request and attention type.""" + page_indices = self.cache_manager.get_cache_indices( + request_id=req.py_request_id, + layer_idx=self.layer_idx, + attn_type=attn_type, + ) + return torch.tensor(page_indices, dtype=torch.int32, device=DEVICE) + + def forward( + self, + x: torch.Tensor, + start_pos: int | torch.Tensor, + freqs_cis: torch.Tensor, + batch_indices: torch.Tensor = None, + seq_lens: torch.Tensor = None, + *, + is_prefill: torch.Tensor | None = None, + ): + """Wrapper forward that matches the reference Compressor interface. + + Supports mixed prefill+decode when seq_lens and start_pos tensor are provided. + """ + ratio = self.compress_ratio + + def normalize_is_prefill( + prefill: torch.Tensor | None, default: torch.Tensor + ) -> torch.Tensor: + if prefill is None: + prefill_tensor = default + else: + assert isinstance(prefill, torch.Tensor), "is_prefill must be a torch.Tensor" + prefill_tensor = prefill.to(device=DEVICE, dtype=torch.bool) + if prefill_tensor.ndim == 0: + prefill_tensor = prefill_tensor.expand(bsz) + assert prefill_tensor.shape == (bsz,), ( + f"is_prefill must have shape ({bsz},), got {tuple(prefill_tensor.shape)}" + ) + return prefill_tensor + + # Handle variable-length sequences + if seq_lens is not None: + # Mixed batch mode with variable-length sequences + seq_lens = seq_lens.to(torch.int32) + bsz = seq_lens.size(0) + if isinstance(start_pos, torch.Tensor): + start_pos_tensor = start_pos.to(torch.int32) + else: + start_pos_tensor = torch.full((bsz,), start_pos, dtype=torch.int32, device=DEVICE) + + # Flatten input tokens + if x.ndim == 3: + x_flat = x.view(-1, DIM) + else: + x_flat = x + total_tokens = int(seq_lens.sum().item()) + x_flat = x_flat[:total_tokens] + + is_prefill_tensor = normalize_is_prefill( + is_prefill, (start_pos_tensor == 0) | (seq_lens > 1) + ) + + # Determine which sequences are context (prefill) vs generation (decode). + is_context = is_prefill_tensor + num_contexts = int(is_context.sum().item()) + num_generations = bsz - num_contexts + + # Reorder: contexts first, then generations + ctx_indices = torch.where(is_context)[0] + gen_indices = torch.where(~is_context)[0] + reorder_indices = torch.cat([ctx_indices, gen_indices]) + + seq_lens_reordered = seq_lens[reorder_indices] + start_pos_reordered = start_pos_tensor[reorder_indices] + + # Compute token offsets for reordering + cu_seq_original = torch.zeros(bsz + 1, dtype=torch.int32, device=DEVICE) + cu_seq_original[1:] = seq_lens.cumsum(0) + + # Reorder tokens: contexts first, then generations + token_indices = [] + for idx in reorder_indices: + start = cu_seq_original[idx].item() + end = cu_seq_original[idx + 1].item() + token_indices.extend(range(start, end)) + x_flat = x_flat[token_indices] + + num_ctx_tokens = ( + int(seq_lens_reordered[:num_contexts].sum().item()) if num_contexts > 0 else 0 + ) + num_gen_tokens = ( + int(seq_lens_reordered[num_contexts:].sum().item()) if num_generations > 0 else 0 + ) + seq_lens = seq_lens_reordered + past_kv_lens = start_pos_reordered + is_prefill = is_context[reorder_indices] + + # Use reorder_indices for request mapping (maps reordered position to original batch index) + batch_indices_for_requests = reorder_indices + start_pos_for_kv = past_kv_lens + + # Get number of compressed tokens + num_comp_per_seq = (past_kv_lens + seq_lens) // ratio - past_kv_lens // ratio + num_ctx_compressed_tokens = int(num_comp_per_seq[:num_contexts].sum().item()) + num_gen_compressed_tokens = int(num_comp_per_seq[num_contexts:].sum().item()) + else: + # Original single-mode logic + bsz, seqlen, _ = x.size() + x_flat = x.reshape(-1, DIM) + if isinstance(start_pos, torch.Tensor): + past_kv_lens = start_pos.to(torch.int32).to(DEVICE) + if past_kv_lens.ndim == 0: + past_kv_lens = past_kv_lens.expand(bsz) + else: + past_kv_lens = torch.full((bsz,), start_pos, dtype=torch.int32, device=DEVICE) + + is_prefill = normalize_is_prefill( + is_prefill, torch.full((bsz,), seqlen > 1, dtype=torch.bool, device=DEVICE) + ) + if not torch.all(is_prefill == is_prefill[0]): + raise ValueError( + "single-shape forward requires uniform is_prefill across the batch" + ) + is_prefill_value = bool(is_prefill[0].item()) + + if not is_prefill_value: + # Decode mode + num_contexts = 0 + num_generations = bsz + num_ctx_tokens = 0 + num_gen_tokens = bsz * seqlen + seq_lens = torch.full((bsz,), seqlen, dtype=torch.int32, device=DEVICE) + num_ctx_compressed_tokens = 0 + kv_lens_local = past_kv_lens + seq_lens + num_gen_compressed_tokens = int( + (kv_lens_local // ratio - past_kv_lens // ratio).sum().item() + ) + else: + # Prefill mode (may be chunked when start_pos > 0) + num_contexts = bsz + num_generations = 0 + num_ctx_tokens = bsz * seqlen + num_gen_tokens = 0 + seq_lens = torch.full((bsz,), seqlen, dtype=torch.int32, device=DEVICE) + kv_lens_local = past_kv_lens + seq_lens + num_ctx_compressed_tokens = int( + (kv_lens_local // ratio - past_kv_lens // ratio).sum().item() + ) + num_gen_compressed_tokens = 0 + # Use batch_indices for request mapping if provided + batch_indices_for_requests = batch_indices + start_pos_for_kv = past_kv_lens + + # Prepare requests for the cache manager + requests, scheduled_batch = self._prepare_requests_for_batch( + bsz, seq_lens, start_pos_for_kv, is_prefill, batch_indices_for_requests + ) + + cu_seq_lens = torch.zeros(bsz + 1, dtype=torch.int32, device=DEVICE) + cu_seq_lens[1:] = seq_lens.cumsum(0) + num_total_compressed_tokens = num_ctx_compressed_tokens + num_gen_compressed_tokens + + # Compute KV lengths (past + current) per sequence + kv_lens = past_kv_lens + seq_lens + + # Compute number of compressed outputs per batch using absolute-aligned formula. + # kv_len // ratio - past // ratio works for all cases: + # fresh prefill (past=0): kv_len // ratio + # chunked prefill (past>0, seqlen>1): correct window boundary counting + # generation (past>0): fires when chunk boundary is crossed + num_comp = (past_kv_lens + seq_lens) // ratio - past_kv_lens // ratio + + cu_new_comp_kv = torch.zeros(bsz + 1, dtype=torch.int32, device=DEVICE) + cu_new_comp_kv[1:] = num_comp.cumsum(0) + max_ctx_comp_kv_lens = num_comp[:num_contexts].max().item() if num_contexts > 0 else 0 + + # Create position IDs for compressed outputs. + # Always use first-token-of-window convention: position = (base_chunk + c) * ratio. + # This matches RefCompressor which uses win_first = pos + 1 - ratio for all cases. + position_ids = torch.zeros(num_total_compressed_tokens, dtype=torch.int32, device=DEVICE) + offset = 0 + for b in range(bsz): + n_out = num_comp[b].item() + sp = past_kv_lens[b].item() + base_chunk = sp // ratio + for c in range(n_out): + position_ids[offset + c] = (base_chunk + c) * ratio + offset += n_out + + # Determine attention types based on is_indexer + if self.is_indexer: + compress_type = DeepseekV4AttentionType.INDEXER_COMPRESS + state_type = DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE + score_type = DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE + else: + compress_type = DeepseekV4AttentionType.COMPRESS + state_type = DeepseekV4AttentionType.COMPRESSOR_STATE + score_type = DeepseekV4AttentionType.COMPRESSOR_SCORE + + # Build block_tables dict keyed by DeepseekV4AttentionType using cache manager + block_table_compress_list = [] + block_table_kv_state_list = [] + block_table_score_state_list = [] + + for b, req in enumerate(requests): + block_table_compress_list.append(self._get_block_table_for_request(req, compress_type)) + block_table_kv_state_list.append(self._get_block_table_for_request(req, state_type)) + block_table_score_state_list.append(self._get_block_table_for_request(req, score_type)) + + # Pad and stack block tables to handle variable-length block indices + max_blocks_compress = max(bt.size(0) for bt in block_table_compress_list) + block_table_compress = torch.zeros( + bsz, max_blocks_compress, dtype=torch.int32, device=DEVICE + ) + for b in range(bsz): + bt = block_table_compress_list[b] + block_table_compress[b, : bt.size(0)] = bt + + max_blocks_state = max(bt.size(0) for bt in block_table_kv_state_list) + block_table_kv_state = torch.zeros(bsz, max_blocks_state, dtype=torch.int32, device=DEVICE) + block_table_score_state = torch.zeros( + bsz, max_blocks_state, dtype=torch.int32, device=DEVICE + ) + for b in range(bsz): + bt_kv = block_table_kv_state_list[b] + bt_score = block_table_score_state_list[b] + block_table_kv_state[b, : bt_kv.size(0)] = bt_kv + block_table_score_state[b, : bt_score.size(0)] = bt_score + + # Update block_offsets for test compatibility + self.block_offsets = block_table_compress + + block_tables = { + (ratio, compress_type): block_table_compress, + (ratio, state_type): block_table_kv_state, + (ratio, score_type): block_table_score_state, + } + + # Both prefill and decode kernels use absolute token positions for the + # state cache, so pass the absolute kv_lens directly. + # Build dicts keyed by compress_ratio + cu_new_comp_kv_dict = {ratio: cu_new_comp_kv} + compressed_position_ids_dict = {ratio: position_ids} + compressed_kv_lens_dict = {ratio: kv_lens} + past_kv_lens_dict = {ratio: past_kv_lens // ratio} + new_comp_kv_lens_cuda_dict = {ratio: num_comp} + num_total_compressed_tokens_dict = {ratio: num_total_compressed_tokens} + max_ctx_compressed_tokens_dict = {ratio: max_ctx_comp_kv_lens} + + # Build per-token compressed_mask + compressed_mask_tokens = torch.zeros( + num_total_compressed_tokens, dtype=torch.bool, device=DEVICE + ) + offset = 0 + for b in range(bsz): + n = int(num_comp[b].item()) + compressed_mask_tokens[offset : offset + n] = True + slot_size = int(cu_new_comp_kv[b + 1].item() - cu_new_comp_kv[b].item()) + offset += slot_size + compressed_mask_cuda_dict = {ratio: compressed_mask_tokens} + + # Build attention metadata using DeepseekV4CacheManager + metadata = DummyAttentionMetadata( + num_contexts=num_contexts, + num_generations=num_generations, + num_ctx_tokens=num_ctx_tokens, + num_tokens=num_ctx_tokens + num_gen_tokens, + kv_cache_manager=self.cache_manager, + block_tables=block_tables, + cu_seq_lens=cu_seq_lens, + cu_new_comp_kv=cu_new_comp_kv_dict, + compressed_position_ids=compressed_position_ids_dict, + compressed_kv_lens=compressed_kv_lens_dict, + past_kv_lens=past_kv_lens_dict, + new_comp_kv_lens_cuda=new_comp_kv_lens_cuda_dict, + num_total_compressed_tokens=num_total_compressed_tokens_dict, + max_ctx_compressed_tokens=max_ctx_compressed_tokens_dict, + compressed_mask_cuda=compressed_mask_cuda_dict, + ) + # kv_lens_cuda_runtime: [num_seqs] total KV length per sequence (past + current) + metadata.kv_lens_cuda_runtime = kv_lens + metadata.cached_token_lens_cuda = past_kv_lens + metadata.num_gen_tokens_per_seq = ( + num_gen_tokens // num_generations if num_generations > 0 else 0 + ) + + # Update kv_cache reference for test compatibility + self._update_kv_cache_reference() + + # Call the compressor forward + result = self.compressor(x=x_flat, metadata=metadata) + + # Update request state and call update_resources after processing + for b, req in enumerate(requests): + token_count = int((start_pos_for_kv[b] + seq_lens[b]).item()) + if is_prefill[b]: + req.context_current_position = token_count + # Call add_new_token for BOTH prefill and generation requests. + req.add_new_token(token_count, 0) + self.cache_manager.update_resources(scheduled_batch) + + # Compressor.forward() returns (kv_comp, scale) tuple. + # For FP8 blockwise (indexer), scale is non-None → return directly. + if isinstance(result, tuple): + kv_comp, scale = result + if scale is not None: + return kv_comp, scale + else: + kv_comp = result + total_outputs = cu_new_comp_kv[-1].item() + if total_outputs == 0: + return None + + # Fused scatter writes postprocessed data to kv_cache but returns raw + # kv_comp. Apply postprocessing inline for test comparison with reference. + if kv_comp is not None and self.kv_cache_dtype == "fp8_pertensor": + # Read FP8 data directly from cache (written by fused kernel) to + # ensure golden cache comparison matches exactly. + all_fp8 = [] + for b in range(bsz): + n = cu_new_comp_kv[b + 1].item() - cu_new_comp_kv[b].item() + if n > 0: + tokens = read_paged_cache_tokens( + self.kv_cache, self.block_offsets, b, n, self.tokens_per_block + ) + all_fp8.append(tokens) + if all_fp8: + kv_fp8 = torch.cat(all_fp8, dim=0).view(torch.float8_e4m3fn) + kv_scale = torch.ones(1, dtype=torch.float32, device=kv_comp.device) + return kv_fp8, kv_scale + return None + elif kv_comp is not None and self.kv_cache_dtype == "default": + # Replay the kernel's fp32-throughout postprocess (RMSNorm + RoPE + # + optional Hadamard) so `out_comp` matches the kernel's cache + # values byte-for-byte. + nope_dim = self.compressor.nope_head_dim + rope_dim = self.compressor.rope_head_dim + kv_proc = kv_comp.clone().float() + var = kv_proc.pow(2).mean(-1, keepdim=True) + kv_proc = kv_proc * torch.rsqrt(var + self.compressor.norm.variance_epsilon) + kv_proc = kv_proc * self.compressor.norm.weight.float() + + pos_ids = metadata.compressed_position_ids_cuda[self.compress_ratio][: kv_proc.shape[0]] + cos_sin = self.compressor.rotary_emb.rotary_cos_sin.float() + half_rope = rope_dim // 2 + cos_v = cos_sin[pos_ids.long(), 0, :] + sin_v = cos_sin[pos_ids.long(), 1, :] + xn = kv_proc[:, :nope_dim] + xp = kv_proc[:, nope_dim:].view(-1, half_rope, 2) + x_even, x_odd = xp[..., 0], xp[..., 1] + xp = torch.stack( + [x_even * cos_v - x_odd * sin_v, x_odd * cos_v + x_even * sin_v], dim=-1 + ).view(-1, rope_dim) + kv_proc = torch.cat([xn, xp], dim=-1) + + if self.compressor.rotate_activation: + kv_proc = rotate_activation(kv_proc) + kv_comp = kv_proc.to(kv_comp.dtype) + + # Split packed output back to per-batch + outputs = [] + for b in range(bsz): + start = cu_new_comp_kv[b].item() + end = cu_new_comp_kv[b + 1].item() + if end > start: + outputs.append(kv_comp[start:end]) + else: + outputs.append(None) + + # If all batches have the same number of outputs, stack them + non_none_outputs = [o for o in outputs if o is not None] + if len(non_none_outputs) == 0: + return None + elif all(o is not None and o.size(0) == outputs[0].size(0) for o in outputs): + return torch.stack(outputs, dim=0) + else: + # Concatenate all non-None outputs along dim=0, then unsqueeze for batch dim + # This handles mixed batches with different compression counts + return torch.cat(non_none_outputs, dim=0).unsqueeze(0) + + def reset_state(self): + """Reset cache manager state for new sequences.""" + # Shutdown existing cache manager to release resources properly + if hasattr(self, "cache_manager") and self.cache_manager is not None: + self.cache_manager.shutdown() + + # Clear active requests and reset request ID counter + self.active_requests.clear() + self.next_request_id = 0 + + # Recreate the cache manager to reset all caches + self.cache_manager = self._create_deepseek_v4_cache_manager(self.compress_ratio) + self._update_kv_cache_reference() + + +def setup_compressors( + compress_ratio: int = 4, + rotate: bool = False, + kv_cache_dtype: str = "default", +): + """Create synced RefCompressor + Compressor with all caches initialized. + + Args: + compress_ratio: Compression ratio (default 4) + rotate: Whether to apply Hadamard rotation + kv_cache_dtype: Cache dtype preset - "default", "fp8_pertensor", + or "fp8_blockwise" + + Returns: + ref: RefCompressor (bf16 reference) + comp: CompressorWrapper (with specified kv_cache_dtype) + """ + args = ModelArgs() + overlap = compress_ratio == 4 + + # For indexer modes with ratio=4, use INDEXER_COMPRESS type and INDEX_HEAD_DIM=128. + # INDEXER_COMPRESS is only registered for sparse layers and uses INDEX_HEAD_DIM=128 + is_indexer = kv_cache_dtype == "fp8_blockwise" + + # Use appropriate head_dim based on is_indexer + target_head_dim = INDEX_HEAD_DIM if is_indexer else HEAD_DIM + + # Reference compressor (bf16) + ref = RefCompressor(args, compress_ratio, target_head_dim, rotate).to(DEVICE) + ref.ape.data.normal_(0, 0.02) + ref.wkv.weight.data.normal_(0, 0.02) + ref.wgate.weight.data.normal_(0, 0.02) + ref.kv_cache = torch.zeros( + MAX_BATCH, MAX_SEQ // compress_ratio, target_head_dim, device=DEVICE, dtype=DTYPE + ) + + # Compressor wrapper with specified kv_cache_dtype + comp = CompressorWrapper( + compress_ratio, rotate, kv_cache_dtype=kv_cache_dtype, is_indexer=is_indexer + ) + + # Copy weights from ref to compressor + coff = 2 if overlap else 1 + comp.compressor.wkv_gate.weight.data[: coff * target_head_dim] = ref.wkv.weight.data.clone() + comp.compressor.wkv_gate.weight.data[coff * target_head_dim :] = ref.wgate.weight.data.clone() + comp.compressor.ape.data.copy_(ref.ape.data) + comp.compressor.norm.weight.data.copy_(ref.norm.weight.data) + + return ref, comp + + +@pytest.fixture(autouse=True) +def seed(): + """Seed RNG for reproducibility and ensure GPU cleanup between tests.""" + torch.manual_seed(42) + torch.cuda.manual_seed(42) + yield + # Force garbage collection so CompressorWrapper.__del__ → cleanup() runs + # before the next test allocates new cache managers. + import gc + + gc.collect() + torch.cuda.empty_cache() + + +# ============================================================================ +# Tests +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (1, 2, 4), + (1, 64, 4), + (2, 256, 4), + (4, 128, 4), # batch/seqlen variations + (1, 256, 128), + (2, 512, 128), + (16, 234, 128), # ratio=128 coverage + ], +) +def test_prefill(batch, seqlen, ratio): + """Test prefill mode.""" + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert_similar(out_ref, out_comp) + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, out_ref[b : b + 1], f"Prefill ref cache[{b}]") + assert_similar(cached_comp, out_comp[b : b + 1], f"Prefill comp cache[{b}]") + assert_similar(cached_ref, cached_comp, f"Prefill cache parity[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_decode(prefill, steps, batch, ratio): + """Test prefill + decode.""" + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + # Prefill + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + assert_similar(ref(x, 0, freqs[:prefill]), comp.forward(x, 0, freqs[:prefill]), "Prefill") + + # Decode + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + out_comp = comp.forward(x, pos, freqs[pos : pos + 1]) + assert_similar(out_ref, out_comp, f"Decode[{step}]") + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar( + cached_ref[:, -1:], out_ref[b : b + 1], f"Decode cache ref[{b}] step{step}" + ) + assert_similar( + cached_comp[:, -1:], + out_comp[b : b + 1], + f"Decode cache comp[{b}] step{step}", + ) + assert_similar(cached_ref, cached_comp, f"Decode cache parity[{b}] step{step}") + + +def test_varlen_batch(): + """Test variable-length prefill batch, compare with reference.""" + seq_lens = [64, 96, 128] + ratio = 4 + + # Test each sequence independently + for i, slen in enumerate(seq_lens): + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(1, slen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + # Reset ref's state for each independent sequence + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + out_ref = ref(x, 0, freqs[:slen]) + out_comp = comp.forward(x, 0, freqs[:slen]) + + assert_similar(out_ref, out_comp, f"Varlen seq{i}") + + +def test_mixed_batch(): + """Test mixed context + generation requests in a single forward call. + + Simulates a realistic mixed batch with: + - 1 context request: 8 tokens, start_pos=0 + - 1 generation request: 1 token, start_pos=127 (triggers compression at 128) + + Generation requests have seqlen=1 and require pre-populated state. + """ + ratio = 4 + # rotate=True required: Compressor unconditionally applies Hadamard rotation + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + # Context request: 8 tokens at start_pos=0 → 2 compressed outputs + ctx_len = 8 + x_ctx = torch.randn(1, ctx_len, DIM, device=DEVICE, dtype=DTYPE) + + # Generation request: 1 token at start_pos=127 → triggers compression at 128 + # (127 + 1) % 4 == 0, so compression is triggered + gen_start_pos = 127 + x_gen = torch.randn(1, 1, DIM, device=DEVICE, dtype=DTYPE) + + # For ref: need to pre-populate state by running prefill of gen_start_pos tokens + # This simulates the generation request having processed gen_start_pos tokens already + x_gen_prefill = torch.randn(1, gen_start_pos, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + # === Reference: run each request separately === + + # Context request on ref (batch 0) + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + out_ref_ctx = ref(x_ctx, 0, freqs[:ctx_len]) + + # Generation request on ref (batch 0, but independent - reset state first) + # Pre-populate state by running prefill of gen_start_pos tokens + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + _ = ref(x_gen_prefill, 0, freqs[:gen_start_pos]) # Sets up state + # Now run the decode token + out_ref_gen = ref(x_gen, gen_start_pos, freqs) + + # Collect non-None outputs + ref_outputs = [o for o in [out_ref_ctx, out_ref_gen] if o is not None] + out_ref = torch.cat(ref_outputs, dim=1) if ref_outputs else None + + # === Compressor: single forward with mixed batch === + # Flatten tokens: [ctx_tokens, gen_token] + x_flat = torch.cat([x_ctx.squeeze(0), x_gen.squeeze(0)], dim=0) + + # Sequence lengths and start positions for each request + seq_lens = torch.tensor([ctx_len, 1], dtype=torch.int32, device=DEVICE) + start_pos_tensor = torch.tensor([0, gen_start_pos], dtype=torch.int32, device=DEVICE) + + # Pre-populate compressor's paged state for generation request (batch idx 1) + # by running prefill on batch 1 first + comp.reset_state() + comp.forward( + x_gen_prefill, + 0, + freqs[:gen_start_pos], + batch_indices=torch.tensor([1], device=DEVICE, dtype=torch.int32), + ) + + # Now run mixed batch + out_comp = comp.forward(x_flat, start_pos_tensor, freqs, seq_lens=seq_lens) + + if out_ref is None: + assert out_comp is None, "Mixed batch: expected no compression output" + else: + assert_similar(out_ref, out_comp, "Mixed batch context+generation single forward") + + +# ============================================================================ +# FP8 Blockwise Quantization Tests +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 64, 4), + (4, 256, 4), + (16, 256, 4), + (16, 512, 4), + ], +) +def test_fp8_blockwise_compressor(batch, seqlen, ratio): + """Test FP8 blockwise Compressor against RefCompressor (bf16 reference). + + FP8 blockwise uses INDEXER_COMPRESS cache type which has INDEX_HEAD_DIM=128. + This test validates FP8 quantization and cache scatter for the indexer path. + """ + num_compressed = seqlen // ratio + + ref, comp = setup_compressors(ratio, rotate=True, kv_cache_dtype="fp8_blockwise") + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert isinstance(out_comp, tuple), f"Expected tuple, got {type(out_comp)}" + kv_fp8, kv_scale = out_comp + + # Verify shapes - INDEXER_COMPRESS uses INDEX_HEAD_DIM=128 + total_comp_tokens = batch * num_compressed + num_scale_blocks = (INDEX_HEAD_DIM + 127) // 128 + assert kv_fp8.shape == (total_comp_tokens, INDEX_HEAD_DIM), ( + f"Expected shape {(total_comp_tokens, INDEX_HEAD_DIM)}, got {kv_fp8.shape}" + ) + assert kv_scale.shape == (total_comp_tokens, num_scale_blocks), ( + f"Expected scale shape {(total_comp_tokens, num_scale_blocks)}, got {kv_scale.shape}" + ) + + # Verify scales are positive and valid + assert (kv_scale > 0).all(), "All scales should be positive" + + # Compare dequantized FP8 with RefCompressor output (both use INDEX_HEAD_DIM=128) + assert_fp8_similar( + out_comp, out_ref.view(-1, INDEX_HEAD_DIM), "fp8_blockwise", "FP8 vs RefCompressor" + ) + + # Verify cache scatter + golden_cache = build_fp8_golden_cache( + kv_fp8, + kv_scale, + comp.kv_cache.shape, + comp.block_offsets, + batch, + num_compressed, + comp.tokens_per_block, + "fp8_blockwise", + head_dim=INDEX_HEAD_DIM, + ) + assert_fp8_cache_match( + comp.kv_cache, + golden_cache, + "fp8_blockwise", + "Blockwise cache layout", + head_dim=INDEX_HEAD_DIM, + ) + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 256, 4), + (1, 256, 128), + (2, 512, 128), # ratio=128 coverage + ], +) +def test_fp8_pertensor_compressor(batch, seqlen, ratio): + """Test per-tensor FP8 Compressor against RefCompressor (bf16 reference).""" + num_compressed = seqlen // ratio + + ref, comp = setup_compressors(ratio, rotate=True, kv_cache_dtype="fp8_pertensor") + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert isinstance(out_comp, tuple), f"Expected tuple, got {type(out_comp)}" + kv_fp8, kv_scale = out_comp + + # Verify shapes + total_comp_tokens = batch * num_compressed + assert kv_fp8.shape == (total_comp_tokens, HEAD_DIM) + assert kv_scale.numel() == 1, f"Per-tensor scale should be scalar, got {kv_scale.shape}" + assert kv_scale.item() > 0, "Scale should be positive" + + # Compare dequantized FP8 with reference + assert_fp8_similar( + out_comp, out_ref.view(-1, HEAD_DIM), "fp8_pertensor", "FP8 vs RefCompressor" + ) + + # Verify scale is fixed at 1.0 (compressor uses static scale following trtllm.py convention) + assert kv_scale.item() == 1.0, ( + f"Per-tensor scale should be fixed at 1.0, got {kv_scale.item():.6f}" + ) + + # Verify cache scatter (cache is FP8 dtype, same as compressor output) + golden_cache = build_fp8_golden_cache( + kv_fp8, + kv_scale, + comp.kv_cache.shape, + comp.block_offsets, + batch, + num_compressed, + comp.tokens_per_block, + "fp8_pertensor", + ) + assert_fp8_cache_match(comp.kv_cache, golden_cache, "fp8_pertensor", "Per-tensor cache layout") + + +# ============================================================================ +# Fused Kernel Tests (RMSNorm + RoPE + Hadamard + Scatter) +# +# The fused CUDA kernel matches the reference model.py pipeline: RMSNorm, +# RoPE, and Hadamard all in bf16 precision (via toBf16 round-trips in CUDA). +# This means fused and unfused paths should produce nearly identical results. +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (1, 64, 4), + (2, 256, 4), + (4, 128, 4), + (1, 256, 128), + (2, 512, 128), + ], +) +def test_fused_prefill(batch, seqlen, ratio): + """Test fused prefill cache against unfused bf16 reference cache. + + Both paths now use bf16 Hadamard so results should match closely. + """ + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + comp.forward(x, 0, freqs) + + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Fused prefill cache[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_fused_decode(prefill, steps, batch, ratio): + """Test fused prefill + decode cache against bf16 reference.""" + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + ref(x, 0, freqs[:prefill]) + comp.forward(x, 0, freqs[:prefill]) + + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + comp.forward(x, pos, freqs[pos : pos + 1]) + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Fused decode cache[{b}] step{step}") + + +def test_fused_mixed_batch(): + """Test fused kernel with mixed context + generation batch.""" + ratio = 4 + ref, comp = setup_compressors(ratio, rotate=True) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + ctx_len = 8 + x_ctx = torch.randn(1, ctx_len, DIM, device=DEVICE, dtype=DTYPE) + gen_start_pos = 127 + x_gen = torch.randn(1, 1, DIM, device=DEVICE, dtype=DTYPE) + x_gen_prefill = torch.randn(1, gen_start_pos, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + ref(x_ctx, 0, freqs[:ctx_len]) + + # Save ctx cache before gen overwrites positions 0..1 + num_ctx_comp = ctx_len // ratio + cached_ref_ctx = ref.kv_cache[0:1, :num_ctx_comp].clone() + + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + _ = ref(x_gen_prefill, 0, freqs[:gen_start_pos]) + ref(x_gen, gen_start_pos, freqs) + + x_flat = torch.cat([x_ctx.squeeze(0), x_gen.squeeze(0)], dim=0) + seq_lens = torch.tensor([ctx_len, 1], dtype=torch.int32, device=DEVICE) + start_pos_tensor = torch.tensor([0, gen_start_pos], dtype=torch.int32, device=DEVICE) + + comp.reset_state() + comp.forward( + x_gen_prefill, + 0, + freqs[:gen_start_pos], + batch_indices=torch.tensor([1], device=DEVICE, dtype=torch.int32), + ) + comp.forward(x_flat, start_pos_tensor, freqs, seq_lens=seq_lens) + + # Compare context request cache (first 2 compressed tokens) + if num_ctx_comp > 0: + cached_ref = cached_ref_ctx + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, 0, num_ctx_comp, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, "Fused mixed ctx cache") + + +# ============================================================================ +# Tests with rotate_activation=False (no Hadamard transform) +# ============================================================================ + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (4, 128, 4), + (1, 256, 128), + (2, 512, 128), + ], +) +def test_prefill_no_rotate(batch, seqlen, ratio): + """Test prefill mode with rotate_activation=False (Hadamard skipped).""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + out_comp = comp.forward(x, 0, freqs) + + assert_similar(out_ref, out_comp) + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Prefill no-rotate cache[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_decode_no_rotate(prefill, steps, batch, ratio): + """Test prefill + decode with rotate_activation=False.""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + assert_similar( + ref(x, 0, freqs[:prefill]), comp.forward(x, 0, freqs[:prefill]), "Prefill no-rotate" + ) + + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + out_comp = comp.forward(x, pos, freqs[pos : pos + 1]) + assert_similar(out_ref, out_comp, f"Decode no-rotate[{step}]") + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar( + cached_ref, cached_comp, f"Decode no-rotate cache[{b}] step{step}" + ) + + +@pytest.mark.parametrize( + "batch,seqlen,ratio", + [ + (1, 128, 4), + (2, 130, 4), + (2, 256, 4), + (1, 256, 128), + (2, 512, 128), + ], +) +def test_fused_prefill_no_rotate(batch, seqlen, ratio): + """Test fused prefill cache with rotate_activation=False.""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE)[:seqlen] + x = torch.randn(batch, seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs) + comp.forward(x, 0, freqs) + + if out_ref is not None: + num_tokens = out_ref.size(1) + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"Fused prefill no-rotate cache[{b}]") + + +@pytest.mark.parametrize( + "prefill,steps,batch,ratio", + [ + (128, 8, 1, 4), + (128, 8, 2, 4), + (128, 24, 1, 4), + (128, 4, 1, 128), + ], +) +def test_fused_decode_no_rotate(prefill, steps, batch, ratio): + """Test fused prefill + decode cache with rotate_activation=False.""" + ref, comp = setup_compressors(ratio, rotate=False) + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x = torch.randn(batch, prefill, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + ref(x, 0, freqs[:prefill]) + comp.forward(x, 0, freqs[:prefill]) + + for step in range(steps): + pos = prefill + step + x = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + out_ref = ref(x, pos, freqs) + comp.forward(x, pos, freqs[pos : pos + 1]) + if out_ref is not None: + num_tokens = pos // ratio + 1 + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, num_tokens, comp.tokens_per_block + ).unsqueeze(0) + assert_similar( + cached_ref, cached_comp, f"Fused decode no-rotate cache[{b}] step{step}" + ) + + +@pytest.mark.parametrize( + "prefill_len, decode_steps, batch, ratio, rotate", + [ + # ratio=4 (overlap mode) + (1, 8, 1, 4, True), + (2, 8, 2, 4, True), + (3, 8, 1, 4, True), + (1, 4, 1, 4, False), + (3, 12, 2, 4, False), + # ratio=128 (non-overlap): these should pass + (64, 128, 1, 128, True), # half-ratio prefill, decode rest + (1, 128, 1, 128, True), # minimal prefill + (127, 4, 1, 128, True), # one token short of compression at prefill + ], +) +def test_short_prefill_then_decode(prefill_len, decode_steps, batch, ratio, rotate): + """Test prefill with fewer tokens than compress_ratio, then decode until compression. + + When prefill_len < compress_ratio, the prefill produces no compressed outputs; + all tokens are saved as remainder state. Subsequent decode tokens accumulate + in the state until compress_ratio is reached, at which point compression fires. + This tests the state handoff from prefill remainder to decode accumulation. + """ + ref, comp = setup_compressors(ratio, rotate=rotate) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + x_prefill = torch.randn(batch, prefill_len, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + # Prefill: should produce no compressed tokens (prefill_len < ratio) + out_ref = ref(x_prefill, 0, freqs[:prefill_len]) + out_comp = comp.forward( + x_prefill, + 0, + freqs[:prefill_len], + is_prefill=torch.ones(batch, dtype=torch.bool, device=DEVICE), + ) + + if prefill_len < ratio: + # No compression expected from prefill + assert out_ref is None, ( + f"Ref should produce no output for prefill_len={prefill_len} < ratio={ratio}" + ) + + # Decode: step one token at a time + for step in range(decode_steps): + pos = prefill_len + step + x_decode = torch.randn(batch, 1, DIM, device=DEVICE, dtype=DTYPE) + out_ref = ref(x_decode, pos, freqs) + out_comp = comp.forward(x_decode, pos, freqs[pos : pos + 1]) + + should_compress = (pos + 1) % ratio == 0 + if should_compress: + assert out_ref is not None, f"Ref should compress at step {step} (pos={pos})" + assert_similar(out_ref, out_comp, f"Short prefill decode step {step}") + + # Verify cache parity + num_comp_tokens = (pos + 1) // ratio + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :num_comp_tokens] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, + comp.block_offsets, + b, + num_comp_tokens, + comp.tokens_per_block, + ).unsqueeze(0) + assert_similar( + cached_ref, + cached_comp, + f"Short prefill cache parity[{b}] step{step}", + ) + else: + assert out_ref is None, f"Ref should NOT compress at step {step} (pos={pos})" + finally: + comp.cleanup() + + +MTP_DECODE_CASES = [ + # No prior compressed output: decode spans the first compression boundary. + (1, 8, 1, 4, 2), + (3, 8, 1, 4, 3), + # Prefill ends exactly on or just after a compression boundary. + (4, 8, 1, 4, 2), + (5, 8, 1, 4, 3), + # Large absolute positions: exercise the same boundary cases after many windows. + (127, 8, 1, 4, 2), + (128, 8, 2, 4, 3), +] + + +@pytest.mark.parametrize("prefill_len, decode_steps, batch, ratio, next_n", MTP_DECODE_CASES) +def test_mtp_decode_overlap(prefill_len, decode_steps, batch, ratio, next_n): + """MTP decode: ref(combined seqlen=n) == ref(chunked seqlen=1 each). + + Verifies that RefCompressor called once with all next_n tokens produces + the same output as calling it one token at a time — i.e., the else-branch + for-loop is equivalent to sequential single-token calls. + """ + ref_combined, comp = setup_compressors(ratio, rotate=True) + try: + # Build ref_chunked with identical weights + args = ModelArgs( + dim=DIM, + head_dim=HEAD_DIM, + rope_head_dim=ROPE_DIM, + max_seq_len=MAX_SEQ, + max_batch_size=MAX_BATCH, + ) + ref_chunked = RefCompressor(args, compress_ratio=ratio, head_dim=HEAD_DIM, rotate=True).to( + DEVICE + ) + ref_chunked.wkv.weight.data.copy_(ref_combined.wkv.weight.data) + ref_chunked.wgate.weight.data.copy_(ref_combined.wgate.weight.data) + ref_chunked.ape.data.copy_(ref_combined.ape.data) + ref_chunked.norm.weight.data.copy_(ref_combined.norm.weight.data) + ref_chunked.kv_cache = torch.zeros_like(ref_combined.kv_cache) + + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x_prefill = torch.randn(batch, prefill_len, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + ref_combined(x_prefill, 0, freqs[:prefill_len]) + ref_chunked(x_prefill, 0, freqs[:prefill_len]) + + step = 0 + while step < decode_steps: + n = min(next_n, decode_steps - step) + pos = prefill_len + step + x_decode = torch.randn(batch, n, DIM, device=DEVICE, dtype=DTYPE) + + # Combined: one forward call with all n tokens + out_combined = ref_combined(x_decode, pos, freqs) + + # Chunked: one token at a time + last_chunked_out = None + for t in range(n): + out_t = ref_chunked(x_decode[:, t : t + 1], pos + t, freqs) + if out_t is not None: + last_chunked_out = out_t + + if out_combined is not None: + assert last_chunked_out is not None, ( + f"MTP step {step}: chunked produced no output but combined compressed" + ) + assert_similar(out_combined, last_chunked_out, f"MTP decode step {step}") + for b in range(batch): + assert_similar( + ref_combined.kv_cache[b : b + 1], + ref_chunked.kv_cache[b : b + 1], + f"MTP cache[{b}] step {step}", + ) + else: + assert last_chunked_out is None, ( + f"MTP step {step}: chunked produced unexpected output" + ) + + step += n + finally: + comp.cleanup() + + +@pytest.mark.parametrize("prefill_len, decode_steps, batch, ratio, next_n", MTP_DECODE_CASES) +def test_mtp_decode_overlap_module(prefill_len, decode_steps, batch, ratio, next_n): + """MTP decode through CompressorWrapper requires explicit generation mode.""" + ref, comp = setup_compressors(ratio, rotate=True) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x_prefill = torch.randn(batch, prefill_len, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + ref(x_prefill, 0, freqs[:prefill_len]) + comp.forward( + x_prefill, + 0, + freqs[:prefill_len], + is_prefill=torch.ones(batch, dtype=torch.bool, device=DEVICE), + ) + + step = 0 + while step < decode_steps: + n = min(next_n, decode_steps - step) + pos = prefill_len + step + x_decode = torch.randn(batch, n, DIM, device=DEVICE, dtype=DTYPE) + + out_ref = ref(x_decode, pos, freqs) + out_comp = comp.forward( + x_decode, + pos, + freqs, + is_prefill=torch.zeros(batch, dtype=torch.bool, device=DEVICE), + ) + + if out_ref is None: + assert out_comp is None, ( + f"MTP module step {step}: expected no compressed output" + ) + else: + assert out_comp is not None, ( + f"MTP module step {step}: wrapper produced no output for generation" + ) + assert_similar(out_ref, out_comp, f"MTP module decode step {step}") + + step += n + finally: + comp.cleanup() + + +@pytest.mark.parametrize( + "batch, ratio, rotate", + [ + (1, 4, True), + (2, 4, True), + (1, 4, False), + (1, 128, True), + ], +) +def test_prefill_exact_ratio(batch, ratio, rotate): + """Prefill with seqlen == compress_ratio: exactly 1 full chunk, no remainder tokens. + + In overlap mode (ratio=4) the single output has no predecessor chunk so the + overlap first-half should be zero-weighted. In non-overlap mode (ratio=128) + this is a standard single-chunk prefill. Both output and cache are verified. + """ + ref, comp = setup_compressors(ratio, rotate=rotate) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x = torch.randn(batch, ratio, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs[:ratio]) + out_comp = comp.forward(x, 0, freqs[:ratio]) + + assert out_ref is not None, "Expected compression for seqlen==ratio" + assert out_comp is not None, "Expected compression for seqlen==ratio" + assert_similar(out_ref, out_comp, "exact_ratio prefill output") + + for b in range(batch): + cached_ref = ref.kv_cache[b : b + 1, :1] + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, b, 1, comp.tokens_per_block + ).unsqueeze(0) + assert_similar(cached_ref, cached_comp, f"exact_ratio cache[{b}]") + finally: + comp.cleanup() + + +@pytest.mark.parametrize( + "seq_lens_list, ratio, rotate", + [ + ([3, 8, 5], 4, True), # short / long / medium — varied output counts + ([1, 4, 7], 4, False), # minimal, exact-ratio, one-beyond-exact + ([64, 3, 128], 128, True), # ratio=128, mixed short / long + ], +) +def test_mixed_seqlen_contexts(seq_lens_list, ratio, rotate): + """Prefill batch with variable seqlens: some < ratio (no compressed output), some >= ratio. + + Verifies that zero-output sequences do not corrupt the cu_new_comp_kv layout + or the compressed token buffer for neighbouring sequences. Each sequence is + also run through the reference independently so that the paged cache content + can be compared token-by-token. + """ + batch = len(seq_lens_list) + max_sl = max(seq_lens_list) + + ref, comp = setup_compressors(ratio, rotate=rotate) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + + xs = [torch.randn(1, sl, DIM, device=DEVICE, dtype=DTYPE) for sl in seq_lens_list] + + with torch.no_grad(): + # ---- Reference: run each sequence independently, save compressed cache ---- + ref_caches = {} + for i, sl in enumerate(seq_lens_list): + n_out = sl // ratio + if n_out > 0: + # Reset ref state so independent runs don't interfere + ref.kv_state.zero_() + ref.score_state.fill_(float("-inf")) + ref(xs[i], 0, freqs[:sl]) + # ref writes to kv_cache[0]; save before next independent run + ref_caches[i] = ref.kv_cache[0, :n_out].clone() + + # ---- Wrapper: run all sequences in one variable-length prefill call ---- + seq_lens_t = torch.tensor(seq_lens_list, dtype=torch.int32, device=DEVICE) + start_pos_t = torch.zeros(batch, dtype=torch.int32, device=DEVICE) + + # Build flat (non-padded) 2D token tensor. A padded 3D tensor would + # interleave padding zeros with real tokens after view(-1, DIM), so we + # concatenate actual token rows directly instead. + x_flat_input = torch.cat([xs[i][0] for i in range(batch)], dim=0) + + comp.forward(x_flat_input, start_pos_t, freqs[:max_sl], seq_lens=seq_lens_t) + + # ---- Compare per-sequence compressed caches ---- + for rank, sl in enumerate(seq_lens_list): + n_out = sl // ratio + if n_out > 0 and rank in ref_caches: + cached_comp = read_paged_cache_tokens( + comp.kv_cache, comp.block_offsets, rank, n_out, comp.tokens_per_block + ) + assert_similar( + ref_caches[rank].unsqueeze(0), + cached_comp.unsqueeze(0), + f"mixed_seqlen cache seq[{rank}] sl={sl}", + ) + + # ---- Decode from remainder state for sequences with leftover tokens ---- + # Find the first sequence that has a non-zero remainder and run it to + # the next compression point; this validates state hand-off. + for rank, sl in enumerate(seq_lens_list): + remainder = sl % ratio + if remainder == 0: + continue + steps_to_compress = ratio - remainder + for step in range(steps_to_compress): + pos = sl + step + x_dec = torch.randn(1, 1, DIM, device=DEVICE, dtype=DTYPE) + pos_tensor = torch.tensor([pos], dtype=torch.int32, device=DEVICE) + seq_lens_dec = torch.ones(1, dtype=torch.int32, device=DEVICE) + out = comp.forward( + x_dec, pos_tensor, freqs[pos : pos + 1], seq_lens=seq_lens_dec + ) + should_compress = (pos + 1) % ratio == 0 + if should_compress: + assert out is not None, ( + f"mixed_seqlen seq[{rank}]: expected compression at pos={pos}" + ) + else: + assert out is None, ( + f"mixed_seqlen seq[{rank}]: unexpected output at pos={pos}" + ) + break # Only test the first remainder sequence for brevity + finally: + comp.cleanup() + + +@pytest.mark.parametrize( + "total_seqlen, split_pos, batch, ratio, rotate", + [ + # overlap mode (ratio=4): aligned split + (128, 20, 1, 4, True), + (128, 48, 2, 4, True), + (128, 20, 1, 4, False), + # overlap mode: unaligned split (split_pos % ratio != 0) + (128, 5, 1, 4, True), + (128, 6, 1, 4, True), + # non-overlap mode (ratio=128): aligned split + (256, 128, 1, 128, True), + # non-overlap mode: unaligned split + (256, 50, 1, 128, True), + ], +) +def test_chunked_prefill_ref(total_seqlen, split_pos, batch, ratio, rotate): + """Validate RefCompressor full prefill == RefCompressor two-step (initial + chunked). + + Pure-Python reference test — no CUDA kernels or cache managers. + Ensures the RefCompressor's sequential token-by-token path produces + identical kv_cache contents as the bulk prefill path. + """ + import copy + + args = ModelArgs() + + ref_full = RefCompressor(args, ratio, HEAD_DIM, rotate).to(DEVICE) + ref_full.ape.data.normal_(0, 0.02) + ref_full.wkv.weight.data.normal_(0, 0.02) + ref_full.wgate.weight.data.normal_(0, 0.02) + ref_full.kv_cache = torch.zeros( + MAX_BATCH, MAX_SEQ // ratio, HEAD_DIM, device=DEVICE, dtype=DTYPE + ) + + ref_split = copy.deepcopy(ref_full) + + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x = torch.randn(batch, total_seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_full = ref_full(x, 0, freqs[:total_seqlen]) + out_1 = ref_split(x[:, :split_pos], 0, freqs[:split_pos]) + out_2 = ref_split(x[:, split_pos:], split_pos, freqs) + + parts = [p for p in [out_1, out_2] if p is not None] + if out_full is None: + assert len(parts) == 0, "Split produced output but full did not" + return + assert len(parts) > 0, "Full produced output but split did not" + combined = torch.cat(parts, dim=1) + + assert_similar(out_full, combined, "Chunked prefill ref: full vs split") + + n_comp_full = out_full.size(1) + cache_full = ref_full.kv_cache[:batch, :n_comp_full] + cache_split = ref_split.kv_cache[:batch, :n_comp_full] + assert_similar(cache_full, cache_split, "Chunked prefill ref: cache parity") + + +@pytest.mark.parametrize( + "total_seqlen, split_pos, batch, ratio", + [ + (128, 20, 1, 4), + (128, 48, 2, 4), + (256, 128, 1, 128), + ], +) +def test_chunked_prefill_module(total_seqlen, split_pos, batch, ratio): + """Validate CompressorWrapper two-step matches RefCompressor full prefill. + + Step 1: RefCompressor processes the full sequence in one call (ground truth). + Step 2: CompressorWrapper processes the sequence in two calls (initial + chunked). + The concatenated CompressorWrapper outputs must match the RefCompressor output. + """ + ref, comp = setup_compressors(ratio, rotate=True) + try: + freqs = precompute_freqs_cis( + ROPE_DIM, MAX_SEQ, ORI_SEQ_LEN, ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW + ).to(DEVICE) + x = torch.randn(batch, total_seqlen, DIM, device=DEVICE, dtype=DTYPE) + + with torch.no_grad(): + out_ref = ref(x, 0, freqs[:total_seqlen]) + + with torch.no_grad(): + out_1 = comp.forward(x[:, :split_pos], 0, freqs[:split_pos]) + out_2 = comp.forward( + x[:, split_pos:], + split_pos, + freqs[:total_seqlen], + is_prefill=torch.ones(batch, dtype=torch.bool, device=DEVICE), + ) + + parts = [p for p in [out_1, out_2] if p is not None] + if out_ref is None: + assert len(parts) == 0, "Comp produced output but ref did not" + return + assert len(parts) > 0, "Ref produced output but comp did not" + combined = ( + torch.cat(parts, dim=1) + if all(p.dim() == out_ref.dim() for p in parts) + else torch.cat(parts, dim=0).unsqueeze(0) + ) + + assert_similar(out_ref, combined, "Chunked prefill module: ref vs comp") + finally: + comp.cleanup() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py new file mode 100644 index 000000000000..56047d62da05 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the BF16 GEMM and FP32 escape-hatch paths in the DeepSeek-V4 +Compressor wkv_gate. + +The wkv_gate weight now matches the upstream V4 checkpoint dtype (bf16). The +DEEPSEEK_V4_COMPRESSOR_FP32=1 escape hatch keeps the GEMM in fp32 for +accuracy debugging. These tests cover both paths against an fp32 reference. +""" + +import os +from unittest import mock + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.modules.linear import Linear + +DEVICE = "cuda" +# Typical DeepSeek-V4 dimensions +DIM = 4096 +HEAD_DIM = 512 +STATE_DIM = 2 * HEAD_DIM # overlap=True for compress_ratio=4 +OUT_DIM = STATE_DIM * 2 # wkv + gate + + +def _create_wkv_gate(dtype: torch.dtype = torch.bfloat16) -> Linear: + """Create a Linear layer matching the compressor's wkv_gate.""" + gate = Linear( + DIM, + OUT_DIM, + bias=False, + dtype=dtype, + quant_config=None, + skip_create_weights_in_init=False, + use_custom_cublas_mm=True, + ).to(DEVICE) + gate.weight.data.normal_(0, 0.02) + return gate + + +def _bf16_path(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: + """Default path: bf16 GEMM via F.linear, then upcast to fp32 for the + downstream compression kernels.""" + return F.linear(x.to(wkv_gate.weight.dtype), wkv_gate.weight).float() + + +def _fp32_path(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: + """FP32 escape hatch: upcast both inputs to fp32 (DEEPSEEK_V4_COMPRESSOR_FP32=1).""" + return F.linear(x.float(), wkv_gate.weight.float()) + + +@pytest.fixture(autouse=True) +def seed(): + torch.manual_seed(42) + torch.cuda.manual_seed(42) + yield + torch.cuda.empty_cache() + + +@pytest.mark.parametrize("num_tokens", [1, 16, 128, 512]) +def test_bf16_path_close_to_fp32_reference(num_tokens): + """BF16 GEMM is close enough to an FP32 reference for the compressor's use. + + BF16 has 8-bit mantissa precision, but the input is already bf16 (matching + upstream V4) so the additional precision loss from a bf16 GEMM is bounded + by the input precision. + """ + bf16_gate = _create_wkv_gate(torch.bfloat16) + fp32_gate = _create_wkv_gate(torch.float32) + # Make weights numerically equivalent (rounded to bf16 in both copies) + fp32_gate.weight.data.copy_(bf16_gate.weight.data.float()) + + x = torch.randn(num_tokens, DIM, device=DEVICE, dtype=torch.bfloat16) + + with torch.no_grad(): + out_bf16 = _bf16_path(x, bf16_gate) + out_fp32 = _fp32_path(x, fp32_gate) + + assert out_bf16.shape == out_fp32.shape + + a, b = out_bf16.flatten(), out_fp32.flatten() + cos_sim = F.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item() + assert cos_sim >= 0.99, f"cos_sim={cos_sim:.6f}" + + max_diff = (a - b).abs().max().item() + scale = max(a.abs().max().item(), b.abs().max().item(), 1e-3) + rel_err = max_diff / scale + assert rel_err <= 0.1, f"rel_err={rel_err:.6f}, max_diff={max_diff:.6f}" + + +def test_env_var_selects_fp32_path(): + """DEEPSEEK_V4_COMPRESSOR_FP32=1 env var causes module-level flag to be True.""" + import importlib + + import tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor as mod + + with mock.patch.dict(os.environ, {"DEEPSEEK_V4_COMPRESSOR_FP32": "1"}): + importlib.reload(mod) + assert mod._USE_FP32_COMPRESSOR is True + + # Restore default + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("DEEPSEEK_V4_COMPRESSOR_FP32", None) + importlib.reload(mod) + assert mod._USE_FP32_COMPRESSOR is False diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py new file mode 100644 index 000000000000..8fd9515d46c4 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -0,0 +1,891 @@ +from typing import Dict, List, Optional, Tuple + +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4AttentionType, +) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._utils import binding_to_torch_dtype +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX + +_RequestCache = Dict[ + Tuple[int, DeepseekV4AttentionType], # (layer index, attention type) + Tuple[torch.Tensor, torch.Tensor | None], # (values tensor, scales tensor) +] + + +def _view_fp8_as_uint8(buffer: torch.Tensor) -> torch.Tensor: + """View an FP8 buffer as uint8. Non-FP8 buffers are returned as-is.""" + if buffer.dtype == torch.float8_e4m3fn: + return buffer.view(torch.uint8) + return buffer + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +class TestDeepseekV4CacheManager: + # deepseek_v4 specific param + head_dim = 512 + index_head_dim = 128 + window_size = 128 + vocab_size = 129280 + sparse_layer_ratio = 4 + overlap_compress_layer_ratio = 4 + + # indexer quantization config + indexer_dtype = DataType.FP8 + indexer_scale_dtype = DataType.FLOAT + indexer_quant_block_size = 128 + + # cache manager specific param + tokens_per_block = 128 + + def _is_compress_layer(self, compress_ratio: int) -> bool: + """Check if a layer uses compression based on its compress ratio. + + Args: + compress_ratio: The compression ratio for the layer + + Returns: + True if the layer uses compression (ratio > 1) + """ + return compress_ratio > 1 + + def _is_sparse_layer(self, compress_ratio: int) -> bool: + """Check if a layer uses sparse attention based on its compress ratio. + + Args: + compress_ratio: The compression ratio for the layer + + Returns: + True if the layer uses sparse attention (ratio == 4) + """ + return compress_ratio == self.sparse_layer_ratio + + def _is_overlap_compressor(self, compress_ratio: int) -> bool: + """Check if a layer uses overlap compressor based on its compress ratio. + + Args: + compress_ratio: The compression ratio for the layer + + Returns: + True if the layer uses overlap compressor (ratio == 4) + """ + return compress_ratio == self.overlap_compress_layer_ratio + + def _get_window_size(self, compress_ratio: int, attn_type: DeepseekV4AttentionType) -> int: + """Get the window size for a layer based on its compress ratio and attention type. + + Args: + compress_ratio: The compression ratio for the layer + attn_type: The attention type + + Returns: + The window size for the layer + """ + state_factor = 2 if self._is_overlap_compressor(compress_ratio) else 1 + if attn_type == DeepseekV4AttentionType.SWA: + return self.window_size + elif attn_type in [ + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ]: + return state_factor * compress_ratio + elif attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + return None + + def _create_deepseek_v4_cache_manager( + self, + tokens_per_block: int, + max_batch_size: int, + max_seq_len: int, + compress_ratios: List[int], + dtype: DataType, + compressor_dtype: DataType, + max_input_len: Optional[int] = None, + ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]: + """Helper to create a DeepseekV4CacheManager for testing.""" + + # Create sparse attention config + sparse_attn_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=self.index_head_dim, + window_size=self.window_size, + compress_ratios=compress_ratios, + ) + + # Create KV cache config + if max_input_len is None: + max_input_len = max_seq_len + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + max_tokens=max_seq_len * max_batch_size, + event_buffer_max_size=0, + ) + + # Create mapping (single GPU, no parallelism) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + # Create cache manager + cache_manager = DeepseekV4CacheManager( + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=len(compress_ratios), + num_kv_heads=1, + head_dim=self.head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + max_input_len=max_input_len, + mapping=mapping, + dtype=dtype, + compressor_dtype=compressor_dtype, + vocab_size=self.vocab_size, + max_num_tokens=max_batch_size * (max_input_len + 1), + sparse_attn_config=sparse_attn_config, + ) + + return cache_manager, sparse_attn_config + + def _create_request(self, request_id: int, prompt_len: int) -> LlmRequest: + """Helper to create a test LlmRequest. + + Args: + request_id: Unique request identifier + prompt_len: Prompt length (number of tokens) + + Returns: + LlmRequest instance + """ + input_tokens = list(range(prompt_len)) + request = LlmRequest( + request_id=request_id, + max_new_tokens=1024, + input_tokens=input_tokens, + sampling_config=SamplingConfig(), + is_streaming=False, + ) + + return request + + def _rand_tensor( + self, + shape: Tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + if dtype in (torch.uint8, torch.float8_e4m3fn): + # Use uint8 for both uint8 and FP8 (same 1-byte layout) + return torch.randint(0, 255, shape, dtype=torch.uint8, device=device) + else: + return torch.randn(shape, dtype=dtype, device=device) * 1000.0 + + def _create_random_cache( + self, + seq_len: int, + head_dim: int, + sparse_attn_config: DeepSeekV4SparseAttentionConfig, + dtype: torch.dtype, + compressor_dtype: torch.dtype, + device: torch.device | None = None, + ) -> _RequestCache: + """Helper to create random cache values for all layers and attention types. + + Args: + seq_len: Sequence length + head_dim: Head dimension for regular attention + sparse_attn_config: Sparse attention configuration + + Returns: + Dictionary mapping (layer_idx, attn_type) to (values, scales) tuples. + scales is None for non-quantized attention types. + """ + device = device or torch.device("cuda") + cache: _RequestCache = {} + + for layer, ratio in enumerate(sparse_attn_config.compress_ratios): + is_overlap = self._is_overlap_compressor(ratio) + + cache[layer, DeepseekV4AttentionType.SWA] = ( + self._rand_tensor((seq_len, head_dim), dtype, device), + None, + ) + + if self._is_compress_layer(ratio): + compressor_dim = 2 * head_dim if is_overlap else head_dim + cache[layer, DeepseekV4AttentionType.COMPRESS] = ( + self._rand_tensor((seq_len // ratio, head_dim), dtype, device), + None, + ) + cache[layer, DeepseekV4AttentionType.COMPRESSOR_STATE] = ( + self._rand_tensor((seq_len, compressor_dim), compressor_dtype, device), + None, + ) + cache[layer, DeepseekV4AttentionType.COMPRESSOR_SCORE] = ( + self._rand_tensor((seq_len, compressor_dim), compressor_dtype, device), + None, + ) + + if self._is_sparse_layer(ratio): + # indexer kv cache is blockwise FP8 quantized + indexer_dim = sparse_attn_config.index_head_dim + indexer_num_tokens = seq_len // ratio + num_scales = indexer_dim // self.indexer_quant_block_size + indexer_values = self._rand_tensor( + (indexer_num_tokens, indexer_dim), torch.uint8, device + ) + indexer_scales = self._rand_tensor( + (indexer_num_tokens, num_scales), torch.float32, device + ) + cache[layer, DeepseekV4AttentionType.INDEXER_COMPRESS] = ( + indexer_values, + indexer_scales, + ) + + indexer_compressor_dim = 2 * indexer_dim if is_overlap else indexer_dim + cache[layer, DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE] = ( + self._rand_tensor((seq_len, indexer_compressor_dim), compressor_dtype, device), + None, + ) + cache[layer, DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE] = ( + self._rand_tensor((seq_len, indexer_compressor_dim), compressor_dtype, device), + None, + ) + + return cache + + def _split_blockwise_buffer(self, buffer: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Split a blockwise FP8 quantized buffer into value and scale buffers. + + Args: + buffer: The blockwise FP8 quantized buffer (shape: [num_blocks, tokens_per_block, bytes_per_token]) + + Returns: + Tuple of (value_buffer, scale_buffer) + """ + num_blocks, tokens_per_block, bytes_per_token = buffer.shape + bytes_per_block = bytes_per_token * tokens_per_block + + # Get value buffer + value_shape = (num_blocks, tokens_per_block, self.index_head_dim) + value_stride = (bytes_per_block, self.index_head_dim, 1) + value_buffer = buffer.as_strided(value_shape, value_stride, 0).view(torch.uint8) + + # Get scale buffer + scale_dim = self.index_head_dim // self.indexer_quant_block_size + scale_bytes = scale_dim * 4 # float32 = 4 bytes + scale_shape = (num_blocks, tokens_per_block, scale_bytes) + scale_stride = (bytes_per_block, scale_bytes, 1) + scale_offset = self.index_head_dim * tokens_per_block + scale_buffer = buffer.as_strided(scale_shape, scale_stride, scale_offset).view( + torch.float32 + ) + + return value_buffer, scale_buffer + + def _prefill_write_paged_cache( + self, + buffer: torch.Tensor, + block_indices: List[int], + values: torch.Tensor, + ) -> None: + """Write context values to a paged cache buffer. + + Args: + buffer: The cache buffer to write to (shape: [num_blocks, tokens_per_block, dim_per_token]) + block_indices: List of block indices to write to + values: Values to write (shape: [seq_len, dim_per_token]) + """ + assert buffer.size(2) == values.size(1), f"{buffer.size(2)=} != {values.size(1)=}" + tokens_per_block = buffer.size(1) + seq_len, dim_per_token = values.shape + + num_blocks = (seq_len + tokens_per_block - 1) // tokens_per_block + assert all(idx != BAD_PAGE_INDEX for idx in block_indices[:num_blocks]), ( + f"{block_indices[:num_blocks]=} contains BAD_PAGE_INDEX" + ) + + if seq_len % tokens_per_block != 0: + # pad the values to the nearest multiple of tokens_per_block + pad_len = tokens_per_block - (seq_len % tokens_per_block) + values = torch.cat( + [ + values, + self._rand_tensor((pad_len, dim_per_token), values.dtype, values.device), + ], + dim=0, + ) + + values_blocks = values.reshape(num_blocks, tokens_per_block, dim_per_token) + buffer[block_indices[:num_blocks]] = values_blocks + + def _decode_write_paged_cache( + self, + buffer: torch.Tensor, + block_indices: List[int], + token_idx: int, + value: torch.Tensor, + ) -> None: + """Simulate the decode phrase. Write one new token to the cache. + + Args: + buffer: The cache buffer to write to (shape: [num_blocks, tokens_per_block, dim_per_token]) + block_indices: List of block indices to write to + token_idx: Index of the new token to write + value: Value to write (shape: [dim_per_token]) + """ + assert buffer.size(2) == value.size(0), f"{buffer.size(2)=} != {value.size(0)=}" + num_blocks, tokens_per_block, _ = buffer.shape + + block_idx = token_idx // tokens_per_block + block_offset = token_idx % tokens_per_block + assert block_idx < num_blocks, f"{block_idx=} >= {num_blocks=}" + assert block_indices[block_idx] != BAD_PAGE_INDEX, ( + f"{block_indices[block_idx]=} == BAD_PAGE_INDEX" + ) + + buffer[block_indices[block_idx], block_offset] = value + + def _read_paged_cache( + self, buffer: torch.Tensor, block_indices: List[int], seq_len: int, window_size: int | None + ) -> torch.Tensor: + """Read values from a paged cache buffer. + + Args: + buffer: The cache buffer to read from (shape: [num_blocks, tokens_per_block, dim_per_token]) + block_indices: List of block/page indices to read from + seq_len: Sequence length + window_size: sliding window size to read from the cache + + Returns: + Tensor containing the read values (shape: [seq_len, dim_per_token] or [window_size, dim_per_token] + if window_size is given and seq_len > window_size) + """ + _, tokens_per_block, dim_per_token = buffer.shape + + # check if all blocks within the window are valid + end_block_idx = (seq_len + tokens_per_block - 1) // tokens_per_block + if window_size is not None: + start_block_idx = (seq_len - window_size + tokens_per_block - 1) // tokens_per_block + else: + start_block_idx = 0 + assert all(idx != BAD_PAGE_INDEX for idx in block_indices[start_block_idx:end_block_idx]), ( + f"{block_indices[start_block_idx:end_block_idx]=} contains BAD_PAGE_INDEX" + ) + + # read values from the cache + values = buffer[block_indices].reshape(-1, dim_per_token)[:seq_len] + if window_size is not None and seq_len > window_size: + values = values[-window_size:] + return values + + def _write_request_prefill( + self, + req: LlmRequest, + prompt_len: int, + cache_manager: DeepseekV4CacheManager, + cache_values: _RequestCache, + ) -> None: + """Write cache values for a request to the cache manager. + + Args: + req: The request to write cache for + prompt_len: Prompt length + cache_manager: The cache manager instance + cache_values: Request's cache to write + """ + compress_ratios = cache_manager._compress_ratios + for (layer_idx, attn_type), (values, scales) in cache_values.items(): + page_indices = cache_manager.get_batch_attn_offset( + [req.py_request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + attn_type=attn_type, + compress_ratio=compress_ratios[layer_idx], + ).squeeze(0) + + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + seq_len = prompt_len // compress_ratios[layer_idx] + else: + seq_len = prompt_len + + buffer = _view_fp8_as_uint8(cache_manager.get_buffers(layer_idx, attn_type)) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + # indexer compress is blockwise FP8 quantized + values_buffer, scales_buffer = self._split_blockwise_buffer(buffer) + self._prefill_write_paged_cache( + buffer=values_buffer, + block_indices=page_indices, + values=values[:seq_len], + ) + self._prefill_write_paged_cache( + buffer=scales_buffer, + block_indices=page_indices, + values=scales[:seq_len], + ) + else: + self._prefill_write_paged_cache( + buffer=buffer, + block_indices=page_indices, + values=values[:seq_len], + ) + + def _write_request_decode( + self, + req: LlmRequest, + token_idx: int, + cache_manager: DeepseekV4CacheManager, + cache_values: _RequestCache, + ) -> None: + """Simulate the decode phrase. Write one new token to the cache. + + Args: + req: The request to write cache for + token_idx: Index of the new token to write + cache_manager: The cache manager instance + cache_values: Request's cache to write + """ + compress_ratios = cache_manager._compress_ratios + for (layer_idx, attn_type), (values, scales) in cache_values.items(): + block_indices = cache_manager.get_batch_attn_offset( + [req.py_request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + attn_type=attn_type, + compress_ratio=compress_ratios[layer_idx], + ).squeeze(0) + + compressed_token_idx = token_idx + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + if (token_idx + 1) % compress_ratios[layer_idx] != 0: + # skip if current token will not trigger compression + continue + compressed_token_idx = token_idx // compress_ratios[layer_idx] + + buffer = _view_fp8_as_uint8(cache_manager.get_buffers(layer_idx, attn_type)) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + values_buffer, scales_buffer = self._split_blockwise_buffer(buffer) + self._decode_write_paged_cache( + buffer=values_buffer, + block_indices=block_indices, + token_idx=compressed_token_idx, + value=values[compressed_token_idx], + ) + self._decode_write_paged_cache( + buffer=scales_buffer, + block_indices=block_indices, + token_idx=compressed_token_idx, + value=scales[compressed_token_idx], + ) + else: + self._decode_write_paged_cache( + buffer=buffer, + block_indices=block_indices, + token_idx=compressed_token_idx, + value=values[compressed_token_idx], + ) + + def _read_request( + self, + req: LlmRequest, + seq_len: int, + cache_manager: DeepseekV4CacheManager, + compress_ratios: List[int], + ) -> _RequestCache: + """Read cache values for a request from the cache manager. + + Args: + req: The request to read cache for + seq_len: Sequence length + cache_manager: The cache manager instance + compress_ratios: Compression ratios for each layer + + Returns: + Request's cache + """ + cache_values: _RequestCache = {} + for layer, ratio in enumerate(compress_ratios): + attn_types = [DeepseekV4AttentionType.SWA] + if self._is_compress_layer(ratio): + attn_types.extend( + [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.COMPRESSOR_STATE, + DeepseekV4AttentionType.COMPRESSOR_SCORE, + ] + ) + if self._is_sparse_layer(ratio): + attn_types.extend( + [ + DeepseekV4AttentionType.INDEXER_COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_STATE, + DeepseekV4AttentionType.INDEXER_COMPRESSOR_SCORE, + ] + ) + + # read cache values for each attention type + for attn_type in attn_types: + page_indices = cache_manager.get_batch_attn_offset( + [req.py_request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + attn_type=attn_type, + compress_ratio=ratio, + ).squeeze(0) + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + attn_len = seq_len // ratio + else: + attn_len = seq_len + window_size = self._get_window_size(ratio, attn_type) + + buffer = _view_fp8_as_uint8(cache_manager.get_buffers(layer, attn_type)) + if attn_type == DeepseekV4AttentionType.INDEXER_COMPRESS: + values_buffer, scales_buffer = self._split_blockwise_buffer(buffer) + values = self._read_paged_cache( + buffer=values_buffer, + block_indices=page_indices, + seq_len=attn_len, + window_size=window_size, + ) + scales = self._read_paged_cache( + buffer=scales_buffer, + block_indices=page_indices, + seq_len=attn_len, + window_size=window_size, + ) + else: + values = self._read_paged_cache( + buffer=buffer, + block_indices=page_indices, + seq_len=attn_len, + window_size=window_size, + ) + scales = None + + cache_values[layer, attn_type] = (values, scales) + + return cache_values + + def _assert_cache_equal( + self, seq_len: int, compress_ratios: List[int], expect: _RequestCache, actual: _RequestCache + ) -> None: + """Assert that two cache dictionaries contain equal values. + + Args: + seq_len: Sequence length + compress_ratios: Compression ratios for each layer + expected: Expected cache values + actual: Actual cache values read from cache manager + """ + # Check that keys match + assert set(expect.keys()) == set(actual.keys()), ( + f"Cache keys don't match. Expected: {set(expect.keys())}, Actual: {set(actual.keys())}" + ) + + # Check each tensor value + for layer_idx, attn_type in expect.keys(): + if attn_type in [ + DeepseekV4AttentionType.COMPRESS, + DeepseekV4AttentionType.INDEXER_COMPRESS, + ]: + attn_len = seq_len // compress_ratios[layer_idx] + else: + attn_len = seq_len + + expect_values, expect_scales = expect[layer_idx, attn_type] + actual_values, actual_scales = actual[layer_idx, attn_type] + + # Slice to attention length + expect_values = expect_values[:attn_len] + if expect_scales is not None: + expect_scales = expect_scales[:attn_len] + + # Apply window size if applicable + window_size = self._get_window_size(compress_ratios[layer_idx], attn_type) + if window_size is not None: + expect_values = expect_values[-window_size:] + if expect_scales is not None: + expect_scales = expect_scales[-window_size:] + + # Assert values match + torch.testing.assert_close( + actual_values, + expect_values, + rtol=1e-5, + atol=1e-5, + msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (values)", + ) + + # Assert scales match (both should be None or both should be tensors) + if expect_scales is None: + assert actual_scales is None, ( + f"Expected no scales for layer {layer_idx}, attention type {attn_type.name}, " + f"but got scales with shape {actual_scales.shape}" + ) + else: + assert actual_scales is not None, ( + f"Expected scales for layer {layer_idx}, attention type {attn_type.name}, " + f"but got None" + ) + torch.testing.assert_close( + actual_scales, + expect_scales, + rtol=1e-5, + atol=1e-5, + msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (scales)", + ) + + def test_indexer_cache_layout_default(self): + """Indexer compressor cache: FP8 blockwise (128 fp8 + per-128 fp32 scale).""" + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=1, + max_seq_len=512, + compress_ratios=[4], + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + ) + try: + buffer = cache_manager.get_buffers(0, DeepseekV4AttentionType.INDEXER_COMPRESS) + assert buffer.dtype == torch.float8_e4m3fn + assert buffer.shape[-1] == 128 + 4 + assert cache_manager.quant_block_size == 128 + finally: + cache_manager.shutdown() + + @pytest.mark.parametrize("compress_ratios", [[1, 4, 128]]) + @pytest.mark.parametrize( + "dtype,compressor_dtype", [(DataType.BF16, DataType.FLOAT), (DataType.FP8, DataType.FLOAT)] + ) + @pytest.mark.parametrize("prompt_lens", [[512, 128, 160], [1024, 2048, 4096]]) + @pytest.mark.parametrize("num_generation_steps", [2, 100]) + def test_write_read_cache( + self, + compress_ratios: List[int], + prompt_lens: List[int], + num_generation_steps: int, + dtype: DataType, + compressor_dtype: DataType, + ): + max_batch_size = len(prompt_lens) + max_seq_len = max(prompt_lens) + num_generation_steps + 1 + max_input_len = max(prompt_lens) + # Create cache manager and sparse attention config + cache_manager, sparse_attn_config = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + compress_ratios=compress_ratios, + dtype=dtype, + compressor_dtype=compressor_dtype, + max_input_len=max_input_len, + ) + + # Create requests and their cache values + requests = list[LlmRequest]() + cache_values = dict[int, _RequestCache]() + for req_id, prompt_len in enumerate(prompt_lens): + req = self._create_request(req_id, prompt_len) + requests.append(req) + + # Generate random cache values for this request + cache_values[req_id] = self._create_random_cache( + seq_len=prompt_len + num_generation_steps + 1, + head_dim=self.head_dim, + sparse_attn_config=sparse_attn_config, + dtype=binding_to_torch_dtype(dtype), + compressor_dtype=binding_to_torch_dtype(compressor_dtype), + ) + + # Simulate the prefill phrase + scheduled_batch = ScheduledRequests() + scheduled_batch.context_requests_last_chunk = requests + for req in requests: + cache_manager.prepare_context(req) + cache_manager.resize_context(req, req.context_chunk_size) + + # Write context to cache + for req in requests: + self._write_request_prefill( + req=req, + prompt_len=prompt_lens[req.py_request_id], + cache_manager=cache_manager, + cache_values=cache_values[req.py_request_id], + ) + + # Update requests state and call update_resources + for req in requests: + req.context_current_position = prompt_lens[req.py_request_id] + req.add_new_token(prompt_lens[req.py_request_id], 0) + cache_manager.update_resources(scheduled_batch) + + # Read context from cache and verify + for req in requests: + actual_cache_values = self._read_request( + req=req, + seq_len=prompt_lens[req.py_request_id], + cache_manager=cache_manager, + compress_ratios=compress_ratios, + ) + self._assert_cache_equal( + seq_len=prompt_lens[req.py_request_id], + compress_ratios=compress_ratios, + expect=cache_values[req.py_request_id], + actual=actual_cache_values, + ) + + # Simulate the decode phrase + for i in range(num_generation_steps): + seq_lens = [prompt_len + i + 1 for prompt_len in prompt_lens] + scheduled_batch = ScheduledRequests() + scheduled_batch.generation_requests = requests + for req in requests: + cache_manager.try_allocate_generation(req) + + # Write new token to cache + for req in requests: + self._write_request_decode( + req=req, + token_idx=seq_lens[req.py_request_id] - 1, + cache_manager=cache_manager, + cache_values=cache_values[req.py_request_id], + ) + + # Read context from cache and verify + for req in requests: + actual_cache_values = self._read_request( + req=req, + seq_len=seq_lens[req.py_request_id], + cache_manager=cache_manager, + compress_ratios=compress_ratios, + ) + self._assert_cache_equal( + seq_len=seq_lens[req.py_request_id], + compress_ratios=compress_ratios, + expect=cache_values[req.py_request_id], + actual=actual_cache_values, + ) + + for req in requests: + req.add_new_token(seq_lens[req.py_request_id], 0) + cache_manager.update_resources(scheduled_batch) + + # Clean up + for req in requests: + cache_manager.free_resources(req) + cache_manager.shutdown() + + @pytest.mark.parametrize("compress_ratios", [[1, 4, 128]]) + @pytest.mark.parametrize( + "dtype,compressor_dtype", [(DataType.BF16, DataType.FLOAT), (DataType.FP8, DataType.FLOAT)] + ) + def test_kv_cache_pool_mapping( + self, compress_ratios: List[int], dtype: DataType, compressor_dtype: DataType + ): + # Create cache manager and sparse attention config + num_layers = len(compress_ratios) + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=4, + max_seq_len=1024, + compress_ratios=compress_ratios, + dtype=dtype, + compressor_dtype=compressor_dtype, + ) + + kv_cache_pool_mapping = cache_manager.kv_cache_pool_mapping + assert kv_cache_pool_mapping.shape == (num_layers, 2) + + assert torch.all(kv_cache_pool_mapping[:, 0] != -1), ( + "all layers should have swa attention pool" + ) + assert torch.all(kv_cache_pool_mapping[:, 1] >= 0), ( + "buffer pointer offset should be non-negative" + ) + assert torch.all(kv_cache_pool_mapping[:, 0] == kv_cache_pool_mapping[0, 0]), ( + "all layers should have the same pool_id" + ) + + @pytest.mark.parametrize("compress_ratios", [[1, 4, 128]]) + @pytest.mark.parametrize( + "dtype,compressor_dtype", [(DataType.BF16, DataType.FLOAT), (DataType.FP8, DataType.FLOAT)] + ) + @pytest.mark.parametrize("invalid", [False, True]) + @pytest.mark.parametrize("fill_with_zero", [False, True]) + def test_check_invalid_values_in_kv_cache( + self, + compress_ratios: List[int], + dtype: DataType, + compressor_dtype: DataType, + invalid: bool, + fill_with_zero: bool, + ): + """Test invalid value detection and optional zero-fill behavior in KV cache.""" + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=4, + max_seq_len=1024, + compress_ratios=compress_ratios, + dtype=dtype, + compressor_dtype=compressor_dtype, + ) + + # Fresh cache (zero-initialized) should have no invalid values + result = cache_manager.check_invalid_values_in_kv_cache() + assert not result, "Fresh cache should have no invalid values" + + if invalid: + # Inject invalid into a float buffer so NaN/Inf checks are supported. + layer_idx = next(i for i, ratio in enumerate(compress_ratios) if ratio > 1) + buffer = cache_manager.get_buffers(layer_idx, DeepseekV4AttentionType.COMPRESSOR_STATE) + buffer[0, 0, 0] = torch.nan + + result = cache_manager.check_invalid_values_in_kv_cache(fill_with_zero=fill_with_zero) + assert result == invalid, ( + f"Expected invalid={invalid} from check_invalid_values_in_kv_cache, got {result}" + ) + + # Verify whether invalid values remain after the check. + post_check = cache_manager.check_invalid_values_in_kv_cache() + expected_post_check = invalid and not fill_with_zero + assert post_check == expected_post_check, ( + f"Expected post-check invalid={expected_post_check}, got {post_check}" + ) + + if expected_post_check: + # Cleanup for shutdown path when zero-fill wasn't requested above. + cache_manager.check_invalid_values_in_kv_cache(fill_with_zero=True) + + cache_manager.shutdown() + + +if __name__ == "__main__": + tester = TestDeepseekV4CacheManager() + print("=== FP8, prompt_lens=[1024, 2048, 4096], steps=100 ===") + tester.test_write_read_cache([1, 4, 128], [1024, 2048, 4096], 100, DataType.FP8, DataType.FLOAT) + print("Test passed") diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py new file mode 100644 index 000000000000..7ef119515652 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py @@ -0,0 +1,540 @@ +""" +Tests for DeepSeek-V4 Index Transform Kernel. +""" + +from dataclasses import dataclass +from typing import List, Optional + +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention_backend.sparse.kernel import deepseek_v4_local_to_global_indices +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import ( + DeepseekV4AttentionType, + DeepseekV4CacheManager, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import get_token_bytes +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, DeepSeekV4SparseAttentionConfig +from tensorrt_llm.mapping import Mapping + + +@dataclass(kw_only=True, frozen=True) +class Scenario: + """Test scenario configuration.""" + + layer_idx: int = 0 + head_dim: int = 512 + index_head_dim: int = 128 + window_size: int = 128 + vocab_size: int = 129280 + tokens_per_block: int = 128 + max_batch_size: int = 16 + max_seq_len: int = 2048 + dtype: DataType = DataType.BF16 + compressor_dtype: DataType = DataType.FLOAT + + compress_ratio: int = 1 + swa_topk: int = 128 + compressed_topk: int = 512 + compressed_attn_type: DeepseekV4AttentionType = DeepseekV4AttentionType.COMPRESS + + +scenarios = [ + Scenario(compress_ratio=1, swa_topk=128, compressed_topk=0, layer_idx=0), + Scenario(compress_ratio=4, swa_topk=128, compressed_topk=512, layer_idx=2), + # Set compressed_topk to 2048 to test all compressed tokens + # It's not a realistic scenario, but it's helpful to test all compressed tokens. + Scenario(compress_ratio=128, swa_topk=128, compressed_topk=2048, layer_idx=3), +] + +batch_configs = [ + [256, 192], + [128, 233, 876], + [1158], +] + +index_transform_cases = [ + pytest.param( + scenario, + context_lengths, + False, + id=f"compress={scenario.compress_ratio}-batch={len(context_lengths)}", + ) + for scenario in scenarios + for context_lengths in batch_configs +] +index_transform_cases.append( + pytest.param( + None, + None, + True, + id="large-address-gap-int32-safe", + ) +) + + +def _create_cache_manager(scenario: Scenario, num_layers: int = 1): + """Create a DeepseekV4CacheManager for testing.""" + base_ratios = [1, 4, 128] + compress_ratios = [base_ratios[i % len(base_ratios)] for i in range(num_layers)] + if scenario.layer_idx < num_layers: + compress_ratios[scenario.layer_idx] = scenario.compress_ratio + sparse_attn_config = DeepSeekV4SparseAttentionConfig( + index_head_dim=scenario.index_head_dim, + window_size=scenario.window_size, + compress_ratios=compress_ratios, + ) + + max_tokens = scenario.max_seq_len * scenario.max_batch_size + cache_manager = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + max_tokens=max_tokens, + event_buffer_max_size=0, + ), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=num_layers, + num_kv_heads=1, + head_dim=scenario.head_dim, + tokens_per_block=scenario.tokens_per_block, + max_seq_len=scenario.max_seq_len, + max_batch_size=scenario.max_batch_size, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=scenario.dtype, + compressor_dtype=scenario.compressor_dtype, + vocab_size=scenario.vocab_size, + max_input_len=scenario.max_seq_len, + max_num_tokens=max_tokens, + sparse_attn_config=sparse_attn_config, + ) + return cache_manager + + +def _local_to_physical_idx(local_idx: int, block_table: List[int], tokens_per_block: int) -> int: + """ + Convert local token index to physical buffer index using block table. + This simulates: buffer_ptr + block_table[block_idx] * tokens_per_block + token_in_block + """ + block_idx = local_idx // tokens_per_block + token_in_block = local_idx % tokens_per_block + page_idx = block_table[block_idx] + return page_idx * tokens_per_block + token_in_block + + +def _build_swa_indices(token_pos: int, window_size: int, swa_topk: int, device) -> torch.Tensor: + """Build SWA local indices for a token at position token_pos.""" + indices = torch.full((swa_topk,), -1, dtype=torch.int32, device=device) + if token_pos < window_size: + indices[: token_pos + 1] = torch.arange(token_pos + 1, dtype=torch.int32, device=device) + else: + indices[:window_size] = torch.arange( + token_pos - window_size + 1, token_pos + 1, dtype=torch.int32, device=device + ) + return indices + + +def _build_compressed_indices( + token_pos: int, compress_ratio: int, compressed_topk: int, device +) -> torch.Tensor: + """Build compressed local indices for a token at position token_pos.""" + indices = torch.full((compressed_topk,), -1, dtype=torch.int32, device=device) + num_valid = (token_pos + 1) // compress_ratio + if compress_ratio == 128: + indices[:num_valid] = torch.arange(num_valid, dtype=torch.int32, device=device) + else: + select_count = min(compressed_topk, num_valid) + indices[:select_count] = torch.randperm(num_valid, device=device)[:select_count].to( + torch.int32 + ) + return indices + + +def _run_test(scenario: Scenario, context_lengths: List[int]): + """Run index transformation test.""" + device = torch.device("cuda") + layer_idx = scenario.layer_idx + has_compressed = scenario.compress_ratio > 1 + total_tokens = sum(context_lengths) + + torch.manual_seed(42) + + # Create cache manager and requests + cache_manager = _create_cache_manager(scenario, num_layers=7) + requests = [ + LlmRequest( + request_id=i, + max_new_tokens=1024, + input_tokens=list(range(ctx_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + for i, ctx_len in enumerate(context_lengths) + ] + + scheduled_batch = ScheduledRequests() + for req in requests: + scheduled_batch.append_context_request(req) + for req in requests: + cache_manager.prepare_context(req) + cache_manager.resize_context(req, req.context_chunk_size) + + # Get pointers and offsets + swa_pool_base_ptr = cache_manager.swa_pool_ptr + swa_buffer_ptr = cache_manager.get_buffers(layer_idx, DeepseekV4AttentionType.SWA).data_ptr() + + # Single token stride for all buffers + has_fp8_kv_cache = scenario.dtype == DataType.FP8 + token_stride = get_token_bytes( + scenario.head_dim, + scenario.index_head_dim, + scenario.compress_ratio, + DeepseekV4AttentionType.SWA, + has_fp8_kv_cache, + ) + swa_offset = (swa_buffer_ptr - swa_pool_base_ptr) // token_stride + + # Get compressed buffer type from scenario property + compressed_attn_type = scenario.compressed_attn_type + + if has_compressed: + compress_pool_base_ptr = cache_manager.compress_pool_ptrs[scenario.compress_ratio] + compressed_buffer_ptr = cache_manager.get_buffers( + layer_idx, compressed_attn_type + ).data_ptr() + compressed_offset = (compressed_buffer_ptr - compress_pool_base_ptr) // token_stride + tokens_per_block_compressed = scenario.tokens_per_block // scenario.compress_ratio + else: + ( + compress_pool_base_ptr, + compressed_buffer_ptr, + compressed_offset, + tokens_per_block_compressed, + ) = ( + 0, + 0, + 0, + scenario.tokens_per_block, + ) + + # Get buffers and write random values + swa_buffer = cache_manager.get_buffers(layer_idx, DeepseekV4AttentionType.SWA) + swa_buffer.copy_(torch.randn_like(swa_buffer)) + + if has_compressed: + compressed_buffer = cache_manager.get_buffers(layer_idx, compressed_attn_type) + # Generate random data in float32 first, then convert to target dtype (for FP8 support) + random_data = torch.randn(compressed_buffer.shape, dtype=torch.float32, device=device) + compressed_buffer.copy_(random_data.to(compressed_buffer.dtype)) + + # Get block tables + block_tables_swa = [ + cache_manager.get_cache_indices(req.py_request_id, layer_idx, DeepseekV4AttentionType.SWA) + for req in requests + ] + block_tables_compressed = ( + [ + cache_manager.get_cache_indices(req.py_request_id, layer_idx, compressed_attn_type) + for req in requests + ] + if has_compressed + else [] + ) + + # Flatten buffers for access + swa_buffer_flat = swa_buffer.view(-1, scenario.head_dim) + if has_compressed: + compressed_buffer_flat = compressed_buffer.view(-1, scenario.head_dim) + + # Pad and convert block tables to tensors + max_blocks_swa = max(len(bt) for bt in block_tables_swa) + block_table_swa_t = torch.tensor( + [bt + [-1] * (max_blocks_swa - len(bt)) for bt in block_tables_swa], + dtype=torch.int32, + device=device, + ) + + if has_compressed: + max_blocks_compressed = max(len(bt) for bt in block_tables_compressed) + block_table_compressed_t = torch.tensor( + [bt + [-1] * (max_blocks_compressed - len(bt)) for bt in block_tables_compressed], + dtype=torch.int32, + device=device, + ) + else: + block_table_compressed_t = None + + # Build inputs for all tokens + req_ids, token_positions = [], [] + for r, ctx_len in enumerate(context_lengths): + for pos in range(ctx_len): + req_ids.append(r) + token_positions.append(pos) + + req_id = torch.tensor(req_ids, dtype=torch.int32, device=device) + + # Build local indices + swa_local_indices = torch.stack( + [ + _build_swa_indices(pos, scenario.window_size, scenario.swa_topk, device) + for pos in token_positions + ] + ) + + if has_compressed: + compressed_local_indices = torch.stack( + [ + _build_compressed_indices( + pos, scenario.compress_ratio, scenario.compressed_topk, device + ) + for pos in token_positions + ] + ) + else: + compressed_local_indices = None + + # Run kernel + global_indices = deepseek_v4_local_to_global_indices( + req_id=req_id, + block_table_swa=block_table_swa_t, + swa_local_indices=swa_local_indices, + swa_pool_base_ptr=swa_pool_base_ptr, + swa_buffer_ptr=swa_buffer_ptr, + tokens_per_block=scenario.tokens_per_block, + token_stride=token_stride, + block_table_compressed=block_table_compressed_t, + compressed_local_indices=compressed_local_indices, + compress_pool_base_ptr=compress_pool_base_ptr, + compressed_buffer_ptr=compressed_buffer_ptr, + compress_ratio=scenario.compress_ratio, + num_compressed_indices=scenario.compressed_topk if has_compressed else 0, + ) + + # Verify dual-pool non-compact layout: + # SWA region [0, window_size) — indices relative to swa_pool_base_ptr + # Compress region [window_size, window_size + compressed_topk) — indices relative to compress_pool_base_ptr + # Invalid positions padded with -1 at their fixed positions. + window_size = scenario.window_size + num_samples = min(32, total_tokens) + sample_indices = torch.randperm(total_tokens)[:num_samples].tolist() + + for t in sample_indices: + r = req_ids[t] + out_row = global_indices[t] + + # Verify SWA region [0, window_size) — at fixed positions, matching input positions + for pos in range(scenario.swa_topk): + local_idx = swa_local_indices[t, pos].item() + global_idx = out_row[pos].item() + + if local_idx < 0: + # Invalid input → -1 in output + assert global_idx == -1, ( + f"Token {t} SWA out[{pos}]: expected -1 for invalid input, got {global_idx}" + ) + else: + # Valid: verify data access + assert global_idx >= 0, ( + f"Token {t} SWA out[{pos}]: expected valid index, got {global_idx}" + ) + pool_based_idx = global_idx - swa_offset + actual = swa_buffer_flat[pool_based_idx] + + physical_idx = _local_to_physical_idx( + local_idx, block_tables_swa[r], scenario.tokens_per_block + ) + expected = swa_buffer_flat[physical_idx] + + torch.testing.assert_close( + actual, expected, msg=f"Token {t} SWA out[{pos}]: value mismatch" + ) + + # Verify compress region [window_size, window_size + compressed_topk) + if has_compressed: + for pos in range(scenario.compressed_topk): + out_pos = window_size + pos + local_idx = compressed_local_indices[t, pos].item() + global_idx = out_row[out_pos].item() + + if local_idx < 0: + assert global_idx == -1, ( + f"Token {t} compress out[{out_pos}]: expected -1 for invalid, got {global_idx}" + ) + else: + assert global_idx >= 0, ( + f"Token {t} compress out[{out_pos}]: expected valid index, got {global_idx}" + ) + pool_based_idx = global_idx - compressed_offset + actual = compressed_buffer_flat[pool_based_idx] + + physical_idx = _local_to_physical_idx( + local_idx, block_tables_compressed[r], tokens_per_block_compressed + ) + expected = compressed_buffer_flat[physical_idx] + + torch.testing.assert_close( + actual, expected, msg=f"Token {t} compress out[{out_pos}]: value mismatch" + ) + + # Cleanup + for req, ctx_len in zip(requests, context_lengths): + req.context_current_position = ctx_len + req.add_new_token(ctx_len, 0) + cache_manager.update_resources(scheduled_batch) + cache_manager.shutdown() + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.parametrize( + "scenario, context_lengths, is_large_gap_case", + index_transform_cases, +) +def test_deepseek_v4_indices_transform( + scenario: Optional[Scenario], + context_lengths: Optional[List[int]], + is_large_gap_case: bool, +): + """Test DeepSeek-V4 index transformation kernel by verifying VALUES.""" + if is_large_gap_case: + _run_large_address_gap_int32_safe_test() + return + + assert scenario is not None and context_lengths is not None + _run_test(scenario, context_lengths) + + +def _run_large_address_gap_int32_safe_test(): + """ + Synthetic test: emulate SWA/compress pool pointer deltas that would overflow int32 + in the old single-base scheme, and verify the new per-pool index outputs remain int32-safe. + """ + device = torch.device("cuda") + window_size = 128 + tokens_per_block = 128 + compressed_topk = 512 + compress_ratio = 4 + tokens_per_block_compressed = tokens_per_block // compress_ratio + num_tokens = 8 + ctx_len = 600 + + # Simulate two pools 100+ TB apart — would overflow int32 with single base ptr + # Using synthetic offsets that stay int32-safe per-pool + fake_swa_pool_base_ptr = 0x7F0000000000 # ~127 TB + fake_swa_buffer_ptr = fake_swa_pool_base_ptr + 1024 * 1024 # 1 MB offset + fake_compress_pool_base_ptr = 0x100000000 # ~4 GB (100+ TB away from SWA pool) + fake_compressed_buffer_ptr = fake_compress_pool_base_ptr + 512 * 1024 + + token_stride = 1024 # bytes per token (synthetic) + swa_offset = (fake_swa_buffer_ptr - fake_swa_pool_base_ptr) // token_stride + compressed_offset = (fake_compressed_buffer_ptr - fake_compress_pool_base_ptr) // token_stride + + # Both offsets should be small (well within int32) + assert abs(swa_offset) < 2**31, f"SWA offset overflows int32: {swa_offset}" + assert abs(compressed_offset) < 2**31, f"Compress offset overflows int32: {compressed_offset}" + + # Build simple block tables (identity: page i = i) + max_blocks_swa = (ctx_len + tokens_per_block - 1) // tokens_per_block + max_blocks_compressed = ( + ctx_len // compress_ratio + tokens_per_block_compressed - 1 + ) // tokens_per_block_compressed + block_table_swa = torch.arange(max_blocks_swa, dtype=torch.int32, device=device).unsqueeze(0) + block_table_compressed = torch.arange( + max_blocks_compressed, dtype=torch.int32, device=device + ).unsqueeze(0) + + req_id = torch.zeros(num_tokens, dtype=torch.int32, device=device) + + # Build SWA indices: positions [ctx_len - num_tokens, ..., ctx_len - 1] + positions = list(range(ctx_len - num_tokens, ctx_len)) + swa_indices_list = [] + for pos in positions: + indices = torch.full((window_size,), -1, dtype=torch.int32, device=device) + start = max(0, pos - window_size + 1) + valid_count = min(window_size, pos + 1) + indices[:valid_count] = torch.arange( + start, start + valid_count, dtype=torch.int32, device=device + ) + swa_indices_list.append(indices) + swa_local_indices = torch.stack(swa_indices_list) + + # Build compressed indices + compressed_indices_list = [] + for pos in positions: + indices = torch.full((compressed_topk,), -1, dtype=torch.int32, device=device) + num_valid = (pos + 1) // compress_ratio + select_count = min(compressed_topk, num_valid) + if select_count > 0: + indices[:select_count] = torch.arange(select_count, dtype=torch.int32, device=device) + compressed_indices_list.append(indices) + compressed_local_indices = torch.stack(compressed_indices_list) + + # Run kernel with separate base pointers + global_indices = deepseek_v4_local_to_global_indices( + req_id=req_id, + block_table_swa=block_table_swa, + swa_local_indices=swa_local_indices, + swa_pool_base_ptr=fake_swa_pool_base_ptr, + swa_buffer_ptr=fake_swa_buffer_ptr, + tokens_per_block=tokens_per_block, + token_stride=token_stride, + block_table_compressed=block_table_compressed, + compressed_local_indices=compressed_local_indices, + compress_pool_base_ptr=fake_compress_pool_base_ptr, + compressed_buffer_ptr=fake_compressed_buffer_ptr, + compress_ratio=compress_ratio, + num_compressed_indices=compressed_topk, + ) + + # Verify: all valid SWA indices are int32-safe and relative to swa_pool + assert global_indices.dtype == torch.int32, f"Expected int32 output, got {global_indices.dtype}" + swa_region = global_indices[:, :window_size] + compress_region = global_indices[:, window_size:] + + # Check SWA region: valid entries should be >= 0 and represent swa_offset + page*tpb + token_in_block + for t in range(num_tokens): + for pos in range(window_size): + local_idx = swa_local_indices[t, pos].item() + global_idx = swa_region[t, pos].item() + if local_idx < 0: + assert global_idx == -1, f"Token {t} SWA[{pos}]: invalid input should give -1" + else: + assert global_idx >= 0, ( + f"Token {t} SWA[{pos}]: valid input gave negative index {global_idx}" + ) + # Verify the index is swa_offset + page_idx * tpb + token_in_block + block_idx = local_idx // tokens_per_block + token_in_block = local_idx % tokens_per_block + page_idx = block_table_swa[0, block_idx].item() + expected = swa_offset + page_idx * tokens_per_block + token_in_block + assert global_idx == expected, ( + f"Token {t} SWA[{pos}]: expected {expected}, got {global_idx}" + ) + + # Check compress region + for t in range(num_tokens): + for pos in range(compressed_topk): + local_idx = compressed_local_indices[t, pos].item() + global_idx = compress_region[t, pos].item() + if local_idx < 0: + assert global_idx == -1 + else: + assert global_idx >= 0 + block_idx = local_idx // tokens_per_block_compressed + token_in_block = local_idx % tokens_per_block_compressed + page_idx = block_table_compressed[0, block_idx].item() + expected = ( + compressed_offset + page_idx * tokens_per_block_compressed + token_in_block + ) + assert global_idx == expected, ( + f"Token {t} compress[{pos}]: expected {expected}, got {global_idx}" + ) + + +if __name__ == "__main__": + _run_test(scenarios[1], [256, 192]) + print("PASSED") diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py new file mode 100644 index 000000000000..ffe72b8c5fcf --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py @@ -0,0 +1,293 @@ +""" +Test for DeepSeek-V4 output projection (_deepseek_v4_o_proj). +""" + +from types import SimpleNamespace + +import pytest +import torch +from _torch.helpers import per_block_cast_to_fp8_e8m0 +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv3 import weight_dequant +from tensorrt_llm._torch.modules.attention import MLA +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from ..test_sparse_mla_forward import RopeConfig, _calc_diff, apply_rotary_emb, precompute_freqs_cis + + +def calculate_reference_deepseek_v4_o_proj( + attn_out_latent, + o_a_proj, + o_b_proj_weight, + freqs_cis, + n_local_groups, + qk_nope_head_dim, + qk_rope_head_dim, + device, + is_fp8: bool = False, +): + """ + Reference implementation for DeepSeek-V4 output projection based on ref/model.py. + + Args: + attn_out_latent: [num_tokens, num_heads, qk_head_dim] attention output + o_a_proj: [n_local_groups, o_lora_rank, num_heads * qk_head_dim // n_groups] + o_b_proj_weight: [hidden_size, n_groups * o_lora_rank] + freqs_cis: [num_toknes, rope_head_dim / 2] rotary embeddings + n_local_groups: Number of local output projection groups + qk_nope_head_dim: Dimension of non-positional part + qk_rope_head_dim: Dimension of positional part + device: Device to run on + is_fp8: Whether test fp8 or bf16 + + Returns: + output: [num_tokens, hidden_size] projected output + """ + num_tokens = attn_out_latent.shape[0] + qk_nope_head_dim + qk_rope_head_dim + + # Apply RoPE to attn_out_pe + attn_out_latent = attn_out_latent.unsqueeze(0) + apply_rotary_emb(attn_out_latent[..., -qk_rope_head_dim:], freqs_cis, inverse=True) + + # Reshape for grouped projection + attn_out_grouped = attn_out_latent.view(num_tokens, n_local_groups, -1) + + # Apply o_a_proj: einsum equivalent to bmm + o_lora = torch.einsum("tgd,grd->tgr", attn_out_grouped, o_a_proj) + + # Flatten and apply o_b_proj, [num_tokens, n_local_groups * o_lora_rank] + o_lora_flat = o_lora.flatten(1) + if is_fp8: + o_lora_flat = o_lora_flat.to(torch.float8_e4m3fn).to(torch.bfloat16) + output = torch.nn.functional.linear(o_lora_flat, o_b_proj_weight) # [num_tokens, hidden_size] + + return output + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +@pytest.mark.parametrize("dtype_str", ["bf16", "fp8"]) +def test_deepseek_v4_o_proj(num_tokens: int, dtype_str: str): + """Test DeepSeek-V4 output projection (_deepseek_v4_o_proj).""" + print( + f"\n{'=' * 80}\nTesting: deepseek_v4_o_proj num_tokens={num_tokens} dtype={dtype_str}\n{'=' * 80}" + ) + + if dtype_str == "fp8" and get_sm_version() < 100: + pytest.skip("FP8 is not supported on pre-Blackwell architectures") + + device = torch.device("cuda") + dtype = torch.bfloat16 + + # Model configuration matching the reference model + num_heads = 64 + q_lora_rank = 1024 + kv_lora_rank = 448 + qk_nope_head_dim = 448 + qk_rope_head_dim = 64 + v_head_dim = 512 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + hidden_size = 4096 + max_position_embeddings = 65536 + o_lora_rank = 1024 + num_groups = 8 + n_local_groups = num_groups # no TP in this test + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # Create RoPE config + rope_config = RopeConfig( + hidden_size=hidden_size, + num_attention_heads=num_heads, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 4, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 65536, + "type": "yarn", + }, + max_position_embeddings=max_position_embeddings, + rope_theta=10000.0, + qk_rope_head_dim=qk_rope_head_dim, + model_type="deepseek_v4", + ) + + # Setup model config with deepseek_v4 sparse attention + mapping = Mapping(world_size=1, tp_size=1, rank=0) + pretrained_config = SimpleNamespace( + rms_norm_eps=1e-6, + ) + + # Create sparse attention config for deepseek_v4 + sparse_config = DeepSeekV4SparseAttentionConfig( + index_n_heads=32, + index_head_dim=128, + index_topk=512, + ) + + quant_config = QuantConfig() + if dtype_str == "fp8": + quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES + quant_config.group_size = 128 + + model_config = ModelConfig( + mapping=mapping, + pretrained_config=pretrained_config, + sparse_attention_config=sparse_config, + quant_config=quant_config, + ) + + # Setup positional embedding params + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, + rope=RopeParams.from_config(rope_config), + is_neox=False, + ) + + # Create MLA module with deepseek_v4 configuration + mla = MLA( + hidden_size=hidden_size, + num_attention_heads=num_heads, + num_key_value_heads=1, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + predicted_tokens_per_seq=1, + max_position_embeddings=max_position_embeddings, + bias=False, + pos_embd_params=pos_embd_params, + layer_idx=0, + dtype=dtype, + config=model_config, + num_groups=num_groups, + o_lora_rank=o_lora_rank, + ).to(device) + + # Initialize weights + nn_init_std = 0.02 + with torch.no_grad(): + # Initialize o_a_proj weights + if dtype_str == "bf16": + mla.o_a_proj.data = ( + torch.randn( + n_local_groups, + o_lora_rank, + num_heads * qk_head_dim // num_groups, + dtype=dtype, + device=device, + ) + * nn_init_std + ) + elif dtype_str == "fp8": + dim = num_heads * qk_head_dim // num_groups + o_a_proj_bf16 = ( + torch.randn(n_local_groups, o_lora_rank, dim, dtype=torch.bfloat16, device=device) + * nn_init_std + ) + + fp8_a_weight, fp8_a_scale = per_block_cast_to_fp8_e8m0(o_a_proj_bf16.reshape(-1, dim)) + fp8_a_weight = fp8_a_weight.reshape(n_local_groups, o_lora_rank, dim) + mla.o_a_proj.data = fp8_a_weight + mla.o_a_proj_scale.data = fp8_a_scale + mla.o_a_proj_dequant.data = o_a_proj_bf16 + + # Initialize o_b_proj weights + if dtype_str == "bf16": + mla.o_b_proj.weight.data = ( + torch.randn(hidden_size, num_groups * o_lora_rank, dtype=dtype, device=device) + * nn_init_std + ) + elif dtype_str == "fp8": + # For FP8, properly quantize using fp8_quantize_1x128_sf_transpose + o_b_proj_weight_bf16 = ( + torch.randn( + hidden_size, num_groups * o_lora_rank, dtype=torch.bfloat16, device=device + ) + * nn_init_std + ) + + # Quantize the weight + fp8_b_weight, fp8_b_scale = per_block_cast_to_fp8_e8m0(o_b_proj_weight_bf16) + fp8_b_weight_dequant = weight_dequant(fp8_b_weight, fp8_b_scale).bfloat16() + mla.o_b_proj.weight.data = fp8_b_weight + mla.o_b_proj.weight_scale.data = fp8_b_scale + + # Generate test inputs + # Note: for deepseek_v4, kv_lora_rank equals qk_head_dim + attn_out_latent = torch.randn(num_tokens, num_heads, qk_head_dim, dtype=dtype, device=device) + position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) + + # Call the deepseek_v4 output projection (mla_rope_inplace modifies attn_out_latent + # in-place, so clone before passing to preserve original for reference) + output = mla._deepseek_v4_o_proj(attn_out_latent.clone(), position_ids) + + # Calculate reference output + if dtype_str == "bf16": + o_a_proj_ref = mla.o_a_proj.data + o_b_proj_weight_ref = mla.o_b_proj.weight.data + else: + # For FP8, convert back to bf16 for reference calculation + o_a_proj_ref = o_a_proj_bf16 + o_b_proj_weight_ref = fp8_b_weight_dequant + + freqs_cis = precompute_freqs_cis( + qk_rope_head_dim, + num_tokens, + max_position_embeddings, + rope_config.rope_theta, + rope_config.rope_scaling["factor"], + rope_config.rope_scaling["beta_fast"], + rope_config.rope_scaling["beta_slow"], + ).to(device) + + reference_output = calculate_reference_deepseek_v4_o_proj( + attn_out_latent=attn_out_latent, + o_a_proj=o_a_proj_ref, + o_b_proj_weight=o_b_proj_weight_ref, + freqs_cis=freqs_cis[0:num_tokens], + n_local_groups=n_local_groups, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + device=device, + is_fp8=dtype_str == "fp8", + ) + + # Validate output shapes + assert output.shape == reference_output.shape, ( + f"Shape mismatch: output {output.shape} vs reference {reference_output.shape}" + ) + assert output.dtype == reference_output.dtype, ( + f"Dtype mismatch: output {output.dtype} vs reference {reference_output.dtype}" + ) + assert torch.isfinite(output).all(), "Output contains non-finite values" + assert torch.isfinite(reference_output).all(), "Reference output contains non-finite values" + + # Compare results + abs_error = (output - reference_output).abs() + max_error = abs_error.max().item() + mean_error = abs_error.mean().item() + + print(f" Max error: {max_error:.6f}") + print(f" Mean error: {mean_error:.6f}") + + if dtype_str == "fp8": + diff = _calc_diff(output, reference_output) + assert diff < 1e-3, f"{diff=}" + else: + torch.testing.assert_close(output, reference_output, rtol=0.1, atol=0.1) + print(f" ✓ Test passed for num_tokens={num_tokens}, dtype={dtype_str}\n") diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py new file mode 100644 index 000000000000..fd23e7032aac --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -0,0 +1,1464 @@ +""" +Tests for DeepSeek-V4 sparse MLA attention. +""" + +import math +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + +import pytest +import torch +from utils.util import skip_pre_blackwell + +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + RopeParams, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import ( + DeepseekV4AttentionType, + DeepseekV4CacheManager, + DeepseekV4TrtllmAttention, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, DeepSeekV4SparseAttentionConfig +from tensorrt_llm.mapping import Mapping + + +@dataclass(kw_only=True, frozen=True) +class Scenario: + dtype: torch.dtype = torch.bfloat16 + kv_cache_dtype: torch.dtype = torch.bfloat16 + num_layers: int = 3 + num_heads: int = 128 + num_kv_heads: int = 1 + q_lora_rank: int = 1536 + kv_lora_rank: int = 512 # raw config value; effective = 448 when rope_append=False + qk_nope_head_dim: int = 128 + qk_rope_head_dim: int = 64 + v_head_dim: int = 512 + rope_append: bool = False # DeepSeek-V4 requires rope_append=False + 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 = 128 + window_size: int = 128 + compress_ratios: List[int] = field(default_factory=lambda: [1, 4, 128]) + index_topk: int = 512 + + +@dataclass(kw_only=True, frozen=True) +class RopeConfig: + hidden_size: int = 7168 + num_attention_heads: int = 128 + rope_scaling: dict = field( + default_factory=lambda: { + "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" + + +# Layers to test: layer 0 → compress_ratio=1, layer 1 → ratio=4, layer 2 → ratio=128 +TEST_LAYERS = [0, 1, 2] + + +# RoPE helpers +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 _rotate_k_pe_for_ctx( + k_pe: torch.Tensor, rope_cos_sin: torch.Tensor, sequence_lengths: List[int] +) -> torch.Tensor: + k_pe_ref_list = [] + total_tokens = 0 + for seq_len in sequence_lengths: + k_pe_seq = k_pe[total_tokens : total_tokens + seq_len].unsqueeze(-2) + cos, sin = rope_cos_sin[:seq_len].chunk(2, dim=-2) + k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + k_pe_seq = ((k_pe_seq * cos) + (rotate_half(k_pe_seq) * sin)).to(dtype=k_pe_seq.dtype) + k_pe_seq = k_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + k_pe_ref_list.append(k_pe_seq) + total_tokens += seq_len + return torch.cat(k_pe_ref_list).squeeze(-2) + + +def _rotate_fused_q_for_ctx( + fused_q: torch.Tensor, + rope_cos_sin: torch.Tensor, + sequence_lengths: List[int], + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> torch.Tensor: + fused_q = fused_q.clone() + fused_head_dim = kv_lora_rank + qk_rope_head_dim + total_tokens = 0 + for seq_len in sequence_lengths: + fused_q_seq = fused_q[total_tokens : total_tokens + seq_len].view( + seq_len, num_heads, fused_head_dim + ) + q_rope = fused_q_seq[..., -qk_rope_head_dim:] + cos, sin = rope_cos_sin[:seq_len].chunk(2, dim=-2) + q_rope = q_rope.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + q_rope = ((q_rope * cos) + (rotate_half(q_rope) * sin)).to(dtype=fused_q.dtype) + q_rope = q_rope.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + fused_q_seq[..., -qk_rope_head_dim:] = q_rope + fused_q[total_tokens : total_tokens + seq_len] = fused_q_seq.view(seq_len, -1) + total_tokens += seq_len + return fused_q + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + 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) + + +def _create_cache_manager(scenario: Scenario, context_lengths: List[int], max_seq_len: int): + sparse_config = DeepSeekV4SparseAttentionConfig( + index_n_heads=64, + index_head_dim=128, + window_size=scenario.window_size, + compress_ratios=scenario.compress_ratios, + index_topk=scenario.index_topk, + skip_indexer_for_short_seqs=False, + ) + batch_size = len(context_lengths) + max_input_len = max(context_lengths) + + cache_manager = DeepseekV4CacheManager( + kv_cache_config=KvCacheConfig( + max_tokens=max_seq_len * batch_size, + enable_block_reuse=False, + event_buffer_max_size=0, + ), + kv_cache_type=CacheTypeCpp.SELFKONLY, + num_layers=scenario.num_layers, + num_kv_heads=1, + # When rope_append=False: effective kv_lora_rank is 448, head_dim = 448 + 64 = 512 + head_dim=( + scenario.kv_lora_rank - scenario.qk_rope_head_dim + if not scenario.rope_append + else scenario.kv_lora_rank + ) + + scenario.qk_rope_head_dim, + tokens_per_block=scenario.kv_cache_tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_input_len=max_input_len, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + vocab_size=129280, + max_num_tokens=max_input_len * batch_size + batch_size, + sparse_attn_config=sparse_config, + ) + return cache_manager, sparse_config + + +def _prefill_compress_buffer( + cache_manager: DeepseekV4CacheManager, + layer_idx: int, + context_lengths: List[int], + request_ids: List[int], + head_dim: int, + device: torch.device, +) -> List[torch.Tensor]: + """Pre-fill COMPRESS buffer with known random data. + + Returns flat reference data per request. + """ + compress_ratio = cache_manager._compress_ratios[layer_idx] + buffer = cache_manager.get_buffers(layer_idx, DeepseekV4AttentionType.COMPRESS) + tokens_per_block_compressed = cache_manager.compressed_block_sizes[layer_idx] + + ref_data = [] + for req_idx, ctx_len in enumerate(context_lengths): + num_compressed = ctx_len // compress_ratio + data = torch.randn(num_compressed, head_dim, device=device, dtype=torch.bfloat16) + ref_data.append(data) + + block_ids = cache_manager.get_cache_indices( + request_ids[req_idx], layer_idx, DeepseekV4AttentionType.COMPRESS + ) + for tok_idx in range(num_compressed): + block_idx = tok_idx // tokens_per_block_compressed + offset = tok_idx % tokens_per_block_compressed + buffer[block_ids[block_idx], offset, :head_dim] = data[tok_idx] + + return ref_data + + +def _grow_compress_buffer_for_generation( + cache_manager: DeepseekV4CacheManager, + layer_idx: int, + request_ids: List[int], + head_dim: int, + device: torch.device, + kv_lens: List[int], + compress_ref_data_per_layer: List[torch.Tensor], +): + """Grow COMPRESS buffer and ref data for generation steps. + + As KV length increases during generation, new compressed tokens appear. + This appends random data for those new entries to both the COMPRESS buffer + and the reference data list (in-place). + """ + compress_ratio = cache_manager._compress_ratios[layer_idx] + buffer = cache_manager.get_buffers(layer_idx, DeepseekV4AttentionType.COMPRESS) + tokens_per_block_compressed = cache_manager.tokens_per_block // compress_ratio + + for req_idx in range(len(request_ids)): + old_count = compress_ref_data_per_layer[req_idx].shape[0] + new_total = kv_lens[req_idx] // compress_ratio + num_new = new_total - old_count + if num_new <= 0: + continue + new_data = torch.randn(num_new, head_dim, device=device, dtype=torch.bfloat16) + block_ids = cache_manager.get_cache_indices( + request_ids[req_idx], layer_idx, DeepseekV4AttentionType.COMPRESS + ) + for j in range(num_new): + tok_idx = old_count + j + block_idx = tok_idx // tokens_per_block_compressed + offset = tok_idx % tokens_per_block_compressed + buffer[block_ids[block_idx], offset, :head_dim] = new_data[j] + compress_ref_data_per_layer[req_idx] = torch.cat( + [compress_ref_data_per_layer[req_idx], new_data], dim=0 + ) + + +def _build_compressed_topk_indices( + token_positions: List[int], + compress_ratio: int, + topk: int, + device: torch.device, +) -> torch.Tensor: + """Build random topk indices into compressed space for testing.""" + num_tokens = len(token_positions) + indices = torch.full((num_tokens, topk), -1, dtype=torch.int32, device=device) + for i, pos in enumerate(token_positions): + num_compressed = (pos + 1) // compress_ratio + if num_compressed > 0: + valid_k = min(topk, num_compressed) + selected = torch.randperm(num_compressed, device=device)[:valid_k].sort().values + indices[i, :valid_k] = selected.to(torch.int32) + return indices + + +def _softmax_with_sink( + logits: torch.Tensor, + attn_sink: Optional[torch.Tensor], + out_dtype: torch.dtype, +) -> torch.Tensor: + """Numerically stable softmax with an optional per-head sink in the denominator. + + The sink behaves like a virtual register key whose value is zero: it + contributes exp(sink - max) to the denominator only, leaving the numerator + untouched. + """ + fp32_logits = logits.float() + if attn_sink is None: + return torch.softmax(fp32_logits, dim=-1).to(out_dtype) + max_val = fp32_logits.amax(dim=-1, keepdim=True) + num = torch.exp(fp32_logits - max_val) + denom = num.sum(dim=-1, keepdim=True) + sink = attn_sink.float().view(-1, 1, 1) + denom = denom + torch.exp(sink - max_val) + return (num / denom).to(out_dtype) + + +def calculate_deepseek_v4_ref_ctx_sparse( + fused_q_rot: torch.Tensor, + latent_cache_ref: torch.Tensor, + compressed_ref_data: Optional[List[torch.Tensor]], + swa_window_size: int, + compressed_topk_indices: Optional[torch.Tensor], + seq_lens: List[int], + num_heads: int, + kv_lora_rank: int, + v_head_dim: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + q_scaling: float, + compress_ratio: int, + attn_sink: Optional[torch.Tensor] = None, +): + """Per-token reference attention for DeepSeek-V4 context phase. + + For compress_ratio==1: only SWA tokens (causal window). + For compress_ratio==4: SWA tokens + indexer topk compressed tokens. + For compress_ratio==128: SWA tokens + all compressed tokens. + """ + fused_head_dim = kv_lora_rank + qk_rope_head_dim + bmm1_scale = 1 / (math.sqrt(qk_nope_head_dim + qk_rope_head_dim) * q_scaling) + + ref_results = [] + token_offset = 0 + for batch_idx, seq_len in enumerate(seq_lens): + per_token_outputs = [] + for token_idx in range(seq_len): + global_token_idx = token_offset + token_idx + pos = token_idx + + # Gather SWA KV + swa_start = max(0, pos - swa_window_size + 1) + swa_end = pos + 1 + swa_kv = latent_cache_ref[token_offset + swa_start : token_offset + swa_end] + + # Gather compressed KV + if compress_ratio > 1 and compressed_ref_data is not None: + if compress_ratio == 4 and compressed_topk_indices is not None: + indices_row = compressed_topk_indices[global_token_idx] + valid = indices_row[indices_row >= 0] + comp_kv = compressed_ref_data[batch_idx][valid.long()] + elif compress_ratio == 128: + num_comp = (pos + 1) // compress_ratio + comp_kv = compressed_ref_data[batch_idx][:num_comp] + else: + comp_kv = torch.empty( + 0, + swa_kv.shape[-1], + device=swa_kv.device, + dtype=swa_kv.dtype, + ) + + if comp_kv.numel() > 0: + all_kv = torch.cat([swa_kv, comp_kv], dim=0) + else: + all_kv = swa_kv + else: + all_kv = swa_kv + + # Compute attention + q_tok = fused_q_rot[global_token_idx].view(num_heads, fused_head_dim) + k_sel = all_kv.unsqueeze(0).expand(num_heads, -1, -1) + v_sel = all_kv[:, :v_head_dim].unsqueeze(0).expand(num_heads, -1, -1) + + attn_w = torch.matmul(q_tok.unsqueeze(1), k_sel.transpose(1, 2)) * bmm1_scale + attn_w = _softmax_with_sink(attn_w, attn_sink, fused_q_rot.dtype) + out = torch.matmul(attn_w, v_sel).squeeze(1) + per_token_outputs.append(out.reshape(1, num_heads * v_head_dim)) + + ref_results.append(torch.cat(per_token_outputs, dim=0)) + token_offset += seq_len + + return torch.cat(ref_results, dim=0) + + +def _rotate_gen_inputs( + fused_q: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + k_pe: torch.Tensor, + rope_cos_sin: torch.Tensor, + seq_lens_kv: List[int], + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + fused_head_dim = kv_lora_rank + qk_rope_head_dim + num_requests = len(seq_lens_kv) + seq_len_q = fused_q.shape[0] // num_requests + + fused_q_rot = fused_q.clone() + new_latent_list = [] + + for i in range(num_requests): + past_len = seq_lens_kv[i] + + fused_q_seq = fused_q_rot[i * seq_len_q : (i + 1) * seq_len_q].unflatten( + -1, [num_heads, fused_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] + k_pe_seq = k_pe[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze(-2) + + cos, sin = rope_cos_sin[past_len : past_len + seq_len_q].chunk(2, dim=-2) + + # Apply RoPE to q_pe + q_pe_seq = q_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) + q_pe_seq = q_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + fused_q_seq[..., -qk_rope_head_dim:] = q_pe_seq + + # Apply RoPE to k_pe + k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + k_pe_seq = ((k_pe_seq * cos) + (rotate_half(k_pe_seq) * sin)).to(dtype=k_pe_seq.dtype) + k_pe_seq = k_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + + # Build new-token latent: [compressed_kv, rotated_k_pe] + new_latent = torch.cat([compressed_kv_seq.unsqueeze(-2), k_pe_seq], dim=-1).squeeze( + -2 + ) # [seq_len_q, head_dim] + new_latent_list.append(new_latent) + + return fused_q_rot, torch.cat(new_latent_list, dim=0) + + +def calculate_deepseek_v4_ref_gen_sparse( + fused_q_rot: torch.Tensor, + new_latent_cache: torch.Tensor, + latent_cache_ref: torch.Tensor, + compressed_ref_data: Optional[List[torch.Tensor]], + compressed_topk_indices: Optional[torch.Tensor], + swa_window_size: int, + seq_lens_kv: List[int], + num_heads: int, + kv_lora_rank: int, + v_head_dim: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + q_scaling: float, + compress_ratio: int, + attn_sink: Optional[torch.Tensor] = None, +): + """Reference attention for DeepSeek-V4 generation phase.""" + fused_head_dim = kv_lora_rank + qk_rope_head_dim + bmm1_scale = 1 / (math.sqrt(qk_nope_head_dim + qk_rope_head_dim) * q_scaling) + num_requests = len(seq_lens_kv) + seq_len_q = fused_q_rot.shape[0] // num_requests + + ref_results = [] + latent_cache_list = [] + total_past_tokens = 0 + for i in range(num_requests): + past_len = seq_lens_kv[i] + + fused_q_seq = fused_q_rot[i * seq_len_q : (i + 1) * seq_len_q].unflatten( + -1, [num_heads, fused_head_dim] + ) + + # New token's latent cache + new_token_latent = new_latent_cache[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze( + -2 + ) # [seq_len_q, 1, head_dim] + + # Past latent cache + past_latent = latent_cache_ref[total_past_tokens : total_past_tokens + past_len].unsqueeze( + -2 + ) + + # Combine past + new + full_latent = torch.cat([past_latent, new_token_latent], dim=0) + latent_cache_list.append(full_latent) + + # For each query token, gather the relevant KV + for qi in range(seq_len_q): + q_tok = fused_q_seq[qi] # [num_heads, fused_head_dim] + current_kv_len = past_len + qi + 1 + current_pos = past_len + qi + + # SWA window + swa_start = max(0, current_pos - swa_window_size + 1) + swa_end = current_pos + 1 + swa_kv = full_latent[swa_start:swa_end].squeeze(-2) + + # Compressed KV + if compress_ratio > 1 and compressed_ref_data is not None: + if compress_ratio == 4 and compressed_topk_indices is not None: + row = i * seq_len_q + qi + indices_row = compressed_topk_indices[row] + valid = indices_row[indices_row >= 0] + comp_kv = compressed_ref_data[i][valid.long()] + elif compress_ratio == 128: + num_comp = current_kv_len // compress_ratio + comp_kv = compressed_ref_data[i][:num_comp] + else: + comp_kv = torch.empty( + 0, + swa_kv.shape[-1], + device=swa_kv.device, + dtype=swa_kv.dtype, + ) + + if comp_kv.numel() > 0: + all_kv = torch.cat([swa_kv, comp_kv], dim=0) + else: + all_kv = swa_kv + else: + all_kv = swa_kv + + k_sel = all_kv.unsqueeze(0).expand(num_heads, -1, -1) + v_sel = all_kv[:, :v_head_dim].unsqueeze(0).expand(num_heads, -1, -1) + + attn_w = torch.matmul(q_tok.unsqueeze(1), k_sel.transpose(1, 2)) * bmm1_scale + attn_w = _softmax_with_sink(attn_w, attn_sink, fused_q_rot.dtype) + out = torch.matmul(attn_w, v_sel).squeeze(1) + ref_results.append(out.reshape(1, num_heads * v_head_dim)) + + total_past_tokens += past_len + + ref_result = torch.cat(ref_results, dim=0) + new_latent_cache_out = torch.cat(latent_cache_list).squeeze(-2) + return ref_result, new_latent_cache_out + + +def _allocate_kv_cache_for_generation(cache_manager, requests: List[LlmRequest]): + for req in requests: + assert cache_manager.try_allocate_generation(req), ( + f"Failed to allocate generation KV cache for request {req.py_request_id}" + ) + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.parametrize("context_lengths", [[4399], [14, 508, 3947], [2, 1406, 3327]]) +@pytest.mark.parametrize("num_generation_steps", [2]) +def test_deepseek_v4_sparse_mla(context_lengths: List[int], num_generation_steps: int): + generation_seq_len_q = 1 + scenario = Scenario() + device = torch.device("cuda") + dtype = scenario.dtype + + 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_append = scenario.rope_append + # When rope_append is False, [448: 512) are used for qk_rope_head_dim + kv_lora_rank = ( + scenario.kv_lora_rank - qk_rope_head_dim if not rope_append else scenario.kv_lora_rank + ) + head_dim = kv_lora_rank + qk_rope_head_dim + # When rope_append is False, use 64 heads + num_heads = 64 if not rope_append else scenario.num_heads + batch_size = len(context_lengths) + max_context_len = max(context_lengths) + max_seq_len = max_context_len + (num_generation_steps + 1) * generation_seq_len_q + total_ctx_tokens = sum(context_lengths) + + torch.manual_seed(42) + + # 1. Setup cache manager + cache_manager, sparse_config = _create_cache_manager(scenario, context_lengths, max_seq_len) + request_ids = list(range(batch_size)) + + requests = [ + LlmRequest( + request_id=i, + max_new_tokens=num_generation_steps + 1, + input_tokens=list(range(ctx_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + for i, ctx_len in enumerate(context_lengths) + ] + scheduled_batch = ScheduledRequests() + for req in requests: + scheduled_batch.append_context_request(req) + for req in requests: + cache_manager.prepare_context(req) + cache_manager.resize_context(req, req.context_chunk_size) + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + + # 2. Create RoPE cos/sin + 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, + ) + 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) + ) + + # 3. Setup attention params + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, + rope=RopeParams.from_config(rope_config), + is_neox=False, + ) + mla_params = MLAParams( + q_lora_rank=scenario.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, + rope_append=rope_append, + predicted_tokens_per_seq=1, + hidden_size=scenario.hidden_size, + ) + + 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) + + # 4. Create attention layers + layers = {} + for layer_idx in TEST_LAYERS: + layer = DeepseekV4TrtllmAttention( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=1, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + sparse_attention_config=sparse_config, + skip_create_weights_in_init=True, + ) + layer.update_quant_config(None) + layers[layer_idx] = layer + + # Install a per-layer attention sink + attn_sinks: Dict[int, torch.Tensor] = {} + for layer_idx in TEST_LAYERS: + sink = torch.randn(num_heads, dtype=torch.float32, device=device).mul_(0.5) + layers[layer_idx].attn_sink = torch.nn.Parameter(sink, requires_grad=False) + attn_sinks[layer_idx] = sink + + # 5. Create random inputs per layer + inputs_per_layer = {} + for layer_idx in TEST_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_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_lengths + ] + ) + ctx_q = torch.cat( + [ + torch.empty( + [ctx_len, num_heads, kv_lora_rank], + dtype=dtype, + device=device, + ).uniform_(-1, 1) + for ctx_len in context_lengths + ] + ) + ctx_q_pe = torch.cat( + [ + torch.empty( + [ctx_len, num_heads, qk_rope_head_dim], + dtype=dtype, + device=device, + ).uniform_(-1, 1) + for ctx_len in context_lengths + ] + ) + ctx_fused_q = torch.cat([ctx_q, ctx_q_pe], dim=-1).view(-1, num_heads * 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_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_lengths + ] + ) + for _ in range(num_generation_steps) + ] + gen_q_list = [ + torch.cat( + [ + torch.empty( + [generation_seq_len_q, num_heads, kv_lora_rank], + dtype=dtype, + device=device, + ).uniform_(-1, 1) + for _ in context_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_lengths + ] + ) + for _ in range(num_generation_steps) + ] + gen_fused_q_list = [ + torch.cat([gen_q_list[i], gen_q_pe_list[i]], dim=-1).view(-1, num_heads * head_dim) + for i in range(num_generation_steps) + ] + + inputs_per_layer[layer_idx] = { + "ctx_compressed_kv": ctx_compressed_kv, + "ctx_k_pe": ctx_k_pe, + "ctx_q_pe": ctx_q_pe, + "ctx_fused_q": ctx_fused_q, + "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, + } + + # 6. Pre-fill COMPRESS buffers for layers with ratio > 1 + compress_ref_data: Dict[int, List[torch.Tensor]] = {} + for layer_idx in TEST_LAYERS: + ratio = scenario.compress_ratios[layer_idx] + if ratio > 1: + compress_ref_data[layer_idx] = _prefill_compress_buffer( + cache_manager, + layer_idx, + context_lengths, + request_ids, + head_dim, + device, + ) + + # 7. Context phase + ctx_seq_lens = torch.tensor(context_lengths, dtype=torch.int) + attn_metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=ctx_seq_lens, + request_ids=request_ids, + max_num_requests=batch_size, + num_contexts=batch_size, + prompt_lens=context_lengths, + max_num_tokens=total_ctx_tokens, + kv_cache_manager=cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0] * batch_size, + ), + mapping=mapping, + sparse_attention_config=sparse_config, + ) + attn_metadata.prepare() + + latent_cache_ref_all: Dict[int, torch.Tensor] = {} + for layer_idx in TEST_LAYERS: + ratio = scenario.compress_ratios[layer_idx] + print(f"\n--- Context phase: layer {layer_idx}, compress_ratio={ratio} ---") + + inp = inputs_per_layer[layer_idx] + fused_q = inp["ctx_fused_q"] + compressed_kv = inp["ctx_compressed_kv"] + k_pe = inp["ctx_k_pe"] + q_pe = inp["ctx_q_pe"] + latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + + # Build topk_indices for ratio=4 + topk_indices = None + if ratio == 4: + token_positions = [] + for ctx_len in context_lengths: + token_positions.extend(range(ctx_len)) + topk_indices = _build_compressed_topk_indices( + token_positions, ratio, scenario.index_topk, device + ) + + # Forward through the attention layer + result = layers[layer_idx].forward( + fused_q.clone(), + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=latent_cache, + q_pe=q_pe, + topk_indices=topk_indices, + is_generation=False, + ) + + # Reference computation + k_pe_ref = _rotate_k_pe_for_ctx(k_pe, rope_cos_sin, context_lengths) + latent_cache_ref = torch.cat([compressed_kv, k_pe_ref], dim=-1) + fused_q_rot = _rotate_fused_q_for_ctx( + fused_q, + rope_cos_sin, + context_lengths, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + ref_result = calculate_deepseek_v4_ref_ctx_sparse( + fused_q_rot, + latent_cache_ref, + compress_ref_data.get(layer_idx), + scenario.window_size, + topk_indices, + context_lengths, + num_heads, + kv_lora_rank, + v_head_dim, + qk_nope_head_dim, + qk_rope_head_dim, + q_scaling, + ratio, + attn_sink=attn_sinks[layer_idx], + ) + + latent_cache_ref_all[layer_idx] = latent_cache_ref + + # Check results + print(f" Output shape: {result.shape}") + print( + f" Output mean: {result.abs().mean().item():.6f}, max: {result.abs().max().item():.6f}" + ) + print( + f" Ref mean: {ref_result.abs().mean().item():.6f}, " + f"max: {ref_result.abs().max().item():.6f}" + ) + diff = (result - ref_result).abs() + print(f" Diff mean: {diff.mean().item():.6f}, max: {diff.max().item():.6f}") + + assert torch.allclose(result, ref_result, atol=0.2, rtol=0.02), ( + f"Context phase mismatch at layer {layer_idx} (ratio={ratio}): " + f"max diff={diff.max().item():.6f}" + ) + print(" PASSED") + + # 8. Generation steps + for step in range(num_generation_steps): + print(f"\n=== Generation step {step + 1} ===") + _allocate_kv_cache_for_generation(cache_manager, requests) + + cached_lens = [ctx_len + step * generation_seq_len_q for ctx_len in context_lengths] + kv_lens = [cl + generation_seq_len_q for cl in cached_lens] + + # Grow compressed ref data: as KV length increases, new compressed + # tokens appear. Simulate this by appending random data. + for layer_idx in TEST_LAYERS: + ratio = scenario.compress_ratios[layer_idx] + if ratio > 1 and layer_idx in compress_ref_data: + _grow_compress_buffer_for_generation( + cache_manager, + layer_idx, + request_ids, + head_dim, + device, + kv_lens, + compress_ref_data[layer_idx], + ) + + gen_seq_lens = torch.tensor([generation_seq_len_q] * batch_size, dtype=torch.int) + total_gen_tokens = batch_size * generation_seq_len_q + gen_metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=gen_seq_lens, + request_ids=request_ids, + max_num_requests=batch_size, + num_contexts=0, + prompt_lens=context_lengths, + max_num_tokens=total_gen_tokens, + kv_cache_manager=cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=cached_lens, + ), + mapping=mapping, + enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), + sparse_attention_config=sparse_config, + ) + gen_metadata.prepare() + + for layer_idx in TEST_LAYERS: + ratio = scenario.compress_ratios[layer_idx] + print(f"\n--- Gen step {step + 1}: layer {layer_idx}, compress_ratio={ratio} ---") + + inp = inputs_per_layer[layer_idx] + fused_q = inp["gen_fused_q_list"][step] + q_pe = inp["gen_q_pe_list"][step] + compressed_kv = inp["gen_compressed_kv_list"][step] + k_pe = inp["gen_k_pe_list"][step] + latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + + # Prepare generation-specific tensors + num_seqs = gen_metadata.kv_lens_cuda_runtime.size(0) + cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) + cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=device) + + layers[layer_idx].mla_rope_generation( + fused_q, + q_pe, + latent_cache, + gen_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + None, # mla_bmm1_scale + None, # mla_bmm2_scale + None, # quant_q_buffer + ) + + # Build topk_indices for ratio=4 + topk_indices = None + if ratio == 4: + topk_indices = _build_compressed_topk_indices( + [kv - 1 for kv in kv_lens], + ratio, + scenario.index_topk, + device, + ) + + result = layers[layer_idx].forward( + fused_q, + None, + None, + gen_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=latent_cache, + q_pe=q_pe, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + topk_indices=topk_indices, + is_generation=True, + ) + + # Reference: apply RoPE separately, then compute attention + fused_q_rot, new_latent = _rotate_gen_inputs( + fused_q, + q_pe, + compressed_kv, + k_pe, + rope_cos_sin, + cached_lens, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + ref_result, new_latent_cache = calculate_deepseek_v4_ref_gen_sparse( + fused_q_rot, + new_latent, + latent_cache_ref_all[layer_idx], + compress_ref_data.get(layer_idx), + topk_indices, + scenario.window_size, + cached_lens, + num_heads, + kv_lora_rank, + v_head_dim, + qk_nope_head_dim, + qk_rope_head_dim, + q_scaling, + ratio, + attn_sink=attn_sinks[layer_idx], + ) + latent_cache_ref_all[layer_idx] = new_latent_cache + + print(f" Output shape: {result.shape}") + print( + f" Output mean: {result.abs().mean().item():.6f}, " + f"max: {result.abs().max().item():.6f}" + ) + print( + f" Ref mean: {ref_result.abs().mean().item():.6f}, " + f"max: {ref_result.abs().max().item():.6f}" + ) + diff = (result - ref_result).abs() + print(f" Diff mean: {diff.mean().item():.6f}, max: {diff.max().item():.6f}") + + assert torch.allclose(result, ref_result, atol=0.2, rtol=0.02), ( + f"Gen step {step + 1} mismatch at layer {layer_idx} " + f"(ratio={ratio}): max diff={diff.max().item():.6f}" + ) + print(" PASSED") + + cache_manager.shutdown() + print("\nAll tests passed!") + + +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.parametrize("context_lengths", [[14, 508, 3947]]) +def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): + scenario = Scenario() + device = torch.device("cuda") + dtype = scenario.dtype + generation_seq_len_q = 1 + + 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_append = scenario.rope_append + kv_lora_rank = ( + scenario.kv_lora_rank - qk_rope_head_dim if not rope_append else scenario.kv_lora_rank + ) + head_dim = kv_lora_rank + qk_rope_head_dim + num_heads = 64 if not rope_append else scenario.num_heads + mscale = scenario.rope_mscale + q_scaling = 1.0 / (mscale * mscale) + + batch_size = len(context_lengths) + assert batch_size >= 2 + num_ctx = 1 + num_gen = batch_size - num_ctx + max_context_len = max(context_lengths) + max_seq_len = max_context_len + 2 * generation_seq_len_q + total_ctx_tokens = context_lengths[0] + total_gen_tokens = num_gen * generation_seq_len_q + total_mixed_tokens = total_ctx_tokens + total_gen_tokens + + torch.manual_seed(42) + + # 1. Setup cache, layers, RoPE + cache_manager, sparse_config = _create_cache_manager(scenario, context_lengths, max_seq_len) + request_ids = list(range(batch_size)) + + requests = [ + LlmRequest( + request_id=i, + max_new_tokens=2, + input_tokens=list(range(ctx_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) + for i, ctx_len in enumerate(context_lengths) + ] + scheduled_batch = ScheduledRequests() + for req in requests: + scheduled_batch.append_context_request(req) + for req in requests: + cache_manager.prepare_context(req) + cache_manager.resize_context(req, req.context_chunk_size) + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + 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, + ) + 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) + ) + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, + rope=RopeParams.from_config(rope_config), + is_neox=False, + ) + mla_params = MLAParams( + q_lora_rank=scenario.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, + rope_append=rope_append, + predicted_tokens_per_seq=1, + ) + + layers = {} + for li in TEST_LAYERS: + layer = DeepseekV4TrtllmAttention( + layer_idx=li, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=scenario.num_kv_heads, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + sparse_attention_config=sparse_config, + dtype=dtype, + ) + layer.update_quant_config(None) + layers[li] = layer + + attn_sinks: Dict[int, torch.Tensor] = {} + for li in TEST_LAYERS: + sink = torch.randn(num_heads, dtype=torch.float32, device=device).mul_(0.5) + layers[li].attn_sink = torch.nn.Parameter(sink, requires_grad=False) + attn_sinks[li] = sink + + # 2. Pre-fill KV cache for gen requests (1..N) — all layers, all ratios. + gen_ctx_lengths = context_lengths[num_ctx:] + gen_request_ids = request_ids[num_ctx:] + gen_total_ctx = sum(gen_ctx_lengths) + + prefill_fused_q = torch.empty( + [gen_total_ctx, num_heads * head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + prefill_k_pe = torch.empty( + [gen_total_ctx, qk_rope_head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + prefill_q_pe = torch.empty( + [gen_total_ctx, num_heads, qk_rope_head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + prefill_compressed_kv = torch.empty( + [gen_total_ctx, kv_lora_rank], dtype=dtype, device=device + ).uniform_(-1, 1) + prefill_latent = torch.cat([prefill_compressed_kv, prefill_k_pe], dim=-1) + + prefill_metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(gen_ctx_lengths, dtype=torch.int), + request_ids=gen_request_ids, + max_num_requests=num_gen, + num_contexts=num_gen, + prompt_lens=gen_ctx_lengths, + max_num_tokens=gen_total_ctx, + kv_cache_manager=cache_manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0] * num_gen), + mapping=mapping, + sparse_attention_config=sparse_config, + ) + prefill_metadata.prepare() + for li in TEST_LAYERS: + ratio = scenario.compress_ratios[li] + prefill_topk = None + if ratio == 4: + positions = [] + for cl in gen_ctx_lengths: + positions.extend(range(cl)) + prefill_topk = _build_compressed_topk_indices( + positions, ratio, scenario.index_topk, device + ) + layers[li].forward( + prefill_fused_q.clone(), + None, + None, + prefill_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=prefill_latent, + q_pe=prefill_q_pe, + topk_indices=prefill_topk, + is_generation=False, + ) + + # Pre-fill COMPRESS buffers for ratio > 1. + compress_ref_data: Dict[int, List[torch.Tensor]] = {} + for li in TEST_LAYERS: + ratio = scenario.compress_ratios[li] + if ratio > 1: + compress_ref_data[li] = _prefill_compress_buffer( + cache_manager, + li, + gen_ctx_lengths, + gen_request_ids, + head_dim, + device, + ) + + # 3. Allocate 1 gen step for gen requests. + gen_requests = requests[1:] + _allocate_kv_cache_for_generation(cache_manager, gen_requests) + gen_cached_lens = [cl + generation_seq_len_q for cl in gen_ctx_lengths] + + # Grow compress buffers for gen step. + gen_kv_lens = [cl + generation_seq_len_q for cl in gen_cached_lens] + for li in TEST_LAYERS: + ratio = scenario.compress_ratios[li] + if ratio > 1 and li in compress_ref_data: + _grow_compress_buffer_for_generation( + cache_manager, + li, + gen_request_ids, + head_dim, + device, + gen_kv_lens, + compress_ref_data[li], + ) + + # 4. Mixed metadata: request 0 = context, requests 1..N = generation. + mixed_seq_lens = [context_lengths[0]] + [generation_seq_len_q] * num_gen + mixed_cached_lens = [0] + gen_cached_lens + mixed_metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor(mixed_seq_lens, dtype=torch.int), + request_ids=request_ids, + max_num_requests=batch_size, + num_contexts=num_ctx, + prompt_lens=context_lengths, + max_num_tokens=total_mixed_tokens, + kv_cache_manager=cache_manager, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=mixed_cached_lens), + mapping=mapping, + enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), + sparse_attention_config=sparse_config, + ) + mixed_metadata.prepare() + + # 5. Per-layer forward + verify (mirrors forward_impl_with_deepseek_v4). + for li in TEST_LAYERS: + ratio = scenario.compress_ratios[li] + print(f"\n--- Mixed: layer {li}, compress_ratio={ratio} ---") + + # Random inputs + ctx_q_nope = torch.empty( + [total_ctx_tokens, num_heads, kv_lora_rank], dtype=dtype, device=device + ).uniform_(-1, 1) + ctx_q_pe = torch.empty( + [total_ctx_tokens, num_heads, qk_rope_head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + ctx_fused_q = torch.cat([ctx_q_nope, ctx_q_pe], dim=-1).view(-1, num_heads * head_dim) + ctx_k_pe = torch.empty( + [total_ctx_tokens, qk_rope_head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + ctx_compressed_kv = torch.empty( + [total_ctx_tokens, kv_lora_rank], dtype=dtype, device=device + ).uniform_(-1, 1) + ctx_latent = torch.cat([ctx_compressed_kv, ctx_k_pe], dim=-1) + + gen_fused_q = torch.empty( + [total_gen_tokens, num_heads * head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + gen_q_pe = torch.empty( + [total_gen_tokens, num_heads, qk_rope_head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + gen_compressed_kv = torch.empty( + [total_gen_tokens, kv_lora_rank], dtype=dtype, device=device + ).uniform_(-1, 1) + gen_k_pe = torch.empty( + [total_gen_tokens, qk_rope_head_dim], dtype=dtype, device=device + ).uniform_(-1, 1) + gen_latent = torch.cat([gen_compressed_kv, gen_k_pe], dim=-1) + + output = torch.empty( + [total_mixed_tokens, num_heads * v_head_dim], dtype=dtype, device=device + ) + + # topk_indices for ratio=4 + ctx_topk = gen_topk = None + if ratio == 4: + ctx_positions = list(range(context_lengths[0])) + ctx_topk = _build_compressed_topk_indices( + ctx_positions, ratio, scenario.index_topk, device + ) + gen_topk = _build_compressed_topk_indices( + [kv - 1 for kv in [cl + generation_seq_len_q for cl in gen_cached_lens]], + ratio, + scenario.index_topk, + device, + ) + + # Context forward → output[:total_ctx_tokens] + layers[li].forward( + ctx_fused_q.clone(), + None, + None, + mixed_metadata, + attention_input_type=AttentionInputType.context_only, + output=output[:total_ctx_tokens], + latent_cache=ctx_latent, + q_pe=ctx_q_pe, + topk_indices=ctx_topk, + is_generation=False, + ) + + # Generation forward → output[total_ctx_tokens:] + num_seqs = mixed_metadata.kv_lens_cuda_runtime.size(0) + cu_q = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) + cu_kv = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) + counter = torch.empty(1, dtype=torch.uint32, device=device) + layers[li].mla_rope_generation( + gen_fused_q, + gen_q_pe, + gen_latent, + mixed_metadata, + cu_q, + cu_kv, + counter, + None, + None, + None, + ) + layers[li].forward( + gen_fused_q, + None, + None, + mixed_metadata, + attention_input_type=AttentionInputType.generation_only, + output=output[total_ctx_tokens:total_mixed_tokens], + latent_cache=gen_latent, + q_pe=gen_q_pe, + cu_q_seqlens=cu_q, + cu_kv_seqlens=cu_kv, + fmha_scheduler_counter=counter, + topk_indices=gen_topk, + is_generation=True, + ) + + # Context reference + ctx_k_pe_ref = _rotate_k_pe_for_ctx(ctx_k_pe, rope_cos_sin, [context_lengths[0]]) + ctx_latent_ref = torch.cat([ctx_compressed_kv, ctx_k_pe_ref], dim=-1) + ctx_q_rot = _rotate_fused_q_for_ctx( + ctx_fused_q, + rope_cos_sin, + [context_lengths[0]], + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + ctx_ref = calculate_deepseek_v4_ref_ctx_sparse( + ctx_q_rot, + ctx_latent_ref, + None, + scenario.window_size, + ctx_topk, + [context_lengths[0]], + num_heads, + kv_lora_rank, + v_head_dim, + qk_nope_head_dim, + qk_rope_head_dim, + q_scaling, + ratio, + attn_sink=attn_sinks[li], + ) + + # Generation reference + prefill_k_pe_ref = _rotate_k_pe_for_ctx(prefill_k_pe, rope_cos_sin, gen_ctx_lengths) + gen_latent_ref = torch.cat([prefill_compressed_kv, prefill_k_pe_ref], dim=-1) + gen_q_rot, new_latent = _rotate_gen_inputs( + gen_fused_q, + gen_q_pe, + gen_compressed_kv, + gen_k_pe, + rope_cos_sin, + gen_cached_lens, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + gen_ref, _ = calculate_deepseek_v4_ref_gen_sparse( + gen_q_rot, + new_latent, + gen_latent_ref, + compress_ref_data.get(li), + gen_topk, + scenario.window_size, + gen_cached_lens, + num_heads, + kv_lora_rank, + v_head_dim, + qk_nope_head_dim, + qk_rope_head_dim, + q_scaling, + ratio, + attn_sink=attn_sinks[li], + ) + + ctx_diff = (output[:total_ctx_tokens] - ctx_ref).abs() + gen_diff = (output[total_ctx_tokens:] - gen_ref).abs() + print( + f" Context: diff mean={ctx_diff.mean().item():.6f}, max={ctx_diff.max().item():.6f}" + ) + print( + f" Generation: diff mean={gen_diff.mean().item():.6f}, max={gen_diff.max().item():.6f}" + ) + assert torch.allclose(output[:total_ctx_tokens], ctx_ref, atol=0.2, rtol=0.02), ( + f"Mixed ctx mismatch layer {li} ratio={ratio}: max diff={ctx_diff.max().item():.6f}" + ) + assert torch.allclose(output[total_ctx_tokens:], gen_ref, atol=0.2, rtol=0.02), ( + f"Mixed gen mismatch layer {li} ratio={ratio}: max diff={gen_diff.max().item():.6f}" + ) + print(" PASSED") + + cache_manager.shutdown() + print("\nMixed batch test PASSED!") + + +if __name__ == "__main__": + test_deepseek_v4_sparse_mla(context_lengths=[4399], num_generation_steps=2) diff --git a/tests/unittest/_torch/attention/sparse/dsa/__init__.py b/tests/unittest/_torch/attention/sparse/dsa/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py similarity index 91% rename from tests/unittest/_torch/attention/sparse/test_dsa_indexer.py rename to tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 67b4d5641565..06ed2743543e 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -17,19 +17,22 @@ from utils.util import check_accuracy, skip_pre_blackwell, skip_pre_hopper from tensorrt_llm import deep_gemm -from tensorrt_llm._torch.attention_backend.interface import ( - PositionalEmbeddingParams, RopeParams) +from tensorrt_llm._torch.attention_backend.interface import PositionalEmbeddingParams, RopeParams from tensorrt_llm._torch.attention_backend.sparse.dsa import ( - DSACacheManager, DSAtrtllmAttentionMetadata, Indexer, - compute_cu_seqlen_kv_bounds_with_cache, split_prefill_chunks) + DSACacheManager, + DSAtrtllmAttentionMetadata, + Indexer, + compute_cu_seqlen_kv_bounds_with_cache, + split_prefill_chunks, +) from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, - restore_attn_metadata_after_draft_replay) + restore_attn_metadata_after_draft_replay, +) from tensorrt_llm._utils import get_sm_version from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.bindings.internal.batch_manager import \ - CacheType as CacheTypeCpp +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.deep_gemm import fp8_paged_mqa_logits from tensorrt_llm.functional import PositionEmbeddingType from tensorrt_llm.mapping import Mapping @@ -195,9 +198,11 @@ def _ref_fp8_paged_mqa_logits( q: torch.Tensor, kv_cache: torch.Tensor, weights: torch.Tensor, - context_lens: torch.Tensor, + num_tokens: torch.Tensor, + num_kv_tokens: torch.Tensor, block_tables: torch.Tensor, max_model_len: int, + compress_ratio: int = 1, ): """ Reference implementation of fp8_paged_mqa_logits (optimized version). @@ -206,9 +211,11 @@ def _ref_fp8_paged_mqa_logits( q: [batch_size, next_n, num_heads, head_dim] kv_cache: [num_blocks, block_size, 1, head_dim] weights: [batch_size * next_n, num_heads] - context_lens: [batch_size] + num_tokens: [batch_size] + num_kv_tokens: [batch_size] block_tables: [batch_size, max_num_blocks] max_model_len: Maximum sequence length + compress_ratio: Compression ratio for the KV cache Returns: logits: [batch_size * next_n, max_model_len] @@ -223,22 +230,22 @@ def _ref_fp8_paged_mqa_logits( dtype=torch.float32, ) - context_lens_list = context_lens.tolist() + num_tokens_list = num_tokens.tolist() + num_kv_tokens_list = num_kv_tokens.tolist() for i in range(batch_size): - context_len = context_lens_list[i] + num_token = num_tokens_list[i] + num_kv_token = num_kv_tokens_list[i] - # Query positions: [context_len - next_n, ..., context_len - 1] - q_offsets = torch.arange(context_len - next_n, - context_len, - device="cuda") + # Query positions: [num_token - next_n, ..., num_token - 1] + q_offsets = torch.arange(num_token - next_n, num_token, device="cuda") # Transpose weights for this sequence: [num_heads, next_n] weight_slice = (weights[i * next_n:(i + 1) * next_n, :].transpose( 0, 1).contiguous()) # Process each block in the sequence - for block_rk in range(cdiv(context_len, block_size)): + for block_rk in range(cdiv(num_kv_token, block_size)): block_idx = block_tables[i][block_rk] qx, kx = q[i], kv_cache[block_idx] @@ -249,9 +256,11 @@ def _ref_fp8_paged_mqa_logits( device="cuda", ) - # Causal mask: k_pos < context_len AND k_pos <= q_pos - mask = (k_offsets[None, :] < context_len) & (k_offsets[None, :] - <= q_offsets[:, None]) + # Causal mask: k_pos < num_token AND k_pos < (q_pos + 1) // compress_ratio + k_mask = k_offsets[None, :] < num_kv_token + causal_mask = k_offsets[None, :] < (q_offsets[:, None] + + 1) // compress_ratio + mask = k_mask & causal_mask # Compute attention scores: [num_heads, next_n, block_size] s = torch.where( @@ -269,8 +278,7 @@ def _ref_fp8_paged_mqa_logits( logits[ i * next_n:(i + 1) * next_n, block_rk * block_size:(block_rk + 1) * block_size, - ] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, - float("-inf")) + ] = torch.where(causal_mask, s, float("-inf")) return logits @@ -320,7 +328,8 @@ def _ref_fp8_mqa_logits( @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper -def test_deepgemm_fp8_mqa_logits_basic(): +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_deepgemm_fp8_mqa_logits_basic(compress_ratio): """ Basic test for deepgemm.fp8_mqa_logits kernel. Tests the disable_cp path with simple validation. @@ -330,6 +339,7 @@ def test_deepgemm_fp8_mqa_logits_basic(): num_heads, head_dim = 64, 128 seq_len = 2048 seq_len_kv = 4096 + seq_len_kv_compressed = seq_len_kv // compress_ratio #[seq_len, num_heads, head_dim] q = torch.randn( seq_len, @@ -340,7 +350,7 @@ def test_deepgemm_fp8_mqa_logits_basic(): ) #[seq_len_kv, head_dim] -> num_head = 1 kv = torch.randn( - seq_len_kv, + seq_len_kv_compressed, head_dim, device="cuda", dtype=torch.bfloat16, @@ -353,9 +363,15 @@ def test_deepgemm_fp8_mqa_logits_basic(): dtype=torch.float32, ) # ks[i] -> ke[i] for each q[i] - ks = torch.zeros(seq_len, dtype=torch.int, device="cuda") - ke = torch.arange(seq_len, dtype=torch.int, - device="cuda") + (seq_len_kv - seq_len) + ks, ke = compute_cu_seqlen_kv_bounds_with_cache( + seq_lens=torch.tensor([seq_len], dtype=torch.int, device="cuda"), + num_contexts=1, + num_ctx_tokens=seq_len, + cached_token_lens=torch.tensor([512], dtype=torch.int, device="cuda"), + kv_lens=torch.tensor([seq_len_kv_compressed], + dtype=torch.int, + device="cuda"), + compress_ratio=compress_ratio) # Convert to FP8 q_fp8 = q.to(torch.float8_e4m3fn) @@ -364,8 +380,8 @@ def test_deepgemm_fp8_mqa_logits_basic(): ke) # -> [seq_len, seq_len_kv] # Basic sanity checks - assert logits.shape == (seq_len, seq_len_kv), \ - f"Expected shape ({seq_len}, {seq_len_kv}), got {logits.shape}" + assert logits.shape == (seq_len, seq_len_kv_compressed), \ + f"Expected shape ({seq_len}, {seq_len_kv_compressed}), got {logits.shape}" assert logits.dtype == torch.float32, \ f"Expected dtype torch.float32, got {logits.dtype}" @@ -405,7 +421,9 @@ def _create_mock_metadata(request_ids, enable_context_mla_with_cached_kv=False, index_topk=2048, enable_indexer_skip=False, - use_cute_dsl_paged_mqa_logits=False): + use_cute_dsl_paged_mqa_logits=False, + compress_ratio=1, + indexer_head_dim=128): """Helper to create mock metadata for testing.""" class MockKVCacheParams: @@ -428,12 +446,22 @@ def __init__(self): self.max_draft_tokens = max_draft_tokens self.num_sparse_topk = index_topk self.enable_indexer_skip = enable_indexer_skip + self.indexer_head_dim = indexer_head_dim + self.indexer_quant_block_size = 128 + self.compress_ratios = [compress_ratio] # Keep seq_lens on CPU for split_prefill_chunks and other CPU operations # CUDA kernels will convert to CUDA as needed self.seq_lens = seq_lens.cpu() if seq_lens.is_cuda else seq_lens self.kv_lens = kv_lens self.kv_cache_params = MockKVCacheParams() self.kv_cache_manager = cache_manager + # Effective tokens-per-block for the indexer k-cache slot mapping, + # mirroring DSAtrtllmAttentionMetadata.__post_init__ (dsa.py). + tpb = self.kv_cache_manager.tokens_per_block + if hasattr(self.kv_cache_manager, 'compressed_block_sizes'): + _cr = 4 if 4 in self.compress_ratios else 1 + tpb = tpb // _cr + self._tokens_per_block = tpb self.kv_lens_cuda_runtime = kv_lens.cuda() self.indexer_k_cache_block_offsets = torch.zeros( [batch_size, cache_manager.max_blocks_per_seq], @@ -729,6 +757,7 @@ def test_indexer_k_cache_scatter_custom_op(): cache_manager=cache_manager, num_ctx_tokens=num_tokens, num_tokens=num_tokens, + indexer_head_dim=head_dim, ) from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer @@ -893,6 +922,7 @@ def test_fp8_k_cache_roundtrip(): cache_manager=cache_manager, num_ctx_tokens=total_tokens, num_tokens=total_tokens, + indexer_head_dim=head_dim, ) Indexer.prepare(metadata) @@ -951,7 +981,9 @@ def test_fp8_k_cache_roundtrip(): f"CuTe DSL FP8 Paged MQA Logits only supports SM 100/103, got SM {get_sm_version()}", )), ]) -def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, + compress_ratio): """ Test FP8 paged KV cache with two-phase workflow and variable context lengths. @@ -983,7 +1015,14 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): # Final lengths after generation phase final_lens = context_lens_context + num_gen_tokens + final_kv_lens = final_lens // compress_ratio + num_ctx_kv_tokens = context_lens_context // compress_ratio + num_gen_kv_tokens = final_kv_lens - num_ctx_kv_tokens max_seq_len = final_lens.max().item() + max_kv_len = final_kv_lens.max().item() + total_context_tokens = context_lens_context.sum().item() + total_context_kv_tokens = num_ctx_kv_tokens.sum().item() + total_gen_kv_tokens = num_gen_kv_tokens.sum().item() print("\n=== Test Config ===") print( @@ -991,36 +1030,38 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): ) print(f" Context lengths: {context_lens_context.tolist()}") print(f" Final lengths: {final_lens.tolist()}") + print(f" Final kv lengths: {final_kv_lens.tolist()}") print(f" Max sequence length: {max_seq_len}") + if compress_ratio != 1: + print(f" Compress ratio: {compress_ratio}") # Setup: Create cache manager and indexer cache_manager, sparse_attn_config = create_dsa_cache_manager( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, - max_seq_len=max_model_len, + max_seq_len=max_kv_len, num_layers=1) indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) # Allocate blocks for all sequences (max final length) request_ids = list(range(batch_size)) cache_manager.add_dummy_requests(request_ids=request_ids, - token_nums=final_lens.tolist(), + token_nums=final_kv_lens.tolist(), is_gen=False, prepare_resource=True) # Generate test data with variable lengths - total_context_tokens = context_lens_context.sum().item() q = torch.randn((batch_size, next_n, heads, head_dim), device="cuda", dtype=torch.bfloat16) weights = torch.randn((batch_size * next_n, heads), device="cuda", dtype=torch.float32) - k_context_bf16 = torch.randn((total_context_tokens, head_dim), + k_context_bf16 = torch.randn((total_context_kv_tokens, head_dim), device="cuda", dtype=torch.bfloat16) - k_gen_bf16 = torch.randn((batch_size * num_gen_tokens, head_dim), + k_gen_bf16 = torch.randn((total_gen_kv_tokens, head_dim), device="cuda", dtype=torch.bfloat16) @@ -1032,13 +1073,15 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): num_contexts=batch_size, num_generations=0, seq_lens=context_lens_context.clone(), - kv_lens=context_lens_context.clone(), + kv_lens=num_ctx_kv_tokens.clone(), num_cached_tokens=[0] * batch_size, cache_manager=cache_manager, num_ctx_tokens=total_context_tokens, num_tokens=total_context_tokens, max_draft_tokens=next_n - 1, use_cute_dsl_paged_mqa_logits=use_dsl, + compress_ratio=compress_ratio, + indexer_head_dim=head_dim, ) Indexer.prepare(metadata_context) @@ -1046,7 +1089,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): k_context_bf16) indexer._update_k_cache(k_context_fp8, k_context_scale, metadata_context) - print(f"✓ Wrote {total_context_tokens} FP8 context tokens to cache") + print(f"✓ Wrote {total_context_kv_tokens} FP8 context tokens to cache") # Phase 2: Write generation tokens (next_n per sequence) as FP8 # Similar to prepare_resources: add_token() for each new token @@ -1059,21 +1102,22 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): seq_lens=torch.tensor([num_gen_tokens] * batch_size, dtype=torch.int32, device='cpu'), - kv_lens=final_lens.clone(), + kv_lens=final_kv_lens.clone(), num_cached_tokens=context_lens_context.tolist(), cache_manager=cache_manager, num_ctx_tokens=0, num_tokens=batch_size * num_gen_tokens, max_draft_tokens=next_n - 1, use_cute_dsl_paged_mqa_logits=use_dsl, + compress_ratio=compress_ratio, + indexer_head_dim=head_dim, ) Indexer.prepare(metadata_gen) k_gen_fp8, k_gen_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( k_gen_bf16) indexer._update_k_cache(k_gen_fp8, k_gen_scale, metadata_gen) - print( - f"✓ Wrote {batch_size * num_gen_tokens} FP8 generation tokens to cache") + print(f"✓ Wrote {total_gen_kv_tokens} FP8 generation tokens to cache") # Run kernel: FP8 paged MQA with actual cache print("\n=== Kernel Execution ===") @@ -1129,7 +1173,8 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): context_offset = 0 gen_offset = 0 for seq_idx in range(batch_size): - seq_context_len = context_lens_context[seq_idx].item() + seq_context_len = num_ctx_kv_tokens[seq_idx].item() + seq_gen_len = num_gen_kv_tokens[seq_idx].item() # Write context tokens for token_pos in range(seq_context_len): @@ -1142,7 +1187,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): k_context_bf16[context_offset + token_pos] # Write generation tokens - for gen_token_idx in range(num_gen_tokens): + for gen_token_idx in range(seq_gen_len): token_pos = seq_context_len + gen_token_idx block_idx = token_pos // block_size pos_in_block = token_pos % block_size @@ -1153,17 +1198,19 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): k_gen_bf16[gen_offset + gen_token_idx] context_offset += seq_context_len - gen_offset += num_gen_tokens + gen_offset += seq_gen_len + num_tokens_cuda = final_lens.cuda() + num_kv_tokens_cuda = metadata_gen.kv_lens_cuda_runtime ref_logits = _ref_fp8_paged_mqa_logits( - q, kv_cache_bf16, weights, - metadata_gen.kv_lens_cuda_runtime[0:batch_size], - metadata_gen.indexer_k_cache_block_offsets, max_model_len) + q, kv_cache_bf16, weights, num_tokens_cuda[0:batch_size], + num_kv_tokens_cuda[0:batch_size], + metadata_gen.indexer_k_cache_block_offsets, max_model_len, + compress_ratio) print(f"✓ Reference output shape: {ref_logits.shape}") # Validate: Compare masked outputs (handle variable lengths and next_n) print("\n=== Validation ===") - context_lens_cuda = metadata_gen.kv_lens_cuda_runtime # [batch_size] # Expand context lens for each query: each sequence has next_n queries # Query at position i (where i = 0..next_n-1) attends to tokens up to (context_len - next_n + i) @@ -1177,12 +1224,12 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend): next_n_offset = torch.arange( batch_size * next_n, device="cuda") % next_n # Query offset within sequence - query_end_positions = context_lens_cuda[ + query_end_positions = num_tokens_cuda[ row_indices] - next_n + next_n_offset # [batch_size * next_n] # Create mask: positions <= query_end_position # Shape: [batch_size * next_n, max_model_len] - mask = positions <= query_end_positions.unsqueeze(1) + mask = positions < (query_end_positions.unsqueeze(1) + 1) // compress_ratio diff = _calc_diff(logits.masked_fill(~mask, 0), ref_logits.masked_fill(~mask, 0)) @@ -1841,6 +1888,90 @@ def test_compute_cu_seqlen_bounds_with_cache_properties(): f"Trial {trial}: Last ke should equal total KV count: {cu_seqlen_ke[-1].item()} != {expected_total_kv}" +def test_compute_cu_seqlen_bounds_nocache_compressed_kv(): + """Simple test case with 2 sequences and compressed KV.""" + seq_lens = torch.tensor([3, 4], dtype=torch.int32, device="cuda") + num_contexts = 2 + num_ctx_tokens = 7 + compress_ratio = 2 + + kv_lens = seq_lens // compress_ratio + + cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( + seq_lens, num_contexts, num_ctx_tokens, None, kv_lens, compress_ratio) + + # Expected results: + # Seq 0: tokens [0,1,2], KV [0] + # Token 0: [0, 0) + # Token 1: [0, 1) + # Token 2: [0, 1) + # Seq 1: tokens [3,4,5,6], KV [1,2] + # Token 3: [1, 1) + # Token 4: [1, 2) + # Token 5: [1, 2) + # Token 6: [1, 3) + + expected_ks = torch.tensor([0, 0, 0, 1, 1, 1, 1], + dtype=torch.int32, + device="cuda") + expected_ke = torch.tensor([0, 1, 1, 1, 2, 2, 3], + dtype=torch.int32, + device="cuda") + + assert torch.equal(cu_seqlen_ks, expected_ks), \ + f"cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" + assert torch.equal(cu_seqlen_ke, expected_ke), \ + f"cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + + +def test_compute_cu_seqlen_bounds_with_cache_compressed_kv(): + """ + Test case with 2 sequences using chunked prefill (with cached tokens) and compressed KV. + + Scenario: + - Seq 0: 2 cached tokens, 3 new tokens being added + - Seq 1: 1 cached token, 4 new tokens being added + """ + compress_ratio = 2 + # New tokens being added in this chunk + seq_lens = torch.tensor([3, 4], dtype=torch.int32, device="cuda") + # Previously cached tokens + num_past_tokens = torch.tensor([2, 1], dtype=torch.int32, device="cuda") + kv_lens = (num_past_tokens + seq_lens) // compress_ratio + + num_contexts = 2 + num_ctx_tokens = 7 # 3 + 4 new tokens total + + cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( + seq_lens, num_contexts, num_ctx_tokens, num_past_tokens, kv_lens, + compress_ratio) + + # Expected results: + # + # Seq 0: req has [past: 0,1] + [new: 2,3,4] = total 5 tokens, compressed to 2 KV tokens, global range [0:2] + # New Q token 0: [0, 1) + # New Q token 1: [0, 2) + # New Q token 2: [0, 2] + # + # Seq 1: req has [past: 0] + [new: 1,2,3,4] = total 5 tokens, compressed to 2 KV tokens, global range [2:4] + # New Q token 0: [2, 3) + # New Q token 1: [2, 3) + # New Q token 2: [2, 4) + # New Q token 3: [2, 4) + + expected_ks = torch.tensor([0, 0, 0, 2, 2, 2, 2], + dtype=torch.int32, + device="cuda") + expected_ke = torch.tensor([1, 2, 2, 3, 3, 4, 4], + dtype=torch.int32, + device="cuda") + + assert torch.equal(cu_seqlen_ks, expected_ks), \ + f"cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" + assert torch.equal(cu_seqlen_ke, expected_ke), \ + f"cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + + @pytest.mark.parametrize( "max_chunk_size,seq_lens,start_idx,expected_specs", [ @@ -1894,6 +2025,7 @@ def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper +@pytest.mark.parametrize("compress_ratio", [1, 4]) @pytest.mark.parametrize( "chunk_size,seq_lens_list,chunking_type", [ @@ -1909,7 +2041,8 @@ def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, ], "two_level"), # Mixed: request 2 needs Q-blocks ], ) -def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type): +def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, + compress_ratio): """ Tests for indexer chunked prefill: 1. Request-level chunking: Multiple small requests packed together @@ -1951,10 +2084,18 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type): total_tokens = seq_lens.sum().item() max_seq_len = seq_lens.max().item() - print(f"\n=== Test Config: {chunking_type} ===") + # Compute compressed KV token counts + kv_lens_compressed = seq_lens // compress_ratio + total_kv_tokens = kv_lens_compressed.sum().item() + + print( + f"\n=== Test Config: {chunking_type}, compress_ratio={compress_ratio} ===" + ) print(f" Batch: {batch_size}, Chunk size: {chunk_size}") print(f" Sequence lengths: {seq_lens_list}") - print(f" Total tokens: {total_tokens}, Max seq len: {max_seq_len}") + print( + f" Total tokens: {total_tokens}, Total KV tokens: {total_kv_tokens}, Max seq len: {max_seq_len}" + ) # Identify large requests for two-level chunking if chunking_type == "two_level": @@ -1990,7 +2131,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type): q = torch.randn((total_tokens, heads, head_dim), device="cuda", dtype=torch.bfloat16) - k = torch.randn((total_tokens, head_dim), + k = torch.randn((total_kv_tokens, head_dim), device="cuda", dtype=torch.bfloat16) weights = torch.randn((total_tokens, heads), @@ -2019,6 +2160,8 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type): num_ctx_tokens=total_tokens, num_tokens=total_tokens, indexer_max_chunk_size=chunk_size, + indexer_head_dim=head_dim, + compress_ratio=compress_ratio, ) Indexer.prepare(metadata_chunked) @@ -2057,6 +2200,8 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type): num_ctx_tokens=total_tokens, num_tokens=total_tokens, indexer_max_chunk_size=max_model_len, + indexer_head_dim=head_dim, + compress_ratio=compress_ratio, ) Indexer.prepare(metadata_baseline) @@ -2287,6 +2432,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, num_ctx_tokens=total_context_tokens, num_tokens=total_context_tokens, max_draft_tokens=next_n - 1, + indexer_head_dim=head_dim, ) Indexer.prepare(metadata_context) indexer._update_k_cache(k_context_fp8, k_context_scale, metadata_context) @@ -2320,6 +2466,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, num_ctx_tokens=0, num_tokens=num_gen_tokens, max_draft_tokens=next_n - 1, + indexer_head_dim=head_dim, ) Indexer.prepare(metadata_gen_write) indexer._update_k_cache(k_fp8, k_scale, metadata_gen_write) @@ -2336,7 +2483,8 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, 0, num_gen_tokens, max_model_len, - max_draft_tokens=next_n - 1) + max_draft_tokens=next_n - 1, + indexer_head_dim=head_dim) Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) @@ -2364,7 +2512,8 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, 0, num_gen_tokens, max_model_len, - max_draft_tokens=next_n - 1) + max_draft_tokens=next_n - 1, + indexer_head_dim=head_dim) Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) @@ -2390,7 +2539,8 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, num_gen_tokens, max_model_len, max_draft_tokens=next_n - 1, - enable_indexer_skip=True) + enable_indexer_skip=True, + indexer_head_dim=head_dim) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) @@ -2499,11 +2649,17 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k) # Test with custom CUDA kernel - metadata_custom = _create_mock_metadata(request_ids, batch_size, - batch_size, 0, seq_lens.clone(), + metadata_custom = _create_mock_metadata(request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), seq_lens.clone(), [0] * batch_size, - cache_manager, total_tokens, - total_tokens, chunk_size) + cache_manager, + total_tokens, + total_tokens, + chunk_size, + indexer_head_dim=head_dim) Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) @@ -2522,12 +2678,18 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, pytest.skip(f"Custom topk not available: {e}") # Test with PyTorch fallback - metadata_fallback = _create_mock_metadata(request_ids, batch_size, - batch_size, 0, seq_lens.clone(), + metadata_fallback = _create_mock_metadata(request_ids, + batch_size, + batch_size, + 0, seq_lens.clone(), - [0] * batch_size, cache_manager, - total_tokens, total_tokens, - chunk_size) + seq_lens.clone(), + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + chunk_size, + indexer_head_dim=head_dim) Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) @@ -2607,11 +2769,17 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k) # Test with custom CUDA kernel - metadata_custom = _create_mock_metadata(request_ids, batch_size, - batch_size, 0, seq_lens.clone(), + metadata_custom = _create_mock_metadata(request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), seq_lens.clone(), [0] * batch_size, - cache_manager, total_tokens, - total_tokens, max_model_len) + cache_manager, + total_tokens, + total_tokens, + max_model_len, + indexer_head_dim=head_dim) Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) @@ -2630,12 +2798,18 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, pytest.skip(f"Custom topk not available: {e}") # Test with PyTorch fallback - metadata_fallback = _create_mock_metadata(request_ids, batch_size, - batch_size, 0, seq_lens.clone(), + metadata_fallback = _create_mock_metadata(request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), seq_lens.clone(), - [0] * batch_size, cache_manager, - total_tokens, total_tokens, - max_model_len) + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + max_model_len, + indexer_head_dim=head_dim) Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) @@ -2661,7 +2835,8 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, total_tokens, total_tokens, max_model_len, - enable_indexer_skip=True) + enable_indexer_skip=True, + indexer_head_dim=head_dim) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) metadata_skip.indexer_prefill_chunks = None @@ -2772,7 +2947,8 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): total_tokens, total_tokens, indexer_max_chunk_size=32768, - enable_context_mla_with_cached_kv=True) + enable_context_mla_with_cached_kv=True, + indexer_head_dim=head_dim) Indexer.prepare(metadata) indexer._update_k_cache(k_fp8, k_scale, metadata) @@ -2810,7 +2986,8 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): total_tokens, indexer_max_chunk_size=32768, enable_context_mla_with_cached_kv=True, - enable_indexer_skip=True) + enable_indexer_skip=True, + indexer_head_dim=head_dim) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) topk_indices_skip = indexer.sparse_attn_indexer(metadata_skip, @@ -2846,8 +3023,9 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): # Check tokens with large windows (>= 2048) should have exactly 2048 valid indices print("\n=== Check: Large windows must have 2048 valid ===") - from tensorrt_llm._torch.attention_backend.sparse.dsa import \ - compute_cu_seqlen_kv_bounds_with_cache + from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + compute_cu_seqlen_kv_bounds_with_cache, + ) host_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device='cpu') host_cached = torch.tensor(cached_tokens, dtype=torch.int32, device='cpu') cu_ks, cu_ke = compute_cu_seqlen_kv_bounds_with_cache( @@ -2911,8 +3089,7 @@ def _make_mock_draft_manager(): def test_prepare_swaps_and_restore_recovers(self): """Test that prepare swaps KV manager and restore recovers original state.""" - from tensorrt_llm._torch.attention_backend.trtllm import \ - TrtllmAttentionMetadata + from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata meta = self._make_mock_metadata() mgr = self._make_mock_draft_manager() diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py new file mode 100644 index 000000000000..533446f257c0 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py @@ -0,0 +1,1009 @@ +""" +Tests for sparse MLA attention using explicit sparse indices. +""" + +import math +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import List, Optional + +import pytest +import torch +from utils.util import skip_pre_blackwell + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.interface import ( + AttentionInputType, + MLAParams, + PositionalEmbeddingParams, + RopeParams, +) +from tensorrt_llm._torch.attention_backend.sparse.dsa import DSACacheManager +from tensorrt_llm._torch.attention_backend.utils import get_attention_backend +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.model_config import ModelConfig +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.llmapi.llm_args import DeepSeekSparseAttentionConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + + +# 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 _rotate_k_pe_for_ctx( + k_pe: torch.Tensor, rope_cos_sin: torch.Tensor, sequence_lengths: List[int] +) -> torch.Tensor: + k_pe_ref_list = [] + total_tokens = 0 + for seq_len in sequence_lengths: + k_pe_seq = k_pe[total_tokens : total_tokens + seq_len].unsqueeze(-2) + cos, sin = rope_cos_sin[:seq_len].chunk(2, dim=-2) + k_pe_seq = k_pe_seq.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + k_pe_seq = ((k_pe_seq * cos) + (rotate_half(k_pe_seq) * sin)).to(dtype=k_pe_seq.dtype) + k_pe_seq = k_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + k_pe_ref_list.append(k_pe_seq) + total_tokens += seq_len + return torch.cat(k_pe_ref_list).squeeze(-2) + + +def _rotate_fused_q_for_ctx( + fused_q: torch.Tensor, + rope_cos_sin: torch.Tensor, + sequence_lengths: List[int], + num_heads: int, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> torch.Tensor: + fused_q = fused_q.clone() + fused_head_dim = kv_lora_rank + qk_rope_head_dim + total_tokens = 0 + for seq_len in sequence_lengths: + fused_q_seq = fused_q[total_tokens : total_tokens + seq_len].view( + seq_len, num_heads, fused_head_dim + ) + q_rope = fused_q_seq[..., -qk_rope_head_dim:] + cos, sin = rope_cos_sin[:seq_len].chunk(2, dim=-2) + q_rope = q_rope.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) + q_rope = ((q_rope * cos) + (rotate_half(q_rope) * sin)).to(dtype=fused_q.dtype) + q_rope = q_rope.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + fused_q_seq[..., -qk_rope_head_dim:] = q_rope + fused_q[total_tokens : total_tokens + seq_len] = fused_q_seq.view(seq_len, -1) + total_tokens += seq_len + return fused_q + + +def calculate_ref_result_ctx_sparse( + fused_q: torch.Tensor, + latent_cache: torch.Tensor, + sequence_lengths: List[int], + num_heads: int, + kv_lora_rank: int, + v_head_dim: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + q_scaling: float, + topk_indices: Optional[torch.Tensor] = None, +): + """ + Reference for sparse MLA context using fused Q and latent cache. + fused_q shape: (total_tokens, num_heads * (kv_lora_rank + qk_rope_head_dim)) + latent_cache shape: (total_tokens, kv_lora_rank + qk_rope_head_dim) + """ + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + bmm1_scale = 1 / (math.sqrt(qk_head_dim) * q_scaling) + fused_head_dim = kv_lora_rank + qk_rope_head_dim + ref_results = [] + total_tokens = 0 + for seq_len in sequence_lengths: + fused_q_seq = fused_q[total_tokens : total_tokens + seq_len].unflatten( + -1, [num_heads, fused_head_dim] + ) + fused_q_seq = fused_q_seq.transpose(0, 1) # (num_heads, seq_len, fused_head_dim) + + latent_seq = latent_cache[total_tokens : total_tokens + seq_len] + k_seq = latent_seq.unsqueeze(0) # (1, seq_len, fused_head_dim) + v_seq = latent_seq[..., :v_head_dim].unsqueeze(0) # (1, seq_len, v_head_dim) + + k_seq = repeat_kv(k_seq.unsqueeze(0), num_heads).squeeze(0) + v_seq = repeat_kv(v_seq.unsqueeze(0), num_heads).squeeze(0) + + if topk_indices is None: + attn_weights = torch.matmul(fused_q_seq, k_seq.transpose(1, 2)) * bmm1_scale + causal_mask = torch.triu( + torch.ones(seq_len, seq_len, device=fused_q.device, dtype=torch.bool), diagonal=1 + ) + attn_weights = attn_weights.masked_fill(causal_mask, float("-inf")) + attn_weights = torch.nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(fused_q.dtype) + attn_output = torch.matmul(attn_weights, v_seq) # (num_heads, seq_len, v_head_dim) + attn_output = ( + attn_output.transpose(0, 1).contiguous().view(seq_len, num_heads * v_head_dim) + ) + ref_results.append(attn_output) + else: + per_token_outputs = [] + token_rows = topk_indices[total_tokens : total_tokens + seq_len] + for token_idx in range(seq_len): + token_indices = token_rows[token_idx] + token_indices = token_indices[token_indices >= 0] + q_tok = fused_q_seq[:, token_idx, :] + k_sel = k_seq[:, token_indices, :] + v_sel = v_seq[:, token_indices, :] + attn_weights = ( + torch.matmul( + q_tok.unsqueeze(1), + k_sel.transpose(1, 2), + ) + * bmm1_scale + ) + attn_weights = torch.nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(fused_q.dtype) + attn_output = torch.matmul(attn_weights, v_sel) # (num_heads, 1, v_head_dim) + per_token_outputs.append( + attn_output.transpose(0, 1).contiguous().view(1, num_heads * v_head_dim) + ) + ref_results.append(torch.cat(per_token_outputs, dim=0)) + total_tokens += seq_len + return torch.cat(ref_results) + + +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, + v_head_dim: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + sequence_lengths: List[int], + q_scaling: float, + topk_indices: Optional[torch.Tensor] = None, +): + """ + 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) + q_pe_seq = q_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + k_pe_seq = k_pe_seq.unflatten(-1, [2, -1]).transpose(-2, -1).flatten(start_dim=-2) + 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[..., :v_head_dim].transpose(0, 1) + ) # (1, seq_len_kv, 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_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, v_head_dim) + + # Handle grouped-query attention + k = repeat_kv(k.unsqueeze(0), num_heads).squeeze(0) + v = repeat_kv(v.unsqueeze(0), num_heads).squeeze(0) + + seq_len_q = q.shape[1] + seq_len_kv = k.shape[1] + if topk_indices is None: + # 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. + 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, v_head_dim) + + # Reshape back to (seq_len_q, num_heads*v_head_dim) + attn_output = attn_output.transpose(0, 1).contiguous().view(-1, num_heads * v_head_dim) + ref_results.append(attn_output) + else: + per_token_outputs = [] + for qi in range(seq_len_q): + row = i * seq_len_q + qi + token_indices = topk_indices[row] + token_indices = token_indices[token_indices >= 0] + q_tok = q[:, qi, :] + k_sel = k[:, token_indices, :] + v_sel = v[:, token_indices, :] + attn_weights = torch.matmul( + q_tok.unsqueeze(1), + k_sel.transpose(1, 2), + ) / (q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim)) + attn_weights = torch.nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(q.dtype) + attn_output = torch.matmul(attn_weights, v_sel) # (num_heads, 1, v_head_dim) + per_token_outputs.append( + attn_output.transpose(0, 1).contiguous().view(1, num_heads * v_head_dim) + ) + ref_results.append(torch.cat(per_token_outputs, dim=0)) + + 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 = 1 + 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 = 512 + rope_append: bool = True + 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 = field( + default_factory=lambda: { + "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) + + +def _build_sparse_topk_indices_context( + seq_lens: List[int], topk: int, device: torch.device +) -> torch.Tensor: + total_tokens = sum(seq_lens) + topk_indices = torch.full((total_tokens, topk), -1, dtype=torch.int32, device=device) + token_offset = 0 + for seq_len in seq_lens: + for token_idx in range(seq_len): + max_index = token_idx + valid_len = min(max_index + 1, topk) + indices = torch.randperm(max_index + 1, device=device)[:valid_len] + indices, _ = torch.sort(indices) + topk_indices[token_offset + token_idx, :valid_len] = indices.to(torch.int32) + token_offset += seq_len + return topk_indices + + +def _build_sparse_topk_indices_generation( + cached_lens: List[int], seq_len_q: int, topk: int, device: torch.device +) -> torch.Tensor: + total_tokens = len(cached_lens) * seq_len_q + topk_indices = torch.full((total_tokens, topk), -1, dtype=torch.int32, device=device) + row = 0 + for cached_len in cached_lens: + for q_idx in range(seq_len_q): + max_index = cached_len + q_idx + valid_len = min(max_index + 1, topk) + indices = torch.randperm(max_index + 1, device=device)[:valid_len] + indices, _ = torch.sort(indices) + topk_indices[row, :valid_len] = indices.to(torch.int32) + row += 1 + return topk_indices + + +def _allocate_kv_cache_for_generation(kv_cache_manager, request_ids, num_tokens: int): + for request_id in request_ids: + for _ in range(num_tokens): + kv_cache_manager.impl.add_token(request_id) + if hasattr(kv_cache_manager, "indexer_k_cache_manager"): + kv_cache_manager.indexer_k_cache_manager.add_tokens(request_id, 1) + + +# Define test data +context_sequence_lengths = [[10], [3000, 3100], [508, 4399, 9981]] +# Use MTP by default if seqlen_q > 1. +generation_seq_len_q = [1, 4] +num_generation_steps = [2] + +tokens_per_block = 64 + +kv_cache_dtype_list = [torch.bfloat16] +# DSA only supports rope_append=True +rope_append_values = [True] +scenarios = [ + Scenario( + kv_cache_dtype=kv_cache_dtype, + num_layers=num_layers, + kv_cache_tokens_per_block=tokens_per_block, + rope_append=rope_append, + ) + for kv_cache_dtype in kv_cache_dtype_list + for num_layers in [1] + for rope_append in rope_append_values +] + +accuracy_dict = { + torch.bfloat16: (0.1, 0.01), + torch.float8_e4m3fn: (0.1, 0.01), +} + +SPARSE_TOPK = 2048 + + +# Convert parameterized tests to pytest parametrize +@skip_pre_blackwell +@pytest.mark.skip_less_device_memory(80000) +@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_sparse_attention_mla( + scenario: Scenario, + context_sequence_lengths: List[int], + generation_seq_len_q: int, + num_generation_steps: int, +): + """Test sparse 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 + 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_append = scenario.rope_append + if rope_append is False: + print("rope_append is False, setting num_heads to 64") + num_heads = 64 + kv_lora_rank = scenario.kv_lora_rank + 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 + + assert SPARSE_TOPK % 128 == 0 + + 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_append, + rope_config, + kv_cache_tokens_per_block, + device, + dtype, + kv_cache_dtype, + context_sequence_lengths, + generation_seq_len_q, + num_generation_steps, + ) + + +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_append, + rope_config, + kv_cache_tokens_per_block, + device, + dtype, + kv_cache_dtype, + context_sequence_lengths, + generation_seq_len_q, + num_generation_steps, +): + sparse_config = DeepSeekSparseAttentionConfig( + index_n_heads=64, + index_head_dim=128, + index_topk=SPARSE_TOPK, + skip_indexer_for_short_seqs=False, + ) + AttentionCls = get_attention_backend(backend_name, sparse_config) + # When rope_append is False, [448: 512) are used for qk_rope_head_dim + kv_lora_rank = kv_lora_rank - qk_rope_head_dim if not rope_append else kv_lora_rank + head_dim = kv_lora_rank + 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, kv_lora_rank], # sparse MLA uses absorption mode + dtype=dtype, + device=device, + ).uniform_(-1, 1) + for ctx_len in context_sequence_lengths + ] + ) + ctx_q_pe = torch.cat( + [ + torch.empty( + [ctx_len, num_heads, qk_rope_head_dim], + dtype=dtype, + device=device, + ).uniform_(-1, 1) + for ctx_len in context_sequence_lengths + ] + ) + ctx_fused_q = torch.cat([ctx_q, ctx_q_pe], dim=-1).view(-1, num_heads * 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_q_list = [ + torch.cat( + [ + torch.empty( + [generation_seq_len_q, num_heads, kv_lora_rank], + 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) + ] + gen_fused_q_list = [ + torch.cat([gen_q_list[i], gen_q_pe_list[i]], dim=-1).view(-1, num_heads * head_dim) + for i in range(num_generation_steps) + ] + + inputs = { + "ctx_compressed_kv": ctx_compressed_kv, + "ctx_k_pe": ctx_k_pe, + "ctx_q_pe": ctx_q_pe, + "ctx_fused_q": ctx_fused_q, + "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, + rope_append=rope_append, + 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=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, + sparse_attention_config=sparse_config, + ) + for layer_idx in range(num_layers) + ] + gen_layers = [ + AttentionCls( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=1, + quant_config=quant_config, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=mla_params, + sparse_attention_config=sparse_config, + ) + 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_context_sequence_length = max(context_sequence_lengths) + max_num_contexts = len(context_sequence_lengths) + 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 + ) + + pretrained_config = SimpleNamespace( + rms_norm_eps=1e-6, + ) + model_config = ModelConfig( + mapping=mapping, + sparse_attention_config=sparse_config, + pretrained_config=pretrained_config, + ) + kv_cache_manager = DSACacheManager( + 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=head_dim, + tokens_per_block=kv_cache_tokens_per_block, + max_seq_len=max_context_sequence_length + (num_generation_steps + 1) * generation_seq_len_q, + max_batch_size=max_num_contexts, + mapping=mapping, + dtype=str_dtype_to_binding(torch_dtype_to_str(kv_cache_dtype)), + sparse_attn_config=sparse_config, + model_config=model_config, + ) + request_ids = list(range(max_num_contexts)) + kv_cache_manager.add_dummy_requests(request_ids, context_sequence_lengths) + + ctx_seq_lens = torch.tensor(context_sequence_lengths, dtype=torch.int) + total_ctx_tokens = sum(context_sequence_lengths) + attn_metadata = AttentionCls.Metadata( + seq_lens=ctx_seq_lens, + request_ids=request_ids, + max_num_requests=max_num_contexts, + num_contexts=max_num_contexts, + prompt_lens=context_sequence_lengths, + max_num_tokens=total_ctx_tokens, + 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, + sparse_attention_config=sparse_config, + ) + 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: + _allocate_kv_cache_for_generation(kv_cache_manager, request_ids, generation_seq_len_q) + gen_seq_lens = torch.tensor([generation_seq_len_q] * max_num_contexts, dtype=torch.int) + total_gen_tokens = max_num_contexts * generation_seq_len_q + attn_metadata = AttentionCls.Metadata( + seq_lens=gen_seq_lens, + request_ids=request_ids, + max_num_requests=max_num_contexts, + num_contexts=0, + prompt_lens=context_sequence_lengths, + max_num_tokens=total_gen_tokens, + 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), + sparse_attention_config=sparse_config, + ) + attn_metadata.prepare() + for layer_idx in range(num_layers): + print( + f"--------------------------------step {step} layer {layer_idx} start--------------------------------" + ) + if step == 0: + fused_q = inputs_per_layer[layer_idx]["ctx_fused_q"] + 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_pe = inputs_per_layer[layer_idx]["ctx_q_pe"] + topk_indices = _build_sparse_topk_indices_context( + context_sequence_lengths, SPARSE_TOPK, device + ) + result = ctx_layers[layer_idx].forward( + fused_q.clone(), + None, + None, + attn_metadata, + attention_input_type=AttentionInputType.context_only, + latent_cache=latent_cache, + q_pe=q_pe, + topk_indices=topk_indices, + is_generation=False, + ) + k_pe_ref = _rotate_k_pe_for_ctx(k_pe, rope_cos_sin, context_sequence_lengths) + latent_cache_ref = torch.cat([compressed_kv, k_pe_ref], dim=-1) + fused_q_rot = _rotate_fused_q_for_ctx( + fused_q, + rope_cos_sin, + context_sequence_lengths, + num_heads, + kv_lora_rank, + qk_rope_head_dim, + ) + ref_result = calculate_ref_result_ctx_sparse( + fused_q_rot, + latent_cache_ref, + context_sequence_lengths, + num_heads, + kv_lora_rank, + v_head_dim, + qk_nope_head_dim, + qk_rope_head_dim, + q_scaling, + topk_indices=topk_indices, + ) + latent_cache_ref_all_list[layer_idx] = latent_cache_ref + 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) + + num_tokens = fused_q.size(0) + num_seqs = attn_metadata.kv_lens_cuda_runtime.size(0) + cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=fused_q.device) + cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=fused_q.device) + fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=fused_q.device) + has_fp8_kv_cache = ( + gen_layers[layer_idx].has_fp8_kv_cache + if hasattr(gen_layers[layer_idx], "has_fp8_kv_cache") + else False + ) + + if has_fp8_kv_cache: + mla_bmm1_scale = torch.empty(2, dtype=torch.float32, device=fused_q.device) + mla_bmm2_scale = torch.empty(1, dtype=torch.float32, device=fused_q.device) + quant_q_buffer = torch.empty( + num_tokens, num_heads * head_dim, dtype=torch.uint8, device=fused_q.device + ) + else: + mla_bmm1_scale = None + mla_bmm2_scale = None + quant_q_buffer = None + + gen_layers[layer_idx].mla_rope_generation( + fused_q, + q_pe, + latent_cache, + attn_metadata, + cu_q_seqlens, + cu_kv_seqlens, + fmha_scheduler_counter, + mla_bmm1_scale, + mla_bmm2_scale, + quant_q_buffer, + ) + + cached_lens = [ + ctx_len + (step - 1) * generation_seq_len_q + for ctx_len in context_sequence_lengths + ] + topk_indices = _build_sparse_topk_indices_generation( + cached_lens, generation_seq_len_q, SPARSE_TOPK, device + ) + + 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, + cu_q_seqlens=cu_q_seqlens, + cu_kv_seqlens=cu_kv_seqlens, + fmha_scheduler_counter=fmha_scheduler_counter, + mla_bmm1_scale=mla_bmm1_scale, + mla_bmm2_scale=mla_bmm2_scale, + quant_q_buffer=quant_q_buffer, + topk_indices=topk_indices, + is_generation=True, + ) + 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, + v_head_dim, + 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, + topk_indices=topk_indices, + ) + 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 sparse MLA in {backend_name} backend don't match reference implementation \ + at layer {layer_idx} in step {step}" + ) + + print( + f"Test for sparse 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 sparse MLA in {backend_name} backend passed") + kv_cache_manager.shutdown() + + +if __name__ == "__main__": + test_sparse_attention_mla( + scenario=scenarios[0], + context_sequence_lengths=context_sequence_lengths[0], + generation_seq_len_q=generation_seq_len_q[0], + num_generation_steps=num_generation_steps[0], + ) diff --git a/tests/unittest/_torch/attention/sparse/test_short_seq_mha.py b/tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py similarity index 99% rename from tests/unittest/_torch/attention/sparse/test_short_seq_mha.py rename to tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py index b9c339f29dc6..e118e5da2f83 100644 --- a/tests/unittest/_torch/attention/sparse/test_short_seq_mha.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py @@ -381,7 +381,7 @@ def _run_forward( ): """Run forward_context_dsa on cloned inputs and return the output tensor.""" output = torch.empty(q.shape[0], NUM_HEADS * V_HEAD_DIM, dtype=q.dtype, device=q.device) - mla.forward_context_dsa( + mla.forward_context_sparse_mla( q=q.clone(), compressed_kv=compressed_kv.clone(), k_pe=k_pe.clone(), diff --git a/tests/unittest/_torch/attention/sparse/kernel/__init__.py b/tests/unittest/_torch/attention/sparse/kernel/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unittest/_torch/attention/sparse/test_flash_mla.py b/tests/unittest/_torch/attention/sparse/kernel/test_flash_mla.py similarity index 100% rename from tests/unittest/_torch/attention/sparse/test_flash_mla.py rename to tests/unittest/_torch/attention/sparse/kernel/test_flash_mla.py diff --git a/tests/unittest/_torch/attention/sparse/test_triton_bmm.py b/tests/unittest/_torch/attention/sparse/kernel/test_triton.py similarity index 66% rename from tests/unittest/_torch/attention/sparse/test_triton_bmm.py rename to tests/unittest/_torch/attention/sparse/kernel/test_triton.py index 93695fb26d63..089513644ee8 100644 --- a/tests/unittest/_torch/attention/sparse/test_triton_bmm.py +++ b/tests/unittest/_torch/attention/sparse/kernel/test_triton.py @@ -6,6 +6,7 @@ from tensorrt_llm._torch.attention_backend.sparse.kernel import ( triton_bmm, triton_rocket_paged_kt_cache_bmm, + triton_topk, ) @@ -217,6 +218,9 @@ def pytorch_reference_paged_kt_cache_bmm( return scores +# ==================== Triton BMM Tests ==================== + + @pytest.mark.parametrize( "batch_size,q_len_per_seq,k_lens,num_q_heads,num_kv_heads,head_dim,causal", [ @@ -372,6 +376,212 @@ def test_triton_rocket_paged_kt_cache_bmm( assert max_diff < 0.05, f"Max difference {max_diff} exceeds threshold" +# ==================== Triton TopK Tests ==================== + + +def pytorch_reference_topk( + input_tensor: torch.Tensor, + kv_offsets: torch.Tensor, + kv_lens: torch.Tensor, + topk: int, +) -> torch.Tensor: + """ + Args: + input_tensor: Input scores [num_kv_heads, sum(kv_lens)] + kv_offsets: KV offsets [batch_size + 1] + kv_lens: KV lengths [batch_size] + topk: TopK parameter + + Returns: + output_indices: Padded indices [num_kv_heads, batch_size, topk] + """ + num_kv_heads = input_tensor.shape[0] + batch_size = len(kv_lens) + device = input_tensor.device + + # Compute max sequence length for padding + max_seq_len = kv_lens.max().item() + + # Ensure padding size >= topk + pad_size = max(max_seq_len, topk) + + # Create padded tensor [num_kv_heads, batch_size, pad_size] + padded_tensor = torch.full( + (num_kv_heads, batch_size, pad_size), float("-inf"), dtype=input_tensor.dtype, device=device + ) + + # Fill in actual values + for batch_idx in range(batch_size): + start = kv_offsets[batch_idx].item() + end = kv_offsets[batch_idx + 1].item() + seq_len = kv_lens[batch_idx].item() + + for head_idx in range(num_kv_heads): + padded_tensor[head_idx, batch_idx, :seq_len] = input_tensor[head_idx, start:end] + + # Perform batch topk: [num_kv_heads, batch_size, pad_size] -> [num_kv_heads, batch_size, topk] + topk_values, topk_indices = torch.topk( + padded_tensor, + k=topk, + dim=-1, + largest=True, + ) + + # Mask out invalid indices based on each batch's seq_len + seq_lens_expanded = kv_lens.to(device).unsqueeze(0).unsqueeze(-1) # [1, batch_size, 1] + # topk_indices: [num_kv_heads, batch_size, topk] + mask = topk_indices >= seq_lens_expanded + topk_indices.masked_fill_(mask, -1) + + return topk_indices + + +def triton_topk_wrapper( + input_tensor: torch.Tensor, + kv_offsets: torch.Tensor, + kv_lens: torch.Tensor, + topk: int, +) -> torch.Tensor: + """ + Args: + input_tensor: Input scores [num_kv_heads, sum(kv_lens)] + kv_offsets: KV offsets [batch_size + 1] + kv_lens: KV lengths [batch_size] + topk: TopK parameter + + Returns: + output_indices: Padded indices [num_kv_heads, batch_size, topk] + """ + num_kv_heads = input_tensor.shape[0] + batch_size = len(kv_lens) + device = input_tensor.device + + sparse_lens = torch.tensor( + [min(topk, seq_len.item()) for seq_len in kv_lens], dtype=torch.int32, device=device + ) + sparse_offsets = torch.cat( + [torch.zeros(1, dtype=torch.int32, device=device), torch.cumsum(sparse_lens, dim=0)] + ).to(device) + total_sparse_indices = sparse_offsets[-1].item() + + output_indices_flat = triton_topk( + input_tensor, kv_offsets, sparse_offsets, total_sparse_indices, topk + ) + + # Convert flat format to padded format [num_kv_heads, batch_size, topk] + output_indices_padded = torch.full( + (num_kv_heads, batch_size, topk), -1, dtype=torch.int32, device=device + ) + + for batch_idx in range(batch_size): + start = sparse_offsets[batch_idx].item() + end = sparse_offsets[batch_idx + 1].item() + actual_len = end - start + + for head_idx in range(num_kv_heads): + output_indices_padded[head_idx, batch_idx, :actual_len] = output_indices_flat[ + head_idx, start:end + ] + + return output_indices_padded + + +def compute_overlap_ratio( + triton_indices: torch.Tensor, + reference_indices: torch.Tensor, + kv_lens: torch.Tensor, +) -> float: + """ + Args: + triton_indices: Triton topk results [num_kv_heads, batch_size, topk] + reference_indices: Reference topk results [num_kv_heads, batch_size, topk] + kv_lens: KV lengths [batch_size] + + Returns: + Average overlap ratio across all batches and heads + """ + num_kv_heads = triton_indices.shape[0] + batch_size = triton_indices.shape[1] + + overlap_ratios = [] + + # Compare batch by batch + for batch_idx in range(batch_size): + for head_idx in range(num_kv_heads): + # Extract indices for this batch and head + triton_batch = triton_indices[head_idx, batch_idx, :].cpu().tolist() + reference_batch = reference_indices[head_idx, batch_idx, :].cpu().tolist() + + # Filter out -1 (invalid/padding indices) + triton_set = set([x for x in triton_batch if x >= 0]) + reference_set = set([x for x in reference_batch if x >= 0]) + + if len(reference_set) > 0: + overlap = len(triton_set & reference_set) + overlap_ratio = overlap / len(reference_set) + overlap_ratios.append(overlap_ratio) + + if len(overlap_ratios) == 0: + return 1.0 + + return sum(overlap_ratios) / len(overlap_ratios) + + +@pytest.mark.parametrize( + "batch_size,seq_lens,num_kv_heads,topk", + [ + # Single batch, seq_len > topk + (1, [3000], 1, 2048), + # Single batch, seq_len < topk + (1, [1000], 8, 2048), + # Multiple batches, mixed seq_len (some < topk, some > topk) + (6, [50, 150, 80, 300, 100, 256], 8, 128), + (6, [1000, 2500, 1500, 3000, 1800, 4000], 8, 2048), + ], +) +def test_topk_kernel(batch_size, seq_lens, num_kv_heads, topk): + device = torch.device("cuda") + dtype = torch.float32 + + kv_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + + kv_offsets = torch.cat( + [torch.zeros(1, dtype=torch.int32, device=device), torch.cumsum(kv_lens, dim=0)] + ).to(device) + + total_tokens = kv_offsets[-1].item() + + input_tensor = torch.randn((num_kv_heads, total_tokens), dtype=dtype, device=device) + + triton_output = triton_topk_wrapper( + input_tensor=input_tensor, + kv_offsets=kv_offsets, + kv_lens=kv_lens, + topk=topk, + ) + + reference_output = pytorch_reference_topk( + input_tensor=input_tensor, + kv_offsets=kv_offsets, + kv_lens=kv_lens, + topk=topk, + ) + + overlap_ratio = compute_overlap_ratio( + triton_output, + reference_output, + kv_lens, + ) + + min_threshold = 0.99 + + print(f"overlap_ratio: {overlap_ratio}") + + assert overlap_ratio >= min_threshold, ( + f"Overlap ratio {overlap_ratio:.4f} is too low (< {min_threshold})" + ) + + if __name__ == "__main__": print("\n" + "=" * 80) print("Testing Triton BMM Kernel") @@ -398,6 +608,21 @@ def test_triton_rocket_paged_kt_cache_bmm( print("\n--- Multiple batches, different kv_len ---") test_triton_rocket_paged_kt_cache_bmm(3, [100, 200, 150], 8, 4, 128, 3, 64) + print("\n" + "=" * 80) + print("Testing Triton TopK Kernel") + print("=" * 80) + + # Single batch tests + print("\n--- Single batch, seq_len > topk ---") + test_topk_kernel(1, [3000], 8, 2048) + + print("\n--- Single batch, seq_len < topk ---") + test_topk_kernel(1, [1000], 8, 2048) + + print("\n--- Multiple batches, mixed seq_len ---") + test_topk_kernel(6, [50, 150, 80, 300, 100, 256], 8, 128) + test_topk_kernel(6, [1000, 2500, 1500, 3000, 1800, 4000], 8, 2048) + print("\n" + "=" * 80) print("All tests passed!") print("=" * 80) diff --git a/tests/unittest/_torch/attention/sparse/rocketkv/__init__.py b/tests/unittest/_torch/attention/sparse/rocketkv/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/unittest/_torch/attention/sparse/test_rocketkv.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py similarity index 99% rename from tests/unittest/_torch/attention/sparse/test_rocketkv.py rename to tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py index e6bfcbb55aa5..f54ec8e0b80a 100644 --- a/tests/unittest/_torch/attention/sparse/test_rocketkv.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py @@ -60,8 +60,8 @@ def test_model(backend, model_name, attention_backend): inputs, references = [], [] current_file = os.path.abspath(__file__) - current_dir = os.path.dirname(os.path.dirname( - os.path.dirname(current_file))) + current_dir = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) input_file = f'{current_dir}/multi_gpu/NIAH_simple_data.jsonl' with open(input_file, 'r') as f: for line in f: diff --git a/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py b/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py index 944986dc2adb..ee727a285710 100644 --- a/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py +++ b/tests/unittest/_torch/attention/sparse/test_dsa_fp4_indexer.py @@ -41,7 +41,7 @@ "check that the correct DeepGEMM version is installed" ) -from test_dsa_indexer import _create_mock_metadata, create_dsa_cache_manager +from .dsa.test_dsa_indexer import _create_mock_metadata, create_dsa_cache_manager from utils.util import skip_pre_blackwell diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index ffd9896fb53a..35fb9d5e1352 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -12,18 +12,25 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.interface import ( - PositionalEmbeddingParams, RopeParams) + AttentionBackend, PositionalEmbeddingParams, RopeParams) from tensorrt_llm._torch.attention_backend.sparse.dsa import DSACacheManager +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + DeepseekV4CacheManager from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.attention import MLA from tensorrt_llm._utils import (get_sm_version, 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.llmapi.llm_args import DeepSeekSparseAttentionConfig +from tensorrt_llm.llmapi.llm_args import (DeepSeekSparseAttentionConfig, + KvCacheConfig, + DeepSeekV4SparseAttentionConfig) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +from .deepseek_v4.test_compressor_module import RefCompressor try: HAS_FLASH_MLA = True @@ -42,6 +49,27 @@ def batch_size(self): return len(self.seq_lens) +@dataclass +class DummyPretrainedConfig: + """Dummy pretrained configuration for testing.""" + num_heads: int + q_lora_rank: int + kv_lora_rank: int + qk_nope_head_dim: int + qk_rope_head_dim: int + v_head_dim: int + qk_head_dim: int + head_size: int + topk_tokens: int + hidden_size: int + max_position_embeddings: int + num_groups: int + o_lora_rank: int + num_layers: int + vocab_size: int + compress_ratios: int + + # Unified batch specifications covering all scenarios # seq_lens = total KV cache length (for decode/mixed) or current sequence length (for pure prefill) # query_lens = number of new query tokens per request @@ -89,12 +117,79 @@ def batch_size(self): } +def apply_rotary_emb(x: torch.Tensor, + freqs_cis: torch.Tensor, + inverse: bool = False) -> torch.Tensor: + """Apply rotary positional embeddings for DeepSeek-V4.""" + y = x + x = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if x.ndim == 3: + freqs_cis = freqs_cis.view(1, x.size(1), x.size(-1)) + else: + freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) + x = torch.view_as_real(x * freqs_cis).flatten(-2) + y.copy_(x) + return y + + +def precompute_freqs_cis(dim, seqlen, original_seq_len, base, factor, beta_fast, + beta_slow) -> torch.Tensor: + """Precompute rotary embeddings for DeepSeek-V4.""" + + def find_correction_dim(num_rotations, dim, base, max_seq_len): + return dim * math.log( + max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): + low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) + high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - + min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + freqs = 1.0 / (base**(torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + if seqlen > original_seq_len: + low, high = find_correction_range(beta_fast, beta_slow, dim, base, + original_seq_len) + smooth = 1 - linear_ramp_factor(low, high, dim // 2) + freqs = freqs / factor * (1 - smooth) + freqs * smooth + + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def _calc_diff(x: torch.Tensor, y: torch.Tensor): + """Return a global difference metric for unit tests. + + When comparing the fp8 results with the bf16 reference, there are + noticeable differences, causing ``torch.testing.assert_close`` to fail. + Instead of checking every element, we compute a cosine-style similarity + over the whole tensor and report ``1 - sim``. + """ + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return 1 - sim.item() + + def apply_rotary_embedding( x: torch.Tensor, # [seq_len, ..., rope_dim] cos_sin: torch.Tensor, # [seq_len, 2, rope_dim] + inverse: bool = False, # If True, apply inverse RoPE ) -> torch.Tensor: """ - Apply rotary position embedding. + Apply rotary position embedding for DSA. + For deepseek_v4, supports inverse rotation by negating sin. """ def rotate_half(x): @@ -107,6 +202,10 @@ def rotate_half(x): cos = cos.squeeze(1) sin = sin.squeeze(1) + # For inverse RoPE, negate sin + if inverse: + sin = -sin + # Convert to interleaved format x_interleaved = x.unflatten(-1, [-1, 2]).transpose(-2, -1).flatten(start_dim=-2) @@ -191,7 +290,7 @@ def calculate_reference_output_generation(q_c, kv_c, k_pe, W_UK, W_UV, num_heads, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, softmax_scale, device): - """Reference for generation (rotated inputs, no RoPE application).""" + """Reference for generation using unrotated inputs and explicit RoPE.""" results, kv_offset = [], 0 for kv_len in kv_cache_lens: q_seq = q_c[len(results):len(results) + @@ -200,15 +299,11 @@ def calculate_reference_output_generation(q_c, kv_c, k_pe, W_UK, W_UV, k_pe_seq = k_pe[kv_offset:kv_offset + kv_len] q_nope, q_pe = q_seq.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1) - # SM90: apply RoPE to q_pe, k_pe_seq - # SM100+: use unrotated q_pe, k_pe_seq - if get_sm_version() >= 100: - cos_sin_q = rope_cos_sin[kv_len - q_seq.shape[0]:kv_len] - cos_sin_pe = rope_cos_sin[:kv_len] - q_pe = apply_rotary_embedding(q_pe, cos_sin_q) - k_pe_rot = apply_rotary_embedding(k_pe_seq, cos_sin_pe) - else: - k_pe_rot = k_pe_seq + # Apply RoPE + cos_sin_q = rope_cos_sin[kv_len - q_seq.shape[0]:kv_len] + cos_sin_pe = rope_cos_sin[:kv_len] + q_pe = apply_rotary_embedding(q_pe, cos_sin_q) + k_pe_rot = apply_rotary_embedding(k_pe_seq, cos_sin_pe) ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK) q_mqa = torch.cat([ql_nope, q_pe], dim=-1) @@ -284,16 +379,882 @@ def extract_kv_slices(indices, start_offset=0): return torch.cat(ref_results, dim=0) +def calculate_reference_output_deepseek_v4_prefill(hidden_states, + indices, + seq_lens, + q, + kv_data, + freqs_cis, + ref_compressor, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + softmax_scale, + device, + o_a_proj, + o_b_proj_weight, + n_local_groups, + window_size=512, + compress_ratio=4): + """Reference for deepseek_v4 prefill. + + Implements attention over: + 1. Sliding window: last window_size tokens + 2. Compressed KV: all compressed tokens (assuming all are selected) + """ + results, offset = [], 0 + device = q.device + + for req_idx in indices: + seq_len = seq_lens[req_idx] + q_seq = q[offset:offset + seq_len] # [seq_len, num_heads, qk_head_dim] + hidden_states_seq = hidden_states[offset:offset + seq_len] + kv_seq = kv_data["latent_kv"][req_idx] # [seq_len, head_dim] + + # Apply RoPE to Q/KV + q_seq = q_seq.unsqueeze(0) + kv_seq = kv_seq.unsqueeze(0) + apply_rotary_emb(q_seq[..., -qk_rope_head_dim:], freqs_cis[:seq_len]) + apply_rotary_emb(kv_seq[..., -qk_rope_head_dim:], freqs_cis[:seq_len]) + + # Compressor + if compress_ratio > 1: + kv_cache = torch.zeros(1, + seq_len // compress_ratio, + ref_compressor.head_dim, + device=device) + ref_compressor.kv_cache = kv_cache + compressed_kv = ref_compressor(hidden_states_seq.unsqueeze(0), + start_pos=0, + freqs_cis=freqs_cis[:seq_len]) + num_compressed = compressed_kv.shape[ + 1] if compressed_kv is not None else 0 + if num_compressed > 0: + # RefCompressor returns fp32 (matches the kernel's + # fp32-throughout postprocess); cast to the attention dtype + # before concatenation. + kv_combined = torch.cat( + [kv_seq, compressed_kv.to(kv_seq.dtype)], + dim=1).squeeze(0) + else: + kv_combined = kv_seq.squeeze(0) + num_compressed = 0 + else: + kv_combined = kv_seq.squeeze(0) + num_compressed = 0 + + # Create attention mask: [seq_len, seq_len + num_compressed] + # KV structure: [0:seq_len] = window_kv, [seq_len:seq_len+num_compressed] = compressed_kv + # Each query at position i can attend to: + # 1. Window tokens: positions [max(0, i - window_size + 1), i] in range [0, seq_len) + # 2. Compressed tokens: positions [seq_len, seq_len + i//compress_ratio) representing original tokens before position i + attn_mask = torch.full((seq_len, seq_len + num_compressed), + float('-inf'), + device=device, + dtype=q_seq.dtype) + + for i in range(seq_len): + # Window KV: query at position i can attend to positions [max(0, i - window_size + 1), i] + window_start = max(0, i - window_size + 1) + attn_mask[i, window_start:i + 1] = 0 + + # Compressed KV: production uses (pos+1)//compress_ratio valid compressed tokens + if num_compressed > 0: + max_compressed_idx = (i + 1) // compress_ratio + if max_compressed_idx > 0: + attn_mask[i, seq_len:seq_len + max_compressed_idx] = 0 + + # Expand KV for multi-head (deepseek_v4 uses MQA - single KV head for all Q heads) + k_combined = kv_combined.unsqueeze(1).expand(-1, num_heads, -1) + # Use full kv_combined as V (not just first v_head_dim dimensions) to match reference + v_combined = kv_combined.unsqueeze(1).expand(-1, num_heads, -1) + + # Prepare for attention: [1, num_heads, seq_len, head_dim] + # q: [1, num_heads, seq_len, qk_head_dim] + # k: [1, num_heads, seq_len + num_compressed, qk_head_dim] + # v: [1, num_heads, seq_len + num_compressed, qk_head_dim] + # attn_mask: [1, 1, seq_len, seq_len + num_compressed] + q_in = q_seq.transpose(1, 2) + k_in = k_combined.unsqueeze(0).transpose(1, 2) + v_in = v_combined.unsqueeze(0).transpose(1, 2) + attn_mask = attn_mask.unsqueeze(0).unsqueeze(0) + + # Compute attention with custom mask, attn_out: # [seq_len, num_heads, qk_head_dim] + attn_out = torch.nn.functional.scaled_dot_product_attention( + q_in, k_in, v_in, attn_mask=attn_mask, + scale=softmax_scale).transpose(1, 2).squeeze(0) + + results.append(attn_out) + offset += seq_len + + return torch.cat(results, dim=0).flatten(1) + + +def calculate_reference_output_deepseek_v4_generation(hidden_states, + gen_indices, + seq_lens, + q, + kv_data, + freqs_cis, + ref_compressor, + num_contexts, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + softmax_scale, + device, + o_a_proj, + o_b_proj_weight, + n_local_groups, + window_size=512, + compress_ratio=4): + """Reference for deepseek_v4 generation with cached window_kv and compressed_kv. + + For generation, we have: + - 1 new query token and its latent KV + - Cached window KV (last window_size tokens) + - Cached compressed KV (compressed tokens from earlier in sequence) + - Compressor state (kv_state, score_state) that needs to be restored + """ + results = [] + qk_nope_head_dim + qk_rope_head_dim + + for gen_req_idx in gen_indices: + req_idx = gen_req_idx - num_contexts + # Total sequence length including cache + new token + seq_len = seq_lens[gen_req_idx] + start_pos = seq_len - 1 # Position of the new token + + q_seq = q[req_idx].unsqueeze(0) # [1, num_heads, qk_head_dim] + kv_seq = kv_data["latent_kv"][gen_req_idx] # [1, head_dim] + hidden_states_seq = hidden_states[req_idx].unsqueeze(0) # [1, head_dim] + + # Get cached data + # window_kv: [window_size, head_dim] + # compressed_kv: [num_compressed, head_dim] or None + # kv_state: Compressor internal KV state + # score_state: Compressor internal score state + window_kv = kv_data["window_kv"][gen_req_idx] + compressed_kv = kv_data["compressed_kv"][gen_req_idx] + # The compressed_kv cache may have been populated from a fp32 + # RefCompressor return; cast back to the attention dtype. + if compressed_kv is not None: + compressed_kv = compressed_kv.to(kv_seq.dtype) + kv_state = kv_data["kv_state"][gen_req_idx] + score_state = kv_data["score_state"][gen_req_idx] + + # Apply RoPE to query at current position + q_seq = q_seq.unsqueeze(0) + kv_seq = kv_seq.unsqueeze(0) + apply_rotary_emb(q_seq[..., -qk_rope_head_dim:], + freqs_cis[start_pos:start_pos + 1]) + apply_rotary_emb(kv_seq[..., -qk_rope_head_dim:], + freqs_cis[start_pos:start_pos + 1]) + + # Restore compressor state and run compression on new token + if compress_ratio > 1: + # Restore compressor's internal state (kv_state, score_state) + ref_compressor.kv_state = kv_state.clone() + ref_compressor.score_state = score_state.clone() + + # Restore compressor's kv_cache with previously compressed tokens + kv_cache = torch.zeros(1, + seq_len // compress_ratio, + ref_compressor.head_dim, + device=device) + ref_compressor.kv_cache = kv_cache + if compressed_kv is not None: + num_compressed = compressed_kv.shape[0] + ref_compressor.kv_cache[0, :num_compressed] = compressed_kv + else: + num_compressed = 0 + + # Run compressor on new token (returns compressed KV only if compression happens at this step) + new_compressed_kv = ref_compressor( + hidden_states_seq.unsqueeze(0), # [1, 1, head_dim] + start_pos=start_pos, + freqs_cis=freqs_cis[start_pos:start_pos + 1]) + + # Update compressed_kv if a new compressed token was produced. + # RefCompressor returns fp32; cast back to the attention dtype. + if new_compressed_kv is not None: + new_compressed_kv = new_compressed_kv.to(kv_seq.dtype) + if compressed_kv is not None: + compressed_kv = torch.cat( + [compressed_kv, + new_compressed_kv.squeeze(0)], dim=0) + else: + compressed_kv = new_compressed_kv.squeeze(0) + + # Update window KV with the new token + # In a circular buffer: new token goes at position (start_pos % window_size) + # For simplicity in reference, we concatenate and take last window_size tokens + window_kv_with_new = torch.cat([window_kv, kv_seq.squeeze(0)], dim=0) + if window_kv_with_new.shape[0] > window_size: + window_kv_with_new = window_kv_with_new[-window_size:] + + # Build combined KV: [window_kv, compressed_kv] + if compressed_kv is not None: + num_compressed = compressed_kv.shape[0] + kv_combined = torch.cat([window_kv_with_new, compressed_kv], dim=0) + else: + num_compressed = 0 + kv_combined = window_kv_with_new + + total_kv_len = kv_combined.shape[0] + + # Create attention mask for generation: [1, total_kv_len] + # The single query can attend to all cached tokens (window + compressed) + attn_mask = torch.zeros((1, total_kv_len), + device=device, + dtype=q_seq.dtype) + + # Expand KV for multi-head (deepseek_v4 uses MQA - single KV head for all Q heads) + k_combined = kv_combined.unsqueeze(1).expand(-1, num_heads, -1) + v_combined = kv_combined.unsqueeze(1).expand(-1, num_heads, -1) + + # Prepare for attention: [1, num_heads, 1, head_dim] + # q: [1, num_heads, 1, qk_head_dim] + # k: [1, num_heads, total_kv_len, qk_head_dim] + # v: [1, num_heads, total_kv_len, qk_head_dim] + # attn_mask: [1, 1, 1, total_kv_len] + q_in = q_seq.transpose(1, 2) + k_in = k_combined.unsqueeze(0).transpose(1, 2) + v_in = v_combined.unsqueeze(0).transpose(1, 2) + attn_mask = attn_mask.unsqueeze(0).unsqueeze(0) + + # Compute attention, attn_out: [1, num_heads, qk_head_dim] + attn_out = torch.nn.functional.scaled_dot_product_attention( + q_in, k_in, v_in, attn_mask, + scale=softmax_scale).transpose(1, 2).squeeze(0) + + results.append(attn_out) + + return torch.cat(results, dim=0).flatten(1) + + +def calculate_reference_output_deepseek_v4_mixed(hidden_states, + q_ctx, + q_gen, + kv_data, + freqs_cis, + ref_compressor, + ctx_indices, + gen_indices, + seq_lens, + query_lens, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + softmax_scale, + device, + o_a_proj, + o_b_proj_weight, + n_local_groups, + window_size=512, + compress_ratio=4): + """Reference for deepseek_v4 mixed batch (combines context and generation).""" + ref_results = [None] * len(seq_lens) + + # Process context requests + if ctx_indices: + ctx_results = calculate_reference_output_deepseek_v4_prefill( + hidden_states, ctx_indices, seq_lens, q_ctx, kv_data, freqs_cis, + ref_compressor, num_heads, qk_nope_head_dim, qk_rope_head_dim, + v_head_dim, softmax_scale, device, o_a_proj, o_b_proj_weight, + n_local_groups, window_size, compress_ratio) + offset = 0 + for req_idx in ctx_indices: + seq_len = seq_lens[req_idx] + ref_results[req_idx] = ctx_results[offset:offset + seq_len] + offset += seq_len + + # Process generation requests + if gen_indices: + gen_results = calculate_reference_output_deepseek_v4_generation( + hidden_states, + gen_indices, seq_lens, q_gen, kv_data, freqs_cis, ref_compressor, + len(ctx_indices), num_heads, qk_nope_head_dim, qk_rope_head_dim, + v_head_dim, softmax_scale, device, o_a_proj, o_b_proj_weight, + n_local_groups, window_size, compress_ratio) + for idx, req_idx in enumerate(gen_indices): + ref_results[req_idx] = gen_results[idx:idx + 1] + + return torch.cat(ref_results, dim=0) + + +def prepare_reference_inputs( + sparse_attn_algo: str, + batch_order: List[int], + ctx_indices: List[int], + batch_query_lens: List[int], + q_original_for_ref: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + k_pe_original_for_ref: torch.Tensor, + cached_lens: List[int], + kv_cache_for_ref: dict, + kv_lora_rank: int, + num_heads: int, + qk_head_dim: int, + device: torch.device, +): + """Prepare reference inputs for both DSA and DeepSeek-V4 sparse attention. + + Args: + sparse_attn_algo: "dsa" or "deepseek_v4" + batch_order: List of request indices in batch order + ctx_indices: List of context request indices + batch_query_lens: Query lengths in batch order + q_original_for_ref: Original Q tensor (unrotated) + q: Q tensor (may be rotated for generation on SM90) + latent_cache: Latent cache [compressed_kv, k_pe] or full KV for deepseek_v4 + k_pe_original_for_ref: Original k_pe (unrotated) + cached_lens: Cached token counts per request + kv_cache_for_ref: Reference cache dict from populate_gen_*_kv_cache + kv_lora_rank: KV lora rank dimension + num_heads: Number of attention heads + qk_head_dim: Head dimension + device: Torch device + + Returns: + Tuple of (q_for_ref, kv_data): + - q_for_ref: [total_tokens, num_heads, qk_head_dim] + - kv_data: For DSA: (all_compressed_kv, all_k_pe_for_ref) + For DeepSeek-V4: all_kv + """ + q_for_ref_list = [] + batch_token_offsets = [0] + [ + sum(batch_query_lens[:i + 1]) for i in range(len(batch_order)) + ] + + if sparse_attn_algo == "dsa": + kv_c_list, k_pe_list = [], [] + + for batch_idx, orig_req_idx in enumerate(batch_order): + batch_start = batch_token_offsets[batch_idx] + batch_end = batch_token_offsets[batch_idx + 1] + + if orig_req_idx in ctx_indices: + # Context: use unrotated q and k_pe from latent_cache + q_req = q_original_for_ref[batch_start:batch_end] + kv_c_list.append( + latent_cache[batch_start:batch_end, :kv_lora_rank]) + k_pe_list.append(k_pe_original_for_ref[batch_start:batch_end]) + else: + # Generation: use original inputs and apply RoPE in the + # reference path, independent of GPU architecture. + cached_len = cached_lens[orig_req_idx] + q_req = q_original_for_ref[batch_start:batch_end] + k_pe_list.append( + torch.cat([ + kv_cache_for_ref["k_pe_original"][orig_req_idx] + [:cached_len], + k_pe_original_for_ref[batch_start:batch_end] + ])) + kv_c_list.append( + torch.cat([ + kv_cache_for_ref["compressed_kv"][orig_req_idx] + [:cached_len], + latent_cache[batch_start:batch_end, :kv_lora_rank] + ])) + + q_for_ref_list.append(q_req.view(-1, num_heads, qk_head_dim)) + + q_for_ref = torch.cat(q_for_ref_list, dim=0) + all_compressed_kv = torch.cat(kv_c_list, dim=0) + all_k_pe_for_ref = torch.cat(k_pe_list, dim=0) + return q_for_ref, (all_compressed_kv, all_k_pe_for_ref) + + elif sparse_attn_algo == "deepseek_v4": + all_latent_kv = {} + + for batch_idx, orig_req_idx in enumerate(batch_order): + batch_start = batch_token_offsets[batch_idx] + batch_end = batch_token_offsets[batch_idx + 1] + q_req = q_original_for_ref[batch_start:batch_end] + q_for_ref_list.append(q_req.view(-1, num_heads, qk_head_dim)) + all_latent_kv[orig_req_idx] = latent_cache[batch_start:batch_end, :] + + q_for_ref = torch.cat(q_for_ref_list, dim=0) + kv_cache_for_ref["latent_kv"] = all_latent_kv + return q_for_ref, kv_cache_for_ref + + else: + raise ValueError( + f"Invalid sparse attention algorithm: {sparse_attn_algo}") + + +def init_layers(mla_layers: List[MLA], sparse_attn_algo: str, + pretrained_config: DummyPretrainedConfig): + if sparse_attn_algo == "dsa": + nn_init_std = 0.02 + with torch.no_grad(): + for mla_layer in mla_layers: + # Initialize kv_b_proj weight + mla_layer.kv_b_proj.weight.normal_(mean=0.0, std=nn_init_std) + + # Extract weights for this layer + kv_b_weight = mla_layer.kv_b_proj.weight.data + kv_b_weight_reshaped = kv_b_weight.view( + pretrained_config.num_heads, + pretrained_config.qk_nope_head_dim + + pretrained_config.v_head_dim, + pretrained_config.kv_lora_rank) + # v_b_proj and k_b_proj_trans + mla_layer.v_b_proj.data = kv_b_weight_reshaped[:, + pretrained_config + . + qk_nope_head_dim:, :].contiguous( + ) + mla_layer.k_b_proj_trans.data = kv_b_weight_reshaped[:, : + pretrained_config + . + qk_nope_head_dim, :].transpose( + 1, 2 + ).contiguous( + ) + # Initialize indexer weights + mla_layer.mqa.indexer.wq_b.weight.normal_(mean=0.0, + std=nn_init_std) + mla_layer.mqa.indexer.wk.weight.normal_(mean=0.0, + std=nn_init_std) + mla_layer.mqa.indexer.weights_proj.weight.normal_( + mean=0.0, std=nn_init_std) + # Build fused wk+weights_proj weight after random init + mla_layer.mqa.indexer.post_load_weights() + elif sparse_attn_algo == "deepseek_v4": + nn_init_std = 0.02 + with torch.no_grad(): + for mla_layer in mla_layers: + if hasattr(mla_layer.mqa, 'indexer'): + # Initialize indexer weights for deepseek_v4 + mla_layer.mqa.indexer.wq_b.weight.normal_(mean=0.0, + std=nn_init_std) + mla_layer.mqa.indexer.weights_proj.weight.normal_( + mean=0.0, std=nn_init_std) + + # Initialize indexer's compressor weights + # Note: wkv_gate is a fused layer that projects to state_dim * 2 + mla_layer.mqa.indexer.compressor.wkv_gate.weight.normal_( + mean=0.0, std=nn_init_std) + mla_layer.mqa.indexer.compressor.ape.normal_( + mean=0.0, std=nn_init_std) + # RMSNorm weight is typically initialized to ones + mla_layer.mqa.indexer.compressor.norm.weight.fill_(1.0) + + # Initialize MLA's compressor weights (for KV cache compression) + if hasattr(mla_layer.mqa, 'compressor'): + mla_layer.mqa.compressor.wkv_gate.weight.normal_( + mean=0.0, std=nn_init_std) + mla_layer.mqa.compressor.ape.normal_(mean=0.0, + std=nn_init_std) + # RMSNorm weight is typically initialized to ones + mla_layer.mqa.compressor.norm.weight.fill_(1.0) + + +def populate_gen_dsa_kv_cache(mla: MLA, AttentionCls: AttentionBackend, + pretrained_config: DummyPretrainedConfig, + kv_cache_manager: DSACacheManager, + mapping: Mapping, + sparse_config: DeepSeekSparseAttentionConfig, + gen_indices: List[int], cached_lens: List[int], + dtype: torch.dtype, device: torch.device): + + all_cached_compressed_kv = {} + all_cached_k_pe_rotated = {} + all_cached_k_pe_original = {} + + if gen_indices: + gen_cached_lens = [ + cached_lens[i] for i in gen_indices if cached_lens[i] > 0 + ] + gen_with_cache = [i for i in gen_indices if cached_lens[i] > 0] + # Generate batched cache data + total_gen_cache_tokens = sum(gen_cached_lens) + batched_latent = torch.empty(total_gen_cache_tokens, + pretrained_config.kv_lora_rank + + pretrained_config.qk_rope_head_dim, + dtype=dtype, + device=device) + + offset = 0 + for req_idx in gen_with_cache: + cached_len = cached_lens[req_idx] + all_cached_compressed_kv[req_idx] = torch.randn( + cached_len, + pretrained_config.kv_lora_rank, + dtype=dtype, + device=device) + cached_k_pe = torch.randn(cached_len, + pretrained_config.qk_rope_head_dim, + dtype=dtype, + device=device) + batched_latent[offset:offset + cached_len] = torch.cat( + [all_cached_compressed_kv[req_idx], cached_k_pe], dim=-1) + offset += cached_len + + # Single batched metadata for all generation cache population + cached_metadata = AttentionCls.Metadata( + seq_lens=torch.tensor(gen_cached_lens, dtype=torch.int), + request_ids=gen_with_cache, + max_num_requests=len(gen_with_cache), + num_contexts=len(gen_with_cache), + prompt_lens=gen_cached_lens, + max_num_tokens=total_gen_cache_tokens, + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams(use_cache=True, + num_cached_tokens_per_seq=[0] * + len(gen_with_cache)), + mapping=mapping, + sparse_attention_config=sparse_config, + ) + cached_metadata.prepare() + + dummy_q = torch.randn(total_gen_cache_tokens, + pretrained_config.num_heads * + pretrained_config.qk_head_dim, + dtype=dtype, + device=device) + batched_latent_original = batched_latent.clone() + mla.mqa.mla_rope_append_paged_kv_assign_q(dummy_q, batched_latent, + cached_metadata) + + # Extract rotated k_pe for each request + offset = 0 + for req_idx in gen_with_cache: + cached_len = cached_lens[req_idx] + all_cached_k_pe_rotated[req_idx] = batched_latent[ + offset:offset + cached_len, + pretrained_config.kv_lora_rank:].clone() + all_cached_k_pe_original[req_idx] = batched_latent_original[ + offset:offset + cached_len, + pretrained_config.kv_lora_rank:].clone() + offset += cached_len + print( + f" - Allocated+populated cache for gen request {req_idx}: {cached_len} cached + 1 new = {cached_len+1} tokens" + ) + + kv_cache_for_ref = {} + kv_cache_for_ref["compressed_kv"] = all_cached_compressed_kv + kv_cache_for_ref["k_pe_rotated"] = all_cached_k_pe_rotated + kv_cache_for_ref["k_pe_original"] = all_cached_k_pe_original + + return kv_cache_for_ref + + +def populate_gen_deepseek_v4_kv_cache( + mla: MLA, ref_compressor: RefCompressor, AttentionCls: AttentionBackend, + pretrained_config: DummyPretrainedConfig, + kv_cache_manager: DeepseekV4CacheManager, mapping: Mapping, + sparse_config: DeepSeekV4SparseAttentionConfig, gen_indices: List[int], + cached_lens: List[int], freqs_cis: torch.Tensor, dtype: torch.dtype, + device: torch.device): + """Populate KV cache for deepseek_v4 generation requests by running prefill forward. + + Unlike DSA which uses mla_rope_append_paged_kv_assign_q to directly write to cache, + deepseek_v4 runs a full prefill forward pass to populate both: + - Window KV cache (recent tokens) + - Compressed KV cache (via compressor) + + This ensures the cache is populated through the same path as production. + """ + + all_ref_window_kv = {} + all_ref_compressed_kv = {} + all_ref_score_state = {} + all_ref_kv_state = {} + + if gen_indices: + gen_cached_lens = [ + cached_lens[i] for i in gen_indices if cached_lens[i] > 0 + ] + gen_with_cache = [i for i in gen_indices if cached_lens[i] > 0] + + if not gen_with_cache: + return {"window_kv": {}, "compressed_kv": {}, "full_kv": {}} + + # Generate batched input data for cached context tokens + total_gen_cache_tokens = sum(gen_cached_lens) + + # Generate input tensors for the prefill forward pass + q = torch.randn(total_gen_cache_tokens, + pretrained_config.num_heads * + pretrained_config.qk_head_dim, + dtype=dtype, + device=device) + qr = torch.randn(total_gen_cache_tokens, + pretrained_config.q_lora_rank, + dtype=dtype, + device=device) + compressed_kv = torch.randn(total_gen_cache_tokens, + pretrained_config.kv_lora_rank, + dtype=dtype, + device=device) + k_pe = torch.randn(total_gen_cache_tokens, + pretrained_config.qk_rope_head_dim, + dtype=dtype, + device=device) + hidden_states = torch.randn(total_gen_cache_tokens, + pretrained_config.hidden_size, + dtype=dtype, + device=device) + latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + position_ids = torch.cat([ + torch.arange(cached_lens[i], dtype=torch.int32, device=device) + for i in gen_with_cache + ]) + window_size = sparse_config.window_size + + # Run reference compressor forward for each request + latent_cache_ref = latent_cache.clone() + + offset = 0 + for batch_idx, req_idx in enumerate(gen_with_cache): + cached_len = cached_lens[req_idx] + req_hidden = hidden_states[offset:offset + cached_len].unsqueeze( + 0) # [1, seq_len, dim] + req_freqs = freqs_cis[:cached_len] + latent_cache_i = latent_cache_ref[offset:offset + cached_len] + + # Run reference compressor to get compressed KV + kv_cache = torch.zeros(1, + cached_len // ref_compressor.compress_ratio, + ref_compressor.head_dim, + device=device) + ref_compressor.kv_cache = kv_cache + ref_compressor.kv_state *= 0 + ref_compressor.score_state.fill_(-float("inf")) + compressed_kv_out = ref_compressor(req_hidden, + start_pos=0, + freqs_cis=req_freqs) + + # Extract kv_state and score_state for this request (batch_idx) + # These are the internal states used during compression + # Shape: [compress_ratio * (1 + overlap), head_dim] for overlap mode + all_ref_kv_state[req_idx] = ref_compressor.kv_state.clone() + all_ref_score_state[req_idx] = ref_compressor.score_state.clone() + + # Apply RoPE to the k_pe and collect the window KV + latent_cache_i = latent_cache_i.unsqueeze(0) + apply_rotary_emb( + latent_cache_i[..., -pretrained_config.qk_rope_head_dim:], + req_freqs) + req_latent = latent_cache_ref[offset:offset + cached_len] + if cached_len > window_size: + kv_window = req_latent[-window_size:] + else: + kv_window = req_latent + all_ref_window_kv[req_idx] = kv_window.clone() + # Concatenate the window KV and the compressed KV and store it to all_ref_kv + if compressed_kv_out is not None: + all_ref_compressed_kv[req_idx] = compressed_kv_out.squeeze( + 0).clone() + else: + all_ref_compressed_kv[req_idx] = None + + offset += cached_len + num_compressed = compressed_kv_out.shape[ + 1] if compressed_kv_out is not None else 0 + print(f" - Computed reference cache for gen request {req_idx}: " + f"{cached_len} tokens -> {num_compressed} compressed tokens") + + # Step 2: Run TRT-LLM forward to populate actual KV cache + cached_metadata = AttentionCls.Metadata( + seq_lens=torch.tensor(gen_cached_lens, dtype=torch.int), + request_ids=gen_with_cache, + max_num_requests=len(gen_with_cache), + num_contexts=len( + gen_with_cache), # All are context (prefill) requests + prompt_lens=gen_cached_lens, + max_num_tokens=total_gen_cache_tokens, + kv_cache_manager=kv_cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[0] * + len(gen_with_cache) # No prior cache + ), + mapping=mapping, + sparse_attention_config=sparse_config, + ) + cached_metadata.prepare() + _ = mla_forward_impl_with_deepseek_v4_wo_linear(mla, cached_metadata, q, qr, + compressed_kv, k_pe, + latent_cache, hidden_states, + position_ids, dtype, device) + + kv_cache_for_ref = {} + kv_cache_for_ref["window_kv"] = all_ref_window_kv # Window KV for reference + kv_cache_for_ref[ + "compressed_kv"] = all_ref_compressed_kv # Compressor's compressed KV + kv_cache_for_ref["kv_state"] = all_ref_kv_state # Compressor's kv_state + kv_cache_for_ref[ + "score_state"] = all_ref_score_state # Compressor's score_state + + return kv_cache_for_ref + + +def mla_forward_impl_with_dsa_wo_linear(mla, attn_metadata, q, qr, + compressed_kv, k_pe, latent_cache, + hidden_states, position_ids, dtype, + device): + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + num_gen_tokens = num_tokens - num_ctx_tokens + + output = torch.empty(num_tokens, + mla.num_heads * mla.v_head_dim, + dtype=dtype, + device=device) + + topk_indices_local = mla.mqa.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + ) + + # Validate indexer output against expected causal indices (since seq_len < topk=2048) + if num_contexts > 0: + # Transform context indices from local to global + ctx_topk_local = topk_indices_local[:num_ctx_tokens] + + mla.forward_context_sparse_mla( + q=q[:num_ctx_tokens], + compressed_kv=compressed_kv[:num_ctx_tokens], + k_pe=k_pe[:num_ctx_tokens], + attn_metadata=attn_metadata, + output=output[:num_ctx_tokens], + latent_cache=latent_cache[:num_ctx_tokens], + topk_indices=ctx_topk_local, # Use global indices + ) + print( + f" ✓ Context forward: {num_ctx_tokens} tokens from {num_contexts} requests" + ) + + if num_generations > 0: + # Transform generation indices from local to global + gen_topk_local = topk_indices_local[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens] + mla.forward_generation_sparse_mla( + q=q[num_ctx_tokens:], + compressed_kv=compressed_kv[num_ctx_tokens:], + k_pe=k_pe[num_ctx_tokens:], + attn_metadata=attn_metadata, + output=output[num_ctx_tokens:], + latent_cache=latent_cache[num_ctx_tokens:], + topk_indices=gen_topk_local, # Use global indices + ) + print( + f" ✓ Generation forward: {num_gen_tokens} tokens from {num_generations} requests" + ) + return output + + +def mla_forward_impl_with_deepseek_v4_wo_linear(mla, attn_metadata, q, qr, + compressed_kv, k_pe, latent_cache, + hidden_states, position_ids, dtype, + device): + """Forward implementation for deepseek_v4 sparse attention. + + For deepseek_v4, compressed_kv is actually the nope part of KV (qk_nope_head_dim), + and k_pe is the rope part (qk_rope_head_dim). Together they form the full KV. + + This includes: + 1. Indexer forward (with its compressor) to get topk indices + 2. MLA's compressor forward to compress KV cache + 3. Sparse attention computation + """ + num_contexts = attn_metadata.num_contexts + num_generations = attn_metadata.num_generations + num_ctx_tokens = attn_metadata.num_ctx_tokens + num_tokens = attn_metadata.num_tokens + num_gen_tokens = num_tokens - num_ctx_tokens + + output = torch.empty(num_tokens, + mla.num_heads * mla.v_head_dim, + dtype=dtype, + device=device) + + # Step 1: Call indexer to get topk indices (indexer internally uses its compressor) + topk_indices_local = mla.mqa.indexer( + qr, + hidden_states, + attn_metadata, + position_ids, + ) + + # Step 2: Call MLA's compressor to compress KV cache + if hasattr(mla.mqa, 'compressor') and mla.mqa.compressor is not None: + mla.mqa.compressor( + hidden_states, # [num_tokens, dim] + attn_metadata, # Contains all cache management info + ) + + print( + f" ✓ Compressor forward: processed {num_contexts} ctx + {num_generations} gen requests" + ) + + # Step 3: Process context requests + if num_contexts > 0: + ctx_topk_local = topk_indices_local[:num_ctx_tokens] + + mla.forward_context_sparse_mla( + q=q[:num_ctx_tokens], + compressed_kv=compressed_kv[:num_ctx_tokens], + k_pe=k_pe[:num_ctx_tokens], + attn_metadata=attn_metadata, + output=output[:num_ctx_tokens], + latent_cache=latent_cache[:num_ctx_tokens], + topk_indices=ctx_topk_local, + position_ids=position_ids[:num_ctx_tokens], + ) + print( + f" ✓ Context forward: {num_ctx_tokens} tokens from {num_contexts} requests" + ) + + # Step 4: Process generation requests + if num_generations > 0: + gen_topk_local = topk_indices_local[num_ctx_tokens:num_tokens] + + mla.forward_generation_sparse_mla( + q=q[num_ctx_tokens:], + compressed_kv=compressed_kv[num_ctx_tokens:], + k_pe=k_pe[num_ctx_tokens:], + attn_metadata=attn_metadata, + output=output[num_ctx_tokens:], + latent_cache=latent_cache[num_ctx_tokens:], + topk_indices=gen_topk_local, + position_ids=position_ids[num_ctx_tokens:num_tokens], + ) + print( + f" ✓ Generation forward: {num_gen_tokens} tokens from {num_generations} requests" + ) + + return output + + @pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") @pytest.mark.skipif(get_sm_version() < 90, reason="FlashMLA requires SM90 (Hopper) or later") @pytest.mark.parametrize("batch_name", list(BATCH_SPECS.keys())) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) -def test_forward_sparse_mla_unified(batch_name, kv_cache_dtype: str): +@pytest.mark.parametrize("sparse_attn_algo", ["dsa", "deepseek_v4"]) +def test_forward_sparse_mla_unified(batch_name, kv_cache_dtype: str, + sparse_attn_algo: str): """Test sparse MLA attention for pure prefill, pure decode, and mixed batches.""" print( - f"\n{'='*80}\nTesting: {batch_name} kv_cache_dtype: {kv_cache_dtype}\n{'='*80}" + f"\n{'='*80}\nTesting: {batch_name}, sparse_attn_algo: {sparse_attn_algo}, kv_cache_dtype: {kv_cache_dtype}\n{'='*80}" ) + if sparse_attn_algo == "deepseek_v4" and get_sm_version() < 100: + pytest.skip( + "DeepSeek-V4 sparse MLA unittest is not supported on pre-Blackwell architectures" + ) if kv_cache_dtype == "fp8" and get_sm_version() < 100: pytest.skip( "FP8 kv cache is not supported on pre-Blackwell architectures") @@ -319,60 +1280,125 @@ def test_forward_sparse_mla_unified(batch_name, kv_cache_dtype: str): ) # Model configuration - # Since topk=2048 > seq_lens, indexer selects all tokens anyway - num_heads = 128 - q_lora_rank = 512 - kv_lora_rank = 512 - qk_nope_head_dim = 128 - qk_rope_head_dim = 64 - v_head_dim = 128 - qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - head_size = kv_lora_rank + qk_rope_head_dim # 576 - topk_tokens = 2048 - hidden_size = 2048 - max_position_embeddings = 4096 - tokens_per_block = 64 - num_layers = 3 # Test multi-layer pool + if sparse_attn_algo == "dsa": + # Since topk=2048 > seq_lens, indexer selects all tokens anyway + num_heads = 128 + q_lora_rank = 512 + kv_lora_rank = 512 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + v_head_dim = 128 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + head_size = kv_lora_rank + qk_rope_head_dim # 576 + topk_tokens = 2048 + hidden_size = 2048 + max_position_embeddings = 4096 + num_groups = 1 + o_lora_rank = 1024 + num_layers = 3 # Test multi-layer pool + vocab_size = 129280 + compress_ratios = [1, 1, 1] + tokens_per_block = 64 + elif sparse_attn_algo == "deepseek_v4": + num_heads = 64 + q_lora_rank = 1024 + kv_lora_rank = 448 + qk_nope_head_dim = 448 + qk_rope_head_dim = 64 + v_head_dim = 512 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + head_size = kv_lora_rank + qk_rope_head_dim # 512 + topk_tokens = 512 + hidden_size = 4096 + max_position_embeddings = 4096 + num_groups = 8 + o_lora_rank = 1024 + num_layers = 3 # Test multi-layer pool + vocab_size = 129280 + compress_ratios = [1, 4, 128] + tokens_per_block = 128 + else: + raise ValueError( + f"Invalid sparse attention algorithm: {sparse_attn_algo}") + + pretrained_config = DummyPretrainedConfig( + num_heads=num_heads, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + qk_head_dim=qk_head_dim, + head_size=head_size, + topk_tokens=topk_tokens, + hidden_size=hidden_size, + max_position_embeddings=max_position_embeddings, + num_groups=num_groups, + o_lora_rank=o_lora_rank, + num_layers=num_layers, + vocab_size=vocab_size, + compress_ratios=compress_ratios, + ) + layer_idx = 1 # Use middle layer to verify layer extraction torch.manual_seed(42) torch.cuda.manual_seed(42) # Create RoPE config - rope_config = RopeConfig( - hidden_size=hidden_size, - num_attention_heads=num_heads, - rope_scaling={ - "beta_fast": 32, - "beta_slow": 1, - "factor": 40, - "mscale": 1.0, - "mscale_all_dim": 1.0, - "original_max_position_embeddings": 4096, - "type": "yarn", - }, - max_position_embeddings=max_position_embeddings, - rope_theta=10000.0, - qk_rope_head_dim=qk_rope_head_dim, - model_type="deepseek_v2", - ) - - # Generate RoPE cos/sin embeddings for reference calculation - 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) + if sparse_attn_algo == "dsa": + rope_config = RopeConfig( + hidden_size=hidden_size, + num_attention_heads=num_heads, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + max_position_embeddings=max_position_embeddings, + rope_theta=10000.0, + qk_rope_head_dim=qk_rope_head_dim, + model_type="deepseek_v2", + ) + # Generate RoPE cos/sin embeddings for reference calculation + 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) + elif sparse_attn_algo == "deepseek_v4": + rope_config = RopeConfig( + hidden_size=hidden_size, + num_attention_heads=num_heads, + rope_scaling={ + "beta_fast": 32, + "beta_slow": 1, + "factor": 4, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + max_position_embeddings=max_position_embeddings, + rope_theta=10000.0, + qk_rope_head_dim=qk_rope_head_dim, + model_type="deepseek_v4", + ) + freqs_cis = None # Will be set after MLA modules are created # Calculate scaling factors (aligned with vLLM) def yarn_get_mscale(scale=1, mscale=1): @@ -393,12 +1419,27 @@ def yarn_get_mscale(scale=1, mscale=1): max_seqlen = max(seq_lens) max_tokens = 16384 - # Create sparse attention config (DSA - DeepSeek Sparse Attention) - sparse_config = DeepSeekSparseAttentionConfig( - index_n_heads=64, # Number of heads for indexer - index_head_dim=128, # Dimension of indexer heads - index_topk=topk_tokens, # Top-k tokens to select (2048) - ) + # Create sparse attention config + if sparse_attn_algo == "dsa": + # (DSA - DeepSeek Sparse Attention) + sparse_config = DeepSeekSparseAttentionConfig( + index_n_heads=64, # Number of heads for indexer + index_head_dim=128, # Dimension of indexer heads + index_topk=topk_tokens, # Top-k tokens to select (2048) + ) + elif sparse_attn_algo == "deepseek_v4": + # (DeepSeek-V4 - DeepSeek-V4 Sparse Attention) + sparse_config = DeepSeekV4SparseAttentionConfig( + index_n_heads=64, # Number of heads for indexer + index_head_dim=128, # Dimension of indexer heads + index_topk=topk_tokens, # Top-k tokens to select (512) + compress_ratios=compress_ratios, + ) + else: + raise ValueError( + f"Invalid sparse attention algorithm: {sparse_attn_algo}") + + print(f"sparse_config: {sparse_config}") if kv_cache_dtype == "auto": cache_dtype = dtype @@ -407,13 +1448,16 @@ def yarn_get_mscale(scale=1, mscale=1): else: cache_dtype = dtype - # Create pretrained_config for ModelConfig - pretrained_config = SimpleNamespace(rms_norm_eps=1e-6, ) + # Configure quantization for FP8 KV cache (per-tensor FP8, no blockwise scales) + quant_config = QuantConfig() + if kv_cache_dtype == "fp8": + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 model_config = ModelConfig( mapping=mapping, sparse_attention_config=sparse_config, - pretrained_config=pretrained_config, + pretrained_config=SimpleNamespace(rms_norm_eps=1e-6, ), + quant_config=quant_config, ) # Setup positional embedding params @@ -442,8 +1486,13 @@ def yarn_get_mscale(scale=1, mscale=1): layer_idx=layer_id, dtype=dtype, config=model_config, + num_groups=num_groups, + o_lora_rank=o_lora_rank, ).to(device) - mla.indexer = mla.mqa.indexer.to(device) + if hasattr(mla.mqa, 'indexer'): + mla.indexer = mla.mqa.indexer.to(device) + if hasattr(mla.mqa, 'compressor'): + mla.compressor = mla.mqa.compressor.to(device) mla_layers.append(mla) # Use the test layer @@ -453,34 +1502,21 @@ def yarn_get_mscale(scale=1, mscale=1): else: print(f" Testing single layer (baseline)") + # For deepseek_v4: derive freqs_cis from the production RotaryEmbedding to + # guarantee matching frequencies (the reference precompute_freqs_cis may + # differ from the production's yarn-scaled frequencies). + if sparse_attn_algo == "deepseek_v4" and freqs_cis is None: + prod_cos_sin = mla.mqa.compressor.rotary_emb.rotary_cos_sin + # prod_cos_sin: [max_positions, 2, rope_dim//2] + prod_cos = prod_cos_sin[:, 0, :] # [max_positions, rope_dim//2] + prod_sin = prod_cos_sin[:, 1, :] # [max_positions, rope_dim//2] + freqs_cis = torch.complex(prod_cos, prod_sin).to(device) + # Initialize weights for all layers - nn_init_std = 0.02 - with torch.no_grad(): - for mla_layer in mla_layers: - # Initialize kv_b_proj weight - mla_layer.kv_b_proj.weight.normal_(mean=0.0, std=nn_init_std) - - # Extract weights for this layer - kv_b_weight = mla_layer.kv_b_proj.weight.data - kv_b_weight_reshaped = kv_b_weight.view( - num_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank) - # v_b_proj and k_b_proj_trans - mla_layer.v_b_proj.data = kv_b_weight_reshaped[:, - qk_nope_head_dim:, :].contiguous( - ) - mla_layer.k_b_proj_trans.data = kv_b_weight_reshaped[:, : - qk_nope_head_dim, :].transpose( - 1, 2 - ).contiguous() - # Initialize indexer weights - mla_layer.mqa.indexer.wq_b.weight.normal_(mean=0.0, std=nn_init_std) - mla_layer.mqa.indexer.wk.weight.normal_(mean=0.0, std=nn_init_std) - mla_layer.mqa.indexer.weights_proj.weight.normal_(mean=0.0, - std=nn_init_std) - # Build fused wk+weights_proj weight after random init - mla_layer.mqa.indexer.post_load_weights() + init_layers(mla_layers, sparse_attn_algo, pretrained_config) - # Extract W_UK and W_UV from test layer for reference calculation + # Extract W_UK and W_UV from test layer for reference calculation (DSA only) + if sparse_attn_algo == "dsa": kv_b_weight = mla.kv_b_proj.weight.data kv_b_weight_reshaped = kv_b_weight.view(num_heads, qk_nope_head_dim + v_head_dim, @@ -489,13 +1525,49 @@ def yarn_get_mscale(scale=1, mscale=1): 2, 0, 1).contiguous() # [kv_lora_rank, num_heads, qk_nope_head_dim] W_UV = kv_b_weight_reshaped[:, qk_nope_head_dim:, :].permute( 2, 0, 1).contiguous() # [kv_lora_rank, num_heads, v_head_dim] + elif sparse_attn_algo == "deepseek_v4": + # DeepSeek-V4 doesn't use W_UK, W_UV expansion - KV is stored at full head_dim + W_UK, W_UV = None, None + # Create compressor for reference calculation + # Create reference compressor with matching config + ref_args = SimpleNamespace( + dim=pretrained_config.hidden_size, + head_dim=pretrained_config.head_size, + rope_head_dim=pretrained_config.qk_rope_head_dim, + norm_eps=1e-6, + max_batch_size=1, + max_seq_len=max_seqlen, + ) + # V4 reference: attention's compressor uses rotate=False (only the + # indexer's compressor rotates). See V4-Pro inference/model.py + # Attention.__init__: ``Compressor(args, ratio, head_dim)`` (default + # rotate=False), Indexer.__init__: ``Compressor(..., rotate=True)``. + ref_compressor = RefCompressor(ref_args, + compress_ratios[layer_idx], + pretrained_config.head_size, + rotate=False).to(device) + + # Initialize reference compressor with same weights as MLA's compressor + weights_wkv_gate = mla.mqa.compressor.wkv_gate.weight.data + weights_wkv = weights_wkv_gate[:mla.mqa.compressor.state_dim] + weights_wgate = weights_wkv_gate[mla.mqa.compressor.state_dim:] + ref_compressor.wkv.weight.data.copy_(weights_wkv) + ref_compressor.wgate.weight.data.copy_(weights_wgate) + ref_compressor.ape.data.copy_(mla.mqa.compressor.ape.data) + ref_compressor.norm.weight.data.copy_( + mla.mqa.compressor.norm.weight.data) + else: + W_UK, W_UV = None, None # Calculate cached token counts (context already in cache) cached_lens = [seq_lens[i] - query_lens[i] for i in range(len(seq_lens))] # Create KV cache manager - kv_cache_manager = DSACacheManager( - KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False), + kv_cache_manager_cls = DSACacheManager if sparse_attn_algo == "dsa" else DeepseekV4CacheManager + kv_cache_manager = kv_cache_manager_cls( + KvCacheConfig(max_tokens=max_tokens, + enable_block_reuse=False, + event_buffer_max_size=0), tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, num_layers=num_layers, num_kv_heads=1, @@ -507,14 +1579,12 @@ def yarn_get_mscale(scale=1, mscale=1): dtype=str_dtype_to_binding(torch_dtype_to_str(cache_dtype)), sparse_attn_config=sparse_config, model_config=model_config, + vocab_size=vocab_size, ) AttentionCls = get_attention_backend("TRTLLM", sparse_config) # Allocate and pre-populate KV cache in batch order [context...][generation...] - all_cached_compressed_kv = {} - all_cached_k_pe_rotated = {} - all_cached_k_pe_original = {} print(f" Allocating and pre-populating cache...") @@ -542,71 +1612,25 @@ def yarn_get_mscale(scale=1, mscale=1): for req_idx in gen_with_cache: kv_cache_manager.add_dummy_requests( request_ids=[req_idx], - token_nums=[cached_lens[req_idx]], + token_nums=[seq_lens[req_idx]], is_gen=False, prepare_resource=True, ) - _allocate_kv_cache_for_generation(kv_cache_manager, [req_idx]) - - # Generate batched cache data - total_gen_cache_tokens = sum(gen_cached_lens) - batched_latent = torch.empty(total_gen_cache_tokens, - kv_lora_rank + qk_rope_head_dim, - dtype=dtype, - device=device) - - offset = 0 - for req_idx in gen_with_cache: - cached_len = cached_lens[req_idx] - all_cached_compressed_kv[req_idx] = torch.randn(cached_len, - kv_lora_rank, - dtype=dtype, - device=device) - cached_k_pe = torch.randn(cached_len, - qk_rope_head_dim, - dtype=dtype, - device=device) - batched_latent[offset:offset + cached_len] = torch.cat( - [all_cached_compressed_kv[req_idx], cached_k_pe], dim=-1) - offset += cached_len - - # Single batched metadata for all generation cache population - cached_metadata = AttentionCls.Metadata( - seq_lens=torch.tensor(gen_cached_lens, dtype=torch.int), - request_ids=gen_with_cache, - max_num_requests=len(gen_with_cache), - num_contexts=len(gen_with_cache), - prompt_lens=gen_cached_lens, - max_num_tokens=total_gen_cache_tokens, - kv_cache_manager=kv_cache_manager, - kv_cache_params=KVCacheParams(use_cache=True, - num_cached_tokens_per_seq=[0] * - len(gen_with_cache)), - mapping=mapping, - sparse_attention_config=sparse_config, - ) - cached_metadata.prepare() - - dummy_q = torch.randn(total_gen_cache_tokens, - num_heads * qk_head_dim, - dtype=dtype, - device=device) - batched_latent_original = batched_latent.clone() - mla.mqa.mla_rope_append_paged_kv_assign_q(dummy_q, batched_latent, - cached_metadata) - # Extract rotated k_pe for each request - offset = 0 - for req_idx in gen_with_cache: - cached_len = cached_lens[req_idx] - all_cached_k_pe_rotated[req_idx] = batched_latent[ - offset:offset + cached_len, kv_lora_rank:].clone() - all_cached_k_pe_original[req_idx] = batched_latent_original[ - offset:offset + cached_len, kv_lora_rank:].clone() - offset += cached_len - print( - f" - Allocated+populated cache for gen request {req_idx}: {cached_len} cached + 1 new = {seq_lens[req_idx]} tokens" - ) + if sparse_attn_algo == "dsa": + kv_cache_for_ref = populate_gen_dsa_kv_cache(mla, AttentionCls, + pretrained_config, + kv_cache_manager, mapping, + sparse_config, gen_indices, + cached_lens, dtype, device) + elif sparse_attn_algo == "deepseek_v4": + kv_cache_for_ref = populate_gen_deepseek_v4_kv_cache( + mla, ref_compressor, AttentionCls, pretrained_config, + kv_cache_manager, mapping, sparse_config, gen_indices, cached_lens, + freqs_cis, dtype, device) + else: + raise ValueError( + f"Invalid sparse attention algorithm: {sparse_attn_algo}") print(f" ✓ KV cache allocated and pre-populated") @@ -644,10 +1668,6 @@ def yarn_get_mscale(scale=1, mscale=1): device=device, dtype=torch.int32) for i in batch_order ]) - output = torch.empty(total_query_tokens, - num_heads * v_head_dim, - dtype=dtype, - device=device) attn_metadata = AttentionCls.Metadata( seq_lens=torch.tensor(batch_query_lens, dtype=torch.int), @@ -669,137 +1689,119 @@ def yarn_get_mscale(scale=1, mscale=1): ) attn_metadata.prepare() + q_original_for_ref = q.clone() + hidden_states_ref = hidden_states.clone() + latent_cache_for_ref = latent_cache.clone() + assert hasattr( mla, 'softmax_scale') and abs(mla.softmax_scale - softmax_scale) < 1e-6 print(f" ✓ Inputs prepared: {total_query_tokens} query tokens") - q_original_for_ref = q.clone() - batch_token_offsets = [0] + [ - sum(batch_query_lens[:i + 1]) for i in range(len(batch_order)) - ] - num_ctx_tokens = sum(query_lens[i] for i in ctx_indices) - topk_indices_local = mla.mqa.indexer( - qr, - hidden_states, - attn_metadata, - position_ids, - ) - - # Validate indexer output against expected causal indices (since seq_len < topk=2048) - if num_contexts > 0: - # Transform context indices from local to global - ctx_topk_local = topk_indices_local[:num_ctx_tokens] - - mla.forward_context_dsa( - q=q[:num_ctx_tokens], - compressed_kv=compressed_kv[:num_ctx_tokens], - k_pe=k_pe[:num_ctx_tokens], - attn_metadata=attn_metadata, - output=output[:num_ctx_tokens], - latent_cache=latent_cache[:num_ctx_tokens], - topk_indices=ctx_topk_local, # Use global indices - ) - print( - f" ✓ Context forward: {num_ctx_tokens} tokens from {num_contexts} requests" - ) - - if num_generations > 0: - # Transform generation indices from local to global - num_gen_tokens = sum(query_lens[i] for i in gen_indices) - gen_topk_local = topk_indices_local[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens] - mla.forward_generation_dsa( - q=q[num_ctx_tokens:], - compressed_kv=compressed_kv[num_ctx_tokens:], - k_pe=k_pe[num_ctx_tokens:], - attn_metadata=attn_metadata, - output=output[num_ctx_tokens:], - latent_cache=latent_cache[num_ctx_tokens:], - topk_indices=gen_topk_local, # Use global indices - ) - print( - f" ✓ Generation forward: {sum(query_lens[i] for i in gen_indices)} tokens from {num_generations} requests" - ) + if sparse_attn_algo == "dsa": + output = mla_forward_impl_with_dsa_wo_linear( + mla, attn_metadata, q, qr, compressed_kv, k_pe, latent_cache, + hidden_states, position_ids, dtype, device) + elif sparse_attn_algo == "deepseek_v4": + output = mla_forward_impl_with_deepseek_v4_wo_linear( + mla, attn_metadata, q, qr, compressed_kv, k_pe, latent_cache, + hidden_states, position_ids, dtype, device) + else: + raise ValueError( + f"Invalid sparse attention algorithm: {sparse_attn_algo}") print(f" ✓ Forward pass complete: output shape {output.shape}") - # Assemble reference in BATCH order (same as output) - q_for_ref_list, kv_c_list, k_pe_list = [], [], [] - - for batch_idx, orig_req_idx in enumerate(batch_order): - batch_start = batch_token_offsets[batch_idx] - batch_end = batch_token_offsets[batch_idx + 1] - - if orig_req_idx in ctx_indices: - # Context: use unrotated q and k_pe from latent_cache - q_req = q_original_for_ref[batch_start:batch_end] - kv_c_list.append(latent_cache[batch_start:batch_end, :kv_lora_rank]) - k_pe_list.append(k_pe_original_for_ref[batch_start:batch_end]) - else: - # Generation: - # SM90: use rotated q, combine rotated k_pe and new KV from latent_cache - # SM100+: use unrotated q, combine unrotated k_pe and new KV from latent_cache - cached_len = cached_lens[orig_req_idx] - if get_sm_version() >= 100: - q_req = q_original_for_ref[batch_start:batch_end] - k_pe_list.append( - torch.cat([ - all_cached_k_pe_original[orig_req_idx][:cached_len], - latent_cache[batch_start:batch_end, kv_lora_rank:] - ])) - else: - q_req = q[batch_start:batch_end] - k_pe_list.append( - torch.cat([ - all_cached_k_pe_rotated[orig_req_idx][:cached_len], - latent_cache[batch_start:batch_end, kv_lora_rank:] - ])) - kv_c_list.append( - torch.cat([ - all_cached_compressed_kv[orig_req_idx][:cached_len], - latent_cache[batch_start:batch_end, :kv_lora_rank] - ])) - - q_for_ref_list.append(q_req.view(-1, num_heads, qk_head_dim)) - - q_for_ref = torch.cat(q_for_ref_list, dim=0) - all_compressed_kv = torch.cat(kv_c_list, dim=0) - all_k_pe_for_ref = torch.cat(k_pe_list, dim=0) + # Assemble reference inputs in BATCH order (same as output) + q_for_ref, kv_data = prepare_reference_inputs( + sparse_attn_algo=sparse_attn_algo, + batch_order=batch_order, + ctx_indices=ctx_indices, + batch_query_lens=batch_query_lens, + q_original_for_ref=q_original_for_ref, + q=q, + latent_cache=latent_cache_for_ref, + k_pe_original_for_ref=k_pe_original_for_ref, + cached_lens=cached_lens, + kv_cache_for_ref=kv_cache_for_ref, + kv_lora_rank=kv_lora_rank, + num_heads=num_heads, + qk_head_dim=qk_head_dim, + device=device, + ) print( f" - Computing reference ({num_contexts} ctx, {num_generations} gen)..." ) - q_ctx_ref = q_for_ref[:num_ctx_tokens] if ctx_indices else torch.empty( - 0, num_heads, qk_head_dim, device=device) - q_gen_ref = q_for_ref[num_ctx_tokens:] if gen_indices else torch.empty( - 0, num_heads, qk_head_dim, device=device) - - reference_output = calculate_reference_output_mixed( - q_ctx=q_ctx_ref, - q_gen=q_gen_ref, - kv_c_all=all_compressed_kv, - k_pe_all=all_k_pe_for_ref, - W_UK=W_UK, - W_UV=W_UV, - rope_cos_sin=rope_cos_sin, - ctx_indices=list(range(num_contexts)), - gen_indices=list(range(num_contexts, len(batch_order))), - seq_lens=[seq_lens[i] for i in batch_order], - query_lens=batch_query_lens, - num_heads=num_heads, - kv_lora_rank=kv_lora_rank, - qk_nope_head_dim=qk_nope_head_dim, - qk_rope_head_dim=qk_rope_head_dim, - v_head_dim=v_head_dim, - softmax_scale=softmax_scale, - device=device, - ) + q_ctx_ref = q_for_ref[:attn_metadata. + num_ctx_tokens] if ctx_indices else torch.empty( + 0, num_heads, qk_head_dim, device=device) + q_gen_ref = q_for_ref[attn_metadata. + num_ctx_tokens:] if gen_indices else torch.empty( + 0, num_heads, qk_head_dim, device=device) + + if sparse_attn_algo == "dsa": + all_compressed_kv, all_k_pe_for_ref = kv_data + reference_output = calculate_reference_output_mixed( + q_ctx=q_ctx_ref, + q_gen=q_gen_ref, + kv_c_all=all_compressed_kv, + k_pe_all=all_k_pe_for_ref, + W_UK=W_UK, + W_UV=W_UV, + rope_cos_sin=rope_cos_sin, + ctx_indices=list(range(num_contexts)), + gen_indices=list(range(num_contexts, len(batch_order))), + seq_lens=[seq_lens[i] for i in batch_order], + query_lens=batch_query_lens, + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + softmax_scale=softmax_scale, + device=device, + ) + elif sparse_attn_algo == "deepseek_v4": + # Extract output projection weights for reference calculation + o_a_proj_ref = mla.o_a_proj.data + o_b_proj_weight_ref = mla.o_b_proj.weight.data + n_local_groups = num_groups # No TP in this test + + reference_output = calculate_reference_output_deepseek_v4_mixed( + hidden_states=hidden_states_ref, + q_ctx=q_ctx_ref, + q_gen=q_gen_ref, + kv_data=kv_data, + freqs_cis=freqs_cis, + ref_compressor=ref_compressor, + ctx_indices=list(range(num_contexts)), + gen_indices=list(range(num_contexts, len(batch_order))), + seq_lens=[seq_lens[i] for i in batch_order], + query_lens=batch_query_lens, + num_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + softmax_scale=softmax_scale, + device=device, + o_a_proj=o_a_proj_ref, + o_b_proj_weight=o_b_proj_weight_ref, + n_local_groups=n_local_groups, + window_size=sparse_config.window_size, + compress_ratio=compress_ratios[layer_idx], + ) + else: + raise ValueError( + f"Invalid sparse attention algorithm: {sparse_attn_algo}") assert output.shape == reference_output.shape and output.dtype == reference_output.dtype - assert torch.isfinite(output).all() and torch.isfinite( - reference_output).all() + assert torch.isfinite(output).all() + assert torch.isfinite(reference_output).all() # Compare directly (both in batch order now) + batch_token_offsets = [0] + [ + sum(batch_query_lens[:i + 1]) for i in range(len(batch_order)) + ] abs_error = (output - reference_output).abs() for batch_idx, orig_req_idx in enumerate(batch_order): req_error = abs_error[ @@ -810,7 +1812,11 @@ def yarn_get_mscale(scale=1, mscale=1): f" ⚠ Request {orig_req_idx} [{req_type}]: max error {req_error.max():.3f}" ) - torch.testing.assert_close(output, reference_output, rtol=0.1, atol=0.1) + if kv_cache_dtype == "auto": + torch.testing.assert_close(output, reference_output, rtol=0.2, atol=0.2) + else: + diff = _calc_diff(output, reference_output) + assert diff < 1e-2, f"{diff=}" print( f" ✓ Validation passed: max_error={abs_error.max():.4f}, mean_error={abs_error.mean():.6f}" ) @@ -819,28 +1825,6 @@ def yarn_get_mscale(scale=1, mscale=1): print(f" ✓ Test '{batch_name}' completed\n") -def _allocate_kv_cache_for_generation(kv_cache_manager, request_ids): - """ - Allocate KV cache blocks for generation phase following PyExecutor's pattern. - - This mimics the production flow: prepare_resources() calls add_token() - for each generation request, which allocates blocks as needed. - - For DSACacheManager, we need to allocate for both: - - Main KV cache (via impl.add_token) - - Indexer K cache (via indexer_k_cache_manager.add_tokens) - """ - for request_id in request_ids: - # Allocate main KV cache block (mimics prepare_resources flow) - kv_cache_manager.impl.add_token(request_id) - - # Allocate indexer K cache block for DSA sparse attention - if hasattr(kv_cache_manager, 'indexer_k_cache_manager'): - kv_cache_manager.indexer_k_cache_manager.add_tokens(request_id, 1) - - -# Old test_forward_sparse_mla_generation removed - unified into test_forward_sparse_mla_unified - if __name__ == "__main__": # Test pure prefill test_forward_sparse_mla_unified(batch_name="small_prefill", diff --git a/tests/unittest/_torch/attention/sparse/test_triton_topk.py b/tests/unittest/_torch/attention/sparse/test_triton_topk.py deleted file mode 100644 index 5375619bb411..000000000000 --- a/tests/unittest/_torch/attention/sparse/test_triton_topk.py +++ /dev/null @@ -1,228 +0,0 @@ -import pytest -import torch - -from tensorrt_llm._torch.attention_backend.sparse.kernel import triton_topk - - -def pytorch_reference_topk( - input_tensor: torch.Tensor, - kv_offsets: torch.Tensor, - kv_lens: torch.Tensor, - topk: int, -) -> torch.Tensor: - """ - Args: - input_tensor: Input scores [num_kv_heads, sum(kv_lens)] - kv_offsets: KV offsets [batch_size + 1] - kv_lens: KV lengths [batch_size] - topk: TopK parameter - - Returns: - output_indices: Padded indices [num_kv_heads, batch_size, topk] - """ - num_kv_heads = input_tensor.shape[0] - batch_size = len(kv_lens) - device = input_tensor.device - - # Compute max sequence length for padding - max_seq_len = kv_lens.max().item() - - # Ensure padding size >= topk - pad_size = max(max_seq_len, topk) - - # Create padded tensor [num_kv_heads, batch_size, pad_size] - padded_tensor = torch.full( - (num_kv_heads, batch_size, pad_size), float("-inf"), dtype=input_tensor.dtype, device=device - ) - - # Fill in actual values - for batch_idx in range(batch_size): - start = kv_offsets[batch_idx].item() - end = kv_offsets[batch_idx + 1].item() - seq_len = kv_lens[batch_idx].item() - - for head_idx in range(num_kv_heads): - padded_tensor[head_idx, batch_idx, :seq_len] = input_tensor[head_idx, start:end] - - # Perform batch topk: [num_kv_heads, batch_size, pad_size] -> [num_kv_heads, batch_size, topk] - topk_values, topk_indices = torch.topk( - padded_tensor, - k=topk, - dim=-1, - largest=True, - ) - - # Mask out invalid indices based on each batch's seq_len - seq_lens_expanded = kv_lens.to(device).unsqueeze(0).unsqueeze(-1) # [1, batch_size, 1] - # topk_indices: [num_kv_heads, batch_size, topk] - mask = topk_indices >= seq_lens_expanded - topk_indices.masked_fill_(mask, -1) - - return topk_indices - - -def triton_topk_wrapper( - input_tensor: torch.Tensor, - kv_offsets: torch.Tensor, - kv_lens: torch.Tensor, - topk: int, -) -> torch.Tensor: - """ - Args: - input_tensor: Input scores [num_kv_heads, sum(kv_lens)] - kv_offsets: KV offsets [batch_size + 1] - kv_lens: KV lengths [batch_size] - topk: TopK parameter - - Returns: - output_indices: Padded indices [num_kv_heads, batch_size, topk] - """ - num_kv_heads = input_tensor.shape[0] - batch_size = len(kv_lens) - device = input_tensor.device - - sparse_lens = torch.tensor( - [min(topk, seq_len.item()) for seq_len in kv_lens], dtype=torch.int32, device=device - ) - sparse_offsets = torch.cat( - [torch.zeros(1, dtype=torch.int32, device=device), torch.cumsum(sparse_lens, dim=0)] - ).to(device) - total_sparse_indices = sparse_offsets[-1].item() - - output_indices_flat = triton_topk( - input_tensor, kv_offsets, sparse_offsets, total_sparse_indices, topk - ) - - # Convert flat format to padded format [num_kv_heads, batch_size, topk] - output_indices_padded = torch.full( - (num_kv_heads, batch_size, topk), -1, dtype=torch.int32, device=device - ) - - for batch_idx in range(batch_size): - start = sparse_offsets[batch_idx].item() - end = sparse_offsets[batch_idx + 1].item() - actual_len = end - start - - for head_idx in range(num_kv_heads): - output_indices_padded[head_idx, batch_idx, :actual_len] = output_indices_flat[ - head_idx, start:end - ] - - return output_indices_padded - - -def compute_overlap_ratio( - triton_indices: torch.Tensor, - reference_indices: torch.Tensor, - kv_lens: torch.Tensor, -) -> float: - """ - Args: - triton_indices: Triton topk results [num_kv_heads, batch_size, topk] - reference_indices: Reference topk results [num_kv_heads, batch_size, topk] - kv_lens: KV lengths [batch_size] - - Returns: - Average overlap ratio across all batches and heads - """ - num_kv_heads = triton_indices.shape[0] - batch_size = triton_indices.shape[1] - - overlap_ratios = [] - - # Compare batch by batch - for batch_idx in range(batch_size): - for head_idx in range(num_kv_heads): - # Extract indices for this batch and head - triton_batch = triton_indices[head_idx, batch_idx, :].cpu().tolist() - reference_batch = reference_indices[head_idx, batch_idx, :].cpu().tolist() - - # Filter out -1 (invalid/padding indices) - triton_set = set([x for x in triton_batch if x >= 0]) - reference_set = set([x for x in reference_batch if x >= 0]) - - if len(reference_set) > 0: - overlap = len(triton_set & reference_set) - overlap_ratio = overlap / len(reference_set) - overlap_ratios.append(overlap_ratio) - - if len(overlap_ratios) == 0: - return 1.0 - - return sum(overlap_ratios) / len(overlap_ratios) - - -@pytest.mark.parametrize( - "batch_size,seq_lens,num_kv_heads,topk", - [ - # Single batch, seq_len > topk - (1, [3000], 1, 2048), - # Single batch, seq_len < topk - (1, [1000], 8, 2048), - # Multiple batches, mixed seq_len (some < topk, some > topk) - (6, [50, 150, 80, 300, 100, 256], 8, 128), - (6, [1000, 2500, 1500, 3000, 1800, 4000], 8, 2048), - ], -) -def test_topk_kernel(batch_size, seq_lens, num_kv_heads, topk): - device = torch.device("cuda") - dtype = torch.float32 - - kv_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) - - kv_offsets = torch.cat( - [torch.zeros(1, dtype=torch.int32, device=device), torch.cumsum(kv_lens, dim=0)] - ).to(device) - - total_tokens = kv_offsets[-1].item() - - input_tensor = torch.randn((num_kv_heads, total_tokens), dtype=dtype, device=device) - - triton_output = triton_topk_wrapper( - input_tensor=input_tensor, - kv_offsets=kv_offsets, - kv_lens=kv_lens, - topk=topk, - ) - - reference_output = pytorch_reference_topk( - input_tensor=input_tensor, - kv_offsets=kv_offsets, - kv_lens=kv_lens, - topk=topk, - ) - - overlap_ratio = compute_overlap_ratio( - triton_output, - reference_output, - kv_lens, - ) - - min_threshold = 0.99 - - print(f"overlap_ratio: {overlap_ratio}") - - assert overlap_ratio >= min_threshold, ( - f"Overlap ratio {overlap_ratio:.4f} is too low (< {min_threshold})" - ) - - -if __name__ == "__main__": - print("\n" + "=" * 80) - print("Testing Triton TopK Kernel") - print("=" * 80) - - # Single batch tests - print("\n--- Single batch, seq_len > topk ---") - test_topk_kernel(1, [3000], 8, 2048) - - print("\n--- Single batch, seq_len < topk ---") - test_topk_kernel(1, [1000], 8, 2048) - - print("\n--- Multiple batches, mixed seq_len ---") - test_topk_kernel(6, [50, 150, 80, 300, 100, 256], 8, 128) - test_topk_kernel(6, [1000, 2500, 1500, 3000, 1800, 4000], 8, 2048) - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py new file mode 100644 index 000000000000..41fa185f926e --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -0,0 +1,765 @@ +import inspect +import json +import struct +import weakref +from copy import deepcopy + +import pytest +import torch +from transformers import PretrainedConfig + +# from utils.util import default_dtype +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.cache_manager import ( + DeepseekV4CacheManager, +) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import Compressor +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4Indexer, + DeepseekV4TrtllmAttention, + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_deepseekv4 import ( + DeepseekV4DecoderLayer, + DeepseekV4ForCausalLM, + DeepseekV4Gate, + DeepseekV4MoE, + DeepseekV4MTP, + _deepseek_v4_pos_embd_params, + _remap_deepseek_v4_checkpoint_keys, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.utils import AuxStreamType, model_extra_attrs +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +DEEPSEEK_V4_TINY_CONFIG = { + "architectures": ["DeepseekV4ForCausalLM"], + "model_type": "deepseek_v4", + "hidden_size": 4096, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "qk_nope_head_dim": 448, + "qk_rope_head_dim": 64, + "v_head_dim": 512, + "q_lora_rank": 1024, + "kv_lora_rank": 448, + "o_groups": 8, + "o_lora_rank": 1024, + "max_position_embeddings": 65536, + "rms_norm_eps": 1e-6, + "dtype": "bfloat16", + "vocab_size": 129280, + "num_hidden_layers": 7, + "n_hash_layers": 3, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "num_experts_per_tok": 6, + "n_group": 1, + "topk_group": 1, + "routed_scaling_factor": 1.5, + "score_func": "sqrtsoftplus", + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "hc_eps": 1e-6, + "compress_rope_theta": 40000.0, + "rope_theta": 10000.0, + "rope_scaling": { + "type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 65536, + "beta_fast": 32, + "beta_slow": 1, + }, + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + }, +} + + +def _write_safetensors_header(path, tensor_name, dtype, shape): + header = { + tensor_name: { + "dtype": dtype, + "shape": shape, + "data_offsets": [0, 0], + } + } + payload = json.dumps(header).encode("utf-8") + path.write_bytes(struct.pack(" torch.Tensor: + x = x.softmax(-1) + eps + x = x / (x.sum(-2, keepdim=True) + eps) + for _ in range(repeat - 1): + x = x / (x.sum(-1, keepdim=True) + eps) + x = x / (x.sum(-2, keepdim=True) + eps) + return x + + +def vanilla_pre_mapping( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + mult: int, + norm_eps: float, + eps: float, + sinkhorn_eps: float, + post_mult_value: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reference pre_mapping implementation in pure PyTorch.""" + assert mult == x.shape[-2] + residual_flat = x.flatten(-2, -1).float() + sqrsum = residual_flat.square().sum(-1) + mixes = residual_flat @ fn.T * (sqrsum.unsqueeze(-1) / fn.shape[-1] + norm_eps).rsqrt() + scale_expanded = torch.cat( + [ + scale[0].expand(mult), + scale[1].expand(mult), + scale[2].expand(mult * mult), + ], + ) + mixes = mixes * scale_expanded + base + pre_mix = mixes[:, :mult].sigmoid().unsqueeze(-1) + eps + post_mix = (mixes[:, mult : 2 * mult].sigmoid() * post_mult_value).unsqueeze(-1) + res_mix = mixes[:, 2 * mult :].view(-1, mult, mult) + res_mix = _sinkhorn_normalize_ref(res_mix, repeat=sinkhorn_iters, eps=sinkhorn_eps) + layer_input = (x * pre_mix).sum(-2).bfloat16() + return post_mix, res_mix, layer_input + + +def vanilla_post_mapping( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, +) -> torch.Tensor: + """Reference post_mapping implementation in pure PyTorch.""" + term2 = torch.bmm(comb_res_mix.mT, residual.float()) + return (x.float().unsqueeze(-2) * post_layer_mix + term2).bfloat16() + + +def vanilla_hc_head( + x: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + norm_eps: float, + eps: float, +) -> torch.Tensor: + """Reference HCHead forward implementation in pure PyTorch.""" + shape, dtype = x.size(), x.dtype + x = x.flatten(-2, -1).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + norm_eps) + mixes = F.linear(x, fn) * rsqrt + pre = torch.sigmoid(mixes * scale + base) + eps + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1) + return y.to(dtype) + + +# --------------------------------------------------------------------------- +# Profiling helpers (from bench_dg_vs_fma_nsys.py) +# --------------------------------------------------------------------------- + + +def profile_fn(fn, warmup=BENCH_WARMUP, iters=BENCH_ITERS): + """Return dict of {kernel_name: avg_us} for all CUDA kernels.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + result = {} + for evt in prof.key_averages(): + if evt.self_device_time_total > 0: + result[evt.key] = evt.self_device_time_total / evt.count + return result + + +def profile_fn_total(fn, warmup: int = BENCH_WARMUP, iters: int = BENCH_ITERS) -> float: + """Return average per-iter kernel time (us) via torch.profiler. + + Sums `self_device_time_total` (microseconds) across every CUDA event + recorded between start() and stop() and divides by `iters`. This + captures the true per-iter kernel time and excludes host-side gaps + between launches (e.g. between post_mapping and pre_mapping in the + unfused path), regardless of how kernel counts differ across paths. + """ + for _ in range(warmup): + fn() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + total_us = sum(evt.self_device_time_total for evt in prof.key_averages()) + return total_us / iters + + +def sum_kernel_times(timings, filters): + """Sum times for kernel names matching any filter substring.""" + total = 0.0 + for name, us in timings.items(): + if any(f in name for f in filters): + total += us + return total + + +def sum_all_kernel_times(timings): + """Sum all GPU kernel times.""" + return sum(timings.values()) + + +# --------------------------------------------------------------------------- +# Test data generators +# --------------------------------------------------------------------------- + + +def generate_pre_data( + n: int, + hc_mult: int, + hidden_size: int, + rms_eps: float = 1e-6, + hc_pre_eps: float = 1e-6, + hc_sinkhorn_eps: float = 1e-6, + hc_post_mult_value: float = 1.0, + sinkhorn_repeat: int = 20, +) -> dict[str, torch.Tensor | float]: + """Generate test data for big fuse operator.""" + torch.random.manual_seed(42) + + hc_mult2 = hc_mult * hc_mult + hc_mult3 = hc_mult * 2 + hc_mult2 + device = "cuda" + + residual = ( + (torch.randn((n, hc_mult, hidden_size), dtype=torch.float, device=device) / hidden_size) + .mul(1 + torch.arange(hc_mult, device=device).mul(0.01).view(1, -1, 1)) + .bfloat16() + ) + + fn = ( + torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float, device=device) + * 1e-4 + * (1 + torch.arange(hc_mult, device=device).mul(0.01).view(1, -1, 1)) + ).flatten(1, 2) + + hc_scale = torch.randn((3,), dtype=torch.float, device=device) * 0.1 + + hc_base = torch.randn((hc_mult3,), dtype=torch.float, device=device) * 0.1 + + return { + "residual": residual, + "fn": fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + "rms_eps": rms_eps, + "hc_pre_eps": hc_pre_eps, + "hc_sinkhorn_eps": hc_sinkhorn_eps, + "hc_post_mult_value": hc_post_mult_value, + "sinkhorn_repeat": sinkhorn_repeat, + } + + +def generate_post_data( + n: int, + hidden_size: int, + hc_mult: int, + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Generate test data for post operator.""" + torch.random.manual_seed(42) + + x = torch.randn((n, hidden_size), dtype=torch.bfloat16, device=device) / hidden_size + residual = torch.randn((n, hc_mult, hidden_size), dtype=torch.bfloat16, device=device) + post_layer_mix = torch.randn((n, hc_mult, 1), dtype=torch.float32, device=device) + comb_res_mix = torch.randn((n, hc_mult, hc_mult), dtype=torch.float32, device=device) + + return { + "x": x, + "residual": residual, + "post_layer_mix": post_layer_mix, + "comb_res_mix": comb_res_mix, + } + + +def generate_head_data( + m: int, + hidden_size: int, + hc_mult: int, + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Generate test data for post operator.""" + torch.random.manual_seed(42) + + x = torch.randn((m, hc_mult, hidden_size), dtype=torch.bfloat16, device=device) / hidden_size + hc_fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32, device=device) + hc_base = torch.randn((hc_mult,), dtype=torch.float32, device=device) + hc_scale = torch.randn((1,), dtype=torch.float32, device=device) + + return { + "x": x, + "hc_fn": hc_fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + } + + +# --------------------------------------------------------------------------- +# Correctness + profiling tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("n", [1, 32, 64, 128, 256, 512, 4096, 8192]) +@pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_pre_mapping(n: int, hidden_size: int, hc_mult: int): + test_data = generate_pre_data( + n=n, + hc_mult=hc_mult, + hidden_size=hidden_size, + ) + + test_module = mHC( + mult=hc_mult, + hidden_size=hidden_size, + sinkhorn_iters=test_data["sinkhorn_repeat"], + dtype=None, + eps=test_data["hc_pre_eps"], + norm_eps=test_data["rms_eps"], + post_mult_value=test_data["hc_post_mult_value"], + ).cuda() + test_module.fn.copy_(test_data["fn"]) + test_module.scale.copy_(test_data["hc_scale"]) + test_module.base.copy_(test_data["hc_base"]) + + residual = test_data["residual"] + + t = profile_fn(lambda: test_module.pre_mapping(residual)) + total_us = sum_all_kernel_times(t) + timing_stats[("pre_mapping", n, hidden_size)]["cuda"] = total_us + + post_mix_cuda, comb_mix_cuda, layer_input_cuda = test_module.pre_mapping(residual) + post_mix_ref, comb_mix_ref, layer_input_ref = vanilla_pre_mapping( + residual, + test_data["fn"], + test_data["hc_scale"], + test_data["hc_base"], + hc_mult, + test_data["rms_eps"], + test_data["hc_pre_eps"], + test_data["hc_sinkhorn_eps"], + test_data["hc_post_mult_value"], + test_data["sinkhorn_repeat"], + ) + torch.testing.assert_close(post_mix_ref, post_mix_cuda, rtol=1e-4, atol=1e-3) + torch.testing.assert_close(comb_mix_ref, comb_mix_cuda, rtol=1e-3, atol=5e-3) + torch.testing.assert_close(layer_input_ref, layer_input_cuda, rtol=1e-4, atol=1e-3) + + +@pytest.mark.parametrize("n", [64, 128, 4096, 8192]) +@pytest.mark.parametrize("hidden_size", [7168]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_post_mapping(n: int, hidden_size: int, hc_mult: int): + test_data = generate_post_data( + n=n, + hc_mult=hc_mult, + hidden_size=hidden_size, + ) + + test_module = mHC(mult=hc_mult, hidden_size=hidden_size, sinkhorn_iters=10) + + t = profile_fn(lambda: test_module.post_mapping(**test_data)) + total_us = sum_all_kernel_times(t) + timing_stats[("post_mapping", n, hidden_size)]["cuda"] = total_us + + output_cuda = test_module.post_mapping(**test_data) + output_ref = vanilla_post_mapping(**test_data) + torch.testing.assert_close(output_ref, output_cuda, rtol=1e-2, atol=0.1) + + +@pytest.mark.parametrize("n", [1, 32, 128, 512, 4096, 8192]) +@pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_fused_hc(n: int, hidden_size: int, hc_mult: int): + """Correctness test for mHC.fused_hc. + + fused_hc(x_prev, residual_prev, post_mix_prev, comb_mix_prev) must be + numerically equivalent to: + residual_cur = post_mapping(x_prev, residual_prev, post_mix_prev, comb_mix_prev) + post_mix_cur, comb_mix_cur, layer_input_cur = pre_mapping(residual_cur) + + Uses two distinct mHC modules so that the 'prev' and 'cur' blocks have + different weights — mirroring the real decoder layer boundary. + """ + # Generate parameters for the 'current' mHC (consumed by pre_mapping part). + pre_data = generate_pre_data(n=n, hc_mult=hc_mult, hidden_size=hidden_size) + + # Generate the incoming (residual_prev, x_prev, post_mix_prev, comb_mix_prev) + # that the 'previous' block would have emitted. + torch.random.manual_seed(7) + device = "cuda" + x_prev = torch.randn((n, hidden_size), dtype=torch.bfloat16, device=device) / hidden_size + residual_prev = ( + torch.randn((n, hc_mult, hidden_size), dtype=torch.float, device=device) / hidden_size + ).bfloat16() + post_mix_prev = torch.randn((n, hc_mult, 1), dtype=torch.float32, device=device) * 0.1 + comb_mix_prev = torch.randn((n, hc_mult, hc_mult), dtype=torch.float32, device=device) * 0.1 + + cur_module = mHC( + mult=hc_mult, + hidden_size=hidden_size, + sinkhorn_iters=pre_data["sinkhorn_repeat"], + dtype=None, + eps=pre_data["hc_pre_eps"], + norm_eps=pre_data["rms_eps"], + post_mult_value=pre_data["hc_post_mult_value"], + ).cuda() + cur_module.fn.copy_(pre_data["fn"]) + cur_module.scale.copy_(pre_data["hc_scale"]) + cur_module.base.copy_(pre_data["hc_base"]) + + # --- fused_hc path --- + ( + residual_cur_f, + post_mix_cur_f, + comb_mix_cur_f, + layer_input_cur_f, + ) = cur_module.fused_hc(x_prev, residual_prev, post_mix_prev, comb_mix_prev) + + # --- two-step reference (post_mapping then pre_mapping via the same module) --- + residual_cur_ref = cur_module.post_mapping(x_prev, residual_prev, post_mix_prev, comb_mix_prev) + post_mix_cur_ref, comb_mix_cur_ref, layer_input_cur_ref = cur_module.pre_mapping( + residual_cur_ref + ) + + # Timing: fused_hc vs separate (post_mapping + pre_mapping). + # Both paths sum every CUDA event's self_device_time_total via + # torch.profiler and divide by the iteration count, so host-side gaps + # between post_mapping and pre_mapping in the unfused path are excluded. + def _unfused(): + residual_cur = cur_module.post_mapping(x_prev, residual_prev, post_mix_prev, comb_mix_prev) + cur_module.pre_mapping(residual_cur) + + fused_us = profile_fn_total( + lambda: cur_module.fused_hc(x_prev, residual_prev, post_mix_prev, comb_mix_prev) + ) + unfused_us = profile_fn_total(_unfused) + + timing_stats[("fused_hc", n, hidden_size)]["cuda"] = fused_us + timing_stats[("fused_hc", n, hidden_size)]["cuda_unfused"] = unfused_us + speedup = (unfused_us / fused_us) if fused_us > 0 else 0.0 + print( + f"[fused_hc benchmark] n={n} hidden={hidden_size} " + f"fused={fused_us:7.2f}us unfused={unfused_us:7.2f}us " + f"speedup={speedup:.2f}x" + ) + + # fused_hc is a Python-level chain of the same kernels that pre_mapping and + # post_mapping use (mhc_post_mapping then the bigfuse pre_mapping pipeline). + # Tolerance matches the baseline post_mapping test (residuals are bf16). + torch.testing.assert_close(residual_cur_ref, residual_cur_f, rtol=1e-2, atol=0.1) + torch.testing.assert_close(post_mix_cur_ref, post_mix_cur_f, rtol=1e-3, atol=5e-3) + torch.testing.assert_close(comb_mix_cur_ref, comb_mix_cur_f, rtol=1e-3, atol=5e-3) + torch.testing.assert_close(layer_input_cur_ref, layer_input_cur_f, rtol=1e-3, atol=5e-3) + + +# Explicit backend coverage. The autotuner picks one tactic per M-bucket at +# warmup; to actually exercise every backend across CI we force each tactic. +# Tactic format mirrors MhcFusedHcRunner: (backend, tile_n, num_k_splits, +# bigfuse_bs, tile_m). +# +# FMA tactics intentionally sweep both ks>1 (cross-CTA atomicAdd into y_acc / +# r_acc) and tile_m>1 (Path F only; multi-token per CTA, which reshapes how +# the atomic accumulation buckets tokens). Keeping ks=1,tm=1 only would leave +# the cross-CTA atomic path uncovered. +_BACKEND_TACTICS_BY_M = { + 64: [ + ("fused_half_mma", 0, 8, 256, 1), + ("fused_all_mma", 0, 1, 0, 1), + ("fused_half_fma", 2, 2, 256, 1), # FMA cross-CTA atomic (ks=2) + ("fused_half_fma", 2, 4, 256, 1), # FMA deeper cross-CTA atomic (ks=4) + ("fused_all_fma", 2, 1, 0, 1), + ("fused_all_fma", 2, 2, 0, 1), # Path F ks=2 atomic + ("fused_all_fma", 2, 1, 0, 2), # Path F tile_m=2 (multi-token CTA) + ], + 256: [ + ("fused_half_mma", 0, 4, 256, 1), + ("fused_all_mma", 0, 1, 0, 1), + ("fused_half_fma", 4, 1, 256, 1), + ("fused_half_fma", 2, 2, 256, 1), # ks=2 atomic + ("fused_all_fma", 4, 1, 0, 1), + ("fused_all_fma", 2, 2, 0, 1), # Path F ks=2 atomic (tn=2 required for ks>1) + ("fused_all_fma", 4, 1, 0, 2), # Path F tile_m=2 + ], + # fused_half_fma is intentionally omitted at M=2048: the runner guards the + # FMA 2-kernel path to M <= 512 (it stops scaling past that M). + 2048: [ + ("fused_half_mma", 0, 2, 128, 1), + ("fused_all_mma", 0, 1, 0, 1), + ("fused_all_fma", 4, 1, 0, 1), + ("fused_all_fma", 2, 2, 0, 1), # Path F ks=2 atomic (tn=2 required for ks>1) + ("fused_all_fma", 4, 1, 0, 4), # Path F tile_m=4 + ], +} + + +@pytest.mark.parametrize("n", list(_BACKEND_TACTICS_BY_M.keys())) +@pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_fused_hc_backends(n: int, hidden_size: int, hc_mult: int): + """Every wired fused_hc backend sees bit-identical input and is checked + against one shared torch reference and one golden backend output. + + Paths D (fused_all_mma) and F (fused_all_fma) are single-kernel + all-in-one variants; fused_half_mma (Path B) and fused_half_fma (Path E) + are the 2-kernel baselines. Each is forced by calling + MhcFusedHcRunner.forward directly with an explicit tactic, bypassing the + autotuner. + + Path C (bigfuse tcgen05) is not covered: its kernel emits (D_next, + sqr_sum_next, layer_input) as the layer-to-layer state carrier and does + not produce post_mix_cur / comb_mix_cur, so it cannot be dropped in + behind the current mhc_fused_hc API without a kernel-side modification + that adds post_mix_out / comb_mix_out stores. + """ + from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner + + pre_data = generate_pre_data(n=n, hc_mult=hc_mult, hidden_size=hidden_size) + + torch.random.manual_seed(13) + device = "cuda" + # Canonical input tensors — generated once, then deep-cloned per consumer + # so the torch ref and each backend each get an independent byte-identical + # copy. Protects the test from any hypothetical in-place mutation inside + # a kernel launcher or a contiguous() call. + x_prev_ref = torch.randn((n, hidden_size), dtype=torch.bfloat16, device=device) / hidden_size + residual_prev_ref = ( + torch.randn((n, hc_mult, hidden_size), dtype=torch.float, device=device) / hidden_size + ).bfloat16() + post_mix_prev_ref = torch.randn((n, hc_mult, 1), dtype=torch.float32, device=device) * 0.1 + comb_mix_prev_ref = torch.randn((n, hc_mult, hc_mult), dtype=torch.float32, device=device) * 0.1 + + cur_module = mHC( + mult=hc_mult, + hidden_size=hidden_size, + sinkhorn_iters=pre_data["sinkhorn_repeat"], + dtype=None, + eps=pre_data["hc_pre_eps"], + norm_eps=pre_data["rms_eps"], + post_mult_value=pre_data["hc_post_mult_value"], + ).cuda() + cur_module.fn.copy_(pre_data["fn"]) + cur_module.scale.copy_(pre_data["hc_scale"]) + cur_module.base.copy_(pre_data["hc_base"]) + + # Torch ground-truth — computed from clones so the ref path cannot + # perturb the canonical input tensors either. + residual_cur_ref = cur_module.post_mapping( + x_prev_ref.clone(), + residual_prev_ref.clone(), + post_mix_prev_ref.clone(), + comb_mix_prev_ref.clone(), + ) + post_mix_ref, comb_mix_ref, layer_input_ref = cur_module.pre_mapping(residual_cur_ref.clone()) + + runner = MhcFusedHcRunner( + n=hc_mult, + hidden_size=hidden_size, + rms_eps=pre_data["rms_eps"], + hc_pre_eps=pre_data["hc_pre_eps"], + hc_sinkhorn_eps=pre_data["hc_sinkhorn_eps"], + hc_post_mult_value=pre_data["hc_post_mult_value"], + sinkhorn_repeat=pre_data["sinkhorn_repeat"], + ) + + def make_runner_inputs(): + return [ + x_prev_ref.clone(), + residual_prev_ref.reshape(n, hc_mult, hidden_size).clone().contiguous(), + post_mix_prev_ref.reshape(n, hc_mult).clone().contiguous(), + comb_mix_prev_ref.reshape(n, hc_mult, hc_mult).clone().contiguous(), + cur_module.fn.detach().clone().contiguous(), + cur_module.scale.detach().clone(), + cur_module.base.detach().clone(), + ] + + backend_outputs = {} + for tactic in _BACKEND_TACTICS_BY_M[n]: + backend = tactic[0] + residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = runner( + inputs=make_runner_inputs(), tactic=tactic + ) + backend_outputs[backend] = ( + residual_cur, + post_mix_cur.view(n, hc_mult, 1), + comb_mix_cur.view(n, hc_mult, hc_mult), + layer_input_cur, + ) + + # Tolerances: bf16 has a 7-bit mantissa so 1 ulp ~ 7.8e-3. For outputs + # near unit scale with fp32-accumulated reductions, rtol=1e-2 atol=1e-2 + # is the expected bf16 parity — tighter than test_mhc_post_mapping's + # atol=0.1 which compared against a pure-bf16 vanilla reference. + bf16_tol = dict(rtol=1e-2, atol=1e-2) + fp32_tol = dict(rtol=1e-3, atol=5e-3) + + # (1) Every backend must match the torch reference. + for backend, ( + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + ) in backend_outputs.items(): + torch.testing.assert_close( + residual_cur_ref, + residual_cur, + **bf16_tol, + msg=f"[vs torch-ref] backend={backend} n={n} residual mismatch", + ) + torch.testing.assert_close( + post_mix_ref, + post_mix_cur, + **fp32_tol, + msg=f"[vs torch-ref] backend={backend} n={n} post_mix mismatch", + ) + torch.testing.assert_close( + comb_mix_ref, + comb_mix_cur, + **fp32_tol, + msg=f"[vs torch-ref] backend={backend} n={n} comb_mix mismatch", + ) + torch.testing.assert_close( + layer_input_ref, + layer_input_cur, + **bf16_tol, + msg=f"[vs torch-ref] backend={backend} n={n} layer_input mismatch", + ) + + # (2) All backends must agree with one golden backend at the same tolerance + # as vs the torch ref. Different backends vary only in tile shape and + # reduction order, so cross-backend divergence would indicate a kernel + # correctness bug rather than expected rounding drift. + gold = "fused_half_mma" if "fused_half_mma" in backend_outputs else next(iter(backend_outputs)) + gr, gpm, gcm, gli = backend_outputs[gold] + for backend, ( + residual_cur, + post_mix_cur, + comb_mix_cur, + layer_input_cur, + ) in backend_outputs.items(): + if backend == gold: + continue + torch.testing.assert_close( + gr, + residual_cur, + **bf16_tol, + msg=f"[vs {gold}] backend={backend} n={n} residual mismatch", + ) + torch.testing.assert_close( + gpm, + post_mix_cur, + **fp32_tol, + msg=f"[vs {gold}] backend={backend} n={n} post_mix mismatch", + ) + torch.testing.assert_close( + gcm, + comb_mix_cur, + **fp32_tol, + msg=f"[vs {gold}] backend={backend} n={n} comb_mix mismatch", + ) + torch.testing.assert_close( + gli, + layer_input_cur, + **bf16_tol, + msg=f"[vs {gold}] backend={backend} n={n} layer_input mismatch", + ) + + +@pytest.mark.parametrize("n", [128, 2048]) +@pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_fused_hc_cuda_graph(n: int, hidden_size: int, hc_mult: int): + """CUDA-graph capture/replay of mHC.fused_hc. + + The decoder uses fused_hc at every non-first layer boundary; the whole + decoder is expected to be traced into a single CUDA graph. This test + verifies that (a) fused_hc can be captured without host syncs, and + (b) replay produces bit-exact results to eager. + + To keep the bit-exact assertion structurally valid, we drive the runner + with an explicit ``num_k_splits=1`` tactic (Path B, fused_half_mma). + That disables split-K atomic accumulation entirely, so none of the four + outputs depend on the non-deterministic FP ordering that pickKSplits(M) + would otherwise introduce (it picks ks=16 at M=128 and ks=4 at M=2048 + for the autotuner fallback — atomics active, not deterministic across + replays). + """ + from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner + + pre_data = generate_pre_data(n=n, hc_mult=hc_mult, hidden_size=hidden_size) + + torch.random.manual_seed(11) + device = "cuda" + x_prev = torch.randn((n, hidden_size), dtype=torch.bfloat16, device=device) / hidden_size + residual_prev = ( + torch.randn((n, hc_mult, hidden_size), dtype=torch.float, device=device) / hidden_size + ).bfloat16() + post_mix_prev = torch.randn((n, hc_mult, 1), dtype=torch.float32, device=device) * 0.1 + comb_mix_prev = torch.randn((n, hc_mult, hc_mult), dtype=torch.float32, device=device) * 0.1 + + cur_module = mHC( + mult=hc_mult, + hidden_size=hidden_size, + sinkhorn_iters=pre_data["sinkhorn_repeat"], + dtype=None, + eps=pre_data["hc_pre_eps"], + norm_eps=pre_data["rms_eps"], + post_mult_value=pre_data["hc_post_mult_value"], + ).cuda() + cur_module.fn.copy_(pre_data["fn"]) + cur_module.scale.copy_(pre_data["hc_scale"]) + cur_module.base.copy_(pre_data["hc_base"]) + + runner = MhcFusedHcRunner( + n=hc_mult, + hidden_size=hidden_size, + rms_eps=pre_data["rms_eps"], + hc_pre_eps=pre_data["hc_pre_eps"], + hc_sinkhorn_eps=pre_data["hc_sinkhorn_eps"], + hc_post_mult_value=pre_data["hc_post_mult_value"], + sinkhorn_repeat=pre_data["sinkhorn_repeat"], + ) + # Pin tactic to Path B with num_k_splits=1 → no atomic accumulation on + # any output. Tactic tuple matches MhcFusedHcRunner.get_tactics(). + tactic = ("fused_half_mma", 0, 1, 128, 1) + assert tactic[2] == 1, "bit-exact assertion requires num_k_splits=1" + + def _inputs(): + return [ + x_prev, + residual_prev.reshape(n, hc_mult, hidden_size).contiguous(), + post_mix_prev.reshape(n, hc_mult).contiguous(), + comb_mix_prev.reshape(n, hc_mult, hc_mult).contiguous(), + cur_module.fn, + cur_module.scale, + cur_module.base, + ] + + # Eager reference — runner's workspace cache reuses output tensors across + # calls with matching shape, so eager_out and graph_out alias the same + # storage. Clone eager_out so we can compare after the graph replay + # overwrites the workspace. + eager_raw = runner(inputs=_inputs(), tactic=tactic) + eager_out = tuple(t.clone() for t in eager_raw) + + # Warm up on a side stream — required for CUDA graph capture. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + runner(inputs=_inputs(), tactic=tactic) + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + + # Capture. + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + graph_out = runner(inputs=_inputs(), tactic=tactic) + + # Replay — outputs should update in place. + g.replay() + torch.cuda.synchronize() + + # With ks=1 the kernel has no atomic accumulation anywhere, so replay must + # be bit-exact against eager on all four outputs. + for ge, ee, name in zip( + graph_out, eager_out, ["residual", "post_mix", "comb_mix", "layer_input"] + ): + torch.testing.assert_close( + ge, ee, rtol=0, atol=0, msg=f"fused_hc CUDA-graph mismatch in {name}" + ) + + # Mutate inputs in-place and replay; result should follow — proves the graph + # is parameterised by input storage, not cached constants. + x_prev.mul_(1.001) + residual_prev.mul_(1.001) + post_mix_prev.mul_(1.001) + comb_mix_prev.mul_(1.001) + eager_raw2 = runner(inputs=_inputs(), tactic=tactic) + eager_out2 = tuple(t.clone() for t in eager_raw2) + g.replay() + torch.cuda.synchronize() + for ge, ee, name in zip( + graph_out, eager_out2, ["residual", "post_mix", "comb_mix", "layer_input"] + ): + torch.testing.assert_close( + ge, ee, rtol=0, atol=0, msg=f"fused_hc CUDA-graph replay mismatch in {name}" + ) + + +@pytest.mark.parametrize("m", [64, 128, 4096, 8192]) +@pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_hc_head(m: int, hidden_size: int, hc_mult: int): + test_data = generate_head_data( + m=m, + hc_mult=hc_mult, + hidden_size=hidden_size, + ) + + test_module = HCHead(mult=hc_mult, hidden_size=hidden_size).cuda() + test_module.fn.copy_(test_data["hc_fn"]) + test_module.scale.copy_(test_data["hc_scale"]) + test_module.base.copy_(test_data["hc_base"]) + + t = profile_fn(lambda: test_module(test_data["x"])) + total_us = sum_all_kernel_times(t) + timing_stats[("hc_head", m, hidden_size)]["cuda"] = total_us + + output_cuda = test_module(test_data["x"]) + output_ref = vanilla_hc_head( + test_data["x"], + test_data["hc_fn"], + test_data["hc_scale"], + test_data["hc_base"], + norm_eps=1e-6, + eps=1e-6, + ) + torch.testing.assert_close(output_ref, output_cuda, rtol=1e-2, atol=0.1) + + +# --------------------------------------------------------------------------- +# Low-level pre_mapping pipeline benchmark: DG / DG-s16 / FMA +# --------------------------------------------------------------------------- + +HC_MULT = 4 +HIDDEN_SIZE = 4096 +_N = HC_MULT * (HC_MULT + 1 + 1) # 24 +_K = HC_MULT * HIDDEN_SIZE # 16384 +_NUM_SPLITS = 16 +_SINKHORN_REPEAT = 20 + + +def _try_import_backends(): + """Return (tf32_hc_prenorm_gemm|None, mhc_gemm_rms_fma_cuda|None, + mhc_big_fuse_cuda|None).""" + tf32_hc_prenorm_gemm = None + try: + from deep_gemm import tf32_hc_prenorm_gemm + except ImportError: + try: + from tensorrt_llm.deep_gemm import tf32_hc_prenorm_gemm + except ImportError: + pass + + mhc_gemm_rms_fma_cuda = mhc_big_fuse_cuda = None + try: + from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( + mhc_big_fuse_cuda, + mhc_gemm_rms_fma_cuda, + ) + except Exception: + pass + + return ( + tf32_hc_prenorm_gemm, + mhc_gemm_rms_fma_cuda, + mhc_big_fuse_cuda, + ) + + +def run_bench_pre_mapping(M: int) -> dict: + """Low-level kernel benchmark for one M: profiles GEMM + BigFuse per backend. + Returns dict like {"DG": (gemm_us, fuse_us), "FMA": (...), ...}. + """ + device = "cuda" + ( + tf32_hc_prenorm_gemm, + mhc_gemm_rms_fma_cuda, + mhc_big_fuse_cuda, + ) = _try_import_backends() + + w_nk = torch.randn(_N, _K, dtype=torch.float32, device=device) * 0.01 + hc_scale = torch.randn(3, dtype=torch.float32, device=device) + hc_base = torch.randn(_N, dtype=torch.float32, device=device) + x = (torch.randn(M, _K, dtype=torch.float32, device=device) * 0.01).bfloat16() + residual = ( + torch.randn(M, HC_MULT, HIDDEN_SIZE, dtype=torch.float32, device=device) / HIDDEN_SIZE + ).bfloat16() + + times = {} + + if tf32_hc_prenorm_gemm is not None and mhc_big_fuse_cuda is not None: + y = torch.empty(M, _N, dtype=torch.float32, device=device) + r = torch.empty(M, dtype=torch.float32, device=device) + pm = torch.empty(M, HC_MULT, dtype=torch.float32, device=device) + cm = torch.empty(M, HC_MULT * HC_MULT, dtype=torch.float32, device=device) + li = torch.empty(M, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + + def dg_fn(): + tf32_hc_prenorm_gemm(x, w_nk, y, r) + mhc_big_fuse_cuda( + y, + r, + residual, + hc_scale, + hc_base, + pm, + cm, + li, + M, + _K, + HIDDEN_SIZE, + 1e-6, + 1e-6, + 1e-6, + 1.0, + _SINKHORN_REPEAT, + num_splits=1, + ) + + t = profile_fn(dg_fn) + times["DG"] = (sum_kernel_times(t, ["hc_prenorm_gemm"]), sum_kernel_times(t, ["BigFuse"])) + + y_s = torch.empty(_NUM_SPLITS, M, _N, dtype=torch.float32, device=device) + r_s = torch.empty(_NUM_SPLITS, M, dtype=torch.float32, device=device) + pm_s = torch.empty(M, HC_MULT, dtype=torch.float32, device=device) + cm_s = torch.empty(M, HC_MULT * HC_MULT, dtype=torch.float32, device=device) + li_s = torch.empty(M, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + + def dg_s16_fn(): + tf32_hc_prenorm_gemm(x, w_nk, y_s, r_s, num_splits=_NUM_SPLITS) + mhc_big_fuse_cuda( + y_s, + r_s, + residual, + hc_scale, + hc_base, + pm_s, + cm_s, + li_s, + M, + _K, + HIDDEN_SIZE, + 1e-6, + 1e-6, + 1e-6, + 1.0, + _SINKHORN_REPEAT, + num_splits=_NUM_SPLITS, + ) + + t = profile_fn(dg_s16_fn) + times["DG-s16"] = ( + sum_kernel_times(t, ["hc_prenorm_gemm"]), + sum_kernel_times(t, ["BigFuse"]), + ) + + if mhc_gemm_rms_fma_cuda is not None and mhc_big_fuse_cuda is not None: + pm_f = torch.empty(M, HC_MULT, dtype=torch.float32, device=device) + cm_f = torch.empty(M, HC_MULT * HC_MULT, dtype=torch.float32, device=device) + li_f = torch.empty(M, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + + def fma_fn(): + y_f, r_f = mhc_gemm_rms_fma_cuda(x, None, M, _N, _K, w_t=w_nk) + mhc_big_fuse_cuda( + y_f, + r_f, + residual, + hc_scale, + hc_base, + pm_f, + cm_f, + li_f, + M, + _K, + HIDDEN_SIZE, + 1e-6, + 1e-6, + 1e-6, + 1.0, + _SINKHORN_REPEAT, + num_splits=1, + ) + + t = profile_fn(fma_fn) + times["FMA"] = (sum_kernel_times(t, ["GemmSqrsumFma"]), sum_kernel_times(t, ["BigFuse"])) + + return times + + +def _print_bench_timing_table(bench_entries: dict): + """Print the pre_mapping pipeline (GEMM + BigFuse) benchmark table.""" + if not bench_entries: + return + all_cols = [] + for v in bench_entries.values(): + for c in v: + if c not in all_cols: + all_cols.append(c) + print("\nPRE_MAPPING PIPELINE (GEMM + BigFuse)") + header = f" {'M':>6s}" + for c in all_cols: + header += f" {c:>16s}" + header += f" {'best':>8s}" + print(header) + print(" " + "-" * (len(header) - 2)) + for key in sorted(bench_entries): + _, M, _ = key + times = bench_entries[key] + totals = {c: times[c][0] + times[c][1] for c in times} + best = min(totals, key=totals.get) if totals else "N/A" + row = f" {M:6d}" + for c in all_cols: + if c in times: + g, f = times[c] + row += f" {g + f:8.1f}({g:4.1f}+{f:4.1f})" + else: + row += f" {'N/A':>16s}" + row += f" {best:>8s}" + print(row) + + +# --------------------------------------------------------------------------- +# Session-scoped fixture: print timing table at end (pytest only) +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="session", autouse=True) +def print_timing_stats(): + """Print collected GPU profiler timings at end of session.""" + yield + + if not timing_stats: + return + + print("\n" + "=" * 90) + print("GPU Kernel Timing (torch.profiler, microseconds)") + print("=" * 90) + + # --- Per-backend correctness/perf tests (pre_mapping, post_mapping, hc_head) --- + for test_type in ("pre_mapping", "post_mapping", "fused_hc", "hc_head"): + entries = { + k: v for k, v in timing_stats.items() if isinstance(k, tuple) and k[0] == test_type + } + if not entries: + continue + + dim_label = "m" if test_type == "hc_head" else "n" + print(f"\n{test_type.upper()}") + + all_backends = sorted({b for d in entries.values() for b in d}) + header = f" {dim_label:>6s} hidden" + for b in all_backends: + header += f" {b:>10s}" + print(header) + print(" " + "-" * (len(header) - 2)) + + for key in sorted(entries): + _, dim_val, hidden = key + row = f" {dim_val:6d} {hidden:6d}" + for b in all_backends: + us = entries[key].get(b) + row += f" {us:10.1f}" if us is not None else f" {'N/A':>10s}" + print(row) + + # --- Low-level pipeline bench table (only populated when run via main()) --- + bench_entries = { + k: v for k, v in timing_stats.items() if isinstance(k, tuple) and k[0] == "bench_pre" + } + _print_bench_timing_table(bench_entries) + + print("\n" + "=" * 90) + + +def main(): + """Run pre_mapping pipeline benchmark (GEMM + BigFuse) for various M. + Invoked when running: python test_mhc.py + """ + torch.manual_seed(42) + bench_M = [1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + bench_stats = {} + for M in bench_M: + bench_stats[("bench_pre", M, HIDDEN_SIZE)] = run_bench_pre_mapping(M) + + print("\n" + "=" * 90) + print("GPU Kernel Timing (torch.profiler, microseconds) — benchmark only") + print("=" * 90) + _print_bench_timing_table(bench_stats) + print("\n" + "=" * 90) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/modules/test_rotary_embedding.py b/tests/unittest/_torch/modules/test_rotary_embedding.py index cc20dee61a45..863f0d9c556b 100644 --- a/tests/unittest/_torch/modules/test_rotary_embedding.py +++ b/tests/unittest/_torch/modules/test_rotary_embedding.py @@ -1,11 +1,11 @@ import pytest import torch +import tensorrt_llm # noqa: F401 - registers torch ops from tensorrt_llm._torch.attention_backend.interface import (RopeParams, RotaryScalingType) from tensorrt_llm._torch.modules.rotary_embedding import (MRotaryEmbedding, RotaryEmbedding) -from tensorrt_llm.functional import RopeEmbeddingUtils class TestRotaryEmbedding: @@ -161,87 +161,183 @@ def test_mrotary_embedding_sanity_1d_position_ids(self, basic_rope_params): assert result[1].shape == k.shape -class TestDuplicateData: - """Tests for duplicate_data support in RoPE table creation.""" - - @pytest.mark.parametrize("dim,max_positions,theta", [ - (64, 4096, 10000.0), - (64, 4096, 1000000.0), - (128, 8192, 10000.0), - ]) - def test_attention_plugin_duplicate_matches_manual_cat( - self, dim, max_positions, theta): - """duplicate_data=True should equal creating without duplication then - manually concatenating, which is what _normalize_mla_rotary_cache_layout did.""" - _, cs_no = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin( - num_pos=max_positions, - dim=dim, - theta=theta, - duplicate_data=False, - ) - _, cs_yes = RopeEmbeddingUtils.create_sinusoidal_positions_for_attention_plugin( - num_pos=max_positions, - dim=dim, - theta=theta, - duplicate_data=True, - ) +def _python_mla_rope_ref(data, position_ids, cos_sin_cache, nope_dim, rope_dim, + inverse, is_neox): + """Python reference for MLA RoPE (neox or interleaved style).""" + num_tokens, num_heads, _ = data.shape + half_rope = rope_dim // 2 + out = data.clone() + for t in range(num_tokens): + pos = position_ids[t].item() + cos_val = cos_sin_cache[pos, 0, :] # [half_rope] + sin_val = cos_sin_cache[pos, 1, :] # [half_rope] + for h in range(num_heads): + if is_neox: + x1 = data[t, h, nope_dim:nope_dim + half_rope].float() + x2 = data[t, h, + nope_dim + half_rope:nope_dim + rope_dim].float() + else: + rope_slice = data[t, h, nope_dim:nope_dim + rope_dim].float() + x1 = rope_slice[0::2] + x2 = rope_slice[1::2] + if inverse: + o1 = x1 * cos_val + x2 * sin_val + o2 = x2 * cos_val - x1 * sin_val + else: + o1 = x1 * cos_val - x2 * sin_val + o2 = x2 * cos_val + x1 * sin_val + if is_neox: + out[t, h, nope_dim:nope_dim + half_rope] = o1.to(out.dtype) + out[t, h, + nope_dim + half_rope:nope_dim + rope_dim] = o2.to(out.dtype) + else: + out[t, h, nope_dim:nope_dim + rope_dim:2] = o1.to(out.dtype) + out[t, h, nope_dim + 1:nope_dim + rope_dim:2] = o2.to(out.dtype) + return out + + +class TestMLARoPEInplace: + """Test suite for the fused mla_rope_inplace torch op.""" - t = torch.tensor(cs_no).view(max_positions, -1, 2) - expected = torch.cat([t, t], dim=1).reshape(1, -1).contiguous() - actual = torch.tensor(cs_yes) - assert actual.shape == expected.shape - torch.testing.assert_close(actual, expected) - - def test_create_rope_const_params_mla_table_size(self): - """With duplicate_data=True and interleave=True, the table should have - dim*2 floats per position (doubled).""" - rp = RopeParams(dim=64, - theta=10000.0, - max_positions=2048, - duplicate_data=True) - _, cos_sin = rp.create_rope_const_params(interleave=True) - floats_per_pos = cos_sin.numel() // rp.max_positions - assert floats_per_pos == rp.dim * 2 - - def test_create_rope_const_params_normal_table_size(self): - """With duplicate_data=False and interleave=True, the table should have - dim floats per position (normal).""" - rp = RopeParams(dim=64, - theta=10000.0, - max_positions=2048, - duplicate_data=False) - _, cos_sin = rp.create_rope_const_params(interleave=True) - floats_per_pos = cos_sin.numel() // rp.max_positions - assert floats_per_pos == rp.dim - - -class TestFromConfigDuplicateData: - """Tests for RopeParams.from_config setting duplicate_data based on qk_rope_head_dim.""" - - @staticmethod - def _make_config(**overrides): - defaults = dict( - hidden_size=4096, - num_attention_heads=32, - max_position_embeddings=4096, - rope_theta=10000.0, - model_type="llama", + @pytest.fixture + def rope_params(self): + return RopeParams( + dim=64, + theta=1000000.0, + alpha=1.0, + scale_type=RotaryScalingType.none, + scale=1.0, + max_positions=32768, + original_max_positions=1024, + beta_fast=32, + beta_slow=1, + mscale=1.0, + mscale_all_dim=0.0, + short_factor=None, + long_factor=None, + max_seq_len=None, + duplicate_data=True, ) - defaults.update(overrides) - class _Cfg: - pass - - cfg = _Cfg() - for k, v in defaults.items(): - setattr(cfg, k, v) - return cfg - - def test_no_qk_rope_head_dim(self): - rp = RopeParams.from_config(self._make_config()) - assert rp.duplicate_data is False + @pytest.mark.parametrize("is_neox", [True, False]) + @pytest.mark.parametrize("inverse", [False, True]) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_mla_rope_inplace_matches_reference(self, rope_params, inverse, + dtype, is_neox): + """Test mla_rope_inplace op matches Python reference implementation.""" + assert torch.cuda.is_available(), "This test requires CUDA" + device = "cuda" + num_tokens = 32 + num_heads = 8 + nope_dim = 128 + rope_dim = 64 + + rope = RotaryEmbedding(rope_params=rope_params, + head_dim=rope_dim, + is_neox=is_neox, + inverse=inverse) + cos_sin_cache = rope.rotary_cos_sin.to(device) # [max_pos, 2, half] + + data = torch.randn(num_tokens, + num_heads, + nope_dim + rope_dim, + dtype=dtype, + device=device) + position_ids = torch.randint(0, + 1024, (num_tokens, ), + dtype=torch.int32, + device=device) + + # Reference + expected = _python_mla_rope_ref(data, position_ids, cos_sin_cache, + nope_dim, rope_dim, inverse, is_neox) + + # Op under test (in-place) + torch.ops.trtllm.mla_rope_inplace(data, position_ids, cos_sin_cache, + num_heads, nope_dim, rope_dim, + inverse, is_neox) + + # nope portion should be unchanged + torch.testing.assert_close(data[:, :, :nope_dim], + expected[:, :, :nope_dim], + atol=0, + rtol=0) + # rope portion should match reference + torch.testing.assert_close(data[:, :, nope_dim:], + expected[:, :, nope_dim:], + atol=1e-2, + rtol=1e-2) + + @pytest.mark.parametrize("is_neox", [True, False]) + def test_mla_rope_inplace_roundtrip(self, rope_params, is_neox): + """Test that forward followed by inverse RoPE recovers original data.""" + assert torch.cuda.is_available(), "This test requires CUDA" + device = "cuda" + num_tokens = 16 + num_heads = 4 + nope_dim = 128 + rope_dim = 64 + dtype = torch.bfloat16 + + rope = RotaryEmbedding(rope_params=rope_params, + head_dim=rope_dim, + is_neox=is_neox) + cos_sin_cache = rope.rotary_cos_sin.to(device) + + data = torch.randn(num_tokens, + num_heads, + nope_dim + rope_dim, + dtype=dtype, + device=device) + original = data.clone() + position_ids = torch.arange(num_tokens, + dtype=torch.int32, + device=device) + + # Forward then inverse should recover original + torch.ops.trtllm.mla_rope_inplace(data, position_ids, cos_sin_cache, + num_heads, nope_dim, rope_dim, False, + is_neox) + torch.ops.trtllm.mla_rope_inplace(data, position_ids, cos_sin_cache, + num_heads, nope_dim, rope_dim, True, + is_neox) + + torch.testing.assert_close(data, original, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("num_tokens,num_heads", [(1, 4), (1, 64), + (16, 64)]) + def test_mla_rope_inplace_edge_cases(self, rope_params, num_tokens, + num_heads): + """Test edge cases: single token, many heads (exercises grid.y > 1).""" + assert torch.cuda.is_available(), "This test requires CUDA" + device = "cuda" + nope_dim = 128 + rope_dim = 64 + dtype = torch.bfloat16 - def test_with_qk_rope_head_dim(self): - rp = RopeParams.from_config(self._make_config(qk_rope_head_dim=64)) - assert rp.duplicate_data is True - assert rp.dim == 64 + rope = RotaryEmbedding(rope_params=rope_params, + head_dim=rope_dim, + is_neox=True) + cos_sin_cache = rope.rotary_cos_sin.to(device) + + data = torch.randn(num_tokens, + num_heads, + nope_dim + rope_dim, + dtype=dtype, + device=device) + expected = _python_mla_rope_ref( + data, torch.arange(num_tokens, dtype=torch.int32, device=device), + cos_sin_cache, nope_dim, rope_dim, False, True) + + torch.ops.trtllm.mla_rope_inplace( + data, torch.arange(num_tokens, dtype=torch.int32, device=device), + cos_sin_cache, num_heads, nope_dim, rope_dim, False, True) + + torch.testing.assert_close(data[:, :, :nope_dim], + expected[:, :, :nope_dim], + atol=0, + rtol=0) + torch.testing.assert_close(data[:, :, nope_dim:], + expected[:, :, nope_dim:], + atol=1e-2, + rtol=1e-2) diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 5d8d5dfef98a..b555a7cee945 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -227,9 +227,10 @@ def generate_seq_lens(batch_size, min_long_seq, num_tokens): @pytest.mark.parametrize("batch_size", [1, 64, 512, 2048]) @pytest.mark.parametrize("next_n", [1, 2]) -@pytest.mark.parametrize("index_topk", [2048, 128]) +@pytest.mark.parametrize("index_topk", [2048, 512, 128]) @pytest.mark.parametrize("num_tokens", [4096, 8192]) -def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens): +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens, compress_ratio): """Verify indexer_topk_decode output matches torch.topk for random logits.""" torch.manual_seed(24) torch.cuda.manual_seed(24) @@ -239,16 +240,30 @@ def test_indexer_topk_decode(batch_size, next_n, index_topk, num_tokens): next_n_offset = torch.arange(num_gen_tokens, device="cuda") % next_n seq_lens = generate_seq_lens(batch_size, index_topk, num_tokens) - row_ends = seq_lens[row_indices] - next_n + next_n_offset + 1 + # Calculate actual KV lengths + actual_kv_lens = seq_lens[row_indices] - next_n + next_n_offset + 1 + + # Apply compression with floor division + row_ends = actual_kv_lens // compress_ratio + + # Generate logits with the compressed size logits = create_random_logits(row_starts, row_ends, torch.float32, 42) indices = torch.empty((num_gen_tokens, index_topk), dtype=torch.int32, device="cuda") - torch.ops.trtllm.indexer_topk_decode(logits, seq_lens, indices, next_n, index_topk) + # Run CUDA implementation with compress_ratio + torch.ops.trtllm.indexer_topk_decode( + logits, seq_lens, indices, next_n, index_topk, compress_ratio=compress_ratio + ) + torch.cuda.synchronize() + # Run reference implementation on compressed row_ends max_row_len = row_ends.max().item() + if max_row_len == 0: + # All rows are empty after compression, skip comparison + return torch_indices = logits.topk(min(index_topk, max_row_len), dim=-1)[1] mask_lo = torch_indices >= 0 mask_hi = (torch_indices - (row_ends - row_starts)[:, None]) < 0 diff --git a/tests/unittest/_torch/thop/serial/test_moe_gate.py b/tests/unittest/_torch/thop/serial/test_moe_gate.py new file mode 100644 index 000000000000..0f5e66ecf0bd --- /dev/null +++ b/tests/unittest/_torch/thop/serial/test_moe_gate.py @@ -0,0 +1,350 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +import os +import sys + +import pytest +import torch +import torch.nn.functional as F + +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) + +import tensorrt_llm # noqa: F401 # Import to load C++ extensions + +# Constants from config (can be adjusted for testing) +N_EXPERTS = 256 +N_EXPERTS_PRO = 384 +TOPK = 6 + + +def pytorch_gate_forward( + scores: torch.Tensor, + bias: torch.Tensor = None, + input_ids: torch.Tensor = None, + tid2eid: torch.Tensor = None, + route_scale: float = 1.5, + is_hash: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Reference PyTorch implementation of gate forward. + + Args: + scores: [batch_size, n_experts] - pre-computed scores + bias: [n_experts] - expert bias (used in non-hash mode) + input_ids: [batch_size] - token IDs (used in hash mode) + tid2eid: [vocab_size, topk] - token to expert mapping (used in hash mode) + route_scale: scalar to multiply final weights + is_hash: whether to use hash routing or topk routing + + Returns: + weights: [batch_size, topk] - normalized routing weights + indices: [batch_size, topk] - selected expert indices + """ + + # Apply score function: softplus + sqrt + scores = F.softplus(scores).sqrt() + original_scores = scores + + # Add bias for topk selection (non-hash mode) + if not is_hash and bias is not None: + scores = scores + bias.float() + + # Select experts + if is_hash: + # Hash mode: directly lookup from tid2eid + indices = tid2eid[input_ids] # [batch_size, topk] + else: + # Topk mode: select top-k experts + indices = scores.topk(TOPK, dim=-1)[1] # [batch_size, topk] + + # Gather original weights (without bias) + weights = original_scores.gather(1, indices.long()) + + # Normalize weights + weights = weights / weights.sum(dim=-1, keepdim=True) + + # Apply route scale + weights = weights * route_scale + + return weights, indices.int() + + +class TestMoeGate: + """Test suite for MOE gate forward operation.""" + + @pytest.mark.parametrize("batch_size", [1, 4, 32, 128]) + @pytest.mark.parametrize("route_scale", [1.0, 1.5, 2.0]) + def test_gate_forward_non_hash(self, batch_size, route_scale): + """Test non-hash (topk) mode gate forward.""" + # Generate random scores + scores = torch.randn(batch_size, N_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.randn(N_EXPERTS, dtype=torch.float32, device="cuda") + + # Pre-allocate outputs + out_weights = torch.empty(batch_size, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(batch_size, TOPK, dtype=torch.int32, device="cuda") + + # Empty tensors for hash mode parameters + input_ids_tensor = torch.empty(0, dtype=torch.int32, device="cuda") + tid2eid_tensor = torch.empty(0, dtype=torch.int32, device="cuda") + + # Call CUDA kernel + torch.ops.trtllm.gate_forward( + scores, + bias, + input_ids_tensor, + tid2eid_tensor, + out_weights, + out_indices, + TOPK, + route_scale, + False, + ) + + # Compute reference + ref_weights, ref_indices = pytorch_gate_forward( + scores, bias=bias, route_scale=route_scale, is_hash=False + ) + + # Verify shapes + assert out_weights.shape == (batch_size, TOPK), ( + f"weights shape mismatch: {out_weights.shape}" + ) + assert out_indices.shape == (batch_size, TOPK), ( + f"indices shape mismatch: {out_indices.shape}" + ) + + # Verify values + assert torch.equal(out_indices, ref_indices), ( + f"Indices mismatch:\nCUDA: {out_indices}\nRef: {ref_indices}" + ) + assert torch.allclose(out_weights, ref_weights, rtol=1e-4, atol=1e-5), ( + f"Weights mismatch (max diff: {(out_weights - ref_weights).abs().max():.6f})" + ) + + @pytest.mark.parametrize("batch_size", [1, 4, 32, 128]) + @pytest.mark.parametrize("route_scale", [1.0, 1.5, 2.0]) + @pytest.mark.parametrize("vocab_size", [1000, 10000]) + def test_gate_forward_hash(self, batch_size, route_scale, vocab_size): + """Test hash mode gate forward.""" + # Generate random scores + scores = torch.randn(batch_size, N_EXPERTS, dtype=torch.float32, device="cuda") + input_ids = torch.randint(0, vocab_size, (batch_size,), dtype=torch.int32, device="cuda") + tid2eid = torch.randint(0, N_EXPERTS, (vocab_size, TOPK), dtype=torch.int32, device="cuda") + + # Pre-allocate outputs + out_weights = torch.empty(batch_size, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(batch_size, TOPK, dtype=torch.int32, device="cuda") + + # Empty tensors for non-hash mode parameters + bias = torch.empty(0, dtype=torch.float32, device="cuda") + + # Call CUDA kernel + torch.ops.trtllm.gate_forward( + scores, bias, input_ids, tid2eid, out_weights, out_indices, TOPK, route_scale, True + ) + + # Compute reference + ref_weights, ref_indices = pytorch_gate_forward( + scores, input_ids=input_ids, tid2eid=tid2eid, route_scale=route_scale, is_hash=True + ) + + # Verify shapes + assert out_weights.shape == (batch_size, TOPK), ( + f"weights shape mismatch: {out_weights.shape}" + ) + assert out_indices.shape == (batch_size, TOPK), ( + f"indices shape mismatch: {out_indices.shape}" + ) + + # Verify values + assert torch.equal(out_indices, ref_indices), ( + f"Indices mismatch:\nCUDA: {out_indices}\nRef: {ref_indices}" + ) + assert torch.allclose(out_weights, ref_weights, rtol=1e-4, atol=1e-5), ( + f"Weights mismatch (max diff: {(out_weights - ref_weights).abs().max():.6f})" + ) + + @pytest.mark.parametrize("is_hash", [False, True]) + def test_gate_forward_384_experts(self, is_hash): + """Test DeepSeek-V4-Pro expert count.""" + batch_size = 8 + route_scale = 1.5 + scores = torch.randn(batch_size, N_EXPERTS_PRO, dtype=torch.float32, device="cuda") + out_weights = torch.empty(batch_size, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(batch_size, TOPK, dtype=torch.int32, device="cuda") + + if is_hash: + vocab_size = 1024 + input_ids = torch.randint(0, vocab_size, (batch_size,), dtype=torch.int32, device="cuda") + tid2eid = torch.randint(0, N_EXPERTS_PRO, (vocab_size, TOPK), dtype=torch.int32, device="cuda") + bias = torch.empty(0, dtype=torch.float32, device="cuda") + ref_weights, ref_indices = pytorch_gate_forward( + scores, input_ids=input_ids, tid2eid=tid2eid, route_scale=route_scale, is_hash=True + ) + else: + bias = torch.randn(N_EXPERTS_PRO, dtype=torch.float32, device="cuda") + input_ids = torch.empty(0, dtype=torch.int32, device="cuda") + tid2eid = torch.empty(0, dtype=torch.int32, device="cuda") + ref_weights, ref_indices = pytorch_gate_forward( + scores, bias=bias, route_scale=route_scale, is_hash=False + ) + + torch.ops.trtllm.gate_forward( + scores, bias, input_ids, tid2eid, out_weights, out_indices, TOPK, route_scale, is_hash + ) + + assert torch.equal(out_indices, ref_indices) + assert torch.allclose(out_weights, ref_weights, rtol=1e-4, atol=1e-5) + + def test_gate_forward_deterministic(self): + """Test that gate forward is deterministic across multiple runs.""" + batch_size = 16 + scores = torch.randn(batch_size, N_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.randn(N_EXPERTS, dtype=torch.float32, device="cuda") + + # Run multiple times + results_weights = [] + results_indices = [] + for _ in range(5): + out_weights = torch.empty(batch_size, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(batch_size, TOPK, dtype=torch.int32, device="cuda") + + torch.ops.trtllm.gate_forward( + scores, + bias, + torch.empty(0, dtype=torch.int32, device="cuda"), + torch.empty(0, dtype=torch.int32, device="cuda"), + out_weights, + out_indices, + TOPK, + 1.5, + False, + ) + + results_weights.append(out_weights.clone()) + results_indices.append(out_indices.clone()) + + # Verify all results are identical + for i in range(1, len(results_weights)): + assert torch.equal(results_indices[0], results_indices[i]), ( + "Gate forward is not deterministic (indices differ)" + ) + assert torch.equal(results_weights[0], results_weights[i]), ( + "Gate forward is not deterministic (weights differ)" + ) + + def test_gate_forward_edge_cases(self): + """Test edge cases.""" + # Test with single token + scores = torch.randn(1, N_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.randn(N_EXPERTS, dtype=torch.float32, device="cuda") + out_weights = torch.empty(1, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(1, TOPK, dtype=torch.int32, device="cuda") + + torch.ops.trtllm.gate_forward( + scores, + bias, + torch.empty(0, dtype=torch.int32, device="cuda"), + torch.empty(0, dtype=torch.int32, device="cuda"), + out_weights, + out_indices, + TOPK, + 1.5, + False, + ) + + # Verify weights sum to approximately route_scale + assert torch.allclose( + out_weights.sum(dim=-1), torch.tensor([1.5], device="cuda"), rtol=1e-4 + ) + + # Test with all equal scores (tie-breaking should be deterministic) + scores = torch.ones(4, N_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.zeros(N_EXPERTS, dtype=torch.float32, device="cuda") + out_weights = torch.empty(4, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(4, TOPK, dtype=torch.int32, device="cuda") + + torch.ops.trtllm.gate_forward( + scores, + bias, + torch.empty(0, dtype=torch.int32, device="cuda"), + torch.empty(0, dtype=torch.int32, device="cuda"), + out_weights, + out_indices, + TOPK, + 1.0, + False, + ) + + # With equal scores, weights should be equal + assert torch.allclose(out_weights, out_weights[0:1].expand(4, TOPK), rtol=1e-4) + + @pytest.mark.parametrize("batch_size", [32, 128, 256, 1024]) + def test_gate_forward_performance(self, batch_size): + """Benchmark gate forward for different batch sizes.""" + scores = torch.randn(batch_size, N_EXPERTS, dtype=torch.float32, device="cuda") + bias = torch.randn(N_EXPERTS, dtype=torch.float32, device="cuda") + out_weights = torch.empty(batch_size, TOPK, dtype=torch.float32, device="cuda") + out_indices = torch.empty(batch_size, TOPK, dtype=torch.int32, device="cuda") + + # Warmup + for _ in range(10): + torch.ops.trtllm.gate_forward( + scores, + bias, + torch.empty(0, dtype=torch.int32, device="cuda"), + torch.empty(0, dtype=torch.int32, device="cuda"), + out_weights, + out_indices, + TOPK, + 1.5, + False, + ) + + torch.cuda.synchronize() + + # Benchmark + import time + + n_runs = 100 + start = time.perf_counter() + for _ in range(n_runs): + torch.ops.trtllm.gate_forward( + scores, + bias, + torch.empty(0, dtype=torch.int32, device="cuda"), + torch.empty(0, dtype=torch.int32, device="cuda"), + out_weights, + out_indices, + TOPK, + 1.5, + False, + ) + torch.cuda.synchronize() + end = time.perf_counter() + + avg_time_ms = (end - start) * 1000.0 / n_runs + print(f"\nBatch size {batch_size}: {avg_time_ms:.4f} ms per iteration") + + # Basic sanity check that it completed successfully + assert out_weights.shape == (batch_size, TOPK) + assert out_indices.shape == (batch_size, TOPK) + + +if __name__ == "__main__": + # Run tests with pytest + pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 87a4a53d2b62..db179f03b1fc 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -228,7 +228,7 @@ methods: default: null status: prototype sparse_attention_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekV4SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, NoneType] default: null status: prototype otlp_traces_endpoint: diff --git a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py new file mode 100644 index 000000000000..a95525fe64fd --- /dev/null +++ b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py @@ -0,0 +1,97 @@ +# 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. + +from tensorrt_llm.inputs.utils import apply_chat_template +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer + + +class _DummyTokenizer: + all_special_tokens = [] + eos_token_id = 1 + pad_token_id = 0 + name_or_path = "dummy" + + def encode(self, text, *args, **kwargs): + self.last_encoded_text = text + return [1, 2, 3] + + +def test_deepseek_v4_chat_template_matches_reference_single_user_prompt(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "Question: 1+1?\nAnswer:", + } + ], + tokenize=False, + add_generation_prompt=True, + ) + + assert prompt == ( + "<｜begin▁of▁sentence｜><｜User｜>Question: 1+1?\nAnswer:<｜Assistant｜>" + ) + + +def test_deepseek_v4_chat_template_tokenize_uses_rendered_prompt(): + dummy = _DummyTokenizer() + tokenizer = DeepseekV4Tokenizer(dummy) + + token_ids = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=True, + add_generation_prompt=True, + ) + + assert token_ids == [1, 2, 3] + assert dummy.last_encoded_text == ( + "<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>" + ) + + +def test_deepseek_v4_custom_tokenizer_reuses_loaded_wrapper(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + args = TorchLlmArgs(model="dummy", tokenizer=tokenizer, custom_tokenizer="deepseek_v4") + + assert args.tokenizer is tokenizer + + +def test_deepseek_v4_server_chat_template_path_uses_custom_tokenizer(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = apply_chat_template( + model_type="deepseek_v4", + tokenizer=tokenizer, + processor=None, + conversation=[ + { + "role": "user", + "content": "hello", + } + ], + add_generation_prompt=True, + mm_placeholder_counts=[{}], + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") From 60ccd524bc88ded4f9e56720072726ef7594d368 Mon Sep 17 00:00:00 2001 From: Yuewei Na <248773860+nv-yna@users.noreply.github.com> Date: Wed, 29 Apr 2026 00:31:07 -0700 Subject: [PATCH 02/58] [TRTLLM-12338][feat] Lift TOKENIZER_ALIASES to module level in llmapi.llm_args (#13568) Signed-off-by: Yuewei Na Co-authored-by: Yuewei Na Signed-off-by: Fanrong Li (cherry picked from commit 085e2e19583d677ec8636fec16979ded536c9042) Signed-off-by: Fanrong Li --- tensorrt_llm/llmapi/llm_args.py | 33 ++++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0f8c9b3cc477..9721c8a17079 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -623,9 +623,7 @@ def needs_separate_short_long_cuda_graphs(self) -> bool: class DeepSeekV4SparseAttentionConfig(DeepSeekSparseAttentionConfig): - """ - Configuration for DeepSeek-V4 Sparse Attention. - """ + """Configuration for DeepSeek-V4 Sparse Attention.""" algorithm: Literal["deepseek_v4"] = "deepseek_v4" skip_indexer_for_short_seqs: bool = Field( default=False, @@ -871,6 +869,7 @@ class MoeConfig(StrictBaseModel): # Maps alias → full import path (module.ClassName). TOKENIZER_ALIASES = { 'deepseek_v32': 'tensorrt_llm.tokenizer.deepseek_v32.DeepseekV32Tokenizer', + 'deepseek_v4': 'tensorrt_llm.tokenizer.deepseek_v4.DeepseekV4Tokenizer', } @@ -3588,15 +3587,27 @@ def validate_and_init_tokenizer(self): "Please specify a tokenizer path or leave it as None to load from model path." ) - from tensorrt_llm.tokenizer import load_custom_tokenizer + # Resolve short aliases via the module-level TOKENIZER_ALIASES. + tokenizer_path = TOKENIZER_ALIASES.get(self.custom_tokenizer, + self.custom_tokenizer) - # Use tokenizer path if specified, otherwise use model path - load_path = self.tokenizer if self.tokenizer else self.model - self.tokenizer = load_custom_tokenizer( - self.custom_tokenizer, - load_path, - trust_remote_code=self.trust_remote_code, - use_fast=self.tokenizer_mode != 'slow') + # Dynamically import and use custom tokenizer + from importlib import import_module + try: + module_path, class_name = tokenizer_path.rsplit('.', 1) + module = import_module(module_path) + tokenizer_class = getattr(module, class_name) + # Use tokenizer path if specified, otherwise use model path + load_path = self.tokenizer if self.tokenizer else self.model + self.tokenizer = tokenizer_class.from_pretrained( + load_path, + trust_remote_code=self.trust_remote_code, + use_fast=self.tokenizer_mode != 'slow') + except (ValueError, ImportError, AttributeError) as e: + raise ValueError( + f"Failed to load custom tokenizer '{self.custom_tokenizer}': {e}. " + "Expected format: 'module.path.ClassName' or a recognized alias." + ) from e else: self.tokenizer = tokenizer_factory( self.tokenizer, From fe2f384dc582d2408ebb5e216f1f9cf3e4427984 Mon Sep 17 00:00:00 2001 From: Mingyang Date: Wed, 29 Apr 2026 16:01:02 +0800 Subject: [PATCH 03/58] [None][fix] Fix fused mHC RMS normalization (#13587) Signed-off-by: Mingyang Hao Signed-off-by: Fanrong Li (cherry picked from commit 3303f1c3c8a1665b4cdff952215bb0c0675cc40c) Signed-off-by: Fanrong Li --- .../mhcKernels/fused_tf32_pmap_gemm.cuh | 3 +- .../kernels/mhcKernels/mhcFusedHcKernel.cu | 6 +- .../_torch/models/modeling_deepseekv4.py | 18 +-- .../modeling/test_modeling_deepseekv4.py | 17 +++ tests/unittest/_torch/modules/test_mhc.py | 111 ++++++++++++++++++ 5 files changed, 142 insertions(+), 13 deletions(-) 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 24e9d3d656a5..ad9b10fe2e55 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -1207,8 +1207,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) for (uint32_t c = 0; c < HC_MULT3; ++c) y_local[c] = y_row[c]; - // Reference bigfuse divides by HIDDEN (per-head count), matching Path A. - float const rstd = rsqrtf(r_val / static_cast(HIDDEN) + rms_eps); + float const rstd = rsqrtf(r_val / static_cast(HC_MULT * HIDDEN) + rms_eps); float const s0 = hc_scale[0]; float const s1 = hc_scale[1]; float const s2 = hc_scale[2]; diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index e85c868509fa..b085c0635adb 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -257,8 +257,8 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual // Delegate to mhcBigFuseLaunch (defined in mhcKernels.cu) to avoid // instantiating the mhcBigFuseKernel template in this TU. mhcBigFuseLaunch(y_acc_workspace, r_acc_workspace, residual_cur, hc_scale, hc_base, post_mix_cur, comb_mix_cur, - layer_input_cur, M, /*K=*/hidden_size, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, - sinkhorn_repeat, /*num_splits=*/1, /*block_size=*/bs, stream); + layer_input_cur, M, /*K=*/static_cast(SHAPE_K), hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, + hc_post_mult_value, sinkhorn_repeat, /*num_splits=*/1, /*block_size=*/bs, stream); } // =================================================================== @@ -336,7 +336,7 @@ void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* resid // ---- Step 2: big-fuse postlogue (reduces ks splits internally) ---- mhcBigFuseLaunch(y_acc_workspace, r_acc_workspace, residual_cur, hc_scale, hc_base, post_mix_cur, comb_mix_cur, - layer_input_cur, M, /*K=*/hidden_size, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, + layer_input_cur, M, /*K=*/K, hidden_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, /*num_splits=*/num_k_splits, bigfuse_block_size, stream); } diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 350d5956c99e..edfd51876076 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -206,6 +206,14 @@ def moe_reduce_add_shared_output(routed_output, shared_output): } +def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool: + """Resolve the DeepSeek-V4 fused HC boundary-fusion knob.""" + env = os.environ.get("TRTLLM_MHC_ENABLE_FUSED_HC") + if env is not None: + return env not in ("0", "false", "False") + return bool(getattr(config, "enable_fused_hc", True)) + + def _remap_deepseek_v4_checkpoint_keys( weights: Dict, num_hidden_layers: int, kv_lora_rank: int = 448 ) -> Dict: @@ -1757,14 +1765,8 @@ def __init__( # the MHC boundary fusion (`mHC.fused_hc`) is used. When False, fall back # to the unfused `post_mapping → pre_mapping` chain (same path engram # layers already take). Env var TRTLLM_MHC_ENABLE_FUSED_HC overrides the - # config attr (set to "1" to force-enable while validating fused_hc). - # Default is disabled because fused_hc must be bit-equivalent to the - # explicit post_mapping -> pre_mapping chain before it can safely replace it. - _env = os.environ.get("TRTLLM_MHC_ENABLE_FUSED_HC") - if _env is not None: - self.enable_fused_hc = _env not in ("0", "false", "False") - else: - self.enable_fused_hc = bool(getattr(config, "enable_fused_hc", False)) + # config attr (set to "0" to force-disable for validation/rollback). + self.enable_fused_hc = _resolve_enable_fused_hc(config) self.next_layer_layernorm: RMSNorm = None # Finalized in DeepseekV4ForCausalLM.post_load_weights once the full layer # list is visible: a layer may defer its hc_ffn.post_mapping only if diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 41fa185f926e..07e2b75a5a84 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -31,6 +31,7 @@ DeepseekV4MTP, _deepseek_v4_pos_embd_params, _remap_deepseek_v4_checkpoint_keys, + _resolve_enable_fused_hc, ) from tensorrt_llm._torch.modules.linear import TensorParallelMode from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig @@ -116,6 +117,22 @@ def test_deepseek_v4_config_aliases(): assert config.swiglu_limit == 9.0 +def test_deepseek_v4_fused_hc_default_enabled(monkeypatch): + monkeypatch.delenv("TRTLLM_MHC_ENABLE_FUSED_HC", raising=False) + config = PretrainedConfig() + + assert _resolve_enable_fused_hc(config) is True + + config.enable_fused_hc = False + assert _resolve_enable_fused_hc(config) is False + + monkeypatch.setenv("TRTLLM_MHC_ENABLE_FUSED_HC", "1") + assert _resolve_enable_fused_hc(config) is True + + monkeypatch.setenv("TRTLLM_MHC_ENABLE_FUSED_HC", "0") + assert _resolve_enable_fused_hc(config) is False + + def test_deepseek_v4_model_defaults_keep_tokens_per_block(): class LlmArgs: pass diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 7eed34ffdb18..7cd05ae0554a 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -196,6 +196,43 @@ def generate_pre_data( } +def generate_realistic_pre_data( + n: int, + hc_mult: int, + hidden_size: int, + rms_eps: float = 1e-6, + hc_pre_eps: float = 1e-6, + hc_sinkhorn_eps: float = 1e-6, + hc_post_mult_value: float = 1.0, + sinkhorn_repeat: int = 20, +) -> dict[str, torch.Tensor | float]: + """Generate real-scale mHC data to catch RMS denominator regressions.""" + torch.random.manual_seed(123) + + hc_mult2 = hc_mult * hc_mult + hc_mult3 = hc_mult * 2 + hc_mult2 + device = "cuda" + + residual = torch.randn((n, hc_mult, hidden_size), dtype=torch.float, device=device).bfloat16() + fn = ( + torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float, device=device) * 0.05 + ).flatten(1, 2) + hc_scale = torch.tensor([0.10, 0.10, 0.30], dtype=torch.float, device=device) + hc_base = torch.randn((hc_mult3,), dtype=torch.float, device=device) * 2.0 + + return { + "residual": residual, + "fn": fn, + "hc_scale": hc_scale, + "hc_base": hc_base, + "rms_eps": rms_eps, + "hc_pre_eps": hc_pre_eps, + "hc_sinkhorn_eps": hc_sinkhorn_eps, + "hc_post_mult_value": hc_post_mult_value, + "sinkhorn_repeat": sinkhorn_repeat, + } + + def generate_post_data( n: int, hidden_size: int, @@ -611,6 +648,80 @@ def make_runner_inputs(): ) +@pytest.mark.parametrize( + "tactic", + [ + ("fused_half_mma", 0, 1, 256, 1), + ("fused_half_fma", 2, 1, 256, 1), + ("fused_all_mma", 0, 1, 0, 1), + ("fused_all_fma", 2, 1, 0, 1), + ], +) +def test_mhc_fused_hc_realistic_scale_regression(tactic): + """Real-scale mHC data catches fused_hc RMS normalization regressions.""" + from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner + + n = 16 + hc_mult = 4 + hidden_size = 4096 + pre_data = generate_realistic_pre_data(n=n, hc_mult=hc_mult, hidden_size=hidden_size) + + torch.random.manual_seed(17) + device = "cuda" + x_prev = torch.randn((n, hidden_size), dtype=torch.float, device=device).bfloat16() + residual_prev = torch.randn( + (n, hc_mult, hidden_size), dtype=torch.float, device=device + ).bfloat16() + post_mix_prev = torch.randn((n, hc_mult, 1), dtype=torch.float32, device=device) * 0.1 + comb_mix_prev = torch.randn((n, hc_mult, hc_mult), dtype=torch.float32, device=device) * 0.1 + + cur_module = mHC( + mult=hc_mult, + hidden_size=hidden_size, + sinkhorn_iters=pre_data["sinkhorn_repeat"], + dtype=None, + eps=pre_data["hc_pre_eps"], + norm_eps=pre_data["rms_eps"], + post_mult_value=pre_data["hc_post_mult_value"], + ).cuda() + cur_module.fn.copy_(pre_data["fn"]) + cur_module.scale.copy_(pre_data["hc_scale"]) + cur_module.base.copy_(pre_data["hc_base"]) + + residual_cur_ref = cur_module.post_mapping(x_prev, residual_prev, post_mix_prev, comb_mix_prev) + post_mix_ref, comb_mix_ref, layer_input_ref = cur_module.pre_mapping(residual_cur_ref) + + runner = MhcFusedHcRunner( + n=hc_mult, + hidden_size=hidden_size, + rms_eps=pre_data["rms_eps"], + hc_pre_eps=pre_data["hc_pre_eps"], + hc_sinkhorn_eps=pre_data["hc_sinkhorn_eps"], + hc_post_mult_value=pre_data["hc_post_mult_value"], + sinkhorn_repeat=pre_data["sinkhorn_repeat"], + ) + + residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = runner( + inputs=[ + x_prev.contiguous(), + residual_prev.contiguous(), + post_mix_prev.view(n, hc_mult).contiguous(), + comb_mix_prev.contiguous(), + cur_module.fn.detach().contiguous(), + cur_module.scale.detach().contiguous(), + cur_module.base.detach().contiguous(), + ], + tactic=tactic, + ) + + torch.testing.assert_close(residual_cur_ref, residual_cur, rtol=1e-2, atol=2e-2) + torch.testing.assert_close(post_mix_ref, post_mix_cur.view(n, hc_mult, 1), rtol=3e-3, atol=5e-3) + torch.testing.assert_close( + comb_mix_ref, comb_mix_cur.view(n, hc_mult, hc_mult), rtol=3e-3, atol=5e-3 + ) + torch.testing.assert_close(layer_input_ref, layer_input_cur, rtol=1e-2, atol=2e-2) + + @pytest.mark.parametrize("n", [128, 2048]) @pytest.mark.parametrize("hidden_size", [4096]) @pytest.mark.parametrize("hc_mult", [4]) From 4d76ccea8da9cb5e7b592a1ec341cfc3c576e390 Mon Sep 17 00:00:00 2001 From: Jiagan Cheng Date: Wed, 29 Apr 2026 18:13:03 +0800 Subject: [PATCH 04/58] [None][fix] Cherry pick KV Cache Manager V2 fixes (#13597) Signed-off-by: Jiagan Cheng Signed-off-by: Fanrong Li (cherry picked from commit 7a7935caaf19aee1d7eddcfcab19d808e8e3a833) Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/pyexecutor/_util.py | 82 ++++++++++++++----------- 1 file changed, 46 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 4f6ec92bd6cd..5d177a9a6a9f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -226,8 +226,6 @@ def __init__( self._max_beam_width = max_beam_width self._kv_connector_manager = kv_connector_manager self._llm_args = llm_args - # For V2 fallback use only, will be removed after V2 is stable - self._cache_transceiver_config = llm_args.cache_transceiver_config self._speculative_config = speculative_config self._sparse_attention_config = sparse_attention_config self._tokens_per_block = tokens_per_block @@ -250,13 +248,38 @@ def _get_model_kv_cache_manager_cls( ): kv_cache_config = (kv_cache_config_override if kv_cache_config_override is not None else self._kv_cache_config) - config = model_engine.model.model_config.pretrained_config cls = get_kv_cache_manager_cls( model_engine.model.model_config, kv_cache_config, is_disagg=self._is_disagg, cache_transceiver_config=self._cache_transceiver_config) - if cls == KVCacheManagerV2: + cls = self._fallback_if_unsupported_kv_cache_manager_v2( + cls, model_engine.model.model_config, kv_cache_config) + # The V1-route hybrid mamba managers (disagg, TRTLLM_USE_CPP_MAMBA, + # TRTLLM_USE_PY_MAMBA, or one-model speculative decoding) keep mamba + # state in a separate cache that doesn't honor block reuse. Warn at + # the routing site so users see the warning where the decision is + # actually made. + if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ + and kv_cache_config.enable_block_reuse: + uses_v1_mamba_route = self._is_disagg \ + or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ + or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \ + or self._speculative_config is not None + if uses_v1_mamba_route: + logger.warning( + "Block reuse does not work with MTP for hybrid linear models " + "when using the legacy MambaCacheManager (TRTLLM_USE_CPP_MAMBA=1)" + ) + return cls + + def _fallback_if_unsupported_kv_cache_manager_v2( + self, + kv_cache_manager_cls, + model_config: Optional[ModelConfig] = None, + kv_cache_config: Optional[KvCacheConfig] = None): + kv_cache_config = kv_cache_config or self._kv_cache_config + if kv_cache_manager_cls == KVCacheManagerV2: if self._kv_connector_manager is not None or ( self._max_beam_width is not None and self._max_beam_width > 1) or kv_cache_config.event_buffer_max_size > 0 or ( @@ -267,7 +290,8 @@ def _get_model_kv_cache_manager_cls( # to max(head_dim), changing per-layer KV byte sizes — which # breaks correctness, not just efficiency. Fail fast here # rather than silently producing wrong outputs. - if is_gemma4_hybrid(config): + if (model_config is not None + and is_gemma4_hybrid(model_config.pretrained_config)): raise NotImplementedError( "Gemma4 hybrid attention requires KVCacheManagerV2, " "which is not yet supported with kv_connector_manager, " @@ -278,24 +302,8 @@ def _get_model_kv_cache_manager_cls( "KVCacheManagerV2 is not supported with kv_connector_manager, beam width > 1, " "event buffer max size > 0, or cache transceiver. Falling back to KVCacheManager." ) - cls = KVCacheManager - # The V1-route hybrid mamba managers (disagg, TRTLLM_USE_CPP_MAMBA, - # TRTLLM_USE_PY_MAMBA, or one-model speculative decoding) keep mamba - # state in a separate cache that doesn't honor block reuse. Warn at - # the routing site so users see the warning where the decision is - # actually made. - if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ - and kv_cache_config.enable_block_reuse: - uses_v1_mamba_route = self._is_disagg \ - or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ - or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \ - or self._speculative_config is not None - if uses_v1_mamba_route: - logger.warning( - "Block reuse does not work with MTP for hybrid linear models " - "when using the legacy MambaCacheManager (TRTLLM_USE_CPP_MAMBA=1)" - ) - return cls + return KVCacheManager + return kv_cache_manager_cls def _per_manager_cache_cost(self, manager_cls, @@ -912,18 +920,8 @@ def _create_one_model_draft_kv_cache_manager( # Get the appropriate KV cache manager class for the draft model draft_kv_cache_manager_cls = get_kv_cache_manager_cls( effective_draft_config, draft_kv_config, is_disagg=self._is_disagg) - - # Use V2 if enabled and the base class is KVCacheManager - if draft_kv_cache_manager_cls == KVCacheManagerV2: - if self._kv_connector_manager is not None or ( - self._max_beam_width is not None and self._max_beam_width - > 1) or draft_kv_config.event_buffer_max_size > 0 or ( - self._cache_transceiver_config is not None - and self._cache_transceiver_config.backend is not None): - logger.warning( - "KVCacheManagerV2 is not supported with disaggregated serving or beam width > 1 or event buffer max size > 0 or disagg config. " - "Falling back to KVCacheManager for draft model.") - draft_kv_cache_manager_cls = KVCacheManager + draft_kv_cache_manager_cls = self._fallback_if_unsupported_kv_cache_manager_v2( + draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config) estimating_kv_cache = estimating_kv_cache and not self._skip_est # For MTP with models using sparse attention (e.g., DeepSeek V3 with DSA), @@ -2053,6 +2051,18 @@ def create_py_executor_instance( if cross_kv_cache_manager is not None else LlmRequestState.CONTEXT_INIT) + # V2 scheduler uses scheduler_capacity as the per-iteration request + # budget (BudgetTracker.max_num_requests). Unlike V1 which has a + # separate CapacityScheduler (needs pp_size * max_batch_size to hold + # requests across PP stages) and MicroBatchScheduler (uses + # max_batch_size for per-forward batch limit), V2 merges both into + # one loop. PP on-the-fly is handled by inflight_request_ids + # filtering, so its budget should be based on max_batch_size, not + # max_num_sequences (which includes the pp_size multiplier). + v2_scheduler_capacity = max_batch_size + if v2_scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: + v2_scheduler_capacity += 1 + if isinstance(kv_cache_manager, KVCacheManagerV2): # V2: interleaved scheduler handles both capacity and budget draft_kv_cache_manager = resources.get( @@ -2068,7 +2078,7 @@ def create_py_executor_instance( ctx_chunk_config=ctx_chunk_config, peft_cache_manager=peft_cache_manager.impl if peft_cache_manager is not None else None, - scheduler_capacity=scheduler_capacity, + scheduler_capacity=v2_scheduler_capacity, draft_kv_cache_manager=draft_kv_cache_manager, cross_kv_cache_manager=cross_kv_cache_manager, no_schedule_until_state=no_schedule_until_state, From ab27f49c4ccff10721f6da2d2e2209733514e212 Mon Sep 17 00:00:00 2001 From: Mingyang Date: Wed, 29 Apr 2026 19:18:58 +0800 Subject: [PATCH 05/58] [None][fix] Remove MHC fused hidden-size guard (#13611) Signed-off-by: Mingyang Hao Co-authored-by: Mingyang Hao Signed-off-by: Fanrong Li (cherry picked from commit b3f45bb608aecca666a451ca5138b81470487f05) Signed-off-by: Fanrong Li --- cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu | 4 ---- 1 file changed, 4 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index b085c0635adb..7954e17c127e 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -209,8 +209,6 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual if (M <= 0) return; - TLLM_CHECK_WITH_INFO(hidden_size == static_cast(FHC_HIDDEN), - "mhcFusedHcLaunch: hidden_size=%d not supported (only %u)", hidden_size, FHC_HIDDEN); TLLM_CHECK_WITH_INFO(hc_mult == static_cast(FHC_HC_MULT), "mhcFusedHcLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT); @@ -393,8 +391,6 @@ void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* if (M <= 0) return; - TLLM_CHECK_WITH_INFO(hidden_size == static_cast(FHC_HIDDEN), - "mhcFusedHcAllInOneLaunch: hidden_size=%d not supported (only %u)", hidden_size, FHC_HIDDEN); TLLM_CHECK_WITH_INFO(hc_mult == static_cast(FHC_HC_MULT), "mhcFusedHcAllInOneLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT); From 299130923424b23d6bd1e00e0498a5f5f33275f4 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Wed, 29 Apr 2026 19:25:53 +0800 Subject: [PATCH 06/58] [None][feat] Enable EPLB for DeepSeek-V4 (#13595) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit eb855283764abb263f79b12f8c81e7e7be4fe1af) Signed-off-by: Fanrong Li --- .../_torch/models/modeling_deepseekv4.py | 32 ++++-- .../modules/fused_moe/moe_load_balancer.py | 1 + .../tokenizer/deepseek_v4/tokenizer.py | 22 ++-- .../defs/accuracy/test_llm_api_pytorch.py | 106 ++++++++++++++++++ .../test_lists/test-db/l0_dgx_b200.yml | 4 + 5 files changed, 140 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index edfd51876076..8529e6b05af6 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -63,7 +63,16 @@ from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.engram import Engram, EngramConfig, EngramHashProvider -from ..modules.fused_moe import DeepSeekV4MoeRoutingMethod, MoE, MoEWeightLoadingMode, create_moe +from ..modules.fused_moe import ( + CutlassFusedMoE, + DeepSeekV4MoeRoutingMethod, + MoE, + MoEWeightLoadingMode, + TritonFusedMoE, + TRTLLMGenFusedMoE, + create_moe, + get_moe_cls, +) from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE from ..modules.linear import Linear from ..modules.mhc.hyper_connection import HCHead, HCState, mHC @@ -1386,17 +1395,16 @@ def __init__( swiglu_limit = getattr(config, "swiglu_limit", None) moe_swiglu_limit = None if swiglu_limit is not None: - supports_swiglu_limit = True - if (model_config.moe_backend or "").upper() == "TRTLLM": - supports_swiglu_limit = False - if experts_quant_config is not None: - mode = experts_quant_config.layer_quant_mode - supports_swiglu_limit = ( - mode.has_nvfp4() - or mode.has_w4a16_mxfp4() - or mode.has_w4a8_mxfp4_fp8() - or mode.has_w4a8_mxfp4_mxfp8() - ) + # `create_moe` only accepts swiglu_limit for these MoE classes; + # resolve via get_moe_cls so backend-string fallbacks (e.g. + # TRTLLM/CUTEDSL/DENSEGEMM dropping back to CutlassFusedMoE on + # unsupported quant) are handled correctly. + moe_cls = get_moe_cls(model_config, override_quant_config=experts_quant_config) + supports_swiglu_limit = moe_cls in ( + CutlassFusedMoE, + TritonFusedMoE, + TRTLLMGenFusedMoE, + ) if supports_swiglu_limit: moe_load_balancer_config = getattr(model_config, "moe_load_balancer", None) num_slots = ( diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py index d87a55ad6651..d225d2ee1fba 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_load_balancer.py @@ -1076,6 +1076,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): moe_model_arch_list = [ 'DeepseekV3ForCausalLM', 'DeepseekV32ForCausalLM', + 'DeepseekV4ForCausalLM', 'GlmMoeDsaForCausalLM', 'GptOssForCausalLM', 'MixtralForCausalLM', diff --git a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py index 7e022ea27158..9149d229dd55 100644 --- a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py +++ b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py @@ -20,11 +20,11 @@ from ..tokenizer import TransformersTokenizer -BOS_TOKEN = "<｜begin▁of▁sentence｜>" -EOS_TOKEN = "<｜end▁of▁sentence｜>" -USER_TOKEN = "<｜User｜>" -ASSISTANT_TOKEN = "<｜Assistant｜>" -THINKING_END_TOKEN = "" +BOS_TOKEN = "<｜begin▁of▁sentence｜>" # nosec B105 +EOS_TOKEN = "<｜end▁of▁sentence｜>" # nosec B105 +USER_TOKEN = "<｜User｜>" # nosec B105 +ASSISTANT_TOKEN = "<｜Assistant｜>" # nosec B105 +THINKING_END_TOKEN = "" # nosec B105 def _message_content_to_text(content: Any) -> str: @@ -66,8 +66,7 @@ def from_pretrained( def apply_chat_template(self, messages, tools=None, **kwargs): if tools: - raise NotImplementedError( - "DeepSeek-V4 tool-call chat formatting is not supported yet.") + raise NotImplementedError("DeepSeek-V4 tool-call chat formatting is not supported yet.") add_generation_prompt = kwargs.get("add_generation_prompt", True) tokenize = kwargs.get("tokenize", False) @@ -76,21 +75,18 @@ def apply_chat_template(self, messages, tools=None, **kwargs): for idx, message in enumerate(messages): role = message.get("role") content = _message_content_to_text(message.get("content")) - next_role = (messages[idx + 1].get("role") - if idx + 1 < len(messages) else None) + next_role = messages[idx + 1].get("role") if idx + 1 < len(messages) else None if role == "system": rendered += content elif role in ("user", "developer"): rendered += USER_TOKEN + content - if next_role == "assistant" or (next_role is None - and add_generation_prompt): + if next_role == "assistant" or (next_role is None and add_generation_prompt): rendered += ASSISTANT_TOKEN + THINKING_END_TOKEN elif role == "assistant": rendered += content + EOS_TOKEN else: - raise NotImplementedError( - f"Unsupported DeepSeek-V4 message role: {role}") + raise NotImplementedError(f"Unsupported DeepSeek-V4 message role: {role}") if tokenize: return self.encode(rendered, add_special_tokens=False) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 41e3d93944ea..0bff4d63eca8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3609,6 +3609,112 @@ def test_nvfp4_multi_gpus_chunked_prefill(self, tp_size, pp_size, ep_size, task.evaluate(llm) +def _make_deepseekv4_eplb_config(model_path, layer_updates_per_iter, ep_size=8): + """Build a MoeLoadBalancerConfig for DeepSeek V4 from the HF config. + + All V4 layers run MoE (no first_k_dense_replace prefix), so static + assignments cover every layer in 0..num_hidden_layers-1. Extra slots + per rank match the TestNemotronV3Super pattern (16 redundant per rank). + """ + with open(f"{model_path}/config.json") as f: + cfg = json.load(f) + num_experts = cfg["n_routed_experts"] + num_slots = num_experts + 16 * ep_size + if layer_updates_per_iter > 0: + return MoeLoadBalancerConfig( + num_slots=num_slots, layer_updates_per_iter=layer_updates_per_iter) + num_hidden_layers = cfg["num_hidden_layers"] + initial_global_assignments = { + i: [(i + j) % num_experts for j in range(num_slots)] + for i in range(num_hidden_layers) + } + return MoeLoadBalancerConfig( + num_slots=num_slots, + initial_global_assignments=initial_global_assignments, + layer_updates_per_iter=0) + + +def _run_deepseekv4_eplb(model_name, + model_path, + moe_backend, + eplb_config, + mtp_nextn=0): + # V4 (~284B) plus EPLB redundant slots leaves little headroom; keep the + # KV-cache fraction conservative and clamp max_seq_len so the default + # 1M context doesn't pre-allocate KV blocks beyond what 180GB B200s have. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) + pytorch_config = dict(cuda_graph_config=CudaGraphConfig(), + moe_config=MoeConfig(backend=moe_backend, + load_balancer=eplb_config)) + mtp_config = (MTPDecodingConfig( + num_nextn_predict_layers=mtp_nextn) if mtp_nextn > 0 else None) + with LLM(model_path, + tensor_parallel_size=8, + moe_expert_parallel_size=8, + kv_cache_config=kv_cache_config, + enable_attention_dp=True, + max_seq_len=4096, + **pytorch_config, + speculative_config=mtp_config) as llm: + task = GSM8K(model_name) + task.evaluate(llm) + + +@pytest.mark.timeout(14400) +@pytest.mark.skip_less_device_memory(140000) +@skip_pre_blackwell +class TestDeepSeekV4Flash(LlmapiAccuracyTestHarness): + MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash" + MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash" + + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("moe_backend", ["WIDEEP", "TRTLLM"]) + def test_nvfp4_8gpus_static_eplb(self, moe_backend): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=0) + _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, + eplb_config) + + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("moe_backend", ["WIDEEP", "TRTLLM"]) + @parametrize_with_ids("mtp_nextn", [0, 1]) + def test_nvfp4_8gpus_online_eplb(self, moe_backend, mtp_nextn): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=2) + _run_deepseekv4_eplb(self.MODEL_NAME, + self.MODEL_PATH, + moe_backend, + eplb_config, + mtp_nextn=mtp_nextn) + + +@pytest.mark.timeout(14400) +@pytest.mark.skip_less_device_memory(140000) +@skip_pre_blackwell +class TestDeepSeekV4FlashBase(LlmapiAccuracyTestHarness): + MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash-Base" + MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash-Base" + + # CUTLASS is omitted: V4 Flash-Base FP8 block-scale weights take a + # Hopper-only kernel path (CutlassFp8BlockScaleGemmRunner::moeGemm) that + # fails on Blackwell. WIDEEP avoids that path and works on B200. + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("moe_backend", ["WIDEEP"]) + def test_fp8_8gpus_static_eplb(self, moe_backend): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=0) + _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, + eplb_config) + + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("moe_backend", ["WIDEEP"]) + def test_fp8_8gpus_online_eplb(self, moe_backend): + eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, + layer_updates_per_iter=2) + _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, + eplb_config) + + @pytest.mark.timeout(10800) @pytest.mark.skip_less_device_memory(100000) class TestKimiK2(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index acb60cffd0d1..cfb24ce83727 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -177,6 +177,10 @@ l0_dgx_b200: - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=3-block_reuse=True-use_py_transceiver=False] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_auto_dtype[mtp_nextn=0-block_reuse=False-use_py_transceiver=False] TIMEOUT (60) + # DeepSeek-V4 EPLB pre-merge sanity (uncomment once DeepSeek-V4-Flash/Flash-Base + # checkpoints are staged under llm_models_root()). + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_8gpus_static_eplb[moe_backend=WIDEEP] TIMEOUT (120) + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_8gpus_static_eplb[moe_backend=WIDEEP] TIMEOUT (120) - accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4 TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] TIMEOUT (60) - accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] TIMEOUT (60) From 0f0afd3de5d442d4c8c7275849fecf338fd66cf3 Mon Sep 17 00:00:00 2001 From: "Qi Zhang (qizh)" <10434017+Tracin@users.noreply.github.com> Date: Wed, 29 Apr 2026 20:00:07 +0800 Subject: [PATCH 07/58] [None][feat] tool template and parser for dsv4 (#13608) Signed-off-by: Tracin <10434017+Tracin@users.noreply.github.com> Signed-off-by: Qi Zhang (qizh) <10434017+Tracin@users.noreply.github.com> Co-authored-by: OpenAI Codex Signed-off-by: Fanrong Li (cherry picked from commit eb97997cf336f0b852cac3f37d7c291dcf2727f0) Signed-off-by: Fanrong Li --- tensorrt_llm/commands/serve.py | 2 +- tensorrt_llm/llmapi/reasoning_parser.py | 50 +++ .../serve/tool_parser/deepseekv32_parser.py | 2 +- .../serve/tool_parser/deepseekv4_parser.py | 25 ++ .../serve/tool_parser/tool_parser_factory.py | 3 + .../tokenizer/deepseek_v4/tokenizer.py | 406 +++++++++++++++++- .../unittest/llmapi/apps/test_tool_parsers.py | 55 +++ .../llmapi/test_deepseek_v4_tokenizer.py | 352 +++++++++++++++ .../unittest/llmapi/test_reasoning_parser.py | 28 ++ 9 files changed, 900 insertions(+), 23 deletions(-) create mode 100644 tensorrt_llm/serve/tool_parser/deepseekv4_parser.py diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 3327af237ee8..26d593d2b3cd 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -971,7 +971,7 @@ def serve( f"Cannot auto-detect reasoning parser for model '{model}'. " f"Supported model types for auto-detection: qwen3, qwen3_moe, " f"qwen3_5, qwen3_5_moe, qwen3_next, deepseek_v3 (R1 only), " - f"deepseek_v32 (R1 only), nemotron_h, gemma4, " + f"deepseek_v32 (R1 only), deepseek_v4, nemotron_h, gemma4, " f"kimi_k2, kimi_k25. " f"Please specify a parser explicitly: " f"{list(ReasoningParserFactory.keys())}", diff --git a/tensorrt_llm/llmapi/reasoning_parser.py b/tensorrt_llm/llmapi/reasoning_parser.py index 439dccceca66..de0abf9e43ef 100644 --- a/tensorrt_llm/llmapi/reasoning_parser.py +++ b/tensorrt_llm/llmapi/reasoning_parser.py @@ -95,6 +95,19 @@ def finish(self) -> ReasoningParserResult: return ReasoningParserResult() +class IdentityReasoningParser(BaseReasoningParser): + """Reasoning parser that treats all model output as visible content.""" + + reasoning_start = "" + reasoning_end = "" + + def parse(self, text: str) -> ReasoningParserResult: + return ReasoningParserResult(content=text) + + def parse_delta(self, delta_text: str) -> ReasoningParserResult: + return ReasoningParserResult(content=delta_text) + + @register_reasoning_parser("deepseek-r1", reasoning_at_start=True) @register_reasoning_parser("laguna") @register_reasoning_parser("qwen3") @@ -198,6 +211,42 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult: "Unreachable code reached in `DeepSeekR1Parser.parse_delta`") +@register_reasoning_parser("deepseek_v4") +class DeepSeekV4ReasoningParser(BaseReasoningParser): + """DeepSeek-V4 parser selected by thinking-mode chat template kwargs.""" + + reasoning_start = "" + reasoning_end = "" + + def __init__( + self, + *, + chat_template_kwargs: Optional[dict[str, Any]] = None, + ) -> None: + super().__init__(chat_template_kwargs=chat_template_kwargs) + chat_template_kwargs = chat_template_kwargs or {} + thinking = bool( + chat_template_kwargs.get("thinking", False) + or chat_template_kwargs.get("enable_thinking", False)) + if thinking: + self._parser = DeepSeekR1Parser( + reasoning_at_start=True, + chat_template_kwargs=chat_template_kwargs, + ) + else: + self._parser = IdentityReasoningParser( + chat_template_kwargs=chat_template_kwargs) + + def parse(self, text: str) -> ReasoningParserResult: + return self._parser.parse(text) + + def parse_delta(self, delta_text: str) -> ReasoningParserResult: + return self._parser.parse_delta(delta_text) + + def finish(self) -> ReasoningParserResult: + return self._parser.finish() + + MODEL_TYPE_TO_REASONING_PARSER: dict[str, str] = { "qwen3": "qwen3", "qwen3_moe": "qwen3", @@ -207,6 +256,7 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult: "deepseek_v3": "deepseek-r1", "deepseek_v32": "deepseek-r1", "laguna": "laguna", + "deepseek_v4": "deepseek_v4", "nemotron_h": "nemotron-v3", "nemotron_h_puzzle": "nemotron-v3", "gemma4": "gemma4", diff --git a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py index 25c49ae2a2cb..3362ab7cb68a 100644 --- a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py +++ b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py @@ -133,7 +133,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult try: # Extract content between function_calls tags function_calls_match = re.search( - r"<｜DSML｜function_calls>(.*?)", + re.escape(self.bot_token) + r"(.*?)" + re.escape(self.eot_token), text, re.DOTALL, ) diff --git a/tensorrt_llm/serve/tool_parser/deepseekv4_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv4_parser.py new file mode 100644 index 000000000000..05896561f839 --- /dev/null +++ b/tensorrt_llm/serve/tool_parser/deepseekv4_parser.py @@ -0,0 +1,25 @@ +# 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. + +from .deepseekv32_parser import DeepSeekV32Parser + + +class DeepSeekV4Parser(DeepSeekV32Parser): + """Tool parser for the DeepSeek V4 DSML tool call format.""" + + def __init__(self) -> None: + super().__init__() + self.bot_token = "<｜DSML｜tool_calls>" # nosec B105 + self.eot_token = "" # nosec B105 diff --git a/tensorrt_llm/serve/tool_parser/tool_parser_factory.py b/tensorrt_llm/serve/tool_parser/tool_parser_factory.py index a0f54fdb2558..42e0ab8ad881 100644 --- a/tensorrt_llm/serve/tool_parser/tool_parser_factory.py +++ b/tensorrt_llm/serve/tool_parser/tool_parser_factory.py @@ -4,6 +4,7 @@ from .base_tool_parser import BaseToolParser from .deepseekv3_parser import DeepSeekV3Parser +from .deepseekv4_parser import DeepSeekV4Parser from .deepseekv31_parser import DeepSeekV31Parser from .deepseekv32_parser import DeepSeekV32Parser from .gemma4_parser import Gemma4ToolParser @@ -24,6 +25,7 @@ "qwen3_next": "qwen3", "deepseek_v3": "deepseek_v3", "deepseek_v32": "deepseek_v32", + "deepseek_v4": "deepseek_v4", "kimi_k2": "kimi_k2", "kimi_k25": "kimi_k2", "glm4": "glm4", @@ -57,6 +59,7 @@ class ToolParserFactory: "deepseek_v3": DeepSeekV3Parser, "deepseek_v31": DeepSeekV31Parser, "deepseek_v32": DeepSeekV32Parser, + "deepseek_v4": DeepSeekV4Parser, "gemma4": Gemma4ToolParser, "glm4": Glm4ToolParser, "glm47": Glm47ToolParser, diff --git a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py index 9149d229dd55..a5600816b072 100644 --- a/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py +++ b/tensorrt_llm/tokenizer/deepseek_v4/tokenizer.py @@ -12,7 +12,10 @@ # 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. +# ruff: noqa: E501 +import copy +import json from pathlib import Path from typing import Any @@ -24,7 +27,64 @@ EOS_TOKEN = "<｜end▁of▁sentence｜>" # nosec B105 USER_TOKEN = "<｜User｜>" # nosec B105 ASSISTANT_TOKEN = "<｜Assistant｜>" # nosec B105 +LATEST_REMINDER_TOKEN = "<｜latest_reminder｜>" # nosec B105 +THINKING_START_TOKEN = "" # nosec B105 THINKING_END_TOKEN = "" # nosec B105 +DSML_TOKEN = "｜DSML｜" # nosec B105 + +TOOL_CALLS_BLOCK_NAME = "tool_calls" +VALID_TASKS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} + +REASONING_EFFORT_MAX = ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose " + "the problem to resolve the root cause, rigorously stress-testing your " + "logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every " + "intermediate step, considered alternative, and rejected hypothesis to " + "ensure absolutely no assumption is left unchecked.\n\n" +) + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +RESPONSE_FORMAT_TEMPLATE = ( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" +) +TOOL_CALL_TEMPLATE = '<{dsml_token}invoke name="{name}">\n{arguments}\n' +TOOL_CALLS_TEMPLATE = "<{dsml_token}{block_name}>\n{tool_calls}\n" +TOOL_OUTPUT_TEMPLATE = "{content}" def _message_content_to_text(content: Any) -> str: @@ -43,6 +103,310 @@ def _message_content_to_text(content: Any) -> str: return str(content) +def _to_json(value: Any) -> str: + try: + return json.dumps(value, ensure_ascii=False) + except TypeError: + return json.dumps(value, ensure_ascii=True) + + +def _tools_from_openai_format(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [tool["function"] for tool in tools] + + +def _tool_calls_from_openai_format(tool_calls: list[dict[str, Any]]) -> list[dict[str, Any]]: + return [ + { + "name": tool_call["function"]["name"], + "arguments": tool_call["function"]["arguments"], + } + for tool_call in tool_calls + ] + + +def _encode_arguments_to_dsml(tool_call: dict[str, Any]) -> str: + raw_args = tool_call["arguments"] + arguments = json.loads(raw_args) if isinstance(raw_args, str) else raw_args + if not isinstance(arguments, dict): + raise ValueError("DeepSeek-V4 tool call arguments must be a JSON object.") + + parameters = [] + for key, value in arguments.items(): + parameters.append( + f'<{DSML_TOKEN}parameter name="{key}" ' + f'string="{"true" if isinstance(value, str) else "false"}">' + f"{value if isinstance(value, str) else _to_json(value)}" + ) + return "\n".join(parameters) + + +def _render_tools(tools: list[dict[str, Any]]) -> str: + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(_to_json(tool) for tool in tools), + dsml_token=DSML_TOKEN, + thinking_start_token=THINKING_START_TOKEN, + thinking_end_token=THINKING_END_TOKEN, + ) + + +def _find_last_user_index(messages: list[dict[str, Any]]) -> int: + for index in range(len(messages) - 1, -1, -1): + if messages[index].get("role") in ("user", "developer"): + return index + return -1 + + +def _merge_tool_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + merged: list[dict[str, Any]] = [] + for message in messages: + message = copy.deepcopy(message) + role = message.get("role") + + if role == "tool": + tool_block = { + "type": "tool_result", + "tool_use_id": message.get("tool_call_id", ""), + "content": message.get("content", ""), + } + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]: + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append({"role": "user", "content_blocks": [tool_block]}) + elif role == "user": + text_block = {"type": "text", "text": _message_content_to_text(message.get("content"))} + if ( + merged + and merged[-1].get("role") == "user" + and "content_blocks" in merged[-1] + and merged[-1].get("task") is None + ): + merged[-1]["content_blocks"].append(text_block) + else: + new_message = { + "role": "user", + "content": message.get("content", ""), + "content_blocks": [text_block], + } + for key in ("task", "wo_eos", "mask"): + if key in message: + new_message[key] = message[key] + merged.append(new_message) + else: + merged.append(message) + + return merged + + +def _sort_tool_results_by_call_order(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + last_tool_call_order: dict[str, int] = {} + + for message in messages: + if message.get("role") == "assistant" and message.get("tool_calls"): + last_tool_call_order = {} + for index, tool_call in enumerate(message["tool_calls"]): + tool_call_id = tool_call.get("id") or tool_call.get("function", {}).get("id", "") + if tool_call_id: + last_tool_call_order[tool_call_id] = index + elif message.get("role") == "user" and message.get("content_blocks"): + tool_blocks = [ + block for block in message["content_blocks"] if block.get("type") == "tool_result" + ] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda block: last_tool_call_order.get(block.get("tool_use_id", ""), 0), + ) + sorted_index = 0 + new_blocks = [] + for block in message["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_index]) + sorted_index += 1 + else: + new_blocks.append(block) + message["content_blocks"] = new_blocks + + return messages + + +def _drop_thinking_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + last_user_index = _find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for index, message in enumerate(messages): + role = message.get("role") + if role in keep_roles or index >= last_user_index: + result.append(message) + elif role == "assistant": + message_without_reasoning = copy.copy(message) + message_without_reasoning.pop("reasoning", None) + message_without_reasoning.pop("reasoning_content", None) + result.append(message_without_reasoning) + + return result + + +def _render_user_content(message: dict[str, Any]) -> str: + content_blocks = message.get("content_blocks") + if not content_blocks: + return _message_content_to_text(message.get("content")) + + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(str(block.get("text", ""))) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for item in tool_content: + if isinstance(item, dict) and item.get("type") == "text": + text_parts.append(str(item.get("text", ""))) + elif isinstance(item, dict): + text_parts.append(f"[Unsupported {item.get('type')}]") + else: + text_parts.append(str(item)) + tool_content = "\n\n".join(text_parts) + parts.append(TOOL_OUTPUT_TEMPLATE.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + return "\n\n".join(parts) + + +def _render_message( + index: int, + messages: list[dict[str, Any]], + thinking_mode: str, + drop_thinking: bool, + add_generation_prompt: bool, + reasoning_effort: str | None, +) -> str: + if thinking_mode not in ("chat", "thinking"): + raise ValueError(f"Invalid thinking_mode: {thinking_mode}") + + message = messages[index] + last_user_index = _find_last_user_index(messages) + role = message.get("role") + content = _message_content_to_text(message.get("content")) + tools = message.get("tools") + response_format = message.get("response_format") + tool_calls = message.get("tool_calls") + reasoning = message.get("reasoning") or message.get("reasoning_content") or "" + prompt = "" + + if tools: + tools = _tools_from_openai_format(tools) + if tool_calls: + tool_calls = _tool_calls_from_openai_format(tool_calls) + + if index == 0 and thinking_mode == "thinking" and reasoning_effort == "max": + prompt += REASONING_EFFORT_MAX + + if role == "system": + prompt += content + if tools: + prompt += "\n\n" + _render_tools(tools) + if response_format: + prompt += "\n\n" + RESPONSE_FORMAT_TEMPLATE.format(schema=_to_json(response_format)) + elif role == "developer": + prompt += USER_TOKEN + content + if tools: + prompt += "\n\n" + _render_tools(tools) + if response_format: + prompt += "\n\n" + RESPONSE_FORMAT_TEMPLATE.format(schema=_to_json(response_format)) + elif role == "user": + prompt += USER_TOKEN + _render_user_content(message) + elif role == "latest_reminder": + prompt += LATEST_REMINDER_TOKEN + content + elif role == "tool": + raise NotImplementedError( + "DeepSeek-V4 merges tool messages into user messages; " + "preprocess with _merge_tool_messages()." + ) + elif role == "assistant": + tool_calls_content = "" + if tool_calls: + rendered_tool_calls = [ + TOOL_CALL_TEMPLATE.format( + dsml_token=DSML_TOKEN, + name=tool_call.get("name"), + arguments=_encode_arguments_to_dsml(tool_call), + ) + for tool_call in tool_calls + ] + tool_calls_content += "\n\n" + TOOL_CALLS_TEMPLATE.format( + dsml_token=DSML_TOKEN, + block_name=TOOL_CALLS_BLOCK_NAME, + tool_calls="\n".join(rendered_tool_calls), + ) + + thinking_part = "" + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_index: + thinking_part = reasoning + THINKING_END_TOKEN + + if message.get("wo_eos", False): + prompt += thinking_part + content + tool_calls_content + else: + prompt += thinking_part + content + tool_calls_content + EOS_TOKEN + else: + raise NotImplementedError(f"Unsupported DeepSeek-V4 message role: {role}") + + next_role = messages[index + 1].get("role") if index + 1 < len(messages) else None + if next_role is not None and next_role not in ("assistant", "latest_reminder"): + return prompt + + task = message.get("task") + if task is not None: + if task not in VALID_TASKS: + raise ValueError(f"Invalid DeepSeek-V4 task: {task}") + if task == "action": + prompt += ASSISTANT_TOKEN + prompt += THINKING_START_TOKEN if thinking_mode == "thinking" else THINKING_END_TOKEN + prompt += VALID_TASKS[task] + elif role in ("user", "developer") and (next_role == "assistant" or add_generation_prompt): + prompt += ASSISTANT_TOKEN + if thinking_mode == "thinking" and (not drop_thinking or index >= last_user_index): + prompt += THINKING_START_TOKEN + else: + prompt += THINKING_END_TOKEN + + return prompt + + +def _encode_messages( + messages: list[dict[str, Any]], + thinking_mode: str, + drop_thinking: bool, + add_generation_prompt: bool, + reasoning_effort: str | None, +) -> str: + messages = _merge_tool_messages(messages) + messages = _sort_tool_results_by_call_order(messages) + + effective_drop_thinking = drop_thinking + if any(message.get("tools") for message in messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + messages = _drop_thinking_messages(messages) + + prompt = BOS_TOKEN + for index in range(len(messages)): + prompt += _render_message( + index, + messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + add_generation_prompt=add_generation_prompt, + reasoning_effort=reasoning_effort, + ) + return prompt + + class DeepseekV4Tokenizer(TransformersTokenizer): """DeepSeek-V4 tokenizer with the checkpoint reference chat format.""" @@ -65,29 +429,29 @@ def from_pretrained( return cls(tokenizer) def apply_chat_template(self, messages, tools=None, **kwargs): - if tools: - raise NotImplementedError("DeepSeek-V4 tool-call chat formatting is not supported yet.") - - add_generation_prompt = kwargs.get("add_generation_prompt", True) tokenize = kwargs.get("tokenize", False) + thinking = kwargs.get("thinking", False) or kwargs.get("enable_thinking", False) + thinking_mode = "thinking" if thinking else "chat" + reasoning_effort = kwargs.get("reasoning_effort") + if reasoning_effort not in ("max", "high"): + reasoning_effort = None - rendered = BOS_TOKEN - for idx, message in enumerate(messages): - role = message.get("role") - content = _message_content_to_text(message.get("content")) - next_role = messages[idx + 1].get("role") if idx + 1 < len(messages) else None - - if role == "system": - rendered += content - elif role in ("user", "developer"): - rendered += USER_TOKEN + content - if next_role == "assistant" or (next_role is None and add_generation_prompt): - rendered += ASSISTANT_TOKEN + THINKING_END_TOKEN - elif role == "assistant": - rendered += content + EOS_TOKEN - else: - raise NotImplementedError(f"Unsupported DeepSeek-V4 message role: {role}") + conversation = kwargs.get("conversation", messages) + messages = list(conversation) + if tools: + messages.insert(0, {"role": "system", "tools": tools}) + + rendered = _encode_messages( + messages=messages, + thinking_mode=thinking_mode, + drop_thinking=kwargs.get("drop_thinking", True), + add_generation_prompt=True, + reasoning_effort=reasoning_effort, + ) if tokenize: - return self.encode(rendered, add_special_tokens=False) + tokenizer_kwargs = { + key: kwargs[key] for key in ("truncation", "max_length") if key in kwargs + } + return self.encode(rendered, add_special_tokens=False, **tokenizer_kwargs) return rendered diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 59d0e8b2814a..3b7e19f3ce55 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -25,6 +25,7 @@ from tensorrt_llm.serve.tool_parser.core_types import (StreamingParseResult, StructureInfo) from tensorrt_llm.serve.tool_parser.deepseekv3_parser import DeepSeekV3Parser +from tensorrt_llm.serve.tool_parser.deepseekv4_parser import DeepSeekV4Parser from tensorrt_llm.serve.tool_parser.deepseekv31_parser import DeepSeekV31Parser from tensorrt_llm.serve.tool_parser.deepseekv32_parser import DeepSeekV32Parser from tensorrt_llm.serve.tool_parser.gemma4_parser import Gemma4ToolParser @@ -1504,6 +1505,60 @@ def test_encode_messages_multi_turn_with_tool_calls(self): assert ">ls<" in result +# ============================================================================ +# DeepSeekV4Parser Tests +# ============================================================================ + + +class TestDeepSeekV4Parser(BaseToolParserTestClass): + """Test suite for DeepSeekV4Parser class.""" + + def make_parser(self): + return DeepSeekV4Parser() + + def make_tool_parser_test_cases(self): + return ToolParserTestCases( + has_tool_call_true= + ('Some text <｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + " "), + detect_and_parse_single_tool=( + ('Normal text<｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + " "), + "Normal text", + "get_weather", + { + "location": "NYC" + }, + ), + detect_and_parse_multiple_tools=( + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + ' <｜DSML｜invoke name="search_web"> ' + '{ "query": "AI" } '), + ("get_weather", "search_web"), + ), + detect_and_parse_malformed_tool= + ('<|DSML|tool_calls> <|DSML|invoke name="get_weather"> ' + '<|DSML|parameter name="location" string="true">NYC ' + " "), + detect_and_parse_with_parameters_key=( + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="search_web"> ' + '{ "query": "test" } '), + "search_web", + { + "query": "test" + }, + ), + parse_streaming_increment_partial_bot_token="<｜DSML｜tool", + undefined_tool= + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="undefined_func"> ' + '<｜DSML｜parameter name="arg" string="true">value ' + " "), + ) + + # ============================================================================ # Glm4ToolParser Tests # ============================================================================ diff --git a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py index a95525fe64fd..dec508d49cbc 100644 --- a/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py +++ b/tests/unittest/llmapi/test_deepseek_v4_tokenizer.py @@ -69,6 +69,326 @@ def test_deepseek_v4_chat_template_tokenize_uses_rendered_prompt(): ) +def test_deepseek_v4_chat_template_supports_thinking_mode(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + enable_thinking=True, + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_supports_thinking_alias(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + thinking=True, + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_matches_vllm_add_generation_prompt_behavior(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + add_generation_prompt=False, + ) + + assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_accepts_openai_reasoning_effort_values(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + for reasoning_effort in ("none", "low", "medium", "high"): + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + enable_thinking=True, + reasoning_effort=reasoning_effort, + ) + + assert prompt.endswith("<｜Assistant｜>") + assert "Reasoning Effort: Absolute maximum" not in prompt + + +def test_deepseek_v4_chat_template_preserves_reference_max_reasoning_effort(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "hello", + } + ], + tokenize=False, + enable_thinking=True, + reasoning_effort="max", + ) + + assert prompt.startswith("<｜begin▁of▁sentence｜>Reasoning Effort: Absolute maximum") + assert prompt.endswith("<｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_drops_historical_thinking_without_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "first", + }, + { + "role": "assistant", + "reasoning": "hidden chain", + "content": "answer", + }, + { + "role": "user", + "content": "second", + }, + ], + tokenize=False, + enable_thinking=True, + ) + + assert "hidden chain" not in prompt + assert "answer<｜end▁of▁sentence｜>" in prompt + assert prompt.endswith("<｜User｜>second<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_keeps_historical_thinking_with_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ] + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "first", + }, + { + "role": "assistant", + "reasoning": "kept chain", + "content": "answer", + }, + { + "role": "user", + "content": "second", + }, + ], + tools=tools, + tokenize=False, + enable_thinking=True, + ) + + assert "kept chainanswer<｜end▁of▁sentence｜>" in prompt + assert prompt.endswith("<｜User｜>second<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_renders_developer_tools_and_latest_reminder(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ] + messages = [ + { + "role": "system", + "content": "sys", + }, + { + "role": "latest_reminder", + "content": "today", + }, + { + "role": "developer", + "content": "dev", + "tools": tools, + }, + { + "role": "assistant", + "reasoning": "need search", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "search", + "arguments": '{"query": "x"}', + }, + } + ], + }, + { + "role": "tool", + "content": "[0]", + }, + ] + + prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + enable_thinking=True, + ) + + assert prompt.startswith("<｜begin▁of▁sentence｜>sys<｜latest_reminder｜>today<｜User｜>dev") + assert "## Tools" in prompt + assert '<｜DSML｜invoke name="search">' in prompt + assert "need search" in prompt + assert "<｜User｜>[0]<｜Assistant｜>" in prompt + + +def test_deepseek_v4_chat_template_renders_action_task_token(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "system", + "content": "sys", + }, + { + "role": "latest_reminder", + "content": "today", + }, + { + "role": "user", + "content": "search this", + "task": "action", + }, + { + "role": "assistant", + "content": "Search", + }, + ], + tokenize=False, + ) + + assert prompt == ( + "<｜begin▁of▁sentence｜>sys<｜latest_reminder｜>today" + "<｜User｜>search this<｜Assistant｜><｜action｜>" + "Search<｜end▁of▁sentence｜>" + ) + + +def test_deepseek_v4_chat_template_uses_v4_tool_prompt_from_request_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + + prompt = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": "Weather?", + } + ], + tools=tools, + tokenize=False, + ) + + assert "## Tools" in prompt + assert "<｜DSML｜tool_calls>" in prompt + assert "" in prompt + assert "function_calls" not in prompt + assert '"name": "get_weather"' in prompt + assert prompt.endswith("<｜User｜>Weather?<｜Assistant｜>") + + +def test_deepseek_v4_chat_template_renders_tool_call_history(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + messages = [ + { + "role": "user", + "content": "List the repo", + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "str_replace_editor", + "arguments": '{"command": "view", "path": "/testbed"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "file list", + }, + ] + + prompt = tokenizer.apply_chat_template(messages, tokenize=False) + + assert '<｜DSML｜invoke name="str_replace_editor">' in prompt + assert '<｜DSML｜parameter name="command" string="true">view' in prompt + assert '<｜DSML｜parameter name="path" string="true">/testbed' in prompt + assert "<｜User｜>file list<｜Assistant｜>" in prompt + assert 'parameter name="arguments"' not in prompt + + def test_deepseek_v4_custom_tokenizer_reuses_loaded_wrapper(): tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) @@ -95,3 +415,35 @@ def test_deepseek_v4_server_chat_template_path_uses_custom_tokenizer(): ) assert prompt == ("<｜begin▁of▁sentence｜><｜User｜>hello<｜Assistant｜>") + + +def test_deepseek_v4_server_chat_template_path_forwards_tools(): + tokenizer = DeepseekV4Tokenizer(_DummyTokenizer()) + tools = [ + { + "type": "function", + "function": { + "name": "search", + "description": "Search", + "parameters": {"type": "object"}, + }, + } + ] + + prompt = apply_chat_template( + model_type="deepseek_v4", + tokenizer=tokenizer, + processor=None, + conversation=[ + { + "role": "user", + "content": "hello", + } + ], + add_generation_prompt=True, + mm_placeholder_counts=[{}], + tools=tools, + ) + + assert "<｜DSML｜tool_calls>" in prompt + assert '"name": "search"' in prompt diff --git a/tests/unittest/llmapi/test_reasoning_parser.py b/tests/unittest/llmapi/test_reasoning_parser.py index 6806428ec5ea..68543c247a5f 100644 --- a/tests/unittest/llmapi/test_reasoning_parser.py +++ b/tests/unittest/llmapi/test_reasoning_parser.py @@ -60,6 +60,34 @@ def test_deepseek_r1_reasoning_parser_stream(delta_texts: list, content: list, assert result.reasoning_content == reasoning_context[i] +@pytest.mark.parametrize("chat_template_kwargs", [{ + "thinking": True +}, { + "enable_thinking": True +}]) +def test_deepseek_v4_reasoning_parser_extracts_when_thinking( + chat_template_kwargs: dict): + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + "deepseek_v4", chat_template_kwargs) + + result = reasoning_parser.parse(f"hidden{R1_END}visible") + + assert result.content == "visible" + assert result.reasoning_content == "hidden" + + +def test_deepseek_v4_reasoning_parser_streams_when_thinking(): + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + "deepseek_v4", {"enable_thinking": True}) + + deltas = ["hid", f"den{R1_END}visible", " tail"] + results = [reasoning_parser.parse_delta(delta) for delta in deltas] + + assert [result.content for result in results] == ["", "visible", " tail"] + assert [result.reasoning_content + for result in results] == ["hid", "den", ""] + + TOOL_START = "<|tool_calls_section_begin|>" From 7e763a61593d28cfdee0a197d8bad9b96e133d75 Mon Sep 17 00:00:00 2001 From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Date: Thu, 21 May 2026 06:52:01 +0000 Subject: [PATCH 08/58] [TRTLLM-12111][feat] Add V2 KV cache event support (#13589) Filtered out disaggregation serving/router test changes for user/fanrongl/dsv4_model. Signed-off-by: Fanrong Li --- .../_torch/pyexecutor/model_engine.py | 11 +- tensorrt_llm/_utils.py | 6 + tensorrt_llm/llmapi/llm_args.py | 10 + tensorrt_llm/runtime/kv_cache_hash.py | 32 + .../runtime/kv_cache_manager_v2/__init__.py | 20 + .../runtime/kv_cache_manager_v2/__init__.pyi | 100 +- .../kv_cache_manager_v2/_block_radix_tree.py | 35 +- .../kv_cache_manager_v2/_core/_kv_cache.py | 6 + .../_core/_kv_cache_manager.py | 27 +- .../kv_cache_manager_v2/_event_manager.py | 536 +++++++++++ .../kv_cache_manager_v2/_storage_manager.py | 54 +- .../kv_cache_manager_v2/setup_mypyc.py | 1 + .../_torch/executor/test_resource_manager.py | 35 +- .../test_kv_cache_event_manager.py | 896 ++++++++++++++++++ tests/unittest/llmapi/test_llm_args.py | 8 + .../llmapi/test_llm_kv_cache_events.py | 48 +- 16 files changed, 1808 insertions(+), 17 deletions(-) create mode 100644 tensorrt_llm/runtime/kv_cache_hash.py create mode 100644 tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 0e3c03c7b15e..ea32c50f3a3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1637,6 +1637,12 @@ def _create_cuda_graph_warmup_request( if requests is None: return None + def free_warmup_requests() -> None: + for r in requests: + kv_cache_manager.free_resources(r) + if draft_kv_cache_manager is not None: + draft_kv_cache_manager.free_resources(r) + # Add one dummy request with the maximum possible sequence length. max_seq_len = min( self.max_seq_len if max_seq_len is None else max_seq_len, @@ -1688,10 +1694,7 @@ def _create_cuda_graph_warmup_request( draft_kv_cache_manager=draft_kv_cache_manager) if max_seq_len_request is None: - for r in requests: - kv_cache_manager.free_resources(r) - if draft_kv_cache_manager is not None: - draft_kv_cache_manager.free_resources(r) + free_warmup_requests() return None else: max_seq_len_request = max_seq_len_request[0] diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 520579efd2e9..9abe26037331 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1178,6 +1178,12 @@ def to_json_str(cls, event): "data": event_serialize_func(event.data), "window_size": event.window_size, } + hash_algo = getattr(event, "hash_algo", None) + if hash_algo is not None: + json_str["hash_algo"] = hash_algo + layer_group_id = getattr(event, "layer_group_id", None) + if layer_group_id is not None: + json_str["layer_group_id"] = layer_group_id if event.attention_dp_rank is not None: json_str["attention_dp_rank"] = event.attention_dp_rank diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9721c8a17079..c70b7ce909c6 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2916,6 +2916,16 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): status="prototype", description="Whether to use the KV cache manager v2 (experimental).") + kv_cache_event_hash_algo: Literal[ + "auto", "v1_block_key", "v2_sha256", "v2_sha256_64"] = Field( + default="auto", + status="prototype", + description= + "The block hash algorithm used by KV cache manager events. " + "'auto' uses the native hash for each KV cache manager. " + "Explicit V2 hash choices are ignored with a warning by the V1 " + "KV cache manager.") + max_util_for_resume: float = Field( default=0.95, ge=0, diff --git a/tensorrt_llm/runtime/kv_cache_hash.py b/tensorrt_llm/runtime/kv_cache_hash.py new file mode 100644 index 000000000000..ae3de666922a --- /dev/null +++ b/tensorrt_llm/runtime/kv_cache_hash.py @@ -0,0 +1,32 @@ +# 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. + +KV_CACHE_HASH_ALGO_AUTO = "auto" +KV_CACHE_HASH_ALGO_V1 = "v1_block_key" +KV_CACHE_HASH_ALGO_V2 = "v2_sha256" +KV_CACHE_HASH_ALGO_V2_SHA256_64 = "v2_sha256_64" +KV_CACHE_HASH_ALGO_DEFAULT = KV_CACHE_HASH_ALGO_V1 + + +def get_effective_kv_cache_event_hash_algo(hash_algo: str, use_kv_cache_manager_v2: bool) -> str: + if hash_algo != KV_CACHE_HASH_ALGO_AUTO: + return hash_algo + if use_kv_cache_manager_v2: + return KV_CACHE_HASH_ALGO_V2 + return KV_CACHE_HASH_ALGO_V1 + + +def truncate_sha256_hash_to_int64(block_hash: bytes) -> int: + return int.from_bytes(block_hash[:8], "big", signed=False) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 4d3607a9cc7c..3187a3407dc2 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -50,6 +50,17 @@ ScratchDesc, _KVCache, ) +from ._event_manager import ( + KVCacheCreatedData, + KVCacheEvent, + KVCacheEventDiff, + KVCacheEventManager, + KVCacheRemovedData, + KVCacheStoredBlockData, + KVCacheStoredData, + KVCacheUpdatedData, + UniqueToken, +) from ._life_cycle_registry import LayerGroupId, LifeCycleId from ._storage import BufferId @@ -60,6 +71,15 @@ "TokenIdExt", "KVCacheManager", "_KVCache", + "KVCacheCreatedData", + "KVCacheEvent", + "KVCacheEventDiff", + "KVCacheEventManager", + "KVCacheRemovedData", + "KVCacheStoredBlockData", + "KVCacheStoredData", + "KVCacheUpdatedData", + "UniqueToken", "BeamIndex", "DEFAULT_BEAM_INDEX", "LayerId", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 451778f7e754..4029775c9e79 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -153,6 +153,98 @@ class KVCacheManagerConfig: @property def enable_swa_scratch_reuse(self) -> bool: ... +# From _event_manager.py +EventBlockHash: TypeAlias = int | str +BlockHashLike: TypeAlias = bytes | EventBlockHash +BlockHashesLike: TypeAlias = BlockHashLike | Iterable[BlockHashLike] +EventTokenId: TypeAlias = int | str +MmKey: TypeAlias = tuple[bytes, int] | tuple[bytes, int, str | None] +AttentionDpGatherFn: TypeAlias = Callable[[list["KVCacheEvent"]], list[list["KVCacheEvent"]]] + +@dataclass(slots=True, frozen=True) +class UniqueToken: + token_id: EventTokenId + token_extra_id: int = ... + +@dataclass(slots=True, frozen=True) +class KVCacheCreatedData: + num_blocks_per_cache_level: list[int] + +@dataclass(slots=True, frozen=True) +class KVCacheStoredBlockData: + block_hash: EventBlockHash + tokens: list[UniqueToken] + cache_level: int + priority: int + mm_keys: list[MmKey] = ... + +@dataclass(slots=True, frozen=True) +class KVCacheStoredData: + parent_hash: EventBlockHash | None + blocks: list[KVCacheStoredBlockData] + +@dataclass(slots=True, frozen=True) +class KVCacheRemovedData: + block_hashes: list[EventBlockHash] + +@dataclass(slots=True, frozen=True) +class KVCacheEventDiff: + old_value: int + new_value: int + +@dataclass(slots=True, frozen=True) +class KVCacheUpdatedData: + block_hash: EventBlockHash + cache_level: KVCacheEventDiff | None + priority: KVCacheEventDiff | None + +@dataclass(slots=True, frozen=True) +class KVCacheEvent: + event_id: int + data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData + window_size: int + hash_algo: str | None = None + attention_dp_rank: int | None = None + layer_group_id: int | None = None + +class KVCacheEventManager: + def __init__( + self, + max_kv_event_entries: int, + *, + window_size: int = ..., + attention_dp_rank: int | None = None, + attention_dp_gather: AttentionDpGatherFn | None = None, + hash_algo: str = ..., + window_size_by_layer_group: dict[int, int] | None = None, + ) -> None: ... + def add_created_event( + self, + num_blocks_per_cache_level: Sequence[int], + layer_group_ids: Sequence[int] | None = None, + ) -> None: ... + def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: ... + def add_stored_event( + self, + parent_hash: EventBlockHash | None, + blocks: Sequence[KVCacheStoredBlockData], + layer_group_id: int | None = None, + ) -> None: ... + def add_stored_block_event_from_block(self, block: Any) -> None: ... + def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: ... + def add_removed_event(self, block_hashes: BlockHashesLike) -> None: ... + def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: ... + def add_updated_event( + self, + block_hash: BlockHashLike, + *, + cache_level: KVCacheEventDiff | None = None, + priority: KVCacheEventDiff | None = None, + layer_group_id: int | None = None, + ) -> None: ... + def flush_iteration_events(self) -> None: ... + def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: ... + # From _block_radix_tree.py def gen_multimodal_cache_key_tokens( id_offset: int, @@ -295,7 +387,11 @@ class PageIndexConverter: ) -> list[int]: ... class KVCacheManager: - def __init__(self, config: KVCacheManagerConfig) -> None: ... + def __init__( + self, + config: KVCacheManagerConfig, + event_manager: KVCacheEventManager | None = None, + ) -> None: ... def __del__(self) -> None: ... def shutdown(self) -> None: ... def clear_reusable_blocks(self) -> None: ... @@ -327,6 +423,8 @@ class KVCacheManager: @property def tokens_per_block(self) -> int: ... @property + def event_manager(self) -> Any | None: ... + @property def allow_seq_rebasing(self) -> bool: ... @property def enable_partial_match(self) -> bool: ... diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index 1024eca91575..e62cb2a30ee3 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -23,6 +23,7 @@ from ._utils import TypedIndexList, chunked, div_up, filled_list, find_index, unwrap_rawref if TYPE_CHECKING: + from ._event_manager import KVCacheEventManager from ._page import CommittedPage BlockKey = bytes @@ -141,12 +142,15 @@ def remove_subtree(root: "RootBlock | Block") -> list[rawref.ref["CommittedPage" # taking O(1) space # remove leaf blocks one by one, in post-order ret: list[rawref.ref["CommittedPage"]] = [] + removed_block_hashes: list[BlockKey] = [] + event_manager = get_tree(root).event_manager block: "RootBlock | Block" = root while True: if block.next: block = next(iter(block.next.values())) else: if isinstance(block, Block): + removed_block_hashes.append(block.key) ret.extend(p for p in block.storage if p is not None) block.storage = filled_list(None, block.num_life_cycles) assert isinstance(block, RootBlock) or all(page is None for page in block.storage), ( @@ -164,6 +168,8 @@ def remove_subtree(root: "RootBlock | Block") -> list[rawref.ref["CommittedPage" break assert not isinstance(prev_block, BlockRadixTree) block = prev_block + if event_manager is not None: + event_manager.add_removed_event(removed_block_hashes) return ret @@ -324,8 +330,11 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N if len(b.tokens) < len(tokens) and tokens[: len(b.tokens)] == b.tokens: assert NDEBUG or (not b.is_full and b is not self and b.key == k and not b.next) to_remove.append(k) + event_manager = get_tree(prev).event_manager if to_remove else None for k in to_remove: b = prev.next.pop(k) + if event_manager is not None: + event_manager.add_removed_event(b.key) assert b.is_orphan # _KVCache may still hold it. # prev.next keeps a strong ref to this _Block, so no need to remove self from prev.next in __del__(). prev.next[self.key] = self @@ -361,6 +370,7 @@ def prev(self) -> "Block | RootBlock": def unset_page(self, lc_idx: LifeCycleId, lc: LifeCycle) -> None: if self.storage[lc_idx] is None: return + event_manager = get_tree(self).event_manager ordinal = self.ordinal self.storage[lc_idx] = None if type(lc) is AttnLifeCycle and (lc.window_size is None or ordinal < lc.num_sink_blocks): @@ -371,6 +381,8 @@ def unset_page(self, lc_idx: LifeCycleId, lc: LifeCycle) -> None: assert page.status == PageStatus.DROPPABLE if page.scheduled_for_eviction: page.manager.exclude_from_eviction(page) + elif event_manager is not None: + event_manager.add_removed_life_cycle_event(self.key, int(lc_idx)) # It's possible to implement more sophisticated logic to remove useless blocks for SWA, e.g. # check if consecutive available blocks is sufficient for window_size. (TRTLLM-8802) # But for simplicity, we leave it for now. @@ -382,6 +394,8 @@ def unset_page(self, lc_idx: LifeCycleId, lc: LifeCycle) -> None: ): if curr.key in curr.prev.next: curr.prev.next.pop(curr.key) + if event_manager is not None: + event_manager.add_removed_event(curr.key) curr = curr.prev @property @@ -400,15 +414,28 @@ def is_orphan(self) -> bool: class BlockRadixTree: - __slots__ = ("_life_cycles", "_tokens_per_block", "next", "__rawref__") + __slots__ = ( + "_life_cycles", + "_tokens_per_block", + "_event_manager", + "next", + "__rawref__", + ) _life_cycles: LifeCycleRegistry _tokens_per_block: int + _event_manager: "KVCacheEventManager | None" next: Children[RootBlock] __rawref__: rawref.ref["BlockRadixTree"] - def __init__(self, life_cycles: LifeCycleRegistry, tokens_per_block: int) -> None: + def __init__( + self, + life_cycles: LifeCycleRegistry, + tokens_per_block: int, + event_manager: "KVCacheEventManager | None" = None, + ) -> None: self._life_cycles = life_cycles self._tokens_per_block = tokens_per_block + self._event_manager = event_manager self.next = {} self.__rawref__ = rawref.NULL @@ -429,6 +456,10 @@ def tokens_per_block(self) -> int: def life_cycles(self) -> LifeCycleRegistry: return self._life_cycles + @property + def event_manager(self) -> "KVCacheEventManager | None": + return self._event_manager + @property def num_life_cycles(self) -> LifeCycleId: return self.life_cycles.size 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 a7fec6eea3e6..057f6a0ce32b 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 @@ -1152,6 +1152,9 @@ def _commit_block(self, ordinal: BlockOrdinal, is_last: bool) -> None: seq_block.tree_block = tree_block assert self._get_tree_block(ordinal) is tree_block self._num_committed_blocks = BlockOrdinal(ordinal + 1) + event_manager = self.manager.event_manager + if event_manager is not None: + event_manager.add_stored_block_event_from_block(tree_block) elif tree_block.is_full and self.manager.allow_seq_rebasing and is_full: # Happens when a concurrent request committed the same tokens before us. # Try to replace our pages with pages from the existing block to save memory. @@ -1168,6 +1171,9 @@ def _commit_block(self, ordinal: BlockOrdinal, is_last: bool) -> None: page = cast(UncommittedPage, cast(_SharedPageLock, beam_block[lc]).page) beam_block[lc] = None p = page.convert_to_committed(tree_block, self.finish_event) + event_manager = self.manager.event_manager + if event_manager is not None: + event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc)) # The page comes from uncommitted page of self, so safe to skip wait. beam_block[lc] = ( p.lock(self, beam_idx, ordinal, lc, skip_wait=True) if locked else p.hold() 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 da065a366158..dd9f8a874a9c 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 @@ -18,7 +18,7 @@ from collections.abc import Callable, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import Iterable, Iterator, cast +from typing import TYPE_CHECKING, Iterable, Iterator, cast from .. import rawref from .._block_radix_tree import BlockRadixTree, ReuseMatch, ReuseScope @@ -58,6 +58,9 @@ from ._kv_cache import _KVCache from ._moving_average import MovingAverage +if TYPE_CHECKING: + from .._event_manager import KVCacheEventManager + @dataclass(slots=True, frozen=True) class MemoryPoolDesc: @@ -194,6 +197,7 @@ class KVCacheManager: "_num_sampled_kv_caches", "_last_adjustment_time", "_last_update_num_sampled_kv_caches", + "_event_manager", ) _init_config: KVCacheManagerConfig _life_cycles: LifeCycleRegistry @@ -216,13 +220,18 @@ class KVCacheManager: _num_sampled_kv_caches: int _last_adjustment_time: float _last_update_num_sampled_kv_caches: int + _event_manager: "KVCacheEventManager | None" - def __init__(self, config: KVCacheManagerConfig) -> None: + def __init__( + self, + config: KVCacheManagerConfig, + event_manager: "KVCacheEventManager | None" = None, + ) -> None: init_cuda_once() config = deepcopy(config) self._init_config = config self._life_cycles = LifeCycleRegistry(config) - self._radix_tree = BlockRadixTree(self._life_cycles, config.tokens_per_block) + self._radix_tree = BlockRadixTree(self._life_cycles, config.tokens_per_block, event_manager) storage_config = create_storage_config(config) self._storage = StorageManager( self._life_cycles, @@ -231,6 +240,7 @@ def __init__(self, config: KVCacheManagerConfig) -> None: config.swa_scratch_reuse, typical_batch=config.typical_step, constraints=config.constraints, + event_manager=event_manager, ) self._living_kv_caches = set[rawref.ref[_KVCache]]() decay = 0.9999 @@ -243,6 +253,7 @@ def __init__(self, config: KVCacheManagerConfig) -> None: self._num_sampled_kv_caches = 0 self._last_adjustment_time = time.monotonic() self._last_update_num_sampled_kv_caches = 0 + self._event_manager = event_manager def __del__(self) -> None: self.shutdown() @@ -425,6 +436,10 @@ def cache_tier_list(self) -> HomoTuple[CacheTier]: def tokens_per_block(self) -> int: return self._radix_tree.tokens_per_block + @property + def event_manager(self) -> "KVCacheEventManager | None": + return self._event_manager + @property def allow_seq_rebasing(self) -> bool: """ @@ -687,7 +702,8 @@ def get_num_slots(seq_len: int) -> TypedIndexList[PoolGroupIndex, int]: for pg in typed_range(num_pool_groups): remaining_slots[pg] -= get_num_slots(1)[pg] * (batch_size - 1) - assert remaining_slots[pg] >= 0 + if remaining_slots[pg] < 0: + return 0 def is_enough(num_blocks: int) -> bool: return all( @@ -695,7 +711,8 @@ def is_enough(num_blocks: int) -> bool: for cnt, rem in zip(get_num_slots(num_blocks * tokens_per_block), remaining_slots) ) - assert is_enough(1) + if not is_enough(1): + return 0 lb = 1 ub = div_up(token_num_upper_bound, tokens_per_block) if is_enough(ub): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py new file mode 100644 index 000000000000..38488505a8b1 --- /dev/null +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -0,0 +1,536 @@ +# 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. + +from __future__ import annotations + +import time +from collections import deque +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field, replace +from threading import Condition +from typing import Any, Callable + +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_hash import ( + KV_CACHE_HASH_ALGO_AUTO, + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + truncate_sha256_hash_to_int64, +) + +from ._common import GPU_LEVEL, PRIORITY_DEFAULT, CacheLevel, Priority, TokenIdExt + +EventBlockHash = int | str +BlockHashLike = bytes | EventBlockHash +BlockHashesLike = BlockHashLike | Iterable[BlockHashLike] +LayerGroupId = int | None +EventTokenId = int | str +MmKey = tuple[bytes, int] | tuple[bytes, int, str | None] +AttentionDpGatherFn = Callable[[list["KVCacheEvent"]], list[list["KVCacheEvent"]]] + + +@dataclass(slots=True, frozen=True) +class UniqueToken: + token_id: EventTokenId + token_extra_id: int = 0 + + +@dataclass(slots=True, frozen=True) +class KVCacheCreatedData: + num_blocks_per_cache_level: list[int] + + +@dataclass(slots=True, frozen=True) +class KVCacheStoredBlockData: + block_hash: EventBlockHash + tokens: list[UniqueToken] + cache_level: int + priority: int + mm_keys: list[MmKey] = field(default_factory=list) + + +@dataclass(slots=True, frozen=True) +class KVCacheStoredData: + parent_hash: EventBlockHash | None + blocks: list[KVCacheStoredBlockData] + + +@dataclass(slots=True, frozen=True) +class KVCacheRemovedData: + block_hashes: list[EventBlockHash] + + +@dataclass(slots=True, frozen=True) +class KVCacheEventDiff: + old_value: int + new_value: int + + +@dataclass(slots=True, frozen=True) +class KVCacheUpdatedData: + block_hash: EventBlockHash + cache_level: KVCacheEventDiff | None + priority: KVCacheEventDiff | None + + +@dataclass(slots=True, frozen=True) +class KVCacheEvent: + event_id: int + data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData + window_size: int + hash_algo: str | None = None + attention_dp_rank: int | None = None + layer_group_id: int | None = None + + +@dataclass(slots=True) +class _StoredBlockState: + block_hash: EventBlockHash + life_cycle_ids: set[int] + + +class KVCacheEventManager: + """Python event queue matching the C++ KV cache event serializer contract.""" + + def __init__( + self, + max_kv_event_entries: int, + *, + window_size: int = 0, + attention_dp_rank: int | None = None, + attention_dp_gather: AttentionDpGatherFn | None = None, + hash_algo: str = KV_CACHE_HASH_ALGO_V2, + window_size_by_layer_group: dict[int, int] | None = None, + ) -> None: + if hash_algo == KV_CACHE_HASH_ALGO_AUTO: + hash_algo = KV_CACHE_HASH_ALGO_V2 + elif hash_algo == KV_CACHE_HASH_ALGO_V1: + logger.warning( + "V1 KV cache event hash algorithm is not supported by V2 " + "KV cache manager events: " + f"{hash_algo}. Falling back to {KV_CACHE_HASH_ALGO_V2_SHA256_64}." + ) + hash_algo = KV_CACHE_HASH_ALGO_V2_SHA256_64 + elif hash_algo not in (KV_CACHE_HASH_ALGO_V2, KV_CACHE_HASH_ALGO_V2_SHA256_64): + raise ValueError(f"Unsupported V2 KV cache event hash algorithm: {hash_algo}") + self._max_kv_event_entries = max_kv_event_entries + self._window_size = window_size + self._window_size_by_layer_group = dict(window_size_by_layer_group or {}) + self._attention_dp_rank = attention_dp_rank + self._attention_dp_gather = attention_dp_gather + self._hash_algo = hash_algo + self._next_event_id = 0 + self._stored_blocks: dict[bytes, _StoredBlockState] = {} + self._latest_stored_events: dict[LayerGroupId, KVCacheEvent] = {} + self._latest_removed_block_hashes: dict[LayerGroupId, list[EventBlockHash]] = {} + self._pending_events: list[KVCacheEvent] = [] + self._events: deque[KVCacheEvent] = deque() + self._condition = Condition() + + def add_created_event( + self, + num_blocks_per_cache_level: Sequence[int], + layer_group_ids: Sequence[int] | None = None, + ) -> None: + data = KVCacheCreatedData(list(num_blocks_per_cache_level)) + if layer_group_ids is None: + self._add_event(data) + return + for layer_group_id in layer_group_ids: + self._add_event(data, layer_group_id=int(layer_group_id)) + + def set_layer_group_window_sizes(self, window_sizes: dict[int, int]) -> None: + with self._condition: + self._window_size_by_layer_group = dict(window_sizes) + + def add_stored_event( + self, + parent_hash: EventBlockHash | None, + blocks: Sequence[KVCacheStoredBlockData], + layer_group_id: int | None = None, + ) -> None: + if not blocks: + return + self._flush_removed_events(layer_group_id) + self._add_stored_event( + KVCacheStoredData(parent_hash, list(blocks)), + layer_group_id=layer_group_id, + ) + + def add_stored_block_event_from_block(self, block: Any) -> None: + life_cycle_ids = self._life_cycle_ids_from_radix_block(block) + if not life_cycle_ids: + return + parent_hash = self._parent_hash_from_radix_block(block) + self._stored_blocks[block.key] = _StoredBlockState( + block_hash=self._normalize_block_hash(block.key), + life_cycle_ids=set(life_cycle_ids), + ) + for life_cycle_id in sorted(life_cycle_ids): + block_data = self._stored_block_from_radix_block(block, life_cycle_ids={life_cycle_id}) + if block_data is not None: + self.add_stored_event(parent_hash, [block_data], life_cycle_id) + + def add_stored_life_cycle_event_from_block(self, block: Any, life_cycle_id: int) -> None: + state = self._stored_blocks.get(block.key) + life_cycle_id = int(life_cycle_id) + if state is not None: + if life_cycle_id in state.life_cycle_ids: + return + block_data = self._stored_block_from_radix_block(block, life_cycle_ids={life_cycle_id}) + if block_data is None: + return + state.life_cycle_ids.add(life_cycle_id) + self.add_stored_event( + self._parent_hash_from_radix_block(block), + [block_data], + layer_group_id=life_cycle_id, + ) + return + self.add_stored_block_event_from_block(block) + + def add_removed_event(self, block_hashes: BlockHashesLike) -> None: + removed_block_hashes_by_layer_group: dict[int, list[EventBlockHash]] = {} + removed_block_hashes_without_layer_group: list[EventBlockHash] = [] + for block_hash in self._iter_block_hashes(block_hashes): + removed_state = self._pop_stored_block_state(block_hash) + if removed_state is None: + continue + normalized_hash, life_cycle_ids = removed_state + if life_cycle_ids: + for life_cycle_id in sorted(life_cycle_ids): + removed_block_hashes_by_layer_group.setdefault(life_cycle_id, []).append( + normalized_hash + ) + else: + removed_block_hashes_without_layer_group.append(normalized_hash) + + if removed_block_hashes_without_layer_group: + self._enqueue_removed_event(removed_block_hashes_without_layer_group) + for layer_group_id, removed_block_hashes in sorted( + removed_block_hashes_by_layer_group.items() + ): + self._enqueue_removed_event(removed_block_hashes, layer_group_id=layer_group_id) + + def add_removed_life_cycle_event(self, block_hash: bytes, life_cycle_id: int) -> None: + removed_state = self._pop_stored_life_cycle_block_state(block_hash, life_cycle_id) + if removed_state is None: + return + normalized_hash, removed_life_cycle_id, _ = removed_state + self._enqueue_removed_event( + [normalized_hash], + layer_group_id=removed_life_cycle_id, + ) + + def add_updated_event( + self, + block_hash: BlockHashLike, + *, + cache_level: KVCacheEventDiff | None = None, + priority: KVCacheEventDiff | None = None, + layer_group_id: int | None = None, + ) -> None: + if cache_level is None and priority is None: + return + normalized_block_hash = self._get_stored_block_hash(block_hash) + if normalized_block_hash is None: + return + self._add_event( + KVCacheUpdatedData( + block_hash=normalized_block_hash, + cache_level=cache_level, + priority=priority, + ), + layer_group_id=layer_group_id, + ) + + def flush_iteration_events(self) -> None: + if self._attention_dp_gather is not None: + with self._condition: + local_events = self._drain_pending_events_unlocked() + local_events = self._trim_events(local_events, self._max_kv_event_entries) + gathered_events = self._attention_dp_gather(local_events) + if self._attention_dp_rank != 0: + return + events = [ + event + for rank_events in gathered_events + for event in self._trim_events(rank_events, self._max_kv_event_entries) + ] + with self._condition: + self._publish_events_unlocked( + events, + max_kv_event_entries=( + self._max_kv_event_entries * max(1, len(gathered_events)) + ), + ) + self._condition.notify_all() + return + + with self._condition: + self._publish_events_unlocked(self._drain_pending_events_unlocked()) + self._condition.notify_all() + + def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEvent]: + with self._condition: + if not self._events and timeout_ms is None: + while not self._events: + self._condition.wait() + elif not self._events and timeout_ms > 0: + deadline = time.monotonic() + timeout_ms / 1000 + while not self._events: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self._condition.wait(timeout=remaining) + + events = list(self._events) + self._events.clear() + return events + + def _add_event( + self, + data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData, + layer_group_id: LayerGroupId = None, + ) -> None: + if self._max_kv_event_entries <= 0: + return + with self._condition: + self._add_event_unlocked(data, layer_group_id) + + def _add_stored_event( + self, + data: KVCacheStoredData, + layer_group_id: LayerGroupId = None, + ) -> None: + if self._max_kv_event_entries <= 0: + return + with self._condition: + has_pending_removed_events = bool(self._latest_removed_block_hashes) + latest_event = self._latest_stored_events.get(layer_group_id) + if ( + not has_pending_removed_events + and latest_event is not None + and isinstance(latest_event.data, KVCacheStoredData) + ): + latest_blocks = latest_event.data.blocks + if latest_blocks and latest_blocks[-1].block_hash == data.parent_hash: + merged_data = replace( + latest_event.data, + blocks=[*latest_blocks, *data.blocks], + ) + merged_event = replace(latest_event, data=merged_data) + self._replace_pending_event_unlocked(latest_event, merged_event) + self._latest_stored_events[layer_group_id] = merged_event + return + + event = self._add_event_unlocked(data, layer_group_id) + self._latest_stored_events[layer_group_id] = event + + def _replace_pending_event_unlocked( + self, + old_event: KVCacheEvent, + new_event: KVCacheEvent, + ) -> None: + for event_idx in range(len(self._pending_events) - 1, -1, -1): + if self._pending_events[event_idx] is old_event: + self._pending_events[event_idx] = new_event + return + raise RuntimeError("Stored event coalescing lost the pending event") + + def _enqueue_removed_event( + self, + block_hashes: Sequence[EventBlockHash], + layer_group_id: LayerGroupId = None, + ) -> None: + if not block_hashes or self._max_kv_event_entries <= 0: + return + with self._condition: + self._latest_removed_block_hashes.setdefault(layer_group_id, []).extend(block_hashes) + self._latest_stored_events.pop(layer_group_id, None) + + def _flush_removed_events(self, layer_group_id: LayerGroupId) -> None: + if self._max_kv_event_entries <= 0: + return + with self._condition: + self._flush_removed_events_unlocked(layer_group_id) + + def _flush_removed_events_unlocked(self, layer_group_id: LayerGroupId) -> None: + block_hashes = self._latest_removed_block_hashes.pop(layer_group_id, None) + if not block_hashes: + return + self._add_event_unlocked( + KVCacheRemovedData(block_hashes), + layer_group_id=layer_group_id, + ) + + def _flush_all_removed_events_unlocked(self) -> None: + layer_group_ids = list(self._latest_removed_block_hashes) + for layer_group_id in layer_group_ids: + self._flush_removed_events_unlocked(layer_group_id) + + def _add_event_unlocked( + self, + data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData, + layer_group_id: LayerGroupId = None, + ) -> KVCacheEvent: + if not isinstance(data, KVCacheRemovedData): + self._flush_all_removed_events_unlocked() + event = KVCacheEvent( + event_id=self._next_event_id, + data=data, + window_size=self._get_window_size(layer_group_id), + hash_algo=self._hash_algo, + attention_dp_rank=self._attention_dp_rank, + layer_group_id=layer_group_id, + ) + self._next_event_id += 1 + self._pending_events.append(event) + if not isinstance(data, KVCacheStoredData): + self._latest_stored_events.pop(layer_group_id, None) + return event + + def _drain_pending_events_unlocked(self) -> list[KVCacheEvent]: + self._flush_all_removed_events_unlocked() + events = self._pending_events + self._pending_events = [] + self._latest_stored_events.clear() + return events + + def _publish_events_unlocked( + self, + events: Sequence[KVCacheEvent], + *, + max_kv_event_entries: int | None = None, + ) -> None: + if not events: + return + if max_kv_event_entries is None: + max_kv_event_entries = self._max_kv_event_entries + self._events.extend(events) + while len(self._events) > max_kv_event_entries: + self._events.popleft() + + @staticmethod + def _trim_events( + events: Sequence[KVCacheEvent], max_kv_event_entries: int + ) -> list[KVCacheEvent]: + if max_kv_event_entries <= 0: + return [] + if len(events) <= max_kv_event_entries: + return list(events) + return list(events[-max_kv_event_entries:]) + + def _get_window_size(self, layer_group_id: LayerGroupId) -> int: + if layer_group_id is None: + return self._window_size + return self._window_size_by_layer_group.get(int(layer_group_id), self._window_size) + + @staticmethod + def _iter_block_hashes(block_hashes: BlockHashesLike) -> Iterable[BlockHashLike]: + if isinstance(block_hashes, (bytes, str, int)): + return (block_hashes,) + return block_hashes + + def _normalize_block_hash(self, block_hash: BlockHashLike) -> EventBlockHash: + if isinstance(block_hash, bytes): + if self._hash_algo == KV_CACHE_HASH_ALGO_V2_SHA256_64: + return truncate_sha256_hash_to_int64(block_hash) + return block_hash.hex() + return block_hash + + def _get_stored_block_hash(self, block_hash: BlockHashLike) -> EventBlockHash | None: + if isinstance(block_hash, bytes): + state = self._stored_blocks.get(block_hash) + return None if state is None else state.block_hash + return block_hash + + def _pop_stored_block_state( + self, block_hash: BlockHashLike + ) -> tuple[EventBlockHash, set[int]] | None: + if isinstance(block_hash, bytes): + state = self._stored_blocks.pop(block_hash, None) + if state is None: + return None + return state.block_hash, set(state.life_cycle_ids) + return block_hash, set() + + def _pop_stored_life_cycle_block_state( + self, block_hash: bytes, life_cycle_id: int + ) -> tuple[EventBlockHash, int, bool] | None: + state = self._stored_blocks.get(block_hash) + if state is None or not state.life_cycle_ids: + return None + + life_cycle_id = int(life_cycle_id) + if life_cycle_id not in state.life_cycle_ids: + return None + + state.life_cycle_ids.remove(life_cycle_id) + is_last_life_cycle = not state.life_cycle_ids + if is_last_life_cycle: + self._stored_blocks.pop(block_hash, None) + return state.block_hash, life_cycle_id, is_last_life_cycle + + @staticmethod + def _normalize_token(token: TokenIdExt) -> UniqueToken: + if isinstance(token, bytes): + return UniqueToken(token.hex()) + return UniqueToken(int(token)) + + def _stored_block_from_radix_block( + self, block: Any, life_cycle_ids: set[int] | None = None + ) -> KVCacheStoredBlockData | None: + cache_level: CacheLevel = GPU_LEVEL + priority: Priority = PRIORITY_DEFAULT + found_page = False + for life_cycle_id, page_ref in enumerate(block.storage): + if life_cycle_ids is not None and life_cycle_id not in life_cycle_ids: + continue + if page_ref is None: + continue + page = page_ref() + if page is None: + continue + cache_level = page.cache_level + priority = page.priority + found_page = True + break + + if life_cycle_ids is not None and not found_page: + return None + + return KVCacheStoredBlockData( + block_hash=self._normalize_block_hash(block.key), + tokens=[self._normalize_token(token) for token in block.tokens], + cache_level=int(cache_level), + priority=int(priority), + mm_keys=[], + ) + + @staticmethod + def _life_cycle_ids_from_radix_block(block: Any) -> set[int]: + return { + life_cycle_id + for life_cycle_id, page_ref in enumerate(block.storage) + if page_ref is not None and page_ref() is not None + } + + def _parent_hash_from_radix_block(self, block: Any) -> EventBlockHash | None: + parent = block.prev + if getattr(parent, "ordinal", -1) == -1: + return None + return self._normalize_block_hash(parent.key) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 9ade62d5dc99..b45d03362374 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -19,7 +19,7 @@ from collections import deque from dataclasses import dataclass from fractions import Fraction -from typing import Iterator, Sequence, cast +from typing import TYPE_CHECKING, Iterator, Sequence, cast from . import rawref from ._common import ( @@ -42,10 +42,11 @@ SwaScratchReuseConfig, ) from ._copy_engine import CopyTask, batched_copy +from ._event_manager import KVCacheEventDiff from ._eviction_controller import EvictablePage, PerLevelEvictionController from ._exceptions import OutOfPagesError from ._life_cycle_registry import LifeCycleId, LifeCycleRegistry, compute_scratch_range -from ._page import Page +from ._page import CommittedPage, Page from ._storage import CacheLevelStorage from ._storage._config import BufferAttr, BufferId, LayerAttr, SlotDesc, StorageConfig from ._storage._core import ( @@ -80,6 +81,9 @@ typed_range, ) +if TYPE_CHECKING: + from ._event_manager import KVCacheEventManager + class CacheLevelManager: __slots__ = ("cache_level", "storage", "controller") @@ -177,6 +181,8 @@ class StorageManager: "_slot_desc_list", "_levels", "_min_slots", + "_execution_stream", + "_event_manager", "__rawref__", ) _life_cycles: LifeCycleRegistry @@ -189,6 +195,8 @@ class StorageManager: _slot_desc_list: TypedIndexList[PoolGroupIndex, SlotDesc] _levels: TypedIndexList[CacheLevel, CacheLevelManager] _min_slots: TypedIndexList[PoolGroupIndex, int] + _execution_stream: CudaStream | None + _event_manager: "KVCacheEventManager | None" __rawref__: rawref.ref["StorageManager"] def __init__( @@ -199,8 +207,11 @@ def __init__( swa_scratch_reuse: SwaScratchReuseConfig | None, typical_batch: BatchDesc | None = None, constraints: list[BatchDesc] | None = None, + event_manager: "KVCacheEventManager | None" = None, ) -> None: self.__rawref__ = rawref.NULL + self._execution_stream: CudaStream | None = None + self._event_manager = event_manager assert config.cache_tiers[GPU_LEVEL].tier == CacheTier.GPU_MEM, ( "The first cache tier must be GPU memory" ) @@ -528,6 +539,13 @@ def _batched_migrate( for pool_idx, tasks in typed_enumerate(tasks_per_pool): batched_copy(dst_tier, src_tier, slot_sizes[pool_idx], tasks, stream.get()) finish_event = stream.take_finish_event() + emit_cache_level_updates = ( + update_src + and not defrag + and src_level != dst_level + and self._event_manager is not None + ) + emitted_update_keys: set[tuple[bytes, LifeCycleId]] = set() for src, dst in zip(src_pages, dst_slots): dst.ready_event = finish_event src.ready_event = ( @@ -540,6 +558,10 @@ def _batched_migrate( src_pool_group.release(src) src.set_slot(dst) src.cache_level = dst_level + if emit_cache_level_updates: + self._emit_cache_level_updated_event( + src, src_level, dst_level, emitted_update_keys + ) if scheduled_for_eviction: self.schedule_for_eviction(src) return None if update_src else dst_slots @@ -548,6 +570,34 @@ def _batched_migrate( dst_pool_group.release(s) raise + def _emit_cache_level_updated_event( + self, + page: Page, + old_level: CacheLevel, + new_level: CacheLevel, + emitted_keys: set[tuple[bytes, LifeCycleId]], + ) -> None: + if self._event_manager is None or not isinstance(page, CommittedPage): + return + + block = page.block() + if block is None or block.is_orphan: + return + + event_key = (block.key, page.life_cycle) + if event_key in emitted_keys: + return + + emitted_keys.add(event_key) + self._event_manager.add_updated_event( + block.key, + cache_level=KVCacheEventDiff( + old_value=int(old_level), + new_value=int(new_level), + ), + layer_group_id=int(page.life_cycle), + ) + def _pool_group( self, cache_level: CacheLevel, pool_group_index: PoolGroupIndex ) -> PoolGroupBase: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py b/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py index cd80e4197bb7..799520206298 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py @@ -76,6 +76,7 @@ "kv_cache_manager_v2/_config.py", "kv_cache_manager_v2/_copy_engine.py", "kv_cache_manager_v2/_cuda_virt_mem.py", + "kv_cache_manager_v2/_event_manager.py", "kv_cache_manager_v2/_exceptions.py", "kv_cache_manager_v2/_life_cycle_registry.py", "kv_cache_manager_v2/_page.py", diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index d45092117f6c..17278f9cb293 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -12,8 +12,9 @@ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._torch.pyexecutor.resource_manager import (KVCacheManager, - PeftCacheManager) +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + KVCacheManager, PeftCacheManager, + _warn_if_unsupported_v1_kv_cache_event_hash_algo) from tensorrt_llm.bindings import LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp from tensorrt_llm.bindings import executor as tllm @@ -24,6 +25,9 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig, PeftCacheConfig from tensorrt_llm.lora_helper import LoraConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_hash import (KV_CACHE_HASH_ALGO_AUTO, + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2) from tensorrt_llm.sampling_params import SamplingParams DataType = tensorrt_llm.bindings.DataType @@ -35,6 +39,33 @@ sys.path.append(str(root_dir / "tests" / "integration")) +def test_v1_kv_cache_event_hash_algo_warning_for_non_v1(): + with patch("tensorrt_llm._torch.pyexecutor.resource_manager.logger.warning" + ) as warning: + _warn_if_unsupported_v1_kv_cache_event_hash_algo(KV_CACHE_HASH_ALGO_V2) + + warning.assert_called_once() + assert KV_CACHE_HASH_ALGO_V1 in warning.call_args.args[0] + assert KV_CACHE_HASH_ALGO_V2 in warning.call_args.args[0] + + +def test_v1_kv_cache_event_hash_algo_no_warning_for_v1(): + with patch("tensorrt_llm._torch.pyexecutor.resource_manager.logger.warning" + ) as warning: + _warn_if_unsupported_v1_kv_cache_event_hash_algo(KV_CACHE_HASH_ALGO_V1) + + warning.assert_not_called() + + +def test_v1_kv_cache_event_hash_algo_no_warning_for_auto(): + with patch("tensorrt_llm._torch.pyexecutor.resource_manager.logger.warning" + ) as warning: + _warn_if_unsupported_v1_kv_cache_event_hash_algo( + KV_CACHE_HASH_ALGO_AUTO) + + warning.assert_not_called() + + class TestResourceManager(unittest.TestCase): CPP_RESOURCES_DIR = os.path.join(str(root_dir), "cpp", "tests", "resources") CPP_DATA_DIR = os.path.join(CPP_RESOURCES_DIR, "data") diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py new file mode 100644 index 000000000000..1dd025caa306 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -0,0 +1,896 @@ +# 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. + +import gc +import os +import threading +import time +from importlib.util import find_spec +from typing import TYPE_CHECKING, cast + +import pytest + +from tensorrt_llm._utils import KVCacheEventSerializer +from tensorrt_llm.runtime.kv_cache_hash import ( + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + truncate_sha256_hash_to_int64, +) +from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import ( + KVCacheEventDiff, + KVCacheEventManager, + KVCacheStoredBlockData, + UniqueToken, +) + +if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2 import CacheLevel, CudaStream, KVCacheManager, TokenId + from kv_cache_manager_v2._block_radix_tree import Block, RootBlock + from kv_cache_manager_v2._utils import CachedCudaStream, init_cuda_once, temporary_sys_path +else: + from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + CacheLevel, + CudaStream, + KVCacheManager, + TokenId, + ) + from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import Block, RootBlock + from tensorrt_llm.runtime.kv_cache_manager_v2._utils import ( + CachedCudaStream, + init_cuda_once, + temporary_sys_path, + ) + +try: + import torch +except ImportError: + torch = None + + +class _FakePage: + def __init__(self, cache_level=CacheLevel(0), priority=0): + self.cache_level = cache_level + self.priority = priority + + +class _FakePageRef: + def __init__(self, page): + self._page = page + + def __call__(self): + return self._page + + +class _FakeRootBlock: + ordinal = -1 + + +class _FakeBlock: + def __init__(self, key, tokens, num_life_cycles=1, prev=None): + self.key = key + self.tokens = tokens + self.prev = prev or _FakeRootBlock() + self.ordinal = getattr(self.prev, "ordinal", -1) + 1 + self.storage = [_FakePageRef(_FakePage()) for _ in range(num_life_cycles)] + + +with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): + from test_kv_cache_manager_v2 import create_config + + +def _create_test_manager( + event_manager, + *, + tokens_per_block=4, + gpu_quota=16 << 20, + host_quota=0, + window_size=None, + kv_buf_size=8192, +): + return KVCacheManager( + create_config( + tokens_per_block=tokens_per_block, + gpu_quota=gpu_quota, + host_quota=host_quota, + disk_quota=0, + num_layers=2, + window_size=window_size, + sink_tokens=0, + kv_buf_size=kv_buf_size, + ), + event_manager=event_manager, + ) + + +def _token_ids(start, end): + return [TokenId(token_id) for token_id in range(start, end)] + + +def _flush_serialized_events(event_manager): + event_manager.flush_iteration_events() + return KVCacheEventSerializer.serialize(event_manager.get_latest_events(0)) + + +def _stored_events(events): + return [event for event in events if event["data"]["type"] == "stored"] + + +def _stored_block_hashes(events): + return [ + block["block_hash"] for event in _stored_events(events) for block in event["data"]["blocks"] + ] + + +def _stored_block_hashes_by_layer_group(events): + return { + event["layer_group_id"]: [block["block_hash"] for block in event["data"]["blocks"]] + for event in _stored_events(events) + } + + +def _commit_and_close(manager, stream, tokens, *, input_tokens=None): + kv_cache = manager.create_kv_cache(input_tokens=input_tokens) + assert kv_cache.resume(stream) + kv_cache.capacity = len(input_tokens or []) + len(tokens) + kv_cache.commit(tokens) + kv_cache.close() + del kv_cache + gc.collect() + + +def test_v2_kv_cache_event_manager_serialization(): + event_manager = KVCacheEventManager(max_kv_event_entries=4, window_size=128) + event_manager.add_created_event([2, 3]) + event_manager.add_stored_event( + parent_hash=None, + blocks=[ + KVCacheStoredBlockData( + block_hash="abcd", + tokens=[UniqueToken(1), UniqueToken(2)], + cache_level=0, + priority=0, + ) + ], + ) + event_manager.add_removed_event("abcd") + + event_manager.flush_iteration_events() + events = KVCacheEventSerializer.serialize(event_manager.get_latest_events()) + + assert [event["event_id"] for event in events] == [0, 1, 2] + assert [event["hash_algo"] for event in events] == ["v2_sha256"] * 3 + assert events[0]["window_size"] == 128 + assert events[0]["data"] == { + "type": "created", + "num_blocks_per_cache_level": [2, 3], + } + assert events[1]["data"]["type"] == "stored" + assert events[1]["data"]["parent_hash"] is None + assert events[1]["data"]["blocks"][0]["block_hash"] == "abcd" + assert events[1]["data"]["blocks"][0]["tokens"][0] == { + "type": "unique_token", + "token_id": 1, + "token_extra_id": 0, + } + assert events[1]["data"]["blocks"][0]["mm_keys"] == [] + assert events[2]["data"] == { + "type": "removed", + "block_hashes": ["abcd"], + } + assert event_manager.get_latest_events(0) == [] + + +def test_v2_kv_cache_event_manager_default_get_latest_events_waits(): + event_manager = KVCacheEventManager(max_kv_event_entries=4, window_size=128) + events = [] + + def get_events(): + events.extend(event_manager.get_latest_events()) + + thread = threading.Thread(target=get_events) + thread.start() + time.sleep(0.05) + + assert thread.is_alive() + + event_manager.add_created_event([1]) + event_manager.flush_iteration_events() + thread.join(timeout=1) + + assert not thread.is_alive() + assert len(events) == 1 + assert events[0].data.num_blocks_per_cache_level == [1] + + +def test_v2_kv_cache_event_manager_uses_layer_group_window_size(): + event_manager = KVCacheEventManager( + max_kv_event_entries=4, + window_size=4096, + window_size_by_layer_group={ + 0: 128, + 1: 4096, + }, + ) + + event_manager.add_created_event([1], layer_group_ids=[0, 1]) + event_manager.flush_iteration_events() + events = event_manager.get_latest_events(0) + + assert [event.layer_group_id for event in events] == [0, 1] + assert [event.window_size for event in events] == [128, 4096] + + +def test_truncate_sha256_hash_to_int64_uses_unsigned_value(): + assert truncate_sha256_hash_to_int64(b"\x80" + b"\x00" * 31) == 1 << 63 + assert truncate_sha256_hash_to_int64(b"\xff" * 32) == (1 << 64) - 1 + + +def test_v2_kv_cache_event_diff_positional_order_matches_v1(): + diff = KVCacheEventDiff(3, 5) + + assert diff.old_value == 3 + assert diff.new_value == 5 + + +def test_v2_kv_cache_event_manager_drops_oldest_events(): + event_manager = KVCacheEventManager(max_kv_event_entries=2, window_size=128) + event_manager.add_created_event([1]) + event_manager.add_stored_event( + parent_hash=None, + blocks=[ + KVCacheStoredBlockData( + block_hash="stored", + tokens=[UniqueToken(1)], + cache_level=0, + priority=0, + ) + ], + ) + event_manager.add_removed_event("a") + event_manager.add_removed_event("b") + + event_manager.flush_iteration_events() + events = KVCacheEventSerializer.serialize(event_manager.get_latest_events()) + + assert [event["event_id"] for event in events] == [1, 2] + assert events[0]["data"]["blocks"][0]["block_hash"] == "stored" + assert events[1]["data"]["block_hashes"] == ["a", "b"] + + +def test_v2_kv_cache_event_manager_accepts_removed_iterables(): + event_manager = KVCacheEventManager(max_kv_event_entries=2, window_size=128) + event_manager.add_removed_event(["a", "b"]) + + event_manager.flush_iteration_events() + events = KVCacheEventSerializer.serialize(event_manager.get_latest_events()) + + assert len(events) == 1 + assert events[0]["data"] == { + "type": "removed", + "block_hashes": ["a", "b"], + } + + +def test_v2_kv_cache_event_manager_coalesces_contiguous_stored_events(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block0 = _FakeBlock(b"\xab\xcd", [1, 2], num_life_cycles=2) + block1 = _FakeBlock(b"\xab\xce", [3, 4], num_life_cycles=2, prev=block0) + + event_manager.add_stored_block_event_from_block(block0) + event_manager.add_stored_block_event_from_block(block1) + + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == ["stored", "stored"] + assert [event["layer_group_id"] for event in events] == [0, 1] + assert [block["block_hash"] for block in events[0]["data"]["blocks"]] == [ + "abcd", + "abce", + ] + assert [block["block_hash"] for block in events[1]["data"]["blocks"]] == [ + "abcd", + "abce", + ] + assert events[0]["data"]["parent_hash"] is None + assert events[1]["data"]["parent_hash"] is None + + +def test_v2_kv_cache_event_manager_serializes_layer_group_id(): + event_manager = KVCacheEventManager(max_kv_event_entries=4, window_size=128) + event_manager.add_created_event([2, 3], [0, 1]) + + events = _flush_serialized_events(event_manager) + + assert [event["layer_group_id"] for event in events] == [0, 1] + assert [event["data"] for event in events] == [ + { + "type": "created", + "num_blocks_per_cache_level": [2, 3], + }, + { + "type": "created", + "num_blocks_per_cache_level": [2, 3], + }, + ] + + +def test_v2_kv_cache_event_manager_sha256_64_compatibility_mode(): + event_manager = KVCacheEventManager( + max_kv_event_entries=8, + window_size=128, + hash_algo=KV_CACHE_HASH_ALGO_V2_SHA256_64, + ) + block0 = _FakeBlock(bytes.fromhex("8000000000000001" + "00" * 24), [1, 2]) + block1 = _FakeBlock(bytes.fromhex("0102030405060708" + "00" * 24), [3, 4], prev=block0) + + event_manager.add_stored_block_event_from_block(block0) + event_manager.add_stored_block_event_from_block(block1) + event_manager.add_removed_event(block0.key) + events = _flush_serialized_events(event_manager) + + expected_block0_hash = truncate_sha256_hash_to_int64(block0.key) + expected_block1_hash = truncate_sha256_hash_to_int64(block1.key) + assert [event["hash_algo"] for event in events] == [ + KV_CACHE_HASH_ALGO_V2_SHA256_64, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + ] + assert events[0]["data"]["type"] == "stored" + assert events[0]["data"]["parent_hash"] is None + assert [block["block_hash"] for block in events[0]["data"]["blocks"]] == [ + expected_block0_hash, + expected_block1_hash, + ] + assert events[1]["data"] == { + "type": "removed", + "block_hashes": [expected_block0_hash], + } + assert all(isinstance(block_hash, int) for block_hash in events[1]["data"]["block_hashes"]) + + +def test_v2_kv_cache_event_manager_v1_hash_algo_falls_back_to_sha256_64(): + event_manager = KVCacheEventManager( + max_kv_event_entries=4, + window_size=128, + hash_algo=KV_CACHE_HASH_ALGO_V1, + ) + block = _FakeBlock(bytes.fromhex("0102030405060708" + "00" * 24), [1, 2]) + + event_manager.add_stored_block_event_from_block(block) + events = _flush_serialized_events(event_manager) + + expected_block_hash = truncate_sha256_hash_to_int64(block.key) + assert events[0]["hash_algo"] == KV_CACHE_HASH_ALGO_V2_SHA256_64 + assert events[0]["data"]["blocks"][0]["block_hash"] == expected_block_hash + + +def test_v2_kv_cache_event_manager_unknown_hash_algo_raises(): + with pytest.raises(ValueError, match="Unsupported V2 KV cache event hash algorithm"): + KVCacheEventManager( + max_kv_event_entries=4, + window_size=128, + hash_algo="unknown_hash_algo", + ) + + +def test_v2_kv_cache_event_manager_gathers_attention_dp_events_on_rank_zero(): + remote_manager = KVCacheEventManager( + max_kv_event_entries=4, + window_size=128, + attention_dp_rank=1, + ) + remote_manager.add_created_event([3]) + remote_manager.flush_iteration_events() + remote_event_objects = remote_manager.get_latest_events() + remote_events = KVCacheEventSerializer.serialize(remote_event_objects) + assert remote_events[0]["attention_dp_rank"] == 1 + + gathered_local_events = [] + + def gather(local_events): + gathered_local_events.extend(local_events) + return [ + local_events, + remote_event_objects, + ] + + event_manager = KVCacheEventManager( + max_kv_event_entries=4, + window_size=128, + attention_dp_rank=0, + attention_dp_gather=gather, + ) + event_manager.add_created_event([2]) + + events = _flush_serialized_events(event_manager) + + assert len(gathered_local_events) == 1 + assert [event["attention_dp_rank"] for event in events] == [0, 1] + assert [event["data"]["num_blocks_per_cache_level"] for event in events] == [ + [2], + [3], + ] + + +def test_v2_kv_cache_event_manager_attention_dp_nonzero_rank_sends_without_publishing(): + gathered_local_events = [] + + def gather(local_events): + gathered_local_events.extend(local_events) + return [ + [], + local_events, + ] + + event_manager = KVCacheEventManager( + max_kv_event_entries=1, + window_size=128, + attention_dp_rank=1, + attention_dp_gather=gather, + ) + event_manager.add_created_event([2]) + event_manager.add_created_event([3]) + + events = _flush_serialized_events(event_manager) + + assert len(gathered_local_events) == 1 + assert gathered_local_events[0].attention_dp_rank == 1 + assert KVCacheEventSerializer.serialize(gathered_local_events)[0]["data"][ + "num_blocks_per_cache_level" + ] == [3] + assert events == [] + + +def test_v2_kv_cache_event_manager_attention_dp_rank_zero_uses_dp_capacity(): + remote_manager = KVCacheEventManager( + max_kv_event_entries=2, + window_size=128, + attention_dp_rank=1, + ) + remote_manager.add_created_event([10]) + remote_manager.add_created_event([11]) + remote_manager.flush_iteration_events() + remote_event_objects = remote_manager.get_latest_events() + + def gather(local_events): + return [ + local_events, + remote_event_objects, + ] + + event_manager = KVCacheEventManager( + max_kv_event_entries=1, + window_size=128, + attention_dp_rank=0, + attention_dp_gather=gather, + ) + event_manager.add_created_event([1]) + event_manager.add_created_event([2]) + + events = _flush_serialized_events(event_manager) + + assert [event["attention_dp_rank"] for event in events] == [0, 1] + assert [event["data"]["num_blocks_per_cache_level"] for event in events] == [ + [2], + [11], + ] + + +def test_v2_kv_cache_event_manager_serializes_updated_event(): + event_manager = KVCacheEventManager(max_kv_event_entries=2, window_size=128) + event_manager.add_updated_event( + "abcd", + cache_level=KVCacheEventDiff(old_value=0, new_value=1), + priority=KVCacheEventDiff(old_value=1, new_value=2), + ) + + event_manager.flush_iteration_events() + events = KVCacheEventSerializer.serialize(event_manager.get_latest_events()) + + assert len(events) == 1 + assert events[0]["hash_algo"] == "v2_sha256" + assert events[0]["data"] == { + "type": "updated", + "block_hash": "abcd", + "cache_level": { + "type": "event_diff", + "new_value": 1, + "old_value": 0, + }, + "priority": { + "type": "event_diff", + "new_value": 2, + "old_value": 1, + }, + } + + +def test_v2_kv_cache_event_manager_uses_stored_registry_for_removed_event(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xcd", [1, 2]) + + event_manager.add_stored_block_event_from_block(block) + block.storage = [] + event_manager.add_removed_event(block.key) + event_manager.add_removed_event(block.key) + + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == ["stored", "removed"] + assert [event["layer_group_id"] for event in events] == [0, 0] + assert events[0]["data"]["blocks"][0]["block_hash"] == "abcd" + assert "layer_groups" not in events[0]["data"]["blocks"][0] + assert events[1]["data"]["block_hashes"] == ["abcd"] + assert "layer_groups" not in events[1]["data"] + + +def test_v2_kv_cache_event_manager_emits_partial_life_cycle_removed_events(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xcd", [1, 2], num_life_cycles=2) + + event_manager.add_stored_block_event_from_block(block) + event_manager.add_removed_life_cycle_event(block.key, 0) + + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == [ + "stored", + "stored", + "removed", + ] + assert [event["layer_group_id"] for event in events] == [0, 1, 0] + assert events[0]["data"]["blocks"][0]["block_hash"] == "abcd" + assert events[1]["data"]["blocks"][0]["block_hash"] == "abcd" + assert "layer_groups" not in events[0]["data"]["blocks"][0] + assert "layer_groups" not in events[1]["data"]["blocks"][0] + assert events[2]["data"]["block_hashes"] == ["abcd"] + assert "layer_groups" not in events[2]["data"] + + event_manager.add_removed_life_cycle_event(block.key, 1) + event_manager.add_removed_life_cycle_event(block.key, 1) + + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == ["removed"] + assert events[0]["layer_group_id"] == 1 + assert events[0]["data"]["block_hashes"] == ["abcd"] + assert "layer_groups" not in events[0]["data"] + + +def test_v2_kv_cache_event_manager_whole_block_removal_clears_life_cycle_state(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xce", [1, 2], num_life_cycles=2) + + event_manager.add_stored_block_event_from_block(block) + event_manager.add_removed_life_cycle_event(block.key, 0) + event_manager.add_removed_event(block.key) + event_manager.add_removed_event(block.key) + + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == [ + "stored", + "stored", + "removed", + "removed", + ] + assert [event["layer_group_id"] for event in events] == [0, 1, 0, 1] + assert events[2]["data"]["block_hashes"] == ["abce"] + assert events[3]["data"]["block_hashes"] == ["abce"] + + +def test_v2_kv_cache_event_manager_readds_life_cycle_emits_stored_event(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xcf", [1, 2], num_life_cycles=2) + + event_manager.add_stored_block_event_from_block(block) + event_types = [event["data"]["type"] for event in _flush_serialized_events(event_manager)] + assert event_types == ["stored", "stored"] + + event_manager.add_removed_life_cycle_event(block.key, 0) + event_manager.add_stored_life_cycle_event_from_block(block, 0) + event_manager.add_removed_life_cycle_event(block.key, 1) + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == [ + "removed", + "stored", + "removed", + ] + assert [event["layer_group_id"] for event in events] == [0, 0, 1] + assert events[0]["data"]["block_hashes"] == ["abcf"] + assert events[1]["data"]["blocks"][0]["block_hash"] == "abcf" + assert "layer_groups" not in events[1]["data"]["blocks"][0] + assert events[2]["data"]["block_hashes"] == ["abcf"] + + event_manager.add_removed_life_cycle_event(block.key, 0) + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == ["removed"] + assert events[0]["layer_group_id"] == 0 + assert events[0]["data"]["block_hashes"] == ["abcf"] + assert "layer_groups" not in events[0]["data"] + + +def test_v2_kv_cache_event_manager_reemits_stored_after_all_life_cycles_were_removed(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xd0", [1, 2]) + + event_manager.add_stored_block_event_from_block(block) + event_manager.add_removed_life_cycle_event(block.key, 0) + event_types = [event["data"]["type"] for event in _flush_serialized_events(event_manager)] + assert event_types == ["stored", "removed"] + + event_manager.add_stored_life_cycle_event_from_block(block, 0) + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == ["stored"] + assert events[0]["layer_group_id"] == 0 + assert events[0]["data"]["blocks"][0]["block_hash"] == "abd0" + assert "layer_groups" not in events[0]["data"]["blocks"][0] + + +def test_v2_kv_cache_event_manager_flushes_removed_before_updated_event(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + removed_block = _FakeBlock(b"\xab\xd1", [1, 2]) + updated_block = _FakeBlock(b"\xab\xd2", [3, 4]) + + event_manager.add_stored_block_event_from_block(removed_block) + event_manager.add_stored_block_event_from_block(updated_block) + _flush_serialized_events(event_manager) + + event_manager.add_removed_event(removed_block.key) + event_manager.add_updated_event( + updated_block.key, + cache_level=KVCacheEventDiff(old_value=0, new_value=1), + layer_group_id=0, + ) + events = _flush_serialized_events(event_manager) + + assert [event["data"]["type"] for event in events] == ["removed", "updated"] + assert events[0]["data"]["block_hashes"] == ["abd1"] + assert events[1]["data"]["block_hash"] == "abd2" + + +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires CUDA") +def test_v2_stored_events_match_block_hash_chain(): + init_cuda_once() + gc.collect() + gc.disable() + + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) + manager = None + try: + tokens_per_block = 4 + manager = _create_test_manager(event_manager, tokens_per_block=tokens_per_block) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + tokens = _token_ids(0, 2 * tokens_per_block) + + _commit_and_close(manager, stream, tokens) + events = _flush_serialized_events(event_manager) + + stored_events = _stored_events(events) + assert len(stored_events) == 1 + + root_key = RootBlock.make_key(None) + block0_key = Block.make_key(root_key, tokens[:tokens_per_block]) + block1_key = Block.make_key(block0_key, tokens[tokens_per_block:]) + expected_hashes = [block0_key.hex(), block1_key.hex()] + + assert _stored_block_hashes(stored_events) == expected_hashes + assert stored_events[0]["data"]["parent_hash"] is None + assert [ + block["block_hash"] for block in stored_events[0]["data"]["blocks"] + ] == expected_hashes + assert [ + token["token_id"] for token in stored_events[0]["data"]["blocks"][0]["tokens"] + ] == list(range(tokens_per_block)) + assert [ + token["token_id"] for token in stored_events[0]["data"]["blocks"][1]["tokens"] + ] == list(range(tokens_per_block, 2 * tokens_per_block)) + finally: + gc.enable() + if manager is not None: + manager.shutdown() + + +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires CUDA") +def test_v2_reused_prefix_does_not_emit_duplicate_stored_events(): + init_cuda_once() + gc.collect() + gc.disable() + + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) + manager = None + try: + tokens_per_block = 4 + manager = _create_test_manager(event_manager, tokens_per_block=tokens_per_block) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prefix_tokens = _token_ids(0, 2 * tokens_per_block) + new_tokens = _token_ids(2 * tokens_per_block, 3 * tokens_per_block) + + _commit_and_close(manager, stream, prefix_tokens) + first_events = _flush_serialized_events(event_manager) + prefix_hashes = _stored_block_hashes(first_events) + assert len(prefix_hashes) == 2 + + _commit_and_close( + manager, + stream, + new_tokens, + input_tokens=prefix_tokens, + ) + reuse_events = _flush_serialized_events(event_manager) + reused_hashes = _stored_block_hashes(reuse_events) + + root_key = RootBlock.make_key(None) + block0_key = Block.make_key(root_key, prefix_tokens[:tokens_per_block]) + block1_key = Block.make_key(block0_key, prefix_tokens[tokens_per_block:]) + block2_key = Block.make_key(block1_key, new_tokens) + + assert prefix_hashes == [block0_key.hex(), block1_key.hex()] + assert reused_hashes == [block2_key.hex()] + assert not (set(prefix_hashes) & set(reused_hashes)) + finally: + gc.enable() + if manager is not None: + manager.shutdown() + + +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires CUDA") +def test_v2_removed_events_match_stored_hashes(): + init_cuda_once() + gc.collect() + gc.disable() + + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=128) + manager = None + try: + tokens_per_block = 4 + manager = _create_test_manager(event_manager, tokens_per_block=tokens_per_block) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + tokens = _token_ids(0, 2 * tokens_per_block) + + _commit_and_close(manager, stream, tokens) + stored_events = _flush_serialized_events(event_manager) + stored_hashes = _stored_block_hashes(stored_events) + assert len(stored_hashes) == 2 + + manager.clear_reusable_blocks() + removal_events = _flush_serialized_events(event_manager) + removed_hashes = [ + block_hash + for event in removal_events + if event["data"]["type"] == "removed" + for block_hash in event["data"]["block_hashes"] + ] + + assert set(removed_hashes) == set(stored_hashes) + assert len(removed_hashes) == len(stored_hashes) + finally: + gc.enable() + if manager is not None: + manager.shutdown() + + +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires CUDA") +def test_v2_removed_event_emitted_when_last_level_page_is_dropped(): + init_cuda_once() + gc.collect() + gc.disable() + + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=4096) + manager = None + try: + tokens_per_block = 8 + manager = _create_test_manager( + event_manager, + tokens_per_block=tokens_per_block, + gpu_quota=8 << 20, + window_size=4096, + kv_buf_size=1 << 20, + ) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + kv_cache = manager.create_kv_cache() + assert kv_cache.resume(stream) + kv_cache.capacity = tokens_per_block * 2 + kv_cache.commit(_token_ids(0, tokens_per_block * 2)) + + stored_events = _flush_serialized_events(event_manager) + stored_hashes_by_layer_group = _stored_block_hashes_by_layer_group(stored_events) + assert set(stored_hashes_by_layer_group) == {0, 1} + assert stored_hashes_by_layer_group[0] == stored_hashes_by_layer_group[1] + + kv_cache.close() + del kv_cache + gc.collect() + + assert manager.resize(CacheLevel(0), 4 << 20) + removal_events = _flush_serialized_events(event_manager) + removed_hashes_by_layer_group = { + event["layer_group_id"]: event["data"]["block_hashes"] + for event in removal_events + if event["data"]["type"] == "removed" + } + + assert set(removed_hashes_by_layer_group) == {0, 1} + assert removed_hashes_by_layer_group[0] == removed_hashes_by_layer_group[1] + for layer_group_id, removed_hashes in removed_hashes_by_layer_group.items(): + assert len(removed_hashes) == 1 + assert removed_hashes[0] in stored_hashes_by_layer_group[layer_group_id] + finally: + gc.enable() + if manager is not None: + manager.shutdown() + + +@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires CUDA") +def test_v2_kv_cache_event_manager_emits_updated_on_level_migration(): + init_cuda_once() + gc.collect() + gc.disable() + + event_manager = KVCacheEventManager(max_kv_event_entries=16, window_size=4096) + manager = None + try: + tokens_per_block = 8 + manager = _create_test_manager( + event_manager, + tokens_per_block=tokens_per_block, + gpu_quota=8 << 20, + host_quota=8 << 20, + window_size=4096, + kv_buf_size=1 << 20, + ) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + kv_cache = manager.create_kv_cache() + assert kv_cache.resume(stream) + kv_cache.capacity = tokens_per_block * 2 + kv_cache.commit([TokenId(token_id) for token_id in range(tokens_per_block * 2)]) + + stored_events = _flush_serialized_events(event_manager) + stored_hashes_by_layer_group = _stored_block_hashes_by_layer_group(stored_events) + assert set(stored_hashes_by_layer_group) == {0, 1} + assert stored_hashes_by_layer_group[0] == stored_hashes_by_layer_group[1] + + kv_cache.close() + del kv_cache + gc.collect() + + assert manager.resize(CacheLevel(0), 4 << 20) + + event_manager.flush_iteration_events() + events = KVCacheEventSerializer.serialize(event_manager.get_latest_events()) + updated_events = [event for event in events if event["data"]["type"] == "updated"] + + assert len(updated_events) == 2 + assert {event["layer_group_id"] for event in updated_events} == {0, 1} + assert len({event["data"]["block_hash"] for event in updated_events}) == 1 + assert updated_events[0]["data"]["block_hash"] in stored_hashes_by_layer_group[0] + for event in updated_events: + assert event["data"]["cache_level"] == { + "type": "event_diff", + "old_value": 0, + "new_value": 1, + } + assert event["data"]["priority"] is None + finally: + gc.enable() + if manager is not None: + manager.shutdown() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 5b73c0e60f7f..d31a7726d682 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -331,6 +331,8 @@ def get_model_defaults(cls, llm_args): def test_KvCacheConfig_declaration(): + assert KvCacheConfig().kv_cache_event_hash_algo == "auto" + config = KvCacheConfig(enable_block_reuse=True, max_tokens=1024, max_attention_window=[1024, 1024, 1024], @@ -341,6 +343,7 @@ def test_KvCacheConfig_declaration(): cross_kv_cache_fraction=0.5, secondary_offload_min_priority=1, event_buffer_max_size=0, + kv_cache_event_hash_algo="v2_sha256_64", enable_partial_reuse=True, copy_on_partial_reuse=True, attention_dp_events_gather_period_ms=10) @@ -356,6 +359,11 @@ def test_KvCacheConfig_declaration(): assert pybind_config.cross_kv_cache_fraction == 0.5 assert pybind_config.secondary_offload_min_priority == 1 assert pybind_config.event_buffer_max_size == 0 + assert config.kv_cache_event_hash_algo == "v2_sha256_64" + assert KvCacheConfig( + kv_cache_event_hash_algo="auto").kv_cache_event_hash_algo == "auto" + assert KvCacheConfig(kv_cache_event_hash_algo="v1_block_key" + ).kv_cache_event_hash_algo == "v1_block_key" assert pybind_config.enable_partial_reuse == True assert pybind_config.copy_on_partial_reuse == True assert pybind_config.attention_dp_events_gather_period_ms == 10 diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index ee002905d26c..268ef50b867e 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -24,6 +24,7 @@ serialize_item) from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_hash import KV_CACHE_HASH_ALGO_V2 from tensorrt_llm.sampling_params import SamplingParams from tensorrt_llm.scheduling_params import SchedulingParams @@ -34,7 +35,9 @@ global_kvcache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, event_buffer_max_size=1024, enable_block_reuse=True, - max_tokens=256) + onboard_blocks=True, + max_tokens=256, + use_kv_cache_manager_v2=False) def create_kv_cache_manager(): @@ -66,6 +69,17 @@ def create_llm(tensor_parallel_size=1): enable_autotuner=False) +def create_v2_llm(): + v2_kvcache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, + event_buffer_max_size=1024, + enable_block_reuse=True, + max_tokens=256, + use_kv_cache_manager_v2=True) + return LLM(model=llama_model_path, + kv_cache_config=v2_kvcache_config, + enable_autotuner=False) + + def create_llm_request(id, input_tokens, new_tokens=1): sampling_params = SamplingParams() req = LlmRequest(request_id=id, @@ -892,6 +906,38 @@ def test_expected_kv_cache_events(): assert event["data"]["type"] == "stored" +def test_expected_v2_kv_cache_events(): + llm = create_v2_llm() + sampling_params = SamplingParams(max_tokens=6, temperature=0.01) + prompt = list(range(127)) + + _ = llm.generate(prompt, sampling_params=sampling_params) + + events = llm.get_kv_cache_events(5) + assert events and len(events) >= 2 + assert all(event["hash_algo"] == KV_CACHE_HASH_ALGO_V2 for event in events + if event) + + created_events = [ + event for event in events if event["data"]["type"] == "created" + ] + stored_events = [ + event for event in events if event["data"]["type"] == "stored" + ] + assert created_events + assert stored_events + assert created_events[0]["event_id"] == 0 + + block_hashes = [ + block["block_hash"] for event in stored_events + for block in event["data"]["blocks"] + ] + assert block_hashes + assert all( + isinstance(block_hash, str) and len(block_hash) == 64 + for block_hash in block_hashes) + + def test_kv_cache_event_async_api(): llm = create_llm() sampling_params = SamplingParams(max_tokens=6, temperature=0.01) From 7f4c270ba0c685748dc48c89d3a0191b6c34210a Mon Sep 17 00:00:00 2001 From: lishicheng1996-nv Date: Thu, 30 Apr 2026 15:34:06 +0800 Subject: [PATCH 09/58] [None][feat] Fuse FP8 1x128 quantize + UE8M0 scale pack on SM100 (#13628) Signed-off-by: Shicheng Li Signed-off-by: Fanrong Li (cherry picked from commit 096f86eca51bfc74e79a604525f78848ba235674) Signed-off-by: Fanrong Li --- .../fp8_blockscale_quant_packed.cu | 196 ++++++++++++++++++ .../fp8_blockscale_quant_packed.h | 52 +++++ cpp/tensorrt_llm/thop/fp8Quantize.cpp | 70 +++++++ .../_torch/custom_ops/cpp_custom_ops.py | 11 + .../_torch/custom_ops/torch_custom_ops.py | 22 +- 5 files changed, 348 insertions(+), 3 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu create mode 100644 cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu new file mode 100644 index 000000000000..f1a24f6de17e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu @@ -0,0 +1,196 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Fused 1x128 FP8 quantize + UE8M0 scale packing. +// +// Replaces the (scale_1x128_kernel + pack_fp32_into_ue8m0) two-kernel sequence +// used by SM100 deep_gemm fp8 block-scale GEMMs. Adapted from the SM120 MoE +// in-kernel packing pattern (`scale_1x128_kernel_sm120` in +// sm120_blockwise_gemm/sm120_fp8_moe_gemm_1d1d.cuh), specialised for the +// non-MoE case (single contiguous batch, no token offsets). + +#include "fp8_blockscale_quant_packed.h" + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::fp8_blockscale_gemm +{ + +namespace +{ + +__device__ __forceinline__ float reciprocal_approximate_ftz_local(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// Each warp consumes one row × 4 quantization blocks (4 × 128 = 512 K elems). +// 32 lanes split into 4 lane-groups of 8: each group covers 1 quant block +// (8 lanes × 16 BF16 elems = 128 elems). After per-block amax, lanes +// 0/8/16/24 each hold one UE8M0 scale byte; lane 0 packs them into a uint32 +// and stores in the deep_gemm-expected MN-major layout. +template +__global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict__ fp8_output, + int32_t* __restrict__ packed_scale_output, __nv_bfloat16 const* __restrict__ input, int const m, int const k, + int const scale_leading_dim_uint32) +{ + int const packed_sf_k_idx = static_cast(blockIdx.x); + int const warp_id = static_cast(threadIdx.x) >> 5; + int const lane_id = static_cast(threadIdx.x) & 31; + int const m_idx = static_cast(blockIdx.y) * WarpsPerBlock + warp_id; + + if (m_idx >= m) + { + return; + } + + int const k_base = packed_sf_k_idx * 512 + lane_id * 16; + + // ---- 1. Load 16 BF16 elements per lane. ---- + auto const* in_ptr = reinterpret_cast(input + static_cast(m_idx) * k + k_base); + constexpr int kLoadNumElems = sizeof(double4) / sizeof(__nv_bfloat16); // 16 + + union LoadTrick + { + double4 pack; + __nv_bfloat16 v[kLoadNumElems]; + }; + + LoadTrick load_trick; + bool const k_in_range = (k_base < k); + load_trick.pack = k_in_range ? in_ptr[0] : double4{}; + + if (k_in_range && k_base + kLoadNumElems > k) + { + int const valid = k - k_base; +#pragma unroll + for (int i = 0; i < kLoadNumElems; ++i) + { + if (i >= valid) + { + load_trick.v[i] = __nv_bfloat16(0.0f); + } + } + } + + // ---- 2. Per-block amax (lanes 0..7 / 8..15 / 16..23 / 24..31 = 4 quant blocks). ---- + __nv_bfloat16 max_elem = __nv_bfloat16(0.0f); +#pragma unroll + for (int i = 0; i < kLoadNumElems; ++i) + { + max_elem = __hmax(max_elem, __habs(load_trick.v[i])); + } + float amax = static_cast(max_elem); + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 4, 8)); + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 2, 8)); + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFFu, amax, 1, 8)); + amax = fmaxf(amax, 1e-10f); + + // ---- 3. UE8M0 dequant scale. ---- + float const dequant_scale_raw = amax * reciprocal_approximate_ftz_local(448.0f); + __nv_fp8_e8m0 ue8m0_scale; + ue8m0_scale.__x = __nv_cvt_float_to_e8m0(dequant_scale_raw, __NV_SATFINITE, cudaRoundPosInf); + + // Recover quant_scale = 1 / 2^(exp - 127) for fp8 conversion. + constexpr uint32_t FP32_EXPONENT_BIAS = 127u; + float const quant_scale = (ue8m0_scale.__x == 0) + ? 1.0f + : exp2f(static_cast(FP32_EXPONENT_BIAS) - static_cast(ue8m0_scale.__x)); + + // ---- 4. Quantize and store FP8 output. ---- + constexpr int kStoreNumElems = sizeof(float4) / sizeof(__nv_fp8_e4m3); // 16 + + union StoreTrick + { + float4 pack; + __nv_fp8_e4m3 v[kStoreNumElems]; + }; + + StoreTrick store_trick; + store_trick.pack = float4{}; +#pragma unroll + for (int i = 0; i < kStoreNumElems; ++i) + { + store_trick.v[i] = __nv_fp8_e4m3(static_cast(load_trick.v[i]) * quant_scale); + } + auto* out_ptr = reinterpret_cast(fp8_output + static_cast(m_idx) * k + k_base); + if (k_in_range) + { + if (k_base + kStoreNumElems > k) + { + int const valid = k - k_base; +#pragma unroll + for (int i = 0; i < kStoreNumElems; ++i) + { + if (i >= valid) + { + store_trick.v[i] = __nv_fp8_e4m3(0.0f); + } + } + } + out_ptr[0] = store_trick.pack; + } + + // ---- 5. Pack 4 UE8M0 scales (lanes 0/8/16/24) and store. ---- + uint32_t const s0 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 0); + uint32_t const s1 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 8); + uint32_t const s2 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 16); + uint32_t const s3 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 24); + if (lane_id == 0) + { + // Mask off scale bytes whose sf_k is past the actual K. + int const num_sf_k = (k + 127) / 128; + int const sf_k_base = packed_sf_k_idx * 4; + uint32_t packed = 0u; + if (sf_k_base + 0 < num_sf_k) + packed |= s0; + if (sf_k_base + 1 < num_sf_k) + packed |= (s1 << 8); + if (sf_k_base + 2 < num_sf_k) + packed |= (s2 << 16); + if (sf_k_base + 3 < num_sf_k) + packed |= (s3 << 24); + // Layout: packed_scale[packed_sf_k_idx, m_idx] + packed_scale_output[static_cast(packed_sf_k_idx) * scale_leading_dim_uint32 + m_idx] = packed; + } +} + +} // namespace + +void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, + __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream) +{ + constexpr int kWarpsPerBlock = 4; + int const num_packed_sf_k = (((k + 127) / 128) + 3) / 4; + int const m_blocks = (m + kWarpsPerBlock - 1) / kWarpsPerBlock; + dim3 const grid(num_packed_sf_k, m_blocks, 1); + dim3 const block(kWarpsPerBlock * 32, 1, 1); + fp8_quantize_1x128_packed_kernel_impl + <<>>(fp8_output, packed_scale_output, input, m, k, scale_leading_dim_uint32); +} + +} // namespace kernels::fp8_blockscale_gemm + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h new file mode 100644 index 000000000000..a4079cd9b054 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// Host-callable launcher for the fused FP8 1x128 quantize + UE8M0-pack kernel. +// Kernel implementation lives in fp8_blockscale_quant_packed.cu and is built +// by nvcc; this header is safe to include from .cpp files compiled by g++. + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::fp8_blockscale_gemm +{ + +// Launches the fused 1x128 FP8 quant + UE8M0 pack kernel. +// +// Inputs: +// input : BF16 [m, k] row-major contiguous +// Outputs: +// fp8_output : E4M3 [m, k] row-major contiguous +// packed_scale_output : uint32 [packed_sf_k, scale_leading_dim_uint32] +// where packed_sf_k = ceil(ceil(k/128)/4) +// +// `scale_leading_dim_uint32` is the (uint32) stride between consecutive +// packed_sf_k rows of the scale tensor; caller is responsible for choosing +// it (typically aligned to 4 uint32 = 16 bytes for TMA alignment). +void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, + __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream); + +} // namespace kernels::fp8_blockscale_gemm + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp index 9e94e02950f7..a05af0842a39 100644 --- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp8Quantize.cpp @@ -16,6 +16,7 @@ #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_gemm.h" +#include "tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h" #include "tensorrt_llm/thop/thUtils.h" #include @@ -141,6 +142,73 @@ std::tuple fp8_batched_quantize_1x128_permute102(at::Ten return {valueE4M3.slice(0, 0, b * m * n).view({b, m, n}), scaleFP8SF}; } + +// Fused 1x128 FP8 quantize + UE8M0 packing (SM100 only). +// +// Drop-in replacement for the (fp8_quantize_1x128 → get_mn_major_tma_aligned_packed_ue8m0_tensor) +// two-kernel sequence used by deep_gemm fp8 block-scale GEMMs. Returns the +// packed UE8M0 scale tensor (int32, MN-major, TMA-aligned) directly so +// deep_gemm's transform_sf_into_required_layout falls through the +// `(INT, 1, gran_k)` branch and skips its own pack call. +std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor const& self) +{ + CHECK_TH_CUDA(self); + CHECK_CONTIGUOUS(self); + + TORCH_CHECK(self.scalar_type() == at::ScalarType::BFloat16, "Input matrix dtype must be BF16."); + TORCH_CHECK(self.dim() == 2, "input must be a matrix"); + TORCH_CHECK(tensorrt_llm::common::isSM100Family(), + "fp8_quantize_1x128_packed_ue8m0 currently only supports SM100 (Blackwell)."); + + auto const m = self.sizes()[0]; + auto const n = self.sizes()[1]; + + TORCH_CHECK(m <= std::numeric_limits::max(), "M must be within int32"); + TORCH_CHECK(n <= std::numeric_limits::max(), "N must be within int32"); + TORCH_CHECK(n % 16 == 0, "self.sizes()[1] must be a multiple of 16, but got ", n); + + // FP8 output is row-major [m, n] with the same alignment used by the legacy path. + // The legacy path pads M to a multiple of 4; replicate that to avoid layout surprises. + auto const m_padded = (m + 4 - 1) / 4 * 4; + + at::Tensor valueE4M3 = at::detail::empty_cuda( + {m_padded, n}, at::ScalarType::Float8_e4m3fn, self.device(), /* stride */ std::nullopt); + + // Packed scale physical layout: [num_packed_sf_k, m_aligned] uint32, MN-contiguous in memory. + // deep_gemm's get_mn_major_tma_aligned_packed_ue8m0_tensor returns a strided VIEW with + // PyTorch shape `[mn, packed_sf_k]` and strides `(1, tma_aligned_mn)`. We build the same + // strided view so deep_gemm's transform_sf_into_required_layout falls into the + // `(INT, 1, gran_k)` branch and skips its own pack call. + auto const num_n_blocks = (n + 127) / 128; + auto const num_packed_sf_k = (num_n_blocks + 3) / 4; + constexpr int kTmaAlignedUint32Elems = 4; // 16 bytes / sizeof(uint32_t) + auto const m_aligned = (m_padded + kTmaAlignedUint32Elems - 1) / kTmaAlignedUint32Elems * kTmaAlignedUint32Elems; + + // Allocate physical buffer [num_packed_sf_k, m_aligned] (K-major in memory). + at::Tensor packedBuf = at::detail::empty_cuda( + {num_packed_sf_k, m_aligned}, at::ScalarType::Int, self.device(), /* stride */ std::nullopt); + // Zero-fill so the [m, m_aligned) tail and out-of-range scale slots are deterministic. + packedBuf.zero_(); + + auto stream = at::cuda::getCurrentCUDAStream(self.get_device()); + + tensorrt_llm::kernels::fp8_blockscale_gemm::launch_fp8_quantize_1x128_packed_bf16_e4m3( + reinterpret_cast<__nv_fp8_e4m3*>(valueE4M3.data_ptr()), reinterpret_cast(packedBuf.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), static_cast(m), static_cast(n), + static_cast(m_aligned), stream); + + // Wrap the [num_packed_sf_k, m_aligned] memory as a [m, num_packed_sf_k] strided tensor + // matching deep_gemm's get_mn_major_tma_aligned_packed_ue8m0_tensor return contract: + // shape = (m, num_packed_sf_k) + // stride = (1, m_aligned) + at::Tensor packedScale = at::from_blob( + packedBuf.data_ptr(), + /* sizes */ {m, num_packed_sf_k}, + /* strides */ {1, m_aligned}, + /* deleter */ [keep = packedBuf](void*) mutable {}, packedBuf.options()); + + return {valueE4M3.slice(0, 0, m), packedScale}; +} } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -149,10 +217,12 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("fp8_quantize_1x128(Tensor input, bool use_ue8m0=False) -> (Tensor, Tensor)"); m.def("fp8_batched_quantize_1x128_permute102(Tensor input) -> (Tensor, Tensor)"); + m.def("fp8_quantize_1x128_packed_ue8m0(Tensor input) -> (Tensor, Tensor)"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("fp8_quantize_1x128", &tensorrt_llm::torch_ext::fp8_quantize_1x128); m.impl("fp8_batched_quantize_1x128_permute102", &tensorrt_llm::torch_ext::fp8_batched_quantize_1x128_permute102); + m.impl("fp8_quantize_1x128_packed_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_packed_ue8m0); } diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index efbb06238225..aa8de08f826c 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -680,6 +680,17 @@ def _(pe: torch.Tensor, nope: torch.Tensor): scale = pe.new_empty((M, 1), dtype=torch.int32) return packed, scale + @torch.library.register_fake("trtllm::fp8_quantize_1x128_packed_ue8m0") + def _(input: torch.Tensor): + # Returns (fp8_e4m3 [m, k], packed_ue8m0_int32 [m, packed_sf_k]) + # matching deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor's return shape. + m, k = input.shape[0], input.shape[1] + num_n_blocks = (k + 127) // 128 + num_packed_sf_k = (num_n_blocks + 3) // 4 + return torch.empty_like(input, + dtype=torch.float8_e4m3fn), input.new_empty( + (m, num_packed_sf_k), dtype=torch.int32) + @torch.library.register_fake("trtllm::causal_conv1d_fwd") def _( x: torch.Tensor, diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f5338c433625..8745400ed7bb 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1687,14 +1687,30 @@ def _( # deep_gemm_gen_tuning_buckets is imported from ..utils +_USE_FUSED_FP8_QUANT_PACK = os.environ.get("TRTLLM_FUSED_FP8_QUANT_PACK", + "0") == "1" + def _fp8_quantize_1x128_ue8m0(input: torch.Tensor, tactic: int): - """Dispatch FP8 1x128 quantization to CUDA or Triton kernel.""" + """Dispatch FP8 1x128 quantization to CUDA or Triton kernel. + + When the CUDA path is selected on SM100 and ``TRTLLM_FUSED_FP8_QUANT_PACK=1`` + is set, the fused ``fp8_quantize_1x128_packed_ue8m0`` op is used and the + follow-on ``get_mn_major_tma_aligned_packed_ue8m0_tensor`` call is skipped: + the new op writes packed-UE8M0 (int32) scales directly in the layout + deep_gemm expects, so deep_gemm's internal layout transform falls into the + pre-packed branch and skips its own pack kernel as well. + """ TACTIC_TRITON = 1 if tactic == TACTIC_TRITON: a, a_sf = fp8_quantize.triton_fp8_quantize_1x128(input, use_ue8m0=True) - else: - a, a_sf = torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0=True) + a_sf = deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor( + a_sf.transpose(0, 1)) + return a, a_sf + if _USE_FUSED_FP8_QUANT_PACK and get_sm_version() >= 100: + a, a_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(input) + return a, a_sf + a, a_sf = torch.ops.trtllm.fp8_quantize_1x128(input, use_ue8m0=True) a_sf = deep_gemm.get_mn_major_tma_aligned_packed_ue8m0_tensor( a_sf.transpose(0, 1)) return a, a_sf From 47f46c546ad2a96110bca26eda88898ad690b18b Mon Sep 17 00:00:00 2001 From: lishicheng1996-nv Date: Thu, 30 Apr 2026 15:34:51 +0800 Subject: [PATCH 10/58] [None][feat] Add MLA dependency-aware overlap on DSv4 (compressor || q_b_proj+norm) (#13629) Signed-off-by: Shicheng Li Signed-off-by: Fanrong Li (cherry picked from commit 7a8c7c9bed18672e1eb22e004250c4d70cd6785c) Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/modules/attention.py | 52 ++++++++++++++++++++---- 1 file changed, 44 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index e71eb48c3beb..2c300cd1e3e5 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -2212,10 +2212,50 @@ def forward_impl_with_deepseek_v4(self, position_ids: Optional[torch.Tensor], qr = q latent_cache = torch.concat([compressed_kv, k_pe], dim=-1) - q = self.q_b_proj(q) - # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm reduces per-head. - q = self.q_b_layernorm(q.view(-1, self.qk_head_dim)).view_as(q) - # Indexer + # TRTLLM_MLA_EXTRA_OVERLAP=1 reorders the V4 attention prologue so the + # outer compressor (writes its own KV-cache slot, only reads + # hidden_states) executes on the auxiliary CUDA stream concurrently + # with q_b_proj + q_b_layernorm on the default stream. Both branches + # are entirely independent (no shared inputs and disjoint writes), + # so this is a pure dependency-aware reorder. Falls back to the + # serial schedule when the env-var is unset, when the aux stream is + # unavailable, or when multi-stream mode is off. + _v4_extra_overlap = (os.environ.get("TRTLLM_MLA_EXTRA_OVERLAP", "0") + == "1" and self.compressor is not None + and self.aux_stream is not None) + + if _v4_extra_overlap: + + def _q_branch(): + q_proj = self.q_b_proj(q) + # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm + # reduces per-head. + return self.q_b_layernorm(q_proj.view( + -1, self.qk_head_dim)).view_as(q_proj) + + def _compressor_branch(): + self.compressor(hidden_states, attn_metadata) + return None + + q, _ = maybe_execute_in_parallel( + _q_branch, + _compressor_branch, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + else: + q = self.q_b_proj(q) + # Per-head RMS: view as [N*n_heads, head_dim] so RMSNorm reduces per-head. + q = self.q_b_layernorm(q.view(-1, self.qk_head_dim)).view_as(q) + if self.compressor is not None: + self.compressor(hidden_states, attn_metadata) + + # Indexer is independent of both q_b_proj and the compressor's KV-cache + # write, so it runs after either schedule. Kept serial because it + # internally reuses self.aux_stream for its own multi-stream q-proj || + # weights-proj split; running it concurrently with q_b_proj would + # create a stream-aliasing hazard. topk_indices = None if self.indexer is not None: topk_indices = self.indexer( @@ -2225,10 +2265,6 @@ def forward_impl_with_deepseek_v4(self, position_ids: Optional[torch.Tensor], position_ids, ) - # Compressor - if self.compressor is not None: - self.compressor(hidden_states, attn_metadata) - assert q.shape[ 0] == num_tokens, f"Expect q.shape[0] to be {num_tokens}, but got {q.shape[0]}" From eb638123ee86494f19bdaa799b86720c712c5ad7 Mon Sep 17 00:00:00 2001 From: Mingyang Date: Thu, 30 Apr 2026 15:58:10 +0800 Subject: [PATCH 11/58] [None][fix] Change compressor linear dtype (#13605) Signed-off-by: Mingyang Hao Co-authored-by: Mingyang Hao Signed-off-by: Fanrong Li (cherry picked from commit 7f0224bc947c124f26d1f28e740388b337721499) Signed-off-by: Fanrong Li --- .../sparse/deepseek_v4/compressor.py | 44 +--- .../sparse/deepseek_v4/deepseek_v4.py | 17 +- .../deepseek_v4/test_compressor_module.py | 234 +++++++++++++++++- .../deepseek_v4/test_compressor_tf32.py | 35 +-- 4 files changed, 248 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py index d8c3d2620413..e382373a491d 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -1,5 +1,3 @@ -import os -from contextlib import contextmanager from enum import IntEnum from typing import TYPE_CHECKING, Optional, Tuple, Union @@ -11,38 +9,10 @@ from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.modules.rms_norm import RMSNorm from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding -from tensorrt_llm._torch.utils import maybe_compile if TYPE_CHECKING: from .deepseek_v4 import DeepseekV4TrtllmAttentionMetadata -# When set to "1", forces wkv_gate to use full FP32 computation (via nn.Linear) -# instead of the default TF32 path (via F.linear + allow_tf32). -_USE_FP32_COMPRESSOR = os.environ.get("DEEPSEEK_V4_COMPRESSOR_FP32", "0") == "1" - - -@maybe_compile(dynamic=True) -def _to_float(x: torch.Tensor) -> torch.Tensor: - """Cast to float32 for TF32 GEMM (following DSA pattern).""" - return x.float() - - -@contextmanager -def _tf32_matmul_enabled(): - """Temporarily enable TF32 for FP32 matmul in this scope. - - Forces PyTorch/cuBLASLt to use CUBLAS_COMPUTE_32F_FAST_TF32 which - guarantees TF32 tensor cores. Plain CUBLAS_COMPUTE_32F (used by - torch.ops.trtllm.cublas_mm) falls back to SIMT SGEMM based on - cuBLASLt heuristics for small M. - """ - prev = torch.backends.cuda.matmul.allow_tf32 - torch.backends.cuda.matmul.allow_tf32 = True - try: - yield - finally: - torch.backends.cuda.matmul.allow_tf32 = prev - class KVCacheDtype(IntEnum): """KV cache dtype/layout preset (values match C++ cache_scale_type parameter). @@ -137,7 +107,7 @@ def __init__( self.dim, self.state_dim * 2, bias=False, - dtype=torch.float32, + dtype=dtype, quant_config=None, skip_create_weights_in_init=skip_create_weights_in_init, use_custom_cublas_mm=True, @@ -213,15 +183,9 @@ def forward( num_comp_tokens = metadata.new_comp_kv_lens_cuda[self.compress_ratio][:bsz] max_ctx_comp_kv_lens = metadata.max_ctx_compressed_tokens[self.compress_ratio] - # Project input to KV and score. - # Default: TF32 via F.linear under allow_tf32 context (explicit - # CUBLAS_COMPUTE_32F_FAST_TF32 -> TF32 tensor cores on Ampere+). - # Fallback: strict FP32 via nn.Linear when DEEPSEEK_V4_COMPRESSOR_FP32=1. - if _USE_FP32_COMPRESSOR: - kv_score = self.wkv_gate(_to_float(x)) - else: - with _tf32_matmul_enabled(): - kv_score = F.linear(_to_float(x), self.wkv_gate.weight) + # Project input to KV and score in the checkpoint dtype, then store the + # result in fp32 for the compressor kernels. + kv_score = F.linear(x.to(self.wkv_gate.weight.dtype), self.wkv_gate.weight).float() # Allocate output buffer kv_comp = torch.empty(total_num_comp_tokens, self.head_dim, device=x.device, dtype=x.dtype) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index bf187391dcb0..d49242856fc1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -984,10 +984,6 @@ def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): ) return q - def _update_k_cache(self, k_fp8, k_scale, metadata): - """Overwrite the fused compressor's FP8 cache with reference-style FP4-QAT values.""" - super()._update_k_cache(k_fp8, k_scale, metadata) - def forward( self, qr: torch.Tensor, @@ -995,15 +991,12 @@ def forward( metadata: DeepseekV4TrtllmAttentionMetadata, position_ids: torch.Tensor, ): - if self.indexer_cache_dtype not in ( - KVCacheDtype.FP8_PERTENSOR, - KVCacheDtype.FP8_BLOCKWISE, - ): + if self.indexer_cache_dtype != KVCacheDtype.FP8_BLOCKWISE: raise NotImplementedError( - "DeepSeek-V4 indexer BMM currently consumes FP8 K. " - f"Indexer cache preset {self.indexer_k_cache_dtype!r} is supported by " - "the compressor/cache scatter path, but needs the matching BF16/FP4 " - "indexer BMM path before end-to-end indexer execution can use it." + "DeepSeek-V4 indexer currently consumes FP8 blockwise K. " + f"Indexer cache preset {self.indexer_k_cache_dtype!r} needs matching " + "cache layout, cache update, and BMM paths before end-to-end indexer " + "execution can use it." ) # compress k k_fp8, k_scale = self.compressor(hidden_states, metadata) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index fa0113a1465d..d95d397cbea8 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Tests comparing Compressor with RefCompressor.""" @@ -19,10 +19,15 @@ RotaryScalingType, ) from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import DeepseekV4CacheManager -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import Compressor +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor import ( + Compressor, + KVCacheDtype, +) from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( DeepseekV4AttentionType, + DeepseekV4Indexer, ) +from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer from tensorrt_llm._torch.modules.rotary_embedding import RopeParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -1577,6 +1582,231 @@ def test_mixed_batch(): assert_similar(out_ref, out_comp, "Mixed batch context+generation single forward") +class _FakeCompressorCacheManager: + def __init__(self, head_dim: int, tokens_per_block: int = 4): + self.tokens_per_block = tokens_per_block + self.compressed_block_sizes = {0: tokens_per_block} + self._buffer = torch.empty(1, tokens_per_block * head_dim, device=DEVICE, dtype=DTYPE) + + def get_buffers(self, layer_idx, attn_type): + return self._buffer + + +def _create_small_compressor(kv_cache_dtype: str, is_indexer: bool) -> Compressor: + head_dim = 128 + rope_dim = 64 + mla_params = MLAParams( + hidden_size=16, + qk_rope_head_dim=rope_dim, + qk_nope_head_dim=head_dim - rope_dim, + ) + rope_params = RopeParams( + dim=rope_dim, + theta=ROPE_THETA, + max_positions=16, + original_max_positions=16, + scale_type=RotaryScalingType.none, + ) + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + is_neox=False, + ) + return Compressor( + mla_params=mla_params, + layer_idx=0, + compress_ratio=4, + norm_eps=1e-6, + skip_create_weights_in_init=False, + pos_embd_params=pos_embd_params, + dtype=DTYPE, + kv_cache_dtype=kv_cache_dtype, + is_indexer=is_indexer, + rotate_activation=True, + ).to(DEVICE) + + +def _create_minimal_metadata(compressor: Compressor, total_compressed_tokens: int = 1): + ratio = compressor.compress_ratio + bsz = 1 + block_table = torch.zeros(bsz, 1, device=DEVICE, dtype=torch.int32) + block_tables = {(ratio, attn_type): block_table for attn_type in DeepseekV4AttentionType} + metadata = DummyAttentionMetadata( + num_contexts=1, + num_generations=0, + num_ctx_tokens=ratio, + num_tokens=ratio, + kv_cache_manager=_FakeCompressorCacheManager(compressor.head_dim), + block_tables=block_tables, + cu_seq_lens=torch.tensor([0, ratio], device=DEVICE, dtype=torch.int32), + cu_new_comp_kv={ + ratio: torch.tensor([0, total_compressed_tokens], device=DEVICE, dtype=torch.int32) + }, + compressed_position_ids={ + ratio: torch.zeros(total_compressed_tokens, device=DEVICE, dtype=torch.int32) + }, + compressed_kv_lens={ + ratio: torch.tensor([total_compressed_tokens], device=DEVICE, dtype=torch.int32) + }, + past_kv_lens={ratio: torch.zeros(bsz, device=DEVICE, dtype=torch.int32)}, + new_comp_kv_lens_cuda={ + ratio: torch.tensor([total_compressed_tokens], device=DEVICE, dtype=torch.int32) + }, + num_total_compressed_tokens={ratio: total_compressed_tokens}, + max_ctx_compressed_tokens={ratio: total_compressed_tokens}, + compressed_mask_cuda={ + ratio: torch.ones(total_compressed_tokens, device=DEVICE, dtype=torch.bool) + }, + ) + metadata.kv_lens_cuda_runtime = torch.tensor([ratio], device=DEVICE, dtype=torch.int32) + metadata.cached_token_lens_cuda = torch.zeros(bsz, device=DEVICE, dtype=torch.int32) + return metadata + + +def test_compressor_wkv_gate_uses_checkpoint_dtype(): + compressor = _create_small_compressor(kv_cache_dtype="default", is_indexer=False) + + assert compressor.wkv_gate.weight.dtype == DTYPE + + +def _run_compressor_with_fake_postprocess(monkeypatch, kv_cache_dtype: str, is_indexer: bool): + compressor = _create_small_compressor(kv_cache_dtype, is_indexer) + metadata = _create_minimal_metadata(compressor) + seen = {} + + def fake_prefill_reduction(*args): + kv_comp = args[6] + kv_comp.fill_(0.25) + + def fake_paged_kv_compress(*args): + raise AssertionError("generation compression path should not run") + + def fake_postprocess_scatter( + kv_comp, + kv_out, + rms_weight, + rms_eps, + rotary_cos_sin, + position_ids, + nope_head_dim, + rope_head_dim, + kv_cache, + num_comp_tokens, + cu_new_comp_kv, + start_pos, + block_table, + compressed_mask, + tokens_per_block, + cache_dtype, + rotate_activation, + quant_output, + scale_output, + ): + seen["kv_out"] = kv_out + seen["quant_output"] = quant_output + seen["scale_output"] = scale_output + seen["cache_dtype"] = cache_dtype + if kv_out is not None: + kv_out.fill_(0.5) + if quant_output is not None: + quant_output.fill_(0x38) + if scale_output.dtype.is_floating_point: + scale_output.fill_(2.0) + else: + scale_output.fill_(0x7F) + + monkeypatch.setattr( + torch.ops.trtllm, + "compressor_prefill_reduction", + fake_prefill_reduction, + raising=False, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "compressor_paged_kv_compress", + fake_paged_kv_compress, + raising=False, + ) + monkeypatch.setattr( + torch.ops.trtllm, + "compressor_postprocess_scatter", + fake_postprocess_scatter, + raising=False, + ) + + x = torch.randn(compressor.compress_ratio, 16, device=DEVICE, dtype=DTYPE) + with torch.no_grad(): + output = compressor(x, metadata) + return output, seen + + +def test_main_compressor_does_not_materialize_postprocess_output(monkeypatch): + """The fused kernel writes main compressor output directly to paged cache. + + Re-materializing kv_out here and doing a Python scatter afterwards defeats + the fused postprocess/scatter kernel and regresses end-to-end performance. + """ + output, seen = _run_compressor_with_fake_postprocess( + monkeypatch, kv_cache_dtype="default", is_indexer=False + ) + + assert seen["cache_dtype"] == int(KVCacheDtype.NONE) + assert seen["kv_out"] is None + assert seen["quant_output"] is None + assert seen["scale_output"] is None + assert output[0].shape == (1, 128) + assert output[1] is None + + +@pytest.mark.parametrize( + "kv_cache_dtype,cache_dtype,expected_dtype,expected_quant_shape,expected_scale_shape", + [ + ( + "fp8_blockwise", + KVCacheDtype.FP8_BLOCKWISE, + torch.float8_e4m3fn, + (1, 128), + (1, 1), + ), + ( + "mxfp4", + KVCacheDtype.MXFP4_BLOCKWISE, + torch.float4_e2m1fn_x2, + (1, 64), + (1, 4), + ), + ], +) +def test_indexer_returns_fused_quant_outputs( + monkeypatch, + kv_cache_dtype, + cache_dtype, + expected_dtype, + expected_quant_shape, + expected_scale_shape, +): + output, seen = _run_compressor_with_fake_postprocess( + monkeypatch, kv_cache_dtype=kv_cache_dtype, is_indexer=True + ) + + quant_output, scale_output = output + assert seen["cache_dtype"] == int(cache_dtype) + assert seen["kv_out"] is None + assert quant_output.dtype == expected_dtype + assert quant_output.shape == expected_quant_shape + assert scale_output.shape == expected_scale_shape + assert torch.equal(quant_output.view(torch.uint8), torch.full_like(seen["quant_output"], 0x38)) + if scale_output.dtype.is_floating_point: + assert torch.equal(scale_output, torch.full_like(scale_output, 2.0)) + else: + assert torch.equal(scale_output, torch.full_like(scale_output, 0x7F)) + + +def test_deepseek_v4_indexer_uses_base_cache_scatter(): + assert "_update_k_cache" not in DeepseekV4Indexer.__dict__ + assert DeepseekV4Indexer._update_k_cache is Indexer._update_k_cache + + # ============================================================================ # FP8 Blockwise Quantization Tests # ============================================================================ diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py index 56047d62da05..2cfd278d4fa0 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py @@ -1,17 +1,12 @@ #!/usr/bin/env python3 -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the BF16 GEMM and FP32 escape-hatch paths in the DeepSeek-V4 -Compressor wkv_gate. +"""Tests for the BF16 GEMM path in the DeepSeek-V4 Compressor wkv_gate. The wkv_gate weight now matches the upstream V4 checkpoint dtype (bf16). The -DEEPSEEK_V4_COMPRESSOR_FP32=1 escape hatch keeps the GEMM in fp32 for -accuracy debugging. These tests cover both paths against an fp32 reference. +GEMM runs in that dtype, then stores the result in fp32 for compressor kernels. """ -import os -from unittest import mock - import pytest import torch import torch.nn.functional as F @@ -47,8 +42,8 @@ def _bf16_path(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: return F.linear(x.to(wkv_gate.weight.dtype), wkv_gate.weight).float() -def _fp32_path(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: - """FP32 escape hatch: upcast both inputs to fp32 (DEEPSEEK_V4_COMPRESSOR_FP32=1).""" +def _fp32_reference(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: + """FP32 reference path.""" return F.linear(x.float(), wkv_gate.weight.float()) @@ -77,9 +72,10 @@ def test_bf16_path_close_to_fp32_reference(num_tokens): with torch.no_grad(): out_bf16 = _bf16_path(x, bf16_gate) - out_fp32 = _fp32_path(x, fp32_gate) + out_fp32 = _fp32_reference(x, fp32_gate) assert out_bf16.shape == out_fp32.shape + assert out_bf16.dtype == torch.float32 a, b = out_bf16.flatten(), out_fp32.flatten() cos_sim = F.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item() @@ -89,20 +85,3 @@ def test_bf16_path_close_to_fp32_reference(num_tokens): scale = max(a.abs().max().item(), b.abs().max().item(), 1e-3) rel_err = max_diff / scale assert rel_err <= 0.1, f"rel_err={rel_err:.6f}, max_diff={max_diff:.6f}" - - -def test_env_var_selects_fp32_path(): - """DEEPSEEK_V4_COMPRESSOR_FP32=1 env var causes module-level flag to be True.""" - import importlib - - import tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.compressor as mod - - with mock.patch.dict(os.environ, {"DEEPSEEK_V4_COMPRESSOR_FP32": "1"}): - importlib.reload(mod) - assert mod._USE_FP32_COMPRESSOR is True - - # Restore default - with mock.patch.dict(os.environ, {}, clear=False): - os.environ.pop("DEEPSEEK_V4_COMPRESSOR_FP32", None) - importlib.reload(mod) - assert mod._USE_FP32_COMPRESSOR is False From a8cd76c2689d41ad529d55d94d0006fd5afe8bab Mon Sep 17 00:00:00 2001 From: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Date: Thu, 30 Apr 2026 09:59:35 +0000 Subject: [PATCH 12/58] [None][fix] Remove duplicate V2 KV cache allocation revert Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit e43a9c849b41c1cfeb5a17ba8e825e6044e65393) Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 73ef1246bc17..bef5f5755fcb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2978,13 +2978,6 @@ def _executor_loop(self): can_queue, can_queue_this_rank = self._can_queue( scheduled_batch) - # Revert spurious KV cache capacity growth (see overlap - # loop for the full comment). - if not can_queue and can_queue_this_rank \ - and self._scheduler_manages_kv_suspend: - for req in scheduled_batch.generation_requests: - self.kv_cache_manager.revert_allocate_generation(req) - if not can_queue: self._revert_gen_alloc(scheduled_batch) @@ -3388,18 +3381,6 @@ def _executor_loop_overlap(self): # we need to delay the update of the previous batch's sample state, # and let the later iteration to update it. should_process_previous_batch = can_queue or not can_queue_this_rank - - # With attention DP, can_queue=False means another rank has - # an empty batch so no forward pass will run. The V2 - # scheduler already grew each generation request's KV cache - # capacity during scheduling; revert that growth so it does - # not accumulate across skipped iterations and overflow the - # host page-index buffer. - if not can_queue and can_queue_this_rank \ - and self._scheduler_manages_kv_suspend: - for req in scheduled_batch.generation_requests: - self.kv_cache_manager.revert_allocate_generation(req) - if can_queue: # The generation requests that do not have batch_idx From 8ff8e428ef8d18989f30644f08c39be95eeb7895 Mon Sep 17 00:00:00 2001 From: Emma Qiao Date: Thu, 30 Apr 2026 17:04:40 +0800 Subject: [PATCH 13/58] [None][infra] Update CI for DS v4 (#13604) Signed-off-by: EmmaQiaoCh Signed-off-by: Emma Qiao Signed-off-by: Fanrong Li (cherry picked from commit 9ceb421f33ebf03950e2fb42270fe160c1f3e2ec) Signed-off-by: Fanrong Li --- jenkins/Build.groovy | 14 ++++++------- jenkins/L0_Test.groovy | 10 +++++----- .../test_lists/test-db/l0_b200_ds.yml | 16 +++++++++++++++ .../test_lists/test-db/l0_dgx_b200_ds.yml | 18 +++++++++++++++++ .../test_lists/test-db/l0_dgx_b300_ds.yml | 20 +++++++++++++++++++ .../test-db/l0_gb200_multi_gpus_ds.yml | 18 +++++++++++++++++ 6 files changed, 84 insertions(+), 12 deletions(-) create mode 100644 tests/integration/test_lists/test-db/l0_b200_ds.yml create mode 100644 tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml create mode 100644 tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml create mode 100644 tests/integration/test_lists/test-db/l0_gb200_multi_gpus_ds.yml diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index 405280496af8..045c010594d2 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -57,28 +57,28 @@ def BUILD_CONFIGS = [ (CONFIG_LINUX_X86_64_VANILLA) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars ENABLE_MULTI_DEVICE=1 --extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl --extra-cmake-vars MOONCAKE_ROOT=/usr/local/Mooncake --micro_benchmarks", (TARNAME) : "TensorRT-LLM.tar.gz", - (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;103-real", ], (CONFIG_LINUX_X86_64_SINGLE_DEVICE) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars ENABLE_MULTI_DEVICE=0 --extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars ENABLE_UCX=0 --micro_benchmarks", (TARNAME) : "single-device-TensorRT-LLM.tar.gz", - (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;103-real", ], (CONFIG_LINUX_X86_64_LLVM) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars ENABLE_MULTI_DEVICE=1 --extra-cmake-vars WARNING_IS_ERROR=ON --micro_benchmarks -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_CUDA_HOST_COMPILER=clang -DCMAKE_LINKER_TYPE=LLD", (TARNAME) : "llvm-TensorRT-LLM.tar.gz", - (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "80-real;86-real;89-real;90-real;100-real;103-real", ], (CONFIG_LINUX_AARCH64): [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl --extra-cmake-vars MOONCAKE_ROOT=/usr/local/Mooncake", (TARNAME) : "TensorRT-LLM-GH200.tar.gz", - (WHEEL_ARCHS): "90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "90-real;100-real;103-real", (BUILD_JOBS_FOR_CONFIG): "8", // TODO: Remove after fix the build OOM issue on SBSA ], (CONFIG_LINUX_AARCH64_LLVM) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_CUDA_HOST_COMPILER=clang -DCMAKE_LINKER_TYPE=LLD", (TARNAME) : "llvm-TensorRT-LLM-GH200.tar.gz", - (WHEEL_ARCHS): "90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "90-real;100-real;103-real", (BUILD_JOBS_FOR_CONFIG): "8", // TODO: Remove after fix the build OOM issue on SBSA ], ] @@ -470,9 +470,9 @@ def buildWheelInContainer(pipeline, libraries=[], triple=X86_64_TRIPLE, clean=fa if (extra_args == "") { if (triple == AARCH64_TRIPLE) { - extra_args = "-a '90-real;100-real;103-real;120-real'" + extra_args = "-a '90-real;100-real;103-real'" } else { - extra_args = "-a '80-real;86-real;89-real;90-real;100-real;103-real;120-real'" + extra_args = "-a '80-real;86-real;89-real;90-real;100-real;103-real'" } } if (pre_cxx11abi) { diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index c41bb35e0ed6..b40994ffdb78 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -4243,6 +4243,9 @@ def launchTestJobs(pipeline, testFilter) "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-3": ["auto:dgx-b200-flex", "l0_b200_multi_gpus_perf_sanity", 3, 5, 8, 1, true], "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-4": ["auto:dgx-b200-flex", "l0_b200_multi_gpus_perf_sanity", 4, 5, 8, 1, true], "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-5": ["auto:dgx-b200-flex", "l0_b200_multi_gpus_perf_sanity", 5, 5, 8, 1, true], + "DGX_B200-PyTorch-DS-1": ["auto:dgx-b200-flex", "l0_b200_ds", 1, 1, 1, 1, true], + "DGX_B200-4_GPUs-PyTorch-DS-1": ["auto:dgx-b200-flex", "l0_dgx_b200_ds", 1, 1, 4, 1, true], + "DGX_B300-4_GPUs-PyTorch-DS-1": ["auto:dgx-b300-flex", "l0_dgx_b300_ds", 1, 1, 4, 1, true], ] // B200 PerfSanity post-merge disaggregated // 2 Nodes @@ -4283,11 +4286,7 @@ def launchTestJobs(pipeline, testFilter) parallelJobs += parallelSlurmJobs // SBSA machines from the Blossom machine pool - SBSATestConfigs = [ - "GH200-TensorRT-Post-Merge-1": ["gh200", "l0_gh200", 1, 1], - // DGX Spark is also named as GB10 Grace Blackwell Superchip. - "GB10-PyTorch-1": ["gb10x", "l0_gb10", 1, 1], - ] + SBSATestConfigs = [:] fullSet += SBSATestConfigs.keySet() SBSASlurmTestConfigs = [ @@ -4321,6 +4320,7 @@ def launchTestJobs(pipeline, testFilter) "GB300-4_GPUs-PyTorch-PerfSanity-Post-Merge-1": ["auto:gb300-x4", "l0_gb300_multi_gpus_perf_sanity", 1, 3, 4], "GB300-4_GPUs-PyTorch-PerfSanity-Post-Merge-2": ["auto:gb300-x4", "l0_gb300_multi_gpus_perf_sanity", 2, 3, 4], "GB300-4_GPUs-PyTorch-PerfSanity-Post-Merge-3": ["auto:gb300-x4", "l0_gb300_multi_gpus_perf_sanity", 3, 3, 4], + "GB200-4_GPUs-PyTorch-DS-1": ["auto:gb200-x4", "l0_gb200_multi_gpus_ds", 1, 1, 4], ] fullSet += SBSASlurmTestConfigs.keySet() diff --git a/tests/integration/test_lists/test-db/l0_b200_ds.yml b/tests/integration/test_lists/test-db/l0_b200_ds.yml new file mode 100644 index 000000000000..05f18dc73115 --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_b200_ds.yml @@ -0,0 +1,16 @@ +version: 0.0.1 +l0_b200_ds: +- condition: + ranges: + system_gpu_count: + gte: 1 + lte: 1 + wildcards: + gpu: + - '*b200*' + linux_distribution_name: ubuntu* + terms: + stage: pre_merge + backend: pytorch + tests: + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml b/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml new file mode 100644 index 000000000000..213d2099074a --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml @@ -0,0 +1,18 @@ +version: 0.0.1 +l0_dgx_b200_ds: +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + wildcards: + gpu: + - '*b200*' + linux_distribution_name: ubuntu* + cpu: x86_64 + terms: + stage: pre_merge + backend: pytorch + orchestrator: mpi + tests: + - unittest/_torch/modeling/test_modeling_deepseekv4.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml b/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml new file mode 100644 index 000000000000..fe4ce1e79798 --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml @@ -0,0 +1,20 @@ +version: 0.0.1 +l0_dgx_b300_ds: +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + wildcards: + gpu: + - '*gb110*' + - '*b300*' + linux_distribution_name: ubuntu* + cpu: x86_64 + terms: + stage: pre_merge + backend: pytorch + tests: + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_kv_cache_v2_nixl_python + - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_kv_cache_v2_nixl_python + - accuracy/test_disaggregated_serving.py::TestGemma3_1BInstruct::test_kv_cache_v2_nixl_python diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus_ds.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus_ds.yml new file mode 100644 index 000000000000..ee43a43a60fe --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus_ds.yml @@ -0,0 +1,18 @@ +version: 0.0.1 +l0_gb200_multi_gpus_ds: +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + wildcards: + gpu: + - '*gb200*' + linux_distribution_name: ubuntu* + cpu: aarch64 + terms: + stage: pre_merge + backend: pytorch + tests: + - unittest/_torch/modules/test_engram.py + - unittest/_torch/modules/test_mhc.py From 6151697d766340de050e38fbc970b6fce223fb11 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Thu, 30 Apr 2026 17:42:47 +0800 Subject: [PATCH 14/58] [TRTLLM-12374][fix] Tighten DSv4 cache constraint floor to match warmup shapes (#13657) Signed-off-by: Lance Liao Signed-off-by: Lanyu Liao Co-authored-by: Lance Liao Signed-off-by: Fanrong Li (cherry picked from commit 8e8f37d066274128fd519c20ae522ad7950e8bd3) Signed-off-by: Fanrong Li --- .../sparse/deepseek_v4/cache_manager.py | 34 ++++++++----------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index c10dd440af25..016a43c2ea6b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -448,7 +448,6 @@ def _add_layer( # Build constraints and typical_step for better pool ratio. max_batch_size = self.max_batch_size max_seq_len = self.max_seq_len - max_input_len = self._max_input_len max_num_tokens = self._max_num_tokens max_draft_len = self._max_draft_len @@ -457,32 +456,29 @@ def _add_layer( # the typical step. An all-generation typical_step over-provisions the # compressed-cache pool at the expense of the SWA pool, starving the # SWA pool and artificially capping the achievable batch size. + ctx_capacity = max_num_tokens if max_num_tokens is not None else max_seq_len typical_step = BatchDesc( kv_caches=[ - KVCacheDesc(capacity=max_seq_len, history_length=0), + KVCacheDesc(capacity=ctx_capacity, history_length=0), ] + [KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - max_draft_len - 1)] * (max_batch_size - 1), ) constraints = [] - # Constraint 1: one context request at max_seq_len, this case is used in warmup stage. - # max_draft_len is already included in max_seq_len, so no need to add it again. - constraints.append(BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=0)])) - - # Constraint 2: when user given max_input_len and max_num_tokens, we can build constraints for context requests. - if max_input_len is not None and max_num_tokens is not None: - # There are at most max_batch_size generation requests, and each generation request consumes - # (max_draft_len + 1) tokens, so the remaining tokens are for context requests. - max_context_tokens = max_num_tokens - max_batch_size * (max_draft_len + 1) - if max_context_tokens > 0: - max_num_context = min(max_context_tokens // max_input_len, max_batch_size) - constraints.append( - BatchDesc( - [KVCacheDesc(capacity=max_input_len + max_draft_len + 1, history_length=0)] - * max_num_context - ) - ) + # Constraint 1: cuda graph generation warmup — one decode request that has + # accumulated to the tail of max_seq_len. Using history_length=max_seq_len-1 + # (instead of 0) lets SWA / SSM pools collapse to their windowed working set, + # while full-cache pools still need max_seq_len/tokens_per_block blocks + # because they don't age. + constraints.append( + BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1)]) + ) + + # Constraint 2: general / chunked-prefill warmup — one fresh context request + # at max_num_tokens (the per-iteration token budget). + if max_num_tokens is not None: + constraints.append(BatchDesc([KVCacheDesc(capacity=max_num_tokens, history_length=0)])) return KVCacheManagerConfigPy( tokens_per_block=tokens_per_block, From e9ecf1cf742a56310292a1e5a1f2aa8d757ffe51 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:39:51 +0800 Subject: [PATCH 15/58] [None][test] Add DeepSeek-V4 CI coverage (#13653) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit fa1e55e4238a337bb1c1b1867670c66bbc206782) Signed-off-by: Fanrong Li --- .../defs/accuracy/references/gsm8k.yaml | 10 ++ .../defs/accuracy/test_llm_api_pytorch.py | 103 ++++++++++++++---- .../test_lists/test-db/l0_b200_ds.yml | 12 ++ .../test_lists/test-db/l0_dgx_b200_ds.yml | 4 + .../test_lists/test-db/l0_dgx_b300_ds.yml | 25 +++++ .../modeling/test_modeling_deepseekv4.py | 55 ++++------ 6 files changed, 153 insertions(+), 56 deletions(-) diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 489beae5acfd..11b7663ccd66 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -143,6 +143,16 @@ deepseek-ai/DeepSeek-V3.2-Exp: kv_cache_quant_algo: FP8 spec_dec_algo: MTP accuracy: 95.6 +deepseek-ai/DeepSeek-V4-Flash: + # GSM8K measurements: + # * 95.11 on 8x B200 178GB at TP=8 with tightened config (fraction=0.15, + # cuda_graph=None, max_num_tokens=2048) — original measurement. + # * 95.38 on 4x B300 275GB at TP=4 with default config (fraction=0.5, + # CudaGraphConfig()) — current CI path (test_nvfp4_4gpus_static_eplb + # in l0_dgx_b300_ds.yml). Drift +0.27 from 95.11 is within sigma; the + # 95.11 reference still holds for the hypothesis test. + - quant_algo: FP8_BLOCK_SCALES + accuracy: 95.11 Qwen3/Qwen3-4B: - spec_dec_algo: Eagle accuracy: 85.823 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0bff4d63eca8..72d767470761 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3638,10 +3638,12 @@ def _run_deepseekv4_eplb(model_name, model_path, moe_backend, eplb_config, - mtp_nextn=0): - # V4 (~284B) plus EPLB redundant slots leaves little headroom; keep the - # KV-cache fraction conservative and clamp max_seq_len so the default - # 1M context doesn't pre-allocate KV blocks beyond what 180GB B200s have. + mtp_nextn=0, + tensor_parallel_size=4): + # Default config targets 4x B300 (~288 GB/GPU): plenty of headroom at + # TP=4 for V4-Flash NVFP4 (~36 GB/rank weights) or V4-Flash-Base FP8 + # (~71 GB/rank weights), even with EPLB redundancy + DeepGemm MoE + # workspace + cuda graph capture. kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) pytorch_config = dict(cuda_graph_config=CudaGraphConfig(), moe_config=MoeConfig(backend=moe_backend, @@ -3649,8 +3651,8 @@ def _run_deepseekv4_eplb(model_name, mtp_config = (MTPDecodingConfig( num_nextn_predict_layers=mtp_nextn) if mtp_nextn > 0 else None) with LLM(model_path, - tensor_parallel_size=8, - moe_expert_parallel_size=8, + tensor_parallel_size=tensor_parallel_size, + moe_expert_parallel_size=tensor_parallel_size, kv_cache_config=kv_cache_config, enable_attention_dp=True, max_seq_len=4096, @@ -3667,20 +3669,61 @@ class TestDeepSeekV4Flash(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash" - @pytest.mark.skip_less_mpi_world_size(8) - @parametrize_with_ids("moe_backend", ["WIDEEP", "TRTLLM"]) - def test_nvfp4_8gpus_static_eplb(self, moe_backend): + @pytest.mark.skip_less_mpi_world_size(4) + def test_auto_dtype(self): + # Aggregate (non-disagg, non-EPLB) smoke test. NVFP4 weights are ~71 + # GB/rank at TP=2, ~36 GB/rank at TP=4 — TP=4 fits comfortably on + # 4x B200 178GB. TRTLLM backend required because V4-Flash MXFP4 + # routed experts are unsupported by WIDEEP (raises "Unsupported + # quantization mode: [65536]"). is_integration_test=True keeps this + # to a 1-sample smoke (no GSM8K reference required). + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) + with LLM(self.MODEL_PATH, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + moe_config=MoeConfig(backend="TRTLLM"), + enable_attention_dp=True, + max_seq_len=4096, + kv_cache_config=kv_cache_config) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", [ + pytest.param( + "WIDEEP", + marks=pytest.mark.skip( + reason= + "V4-Flash MXFP4 routed experts: WIDEEP _get_quant_method has " + "no MXFP4 branch (raises 'Unsupported quantization mode: " + "[65536]'). Re-enable once fused_moe_wide_ep.py supports MXFP4." + )), + "TRTLLM", + ]) + def test_nvfp4_4gpus_static_eplb(self, moe_backend): eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, - layer_updates_per_iter=0) + layer_updates_per_iter=0, + ep_size=4) _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, eplb_config) - @pytest.mark.skip_less_mpi_world_size(8) - @parametrize_with_ids("moe_backend", ["WIDEEP", "TRTLLM"]) + @pytest.mark.skip_less_mpi_world_size(4) + @parametrize_with_ids("moe_backend", [ + pytest.param( + "WIDEEP", + marks=pytest.mark.skip( + reason= + "V4-Flash MXFP4 routed experts: WIDEEP _get_quant_method has " + "no MXFP4 branch (raises 'Unsupported quantization mode: " + "[65536]'). Re-enable once fused_moe_wide_ep.py supports MXFP4." + )), + "TRTLLM", + ]) @parametrize_with_ids("mtp_nextn", [0, 1]) - def test_nvfp4_8gpus_online_eplb(self, moe_backend, mtp_nextn): + def test_nvfp4_4gpus_online_eplb(self, moe_backend, mtp_nextn): eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, - layer_updates_per_iter=2) + layer_updates_per_iter=2, + ep_size=4) _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, @@ -3695,22 +3738,40 @@ class TestDeepSeekV4FlashBase(LlmapiAccuracyTestHarness): MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash-Base" MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash-Base" + @pytest.mark.skip_less_mpi_world_size(4) + def test_auto_dtype(self): + # Aggregate (non-disagg, non-EPLB) smoke test. FP8 weights ~71 GB/rank + # at TP=4 — fits on 4x B300 (~288 GB/GPU). WIDEEP required (CUTLASS + # path is Hopper-only on Blackwell). 1-sample smoke. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) + with LLM(self.MODEL_PATH, + tensor_parallel_size=4, + moe_expert_parallel_size=4, + moe_config=MoeConfig(backend="WIDEEP"), + enable_attention_dp=True, + max_seq_len=4096, + kv_cache_config=kv_cache_config) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, is_integration_test=True) + # CUTLASS is omitted: V4 Flash-Base FP8 block-scale weights take a # Hopper-only kernel path (CutlassFp8BlockScaleGemmRunner::moeGemm) that - # fails on Blackwell. WIDEEP avoids that path and works on B200. - @pytest.mark.skip_less_mpi_world_size(8) + # fails on Blackwell. WIDEEP avoids that path and works on B200/B300. + @pytest.mark.skip_less_mpi_world_size(4) @parametrize_with_ids("moe_backend", ["WIDEEP"]) - def test_fp8_8gpus_static_eplb(self, moe_backend): + def test_fp8_4gpus_static_eplb(self, moe_backend): eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, - layer_updates_per_iter=0) + layer_updates_per_iter=0, + ep_size=4) _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, eplb_config) - @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.skip_less_mpi_world_size(4) @parametrize_with_ids("moe_backend", ["WIDEEP"]) - def test_fp8_8gpus_online_eplb(self, moe_backend): + def test_fp8_4gpus_online_eplb(self, moe_backend): eplb_config = _make_deepseekv4_eplb_config(self.MODEL_PATH, - layer_updates_per_iter=2) + layer_updates_per_iter=2, + ep_size=4) _run_deepseekv4_eplb(self.MODEL_NAME, self.MODEL_PATH, moe_backend, eplb_config) diff --git a/tests/integration/test_lists/test-db/l0_b200_ds.yml b/tests/integration/test_lists/test-db/l0_b200_ds.yml index 05f18dc73115..4f14d87cc1c6 100644 --- a/tests/integration/test_lists/test-db/l0_b200_ds.yml +++ b/tests/integration/test_lists/test-db/l0_b200_ds.yml @@ -14,3 +14,15 @@ l0_b200_ds: backend: pytorch tests: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp + # ------------- DeepSeek-V4 unit tests (single B200) --------------- + # Sparse-attention kernels + cache manager + compressor + indices_transform + # + o_proj. Only path to V4 sparse-attention kernel coverage in CI. + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py TIMEOUT (30) + - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py TIMEOUT (90) + - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py TIMEOUT (90) + - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py TIMEOUT (15) + # Tokenizer chat-template tests. + - unittest/llmapi/test_deepseek_v4_tokenizer.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml b/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml index 213d2099074a..1ba22ffaf43f 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200_ds.yml @@ -16,3 +16,7 @@ l0_dgx_b200_ds: orchestrator: mpi tests: - unittest/_torch/modeling/test_modeling_deepseekv4.py + # V4-Flash NVFP4 aggregate smoke (TP=4, no EPLB): integration-test mode + # (1 MMLU sample, no GSM8K reference required). Validates the basic + # forward path without disagg or EPLB complexity. + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60) diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml b/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml index fe4ce1e79798..58a05389bc1a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml @@ -14,7 +14,32 @@ l0_dgx_b300_ds: terms: stage: pre_merge backend: pytorch + orchestrator: mpi tests: - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_kv_cache_v2_nixl_python - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_kv_cache_v2_nixl_python - accuracy/test_disaggregated_serving.py::TestGemma3_1BInstruct::test_kv_cache_v2_nixl_python + # ------------- DeepSeek-V4 (4x B300, pre_merge) --------------- + # B300's larger per-GPU memory (~288 GB) lets V4 EPLB and FP8 aggregate + # smoke fit at TP=4. (Same tests at TP=8 would also fit on 8x B200, but a + # dedicated 8-GPU B200 stage is not currently provisioned.) + # V4-Flash-Base FP8 aggregate smoke (TP=4, no EPLB): 1 MMLU sample. + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_auto_dtype TIMEOUT (60) + # V4-Flash NVFP4 static EPLB sanity. WIDEEP backend is skipped at the + # test level because V4-Flash MXFP4 routed experts are unsupported by + # WIDEEP (raises "Unsupported quantization mode: [65536]"); TRTLLM is the + # supported backend. Measured GSM8K accuracy 95.11 (see + # references/gsm8k.yaml; reference was captured at TP=8 / 8x B200, may + # drift slightly at TP=4 / 4x B300 — re-measure if hypothesis test flaps). + - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_static_eplb[moe_backend=TRTLLM] TIMEOUT (120) + # Known issues (uncomment when unblocked): + # * V4-Flash-Base FP8 + EPLB hits CUBLAS_STATUS_EXECUTION_FAILED in + # forward pass with WIDEEP backend (memory pressure aside). + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4gpus_static_eplb[moe_backend=WIDEEP] + # * Online EPLB (layer_updates_per_iter>0) requires gdrcopy/gdrdrv for + # HostAccessibleDeviceAllocator (moeLoadBalancer.cpp:846); without + # gdrcopy the executor crashes during initialization. Static EPLB does + # not hit this path. + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_online_eplb[mtp_nextn=0-moe_backend=TRTLLM] + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_nvfp4_4gpus_online_eplb[mtp_nextn=1-moe_backend=TRTLLM] + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_fp8_4gpus_online_eplb[moe_backend=WIDEEP] diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index 07e2b75a5a84..7a0b95f95da9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -27,7 +27,6 @@ DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, DeepseekV4Gate, - DeepseekV4MoE, DeepseekV4MTP, _deepseek_v4_pos_embd_params, _remap_deepseek_v4_checkpoint_keys, @@ -336,19 +335,6 @@ def test_deepseek_v4_compressor_rotate_and_indexer_rope_contracts(): assert "rotate_activation=False" in attention_init -def test_deepseek_v4_moe_swiglu_limit_applies_to_routed_and_shared_experts(): - moe_init = inspect.getsource(DeepseekV4MoE.__init__) - assert "moe_swiglu_limit = None" in moe_init - assert "supports_swiglu_limit = False" in moe_init - assert "mode.has_w4a8_mxfp4_mxfp8()" in moe_init - assert "swiglu_limit=moe_swiglu_limit" in moe_init - - shared_expert_block = moe_init.split("self.shared_experts = GatedMLP", 1)[1].split( - "self.allreduce", 1 - )[0] - assert "swiglu_limit=swiglu_limit" in shared_expert_block - - def test_deepseek_v4_attention_forward_injects_attn_sink(monkeypatch): captured = {} @@ -534,17 +520,17 @@ def test_deepseek_v4_rope_params_follow_layer_compress_ratio(): assert active_config_rope.rope.scale_type == RotaryScalingType.none assert active_config_rope.rope.theta == 10000.0 -def test_deepseek_v4_sparse_ratios_prefer_checkpoint_defaults( - tmp_path, monkeypatch): + +def test_deepseek_v4_sparse_ratios_prefer_checkpoint_defaults(tmp_path, monkeypatch): checkpoint_ratios = [128, 128, 4, 128, 4, 128, 0, 4] config = DeepseekV4Config( architectures=["DeepseekV4ForCausalLM"], num_hidden_layers=len(checkpoint_ratios), compress_ratios=checkpoint_ratios, - sliding_window=256, ) - monkeypatch.setattr("tensorrt_llm._torch.model_config.load_pretrained_config", - lambda *args, **kwargs: config) + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", lambda *args, **kwargs: config + ) sparse_attn_config = DeepSeekV4SparseAttentionConfig( compress_ratios=[1, 1, 4, 128, 4, 128, 4], q_split_threshold=2048, @@ -558,22 +544,23 @@ def test_deepseek_v4_sparse_ratios_prefer_checkpoint_defaults( moe_backend="TRTLLM", ) - assert model_config.sparse_attention_config.compress_ratios == [ - 128, 128, 4, 128, 4, 128, 1, 4 - ] - assert model_config.sparse_attention_config.window_size == 256 + assert model_config.sparse_attention_config.compress_ratios == [128, 128, 4, 128, 4, 128, 1, 4] + # V4 sparse MLA hardcodes window_size==128 (FMHA kernel TileSizeKV; see + # the runtime assertion in deepseek_v4.py:DeepseekV4TrtllmAttentionMetadata + # __post_init__), so this is the only legal value here. + assert model_config.sparse_attention_config.window_size == 128 -def test_deepseek_v4_sparse_ratios_keep_checkpoint_length_without_mtp( - tmp_path, monkeypatch): +def test_deepseek_v4_sparse_ratios_keep_checkpoint_length_without_mtp(tmp_path, monkeypatch): checkpoint_ratios = [128, 128] + [4, 128] * 29 + [0, 4] config = DeepseekV4Config( architectures=["DeepseekV4ForCausalLM"], num_hidden_layers=61, compress_ratios=checkpoint_ratios, ) - monkeypatch.setattr("tensorrt_llm._torch.model_config.load_pretrained_config", - lambda *args, **kwargs: config) + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", lambda *args, **kwargs: config + ) sparse_attn_config = DeepSeekV4SparseAttentionConfig( compress_ratios=[1, 1, 4, 128, 4, 128, 4], ) @@ -585,15 +572,12 @@ def test_deepseek_v4_sparse_ratios_keep_checkpoint_length_without_mtp( moe_backend="TRTLLM", ) - assert len(model_config.sparse_attention_config.compress_ratios) == len( - checkpoint_ratios) - assert model_config.sparse_attention_config.compress_ratios[:-2] == ( - checkpoint_ratios[:-2]) + assert len(model_config.sparse_attention_config.compress_ratios) == len(checkpoint_ratios) + assert model_config.sparse_attention_config.compress_ratios[:-2] == (checkpoint_ratios[:-2]) assert model_config.sparse_attention_config.compress_ratios[-2:] == [1, 4] -def test_deepseek_v4_sparse_ratios_keep_explicit_override( - tmp_path, monkeypatch): +def test_deepseek_v4_sparse_ratios_keep_explicit_override(tmp_path, monkeypatch): checkpoint_ratios = [128, 128, 4, 128, 4, 128, 0, 4] explicit_ratios = [1, 4, 1, 4, 1, 4, 1, 4] config = DeepseekV4Config( @@ -601,8 +585,9 @@ def test_deepseek_v4_sparse_ratios_keep_explicit_override( num_hidden_layers=len(checkpoint_ratios), compress_ratios=checkpoint_ratios, ) - monkeypatch.setattr("tensorrt_llm._torch.model_config.load_pretrained_config", - lambda *args, **kwargs: config) + monkeypatch.setattr( + "tensorrt_llm._torch.model_config.load_pretrained_config", lambda *args, **kwargs: config + ) sparse_attn_config = DeepSeekV4SparseAttentionConfig( compress_ratios=explicit_ratios, ) From c5b467e6093470f3fcdf630bd2faa80f04d8898e Mon Sep 17 00:00:00 2001 From: Mingyang Date: Thu, 30 Apr 2026 18:40:17 +0800 Subject: [PATCH 16/58] [TRTLLM-12383][fix] limit MHC TF32 pmap GEMM to SM100 (#13660) Signed-off-by: Mingyang Hao Co-authored-by: Mingyang Hao Signed-off-by: Fanrong Li (cherry picked from commit d2ea57d32568941468d5a268fe530ca3a5ee28f8) Signed-off-by: Fanrong Li --- cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 ad9b10fe2e55..a34350deddfe 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -113,7 +113,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0) float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum) { // [M] (caller memsets to 0) -#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; constexpr uint32_t SHAPE_K = HC_MULT * HIDDEN; @@ -637,7 +637,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) float* __restrict__ comb_mix_out, // [M, HC_MULT, HC_MULT] float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, uint32_t sinkhorn_repeat) { -#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) and (__CUDA_ARCH__ < 1100)) or defined(__CLION_IDE__) using Barrier = cutlass::arch::ClusterTransactionBarrier; constexpr uint32_t HC_MULT2 = HC_MULT * HC_MULT; From 22940240c788d125aa5416d8bde1de581b6ac575 Mon Sep 17 00:00:00 2001 From: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Date: Tue, 5 May 2026 05:50:43 +0000 Subject: [PATCH 17/58] [None][chore] Apply pre-commit fixes Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit 3f1313b3c03af4260b242686a916a8d25efa4ca7) Signed-off-by: Fanrong Li --- .pre-commit-config.yaml | 4 +- .../kernels/customMoeRoutingKernels.cu | 3 +- .../kernels/mhcKernels/CMakeLists.txt | 24 +- .../mhcKernels/fused_tf32_pmap_gemm.cuh | 3 +- .../kernels/mhcKernels/mhcFusedHcKernel.cu | 10 +- .../kernels/mhcKernels/mhc_fused_fma.cuh | 3 +- cpp/tensorrt_llm/thop/compressorOp.cpp | 31 +- tensorrt_llm/_torch/model_config.py | 11 +- tensorrt_llm/_torch/models/__init__.py | 2 +- tensorrt_llm/_torch/modules/attention.py | 34 +- .../_torch/modules/fused_moe/__init__.py | 6 +- tensorrt_llm/_torch/modules/gated_mlp.py | 3 +- .../_torch/modules/mhc/hyper_connection.py | 5 +- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 5 +- tensorrt_llm/llmapi/tokenizer.py | 2 +- .../test_deepseek_v4_indices_transform.py | 4 +- .../test_deepseek_v4_sparse_mla.py | 6 +- .../attention/sparse/dsa/test_dsa_indexer.py | 1424 ++++++++--------- .../attention/sparse/kernel/test_flash_mla.py | 61 +- .../sparse/rocketkv/test_rocketkv.py | 390 +++-- .../sparse/test_sparse_mla_forward.py | 132 +- .../_torch/thop/serial/test_moe_gate.py | 8 +- 22 files changed, 1066 insertions(+), 1105 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5b0072c49fad..193bf64671e8 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2726,8 +2726,8 @@ static-analysis-files: &static_analysis_files | tests/unittest/_torch/sampler/test_beam_search_util.py | )$ -# Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx) -exclude: '(^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)' +# Global exclude: vendored code + trtllm-gen generated artifacts (cubin pointers, export headers, cuda_ptx) +exclude: '(^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)' default_install_hook_types: [pre-commit, commit-msg] repos: diff --git a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu index 6c2c32cf4be4..f76db261e5a3 100644 --- a/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu +++ b/cpp/tensorrt_llm/kernels/customMoeRoutingKernels.cu @@ -475,8 +475,7 @@ void gate_forward(void* scores_in, // [batch_size, nExperts] - pre-computed from scores, bias_ptr, nullptr, nullptr, weights, indices, batch_size, route_scale, stream); } break; - default: - TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384"); + default: TLLM_CHECK_WITH_INFO(false, "gate_forward only supports n_experts 256 or 384"); } sync_check_cuda_error(stream); } diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt index ae7ed2d085c8..58dc760fcf9d 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt @@ -15,23 +15,21 @@ # the License. # -# Library sources are listed explicitly. Profiling/benchmark/test harnesses -# in this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are +# Library sources are listed explicitly. Profiling/benchmark/test harnesses in +# this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are # standalone binaries with their own main() and must NOT be compiled into the # trtllm shared library. -set(SRC_CU - mhcKernels.cu - mhcFusedHcKernel.cu -) +set(SRC_CU mhcKernels.cu mhcFusedHcKernel.cu) add_library(mhcKernels_src OBJECT ${SRC_CU}) set_property(TARGET mhcKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET mhcKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) -target_compile_options(mhcKernels_src PRIVATE $<$:--use_fast_math>) +target_compile_options(mhcKernels_src + PRIVATE $<$:--use_fast_math>) -# Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh, reduction.cuh) -# Used by the tcgen05.mma-based fused post-mapping + GEMM kernel on SM100. -target_include_directories(mhcKernels_src - PRIVATE - ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include -) +# Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh, +# reduction.cuh) Used by the tcgen05.mma-based fused post-mapping + GEMM kernel +# on SM100. +target_include_directories( + mhcKernels_src + PRIVATE ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include) 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 a34350deddfe..37b45ddf3457 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -1229,8 +1229,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) + hc_base[2 * HC_MULT + lane_bf * HC_MULT + k]; constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; - float const rowMax - = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); + float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); #pragma unroll for (uint32_t k = 0; k < HC_MULT; ++k) cm_vals[k] = __expf(cm_vals[k] - rowMax); diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index 7954e17c127e..334961bdfced 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -51,9 +51,8 @@ namespace kernels::mhc namespace { -__global__ void fhcZeroWorkspacesKernel( - float* __restrict__ y_acc, uint32_t y_elems, float* __restrict__ r_acc, uint32_t r_elems, - int* __restrict__ done_counter, uint32_t done_elems) +__global__ void fhcZeroWorkspacesKernel(float* __restrict__ y_acc, uint32_t y_elems, float* __restrict__ r_acc, + uint32_t r_elems, int* __restrict__ done_counter, uint32_t done_elems) { uint32_t const tid = blockIdx.x * blockDim.x + threadIdx.x; uint32_t const stride = gridDim.x * blockDim.x; @@ -74,8 +73,8 @@ __global__ void fhcZeroWorkspacesKernel( } } -inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint32_t r_elems, - int* done_counter, uint32_t done_elems, cudaStream_t stream) +inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint32_t r_elems, int* done_counter, + uint32_t done_elems, cudaStream_t stream) { uint32_t const total = y_elems + r_elems + (done_counter != nullptr ? done_elems : 0u); if (total == 0u) @@ -91,7 +90,6 @@ inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint } // namespace - // ---- mHC fused kernel shape constants (mirrors the Python module) ---- static constexpr uint32_t FHC_SHAPE_N = 24; // HC_MULT * (2 + HC_MULT) = 4 * 6 = 24 static constexpr uint32_t FHC_HIDDEN = 4096; // only this hidden size is currently wired up diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh index e3b01fab6b7c..23a3273637ca 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhc_fused_fma.cuh @@ -814,8 +814,7 @@ __launch_bounds__(256) __global__ void fused_pmap_gemm_fma_allinone(__nv_bfloat1 = y_local[2 * HC_MULT + lane * HC_MULT + k] * rstd * s2 + hc_base[2 * HC_MULT + lane * HC_MULT + k]; constexpr unsigned LANE_MASK = (1u << HC_MULT) - 1; - float const rowMax - = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); + float const rowMax = fmaxf(fmaxf(cm_vals[0], cm_vals[1]), fmaxf(cm_vals[2], cm_vals[3])); #pragma unroll for (int k = 0; k < HC_MULT; k++) cm_vals[k] = expf(cm_vals[k] - rowMax); diff --git a/cpp/tensorrt_llm/thop/compressorOp.cpp b/cpp/tensorrt_llm/thop/compressorOp.cpp index 3de3a9821e95..b73fbd4469f8 100644 --- a/cpp/tensorrt_llm/thop/compressorOp.cpp +++ b/cpp/tensorrt_llm/thop/compressorOp.cpp @@ -46,9 +46,9 @@ void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf1 tk::pagedKvCompressLaunch(kv_score.data_ptr(), ape.data_ptr(), paged_kv.data_ptr(), paged_score.data_ptr(), block_table_kv.data_ptr(), block_table_score.data_ptr(), output.data_ptr(), kv_lens.data_ptr(), start_pos.data_ptr(), cu_seq_lens.data_ptr(), - cu_kv_comp.data_ptr(), - static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), - static_cast(head_dim), static_cast(compress_ratio), static_cast(next_n), io_eb, out_eb, stream); + cu_kv_comp.data_ptr(), static_cast(batch_size), static_cast(page_size), + static_cast(block_table_kv.size(1)), static_cast(head_dim), static_cast(compress_ratio), + static_cast(next_n), io_eb, out_eb, stream); } // Prefill kernel: bulk compression with state update @@ -64,10 +64,9 @@ void compressorPrefillReductionOp(torch::Tensor kv_score, torch::Tensor ape, tor tk::prefillReductionLaunch(kv_score.data_ptr(), ape.data_ptr(), paged_kv.data_ptr(), paged_score.data_ptr(), block_table_kv.data_ptr(), block_table_score.data_ptr(), output.data_ptr(), kv_lens.data_ptr(), start_pos.data_ptr(), cu_seq_lens.data_ptr(), - cu_kv_comp.data_ptr(), - static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), - static_cast(head_dim), static_cast(compress_ratio), static_cast(max_outputs), io_eb, out_eb, - stream); + cu_kv_comp.data_ptr(), static_cast(batch_size), static_cast(page_size), + static_cast(block_table_kv.size(1)), static_cast(head_dim), static_cast(compress_ratio), + static_cast(max_outputs), io_eb, out_eb, stream); } // Fused postprocess + scatter: RMSNorm + RoPE + Hadamard + paged scatter in one kernel @@ -89,21 +88,19 @@ void compressorPostProcessScatterOp(torch::Tensor kv_comp, // [total_tokens, hea { auto stream = at::cuda::getCurrentCUDAStream(); - TORCH_CHECK(cos_sin_table.scalar_type() == at::kFloat, - "cos_sin_table must be float32, got ", cos_sin_table.scalar_type()); + TORCH_CHECK( + cos_sin_table.scalar_type() == at::kFloat, "cos_sin_table must be float32, got ", cos_sin_table.scalar_type()); TORCH_CHECK(cos_sin_table.is_contiguous(), "cos_sin_table must be contiguous"); - TORCH_CHECK(position_ids.scalar_type() == at::kInt, - "position_ids must be int32, got ", position_ids.scalar_type()); + TORCH_CHECK(position_ids.scalar_type() == at::kInt, "position_ids must be int32, got ", position_ids.scalar_type()); TORCH_CHECK(position_ids.is_contiguous(), "position_ids must be contiguous"); - TORCH_CHECK(compressed_mask.scalar_type() == at::kBool, - "compressed_mask must be bool, got ", compressed_mask.scalar_type()); + TORCH_CHECK(compressed_mask.scalar_type() == at::kBool, "compressed_mask must be bool, got ", + compressed_mask.scalar_type()); tk::postProcessScatterLaunch(kv_comp.data_ptr(), kv_out.has_value() ? kv_out->data_ptr() : nullptr, rms_weight.data_ptr(), static_cast(rms_eps), cos_sin_table.data_ptr(), - position_ids.data_ptr(), static_cast(nope_dim), static_cast(rope_dim), - kv_cache.data_ptr(), num_outputs.data_ptr(), cu_kv_comp.data_ptr(), - start_pos.data_ptr(), block_offsets.data_ptr(), - reinterpret_cast(compressed_mask.data_ptr()), + position_ids.data_ptr(), static_cast(nope_dim), static_cast(rope_dim), kv_cache.data_ptr(), + num_outputs.data_ptr(), cu_kv_comp.data_ptr(), start_pos.data_ptr(), + block_offsets.data_ptr(), reinterpret_cast(compressed_mask.data_ptr()), static_cast(num_outputs.size(0)), // batch_size static_cast(tokens_per_block), static_cast(kv_comp.size(1)), // head_dim diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index e7830d8f9fb0..d1e6e4b77e08 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -772,9 +772,8 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): pretrained_config, 'compress_ratios', None) num_base_layers = pretrained_config.num_hidden_layers spec_config = kwargs.get('spec_config', None) - mtp_enabled = ( - spec_config is not None - and spec_config.spec_dec_mode.is_mtp_one_model()) + mtp_enabled = (spec_config is not None and + spec_config.spec_dec_mode.is_mtp_one_model()) sparse_attention_config = kwargs.get( 'sparse_attention_config') checkpoint_window_size = getattr( @@ -793,12 +792,14 @@ def update_sparse_attention_indexer_config(pretrained_config, kwargs): if (checkpoint_compress_ratios is not None and (compress_ratios is None - or len(checkpoint_compress_ratios) > - len(compress_ratios))): + or len(checkpoint_compress_ratios) + > len(compress_ratios))): compress_ratios = checkpoint_compress_ratios if window_size is None: window_size = checkpoint_window_size + if window_size is None: + window_size = pretrained_config.sliding_window # Normalize checkpoint-facing ratio 0 (SWA-only/uncompressed) # to 1 internally so cache allocation math works. The diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 463ec8ff344e..e0ba8b0179b2 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -13,6 +13,7 @@ from .modeling_clip import CLIPVisionModel from .modeling_cohere2 import Cohere2ForCausalLM from .modeling_deepseekv3 import DeepseekV3ForCausalLM +from .modeling_deepseekv4 import DeepseekV4ForCausalLM from .modeling_exaone4 import Exaone4ForCausalLM from .modeling_exaone4_5 import Exaone4_5_ForConditionalGeneration from .modeling_exaone_moe import ExaoneMoeForCausalLM @@ -27,7 +28,6 @@ from .modeling_laguna import LagunaForCausalLM from .modeling_llama import LlamaForCausalLM from .modeling_llava_next import LlavaNextModel -from .modeling_deepseekv4 import DeepseekV4ForCausalLM from .modeling_minimaxm2 import MiniMaxM2ForCausalLM from .modeling_mistral import Mistral3VLM, MistralForCausalLM from .modeling_mixtral import MixtralForCausalLM diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 2c300cd1e3e5..7967cc807f28 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -1092,9 +1092,9 @@ def mla_custom_op_inplace( metadata, mla_layer = extract_extra_attrs(layer_idx, "mla") if mla_layer.is_deepseek_v4: mla_layer.forward_impl_with_deepseek_v4(position_ids, - hidden_states, - metadata, - output=output) + hidden_states, + metadata, + output=output) else: mla_layer.forward_impl(position_ids, hidden_states, @@ -1564,7 +1564,8 @@ def yarn_get_mscale(scale=1, mscale=1): kv_lora_rank=self.kv_lora_rank, qk_nope_head_dim=self.qk_nope_head_dim, qk_rope_head_dim=self.qk_rope_head_dim, - v_head_dim=self.v_head_dim if self.is_deepseek_v4 else self.kv_lora_rank, + v_head_dim=self.v_head_dim + if self.is_deepseek_v4 else self.kv_lora_rank, hidden_size=self.hidden_size, predicted_tokens_per_seq=self.predicted_tokens_per_seq, skip_create_weights_in_init=config.skip_create_weights_in_init, @@ -1823,7 +1824,7 @@ def _get_attn_scale(position_ids: torch.Tensor) -> torch.Tensor: return q def _deepseek_v4_o_proj(self, attn_out_latent: torch.Tensor, - position_ids: torch.Tensor) -> torch.Tensor: + position_ids: torch.Tensor) -> torch.Tensor: num_tokens = attn_out_latent.shape[0] attn_out_latent = attn_out_latent.view(num_tokens, self.num_heads_tp, -1) @@ -2033,8 +2034,7 @@ def forward_dsa_proj( assert self.mqa is not None, "DSA is only supported in MQA mode" q, compressed_kv, k_pe = self.kv_a_proj_with_mqa(hidden_states).split( - [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], - -1) + [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim], -1) q, compressed_kv = maybe_execute_in_parallel( lambda: self.q_a_layernorm(q), @@ -2170,10 +2170,11 @@ def forward_dsa_attn( topk_indices=topk_indices[num_ctx_tokens:num_tokens, :], ) - def forward_impl_with_deepseek_v4(self, position_ids: Optional[torch.Tensor], - hidden_states: torch.Tensor, - attn_metadata: AttentionMetadata, - output: torch.Tensor) -> None: + def forward_impl_with_deepseek_v4(self, + position_ids: Optional[torch.Tensor], + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + output: torch.Tensor) -> None: """ Forward pass for the MLA module with DeepSeek-V4 (always in MQA mode). @@ -2196,9 +2197,8 @@ def forward_impl_with_deepseek_v4(self, position_ids: Optional[torch.Tensor], if position_ids is not None: position_ids = position_ids[..., :num_tokens] - q, kv = self.kv_a_proj_with_mqa(hidden_states).split([ - self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim - ], -1) + q, kv = self.kv_a_proj_with_mqa(hidden_states).split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], -1) q, kv = maybe_execute_in_parallel( lambda: self.q_a_layernorm(q), @@ -3337,9 +3337,9 @@ def forward( output=attn_output) elif self.is_deepseek_v4: self.forward_impl_with_deepseek_v4(position_ids, - hidden_states, - attn_metadata, - output=attn_output) + hidden_states, + attn_metadata, + output=attn_output) else: self.forward_impl(position_ids, hidden_states, diff --git a/tensorrt_llm/_torch/modules/fused_moe/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/__init__.py index 293e695c452b..a7467f7abbd7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/__init__.py @@ -12,10 +12,10 @@ from .quantization import FusedMoEQuantScalesFP8 # yapf: disable from .routing import (BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod, - DefaultMoeRoutingMethod, + DeepSeekV4MoeRoutingMethod, DefaultMoeRoutingMethod, Llama4RenormalizeMoeRoutingMethod, - LoadBalancedMoeRoutingMethod, DeepSeekV4MoeRoutingMethod, - MiniMaxM2MoeRoutingMethod, RenormalizeMoeRoutingMethod, + LoadBalancedMoeRoutingMethod, MiniMaxM2MoeRoutingMethod, + RenormalizeMoeRoutingMethod, RenormalizeNaiveMoeRoutingMethod, RoutingMethodType, SigmoidRenormMoeRoutingMethod, SparseMixerMoeRoutingMethod, StaticMoeRoutingMethod, diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index 32965bcf3263..758e6a624e1f 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -43,7 +43,8 @@ def __init__( self.intermediate_size = intermediate_size self.activation = activation self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm - self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None + self.swiglu_limit = float( + swiglu_limit) if swiglu_limit is not None else None config = config or ModelConfig() self.mapping = config.mapping diff --git a/tensorrt_llm/_torch/modules/mhc/hyper_connection.py b/tensorrt_llm/_torch/modules/mhc/hyper_connection.py index 051eeaa95a48..146278f4e2e0 100644 --- a/tensorrt_llm/_torch/modules/mhc/hyper_connection.py +++ b/tensorrt_llm/_torch/modules/mhc/hyper_connection.py @@ -49,11 +49,10 @@ def deferred( ) -> "HCState": return cls(residual=residual, post_mix=post_mix, comb_mix=comb_mix, x_prev=x_prev) + try: + from tensorrt_llm._torch.modules.mhc.mhc_cuda import mhc_fused_hc as mhc_fused_hc_cuda from tensorrt_llm._torch.modules.mhc.mhc_cuda import mhc_hc_head_cuda, mhc_post_mapping_cuda - from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( - mhc_fused_hc as mhc_fused_hc_cuda, - ) from tensorrt_llm._torch.modules.mhc.mhc_cuda import ( mhc_pre_mapping_fused as mhc_pre_mapping_fused_cuda, ) diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 125c5d6adf4d..881d672d0993 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -43,6 +43,7 @@ def _fused_hc_mma_supported() -> bool: except Exception: return False + _DG_NUM_SPLITS = 16 @@ -719,9 +720,7 @@ def _alloc_fused_hc_outputs( B: int, n: int, hidden_size: int, num_k_splits: int, tile_m: int, device ): """Uncached fallback (kept for API compatibility).""" - return _FusedHcWorkspaceCache(n=n, hidden_size=hidden_size).get( - B, num_k_splits, tile_m, device - ) + return _FusedHcWorkspaceCache(n=n, hidden_size=hidden_size).get(B, num_k_splits, tile_m, device) # Fallback tactic: backend, tile_n, num_k_splits, bigfuse_bs, tile_m. diff --git a/tensorrt_llm/llmapi/tokenizer.py b/tensorrt_llm/llmapi/tokenizer.py index 5f593d5d6000..67a7bf5c227f 100644 --- a/tensorrt_llm/llmapi/tokenizer.py +++ b/tensorrt_llm/llmapi/tokenizer.py @@ -6,8 +6,8 @@ _llguidance_tokenizer_info, _xgrammar_tokenizer_info, load_hf_tokenizer, tokenizer_factory) -from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer +from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer __all__ = [ "TLLM_INCREMENTAL_DETOKENIZATION_BACKEND", diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py index 7ef119515652..d6bf75f3d19d 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py @@ -9,17 +9,17 @@ import torch from utils.util import skip_pre_blackwell -from tensorrt_llm._torch.attention_backend.sparse.kernel import deepseek_v4_local_to_global_indices from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import ( DeepseekV4AttentionType, DeepseekV4CacheManager, ) from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import get_token_bytes +from tensorrt_llm._torch.attention_backend.sparse.kernel import deepseek_v4_local_to_global_indices from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, DeepSeekV4SparseAttentionConfig +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index fd23e7032aac..bc7a69781da6 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -21,14 +21,16 @@ DeepseekV4CacheManager, DeepseekV4TrtllmAttention, ) -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( + DeepseekV4TrtllmAttentionMetadata, +) from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, DeepSeekV4SparseAttentionConfig +from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 06ed2743543e..c1aa41102c82 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -123,13 +123,13 @@ def create_indexer(sparse_attn_config, layer_idx=0): rope_params = RopeParams( dim=64, # qk_rope_head_dim theta=10000.0, - max_positions=4096) + max_positions=4096, + ) # Create PositionalEmbeddingParams with required 'type' argument pos_embd_params = PositionalEmbeddingParams( - type=PositionEmbeddingType.rope_gpt_neox, - rope=rope_params, - is_neox=True) + type=PositionEmbeddingType.rope_gpt_neox, rope=rope_params, is_neox=True + ) # Create MLAParams class MLAParams: @@ -144,13 +144,10 @@ def __init__(self, head_dim): mla_params = MLAParams(sparse_attn_config.index_head_dim) # Mock RotaryEmbedding since we're only testing cache management, not rope functionality - with patch( - 'tensorrt_llm._torch.attention_backend.sparse.dsa.RotaryEmbedding' - ) as mock_rope: + with patch("tensorrt_llm._torch.attention_backend.sparse.dsa.RotaryEmbedding") as mock_rope: # Create a mock instance with a simple forward method mock_rope_instance = Mock() - mock_rope_instance.forward = Mock( - side_effect=lambda pos_ids, tensors: tensors) + mock_rope_instance.forward = Mock(side_effect=lambda pos_ids, tensors: tensors) mock_rope.return_value = mock_rope_instance indexer = Indexer( @@ -160,7 +157,8 @@ def __init__(self, head_dim): skip_create_weights_in_init=True, # Skip weight creation for test sparse_attention_config=sparse_attn_config, dtype=torch.bfloat16, - layer_idx=layer_idx) + layer_idx=layer_idx, + ) return indexer @@ -241,8 +239,7 @@ def _ref_fp8_paged_mqa_logits( q_offsets = torch.arange(num_token - next_n, num_token, device="cuda") # Transpose weights for this sequence: [num_heads, next_n] - weight_slice = (weights[i * next_n:(i + 1) * next_n, :].transpose( - 0, 1).contiguous()) + weight_slice = weights[i * next_n : (i + 1) * next_n, :].transpose(0, 1).contiguous() # Process each block in the sequence for block_rk in range(cdiv(num_kv_token, block_size)): @@ -258,15 +255,13 @@ def _ref_fp8_paged_mqa_logits( # Causal mask: k_pos < num_token AND k_pos < (q_pos + 1) // compress_ratio k_mask = k_offsets[None, :] < num_kv_token - causal_mask = k_offsets[None, :] < (q_offsets[:, None] + - 1) // compress_ratio + causal_mask = k_offsets[None, :] < (q_offsets[:, None] + 1) // compress_ratio mask = k_mask & causal_mask # Compute attention scores: [num_heads, next_n, block_size] s = torch.where( mask[None, :, :], - (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to( - logits.dtype), + (qx.transpose(0, 1) @ kx.transpose(0, 1).transpose(1, 2)).to(logits.dtype), float("-inf"), ) @@ -276,8 +271,8 @@ def _ref_fp8_paged_mqa_logits( # Write to output with additional causal mask logits[ - i * next_n:(i + 1) * next_n, - block_rk * block_size:(block_rk + 1) * block_size, + i * next_n : (i + 1) * next_n, + block_rk * block_size : (block_rk + 1) * block_size, ] = torch.where(causal_mask, s, float("-inf")) return logits @@ -313,10 +308,8 @@ def _ref_fp8_mqa_logits( q = q.float() k = k.float() - mask_lo = (torch.arange(0, seq_len_kv, device="cuda")[None, :] - >= cu_seqlen_ks[:, None]) - mask_hi = (torch.arange(0, seq_len_kv, device="cuda")[None, :] - < cu_seqlen_ke[:, None]) + mask_lo = torch.arange(0, seq_len_kv, device="cuda")[None, :] >= cu_seqlen_ks[:, None] + mask_hi = torch.arange(0, seq_len_kv, device="cuda")[None, :] < cu_seqlen_ke[:, None] mask = mask_lo & mask_hi score = torch.einsum("mhd,nd->hmn", q, k) @@ -340,7 +333,7 @@ def test_deepgemm_fp8_mqa_logits_basic(compress_ratio): seq_len = 2048 seq_len_kv = 4096 seq_len_kv_compressed = seq_len_kv // compress_ratio - #[seq_len, num_heads, head_dim] + # [seq_len, num_heads, head_dim] q = torch.randn( seq_len, num_heads, @@ -348,14 +341,14 @@ def test_deepgemm_fp8_mqa_logits_basic(compress_ratio): device="cuda", dtype=torch.bfloat16, ) - #[seq_len_kv, head_dim] -> num_head = 1 + # [seq_len_kv, head_dim] -> num_head = 1 kv = torch.randn( seq_len_kv_compressed, head_dim, device="cuda", dtype=torch.bfloat16, ) - #[seq_len, num_heads] + # [seq_len, num_heads] weights = torch.randn( seq_len, num_heads, @@ -368,22 +361,20 @@ def test_deepgemm_fp8_mqa_logits_basic(compress_ratio): num_contexts=1, num_ctx_tokens=seq_len, cached_token_lens=torch.tensor([512], dtype=torch.int, device="cuda"), - kv_lens=torch.tensor([seq_len_kv_compressed], - dtype=torch.int, - device="cuda"), - compress_ratio=compress_ratio) + kv_lens=torch.tensor([seq_len_kv_compressed], dtype=torch.int, device="cuda"), + compress_ratio=compress_ratio, + ) # Convert to FP8 q_fp8 = q.to(torch.float8_e4m3fn) - kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0, ), False) - logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, - ke) # -> [seq_len, seq_len_kv] + kv_fp8 = per_custom_dims_cast_to_fp8(kv, (0,), False) + logits = deep_gemm.fp8_mqa_logits(q_fp8, kv_fp8, weights, ks, ke) # -> [seq_len, seq_len_kv] # Basic sanity checks - assert logits.shape == (seq_len, seq_len_kv_compressed), \ + assert logits.shape == (seq_len, seq_len_kv_compressed), ( f"Expected shape ({seq_len}, {seq_len_kv_compressed}), got {logits.shape}" - assert logits.dtype == torch.float32, \ - f"Expected dtype torch.float32, got {logits.dtype}" + ) + assert logits.dtype == torch.float32, f"Expected dtype torch.float32, got {logits.dtype}" ref_logits = _ref_fp8_mqa_logits( q=q, @@ -402,28 +393,31 @@ def test_deepgemm_fp8_mqa_logits_basic(compress_ratio): diff = _calc_diff(logits, ref_logits) assert diff < 1e-3, f"{diff=}" # check for cosine similarity - check_accuracy(logits, ref_logits, atol=1e-2, rtol=2e-1, - percent=0.9) # double check for per-element similarity - - -def _create_mock_metadata(request_ids, - batch_size, - num_contexts, - num_generations, - seq_lens, - kv_lens, - num_cached_tokens, - cache_manager, - num_ctx_tokens, - num_tokens, - indexer_max_chunk_size=8194, - max_draft_tokens=0, - enable_context_mla_with_cached_kv=False, - index_topk=2048, - enable_indexer_skip=False, - use_cute_dsl_paged_mqa_logits=False, - compress_ratio=1, - indexer_head_dim=128): + check_accuracy( + logits, ref_logits, atol=1e-2, rtol=2e-1, percent=0.9 + ) # double check for per-element similarity + + +def _create_mock_metadata( + request_ids, + batch_size, + num_contexts, + num_generations, + seq_lens, + kv_lens, + num_cached_tokens, + cache_manager, + num_ctx_tokens, + num_tokens, + indexer_max_chunk_size=8194, + max_draft_tokens=0, + enable_context_mla_with_cached_kv=False, + index_topk=2048, + enable_indexer_skip=False, + use_cute_dsl_paged_mqa_logits=False, + compress_ratio=1, + indexer_head_dim=128, +): """Helper to create mock metadata for testing.""" class MockKVCacheParams: @@ -458,29 +452,29 @@ def __init__(self): # Effective tokens-per-block for the indexer k-cache slot mapping, # mirroring DSAtrtllmAttentionMetadata.__post_init__ (dsa.py). tpb = self.kv_cache_manager.tokens_per_block - if hasattr(self.kv_cache_manager, 'compressed_block_sizes'): + if hasattr(self.kv_cache_manager, "compressed_block_sizes"): _cr = 4 if 4 in self.compress_ratios else 1 tpb = tpb // _cr self._tokens_per_block = tpb self.kv_lens_cuda_runtime = kv_lens.cuda() self.indexer_k_cache_block_offsets = torch.zeros( [batch_size, cache_manager.max_blocks_per_seq], - device='cuda', + device="cuda", dtype=torch.int32, ) self.host_indexer_k_cache_block_offsets = torch.zeros_like( self.indexer_k_cache_block_offsets, - device='cpu', + device="cpu", pin_memory=True, ) block_ids = cache_manager.get_batch_cache_indices(request_ids) for i in range(len(block_ids)): - self.host_indexer_k_cache_block_offsets[ - i, :len(block_ids[i])] = torch.tensor(block_ids[i], - dtype=torch.int32) + self.host_indexer_k_cache_block_offsets[i, : len(block_ids[i])] = torch.tensor( + block_ids[i], dtype=torch.int32 + ) self.indexer_k_cache_block_offsets[:batch_size].copy_( - self.host_indexer_k_cache_block_offsets[:batch_size], - non_blocking=True) + self.host_indexer_k_cache_block_offsets[:batch_size], non_blocking=True + ) self.slot_mapping_fp8 = torch.zeros((num_tokens, ), device='cuda', @@ -528,21 +522,18 @@ def __init__(self): device='cpu', pin_memory=True) self.host_slot_mapping_scale = torch.zeros_like( - self.slot_mapping_scale, device='cpu', pin_memory=True) - self.host_ctx_kv_indptr = torch.zeros((num_contexts + 1, ), - device='cpu', - pin_memory=True, - dtype=torch.int64) - self.host_gen_kv_indptr = torch.zeros((num_generations + 1, ), - device='cpu', - pin_memory=True, - dtype=torch.int64) + self.slot_mapping_scale, device="cpu", pin_memory=True + ) + self.host_ctx_kv_indptr = torch.zeros( + (num_contexts + 1,), device="cpu", pin_memory=True, dtype=torch.int64 + ) + self.host_gen_kv_indptr = torch.zeros( + (num_generations + 1,), device="cpu", pin_memory=True, dtype=torch.int64 + ) # Add host_ctx_cached_token_indptr for prepare_one_prefill_chunk self.host_ctx_cached_token_indptr = torch.zeros( - (num_contexts + 1, ), - device='cpu', - pin_memory=True, - dtype=torch.int64) + (num_contexts + 1,), device="cpu", pin_memory=True, dtype=torch.int64 + ) self._num_ctx_tokens = num_ctx_tokens self._num_tokens = num_tokens self.num_gen_tokens = num_tokens - num_ctx_tokens @@ -552,23 +543,27 @@ def __init__(self): self._num_ctx_tokens = num_ctx_tokens self._num_tokens = num_tokens - torch.cumsum(kv_lens[:num_contexts], - dim=0, - dtype=torch.int64, - out=self.host_ctx_kv_indptr[1:num_contexts + 1]) - torch.cumsum(kv_lens[num_contexts:num_contexts + num_generations], - dim=0, - dtype=torch.int64, - out=self.host_gen_kv_indptr[1:num_generations + 1]) + torch.cumsum( + kv_lens[:num_contexts], + dim=0, + dtype=torch.int64, + out=self.host_ctx_kv_indptr[1 : num_contexts + 1], + ) + torch.cumsum( + kv_lens[num_contexts : num_contexts + num_generations], + dim=0, + dtype=torch.int64, + out=self.host_gen_kv_indptr[1 : num_generations + 1], + ) # Compute host_ctx_cached_token_indptr from num_cached_tokens if num_contexts > 0: - cached_lens = torch.tensor(num_cached_tokens[:num_contexts], - dtype=torch.int64) + cached_lens = torch.tensor(num_cached_tokens[:num_contexts], dtype=torch.int64) torch.cumsum( cached_lens, dim=0, dtype=torch.int64, - out=self.host_ctx_cached_token_indptr[1:num_contexts + 1]) + out=self.host_ctx_cached_token_indptr[1 : num_contexts + 1], + ) # Add indexer-specific attributes self.indexer_max_chunk_size = indexer_max_chunk_size @@ -598,74 +593,79 @@ def __init__(self): ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) and get_sm_version() >= 100))) self.kv_lens_expanded_cuda = torch.zeros( - (self.num_seqs * (1 + self.max_draft_tokens), ), - device='cuda', - dtype=torch.int32) + (self.num_seqs * (1 + self.max_draft_tokens),), device="cuda", dtype=torch.int32 + ) self.kv_lens_expanded_host = torch.zeros_like( - self.kv_lens_expanded_cuda, device='cpu', pin_memory=True) + self.kv_lens_expanded_cuda, device="cpu", pin_memory=True + ) self.block_table_expanded = torch.zeros( - (self.num_seqs * (1 + self.max_draft_tokens), - self.kv_cache_manager.max_blocks_per_seq), - device='cuda', - dtype=torch.int32) + ( + self.num_seqs * (1 + self.max_draft_tokens), + self.kv_cache_manager.max_blocks_per_seq, + ), + device="cuda", + dtype=torch.int32, + ) self.host_block_table_expanded = torch.zeros_like( - self.block_table_expanded, device='cpu', pin_memory=True) + self.block_table_expanded, device="cpu", pin_memory=True + ) self.scheduler_metadata_buffer_expanded = torch.zeros( - (self.num_sms + 1, 2), device='cuda', dtype=torch.int32) + (self.num_sms + 1, 2), device="cuda", dtype=torch.int32 + ) + if self.max_draft_tokens == 3: + self.scheduler_metadata_buffer_mtp3 = torch.zeros( + (self.num_sms // 2 + 1, 2), device="cuda", dtype=torch.int32 + ) if self.use_expanded_buffers_for_mtp: - gen_kv_lens = kv_lens[num_contexts:self.num_seqs] - gen_kv_lens_expanded = torch.stack([gen_kv_lens] * - (1 + self.max_draft_tokens), - dim=0) - gen_kv_lens_expanded = gen_kv_lens_expanded.transpose( - 0, 1).contiguous().flatten() - self.kv_lens_expanded_host[:self.num_gen_tokens].copy_( - gen_kv_lens_expanded) - self.kv_lens_expanded_cuda[:self.num_gen_tokens].copy_( - self.kv_lens_expanded_host[:self.num_gen_tokens], - non_blocking=True) + gen_kv_lens = kv_lens[num_contexts : self.num_seqs] + gen_kv_lens_expanded = torch.stack( + [gen_kv_lens] * (1 + self.max_draft_tokens), dim=0 + ) + gen_kv_lens_expanded = gen_kv_lens_expanded.transpose(0, 1).contiguous().flatten() + self.kv_lens_expanded_host[: self.num_gen_tokens].copy_(gen_kv_lens_expanded) + self.kv_lens_expanded_cuda[: self.num_gen_tokens].copy_( + self.kv_lens_expanded_host[: self.num_gen_tokens], non_blocking=True + ) if self.kv_cache_manager is not None: - block_ids = self.kv_cache_manager.get_batch_cache_indices( - self.request_ids) - gen_block_ids = block_ids[self.num_contexts:] + block_ids = self.kv_cache_manager.get_batch_cache_indices(self.request_ids) + gen_block_ids = block_ids[self.num_contexts :] if len(gen_block_ids) > 0: # Find max length and create padded tensor max_len = max(len(bid) for bid in gen_block_ids) gen_block_tensor = self.host_indexer_k_cache_block_offsets[ - self.num_contexts:self.num_seqs, :max_len] + self.num_contexts : self.num_seqs, :max_len + ] expanded_blocks = gen_block_tensor.repeat_interleave( - 1 + self.max_draft_tokens, dim=0) - self.host_block_table_expanded[:self.num_gen_tokens, : - max_len].copy_( - expanded_blocks, - non_blocking=True) - self.block_table_expanded[:self.num_gen_tokens].copy_( - self.host_block_table_expanded[:self. - num_gen_tokens], - non_blocking=True) + 1 + self.max_draft_tokens, dim=0 + ) + self.host_block_table_expanded[: self.num_gen_tokens, :max_len].copy_( + expanded_blocks, non_blocking=True + ) + self.block_table_expanded[: self.num_gen_tokens].copy_( + self.host_block_table_expanded[: self.num_gen_tokens], non_blocking=True + ) # Add skip indexer attributes self.topk_indices_buffer = torch.zeros( - (num_tokens, self.num_sparse_topk), - device='cuda', - dtype=torch.int32) + (num_tokens, self.num_sparse_topk), device="cuda", dtype=torch.int32 + ) if self.num_contexts > 0 and self.enable_indexer_skip: - self.skip_indexer_for_ctx_reqs = kv_lens[:self.num_contexts].max( - ).item() <= self.num_sparse_topk + self.skip_indexer_for_ctx_reqs = ( + kv_lens[: self.num_contexts].max().item() <= self.num_sparse_topk + ) else: self.skip_indexer_for_ctx_reqs = False if self.num_generations > 0 and self.enable_indexer_skip: self.max_draft_tokens + 1 - self.skip_indexer_for_gen_reqs = kv_lens[ - self.num_contexts:self.num_seqs].max().item( - ) <= self.num_sparse_topk + self.skip_indexer_for_gen_reqs = ( + kv_lens[self.num_contexts : self.num_seqs].max().item() <= self.num_sparse_topk + ) else: self.skip_indexer_for_gen_reqs = False - self.prepare_dense_topk_indices(self.kv_lens_cuda_runtime, - device=True) + self.prepare_dense_topk_indices(self.kv_lens_cuda_runtime, device=True) @property def num_seqs(self) -> int: @@ -735,15 +735,15 @@ def test_indexer_k_cache_scatter_custom_op(): head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_seq_len, - num_layers=3) # Multi-layer pool for non-contiguous test + num_layers=3, + ) # Multi-layer pool for non-contiguous test # Allocate blocks request_ids = list(range(batch_size)) tokens_per_req = [32, 32, 32] - cache_manager.add_dummy_requests(request_ids, - tokens_per_req, - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids, tokens_per_req, is_gen=False, prepare_resource=True + ) # Create metadata metadata = _create_mock_metadata( @@ -761,12 +761,11 @@ def test_indexer_k_cache_scatter_custom_op(): ) from tensorrt_llm._torch.attention_backend.sparse.dsa import Indexer + Indexer.prepare(metadata) # Generate test data - k_original = torch.randn((num_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) + k_original = torch.randn((num_tokens, head_dim), device="cuda", dtype=torch.bfloat16) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_original) # Prepare byte-level data for the Python reference path @@ -774,9 +773,9 @@ def test_indexer_k_cache_scatter_custom_op(): k_fp8_bytes = k_fp8.view(-1).view(torch.uint8).view(num_tokens, head_dim) k_scale_flat = k_scale.view(-1) if k_scale_flat.stride(-1) != 1: - k_scale_flat = torch.as_strided(k_scale_flat.contiguous(), - size=(k_scale_flat.numel(), ), - stride=(1, )) + k_scale_flat = torch.as_strided( + k_scale_flat.contiguous(), size=(k_scale_flat.numel(),), stride=(1,) + ) k_scale_bytes = k_scale_flat.view(torch.uint8).view(num_tokens, scale_size) flat_indices_fp8 = metadata.slot_mapping_fp8[:num_tokens] @@ -808,10 +807,14 @@ def test_indexer_k_cache_scatter_custom_op(): # ========== Path 1: CUDA Kernel ========== print("\n=== Path 1: CUDA Kernel ===") - torch.ops.trtllm.indexer_k_cache_scatter_op(k_fp8, k_scale, k_cache_cuda, - metadata.slot_mapping_fp8, - metadata.slot_mapping_scale, - num_tokens) + torch.ops.trtllm.indexer_k_cache_scatter_op( + k_fp8, + k_scale, + k_cache_cuda, + metadata.slot_mapping_fp8, + metadata.slot_mapping_scale, + num_tokens, + ) torch.cuda.synchronize() print("✓ CUDA kernel completed") @@ -833,19 +836,15 @@ def _unravel_indices(flat_indices, shape): return i0, i1, i2, i3 # Scatter FP8 data - byte_offsets = torch.arange(head_dim, - device=k_cache_python.device).unsqueeze(0) + byte_offsets = torch.arange(head_dim, device=k_cache_python.device).unsqueeze(0) scatter_indices_fp8 = flat_indices_fp8.unsqueeze(1) + byte_offsets - scatter_indices_fp8 = _unravel_indices(scatter_indices_fp8, - k_cache_python.shape) + scatter_indices_fp8 = _unravel_indices(scatter_indices_fp8, k_cache_python.shape) k_cache_python[scatter_indices_fp8] = k_fp8_bytes # Scatter scale data - byte_offsets = torch.arange(scale_size, - device=k_cache_python.device).unsqueeze(0) + byte_offsets = torch.arange(scale_size, device=k_cache_python.device).unsqueeze(0) scatter_indices_scale = flat_indices_scale.unsqueeze(1) + byte_offsets - scatter_indices_scale = _unravel_indices(scatter_indices_scale, - k_cache_python.shape) + scatter_indices_scale = _unravel_indices(scatter_indices_scale, k_cache_python.shape) k_cache_python[scatter_indices_scale] = k_scale_bytes # ========== Validation: Byte-for-Byte Comparison ========== @@ -857,16 +856,14 @@ def _unravel_indices(flat_indices, shape): if torch.equal(k_cache_cuda, k_cache_python): print("✅ PERFECT MATCH! CUDA and Python produce identical cache") print(f" Total bytes compared: {total_bytes}") - print( - f" Tokens: {num_tokens}, head_dim: {head_dim}, block_size: {block_size}" - ) + print(f" Tokens: {num_tokens}, head_dim: {head_dim}, block_size: {block_size}") else: # Find differences diff_mask = k_cache_cuda != k_cache_python num_diffs = diff_mask.sum().item() print( - f"⚠️ Found {num_diffs}/{total_bytes} byte differences ({100*num_diffs/total_bytes:.4f}%)" + f"⚠️ Found {num_diffs}/{total_bytes} byte differences ({100 * num_diffs / total_bytes:.4f}%)" ) # Show first few differences @@ -875,11 +872,11 @@ def _unravel_indices(flat_indices, shape): flat_idx = idx.item() print( f" Byte {flat_idx}: CUDA={k_cache_cuda.view(-1)[flat_idx].item()}, " - f"Python={k_cache_python.view(-1)[flat_idx].item()}") + f"Python={k_cache_python.view(-1)[flat_idx].item()}" + ) # Fail the test - raise AssertionError( - "CUDA kernel produced different results than Python reference") + raise AssertionError("CUDA kernel produced different results than Python reference") @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -899,15 +896,15 @@ def test_fp8_k_cache_roundtrip(): head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_seq_len, - num_layers=1) + num_layers=1, + ) indexer = create_indexer(sparse_attn_config, layer_idx=0) # Allocate blocks for both requests request_ids = [0, 1] - cache_manager.add_dummy_requests(request_ids, - num_tokens_per_req, - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids, num_tokens_per_req, is_gen=False, prepare_resource=True + ) # Prepare and write data for each request total_tokens = sum(num_tokens_per_req) @@ -927,17 +924,14 @@ def test_fp8_k_cache_roundtrip(): Indexer.prepare(metadata) # Generate unique patterns for each request and quantize - k_original = torch.randn((total_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) + k_original = torch.randn((total_tokens, head_dim), device="cuda", dtype=torch.bfloat16) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_original) # Write to cache indexer._update_k_cache(k_fp8, k_scale, metadata) # Verify scales for both requests - cache_flat = cache_manager.indexer_k_cache_pool_per_layer[ - 0] # [num_blocks, flat_bytes] + cache_flat = cache_manager.indexer_k_cache_pool_per_layer[0] # [num_blocks, flat_bytes] scale_offset = block_size * head_dim # Scales start after FP8 data scale_size = 4 # float32 @@ -950,20 +944,22 @@ def test_fp8_k_cache_roundtrip(): # Compute block location block_idx_in_seq = local_token_idx // block_size pos_in_block = local_token_idx % block_size - block_id = metadata.host_indexer_k_cache_block_offsets[ - req_idx, block_idx_in_seq].item() + block_id = metadata.host_indexer_k_cache_block_offsets[req_idx, block_idx_in_seq].item() # Extract stored scale - scale_bytes = cache_flat[block_id, scale_offset + - pos_in_block * scale_size:scale_offset + - (pos_in_block + 1) * scale_size] + scale_bytes = cache_flat[ + block_id, + scale_offset + pos_in_block * scale_size : scale_offset + + (pos_in_block + 1) * scale_size, + ] stored_scale = scale_bytes.view(torch.float32).item() # Compare with original orig_scale = original_scales[global_token_idx] - assert abs(orig_scale - stored_scale) < 1e-6, \ - f"Request {req_idx}, token {local_token_idx} (block {block_idx_in_seq}, pos {pos_in_block}): " \ + assert abs(orig_scale - stored_scale) < 1e-6, ( + f"Request {req_idx}, token {local_token_idx} (block {block_idx_in_seq}, pos {pos_in_block}): " f"scale mismatch (orig={orig_scale:.6f}, stored={stored_scale:.6f})" + ) global_token_idx += 1 @@ -1007,11 +1003,13 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, layer_idx = 0 # Generate variable context lengths per sequence - context_lens_context = torch.randint(int(0.7 * avg_context_len), - int(1.4 * avg_context_len), - (batch_size, ), - dtype=torch.int32, - device="cpu") + context_lens_context = torch.randint( + int(0.7 * avg_context_len), + int(1.4 * avg_context_len), + (batch_size,), + dtype=torch.int32, + device="cpu", + ) # Final lengths after generation phase final_lens = context_lens_context + num_gen_tokens @@ -1025,9 +1023,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, total_gen_kv_tokens = num_gen_kv_tokens.sum().item() print("\n=== Test Config ===") - print( - f" Batch: {batch_size}, Next_N: {next_n}, Heads: {heads}, Head_dim: {head_dim}" - ) + print(f" Batch: {batch_size}, Next_N: {next_n}, Heads: {heads}, Head_dim: {head_dim}") print(f" Context lengths: {context_lens_context.tolist()}") print(f" Final lengths: {final_lens.tolist()}") print(f" Final kv lengths: {final_kv_lens.tolist()}") @@ -1041,29 +1037,26 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_kv_len, - num_layers=1) + num_layers=1, + ) indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) # Allocate blocks for all sequences (max final length) request_ids = list(range(batch_size)) - cache_manager.add_dummy_requests(request_ids=request_ids, - token_nums=final_kv_lens.tolist(), - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids=request_ids, + token_nums=final_kv_lens.tolist(), + is_gen=False, + prepare_resource=True, + ) # Generate test data with variable lengths - q = torch.randn((batch_size, next_n, heads, head_dim), - device="cuda", - dtype=torch.bfloat16) - weights = torch.randn((batch_size * next_n, heads), - device="cuda", - dtype=torch.float32) - k_context_bf16 = torch.randn((total_context_kv_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - k_gen_bf16 = torch.randn((total_gen_kv_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) + q = torch.randn((batch_size, next_n, heads, head_dim), device="cuda", dtype=torch.bfloat16) + weights = torch.randn((batch_size * next_n, heads), device="cuda", dtype=torch.float32) + k_context_bf16 = torch.randn( + (total_context_kv_tokens, head_dim), device="cuda", dtype=torch.bfloat16 + ) + k_gen_bf16 = torch.randn((total_gen_kv_tokens, head_dim), device="cuda", dtype=torch.bfloat16) # Phase 1: Write context tokens (variable per sequence) as FP8 print("\n=== Phase 1: Context (variable tokens/seq) ===") @@ -1085,8 +1078,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, ) Indexer.prepare(metadata_context) - k_context_fp8, k_context_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( - k_context_bf16) + k_context_fp8, k_context_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_context_bf16) indexer._update_k_cache(k_context_fp8, k_context_scale, metadata_context) print(f"✓ Wrote {total_context_kv_tokens} FP8 context tokens to cache") @@ -1099,9 +1091,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, batch_size, num_contexts=0, num_generations=batch_size, - seq_lens=torch.tensor([num_gen_tokens] * batch_size, - dtype=torch.int32, - device='cpu'), + seq_lens=torch.tensor([num_gen_tokens] * batch_size, dtype=torch.int32, device="cpu"), kv_lens=final_kv_lens.clone(), num_cached_tokens=context_lens_context.tolist(), cache_manager=cache_manager, @@ -1114,8 +1104,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, ) Indexer.prepare(metadata_gen) - k_gen_fp8, k_gen_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( - k_gen_bf16) + k_gen_fp8, k_gen_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_gen_bf16) indexer._update_k_cache(k_gen_fp8, k_gen_scale, metadata_gen) print(f"✓ Wrote {total_gen_kv_tokens} FP8 generation tokens to cache") @@ -1165,9 +1154,9 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, # Reference: Reconstruct BF16 cache from original values print("\n=== Reference Computation ===") num_blocks = kv_cache_fp8_pool.shape[0] - kv_cache_bf16 = torch.zeros((num_blocks, block_size, 1, head_dim), - device="cuda", - dtype=torch.bfloat16) + kv_cache_bf16 = torch.zeros( + (num_blocks, block_size, 1, head_dim), device="cuda", dtype=torch.bfloat16 + ) # Populate cache with variable-length sequences context_offset = 0 @@ -1181,10 +1170,12 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, block_idx = token_pos // block_size pos_in_block = token_pos % block_size physical_block_id = metadata_gen.indexer_k_cache_block_offsets[ - seq_idx, block_idx].item() + seq_idx, block_idx + ].item() if physical_block_id >= 0: - kv_cache_bf16[physical_block_id, pos_in_block, 0, :] = \ - k_context_bf16[context_offset + token_pos] + kv_cache_bf16[physical_block_id, pos_in_block, 0, :] = k_context_bf16[ + context_offset + token_pos + ] # Write generation tokens for gen_token_idx in range(seq_gen_len): @@ -1192,10 +1183,12 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, block_idx = token_pos // block_size pos_in_block = token_pos % block_size physical_block_id = metadata_gen.indexer_k_cache_block_offsets[ - seq_idx, block_idx].item() + seq_idx, block_idx + ].item() if physical_block_id >= 0: - kv_cache_bf16[physical_block_id, pos_in_block, 0, :] = \ - k_gen_bf16[gen_offset + gen_token_idx] + kv_cache_bf16[physical_block_id, pos_in_block, 0, :] = k_gen_bf16[ + gen_offset + gen_token_idx + ] context_offset += seq_context_len gen_offset += seq_gen_len @@ -1203,10 +1196,15 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, num_tokens_cuda = final_lens.cuda() num_kv_tokens_cuda = metadata_gen.kv_lens_cuda_runtime ref_logits = _ref_fp8_paged_mqa_logits( - q, kv_cache_bf16, weights, num_tokens_cuda[0:batch_size], + q, + kv_cache_bf16, + weights, + num_tokens_cuda[0:batch_size], num_kv_tokens_cuda[0:batch_size], - metadata_gen.indexer_k_cache_block_offsets, max_model_len, - compress_ratio) + metadata_gen.indexer_k_cache_block_offsets, + max_model_len, + compress_ratio, + ) print(f"✓ Reference output shape: {ref_logits.shape}") # Validate: Compare masked outputs (handle variable lengths and next_n) @@ -1214,31 +1212,27 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, # Expand context lens for each query: each sequence has next_n queries # Query at position i (where i = 0..next_n-1) attends to tokens up to (context_len - next_n + i) - positions = torch.arange(max_model_len, - device="cuda").unsqueeze(0) # [1, max_model_len] + positions = torch.arange(max_model_len, device="cuda").unsqueeze(0) # [1, max_model_len] # For each query, compute its end position # Shape: [batch_size * next_n] - row_indices = torch.arange(batch_size * next_n, - device="cuda") // next_n # Which sequence - next_n_offset = torch.arange( - batch_size * next_n, - device="cuda") % next_n # Query offset within sequence - query_end_positions = num_tokens_cuda[ - row_indices] - next_n + next_n_offset # [batch_size * next_n] + row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n # Which sequence + next_n_offset = ( + torch.arange(batch_size * next_n, device="cuda") % next_n + ) # Query offset within sequence + query_end_positions = ( + num_tokens_cuda[row_indices] - next_n + next_n_offset + ) # [batch_size * next_n] # Create mask: positions <= query_end_position # Shape: [batch_size * next_n, max_model_len] mask = positions < (query_end_positions.unsqueeze(1) + 1) // compress_ratio - diff = _calc_diff(logits.masked_fill(~mask, 0), - ref_logits.masked_fill(~mask, 0)) + diff = _calc_diff(logits.masked_fill(~mask, 0), ref_logits.masked_fill(~mask, 0)) assert diff < 1e-3, f"Accuracy check failed: {diff=}" print(f"✅ Test passed! Accuracy: {diff:.6f} < 1e-3") - print( - f" Total cache tokens: {final_lens.sum().item()}, Avg: {final_lens.float().mean():.1f}" - ) + print(f" Total cache tokens: {final_lens.sum().item()}, Avg: {final_lens.float().mean():.1f}") @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -1650,7 +1644,8 @@ def test_compute_cu_seqlen_bounds_nocache(): num_ctx_tokens = 7 cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, None) + seq_lens, num_contexts, num_ctx_tokens, None + ) # Expected results: # Seq 0: tokens [0,1,2], KV [0,1,2] @@ -1663,17 +1658,15 @@ def test_compute_cu_seqlen_bounds_nocache(): # Token 5: [3, 6) # Token 6: [3, 7) - expected_ks = torch.tensor([0, 0, 0, 3, 3, 3, 3], - dtype=torch.int32, - device="cuda") - expected_ke = torch.tensor([1, 2, 3, 4, 5, 6, 7], - dtype=torch.int32, - device="cuda") + expected_ks = torch.tensor([0, 0, 0, 3, 3, 3, 3], dtype=torch.int32, device="cuda") + expected_ke = torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.int32, device="cuda") - assert torch.equal(cu_seqlen_ks, expected_ks), \ + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) def test_compute_cu_seqlen_bounds_with_cache(): @@ -1694,7 +1687,8 @@ def test_compute_cu_seqlen_bounds_with_cache(): num_ctx_tokens = 7 # 3 + 4 new tokens total cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, cached_token_lens) + seq_lens, num_contexts, num_ctx_tokens, cached_token_lens + ) # Expected results: # @@ -1709,17 +1703,15 @@ def test_compute_cu_seqlen_bounds_with_cache(): # New Q token 2 (local pos 2, global Q pos 5): attends to cached(1) + 0,1,2(3) = [5, 9) # New Q token 3 (local pos 3, global Q pos 6): attends to cached(1) + 0,1,2,3(4) = [5, 10) - expected_ks = torch.tensor([0, 0, 0, 5, 5, 5, 5], - dtype=torch.int32, - device="cuda") - expected_ke = torch.tensor([3, 4, 5, 7, 8, 9, 10], - dtype=torch.int32, - device="cuda") + expected_ks = torch.tensor([0, 0, 0, 5, 5, 5, 5], dtype=torch.int32, device="cuda") + expected_ke = torch.tensor([3, 4, 5, 7, 8, 9, 10], dtype=torch.int32, device="cuda") - assert torch.equal(cu_seqlen_ks, expected_ks), \ + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) def test_compute_cu_seqlen_bounds_with_cache_edge_cases(): @@ -1732,21 +1724,20 @@ def test_compute_cu_seqlen_bounds_with_cache_edge_cases(): num_ctx_tokens = 5 cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, cached_token_lens) + seq_lens, num_contexts, num_ctx_tokens, cached_token_lens + ) # Seq 0: KV [0:2], Q tokens attend to [0,1), [0,2) # Seq 1: KV [2:5], Q tokens attend to [2,3), [2,4), [2,5) - expected_ks = torch.tensor([0, 0, 2, 2, 2], - dtype=torch.int32, - device="cuda") - expected_ke = torch.tensor([1, 2, 3, 4, 5], - dtype=torch.int32, - device="cuda") - - assert torch.equal(cu_seqlen_ks, expected_ks), \ + expected_ks = torch.tensor([0, 0, 2, 2, 2], dtype=torch.int32, device="cuda") + expected_ke = torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32, device="cuda") + + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"Case 1 - cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"Case 1 - cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) # Case 2: Single new token per sequence seq_lens = torch.tensor([1, 1], dtype=torch.int32, device="cuda") @@ -1755,7 +1746,8 @@ def test_compute_cu_seqlen_bounds_with_cache_edge_cases(): num_ctx_tokens = 2 cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, cached_token_lens) + seq_lens, num_contexts, num_ctx_tokens, cached_token_lens + ) # Seq 0: 5 cached + 1 new = 6 total KV, global [0:6] # New Q token 0: attends to cached(5) + self(1) = [0, 6) @@ -1764,21 +1756,22 @@ def test_compute_cu_seqlen_bounds_with_cache_edge_cases(): expected_ks = torch.tensor([0, 6], dtype=torch.int32, device="cuda") expected_ke = torch.tensor([6, 10], dtype=torch.int32, device="cuda") - assert torch.equal(cu_seqlen_ks, expected_ks), \ + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"Case 2 - cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"Case 2 - cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) # Case 3: Different cached amounts across sequences seq_lens = torch.tensor([2, 1, 3], dtype=torch.int32, device="cuda") - cached_token_lens = torch.tensor([10, 0, 5], - dtype=torch.int32, - device="cuda") + cached_token_lens = torch.tensor([10, 0, 5], dtype=torch.int32, device="cuda") num_contexts = 3 num_ctx_tokens = 6 cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, cached_token_lens) + seq_lens, num_contexts, num_ctx_tokens, cached_token_lens + ) # Seq 0: 10 cached + 2 new = 12 total KV, global [0:12] # Q token 0 (local 0): cached(10) + self(1) = [0, 11) @@ -1789,17 +1782,15 @@ def test_compute_cu_seqlen_bounds_with_cache_edge_cases(): # Q token 0 (local 0): cached(5) + self(1) = [13, 19) # Q token 1 (local 1): cached(5) + 0,1(2) = [13, 20) # Q token 2 (local 2): cached(5) + 0,1,2(3) = [13, 21) - expected_ks = torch.tensor([0, 0, 12, 13, 13, 13], - dtype=torch.int32, - device="cuda") - expected_ke = torch.tensor([11, 12, 13, 19, 20, 21], - dtype=torch.int32, - device="cuda") - - assert torch.equal(cu_seqlen_ks, expected_ks), \ + expected_ks = torch.tensor([0, 0, 12, 13, 13, 13], dtype=torch.int32, device="cuda") + expected_ke = torch.tensor([11, 12, 13, 19, 20, 21], dtype=torch.int32, device="cuda") + + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"Case 3 - cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"Case 3 - cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) # Case 4: All tokens cached, single new token seq_lens = torch.tensor([1], dtype=torch.int32, device="cuda") @@ -1808,17 +1799,20 @@ def test_compute_cu_seqlen_bounds_with_cache_edge_cases(): num_ctx_tokens = 1 cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, cached_token_lens) + seq_lens, num_contexts, num_ctx_tokens, cached_token_lens + ) # Seq 0: 100 cached + 1 new = 101 total KV # Q token 0: attends to all = [0, 101) expected_ks = torch.tensor([0], dtype=torch.int32, device="cuda") expected_ke = torch.tensor([101], dtype=torch.int32, device="cuda") - assert torch.equal(cu_seqlen_ks, expected_ks), \ + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"Case 4 - cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"Case 4 - cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) def test_compute_cu_seqlen_bounds_with_cache_properties(): @@ -1826,40 +1820,39 @@ def test_compute_cu_seqlen_bounds_with_cache_properties(): for trial in range(10): num_contexts = random.randint(1, 5) - seq_lens = torch.randint(1, - 10, (num_contexts, ), - dtype=torch.int32, - device="cuda") - cached_token_lens = torch.randint(0, - 20, (num_contexts, ), - dtype=torch.int32, - device="cuda") + seq_lens = torch.randint(1, 10, (num_contexts,), dtype=torch.int32, device="cuda") + cached_token_lens = torch.randint(0, 20, (num_contexts,), dtype=torch.int32, device="cuda") num_ctx_tokens = seq_lens.sum().item() cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, cached_token_lens) + seq_lens, num_contexts, num_ctx_tokens, cached_token_lens + ) # Property 1: Output length matches num_ctx_tokens - assert cu_seqlen_ks.shape[0] == num_ctx_tokens, \ + assert cu_seqlen_ks.shape[0] == num_ctx_tokens, ( f"Trial {trial}: ks length mismatch: {cu_seqlen_ks.shape[0]} != {num_ctx_tokens}" - assert cu_seqlen_ke.shape[0] == num_ctx_tokens, \ + ) + assert cu_seqlen_ke.shape[0] == num_ctx_tokens, ( f"Trial {trial}: ke length mismatch: {cu_seqlen_ke.shape[0]} != {num_ctx_tokens}" + ) # Property 2: End > Start for all tokens - assert torch.all(cu_seqlen_ke > cu_seqlen_ks), \ + assert torch.all(cu_seqlen_ke > cu_seqlen_ks), ( f"Trial {trial}: End indices must be greater than start indices" + ) # Property 3: End indices are strictly increasing within each sequence offset = 0 for seq_idx in range(num_contexts): seq_len = seq_lens[seq_idx].item() - seq_ke = cu_seqlen_ke[offset:offset + seq_len] + seq_ke = cu_seqlen_ke[offset : offset + seq_len] # Each token should attend to one more KV position than previous if seq_len > 1: diffs = seq_ke[1:] - seq_ke[:-1] - assert torch.all(diffs == 1), \ + assert torch.all(diffs == 1), ( f"Trial {trial}: Seq {seq_idx} - consecutive ke diffs should be 1, got {diffs.tolist()}" + ) offset += seq_len @@ -1871,11 +1864,12 @@ def test_compute_cu_seqlen_bounds_with_cache_properties(): # First new Q token should attend to all cached + itself expected_window_size = cached_len + 1 - actual_window_size = (cu_seqlen_ke[offset] - - cu_seqlen_ks[offset]).item() + actual_window_size = (cu_seqlen_ke[offset] - cu_seqlen_ks[offset]).item() - assert actual_window_size == expected_window_size, \ - f"Trial {trial}: Seq {seq_idx} first token: expected window {expected_window_size}, got {actual_window_size}" + assert actual_window_size == expected_window_size, ( + f"Trial {trial}: Seq {seq_idx} first token: " + f"expected window {expected_window_size}, got {actual_window_size}" + ) offset += seq_len @@ -1884,8 +1878,9 @@ def test_compute_cu_seqlen_bounds_with_cache_properties(): expected_total_kv = kv_lens.sum().item() # Last token of last sequence should end at total KV count - assert cu_seqlen_ke[-1].item() == expected_total_kv, \ + assert cu_seqlen_ke[-1].item() == expected_total_kv, ( f"Trial {trial}: Last ke should equal total KV count: {cu_seqlen_ke[-1].item()} != {expected_total_kv}" + ) def test_compute_cu_seqlen_bounds_nocache_compressed_kv(): @@ -1898,7 +1893,8 @@ def test_compute_cu_seqlen_bounds_nocache_compressed_kv(): kv_lens = seq_lens // compress_ratio cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, None, kv_lens, compress_ratio) + seq_lens, num_contexts, num_ctx_tokens, None, kv_lens, compress_ratio + ) # Expected results: # Seq 0: tokens [0,1,2], KV [0] @@ -1911,17 +1907,15 @@ def test_compute_cu_seqlen_bounds_nocache_compressed_kv(): # Token 5: [1, 2) # Token 6: [1, 3) - expected_ks = torch.tensor([0, 0, 0, 1, 1, 1, 1], - dtype=torch.int32, - device="cuda") - expected_ke = torch.tensor([0, 1, 1, 1, 2, 2, 3], - dtype=torch.int32, - device="cuda") + expected_ks = torch.tensor([0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device="cuda") + expected_ke = torch.tensor([0, 1, 1, 1, 2, 2, 3], dtype=torch.int32, device="cuda") - assert torch.equal(cu_seqlen_ks, expected_ks), \ + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) def test_compute_cu_seqlen_bounds_with_cache_compressed_kv(): @@ -1943,8 +1937,8 @@ def test_compute_cu_seqlen_bounds_with_cache_compressed_kv(): num_ctx_tokens = 7 # 3 + 4 new tokens total cu_seqlen_ks, cu_seqlen_ke = compute_cu_seqlen_kv_bounds_with_cache( - seq_lens, num_contexts, num_ctx_tokens, num_past_tokens, kv_lens, - compress_ratio) + seq_lens, num_contexts, num_ctx_tokens, num_past_tokens, kv_lens, compress_ratio + ) # Expected results: # @@ -1959,48 +1953,52 @@ def test_compute_cu_seqlen_bounds_with_cache_compressed_kv(): # New Q token 2: [2, 4) # New Q token 3: [2, 4) - expected_ks = torch.tensor([0, 0, 0, 2, 2, 2, 2], - dtype=torch.int32, - device="cuda") - expected_ke = torch.tensor([1, 2, 2, 3, 3, 4, 4], - dtype=torch.int32, - device="cuda") + expected_ks = torch.tensor([0, 0, 0, 2, 2, 2, 2], dtype=torch.int32, device="cuda") + expected_ke = torch.tensor([1, 2, 2, 3, 3, 4, 4], dtype=torch.int32, device="cuda") - assert torch.equal(cu_seqlen_ks, expected_ks), \ + assert torch.equal(cu_seqlen_ks, expected_ks), ( f"cu_seqlen_ks mismatch:\nGot: {cu_seqlen_ks.tolist()}\nExpected: {expected_ks.tolist()}" - assert torch.equal(cu_seqlen_ke, expected_ke), \ + ) + assert torch.equal(cu_seqlen_ke, expected_ke), ( f"cu_seqlen_ke mismatch:\nGot: {cu_seqlen_ke.tolist()}\nExpected: {expected_ke.tolist()}" + ) @pytest.mark.parametrize( "max_chunk_size,seq_lens,start_idx,expected_specs", [ # Small requests - single chunk - (1000, [200, 300, 400], 0, [[(0, 0, 200, 0), (1, 0, 300, 200), - (2, 0, 400, 500)]]), + (1000, [200, 300, 400], 0, [[(0, 0, 200, 0), (1, 0, 300, 200), (2, 0, 400, 500)]]), # Small requests - multiple chunks - (1000, [400, 500, 300, 600], 0, [[(0, 0, 400, 0), (1, 0, 500, 400)], - [(2, 0, 300, 900), - (3, 0, 600, 1200)]]), + ( + 1000, + [400, 500, 300, 600], + 0, + [[(0, 0, 400, 0), (1, 0, 500, 400)], [(2, 0, 300, 900), (3, 0, 600, 1200)]], + ), # Large request - intra-request chunking - (1000, [2500], 0, [[(0, 0, 1000, 0)], [(0, 1000, 2000, 0)], - [(0, 2000, 2500, 0)]]), + (1000, [2500], 0, [[(0, 0, 1000, 0)], [(0, 1000, 2000, 0)], [(0, 2000, 2500, 0)]]), # Mixed: small + large + small - (1000, [400, 2500, 300], 0, [[ - (0, 0, 400, 0) - ], [(1, 0, 1000, 400)], [(1, 1000, 2000, 400)], [(1, 2000, 2500, 400)], - [(2, 0, 300, 2900)]]), + ( + 1000, + [400, 2500, 300], + 0, + [ + [(0, 0, 400, 0)], + [(1, 0, 1000, 400)], + [(1, 1000, 2000, 400)], + [(1, 2000, 2500, 400)], + [(2, 0, 300, 2900)], + ], + ), # Edge case: exact chunk size (1000, [1000], 0, [[(0, 0, 1000, 0)]]), # Edge case: non-zero start index (1000, [100, 200, 300], 1, [[(1, 0, 200, 100), (2, 0, 300, 300)]]), ], - ids=[ - "small_single", "small_multi", "large_chunked", "mixed", "exact_size", - "non_zero_start" - ]) -def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, - expected_specs): + ids=["small_single", "small_multi", "large_chunked", "mixed", "exact_size", "non_zero_start"], +) +def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, expected_specs): """ Test split_prefill_chunks covering: - Request-level chunking (small requests) @@ -2009,16 +2007,16 @@ def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, - Edge cases """ seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32) - chunk_groups = split_prefill_chunks(seq_lens_tensor, - max_chunk_size, - start_idx=start_idx) + chunk_groups = split_prefill_chunks(seq_lens_tensor, max_chunk_size, start_idx=start_idx) - assert len(chunk_groups) == len(expected_specs), \ + assert len(chunk_groups) == len(expected_specs), ( f"Expected {len(expected_specs)} chunks, got {len(chunk_groups)}" + ) for i, expected in enumerate(expected_specs): - assert chunk_groups[i] == expected, \ + assert chunk_groups[i] == expected, ( f"Chunk {i} mismatch:\nGot: {chunk_groups[i]}\nExpected: {expected}" + ) print("✅ test_split_prefill_chunks passed") @@ -2033,16 +2031,13 @@ def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, (512, None, "request_level"), # Random seq lens, 5 requests (1024, None, "request_level"), # Random seq lens, 4 requests (2048, None, "request_level"), # Random seq lens, 4 requests - # Two-level chunking: intra-request Q-block splitting for large requests (1024, [512, 2500, 800], "two_level"), # Middle request needs Q-blocks (1500, [3200, 1200], "two_level"), # First request needs 3 Q-blocks - (1000, [400, 600, 2200, 500 - ], "two_level"), # Mixed: request 2 needs Q-blocks + (1000, [400, 600, 2200, 500], "two_level"), # Mixed: request 2 needs Q-blocks ], ) -def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, - compress_ratio): +def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compress_ratio): """ Tests for indexer chunked prefill: 1. Request-level chunking: Multiple small requests packed together @@ -2073,14 +2068,13 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, batch_size = 5 if chunk_size == 512 else 4 seq_lens_list = [] for _ in range(batch_size): - seq_len = random.randint(int(0.3 * chunk_size), - int(0.9 * chunk_size)) + seq_len = random.randint(int(0.3 * chunk_size), int(0.9 * chunk_size)) seq_lens_list.append(seq_len) else: # Two-level chunking: use provided seq_lens with large requests batch_size = len(seq_lens_list) - seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device='cpu') + seq_lens = torch.tensor(seq_lens_list, dtype=torch.int32, device="cpu") total_tokens = seq_lens.sum().item() max_seq_len = seq_lens.max().item() @@ -2088,9 +2082,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, kv_lens_compressed = seq_lens // compress_ratio total_kv_tokens = kv_lens_compressed.sum().item() - print( - f"\n=== Test Config: {chunking_type}, compress_ratio={compress_ratio} ===" - ) + print(f"\n=== Test Config: {chunking_type}, compress_ratio={compress_ratio} ===") print(f" Batch: {batch_size}, Chunk size: {chunk_size}") print(f" Sequence lengths: {seq_lens_list}") print( @@ -2099,16 +2091,14 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, # Identify large requests for two-level chunking if chunking_type == "two_level": - large_requests = [(i, seq_len) - for i, seq_len in enumerate(seq_lens_list) - if seq_len > chunk_size] + large_requests = [ + (i, seq_len) for i, seq_len in enumerate(seq_lens_list) if seq_len > chunk_size + ] if large_requests: print(" Large requests (Q-block splitting):") for req_idx, seq_len in large_requests: num_q_blocks = (seq_len + chunk_size - 1) // chunk_size - print( - f" Request {req_idx}: {seq_len} tokens → {num_q_blocks} Q-blocks" - ) + print(f" Request {req_idx}: {seq_len} tokens → {num_q_blocks} Q-blocks") # Create cache manager and indexer cache_manager, sparse_attn_config = create_dsa_cache_manager( @@ -2116,30 +2106,22 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_model_len, - num_layers=1) + num_layers=1, + ) sparse_attn_config.index_topk = index_topk indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) # Allocate blocks for all sequences request_ids = list(range(batch_size)) - cache_manager.add_dummy_requests(request_ids=request_ids, - token_nums=seq_lens_list, - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=seq_lens_list, is_gen=False, prepare_resource=True + ) # Generate test data - q = torch.randn((total_tokens, heads, head_dim), - device="cuda", - dtype=torch.bfloat16) - k = torch.randn((total_kv_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - weights = torch.randn((total_tokens, heads), - device="cuda", - dtype=torch.float32) - hidden_states = torch.randn((total_tokens, 4096), - device="cuda", - dtype=torch.bfloat16) + q = torch.randn((total_tokens, heads, head_dim), device="cuda", dtype=torch.bfloat16) + k = torch.randn((total_kv_tokens, head_dim), device="cuda", dtype=torch.bfloat16) + weights = torch.randn((total_tokens, heads), device="cuda", dtype=torch.float32) + hidden_states = torch.randn((total_tokens, 4096), device="cuda", dtype=torch.bfloat16) # Quantize inputs q_fp8 = q.to(torch.float8_e4m3fn) @@ -2176,12 +2158,13 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, num_k = chunk.k_token_end - chunk.k_token_start print( f" Chunk {i}: Q[{chunk.token_start}:{chunk.token_end}] ({num_q} tokens), " - f"K[{chunk.k_token_start}:{chunk.k_token_end}] ({num_k} tokens)") + f"K[{chunk.k_token_start}:{chunk.k_token_end}] ({num_k} tokens)" + ) indexer._update_k_cache(k_fp8, k_scale, metadata_chunked) - topk_indices_chunked = indexer.sparse_attn_indexer(metadata_chunked, - hidden_states, q_fp8, - k_fp8, k_scale, weights) + topk_indices_chunked = indexer.sparse_attn_indexer( + metadata_chunked, hidden_states, q_fp8, k_fp8, k_scale, weights + ) print(f"✓ Chunked execution completed, shape: {topk_indices_chunked.shape}") @@ -2208,19 +2191,15 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, if metadata_baseline.indexer_prefill_chunks is not None: num_baseline_chunks = len(metadata_baseline.indexer_prefill_chunks) - print( - f"✓ Created {num_baseline_chunks} chunk(s) (effectively non-chunked)" - ) + print(f"✓ Created {num_baseline_chunks} chunk(s) (effectively non-chunked)") indexer._update_k_cache(k_fp8, k_scale, metadata_baseline) - topk_indices_baseline = indexer.sparse_attn_indexer(metadata_baseline, - hidden_states, q_fp8, - k_fp8, k_scale, weights) - - print( - f"✓ Non-chunked execution completed, shape: {topk_indices_baseline.shape}" + topk_indices_baseline = indexer.sparse_attn_indexer( + metadata_baseline, hidden_states, q_fp8, k_fp8, k_scale, weights ) + print(f"✓ Non-chunked execution completed, shape: {topk_indices_baseline.shape}") + # ========== Validation ========== print("\n=== Validation ===") @@ -2259,13 +2238,10 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, # Calculate statistics avg_similarity = total_similarity / total_tokens exact_match_ratio = num_exact_matches / total_tokens - high_similarity_ratio = (num_exact_matches + - num_high_similarity) / total_tokens + high_similarity_ratio = (num_exact_matches + num_high_similarity) / total_tokens print(" Results:") - print( - f" Exact matches: {num_exact_matches}/{total_tokens} ({exact_match_ratio:.1%})" - ) + print(f" Exact matches: {num_exact_matches}/{total_tokens} ({exact_match_ratio:.1%})") print(f" High similarity (>=95%): {num_high_similarity} additional") print(f" Overall high similarity ratio: {high_similarity_ratio:.1%}") print(f" Average Jaccard similarity: {avg_similarity:.4f}") @@ -2292,36 +2268,32 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, if similarity < 0.9: if low_sim_count < 5: # Show first 5 low similarity cases # Find which request this token belongs to - req_idx = (cumulative_lens - <= token_idx).sum().item() - 1 - local_token_idx = token_idx - cumulative_lens[ - req_idx].item() + req_idx = (cumulative_lens <= token_idx).sum().item() - 1 + local_token_idx = token_idx - cumulative_lens[req_idx].item() print( f" Token {token_idx} (req {req_idx}, local pos {local_token_idx}): " - f"similarity {similarity:.3f}") - print( - f" Chunked size: {len(chunked_set)}, Baseline size: {len(baseline_set)}" + f"similarity {similarity:.3f}" ) print( - f" Intersection: {intersection}, Union: {union}" + f" Chunked size: {len(chunked_set)}, Baseline size: {len(baseline_set)}" ) + print(f" Intersection: {intersection}, Union: {union}") low_sim_count += 1 if low_sim_count > 5: - print( - f" ... and {low_sim_count - 5} more tokens with low similarity" - ) + print(f" ... and {low_sim_count - 5} more tokens with low similarity") # Use Jaccard similarity threshold instead of exact match - assert avg_similarity >= 0.9, \ + assert avg_similarity >= 0.9, ( f"Chunked and non-chunked results differ significantly: avg similarity {avg_similarity:.4f} < 0.9" + ) + print(f"\n✅ Test passed! {chunking_type} chunking produces highly similar results") print( - f"\n✅ Test passed! {chunking_type} chunking produces highly similar results" + f" Config: chunk_size={chunk_size}, num_chunks={num_chunks}, " + f"batch={batch_size}, seq_lens={seq_lens_list}" ) - print(f" Config: chunk_size={chunk_size}, num_chunks={num_chunks}, " - f"batch={batch_size}, seq_lens={seq_lens_list}") @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -2330,8 +2302,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, @pytest.mark.parametrize("next_n", [1, 2, 4]) @pytest.mark.parametrize("index_topk", [2048]) @pytest.mark.parametrize("seq_len_range", [(2048, 8192), (512, 1024)]) -def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, - seq_len_range): +def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_len_range): """ Test that use_custom_topk=True and use_custom_topk=False produce identical results in the decode phase of sparse_attn_indexer. @@ -2361,9 +2332,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, # Generate KV cache lengths if enable_indexer_skip: - kv_lens = torch.randint(min_seq_len, - max_seq_len, (batch_size, ), - dtype=torch.int32) + kv_lens = torch.randint(min_seq_len, max_seq_len, (batch_size,), dtype=torch.int32) else: # (90% >= 2048 to test realistic scenarios) kv_lens = torch.zeros(batch_size, dtype=torch.int32) @@ -2373,24 +2342,21 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, if num_long > 0: long_min = max(2048, min_seq_len) long_max = max(long_min + 1, max_seq_len) - kv_lens[is_long] = torch.randint(long_min, - long_max, (num_long, ), - dtype=torch.int32) + kv_lens[is_long] = torch.randint(long_min, long_max, (num_long,), dtype=torch.int32) num_short = (~is_long).sum().item() if num_short > 0: short_max = min(2048, max_seq_len) if short_max > min_seq_len: - kv_lens[~is_long] = torch.randint(min_seq_len, - short_max, (num_short, ), - dtype=torch.int32) + kv_lens[~is_long] = torch.randint( + min_seq_len, short_max, (num_short,), dtype=torch.int32 + ) else: - kv_lens[~is_long] = torch.randint(max(2048, min_seq_len), - max(2049, max_seq_len), - (num_short, ), - dtype=torch.int32) + kv_lens[~is_long] = torch.randint( + max(2048, min_seq_len), max(2049, max_seq_len), (num_short,), dtype=torch.int32 + ) - seq_lens = torch.full((batch_size, ), next_n, dtype=torch.int32) + seq_lens = torch.full((batch_size,), next_n, dtype=torch.int32) num_gen_tokens = batch_size * next_n num_cached_tokens = kv_lens.tolist() @@ -2400,25 +2366,24 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_model_len, - num_layers=1) + num_layers=1, + ) sparse_attn_config.index_topk = index_topk indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) # Allocate blocks for all sequences (including historical + new tokens) request_ids = list(range(batch_size)) final_lens = kv_lens + next_n # Historical + new decode tokens - cache_manager.add_dummy_requests(request_ids=request_ids, - token_nums=final_lens.tolist(), - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=final_lens.tolist(), is_gen=False, prepare_resource=True + ) # Populate KV cache with historical context total_context_tokens = kv_lens.sum().item() - k_context_bf16 = torch.randn((total_context_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - k_context_fp8, k_context_scale = fp8_utils.fp8_quantize_1x128_sf_transpose( - k_context_bf16) + k_context_bf16 = torch.randn( + (total_context_tokens, head_dim), device="cuda", dtype=torch.bfloat16 + ) + k_context_fp8, k_context_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_context_bf16) metadata_context = _create_mock_metadata( request_ids=request_ids, @@ -2438,18 +2403,10 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, indexer._update_k_cache(k_context_fp8, k_context_scale, metadata_context) # Generate decode phase test data - q = torch.randn((num_gen_tokens, heads, head_dim), - device="cuda", - dtype=torch.bfloat16) - k_gen_bf16 = torch.randn((num_gen_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - weights = torch.randn((num_gen_tokens, heads), - device="cuda", - dtype=torch.float32) - hidden_states = torch.randn((num_gen_tokens, 4096), - device="cuda", - dtype=torch.bfloat16) + q = torch.randn((num_gen_tokens, heads, head_dim), device="cuda", dtype=torch.bfloat16) + k_gen_bf16 = torch.randn((num_gen_tokens, head_dim), device="cuda", dtype=torch.bfloat16) + weights = torch.randn((num_gen_tokens, heads), device="cuda", dtype=torch.float32) + hidden_states = torch.randn((num_gen_tokens, 4096), device="cuda", dtype=torch.bfloat16) q_fp8 = q.to(torch.float8_e4m3fn) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_gen_bf16) @@ -2472,112 +2429,106 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, indexer._update_k_cache(k_fp8, k_scale, metadata_gen_write) # Test with custom CUDA kernel - metadata_custom = _create_mock_metadata(request_ids, - batch_size, - 0, - batch_size, - seq_lens.clone(), - final_lens.clone(), - num_cached_tokens, - cache_manager, - 0, - num_gen_tokens, - max_model_len, - max_draft_tokens=next_n - 1, - indexer_head_dim=head_dim) + metadata_custom = _create_mock_metadata( + request_ids, + batch_size, + 0, + batch_size, + seq_lens.clone(), + final_lens.clone(), + num_cached_tokens, + cache_manager, + 0, + num_gen_tokens, + max_model_len, + max_draft_tokens=next_n - 1, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) try: - topk_indices_custom = indexer.sparse_attn_indexer(metadata_custom, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + topk_indices_custom = indexer.sparse_attn_indexer( + metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) except Exception as e: pytest.skip(f"Custom topk not available: {e}") # Test with PyTorch fallback - metadata_fallback = _create_mock_metadata(request_ids, - batch_size, - 0, - batch_size, - seq_lens.clone(), - final_lens.clone(), - num_cached_tokens, - cache_manager, - 0, - num_gen_tokens, - max_model_len, - max_draft_tokens=next_n - 1, - indexer_head_dim=head_dim) + metadata_fallback = _create_mock_metadata( + request_ids, + batch_size, + 0, + batch_size, + seq_lens.clone(), + final_lens.clone(), + num_cached_tokens, + cache_manager, + 0, + num_gen_tokens, + max_model_len, + max_draft_tokens=next_n - 1, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) - topk_indices_fallback = indexer.sparse_attn_indexer(metadata_fallback, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=False) + topk_indices_fallback = indexer.sparse_attn_indexer( + metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + ) # Test with indexer skip enabled if enable_indexer_skip: - metadata_skip = _create_mock_metadata(request_ids, - batch_size, - 0, - batch_size, - seq_lens.clone(), - final_lens.clone(), - num_cached_tokens, - cache_manager, - 0, - num_gen_tokens, - max_model_len, - max_draft_tokens=next_n - 1, - enable_indexer_skip=True, - indexer_head_dim=head_dim) + metadata_skip = _create_mock_metadata( + request_ids, + batch_size, + 0, + batch_size, + seq_lens.clone(), + final_lens.clone(), + num_cached_tokens, + cache_manager, + 0, + num_gen_tokens, + max_model_len, + max_draft_tokens=next_n - 1, + enable_indexer_skip=True, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) try: topk_indices_skip = indexer.sparse_attn_indexer( - metadata_skip, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) except Exception as e: raise RuntimeError(f"Error when testing indexer skip: {e}") # Validation ## Custom vs fallback num_ctx_tokens = 0 - custom_decode = topk_indices_custom[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :] - fallback_decode = topk_indices_fallback[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :] + custom_decode = topk_indices_custom[num_ctx_tokens : num_ctx_tokens + num_gen_tokens, :] + fallback_decode = topk_indices_fallback[num_ctx_tokens : num_ctx_tokens + num_gen_tokens, :] num_exact_matches, total_similarity, _ = validate_topk_indices( - custom_decode, fallback_decode, num_gen_tokens) + custom_decode, fallback_decode, num_gen_tokens + ) avg_similarity = total_similarity / num_gen_tokens - assert avg_similarity >= 0.95, \ + assert avg_similarity >= 0.95, ( f"Decode custom vs fallback differ: avg similarity {avg_similarity:.4f} < 0.95" + ) ## Custom vs skip if enable_indexer_skip: - skip_decode = topk_indices_skip[num_ctx_tokens:num_ctx_tokens + - num_gen_tokens, :] + skip_decode = topk_indices_skip[num_ctx_tokens : num_ctx_tokens + num_gen_tokens, :] num_exact_matches, total_similarity, _ = validate_topk_indices( - custom_decode, skip_decode, num_gen_tokens) + custom_decode, skip_decode, num_gen_tokens + ) avg_similarity = total_similarity / num_gen_tokens - assert avg_similarity >= 0.95, \ + assert avg_similarity >= 0.95, ( f"Decode custom vs skip differ: avg similarity {avg_similarity:.4f} < 0.95" + ) @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -2585,8 +2536,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, @pytest.mark.parametrize("batch_size", [4, 16]) @pytest.mark.parametrize("index_topk", [2048]) @pytest.mark.parametrize("chunk_size", [1024, 2048]) -def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, - chunk_size): +def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chunk_size): """ Test chunked prefill: use_custom_topk=True vs use_custom_topk=False with metadata.indexer_prefill_chunks != None. @@ -2609,9 +2559,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, # Generate variable sequence lengths to trigger chunking min_seq_len = chunk_size // 2 max_seq_len = chunk_size * 3 - seq_lens = torch.randint(min_seq_len, - max_seq_len, (batch_size, ), - dtype=torch.int32) + seq_lens = torch.randint(min_seq_len, max_seq_len, (batch_size,), dtype=torch.int32) total_tokens = seq_lens.sum().item() # Create cache manager and indexer @@ -2620,46 +2568,41 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_model_len, - num_layers=1) + num_layers=1, + ) sparse_attn_config.index_topk = index_topk indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) # Allocate cache blocks request_ids = list(range(batch_size)) - cache_manager.add_dummy_requests(request_ids=request_ids, - token_nums=seq_lens.tolist(), - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=seq_lens.tolist(), is_gen=False, prepare_resource=True + ) # Generate test data - q = torch.randn((total_tokens, heads, head_dim), - device="cuda", - dtype=torch.bfloat16) - k = torch.randn((total_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - weights = torch.randn((total_tokens, heads), - device="cuda", - dtype=torch.float32) - hidden_states = torch.randn((total_tokens, 4096), - device="cuda", - dtype=torch.bfloat16) + q = torch.randn((total_tokens, heads, head_dim), device="cuda", dtype=torch.bfloat16) + k = torch.randn((total_tokens, head_dim), device="cuda", dtype=torch.bfloat16) + weights = torch.randn((total_tokens, heads), device="cuda", dtype=torch.float32) + hidden_states = torch.randn((total_tokens, 4096), device="cuda", dtype=torch.bfloat16) q_fp8 = q.to(torch.float8_e4m3fn) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k) # Test with custom CUDA kernel - metadata_custom = _create_mock_metadata(request_ids, - batch_size, - batch_size, - 0, - seq_lens.clone(), - seq_lens.clone(), [0] * batch_size, - cache_manager, - total_tokens, - total_tokens, - chunk_size, - indexer_head_dim=head_dim) + metadata_custom = _create_mock_metadata( + request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), + seq_lens.clone(), + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + chunk_size, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) @@ -2667,46 +2610,42 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, assert metadata_custom.indexer_prefill_chunks is not None try: - topk_indices_custom = indexer.sparse_attn_indexer(metadata_custom, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + topk_indices_custom = indexer.sparse_attn_indexer( + metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) except Exception as e: pytest.skip(f"Custom topk not available: {e}") # Test with PyTorch fallback - metadata_fallback = _create_mock_metadata(request_ids, - batch_size, - batch_size, - 0, - seq_lens.clone(), - seq_lens.clone(), - [0] * batch_size, - cache_manager, - total_tokens, - total_tokens, - chunk_size, - indexer_head_dim=head_dim) + metadata_fallback = _create_mock_metadata( + request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), + seq_lens.clone(), + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + chunk_size, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) - topk_indices_fallback = indexer.sparse_attn_indexer(metadata_fallback, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=False) + topk_indices_fallback = indexer.sparse_attn_indexer( + metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + ) # Validation num_exact_matches, total_similarity, _ = validate_topk_indices( - topk_indices_custom, topk_indices_fallback, total_tokens) + topk_indices_custom, topk_indices_fallback, total_tokens + ) avg_similarity = total_similarity / total_tokens - assert avg_similarity >= 0.95, \ + assert avg_similarity >= 0.95, ( f"Chunked prefill differ: avg similarity {avg_similarity:.4f} < 0.95" + ) @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -2714,8 +2653,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, @pytest.mark.parametrize("batch_size", [4, 16]) @pytest.mark.parametrize("index_topk", [2048]) @pytest.mark.parametrize("seq_len_range", [(1, 512)]) -def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, - seq_len_range): +def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, seq_len_range): """ Test single-pass prefill: use_custom_topk=True vs use_custom_topk=False with metadata.indexer_prefill_chunks == None (else branch). @@ -2730,9 +2668,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, min_seq_len, max_seq_len = seq_len_range # Generate variable context lengths per sequence - seq_lens = torch.randint(min_seq_len, - max_seq_len, (batch_size, ), - dtype=torch.int32) + seq_lens = torch.randint(min_seq_len, max_seq_len, (batch_size,), dtype=torch.int32) total_tokens = seq_lens.sum().item() # Create cache manager and indexer @@ -2741,45 +2677,40 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_model_len, - num_layers=1) + num_layers=1, + ) sparse_attn_config.index_topk = index_topk indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) request_ids = list(range(batch_size)) - cache_manager.add_dummy_requests(request_ids=request_ids, - token_nums=seq_lens.tolist(), - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids=request_ids, token_nums=seq_lens.tolist(), is_gen=False, prepare_resource=True + ) # Generate test data - q = torch.randn((total_tokens, heads, head_dim), - device="cuda", - dtype=torch.bfloat16) - k = torch.randn((total_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - weights = torch.randn((total_tokens, heads), - device="cuda", - dtype=torch.float32) - hidden_states = torch.randn((total_tokens, 4096), - device="cuda", - dtype=torch.bfloat16) + q = torch.randn((total_tokens, heads, head_dim), device="cuda", dtype=torch.bfloat16) + k = torch.randn((total_tokens, head_dim), device="cuda", dtype=torch.bfloat16) + weights = torch.randn((total_tokens, heads), device="cuda", dtype=torch.float32) + hidden_states = torch.randn((total_tokens, 4096), device="cuda", dtype=torch.bfloat16) q_fp8 = q.to(torch.float8_e4m3fn) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k) # Test with custom CUDA kernel - metadata_custom = _create_mock_metadata(request_ids, - batch_size, - batch_size, - 0, - seq_lens.clone(), - seq_lens.clone(), [0] * batch_size, - cache_manager, - total_tokens, - total_tokens, - max_model_len, - indexer_head_dim=head_dim) + metadata_custom = _create_mock_metadata( + request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), + seq_lens.clone(), + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + max_model_len, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) @@ -2787,84 +2718,81 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, metadata_custom.indexer_prefill_chunks = None try: - topk_indices_custom = indexer.sparse_attn_indexer(metadata_custom, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + topk_indices_custom = indexer.sparse_attn_indexer( + metadata_custom, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) except Exception as e: pytest.skip(f"Custom topk not available: {e}") # Test with PyTorch fallback - metadata_fallback = _create_mock_metadata(request_ids, - batch_size, - batch_size, - 0, - seq_lens.clone(), - seq_lens.clone(), - [0] * batch_size, - cache_manager, - total_tokens, - total_tokens, - max_model_len, - indexer_head_dim=head_dim) + metadata_fallback = _create_mock_metadata( + request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), + seq_lens.clone(), + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + max_model_len, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) # Force single-pass path by setting indexer_prefill_chunks to None metadata_fallback.indexer_prefill_chunks = None - topk_indices_fallback = indexer.sparse_attn_indexer(metadata_fallback, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=False) + topk_indices_fallback = indexer.sparse_attn_indexer( + metadata_fallback, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + ) # Test with indexer skip enabled - metadata_skip = _create_mock_metadata(request_ids, - batch_size, - batch_size, - 0, - seq_lens.clone(), - seq_lens.clone(), [0] * batch_size, - cache_manager, - total_tokens, - total_tokens, - max_model_len, - enable_indexer_skip=True, - indexer_head_dim=head_dim) + metadata_skip = _create_mock_metadata( + request_ids, + batch_size, + batch_size, + 0, + seq_lens.clone(), + seq_lens.clone(), + [0] * batch_size, + cache_manager, + total_tokens, + total_tokens, + max_model_len, + enable_indexer_skip=True, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) metadata_skip.indexer_prefill_chunks = None try: - topk_indices_skip = indexer.sparse_attn_indexer(metadata_skip, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + topk_indices_skip = indexer.sparse_attn_indexer( + metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) except Exception as e: raise RuntimeError(f"Indexer skip not available: {e}") # Validation ## Custom vs fallback num_exact_matches, total_similarity, _ = validate_topk_indices( - topk_indices_custom, topk_indices_fallback, total_tokens) + topk_indices_custom, topk_indices_fallback, total_tokens + ) avg_similarity = total_similarity / total_tokens - assert avg_similarity >= 0.95, \ + assert avg_similarity >= 0.95, ( f"Single-pass prefill differ: avg similarity {avg_similarity:.4f} < 0.95" + ) ## Custom vs skip num_exact_matches, total_similarity, _ = validate_topk_indices( - topk_indices_custom, topk_indices_skip, total_tokens) + topk_indices_custom, topk_indices_skip, total_tokens + ) avg_similarity = total_similarity / total_tokens - assert avg_similarity >= 0.95, \ + assert avg_similarity >= 0.95, ( f"Single-pass prefill differ: avg similarity {avg_similarity:.4f} < 0.95" + ) @skip_pre_hopper @@ -2893,12 +2821,8 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): total_tokens = sum(seq_lens) print("\n=== Test: Multi-request with different cache ===") - print( - f" Req0: {cached_tokens[0]} cached + {seq_lens[0]} new = {total_kv_lens[0]} total" - ) - print( - f" Req1: {cached_tokens[1]} cached + {seq_lens[1]} new = {total_kv_lens[1]} total" - ) + print(f" Req0: {cached_tokens[0]} cached + {seq_lens[0]} new = {total_kv_lens[0]} total") + print(f" Req1: {cached_tokens[1]} cached + {seq_lens[1]} new = {total_kv_lens[1]} total") # Create cache manager cache_manager, sparse_attn_config = create_dsa_cache_manager( @@ -2906,70 +2830,55 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): head_dim=head_dim, tokens_per_block=block_size, max_seq_len=max_model_len, - num_layers=1) + num_layers=1, + ) sparse_attn_config.index_topk = index_topk indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx) # Allocate blocks request_ids = [0, 1] - cache_manager.add_dummy_requests(request_ids, - total_kv_lens, - is_gen=False, - prepare_resource=True) + cache_manager.add_dummy_requests( + request_ids, total_kv_lens, is_gen=False, prepare_resource=True + ) # Generate test data - q = torch.randn((total_tokens, heads, head_dim), - device="cuda", - dtype=torch.bfloat16) - k = torch.randn((total_tokens, head_dim), - device="cuda", - dtype=torch.bfloat16) - weights = torch.randn((total_tokens, heads), - device="cuda", - dtype=torch.float32) - hidden_states = torch.randn((total_tokens, 4096), - device="cuda", - dtype=torch.bfloat16) + q = torch.randn((total_tokens, heads, head_dim), device="cuda", dtype=torch.bfloat16) + k = torch.randn((total_tokens, head_dim), device="cuda", dtype=torch.bfloat16) + weights = torch.randn((total_tokens, heads), device="cuda", dtype=torch.float32) + hidden_states = torch.randn((total_tokens, 4096), device="cuda", dtype=torch.bfloat16) q_fp8 = q.to(torch.float8_e4m3fn) k_fp8, k_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k) # Create metadata with chunked prefill enabled - metadata = _create_mock_metadata(request_ids, - batch_size, - batch_size, - 0, - torch.tensor(seq_lens, dtype=torch.int32), - torch.tensor(total_kv_lens, - dtype=torch.int32), - cached_tokens, - cache_manager, - total_tokens, - total_tokens, - indexer_max_chunk_size=32768, - enable_context_mla_with_cached_kv=True, - indexer_head_dim=head_dim) + metadata = _create_mock_metadata( + request_ids, + batch_size, + batch_size, + 0, + torch.tensor(seq_lens, dtype=torch.int32), + torch.tensor(total_kv_lens, dtype=torch.int32), + cached_tokens, + cache_manager, + total_tokens, + total_tokens, + indexer_max_chunk_size=32768, + enable_context_mla_with_cached_kv=True, + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata) indexer._update_k_cache(k_fp8, k_scale, metadata) # Test custom kernel - topk_custom = indexer.sparse_attn_indexer(metadata, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + topk_custom = indexer.sparse_attn_indexer( + metadata, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) # Test fallback - topk_fallback = indexer.sparse_attn_indexer(metadata, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=False) + topk_fallback = indexer.sparse_attn_indexer( + metadata, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=False + ) # Test with indexer skip enabled if enable_indexer_skip: @@ -2987,16 +2896,13 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): indexer_max_chunk_size=32768, enable_context_mla_with_cached_kv=True, enable_indexer_skip=True, - indexer_head_dim=head_dim) + indexer_head_dim=head_dim, + ) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) - topk_indices_skip = indexer.sparse_attn_indexer(metadata_skip, - hidden_states, - q_fp8, - k_fp8, - k_scale, - weights, - use_custom_topk=True) + topk_indices_skip = indexer.sparse_attn_indexer( + metadata_skip, hidden_states, q_fp8, k_fp8, k_scale, weights, use_custom_topk=True + ) # Validate: custom and fallback should match print("\n=== Validation ===") @@ -3010,60 +2916,59 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): has_invalid = min_val_custom < -1 if has_invalid or num_valid_custom != num_valid_fallback: - print( - f" Token {tok_id}: custom={num_valid_custom}, fallback={num_valid_fallback}" - ) - print( - f" Custom min={min_val_custom}, Fallback min={min_val_fallback}" - ) + print(f" Token {tok_id}: custom={num_valid_custom}, fallback={num_valid_fallback}") + print(f" Custom min={min_val_custom}, Fallback min={min_val_fallback}") if has_invalid: - print( - " ⚠️ INVALID: Custom has negative indices < -1 (kernel bug!)" - ) + print(" ⚠️ INVALID: Custom has negative indices < -1 (kernel bug!)") # Check tokens with large windows (>= 2048) should have exactly 2048 valid indices print("\n=== Check: Large windows must have 2048 valid ===") from tensorrt_llm._torch.attention_backend.sparse.dsa import ( compute_cu_seqlen_kv_bounds_with_cache, ) - host_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device='cpu') - host_cached = torch.tensor(cached_tokens, dtype=torch.int32, device='cpu') + host_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device="cpu") + host_cached = torch.tensor(cached_tokens, dtype=torch.int32, device="cpu") cu_ks, cu_ke = compute_cu_seqlen_kv_bounds_with_cache( - host_seq_lens, batch_size, total_tokens, host_cached) + host_seq_lens, batch_size, total_tokens, host_cached + ) for tok_id in range(total_tokens): window_size = (cu_ke[tok_id] - cu_ks[tok_id]).item() if window_size >= index_topk: num_valid = (topk_custom[tok_id] >= 0).sum().item() num_valid_fallback = (topk_fallback[tok_id] >= 0).sum().item() - assert num_valid_fallback == index_topk, \ - f"[Fallback] Token {tok_id}: window={window_size}, but only {num_valid_fallback}/{index_topk} valid indices" - assert num_valid == index_topk, \ + assert num_valid_fallback == index_topk, ( + f"[Fallback] Token {tok_id}: window={window_size}, " + f"but only {num_valid_fallback}/{index_topk} valid indices" + ) + assert num_valid == index_topk, ( f"[Custom Topk] Token {tok_id}: window={window_size}, but only {num_valid}/{index_topk} valid indices" + ) print(f" ✓ All large-window tokens have {index_topk} valid indices") # Validation num_exact_matches, total_similarity, min_similarity = validate_topk_indices( - topk_custom, topk_fallback, total_tokens) + topk_custom, topk_fallback, total_tokens + ) avg_similarity = total_similarity / total_tokens print(f" Exact matches: {num_exact_matches}/{total_tokens}") - print( - f" Similarity - Min: {min_similarity:.4f}, Avg: {avg_similarity:.4f}") + print(f" Similarity - Min: {min_similarity:.4f}, Avg: {avg_similarity:.4f}") - assert avg_similarity >= 0.95, \ + assert avg_similarity >= 0.95, ( f"Custom vs fallback differ: avg similarity {avg_similarity:.4f} < 0.95" + ) if enable_indexer_skip: num_exact_matches, total_similarity, min_similarity = validate_topk_indices( - topk_custom, topk_indices_skip, total_tokens) + topk_custom, topk_indices_skip, total_tokens + ) avg_similarity = total_similarity / total_tokens print(f" Exact matches: {num_exact_matches}/{total_tokens}") - print( - f" Similarity - Min: {min_similarity:.4f}, Avg: {avg_similarity:.4f}" - ) - assert avg_similarity >= 0.95, \ + print(f" Similarity - Min: {min_similarity:.4f}, Avg: {avg_similarity:.4f}") + assert avg_similarity >= 0.95, ( f"Custom vs indexer skip differ: avg similarity {avg_similarity:.4f} < 0.95" + ) class TestPrepareRestoreAttnMetadataForDraftReplay: @@ -3097,22 +3002,25 @@ def test_prepare_swaps_and_restore_recovers(self): original_offsets = meta.kv_cache_block_offsets.clone() original_host_offsets = meta.host_kv_cache_block_offsets.clone() - with patch('tensorrt_llm._torch.speculative.interface.isinstance', - side_effect=lambda obj, cls: - (obj is meta if cls is TrtllmAttentionMetadata else False - if cls.__name__ == 'DSAtrtllmAttentionMetadata' else - builtins.isinstance(obj, cls))): + with patch( + "tensorrt_llm._torch.speculative.interface.isinstance", + side_effect=lambda obj, cls: ( + obj is meta + if cls is TrtllmAttentionMetadata + else False + if cls.__name__ == "DSAtrtllmAttentionMetadata" + else builtins.isinstance(obj, cls) + ), + ): saved = prepare_attn_metadata_for_draft_replay(meta, mgr) assert saved is not None - assert saved['target_kv_cache_manager'] is original_kv_mgr + assert saved["target_kv_cache_manager"] is original_kv_mgr assert meta.kv_cache_manager is mgr - assert 'saved_dsa_state' not in saved + assert "saved_dsa_state" not in saved restore_attn_metadata_after_draft_replay(meta, saved) assert meta.kv_cache_manager is original_kv_mgr - torch.testing.assert_close(meta.kv_cache_block_offsets, - original_offsets) - torch.testing.assert_close(meta.host_kv_cache_block_offsets, - original_host_offsets) + torch.testing.assert_close(meta.kv_cache_block_offsets, original_offsets) + torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets) diff --git a/tests/unittest/_torch/attention/sparse/kernel/test_flash_mla.py b/tests/unittest/_torch/attention/sparse/kernel/test_flash_mla.py index 3caabefa0250..a9686877b2b9 100644 --- a/tests/unittest/_torch/attention/sparse/kernel/test_flash_mla.py +++ b/tests/unittest/_torch/attention/sparse/kernel/test_flash_mla.py @@ -13,6 +13,7 @@ def has_flash_mla(): """Check if FlashMLA module is available.""" try: from tensorrt_llm.flash_mla import flash_mla_sparse_fwd # noqa: F401 + return True except ImportError: return False @@ -20,15 +21,16 @@ def has_flash_mla(): @pytest.mark.skipif(not has_flash_mla(), reason="FlashMLA not available") @pytest.mark.skipif( - getSMVersion() < 90, - reason="FlashMLA requires SM90 (Hopper) or SM100 (Blackwell)") + getSMVersion() < 90, reason="FlashMLA requires SM90 (Hopper) or SM100 (Blackwell)" +) @pytest.mark.parametrize( "seq_len_q,seq_len_kv,topk", [ (62, 128, 128), # Small test case (128, 256, 128), # Medium (128, 512, 256), # Larger topk - ]) + ], +) def test_flash_mla_sparse_fwd(seq_len_q, seq_len_kv, topk): """ Test FlashMLA sparse attention forward kernel. @@ -53,29 +55,27 @@ def test_flash_mla_sparse_fwd(seq_len_q, seq_len_kv, topk): # Generate test inputs # Q: [b, s_q, h_q, d_qk] - q = torch.randn(batch_size, - seq_len_q, - num_heads_q, - head_dim_qk, - dtype=torch.bfloat16, - device='cuda') / 10.0 + q = ( + torch.randn( + batch_size, seq_len_q, num_heads_q, head_dim_qk, dtype=torch.bfloat16, device="cuda" + ) + / 10.0 + ) q.clamp_(-10, 10) # KV: [b, s_kv, h_kv, d_qk] - kv = torch.randn(batch_size, - seq_len_kv, - num_heads_kv, - head_dim_qk, - dtype=torch.bfloat16, - device='cuda') / 10.0 + kv = ( + torch.randn( + batch_size, seq_len_kv, num_heads_kv, head_dim_qk, dtype=torch.bfloat16, device="cuda" + ) + / 10.0 + ) kv.clamp_(-10, 10) # Indices: [b, s_q, h_kv, topk] - which KV tokens each Q attends to - indices = torch.randint(0, - seq_len_kv, - (batch_size, seq_len_q, num_heads_kv, topk), - dtype=torch.int32, - device='cuda') + indices = torch.randint( + 0, seq_len_kv, (batch_size, seq_len_q, num_heads_kv, topk), dtype=torch.int32, device="cuda" + ) softmax_scale = 1.0 / math.sqrt(head_dim_qk) @@ -84,23 +84,24 @@ def test_flash_mla_sparse_fwd(seq_len_q, seq_len_kv, topk): q.squeeze(0), # [s_q, h_q, d_qk] kv.squeeze(0), # [s_kv, h_kv, d_qk] indices.squeeze(0), # [s_q, h_kv, topk] - sm_scale=softmax_scale) + sm_scale=softmax_scale, + ) # Validate outputs - assert output.shape == (seq_len_q, num_heads_q, head_dim_v), \ + assert output.shape == (seq_len_q, num_heads_q, head_dim_v), ( f"Output shape mismatch: expected [{seq_len_q}, {num_heads_q}, {head_dim_v}], got {output.shape}" - assert output.dtype == torch.bfloat16, \ + ) + assert output.dtype == torch.bfloat16, ( f"Output dtype mismatch: expected torch.bfloat16, got {output.dtype}" + ) - assert max_logits.shape == (seq_len_q, num_heads_q), \ + assert max_logits.shape == (seq_len_q, num_heads_q), ( f"Max logits shape mismatch: got {max_logits.shape}" - assert max_logits.dtype == torch.float32, \ - f"Max logits dtype mismatch: got {max_logits.dtype}" + ) + assert max_logits.dtype == torch.float32, f"Max logits dtype mismatch: got {max_logits.dtype}" - assert lse.shape == (seq_len_q, num_heads_q), \ - f"LSE shape mismatch: got {lse.shape}" - assert lse.dtype == torch.float32, \ - f"LSE dtype mismatch: got {lse.dtype}" + assert lse.shape == (seq_len_q, num_heads_q), f"LSE shape mismatch: got {lse.shape}" + assert lse.dtype == torch.float32, f"LSE dtype mismatch: got {lse.dtype}" # Numerical validity checks assert not torch.isnan(output).any(), "Output contains NaN" diff --git a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py index f54ec8e0b80a..81bf2114fd35 100644 --- a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py @@ -10,28 +10,28 @@ from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention_backend.interface import AttentionForwardArgs from tensorrt_llm._torch.attention_backend.sparse.rocket import ( - RocketKVCacheManager, RocketTrtllmAttention, RocketTrtllmAttentionMetadata, - RocketVanillaAttention, RocketVanillaAttentionMetadata) + RocketKVCacheManager, + RocketTrtllmAttention, + RocketTrtllmAttentionMetadata, + RocketVanillaAttention, + RocketVanillaAttentionMetadata, +) from tensorrt_llm._torch.metadata import KVCacheParams -from tensorrt_llm.llmapi import (CudaGraphConfig, KvCacheConfig, - RocketSparseAttentionConfig) +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, RocketSparseAttentionConfig from tensorrt_llm.mapping import Mapping -@pytest.mark.skipif(getSMVersion() < 100, - reason="RocketKV requires SM100 (Blackwell)") +@pytest.mark.skipif(getSMVersion() < 100, reason="RocketKV requires SM100 (Blackwell)") @pytest.mark.parametrize("backend", ["pytorch"]) -@pytest.mark.parametrize("model_name", - ["llama-3.1-model/Llama-3.1-8B-Instruct"]) +@pytest.mark.parametrize("model_name", ["llama-3.1-model/Llama-3.1-8B-Instruct"]) @pytest.mark.parametrize("attention_backend", ["VANILLA", "TRTLLM"]) def test_model(backend, model_name, attention_backend): model_dir = str(llm_models_root() / model_name) max_batch_size = 16 max_output_tokens = 128 - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, - enable_block_reuse=False) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, enable_block_reuse=False) - kt_cache_dtype = 'float8_e5m2' if attention_backend == "TRTLLM" else 'bfloat16' + kt_cache_dtype = "float8_e5m2" if attention_backend == "TRTLLM" else "bfloat16" sparse_attention_config = RocketSparseAttentionConfig( window_size=32, @@ -54,32 +54,30 @@ def test_model(backend, model_name, attention_backend): max_batch_size=max_batch_size, max_seq_len=20480, max_num_tokens=81920, - cuda_graph_config=None - if attention_backend == "VANILLA" else cuda_graph_config, + cuda_graph_config=None if attention_backend == "VANILLA" else cuda_graph_config, ) inputs, references = [], [] current_file = os.path.abspath(__file__) - current_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) - input_file = f'{current_dir}/multi_gpu/NIAH_simple_data.jsonl' - with open(input_file, 'r') as f: + current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) + input_file = f"{current_dir}/multi_gpu/NIAH_simple_data.jsonl" + with open(input_file, "r") as f: for line in f: sample = json.loads(line) - inputs.append({ - 'prompt': - sample['input_context'] + sample['input_query'], - }) - references.append(sample['outputs'][0]) + inputs.append( + { + "prompt": sample["input_context"] + sample["input_query"], + } + ) + references.append(sample["outputs"][0]) with llm: outputs = llm.generate( inputs, use_tqdm=True, - sampling_params=SamplingParams(add_special_tokens=False, - max_tokens=max_output_tokens, - temperature=0.8, - top_p=0.95), + sampling_params=SamplingParams( + add_special_tokens=False, max_tokens=max_output_tokens, temperature=0.8, top_p=0.95 + ), ) count = 0 @@ -87,21 +85,29 @@ def test_model(backend, model_name, attention_backend): print(f"ret: {ret.outputs[0].text}") print(f"ref: {ref}") if ref not in ret.outputs[0].text: - print(f'reference {ref} is not in the output {ret.outputs[0].text}') + print(f"reference {ref} is not in the output {ret.outputs[0].text}") else: count = count + 1 acc = count / len(outputs) - assert acc >= 0.9, 'accuracy test of rocketkv sparse attention failed' + assert acc >= 0.9, "accuracy test of rocketkv sparse attention failed" -def create_rocket_kv_cache_manager(num_layers, num_kv_heads, head_dim, - tokens_per_block, max_seq_len, - max_batch_size, dtype, sparse_attn_config): +def create_rocket_kv_cache_manager( + num_layers, + num_kv_heads, + head_dim, + tokens_per_block, + max_seq_len, + max_batch_size, + dtype, + sparse_attn_config, +): mapping = Mapping(world_size=1, tp_size=1, rank=0) num_blocks = 100 - kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block, - enable_block_reuse=False) + kv_cache_config = KvCacheConfig( + max_tokens=num_blocks * tokens_per_block, enable_block_reuse=False + ) kv_cache_manager = RocketKVCacheManager( kv_cache_config, @@ -119,8 +125,15 @@ def create_rocket_kv_cache_manager(num_layers, num_kv_heads, head_dim, return kv_cache_manager -def create_test_metadata(seq_lens, num_contexts, past_seen_tokens, request_ids, - kv_cache_manager, sparse_attn_config, metadata_cls): +def create_test_metadata( + seq_lens, + num_contexts, + past_seen_tokens, + request_ids, + kv_cache_manager, + sparse_attn_config, + metadata_cls, +): prompt_lens = [] for i, (seq_len, past_token) in enumerate(zip(seq_lens, past_seen_tokens)): if i < num_contexts: @@ -131,8 +144,7 @@ def create_test_metadata(seq_lens, num_contexts, past_seen_tokens, request_ids, metadata = metadata_cls( seq_lens=torch.tensor(seq_lens, dtype=torch.int), num_contexts=num_contexts, - kv_cache_params=KVCacheParams( - use_cache=True, num_cached_tokens_per_seq=past_seen_tokens), + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=past_seen_tokens), max_num_requests=len(seq_lens), max_num_sequences=len(seq_lens), max_num_tokens=8192, @@ -151,7 +163,8 @@ def create_test_metadata(seq_lens, num_contexts, past_seen_tokens, request_ids, (1, 1), # bs=1 (4, 4), # bs=2, context only (2 contexts) (6, 3), # bs=6, mixed (3 contexts + 3 generations) - ]) + ], +) def test_sparse_kv_predict(batch_size, num_contexts): """ Test sparse_kv_predict against vanilla _get_snapkv_indices. @@ -164,7 +177,7 @@ def test_sparse_kv_predict(batch_size, num_contexts): num_heads = 32 num_kv_heads = 8 head_dim = 128 - device = torch.device('cuda') + device = torch.device("cuda") dtype = torch.bfloat16 sparse_attn_config = RocketSparseAttentionConfig( @@ -172,7 +185,7 @@ def test_sparse_kv_predict(batch_size, num_contexts): kernel_size=3, prompt_budget=256, page_size=4, - kt_cache_dtype='bfloat16', + kt_cache_dtype="bfloat16", ) # Create sequence lengths - mix short and long sequences in context phase @@ -184,20 +197,26 @@ def test_sparse_kv_predict(batch_size, num_contexts): if i % 2 == 1 and batch_size > 1: # Short sequence: seq_len < prompt_budget seq_lens.append( - torch.randint(sparse_attn_config.prompt_budget // 2, - sparse_attn_config.prompt_budget - 10, - (1, )).item()) + torch.randint( + sparse_attn_config.prompt_budget // 2, + sparse_attn_config.prompt_budget - 10, + (1,), + ).item() + ) else: # Long sequence: seq_len > prompt_budget seq_lens.append( - torch.randint(sparse_attn_config.prompt_budget, - sparse_attn_config.prompt_budget + 200, - (1, )).item()) + torch.randint( + sparse_attn_config.prompt_budget, + sparse_attn_config.prompt_budget + 200, + (1,), + ).item() + ) past_seen_tokens.append(0) else: # Generation phase: single token seq_lens.append(1) - past_seen_tokens.append(torch.randint(100, 200, (1, )).item()) + past_seen_tokens.append(torch.randint(100, 200, (1,)).item()) request_ids = list(range(batch_size)) @@ -214,17 +233,28 @@ def test_sparse_kv_predict(batch_size, num_contexts): vanilla_tokens_per_block = max_seq_len # Each sequence in one block trtllm_kv_cache_manager = create_rocket_kv_cache_manager( - num_layers, num_kv_heads, head_dim, tokens_per_block, max_seq_len, - batch_size, kv_cache_dtype, sparse_attn_config) + num_layers, + num_kv_heads, + head_dim, + tokens_per_block, + max_seq_len, + batch_size, + kv_cache_dtype, + sparse_attn_config, + ) vanilla_kv_cache_manager = create_rocket_kv_cache_manager( - num_layers, num_kv_heads, head_dim, vanilla_tokens_per_block, - max_seq_len, batch_size, kv_cache_dtype, sparse_attn_config) + num_layers, + num_kv_heads, + head_dim, + vanilla_tokens_per_block, + max_seq_len, + batch_size, + kv_cache_dtype, + sparse_attn_config, + ) # Add dummy requests to both cache managers - token_nums = [ - seq_len + past_token - for seq_len, past_token in zip(seq_lens, past_seen_tokens) - ] + token_nums = [seq_len + past_token for seq_len, past_token in zip(seq_lens, past_seen_tokens)] trtllm_kv_cache_manager.add_dummy_requests(request_ids, token_nums) vanilla_kv_cache_manager.add_dummy_requests(request_ids, token_nums) @@ -244,21 +274,31 @@ def test_sparse_kv_predict(batch_size, num_contexts): sparse_attention_config=sparse_attn_config, ) - trtllm_metadata = create_test_metadata(seq_lens, num_contexts, - past_seen_tokens, request_ids, - trtllm_kv_cache_manager, - sparse_attn_config, - RocketTrtllmAttentionMetadata) - vanilla_metadata = create_test_metadata(seq_lens, num_contexts, - past_seen_tokens, request_ids, - vanilla_kv_cache_manager, - sparse_attn_config, - RocketVanillaAttentionMetadata) + trtllm_metadata = create_test_metadata( + seq_lens, + num_contexts, + past_seen_tokens, + request_ids, + trtllm_kv_cache_manager, + sparse_attn_config, + RocketTrtllmAttentionMetadata, + ) + _ = create_test_metadata( + seq_lens, + num_contexts, + past_seen_tokens, + request_ids, + vanilla_kv_cache_manager, + sparse_attn_config, + RocketVanillaAttentionMetadata, + ) total_tokens = sum(seq_lens) - qkv = torch.randn(total_tokens, (num_heads + 2 * num_kv_heads) * head_dim, - dtype=dtype, - device=device) + qkv = torch.randn( + total_tokens, + (num_heads + 2 * num_kv_heads) * head_dim, + dtype=dtype, + device=device) forward_args = AttentionForwardArgs() trtllm_sparse_kv_indices, trtllm_sparse_kv_offsets = trtllm_attn.sparse_kv_predict( @@ -268,35 +308,32 @@ def test_sparse_kv_predict(batch_size, num_contexts): offset = 0 for i in range(num_contexts): seq_len = seq_lens[i] - single_qkv = qkv[offset:offset + seq_len] - q, k, _ = single_qkv.split([ - num_heads * head_dim, num_kv_heads * head_dim, - num_kv_heads * head_dim - ], - dim=-1) + single_qkv = qkv[offset : offset + seq_len] + q, k, _ = single_qkv.split( + [num_heads * head_dim, num_kv_heads * head_dim, num_kv_heads * head_dim], dim=-1 + ) q = q.view(1, seq_len, num_heads, head_dim).transpose(1, 2) k = k.view(1, seq_len, num_kv_heads, head_dim) if seq_len <= sparse_attn_config.prompt_budget: # Short sequences: vanilla returns None, but trtllm returns [0, 1, ..., seq_len-1] # Generate expected indices for comparison - short_indices = torch.arange(seq_len, - device=device, - dtype=torch.int32).unsqueeze(0).expand( - num_kv_heads, -1) + short_indices = ( + torch.arange(seq_len, device=device, dtype=torch.int32) + .unsqueeze(0) + .expand(num_kv_heads, -1) + ) vanilla_sparse_kv_indices_list.append(short_indices) else: vanilla_indices = vanilla_attn._get_snapkv_indices(q, k, i) if vanilla_indices is not None: - vanilla_indices = vanilla_indices.squeeze(0).transpose( - 0, 1).contiguous() + vanilla_indices = vanilla_indices.squeeze(0).transpose(0, 1).contiguous() vanilla_sparse_kv_indices_list.append(vanilla_indices) offset += seq_len if len(vanilla_sparse_kv_indices_list) > 0: - vanilla_sparse_kv_indices = torch.cat(vanilla_sparse_kv_indices_list, - dim=-1).contiguous() + vanilla_sparse_kv_indices = torch.cat(vanilla_sparse_kv_indices_list, dim=-1).contiguous() else: vanilla_sparse_kv_indices = None @@ -304,8 +341,9 @@ def test_sparse_kv_predict(batch_size, num_contexts): if trtllm_sparse_kv_indices is not None: assert vanilla_sparse_kv_indices is not None, "Vanilla should also produce indices" - assert trtllm_sparse_kv_indices.shape == vanilla_sparse_kv_indices.shape, \ + assert trtllm_sparse_kv_indices.shape == vanilla_sparse_kv_indices.shape, ( f"Shape mismatch: {trtllm_sparse_kv_indices.shape} vs {vanilla_sparse_kv_indices.shape}" + ) # Check indices overlap per batch and per head num_kv_heads = trtllm_sparse_kv_indices.shape[0] @@ -323,33 +361,34 @@ def test_sparse_kv_predict(batch_size, num_contexts): batch_overlaps = [] for head_idx in range(num_kv_heads): - trtllm_batch = trtllm_sparse_kv_indices[ - head_idx, start_idx:end_idx].cpu().tolist() - vanilla_batch = vanilla_sparse_kv_indices[ - head_idx, start_idx:end_idx].cpu().tolist() + trtllm_batch = trtllm_sparse_kv_indices[head_idx, start_idx:end_idx].cpu().tolist() + vanilla_batch = ( + vanilla_sparse_kv_indices[head_idx, start_idx:end_idx].cpu().tolist() + ) trtllm_set = set(trtllm_batch) vanilla_set = set(vanilla_batch) # Calculate overlap overlap = len(vanilla_set & trtllm_set) - overlap_ratio = overlap / len(vanilla_set) if len( - vanilla_set) > 0 else 1.0 + overlap_ratio = overlap / len(vanilla_set) if len(vanilla_set) > 0 else 1.0 batch_overlaps.append(overlap_ratio) overlap_ratios.append(overlap_ratio) avg_batch_overlap = sum(batch_overlaps) / len(batch_overlaps) - batch_overlap_details.append( - f"Batch {batch_idx}: {avg_batch_overlap:.4f}") + batch_overlap_details.append(f"Batch {batch_idx}: {avg_batch_overlap:.4f}") avg_overlap_ratio = sum(overlap_ratios) / len(overlap_ratios) print(f"Average overlap ratio: {avg_overlap_ratio:.4f}") print(f"Per-batch average: {batch_overlap_details}") - assert avg_overlap_ratio >= 0.98, \ + assert avg_overlap_ratio >= 0.98, ( f"Indices overlap ratio {avg_overlap_ratio:.4f} is too low (< 0.98)" + ) else: - assert vanilla_sparse_kv_indices is None, "Both should return None when no sparse attention is needed" + assert vanilla_sparse_kv_indices is None, ( + "Both should return None when no sparse attention is needed" + ) @pytest.mark.parametrize( @@ -360,7 +399,8 @@ def test_sparse_kv_predict(batch_size, num_contexts): (3, 0), # bs=3, generation only (3 generations) (5, 3), # bs=5, mixed (3 contexts + 2 generations) (6, 2), # bs=6, mixed (2 ctx + 4 gen) - ]) + ], +) def test_sparse_attn_predict(batch_size, num_contexts): """Test sparse_attn_predict against vanilla _rocketkv_selection.""" num_generations = batch_size - num_contexts @@ -369,7 +409,7 @@ def test_sparse_attn_predict(batch_size, num_contexts): num_heads = 32 num_kv_heads = 8 head_dim = 128 - device = torch.device('cuda') + device = torch.device("cuda") dtype = torch.bfloat16 sparse_attn_config_vanilla = RocketSparseAttentionConfig( @@ -379,7 +419,7 @@ def test_sparse_attn_predict(batch_size, num_contexts): page_size=2, topk=128, topr=96, - kt_cache_dtype='bfloat16', + kt_cache_dtype="bfloat16", ) sparse_attn_config_trtllm = RocketSparseAttentionConfig( window_size=32, @@ -388,7 +428,7 @@ def test_sparse_attn_predict(batch_size, num_contexts): page_size=2, topk=64, topr=96, - kt_cache_dtype='bfloat16', + kt_cache_dtype="bfloat16", ) # Create sequence lengths @@ -397,13 +437,13 @@ def test_sparse_attn_predict(batch_size, num_contexts): for i in range(batch_size): if i < num_contexts: # Context phase: longer sequences - seq_lens.append(torch.randint(300, 400, (1, )).item()) + seq_lens.append(torch.randint(300, 400, (1,)).item()) past_seen_tokens.append(0) else: # Generation phase: single token seq_lens.append(1) # 128 is the minimum number of tokens for shape alignment - past_seen_tokens.append(torch.randint(128, 300, (1, )).item()) + past_seen_tokens.append(torch.randint(128, 300, (1,)).item()) request_ids = list(range(batch_size)) @@ -420,17 +460,28 @@ def test_sparse_attn_predict(batch_size, num_contexts): vanilla_tokens_per_block = max_seq_len # Each sequence in one block trtllm_kv_cache_manager = create_rocket_kv_cache_manager( - num_layers, num_kv_heads, head_dim, tokens_per_block, max_seq_len, - batch_size, kv_cache_dtype, sparse_attn_config_trtllm) + num_layers, + num_kv_heads, + head_dim, + tokens_per_block, + max_seq_len, + batch_size, + kv_cache_dtype, + sparse_attn_config_trtllm, + ) vanilla_kv_cache_manager = create_rocket_kv_cache_manager( - num_layers, num_kv_heads, head_dim, vanilla_tokens_per_block, - max_seq_len, batch_size, kv_cache_dtype, sparse_attn_config_vanilla) + num_layers, + num_kv_heads, + head_dim, + vanilla_tokens_per_block, + max_seq_len, + batch_size, + kv_cache_dtype, + sparse_attn_config_vanilla, + ) # Add dummy requests to both cache managers - token_nums = [ - seq_len + past_token - for seq_len, past_token in zip(seq_lens, past_seen_tokens) - ] + token_nums = [seq_len + past_token for seq_len, past_token in zip(seq_lens, past_seen_tokens)] trtllm_kv_cache_manager.add_dummy_requests(request_ids, token_nums) vanilla_kv_cache_manager.add_dummy_requests(request_ids, token_nums) @@ -450,21 +501,29 @@ def test_sparse_attn_predict(batch_size, num_contexts): sparse_attention_config=sparse_attn_config_vanilla, ) - trtllm_metadata = create_test_metadata(seq_lens, num_contexts, - past_seen_tokens, request_ids, - trtllm_kv_cache_manager, - sparse_attn_config_trtllm, - RocketTrtllmAttentionMetadata) - vanilla_metadata = create_test_metadata(seq_lens, num_contexts, - past_seen_tokens, request_ids, - vanilla_kv_cache_manager, - sparse_attn_config_vanilla, - RocketVanillaAttentionMetadata) + trtllm_metadata = create_test_metadata( + seq_lens, + num_contexts, + past_seen_tokens, + request_ids, + trtllm_kv_cache_manager, + sparse_attn_config_trtllm, + RocketTrtllmAttentionMetadata, + ) + vanilla_metadata = create_test_metadata( + seq_lens, + num_contexts, + past_seen_tokens, + request_ids, + vanilla_kv_cache_manager, + sparse_attn_config_vanilla, + RocketVanillaAttentionMetadata, + ) total_tokens = sum(seq_lens) - qkv = torch.randn(total_tokens, (num_heads + 2 * num_kv_heads) * head_dim, - dtype=dtype, - device=device) + qkv = torch.randn( + total_tokens, (num_heads + 2 * num_kv_heads) * head_dim, dtype=dtype, device=device + ) for layer_idx in range(num_layers): trtllm_kt_buf = trtllm_kv_cache_manager.get_kt_buffers(layer_idx) @@ -489,9 +548,9 @@ def test_sparse_attn_predict(batch_size, num_contexts): for req_idx in range(num_contexts, batch_size): # Get the number of KT tokens for this request past_token = past_seen_tokens[req_idx] - num_kt_tokens = (past_token + 1 + - sparse_attn_config_trtllm.page_size - - 1) // sparse_attn_config_trtllm.page_size + num_kt_tokens = ( + past_token + 1 + sparse_attn_config_trtllm.page_size - 1 + ) // sparse_attn_config_trtllm.page_size # Get block offsets for this request trtllm_blocks = trtllm_block_offsets[req_idx] @@ -511,17 +570,20 @@ def test_sparse_attn_predict(batch_size, num_contexts): break # How many KT tokens in this trtllm block - kt_tokens_in_this_block = min(trtllm_kt_tokens_per_block, - num_kt_tokens - kt_token_idx) + kt_tokens_in_this_block = min( + trtllm_kt_tokens_per_block, num_kt_tokens - kt_token_idx + ) # Copy to vanilla buffer vanilla_block = vanilla_blocks[vanilla_block_idx] if vanilla_block >= 0: vanilla_kt_buf[ - vanilla_block, kt_token_idx:kt_token_idx + - kt_tokens_in_this_block].copy_(trtllm_kt_buf[ - trtllm_block, :kt_tokens_in_this_block].to( - vanilla_kt_buf.dtype)) + vanilla_block, kt_token_idx : kt_token_idx + kt_tokens_in_this_block + ].copy_( + trtllm_kt_buf[trtllm_block, :kt_tokens_in_this_block].to( + vanilla_kt_buf.dtype + ) + ) kt_token_idx += kt_tokens_in_this_block @@ -535,28 +597,26 @@ def test_sparse_attn_predict(batch_size, num_contexts): for i in range(num_contexts, batch_size): seq_len = seq_lens[i] - single_qkv = qkv[offset:offset + seq_len] - q, k, _ = single_qkv.split([ - num_heads * head_dim, num_kv_heads * head_dim, - num_kv_heads * head_dim - ], - dim=-1) + single_qkv = qkv[offset : offset + seq_len] + q, k, _ = single_qkv.split( + [num_heads * head_dim, num_kv_heads * head_dim, num_kv_heads * head_dim], dim=-1 + ) q = q.view(1, seq_len, num_heads, head_dim).transpose(1, 2) k = k.view(1, seq_len, num_kv_heads, head_dim) past_seen_token = past_seen_tokens[i] vanilla_indices = vanilla_attn._rocketkv_selection( - q, k, vanilla_metadata, past_seen_token, i) + q, k, vanilla_metadata, past_seen_token, i + ) vanilla_sparse_attn_indices_list.append(vanilla_indices.squeeze(0)) offset += seq_len if trtllm_sparse_attn_indices is not None: - assert len(vanilla_sparse_attn_indices_list - ) > 0, "Vanilla should also produce indices" + assert len(vanilla_sparse_attn_indices_list) > 0, "Vanilla should also produce indices" - vanilla_sparse_attn_indices = torch.cat( - vanilla_sparse_attn_indices_list, dim=0).transpose(0, - 1).contiguous() + vanilla_sparse_attn_indices = ( + torch.cat(vanilla_sparse_attn_indices_list, dim=0).transpose(0, 1).contiguous() + ) # Apply interleave operation to trtllm indices # For each head, multiply indices by page_size and expand to include all tokens in each page @@ -565,16 +625,15 @@ def test_sparse_attn_predict(batch_size, num_contexts): interleaved_indices_list = [] for head_idx in range(num_kv_heads): - head_indices = trtllm_sparse_attn_indices[ - head_idx] # Shape: [total_indices] + head_indices = trtllm_sparse_attn_indices[head_idx] # Shape: [total_indices] page_starts = head_indices * page_size # Shape: [total_indices] expanded_indices = [] for page_start in page_starts: - page_indices = torch.arange(page_start, - page_start + page_size, - device=page_starts.device) + page_indices = torch.arange( + page_start, page_start + page_size, device=page_starts.device + ) expanded_indices.append(page_indices) head_interleaved = torch.cat(expanded_indices, dim=0) @@ -586,16 +645,14 @@ def test_sparse_attn_predict(batch_size, num_contexts): interleaved_indices_list.append(head_interleaved) # Stack all heads - trtllm_sparse_attn_indices = torch.stack(interleaved_indices_list, - dim=0) + trtllm_sparse_attn_indices = torch.stack(interleaved_indices_list, dim=0) - assert trtllm_sparse_attn_indices.shape == vanilla_sparse_attn_indices.shape, \ + assert trtllm_sparse_attn_indices.shape == vanilla_sparse_attn_indices.shape, ( f"Shape mismatch: {trtllm_sparse_attn_indices.shape} vs {vanilla_sparse_attn_indices.shape}" + ) - trtllm_sparse_attn_indices = trtllm_sparse_attn_indices.sort( - dim=-1).values - vanilla_sparse_attn_indices = vanilla_sparse_attn_indices.sort( - dim=-1).values + trtllm_sparse_attn_indices = trtllm_sparse_attn_indices.sort(dim=-1).values + vanilla_sparse_attn_indices = vanilla_sparse_attn_indices.sort(dim=-1).values # Check indices overlap per batch and per head num_kv_heads = trtllm_sparse_attn_indices.shape[0] @@ -613,39 +670,40 @@ def test_sparse_attn_predict(batch_size, num_contexts): batch_overlaps = [] for head_idx in range(num_kv_heads): - trtllm_batch = trtllm_sparse_attn_indices[ - head_idx, start_idx:end_idx].cpu().tolist() - vanilla_batch = vanilla_sparse_attn_indices[ - head_idx, start_idx:end_idx].cpu().tolist() + trtllm_batch = ( + trtllm_sparse_attn_indices[head_idx, start_idx:end_idx].cpu().tolist() + ) + vanilla_batch = ( + vanilla_sparse_attn_indices[head_idx, start_idx:end_idx].cpu().tolist() + ) trtllm_set = set(trtllm_batch) vanilla_set = set(vanilla_batch) # Calculate overlap overlap = len(vanilla_set & trtllm_set) - overlap_ratio = overlap / len(vanilla_set) if len( - vanilla_set) > 0 else 1.0 + overlap_ratio = overlap / len(vanilla_set) if len(vanilla_set) > 0 else 1.0 batch_overlaps.append(overlap_ratio) overlap_ratios.append(overlap_ratio) avg_batch_overlap = sum(batch_overlaps) / len(batch_overlaps) - batch_overlap_details.append( - f"Batch {batch_idx}: {avg_batch_overlap:.4f}") + batch_overlap_details.append(f"Batch {batch_idx}: {avg_batch_overlap:.4f}") avg_overlap_ratio = sum(overlap_ratios) / len(overlap_ratios) print(f"Average overlap ratio: {avg_overlap_ratio:.4f}") print(f"Per-batch average: {batch_overlap_details}") threshold = 0.94 - assert avg_overlap_ratio >= threshold, \ + assert avg_overlap_ratio >= threshold, ( f"Indices overlap ratio {avg_overlap_ratio:.4f} is too low (< {threshold})" + ) else: - assert len( - vanilla_sparse_attn_indices_list - ) == 0, "Both should return None when no sparse attention is needed" + assert len(vanilla_sparse_attn_indices_list) == 0, ( + "Both should return None when no sparse attention is needed" + ) -if __name__ == '__main__': +if __name__ == "__main__": # RocketKV e2e tests print("=== Testing RocketKV E2E tests ===") test_model("pytorch", "llama-3.1-model/Llama-3.1-8B-Instruct", "VANILLA") diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index 35fb9d5e1352..25d23ebfa686 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -13,9 +13,9 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.interface import ( AttentionBackend, PositionalEmbeddingParams, RopeParams) -from tensorrt_llm._torch.attention_backend.sparse.dsa import DSACacheManager from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ DeepseekV4CacheManager +from tensorrt_llm._torch.attention_backend.sparse.dsa import DSACacheManager from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -24,8 +24,8 @@ torch_dtype_to_str) from tensorrt_llm.functional import PositionEmbeddingType, RopeEmbeddingUtils from tensorrt_llm.llmapi.llm_args import (DeepSeekSparseAttentionConfig, - KvCacheConfig, - DeepSeekV4SparseAttentionConfig) + DeepSeekV4SparseAttentionConfig, + KvCacheConfig) from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -380,23 +380,23 @@ def extract_kv_slices(indices, start_offset=0): def calculate_reference_output_deepseek_v4_prefill(hidden_states, - indices, - seq_lens, - q, - kv_data, - freqs_cis, - ref_compressor, - num_heads, - qk_nope_head_dim, - qk_rope_head_dim, - v_head_dim, - softmax_scale, - device, - o_a_proj, - o_b_proj_weight, - n_local_groups, - window_size=512, - compress_ratio=4): + indices, + seq_lens, + q, + kv_data, + freqs_cis, + ref_compressor, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + softmax_scale, + device, + o_a_proj, + o_b_proj_weight, + n_local_groups, + window_size=512, + compress_ratio=4): """Reference for deepseek_v4 prefill. Implements attention over: @@ -435,8 +435,7 @@ def calculate_reference_output_deepseek_v4_prefill(hidden_states, # fp32-throughout postprocess); cast to the attention dtype # before concatenation. kv_combined = torch.cat( - [kv_seq, compressed_kv.to(kv_seq.dtype)], - dim=1).squeeze(0) + [kv_seq, compressed_kv.to(kv_seq.dtype)], dim=1).squeeze(0) else: kv_combined = kv_seq.squeeze(0) num_compressed = 0 @@ -492,24 +491,24 @@ def calculate_reference_output_deepseek_v4_prefill(hidden_states, def calculate_reference_output_deepseek_v4_generation(hidden_states, - gen_indices, - seq_lens, - q, - kv_data, - freqs_cis, - ref_compressor, - num_contexts, - num_heads, - qk_nope_head_dim, - qk_rope_head_dim, - v_head_dim, - softmax_scale, - device, - o_a_proj, - o_b_proj_weight, - n_local_groups, - window_size=512, - compress_ratio=4): + gen_indices, + seq_lens, + q, + kv_data, + freqs_cis, + ref_compressor, + num_contexts, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + softmax_scale, + device, + o_a_proj, + o_b_proj_weight, + n_local_groups, + window_size=512, + compress_ratio=4): """Reference for deepseek_v4 generation with cached window_kv and compressed_kv. For generation, we have: @@ -636,26 +635,26 @@ def calculate_reference_output_deepseek_v4_generation(hidden_states, def calculate_reference_output_deepseek_v4_mixed(hidden_states, - q_ctx, - q_gen, - kv_data, - freqs_cis, - ref_compressor, - ctx_indices, - gen_indices, - seq_lens, - query_lens, - num_heads, - qk_nope_head_dim, - qk_rope_head_dim, - v_head_dim, - softmax_scale, - device, - o_a_proj, - o_b_proj_weight, - n_local_groups, - window_size=512, - compress_ratio=4): + q_ctx, + q_gen, + kv_data, + freqs_cis, + ref_compressor, + ctx_indices, + gen_indices, + seq_lens, + query_lens, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + softmax_scale, + device, + o_a_proj, + o_b_proj_weight, + n_local_groups, + window_size=512, + compress_ratio=4): """Reference for deepseek_v4 mixed batch (combines context and generation).""" ref_results = [None] * len(seq_lens) @@ -1081,10 +1080,9 @@ def populate_gen_deepseek_v4_kv_cache( sparse_attention_config=sparse_config, ) cached_metadata.prepare() - _ = mla_forward_impl_with_deepseek_v4_wo_linear(mla, cached_metadata, q, qr, - compressed_kv, k_pe, - latent_cache, hidden_states, - position_ids, dtype, device) + _ = mla_forward_impl_with_deepseek_v4_wo_linear( + mla, cached_metadata, q, qr, compressed_kv, k_pe, latent_cache, + hidden_states, position_ids, dtype, device) kv_cache_for_ref = {} kv_cache_for_ref["window_kv"] = all_ref_window_kv # Window KV for reference @@ -1157,9 +1155,9 @@ def mla_forward_impl_with_dsa_wo_linear(mla, attn_metadata, q, qr, def mla_forward_impl_with_deepseek_v4_wo_linear(mla, attn_metadata, q, qr, - compressed_kv, k_pe, latent_cache, - hidden_states, position_ids, dtype, - device): + compressed_kv, k_pe, + latent_cache, hidden_states, + position_ids, dtype, device): """Forward implementation for deepseek_v4 sparse attention. For deepseek_v4, compressed_kv is actually the nope part of KV (qk_nope_head_dim), diff --git a/tests/unittest/_torch/thop/serial/test_moe_gate.py b/tests/unittest/_torch/thop/serial/test_moe_gate.py index 0f5e66ecf0bd..d22b1b5de654 100644 --- a/tests/unittest/_torch/thop/serial/test_moe_gate.py +++ b/tests/unittest/_torch/thop/serial/test_moe_gate.py @@ -189,8 +189,12 @@ def test_gate_forward_384_experts(self, is_hash): if is_hash: vocab_size = 1024 - input_ids = torch.randint(0, vocab_size, (batch_size,), dtype=torch.int32, device="cuda") - tid2eid = torch.randint(0, N_EXPERTS_PRO, (vocab_size, TOPK), dtype=torch.int32, device="cuda") + input_ids = torch.randint( + 0, vocab_size, (batch_size,), dtype=torch.int32, device="cuda" + ) + tid2eid = torch.randint( + 0, N_EXPERTS_PRO, (vocab_size, TOPK), dtype=torch.int32, device="cuda" + ) bias = torch.empty(0, dtype=torch.float32, device="cuda") ref_weights, ref_indices = pytorch_gate_forward( scores, input_ids=input_ids, tid2eid=tid2eid, route_scale=route_scale, is_hash=True From d7beb0bd08f38c3782f6b286138c974cf0ab9c3e Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Thu, 7 May 2026 08:59:10 +0800 Subject: [PATCH 18/58] [None][fix] Plumb swiglu_limit through DeepGEMM and TRTLLMGen FP8 fused MoE (#13767) Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> (cherry picked from commit 1a52b727bf7dc3d0ee97490fae1e635e9e04e57a) Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit 7a9b0caecf4040b54e531f001281c7fd06761ece) Signed-off-by: Fanrong Li --- .../blockScaleMoe/DevKernel.cu | 34 ++++++++++++++++--- .../blockScaleMoe/DevKernel.h | 16 +++++++++ .../trtllmGenKernels/blockScaleMoe/runner.cu | 5 +++ .../trtllmGenKernels/blockScaleMoe/runner.h | 8 +++++ cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp | 21 ++++++++++-- .../custom_ops/trtllm_gen_custom_ops.py | 11 +++++- .../_torch/models/modeling_deepseekv4.py | 4 +++ .../modules/fused_moe/configurable_moe.py | 1 + .../_torch/modules/fused_moe/create_moe.py | 24 ++++++++++--- .../modules/fused_moe/fused_moe_cutlass.py | 2 ++ .../modules/fused_moe/fused_moe_deepgemm.py | 7 +++- .../modules/fused_moe/fused_moe_trtllm_gen.py | 25 ++++++++++++-- .../modules/fused_moe/fused_moe_wide_ep.py | 4 +++ .../_torch/modules/fused_moe/interface.py | 7 ++++ .../modules/fused_moe/moe_op_backend.py | 9 +++++ .../modules/fused_moe/ops/moe_op_deepgemm.py | 3 +- tensorrt_llm/quantization/utils/fp8_utils.py | 21 ++++++++++++ .../defs/accuracy/test_llm_api_pytorch.py | 9 ++--- .../test_lists/test-db/l0_dgx_b300_ds.yml | 10 +++++- 19 files changed, 199 insertions(+), 22 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu index 15e642e5b9be..22f1d38b83f7 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.cu @@ -69,6 +69,13 @@ __global__ void activationKernel(KernelParams params) } #endif + // FP8 separate-activation path: swiglu_limit is uniform across experts, + // passed by value via KernelParams::swigluLimit (gated by hasSwigluLimit). + // Apply gate.clamp(max=limit) / up.clamp(-limit, limit) before silu/mul. + // Per-expert non-uniform limits are not supported here. + bool const hasSwigluLimit = params.hasSwigluLimit; + float const swigluLimit = params.swigluLimit; + for (int tokenIdx = blockIdx.z; tokenIdx < params.numTokens; tokenIdx += gridDim.z) { // Look over experts per token @@ -85,8 +92,14 @@ __global__ void activationKernel(KernelParams params) { int const baseIdx = permutedIdx * params.innerDim + hiddenIdx; - float x1 = (float) params.inPtr[baseIdx]; - float x2 = (float) params.inPtr[baseIdx + params.innerDim / 2]; + float x1 = (float) params.inPtr[baseIdx]; // up (linear) + float x2 = (float) params.inPtr[baseIdx + params.innerDim / 2]; // gate (silu input) + + if (hasSwigluLimit) + { + x2 = fminf(x2, swigluLimit); + x1 = fmaxf(fminf(x1, swigluLimit), -swigluLimit); + } float act = silu(x2); Type out = (Type) (act * x1); @@ -243,6 +256,14 @@ __global__ void activationDeepSeekKernel(KernelParams params) // The largest (finite) value that can be represented using E4m3. float constexpr E4m3MaxVal{448.f}; + // FP8 separate-activation path: swiglu_limit is uniform across experts, + // passed by value via KernelParams::swigluLimit (gated by hasSwigluLimit). + // Apply gate.clamp(max=limit) / up.clamp(-limit, limit) AFTER + // dequantization but BEFORE silu/mul. Per-expert non-uniform limits are + // not supported here. + bool const hasSwigluLimit = params.hasSwigluLimit; + float const swigluLimit = params.swigluLimit; + int const totalNumPaddedTokens = params.totalNumPaddedTokens[0]; // Loop over tokens float scale1Arr[NumTokensPerCta]; @@ -305,8 +326,13 @@ __global__ void activationDeepSeekKernel(KernelParams params) #pragma unroll for (int tokenInCtaIdx = 0; tokenInCtaIdx < NumTokensPerCta; tokenInCtaIdx++) { - float x1 = scale1Arr[tokenInCtaIdx] * dataX1Arr[tokenInCtaIdx]; - float x2 = scale2Arr[tokenInCtaIdx] * dataX2Arr[tokenInCtaIdx]; + float x1 = scale1Arr[tokenInCtaIdx] * dataX1Arr[tokenInCtaIdx]; // up (linear) + float x2 = scale2Arr[tokenInCtaIdx] * dataX2Arr[tokenInCtaIdx]; // gate (silu input) + if (hasSwigluLimit) + { + x2 = fminf(x2, swigluLimit); + x1 = fmaxf(fminf(x1, swigluLimit), -swigluLimit); + } float act = silu(x2); float out = act * x1; outArr[tokenInCtaIdx] = out; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h index 72996839037d..171ceaaf12ea 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/DevKernel.h @@ -260,6 +260,16 @@ struct Data int32_t* expandedIdxToPermutedIdx; int32_t const* totalNumPaddedTokens; + + // Optional swiglu clamp limit (fp32, uniform across experts on the FP8 + // separate-activation path). When `hasSwigluLimit` is true, the kernel + // applies gate.clamp(max=swigluLimit) and up.clamp(-swigluLimit, + // swigluLimit) before silu/mul, matching the swiglu_torch reference. + // NVFP4 still consumes a per-expert tensor via fused-activation GEMM + // cubins (see MoERunnerArgs::gemm1_clamp_limit); this scalar is only + // for the FP8 separate-activation kernel. + float swigluLimit = 0.0f; + bool hasSwigluLimit = false; }; template @@ -282,6 +292,9 @@ struct KernelParams int32_t const* totalNumPaddedTokens; + float swigluLimit = 0.0f; + bool hasSwigluLimit = false; + static KernelParams setKernelParams(Data const& data) { KernelParams params; @@ -298,6 +311,9 @@ struct KernelParams params.topK = data.topK; params.totalNumPaddedTokens = data.totalNumPaddedTokens; + params.swigluLimit = data.swigluLimit; + params.hasSwigluLimit = data.hasSwigluLimit; + return params; } }; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu index d47ea872cb1f..39021ce642b1 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu @@ -670,6 +670,11 @@ void Runner::setOpsData(MoERunnerArgs const& args, MoEWorkspace const& workspace activationData.topK = args.top_k; activationData.numTokens = args.num_tokens; activationData.expandedIdxToPermutedIdx = workspace.expanded_idx_to_permuted_idx; + // For DeepSeek FP8 the activation runs as a separate kernel rather than + // fused into the FC1 GEMM cubin; forward the scalar swiglu_limit so it + // can honor swiglu_limit (uniform across experts; see DevKernel.h note). + activationData.swigluLimit = args.gemm1_clamp_limit_value; + activationData.hasSwigluLimit = args.has_gemm1_clamp_limit_value; activationData.totalNumPaddedTokens = workspace.total_num_padded_tokens; diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h index 30554747d820..97d77eddad0a 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h @@ -284,7 +284,15 @@ struct MoERunnerArgs float* gemm1_bias = nullptr; float* gemm1_alpha = nullptr; float* gemm1_beta = nullptr; + // Per-expert clamp tensor consumed by fused-activation GEMM cubins on the + // NVFP4 / MXFP4 paths. The FP8 separate-activation kernel uses the scalar + // variants below instead — see DevKernel.h::activation::Data::swigluLimit. float* gemm1_clamp_limit = nullptr; + // Uniform-across-experts scalar variant, only consumed by the FP8 + // separate-activation kernel (DevKernel.cu::activationKernel / + // activationDeepSeekKernel). Avoids a per-CTA global load. + float gemm1_clamp_limit_value = 0.0f; + bool has_gemm1_clamp_limit_value = false; float* gemm2_bias = nullptr; int32_t num_tokens{0}; diff --git a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp index 2a403ae868e4..e3079b3e71ee 100644 --- a/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp +++ b/cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp @@ -44,7 +44,8 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit int64_t const intermediate_size, int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t const tile_tokens_dim, int64_t const routing_method_type, MoeRunnerType& moe_runner, int64_t moeConfigIndex, std::optional const& topk_weights, - std::optional const& topk_ids, std::optional const& out_tensor = std::nullopt) + std::optional const& topk_ids, std::optional const& gemm1_clamp_limit = std::nullopt, + std::optional const& out_tensor = std::nullopt) { TORCH_CHECK(tensorrt_llm::common::isSM100Family(), "Only SM100f is supported by FP8 block scale MOE"); @@ -168,6 +169,19 @@ at::Tensor run_fp8_block_scale_moe(at::optional const& routing_logit args.intermediate_size = intermediate_size; args.mUseDeepSeekFp8 = true; + if (gemm1_clamp_limit.has_value()) + { + // FP8 path's separate activation kernel honors a single uniform clamp; + // see DevKernel.h::activation::Data::swigluLimit. NVFP4 path's + // fused-activation cubins consume a per-expert tensor via + // args.gemm1_clamp_limit (kept for API symmetry, populated by + // run_fp4_block_scale_moe). If non-uniform usage becomes necessary on + // the FP8 path, extend the activation kernel with a + // permutedIdx -> expertIdx lookup and surface a tensor variant here. + args.gemm1_clamp_limit_value = static_cast(gemm1_clamp_limit.value()); + args.has_gemm1_clamp_limit_value = true; + } + // allocate workspace for routing kernel if (routing_logits.has_value() && topk_ids.has_value()) { @@ -388,7 +402,8 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder int64_t const local_expert_offset, int64_t const local_num_experts, std::optional const routed_scaling_factor, int64_t routing_method_type, std::vector tile_config_pair, std::optional const& topk_weights, - std::optional const& topk_ids, std::optional const& output = std::nullopt) + std::optional const& topk_ids, std::optional const& gemm1_clamp_limit = std::nullopt, + std::optional const& output = std::nullopt) { // tile_config_pair corresponds to pair (tileN, config) auto [tileN, config] = std::tie(tile_config_pair[0], tile_config_pair[1]); @@ -409,7 +424,7 @@ class FP8BlockScaleMoeRunner : public torch::CustomClassHolder return run_fp8_block_scale_moe(routing_logits, routing_bias, hidden_states, hidden_states_scale, gemm1_weights, gemm1_weights_scale, gemm2_weights, gemm2_weights_scale, num_experts, top_k, n_group, topk_group, intermediate_size, local_expert_offset, local_num_experts, routed_scaling_factor, tileN, - routing_method_type, *mRunners.at(tileN), config, topk_weights, topk_ids, output); + routing_method_type, *mRunners.at(tileN), config, topk_weights, topk_ids, gemm1_clamp_limit, output); } private: diff --git a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py index 0f9eb4a5e618..a8bbb4cb2f09 100644 --- a/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py @@ -830,6 +830,7 @@ def __init__( act_type: int, tune_max_num_tokens: int = 8192, use_dp: bool = False, + gemm1_clamp_limit_value: Optional[float] = None, ): self.num_experts = num_experts @@ -842,6 +843,10 @@ def __init__( self.routed_scaling_factor = routed_scaling_factor self.routing_method_type = routing_method_type self.act_type = 0 + # Uniform-across-experts swiglu_limit. Passed by value to the C++ + # binding (not part of the autotuner's tensor input list, since it + # doesn't influence tactic validity). + self.gemm1_clamp_limit_value = gemm1_clamp_limit_value self.tuning_config = FP8BlockScaleMoERunner.get_tuning_config( self.num_experts // self.local_num_experts, tune_max_num_tokens=tune_max_num_tokens, @@ -880,7 +885,8 @@ def forward( self.n_group, self.topk_group, self.intermediate_size, self.local_expert_offset, self.local_num_experts, self.routed_scaling_factor, self.routing_method_type, tactic, - args.topk_weights, args.topk_ids, output) + args.topk_weights, args.topk_ids, self.gemm1_clamp_limit_value, + output) def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, @@ -1011,6 +1017,7 @@ def fp8_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], topk_weights: Optional[torch.Tensor] = None, topk_ids: Optional[torch.Tensor] = None, act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, use_dp: bool = False) -> torch.Tensor: @@ -1029,6 +1036,7 @@ def fp8_block_scale_moe_runner(routing_logits: Optional[torch.Tensor], act_type, tune_max_num_tokens=tune_max_num_tokens, use_dp=use_dp, + gemm1_clamp_limit_value=gemm1_clamp_limit, ) # Prepare dummy topk tensors and hook for AutoTuner profiling @@ -1110,6 +1118,7 @@ def _(routing_logits: torch.Tensor, topk_weights: Optional[torch.Tensor] = None, topk_ids: Optional[torch.Tensor] = None, act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, output: Optional[torch.Tensor] = None, tune_max_num_tokens: int = 8192, use_dp: bool = False) -> torch.Tensor: diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 8529e6b05af6..415f10c08f3d 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -73,6 +73,7 @@ create_moe, get_moe_cls, ) +from ..modules.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE from ..modules.fused_moe.fused_moe_wide_ep import WideEPMoE from ..modules.linear import Linear from ..modules.mhc.hyper_connection import HCHead, HCState, mHC @@ -1404,6 +1405,8 @@ def __init__( CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE, + WideEPMoE, + DeepGemmFusedMoE, ) if supports_swiglu_limit: moe_load_balancer_config = getattr(model_config, "moe_load_balancer", None) @@ -1437,6 +1440,7 @@ def __init__( else MoEWeightLoadingMode.VANILLA ), swiglu_limit=moe_swiglu_limit, + swiglu_limit_scalar=(float(swiglu_limit) if swiglu_limit is not None else None), ) self.mapping = model_config.mapping diff --git a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py index 0b28df3156ea..51110bbe3f4a 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py @@ -323,6 +323,7 @@ def _create_and_sync_backend( swiglu_alpha=kwargs.get("swiglu_alpha"), swiglu_beta=kwargs.get("swiglu_beta"), swiglu_limit=kwargs.get("swiglu_limit"), + swiglu_limit_scalar=kwargs.get("swiglu_limit_scalar"), init_load_balancer=False, without_comm=True, activation_type=self.activation_type, diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 67142b6e02d2..e12849b2483c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -221,6 +221,7 @@ def create_moe_backend( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = True, without_comm: bool = False, activation_type: ActivationType = ActivationType.Swiglu, @@ -244,7 +245,8 @@ def create_moe_backend( layer_idx: Layer index swiglu_alpha: SwiGLU alpha parameter swiglu_beta: SwiGLU beta parameter - swiglu_limit: SwiGLU limit parameter + swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4) + swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8) activation_type: Activation type Returns: @@ -295,9 +297,11 @@ def create_moe_backend( assert swiglu_alpha is not None and swiglu_beta is not None, \ "Both swiglu_alpha and swiglu_beta must be provided." - if swiglu_limit is not None: - assert moe_cls in [CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE], \ - f"swiglu_limit is only supported in CutlassFusedMoE, TritonFusedMoE and TRTLLMGenFusedMoE, not in {moe_cls.__name__}." + if swiglu_limit is not None or swiglu_limit_scalar is not None: + assert moe_cls in [ + CutlassFusedMoE, TritonFusedMoE, TRTLLMGenFusedMoE, WideEPMoE, + DeepGemmFusedMoE + ], f"swiglu_limit is not supported in {moe_cls.__name__}." if moe_cls == TRTLLMGenFusedMoE: assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in TRTLLMGenFusedMoE." @@ -317,6 +321,7 @@ def create_moe_backend( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, without_comm=without_comm, activation_type=activation_type, @@ -341,6 +346,7 @@ def create_moe_backend( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, without_comm=without_comm, activation_type=activation_type, @@ -358,6 +364,8 @@ def create_moe_backend( weight_loading_mode=weight_loading_mode, apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, + swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, activation_type=activation_type) elif moe_cls == VanillaMoE: assert not apply_router_weight_on_input, "apply_router_weight_on_input is not supported in VanillaMoE." @@ -408,6 +416,8 @@ def create_moe_backend( apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, init_load_balancer=init_load_balancer, + swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, without_comm=without_comm, ) elif moe_cls == TritonFusedMoE: @@ -487,6 +497,7 @@ def create_moe( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, activation_type: ActivationType = ActivationType.Swiglu, ) -> MoE: """ @@ -508,7 +519,8 @@ def create_moe( layer_idx: Layer index swiglu_alpha: SwiGLU alpha parameter swiglu_beta: SwiGLU beta parameter - swiglu_limit: SwiGLU limit parameter + swiglu_limit: SwiGLU limit parameter (per-expert tensor; for NVFP4) + swiglu_limit_scalar: SwiGLU limit scalar (uniform across experts; for FP8) activation_type: Activation type Returns: @@ -560,6 +572,7 @@ def create_moe( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, activation_type=activation_type, ) else: @@ -599,5 +612,6 @@ def create_moe( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, activation_type=activation_type, ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 93eb1b931955..acae5cd37ba3 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -225,6 +225,7 @@ def __init__( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = True, without_comm: bool = False, activation_type: ActivationType = ActivationType.Swiglu, @@ -243,6 +244,7 @@ def __init__( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, layer_idx=layer_idx, init_load_balancer=init_load_balancer, activation_type=activation_type, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py index 6e39e22a8c16..601bc0abed74 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_deepgemm.py @@ -797,6 +797,8 @@ def __init__( VANILLA, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, + swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = True, without_comm: bool = False, ): @@ -826,6 +828,8 @@ def __init__( weight_loading_mode=weight_loading_mode, apply_router_weight_on_input=apply_router_weight_on_input, layer_idx=layer_idx, + swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, init_load_balancer=init_load_balancer, without_comm=without_comm, ) @@ -1072,7 +1076,8 @@ def run_moe( input=h1, quant_group_size=128, masked_m=masked_m, - scale_ue8m0=True) + scale_ue8m0=True, + swiglu_limit=self.swiglu_limit_scalar) # Grouped gemm 2 h3 = set_strides(workspace["workspace_1"], diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 83e0342d90d2..d92a38734040 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -104,11 +104,18 @@ class TRTLLMGenFusedMoE(MoE): } # Quantization algorithms that support swiglu_gptoss_style + # FP8_BLOCK_SCALES is included for swiglu_limit only via the + # DeepSeek FP8 separate-activation path + # (DevKernel.cu::activationDeepSeekKernel reads gemm1_clamp_limit + # when MoERunnerArgs supplies it). bias and swiglu_alpha/beta still + # require GEMM-fused activation cubins, which are NVFP4/MXFP4 only + # — _check_configs gates those separately. _GPTOSS_SUPPORTED_ALGOS = { QuantAlgo.NVFP4, QuantAlgo.W4A16_MXFP4, QuantAlgo.W4A8_MXFP4_FP8, QuantAlgo.W4A8_MXFP4_MXFP8, + QuantAlgo.FP8_BLOCK_SCALES, } # Activations supported by the FlashInfer BF16 kernels: Swiglu and Relu2. @@ -207,6 +214,7 @@ def __init__( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, init_load_balancer: bool = True, without_comm: bool = False, activation_type: ActivationType = ActivationType.Swiglu, @@ -225,6 +233,7 @@ def __init__( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, layer_idx=layer_idx, init_load_balancer=init_load_balancer, activation_type=activation_type, @@ -461,8 +470,19 @@ def _check_configs(self): assert not self.bias and self.swiglu_alpha is None and self.swiglu_beta is None and self.swiglu_limit is None, \ "TRTLLMGenFusedMoE BF16 path does not support bias/swiglu custom parameters." - if self.bias or self.swiglu_alpha is not None or self.swiglu_beta is not None or self.swiglu_limit is not None: - assert self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8, "TRTLLMGenFusedMoE supports bias/swiglu only for nvfp4 and mxfp4 variants." + if self.bias or self.swiglu_alpha is not None or self.swiglu_beta is not None: + assert self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8, \ + "TRTLLMGenFusedMoE supports bias/swiglu_alpha/swiglu_beta only for nvfp4 and mxfp4 variants." + if self.swiglu_limit is not None or self.swiglu_limit_scalar is not None: + # swiglu_limit additionally goes through the DeepSeek FP8 + # separate-activation path + # (DevKernel.cu::activationDeepSeekKernel) when + # has_deepseek_fp8_block_scales. The FP8 path consumes the scalar + # variant (uniform across experts); NVFP4/MXFP4 fused-activation + # cubins consume the per-expert tensor. + assert self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_mxfp4_fp8 \ + or self.has_w4a8_mxfp4_mxfp8 or self.has_deepseek_fp8_block_scales, \ + "TRTLLMGenFusedMoE supports swiglu_limit only for nvfp4, mxfp4, and fp8_block_scale variants." def _get_quant_method(self): if self.quant_config is not None and self.quant_config.layer_quant_mode.has_any_quant( @@ -756,6 +776,7 @@ def run_moe( self.routing_method.routing_method_type, topk_weights=token_final_scales, topk_ids=token_selected_experts, + gemm1_clamp_limit=self.swiglu_limit_scalar, output=moe_output, tune_max_num_tokens=self.max_num_tokens, use_dp=self.use_dp, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py index 0f207ceff0e7..dcfbaca76f88 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py @@ -62,6 +62,8 @@ def __init__( VANILLA, apply_router_weight_on_input: bool = False, layer_idx: Optional[int] = None, + swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, activation_type: ActivationType = ActivationType.Swiglu, ): @@ -75,6 +77,8 @@ def __init__( model_config=model_config, aux_stream_dict=aux_stream_dict, weight_loading_mode=weight_loading_mode, + swiglu_limit=swiglu_limit, + swiglu_limit_scalar=swiglu_limit_scalar, layer_idx=layer_idx, activation_type=activation_type, ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index c57386fae2d9..6b477eb882a7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -283,6 +283,7 @@ def __init__( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + swiglu_limit_scalar: Optional[float] = None, layer_idx: Optional[int] = None, activation_type: ActivationType = ActivationType.Swiglu, init_load_balancer: bool = True, @@ -301,6 +302,12 @@ def __init__( self.swiglu_alpha = swiglu_alpha self.swiglu_beta = swiglu_beta self.swiglu_limit = swiglu_limit + # Uniform-across-experts scalar variant of swiglu_limit, consumed by + # FP8 paths (DeepGEMM Triton kernel, TRTLLMGen FP8 separate-activation + # kernel) that don't actually need a per-expert tensor. Distinct from + # `swiglu_limit` (kept for NVFP4 fused-activation cubins that *do* + # consume per-expert values via fc31_alpha rescaling). + self.swiglu_limit_scalar = swiglu_limit_scalar self.layer_idx = layer_idx self.layer_idx_str = str(layer_idx) if layer_idx is not None else None self.activation_type = int(activation_type) diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py index 775590ec2434..cf0e1ad3c6b1 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_op_backend.py @@ -115,6 +115,7 @@ def run_fp8_block_scale_moe( topk_weights: Optional[torch.Tensor] = None, topk_ids: Optional[torch.Tensor] = None, gated_act_type: int = 0, + gemm1_clamp_limit: Optional[float] = None, output: Optional[torch.Tensor] = None, use_shuffled_weight: bool = False, weight_layout: int = 0, @@ -256,6 +257,7 @@ def run_fp8_block_scale_moe( topk_weights=None, topk_ids=None, gated_act_type=0, + gemm1_clamp_limit=None, output=None, use_shuffled_weight=False, weight_layout=0, @@ -284,6 +286,7 @@ def run_fp8_block_scale_moe( topk_weights=topk_weights, topk_ids=topk_ids, act_type=gated_act_type, + gemm1_clamp_limit=gemm1_clamp_limit, output=output, tune_max_num_tokens=tune_max_num_tokens, use_dp=use_dp, @@ -599,6 +602,7 @@ def run_fp8_block_scale_moe( topk_weights=None, topk_ids=None, gated_act_type=0, + gemm1_clamp_limit=None, output=None, use_shuffled_weight=False, weight_layout=0, @@ -606,6 +610,11 @@ def run_fp8_block_scale_moe( tune_max_num_tokens=8192, use_dp=False, ): + if gemm1_clamp_limit is not None: + raise NotImplementedError( + "FlashinferOpBackend.run_fp8_block_scale_moe does not yet " + "forward gemm1_clamp_limit; use the trtllm op backend." + ) if router_logits is not None: return self._fused_moe.trtllm_fp8_block_scale_moe( router_logits, diff --git a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py index ab8936374530..9b86c02673f2 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ops/moe_op_deepgemm.py @@ -272,7 +272,8 @@ def compute_moe( input=h1, quant_group_size=128, masked_m=masked_m, - scale_ue8m0=True) + scale_ue8m0=True, + swiglu_limit=getattr(module, "swiglu_limit_scalar", None)) # Second grouped GEMM (w2) h3 = set_strides(workspace["workspace_1"], expert_size_per_partition, diff --git a/tensorrt_llm/quantization/utils/fp8_utils.py b/tensorrt_llm/quantization/utils/fp8_utils.py index a2f5e6942857..19c0a336d17b 100644 --- a/tensorrt_llm/quantization/utils/fp8_utils.py +++ b/tensorrt_llm/quantization/utils/fp8_utils.py @@ -369,6 +369,8 @@ def _silu_and_mul_post_quant_kernel( BLOCK: tl.constexpr, NUM_STAGE: tl.constexpr, SCALE_UE8M0: tl.constexpr, + HAS_SWIGLU_LIMIT: tl.constexpr, + SWIGLU_LIMIT: tl.constexpr, ): expert_id = tl.program_id(2) token_id = tl.program_id(1) @@ -408,6 +410,15 @@ def _silu_and_mul_post_quant_kernel( mask=local_mask < size_k, other=0.0, ).to(tl.float32) + if HAS_SWIGLU_LIMIT: + # gate is fp32; clamp directly. up is in input dtype (bf16/fp16); + # cast the limit constant to that dtype to avoid a fp32 round-trip + # on every element. + gate = tl.minimum(gate, SWIGLU_LIMIT) + limit_native = tl.cast(SWIGLU_LIMIT, input_ptr.dtype.element_ty) + neg_limit_native = tl.cast(-SWIGLU_LIMIT, + input_ptr.dtype.element_ty) + up = tl.maximum(tl.minimum(up, limit_native), neg_limit_native) gate = gate / (1 + tl.exp(-gate)) gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate @@ -438,6 +449,7 @@ def silu_and_mul_masked_post_quant_fwd( quant_group_size: int, masked_m: torch.Tensor, scale_ue8m0: bool = False, + swiglu_limit: Optional[float] = None, ): """ input shape [g, m, k] @@ -445,6 +457,11 @@ def silu_and_mul_masked_post_quant_fwd( output_scale [g, k // 4, m // 2 // 128], dtype int32 quant_group_size int masked_m shape [g] + swiglu_limit optional Python float (uniform across experts); when provided, + the gate input is clamped to (-inf, limit] before silu and the up + branch is clamped to [-limit, limit] before the multiply. Baked into + the kernel as a constexpr — caller supplies a host-side float, no + per-expert tensor / global load. """ assert input.is_contiguous() @@ -477,6 +494,8 @@ def silu_and_mul_masked_post_quant_fwd( BLOCK_NUM_PER_EXPERT, expert_num, ) + has_swiglu_limit = swiglu_limit is not None + swiglu_limit_value = float(swiglu_limit) if has_swiglu_limit else 0.0 _silu_and_mul_post_quant_kernel[grid]( input, *input.stride(), @@ -492,6 +511,8 @@ def silu_and_mul_masked_post_quant_fwd( NUM_STAGE=NUM_STAGES, num_warps=num_warps, SCALE_UE8M0=scale_ue8m0, + HAS_SWIGLU_LIMIT=has_swiglu_limit, + SWIGLU_LIMIT=swiglu_limit_value, ) output_scale = output_scale.transpose(1, 2)[:, :m, :] check_sf_layout( diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 72d767470761..f15aaf2928fa 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -3739,15 +3739,16 @@ class TestDeepSeekV4FlashBase(LlmapiAccuracyTestHarness): MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Flash-Base" @pytest.mark.skip_less_mpi_world_size(4) - def test_auto_dtype(self): + @parametrize_with_ids("moe_backend", ["WIDEEP", "TRTLLM"]) + def test_auto_dtype(self, moe_backend): # Aggregate (non-disagg, non-EPLB) smoke test. FP8 weights ~71 GB/rank - # at TP=4 — fits on 4x B300 (~288 GB/GPU). WIDEEP required (CUTLASS - # path is Hopper-only on Blackwell). 1-sample smoke. + # at TP=4 — fits on 4x B300 (~288 GB/GPU). 1-sample smoke. CUTLASS is + # Hopper-only on Blackwell and skipped here. kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.5) with LLM(self.MODEL_PATH, tensor_parallel_size=4, moe_expert_parallel_size=4, - moe_config=MoeConfig(backend="WIDEEP"), + moe_config=MoeConfig(backend=moe_backend), enable_attention_dp=True, max_seq_len=4096, kv_cache_config=kv_cache_config) as llm: diff --git a/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml b/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml index 58a05389bc1a..281e993ff12b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b300_ds.yml @@ -24,7 +24,15 @@ l0_dgx_b300_ds: # smoke fit at TP=4. (Same tests at TP=8 would also fit on 8x B200, but a # dedicated 8-GPU B200 stage is not currently provisioned.) # V4-Flash-Base FP8 aggregate smoke (TP=4, no EPLB): 1 MMLU sample. - - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_auto_dtype TIMEOUT (60) + # TRTLLMGen FP8 backend selected here: it is the user-facing default on + # Blackwell (model_config.py::resolve_moe_backend) and ~7% faster per step + # than WIDEEP in our measurements. WIDEEP variant of the same test is + # currently held out of CI (see test_auto_dtype[moe_backend=WIDEEP]). + # TODO: re-enable once DeepSeek-V4-Flash-Base is uploaded to the CI model + # share at /scratch.trt_llm_data/llm-models/DeepSeek-V4-Flash-Base. The + # model is currently missing on B300 CI workers, so LLM(...) treats the + # path as an HF repo id and HFValidationError fires. + # - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4FlashBase::test_auto_dtype[moe_backend=TRTLLM] TIMEOUT (60) # V4-Flash NVFP4 static EPLB sanity. WIDEEP backend is skipped at the # test level because V4-Flash MXFP4 routed experts are unsupported by # WIDEEP (raises "Unsupported quantization mode: [65536]"); TRTLLM is the From f450e160d8633a14165e1ad5296505c2a8687772 Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 7 May 2026 10:34:06 +0800 Subject: [PATCH 19/58] [None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4 (#13646) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-authored-by: peihengh <259410613+peihu-nv@users.noreply.github.com> (cherry picked from commit a3f37757bdf1920d7ab2455ec0cbb86bbd98f171) Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit aa67e3ed3795e34ca1598d4032c70df138fcd1eb) Signed-off-by: Fanrong Li --- .../dsv3MinLatencyKernels/dsv3RouterGemm.cu | 49 +++++++++++++++++++ cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp | 11 ++++- .../_torch/models/modeling_deepseekv4.py | 4 +- .../thop/parallel/test_dsv3_router_gemm.py | 3 +- 4 files changed, 62 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index e33f46e1e152..102866de363f 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu @@ -291,6 +291,55 @@ template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__n template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 6144>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); +// hidden_dim=4096 instantiations (DeepSeek-V4). +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 1, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 2, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 3, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 4, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 5, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 6, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 7, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 8, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 9, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 10, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 11, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 12, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 13, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 14, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 15, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 4096>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + } // namespace kernels::dsv3MinLatencyKernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp index 2e6e0822debe..76760710b959 100644 --- a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp @@ -79,15 +79,16 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, constexpr int kNumExperts = 256; constexpr int kHiddenDim7168 = 7168; // DeepSeek-V3 / DeepSeek-V3.2 constexpr int kHiddenDim6144 = 6144; // GLM-5 + constexpr int kHiddenDim4096 = 4096; // DeepSeek-V4 std::vector output_size = {mat_a.sizes()[0], mat_b.sizes()[1]}; th::Tensor out = th::empty(output_size, mat_a.options().dtype(out_dtype_)); TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2); TORCH_CHECK(mat_a.strides()[1] == 1 && out.strides()[1] == 1); // Row-major TORCH_CHECK(mat_b.strides()[0] == 1); // Column-major - TORCH_CHECK(!bias.has_value(), "bias is not support yet"); auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); bool const shape_ok = (num_tokens >= 1 && num_tokens <= 16 && num_experts == kNumExperts - && mat_b.sizes()[0] == hidden_dim && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32); + && mat_b.sizes()[0] == hidden_dim && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32 + && !bias.has_value()); if (shape_ok && hidden_dim == kHiddenDim7168) { @@ -101,6 +102,12 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); } + else if (shape_ok && hidden_dim == kHiddenDim4096) + { + LoopUnroller<1, 16, kNumExperts, kHiddenDim4096>::unroll(num_tokens, + reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); + } else // fallback to cublas, can be slow { cublas_mm_out(mat_a, mat_b, bias, out); diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 415f10c08f3d..d150ba5d28f9 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1326,7 +1326,9 @@ def fetch_e_score_correction_bias(): ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return torch.nn.functional.linear(hidden_states.float(), self.weight.float()) + return torch.ops.trtllm.dsv3_router_gemm_op( + hidden_states, self.weight.t(), bias=None, out_dtype=torch.float32 + ) def load_weights(self, weights: List[Dict]): assert len(weights) == 1 diff --git a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py index abafe6a3e271..1e03da59024a 100644 --- a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py @@ -25,8 +25,7 @@ def router_gemm_ref(input, weight, bias, dtype): @pytest.mark.parametrize( "num_tokens", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) @pytest.mark.parametrize("num_experts", [256]) -@pytest.mark.parametrize("hidden_size", - [7168, 6144, 4096]) # 4096 will enter fallback kernel +@pytest.mark.parametrize("hidden_size", [7168, 6144, 4096]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) def test_router_gemm_run(num_tokens, num_experts, hidden_size, dtype): torch.manual_seed(24) From 921f99419557917c80840938ecb5ce928bc7823d Mon Sep 17 00:00:00 2001 From: Mingyang Date: Thu, 7 May 2026 11:12:04 +0800 Subject: [PATCH 20/58] [None][perf] Optimize DeepSeek-V4 compressor BF16 input (#13761) Signed-off-by: Mingyang Hao Co-authored-by: Mingyang Hao (cherry picked from commit b33d2dc528aa352d30561f54352654c2fca23c2a) Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit 0956712c0642dfa4f7a998abecf2a3568447fdf5) Signed-off-by: Fanrong Li --- .../compressorKernels/compressorKernels.cu | 458 +++++++++--------- .../compressorKernels/compressorKernels.h | 8 +- cpp/tensorrt_llm/thop/compressorOp.cpp | 29 +- .../sparse/deepseek_v4/compressor.py | 8 +- .../deepseek_v4/test_compressor_kernel.py | 165 ++++++- .../deepseek_v4/test_compressor_module.py | 2 + .../deepseek_v4/test_compressor_tf32.py | 11 +- 7 files changed, 436 insertions(+), 245 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index 7fbe9034d50e..cf48dd04b720 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -55,9 +55,10 @@ // per position but improves compression quality. // // Template parameters: -// HEAD_DIM — Head dimension (128 or 512) -// ELEM_BYTES — Element size in bytes (2=bf16, 4=fp32) -// SCALE_TYPE — Output cache scale/dtype for postProcessScatterKernel +// HEAD_DIM — Head dimension (128 or 512) +// KV_SCORE_ELEM_BYTES — kv_score element size (2=bf16, 4=fp32) +// STATE_ELEM_BYTES — compressor state element size (2=bf16, 4=fp32) +// SCALE_TYPE — Output cache scale/dtype for postProcessScatterKernel // ============================================================================ #include "tensorrt_llm/kernels/compressorKernels/compressorKernels.h" @@ -211,7 +212,7 @@ enum class CacheScaleType // ============================================================================ // Decode Kernel: pagedKvCompressKernel // -// Template: +// Template: // NEXT_N: number of new tokens per sequence in this decode step (1-4) // // Grid: (batch_size) — one block per batch element @@ -241,7 +242,7 @@ enum class CacheScaleType // online softmax accumulators (rmax, rsum, rwsum) per element. // APE is already baked into paged_score (added during Phase 1), so no APE // addition is performed here. -template +template __device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_kv_raw, void const* __restrict__ paged_score_raw, int64_t page_sd, // page_size * state_dim (in elements) @@ -252,21 +253,19 @@ __device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_ int kv_col_off, // column offset (0 or HEAD_DIM) int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) { - using IoElemT = typename std::conditional::type; - constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; - constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - using IoVecT = typename VecType::type; + using StateElemT = typename std::conditional::type; + using StateVecT = typename VecType::type; - auto const* kv = reinterpret_cast(paged_kv_raw); - auto const* sc = reinterpret_cast(paged_score_raw); + auto const* kv = reinterpret_cast(paged_kv_raw); + auto const* sc = reinterpret_cast(paged_score_raw); int64_t base_kv = static_cast(phys_kv) * page_sd + blk_off * state_dim + kv_col_off; int64_t base_sc = static_cast(phys_sc) * page_sd + blk_off * state_dim + kv_col_off; - IoVecT k_raw = reinterpret_cast(&kv[base_kv])[tid]; - IoVecT s_raw = reinterpret_cast(&sc[base_sc])[tid]; - IoElemT const* ke = reinterpret_cast(&k_raw); - IoElemT const* se = reinterpret_cast(&s_raw); + StateVecT k_raw = reinterpret_cast(&kv[base_kv])[tid]; + StateVecT s_raw = reinterpret_cast(&sc[base_sc])[tid]; + StateElemT const* ke = reinterpret_cast(&k_raw); + StateElemT const* se = reinterpret_cast(&s_raw); #pragma unroll for (int i = 0; i < VEC; i += 4) @@ -292,26 +291,31 @@ __device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_ } } -template +template __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, - int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ cu_seq_lens, - int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, int out_elem_bytes) + int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, + int out_elem_bytes) { - using IoElemT = typename std::conditional::type; + 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"); constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); - constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; + constexpr int ELEM_BYTES_FOR_VEC + = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; + constexpr int MAX_VEC = 16 / ELEM_BYTES_FOR_VEC; constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - using IoVecT = typename VecType::type; + using KvScoreVecT = typename VecType::type; + using StateVecT = typename VecType::type; static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); // HEAD_BLOCKS: split head_dim across blockIdx.y for better SM utilisation. - // For HD=512 bf16: NTHRD_BASE=64 → HEAD_BLOCKS=2, NTHRD_INNER=32. - // For HD=128 bf16: NTHRD_BASE=32 → HEAD_BLOCKS=1, NTHRD_INNER=32. - // For HD=512 fp32: NTHRD_BASE=128 → HEAD_BLOCKS=4, NTHRD_INNER=32. + // For HD=512 and max elem size 2: NTHRD_BASE=64 → HEAD_BLOCKS=2, NTHRD_INNER=32. + // For HD=128 and max elem size 2/4: NTHRD_BASE=32 → HEAD_BLOCKS=1, NTHRD_INNER=32. + // For HD=512 and max elem size 4: NTHRD_BASE=128 → HEAD_BLOCKS=4, NTHRD_INNER=32. constexpr int NTHRD_BASE = HEAD_DIM / VEC; constexpr int HEAD_BLOCKS = (NTHRD_BASE > 32) ? (NTHRD_BASE / 32) : 1; constexpr int NTHRD_INNER = NTHRD_BASE / HEAD_BLOCKS; // always <= 32 @@ -332,15 +336,15 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int const head_blk = blockIdx.y; int const eff_tid = head_blk * NTHRD_INNER + tid; - int const sp = start_pos_arr[batch_idx]; int const kv_len = kv_lens[batch_idx]; + int const sp = kv_len - NEXT_N; int const in_off = cu_seq_lens[batch_idx]; int const out_off = cu_kv_comp[batch_idx]; int64_t const page_sd = static_cast(page_size) * STATE_DIM; - auto const* kv_score = reinterpret_cast(kv_score_raw); - auto* paged_kv = reinterpret_cast(paged_kv_raw); - auto* paged_score = reinterpret_cast(paged_score_raw); + auto const* kv_score = reinterpret_cast(kv_score_raw); + auto* paged_kv = reinterpret_cast(paged_kv_raw); + auto* paged_score = reinterpret_cast(paged_score_raw); // ================================================================ // Phase 1: Write NEXT_N new tokens' KV and score state to paged cache. @@ -369,25 +373,33 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * STATE_DIM + col; int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * STATE_DIM + col; - IoVecT kv_raw = reinterpret_cast(&kv_score[src])[eff_tid]; - IoVecT sc_raw = reinterpret_cast(&kv_score[src + STATE_DIM])[eff_tid]; + KvScoreVecT kv_raw = reinterpret_cast(&kv_score[src])[eff_tid]; + KvScoreVecT sc_raw = reinterpret_cast(&kv_score[src + STATE_DIM])[eff_tid]; - reinterpret_cast(&paged_kv[dkv])[eff_tid] = kv_raw; + KvScoreElemT const* kv_e = reinterpret_cast(&kv_raw); + StateVecT kv_out; + StateElemT* kv_o = reinterpret_cast(&kv_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + { + kv_o[i] = static_cast(static_cast(kv_e[i])); + } + reinterpret_cast(&paged_kv[dkv])[eff_tid] = kv_out; - IoElemT const* sc_e = reinterpret_cast(&sc_raw); - IoVecT sc_out; - IoElemT* sc_o = reinterpret_cast(&sc_out); + KvScoreElemT const* sc_e = reinterpret_cast(&sc_raw); + StateVecT sc_out; + StateElemT* sc_o = reinterpret_cast(&sc_out); #pragma unroll for (int i = 0; i < VEC; i += 4) { float4 av = *reinterpret_cast(&ape[ape_idx * STATE_DIM + col + eff_tid * VEC + i]); - sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); - sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); - sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); - sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); + sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); + sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); + sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); + sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); } - reinterpret_cast(&paged_score[dsc])[eff_tid] = sc_out; + reinterpret_cast(&paged_score[dsc])[eff_tid] = sc_out; } } } @@ -446,8 +458,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int chunk_off_prev = prev_start % page_size; for (int r = my_r_start; r < my_r_end; r++) { - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, - phys_kv_prev, phys_sc_prev, chunk_off_prev + r, 0, eff_tid, rmax, rsum, rwsum); + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, + STATE_DIM, phys_kv_prev, phys_sc_prev, chunk_off_prev + r, 0, eff_tid, rmax, rsum, rwsum); } } else @@ -459,8 +471,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int blk_off = pos % page_size; int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, - phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, + STATE_DIM, phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); } } } @@ -473,7 +485,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int chunk_off_cur = curr_chunk_start % page_size; for (int r = my_r_start; r < my_r_end; r++) { - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, phys_kv_cur, phys_sc_cur, chunk_off_cur + r, HEAD_DIM, eff_tid, rmax, rsum, rwsum); } } @@ -486,7 +498,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int blk_off = pos % page_size; int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, phys_kv, phys_sc, blk_off, HEAD_DIM, eff_tid, rmax, rsum, rwsum); } } @@ -501,7 +513,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int chunk_off = curr_chunk_start % page_size; for (int r = my_r_start; r < my_r_end; r++) { - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, phys_kv, phys_sc, chunk_off + r, 0, eff_tid, rmax, rsum, rwsum); } } @@ -514,7 +526,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int blk_off = pos % page_size; int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, STATE_DIM, phys_kv, phys_sc, blk_off, 0, eff_tid, rmax, rsum, rwsum); } } @@ -599,38 +611,47 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo } // Explicit instantiations for decode kernel (NUM_RED_WARPS defaults to 1) -#define INST_DECODE(HD, EB, CR, NN, NRW) \ - template __global__ void pagedKvCompressKernel(void const*, float const*, void*, void*, \ - int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, int, \ - int, int); +#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW) \ + template __global__ void pagedKvCompressKernel(void const*, float const*, void*, \ + void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int); + +#define INST_DECODE_NN(HD, KV_EB, STATE_EB, CR) \ + INST_DECODE(HD, KV_EB, STATE_EB, CR, 1, 1) \ + INST_DECODE(HD, KV_EB, STATE_EB, CR, 2, 1) \ + INST_DECODE(HD, KV_EB, STATE_EB, CR, 3, 1) INST_DECODE(HD, KV_EB, STATE_EB, CR, 4, 1) -#define INST_DECODE_NN(HD, EB, CR) \ - INST_DECODE(HD, EB, CR, 1, 1) \ - INST_DECODE(HD, EB, CR, 2, 1) INST_DECODE(HD, EB, CR, 3, 1) INST_DECODE(HD, EB, CR, 4, 1) +#define INST_DECODE_DTYPES(HD, CR) \ + INST_DECODE_NN(HD, 2, 2, CR) \ + INST_DECODE_NN(HD, 2, 4, CR) INST_DECODE_NN(HD, 4, 2, CR) INST_DECODE_NN(HD, 4, 4, CR) -INST_DECODE_NN(128, 2, 4) -INST_DECODE_NN(128, 2, 128) -INST_DECODE_NN(128, 4, 4) -INST_DECODE_NN(128, 4, 128) -INST_DECODE_NN(512, 2, 4) -INST_DECODE_NN(512, 2, 128) -INST_DECODE_NN(512, 4, 4) -INST_DECODE_NN(512, 4, 128) +INST_DECODE_DTYPES(128, 4) +INST_DECODE_DTYPES(128, 128) +INST_DECODE_DTYPES(512, 4) +INST_DECODE_DTYPES(512, 128) // 4-warp parallel reduction variants for large compress_ratio. // HD=128: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. // HD=512 bf16: ELEM_PER_BLOCK=256, smem = 3*4*256*4 = 12 KB. // HD=512 fp32: ELEM_PER_BLOCK=128, smem = 3*4*128*4 = 6 KB. // NEXT_N=1 (single-token decode) and NEXT_N=2 (MTP speculative decode). -INST_DECODE(128, 2, 128, 1, 4) -INST_DECODE(128, 4, 128, 1, 4) -INST_DECODE(128, 2, 128, 2, 4) -INST_DECODE(128, 4, 128, 2, 4) -INST_DECODE(512, 2, 128, 1, 4) -INST_DECODE(512, 4, 128, 1, 4) -INST_DECODE(512, 2, 128, 2, 4) -INST_DECODE(512, 4, 128, 2, 4) - +INST_DECODE(128, 2, 2, 128, 1, 4) +INST_DECODE(128, 2, 4, 128, 1, 4) +INST_DECODE(128, 4, 2, 128, 1, 4) +INST_DECODE(128, 4, 4, 128, 1, 4) +INST_DECODE(128, 2, 2, 128, 2, 4) +INST_DECODE(128, 2, 4, 128, 2, 4) +INST_DECODE(128, 4, 2, 128, 2, 4) +INST_DECODE(128, 4, 4, 128, 2, 4) +INST_DECODE(512, 2, 2, 128, 1, 4) +INST_DECODE(512, 2, 4, 128, 1, 4) +INST_DECODE(512, 4, 2, 128, 1, 4) +INST_DECODE(512, 4, 4, 128, 1, 4) +INST_DECODE(512, 2, 2, 128, 2, 4) +INST_DECODE(512, 2, 4, 128, 2, 4) +INST_DECODE(512, 4, 2, 128, 2, 4) +INST_DECODE(512, 4, 4, 128, 2, 4) + +#undef INST_DECODE_DTYPES #undef INST_DECODE_NN #undef INST_DECODE @@ -644,22 +665,26 @@ INST_DECODE(512, 4, 128, 2, 4) // ============================================================================ // Forward declaration (defined in prefill section below). -static inline int prefillNthreads(int head_dim, int io_elem_bytes); +static inline int prefillNthreads(int head_dim, int elem_bytes_for_vec); void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, - int32_t const* start_pos, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, - int max_blocks, int head_dim, int compress_ratio, int next_n, int io_elem_bytes, int out_elem_bytes, + int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim, + int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, 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( + (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"); // Compute HEAD_BLOCKS: mirrors the compile-time constant in the kernel. // VEC = max_vec if HEAD_DIM/max_vec >= 32, else HEAD_DIM/32. // NTHRD_BASE = HEAD_DIM / VEC; HEAD_BLOCKS = NTHRD_BASE / 32 (or 1 if <= 32). // This spreads the head_dim across multiple blocks for better SM utilisation. - int const max_vec_elem = 16 / io_elem_bytes; + int const elem_bytes_for_vec = max(kv_score_elem_bytes, state_elem_bytes); + int const max_vec_elem = 16 / elem_bytes_for_vec; int const vec = (head_dim / max_vec_elem >= 32) ? max_vec_elem : (head_dim / 32); int const nthrd_base = head_dim / vec; // mirrors kernel's NTHRD_BASE = HEAD_DIM / VEC int const head_blocks = (nthrd_base > 32) ? (nthrd_base / 32) : 1; @@ -672,8 +697,8 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k // smem per block = 3 * MULTI_WARP * ELEM_PER_BLOCK * sizeof(float) // where ELEM_PER_BLOCK = nthreads_inner * vec = HEAD_DIM / HEAD_BLOCKS. // HD=128: ELEM_PER_BLOCK=128 → 6 KB. - // HD=512 bf16 (vec=8, HEAD_BLOCKS=2): ELEM_PER_BLOCK=256 → 12 KB. - // HD=512 fp32 (vec=4, HEAD_BLOCKS=4): ELEM_PER_BLOCK=128 → 6 KB. + // HD=512 with max elem size 2 (vec=8, HEAD_BLOCKS=2): ELEM_PER_BLOCK=256 → 12 KB. + // HD=512 with max elem size 4 (vec=4, HEAD_BLOCKS=4): ELEM_PER_BLOCK=128 → 6 KB. constexpr int MULTI_WARP = 4; bool const use_multi_warp = (compress_ratio == 128 && next_n <= 2); int const num_red_warps = use_multi_warp ? MULTI_WARP : 1; @@ -683,106 +708,89 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k dim3 grid(batch_size, head_blocks); -#define LAUNCH_DECODE(HD, EB, CR, NN) \ - pagedKvCompressKernel<<>>(kv_score, ape, paged_kv, \ - paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ - page_size, max_blocks, out_elem_bytes) +#define LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, NN) \ + pagedKvCompressKernel<<>>(kv_score, ape, \ + paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, page_size, \ + max_blocks, out_elem_bytes) -#define LAUNCH_DECODE_MW(HD, EB, CR, NN) \ - pagedKvCompressKernel<<>>(kv_score, ape, paged_kv, \ - paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ +#define LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, NN) \ + pagedKvCompressKernel<<>>(kv_score, \ + ape, paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ page_size, max_blocks, out_elem_bytes) -#define DISPATCH_NN_MW(HD, EB, CR) \ +#define DISPATCH_NN_MW(HD, KV_EB, STATE_EB, CR) \ switch (next_n) \ { \ - case 2: LAUNCH_DECODE_MW(HD, EB, CR, 2); break; \ - default: LAUNCH_DECODE_MW(HD, EB, CR, 1); break; \ + case 2: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 2); break; \ + default: LAUNCH_DECODE_MW(HD, KV_EB, STATE_EB, CR, 1); break; \ } -#define DISPATCH_NN(HD, EB, CR) \ +#define DISPATCH_NN(HD, KV_EB, STATE_EB, CR) \ switch (next_n) \ { \ - case 1: LAUNCH_DECODE(HD, EB, CR, 1); break; \ - case 2: LAUNCH_DECODE(HD, EB, CR, 2); break; \ - case 3: LAUNCH_DECODE(HD, EB, CR, 3); break; \ - default: LAUNCH_DECODE(HD, EB, CR, 4); break; \ + case 1: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 1); break; \ + case 2: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 2); break; \ + case 3: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 3); break; \ + default: LAUNCH_DECODE(HD, KV_EB, STATE_EB, CR, 4); break; \ } +#define DISPATCH_DTYPE(HD, CR, DISPATCH_MACRO) \ + do \ + { \ + if (kv_score_elem_bytes == 4 && state_elem_bytes == 4) \ + { \ + DISPATCH_MACRO(HD, 4, 4, CR); \ + } \ + else if (kv_score_elem_bytes == 2 && state_elem_bytes == 4) \ + { \ + DISPATCH_MACRO(HD, 2, 4, CR); \ + } \ + else if (kv_score_elem_bytes == 4 && state_elem_bytes == 2) \ + { \ + DISPATCH_MACRO(HD, 4, 2, CR); \ + } \ + else \ + { \ + DISPATCH_MACRO(HD, 2, 2, CR); \ + } \ + } while (false) + if (use_multi_warp) { // Multi-warp path: HD=128 or HD=512, NEXT_N=1 or NEXT_N=2, CR=128. - if (io_elem_bytes == 4) + if (head_dim == 512) { - if (head_dim == 512) - { - DISPATCH_NN_MW(512, 4, 128); - } - else - { - DISPATCH_NN_MW(128, 4, 128); - } + DISPATCH_DTYPE(512, 128, DISPATCH_NN_MW); } else { - if (head_dim == 512) - { - DISPATCH_NN_MW(512, 2, 128); - } - else - { - DISPATCH_NN_MW(128, 2, 128); - } + DISPATCH_DTYPE(128, 128, DISPATCH_NN_MW); } } else if (compress_ratio == 4) - { - if (io_elem_bytes == 4) - { - if (head_dim == 512) - { - DISPATCH_NN(512, 4, 4); - } - else - { - DISPATCH_NN(128, 4, 4); - } - } - else - { - if (head_dim == 512) - { - DISPATCH_NN(512, 2, 4); - } - else - { - DISPATCH_NN(128, 2, 4); - } - } - } - else if (io_elem_bytes == 4) { if (head_dim == 512) { - DISPATCH_NN(512, 4, 128); + DISPATCH_DTYPE(512, 4, DISPATCH_NN); } else { - DISPATCH_NN(128, 4, 128); + DISPATCH_DTYPE(128, 4, DISPATCH_NN); } } else { if (head_dim == 512) { - DISPATCH_NN(512, 2, 128); + DISPATCH_DTYPE(512, 128, DISPATCH_NN); } else { - DISPATCH_NN(128, 2, 128); + DISPATCH_DTYPE(128, 128, DISPATCH_NN); } } +#undef DISPATCH_DTYPE #undef DISPATCH_NN #undef DISPATCH_NN_MW #undef LAUNCH_DECODE_MW @@ -792,7 +800,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k // ============================================================================ // Prefill Kernel: prefillReductionKernel // -// Template: +// Template: // // Grid: (batch_size, max_outputs_per_batch) — one block per compressed output // Block: (NTHRD) where NTHRD = HEAD_DIM / VEC (>= 32 threads) @@ -816,24 +824,22 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k // Per-element online softmax step on VEC elements via 128-bit vectorized loads. // Reads directly from the kv_score input buffer (not paged state) since prefill // has the full sequence available. -template +template __device__ __forceinline__ void prefillSoftmaxVec(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, int64_t row_elem, // (input_offset + row_idx) * two_sd int kv_col_off, // column offset into kv_score row (0 or HEAD_DIM) int ape_base, // r * state_dim + ape_col_off int state_dim, int tid, float* __restrict__ rmax, float* __restrict__ rsum, float* __restrict__ rwsum) { - using IoElemT = typename std::conditional::type; - constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; - constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - using IoVecT = typename VecType::type; + using KvScoreElemT = typename std::conditional::type; + using KvScoreVecT = typename VecType::type; - auto const* kv = reinterpret_cast(kv_score_raw); + auto const* kv = reinterpret_cast(kv_score_raw); - IoVecT k_raw = reinterpret_cast(&kv[row_elem + kv_col_off])[tid]; - IoVecT s_raw = reinterpret_cast(&kv[row_elem + state_dim + kv_col_off])[tid]; - IoElemT const* ke = reinterpret_cast(&k_raw); - IoElemT const* se = reinterpret_cast(&s_raw); + KvScoreVecT k_raw = reinterpret_cast(&kv[row_elem + kv_col_off])[tid]; + KvScoreVecT s_raw = reinterpret_cast(&kv[row_elem + state_dim + kv_col_off])[tid]; + KvScoreElemT const* ke = reinterpret_cast(&k_raw); + KvScoreElemT const* se = reinterpret_cast(&s_raw); #pragma unroll for (int i = 0; i < VEC; i += 4) @@ -856,20 +862,24 @@ __device__ __forceinline__ void prefillSoftmaxVec(void const* __restrict__ kv_sc } } -template +template __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, int32_t const* __restrict__ start_pos_arr, int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int state_dim, int max_blocks, int out_elem_bytes) { - using IoElemT = typename std::conditional::type; + using KvScoreElemT = typename std::conditional::type; + using StateElemT = typename std::conditional::type; static_assert(COMPRESS_RATIO == 4 || COMPRESS_RATIO == 128, "Unsupported COMPRESS_RATIO"); constexpr bool IS_OVERLAP = (COMPRESS_RATIO == 4); - constexpr int MAX_VEC = 16 / IO_ELEM_BYTES; + constexpr int ELEM_BYTES_FOR_VEC + = (KV_SCORE_ELEM_BYTES > STATE_ELEM_BYTES) ? KV_SCORE_ELEM_BYTES : STATE_ELEM_BYTES; + constexpr int MAX_VEC = 16 / ELEM_BYTES_FOR_VEC; constexpr int VEC = (HEAD_DIM / MAX_VEC >= 32) ? MAX_VEC : (HEAD_DIM / 32); - using IoVecT = typename VecType::type; + using KvScoreVecT = typename VecType::type; + using StateVecT = typename VecType::type; static_assert(VEC >= 4, "VEC must be >= 4 for float4 ape loads"); int const tid = threadIdx.x; @@ -900,9 +910,9 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl int const abs_idx = first_abs_idx + local_output_idx; int const win_start = abs_idx * COMPRESS_RATIO; - auto const* kv_score = reinterpret_cast(kv_score_raw); - auto* paged_kv = reinterpret_cast(paged_kv_raw); - auto* paged_score = reinterpret_cast(paged_score_raw); + auto const* kv_score = reinterpret_cast(kv_score_raw); + auto* paged_kv = reinterpret_cast(paged_kv_raw); + auto* paged_score = reinterpret_cast(paged_score_raw); int64_t const two_sd = 2 * state_dim; int64_t const page_sd = static_cast(page_size) * state_dim; @@ -939,23 +949,32 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl int64_t const dkv = static_cast(phys_kv) * page_sd + blk_off * state_dim + col; int64_t const dsc = static_cast(phys_sc) * page_sd + blk_off * state_dim + col; - IoVecT kv_raw = reinterpret_cast(&kv_score[src])[tid]; - IoVecT sc_raw = reinterpret_cast(&kv_score[src + state_dim])[tid]; - reinterpret_cast(&paged_kv[dkv])[tid] = kv_raw; + KvScoreVecT kv_raw = reinterpret_cast(&kv_score[src])[tid]; + KvScoreVecT sc_raw = reinterpret_cast(&kv_score[src + state_dim])[tid]; + + KvScoreElemT const* kv_e = reinterpret_cast(&kv_raw); + StateVecT kv_out; + StateElemT* kv_o = reinterpret_cast(&kv_out); +#pragma unroll + for (int i = 0; i < VEC; i++) + { + kv_o[i] = static_cast(static_cast(kv_e[i])); + } + reinterpret_cast(&paged_kv[dkv])[tid] = kv_out; - IoElemT const* sc_e = reinterpret_cast(&sc_raw); - IoVecT sc_out; - IoElemT* sc_o = reinterpret_cast(&sc_out); + KvScoreElemT const* sc_e = reinterpret_cast(&sc_raw); + StateVecT sc_out; + StateElemT* sc_o = reinterpret_cast(&sc_out); #pragma unroll for (int i = 0; i < VEC; i += 4) { float4 av = *reinterpret_cast(&ape[r * state_dim + col + tid * VEC + i]); - sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); - sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); - sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); - sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); + sc_o[i] = static_cast(static_cast(sc_e[i]) + av.x); + sc_o[i + 1] = static_cast(static_cast(sc_e[i + 1]) + av.y); + sc_o[i + 2] = static_cast(static_cast(sc_e[i + 2]) + av.z); + sc_o[i + 3] = static_cast(static_cast(sc_e[i + 3]) + av.w); } - reinterpret_cast(&paged_score[dsc])[tid] = sc_out; + reinterpret_cast(&paged_score[dsc])[tid] = sc_out; } } }; @@ -1022,8 +1041,8 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl int blk_off = pos % page_size; int phys_kv = block_table_kv[batch_idx * max_blocks + log_blk]; int phys_sc = block_table_score[batch_idx * max_blocks + log_blk]; - decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, state_dim, phys_kv, - phys_sc, blk_off, kv_col_off, tid, rmax, rsum, rwsum); + decodeSoftmaxVec(paged_kv_raw, paged_score_raw, page_sd, state_dim, + phys_kv, phys_sc, blk_off, kv_col_off, tid, rmax, rsum, rwsum); } // New-token portion — read from kv_score and add APE live. @@ -1032,7 +1051,7 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl int const pos = range_start + r; int const input_row = pos - sp; int64_t const row = static_cast(input_offset + input_row) * two_sd; - prefillSoftmaxVec( + prefillSoftmaxVec( kv_score_raw, ape, row, kv_col_off, r * state_dim + ape_col_off, state_dim, tid, rmax, rsum, rwsum); } }; @@ -1083,19 +1102,20 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl } // Explicit instantiations -#define INST_PREFILL(HD, EB, CR) \ - template __global__ void prefillReductionKernel(void const*, float const*, void*, void*, \ +#define INST_PREFILL(HD, KV_EB, STATE_EB, CR) \ + template __global__ void prefillReductionKernel(void const*, float const*, void*, void*, \ int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, int, \ int, int, int); -INST_PREFILL(128, 2, 4) -INST_PREFILL(128, 2, 128) -INST_PREFILL(128, 4, 4) -INST_PREFILL(128, 4, 128) -INST_PREFILL(512, 2, 4) -INST_PREFILL(512, 2, 128) -INST_PREFILL(512, 4, 4) -INST_PREFILL(512, 4, 128) +#define INST_PREFILL_DTYPES(HD, CR) \ + INST_PREFILL(HD, 2, 2, CR) \ + INST_PREFILL(HD, 2, 4, CR) INST_PREFILL(HD, 4, 2, CR) INST_PREFILL(HD, 4, 4, CR) + +INST_PREFILL_DTYPES(128, 4) +INST_PREFILL_DTYPES(128, 128) +INST_PREFILL_DTYPES(512, 4) +INST_PREFILL_DTYPES(512, 128) +#undef INST_PREFILL_DTYPES #undef INST_PREFILL // ============================================================================ @@ -1106,9 +1126,9 @@ INST_PREFILL(512, 4, 128) // ============================================================================ // Compute threads per block: mirrors compile-time NTHRD = HEAD_DIM / VEC. -static inline int prefillNthreads(int head_dim, int io_elem_bytes) +static inline int prefillNthreads(int head_dim, int elem_bytes_for_vec) { - int max_vec = 16 / io_elem_bytes; + int max_vec = 16 / elem_bytes_for_vec; int vec = (head_dim / max_vec >= 32) ? max_vec : (head_dim / 32); return head_dim / vec; } @@ -1116,57 +1136,63 @@ static inline int prefillNthreads(int head_dim, int io_elem_bytes) void prefillReductionLaunch(void const* kv_score, float const* ape, void* paged_kv, void* paged_score, int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, int32_t const* start_pos, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, - int max_blocks, int head_dim, int compress_ratio, int max_outputs, int io_elem_bytes, int out_elem_bytes, - cudaStream_t stream) + int max_blocks, int head_dim, int compress_ratio, int max_outputs, int kv_score_elem_bytes, int state_elem_bytes, + 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"); - int const nthreads = prefillNthreads(head_dim, io_elem_bytes); + 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"); + int const elem_bytes_for_vec = max(kv_score_elem_bytes, state_elem_bytes); + int const nthreads = prefillNthreads(head_dim, elem_bytes_for_vec); int const coff = overlap ? 2 : 1; int const state_dim = coff * head_dim; dim3 grid(batch_size, max(max_outputs, 1)); -#define LAUNCH_PREFILL(HD, EB, CR) \ - prefillReductionKernel<<>>(kv_score, ape, paged_kv, paged_score, \ - block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, page_size, state_dim, \ - max_blocks, out_elem_bytes) +#define LAUNCH_PREFILL(HD, KV_EB, STATE_EB, CR) \ + prefillReductionKernel<<>>(kv_score, ape, paged_kv, \ + paged_score, block_table_kv, block_table_score, output, kv_lens, start_pos, cu_seq_lens, cu_kv_comp, \ + page_size, state_dim, max_blocks, out_elem_bytes) + +#define DISPATCH_PREFILL_DTYPE(HD, CR) \ + do \ + { \ + if (kv_score_elem_bytes == 4 && state_elem_bytes == 4) \ + { \ + LAUNCH_PREFILL(HD, 4, 4, CR); \ + } \ + else if (kv_score_elem_bytes == 2 && state_elem_bytes == 4) \ + { \ + LAUNCH_PREFILL(HD, 2, 4, CR); \ + } \ + else if (kv_score_elem_bytes == 4 && state_elem_bytes == 2) \ + { \ + LAUNCH_PREFILL(HD, 4, 2, CR); \ + } \ + else \ + { \ + LAUNCH_PREFILL(HD, 2, 2, CR); \ + } \ + } while (false) - if (io_elem_bytes == 4) + if (head_dim == 512) { - if (head_dim == 512) - { - if (compress_ratio == 4) - LAUNCH_PREFILL(512, 4, 4); - else - LAUNCH_PREFILL(512, 4, 128); - } + if (compress_ratio == 4) + DISPATCH_PREFILL_DTYPE(512, 4); else - { - if (compress_ratio == 4) - LAUNCH_PREFILL(128, 4, 4); - else - LAUNCH_PREFILL(128, 4, 128); - } + DISPATCH_PREFILL_DTYPE(512, 128); } else { - if (head_dim == 512) - { - if (compress_ratio == 4) - LAUNCH_PREFILL(512, 2, 4); - else - LAUNCH_PREFILL(512, 2, 128); - } + if (compress_ratio == 4) + DISPATCH_PREFILL_DTYPE(128, 4); else - { - if (compress_ratio == 4) - LAUNCH_PREFILL(128, 2, 4); - else - LAUNCH_PREFILL(128, 2, 128); - } + DISPATCH_PREFILL_DTYPE(128, 128); } +#undef DISPATCH_PREFILL_DTYPE #undef LAUNCH_PREFILL } diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h index 2f5ea155447e..ca8424f64fc8 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h @@ -41,11 +41,11 @@ void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or f int32_t const* block_table_score, // [bsz, max_blocks] void* output, // [total_outputs, head_dim] int32_t const* kv_lens, // [bsz] - int32_t const* start_pos, // [bsz] int32_t const* cu_seq_lens, // [bsz+1] int32_t const* cu_kv_comp, // [bsz+1] int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int next_n, - int io_elem_bytes, // bytes per element for kv_score/paged (2=bf16, 4=fp32) + int kv_score_elem_bytes, // bytes per element for kv_score (2=bf16, 4=fp32) + int state_elem_bytes, // bytes per element for paged state (2=bf16, 4=fp32) int out_elem_bytes, // bytes per element for output cudaStream_t stream); @@ -65,8 +65,8 @@ void prefillReductionLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or int32_t const* start_pos, // [bsz] int32_t const* cu_seq_lens, // [bsz+1] int32_t const* cu_kv_comp, // [bsz+1] - int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int max_outputs, int io_elem_bytes, - int out_elem_bytes, cudaStream_t stream); + int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int max_outputs, + int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, cudaStream_t stream); // RMSNorm + RoPE + Hadamard + paged scatter in a single kernel launch. // Optionally writes postprocessed result to kv_out (nullptr to skip). diff --git a/cpp/tensorrt_llm/thop/compressorOp.cpp b/cpp/tensorrt_llm/thop/compressorOp.cpp index b73fbd4469f8..4c57e47a15aa 100644 --- a/cpp/tensorrt_llm/thop/compressorOp.cpp +++ b/cpp/tensorrt_llm/thop/compressorOp.cpp @@ -26,29 +26,31 @@ namespace { // Decode kernel: write tokens to paged cache + conditional compression -void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf16 +void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf16/fp32 torch::Tensor ape, // [compress_ratio, state_dim] fp32 - torch::Tensor paged_kv, // [num_blocks, page_size, state_dim] bf16 - torch::Tensor paged_score, // [num_blocks, page_size, state_dim] bf16 + torch::Tensor paged_kv, // [num_blocks, page_size, state_dim] bf16/fp32 + torch::Tensor paged_score, // [num_blocks, page_size, state_dim] bf16/fp32 torch::Tensor block_table_kv, // [bsz, max_blocks] int32 torch::Tensor block_table_score, // [bsz, max_blocks] int32 torch::Tensor output, // [total_outputs, head_dim] bf16 torch::Tensor kv_lens, // [bsz] int32 - torch::Tensor start_pos, // [bsz] int32 torch::Tensor cu_seq_lens, // [bsz+1] int32 torch::Tensor cu_kv_comp, // [bsz+1] int32 int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t next_n) { auto stream = at::cuda::getCurrentCUDAStream(); - int io_eb = static_cast(kv_score.element_size()); + int kv_score_eb = static_cast(kv_score.element_size()); + int state_eb = static_cast(paged_kv.element_size()); int out_eb = static_cast(output.element_size()); + TORCH_CHECK( + paged_score.element_size() == paged_kv.element_size(), "paged_kv and paged_score must use the same dtype"); tk::pagedKvCompressLaunch(kv_score.data_ptr(), ape.data_ptr(), paged_kv.data_ptr(), paged_score.data_ptr(), block_table_kv.data_ptr(), block_table_score.data_ptr(), output.data_ptr(), - kv_lens.data_ptr(), start_pos.data_ptr(), cu_seq_lens.data_ptr(), - cu_kv_comp.data_ptr(), static_cast(batch_size), static_cast(page_size), - static_cast(block_table_kv.size(1)), static_cast(head_dim), static_cast(compress_ratio), - static_cast(next_n), io_eb, out_eb, stream); + kv_lens.data_ptr(), cu_seq_lens.data_ptr(), cu_kv_comp.data_ptr(), + static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), + static_cast(head_dim), static_cast(compress_ratio), static_cast(next_n), kv_score_eb, state_eb, + out_eb, stream); } // Prefill kernel: bulk compression with state update @@ -58,15 +60,18 @@ void compressorPrefillReductionOp(torch::Tensor kv_score, torch::Tensor ape, tor int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t max_outputs) { auto stream = at::cuda::getCurrentCUDAStream(); - int io_eb = static_cast(kv_score.element_size()); + int kv_score_eb = static_cast(kv_score.element_size()); + int state_eb = static_cast(paged_kv.element_size()); int out_eb = static_cast(output.element_size()); + TORCH_CHECK( + paged_score.element_size() == paged_kv.element_size(), "paged_kv and paged_score must use the same dtype"); tk::prefillReductionLaunch(kv_score.data_ptr(), ape.data_ptr(), paged_kv.data_ptr(), paged_score.data_ptr(), block_table_kv.data_ptr(), block_table_score.data_ptr(), output.data_ptr(), kv_lens.data_ptr(), start_pos.data_ptr(), cu_seq_lens.data_ptr(), cu_kv_comp.data_ptr(), static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), static_cast(head_dim), static_cast(compress_ratio), - static_cast(max_outputs), io_eb, out_eb, stream); + static_cast(max_outputs), kv_score_eb, state_eb, out_eb, stream); } // Fused postprocess + scatter: RMSNorm + RoPE + Hadamard + paged scatter in one kernel @@ -122,7 +127,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor(a!) paged_kv, Tensor(b!) paged_score, " "Tensor block_table_kv, Tensor block_table_score, " "Tensor(c!) output, " - "Tensor kv_lens, Tensor start_pos, " + "Tensor kv_lens, " "Tensor cu_seq_lens, Tensor cu_kv_comp, " "int batch_size, int page_size, " "int head_dim, int compress_ratio, " diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py index e382373a491d..2ee728a6bfbc 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -183,9 +183,10 @@ def forward( num_comp_tokens = metadata.new_comp_kv_lens_cuda[self.compress_ratio][:bsz] max_ctx_comp_kv_lens = metadata.max_ctx_compressed_tokens[self.compress_ratio] - # Project input to KV and score in the checkpoint dtype, then store the - # result in fp32 for the compressor kernels. - kv_score = F.linear(x.to(self.wkv_gate.weight.dtype), self.wkv_gate.weight).float() + # Project input to KV and score in the checkpoint dtype. The compressor + # kernels accept bf16 or fp32 kv_score and convert values to fp32 + # internally for state updates and online-softmax accumulation. + kv_score = F.linear(x.to(self.wkv_gate.weight.dtype), self.wkv_gate.weight) # Allocate output buffer kv_comp = torch.empty(total_num_comp_tokens, self.head_dim, device=x.device, dtype=x.dtype) @@ -225,7 +226,6 @@ def forward( block_table_score_state[num_contexts:], kv_comp, gen_kv_lens, - gen_kv_lens - next_n, metadata.cu_seq_lens_cuda[num_contexts:], cu_new_comp_kv[num_contexts:], num_generations, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py index 02cba87f3f5e..6255cf06b0dd 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_kernel.py @@ -110,8 +110,6 @@ def decode_kernel( ) if block_table_score is None: block_table_score = block_table_kv - if start_pos is None: - start_pos = kv_lens - next_n torch.ops.trtllm.compressor_paged_kv_compress( kv_score, @@ -122,7 +120,6 @@ def decode_kernel( block_table_score, kv_comp, kv_lens, - start_pos, cu_seq_lens, cu_new_comp_kv, kv_lens.shape[0], @@ -560,6 +557,168 @@ def test_decode_corner_cases(batch_size, compress_ratio, head_dim, overlap, num_ ), f"Step {step}, Batch {b}: mismatch diff={(diff).abs().max():.6f}" +def test_prefill_accepts_bf16_kv_score_with_fp32_state(): + """Prefill accepts bf16 kv_score while preserving fp32 compressor state.""" + batch_size, seqlen, compress_ratio, head_dim, overlap = 1, 4, 4, 128, True + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + + kv_bf16 = torch.randn(batch_size, seqlen, state_dim, device="cuda").bfloat16() + score_bf16 = torch.randn(batch_size, seqlen, state_dim, device="cuda").bfloat16() + kv = kv_bf16.float() + score = score_bf16.float() + ape = torch.randn(compress_ratio, state_dim, device="cuda") + + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + out_py = run_pytorch_prefill_reference( + kv.clone(), + score.clone(), + ape, + kv_state_py, + score_state_py, + compress_ratio, + head_dim, + overlap, + ) + + paged_kv, paged_score, block_table, page_size, _ = create_paged_cache( + batch_size, seqlen, compress_ratio, head_dim, overlap, page_size=8 + ) + assert paged_kv.dtype == torch.float32 + + kv_score = fuse_kv_score(kv_bf16.view(-1, state_dim), score_bf16.view(-1, state_dim)) + assert kv_score.dtype == torch.bfloat16 + + kv_lens = torch.full((batch_size,), seqlen, device="cuda", dtype=torch.int32) + start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + cu_seq_lens, cu_outputs = prepare_prefill_metadata( + kv_lens, start_pos, compress_ratio, head_dim, kv_score.device + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, kv_score.device, torch.bfloat16 + ) + + prefill_kernel( + kv_score, + ape, + kv_lens, + start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + max_outputs=1, + block_table_score=block_table, + compress_ratio=compress_ratio, + head_dim=head_dim, + page_size=page_size, + ) + + assert out_py is not None + assert torch.allclose(out_py.to(kv_comp.dtype).view_as(kv_comp), kv_comp, rtol=2e-3, atol=5e-3) + + for p in range(seqlen): + phys = block_table[0, p // page_size].item() + off = p % page_size + assert torch.allclose(paged_kv[phys, off], kv[0, p], atol=1e-5) + assert torch.allclose( + paged_score[phys, off], score[0, p] + ape[p % compress_ratio], atol=1e-5 + ) + + +def test_decode_accepts_bf16_kv_score_with_fp32_state(): + """Decode accepts bf16 kv_score and computes start_pos as kv_len - next_n.""" + batch_size, compress_ratio, head_dim, overlap = 1, 4, 128, True + coff = 2 if overlap else 1 + state_len, state_dim = coff * compress_ratio, coff * head_dim + page_size = 8 + decode_steps = 4 + max_blocks = 1 + + ape = torch.randn(compress_ratio, state_dim, device="cuda") + paged_kv = torch.zeros(batch_size * max_blocks, page_size, state_dim, device="cuda") + paged_score = torch.zeros_like(paged_kv) + assert paged_kv.dtype == torch.float32 + block_table = torch.zeros(batch_size, max_blocks, device="cuda", dtype=torch.int32) + + kv_state_py = torch.zeros(batch_size, state_len, state_dim, device="cuda") + score_state_py = torch.full((batch_size, state_len, state_dim), float("-inf"), device="cuda") + init_kv = torch.randn(compress_ratio, state_dim, device="cuda") + init_score = torch.randn(compress_ratio, state_dim, device="cuda") + kv_state_py[:, :compress_ratio] = init_kv + score_state_py[:, :compress_ratio] = init_score + paged_kv[0, :compress_ratio] = init_kv + paged_score[0, :compress_ratio] = init_score + + cu_seq_lens, cu_outputs = prepare_decode_metadata( + batch_size, compress_ratio, head_dim, torch.device("cuda"), next_n=1 + ) + kv_comp = prepare_compress_output( + cu_outputs, batch_size, head_dim, torch.device("cuda"), torch.bfloat16 + ) + + for step in range(decode_steps): + token_idx = compress_ratio + step + total_tokens = token_idx + 1 + new_kv_bf16 = torch.randn(batch_size, state_dim, device="cuda").bfloat16() + new_score_bf16 = torch.randn(batch_size, state_dim, device="cuda").bfloat16() + new_kv = new_kv_bf16.float() + new_score = new_score_bf16.float() + + out_py = run_pytorch_reference( + new_kv.unsqueeze(1), + new_score.unsqueeze(1), + ape, + kv_state_py, + score_state_py, + token_idx, + compress_ratio, + head_dim, + overlap, + ) + + kv_score = fuse_kv_score(new_kv_bf16, new_score_bf16) + assert kv_score.dtype == torch.bfloat16 + kv_lens = torch.full((batch_size,), total_tokens, device="cuda", dtype=torch.int32) + stale_start_pos = torch.zeros(batch_size, device="cuda", dtype=torch.int32) + + decode_kernel( + kv_score, + ape, + kv_lens, + stale_start_pos, + cu_seq_lens, + cu_outputs, + kv_comp, + paged_kv, + paged_score, + block_table, + block_table, + compress_ratio, + head_dim, + page_size, + next_n=1, + ) + + phys = block_table[0, token_idx // page_size].item() + off = token_idx % page_size + assert torch.allclose(paged_kv[phys, off], new_kv[0], atol=1e-5) + assert torch.allclose( + paged_score[phys, off], new_score[0] + ape[token_idx % compress_ratio], atol=1e-5 + ) + + if out_py is not None: + assert torch.allclose( + out_py[0, 0, :head_dim].to(kv_comp.dtype), + kv_comp[0], + rtol=2e-3, + atol=5e-3, + ) + + STATE_UPDATE_CONFIGS = [ pytest.param(1, 12, 4, 128, True, id="overlap_hd128_3chunks"), pytest.param(1, 13, 4, 512, True, id="overlap_hd512_3chunks_1remainder"), diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index d95d397cbea8..f277ba15d342 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -1675,6 +1675,7 @@ def _run_compressor_with_fake_postprocess(monkeypatch, kv_cache_dtype: str, is_i seen = {} def fake_prefill_reduction(*args): + seen["kv_score_dtype"] = args[0].dtype kv_comp = args[6] kv_comp.fill_(0.25) @@ -1751,6 +1752,7 @@ def test_main_compressor_does_not_materialize_postprocess_output(monkeypatch): ) assert seen["cache_dtype"] == int(KVCacheDtype.NONE) + assert seen["kv_score_dtype"] == torch.bfloat16 assert seen["kv_out"] is None assert seen["quant_output"] is None assert seen["scale_output"] is None diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py index 2cfd278d4fa0..b052cee0171c 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py @@ -4,7 +4,7 @@ """Tests for the BF16 GEMM path in the DeepSeek-V4 Compressor wkv_gate. The wkv_gate weight now matches the upstream V4 checkpoint dtype (bf16). The -GEMM runs in that dtype, then stores the result in fp32 for compressor kernels. +GEMM runs in that dtype and the compressor kernels consume bf16 kv_score input. """ import pytest @@ -37,9 +37,8 @@ def _create_wkv_gate(dtype: torch.dtype = torch.bfloat16) -> Linear: def _bf16_path(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: - """Default path: bf16 GEMM via F.linear, then upcast to fp32 for the - downstream compression kernels.""" - return F.linear(x.to(wkv_gate.weight.dtype), wkv_gate.weight).float() + """Default path: bf16 GEMM via F.linear.""" + return F.linear(x.to(wkv_gate.weight.dtype), wkv_gate.weight) def _fp32_reference(x: torch.Tensor, wkv_gate: Linear) -> torch.Tensor: @@ -75,9 +74,9 @@ def test_bf16_path_close_to_fp32_reference(num_tokens): out_fp32 = _fp32_reference(x, fp32_gate) assert out_bf16.shape == out_fp32.shape - assert out_bf16.dtype == torch.float32 + assert out_bf16.dtype == torch.bfloat16 - a, b = out_bf16.flatten(), out_fp32.flatten() + a, b = out_bf16.float().flatten(), out_fp32.flatten() cos_sim = F.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item() assert cos_sim >= 0.99, f"cos_sim={cos_sim:.6f}" From f7d4c5f3389cf2b9577a3f8498ae2f98a007d40b Mon Sep 17 00:00:00 2001 From: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Date: Wed, 6 May 2026 21:22:28 -0700 Subject: [PATCH 21/58] [TRTLLM-12478][fix] Fix IMA about KVCacheManagerV2's overlap scheduler mode Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit 044c13c41eeac164e40b087208662dab0b1b9e90) Signed-off-by: Fanrong Li --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 +++- .../kv_cache_manager_v2/_core/_kv_cache.py | 23 ++++++++++++++++++- 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bef5f5755fcb..2c39af8bc5bc 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2929,6 +2929,7 @@ def _executor_loop(self): can_forward, should_retry = self._check_benchmark_disagg_gate( scheduled_batch, can_forward) if should_retry: + self._revert_gen_alloc(scheduled_batch) continue if not self._is_kv_manager_v2: @@ -3328,6 +3329,7 @@ def _executor_loop_overlap(self): can_forward, should_retry = self._check_benchmark_disagg_gate( scheduled_batch, can_forward) if should_retry: + self._revert_gen_alloc(scheduled_batch) continue if not self._is_kv_manager_v2: @@ -3505,7 +3507,8 @@ def _executor_loop_overlap(self): # blocks freed by this chunk are visible to the next # iteration's scheduler. # Only applies to KV cache manager V2 + scheduler V2. - if self._scheduler_manages_kv_suspend: + if (self._scheduler_manages_kv_suspend + and scheduled_batch.context_requests): self.kv_cache_manager.update_context_resources( scheduled_batch) 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 057f6a0ce32b..fa61125b22a6 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 @@ -619,6 +619,8 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo slot = slots[lc].pop() self._scratch_slots[lc].append(ScratchSlotLock(slot, self, lc, skip_wait=True)) + normal_slots = tuple(chain.from_iterable(slots)) + new_slot_ready_events = tuple(slot.ready_event for slot in normal_slots) for beam_indices in self._base_page_indices: for indices in beam_indices: if type(indices) is array.array: @@ -627,6 +629,23 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo else: if len(indices) < new_num_blocks: raise ValueError("User-provided base page indices is too short") + if any( + type(indices) is memoryview + for beam_indices in self._base_page_indices + for indices in beam_indices + ): + # User-provided base page indices are host-visible. Do not publish + # a recycled slot id there until the slot is actually ready; a + # stream wait only orders GPU work and does not protect CPU-side + # metadata readers. + pending_ready_events = [ + event for event in set(new_slot_ready_events) if not event.query_complete() + ] + for event in pending_ready_events: + event.synchronize() + for slot in normal_slots: + slot.ready_event = CachedCudaEvent.NULL + stream_wait_events(self.cuda_stream, new_slot_ready_events) for ordinal in typed_range(old_num_blocks, new_num_blocks): block = make_typed( lambda _: filled_list(cast(BlockPage, None), num_life_cycles), beam_width @@ -644,7 +663,9 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo if stale_beg <= ordinal < stale_end: continue slot = slots[lc].pop() - # We have already waited for ready_event of the slots. + # Slot ready events are already handled above: either + # synchronized for host-visible page-index buffers, or + # enqueued as stream waits for internal buffers. block[beam_index][lc] = UncommittedPage( self, ordinal, lc, GPU_LEVEL, slot, beam_index ).lock(self, beam_index, ordinal, lc, skip_wait=True) From d1c27f5fc4ab2a17bfaea7df0ae068d213f2f699 Mon Sep 17 00:00:00 2001 From: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com> Date: Tue, 5 May 2026 11:13:08 -0700 Subject: [PATCH 22/58] [None][fix] Fix fused MHC for DeepSeek-V4-Pro hidden size (#13771) Signed-off-by: Oseltamivir Signed-off-by: Mingyang Hao Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com> Co-authored-by: Oseltamivir Co-authored-by: Mingyang Hao (cherry picked from commit 9aa37159f8958f16be6e97ddd76f6925e99d2489) Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> Signed-off-by: Fanrong Li (cherry picked from commit e62b6c26e53916db302c8e11e66995766a0c98b3) Signed-off-by: Fanrong Li --- .../kernels/mhcKernels/mhcFusedHcKernel.cu | 215 +++++++++++++++--- .../kernels/mhcKernels/mhcKernels.h | 10 +- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 62 +++-- tests/unittest/_torch/modules/test_mhc.py | 109 +++++++-- 4 files changed, 320 insertions(+), 76 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index 334961bdfced..7d5d2b6daf45 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -91,9 +91,11 @@ inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint } // namespace // ---- mHC fused kernel shape constants (mirrors the Python module) ---- -static constexpr uint32_t FHC_SHAPE_N = 24; // HC_MULT * (2 + HC_MULT) = 4 * 6 = 24 -static constexpr uint32_t FHC_HIDDEN = 4096; // only this hidden size is currently wired up +// HC_MULT * (2 + HC_MULT) = 4 * 6 = 24. +static constexpr uint32_t FHC_SHAPE_N = 24; static constexpr uint32_t FHC_HC_MULT = 4; +static constexpr uint32_t FHC_HIDDEN_FLASH = 4096; +static constexpr uint32_t FHC_HIDDEN_PRO = 7168; static constexpr uint32_t FHC_BLOCK_M = 64; static constexpr uint32_t FHC_BLOCK_N = 32; static constexpr uint32_t FHC_BLOCK_K = 64; @@ -103,6 +105,38 @@ static constexpr uint32_t FHC_N_INPUT_STG = 2; static constexpr uint32_t FHC_NUM_MMA_TH = 128; static constexpr uint32_t FHC_NUM_PMAP_TH = 128; +template